From 626ddad8341d44d849318c83499b064f133bfcad Mon Sep 17 00:00:00 2001 From: Daphne Date: Mon, 6 Jul 2026 18:14:50 -0400 Subject: [PATCH] feat: DeepSeek V4 Flash DSpark C12 NVFP4 recipe + mod spark-vllm-docker v1 style recipe for serving DeepSeek-V4-Flash-DSpark on 2x DGX Spark (GB10) with vLLM TP=2, DSpark speculative decoding, NVFP4 KV cache (nvfp4_ds_mla), B12X MoE backend, and the 2026-07-03 garble fix. Recipe: recipes/deepseek-v4-flash-dspark.yaml Mod: mods/deepseek-v4-flash-dspark/ (18 overlay files + 3 NVFP4 patches) Based on tonyd2wild C12 NVFP4 default config (2026-07-04). Includes Keys DSpark concurrency patch via overlay files. NVFP4 KV cache implementation via 3-stage patch pipeline. Default: gpu_memory_utilization=0.77, max_model_len=350000, nvfp4_ds_mla. Distributed executor backend: mp (multiprocessing). --- .../overlay/vllm/config/speculative.py | 1117 +++ .../overlay/vllm/envs.py | 2314 +++++ .../layers/fused_moe/b12x_moe.py | 785 ++ .../vllm/model_executor/models/registry.py | 1412 +++ .../model_executor/warmup/kernel_warmup.py | 851 ++ .../vllm/models/deepseek_v4/__init__.py | 32 + .../deepseek_v4/common/ops/cache_utils.py | 623 ++ .../vllm/models/deepseek_v4/nvidia/dspark.py | 1192 +++ .../deepseek_v4/nvidia/dspark_kernels.py | 822 ++ .../vllm/models/deepseek_v4/nvidia/model.py | 2187 +++++ .../vllm/models/deepseek_v4/nvidia/sm120.py | 519 ++ .../attention/backends/mla/b12x_mla_sparse.py | 501 ++ .../vllm/v1/attention/backends/registry.py | 264 + .../overlay/vllm/v1/core/sched/scheduler.py | 2369 +++++ .../overlay/vllm/v1/outputs.py | 339 + .../overlay/vllm/v1/spec_decode/dspark.py | 683 ++ .../vllm/v1/spec_decode/dspark_proposer.py | 1021 +++ .../vllm/v1/worker/gpu_model_runner.py | 7728 +++++++++++++++++ .../patch-nvfp4-stage-a.py | 46 + .../patch-nvfp4-stage-b.py | 34 + .../patch-nvfp4-stage-c.py | 34 + mods/deepseek-v4-flash-dspark/run.sh | 72 + recipes/deepseek-v4-flash-dspark.yaml | 134 + 23 files changed, 25079 insertions(+) create mode 100644 mods/deepseek-v4-flash-dspark/overlay/vllm/config/speculative.py create mode 100755 mods/deepseek-v4-flash-dspark/overlay/vllm/envs.py create mode 100644 mods/deepseek-v4-flash-dspark/overlay/vllm/model_executor/layers/fused_moe/b12x_moe.py create mode 100644 mods/deepseek-v4-flash-dspark/overlay/vllm/model_executor/models/registry.py create mode 100644 mods/deepseek-v4-flash-dspark/overlay/vllm/model_executor/warmup/kernel_warmup.py create mode 100644 mods/deepseek-v4-flash-dspark/overlay/vllm/models/deepseek_v4/__init__.py create mode 100644 mods/deepseek-v4-flash-dspark/overlay/vllm/models/deepseek_v4/common/ops/cache_utils.py create mode 100644 mods/deepseek-v4-flash-dspark/overlay/vllm/models/deepseek_v4/nvidia/dspark.py create mode 100644 mods/deepseek-v4-flash-dspark/overlay/vllm/models/deepseek_v4/nvidia/dspark_kernels.py create mode 100644 mods/deepseek-v4-flash-dspark/overlay/vllm/models/deepseek_v4/nvidia/model.py create mode 100644 mods/deepseek-v4-flash-dspark/overlay/vllm/models/deepseek_v4/nvidia/sm120.py create mode 100644 mods/deepseek-v4-flash-dspark/overlay/vllm/v1/attention/backends/mla/b12x_mla_sparse.py create mode 100644 mods/deepseek-v4-flash-dspark/overlay/vllm/v1/attention/backends/registry.py create mode 100644 mods/deepseek-v4-flash-dspark/overlay/vllm/v1/core/sched/scheduler.py create mode 100644 mods/deepseek-v4-flash-dspark/overlay/vllm/v1/outputs.py create mode 100644 mods/deepseek-v4-flash-dspark/overlay/vllm/v1/spec_decode/dspark.py create mode 100644 mods/deepseek-v4-flash-dspark/overlay/vllm/v1/spec_decode/dspark_proposer.py create mode 100644 mods/deepseek-v4-flash-dspark/overlay/vllm/v1/worker/gpu_model_runner.py create mode 100755 mods/deepseek-v4-flash-dspark/patch-nvfp4-stage-a.py create mode 100755 mods/deepseek-v4-flash-dspark/patch-nvfp4-stage-b.py create mode 100755 mods/deepseek-v4-flash-dspark/patch-nvfp4-stage-c.py create mode 100755 mods/deepseek-v4-flash-dspark/run.sh create mode 100644 recipes/deepseek-v4-flash-dspark.yaml diff --git a/mods/deepseek-v4-flash-dspark/overlay/vllm/config/speculative.py b/mods/deepseek-v4-flash-dspark/overlay/vllm/config/speculative.py new file mode 100644 index 00000000..52dfba7f --- /dev/null +++ b/mods/deepseek-v4-flash-dspark/overlay/vllm/config/speculative.py @@ -0,0 +1,1117 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import copy +from typing import TYPE_CHECKING, Any, Literal, get_args + +from pydantic import Field, SkipValidation, field_validator, model_validator +from typing_extensions import Self + +from vllm.config import LoadConfig +from vllm.config.kernel import MoEBackend +from vllm.config.model import ModelConfig +from vllm.config.parallel import ParallelConfig +from vllm.config.utils import config +from vllm.logger import init_logger +from vllm.transformers_utils.config import get_hf_text_config +from vllm.utils.hashing import safe_hash +from vllm.utils.import_utils import LazyLoader, has_arctic_inference +from vllm.v1.attention.backends.registry import AttentionBackendEnum + +if TYPE_CHECKING: + from transformers import PretrainedConfig + + import vllm.model_executor.layers.quantization as me_quant +else: + PretrainedConfig = Any + + me_quant = LazyLoader( + "model_executor", globals(), "vllm.model_executor.layers.quantization" + ) + +logger = init_logger(__name__) + +MTPModelTypes = Literal[ + "deepseek_mtp", + "mimo_mtp", + "mimo_v2_mtp", + "glm4_moe_mtp", + "glm4_moe_lite_mtp", + "glm_ocr_mtp", + "ernie_mtp", + "nemotron_h_mtp", + "exaone_moe_mtp", + "exaone4_5_mtp", + "qwen3_next_mtp", + "qwen3_5_mtp", + "longcat_flash_mtp", + "mtp", + "pangu_ultra_moe_mtp", + "step3p5_mtp", + "hy_v3_mtp", + "gemma4_mtp", +] +NgramGPUTypes = Literal["ngram_gpu"] +DFlashModelTypes = Literal["dflash"] +DSparkModelTypes = Literal["dspark"] +EagleModelTypes = Literal[ + "eagle", + "eagle3", + "extract_hidden_states", + MTPModelTypes, + DFlashModelTypes, + DSparkModelTypes, +] +SpeculativeMethod = Literal[ + "ngram", + "medusa", + "mlp_speculator", + "draft_model", + "suffix", + "custom_class", + EagleModelTypes, + NgramGPUTypes, +] +RejectionSampleMethod = Literal["standard", "synthetic"] +DraftSampleMethod = Literal["greedy", "probabilistic"] + + +@config +class SpeculativeConfig: + """Configuration for speculative decoding.""" + + enforce_eager: bool | None = None + """Override the default enforce_eager from model_config""" + # General speculative decoding control + num_speculative_tokens: int = Field(default=None, gt=0) # type: ignore[assignment] + """The number of speculative tokens, if provided. It will default to the + number in the draft model config if present, otherwise, it is required.""" + model: str | None = None + """The name of the draft model, eagle head, or additional weights, if + provided.""" + method: SpeculativeMethod | None = None + """The name of the speculative method to use. If users provide and set the + `model` param, the speculative method type will be detected automatically + if possible, if `model` param is not provided, the method name must be + provided. + + If using `ngram` method, the related configuration `prompt_lookup_max` and + `prompt_lookup_min` should be considered.""" + draft_tensor_parallel_size: int | None = Field(default=None, ge=1) + """The degree of the tensor parallelism for the draft model. Can only be 1 + or the same as the target model's tensor parallel size.""" + tensor_parallel_size: int | None = None + """Users should pass "draft_tensor_parallel_size". This parameter's purpose is to + warn users when they mistakenly provide the wrong argument.""" + + # Draft model configuration + quantization: me_quant.QuantizationMethods | str | None = None + """Quantization method that was used to quantize the draft model weights. + If `None`, we assume the model weights are not quantized. Note that it only + takes effect when using the draft model-based speculative method.""" + moe_backend: MoEBackend | None = None + """MoE backend to use for the draft model. When `None`, the draft model + inherits the target model's `--moe-backend` setting. Useful when the + drafter and generator require different MoE kernels (e.g. quantized + generator with unquantized drafter).""" + attention_backend: AttentionBackendEnum | None = None + """Attention backend to use for the draft model. When `None`, the backend is + automatically selected. Useful when the drafter requires a different attention + backend (e.g. DFlash needs a non-causal-capable backend like FLASH_ATTN).""" + max_model_len: int | None = Field(default=None, ge=1) + """The maximum model length of the draft model. Used when testing the + ability to skip speculation for some sequences.""" + revision: str | None = None + """The specific model version to use for the draft model. It can be a + branch name, a tag name, or a commit id. If unspecified, will use the + default version.""" + code_revision: str | None = None + """The specific revision to use for the draft model code on Hugging Face + Hub. It can be a branch name, a tag name, or a commit id. If unspecified, + will use the default version.""" + + # Advanced control + disable_padded_drafter_batch: bool = False + """Disable input padding for speculative decoding. If set to True, + speculative input batches can contain sequences of different lengths, + which may only be supported by certain attention backends. This currently + only affects the EAGLE method of speculation.""" + use_local_argmax_reduction: bool = False + """Use vocab-parallel local argmax instead of all-gathering full logits + for draft token generation. Reduces communication from O(vocab_size) to + O(2 * tp_size) per token. Only applies to greedy draft selection in + non-tree speculation.""" + + # Ngram proposer configuration + prompt_lookup_max: int | None = Field(default=None, ge=1) + """Maximum size of ngram token window when using Ngram proposer, required + when method is set to ngram.""" + prompt_lookup_min: int | None = Field(default=None, ge=1) + """Minimum size of ngram token window when using Ngram proposer, if + provided. Defaults to 1.""" + + # Alternative drafting strategies + parallel_drafting: bool = False + """Enable parallel drafting, where all speculative tokens are generated + in parallel rather than sequentially. This can improve performance but + requires the speculative model be trained to support parallel drafting. + Only compatible with EAGLE and draft model methods.""" + + # required configuration params passed from engine + target_model_config: SkipValidation[ModelConfig] = None # type: ignore + """The configuration of the target model.""" + target_parallel_config: SkipValidation[ParallelConfig] = None # type: ignore + """The parallel configuration for the target model.""" + + # params generated in the post-init stage + draft_model_config: SkipValidation[ModelConfig] = None # type: ignore + """The configuration of the draft model initialized internal.""" + draft_parallel_config: SkipValidation[ParallelConfig] = None # type: ignore + """The parallel configuration for the draft model initialized internal.""" + + # Suffix decoding configuration + suffix_decoding_max_tree_depth: int = 24 + """The maximum depth of the suffix decoding global and prompt trees. The + tree depth limits the sum of the prefix match and speculation lengths.""" + + suffix_decoding_max_cached_requests: int = 10000 + """The maximum number of requests to cache in the global suffix tree. If + exceeded, will trigger eviction in FIFO order. If set to 0, the global + suffix tree is disabled and past responses are not cached (prompt trees + are still used).""" + + suffix_decoding_max_spec_factor: float = 1.0 + """The maximum spec factor for suffix decoding. The spec factor controls + speculation lengths based on the prefix match length: max_spec_tokens = + max_spec_factor * prefix_match_length.""" + + suffix_decoding_min_token_prob: float = 0.1 + """The minimum token probability for suffix decoding. Will only speculate + tokens with estimated probability (based on frequency counts) greater than + or equal to this value.""" + + draft_load_config: LoadConfig | None = None + """Load config for the draft model. If not specified, will use the load + config from the target model.""" + + rejection_sample_method: RejectionSampleMethod = "standard" + """The rejection sampling method to use. 'standard' uses probabilistic + rejection sampling (with or without cached draft logits, controlled by + draft_sample_method). 'synthetic' accepts draft tokens with a decaying + probability calibrated to synthetic_acceptance_rate.""" + + synthetic_acceptance_rates: list[float] | None = None + """Per-position *unconditional* acceptance rates for synthetic rejection + sampling. Position i's entry is the marginal probability that the first + i+1 draft tokens are all accepted; the list must have length + num_speculative_tokens, each entry in [0, 1], and be monotonically + non-increasing. Only valid when rejection_sample_method is 'synthetic'. + Mutually exclusive with synthetic_acceptance_length.""" + + synthetic_acceptance_length: float | None = None + """Target mean acceptance length for synthetic rejection sampling, in + [1, num_speculative_tokens + 1]. Resolved internally to + synthetic_acceptance_rates. Only valid when rejection_sample_method is 'synthetic'. + Mutually exclusive with synthetic_acceptance_rates.""" + + @staticmethod + def _acceptance_length_to_rates(length: float, n: int) -> list[float]: + """Mean acceptance length to unconditional per-position rates, using + the minimum-variance schedule.""" + num_drafts = length - 1 # expected number of accepted draft tokens + num_full = int(num_drafts) + return ( + [1.0] * num_full + [num_drafts - num_full] + [0.0] * (n - num_full - 1) + )[:n] + + @staticmethod + def _resolve_synthetic_acceptance_rates( + n: int, + rates: list[float] | None, + length: float | None, + ) -> list[float]: + """Return per-position unconditional acceptance rates from exactly one + of `rates` or `length` (validates range, length, and monotonicity).""" + if (rates is None) == (length is None): + raise ValueError( + "rejection_sample_method='synthetic' requires exactly one of " + "synthetic_acceptance_rates or synthetic_acceptance_length." + ) + if rates is not None: + if len(rates) != n: + raise ValueError( + f"synthetic_acceptance_rates must have length {n}, got {rates}." + ) + if not all(0.0 <= r <= 1.0 for r in rates): + raise ValueError( + f"synthetic_acceptance_rates entries must be in [0, 1], " + f"got {rates}." + ) + if any(rates[i] > rates[i - 1] for i in range(1, n)): + raise ValueError( + f"synthetic_acceptance_rates must be non-increasing, got {rates}." + ) + return list(rates) + assert length is not None + if not 1.0 <= length <= float(n + 1): + raise ValueError( + f"synthetic_acceptance_length must be in [1, {n + 1}], got {length}." + ) + return SpeculativeConfig._acceptance_length_to_rates(length, n) + + draft_sample_method: DraftSampleMethod = "greedy" + """How the draft model samples tokens. 'greedy' always picks the argmax + token, and the draft probabilities are treated as one-hot during rejection + sampling. 'probabilistic' samples stochastically from the draft + distribution and uses the full draft logits for the probability ratio test + during rejection sampling. This comes at the cost of additional GPU memory + usage.""" + + 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 + it affects the computation graph. + + Provide a hash that uniquely identifies all the configs + that affect the structure of the computation + graph from input ids/embeddings to the final hidden states, + excluding anything before input ids/embeddings and after + the final hidden states. + """ + factors: list[Any] = [] + # Eagle3 and extract_hidden_states affect the computation graph because + # they return intermediate hidden states in addition to the final hidden state. + uses_aux_hidden_states = self.method in ( + "eagle3", + "extract_hidden_states", + "dflash", + "dspark", + ) + factors.append(uses_aux_hidden_states) + + # The specific layers used also affect the computation graph + if uses_aux_hidden_states and self.draft_model_config is not None: + layer_ids = getattr( + self.draft_model_config.hf_config, + "eagle_aux_hidden_state_layer_ids", + None, + ) + if layer_ids is not None: + # Convert to tuple to make it hashable + factors.append(tuple(layer_ids)) + + hash_str = safe_hash(str(factors).encode(), usedforsecurity=False).hexdigest() + return hash_str + + @staticmethod + def hf_config_override(hf_config: PretrainedConfig) -> PretrainedConfig: + initial_architecture = hf_config.architectures[0] + if hf_config.model_type in ( + "deepseek_v3", + "deepseek_v32", + "glm_moe_dsa", + ): + hf_config.model_type = "deepseek_mtp" + if hf_config.model_type == "deepseek_mtp": + n_predict = getattr(hf_config, "num_nextn_predict_layers", None) + hf_config.update( + {"n_predict": n_predict, "architectures": ["DeepSeekMTPModel"]} + ) + if hf_config.model_type == "deepseek_v4" and getattr( + hf_config, "dspark_block_size", 0 + ): + n_draft_layers = len(getattr(hf_config, "dspark_target_layer_ids", ())) + n_draft_layers = max(1, n_draft_layers) + n_predict = getattr(hf_config, "dspark_block_size", None) + hf_config.model_type = "deepseek_v4_dspark" + hf_config.update( + { + "n_predict": n_predict, + "num_nextn_predict_layers": n_draft_layers, + "dspark_num_draft_layers": n_draft_layers, + "architectures": ["DeepSeekV4DSparkModel"], + } + ) + elif hf_config.model_type == "deepseek_v4": + hf_config.model_type = "deepseek_mtp" + n_predict = getattr(hf_config, "num_nextn_predict_layers", None) + hf_config.update( + {"n_predict": n_predict, "architectures": ["DeepSeekV4MTPModel"]} + ) + if hf_config.model_type in ("pangu_ultra_moe"): + hf_config.model_type = "pangu_ultra_moe_mtp" + if hf_config.model_type == "pangu_ultra_moe_mtp": + n_predict = getattr(hf_config, "num_nextn_predict_layers", None) + hf_config.update( + {"n_predict": n_predict, "architectures": ["OpenPanguMTPModel"]} + ) + + if hf_config.architectures[0] == "MiMoForCausalLM": + hf_config.model_type = "mimo_mtp" + n_predict = getattr(hf_config, "num_nextn_predict_layers", None) + hf_config.update( + { + "num_hidden_layers": 0, + "n_predict": n_predict, + "architectures": ["MiMoMTPModel"], + } + ) + + if (arch := hf_config.architectures[0]) in ( + "MiMoV2ForCausalLM", + "MiMoV2OmniForCausalLM", + ): + from vllm.model_executor.models.mimo_v2_mtp import ( + _MIMO_V2_PRO_NUM_MTP_LAYERS, + ) + + mtp_arch_maps = { + "MiMoV2ForCausalLM": "MiMoV2MTPModel", + "MiMoV2OmniForCausalLM": "MiMoV2OmniMTPModel", + } + + hf_config.model_type = "mimo_v2_mtp" + # vLLM currently supports only the first MiMo-V2 MTP layer. + n_predict = _MIMO_V2_PRO_NUM_MTP_LAYERS + hf_config.update( + { + "num_hidden_layers": 0, + "n_predict": n_predict, + "num_nextn_predict_layers": n_predict, + "architectures": [mtp_arch_maps[arch]], + } + ) + + if hf_config.architectures[0] == "MiMoV2FlashForCausalLM": + from vllm.model_executor.models.mimo_v2_mtp import ( + _MIMO_V2_FLASH_NUM_MTP_LAYERS, + ) + + hf_config.model_type = "mimo_v2_mtp" + # vLLM currently supports only the first MiMo-V2 MTP layer. + n_predict = _MIMO_V2_FLASH_NUM_MTP_LAYERS + hf_config.update( + { + "num_hidden_layers": 0, + "n_predict": n_predict, + "num_nextn_predict_layers": n_predict, + "architectures": ["MiMoV2MTPModel"], + } + ) + + if hf_config.architectures[0] == "Glm4MoeForCausalLM": + hf_config.model_type = "glm4_moe_mtp" + n_predict = getattr(hf_config, "num_nextn_predict_layers", None) + hf_config.update( + { + "n_predict": n_predict, + "architectures": ["Glm4MoeMTPModel"], + } + ) + + if hf_config.architectures[0] == "Glm4MoeLiteForCausalLM": + hf_config.model_type = "glm4_moe_lite_mtp" + n_predict = getattr(hf_config, "num_nextn_predict_layers", None) + hf_config.update( + { + "num_hidden_layers": 0, + "n_predict": n_predict, + "architectures": ["Glm4MoeLiteMTPModel"], + } + ) + + if hf_config.architectures[0] == "GlmOcrForConditionalGeneration": + hf_config.model_type = "glm_ocr_mtp" + n_predict = getattr(hf_config, "num_nextn_predict_layers", None) + hf_config.update( + { + "num_hidden_layers": 0, + "n_predict": n_predict, + "architectures": ["GlmOcrMTPModel"], + } + ) + + if hf_config.model_type == "ernie4_5_moe": + hf_config.model_type = "ernie_mtp" + if hf_config.model_type == "ernie_mtp": + n_predict = getattr(hf_config, "num_nextn_predict_layers", None) + hf_config.update( + {"n_predict": n_predict, "architectures": ["ErnieMTPModel"]} + ) + + if hf_config.architectures[0] == "NemotronH_Super_Omni_Reasoning_V3": + # Promote VLM's text_config so MTP detection below fires correctly + hf_config = hf_config.text_config + + if ( + hf_config.model_type in {"nemotron_h", "nemotron_h_puzzle"} + and hasattr(hf_config, "num_nextn_predict_layers") + and hf_config.num_nextn_predict_layers > 0 + ): + # Check if this is an MTP variant + hf_config.model_type = "nemotron_h_mtp" + if hf_config.model_type == "nemotron_h_mtp": + n_predict = getattr(hf_config, "num_nextn_predict_layers", 1) + hf_config.update( + {"n_predict": n_predict, "architectures": ["NemotronHMTPModel"]} + ) + + if hf_config.model_type == "qwen3_next": + hf_config.model_type = "qwen3_next_mtp" + if hf_config.model_type == "qwen3_next_mtp": + n_predict = getattr(hf_config, "num_nextn_predict_layers", None) + hf_config.update( + {"n_predict": n_predict, "architectures": ["Qwen3NextMTP"]} + ) + + if hf_config.model_type == "exaone_moe": + hf_config.model_type = "exaone_moe_mtp" + if hf_config.model_type == "exaone_moe_mtp": + n_predict = getattr(hf_config, "num_nextn_predict_layers", None) + hf_config.update( + {"n_predict": n_predict, "architectures": ["ExaoneMoeMTP"]} + ) + if "exaone4_5" in hf_config.model_type: + hf_config.model_type = "exaone4_5_mtp" + if hf_config.model_type == "exaone4_5_mtp": + n_predict = getattr(hf_config, "num_nextn_predict_layers", None) + hf_config.update( + {"n_predict": n_predict, "architectures": ["Exaone4_5_MTP"]} + ) + if hf_config.model_type in ("qwen3_5", "qwen3_5_moe"): + is_moe = hf_config.model_type == "qwen3_5_moe" + hf_config.model_type = "qwen3_5_mtp" + n_predict = getattr(hf_config, "mtp_num_hidden_layers", None) + hf_config.update( + { + "n_predict": n_predict, + "architectures": ["Qwen3_5MoeMTP" if is_moe else "Qwen3_5MTP"], + } + ) + if hf_config.model_type == "intern_s2_preview": + text_config = getattr(hf_config, "text_config", None) + is_moe = getattr(text_config, "model_type", None) == "qwen3_5_moe_text" + hf_config.model_type = "qwen3_5_mtp" + n_predict = getattr(text_config, "mtp_num_hidden_layers", None) + hf_config.update( + { + "n_predict": n_predict, + "architectures": ["Qwen3_5MoeMTP" if is_moe else "Qwen3_5MTP"], + } + ) + if hf_config.model_type == "longcat_flash": + hf_config.model_type = "longcat_flash_mtp" + n_predict = getattr(hf_config, "num_nextn_predict_layers", 1) + hf_config.update( + {"n_predict": n_predict, "architectures": ["LongCatFlashMTPModel"]} + ) + + if hf_config.model_type == "step3p5": + hf_config.model_type = "step3p5_mtp" + n_predict = getattr(hf_config, "num_nextn_predict_layers", 1) + hf_config.update({"n_predict": n_predict, "architectures": ["Step3p5MTP"]}) + + if initial_architecture == "MistralLarge3ForCausalLM": + hf_config.update({"architectures": ["EagleMistralLarge3ForCausalLM"]}) + + if hf_config.model_type == "hy_v3": + hf_config.model_type = "hy_v3_mtp" + n_predict = getattr(hf_config, "num_nextn_predict_layers", None) + hf_config.update( + {"n_predict": n_predict, "architectures": ["HYV3MTPModel"]} + ) + + if hf_config.model_type == "gemma4_assistant": + hf_config.model_type = "gemma4_mtp" + text_config = getattr(hf_config, "text_config", hf_config) + # The assistant runs all decoder layers in a single forward + # call to produce one draft token, so n_predict=1. + # num_kv_shared_layers must be 0: cross-model KV sharing is + # set up by the proposer after model construction. + if hasattr(text_config, "num_kv_shared_layers"): + text_config.num_kv_shared_layers = 0 + hf_config.update({"n_predict": 1, "architectures": ["Gemma4MTPModel"]}) + + return hf_config + + def __post_init__(self): + # Note: "method" is a new parameter that helps to extend the + # configuration of non-model-based proposers, and the "model" parameter + # will be used to set the draft model, eagle head, or additional weight + # when needed. If users do not specify "method", the speculative method + # will be detected automatically if possible. If the speculative method + # can not be detected, it will be considered as the "draft_model" by + # default. + + # infer method from user args + # Check if the model field contains a custom module path (e.g., 'pkg.Mod') + if ( + self.model is not None + and "." in self.model + and not self.model.startswith(("http://", "https://", "file://")) + and "/" not in self.model # not a HuggingFace repo (org/model) + ): + # Treat as a custom class path + self.method = "custom_class" + elif self.method is None: + if self.model in ("ngram", "[ngram]"): + self.method = "ngram" + else: + self.method = "draft_model" + + if self.method in get_args(MTPModelTypes) and self.method != "mtp": + logger.warning( + "method `%s` is deprecated and replaced with mtp.", self.method + ) + self.method = "mtp" + + if self.model is None and self.num_speculative_tokens is not None: + if self.method in ("mtp", "dspark"): + if self.target_model_config is None: + raise ValueError( + f"target_model_config must be present for {self.method}" + ) + if self.target_model_config.hf_text_config.model_type == "deepseek_v32": + # FIXME(luccafong): cudagraph with v32 MTP is not supported, + # remove this when the issue is fixed. + self.enforce_eager = True + # use the draft model from the same model: + self.model = self.target_model_config.model + # Align the quantization of draft model for cases such as + # --quantization fp8 with a bf16 checkpoint. + if not self.quantization: + self.quantization = self.target_model_config.quantization + elif self.method in ("ngram", "[ngram]"): + self.model = "ngram" + elif self.method == "ngram_gpu": + self.model = "ngram_gpu" + elif self.method == "suffix": + self.model = "suffix" + elif self.method == "extract_hidden_states": + self.model = "extract_hidden_states" + elif self.method == "custom_class": + # method was set explicitly, but model should already contain the + # custom module path. If not, this is a configuration error. + if self.model is None: + raise ValueError( + "method='custom_class' requires 'model' to contain the " + "custom proposer module path (e.g., 'my_module.MyProposer')." + ) + else: + raise ValueError( + "num_speculative_tokens was provided but without speculative model." + ) + + if self.method in ("ngram", "[ngram]"): + self.method = "ngram" + + if self.method in ("ngram", "ngram_gpu"): + # Set default values if not provided + if self.prompt_lookup_min is None and self.prompt_lookup_max is None: + # TODO(woosuk): Tune these values. They are arbitrarily chosen. + self.prompt_lookup_min = 5 + self.prompt_lookup_max = 5 + elif self.prompt_lookup_min is None: + if self.prompt_lookup_max is None: + raise ValueError( + "Either prompt_lookup_max or prompt_lookup_min must be " + "provided when using the ngram method." + ) + self.prompt_lookup_min = self.prompt_lookup_max + elif self.prompt_lookup_max is None: + if self.prompt_lookup_min is None: + raise ValueError( + "Either prompt_lookup_max or prompt_lookup_min must be " + "provided when using the ngram method." + ) + self.prompt_lookup_max = self.prompt_lookup_min + + # Validate values + if self.prompt_lookup_min > self.prompt_lookup_max: + raise ValueError( + f"prompt_lookup_min={self.prompt_lookup_min} must " + f"be <= prompt_lookup_max={self.prompt_lookup_max}" + ) + + # TODO: current we still need extract vocab_size from target model + # config, in future, we may try refactor it out, and set + # draft related config as None here. + self.draft_model_config = self.target_model_config + self.draft_parallel_config = self.target_parallel_config + elif self.method == "suffix": + self._validate_suffix_decoding() + elif self.method == "custom_class": + # Custom class proposer does not need a draft model. + # It will dynamically load the user-provided class at runtime. + logger.warning_once( + "Using a custom class-based proposer backend. This is an " + "experimental feature and the proposer interface is subject to " + "breaking changes in future vLLM releases." + ) + self.prompt_lookup_max = 0 + self.prompt_lookup_min = 0 + self.draft_model_config = self.target_model_config + self.draft_parallel_config = self.target_parallel_config + elif self.method == "extract_hidden_states": + from vllm.transformers_utils.configs.extract_hidden_states import ( + ExtractHiddenStatesConfig, + ) + + # ExtractHiddenStatesModel is instantiated manually in load_model() + # We just need to store the target model config for KV cache shape info + self.model = "extract_hidden_states" + self.prompt_lookup_max = 0 + self.prompt_lookup_min = 0 + + if hasattr(self.draft_model_config, "hf_config"): + hf_config = self.draft_model_config.hf_config.to_dict() + elif ( + isinstance(self.draft_model_config, dict) + and "hf_config" in self.draft_model_config + ): + hf_config = self.draft_model_config["hf_config"] + else: + hf_config = {} + + self.draft_model_config = copy.copy(self.target_model_config) + self.draft_model_config.hf_config = ExtractHiddenStatesConfig( + self.draft_model_config.hf_config, **hf_config + ) + self.update_arch_() + self.draft_parallel_config = self.target_parallel_config + + else: + self.prompt_lookup_max = 0 + self.prompt_lookup_min = 0 + + if self.model is not None: + self.draft_model_config = ModelConfig( + model=self.model, + runner="draft", + tokenizer=self.target_model_config.tokenizer, + tokenizer_mode=self.target_model_config.tokenizer_mode, + trust_remote_code=self.target_model_config.trust_remote_code, + allowed_local_media_path=self.target_model_config.allowed_local_media_path, + allowed_media_domains=self.target_model_config.allowed_media_domains, + dtype=self.target_model_config.dtype, + seed=self.target_model_config.seed, + revision=self.revision, + code_revision=self.code_revision, + tokenizer_revision=self.target_model_config.tokenizer_revision, + max_model_len=self.max_model_len, # type: ignore[arg-type] + spec_target_max_model_len=self.target_model_config.max_model_len, + quantization=self.quantization, + enforce_eager=self.target_model_config.enforce_eager, + max_logprobs=self.target_model_config.max_logprobs, + hf_overrides=SpeculativeConfig.hf_config_override, + config_format=self.target_model_config.config_format, + ) + + # Automatically detect the method + if self.method in ("eagle", "eagle3", "dflash", "dspark"): + pass + # examples: + # yuhuili/EAGLE-LLaMA3-Instruct-8B + # yuhuili/EAGLE3-LLaMA3.1-Instruct-8B + # AngelSlim/Qwen3-8B_eagle3 + elif "eagle-" in self.draft_model_config.model.lower(): + self.method = "eagle" + elif "eagle3" in self.draft_model_config.model.lower(): + self.method = "eagle3" + elif "dflash" in self.draft_model_config.model.lower(): + self.method = "dflash" + elif "dspark" in self.draft_model_config.model.lower(): + self.method = "dspark" + elif self.draft_model_config.hf_config.model_type == "medusa": + self.method = "medusa" + elif self.draft_model_config.hf_config.model_type == "mlp_speculator": + self.method = "mlp_speculator" + elif self.draft_model_config.hf_config.model_type in get_args( + MTPModelTypes + ): + self.method = "mtp" + if self.num_speculative_tokens > 1: + logger.warning( + "Enabling num_speculative_tokens > 1 will run " + "multiple times of forward on same MTP layer" + ",which may result in lower acceptance rate" + ) + elif self.draft_model_config.hf_config.model_type in ( + "longcat_flash_mtp" + ): + self.method = "longcat_flash_mtp" + if self.num_speculative_tokens > 1: + logger.warning( + "LongCat MTP models only have " + "one layer. Might need some code changes " + "to support multiple layers." + ) + elif self.method == "draft_model": + pass + else: + raise NotImplementedError( + f"Unsupported speculative method: '{self.method}'" + ) + + # Replace hf_config for EAGLE draft_model + if self.method in ("eagle", "eagle3", "dflash"): + from vllm.transformers_utils.configs.eagle import EAGLEConfig + from vllm.transformers_utils.configs.speculators import ( + SpeculatorsConfig, + ) + + if isinstance( + self.draft_model_config.hf_config, + (EAGLEConfig, SpeculatorsConfig), + ): + pass + else: + eagle_config = EAGLEConfig( + self.draft_model_config.hf_config, + method=self.method, + model_type="eagle", + ) + self.draft_model_config.hf_config = eagle_config + self.update_arch_() + + if self.method == "dflash": + self.parallel_drafting = True + + if self.num_speculative_tokens is not None and hasattr( + self.draft_model_config.hf_config, "num_lookahead_tokens" + ): + self.draft_model_config.hf_config.num_lookahead_tokens = ( + self.num_speculative_tokens + ) + + n_predict = getattr( + self.draft_model_config.hf_config, "n_predict", None + ) + if n_predict is not None: + if self.num_speculative_tokens is None: + # Default to max value defined in draft model config. + self.num_speculative_tokens = n_predict + elif ( + self.num_speculative_tokens > n_predict + and self.num_speculative_tokens % n_predict != 0 + ): + # Ensure divisibility for MTP module reuse. + raise ValueError( + f"num_speculative_tokens:{self.num_speculative_tokens}" + f" must be divisible by {n_predict=}" + ) + + if self.num_speculative_tokens is None: + raise ValueError( + "A speculative model was provided, but " + "`num_speculative_tokens` was not provided" + ) + + self.draft_tensor_parallel_size = ( + SpeculativeConfig._verify_and_get_draft_tp( + self.target_parallel_config, + self.draft_tensor_parallel_size, + self.draft_model_config.hf_config, + ) + ) + + self.draft_model_config.max_model_len = ( + SpeculativeConfig._maybe_override_draft_max_model_len( + self.max_model_len, + self.draft_model_config.max_model_len, + self.target_model_config.max_model_len, + ) + ) + + self.draft_parallel_config = ( + SpeculativeConfig.create_draft_parallel_config( + self.target_parallel_config, self.draft_tensor_parallel_size + ) + ) + return self + + def _validate_suffix_decoding(self): + if not has_arctic_inference(): + raise ImportError( + "Arctic Inference is required for suffix decoding. " + "Install via `pip install arctic-inference==0.1.1`." + ) + if self.num_speculative_tokens is None: + # Suffix decoding decides the actual number of speculative tokens + # dynamically and treats num_speculative_tokens as a maximum limit. + self.num_speculative_tokens = self.suffix_decoding_max_tree_depth + logger.warning( + "Defaulted num_speculative_tokens to %s for suffix decoding.", + self.num_speculative_tokens, + ) + # Validate values + if self.suffix_decoding_max_tree_depth < 1: + raise ValueError( + f"suffix_decoding_max_tree_depth=" + f"{self.suffix_decoding_max_tree_depth} must be >= 1" + ) + if self.suffix_decoding_max_cached_requests < 0: + raise ValueError( + f"suffix_decoding_max_cached_requests=" + f"{self.suffix_decoding_max_cached_requests} must be >= 0" + ) + if self.suffix_decoding_max_spec_factor < 0: + raise ValueError( + f"suffix_decoding_max_spec_factor=" + f"{self.suffix_decoding_max_spec_factor} must be >= 0" + ) + if not 0 <= self.suffix_decoding_min_token_prob <= 1: + raise ValueError( + f"suffix_decoding_min_token_prob=" + f"{self.suffix_decoding_min_token_prob} must be in [0, 1]" + ) + + @staticmethod + def _maybe_override_draft_max_model_len( + speculative_max_model_len: int | None, + draft_max_model_len: int, + target_max_model_len: int, + ) -> int: + """Determine the max sequence len for the draft model. This is usually + the draft_max_model_len, but may be the target_max_model_len if it is + less than the draft_max_model_len, or may be speculative_max_model_len + if it is specified. + + This is necessary so that sequences do not exceed the capacity of the + draft model or the target model. + + speculative_max_model_len is mainly used for testing that sequences can + skip speculation. + """ + + if speculative_max_model_len is not None: + if speculative_max_model_len > draft_max_model_len: + raise ValueError( + f"{speculative_max_model_len=} cannot be " + f"larger than {draft_max_model_len=}" + ) + + if speculative_max_model_len > target_max_model_len: + raise ValueError( + f"{speculative_max_model_len=} cannot be " + f"larger than {target_max_model_len=}" + ) + + return speculative_max_model_len + + result = min( + draft_max_model_len, + target_max_model_len, + ) + if result != draft_max_model_len: + logger.info( + "Overriding draft model max model len from %d to %d", + draft_max_model_len, + result, + ) + return result + + @staticmethod + def _verify_and_get_draft_tp( + target_parallel_config: ParallelConfig, + speculative_draft_tensor_parallel_size: int | None, + draft_hf_config: PretrainedConfig, + ) -> int: + """ + Verifies and adjusts the tensor parallel size for a draft model + specified using speculative_draft_tensor_parallel_size. + """ + # If speculative_draft_tensor_parallel_size is unset then set it + # appropriately else verify that it is set correctly. + if speculative_draft_tensor_parallel_size is None: + if draft_hf_config.model_type == "mlp_speculator": + speculative_draft_tensor_parallel_size = 1 + if target_parallel_config.tensor_parallel_size > 1: + logger.warning( + "%s cannot currently be run with tp>1; " + "setting speculative_draft_tensor_parallel_size=1", + draft_hf_config.model_type, + ) + else: + speculative_draft_tensor_parallel_size = ( + target_parallel_config.tensor_parallel_size + ) + elif speculative_draft_tensor_parallel_size not in ( + 1, + target_parallel_config.tensor_parallel_size, + ): + raise ValueError( + f"{speculative_draft_tensor_parallel_size=} cannot be " + f"other value than 1 or target model tensor_parallel_size" + ) + return speculative_draft_tensor_parallel_size + + def update_arch_(self): + """ + EagleConfig and ExtractHiddenStatesConfig update architectures, so update all + architectures-related fields in self.draft_model_config + """ + self.draft_model_config.hf_text_config = get_hf_text_config( + self.draft_model_config.hf_config + ) + self.draft_model_config.model_arch_config = ( + self.draft_model_config.get_model_arch_config() + ) + model_info, arch = self.draft_model_config.registry.inspect_model_cls( + self.draft_model_config.architectures, + self.draft_model_config, + ) + self.draft_model_config._model_info = model_info + self.draft_model_config._architecture = arch + + @staticmethod + def create_draft_parallel_config( + target_parallel_config: ParallelConfig, + speculative_draft_tensor_parallel_size: int, + ) -> ParallelConfig: + """Create a parallel config for use by the draft worker. + + This is mostly a copy of the target parallel config, except the tp_size. + """ + draft_parallel_config = ParallelConfig( + pipeline_parallel_size=target_parallel_config.pipeline_parallel_size, + tensor_parallel_size=speculative_draft_tensor_parallel_size, + distributed_executor_backend=target_parallel_config.distributed_executor_backend, + max_parallel_loading_workers=target_parallel_config.max_parallel_loading_workers, + disable_custom_all_reduce=target_parallel_config.disable_custom_all_reduce, + ray_workers_use_nsight=target_parallel_config.ray_workers_use_nsight, + placement_group=target_parallel_config.placement_group, + ) + + return draft_parallel_config + + @field_validator("attention_backend", mode="before") + @classmethod + def _parse_attention_backend(cls, value: Any) -> Any: + if isinstance(value, str): + if value.lower() == "auto": + return None + return AttentionBackendEnum[value.upper()] + return value + + @model_validator(mode="after") + def _verify_args(self) -> Self: + if self.tensor_parallel_size is not None: + raise ValueError( + "'tensor_parallel_size' is not a valid argument in the " + "speculative_config. Please pass 'draft_tensor_parallel_size' instead." + ) + + if self.num_speculative_tokens is None: + raise ValueError( + "num_speculative_tokens must be provided with " + "speculative model unless the draft model config contains an " + "n_predict parameter." + ) + + if self.num_speculative_tokens <= 0: + raise ValueError( + "Expected num_speculative_tokens to be greater " + f"than zero ({self.num_speculative_tokens})." + ) + + if self.rejection_sample_method == "synthetic": + # Consolidate to per-position rates + self.synthetic_acceptance_rates = self._resolve_synthetic_acceptance_rates( + self.num_speculative_tokens, + self.synthetic_acceptance_rates, + self.synthetic_acceptance_length, + ) + self.synthetic_acceptance_length = None + elif ( + self.synthetic_acceptance_rates is not None + or self.synthetic_acceptance_length is not None + ): + raise ValueError( + "synthetic_acceptance_rates / synthetic_acceptance_length " + "are only valid with rejection_sample_method='synthetic'." + ) + + if self.draft_model_config: + self.draft_model_config.verify_with_parallel_config( + self.draft_parallel_config + ) + + self.verify_equal_vocab_size_if_draft_model() + return self + + def verify_equal_vocab_size_if_draft_model(self): + if ( + self.method == "draft_model" + and self.target_model_config is not None + and self.draft_model_config is not None + ): + target_vocab_size = self.target_model_config.get_vocab_size() + draft_vocab_size = self.draft_model_config.get_vocab_size() + if target_vocab_size != draft_vocab_size: + raise ValueError( + f"Target and draft model should have the same vocabulary size. " + f"Target model vocab_size={target_vocab_size}. " + f"Draft model vocab_size={draft_vocab_size}. " + f"Using models with different tokenizers can cause out-of-bounds " + f"errors during speculative decoding." + ) + + @property + def max_num_new_slots_for_drafting(self) -> int: + """ + Calculate the maximum number of new slots that might be added to the batch + when drafting. + """ + slots_per_req = 0 # for serial non-draft-model methods, no change needed + if self.parallel_drafting: + # For parallel drafting, we need one new slot per 'masked' token + slots_per_req = self.num_speculative_tokens - 1 + if self.uses_draft_model(): + # For draft model-based speculation, we need one new slot per request + # Since we do not slice the draft tokens + slots_per_req += 1 + return slots_per_req + + def use_gemma4_mtp(self) -> bool: + return ( + self.method == "mtp" + and self.draft_model_config is not None + and getattr(self.draft_model_config.hf_config, "model_type", None) + == "gemma4_mtp" + ) + + def use_eagle(self) -> bool: + return self.method in ("eagle", "eagle3", "mtp", "dflash", "dspark") + + def use_dflash(self) -> bool: + return self.method == "dflash" + + def use_dspark(self) -> bool: + return self.method == "dspark" + + def uses_draft_model(self) -> bool: + return self.method == "draft_model" + + def uses_extract_hidden_states(self) -> bool: + return self.method == "extract_hidden_states" + + def use_ngram_gpu(self) -> bool: + return self.method == "ngram_gpu" + + def __repr__(self) -> str: + method = self.method + model = ( + None + if method + in ( + "ngram", + "suffix", + "extract_hidden_states", + "custom_class", + ) + else self.draft_model_config.model + ) + num_spec_tokens = self.num_speculative_tokens + return f"SpeculativeConfig({method=}, {model=}, {num_spec_tokens=})" diff --git a/mods/deepseek-v4-flash-dspark/overlay/vllm/envs.py b/mods/deepseek-v4-flash-dspark/overlay/vllm/envs.py new file mode 100755 index 00000000..da651f15 --- /dev/null +++ b/mods/deepseek-v4-flash-dspark/overlay/vllm/envs.py @@ -0,0 +1,2314 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import functools +import json +import logging +import os +import sys +import tempfile +import uuid +import warnings +from collections.abc import Callable +from typing import TYPE_CHECKING, Any, Literal + +if TYPE_CHECKING: + VLLM_HOST_IP: str = "" + VLLM_PORT: int | None = None + VLLM_RPC_BASE_PATH: str = tempfile.gettempdir() + VLLM_USE_MODELSCOPE: bool = False + VLLM_USE_FASTOKENS: bool = False + VLLM_RINGBUFFER_WARNING_INTERVAL: int = 60 + VLLM_NCCL_SO_PATH: str | None = None + LD_LIBRARY_PATH: str | None = None + VLLM_ROCM_SLEEP_MEM_CHUNK_SIZE: int = 256 + LOCAL_RANK: int = 0 + CUDA_VISIBLE_DEVICES: str | None = None + VLLM_ENGINE_ITERATION_TIMEOUT_S: int = 60 + VLLM_ENGINE_READY_TIMEOUT_S: int = 600 + VLLM_API_KEY: str | None = None + VLLM_DEBUG_LOG_API_SERVER_RESPONSE: bool = False + S3_ACCESS_KEY_ID: str | None = None + S3_SECRET_ACCESS_KEY: str | None = None + S3_ENDPOINT_URL: str | None = None + VLLM_MODEL_REDIRECT_PATH: str | None = None + VLLM_CACHE_ROOT: str = os.path.expanduser("~/.cache/vllm") + VLLM_CONFIG_ROOT: str = os.path.expanduser("~/.config/vllm") + VLLM_USAGE_STATS_SERVER: str = "https://stats.vllm.ai" + VLLM_NO_USAGE_STATS: bool = False + VLLM_DO_NOT_TRACK: bool = False + VLLM_USAGE_SOURCE: str = "production" + VLLM_CONFIGURE_LOGGING: bool = True + VLLM_LOGGING_LEVEL: str = "INFO" + VLLM_LOGGING_PREFIX: str = "" + VLLM_LOGGING_STREAM: str = "ext://sys.stdout" + VLLM_LOGGING_CONFIG_PATH: str | None = None + VLLM_LOGGING_COLOR: str = "auto" + NO_COLOR: bool = False + VLLM_LOG_STATS_INTERVAL: float = 10.0 + VLLM_TRACE_FUNCTION: int = 0 + VLLM_USE_FLASHINFER_SAMPLER: bool = True + VLLM_PP_LAYER_PARTITION: str | None = None + VLLM_CPU_KVCACHE_SPACE: int | None = 0 + VLLM_CPU_OMP_THREADS_BIND: str = "auto" + VLLM_CPU_NUM_OF_RESERVED_CPU: int | None = None + VLLM_CPU_SGL_KERNEL: bool = False + VLLM_CPU_ATTN_SPLIT_KV: bool = True + VLLM_ZENTORCH_WEIGHT_PREPACK: bool = True + VLLM_CPU_INT4_W4A8: bool = True + VLLM_XLA_CACHE_PATH: str = os.path.join(VLLM_CACHE_ROOT, "xla_cache") + VLLM_XLA_CHECK_RECOMPILATION: bool = False + VLLM_SPARSE_INDEXER_MAX_LOGITS_MB: int = 512 + VLLM_USE_B12X_SPARSE_INDEXER: bool = False + VLLM_USE_B12X_MHC: bool = False + VLLM_USE_B12X_FP8_GEMM: bool = False + VLLM_USE_B12X_WO_PROJECTION: bool = False + VLLM_USE_B12X_MOE: bool = False + VLLM_B12X_W4A16_FORCE_BLOCKS_PER_SM: int = 0 + VLLM_B12X_W4A16_FORCE_BLOCKS_MAX_M: int = 16 + VLLM_B12X_W4A16_FORCE_TILE_CONFIG: str = "" + VLLM_USE_RAY_COMPILED_DAG_CHANNEL_TYPE: Literal["auto", "nccl", "shm"] = "auto" + VLLM_USE_RAY_COMPILED_DAG_OVERLAP_COMM: bool = False + VLLM_USE_RAY_WRAPPED_PP_COMM: bool = True + VLLM_USE_RAY_V2_EXECUTOR_BACKEND: bool = False + VLLM_XLA_USE_SPMD: bool = False + VLLM_WORKER_MULTIPROC_METHOD: Literal["fork", "spawn"] = "fork" + VLLM_ASSETS_CACHE: str = os.path.join(VLLM_CACHE_ROOT, "assets") + VLLM_ASSETS_CACHE_MODEL_CLEAN: bool = False + VLLM_IMAGE_FETCH_TIMEOUT: int = 5 + VLLM_VIDEO_FETCH_TIMEOUT: int = 30 + VLLM_AUDIO_FETCH_TIMEOUT: int = 10 + VLLM_MEDIA_CACHE: str = "" + VLLM_MEDIA_CACHE_MAX_SIZE_MB: int = 5120 + VLLM_MEDIA_CACHE_TTL_HOURS: float = 24 + VLLM_MEDIA_FETCH_MAX_RETRIES: int = 3 + VLLM_MEDIA_URL_ALLOW_REDIRECTS: bool = True + VLLM_MEDIA_LOADING_THREAD_COUNT: int = 8 + VLLM_MAX_AUDIO_CLIP_FILESIZE_MB: int = 25 + VLLM_VIDEO_LOADER_BACKEND: str = "opencv" + VLLM_MEDIA_CONNECTOR: str = "http" + VLLM_MM_HASHER_ALGORITHM: str = "blake3" + VLLM_TARGET_DEVICE: str = "cuda" + VLLM_MAIN_CUDA_VERSION: str = "13.0" + VLLM_FLOAT32_MATMUL_PRECISION: Literal["highest", "high", "medium"] = "highest" + VLLM_BATCH_INVARIANT: bool = False + VLLM_TRITON_ATTN_USE_TD: bool | None = None + MAX_JOBS: str | None = None + NVCC_THREADS: str | None = None + VLLM_USE_PRECOMPILED: bool = False + VLLM_USE_PRECOMPILED_RUST: bool = False + VLLM_SKIP_PRECOMPILED_VERSION_SUFFIX: bool = False + VLLM_DOCKER_BUILD_CONTEXT: bool = False + VLLM_KEEP_ALIVE_ON_ENGINE_DEATH: bool = False + CMAKE_BUILD_TYPE: Literal["Debug", "Release", "RelWithDebInfo"] | None = None + VERBOSE: bool = False + VLLM_ALLOW_LONG_MAX_MODEL_LEN: bool = False + VLLM_RPC_TIMEOUT: int = 10000 # ms + VLLM_HTTP_TIMEOUT_KEEP_ALIVE: int = 5 # seconds + VLLM_MAX_N_SEQUENCES: int = 16384 + VLLM_PLUGINS: list[str] | None = None + VLLM_LORA_RESOLVER_CACHE_DIR: str | None = None + VLLM_LORA_RESOLVER_HF_REPO_LIST: str | None = None + VLLM_USE_AOT_COMPILE: bool = False + VLLM_USE_BYTECODE_HOOK: bool = True + VLLM_FORCE_AOT_LOAD: bool = False + VLLM_USE_MEGA_AOT_ARTIFACT: bool = False + VLLM_USE_TRITON_AWQ: bool = False + VLLM_ALLOW_RUNTIME_LORA_UPDATING: bool = False + VLLM_SKIP_P2P_CHECK: bool = False + VLLM_DISABLED_KERNELS: list[str] = [] + VLLM_ENABLE_FLA_PACKED_RECURRENT_DECODE: bool = True + VLLM_DISABLE_PYNCCL: bool = False + VLLM_USE_OINK_OPS: bool = False + VLLM_ROCM_USE_AITER: bool = False + VLLM_ROCM_USE_AITER_PAGED_ATTN: bool = False + VLLM_ROCM_USE_AITER_LINEAR: bool = True + VLLM_ROCM_USE_AITER_MOE: bool = True + VLLM_ROCM_AITER_MOE_DISPATCH_POLICY: int = 0 + VLLM_ROCM_USE_AITER_RMSNORM: bool = True + VLLM_ROCM_USE_AITER_MLA: bool = True + VLLM_ROCM_USE_AITER_MHA: bool = True + VLLM_ROCM_USE_AITER_FP4_ASM_GEMM: bool = False + VLLM_ROCM_USE_AITER_TRITON_ROPE: bool = False + VLLM_ROCM_USE_AITER_FP8BMM: bool = True + 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_USE_AITER_TRITON_GEMM: bool = True + VLLM_ROCM_USE_SKINNY_GEMM: bool = True + VLLM_ROCM_FP8_PADDING: bool = True + VLLM_ROCM_MOE_PADDING: bool = True + VLLM_ROCM_SHUFFLE_KV_CACHE_LAYOUT: bool = False + VLLM_ENABLE_V1_MULTIPROCESSING: bool = True + VLLM_LOG_BATCHSIZE_INTERVAL: float = -1 + VLLM_DISABLE_COMPILE_CACHE: bool = False + VLLM_USE_LAYERNAME: bool = True + Q_SCALE_CONSTANT: int = 200 + K_SCALE_CONSTANT: int = 200 + V_SCALE_CONSTANT: int = 100 + VLLM_USE_RUST_FRONTEND: bool = False + VLLM_RUST_FRONTEND_PATH: str | None = "auto" + VLLM_SERVER_DEV_MODE: bool = False + VLLM_V1_OUTPUT_PROC_CHUNK_SIZE: int = 128 + VLLM_MLA_DISABLE: bool = False + VLLM_RAY_PER_WORKER_GPUS: float = 1.0 + VLLM_RAY_BUNDLE_INDICES: str = "" + VLLM_CUDART_SO_PATH: str | None = None + VLLM_DP_RANK: int = 0 + VLLM_DP_RANK_LOCAL: int = -1 + VLLM_DP_SIZE: int = 1 + VLLM_USE_STANDALONE_COMPILE: bool = True + VLLM_ENABLE_PREGRAD_PASSES: bool = True + VLLM_USE_BREAKABLE_CUDAGRAPH: bool = False + VLLM_DP_MASTER_IP: str = "" + VLLM_DP_MASTER_PORT: int = 0 + VLLM_RANDOMIZE_DP_DUMMY_INPUTS: bool = False + VLLM_RAY_DP_PACK_STRATEGY: Literal["strict", "fill", "span"] = "strict" + VLLM_RAY_EXTRA_ENV_VAR_PREFIXES_TO_COPY: str = "" + VLLM_RAY_EXTRA_ENV_VARS_TO_COPY: str = "" + VLLM_MARLIN_USE_ATOMIC_ADD: bool = False + VLLM_MARLIN_INPUT_DTYPE: Literal["int8", "fp8"] | None = None + VLLM_HUMMING_ONLINE_QUANT_CONFIG: dict[str, Any] | None = None + VLLM_HUMMING_INPUT_QUANT_CONFIG: dict[str, Any] | None = None + VLLM_HUMMING_USE_F16_ACCUM: bool = False + VLLM_HUMMING_MOE_GEMM_TYPE: Literal["indexed", "grouped", "auto"] | None = None + VLLM_MXFP4_USE_MARLIN: bool | None = None + VLLM_DEEPEPLL_NVFP4_DISPATCH: bool = False + VLLM_V1_USE_OUTLINES_CACHE: bool = False + VLLM_TPU_BUCKET_PADDING_GAP: int = 0 + VLLM_TPU_MOST_MODEL_LEN: int | None = None + VLLM_TPU_USING_PATHWAYS: bool = False + VLLM_USE_DEEP_GEMM: bool = True + VLLM_MOE_USE_DEEP_GEMM: bool = True + VLLM_USE_DEEP_GEMM_E8M0: bool = True + VLLM_USE_DEEP_GEMM_TMA_ALIGNED_SCALES: bool = True + VLLM_DEEP_GEMM_WARMUP: Literal[ + "skip", + "full", + "relax", + ] = "relax" + VLLM_ENABLE_DEEPSEEK_V4_SPARSE_MLA_WARMUP: bool = True + VLLM_DSPARK_CONFIDENCE_THRESHOLD: str = "0.0" + VLLM_DSPARK_FORCE_DRAFT_LENGTH: str = "" + VLLM_DSPARK_REPLICATE_MARKOV_W1: bool = False + VLLM_DSPARK_STAGE_TIMING: bool = False + VLLM_DSPARK_STAGE_TIMING_LOG_EVERY: int = 20 + VLLM_DSPARK_ITER_TIMING: bool = False + VLLM_DSPARK_ITER_TIMING_LOG_EVERY: int = 20 + VLLM_DSPARK_TARGET_TIMING: bool = False + VLLM_DSPARK_TARGET_TIMING_LOG_EVERY: int = 20 + VLLM_DSV4_B12X_COMPRESSED_MLA: bool = False + VLLM_DSV4_DSPARK_DEFER_TARGET_CAPTURE: bool = False + VLLM_DSV4_DSPARK_DEFER_TARGET_CAPTURE_EXACT: bool = False + VLLM_USE_FUSED_MOE_GROUPED_TOPK: bool = True + VLLM_BLOCKSCALE_FP8_GEMM_FLASHINFER: bool = True + VLLM_USE_FLASHINFER_MOE_FP16: bool = False + VLLM_USE_FLASHINFER_MOE_FP8: bool = False + VLLM_USE_FLASHINFER_MOE_FP4: bool = False + VLLM_USE_FLASHINFER_MOE_INT4: bool = False + VLLM_FLASHINFER_MOE_BACKEND: Literal["throughput", "latency", "masked_gemm"] = ( + "latency" + ) + VLLM_FLASHINFER_AUTOTUNE_CACHE_DIR: str | None = None + VLLM_FLASHINFER_ALLREDUCE_BACKEND: Literal["auto", "trtllm", "mnnvl"] = "auto" + VLLM_ENABLE_PCIE_ALLREDUCE: bool = False + VLLM_FLASHINFER_WORKSPACE_BUFFER_SIZE: int = 394 * 1024 * 1024 + VLLM_XGRAMMAR_CACHE_MB: int = 0 + VLLM_MSGPACK_ZERO_COPY_THRESHOLD: int = 256 + VLLM_ALLOW_INSECURE_SERIALIZATION: bool = False + VLLM_DISABLE_REQUEST_ID_RANDOMIZATION: bool = False + VLLM_NIXL_SIDE_CHANNEL_HOST: str = "localhost" + VLLM_NIXL_SIDE_CHANNEL_PORT: int = 5600 + VLLM_MOONCAKE_BOOTSTRAP_PORT: int = 8998 + VLLM_MOONCAKE_STORE_TIER_LOG: bool = False + VLLM_MOONCAKE_DISK_STAGING_USABLE_RATIO: float = 0.9 + MOONCAKE_PREFERRED_SEGMENT: str | None = None + MOONCAKE_REQUESTER_LOCAL_HOSTNAME: str | None = None + VLLM_MAX_TOKENS_PER_EXPERT_FP4_MOE: int = 163840 + VLLM_TOOL_PARSE_REGEX_TIMEOUT_SECONDS: int = 1 + VLLM_MQ_MAX_CHUNK_BYTES_MB: int = 16 + VLLM_EXECUTE_MODEL_TIMEOUT_SECONDS: int = 300 + VLLM_KV_CACHE_LAYOUT: Literal["NHD", "HND"] | None = None + VLLM_SSM_CONV_STATE_LAYOUT: Literal["SD", "DS"] | None = None + VLLM_COMPUTE_NANS_IN_LOGITS: bool = False + VLLM_USE_NVFP4_CT_EMULATIONS: bool = False + VLLM_ROCM_QUICK_REDUCE_QUANTIZATION: Literal[ + "FP", "INT8", "INT6", "INT4", "NONE" + ] = "NONE" + VLLM_ROCM_QUICK_REDUCE_CAST_BF16_TO_FP16: bool = True + VLLM_ROCM_QUICK_REDUCE_MAX_SIZE_BYTES_MB: int | None = None + VLLM_ROCM_QUICK_REDUCE_MIN_SIZE_BYTES_MB: int | None = None + VLLM_ROCM_QUICK_REDUCE_QUANTIZATION_MIN_SIZE_KB: int | None = None + VLLM_MOONCAKE_ABORT_REQUEST_TIMEOUT: int = 480 + VLLM_ENABLE_CUDAGRAPH_GC: bool = False + VLLM_LOOPBACK_IP: str = "" + VLLM_ALLOW_CHUNKED_LOCAL_ATTN_WITH_HYBRID_KV_CACHE: bool = True + VLLM_ENABLE_RESPONSES_API_STORE: bool = False + VLLM_NVFP4_GEMM_BACKEND: str | None = None + VLLM_HAS_FLASHINFER_CUBIN: bool = False + VLLM_USE_FLASHINFER_MOE_MXFP4_MXFP8: bool = False + VLLM_USE_FLASHINFER_MOE_MXFP4_BF16: bool = False + VLLM_ROCM_FP8_MFMA_PAGE_ATTN: bool = False + VLLM_USE_FLASHINFER_MOE_MXFP4_MXFP8_CUTLASS: bool = False + VLLM_ALLREDUCE_USE_SYMM_MEM: bool = True + VLLM_ALLREDUCE_USE_FLASHINFER: bool = False + VLLM_TUNED_CONFIG_FOLDER: str | None = None + VLLM_GPT_OSS_SYSTEM_TOOL_MCP_LABELS: set[str] = set() + VLLM_USE_EXPERIMENTAL_PARSER_CONTEXT: bool = False + VLLM_GPT_OSS_HARMONY_SYSTEM_INSTRUCTIONS: bool = False + VLLM_SYSTEM_START_DATE: str | None = None + VLLM_TOOL_JSON_ERROR_AUTOMATIC_RETRY: bool = False + VLLM_ENFORCE_STRICT_TOOL_CALLING: bool = False + VLLM_CUSTOM_SCOPES_FOR_PROFILING: bool = False + VLLM_NVTX_SCOPES_FOR_PROFILING: bool = False + VLLM_KV_EVENTS_USE_INT_BLOCK_HASHES: bool = True + VLLM_OBJECT_STORAGE_SHM_BUFFER_NAME: str = "VLLM_OBJECT_STORAGE_SHM_BUFFER" + VLLM_DEEPEP_BUFFER_SIZE_MB: int = 1024 + VLLM_DEEPEP_HIGH_THROUGHPUT_FORCE_INTRA_NODE: bool = False + VLLM_DEEPEP_LOW_LATENCY_USE_MNNVL: bool = False + VLLM_DBO_COMM_SMS: int = 20 + VLLM_PATTERN_MATCH_DEBUG: str | None = None + VLLM_DEBUG_DUMP_PATH: str | None = None + VLLM_ENABLE_INDUCTOR_MAX_AUTOTUNE: bool = True + VLLM_ENABLE_INDUCTOR_COORDINATE_DESCENT_TUNING: bool = True + VLLM_USE_NCCL_SYMM_MEM: bool = False + VLLM_NCCL_INCLUDE_PATH: str | None = None + VLLM_USE_FBGEMM: bool = False + VLLM_GC_DEBUG: str = "" + VLLM_DEBUG_WORKSPACE: bool = False + VLLM_DISABLE_SHARED_EXPERTS_STREAM: bool = False + VLLM_SHARED_EXPERTS_STREAM_TOKEN_THRESHOLD: int = 256 + VLLM_MULTI_STREAM_GEMM_TOKEN_THRESHOLD: int = 1024 + VLLM_COMPILE_CACHE_SAVE_FORMAT: Literal["binary", "unpacked"] = "binary" + VLLM_USE_V2_MODEL_RUNNER: bool | None = None + VLLM_LOG_MODEL_INSPECTION: bool = False + VLLM_DEBUG_MFU_METRICS: bool = False + VLLM_WEIGHT_OFFLOADING_DISABLE_PIN_MEMORY: bool = False + VLLM_WEIGHT_OFFLOADING_DISABLE_UVA: bool = False + VLLM_DISABLE_LOG_LOGO: bool = False + VLLM_LORA_DISABLE_PDL: bool = False + VLLM_ENABLE_CUDA_COMPATIBILITY: bool = False + VLLM_CUDA_COMPATIBILITY_PATH: str | None = None + VLLM_SKIP_MODEL_NAME_VALIDATION: bool = False + """If set, vLLM will skip model name validation in API requests. + This allows any model name to be accepted in the 'model' field of requests, + making the server model-name agnostic. Useful for proxy/gateway scenarios.""" + VLLM_ELASTIC_EP_SCALE_UP_LAUNCH: bool = False + VLLM_ELASTIC_EP_DRAIN_REQUESTS: bool = False + VLLM_MEMORY_PROFILER_ESTIMATE_CUDAGRAPHS: bool = True + VLLM_NIXL_EP_MAX_NUM_RANKS: int = 32 + VLLM_XPU_ENABLE_XPU_GRAPH: bool = False + VLLM_XPU_USE_SAMPLER_KERNEL: bool = True + VLLM_LORA_ENABLE_DUAL_STREAM: bool = False + + +def get_default_cache_root(): + return os.getenv( + "XDG_CACHE_HOME", + os.path.join(os.path.expanduser("~"), ".cache"), + ) + + +def get_default_config_root(): + return os.getenv( + "XDG_CONFIG_HOME", + os.path.join(os.path.expanduser("~"), ".config"), + ) + + +def maybe_convert_int(value: str | None) -> int | None: + if value is None: + return None + return int(value) + + +def maybe_convert_bool(value: str | None) -> bool | None: + if value is None: + return None + return bool(int(value)) + + +def maybe_convert_json_str_or_file(value: str | None) -> dict[str, Any] | None: + if value is None: + return None + if os.path.exists(value): + with open(value) as f: + return json.load(f) + return json.loads(value) + + +def disable_compile_cache() -> bool: + return bool(int(os.getenv("VLLM_DISABLE_COMPILE_CACHE", "0"))) + + +def use_aot_compile() -> bool: + from vllm.utils.torch_utils import is_torch_equal_or_newer + + default_value = ( + "1" + if is_torch_equal_or_newer("2.10.0") and not disable_compile_cache() + else "0" + ) + + return os.environ.get("VLLM_USE_AOT_COMPILE", default_value) == "1" + + +def use_mega_aot_artifact(): + from vllm.utils.torch_utils import is_torch_equal_or_newer + + default_value = ( + "1" if is_torch_equal_or_newer("2.12.0.dev") and use_aot_compile() else "0" + ) + + return os.environ.get("VLLM_USE_MEGA_AOT_ARTIFACT", default_value) == "1" + + +def deprecated_env( + env_name: str, + removal_version: str, + replacement: str, + getter: Callable[[], Any], +) -> Callable[[], Any]: + """Wrap an env-var getter to emit a FutureWarning when the var is set.""" + + def _read() -> Any: + if env_name in os.environ: + warnings.warn( + f"{env_name} is deprecated and will be removed in " + f"{removal_version}. {replacement}", + FutureWarning, + stacklevel=2, + ) + return getter() + + return _read + + +def env_with_choices( + env_name: str, + default: str | None, + choices: list[str] | Callable[[], list[str]], + case_sensitive: bool = True, +) -> Callable[[], str | None]: + """ + Create a lambda that validates environment variable against allowed choices + + Args: + env_name: Name of the environment variable + default: Default value if not set (can be None) + choices: List of valid string options or callable that returns list + case_sensitive: Whether validation should be case sensitive + + Returns: + Lambda function for environment_variables dict + """ + + def _get_validated_env() -> str | None: + value = os.getenv(env_name) + if value is None: + return default + + # Resolve choices if it's a callable (for lazy loading) + actual_choices = choices() if callable(choices) else choices + + if not case_sensitive: + check_value = value.lower() + check_choices = [choice.lower() for choice in actual_choices] + else: + check_value = value + check_choices = actual_choices + + if check_value not in check_choices: + raise ValueError( + f"Invalid value '{value}' for {env_name}. " + f"Valid options: {actual_choices}." + ) + + return value + + return _get_validated_env + + +def env_list_with_choices( + env_name: str, + default: list[str], + choices: list[str] | Callable[[], list[str]], + case_sensitive: bool = True, +) -> Callable[[], list[str]]: + """ + Create a lambda that validates environment variable + containing comma-separated values against allowed choices + + Args: + env_name: Name of the environment variable + default: Default list of values if not set + choices: List of valid string options or callable that returns list + case_sensitive: Whether validation should be case sensitive + + Returns: + Lambda function for environment_variables + dict that returns list of strings + """ + + def _get_validated_env_list() -> list[str]: + value = os.getenv(env_name) + if value is None: + return default + + # Split comma-separated values and strip whitespace + values = [v.strip() for v in value.split(",") if v.strip()] + + if not values: + return default + + # Resolve choices if it's a callable (for lazy loading) + actual_choices = choices() if callable(choices) else choices + + # Validate each value + for val in values: + if not case_sensitive: + check_value = val.lower() + check_choices = [choice.lower() for choice in actual_choices] + else: + check_value = val + check_choices = actual_choices + + if check_value not in check_choices: + raise ValueError( + f"Invalid value '{val}' in {env_name}. " + f"Valid options: {actual_choices}." + ) + + return values + + return _get_validated_env_list + + +def env_set_with_choices( + env_name: str, + default: list[str], + choices: list[str] | Callable[[], list[str]], + case_sensitive: bool = True, +) -> Callable[[], set[str]]: + """ + Creates a lambda which that validates environment variable + containing comma-separated values against allowed choices which + returns choices as a set. + """ + + def _get_validated_env_set() -> set[str]: + return set(env_list_with_choices(env_name, default, choices, case_sensitive)()) + + return _get_validated_env_set + + +def get_vllm_port() -> int | None: + """Get the port from VLLM_PORT environment variable. + + Returns: + The port number as an integer if VLLM_PORT is set, None otherwise. + + Raises: + ValueError: If VLLM_PORT is a URI, suggest k8s service discovery issue. + """ + if "VLLM_PORT" not in os.environ: + return None + + port = os.getenv("VLLM_PORT", "0") + + try: + return int(port) + except ValueError as err: + from urllib3.util import parse_url + + parsed = parse_url(port) + if parsed.scheme: + raise ValueError( + f"VLLM_PORT '{port}' appears to be a URI. " + "This may be caused by a Kubernetes service discovery issue," + "check the warning in: https://docs.vllm.ai/en/stable/serving/env_vars.html" + ) from None + raise ValueError(f"VLLM_PORT '{port}' must be a valid integer") from err + + +def get_env_or_set_default( + env_name: str, + default_factory: Callable[[], str], +) -> Callable[[], str]: + """ + Create a lambda that returns an environment variable value if set, + or generates and sets a default value using the provided factory function. + """ + + def _get_or_set_default() -> str: + value = os.getenv(env_name) + if value is not None: + return value + + default_value = default_factory() + os.environ[env_name] = default_value + return default_value + + return _get_or_set_default + + +# The start-* and end* here are used by the documentation generator +# to extract the used env vars. + +# --8<-- [start:env-vars-definition] + +logger = logging.getLogger(__name__) + + +def _resolve_rust_frontend_path() -> str | None: + """Resolve the Rust frontend binary path. + + Returns None if VLLM_USE_RUST_FRONTEND is not enabled. + When enabled, resolves VLLM_RUST_FRONTEND_PATH ("auto" by default) + to the actual binary path. + """ + use_rust = bool(int(os.environ.get("VLLM_USE_RUST_FRONTEND", "0"))) + raw = os.environ.get("VLLM_RUST_FRONTEND_PATH", "auto") + + if not use_rust: + if os.environ.get("VLLM_RUST_FRONTEND_PATH") is not None: + logger.warning( + "VLLM_RUST_FRONTEND_PATH is set but VLLM_USE_RUST_FRONTEND " + "is not enabled. The Rust frontend will not be used. " + "Set VLLM_USE_RUST_FRONTEND=1 to enable it." + ) + return None + + if raw.lower() in ("auto", "1", "true"): + pkg_dir = os.path.dirname(os.path.abspath(__file__)) + candidate = os.path.join(pkg_dir, "vllm-rs") + if os.path.isfile(candidate) and os.access(candidate, os.X_OK): + return candidate + + raise FileNotFoundError( + "VLLM_RUST_FRONTEND_PATH=auto but the vllm-rs binary was " + f"not found at {candidate}. " + "Build with setuptools-rust or set the path explicitly." + ) + return raw + + +environment_variables: dict[str, Callable[[], Any]] = { + # ================== Installation Time Env Vars ================== + # Target device of vLLM, supporting [cuda (by default), + # rocm, cpu] + "VLLM_TARGET_DEVICE": lambda: os.getenv("VLLM_TARGET_DEVICE", "cuda").lower(), + # Main CUDA version of vLLM. This follows PyTorch but can be overridden. + "VLLM_MAIN_CUDA_VERSION": lambda: ( + os.getenv("VLLM_MAIN_CUDA_VERSION", "").lower() or "13.0" + ), + # Controls PyTorch float32 matmul precision mode within vLLM workers. + # Valid options mirror torch.set_float32_matmul_precision + "VLLM_FLOAT32_MATMUL_PRECISION": env_with_choices( + "VLLM_FLOAT32_MATMUL_PRECISION", + "highest", + ["highest", "high", "medium"], + case_sensitive=False, + ), + # Enable batch-invariant mode: deterministic results regardless of + # batch composition. Requires NVIDIA GPU with compute capability >= 9.0. + "VLLM_BATCH_INVARIANT": lambda: bool(int(os.getenv("VLLM_BATCH_INVARIANT", "0"))), + # Use tensor descriptors for Q/K/V loads and output stores in the + # Triton unified-attention kernel. Enables HW 2D block reads on + # Intel Xe2/Xe3; the non-TD branch is dead-code-eliminated at Triton + # compile time so other platforms see no overhead. Tri-state override: + # unset (default) lets the `triton_attn` backend auto-select per + # platform (currently auto-enabled on XPU only); ``1`` forces TD on; + # ``0`` forces TD off. Useful for A/B benchmarking the TD path. + "VLLM_TRITON_ATTN_USE_TD": lambda: {"1": True, "0": False}.get( + os.getenv("VLLM_TRITON_ATTN_USE_TD", "").strip() + ), + # Maximum number of compilation jobs to run in parallel. + # By default this is the number of CPUs + "MAX_JOBS": lambda: os.getenv("MAX_JOBS", None), + # Number of threads to use for nvcc + # By default this is 1. + # If set, `MAX_JOBS` will be reduced to avoid oversubscribing the CPU. + "NVCC_THREADS": lambda: os.getenv("NVCC_THREADS", None), + # If set, vllm will use precompiled native binaries (*.so and vllm-rs). + "VLLM_USE_PRECOMPILED": lambda: ( + os.environ.get("VLLM_USE_PRECOMPILED", "").strip().lower() in ("1", "true") + or bool(os.environ.get("VLLM_PRECOMPILED_WHEEL_LOCATION")) + ), + # If set, vllm will use the precompiled Rust frontend binary (vllm-rs). + "VLLM_USE_PRECOMPILED_RUST": lambda: ( + os.environ.get("VLLM_USE_PRECOMPILED_RUST", "").strip().lower() in ("1", "true") + ), + # If set, skip adding +precompiled suffix to version string + "VLLM_SKIP_PRECOMPILED_VERSION_SUFFIX": lambda: bool( + int(os.environ.get("VLLM_SKIP_PRECOMPILED_VERSION_SUFFIX", "0")) + ), + # Used to mark that setup.py is running in a Docker build context, + # in order to force the use of precompiled binaries. + "VLLM_DOCKER_BUILD_CONTEXT": lambda: ( + os.environ.get("VLLM_DOCKER_BUILD_CONTEXT", "").strip().lower() in ("1", "true") + ), + # CMake build type + # If not set, defaults to "Debug" or "RelWithDebInfo" + # Available options: "Debug", "Release", "RelWithDebInfo" + "CMAKE_BUILD_TYPE": env_with_choices( + "CMAKE_BUILD_TYPE", None, ["Debug", "Release", "RelWithDebInfo"] + ), + # If set, vllm will print verbose logs during installation + "VERBOSE": lambda: bool(int(os.getenv("VERBOSE", "0"))), + # Root directory for vLLM configuration files + # Defaults to `~/.config/vllm` unless `XDG_CONFIG_HOME` is set + # Note that this not only affects how vllm finds its configuration files + # during runtime, but also affects how vllm installs its configuration + # files during **installation**. + "VLLM_CONFIG_ROOT": lambda: os.path.expanduser( + os.getenv( + "VLLM_CONFIG_ROOT", + os.path.join(get_default_config_root(), "vllm"), + ) + ), + # ================== Runtime Env Vars ================== + # Root directory for vLLM cache files + # Defaults to `~/.cache/vllm` unless `XDG_CACHE_HOME` is set + "VLLM_CACHE_ROOT": lambda: os.path.expanduser( + os.getenv( + "VLLM_CACHE_ROOT", + os.path.join(get_default_cache_root(), "vllm"), + ) + ), + # used in distributed environment to determine the ip address + # of the current node, when the node has multiple network interfaces. + # If you are using multi-node inference, you should set this differently + # on each node. + "VLLM_HOST_IP": lambda: os.getenv("VLLM_HOST_IP", ""), + # used in distributed environment to manually set the communication port + # Note: if VLLM_PORT is set, and some code asks for multiple ports, the + # VLLM_PORT will be used as the first port, and the rest will be generated + # by incrementing the VLLM_PORT value. + "VLLM_PORT": get_vllm_port, + # path used for ipc when the frontend api server is running in + # multi-processing mode to communicate with the backend engine process. + "VLLM_RPC_BASE_PATH": lambda: os.getenv( + "VLLM_RPC_BASE_PATH", tempfile.gettempdir() + ), + # If true, will load models from ModelScope instead of Hugging Face Hub. + # note that the value is true or false, not numbers + "VLLM_USE_MODELSCOPE": lambda: ( + os.environ.get("VLLM_USE_MODELSCOPE", "False").lower() == "true" + ), + # If true, replace the Rust BPE backend that powers HF fast tokenizers + # with the `fastokens` (https://github.com/crusoecloud/fastokens) shim. + # Applies to any tokenizer mode that loads an HF fast tokenizer + # (`hf`, `deepseek_v32`, `deepseek_v4`, `qwen_vl`, …). The `fastokens` + # Python package must be installed. + "VLLM_USE_FASTOKENS": lambda: bool(int(os.getenv("VLLM_USE_FASTOKENS", "0"))), + # Interval in seconds to log a warning message when the ring buffer is full + "VLLM_RINGBUFFER_WARNING_INTERVAL": lambda: int( + os.environ.get("VLLM_RINGBUFFER_WARNING_INTERVAL", "60") + ), + # path to cudatoolkit home directory, under which should be bin, include, + # and lib directories. + "CUDA_HOME": lambda: os.environ.get("CUDA_HOME", None), + # Path to the NCCL library file. It is needed because nccl>=2.19 brought + # by PyTorch contains a bug: https://github.com/NVIDIA/nccl/issues/1234 + "VLLM_NCCL_SO_PATH": lambda: os.environ.get("VLLM_NCCL_SO_PATH", None), + # when `VLLM_NCCL_SO_PATH` is not set, vllm will try to find the nccl + # library file in the locations specified by `LD_LIBRARY_PATH` + "LD_LIBRARY_PATH": lambda: os.environ.get("LD_LIBRARY_PATH", None), + # flag to control the chunk size (in MB) for sleeping memory allocations under ROCm + "VLLM_ROCM_SLEEP_MEM_CHUNK_SIZE": lambda: int( + os.environ.get("VLLM_ROCM_SLEEP_MEM_CHUNK_SIZE", "256") + ), + # Feature flag to enable/disable Inductor standalone compile. + # In torch <= 2.7 we ignore this flag; in torch >= 2.9 this is + # enabled by default. + "VLLM_USE_STANDALONE_COMPILE": lambda: ( + os.environ.get("VLLM_USE_STANDALONE_COMPILE", "1") == "1" + ), + # Inductor's pre-grad passes don't do anything for vLLM. + # The pre-grad passes get run even on cache-hit and negatively impact + # vllm cold compile times by O(1s) + # Can remove this after the following issue gets fixed + # TODO(luka): maybe_inplace requires this + # https://github.com/pytorch/pytorch/issues/174502 + "VLLM_ENABLE_PREGRAD_PASSES": lambda: ( + os.environ.get("VLLM_ENABLE_PREGRAD_PASSES", "1") == "1" + ), + # Experimental: breakable cudagraph does not rely on torch.compile + "VLLM_USE_BREAKABLE_CUDAGRAPH": lambda: ( + os.environ.get("VLLM_USE_BREAKABLE_CUDAGRAPH", "0") == "1" + ), + # Debug pattern matching inside custom passes. + # Should be set to the fx.Node name (e.g. 'getitem_34' or 'scaled_mm_3'). + "VLLM_PATTERN_MATCH_DEBUG": lambda: os.environ.get( + "VLLM_PATTERN_MATCH_DEBUG", None + ), + # Dump fx graphs to the given directory. + # It will override CompilationConfig.debug_dump_path if set. + "VLLM_DEBUG_DUMP_PATH": lambda: os.environ.get("VLLM_DEBUG_DUMP_PATH", None), + # Feature flag to enable/disable AOT compilation. This will ensure + # compilation is done in warmup phase and the compilation will be + # reused in subsequent calls. + "VLLM_USE_AOT_COMPILE": use_aot_compile, + # Feature flag to enable/disable bytecode in + # TorchCompileWithNoGuardsWrapper. + "VLLM_USE_BYTECODE_HOOK": lambda: bool( + int(os.environ.get("VLLM_USE_BYTECODE_HOOK", "1")) + ), + # Force vllm to always load AOT compiled models from disk. Failure + # to load will result in a hard error when this is enabled. + # Will be ignored when VLLM_USE_AOT_COMPILE is disabled. + "VLLM_FORCE_AOT_LOAD": lambda: os.environ.get("VLLM_FORCE_AOT_LOAD", "0") == "1", + # Enable loading compiled models directly from cached standalone compile artifacts + # without re-splitting graph modules. This reduces overhead during model + # loading by using reconstruct_serializable_fn_from_mega_artifact. + "VLLM_USE_MEGA_AOT_ARTIFACT": use_mega_aot_artifact, + # local rank of the process in the distributed setting, used to determine + # the GPU device id + "LOCAL_RANK": lambda: int(os.environ.get("LOCAL_RANK", "0")), + # used to control the visible devices in the distributed setting + "CUDA_VISIBLE_DEVICES": lambda: os.environ.get("CUDA_VISIBLE_DEVICES", None), + # timeout for each iteration in the engine + "VLLM_ENGINE_ITERATION_TIMEOUT_S": lambda: int( + os.environ.get("VLLM_ENGINE_ITERATION_TIMEOUT_S", "60") + ), + # Timeout in seconds for waiting for engine cores to become ready + # during startup. Default is 600 seconds (10 minutes). + "VLLM_ENGINE_READY_TIMEOUT_S": lambda: int( + os.environ.get("VLLM_ENGINE_READY_TIMEOUT_S", "600") + ), + # API key for vLLM API server + "VLLM_API_KEY": lambda: os.environ.get("VLLM_API_KEY", None), + # Whether to log responses from API Server for debugging + "VLLM_DEBUG_LOG_API_SERVER_RESPONSE": lambda: ( + os.environ.get("VLLM_DEBUG_LOG_API_SERVER_RESPONSE", "False").lower() == "true" + ), + # S3 access information, used for tensorizer to load model from S3 + "S3_ACCESS_KEY_ID": lambda: os.environ.get("S3_ACCESS_KEY_ID", None), + "S3_SECRET_ACCESS_KEY": lambda: os.environ.get("S3_SECRET_ACCESS_KEY", None), + "S3_ENDPOINT_URL": lambda: os.environ.get("S3_ENDPOINT_URL", None), + # Usage stats collection + "VLLM_USAGE_STATS_SERVER": lambda: os.environ.get( + "VLLM_USAGE_STATS_SERVER", "https://stats.vllm.ai" + ), + "VLLM_NO_USAGE_STATS": lambda: os.environ.get("VLLM_NO_USAGE_STATS", "0") == "1", + "VLLM_DO_NOT_TRACK": lambda: ( + ( + os.environ.get("VLLM_DO_NOT_TRACK", None) + or os.environ.get("DO_NOT_TRACK", None) + or "0" + ) + == "1" + ), + "VLLM_USAGE_SOURCE": lambda: os.environ.get("VLLM_USAGE_SOURCE", "production"), + # Logging configuration + # If set to 0, vllm will not configure logging + # If set to 1, vllm will configure logging using the default configuration + # or the configuration file specified by VLLM_LOGGING_CONFIG_PATH + "VLLM_CONFIGURE_LOGGING": lambda: bool( + int(os.getenv("VLLM_CONFIGURE_LOGGING", "1")) + ), + "VLLM_LOGGING_CONFIG_PATH": lambda: os.getenv("VLLM_LOGGING_CONFIG_PATH"), + # this is used for configuring the default logging level + "VLLM_LOGGING_LEVEL": lambda: os.getenv("VLLM_LOGGING_LEVEL", "INFO").upper(), + # this is used for configuring the default logging stream + "VLLM_LOGGING_STREAM": lambda: os.getenv("VLLM_LOGGING_STREAM", "ext://sys.stdout"), + # if set, VLLM_LOGGING_PREFIX will be prepended to all log messages + "VLLM_LOGGING_PREFIX": lambda: os.getenv("VLLM_LOGGING_PREFIX", ""), + # Controls colored logging output. Options: "auto" (default, colors when terminal), + # "1" (always use colors), "0" (never use colors) + "VLLM_LOGGING_COLOR": lambda: os.getenv("VLLM_LOGGING_COLOR", "auto"), + # Standard unix flag for disabling ANSI color codes + "NO_COLOR": lambda: os.getenv("NO_COLOR", "0") != "0", + # If set, vllm will log stats at this interval in seconds + # If not set, vllm will log stats every 10 seconds. + "VLLM_LOG_STATS_INTERVAL": lambda: ( + val + if (val := float(os.getenv("VLLM_LOG_STATS_INTERVAL", "10."))) > 0.0 + else 10.0 + ), + # Trace function calls + # If set to 1, vllm will trace function calls + # Useful for debugging + "VLLM_TRACE_FUNCTION": lambda: int(os.getenv("VLLM_TRACE_FUNCTION", "0")), + # Whether to use the FlashInfer top-k / top-p sampler on CUDA. Enabled + # by default when the hardware supports it — set to 0 to opt out + # explicitly, which forces the PyTorch-native (Triton for bs>=8) path. + "VLLM_USE_FLASHINFER_SAMPLER": lambda: ( + bool(int(os.environ["VLLM_USE_FLASHINFER_SAMPLER"])) + if "VLLM_USE_FLASHINFER_SAMPLER" in os.environ + else True + ), + # Pipeline stage partition strategy + "VLLM_PP_LAYER_PARTITION": lambda: os.getenv("VLLM_PP_LAYER_PARTITION", None), + # (CPU backend only) CPU key-value cache space. + # default is None and will be set as 4 GB + "VLLM_CPU_KVCACHE_SPACE": lambda: ( + int(os.getenv("VLLM_CPU_KVCACHE_SPACE", "0")) + if "VLLM_CPU_KVCACHE_SPACE" in os.environ + else None + ), + # (CPU backend only) CPU core ids bound by OpenMP threads, e.g., "0-31", + # "0,1,2", "0-31,33". CPU cores of different ranks are separated by '|'. + "VLLM_CPU_OMP_THREADS_BIND": lambda: os.getenv("VLLM_CPU_OMP_THREADS_BIND", "auto"), + # (CPU backend only) CPU cores not used by OMP threads . + # Those CPU cores will not be used by OMP threads of a rank. + "VLLM_CPU_NUM_OF_RESERVED_CPU": lambda: ( + int(os.getenv("VLLM_CPU_NUM_OF_RESERVED_CPU", "0")) + if "VLLM_CPU_NUM_OF_RESERVED_CPU" in os.environ + else None + ), + # (CPU backend only) whether to use SGL kernels, optimized for small batch. + "VLLM_CPU_SGL_KERNEL": lambda: bool(int(os.getenv("VLLM_CPU_SGL_KERNEL", "0"))), + # (CPU backend only) whether to enable attention spilt KV. + "VLLM_CPU_ATTN_SPLIT_KV": lambda: bool( + int(os.getenv("VLLM_CPU_ATTN_SPLIT_KV", "1")) + ), + # (Zen CPU backend) eagerly prepack weights into ZenDNN blocked layout + # at model load time. Eliminates per-inference layout conversion overhead. + "VLLM_ZENTORCH_WEIGHT_PREPACK": lambda: bool( + int(os.getenv("VLLM_ZENTORCH_WEIGHT_PREPACK", "1")) + ), + # (CPU backend only) whether to use SGLang INT4 W4A8 kernels for AWQ. + "VLLM_CPU_INT4_W4A8": lambda: bool(int(os.getenv("VLLM_CPU_INT4_W4A8", "1"))), + # If the env var is set, Ray Compiled Graph uses the specified + # channel type to communicate between workers belonging to + # different pipeline-parallel stages. + # Available options: + # - "auto": use the default channel type + # - "nccl": use NCCL for communication + # - "shm": use shared memory and gRPC for communication + "VLLM_USE_RAY_COMPILED_DAG_CHANNEL_TYPE": env_with_choices( + "VLLM_USE_RAY_COMPILED_DAG_CHANNEL_TYPE", "auto", ["auto", "nccl", "shm"] + ), + # If the env var is set, it enables GPU communication overlap + # (experimental feature) in Ray's Compiled Graph. + "VLLM_USE_RAY_COMPILED_DAG_OVERLAP_COMM": lambda: bool( + int(os.getenv("VLLM_USE_RAY_COMPILED_DAG_OVERLAP_COMM", "0")) + ), + # If the env var is set, it uses a Ray Communicator wrapping + # vLLM's pipeline parallelism communicator to interact with Ray's + # Compiled Graph. Otherwise, it uses Ray's NCCL communicator. + "VLLM_USE_RAY_WRAPPED_PP_COMM": lambda: bool( + int(os.getenv("VLLM_USE_RAY_WRAPPED_PP_COMM", "1")) + ), + # When True and distributed_executor_backend="ray", use RayExecutorV2 + # (MQ-based) instead of RayDistributedExecutor (compiled-graph backend). + "VLLM_USE_RAY_V2_EXECUTOR_BACKEND": lambda: bool( + int(os.getenv("VLLM_USE_RAY_V2_EXECUTOR_BACKEND", "1")) + ), + # Use dedicated multiprocess context for workers. + # Both spawn and fork work + "VLLM_WORKER_MULTIPROC_METHOD": env_with_choices( + "VLLM_WORKER_MULTIPROC_METHOD", "fork", ["spawn", "fork"] + ), + # Path to the cache for storing downloaded assets + "VLLM_ASSETS_CACHE": lambda: os.path.expanduser( + os.getenv( + "VLLM_ASSETS_CACHE", + os.path.join(get_default_cache_root(), "vllm", "assets"), + ) + ), + # If the env var is set, we will clean model file in + # this path $VLLM_ASSETS_CACHE/model_streamer/$model_name + "VLLM_ASSETS_CACHE_MODEL_CLEAN": lambda: bool( + int(os.getenv("VLLM_ASSETS_CACHE_MODEL_CLEAN", "0")) + ), + # Timeout for fetching images when serving multimodal models + # Default is 5 seconds + "VLLM_IMAGE_FETCH_TIMEOUT": lambda: int(os.getenv("VLLM_IMAGE_FETCH_TIMEOUT", "5")), + # Timeout for fetching videos when serving multimodal models + # Default is 30 seconds + "VLLM_VIDEO_FETCH_TIMEOUT": lambda: int( + os.getenv("VLLM_VIDEO_FETCH_TIMEOUT", "30") + ), + # Timeout for fetching audio when serving multimodal models + # Default is 10 seconds + "VLLM_AUDIO_FETCH_TIMEOUT": lambda: int( + os.getenv("VLLM_AUDIO_FETCH_TIMEOUT", "10") + ), + # Directory for caching media downloads (images, video, audio fetched + # from URLs during inference). Empty string disables caching. + "VLLM_MEDIA_CACHE": lambda: os.getenv("VLLM_MEDIA_CACHE", ""), + # Maximum cache size in MB. When exceeded, least-recently-used entries + # are evicted. Default is 5120 (5 GB). + "VLLM_MEDIA_CACHE_MAX_SIZE_MB": lambda: int( + os.getenv("VLLM_MEDIA_CACHE_MAX_SIZE_MB", "5120") + ), + # Time-to-live in hours for cached media files. Entries older than this + # are evicted regardless of cache size. Default is 24 hours. + "VLLM_MEDIA_CACHE_TTL_HOURS": lambda: float( + os.getenv("VLLM_MEDIA_CACHE_TTL_HOURS", "24") + ), + # Maximum number of retries for fetching media (images, audio, video) + # from URLs. Each retry quadruples the timeout. Default is 3. + "VLLM_MEDIA_FETCH_MAX_RETRIES": lambda: int( + os.getenv("VLLM_MEDIA_FETCH_MAX_RETRIES", "3") + ), + # Whether to allow HTTP redirects when fetching from media URLs. + # Default to True + "VLLM_MEDIA_URL_ALLOW_REDIRECTS": lambda: bool( + int(os.getenv("VLLM_MEDIA_URL_ALLOW_REDIRECTS", "1")) + ), + # Max number of workers for the thread pool handling + # media bytes loading. Set to 1 to disable parallel processing. + # Default is 8 + "VLLM_MEDIA_LOADING_THREAD_COUNT": lambda: int( + os.getenv("VLLM_MEDIA_LOADING_THREAD_COUNT", "8") + ), + # Maximum filesize in MB for a single audio file when processing + # speech-to-text requests. Files larger than this will be rejected. + # Default is 25 MB + "VLLM_MAX_AUDIO_CLIP_FILESIZE_MB": lambda: int( + os.getenv("VLLM_MAX_AUDIO_CLIP_FILESIZE_MB", "25") + ), + # Backend for Video IO — selects the frame-sampling algorithm. + # - "opencv": uniform sampling. + # - "opencv_dynamic": duration-aware dynamic sampling. + # - "identity": returns raw video bytes for model processor to handle. + # + # Custom backend implementations can be registered + # via `@VIDEO_LOADER_REGISTRY.register("my_custom_video_loader")` and + # imported at runtime. + # If a non-existing backend is used, an AssertionError will be thrown. + "VLLM_VIDEO_LOADER_BACKEND": lambda: os.getenv( + "VLLM_VIDEO_LOADER_BACKEND", "opencv" + ), + # Media connector implementation. + # - "http": Default connector that supports fetching media via HTTP. + # + # Custom implementations can be registered + # via `@MEDIA_CONNECTOR_REGISTRY.register("my_custom_media_connector")` and + # imported at runtime. + # If a non-existing backend is used, an AssertionError will be thrown. + "VLLM_MEDIA_CONNECTOR": lambda: os.getenv("VLLM_MEDIA_CONNECTOR", "http"), + # Hash algorithm for multimodal content hashing. + # - "blake3": Default, fast cryptographic hash (not FIPS 140-3 compliant) + # - "sha256": FIPS 140-3 compliant, widely supported + # - "sha512": FIPS 140-3 compliant, faster on 64-bit systems + # Use sha256 or sha512 for FIPS compliance in government/enterprise deployments + "VLLM_MM_HASHER_ALGORITHM": env_with_choices( + "VLLM_MM_HASHER_ALGORITHM", + "blake3", + ["blake3", "sha256", "sha512"], + case_sensitive=False, + ), + # Path to the XLA persistent cache directory. + # Only used for XLA devices such as TPUs. + "VLLM_XLA_CACHE_PATH": lambda: os.path.expanduser( + os.getenv( + "VLLM_XLA_CACHE_PATH", + os.path.join(get_default_cache_root(), "vllm", "xla_cache"), + ) + ), + # If set, assert on XLA recompilation after each execution step. + "VLLM_XLA_CHECK_RECOMPILATION": lambda: bool( + int(os.getenv("VLLM_XLA_CHECK_RECOMPILATION", "0")) + ), + # Enable SPMD mode for TPU backend. + "VLLM_XLA_USE_SPMD": lambda: bool(int(os.getenv("VLLM_XLA_USE_SPMD", "0"))), + # Maximum size (in MB) for logits tensor in sparse MLA indexer prefill chunks. + # Bounds the [M, N] float32 logits tensor to prevent CUDA OOM. + # Default: 512 MB + "VLLM_SPARSE_INDEXER_MAX_LOGITS_MB": lambda: int( + os.getenv("VLLM_SPARSE_INDEXER_MAX_LOGITS_MB", "512") + ), + # Use b12x for the DeepSeek V4 C4 sparse indexer and its top-k selection. + # This is opt-in while the b12x subsystems are brought over one at a time. + "VLLM_USE_B12X_SPARSE_INDEXER": lambda: bool( + int(os.getenv("VLLM_USE_B12X_SPARSE_INDEXER", "0")) + ), + # Use b12x for DeepSeek V4 mHC pre/post residual mixing. + # This is opt-in while the b12x subsystems are brought over one at a time. + "VLLM_USE_B12X_MHC": lambda: bool(int(os.getenv("VLLM_USE_B12X_MHC", "0"))), + # Use b12x for block-scaled FP8 linear GEMMs. + # This is opt-in while the b12x subsystems are brought over one at a time. + "VLLM_USE_B12X_FP8_GEMM": lambda: bool( + int(os.getenv("VLLM_USE_B12X_FP8_GEMM", "0")) + ), + # Use b12x for the DeepSeek V4 WO-A/WO-B fused projection. + # This is separate from the generic FP8 linear switch for perf isolation. + "VLLM_USE_B12X_WO_PROJECTION": lambda: bool( + int(os.getenv("VLLM_USE_B12X_WO_PROJECTION", "0")) + ), + # Use b12x for DeepSeek V4 MXFP4 MoE experts. + # This is opt-in while the b12x subsystems are brought over one at a time. + "VLLM_USE_B12X_MOE": lambda: bool(int(os.getenv("VLLM_USE_B12X_MOE", "0"))), + # Experimental B12X W4A16 MoE selector override for DSpark decode A/Bs. + # 0 keeps the upstream selector result unchanged. + "VLLM_B12X_W4A16_FORCE_BLOCKS_PER_SM": lambda: int( + os.getenv("VLLM_B12X_W4A16_FORCE_BLOCKS_PER_SM", "0") + ), + "VLLM_B12X_W4A16_FORCE_BLOCKS_MAX_M": lambda: int( + os.getenv("VLLM_B12X_W4A16_FORCE_BLOCKS_MAX_M", "16") + ), + "VLLM_B12X_W4A16_FORCE_TILE_CONFIG": lambda: os.getenv( + "VLLM_B12X_W4A16_FORCE_TILE_CONFIG", "" + ), + # If set, the OpenAI API server will stay alive even after the underlying + # AsyncLLMEngine errors and stops serving requests + "VLLM_KEEP_ALIVE_ON_ENGINE_DEATH": lambda: bool( + int(os.getenv("VLLM_KEEP_ALIVE_ON_ENGINE_DEATH", "0")) + ), + # If the env var VLLM_ALLOW_LONG_MAX_MODEL_LEN is set, it allows + # the user to specify a max sequence length greater than + # the max length derived from the model's config.json. + # To enable this, set VLLM_ALLOW_LONG_MAX_MODEL_LEN=1. + "VLLM_ALLOW_LONG_MAX_MODEL_LEN": lambda: ( + os.environ.get("VLLM_ALLOW_LONG_MAX_MODEL_LEN", "0").strip().lower() + in ("1", "true") + ), + # If set, forces FP8 Marlin to be used for FP8 quantization regardless + # of the hardware support for FP8 compute. + "VLLM_TEST_FORCE_FP8_MARLIN": lambda: ( + os.environ.get("VLLM_TEST_FORCE_FP8_MARLIN", "0").strip().lower() + in ("1", "true") + ), + "VLLM_TEST_FORCE_LOAD_FORMAT": lambda: os.getenv( + "VLLM_TEST_FORCE_LOAD_FORMAT", "dummy" + ), + # Time in ms for the zmq client to wait for a response from the backend + # server for simple data operations + "VLLM_RPC_TIMEOUT": lambda: int(os.getenv("VLLM_RPC_TIMEOUT", "10000")), + # Timeout in seconds for keeping HTTP connections alive in API server + "VLLM_HTTP_TIMEOUT_KEEP_ALIVE": lambda: int( + os.environ.get("VLLM_HTTP_TIMEOUT_KEEP_ALIVE", "5") + ), + # Maximum allowed value for the `n` sampling parameter (number of output + # sequences per request). Limits resource consumption to prevent + # denial-of-service via excessively large fan-out. Default: 16384. + "VLLM_MAX_N_SEQUENCES": lambda: int( + os.environ.get("VLLM_MAX_N_SEQUENCES", "16384") + ), + # a list of plugin names to load, separated by commas. + # if this is not set, it means all plugins will be loaded + # if this is set to an empty string, no plugins will be loaded + "VLLM_PLUGINS": lambda: ( + None + if "VLLM_PLUGINS" not in os.environ + else os.environ["VLLM_PLUGINS"].split(",") + ), + # a local directory to look in for unrecognized LoRA adapters. + # only works if plugins are enabled and + # VLLM_ALLOW_RUNTIME_LORA_UPDATING is enabled. + "VLLM_LORA_RESOLVER_CACHE_DIR": lambda: os.getenv( + "VLLM_LORA_RESOLVER_CACHE_DIR", None + ), + # A remote HF repo(s) containing one or more LoRA adapters, which + # may be downloaded and leveraged as needed. Only works if plugins + # are enabled and VLLM_ALLOW_RUNTIME_LORA_UPDATING is enabled. + # Values should be comma separated. + "VLLM_LORA_RESOLVER_HF_REPO_LIST": lambda: os.getenv( + "VLLM_LORA_RESOLVER_HF_REPO_LIST", None + ), + # If set, vLLM will use Triton implementations of AWQ. + "VLLM_USE_TRITON_AWQ": lambda: bool(int(os.getenv("VLLM_USE_TRITON_AWQ", "0"))), + # If set, allow loading or unloading lora adapters in runtime, + "VLLM_ALLOW_RUNTIME_LORA_UPDATING": lambda: ( + os.environ.get("VLLM_ALLOW_RUNTIME_LORA_UPDATING", "0").strip().lower() + in ("1", "true") + ), + # We assume drivers can report p2p status correctly. + # If the program hangs when using custom allreduce, + # potantially caused by a bug in the driver (535 series), + # if might be helpful to set VLLM_SKIP_P2P_CHECK=0 + # so that vLLM can verify if p2p is actually working. + # See https://github.com/vllm-project/vllm/blob/a9b15c606fea67a072416ea0ea115261a2756058/vllm/distributed/device_communicators/custom_all_reduce_utils.py#L101-L108 for details. # noqa + "VLLM_SKIP_P2P_CHECK": lambda: os.getenv("VLLM_SKIP_P2P_CHECK", "1") == "1", + # List of quantization kernels that should be disabled, used for testing + # and performance comparisons. Currently only affects MPLinearKernel + # selection + # (kernels: MacheteLinearKernel, MarlinLinearKernel, ExllamaLinearKernel) + "VLLM_DISABLED_KERNELS": lambda: ( + [] + if "VLLM_DISABLED_KERNELS" not in os.environ + else os.environ["VLLM_DISABLED_KERNELS"].split(",") + ), + "VLLM_ENABLE_FLA_PACKED_RECURRENT_DECODE": lambda: bool( + int(os.getenv("VLLM_ENABLE_FLA_PACKED_RECURRENT_DECODE", "1")) + ), + # Disable pynccl (using torch.distributed instead) + "VLLM_DISABLE_PYNCCL": lambda: ( + os.getenv("VLLM_DISABLE_PYNCCL", "False").lower() in ("true", "1") + ), + # Optional: enable external Oink custom ops (e.g., Blackwell RMSNorm). + # Disabled by default. + "VLLM_USE_OINK_OPS": lambda: ( + os.getenv("VLLM_USE_OINK_OPS", "False").lower() in ("true", "1") + ), + # Disable aiter ops unless specifically enabled. + # Acts as a parent switch to enable the rest of the other operations. + "VLLM_ROCM_USE_AITER": lambda: ( + os.getenv("VLLM_ROCM_USE_AITER", "False").lower() in ("true", "1") + ), + # Whether to use aiter paged attention. + # By default is disabled. + "VLLM_ROCM_USE_AITER_PAGED_ATTN": lambda: ( + os.getenv("VLLM_ROCM_USE_AITER_PAGED_ATTN", "False").lower() in ("true", "1") + ), + # use aiter linear op if aiter ops are enabled + # The following list of related ops + # - scaled_mm (per-tensor / rowwise) + # - use aiter tuned gemms for unquantized gemms + "VLLM_ROCM_USE_AITER_LINEAR": lambda: ( + os.getenv("VLLM_ROCM_USE_AITER_LINEAR", "True").lower() in ("true", "1") + ), + # Whether to use aiter moe ops. + # By default is enabled. + "VLLM_ROCM_USE_AITER_MOE": lambda: ( + os.getenv("VLLM_ROCM_USE_AITER_MOE", "True").lower() in ("true", "1") + ), + # MoE sorting dispatch policy for AITER fused MoE kernels. + # 0 = auto (default): single-pass for small batches, multi-pass + # for large batches + # 1 = always single-pass: one kernel launch, no workspace, + # may be preferred for low-concurrency decode workloads + # 2 = always multi-pass: can be faster for MoE-heavy models + # (e.g., +2-5% on Qwen3-Next, +1.5% on DeepSeek-V3 at TP4, + # see PR #39177 for benchmarks) + "VLLM_ROCM_AITER_MOE_DISPATCH_POLICY": lambda: int( + os.getenv("VLLM_ROCM_AITER_MOE_DISPATCH_POLICY", "0") + ), + # use aiter rms norm op if aiter ops are enabled. + "VLLM_ROCM_USE_AITER_RMSNORM": lambda: ( + os.getenv("VLLM_ROCM_USE_AITER_RMSNORM", "True").lower() in ("true", "1") + ), + # Whether to use aiter mla ops. + # By default is enabled. + "VLLM_ROCM_USE_AITER_MLA": lambda: ( + os.getenv("VLLM_ROCM_USE_AITER_MLA", "True").lower() in ("true", "1") + ), + # Whether to use aiter mha ops. + # By default is enabled. + "VLLM_ROCM_USE_AITER_MHA": lambda: ( + os.getenv("VLLM_ROCM_USE_AITER_MHA", "True").lower() in ("true", "1") + ), + # Whether to use aiter fp4 gemm asm. + # By default is disabled. + "VLLM_ROCM_USE_AITER_FP4_ASM_GEMM": lambda: ( + os.getenv("VLLM_ROCM_USE_AITER_FP4_ASM_GEMM", "False").lower() in ("true", "1") + ), + # Whether to use aiter rope. + # By default is disabled. + "VLLM_ROCM_USE_AITER_TRITON_ROPE": lambda: ( + os.getenv("VLLM_ROCM_USE_AITER_TRITON_ROPE", "False").lower() in ("true", "1") + ), + # Whether to use aiter triton fp8 bmm kernel + # By default is enabled. + "VLLM_ROCM_USE_AITER_FP8BMM": lambda: ( + os.getenv("VLLM_ROCM_USE_AITER_FP8BMM", "True").lower() in ("true", "1") + ), + # Whether to use aiter triton fp4 bmm kernel + # By default is enabled. + "VLLM_ROCM_USE_AITER_FP4BMM": lambda: ( + os.getenv("VLLM_ROCM_USE_AITER_FP4BMM", "True").lower() in ("true", "1") + ), + # Use AITER triton unified attention for V1 attention + "VLLM_ROCM_USE_AITER_UNIFIED_ATTENTION": lambda: ( + os.getenv("VLLM_ROCM_USE_AITER_UNIFIED_ATTENTION", "False").lower() + in ("true", "1") + ), + # Whether to use aiter fusion shared experts ops. + # By default is disabled. + "VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS": lambda: ( + os.getenv("VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS", "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: ( + os.getenv("VLLM_ROCM_USE_AITER_TRITON_GEMM", "True").lower() in ("true", "1") + ), + # use rocm skinny gemms + "VLLM_ROCM_USE_SKINNY_GEMM": lambda: ( + os.getenv("VLLM_ROCM_USE_SKINNY_GEMM", "True").lower() in ("true", "1") + ), + # Pad the fp8 weights to 256 bytes for ROCm + "VLLM_ROCM_FP8_PADDING": lambda: bool(int(os.getenv("VLLM_ROCM_FP8_PADDING", "1"))), + # Pad the weights for the moe kernel + "VLLM_ROCM_MOE_PADDING": lambda: bool(int(os.getenv("VLLM_ROCM_MOE_PADDING", "1"))), + # Whether to use the shuffled kv cache layout + "VLLM_ROCM_SHUFFLE_KV_CACHE_LAYOUT": lambda: ( + os.getenv("VLLM_ROCM_SHUFFLE_KV_CACHE_LAYOUT", "False").lower() in ("true", "1") + ), + # Custom quick allreduce kernel for MI3* cards + # Choice of quantization level: FP, INT8, INT6, INT4 or NONE + # Recommended for large models to get allreduce + "VLLM_ROCM_QUICK_REDUCE_QUANTIZATION": env_with_choices( + "VLLM_ROCM_QUICK_REDUCE_QUANTIZATION", + "NONE", + ["FP", "INT8", "INT6", "INT4", "NONE"], + ), + # Custom quick allreduce kernel for MI3* cards + # Due to the lack of the bfloat16 asm instruction, bfloat16 + # kernels are slower than fp16, + # If environment variable is set to 1, the input is converted to fp16 + "VLLM_ROCM_QUICK_REDUCE_CAST_BF16_TO_FP16": lambda: ( + os.getenv("VLLM_ROCM_QUICK_REDUCE_CAST_BF16_TO_FP16", "True").lower() + in ("true", "1") + ), + # Custom quick allreduce kernel for MI3* cards. + # Controls the maximum allowed number of data bytes(MB) for custom quick + # allreduce communication. + # Default: 2048 MB. + # Data exceeding this size will use either custom allreduce or RCCL + # communication. + "VLLM_ROCM_QUICK_REDUCE_MAX_SIZE_BYTES_MB": lambda: maybe_convert_int( + os.environ.get("VLLM_ROCM_QUICK_REDUCE_MAX_SIZE_BYTES_MB", None) + ), + # Custom quick allreduce kernel for MI3* cards. + # Controls the minimum allowed number of data bytes(MB) required to use + # custom quick allreduce communication. + # If unset, use the built-in threshold table. + "VLLM_ROCM_QUICK_REDUCE_MIN_SIZE_BYTES_MB": lambda: maybe_convert_int( + os.environ.get("VLLM_ROCM_QUICK_REDUCE_MIN_SIZE_BYTES_MB", None) + ), + # Controls the minimum tensor size (KB, where 1 KB = 1024 bytes) required + # to use the configured QuickReduce codec. Smaller tensors use FP + # QuickReduce. This does not affect QuickReduce eligibility. + "VLLM_ROCM_QUICK_REDUCE_QUANTIZATION_MIN_SIZE_KB": lambda: maybe_convert_int( + os.environ.get("VLLM_ROCM_QUICK_REDUCE_QUANTIZATION_MIN_SIZE_KB", None) + ), + # Divisor for dynamic query scale factor calculation for FP8 KV Cache + "Q_SCALE_CONSTANT": lambda: int(os.getenv("Q_SCALE_CONSTANT", "200")), + # Divisor for dynamic key scale factor calculation for FP8 KV Cache + "K_SCALE_CONSTANT": lambda: int(os.getenv("K_SCALE_CONSTANT", "200")), + # Divisor for dynamic value scale factor calculation for FP8 KV Cache + "V_SCALE_CONSTANT": lambda: int(os.getenv("V_SCALE_CONSTANT", "100")), + # If set, enable multiprocessing in LLM for the V1 code path. + "VLLM_ENABLE_V1_MULTIPROCESSING": lambda: bool( + int(os.getenv("VLLM_ENABLE_V1_MULTIPROCESSING", "1")) + ), + "VLLM_LOG_BATCHSIZE_INTERVAL": lambda: float( + os.getenv("VLLM_LOG_BATCHSIZE_INTERVAL", "-1") + ), + "VLLM_DISABLE_COMPILE_CACHE": disable_compile_cache, + # If set to "0", disable LayerName opaque type for layer_name + # parameters in custom ops. Defaults to enabled on torch >= 2.11. + "VLLM_USE_LAYERNAME": lambda: bool(int(os.getenv("VLLM_USE_LAYERNAME", "1"))), + # If set, use the Rust frontend binary instead of the Python API server + # process(es). + "VLLM_USE_RUST_FRONTEND": lambda: bool( + int(os.getenv("VLLM_USE_RUST_FRONTEND", "0")) + ), + # Path to the Rust frontend binary. Defaults to "auto" which discovers + # the binary installed with the vllm package. Only used when + # VLLM_USE_RUST_FRONTEND=1. + "VLLM_RUST_FRONTEND_PATH": lambda: _resolve_rust_frontend_path(), + # If set, vllm will run in development mode, which will enable + # some additional endpoints for developing and debugging, + # e.g. `/reset_prefix_cache` + "VLLM_SERVER_DEV_MODE": lambda: bool(int(os.getenv("VLLM_SERVER_DEV_MODE", "0"))), + # Controls the maximum number of requests to handle in a + # single asyncio task when processing per-token outputs in the + # V1 AsyncLLM interface. It is applicable when handling a high + # concurrency of streaming requests. + # Setting this too high can result in a higher variance of + # inter-message latencies. Setting it too low can negatively impact + # TTFT and overall throughput. + "VLLM_V1_OUTPUT_PROC_CHUNK_SIZE": lambda: int( + os.getenv("VLLM_V1_OUTPUT_PROC_CHUNK_SIZE", "128") + ), + # If set, vLLM will disable the MLA attention optimizations. + "VLLM_MLA_DISABLE": lambda: bool(int(os.getenv("VLLM_MLA_DISABLE", "0"))), + # If set, vLLM will pick up the provided Flash Attention MLA + # Number of GPUs per worker in Ray, if it is set to be a fraction, + # it allows ray to schedule multiple actors on a single GPU, + # so that users can colocate other actors on the same GPUs as vLLM. + "VLLM_RAY_PER_WORKER_GPUS": lambda: float( + os.getenv("VLLM_RAY_PER_WORKER_GPUS", "1.0") + ), + # Bundle indices for Ray, if it is set, it can control precisely + # 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", ""), + # 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), + # Rank of the process in the data parallel setting + "VLLM_DP_RANK": lambda: int(os.getenv("VLLM_DP_RANK", "0")), + # Rank of the process in the data parallel setting. + # Defaults to VLLM_DP_RANK when not set. + "VLLM_DP_RANK_LOCAL": lambda: int( + os.getenv("VLLM_DP_RANK_LOCAL", sys.modules[__name__].VLLM_DP_RANK) + ), + # World size of the data parallel setting + "VLLM_DP_SIZE": lambda: int(os.getenv("VLLM_DP_SIZE", "1")), + # IP address of the master node in the data parallel setting + "VLLM_DP_MASTER_IP": lambda: os.getenv("VLLM_DP_MASTER_IP", "127.0.0.1"), + # Port of the master node in the data parallel setting + "VLLM_DP_MASTER_PORT": lambda: int(os.getenv("VLLM_DP_MASTER_PORT", "0")), + # Randomize inputs during dummy runs when using Data Parallel + "VLLM_RANDOMIZE_DP_DUMMY_INPUTS": lambda: ( + os.environ.get("VLLM_RANDOMIZE_DP_DUMMY_INPUTS", "0") == "1" + ), + # Strategy to pack the data parallel ranks for Ray. + # Available options: + # - "fill": + # for DP master node, allocate exactly data-parallel-size-local DP ranks, + # for non-master nodes, allocate as many DP ranks as can fit; + # - "strict": + # allocate exactly data-parallel-size-local DP ranks to each picked node; + # - "span": + # Should be used only when a single DP rank requires multiple nodes. + # allocate one DP rank over as many nodes as required for set world_size; + # This environment variable is ignored if data-parallel-backend is not Ray. + "VLLM_RAY_DP_PACK_STRATEGY": lambda: os.getenv( + "VLLM_RAY_DP_PACK_STRATEGY", "strict" + ), + # Comma-separated *additional* prefixes of env vars to copy from the + # driver to Ray workers. These are merged with the built-in defaults + # defined in ``vllm.ray.ray_env`` (VLLM_, etc.). Example: "MYLIB_,OTHER_" + "VLLM_RAY_EXTRA_ENV_VAR_PREFIXES_TO_COPY": lambda: os.getenv( + "VLLM_RAY_EXTRA_ENV_VAR_PREFIXES_TO_COPY", "" + ), + # Comma-separated *additional* individual env var names to copy from + # the driver to Ray workers. Merged with the built-in defaults + # defined in ``vllm.ray.ray_env`` (PYTHONHASHSEED). + # Example: "MY_SECRET,MY_FLAG" + "VLLM_RAY_EXTRA_ENV_VARS_TO_COPY": lambda: os.getenv( + "VLLM_RAY_EXTRA_ENV_VARS_TO_COPY", "" + ), + # Whether to use S3 path for model loading in CI via RunAI Streamer + "VLLM_CI_USE_S3": lambda: os.environ.get("VLLM_CI_USE_S3", "0") == "1", + # Use model_redirect to redirect the model name to a local folder. + # `model_redirect` can be a json file mapping the model between + # repo_id and local folder: + # {"meta-llama/Llama-3.2-1B": "/tmp/Llama-3.2-1B"} + # or a space separated values table file: + # meta-llama/Llama-3.2-1B /tmp/Llama-3.2-1B + "VLLM_MODEL_REDIRECT_PATH": lambda: os.environ.get( + "VLLM_MODEL_REDIRECT_PATH", None + ), + # Whether to use atomicAdd reduce in gptq/awq marlin kernel. + "VLLM_MARLIN_USE_ATOMIC_ADD": lambda: ( + os.environ.get("VLLM_MARLIN_USE_ATOMIC_ADD", "0") == "1" + ), + # Whether to use marlin kernel in mxfp4 quantization method + # Deprecated: use --moe-backend marlin (MoE) or --linear-backend marlin + # (linear) instead. + "VLLM_MXFP4_USE_MARLIN": deprecated_env( + "VLLM_MXFP4_USE_MARLIN", + "v0.23", + "Use --moe-backend marlin or --linear-backend marlin.", + lambda: maybe_convert_bool(os.environ.get("VLLM_MXFP4_USE_MARLIN", None)), + ), + # The activation dtype for marlin kernel + "VLLM_MARLIN_INPUT_DTYPE": env_with_choices( + "VLLM_MARLIN_INPUT_DTYPE", None, ["int8", "fp8"] + ), + # The online quantization dtype for humming kernel + "VLLM_HUMMING_ONLINE_QUANT_CONFIG": lambda: maybe_convert_json_str_or_file( + os.environ.get("VLLM_HUMMING_ONLINE_QUANT_CONFIG", None) + ), + # The activation dtype config for humming kernel + "VLLM_HUMMING_INPUT_QUANT_CONFIG": lambda: maybe_convert_json_str_or_file( + os.environ.get("VLLM_HUMMING_INPUT_QUANT_CONFIG", None) + ), + # Whether to use fp16 accumulator mma + "VLLM_HUMMING_USE_F16_ACCUM": lambda: maybe_convert_bool( + os.environ.get("VLLM_HUMMING_USE_F16_ACCUM", "0") + ), + # Whether to use indexed gemm for humming moe + # if 1, force use indexed gemm + # if 0, force use grouped gemm + # if None, choose better gemm type automatically + "VLLM_HUMMING_MOE_GEMM_TYPE": lambda: os.environ.get( + "VLLM_HUMMING_MOE_GEMM_TYPE", None + ), + # Whether to use DeepEPLL kernels for NVFP4 quantization and dispatch method + # only supported on Blackwell GPUs and with + # https://github.com/deepseek-ai/DeepEP/pull/341 + "VLLM_DEEPEPLL_NVFP4_DISPATCH": lambda: bool( + int(os.getenv("VLLM_DEEPEPLL_NVFP4_DISPATCH", "0")) + ), + # Whether to turn on the outlines cache for V1 + # This cache is unbounded and on disk, so it's not safe to use in + # an environment with potentially malicious users. + "VLLM_V1_USE_OUTLINES_CACHE": lambda: ( + os.environ.get("VLLM_V1_USE_OUTLINES_CACHE", "0") == "1" + ), + # Gap between padding buckets for the forward pass. So we have + # 8, we will run forward pass with [16, 24, 32, ...]. + "VLLM_TPU_BUCKET_PADDING_GAP": lambda: ( + int(os.environ["VLLM_TPU_BUCKET_PADDING_GAP"]) + if "VLLM_TPU_BUCKET_PADDING_GAP" in os.environ + else 0 + ), + "VLLM_TPU_MOST_MODEL_LEN": lambda: maybe_convert_int( + os.environ.get("VLLM_TPU_MOST_MODEL_LEN", None) + ), + # Whether using Pathways + "VLLM_TPU_USING_PATHWAYS": lambda: bool( + "proxy" in os.getenv("JAX_PLATFORMS", "").lower() + ), + # Pre-JIT DSv4 sparse-MLA input-prep kernels to close a cold-JIT IMA + # window on SM12x. + "VLLM_ENABLE_DEEPSEEK_V4_SPARSE_MLA_WARMUP": lambda: bool( + int(os.getenv("VLLM_ENABLE_DEEPSEEK_V4_SPARSE_MLA_WARMUP", "1")) + ), + # Allow use of DeepGemm kernels for fused moe ops. + "VLLM_USE_DEEP_GEMM": lambda: bool(int(os.getenv("VLLM_USE_DEEP_GEMM", "1"))), + # Allow use of DeepGemm specifically for MoE fused ops (overrides only MoE). + "VLLM_MOE_USE_DEEP_GEMM": lambda: bool( + int(os.getenv("VLLM_MOE_USE_DEEP_GEMM", "1")) + ), + # Whether to use E8M0 scaling when DeepGEMM is used on Blackwell GPUs. + "VLLM_USE_DEEP_GEMM_E8M0": lambda: bool( + int(os.getenv("VLLM_USE_DEEP_GEMM_E8M0", "1")) + ), + # Whether to create TMA-aligned scale tensor when DeepGEMM is used. + "VLLM_USE_DEEP_GEMM_TMA_ALIGNED_SCALES": lambda: bool( + int(os.getenv("VLLM_USE_DEEP_GEMM_TMA_ALIGNED_SCALES", "1")) + ), + # DeepGemm JITs the kernels on-demand. The warmup attempts to make DeepGemm + # JIT all the required kernels before model execution so there is no + # JIT'ing in the hot-path. However, this warmup increases the engine + # startup time by a couple of minutes. + # Available options: + # - "skip" : Skip warmup. + # - "full" : Warmup deepgemm by running all possible gemm shapes the + # engine could encounter. + # - "relax" : Select gemm shapes to run based on some heuristics. The + # heuristic aims to have the same effect as running all possible gemm + # shapes, but provides no guarantees. + "VLLM_DEEP_GEMM_WARMUP": env_with_choices( + "VLLM_DEEP_GEMM_WARMUP", + "relax", + [ + "skip", + "full", + "relax", + ], + ), + # Experimental DSpark confidence-scheduled verification threshold. + # Parsed and range-checked by DSparkProposer. + "VLLM_DSPARK_CONFIDENCE_THRESHOLD": lambda: os.getenv( + "VLLM_DSPARK_CONFIDENCE_THRESHOLD", "0.0" + ), + "VLLM_DSPARK_CONFIDENCE_SCHEDULER": lambda: os.getenv( + "VLLM_DSPARK_CONFIDENCE_SCHEDULER", "auto" + ), + "VLLM_DSPARK_SPS_CURVE": lambda: os.getenv("VLLM_DSPARK_SPS_CURVE", ""), + "VLLM_DSPARK_HARDWARE_SCHEDULER_EARLY_STOP": lambda: os.getenv( + "VLLM_DSPARK_HARDWARE_SCHEDULER_EARLY_STOP", "1" + ), + "VLLM_DSPARK_FORCE_DRAFT_LENGTH": lambda: os.getenv( + "VLLM_DSPARK_FORCE_DRAFT_LENGTH", "" + ), + "VLLM_DSPARK_EXPORT_DRAFT_PROBS": lambda: ( + os.getenv("VLLM_DSPARK_EXPORT_DRAFT_PROBS", "0").strip().lower() + in ("1", "true", "yes", "on") + ), + "VLLM_DSPARK_COLLECT_CONFIDENCE_DIAGNOSTICS": lambda: ( + os.getenv("VLLM_DSPARK_COLLECT_CONFIDENCE_DIAGNOSTICS", "0").strip().lower() + in ("1", "true", "yes", "on") + ), + "VLLM_DSPARK_POSITION0_DIAGNOSTICS": lambda: ( + os.getenv("VLLM_DSPARK_POSITION0_DIAGNOSTICS", "0").strip().lower() + in ("1", "true", "yes", "on") + ), + "VLLM_DSPARK_LOCAL_ARGMAX": lambda: ( + os.getenv("VLLM_DSPARK_LOCAL_ARGMAX", "0").strip().lower() + in ("1", "true", "yes", "on") + ), + "VLLM_DSPARK_FUSED_MARKOV_ARGMAX": lambda: ( + os.getenv("VLLM_DSPARK_FUSED_MARKOV_ARGMAX", "0").strip().lower() + in ("1", "true", "yes", "on") + ), + "VLLM_DSPARK_GPU_REJECTED_CONTEXT_MASK": lambda: ( + os.getenv("VLLM_DSPARK_GPU_REJECTED_CONTEXT_MASK", "0").strip().lower() + in ("1", "true", "yes", "on") + ), + "VLLM_DSPARK_REFERENCE_KV_QUANT_DEQUANT": lambda: ( + os.getenv("VLLM_DSPARK_REFERENCE_KV_QUANT_DEQUANT", "0").strip().lower() + in ("1", "true", "yes", "on") + ), + "VLLM_DSPARK_REPLICATE_MARKOV_W1": lambda: ( + os.getenv("VLLM_DSPARK_REPLICATE_MARKOV_W1", "0").strip().lower() + in ("1", "true", "yes", "on") + ), + "VLLM_DSPARK_STAGE_TIMING": lambda: ( + os.getenv("VLLM_DSPARK_STAGE_TIMING", "0").strip().lower() + in ("1", "true", "yes", "on") + ), + "VLLM_DSPARK_STAGE_TIMING_LOG_EVERY": lambda: int( + os.getenv("VLLM_DSPARK_STAGE_TIMING_LOG_EVERY", "20") + ), + "VLLM_DSPARK_ITER_TIMING": lambda: ( + os.getenv("VLLM_DSPARK_ITER_TIMING", "0").strip().lower() + in ("1", "true", "yes", "on") + ), + "VLLM_DSPARK_ITER_TIMING_LOG_EVERY": lambda: int( + os.getenv("VLLM_DSPARK_ITER_TIMING_LOG_EVERY", "20") + ), + "VLLM_DSPARK_TARGET_TIMING": lambda: ( + os.getenv("VLLM_DSPARK_TARGET_TIMING", "0").strip().lower() + in ("1", "true", "yes", "on") + ), + "VLLM_DSPARK_TARGET_TIMING_LOG_EVERY": lambda: int( + os.getenv("VLLM_DSPARK_TARGET_TIMING_LOG_EVERY", "20") + ), + "VLLM_DSV4_B12X_COMPRESSED_MLA": lambda: ( + os.getenv("VLLM_DSV4_B12X_COMPRESSED_MLA", "0").strip().lower() + in ("1", "true", "yes", "on") + ), + "VLLM_DSV4_DSPARK_DEFER_TARGET_CAPTURE": lambda: ( + os.getenv("VLLM_DSV4_DSPARK_DEFER_TARGET_CAPTURE", "0").strip().lower() + in ("1", "true", "yes", "on") + ), + "VLLM_DSV4_DSPARK_DEFER_TARGET_CAPTURE_EXACT": lambda: ( + os.getenv("VLLM_DSV4_DSPARK_DEFER_TARGET_CAPTURE_EXACT", "0").strip().lower() + in ("1", "true", "yes", "on") + ), + # Whether to use fused grouped_topk used for MoE expert selection. + "VLLM_USE_FUSED_MOE_GROUPED_TOPK": lambda: bool( + int(os.getenv("VLLM_USE_FUSED_MOE_GROUPED_TOPK", "1")) + ), + # Allow use of FlashInfer FP8 block-scale GEMM for linear layers. + # This uses TensorRT-LLM kernels and requires SM90+ (Hopper). + "VLLM_BLOCKSCALE_FP8_GEMM_FLASHINFER": lambda: bool( + int(os.getenv("VLLM_BLOCKSCALE_FP8_GEMM_FLASHINFER", "1")) + ), + # Allow use of FlashInfer BF16 MoE kernels for fused moe ops. + # Deprecated: use --moe-backend to select a kernel explicitly. + "VLLM_USE_FLASHINFER_MOE_FP16": deprecated_env( + "VLLM_USE_FLASHINFER_MOE_FP16", + "v0.23", + "Use --moe-backend (e.g. flashinfer_trtllm, flashinfer_cutlass).", + lambda: bool(int(os.getenv("VLLM_USE_FLASHINFER_MOE_FP16", "0"))), + ), + # Allow use of FlashInfer FP8 MoE kernels for fused moe ops. + # Deprecated: use --moe-backend to select a kernel explicitly. + "VLLM_USE_FLASHINFER_MOE_FP8": deprecated_env( + "VLLM_USE_FLASHINFER_MOE_FP8", + "v0.23", + "Use --moe-backend (e.g. flashinfer_trtllm, flashinfer_cutlass).", + lambda: bool(int(os.getenv("VLLM_USE_FLASHINFER_MOE_FP8", "0"))), + ), + # Allow use of FlashInfer NVFP4 MoE kernels for fused moe ops. + # Deprecated: use --moe-backend to select a kernel explicitly. + "VLLM_USE_FLASHINFER_MOE_FP4": deprecated_env( + "VLLM_USE_FLASHINFER_MOE_FP4", + "v0.23", + "Use --moe-backend (e.g. flashinfer_trtllm, flashinfer_cutlass, " + "flashinfer_cutedsl).", + lambda: bool(int(os.getenv("VLLM_USE_FLASHINFER_MOE_FP4", "0"))), + ), + # Allow use of FlashInfer MxInt4 MoE kernels for fused moe ops. + "VLLM_USE_FLASHINFER_MOE_INT4": lambda: bool( + int(os.getenv("VLLM_USE_FLASHINFER_MOE_INT4", "0")) + ), + # If set to 1, use the FlashInfer + # MXFP8 (activation) x MXFP4 (weight) MoE backend. + # Deprecated: use --moe-backend flashinfer_trtllm combined with + # --quantization_config.moe.activation mxfp8. + "VLLM_USE_FLASHINFER_MOE_MXFP4_MXFP8": deprecated_env( + "VLLM_USE_FLASHINFER_MOE_MXFP4_MXFP8", + "v0.23", + "Use --moe-backend flashinfer_trtllm with " + "--quantization_config.moe.activation mxfp8.", + lambda: bool(int(os.getenv("VLLM_USE_FLASHINFER_MOE_MXFP4_MXFP8", "0"))), + ), + # If set to 1, use the FlashInfer CUTLASS backend for + # MXFP8 (activation) x MXFP4 (weight) MoE. + # Deprecated: use --moe-backend flashinfer_cutlass combined with + # --quantization_config.moe.activation mxfp8. + "VLLM_USE_FLASHINFER_MOE_MXFP4_MXFP8_CUTLASS": deprecated_env( + "VLLM_USE_FLASHINFER_MOE_MXFP4_MXFP8_CUTLASS", + "v0.23", + "Use --moe-backend flashinfer_cutlass with " + "--quantization_config.moe.activation mxfp8.", + lambda: bool( + int(os.getenv("VLLM_USE_FLASHINFER_MOE_MXFP4_MXFP8_CUTLASS", "0")) + ), + ), + # If set to 1, use the FlashInfer + # BF16 (activation) x MXFP4 (weight) MoE backend. + # Deprecated: use --moe-backend to select a kernel explicitly. + "VLLM_USE_FLASHINFER_MOE_MXFP4_BF16": deprecated_env( + "VLLM_USE_FLASHINFER_MOE_MXFP4_BF16", + "v0.23", + "Use --moe-backend (e.g. flashinfer_trtllm, flashinfer_cutlass).", + lambda: bool(int(os.getenv("VLLM_USE_FLASHINFER_MOE_MXFP4_BF16", "0"))), + ), + # Control the cache sized used by the xgrammar compiler. The default + # of 512 MB should be enough for roughly 1000 JSON schemas. + # It can be changed with this variable if needed for some reason. + "VLLM_XGRAMMAR_CACHE_MB": lambda: int(os.getenv("VLLM_XGRAMMAR_CACHE_MB", "512")), + # Control the threshold for msgspec to use 'zero copy' for + # serialization/deserialization of tensors. Tensors below + # this limit will be encoded into the msgpack buffer, and + # tensors above will instead be sent via a separate message. + # While the sending side still actually copies the tensor + # in all cases, on the receiving side, tensors above this + # limit will actually be zero-copy decoded. + "VLLM_MSGPACK_ZERO_COPY_THRESHOLD": lambda: int( + os.getenv("VLLM_MSGPACK_ZERO_COPY_THRESHOLD", "256") + ), + # If set, allow insecure serialization using pickle. + # This is useful for environments where it is deemed safe to use the + # insecure method and it is needed for some reason. + "VLLM_ALLOW_INSECURE_SERIALIZATION": lambda: bool( + int(os.getenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "0")) + ), + # Temporary: skip adding random suffix to internal request IDs. May be + # needed for KV connectors that match request IDs across instances. + "VLLM_DISABLE_REQUEST_ID_RANDOMIZATION": lambda: bool( + int(os.getenv("VLLM_DISABLE_REQUEST_ID_RANDOMIZATION", "0")) + ), + # IP address used for NIXL handshake between remote agents. + "VLLM_NIXL_SIDE_CHANNEL_HOST": lambda: os.getenv( + "VLLM_NIXL_SIDE_CHANNEL_HOST", "localhost" + ), + # Port used for NIXL handshake between remote agents. + "VLLM_NIXL_SIDE_CHANNEL_PORT": lambda: int( + os.getenv("VLLM_NIXL_SIDE_CHANNEL_PORT", "5600") + ), + # Port used for Mooncake handshake between remote agents. + "VLLM_MOONCAKE_BOOTSTRAP_PORT": lambda: int( + os.getenv("VLLM_MOONCAKE_BOOTSTRAP_PORT", "8998") + ), + # Log per-batch memory/disk tier breakdown on external GETs. + "VLLM_MOONCAKE_STORE_TIER_LOG": lambda: ( + os.getenv("VLLM_MOONCAKE_STORE_TIER_LOG", "False").lower() in ("true", "1") + ), + # Fraction of the owner's DirectIO staging buffer to fill per GET batch. + "VLLM_MOONCAKE_DISK_STAGING_USABLE_RATIO": lambda: float( + os.getenv("VLLM_MOONCAKE_DISK_STAGING_USABLE_RATIO", "0.9") + ), + # Pin this rank to a specific owner segment ("host:port"). + "MOONCAKE_PREFERRED_SEGMENT": lambda: os.getenv("MOONCAKE_PREFERRED_SEGMENT"), + # Override the hostname the rank registers as a Mooncake requester. + "MOONCAKE_REQUESTER_LOCAL_HOSTNAME": lambda: os.getenv( + "MOONCAKE_REQUESTER_LOCAL_HOSTNAME" + ), + # Flashinfer MoE backend for vLLM's fused Mixture-of-Experts support. + # Both require compute capability 10.0 or above. + # Available options: + # - "throughput": [default] + # Uses CUTLASS kernels optimized for high-throughput batch inference. + # - "latency": + # Uses TensorRT-LLM kernels optimized for low-latency inference. + # Deprecated: pass --moe-backend flashinfer_{trtllm,cutlass,cutedsl} directly. + "VLLM_FLASHINFER_MOE_BACKEND": deprecated_env( + "VLLM_FLASHINFER_MOE_BACKEND", + "v0.23", + "Use --moe-backend flashinfer_trtllm, flashinfer_cutlass, or " + "flashinfer_cutedsl.", + env_with_choices( + "VLLM_FLASHINFER_MOE_BACKEND", + "latency", + ["throughput", "latency", "masked_gemm"], + ), + ), + # Override the directory for the FlashInfer autotune config cache. + "VLLM_FLASHINFER_AUTOTUNE_CACHE_DIR": lambda: os.getenv( + "VLLM_FLASHINFER_AUTOTUNE_CACHE_DIR", None + ), + # Flashinfer fused allreduce backend. + "VLLM_FLASHINFER_ALLREDUCE_BACKEND": env_with_choices( + "VLLM_FLASHINFER_ALLREDUCE_BACKEND", + "auto", + ["auto", "trtllm", "mnnvl"], + ), + # Opt in to the b12x PCIe oneshot custom allreduce path on PCIe-only GPUs. + "VLLM_ENABLE_PCIE_ALLREDUCE": lambda: bool( + int(os.getenv("VLLM_ENABLE_PCIE_ALLREDUCE", "0")) + ), + # Control the workspace buffer size for the FlashInfer backend. + "VLLM_FLASHINFER_WORKSPACE_BUFFER_SIZE": lambda: int( + os.getenv("VLLM_FLASHINFER_WORKSPACE_BUFFER_SIZE", str(394 * 1024 * 1024)) + ), + # Control the maximum number of tokens per expert supported by the + # NVFP4 MoE CUTLASS Kernel. This value is used to create a buffer for + # the blockscale tensor of activations NVFP4 Quantization. + # This is used to prevent the kernel from running out of memory. + "VLLM_MAX_TOKENS_PER_EXPERT_FP4_MOE": lambda: int( + os.getenv("VLLM_MAX_TOKENS_PER_EXPERT_FP4_MOE", "163840") + ), + # Specifies the thresholds of the communicated tensor sizes under which + # vllm should use flashinfer fused allreduce. The variable should be a + # JSON with the following format: + # { : } + # Unspecified world sizes will fall back to + # { 2: 64, 4: 1, : 0.5 } + "VLLM_FLASHINFER_ALLREDUCE_FUSION_THRESHOLDS_MB": lambda: json.loads( + os.getenv("VLLM_FLASHINFER_ALLREDUCE_FUSION_THRESHOLDS_MB", "{}") + ), + # MoE routing strategy selector. + # See `RoutingSimulator.get_available_strategies()` # for available + # strategies. + # Custom routing strategies can be registered by + # RoutingSimulator.register_strategy() + # Note: custom strategies may not produce correct model outputs + "VLLM_MOE_ROUTING_SIMULATION_STRATEGY": lambda: os.environ.get( + "VLLM_MOE_ROUTING_SIMULATION_STRATEGY", "" + ).lower(), + # Regex timeout for use by the vLLM tool parsing plugins. + "VLLM_TOOL_PARSE_REGEX_TIMEOUT_SECONDS": lambda: int( + os.getenv("VLLM_TOOL_PARSE_REGEX_TIMEOUT_SECONDS", "1") + ), + # Control the max chunk bytes (in MB) for the rpc message queue. + # Object larger than this threshold will be broadcast to worker + # processes via zmq. + "VLLM_MQ_MAX_CHUNK_BYTES_MB": lambda: int( + os.getenv("VLLM_MQ_MAX_CHUNK_BYTES_MB", "16") + ), + # Timeout in seconds for execute_model RPC calls in multiprocessing + # executor (only applies when TP > 1). + "VLLM_EXECUTE_MODEL_TIMEOUT_SECONDS": lambda: int( + os.getenv("VLLM_EXECUTE_MODEL_TIMEOUT_SECONDS", "300") + ), + # KV Cache layout used throughout vllm. + # Some common values are: + # - NHD + # - HND + # Where N=num_blocks, H=num_heads and D=head_size. The default value will + # leave the layout choice to the backend. Mind that backends may only + # implement and support a subset of all possible layouts. + "VLLM_KV_CACHE_LAYOUT": env_with_choices( + "VLLM_KV_CACHE_LAYOUT", None, ["NHD", "HND"] + ), + # SSM conv state layout used for Mamba models. + # - SD: (state_len, dim) — dim contiguous (default) + # - DS: (dim, state_len) — TP-sharded dim on dim1, + # consistent with SSM temporal state and HND KV cache layout. + "VLLM_SSM_CONV_STATE_LAYOUT": env_with_choices( + "VLLM_SSM_CONV_STATE_LAYOUT", None, ["SD", "DS"] + ), + # Enable checking whether the generated logits contain NaNs, + # indicating corrupted output. Useful for debugging low level bugs + # or bad hardware but it may add compute overhead. + "VLLM_COMPUTE_NANS_IN_LOGITS": lambda: bool( + int(os.getenv("VLLM_COMPUTE_NANS_IN_LOGITS", "0")) + ), + # Controls whether or not emulations are used for NVFP4 + # generations on machines < 100 for compressed-tensors + # models + # Deprecated: use --linear-backend emulation instead. + "VLLM_USE_NVFP4_CT_EMULATIONS": deprecated_env( + "VLLM_USE_NVFP4_CT_EMULATIONS", + "v0.23", + "Use --linear-backend emulation.", + lambda: bool(int(os.getenv("VLLM_USE_NVFP4_CT_EMULATIONS", "0"))), + ), + # Timeout (in seconds) for MooncakeConnector in PD disaggregated setup. + "VLLM_MOONCAKE_ABORT_REQUEST_TIMEOUT": lambda: int( + os.getenv("VLLM_MOONCAKE_ABORT_REQUEST_TIMEOUT", "480") + ), + # If set, it means we pre-downloaded cubin files and flashinfer will + # read the cubin files directly. + "VLLM_HAS_FLASHINFER_CUBIN": lambda: bool( + int(os.getenv("VLLM_HAS_FLASHINFER_CUBIN", "0")) + ), + # Supported options: + # - "flashinfer-cudnn": use flashinfer cudnn GEMM backend + # - "flashinfer-trtllm": use flashinfer trtllm GEMM backend + # - "flashinfer-cutlass": use flashinfer cutlass GEMM backend + # - "marlin": use marlin GEMM backend (for GPUs without native FP4 support) + # - "emulation": + # use BF16/FP16 GEMM, dequantizing weights and running QDQ on activations. + # This is only meant for research purposes to run on devices where NVFP4 + # GEMM kernels are not available. + # - : automatically pick an available backend + # Deprecated: use --linear-backend instead. + "VLLM_NVFP4_GEMM_BACKEND": deprecated_env( + "VLLM_NVFP4_GEMM_BACKEND", + "v0.23", + "Use --linear-backend.", + env_with_choices( + "VLLM_NVFP4_GEMM_BACKEND", + None, + [ + "flashinfer-b12x", + "flashinfer-cudnn", + "flashinfer-trtllm", + "flashinfer-cutlass", + "cutlass", + "marlin", + "emulation", + ], + ), + ), + # Controls garbage collection during CUDA graph capture. + # If set to 0 (default), enables GC freezing to speed up capture time. + # If set to 1, allows GC to run during capture. + "VLLM_ENABLE_CUDAGRAPH_GC": lambda: bool( + int(os.getenv("VLLM_ENABLE_CUDAGRAPH_GC", "0")) + ), + # Used to force set up loopback IP + "VLLM_LOOPBACK_IP": lambda: os.getenv("VLLM_LOOPBACK_IP", ""), + # Used to set the process name prefix for vLLM processes. + # This is useful for debugging and monitoring purposes. + # The default value is "VLLM". + "VLLM_PROCESS_NAME_PREFIX": lambda: os.getenv("VLLM_PROCESS_NAME_PREFIX", "VLLM"), + # Allow chunked local attention with hybrid kv cache manager. + # Currently using the Hybrid KV cache manager with chunked local attention + # in the Llama4 models (the only models currently using chunked local attn) + # causes a latency regression. For this reason, we disable it by default. + # This flag is used to allow users to enable it if they want to (to save on + # kv-cache memory usage and enable longer contexts) + # TODO(lucas): Remove this flag once latency regression is resolved. + "VLLM_ALLOW_CHUNKED_LOCAL_ATTN_WITH_HYBRID_KV_CACHE": lambda: bool( + int(os.getenv("VLLM_ALLOW_CHUNKED_LOCAL_ATTN_WITH_HYBRID_KV_CACHE", "1")) + ), + # Enables support for the "store" option in the OpenAI Responses API. + # When set to 1, vLLM's OpenAI server will retain the input and output + # messages for those requests in memory. By default, this is disabled (0), + # and the "store" option is ignored. + # NOTE/WARNING: + # 1. Messages are kept in memory only (not persisted to disk) and will be + # lost when the vLLM server shuts down. + # 2. Enabling this option will cause a memory leak, as stored messages are + # never removed from memory until the server terminates. + "VLLM_ENABLE_RESPONSES_API_STORE": lambda: bool( + int(os.getenv("VLLM_ENABLE_RESPONSES_API_STORE", "0")) + ), + # If set, use the fp8 mfma in rocm paged attention. + "VLLM_ROCM_FP8_MFMA_PAGE_ATTN": lambda: bool( + int(os.getenv("VLLM_ROCM_FP8_MFMA_PAGE_ATTN", "0")) + ), + # Whether to use pytorch symmetric memory for allreduce + "VLLM_ALLREDUCE_USE_SYMM_MEM": lambda: bool( + int(os.getenv("VLLM_ALLREDUCE_USE_SYMM_MEM", "1")) + ), + # Whether to use FlashInfer allreduce + "VLLM_ALLREDUCE_USE_FLASHINFER": lambda: bool( + int(os.getenv("VLLM_ALLREDUCE_USE_FLASHINFER", "0")) + ), + # Experimental: use this to enable MCP tool calling for non harmony models + "VLLM_USE_EXPERIMENTAL_PARSER_CONTEXT": lambda: bool( + int(os.getenv("VLLM_USE_EXPERIMENTAL_PARSER_CONTEXT", "0")) + ), + # User override folder for tuned Triton-kernel configs. Shared by MoE, + # Mamba SSU, and LoRA. Filenames are distinct so one folder can hold all. + # Each component first checks this folder, then the configs shipped with + # vLLM (if any). If no JSON matches, it uses a hard-coded heuristic. + "VLLM_TUNED_CONFIG_FOLDER": lambda: os.getenv("VLLM_TUNED_CONFIG_FOLDER", None), + # Valid values are container,code_interpreter,web_search_preview + # ex VLLM_GPT_OSS_SYSTEM_TOOL_MCP_LABELS=container,code_interpreter + # If the server_label of your mcp tool is not in this list it will + # be completely ignored. + "VLLM_GPT_OSS_SYSTEM_TOOL_MCP_LABELS": env_set_with_choices( + "VLLM_GPT_OSS_SYSTEM_TOOL_MCP_LABELS", + default=[], + choices=["container", "code_interpreter", "web_search_preview"], + ), + # Allows harmony instructions to be injected on system messages + "VLLM_GPT_OSS_HARMONY_SYSTEM_INSTRUCTIONS": lambda: bool( + int(os.getenv("VLLM_GPT_OSS_HARMONY_SYSTEM_INSTRUCTIONS", "0")) + ), + # Pin the conversation start date injected into the Harmony system + # message. When unset the current date is used, which introduces + # non-determinism (different tokens -> different model behaviour at + # temperature=0). Set to an ISO date string, e.g. "2023-09-12", + # for reproducible inference or testing. + "VLLM_SYSTEM_START_DATE": lambda: os.getenv("VLLM_SYSTEM_START_DATE", None), + # Enable automatic retry when tool call JSON parsing fails + # If enabled, returns an error message to the model to retry + # If disabled (default), raises an exception and fails the request + "VLLM_TOOL_JSON_ERROR_AUTOMATIC_RETRY": lambda: bool( + int(os.getenv("VLLM_TOOL_JSON_ERROR_AUTOMATIC_RETRY", "0")) + ), + # When 1,the model structural tags will be used to enforce the model + # output conforming to the model's tool-calling format and schema. + # Default 0 (off). + "VLLM_ENFORCE_STRICT_TOOL_CALLING": lambda: bool( + int(os.getenv("VLLM_ENFORCE_STRICT_TOOL_CALLING", "0")) + ), + # Add optional custom scopes for profiling, disable to avoid overheads + "VLLM_CUSTOM_SCOPES_FOR_PROFILING": lambda: bool( + int(os.getenv("VLLM_CUSTOM_SCOPES_FOR_PROFILING", "0")) + ), + # Add optional nvtx scopes for profiling, disable to avoid overheads + "VLLM_NVTX_SCOPES_FOR_PROFILING": lambda: bool( + int(os.getenv("VLLM_NVTX_SCOPES_FOR_PROFILING", "0")) + ), + # Represent block hashes in KV cache events as 64-bit integers instead of + # raw bytes. Defaults to True for backward compatibility. + "VLLM_KV_EVENTS_USE_INT_BLOCK_HASHES": lambda: bool( + int(os.getenv("VLLM_KV_EVENTS_USE_INT_BLOCK_HASHES", "1")) + ), + # Name of the shared memory buffer used for object storage. + # Only effective when mm_config.mm_processor_cache_type == "shm". + # Automatically generates a unique UUID-based name per process tree + # if not explicitly set. + "VLLM_OBJECT_STORAGE_SHM_BUFFER_NAME": get_env_or_set_default( + "VLLM_OBJECT_STORAGE_SHM_BUFFER_NAME", + lambda: f"VLLM_OBJECT_STORAGE_SHM_BUFFER_{uuid.uuid4().hex}", + ), + # The size in MB of the buffers (NVL and RDMA) used by DeepEP + "VLLM_DEEPEP_BUFFER_SIZE_MB": lambda: int( + os.getenv("VLLM_DEEPEP_BUFFER_SIZE_MB", "1024") + ), + # Force DeepEP to use intranode kernel for inter-node communication in + # high throughput mode. This is useful archive higher prefill throughput + # on system supports multi-node nvlink (e.g GB200). + "VLLM_DEEPEP_HIGH_THROUGHPUT_FORCE_INTRA_NODE": lambda: bool( + int(os.getenv("VLLM_DEEPEP_HIGH_THROUGHPUT_FORCE_INTRA_NODE", "0")) + ), + # Allow DeepEP to use MNNVL (multi-node nvlink) for internode_ll kernel, + # turn this for better latency on GB200 like system + "VLLM_DEEPEP_LOW_LATENCY_USE_MNNVL": lambda: bool( + int(os.getenv("VLLM_DEEPEP_LOW_LATENCY_USE_MNNVL", "0")) + ), + # The number of SMs/CUs to allocate for communication kernels when + # running DBO; the rest will be allocated to compute. + # Default: 20 on CUDA (SMs), 64 on ROCm (CUs). + "VLLM_DBO_COMM_SMS": lambda: int( + os.getenv( + "VLLM_DBO_COMM_SMS", + "64" + if hasattr(__import__("torch").version, "hip") + and __import__("torch").version.hip is not None + else "20", + ) + ), + # Enable max_autotune & coordinate_descent_tuning in inductor_config + # to compile static shapes passed from compile_sizes in compilation_config + # If set to 1, enable max_autotune; By default, this is enabled (1) + "VLLM_ENABLE_INDUCTOR_MAX_AUTOTUNE": lambda: bool( + int(os.getenv("VLLM_ENABLE_INDUCTOR_MAX_AUTOTUNE", "1")) + ), + # If set to 1, enable coordinate_descent_tuning; + # By default, this is enabled (1) + "VLLM_ENABLE_INDUCTOR_COORDINATE_DESCENT_TUNING": lambda: bool( + int(os.getenv("VLLM_ENABLE_INDUCTOR_COORDINATE_DESCENT_TUNING", "1")) + ), + # Flag to enable NCCL symmetric memory allocation and registration + "VLLM_USE_NCCL_SYMM_MEM": lambda: bool( + int(os.getenv("VLLM_USE_NCCL_SYMM_MEM", "0")) + ), + # NCCL header path + "VLLM_NCCL_INCLUDE_PATH": lambda: os.environ.get("VLLM_NCCL_INCLUDE_PATH", None), + # Flag to enable FBGemm kernels on model execution + # Deprecated: use --linear-backend fbgemm instead. + "VLLM_USE_FBGEMM": deprecated_env( + "VLLM_USE_FBGEMM", + "v0.23", + "Use --linear-backend fbgemm.", + lambda: bool(int(os.getenv("VLLM_USE_FBGEMM", "0"))), + ), + # GC debug config + # - VLLM_GC_DEBUG=0: disable GC debugger + # - VLLM_GC_DEBUG=1: enable GC debugger with gc.collect elpased times + # - VLLM_GC_DEBUG='{"top_objects":5}': enable GC debugger with + # top 5 collected objects + "VLLM_GC_DEBUG": lambda: os.getenv("VLLM_GC_DEBUG", ""), + # Debug workspace allocations. + # logging of workspace resize operations. + "VLLM_DEBUG_WORKSPACE": lambda: bool(int(os.getenv("VLLM_DEBUG_WORKSPACE", "0"))), + # Disables parallel execution of shared_experts via separate cuda stream + "VLLM_DISABLE_SHARED_EXPERTS_STREAM": lambda: bool( + int(os.getenv("VLLM_DISABLE_SHARED_EXPERTS_STREAM", "0")) + ), + # Limits when we run shared_experts in a separate stream. + # We found out that for large batch sizes, the separate stream + # execution is not beneficial (most likely because of the input clone) + # TODO(alexm-redhat): Tune to be more dynamic based on GPU type + "VLLM_SHARED_EXPERTS_STREAM_TOKEN_THRESHOLD": lambda: int( + int(os.getenv("VLLM_SHARED_EXPERTS_STREAM_TOKEN_THRESHOLD", 256)) + ), + # Token-count cutoff for multi-stream overlap of the attention input + # GEMM with auxiliary GEMMs (e.g. fused_wqa_wkv overlapped with indexer + # weights / kv-score projections in DeepSeek-V4). At or below this many + # tokens the FP8 main GEMM has idle SMs to share with the bf16 aux GEMMs + # and overlap is a 5-45% win; above it the FP8 GEMM saturates the device + # and the cross-stream sync becomes pure overhead. Set to 0 to disable + # the multi-stream path entirely. See #PR 41526 for the empirical result + # for the default value of 1024 tokens. + "VLLM_MULTI_STREAM_GEMM_TOKEN_THRESHOLD": lambda: int( + os.getenv("VLLM_MULTI_STREAM_GEMM_TOKEN_THRESHOLD", "1024") + ), + # Format for saving torch.compile cache artifacts + # - "binary": saves as binary file + # Safe for multiple vllm serve processes accessing the same torch compile cache. + # - "unpacked": saves as directory structure (for inspection/debugging) + # NOT multiprocess safe - race conditions may occur with multiple processes. + # Allows viewing and setting breakpoints in Inductor's code output files. + "VLLM_COMPILE_CACHE_SAVE_FORMAT": env_with_choices( + "VLLM_COMPILE_CACHE_SAVE_FORMAT", "binary", ["binary", "unpacked"] + ), + # Flag to control the v2 model runner. If unset, use config defaults. + "VLLM_USE_V2_MODEL_RUNNER": lambda: maybe_convert_bool( + os.getenv("VLLM_USE_V2_MODEL_RUNNER", None) + ), + # Log model inspection after loading. + # If enabled, logs a transformers-style hierarchical view of the model + # with quantization methods and attention backends. + "VLLM_LOG_MODEL_INSPECTION": lambda: bool( + int(os.getenv("VLLM_LOG_MODEL_INSPECTION", "0")) + ), + # Debug logging for --enable-mfu-metrics + "VLLM_DEBUG_MFU_METRICS": lambda: bool( + int(os.getenv("VLLM_DEBUG_MFU_METRICS", "0")) + ), + # Disable using pytorch's pin memory for CPU offloading. + "VLLM_WEIGHT_OFFLOADING_DISABLE_PIN_MEMORY": lambda: bool( + int(os.getenv("VLLM_WEIGHT_OFFLOADING_DISABLE_PIN_MEMORY", "0")) + ), + # Disable using UVA (Unified Virtual Addressing) for CPU offloading. + "VLLM_WEIGHT_OFFLOADING_DISABLE_UVA": lambda: bool( + int(os.getenv("VLLM_WEIGHT_OFFLOADING_DISABLE_UVA", "0")) + ), + # Disable logging of vLLM logo at server startup time. + "VLLM_DISABLE_LOG_LOGO": lambda: bool(int(os.getenv("VLLM_DISABLE_LOG_LOGO", "0"))), + # Disable PDL for LoRA, as enabling PDL with LoRA on SM100 causes + # Triton compilation to fail. + "VLLM_LORA_DISABLE_PDL": lambda: bool(int(os.getenv("VLLM_LORA_DISABLE_PDL", "0"))), + # Enable CUDA compatibility mode for datacenter GPUs with older + # driver versions than the CUDA toolkit major version of vLLM. + "VLLM_ENABLE_CUDA_COMPATIBILITY": lambda: ( + os.environ.get("VLLM_ENABLE_CUDA_COMPATIBILITY", "0").strip().lower() + in ("1", "true") + ), + # Path to the CUDA compatibility libraries when CUDA compatibility is enabled. + "VLLM_CUDA_COMPATIBILITY_PATH": lambda: os.environ.get( + "VLLM_CUDA_COMPATIBILITY_PATH", None + ), + # Skip model name validation in OpenAI API requests. + # When set to 1, any model name will be accepted in the 'model' field + # of API requests. This is useful for proxy/gateway scenarios where + # the actual model is served but different names may be used in requests. + "VLLM_SKIP_MODEL_NAME_VALIDATION": lambda: ( + os.getenv("VLLM_SKIP_MODEL_NAME_VALIDATION", "0").strip().lower() + in ("1", "true") + ), + # Whether it is a scale up launch engine for elastic EP, + # Should only be set by EngineCoreClient. + "VLLM_ELASTIC_EP_SCALE_UP_LAUNCH": lambda: bool( + int(os.getenv("VLLM_ELASTIC_EP_SCALE_UP_LAUNCH", "0")) + ), + # Whether to wait for all requests to drain before sending the + # scaling command in elastic EP. + "VLLM_ELASTIC_EP_DRAIN_REQUESTS": lambda: bool( + int(os.getenv("VLLM_ELASTIC_EP_DRAIN_REQUESTS", "0")) + ), + # If set to 1, enable CUDA graph memory estimation during memory profiling. + # This profiles CUDA graph memory usage to provide more accurate KV cache + # memory allocation. Enabled by default as of v0.21.0 + "VLLM_MEMORY_PROFILER_ESTIMATE_CUDAGRAPHS": lambda: bool( + int(os.getenv("VLLM_MEMORY_PROFILER_ESTIMATE_CUDAGRAPHS", "1")) + ), + # NIXL EP environment variables + "VLLM_NIXL_EP_MAX_NUM_RANKS": lambda: int( + os.getenv("VLLM_NIXL_EP_MAX_NUM_RANKS", "32") + ), + # Whether enable XPU graph on Intel GPU + "VLLM_XPU_ENABLE_XPU_GRAPH": lambda: bool( + int(os.getenv("VLLM_XPU_ENABLE_XPU_GRAPH", "0")) + ), + # whether use xpu specific sample kernel + "VLLM_XPU_USE_SAMPLER_KERNEL": lambda: bool( + int(os.getenv("VLLM_XPU_USE_SAMPLER_KERNEL", "1")) + ), + # Enable simple KV offload. + "VLLM_USE_SIMPLE_KV_OFFLOAD": lambda: bool( + int(os.getenv("VLLM_USE_SIMPLE_KV_OFFLOAD", "0")) + ), + # Whether to enable dual cuda streams for LoRA computation + # (used by both BaseLinearLayerWithLoRA and FusedMoEWithLoRA to + # overlap the base layer compute with the LoRA fast path). + "VLLM_LORA_ENABLE_DUAL_STREAM": lambda: bool( + int(os.getenv("VLLM_LORA_ENABLE_DUAL_STREAM", "0")) + ), + # If set to 1, use Python spinloop extension to poll in a more efficient + # way when using the mp backend. + "VLLM_USE_SPINLOOP_EXT": lambda: bool(int(os.getenv("VLLM_USE_SPINLOOP_EXT", "0"))), +} + + +# --8<-- [end:env-vars-definition] + + +def __getattr__(name: str): + """ + Gets environment variables lazily. + + NOTE: After enable_envs_cache() invocation (which triggered after service + initialization), all environment variables will be cached. + """ + if name in environment_variables: + return environment_variables[name]() + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") + + +def _is_envs_cache_enabled() -> bool: + """Checked if __getattr__ is wrapped with functools.cache""" + global __getattr__ + return hasattr(__getattr__, "cache_clear") + + +def enable_envs_cache() -> None: + """ + Enables caching of environment variables. This is useful for performance + reasons, as it avoids the need to re-evaluate environment variables on + every call. + + NOTE: Currently, it's invoked after service initialization to reduce + runtime overhead. This also means that environment variables should NOT + be updated after the service is initialized. + """ + if _is_envs_cache_enabled(): + # Avoid wrapping functools.cache multiple times + return + # Tag __getattr__ with functools.cache + global __getattr__ + __getattr__ = functools.cache(__getattr__) + + # Cache all environment variables + for key in environment_variables: + __getattr__(key) + + +def disable_envs_cache() -> None: + """ + Resets the environment variables cache. It could be used to isolate environments + between unit tests. + """ + global __getattr__ + # If __getattr__ is wrapped by functions.cache, unwrap the caching layer. + if _is_envs_cache_enabled(): + assert hasattr(__getattr__, "__wrapped__") + __getattr__ = __getattr__.__wrapped__ + + +def __dir__(): + return list(environment_variables.keys()) + + +def is_set(name: str): + """Check if an environment variable is explicitly set.""" + if name in environment_variables: + return name in os.environ + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") + + +def validate_environ(hard_fail: bool) -> None: + for env in os.environ: + if env.startswith("VLLM_") and env not in environment_variables: + if hard_fail: + raise ValueError(f"Unknown vLLM environment variable detected: {env}") + else: + logger.warning("Unknown vLLM environment variable detected: %s", env) + + +def compile_factors() -> dict[str, object]: + """Return env vars used for torch.compile cache keys. + + Start with every known vLLM env var; drop entries in `ignored_factors`; + hash everything else. This keeps the cache key aligned across workers.""" + + ignored_factors: set[str] = { + "MAX_JOBS", + "VLLM_RPC_BASE_PATH", + "VLLM_USE_MODELSCOPE", + "VLLM_RINGBUFFER_WARNING_INTERVAL", + "VLLM_DEBUG_DUMP_PATH", + "VLLM_PORT", + "VLLM_CACHE_ROOT", + "LD_LIBRARY_PATH", + "VLLM_SERVER_DEV_MODE", + "VLLM_DP_MASTER_IP", + "VLLM_DP_MASTER_PORT", + "VLLM_NIXL_SIDE_CHANNEL_HOST", + "VLLM_RANDOMIZE_DP_DUMMY_INPUTS", + "VLLM_CI_USE_S3", + "VLLM_MODEL_REDIRECT_PATH", + "VLLM_HOST_IP", + "VLLM_FORCE_AOT_LOAD", + "S3_ACCESS_KEY_ID", + "S3_SECRET_ACCESS_KEY", + "S3_ENDPOINT_URL", + "VLLM_USAGE_STATS_SERVER", + "VLLM_NO_USAGE_STATS", + "VLLM_DO_NOT_TRACK", + "VLLM_LOGGING_LEVEL", + "VLLM_LOGGING_PREFIX", + "VLLM_LOGGING_STREAM", + "VLLM_LOGGING_CONFIG_PATH", + "VLLM_LOGGING_COLOR", + "VLLM_LOG_STATS_INTERVAL", + "VLLM_DEBUG_LOG_API_SERVER_RESPONSE", + "VLLM_TUNED_CONFIG_FOLDER", + "VLLM_FLASHINFER_AUTOTUNE_CACHE_DIR", + "VLLM_ENGINE_ITERATION_TIMEOUT_S", + "VLLM_HTTP_TIMEOUT_KEEP_ALIVE", + "VLLM_EXECUTE_MODEL_TIMEOUT_SECONDS", + "VLLM_KEEP_ALIVE_ON_ENGINE_DEATH", + "VLLM_IMAGE_FETCH_TIMEOUT", + "VLLM_VIDEO_FETCH_TIMEOUT", + "VLLM_AUDIO_FETCH_TIMEOUT", + "VLLM_MEDIA_CACHE", + "VLLM_MEDIA_CACHE_MAX_SIZE_MB", + "VLLM_MEDIA_CACHE_TTL_HOURS", + "VLLM_MEDIA_FETCH_MAX_RETRIES", + "VLLM_MEDIA_URL_ALLOW_REDIRECTS", + "VLLM_MEDIA_LOADING_THREAD_COUNT", + "VLLM_MAX_AUDIO_CLIP_FILESIZE_MB", + "VLLM_VIDEO_LOADER_BACKEND", + "VLLM_MEDIA_CONNECTOR", + "VLLM_OBJECT_STORAGE_SHM_BUFFER_NAME", + "VLLM_ASSETS_CACHE", + "VLLM_ASSETS_CACHE_MODEL_CLEAN", + "VLLM_WORKER_MULTIPROC_METHOD", + "VLLM_ENABLE_V1_MULTIPROCESSING", + "VLLM_V1_OUTPUT_PROC_CHUNK_SIZE", + "VLLM_CPU_KVCACHE_SPACE", + "VLLM_CPU_MOE_PREPACK", + "VLLM_ZENTORCH_WEIGHT_PREPACK", + "VLLM_TEST_FORCE_LOAD_FORMAT", + "VLLM_ENABLE_CUDA_COMPATIBILITY", + "VLLM_CUDA_COMPATIBILITY_PATH", + "VLLM_SKIP_MODEL_NAME_VALIDATION", + "LOCAL_RANK", + "CUDA_VISIBLE_DEVICES", + "NO_COLOR", + } + + from vllm.config.utils import normalize_value + + factors: dict[str, object] = {} + for factor, getter in environment_variables.items(): + if factor in ignored_factors: + continue + + try: + raw = getter() + except Exception as exc: # pragma: no cover - defensive logging + logger.warning( + "Skipping environment variable %s while hashing compile factors: %s", + factor, + exc, + ) + continue + + factors[factor] = normalize_value(raw) + + ray_noset_env_vars = [ + # Refer to + # https://github.com/ray-project/ray/blob/c584b1ea97b00793d1def71eaf81537d70efba42/python/ray/_private/accelerators/nvidia_gpu.py#L11 + # https://github.com/ray-project/ray/blob/c584b1ea97b00793d1def71eaf81537d70efba42/python/ray/_private/accelerators/amd_gpu.py#L11 + # https://github.com/ray-project/ray/blob/b97d21dab233c2bd8ed7db749a82a1e594222b5c/python/ray/_private/accelerators/amd_gpu.py#L10 + # https://github.com/ray-project/ray/blob/c584b1ea97b00793d1def71eaf81537d70efba42/python/ray/_private/accelerators/npu.py#L12 + # https://github.com/ray-project/ray/blob/c584b1ea97b00793d1def71eaf81537d70efba42/python/ray/_private/accelerators/hpu.py#L12 + # https://github.com/ray-project/ray/blob/c584b1ea97b00793d1def71eaf81537d70efba42/python/ray/_private/accelerators/neuron.py#L14 + # https://github.com/ray-project/ray/blob/c584b1ea97b00793d1def71eaf81537d70efba42/python/ray/_private/accelerators/tpu.py#L38 + # https://github.com/ray-project/ray/blob/c584b1ea97b00793d1def71eaf81537d70efba42/python/ray/_private/accelerators/intel_gpu.py#L10 + # https://github.com/ray-project/ray/blob/c584b1ea97b00793d1def71eaf81537d70efba42/python/ray/_private/accelerators/rbln.py#L10 + "RAY_EXPERIMENTAL_NOSET_CUDA_VISIBLE_DEVICES", + "RAY_EXPERIMENTAL_NOSET_ROCR_VISIBLE_DEVICES", + "RAY_EXPERIMENTAL_NOSET_HIP_VISIBLE_DEVICES", + "RAY_EXPERIMENTAL_NOSET_ASCEND_RT_VISIBLE_DEVICES", + "RAY_EXPERIMENTAL_NOSET_HABANA_VISIBLE_MODULES", + "RAY_EXPERIMENTAL_NOSET_NEURON_RT_VISIBLE_CORES", + "RAY_EXPERIMENTAL_NOSET_TPU_VISIBLE_CHIPS", + "RAY_EXPERIMENTAL_NOSET_ONEAPI_DEVICE_SELECTOR", + "RAY_EXPERIMENTAL_NOSET_RBLN_RT_VISIBLE_DEVICES", + ] + + for var in ray_noset_env_vars: + factors[var] = normalize_value(os.getenv(var)) + + return factors diff --git a/mods/deepseek-v4-flash-dspark/overlay/vllm/model_executor/layers/fused_moe/b12x_moe.py b/mods/deepseek-v4-flash-dspark/overlay/vllm/model_executor/layers/fused_moe/b12x_moe.py new file mode 100644 index 00000000..752480cc --- /dev/null +++ b/mods/deepseek-v4-flash-dspark/overlay/vllm/model_executor/layers/fused_moe/b12x_moe.py @@ -0,0 +1,785 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""B12X modular fused-MoE backend for DeepSeek V4 native MXFP4 weights.""" + +from collections.abc import Callable +from typing import Any, cast + +import torch + +import vllm.envs as envs +import vllm.model_executor.layers.fused_moe.modular_kernel as mk +from vllm.logger import init_logger +from vllm.model_executor.layers.fused_moe.activation import MoEActivation +from vllm.model_executor.layers.fused_moe.config import ( + FusedMoEParallelConfig, + FusedMoEQuantConfig, + RoutingMethodType, +) +from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import ( + TopKWeightAndReduceNoOP, +) +from vllm.model_executor.layers.quantization.utils.quant_utils import ( + QuantKey, + kMxfp4Static, +) +from vllm.model_executor.utils import replace_parameter +from vllm.platforms import current_platform + +logger = init_logger(__name__) + + +def _dtype_element_size(dtype: torch.dtype) -> int: + return torch.empty((), dtype=dtype).element_size() + + +def _ceil_div(a: int, b: int) -> int: + return (int(a) + int(b) - 1) // int(b) + + +def _plan_b12x_moe_fp4_scratch( + *, + tokens: int, + weight_E: int, + k: int, + n: int, + topk: int, + device: torch.device, + dtype: torch.dtype, + activation: str, + quant_mode: str, + source_format: str, + w13_layout: str, + apply_router_weight_on_input: bool = False, + swiglu_limit: float | None = None, +): + from b12x.integration.tp_moe import TPMoEScratchCaps, plan_tp_moe_scratch + + return plan_tp_moe_scratch( + TPMoEScratchCaps( + max_tokens=max(int(tokens), 1), + weight_E=int(weight_E), + k=int(k), + n=int(n), + num_topk=int(topk), + device=device, + dtype=dtype, + core_token_counts=(max(int(tokens), 1),), + route_num_experts=0, + quant_mode=quant_mode, + activation=activation, + apply_router_weight_on_input=apply_router_weight_on_input, + swiglu_limit=swiglu_limit, + source_format=source_format, + w13_layout=w13_layout, + frozen=True, + ) + ) + + +def _b12x_scratch_nbytes(plan: Any) -> int: + specs = plan.scratch_specs() + if len(specs) != 1: + raise RuntimeError(f"expected one b12x MoE scratch buffer, got {len(specs)}") + spec = specs[0] + if spec.dtype != torch.uint8: + raise TypeError(f"expected b12x MoE scratch dtype uint8, got {spec.dtype}") + return int(spec.shape[0]) + + +def _workspace2_as_b12x_scratch( + workspace2: torch.Tensor | None, + plan: Any, +) -> torch.Tensor: + if workspace2 is None: + raise RuntimeError("B12X MoE requires vLLM workspace2 scratch") + if not workspace2.is_contiguous(): + raise ValueError("B12X MoE workspace2 must be contiguous") + scratch = workspace2.view(-1).view(torch.uint8) + required_nbytes = _b12x_scratch_nbytes(plan) + if int(scratch.numel()) < required_nbytes: + raise ValueError( + "B12X MoE workspace2 is too small for planned scratch: " + f"have={int(scratch.numel())} bytes, need={required_nbytes} bytes" + ) + return scratch + + +def _run_b12x_moe_fp4( + *, + a: torch.Tensor, + a1_gscale: torch.Tensor, + w1_fp4: torch.Tensor, + w1_blockscale: torch.Tensor, + w1_alphas: torch.Tensor, + a2_gscale: torch.Tensor, + w2_fp4: torch.Tensor, + w2_blockscale: torch.Tensor, + w2_alphas: torch.Tensor, + output: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + apply_router_weight_on_input: bool, + input_scales_are_reciprocal: bool, + input_scales_static: bool, + activation: str, + quant_mode: str, + unit_scale_contract: bool, + source_format: str, + w13_layout: str, + prepared_w4a16: Any, + swiglu_limit: float | None, + plan: Any, + scratch: torch.Tensor, +) -> None: + """Call b12x MoE with caller-owned live scratch.""" + from b12x.integration.tp_moe import b12x_moe_fp4 + + binding = plan.bind( + scratch=scratch, + a=a, + a1_gscale=a1_gscale, + w1_fp4=w1_fp4, + w1_blockscale=w1_blockscale, + w1_alphas=w1_alphas, + a2_gscale=a2_gscale, + w2_fp4=w2_fp4, + w2_blockscale=w2_blockscale, + w2_alphas=w2_alphas, + topk_weights=topk_weights, + topk_ids=topk_ids, + apply_router_weight_on_input=apply_router_weight_on_input, + output=output, + input_scales_are_reciprocal=input_scales_are_reciprocal, + input_scales_static=input_scales_static, + activation=activation, + quant_mode=quant_mode, + unit_scale_contract=unit_scale_contract, + source_format=source_format, + w13_layout=w13_layout, + prepared_w4a16=prepared_w4a16, + swiglu_limit=swiglu_limit, + ) + b12x_moe_fp4(binding=binding) + + +def _b12x_activation_name(activation: MoEActivation) -> str: + if activation in (MoEActivation.SILU, MoEActivation.SWIGLUOAI): + return "silu" + if activation == MoEActivation.RELU2: + return "relu2" + return activation.value + + +def _parse_b12x_w4a16_tile_config() -> tuple[int, int, int] | None: + raw_config = str(envs.VLLM_B12X_W4A16_FORCE_TILE_CONFIG).strip() + if not raw_config: + return None + parts = [part.strip() for part in raw_config.split(",")] + if len(parts) != 3: + raise ValueError( + "VLLM_B12X_W4A16_FORCE_TILE_CONFIG must be " + "TILE_K,TILE_N,CTA_THREADS, got " + f"{raw_config!r}" + ) + return tuple(int(part) for part in parts) + + +def _forced_b12x_w4a16_tile_blocks_per_sm( + w4a16_kernel: Any, + kwargs: dict[str, Any], + tile_config: tuple[int, int, int], +) -> int | None: + required_names = ( + "problem_m", + "problem_n", + "problem_k", + "top_k", + "moe_block_size", + "sms", + "max_shared_mem", + ) + if any(name not in kwargs for name in required_names): + return None + + tile_k, tile_n, cta_threads = tile_config + try: + cta_m_blocks = w4a16_kernel._covering_count(int(kwargs["moe_block_size"]), 16) + tile_fits = w4a16_kernel._candidate_tile_fits( + problem_n=int(kwargs["problem_n"]), + problem_k=int(kwargs["problem_k"]), + cta_m_blocks=cta_m_blocks, + tile_n=tile_n, + tile_k=tile_k, + cta_threads=cta_threads, + max_shared_mem=int(kwargs["max_shared_mem"]) - 512, + scale_format=kwargs.get("scale_format", "e4m3_k16"), + ) + except Exception: + return None + if not tile_fits: + return None + + try: + return int( + w4a16_kernel._determine_blocks_per_sm( + problem_m=int(kwargs["problem_m"]), + problem_n=int(kwargs["problem_n"]), + top_k=int(kwargs["top_k"]), + cta_threads=cta_threads, + cta_m_blocks=cta_m_blocks, + tile_n=tile_n, + tile_k=tile_k, + uses_m_block_8=int(kwargs["moe_block_size"]) == 8, + sms=int(kwargs["sms"]), + max_shared_mem=int(kwargs["max_shared_mem"]), + scale_format=kwargs.get("scale_format", "e4m3_k16"), + ) + ) + except Exception: + return None + + +def _maybe_apply_b12x_w4a16_selector_override() -> None: + forced_blocks_per_sm = int(envs.VLLM_B12X_W4A16_FORCE_BLOCKS_PER_SM) + forced_tile_config = _parse_b12x_w4a16_tile_config() + max_problem_m = int(envs.VLLM_B12X_W4A16_FORCE_BLOCKS_MAX_M) + if forced_blocks_per_sm < 0: + raise ValueError( + "VLLM_B12X_W4A16_FORCE_BLOCKS_PER_SM must be >= 0, got " + f"{forced_blocks_per_sm}" + ) + if max_problem_m < 0: + raise ValueError( + f"VLLM_B12X_W4A16_FORCE_BLOCKS_MAX_M must be >= 0, got {max_problem_m}" + ) + if forced_blocks_per_sm == 0 and forced_tile_config is None: + return + + try: + from b12x.moe.fused.w4a16 import kernel as w4a16_kernel + except Exception: + logger.warning( + "Could not install B12X W4A16 MoE selector override; b12x " + "kernel module is unavailable.", + exc_info=True, + ) + return + + original_attr = "_vllm_original_select_tile_config" + if hasattr(w4a16_kernel, original_attr): + return + + original_select_tile_config = getattr(w4a16_kernel, "_select_tile_config", None) + if not callable(original_select_tile_config): + logger.warning( + "Could not install B12X W4A16 MoE selector override; " + "_select_tile_config is missing." + ) + return + setattr(w4a16_kernel, original_attr, original_select_tile_config) + + def _vllm_select_tile_config( + *args: Any, **kwargs: Any + ) -> tuple[int, int, int, int]: + selected = original_select_tile_config(*args, **kwargs) + if len(selected) != 4: + return selected + tile_k, tile_n, cta_threads, _blocks_per_sm = selected + problem_m = kwargs.get("problem_m") + if problem_m is None: + return selected + active_max_problem_m = int(envs.VLLM_B12X_W4A16_FORCE_BLOCKS_MAX_M) + if active_max_problem_m > 0 and int(problem_m) > active_max_problem_m: + return selected + active_tile_config = _parse_b12x_w4a16_tile_config() + if active_tile_config is not None: + forced_tile_blocks_per_sm = _forced_b12x_w4a16_tile_blocks_per_sm( + w4a16_kernel, + kwargs, + active_tile_config, + ) + if forced_tile_blocks_per_sm is not None: + tile_k, tile_n, cta_threads = active_tile_config + selected = ( + tile_k, + tile_n, + cta_threads, + forced_tile_blocks_per_sm, + ) + active_forced_blocks_per_sm = int(envs.VLLM_B12X_W4A16_FORCE_BLOCKS_PER_SM) + if active_forced_blocks_per_sm <= 0: + return selected + tile_k, tile_n, cta_threads, _blocks_per_sm = selected + return tile_k, tile_n, cta_threads, active_forced_blocks_per_sm + + w4a16_kernel._select_tile_config = cast( + Callable[..., tuple[int, int, int, int]], _vllm_select_tile_config + ) + logger.info( + "Enabled B12X W4A16 MoE selector override: preserving selected tile " + "unless a tile is forced, tile_config=%s, blocks_per_sm=%d, " + "problem_m<=%d", + forced_tile_config, + forced_blocks_per_sm, + max_problem_m, + ) + + +_maybe_apply_b12x_w4a16_selector_override() + + +def _prepare_b12x_fp4_moe_weights(**kwargs): + _maybe_apply_b12x_w4a16_selector_override() + from b12x.integration import prepare_b12x_fp4_moe_weights + + return prepare_b12x_fp4_moe_weights(**kwargs) + + +def _replace_parameter_with_empty( + layer: torch.nn.Module, + param_name: str, +) -> torch.Tensor | None: + param = getattr(layer, param_name, None) + if not isinstance(param, torch.Tensor): + return None + empty = torch.empty((0,), dtype=param.dtype, device=param.device) + replace_parameter(layer, param_name, empty) + return getattr(layer, param_name) + + +def _set_quant_config_weight_scale( + quant_config: FusedMoEQuantConfig, + weight_name: str, + scale: torch.Tensor, +) -> None: + desc = getattr(quant_config, weight_name, None) + if desc is not None and hasattr(desc, "scale"): + desc.scale = scale + return + + public_name = "w1_scale" if weight_name == "_w1" else "w2_scale" + if hasattr(quant_config, public_name): + setattr(quant_config, public_name, scale) + + +def _maybe_release_cuda_cache(device: torch.device) -> None: + if device.type != "cuda" or _is_current_stream_capturing(): + return + accelerator = getattr(torch, "accelerator", None) + if accelerator is not None: + accelerator.empty_cache() + else: + torch.cuda.empty_cache() + + +def _raise_if_capture_copy_required(tensor: torch.Tensor, description: str) -> None: + if tensor.device.type != "cuda" or not _is_current_stream_capturing(): + return + raise RuntimeError( + f"B12X MoE {description} would allocate during CUDA graph capture" + ) + + +def _is_current_stream_capturing() -> bool: + cuda = getattr(torch, "cuda", None) + if cuda is None: + return False + is_capturing = getattr(cuda, "is_current_stream_capturing", None) + return bool(is_capturing is not None and is_capturing()) + + +def _normalize_b12x_moe_topk_ids(topk_ids: torch.Tensor) -> torch.Tensor: + if topk_ids.dtype != torch.int32: + _raise_if_capture_copy_required(topk_ids, "topk_ids dtype normalization") + topk_ids = topk_ids.to(torch.int32) + if not topk_ids.is_contiguous(): + _raise_if_capture_copy_required(topk_ids, "topk_ids contiguity normalization") + topk_ids = topk_ids.contiguous() + return topk_ids + + +def _normalize_b12x_moe_topk_weights(topk_weights: torch.Tensor) -> torch.Tensor: + if topk_weights.dtype != torch.float32: + _raise_if_capture_copy_required( + topk_weights, + "topk_weights dtype normalization", + ) + topk_weights = topk_weights.to(torch.float32) + if not topk_weights.is_contiguous(): + _raise_if_capture_copy_required( + topk_weights, + "topk_weights contiguity normalization", + ) + topk_weights = topk_weights.contiguous() + return topk_weights + + +def _has_b12x() -> bool: + try: + from b12x.integration.tp_moe import b12x_moe_fp4 # noqa: F401 + + return True + except ImportError: + return False + + +class B12xExperts(mk.FusedMoEExpertsModular): + """Native DeepSeek V4 MXFP4 MoE backend powered by b12x kernels.""" + + def __init__( + self, + moe_config: mk.FusedMoEConfig, + quant_config: FusedMoEQuantConfig, + ): + super().__init__(moe_config, quant_config) + + assert quant_config.weight_quant_dtype == "mxfp4", ( + "B12xExperts only supports native MXFP4 weights, got " + f"{quant_config.weight_quant_dtype}" + ) + + self._prepared_fp4_moe_by_dtype: dict[torch.dtype, Any] = {} + self._released_w4a16_source_scales = False + self._unit_scale_by_device: dict[torch.device, torch.Tensor] = {} + + def _source_format(self) -> str: + return "fp4_e8m0_k32" + + def _w13_layout(self) -> str: + # vLLM DSV4 loading stores fused W13 as [w1/gate, w3/up], which is the + # row order consumed by b12x for the runtime SwiGLU path. + return "w31" + + def _unit_expert_scale( + self, device: torch.device, num_experts: int + ) -> torch.Tensor: + scale = self._unit_scale_by_device.get(device) + if scale is None or scale.numel() != num_experts: + scale = torch.ones(num_experts, dtype=torch.float32, device=device) + self._unit_scale_by_device[device] = scale + return scale + + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + """Prepare b12x-owned W4A16 weights and release one-way sources.""" + device = layer.w13_weight.device + moe_config = getattr(self, "moe_config", None) + params_dtype = getattr(moe_config, "in_dtype", torch.bfloat16) + activation = getattr(layer, "activation", None) + if activation is None: + activation = getattr(moe_config, "activation", MoEActivation.SILU) + activation = cast(MoEActivation, activation) + + self._get_or_prepare_fp4_moe_weights( + w1=layer.w13_weight, + w2=layer.w2_weight, + activation=activation, + params_dtype=params_dtype, + ) + self._release_w4a16_source_scales(layer) + self._release_w4a16_source_weights(layer) + _maybe_release_cuda_cache(device) + + @staticmethod + def _supports_current_device() -> bool: + p = current_platform + return p.is_cuda() and p.is_device_capability_family(120) and _has_b12x() + + @staticmethod + def _supports_no_act_and_mul() -> bool: + return True + + @staticmethod + def _supports_quant_scheme( + weight_key: QuantKey | None, + activation_key: QuantKey | None, + ) -> bool: + return (weight_key, activation_key) == (kMxfp4Static, None) + + @staticmethod + def _supports_activation(activation: MoEActivation) -> bool: + return activation in (MoEActivation.SILU, MoEActivation.SWIGLUOAI) + + @staticmethod + def _supports_parallel_config( + moe_parallel_config: FusedMoEParallelConfig, + ) -> bool: + return ( + not moe_parallel_config.use_ep + and moe_parallel_config.ep_size <= 1 + and not moe_parallel_config.use_all2all_kernels + and not moe_parallel_config.enable_eplb + ) + + @staticmethod + def _supports_routing_method( + routing_method: RoutingMethodType, + weight_key: QuantKey | None, + activation_key: QuantKey | None, + ) -> bool: + return routing_method == RoutingMethodType.DeepseekV4 + + @staticmethod + def activation_format() -> mk.FusedMoEActivationFormat: + return mk.FusedMoEActivationFormat.Standard + + @property + def expects_unquantized_inputs(self) -> bool: + return True + + def supports_expert_map(self) -> bool: + return False + + def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce: + return TopKWeightAndReduceNoOP() + + def _get_or_prepare_fp4_moe_weights( + self, + *, + w1: torch.Tensor, + w2: torch.Tensor, + activation: MoEActivation, + params_dtype: torch.dtype, + ): + prepared = self._prepared_fp4_moe_by_dtype.get(params_dtype) + if prepared is not None and getattr(prepared, "w4a16", None) is not None: + return prepared + + if self._released_w4a16_source_scales: + prepared_dtypes = ", ".join( + str(dtype) for dtype in self._prepared_fp4_moe_by_dtype + ) + raise RuntimeError( + "B12X W4A16 source block scales were already released; " + f"cannot prepare FP4 MoE weights for dtype {params_dtype}. " + f"Prepared dtypes: {prepared_dtypes or 'none'}." + ) + + if w1.device.type == "cuda" and _is_current_stream_capturing(): + raise RuntimeError( + "B12X FP4 MoE weights were not prepared before CUDA " + f"graph capture for dtype {params_dtype}." + ) + assert self.w1_scale is not None and self.w2_scale is not None, ( + "w1_scale and w2_scale must not be None for B12xExperts" + ) + + unit_scale = self._unit_expert_scale(w1.device, int(w1.shape[0])) + prepared = _prepare_b12x_fp4_moe_weights( + source_format=self._source_format(), + w13_layout=self._w13_layout(), + w1_fp4=w1, + w1_blockscale=self.w1_scale, + w1_global_scale=unit_scale, + a1_gscale=unit_scale, + w2_fp4=w2, + w2_blockscale=self.w2_scale, + w2_global_scale=unit_scale, + a2_gscale=unit_scale, + activation=_b12x_activation_name(activation), + params_dtype=params_dtype, + prepare_runtime_alphas=False, + prepare_w4a16=True, + reuse_input_storage=True, + ) + self._prepared_fp4_moe_by_dtype[params_dtype] = prepared + return prepared + + def _lookup_prepared_w4a16(self) -> Any | None: + for prepared in self._prepared_fp4_moe_by_dtype.values(): + w4a16 = getattr(prepared, "w4a16", None) + if w4a16 is not None: + return w4a16 + return None + + def _release_w4a16_source_scales(self, layer: torch.nn.Module) -> None: + if self._released_w4a16_source_scales: + return + + w1_scale = _replace_parameter_with_empty(layer, "w13_weight_scale") + w2_scale = _replace_parameter_with_empty(layer, "w2_weight_scale") + if w1_scale is not None: + _set_quant_config_weight_scale(self.quant_config, "_w1", w1_scale) + if w2_scale is not None: + _set_quant_config_weight_scale(self.quant_config, "_w2", w2_scale) + + self._released_w4a16_source_scales = True + + def _release_w4a16_source_weights(self, layer: torch.nn.Module) -> None: + _replace_parameter_with_empty(layer, "w13_weight") + _replace_parameter_with_empty(layer, "w2_weight") + + def moe_problem_size( + self, + a1: torch.Tensor, + w1: torch.Tensor, + w2: torch.Tensor, + topk_ids: torch.Tensor, + ) -> tuple[int, int, int, int, int]: + if w1.numel() != 0 and w2.numel() != 0: + return super().moe_problem_size(a1, w1, w2, topk_ids) + + prepared_w4a16 = self._lookup_prepared_w4a16() + if prepared_w4a16 is None: + return super().moe_problem_size(a1, w1, w2, topk_ids) + + if a1.dim() == 2: + assert topk_ids.size(0) == a1.size(0), f"{topk_ids.size(0)} != {a1.size(0)}" + m = a1.size(0) + else: + assert a1.dim() == 3 + m = a1.size(1) + + intermediate_size = int(prepared_w4a16.intermediate_size) + n = intermediate_size * 2 + return ( + int(prepared_w4a16.num_experts), + m, + n, + a1.size(-1), + topk_ids.size(1), + ) + + def workspace_shapes( + self, + M: int, + N: int, + K: int, + topk: int, + global_num_experts: int, + local_num_experts: int, + expert_tokens_meta: mk.ExpertTokensMetadata | None, + activation: MoEActivation, + ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]: + prepared_w4a16 = self._lookup_prepared_w4a16() + if prepared_w4a16 is None: + weight_E = int(local_num_experts) + n = max(int(N) // 2, 1) + device = ( + torch.device("cuda", torch.cuda.current_device()) + if torch.cuda.is_available() + else torch.device("cpu") + ) + else: + weight_E = int(prepared_w4a16.num_experts) + n = int(prepared_w4a16.intermediate_size) + w13 = getattr(prepared_w4a16, "w13", None) + device = ( + w13.device + if isinstance(w13, torch.Tensor) + else torch.device("cuda", torch.cuda.current_device()) + ) + workspace_dtype = getattr(self.moe_config, "in_dtype", torch.bfloat16) + plan = _plan_b12x_moe_fp4_scratch( + tokens=max(int(M), 1), + weight_E=weight_E, + k=int(K), + n=n, + topk=int(topk), + device=device, + dtype=workspace_dtype, + activation=_b12x_activation_name(activation), + quant_mode="w4a16", + source_format=self._source_format(), + w13_layout=self._w13_layout(), + swiglu_limit=getattr(self.quant_config, "gemm1_clamp_limit", None), + ) + scratch_elements = max( + 1, + _ceil_div(_b12x_scratch_nbytes(plan), _dtype_element_size(workspace_dtype)), + ) + return (0,), (scratch_elements,), (M, K) + + def apply( + self, + output: torch.Tensor, + hidden_states: torch.Tensor, + w1: torch.Tensor, + w2: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + activation: MoEActivation, + global_num_experts: int, + expert_map: torch.Tensor | None, + a1q_scale: torch.Tensor | None, + a2_scale: torch.Tensor | None, + workspace13: torch.Tensor | None, + workspace2: torch.Tensor | None, + expert_tokens_meta: mk.ExpertTokensMetadata | None, + apply_router_weight_on_input: bool | None, + ) -> None: + prepared = self._get_or_prepare_fp4_moe_weights( + w1=w1, + w2=w2, + activation=activation, + params_dtype=hidden_states.dtype, + ) + prepared_w4a16 = prepared.w4a16 + assert prepared_w4a16 is not None + assert self.w1_scale is not None and self.w2_scale is not None, ( + "w1_scale and w2_scale must not be None for B12xExperts" + ) + + if expert_map is not None: + raise RuntimeError( + "B12X MoE does not support expert_map with the current b12x_moe_fp4 API" + ) + + num_experts = int(prepared_w4a16.num_experts) + unit_scale = self._unit_expert_scale(hidden_states.device, num_experts) + topk_ids = _normalize_b12x_moe_topk_ids(topk_ids) + topk_weights = _normalize_b12x_moe_topk_weights(topk_weights) + plan = _plan_b12x_moe_fp4_scratch( + tokens=int(hidden_states.shape[0]), + weight_E=num_experts, + k=int(hidden_states.shape[1]), + n=int(prepared_w4a16.intermediate_size), + topk=int(topk_ids.shape[1]), + device=hidden_states.device, + dtype=hidden_states.dtype, + activation=_b12x_activation_name(activation), + quant_mode="w4a16", + source_format=self._source_format(), + w13_layout=self._w13_layout(), + apply_router_weight_on_input=( + apply_router_weight_on_input + if apply_router_weight_on_input is not None + else False + ), + swiglu_limit=getattr(self.quant_config, "gemm1_clamp_limit", None), + ) + scratch = _workspace2_as_b12x_scratch(workspace2, plan) + + _run_b12x_moe_fp4( + a=hidden_states, + a1_gscale=unit_scale, + w1_fp4=w1, + w1_blockscale=self.w1_scale, + w1_alphas=unit_scale, + a2_gscale=unit_scale, + w2_fp4=w2, + w2_blockscale=self.w2_scale, + w2_alphas=unit_scale, + topk_weights=topk_weights, + topk_ids=topk_ids, + apply_router_weight_on_input=( + apply_router_weight_on_input + if apply_router_weight_on_input is not None + else False + ), + output=output, + input_scales_are_reciprocal=True, + input_scales_static=True, + activation=_b12x_activation_name(activation), + quant_mode="w4a16", + unit_scale_contract=True, + source_format=self._source_format(), + w13_layout=self._w13_layout(), + prepared_w4a16=prepared_w4a16, + swiglu_limit=getattr(self.quant_config, "gemm1_clamp_limit", None), + plan=plan, + scratch=scratch, + ) + + def moe_sum(self, input: torch.Tensor, output: torch.Tensor) -> None: + raise NotImplementedError("LoRA is not supported for B12xExperts") diff --git a/mods/deepseek-v4-flash-dspark/overlay/vllm/model_executor/models/registry.py b/mods/deepseek-v4-flash-dspark/overlay/vllm/model_executor/models/registry.py new file mode 100644 index 00000000..ff4196ec --- /dev/null +++ b/mods/deepseek-v4-flash-dspark/overlay/vllm/model_executor/models/registry.py @@ -0,0 +1,1412 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +""" +Whenever you add an architecture to this page, please also update +`tests/models/registry.py` with example HuggingFace models for it. +""" + +import importlib +import importlib.util +import json +import os +import pickle +import subprocess +import sys +import tempfile +from abc import ABC, abstractmethod +from collections.abc import Callable, Set +from dataclasses import asdict, dataclass, field +from functools import lru_cache +from pathlib import Path +from typing import TYPE_CHECKING, Any, TypeVar + +import torch.nn as nn +import transformers + +from vllm import envs +from vllm.config import ( + ModelConfig, + iter_architecture_defaults, + try_match_architecture_defaults, +) +from vllm.logger import init_logger +from vllm.logging_utils import logtime +from vllm.tasks import ScoreType +from vllm.transformers_utils.dynamic_module import try_get_class_from_dynamic_module +from vllm.utils.hashing import safe_hash + +if TYPE_CHECKING: + from vllm.config.model import AttnTypeStr + from vllm.config.pooler import SequencePoolingType, TokenPoolingType +else: + AttnTypeStr = Any + SequencePoolingType = Any + TokenPoolingType = Any + + +from .interfaces import ( + has_inner_state, + has_noops, + is_attention_free, + is_hybrid, + requires_raw_input_tokens, + supports_mamba_prefix_caching, + supports_multimodal, + supports_multimodal_encoder_tp_data, + supports_multimodal_raw_input_only, + supports_pp, + supports_transcription, +) +from .interfaces_base import ( + get_attn_type, + get_default_seq_pooling_type, + get_default_tok_pooling_type, + get_score_type, + is_pooling_model, + is_text_generation_model, +) + +logger = init_logger(__name__) + +_TEXT_GENERATION_MODELS = { + # [Decoder-only] + "AfmoeForCausalLM": ("afmoe", "AfmoeForCausalLM"), + "ApertusForCausalLM": ("apertus", "ApertusForCausalLM"), + "AquilaModel": ("llama", "LlamaForCausalLM"), + "AquilaForCausalLM": ("llama", "LlamaForCausalLM"), # AquilaChat2 + "ArceeForCausalLM": ("arcee", "ArceeForCausalLM"), + "ArcticForCausalLM": ("arctic", "ArcticForCausalLM"), + "AXK1ForCausalLM": ("AXK1", "AXK1ForCausalLM"), + # baichuan-7b, upper case 'C' in the class name + "BaiChuanForCausalLM": ("baichuan", "BaiChuanForCausalLM"), + # baichuan-13b, lower case 'c' in the class name + "BaichuanForCausalLM": ("baichuan", "BaichuanForCausalLM"), + "BailingMoeForCausalLM": ("bailing_moe", "BailingMoeForCausalLM"), + "BailingMoeV2ForCausalLM": ("bailing_moe", "BailingMoeV2ForCausalLM"), + "BailingMoeV2_5ForCausalLM": ("bailing_moe_linear", "BailingMoeV25ForCausalLM"), + "BambaForCausalLM": ("bamba", "BambaForCausalLM"), + "BloomForCausalLM": ("bloom", "BloomForCausalLM"), + "ChatGLMModel": ("chatglm", "ChatGLMForCausalLM"), + "ChatGLMForConditionalGeneration": ("chatglm", "ChatGLMForCausalLM"), + "CohereForCausalLM": ("commandr", "CohereForCausalLM"), + "Cohere2ForCausalLM": ("commandr", "CohereForCausalLM"), + "Cohere2MoeForCausalLM": ("cohere2_moe", "Cohere2MoeForCausalLM"), + "CwmForCausalLM": ("llama", "LlamaForCausalLM"), + "DbrxForCausalLM": ("dbrx", "DbrxForCausalLM"), + "DeciLMForCausalLM": ("nemotron_nas", "DeciLMForCausalLM"), + "DeepseekForCausalLM": ("deepseek_v2", "DeepseekForCausalLM"), + "DeepseekV2ForCausalLM": ("deepseek_v2", "DeepseekV2ForCausalLM"), + "DeepseekV3ForCausalLM": ("deepseek_v2", "DeepseekV3ForCausalLM"), + "DeepseekV32ForCausalLM": ("deepseek_v2", "DeepseekV3ForCausalLM"), + "DeepseekV4ForCausalLM": ("vllm.models.deepseek_v4", "DeepseekV4ForCausalLM"), + "Dots1ForCausalLM": ("dots1", "Dots1ForCausalLM"), + "Ernie4_5ForCausalLM": ("ernie45", "Ernie4_5ForCausalLM"), + "Ernie4_5_MoeForCausalLM": ("ernie45_moe", "Ernie4_5_MoeForCausalLM"), + "ExaoneForCausalLM": ("exaone", "ExaoneForCausalLM"), + "Exaone4ForCausalLM": ("exaone4", "Exaone4ForCausalLM"), + "ExaoneMoEForCausalLM": ("exaone_moe", "ExaoneMoeForCausalLM"), + "Fairseq2LlamaForCausalLM": ("fairseq2_llama", "Fairseq2LlamaForCausalLM"), + "FalconForCausalLM": ("falcon", "FalconForCausalLM"), + "FalconMambaForCausalLM": ("mamba", "MambaForCausalLM"), + "FalconH1ForCausalLM": ("falcon_h1", "FalconH1ForCausalLM"), + "FlexOlmoForCausalLM": ("flex_olmo", "FlexOlmoForCausalLM"), + "GemmaForCausalLM": ("gemma", "GemmaForCausalLM"), + "Gemma2ForCausalLM": ("gemma2", "Gemma2ForCausalLM"), + "Gemma3ForCausalLM": ("gemma3", "Gemma3ForCausalLM"), + "Rnj1ForCausalLM": ("rnj1", "Rnj1ForCausalLM"), + "Gemma3nForCausalLM": ("gemma3n", "Gemma3nForCausalLM"), + "Gemma4ForCausalLM": ("gemma4", "Gemma4ForCausalLM"), + "Qwen3NextForCausalLM": ("qwen3_next", "Qwen3NextForCausalLM"), + "GlmForCausalLM": ("glm", "GlmForCausalLM"), + "Glm4ForCausalLM": ("glm4", "Glm4ForCausalLM"), + "Glm4MoeForCausalLM": ("glm4_moe", "Glm4MoeForCausalLM"), + "Glm4MoeLiteForCausalLM": ("glm4_moe_lite", "Glm4MoeLiteForCausalLM"), + "GlmMoeDsaForCausalLM": ("deepseek_v2", "GlmMoeDsaForCausalLM"), + "GptOssForCausalLM": ("gpt_oss", "GptOssForCausalLM"), + "GPT2LMHeadModel": ("gpt2", "GPT2LMHeadModel"), + "GPTBigCodeForCausalLM": ("gpt_bigcode", "GPTBigCodeForCausalLM"), + "GPTJForCausalLM": ("gpt_j", "GPTJForCausalLM"), + "GPTNeoXForCausalLM": ("gpt_neox", "GPTNeoXForCausalLM"), + "GraniteForCausalLM": ("granite", "GraniteForCausalLM"), + "GraniteMoeForCausalLM": ("granitemoe", "GraniteMoeForCausalLM"), + "GraniteMoeHybridForCausalLM": ("granitemoehybrid", "GraniteMoeHybridForCausalLM"), + "GraniteMoeSharedForCausalLM": ("granitemoeshared", "GraniteMoeSharedForCausalLM"), + "GritLM": ("gritlm", "GritLM"), + "Grok1ModelForCausalLM": ("grok1", "GrokForCausalLM"), + "Grok1ForCausalLM": ("grok1", "GrokForCausalLM"), + "HunYuanMoEV1ForCausalLM": ("hunyuan_v1", "HunYuanMoEV1ForCausalLM"), + "HunYuanDenseV1ForCausalLM": ("hunyuan_v1", "HunYuanDenseV1ForCausalLM"), + "HYV3ForCausalLM": ("hy_v3", "HYV3ForCausalLM"), + "HCXVisionForCausalLM": ("hyperclovax_vision", "HCXVisionForCausalLM"), + "HCXVisionV2ForCausalLM": ("hyperclovax_vision_v2", "HCXVisionV2ForCausalLM"), + "HyperCLOVAXForCausalLM": ("hyperclovax", "HyperCLOVAXForCausalLM"), + "InternLMForCausalLM": ("llama", "LlamaForCausalLM"), + "InternLM2ForCausalLM": ("internlm2", "InternLM2ForCausalLM"), + "InternLM2VEForCausalLM": ("internlm2_ve", "InternLM2VEForCausalLM"), + "InternLM3ForCausalLM": ("llama", "LlamaForCausalLM"), + "IQuestCoderForCausalLM": ("llama", "LlamaForCausalLM"), + "IQuestLoopCoderForCausalLM": ("iquest_loopcoder", "IQuestLoopCoderForCausalLM"), + "JAISLMHeadModel": ("jais", "JAISLMHeadModel"), + "Jais2ForCausalLM": ("jais2", "Jais2ForCausalLM"), + "JambaForCausalLM": ("jamba", "JambaForCausalLM"), + "KimiLinearForCausalLM": ("kimi_linear", "KimiLinearForCausalLM"), + "Lfm2ForCausalLM": ("lfm2", "Lfm2ForCausalLM"), + "Lfm2MoeForCausalLM": ("lfm2_moe", "Lfm2MoeForCausalLM"), + "LagunaForCausalLM": ("laguna", "LagunaForCausalLM"), + "LlamaForCausalLM": ("llama", "LlamaForCausalLM"), + "Llama4ForCausalLM": ("llama4", "Llama4ForCausalLM"), + # For decapoda-research/llama-* + "LLaMAForCausalLM": ("llama", "LlamaForCausalLM"), + "LongcatFlashForCausalLM": ("longcat_flash", "LongcatFlashForCausalLM"), + "MambaForCausalLM": ("mamba", "MambaForCausalLM"), + "Mamba2ForCausalLM": ("mamba2", "Mamba2ForCausalLM"), + "MiniCPMForCausalLM": ("minicpm", "MiniCPMForCausalLM"), + "MiniCPM3ForCausalLM": ("minicpm3", "MiniCPM3ForCausalLM"), + "MiniMaxForCausalLM": ("minimax_text_01", "MiniMaxText01ForCausalLM"), + "MiniMaxText01ForCausalLM": ("minimax_text_01", "MiniMaxText01ForCausalLM"), + "MiniMaxM1ForCausalLM": ("minimax_text_01", "MiniMaxText01ForCausalLM"), + "MiniMaxM2ForCausalLM": ("minimax_m2", "MiniMaxM2ForCausalLM"), + "Ministral3ForCausalLM": ("mistral", "MistralForCausalLM"), + "MistralForCausalLM": ("mistral", "MistralForCausalLM"), + "MistralLarge3ForCausalLM": ("mistral_large_3", "MistralLarge3ForCausalLM"), + "MixtralForCausalLM": ("mixtral", "MixtralForCausalLM"), + # transformers's mpt class has lower case + "MptForCausalLM": ("mpt", "MPTForCausalLM"), + "MPTForCausalLM": ("mpt", "MPTForCausalLM"), + "MiMoForCausalLM": ("mimo", "MiMoForCausalLM"), + "MiMoV2FlashForCausalLM": ("mimo_v2", "MiMoV2FlashForCausalLM"), + "MiMoV2ForCausalLM": ("mimo_v2", "MiMoV2ForCausalLM"), + "NemotronForCausalLM": ("nemotron", "NemotronForCausalLM"), + "NemotronHForCausalLM": ("nemotron_h", "NemotronHForCausalLM"), + "NemotronHPuzzleForCausalLM": ("nemotron_h", "NemotronHForCausalLM"), + "OlmoForCausalLM": ("olmo", "OlmoForCausalLM"), + "Olmo2ForCausalLM": ("olmo2", "Olmo2ForCausalLM"), + "Olmo3ForCausalLM": ("olmo2", "Olmo2ForCausalLM"), + "OlmoHybridForCausalLM": ("olmo_hybrid", "OlmoHybridForCausalLM"), + "OlmoeForCausalLM": ("olmoe", "OlmoeForCausalLM"), + "OPTForCausalLM": ("opt", "OPTForCausalLM"), + "OrionForCausalLM": ("orion", "OrionForCausalLM"), + "OuroForCausalLM": ("ouro", "OuroForCausalLM"), + "PanguEmbeddedForCausalLM": ("openpangu", "PanguEmbeddedForCausalLM"), + "PanguProMoEV2ForCausalLM": ("openpangu", "PanguProMoEV2ForCausalLM"), + "PanguUltraMoEForCausalLM": ("openpangu", "PanguUltraMoEForCausalLM"), + "Param2MoEForCausalLM": ("param2moe", "Param2MoEForCausalLM"), + "PersimmonForCausalLM": ("persimmon", "PersimmonForCausalLM"), + "PhiForCausalLM": ("phi", "PhiForCausalLM"), + "Phi3ForCausalLM": ("phi3", "Phi3ForCausalLM"), + "PhiMoEForCausalLM": ("phimoe", "PhiMoEForCausalLM"), + "Plamo2ForCausalLM": ("plamo2", "Plamo2ForCausalLM"), + "Plamo3ForCausalLM": ("plamo3", "Plamo3ForCausalLM"), + "QWenLMHeadModel": ("qwen", "QWenLMHeadModel"), + "Qwen2ForCausalLM": ("qwen2", "Qwen2ForCausalLM"), + "Qwen2MoeForCausalLM": ("qwen2_moe", "Qwen2MoeForCausalLM"), + "Qwen3ForCausalLM": ("qwen3", "Qwen3ForCausalLM"), + "Qwen3MoeForCausalLM": ("qwen3_moe", "Qwen3MoeForCausalLM"), + "RWForCausalLM": ("falcon", "FalconForCausalLM"), + "SarvamMoEForCausalLM": ("sarvam", "SarvamMoEForCausalLM"), + "SarvamMLAForCausalLM": ("sarvam", "SarvamMLAForCausalLM"), + "SeedOssForCausalLM": ("seed_oss", "SeedOssForCausalLM"), + "Step1ForCausalLM": ("step1", "Step1ForCausalLM"), + "Step3TextForCausalLM": ("step3_text", "Step3TextForCausalLM"), + "Step3p5ForCausalLM": ("step3p5", "Step3p5ForCausalLM"), + "StableLMEpochForCausalLM": ("stablelm", "StablelmForCausalLM"), + "StableLmForCausalLM": ("stablelm", "StablelmForCausalLM"), + "Starcoder2ForCausalLM": ("starcoder2", "Starcoder2ForCausalLM"), + "SolarForCausalLM": ("solar", "SolarForCausalLM"), + "TeleChatForCausalLM": ("telechat2", "TeleChat2ForCausalLM"), + "TeleChat2ForCausalLM": ("telechat2", "TeleChat2ForCausalLM"), + "TeleChat3ForCausalLM": ("llama", "LlamaForCausalLM"), + "TeleFLMForCausalLM": ("teleflm", "TeleFLMForCausalLM"), + "XverseForCausalLM": ("llama", "LlamaForCausalLM"), + "Zamba2ForCausalLM": ("zamba2", "Zamba2ForCausalLM"), +} + +_EMBEDDING_MODELS = { + # [Text-only] + "BertModel": ("bert", "BertEmbeddingModel"), + "BertSpladeSparseEmbeddingModel": ("bert", "BertSpladeSparseEmbeddingModel"), + "ErnieModel": ("ernie", "ErnieEmbeddingModel"), + "BgeM3EmbeddingModel": ("roberta", "BgeM3EmbeddingModel"), + "DeciLMForCausalLM": ("nemotron_nas", "DeciLMForCausalLM"), + "Gemma2Model": ("gemma2", "Gemma2ForCausalLM"), + "Gemma3TextModel": ("gemma3", "Gemma3Model"), + "GlmForCausalLM": ("glm", "GlmForCausalLM"), + "GritLM": ("gritlm", "GritLM"), + "GteModel": ("bert_with_rope", "SnowflakeGteNewModel"), + "GteNewModel": ("bert_with_rope", "GteNewModel"), + "JinaEmbeddingsV5Model": ("jina", "JinaEmbeddingsV5Model"), + "LlamaBidirectionalModel": ("llama", "LlamaBidirectionalModel"), + "LlamaModel": ("llama", "LlamaForCausalLM"), + **{ + # Multiple models share the same architecture, so we include them all + k: (mod, arch) + for k, (mod, arch) in _TEXT_GENERATION_MODELS.items() + if arch == "LlamaForCausalLM" + }, + "MistralModel": ("llama", "LlamaForCausalLM"), + "ModernBertModel": ("modernbert", "ModernBertModel"), + "NomicBertModel": ("bert_with_rope", "NomicBertModel"), + "Phi3ForCausalLM": ("phi3", "Phi3ForCausalLM"), + "Qwen2Model": ("qwen2", "Qwen2ForCausalLM"), + "Qwen2ForCausalLM": ("qwen2", "Qwen2ForCausalLM"), + "RobertaForMaskedLM": ("roberta", "RobertaEmbeddingModel"), + "RobertaModel": ("roberta", "RobertaEmbeddingModel"), + "TeleChatForCausalLM": ("telechat2", "TeleChat2ForCausalLM"), + "TeleChat2ForCausalLM": ("telechat2", "TeleChat2ForCausalLM"), + "VoyageQwen3BidirectionalEmbedModel": ( + "voyage", + "VoyageQwen3BidirectionalEmbedModel", + ), + "XLMRobertaModel": ("roberta", "RobertaEmbeddingModel"), + # [Multimodal] + "CLIPModel": ("clip", "CLIPEmbeddingModel"), + "ColPaliForRetrieval": ("colpali", "ColPaliModel"), + "LlamaNemotronVLModel": ("nemotron_vl", "LlamaNemotronVLForEmbedding"), + "LlavaNextForConditionalGeneration": ( + "llava_next", + "LlavaNextForConditionalGeneration", + ), + "Phi3VForCausalLM": ("phi3v", "Phi3VForCausalLM"), + "Qwen2VLForConditionalGeneration": ("qwen2_vl", "Qwen2VLForConditionalGeneration"), + "SiglipModel": ("siglip", "SiglipEmbeddingModel"), + # Technically Terratorch models work on images, both in + # input and output. I am adding it here because it piggy-backs on embedding + # models for the time being. + "PrithviGeoSpatialMAE": ("terratorch", "Terratorch"), + "Terratorch": ("terratorch", "Terratorch"), +} + +_LATE_INTERACTION_MODELS = { + # [Text-only] + "HF_ColBERT": ("colbert", "ColBERTModel"), + "ColBERTModernBertModel": ("colbert", "ColBERTModernBertModel"), + "ColBERTJinaRobertaModel": ("colbert", "ColBERTJinaRobertaModel"), + "ColBERTLfm2Model": ("colbert", "ColBERTLfm2Model"), + "JinaForRanking": ("jina", "JinaForRanking"), + # [Multimodal] + "ColModernVBertForRetrieval": ("colmodernvbert", "ColModernVBertForRetrieval"), + "ColPaliForRetrieval": ("colpali", "ColPaliModel"), + "ColQwen3": ("colqwen3", "ColQwen3Model"), + "OpsColQwen3Model": ("colqwen3", "ColQwen3Model"), + "ColQwen3_5": ("colqwen3_5", "ColQwen3_5Model"), + "Qwen3VLNemotronEmbedModel": ("colqwen3", "ColQwen3Model"), +} + +_REWARD_MODELS = { + "InternLM2ForRewardModel": ("internlm2", "InternLM2ForRewardModel"), + "Qwen2ForRewardModel": ("qwen2_rm", "Qwen2ForRewardModel"), + "Qwen2ForProcessRewardModel": ("qwen2_rm", "Qwen2ForProcessRewardModel"), +} + +_TOKEN_CLASSIFICATION_MODELS = { + "BertForTokenClassification": ("bert", "BertForTokenClassification"), + "ErnieForTokenClassification": ("ernie", "ErnieForTokenClassification"), + "ModernBertForTokenClassification": ( + "modernbert", + "ModernBertForTokenClassification", + ), + "Qwen3ASRForcedAlignerForTokenClassification": ( + "qwen3_asr_forced_aligner", + "Qwen3ASRForcedAlignerForTokenClassification", + ), +} + +_SEQUENCE_CLASSIFICATION_MODELS = { + "BertForSequenceClassification": ("bert", "BertForSequenceClassification"), + "GPT2ForSequenceClassification": ("gpt2", "GPT2ForSequenceClassification"), + "ErnieForSequenceClassification": ("ernie", "ErnieForSequenceClassification"), + "GteNewForSequenceClassification": ( + "bert_with_rope", + "GteNewForSequenceClassification", + ), + "JambaForSequenceClassification": ("jamba", "JambaForSequenceClassification"), + "LlamaBidirectionalForSequenceClassification": ( + "llama", + "LlamaBidirectionalForSequenceClassification", + ), + "ModernBertForSequenceClassification": ( + "modernbert", + "ModernBertForSequenceClassification", + ), + "RobertaForSequenceClassification": ("roberta", "RobertaForSequenceClassification"), + "XLMRobertaForSequenceClassification": ( + "roberta", + "RobertaForSequenceClassification", + ), + # [Multimodal] + "JinaVLForRanking": ("jina_vl", "JinaVLForSequenceClassification"), + "LlamaNemotronVLForSequenceClassification": ( + "nemotron_vl", + "LlamaNemotronVLForSequenceClassification", + ), +} + +_MULTIMODAL_MODELS = { + # [Decoder-only] + "AriaForConditionalGeneration": ("aria", "AriaForConditionalGeneration"), + "AudioFlamingo3ForConditionalGeneration": ( + "audioflamingo3", + "AudioFlamingo3ForConditionalGeneration", + ), + "MusicFlamingoForConditionalGeneration": ( + "musicflamingo", + "MusicFlamingoForConditionalGeneration", + ), + "AyaVisionForConditionalGeneration": ( + "aya_vision", + "AyaVisionForConditionalGeneration", + ), + "BagelForConditionalGeneration": ("bagel", "BagelForConditionalGeneration"), + "BeeForConditionalGeneration": ("bee", "BeeForConditionalGeneration"), + "Blip2ForConditionalGeneration": ("blip2", "Blip2ForConditionalGeneration"), + "ChameleonForConditionalGeneration": ( + "chameleon", + "ChameleonForConditionalGeneration", + ), + "Cheers": ("cheers", "CheersForConditionalGeneration"), + "CheersForConditionalGeneration": ("cheers", "CheersForConditionalGeneration"), + "Cohere2VisionForConditionalGeneration": ( + "cohere2_vision", + "Cohere2VisionForConditionalGeneration", + ), + "DeepseekVLV2ForCausalLM": ("deepseek_vl2", "DeepseekVLV2ForCausalLM"), + "DeepseekOCRForCausalLM": ("deepseek_ocr", "DeepseekOCRForCausalLM"), + "DeepseekOCR2ForCausalLM": ("deepseek_ocr2", "DeepseekOCR2ForCausalLM"), + "DotsOCRForCausalLM": ("dots_ocr", "DotsOCRForCausalLM"), + "Eagle2_5_VLForConditionalGeneration": ( + "eagle2_5_vl", + "Eagle2_5_VLForConditionalGeneration", + ), + "Ernie4_5_VLMoeForConditionalGeneration": ( + "ernie45_vl", + "Ernie4_5_VLMoeForConditionalGeneration", + ), + "Exaone4_5_ForConditionalGeneration": ( + "exaone4_5", + "Exaone4_5_ForConditionalGeneration", + ), # noqa: E501 + "FireRedASR2ForConditionalGeneration": ( + "fireredasr2", + "FireRedASR2ForConditionalGeneration", + ), + "FunASRForConditionalGeneration": ("funasr", "FunASRForConditionalGeneration"), + "FireRedLIDForConditionalGeneration": ( + "fireredlid", + "FireRedLIDForConditionalGeneration", + ), + "FunAudioChatForConditionalGeneration": ( + "funaudiochat", + "FunAudioChatForConditionalGeneration", + ), + "FuyuForCausalLM": ("fuyu", "FuyuForCausalLM"), + "Gemma3ForConditionalGeneration": ("gemma3_mm", "Gemma3ForConditionalGeneration"), + "Gemma3nForConditionalGeneration": ( + "gemma3n_mm", + "Gemma3nForConditionalGeneration", + ), + "Gemma4ForConditionalGeneration": ("gemma4_mm", "Gemma4ForConditionalGeneration"), + "GlmAsrForConditionalGeneration": ("glmasr", "GlmAsrForConditionalGeneration"), + "GLM4VForCausalLM": ("glm4v", "GLM4VForCausalLM"), + "Glm4vForConditionalGeneration": ("glm4_1v", "Glm4vForConditionalGeneration"), + "Glm4vMoeForConditionalGeneration": ("glm4_1v", "Glm4vMoeForConditionalGeneration"), + "GlmOcrForConditionalGeneration": ("glm_ocr", "GlmOcrForConditionalGeneration"), + "GraniteSpeechForConditionalGeneration": ( + "granite_speech", + "GraniteSpeechForConditionalGeneration", + ), + "Granite4VisionForConditionalGeneration": ( + "granite4_vision", + "Granite4VisionForConditionalGeneration", + ), + "H2OVLChatModel": ("h2ovl", "H2OVLChatModel"), + "HunYuanVLForConditionalGeneration": ( + "hunyuan_vision", + "HunYuanVLForConditionalGeneration", + ), + "InternVLChatModel": ("internvl", "InternVLChatModel"), + "InternS1ForConditionalGeneration": ( + "interns1", + "InternS1ForConditionalGeneration", + ), + "InternVLForConditionalGeneration": ( + "interns1", + "InternS1ForConditionalGeneration", + ), + "InternS1ProForConditionalGeneration": ( + "interns1_pro", + "InternS1ProForConditionalGeneration", + ), + "InternS2PreviewForConditionalGeneration": ( + "interns2_preview", + "InternS2PreviewForConditionalGeneration", + ), + "Idefics3ForConditionalGeneration": ( + "idefics3", + "Idefics3ForConditionalGeneration", + ), + "IsaacForConditionalGeneration": ("isaac", "IsaacForConditionalGeneration"), + "KananaVForConditionalGeneration": ("kanana_v", "KananaVForConditionalGeneration"), + "KeyeForConditionalGeneration": ("keye", "KeyeForConditionalGeneration"), + "KeyeVL1_5ForConditionalGeneration": ( + "keye_vl1_5", + "KeyeVL1_5ForConditionalGeneration", + ), + "KimiVLForConditionalGeneration": ("kimi_vl", "KimiVLForConditionalGeneration"), + "KimiK25ForConditionalGeneration": ("kimi_k25", "KimiK25ForConditionalGeneration"), + "MoonshotKimiaForCausalLM": ("kimi_audio", "KimiAudioForConditionalGeneration"), + "LightOnOCRForConditionalGeneration": ( + "lightonocr", + "LightOnOCRForConditionalGeneration", + ), + "Lfm2VlForConditionalGeneration": ("lfm2_vl", "Lfm2VLForConditionalGeneration"), + "Llama4ForConditionalGeneration": ("mllama4", "Llama4ForConditionalGeneration"), + "Llama_Nemotron_Nano_VL": ("nemotron_vl", "LlamaNemotronVLChatModel"), + "LlavaForConditionalGeneration": ("llava", "LlavaForConditionalGeneration"), + "LlavaNextForConditionalGeneration": ( + "llava_next", + "LlavaNextForConditionalGeneration", + ), + "LlavaNextVideoForConditionalGeneration": ( + "llava_next_video", + "LlavaNextVideoForConditionalGeneration", + ), + "LlavaOnevisionForConditionalGeneration": ( + "llava_onevision", + "LlavaOnevisionForConditionalGeneration", + ), + "MantisForConditionalGeneration": ("llava", "MantisForConditionalGeneration"), + "MiDashengLMModel": ("midashenglm", "MiDashengLMModel"), + "MiMoV2OmniForCausalLM": ("mimo_v2_omni", "MiMoV2OmniForCausalLM"), + "MiniMaxVL01ForConditionalGeneration": ( + "minimax_vl_01", + "MiniMaxVL01ForConditionalGeneration", + ), + "MiniCPMO": ("minicpmo", "MiniCPMO"), + "MiniCPMV": ("minicpmv", "MiniCPMV"), + "MiniCPMV4_6ForConditionalGeneration": ( + "minicpmv4_6", + "MiniCPMV4_6ForConditionalGeneration", + ), + "Mistral3ForConditionalGeneration": ( + "mistral3", + "Mistral3ForConditionalGeneration", + ), + "MolmoForCausalLM": ("molmo", "MolmoForCausalLM"), + "Molmo2ForConditionalGeneration": ("molmo2", "Molmo2ForConditionalGeneration"), + "Moondream3ForCausalLM": ("moondream3", "Moondream3ForCausalLM"), + "HfMoondream": ("moondream3", "Moondream3ForCausalLM"), + "NemotronH_Nano_VL_V2": ("nano_nemotron_vl", "NemotronH_Nano_VL_V2"), + "NemotronH_Nano_Omni_Reasoning_V3": ("nano_nemotron_vl", "NemotronH_Nano_VL_V2"), + "NemotronH_Super_Omni_Reasoning_V3": ("nano_nemotron_vl", "NemotronH_Nano_VL_V2"), + "NVLM_D": ("nvlm_d", "NVLM_D_Model"), + "OpenCUAForConditionalGeneration": ("opencua", "OpenCUAForConditionalGeneration"), + "OpenPanguVLForConditionalGeneration": ( + "openpangu_vl", + "OpenPanguVLForConditionalGeneration", + ), + "OpenVLAForActionPrediction": ("openvla", "OpenVLAForActionPrediction"), + "Ovis": ("ovis", "Ovis"), + "Ovis2_5": ("ovis2_5", "Ovis2_5"), + "Ovis2_6ForCausalLM": ("ovis2_5", "Ovis2_5"), + "Ovis2_6_MoeForCausalLM": ("ovis2_5", "Ovis2_5"), + "PaddleOCRVLForConditionalGeneration": ( + "paddleocr_vl", + "PaddleOCRVLForConditionalGeneration", + ), + "PaliGemmaForConditionalGeneration": ( + "paligemma", + "PaliGemmaForConditionalGeneration", + ), + "Phi3VForCausalLM": ("phi3v", "Phi3VForCausalLM"), + "Phi4ForCausalLMV": ("phi4siglip", "Phi4ForCausalLMV"), + "Phi4MMForCausalLM": ("phi4mm", "Phi4MMForCausalLM"), + "PixtralForConditionalGeneration": ("pixtral", "PixtralForConditionalGeneration"), + "QianfanOCRForConditionalGeneration": ( + "qianfan_ocr", + "QianfanOCRForConditionalGeneration", + ), + "QwenVLForConditionalGeneration": ("qwen_vl", "QwenVLForConditionalGeneration"), + "Qwen2VLForConditionalGeneration": ("qwen2_vl", "Qwen2VLForConditionalGeneration"), + "Qwen2_5_VLForConditionalGeneration": ( + "qwen2_5_vl", + "Qwen2_5_VLForConditionalGeneration", + ), + "Qwen2AudioForConditionalGeneration": ( + "qwen2_audio", + "Qwen2AudioForConditionalGeneration", + ), + "Qwen2_5OmniModel": ( + "qwen2_5_omni_thinker", + "Qwen2_5OmniThinkerForConditionalGeneration", + ), + "Qwen2_5OmniForConditionalGeneration": ( + "qwen2_5_omni_thinker", + "Qwen2_5OmniThinkerForConditionalGeneration", + ), + "Qwen3OmniMoeForConditionalGeneration": ( + "qwen3_omni_moe_thinker", + "Qwen3OmniMoeThinkerForConditionalGeneration", + ), + "Qwen3ASRForConditionalGeneration": ( + "qwen3_asr", + "Qwen3ASRForConditionalGeneration", + ), + "Qwen3ASRRealtimeGeneration": ("qwen3_asr_realtime", "Qwen3ASRRealtimeGeneration"), + "Qwen3VLForConditionalGeneration": ("qwen3_vl", "Qwen3VLForConditionalGeneration"), + "Qwen3VLMoeForConditionalGeneration": ( + "qwen3_vl_moe", + "Qwen3VLMoeForConditionalGeneration", + ), + "Qwen3_5ForConditionalGeneration": ("qwen3_5", "Qwen3_5ForConditionalGeneration"), + "Qwen3_5MoeForConditionalGeneration": ( + "qwen3_5", + "Qwen3_5MoeForConditionalGeneration", + ), + "RForConditionalGeneration": ("rvl", "RForConditionalGeneration"), + "SkyworkR1VChatModel": ("skyworkr1v", "SkyworkR1VChatModel"), + "SmolVLMForConditionalGeneration": ("smolvlm", "SmolVLMForConditionalGeneration"), + "StepVLForConditionalGeneration": ("step_vl", "StepVLForConditionalGeneration"), + "Step3VLForConditionalGeneration": ("step3_vl", "Step3VLForConditionalGeneration"), + "TarsierForConditionalGeneration": ("tarsier", "TarsierForConditionalGeneration"), + "Tarsier2ForConditionalGeneration": ( + "qwen2_vl", + "Tarsier2ForConditionalGeneration", + ), + "UltravoxModel": ("ultravox", "UltravoxModel"), + "VoxtralForConditionalGeneration": ("voxtral", "VoxtralForConditionalGeneration"), + "VoxtralRealtimeGeneration": ("voxtral_realtime", "VoxtralRealtimeGeneration"), + # [Encoder-decoder] + "CohereAsrForConditionalGeneration": ( + "cohere_asr", + "CohereAsrForConditionalGeneration", + ), + "NemotronParseForConditionalGeneration": ( + "nemotron_parse", + "NemotronParseForConditionalGeneration", + ), + "WhisperForConditionalGeneration": ("whisper", "WhisperForConditionalGeneration"), +} + +_SPECULATIVE_DECODING_MODELS = { + "ExtractHiddenStatesModel": ("extract_hidden_states", "ExtractHiddenStatesModel"), + "MiMoMTPModel": ("mimo_mtp", "MiMoMTP"), + "MiMoV2MTPModel": ("mimo_v2_mtp", "MiMoV2MTP"), + "MiMoV2OmniMTPModel": ("mimo_v2_mtp", "MiMoV2OmniMTP"), + "EagleCohereForCausalLM": ("cohere_eagle", "EagleCohereForCausalLM"), + "EagleLlamaForCausalLM": ("llama_eagle", "EagleLlamaForCausalLM"), + "EagleLlama4ForCausalLM": ("llama4_eagle", "EagleLlama4ForCausalLM"), + "EagleMiniCPMForCausalLM": ("minicpm_eagle", "EagleMiniCPMForCausalLM"), + "DFlashDraftModel": ("qwen3_dflash", "DFlashQwen3ForCausalLM"), + "DeepSeekV4DSparkModel": ("vllm.models.deepseek_v4", "DeepSeekV4DSpark"), + "PEagleDraftModel": ("llama_eagle3", "Eagle3LlamaForCausalLM"), + "PeagleLlamaForCausalLM": ("llama_eagle3", "Eagle3LlamaForCausalLM"), + "Eagle3LlamaForCausalLM": ("llama_eagle3", "Eagle3LlamaForCausalLM"), + "Eagle3MiniMaxM2ForCausalLM": ("llama_eagle3", "Eagle3LlamaForCausalLM"), + "LlamaForCausalLMEagle3": ("llama_eagle3", "Eagle3LlamaForCausalLM"), + "Eagle3Qwen2_5vlForCausalLM": ("llama_eagle3", "Eagle3LlamaForCausalLM"), + "Eagle3Qwen3vlForCausalLM": ("llama_eagle3", "Eagle3LlamaForCausalLM"), + "EagleMistralForCausalLM": ("mistral_eagle", "EagleMistralForCausalLM"), + "EagleMistralLarge3ForCausalLM": ( + "mistral_large_3_eagle", + "EagleMistralLarge3ForCausalLM", + ), + "Eagle3DeepseekV2ForCausalLM": ("deepseek_eagle3", "Eagle3DeepseekV2ForCausalLM"), + "Eagle3DeepseekV3ForCausalLM": ("deepseek_eagle3", "Eagle3DeepseekV2ForCausalLM"), + "EagleDeepSeekMTPModel": ("deepseek_eagle", "EagleDeepseekV3ForCausalLM"), + "DeepSeekMTPModel": ("deepseek_mtp", "DeepSeekMTP"), + "DeepSeekV4MTPModel": ("vllm.models.deepseek_v4", "DeepSeekV4MTP"), + "Gemma4MTPModel": ("gemma4_mtp", "Gemma4MTP"), + "ErnieMTPModel": ("ernie_mtp", "ErnieMTP"), + "ExaoneMoeMTP": ("exaone_moe_mtp", "ExaoneMoeMTP"), + "Exaone4_5_MTP": ("exaone4_5_mtp", "Exaone4_5_MTP"), + "NemotronHMTPModel": ("nemotron_h_mtp", "NemotronHMTP"), + "LongCatFlashMTPModel": ("longcat_flash_mtp", "LongCatFlashMTP"), + "Glm4MoeMTPModel": ("glm4_moe_mtp", "Glm4MoeMTP"), + "Glm4MoeLiteMTPModel": ("glm4_moe_lite_mtp", "Glm4MoeLiteMTP"), + "GlmOcrMTPModel": ("glm_ocr_mtp", "GlmOcrMTP"), + "MedusaModel": ("medusa", "Medusa"), + "OpenPanguMTPModel": ("openpangu_mtp", "OpenPanguMTP"), + "Qwen3NextMTP": ("qwen3_next_mtp", "Qwen3NextMTP"), + "Step3p5MTP": ("step3p5_mtp", "Step3p5MTP"), + "Qwen3_5MTP": ("qwen3_5_mtp", "Qwen3_5MTP"), + "Qwen3_5MoeMTP": ("qwen3_5_mtp", "Qwen3_5MoeMTP"), + "HYV3MTPModel": ("hy_v3_mtp", "HYV3MTP"), + # Temporarily disabled. + # # TODO(woosuk): Re-enable this once the MLP Speculator is supported in V1. + # "MLPSpeculatorPreTrainedModel": ("mlp_speculator", "MLPSpeculator"), +} + +_TRANSFORMERS_SUPPORTED_MODELS = { + # Text generation models + "SmolLM3ForCausalLM": ("transformers", "TransformersForCausalLM"), + # Multimodal models + "Emu3ForConditionalGeneration": ( + "transformers", + "TransformersMultiModalForCausalLM", + ), +} + +_TRANSFORMERS_BACKEND_MODELS = { + # Text generation models + "TransformersForCausalLM": ("transformers", "TransformersForCausalLM"), + "TransformersMoEForCausalLM": ("transformers", "TransformersMoEForCausalLM"), + # Multimodal models + "TransformersMultiModalForCausalLM": ( + "transformers", + "TransformersMultiModalForCausalLM", + ), + "TransformersMultiModalMoEForCausalLM": ( + "transformers", + "TransformersMultiModalMoEForCausalLM", + ), + # Embedding models + "TransformersEmbeddingModel": ("transformers", "TransformersEmbeddingModel"), + "TransformersMoEEmbeddingModel": ("transformers", "TransformersMoEEmbeddingModel"), + "TransformersMultiModalEmbeddingModel": ( + "transformers", + "TransformersMultiModalEmbeddingModel", + ), + # Sequence classification models + "TransformersForSequenceClassification": ( + "transformers", + "TransformersForSequenceClassification", + ), + "TransformersMoEForSequenceClassification": ( + "transformers", + "TransformersMoEForSequenceClassification", + ), + "TransformersMultiModalForSequenceClassification": ( + "transformers", + "TransformersMultiModalForSequenceClassification", + ), +} + +_VLLM_MODELS = { + **_TEXT_GENERATION_MODELS, + **_EMBEDDING_MODELS, + **_LATE_INTERACTION_MODELS, + **_REWARD_MODELS, + **_TOKEN_CLASSIFICATION_MODELS, + **_SEQUENCE_CLASSIFICATION_MODELS, + **_MULTIMODAL_MODELS, + **_SPECULATIVE_DECODING_MODELS, + **_TRANSFORMERS_SUPPORTED_MODELS, + **_TRANSFORMERS_BACKEND_MODELS, +} + +# This variable is used as the args for subprocess.run(). We +# can modify this variable to alter the args if needed. e.g. +# when we use par format to pack things together, sys.executable +# might not be the target we want to run. +_SUBPROCESS_COMMAND = [sys.executable, "-m", "vllm.model_executor.models.registry"] + +_PREVIOUSLY_SUPPORTED_MODELS = { + "MotifForCausalLM": "0.10.2", + "Phi3SmallForCausalLM": "0.9.2", + "Phi4FlashForCausalLM": "0.10.2", + "Phi4MultimodalForCausalLM": "0.12.0", + # encoder-decoder models except whisper + # have been removed for V0 deprecation. + "DonutForConditionalGeneration": "0.10.2", + "MllamaForConditionalGeneration": "0.10.2", +} + +_OOT_SUPPORTED_MODELS = { + "BartModel": "https://github.com/vllm-project/bart-plugin", + "BartForConditionalGeneration": "https://github.com/vllm-project/bart-plugin", + "Florence2ForConditionalGeneration": "https://github.com/vllm-project/bart-plugin", + "MBartForConditionalGeneration": "https://github.com/vllm-project/bart-plugin", +} + + +@dataclass(frozen=True) +class _ModelInfo: + architecture: str + is_text_generation_model: bool + is_pooling_model: bool + attn_type: AttnTypeStr + default_seq_pooling_type: SequencePoolingType + default_tok_pooling_type: TokenPoolingType + score_type: ScoreType + supports_multimodal: bool + supports_multimodal_raw_input_only: bool + requires_raw_input_tokens: bool + supports_multimodal_encoder_tp_data: bool + supports_pp: bool + has_inner_state: bool + is_attention_free: bool + is_hybrid: bool + has_noops: bool + supports_mamba_prefix_caching: bool + supports_transcription: bool + supports_transcription_only: bool + + @staticmethod + def from_model_cls(model: type[nn.Module]) -> "_ModelInfo": + return _ModelInfo( + architecture=model.__name__, + is_text_generation_model=is_text_generation_model(model), + is_pooling_model=is_pooling_model(model), + default_seq_pooling_type=get_default_seq_pooling_type(model), + default_tok_pooling_type=get_default_tok_pooling_type(model), + attn_type=get_attn_type(model), + score_type=get_score_type(model), + supports_multimodal=supports_multimodal(model), + supports_multimodal_raw_input_only=supports_multimodal_raw_input_only( + model + ), + requires_raw_input_tokens=requires_raw_input_tokens(model), + supports_multimodal_encoder_tp_data=supports_multimodal_encoder_tp_data( + model + ), + supports_pp=supports_pp(model), + 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_transcription=supports_transcription(model), + supports_transcription_only=( + supports_transcription(model) and model.supports_transcription_only + ), + has_noops=has_noops(model), + ) + + +class _BaseRegisteredModel(ABC): + @abstractmethod + def inspect_model_cls(self) -> _ModelInfo: + raise NotImplementedError + + @abstractmethod + def load_model_cls(self) -> type[nn.Module]: + raise NotImplementedError + + +@dataclass(frozen=True) +class _RegisteredModel(_BaseRegisteredModel): + """ + Represents a model that has already been imported in the main process. + """ + + interfaces: _ModelInfo + model_cls: type[nn.Module] + + @staticmethod + def from_model_cls(model_cls: type[nn.Module]): + return _RegisteredModel( + interfaces=_ModelInfo.from_model_cls(model_cls), + model_cls=model_cls, + ) + + def inspect_model_cls(self) -> _ModelInfo: + return self.interfaces + + def load_model_cls(self) -> type[nn.Module]: + return self.model_cls + + +@dataclass(frozen=True) +class _LazyRegisteredModel(_BaseRegisteredModel): + """ + Represents a model that has not been imported in the main process. + """ + + module_name: str + class_name: str + + @staticmethod + def _get_cache_dir() -> Path: + return Path(envs.VLLM_CACHE_ROOT) / "modelinfos" + + def _get_cache_filename(self) -> str: + cls_name = f"{self.module_name}-{self.class_name}".replace(".", "-") + return f"{cls_name}.json" + + def _load_modelinfo_from_cache(self, module_hash: str) -> _ModelInfo | None: + try: + try: + modelinfo_path = self._get_cache_dir() / self._get_cache_filename() + with open(modelinfo_path, encoding="utf-8") as file: + mi_dict = json.load(file) + except FileNotFoundError: + logger.debug( + "Cached model info file for class %s.%s not found", + self.module_name, + self.class_name, + ) + return None + + if mi_dict["hash"] != module_hash: + logger.debug( + "Cached model info file for class %s.%s is stale", + self.module_name, + self.class_name, + ) + return None + + # file not changed, use cached _ModelInfo properties + return _ModelInfo(**mi_dict["modelinfo"]) + except Exception: + logger.debug( + "Cached model info for class %s.%s error. ", + self.module_name, + self.class_name, + ) + return None + + def _save_modelinfo_to_cache(self, mi: _ModelInfo, module_hash: str) -> None: + """save dictionary json file to cache""" + from vllm.model_executor.model_loader.weight_utils import atomic_writer + + try: + modelinfo_dict = { + "hash": module_hash, + "modelinfo": asdict(mi), + } + cache_dir = self._get_cache_dir() + cache_dir.mkdir(parents=True, exist_ok=True) + modelinfo_path = cache_dir / self._get_cache_filename() + with atomic_writer(modelinfo_path, encoding="utf-8") as f: + json.dump(modelinfo_dict, f, indent=2) + except Exception: + logger.exception("Error saving model info cache.") + + @logtime(logger=logger, msg="Registry inspect model class") + def inspect_model_cls(self) -> _ModelInfo: + # Modules registered with a non-default location (e.g. the + # hardware-isolated ``vllm.models.`` layout) live outside + # ``vllm/model_executor/models``. Resolve the module spec directly + # so the file-hash cache stays warm for them. + if self.module_name.startswith("vllm.model_executor.models."): + model_path = Path(__file__).parent / f"{self.module_name.split('.')[-1]}.py" + else: + try: + spec = importlib.util.find_spec(self.module_name) + except (ImportError, ValueError): + spec = None + model_path = Path(spec.origin) if spec is not None and spec.origin else None + module_hash = None + + if model_path is not None and model_path.exists(): + with open(model_path, "rb") as f: + module_hash = safe_hash(f.read(), usedforsecurity=False).hexdigest() + + mi = self._load_modelinfo_from_cache(module_hash) + if mi is not None: + logger.debug( + "Loaded model info for class %s.%s from cache", + self.module_name, + self.class_name, + ) + return mi + else: + logger.debug( + "Cache model info for class %s.%s miss. Loading model instead.", + self.module_name, + self.class_name, + ) + + # Performed in another process to avoid initializing CUDA + mi = _run_in_subprocess( + lambda: _ModelInfo.from_model_cls(self.load_model_cls()) + ) + logger.debug( + "Loaded model info for class %s.%s", self.module_name, self.class_name + ) + + # save cache file + if module_hash is not None: + self._save_modelinfo_to_cache(mi, module_hash) + + return mi + + def load_model_cls(self) -> type[nn.Module]: + mod = importlib.import_module(self.module_name) + return getattr(mod, self.class_name) + + +@lru_cache(maxsize=128) +def _try_load_model_cls( + model_arch: str, + model: _BaseRegisteredModel, +) -> type[nn.Module] | None: + from vllm.platforms import current_platform + + current_platform.verify_model_arch(model_arch) + try: + return model.load_model_cls() + except Exception: + logger.exception("Error in loading model architecture '%s'", model_arch) + return None + + +@lru_cache(maxsize=128) +def _try_inspect_model_cls( + model_arch: str, + model: _BaseRegisteredModel, +) -> _ModelInfo | None: + try: + return model.inspect_model_cls() + except Exception: + logger.exception("Error in inspecting model architecture '%s'", model_arch) + return None + + +@dataclass +class _ModelRegistry: + # Keyed by model_arch + models: dict[str, _BaseRegisteredModel] = field(default_factory=dict) + + def get_supported_archs(self) -> Set[str]: + return self.models.keys() + + def register_model( + self, + model_arch: str, + model_cls: type[nn.Module] | str, + ) -> None: + """ + Register an external model to be used in vLLM. + + `model_cls` can be either: + + - A [`torch.nn.Module`][] class directly referencing the model. + - A string in the format `:` which can be used to + lazily import the model. This is useful to avoid initializing CUDA + when importing the model and thus the related error + `RuntimeError: Cannot re-initialize CUDA in forked subprocess`. + """ + if not isinstance(model_arch, str): + msg = f"`model_arch` should be a string, not a {type(model_arch)}" + raise TypeError(msg) + + if model_arch in self.models: + logger.debug( + "Model architecture %s is already registered, and will be " + "overwritten by the new model class %s.", + model_arch, + model_cls, + ) + + if isinstance(model_cls, str): + split_str = model_cls.split(":") + if len(split_str) != 2: + msg = "Expected a string in the format `:`" + raise ValueError(msg) + + model = _LazyRegisteredModel(*split_str) + elif isinstance(model_cls, type) and issubclass(model_cls, nn.Module): + model = _RegisteredModel.from_model_cls(model_cls) + else: + msg = ( + "`model_cls` should be a string or PyTorch model class, " + f"not a {type(model_arch)}" + ) + raise TypeError(msg) + + self.models[model_arch] = model + + def _raise_for_unsupported(self, architectures: list[str]): + all_supported_archs = self.get_supported_archs() + + if any(arch in all_supported_archs for arch in architectures): + raise ValueError( + f"Model architectures {architectures} failed " + "to be inspected. Please check the logs for more details." + ) + + for arch in architectures: + if arch in _PREVIOUSLY_SUPPORTED_MODELS: + previous_version = _PREVIOUSLY_SUPPORTED_MODELS[arch] + + raise ValueError( + f"Model architecture {arch} was supported in vLLM until " + f"v{previous_version}, and is not supported anymore. " + "Please use an older version of vLLM if you want to " + "use this model architecture." + ) + if arch in _OOT_SUPPORTED_MODELS: + plugin_url = _OOT_SUPPORTED_MODELS[arch] + + raise ValueError( + f"Model architecture {arch} is not supported in-tree anymore. " + f"Please install the plugin at {plugin_url} if you want to " + "use this model architecture." + ) + + raise ValueError( + f"Model architectures {architectures} are not supported for now. " + f"Supported architectures: {all_supported_archs}" + ) + + def _try_load_model_cls(self, model_arch: str) -> type[nn.Module] | None: + if model_arch not in self.models: + return None + + return _try_load_model_cls(model_arch, self.models[model_arch]) + + def _try_inspect_model_cls(self, model_arch: str) -> _ModelInfo | None: + if model_arch not in self.models: + return None + + return _try_inspect_model_cls(model_arch, self.models[model_arch]) + + def _try_resolve_transformers( + self, + architecture: str, + model_config: ModelConfig, + ) -> str | None: + if architecture in _TRANSFORMERS_BACKEND_MODELS: + return architecture + + auto_map: dict[str, str] = ( + getattr(model_config.hf_config, "auto_map", None) or dict() + ) + + # Make sure that config class is always initialized before model class, + # otherwise the model class won't be able to access the config class, + # the expected auto_map should have correct order like: + # "auto_map": { + # "AutoConfig": "--", + # "AutoModel": "--", + # "AutoModelFor": "--", + # }, + for prefix in ("AutoConfig", "AutoModel"): + for name, module in auto_map.items(): + if name.startswith(prefix): + try_get_class_from_dynamic_module( + module, + model_config.model, + revision=model_config.revision, + code_revision=model_config.code_revision, + trust_remote_code=model_config.trust_remote_code, + warn_on_fail=False, + ) + + model_module = getattr(transformers, architecture, None) + + if model_module is None: + for name, module in auto_map.items(): + if name.startswith("AutoModel"): + model_module = try_get_class_from_dynamic_module( + module, + model_config.model, + revision=model_config.revision, + code_revision=model_config.code_revision, + trust_remote_code=model_config.trust_remote_code, + warn_on_fail=True, + ) + if model_module is not None: + break + else: + if model_config.model_impl != "transformers": + return None + + raise ValueError( + f"Cannot find model module. {architecture!r} is not a " + "registered model in the Transformers library (only " + "relevant if the model is meant to be in Transformers) " + "and 'AutoModel' is not present in the model config's " + "'auto_map' (relevant if the model is custom)." + ) + + if not model_module.is_backend_compatible(): + if model_config.model_impl != "transformers": + return None + + raise ValueError( + f"The Transformers implementation of {architecture!r} " + "is not compatible with vLLM." + ) + + return model_config._get_transformers_backend_cls() + + def _normalize_arch( + self, + architecture: str, + model_config: ModelConfig, + ) -> str: + if architecture in self.models: + return architecture + + # This may be called in order to resolve runner_type and convert_type + # in the first place, in which case we consider the default match + match = try_match_architecture_defaults( + architecture, + runner_type=getattr(model_config, "runner_type", None), + convert_type=getattr(model_config, "convert_type", None), + ) + if match: + suffix, _ = match + + # Get the name of the base model to convert + for repl_suffix, _ in iter_architecture_defaults(): + base_arch = architecture.replace(suffix, repl_suffix) + if base_arch in self.models: + return base_arch + + return architecture + + def inspect_model_cls( + self, + architectures: str | list[str], + model_config: ModelConfig, + ) -> tuple[_ModelInfo, str]: + if isinstance(architectures, str): + architectures = [architectures] + if not architectures: + raise ValueError("No model architectures are specified") + + # Require transformers impl + if model_config.model_impl == "transformers": + arch = self._try_resolve_transformers(architectures[0], model_config) + if arch is not None: + model_info = self._try_inspect_model_cls(arch) + if model_info is not None: + return (model_info, arch) + elif model_config.model_impl == "terratorch": + model_info = self._try_inspect_model_cls("Terratorch") + return (model_info, "Terratorch") + + # Fallback to transformers impl (after resolving convert_type) + if ( + all(arch not in self.models for arch in architectures) + and model_config.model_impl == "auto" + and getattr(model_config, "convert_type", "none") == "none" + ): + arch = self._try_resolve_transformers(architectures[0], model_config) + if arch is not None: + model_info = self._try_inspect_model_cls(arch) + if model_info is not None: + return (model_info, arch) + + for arch in architectures: + normalized_arch = self._normalize_arch(arch, model_config) + model_info = self._try_inspect_model_cls(normalized_arch) + if model_info is not None: + return (model_info, arch) + + # Fallback to transformers impl (before resolving runner_type) + if ( + all(arch not in self.models for arch in architectures) + and model_config.model_impl == "auto" + ): + arch = self._try_resolve_transformers(architectures[0], model_config) + if arch is not None: + model_info = self._try_inspect_model_cls(arch) + if model_info is not None: + return (model_info, arch) + + return self._raise_for_unsupported(architectures) + + def resolve_model_cls( + self, + architectures: str | list[str], + model_config: ModelConfig, + ) -> tuple[type[nn.Module], str]: + if isinstance(architectures, str): + architectures = [architectures] + if not architectures: + raise ValueError("No model architectures are specified") + + # Require transformers impl + if model_config.model_impl == "transformers": + arch = self._try_resolve_transformers(architectures[0], model_config) + if arch is not None: + model_cls = self._try_load_model_cls(arch) + if model_cls is not None: + return (model_cls, arch) + elif model_config.model_impl == "terratorch": + arch = "Terratorch" + model_cls = self._try_load_model_cls(arch) + if model_cls is not None: + return (model_cls, arch) + + # Fallback to transformers impl (after resolving convert_type) + if ( + all(arch not in self.models for arch in architectures) + and model_config.model_impl == "auto" + and getattr(model_config, "convert_type", "none") == "none" + ): + arch = self._try_resolve_transformers(architectures[0], model_config) + if arch is not None: + model_cls = self._try_load_model_cls(arch) + if model_cls is not None: + return (model_cls, arch) + + for arch in architectures: + normalized_arch = self._normalize_arch(arch, model_config) + model_cls = self._try_load_model_cls(normalized_arch) + if model_cls is not None: + return (model_cls, arch) + + # Fallback to transformers impl (before resolving runner_type) + if ( + all(arch not in self.models for arch in architectures) + and model_config.model_impl == "auto" + ): + arch = self._try_resolve_transformers(architectures[0], model_config) + if arch is not None: + model_cls = self._try_load_model_cls(arch) + if model_cls is not None: + return (model_cls, arch) + + return self._raise_for_unsupported(architectures) + + def is_text_generation_model( + self, + architectures: str | list[str], + model_config: ModelConfig, + ) -> bool: + model_cls, _ = self.inspect_model_cls(architectures, model_config) + return model_cls.is_text_generation_model + + def is_pooling_model( + self, + architectures: str | list[str], + model_config: ModelConfig, + ) -> bool: + model_cls, _ = self.inspect_model_cls(architectures, model_config) + return model_cls.is_pooling_model + + def is_multimodal_model( + self, + architectures: str | list[str], + model_config: ModelConfig, + ) -> bool: + model_cls, _ = self.inspect_model_cls(architectures, model_config) + return model_cls.supports_multimodal + + def is_multimodal_raw_input_only_model( + self, + architectures: str | list[str], + model_config: ModelConfig, + ) -> bool: + model_cls, _ = self.inspect_model_cls(architectures, model_config) + return model_cls.supports_multimodal_raw_input_only + + def is_pp_supported_model( + self, + architectures: str | list[str], + model_config: ModelConfig, + ) -> bool: + model_cls, _ = self.inspect_model_cls(architectures, model_config) + return model_cls.supports_pp + + def model_has_inner_state( + self, + architectures: str | list[str], + model_config: ModelConfig, + ) -> bool: + model_cls, _ = self.inspect_model_cls(architectures, model_config) + return model_cls.has_inner_state + + def is_attention_free_model( + self, + architectures: str | list[str], + model_config: ModelConfig, + ) -> bool: + model_cls, _ = self.inspect_model_cls(architectures, model_config) + return model_cls.is_attention_free + + def is_hybrid_model( + self, + architectures: str | list[str], + model_config: ModelConfig, + ) -> bool: + model_cls, _ = self.inspect_model_cls(architectures, model_config) + return model_cls.is_hybrid + + def is_noops_model( + self, + architectures: str | list[str], + model_config: ModelConfig, + ) -> bool: + model_cls, _ = self.inspect_model_cls(architectures, model_config) + return model_cls.has_noops + + def is_transcription_model( + self, + architectures: str | list[str], + model_config: ModelConfig, + ) -> bool: + model_cls, _ = self.inspect_model_cls(architectures, model_config) + return model_cls.supports_transcription + + def is_transcription_only_model( + self, + architectures: str | list[str], + model_config: ModelConfig, + ) -> bool: + model_cls, _ = self.inspect_model_cls(architectures, model_config) + return model_cls.supports_transcription_only + + +def _resolve_module_name(mod_relname: str) -> str: + # Allow registry entries to point at fully-qualified module paths (e.g. + # ``vllm.models.deepseek_v4``) for models that live outside the legacy + # ``vllm.model_executor.models`` flat layout. + if mod_relname.startswith("vllm."): + return mod_relname + return f"vllm.model_executor.models.{mod_relname}" + + +ModelRegistry = _ModelRegistry( + { + model_arch: _LazyRegisteredModel( + module_name=_resolve_module_name(mod_relname), + class_name=cls_name, + ) + for model_arch, (mod_relname, cls_name) in _VLLM_MODELS.items() + } +) + +_T = TypeVar("_T") + + +def _run_in_subprocess(fn: Callable[[], _T]) -> _T: + # NOTE: We use a temporary directory instead of a temporary file to avoid + # issues like https://stackoverflow.com/questions/23212435/permission-denied-to-write-to-my-temporary-file + with tempfile.TemporaryDirectory() as tempdir: + output_filepath = os.path.join(tempdir, "registry_output.tmp") + + # `cloudpickle` allows pickling lambda functions directly + import cloudpickle + + input_bytes = cloudpickle.dumps((fn, output_filepath)) + + # cannot use `sys.executable __file__` here because the script + # contains relative imports + returned = subprocess.run( + _SUBPROCESS_COMMAND, input=input_bytes, capture_output=True + ) + + # check if the subprocess is successful + try: + returned.check_returncode() + except Exception as e: + # wrap raised exception to provide more information + raise RuntimeError( + f"Error raised in subprocess:\n{returned.stderr.decode()}" + ) from e + + with open(output_filepath, "rb") as f: + return pickle.load(f) + + +def _run() -> None: + # Setup plugins + from vllm.plugins import load_general_plugins + + load_general_plugins() + + fn, output_file = pickle.loads(sys.stdin.buffer.read()) + + result = fn() + + with open(output_file, "wb") as f: + f.write(pickle.dumps(result)) + + +if __name__ == "__main__": + _run() diff --git a/mods/deepseek-v4-flash-dspark/overlay/vllm/model_executor/warmup/kernel_warmup.py b/mods/deepseek-v4-flash-dspark/overlay/vllm/model_executor/warmup/kernel_warmup.py new file mode 100644 index 00000000..acbb1620 --- /dev/null +++ b/mods/deepseek-v4-flash-dspark/overlay/vllm/model_executor/warmup/kernel_warmup.py @@ -0,0 +1,851 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +""" +Warmup kernels used during model execution. +This is useful specifically for JIT'ed kernels as we don't want JIT'ing to +happen during model execution. +""" + +import hashlib +from pathlib import Path +from typing import TYPE_CHECKING + +import torch + +import vllm.envs as envs +from vllm.compilation.caching import aot_compile_hash_factors +from vllm.logger import init_logger +from vllm.model_executor.warmup.deep_gemm_warmup import deep_gemm_warmup +from vllm.model_executor.warmup.deepseek_v4_mhc_warmup import ( + deepseek_v4_mhc_warmup, +) +from vllm.platforms import current_platform +from vllm.utils.deep_gemm import is_deep_gemm_supported +from vllm.utils.flashinfer import has_flashinfer + +if TYPE_CHECKING: + from vllm.v1.worker.gpu_model_runner import GPUModelRunner + from vllm.v1.worker.gpu_worker import Worker + +logger = init_logger(__name__) + +_DEEPSEEK_V4_SPARSE_MLA_BACKENDS = frozenset( + { + "FLASHMLA_SPARSE", + "DEEPSEEK_SPARSE_SWA", + } +) + +_DEEPSEEK_V4_SPARSE_MLA_MIXED_WARMUP_TOKENS = 16 +_DEEPSEEK_V4_SPARSE_MLA_PREFILL_WARMUP_TOKENS = 8192 +_DEEPSEEK_V4_DSPARK_DECODE_AUTOTUNE_SEQ_LENS = (512, 2048) +_DEEPSEEK_V4_DSPARK_SHORT_PREFILL_WARMUP_TOKENS = (12, 16, 20, 32) +_DEEPSEEK_V4_DSPARK_ROUTE_PACK_PREFILL_TOKENS = ( + # The single-stream coding benchmark targets a 512-token prompt, but the + # chat template makes the served prompt land near 415 tokens. B12X route-pack + # kernels specialize on exact packed-route workspace sizes, not only the + # next power-of-two capacity, so warm the observed prompt neighborhood too. + 411, + 415, + 416, + 512, + 513, + 1024, +) + +# Fan of num_tokens specializations to pre-JIT for +# `_compute_slot_mapping_kernel`. On SM12x cold JIT can emit +# non-deterministic codegen that writes wrong slot_mapping → KV corruption +# → downstream sparse-MLA IMA. +_DEEPSEEK_V4_SLOT_MAPPING_WARMUP_TOKENS = tuple(range(1, 17)) + ( + 32, + 64, + 128, + 256, + 512, +) + + +def _attention_backend_name(backend: object) -> str | None: + get_name = getattr(backend, "get_name", None) + if get_name is None: + return None + try: + return get_name() + except NotImplementedError: + return None + + +def _has_deepseek_v4_sparse_mla_backend(runner: "GPUModelRunner") -> bool: + for groups in getattr(runner, "attn_groups", []) or (): + for group in groups: + name = _attention_backend_name(getattr(group, "backend", None)) + if name in _DEEPSEEK_V4_SPARSE_MLA_BACKENDS: + return True + return False + + +def _clamp_warmup_tokens(num_tokens: int, max_tokens: int) -> int: + return max(0, min(num_tokens, max_tokens)) + + +def _runner_max_num_tokens(runner: "GPUModelRunner") -> int: + max_num_tokens = getattr(runner, "max_num_tokens", None) + if max_num_tokens is not None: + return int(max_num_tokens) + + scheduler_config = getattr(runner, "scheduler_config", None) + max_num_batched_tokens = getattr(scheduler_config, "max_num_batched_tokens", 1) + return int(max_num_batched_tokens) + + +def _runner_vocab_size(runner: "GPUModelRunner") -> int: + vocab_size = getattr(runner, "vocab_size", None) + if vocab_size is not None: + return int(vocab_size) + + model_config = getattr(runner, "model_config", None) + get_vocab_size = getattr(model_config, "get_vocab_size", None) + if get_vocab_size is not None: + return int(get_vocab_size()) + + input_batch = getattr(runner, "input_batch", None) + vocab_size = getattr(input_batch, "vocab_size", None) + if vocab_size is not None: + return int(vocab_size) + + return 1 + + +def _deepseek_v4_hf_config(worker: "Worker") -> object | None: + model_config = getattr(worker.model_runner, "model_config", None) + return getattr(model_config, "hf_text_config", None) or getattr( + model_config, "hf_config", None + ) + + +def _dspark_spec_decode_query_len(worker: "Worker") -> int | None: + spec_config = getattr(worker.vllm_config, "speculative_config", None) + if spec_config is None: + return None + is_dspark = getattr(spec_config, "is_dspark", None) + if is_dspark is not None: + if not is_dspark(): + return None + elif getattr(spec_config, "method", None) != "dspark": + return None + + num_spec_tokens = getattr(spec_config, "num_speculative_tokens", None) + if num_spec_tokens is None: + return None + query_len = int(num_spec_tokens) + 1 + if query_len <= 1: + return None + return query_len + + +def _dspark_uniform_decode_autotune_kwargs( + worker: "Worker", +) -> list[dict[str, object]]: + query_len = _dspark_spec_decode_query_len(worker) + if query_len is None: + return [] + max_tokens = _runner_max_num_tokens(worker.model_runner) + if query_len > max_tokens: + return [] + + max_model_len = int(getattr(worker.model_runner, "max_model_len", 0) or 0) + seq_lens = [ + seq_len + for seq_len in _DEEPSEEK_V4_DSPARK_DECODE_AUTOTUNE_SEQ_LENS + if max_model_len <= 0 or seq_len <= max_model_len + ] + if not seq_lens: + seq_lens = [query_len] + + return [ + dict( + num_tokens=query_len, + skip_eplb=True, + is_profile=True, + force_attention=True, + uniform_decode=True, + profile_seq_lens=seq_len, + ) + for seq_len in seq_lens + ] + + +def _dspark_route_pack_token_counts(worker: "Worker") -> tuple[int, ...]: + query_len = _dspark_spec_decode_query_len(worker) + if query_len is None: + return () + + max_tokens = _runner_max_num_tokens(worker.model_runner) + token_counts = [query_len] + token_counts.extend(_DEEPSEEK_V4_DSPARK_SHORT_PREFILL_WARMUP_TOKENS) + token_counts.extend(_DEEPSEEK_V4_DSPARK_ROUTE_PACK_PREFILL_TOKENS) + return tuple( + sorted( + {token_count for token_count in token_counts if token_count <= max_tokens} + ) + ) + + +def _dspark_warmup_request_counts(worker: "Worker") -> tuple[int, ...]: + max_num_seqs = max(1, int(worker.scheduler_config.max_num_seqs)) + return tuple(sorted({1, min(max_num_seqs, 4)})) + + +@torch.inference_mode() +def _deepseek_v4_b12x_route_pack_warmup(worker: "Worker") -> None: + """Pre-JIT B12X/FlashInfer W4A16 MoE route-packing kernels. + + DSpark's first live request is usually a short prompt plus speculative + decode width 6. Those shapes are too small to be covered reliably by the + broader DeepGEMM warmup, but they still hit Triton route-pack prefix kernels. + """ + token_counts = _dspark_route_pack_token_counts(worker) + if not token_counts: + return + + hf_config = _deepseek_v4_hf_config(worker) + num_experts = int(getattr(hf_config, "n_routed_experts", 0) or 0) + top_k = int(getattr(hf_config, "num_experts_per_tok", 0) or 0) + if num_experts <= 0 or top_k <= 0: + return + + try: + from b12x.moe.fused.w4a16.host import ( + max_packed_route_slots, + select_route_block_size_m, + ) + from b12x.moe.fused.w4a16.kernel import pack_topk_routes_by_expert + except ImportError: + logger.debug("Skipping B12X route-pack warmup: package is unavailable.") + return + + device = worker.model_runner.device + for token_count in token_counts: + block_size_m = select_route_block_size_m(token_count, top_k, num_experts) + topk_ids = torch.zeros( + (token_count, top_k), dtype=torch.int32, device=device + ) + route_id_shapes = (topk_ids, topk_ids.view(-1)) + for route_ids in route_id_shapes: + pack_topk_routes_by_expert(route_ids, block_size_m, num_experts) + + routed_rows = int(route_ids.numel()) + route_slots = max( + 1, + max_packed_route_slots(routed_rows, block_size_m, num_experts), + ) + route_blocks = max(1, (route_slots + block_size_m - 1) // block_size_m) + pack_topk_routes_by_expert( + route_ids, + block_size_m, + num_experts, + packed_route_indices=torch.empty( + (route_slots,), dtype=torch.int32, device=device + ), + block_expert_ids=torch.empty( + (route_blocks,), dtype=torch.int32, device=device + ), + packed_route_count=torch.empty(1, dtype=torch.int32, device=device), + expert_offsets=torch.empty( + (num_experts + 1,), dtype=torch.int32, device=device + ), + ) + + +@torch.inference_mode() +def _deepseek_v4_spec_decode_padded_kernel_warmup(worker: "Worker") -> None: + """Pre-JIT padded speculative decode input-prep kernels. + + DSpark uses the padded speculative path with query length + `1 + num_speculative_tokens`; the first real request otherwise pays Triton + JIT for these small kernels. + """ + query_len = _dspark_spec_decode_query_len(worker) + if query_len is None: + return + + runner = worker.model_runner + device = runner.device + vocab_size = _runner_vocab_size(runner) + next_token_kernel, inputs_kernel = _spec_decode_padded_warmup_kernels() + + for num_reqs in _dspark_warmup_request_counts(worker): + discard_buffer = getattr( + getattr(runner, "discard_request_mask", None), "gpu", None + ) + if discard_buffer is not None and discard_buffer.numel() >= num_reqs: + discard_request_mask = discard_buffer[:num_reqs] + discard_request_mask.zero_() + else: + discard_request_mask = torch.zeros( + num_reqs, dtype=torch.bool, device=device + ) + + drafter = getattr(runner, "drafter", None) + backup_buffer = getattr( + getattr(drafter, "backup_next_token_ids", None), "gpu", None + ) + if backup_buffer is not None and backup_buffer.numel() >= num_reqs: + backup_tokens = backup_buffer[:num_reqs] + backup_tokens.zero_() + else: + backup_tokens = torch.zeros(num_reqs, dtype=torch.int32, device=device) + + valid_sampled_tokens_count = None + for sample_width in range(1, query_len + 1): + sampled_token_ids = torch.zeros( + (num_reqs, sample_width), dtype=torch.int32, device=device + ) + if sample_width > 1: + sampled_token_ids[:, -1] = -1 + + next_token_ids = torch.empty(num_reqs, dtype=torch.int32, device=device) + valid_sampled_tokens_count = torch.empty( + num_reqs, dtype=torch.int32, device=device + ) + block_size_tokens = 1 << (sample_width - 1).bit_length() + + next_token_kernel[(num_reqs,)]( + sampled_token_ids, + discard_request_mask, + backup_tokens, + next_token_ids, + valid_sampled_tokens_count, + vocab_size, + sample_width, + num_reqs, + sampled_token_ids.stride(0), + BLOCK_SIZE_TOKENS=block_size_tokens, + ) + + assert valid_sampled_tokens_count is not None + + cu_num_draft_tokens = torch.arange( + query_len - 1, + (query_len - 1) * num_reqs + 1, + query_len - 1, + dtype=torch.int32, + device=device, + ) + query_start_loc = torch.arange( + 0, + query_len * num_reqs + 1, + query_len, + dtype=torch.int32, + device=device, + ) + token_indices_to_sample = torch.empty( + num_reqs, dtype=torch.int32, device=device + ) + num_rejected_tokens = torch.empty(num_reqs, dtype=torch.int32, device=device) + + inputs_kernel[(num_reqs,)]( + cu_num_draft_tokens, + valid_sampled_tokens_count, + query_start_loc, + token_indices_to_sample, + num_rejected_tokens, + num_reqs, + ) + + +@torch.inference_mode() +def _deepseek_v4_rejection_sampler_warmup(worker: "Worker") -> None: + """Pre-JIT greedy rejection sampling for the DSpark draft width.""" + query_len = _dspark_spec_decode_query_len(worker) + if query_len is None: + return + + num_draft_tokens = query_len - 1 + if num_draft_tokens <= 0: + return + + from vllm.v1.sample.rejection_sampler import rejection_greedy_sample_kernel + + device = worker.model_runner.device + for batch_size in _dspark_warmup_request_counts(worker): + total_draft_tokens = batch_size * num_draft_tokens + output_token_ids = torch.empty( + (batch_size, query_len), dtype=torch.int32, device=device + ) + cu_num_draft_tokens = torch.arange( + num_draft_tokens, + total_draft_tokens + 1, + num_draft_tokens, + dtype=torch.int32, + device=device, + ) + draft_token_ids = torch.zeros( + total_draft_tokens, dtype=torch.int32, device=device + ) + target_argmax = torch.zeros( + total_draft_tokens, dtype=torch.int64, device=device + ) + bonus_token_ids = torch.zeros((batch_size, 1), dtype=torch.int32, device=device) + + rejection_greedy_sample_kernel[(batch_size,)]( + output_token_ids, + cu_num_draft_tokens, + draft_token_ids, + target_argmax, + bonus_token_ids, + None, + num_draft_tokens, + None, + None, + SYNTHETIC_MODE=False, + ) + + +def _spec_decode_padded_warmup_kernels(): + from vllm.v1.spec_decode.utils import ( + eagle_prepare_inputs_padded_kernel, + eagle_prepare_next_token_padded_kernel, + ) + + return eagle_prepare_next_token_padded_kernel, eagle_prepare_inputs_padded_kernel + + +def _deepseek_v4_slot_mapping_warmup(runner: "GPUModelRunner") -> None: + """Pre-JIT `_compute_slot_mapping_kernel` across decode-shaped sizes.""" + max_tokens = _runner_max_num_tokens(runner) + input_batch = getattr(runner, "input_batch", None) + legacy_block_table = getattr(input_batch, "block_table", None) + v2_block_tables = getattr(runner, "block_tables", None) + if legacy_block_table is None and v2_block_tables is None: + logger.debug("Skipping DeepSeek V4 slot-mapping warmup: no block tables.") + return + + # Snapshot the runner buffers we mutate so warmup doesn't leak state. + saved_query_start_loc_np = None + saved_query_start_loc_gpu = None + if hasattr(runner, "query_start_loc"): + saved_query_start_loc_np = runner.query_start_loc.np[:2].copy() + saved_query_start_loc_gpu = runner.query_start_loc.gpu[:2].clone() + + try: + for requested_tokens in _DEEPSEEK_V4_SLOT_MAPPING_WARMUP_TOKENS: + num_tokens = _clamp_warmup_tokens(requested_tokens, max_tokens) + if num_tokens <= 0: + continue + + positions_source = torch.arange( + num_tokens, dtype=torch.int64, device=runner.device + ) + if hasattr(runner, "query_start_loc"): + runner.query_start_loc.np[0] = 0 + runner.query_start_loc.np[1] = num_tokens + runner.query_start_loc.copy_to_gpu(2) + query_start_loc = runner.query_start_loc.gpu[:2] + else: + query_start_loc = torch.tensor( + [0, num_tokens], dtype=torch.int32, device=runner.device + ) + + if hasattr(runner, "positions"): + saved_positions = runner.positions[:num_tokens].clone() + runner.positions[:num_tokens].copy_(positions_source) + positions = runner.positions[:num_tokens] + else: + saved_positions = None + positions = positions_source + + try: + if legacy_block_table is not None: + legacy_block_table.commit_block_table(1) + legacy_block_table.compute_slot_mapping( + 1, query_start_loc, positions + ) + else: + idx_mapping = torch.zeros( + 1, dtype=torch.int32, device=runner.device + ) + assert v2_block_tables is not None + v2_block_tables.compute_slot_mappings( + idx_mapping, + query_start_loc, + positions, + num_tokens_padded=num_tokens, + ) + finally: + if saved_positions is not None: + runner.positions[:num_tokens].copy_(saved_positions) + finally: + if saved_query_start_loc_np is not None: + runner.query_start_loc.np[:2] = saved_query_start_loc_np + assert saved_query_start_loc_gpu is not None + runner.query_start_loc.gpu[:2].copy_(saved_query_start_loc_gpu) + + +@torch.inference_mode() +def _deepseek_v4_request_prep_warmup(worker: "Worker") -> None: + """Pre-JIT the slot-mapping kernel before the first real request.""" + if not envs.VLLM_ENABLE_DEEPSEEK_V4_SPARSE_MLA_WARMUP: + return + + runner = worker.model_runner + if runner.is_pooling_model or not _has_deepseek_v4_sparse_mla_backend(runner): + return + if not current_platform.is_cuda_alike(): + return + + logger.info("Warming up DeepSeek V4 request preparation kernels.") + _deepseek_v4_slot_mapping_warmup(runner) + _deepseek_v4_b12x_route_pack_warmup(worker) + _deepseek_v4_spec_decode_padded_kernel_warmup(worker) + _deepseek_v4_rejection_sampler_warmup(worker) + torch.accelerator.synchronize() + + +def _deepseek_v4_sparse_mla_decode_autotune( + worker: "Worker", + num_tokens: int, +) -> bool: + """Autotune FlashInfer's DSv4 SM120 sparse-MLA decode path. + + Returns True when this function consumed the mixed attention warmup shape. + """ + if worker.vllm_config.kernel_config.enable_flashinfer_autotune is not True: + return False + if not has_flashinfer() or not current_platform.is_device_capability_family(120): + return False + + try: + from flashinfer import sparse_mla_sm120_decode_dsv4_autotune + from flashinfer.autotuner import AutoTuner + except ImportError: + logger.warning( + "Skipping DeepSeek V4 sparse MLA decode autotune because this " + "FlashInfer build does not expose sparse_mla_sm120_decode_dsv4_autotune." + ) + return False + + from vllm.distributed.parallel_state import get_world_group + + runner = worker.model_runner + world = get_world_group() + is_leader = world.rank_in_group == 0 + cache_path = _resolve_flashinfer_autotune_file(runner) + + dummy_run_kwargs: list[dict[str, object]] = [ + dict( + num_tokens=num_tokens, + skip_eplb=True, + is_profile=True, + force_attention=True, + create_mixed_batch=True, + ) + ] + dummy_run_kwargs.extend(_dspark_uniform_decode_autotune_kwargs(worker)) + + if is_leader and len(dummy_run_kwargs) > 1: + logger.info( + "Including %d DSpark uniform-decode sparse MLA autotune shapes.", + len(dummy_run_kwargs) - 1, + ) + + def run_autotune_shapes() -> None: + for kwargs in dummy_run_kwargs: + runner._dummy_run(**kwargs) + + if is_leader: + logger.info( + "Autotuning DeepSeek V4 SM120 sparse MLA decode with FlashInfer " + "cache file: %s", + cache_path, + ) + + with torch.inference_mode(): + if is_leader: + with sparse_mla_sm120_decode_dsv4_autotune(cache_path=str(cache_path)): + run_autotune_shapes() + else: + run_autotune_shapes() + + tune_results: bytes | None = None + if is_leader and cache_path.exists(): + with open(cache_path, "rb") as f: + tune_results = f.read() + + tune_results = world.broadcast_object(tune_results, src=0) + if tune_results is None: + logger.warning( + "No DeepSeek V4 sparse MLA decode autotune cache entries found. " + "Falling back to FlashInfer's default tactic heuristic." + ) + world.barrier() + return True + + if not is_leader and world.local_rank == 0: + cache_path.parent.mkdir(parents=True, exist_ok=True) + with open(cache_path, "wb") as f: + f.write(tune_results) + world.barrier() + + AutoTuner.get().load_configs(str(cache_path)) + logger.info( + "DeepSeek V4 sparse MLA decode autotune cache loaded on rank %d from %s.", + world.rank_in_group, + cache_path, + ) + return True + + +def _deepseek_v4_sparse_mla_attention_warmup(worker: "Worker") -> None: + """Warm sparse-MLA attention shapes via `_dummy_run`. + + Three shapes: mixed prefill+decode, single max-chunk prefill, and a + second-chunk prefill (prior context) — the last covers + `_build_prefill_chunk_metadata_kernel`'s alt-shape specialization. + """ + if not envs.VLLM_ENABLE_DEEPSEEK_V4_SPARSE_MLA_WARMUP: + return + + runner = worker.model_runner + if runner.is_pooling_model or not _has_deepseek_v4_sparse_mla_backend(runner): + return + + max_tokens = worker.scheduler_config.max_num_batched_tokens + mixed_tokens = _clamp_warmup_tokens( + _DEEPSEEK_V4_SPARSE_MLA_MIXED_WARMUP_TOKENS, max_tokens + ) + prefill_tokens = _clamp_warmup_tokens( + _DEEPSEEK_V4_SPARSE_MLA_PREFILL_WARMUP_TOKENS, max_tokens + ) + if mixed_tokens <= 0 and prefill_tokens <= 0: + return + + logger.info( + "Warming up DeepSeek V4 sparse MLA attention " + "for mixed tokens=%s and prefill tokens=%s.", + mixed_tokens, + prefill_tokens, + ) + if mixed_tokens > 0: + mixed_warmup_done = _deepseek_v4_sparse_mla_decode_autotune( + worker, mixed_tokens + ) + if not mixed_warmup_done: + runner._dummy_run( + num_tokens=mixed_tokens, + skip_eplb=True, + is_profile=True, + force_attention=True, + create_mixed_batch=True, + ) + if prefill_tokens > 0: + for short_prefill_tokens in _DEEPSEEK_V4_DSPARK_SHORT_PREFILL_WARMUP_TOKENS: + short_prefill_tokens = _clamp_warmup_tokens( + short_prefill_tokens, max_tokens + ) + if short_prefill_tokens <= 0: + continue + runner._dummy_run( + num_tokens=short_prefill_tokens, + skip_eplb=True, + is_profile=True, + force_attention=True, + create_single_prefill=True, + ) + runner._dummy_run( + num_tokens=prefill_tokens, + skip_eplb=True, + is_profile=True, + force_attention=True, + create_single_prefill=True, + ) + # Second-chunk shape: indexer sees prior context, hits the alt + # specialization of `_build_prefill_chunk_metadata_kernel`. + runner._dummy_run( + num_tokens=prefill_tokens, + skip_eplb=True, + is_profile=True, + force_attention=True, + create_single_prefill=True, + profile_seq_lens=prefill_tokens * 2, + ) + + +def _flashinfer_autotune_cache_hash(runner: "GPUModelRunner") -> str: + factors = aot_compile_hash_factors(runner.vllm_config) + return hashlib.sha256(str(factors).encode()).hexdigest() + + +def _resolve_flashinfer_autotune_file(runner: "GPUModelRunner") -> Path: + override_dir = envs.VLLM_FLASHINFER_AUTOTUNE_CACHE_DIR + if override_dir: + root = Path(override_dir).expanduser() + else: + from flashinfer.jit import env as flashinfer_jit_env + + flashinfer_workspace = flashinfer_jit_env.FLASHINFER_WORKSPACE_DIR + root = ( + Path(envs.VLLM_CACHE_ROOT) + / "flashinfer_autotune_cache" + / flashinfer_workspace.parent.name + / flashinfer_workspace.name + ) + + output_dir = root / _flashinfer_autotune_cache_hash(runner) + output_dir.mkdir(parents=True, exist_ok=True) + return output_dir / "autotune_configs.json" + + +def kernel_warmup(worker: "Worker"): + # DSv4 mHC TileLang kernels run every decoder layer per token; warm them + # across token sizes first so the first real request doesn't pay JIT cost. + # No-op for non-DSv4 models and for the b12x mHC path (gated inside). + deepseek_v4_mhc_warmup( + worker.get_model(), + max_tokens=worker.scheduler_config.max_num_batched_tokens, + cudagraph_capture_sizes=( + worker.vllm_config.compilation_config.cudagraph_capture_sizes or [] + ), + ) + + # Run next so input-prep kernels JIT against pristine runner state. + _deepseek_v4_sparse_mla_attention_warmup(worker) + _deepseek_v4_request_prep_warmup(worker) + + # Deep GEMM warmup + do_deep_gemm_warmup = ( + envs.VLLM_USE_DEEP_GEMM + and is_deep_gemm_supported() + and envs.VLLM_DEEP_GEMM_WARMUP != "skip" + ) + if do_deep_gemm_warmup: + model = worker.get_model() + max_tokens = worker.scheduler_config.max_num_batched_tokens + deep_gemm_warmup(model, max_tokens) + + enable_flashinfer_autotune = ( + worker.vllm_config.kernel_config.enable_flashinfer_autotune + ) + # FlashInfer autotune for Hopper (SM 9.0) and Blackwell (SM 10.0) GPUs + if enable_flashinfer_autotune is False: + logger.info("Skipping FlashInfer autotune because it is disabled.") + elif has_flashinfer() and current_platform.has_device_capability(90): + flashinfer_autotune(worker.model_runner) + + # FlashInfer attention warmup + # Only warmup if the model has FlashInfer attention groups + # and is not a pooling model + def _is_flashinfer_backend(backend): + try: + return backend.get_name() == "FLASHINFER" + except NotImplementedError: + return False + + if ( + not worker.model_runner.is_pooling_model + and worker.model_runner.attn_groups + # NOTE: This should be `any` instead of `all` but other hybrid attention + # backends don't support this dummy run. Once we remove + # `build_for_cudagraph_capture`, we can change it to `any`. + and all( + _is_flashinfer_backend(group.backend) + for groups in worker.model_runner.attn_groups + for group in groups + ) + ): + logger.info("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, + ) + + +# TODO: remove once FlashInfer upstream fixes the persistent file cache +# to resolve collisions like `use_8x4_sf_layout=True/False`, which causes +# invalid tactics to be chosen +_FLASHINFER_USE_PERSISTENT_CACHE = False + + +def flashinfer_autotune(runner: "GPUModelRunner") -> None: + """ + Autotune FlashInfer operations. + FlashInfer have many implementations for the same operation, + autotuning runs benchmarks for each implementation and stores + the results. The results are cached transparently and + future calls to FlashInfer will use the best implementation. + Without autotuning, FlashInfer will rely on heuristics, which may + be significantly slower. + + Tuning is performed only on rank 0. The resulting cache is broadcast + to every rank so all ranks dispatch the same kernel tactic. + """ + import vllm.utils.flashinfer as fi_utils + from vllm.distributed.parallel_state import get_world_group + + if not _FLASHINFER_USE_PERSISTENT_CACHE: + with torch.inference_mode(), fi_utils.autotune(): + runner._dummy_run( + num_tokens=runner.scheduler_config.max_num_batched_tokens, + skip_eplb=True, + is_profile=True, + ) + get_world_group().barrier() + return + + world = get_world_group() + is_leader = world.rank_in_group == 0 + + cache_path = _resolve_flashinfer_autotune_file(runner) + if is_leader: + logger.info("Using FlashInfer autotune cache file: %s", cache_path) + + # We skip EPLB here since we don't want to record dummy metrics. + # When autotuning with number of tokens m, flashinfer will autotune + # operations for all number of tokens up to m, so we only need to + # run with the max number of tokens. + dummy_run_kwargs = dict( + num_tokens=runner.scheduler_config.max_num_batched_tokens, + skip_eplb=True, + is_profile=True, + ) + + with torch.inference_mode(): + if is_leader: + with fi_utils.autotune(tune_mode=True, cache=str(cache_path)): + runner._dummy_run(**dummy_run_kwargs) + else: + runner._dummy_run(**dummy_run_kwargs) + + # Broadcast autotune cache from rank 0 to all other ranks so every + # rank loads the same set of chosen tactics. + tune_results: bytes | None = None + if is_leader and cache_path.exists(): + with open(cache_path, "rb") as f: + tune_results = f.read() + + tune_results = world.broadcast_object(tune_results, src=0) + + if tune_results is None: + logger.warning( + "No FlashInfer autotune cache entries found." + "Falling back to default tactics." + ) + else: + if not is_leader and world.local_rank == 0: + with open(cache_path, "wb") as f: + f.write(tune_results) + world.barrier() + from flashinfer.autotuner import AutoTuner + + AutoTuner.get().load_configs(str(cache_path)) + logger.info( + "FlashInfer autotune cache loaded on rank %d from %s.", + world.rank_in_group, + cache_path, + ) diff --git a/mods/deepseek-v4-flash-dspark/overlay/vllm/models/deepseek_v4/__init__.py b/mods/deepseek-v4-flash-dspark/overlay/vllm/models/deepseek_v4/__init__.py new file mode 100644 index 00000000..52c2836d --- /dev/null +++ b/mods/deepseek-v4-flash-dspark/overlay/vllm/models/deepseek_v4/__init__.py @@ -0,0 +1,32 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""DeepSeek V4 model — hardware-isolated entry point. + +The actual implementation lives under ``nvidia/`` and ``amd/``; this module +picks the right one for the current platform and re-exports the public +classes used by the model registry and quantization config lookup. +""" + +from typing import TYPE_CHECKING + +from vllm.platforms import current_platform + +from .quant_config import DeepseekV4FP8Config + +# Pick the per-platform implementation. The NVIDIA branch is the static +# default that mypy sees; the ROCm branch overrides it at runtime and is +# kept type-compatible via ``# type: ignore[assignment]``. +if TYPE_CHECKING or not current_platform.is_rocm(): + from .nvidia.dspark import DeepSeekV4DSpark + from .nvidia.model import DeepseekV4ForCausalLM + from .nvidia.mtp import DeepSeekV4MTP +else: + from .amd.model import DeepseekV4ForCausalLM # type: ignore[assignment] + from .amd.mtp import DeepSeekV4MTP # type: ignore[assignment] + +__all__ = [ + "DeepSeekV4DSpark", + "DeepSeekV4MTP", + "DeepseekV4FP8Config", + "DeepseekV4ForCausalLM", +] diff --git a/mods/deepseek-v4-flash-dspark/overlay/vllm/models/deepseek_v4/common/ops/cache_utils.py b/mods/deepseek-v4-flash-dspark/overlay/vllm/models/deepseek_v4/common/ops/cache_utils.py new file mode 100644 index 00000000..d297de7b --- /dev/null +++ b/mods/deepseek-v4-flash-dspark/overlay/vllm/models/deepseek_v4/common/ops/cache_utils.py @@ -0,0 +1,623 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +""" +Triton kernels for DeepseekV4 paged K-cache management and sparse-attention index +preparation. + +- quantize_and_insert_k_cache: quantize bf16 K to UE8M0 FP8 and insert into + the paged cache. +- dequantize_and_gather_k_cache: gather and dequantize FP8 K from the paged + cache for sparse/SWA prefill. +- compute_global_topk_indices_and_lens: map local topk indices to global KV + cache slots and count valid entries. +- combine_topk_swa_indices: concatenate topk compressed indices with SWA + window indices for sparse prefill. +""" + +import torch + +from vllm.triton_utils import tl, triton +from vllm.utils.import_utils import has_cutedsl + + +@triton.jit +def quantize_and_insert_k_kernel( + # Input tensors + k_ptr, # [num_tokens, 512] bf16 + slot_mapping_ptr, # [num_tokens] int64 + # Output tensor + k_cache_ptr, # [num_blocks, block_bytes] as uint8 (flattened view) + # Dimensions + num_tokens, + input_dim: tl.constexpr, # 512 + 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 (paged cache block size) + token_data_size: tl.constexpr, # 576 bytes per token data + block_stride: tl.constexpr, # total bytes per block (padded) + fp8_max: tl.constexpr, + n_quant_blocks: tl.constexpr, # 8 (7 real + 1 padding) +): + """ + Quantize K tensor and insert into paged K cache. + + K Cache block layout (block_size=64 tokens): + - [0, 64*576): Token data, each token has 448 fp8 + 128 bf16 + - [64*576, 64*576 + 64*8): Scales, each token has 8 uint8 scales + - [64*576 + 64*8, block_stride): Padding + + One program per token. + """ + pid = tl.program_id(0) + + if pid >= num_tokens: + return + + # Get slot mapping + slot_idx = tl.load(slot_mapping_ptr + pid) + if slot_idx == -1: + return + + block_idx = slot_idx // cache_block_size + pos_in_block = slot_idx % cache_block_size + + # Input pointer for this token + input_row_ptr = k_ptr + pid * input_dim + + # int64: block_idx * block_stride can exceed 2^31 with many KV-cache blocks + # (e.g. >= 57K at block_stride ~37K). Matches gather path below. + cache_block_ptr = k_cache_ptr + block_idx.to(tl.int64) * block_stride + + # Token data pointer: token data is stored contiguously at start of block + # Each token's data is at offset pos_in_block * token_data_size + token_data_ptr = cache_block_ptr + pos_in_block * token_data_size + + # Scale pointer: scales are stored after ALL token data in the block + # Scale for this token is at offset (64 * 576) + pos_in_block * 8 + token_scale_ptr = ( + cache_block_ptr + cache_block_size * token_data_size + pos_in_block * scale_dim + ) + + # Token data layout: [0:448] fp8, [448:576] bf16 + token_fp8_ptr = token_data_ptr + token_bf16_ptr = token_data_ptr + fp8_dim + + # ========== Quantize and store FP8 portion (first 448 elements) ========== + # Using UE8M0 quantization strategy (scale is power of 2, stored as uint8 exponent) + for qblock_idx in tl.static_range(n_quant_blocks): + qblock_start = qblock_idx * quant_block + + if qblock_start < fp8_dim: + offsets = qblock_start + tl.arange(0, quant_block) + mask = offsets < fp8_dim + + # Load bf16 input + x = tl.load(input_row_ptr + offsets, mask=mask, other=0.0) + + # Compute absmax scale (same as CUDA kernel) + abs_x = tl.abs(x) + block_max = tl.max(abs_x, axis=0) + block_max = tl.maximum(block_max, 1e-4) # Match CUDA: fmaxf(amax, 1e-4) + + # UE8M0: Round scale UP to next power of 2 + # scale = 2^ceil(log2(block_max / fp8_max)) + raw_scale = block_max / fp8_max + log_scale = tl.log2(raw_scale) + exponent = tl.ceil(log_scale) # Round UP to next integer exponent + scale = tl.exp2(exponent) # scale = 2^exponent (power of 2) + + # Quantize to fp8: fp8_value = bf16_value / scale + x_scaled = x / scale + x_clamped = tl.clamp(x_scaled, -fp8_max, fp8_max) + + # Convert to fp8, then bitcast to uint8 for storage + x_fp8 = x_clamped.to(tl.float8e4nv) + x_uint8 = x_fp8.to(tl.uint8, bitcast=True) + + # Store as uint8 (1 byte each) + tl.store(token_fp8_ptr + offsets, x_uint8, mask=mask) + + # UE8M0 scale encoding: stored_value = exponent + 127 (bias) + # During dequant: scale = 2^(stored_value - 127) + encoded_scale = exponent + 127.0 + encoded_scale = tl.maximum(tl.minimum(encoded_scale, 255.0), 0.0) + tl.store(token_scale_ptr + qblock_idx, encoded_scale.to(tl.uint8)) + + # Padding scale at index 7 + tl.store(token_scale_ptr + 7, tl.zeros((), dtype=tl.uint8)) + + # ========== Store BF16 portion (last 64 elements, no quantization) ========== + bf16_input_offset = fp8_dim + + # Process bf16 in chunks of 16 + bf16_out_ptr = token_bf16_ptr.to(tl.pointer_type(tl.bfloat16)) + for i in tl.static_range(bf16_dim // 16): + chunk_offsets = i * 16 + tl.arange(0, 16) + bf16_vals = tl.load(input_row_ptr + bf16_input_offset + chunk_offsets) + tl.store(bf16_out_ptr + chunk_offsets, bf16_vals) + + +def quantize_and_insert_k_cache( + k: torch.Tensor, # [num_tokens, 512] bf16 + k_cache: torch.Tensor, # [num_blocks, block_bytes] uint8 + slot_mapping: torch.Tensor, # [num_tokens] int64 + block_size: int = 64, + is_ue8m0: bool = True, +): + """ + Quantize K tensor and insert into paged K cache. + + K Cache block layout (block_size=64 tokens): + - First 64 * 576 = 36864 bytes: Token data + - Each token: 448 bytes (fp8) + 128 bytes (bf16) + - Next 64 * 8 = 512 bytes: Scales + - Each token: 8 bytes (uint8 scales, 7 real + 1 padding) + - Padded to multiple of 576 + """ + assert k.dim() == 2 and k.shape[1] == 512, ( + f"K must be [num_tokens, 512], got {k.shape}" + ) + assert k.dtype == torch.bfloat16, f"K must be bf16, got {k.dtype}" + assert is_ue8m0, "Only support ue8m0 quantization." + + # NOTE: When using DP, slot_mapping.shape[0] can be less than k.shape[0] due to + # padding. Always use slot_mapping.shape[0] as the token count. + num_tokens = slot_mapping.shape[0] + block_stride = k_cache.stride(0) # bytes per block + + TOKEN_FP8_DIM = 448 + TOKEN_BF16_DIM = 64 + TOKEN_SCALE_DIM = 8 + QUANT_BLOCK_SIZE = 64 + FP8_MAX = 448.0 + TOKEN_DATA_SIZE = TOKEN_FP8_DIM + TOKEN_BF16_DIM * 2 + + grid = (num_tokens,) + + quantize_and_insert_k_kernel[grid]( + k, + slot_mapping, + k_cache, + num_tokens, + input_dim=512, + fp8_dim=TOKEN_FP8_DIM, + bf16_dim=TOKEN_BF16_DIM, + scale_dim=TOKEN_SCALE_DIM, + quant_block=QUANT_BLOCK_SIZE, + cache_block_size=block_size, + token_data_size=TOKEN_DATA_SIZE, + block_stride=block_stride, + fp8_max=FP8_MAX, + n_quant_blocks=8, + ) + + +@triton.jit +def _dequantize_and_gather_k_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 +): + 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 all tokens + gather_len = seq_len + start_pos = seq_len - gather_len + + for i in range(worker_id, gather_len, num_workers): + # Calculate the actual token index in the sequence + pos = start_pos + i + + # Calculate which block and position within block + block_in_seq = pos // cache_block_size + pos_in_block = pos % cache_block_size + + # Get physical block index from block table + 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) # int32 + + # int64: physical_block_idx * block_stride can exceed 2^31 with many + # KV-cache blocks (e.g. >= 57K at block_stride ~37K). + cache_block_ptr = k_cache_ptr + physical_block_idx.to(tl.int64) * block_stride + + # Token data pointer + token_data_ptr = cache_block_ptr + pos_in_block * token_data_size + + # Scale pointer: after all token data + token_scale_ptr = ( + cache_block_ptr + + cache_block_size * token_data_size + + pos_in_block * scale_dim + ) + + # Token data layout: [0:448] fp8, [448:576] bf16 + token_fp8_ptr = token_data_ptr + token_bf16_ptr = token_data_ptr + fp8_dim + + # Output pointer for this token (flattened). + # int64: batch_idx * out_stride0 can exceed 2^31 when the gather buffer + # is sized for very long sequences (out_stride0 = max_num_tokens * 512; + # wraps at >= ~819K tokens with batch_idx >= 4) — same overflow class + # as the guarded physical_block_idx multiply above. + output_row_ptr = ( + out_ptr + + batch_idx.to(tl.int64) * out_stride0 + + (offset + i) * out_stride1 + ) + + # ========== Dequantize FP8 portion using UE8M0 ========== + for qblock_idx in tl.static_range(n_quant_blocks): + qblock_start = qblock_idx * quant_block + + if qblock_start < fp8_dim: + offsets = qblock_start + tl.arange(0, quant_block) + mask = offsets < fp8_dim + + # Load quantized fp8 values (stored as uint8) + x_uint8 = tl.load(token_fp8_ptr + offsets, mask=mask, other=0) + + # Bitcast uint8 back to fp8 + x_fp8 = x_uint8.to(tl.float8e4nv, bitcast=True) + + # Convert fp8 to float32 for computation + x_float = x_fp8.to(tl.float32) + + # Load and decode UE8M0 scale + # UE8M0: scale = 2^(stored_value - 127) + encoded_scale = tl.load(token_scale_ptr + qblock_idx) + exponent = encoded_scale.to(tl.float32) - 127.0 + scale = tl.exp2(exponent) + + # Dequantize: bf16_value = fp8_value * scale + x_dequant = x_float * scale + + # Store as bf16 + tl.store(output_row_ptr + offsets, x_dequant.to(tl.bfloat16), mask=mask) + + # ========== Copy BF16 portion directly ========== + bf16_output_offset = fp8_dim # After 448 elements in output + + # Read bf16 from cache + bf16_cache_ptr = token_bf16_ptr.to(tl.pointer_type(tl.bfloat16)) + + # Process in chunks of 16 + for j in tl.static_range(bf16_dim // 16): + chunk_offsets = j * 16 + tl.arange(0, 16) + bf16_vals = tl.load(bf16_cache_ptr + chunk_offsets) + tl.store(output_row_ptr + bf16_output_offset + chunk_offsets, bf16_vals) + + +def dequantize_and_gather_k_cache_triton( + # [num_reqs, max_num_tokens, head_size] + out: torch.Tensor, + # [num_blocks, block_size, head_bytes] + k_cache: torch.Tensor, + # [num_reqs] + seq_lens: torch.Tensor, + # [num_reqs] + gather_lens: torch.Tensor | None, + # [num_reqs, max_blocks_per_seq] + block_table: torch.Tensor, + block_size: int, + offset: int, +) -> None: + TOKEN_FP8_DIM = 448 + TOKEN_BF16_DIM = 64 + TOKEN_SCALE_DIM = 8 + QUANT_BLOCK_SIZE = 64 + FP8_MAX = 448.0 + TOKEN_DATA_SIZE = TOKEN_FP8_DIM + TOKEN_BF16_DIM * 2 + + num_reqs = seq_lens.shape[0] + NUM_WORKERS = 128 + _dequantize_and_gather_k_kernel[(num_reqs, NUM_WORKERS)]( + out, + out.stride(0), + out.stride(1), + k_cache, + seq_lens, + block_table, + offset, + gather_lens, + max_blocks_per_seq=block_table.shape[-1], + fp8_dim=TOKEN_FP8_DIM, + bf16_dim=TOKEN_BF16_DIM, + scale_dim=TOKEN_SCALE_DIM, + quant_block=QUANT_BLOCK_SIZE, + cache_block_size=block_size, + token_data_size=TOKEN_DATA_SIZE, + block_stride=k_cache.stride(0), + output_dim=512, + fp8_max=FP8_MAX, + n_quant_blocks=7, + ) + + +def dequantize_and_gather_k_cache( + # [num_reqs, max_num_tokens, head_size] + out: torch.Tensor, + # [num_blocks, block_size, head_bytes] + k_cache: torch.Tensor, + # [num_reqs] + seq_lens: torch.Tensor, + # [num_reqs] + gather_lens: torch.Tensor | None, + # [num_reqs, max_blocks_per_seq] + block_table: torch.Tensor, + block_size: int, + offset: int, +) -> None: + if has_cutedsl(): + # lazily import, otherwise some tests fail due to CUDA driver init failure. + from vllm.models.deepseek_v4.nvidia.ops import ( + dequantize_and_gather_k_cache_cutedsl, + ) + + dequantize_and_gather_k_cache_cutedsl( + out, k_cache, seq_lens, gather_lens, block_table, block_size, offset + ) + return + + dequantize_and_gather_k_cache_triton( + out, k_cache, seq_lens, gather_lens, block_table, block_size, offset + ) + + +def compute_global_topk_indices_and_lens( + topk_indices: torch.Tensor, + token_to_req_indices: torch.Tensor, + block_table: torch.Tensor, + block_size: int, + is_valid_token: torch.Tensor, + num_blocks: int | None = None, +) -> tuple[torch.Tensor, torch.Tensor]: + """Map local topk indices to global KV cache slots and count valid entries. + + Fuses three operations into a single kernel: + 1. Block-table lookup (local index → global slot id) + 2. Valid-entry counting (topk_lens per token) + 3. Masking padding tokens to length 0 + + ``num_blocks`` is the block count of the KV cache these slots index into; + slots resolving to a block outside it reference nonexistent KV (stale + block-table entries) and are dropped rather than gathered out of bounds. + When ``None`` the bound is disabled (no-op sentinel). + """ + num_tokens = topk_indices.shape[0] + global_topk_indices = torch.empty_like(topk_indices) + topk_lens = torch.empty(num_tokens, dtype=torch.int32, device=topk_indices.device) + num_blocks_bound = (1 << 31) - 1 if num_blocks is None else num_blocks + _compute_global_topk_indices_and_lens_kernel[(num_tokens,)]( + global_topk_indices, + global_topk_indices.stride(0), + topk_lens, + topk_indices, + topk_indices.stride(0), + topk_indices.shape[-1], + token_to_req_indices, + block_table, + block_table.stride(0), + block_size, + is_valid_token, + num_blocks_bound, + TRITON_BLOCK_SIZE=1024, + ) + return global_topk_indices, topk_lens + + +@triton.jit +def _compute_global_topk_indices_and_lens_kernel( + global_topk_indices_ptr, + global_topk_indices_stride, + topk_lens_ptr, + topk_indices_ptr, + topk_indices_stride, + topk, + token_to_req_indices_ptr, + block_table_ptr, + block_table_stride, + block_size, + is_valid_token_ptr, + num_blocks, + TRITON_BLOCK_SIZE: tl.constexpr, +): + token_idx = tl.program_id(0) + is_valid_token = tl.load(is_valid_token_ptr + token_idx) + req_idx = tl.load(token_to_req_indices_ptr + token_idx) + + count = tl.zeros((), dtype=tl.int32) + for i in range(0, topk, TRITON_BLOCK_SIZE): + offset = i + tl.arange(0, TRITON_BLOCK_SIZE) + mask = offset < topk + + local_idx = tl.load( + topk_indices_ptr + token_idx * topk_indices_stride + offset, + mask=mask, + other=-1, + ) + # Padding tokens (is_valid_token == 0) keep uninitialized entries in the + # persistent topk_indices buffer; those stale indices must not drive the + # block-table gather or it reads out of bounds. Gate the gather on token + # validity as well as the per-slot sentinel. + is_valid = (local_idx >= 0) & (is_valid_token != 0) + + block_indices = local_idx // block_size + # A stale local_idx can exceed the request's block-table row; don't read + # past it. + is_valid = is_valid & (block_indices < block_table_stride) + block_numbers = tl.load( + block_table_ptr + req_idx * block_table_stride + block_indices, + mask=mask & is_valid, + other=-1, + ) + # A block number outside the KV cache is a stale/unallocated block-table + # entry pointing at nonexistent KV. Dropping it (vs. gathering out of + # bounds) is correct: there is nothing valid to attend at that slot. + is_valid = is_valid & (block_numbers >= 0) & (block_numbers < num_blocks) + block_offsets = local_idx % block_size + + slot_ids = block_numbers * block_size + block_offsets + slot_ids = tl.where(is_valid, slot_ids, -1) + tl.store( + global_topk_indices_ptr + token_idx * global_topk_indices_stride + offset, + slot_ids, + mask=mask, + ) + count += tl.sum(is_valid.to(tl.int32), axis=0) + + # Zero out length for padding tokens. + tl.store(topk_lens_ptr + token_idx, tl.where(is_valid_token, count, 0)) + + +# FlashMLA sparse prefill asserts `params.topk % B_TOPK == 0` (see +# flashmla/csrc/sm100/prefill/sparse/fwd/head{64,128}/phase1.cuh). B_TOPK is +# 64 for the h_q=64 kernel and 128 for h_q=128; pad to 128 to satisfy both. +# The extra slots stay as -1 sentinels and `combined_lens` caps the valid +# range via `topk_length`, so padding is a no-op at kernel level. +_SPARSE_PREFILL_TOPK_ALIGNMENT = 128 + + +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]: + num_tokens = topk_indices.shape[0] + num_reqs = seq_lens.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 + ) + + NUM_WORKERS = 128 + _combine_topk_swa_indices_kernel[(num_reqs, NUM_WORKERS)]( + combined_indices, + combined_indices.stride(0), + combined_lens, + topk_indices, + topk_indices.stride(0), + query_start_loc, + seq_lens, + gather_lens, + M, + N, + TOP_K=topk, + COMPRESS_RATIO=compress_ratio, + WINDOW_SIZE=window_size, + PADDED_TOP_K=triton.next_power_of_2(topk_indices.shape[-1]), + ) + return combined_indices, combined_lens + + +@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, + PADDED_TOP_K: tl.constexpr, +): + batch_idx = tl.program_id(0) + worker_id = tl.program_id(1) + num_workers = tl.num_programs(1) + + # query_start_loc is a global tensor; rebase to chunk-local offsets + # by subtracting the chunk's starting value. + 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 + # The SWA portion of the gathered buffer starts from position + # (seq_len - gather_len), not position 0. We need this offset + # to correctly index into the gathered buffer. + gather_start = seq_len - gather_len + + for token_idx in range(query_start + worker_id, query_end, num_workers): + # topk_len is fully determined by the query token's absolute position: + # both the C4A indexer and the C128A metadata builder emit + # min((pos + 1) // compress_ratio, topk_tokens) valid entries. + # Caller passes TOP_K=0 for SWA-only layers to zero this out. + 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) + + offset = tl.arange(0, PADDED_TOP_K) + mask = offset < topk_len + topk_indices = tl.load( + topk_indices_ptr + token_idx * topk_indices_stride + offset, + mask=mask, + ) + tl.store( + combined_indices_ptr + token_idx * combined_indices_stride + offset, + topk_indices + M * batch_idx, + mask=mask, + ) + offset = tl.arange(0, WINDOW_SIZE) + # Index into gathered buffer: N + (position - gather_start) + # For positions [pos - swa_len + 1, pos], the buffer indices are: + # [N + pos - swa_len + 1 - gather_start, N + pos - gather_start] + tl.store( + combined_indices_ptr + + token_idx * combined_indices_stride + + topk_len + + offset, + M * batch_idx + N + offset + pos - swa_len + 1 - gather_start, + mask=offset < swa_len, + ) + + combined_len = topk_len + swa_len + tl.store(combined_lens_ptr + token_idx, combined_len) diff --git a/mods/deepseek-v4-flash-dspark/overlay/vllm/models/deepseek_v4/nvidia/dspark.py b/mods/deepseek-v4-flash-dspark/overlay/vllm/models/deepseek_v4/nvidia/dspark.py new file mode 100644 index 00000000..29c3cc3c --- /dev/null +++ b/mods/deepseek-v4-flash-dspark/overlay/vllm/models/deepseek_v4/nvidia/dspark.py @@ -0,0 +1,1192 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Experimental DSpark draft model for DeepSeek V4 Flash. + +This module follows the reference implementation shipped with +DeepSeek-V4-Flash-DSpark. It intentionally keeps the draft-side DSpark +attention cache internal to the draft model instead of registering more vLLM +KV-cache layers; DSpark uses a small sliding window over target features and +draft block tokens, which is different from the normal MTP cache contract. +""" + +from __future__ import annotations + +import os +import typing +from collections.abc import Callable, Iterable + +import regex as re +import torch +import torch.nn as nn + +from vllm.config import VllmConfig +from vllm.distributed import ( + get_tensor_model_parallel_rank, + get_tensor_model_parallel_world_size, + tensor_model_parallel_all_gather, +) +from vllm.logger import init_logger +from vllm.model_executor.layers.fused_moe import FusedMoE +from vllm.model_executor.layers.layernorm import RMSNorm +from vllm.model_executor.layers.linear import ( + ColumnParallelLinear, + MergedColumnParallelLinear, + ReplicatedLinear, + RowParallelLinear, +) +from vllm.model_executor.layers.logits_processor import LogitsProcessor +from vllm.model_executor.layers.mhc import HCHeadOp +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.model_loader.weight_utils import default_weight_loader +from vllm.model_executor.models.utils import maybe_prefix +from vllm.models.deepseek_v4.common.ops import fused_inv_rope_fp8_quant +from vllm.platforms import current_platform +from vllm.v1.spec_decode.dspark import ( + map_dspark_stacked_param_name, + unpack_mhc_pre_outputs, +) + +from .dspark_kernels import ( + dspark_markov_argmax, + dspark_quant_dequant_nope, + dspark_sparse_attention, +) +from .model import ( + DeepseekV4MoE, + make_deepseek_v4_expert_params_mapping, +) + +logger = init_logger(__name__) + + +# Kill switch for A/B diagnosis: DSPARK_SLOT_CLAMP=0 reverts the protective +# clamp (detect+log only, return the index unchanged) so operators can verify +# the clamp is load-bearing on their rig. Default on = the fix. Read once at +# import (fixed per container start). +_SLOT_CLAMP_ENABLED = os.environ.get("DSPARK_SLOT_CLAMP", "1") != "0" + + +def _guard_slot_index( + slot_index: torch.Tensor | None, num_rows: int, site: str +) -> torch.Tensor | None: + """Bounds-check DSpark ring-buffer slot ids before a KV gather. + + A stale/out-of-range ``slot_index`` (observed after request condensation at + long context) gathering row >= ``num_rows`` triggers a device-side + ``indexSelectSmallIndex`` assert (``srcIndex < srcSelectDimSize``) that kills + the worker. Clamp into range on-device so the gather degrades to a rejected + speculation instead of a crash: this is the draft path only, so a clamped + (wrong) slot yields a bad draft token that the target model simply rejects + at verification — it cannot corrupt output. + + The clamp is graph-safe (pure device op, baked into the captured graph and + active at replay). The loud host-side value log requires a device sync, + which is illegal mid CUDA-graph-capture, so it only runs on eager paths + (``is_current_stream_capturing()`` is a driver query, no sync). + """ + if slot_index is None or slot_index.numel() == 0: + return slot_index + clamped = slot_index.clamp(0, num_rows - 1) + if not torch.cuda.is_current_stream_capturing(): + hi = int(slot_index.max()) + lo = int(slot_index.min()) + if hi >= num_rows or lo < 0: + logger.error( + "DSPARK_STALE_SLOT_INDEX site=%s num_rows=%d min=%d max=%d " + "slot_index=%s", + site, + num_rows, + lo, + hi, + slot_index.tolist(), + ) + return clamped if _SLOT_CLAMP_ENABLED else slot_index + + +_EXPERT_SCALE_RE = re.compile(r"\.experts\.\d+\.w[123]\.scale$") + + +def _read_bool_env(name: str, default: str = "0") -> bool: + return os.getenv(name, default).strip().lower() in {"1", "true", "yes", "on"} + + +def _linear_no_bias( + linear: nn.Module, + x: torch.Tensor, +) -> torch.Tensor: + out = linear(x) + if isinstance(out, tuple): + y, bias = out + assert bias is None + return y + return out + + +def _vocab_parallel_argmax( + local_logits: torch.Tensor, + lm_head: VocabParallelEmbedding, +) -> torch.Tensor: + """Return global greedy token ids from local vocab-parallel logits.""" + num_pad = lm_head.shard_indices.num_org_vocab_padding + if num_pad > 0: + local_logits[..., -num_pad:] = -float("inf") + + local_max_vals, local_max_indices = local_logits.max(dim=-1) + global_indices = local_max_indices + lm_head.shard_indices.org_vocab_start_index + return _vocab_parallel_argmax_from_local(local_max_vals, global_indices) + + +def _vocab_parallel_argmax_from_local( + local_max_vals: torch.Tensor, + global_indices: torch.Tensor, +) -> torch.Tensor: + """Return global token ids from per-rank local top-1 candidates.""" + + tp_size = get_tensor_model_parallel_world_size() + if tp_size == 1: + return global_indices.to(torch.long) + + local_pair = torch.stack( + [local_max_vals.float(), global_indices.float()], + dim=-1, + ) + gathered = tensor_model_parallel_all_gather(local_pair, dim=-1) + gathered = gathered.view(local_max_vals.shape[0], tp_size, 2) + max_rank_idx = gathered[:, :, 0].argmax(dim=-1, keepdim=True) + top_tokens = gathered[:, :, 1].gather(dim=-1, index=max_rank_idx) + return top_tokens.squeeze(-1).to(torch.long) + + +def _vocab_parallel_markov_argmax( + base_logits: torch.Tensor, + markov_embed: torch.Tensor, + markov_w2: ParallelLMHead, + lm_head: VocabParallelEmbedding, +) -> torch.Tensor: + """Return global greedy ids for base logits plus DSpark Markov bias.""" + + num_pad = lm_head.shard_indices.num_org_vocab_padding + local_max_vals, local_max_indices = dspark_markov_argmax( + base_logits, + markov_embed, + markov_w2.weight, + num_pad=num_pad, + ) + global_indices = local_max_indices + lm_head.shard_indices.org_vocab_start_index + return _vocab_parallel_argmax_from_local(local_max_vals, global_indices) + + +class DeepSeekV4DSparkAttention(nn.Module): + """DSpark sparse MLA attention with an internal main-token KV window.""" + + def __init__( + self, + vllm_config: VllmConfig, + *, + prefix: str, + ) -> None: + super().__init__() + config = vllm_config.model_config.hf_config + quant_config = vllm_config.quant_config + + self.hidden_size = config.hidden_size + self.dtype = vllm_config.model_config.dtype + self.n_heads = config.num_attention_heads + tp_size = get_tensor_model_parallel_world_size() + assert self.n_heads % tp_size == 0 + self.n_local_heads = self.n_heads // tp_size + self.q_lora_rank = config.q_lora_rank + self.o_lora_rank = config.o_lora_rank + self.head_dim = config.head_dim + self.rope_head_dim = config.qk_rope_head_dim + self.nope_head_dim = self.head_dim - self.rope_head_dim + self.n_groups = config.o_groups + self.n_local_groups = self.n_groups // tp_size + self.window_size = config.sliding_window + self.block_size = config.dspark_block_size + self.eps = config.rms_norm_eps + self.softmax_scale = self.head_dim**-0.5 + self._reference_kv_quant_dequant = _read_bool_env( + "VLLM_DSPARK_REFERENCE_KV_QUANT_DEQUANT" + ) + cap = current_platform.get_device_capability() + assert cap is not None, "DSpark attention requires a CUDA device" + self._einsum_recipe = (1, 128, 128) if cap.major <= 9 else (1, 1, 128) + self._tma_aligned_scales = cap.major >= 10 + + self.attn_sink = nn.Parameter( + torch.full((self.n_local_heads,), -float("inf"), dtype=torch.float32), + requires_grad=False, + ) + self.fused_wqa_wkv = MergedColumnParallelLinear( + self.hidden_size, + [self.q_lora_rank, self.head_dim], + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.fused_wqa_wkv", + disable_tp=True, + ) + self.q_norm = RMSNorm(self.q_lora_rank, self.eps) + self.wq_b = ColumnParallelLinear( + self.q_lora_rank, + self.n_heads * self.head_dim, + bias=False, + quant_config=quant_config, + return_bias=False, + prefix=f"{prefix}.wq_b", + ) + self.kv_norm = RMSNorm(self.head_dim, self.eps) + self.wo_a = ColumnParallelLinear( + self.n_heads * self.head_dim // self.n_groups, + self.n_groups * self.o_lora_rank, + bias=False, + quant_config=quant_config, + return_bias=False, + prefix=f"{prefix}.wo_a", + ) + self.wo_a.is_bmm = True + self.wo_a.bmm_batch_size = self.n_local_groups + self.wo_b = RowParallelLinear( + self.n_groups * self.o_lora_rank, + self.hidden_size, + bias=False, + quant_config=quant_config, + return_bias=False, + prefix=f"{prefix}.wo_b", + ) + + rope_parameters = config.rope_parameters + rope_parameters["rope_theta"] = config.rope_theta + if config.rope_parameters["rope_type"] != "default": + config.rope_parameters["rope_type"] = ( + "deepseek_yarn" + if config.rope_parameters.get("apply_yarn_scaling", True) + else "deepseek_llama_scaling" + ) + rope_parameters["mscale"] = 0 + rope_parameters["mscale_all_dim"] = 0 + rope_parameters["is_deepseek_v4"] = True + rope_parameters["rope_dim"] = self.rope_head_dim + self.rotary_emb = get_rope( + self.head_dim, + max_position=config.max_position_embeddings, + rope_parameters=rope_parameters, + is_neox_style=False, + ) + + max_batch_size = vllm_config.scheduler_config.max_num_seqs + self.register_buffer( + "main_kv_cache", + torch.zeros( + max_batch_size, + self.window_size, + self.head_dim, + dtype=vllm_config.model_config.dtype, + device=current_platform.device_type, + ), + persistent=False, + ) + self.register_buffer( + "sparse_scores", + torch.empty( + max_batch_size, + self.block_size, + self.n_local_heads, + self.window_size + self.block_size, + dtype=torch.float32, + device=current_platform.device_type, + ), + persistent=False, + ) + if self._reference_kv_quant_dequant: + logger.info( + "DSpark reference KV FP8 quant-dequant enabled for no-RoPE " + "dimensions in %s.", + prefix, + ) + + def _maybe_quant_dequant_kv(self, kv: torch.Tensor) -> torch.Tensor: + if self._reference_kv_quant_dequant: + return dspark_quant_dequant_nope( + kv, + rope_dim=self.rope_head_dim, + group_size=64, + ) + return kv + + def _project_kv( + self, + hidden_states: torch.Tensor, + positions: torch.Tensor, + ) -> torch.Tensor: + qra_kv = _linear_no_bias(self.fused_wqa_wkv, hidden_states) + _, kv = qra_kv.split([self.q_lora_rank, self.head_dim], dim=-1) + kv = self.kv_norm(kv) + kv, _ = self.rotary_emb(positions, kv.unsqueeze(1), None) + kv = kv.squeeze(1) + return self._maybe_quant_dequant_kv(kv) + + def _store_main_kv_ragged( + self, + main_x: torch.Tensor, + main_positions: torch.Tensor, + query_start_loc: list[int], + num_rejected_tokens: torch.Tensor | None, + slot_index: list[int] | torch.Tensor | None, + ) -> None: + """Scatter ragged (mixed prefill+decode) per-request rows. + + ``main_x`` is flat [total_rows, hidden_size] and ``main_positions`` is + flat [total_rows]; ``query_start_loc`` gives per-request offsets. Each + request's rows are scattered into ITS slot's ring buffer using absolute + positions (``positions % window``), so no rectangular view / uniform + length is required. Runs eager (mixed steps are never cudagraphed). + """ + flat_positions = main_positions.reshape(-1) + flat_kv = self._project_kv(main_x, flat_positions) + batch_size = len(query_start_loc) - 1 + rejected = None + if num_rejected_tokens is not None: + rejected = num_rejected_tokens.to( + device=main_x.device, + dtype=torch.long, + non_blocking=True, + ).view(batch_size) + for i in range(batch_size): + start = int(query_start_loc[i]) + end = int(query_start_loc[i + 1]) + seg_kv = flat_kv[start:end] + seg_positions = flat_positions[start:end].to(torch.long) + if seg_kv.shape[0] > self.window_size: + seg_kv = seg_kv[-self.window_size :] + seg_positions = seg_positions[-self.window_size :] + seg_len = seg_kv.shape[0] + if seg_len == 0: + continue + slots = seg_positions.remainder(self.window_size) + row = i if slot_index is None else int(slot_index[i]) + num_rows = self.main_kv_cache.shape[0] + if slot_index is not None and not 0 <= row < num_rows: + logger.error( + "DSPARK_STALE_SLOT_INDEX site=ragged row=%d num_rows=%d i=%d", + row, + num_rows, + i, + ) + if _SLOT_CLAMP_ENABLED: + row = min(max(row, 0), num_rows - 1) + cache_row = self.main_kv_cache[row] + values = seg_kv + if rejected is not None: + valid_len = (seg_len - rejected[i]).clamp(min=1, max=seg_len) + token_offsets = torch.arange( + seg_len, + device=seg_kv.device, + dtype=torch.long, + ) + valid_mask = (token_offsets < valid_len).unsqueeze(-1) + old_values = cache_row.index_select(0, slots) + values = torch.where(valid_mask, seg_kv, old_values) + cache_row.index_copy_(0, slots, values) + + def store_main_kv( + self, + main_x: torch.Tensor, + main_positions: torch.Tensor, + num_rejected_tokens: torch.Tensor | None = None, + slot_index: torch.Tensor | None = None, + query_start_loc: list[int] | None = None, + ) -> None: + if query_start_loc is not None: + self._store_main_kv_ragged( + main_x, + main_positions, + query_start_loc, + num_rejected_tokens, + slot_index, + ) + return + if main_x.shape[1] > self.window_size: + main_x = main_x[:, -self.window_size :] + main_positions = main_positions[:, -self.window_size :] + batch_size, seq_len, _ = main_x.shape + flat_kv = self._project_kv( + main_x.reshape(batch_size * seq_len, self.hidden_size), + main_positions.reshape(batch_size * seq_len), + ).view(batch_size, seq_len, self.head_dim) + slots = main_positions.to(torch.long).remainder(self.window_size) + slots_index = slots.unsqueeze(-1).expand(-1, -1, self.head_dim) + values = flat_kv + # ``slot_index is None`` keeps the original single-stream behaviour: + # operate in place on the leading ``batch_size`` rows (byte-for-byte + # identical). Otherwise gather the per-request slot rows, update them, + # and scatter them back, so the stateful sliding-window KV follows the + # request id rather than the (condensable) batch-row position. + if slot_index is None: + cache_rows = self.main_kv_cache[:batch_size] + else: + slot_index = _guard_slot_index( + slot_index, self.main_kv_cache.shape[0], "store_main_kv" + ) + cache_rows = self.main_kv_cache.index_select(0, slot_index) + if num_rejected_tokens is not None: + rejected = num_rejected_tokens.to( + device=main_x.device, + dtype=torch.long, + non_blocking=True, + ).view(batch_size) + valid_lengths = (seq_len - rejected).clamp(min=1, max=seq_len) + token_offsets = torch.arange( + seq_len, + device=main_x.device, + dtype=torch.long, + ).view(1, seq_len) + valid_mask = token_offsets < valid_lengths.view(batch_size, 1) + old_values = cache_rows.gather(1, slots_index) + values = torch.where(valid_mask.unsqueeze(-1), flat_kv, old_values) + cache_rows.scatter_(1, slots_index, values) + if slot_index is not None: + self.main_kv_cache.index_copy_(0, slot_index, cache_rows) + + def _project_q_and_draft_kv( + self, + hidden_states: torch.Tensor, + positions: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + qra_kv = _linear_no_bias(self.fused_wqa_wkv, hidden_states) + qra, kv = qra_kv.split([self.q_lora_rank, self.head_dim], dim=-1) + qra = self.q_norm(qra) + q = _linear_no_bias(self.wq_b, qra).view(-1, self.n_local_heads, self.head_dim) + q = q * torch.rsqrt(q.square().mean(-1, keepdim=True) + self.eps) + kv = self.kv_norm(kv) + q, _ = self.rotary_emb(positions, q, None) + kv, _ = self.rotary_emb(positions, kv.unsqueeze(1), None) + kv = kv.squeeze(1) + return q, self._maybe_quant_dequant_kv(kv) + + def forward_dspark( + self, + hidden_states: torch.Tensor, + positions: torch.Tensor, + *, + batch_size: int, + block_size: int, + main_x: torch.Tensor, + main_positions: torch.Tensor, + store_main_kv: bool = True, + slot_index: torch.Tensor | None = None, + ) -> torch.Tensor: + if store_main_kv: + self.store_main_kv(main_x, main_positions, slot_index=slot_index) + + q, draft_kv = self._project_q_and_draft_kv(hidden_states, positions) + q = q.view(batch_size, block_size, self.n_local_heads, self.head_dim) + draft_kv = draft_kv.view(batch_size, block_size, self.head_dim) + current_positions = main_positions[:, -1] + valid_main_lengths = torch.minimum( + current_positions + 1, + torch.full_like(current_positions, self.window_size), + ) + + # The kernel indexes contiguous batch rows 0..B-1. ``slot_index is + # None`` passes the full persistent cache (rows 0..B-1 == identity); + # otherwise gather the per-request slot rows into a contiguous + # [B, window, head_dim] view in the same order as ``valid_main_lengths``. + if slot_index is None: + main_kv_cache = self.main_kv_cache + else: + slot_index = _guard_slot_index( + slot_index, self.main_kv_cache.shape[0], "forward_dspark" + ) + main_kv_cache = self.main_kv_cache.index_select(0, slot_index) + + out = dspark_sparse_attention( + q, + draft_kv, + main_kv_cache, + valid_main_lengths, + self.attn_sink, + self.softmax_scale, + self.sparse_scores[:batch_size, :block_size], + ).to(self.dtype) + out_fp8, out_scale = fused_inv_rope_fp8_quant( + out, + positions, + self.rotary_emb.cos_sin_cache, + n_groups=self.n_local_groups, + heads_per_group=self.n_local_heads // self.n_local_groups, + nope_dim=self.nope_head_dim, + rope_dim=self.rope_head_dim, + tma_aligned_scales=self._tma_aligned_scales, + ) + projected = torch.empty( + (batch_size * block_size, self.n_local_groups, self.o_lora_rank), + dtype=self.dtype, + device=out.device, + ) + torch.ops.vllm.deepseek_v4_fp8_einsum( + out_fp8, + out_scale, + self.wo_a.weight, + self.wo_a.weight_scale_inv, + projected, + "bhr,hdr->bhd", + list(self._einsum_recipe), + ) + return _linear_no_bias(self.wo_b, projected.flatten(1).to(self.dtype)) + + +class DeepSeekV4DSparkMarkovHead(nn.Module): + def __init__(self, vllm_config: VllmConfig, *, prefix: str) -> None: + super().__init__() + config = vllm_config.model_config.hf_config + self._replicated_w1 = _read_bool_env("VLLM_DSPARK_REPLICATE_MARKOV_W1") + if self._replicated_w1: + self.markov_w1 = nn.Embedding( + config.vocab_size, + config.dspark_markov_rank, + dtype=vllm_config.model_config.dtype, + ) + self.markov_w1.weight.requires_grad_(False) + logger.info( + "DSpark replicated Markov W1 enabled for %s. This removes " + "the per-position vocab-parallel embedding all-reduce.", + prefix, + ) + else: + self.markov_w1 = VocabParallelEmbedding( + config.vocab_size, + config.dspark_markov_rank, + prefix=f"{prefix}.markov_w1", + ) + self.markov_w2 = ParallelLMHead( + config.vocab_size, + config.dspark_markov_rank, + prefix=f"{prefix}.markov_w2", + ) + self.logits_processor = LogitsProcessor(config.vocab_size) + + def forward(self, token_ids: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + markov_embed = self.markov_w1(token_ids) + markov_logits = self.logits_processor(self.markov_w2, markov_embed) + return markov_logits, markov_embed + + def forward_local( + self, + token_ids: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + markov_embed = self.markov_w1(token_ids) + markov_logits = self.markov_w2.quant_method.apply( + self.markov_w2, + markov_embed, + bias=None, + ) + return markov_logits, markov_embed + + +class DeepSeekV4DSparkConfidenceHead(nn.Module): + def __init__(self, vllm_config: VllmConfig, *, prefix: str) -> None: + super().__init__() + config = vllm_config.model_config.hf_config + self.proj = ReplicatedLinear( + config.hidden_size + config.dspark_markov_rank, + 1, + bias=False, + params_dtype=torch.float32, + quant_config=None, + return_bias=False, + prefix=f"{prefix}.proj", + ) + + def forward( + self, + hidden_states: torch.Tensor, + markov_embed: torch.Tensor, + ) -> torch.Tensor: + features = torch.cat([hidden_states, markov_embed], dim=-1) + return _linear_no_bias(self.proj, features.float()).squeeze(-1) + + +class DeepSeekV4DSparkLayer(nn.Module): + def __init__( + self, + vllm_config: VllmConfig, + *, + stage_id: int, + prefix: str, + ) -> None: + super().__init__() + config = vllm_config.model_config.hf_config + quant_config = vllm_config.quant_config + self.stage_id = stage_id + self.hidden_size = config.hidden_size + self.dtype = vllm_config.model_config.dtype + self.hc_mult = config.hc_mult + self.hc_dim = self.hc_mult * self.hidden_size + self.block_size = config.dspark_block_size + self.rms_norm_eps = config.rms_norm_eps + self.hc_head_op = HCHeadOp() + + if stage_id == 0: + self.main_proj = ReplicatedLinear( + config.hidden_size * len(config.dspark_target_layer_ids), + config.hidden_size, + bias=False, + quant_config=quant_config, + return_bias=False, + prefix=f"{prefix}.main_proj", + ) + self.main_norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + + self.attn = DeepSeekV4DSparkAttention( + vllm_config, + prefix=f"{prefix}.attn", + ) + self.ffn = DeepseekV4MoE(vllm_config, prefix=f"{prefix}.ffn") + self.attn_norm = RMSNorm(self.hidden_size, self.rms_norm_eps) + self.ffn_norm = RMSNorm(self.hidden_size, self.rms_norm_eps) + + # Reuse the target decoder's MHC kernels/parameters. + from vllm.model_executor.layers.mhc import MHCPostOp, MHCPreOp + + self.hc_sinkhorn_iters = config.hc_sinkhorn_iters + self.hc_eps = config.hc_eps + self.hc_post_alpha = 2.0 + mix_hc = (2 + self.hc_mult) * self.hc_mult + self.hc_attn_fn = nn.Parameter( + torch.empty((mix_hc, self.hc_dim), dtype=torch.float32), + requires_grad=False, + ) + self.hc_ffn_fn = nn.Parameter( + torch.empty((mix_hc, self.hc_dim), dtype=torch.float32), + requires_grad=False, + ) + self.hc_attn_base = nn.Parameter( + torch.empty(mix_hc, dtype=torch.float32), + requires_grad=False, + ) + self.hc_ffn_base = nn.Parameter( + torch.empty(mix_hc, dtype=torch.float32), + requires_grad=False, + ) + self.hc_attn_scale = nn.Parameter( + torch.empty(3, dtype=torch.float32), + requires_grad=False, + ) + self.hc_ffn_scale = nn.Parameter( + torch.empty(3, dtype=torch.float32), + requires_grad=False, + ) + self.mhc_pre = MHCPreOp() + self.mhc_post = MHCPostOp() + + if stage_id == config.dspark_num_draft_layers - 1: + self.norm = RMSNorm(config.hidden_size, config.rms_norm_eps) + self.markov_head = DeepSeekV4DSparkMarkovHead( + vllm_config, prefix=f"{prefix}.markov_head" + ) + self.confidence_head = DeepSeekV4DSparkConfidenceHead( + vllm_config, prefix=f"{prefix}.confidence_head" + ) + self.hc_head_fn = nn.Parameter( + torch.empty(self.hc_mult, self.hc_dim, dtype=torch.float32), + requires_grad=False, + ) + self.hc_head_base = nn.Parameter( + torch.empty(self.hc_mult, dtype=torch.float32), + requires_grad=False, + ) + self.hc_head_scale = nn.Parameter( + torch.empty(1, dtype=torch.float32), + requires_grad=False, + ) + + def project_main(self, main_hidden: torch.Tensor) -> torch.Tensor: + assert hasattr(self, "main_proj") + main_x = _linear_no_bias(self.main_proj, main_hidden) + return self.main_norm(main_x) + + def hc_pre( + self, + x: torch.Tensor, + hc_fn: torch.Tensor, + hc_scale: torch.Tensor, + hc_base: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + return self.mhc_pre( + residual=x, + fn=hc_fn, + hc_scale=hc_scale, + hc_base=hc_base, + rms_eps=self.rms_norm_eps, + hc_pre_eps=self.hc_eps, + hc_sinkhorn_eps=self.hc_eps, + hc_post_mult_value=self.hc_post_alpha, + sinkhorn_repeat=self.hc_sinkhorn_iters, + ) + + def hc_post( + self, + x: torch.Tensor, + residual: torch.Tensor, + post: torch.Tensor, + comb: torch.Tensor, + ) -> torch.Tensor: + return self.mhc_post(x, residual, post, comb) + + def forward_dspark( + self, + x: torch.Tensor, + positions: torch.Tensor, + input_ids: torch.Tensor, + *, + batch_size: int, + block_size: int, + main_x: torch.Tensor, + main_positions: torch.Tensor, + store_main_kv: bool = True, + slot_index: torch.Tensor | None = None, + ) -> torch.Tensor: + residual = x + attn_in, post, comb = unpack_mhc_pre_outputs( + self.hc_pre(x, self.hc_attn_fn, self.hc_attn_scale, self.hc_attn_base) + ) + attn_in = self.attn_norm(attn_in).to(self.dtype) + if attn_in.ndim == 3: + attn_in = attn_in.mean(dim=1).to(self.dtype) + attn_out = self.attn.forward_dspark( + attn_in, + positions, + batch_size=batch_size, + block_size=block_size, + main_x=main_x, + main_positions=main_positions, + store_main_kv=store_main_kv, + slot_index=slot_index, + ) + x = self.hc_post(attn_out.to(self.dtype), residual, post, comb).to(self.dtype) + + residual = x + ffn_in, post, comb = unpack_mhc_pre_outputs( + self.hc_pre(x, self.hc_ffn_fn, self.hc_ffn_scale, self.hc_ffn_base) + ) + ffn_in = self.ffn_norm(ffn_in).to(self.dtype) + ffn_out = self.ffn(ffn_in, input_ids) + return self.hc_post(ffn_out.to(self.dtype), residual, post, comb).to(self.dtype) + + def forward_head(self, x: torch.Tensor) -> torch.Tensor: + return self.hc_head_op( + x, + self.hc_head_fn, + self.hc_head_scale, + self.hc_head_base, + self.rms_norm_eps, + self.hc_eps, + ) + + +class DeepSeekV4DSparkModel(nn.Module): + def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: + super().__init__() + config = vllm_config.model_config.hf_config + self.config = config + self.block_size = config.dspark_block_size + self.noise_token_id = config.dspark_noise_token_id + self.num_draft_layers = config.dspark_num_draft_layers + self._local_argmax = _read_bool_env("VLLM_DSPARK_LOCAL_ARGMAX") + self._fused_markov_argmax = _read_bool_env( + "VLLM_DSPARK_FUSED_MARKOV_ARGMAX" + ) + self.dspark_start_layer_idx = max( + config.num_hidden_layers, + max(getattr(config, "dspark_target_layer_ids", [-1])) + 1, + ) + self.embed_tokens = VocabParallelEmbedding( + config.vocab_size, + config.hidden_size, + prefix=maybe_prefix(prefix, "embed_tokens"), + ) + self.stage_layer_keys = [ + str(self.dspark_start_layer_idx + stage_id) + for stage_id in range(self.num_draft_layers) + ] + self.layers = nn.ModuleDict( + { + layer_key: DeepSeekV4DSparkLayer( + vllm_config, + stage_id=stage_id, + prefix=maybe_prefix(prefix, f"layers.{layer_key}"), + ) + for stage_id, layer_key in enumerate(self.stage_layer_keys) + } + ) + if self._local_argmax: + logger.info( + "DSpark local vocab-parallel argmax is enabled. This is " + "experimental and may add per-position synchronization overhead." + ) + if self._fused_markov_argmax: + logger.info( + "DSpark fused Markov argmax is enabled. This keeps the paper's " + "low-rank Markov bias but avoids materializing Markov logits " + "on the greedy no-confidence draft path." + ) + + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: + return self.embed_tokens(input_ids) + + def project_main(self, main_hidden: torch.Tensor) -> torch.Tensor: + first = self.layers[self.stage_layer_keys[0]] + return first.project_main(main_hidden) + + def prefill_main( + self, + main_hidden: torch.Tensor, + main_positions: torch.Tensor, + num_rejected_tokens: torch.Tensor | None = None, + slot_index: torch.Tensor | None = None, + query_start_loc: list[int] | None = None, + ) -> None: + main_x = self.project_main(main_hidden.reshape(-1, main_hidden.shape[-1])) + if query_start_loc is None: + # Rectangular [B, seq, hidden_size] fast-path (uniform / static). + main_x = main_x.view(*main_hidden.shape[:-1], self.config.hidden_size) + # else: keep flat [total_rows, hidden_size] for the ragged path. + for layer in self.layers.values(): + layer.attn.store_main_kv( + main_x, + main_positions, + num_rejected_tokens=num_rejected_tokens, + slot_index=slot_index, + query_start_loc=query_start_loc, + ) + + def draft( + self, + input_ids: torch.Tensor, + main_hidden: torch.Tensor, + main_positions: torch.Tensor, + lm_head: ParallelLMHead, + logits_processor: LogitsProcessor, + *, + return_logits: bool = True, + return_confidence: bool = True, + store_main_kv: bool = True, + slot_index: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + batch_size = input_ids.shape[0] + block_size = self.block_size + main_positions = main_positions.view(batch_size, 1) + main_x = self.project_main(main_hidden).view( + batch_size, 1, self.config.hidden_size + ) + + draft_input_ids = input_ids.new_full( + (batch_size, block_size), self.noise_token_id + ) + draft_input_ids[:, 0] = input_ids + x = self.embed_tokens(draft_input_ids).view( + batch_size * block_size, self.config.hidden_size + ) + x = x.unsqueeze(1).repeat(1, self.config.hc_mult, 1) + + # DeepSpec trains/evaluates DSpark with the anchor token itself at the + # first draft position: anchor + [0, gamma). The hidden state at that + # anchor position predicts the next token. + offsets = torch.arange( + 0, + block_size, + dtype=main_positions.dtype, + device=main_positions.device, + ) + draft_positions = (main_positions[:, -1:] + offsets).reshape(-1) + flat_draft_input_ids = draft_input_ids.reshape(-1) + for layer in self.layers.values(): + x = layer.forward_dspark( + x, + draft_positions, + flat_draft_input_ids, + batch_size=batch_size, + block_size=block_size, + main_x=main_x, + main_positions=main_positions, + store_main_kv=store_main_kv, + slot_index=slot_index, + ) + + final_layer = self.layers[self.stage_layer_keys[-1]] + dense = final_layer.forward_head(x).view( + batch_size, block_size, self.config.hidden_size + ) + normed = final_layer.norm(dense.reshape(batch_size * block_size, -1)) + + if not return_logits and getattr(self, "_local_argmax", False): + local_logits = lm_head.quant_method.apply( + lm_head, + normed, + bias=None, + ).view(batch_size, block_size, -1) + output_ids = input_ids.new_empty(batch_size, block_size + 1) + output_ids[:, 0] = input_ids + markov_embeds = [] if return_confidence else None + for pos in range(block_size): + if markov_embeds is not None: + markov_logits, markov_embed = ( + final_layer.markov_head.forward_local(output_ids[:, pos]) + ) + markov_embeds.append(markov_embed) + step_logits = local_logits[:, pos] + markov_logits + output_ids[:, pos + 1] = _vocab_parallel_argmax( + step_logits, + lm_head, + ) + elif getattr(self, "_fused_markov_argmax", False): + markov_embed = final_layer.markov_head.markov_w1(output_ids[:, pos]) + output_ids[:, pos + 1] = _vocab_parallel_markov_argmax( + local_logits[:, pos], + markov_embed, + final_layer.markov_head.markov_w2, + lm_head, + ) + else: + markov_logits, _ = final_layer.markov_head.forward_local( + output_ids[:, pos] + ) + step_logits = local_logits[:, pos] + markov_logits + output_ids[:, pos + 1] = _vocab_parallel_argmax( + step_logits, + lm_head, + ) + logits = normed.new_empty((0, 0, 0)) + if return_confidence: + assert markov_embeds is not None + markov_embed = torch.stack(markov_embeds, dim=1) + confidence = final_layer.confidence_head(dense, markov_embed).sigmoid() + else: + confidence = dense.new_empty((batch_size, 0), dtype=torch.float32) + return output_ids[:, 1:], logits, confidence + + logits = logits_processor(lm_head, normed).view( + batch_size, block_size, self.config.vocab_size + ) + + output_ids = input_ids.new_empty(batch_size, block_size + 1) + output_ids[:, 0] = input_ids + markov_embeds = [] if return_confidence else None + for pos in range(block_size): + markov_logits, markov_embed = final_layer.markov_head(output_ids[:, pos]) + logits[:, pos].add_(markov_logits) + if markov_embeds is not None: + markov_embeds.append(markov_embed) + output_ids[:, pos + 1] = logits[:, pos].argmax(dim=-1) + + if return_confidence: + assert markov_embeds is not None + markov_embed = torch.stack(markov_embeds, dim=1) + confidence = final_layer.confidence_head(dense, markov_embed).sigmoid() + else: + confidence = dense.new_empty((batch_size, 0), dtype=torch.float32) + if not return_logits: + logits = logits.new_empty((0, 0, 0)) + return output_ids[:, 1:], logits, confidence + + def finalize_mega_moe_weights(self) -> None: + for layer in self.layers.values(): + if layer.ffn.use_mega_moe: + layer.ffn.finalize_mega_moe_weights() + + +class DeepSeekV4DSpark(nn.Module): + def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: + super().__init__() + self.config = vllm_config.speculative_config.draft_model_config.hf_config + self.model = DeepSeekV4DSparkModel( + vllm_config=vllm_config, + prefix=maybe_prefix(prefix, "model"), + ) + self.lm_head = ParallelLMHead( + self.config.vocab_size, + self.config.hidden_size, + prefix=maybe_prefix(prefix, "lm_head"), + ) + self.logits_processor = LogitsProcessor(self.config.vocab_size) + self._last_confidence: torch.Tensor | None = None + + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: + return self.model.embed_input_ids(input_ids) + + def prefill_main( + self, + main_hidden: torch.Tensor, + main_positions: torch.Tensor, + num_rejected_tokens: torch.Tensor | None = None, + slot_index: torch.Tensor | None = None, + query_start_loc: list[int] | None = None, + ) -> None: + self.model.prefill_main( + main_hidden, + main_positions, + num_rejected_tokens=num_rejected_tokens, + slot_index=slot_index, + query_start_loc=query_start_loc, + ) + + def draft( + self, + input_ids: torch.Tensor, + main_hidden: torch.Tensor, + main_positions: torch.Tensor, + *, + store_main_kv: bool = True, + slot_index: torch.Tensor | None = None, + ) -> torch.Tensor: + draft_ids, _logits, confidence = self.model.draft( + input_ids, + main_hidden, + main_positions, + self.lm_head, + self.logits_processor, + store_main_kv=store_main_kv, + slot_index=slot_index, + ) + self._last_confidence = confidence + return draft_ids + + def draft_with_confidence( + self, + input_ids: torch.Tensor, + main_hidden: torch.Tensor, + main_positions: torch.Tensor, + *, + return_logits: bool = True, + return_confidence: bool = True, + store_main_kv: bool = True, + slot_index: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + draft_ids, logits, confidence = self.model.draft( + input_ids, + main_hidden, + main_positions, + self.lm_head, + self.logits_processor, + return_logits=return_logits, + return_confidence=return_confidence, + store_main_kv=store_main_kv, + slot_index=slot_index, + ) + return draft_ids, logits, confidence + + def take_last_confidence(self) -> torch.Tensor | None: + confidence = self._last_confidence + self._last_confidence = None + return confidence + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + params_dict = dict(self.named_parameters()) + loaded_params: set[str] = set() + tp_size = get_tensor_model_parallel_world_size() + tp_rank = get_tensor_model_parallel_rank() + n_head = self.config.num_attention_heads + n_local_head = n_head // tp_size + head_rank_start = n_local_head * tp_rank + head_rank_end = n_local_head * (tp_rank + 1) + + first_layer = next(iter(self.model.layers.values())) + if first_layer.ffn.use_mega_moe: + expert_mapping = make_deepseek_v4_expert_params_mapping( + self.config.n_routed_experts + ) + else: + expert_mapping = FusedMoE.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, + ) + expert_scale_suffix = ( + ".weight_scale" + if getattr(self.config, "expert_dtype", "fp4") == "fp4" + else ".weight_scale_inv" + ) + + for name, loaded_weight in weights: + if not name.startswith("mtp."): + continue + stage_id = int(name.split(".", 2)[1]) + if stage_id >= self.config.dspark_num_draft_layers: + continue + virtual_layer_id = self.model.dspark_start_layer_idx + stage_id + name = name.replace( + f"mtp.{stage_id}.", f"model.layers.{virtual_layer_id}.", 1 + ) + name = name.replace(".shared_experts.w2", ".shared_experts.down_proj") + name = name.replace(".ffn.gate.bias", ".ffn.gate.e_score_correction_bias") + + if ".attn.attn_sink" in name: + param = params_dict[name] + param.data.copy_(loaded_weight[head_rank_start:head_rank_end]) + loaded_params.add(name) + continue + if name.endswith(".scale"): + suffix = ( + expert_scale_suffix + if _EXPERT_SCALE_RE.search(name) + else ".weight_scale_inv" + ) + name = name.removesuffix(".scale") + suffix + + mapped_stacked = map_dspark_stacked_param_name(name) + if mapped_stacked is not None: + mapped_name, shard_id = mapped_stacked + param = params_dict.get(mapped_name) + if param is None: + raise KeyError(mapped_name) + else: + weight_loader = param.weight_loader + weight_loader(param, loaded_weight, shard_id) + loaded_params.add(mapped_name) + continue + + if ".experts." in name: + if ( + "weight_scale" in name + and loaded_weight.dtype == torch.float8_e8m0fnu + ): + loaded_weight = loaded_weight.view(torch.uint8) + for mapping in expert_mapping: + param_name, weight_name, expert_id, expert_shard_id = mapping + if weight_name not in name: + continue + mapped_name = name.replace(weight_name, param_name) + param = params_dict[mapped_name] + weight_loader = typing.cast( + Callable[..., bool], param.weight_loader + ) + success = weight_loader( + param, + loaded_weight, + mapped_name, + shard_id=expert_shard_id, + expert_id=expert_id, + return_success=True, + ) + if success: + loaded_params.add(mapped_name) + break + continue + + param = params_dict.get(name) + if param is None: + logger.debug("Skipping unknown DSpark weight %s", name) + continue + weight_loader = getattr(param, "weight_loader", default_weight_loader) + weight_loader(param, loaded_weight) + loaded_params.add(name) + + self.model.finalize_mega_moe_weights() + return loaded_params diff --git a/mods/deepseek-v4-flash-dspark/overlay/vllm/models/deepseek_v4/nvidia/dspark_kernels.py b/mods/deepseek-v4-flash-dspark/overlay/vllm/models/deepseek_v4/nvidia/dspark_kernels.py new file mode 100644 index 00000000..d2d8c75a --- /dev/null +++ b/mods/deepseek-v4-flash-dspark/overlay/vllm/models/deepseek_v4/nvidia/dspark_kernels.py @@ -0,0 +1,822 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from __future__ import annotations + +import torch + +from vllm.triton_utils import HAS_TRITON, tl, triton + +_NEG_INF = -3.4028234663852886e38 +_DSPARK_SCORE_K_BLOCK = 8 +_FP8_E4M3_MAX = float(torch.finfo(torch.float8_e4m3fn).max) +_DSPARK_MARKOV_V_BLOCK = 256 +_DSPARK_MARKOV_R_BLOCK = 32 +_DSPARK_HC_POST_H_BLOCK = 64 + + +@triton.jit +def _dspark_markov_block_argmax_kernel( + base_logits_ptr, + markov_embed_ptr, + markov_w2_ptr, + block_vals_ptr, + block_indices_ptr, + batch_size, + local_vocab_size, + markov_rank: tl.constexpr, + num_pad, + base_stride_b, + base_stride_v, + embed_stride_b, + embed_stride_r, + w2_stride_v, + w2_stride_r, + out_stride_b, + out_stride_block, + V_BLOCK: tl.constexpr, + R_BLOCK: tl.constexpr, + NEG_INF: tl.constexpr, +): + pid_b = tl.program_id(0).to(tl.int64) + pid_block = tl.program_id(1).to(tl.int64) + + offs_v = pid_block * V_BLOCK + tl.arange(0, V_BLOCK) + valid_vocab = offs_v < local_vocab_size + padded_vocab = offs_v >= (local_vocab_size - num_pad) + valid_vocab = valid_vocab & ~padded_vocab + + acc = tl.zeros((V_BLOCK,), dtype=tl.float32) + for r_start in tl.static_range(0, markov_rank, R_BLOCK): + offs_r = r_start + tl.arange(0, R_BLOCK) + embed = tl.load( + markov_embed_ptr + pid_b * embed_stride_b + offs_r * embed_stride_r, + mask=offs_r < markov_rank, + other=0.0, + ).to(tl.float32) + w2 = tl.load( + markov_w2_ptr + + offs_v[None, :] * w2_stride_v + + offs_r[:, None] * w2_stride_r, + mask=(offs_v[None, :] < local_vocab_size) & (offs_r[:, None] < markov_rank), + other=0.0, + ).to(tl.float32) + acc += tl.sum(w2 * embed[:, None], axis=0) + + base = tl.load( + base_logits_ptr + pid_b * base_stride_b + offs_v * base_stride_v, + mask=offs_v < local_vocab_size, + other=NEG_INF, + ).to(tl.float32) + scores = tl.where(valid_vocab, base + acc, NEG_INF) + max_val = tl.max(scores, axis=0) + local_idx = tl.argmax(scores, axis=0) + token_idx = pid_block * V_BLOCK + local_idx + + tl.store( + block_vals_ptr + pid_b * out_stride_b + pid_block * out_stride_block, + max_val, + mask=pid_b < batch_size, + ) + tl.store( + block_indices_ptr + pid_b * out_stride_b + pid_block * out_stride_block, + token_idx, + mask=pid_b < batch_size, + ) + + +@triton.jit +def _dspark_quant_dequant_nope_kernel( + kv_ptr, + num_rows, + kv_stride_row, + kv_stride_d, + NOPE_DIM: tl.constexpr, + GROUP_SIZE: tl.constexpr, + FP8_MAX: tl.constexpr, + EPS: tl.constexpr, +): + row = tl.program_id(0).to(tl.int64) + group = tl.program_id(1).to(tl.int64) + + offsets = group * GROUP_SIZE + tl.arange(0, GROUP_SIZE) + mask = (row < num_rows) & (offsets < NOPE_DIM) + vals = tl.load( + kv_ptr + row * kv_stride_row + offsets * kv_stride_d, + mask=mask, + other=0.0, + ).to(tl.float32) + abs_vals = tl.where(offsets < NOPE_DIM, tl.abs(vals), 0.0) + amax = tl.maximum(tl.max(abs_vals, axis=0), EPS) + scale = tl.math.exp2(tl.ceil(tl.log2(amax * (1.0 / FP8_MAX)))) + quantized = tl.clamp(vals * (1.0 / scale), -FP8_MAX, FP8_MAX).to(tl.float8e4nv) + dequantized = quantized.to(tl.float32) * scale + tl.store( + kv_ptr + row * kv_stride_row + offsets * kv_stride_d, + dequantized, + mask=mask, + ) + + +@triton.jit +def _dspark_sparse_scores_kernel( + q_ptr, + draft_kv_ptr, + main_kv_ptr, + valid_main_lengths_ptr, + scores_ptr, + softmax_scale: tl.constexpr, + q_stride_b, + q_stride_q, + q_stride_h, + q_stride_d, + draft_stride_b, + draft_stride_k, + draft_stride_d, + main_stride_b, + main_stride_k, + main_stride_d, + scores_stride_b, + scores_stride_q, + scores_stride_h, + scores_stride_k, + BLOCK_SIZE: tl.constexpr, + NUM_HEADS: tl.constexpr, + HEAD_DIM: tl.constexpr, + WINDOW_SIZE: tl.constexpr, + KV_TOKENS: tl.constexpr, + K_BLOCK: tl.constexpr, + D_BLOCK: tl.constexpr, + NEG_INF: tl.constexpr, +): + pid_bqh = tl.program_id(0).to(tl.int64) + pid_k = tl.program_id(1).to(tl.int64) + + h = pid_bqh % NUM_HEADS + tmp = pid_bqh // NUM_HEADS + q_idx = tmp % BLOCK_SIZE + batch_idx = tmp // BLOCK_SIZE + + offs_k = pid_k * K_BLOCK + tl.arange(0, K_BLOCK) + valid_main_len = tl.load(valid_main_lengths_ptr + batch_idx).to(tl.int64) + is_main = offs_k < WINDOW_SIZE + is_draft = (offs_k >= WINDOW_SIZE) & (offs_k < KV_TOKENS) + is_valid = (is_main & (offs_k < valid_main_len)) | is_draft + + acc = tl.zeros((K_BLOCK,), dtype=tl.float32) + for d_start in tl.static_range(0, HEAD_DIM, D_BLOCK): + offs_d = d_start + tl.arange(0, D_BLOCK) + q_vals = tl.load( + q_ptr + + batch_idx * q_stride_b + + q_idx * q_stride_q + + h * q_stride_h + + offs_d * q_stride_d + ).to(tl.float32) + + main_vals = tl.load( + main_kv_ptr + + batch_idx * main_stride_b + + offs_k[:, None] * main_stride_k + + offs_d[None, :] * main_stride_d, + mask=(offs_k[:, None] < WINDOW_SIZE), + other=0.0, + ) + draft_k = offs_k - WINDOW_SIZE + draft_vals = tl.load( + draft_kv_ptr + + batch_idx * draft_stride_b + + draft_k[:, None] * draft_stride_k + + offs_d[None, :] * draft_stride_d, + mask=(draft_k[:, None] >= 0) & (draft_k[:, None] < BLOCK_SIZE), + other=0.0, + ) + kv_vals = tl.where(is_main[:, None], main_vals, draft_vals).to(tl.float32) + acc += tl.sum(kv_vals * q_vals[None, :], axis=1) + + scores = acc * softmax_scale + scores = tl.where(is_valid, scores, NEG_INF) + tl.store( + scores_ptr + + batch_idx * scores_stride_b + + q_idx * scores_stride_q + + h * scores_stride_h + + offs_k * scores_stride_k, + scores, + mask=offs_k < KV_TOKENS, + ) + + +@triton.jit +def _dspark_sparse_out_kernel( + scores_ptr, + draft_kv_ptr, + main_kv_ptr, + attn_sink_ptr, + out_ptr, + draft_stride_b, + draft_stride_k, + draft_stride_d, + main_stride_b, + main_stride_k, + main_stride_d, + scores_stride_b, + scores_stride_q, + scores_stride_h, + scores_stride_k, + out_stride_b, + out_stride_q, + out_stride_h, + out_stride_d, + BLOCK_SIZE: tl.constexpr, + NUM_HEADS: tl.constexpr, + HEAD_DIM: tl.constexpr, + WINDOW_SIZE: tl.constexpr, + KV_TOKENS: tl.constexpr, + K_BLOCK: tl.constexpr, + D_BLOCK: tl.constexpr, + NEG_INF: tl.constexpr, +): + pid_bqh = tl.program_id(0).to(tl.int64) + pid_d = tl.program_id(1).to(tl.int64) + + h = pid_bqh % NUM_HEADS + tmp = pid_bqh // NUM_HEADS + q_idx = tmp % BLOCK_SIZE + batch_idx = tmp // BLOCK_SIZE + + offs_k = tl.arange(0, K_BLOCK) + scores = tl.load( + scores_ptr + + batch_idx * scores_stride_b + + q_idx * scores_stride_q + + h * scores_stride_h + + offs_k * scores_stride_k, + mask=offs_k < KV_TOKENS, + other=NEG_INF, + ).to(tl.float32) + + sink = tl.load(attn_sink_ptr + h).to(tl.float32) + normalizer = tl.maximum(tl.max(scores, axis=0), sink) + weights = tl.exp(scores - normalizer) + denom = tl.sum(weights, axis=0) + tl.exp(sink - normalizer) + + offs_d = pid_d * D_BLOCK + tl.arange(0, D_BLOCK) + main_vals = tl.load( + main_kv_ptr + + batch_idx * main_stride_b + + offs_k[:, None] * main_stride_k + + offs_d[None, :] * main_stride_d, + mask=(offs_k[:, None] < WINDOW_SIZE) & (offs_d[None, :] < HEAD_DIM), + other=0.0, + ) + draft_k = offs_k - WINDOW_SIZE + draft_vals = tl.load( + draft_kv_ptr + + batch_idx * draft_stride_b + + draft_k[:, None] * draft_stride_k + + offs_d[None, :] * draft_stride_d, + mask=( + (draft_k[:, None] >= 0) + & (draft_k[:, None] < BLOCK_SIZE) + & (offs_d[None, :] < HEAD_DIM) + ), + other=0.0, + ) + vals = tl.where((offs_k < WINDOW_SIZE)[:, None], main_vals, draft_vals).to( + tl.float32 + ) + out = tl.sum(weights[:, None] * vals, axis=0) / denom + tl.store( + out_ptr + + batch_idx * out_stride_b + + q_idx * out_stride_q + + h * out_stride_h + + offs_d * out_stride_d, + out, + mask=offs_d < HEAD_DIM, + ) + + +@triton.jit +def _dspark_hc_post_mean_kernel( + x_ptr, + residual_ptr, + post_ptr, + comb_ptr, + out_ptr, + num_tokens, + hidden_size: tl.constexpr, + hc_mult: tl.constexpr, + x_stride_t, + x_stride_h, + residual_stride_t, + residual_stride_c, + residual_stride_h, + post_stride_t, + post_stride_c, + comb_stride_t, + comb_stride_i, + comb_stride_j, + out_stride_t, + out_stride_h, + H_BLOCK: tl.constexpr, +): + token = tl.program_id(0).to(tl.int64) + h_start = tl.program_id(1).to(tl.int64) * H_BLOCK + offs_h = h_start + tl.arange(0, H_BLOCK) + mask_h = (token < num_tokens) & (offs_h < hidden_size) + + inv_hc = 1.0 / hc_mult + post_sum = tl.full((), 0.0, dtype=tl.float32) + for j in tl.static_range(0, hc_mult): + post_sum += tl.load( + post_ptr + token * post_stride_t + j * post_stride_c, + mask=token < num_tokens, + other=0.0, + ).to(tl.float32) + + x_vals = tl.load( + x_ptr + token * x_stride_t + offs_h * x_stride_h, + mask=mask_h, + other=0.0, + ).to(tl.float32) + acc = x_vals * post_sum * inv_hc + + for i in tl.static_range(0, hc_mult): + comb_sum = tl.full((), 0.0, dtype=tl.float32) + for j in tl.static_range(0, hc_mult): + comb_sum += tl.load( + comb_ptr + + token * comb_stride_t + + i * comb_stride_i + + j * comb_stride_j, + mask=token < num_tokens, + other=0.0, + ).to(tl.float32) + residual_vals = tl.load( + residual_ptr + + token * residual_stride_t + + i * residual_stride_c + + offs_h * residual_stride_h, + mask=mask_h, + other=0.0, + ).to(tl.float32) + acc += residual_vals * comb_sum * inv_hc + + tl.store( + out_ptr + token * out_stride_t + offs_h * out_stride_h, + acc, + mask=mask_h, + ) + + +def _next_power_of_2(value: int) -> int: + return 1 << (value - 1).bit_length() + + +def _post_strides(post: torch.Tensor) -> tuple[int, int]: + if post.ndim == 2: + return post.stride(0), post.stride(1) + if post.ndim == 3 and post.shape[-1] == 1: + return post.stride(0), post.stride(1) + raise ValueError(f"Unsupported DSpark post mix shape: {tuple(post.shape)}") + + +def _post_2d(post: torch.Tensor) -> torch.Tensor: + if post.ndim == 2: + return post + if post.ndim == 3 and post.shape[-1] == 1: + return post.squeeze(-1) + raise ValueError(f"Unsupported DSpark post mix shape: {tuple(post.shape)}") + + +def dspark_hc_post_mean_torch( + x: torch.Tensor, + residual: torch.Tensor, + post: torch.Tensor, + comb: torch.Tensor, +) -> torch.Tensor: + """Reference ``MHCPostOp(...).mean(dim=1)`` without materializing streams.""" + + post_2d = _post_2d(post).to(torch.float32) + residual_f = residual.to(torch.float32) + comb_f = comb.to(torch.float32) + x_f = x.to(torch.float32) + hc_mult = residual.shape[1] + post_term = post_2d.sum(dim=1, keepdim=True) * (1.0 / hc_mult) * x_f + residual_weights = comb_f.sum(dim=-1) * (1.0 / hc_mult) + residual_term = torch.sum(residual_weights.unsqueeze(-1) * residual_f, dim=1) + return (post_term + residual_term).to(residual.dtype) + + +def dspark_hc_post_mean( + x: torch.Tensor, + residual: torch.Tensor, + post: torch.Tensor, + comb: torch.Tensor, + out: torch.Tensor | None = None, +) -> torch.Tensor: + """Compute the DSpark target feature ``MHCPostOp(...).mean(dim=1)``. + + This is used by deferred DSpark target-layer capture. The target forward + still uses the fused post/pre path, while this kernel writes only the + reduced hidden feature required by the DSpark drafter. + """ + + if out is None: + out = torch.empty( + (x.shape[0], x.shape[-1]), + dtype=residual.dtype, + device=x.device, + ) + + if ( + not x.is_cuda + or not residual.is_cuda + or not post.is_cuda + or not comb.is_cuda + or not out.is_cuda + or not HAS_TRITON + ): + out.copy_(dspark_hc_post_mean_torch(x, residual, post, comb)) + return out + + if residual.ndim != 3 or x.ndim != 2 or comb.ndim != 3: + out.copy_(dspark_hc_post_mean_torch(x, residual, post, comb)) + return out + if x.shape[0] != residual.shape[0] or comb.shape[:2] != residual.shape[:2]: + out.copy_(dspark_hc_post_mean_torch(x, residual, post, comb)) + return out + if comb.shape[2] != residual.shape[1] or x.shape[1] != residual.shape[2]: + out.copy_(dspark_hc_post_mean_torch(x, residual, post, comb)) + return out + if out.shape != x.shape: + out.copy_(dspark_hc_post_mean_torch(x, residual, post, comb)) + return out + + num_tokens, hidden_size = x.shape + hc_mult = residual.shape[1] + if hc_mult <= 0 or hidden_size <= 0: + return out + + post_stride_t, post_stride_c = _post_strides(post) + grid = (num_tokens, triton.cdiv(hidden_size, _DSPARK_HC_POST_H_BLOCK)) + _dspark_hc_post_mean_kernel[grid]( + x, + residual, + post, + comb, + out, + num_tokens, + hidden_size, + hc_mult, + x.stride(0), + x.stride(1), + residual.stride(0), + residual.stride(1), + residual.stride(2), + post_stride_t, + post_stride_c, + comb.stride(0), + comb.stride(1), + comb.stride(2), + out.stride(0), + out.stride(1), + H_BLOCK=_DSPARK_HC_POST_H_BLOCK, + num_warps=4, + num_stages=4, + ) + return out + + +def dspark_markov_argmax_torch( + base_logits: torch.Tensor, + markov_embed: torch.Tensor, + markov_w2_weight: torch.Tensor, + *, + num_pad: int = 0, +) -> tuple[torch.Tensor, torch.Tensor]: + """Reference local top-1 for ``base_logits + W1[token] @ W2``.""" + + scores = base_logits.float() + torch.matmul( + markov_embed.float(), markov_w2_weight.float().t() + ) + if num_pad > 0: + scores[..., -num_pad:] = -float("inf") + local_max_vals, local_max_indices = scores.max(dim=-1) + return local_max_vals, local_max_indices.to(torch.long) + + +def dspark_markov_argmax( + base_logits: torch.Tensor, + markov_embed: torch.Tensor, + markov_w2_weight: torch.Tensor, + *, + num_pad: int = 0, +) -> tuple[torch.Tensor, torch.Tensor]: + """Fused local top-1 for the DSpark Markov-head greedy step. + + This computes only the local maximum of the Markov-corrected logits instead + of materializing the full Markov-logit vector. The caller still performs the + tensor-parallel global top-1 reduction. + """ + + if ( + not base_logits.is_cuda + or not markov_embed.is_cuda + or not markov_w2_weight.is_cuda + or not HAS_TRITON + ): + return dspark_markov_argmax_torch( + base_logits, + markov_embed, + markov_w2_weight, + num_pad=num_pad, + ) + + if base_logits.ndim != 2 or markov_embed.ndim != 2 or markov_w2_weight.ndim != 2: + return dspark_markov_argmax_torch( + base_logits, + markov_embed, + markov_w2_weight, + num_pad=num_pad, + ) + + batch_size, local_vocab_size = base_logits.shape + markov_rank = markov_embed.shape[-1] + if markov_w2_weight.shape != (local_vocab_size, markov_rank): + return dspark_markov_argmax_torch( + base_logits, + markov_embed, + markov_w2_weight, + num_pad=num_pad, + ) + if markov_rank % _DSPARK_MARKOV_R_BLOCK != 0: + return dspark_markov_argmax_torch( + base_logits, + markov_embed, + markov_w2_weight, + num_pad=num_pad, + ) + + num_blocks = triton.cdiv(local_vocab_size, _DSPARK_MARKOV_V_BLOCK) + block_vals = torch.empty( + (batch_size, num_blocks), + device=base_logits.device, + dtype=torch.float32, + ) + block_indices = torch.empty( + (batch_size, num_blocks), + device=base_logits.device, + dtype=torch.int64, + ) + _dspark_markov_block_argmax_kernel[(batch_size, num_blocks)]( + base_logits, + markov_embed, + markov_w2_weight, + block_vals, + block_indices, + batch_size, + local_vocab_size, + markov_rank, + int(num_pad), + base_logits.stride(0), + base_logits.stride(1), + markov_embed.stride(0), + markov_embed.stride(1), + markov_w2_weight.stride(0), + markov_w2_weight.stride(1), + block_vals.stride(0), + block_vals.stride(1), + V_BLOCK=_DSPARK_MARKOV_V_BLOCK, + R_BLOCK=_DSPARK_MARKOV_R_BLOCK, + NEG_INF=_NEG_INF, + num_warps=8, + num_stages=4, + ) + max_block = block_vals.argmax(dim=-1, keepdim=True) + local_max_vals = block_vals.gather(dim=-1, index=max_block).squeeze(-1) + local_max_indices = block_indices.gather(dim=-1, index=max_block).squeeze(-1) + return local_max_vals, local_max_indices + + +def dspark_quant_dequant_nope_torch( + kv: torch.Tensor, + rope_dim: int, + group_size: int = 64, +) -> torch.Tensor: + """Reference in-place FP8 quant-dequant for DSpark no-RoPE KV dims.""" + + head_dim = kv.shape[-1] + nope_dim = head_dim - rope_dim + assert nope_dim >= 0 + if nope_dim == 0: + return kv + assert nope_dim % group_size == 0 + + nope = kv[..., :nope_dim] + original_shape = nope.shape + groups = nope.reshape(-1, nope_dim // group_size, group_size).float() + amax = groups.abs().amax(dim=-1, keepdim=True).clamp_min_(1.0e-4) + scale = torch.pow( + torch.full((), 2.0, device=kv.device, dtype=torch.float32), + torch.ceil(torch.log2(amax / _FP8_E4M3_MAX)), + ) + quantized = torch.clamp( + groups / scale, + min=-_FP8_E4M3_MAX, + max=_FP8_E4M3_MAX, + ).to(torch.float8_e4m3fn) + nope.copy_((quantized.float() * scale).reshape(original_shape).to(kv.dtype)) + return kv + + +def dspark_quant_dequant_nope( + kv: torch.Tensor, + rope_dim: int, + group_size: int = 64, +) -> torch.Tensor: + """In-place reference-parity FP8 quant-dequant for no-RoPE KV dims. + + The released DeepSeek V4 Flash path applies `act_quant(..., inplace=True)` + after KV norm and RoPE, but only to the non-RoPE dimensions. This helper + mirrors that QAT simulation while keeping the RoPE slice in bf16. + """ + + head_dim = kv.shape[-1] + nope_dim = head_dim - rope_dim + assert nope_dim >= 0 + if nope_dim == 0: + return kv + assert nope_dim % group_size == 0 + + if not kv.is_cuda or not HAS_TRITON: + return dspark_quant_dequant_nope_torch(kv, rope_dim, group_size) + + if not kv.is_contiguous(): + return dspark_quant_dequant_nope_torch(kv, rope_dim, group_size) + + flat = kv.view(-1, head_dim) + grid = (flat.shape[0], triton.cdiv(nope_dim, group_size)) + _dspark_quant_dequant_nope_kernel[grid]( + flat, + flat.shape[0], + flat.stride(0), + flat.stride(1), + NOPE_DIM=nope_dim, + GROUP_SIZE=group_size, + FP8_MAX=_FP8_E4M3_MAX, + EPS=1.0e-4, + num_warps=2, + num_stages=4, + ) + return kv + + +def dspark_sparse_attention_torch( + q: torch.Tensor, + draft_kv: torch.Tensor, + main_kv_cache: torch.Tensor, + valid_main_lengths: torch.Tensor, + attn_sink: torch.Tensor, + softmax_scale: float, +) -> torch.Tensor: + """Reference DSpark sparse attention matching the CUDA kernel contract.""" + + batch_size, block_size, num_heads, head_dim = q.shape + window_size = main_kv_cache.shape[1] + main_kv = main_kv_cache[:batch_size] + kv = torch.cat([main_kv, draft_kv], dim=1) + kv_tokens = window_size + block_size + + kv_idx = torch.arange(kv_tokens, device=q.device) + valid_main = kv_idx.unsqueeze(0) < valid_main_lengths.to(torch.long).unsqueeze(1) + valid = torch.where( + kv_idx.unsqueeze(0) < window_size, + valid_main, + torch.ones((batch_size, kv_tokens), dtype=torch.bool, device=q.device), + ) + + scores = torch.einsum("bqhd,bkd->bqhk", q.float(), kv.float()) + scores.mul_(softmax_scale) + scores.masked_fill_(~valid[:, None, None, :], _NEG_INF) + normalizer = torch.maximum( + scores.max(dim=-1, keepdim=True).values, + attn_sink[:num_heads].view(1, 1, num_heads, 1), + ) + weights = torch.exp(scores - normalizer) + denom = weights.sum(dim=-1, keepdim=True) + torch.exp( + attn_sink[:num_heads].view(1, 1, num_heads, 1) - normalizer + ) + out = torch.einsum("bqhk,bkd->bqhd", weights.to(kv.dtype), kv) / denom.to(kv.dtype) + return out.reshape(batch_size * block_size, num_heads, head_dim) + + +def dspark_sparse_attention( + q: torch.Tensor, + draft_kv: torch.Tensor, + main_kv_cache: torch.Tensor, + valid_main_lengths: torch.Tensor, + attn_sink: torch.Tensor, + softmax_scale: float, + scores_buffer: torch.Tensor, +) -> torch.Tensor: + """Run DSpark sparse attention with a Triton CUDA kernel when available.""" + + if not q.is_cuda or not HAS_TRITON: + return dspark_sparse_attention_torch( + q, + draft_kv, + main_kv_cache, + valid_main_lengths, + attn_sink, + softmax_scale, + ) + + batch_size, block_size, num_heads, head_dim = q.shape + window_size = main_kv_cache.shape[1] + kv_tokens = window_size + block_size + assert scores_buffer.shape[:4] == (batch_size, block_size, num_heads, kv_tokens) + assert head_dim % 64 == 0 + + scores = scores_buffer + out = torch.empty_like(q) + k_score_block = _DSPARK_SCORE_K_BLOCK + k_out_block = _next_power_of_2(kv_tokens) + d_score_block = 64 + d_out_block = 32 + + grid_scores = ( + batch_size * block_size * num_heads, + triton.cdiv(kv_tokens, k_score_block), + ) + _dspark_sparse_scores_kernel[grid_scores]( + q, + draft_kv, + main_kv_cache, + valid_main_lengths, + scores, + softmax_scale, + q.stride(0), + q.stride(1), + q.stride(2), + q.stride(3), + draft_kv.stride(0), + draft_kv.stride(1), + draft_kv.stride(2), + main_kv_cache.stride(0), + main_kv_cache.stride(1), + main_kv_cache.stride(2), + scores.stride(0), + scores.stride(1), + scores.stride(2), + scores.stride(3), + BLOCK_SIZE=block_size, + NUM_HEADS=num_heads, + HEAD_DIM=head_dim, + WINDOW_SIZE=window_size, + KV_TOKENS=kv_tokens, + K_BLOCK=k_score_block, + D_BLOCK=d_score_block, + NEG_INF=_NEG_INF, + num_warps=2, + num_stages=4, + ) + + grid_out = ( + batch_size * block_size * num_heads, + triton.cdiv(head_dim, d_out_block), + ) + _dspark_sparse_out_kernel[grid_out]( + scores, + draft_kv, + main_kv_cache, + attn_sink, + out, + draft_kv.stride(0), + draft_kv.stride(1), + draft_kv.stride(2), + main_kv_cache.stride(0), + main_kv_cache.stride(1), + main_kv_cache.stride(2), + scores.stride(0), + scores.stride(1), + scores.stride(2), + scores.stride(3), + out.stride(0), + out.stride(1), + out.stride(2), + out.stride(3), + BLOCK_SIZE=block_size, + NUM_HEADS=num_heads, + HEAD_DIM=head_dim, + WINDOW_SIZE=window_size, + KV_TOKENS=kv_tokens, + K_BLOCK=k_out_block, + D_BLOCK=d_out_block, + NEG_INF=_NEG_INF, + num_warps=8, + num_stages=4, + ) + return out.reshape(batch_size * block_size, num_heads, head_dim) diff --git a/mods/deepseek-v4-flash-dspark/overlay/vllm/models/deepseek_v4/nvidia/model.py b/mods/deepseek-v4-flash-dspark/overlay/vllm/models/deepseek_v4/nvidia/model.py new file mode 100644 index 00000000..4ce275ea --- /dev/null +++ b/mods/deepseek-v4-flash-dspark/overlay/vllm/models/deepseek_v4/nvidia/model.py @@ -0,0 +1,2187 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import os +import time +import typing +from collections import defaultdict +from collections.abc import Callable, Iterable +from itertools import islice + +import regex as re +import torch +import torch.nn as nn + +import vllm.envs as envs +from vllm.compilation.decorators import support_torch_compile +from vllm.config import VllmConfig +from vllm.distributed import ( + get_ep_group, + get_pp_group, + get_tensor_model_parallel_rank, + get_tensor_model_parallel_world_size, +) +from vllm.forward_context import get_forward_context, is_forward_context_available +from vllm.logger import init_logger +from vllm.model_executor.layers.activation import SiluAndMul, SiluAndMulWithClamp +from vllm.model_executor.layers.fused_moe import FusedMoE +from vllm.model_executor.layers.fused_moe.router.fused_topk_bias_router import ( + fused_topk_bias, +) +from vllm.model_executor.layers.fused_moe.router.gate_linear import GateLinear +from vllm.model_executor.layers.layernorm import RMSNorm +from vllm.model_executor.layers.linear import ( + ColumnParallelLinear, + MergedColumnParallelLinear, + RowParallelLinear, +) +from vllm.model_executor.layers.logits_processor import LogitsProcessor +from vllm.model_executor.layers.mhc import ( + HCHeadOp, +) +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 ( + ParallelLMHead, + VocabParallelEmbedding, +) +from vllm.model_executor.model_loader.weight_utils import default_weight_loader +from vllm.model_executor.models.interfaces import SupportsPP +from vllm.model_executor.models.utils import ( + AutoWeightsLoader, + PPMissingLayer, + WeightsMapper, + extract_layer_index, + is_pp_missing_parameter, + make_layers, + maybe_prefix, +) +from vllm.model_executor.utils import set_weight_attrs +from vllm.models.deepseek_v4.attention import ( + DeepseekV4Indexer, + DeepseekV4MLAModules, + DeepseekV4MultiHeadLatentAttentionWrapper, + get_deepseek_v4_padded_num_q_heads, +) +from vllm.models.deepseek_v4.nvidia.ops import prepare_megamoe_inputs +from vllm.platforms import current_platform +from vllm.sequence import IntermediateTensors +from vllm.utils.torch_utils import direct_register_custom_op + +logger = init_logger(__name__) + + +def _env_flag(*names: str) -> bool: + return any( + os.environ.get(name, "").lower() not in ("", "0", "false", "no") + for name in names + ) + + +_B12X_MHC_TRACE = _env_flag("B12X_TRACE_MHC", "VLLM_TRACE_B12X_MHC") +_B12X_MHC_TRACE_LIMIT = int( + os.environ.get( + "B12X_TRACE_MHC_LIMIT", + os.environ.get("VLLM_TRACE_B12X_MHC_LIMIT", "240"), + ) +) +_B12X_MHC_TRACE_COUNTS: dict[str, int] = {} + + +def _env_enabled(name: str, default: str = "0") -> bool: + return os.environ.get(name, default).strip().lower() not in ( + "", + "0", + "false", + "no", + ) + + +_DSPARK_DEFER_TARGET_CAPTURE = _env_enabled( + "VLLM_DSV4_DSPARK_DEFER_TARGET_CAPTURE" +) +_DSPARK_DEFER_TARGET_CAPTURE_EXACT = _env_enabled( + "VLLM_DSV4_DSPARK_DEFER_TARGET_CAPTURE_EXACT" +) +_DSPARK_TARGET_TIMING = _env_enabled("VLLM_DSPARK_TARGET_TIMING") +_DSPARK_TARGET_TIMING_LOG_EVERY = int( + os.environ.get("VLLM_DSPARK_TARGET_TIMING_LOG_EVERY", "20") +) +_DSPARK_TARGET_TIMING_TOTALS: defaultdict[str, float] = defaultdict(float) +_DSPARK_TARGET_TIMING_COUNTS: defaultdict[str, int] = defaultdict(int) +_DSPARK_TARGET_TIMING_FORWARDS = 0 + + +def _dspark_target_timing_active() -> bool: + if not _DSPARK_TARGET_TIMING: + return False + compiler = getattr(torch, "compiler", None) + is_compiling = getattr(compiler, "is_compiling", None) + if is_compiling is not None and is_compiling(): + return False + return True + + +def _dspark_target_timing_start() -> float: + if not _dspark_target_timing_active(): + return 0.0 + if current_platform.is_cuda(): + torch.cuda.synchronize() + return time.perf_counter() + + +def _dspark_target_timing_record(stage: str, started: float) -> None: + if not _dspark_target_timing_active() or started == 0.0: + return + if current_platform.is_cuda(): + torch.cuda.synchronize() + _DSPARK_TARGET_TIMING_TOTALS[stage] += (time.perf_counter() - started) * 1000.0 + _DSPARK_TARGET_TIMING_COUNTS[stage] += 1 + + +def _dspark_target_timing_finish(total_started: float, tokens: int, layers: int) -> None: + global _DSPARK_TARGET_TIMING_FORWARDS + + if not _dspark_target_timing_active(): + return + _dspark_target_timing_record("forward_total", total_started) + _DSPARK_TARGET_TIMING_FORWARDS += 1 + every = max(1, _DSPARK_TARGET_TIMING_LOG_EVERY) + if _DSPARK_TARGET_TIMING_FORWARDS % every != 0: + return + + forwards = max(1, _DSPARK_TARGET_TIMING_FORWARDS) + stages = ( + "embed_or_input", + "layers_total", + "layer_attn_mhc", + "layer_attn", + "layer_ffn_mhc", + "layer_ffn", + "dspark_capture", + "final_hc_post", + "mtp_hidden_copy", + "hc_head_norm", + "forward_total", + ) + avg = { + stage: _DSPARK_TARGET_TIMING_TOTALS.get(stage, 0.0) / forwards + for stage in stages + } + logger.info( + "DSpark target timing forwards=%d last_tokens=%d layers=%d avg_ms=%s", + _DSPARK_TARGET_TIMING_FORWARDS, + tokens, + layers, + ", ".join(f"{stage}:{avg[stage]:.3f}" for stage in stages), + ) + + +def _trace_b12x_mhc_call( + op_name: str, + layer_name: str, + tokens: int, + run: Callable[[], typing.Any], +) -> typing.Any: + if not _B12X_MHC_TRACE: + return run() + + call_count = _B12X_MHC_TRACE_COUNTS.get(op_name, 0) + 1 + _B12X_MHC_TRACE_COUNTS[op_name] = call_count + if call_count > _B12X_MHC_TRACE_LIMIT: + if call_count == _B12X_MHC_TRACE_LIMIT + 1: + logger.info( + "b12x mHC trace limit reached for %s at %d calls.", + op_name, + _B12X_MHC_TRACE_LIMIT, + ) + return run() + + is_capturing = False + is_current_stream_capturing = getattr( + torch.cuda, "is_current_stream_capturing", None + ) + if is_current_stream_capturing is not None: + try: + is_capturing = bool(is_current_stream_capturing()) + except Exception: + is_capturing = False + + started = time.perf_counter() + try: + return run() + finally: + logger.info( + "b12x mHC %s call=%d layer=%s tokens=%d capturing=%s elapsed=%.3f ms", + op_name, + call_count, + layer_name, + tokens, + is_capturing, + (time.perf_counter() - started) * 1000.0, + ) + + +def _use_b12x_mhc() -> bool: + if not envs.VLLM_USE_B12X_MHC: + return False + if not current_platform.is_cuda(): + raise RuntimeError("VLLM_USE_B12X_MHC requires CUDA.") + if not current_platform.is_device_capability_family(120): + raise RuntimeError("VLLM_USE_B12X_MHC currently requires an SM120 GPU.") + return True + + +def _b12x_mhc_max_tokens() -> int: + raw = os.environ.get("B12X_MHC_MAX_TOKENS", "16") + try: + return int(raw) + except ValueError as exc: + raise RuntimeError( + f"B12X_MHC_MAX_TOKENS must be an integer, got {raw!r}" + ) from exc + + +def _empty_b12x_plan_scratch( + plan: object, + device: torch.device, +) -> torch.Tensor | tuple[torch.Tensor, ...]: + specs = plan.shapes_and_dtypes() + if not specs: + raise ValueError("b12x scratch plan did not provide any scratch specs") + buffers = tuple( + torch.empty(shape, dtype=dtype, device=device) for shape, dtype in specs + ) + if len(buffers) == 1: + return buffers[0] + return buffers + + +class DeepseekV4MLP(nn.Module): + def __init__( + self, + hidden_size: int, + intermediate_size: int, + hidden_act: str, + swiglu_limit: float | None = None, + quant_config: QuantizationConfig | None = None, + reduce_results: bool = True, + is_sequence_parallel: bool = False, + prefix: str = "", + ) -> None: + super().__init__() + + # If is_sequence_parallel, the input and output tensors are sharded + # across the ranks within the tp_group. In this case the weights are + # replicated and no collective ops are needed. + # Otherwise we use standard TP with an allreduce at the end. + self.gate_up_proj = MergedColumnParallelLinear( + hidden_size, + [intermediate_size] * 2, + bias=False, + quant_config=quant_config, + disable_tp=is_sequence_parallel, + prefix=f"{prefix}.gate_up_proj", + ) + self.down_proj = RowParallelLinear( + intermediate_size, + hidden_size, + bias=False, + quant_config=quant_config, + reduce_results=reduce_results, + disable_tp=is_sequence_parallel, + prefix=f"{prefix}.down_proj", + ) + if hidden_act != "silu": + raise ValueError( + f"Unsupported activation: {hidden_act}. Only silu is supported for now." + ) + if swiglu_limit is not None: + self.act_fn = SiluAndMulWithClamp(swiglu_limit) + else: + self.act_fn = SiluAndMul() + + def forward(self, x): + gate_up, _ = self.gate_up_proj(x) + x = self.act_fn(gate_up) + x, _ = self.down_proj(x) + return x + + +def make_deepseek_v4_expert_params_mapping( + num_experts: int, +) -> list[tuple[str, str, int, str]]: + return [ + ( + "experts.w13_" if shard_id in ("w1", "w3") else "experts.w2_", + f"experts.{expert_id}.{weight_name}.", + expert_id, + shard_id, + ) + for expert_id in range(num_experts) + for shard_id, weight_name in [ + ("w1", "w1"), + ("w2", "w2"), + ("w3", "w3"), + ] + ] + + +class DeepseekV4MegaMoEExperts(nn.Module): + _symm_buffer_cache: dict[tuple[int, int, int, int, int, int, int], object] = {} + + def __init__( + self, + vllm_config: VllmConfig, + *, + num_experts: int, + num_local_experts: int, + experts_start_idx: int, + top_k: int, + hidden_size: int, + intermediate_size: int, + prefix: str = "", + ): + super().__init__() + self.prefix = prefix + self.num_experts = num_experts + self.num_local_experts = num_local_experts + self.experts_start_idx = experts_start_idx + self.experts_end_idx = experts_start_idx + num_local_experts + self.top_k = top_k + self.hidden_size = hidden_size + self.intermediate_size = intermediate_size + self.max_num_tokens = vllm_config.scheduler_config.max_num_batched_tokens + + weight_attrs = {"weight_loader": self.weight_loader} + self.w13_weight = nn.Parameter( + torch.zeros( + num_local_experts, + 2 * intermediate_size, + hidden_size // 2, + dtype=torch.uint8, + ), + requires_grad=False, + ) + set_weight_attrs(self.w13_weight, weight_attrs) + + self.w13_weight_scale = nn.Parameter( + torch.zeros( + num_local_experts, + 2 * intermediate_size, + hidden_size // 32, + dtype=torch.uint8, + ), + requires_grad=False, + ) + set_weight_attrs(self.w13_weight_scale, weight_attrs) + self.w13_weight_scale.quant_method = "block" + + self.w2_weight = nn.Parameter( + torch.zeros( + num_local_experts, + hidden_size, + intermediate_size // 2, + dtype=torch.uint8, + ), + requires_grad=False, + ) + set_weight_attrs(self.w2_weight, weight_attrs) + + self.w2_weight_scale = nn.Parameter( + torch.zeros( + num_local_experts, + hidden_size, + intermediate_size // 32, + dtype=torch.uint8, + ), + requires_grad=False, + ) + set_weight_attrs(self.w2_weight_scale, weight_attrs) + self.w2_weight_scale.quant_method = "block" + + self._transformed_l1_weights: tuple[torch.Tensor, torch.Tensor] | None = None + self._transformed_l2_weights: tuple[torch.Tensor, torch.Tensor] | 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 + if prefix in compilation_config.static_forward_context: + raise ValueError(f"Duplicate layer name: {prefix}") + compilation_config.static_forward_context[prefix] = self + + def _map_global_expert_id(self, expert_id: int) -> int: + if expert_id < self.experts_start_idx or expert_id >= self.experts_end_idx: + return -1 + return expert_id - self.experts_start_idx + + 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 = self._map_global_expert_id(expert_id) + if local_expert_id == -1: + return False if return_success else None + + expert_data = param.data[local_expert_id] + if shard_id in ("w1", "w3"): + if "w13_" not in weight_name: + return False if return_success else None + 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": + if "w2_" not in weight_name: + return False if return_success else None + else: + raise ValueError(f"Unsupported expert shard id: {shard_id}") + + if expert_data.shape != loaded_weight.shape: + raise ValueError( + f"DeepSeek V4 MegaMoE expert weight shape mismatch for " + f"{weight_name}: parameter shard {tuple(expert_data.shape)} " + f"vs checkpoint {tuple(loaded_weight.shape)}" + ) + expert_data.copy_(loaded_weight) + return True if return_success else None + + @staticmethod + def _ue8m0_uint8_to_float(sf: torch.Tensor) -> torch.Tensor: + return (sf.to(torch.int32) << 23).view(torch.float32) + + def _check_runtime_supported(self) -> None: + if not torch.cuda.is_available(): + raise NotImplementedError("DeepSeek V4 MegaMoE requires CUDA.") + device = self.w13_weight.device + if device.type != "cuda": + raise NotImplementedError( + "DeepSeek V4 MegaMoE expert weights must be loaded on CUDA." + ) + 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." + ) + + def finalize_weights(self) -> None: + if self._transformed_l1_weights is not None: + return + + self._check_runtime_supported() + import vllm.third_party.deep_gemm as 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, + ) + 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), + ) + ) + # Drop the original loader-side parameters: the MegaMoE kernels only + # consume the transformed views above. transform_weights_for_mega_moe + # allocates a fresh tensor for the L1 weight (see _interleave_l1_weights) + # and fresh SF tensors for L1/L2; the L2 weight is the only tensor that + # aliases the original storage, and _transformed_l2_weights still holds + # it, so the storage stays live after we drop the Parameter. + self.w13_weight = None + self.w13_weight_scale = None + self.w2_weight = None + self.w2_weight_scale = None + + def get_symm_buffer(self): + import vllm.third_party.deep_gemm as deep_gemm + + group = get_ep_group().device_group + device = torch.accelerator.current_device_index() + key = ( + id(group), + device, + self.num_experts, + self.max_num_tokens, + self.top_k, + self.hidden_size, + self.intermediate_size, + ) + symm_buffer = self._symm_buffer_cache.get(key) + if symm_buffer is None: + symm_buffer = deep_gemm.get_symm_buffer_for_mega_moe( + group, + self.num_experts, + self.max_num_tokens, + self.top_k, + self.hidden_size, + self.intermediate_size, + ) + self._symm_buffer_cache[key] = symm_buffer + return symm_buffer + + def forward( + self, + hidden_states: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + *, + activation_clamp: float | None, + fast_math: bool = True, + ) -> torch.Tensor: + if hidden_states.shape[0] > self.max_num_tokens: + raise ValueError( + f"DeepSeek V4 MegaMoE got {hidden_states.shape[0]} tokens, " + f"but the symmetric buffer was sized for {self.max_num_tokens}." + ) + y = torch.empty_like(hidden_states, dtype=torch.bfloat16) + torch.ops.vllm.deepseek_v4_mega_moe_experts( + hidden_states, + topk_weights, + topk_ids, + y, + self.prefix, + activation_clamp, + fast_math, + ) + return y + + def _run_mega_moe( + self, + hidden_states: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + y: torch.Tensor, + activation_clamp: float | None, + fast_math: bool, + ) -> None: + import vllm.third_party.deep_gemm as deep_gemm + + symm_buffer = self.get_symm_buffer() + num_tokens = hidden_states.shape[0] + prepare_megamoe_inputs( + hidden_states, + topk_weights, + topk_ids, + symm_buffer.x[:num_tokens], + symm_buffer.x_sf[:num_tokens], + symm_buffer.topk_idx[:num_tokens], + symm_buffer.topk_weights[:num_tokens], + ) + + # This method must have been already called during the weight loading phase. + # We call it again here to cover the dummy weight loading case. + 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, + activation_clamp=activation_clamp, + fast_math=fast_math, + ) + + +DeepseekV4MegaMoEExperts.weight_loader.supports_moe_loading = True # type: ignore[attr-defined] + + +def _deepseek_v4_mega_moe_experts_op( + hidden_states: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + out: torch.Tensor, + layer_name: str, + activation_clamp: float | None, + fast_math: bool, +) -> None: + self = get_forward_context().no_compile_layers[layer_name] + self._run_mega_moe( + hidden_states, + topk_weights, + topk_ids, + out, + activation_clamp, + fast_math, + ) + + +def _deepseek_v4_mega_moe_experts_op_fake( + hidden_states: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + out: torch.Tensor, + layer_name: str, + activation_clamp: float | None, + fast_math: bool, +) -> None: + return None + + +direct_register_custom_op( + op_name="deepseek_v4_mega_moe_experts", + op_func=_deepseek_v4_mega_moe_experts_op, + mutates_args=["out"], + fake_impl=_deepseek_v4_mega_moe_experts_op_fake, +) + + +def _deepseek_v4_b12x_mhc_post_pre_op( + x: torch.Tensor, + residual: torch.Tensor, + post: torch.Tensor, + comb: torch.Tensor, + hc_fn: torch.Tensor, + hc_scale: torch.Tensor, + hc_base: torch.Tensor, + norm_weight: torch.Tensor, + norm_eps: float, + layer_name: str, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + self = get_forward_context().no_compile_layers[layer_name] + return _trace_b12x_mhc_call( + "post_pre", + layer_name, + int(residual.shape[0]), + lambda: self._run_b12x_mhc_post_pre( + x, + residual, + post, + comb, + hc_fn, + hc_scale, + hc_base, + norm_weight, + norm_eps, + ), + ) + + +def _deepseek_v4_b12x_mhc_post_pre_op_fake( + x: torch.Tensor, + residual: torch.Tensor, + post: torch.Tensor, + comb: torch.Tensor, + hc_fn: torch.Tensor, + hc_scale: torch.Tensor, + hc_base: torch.Tensor, + norm_weight: torch.Tensor, + norm_eps: float, + layer_name: str, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + del x, post, comb, hc_fn, hc_scale, hc_base, norm_weight, norm_eps, layer_name + tokens, hc_mult, hidden_size = residual.shape + residual_out = torch.empty_like(residual) + post_out = torch.empty( + (tokens, hc_mult), dtype=torch.float32, device=residual.device + ) + comb_out = torch.empty( + (tokens, hc_mult, hc_mult), dtype=torch.float32, device=residual.device + ) + y_out = torch.empty( + (tokens, hidden_size), dtype=residual.dtype, device=residual.device + ) + return residual_out, post_out, comb_out, y_out + + +direct_register_custom_op( + op_name="deepseek_v4_b12x_mhc_post_pre", + op_func=_deepseek_v4_b12x_mhc_post_pre_op, + fake_impl=_deepseek_v4_b12x_mhc_post_pre_op_fake, +) + + +class DeepseekV4MoE(nn.Module): + def __init__( + self, + vllm_config: VllmConfig, + prefix: str = "", + ): + super().__init__() + + self.tp_size = get_tensor_model_parallel_world_size() + config = vllm_config.model_config.hf_config + quant_config = vllm_config.quant_config + self.prefix = prefix + self.use_mega_moe = ( + vllm_config.kernel_config.moe_backend == "deep_gemm_mega_moe" + ) + if self.use_mega_moe and not vllm_config.parallel_config.enable_expert_parallel: + raise NotImplementedError( + "DeepSeek V4 MegaMoE currently requires expert parallel. " + "Enable it with --enable-expert-parallel, or pick a different " + "moe backend." + ) + + self.routed_scaling_factor = getattr(config, "routed_scaling_factor", 1.0) + self.hidden_size = config.hidden_size + + self.n_routed_experts = config.n_routed_experts + self.n_activated_experts = config.num_experts_per_tok + self.moe_intermediate_size = config.moe_intermediate_size + self.swiglu_limit = config.swiglu_limit + self.renormalize = config.norm_topk_prob + self.scoring_func = getattr(config, "scoring_func", "sqrtsoftplus") + if self.use_mega_moe and self.scoring_func != "sqrtsoftplus": + raise NotImplementedError( + "DeepSeek V4 MegaMoE currently supports sqrtsoftplus routing only." + ) + if self.use_mega_moe and getattr(config, "expert_dtype", "fp4") != "fp4": + raise NotImplementedError( + "DeepSeek V4 MegaMoE only supports fp4 experts; got expert_dtype=" + f"{config.expert_dtype!r}. Drop --kernel-config moe_backend=" + "deep_gemm_mega_moe for this checkpoint." + ) + + self.gate = GateLinear( + input_size=config.hidden_size, + output_size=config.n_routed_experts, + bias=False, + out_dtype=torch.float32, + prefix=f"{prefix}.gate", + ) + + self.gate.e_score_correction_bias = None + self.gate.tid2eid = None + is_hash_moe = extract_layer_index(prefix) < config.num_hash_layers + self.hash_indices_dtype = torch.int64 if self.use_mega_moe else torch.int32 + if is_hash_moe: + # hash MoE doesn't use e_score_correction_bias + # Use randint instead of empty to avoid garbage values causing + # invalid memory access in dummy mode (--load-format="dummy") + self.gate.tid2eid = nn.Parameter( + torch.randint( + 0, + config.n_routed_experts, + (config.vocab_size, config.num_experts_per_tok), + dtype=self.hash_indices_dtype, + ), + requires_grad=False, + ) + elif getattr(config, "topk_method", None) == "noaux_tc": + self.gate.e_score_correction_bias = nn.Parameter( + torch.empty(config.n_routed_experts, dtype=torch.float32), + requires_grad=False, + ) + + if config.n_shared_experts is None: + self.shared_experts = None + else: + intermediate_size = config.moe_intermediate_size * config.n_shared_experts + + self.shared_experts = DeepseekV4MLP( + hidden_size=config.hidden_size, + intermediate_size=intermediate_size, + hidden_act=config.hidden_act, + swiglu_limit=self.swiglu_limit, + quant_config=quant_config, + reduce_results=self.use_mega_moe, + prefix=f"{prefix}.shared_experts", + ) + + if self.use_mega_moe: + self._init_mega_moe_experts(vllm_config, config, prefix) + else: + self._init_fused_moe_experts(config, quant_config, prefix) + + def _init_mega_moe_experts( + self, + vllm_config: VllmConfig, + config, + prefix: str, + ) -> None: + self.ep_group = get_ep_group() + self.ep_size = self.ep_group.world_size + self.ep_rank = self.ep_group.rank_in_group + assert config.n_routed_experts % self.ep_size == 0 + + self.n_local_experts = config.n_routed_experts // self.ep_size + self.experts_start_idx = self.ep_rank * self.n_local_experts + self.experts_end_idx = self.experts_start_idx + self.n_local_experts + + self.experts = DeepseekV4MegaMoEExperts( + vllm_config, + num_experts=config.n_routed_experts, + num_local_experts=self.n_local_experts, + experts_start_idx=self.experts_start_idx, + top_k=config.num_experts_per_tok, + hidden_size=config.hidden_size, + intermediate_size=config.moe_intermediate_size, + prefix=f"{prefix}.experts", + ) + + def _init_fused_moe_experts( + self, + config, + quant_config, + prefix: str, + ) -> None: + self.tp_rank = get_tensor_model_parallel_rank() + assert config.n_routed_experts % self.tp_size == 0 + + self.n_local_experts = config.n_routed_experts // self.tp_size + self.experts_start_idx = self.tp_rank * self.n_local_experts + self.experts_end_idx = self.experts_start_idx + self.n_local_experts + + self.experts = FusedMoE( + shared_experts=self.shared_experts, + gate=self.gate, + num_experts=config.n_routed_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, + prefix=f"{prefix}.experts", + scoring_func=self.scoring_func, + routed_scaling_factor=self.routed_scaling_factor, + e_score_correction_bias=self.gate.e_score_correction_bias, + hash_indices_table=self.gate.tid2eid, + swiglu_limit=self.swiglu_limit, + router_logits_dtype=torch.float32, + ) + + def forward( + self, hidden_states: torch.Tensor, input_ids: torch.Tensor | None = None + ) -> torch.Tensor: + if self.gate.tid2eid is not None and input_ids is None: + raise ValueError("DeepSeek V4 hash MoE routing requires input_ids.") + + if not self.use_mega_moe: + return self._forward_fused_moe(hidden_states, input_ids) + + org_shape = hidden_states.shape + 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.data + if self.gate.e_score_correction_bias is not None + else None, + topk=self.n_activated_experts, + renormalize=self.renormalize, + indices_type=self.hash_indices_dtype, + input_tokens=input_ids, + hash_indices_table=self.gate.tid2eid, + routed_scaling_factor=self.routed_scaling_factor, + ) + activation_clamp = ( + float(self.swiglu_limit) if self.swiglu_limit is not None else None + ) + final_hidden_states = self.experts( + hidden_states, + topk_weights, + topk_ids, + activation_clamp=activation_clamp, + ) + + if self.shared_experts is not None: + shared_output = self.shared_experts(hidden_states) + final_hidden_states += shared_output + + return final_hidden_states.view(org_shape) + + def _forward_fused_moe( + self, hidden_states: torch.Tensor, input_ids: torch.Tensor | None = None + ) -> torch.Tensor: + org_shape = hidden_states.shape + if self.experts.is_internal_router: + # In this case, the gate/router runs inside the FusedMoE class + final_hidden_states = self.experts( + hidden_states=hidden_states, + router_logits=hidden_states, + input_ids=input_ids, + ) + else: + router_logits, _ = self.gate(hidden_states) + final_hidden_states = self.experts( + hidden_states=hidden_states, + router_logits=router_logits, + input_ids=input_ids, + ) + + return final_hidden_states.view(org_shape) + + def finalize_mega_moe_weights(self) -> None: + if self.use_mega_moe: + self.experts.finalize_weights() + + +class DeepseekV4Attention(nn.Module): + def __init__( + self, + vllm_config: VllmConfig, + prefix: str, + topk_indices_buffer: torch.Tensor | None = None, + aux_stream_list: list[torch.cuda.Stream] | None = None, + ): + super().__init__() + config = vllm_config.model_config.hf_config + quant_config = vllm_config.quant_config + layer_id = extract_layer_index(prefix) + + self.layer_id = layer_id + self.hidden_size = config.hidden_size + self.n_heads = config.num_attention_heads + tp_size = get_tensor_model_parallel_world_size() + assert self.n_heads % tp_size == 0 + + self.n_local_heads = self.n_heads // tp_size + self.q_lora_rank = config.q_lora_rank + self.o_lora_rank = config.o_lora_rank + self.head_dim = config.head_dim + self.rope_head_dim = config.qk_rope_head_dim + self.nope_head_dim = self.head_dim - self.rope_head_dim + self.n_groups = config.o_groups + self.n_local_groups = self.n_groups // tp_size + self.window_size = config.sliding_window + # NOTE(zyongye) Compress ratio can't be 0 + # we do this for because MTP layer is not included + # in the compress ratio list + if layer_id < config.num_hidden_layers: + self.compress_ratio = max(1, config.compress_ratios[layer_id]) + else: + self.compress_ratio = 1 + self.eps = config.rms_norm_eps + self.max_position_embeddings = config.max_position_embeddings + + # Must match DeepseekV4MLAAttention.padded_heads; padded entries stay + # -inf so they contribute no sink effect. + padded_heads = get_deepseek_v4_padded_num_q_heads(self.n_local_heads) + self.attn_sink = nn.Parameter( + torch.full((padded_heads,), -float("inf"), dtype=torch.float32), + requires_grad=False, + ) + + self.fused_wqa_wkv = MergedColumnParallelLinear( + self.hidden_size, + [self.q_lora_rank, self.head_dim], + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.fused_wqa_wkv", + disable_tp=True, # fused ReplicatedLinear + ) + self.q_norm = RMSNorm(self.q_lora_rank, self.eps) + self.wq_b = ColumnParallelLinear( + self.q_lora_rank, + self.n_heads * self.head_dim, + bias=False, + quant_config=quant_config, + return_bias=False, + prefix=f"{prefix}.wq_b", + ) + + self.kv_norm = RMSNorm(self.head_dim, self.eps) + self.wo_a = ColumnParallelLinear( + self.n_heads * self.head_dim // self.n_groups, + self.n_groups * self.o_lora_rank, + bias=False, + quant_config=quant_config, + return_bias=False, + prefix=f"{prefix}.wo_a", + ) + self.wo_a.is_bmm = True + self.wo_a.bmm_batch_size = self.n_local_groups + if envs.VLLM_USE_B12X_WO_PROJECTION: + if not hasattr(self.wo_a, "weight_scale_inv"): + raise RuntimeError( + "VLLM_USE_B12X_WO_PROJECTION requires FP8 wo_a.weight_scale_inv" + ) + # Preserve checkpoint UE8M0 scales for the fused b12x WO kernel. + self.wo_a.weight_scale_inv.format_ue8m0 = True + self.wo_b = RowParallelLinear( + self.n_groups * self.o_lora_rank, + self.hidden_size, + bias=False, + quant_config=quant_config, + return_bias=False, + prefix=f"{prefix}.wo_b", + ) + if envs.VLLM_USE_B12X_WO_PROJECTION: + if not hasattr(self.wo_b, "weight_scale_inv"): + raise RuntimeError( + "VLLM_USE_B12X_WO_PROJECTION requires FP8 wo_b.weight_scale_inv" + ) + self.wo_a.b12x_skip_generic_block_fp8_linear = True + self.wo_b.b12x_skip_generic_block_fp8_linear = True + + self.softmax_scale = self.head_dim**-0.5 + self.scale_fmt = config.quantization_config["scale_fmt"] + + self.rope_parameters = config.rope_scaling + + # Initialize rotary embedding BEFORE DeepseekV4MLAModules (which needs it) + rope_parameters = config.rope_parameters + rope_parameters["rope_theta"] = ( + config.compress_rope_theta if self.compress_ratio > 1 else config.rope_theta + ) + if config.rope_parameters["rope_type"] != "default": + config.rope_parameters["rope_type"] = ( + "deepseek_yarn" + if config.rope_parameters.get("apply_yarn_scaling", True) + else "deepseek_llama_scaling" + ) + rope_parameters["mscale"] = 0 # Disable mscale + rope_parameters["mscale_all_dim"] = 0 # Disable mscale + rope_parameters["is_deepseek_v4"] = True + rope_parameters["rope_dim"] = self.rope_head_dim + self.rotary_emb = get_rope( + self.head_dim, + max_position=self.max_position_embeddings, + rope_parameters=rope_parameters, + is_neox_style=False, + ) + + self.indexer = None + if self.compress_ratio == 4: + # Only C4A uses sparse attention and hence has indexer. + # aux_stream_list[0] runs indexer.forward() in the wrapper; [2] is + # free here (outer GEMMs joined) for the inner overlap of + # wq_b+fused_indexer_q_rope_quant vs compressor. + indexer_aux_stream = ( + aux_stream_list[2] if aux_stream_list is not None else None + ) + self.indexer = DeepseekV4Indexer( + vllm_config, + config=config, + hidden_size=self.hidden_size, + q_lora_rank=self.q_lora_rank, + quant_config=quant_config, + cache_config=vllm_config.cache_config, + topk_indices_buffer=topk_indices_buffer, + compress_ratio=self.compress_ratio, + prefix=f"{prefix}.indexer", + aux_stream=indexer_aux_stream, + ) + + mla_modules = DeepseekV4MLAModules( + vllm_config=vllm_config, + fused_wqa_wkv=self.fused_wqa_wkv, + q_norm=self.q_norm, + wq_b=self.wq_b, + kv_norm=self.kv_norm, + wo_a=self.wo_a, + wo_b=self.wo_b, + attn_sink=self.attn_sink, + rotary_emb=self.rotary_emb, + indexer=self.indexer, + indexer_rotary_emb=self.rotary_emb, + topk_indices_buffer=topk_indices_buffer, + aux_stream_list=aux_stream_list, + ) + self.mla_attn = DeepseekV4MultiHeadLatentAttentionWrapper( + hidden_size=self.hidden_size, + num_heads=self.n_local_heads, + head_dim=self.head_dim, + scale=self.softmax_scale, + qk_nope_head_dim=self.nope_head_dim, + qk_rope_head_dim=self.rope_head_dim, + v_head_dim=self.head_dim, + q_lora_rank=self.q_lora_rank, + kv_lora_rank=self.head_dim, + o_lora_rank=self.o_lora_rank, + mla_modules=mla_modules, + window_size=self.window_size, + compress_ratio=self.compress_ratio, + cache_config=vllm_config.cache_config, + quant_config=quant_config, + prefix=prefix, + ) + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + llama_4_scaling: torch.Tensor | None, + ): + return self.mla_attn(positions, hidden_states, llama_4_scaling) + + def setup_b12x_wo_projection(self) -> None: + self.mla_attn.setup_b12x_wo_projection() + + +class DeepseekV4DecoderLayer(nn.Module): + def __init__( + self, + vllm_config, + prefix, + topk_indices_buffer: torch.Tensor | None = None, + aux_stream_list: list[torch.cuda.Stream] | None = None, + ): + super().__init__() + + config = vllm_config.model_config.hf_config + self.layer_name = prefix + self._dspark_prev_capture_buffer: torch.Tensor | None = None + self._dspark_prev_capture_start = 0 + self._dspark_prev_capture_end = 0 + self._dspark_prev_capture_exact = False + self._use_b12x_mhc = _use_b12x_mhc() + self._b12x_mhc_max_tokens = _b12x_mhc_max_tokens() if self._use_b12x_mhc else 0 + if self._use_b12x_mhc: + if not prefix: + raise RuntimeError("DeepSeek V4 b12x mHC decoder layer needs a prefix") + compilation_config = vllm_config.compilation_config + if prefix in compilation_config.static_forward_context: + raise ValueError(f"Duplicate layer name: {prefix}") + compilation_config.static_forward_context[prefix] = self + + if self._b12x_mhc_max_tokens <= 0: + logger.info_once("DeepSeek V4 b12x mHC enabled for all token counts.") + else: + logger.info_once( + "DeepSeek V4 b12x mHC enabled for token counts <= %d; " + "using TileLang mHC above that.", + self._b12x_mhc_max_tokens, + ) + + # Registers torch.ops.vllm.mhc_* and provides the fallback path for + # mixed/prefill capture sizes when b12x mHC is decode-limited. + from vllm.model_executor.layers.mhc import ( + MHCFusedPostPreOp, + MHCPostOp, + MHCPreOp, + ) + + self.hidden_size = config.hidden_size + + self.rms_norm_eps = config.rms_norm_eps + self.attn = DeepseekV4Attention( + vllm_config, + prefix=f"{prefix}.attn", + topk_indices_buffer=topk_indices_buffer, + aux_stream_list=aux_stream_list, + ) + self.ffn = DeepseekV4MoE(vllm_config, prefix=f"{prefix}.ffn") + + self.attn_norm = RMSNorm(self.hidden_size, self.rms_norm_eps) + self.ffn_norm = RMSNorm(self.hidden_size, self.rms_norm_eps) + self.hc_mult = config.hc_mult + self.hc_sinkhorn_iters = config.hc_sinkhorn_iters + self.hc_eps = config.hc_eps + self.hc_post_alpha = 2.0 + mix_hc = (2 + self.hc_mult) * self.hc_mult + hc_dim = self.hc_mult * self.hidden_size + self.hc_attn_fn = nn.Parameter( + torch.empty( + (mix_hc, hc_dim), + dtype=torch.float32, + ), + requires_grad=False, + ) + self.hc_ffn_fn = nn.Parameter( + torch.empty( + (mix_hc, hc_dim), + dtype=torch.float32, + ), + requires_grad=False, + ) + self.hc_attn_base = nn.Parameter( + torch.empty( + mix_hc, + dtype=torch.float32, + ), + requires_grad=False, + ) + self.hc_ffn_base = nn.Parameter( + torch.empty( + mix_hc, + dtype=torch.float32, + ), + requires_grad=False, + ) + self.hc_attn_scale = nn.Parameter( + torch.empty( + 3, + dtype=torch.float32, + ), + requires_grad=False, + ) + self.hc_ffn_scale = nn.Parameter( + torch.empty( + 3, + dtype=torch.float32, + ), + requires_grad=False, + ) + if self._use_b12x_mhc: + from b12x.integration.residual import ( + MHC_DEFAULT_BLOCK_K, + MHC_DEFAULT_SPLIT_K, + MHC_MULT, + ) + + if self.hc_mult != MHC_MULT: + raise NotImplementedError( + f"DeepSeek V4 b12x mHC requires hc_mult={MHC_MULT}, " + f"got {self.hc_mult}." + ) + if self.hidden_size != 4096: + raise NotImplementedError( + "DeepSeek V4 b12x mHC currently requires hidden_size=4096, " + f"got {self.hidden_size}." + ) + self._b12x_mhc_block_k = int(MHC_DEFAULT_BLOCK_K) + total_k = self.hc_mult * self.hidden_size + if total_k % self._b12x_mhc_block_k != 0: + raise ValueError( + "DeepSeek V4 b12x mHC requires hc_mult * hidden_size to " + f"be divisible by block_k={self._b12x_mhc_block_k}, got {total_k}." + ) + self._b12x_mhc_split_k = total_k // self._b12x_mhc_block_k + if self._b12x_mhc_split_k != MHC_DEFAULT_SPLIT_K: + raise NotImplementedError( + "DeepSeek V4 b12x mHC currently requires " + f"split_k={MHC_DEFAULT_SPLIT_K}, got {self._b12x_mhc_split_k}." + ) + else: + self._b12x_mhc_block_k = 0 + self._b12x_mhc_split_k = 0 + self.mhc_pre = MHCPreOp() + self.mhc_post = MHCPostOp() + self.mhc_fused_post_pre = MHCFusedPostPreOp() + + def set_dspark_previous_layer_capture( + self, + buffer: torch.Tensor, + start: int, + end: int, + *, + exact: bool = False, + ) -> None: + self._dspark_prev_capture_buffer = buffer + self._dspark_prev_capture_start = int(start) + self._dspark_prev_capture_end = int(end) + self._dspark_prev_capture_exact = bool(exact) + + def _maybe_capture_dspark_previous_layer_exact( + self, + x: torch.Tensor, + residual: torch.Tensor, + post: torch.Tensor, + comb: torch.Tensor, + ) -> None: + buffer = self._dspark_prev_capture_buffer + if buffer is None or not self._dspark_prev_capture_exact: + return + start = self._dspark_prev_capture_start + end = self._dspark_prev_capture_end + from vllm.models.deepseek_v4.nvidia.dspark_kernels import ( + dspark_hc_post_mean, + ) + + dspark_hc_post_mean( + x, + residual, + post, + comb, + buffer[: x.shape[0], start:end], + ) + + def _maybe_capture_dspark_previous_layer( + self, + residual: torch.Tensor, + ) -> None: + buffer = self._dspark_prev_capture_buffer + if buffer is None: + return + if self._dspark_prev_capture_exact: + return + start = self._dspark_prev_capture_start + end = self._dspark_prev_capture_end + buffer[: residual.shape[0], start:end].copy_(residual.mean(dim=1)) + + def _should_run_b12x_mhc(self, tokens: int) -> bool: + if not self._use_b12x_mhc: + return False + max_tokens = self._b12x_mhc_max_tokens + return max_tokens <= 0 or int(tokens) <= max_tokens + + def _get_b12x_mhc_binding( + self, + x: torch.Tensor, + y: torch.Tensor | None = None, + post: torch.Tensor | None = None, + comb: torch.Tensor | None = None, + out: torch.Tensor | None = None, + ) -> object: + from b12x.integration.residual import B12XMHCScratchCaps, plan_mhc_scratch + + tokens = int(x.shape[0]) + plan = plan_mhc_scratch( + B12XMHCScratchCaps( + device=x.device, + dtype=x.dtype, + max_tokens=max(1, tokens), + hidden_size=self.hidden_size, + split_k=self._b12x_mhc_split_k, + ) + ) + scratch = _empty_b12x_plan_scratch(plan, x.device) + return plan.bind( + scratch=scratch, + tokens=tokens, + y=y, + post=post, + comb=comb, + out=out, + ) + + def _run_b12x_mhc_post_pre( + self, + x: torch.Tensor, + residual: torch.Tensor, + post: torch.Tensor, + comb: torch.Tensor, + hc_fn: torch.Tensor, + hc_scale: torch.Tensor, + hc_base: torch.Tensor, + norm_weight: torch.Tensor | None = None, + norm_eps: float = 0.0, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + from b12x.integration.residual import b12x_mhc_post_pre + + tokens, hc_mult, hidden_size = residual.shape + residual_out = torch.empty_like(residual) + y_out = torch.empty( + (tokens, hidden_size), dtype=residual.dtype, device=residual.device + ) + post_out = torch.empty( + (tokens, hc_mult), dtype=torch.float32, device=residual.device + ) + comb_out = torch.empty( + (tokens, hc_mult, hc_mult), dtype=torch.float32, device=residual.device + ) + binding = self._get_b12x_mhc_binding( + residual, + y=y_out, + post=post_out, + comb=comb_out, + out=residual_out, + ) + return b12x_mhc_post_pre( + x, + residual, + post, + comb, + hc_fn, + hc_scale, + hc_base, + rms_eps=self.rms_norm_eps, + hc_eps=self.hc_eps, + sinkhorn_iters=self.hc_sinkhorn_iters, + norm_weight=norm_weight, + norm_eps=norm_eps, + binding=binding, + block_k=self._b12x_mhc_block_k, + ) + + def hc_pre( + self, + x: torch.Tensor, + hc_fn: torch.Tensor, + hc_scale: torch.Tensor, + hc_base: torch.Tensor, + norm_weight: torch.Tensor | None = None, + norm_eps: float = 1e-6, + ): + assert self.mhc_pre is not None + post_mix, res_mix, layer_input = self.mhc_pre( + residual=x, + fn=hc_fn, + hc_scale=hc_scale, + hc_base=hc_base, + rms_eps=self.rms_norm_eps, + hc_pre_eps=self.hc_eps, + hc_sinkhorn_eps=self.hc_eps, + hc_post_mult_value=self.hc_post_alpha, + sinkhorn_repeat=self.hc_sinkhorn_iters, + norm_weight=norm_weight, + norm_eps=norm_eps, + ) + return layer_input, post_mix, res_mix + + def hc_post( + self, + x: torch.Tensor, + residual: torch.Tensor, + post: torch.Tensor, + comb: torch.Tensor, + ): + assert self.mhc_post is not None + return self.mhc_post(x, residual, post, comb) + + def hc_post_pre( + self, + x: torch.Tensor, + residual: torch.Tensor, + post: torch.Tensor, + comb: torch.Tensor, + hc_fn: torch.Tensor, + hc_scale: torch.Tensor, + hc_base: torch.Tensor, + norm_weight: torch.Tensor | None = None, + norm_eps: float = 1e-6, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + if self._should_run_b12x_mhc(int(residual.shape[0])): + if not is_forward_context_available(): + return self._run_b12x_mhc_post_pre( + x, + residual, + post, + comb, + hc_fn, + hc_scale, + hc_base, + norm_weight, + norm_eps, + ) + return torch.ops.vllm.deepseek_v4_b12x_mhc_post_pre( + x, + residual, + post, + comb, + hc_fn, + hc_scale, + hc_base, + norm_weight, + norm_eps, + self.layer_name, + ) + + assert self.mhc_fused_post_pre is not None + return self.mhc_fused_post_pre( + x, + residual, + post, + comb, + hc_fn, + hc_scale, + hc_base, + self.rms_norm_eps, + self.hc_eps, + self.hc_eps, + self.hc_post_alpha, + self.hc_sinkhorn_iters, + n_splits=1, + tile_n=1, + norm_weight=norm_weight, + norm_eps=norm_eps, + ) + + def _forward_cuda( + self, + x: torch.Tensor, + positions: torch.Tensor, + input_ids: torch.Tensor | None, + post_mix: torch.Tensor | None = None, + res_mix: torch.Tensor | None = None, + residual: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + if self._should_run_b12x_mhc(int(x.shape[0])): + attn_norm_weight = self.attn_norm.weight.data + attn_norm_eps = self.attn_norm.variance_epsilon + stage_start = _dspark_target_timing_start() + if residual is None: + residual = x + x, post_mix, res_mix = self.hc_pre( + residual, + self.hc_attn_fn, + self.hc_attn_scale, + self.hc_attn_base, + norm_weight=attn_norm_weight, + norm_eps=attn_norm_eps, + ) + else: + assert post_mix is not None + assert res_mix is not None + self._maybe_capture_dspark_previous_layer_exact( + x, residual, post_mix, res_mix + ) + residual, post_mix, res_mix, x = self.hc_post_pre( + x, + residual, + post_mix, + res_mix, + self.hc_attn_fn, + self.hc_attn_scale, + self.hc_attn_base, + norm_weight=attn_norm_weight, + norm_eps=attn_norm_eps, + ) + self._maybe_capture_dspark_previous_layer(residual) + _dspark_target_timing_record("layer_attn_mhc", stage_start) + + stage_start = _dspark_target_timing_start() + x = self.attn(positions, x, None) + _dspark_target_timing_record("layer_attn", stage_start) + + ffn_norm_weight = self.ffn_norm.weight.data + ffn_norm_eps = self.ffn_norm.variance_epsilon + stage_start = _dspark_target_timing_start() + residual, post_mix, res_mix, x = self.hc_post_pre( + x, + residual, + post_mix, + res_mix, + self.hc_ffn_fn, + self.hc_ffn_scale, + self.hc_ffn_base, + norm_weight=ffn_norm_weight, + norm_eps=ffn_norm_eps, + ) + _dspark_target_timing_record("layer_ffn_mhc", stage_start) + stage_start = _dspark_target_timing_start() + x = self.ffn(x, input_ids) + _dspark_target_timing_record("layer_ffn", stage_start) + return x, residual, post_mix, res_mix + + assert self.mhc_fused_post_pre is not None + attn_norm_weight = self.attn_norm.weight.data + attn_norm_eps = self.attn_norm.variance_epsilon + stage_start = _dspark_target_timing_start() + if residual is None: + # Run standalone hc_pre on first layer + residual = x + x, post_mix, res_mix = self.hc_pre( + x, + self.hc_attn_fn, + self.hc_attn_scale, + self.hc_attn_base, + norm_weight=attn_norm_weight, + norm_eps=attn_norm_eps, + ) + else: + assert post_mix is not None + assert res_mix is not None + self._maybe_capture_dspark_previous_layer_exact( + x, residual, post_mix, res_mix + ) + residual, post_mix, res_mix, x = self.mhc_fused_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, + n_splits=1, + tile_n=1, + norm_weight=attn_norm_weight, + norm_eps=attn_norm_eps, + ) + self._maybe_capture_dspark_previous_layer(residual) + _dspark_target_timing_record("layer_attn_mhc", stage_start) + + # attn_norm is fused into hc_pre / mhc_fused_post_pre above. + stage_start = _dspark_target_timing_start() + x = self.attn(positions, x, None) + _dspark_target_timing_record("layer_attn", stage_start) + + ffn_norm_weight = self.ffn_norm.weight.data + ffn_norm_eps = self.ffn_norm.variance_epsilon + stage_start = _dspark_target_timing_start() + residual, post_mix, res_mix, x = self.mhc_fused_post_pre( + x, + residual, + post_mix, + res_mix, + self.hc_ffn_fn, + self.hc_ffn_scale, + self.hc_ffn_base, + self.rms_norm_eps, + self.hc_eps, + self.hc_eps, + self.hc_post_alpha, + self.hc_sinkhorn_iters, + n_splits=1, + tile_n=1, + norm_weight=ffn_norm_weight, + norm_eps=ffn_norm_eps, + ) + _dspark_target_timing_record("layer_ffn_mhc", stage_start) + + stage_start = _dspark_target_timing_start() + x = self.ffn(x, input_ids) + _dspark_target_timing_record("layer_ffn", stage_start) + return x, residual, post_mix, res_mix + + def _forward_native( + self, + x: torch.Tensor, + positions: torch.Tensor, + input_ids: torch.Tensor | None, + post_mix: torch.Tensor | None = None, + res_mix: torch.Tensor | None = None, + residual: torch.Tensor | None = None, + ) -> tuple[ + torch.Tensor, torch.Tensor | None, torch.Tensor | None, torch.Tensor | None + ]: + residual = x + x, post, comb = self.hc_pre( + x, self.hc_attn_fn, self.hc_attn_scale, self.hc_attn_base + ) + x = self.attn_norm(x) + x = self.attn(positions, x, None) + x = self.hc_post(x, residual, post, comb) + + residual = x + x, post, comb = self.hc_pre( + x, self.hc_ffn_fn, self.hc_ffn_scale, self.hc_ffn_base + ) + x = self.ffn_norm(x) + x = self.ffn(x, input_ids) + x = self.hc_post(x, residual, post, comb) + return x, None, None, None + + def forward( + self, + x: torch.Tensor, + positions: torch.Tensor, + input_ids: torch.Tensor | None, + post_mix: torch.Tensor | None = None, + res_mix: torch.Tensor | None = None, + residual: torch.Tensor | None = None, + ) -> tuple[ + torch.Tensor, torch.Tensor | None, torch.Tensor | None, torch.Tensor | None + ]: + if current_platform.is_rocm() or current_platform.is_xpu(): + return self._forward_native( + x, positions, input_ids, post_mix, res_mix, residual + ) + + return self._forward_cuda(x, positions, input_ids, post_mix, res_mix, residual) + + +@support_torch_compile +class DeepseekV4Model(nn.Module): + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): + super().__init__() + + config = vllm_config.model_config.hf_config + quant_config = vllm_config.quant_config + self.config = config + self.use_mega_moe = ( + vllm_config.kernel_config.moe_backend == "deep_gemm_mega_moe" + ) + if self.use_mega_moe and not vllm_config.parallel_config.enable_expert_parallel: + raise NotImplementedError( + "DeepSeek V4 MegaMoE currently requires expert parallel. " + "Enable it with --enable-expert-parallel, or pick a different " + "moe backend." + ) + self.vocab_size = config.vocab_size + self.hc_eps = config.hc_eps + self.hc_mult = config.hc_mult + self.hc_dim = self.hc_mult * config.hidden_size + self.rms_norm_eps = config.rms_norm_eps + + # Three aux streams: one per non-default input GEMM in + # DeepseekV4MultiHeadLatentAttentionWrapper.attn_gemm_parallel_execute + # (compressor kv_score, indexer.weights_proj, indexer.compressor + # kv_score). fused_wqa_wkv stays on the default stream. + # Disable them on ROCm / XPU because of hang issues / no overlap. + aux_stream_list = ( + None + if current_platform.is_rocm() or current_platform.is_xpu() + else [torch.cuda.Stream() for _ in range(3)] + ) + + self.device = current_platform.device_type + # Reserved topk indices buffer for all Indexer layers to reuse. + self.topk_indices_buffer = torch.empty( + vllm_config.scheduler_config.max_num_batched_tokens, + config.index_topk, + dtype=torch.int32, + device=self.device, + ) + + 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: DeepseekV4DecoderLayer( + vllm_config, + prefix=prefix, + topk_indices_buffer=self.topk_indices_buffer, + aux_stream_list=aux_stream_list, + ), + prefix=f"{prefix}.layers", + ) + + if get_pp_group().is_last_rank: + self.norm = RMSNorm(config.hidden_size, self.rms_norm_eps) + else: + self.norm = PPMissingLayer() + + self.hc_head_fn = nn.Parameter( + torch.empty( + self.hc_mult, + self.hc_dim, + dtype=torch.float32, + ), + requires_grad=False, + ) + self.hc_head_base = nn.Parameter( + torch.empty( + self.hc_mult, + dtype=torch.float32, + ), + requires_grad=False, + ) + self.hc_head_scale = nn.Parameter( + torch.empty(1, dtype=torch.float32), + requires_grad=False, + ) + self.hc_head_op = HCHeadOp() + # Pre-hc_head residual stream buffer for the MTP draft. Stable + # address (outside the cudagraph pool) so the copy_ in forward() + # refreshes it correctly across captured shapes. + # refreshes it correctly across captured shapes. Only allocated on + # the last PP rank — that's where MTP target hidden states are + # produced. + if get_pp_group().is_last_rank: + self._mtp_hidden_buffer = torch.empty( + vllm_config.scheduler_config.max_num_batched_tokens, + self.hc_dim, + dtype=vllm_config.model_config.dtype, + device=self.device, + ) + else: + self._mtp_hidden_buffer = None + + self._dspark_target_layer_ids = tuple( + int(layer_id) for layer_id in getattr(config, "dspark_target_layer_ids", ()) + ) + if get_pp_group().is_last_rank and self._dspark_target_layer_ids: + self._dspark_hidden_buffer = torch.empty( + vllm_config.scheduler_config.max_num_batched_tokens, + len(self._dspark_target_layer_ids) * config.hidden_size, + dtype=vllm_config.model_config.dtype, + device=self.device, + ) + self._dspark_layer_to_buffer_index = { + layer_id: idx + for idx, layer_id in enumerate(self._dspark_target_layer_ids) + } + self._dspark_deferred_capture_layer_ids = set() + if _DSPARK_DEFER_TARGET_CAPTURE: + self._setup_dspark_deferred_target_capture(config.hidden_size) + else: + self._dspark_hidden_buffer = None + self._dspark_layer_to_buffer_index = {} + self._dspark_deferred_capture_layer_ids = set() + + def _setup_dspark_deferred_target_capture(self, hidden_size: int) -> None: + assert self._dspark_hidden_buffer is not None + layers_by_id = { + extract_layer_index(layer.layer_name): layer + for layer in self.layers + if hasattr(layer, "layer_name") + } + for target_layer_id, buffer_idx in self._dspark_layer_to_buffer_index.items(): + if target_layer_id < 0: + continue + next_layer = layers_by_id.get(target_layer_id + 1) + if next_layer is None or not hasattr( + next_layer, "set_dspark_previous_layer_capture" + ): + continue + start = buffer_idx * hidden_size + end = start + hidden_size + next_layer.set_dspark_previous_layer_capture( + self._dspark_hidden_buffer, + start, + end, + exact=_DSPARK_DEFER_TARGET_CAPTURE_EXACT, + ) + self._dspark_deferred_capture_layer_ids.add(target_layer_id) + if self._dspark_deferred_capture_layer_ids: + mode = "exact" if _DSPARK_DEFER_TARGET_CAPTURE_EXACT else "fused" + logger.info_once( + "DeepSeek V4 DSpark deferred target-layer capture enabled for " + "layers %s with %s capture.", + tuple(sorted(self._dspark_deferred_capture_layer_ids)), + mode, + ) + + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: + return self.embed_tokens(input_ids) + + def make_empty_intermediate_tensors( + self, + batch_size: int, + dtype: torch.dtype, + device: torch.device, + ) -> IntermediateTensors: + # PP intermediate tensors carry the multi-stream hidden_states + # of shape (num_tokens, hc_mult, hidden_size) — V4 expands the + # token embedding to hc_mult streams before the first decoder + # layer and keeps that shape until hc_head() collapses it. + return IntermediateTensors( + { + "hidden_states": torch.zeros( + (batch_size, self.hc_mult, self.config.hidden_size), + dtype=dtype, + device=device, + ), + } + ) + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + intermediate_tensors: IntermediateTensors | None, + inputs_embeds: torch.Tensor | None = None, + ) -> torch.Tensor | IntermediateTensors: + total_start = _dspark_target_timing_start() + stage_start = _dspark_target_timing_start() + if get_pp_group().is_first_rank: + if inputs_embeds is not None: + hidden_states = inputs_embeds + else: + hidden_states = self.embed_input_ids(input_ids) + hidden_states = hidden_states.unsqueeze(-2).repeat(1, self.hc_mult, 1) + else: + assert intermediate_tensors is not None + hidden_states = intermediate_tensors["hidden_states"] + _dspark_target_timing_record("embed_or_input", stage_start) + + if self.use_mega_moe: + input_ids = input_ids.to(torch.int64) + + residual, post_mix, res_mix = None, None, None + layer = None + num_layers = 0 + for layer in islice(self.layers, self.start_layer, self.end_layer): + num_layers += 1 + layer_start = _dspark_target_timing_start() + hidden_states, residual, post_mix, res_mix = layer( + hidden_states, + positions, + input_ids, + post_mix, + res_mix, + residual, + ) + _dspark_target_timing_record("layers_total", layer_start) + layer_idx = getattr(layer, "layer_name", "") + layer_id = extract_layer_index(layer_idx) if layer_idx else None + if ( + layer_id in self._dspark_layer_to_buffer_index + and layer_id not in self._dspark_deferred_capture_layer_ids + ): + stage_start = _dspark_target_timing_start() + if current_platform.is_cuda() and residual is not None: + hidden_states = layer.hc_post( + hidden_states, residual, post_mix, res_mix + ) + residual, post_mix, res_mix = None, None, None + buffer_idx = self._dspark_layer_to_buffer_index[layer_id] + if self._dspark_hidden_buffer is not None: + start = buffer_idx * self.config.hidden_size + end = start + self.config.hidden_size + self._dspark_hidden_buffer[ + : hidden_states.shape[0], start:end + ].copy_(hidden_states.mean(dim=1)) + _dspark_target_timing_record("dspark_capture", stage_start) + if layer is not None and current_platform.is_cuda() and residual is not None: + stage_start = _dspark_target_timing_start() + hidden_states = layer.hc_post(hidden_states, residual, post_mix, res_mix) + _dspark_target_timing_record("final_hc_post", stage_start) + + if not get_pp_group().is_last_rank: + _dspark_target_timing_finish( + total_start, int(hidden_states.shape[0]), num_layers + ) + return IntermediateTensors({"hidden_states": hidden_states}) + + # Stash pre-hc_head residual for the MTP draft (captured copy_). + num_tokens = hidden_states.shape[0] + stage_start = _dspark_target_timing_start() + self._mtp_hidden_buffer[:num_tokens].copy_(hidden_states.flatten(1)) + _dspark_target_timing_record("mtp_hidden_copy", stage_start) + + stage_start = _dspark_target_timing_start() + hidden_states = self.hc_head_op( + hidden_states, + self.hc_head_fn, + self.hc_head_scale, + self.hc_head_base, + self.rms_norm_eps, + self.hc_eps, + ) + hidden_states = self.norm(hidden_states) + _dspark_target_timing_record("hc_head_norm", stage_start) + _dspark_target_timing_finish(total_start, int(num_tokens), num_layers) + return hidden_states + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + stacked_params_mapping = [ + # (param_name, shard_name, shard_id) + ("gate_up_proj", "w1", 0), + ("gate_up_proj", "w3", 1), + ("attn.fused_wqa_wkv", "attn.wq_a", 0), + ("attn.fused_wqa_wkv", "attn.wkv", 1), + ("compressor.fused_wkv_wgate", "compressor.wkv", 0), + ("compressor.fused_wkv_wgate", "compressor.wgate", 1), + ] + params_dict = dict(self.named_parameters()) + loaded_params: set[str] = set() + + # TP for attention + tp_size = get_tensor_model_parallel_world_size() + tp_rank = get_tensor_model_parallel_rank() + n_head = self.config.num_attention_heads + n_local_head = n_head // tp_size + head_rank_start = n_local_head * tp_rank + head_rank_end = n_local_head * (tp_rank + 1) + + # Pre-compute expert mapping ONCE. + expert_mapping = self.get_expert_mapping() + + for name, loaded_weight in weights: + for param_name, weight_name, shard_id in stacked_params_mapping: + # Skip non-stacked layers and experts (experts handled below). + if ".experts." in name: + continue + if weight_name not in name: + continue + name = name.replace(weight_name, param_name) + + if is_pp_missing_parameter(name, self): + break + param = params_dict[name] + weight_loader = param.weight_loader + weight_loader(param, loaded_weight, shard_id) + loaded_params.add(name) + break + else: + if ".experts." in name: + # E8M0 scales are stored as float8_e8m0fnu in + # checkpoints but the MoE param is uint8. copy_() + # would do a numeric conversion (e.g. 2^-7 → 0), + # destroying the raw exponent bytes. + if ( + "weight_scale" in name + and loaded_weight.dtype == torch.float8_e8m0fnu + ): + loaded_weight = loaded_weight.view(torch.uint8) + for mapping in expert_mapping: + param_name, weight_name, expert_id, expert_shard_id = mapping + if weight_name not in name: + continue + name_mapped = name.replace(weight_name, param_name) + if is_pp_missing_parameter(name_mapped, self): + continue + param = params_dict[name_mapped] + # We should ask the weight loader to return success or not + # here since otherwise we may skip experts with other + # available replicas. + weight_loader = typing.cast( + Callable[..., bool], param.weight_loader + ) + success = weight_loader( + param, + loaded_weight, + name_mapped, + shard_id=expert_shard_id, + expert_id=expert_id, + return_success=True, + ) + if success: + name = name_mapped + break + loaded_params.add(name_mapped) + continue + elif "attn_sink" in name: + if is_pp_missing_parameter(name, self): + continue + narrow_weight = loaded_weight[head_rank_start:head_rank_end] + n = narrow_weight.shape[0] + params_dict[name][:n].copy_(narrow_weight) + loaded_params.add(name) + continue + else: + if is_pp_missing_parameter(name, self): + continue + param = params_dict[name] + weight_loader = getattr( + param, "weight_loader", default_weight_loader + ) + weight_loader(param, loaded_weight) + loaded_params.add(name) + continue + + return loaded_params + + def get_expert_mapping(self) -> list[tuple[str, str, int, str]]: + first_layer = next(iter(islice(self.layers, self.start_layer, self.end_layer))) + if first_layer.ffn.use_mega_moe: + return make_deepseek_v4_expert_params_mapping(self.config.n_routed_experts) + # Params for weights, fp8 weight scales, fp8 activation scales + # (param_name, weight_name, expert_id, shard_id) + return FusedMoE.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, + ) + + 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 setup_b12x_wo_projection(self) -> None: + if not envs.VLLM_USE_B12X_WO_PROJECTION: + return + for layer in islice(self.layers, self.start_layer, self.end_layer): + layer.attn.setup_b12x_wo_projection() + + +def _make_deepseek_v4_weights_mapper(expert_dtype: str) -> WeightsMapper: + if expert_dtype == "fp4": + # MXFP4 experts use Mxfp4MoEMethod, which registers scales as + # ``w{1,2,3}_weight_scale`` (no _inv suffix). FP8 linear and + # shared experts use Fp8LinearMethod's block scales, which + # register as ``weight_scale_inv``. + scale_regex = { + re.compile(r"(\.experts\.\d+\.w[123])\.scale$"): r"\1.weight_scale", + re.compile(r"\.scale$"): ".weight_scale_inv", + } + else: + # FP8 experts use Fp8MoEMethod (block_quant=True), which registers + # scales as ``w{13,2}_weight_scale_inv``. Map all ``.scale`` keys + # there. + scale_regex = { + re.compile(r"\.scale$"): ".weight_scale_inv", + } + return WeightsMapper( + orig_to_new_prefix={ + "layers.": "model.layers.", + "embed.": "model.embed.", + "norm.": "model.norm.", + "hc_head": "model.hc_head", + "mtp.": "model.mtp.", + }, + orig_to_new_regex=scale_regex, + orig_to_new_suffix={ + "head.weight": "lm_head.weight", + "embed.weight": "embed_tokens.weight", + ".ffn.gate.bias": ".ffn.gate.e_score_correction_bias", + }, + orig_to_new_substr={ + ".attn.compressor.": ".attn.mla_attn.compressor.", + ".shared_experts.w2": ".shared_experts.down_proj", + }, + ) + + +class DeepseekV4ForCausalLM(nn.Module, SupportsPP): + model_cls = DeepseekV4Model + + # Default mapper assumes the original FP4-expert checkpoint layout. + # Overridden per-instance in __init__ when expert_dtype != "fp4". + hf_to_vllm_mapper = _make_deepseek_v4_weights_mapper("fp4") + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): + super().__init__() + + config = vllm_config.model_config.hf_config + self.config = config + expert_dtype = getattr(config, "expert_dtype", "fp4") + if expert_dtype != "fp4": + self.hf_to_vllm_mapper = _make_deepseek_v4_weights_mapper(expert_dtype) + + self.model = self.model_cls( + vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model") + ) + if get_pp_group().is_last_rank: + self.lm_head = ParallelLMHead( + config.vocab_size, + config.hidden_size, + prefix=maybe_prefix(prefix, "lm_head"), + ) + else: + self.lm_head = PPMissingLayer() + self.logits_processor = LogitsProcessor(config.vocab_size) + self.make_empty_intermediate_tensors = ( # type: ignore[method-assign] + self.model.make_empty_intermediate_tensors + ) + + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: + return self.model.embed_input_ids(input_ids) + + def compute_logits( + self, + hidden_states: torch.Tensor, + ) -> torch.Tensor | None: + logits = self.logits_processor(self.lm_head, hidden_states) + return logits + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + intermediate_tensors: IntermediateTensors | None = None, + inputs_embeds: torch.Tensor | None = None, + ) -> torch.Tensor | IntermediateTensors: + hidden_states = self.model( + input_ids, positions, intermediate_tensors, inputs_embeds + ) + return hidden_states + + def get_mtp_target_hidden_states(self) -> torch.Tensor | None: + """Pre-hc_head residual stream buffer (max_num_batched_tokens, + hc_mult * hidden_size) for the MTP draft model. Populated by + forward(); valid after each target step.""" + return getattr(self.model, "_mtp_hidden_buffer", None) + + def get_dspark_target_hidden_states(self) -> torch.Tensor | None: + """Concatenated DSpark target-layer features for the draft model. + + Shape is (max_num_batched_tokens, + len(dspark_target_layer_ids) * hidden_size). Populated by forward(). + """ + return getattr(self.model, "_dspark_hidden_buffer", None) + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + loader = AutoWeightsLoader(self, skip_substrs=["mtp."]) + loaded_params = loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) + self.model.finalize_mega_moe_weights() + self.model.setup_b12x_wo_projection() + return loaded_params + + def get_expert_mapping(self) -> list[tuple[str, str, int, str]]: + return self.model.get_expert_mapping() diff --git a/mods/deepseek-v4-flash-dspark/overlay/vllm/models/deepseek_v4/nvidia/sm120.py b/mods/deepseek-v4-flash-dspark/overlay/vllm/models/deepseek_v4/nvidia/sm120.py new file mode 100644 index 00000000..2e3100da --- /dev/null +++ b/mods/deepseek-v4-flash-dspark/overlay/vllm/models/deepseek_v4/nvidia/sm120.py @@ -0,0 +1,519 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""SM120 (consumer Blackwell) sparse-MLA impl for DeepSeek-V4. + +Counterpart to :class:`DeepseekV4FlashMLASparseImpl` (Hopper / SM10x). The +forward path is driven by flashinfer's :class:`BatchSparseMLAPagedAttention +Wrapper` — the same wrapper used by the V32-family SPARSE_MLA_SM120 backend — +which auto-dispatches decode (num_tokens <= 64) and prefill internally and +accepts the SWA + compressed-indexer dual cache through its ``extra_kv_cache`` +parameter. Decode scratch is borrowed from vLLM's shared workspace so large +C128A contexts do not allocate per-layer split-K buffers. + +Selected by ``_select_v4_sparse_impl()`` in :mod:`vllm.models.deepseek_v4 +.attention` when the runtime compute capability is SM120; the +flashinfer wrapper itself lives on the layer (``layer._sparse_mla_wrapper``) +only for its reusable LSE buffer; split-K decode scratch is supplied per call. +""" + +import os +from typing import TYPE_CHECKING, ClassVar, cast + +import torch + +from vllm.forward_context import get_forward_context +from vllm.logger import init_logger +from vllm.models.deepseek_v4.common.ops import ( + compute_global_topk_indices_and_lens, +) +from vllm.models.deepseek_v4.nvidia.flashmla import ( + DeepseekV4FlashMLASparseBackend, + DeepseekV4SparseMLAAttentionImpl, +) +from vllm.v1.attention.backend import AttentionBackend +from vllm.v1.attention.backends.mla.flashmla_sparse import FlashMLASparseMetadata +from vllm.v1.worker.workspace import current_workspace_manager + +if TYPE_CHECKING: + from vllm.models.deepseek_v4.attention import DeepseekV4MLAAttention + from vllm.v1.attention.backends.mla.sparse_swa import DeepseekSparseSWAMetadata + +logger = init_logger(__name__) + + +_DECODE_MAX_TOKENS = 64 +_DECODE_SPLIT_TILE = 64 +_C128A_TOPK_ALIGNMENT = 128 + + +def _cdiv(x: int, y: int) -> int: + return (int(x) + int(y) - 1) // int(y) + + +def _decode_num_splits(topk: int, extra_topk: int = 0) -> int: + return _cdiv(topk, _DECODE_SPLIT_TILE) + _cdiv(extra_topk, _DECODE_SPLIT_TILE) + + +def _max_decode_workspace_tokens(max_num_batched_tokens: int) -> int: + return min(int(max_num_batched_tokens), _DECODE_MAX_TOKENS) + + +def _c128a_max_compressed(max_model_len: int, compress_ratio: int) -> int: + return ( + _cdiv( + _cdiv(max_model_len, compress_ratio), + _C128A_TOPK_ALIGNMENT, + ) + * _C128A_TOPK_ALIGNMENT + ) + + +def _env_enabled(name: str, default: str = "0") -> bool: + value = os.getenv(name, default).strip().lower() + return value not in ("0", "false", "no", "off", "") + + +def _use_b12x_compressed_mla() -> bool: + return _env_enabled("VLLM_DSV4_B12X_COMPRESSED_MLA") + + +def _extra_topk_capacity(layer: "DeepseekV4MLAAttention") -> int: + if layer.compress_ratio <= 1: + return 0 + if layer.compress_ratio == 4: + assert layer.topk_indices_buffer is not None + return int(layer.topk_indices_buffer.shape[-1]) + if layer.compress_ratio == 128: + return _c128a_max_compressed(layer.max_model_len, layer.compress_ratio) + raise ValueError( + f"Unsupported compress_ratio={layer.compress_ratio}; " + "expected 1, 4, or 128." + ) + + +def _get_decode_scratch( + num_tokens: int, + num_heads: int, + d_v: int, + topk: int, + extra_topk: int, +) -> tuple[torch.Tensor, torch.Tensor]: + num_splits = _decode_num_splits(topk, extra_topk) + mid_out, mid_lse = current_workspace_manager().get_simultaneous( + ((num_tokens, num_heads, num_splits, d_v), torch.bfloat16), + ((num_tokens, num_heads, num_splits), torch.float32), + ) + return mid_out, mid_lse + + +def _b12x_index_matrix(indices: torch.Tensor | None) -> torch.Tensor | None: + if indices is None: + return None + if indices.ndim == 3: + assert indices.shape[1] == 1 + return indices.squeeze(1) + return indices + + +def _get_b12x_decode_workspace( + layer: "DeepseekV4MLAAttention", + *, + extra_topk: int, +): + from b12x.attention.workspace import B12XAttentionWorkspace + + total_topk = int(layer.window_size) + int(extra_topk) + max_rows = _max_decode_workspace_tokens(layer.max_num_batched_tokens) + max_chunks = _decode_num_splits(layer.window_size, extra_topk) + + workspace = getattr(layer, "_b12x_compressed_mla_workspace", None) + if ( + workspace is None + or int(workspace.topk) < total_topk + or int(workspace.max_total_q) < max_rows + or int(workspace.max_chunks_per_row) < max_chunks + or int(workspace.num_q_heads) != int(layer.padded_heads) + ): + device = layer.attn_sink.device + if device.type != "cuda": + device = torch.device(f"cuda:{torch.cuda.current_device()}") + workspace = B12XAttentionWorkspace( + mode="decode", + device=device, + dtype=torch.bfloat16, + kv_dtype=torch.uint8, + num_q_heads=int(layer.padded_heads), + head_dim=512, + v_head_dim=512, + topk=total_topk, + max_total_q=max_rows, + max_batch=max_rows, + max_page_table_width=total_topk, + max_paged_q_rows=max_rows, + page_size=int(layer.swa_cache_layer.block_size), + padded_heads=int(layer.padded_heads), + max_chunks_per_row=max_chunks, + ) + workspace.kv_chunk_size_ptr = torch.empty( + (1,), dtype=torch.int32, device=device + ) + workspace.num_chunks_ptr = torch.empty((1,), dtype=torch.int32, device=device) + layer._b12x_compressed_mla_workspace = workspace + logger.info_once( + "DeepSeek V4 SM120 b12x compressed MLA decode enabled " + "(topk=%d, max_rows=%d, max_chunks=%d).", + total_topk, + max_rows, + max_chunks, + ) + return workspace + + +class DeepseekV4SM120SparseBackend(DeepseekV4FlashMLASparseBackend): + """SM120 variant. Geometry is identical to the FlashMLA parent (same KV + layout, head size, block size); the only thing that changes is the impl + class returned by ``get_impl_cls``.""" + + @staticmethod + def get_name() -> str: + return "DSV4_SPARSE_MLA_SM120" + + @staticmethod + def get_impl_cls() -> type["DeepseekV4SM120SparseImpl"]: + return DeepseekV4SM120SparseImpl + + +class DeepseekV4SM120SparseImpl(DeepseekV4SparseMLAAttentionImpl): + """SM120 flashinfer-wrapper-driven sparse-MLA impl for DeepseekV4. + + The wrapper auto-dispatches decode (num_tokens <= 64) and prefill on + num_tokens, so this impl issues a single ``wrapper.run`` per chunk — + no separate prefill kernel call, no plan() step. + """ + + backend_cls: ClassVar[type[AttentionBackend]] = DeepseekV4SM120SparseBackend + + @classmethod + def get_padded_num_q_heads(cls, num_heads: int) -> int: + if num_heads <= 16: + return 16 + if num_heads <= 32: + return 32 + if num_heads <= 64: + return 64 + if num_heads <= 128: + return 128 + raise ValueError( + f"DeepseekV4 SM120 sparse MLA does not support {num_heads} heads " + "(kernel requires h_q in {16, 32, 64, 128})." + ) + + @classmethod + def forward_mqa( # type: ignore[override] + cls, + layer: "DeepseekV4MLAAttention", + q: torch.Tensor, + kv: torch.Tensor, + positions: torch.Tensor, + output: torch.Tensor, + ) -> None: + assert output.shape == q.shape, ( + f"output buffer shape {output.shape} must match q shape {q.shape}" + ) + assert output.dtype == q.dtype, ( + f"output buffer dtype {output.dtype} must match q dtype {q.dtype}" + ) + + forward_context = get_forward_context() + attn_metadata = forward_context.attn_metadata + if attn_metadata is None: + cls._reserve_decode_workspace(layer) + output.zero_() + return + + assert isinstance(attn_metadata, dict) + flashmla_metadata = cast( + FlashMLASparseMetadata | None, attn_metadata.get(layer.prefix) + ) + swa_metadata = cast( + "DeepseekSparseSWAMetadata | None", + attn_metadata.get(layer.swa_cache_layer.prefix), + ) + assert swa_metadata is not None + + swa_only = layer.compress_ratio <= 1 + # SWA-only layers (compress_ratio <= 1) don't have their own KV cache + # allocation; layer.kv_cache may be empty after profiling cleanup. + self_kv_cache = layer.kv_cache if not swa_only else None + swa_kv_cache = layer.swa_cache_layer.kv_cache + + num_decodes = swa_metadata.num_decodes + num_prefills = swa_metadata.num_prefills + num_decode_tokens = swa_metadata.num_decode_tokens + + if num_prefills > 0: + cls._forward_prefill( + layer=layer, + q=q[num_decode_tokens:], + compressed_k_cache=self_kv_cache, + swa_k_cache=swa_kv_cache, + output=output[num_decode_tokens:], + attn_metadata=flashmla_metadata, + swa_metadata=swa_metadata, + ) + if num_decodes > 0: + cls._forward_decode( + layer=layer, + q=q[:num_decode_tokens], + kv_cache=self_kv_cache, + swa_metadata=swa_metadata, + attn_metadata=flashmla_metadata, + swa_only=swa_only, + output=output[:num_decode_tokens], + ) + + @classmethod + def _reserve_decode_workspace(cls, layer: "DeepseekV4MLAAttention") -> None: + extra_topk = _extra_topk_capacity(layer) + _get_decode_scratch( + _max_decode_workspace_tokens(layer.max_num_batched_tokens), + layer.padded_heads, + 512, + layer.window_size, + extra_topk, + ) + if _use_b12x_compressed_mla(): + _get_b12x_decode_workspace(layer, extra_topk=extra_topk) + + @classmethod + def _forward_decode( + cls, + layer: "DeepseekV4MLAAttention", + q: torch.Tensor, + kv_cache: torch.Tensor | None, # only used when compress_ratio > 1 + swa_metadata: "DeepseekSparseSWAMetadata", + attn_metadata: FlashMLASparseMetadata | None, + swa_only: bool, + output: torch.Tensor, + ) -> None: + num_decodes = swa_metadata.num_decodes + num_decode_tokens = swa_metadata.num_decode_tokens + + topk_indices = None + topk_lens = None + if not swa_only: + assert attn_metadata is not None + assert swa_metadata.is_valid_token is not None + block_size = attn_metadata.block_size // layer.compress_ratio + is_valid = swa_metadata.is_valid_token[:num_decode_tokens] + if layer.compress_ratio == 4: + # C4A: local indices differ per layer (filled by Indexer). + assert layer.topk_indices_buffer is not None + global_indices, topk_lens = compute_global_topk_indices_and_lens( + layer.topk_indices_buffer[:num_decode_tokens], + swa_metadata.token_to_req_indices, + attn_metadata.block_table[:num_decodes], + block_size, + is_valid, + kv_cache.shape[0], + ) + topk_indices = global_indices.view(num_decode_tokens, 1, -1) + else: + # C128A: pre-computed during metadata build. + topk_indices = attn_metadata.c128a_global_decode_topk_indices + topk_lens = attn_metadata.c128a_decode_topk_lens + + swa_indices = swa_metadata.decode_swa_indices + swa_lens = swa_metadata.decode_swa_lens + assert swa_indices is not None + assert swa_lens is not None + extra_topk = topk_indices.shape[-1] if topk_indices is not None else 0 + mid_out, mid_lse = _get_decode_scratch( + num_decode_tokens, + q.shape[1], + output.shape[-1], + swa_indices.shape[-1], + extra_topk, + ) + + # Treat queries in the same seq as independent queries (attended + # purely by the generated indices). q arrives pre-padded to + # layer.padded_heads by the outer wrapper. + if _use_b12x_compressed_mla(): + from b12x.attention.mla.compressed_api import ( + compressed_mla_decode_forward, + ) + + workspace = _get_b12x_decode_workspace(layer, extra_topk=extra_topk) + workspace.tmp_output = mid_out + workspace.tmp_lse = mid_lse + workspace.output_buffer = output + result = compressed_mla_decode_forward( + q_all=q, + swa_k_cache=layer.swa_cache_layer.kv_cache, + swa_indices=_b12x_index_matrix(swa_indices), + swa_topk_lengths=swa_lens, + workspace=workspace, + sm_scale=layer.scale, + swa_page_size=swa_metadata.block_size, + indexed_k_cache=kv_cache, + indexed_indices=_b12x_index_matrix(topk_indices), + indexed_topk_lengths=topk_lens, + indexed_page_size=block_size if kv_cache is not None else None, + attn_sink=layer.attn_sink, + expected_num_q_heads=q.shape[1], + backend="sm120_unified", + ) + if result.data_ptr() != output.data_ptr(): + output.copy_(result) + return + + q = q.unsqueeze(1) + swa_cache = layer.swa_cache_layer.kv_cache.unsqueeze(-2) + if kv_cache is not None: + kv_cache = kv_cache.unsqueeze(-2) + + assert layer._sparse_mla_wrapper is not None, ( + "DeepseekV4SM120SparseImpl requires layer._sparse_mla_wrapper; " + "the flashinfer wrapper must be constructed in the layer __init__." + ) + layer._sparse_mla_wrapper.run( + q=q, + kv_cache=swa_cache, + indices=swa_indices, + output=output, + sm_scale=layer.scale, + topk_length=swa_lens, + attn_sink=layer.attn_sink, + extra_kv_cache=kv_cache if not swa_only else None, + extra_indices=topk_indices, + extra_topk_length=topk_lens, + mid_out=mid_out, + mid_lse=mid_lse, + ) + + @classmethod + def _forward_prefill( + cls, + layer: "DeepseekV4MLAAttention", + q: torch.Tensor, + compressed_k_cache: torch.Tensor | None, + swa_k_cache: torch.Tensor, + output: torch.Tensor, + attn_metadata: FlashMLASparseMetadata | None, + swa_metadata: "DeepseekSparseSWAMetadata", + ) -> None: + # `_dummy_run` passes synthetic non-None attn_metadata for swa-only + # layers during cudagraph capture, so check compress_ratio directly. + swa_only = layer.compress_ratio <= 1 + + num_prefills = swa_metadata.num_prefills + num_decodes = swa_metadata.num_decodes + num_decode_tokens = swa_metadata.num_decode_tokens + num_prefill_tokens = swa_metadata.num_prefill_tokens + + # Derive prefill-local token offsets from the full query_start_loc_cpu. + query_start_loc_cpu = swa_metadata.query_start_loc_cpu + assert query_start_loc_cpu is not None + prefill_token_base = query_start_loc_cpu[num_decodes] + + local_topk_indices: torch.Tensor | None + if swa_only: + local_topk_indices = None + elif layer.compress_ratio == 4: + assert layer.topk_indices_buffer is not None + local_topk_indices = layer.topk_indices_buffer[ + num_decode_tokens : num_decode_tokens + num_prefill_tokens + ] + else: + # C128A: pre-computed during metadata build. + assert attn_metadata is not None + local_topk_indices = attn_metadata.c128a_prefill_topk_indices + + extra_topk_indices: torch.Tensor | None = None + extra_topk_lens: torch.Tensor | None = None + if local_topk_indices is not None: + assert attn_metadata is not None + assert swa_metadata.token_to_req_indices is not None + assert swa_metadata.is_valid_token is not None + prefill_token_slice = slice( + num_decode_tokens, num_decode_tokens + num_prefill_tokens + ) + # FlashInfer prefill expects physical KV slots; keep padding rows + # masked through the metadata validity mask. + block_size = attn_metadata.block_size // layer.compress_ratio + extra_topk_indices, extra_topk_lens = compute_global_topk_indices_and_lens( + local_topk_indices, + swa_metadata.token_to_req_indices[prefill_token_slice], + attn_metadata.block_table, + block_size, + swa_metadata.is_valid_token[prefill_token_slice], + compressed_k_cache.shape[0], + ) + + assert swa_metadata.prefill_swa_indices is not None + assert swa_metadata.prefill_swa_lens is not None + assert layer._sparse_mla_wrapper is not None + + # unsqueeze(-2) adds the h_kv=1 axis without copying. + swa_kv_paged = swa_k_cache.unsqueeze(-2) + if swa_only: + extra_kv_paged = None + else: + assert compressed_k_cache is not None + extra_kv_paged = compressed_k_cache.unsqueeze(-2) + + num_chunks = ( + num_prefills + cls.PREFILL_CHUNK_SIZE - 1 + ) // cls.PREFILL_CHUNK_SIZE + for chunk_idx in range(num_chunks): + chunk_start = chunk_idx * cls.PREFILL_CHUNK_SIZE + chunk_end = min(chunk_start + cls.PREFILL_CHUNK_SIZE, num_prefills) + query_start = ( + query_start_loc_cpu[num_decodes + chunk_start] - prefill_token_base + ) + query_end = ( + query_start_loc_cpu[num_decodes + chunk_end] - prefill_token_base + ) + + extra_indices_chunk = ( + extra_topk_indices[query_start:query_end] + if extra_topk_indices is not None + else None + ) + extra_topk_length_chunk = ( + extra_topk_lens[query_start:query_end] + if extra_topk_lens is not None + else None + ) + chunk_tokens = query_end - query_start + mid_out = None + mid_lse = None + if chunk_tokens <= _DECODE_MAX_TOKENS: + extra_topk = ( + extra_indices_chunk.shape[-1] + if extra_indices_chunk is not None + else 0 + ) + mid_out, mid_lse = _get_decode_scratch( + chunk_tokens, + q.shape[1], + output.shape[-1], + swa_metadata.prefill_swa_indices.shape[-1], + extra_topk, + ) + + layer._sparse_mla_wrapper.run( + q=q[query_start:query_end], + kv_cache=swa_kv_paged, + indices=swa_metadata.prefill_swa_indices[query_start:query_end], + output=output[query_start:query_end], + sm_scale=layer.scale, + topk_length=swa_metadata.prefill_swa_lens[query_start:query_end], + attn_sink=layer.attn_sink, + extra_kv_cache=extra_kv_paged, + extra_indices=extra_indices_chunk, + extra_topk_length=extra_topk_length_chunk, + mid_out=mid_out, + mid_lse=mid_lse, + ) diff --git a/mods/deepseek-v4-flash-dspark/overlay/vllm/v1/attention/backends/mla/b12x_mla_sparse.py b/mods/deepseek-v4-flash-dspark/overlay/vllm/v1/attention/backends/mla/b12x_mla_sparse.py new file mode 100644 index 00000000..f40969d7 --- /dev/null +++ b/mods/deepseek-v4-flash-dspark/overlay/vllm/v1/attention/backends/mla/b12x_mla_sparse.py @@ -0,0 +1,501 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""b12x sparse-MLA backend for SM120 / SM121 (consumer Blackwell). + +Counterpart to ``SparseMLASm120Backend`` (FlashInfer V32 v2). Same envelope -- +``fp8_ds_mla`` KV cache (656 B/token), head_size = 576, paged block_size = 64, +V32-family models with an ``index_topk`` config (DeepSeek V3.2, GLM-5.1, Kimi +K2.5) -- but the decode/extend kernels come from b12x's unified SM120 backend +via the ``b12x.integration.mla`` front door (``sparse_mla_decode_forward`` / +``sparse_mla_extend_forward``). On SM120+ CUDA those front-door functions route +to ``b12x/attention/mla/unified_sm120`` automatically (GLM_NSA q_head_dim==576 +contract). + +This backend is **opt-in**: it is not in the platform auto-selection priority +list, so it only runs when explicitly requested via +``VLLM_ATTENTION_BACKEND=B12X_MLA_SPARSE``. The FlashInfer ``SPARSE_MLA_SM120`` +backend is left intact for A/B comparison. + +Workspace philosophy (the idiomatic, no-arena path): b12x's kernels take a +``B12XAttentionWorkspace`` object, but they only read it as a bag of tensor +attributes (``tmp_output`` / ``tmp_lse`` / ``output_buffer`` + control pointers) +plus ``set_split_chunk_config``. We therefore construct a bare workspace +dataclass (which allocates nothing) and back its split-K scratch with tensors +borrowed per-call from vLLM's shared ``current_workspace_manager()`` -- exactly +how ``SparseMLASm120Impl`` borrows ``mid_out``/``mid_lse``. No per-layer b12x +arena is allocated. +""" + +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any, ClassVar, cast + +import numpy as np +import torch + +from vllm.config import VllmConfig +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.platforms.interface import DeviceCapability +from vllm.v1.attention.backend import ( + AttentionBackend, + AttentionCGSupport, + AttentionLayer, + AttentionMetadata, + AttentionMetadataBuilder, + AttentionType, + CommonAttentionMetadata, + MultipleOf, + SparseMLAAttentionImpl, +) +from vllm.v1.attention.backends.mla.sparse_utils import ( + triton_convert_req_index_to_global_index, +) +from vllm.v1.kv_cache_interface import AttentionSpec +from vllm.v1.worker.workspace import current_workspace_manager + +if TYPE_CHECKING: + from vllm.model_executor.models.deepseek_v2 import Indexer + +logger = init_logger(__name__) + +# Split-K tile width. Mirrors SparseMLASm120's _DECODE_SPLIT_TILE: the number of +# split-K chunks is ceil(topk / tile). This bounds the chunk dim of the borrowed +# mid_out/mid_lse scratch and the workspace ``max_chunks_per_row`` cap; b12x's +# wave-balanced planner picks num_splits <= this cap. +_DECODE_SPLIT_TILE = 64 + + +def _cdiv(x: int, y: int) -> int: + return (int(x) + int(y) - 1) // int(y) + + +class B12xMLASparseBackend(AttentionBackend): + """b12x unified sparse-MLA backend (SM120 / SM121). + + Same envelope as ``SparseMLASm120Backend`` (head 576, fp8_ds_mla, block 64, + index_topk) but driven by b12x's unified decode/extend kernels. + """ + + supported_dtypes: ClassVar[list[torch.dtype]] = [torch.bfloat16] + supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [ + "auto", + "bfloat16", + "fp8_ds_mla", + "fp8", # alias for fp8_ds_mla on this backend (auto-converted by MLAAttention) + ] + + @staticmethod + def get_supported_kernel_block_sizes() -> list[int | MultipleOf]: + # Must equal DeepseekV32IndexerBackend.get_supported_kernel_block_sizes + # on CUDA (= [64]); the unified b12x decode/extend kernels dispatch + # page_block_size == 64 natively (matches the fp8_ds_mla layout). + return [64] + + @staticmethod + def get_name() -> str: + return "B12X_MLA_SPARSE" + + @staticmethod + def get_impl_cls() -> type["B12xMLASparseImpl"]: + return B12xMLASparseImpl + + @staticmethod + def get_builder_cls() -> type["B12xMLASparseMetadataBuilder"]: + return B12xMLASparseMetadataBuilder + + @classmethod + def get_supported_head_sizes(cls) -> list[int]: + # GLM_NSA contract: q_head_dim = kv_lora_rank (512) + qk_rope_head_dim + # (64) = 576. The unified decode raises on any other q_head_dim. + return [576] + + @classmethod + def is_mla(cls) -> bool: + return True + + @classmethod + def is_sparse(cls) -> bool: + return True + + @classmethod + def supports_compute_capability(cls, capability: DeviceCapability) -> bool: + # Consumer Blackwell SM120 / SM121. The unified b12x kernels gate on + # get_sm_version(device) >= 120 internally. + return capability.major == 12 + + @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, + device_capability: DeviceCapability, + ) -> str | None: + # Require an indexer-equipped (index_topk) model, same as SPARSE_MLA_SM120. + from vllm.config import get_current_vllm_config + + vllm_config = get_current_vllm_config() + if vllm_config.model_config is not None: + hf_text_config = vllm_config.model_config.hf_text_config + if not hasattr(hf_text_config, "index_topk"): + return "B12X_MLA_SPARSE requires a model with index_topk config" + return None + + @staticmethod + def get_kv_cache_shape( + num_blocks: int, + block_size: int, + num_kv_heads: int, # = 1 for MLA + head_size: int, + cache_dtype_str: str = "auto", + ) -> tuple[int, ...]: + if cache_dtype_str == "fp8_ds_mla": + # V32 fp8_ds_mla packed: 656 B/token (512 NoPE + 16 inline FP32 + # scales + 128 BF16 RoPE). Mirrors the FlashMLA / SPARSE_MLA_SM120 + # layout; b12x's GLM_NSA decode reads the same record. + return (num_blocks, block_size, 656) + return (num_blocks, block_size, head_size) + + +@dataclass +class B12xMLASparseMetadata(AttentionMetadata): + """Attention metadata for the B12X_MLA_SPARSE backend.""" + + num_reqs: int + max_query_len: int + max_seq_len: int + num_actual_tokens: int + + query_start_loc: torch.Tensor + slot_mapping: torch.Tensor + block_table: torch.Tensor + req_id_per_token: torch.Tensor + # Per-request computed KV length (decode cache_seqlens_int32). + seq_lens: torch.Tensor + # Per-token causal KV length; clamped to topk to form nsa_cache_seqlens. + # For pure decode this equals ``seq_lens`` (one token per request). + cache_seq_lens_per_token: torch.Tensor + + block_size: int = 64 + topk_tokens: int = 2048 + + +class B12xMLASparseMetadataBuilder(AttentionMetadataBuilder[B12xMLASparseMetadata]): + """Builder for B12X_MLA_SPARSE attention metadata.""" + + _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.UNIFORM_BATCH + + def __init__( + self, + kv_cache_spec: AttentionSpec, + layer_names: list[str], + vllm_config: VllmConfig, + device: torch.device, + ) -> None: + self.vllm_config = vllm_config + self.layer_names = layer_names + self.kv_cache_spec = kv_cache_spec + self.model_config = vllm_config.model_config + self.device = device + + self.mla_dims = get_mla_dims(self.model_config) + self.topk_tokens = vllm_config.model_config.hf_config.index_topk + + max_tokens = vllm_config.scheduler_config.max_num_batched_tokens + # Max-batched-token scratch buffers so cudagraph capture sees stable + # allocations (sliced per build()). + self.req_id_per_token_buffer = torch.empty( + (max_tokens,), dtype=torch.int32, device=device + ) + self.cache_seq_lens_per_token_buffer = torch.empty( + (max_tokens,), dtype=torch.int32, device=device + ) + + def build( + self, + common_prefix_len: int, + common_attn_metadata: CommonAttentionMetadata, + fast_build: bool = False, + ) -> B12xMLASparseMetadata: + cm = common_attn_metadata + num_tokens = cm.num_actual_tokens + + starts = np.asarray(cm.query_start_loc_cpu, dtype=np.int32) + seg_lengths = np.diff(starts) + req_id_per_token = np.repeat( + np.arange(seg_lengths.shape[0], dtype=np.int32), seg_lengths + ) + + self.req_id_per_token_buffer.fill_(0) + self.req_id_per_token_buffer[: req_id_per_token.shape[0]].copy_( + torch.from_numpy(req_id_per_token), non_blocking=True + ) + req_id_per_token_tensor = self.req_id_per_token_buffer[:num_tokens] + + # Per-token causal KV length. Hot path (pure decode, one token per req): + # the per-token length is just the per-request seq_len -- no expansion. + if cm.max_query_len <= 1 and num_tokens == cm.num_reqs: + cache_seq_lens_per_token = cm.seq_lens[:num_tokens] + else: + # Prefill / mixed: token at within-query offset i in a request with + # ``num_computed`` already-cached tokens has causal KV length + # ``num_computed + i + 1``. Computed entirely on device (no H<->D + # sync); prefill is not cudagraph-captured but this is capture-safe + # regardless. + num_computed = cm.compute_num_computed_tokens() # (num_reqs,) device + req = req_id_per_token_tensor.to(torch.long) + arange = torch.arange(num_tokens, device=self.device, dtype=torch.int32) + within = arange - cm.query_start_loc[:-1].to(torch.int32)[req] + per_token = num_computed.to(torch.int32)[req] + within + 1 + self.cache_seq_lens_per_token_buffer[:num_tokens].copy_( + per_token, non_blocking=True + ) + cache_seq_lens_per_token = self.cache_seq_lens_per_token_buffer[:num_tokens] + + return B12xMLASparseMetadata( + num_reqs=cm.num_reqs, + max_query_len=cm.max_query_len, + max_seq_len=cm.max_seq_len, + num_actual_tokens=num_tokens, + query_start_loc=cm.query_start_loc, + slot_mapping=cm.slot_mapping, + block_table=cm.block_table_tensor, + req_id_per_token=req_id_per_token_tensor, + seq_lens=cm.seq_lens, + cache_seq_lens_per_token=cache_seq_lens_per_token, + block_size=self.kv_cache_spec.block_size, + topk_tokens=self.topk_tokens, + ) + + +class B12xMLASparseImpl(SparseMLAAttentionImpl[B12xMLASparseMetadata]): + """b12x unified sparse-MLA implementation (decode + extend/prefill).""" + + 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, + attn_type: str, + kv_sharing_target_layer_name: str | None, + topk_indice_buffer: torch.Tensor | None = None, + indexer: "Indexer | None" = None, + **mla_args, + ) -> None: + if any([alibi_slopes, sliding_window, logits_soft_cap]): + raise NotImplementedError( + "B12X_MLA_SPARSE does not support alibi_slopes / sliding_window " + "/ logits_soft_cap" + ) + if attn_type != AttentionType.DECODER: + raise NotImplementedError( + "B12X_MLA_SPARSE only supports decoder self-attention" + ) + + self.num_heads = num_heads + self.head_size = head_size + self.scale = float(scale) + self.num_kv_heads = num_kv_heads + self.kv_cache_dtype = kv_cache_dtype + + # MLA dims (absorbed: Q post-projection is [T, H, kv_lora_rank + rope]). + self.kv_lora_rank: int = mla_args["kv_lora_rank"] + self.qk_nope_head_dim: int = mla_args["qk_nope_head_dim"] + self.qk_rope_head_dim: int = mla_args["qk_rope_head_dim"] + self.v_head_dim: int = mla_args.get("v_head_dim", 512) + # GLM_NSA contract: q_head_dim = kv_lora_rank (512) + qk_rope (64) = 576. + self.q_head_dim = self.kv_lora_rank + self.qk_rope_head_dim + + assert indexer is not None, ( + "B12X_MLA_SPARSE requires a sparse-MLA indexer (model with " + "index_topk in its config)." + ) + self.topk_indices_buffer: torch.Tensor | None = indexer.topk_indices_buffer + assert self.topk_indices_buffer is not None + self.topk_tokens = int(self.topk_indices_buffer.shape[-1]) + + from vllm.config import get_current_vllm_config + + vllm_config = get_current_vllm_config() + scheduler_config = vllm_config.scheduler_config + self.device = torch.device(f"cuda:{torch.cuda.current_device()}") + max_batched = int(scheduler_config.max_num_batched_tokens) + max_num_seqs = int(scheduler_config.max_num_seqs) + self.block_size = 64 + + # Split-K cap: ceil(topk / tile). Bounds the borrowed mid_out/mid_lse + # chunk dim and the workspace max_chunks_per_row. + self._num_splits_cap = max(1, _cdiv(self.topk_tokens, _DECODE_SPLIT_TILE)) + + # Decode query rows per request (1, plus speculative draft tokens). + q_per_req = 1 + spec = getattr(vllm_config, "speculative_config", None) + if spec is not None and getattr(spec, "num_speculative_tokens", None): + q_per_req = 1 + int(spec.num_speculative_tokens) + self._decode_max_rows = min(max_num_seqs * q_per_req, max_batched) + + # Lazily import b12x only on this opt-in path. + from b12x.attention.workspace import B12XAttentionWorkspace + from b12x.integration.mla import ( + sparse_mla_decode_forward, + sparse_mla_extend_forward, + ) + + self._sparse_mla_decode_forward = sparse_mla_decode_forward + self._sparse_mla_extend_forward = sparse_mla_extend_forward + + # Persistent (1,) int32 split-K control pointers, shared by every decode + # call on this layer (filled by workspace.set_split_chunk_config). + self._num_chunks_ptr = torch.empty((1,), dtype=torch.int32, device=self.device) + self._kv_chunk_size_ptr = torch.empty( + (1,), dtype=torch.int32, device=self.device + ) + + def _make_workspace(mode: str, max_total_q: int) -> Any: + # Bare dataclass: __post_init__ only canonicalizes scalars (allocates + # nothing). tmp_output/tmp_lse/output_buffer are assigned per-call + # from the shared workspace manager / a fresh output, so b12x's + # _allocate_split_buffers (called by set_split_chunk_config) is a + # no-op (it only fills None fields). + ws = B12XAttentionWorkspace( + mode=mode, + device=self.device, + dtype=torch.bfloat16, + kv_dtype=torch.uint8, + num_q_heads=self.num_heads, + head_dim=self.q_head_dim, + v_head_dim=self.kv_lora_rank, + topk=self.topk_tokens, + max_total_q=int(max_total_q), + max_batch=max_num_seqs, + page_size=self.block_size, + max_chunks_per_row=self._num_splits_cap, + ) + return ws + + self._decode_workspace = _make_workspace("decode", self._decode_max_rows) + self._decode_workspace.num_chunks_ptr = self._num_chunks_ptr + self._decode_workspace.kv_chunk_size_ptr = self._kv_chunk_size_ptr + self._extend_workspace = _make_workspace("extend", max_batched) + + # Pre-touch the shared decode scratch at the max decode batch so the + # workspace manager grows during warmup, before lock_workspace() runs + # post-cudagraph-capture. Mirrors SparseMLASm120Impl.__init__. + self._borrow_decode_scratch(self._decode_max_rows) + + # Q arrives BF16; the unified kernel quantizes inside. + self.supports_quant_query_input = False + + def _borrow_decode_scratch( + self, num_tokens: int + ) -> tuple[torch.Tensor, torch.Tensor]: + """Borrow split-K mid_out / mid_lse from the shared workspace manager.""" + return tuple( # type: ignore[return-value] + current_workspace_manager().get_simultaneous( + ( + (num_tokens, self.num_heads, self._num_splits_cap, self.kv_lora_rank), + torch.bfloat16, + ), + ((num_tokens, self.num_heads, self._num_splits_cap), torch.float32), + ) + ) + + def forward_mqa( + self, + q: torch.Tensor | tuple[torch.Tensor, torch.Tensor], + kv_c_and_k_pe_cache: torch.Tensor, + attn_metadata: B12xMLASparseMetadata, + layer: AttentionLayer, + ) -> tuple[torch.Tensor, torch.Tensor | None]: + # q arrives as (mqa_ql_nope[T, H, kv_lora_rank], mqa_q_pe[T, H, rope]); + # b12x's GLM_NSA contract wants a single contiguous [T, H, 576] tensor. + if isinstance(q, tuple): + q = torch.cat(q, dim=-1) + q = q.contiguous() + + num_actual_toks = q.shape[0] + + assert self.topk_indices_buffer is not None + topk_indices = self.topk_indices_buffer[:num_actual_toks] + + # Per-request topk indices -> physical cache slot ids. Identical + # conversion to FlashMLASparseImpl / SparseMLASm120Impl. + page_table_1 = cast( + torch.Tensor, + 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, + NUM_TOPK_TOKENS=topk_indices.shape[1], + ), + ) + page_table_1 = page_table_1.to(torch.int32).contiguous() + topk_width = page_table_1.shape[1] + + # nsa_cache_seqlens: per-token count of KV rows to attend = min(causal + # KV length, topk). The indexer -1-pads beyond the valid prefix, so a + # short row's selected indices occupy [0, nsa) and the kernel masks the + # rest (per-token section length + idx<0). + per_token_cache = attn_metadata.cache_seq_lens_per_token[:num_actual_toks] + nsa_cache_seqlens = ( + torch.clamp(per_token_cache, max=topk_width).to(torch.int32).contiguous() + ) + # Per-request KV length (validated but unused by the unified kernels). + cache_seqlens = attn_metadata.seq_lens.to(torch.int32).contiguous() + + # KV cache -> flat (num_slots, 1, nbytes) uint8 (b12x requires rank-3 + # uint8; page_size tells it the per-block stride). page_table_1 are + # physical slot ids consistent with block_size == page_size. + kv_u8 = kv_c_and_k_pe_cache.view(torch.uint8) + kv_cache = kv_u8.reshape(-1, 1, kv_u8.shape[-1]) + if not kv_cache.is_contiguous(): + kv_cache = kv_cache.contiguous() + + output = q.new_empty( + (num_actual_toks, self.num_heads, self.kv_lora_rank), dtype=q.dtype + ) + + is_decode = attn_metadata.max_query_len <= 1 + if is_decode: + mid_out, mid_lse = self._borrow_decode_scratch(num_actual_toks) + ws = self._decode_workspace + ws.tmp_output = mid_out + ws.tmp_lse = mid_lse + ws.output_buffer = output + out = self._sparse_mla_decode_forward( + q_all=q, + kv_cache=kv_cache, + page_table_1=page_table_1, + cache_seqlens_int32=cache_seqlens, + nsa_cache_seqlens_int32=nsa_cache_seqlens, + workspace=ws, + sm_scale=self.scale, + v_head_dim=self.kv_lora_rank, + ) + else: + # Extend / prefill -> single-pass unified prefill (no split-K + # scratch needed; only output_buffer is read). + ws = self._extend_workspace + ws.output_buffer = output + out = self._sparse_mla_extend_forward( + q_all=q, + kv_cache=kv_cache, + selected_token_offsets=page_table_1, + cache_seqlens_int32=cache_seqlens, + nsa_cache_seqlens_int32=nsa_cache_seqlens, + workspace=ws, + sm_scale=self.scale, + v_head_dim=self.kv_lora_rank, + ) + return out, None diff --git a/mods/deepseek-v4-flash-dspark/overlay/vllm/v1/attention/backends/registry.py b/mods/deepseek-v4-flash-dspark/overlay/vllm/v1/attention/backends/registry.py new file mode 100644 index 00000000..e0fac76a --- /dev/null +++ b/mods/deepseek-v4-flash-dspark/overlay/vllm/v1/attention/backends/registry.py @@ -0,0 +1,264 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Attention backend registry""" + +from collections.abc import Callable +from enum import Enum, EnumMeta +from typing import TYPE_CHECKING, cast + +from vllm.logger import init_logger +from vllm.utils.import_utils import resolve_obj_by_qualname + +if TYPE_CHECKING: + from vllm.v1.attention.backend import AttentionBackend + +logger = init_logger(__name__) + + +class _AttentionBackendEnumMeta(EnumMeta): + """Metaclass for AttentionBackendEnum to provide better error messages.""" + + def __getitem__(cls, name: str): + """Get backend by name with helpful error messages.""" + try: + return super().__getitem__(name) + except KeyError: + members = cast("dict[str, Enum]", cls.__members__).keys() + valid_backends = ", ".join(members) + raise ValueError( + f"Unknown attention backend: '{name}'. " + f"Valid options are: {valid_backends}" + ) from None + + +class AttentionBackendEnum(Enum, metaclass=_AttentionBackendEnumMeta): + """Enumeration of all supported attention backends. + + The enum value is the default class path, but this can be overridden + at runtime using register_backend(). + + To get the actual backend class (respecting overrides), use: + backend.get_class() + """ + + FLASH_ATTN = "vllm.v1.attention.backends.flash_attn.FlashAttentionBackend" + FLASH_ATTN_DIFFKV = ( + "vllm.v1.attention.backends.flash_attn_diffkv.FlashAttentionDiffKVBackend" + ) + TRITON_ATTN = "vllm.v1.attention.backends.triton_attn.TritonAttentionBackend" + ROCM_ATTN = "vllm.v1.attention.backends.rocm_attn.RocmAttentionBackend" + ROCM_AITER_MLA = "vllm.v1.attention.backends.mla.rocm_aiter_mla.AiterMLABackend" + ROCM_AITER_TRITON_MLA = ( + "vllm.v1.attention.backends.mla.aiter_triton_mla.AiterTritonMLABackend" + ) + ROCM_AITER_FA = ( + "vllm.v1.attention.backends.rocm_aiter_fa.AiterFlashAttentionBackend" + ) + ROCM_AITER_MLA_SPARSE = ( + "vllm.v1.attention.backends.mla.rocm_aiter_mla_sparse.ROCMAiterMLASparseBackend" + ) + XPU_MLA_SPARSE = "vllm.v1.attention.backends.mla.xpu_mla_sparse.XPUMLASparseBackend" + TORCH_SDPA = "" # this tag is only used for ViT + FLASHINFER = "vllm.v1.attention.backends.flashinfer.FlashInferBackend" + FLASHINFER_MLA = ( + "vllm.v1.attention.backends.mla.flashinfer_mla.FlashInferMLABackend" + ) + TOKENSPEED_MLA = ( + "vllm.v1.attention.backends.mla.tokenspeed_mla.TokenspeedMLABackend" + ) + FLASHINFER_MLA_SPARSE = ( + "vllm.v1.attention.backends.mla.flashinfer_mla_sparse." + "FlashInferMLASparseBackend" + ) + TRITON_MLA = "vllm.v1.attention.backends.mla.triton_mla.TritonMLABackend" + CUTLASS_MLA = "vllm.v1.attention.backends.mla.cutlass_mla.CutlassMLABackend" + FLASHMLA = "vllm.v1.attention.backends.mla.flashmla.FlashMLABackend" + FLASHMLA_SPARSE = ( + "vllm.v1.attention.backends.mla.flashmla_sparse.FlashMLASparseBackend" + ) + SPARSE_MLA_SM120 = ( + "vllm.v1.attention.backends.mla.sparse_mla_sm120.SparseMLASm120Backend" + ) + # Opt-in b12x unified sparse-MLA backend (same SM120 envelope as + # SPARSE_MLA_SM120). Not in the platform auto-selection priority list; select + # it explicitly via VLLM_ATTENTION_BACKEND=B12X_MLA_SPARSE. + B12X_MLA_SPARSE = ( + "vllm.v1.attention.backends.mla.b12x_mla_sparse.B12xMLASparseBackend" + ) + FLASH_ATTN_MLA = "vllm.v1.attention.backends.mla.flashattn_mla.FlashAttnMLABackend" + NO_ATTENTION = "vllm.v1.attention.backends.no_attention.NoAttentionBackend" + FLEX_ATTENTION = "vllm.v1.attention.backends.flex_attention.FlexAttentionBackend" + ROCM_AITER_UNIFIED_ATTN = ( + "vllm.v1.attention.backends.rocm_aiter_unified_attn." + "RocmAiterUnifiedAttentionBackend" + ) + CPU_ATTN = "vllm.v1.attention.backends.cpu_attn.CPUAttentionBackend" + TURBOQUANT = "vllm.v1.attention.backends.turboquant_attn.TurboQuantAttentionBackend" + # 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 + + def get_path(self, include_classname: bool = True) -> str: + """Get the class path for this backend (respects overrides). + + Returns: + The fully qualified class path string + + Raises: + ValueError: If Backend.CUSTOM is used without being registered + """ + path = _ATTN_OVERRIDES.get(self, self.value) + if not path: + raise ValueError( + f"Backend {self.name} must be registered before use. " + f"Use register_backend(Backend.{self.name}, 'your.module.YourClass')" + ) + if not include_classname: + path = path.rsplit(".", 1)[0] + return path + + def get_class(self) -> "type[AttentionBackend]": + """Get the backend class (respects overrides). + + Returns: + The backend class + + Raises: + ImportError: If the backend class cannot be imported + ValueError: If Backend.CUSTOM is used without being registered + """ + return resolve_obj_by_qualname(self.get_path()) + + def is_overridden(self) -> bool: + """Check if this backend has been overridden. + + Returns: + True if the backend has a registered override + """ + return self in _ATTN_OVERRIDES + + def clear_override(self) -> None: + """Clear any override for this backend, reverting to the default.""" + _ATTN_OVERRIDES.pop(self, None) + + +class MambaAttentionBackendEnum(Enum, metaclass=_AttentionBackendEnumMeta): + """Enumeration of all supported mamba attention backends. + + The enum value is the default class path, but this can be overridden + at runtime using register_backend(). + + To get the actual backend class (respecting overrides), use: + backend.get_class() + """ + + 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" + LINEAR = "vllm.v1.attention.backends.linear_attn.LinearAttentionBackend" + GDN_ATTN = "vllm.v1.attention.backends.gdn_attn.GDNAttentionBackend" + # 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 + + def get_path(self, include_classname: bool = True) -> str: + """Get the class path for this backend (respects overrides). + + Returns: + The fully qualified class path string + + Raises: + ValueError: If Backend.CUSTOM is used without being registered + """ + path = _MAMBA_ATTN_OVERRIDES.get(self, self.value) + if not path: + raise ValueError( + f"Backend {self.name} must be registered before use. " + f"Use register_backend(Backend.{self.name}, 'your.module.YourClass')" + ) + if not include_classname: + path = path.rsplit(".", 1)[0] + return path + + def get_class(self) -> "type[AttentionBackend]": + """Get the backend class (respects overrides). + + Returns: + The backend class + + Raises: + ImportError: If the backend class cannot be imported + ValueError: If Backend.CUSTOM is used without being registered + """ + return resolve_obj_by_qualname(self.get_path()) + + def is_overridden(self) -> bool: + """Check if this backend has been overridden. + + Returns: + True if the backend has a registered override + """ + return self in _MAMBA_ATTN_OVERRIDES + + def clear_override(self) -> None: + """Clear any override for this backend, reverting to the default.""" + _MAMBA_ATTN_OVERRIDES.pop(self, None) + + +_ATTN_OVERRIDES: dict[AttentionBackendEnum, str] = {} +_MAMBA_ATTN_OVERRIDES: dict[MambaAttentionBackendEnum, str] = {} + + +def register_backend( + backend: AttentionBackendEnum | MambaAttentionBackendEnum, + 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 + class_path: Optional class path. If not provided and used as + decorator, will be auto-generated from the class. + + Returns: + Decorator function if class_path is None, otherwise a no-op + + Examples: + # Override an existing attention backend + @register_backend(AttentionBackendEnum.FLASH_ATTN) + class MyCustomFlashAttn: + ... + + # Override an existing mamba attention backend + @register_backend(MambaAttentionBackendEnum.LINEAR, is_mamba=True) + class MyCustomMambaAttn: + ... + + # Register a custom third-party attention backend + @register_backend(AttentionBackendEnum.CUSTOM) + class MyCustomBackend: + ... + + # Direct registration + register_backend( + AttentionBackendEnum.CUSTOM, + "my.module.MyCustomBackend" + ) + """ + + def decorator(cls: type) -> type: + if is_mamba: + _MAMBA_ATTN_OVERRIDES[backend] = f"{cls.__module__}.{cls.__qualname__}" # type: ignore[index] + else: + _ATTN_OVERRIDES[backend] = f"{cls.__module__}.{cls.__qualname__}" # type: ignore[index] + return cls + + if class_path is not None: + if is_mamba: + _MAMBA_ATTN_OVERRIDES[backend] = class_path # type: ignore[index] + else: + _ATTN_OVERRIDES[backend] = class_path # type: ignore[index] + return lambda x: x + + return decorator diff --git a/mods/deepseek-v4-flash-dspark/overlay/vllm/v1/core/sched/scheduler.py b/mods/deepseek-v4-flash-dspark/overlay/vllm/v1/core/sched/scheduler.py new file mode 100644 index 00000000..eab30271 --- /dev/null +++ b/mods/deepseek-v4-flash-dspark/overlay/vllm/v1/core/sched/scheduler.py @@ -0,0 +1,2369 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import itertools +import time +from collections import defaultdict, deque +from collections.abc import Iterable +from dataclasses import replace +from typing import Any + +from vllm.compilation.cuda_graph import CUDAGraphStat +from vllm.config import VllmConfig +from vllm.distributed.ec_transfer.ec_connector.base import ( + ECConnectorMetadata, + ECConnectorRole, +) +from vllm.distributed.ec_transfer.ec_connector.factory import ECConnectorFactory +from vllm.distributed.kv_events import EventPublisherFactory, KVEventBatch +from vllm.distributed.kv_transfer.kv_connector.factory import KVConnectorFactory +from vllm.distributed.kv_transfer.kv_connector.v1 import ( + KVConnectorBase_V1, + KVConnectorRole, + SupportsHMA, +) +from vllm.distributed.kv_transfer.kv_connector.v1.base import KVConnectorMetadata +from vllm.distributed.kv_transfer.kv_connector.v1.metrics import KVConnectorStats +from vllm.logger import init_logger +from vllm.model_executor.layers.fused_moe.routed_experts_capturer import ( + RoutedExpertsManager, +) +from vllm.multimodal import MULTIMODAL_REGISTRY, MultiModalRegistry +from vllm.multimodal.encoder_budget import MultiModalBudget +from vllm.v1.core.encoder_cache_manager import ( + EncoderCacheManager, + EncoderDecoderCacheManager, +) +from vllm.v1.core.kv_cache_manager import KVCacheBlocks, KVCacheManager +from vllm.v1.core.kv_cache_metrics import KVCacheMetricsCollector +from vllm.v1.core.sched.interface import PauseState, SchedulerInterface +from vllm.v1.core.sched.output import ( + CachedRequestData, + GrammarOutput, + NewRequestData, + SchedulerOutput, +) +from vllm.v1.core.sched.request_queue import ( + RequestQueue, + SchedulingPolicy, + create_request_queue, +) +from vllm.v1.core.sched.utils import check_stop, remove_all +from vllm.v1.engine import EngineCoreEventType, EngineCoreOutput, EngineCoreOutputs +from vllm.v1.kv_cache_interface import KVCacheConfig +from vllm.v1.metrics.perf import ModelMetrics, PerfStats +from vllm.v1.metrics.stats import PrefixCacheStats, SchedulerStats +from vllm.v1.outputs import DraftTokenIds, KVConnectorOutput, ModelRunnerOutput +from vllm.v1.request import Request, RequestStatus, StreamingUpdate +from vllm.v1.spec_decode.metrics import SpecDecodingStats +from vllm.v1.structured_output import StructuredOutputManager +from vllm.v1.utils import record_function_or_nullcontext + +logger = init_logger(__name__) + + +class Scheduler(SchedulerInterface): + def __init__( + self, + vllm_config: VllmConfig, + kv_cache_config: KVCacheConfig, + structured_output_manager: StructuredOutputManager, + block_size: int, + hash_block_size: int | None = None, + mm_registry: MultiModalRegistry = MULTIMODAL_REGISTRY, + include_finished_set: bool = False, + log_stats: bool = False, + ) -> None: + self.vllm_config = vllm_config + self.scheduler_config = vllm_config.scheduler_config + self.cache_config = vllm_config.cache_config + self.lora_config = vllm_config.lora_config + self.kv_cache_config = kv_cache_config + self.kv_events_config = vllm_config.kv_events_config + self.parallel_config = vllm_config.parallel_config + self.log_stats = log_stats + self.observability_config = vllm_config.observability_config + self.kv_metrics_collector: KVCacheMetricsCollector | None = None + if self.observability_config.kv_cache_metrics: + self.kv_metrics_collector = KVCacheMetricsCollector( + self.observability_config.kv_cache_metrics_sample, + ) + self.structured_output_manager = structured_output_manager + self.is_encoder_decoder = vllm_config.model_config.is_encoder_decoder + + # include_finished_set controls whether a separate set of finished + # request ids should be included in the EngineCoreOutputs returned + # by update_from_outputs(). This is currently used in the multi-engine + # case to track request lifetimes efficiently. + self.finished_req_ids_dict: dict[int, set[str]] | None = ( + defaultdict(set) if include_finished_set else None + ) + self.prev_step_scheduled_req_ids: set[str] = set() + + # Scheduling constraints. + self.max_num_running_reqs = self.scheduler_config.max_num_seqs + self.max_num_scheduled_tokens = ( + self.scheduler_config.max_num_scheduled_tokens + if self.scheduler_config.max_num_scheduled_tokens + else self.scheduler_config.max_num_batched_tokens + ) + self.max_model_len = vllm_config.model_config.max_model_len + self.enable_kv_cache_events = ( + self.kv_events_config is not None + and self.kv_events_config.enable_kv_cache_events + ) + + # Create KVConnector for the Scheduler. Note that each Worker + # will have a corresponding KVConnector with Role=WORKER. + # KV Connector pushes/pull of remote KVs for P/D and offloading. + self.connector = None + self.connector_prefix_cache_stats: PrefixCacheStats | None = None + self.recompute_kv_load_failures = True + if self.vllm_config.kv_transfer_config is not None: + assert not self.is_encoder_decoder, ( + "Encoder-decoder models are not currently supported with KV connectors" + ) + self.connector = KVConnectorFactory.create_connector( + config=self.vllm_config, + role=KVConnectorRole.SCHEDULER, + kv_cache_config=self.kv_cache_config, + ) + if self.log_stats: + self.connector_prefix_cache_stats = PrefixCacheStats() + kv_load_failure_policy = ( + self.vllm_config.kv_transfer_config.kv_load_failure_policy + ) + self.recompute_kv_load_failures = kv_load_failure_policy == "recompute" + + self.kv_event_publisher = EventPublisherFactory.create( + self.kv_events_config, + self.parallel_config.data_parallel_index, + ) + self.ec_connector = None + if self.vllm_config.ec_transfer_config is not None: + self.ec_connector = ECConnectorFactory.create_connector( + config=self.vllm_config, role=ECConnectorRole.SCHEDULER + ) + + num_gpu_blocks = self.cache_config.num_gpu_blocks + assert num_gpu_blocks is not None and num_gpu_blocks > 0 + + self.block_size = block_size + self.dcp_world_size = vllm_config.parallel_config.decode_context_parallel_size + self.pcp_world_size = vllm_config.parallel_config.prefill_context_parallel_size + + # req_id -> Request + self.requests: dict[str, Request] = {} + # Scheduling policy + try: + self.policy = SchedulingPolicy(self.scheduler_config.policy) + except ValueError as e: + raise ValueError( + f"Unknown scheduling policy: {self.scheduler_config.policy}" + ) from e + # Priority queues for requests. + self.waiting = create_request_queue(self.policy) + # requests skipped in waiting flow due async deps or constraints. + self.skipped_waiting = create_request_queue(self.policy) + self.running: list[Request] = [] + + # The request IDs that are finished in between the previous and the + # current steps. This is used to notify the workers about the finished + # requests so that they can free the cached states for those requests. + # This is flushed at the end of each scheduling step. + self.finished_req_ids: set[str] = set() + + # Counter for requests waiting for streaming input. Used to calculate + # number of unfinished requests + self.num_waiting_for_streaming_input: int = 0 + + # KV Connector: requests in process of async KV loading or recving + self.finished_recving_kv_req_ids: set[str] = set() + self.failed_recving_kv_req_ids: set[str] = set() + + # Encoder-related. + # Calculate encoder cache size if applicable + supports_mm_inputs = mm_registry.supports_multimodal_inputs( + vllm_config.model_config + ) + mm_budget = ( + MultiModalBudget(vllm_config, mm_registry) if supports_mm_inputs else None + ) + + # NOTE: Text-only encoder-decoder models are implemented as + # multi-modal models for convenience + # Example: https://github.com/vllm-project/bart-plugin + if self.is_encoder_decoder: + assert mm_budget and len(mm_budget.mm_max_toks_per_item) <= 1, ( + "Encoder-decoder models are expected to implement the " + "multimodal interface with at most one modality." + ) + + self.max_num_encoder_input_tokens = ( + mm_budget.encoder_compute_budget if mm_budget else 0 + ) + encoder_cache_size = mm_budget.encoder_cache_size if mm_budget else 0 + self.encoder_cache_manager = ( + EncoderDecoderCacheManager(cache_size=encoder_cache_size) + if self.is_encoder_decoder + else EncoderCacheManager(cache_size=encoder_cache_size) + ) + + speculative_config = vllm_config.speculative_config + self.use_eagle = False + self.num_spec_tokens = self.num_lookahead_tokens = 0 + if speculative_config: + self.num_spec_tokens = speculative_config.num_speculative_tokens + if speculative_config.use_eagle(): + self.use_eagle = True + self.num_lookahead_tokens = self.num_spec_tokens + if speculative_config.uses_draft_model(): + self.num_lookahead_tokens = self.num_spec_tokens + + # Create the KV cache manager. + if hash_block_size is None: + hash_block_size = block_size + self.kv_cache_manager = KVCacheManager( + kv_cache_config=kv_cache_config, + max_model_len=self.max_model_len, + max_num_batched_tokens=self.scheduler_config.max_num_batched_tokens, + enable_caching=self.cache_config.enable_prefix_caching, + use_eagle=self.use_eagle, + log_stats=self.log_stats, + enable_kv_cache_events=self.enable_kv_cache_events, + dcp_world_size=self.dcp_world_size, + pcp_world_size=self.pcp_world_size, + hash_block_size=hash_block_size, + metrics_collector=self.kv_metrics_collector, + ) + # Bind GPU block pool to the KV connector. This must happen after + # kv_cache_manager is constructed so block_pool is available. + if self.connector is not None: + self.connector.bind_gpu_block_pool(self.kv_cache_manager.block_pool) + + self.use_pp = self.parallel_config.pipeline_parallel_size > 1 + self.use_v2_model_runner = vllm_config.use_v2_model_runner + self.scheduler_reserve_full_isl = ( + self.scheduler_config.scheduler_reserve_full_isl + ) + + self.has_mamba_layers = kv_cache_config.has_mamba_layers + self.needs_kv_cache_zeroing = kv_cache_config.needs_kv_cache_zeroing + self.need_mamba_block_aligned_split = ( + self.has_mamba_layers and self.cache_config.mamba_cache_mode == "align" + ) + self.perf_metrics: ModelMetrics | None = None + if self.log_stats and vllm_config.observability_config.enable_mfu_metrics: + self.perf_metrics = ModelMetrics(vllm_config) + + self.enable_return_routed_experts = ( + vllm_config.model_config.enable_return_routed_experts + ) + + if self.enable_return_routed_experts: + assert self.dcp_world_size == 1 and self.pcp_world_size == 1, ( + "enable_return_routed_experts does not support context parallelism " + "(dcp_world_size > 1 or pcp_world_size > 1)" + ) + + self.routed_experts_mgr = RoutedExpertsManager( + vllm_config=vllm_config, + kv_cache_config=kv_cache_config, + ) + # Block-ID snapshot taken at schedule time (before forward), + # so update_from_output can read slot data even if a later + # schedule() frees the blocks (async scheduling race). + self._re_block_ids: dict[str, list[int]] = {} + + self._pause_state: PauseState = PauseState.UNPAUSED + + def _mamba_block_aligned_split( + self, + request: Request, + num_new_tokens: int, + num_new_local_computed_tokens: int = 0, + num_external_computed_tokens: int = 0, + ) -> int: + assert num_external_computed_tokens == 0, ( + "External KV connector is not verified yet" + ) + num_computed_tokens = ( + request.num_computed_tokens + + num_new_local_computed_tokens + + num_external_computed_tokens + ) + # Perform block-aligned splitting at prefill phase, including: + # * non-resumed requests: num_computed_tokens < num_prompt_tokens + 0 + # * resumed requests: num_computed_tokens < ( + # num_prompt_tokens + num_output_tokens + # ) + # NOTE: Use `request.num_tokens - 1` to bypass normal decoding. + if num_computed_tokens < max(request.num_prompt_tokens, request.num_tokens - 1): + # To enable block-aligned caching of the Mamba state, `num_new_tokens` + # must be a multiple of `block_size`. + # As an exception, if `num_new_tokens` is less than `block_size`, the + # state is simply not cached, requiring no special handling. + # Additionally, when Eagle mode is enabled, FullAttn prunes the last + # matching block. To prevent this from causing a Mamba cache miss, the + # last chunk must be not smaller than `block_size`. + block_size = self.cache_config.block_size + last_cache_position = request.num_tokens - request.num_tokens % block_size + # eagle prune + if self.use_eagle: + last_cache_position = max(last_cache_position - block_size, 0) + num_computed_tokens_after_sched = num_computed_tokens + num_new_tokens + if num_computed_tokens_after_sched < last_cache_position: + # align to block_size + num_new_tokens = num_new_tokens // block_size * block_size + elif ( + num_computed_tokens + < last_cache_position + < num_computed_tokens_after_sched + ): + # force to cache the last chunk + num_new_tokens = last_cache_position - num_computed_tokens + else: + # prefill the last few tokens + pass + return num_new_tokens + + def schedule(self) -> SchedulerOutput: + # NOTE(woosuk) on the scheduling algorithm: + # There's no "decoding phase" nor "prefill phase" in the scheduler. + # Each request just has the num_computed_tokens and + # num_tokens_with_spec. num_tokens_with_spec = + # len(prompt_token_ids) + len(output_token_ids) + len(spec_token_ids). + # At each step, the scheduler tries to assign tokens to the requests + # so that each request's num_computed_tokens can catch up its + # num_tokens_with_spec. This is general enough to cover + # chunked prefills, prefix caching, speculative decoding, + # and the "jump decoding" optimization in the future. + + scheduled_new_reqs: list[Request] = [] + scheduled_resumed_reqs: list[Request] = [] + scheduled_running_reqs: list[Request] = [] + preempted_reqs: list[Request] = [] + + req_to_new_blocks: dict[str, KVCacheBlocks] = {} + num_scheduled_tokens: dict[str, int] = {} + token_budget = self.max_num_scheduled_tokens + if self._pause_state == PauseState.PAUSED_ALL: + # Do not schedule any requests when paused. + token_budget = 0 + + # Encoder-related. + scheduled_encoder_inputs: dict[str, list[int]] = {} + encoder_compute_budget = self.max_num_encoder_input_tokens + # Spec decode-related. + scheduled_spec_decode_tokens: dict[str, list[int]] = {} + + # For logging. + scheduled_timestamp = time.monotonic() + + self.kv_cache_manager.new_step_starts() + + # First, schedule the RUNNING requests. + req_index = 0 + while req_index < len(self.running) and token_budget > 0: + request = self.running[req_index] + + if ( + request.num_output_placeholders > 0 + # This is (num_computed_tokens + 1) - (num_output_placeholders - 1). + # Since output placeholders are also included in the computed tokens + # count, we subtract (num_output_placeholders - 1) to remove any draft + # tokens, so that we can be sure no further steps are needed even if + # they are all rejected. + and request.num_computed_tokens + 2 - request.num_output_placeholders + >= request.num_prompt_tokens + request.max_tokens + ): + # Async scheduling: Avoid scheduling an extra step when we are sure that + # the previous step has reached request.max_tokens. We don't schedule + # partial draft tokens since this prevents uniform decode optimizations. + req_index += 1 + continue + + num_new_tokens = ( + request.num_tokens_with_spec + + request.num_output_placeholders + - request.num_computed_tokens + ) + if 0 < self.scheduler_config.long_prefill_token_threshold < num_new_tokens: + num_new_tokens = self.scheduler_config.long_prefill_token_threshold + num_new_tokens = min(num_new_tokens, token_budget) + + # Make sure the input position does not exceed the max model len. + # This is necessary when using spec decoding. + num_new_tokens = min( + num_new_tokens, self.max_model_len - 1 - request.num_computed_tokens + ) + + # Schedule encoder inputs. + encoder_inputs_to_schedule = None + external_load_encoder_input: list[int] = [] + new_encoder_compute_budget = encoder_compute_budget + if request.has_encoder_inputs: + ( + encoder_inputs_to_schedule, + num_new_tokens, + new_encoder_compute_budget, + external_load_encoder_input, + ) = self._try_schedule_encoder_inputs( + request, + request.num_computed_tokens, + num_new_tokens, + encoder_compute_budget, + shift_computed_tokens=1 if self.use_eagle else 0, + ) + + if self.need_mamba_block_aligned_split: + num_new_tokens = self._mamba_block_aligned_split( + request, num_new_tokens + ) + + if num_new_tokens == 0: + # The request cannot be scheduled because one of the following + # reasons: + # 1. No new tokens to schedule. This may happen when + # (1) PP>1 and we have already scheduled all prompt tokens + # but they are not finished yet. + # (2) Async scheduling and the request has reached to either + # its max_total_tokens or max_model_len. + # 2. The encoder budget is exhausted. + # 3. The encoder cache is exhausted. + # 4. Insufficient budget for a block-aligned chunk in hybrid + # models with mamba cache mode \"align\". + # NOTE(woosuk): Here, by doing `continue` instead of `break`, + # we do not strictly follow the FCFS scheduling policy and + # allow the lower-priority requests to be scheduled. + req_index += 1 + continue + + # Schedule newly needed KV blocks for the request. + with record_function_or_nullcontext("schedule: allocate_slots"): + while True: + new_blocks = self.kv_cache_manager.allocate_slots( + request, + num_new_tokens, + num_lookahead_tokens=self.num_lookahead_tokens, + ) + + if new_blocks is not None: + # The request can be scheduled. + break + + # The request cannot be scheduled. + # Preempt the lowest-priority request. + if self.policy == SchedulingPolicy.PRIORITY: + preempted_req = max( + self.running, + key=lambda r: (r.priority, r.arrival_time), + ) + self.running.remove(preempted_req) + if preempted_req in scheduled_running_reqs: + preempted_req_id = preempted_req.request_id + scheduled_running_reqs.remove(preempted_req) + token_budget += num_scheduled_tokens.pop(preempted_req_id) + req_to_new_blocks.pop(preempted_req_id) + scheduled_spec_decode_tokens.pop(preempted_req_id, None) + preempted_encoder_inputs = scheduled_encoder_inputs.pop( + preempted_req_id, None + ) + if preempted_encoder_inputs: + # Restore encoder compute budget if the preempted + # request had encoder inputs scheduled in this step. + num_embeds_to_restore = sum( + preempted_req.get_num_encoder_embeds(i) + for i in preempted_encoder_inputs + ) + encoder_compute_budget += num_embeds_to_restore + req_index -= 1 + else: + preempted_req = self.running.pop() + + self._preempt_request(preempted_req, scheduled_timestamp) + preempted_reqs.append(preempted_req) + if preempted_req == request: + # No more request to preempt. Cannot schedule this request. + break + + if new_blocks is None: + # Cannot schedule this request. + break + + # Schedule the request. + scheduled_running_reqs.append(request) + request_id = request.request_id + req_to_new_blocks[request_id] = new_blocks + num_scheduled_tokens[request_id] = num_new_tokens + token_budget -= num_new_tokens + req_index += 1 + + # Speculative decode related. + if request.spec_token_ids: + num_scheduled_spec_tokens = ( + num_new_tokens + + request.num_computed_tokens + - request.num_tokens + - request.num_output_placeholders + ) + if num_scheduled_spec_tokens > 0: + spec_token_ids = request.spec_token_ids + if len(spec_token_ids) > num_scheduled_spec_tokens: + spec_token_ids = spec_token_ids[:num_scheduled_spec_tokens] + scheduled_spec_decode_tokens[request.request_id] = spec_token_ids + + # New spec tokens will be set in `update_draft_token_ids` before the + # next step when applicable. + request.spec_token_ids = [] + + # 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) + + # Record the LoRAs in scheduled_running_reqs + scheduled_loras: set[int] = set() + if self.lora_config: + scheduled_loras = set( + req.lora_request.lora_int_id + for req in scheduled_running_reqs + if req.lora_request and req.lora_request.lora_int_id > 0 + ) + assert len(scheduled_loras) <= self.lora_config.max_loras + + # Next, schedule the WAITING requests. + if not preempted_reqs and self._pause_state == PauseState.UNPAUSED: + step_skipped_waiting = create_request_queue(self.policy) + + while (self.waiting or self.skipped_waiting) and token_budget > 0: + if len(self.running) == self.max_num_running_reqs: + break + + request_queue = self._select_waiting_queue_for_scheduling() + assert request_queue is not None + + request = request_queue.peek_request() + request_id = request.request_id + + # try to promote blocked statuses while traversing skipped queue. + if self._is_blocked_waiting_status( + request.status + ) and not self._try_promote_blocked_waiting_request(request): + if request.status == RequestStatus.WAITING_FOR_REMOTE_KVS: + logger.debug( + "%s is still in WAITING_FOR_REMOTE_KVS state.", + request_id, + ) + request_queue.pop_request() + step_skipped_waiting.prepend_request(request) + continue + + # Check that adding the request still respects the max_loras + # constraint. + if ( + self.lora_config + and request.lora_request + and ( + len(scheduled_loras) == self.lora_config.max_loras + and request.lora_request.lora_int_id not in scheduled_loras + ) + ): + # Scheduling would exceed max_loras, skip. + request_queue.pop_request() + step_skipped_waiting.prepend_request(request) + continue + + num_external_computed_tokens = 0 + load_kv_async = False + connector_prefix_cache_queries, connector_prefix_cache_hits = 0, 0 + + # Get already-cached tokens. + if request.num_computed_tokens == 0: + # Get locally-cached tokens. + new_computed_blocks, num_new_local_computed_tokens = ( + self.kv_cache_manager.get_computed_blocks(request) + ) + + # Get externally-cached tokens if using a KVConnector. + if self.connector is not None: + ext_tokens, load_kv_async = ( + self.connector.get_num_new_matched_tokens( + request, num_new_local_computed_tokens + ) + ) + + if ext_tokens is None: + # The request cannot be scheduled because + # the KVConnector couldn't determine + # the number of matched tokens. + request_queue.pop_request() + step_skipped_waiting.prepend_request(request) + continue + + num_external_computed_tokens = ext_tokens + + connector_prefix_cache_queries = ( + request.num_tokens - num_new_local_computed_tokens + ) + connector_prefix_cache_hits = num_external_computed_tokens + + # Total computed tokens (local + external). + num_computed_tokens = ( + num_new_local_computed_tokens + num_external_computed_tokens + ) + assert num_computed_tokens <= request.num_tokens + + # Track first scheduled prefill, not post-preemption repeat prefills + if request.prefill_stats is not None: + assert num_computed_tokens <= request.num_prompt_tokens + request.prefill_stats.set( + num_prompt_tokens=request.num_prompt_tokens, + num_local_cached_tokens=num_new_local_computed_tokens, + num_external_cached_tokens=num_external_computed_tokens, + ) + else: + # KVTransfer: WAITING reqs have num_computed_tokens > 0 + # after async KV recvs are completed. + new_computed_blocks = self.kv_cache_manager.empty_kv_cache_blocks + num_new_local_computed_tokens = 0 + num_computed_tokens = request.num_computed_tokens + + encoder_inputs_to_schedule = None + external_load_encoder_input = [] + new_encoder_compute_budget = encoder_compute_budget + + if load_kv_async: + # KVTransfer: loading remote KV, do not allocate for new work. + assert num_external_computed_tokens > 0 + num_new_tokens = 0 + else: + # Number of tokens to be scheduled. + # We use `request.num_tokens` instead of + # `request.num_prompt_tokens` to consider the resumed + # requests, which have output tokens. + num_new_tokens = request.num_tokens - num_computed_tokens + threshold = self.scheduler_config.long_prefill_token_threshold + if 0 < threshold < num_new_tokens: + num_new_tokens = threshold + + # chunked prefill has to be enabled explicitly to allow + # pooling requests to be chunked + if ( + not self.scheduler_config.enable_chunked_prefill + and num_new_tokens > token_budget + ): + # If chunked_prefill is disabled, + # we can stop the scheduling here. + break + + num_new_tokens = min(num_new_tokens, token_budget) + assert num_new_tokens > 0 + + # Schedule encoder inputs. + if request.has_encoder_inputs: + ( + encoder_inputs_to_schedule, + num_new_tokens, + new_encoder_compute_budget, + external_load_encoder_input, + ) = self._try_schedule_encoder_inputs( + request, + num_computed_tokens, + num_new_tokens, + encoder_compute_budget, + shift_computed_tokens=1 if self.use_eagle else 0, + ) + if num_new_tokens == 0: + # The request cannot be scheduled. + break + + if self.need_mamba_block_aligned_split: + num_new_tokens = self._mamba_block_aligned_split( + request, + num_new_tokens, + num_new_local_computed_tokens, + num_external_computed_tokens, + ) + if num_new_tokens == 0: + break + + # Handles an edge case when P/D Disaggregation + # is used with Spec Decoding where an + # extra block gets allocated which + # creates a mismatch between the number + # of local and remote blocks. + effective_lookahead_tokens = ( + 0 if request.num_computed_tokens == 0 else self.num_lookahead_tokens + ) + + # Determine if we need to allocate cross-attention blocks. + num_encoder_tokens = 0 + if ( + self.is_encoder_decoder + and request.has_encoder_inputs + and encoder_inputs_to_schedule + ): + num_encoder_tokens = sum( + request.get_num_encoder_embeds(i) + for i in encoder_inputs_to_schedule + ) + + new_blocks = self.kv_cache_manager.allocate_slots( + request, + num_new_tokens, + num_new_computed_tokens=num_new_local_computed_tokens, + new_computed_blocks=new_computed_blocks, + num_lookahead_tokens=effective_lookahead_tokens, + num_external_computed_tokens=num_external_computed_tokens, + delay_cache_blocks=load_kv_async, + num_encoder_tokens=num_encoder_tokens, + full_sequence_must_fit=self.scheduler_reserve_full_isl, + ) + + if new_blocks is None: + # The request cannot be scheduled. + + # NOTE: we need to untouch the request from the encode cache + # manager + if request.has_encoder_inputs: + self.encoder_cache_manager.free(request) + break + + # KVTransfer: the connector uses this info to determine + # if a load is needed. Note that + # This information is used to determine if a load is + # needed for this request. + if self.connector is not None: + self.connector.update_state_after_alloc( + request, + self.kv_cache_manager.get_blocks(request_id), + num_external_computed_tokens, + ) + if ( + self.connector_prefix_cache_stats is not None + and connector_prefix_cache_queries != 0 + ): + self.connector_prefix_cache_stats.record( + num_tokens=connector_prefix_cache_queries, + num_hits=connector_prefix_cache_hits, + preempted=request.num_preemptions > 0, + ) + + request = request_queue.pop_request() + if load_kv_async: + # If loading async, allocate memory and put request + # into the WAITING_FOR_REMOTE_KV state. + request.status = RequestStatus.WAITING_FOR_REMOTE_KVS + step_skipped_waiting.prepend_request(request) + # Set num_computed_tokens even though KVs are not yet loaded. + # request.num_computed_tokens will not be used anywhere until + # the request finished the KV transfer. + # + # If a transfer error is reported by the connector, + # request.num_computed_tokens will be re-set accordingly in + # _update_requests_with_invalid_blocks. + # + # When the transfer is finished, either successfully or not, + # request.num_computed_tokens will correctly reflect the number + # of computed tokens. + # _update_waiting_for_remote_kv will then cache + # only the successfully loaded tokens. + request.num_computed_tokens = num_computed_tokens + continue + + self.running.append(request) + if self.log_stats: + request.record_event( + EngineCoreEventType.SCHEDULED, scheduled_timestamp + ) + if request.status == RequestStatus.WAITING: + scheduled_new_reqs.append(request) + elif request.status == RequestStatus.PREEMPTED: + scheduled_resumed_reqs.append(request) + else: + raise RuntimeError(f"Invalid request status: {request.status}") + + if self.lora_config and request.lora_request: + scheduled_loras.add(request.lora_request.lora_int_id) + req_to_new_blocks[request_id] = self.kv_cache_manager.get_blocks( + request_id + ) + num_scheduled_tokens[request_id] = num_new_tokens + token_budget -= num_new_tokens + request.status = RequestStatus.RUNNING + request.num_computed_tokens = num_computed_tokens + # 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) + + # re-queue requests skipped in this pass ahead of older skipped items. + if step_skipped_waiting: + self.skipped_waiting.prepend_requests(step_skipped_waiting) + + # Check if the scheduling constraints are satisfied. + total_num_scheduled_tokens = sum(num_scheduled_tokens.values()) + assert total_num_scheduled_tokens <= self.max_num_scheduled_tokens + + assert token_budget >= 0 + assert len(self.running) <= self.max_num_running_reqs + # Since some requests in the RUNNING queue may not be scheduled in + # this step, the total number of scheduled requests can be smaller than + # len(self.running). + assert len(scheduled_new_reqs) + len(scheduled_resumed_reqs) + len( + 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 = scheduled_new_reqs + scheduled_resumed_reqs + scheduled_resumed_reqs = [] + new_reqs_data = [ + NewRequestData.from_request( + req, + req_to_new_blocks[req.request_id].get_block_ids(), + req._all_token_ids, + ) + for req in scheduled_new_reqs + ] + else: + new_reqs_data = [ + NewRequestData.from_request( + req, req_to_new_blocks[req.request_id].get_block_ids() + ) + for req in scheduled_new_reqs + ] + + with record_function_or_nullcontext("schedule: make_cached_request_data"): + cached_reqs_data = self._make_cached_request_data( + scheduled_running_reqs, + scheduled_resumed_reqs, + num_scheduled_tokens, + scheduled_spec_decode_tokens, + req_to_new_blocks, + ) + + # Record the request ids that were scheduled in this step. + self.prev_step_scheduled_req_ids.clear() + self.prev_step_scheduled_req_ids.update(num_scheduled_tokens.keys()) + + new_block_ids_to_zero = ( + (self.kv_cache_manager.take_new_block_ids() or None) + if self.needs_kv_cache_zeroing + else None + ) + + scheduler_output = SchedulerOutput( + scheduled_new_reqs=new_reqs_data, + scheduled_cached_reqs=cached_reqs_data, + num_scheduled_tokens=num_scheduled_tokens, + total_num_scheduled_tokens=total_num_scheduled_tokens, + scheduled_spec_decode_tokens=scheduled_spec_decode_tokens, + scheduled_encoder_inputs=scheduled_encoder_inputs, + num_common_prefix_blocks=num_common_prefix_blocks, + preempted_req_ids={req.request_id for req in preempted_reqs}, + # finished_req_ids is an existing state in the scheduler, + # instead of being newly scheduled in this step. + # It contains the request IDs that are finished in between + # the previous and the current steps. + finished_req_ids=self.finished_req_ids, + free_encoder_mm_hashes=self.encoder_cache_manager.get_freed_mm_hashes(), + new_block_ids_to_zero=new_block_ids_to_zero, + ) + + # NOTE(Kuntai): this function is designed for multiple purposes: + # 1. Plan the KV cache store + # 2. Wrap up all the KV cache load / save ops into an opaque object + # 3. Clear the internal states of the connector + if self.connector is not None: + meta = self._build_kv_connector_meta(self.connector, scheduler_output) + scheduler_output.kv_connector_metadata = meta + + # Build the connector meta for ECConnector + if self.ec_connector is not None: + ec_meta: ECConnectorMetadata = self.ec_connector.build_connector_meta( + scheduler_output + ) + scheduler_output.ec_connector_metadata = ec_meta + + with record_function_or_nullcontext("schedule: update_after_schedule"): + self._update_after_schedule(scheduler_output) + return scheduler_output + + def _build_kv_connector_meta( + self, connector: KVConnectorBase_V1, scheduler_output: SchedulerOutput + ) -> KVConnectorMetadata: + return connector.build_connector_meta(scheduler_output) + + def _preempt_request(self, request: Request, timestamp: float) -> None: + """Preempt a request and put it back to the waiting queue. + + NOTE: The request should be popped from the running queue outside of this + method. + """ + assert request.status == RequestStatus.RUNNING, ( + "Only running requests can be preempted" + ) + self.kv_cache_manager.free(request) + self.encoder_cache_manager.free(request) + request.status = RequestStatus.PREEMPTED + request.num_computed_tokens = 0 + if request.spec_token_ids: + request.spec_token_ids = [] + request.num_preemptions += 1 + if self.log_stats: + request.record_event(EngineCoreEventType.PREEMPTED, timestamp) + + # Put the request back to the waiting queue. + self.waiting.prepend_request(request) + + def _update_after_schedule(self, scheduler_output: SchedulerOutput) -> None: + # Advance the number of computed tokens for the request AFTER + # the request is scheduled. + # 1. The scheduler_output of the current step has to include the + # original number of scheduled tokens to determine input IDs. + # 2. Advance the number of computed tokens here allowing us to + # schedule the prefill request again immediately in the next + # scheduling step. + # 3. If some tokens (e.g. spec tokens) are rejected later, the number of + # computed tokens will be adjusted in update_from_output. + num_scheduled_tokens = scheduler_output.num_scheduled_tokens + for req_id, num_scheduled_token in num_scheduled_tokens.items(): + request = self.requests[req_id] + request.num_computed_tokens += num_scheduled_token + request.is_prefill_chunk = request.num_computed_tokens < ( + request.num_tokens + request.num_output_placeholders + ) + scheduler_output.has_structured_output_requests |= ( + request.use_structured_output and not request.is_prefill_chunk + ) + + # Snapshot block IDs for routed experts before forward starts. + # A concurrent schedule() may preempt requests and free blocks + # before update_from_output runs; the snapshot survives that. + # Use update() to preserve entries from the previous step that + # have not yet been consumed by update_from_output (async + # scheduling may call _update_after_schedule again before the + # prior update_from_output runs). + if self.enable_return_routed_experts: + gid = self.routed_experts_mgr.attn_gid + self._re_block_ids.update( + { + rid: self.kv_cache_manager.get_blocks(rid).get_block_ids()[gid] + for rid in num_scheduled_tokens + } + ) + + # Clear the finished request IDs. + # NOTE: We shouldn't do self.finished_req_ids.clear() here because + # it will also affect the scheduler output. + self.finished_req_ids = set() + + def _update_request_as_session( + self, session: Request, update: StreamingUpdate + ) -> None: + """ + Updates the waiting session with the next streaming update. + + Discards the last sampled output token from the prior input chunk. + """ + + # Current streaming input behaviour: Keep only computed output tokens + # (discard final sampled output token). + num_computed_tokens = session.num_computed_tokens + kept_output_tokens = session._all_token_ids[ + session.num_prompt_tokens : num_computed_tokens + ] + del session._all_token_ids[num_computed_tokens:] + session._output_token_ids.clear() + assert session.prompt_token_ids is not None + # Extend prompt with kept output tokens. + session.prompt_token_ids.extend(kept_output_tokens) + + if update.mm_features: + base = session.num_tokens + for mm_feature in update.mm_features: + mm_feature.mm_position = replace( + mm_feature.mm_position, offset=mm_feature.mm_position.offset + base + ) + session.mm_features.extend(update.mm_features) + + session._all_token_ids.extend(update.prompt_token_ids or ()) + session.prompt_token_ids.extend(update.prompt_token_ids or ()) + # Update block hashes for the new tokens. + session.update_block_hashes() + session.num_prompt_tokens = len(session.prompt_token_ids) + session.arrival_time = update.arrival_time + session.sampling_params = update.sampling_params + if session.status == RequestStatus.WAITING_FOR_STREAMING_REQ: + self.num_waiting_for_streaming_input -= 1 + session.status = RequestStatus.WAITING + + if self.log_stats: + session.record_event(EngineCoreEventType.QUEUED) + + def _make_cached_request_data( + self, + running_reqs: list[Request], + resumed_reqs: list[Request], + num_scheduled_tokens: dict[str, int], + spec_decode_tokens: dict[str, list[int]], + req_to_new_blocks: dict[str, KVCacheBlocks], + ) -> CachedRequestData: + req_ids: list[str] = [] + new_token_ids: list[list[int]] = [] + new_block_ids: list[tuple[list[int], ...] | None] = [] + all_token_ids: dict[str, list[int]] = {} + num_computed_tokens: list[int] = [] + num_output_tokens: list[int] = [] + resumed_req_ids = set() + + num_running_reqs = len(running_reqs) + for idx, req in enumerate(itertools.chain(running_reqs, resumed_reqs)): + req_id = req.request_id + req_ids.append(req_id) + # NOTE: In PP+async scheduling, we consume token ids via a direct GPU + # broadcast path (`input_batch.prev_sampled_token_ids`), so we can + # omit this payload. + if self.use_pp and not self.scheduler_config.async_scheduling: + # When using PP, the scheduler sends the sampled tokens back, + # because there's no direct communication between the first- + # stage worker and the last-stage worker. Otherwise, we don't + # need to send the sampled tokens back because the model runner + # will cache them. + num_tokens = num_scheduled_tokens[req_id] - len( + spec_decode_tokens.get(req_id, ()) + ) + token_ids = req.all_token_ids[ + req.num_computed_tokens : req.num_computed_tokens + num_tokens + ] + new_token_ids.append(token_ids) + scheduled_in_prev_step = req_id in self.prev_step_scheduled_req_ids + if idx >= num_running_reqs: + assert not scheduled_in_prev_step + resumed_req_ids.add(req_id) + if not scheduled_in_prev_step: + all_token_ids[req_id] = req.all_token_ids.copy() + new_block_ids.append( + req_to_new_blocks[req_id].get_block_ids(allow_none=True) + ) + num_computed_tokens.append(req.num_computed_tokens) + num_output_tokens.append( + req.num_output_tokens + req.num_output_placeholders + ) + + return CachedRequestData( + req_ids=req_ids, + resumed_req_ids=resumed_req_ids, + new_token_ids=new_token_ids, + all_token_ids=all_token_ids, + new_block_ids=new_block_ids, + num_computed_tokens=num_computed_tokens, + num_output_tokens=num_output_tokens, + ) + + def _try_schedule_encoder_inputs( + self, + request: Request, + num_computed_tokens: int, + num_new_tokens: int, + encoder_compute_budget: int, + shift_computed_tokens: int = 0, + ) -> tuple[list[int], int, 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. + + An encoder input will be scheduled if: + - Its output tokens overlap with the range of tokens being computed + in this step, i.e., + [num_computed_tokens, num_computed_tokens + num_new_tokens). + - It is not already computed and stored in the encoder cache. + - It is not exist on remote encoder cache (via ECConnector) + - There is sufficient encoder token budget to process it. + - The encoder cache has space to store it. + + If an encoder input cannot be scheduled due to cache or budget + limitations, the method adjusts `num_new_tokens` to schedule only the + decoder tokens up to just before the unschedulable encoder input. + + Note that num_computed_tokens includes both locally cached + 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, [] + 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 = [] + + # NOTE: since scheduler operates on the request level (possibly with + # multiple encoder inputs per request), we need to create temporary + # trackers for accounting at the encoder input level. + mm_hashes_to_schedule = set() + num_embeds_to_schedule = 0 + for i, mm_feature in enumerate(mm_features): + start_pos = mm_feature.mm_position.offset + num_encoder_tokens = mm_feature.mm_position.length + num_encoder_embeds = mm_feature.mm_position.get_num_embeds() + item_identifier = mm_feature.identifier + + # The encoder output is needed if the two ranges overlap: + # [num_computed_tokens, num_computed_tokens + num_new_tokens) and + # [start_pos, start_pos + num_encoder_tokens) + if ( + start_pos + >= num_computed_tokens + num_new_tokens + shift_computed_tokens + ): + # The encoder input is not needed in this step. + break + + if self.is_encoder_decoder and num_computed_tokens > 0: + assert start_pos == 0, ( + "Encoder input should be processed at the beginning of " + "the sequence when encoder-decoder models are used." + ) + # Encoder input has already been computed + # The calculation here is a bit different. We don't turn encoder + # output into tokens that get processed by the decoder and + # reflected in num_computed_tokens. Instead, start_pos reflects + # the position where we need to ensure we calculate encoder + # inputs. This should always be 0 to ensure we calculate encoder + # inputs before running the decoder. Once we've calculated some + # decoder tokens (num_computed_tokens > 0), then we know we + # already calculated encoder inputs and can skip here. + continue + elif start_pos + num_encoder_tokens <= num_computed_tokens: + # The encoder input is already computed and stored + # in the decoder's KV cache. + continue + + if not self.is_encoder_decoder: + # We are not using the encoder cache for encoder-decoder models, + # yet. + if item_identifier in mm_hashes_to_schedule: + # The same encoder input has already been scheduled in the + # current step. + continue + + if self.encoder_cache_manager.check_and_update_cache(request, i): + # The encoder input is already computed and cached from a + # 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 + ): + # The encoder cache is full or the encoder budget is exhausted. + # NOTE(woosuk): We assume that the encoder input tokens should + # be processed altogether, as the encoder usually uses + # bidirectional attention. + if num_computed_tokens + shift_computed_tokens < start_pos: + # We only schedule the decoder tokens just before the + # encoder input. + num_new_tokens = start_pos - ( + num_computed_tokens + shift_computed_tokens + ) + else: + # Because of prefix caching, num_computed_tokens is greater + # than start_pos even though its encoder input is not + # available. In this case, we can't schedule any token for + # the request in this step. + num_new_tokens = 0 + break + + # Calculate the number of embeddings to schedule in the current range + # of scheduled encoder placeholder tokens. + start_idx_rel = max(0, num_computed_tokens - start_pos) + end_idx_rel = min( + num_encoder_tokens, num_computed_tokens + num_new_tokens - start_pos + ) + curr_embeds_start, curr_embeds_end = ( + mm_feature.mm_position.get_embeds_indices_in_range( + start_idx_rel, end_idx_rel + ) + ) + # There's no embeddings in the current range of encoder placeholder tokens + # so we can skip the encoder input. + if curr_embeds_end - curr_embeds_start == 0: + continue + + if self.ec_connector is not None and self.ec_connector.has_cache_item( + item_identifier + ): + mm_hashes_to_schedule.add(item_identifier) + external_load_encoder_input.append(i) + num_embeds_to_schedule += num_encoder_embeds + continue + + num_embeds_to_schedule += num_encoder_embeds + encoder_compute_budget -= num_encoder_embeds + mm_hashes_to_schedule.add(item_identifier) + encoder_inputs_to_schedule.append(i) + + return ( + encoder_inputs_to_schedule, + num_new_tokens, + encoder_compute_budget, + external_load_encoder_input, + ) + + def get_grammar_bitmask( + self, scheduler_output: SchedulerOutput + ) -> GrammarOutput | None: + # Collect list of scheduled request ids that use structured output. + # The corresponding rows of the bitmask will be in this order. + if not scheduler_output.has_structured_output_requests: + return None + + structured_output_request_ids = [ + req_id + for req_id in scheduler_output.num_scheduled_tokens + if (req := self.requests.get(req_id)) + and (req.use_structured_output and not req.is_prefill_chunk) + ] + if not structured_output_request_ids: + return None + + bitmask = self.structured_output_manager.grammar_bitmask( + self.requests, + structured_output_request_ids, + scheduler_output.scheduled_spec_decode_tokens, + ) + return GrammarOutput(structured_output_request_ids, bitmask) + + def update_from_output( + self, + scheduler_output: SchedulerOutput, + model_runner_output: ModelRunnerOutput, + ) -> dict[int, EngineCoreOutputs]: + sampled_token_ids = model_runner_output.sampled_token_ids + logprobs = model_runner_output.logprobs + prompt_logprobs_dict = model_runner_output.prompt_logprobs_dict + num_scheduled_tokens = scheduler_output.num_scheduled_tokens + pooler_outputs = model_runner_output.pooler_output + num_nans_in_logits = model_runner_output.num_nans_in_logits + kv_connector_output = model_runner_output.kv_connector_output + cudagraph_stats = model_runner_output.cudagraph_stats + draft_token_lengths = model_runner_output.draft_token_lengths + + perf_stats: PerfStats | None = None + if self.perf_metrics and self.perf_metrics.is_enabled(): + perf_stats = self.perf_metrics.get_step_perf_stats_per_gpu(scheduler_output) + + outputs: dict[int, list[EngineCoreOutput]] = defaultdict(list) + spec_decoding_stats: SpecDecodingStats | None = None + kv_connector_stats: KVConnectorStats | None = ( + kv_connector_output.kv_connector_stats if kv_connector_output else None + ) + if kv_connector_stats and self.connector: + kv_stats = self.connector.get_kv_connector_stats() + if kv_stats: + kv_connector_stats = kv_connector_stats.aggregate(kv_stats) + + failed_kv_load_req_ids = None + if kv_connector_output and kv_connector_output.invalid_block_ids: + # These blocks contain externally computed tokens that failed to + # load. Identify affected requests and adjust their computed token + # count to trigger recomputation of the invalid blocks. + failed_kv_load_req_ids = self._handle_invalid_blocks( + kv_connector_output.invalid_block_ids, + num_scheduled_tokens, + ) + + # Persist per-step routed experts into the scheduler-side slot + # buffer (CPU->CPU fancy-index assign; ~few MB per step). + # MUST precede the per-request routing reads below: stopped + # requests may terminate on tokens generated in this very step, + # whose routing was just D2H'd into model_runner_output. + routing_data = None + routing_offsets: dict[str, int] = {} + if model_runner_output.routed_experts is not None: + re = model_runner_output.routed_experts + self.routed_experts_mgr.store_batch(re.routing_data, re.slot_mapping) + routing_data = re.routing_data.astype( + self.routed_experts_mgr.routed_experts_by_slot.dtype, + copy=False, + ) + # Build offset map using model runner's request order + # (input_batch ordering), NOT scheduler dict order. + offset = 0 + for rid in model_runner_output.req_ids: + routing_offsets[rid] = offset + offset += num_scheduled_tokens[rid] + + # NOTE(woosuk): As len(num_scheduled_tokens) can be up to 1K or more, + # the below loop can be a performance bottleneck. We should do our best + # to avoid expensive operations inside the loop. + stopped_running_reqs: set[Request] = set() + stopped_preempted_reqs: set[Request] = set() + for req_id, num_tokens_scheduled in num_scheduled_tokens.items(): + assert num_tokens_scheduled > 0 + if failed_kv_load_req_ids and req_id in failed_kv_load_req_ids: + # skip failed or rescheduled requests from KV load failure + continue + request = self.requests.get(req_id) + if request is None or request.is_finished(): + # The request is already finished. This can happen if the + # request is aborted while the model is executing it (e.g., + # in pipeline parallelism or in async scheduling). + # NOTE(Kuntai): When delay_free_blocks=True (for async KV + # cache transfer in KV connector), the aborted request will not + # be set to None (in order to finish async KV transfer). + # In this case, we use is_finished() to check. + continue + + req_index = model_runner_output.req_id_to_index[req_id] + generated_token_ids = ( + sampled_token_ids[req_index] if sampled_token_ids else [] + ) + + scheduled_spec_token_ids = ( + scheduler_output.scheduled_spec_decode_tokens.get(req_id) + ) + if scheduled_spec_token_ids and generated_token_ids: + num_draft_tokens = len(scheduled_spec_token_ids) + num_accepted = len(generated_token_ids) - 1 + num_rejected = num_draft_tokens - num_accepted + # num_computed_tokens represents the number of tokens + # processed in the current step, considering scheduled + # tokens and rejections. If some tokens are rejected, + # num_computed_tokens is decreased by the number of rejected + # tokens. + if request.num_computed_tokens > 0: + request.num_computed_tokens -= num_rejected + # If async scheduling, num_output_placeholders also includes + # the scheduled spec tokens count and so is similarly adjusted. + if request.num_output_placeholders > 0: + request.num_output_placeholders -= num_rejected + spec_decoding_stats = self.make_spec_decoding_stats( + spec_decoding_stats, + num_draft_tokens=num_draft_tokens, + num_accepted_tokens=num_accepted, + num_invalid_spec_tokens=scheduler_output.num_invalid_spec_tokens, + request_id=req_id, + ) + + # Free encoder inputs only after the step has actually executed. + if request.has_encoder_inputs: + self._free_encoder_inputs(request) + + stopped = False + new_logprobs = None + new_token_ids = generated_token_ids + pooler_output = pooler_outputs[req_index] if pooler_outputs else None + kv_transfer_params = None + status_before_stop = request.status + num_output_tokens_before = len(request._output_token_ids) + + # Check for stop and update request status. + if new_token_ids: + new_token_ids, stopped = self._update_request_with_output( + request, new_token_ids + ) + elif request.pooling_params and pooler_output is not None: + # Pooling stops as soon as there is output. + request.status = RequestStatus.FINISHED_STOPPED + stopped = True + + if new_token_ids and self.structured_output_manager.should_advance(request): + struct_output_request = request.structured_output_request + assert struct_output_request is not None + assert struct_output_request.grammar is not None + if not struct_output_request.grammar.accept_tokens( # type: ignore[union-attr] + req_id, new_token_ids + ): + logger.error( + "Unexpected: grammar rejected tokens %s for request %s. " + "Terminating request.", + new_token_ids, + req_id, + ) + request.status = RequestStatus.FINISHED_ERROR + request.resumable = False + stopped = True + + routed_experts = None + if ( + self.enable_return_routed_experts + and routing_data is not None + and new_token_ids + ): + req_offset = routing_offsets[req_id] + end = req_offset + num_tokens_scheduled + block_ids = self._re_block_ids.pop(req_id, []) + if num_output_tokens_before == 0: + # Prefill completed: read full prompt routing from + # slot buffer using the block-ID snapshot taken at + # schedule time (immune to async preemption). + if ( + request.sampling_params is not None + and request.sampling_params.routed_experts_prompt_start + is not None + ): + prompt_start = ( + request.sampling_params.routed_experts_prompt_start + ) + assert prompt_start < request.num_prompt_tokens + else: + prompt_start = 0 + routed_experts = self.routed_experts_mgr.get( + block_ids, + request.num_prompt_tokens, + token_start=prompt_start, + ) + else: + if scheduled_spec_token_ids: + # Spec decode: accepted tokens at the START of + # the scheduled range, rejected at the end. + routed_experts = routing_data[ + req_offset : req_offset + len(new_token_ids) + ] + else: + # Normal decode / re-prefill: token(s) at the END. + routed_experts = routing_data[end - len(new_token_ids) : end] + + finish_reason = None + if stopped: + # Capture finish_reason BEFORE _handle_stopped_request, which may + # reset the status to WAITING for streaming requests that continue. + finish_reason = request.get_finished_reason() + finished = self._handle_stopped_request(request) + if finished: + kv_transfer_params = self._free_request(request) + + if status_before_stop == RequestStatus.RUNNING: + stopped_running_reqs.add(request) + else: + stopped_preempted_reqs.add(request) + elif ( + draft_token_lengths is not None + and self.scheduler_config.async_scheduling + # Patch 3 (root-cause fix, credit: roady001, issue #3): only resize + # spec placeholders for requests the AsyncScheduler itself would give + # placeholders. Without these guards the resize ran for EVERY running + # request — including ones mid chunked-prefill (upstream never installs + # spec placeholders on prefill chunks) — which attached spec tokens to + # the final prompt chunk of a long COLD resume, corrupting the prompt + # tail: prompt echo / leaked tool-schema text at the start of the reply, + # then recovering once pure decode began. Guard it to decode steps only. + and new_token_ids + and not request.is_prefill_chunk + and request.status == RequestStatus.RUNNING + ): + # Async speculative scheduling normally installs a fixed + # placeholder list before the worker proposes the next draft. + # DSpark can produce a shorter confidence-scheduled prefix, so + # resize the placeholder list here for the next scheduler step. + draft_len = draft_token_lengths.get(req_id) + if draft_len is not None: + draft_len = max(0, min(int(draft_len), self.num_spec_tokens)) + request.spec_token_ids = [-1] * draft_len + + # Extract sample logprobs if needed. + if ( + request.sampling_params is not None + and request.sampling_params.num_logprobs is not None + and logprobs + ): + new_logprobs = logprobs.slice_request(req_index, len(new_token_ids)) + + if num_nans_in_logits is not None and req_id in num_nans_in_logits: + request.num_nans_in_logits = num_nans_in_logits[req_id] + + # Get prompt logprobs for this request. + prompt_logprobs_tensors = prompt_logprobs_dict.get(req_id) + if ( + new_token_ids + or pooler_output is not None + or kv_transfer_params + or stopped + ): + # Add EngineCoreOutput for this Request. + outputs[request.client_index].append( + EngineCoreOutput( + request_id=req_id, + new_token_ids=new_token_ids, + finish_reason=finish_reason, + new_logprobs=new_logprobs, + new_prompt_logprobs_tensors=prompt_logprobs_tensors, + pooling_output=pooler_output, + stop_reason=request.stop_reason, + events=request.take_events(), + prefill_stats=request.take_prefill_stats(), + kv_transfer_params=kv_transfer_params, + trace_headers=request.trace_headers, + routed_experts=routed_experts, + num_nans_in_logits=request.num_nans_in_logits, + ) + ) + else: + # Invariant: EngineCore returns no partial prefill outputs. + assert not prompt_logprobs_tensors + + # Remove the stopped requests from the running and waiting queues. + if stopped_running_reqs: + self.running = remove_all(self.running, stopped_running_reqs) + if stopped_preempted_reqs: + # This is a rare case and unlikely to impact performance. + self.waiting.remove_requests(stopped_preempted_reqs) + + if failed_kv_load_req_ids and not self.recompute_kv_load_failures: + requests = [self.requests[req_id] for req_id in failed_kv_load_req_ids] + self.finish_requests(failed_kv_load_req_ids, RequestStatus.FINISHED_ERROR) + for request in requests: + outputs[request.client_index].append( + EngineCoreOutput( + request_id=request.request_id, + new_token_ids=[], + finish_reason=request.get_finished_reason(), + events=request.take_events(), + trace_headers=request.trace_headers, + ) + ) + + # KV Connector: update state for finished KV Transfers. + if kv_connector_output: + self._update_from_kv_xfer_finished(kv_connector_output) + + # collect KV cache events from KV cache manager + events = self.kv_cache_manager.take_events() + + # collect KV cache events from connector + if self.connector is not None: + connector_events = self.connector.take_events() + if connector_events: + if events is None: + events = list(connector_events) + else: + events.extend(connector_events) + + # publish collected KV cache events + if events: + batch = KVEventBatch(ts=time.time(), events=events) + self.kv_event_publisher.publish(batch) + + # Create EngineCoreOutputs for all clients that have requests with + # outputs in this step. + engine_core_outputs = { + client_index: EngineCoreOutputs(outputs=outs) + for client_index, outs in outputs.items() + } + + finished_req_ids = self.finished_req_ids_dict + if finished_req_ids: + # Include ids of requests that finished since last outputs + # were sent. + for client_index, finished_set in finished_req_ids.items(): + # Set finished request set in EngineCoreOutputs for this client. + if (eco := engine_core_outputs.get(client_index)) is not None: + eco.finished_requests = finished_set + else: + engine_core_outputs[client_index] = EngineCoreOutputs( + finished_requests=finished_set + ) + finished_req_ids.clear() + + if ( + stats := self.make_stats( + spec_decoding_stats, kv_connector_stats, cudagraph_stats, perf_stats + ) + ) is not None: + # Return stats to only one of the front-ends. + if (eco := next(iter(engine_core_outputs.values()), None)) is None: + # We must return the stats even if there are no request + # outputs this step. + engine_core_outputs[0] = eco = EngineCoreOutputs() + eco.scheduler_stats = stats + + return engine_core_outputs + + @staticmethod + def _is_blocked_waiting_status(status: RequestStatus) -> bool: + return status in ( + RequestStatus.WAITING_FOR_STRUCTURED_OUTPUT_GRAMMAR, + RequestStatus.WAITING_FOR_REMOTE_KVS, + RequestStatus.WAITING_FOR_STREAMING_REQ, + ) + + def _enqueue_waiting_request(self, request: Request) -> None: + if self._is_blocked_waiting_status(request.status): + self.skipped_waiting.add_request(request) + else: + self.waiting.add_request(request) + + def _select_waiting_queue_for_scheduling(self) -> RequestQueue | None: + if self.policy == SchedulingPolicy.FCFS: + return self.skipped_waiting or self.waiting or None + + # PRIORITY mode: compare queue heads when both queues are non-empty. + if self.waiting and self.skipped_waiting: + waiting_req = self.waiting.peek_request() + skipped_req = self.skipped_waiting.peek_request() + return self.waiting if waiting_req < skipped_req else self.skipped_waiting + + return self.waiting or self.skipped_waiting or None + + def _handle_stopped_request(self, request: Request) -> bool: + """Return True if finished (can be False for resumable requests).""" + if not request.resumable: + return True + + if request.streaming_queue: + update = request.streaming_queue.popleft() + if update is None: + # Streaming request finished. + return True + self._update_request_as_session(request, update) + else: + request.status = RequestStatus.WAITING_FOR_STREAMING_REQ + self.num_waiting_for_streaming_input += 1 + + self._enqueue_waiting_request(request) + return False + + def _update_request_with_output( + self, request: Request, new_token_ids: list[int] + ) -> tuple[list[int], bool]: + # Append generated tokens and check for stop. Note that if + # a request is still being prefilled, we expect the model runner + # to return empty token ids for the request. + stopped = False + for num_new, output_token_id in enumerate(new_token_ids, 1): + request.append_output_token_ids(output_token_id) + + # Check for stop and update request state. + # This must be called before we make the EngineCoreOutput. + stopped = check_stop(request, self.max_model_len) + if stopped: + del new_token_ids[num_new:] # Trim new tokens if needed. + break + return new_token_ids, stopped + + def _free_encoder_inputs(self, request: Request) -> None: + cached_encoder_input_ids = self.encoder_cache_manager.get_cached_input_ids( + request + ) + # OPTIMIZATION: Avoid list(set) if the set is empty. + if not cached_encoder_input_ids: + return + + # Here, we use list(set) to avoid modifying the set while iterating + # over it. + for input_id in list(cached_encoder_input_ids): + mm_feature = request.mm_features[input_id] + start_pos = mm_feature.mm_position.offset + num_tokens = mm_feature.mm_position.length + if self.is_encoder_decoder and request.num_computed_tokens > 0: + # With Whisper, as soon as we've generated a single token, + # we know we're done with the encoder input. Cross Attention + # KVs have been calculated and cached already. + self.encoder_cache_manager.free_encoder_input(request, input_id) + elif start_pos + num_tokens <= request.num_computed_tokens: + # The encoder output is already processed and stored + # in the decoder's KV cache. + self.encoder_cache_manager.free_encoder_input(request, input_id) + + def update_draft_token_ids(self, draft_token_ids: DraftTokenIds) -> None: + for req_id, spec_token_ids in zip( + draft_token_ids.req_ids, + draft_token_ids.draft_token_ids, + ): + request = self.requests.get(req_id) + if request is None or request.is_finished(): + # The request may have been finished. Skip. + continue + + if request.is_prefill_chunk: + # Ignore draft tokens for prefill chunks. + if request.spec_token_ids: + request.spec_token_ids = [] + continue + + # Add newly generated spec token ids to the request. + if self.structured_output_manager.should_advance(request): + metadata = request.structured_output_request + spec_token_ids = metadata.grammar.validate_tokens(spec_token_ids) # type: ignore[union-attr] + request.spec_token_ids = spec_token_ids + + def update_draft_token_ids_in_output( + self, draft_token_ids: DraftTokenIds, scheduler_output: SchedulerOutput + ) -> None: + num_invalid_spec_tokens: dict[str, int] = {} + + sched_spec_tokens = scheduler_output.scheduled_spec_decode_tokens + for req_id, spec_token_ids in zip( + draft_token_ids.req_ids, + draft_token_ids.draft_token_ids, + ): + request = self.requests.get(req_id) + if request is None or request.is_finished(): + # The request may have been finished. Skip. + continue + + placeholder_spec_tokens = sched_spec_tokens.get(req_id) + if not placeholder_spec_tokens: + continue + + orig_num_spec_tokens = len(placeholder_spec_tokens) + # Trim drafts to scheduled number of spec tokens + # (needed for chunked prefill case for example). + del spec_token_ids[orig_num_spec_tokens:] + # Filter out spec tokens which do not adhere to the grammar. + if self.structured_output_manager.should_advance(request): + metadata = request.structured_output_request + assert metadata is not None and metadata.grammar is not None + spec_token_ids = metadata.grammar.validate_tokens(spec_token_ids) + # Pad to original number of spec tokens. + num_invalid_tokens = orig_num_spec_tokens - len(spec_token_ids) + if num_invalid_tokens: + spec_token_ids.extend([-1] * num_invalid_tokens) + num_invalid_spec_tokens[req_id] = num_invalid_tokens + + sched_spec_tokens[req_id] = spec_token_ids + + scheduler_output.num_invalid_spec_tokens = num_invalid_spec_tokens + + def get_request_counts(self) -> tuple[int, int]: + """Returns (num_running_reqs, num_waiting_reqs).""" + return len(self.running), len(self.waiting) + len(self.skipped_waiting) + + def add_request(self, request: Request) -> None: + existing = self.requests.get(request.request_id) + if existing is not None: + update = StreamingUpdate.from_request(request) + if existing.status != RequestStatus.WAITING_FOR_STREAMING_REQ: + assert existing.streaming_queue is not None, "duplicate request id" + # Queue next input chunk (or finished sentinel). + existing.streaming_queue.append(update) + elif update is not None: + # Commence next input chunk. + self._update_request_as_session(existing, update) + else: + # Streaming-input session finished. + self.finish_requests(request.request_id, RequestStatus.FINISHED_ABORTED) + else: + if request.resumable: + request.streaming_queue = deque() + self._enqueue_waiting_request(request) + self.requests[request.request_id] = request + if self.connector is not None: + self.connector.on_new_request(request) + if self.log_stats: + request.record_event(EngineCoreEventType.QUEUED) + + def finish_requests( + self, request_ids: str | Iterable[str] | None, finished_status: RequestStatus + ) -> list[tuple[str, int]]: + """Handles the finish signal from outside the scheduler. + + For example, the API server can abort a request when the client + disconnects. + + If request_ids is None, all requests will be finished. + + Returns: + Tuple of (req_id, client_index) for requests that were aborted. Will not + include any that were already finished. + """ + assert RequestStatus.is_finished(finished_status) + if isinstance(request_ids, str): + request_ids = (request_ids,) + elif request_ids is not None: + request_ids = set(request_ids) + else: + request_ids = self.requests.keys() + + running_requests_to_remove = set() + waiting_requests_to_remove = [] + valid_requests = [] + + # First pass: collect requests to remove from queues + for req_id in request_ids: + request = self.requests.get(req_id) + if request is None or request.is_finished(): + # Invalid request ID. + continue + + valid_requests.append(request) + if request.status == RequestStatus.RUNNING: + running_requests_to_remove.add(request) + else: + if request.status == RequestStatus.WAITING_FOR_STREAMING_REQ: + self.num_waiting_for_streaming_input -= 1 + waiting_requests_to_remove.append(request) + + # Remove all requests from queues at once for better efficiency + if running_requests_to_remove: + self.running = remove_all(self.running, running_requests_to_remove) + if waiting_requests_to_remove: + self.waiting.remove_requests(waiting_requests_to_remove) + self.skipped_waiting.remove_requests(waiting_requests_to_remove) + + # Second pass: set status and free requests + for request in valid_requests: + delay_free_blocks = False + if request.status == RequestStatus.WAITING_FOR_REMOTE_KVS: + delay_free_blocks = ( + request.request_id not in self.finished_recving_kv_req_ids + ) + self.finished_recving_kv_req_ids.discard(request.request_id) + self.failed_recving_kv_req_ids.discard(request.request_id) + + request.status = finished_status + self._free_request(request, delay_free_blocks=delay_free_blocks) + + return [(r.request_id, r.client_index) for r in valid_requests] + + def _free_request( + self, request: Request, delay_free_blocks: bool = False + ) -> dict[str, Any] | None: + assert request.is_finished() + + connector_delay_free_blocks, kv_xfer_params = self._connector_finished(request) + self.encoder_cache_manager.free(request) + request_id = request.request_id + self.finished_req_ids.add(request_id) + if self.finished_req_ids_dict is not None: + self.finished_req_ids_dict[request.client_index].add(request_id) + + delay_free_blocks |= connector_delay_free_blocks + if not delay_free_blocks: + self._free_blocks(request) + + return kv_xfer_params + + def _free_blocks(self, request: Request): + assert request.is_finished() + self.kv_cache_manager.free(request) + del self.requests[request.request_id] + + @property + def pause_state(self) -> PauseState: + return self._pause_state + + def set_pause_state(self, pause_state: PauseState) -> None: + self._pause_state = pause_state + + def get_num_unfinished_requests(self) -> int: + if self._pause_state == PauseState.PAUSED_ALL: + return 0 + if self._pause_state == PauseState.PAUSED_NEW: + return len(self.running) + num_waiting = ( + len(self.waiting) + + len(self.skipped_waiting) + - self.num_waiting_for_streaming_input + ) + return num_waiting + len(self.running) + + def has_finished_requests(self) -> bool: + if self.finished_req_ids: + return True + if self.connector is None: + return False + # Finished requests waiting on delayed connector cleanup remain in + # self.requests after they have been removed from scheduling queues. + num_in_queues = ( + len(self.waiting) + len(self.skipped_waiting) + len(self.running) + ) + return len(self.requests) > num_in_queues + + def reset_prefix_cache( + self, reset_running_requests: bool = False, reset_connector: bool = False + ) -> bool: + """Reset the KV prefix cache. + + If reset_running_requests is True, all the running requests will be + preempted and moved to the waiting queue. + Otherwise, this method will only reset the KV prefix cache when there + is no running requests taking KV cache. + """ + if reset_running_requests: + # For logging. + timestamp = time.monotonic() + # Invalidate all the current running requests KV's by pushing them to + # the waiting queue. In this case, we can reduce the ref count of all + # the kv blocks to 0 and thus we can make sure the reset is successful. + # Preempt in reverse order so the requests will be added back to the + # running queue in FIFO order. + while self.running: + request = self.running.pop() + self._preempt_request(request, timestamp) + # For async scheduling, any output frames already in flight at + # preemption time are now stale and must be discarded when they + # return. num_output_placeholders is exactly that count: 0 if + # the engine has drained (e.g. pause_generation(keep) waited + # for idle), 1 for vanilla async mid-step, or 1 + spec/PP frames + # otherwise. + request.async_tokens_to_discard = request.num_output_placeholders + request.num_output_placeholders = 0 + + # Clear scheduled request ids cache. Since we are forcing preemption + # + resumption in the same step, we must act as if these requests were + # not scheduled in the prior step. They will be flushed from the + # persistent batch in the model runner. + self.prev_step_scheduled_req_ids.clear() + + reset_successful = self.kv_cache_manager.reset_prefix_cache() + if reset_running_requests and not reset_successful: + raise RuntimeError( + "Failed to reset KV cache even when all the running requests are " + "preempted and moved to the waiting queue. This is likely due to " + "the presence of running requests waiting for remote KV transfer, " + "which is not supported yet." + ) + + if reset_connector: + reset_successful = self.reset_connector_cache() and reset_successful + + return reset_successful + + def reset_connector_cache(self) -> bool: + if self.connector is None: + # No connector attached -> nothing to reset, treat as success so + # callers that unconditionally request a connector reset (e.g. as + # part of a cache-clearing cascade after a weight update) don't + # see reset_prefix_cache() flip to False purely because they + # didn't configure a connector. + logger.debug( + "reset_connector requested but no KV connector is configured; " + "treating as no-op success." + ) + return True + + if self.connector.reset_cache() is False: + return False + + if self.log_stats: + assert self.connector_prefix_cache_stats is not None + self.connector_prefix_cache_stats.reset = True + + return True + + def reset_encoder_cache(self) -> None: + """Reset the encoder cache to invalidate all cached encoder outputs. + + This should be called when model weights are updated to ensure + stale vision embeddings are not reused. + """ + self.encoder_cache_manager.reset() + + def make_stats( + self, + spec_decoding_stats: SpecDecodingStats | None = None, + kv_connector_stats: KVConnectorStats | None = None, + cudagraph_stats: CUDAGraphStat | None = None, + perf_stats: PerfStats | None = None, + ) -> SchedulerStats | None: + if not self.log_stats: + return None + prefix_cache_stats = self.kv_cache_manager.make_prefix_cache_stats() + assert prefix_cache_stats is not None + connector_prefix_cache_stats: PrefixCacheStats | None = None + if self.connector_prefix_cache_stats is not None: + connector_prefix_cache_stats = self.connector_prefix_cache_stats + self.connector_prefix_cache_stats = PrefixCacheStats() + eviction_events = ( + self.kv_metrics_collector.drain_events() + if self.kv_metrics_collector is not None + else [] + ) + spec_stats = spec_decoding_stats + connector_stats_payload = ( + kv_connector_stats.data if kv_connector_stats else None + ) + return SchedulerStats( + num_running_reqs=len(self.running), + num_waiting_reqs=len(self.waiting), + num_skipped_waiting_reqs=len(self.skipped_waiting), + kv_cache_usage=self.kv_cache_manager.usage, + prefix_cache_stats=prefix_cache_stats, + connector_prefix_cache_stats=connector_prefix_cache_stats, + kv_cache_eviction_events=eviction_events, + spec_decoding_stats=spec_stats, + kv_connector_stats=connector_stats_payload, + cudagraph_stats=cudagraph_stats, + perf_stats=perf_stats, + ) + + def make_spec_decoding_stats( + self, + spec_decoding_stats: SpecDecodingStats | None, + num_draft_tokens: int, + num_accepted_tokens: int, + num_invalid_spec_tokens: dict[str, int] | None, + request_id: str, + ) -> SpecDecodingStats | None: + if not self.log_stats or not num_draft_tokens: + return None + if spec_decoding_stats is None: + spec_decoding_stats = SpecDecodingStats.new(self.num_spec_tokens) + if num_invalid_spec_tokens: + num_draft_tokens -= num_invalid_spec_tokens.get(request_id, 0) + spec_decoding_stats.observe_draft( + num_draft_tokens=num_draft_tokens, num_accepted_tokens=num_accepted_tokens + ) + return spec_decoding_stats + + def shutdown(self) -> None: + if self.kv_event_publisher: + self.kv_event_publisher.shutdown() + if self.connector is not None: + self.connector.shutdown() + + ######################################################################## + # KV Connector Related Methods + ######################################################################## + + def get_kv_connector(self) -> KVConnectorBase_V1 | None: + return self.connector + + def _connector_finished( + self, request: Request + ) -> tuple[bool, dict[str, Any] | None]: + """ + Invoke the KV connector request_finished() method if applicable. + + Returns optional kv transfer parameters to be included with the + request outputs. + """ + if self.connector is None: + return False, None + + # Free any out-of-window prefix blocks before we hand the block table to + # the connector. + self.kv_cache_manager.remove_skipped_blocks( + request_id=request.request_id, + total_computed_tokens=request.num_computed_tokens, + ) + + block_ids = self.kv_cache_manager.get_block_ids(request.request_id) + + if not isinstance(self.connector, SupportsHMA): + # NOTE(Kuntai): We should deprecate this code path after we enforce + # all connectors to support HMA. + # Hybrid memory allocator should be already turned off for this + # code path, but let's double-check here. + assert len(self.kv_cache_config.kv_cache_groups) == 1 + return self.connector.request_finished(request, block_ids[0]) + + return self.connector.request_finished_all_groups(request, block_ids) + + def _update_waiting_for_remote_kv(self, request: Request) -> None: + """ + KV Connector: update request state after async recv is finished. + + When the kv transfer is ready, we cache the blocks + and the request state will be moved back to WAITING from + WAITING_FOR_REMOTE_KV. + """ + assert self.connector is not None + + if request.request_id in self.failed_recving_kv_req_ids: + # Request had KV load failures; num_computed_tokens was already + # updated in _update_requests_with_invalid_blocks + if request.num_computed_tokens: + # Cache any valid computed tokens. + self.kv_cache_manager.cache_blocks(request, request.num_computed_tokens) + else: + # No valid computed tokens, release allocated blocks. + # There may be a local cache hit on retry. + self.kv_cache_manager.free(request) + + self.failed_recving_kv_req_ids.remove(request.request_id) + else: + # Now that the blocks are ready, actually cache them. + # This will cache the blocks iff caching is enabled. + self.kv_cache_manager.cache_blocks(request, request.num_computed_tokens) + + # on a full prompt hit, we need to re-compute the last token + # in order to be able to sample the next token + if request.num_computed_tokens == request.num_tokens: + request.num_computed_tokens = request.num_tokens - 1 + + self.finished_recving_kv_req_ids.remove(request.request_id) + + def _try_promote_blocked_waiting_request(self, request: Request) -> bool: + """ + Try to promote a blocked waiting request back to schedulable states. + """ + if request.status == RequestStatus.WAITING_FOR_REMOTE_KVS: + # finished_recving_kv_req_ids is populated during + # update_from_output(), based on worker-side connector signals + # in KVConnectorOutput.finished_recving + if request.request_id not in self.finished_recving_kv_req_ids: + return False + self._update_waiting_for_remote_kv(request) + if request.num_preemptions: + request.status = RequestStatus.PREEMPTED + else: + request.status = RequestStatus.WAITING + return True + + if request.status == RequestStatus.WAITING_FOR_STRUCTURED_OUTPUT_GRAMMAR: + structured_output_req = request.structured_output_request + if not (structured_output_req and structured_output_req.grammar): + return False + request.status = RequestStatus.WAITING + return True + + if request.status == RequestStatus.WAITING_FOR_STREAMING_REQ: + assert not request.streaming_queue + return False + + raise AssertionError( + "Unexpected blocked waiting status in promotion: " + f"{request.status.name} for request {request.request_id}" + ) + + def _update_from_kv_xfer_finished(self, kv_connector_output: KVConnectorOutput): + """ + KV Connector: update the scheduler state based on the output. + + The Worker side connectors add finished_recving and + finished_sending reqs to the output. + * if finished_sending: free the blocks + # if finished_recving: add to state so we can + schedule the request during the next step. + """ + + if self.connector is not None: + self.connector.update_connector_output(kv_connector_output) + + # KV Connector:: update recv and send status from last step. + for req_id in kv_connector_output.finished_recving or (): + logger.debug("Finished recving KV transfer for request %s", req_id) + assert req_id in self.requests + req = self.requests[req_id] + if req.status == RequestStatus.WAITING_FOR_REMOTE_KVS: + self.finished_recving_kv_req_ids.add(req_id) + else: + assert RequestStatus.is_finished(req.status) + self._free_blocks(self.requests[req_id]) + for req_id in kv_connector_output.finished_sending or (): + logger.debug("Finished sending KV transfer for request %s", req_id) + assert req_id in self.requests + self._free_blocks(self.requests[req_id]) + + def _update_requests_with_invalid_blocks( + self, + requests: Iterable[Request], + invalid_block_ids: set[int], + num_scheduled_tokens: dict[str, int], + evict_blocks: bool = True, + ) -> tuple[set[str], int, set[int]]: + """ + Identify and update requests affected by invalid KV cache blocks. + + This method scans the given requests, detects those with invalid blocks + and adjusts their `num_computed_tokens` to the longest valid prefix. + For observability, it also accumulates the total number of tokens that + will need to be recomputed across all affected requests. + + Args: + requests: The set of requests to scan for invalid blocks. + invalid_block_ids: IDs of invalid blocks. + num_scheduled_tokens: req_id -> number of scheduled tokens. + evict_blocks: Whether to collect blocks for eviction (False for + async requests which aren't cached yet). + + Returns: + tuple: + - affected_req_ids (set[str]): IDs of requests impacted by + invalid blocks. + - total_affected_tokens (int): Total number of tokens that must + be recomputed across all affected requests. + - blocks_to_evict (set[int]): Block IDs to evict from cache, + including invalid blocks and downstream dependent blocks. + """ + affected_req_ids: set[str] = set() + total_affected_tokens = 0 + blocks_to_evict: set[int] = set() + # If a block is invalid and shared by multiple requests in the batch, + # these requests must be rescheduled, but only the first will recompute + # it. This set tracks blocks already marked for recomputation. + marked_invalid_block_ids: set[int] = set() + for request in requests: + is_affected = False + marked_invalid_block = False + req_id = request.request_id + # TODO (davidb): add support for hybrid memory allocator + (req_block_ids,) = self.kv_cache_manager.get_block_ids(req_id) + # We iterate only over blocks that may contain externally computed + # tokens + req_num_computed_tokens = ( + request.num_computed_tokens - num_scheduled_tokens.get(req_id, 0) + ) + + req_num_computed_blocks = ( + req_num_computed_tokens + self.block_size - 1 + ) // self.block_size + for idx, block_id in zip(range(req_num_computed_blocks), req_block_ids): + if block_id not in invalid_block_ids: + continue + + is_affected = True + + if block_id in marked_invalid_block_ids: + # This invalid block is shared with a previous request + # and was already marked for recomputation. + # This means this request can still consider this block + # as computed when rescheduled. + # Currently this only applies to sync loading; Async + # loading does not yet support block sharing + continue + + marked_invalid_block_ids.add(block_id) + + if marked_invalid_block: + # This request has already marked an invalid block for + # recomputation and updated its num_computed_tokens. + continue + + marked_invalid_block = True + # Truncate the computed tokens at the first failed block + request.num_computed_tokens = idx * self.block_size + num_affected_tokens = ( + req_num_computed_tokens - request.num_computed_tokens + ) + total_affected_tokens += num_affected_tokens + + # collect invalid block and all downstream dependent blocks + if evict_blocks: + blocks_to_evict.update(req_block_ids[idx:]) + + if is_affected: + if not marked_invalid_block: + # All invalid blocks of this request are shared with + # previous requests and will be recomputed by them. + # Revert to considering only cached tokens as computed. + # Currently this only applies to sync loading; Async + # loading does not yet support block sharing + total_affected_tokens += ( + request.num_computed_tokens - req_num_computed_tokens + ) + request.num_computed_tokens = req_num_computed_tokens + + affected_req_ids.add(request.request_id) + + return affected_req_ids, total_affected_tokens, blocks_to_evict + + def _handle_invalid_blocks( + self, invalid_block_ids: set[int], num_scheduled_tokens: dict[str, int] + ) -> set[str]: + """ + Handle requests affected by invalid KV cache blocks. + + Returns: + Set of affected request IDs to skip in update_from_output main loop. + """ + should_fail = not self.recompute_kv_load_failures + + # handle async KV loads (not cached yet, evict_blocks=False) + async_load_reqs = ( + req + for req in self.skipped_waiting + if req.status == RequestStatus.WAITING_FOR_REMOTE_KVS + ) + async_failed_req_ids, num_failed_tokens, _ = ( + self._update_requests_with_invalid_blocks( + async_load_reqs, + invalid_block_ids, + num_scheduled_tokens, + evict_blocks=False, + ) + ) + + total_failed_requests = len(async_failed_req_ids) + total_failed_tokens = num_failed_tokens + + # handle sync loads (may be cached, collect blocks for eviction) + sync_failed_req_ids, num_failed_tokens, sync_blocks_to_evict = ( + self._update_requests_with_invalid_blocks( + self.running, invalid_block_ids, num_scheduled_tokens, evict_blocks=True + ) + ) + + total_failed_requests += len(sync_failed_req_ids) + total_failed_tokens += num_failed_tokens + + if not total_failed_requests: + return set() + + # evict invalid blocks and downstream dependent blocks from cache + # only when not using recompute policy (where blocks will be recomputed + # and reused by other requests sharing them) + if sync_blocks_to_evict and not self.recompute_kv_load_failures: + self.kv_cache_manager.evict_blocks(sync_blocks_to_evict) + + if should_fail: + all_failed_req_ids = async_failed_req_ids | sync_failed_req_ids + logger.error( + "Failing %d request(s) due to KV load failure " + "(failure_policy=fail, %d tokens affected). Request IDs: %s", + total_failed_requests, + total_failed_tokens, + all_failed_req_ids, + ) + return all_failed_req_ids + + logger.warning( + "Recovered from KV load failure: " + "%d request(s) rescheduled (%d tokens affected).", + total_failed_requests, + total_failed_tokens, + ) + + # Mark async requests with KV load failures for retry once loading completes + self.failed_recving_kv_req_ids |= async_failed_req_ids + # Return sync affected IDs to skip in update_from_output + return sync_failed_req_ids diff --git a/mods/deepseek-v4-flash-dspark/overlay/vllm/v1/outputs.py b/mods/deepseek-v4-flash-dspark/overlay/vllm/v1/outputs.py new file mode 100644 index 00000000..a66afe5f --- /dev/null +++ b/mods/deepseek-v4-flash-dspark/overlay/vllm/v1/outputs.py @@ -0,0 +1,339 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from abc import ABC, abstractmethod +from dataclasses import dataclass, field +from typing import TYPE_CHECKING, NamedTuple, TypeAlias + +import numpy as np +import torch + +from vllm.compilation.cuda_graph import CUDAGraphStat +from vllm.v1.core.sched.output import SchedulerOutput + +if TYPE_CHECKING: + from vllm.distributed.kv_events import KVConnectorKVEvents + from vllm.distributed.kv_transfer.kv_connector.v1.base import ( + KVConnectorWorkerMetadata, + ) + from vllm.distributed.kv_transfer.kv_connector.v1.metrics import KVConnectorStats +else: + KVConnectorStats = object + KVConnectorWorkerMetadata = object + KVConnectorKVEvents = object + + +class LogprobsLists(NamedTuple): + # [num_reqs x num_generated_tokens, max_num_logprobs + 1] + logprob_token_ids: np.ndarray + # [num_reqs x num_generated_tokens, max_num_logprobs + 1] + logprobs: np.ndarray + # [num_reqs x num_generated_tokens] + sampled_token_ranks: np.ndarray + # [num_reqs] + # Used for slicing the logprobs in cases like speculative + # decoding where the number of generated tokens may be + # different for each request. + cu_num_generated_tokens: list[int] | None = None + + def slice_request(self, req_idx: int, num_positions: int): + if self.cu_num_generated_tokens is not None: + req_idx = self.cu_num_generated_tokens[req_idx] + end_idx = req_idx + num_positions + return LogprobsLists( + self.logprob_token_ids[req_idx:end_idx], + self.logprobs[req_idx:end_idx], + self.sampled_token_ranks[req_idx:end_idx], + None, + ) + + +class LogprobsTensors(NamedTuple): + # [num_reqs x num_generated_tokens, max_num_logprobs + 1] + logprob_token_ids: torch.Tensor + # [num_reqs x num_generated_tokens, max_num_logprobs + 1] + logprobs: torch.Tensor + # [num_reqs x num_generated_tokens] + selected_token_ranks: torch.Tensor + # [num_reqs] + cu_num_generated_tokens: list[int] | None = None + + def tolists(self, cu_num_generated_tokens: list[int] | None = None): + return LogprobsLists( + self.logprob_token_ids.cpu().numpy(), + self.logprobs.cpu().numpy(), + self.selected_token_ranks.cpu().numpy(), + cu_num_generated_tokens + if cu_num_generated_tokens is not None + else self.cu_num_generated_tokens, + ) + + def to_cpu_nonblocking(self) -> "LogprobsTensors": + if self.logprob_token_ids.device.type == "cpu": + return self + return LogprobsTensors( + self.logprob_token_ids.to("cpu", non_blocking=True), + self.logprobs.to("cpu", non_blocking=True), + self.selected_token_ranks.to("cpu", non_blocking=True), + self.cu_num_generated_tokens, + ) + + def filter(self, mask: torch.Tensor) -> "LogprobsTensors": + """Filter the logprobs tensors with the given bool mask.""" + assert self.cu_num_generated_tokens is None, ( + "filter can't be used with cu_num_generated_tokens" + ) + return LogprobsTensors( + self.logprob_token_ids[mask], + self.logprobs[mask], + self.selected_token_ranks[mask], + ) + + @staticmethod + def empty_cpu( + num_positions: int, num_tokens_per_position: int + ) -> "LogprobsTensors": + """Create empty LogprobsTensors on CPU.""" + + logprob_token_ids = torch.empty( + (num_positions, num_tokens_per_position), dtype=torch.int32, device="cpu" + ) + logprobs = torch.empty_like(logprob_token_ids, dtype=torch.float32) + selected_token_ranks = torch.empty( + num_positions, dtype=torch.int32, device="cpu" + ) + return LogprobsTensors( + logprob_token_ids=logprob_token_ids, + logprobs=logprobs, + selected_token_ranks=selected_token_ranks, + ) + + +class RoutedExpertsTensors(NamedTuple): + """Device-side snapshot of routed experts data, pending async D2H. + + Produced by :class:`GPUModelRunner` at the end of each async-scheduled + step. The copy stream waits on the default stream, then issues + non-blocking D2H via :meth:`to_cpu_nonblocking` into a pinned CPU + buffer; :class:`AsyncGPUModelRunnerOutput.get_output` synchronizes + the copy before the scheduler reads it. + + Sliced to ``total_num_scheduled_tokens`` (step-level, across all + requests — NOT per-request). Both ``routing_data`` and + ``slot_mapping`` must be private clones when sourced from shared + capturer / prepare-input buffers, so the next forward pass / + ``_prepare_inputs`` on the default stream does not race with a + D2H still pending on the copy stream. + """ + + # (num_scheduled_tokens, num_layers, num_experts_per_tok) + routing_data: torch.Tensor + # (num_scheduled_tokens,) + slot_mapping: torch.Tensor + + def to_cpu_nonblocking(self) -> "RoutedExpertsTensors": + """Issue non-blocking D2H on the current stream. + + NOTE: ``non_blocking=True`` only delivers true overlap when the + CPU target is pinned. The current fallback here allocates a + new pageable CPU tensor per call, which silently degrades to a + synchronous copy; acceptable because the sync happens on the + dedicated copy stream, not the default stream. + """ + if self.routing_data.device.type == "cpu": + return self + return RoutedExpertsTensors( + self.routing_data.to("cpu", non_blocking=True), + self.slot_mapping.to("cpu", non_blocking=True), + ) + + def tolists(self) -> "RoutedExpertsLists": + """Convert to the numpy-backed form consumed by the scheduler. + + ``.cpu()`` is a no-op when the tensor is already on CPU, so this + is cheap for the post-D2H case; for raw device tensors it will + synchronously block, which is only reached in tests. + """ + return RoutedExpertsLists( + self.routing_data.cpu().numpy(), + self.slot_mapping.cpu().numpy(), + ) + + +class RoutedExpertsLists(NamedTuple): + """CPU-side routed experts, the form :meth:`RoutedExpertsManager.store_batch` + consumes. + + Batched per scheduler step: the leading dim is the number of tokens + scheduled across all requests in this step (``total_num_scheduled_tokens``), + not per-request tokens. ``slot_mapping[i]`` tells the scheduler which + physical KV-cache slot row ``i`` of ``routing_data`` belongs to. + """ + + # (num_scheduled_tokens, num_layers, num_experts_per_tok) + routing_data: np.ndarray + # (num_scheduled_tokens,) + slot_mapping: np.ndarray + + +# [num_reqs, ] +# The shape of each element depends on the pooler used +PoolerOutput: TypeAlias = torch.Tensor | list[torch.Tensor] | list[torch.Tensor | None] + + +@dataclass +class SamplerOutput: + # [num_reqs, max_num_generated_tokens] + # Different requests can have different number of generated tokens. + # All requests are padded to max_num_generated_tokens. + # PLACEHOLDER_TOKEN_ID (-1 by default) is used for padding. + sampled_token_ids: torch.Tensor + logprobs_tensors: LogprobsTensors | None + + +@dataclass +class KVConnectorOutput: + # [req_ids] + finished_sending: set[str] | None = None + finished_recving: set[str] | None = None + kv_connector_stats: KVConnectorStats | None = None + kv_cache_events: KVConnectorKVEvents | None = None + kv_connector_worker_meta: KVConnectorWorkerMetadata | None = None + # IDs of externally computed KV blocks that failed to load. + # Requests referencing these blocks should be rescheduled to recompute them + invalid_block_ids: set[int] = field(default_factory=set) + # Configuration describing how many finished sending/receiving + # notifications should be expected for each request. This allows + # handshake-based connectors like Nixl to update the KVOutputAggregator. + # It captures a static setup info and should almost always remain constant + # for a given connector after discovery. Default value entails no change. + expected_finished_count: int = 0 + + def is_empty(self): + return ( + not self.finished_sending + and not self.finished_recving + and not self.kv_connector_stats + and not self.kv_cache_events + and not self.invalid_block_ids + and not self.kv_connector_worker_meta + ) + + +@dataclass +class ECConnectorOutput: + # [mm_hash] + finished_sending: set[str] | None = None + finished_recving: set[str] | None = None + + +# ModelRunnerOutput is serialized and sent to the scheduler process. +# This is expensive for torch.Tensor so prefer to use list instead. +@dataclass +class ModelRunnerOutput: + # [num_reqs] + req_ids: list[str] + # req_id -> index + req_id_to_index: dict[str, int] + + # num_reqs x num_generated_tokens + # num_generated_tokens is the number of tokens + # generated in the current step. It can be different for + # each request due to speculative/jump decoding. + sampled_token_ids: list[list[int]] = field(default_factory=list) + + # [num_reqs, max_num_logprobs + 1] + # [num_reqs, max_num_logprobs + 1] + # [num_reqs] + logprobs: LogprobsLists | None = None + + # req_id -> (token_ids, logprobs, ranks) + # [prompt_len, num_prompt_logprobs] + # [prompt_len, num_prompt_logprobs] + # [prompt_len] + prompt_logprobs_dict: dict[str, LogprobsTensors | None] = field( + default_factory=dict + ) + + # [num_reqs, hidden_size] + pooler_output: list[torch.Tensor | None] | None = None + + kv_connector_output: KVConnectorOutput | None = None + + ec_connector_output: ECConnectorOutput | None = None + + # req_id -> num_nans_in_logits + num_nans_in_logits: dict[str, int] | None = None + + # information related to cudagraph execution + cudagraph_stats: CUDAGraphStat | None = None + + # Per-step routed experts data captured by the worker. + # ``routing_data`` shape: (num_scheduled_tokens, num_layers, + # num_experts_per_tok); expert IDs as uint8/uint16. + # ``slot_mapping`` shape: (num_scheduled_tokens,); physical KV-cache + # slot for each row of routing_data. + # ``num_scheduled_tokens`` is step-level (total across all requests + # in this step), not per-request. The scheduler persists this into + # its slot buffer via ``slot_buffer[slot_mapping] = routing_data``. + # ``None`` when ``enable_return_routed_experts`` is off. + routed_experts: RoutedExpertsLists | None = None + + # Optional per-request draft lengths produced by async speculative + # proposers. Async schedulers use this to size the next placeholder list + # before real draft token ids are available on CPU. + draft_token_lengths: dict[str, int] | None = None + + +# ModelRunnerOutput wrapper for async scheduling. +class AsyncModelRunnerOutput(ABC): + @abstractmethod + def get_output(self) -> ModelRunnerOutput: + """Get the ModelRunnerOutput for this async output. + + This is a blocking call that waits until the results are ready, which + might involve copying device tensors to the host. + This method should only be called once per AsyncModelRunnerOutput. + """ + pass + + +@dataclass +class DraftTokenIds: + # [num_reqs] + req_ids: list[str] + # num_reqs x num_draft_tokens + draft_token_ids: list[list[int]] + + +def make_empty_encoder_model_runner_output( + scheduler_output: "SchedulerOutput", +) -> ModelRunnerOutput: + """ + Create a ModelRunnerOutput stub that contains the correct + per-request bookkeeping but no generated data yet. + """ + if not scheduler_output.num_scheduled_tokens: + return EMPTY_MODEL_RUNNER_OUTPUT + + # Convert to list so we get a deterministic, indexable sequence + req_ids: list[str] = list(scheduler_output.num_scheduled_tokens.keys()) + + # Give every request its own contiguous index + req_id_to_index: dict[str, int] = {rid: idx for idx, rid in enumerate(req_ids)} + + # No tokens generated yet ⇒ one empty list per request + sampled_token_ids: list[list[int]] = [[0] for _ in req_ids] + + # Pooler outputs are not available yet ⇒ use None placeholders + pooler_output: list[torch.Tensor | None] = [None for _ in req_ids] + + return ModelRunnerOutput( + req_ids=req_ids, + req_id_to_index=req_id_to_index, + sampled_token_ids=sampled_token_ids, + pooler_output=pooler_output, + ) + + +EMPTY_MODEL_RUNNER_OUTPUT = ModelRunnerOutput(req_ids=[], req_id_to_index={}) diff --git a/mods/deepseek-v4-flash-dspark/overlay/vllm/v1/spec_decode/dspark.py b/mods/deepseek-v4-flash-dspark/overlay/vllm/v1/spec_decode/dspark.py new file mode 100644 index 00000000..8c1d042d --- /dev/null +++ b/mods/deepseek-v4-flash-dspark/overlay/vllm/v1/spec_decode/dspark.py @@ -0,0 +1,683 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from __future__ import annotations + +import math +from collections.abc import Callable, Sequence +from dataclasses import dataclass, field +from typing import Any + +import torch + +StepCurve = Callable[[int], float] + +_STACKED_PARAM_NAME_MAPPING = ( + ("attn.fused_wqa_wkv", ".attn.wq_a", 0), + ("attn.fused_wqa_wkv", ".attn.wkv", 1), +) + + +def map_dspark_stacked_param_name(name: str) -> tuple[str, int] | None: + """Map checkpoint names that load into stacked vLLM parameters. + + Keep this segment-aware: the DSpark checkpoint also has names such as + ``markov_w1`` that must not be treated as FFN ``w1`` shards. + """ + + if ".experts." in name: + return None + for param_name, weight_name, shard_id in _STACKED_PARAM_NAME_MAPPING: + if weight_name in name: + return name.replace(weight_name, f".{param_name}"), shard_id + return None + + +def make_dspark_warmup_draft_token_ids( + *, + batch_size: int, + num_speculative_tokens: int, + noise_token_id: int, + device: torch.device, +) -> torch.Tensor: + """Return a valid synthetic draft block for DSpark's cache-warm step. + + The first DSpark proposal call has prompt target features available, but no + generated-token target feature yet. vLLM's async speculative path still + expects a fixed tensor of draft ids, so use the model's valid noise token as + a conservative one-step proposal rather than Python empty lists or -1 + placeholders. + """ + + if batch_size <= 0: + raise ValueError(f"batch_size must be positive, got {batch_size}") + if num_speculative_tokens <= 0: + raise ValueError( + f"num_speculative_tokens must be positive, got {num_speculative_tokens}" + ) + if noise_token_id < 0: + raise ValueError(f"noise_token_id must be non-negative, got {noise_token_id}") + return torch.full( + (batch_size, num_speculative_tokens), + int(noise_token_id), + dtype=torch.int32, + device=device, + ) + + +def unpack_mhc_pre_outputs( + outputs: tuple[torch.Tensor, torch.Tensor, torch.Tensor], +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Return MHC pre outputs as ``(layer_input, post_mix, comb_mix)``. + + The direct mHC pre op returns ``(post_mix, comb_mix, layer_input)`` while + DSpark's draft layer wants to mirror the decoder block local names. + """ + + post_mix, comb_mix, layer_input = outputs + return layer_input, post_mix, comb_mix + + +@dataclass(frozen=True) +class DSparkModelSpec: + """Normalized DSpark fields from a model config.""" + + block_size: int + noise_token_id: int + target_layer_ids: tuple[int, ...] + markov_rank: int + markov_head_type: str + confidence_head_with_markov: bool + weight_prefix: str | None = None + num_draft_layers: int | None = None + + @classmethod + def from_hf_config(cls, hf_config: Any) -> DSparkModelSpec: + if not has_dspark_config(hf_config): + raise ValueError("hf_config does not contain DSpark fields") + + block_size = int(_get_config_value(hf_config, "dspark_block_size")) + if block_size <= 0: + raise ValueError(f"dspark_block_size must be positive, got {block_size}") + + noise_token_id = int(_get_config_value(hf_config, "dspark_noise_token_id")) + if noise_token_id < 0: + raise ValueError( + f"dspark_noise_token_id must be non-negative, got {noise_token_id}" + ) + + raw_layer_ids = _get_config_value(hf_config, "dspark_target_layer_ids") + target_layer_ids = tuple(int(layer_id) for layer_id in raw_layer_ids) + if not target_layer_ids: + raise ValueError("dspark_target_layer_ids must not be empty") + if any( + layer_id <= previous + for previous, layer_id in zip(target_layer_ids, target_layer_ids[1:]) + ): + raise ValueError("dspark_target_layer_ids must be strictly increasing") + + num_hidden_layers = _get_config_value(hf_config, "num_hidden_layers", None) + if num_hidden_layers is not None: + upper = int(num_hidden_layers) - 1 + for layer_id in target_layer_ids: + if layer_id != -1 and not 0 <= layer_id <= upper: + raise ValueError( + "dspark_target_layer_ids contains layer " + f"{layer_id}, outside {{-1}} U [0, {upper}]" + ) + + markov_rank = int(_get_config_value(hf_config, "dspark_markov_rank", 0)) + if markov_rank < 0: + raise ValueError( + f"dspark_markov_rank must be non-negative, got {markov_rank}" + ) + + markov_head_type = _get_config_value(hf_config, "dspark_markov_head_type", None) + if markov_head_type is None: + markov_head_type = "vanilla" + markov_head_type = str(markov_head_type).lower() + + confidence_head_with_markov = _get_config_value( + hf_config, "dspark_confidence_head_with_markov", None + ) + if confidence_head_with_markov is None: + confidence_head_with_markov = markov_rank > 0 + + return cls( + block_size=block_size, + noise_token_id=noise_token_id, + target_layer_ids=target_layer_ids, + markov_rank=markov_rank, + markov_head_type=markov_head_type, + confidence_head_with_markov=bool(confidence_head_with_markov), + weight_prefix=infer_dspark_weight_prefix( + _get_config_value(hf_config, "_weight_names", ()) + ), + num_draft_layers=infer_dspark_num_draft_layers( + _get_config_value(hf_config, "_weight_names", ()) + ), + ) + + def confidence_input_dim(self, hidden_size: int) -> int: + hidden_size = int(hidden_size) + if hidden_size <= 0: + raise ValueError(f"hidden_size must be positive, got {hidden_size}") + if self.confidence_head_with_markov: + return hidden_size + self.markov_rank + return hidden_size + + +@dataclass(frozen=True) +class DSparkScheduleResult: + """Selected DSpark verification lengths and their profiled throughput.""" + + lengths: tuple[int, ...] + expected_accepted_tokens: float + batch_tokens: int + expected_tokens_per_second: float + + +@dataclass(frozen=True) +class DSparkDiagnosticsSnapshot: + """Aggregated DSpark scheduling signals for logs or metrics exporters.""" + + num_steps: int + num_requests: int + num_possible_draft_tokens: int + num_scheduled_draft_tokens: int + scheduled_length_histogram: tuple[int, ...] + avg_scheduled_length: float + draft_token_prune_rate: float + expected_acceptance_length: float + avg_expected_tokens_per_second: float + avg_confidence_per_pos: tuple[float, ...] + avg_survival_per_pos: tuple[float, ...] + scheduled_fraction_per_pos: tuple[float, ...] + + +@dataclass(frozen=True) +class DSparkPosition0DiagnosticsSnapshot: + """Aggregated position-0 draft quality signals.""" + + num_tokens: int + num_matches: int + match_rate: float + avg_confidence: float | None + avg_confidence_when_matched: float | None + avg_confidence_when_missed: float | None + num_confidence_logits_normalized: int + + +@dataclass +class DSparkDiagnostics: + """Accumulate DSpark confidence-scheduler diagnostics. + + This intentionally sits below Prometheus/logging integration so the DSpark + proposer can collect useful signals before we settle the runtime wiring. + """ + + max_spec_tokens: int + num_steps: int = 0 + num_requests: int = 0 + num_possible_draft_tokens: int = 0 + num_scheduled_draft_tokens: int = 0 + total_expected_accepted_tokens: float = 0.0 + total_expected_tokens_per_second: float = 0.0 + scheduled_length_histogram: list[int] = field(default_factory=list) + confidence_sums: list[float] = field(default_factory=list) + confidence_counts: list[int] = field(default_factory=list) + survival_sums: list[float] = field(default_factory=list) + survival_counts: list[int] = field(default_factory=list) + scheduled_counts: list[int] = field(default_factory=list) + + def __post_init__(self) -> None: + self.max_spec_tokens = int(self.max_spec_tokens) + if self.max_spec_tokens <= 0: + raise ValueError( + f"max_spec_tokens must be positive, got {self.max_spec_tokens}" + ) + self.scheduled_length_histogram = [0] * (self.max_spec_tokens + 1) + self.confidence_sums = [0.0] * self.max_spec_tokens + self.confidence_counts = [0] * self.max_spec_tokens + self.survival_sums = [0.0] * self.max_spec_tokens + self.survival_counts = [0] * self.max_spec_tokens + self.scheduled_counts = [0] * self.max_spec_tokens + + def observe( + self, + confidence_rows: Sequence[Sequence[float]], + schedule_result: DSparkScheduleResult, + ) -> None: + if len(confidence_rows) != len(schedule_result.lengths): + raise ValueError( + "confidence_rows and schedule_result.lengths must have the same length" + ) + + self.num_steps += 1 + self.num_requests += len(confidence_rows) + self.num_possible_draft_tokens += sum(len(row) for row in confidence_rows) + self.num_scheduled_draft_tokens += sum(schedule_result.lengths) + self.total_expected_accepted_tokens += schedule_result.expected_accepted_tokens + self.total_expected_tokens_per_second += ( + schedule_result.expected_tokens_per_second + ) + + for request_index, (confidences, scheduled_length) in enumerate( + zip(confidence_rows, schedule_result.lengths, strict=True) + ): + if len(confidences) > self.max_spec_tokens: + raise ValueError( + f"confidence_rows[{request_index}] has {len(confidences)} " + f"tokens, exceeding max_spec_tokens={self.max_spec_tokens}" + ) + if scheduled_length < 0 or scheduled_length > len(confidences): + raise ValueError( + f"scheduled length {scheduled_length} is invalid for " + f"confidence_rows[{request_index}]" + ) + + self.scheduled_length_histogram[scheduled_length] += 1 + survivals = cumulative_survival(confidences) + for position, confidence in enumerate(confidences): + self.confidence_sums[position] += float(confidence) + self.confidence_counts[position] += 1 + self.survival_sums[position] += survivals[position] + self.survival_counts[position] += 1 + if position < scheduled_length: + self.scheduled_counts[position] += 1 + + def snapshot(self) -> DSparkDiagnosticsSnapshot: + avg_confidence_per_pos = _safe_average_tuple( + self.confidence_sums, + self.confidence_counts, + ) + avg_survival_per_pos = _safe_average_tuple( + self.survival_sums, + self.survival_counts, + ) + scheduled_fraction_per_pos = _safe_average_tuple( + [float(value) for value in self.scheduled_counts], + self.confidence_counts, + ) + + avg_scheduled_length = ( + self.num_scheduled_draft_tokens / self.num_requests + if self.num_requests > 0 + else 0.0 + ) + draft_token_prune_rate = ( + 1.0 - (self.num_scheduled_draft_tokens / self.num_possible_draft_tokens) + if self.num_possible_draft_tokens > 0 + else 0.0 + ) + expected_acceptance_length = ( + self.total_expected_accepted_tokens / self.num_requests + if self.num_requests > 0 + else 0.0 + ) + avg_expected_tokens_per_second = ( + self.total_expected_tokens_per_second / self.num_steps + if self.num_steps > 0 + else 0.0 + ) + + return DSparkDiagnosticsSnapshot( + num_steps=self.num_steps, + num_requests=self.num_requests, + num_possible_draft_tokens=self.num_possible_draft_tokens, + num_scheduled_draft_tokens=self.num_scheduled_draft_tokens, + scheduled_length_histogram=tuple(self.scheduled_length_histogram), + avg_scheduled_length=avg_scheduled_length, + draft_token_prune_rate=draft_token_prune_rate, + expected_acceptance_length=expected_acceptance_length, + avg_expected_tokens_per_second=avg_expected_tokens_per_second, + avg_confidence_per_pos=avg_confidence_per_pos, + avg_survival_per_pos=avg_survival_per_pos, + scheduled_fraction_per_pos=scheduled_fraction_per_pos, + ) + + +@dataclass +class DSparkPosition0Diagnostics: + """Accumulate first-draft-token agreement and confidence diagnostics.""" + + num_tokens: int = 0 + num_matches: int = 0 + confidence_sum: float = 0.0 + confidence_count: int = 0 + matched_confidence_sum: float = 0.0 + matched_confidence_count: int = 0 + missed_confidence_sum: float = 0.0 + missed_confidence_count: int = 0 + num_confidence_logits_normalized: int = 0 + + def observe( + self, + matches: Sequence[bool], + confidences: Sequence[float] | None = None, + ) -> None: + if confidences is not None and len(confidences) != len(matches): + raise ValueError("confidences and matches must have the same length") + + for index, match in enumerate(matches): + self.num_tokens += 1 + matched = bool(match) + if matched: + self.num_matches += 1 + + if confidences is None: + continue + + confidence = float(confidences[index]) + if not math.isfinite(confidence): + continue + if confidence < 0.0 or confidence > 1.0: + # Diagnostic-only path: preserve the observation without + # turning an unexpected raw confidence logit into an engine + # failure. Production schedulers still validate strictly. + confidence = _sigmoid_scalar(confidence) + self.num_confidence_logits_normalized += 1 + self.confidence_sum += confidence + self.confidence_count += 1 + if matched: + self.matched_confidence_sum += confidence + self.matched_confidence_count += 1 + else: + self.missed_confidence_sum += confidence + self.missed_confidence_count += 1 + + @staticmethod + def _avg(total: float, count: int) -> float | None: + return total / count if count > 0 else None + + def snapshot(self) -> DSparkPosition0DiagnosticsSnapshot: + return DSparkPosition0DiagnosticsSnapshot( + num_tokens=self.num_tokens, + num_matches=self.num_matches, + match_rate=( + self.num_matches / self.num_tokens if self.num_tokens > 0 else 0.0 + ), + avg_confidence=self._avg(self.confidence_sum, self.confidence_count), + avg_confidence_when_matched=self._avg( + self.matched_confidence_sum, + self.matched_confidence_count, + ), + avg_confidence_when_missed=self._avg( + self.missed_confidence_sum, + self.missed_confidence_count, + ), + num_confidence_logits_normalized=self.num_confidence_logits_normalized, + ) + + +def _get_config_value(hf_config: Any, name: str, default: Any = ...): + if isinstance(hf_config, dict): + if default is ...: + return hf_config[name] + return hf_config.get(name, default) + if default is ...: + return getattr(hf_config, name) + return getattr(hf_config, name, default) + + +def has_dspark_config(hf_config: Any) -> bool: + return _get_config_value(hf_config, "dspark_block_size", None) is not None + + +def infer_dspark_weight_prefix(weight_names: Sequence[str]) -> str | None: + prefixes = set() + suffixes = ( + ".markov_head.markov_w1.weight", + ".markov_head.markov_w2.weight", + ".confidence_head.proj.weight", + ) + for name in weight_names: + for suffix in suffixes: + if name.endswith(suffix): + prefixes.add(name.removesuffix(suffix)) + + if not prefixes: + return None + if len(prefixes) > 1: + raise ValueError( + f"DSpark weights were found under multiple prefixes: {sorted(prefixes)}" + ) + return next(iter(prefixes)) + + +def infer_dspark_num_draft_layers(weight_names: Sequence[str]) -> int | None: + """Infer how many DSpark draft stages are present in a checkpoint. + + The DeepSeek-V4-Flash-DSpark release stores the DSpark draft stages under + the historical `mtp.N.*` namespace. Unlike regular MTP, the HF config still + reports `num_nextn_predict_layers=1`, so the weight map is the reliable + source for the DSpark draft depth. + """ + + layer_ids: set[int] = set() + for name in weight_names: + if not name.startswith("mtp."): + continue + parts = name.split(".", 2) + if len(parts) < 3: + continue + try: + layer_ids.add(int(parts[1])) + except ValueError: + continue + + if not layer_ids: + return None + expected = set(range(max(layer_ids) + 1)) + if layer_ids != expected: + raise ValueError( + "DSpark MTP namespace must be contiguous from mtp.0; " + f"found {sorted(layer_ids)}" + ) + return max(layer_ids) + 1 + + +def _safe_average_tuple( + values: Sequence[float], counts: Sequence[int] +) -> tuple[float, ...]: + if len(values) != len(counts): + raise ValueError("values and counts must have the same length") + return tuple( + float(value) / count if count > 0 else 0.0 + for value, count in zip(values, counts, strict=True) + ) + + +def _validate_probability(value: float, name: str) -> float: + value = float(value) + if value < 0.0 or value > 1.0: + raise ValueError(f"{name} must be in [0, 1], got {value}") + return value + + +def _sigmoid_scalar(value: float) -> float: + if value >= 0: + z = math.exp(-value) + return 1.0 / (1.0 + z) + z = math.exp(value) + return z / (1.0 + z) + + +def speculative_acceptance_confidence( + draft_probs: Sequence[float], + target_probs: Sequence[float], +) -> float: + """Return the exact one-step speculative acceptance probability. + + For normalized draft distribution q and target distribution p, the expected + acceptance probability is sum_x min(p(x), q(x)), equivalently + 1 - TV(p, q). + """ + + if len(draft_probs) != len(target_probs): + raise ValueError( + "draft_probs and target_probs must have the same vocabulary size" + ) + if len(draft_probs) == 0: + raise ValueError("probability vectors must not be empty") + + draft_total = sum(float(prob) for prob in draft_probs) + target_total = sum(float(prob) for prob in target_probs) + if draft_total <= 0.0 or target_total <= 0.0: + raise ValueError("probability vectors must have positive total mass") + + draft = [float(prob) / draft_total for prob in draft_probs] + target = [float(prob) / target_total for prob in target_probs] + if any(prob < 0.0 for prob in (*draft, *target)): + raise ValueError("probabilities must be non-negative") + + return sum( + min(draft_prob, target_prob) + for draft_prob, target_prob in zip(draft, target, strict=True) + ) + + +def cumulative_survival(confidences: Sequence[float]) -> tuple[float, ...]: + """Convert conditional per-token confidences to prefix survival rates.""" + + survivals: list[float] = [] + survival = 1.0 + for index, confidence in enumerate(confidences): + survival *= _validate_probability(confidence, f"confidences[{index}]") + survivals.append(survival) + return tuple(survivals) + + +def confidence_threshold_prefix_length( + confidences: Sequence[float], + threshold: float, +) -> int: + """Choose the longest prefix whose cumulative confidence stays above a threshold. + + The decision is non-anticipating: position n is admitted using only + confidences from positions <= n. This keeps DSpark's confidence scheduling + compatible with rejection sampling. + """ + + threshold = _validate_probability(threshold, "threshold") + admitted = 0 + survival = 1.0 + for index, confidence in enumerate(confidences): + survival *= _validate_probability(confidence, f"confidences[{index}]") + if survival < threshold: + break + admitted += 1 + return admitted + + +def score_prefix_lengths( + confidence_rows: Sequence[Sequence[float]], + lengths: Sequence[int], + *, + steps_per_second: StepCurve, +) -> DSparkScheduleResult: + """Score a fixed per-request prefix schedule against a profiled step curve.""" + + if len(confidence_rows) != len(lengths): + raise ValueError("confidence_rows and lengths must have the same length") + + expected_accepted_tokens = float(len(confidence_rows)) + batch_tokens = len(confidence_rows) + normalized_lengths: list[int] = [] + + for request_index, (confidences, length) in enumerate( + zip(confidence_rows, lengths, strict=True) + ): + length = int(length) + if length < 0 or length > len(confidences): + raise ValueError( + f"lengths[{request_index}] must be in [0, {len(confidences)}], " + f"got {length}" + ) + + expected_accepted_tokens += sum(cumulative_survival(confidences)[:length]) + batch_tokens += length + normalized_lengths.append(length) + + step_rate = float(steps_per_second(batch_tokens)) + if step_rate < 0.0: + raise ValueError("steps_per_second must return a non-negative value") + + return DSparkScheduleResult( + lengths=tuple(normalized_lengths), + expected_accepted_tokens=expected_accepted_tokens, + batch_tokens=batch_tokens, + expected_tokens_per_second=expected_accepted_tokens * step_rate, + ) + + +def hardware_aware_prefix_schedule( + confidence_rows: Sequence[Sequence[float]], + *, + steps_per_second: StepCurve, + early_stop: bool = True, +) -> DSparkScheduleResult: + """Choose DSpark verification prefix lengths for a batch. + + `confidence_rows[r][j]` is the conditional acceptance probability for + request r at draft position j. The planner greedily admits prefix tokens by + descending cumulative survival, then keeps the best point along that path + after applying the supplied profiled engine step-rate curve. + + `early_stop=True` is intended for smooth capacity curves where adding more + verification tokens cannot recover from the first throughput drop. Set it to + `False` for exhaustive traversal of the greedy prefix path when the profiled + curve has jagged capacity cliffs. + """ + + request_count = len(confidence_rows) + lengths = [0] * request_count + best_lengths = tuple(lengths) + batch_tokens = request_count + + base_step_rate = float(steps_per_second(batch_tokens)) + if base_step_rate < 0.0: + raise ValueError("steps_per_second must return a non-negative value") + + expected_accepted_tokens = float(request_count) + best_batch_tokens = batch_tokens + best_expected_accepted_tokens = expected_accepted_tokens + best_tokens_per_second = expected_accepted_tokens * base_step_rate + + candidates: list[tuple[float, int, int]] = [] + for request_index, confidences in enumerate(confidence_rows): + for position, survival in enumerate(cumulative_survival(confidences), start=1): + if survival > 0.0: + candidates.append((survival, request_index, position)) + + candidates.sort(key=lambda item: (-item[0], item[1], item[2])) + + for survival, request_index, position in candidates: + if position != lengths[request_index] + 1: + continue + + lengths[request_index] = position + batch_tokens += 1 + expected_accepted_tokens += survival + + step_rate = float(steps_per_second(batch_tokens)) + if step_rate < 0.0: + raise ValueError("steps_per_second must return a non-negative value") + tokens_per_second = expected_accepted_tokens * step_rate + + if tokens_per_second > best_tokens_per_second: + best_lengths = tuple(lengths) + best_batch_tokens = batch_tokens + best_expected_accepted_tokens = expected_accepted_tokens + best_tokens_per_second = tokens_per_second + continue + + if early_stop: + break + + return DSparkScheduleResult( + lengths=best_lengths, + expected_accepted_tokens=best_expected_accepted_tokens, + batch_tokens=best_batch_tokens, + expected_tokens_per_second=best_tokens_per_second, + ) diff --git a/mods/deepseek-v4-flash-dspark/overlay/vllm/v1/spec_decode/dspark_proposer.py b/mods/deepseek-v4-flash-dspark/overlay/vllm/v1/spec_decode/dspark_proposer.py new file mode 100644 index 00000000..dcd3208d --- /dev/null +++ b/mods/deepseek-v4-flash-dspark/overlay/vllm/v1/spec_decode/dspark_proposer.py @@ -0,0 +1,1021 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from __future__ import annotations + +import os +import time +from typing import Any + +import torch +from typing_extensions import override + +from vllm.compilation.cuda_graph import CUDAGraphWrapper +from vllm.config import CUDAGraphMode, VllmConfig +from vllm.forward_context import BatchDescriptor, set_forward_context +from vllm.logger import init_logger +from vllm.v1.attention.backend import CommonAttentionMetadata +from vllm.v1.kv_cache_interface import KVCacheConfig +from vllm.v1.sample.metadata import SamplingMetadata +from vllm.v1.spec_decode.dspark import ( + DSparkDiagnostics, + confidence_threshold_prefix_length, + hardware_aware_prefix_schedule, + make_dspark_warmup_draft_token_ids, + score_prefix_lengths, +) +from vllm.v1.spec_decode.llm_base_proposer import SpecDecodeBaseProposer + +logger = init_logger(__name__) + + +class DSparkProposer(SpecDecodeBaseProposer): + """DSpark proposer for DeepSeek V4 Flash DSpark. + + DSpark's draft model owns a small internal sliding-window cache over + target-layer features. It does not allocate draft KV blocks through vLLM's + normal speculative-decoding KV cache path. + """ + + def __init__( + self, + vllm_config: VllmConfig, + device: torch.device, + runner=None, + ) -> None: + assert vllm_config.speculative_config is not None + assert vllm_config.speculative_config.method == "dspark" + super().__init__( + vllm_config=vllm_config, + device=device, + pass_hidden_states_to_model=True, + runner=runner, + ) + hf_config = self.draft_model_config.hf_config + self.target_hidden_size = hf_config.hidden_size * len( + hf_config.dspark_target_layer_ids + ) + self.noise_token_id = int(hf_config.dspark_noise_token_id) + self._prefilled = False + self._runner = runner + self._draft_graph_runner: CUDAGraphWrapper | None = None + self._draft_graph_batch_size = 0 + self._draft_input_ids_buffer = torch.zeros( + self.max_batch_size, + dtype=torch.long, + device=device, + ) + self._draft_hidden_buffer = torch.zeros( + self.max_batch_size, + self.target_hidden_size, + dtype=self.dtype, + device=device, + ) + self._draft_positions_buffer = torch.zeros( + self.max_batch_size, + dtype=torch.long, + device=device, + ) + # Persistent, cudagraph-captured slot index for the draft read path. + # Filled before every replay; defaults to identity so warmup/capture + # and the single-stream path stay byte-for-byte unchanged. + self._draft_slot_index_buffer = torch.arange( + self.max_batch_size, + dtype=torch.long, + device=device, + ) + # Stable per-request -> persistent-KV-slot map. The draft model's only + # cross-step state (main_kv_cache) must follow the request id, not the + # batch-row position, which vLLM-v1 condenses as requests finish. + self._req_id_to_slot: dict[str, int] = {} + self._free_slots: list[int] = list(range(self.max_batch_size)) + self.diagnostics = DSparkDiagnostics( + max_spec_tokens=self.num_speculative_tokens + ) + self.confidence_threshold = self._read_confidence_threshold() + self.confidence_scheduler = self._read_confidence_scheduler( + self.confidence_threshold + ) + self._forced_draft_length = self._read_forced_draft_length( + self.num_speculative_tokens + ) + self._sps_curve = self._read_sps_curve() + self._hardware_scheduler_early_stop = ( + self._read_hardware_scheduler_early_stop() + ) + self._last_draft_lengths: list[int] | None = None + self._last_draft_probs: torch.Tensor | None = None + self._export_draft_probs = self._read_export_draft_probs() + self._collect_confidence_diagnostics = ( + self._read_collect_confidence_diagnostics() + ) + self._collect_position0_diagnostics = ( + self._read_position0_diagnostics() + ) + self._gpu_rejected_context_mask = ( + self._read_gpu_rejected_context_mask() + ) + self._stage_timing = self._read_stage_timing() + self._stage_timing_log_every = self._read_stage_timing_log_every() + self._stage_timing_count = 0 + self._stage_timing_totals_ms: dict[str, float] = {} + self._last_confidence: torch.Tensor | None = None + if self.confidence_threshold > 0.0: + logger.info( + "DSpark confidence-scheduled verification enabled with " + "threshold %.4f.", + self.confidence_threshold, + ) + if self.confidence_scheduler == "hardware": + logger.info( + "DSpark hardware-aware confidence scheduler enabled with " + "early_stop=%s and SPS curve=%s.", + self._hardware_scheduler_early_stop, + self._sps_curve or "constant", + ) + if self._forced_draft_length is not None: + logger.info( + "DSpark forced draft verification length enabled for " + "profiling: %d.", + self._forced_draft_length, + ) + if self._export_draft_probs: + logger.info( + "DSpark draft probability export enabled for quality profiling. " + "This adds a draft-logit softmax on greedy requests." + ) + if self._collect_confidence_diagnostics: + logger.info( + "DSpark confidence diagnostics enabled. This copies confidence " + "scores to CPU on every draft step." + ) + if self._collect_position0_diagnostics: + logger.info( + "DSpark position-0 diagnostics enabled. The confidence head " + "runs on every draft step; the runner logs first-token " + "target-argmax agreement." + ) + if self._gpu_rejected_context_mask: + logger.info( + "DSpark GPU rejected-context mask enabled. Rejected target " + "suffix rows are masked during draft main-KV cache update " + "without synchronizing rejection counts to CPU." + ) + if self._stage_timing: + logger.info( + "DSpark stage timing enabled. This synchronizes CUDA work and " + "is intended for diagnostics, not speed-gate benchmarks." + ) + if not self._needs_draft_logits() and not self._needs_confidence(): + logger.info( + "DSpark fast draft-output mode enabled: confidence head and " + "returned draft logits are skipped on the hot path." + ) + + @staticmethod + def _read_confidence_threshold() -> float: + raw = os.getenv("VLLM_DSPARK_CONFIDENCE_THRESHOLD", "0.0") + try: + threshold = float(raw) + except ValueError as exc: + raise ValueError( + "VLLM_DSPARK_CONFIDENCE_THRESHOLD must be a float in [0, 1], " + f"got {raw!r}" + ) from exc + if threshold < 0.0 or threshold > 1.0: + raise ValueError( + "VLLM_DSPARK_CONFIDENCE_THRESHOLD must be in [0, 1], " + f"got {threshold}" + ) + return threshold + + @staticmethod + def _read_confidence_scheduler(confidence_threshold: float) -> str: + raw = os.getenv("VLLM_DSPARK_CONFIDENCE_SCHEDULER", "auto") + scheduler = raw.strip().lower() + if scheduler in {"", "auto"}: + return "threshold" if confidence_threshold > 0.0 else "off" + if scheduler not in {"off", "threshold", "hardware"}: + raise ValueError( + "VLLM_DSPARK_CONFIDENCE_SCHEDULER must be one of " + "'off', 'threshold', 'hardware', or 'auto', " + f"got {raw!r}" + ) + return scheduler + + @staticmethod + def _read_forced_draft_length(max_draft_length: int) -> int | None: + raw = os.getenv("VLLM_DSPARK_FORCE_DRAFT_LENGTH", "").strip() + if raw == "": + return None + try: + value = int(raw) + except ValueError as exc: + raise ValueError( + "VLLM_DSPARK_FORCE_DRAFT_LENGTH must be an integer in " + f"[0, {max_draft_length}] or empty, got {raw!r}" + ) from exc + if value < 0 or value > max_draft_length: + raise ValueError( + "VLLM_DSPARK_FORCE_DRAFT_LENGTH must be in " + f"[0, {max_draft_length}], got {value}" + ) + return value + + @staticmethod + def _read_sps_curve() -> tuple[tuple[int, float], ...]: + raw = os.getenv("VLLM_DSPARK_SPS_CURVE", "").strip() + if not raw: + return () + + entries: dict[int, float] = {} + for item in raw.split(","): + item = item.strip() + if not item: + continue + try: + batch_tokens_raw, rate_raw = item.split(":", 1) + batch_tokens = int(batch_tokens_raw) + rate = float(rate_raw) + except ValueError as exc: + raise ValueError( + "VLLM_DSPARK_SPS_CURVE must be a comma-separated table " + "of ':' entries, " + f"got {raw!r}" + ) from exc + if batch_tokens <= 0: + raise ValueError( + "VLLM_DSPARK_SPS_CURVE batch-token keys must be positive, " + f"got {batch_tokens}" + ) + if rate < 0.0: + raise ValueError( + "VLLM_DSPARK_SPS_CURVE rates must be non-negative, " + f"got {rate}" + ) + entries[batch_tokens] = rate + return tuple(sorted(entries.items())) + + @staticmethod + def _read_hardware_scheduler_early_stop() -> bool: + raw = os.getenv("VLLM_DSPARK_HARDWARE_SCHEDULER_EARLY_STOP", "1") + return raw.strip().lower() in {"1", "true", "yes", "on"} + + def _steps_per_second(self, batch_tokens: int) -> float: + curve = getattr(self, "_sps_curve", ()) + if not curve: + return 1.0 + + batch_tokens = int(batch_tokens) + selected_rate = curve[0][1] + for profiled_tokens, rate in curve: + if batch_tokens < profiled_tokens: + break + selected_rate = rate + return selected_rate + + def _effective_confidence_scheduler(self) -> str: + scheduler = getattr(self, "confidence_scheduler", None) + if scheduler is not None: + return scheduler + threshold = getattr(self, "confidence_threshold", 0.0) + return "threshold" if threshold > 0.0 else "off" + + @staticmethod + def _read_export_draft_probs() -> bool: + raw = os.getenv("VLLM_DSPARK_EXPORT_DRAFT_PROBS", "0") + return raw.strip().lower() in {"1", "true", "yes", "on"} + + @staticmethod + def _read_collect_confidence_diagnostics() -> bool: + raw = os.getenv("VLLM_DSPARK_COLLECT_CONFIDENCE_DIAGNOSTICS", "0") + return raw.strip().lower() in {"1", "true", "yes", "on"} + + @staticmethod + def _read_position0_diagnostics() -> bool: + raw = os.getenv("VLLM_DSPARK_POSITION0_DIAGNOSTICS", "0") + return raw.strip().lower() in {"1", "true", "yes", "on"} + + @staticmethod + def _read_gpu_rejected_context_mask() -> bool: + raw = os.getenv("VLLM_DSPARK_GPU_REJECTED_CONTEXT_MASK", "0") + return raw.strip().lower() in {"1", "true", "yes", "on"} + + @staticmethod + def _read_stage_timing() -> bool: + raw = os.getenv("VLLM_DSPARK_STAGE_TIMING", "0") + return raw.strip().lower() in {"1", "true", "yes", "on"} + + @staticmethod + def _read_stage_timing_log_every() -> int: + raw = os.getenv("VLLM_DSPARK_STAGE_TIMING_LOG_EVERY", "20") + try: + value = int(raw) + except ValueError as exc: + raise ValueError( + "VLLM_DSPARK_STAGE_TIMING_LOG_EVERY must be an integer, " + f"got {raw!r}" + ) from exc + return max(1, value) + + def _record_stage_timing(self, name: str, elapsed_ms: float) -> None: + self._stage_timing_totals_ms[name] = ( + self._stage_timing_totals_ms.get(name, 0.0) + float(elapsed_ms) + ) + + def _timed_stage(self, name: str, fn): + if not getattr(self, "_stage_timing", False): + return fn() + if self.device.type != "cuda": + started = time.perf_counter() + result = fn() + self._record_stage_timing(name, (time.perf_counter() - started) * 1000.0) + return result + + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + start.record() + result = fn() + end.record() + end.synchronize() + self._record_stage_timing(name, start.elapsed_time(end)) + return result + + def _maybe_log_stage_timing(self) -> None: + if not getattr(self, "_stage_timing", False): + return + self._stage_timing_count += 1 + if self._stage_timing_count % self._stage_timing_log_every != 0: + return + + names = ( + "context_prepare", + "prefill_main", + "graph_prepare", + "draft", + "postprocess", + "total", + ) + parts = [] + for name in names: + total_ms = self._stage_timing_totals_ms.get(name, 0.0) + parts.append(f"{name}={total_ms / self._stage_timing_count:.3f}ms") + logger.info( + "DSpark stage timing avg over %d proposals: %s", + self._stage_timing_count, + ", ".join(parts), + ) + + @override + def initialize_attn_backend( + self, + kv_cache_config: KVCacheConfig, + kernel_block_sizes: list[int] | None = None, + ) -> None: + del kv_cache_config, kernel_block_sizes + self.block_size = 1 + + @override + def dummy_run( + self, + num_tokens: int, + use_cudagraphs: bool = True, + is_graph_capturing: bool = False, + slot_mappings: dict[str, torch.Tensor] | None = None, + ) -> None: + del is_graph_capturing, slot_mappings + batch_size = max(1, min(int(num_tokens), self.max_batch_size)) + ( + cudagraph_runtime_mode, + padded_batch_size, + num_tokens_across_dp, + batch_descriptor, + ) = self._determine_graph_batch(batch_size, use_cudagraphs=use_cudagraphs) + self._prepare_draft_buffers( + input_ids=torch.zeros(batch_size, dtype=torch.long, device=self.device), + hidden_states=torch.zeros( + batch_size, + self.target_hidden_size, + dtype=self.dtype, + device=self.device, + ), + positions=torch.arange(batch_size, dtype=torch.long, device=self.device), + padded_batch_size=padded_batch_size, + ) + with set_forward_context( + None, + self.vllm_config, + num_tokens=padded_batch_size * self.num_speculative_tokens, + num_tokens_across_dp=num_tokens_across_dp, + cudagraph_runtime_mode=cudagraph_runtime_mode, + batch_descriptor=batch_descriptor, + ): + self._run_draft_for_current_context() + + @override + def initialize_cudagraph_keys(self, cudagraph_mode: CUDAGraphMode) -> None: + if ( + not self.speculative_config.enforce_eager + and cudagraph_mode.mixed_mode() + in [CUDAGraphMode.PIECEWISE, CUDAGraphMode.FULL] + ): + dspark_cudagraph_mode = CUDAGraphMode.PIECEWISE + else: + dspark_cudagraph_mode = CUDAGraphMode.NONE + self.cudagraph_dispatcher.initialize_cudagraph_keys(dspark_cudagraph_mode) + if dspark_cudagraph_mode != CUDAGraphMode.NONE and self.device.type == "cuda": + self._draft_graph_runner = CUDAGraphWrapper( + self._run_draft_from_buffers, + self.vllm_config, + runtime_mode=CUDAGraphMode.PIECEWISE, + ) + + def _determine_graph_batch( + self, + batch_size: int, + *, + use_cudagraphs: bool = True, + ) -> tuple[CUDAGraphMode, int, torch.Tensor | None, BatchDescriptor]: + cudagraph_mode, batch_descriptor = self.cudagraph_dispatcher.dispatch( + batch_size, + valid_modes=({CUDAGraphMode.NONE} if not use_cudagraphs else None), + ) + padded_batch_size = batch_descriptor.num_tokens + num_tokens_across_dp = None + if self.vllm_config.parallel_config.data_parallel_size > 1: + from vllm.v1.worker.dp_utils import coordinate_batch_across_dp + + should_ubatch, num_tokens_across_dp, synced_cudagraph_mode = ( + coordinate_batch_across_dp( + num_tokens_unpadded=batch_size, + parallel_config=self.vllm_config.parallel_config, + allow_microbatching=False, + num_tokens_padded=padded_batch_size, + cudagraph_mode=cudagraph_mode.value, + ) + ) + assert not should_ubatch, "DBO ubatching not implemented for DSpark" + if num_tokens_across_dp is not None: + dp_rank = self.dp_rank + padded_batch_size = int(num_tokens_across_dp[dp_rank].item()) + cudagraph_mode, batch_descriptor = self.cudagraph_dispatcher.dispatch( + padded_batch_size, + valid_modes={CUDAGraphMode(synced_cudagraph_mode)}, + ) + assert batch_descriptor.num_tokens == padded_batch_size + num_tokens_across_dp[dp_rank] = padded_batch_size + return ( + cudagraph_mode, + padded_batch_size, + num_tokens_across_dp, + batch_descriptor, + ) + + def _prepare_draft_buffers( + self, + *, + input_ids: torch.Tensor, + hidden_states: torch.Tensor, + positions: torch.Tensor, + padded_batch_size: int, + slot_index: torch.Tensor | None = None, + ) -> None: + batch_size = input_ids.shape[0] + self._draft_graph_batch_size = padded_batch_size + self._draft_input_ids_buffer[:batch_size].copy_(input_ids.to(torch.long)) + self._draft_hidden_buffer[:batch_size].copy_(hidden_states.to(self.dtype)) + self._draft_positions_buffer[:batch_size].copy_(positions.to(torch.long)) + if slot_index is None: + self._draft_slot_index_buffer[:batch_size].copy_( + torch.arange(batch_size, device=self.device, dtype=torch.long) + ) + else: + self._draft_slot_index_buffer[:batch_size].copy_( + slot_index.to(torch.long) + ) + if padded_batch_size > batch_size: + pad_slice = slice(batch_size, padded_batch_size) + self._draft_input_ids_buffer[pad_slice].fill_(self.noise_token_id) + self._draft_hidden_buffer[pad_slice].zero_() + self._draft_positions_buffer[pad_slice].zero_() + # Padding rows only ever gather (read) the persistent cache, so any + # valid in-range slot is safe; 0 keeps the captured indices in range. + self._draft_slot_index_buffer[pad_slice].zero_() + + def _run_draft_from_buffers( + self, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + batch_size = self._draft_graph_batch_size + return self.model.draft_with_confidence( + self._draft_input_ids_buffer[:batch_size], + self._draft_hidden_buffer[:batch_size], + self._draft_positions_buffer[:batch_size], + return_logits=self._needs_draft_logits(), + return_confidence=self._needs_confidence(), + store_main_kv=False, + slot_index=self._draft_slot_index_buffer[:batch_size], + ) + + def _run_draft_for_current_context( + self, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + if self._draft_graph_runner is not None: + return self._draft_graph_runner() + return self._run_draft_from_buffers() + + def _batch_size(self, next_token_ids: torch.Tensor) -> int: + return int(next_token_ids.shape[0]) + + def _row_to_slot(self, req_ids: list[str]) -> list[int]: + """Map current batch rows to stable persistent-KV slots. + + Reclaims slots whose requests are no longer live, assigns a free slot + to any newly seen request, and returns the slot for each request in the + SAME ORDER as ``req_ids`` (i.e. batch-row order). Free slots are reused + lowest-first so a server that only ever runs one request at a time keeps + slot 0 -> the single-stream path stays the identity permutation. + """ + live = set(req_ids) + for stale in [r for r in self._req_id_to_slot if r not in live]: + self._free_slots.append(self._req_id_to_slot.pop(stale)) + self._free_slots.sort() + slots: list[int] = [] + for req_id in req_ids: + slot = self._req_id_to_slot.get(req_id) + if slot is None: + slot = self._free_slots.pop(0) + self._req_id_to_slot[req_id] = slot + slots.append(slot) + return slots + + def _view_by_request( + self, + values: torch.Tensor, + batch_size: int, + ) -> torch.Tensor: + if values.shape[0] % batch_size != 0: + raise ValueError( + "DSpark currently requires uniform flattened per-request inputs; " + f"got {values.shape[0]} rows for batch_size={batch_size}." + ) + seq_len = values.shape[0] // batch_size + return values.view(batch_size, seq_len, values.shape[-1]) + + def _positions_by_request( + self, + positions: torch.Tensor, + batch_size: int, + ) -> torch.Tensor: + if positions.ndim != 1: + positions = positions.reshape(-1) + if positions.shape[0] % batch_size != 0: + raise ValueError( + "DSpark currently requires uniform flattened positions; " + f"got {positions.shape[0]} rows for batch_size={batch_size}." + ) + return positions.view(batch_size, positions.shape[0] // batch_size) + + def _trim_rejected_target_context( + self, + target_hidden_states: torch.Tensor, + target_positions: torch.Tensor, + common_attn_metadata: CommonAttentionMetadata | None, + num_rejected_tokens_gpu: torch.Tensor | None, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Drop rejected verification suffixes before updating DSpark context. + + Padded speculative decoding keeps rejected tokens in the target forward + as padding and expects proposers to ignore them. DSpark stores target + hidden states in an internal context cache, so rejected suffix states + must be removed before `prefill_main()`. + """ + if num_rejected_tokens_gpu is None or common_attn_metadata is None: + return target_hidden_states, target_positions + + rejected = num_rejected_tokens_gpu.detach().cpu().tolist() + if not any(int(value) for value in rejected): + return target_hidden_states, target_positions + + query_start_loc_cpu = common_attn_metadata.query_start_loc_cpu + if query_start_loc_cpu is None: + query_start_loc_cpu = common_attn_metadata.query_start_loc.detach().cpu() + query_starts = query_start_loc_cpu.tolist() + if len(query_starts) != len(rejected) + 1: + raise ValueError( + "DSpark rejected-context trimming requires query_start_loc " + "to have batch_size + 1 entries; got " + f"{len(query_starts)} starts for {len(rejected)} requests." + ) + + hidden_chunks: list[torch.Tensor] = [] + position_chunks: list[torch.Tensor] = [] + effective_lengths: list[int] = [] + flat_positions = target_positions.reshape(-1) + for req_index, num_rejected in enumerate(rejected): + start = int(query_starts[req_index]) + end = int(query_starts[req_index + 1]) - int(num_rejected) + if end <= start: + raise ValueError( + "DSpark rejected-context trimming removed every token for " + f"request {req_index}: start={start}, end={end}." + ) + hidden_chunks.append(target_hidden_states[start:end]) + position_chunks.append(flat_positions[start:end]) + effective_lengths.append(end - start) + + if len(set(effective_lengths)) != 1: + raise ValueError( + "DSpark currently requires uniform effective per-request " + "target context lengths after rejection trimming; got " + f"{effective_lengths}." + ) + + return torch.cat(hidden_chunks, dim=0), torch.cat(position_chunks, dim=0) + + def _warmup_drafts(self, batch_size: int) -> torch.Tensor: + self._last_draft_lengths = [self.num_speculative_tokens] * batch_size + return make_dspark_warmup_draft_token_ids( + batch_size=batch_size, + num_speculative_tokens=self.num_speculative_tokens, + noise_token_id=self.noise_token_id, + device=self.device, + ) + + def _draft_lengths_from_confidence( + self, + confidence_rows: list[list[float]], + ) -> list[int]: + return list(self._schedule_from_confidence(confidence_rows).lengths) + + def _schedule_from_confidence( + self, + confidence_rows: list[list[float]], + ): + confidence_rows = [ + row[: self.num_speculative_tokens] for row in confidence_rows + ] + forced_length = getattr(self, "_forced_draft_length", None) + if forced_length is not None: + lengths = [ + min(int(forced_length), self.num_speculative_tokens, len(row)) + for row in confidence_rows + ] + return score_prefix_lengths( + confidence_rows, + lengths, + steps_per_second=self._steps_per_second, + ) + + scheduler = self._effective_confidence_scheduler() + + if scheduler == "hardware": + return hardware_aware_prefix_schedule( + confidence_rows, + steps_per_second=self._steps_per_second, + early_stop=getattr(self, "_hardware_scheduler_early_stop", True), + ) + + if scheduler == "off" or self.confidence_threshold <= 0.0: + lengths = [ + min(self.num_speculative_tokens, len(row)) + for row in confidence_rows + ] + else: + lengths = [ + confidence_threshold_prefix_length( + row, + self.confidence_threshold, + ) + for row in confidence_rows + ] + return score_prefix_lengths( + confidence_rows, + lengths, + steps_per_second=self._steps_per_second, + ) + + def _observe_confidence(self, confidence: torch.Tensor) -> list[int]: + confidence_rows = confidence.detach().float().cpu().tolist() + schedule = self._schedule_from_confidence(confidence_rows) + self.diagnostics.observe(confidence_rows, schedule) + return list(schedule.lengths) + + def _should_observe_confidence(self) -> bool: + return ( + self._effective_confidence_scheduler() != "off" + or getattr(self, "_collect_confidence_diagnostics", False) + ) + + def _needs_confidence(self) -> bool: + return ( + self._should_observe_confidence() + or getattr(self, "_collect_position0_diagnostics", False) + ) + + def _needs_draft_logits(self) -> bool: + return self._export_draft_probs + + def take_last_draft_lengths(self) -> list[int] | None: + lengths = self._last_draft_lengths + self._last_draft_lengths = None + return lengths + + def take_last_draft_probs(self) -> torch.Tensor | None: + draft_probs = self._last_draft_probs + self._last_draft_probs = None + return draft_probs + + def take_last_confidence(self) -> torch.Tensor | None: + confidence = self._last_confidence + self._last_confidence = None + return confidence + + def _maybe_store_draft_probs( + self, + draft_logits: torch.Tensor, + sampling_metadata: SamplingMetadata, + batch_size: int, + ) -> None: + self._last_draft_probs = None + if not getattr(self, "_export_draft_probs", False): + return + if draft_logits.numel() == 0: + logger.warning_once( + "DSpark draft probability export requested but draft logits " + "were not returned by the draft model." + ) + return + if not sampling_metadata.all_greedy: + logger.warning_once( + "VLLM_DSPARK_EXPORT_DRAFT_PROBS is currently limited to " + "greedy DSpark requests because DSpark non-greedy drafting " + "must sample the Markov-corrected block left-to-right." + ) + return + self._last_draft_probs = draft_logits[ + :batch_size, : self.num_speculative_tokens + ].softmax(dim=-1, dtype=torch.float32) + + @override + @torch.inference_mode() + def propose( + self, + target_token_ids: torch.Tensor, + target_positions: torch.Tensor, + target_hidden_states: torch.Tensor, + next_token_ids: torch.Tensor, + token_indices_to_sample: torch.Tensor | None, + common_attn_metadata: CommonAttentionMetadata, + sampling_metadata: SamplingMetadata, + mm_embed_inputs: tuple[list[torch.Tensor], torch.Tensor] | None = None, + num_rejected_tokens_gpu: torch.Tensor | None = None, + slot_mappings: dict[str, torch.Tensor] + | list[dict[str, torch.Tensor]] + | None = None, + req_ids: list[str] | None = None, + ) -> torch.Tensor: + del ( + target_token_ids, + token_indices_to_sample, + mm_embed_inputs, + slot_mappings, + ) + self._last_draft_probs = None + self._last_confidence = None + total_started = time.perf_counter() + batch_size = self._batch_size(next_token_ids) + + # Resolve the stable per-request KV slot for this step. The map is + # always advanced (reclaim/assign) so it stays consistent across steps, + # but when the resulting permutation is the identity we pass + # ``slot_index=None`` downstream so the single-stream / non-condensed + # path stays byte-for-byte the original in-place behaviour. + slot_list: list[int] | None = None + slot_index: torch.Tensor | None = None + buffer_slot_index: torch.Tensor | None = None + if req_ids is not None and len(req_ids) == batch_size: + slot_list = self._row_to_slot(req_ids) + if slot_list != list(range(batch_size)): + slot_index = torch.tensor( + slot_list, device=self.device, dtype=torch.long + ) + buffer_slot_index = slot_index + + def prepare_context(): + nonlocal target_hidden_states, target_positions + rejected_for_gpu_mask = None + prefill_query_start_loc: list[int] | None = None + ragged = False + query_starts: list[int] | None = None + if getattr(self, "_gpu_rejected_context_mask", False): + rejected_for_gpu_mask = num_rejected_tokens_gpu + # Detect non-uniform per-request query rows (mixed prefill + + # decode under chunked prefill). Only then do we need the ragged + # path; uniform/static batches keep the rectangular fast-path. + if common_attn_metadata is not None: + qsl_cpu = common_attn_metadata.query_start_loc_cpu + if qsl_cpu is None: + qsl_cpu = ( + common_attn_metadata.query_start_loc.detach().cpu() + ) + query_starts = qsl_cpu.tolist() + if len(query_starts) == batch_size + 1: + seg_lengths = [ + query_starts[i + 1] - query_starts[i] + for i in range(batch_size) + ] + ragged = len(set(seg_lengths)) != 1 + else: + target_hidden_states, target_positions = ( + self._trim_rejected_target_context( + target_hidden_states, + target_positions, + common_attn_metadata, + num_rejected_tokens_gpu, + ) + ) + + if ragged: + # Segment ragged-ly via query_start_loc (same pattern as + # _trim_rejected_target_context) and index each request's + # last non-rejected row to build the [B, H]/[B] draft anchors, + # with no rectangular view. The full ragged-flat hidden and + # positions flow to prefill_main, which scatters each segment + # into its slot's ring buffer. + assert query_starts is not None + flat_positions = target_positions.reshape(-1) + device = target_hidden_states.device + starts = torch.tensor( + query_starts[:batch_size], device=device, dtype=torch.long + ) + lengths = torch.tensor( + [ + query_starts[i + 1] - query_starts[i] + for i in range(batch_size) + ], + device=device, + dtype=torch.long, + ) + if rejected_for_gpu_mask is None: + rejected = torch.zeros( + batch_size, device=device, dtype=torch.long + ) + else: + rejected = rejected_for_gpu_mask.to( + device=device, dtype=torch.long, non_blocking=True + ).view(batch_size) + last_offsets = (lengths - rejected - 1).clamp(min=0) + last_offsets = torch.minimum(last_offsets, lengths - 1) + anchor_idx = starts + last_offsets + last_hidden = target_hidden_states.index_select( + 0, anchor_idx + ).contiguous() + last_positions = flat_positions.index_select( + 0, anchor_idx + ).contiguous() + return ( + target_hidden_states, + flat_positions, + last_hidden, + last_positions, + rejected_for_gpu_mask, + query_starts, + ) + + hidden_by_req = self._view_by_request(target_hidden_states, batch_size) + positions_by_req = self._positions_by_request(target_positions, batch_size) + + if rejected_for_gpu_mask is not None: + rejected = rejected_for_gpu_mask.to( + device=hidden_by_req.device, + dtype=torch.long, + non_blocking=True, + ).view(batch_size) + last_indices = ( + hidden_by_req.shape[1] - rejected - 1 + ).clamp(min=0) + last_hidden = hidden_by_req.gather( + 1, + last_indices.view(batch_size, 1, 1).expand( + -1, + -1, + hidden_by_req.shape[-1], + ), + ).squeeze(1).contiguous() + last_positions = positions_by_req.gather( + 1, + last_indices.view(batch_size, 1), + ).squeeze(1).contiguous() + else: + last_hidden = hidden_by_req[:, -1].contiguous() + last_positions = positions_by_req[:, -1].contiguous() + return ( + hidden_by_req, + positions_by_req, + last_hidden, + last_positions, + rejected_for_gpu_mask, + prefill_query_start_loc, + ) + + ( + prefill_hidden, + prefill_positions, + last_hidden, + last_positions, + rejected_for_gpu_mask, + prefill_query_start_loc, + ) = self._timed_stage("context_prepare", prepare_context) + + self._timed_stage( + "prefill_main", + lambda: self.model.prefill_main( + prefill_hidden, + prefill_positions, + num_rejected_tokens=rejected_for_gpu_mask, + slot_index=( + slot_list + if prefill_query_start_loc is not None + else slot_index + ), + query_start_loc=prefill_query_start_loc, + ), + ) + if not self._prefilled: + self._prefilled = True + return self._warmup_drafts(batch_size) + + def prepare_graph(): + ( + cudagraph_runtime_mode, + padded_batch_size, + num_tokens_across_dp, + batch_descriptor, + ) = self._determine_graph_batch(batch_size) + self._prepare_draft_buffers( + input_ids=next_token_ids, + hidden_states=last_hidden, + positions=last_positions, + padded_batch_size=padded_batch_size, + slot_index=buffer_slot_index, + ) + return ( + cudagraph_runtime_mode, + padded_batch_size, + num_tokens_across_dp, + batch_descriptor, + ) + + ( + cudagraph_runtime_mode, + padded_batch_size, + num_tokens_across_dp, + batch_descriptor, + ) = self._timed_stage("graph_prepare", prepare_graph) + with set_forward_context( + None, + self.vllm_config, + num_tokens=padded_batch_size * self.num_speculative_tokens, + num_tokens_across_dp=num_tokens_across_dp, + cudagraph_runtime_mode=cudagraph_runtime_mode, + batch_descriptor=batch_descriptor, + ): + draft_token_ids, draft_logits, confidence = ( + self._timed_stage("draft", self._run_draft_for_current_context) + ) + + def postprocess(): + self._maybe_store_draft_probs(draft_logits, sampling_metadata, batch_size) + confidence_for_batch = None + if confidence is not None and confidence.numel() > 0: + confidence_for_batch = confidence[ + :batch_size, : self.num_speculative_tokens + ] + if getattr(self, "_collect_position0_diagnostics", False): + self._last_confidence = confidence_for_batch.detach().clone() + forced_length = getattr(self, "_forced_draft_length", None) + if forced_length is not None: + self._last_draft_lengths = [int(forced_length)] * batch_size + elif ( + confidence_for_batch is not None + and self._should_observe_confidence() + ): + self._last_draft_lengths = self._observe_confidence( + confidence_for_batch + ) + else: + self._last_draft_lengths = [self.num_speculative_tokens] * batch_size + return draft_token_ids[:batch_size, : self.num_speculative_tokens].to( + torch.int32 + ) + + result = self._timed_stage("postprocess", postprocess) + self._record_stage_timing( + "total", + (time.perf_counter() - total_started) * 1000.0, + ) + self._maybe_log_stage_timing() + return result + + def get_diagnostics_snapshot(self) -> Any: + return self.diagnostics.snapshot() diff --git a/mods/deepseek-v4-flash-dspark/overlay/vllm/v1/worker/gpu_model_runner.py b/mods/deepseek-v4-flash-dspark/overlay/vllm/v1/worker/gpu_model_runner.py new file mode 100644 index 00000000..61577220 --- /dev/null +++ b/mods/deepseek-v4-flash-dspark/overlay/vllm/v1/worker/gpu_model_runner.py @@ -0,0 +1,7728 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import functools +import gc +import itertools +import os +import threading +import time +from collections import defaultdict +from collections.abc import Callable, Iterable, Iterator, Sequence +from contextlib import contextmanager +from copy import copy, deepcopy +from dataclasses import dataclass, replace +from functools import reduce +from typing import TYPE_CHECKING, Any, NamedTuple, TypeAlias, cast + +import numpy as np +import torch +import torch.distributed +import torch.nn as nn +from tqdm import tqdm + +import vllm.envs as envs +from vllm.compilation.breakable_cudagraph import ( + BreakableCUDAGraphWrapper, + is_breakable_cudagraph_enabled, +) +from vllm.compilation.counter import compilation_counter +from vllm.compilation.cuda_graph import CUDAGraphStat, CUDAGraphWrapper +from vllm.compilation.monitor import set_cudagraph_capturing_enabled +from vllm.config import ( + CompilationMode, + CUDAGraphMode, + VllmConfig, + get_layers_from_vllm_config, + set_current_vllm_config, + update_config, +) +from vllm.config.cache import CacheConfig +from vllm.distributed.ec_transfer import get_ec_transfer, has_ec_transfer +from vllm.distributed.eplb.eplb_state import EplbState +from vllm.distributed.kv_transfer import get_kv_transfer_group, has_kv_transfer_group +from vllm.distributed.kv_transfer.kv_connector.utils import copy_kv_blocks +from vllm.distributed.parallel_state import ( + get_dcp_group, + get_pp_group, + get_tp_group, + graph_capture, + is_global_first_rank, + prepare_communication_buffer_for_model, +) +from vllm.forward_context import ( + BatchDescriptor, + set_forward_context, +) +from vllm.logger import init_logger +from vllm.lora.layers import LoRAMapping, LoRAMappingType +from vllm.model_executor.layers.attention import Attention, MLAAttention +from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase +from vllm.model_executor.layers.fused_moe.routed_experts_capturer import ( + RoutedExpertsCapturer, +) +from vllm.model_executor.layers.mamba.ops.ssu_dispatch import ( + initialize_mamba_ssu_backend, +) +from vllm.model_executor.layers.rotary_embedding import ( + MRotaryEmbedding, + XDRotaryEmbedding, +) +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.models.interfaces import ( + MixtureOfExperts, + MultiModalEmbeddings, + SupportsMRoPE, + SupportsMultiModal, + SupportsXDRoPE, + is_mixture_of_experts, + supports_eagle3, + supports_mrope, + supports_multimodal_pruning, + supports_realtime, + supports_transcription, + supports_xdrope, +) +from vllm.model_executor.models.interfaces_base import ( + VllmModelForPooling, + is_pooling_model, + is_text_generation_model, +) +from vllm.model_executor.offloader import ( + create_offloader, + get_offloader, + set_offloader, +) +from vllm.multimodal import MULTIMODAL_REGISTRY +from vllm.multimodal.encoder_budget import MultiModalBudget +from vllm.multimodal.inputs import ( + BatchedTensorInputs, + MultiModalKwargsItem, + PlaceholderRange, +) +from vllm.multimodal.utils import group_and_batch_mm_kwargs +from vllm.platforms import current_platform +from vllm.pooling_params import PoolingParams +from vllm.sampling_params import SamplingType +from vllm.sequence import IntermediateTensors +from vllm.tasks import GenerationTask, PoolingTask, SupportedTask +from vllm.tracing import instrument +from vllm.utils import length_from_prompt_token_ids_or_embeds +from vllm.utils.math_utils import cdiv, round_up +from vllm.utils.mem_utils import DeviceMemoryProfiler, format_gib +from vllm.utils.nvtx_pytorch_hooks import PytHooks +from vllm.utils.platform_utils import is_pin_memory_available, num_compute_units +from vllm.utils.torch_utils import ( + get_dtype_size, + is_quantized_kv_cache, + kv_cache_dtype_str_to_dtype, +) +from vllm.v1.attention.backend import ( + AttentionBackend, + AttentionCGSupport, + AttentionMetadata, + AttentionMetadataBuilder, + AttentionType, + CommonAttentionMetadata, +) +from vllm.v1.attention.backends.gdn_attn import GDNAttentionMetadataBuilder +from vllm.v1.attention.backends.mamba2_attn import Mamba2AttentionMetadataBuilder +from vllm.v1.attention.backends.utils import ( + NULL_BLOCK_ID, + create_fast_prefill_custom_backend, + get_dcp_local_seq_lens, + reorder_batch_to_split_decodes_and_prefills, +) +from vllm.v1.core.sched.output import NewRequestData +from vllm.v1.cudagraph_dispatcher import CudagraphDispatcher +from vllm.v1.kv_cache_interface import ( + AttentionSpec, + ChunkedLocalAttentionSpec, + CrossAttentionSpec, + EncoderOnlyAttentionSpec, + FullAttentionSpec, + KVCacheConfig, + KVCacheGroupSpec, + KVCacheSpec, + MambaSpec, + SlidingWindowSpec, + UniformTypeKVCacheSpecs, +) +from vllm.v1.outputs import ( + EMPTY_MODEL_RUNNER_OUTPUT, + AsyncModelRunnerOutput, + DraftTokenIds, + ECConnectorOutput, + KVConnectorOutput, + LogprobsLists, + LogprobsTensors, + ModelRunnerOutput, + PoolerOutput, + RoutedExpertsLists, + RoutedExpertsTensors, + SamplerOutput, + make_empty_encoder_model_runner_output, +) +from vllm.v1.pool.metadata import PoolingMetadata, PoolingStates +from vllm.v1.sample.logits_processor import LogitsProcessors, build_logitsprocs +from vllm.v1.sample.logits_processor.interface import LogitsProcessor +from vllm.v1.sample.metadata import SamplingMetadata +from vllm.v1.sample.rejection_sampler import RejectionSampler +from vllm.v1.sample.sampler import Sampler +from vllm.v1.spec_decode.custom_class_proposer import create_custom_proposer +from vllm.v1.spec_decode.dflash import DFlashProposer +from vllm.v1.spec_decode.draft_model import DraftModelProposer +from vllm.v1.spec_decode.dspark import DSparkPosition0Diagnostics +from vllm.v1.spec_decode.dspark_proposer import DSparkProposer +from vllm.v1.spec_decode.eagle import EagleProposer +from vllm.v1.spec_decode.extract_hidden_states import ExtractHiddenStatesProposer +from vllm.v1.spec_decode.gemma4 import Gemma4Proposer +from vllm.v1.spec_decode.medusa import MedusaProposer +from vllm.v1.spec_decode.metadata import SpecDecodeMetadata +from vllm.v1.spec_decode.ngram_proposer_gpu import ( + NgramProposerGPU, + copy_num_valid_draft_tokens, + update_ngram_gpu_tensors_incremental, + update_scheduler_for_invalid_drafts, +) +from vllm.v1.spec_decode.suffix_decoding import SuffixDecodingProposer +from vllm.v1.spec_decode.utils import update_num_computed_tokens_for_batch_change +from vllm.v1.structured_output.utils import apply_grammar_bitmask +from vllm.v1.utils import CpuGpuBuffer, record_function_or_nullcontext +from vllm.v1.worker import mamba_utils +from vllm.v1.worker.cp_utils import ( + check_attention_cp_compatibility, + get_total_cp_world_size, +) +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.pool.late_interaction_runner import LateInteractionRunner +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 +from vllm.v1.worker.lora_model_runner_mixin import LoRAModelRunnerMixin +from vllm.v1.worker.ubatch_utils import ( + UBatchSlices, + check_ubatch_thresholds, + maybe_create_ubatch_slices, + split_attn_metadata, +) +from vllm.v1.worker.utils import is_residual_scattered_for_sp +from vllm.v1.worker.workspace import lock_workspace + +from .utils import ( + AttentionGroup, + KVBlockZeroer, + add_kv_sharing_layers_to_kv_cache_groups, + bind_kv_cache, + prepare_kernel_block_sizes, + sanity_check_mm_encoder_outputs, +) + +if TYPE_CHECKING: + from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput + from vllm.v1.spec_decode.ngram_proposer import NgramProposer + from vllm.v1.worker.encoder_cudagraph import EncoderCudaGraphManager + +logger = init_logger(__name__) + +AttnMetadataDict: TypeAlias = dict[str, AttentionMetadata] +# list when ubatching is enabled +PerLayerAttnMetadata: TypeAlias = list[AttnMetadataDict] | AttnMetadataDict + + +# Wrapper for ModelRunnerOutput to support overlapped execution. +class AsyncGPUModelRunnerOutput(AsyncModelRunnerOutput): + def __init__( + self, + model_runner_output: ModelRunnerOutput, + sampled_token_ids: torch.Tensor, + logprobs_tensors: LogprobsTensors | None, + invalid_req_indices: list[int], + async_output_copy_stream: torch.cuda.Stream, + vocab_size: int, + routed_experts: RoutedExpertsTensors | None = None, + ): + self._model_runner_output = model_runner_output + self._invalid_req_indices = invalid_req_indices + + # Event on the copy stream so we can synchronize the non-blocking copy. + self.async_copy_ready_event = torch.Event() + + # Keep a reference to the device tensor to avoid it being + # deallocated until we finish copying it to the host. + self._sampled_token_ids = sampled_token_ids + self.vocab_size = vocab_size + self._logprobs_tensors = logprobs_tensors + self._routed_experts = routed_experts + + # Initiate the copy on a separate stream, but do not synchronize it. + default_stream = torch.cuda.current_stream() + with torch.cuda.stream(async_output_copy_stream): + async_output_copy_stream.wait_stream(default_stream) + self.sampled_token_ids_cpu = self._sampled_token_ids.to( + "cpu", non_blocking=True + ) + self._logprobs_tensors_cpu = ( + self._logprobs_tensors.to_cpu_nonblocking() + if self._logprobs_tensors + else None + ) + self._routed_experts_cpu = ( + self._routed_experts.to_cpu_nonblocking() + if self._routed_experts is not None + else None + ) + self.async_copy_ready_event.record() + + def get_output(self) -> ModelRunnerOutput: + """Copy the device tensors to the host and return a ModelRunnerOutput. + + This function blocks until the copy is finished. + """ + max_gen_len = self.sampled_token_ids_cpu.shape[-1] + self.async_copy_ready_event.synchronize() + + # Release the device tensors once the copy has completed. + del self._logprobs_tensors + del self._sampled_token_ids + if max_gen_len == 1: + valid_sampled_token_ids = self.sampled_token_ids_cpu.tolist() + for i in self._invalid_req_indices: + valid_sampled_token_ids[i].clear() + logprobs_lists = None + if self._logprobs_tensors_cpu is not None: + logprobs_lists = self._logprobs_tensors_cpu.tolists() + else: + valid_sampled_token_ids, logprobs_lists = RejectionSampler.parse_output( + self.sampled_token_ids_cpu, + self.vocab_size, + self._invalid_req_indices, + logprobs_tensors=self._logprobs_tensors_cpu, + ) + + output = self._model_runner_output + output.sampled_token_ids = valid_sampled_token_ids + output.logprobs = logprobs_lists + + if self._routed_experts_cpu is not None: + output.routed_experts = self._routed_experts_cpu.tolists() + del self._routed_experts + + return output + + +def _copy_pooler_output_to_cpu( + raw_pooler_output: PoolerOutput, finished_mask: list[bool] +) -> list[torch.Tensor | None]: + num_reqs = len(finished_mask) + + if isinstance(raw_pooler_output, torch.Tensor): + if raw_pooler_output.shape[0] != num_reqs: + raise ValueError( + "Pooler output batch size does not match finished mask size: " + f"{raw_pooler_output.shape[0]} != {num_reqs}." + ) + + num_finished = sum(finished_mask) + if num_finished == 0: + return [None] * num_reqs + if num_finished == num_reqs: + return list(raw_pooler_output.to("cpu", non_blocking=True)) + + # partial finished + finished_indices = [i for i, include in enumerate(finished_mask) if include] + index_tensor = torch.tensor( + finished_indices, device=raw_pooler_output.device, dtype=torch.long + ) + finished_outputs = raw_pooler_output.index_select(0, index_tensor).to( + "cpu", non_blocking=True + ) + partial_pooler_output: list[torch.Tensor | None] = [None] * num_reqs + for i, out in zip(finished_indices, finished_outputs): + partial_pooler_output[i] = out + return partial_pooler_output + + assert isinstance(raw_pooler_output, list) + if len(raw_pooler_output) != num_reqs: + raise ValueError( + "Pooler output batch size does not match finished mask size: " + f"{len(raw_pooler_output)} != {num_reqs}." + ) + + pooler_output: list[torch.Tensor | None] = [None] * num_reqs + for i, (out, include) in enumerate(zip(raw_pooler_output, finished_mask)): + if include and out is not None: + pooler_output[i] = out.to("cpu", non_blocking=True) + return pooler_output + + +def _format_optional_float(value: float | None) -> str: + return "nan" if value is None else f"{value:.3f}" + + +class AsyncGPUPoolingModelRunnerOutput(AsyncModelRunnerOutput): + def __init__( + self, + model_runner_output: ModelRunnerOutput, + raw_pooler_output: PoolerOutput, + finished_mask: list[bool], + async_output_copy_stream: torch.cuda.Stream, + ): + self._model_runner_output = model_runner_output + + # Event on the copy stream so we can synchronize the non-blocking copy. + self.async_copy_ready_event = torch.Event() + + # Keep a reference to the device tensors to avoid them being + # deallocated until we finish copying it to the host. + self._raw_pooler_output = raw_pooler_output + + # Initiate the copy on a separate stream, but do not synchronize it. + default_stream = torch.cuda.current_stream() + with torch.cuda.stream(async_output_copy_stream): + async_output_copy_stream.wait_stream(default_stream) + self._model_runner_output.pooler_output = _copy_pooler_output_to_cpu( + raw_pooler_output=self._raw_pooler_output, + finished_mask=finished_mask, + ) + self.async_copy_ready_event.record() + + def get_output(self) -> ModelRunnerOutput: + """Copy the device tensors to the host and return a ModelRunnerOutput. + This function blocks until the copy is finished. + """ + self.async_copy_ready_event.synchronize() + + # Release the device tensors once the copy has completed. + del self._raw_pooler_output + return self._model_runner_output + + +class ExecuteModelState(NamedTuple): + """Ephemeral cached state transferred between execute_model() and + sample_tokens(), after execute_model() returns None.""" + + scheduler_output: "SchedulerOutput" + logits: torch.Tensor + spec_decode_metadata: SpecDecodeMetadata | None + spec_decode_common_attn_metadata: CommonAttentionMetadata | None + hidden_states: torch.Tensor + sample_hidden_states: torch.Tensor + aux_hidden_states: list[torch.Tensor] | None + ec_connector_output: ECConnectorOutput | None + cudagraph_stats: CUDAGraphStat | None + slot_mappings: dict[str, torch.Tensor] | list[dict[str, torch.Tensor]] | None + + +class GPUModelRunner( + LoRAModelRunnerMixin, KVConnectorModelRunnerMixin, ECConnectorModelRunnerMixin +): + def __init__( + self, + vllm_config: VllmConfig, + device: torch.device, + ): + self.vllm_config = vllm_config + self.model_config = vllm_config.model_config + self.cache_config = vllm_config.cache_config + self.offload_config = vllm_config.offload_config + self.compilation_config = vllm_config.compilation_config + self.lora_config = vllm_config.lora_config + self.load_config = vllm_config.load_config + self.parallel_config = vllm_config.parallel_config + self.scheduler_config = vllm_config.scheduler_config + self.speculative_config = vllm_config.speculative_config + self.observability_config = vllm_config.observability_config + self._dspark_iter_timing = os.getenv("VLLM_DSPARK_ITER_TIMING", "0").lower() in { + "1", + "true", + "yes", + "on", + } + try: + self._dspark_iter_timing_log_every = max( + 1, int(os.getenv("VLLM_DSPARK_ITER_TIMING_LOG_EVERY", "20")) + ) + except ValueError: + self._dspark_iter_timing_log_every = 20 + self._dspark_iter_timing_count = 0 + self._dspark_iter_timing_totals_ms: defaultdict[str, float] = defaultdict(float) + self._dspark_iter_timing_started = 0.0 + + model_config = self.model_config + cache_config = self.cache_config + scheduler_config = self.scheduler_config + parallel_config = self.parallel_config + self.device = device + self.pin_memory = is_pin_memory_available() + self.dtype = self.model_config.dtype + if self._dspark_iter_timing: + logger.info( + "DSpark iteration timing enabled. CUDA is synchronized around " + "measured stages; use for diagnostics only." + ) + + self.kv_cache_dtype = kv_cache_dtype_str_to_dtype( + cache_config.cache_dtype, self.model_config + ) + + self.is_pooling_model = model_config.runner_type == "pooling" + self.enable_prompt_embeds = model_config.enable_prompt_embeds + self.is_multimodal_raw_input_only_model = ( + model_config.is_multimodal_raw_input_only_model + ) + # These will be overridden in load_model() + self.is_multimodal_pruning_enabled = False + self.requires_sequential_video_encoding = False + # Set to True after init_routed_experts_capturer() completes. + # Prevents routed experts code from running during profiling/dummy run. + self.routed_experts_initialized = False + self.max_model_len = model_config.max_model_len + + # Always set to false after the first forward pass + self.calculate_kv_scales = self.cache_config.calculate_kv_scales + self.dcp_world_size = self.parallel_config.decode_context_parallel_size + self.dcp_rank = 0 if self.dcp_world_size <= 1 else get_dcp_group().rank_in_group + self.max_num_tokens = scheduler_config.max_num_batched_tokens + self.max_num_reqs = scheduler_config.max_num_seqs + + # Broadcast PP output for external_launcher (torchrun) + # to make sure we are synced across pp ranks + # TODO: Support overlapping micro-batches + # https://github.com/vllm-project/vllm/issues/18019 + self.broadcast_pp_output = ( + self.parallel_config.distributed_executor_backend == "external_launcher" + and len(get_pp_group().ranks) > 1 + ) + + # 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 + + # Multi-modal data support + self.mm_registry = MULTIMODAL_REGISTRY + self.uses_mrope = model_config.uses_mrope + self.uses_xdrope_dim = model_config.uses_xdrope_dim + self.supports_mm_inputs = self.mm_registry.supports_multimodal_inputs( + model_config + ) + + if self.model_config.is_encoder_decoder: + # Maximum length of the encoder input, only for encoder-decoder + # models. + self.max_encoder_len = scheduler_config.max_num_encoder_input_tokens + else: + self.max_encoder_len = 0 + + # Async scheduling + self.use_async_scheduling = self.scheduler_config.async_scheduling + + # Sampler + self.sampler = Sampler(logprobs_mode=self.model_config.logprobs_mode) + + self.eplb_state: EplbState | None = None + self._moe_model: MixtureOfExperts | None = None + # NOTE(yongji): flag to temporarily disable EPLB during scaling up/down + self.eep_eplb_suppressed = False + """ + State of the expert parallelism load balancer. + + Will be lazily initialized when the model is loaded. + """ + + # Lazy initializations + # self.model: nn.Module # Set after load_model + # Initialize in initialize_kv_cache + self.kv_caches: list[torch.Tensor] = [] + # Initialize in initialize_kv_cache_tensors + self.cross_layers_kv_cache: torch.Tensor | None = None + self.cross_layers_attn_backend: type[AttentionBackend] | None = None + # indexes: [kv_cache_group_id][attn_group] + self.attn_groups: list[list[AttentionGroup]] = [] + # self.kv_cache_config: KVCacheConfig + + # mm_hash -> encoder_output + self.encoder_cache: dict[str, torch.Tensor] = {} + self.late_interaction_runner = LateInteractionRunner() + + # Encoder CUDA graph manager (initialized after model load if enabled) + self.encoder_cudagraph_manager: EncoderCudaGraphManager | None = None + + self.use_aux_hidden_state_outputs = False + # Set up speculative decoding. + # NOTE(Jiayi): currently we put the entire draft model on + # the last PP rank. This is not ideal if there are many + # layers in the draft model. + if self.speculative_config and get_pp_group().is_last_rank: + self.drafter: ( + NgramProposer # noqa: F823 + | NgramProposerGPU + | SuffixDecodingProposer + | EagleProposer + | DFlashProposer + | DSparkProposer + | DraftModelProposer + | MedusaProposer + | ExtractHiddenStatesProposer + | Gemma4Proposer + ) + if self.speculative_config.method == "custom_class": + self.drafter = create_custom_proposer( # type: ignore[assignment] + self.vllm_config + ) + elif self.speculative_config.method == "ngram": + from vllm.v1.spec_decode.ngram_proposer import NgramProposer + + self.drafter = NgramProposer(self.vllm_config) + elif self.speculative_config.uses_draft_model(): + self.drafter = DraftModelProposer( + vllm_config=self.vllm_config, + device=self.device, + runner=self, + ) + elif self.speculative_config.use_ngram_gpu(): + self.drafter = NgramProposerGPU(self.vllm_config, self.device, self) + self.num_tokens_no_spec_gpu = torch.zeros( + self.max_num_reqs, dtype=torch.int32, device=device + ) + self.token_ids_gpu_tensor = torch.zeros( + self.max_num_reqs, + self.max_model_len, + dtype=torch.int32, + device=device, + ) + self._ngram_pinned_idx_buf = torch.zeros( + self.max_num_reqs, dtype=torch.long, pin_memory=True + ) + self._ngram_pinned_val_buf = torch.zeros( + self.max_num_reqs, dtype=torch.int32, pin_memory=True + ) + elif self.speculative_config.use_gemma4_mtp(): + self.drafter = Gemma4Proposer(self.vllm_config, self.device, self) + elif self.speculative_config.use_dflash(): + self.drafter = DFlashProposer(self.vllm_config, self.device, self) + self.use_aux_hidden_state_outputs = True + elif self.speculative_config.use_dspark(): + self.drafter = DSparkProposer(self.vllm_config, self.device, self) + elif self.speculative_config.method == "suffix": + self.drafter = SuffixDecodingProposer(self.vllm_config) + elif self.speculative_config.use_eagle(): + self.drafter = EagleProposer(self.vllm_config, self.device, self) + if self.speculative_config.method == "eagle3": + self.use_aux_hidden_state_outputs = ( + self.drafter.eagle3_use_aux_hidden_state + ) + elif self.speculative_config.method == "medusa": + self.drafter = MedusaProposer( + vllm_config=self.vllm_config, device=self.device + ) + elif self.speculative_config.method == "extract_hidden_states": + self.drafter = ExtractHiddenStatesProposer( + vllm_config=self.vllm_config, device=self.device + ) + self.use_aux_hidden_state_outputs = True + else: + raise ValueError( + "Unknown speculative decoding method: " + f"{self.speculative_config.method}" + ) + self.rejection_sampler = RejectionSampler( + self.sampler, self.speculative_config, self.device + ) + + self.num_spec_tokens = 0 + self.valid_sampled_token_count_gpu: torch.Tensor | None = None + if self.speculative_config: + self.num_spec_tokens = self.speculative_config.num_speculative_tokens + draft_config = self.speculative_config.draft_model_config + if draft_config is not None and draft_config.max_model_len is not None: + self.effective_drafter_max_model_len = draft_config.max_model_len + else: + self.effective_drafter_max_model_len = self.max_model_len + self.use_async_spec_decode = ( + self.use_async_scheduling and self.num_spec_tokens > 0 + ) + + # Request states. + self.requests: dict[str, CachedRequestState] = {} + # NOTE(rob): num_prompt_logprobs only includes reqs + # that are currently in the prefill phase. + self.num_prompt_logprobs: dict[str, int] = {} + + # Input Batch + # NOTE(Chen): Ideally, we should initialize the input batch inside + # `initialize_kv_cache` based on the kv cache config. However, as in + # https://github.com/vllm-project/vllm/pull/18298, due to some unknown + # reasons, we have to initialize the input batch before `load_model`, + # quantization + weight offloading will fail otherwise. As a temporary + # solution, we initialize the input batch here, and re-initialize it + # in `initialize_kv_cache` if the block_sizes here is different from + # the block_sizes in the kv cache config. + logits_processors = model_config.logits_processors + custom_logitsprocs: Sequence[str | type[LogitsProcessor]] = ( + tuple(logits_processors) if logits_processors is not None else () + ) + placeholder_block_size = ( + self.cache_config.block_size or CacheConfig.DEFAULT_BLOCK_SIZE + ) + self._init_block_sizes = [placeholder_block_size] + self._init_kernel_block_sizes = [placeholder_block_size] + self.input_batch = InputBatch( + max_num_reqs=self.max_num_reqs, + # We need to use the encoder length for encoder-decoder + # because of KV cache for cross-attention. + max_model_len=max(self.max_model_len, self.max_encoder_len), + max_num_batched_tokens=self.max_num_tokens, + device=self.device, + pin_memory=self.pin_memory, + vocab_size=self.model_config.get_vocab_size(), + block_sizes=[placeholder_block_size], + kernel_block_sizes=[placeholder_block_size], + num_spec_tokens=self.num_spec_tokens, + logitsprocs=build_logitsprocs( + self.vllm_config, + self.device, + self.pin_memory, + self.is_pooling_model, + custom_logitsprocs, + ), + # We currently don't know whether a particular custom logits processor + # uses output token ids so we set this conservatively. + # ThinkingTokenBudgetLogitsProcessor also needs output token ids to + # correctly track think start/end token sequences in async scheduling. + logitsprocs_need_output_token_ids=bool(custom_logitsprocs) + or self.vllm_config.reasoning_config is not None, + is_pooling_model=self.is_pooling_model, + cp_kv_cache_interleave_size=self.parallel_config.cp_kv_cache_interleave_size, + reasoning_config=self.vllm_config.reasoning_config, + ) + + # Separate cuda stream for overlapping transfer of sampled token ids from + # GPU to CPU when async scheduling is enabled. + self.async_output_copy_stream: torch.cuda.Stream | None = None + # cuda event to synchronize use of reused CPU tensors between steps + # when async scheduling is enabled. + self.prepare_inputs_event: torch.Event | None = None + if self.use_async_scheduling: + self.async_output_copy_stream = torch.cuda.Stream() + self.prepare_inputs_event = torch.Event() + + # self.cudagraph_batch_sizes sorts in ascending order. + if ( + self.compilation_config.cudagraph_capture_sizes + and self.compilation_config.cudagraph_mode != CUDAGraphMode.NONE + ): + self.cudagraph_batch_sizes = sorted( + self.compilation_config.cudagraph_capture_sizes + ) + 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() + + # Persistent buffers for CUDA graphs. + self.input_ids = self._make_buffer(self.max_num_tokens, dtype=torch.int32) + self.positions = torch.zeros( + self.max_num_tokens, dtype=torch.int64, device=self.device + ) + self.query_start_loc = self._make_buffer( + self.max_num_reqs + 1, dtype=torch.int32 + ) + self.seq_lens = torch.zeros( + self.max_num_reqs, dtype=torch.int32, device=self.device + ) + self.optimistic_seq_lens_cpu = torch.zeros( + self.max_num_reqs, dtype=torch.int32, pin_memory=self.pin_memory + ) + self.num_computed_tokens = torch.zeros( + self.max_num_reqs, dtype=torch.int32, device=self.device + ) + self.prev_num_draft_tokens = self._make_buffer( + self.max_num_reqs, dtype=torch.int32 + ) + self.req_indices = self._make_buffer(self.max_num_tokens, dtype=torch.int64) + # Maps current batch position -> previous batch position (-1 for new reqs) + self.prev_positions = self._make_buffer(self.max_num_reqs, dtype=torch.int64) + self.num_scheduled_tokens = self._make_buffer( + self.max_num_reqs, dtype=torch.int32 + ) + + self.encoder_seq_lens = self._make_buffer(self.max_num_reqs, dtype=torch.int32) + if self.dcp_world_size > 1: + self.dcp_local_seq_lens = self._make_buffer( + self.max_num_reqs, dtype=torch.int32 + ) + # Because inputs_embeds may be bfloat16 and we don't need a numpy + # version of this tensor, avoid a RuntimeError by not creating a + # numpy buffer. + self.inputs_embeds = self._make_buffer( + self.max_num_tokens, self.inputs_embeds_size, dtype=self.dtype, numpy=False + ) + self.is_token_ids = self._make_buffer(self.max_num_tokens, dtype=torch.bool) + self.discard_request_mask = self._make_buffer( + self.max_num_reqs, dtype=torch.bool + ) + self.num_decode_draft_tokens = self._make_buffer( + self.max_num_reqs, dtype=torch.int32 + ) + self.num_accepted_tokens = self._make_buffer( + self.max_num_reqs, dtype=torch.int32 + ) + + # Only relevant for models using M-RoPE (e.g, Qwen2-VL) + if self.uses_mrope: + # NOTE: `mrope_positions` is implemented with one additional dummy + # position on purpose to make it non-contiguous so that it can work + # with torch compile. + # See detailed explanation in https://github.com/vllm-project/vllm/pull/12128#discussion_r1926431923 + + # NOTE: When M-RoPE is enabled, position ids are 3D regardless of + # the modality of inputs. For text-only inputs, each dimension has + # identical position IDs, making M-RoPE functionally equivalent to + # 1D-RoPE. + # See page 5 of https://arxiv.org/abs/2409.12191 + self.mrope_positions = self._make_buffer( + (3, self.max_num_tokens + 1), dtype=torch.int64 + ) + + # Only relevant for models using XD-RoPE (e.g, HunYuan-VL) + if self.uses_xdrope_dim > 0: + # Similar to mrope but use assigned dimension number for RoPE, 4 as default. + self.xdrope_positions = self._make_buffer( + (self.uses_xdrope_dim, self.max_num_tokens + 1), dtype=torch.int64 + ) + + # None in the first PP rank. The rest are set after load_model. + self.intermediate_tensors: IntermediateTensors | None = None + + # OPTIMIZATION: Cache the arange tensors rather than creating them + # every step. Keep in int64 to avoid overflow with long context. + # - arange_np: immutable [0, 1, 2, ...] used as source for batched computation + # - query_pos: CpuGpuBuffer for the computed batched arange result + arange_size = max(self.max_num_reqs + 1, self.max_num_tokens) + self.arange_np = np.arange(arange_size, dtype=np.int64) + self.query_pos = self._make_buffer(arange_size, dtype=torch.int64) + self._arange_scratch = np.empty(arange_size, dtype=np.int64) + + # Layer pairings for cross-layer KV sharing. + # If an Attention layer `layer_name` is in the keys of this dict, it + # means this layer will perform attention using the keys and values + # from the KV cache of `shared_kv_cache_layers[layer_name]`. + self.shared_kv_cache_layers: dict[str, str] = {} + self.kv_sharing_fast_prefill_eligible_layers: set[str] = set() + + self.kv_sharing_fast_prefill_logits_indices = None + if self.cache_config.kv_sharing_fast_prefill: + self.kv_sharing_fast_prefill_logits_indices = torch.zeros( + self.max_num_tokens, dtype=torch.int32, device=self.device + ) + + self.uniform_decode_query_len = 1 + self.num_spec_tokens + + # Cudagraph dispatcher for runtime cudagraph dispatching. + self.cudagraph_dispatcher = CudagraphDispatcher(self.vllm_config) + + self.mm_budget = ( + MultiModalBudget(self.vllm_config, self.mm_registry) + if self.supports_mm_inputs + else None + ) + + self.reorder_batch_threshold: int | None = None + + # Attention layers that are only in the KVCacheConfig of the runner + # (e.g., KV sharing, encoder-only attention), but not in the + # KVCacheConfig of the scheduler. + self.runner_only_attn_layers: set[str] = set() + + # Cached outputs. + self._draft_token_ids: list[list[int]] | torch.Tensor | None = None + self._draft_token_lengths_cpu: list[int] | None = None + self._draft_token_length_req_ids: list[str] | None = None + self._draft_probs: torch.Tensor | None = None + self._draft_prob_req_ids: list[str] | None = None + self._draft_confidence: torch.Tensor | None = None + self._draft_confidence_req_ids: list[str] | None = None + self._dspark_position0_diagnostics = ( + DSparkPosition0Diagnostics() + if ( + self.speculative_config is not None + and self.speculative_config.use_dspark() + and envs.VLLM_DSPARK_POSITION0_DIAGNOSTICS + ) + else None + ) + self._dspark_position0_log_next = 1 + if self._dspark_position0_diagnostics is not None: + logger.info( + "DSpark position-0 quality diagnostics enabled. This copies " + "one scalar agreement and optional confidence per request to " + "CPU on speculative decode steps." + ) + # N-gram GPU path: async D2H buffer/event for per-request valid draft counts. + self._num_valid_draft_tokens: torch.Tensor | None = None + self._num_valid_draft_tokens_cpu: torch.Tensor | None = None + self._num_valid_draft_tokens_event: torch.cuda.Event | None = None + self._num_valid_draft_tokens_copy_stream: torch.cuda.Stream | None = None + if ( + self.speculative_config is not None + and self.speculative_config.use_ngram_gpu() + ): + self._num_valid_draft_tokens_cpu = torch.empty( + self.max_num_reqs, dtype=torch.int32, pin_memory=self.pin_memory + ) + self._num_valid_draft_tokens_event = torch.cuda.Event() + self._num_valid_draft_tokens_copy_stream = torch.cuda.Stream() + + self._draft_token_req_ids: list[str] | None = None + self.transfer_event = torch.Event() + self.sampled_token_ids_pinned_cpu = torch.empty( + (self.max_num_reqs, 1), + dtype=torch.int64, + device="cpu", + pin_memory=self.pin_memory, + ) + + # Pre-allocated tensor for copying valid sampled token counts to CPU, + # with dedicated stream for overlapping and event for coordination. + self.valid_sampled_token_count_event: torch.Event | None = None + self.valid_sampled_token_count_copy_stream: torch.cuda.Stream | None = None + # We also copy the drafted tokens to the CPU asynchronously, + # in case we need them for structured outputs. + self.draft_token_ids_event: torch.Event | None = None + self.draft_token_ids_copy_stream: torch.cuda.Stream | None = None + self.valid_sampled_token_count_cpu: torch.Tensor | None = None + self.draft_token_ids_cpu: torch.Tensor | None = None + self.num_accepted_tokens_event: torch.Event | None = None + if self.num_spec_tokens: + self.draft_token_ids_event = torch.Event() + self.num_accepted_tokens_event = torch.Event() + self.draft_token_ids_copy_stream = torch.cuda.Stream() + self.draft_token_ids_cpu = torch.empty( + (self.max_num_reqs, self.num_spec_tokens), + dtype=torch.int64, + device="cpu", + pin_memory=self.pin_memory, + ) + if self.use_async_scheduling: + self.valid_sampled_token_count_event = torch.Event() + self.valid_sampled_token_count_copy_stream = torch.cuda.Stream() + self.valid_sampled_token_count_cpu = torch.empty( + self.max_num_reqs, + dtype=torch.int32, + device="cpu", + pin_memory=self.pin_memory, + ) + + # Model weight offloader + # Make sure this is called before any get_offloader call + set_offloader(create_offloader(self.offload_config)) + + # Ephemeral state transferred between execute_model() and sample_tokens(). + self.execute_model_state: ExecuteModelState | None = None + self.kv_connector_output: KVConnectorOutput | None = None + self.mamba_state_idx: dict[str, int] = {} + self._mamba_bufs: mamba_utils.MambaBuffers | None = None + self.mamba_prev_last_scheduled_idx: CpuGpuBuffer | None = None + if self.cache_config.mamba_cache_mode == "all" and self.num_spec_tokens > 0: + self.mamba_prev_last_scheduled_idx = self._make_buffer( + self.max_num_reqs, dtype=torch.int32 + ) + self.layerwise_nvtx_hooks_registered = False + + def update_max_model_len(self, max_model_len: int) -> None: + self.max_model_len = max_model_len + if self.speculative_config: + draft_config = self.speculative_config.draft_model_config + if draft_config is None or draft_config.max_model_len is None: + self.effective_drafter_max_model_len = self.max_model_len + + def reset_mm_cache(self) -> None: + """ + Clear the multi-modal cache that was used during profiling, + but no longer needed during inference. + """ + if self.mm_budget: + self.mm_budget.reset_cache() + self.late_interaction_runner.clear() + + def reset_encoder_cache(self) -> None: + """Clear the GPU-side encoder cache storing vision embeddings. + + This should be called when model weights are updated to ensure + stale embeddings computed with old weights are not reused. + """ + self.encoder_cache.clear() + self.late_interaction_runner.clear() + + def post_kv_cache_wake_up(self) -> None: + self.init_fp8_kv_scales() + + @torch.inference_mode() + def init_fp8_kv_scales(self) -> None: + """ + Re-initialize the KV cache and FP8 scales after waking from sleep. + 1. Zero out the KV cache tensors to remove garbage data from re-allocation. + 2. Reset Attention layer scaling factors (_k_scale, _v_scale) to 1.0. + If these are left at 0.0 (default after wake_up), all KV cache values + become effectively zero, causing gibberish output. + """ + if not is_quantized_kv_cache(self.cache_config.cache_dtype): + return + + kv_caches = getattr(self, "kv_caches", []) + for cache_tensor in kv_caches: + if cache_tensor is not None: + cache_tensor.zero_() + + k_attr_names = ("_k_scale", "k_scale") + v_attr_names = ("_v_scale", "v_scale") + + attn_layers = self.compilation_config.static_forward_context + for name, module in attn_layers.items(): + if isinstance(module, (Attention, MLAAttention)): + # TODO: Generally, scale is 1.0 if user uses on-the-fly fp8 + # kvcache quant. However, to get better accuracy, compression + # frameworks like llm-compressors allow users to tune the + # scale. We may need to restore the specific calibrated scales + # here in the future. + k_scale_val, v_scale_val = 1.0, 1.0 + + # Processing K Scale + for attr in k_attr_names: + if hasattr(module, attr): + param = getattr(module, attr) + if isinstance(param, torch.Tensor): + param.fill_(k_scale_val) + + # Processing V Scale + for attr in v_attr_names: + if hasattr(module, attr): + param = getattr(module, attr) + if isinstance(param, torch.Tensor): + param.fill_(v_scale_val) + + def _get_positions(self, num_tokens: Any): + if isinstance(num_tokens, int): + if self.uses_mrope: + return self.mrope_positions.gpu[:, :num_tokens] + if self.uses_xdrope_dim > 0: + return self.xdrope_positions.gpu[:, :num_tokens] + return self.positions[:num_tokens] + else: + if self.uses_mrope: + return self.mrope_positions.gpu[:, num_tokens] + if self.uses_xdrope_dim > 0: + return self.xdrope_positions.gpu[:, num_tokens] + return self.positions[num_tokens] + + def _make_buffer( + self, *size: int | torch.SymInt, dtype: torch.dtype, numpy: bool = True + ) -> CpuGpuBuffer: + return CpuGpuBuffer( + *size, + dtype=dtype, + device=self.device, + pin_memory=self.pin_memory, + with_numpy=numpy, + ) + + def _get_mamba_bufs(self) -> mamba_utils.MambaBuffers: + # Only reachable on the ``mamba_cache_mode == "align"`` path. + # The postprocess sub-object is additionally gated on spec + # decode + hybrid model. + assert self.cache_config.mamba_cache_mode == "align" + if self._mamba_bufs is None: + self._mamba_bufs = mamba_utils.MambaBuffers.create( + max_num_reqs=self.max_num_reqs, + kv_cache_config=self.kv_cache_config, + copy_funcs=self.model.get_mamba_state_copy_func(), + make_buffer=self._make_buffer, + device=self.device, + with_postprocess_align=( + self.speculative_config is not None and self.model_config.is_hybrid + ), + ) + return self._mamba_bufs + + def _init_model_kwargs(self): + model_kwargs = dict[str, Any]() + + if not self.is_pooling_model: + return model_kwargs + + num_reqs = self.input_batch.num_reqs + pooling_params = self.input_batch.get_pooling_params() + + token_type_id_requests = dict[int, Any]() + for i, param in enumerate(pooling_params): + if ( + param.extra_kwargs is not None + and (token_types := param.extra_kwargs.get("compressed_token_type_ids")) + is not None + ): + token_type_id_requests[i] = token_types + + if len(token_type_id_requests) == 0: + return model_kwargs + + # Build ids on CPU using the CPU-resident upper bound for seq_lens; + # `torch.arange(seq_lens[i])` with a GPU scalar would force a sync. + seq_lens_cpu = self.optimistic_seq_lens_cpu[:num_reqs].tolist() + token_type_ids = [] + + for i in range(num_reqs): + seq_len_i = seq_lens_cpu[i] + pos = token_type_id_requests.get(i, seq_len_i) + ids = (torch.arange(seq_len_i) >= pos).int() + token_type_ids.append(ids) + + token_type_ids_cpu = torch.empty( + sum(seq_lens_cpu), dtype=torch.int32, pin_memory=self.pin_memory + ) + torch.cat(token_type_ids, out=token_type_ids_cpu) + model_kwargs["token_type_ids"] = token_type_ids_cpu.to( + device=self.device, non_blocking=True + ) + return model_kwargs + + def _may_reorder_batch(self, scheduler_output: "SchedulerOutput") -> None: + """ + Update the order of requests in the batch based on the attention + backend's needs. For example, some attention backends (namely MLA) may + want to separate requests based on if the attention computation will be + compute-bound or memory-bound. + + Args: + scheduler_output: The scheduler output. + """ + # Attention free models have zero kv_cache_groups, however models + # like Mamba are also attention free but use the kv_cache for + # keeping its internal state. This is why we check the number + # of kv_cache groups instead of solely checking + # for self.model_config.is_attention_free. + if len(self.kv_cache_config.kv_cache_groups) == 0: + return + + if self.reorder_batch_threshold is not None: + reorder_batch_to_split_decodes_and_prefills( + self.input_batch, + scheduler_output, + decode_threshold=self.reorder_batch_threshold, + ) + + def _init_kv_zero_meta(self) -> None: + """One-time precomputation for _zero_block_ids. + + Delegates to KVBlockZeroer.init_meta with the runner's state. + Called from gpu_worker.py outside the CuMem pool context. + """ + self._kv_block_zeroer = KVBlockZeroer(self.device, self.pin_memory) + self._kv_block_zeroer.init_meta( + attn_groups_iter=self._kv_cache_spec_attn_group_iterator(), + kernel_block_sizes=self._kernel_block_sizes, + cache_dtype=self.cache_config.cache_dtype, + runner_only_attn_layers=self.runner_only_attn_layers, + static_forward_context=(self.compilation_config.static_forward_context), + ) + + def _zero_block_ids(self, block_ids: list[int]) -> None: + """Zero the KV cache memory for the given block IDs.""" + if hasattr(self, "_kv_block_zeroer"): + self._kv_block_zeroer.zero_block_ids(block_ids) + + # Note: used for model runner override. + def _init_device_properties(self) -> None: + """Initialize attributes from torch.cuda.get_device_properties""" + + self.num_sms = num_compute_units(self.device.index) + + # Note: used for model runner override. + def _sync_device(self) -> None: + torch.accelerator.synchronize() + + def _get_or_create_async_output_copy_stream(self) -> torch.cuda.Stream: + stream = self.async_output_copy_stream + if stream is None: + stream = torch.cuda.Stream() + self.async_output_copy_stream = stream + return stream + + def _update_states(self, scheduler_output: "SchedulerOutput") -> Callable | None: + """Update the cached states and the persistent batch with the scheduler + output. + + The updated states are used by the `_prepare_inputs` function to create + the input GPU tensors for the model. + + The SamplingMetadata is updated and copied to the GPU if there is a + new/resumed/paused/finished request in the batch. + """ + # Remove finished requests from the cached states. + for req_id in scheduler_output.finished_req_ids: + self.requests.pop(req_id, None) + self.num_prompt_logprobs.pop(req_id, None) + self.late_interaction_runner.on_requests_finished( + scheduler_output.finished_req_ids + ) + # Remove the finished requests from the persistent batch. + # NOTE(woosuk): There could be an edge case where finished_req_ids and + # scheduled_req_ids overlap. This happens when a request is aborted and + # then resubmitted with the same ID. In this case, we treat them as two + # distinct requests - clearing the cached states for the first request + # and handling the second as a new request. + for req_id in scheduler_output.finished_req_ids: + self.input_batch.remove_request(req_id) + + # Zero GPU memory for freshly allocated cache blocks to prevent + # stale NaN/data from corrupting attention or SSM computation. + if scheduler_output.new_block_ids_to_zero: + self._zero_block_ids(scheduler_output.new_block_ids_to_zero) + + # Free the cached encoder outputs. + for mm_hash in scheduler_output.free_encoder_mm_hashes: + self.encoder_cache.pop(mm_hash, None) + + # Remove the unscheduled requests from the persistent batch. + # NOTE(woosuk): The unscheduled requests are either preempted requests + # or running requests that are not scheduled in this step. We remove + # them from the persistent batch but keep their cached states since + # they will be scheduled again sometime in the future. + scheduled_req_ids = scheduler_output.num_scheduled_tokens.keys() + cached_req_ids = self.input_batch.req_id_to_index.keys() + resumed_req_ids = scheduler_output.scheduled_cached_reqs.resumed_req_ids + # NOTE(zhuohan): cached_req_ids and resumed_req_ids are usually disjoint, + # so `(scheduled_req_ids - resumed_req_ids) == scheduled_req_ids` holds + # apart from the forced-preemption case in reset_prefix_cache. And in + # that case we include the resumed_req_ids in the unscheduled set so + # that they get cleared from the persistent batch before being re-scheduled + # in the normal resumed request path. + unscheduled_req_ids = cached_req_ids - (scheduled_req_ids - resumed_req_ids) + # NOTE(woosuk): The persistent batch optimization assumes that + # consecutive batches contain mostly the same requests. If batches + # have low request overlap (e.g., alternating between two distinct + # sets of requests), this optimization becomes very inefficient. + for req_id in unscheduled_req_ids: + self.input_batch.remove_request(req_id) + + is_ngram_gpu = ( + self.speculative_config is not None + and self.speculative_config.use_ngram_gpu() + ) + if is_ngram_gpu: + ngram_gpu_new_reqs: list[CachedRequestState] = [] + + reqs_to_add: list[CachedRequestState] = [] + deferred_spec_decode_corrections = [] + + # Add new requests to the cached states. + for new_req_data in scheduler_output.scheduled_new_reqs: + req_id = new_req_data.req_id + if req_id in self.requests: + # For streaming case only. + req_state = self._update_streaming_request(req_id, new_req_data) + reqs_to_add.append(req_state) + continue + + sampling_params = new_req_data.sampling_params + pooling_params = new_req_data.pooling_params + + if ( + sampling_params + and sampling_params.sampling_type == SamplingType.RANDOM_SEED + ): + generator = torch.Generator(device=self.device) + generator.manual_seed(sampling_params.seed) + else: + generator = None + + if self.is_pooling_model: + assert pooling_params is not None + task = pooling_params.task + assert task is not None, "You did not set `task` in the API" + + model = cast(VllmModelForPooling, self.get_model()) + to_update = model.pooler.get_pooling_updates(task) + to_update.apply(pooling_params) + + req_state = CachedRequestState( + req_id=req_id, + prompt_token_ids=new_req_data.prompt_token_ids, + prompt_embeds=new_req_data.prompt_embeds, + prompt_is_token_ids=new_req_data.prompt_is_token_ids, + mm_features=new_req_data.mm_features, + sampling_params=sampling_params, + pooling_params=pooling_params, + generator=generator, + block_ids=new_req_data.block_ids, + num_computed_tokens=new_req_data.num_computed_tokens, + output_token_ids=[], + lora_request=new_req_data.lora_request, + ) + self.requests[req_id] = req_state + self.late_interaction_runner.register_request(req_id, pooling_params) + + if sampling_params and sampling_params.prompt_logprobs is not None: + self.num_prompt_logprobs[req_id] = ( + self.input_batch.vocab_size + if sampling_params.prompt_logprobs == -1 + else sampling_params.prompt_logprobs + ) + + # Only relevant for models using M-RoPE (e.g, Qwen2-VL) + if self.uses_mrope: + self._init_mrope_positions(req_state) + + # Only relevant for models using XD-RoPE (e.g, HunYuan-VL) + if self.uses_xdrope_dim > 0: + self._init_xdrope_positions(req_state) + + reqs_to_add.append(req_state) + # Track new requests for ngram_gpu full tensor copy + if is_ngram_gpu: + ngram_gpu_new_reqs.append(req_state) + + # Update the states of the running/resumed requests. + is_last_rank = get_pp_group().is_last_rank + req_data = scheduler_output.scheduled_cached_reqs + scheduled_spec_tokens = scheduler_output.scheduled_spec_decode_tokens + + # Save scheduler-allocated spec lengths before trimming so + # prev_num_draft_len keeps the optimistic count for rejection correction. + original_num_spec_per_req: dict[str, int] = {} + if ( + self.speculative_config is not None + and self.speculative_config.use_ngram_gpu() + ): + for req_id, toks in scheduled_spec_tokens.items(): + original_num_spec_per_req[req_id] = len(toks) + update_scheduler_for_invalid_drafts( + self._num_valid_draft_tokens_event, + self._num_valid_draft_tokens_cpu, + scheduler_output, + self.input_batch.req_id_to_index, + ) + if self.use_async_spec_decode: + self.prev_num_draft_tokens.np.fill(0) + + for i, req_id in enumerate(req_data.req_ids): + req_state = self.requests[req_id] + num_computed_tokens = req_data.num_computed_tokens[i] + new_block_ids = req_data.new_block_ids[i] + resumed_from_preemption = req_id in req_data.resumed_req_ids + num_output_tokens = req_data.num_output_tokens[i] + req_index = self.input_batch.req_id_to_index.get(req_id) + + if req_state.prev_num_draft_len and self.use_async_scheduling: + # prev_num_draft_len is used in async scheduling mode with + # spec decode. it indicates if need to update num_computed_tokens + # of the request. for example: + # first step: num_computed_tokens = 0, spec_tokens = [], + # prev_num_draft_len = 0. + # second step: num_computed_tokens = 100(prompt length), + # spec_tokens = [a,b], prev_num_draft_len = 0. + # third step: num_computed_tokens = 100 + 2, spec_tokens = [c,d], + # prev_num_draft_len = 2. + # num_computed_tokens in first step and second step doesn't contain + # the spec tokens length, but in third step it contains the + # spec tokens length. we only need to update num_computed_tokens + # when prev_num_draft_len > 0. + if req_index is None: + req_state.prev_num_draft_len = 0 + else: + # Optimistically assume all accepted; queue up a correction + # to be called after the model forward to preserve async + # scheduling. Corrected on GPU in _prepare_inputs. + optimistic_num_accepted = req_state.prev_num_draft_len + req_state.output_token_ids.extend([-1] * optimistic_num_accepted) + + deferred_spec_decode_corrections.append( + (req_id, optimistic_num_accepted, req_state) + ) + + prev_req_index = ( + self.input_batch.prev_req_id_to_index.get(req_id) + if self.input_batch.prev_req_id_to_index + else None + ) + if prev_req_index is not None: + self.prev_num_draft_tokens.np[prev_req_index] = ( + optimistic_num_accepted + ) + + if is_ngram_gpu and optimistic_num_accepted > 0: + self.input_batch.num_tokens_no_spec[req_index] += ( + optimistic_num_accepted + ) + + # Update the cached states. + req_state.num_computed_tokens = num_computed_tokens + + if not is_last_rank: + if not req_data.new_token_ids: + # Async scheduled PP: Sampled tokens propagated via GPU broadcast. + new_token_ids: list[int] = [] + else: + # Non-async scheduling with PP: The scheduler sends + # sampled token ids back because there's no direct communication + # between the first-stage worker and the last-stage worker. + new_token_ids = req_data.new_token_ids[i] + # Add the sampled token(s) from the previous step (if any). + # This doesn't include "unverified" tokens like spec tokens. + num_new_tokens = ( + num_computed_tokens + len(new_token_ids) - req_state.num_tokens + ) + if num_new_tokens == 1: + # Avoid slicing list in most common case. + req_state.output_token_ids.append(new_token_ids[-1]) + elif num_new_tokens > 0: + req_state.output_token_ids.extend( + new_token_ids[-num_new_tokens:] + ) + elif num_output_tokens < len(req_state.output_token_ids): + # Some output tokens were discarded due to a sync-KV-load + # failure, or output_token_ids was inflated by the optimistic + # extend above (async spec decode). Align the cached state. + del req_state.output_token_ids[num_output_tokens:] + if req_index is not None: + end_idx = ( + self.input_batch.num_prompt_tokens[req_index] + + num_output_tokens + ) + self.input_batch.num_tokens_no_spec[req_index] = end_idx + + # Update the block IDs. + if not resumed_from_preemption: + if new_block_ids is not None: + # Append the new blocks to the existing block IDs. + for block_ids, new_ids in zip(req_state.block_ids, new_block_ids): + block_ids.extend(new_ids) + else: + assert req_index is None + assert new_block_ids is not None + # The request is resumed from preemption. + # Replace the existing block IDs with the new ones. + req_state.block_ids = new_block_ids + + if req_index is None: + # The request is not in the persistent batch. + # The request was either preempted and resumed later, or was not + # scheduled in the previous step and needs to be added again. + + if self.use_async_scheduling and num_output_tokens > 0: + # We must recover the output token ids for resumed requests in the + # async scheduling case, so that correct input_ids are obtained. + resumed_token_ids = req_data.all_token_ids[req_id] + req_state.output_token_ids = resumed_token_ids[-num_output_tokens:] + + reqs_to_add.append(req_state) + # Track resumed requests for ngram_gpu full tensor copy + if is_ngram_gpu: + ngram_gpu_new_reqs.append(req_state) + continue + + # Update the persistent batch. + self.input_batch.num_computed_tokens_cpu[req_index] = num_computed_tokens + if new_block_ids is not None: + self.input_batch.block_table.append_row(new_block_ids, req_index) + + # For the last rank, we don't need to update the token_ids_cpu + # because the sampled tokens are already cached. + if not is_last_rank: + start_token_index = self.input_batch.num_tokens_no_spec[req_index] + # For chunked prefill, num_computed_tokens may less + # than num_tokens_no_spec. + # Async scheduled PP: no new_token_ids, advance num_tokens_no_spec + # according to num_computed_tokens. + end_token_index = max( + start_token_index, + num_computed_tokens + len(new_token_ids), + ) + if end_token_index > start_token_index: + if new_token_ids: + # Add new_token_ids to token_ids_cpu. + num_new_tokens = end_token_index - start_token_index + tokens_to_append = new_token_ids[-num_new_tokens:] + self.input_batch.token_ids_cpu[ + req_index, start_token_index:end_token_index + ] = tokens_to_append + self.input_batch.is_token_ids[ + req_index, start_token_index:end_token_index + ] = True + self.input_batch.num_tokens_no_spec[req_index] = end_token_index + + # Add spec_token_ids to token_ids_cpu. + self.input_batch.update_req_spec_token_ids(req_state, scheduled_spec_tokens) + # Restore scheduler-side draft count after ngram trimming. + if original_num_spec_per_req: + orig = original_num_spec_per_req.get(req_id, 0) + if orig != req_state.prev_num_draft_len: + req_state.prev_num_draft_len = orig + + # Add the new or resumed requests to the persistent batch. + # The smaller empty indices are filled first. + for request in reqs_to_add: + self.input_batch.add_request(request) + self.input_batch.update_req_spec_token_ids(request, scheduled_spec_tokens) + + # Condense the batched states if there are gaps left by removed requests + self.input_batch.condense() + # Allow attention backend to reorder the batch, potentially + self._may_reorder_batch(scheduler_output) + # Refresh batch metadata with any pending updates. + self.input_batch.refresh_metadata() + + # Incrementally update ngram_gpu tensors after batch is stable + if is_ngram_gpu: + update_ngram_gpu_tensors_incremental( + self.input_batch, + self.token_ids_gpu_tensor, + self.num_tokens_no_spec_gpu, + ngram_gpu_new_reqs, + self.device, + _pinned_idx_buf=self._ngram_pinned_idx_buf, + _pinned_val_buf=self._ngram_pinned_val_buf, + ) + + if deferred_spec_decode_corrections: + + def correct_spec_decode_token_counts(): + valid_sampled_token_count = self._get_valid_sampled_token_count() + if not valid_sampled_token_count: + return + prev_req_id_to_index = self.input_batch.prev_req_id_to_index + if not prev_req_id_to_index: + return + for ( + req_id, + optimistic_num_accepted, + req_state, + ) in deferred_spec_decode_corrections: + prev_req_index = prev_req_id_to_index.get(req_id) + if prev_req_index is None: + continue + num_accepted = valid_sampled_token_count[prev_req_index] - 1 + correction = optimistic_num_accepted - num_accepted + req_state.num_computed_tokens -= correction + cur_req_index = self.input_batch.req_id_to_index.get(req_id) + if cur_req_index is None: + continue + self.input_batch.num_computed_tokens_cpu[cur_req_index] -= ( + correction + ) + if is_ngram_gpu and correction > 0: + self.input_batch.num_tokens_no_spec[cur_req_index] -= correction + self.num_tokens_no_spec_gpu[cur_req_index] -= correction + + return correct_spec_decode_token_counts + else: + return None + + def _update_states_after_model_execute( + self, output_token_ids: torch.Tensor, scheduler_output: "SchedulerOutput" + ) -> None: + """Update the cached states after model execution. + + This is used for MTP/EAGLE for hybrid models, as in linear attention, + only the last token's state is kept. In MTP/EAGLE, for draft tokens + the state are kept util we decide how many tokens are accepted for + each sequence, and a shifting is done during the next iteration + based on the number of accepted tokens. + """ + if not self.speculative_config or not self.model_config.is_hybrid: + return + + # Count the number of accepted tokens for each sequence. + # Valid tokens are contiguous from position 0, so counting non-(-1) + # tokens gives us the first -1 position (i.e., number of accepted). + num_reqs = output_token_ids.size(0) + self.num_accepted_tokens.gpu[:num_reqs] = (output_token_ids != -1).sum(dim=1) + + if self.cache_config.mamba_cache_mode == "align": + # Fused GPU postprocess: state copies + per-request accepted-token + # update without CPU-GPU sync. The metadata + # (num_scheduled_tokens, num_draft_tokens, num_computed_tokens) is + # pre-staged to GPU buffers in _prepare_inputs. + mamba_utils.postprocess_mamba_align_gpu( + bufs=self._get_mamba_bufs(), + num_reqs=num_reqs, + num_accepted_tokens_gpu=self.num_accepted_tokens.gpu, + num_accepted_tokens_cpu_tensor=( + self.input_batch.num_accepted_tokens_cpu_tensor + ), + input_batch=self.input_batch, + kv_cache_config=self.kv_cache_config, + forward_context=self.compilation_config.static_forward_context, + mamba_state_copy_funcs=self.model.get_mamba_state_copy_func(), + ) + + assert self.num_accepted_tokens_event is not None + self.num_accepted_tokens_event.record() + else: + self.input_batch.num_accepted_tokens_cpu_tensor[:num_reqs].copy_( + self.num_accepted_tokens.gpu[:num_reqs], non_blocking=True + ) + assert self.num_accepted_tokens_event is not None + self.num_accepted_tokens_event.record() + + if self.cache_config.mamba_cache_mode == "all": + mamba_utils.postprocess_mamba_all( + scheduler_output, + self.kv_cache_config, + self.input_batch, + self.requests, + self.mamba_state_idx, + self.num_spec_tokens, + num_reqs, + ) + + def _update_streaming_request( + self, req_id: str, new_req_data: NewRequestData + ) -> CachedRequestState: + """Updates streaming session request from `scheduled_new_reqs`. + + Removes the request from InputBatch (if present), updates the cached + state, and prepares it for re-addition to the batch. + + NOTE: prompt_token_ids includes intermediate output tokens - tokens + previously generated but now are input context (part of the prompt). + """ + self.input_batch.remove_request(req_id) + req_state = self.requests[req_id] + + req_state.prompt_token_ids = new_req_data.prompt_token_ids + req_state.mm_features = new_req_data.mm_features + req_state.prompt_embeds = new_req_data.prompt_embeds + req_state.sampling_params = new_req_data.sampling_params + req_state.pooling_params = new_req_data.pooling_params + self.late_interaction_runner.register_request(req_id, req_state.pooling_params) + req_state.block_ids = new_req_data.block_ids + req_state.num_computed_tokens = new_req_data.num_computed_tokens + req_state.num_prompt_tokens = length_from_prompt_token_ids_or_embeds( + req_state.prompt_token_ids, req_state.prompt_embeds + ) + + # Clear `output_token_ids` as previous output tokens are now part of + # `prompt_token_ids`. + req_state.output_token_ids.clear() + + if self.uses_mrope: + self._init_mrope_positions(req_state) + + return req_state + + def _init_mrope_positions(self, req_state: CachedRequestState): + model = self.get_model() + assert supports_mrope(model), "M-RoPE support is not implemented." + assert req_state.prompt_token_ids is not None, ( + "M-RoPE requires prompt_token_ids to be available." + ) + mrope_model = cast(SupportsMRoPE, model) + + # `prompt_embeds` is a passthrough modality (no grid_thw), models' + # M-RoPE code assumes per-feature grid info, so filter it out. The + # prompt_embeds positions are treated as text positions for M-RoPE. + mrope_features = [ + f for f in req_state.mm_features if f.modality != "prompt_embeds" + ] + req_state.mrope_positions, req_state.mrope_position_delta = ( + mrope_model.get_mrope_input_positions( + req_state.prompt_token_ids, + mrope_features, + ) + ) + + def _init_xdrope_positions(self, req_state: CachedRequestState): + model = self.get_model() + xdrope_model = cast(SupportsXDRoPE, model) + assert req_state.prompt_token_ids is not None, ( + "XD-RoPE requires prompt_token_ids to be available." + ) + assert supports_xdrope(model), "XD-RoPE support is not implemented." + + req_state.xdrope_positions = xdrope_model.get_xdrope_input_positions( + req_state.prompt_token_ids, + req_state.mm_features, + ) + + def _extract_mm_kwargs( + self, + scheduler_output: "SchedulerOutput", + ) -> BatchedTensorInputs: + if not scheduler_output or not self.is_multimodal_raw_input_only_model: + return {} + + mm_kwargs = list[tuple[str, MultiModalKwargsItem]]() + for req in scheduler_output.scheduled_new_reqs: + for feature in req.mm_features: + if feature.data is not None: + mm_kwargs.append((feature.modality, feature.data)) + + # Input all modalities at once + mm_kwargs_combined: BatchedTensorInputs = {} + for _, _, mm_kwargs_batch in group_and_batch_mm_kwargs( + mm_kwargs, + device=self.device, + pin_memory=self.pin_memory, + ): + mm_kwargs_combined.update(mm_kwargs_batch) + + return mm_kwargs_combined + + def _dummy_mm_kwargs(self, num_seqs: int) -> BatchedTensorInputs: + if not self.is_multimodal_raw_input_only_model: + return {} + + mm_budget = self.mm_budget + assert mm_budget is not None + + if not mm_budget.mm_max_toks_per_item: + return {} # No tower modalities (embed-only mode) + + dummy_modality = mm_budget.get_modality_with_max_tokens() + return self._get_mm_dummy_batch(dummy_modality, num_seqs) + + def _get_cumsum_and_arange( + self, + num_tokens: np.ndarray, + arange_out: np.ndarray, + cumsum_dtype: np.dtype | None = None, + ) -> np.ndarray: + """Get the cumulative sum and batched arange of the given array. + E.g., [2, 5, 3] -> [2, 7, 10], arange written to + arange_out[:10] as [0, 1, 0, 1, 2, 3, 4, 0, 1, 2]. + Equivalent to but faster than: + np.concatenate([np.arange(n) for n in num_tokens]) + """ + # Step 1. [2, 5, 3] -> [2, 7, 10] + cu_num_tokens = np.cumsum(num_tokens, dtype=cumsum_dtype) + total_num_tokens = cu_num_tokens[-1] + # Step 2. [2, 7, 10] -> [0, 0, 2, 2, 2, 2, 2, 7, 7, 7] + cumsums_offsets = np.repeat(cu_num_tokens - num_tokens, num_tokens) + # Step 3. [0, 1, 0, 1, 2, 3, 4, 0, 1, 2] + np.subtract( + self.arange_np[:total_num_tokens], + cumsums_offsets, + out=arange_out[:total_num_tokens], + ) + + return cu_num_tokens + + def _compute_prev_positions(self, num_reqs: int) -> None: + """Build prev_positions mapping: current pos -> previous pos (-1 if new). + + Populates self.prev_positions.np[:num_reqs] with the mapping. + """ + prev_req_id_to_index = self.input_batch.prev_req_id_to_index + prev_positions = self.prev_positions.np[:num_reqs] + + if not prev_req_id_to_index: + prev_positions.fill(-1) + return + + for i, req_id in enumerate(self.input_batch.req_ids[:num_reqs]): + prev_positions[i] = prev_req_id_to_index.get(req_id, -1) + + def _prepare_input_ids( + self, + scheduler_output: "SchedulerOutput", + num_reqs: int, + total_num_scheduled_tokens: int, + cu_num_tokens: np.ndarray, + ) -> None: + """Prepare the input IDs for the current batch. + + Carefully handles the `prev_sampled_token_ids` which can be cached + from the previous engine iteration, in which case those tokens on the + GPU need to be copied into the corresponding slots into input_ids. + + Uses self.prev_positions[:num_reqs] which maps current pos -> prev pos + (-1 for new requests). + """ + + if self.input_batch.prev_sampled_token_ids is None: + # Normal scheduling case + self.input_ids.copy_to_gpu(total_num_scheduled_tokens) + if self.enable_prompt_embeds: + self.inputs_embeds.copy_to_gpu(total_num_scheduled_tokens) + self.is_token_ids.copy_to_gpu(total_num_scheduled_tokens) + return + + # Async scheduling case, where some decode requests from the previous + # iteration won't have entries in input_ids_cpu and need to be copied + # on the GPU from prev_sampled_token_ids. + prev_positions = self.prev_positions.np[:num_reqs] + scheduled_spec_tokens = scheduler_output.scheduled_spec_decode_tokens + sample_flattened_indices: list[int] = [] + spec_flattened_indices: list[int] = [] + prev_draft_token_indices: list[int] = [] + prev_indices: list[int] = [] + common_indices_match = True + max_flattened_index = -1 + total_num_spec_tokens = 0 + + for cur_index in range(num_reqs): + prev_index = prev_positions[cur_index] + if prev_index < 0: + continue + prev_indices.append(prev_index) + req_id = self.input_batch.req_ids[cur_index] + # We need to compute the flattened input_ids index of the + # last token in each common request. + draft_len = len(scheduled_spec_tokens.get(req_id, ())) + total_num_spec_tokens += draft_len + flattened_index = cu_num_tokens[cur_index].item() - 1 + # example: cu_num_tokens = [2, 5, 8], draft_tokens = [1, 2, 2] + # sample_flattened_indices = [0, 2, 5] + # spec_flattened_indices = [1, 3, 4, 6, 7] + sample_flattened_indices.append(flattened_index - draft_len) + spec_flattened_indices.extend( + range(flattened_index - draft_len + 1, flattened_index + 1) + ) + start = prev_index * self.num_spec_tokens + # prev_draft_token_indices is used to find which draft_tokens_id + # should be copied to input_ids + # example: prev draft_tokens_id [[1,2], [3,4], [5, 6]] + # flatten draft_tokens_id [1,2,3,4,5,6] + # draft_len of each request [1, 2, 1] + # then prev_draft_token_indices is [0, 2, 3, 4] + prev_draft_token_indices.extend(range(start, start + draft_len)) + common_indices_match &= prev_index == flattened_index + max_flattened_index = max(max_flattened_index, flattened_index) + + num_common_tokens = len(sample_flattened_indices) + total_without_spec = total_num_scheduled_tokens - total_num_spec_tokens + if num_common_tokens < total_without_spec: + # If not all requests are decodes from the last iteration, + # we need to copy the input_ids_cpu to the GPU first. + self.input_ids.copy_to_gpu(total_num_scheduled_tokens) + if self.enable_prompt_embeds: + self.inputs_embeds.copy_to_gpu(total_num_scheduled_tokens) + self.is_token_ids.copy_to_gpu(total_num_scheduled_tokens) + if num_common_tokens == 0: + # No requests in common with the previous iteration + # So input_ids.cpu will have all the input ids. + return + if common_indices_match and max_flattened_index == (num_common_tokens - 1): + # Common-case optimization: the batch is unchanged + # and no reordering happened. + # The indices are both the same permutation of 0..N-1 so + # we can copy directly using a single slice. + self.input_ids.gpu[:num_common_tokens].copy_( + self.input_batch.prev_sampled_token_ids[:num_common_tokens, 0], + non_blocking=True, + ) + return + # Upload the index tensors asynchronously so the scatter can be non-blocking. + sampled_tokens_index_tensor = torch.tensor( + sample_flattened_indices, dtype=torch.int64, pin_memory=self.pin_memory + ).to(self.device, non_blocking=True) + prev_common_req_indices_tensor = torch.tensor( + prev_indices, dtype=torch.int64, pin_memory=self.pin_memory + ).to(self.device, non_blocking=True) + self.input_ids.gpu.scatter_( + dim=0, + index=sampled_tokens_index_tensor, + src=self.input_batch.prev_sampled_token_ids[ + prev_common_req_indices_tensor, 0 + ], + ) + + # Scatter the draft tokens after the sampled tokens are scattered. + if self._draft_token_ids is None or not spec_flattened_indices: + return + + assert isinstance(self._draft_token_ids, torch.Tensor) + draft_tokens_index_tensor = torch.tensor( + spec_flattened_indices, dtype=torch.int64, pin_memory=self.pin_memory + ).to(self.device, non_blocking=True) + prev_draft_token_indices_tensor = torch.tensor( + prev_draft_token_indices, dtype=torch.int64, pin_memory=self.pin_memory + ).to(self.device, non_blocking=True) + + # because input_ids dtype is torch.int32, + # so convert draft_token_ids to torch.int32 here. + draft_token_ids = self._draft_token_ids.to(dtype=torch.int32) + + self.input_ids.gpu.scatter_( + dim=0, + index=draft_tokens_index_tensor, + src=draft_token_ids.flatten()[prev_draft_token_indices_tensor], + ) + + def _get_encoder_seq_lens( + self, + num_scheduled_tokens: dict[str, int], + kv_cache_spec: KVCacheSpec, + num_reqs: int, + for_cudagraph_capture: bool = False, + ) -> tuple[torch.Tensor | None, np.ndarray | None]: + if not isinstance(kv_cache_spec, CrossAttentionSpec): + return None, None + + # Zero out buffer for padding requests that are not actually scheduled (CGs) + self.encoder_seq_lens.np[:num_reqs] = 0 + + # Build encoder_seq_lens array mapping request indices to + # encoder lengths for inputs scheduled in this batch + for req_id in num_scheduled_tokens: + req_index = self.input_batch.req_id_to_index[req_id] + req_state = self.requests[req_id] + if req_state.mm_features is None: + self.encoder_seq_lens.np[req_index] = 0 + continue + + # Get the total number of encoder input tokens for running encoder requests + # whether encoding is finished or not so that cross-attention knows how + # many encoder tokens to attend to. + encoder_input_tokens = sum( + feature.mm_position.length for feature in req_state.mm_features + ) + self.encoder_seq_lens.np[req_index] = encoder_input_tokens + if for_cudagraph_capture: + # During CUDA graph capture, we need to use realistic encoder lengths + # so that max_seqlen_k is captured with the correct value. + max_encoder_len = getattr( + self.model_config.hf_config, + "max_source_positions", + self.max_encoder_len, + ) + self.encoder_seq_lens.np[:num_reqs] = max_encoder_len + + self.encoder_seq_lens.copy_to_gpu(num_reqs) + encoder_seq_lens = self.encoder_seq_lens.gpu[:num_reqs] + encoder_seq_lens_cpu = self.encoder_seq_lens.np[:num_reqs] + + return encoder_seq_lens, encoder_seq_lens_cpu + + def _prepare_inputs( + self, + scheduler_output: "SchedulerOutput", + num_scheduled_tokens: np.ndarray, + ) -> tuple[ + torch.Tensor, + SpecDecodeMetadata | None, + ]: + """ + :return: tuple[ + logits_indices, spec_decode_metadata, + ] + """ + total_num_scheduled_tokens = scheduler_output.total_num_scheduled_tokens + assert total_num_scheduled_tokens > 0 + num_reqs = self.input_batch.num_reqs + assert num_reqs > 0 + + # OPTIMIZATION: Start copying the block table first. + # This way, we can overlap the copy with the following CPU operations. + self.input_batch.block_table.commit_block_table(num_reqs) + + # Get request indices. + # E.g., [2, 5, 3] -> [0, 0, 1, 1, 1, 1, 1, 2, 2, 2] + req_indices = np.repeat(self.arange_np[:num_reqs], num_scheduled_tokens) + + # cu_num_tokens: [2, 5, 3] -> [2, 7, 10] + # self.query_pos.np[:10]: [0, 1, 0, 1, 2, 3, 4, 0, 1, 2] + cu_num_tokens = self._get_cumsum_and_arange( + num_scheduled_tokens, self.query_pos.np + ) + + # Get positions. + positions_np = ( + self.input_batch.num_computed_tokens_cpu[req_indices] + + self.query_pos.np[: cu_num_tokens[-1]] + ) + + # Calculate M-RoPE positions. + # Only relevant for models using M-RoPE (e.g, Qwen2-VL) + if self.uses_mrope: + self._calc_mrope_positions(scheduler_output) + + # Calculate XD-RoPE positions. + # Only relevant for models using XD-RoPE (e.g, HunYuan-VL) + if self.uses_xdrope_dim > 0: + self._calc_xdrope_positions(scheduler_output) + + # Get token indices. + # E.g., [0, 1, 0, 1, 2, 3, 4, 0, 1, 2] + # -> [0, 1, M, M + 1, M + 2, M + 3, M + 4, 2 * M, 2 * M + 1, 2 * M + 2] + # where M is the max_model_len. + token_indices = ( + positions_np + req_indices * self.input_batch.token_ids_cpu.shape[1] + ) + token_indices_tensor = torch.from_numpy(token_indices) + + # NOTE(woosuk): We use torch.index_select instead of np.take here + # because torch.index_select is much faster than np.take for large + # tensors. + torch.index_select( + self.input_batch.token_ids_cpu_tensor.flatten(), + 0, + token_indices_tensor, + out=self.input_ids.cpu[:total_num_scheduled_tokens], + ) + if self.enable_prompt_embeds: + is_token_ids = self.input_batch.is_token_ids_tensor.flatten() + torch.index_select( + is_token_ids, + 0, + token_indices_tensor, + out=self.is_token_ids.cpu[:total_num_scheduled_tokens], + ) + + # Because we did not pre-allocate a massive prompt_embeds CPU tensor on + # the InputBatch, we need to fill in the prompt embeds into the expected + # spots in the GpuModelRunner's pre-allocated prompt_embeds tensor. + if self.input_batch.req_prompt_embeds: + output_idx = 0 + for req_idx in range(num_reqs): + num_sched = num_scheduled_tokens[req_idx] + + # Skip if this request doesn't have embeddings + if req_idx not in self.input_batch.req_prompt_embeds: + output_idx += num_sched + continue + + # Skip if no tokens scheduled + if num_sched <= 0: + output_idx += num_sched + continue + + req_embeds = self.input_batch.req_prompt_embeds[req_idx] + start_pos = self.input_batch.num_computed_tokens_cpu[req_idx] + + # Skip if trying to read beyond available embeddings + if start_pos >= req_embeds.shape[0]: + output_idx += num_sched + continue + + # Copy available embeddings + end_pos = start_pos + num_sched + actual_end = min(end_pos, req_embeds.shape[0]) + actual_num_sched = actual_end - start_pos + + if actual_num_sched > 0: + self.inputs_embeds.cpu[ + output_idx : output_idx + actual_num_sched + ].copy_(req_embeds[start_pos:actual_end]) + + output_idx += num_sched + + # Prepare the attention metadata. + self.query_start_loc.np[0] = 0 + self.query_start_loc.np[1 : num_reqs + 1] = cu_num_tokens + # Note: pad query_start_loc to be non-decreasing, as kernels + # like FlashAttention requires that + self.query_start_loc.np[num_reqs + 1 :].fill(cu_num_tokens[-1]) + self.query_start_loc.copy_to_gpu() + query_start_loc = self.query_start_loc.gpu[: num_reqs + 1] + + # Compute optimistic seq_lens (assumes all draft tokens from previous + # iteration accepted). Store in optimistic_seq_lens_cpu for use by + # _build_attention_metadata (max_seq_len) and discard_request_mask. + # seq_lens (GPU) will be computed later using the same optimistic values. + torch.add( + self.input_batch.num_computed_tokens_cpu_tensor[:num_reqs], + torch.from_numpy(num_scheduled_tokens), + out=self.optimistic_seq_lens_cpu[:num_reqs], + ) + self.optimistic_seq_lens_cpu[num_reqs:].fill_(0) + + # Build prev_positions mapping: current pos -> prev pos (-1 if new). + # Used for gathering from previous iteration's GPU tensors. + prev_req_id_to_index = self.input_batch.prev_req_id_to_index + self._compute_prev_positions(num_reqs) + + num_tokens = [self.requests[r].num_tokens for r in self.input_batch.req_ids] + num_tokens_np = np.array(num_tokens, dtype=np.int32) + + # Record which requests should not be sampled, + # so that we could clear the sampled tokens before returning + self.discard_request_mask.np[:num_reqs] = ( + self.optimistic_seq_lens_cpu[:num_reqs].numpy() < num_tokens_np + ) + self.discard_request_mask.copy_to_gpu(num_reqs) + + # Sync num_accepted_tokens from CPU (set by + # _update_states_after_model_execute for hybrid models). + if self.num_accepted_tokens_event is not None: + self.num_accepted_tokens_event.synchronize() + # Async mode: condense() reordered indices, use prev_positions mapping + if self.use_async_scheduling and prev_req_id_to_index: + prev_idx = self.prev_positions.np[:num_reqs] + new_mask = prev_idx < 0 + self.num_accepted_tokens.np[:num_reqs] = ( + self.input_batch.num_accepted_tokens_cpu[ + np.where(new_mask, 0, prev_idx) + ] + ) + self.num_accepted_tokens.np[:num_reqs][new_mask] = 1 + self.input_batch.num_accepted_tokens_cpu[:num_reqs] = ( + self.num_accepted_tokens.np[:num_reqs] + ) + else: + # Non-async mode: use values directly + self.num_accepted_tokens.np[:num_reqs] = ( + self.input_batch.num_accepted_tokens_cpu[:num_reqs] + ) + self.num_accepted_tokens.np[num_reqs:].fill(1) + self.num_accepted_tokens.copy_to_gpu() + else: + self.num_accepted_tokens.np.fill(1) + self.num_accepted_tokens.gpu.fill_(1) + + if self.mamba_prev_last_scheduled_idx is not None: + mamba_utils.preprocess_mamba_all_specdec( + scheduler_output, + self.input_batch, + self.mamba_state_idx, + num_reqs, + self.mamba_prev_last_scheduled_idx, + ) + + # Update num_computed_tokens on GPU. In async spec decode, + # CPU values are optimistic (all drafts accepted). The kernel + # corrects on GPU using the previous step's + # valid_sampled_token_count_gpu. Otherwise, just copy from CPU. + if ( + self.use_async_spec_decode + and self.valid_sampled_token_count_gpu is not None + and prev_req_id_to_index + ): + self.prev_positions.copy_to_gpu(num_reqs) + self.prev_num_draft_tokens.copy_to_gpu() + cpu_values = self.input_batch.num_computed_tokens_cpu_tensor[:num_reqs].to( + device=self.device, non_blocking=True + ) + update_num_computed_tokens_for_batch_change( + self.num_computed_tokens, + self.num_accepted_tokens.gpu[:num_reqs], + self.prev_positions.gpu[:num_reqs], + self.valid_sampled_token_count_gpu, + self.prev_num_draft_tokens.gpu, + cpu_values, + ) + else: + self.num_computed_tokens[:num_reqs].copy_( + self.input_batch.num_computed_tokens_cpu_tensor[:num_reqs], + non_blocking=True, + ) + + self.req_indices.np[:total_num_scheduled_tokens] = req_indices + self.req_indices.copy_to_gpu(total_num_scheduled_tokens) + req_indices_gpu = self.req_indices.gpu[:total_num_scheduled_tokens] + + self.query_pos.copy_to_gpu(total_num_scheduled_tokens) + self.num_scheduled_tokens.np[:num_reqs] = num_scheduled_tokens + self.num_scheduled_tokens.copy_to_gpu(num_reqs) + num_scheduled_tokens_gpu = self.num_scheduled_tokens.gpu[:num_reqs] + self.positions[:total_num_scheduled_tokens] = ( + self.num_computed_tokens[req_indices_gpu].to(torch.int64) + + self.query_pos.gpu[:total_num_scheduled_tokens] + ) + self.seq_lens[:num_reqs] = ( + self.num_computed_tokens[:num_reqs] + num_scheduled_tokens_gpu + ) + self.seq_lens[num_reqs:].fill_(0) + + self.input_batch.block_table.compute_slot_mapping( + num_reqs, + self.query_start_loc.gpu[: num_reqs + 1], + self.positions[:total_num_scheduled_tokens], + ) + + # Copy the tensors to the GPU. + self._prepare_input_ids( + scheduler_output, + num_reqs, + total_num_scheduled_tokens, + cu_num_tokens, + ) + + if self.uses_mrope: + # Only relevant for models using M-RoPE (e.g, Qwen2-VL) + self.mrope_positions.gpu[:, :total_num_scheduled_tokens].copy_( + self.mrope_positions.cpu[:, :total_num_scheduled_tokens], + non_blocking=True, + ) + elif self.uses_xdrope_dim > 0: + # Only relevant for models using XD-RoPE (e.g, HunYuan-VL) + self.xdrope_positions.gpu[:, :total_num_scheduled_tokens].copy_( + self.xdrope_positions.cpu[:, :total_num_scheduled_tokens], + non_blocking=True, + ) + if self.use_async_spec_decode and (self.uses_mrope or self.uses_xdrope_dim > 0): + drift = self.num_computed_tokens[req_indices_gpu].to( + torch.int64 + ) - self.input_batch.num_computed_tokens_cpu_tensor[req_indices].to( + device=self.device, dtype=torch.int64, non_blocking=True + ) + target = self.mrope_positions if self.uses_mrope else self.xdrope_positions + target.gpu[:, :total_num_scheduled_tokens] += drift + + use_spec_decode = len(scheduler_output.scheduled_spec_decode_tokens) > 0 + if not use_spec_decode: + # NOTE(woosuk): Due to chunked prefills, the batch may contain + # partial requests. While we should not sample any token + # from these partial requests, we do so for simplicity. + # We will ignore the sampled tokens from the partial requests. + # TODO: Support prompt logprobs. + logits_indices = query_start_loc[1:] - 1 + spec_decode_metadata = None + num_sampled_tokens = np.ones(num_reqs, dtype=np.int32) + else: + # Get the number of draft tokens for each request. + # Iterate over the dictionary rather than all requests since not all + # requests have draft tokens. + num_draft_tokens = np.zeros(num_reqs, dtype=np.int32) + # For chunked prefills, use -1 as mask rather than 0, as guided + # decoding may rollback speculative tokens. + num_decode_draft_tokens = np.full(num_reqs, -1, dtype=np.int32) + for ( + req_id, + draft_token_ids, + ) in scheduler_output.scheduled_spec_decode_tokens.items(): + req_idx = self.input_batch.req_id_to_index[req_id] + draft_len = len(draft_token_ids) + num_draft_tokens[req_idx] = draft_len + if ( + self.input_batch.num_computed_tokens_cpu[req_idx] + >= self.input_batch.num_prompt_tokens[req_idx] + ): + num_decode_draft_tokens[req_idx] = draft_len + spec_decode_metadata = self._calc_spec_decode_metadata( + num_draft_tokens, cu_num_tokens + ) + logits_indices = spec_decode_metadata.logits_indices + num_sampled_tokens = num_draft_tokens + 1 + # For DECODE only cuda graph of some attention backends (e.g., GDN). + self.num_decode_draft_tokens.np[:num_reqs] = num_decode_draft_tokens + self.num_decode_draft_tokens.np[num_reqs:].fill(-1) + self.num_decode_draft_tokens.copy_to_gpu() + + # Hot-Swap lora model + if self.lora_config: + assert ( + np.sum(num_sampled_tokens) + <= self.vllm_config.scheduler_config.max_num_batched_tokens + ) + self.set_active_loras( + self.input_batch, num_scheduled_tokens, num_sampled_tokens + ) + + return ( + logits_indices, + spec_decode_metadata, + ) + + def _build_attention_metadata( + self, + num_tokens: int, + num_reqs: int, + max_query_len: int, + num_tokens_padded: int | None = None, + num_reqs_padded: int | None = None, + ubatch_slices: UBatchSlices | None = None, + logits_indices: torch.Tensor | None = None, + 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]: + """ + :return: tuple[attn_metadata, spec_decode_common_attn_metadata] + """ + # Attention metadata is not needed for attention free models + if len(self.kv_cache_config.kv_cache_groups) == 0: + return {}, None + + num_tokens_padded = num_tokens_padded or num_tokens + num_reqs_padded = num_reqs_padded or num_reqs + assert num_reqs_padded is not None and num_tokens_padded is not None + + attn_metadata: PerLayerAttnMetadata = {} + if ubatch_slices is not None: + attn_metadata = [dict() for _ in range(len(ubatch_slices))] + + if for_cudagraph_capture: + # For some attention backends (e.g. FA) with sliding window models we need + # to make sure the backend see a max_seq_len that is larger to the sliding + # window size when capturing to make sure the correct kernel is selected. + max_seq_len = self.max_model_len + else: + max_seq_len = self.optimistic_seq_lens_cpu.numpy()[:num_reqs].max().item() + + kv_cache_groups = self.kv_cache_config.kv_cache_groups + + def _get_block_table(kv_cache_gid: int): + assert num_reqs_padded is not None and num_tokens_padded is not None + kv_cache_spec = kv_cache_groups[kv_cache_gid].kv_cache_spec + if isinstance(kv_cache_spec, EncoderOnlyAttentionSpec): + blk_table_tensor = torch.zeros( + (num_reqs_padded, 1), + dtype=torch.int32, + device=self.device, + ) + else: + blk_table = self.input_batch.block_table[kv_cache_gid] + blk_table_tensor = blk_table.get_device_tensor(num_reqs_padded) + + # Fill unused block table entries with NULL_BLOCK_ID (null block) + # for CUDAGraph padding. Block 0 is reserved for padding. + blk_table_tensor[num_reqs:num_reqs_padded].fill_(NULL_BLOCK_ID) + return blk_table_tensor + + assert slot_mappings is not None + block_table_gid_0 = _get_block_table(0) + slot_mapping_gid_0 = slot_mappings[0] + + if self.routed_experts_initialized: + # Copy this step's attention slot_mapping into our private + # device buffer. The shared ``slot_mappings[attn_gid]`` is + # owned by the attention block table and will be overwritten + # by the next ``_prepare_inputs``; we need a stable snapshot + # because the async D2H may still be in flight on the copy + # stream when the next step runs. + attn_gid = self.routed_experts_attn_gid + slot_mapping_attn = slot_mappings[attn_gid] + self.routed_experts_slot_mapping_device[:num_tokens].copy_( + slot_mapping_attn[:num_tokens] + ) + + num_computed_tokens_cpu = self.input_batch.num_computed_tokens_cpu_tensor[ + :num_reqs_padded + ] + num_prompt_tokens_cpu = self.input_batch.num_prompt_tokens_cpu_tensor[ + :num_reqs_padded + ] + seq_lens_cpu = self.optimistic_seq_lens_cpu[:num_reqs_padded] + seq_lens_cpu_upper_bound = seq_lens_cpu + + # is_prefilling: True if request is still in prefill phase. + # Used by mamba backends to distinguish actual decodes from + # short extends. + is_prefilling = num_computed_tokens_cpu < num_prompt_tokens_cpu + # Zero out padded rows so stale data from condense() doesn't + # misclassify padding as prefill in CUDA graph mode. + is_prefilling[num_reqs:] = False + + if self.use_async_spec_decode: + # GPU tensors are authoritative in async mode. + seq_lens_cpu = None + num_computed_tokens_cpu = None + + cm_base = CommonAttentionMetadata( + query_start_loc=self.query_start_loc.gpu[: num_reqs_padded + 1], + query_start_loc_cpu=self.query_start_loc.cpu[: num_reqs_padded + 1], + seq_lens=self.seq_lens[:num_reqs_padded], + _seq_lens_cpu=seq_lens_cpu, + _num_computed_tokens_cpu=num_computed_tokens_cpu, + seq_lens_cpu_upper_bound=seq_lens_cpu_upper_bound, + num_reqs=num_reqs_padded, + num_actual_tokens=num_tokens_padded, + max_query_len=max_query_len, + max_seq_len=max_seq_len, + block_table_tensor=block_table_gid_0, + slot_mapping=slot_mapping_gid_0, + causal=True, + is_prefilling=is_prefilling, + positions=self.positions[:num_tokens_padded], + ) + + if self.dcp_world_size > 1: + self.dcp_local_seq_lens.cpu[:num_reqs] = get_dcp_local_seq_lens( + self.optimistic_seq_lens_cpu[:num_reqs], + self.dcp_world_size, + self.dcp_rank, + self.parallel_config.cp_kv_cache_interleave_size, + ) + self.dcp_local_seq_lens.cpu[num_reqs:].fill_(0) + self.dcp_local_seq_lens.copy_to_gpu(num_reqs_padded) + + cm_base.dcp_local_seq_lens = self.dcp_local_seq_lens.gpu[:num_reqs_padded] + cm_base.dcp_local_seq_lens_cpu = self.dcp_local_seq_lens.cpu[ + :num_reqs_padded + ] + + if logits_indices is not None and self.cache_config.kv_sharing_fast_prefill: + cm_base.num_logits_indices = logits_indices.size(0) + cm_base.logits_indices_padded = self._prepare_kv_sharing_fast_prefill( + logits_indices + ) + + # Cache attention metadata builds across hybrid KV-cache groups + # The only thing that changes between different hybrid KV-cache groups when the + # same metadata builder and KVCacheSpec is the same is the block table, so we + # can cache the attention metadata builds and just update the block table using + # `builder.update_block_table` if the builder supports it. + cached_attn_metadata: dict[ + tuple[KVCacheSpec, type[AttentionMetadataBuilder]], AttentionMetadata + ] = {} + + def _build_attn_group_metadata( + kv_cache_gid: int, + attn_gid: int, + common_attn_metadata: CommonAttentionMetadata, + ubid: int | None = None, + ) -> None: + attn_group = self.attn_groups[kv_cache_gid][attn_gid] + builder = attn_group.get_metadata_builder(ubid or 0) + kv_cache_spec = kv_cache_groups[kv_cache_gid].kv_cache_spec + if isinstance(kv_cache_spec, UniformTypeKVCacheSpecs): + 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, (Mamba2AttentionMetadataBuilder, GDNAttentionMetadataBuilder) + ): + assert ubid is None, "UBatching not supported with GDN yet" + extra_attn_metadata_args = dict( + num_accepted_tokens=self.num_accepted_tokens.gpu[:num_reqs_padded], + num_decode_draft_tokens_cpu=self.num_decode_draft_tokens.cpu[ + :num_reqs_padded + ], + ) + if ( + isinstance(builder, Mamba2AttentionMetadataBuilder) + and self.mamba_prev_last_scheduled_idx is not None + ): + extra_attn_metadata_args["prev_last_scheduled_idx"] = ( + self.mamba_prev_last_scheduled_idx.gpu[:num_reqs_padded] + ) + + if for_cudagraph_capture: + attn_metadata_i = builder.build_for_cudagraph_capture( + common_attn_metadata + ) + elif ( + cache_key in cached_attn_metadata + and builder.supports_update_block_table + ): + attn_metadata_i = builder.update_block_table( + cached_attn_metadata[cache_key], + common_attn_metadata.block_table_tensor, + common_attn_metadata.slot_mapping, + ) + else: + attn_metadata_i = builder.build( + common_prefix_len=cascade_attn_prefix_len, + common_attn_metadata=common_attn_metadata, + **extra_attn_metadata_args, + ) + if builder.supports_update_block_table: + cached_attn_metadata[cache_key] = attn_metadata_i + + if ubid is None: + assert isinstance(attn_metadata, dict) + attn_metadata_dict = attn_metadata + else: + assert isinstance(attn_metadata, list) + attn_metadata_dict = attn_metadata[ubid] + + for layer_name in attn_group.layer_names: + attn_metadata_dict[layer_name] = attn_metadata_i + + # Prepare the attention metadata for each KV cache group and make layers + # in the same group share the same metadata. + spec_decode_common_attn_metadata = None + for kv_cache_gid, kv_cache_group in enumerate(kv_cache_groups): + cm = copy(cm_base) # shallow copy + + # Basically only the encoder seq_lens, block_table and slot_mapping change + # for each kv_cache_group. + cm.encoder_seq_lens, cm.encoder_seq_lens_cpu = self._get_encoder_seq_lens( + num_scheduled_tokens or {}, + kv_cache_group.kv_cache_spec, + num_reqs_padded, + for_cudagraph_capture=for_cudagraph_capture, + ) + if kv_cache_gid > 0: + cm.block_table_tensor = _get_block_table(kv_cache_gid) + cm.slot_mapping = slot_mappings[kv_cache_gid] + + if self.speculative_config and spec_decode_common_attn_metadata is None: + if isinstance( + self.drafter, + ( + EagleProposer, + DFlashProposer, + Gemma4Proposer, + ExtractHiddenStatesProposer, + ), + ): + if self.drafter.kv_cache_gid == kv_cache_gid: + spec_decode_common_attn_metadata = cm + else: + spec_decode_common_attn_metadata = cm + # Capture per-group block tables for multi-group proposers. + if self.speculative_config and isinstance(self.drafter, Gemma4Proposer): + self.drafter.set_per_group_block_table( + kv_cache_gid, cm.block_table_tensor + ) + + for attn_gid in range(len(self.attn_groups[kv_cache_gid])): + if ubatch_slices is not None: + for ubid, _cm in enumerate(split_attn_metadata(ubatch_slices, cm)): + _build_attn_group_metadata(kv_cache_gid, attn_gid, _cm, ubid) + + else: + _build_attn_group_metadata(kv_cache_gid, attn_gid, cm) + + if self.is_mm_prefix_lm: + req_doc_ranges = {} + + # Gemma4 bidi: skip ranges that exceed the sliding + # window. When image tokens > sliding_window, bidi causes + # early image tokens to attend to the entire image + # (e.g. 6 → 1092 targets), degrading spatial precision. + # Per-range filtering keeps bidi for small images/video + # frames while skipping oversized images. + hf_text_config = self.model_config.hf_text_config + _bidi_sw = getattr(hf_text_config, "sliding_window", None) + + for req_id in self.input_batch.req_ids: + image_doc_ranges = [] + req_state = self.requests[req_id] + for mm_feature in req_state.mm_features: + pos_info = mm_feature.mm_position + img_doc_range = pos_info.extract_embeds_range() + for r in img_doc_range: + if _bidi_sw is not None and (r[1] - r[0] + 1) > _bidi_sw: + continue + image_doc_ranges.append(r) + req_idx = self.input_batch.req_id_to_index[req_id] + req_doc_ranges[req_idx] = image_doc_ranges + + # Set mm_prefix_range for all attention metadata + self._set_mm_prefix_range_for_metadata(attn_metadata, req_doc_ranges) + + if spec_decode_common_attn_metadata is not None and ( + num_reqs != num_reqs_padded or num_tokens != num_tokens_padded + ): + # Currently the drafter still only uses piecewise cudagraphs (and modifies + # the attention metadata in directly), and therefore does not want to use + # padded attention metadata. + spec_decode_common_attn_metadata = ( + spec_decode_common_attn_metadata.unpadded(num_tokens, num_reqs) + ) + + 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: + """ + :return: 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_common_prefix_blocks: Number of shared KV cache blocks. + + 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): + req = self.requests[req_id] + assert req.mrope_positions is not None + + num_computed_tokens = self.input_batch.num_computed_tokens_cpu[index] + num_scheduled_tokens = scheduler_output.num_scheduled_tokens[req_id] + num_prompt_tokens = length_from_prompt_token_ids_or_embeds( + req.prompt_token_ids, req.prompt_embeds + ) + + if num_computed_tokens + num_scheduled_tokens > num_prompt_tokens: + prompt_part_len = max(0, num_prompt_tokens - num_computed_tokens) + completion_part_len = max(0, num_scheduled_tokens - prompt_part_len) + else: + prompt_part_len = num_scheduled_tokens + completion_part_len = 0 + + assert num_scheduled_tokens == prompt_part_len + completion_part_len + + if prompt_part_len > 0: + # prompt's mrope_positions are pre-computed + dst_start = mrope_pos_ptr + dst_end = mrope_pos_ptr + prompt_part_len + src_start = num_computed_tokens + src_end = num_computed_tokens + prompt_part_len + + self.mrope_positions.cpu[:, dst_start:dst_end] = req.mrope_positions[ + :, src_start:src_end + ] + mrope_pos_ptr += prompt_part_len + + if completion_part_len > 0: + # compute completion's mrope_positions on-the-fly + dst_start = mrope_pos_ptr + dst_end = mrope_pos_ptr + completion_part_len + + assert req.mrope_position_delta is not None + MRotaryEmbedding.get_next_input_positions_tensor( + out=self.mrope_positions.np, + out_offset=dst_start, + mrope_position_delta=req.mrope_position_delta, + context_len=num_computed_tokens + prompt_part_len, + num_new_tokens=completion_part_len, + ) + + mrope_pos_ptr += completion_part_len + + def _calc_xdrope_positions(self, scheduler_output: "SchedulerOutput"): + xdrope_pos_ptr = 0 + for index, req_id in enumerate(self.input_batch.req_ids): + req = self.requests[req_id] + assert req.xdrope_positions is not None + + num_computed_tokens = self.input_batch.num_computed_tokens_cpu[index] + num_scheduled_tokens = scheduler_output.num_scheduled_tokens[req_id] + num_prompt_tokens = length_from_prompt_token_ids_or_embeds( + req.prompt_token_ids, req.prompt_embeds + ) + + if num_computed_tokens + num_scheduled_tokens > num_prompt_tokens: + prompt_part_len = max(0, num_prompt_tokens - num_computed_tokens) + completion_part_len = max(0, num_scheduled_tokens - prompt_part_len) + else: + prompt_part_len = num_scheduled_tokens + completion_part_len = 0 + + assert num_scheduled_tokens == prompt_part_len + completion_part_len + + if prompt_part_len > 0: + # prompt's xdrope_positions are pre-computed + dst_start = xdrope_pos_ptr + dst_end = xdrope_pos_ptr + prompt_part_len + src_start = num_computed_tokens + src_end = num_computed_tokens + prompt_part_len + + self.xdrope_positions.cpu[:, dst_start:dst_end] = req.xdrope_positions[ + :, src_start:src_end + ] + xdrope_pos_ptr += prompt_part_len + + if completion_part_len > 0: + # compute completion's xdrope_positions on-the-fly + dst_start = xdrope_pos_ptr + dst_end = xdrope_pos_ptr + completion_part_len + + XDRotaryEmbedding.get_next_input_positions_tensor( + out=self.xdrope_positions.np, + out_offset=dst_start, + context_len=num_computed_tokens + prompt_part_len, + num_new_tokens=completion_part_len, + ) + + xdrope_pos_ptr += completion_part_len + + def _calc_spec_decode_metadata( + self, + num_draft_tokens: np.ndarray, + cu_num_scheduled_tokens: np.ndarray, + ) -> SpecDecodeMetadata: + # Inputs: + # cu_num_scheduled_tokens: [ 4, 104, 107, 207, 209] + # num_draft_tokens: [ 3, 0, 2, 0, 1] + # Outputs: + # cu_num_draft_tokens: [ 3, 3, 5, 5, 6] + # logits_indices: [ 0, 1, 2, 3, 103, 104, 105, 106, + # 206, 207, 208] + # target_logits_indices: [ 0, 1, 2, 5, 6, 9] + # bonus_logits_indices: [ 3, 4, 7, 8, 10] + + # Compute the logits indices. + # [4, 1, 3, 1, 2] + num_sampled_tokens = num_draft_tokens + 1 + + # Step 1. + # cu_num_sampled_tokens: [4, 5, 8, 9, 11] + # _arange_scratch[:11]: [0, 1, 2, 3, 0, 0, 1, 2, 0, 0, 1] + cu_num_sampled_tokens = self._get_cumsum_and_arange( + num_sampled_tokens, self._arange_scratch, cumsum_dtype=np.int32 + ) + # Step 2. [0, 0, 0, 0, 103, 104, 104, 104, 206, 207, 207] + logits_indices = np.repeat( + cu_num_scheduled_tokens - num_sampled_tokens, num_sampled_tokens + ) + # Step 3. [0, 1, 2, 3, 103, 104, 105, 106, 206, 207, 208] + logits_indices += self._arange_scratch[: cu_num_sampled_tokens[-1]] + + # Compute the bonus logits indices. + bonus_logits_indices = cu_num_sampled_tokens - 1 + + # Compute the draft logits indices. + # cu_num_draft_tokens: [3, 3, 5, 5, 6] + # _arange_scratch[:6]: [0, 1, 2, 0, 1, 0] + cu_num_draft_tokens = self._get_cumsum_and_arange( + num_draft_tokens, self._arange_scratch, cumsum_dtype=np.int32 + ) + # [0, 0, 0, 5, 5, 9] + target_logits_indices = np.repeat( + cu_num_sampled_tokens - num_sampled_tokens, num_draft_tokens + ) + # [0, 1, 2, 5, 6, 9] + target_logits_indices += self._arange_scratch[: cu_num_draft_tokens[-1]] + + # TODO: Optimize the CPU -> GPU copy. + cu_num_draft_tokens = torch.from_numpy(cu_num_draft_tokens).to( + self.device, non_blocking=True + ) + cu_num_sampled_tokens = torch.from_numpy(cu_num_sampled_tokens).to( + self.device, non_blocking=True + ) + logits_indices = torch.from_numpy(logits_indices).to( + self.device, non_blocking=True + ) + target_logits_indices = torch.from_numpy(target_logits_indices).to( + self.device, non_blocking=True + ) + bonus_logits_indices = torch.from_numpy(bonus_logits_indices).to( + self.device, non_blocking=True + ) + + # Compute the draft token ids. + # draft_token_indices: [ 1, 2, 3, 105, 106, 208] + draft_token_ids = self.input_ids.gpu[logits_indices] + draft_token_ids = draft_token_ids[target_logits_indices + 1] + + return SpecDecodeMetadata( + draft_token_ids=draft_token_ids, + num_draft_tokens=num_draft_tokens.tolist(), + cu_num_draft_tokens=cu_num_draft_tokens, + cu_num_sampled_tokens=cu_num_sampled_tokens, + target_logits_indices=target_logits_indices, + bonus_logits_indices=bonus_logits_indices, + logits_indices=logits_indices, + ) + + def _prepare_kv_sharing_fast_prefill( + self, + logits_indices: torch.Tensor, + ) -> torch.Tensor: + assert self.kv_sharing_fast_prefill_logits_indices is not None + num_logits = logits_indices.shape[0] + assert num_logits > 0 + self.kv_sharing_fast_prefill_logits_indices[:num_logits].copy_(logits_indices) + # There might have leftover indices in logits_indices[num_logits:] + # from previous iterations, whose values may be greater than the + # batch size in the current iteration. To ensure indices are always + # valid, fill the padded indices with the last index. Broadcast the + # scalar GPU-side to avoid a D2H sync on `.item()`. + self.kv_sharing_fast_prefill_logits_indices[num_logits:] = logits_indices[-1] + # Dispatch for the decoder portion of the model. + _, batch_desc = self.cudagraph_dispatcher.dispatch( + num_logits, invalid_modes={CUDAGraphMode.FULL} + ) + num_logits_padded = batch_desc.num_tokens + logits_indices_padded = self.kv_sharing_fast_prefill_logits_indices[ + :num_logits_padded + ] + return logits_indices_padded + + def _batch_mm_inputs_from_scheduler( + self, + scheduler_output: "SchedulerOutput", + ) -> tuple[ + list[str], + list[tuple[str, MultiModalKwargsItem]], + list[tuple[str, PlaceholderRange]], + ]: + """Batch multimodal inputs from scheduled encoder inputs. + + Args: + scheduler_output: The scheduler output containing scheduled encoder + inputs. + + Returns: + A tuple of (mm_hashes, mm_kwargs, mm_lora_refs) where: + - mm_hashes: List of multimodal hashes for each item + - mm_kwargs: List of multimodal kwargs for each item + - mm_lora_refs: List of (req_id, placeholder_range) for each item + """ + scheduled_encoder_inputs = scheduler_output.scheduled_encoder_inputs + if not scheduled_encoder_inputs: + return [], [], [] + + mm_hashes = list[str]() + mm_kwargs = list[tuple[str, MultiModalKwargsItem]]() + # Multimodal LoRA reference info to map each multimodal item + # back to its request & position + mm_lora_refs = list[tuple[str, PlaceholderRange]]() + for req_id, encoder_input_ids in scheduled_encoder_inputs.items(): + req_state = self.requests[req_id] + + for mm_input_id in encoder_input_ids: + mm_feature = req_state.mm_features[mm_input_id] + if mm_feature.data is None: + continue + + mm_hashes.append(mm_feature.identifier) + 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 + + def _execute_mm_encoder( + self, scheduler_output: "SchedulerOutput" + ) -> list[torch.Tensor]: + mm_hashes, mm_kwargs, mm_lora_refs = self._batch_mm_inputs_from_scheduler( + scheduler_output + ) + + if not mm_kwargs: + return [] + + # `prompt_embeds` is a passthrough modality, the tensor is already in + # the model embedding space, so no encoder runs. Inject each + # `prompt_embeds` tensor directly into the encoder cache here so that + # `_gather_mm_embeddings` can splice it via the standard `is_mm_embed` + # path. + pe_indices = [ + i + for i, (modality, _) in enumerate(mm_kwargs) + if modality == "prompt_embeds" + ] + if pe_indices: + for i in pe_indices: + pe_tensor = mm_kwargs[i][1]["embedding"].data + assert isinstance(pe_tensor, torch.Tensor) + + self.encoder_cache[mm_hashes[i]] = pe_tensor.to(self.device) + self.maybe_save_ec_to_connector(self.encoder_cache, mm_hashes[i]) + # Filter out `prompt_embeds` items from mm_kwargs/mm_hashes/mm_lora_refs + # since they don't require further encoder processing. + mm_hashes = [h for i, h in enumerate(mm_hashes) if i not in pe_indices] + mm_kwargs = [k for i, k in enumerate(mm_kwargs) if i not in pe_indices] + mm_lora_refs = [ + r for i, r in enumerate(mm_lora_refs) if i not in pe_indices + ] + if not mm_kwargs: + return [] # nothing left to encode after filtering out `prompt_embeds` + + should_time = bool( + self.observability_config + and self.observability_config.enable_mm_processor_stats + 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(): + # Build LoRA mappings independently for encoder inputs + # (encoder batch structure is different from main batch) + prompt_lora_mapping = [] + token_lora_mapping = [] + lora_requests = set() + encoder_token_counts = [] + + for req_id, pos_info in mm_lora_refs: + req_idx = self.input_batch.req_id_to_index[req_id] + lora_id = int(self.input_batch.request_lora_mapping[req_idx]) + + # Prefer pos_info.get_num_embeds to count precise MM embedding tokens. + num_tokens = self.model.get_num_mm_encoder_tokens( # type: ignore[attr-defined] + pos_info.get_num_embeds() + ) + prompt_lora_mapping.append(lora_id) + token_lora_mapping.extend([lora_id] * num_tokens) + encoder_token_counts.append(num_tokens) + + if lora_id > 0: + lora_request = self.input_batch.lora_id_to_lora_request.get(lora_id) + if lora_request is not None: + lora_requests.add(lora_request) + + # Set tower adapter mapping + tower_mapping = LoRAMapping( + tuple(token_lora_mapping), + tuple(prompt_lora_mapping), + is_prefill=True, + type=LoRAMappingType.TOWER, + ) + self.lora_manager.set_active_adapters(lora_requests, tower_mapping) + + # Only set connector mapping if the model actually has a connector. + # Some multimodal models inherit a stub `get_num_mm_connector_tokens` + # from `SupportsMultiModal`, which returns None and should not be + # treated as a signal that connector LoRA is supported. + mm_mapping = ( + self.model.get_mm_mapping() # type: ignore[attr-defined] + if hasattr(self.model, "get_mm_mapping") + else None + ) + if ( + mm_mapping is not None + and mm_mapping.connector + and hasattr(self.model, "get_num_mm_connector_tokens") + ): + post_op_counts = [ + self.model.get_num_mm_connector_tokens(num_tokens) # type: ignore[attr-defined] + for num_tokens in encoder_token_counts + ] + + connector_token_mapping = np.repeat( + np.array(prompt_lora_mapping, dtype=np.int32), + np.array(post_op_counts, dtype=np.int32), + ) + connector_mapping = LoRAMapping( + index_mapping=tuple(connector_token_mapping.tolist()), + prompt_mapping=tuple(prompt_lora_mapping), + is_prefill=True, + type=LoRAMappingType.CONNECTOR, + ) + + self.lora_manager.set_active_adapters( + lora_requests, + connector_mapping, + ) + + encoder_outputs: list[torch.Tensor] = [] + # Track the current index in mm_kwargs/mm_lora_refs to map groups to request IDs + current_item_idx = 0 + for modality, num_items, mm_kwargs_batch in group_and_batch_mm_kwargs( + mm_kwargs, + device=self.device, + pin_memory=self.pin_memory, + ): + batch_outputs: MultiModalEmbeddings + + # EVS and dynamic res video related change. + # (ekhvedchenia): Temporary hack to limit peak memory usage when + # processing multimodal data. This solves the issue with scheduler + # putting too many video samples into a single batch. Scheduler + # uses pruned vision tokens count to compare it versus compute + # budget which is incorrect (Either input media size or non-pruned + # output vision tokens count should be considered) + # dynamic res video for nemotron temporarily uses this hack via + # requires_sequential_video_encoding + # because it doesn't yet support video batching. + # TODO(ywang96): Fix memory profiling to take EVS into account and + # remove this hack. + if ( + ( + self.is_multimodal_pruning_enabled + or self.requires_sequential_video_encoding + ) + and modality == "video" + and num_items > 1 + ): + batch_outputs_lst = list[torch.Tensor]() + for video_idx in range(num_items): + video_mm_kwargs_item = mm_kwargs[current_item_idx + video_idx] + with self.timed_encoder_operation( + should_time, mm_lora_refs, current_item_idx + video_idx, 1 + ): + _, _, micro_batch_mm_inputs = next( + group_and_batch_mm_kwargs( + [video_mm_kwargs_item], + device=self.device, + pin_memory=self.pin_memory, + ) + ) + + micro_batch_outputs = model.embed_multimodal( + **micro_batch_mm_inputs + ) + + batch_outputs_lst.extend(micro_batch_outputs) + + batch_outputs = batch_outputs_lst + else: + # Run the encoder. + # `batch_outputs` is either of the following: + # 1. A tensor of shape (num_items, feature_size, hidden_size) + # in case feature_size is fixed across all multimodal items. + # 2. A list or tuple (length: num_items) of tensors, + # each of shape (feature_size, hidden_size) in case the feature + # size is dynamic depending on the input multimodal items. + + with self.timed_encoder_operation( + should_time, mm_lora_refs, current_item_idx, num_items + ): + cudagraph_output = None + if ( + self.encoder_cudagraph_manager is not None + and self.encoder_cudagraph_manager.supports_modality(modality) + ): + cudagraph_output = self.encoder_cudagraph_manager.execute( + mm_kwargs_batch, + ) + + if cudagraph_output is not None: + batch_outputs = cudagraph_output + else: + batch_outputs = model.embed_multimodal(**mm_kwargs_batch) + + sanity_check_mm_encoder_outputs(batch_outputs, expected_num_items=num_items) + encoder_outputs.extend(batch_outputs) + + current_item_idx += num_items + + # Cache the encoder outputs by mm_hash + for mm_hash, output in zip(mm_hashes, encoder_outputs): + self.encoder_cache[mm_hash] = output + logger.debug("Finish execute for mm hash %s", mm_hash) + self.maybe_save_ec_to_connector(self.encoder_cache, mm_hash) + + return encoder_outputs + + def _gather_mm_embeddings( + self, + scheduler_output: "SchedulerOutput", + shift_computed_tokens: int = 0, + ) -> tuple[list[torch.Tensor], torch.Tensor]: + total_num_scheduled_tokens = scheduler_output.total_num_scheduled_tokens + + mm_embeds = list[torch.Tensor]() + is_mm_embed = torch.zeros( + total_num_scheduled_tokens, dtype=torch.bool, device="cpu" + ) + + req_start_idx = 0 + should_sync_mrope_positions = False + should_sync_xdrope_positions = False + + for req_id in self.input_batch.req_ids: + mm_embeds_req: list[torch.Tensor] = [] + + num_scheduled_tokens = scheduler_output.num_scheduled_tokens[req_id] + req_state = self.requests[req_id] + num_computed_tokens = req_state.num_computed_tokens + shift_computed_tokens + + for mm_feature in req_state.mm_features: + pos_info = mm_feature.mm_position + start_pos = pos_info.offset + num_encoder_tokens = pos_info.length + + # The encoder output is needed if the two ranges overlap: + # [num_computed_tokens, + # num_computed_tokens + num_scheduled_tokens) and + # [start_pos, start_pos + num_encoder_tokens) + if start_pos >= num_computed_tokens + num_scheduled_tokens: + # The encoder output is not needed in this step. + break + if start_pos + num_encoder_tokens <= num_computed_tokens: + # The encoder output is already processed and stored + # in the decoder's KV cache. + continue + + start_idx = max(num_computed_tokens - start_pos, 0) + end_idx = min( + num_computed_tokens - start_pos + num_scheduled_tokens, + num_encoder_tokens, + ) + assert start_idx < end_idx + curr_embeds_start, curr_embeds_end = ( + pos_info.get_embeds_indices_in_range(start_idx, end_idx) + ) + # If there are no embeddings in the current range, we skip + # gathering the embeddings. + if curr_embeds_start == curr_embeds_end: + continue + + mm_hash = mm_feature.identifier + encoder_output = self.encoder_cache.get(mm_hash, None) + assert encoder_output is not None, f"Encoder cache miss for {mm_hash}." + + if (is_embed := pos_info.is_embed) is not None: + is_embed = is_embed[start_idx:end_idx] + mm_embeds_item = encoder_output[curr_embeds_start:curr_embeds_end] + else: + mm_embeds_item = encoder_output[start_idx:end_idx] + + req_start_pos = req_start_idx + start_pos - num_computed_tokens + # OR mask for overlapping mm_features (use_audio_in_video) + if is_embed is None: + is_mm_embed[req_start_pos + start_idx : req_start_pos + end_idx] = ( + True + ) + else: + is_mm_embed[ + req_start_pos + start_idx : req_start_pos + end_idx + ] |= is_embed + mm_embeds_req.append(mm_embeds_item) + + if self.is_multimodal_pruning_enabled and self.uses_mrope: + assert req_state.mrope_positions is not None + should_sync_mrope_positions = True + mm_embeds_req, new_mrope_positions, new_delta = ( + self.model.recompute_mrope_positions( + input_ids=req_state.prompt_token_ids, + multimodal_embeddings=mm_embeds_req, + mrope_positions=req_state.mrope_positions, + num_computed_tokens=req_state.num_computed_tokens, + ) + ) + req_state.mrope_positions.copy_(new_mrope_positions) + req_state.mrope_position_delta = new_delta + + mm_embeds.extend(mm_embeds_req) + req_start_idx += num_scheduled_tokens + + if should_sync_mrope_positions: + self._calc_mrope_positions(scheduler_output) + self.mrope_positions.copy_to_gpu(total_num_scheduled_tokens) + + if should_sync_xdrope_positions: + self._calc_xdrope_positions(scheduler_output) + self.xdrope_positions.copy_to_gpu(total_num_scheduled_tokens) + + return mm_embeds, is_mm_embed + + def get_model(self) -> nn.Module: + if not hasattr(self, "model"): + raise ValueError("Cannot get model before model has been initialized") + if isinstance( + self.model, (CUDAGraphWrapper, UBatchWrapper, BreakableCUDAGraphWrapper) + ): + # get raw model out of the cudagraph wrapper. + return self.model.unwrap() + return self.model + + def get_supported_generation_tasks(self) -> list[GenerationTask]: + model = self.get_model() + supported_tasks = list[GenerationTask]() + + if is_text_generation_model(model): + supported_tasks.append("generate") + + if supports_transcription(model): + if model.supports_transcription_only: + return ["transcription"] + + supported_tasks.append("transcription") + + if supports_realtime(model): + supported_tasks.append("realtime") + + return supported_tasks + + def get_supported_pooling_tasks(self) -> list[PoolingTask]: + model = self.get_model() + if not is_pooling_model(model): + return [] + + return list(model.pooler.get_supported_tasks()) + + def get_supported_tasks(self) -> tuple[SupportedTask, ...]: + tasks = list[SupportedTask]() + + if self.model_config.runner_type == "generate": + tasks.extend(self.get_supported_generation_tasks()) + if self.model_config.runner_type == "pooling": + tasks.extend(self.get_supported_pooling_tasks()) + + return tuple(tasks) + + def sync_and_gather_intermediate_tensors( + self, + num_tokens: int, + intermediate_tensors: IntermediateTensors | None, + sync_self: bool, + ) -> 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 + ) + + return IntermediateTensors( + {k: v[:num_tokens] for k, v in self.intermediate_tensors.items()} + ) + + def eplb_step(self, is_dummy: bool = False, is_profile: bool = False) -> None: + """ + Step for the EPLB (Expert Parallelism Load Balancing) state. + """ + if not self.parallel_config.enable_eplb or self.eep_eplb_suppressed: + return + + assert self.eplb_state is not None + assert self._moe_model is not None + self.eplb_state.step( + is_dummy, + is_profile, + log_stats=self.parallel_config.eplb_config.log_balancedness, + ) + + def setup_eplb_from_mapping( + self, + expanded_physical_to_logical: torch.Tensor, + old_num_physical_experts: int, + ) -> None: + assert self._moe_model is not None + + self.eplb_state = EplbState.from_mapping( + model=self._moe_model, + model_config=self.model_config, + device=self.device, + parallel_config=self.parallel_config, + expanded_physical_to_logical=expanded_physical_to_logical, + num_valid_physical_experts=old_num_physical_experts, + ) + + def _pool( + self, + hidden_states: torch.Tensor, + num_scheduled_tokens: int, + num_scheduled_tokens_np: np.ndarray, + kv_connector_output: KVConnectorOutput | None, + ) -> ModelRunnerOutput | AsyncModelRunnerOutput: + num_reqs = self.input_batch.num_reqs + assert num_reqs == len(self.input_batch.pooling_params), ( + "Either all or none of the requests in a batch must be pooling request" + ) + + hidden_states = hidden_states[:num_scheduled_tokens] + seq_lens_cpu = self.optimistic_seq_lens_cpu[:num_reqs] + + pooling_metadata = self.input_batch.get_pooling_metadata() + pooling_metadata.build_pooling_cursor( + num_scheduled_tokens_np, + seq_lens_cpu, + device=hidden_states.device, + query_start_loc_gpu=self.query_start_loc.gpu[: num_reqs + 1], + ) + + model = cast(VllmModelForPooling, self.model) + raw_pooler_output: PoolerOutput = model.pooler( + hidden_states=hidden_states, pooling_metadata=pooling_metadata + ) + + finished_mask = [ + seq_len == prompt_len + for seq_len, prompt_len in zip(seq_lens_cpu, pooling_metadata.prompt_lens) + ] + raw_pooler_output = self.late_interaction_runner.postprocess_pooler_output( + raw_pooler_output=raw_pooler_output, + pooling_params=pooling_metadata.pooling_params, + req_ids=self.input_batch.req_ids, + finished_mask=finished_mask, + ) + + model_runner_output = ModelRunnerOutput( + req_ids=self.input_batch.req_ids.copy(), + req_id_to_index=self.input_batch.req_id_to_index.copy(), + kv_connector_output=kv_connector_output, + ) + + if raw_pooler_output is None or not any(finished_mask): + model_runner_output.pooler_output = [None] * num_reqs + return model_runner_output + + if not current_platform.is_cuda_alike(): + # cpu/xpu runners cannot use the CUDA stream/event-based wrapper. + model_runner_output.pooler_output = _copy_pooler_output_to_cpu( + raw_pooler_output=raw_pooler_output, + finished_mask=finished_mask, + ) + self._sync_device() + return model_runner_output + + return AsyncGPUPoolingModelRunnerOutput( + model_runner_output=model_runner_output, + raw_pooler_output=raw_pooler_output, + finished_mask=finished_mask, + 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_mm_inputs( + self, num_tokens: int + ) -> tuple[torch.Tensor | None, torch.Tensor]: + if self.model.requires_raw_input_tokens: + input_ids = self.input_ids.gpu[:num_tokens] + else: + input_ids = None + + inputs_embeds = self.inputs_embeds.gpu[:num_tokens] + return input_ids, inputs_embeds + + def _preprocess( + self, + scheduler_output: "SchedulerOutput", + num_input_tokens: int, # Padded + intermediate_tensors: IntermediateTensors | None = None, + ) -> tuple[ + torch.Tensor | None, + torch.Tensor | None, + torch.Tensor, + IntermediateTensors | None, + dict[str, Any], + ECConnectorOutput | None, + ]: + num_scheduled_tokens = scheduler_output.total_num_scheduled_tokens + is_first_rank = get_pp_group().is_first_rank + is_encoder_decoder = self.model_config.is_encoder_decoder + + # _prepare_inputs may reorder the batch, so we must gather multi + # modal outputs after that to ensure the correct order + ec_connector_output = None + + if self.supports_mm_inputs and is_first_rank and not is_encoder_decoder: + # Run the multimodal encoder if any. + with self.maybe_get_ec_connector_output( + scheduler_output, + encoder_cache=self.encoder_cache, + ) as ec_connector_output: + self._execute_mm_encoder(scheduler_output) + mm_embeds, is_mm_embed = self._gather_mm_embeddings(scheduler_output) + + # NOTE(woosuk): To unify token ids and soft tokens (vision + # embeddings), we always use embeddings (rather than token ids) + # as input to the multimodal model, even when the input is text. + inputs_embeds_scheduled = self.model.embed_input_ids( + self.input_ids.gpu[:num_scheduled_tokens], + multimodal_embeddings=mm_embeds, + is_multimodal=is_mm_embed, + ) + + # TODO(woosuk): Avoid the copy. Optimize. + self.inputs_embeds.gpu[:num_scheduled_tokens].copy_(inputs_embeds_scheduled) + + input_ids, inputs_embeds = self._prepare_mm_inputs(num_input_tokens) + model_kwargs = { + **self._init_model_kwargs(), + **self._extract_mm_kwargs(scheduler_output), + } + elif self.enable_prompt_embeds and is_first_rank: + # Get the input embeddings for the tokens that are not input embeds, + # then put them into the appropriate positions. + # TODO(qthequartermasterman): Since even when prompt embeds are + # enabled, (a) not all requests will use prompt embeds, and (b) + # after the initial prompt is processed, the rest of the generated + # tokens will be token ids, it is not desirable to have the + # embedding layer outside of the CUDA graph all the time. The v0 + # engine avoids this by "double compiling" the CUDA graph, once + # with input_ids and again with inputs_embeds, for all num_tokens. + # If a batch only has token ids, then including the embedding layer + # in the CUDA graph will be more performant (like in the else case + # below). + is_token_ids = self.is_token_ids.np[:num_scheduled_tokens] + token_ids_idx_np = np.nonzero(is_token_ids)[0] + # Some tokens ids may need to become embeds + if token_ids_idx_np.size > 0: + token_ids_idx = torch.from_numpy(token_ids_idx_np) + token_ids_idx = token_ids_idx.to(self.device, non_blocking=True) + token_ids = self.input_ids.gpu[token_ids_idx] + tokens_to_embeds = self.model.embed_input_ids(input_ids=token_ids) + self.inputs_embeds.gpu[token_ids_idx] = tokens_to_embeds + + inputs_embeds = self.inputs_embeds.gpu[:num_input_tokens] + model_kwargs = self._init_model_kwargs() + input_ids = None + else: + # For text-only models, we use token ids as input. + # While it is possible to use embeddings as input just like the + # multimodal models, it is not desirable for performance since + # then the embedding layer is not included in the CUDA graph. + input_ids = self.input_ids.gpu[:num_input_tokens] + inputs_embeds = None + model_kwargs = self._init_model_kwargs() + + if self.uses_mrope: + positions = self.mrope_positions.gpu[:, :num_input_tokens] + elif self.uses_xdrope_dim > 0: + positions = self.xdrope_positions.gpu[:, :num_input_tokens] + else: + positions = self.positions[:num_input_tokens] + if num_input_tokens > num_scheduled_tokens: + self.positions[num_scheduled_tokens:num_input_tokens].zero_() + + if is_first_rank: + intermediate_tensors = None + else: + assert intermediate_tensors is not None + intermediate_tensors = self.sync_and_gather_intermediate_tensors( + num_input_tokens, intermediate_tensors, True + ) + + if is_encoder_decoder and scheduler_output.scheduled_encoder_inputs: + # Run the encoder, just like we do with other multimodal inputs. + # For an encoder-decoder model, our processing here is a bit + # simpler, because the outputs are just passed to the decoder. + # We are not doing any prompt replacement. We also will only + # ever have a single encoder input. + encoder_outputs = self._execute_mm_encoder(scheduler_output) + model_kwargs.update({"encoder_outputs": encoder_outputs}) + + return ( + input_ids, + inputs_embeds, + positions, + intermediate_tensors, + model_kwargs, + ec_connector_output, + ) + + def _sample( + self, + logits: torch.Tensor | None, + spec_decode_metadata: SpecDecodeMetadata | None, + ) -> SamplerOutput: + # Sample the next token and get logprobs if needed. + sampling_metadata = self.input_batch.sampling_metadata + # Update output token ids with tokens sampled in last step + # if async scheduling and required by current sampling params. + self.input_batch.update_async_output_token_ids() + if spec_decode_metadata is None: + return self.sampler( + logits=logits, + sampling_metadata=sampling_metadata, + ) + + # Update spec_token_ids with real draft tokens from pre step only when + # output_token_ids is needed (penalties or bad_words are in use). + if self.use_async_scheduling and self._draft_token_req_ids is not None: + draft_token_ids_cpu, _ = self._get_draft_token_ids_cpu() + self.input_batch.update_async_spec_token_ids(draft_token_ids_cpu) + + self._maybe_observe_dspark_position0_quality( + spec_decode_metadata, + logits, + ) + draft_probs = self._get_spec_decode_draft_probs(spec_decode_metadata) + sampler_output = self.rejection_sampler( + spec_decode_metadata, + draft_probs, + logits, + sampling_metadata, + ) + return sampler_output + + def _get_dspark_position0_confidence( + self, + req_ids: list[str], + ) -> torch.Tensor | None: + if self._draft_confidence is None or self._draft_confidence_req_ids is None: + return None + + row_by_req_id = { + req_id: idx for idx, req_id in enumerate(self._draft_confidence_req_ids) + } + confidence_rows: list[torch.Tensor] = [] + for req_id in req_ids: + row_idx = row_by_req_id.get(req_id) + if row_idx is None: + return None + confidence_rows.append(self._draft_confidence[row_idx, 0]) + if not confidence_rows: + return None + return torch.stack(confidence_rows).float() + + def _maybe_observe_dspark_position0_quality( + self, + spec_decode_metadata: SpecDecodeMetadata, + logits: torch.Tensor | None, + ) -> None: + diagnostics = self._dspark_position0_diagnostics + if diagnostics is None or logits is None: + return + + first_offsets: list[int] = [] + active_req_ids: list[str] = [] + offset = 0 + for req_id, num_draft in zip( + self.input_batch.req_ids, + spec_decode_metadata.num_draft_tokens, + ): + if num_draft > 0: + first_offsets.append(offset) + active_req_ids.append(req_id) + offset += num_draft + if not first_offsets: + return + + first_offsets_tensor = torch.tensor( + first_offsets, + dtype=torch.long, + device=spec_decode_metadata.draft_token_ids.device, + ) + first_target_indices = spec_decode_metadata.target_logits_indices.index_select( + 0, + first_offsets_tensor, + ).long() + first_draft_ids = spec_decode_metadata.draft_token_ids.index_select( + 0, + first_offsets_tensor, + ).long() + target_argmax = logits.index_select( + 0, + first_target_indices, + ).argmax(dim=-1) + matches = target_argmax.eq(first_draft_ids) + confidence = self._get_dspark_position0_confidence(active_req_ids) + + confidences = ( + None if confidence is None else confidence.detach().cpu().tolist() + ) + diagnostics.observe( + matches.detach().cpu().tolist(), + confidences, + ) + snapshot = diagnostics.snapshot() + if snapshot.num_tokens < self._dspark_position0_log_next: + return + logger.info( + "DSpark position-0 diagnostics: samples=%d, " + "target_argmax_match_rate=%.3f, avg_confidence=%s, " + "avg_confidence_matched=%s, avg_confidence_missed=%s, " + "confidence_logits_normalized=%d", + snapshot.num_tokens, + snapshot.match_rate, + _format_optional_float(snapshot.avg_confidence), + _format_optional_float(snapshot.avg_confidence_when_matched), + _format_optional_float(snapshot.avg_confidence_when_missed), + snapshot.num_confidence_logits_normalized, + ) + self._dspark_position0_log_next = snapshot.num_tokens + 64 + + def _bookkeeping_sync( + self, + scheduler_output: "SchedulerOutput", + sampler_output: SamplerOutput, + logits: torch.Tensor | None, + hidden_states: torch.Tensor, + num_scheduled_tokens: int, + ) -> tuple[ + dict[str, int], + LogprobsLists | None, + list[list[int]], + dict[str, LogprobsTensors | None], + list[str], + dict[str, int], + list[int], + ]: + num_nans_in_logits = {} + if envs.VLLM_COMPUTE_NANS_IN_LOGITS: + num_nans_in_logits = self._get_nans_in_logits(logits) + + num_reqs = self.input_batch.num_reqs + discard_sampled_tokens_req_indices = np.nonzero( + self.discard_request_mask.np[:num_reqs] + )[0] + for i in discard_sampled_tokens_req_indices: + gen = self.input_batch.generators.get(int(i)) + if gen is not None: + gen.set_offset(gen.get_offset() - 4) + + # Copy some objects so they don't get modified after returning. + # This is important when using async scheduling. + req_ids_output_copy = self.input_batch.req_ids.copy() + req_id_to_index_output_copy = self.input_batch.req_id_to_index.copy() + + num_sampled_tokens = sampler_output.sampled_token_ids.shape[0] + sampled_token_ids = sampler_output.sampled_token_ids + logprobs_tensors = sampler_output.logprobs_tensors + invalid_req_indices = [] + logprobs_lists = None + if not self.use_async_scheduling: + # Sync scheduling: issue routed experts D2H into the pinned + # CPU buffer BEFORE ``_to_list`` below. ``_to_list`` does + # ``event.synchronize()`` on the async copy stream which + # waits for every D2H queued on the default stream since + # the last sync, so this enqueue is naturally covered + # without requiring its own synchronize. + if self.routed_experts_initialized: + buf = self.routed_experts_capturer.get_device_buffer() + total = scheduler_output.total_num_scheduled_tokens + self.routed_experts_cpu[:total].copy_(buf[:total], non_blocking=True) + self.routed_experts_slot_mapping_cpu[:total].copy_( + self.routed_experts_slot_mapping_device[:total], + non_blocking=True, + ) + + # Get the valid generated tokens. + max_gen_len = sampled_token_ids.shape[-1] + if max_gen_len == 1: + # No spec decode tokens. + valid_sampled_token_ids = self._to_list(sampled_token_ids) + # Mask out the sampled tokens that should not be sampled. + for i in discard_sampled_tokens_req_indices: + valid_sampled_token_ids[int(i)].clear() + + if logprobs_tensors is not None: + logprobs_lists = logprobs_tensors.tolists() + else: + # Includes spec decode tokens. + valid_sampled_token_ids, logprobs_lists = RejectionSampler.parse_output( + sampled_token_ids, + self.input_batch.vocab_size, + discard_sampled_tokens_req_indices, + logprobs_tensors=logprobs_tensors, + ) + else: + valid_sampled_token_ids = [] + invalid_req_indices = discard_sampled_tokens_req_indices.tolist() + invalid_req_indices_set = set(invalid_req_indices) + + # Cache the sampled tokens on the GPU and avoid CPU sync. + # These will be copied into input_ids in the next step + # when preparing inputs. + # With spec decoding, this is done in propose_draft_token_ids(). + if self.input_batch.prev_sampled_token_ids is None: + assert sampled_token_ids.shape[-1] == 1 + self.input_batch.prev_sampled_token_ids = sampled_token_ids + self.input_batch.prev_req_id_to_index = { + req_id: i + for i, req_id in enumerate(self.input_batch.req_ids) + if i not in invalid_req_indices_set + } + + # Cache the sampled tokens in the model runner, so that the scheduler + # doesn't need to send them back. + # NOTE(woosuk): As an exception, when using PP, the scheduler sends + # the sampled tokens back, because there's no direct communication + # between the first-stage worker and the last-stage worker. + req_ids = self.input_batch.req_ids + for req_idx in range(num_sampled_tokens): + if self.use_async_scheduling: + sampled_ids = [-1] if req_idx not in invalid_req_indices_set else None + else: + sampled_ids = valid_sampled_token_ids[req_idx] + + num_sampled_ids: int = len(sampled_ids) if sampled_ids else 0 + + if not sampled_ids: + continue + + start_idx = self.input_batch.num_tokens_no_spec[req_idx] + end_idx = start_idx + num_sampled_ids + assert end_idx <= self.max_model_len, ( + "Sampled token IDs exceed the max model length. " + f"Total number of tokens: {end_idx} > max_model_len: " + f"{self.max_model_len}" + ) + + self.input_batch.token_ids_cpu[req_idx, start_idx:end_idx] = sampled_ids + self.input_batch.is_token_ids[req_idx, start_idx:end_idx] = True + self.input_batch.num_tokens_no_spec[req_idx] = end_idx + + req_id = req_ids[req_idx] + req_state = self.requests[req_id] + req_state.output_token_ids.extend(sampled_ids) + + # Compute prompt logprobs if needed. + prompt_logprobs_dict = self._get_prompt_logprobs_dict( + hidden_states[:num_scheduled_tokens], + scheduler_output.num_scheduled_tokens, + ) + + return ( + num_nans_in_logits, + logprobs_lists, + valid_sampled_token_ids, + prompt_logprobs_dict, + req_ids_output_copy, + req_id_to_index_output_copy, + invalid_req_indices, + ) + + @contextmanager + def synchronize_input_prep(self): + if self.prepare_inputs_event is None: + yield + return + + # Ensure prior step has finished with reused CPU tensors. + # This is required in the async scheduling case because + # the CPU->GPU transfer happens async. + self.prepare_inputs_event.synchronize() + try: + yield + finally: + self.prepare_inputs_event.record() + + def _model_forward( + self, + input_ids: torch.Tensor | None = None, + positions: torch.Tensor | None = None, + intermediate_tensors: IntermediateTensors | None = None, + inputs_embeds: torch.Tensor | None = None, + **model_kwargs: dict[str, Any], + ) -> Any: + """Helper method to call the model forward pass. + + This method can be overridden by subclasses for model execution. + Motivation: We can inspect only this method versus + the whole execute_model, which has additional logic. + + Args: + input_ids: Input token IDs + positions: Token positions + intermediate_tensors: Tensors from previous pipeline stages + inputs_embeds: Input embeddings (alternative to input_ids) + **model_kwargs: Additional model arguments + + Returns: + Model output tensor + """ + return self.model( + input_ids=input_ids, + positions=positions, + intermediate_tensors=intermediate_tensors, + inputs_embeds=inputs_embeds, + **model_kwargs, + ) + + @staticmethod + def _is_uniform_decode( + max_num_scheduled_tokens: int, + uniform_decode_query_len: int, + num_tokens: int, + num_reqs: int, + force_uniform_decode: bool | None = None, + ) -> bool: + """ + Checks if it's a decode batch with same amount scheduled tokens + across all requests. + """ + return ( + ( + (max_num_scheduled_tokens == uniform_decode_query_len) + and (num_tokens == max_num_scheduled_tokens * num_reqs) + ) + if force_uniform_decode is None + else force_uniform_decode + ) + + def _determine_batch_execution_and_padding( + self, + num_tokens: int, + num_reqs: int, + num_scheduled_tokens_np: np.ndarray, + max_num_scheduled_tokens: int, + use_cascade_attn: bool, + allow_microbatching: bool = True, + force_eager: bool = False, + # For cudagraph capture TODO(lucas): Refactor how we capture cudagraphs (will + # be improved in model runner v2) + force_uniform_decode: bool | None = None, + force_has_lora: bool | None = None, + force_num_active_loras: int | None = None, + num_encoder_reqs: int = 0, + ) -> tuple[ + CUDAGraphMode, + BatchDescriptor, + bool, + torch.Tensor | None, + CUDAGraphStat | None, + ]: + uniform_decode = self._is_uniform_decode( + max_num_scheduled_tokens=max_num_scheduled_tokens, + uniform_decode_query_len=self.uniform_decode_query_len, + num_tokens=num_tokens, + num_reqs=num_reqs, + force_uniform_decode=force_uniform_decode, + ) + # Encoder-decoder models only support CG for decoder_step > 0 (no enc_output + # is present). Also, chunked-prefill is disabled, so batch are uniform. + has_encoder_output = ( + self.model_config.is_encoder_decoder and num_encoder_reqs > 0 + ) + + # Compute LoRA state for cudagraph dispatch + num_active_loras = ( + force_num_active_loras + if force_num_active_loras is not None + else len(self.input_batch.lora_id_to_lora_request) + ) + 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, + has_lora=has_lora, + uniform_decode=uniform_decode, + num_active_loras=num_active_loras, + valid_modes={CUDAGraphMode.NONE} if force_eager else valid_modes, + invalid_modes={CUDAGraphMode.FULL} if disable_full else None, + ) + + cudagraph_mode, batch_descriptor = dispatch_cudagraph( + num_tokens_padded, 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 + should_ubatch, num_tokens_across_dp = False, None + if self.vllm_config.parallel_config.data_parallel_size > 1: + should_ubatch, num_tokens_across_dp, synced_cudagraph_mode = ( + coordinate_batch_across_dp( + num_tokens_unpadded=num_tokens, + parallel_config=self.parallel_config, + allow_microbatching=allow_microbatching, + num_tokens_padded=num_tokens_padded, + uniform_decode=uniform_decode, + cudagraph_mode=cudagraph_mode.value, + ) + ) + + # Extract DP-synced values + if num_tokens_across_dp is not None: + dp_rank = self.parallel_config.data_parallel_rank + num_tokens_padded = int(num_tokens_across_dp[dp_rank].item()) + # Re-dispatch with DP padding so we have the correct batch_descriptor + cudagraph_mode, batch_descriptor = dispatch_cudagraph( + num_tokens_padded, + valid_modes={CUDAGraphMode(synced_cudagraph_mode)}, + ) + # Assert to make sure the agreed upon token count is correct otherwise + # num_tokens_across_dp will no-longer be valid + assert batch_descriptor.num_tokens == num_tokens_padded + + cudagraph_stats = None + if self.vllm_config.observability_config.cudagraph_metrics: + cudagraph_stats = CUDAGraphStat( + num_unpadded_tokens=num_tokens, + num_padded_tokens=batch_descriptor.num_tokens, + num_paddings=batch_descriptor.num_tokens - num_tokens, + runtime_mode=str(cudagraph_mode), + ) + + return ( + cudagraph_mode, + batch_descriptor, + should_ubatch, + num_tokens_across_dp, + cudagraph_stats, + ) + + def _register_layerwise_nvtx_hooks(self) -> None: + """ + Register layerwise NVTX hooks if --enable-layerwise-nvtx-tracing is enabled + to trace detailed information of each layer or module in the model. + """ + + if ( + self.vllm_config.observability_config.enable_layerwise_nvtx_tracing + and not self.layerwise_nvtx_hooks_registered + ): + if self.compilation_config.cudagraph_mode != CUDAGraphMode.NONE: + logger.debug_once( + "layerwise NVTX tracing is not supported when CUDA graph is " + "turned off; you may observe part or all of the model " + "missing NVTX markers" + ) + + # In STOCK_TORCH_COMPILE mode, after registering hooks here, + # the __call__ function of nn.module will be recompiled with + # fullgraph=True. Since nvtx.range_push/pop are not traceable + # by torch dynamo, we can't register hook functions here + # because hook functions will also be traced by torch dynamo. + if ( + self.vllm_config.compilation_config.mode + == CompilationMode.STOCK_TORCH_COMPILE + ): + logger.debug_once( + "layerwise NVTX tracing is not supported when " + "CompilationMode is STOCK_TORCH_COMPILE, skipping " + "function hooks registration" + ) + else: + pyt_hooks = PytHooks() + pyt_hooks.register_hooks(self.model, self.model.__class__.__name__) + self.layerwise_nvtx_hooks_registered = True + + def _get_slot_mappings( + self, + num_tokens_padded: int, + num_reqs_padded: int, + num_tokens_unpadded: int, + ubatch_slices: "UBatchSlices | None" = None, + ) -> tuple[ + dict[int, torch.Tensor] | None, + dict[str, torch.Tensor] | list[dict[str, torch.Tensor]] | None, + ]: + """ + Build slot mappings in both formats needed by the system. + + Args: + num_tokens_padded: Total number of tokens (padded) + num_reqs_padded: Total number of requests (padded) + num_tokens_unpadded: Actual number of tokens (unpadded) + ubatch_slices: Optional ubatch slicing info for DBO + + Returns: + A tuple of: + - slot_mappings_by_gid: dict[int, torch.Tensor] for attention metadata + - slot_mappings_by_layer: dict[str, torch.Tensor] or list for ForwardContext + """ + if not ( + hasattr(self, "kv_cache_config") + and self.kv_cache_config is not None + and len(self.kv_cache_config.kv_cache_groups) > 0 + ): + return None, None + + def _get_slot_mapping(kv_cache_gid: int): + assert num_reqs_padded is not None and num_tokens_padded is not None + kv_cache_spec = self.kv_cache_config.kv_cache_groups[ + kv_cache_gid + ].kv_cache_spec + if isinstance(kv_cache_spec, EncoderOnlyAttentionSpec): + slot_mapping = torch.zeros( + (num_tokens_padded,), + dtype=torch.int64, + device=self.device, + ) + else: + blk_table = self.input_batch.block_table[kv_cache_gid] + slot_mapping = blk_table.slot_mapping.gpu[:num_tokens_padded] + + # Fill unused with -1. Needed for reshape_and_cache in full cuda + # graph mode. `blk_table_tensor` -1 to match mamba PAD_SLOT_ID + slot_mapping[num_tokens_unpadded:num_tokens_padded].fill_(-1) + + return slot_mapping + + slot_mappings_by_gid = { + gid: _get_slot_mapping(gid) + for gid, _ in enumerate(self.kv_cache_config.kv_cache_groups) + } + + slot_mappings_by_layer: dict[str, torch.Tensor] = {} + for gid, kv_cache_group in enumerate(self.kv_cache_config.kv_cache_groups): + slot_mapping = slot_mappings_by_gid[gid] + for layer_name in kv_cache_group.layer_names: + slot_mappings_by_layer[layer_name] = slot_mapping + + if ubatch_slices is not None: + result: list[dict[str, torch.Tensor]] = [] + for ubatch in ubatch_slices: + sliced_mappings: dict[str, torch.Tensor] = {} + for layer_name, slot_mapping in slot_mappings_by_layer.items(): + sliced_mappings[layer_name] = slot_mapping[ubatch.token_slice] + result.append(sliced_mappings) + return slot_mappings_by_gid, result + + return slot_mappings_by_gid, slot_mappings_by_layer + + def _is_all_reqs_chunked_prefill(self) -> bool: + """Check if all scheduled requests are marked to discard sampled tokens. + + This is true when `discard_request_mask` is set for every scheduled + request (e.g., for chunked prefill requests that are not the last + prefill chunk).""" + num_reqs = self.input_batch.num_reqs + return bool(self.discard_request_mask.np[:num_reqs].all()) + + def _dspark_timing_start(self) -> float: + if not self._dspark_iter_timing: + return 0.0 + if self.device.type == "cuda": + torch.cuda.synchronize() + return time.perf_counter() + + def _dspark_timing_record(self, name: str, started: float) -> None: + if not self._dspark_iter_timing or started == 0.0: + return + if self.device.type == "cuda": + torch.cuda.synchronize() + self._dspark_iter_timing_totals_ms[name] += ( + time.perf_counter() - started + ) * 1000.0 + + def _dspark_timing_finish_iteration(self) -> None: + if not self._dspark_iter_timing: + return + self._dspark_timing_record("iter_total", self._dspark_iter_timing_started) + self._dspark_iter_timing_started = 0.0 + self._dspark_iter_timing_count += 1 + if self._dspark_iter_timing_count % self._dspark_iter_timing_log_every != 0: + return + + names = ( + "execute_preprocess", + "target_forward", + "target_postprocess_logits", + "sample_reject", + "state_update", + "draft_propose", + "bookkeeping", + "finalize_output", + "iter_total", + ) + parts = [] + for name in names: + total_ms = self._dspark_iter_timing_totals_ms.get(name, 0.0) + parts.append(f"{name}={total_ms / self._dspark_iter_timing_count:.3f}ms") + logger.info( + "DSpark iteration timing avg over %d iterations: %s", + self._dspark_iter_timing_count, + ", ".join(parts), + ) + + @torch.inference_mode() + def execute_model( + self, + scheduler_output: "SchedulerOutput", + intermediate_tensors: IntermediateTensors | None = None, + ) -> ModelRunnerOutput | AsyncModelRunnerOutput | IntermediateTensors | None: + if self.execute_model_state is not None: + raise RuntimeError( + "State error: sample_tokens() must be called " + "after execute_model() returns None." + ) + + if self.routed_experts_initialized: + self.routed_experts_capturer.clear_buffer() + + self._dspark_iter_timing_started = self._dspark_timing_start() + + # If ngram_gpu is used, we need to copy the scheduler_output to avoid + # the modification has influence on the scheduler_output in engine core process. + # The replace is much faster than deepcopy. + if ( + self.speculative_config is not None + and self.speculative_config.use_ngram_gpu() + ): + num_scheduled_tokens_copy = scheduler_output.num_scheduled_tokens.copy() + spec_decode_tokens_copy = ( + scheduler_output.scheduled_spec_decode_tokens.copy() + ) + scheduler_output = replace( + scheduler_output, + num_scheduled_tokens=num_scheduled_tokens_copy, + scheduled_spec_decode_tokens=spec_decode_tokens_copy, + ) + + if has_kv_transfer_group(): + kv_connector_metadata = scheduler_output.kv_connector_metadata + assert kv_connector_metadata is not None + get_kv_transfer_group().handle_preemptions(kv_connector_metadata) + + num_scheduled_tokens = scheduler_output.total_num_scheduled_tokens + stage_started = self._dspark_timing_start() + with ( + record_function_or_nullcontext("gpu_model_runner: preprocess"), + self.synchronize_input_prep(), + ): + # Update persistent batch states. + deferred_state_corrections_fn = self._update_states(scheduler_output) + + if has_ec_transfer() and not get_ec_transfer().is_consumer: + with self.maybe_get_ec_connector_output( + scheduler_output, + encoder_cache=self.encoder_cache, + ) as ec_connector_output: + self._execute_mm_encoder(scheduler_output) + return make_empty_encoder_model_runner_output(scheduler_output) + + if not num_scheduled_tokens: + if ( + self.parallel_config.distributed_executor_backend + == "external_launcher" + and self.parallel_config.data_parallel_size > 1 + ): + # this is a corner case when both external launcher + # and DP are enabled, num_scheduled_tokens could be + # 0, and has_unfinished_requests in the outer loop + # returns True. before returning early here we call + # dummy run to ensure coordinate_batch_across_dp + # is called into to avoid out of sync issues. + self._dummy_run(1) + if not has_kv_transfer_group(): + # Return empty ModelRunnerOutput if no work to do. + return EMPTY_MODEL_RUNNER_OUTPUT + return self.kv_connector_no_forward(scheduler_output, self.vllm_config) + + if self.cache_config.kv_sharing_fast_prefill: + assert not self.num_prompt_logprobs, ( + "--kv-sharing-fast-prefill produces incorrect " + "logprobs for prompt tokens, tokens, please disable " + "it when the requests need prompt logprobs" + ) + + num_reqs = self.input_batch.num_reqs + req_ids = self.input_batch.req_ids + tokens = [scheduler_output.num_scheduled_tokens[i] for i in req_ids] + num_scheduled_tokens_np = np.array(tokens, dtype=np.int32) + max_num_scheduled_tokens = int(num_scheduled_tokens_np.max()) + num_tokens_unpadded = scheduler_output.total_num_scheduled_tokens + + logits_indices, spec_decode_metadata = 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, + should_ubatch, + num_tokens_across_dp, + cudagraph_stats, + ) = self._determine_batch_execution_and_padding( + num_tokens=num_tokens_unpadded, + 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, + num_encoder_reqs=len(scheduler_output.scheduled_encoder_inputs), + ) + + logger.debug( + "Running batch with cudagraph_mode: %s, batch_descriptor: %s, " + "should_ubatch: %s, num_tokens_across_dp: %s", + cudagraph_mode, + batch_desc, + should_ubatch, + num_tokens_across_dp, + ) + + num_tokens_padded = batch_desc.num_tokens + num_reqs_padded = ( + batch_desc.num_reqs if batch_desc.num_reqs is not None else num_reqs + ) + ubatch_slices, ubatch_slices_padded = maybe_create_ubatch_slices( + should_ubatch, + num_scheduled_tokens_np, + num_tokens_padded, + num_reqs_padded, + self.parallel_config.num_ubatches, + ) + + logger.debug( + "ubatch_slices: %s, ubatch_slices_padded: %s", + ubatch_slices, + ubatch_slices_padded, + ) + + # True if any attention backend handles KV cache update separately + # from forward() (i.e., forward_includes_kv_cache_update=False). When true, + # slot_mappings must use padded dimensions to match the key/value tensors. + has_separate_kv_update = not all( + all( + g.backend.forward_includes_kv_cache_update + for g in self.attn_groups[id] + ) + for id, spec in enumerate(self.kv_cache_config.kv_cache_groups) + if not isinstance(spec.kv_cache_spec, EncoderOnlyAttentionSpec) + ) + pad_attn = cudagraph_mode == CUDAGraphMode.FULL + + if self.cache_config.mamba_cache_mode == "align": + # preprocess_mamba reads req_state.num_computed_tokens (CPU) + # to decide copy operations, so we must apply deferred + # corrections before it runs. + if deferred_state_corrections_fn: + deferred_state_corrections_fn() + deferred_state_corrections_fn = None + mamba_bufs = self._get_mamba_bufs() + mamba_utils.preprocess_mamba( + scheduler_output, + self.kv_cache_config, + self.cache_config, + self.mamba_state_idx, + self.input_batch, + self.requests, + self.compilation_config.static_forward_context, + self.model.get_mamba_state_copy_func(), + mamba_bufs.preprocess, + ) + # preprocess_mamba resets num_accepted_tokens_cpu to 1 + # for requests whose state was copied to a new block. + # Re-sync to GPU so the mamba kernel reads from the + # correct initial state slot (init_token_idx = 0). + self.num_accepted_tokens.np[:num_reqs] = ( + self.input_batch.num_accepted_tokens_cpu[:num_reqs] + ) + self.num_accepted_tokens.copy_to_gpu(num_reqs) + + # Stage per-request inputs for the fused postprocess kernel + # only when that kernel will actually run. The kernel is + # gated on spec-decode + hybrid (see MambaBuffers.create); + # without it, ``mamba_bufs.postprocess_align`` is None and + # the staging buffers don't exist. + if mamba_bufs.postprocess_align is not None: + mamba_utils.stage_postprocess_inputs_to_gpu( + mamba_bufs.postprocess_align, + scheduler_output, + self.input_batch.req_ids, + num_reqs, + self.requests, + self.mamba_state_idx, + ) + + use_spec_decode = len(scheduler_output.scheduled_spec_decode_tokens) > 0 + ubatch_slices_attn = ubatch_slices_padded if pad_attn else ubatch_slices + + slot_mappings_by_group, slot_mappings = self._get_slot_mappings( + num_tokens_padded=num_tokens_padded + if pad_attn or has_separate_kv_update + else num_tokens_unpadded, + num_reqs_padded=( + num_reqs_padded if pad_attn or has_separate_kv_update else num_reqs + ), + num_tokens_unpadded=num_tokens_unpadded, + ubatch_slices=ubatch_slices_padded, + ) + + attn_metadata, spec_decode_common_attn_metadata = ( + self._build_attention_metadata( + num_tokens=num_tokens_unpadded, + num_tokens_padded=num_tokens_padded if pad_attn else None, + num_reqs=num_reqs, + num_reqs_padded=num_reqs_padded if pad_attn else None, + max_query_len=max_num_scheduled_tokens, + ubatch_slices=ubatch_slices_attn, + logits_indices=logits_indices, + 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, + ) + ) + + ( + input_ids, + inputs_embeds, + positions, + intermediate_tensors, + model_kwargs, + ec_connector_output, + ) = self._preprocess( + scheduler_output, num_tokens_padded, intermediate_tensors + ) + self._dspark_timing_record("execute_preprocess", stage_started) + + # Set cudagraph mode to none if calc_kv_scales is true. + # KV scales calculation involves dynamic operations that are incompatible + # with CUDA graph capture. + if self.calculate_kv_scales: + cudagraph_mode = CUDAGraphMode.NONE + # Mark KV scales as calculated after the first forward pass + self.calculate_kv_scales = False + + # Encoder-decoder models can only compile the pure decode steps where no + # encoder inputs are present. Use eager for the first pass. + num_encoder_reqs = len(scheduler_output.scheduled_encoder_inputs) + has_encoder_input = ( + self.model_config.is_encoder_decoder and num_encoder_reqs > 0 + ) + + # Run the model. + # Use persistent buffers for CUDA graphs. + # When spec decode is enabled, defer connector finalization + # (wait_for_save + clear metadata) until after draft model runs. + defer_kv_connector_finalize = self.speculative_config is not None + stage_started = self._dspark_timing_start() + with ( + set_forward_context( + attn_metadata, + self.vllm_config, + num_tokens=num_tokens_padded, + num_tokens_across_dp=num_tokens_across_dp, + cudagraph_runtime_mode=cudagraph_mode, + batch_descriptor=batch_desc, + ubatch_slices=ubatch_slices_padded, + slot_mapping=slot_mappings, + skip_compiled=has_encoder_input, + ), + record_function_or_nullcontext("gpu_model_runner: forward"), + self.maybe_get_kv_connector_output( + scheduler_output, + defer_finalize=defer_kv_connector_finalize, + ) as kv_connector_output, + ): + model_output = self._model_forward( + input_ids=input_ids, + positions=positions, + intermediate_tensors=intermediate_tensors, + inputs_embeds=inputs_embeds, + **model_kwargs, + ) + self._dspark_timing_record("target_forward", stage_started) + + stage_started = self._dspark_timing_start() + with record_function_or_nullcontext("gpu_model_runner: postprocess"): + if self.use_aux_hidden_state_outputs: + # True when EAGLE 3 is used. + hidden_states, aux_hidden_states = model_output + else: + # Common case. + hidden_states = model_output + aux_hidden_states = None + + if not self.broadcast_pp_output: + # Common case. + if not get_pp_group().is_last_rank: + # Return the intermediate tensors. + assert isinstance(hidden_states, IntermediateTensors) + hidden_states.kv_connector_output = kv_connector_output + self.kv_connector_output = kv_connector_output + return hidden_states + + if self.is_pooling_model: + # Return the pooling output. + return self._pool( + hidden_states, + num_scheduled_tokens, + num_scheduled_tokens_np, + kv_connector_output, + ) + + sample_hidden_states = hidden_states[logits_indices] + logits = self.model.compute_logits(sample_hidden_states) + else: + # Rare case. + assert not self.is_pooling_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: + logits = self.model.compute_logits(sample_hidden_states) + + model_output_broadcast_data: dict[str, Any] = {} + if logits is not None: + model_output_broadcast_data["logits"] = logits.contiguous() + + broadcasted = get_pp_group().broadcast_tensor_dict( + model_output_broadcast_data, src=len(get_pp_group().ranks) - 1 + ) + assert broadcasted is not None + logits = broadcasted["logits"] + self._dspark_timing_record("target_postprocess_logits", stage_started) + + self.execute_model_state = ExecuteModelState( + scheduler_output, + logits, + spec_decode_metadata, + spec_decode_common_attn_metadata, + hidden_states, + sample_hidden_states, + aux_hidden_states, + ec_connector_output, + cudagraph_stats, + slot_mappings, + ) + self.kv_connector_output = kv_connector_output + + # Now the batch has been launched we can wait for corrections from the + # previous model forward without breaking async scheduling. + if deferred_state_corrections_fn: + deferred_state_corrections_fn() + + return None + + @torch.inference_mode + def sample_tokens( + self, grammar_output: "GrammarOutput | None" + ) -> ModelRunnerOutput | AsyncModelRunnerOutput | IntermediateTensors: + if self.execute_model_state is None: + kv_connector_output = self.kv_connector_output + self.kv_connector_output = None + # receive sampled token ids from the last PP rank. + if self.use_async_scheduling and not get_pp_group().is_last_rank: + self._pp_receive_prev_sampled_token_ids_to_input_batch() + if not kv_connector_output: + return None # type: ignore[return-value] + + # In case of PP with kv transfer, we need to pass through the + # kv_connector_output + if kv_connector_output.is_empty(): + return EMPTY_MODEL_RUNNER_OUTPUT + + output = copy(EMPTY_MODEL_RUNNER_OUTPUT) + output.kv_connector_output = kv_connector_output + return output + + # Unpack ephemeral state. + ( + scheduler_output, + logits, + spec_decode_metadata, + spec_decode_common_attn_metadata, + hidden_states, + sample_hidden_states, + aux_hidden_states, + ec_connector_output, + cudagraph_stats, + slot_mappings, + ) = self.execute_model_state + # Clear ephemeral state. + self.execute_model_state = None + + # Apply structured output bitmasks if present. + if grammar_output is not None: + apply_grammar_bitmask( + scheduler_output, grammar_output, self.input_batch, logits + ) + + stage_started = self._dspark_timing_start() + with record_function_or_nullcontext("gpu_model_runner: sample"): + sampler_output = self._sample(logits, spec_decode_metadata) + self._dspark_timing_record("sample_reject", stage_started) + + stage_started = self._dspark_timing_start() + self._update_states_after_model_execute( + sampler_output.sampled_token_ids, scheduler_output + ) + if self.use_async_scheduling: + pp = get_pp_group() + # For torchrun external_launcher PP mode with broadcast_pp_output=True, + # PP outputs have been broadcasted to all ranks at logits computation. + # Therefore, here is no need to send sampled token ids again in this case. + if not self.broadcast_pp_output and pp.world_size > 1 and pp.is_last_rank: + self._pp_broadcast_prev_sampled_token_ids( + sampler_output.sampled_token_ids + ) + self._dspark_timing_record("state_update", stage_started) + + self._draft_token_ids = None + self._draft_token_lengths_cpu = None + self._draft_token_length_req_ids = None + self._draft_probs = None + self._draft_prob_req_ids = None + self._draft_confidence = None + self._draft_confidence_req_ids = None + self._draft_token_req_ids = None + self.valid_sampled_token_count_gpu = None + self.input_batch.prev_sampled_token_ids = None + + def propose_draft_token_ids(sampled_token_ids): + assert spec_decode_common_attn_metadata is not None + stage_started = self._dspark_timing_start() + with record_function_or_nullcontext("gpu_model_runner: draft"): + self._draft_token_ids = self.propose_draft_token_ids( + scheduler_output, + sampled_token_ids, + self.input_batch.sampling_metadata, + hidden_states, + sample_hidden_states, + aux_hidden_states, + spec_decode_metadata, + spec_decode_common_attn_metadata, + slot_mappings, + ) + self._copy_draft_token_ids_to_cpu(scheduler_output) + self._dspark_timing_record("draft_propose", stage_started) + + spec_config = self.speculative_config + propose_drafts_after_bookkeeping = False + if spec_config is not None: + # Decide whether to run the drafter or zero out draft tokens. + input_fits_in_drafter = spec_decode_common_attn_metadata is not None and ( + spec_decode_common_attn_metadata.max_seq_len + self.num_spec_tokens + <= self.effective_drafter_max_model_len + ) + use_gpu_toks = ( + spec_config.use_eagle() + or spec_config.uses_draft_model() + or spec_config.uses_extract_hidden_states() + or spec_config.use_dspark() + ) and not spec_config.disable_padded_drafter_batch + if use_gpu_toks: + # EAGLE/DraftModel speculative decoding can use the GPU sampled tokens + # as inputs, and does not need to wait for bookkeeping to finish. + assert isinstance( + self.drafter, + EagleProposer + | DFlashProposer + | DSparkProposer + | DraftModelProposer + | ExtractHiddenStatesProposer + | Gemma4Proposer, + ) + sampled_token_ids = sampler_output.sampled_token_ids + if input_fits_in_drafter: + propose_draft_token_ids(sampled_token_ids) + elif self.valid_sampled_token_count_event is not None: + assert spec_decode_common_attn_metadata is not None + next_token_ids, valid_sampled_tokens_count = ( + self.drafter.prepare_next_token_ids_padded( + sampled_token_ids, + self.requests, + self.input_batch, + self.discard_request_mask.gpu, + ) + ) + self._copy_valid_sampled_token_count( + next_token_ids, valid_sampled_tokens_count + ) + elif ( + spec_config.use_ngram_gpu() + and not spec_config.disable_padded_drafter_batch + ): + assert isinstance(self.drafter, NgramProposerGPU) + sampled_token_ids = sampler_output.sampled_token_ids + if input_fits_in_drafter: + propose_draft_token_ids(sampled_token_ids) + elif self.valid_sampled_token_count_event is not None: + assert spec_decode_common_attn_metadata is not None + next_token_ids, valid_sampled_tokens_count, _ = ( + self.drafter.update_token_ids_ngram( + sampled_token_ids, + self.input_batch, + self.token_ids_gpu_tensor, + self.num_tokens_no_spec_gpu, + self.discard_request_mask.gpu, + ) + ) + self._copy_valid_sampled_token_count( + next_token_ids, valid_sampled_tokens_count + ) + else: + propose_drafts_after_bookkeeping = input_fits_in_drafter + + if not input_fits_in_drafter: + # Zero out draft tokens so the scheduler doesn't schedule + # stale drafts from the previous step. + # For Nemotron-H: it is necessary to zero out the draft tokens, + # otherwise the stale tokens will corrupt Mamba recurrent + # state and logprobs for sequences near max_model_len. + self._draft_token_ids = torch.zeros( + 1, device=self.device, dtype=torch.int32 + ).expand(len(self.input_batch.req_ids), self.num_spec_tokens) + self._draft_probs = None + self._draft_prob_req_ids = None + self._draft_confidence = None + self._draft_confidence_req_ids = None + self._copy_draft_token_ids_to_cpu(scheduler_output, zeros_only=True) + + stage_started = self._dspark_timing_start() + with record_function_or_nullcontext("gpu_model_runner: bookkeep"): + ( + num_nans_in_logits, + logprobs_lists, + valid_sampled_token_ids, + prompt_logprobs_dict, + req_ids_output_copy, + req_id_to_index_output_copy, + invalid_req_indices, + ) = self._bookkeeping_sync( + scheduler_output, + sampler_output, + logits, + hidden_states, + scheduler_output.total_num_scheduled_tokens, + ) + self._dspark_timing_record("bookkeeping", stage_started) + + if propose_drafts_after_bookkeeping: + # ngram and other speculative decoding methods use the sampled + # tokens on the CPU, so they are run after bookkeeping. + propose_draft_token_ids(valid_sampled_token_ids) + + stage_started = self._dspark_timing_start() + # Finalize KV connector (wait_for_save + clear metadata) after + # draft model runs. Deferred from target model forward to allow + # draft model to also save its KV cache. + if spec_config is not None: + self.finalize_kv_connector() + + with record_function_or_nullcontext("gpu_model_runner: eplb"): + self.eplb_step() + + # self.kv_connector_output may be modified during drafting + kv_connector_output = self.kv_connector_output + self.kv_connector_output = None + draft_token_lengths = None + if ( + self.use_async_scheduling + and self._draft_token_lengths_cpu is not None + and self._draft_token_length_req_ids is not None + ): + draft_token_lengths = { + req_id: int(length) + for req_id, length in zip( + self._draft_token_length_req_ids, + self._draft_token_lengths_cpu, + strict=True, + ) + } + + with record_function_or_nullcontext("gpu_model_runner: ModelRunnerOutput"): + output = ModelRunnerOutput( + req_ids=req_ids_output_copy, + req_id_to_index=req_id_to_index_output_copy, + sampled_token_ids=valid_sampled_token_ids, + logprobs=logprobs_lists, + prompt_logprobs_dict=prompt_logprobs_dict, + kv_connector_output=kv_connector_output, + ec_connector_output=ec_connector_output + if self.supports_mm_inputs + else None, + num_nans_in_logits=num_nans_in_logits, + cudagraph_stats=cudagraph_stats, + routed_experts=None, + draft_token_lengths=draft_token_lengths, + ) + self._dspark_timing_record("finalize_output", stage_started) + + if not self.use_async_scheduling: + if self.routed_experts_initialized: + # Sync path: D2H was issued in ``_bookkeeping_sync`` and + # synchronized by ``_to_list``'s event.synchronize(), so + # the pinned buffers are ready to be wrapped as numpy. + total = scheduler_output.total_num_scheduled_tokens + output.routed_experts = RoutedExpertsLists( + routing_data=self.routed_experts_cpu[:total].numpy(), + slot_mapping=self.routed_experts_slot_mapping_cpu[:total].numpy(), + ) + self._dspark_timing_finish_iteration() + return output + + with record_function_or_nullcontext( + "gpu_model_runner: AsyncGPUModelRunnerOutput" + ): + # Async path: produce a device-side snapshot that the async + # copy stream can D2H later. Both tensors must be private + # clones because: + # - ``routing_data`` source is the shared capturer buffer, + # which is ``clear_buffer()``-ed at the start of the + # next step on the default stream. + # - ``slot_mapping`` source is our own + # ``routed_experts_slot_mapping_device``, which the + # next ``_prepare_inputs`` overwrites on the default + # stream while the D2H is still pending on the copy + # stream. + # Without clones, the copy stream would read torn data. + routed_experts_snapshot = None + if self.routed_experts_initialized: + buf = self.routed_experts_capturer.get_device_buffer() + total = scheduler_output.total_num_scheduled_tokens + routed_experts_snapshot = RoutedExpertsTensors( + routing_data=buf[:total].clone(), + slot_mapping=self.routed_experts_slot_mapping_device[ + :total + ].clone(), + ) + + async_output = AsyncGPUModelRunnerOutput( + model_runner_output=output, + sampled_token_ids=sampler_output.sampled_token_ids, + logprobs_tensors=sampler_output.logprobs_tensors, + invalid_req_indices=invalid_req_indices, + async_output_copy_stream=self._get_or_create_async_output_copy_stream(), + vocab_size=self.input_batch.vocab_size, + routed_experts=routed_experts_snapshot, + ) + with record_function_or_nullcontext( + "gpu_model_runner: set_async_sampled_token_ids" + ): + # Save ref of sampled_token_ids CPU tensor if the batch contains + # any requests with sampling params that require output ids. + self.input_batch.set_async_sampled_token_ids( + async_output.sampled_token_ids_cpu, + async_output.async_copy_ready_event, + ) + + self._dspark_timing_finish_iteration() + return async_output + + def _pp_broadcast_prev_sampled_token_ids( + self, sampled_token_ids: torch.Tensor + ) -> None: + """Broadcast sampled token ids (GPU) from last PP stage""" + pp = get_pp_group() + assert pp.is_last_rank + # `prev_sampled_token_ids` is expected to have shape [num_reqs, 1]. + assert sampled_token_ids.dim() == 2 and sampled_token_ids.shape[-1] == 1, ( + "PP+async expects sampled_token_ids to have shape [num_reqs, 1]" + ) + # Skip for chunked prefill: sampled tokens are dummy + # and will be discarded, no need to broadcast. + if not self._is_all_reqs_chunked_prefill(): + torch.distributed.broadcast( + sampled_token_ids, src=pp.rank, group=pp.device_group + ) + + def _pp_receive_prev_sampled_token_ids_to_input_batch(self) -> None: + """Receive sampled token ids broadcast from last PP stage""" + pp = get_pp_group() + assert not pp.is_last_rank + num_reqs = self.input_batch.num_reqs + # `prev_sampled_token_ids` is expected to have shape [num_reqs, 1]. + recv = torch.empty((num_reqs, 1), dtype=torch.int32, device=self.device) + # skip for chunked prefill. + if not self._is_all_reqs_chunked_prefill(): + torch.distributed.broadcast(recv, src=pp.last_rank, group=pp.device_group) + self.input_batch.prev_sampled_token_ids = recv + + # construct `prev_req_id_to_index` here so `_prepare_input_ids` + # can map req_id -> previous batch row + discard_req_indices = np.nonzero(self.discard_request_mask.np[:num_reqs])[0] + discard_req_indices_set = set(discard_req_indices) + prev_req_id_to_index: dict[str, int] = {} + for i, req_id in enumerate(self.input_batch.req_ids): + if i in discard_req_indices_set: + continue + prev_req_id_to_index[req_id] = i + # PP+async scheduling: advance per-request local cached output length by + # appending a placeholder (-1) token id. + if (req_state := self.requests.get(req_id)) is not None: + req_state.output_token_ids.append(-1) + pos = self.input_batch.num_tokens_no_spec[i] + self.input_batch.is_token_ids[i, pos] = True + self.input_batch.num_tokens_no_spec[i] = pos + 1 + self.input_batch.prev_req_id_to_index = prev_req_id_to_index + + def take_draft_token_ids(self) -> DraftTokenIds | None: + if not self.num_spec_tokens or not self._draft_token_req_ids: + return None + draft_token_ids, req_ids = self._get_draft_token_ids_cpu() + return DraftTokenIds(req_ids, draft_token_ids) + + def _copy_draft_token_ids_to_cpu( + self, scheduler_output: "SchedulerOutput", zeros_only: bool = False + ) -> None: + # Check if we need to copy draft tokens to CPU. In async scheduling, + # we only copy when needed for structured output, penalties or bad_words. + if self.use_async_scheduling and not ( + scheduler_output.has_structured_output_requests + or self.input_batch.sampling_metadata.output_token_ids + ): + return + # We must also set the corresponding request ids. + self._draft_token_req_ids = self.input_batch.req_ids.copy() + + draft_token_ids: torch.Tensor = self._draft_token_ids + if not torch.is_tensor(draft_token_ids): + return + assert self.draft_token_ids_event is not None + assert self.draft_token_ids_copy_stream is not None + assert self.draft_token_ids_cpu is not None + default_stream = torch.cuda.current_stream() + num_reqs = draft_token_ids.shape[0] + with torch.cuda.stream(self.draft_token_ids_copy_stream): + if not zeros_only: + # Trigger async copy of draft token ids to cpu. + self.draft_token_ids_copy_stream.wait_stream(default_stream) + self.draft_token_ids_cpu[:num_reqs].copy_( + draft_token_ids, non_blocking=True + ) + else: + # No copy needed, just zero-out cpu tensor. + self.draft_token_ids_cpu[:num_reqs] = 0 + self.draft_token_ids_event.record() + + def _get_draft_token_ids_cpu(self) -> tuple[list[list[int]], list[str]]: + if isinstance(self._draft_token_ids, list): + return self._draft_token_ids, self.input_batch.req_ids + req_ids = self._draft_token_req_ids + if req_ids is None: + return [], [] + assert self.draft_token_ids_event is not None + assert self.draft_token_ids_cpu is not None + self.draft_token_ids_event.synchronize() + draft_token_ids = self.draft_token_ids_cpu[: len(req_ids)].tolist() + if self._draft_token_lengths_cpu is not None: + draft_token_ids = [ + row[: max(0, min(len(row), int(length)))] + for row, length in zip( + draft_token_ids, + self._draft_token_lengths_cpu, + strict=True, + ) + ] + return draft_token_ids, req_ids + + def _copy_valid_sampled_token_count( + self, next_token_ids: torch.Tensor, valid_sampled_tokens_count: torch.Tensor + ) -> None: + if self.valid_sampled_token_count_event is None: + return + + default_stream = torch.cuda.current_stream() + # Initialize a new stream to overlap the copy operation with + # prepare_input of draft model. + with torch.cuda.stream(self.valid_sampled_token_count_copy_stream): + self.valid_sampled_token_count_copy_stream.wait_stream(default_stream) # type: ignore + counts = valid_sampled_tokens_count + counts_cpu = self.valid_sampled_token_count_cpu + assert counts_cpu is not None + counts_cpu[: counts.shape[0]].copy_(counts, non_blocking=True) + self.valid_sampled_token_count_event.record() + + if self.use_async_spec_decode: + # Stash for GPU-side correction in _prepare_inputs. + self.valid_sampled_token_count_gpu = valid_sampled_tokens_count + self.input_batch.prev_sampled_token_ids = next_token_ids.unsqueeze(1) + + def _get_valid_sampled_token_count(self) -> list[int]: + # Wait until valid_sampled_tokens_count is copied to cpu, + prev_sampled_token_ids = self.input_batch.prev_sampled_token_ids + sampled_count_event = self.valid_sampled_token_count_event + if sampled_count_event is None or prev_sampled_token_ids is None: + return [] + + counts_cpu = self.valid_sampled_token_count_cpu + assert counts_cpu is not None + sampled_count_event.synchronize() + return counts_cpu[: prev_sampled_token_ids.shape[0]].tolist() + + def _get_spec_decode_draft_probs( + self, spec_decode_metadata: SpecDecodeMetadata + ) -> torch.Tensor | None: + if self._draft_probs is None or self._draft_prob_req_ids is None: + return None + + row_by_req_id = { + req_id: idx for idx, req_id in enumerate(self._draft_prob_req_ids) + } + draft_probs_rows: list[torch.Tensor] = [] + for req_id, num_draft in zip( + self.input_batch.req_ids, spec_decode_metadata.num_draft_tokens + ): + if num_draft == 0: + continue + row_idx = row_by_req_id.get(req_id) + if row_idx is None: + logger.warning( + "Missing cached draft probabilities for request %s; " + "falling back to legacy speculative rejection behavior.", + req_id, + ) + return None + draft_probs_rows.append(self._draft_probs[row_idx, :num_draft]) + + if not draft_probs_rows: + return None + return torch.cat(draft_probs_rows, dim=0).contiguous() + + def propose_draft_token_ids( + self, + scheduler_output: "SchedulerOutput", + sampled_token_ids: torch.Tensor | list[list[int]], + sampling_metadata: SamplingMetadata, + hidden_states: torch.Tensor, + sample_hidden_states: torch.Tensor, + aux_hidden_states: list[torch.Tensor] | None, + spec_decode_metadata: SpecDecodeMetadata | None, + common_attn_metadata: CommonAttentionMetadata, + slot_mappings: dict[str, torch.Tensor] | list[dict[str, torch.Tensor]] | None, + ) -> list[list[int]] | torch.Tensor: + num_scheduled_tokens = scheduler_output.total_num_scheduled_tokens + spec_config = self.speculative_config + assert spec_config is not None + self._draft_probs = None + self._draft_prob_req_ids = None + self._draft_confidence = None + self._draft_confidence_req_ids = None + if spec_config.method == "ngram": + from vllm.v1.spec_decode.ngram_proposer import NgramProposer + + assert isinstance(sampled_token_ids, list) + assert isinstance(self.drafter, NgramProposer) + draft_token_ids = self.drafter.propose( + sampled_token_ids, + self.input_batch.num_tokens_no_spec, + self.input_batch.token_ids_cpu, + slot_mappings=slot_mappings, + ) + elif spec_config.method == "custom_class": + assert isinstance(sampled_token_ids, list) + draft_token_ids = cast(Any, self.drafter).propose( + sampled_token_ids, + self.input_batch.num_tokens_no_spec, + self.input_batch.token_ids_cpu, + slot_mappings=slot_mappings, + ) + elif spec_config.use_ngram_gpu(): + assert isinstance(self.drafter, NgramProposerGPU) + ( + next_token_ids, + valid_sampled_tokens_count, + valid_sampled_token_ids_gpu, + ) = self.drafter.update_token_ids_ngram( + sampled_token_ids, + self.input_batch, + self.token_ids_gpu_tensor, + self.num_tokens_no_spec_gpu, + self.discard_request_mask.gpu, + ) + self._copy_valid_sampled_token_count( + next_token_ids, valid_sampled_tokens_count + ) + + batch_size = next_token_ids.shape[0] + + draft_token_ids, num_valid_draft_tokens = self.drafter.propose( + self.num_tokens_no_spec_gpu[:batch_size], + self.token_ids_gpu_tensor[:batch_size], + valid_sampled_token_ids_gpu, + valid_sampled_tokens_count, + ) + + # Cache valid draft counts for scheduler-side trimming. + self._num_valid_draft_tokens = num_valid_draft_tokens + + # Async D2H copy on a dedicated stream. + copy_num_valid_draft_tokens( + self._num_valid_draft_tokens_cpu, + self._num_valid_draft_tokens_copy_stream, + self._num_valid_draft_tokens_event, + self._num_valid_draft_tokens, + self.input_batch.num_reqs, + ) + elif spec_config.method == "suffix": + assert isinstance(sampled_token_ids, list) + assert isinstance(self.drafter, SuffixDecodingProposer) + draft_token_ids = self.drafter.propose( + self.input_batch, sampled_token_ids, slot_mappings=slot_mappings + ) + elif spec_config.method == "medusa": + assert isinstance(sampled_token_ids, list) + assert isinstance(self.drafter, MedusaProposer) + + if sample_hidden_states.shape[0] == len(sampled_token_ids): + # The input to the target model does not include draft tokens. + hidden_states = sample_hidden_states + else: + indices = [] + offset = 0 + assert spec_decode_metadata is not None, ( + "No spec decode metadata for medusa" + ) + for num_draft, tokens in zip( + spec_decode_metadata.num_draft_tokens, sampled_token_ids + ): + indices.append(offset + len(tokens) - 1) + offset += num_draft + 1 + indices = torch.tensor(indices, device=self.device) + hidden_states = sample_hidden_states[indices] + + draft_token_ids = self.drafter.propose( + target_hidden_states=hidden_states, + sampling_metadata=sampling_metadata, + slot_mappings=slot_mappings, + ) + elif spec_config.uses_extract_hidden_states(): + assert isinstance(self.drafter, ExtractHiddenStatesProposer) + assert isinstance(sampled_token_ids, torch.Tensor), ( + "sampled_token_ids should be a torch.Tensor for " + "extract_hidden_states method." + ) + if not self.use_aux_hidden_state_outputs or aux_hidden_states is None: + raise ValueError( + "aux_hidden_states are required when using `extract_hidden_states`" + ) + target_hidden_states = [h[:num_scheduled_tokens] for h in aux_hidden_states] + + draft_token_ids = self.drafter.propose( + sampled_token_ids=sampled_token_ids, + target_hidden_states=target_hidden_states, + common_attn_metadata=common_attn_metadata, + slot_mappings=slot_mappings, + ) + next_token_ids, valid_sampled_tokens_count = ( + self.drafter.prepare_next_token_ids_padded( + sampled_token_ids, + self.requests, + self.input_batch, + self.discard_request_mask.gpu, + ) + ) + self._copy_valid_sampled_token_count( + next_token_ids, valid_sampled_tokens_count + ) + + elif ( + spec_config.use_eagle() + or spec_config.use_dflash() + or spec_config.use_dspark() + or spec_config.uses_draft_model() + ): + assert isinstance( + self.drafter, + EagleProposer + | DFlashProposer + | DSparkProposer + | DraftModelProposer + | Gemma4Proposer, + ) + + if spec_config.disable_padded_drafter_batch: + # When padded-batch is disabled, the sampled_token_ids should be + # the cpu-side list[list[int]] of valid sampled tokens for each + # request, with invalid requests having empty lists. + assert isinstance(sampled_token_ids, list), ( + "sampled_token_ids should be a python list when" + "padded-batch is disabled." + ) + next_token_ids = self.drafter.prepare_next_token_ids_cpu( + sampled_token_ids, + self.requests, + self.input_batch, + scheduler_output.num_scheduled_tokens, + ) + else: + # When using padded-batch, the sampled_token_ids should be + # the gpu tensor of sampled tokens for each request, of shape + # (num_reqs, num_spec_tokens + 1) with rejected tokens having + # value -1. + assert isinstance(sampled_token_ids, torch.Tensor), ( + "sampled_token_ids should be a torch.Tensor when" + "padded-batch is enabled." + ) + next_token_ids, valid_sampled_tokens_count = ( + self.drafter.prepare_next_token_ids_padded( + sampled_token_ids, + self.requests, + self.input_batch, + self.discard_request_mask.gpu, + ) + ) + self._copy_valid_sampled_token_count( + next_token_ids, valid_sampled_tokens_count + ) + + # Let the target override the hidden state fed to the drafter + # (e.g. DeepSeek V4 MTP needs the pre-hc_head residual). Safe to + # rebind here: hidden_states was already consumed for sampling + # above and is not used again in this branch. + target_hidden_getter_name = ( + "get_dspark_target_hidden_states" + if spec_config.use_dspark() + else "get_mtp_target_hidden_states" + ) + alt = getattr(self.get_model(), target_hidden_getter_name, lambda: None)() + if alt is not None: + hidden_states = alt + + num_rejected_tokens_gpu = None + if spec_decode_metadata is None: + token_indices_to_sample = None + # input_ids can be None for multimodal models. + target_token_ids = self.input_ids.gpu[:num_scheduled_tokens] + target_positions = self._get_positions(num_scheduled_tokens) + if self.use_aux_hidden_state_outputs: + assert aux_hidden_states is not None + target_hidden_states = torch.cat( + [h[:num_scheduled_tokens] for h in aux_hidden_states], dim=-1 + ) + else: + target_hidden_states = hidden_states[:num_scheduled_tokens] + else: + if spec_config.disable_padded_drafter_batch: + token_indices_to_sample = None + common_attn_metadata, token_indices = self.drafter.prepare_inputs( + common_attn_metadata, + sampled_token_ids, + spec_decode_metadata.num_draft_tokens, + ) + target_token_ids = self.input_ids.gpu[token_indices] + target_positions = self._get_positions(token_indices) + if self.use_aux_hidden_state_outputs: + assert aux_hidden_states is not None + target_hidden_states = torch.cat( + [h[token_indices] for h in aux_hidden_states], dim=-1 + ) + else: + target_hidden_states = hidden_states[token_indices] + else: + ( + common_attn_metadata, + token_indices_to_sample, + num_rejected_tokens_gpu, + ) = self.drafter.prepare_inputs_padded( + common_attn_metadata, + spec_decode_metadata, + valid_sampled_tokens_count, + ) + total_num_tokens = common_attn_metadata.num_actual_tokens + # When padding the batch, token_indices is just a range + target_token_ids = self.input_ids.gpu[:total_num_tokens] + target_positions = self._get_positions(total_num_tokens) + if self.use_aux_hidden_state_outputs: + assert aux_hidden_states is not None + target_hidden_states = torch.cat( + [h[:total_num_tokens] for h in aux_hidden_states], dim=-1 + ) + else: + target_hidden_states = hidden_states[:total_num_tokens] + + if self.supports_mm_inputs and self.drafter.supports_mm_inputs: + mm_embed_inputs = self._gather_mm_embeddings( + scheduler_output, + shift_computed_tokens=1, + ) + else: + mm_embed_inputs = None + + # DSpark keys its persistent draft KV by a stable per-request slot, + # so it needs the request ids in the same order used to build + # next_token_ids (the running input_batch order). Other proposers + # do not accept this argument. + dspark_propose_kwargs = ( + {"req_ids": self.input_batch.req_ids} + if spec_config.use_dspark() + else {} + ) + draft_token_ids = self.drafter.propose( + target_token_ids=target_token_ids, + target_positions=target_positions, + target_hidden_states=target_hidden_states, + next_token_ids=next_token_ids, + token_indices_to_sample=token_indices_to_sample, + sampling_metadata=sampling_metadata, + common_attn_metadata=common_attn_metadata, + mm_embed_inputs=mm_embed_inputs, + num_rejected_tokens_gpu=num_rejected_tokens_gpu, + slot_mappings=slot_mappings, + **dspark_propose_kwargs, + ) + if hasattr(self.drafter, "take_last_draft_probs"): + draft_probs = self.drafter.take_last_draft_probs() + if draft_probs is not None: + self._draft_probs = draft_probs + self._draft_prob_req_ids = self.input_batch.req_ids.copy() + if hasattr(self.drafter, "take_last_confidence"): + confidence = self.drafter.take_last_confidence() + if confidence is not None: + self._draft_confidence = confidence + self._draft_confidence_req_ids = ( + self.input_batch.req_ids.copy() + ) + if spec_config.use_dspark() and hasattr( + self.drafter, "take_last_draft_lengths" + ): + self._draft_token_lengths_cpu = self.drafter.take_last_draft_lengths() + self._draft_token_length_req_ids = ( + self.input_batch.req_ids.copy() + if self._draft_token_lengths_cpu is not None + else None + ) + + 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(): + assert config_name in allowed_config_names, ( + f"Config `{config_name}` not supported. " + f"Allowed configs: {allowed_config_names}" + ) + 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, load_dummy_weights: bool = False) -> None: + """ + Args: + load_dummy_weights: load dummy weights instead of real weights. + """ + logger.info_once( + "Starting to load model %s...", + self.model_config.model, + scope="global", + ) + + if self.parallel_config.enable_eplb: + self.eplb_state = EplbState(self.parallel_config, self.device) + eplb_models = 0 + + try: + with DeviceMemoryProfiler() as m: + time_before_load = time.perf_counter() + if load_dummy_weights: + self.load_config.load_format = "dummy" + model_loader = get_model_loader(self.load_config) + self.model = model_loader.load_model( + vllm_config=self.vllm_config, model_config=self.model_config + ) + if self.lora_config: + self.model = self.load_lora_model( + self.model, self.vllm_config, self.device + ) + if hasattr(self, "drafter"): + logger.info_once("Loading drafter model...") + if hasattr(self.drafter, "load_model"): + self.drafter.load_model(self.model) + if ( + hasattr(self.drafter, "model") + and is_mixture_of_experts(self.drafter.model) + and self.parallel_config.enable_eplb + ): + 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 + logger.info_once( + "EPLB is enabled for drafter model %s.", + spec_config.draft_model_config.model, + ) + if self.eplb_state is None: + self.eplb_state = EplbState( + self.parallel_config, self.device + ) + self.eplb_state.add_model( + self.drafter.model, + spec_config.draft_model_config, + ) + eplb_models += 1 + + self._setup_eagle3_aux_hidden_state_outputs() + + # Resolve the MoE model, unwrapping VLM wrappers if needed. + # VLM models (e.g. KimiK25ForConditionalGeneration) wrap the + # actual MoE language model but don't implement + # MixtureOfExperts themselves. + moe_candidate = self.model + if not is_mixture_of_experts(moe_candidate) and isinstance( + moe_candidate, SupportsMultiModal + ): + moe_candidate = moe_candidate.get_language_model() + if is_mixture_of_experts(moe_candidate): + self._moe_model = moe_candidate + + if ( + self._moe_model is not None + and self.parallel_config.enable_eplb + and not load_dummy_weights + ): + logger.info_once( + "EPLB is enabled for model %s.", + self.model_config.model, + ) + assert self.eplb_state is not None + self.eplb_state.add_model( + self._moe_model, + self.model_config, + ) + eplb_models += 1 + + time_after_load = time.perf_counter() + self.model_memory_usage = m.consumed_memory + except torch.cuda.OutOfMemoryError as e: + msg = ( + "Failed to load model - not enough GPU memory. " + "Try lowering --gpu-memory-utilization to free memory for weights, " + "increasing --tensor-parallel-size, or using --quantization. " + "See https://docs.vllm.ai/en/latest/configuration/conserving_memory/ " + "for more tips." + ) + combined_msg = f"{msg} (original error: {e})" + logger.error(combined_msg) + raise e + logger.info_once( + "Model loading took %s GiB memory and %.6f seconds", + format_gib(self.model_memory_usage), + time_after_load - time_before_load, + ) + if not load_dummy_weights: + prepare_communication_buffer_for_model(self.model) + if (drafter := getattr(self, "drafter", None)) and ( + drafter_model := getattr(drafter, "model", None) + ): + prepare_communication_buffer_for_model(drafter_model) + mm_config = self.model_config.multimodal_config + self.is_multimodal_pruning_enabled = ( + supports_multimodal_pruning(self.get_model()) + and mm_config is not None + and mm_config.is_multimodal_pruning_enabled() + ) + self.requires_sequential_video_encoding = hasattr( + self.get_model(), "requires_sequential_video_encoding" + ) # Temporary hack for dynamic res video w/o support for bs>1 yet + + if ( + self._moe_model is not None + and self.parallel_config.enable_eplb + and not load_dummy_weights + and self.eplb_state is not None + and self.eplb_state.is_async + ): + self.eplb_state.start_async_loop() + + if ( + self.vllm_config.compilation_config.mode + == CompilationMode.STOCK_TORCH_COMPILE + ): + from vllm.env_override import _apply_constrain_to_fx_strides_patch + + _apply_constrain_to_fx_strides_patch() + backend = self.vllm_config.compilation_config.init_backend(self.vllm_config) + compilation_counter.stock_torch_compile_count += 1 + self.model.compile(fullgraph=True, backend=backend) + return + # for other compilation modes, cudagraph behavior is controlled by + # CudagraphWrapper and CudagraphDispatcher of vllm. + + # wrap the model with full cudagraph wrapper if needed. + cudagraph_mode = self.compilation_config.cudagraph_mode + assert cudagraph_mode is not None + if ( + is_breakable_cudagraph_enabled() + and cudagraph_mode != CUDAGraphMode.NONE + and not self.parallel_config.use_ubatching + ): + self.model = BreakableCUDAGraphWrapper(self.model, self.vllm_config) + elif ( + cudagraph_mode.has_full_cudagraphs() + and not self.parallel_config.use_ubatching + ): + self.model = CUDAGraphWrapper( + self.model, self.vllm_config, runtime_mode=CUDAGraphMode.FULL + ) + elif self.parallel_config.use_ubatching: + if cudagraph_mode.has_full_cudagraphs(): + self.model = UBatchWrapper( + self.model, self.vllm_config, CUDAGraphMode.FULL, self.device + ) + else: + self.model = UBatchWrapper( + self.model, self.vllm_config, CUDAGraphMode.NONE, self.device + ) + + get_offloader().post_init() + + def _setup_eagle3_aux_hidden_state_outputs(self) -> None: + if not self.use_aux_hidden_state_outputs: + return + + if not supports_eagle3(self.get_model()): + raise RuntimeError( + "Model does not support EAGLE3 interface but " + "aux_hidden_state_outputs was requested" + ) + # Try to get auxiliary layers from speculative config, + # otherwise use model's default layers + aux_layers = self._get_eagle3_aux_layers_from_config() + if aux_layers: + logger.info( + "Using auxiliary layers from speculative config: %s", aux_layers + ) + else: + aux_layers = self.model.get_eagle3_default_aux_hidden_state_layers() + + self.model.set_aux_hidden_state_layers(aux_layers) + + def _get_eagle3_aux_layers_from_config(self) -> tuple[int, ...] | None: + """Extract Eagle3 auxiliary layer indices from speculative config. + + These indices specify which hidden states from the base model should + be used as auxiliary inputs for the Eagle3 drafter model during + speculative decoding. + + Returns: + Tuple of layer indices if found in draft model config, + None otherwise. + """ + if not (self.speculative_config and self.speculative_config.draft_model_config): + return None + + hf_config = self.speculative_config.draft_model_config.hf_config + + layer_ids = getattr(hf_config, "eagle_aux_hidden_state_layer_ids", None) + if not layer_ids: + dflash_config = getattr(hf_config, "dflash_config", None) + eagle_config = getattr(hf_config, "eagle_config", None) + + if dflash_config and isinstance(dflash_config, dict): + # Add 1 to convert DFlash's aux layer id semantics + layer_ids = [ + i + 1 for i in (dflash_config.get("target_layer_ids") or []) + ] + + if eagle_config and isinstance(eagle_config, dict): + layer_ids = eagle_config.get("eagle_aux_hidden_state_layer_ids") + + if layer_ids and isinstance(layer_ids, (list, tuple)): + return tuple(layer_ids) + + return None + + 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 from a weights iterator or from disk + + :param weights_iterator: weights to load into model + :param weights_path: path to load weights from if weights_iterator is not + provided. Use path of original model if neither is provided. + :param 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 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: + self.model_config.model = weights_path + 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 = model.get_parameter(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 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, + ) + + def _get_prompt_logprobs_dict( + self, + hidden_states: torch.Tensor, + num_scheduled_tokens: dict[str, int], + ) -> dict[str, LogprobsTensors | None]: + num_prompt_logprobs_dict = self.num_prompt_logprobs + if not num_prompt_logprobs_dict: + return {} + + prompt_logprobs_dict: dict[str, LogprobsTensors | None] = {} + + # Since prompt logprobs are a rare feature, prioritize simple, + # maintainable loop over optimal performance. + completed_prefill_reqs = [] + for req_id, num_prompt_logprobs in num_prompt_logprobs_dict.items(): + num_tokens = num_scheduled_tokens.get(req_id) + if num_tokens is None: + # This can happen if the request was preempted in prefill stage. + continue + + # Get metadata for this request. + request = self.requests[req_id] + if request.prompt_token_ids is None: + # Prompt logprobs is incompatible with prompt embeddings + continue + + num_prompt_tokens = len(request.prompt_token_ids) + prompt_token_ids = torch.tensor(request.prompt_token_ids).to( + self.device, non_blocking=True + ) + + # Set up target LogprobsTensors object. + logprobs_tensors = request.in_progress_prompt_logprobs_cpu + if logprobs_tensors is None: + # Create empty logprobs CPU tensors for the entire prompt. + # If chunked, we'll copy in slice by slice. + logprobs_tensors = LogprobsTensors.empty_cpu( + num_prompt_tokens - 1, num_prompt_logprobs + 1 + ) + request.in_progress_prompt_logprobs_cpu = logprobs_tensors + + # Determine number of logits to retrieve. + start_idx = request.num_computed_tokens + start_tok = start_idx + 1 + num_remaining_tokens = num_prompt_tokens - start_tok + if num_tokens <= num_remaining_tokens: + # This is a chunk, more tokens remain. + # In the == case, there are no more prompt logprobs to produce + # but we want to defer returning them to the next step where we + # have new generated tokens to return. + num_logits = num_tokens + else: + # This is the last chunk of prompt tokens to return. + num_logits = num_remaining_tokens + completed_prefill_reqs.append(req_id) + prompt_logprobs_dict[req_id] = logprobs_tensors + + if num_logits <= 0: + # This can happen for the final chunk if we prefilled exactly + # (num_prompt_tokens - 1) tokens for this request in the prior + # step. There are no more prompt logprobs to produce. + continue + + # Get the logits corresponding to this req's prompt tokens. + # If this is a partial request (i.e. chunked prefill), + # then there is prompt logprob generated for each index. + req_idx = self.input_batch.req_id_to_index[req_id] + offset = self.query_start_loc.np[req_idx].item() + prompt_hidden_states = hidden_states[offset : offset + num_logits] + logits = self.model.compute_logits(prompt_hidden_states) + + # Get the "target" tokens for each index. For prompt at index i, + # the token at prompt index i+1 is the "sampled" token we want + # to gather the logprob for. + tgt_token_ids = prompt_token_ids[start_tok : start_tok + num_logits] + + # Compute prompt logprobs. + logprobs = self.sampler.compute_logprobs(logits) + token_ids, logprobs, ranks, _ = self.sampler.gather_logprobs( + logprobs, num_prompt_logprobs, tgt_token_ids + ) + + # Transfer GPU->CPU async. + chunk_slice = slice(start_idx, start_idx + num_logits) + logprobs_tensors.logprob_token_ids[chunk_slice].copy_( + token_ids, non_blocking=True + ) + logprobs_tensors.logprobs[chunk_slice].copy_(logprobs, non_blocking=True) + logprobs_tensors.selected_token_ranks[chunk_slice].copy_( + ranks, non_blocking=True + ) + + # Remove requests that have completed prefill from the batch + # num_prompt_logprobs_dict. + for req_id in completed_prefill_reqs: + del num_prompt_logprobs_dict[req_id] + self.requests[req_id].in_progress_prompt_logprobs_cpu = None + + # Must synchronize the non-blocking GPU->CPU transfers. + if prompt_logprobs_dict: + self._sync_device() + + return prompt_logprobs_dict + + def _get_nans_in_logits( + self, + logits: torch.Tensor | None, + ) -> dict[str, int]: + try: + if logits is None: + return {req_id: 0 for req_id in self.input_batch.req_ids} + + num_nans_in_logits = {} + num_nans_for_index = logits.isnan().sum(dim=-1).cpu().numpy() + for req_id in self.input_batch.req_ids: + req_index = self.input_batch.req_id_to_index[req_id] + num_nans_in_logits[req_id] = ( + int(num_nans_for_index[req_index]) + if num_nans_for_index is not None and req_index < logits.shape[0] + else 0 + ) + return num_nans_in_logits + except IndexError: + return {} + + @contextmanager + def maybe_randomize_inputs( + self, input_ids: torch.Tensor | None, inputs_embeds: torch.Tensor | None + ): + """ + Randomize input_ids if VLLM_RANDOMIZE_DP_DUMMY_INPUTS is set. + This is to help balance expert-selection + - during profile_run + - during DP rank dummy run + """ + + dp_size = self.vllm_config.parallel_config.data_parallel_size + randomize_inputs = envs.VLLM_RANDOMIZE_DP_DUMMY_INPUTS and dp_size > 1 + if not randomize_inputs: + yield + elif input_ids is not None: + + @functools.cache + def rand_input_ids() -> torch.Tensor: + return torch.randint_like( + self.input_ids.gpu, + low=0, + high=self.model_config.get_vocab_size(), + ) + + logger.debug_once("Randomizing dummy input_ids for DP Rank") + input_ids.copy_(rand_input_ids()[: input_ids.size(0)], non_blocking=True) + yield + input_ids.fill_(0) + else: + + @functools.cache + def rand_inputs_embeds() -> torch.Tensor: + return torch.randn_like( + self.inputs_embeds.gpu, + ) + + assert inputs_embeds is not None + logger.debug_once("Randomizing dummy inputs_embeds for DP Rank") + inputs_embeds.copy_( + rand_inputs_embeds()[: inputs_embeds.size(0)], non_blocking=True + ) + yield + inputs_embeds.fill_(0) + + def _get_mm_dummy_batch( + self, + modality: str, + max_items_per_batch: int, + ) -> BatchedTensorInputs: + """Dummy data for profiling and precompiling multimodal models.""" + assert self.mm_budget is not None + + # Don't use `max_items_per_batch` here to avoid redundant computation + dummy_mm_inputs = self.mm_registry.get_dummy_mm_inputs( + self.model_config, + mm_counts={modality: 1}, + cache=self.mm_budget.cache, + ) + dummy_mm_item = dummy_mm_inputs["mm_kwargs"][modality][0] + + # We use the cache so that the item is saved to the cache, + # but not read from the cache + assert dummy_mm_item is not None, "Item should not already be cached" + + return next( + mm_kwargs_batch + for _, _, mm_kwargs_batch in group_and_batch_mm_kwargs( + [(modality, dummy_mm_item)] * max_items_per_batch, + device=self.device, + pin_memory=self.pin_memory, + ) + ) + + @torch.inference_mode() + def _dummy_run( + self, + num_tokens: int, + cudagraph_runtime_mode: CUDAGraphMode | None = None, + force_attention: bool = False, + uniform_decode: bool = False, + allow_microbatching: bool = True, + skip_eplb: bool = False, + is_profile: bool = False, + create_mixed_batch: bool = False, + create_single_prefill: bool = False, + remove_lora: bool = True, + is_graph_capturing: bool = False, + num_active_loras: int = 0, + profile_seq_lens: int | None = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + """ + Run a dummy forward pass to warm up/profile run or capture the + CUDA graph for the model. + + Args: + num_tokens: Number of tokens to run the dummy forward pass. + cudagraph_runtime_mode: used to control the behavior. + - if not set will determine the cudagraph mode based on using + the self.cudagraph_dispatcher. + - CUDAGraphMode.NONE: No cudagraph, for warm up and profile run + - CUDAGraphMode.PIECEWISE: Piecewise cudagraph. + - CUDAGraphMode.FULL: Full cudagraph, attention metadata is + needed. + force_attention: If True, always create attention metadata. Used to + warm up attention backend when mode is NONE. + uniform_decode: If True, the batch is a uniform decode batch. + skip_eplb: If True, skip EPLB state update. + is_profile: If True, this is a profile run. + create_mixed_batch: If True, create a mixed batch with both decode + (1 token) and prefill (multiple tokens) requests. + remove_lora: If False, dummy LoRAs are not destroyed after the run + num_active_loras: Number of distinct active LoRAs to capture for. + LoRA is activated when num_active_loras > 0. + profile_seq_lens: If provided, use this value for seq_lens instead + of max_query_len. Used to profile attention workspace that + scales with context length. + """ + mm_config = self.vllm_config.model_config.multimodal_config + if mm_config and mm_config.mm_encoder_only: + # The current dummy run only covers LM execution, so we can skip it. + # mm encoder dummy run may need to add in the future. + return torch.tensor([]), torch.tensor([]) + + assert ( + cudagraph_runtime_mode is None + or cudagraph_runtime_mode.is_valid_runtime_mode() + ) + + # If cudagraph_mode.decode_mode() == FULL and + # cudagraph_mode.separate_routine(). This means that we are using + # different graphs and/or modes for mixed prefill-decode batches vs. + # uniform decode batches. A uniform decode batch means that all + # requests have identical query length, except a potential virtual + # request (shorter) in the batch account for padding. + # Uniform decode batch could either be common pure decode, where + # max_query_len == 1, or speculative decode, where + # max_query_len == 1 + num_spec_decode_tokens. + + # When setting max_query_len = 1, we switch to and capture the optimized + # routine of FA2 for pure decode, i.e., Flashdecode + an optimization + # for GQA/MQA. + max_query_len = self.uniform_decode_query_len if uniform_decode else num_tokens + + # Set num_scheduled_tokens based on num_tokens and max_num_seqs + # for dummy run with LoRA so that the num_reqs collectively + # has num_tokens in total. + assert num_tokens <= self.max_num_tokens + max_num_reqs = self.scheduler_config.max_num_seqs + if create_single_prefill: + assert not uniform_decode + assert not create_mixed_batch + # Single-prefill batch (max_query_len == num_tokens) — warms + # specializations the mixed-batch shape misses. + num_reqs = 1 + num_scheduled_tokens_list = [num_tokens] + max_query_len = num_tokens + elif create_mixed_batch: + assert not uniform_decode + assert not create_single_prefill + # Create mixed batch: + # first half decode tokens, second half one prefill + num_decode_tokens = min(max_num_reqs - 1, num_tokens // 2) + num_prefill_tokens = num_tokens - num_decode_tokens + num_reqs = num_decode_tokens + 1 + + # Create decode requests (1 token each) followed by prefill request + num_scheduled_tokens_list = [1] * num_decode_tokens + [num_prefill_tokens] + # Note: Overriding max_query_len to be the prefill tokens + max_query_len = num_prefill_tokens + elif uniform_decode: + assert not create_mixed_batch + assert not create_single_prefill + num_reqs = min(max_num_reqs, cdiv(num_tokens, max_query_len)) + num_scheduled_tokens_list = [max_query_len] * num_reqs + if num_tokens % max_query_len != 0: + num_scheduled_tokens_list[-1] = num_tokens % max_query_len + else: + num_reqs = min(num_tokens, max_num_reqs) + min_tokens_per_req = num_tokens // num_reqs + num_scheduled_tokens_list = [min_tokens_per_req] * num_reqs + num_scheduled_tokens_list[-1] += num_tokens % num_reqs + + assert sum(num_scheduled_tokens_list) == num_tokens + assert len(num_scheduled_tokens_list) == num_reqs + num_scheduled_tokens = np.array(num_scheduled_tokens_list, dtype=np.int32) + num_tokens_unpadded = int(num_scheduled_tokens.sum()) + + num_sampled_tokens = np.ones(num_reqs, dtype=np.int32) + + _cudagraph_mode, batch_desc, should_ubatch, num_tokens_across_dp, _ = ( + self._determine_batch_execution_and_padding( + num_tokens=num_tokens_unpadded, + num_reqs=num_reqs, + num_scheduled_tokens_np=num_scheduled_tokens, + max_num_scheduled_tokens=max_query_len, + use_cascade_attn=False, + allow_microbatching=allow_microbatching, + force_eager=is_profile + or (cudagraph_runtime_mode == CUDAGraphMode.NONE), + # `force_uniform_decode` is used for cudagraph capture; because for + # capturing mixed prefill-decode batches, we sometimes use + # num_tokens == num_reqs which looks like a uniform decode batch to the + # dispatcher; but we actually want to capture a piecewise cudagraph + force_uniform_decode=uniform_decode, + # `force_has_lora` is used for cudagraph capture; because LoRA is + # activated later in the context manager, but we need to know the + # LoRA state when determining the batch descriptor for capture + force_has_lora=num_active_loras > 0, + # `force_num_active_loras` is used for cudagraph capture; because we + # need to capture graphs for specific num_active_loras counts + force_num_active_loras=num_active_loras, + ) + ) + + if cudagraph_runtime_mode is None: + cudagraph_runtime_mode = _cudagraph_mode + else: + assert cudagraph_runtime_mode == _cudagraph_mode, ( + f"Cudagraph runtime mode mismatch in dummy_run. " + f"Expected {_cudagraph_mode}, but got {cudagraph_runtime_mode}." + ) + + num_tokens_padded = batch_desc.num_tokens + num_reqs_padded = ( + batch_desc.num_reqs if batch_desc.num_reqs is not None else num_reqs + ) + ubatch_slices, ubatch_slices_padded = maybe_create_ubatch_slices( + should_ubatch, + num_scheduled_tokens, + num_tokens_padded, + num_reqs_padded, + self.vllm_config.parallel_config.num_ubatches, + ) + logger.debug( + "ubatch_slices: %s, ubatch_slices_padded: %s", + ubatch_slices, + ubatch_slices_padded, + ) + + attn_metadata: PerLayerAttnMetadata | None = None + + slot_mappings_by_group, slot_mappings = self._get_slot_mappings( + num_tokens_padded=num_tokens_padded, + num_reqs_padded=num_reqs_padded, + num_tokens_unpadded=num_tokens_unpadded, + ubatch_slices=ubatch_slices_padded, + ) + + # Dummy runs have no real slot assignments — fill with -1 so + # concat_and_cache kernels skip the KV write. + if slot_mappings_by_group is not None: + for sm in slot_mappings_by_group.values(): + sm.fill_(-1) + + # _dummy_run shares pinned CPU buffers (seq_lens, query_start_loc, + # etc.) with execute_model. It must participate in the same event + # protocol so that back-to-back dummy/real steps don't overwrite + # pinned memory while a prior non_blocking H2D DMA is still reading. + with self.synchronize_input_prep(): + # If force_attention is True, we always capture attention. + # Otherwise, it only happens for cudagraph_runtime_mode=FULL. + if force_attention or cudagraph_runtime_mode == CUDAGraphMode.FULL: + if profile_seq_lens is not None: + seq_lens = profile_seq_lens # type: ignore[assignment] + elif create_mixed_batch: + # In the mixed batch mode (used for FI warmup), we use + # shorter sequence lengths to run faster. + # TODO(luka) better system for describing dummy batches + seq_lens = torch.tensor( # type: ignore[assignment] + [1] * num_decode_tokens + [num_prefill_tokens + 1], + dtype=torch.int, + ) + else: + seq_lens = max_query_len # type: ignore[assignment] + self.optimistic_seq_lens_cpu[:num_reqs] = seq_lens + self.optimistic_seq_lens_cpu[num_reqs:].fill_(0) + self.seq_lens.copy_(self.optimistic_seq_lens_cpu, non_blocking=True) + + cum_num_tokens = self._get_cumsum_and_arange( + num_scheduled_tokens, self.query_pos.np + ) + self.query_start_loc.np[1 : num_reqs + 1] = cum_num_tokens + self.query_start_loc.copy_to_gpu() + + # Sync block table CPU->GPU so cleared rows from + # remove_request() are visible to the attention metadata + # builder. Without this, stale block IDs from finished + # requests can corrupt Mamba state. + self.input_batch.block_table.commit_block_table(num_reqs_padded) + + pad_attn = cudagraph_runtime_mode == CUDAGraphMode.FULL + attn_metadata, _ = self._build_attention_metadata( + num_tokens=num_tokens_unpadded, + num_tokens_padded=num_tokens_padded if pad_attn else None, + num_reqs=num_reqs_padded, + max_query_len=max_query_len, + ubatch_slices=(ubatch_slices_padded if pad_attn else ubatch_slices), + for_cudagraph_capture=is_graph_capturing, + slot_mappings=slot_mappings_by_group, + use_spec_decode=self.speculative_config is not None, + ) + + with self.maybe_dummy_run_with_lora( + self.lora_config, + num_scheduled_tokens, + num_sampled_tokens, + remove_lora, + num_active_loras, + ): + # Make sure padding doesn't exceed max_num_tokens + assert num_tokens_padded <= self.max_num_tokens + model_kwargs = self._init_model_kwargs() + if self.supports_mm_inputs and not self.model_config.is_encoder_decoder: + input_ids, inputs_embeds = self._prepare_mm_inputs(num_tokens_padded) + + model_kwargs = { + **model_kwargs, + **self._dummy_mm_kwargs(num_reqs), + } + elif self.enable_prompt_embeds: + input_ids = None + inputs_embeds = self.inputs_embeds.gpu[:num_tokens_padded] + model_kwargs = self._init_model_kwargs() + else: + input_ids = self.input_ids.gpu[:num_tokens_padded] + inputs_embeds = None + + if self.uses_mrope: + positions = self.mrope_positions.gpu[:, :num_tokens_padded] + elif self.uses_xdrope_dim > 0: + positions = self.xdrope_positions.gpu[:, :num_tokens_padded] + else: + positions = self.positions[:num_tokens_padded] + + if get_pp_group().is_first_rank: + intermediate_tensors = None + else: + if self.intermediate_tensors is None: + self.intermediate_tensors = ( + self.model.make_empty_intermediate_tensors( + batch_size=self.max_num_tokens, + dtype=self.model_config.dtype, + device=self.device, + ) + ) + + intermediate_tensors = self.sync_and_gather_intermediate_tensors( + num_tokens_padded, None, False + ) + + if ubatch_slices_padded is not None: + # Adjust values to reflect a single ubatch. + # TODO(sage,lucas): this is cruft that should be addressed in + # the padding refactor. + num_tokens_padded = ubatch_slices_padded[0].num_tokens + if num_tokens_across_dp is not None: + num_tokens_across_dp[:] = num_tokens_padded + + with ( + self.maybe_randomize_inputs(input_ids, inputs_embeds), + set_forward_context( + attn_metadata, + self.vllm_config, + num_tokens=num_tokens_padded, + num_tokens_across_dp=num_tokens_across_dp, + cudagraph_runtime_mode=cudagraph_runtime_mode, + batch_descriptor=batch_desc, + ubatch_slices=ubatch_slices_padded, + slot_mapping=slot_mappings, + ), + ): + outputs = self.model( + input_ids=input_ids, + positions=positions, + intermediate_tensors=intermediate_tensors, + inputs_embeds=inputs_embeds, + **model_kwargs, + ) + + if self.use_aux_hidden_state_outputs: + hidden_states, _ = outputs + else: + hidden_states = outputs + + if self.speculative_config and ( + self.speculative_config.use_eagle() + or self.speculative_config.uses_draft_model() + or self.speculative_config.uses_extract_hidden_states() + or self.speculative_config.use_dspark() + ): + assert isinstance( + self.drafter, + EagleProposer + | DFlashProposer + | DSparkProposer + | DraftModelProposer + | ExtractHiddenStatesProposer + | Gemma4Proposer, + ) + assert self.speculative_config is not None + # Eagle currently only supports PIECEWISE cudagraphs. + # Therefore only use cudagraphs if the main model uses PIECEWISE + # NOTE(lucas): this is a hack, need to clean up. + use_cudagraphs = ( + ( + is_graph_capturing + and cudagraph_runtime_mode == CUDAGraphMode.PIECEWISE + ) + or ( + not is_graph_capturing + and cudagraph_runtime_mode != CUDAGraphMode.NONE + ) + ) and not self.speculative_config.enforce_eager + + # Note(gnovack) - We need to disable cudagraphs for one of the two + # lora cases when cudagraph_specialize_lora is enabled. This is a + # short term mitigation for issue mentioned in + # https://github.com/vllm-project/vllm/issues/28334 + if ( + self.compilation_config.cudagraph_specialize_lora + and num_active_loras > 0 + ): + use_cudagraphs = False + + self.drafter.dummy_run( + num_tokens, + use_cudagraphs=use_cudagraphs, + is_graph_capturing=is_graph_capturing, + slot_mappings=slot_mappings, + ) + + # We register layerwise NVTX hooks here after the first dynamo tracing is + # done to avoid nvtx operations in hook functions being traced by + # torch dynamo and causing graph breaks. + # Note that for DYNAMO_ONCE and VLLM_COMPILE mode, + # compiled model's dynamo tracing is only done once and the compiled model's + # __call__ function is replaced by calling the compiled function. + # So it's safe to register hooks here. Hooks will be registered to + # both compiled and uncompiled models but they will never + # be called on the compiled model execution path. + self._register_layerwise_nvtx_hooks() + + # This is necessary to avoid blocking DP. + # For dummy runs, we typically skip EPLB since we don't have any real + # requests to process. + # However, in DP settings, there may be cases when some DP ranks do + # not have any requests to process, so they're executing dummy batches. + # In such cases, we still have to trigger EPLB to make sure + # ranks execute the rearrangement in synchronization. + if not skip_eplb: + self.eplb_step(is_dummy=True, is_profile=is_profile) + + logit_indices = np.cumsum(num_scheduled_tokens) - 1 + logit_indices_device = torch.from_numpy(logit_indices).to( + self.device, non_blocking=True + ) + return hidden_states, hidden_states[logit_indices_device] + + @torch.inference_mode() + def _dummy_sampler_run( + self, + hidden_states: torch.Tensor, + ) -> torch.Tensor: + # The dummy hidden states may contain special values, + # like `inf` or `nan`. + # To avoid breaking the sampler, we use a random tensor here instead. + + mm_config = self.vllm_config.model_config.multimodal_config + if mm_config and mm_config.mm_encoder_only: + # MM Encoder only model no need to run sampler. + return torch.tensor([]) + + hidden_states = torch.rand_like(hidden_states) + + logits = self.model.compute_logits(hidden_states) + num_reqs = logits.size(0) + + dummy_tensors = lambda v: torch.full((num_reqs,), v, device=self.device) + + dummy_metadata = SamplingMetadata( + temperature=dummy_tensors(0.5), + all_greedy=False, + all_random=False, + top_p=dummy_tensors(0.9), + top_k=dummy_tensors(logits.size(1) - 1), + generators={}, + max_num_logprobs=None, + logprob_token_ids=None, + no_penalties=True, + prompt_token_ids=None, + frequency_penalties=dummy_tensors(0.1), + presence_penalties=dummy_tensors(0.1), + repetition_penalties=dummy_tensors(0.1), + output_token_ids=[[] for _ in range(num_reqs)], + spec_token_ids=[[] for _ in range(num_reqs)], + allowed_token_ids_mask=None, + bad_words_token_ids={}, + logitsprocs=LogitsProcessors(), + ) + try: + sampler_output = self.sampler( + logits=logits, sampling_metadata=dummy_metadata + ) + # Also warm forward_native (taken when generators dict is non-empty), + # but skip the extra call in 'processed_logits' / 'processed_logprobs' + # modes — there TopKTopPSampler binds forward = forward_native at + # init time, so the warmup call is redundant and only inflates peak + # memory during profile_run. + # No .clone() of logits: warmup output is discarded, so any in-place + # mutation by forward_native does not affect correctness. + if self.sampler.logprobs_mode not in ( + "processed_logits", + "processed_logprobs", + ): + self.sampler( + logits=logits, + sampling_metadata=replace( + dummy_metadata, + generators={ + 0: torch.Generator(device=self.device).manual_seed(0) + }, + ), + ) + except RuntimeError as e: + if "out of memory" in str(e): + raise RuntimeError( + "CUDA out of memory occurred when warming up sampler with " + f"{num_reqs} dummy requests. Please try lowering " + "`max_num_seqs` or `gpu_memory_utilization` when " + "initializing the engine." + ) from e + else: + raise e + if self.speculative_config: + draft_token_ids = [[0] for _ in range(num_reqs)] + dummy_spec_decode_metadata = SpecDecodeMetadata.make_dummy( + draft_token_ids, self.device + ) + + num_tokens = sum(len(ids) for ids in draft_token_ids) + draft_probs = None + if ( + self.speculative_config.rejection_sample_method == "standard" + and self.speculative_config.draft_sample_method == "probabilistic" + ): + draft_probs = torch.rand( + num_tokens, + logits.shape[-1], + device=self.device, + dtype=torch.float32, + ) + draft_probs = torch.softmax(draft_probs, dim=-1) + logits = torch.randn( + num_tokens + num_reqs, + logits.shape[-1], + device=self.device, + dtype=logits.dtype, + ) + self.rejection_sampler( + dummy_spec_decode_metadata, + draft_probs, + logits, + dummy_metadata, + ) + return sampler_output + + def _dummy_pooler_run_task( + self, + hidden_states: torch.Tensor, + task: PoolingTask, + ) -> PoolerOutput: + num_tokens = hidden_states.shape[0] + max_num_reqs = self.scheduler_config.max_num_seqs + num_reqs = min(num_tokens, max_num_reqs) + min_tokens_per_req = num_tokens // num_reqs + num_scheduled_tokens_np = np.full(num_reqs, min_tokens_per_req) + num_scheduled_tokens_np[-1] += num_tokens % num_reqs + assert np.sum(num_scheduled_tokens_np) == num_tokens + assert len(num_scheduled_tokens_np) == num_reqs + + req_num_tokens = num_tokens // num_reqs + + dummy_prompt_lens = torch.from_numpy(num_scheduled_tokens_np) + dummy_token_ids = torch.zeros( + (num_reqs, req_num_tokens), dtype=torch.int32, device=self.device + ) + + model = cast(VllmModelForPooling, self.get_model()) + dummy_pooling_params = PoolingParams(task=task) + dummy_pooling_params.verify(self.model_config) + to_update = model.pooler.get_pooling_updates(task) + to_update.apply(dummy_pooling_params) + + dummy_metadata = PoolingMetadata( + prompt_lens=dummy_prompt_lens, + prompt_token_ids=dummy_token_ids, + prompt_token_ids_cpu=dummy_token_ids.cpu(), + pooling_params=[dummy_pooling_params] * num_reqs, + pooling_states=[PoolingStates() for i in range(num_reqs)], + ) + + dummy_metadata.build_pooling_cursor( + num_scheduled_tokens_np, + seq_lens_cpu=dummy_prompt_lens, + device=hidden_states.device, + ) + + try: + return model.pooler( + hidden_states=hidden_states, pooling_metadata=dummy_metadata + ) + except RuntimeError as e: + if "out of memory" in str(e): + raise RuntimeError( + "CUDA out of memory occurred when warming up pooler " + f"({task=}) with {num_reqs} dummy requests. Please try " + "lowering `max_num_seqs` or `gpu_memory_utilization` when " + "initializing the engine." + ) from e + else: + raise e + + @torch.inference_mode() + def _dummy_pooler_run( + self, + hidden_states: torch.Tensor, + ) -> PoolerOutput: + mm_config = self.vllm_config.model_config.multimodal_config + if mm_config and mm_config.mm_encoder_only: + # MM Encoder only model not need to run pooler. + return torch.tensor([]) + + # Find the task that has the largest output for subsequent steps + supported_pooling_tasks = self.get_supported_pooling_tasks() + + if not supported_pooling_tasks: + raise RuntimeError( + f"Model {self.model_config.model} does not support " + "any pooling tasks. See " + "https://docs.vllm.ai/en/latest/models/pooling_models.html " + "to learn more." + ) + + output_size = dict[PoolingTask, float]() + for task in supported_pooling_tasks: + # Run a full batch with each task to ensure none of them OOMs + output = self._dummy_pooler_run_task(hidden_states, task) + output_size[task] = sum(o.nbytes for o in output if o is not None) + del output # Allow GC + + max_task = max(output_size.items(), key=lambda x: x[1])[0] + return self._dummy_pooler_run_task(hidden_states, max_task) + + def profile_run(self) -> None: + # Profile with multimodal encoder & encoder cache. + if self.supports_mm_inputs: + mm_config = self.model_config.multimodal_config + if mm_config is not None and mm_config.skip_mm_profiling: + logger.info( + "Skipping memory profiling for multimodal encoder and " + "encoder cache." + ) + else: + mm_budget = self.mm_budget + assert mm_budget is not None + + if (encoder_budget := mm_budget.get_encoder_budget()) > 0: + if not mm_budget.mm_max_toks_per_item: + # All modality limits are 0 — embedding-only mode. + # Budget is non-zero for embedding storage, but + # there's no encoder to profile. + logger.info( + "Skipping encoder profiling for embedding-only " + "mode (all modality limits=0 with " + "enable_mm_embeds=True).", + ) + else: + # NOTE: Currently model is profiled with a single + # non-text modality with the max possible input + # tokens even when it supports multiple. + dummy_modality = mm_budget.get_modality_with_max_tokens() + max_mm_items_per_batch = mm_budget.mm_max_items_per_batch[ + dummy_modality + ] + + logger.info_once( + "Encoder cache will be initialized with a " + "budget of %s tokens, and profiled with " + "%s %s items of the maximum feature size.", + encoder_budget, + max_mm_items_per_batch, + dummy_modality, + ) + + # Create dummy batch of multimodal inputs. + batched_dummy_mm_inputs = self._get_mm_dummy_batch( + dummy_modality, + max_mm_items_per_batch, + ) + + # Run multimodal encoder. + dummy_encoder_outputs = self.model.embed_multimodal( + **batched_dummy_mm_inputs + ) + + sanity_check_mm_encoder_outputs( + dummy_encoder_outputs, + expected_num_items=max_mm_items_per_batch, + ) + for i, output in enumerate(dummy_encoder_outputs): + self.encoder_cache[f"tmp_{i}"] = output + + # Add `is_profile` here to pre-allocate communication buffers + hidden_states, last_hidden_states = self._dummy_run( + self.max_num_tokens, is_profile=True + ) + if get_pp_group().is_last_rank: + if self.is_pooling_model: + output = self._dummy_pooler_run(hidden_states) + else: + output = self._dummy_sampler_run(last_hidden_states) + else: + output = None + self._sync_device() + del hidden_states, output + self.encoder_cache.clear() + gc.collect() + + def _init_minimal_kv_cache_for_profiling(self) -> None: + from vllm.v1.core.kv_cache_utils import ( + get_kv_cache_config_from_groups, + get_kv_cache_groups, + ) + + kv_cache_spec = self.get_kv_cache_spec() + kv_cache_groups = get_kv_cache_groups(self.vllm_config, kv_cache_spec) + min_blocks = self.compilation_config.max_cudagraph_capture_size or 1 + + # Temporarily change num_gpu_blocks_override to allocate a minimal KV cache + saved_override = self.cache_config.num_gpu_blocks_override + self.cache_config.num_gpu_blocks_override = min_blocks + minimal_config = get_kv_cache_config_from_groups( + self.vllm_config, kv_cache_groups, available_memory=0 + ) + self.cache_config.num_gpu_blocks_override = saved_override + + self.initialize_kv_cache(minimal_config, is_profiling=True) + self.cache_config.num_gpu_blocks = minimal_config.num_blocks + + logger.debug("Initialized minimal KV cache for CUDA graph profiling") + + @staticmethod + @contextmanager + def _freeze_gc(): + gc.collect() + should_freeze = not envs.VLLM_ENABLE_CUDAGRAPH_GC + if should_freeze: + gc.freeze() + try: + yield + finally: + if should_freeze: + gc.unfreeze() + gc.collect() + + def shutdown(self) -> None: + """Release GPU tensors (model weights, KV caches, workspace) so that + memory is reclaimable when running in the same process.""" + from vllm.model_executor.layers.rotary_embedding import _ROPE_DICT + from vllm.v1.worker.workspace import reset_workspace_manager + + # Calls torch.accelerator.synchronize() + self._cleanup_profiling_kv_cache() + self.compilation_config.static_forward_context.clear() + self.model = None # type: ignore[assignment] + _ROPE_DICT.clear() + + reset_workspace_manager() + + def _cleanup_profiling_kv_cache(self) -> None: + torch.accelerator.synchronize() + if hasattr(self, "kv_caches") and self.kv_caches: + for i in range(len(self.kv_caches)): + self.kv_caches[i] = None # type: ignore + self.kv_caches.clear() + if hasattr(self, "cross_layers_kv_cache"): + self.cross_layers_kv_cache = None + self.cross_layers_attn_backend = None + if hasattr(self, "attn_groups"): + self.attn_groups.clear() + if hasattr(self, "kv_cache_config"): + delattr(self, "kv_cache_config") + self.cache_config.num_gpu_blocks = None + + for layer in self.compilation_config.static_forward_context.values(): + if hasattr(layer, "kv_cache"): + kv_cache = layer.kv_cache + layer.kv_cache = ( + torch.tensor([]) if isinstance(kv_cache, torch.Tensor) else [] + ) + # Clean up quantized KV cache scale views + # (int8_per_token_head, fp8_per_token_head) + if hasattr(layer, "impl"): + if hasattr(layer.impl, "_k_scale_cache"): + layer.impl._k_scale_cache = None + if hasattr(layer.impl, "_v_scale_cache"): + layer.impl._v_scale_cache = None + + gc.collect() + torch.accelerator.empty_cache() + + logger.debug("Cleaned up profiling KV cache and CUDA graphs") + + @torch.inference_mode() + def profile_cudagraph_memory(self) -> int: + with set_current_vllm_config(self.vllm_config): + self._init_minimal_kv_cache_for_profiling() + + saved_num_cudagraph_captured = compilation_counter.num_cudagraph_captured + + capture_descs = self.cudagraph_dispatcher.get_capture_descs() + + total_graphs = sum(len(descs) for _, descs in capture_descs) + if total_graphs == 0: + logger.debug("No CUDA graphs will be captured, skipping profiling") + self._cleanup_profiling_kv_cache() + return 0 + + logger.info( + "Profiling CUDA graph memory: %s", + ", ".join( + f"{mode.name}={len(descs)} (largest={descs[0].num_tokens})" + for mode, descs in capture_descs + if descs + ), + ) + + # Use a temporary pool for profiling to avoid fragmentation in the main pool. + profiling_pool = current_platform.graph_pool_handle() + original_pools: dict[int, Any] = {} + all_wrappers = list(CUDAGraphWrapper._all_instances) + list( + BreakableCUDAGraphWrapper._all_instances + ) + for instance in all_wrappers: + original_pools[id(instance)] = instance.graph_pool + instance.graph_pool = profiling_pool + + set_cudagraph_capturing_enabled(True) + with self._freeze_gc(), graph_capture(device=self.device): + shared_memory_estimate = {} + per_graph_estimate = {} + torch.accelerator.synchronize() + torch.accelerator.empty_cache() + + for mode, descs in capture_descs: + profile_descs = descs[:2] + mem_samples: list[int] = [] + + for i, desc in enumerate(profile_descs): + mem_before = torch.cuda.mem_get_info()[0] + self._warmup_and_capture( + desc, + cudagraph_runtime_mode=mode, + profile_seq_lens=( + min( + self.max_model_len, + self.max_num_tokens // desc.num_tokens, + ) + if mode == CUDAGraphMode.FULL and i == 0 + else None + ), + ) + torch.accelerator.synchronize() + free_after = torch.cuda.mem_get_info()[0] + mem_samples.append(mem_before - free_after) + + first_capture = mem_samples[0] + # Use at least 1 MiB per graph for driver overhead + per_graph = max(mem_samples[1] if len(mem_samples) > 1 else 0, 1 << 20) + + shared_memory_estimate[mode] = first_capture + per_graph_estimate[mode] = per_graph * (len(descs) - 1) + + logger.debug( + "Estimated %s CUDA graph memory: " + "%.2f MiB first-capture + (%d-1) × %.2f MiB per-graph", + mode.name, + first_capture / (1 << 20), + len(descs), + per_graph / (1 << 20), + ) + + set_cudagraph_capturing_enabled(False) + CUDAGraphWrapper.clear_all_graphs() + BreakableCUDAGraphWrapper.clear_all_graphs() + all_wrappers = list(CUDAGraphWrapper._all_instances) + list( + BreakableCUDAGraphWrapper._all_instances + ) + for instance in all_wrappers: + if id(instance) in original_pools: + instance.graph_pool = original_pools[id(instance)] + for key_set in self.cudagraph_dispatcher.cudagraph_keys.values(): + key_set.clear() + self.cudagraph_dispatcher.keys_initialized = False + self.maybe_remove_all_loras(self.lora_config) + self._cleanup_profiling_kv_cache() + compilation_counter.num_cudagraph_captured = saved_num_cudagraph_captured + + # FULL and PIECEWISE graphs share the global pool at runtime and are + # never replayed concurrently, so the pool overlays their memory. + # Take the max to avoid double-counting the overlap. + total_estimate = max(shared_memory_estimate.values()) + sum( + per_graph_estimate.values() + ) + logger.info( + "Estimated CUDA graph memory: %.2f GiB total", + total_estimate / (1 << 30), + ) + + return int(total_estimate) + + @instrument(span_name="Capture model") + def capture_model(self) -> int: + if self.compilation_config.cudagraph_mode == CUDAGraphMode.NONE: + logger.warning( + "Skipping CUDA graph capture. To turn on CUDA graph capture, " + "ensure `cudagraph_mode` was not manually set to `NONE`" + ) + return 0 + + # Initialize encoder CUDA graph manager if enabled. + # Use get_model() to unwrap CUDAGraphWrapper/UBatchWrapper, + # because @runtime_checkable Protocol isinstance() checks do not + # work through __getattr__ forwarding. + if ( + self.compilation_config.cudagraph_mm_encoder + and self.supports_mm_inputs + and self.encoder_cudagraph_manager is None + ): + from vllm.model_executor.models.interfaces import ( + SupportsEncoderCudaGraph, + supports_encoder_cudagraph, + ) + from vllm.v1.worker.encoder_cudagraph import ( + EncoderCudaGraphManager, + ) + + raw_model = self.get_model() + if supports_encoder_cudagraph(raw_model): + self.encoder_cudagraph_manager = EncoderCudaGraphManager( + vllm_config=self.vllm_config, + device=self.device, + dtype=self.dtype, + model=cast(SupportsEncoderCudaGraph, raw_model), + ) + logger.info("Initialized EncoderCudaGraphManager for vision encoder") + + compilation_counter.num_gpu_runner_capture_triggers += 1 + + start_time = time.perf_counter() + + # Trigger CUDA graph capture for specific shapes. + # Capture the large shapes first so that the smaller shapes + # can reuse the memory pool allocated for the large shapes. + set_cudagraph_capturing_enabled(True) + with self._freeze_gc(), graph_capture(device=self.device): + torch.accelerator.synchronize() + torch.accelerator.empty_cache() + start_free_gpu_memory = torch.cuda.mem_get_info()[0] + + for ( + runtime_mode, + batch_descs, + ) in self.cudagraph_dispatcher.get_capture_descs(): + self._capture_cudagraphs( + batch_descriptors=batch_descs, + cudagraph_runtime_mode=runtime_mode, + ) + torch.accelerator.synchronize() + + # Capture encoder CUDA graphs if enabled + if self.encoder_cudagraph_manager is not None: + self.encoder_cudagraph_manager.capture() + + torch.accelerator.synchronize() + end_free_gpu_memory = torch.cuda.mem_get_info()[0] + + # Disable cudagraph capturing globally, so any unexpected cudagraph + # capturing will be detected and raise an error after here. + # Note: We don't put it into graph_capture context manager because + # we may do lazy capturing in future that still allows capturing + # after here. + set_cudagraph_capturing_enabled(False) + + torch.accelerator.synchronize() + torch.accelerator.empty_cache() + + # Lock workspace to prevent resizing during execution. + # Max workspace sizes should have been captured during warmup/profiling. + lock_workspace() + + end_time = time.perf_counter() + elapsed_time = end_time - start_time + cuda_graph_size = start_free_gpu_memory - end_free_gpu_memory + # This usually takes 5~20 seconds. + logger.info_once( + "Graph capturing finished in %.0f secs, took %.2f GiB", + elapsed_time, + cuda_graph_size / (1 << 30), + ) + return cuda_graph_size + + def _warmup_and_capture( + self, + desc: BatchDescriptor, + cudagraph_runtime_mode: CUDAGraphMode, + profile_seq_lens: int | None = None, + allow_microbatching: bool = False, + num_warmups: int | None = None, + ): + if num_warmups is None: + num_warmups = self.compilation_config.cudagraph_num_of_warmups + force_attention = cudagraph_runtime_mode == CUDAGraphMode.FULL + for _ in range(num_warmups): + self._dummy_run( + desc.num_tokens, + cudagraph_runtime_mode=CUDAGraphMode.NONE, + force_attention=force_attention, + uniform_decode=desc.uniform, + allow_microbatching=allow_microbatching, + skip_eplb=True, + remove_lora=False, + num_active_loras=desc.num_active_loras, + profile_seq_lens=profile_seq_lens, + ) + self._dummy_run( + desc.num_tokens, + cudagraph_runtime_mode=cudagraph_runtime_mode, + uniform_decode=desc.uniform, + allow_microbatching=allow_microbatching, + skip_eplb=True, + remove_lora=False, + num_active_loras=desc.num_active_loras, + is_graph_capturing=True, + profile_seq_lens=profile_seq_lens, + ) + + def _capture_cudagraphs( + self, + batch_descriptors: list[BatchDescriptor], + cudagraph_runtime_mode: CUDAGraphMode, + ): + assert ( + cudagraph_runtime_mode != CUDAGraphMode.NONE + and cudagraph_runtime_mode.is_valid_runtime_mode() + ), f"Invalid cudagraph runtime mode: {cudagraph_runtime_mode}" + + if not batch_descriptors: + return + + uniform_decode = batch_descriptors[0].uniform + + # Only rank 0 should print progress bar during capture + if is_global_first_rank(): + batch_descriptors = tqdm( + batch_descriptors, + disable=not self.load_config.use_tqdm_on_load, + desc="Capturing CUDA graphs ({}, {})".format( + "decode" if uniform_decode else "mixed prefill-decode", + cudagraph_runtime_mode.name, + ), + ) + + # We skip EPLB here since we don't want to record dummy metrics + for batch_desc in batch_descriptors: + # We currently only capture ubatched graphs when its a FULL + # cudagraph, a uniform decode batch, and the number of tokens + # is above the threshold. Otherwise we just capture a non-ubatched + # version of the graph + allow_microbatching = ( + self.parallel_config.use_ubatching + and cudagraph_runtime_mode == CUDAGraphMode.FULL + and uniform_decode + and check_ubatch_thresholds( + config=self.vllm_config.parallel_config, + num_tokens=batch_desc.num_tokens, + uniform_decode=uniform_decode, + ) + ) + self._warmup_and_capture( + batch_desc, + cudagraph_runtime_mode=cudagraph_runtime_mode, + allow_microbatching=allow_microbatching, + ) + torch.accelerator.synchronize() + self.maybe_remove_all_loras(self.lora_config) + + def initialize_attn_backend( + self, + kv_cache_config: KVCacheConfig, + is_profiling: bool = False, + ) -> None: + """ + Initialize the attention backends and attention metadata builders. + """ + assert len(self.attn_groups) == 0, "Attention backends are already initialized" + + class AttentionGroupKey(NamedTuple): + """Deduplication key for attention groups within a KV cache group. + + Splits on per-rank ``num_heads_q`` in addition to backend + spec + so layers with different Q-head counts (e.g. a spec-decode draft + with fewer attention heads than its target) get separate metadata + builders. The builders' scratch (e.g. ``softmax_segm_*`` in + ``triton_attn``, ``num_qo_heads`` in FlashInfer) is sized by + ``num_heads_q`` and assumes uniformity within the group; see + ``get_num_attention_heads_from_layers`` in + ``vllm/v1/attention/backends/utils.py``. + """ + + attn_backend: type[AttentionBackend] + kv_cache_spec: KVCacheSpec + num_heads_q: int + + def get_attn_backends_for_group( + kv_cache_group_spec: KVCacheGroupSpec, + ) -> tuple[dict[AttentionGroupKey, list[str]], set[type[AttentionBackend]]]: + layer_type = cast(type[Any], AttentionLayerBase) + layers = get_layers_from_vllm_config( + self.vllm_config, layer_type, kv_cache_group_spec.layer_names + ) + attn_backends = {} + attn_backend_layers = defaultdict(list) + # Dedupe based on full class name; this is a bit safer than + # using the class itself as the key because when we create dynamic + # attention backend subclasses (e.g. ChunkedLocalAttention) unless + # they are cached correctly, there will be different objects per + # layer. + for layer_name in kv_cache_group_spec.layer_names: + attn_backend = layers[layer_name].get_attn_backend() + + if layer_name in self.kv_sharing_fast_prefill_eligible_layers: + attn_backend = create_fast_prefill_custom_backend( + "FastPrefill", + attn_backend, # type: ignore[arg-type] + ) + + full_cls_name = attn_backend.full_cls_name() + layer_kv_cache_spec = kv_cache_group_spec.kv_cache_spec + if isinstance(layer_kv_cache_spec, UniformTypeKVCacheSpecs): + layer_kv_cache_spec = layer_kv_cache_spec.kv_cache_specs[layer_name] + # Non-Attention layer types (e.g. Mamba1, ShortConv) do not + # expose ``num_heads``; fall back to 0 so they cluster as + # before. Such layers never coexist with Attention in a + # single KV cache group (different KVCacheSpec), so the + # fallback can never spuriously merge them with attention + # layers. + num_heads_q = getattr(layers[layer_name], "num_heads", 0) + key = (full_cls_name, layer_kv_cache_spec, num_heads_q) + attn_backends[key] = AttentionGroupKey( + attn_backend, layer_kv_cache_spec, num_heads_q + ) + attn_backend_layers[key].append(layer_name) + return ( + {attn_backends[k]: v for k, v in attn_backend_layers.items()}, + set(group_key.attn_backend for group_key in attn_backends.values()), + ) + + def create_attn_groups( + attn_backends_map: dict[AttentionGroupKey, list[str]], + kv_cache_group_id: int, + ) -> list[AttentionGroup]: + attn_groups: list[AttentionGroup] = [] + for key, layer_names in attn_backends_map.items(): + attn_group = AttentionGroup( + key.attn_backend, + layer_names, + key.kv_cache_spec, + kv_cache_group_id, + ) + + attn_groups.append(attn_group) + return attn_groups + + attention_backend_maps = [] + attention_backend_list = [] + for kv_cache_group_spec in kv_cache_config.kv_cache_groups: + attn_backends = get_attn_backends_for_group(kv_cache_group_spec) + attention_backend_maps.append(attn_backends[0]) + attention_backend_list.append(attn_backends[1]) + + # Resolve cudagraph_mode before actually initialize metadata_builders + self._check_and_update_cudagraph_mode( + attention_backend_list, + kv_cache_config.kv_cache_groups, + is_profiling=is_profiling, + ) + + # Check if attention backend supports PCP&DCP and related features. + check_attention_cp_compatibility(self.vllm_config) + + for i, attn_backend_map in enumerate(attention_backend_maps): + self.attn_groups.append(create_attn_groups(attn_backend_map, i)) + + def initialize_metadata_builders( + self, kv_cache_config: KVCacheConfig, kernel_block_sizes: list[int] + ) -> None: + """ + Create the metadata builders for all KV cache groups and attn groups. + """ + for kv_cache_group_id in range(len(kv_cache_config.kv_cache_groups)): + for attn_group in self.attn_groups[kv_cache_group_id]: + attn_group.create_metadata_builders( + self.vllm_config, + self.device, + kernel_block_sizes[kv_cache_group_id] + if kv_cache_group_id < len(kernel_block_sizes) + else None, + num_metadata_builders=1 + if not self.parallel_config.use_ubatching + else self.parallel_config.num_ubatches, + ) + # Calculate reorder batch threshold (if needed) + # Note (tdoublep): do this *after* constructing builders, + # because some of them change the threshold at init time. + self.calculate_reorder_batch_threshold() + + # Initialize drafter attention backend + if self.speculative_config and ( + self.speculative_config.use_eagle() + or self.speculative_config.uses_draft_model() + or self.speculative_config.use_dspark() + ): + assert isinstance( + self.drafter, + EagleProposer + | DFlashProposer + | DSparkProposer + | DraftModelProposer + | Gemma4Proposer, + ) + self.drafter.initialize_attn_backend(kv_cache_config, kernel_block_sizes) + + def _check_and_update_cudagraph_mode( + self, + attention_backends: list[set[type[AttentionBackend]]], + kv_cache_groups: list[KVCacheGroupSpec], + is_profiling: bool = False, + ) -> None: + """ + Resolve the cudagraph_mode when there are multiple attention + groups with potential conflicting CUDA graph support. + Then initialize the cudagraph_dispatcher based on the resolved + cudagraph_mode. + """ + min_cg_support = AttentionCGSupport.ALWAYS + min_cg_attn_backend = None + + for attn_backend_set, kv_cache_group in zip( + attention_backends, kv_cache_groups + ): + for attn_backend in attn_backend_set: + builder_cls = attn_backend.get_builder_cls() + + cg_support = builder_cls.get_cudagraph_support( + self.vllm_config, kv_cache_group.kv_cache_spec + ) + if cg_support.value < min_cg_support.value: + min_cg_support = cg_support + min_cg_attn_backend = attn_backend.__name__ + cudagraph_mode = self.compilation_config.resolve_cudagraph_mode_and_sizes( + min_cg_support, + min_cg_attn_backend, + self.uniform_decode_query_len, + self.parallel_config.tensor_parallel_size, + self.kv_cache_config, + self.max_num_reqs, + is_profiling=is_profiling, + ) + # Trigger cudagraph dispatching keys initialization after + # resolved cudagraph mode. + self.cudagraph_dispatcher.initialize_cudagraph_keys( + cudagraph_mode, self.uniform_decode_query_len + ) + + # Initialize drafter's cudagraph dispatcher if using spec decode. + if self.speculative_config and ( + self.speculative_config.use_eagle() + or self.speculative_config.uses_extract_hidden_states() + or self.speculative_config.use_dspark() + ): + assert isinstance( + self.drafter, + EagleProposer + | DFlashProposer + | DSparkProposer + | ExtractHiddenStatesProposer + | Gemma4Proposer, + ) + self.drafter.initialize_cudagraph_keys(cudagraph_mode) + + def calculate_reorder_batch_threshold(self) -> None: + """ + Choose the minimum reorder batch threshold from all attention groups. + Backends should be able to support lower threshold then what they request + just may have a performance penalty due to that backend treating decodes + as prefills. + """ + min_none_high = lambda a, b: a if b is None else b if a is None else min(a, b) + + reorder_batch_thresholds: list[int | None] = [ + group.get_metadata_builder().reorder_batch_threshold + for group in self._attn_group_iterator() + ] + # If there are no attention groups (attention-free model) or no backend + # reports a threshold, leave reordering disabled. + if len(reorder_batch_thresholds) == 0: + self.reorder_batch_threshold = None + return + self.reorder_batch_threshold = reduce(min_none_high, reorder_batch_thresholds) # type: ignore[assignment] + + def _set_mm_prefix_range_for_metadata( + self, + attn_metadata: Any, + req_doc_ranges: dict[int, list[tuple[int, int]]], + ) -> None: + """Set mm_prefix_range for all attention metadata objects. + + This method handles both list and non-list attention metadata, + computing mm_prefix_range_tensor once and sharing it across all + metadata objects to avoid redundant host-to-device transfers. + """ + from vllm.v1.attention.backends.triton_attn import ( + TritonAttentionMetadata, + ) + + # Get all metadata objects from either list or dict structure + metadata_list = [] + if isinstance(attn_metadata, list): + for ub_metadata in attn_metadata: + metadata_list.extend(ub_metadata.values()) + else: + metadata_list.extend(attn_metadata.values()) + + # Set mm_prefix_range for all metadata and compute tensor once + shared_tensor = None + for metadata in metadata_list: + metadata.mm_prefix_range = req_doc_ranges # type: ignore[attr-defined] + + # Only compute tensor for TritonAttentionMetadata + if isinstance(metadata, TritonAttentionMetadata): + if shared_tensor is None: + shared_tensor = ( + TritonAttentionMetadata.compute_mm_prefix_range_tensor( + req_doc_ranges, + metadata.seq_lens.shape[0], # type: ignore[attr-defined] + metadata.seq_lens.device, # type: ignore[attr-defined] + ) + ) + metadata.mm_prefix_range_tensor = shared_tensor + + def may_reinitialize_input_batch( + self, kv_cache_config: KVCacheConfig, kernel_block_sizes: list[int] + ) -> 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 + placeholder used during __init__, or when there are multiple + KV cache groups. + + Args: + kv_cache_config: The KV cache configuration. + kernel_block_sizes: The kernel block sizes for each KV cache group. + """ + block_sizes = [] + max_num_blocks = [] + max_model_len = max(self.max_model_len, self.max_encoder_len) + for kv_cache_group in kv_cache_config.kv_cache_groups: + if isinstance(kv_cache_group.kv_cache_spec, EncoderOnlyAttentionSpec): + continue + block_size = kv_cache_group.kv_cache_spec.block_size + block_sizes.append(block_size) + max_num_blocks_per_req = cdiv( + max_model_len, block_size * get_total_cp_world_size() + ) + if isinstance(kv_cache_group.kv_cache_spec, MambaSpec): + max_num_blocks_per_req = ( + max_num_blocks_per_req + if self.cache_config.enable_prefix_caching + else 1 + ) + kv_cache_group.kv_cache_spec.num_speculative_blocks + max_num_blocks.append(max_num_blocks_per_req) + + if ( + block_sizes != self._init_block_sizes + or kernel_block_sizes != self._init_kernel_block_sizes + ): + self._init_block_sizes = block_sizes + self._init_kernel_block_sizes = kernel_block_sizes + self.input_batch = InputBatch( + max_num_reqs=self.max_num_reqs, + max_model_len=max_model_len, + max_num_batched_tokens=self.max_num_tokens, + device=self.device, + pin_memory=self.pin_memory, + 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, + logitsprocs_need_output_token_ids=self.input_batch.logitsprocs_need_output_token_ids, + is_pooling_model=self.is_pooling_model, + reasoning_config=self.vllm_config.reasoning_config, + ) + + assert self._init_block_sizes == block_sizes, ( + 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 _allocate_kv_cache_tensors( + self, kv_cache_config: KVCacheConfig + ) -> dict[str, torch.Tensor]: + """ + Initializes the KV cache buffer with the correct size. The buffer needs + to be reshaped to the desired shape before being used by the models. + + Args: + kv_cache_config: The KV cache config + Returns: + dict[str, torch.Tensor]: A map between layer names to their + corresponding memory buffer for KV cache. + """ + kv_cache_raw_tensors: dict[str, torch.Tensor] = {} + for kv_cache_tensor in kv_cache_config.kv_cache_tensors: + tensor = torch.zeros( + kv_cache_tensor.size, dtype=torch.int8, device=self.device + ) + for layer_name in kv_cache_tensor.shared_by: + kv_cache_raw_tensors[layer_name] = tensor + + layer_names = set() + for group in kv_cache_config.kv_cache_groups: + for layer_name in group.layer_names: + if layer_name in self.runner_only_attn_layers: + continue + layer_names.add(layer_name) + assert layer_names == set(kv_cache_raw_tensors.keys()), ( + "Some layers are not correctly initialized" + ) + return kv_cache_raw_tensors + + def _attn_group_iterator(self) -> Iterator[AttentionGroup]: + return itertools.chain.from_iterable(self.attn_groups) + + def _kv_cache_spec_attn_group_iterator(self) -> Iterator[AttentionGroup]: + if not self.kv_cache_config.kv_cache_groups: + return + for attn_groups in self.attn_groups: + yield from attn_groups + + def _reshape_kv_cache_tensors( + self, + kv_cache_raw_tensors: dict[str, torch.Tensor], + kernel_block_sizes: list[int], + ) -> dict[str, torch.Tensor]: + """ + Reshape the KV cache tensors to the desired shape and dtype. + + Args: + kv_cache_raw_tensors: The KV cache buffer of each layer, with + correct size but uninitialized shape. + kernel_block_sizes: The kernel block sizes for each KV cache group. + Returns: + Dict[str, torch.Tensor]: A map between layer names to their + corresponding memory buffer for KV cache. + """ + kv_caches: dict[str, torch.Tensor] = {} + has_attn, has_mamba = False, False + for group in self._kv_cache_spec_attn_group_iterator(): + kv_cache_spec = group.kv_cache_spec + attn_backend = group.backend + if group.kv_cache_group_id == len(kernel_block_sizes): + # There may be a last group for layers without kv cache. + continue + kernel_block_size = kernel_block_sizes[group.kv_cache_group_id] + for layer_name in group.layer_names: + if layer_name in self.runner_only_attn_layers: + continue + raw_tensor = kv_cache_raw_tensors[layer_name] + assert raw_tensor.numel() % kv_cache_spec.page_size_bytes == 0 + num_blocks = raw_tensor.numel() // kv_cache_spec.page_size_bytes + if isinstance(kv_cache_spec, AttentionSpec): + has_attn = True + num_blocks_per_kv_block = ( + kv_cache_spec.block_size // kernel_block_size + ) + kernel_num_blocks = num_blocks * num_blocks_per_kv_block + + # For MLA with compression, storage_block_size != block_size + if kv_cache_spec.storage_block_size != kv_cache_spec.block_size: + shape_block_size = kv_cache_spec.storage_block_size + else: + shape_block_size = kernel_block_size + + kv_cache_shape = attn_backend.get_kv_cache_shape( + kernel_num_blocks, + shape_block_size, + kv_cache_spec.num_kv_heads, + kv_cache_spec.head_size, + cache_dtype_str=self.cache_config.cache_dtype, + ) + dtype = kv_cache_spec.dtype + try: + kv_cache_stride_order = attn_backend.get_kv_cache_stride_order() + assert len(kv_cache_stride_order) == len(kv_cache_shape) + except (AttributeError, NotImplementedError): + kv_cache_stride_order = tuple(range(len(kv_cache_shape))) + # The allocation respects the backend-defined stride order + # to ensure the semantic remains consistent for each + # backend. We first obtain the generic kv cache shape and + # then permute it according to the stride order which could + # result in a non-contiguous tensor. + kv_cache_shape = tuple( + kv_cache_shape[i] for i in kv_cache_stride_order + ) + # Maintain original KV shape view. + inv_order = [ + kv_cache_stride_order.index(i) + for i in range(len(kv_cache_stride_order)) + ] + + raw_tensor = kv_cache_raw_tensors[layer_name].view(dtype) + if kv_cache_spec.page_size_padded is not None: + # Use strided view to handle page_size_bytes that + # include padding. This follows + # the same pattern as MambaSpec handling below. + # NOTE: This assumes kv_cache_shape[0] == num_blocks + # (i.e. the first physical dimension is the block + # index), which holds for MLA backends but NOT for + # standard attention backends whose shape starts with + # a K/V dimension of size 2. + dtype_size = get_dtype_size(dtype) + page_stride = kv_cache_spec.page_size_bytes // dtype_size + strides = list(torch.empty(kv_cache_shape).stride()) + strides[inv_order[0]] = page_stride + kv_cache = torch.as_strided( + raw_tensor, + size=kv_cache_shape, + stride=tuple(strides), + ) + else: + # No padding — safe to use a contiguous view. + kv_cache = raw_tensor.view(kv_cache_shape) + kv_caches[layer_name] = kv_cache.permute(*inv_order) + + elif isinstance(kv_cache_spec, MambaSpec): + has_mamba = True + raw_tensor = kv_cache_raw_tensors[layer_name] + state_tensors = [] + storage_offset_bytes = 0 + for shape, dtype in zip(kv_cache_spec.shapes, kv_cache_spec.dtypes): + dtype_size = get_dtype_size(dtype) + num_element_per_page = ( + kv_cache_spec.page_size_bytes // dtype_size + ) + target_shape = (num_blocks, *shape) + stride = torch.empty(target_shape).stride() + target_stride = (num_element_per_page, *stride[1:]) + assert storage_offset_bytes % dtype_size == 0 + tensor = torch.as_strided( + raw_tensor.view(dtype), + size=target_shape, + stride=target_stride, + storage_offset=storage_offset_bytes // dtype_size, + ) + state_tensors.append(tensor) + storage_offset_bytes += stride[0] * dtype_size + + kv_caches[layer_name] = state_tensors + else: + raise NotImplementedError + + if has_attn and has_mamba: + self._update_hybrid_attention_mamba_layout(kv_caches, kernel_block_sizes) + + return kv_caches + + def _update_hybrid_attention_mamba_layout( + self, kv_caches: dict[str, torch.Tensor], kernel_block_sizes: list[int] + ) -> None: + """ + Update the layout of attention layers from (2, num_blocks, ...) to + (num_blocks, 2, ...). + + Args: + kv_caches: The KV cache buffer of each layer. + kernel_block_sizes: The kernel block sizes for each KV cache group. + """ + + for group in self._kv_cache_spec_attn_group_iterator(): + kv_cache_spec = group.kv_cache_spec + if not isinstance(kv_cache_spec, AttentionSpec): + continue + block_dim = group.backend.get_kv_cache_block_dim( + kernel_block_sizes[group.kv_cache_group_id], + kv_cache_spec.num_kv_heads, + kv_cache_spec.head_size, + cache_dtype_str=self.cache_config.cache_dtype, + ) + # block_dim: 0 means (num_blocks, 2, ...); 1 means (2, num_blocks, ...). + if block_dim == 0: + continue + assert block_dim == 1 + for layer_name in group.layer_names: + kv_cache = kv_caches[layer_name] + hidden_size = kv_cache.shape[2:].numel() + kv_cache.as_strided_( + size=kv_cache.shape, + stride=(hidden_size, 2 * hidden_size, *kv_cache.stride()[2:]), + ) + + def initialize_kv_cache_tensors( + self, kv_cache_config: KVCacheConfig, kernel_block_sizes: list[int] + ) -> dict[str, torch.Tensor]: + """ + Initialize the memory buffer for KV cache. + + Args: + kv_cache_config: The KV cache config + kernel_block_sizes: The kernel block sizes for each KV cache group. + + Returns: + Dict[str, torch.Tensor]: A map between layer names to their + corresponding memory buffer for KV cache. + """ + + # Try creating KV caches optimized for kv-connector transfers + cache_dtype = self.cache_config.cache_dtype + if self.use_uniform_kv_cache(self.attn_groups, cache_dtype): + kv_caches, cross_layers_kv_cache, attn_backend = ( + self.allocate_uniform_kv_caches( + kv_cache_config, + self.attn_groups, + cache_dtype, + self.device, + kernel_block_sizes, + ) + ) + self.cross_layers_kv_cache = cross_layers_kv_cache + self.cross_layers_attn_backend = attn_backend + else: + # Fallback to the general case + # Initialize the memory buffer for KV cache + kv_cache_raw_tensors = self._allocate_kv_cache_tensors(kv_cache_config) + + # Change the memory buffer to the desired shape + kv_caches = self._reshape_kv_cache_tensors( + kv_cache_raw_tensors, kernel_block_sizes + ) + + # Set up cross-layer KV cache sharing + for layer_name, target_layer_name in self.shared_kv_cache_layers.items(): + logger.debug("%s reuses KV cache of %s", layer_name, target_layer_name) + kv_caches[layer_name] = kv_caches[target_layer_name] + + num_attn_module = ( + 2 if self.model_config.hf_config.model_type == "longcat_flash" else 1 + ) + bind_kv_cache( + kv_caches, + self.compilation_config.static_forward_context, + self.kv_caches, + num_attn_module, + ) + return kv_caches + + def maybe_add_kv_sharing_layers_to_kv_cache_groups( + self, kv_cache_config: KVCacheConfig + ) -> None: + """ + Add layers that re-use KV cache to KV cache group of its target layer. + Mapping of KV cache tensors happens in `initialize_kv_cache_tensors()` + """ + if not self.shared_kv_cache_layers: + # No cross-layer KV sharing, return + return + + add_kv_sharing_layers_to_kv_cache_groups( + self.shared_kv_cache_layers, + kv_cache_config.kv_cache_groups, + self.runner_only_attn_layers, + ) + + if self.cache_config.kv_sharing_fast_prefill: + # In You Only Cache Once (https://arxiv.org/abs/2405.05254) or other + # similar KV sharing setups, only the layers that generate KV caches + # are involved in the prefill phase, enabling prefill to early exit. + attn_layers = get_layers_from_vllm_config(self.vllm_config, Attention) + for layer_name in reversed(attn_layers): + if layer_name in self.shared_kv_cache_layers: + self.kv_sharing_fast_prefill_eligible_layers.add(layer_name) + else: + break + + def initialize_kv_cache( + self, + kv_cache_config: KVCacheConfig, + is_profiling: bool = False, + ) -> None: + """ + Initialize KV cache based on `kv_cache_config`. + Args: + kv_cache_config: Configuration for the KV cache, including the KV + cache size of each layer + """ + kv_cache_config = deepcopy(kv_cache_config) + self.kv_cache_config = kv_cache_config + self._mamba_bufs = None + self.may_add_encoder_only_layers_to_kv_cache_config() + self.maybe_add_kv_sharing_layers_to_kv_cache_groups(kv_cache_config) + self.initialize_attn_backend(kv_cache_config, is_profiling=is_profiling) + initialize_mamba_ssu_backend( + self.vllm_config.mamba_config, self.kv_cache_config + ) + # The kernel block size for all KV cache groups. For example, if + # kv_cache_manager uses block_size 256 for a given group, but the attention + # 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. + kernel_block_sizes = prepare_kernel_block_sizes( + kv_cache_config, self.attn_groups + ) + self._kernel_block_sizes = kernel_block_sizes + + # create metadata builders + self.initialize_metadata_builders(kv_cache_config, kernel_block_sizes) + + # Reinitialize need to after initialize_attn_backend + self.may_reinitialize_input_batch(kv_cache_config, kernel_block_sizes) + kv_caches = self.initialize_kv_cache_tensors( + kv_cache_config, kernel_block_sizes + ) + + if ( + self.speculative_config + and self.speculative_config.uses_extract_hidden_states() + ): + assert isinstance(self.drafter, ExtractHiddenStatesProposer) + # validate all draft model layers belong to the same kv cache + # group + self.drafter.validate_same_kv_cache_group(kv_cache_config) + + if has_kv_transfer_group() and not is_profiling: + kv_transfer_group = get_kv_transfer_group() + if self.cross_layers_kv_cache is not None: + assert self.cross_layers_attn_backend is not None + kv_transfer_group.register_cross_layers_kv_cache( + self.cross_layers_kv_cache, self.cross_layers_attn_backend + ) + else: + kv_transfer_group.register_kv_caches(kv_caches) + kv_transfer_group.set_host_xfer_buffer_ops(copy_kv_blocks) + + def _get_attention_kv_cache_gid(self) -> int: + """Find the KV cache group index for attention layers. + + Must match :attr:`RoutedExpertsManager.attn_gid` in the scheduler: + both pick the first ``FullAttentionSpec`` group so hybrid models + (Mamba / linear-attention layers that use other AttentionSpec + subclasses) end up indexing the same slot layout on both sides. + Falls back to 0 only for legacy single-group configs. + """ + for gid, group in enumerate(self.kv_cache_config.kv_cache_groups): + if isinstance(group.kv_cache_spec, FullAttentionSpec): + return gid + return 0 + + def init_routed_experts_capturer(self): + logger.info( + "Initializing routed experts capturer, enable_return_routed_experts: %s", + self.model_config.enable_return_routed_experts, + ) + self.routed_experts_capturer = RoutedExpertsCapturer( + max_num_batched_tokens=self.scheduler_config.max_num_batched_tokens, + vllm_config=self.vllm_config, + ) + self.routed_experts_attn_gid = self._get_attention_kv_cache_gid() + self._bind_routed_experts_capturer(self.routed_experts_capturer) + + # Pinned CPU buffer for non-blocking D2H of ``routing_data`` on + # the sync scheduling path. Shape / dtype mirror the device + # capturer exactly so ``copy_`` is a straight memcpy. + self.routed_experts_cpu = torch.empty( + self.routed_experts_capturer.device_buffer.shape, + dtype=self.routed_experts_capturer.device_buffer.dtype, + device="cpu", + pin_memory=self.pin_memory, + ) + # ``slot_mapping`` dtype is fixed to int64 by + # ``block_table.slot_mapping``; we mirror that here. + max_tokens = self.scheduler_config.max_num_batched_tokens + self.routed_experts_slot_mapping_cpu = torch.empty( + (max_tokens,), + dtype=torch.int64, + device="cpu", + pin_memory=self.pin_memory, + ) + # Private device buffer so the shared ``block_table.slot_mapping`` + # can be overwritten by the next ``_prepare_inputs`` while the + # D2H is still pending on the copy stream. Written in + # ``_prepare_inputs``, read in ``_bookkeeping_sync`` (sync path) + # or cloned into a snapshot (async path). + self.routed_experts_slot_mapping_device = torch.empty( + (max_tokens,), + dtype=torch.int64, + device=self.device, + ) + self.routed_experts_initialized = True + + def _bind_routed_experts_capturer(self, capturer: RoutedExpertsCapturer) -> None: + from vllm.model_executor.layers.fused_moe.layer import FusedMoE + from vllm.model_executor.layers.fused_moe.router.base_router import ( + BaseRouter, + ) + + for module in self.compilation_config.static_forward_context.values(): + if isinstance(module, FusedMoE) and isinstance(module.router, BaseRouter): + layer_id = module.layer_id + + def _capture_fn(topk_ids, _layer_id=layer_id, _capturer=capturer): + _capturer.capture(_layer_id, topk_ids) + + module.router.set_capture_fn(_capture_fn) + + def may_add_encoder_only_layers_to_kv_cache_config(self) -> None: + """ + Add encoder-only layers to the KV cache config. + """ + block_size = self.vllm_config.cache_config.block_size + encoder_only_attn_specs: dict[AttentionSpec, list[str]] = defaultdict(list) + attn_layers = get_layers_from_vllm_config(self.vllm_config, Attention) + for layer_name, attn_module in attn_layers.items(): + if attn_module.attn_type == AttentionType.ENCODER_ONLY: + attn_spec: AttentionSpec = EncoderOnlyAttentionSpec( + block_size=block_size, + num_kv_heads=attn_module.num_kv_heads, + head_size=attn_module.head_size, + dtype=self.kv_cache_dtype, + ) + encoder_only_attn_specs[attn_spec].append(layer_name) + self.runner_only_attn_layers.add(layer_name) + if len(encoder_only_attn_specs) > 0: + assert len(encoder_only_attn_specs) == 1, ( + "Only support one encoder-only attention spec now" + ) + spec, layer_names = encoder_only_attn_specs.popitem() + self.kv_cache_config.kv_cache_groups.append( + KVCacheGroupSpec(layer_names=layer_names, kv_cache_spec=spec) + ) + + def get_kv_cache_spec(self) -> dict[str, KVCacheSpec]: + """ + Generates the KVCacheSpec by parsing the kv cache format from each + Attention module in the static forward context. + Returns: + KVCacheSpec: A dictionary mapping layer names to their KV cache + format. Layers that do not need KV cache are not included. + """ + if has_ec_transfer() and not get_ec_transfer().is_consumer: + return {} + kv_cache_spec: dict[str, KVCacheSpec] = {} + layer_type = cast(type[Any], AttentionLayerBase) + attn_layers = get_layers_from_vllm_config(self.vllm_config, layer_type) + for layer_name, attn_module in attn_layers.items(): + if isinstance(attn_module, Attention) and ( + kv_tgt_layer := attn_module.kv_sharing_target_layer_name + ): + # The layer doesn't need its own KV cache and will use that of + # the target layer. We skip creating a KVCacheSpec for it, so + # that KV cache management logic will act as this layer does + # not exist, and doesn't allocate KV cache for the layer. This + # enables the memory saving of cross-layer kv sharing, allowing + # a given amount of memory to accommodate longer context lengths + # or enable more requests to be processed simultaneously. + self.shared_kv_cache_layers[layer_name] = kv_tgt_layer + continue + # Skip modules that don't need KV cache (eg encoder-only attention) + if spec := attn_module.get_kv_cache_spec(self.vllm_config): + kv_cache_spec[layer_name] = spec + + return kv_cache_spec + + def _to_list(self, sampled_token_ids: torch.Tensor) -> list[list[int]]: + # This is a short term mitigation for issue mentioned in + # https://github.com/vllm-project/vllm/issues/22754. + # `tolist` would trigger a cuda wise stream sync, which + # would block other copy ops from other cuda streams. + # A cuda event sync would avoid such a situation. Since + # this is in the critical path of every single model + # forward loop, this has caused perf issue for a disagg + # setup. + pinned = self.sampled_token_ids_pinned_cpu[: sampled_token_ids.shape[0]] + pinned.copy_(sampled_token_ids, non_blocking=True) + self.transfer_event.record() + self.transfer_event.synchronize() + return pinned.tolist() + + def get_encoder_timing_stats(self) -> dict[str, dict[str, float | int]]: + """ + Get encoder timing stats for all requests and clear the registry. + + Returns: + Dictionary mapping request_id to stats dict. + """ + with self._encoder_timing_lock: + stats = { + req_id: stats_obj.to_dict() + for req_id, stats_obj in self.encoder_timing_registry.items() + } + self.encoder_timing_registry.clear() + return stats + + @contextmanager + def timed_encoder_operation( + self, + should_time: bool, + group_lora_refs: list[tuple[str, Any]], + current_item_idx: int, + num_items: int, + ): + """ + Context manager to time encoder forward operations. + + Args: + should_time: Whether timing is enabled + group_lora_refs: Full list of (request_id, pos_info) tuples + current_item_idx: Starting index for this group + num_items: Number of items in this group + """ + if not should_time: + yield + return + + group_refs = group_lora_refs[current_item_idx : current_item_idx + num_items] + group_request_ids = {req_id for req_id, _ in group_refs} + + torch.accelerator.synchronize() + start_time = time.perf_counter() + + try: + yield + finally: + torch.accelerator.synchronize() + elapsed = time.perf_counter() - start_time + + per_request_time = elapsed / max(len(group_request_ids), 1) + + with self._encoder_timing_lock: + for req_id in group_request_ids: + if req_id not in self.encoder_timing_registry: + self.encoder_timing_registry[req_id] = EncoderTimingStats() + + stats = self.encoder_timing_registry[req_id] + stats.encoder_forward_secs += per_request_time + stats.num_encoder_calls += 1 + + +@dataclass +class EncoderTimingStats: + """Per-request timing statistics for encoder forward pass.""" + + encoder_forward_secs: float = 0.0 + """Time spent in vision encoder forward pass (seconds).""" + + num_encoder_calls: int = 0 + """Number of times encoder was called for this request.""" + + def to_dict(self) -> dict[str, float | int]: + return { + "encoder_forward_secs": self.encoder_forward_secs, + "num_encoder_calls": self.num_encoder_calls, + } diff --git a/mods/deepseek-v4-flash-dspark/patch-nvfp4-stage-a.py b/mods/deepseek-v4-flash-dspark/patch-nvfp4-stage-a.py new file mode 100755 index 00000000..3080dd0e --- /dev/null +++ b/mods/deepseek-v4-flash-dspark/patch-nvfp4-stage-a.py @@ -0,0 +1,46 @@ +#!/usr/bin/env python3 +"""NVFP4 Stage A: dtype plumbing — auto-detect vLLM site-packages.""" +import os, sys +from pathlib import Path + +site_packages = os.environ.get("VLLM_SITE_PACKAGES") +if not site_packages: + for p in ["/opt/env/lib/python3.12/site-packages", "/usr/local/lib/python3.12/dist-packages"]: + if Path(p, "vllm").exists(): + site_packages = p + break +if not site_packages: + sys.exit(0) + +root = Path(site_packages) / "vllm" + +def replace(path, old, new): + p = root / path + if not p.exists(): + return + text = p.read_text() + if new in text: + return + if old not in text: + print(f" [SKIP] {path}: anchor not found") + return + p.write_text(text.replace(old, new, 1)) + print(f" [OK] {path}: patched") + +print("NVFP4 Stage A: dtype plumbing") +replace("config/cache.py", + ' "fp8_ds_mla",\n "turboquant_k8v4",', + ' "fp8_ds_mla",\n "nvfp4_ds_mla",\n "turboquant_k8v4",') +replace("utils/torch_utils.py", + ' "fp8_ds_mla": torch.uint8,\n "turboquant_k8v4": torch.uint8,', + ' "fp8_ds_mla": torch.uint8,\n "nvfp4_ds_mla": torch.uint8,\n "turboquant_k8v4": torch.uint8,') +replace("utils/torch_utils.py", + ' or kv_cache_dtype == "nvfp4"\n', + ' or kv_cache_dtype == "nvfp4"\n or kv_cache_dtype == "nvfp4_ds_mla"\n') +replace("v1/kv_cache_interface.py", + ' if kv_cache_dtype == "nvfp4":\n return KVQuantMode.NVFP4\n', + ' if kv_cache_dtype == "nvfp4":\n return KVQuantMode.NVFP4\n if kv_cache_dtype == "nvfp4_ds_mla":\n return KVQuantMode.NVFP4\n') +replace("v1/kv_cache_interface.py", + ' if self.cache_dtype_str == "fp8_ds_mla":\n', + ' if self.cache_dtype_str == "nvfp4_ds_mla":\n return self.storage_block_size * 416\n if self.cache_dtype_str == "fp8_ds_mla":\n') +print("Stage A complete") diff --git a/mods/deepseek-v4-flash-dspark/patch-nvfp4-stage-b.py b/mods/deepseek-v4-flash-dspark/patch-nvfp4-stage-b.py new file mode 100755 index 00000000..16a83524 --- /dev/null +++ b/mods/deepseek-v4-flash-dspark/patch-nvfp4-stage-b.py @@ -0,0 +1,34 @@ +#!/usr/bin/env python3 +"""NVFP4 Stage B: probe path — auto-detect vLLM site-packages.""" +import os, sys +from pathlib import Path + +site_packages = os.environ.get("VLLM_SITE_PACKAGES") +if not site_packages: + for p in ["/opt/env/lib/python3.12/site-packages", "/usr/local/lib/python3.12/dist-packages"]: + if Path(p, "vllm").exists(): + site_packages = p + break +if not site_packages: + sys.exit(0) + +root = Path(site_packages) / "vllm" + +def replace(path, old, new): + p = root / path + if not p.exists(): + return + text = p.read_text() + if new in text: + return + if old not in text: + print(f" [SKIP] {path}: anchor not found") + return + p.write_text(text.replace(old, new, 1)) + print(f" [OK] {path}: patched") + +print("NVFP4 Stage B: probe path") +replace("models/deepseek_v4/attention.py", + ' # TODO(yifan): currently hardcoded for FP8 sparse, make it more generic\n head_bytes = (\n self.nope_head_dim # 448 fp8 NoPE\n + self.rope_head_dim * 2 # 64 bf16 RoPE\n + self.nope_head_dim // 64 # 7B scale factors\n + 1 # 1B pad\n )\n', + ' # TODO(yifan): currently hardcoded for FP8 sparse, make it more generic\n head_bytes = (\n self.nope_head_dim # 448 fp8 NoPE\n + self.rope_head_dim * 2 # 64 bf16 RoPE\n + self.nope_head_dim // 64 # 7B scale factors\n + 1 # 1B pad\n )\n if (\n cache_config is not None\n and cache_config.cache_dtype in ("nvfp4", "nvfp4_ds_mla")\n ):\n # Probe layout from the GLM-5.2 NVFP4 sparse-MLA path.\n head_bytes = 416\n') +print("Stage B complete") diff --git a/mods/deepseek-v4-flash-dspark/patch-nvfp4-stage-c.py b/mods/deepseek-v4-flash-dspark/patch-nvfp4-stage-c.py new file mode 100755 index 00000000..b9edf647 --- /dev/null +++ b/mods/deepseek-v4-flash-dspark/patch-nvfp4-stage-c.py @@ -0,0 +1,34 @@ +#!/usr/bin/env python3 +"""NVFP4 Stage C: padded envelope — auto-detect vLLM site-packages.""" +import os, sys +from pathlib import Path + +site_packages = os.environ.get("VLLM_SITE_PACKAGES") +if not site_packages: + for p in ["/opt/env/lib/python3.12/site-packages", "/usr/local/lib/python3.12/dist-packages"]: + if Path(p, "vllm").exists(): + site_packages = p + break +if not site_packages: + sys.exit(0) + +root = Path(site_packages) / "vllm" + +def replace(path, old, new): + p = root / path + if not p.exists(): + return + text = p.read_text() + if new in text: + return + if old not in text: + print(f" [SKIP] {path}: anchor not found") + return + p.write_text(text.replace(old, new, 1)) + print(f" [OK] {path}: patched") + +print("NVFP4 Stage C: padded envelope") +replace("models/deepseek_v4/attention.py", + ' if (\n cache_config is not None\n and cache_config.cache_dtype in ("nvfp4", "nvfp4_ds_mla")\n ):\n # Probe layout from the GLM-5.2 NVFP4 sparse-MLA path.\n head_bytes = 416\n', + ' if (\n cache_config is not None\n and cache_config.cache_dtype in ("nvfp4", "nvfp4_ds_mla")\n ):\n # Stage C: keep DeepSeek V4 proven 584-byte cache envelope.\n head_bytes = 584\n') +print("Stage C complete") diff --git a/mods/deepseek-v4-flash-dspark/run.sh b/mods/deepseek-v4-flash-dspark/run.sh new file mode 100755 index 00000000..5750dd13 --- /dev/null +++ b/mods/deepseek-v4-flash-dspark/run.sh @@ -0,0 +1,72 @@ +#!/bin/bash +set -euo pipefail +# DeepSeek V4 Flash DSpark NVFP4 overlay mod +# Auto-detects vLLM site-packages location (handles both /opt/env/ and /usr/local/ layouts) + +# Auto-detect vLLM site-packages +if [ -d "/opt/env/lib/python3.12/site-packages/vllm" ]; then + SITE_PACKAGES="/opt/env/lib/python3.12/site-packages" +elif [ -d "/usr/local/lib/python3.12/dist-packages/vllm" ]; then + SITE_PACKAGES="/usr/local/lib/python3.12/dist-packages" +elif [ -n "${PYTHON_ROOT:-}" ] && [ -d "$PYTHON_ROOT/vllm" ]; then + SITE_PACKAGES="$PYTHON_ROOT" +else + # Try python3 -c + SITE_PACKAGES=$(python3 -c "import vllm; import os; print(os.path.dirname(os.path.dirname(vllm.__file__)))" 2>/dev/null || echo "") + if [ -z "$SITE_PACKAGES" ] || [ ! -d "$SITE_PACKAGES/vllm" ]; then + echo "[dsv4-dspark] ERROR: Cannot find vLLM installation" >&2 + python3 -c "import vllm; print(vllm.__file__)" 2>/dev/null || echo "vLLM not importable" + exit 1 + fi +fi +echo "[dsv4-dspark] Detected vLLM at: $SITE_PACKAGES/vllm" + +MOD_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +OVERLAY_DIR="$MOD_DIR/overlay" + +# 1. Copy all overlay files into site-packages +echo "[dsv4-dspark] Copying overlay files..." +OVERLAY_FILES=$(find "$OVERLAY_DIR" -type f -name "*.py" | sort) +COUNT=0 +for src in $OVERLAY_FILES; do + rel="${src#$OVERLAY_DIR/}" + dst="$SITE_PACKAGES/$rel" + mkdir -p "$(dirname "$dst")" + cp "$src" "$dst" + COUNT=$((COUNT + 1)) +done +echo "[dsv4-dspark] Copied $COUNT overlay files" + +# 2. Run NVFP4 patch stages if available (pass site-packages path) +export VLLM_SITE_PACKAGES="$SITE_PACKAGES" +for stage in stage-a stage-b stage-c; do + patch_script="$MOD_DIR/patch-nvfp4-$stage.py" + if [ -f "$patch_script" ]; then + echo "[dsv4-dspark] Running NVFP4 Stage $stage..." + python3 "$patch_script" || { + echo "[dsv4-dspark] WARNING: Stage $stage patch failed (may be pre-applied)" + } + fi +done + +# 3. Compile overlaid files +echo "[dsv4-dspark] Compiling overlaid files..." +cd "$SITE_PACKAGES" +for f in $(find "$OVERLAY_DIR" -type f -name "*.py"); do + rel="${f#$OVERLAY_DIR/}" + python3 -m py_compile "$SITE_PACKAGES/$rel" 2>/dev/null || true +done + +# 4. Verify imports +echo "[dsv4-dspark] Verifying imports..." +python3 -c " +from vllm.v1.spec_decode import dspark, dspark_proposer +print(f"DSpark: OK") +from typing import get_args +from vllm.config.cache import CacheDType +if \"nvfp4_ds_mla\" in get_args(CacheDType): + print(f"NVFP4 dtype: present") +else: + print(f"WARNING: nvfp4_ds_mla not in CacheDType") +print(f"DeepSeek V4 Flash DSpark overlay applied successfully") +" 2>&1 && echo "[dsv4-dspark] Overlay applied successfully" || echo "[dsv4-dspark] WARNING: Verification had issues" diff --git a/recipes/deepseek-v4-flash-dspark.yaml b/recipes/deepseek-v4-flash-dspark.yaml new file mode 100644 index 00000000..b1703d9b --- /dev/null +++ b/recipes/deepseek-v4-flash-dspark.yaml @@ -0,0 +1,134 @@ +# Recipe: DeepSeek V4 Flash DSpark C12 NVFP4 +# DSpark speculative decoding + NVFP4 KV cache on 2x DGX Spark (GB10) +# +# Based on tonyd2wild's verified C12 NVFP4 profile (2026-07-04). +# https://github.com/tonyd2wild/DeepSeek-v4-Flash-DSpark-1M-NVFP4-KV-2x-DGX-Spark +# +# Key community sources: +# - tonyd2wild C12 NVFP4 default config (canonical 350K verified profile) +# - MiaAI-Lab DSpark concurrency patch (Keys C12) +# - 0rand, renek: DSpark cold-start garble root cause analysis +# - Aiden Le: production-3.2 image reference, fp8 fallback guidance +# +# The DSpark speculative decode garble fix (probabilistic draft, graph-capture +# sizing, async-scheduling, chunked-prefill) is incorporated — see the AGENT_GARBLE_FIX +# discussion in tonyd2wild's repo for the full diagnosis. +# +# Requires: +# - build-and-copy.sh with --apply-vllm-pr 46995 (DSpark PR) +# - mods/deepseek-v4-flash-dspark mod for overlay files + NVFP4 patches +# - Two-node cluster (back-to-back or switch), user apollo11 +# - Model cached at ~/.cache/huggingface (see prepare-dspark-model-cache.sh) +# +# Usage: +# ./run-recipe.py deepseek-v4-flash-dspark --solo # single node debug +# ./run-recipe.py deepseek-v4-flash-dspark # 2-node cluster +# +# Build flags (pass before recipe name): +# ./run-recipe.py --build-args "--apply-vllm-pr 46995" deepseek-v4-flash-dspark +# +# Uses HF repo ID for model resolution via HF_HOME at container runtime. +# The model cache at ~/.cache/huggingface is bind-mounted into the container. + +recipe_version: "1" +name: DeepSeek-V4-Flash-DSpark +description: "DeepSeek V4 Flash DSpark C12 NVFP4 — DSpark spec decode, NVFP4 KV cache, B12X MoE on 2x DGX Spark" + +model: fraserprice/DeepSeek-V4-Flash-DSpark +container: vllm-dspark-runtime:dspark-nvfp4-stage-c + +# Two-node cluster required (back-to-back or switch) +cluster_only: true + +# DSpark overlay mod — copies overlay files + runs NVFP4 patches at container start +mods: + - mods/deepseek-v4-flash-dspark + +# Default settings (override via CLI) +# Community-verified (tonyd2wild C12 NVFP4, 2026-07-04): gpu_mem=0.80, +# max_model_len=350000, seqs=12, btokens=8192, spec_tokens=5. +# Gigabyte Atom (DGX Spark head, less free GPU mem): use gpu_mem=0.77, +# max_model_len=450000 to account for lower available memory (~94 GiB free +# vs ~103 GiB on a clean Spark). +defaults: + port: 8000 + host: 0.0.0.0 + tensor_parallel: 2 + model: fraserprice/DeepSeek-V4-Flash-DSpark + gpu_memory_utilization: 0.77 + max_model_len: 350000 + block_size: 256 + max_num_seqs: 12 + max_num_batched_tokens: 8192 + num_speculative_tokens: 5 + kv_cache_dtype: nvfp4_ds_mla + +# Environment variables — DSpark + B12X + NVFP4 config +# Set HF_HOME to match the container's HF cache bind mount path. +env: + HF_HOME: "/root/.cache/huggingface" + HF_HUB_OFFLINE: "1" + TRANSFORMERS_OFFLINE: "1" + HF_HUB_DISABLE_XET: "1" + VLLM_CACHE_ROOT: "/root/.cache/huggingface/vllm-cache" + VLLM_ALLOW_LONG_MAX_MODEL_LEN: "1" + VLLM_USE_FLASHINFER_SAMPLER: "1" + VLLM_USE_B12X_MOE: "1" + VLLM_USE_B12X_WO_PROJECTION: "1" + VLLM_DSPARK_CONFIDENCE_THRESHOLD: "0.0" + VLLM_DSPARK_CONFIDENCE_SCHEDULER: "off" + VLLM_DSPARK_LOCAL_ARGMAX: "1" + VLLM_DSPARK_REPLICATE_MARKOV_W1: "1" + VLLM_DSPARK_FUSED_MARKOV_ARGMAX: "0" + VLLM_DSPARK_GPU_REJECTED_CONTEXT_MASK: "1" + VLLM_DSPARK_REFERENCE_KV_QUANT_DEQUANT: "0" + VLLM_DSPARK_HARDWARE_SCHEDULER_EARLY_STOP: "1" + VLLM_DSV4_B12X_COMPRESSED_MLA: "0" + VLLM_DSV4_DSPARK_DEFER_TARGET_CAPTURE: "0" + VLLM_DSV4_DSPARK_DEFER_TARGET_CAPTURE_EXACT: "0" + VLLM_TRITON_MLA_SPARSE: "1" + VLLM_SPARSE_INDEXER_MAX_LOGITS_MB: "256" + VLLM_MEMORY_PROFILER_ESTIMATE_CUDAGRAPHS: "0" + VLLM_SKIP_INIT_MEMORY_CHECK: "1" + TORCH_CUDA_ARCH_LIST: "12.1a" + FLASHINFER_CUDA_ARCH_LIST: "12.1a" + FLASHINFER_DISABLE_VERSION_CHECK: "1" + TILELANG_CLEANUP_TEMP_FILES: "1" + NCCL_NET: "IB" + NCCL_IB_DISABLE: "0" + NCCL_CUMEM_ENABLE: "0" + NCCL_IGNORE_CPU_AFFINITY: "1" + NCCL_NVLS_ENABLE: "0" + PYTORCH_CUDA_ALLOC_CONF: "expandable_segments:True" + +# The vLLM serve command template +# Note: --nnodes, --node-rank, --master-addr, --master-port, --headless are +# automatically appended by launch-cluster.sh in cluster mode. +command: |- + vllm serve {model} \ + --host {host} \ + --port {port} \ + --trust-remote-code \ + --tensor-parallel-size {tensor_parallel} \ + --pipeline-parallel-size 1 \ + --kv-cache-dtype {kv_cache_dtype} \ + --block-size {block_size} \ + --max-model-len {max_model_len} \ + --max-num-seqs {max_num_seqs} \ + --max-num-batched-tokens {max_num_batched_tokens} \ + --max-cudagraph-capture-size {max_num_seqs} \ + --gpu-memory-utilization {gpu_memory_utilization} \ + --enable-prefix-caching \ + --async-scheduling \ + --enable-chunked-prefill \ + --speculative-config '{{"method":"dspark","num_speculative_tokens":{num_speculative_tokens},"draft_sample_method":"probabilistic"}}' \ + --tokenizer-mode deepseek_v4 \ + --tool-call-parser deepseek_v4 \ + --enable-auto-tool-choice \ + --reasoning-parser deepseek_v4 \ + --reasoning-config '{{"reasoning_parser":"deepseek_v4","reasoning_start_str":"","reasoning_end_str":""}}' \ + --default-chat-template-kwargs '{{"thinking":false}}' \ + --generation-config vllm \ + --enable-flashinfer-autotune \ + --distributed-executor-backend mp \ + --served-model-name deepseek-v4-flash-dspark