iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[Mini-VLLM] Remove unreferenced definitions found by static scan

Browse Source 
A symbol scan over every tracked file found 123 top-level definitions in
vllm/ with zero references, including references inside strings. All of
them came from features this fork deleted earlier.

- Delete dead attention helpers: MLADCPManager, split_prefill_chunks, the
  reshape helpers, compute_varlen_chunk_metadata, the LINEAR backend's
  unused BailingLinearAttentionBackend, and the sparse flash-attn entry
  points.
- Delete the CPU and AMX op wrappers from _custom_ops.py: cpu_mla_decode,
  cpu_mla_extend, cpu_mla_get_scheduler_metadata,
  gather_mla_context_cache_cpu.
- Delete dead quantization helpers: awq_pack, gptq_pack,
  gptq_quantize_weights, get_and_maybe_dequant_weights,
  marlin_quant_fp8_torch, the rand_marlin_weight_*_like pair,
  ocp_mx_weight_dtype_and_rows, get_fp8_block_weight_scale,
  input_to_float8.
- Delete dead model helpers: resolve_visual_encoder_outputs,
  run_dp_sharded_vision_model, fast_topk, get_layer_index,
  is_pp_missing_parameter, skip_spec_layers, KimiK3MergedQKVGateLinear,
  selective_scan_fn, causal_conv1d_update_torch.
- Delete dead utils: format_kib, format_mib, chunk_list, prev_power_of_2,
  split_host_port, join_host_port, make_zmq_path, make_async_with_semaphore,
  is_meta_module, is_non_overlapping_and_dense, supports_xpu_graph,
  create_kv_caches_with_random*, and the unused has_* capability probes.
- Delete the unreachable deep-gemm mega-MoE warmup path, which imported
  vllm.models.deepseek_v4.nvidia.model and vllm.models.kimi_k3.nvidia.model,
  neither of which exists anymore.
- Drop the imports and the two now-empty guards left behind.

Verified statically: every changed file compiles, no unused import
remains, and no deleted name is referenced in vllm, tests or csrc. The
runtime gates still need the Python environment rebuilt.

Co-authored-by: Hinata <hinata@idihore.id>
Signed-off-by: MOVZX <movzx@yahoo.com>
This commit is contained in:
MOVZX
2026-10-19 21:39:56 +07:00
co-authored by Hinata
parent 5aa2ac64f8
commit 61fbe0b865
72 changed files with 15 additions and 3385 deletions
Download Patch File Download Diff File Expand all files Collapse all files
requirements/common.txt
+2 -3
View File
@@ -13,7 +13,7 @@ tokenizers >= 0.23.2 # Required for fast incremental detokenization.
safetensors >= 0.6.2 # MXFP4/MXFP6 dtype support (F8_E8M0, F4) added in 0.6.0: https://github.com/huggingface/safetensors/pull/611
protobuf >= 5.29.6, !=6.30.*, !=6.31.*, !=6.32.*, !=6.33.0.*, !=6.33.1.*, !=6.33.2.*, !=6.33.3.*, !=6.33.4.*; platform_machine != "s390x"
protobuf == 7.36.1; platform_machine == "s390x"
fastapi[standard] >= 0.133.0, < 0.137.0 # First version supporting Starlette 1.0; < 0.137.0 avoids route-tree change that breaks model-hosting-container-standards handler overrides.
fastapi[standard] >= 0.133.0, < 0.137.0 # First version supporting Starlette 1.0
starlette >= 1.0.1 # CVE-2026-48710: Host header injection in < 1.0.1
aiohttp >= 3.14.3
openai >= 2.25.0 # For Responses API with namespace tools
@@ -40,13 +40,12 @@ compressed-tensors == 0.17.0 # required for compressed-tensors
depyf==0.20.0 # required for profiling and debugging with compilation config
cloudpickle # allows pickling lambda functions in model_executor/models/registry.py
watchfiles # required for http server to monitor the updates of TLS files
python-json-logger # Used by logging as per examples/features/logging_configuration.md
python-json-logger # Used by the logging formatters in vllm/logger.py
ninja # Required for xgrammar, rocm, tpu, xpu
pybase64 # fast base64 implementation
cbor2 # Required for cross-language serialization of hashable objects
setproctitle # Used to set process names for better debugging and monitoring
oss-harmony >= 0.0.11 # Required for gpt-oss; includes embedded vocabularies
model-hosting-container-standards >= 0.1.14, < 1.0.0
mcp >= 2.0.0, < 3.0.0
opentelemetry-sdk >= 1.27.0
opentelemetry-api >= 1.27.0
vllm/_custom_ops.py
-163
View File
@@ -100,43 +100,6 @@ if hasattr(torch.ops, "_C") and hasattr(torch.ops._C, "scaled_fp4_quant"):
return create_fp4_output_tensors(m, n, input.device, is_sf_swizzled_layout)
# page attention ops
def mla_decode_kvcache_cpu(
out: torch.Tensor,
query: torch.Tensor,
kv_cache: torch.Tensor,
scale: float,
block_tables: torch.Tensor,
seq_lens: torch.Tensor,
) -> None:
torch.ops._C.mla_decode_kvcache(out, query, kv_cache, scale, block_tables, seq_lens)
def gather_mla_context_cache_cpu(
*,
src_cache: torch.Tensor,
dst: torch.Tensor,
block_table: torch.Tensor,
starts: torch.Tensor,
cu_seq_lens: torch.Tensor,
) -> None:
page_size = src_cache.shape[1]
flat_cache = src_cache.view(-1, src_cache.shape[-1])
seq_lens = cu_seq_lens[1:] - cu_seq_lens[:-1]
out_start = 0
for req_idx, (start, seq_len) in enumerate(zip(starts.tolist(), seq_lens.tolist())):
if seq_len <= 0:
continue
positions = torch.arange(start, start + seq_len, device=block_table.device)
block_ids = block_table[req_idx, positions // page_size].to(torch.long)
offsets = positions % page_size
slots = block_ids * page_size + offsets
out_end = out_start + seq_len
dst[out_start:out_end].copy_(flat_cache[slots])
out_start = out_end
# merge attn states ops
def merge_attn_states(
output: torch.Tensor,
@@ -3554,39 +3517,6 @@ def onednn_scaled_mm(
return output
def cpu_attn_get_scheduler_metadata(
num_reqs: int,
num_heads: int,
num_kv_heads: int,
head_dim: int,
seq_lens: torch.Tensor,
dtype: torch.dtype,
query_start_loc: torch.Tensor,
causal: bool,
sliding_window_size: int,
isa: str,
enable_kv_split: bool,
dynamic_causal: torch.Tensor | None = None,
kv_cache_dtype: str = "auto",
) -> torch.Tensor:
scheduler_metadata = torch.ops._C.get_scheduler_metadata(
num_reqs,
num_heads,
num_kv_heads,
head_dim,
seq_lens,
dtype,
query_start_loc,
causal,
sliding_window_size,
isa,
enable_kv_split,
dynamic_causal,
kv_cache_dtype,
)
return scheduler_metadata
def cpu_attn_reshape_and_cache(
key: torch.Tensor,
value: torch.Tensor,
@@ -3653,90 +3583,6 @@ def cpu_attention_with_kv_cache(
)
def cpu_mla_decode(
query: torch.Tensor,
k_buffer: torch.Tensor,
v_buffer: torch.Tensor,
output: torch.Tensor,
key: torch.Tensor | None,
value: torch.Tensor | None,
loc: torch.Tensor | None,
attn_logits: torch.Tensor,
req_to_token: torch.Tensor,
req_pool_indices: torch.Tensor,
seq_lens: torch.Tensor,
sm_scale: float,
logit_cap: float,
is_cross_attn: bool,
sliding_window_size: int,
encoder_lens: torch.Tensor | None,
sinks: torch.Tensor | None,
) -> None:
torch.ops._C.decode_attention_cpu(
query,
k_buffer,
v_buffer,
output,
key,
value,
loc,
attn_logits,
req_to_token,
req_pool_indices,
seq_lens,
sm_scale,
logit_cap,
is_cross_attn,
sliding_window_size,
encoder_lens,
sinks,
)
def cpu_mla_extend(
q_extend: torch.Tensor,
k_extend: torch.Tensor | None,
v_extend: torch.Tensor | None,
o_extend: torch.Tensor,
k_buffer: torch.Tensor,
v_buffer: torch.Tensor,
req_to_token: torch.Tensor,
req_pool_indices: torch.Tensor,
seq_lens: torch.Tensor,
extend_seq_lens: torch.Tensor,
extend_start_loc: torch.Tensor,
max_len_extend: int,
sm_scale: float,
logit_cap: float,
is_cross_attn: bool,
sliding_window_size: int,
encoder_lens: torch.Tensor | None,
sinks: torch.Tensor | None,
tree_mask: torch.Tensor | None = None,
) -> None:
torch.ops._C.extend_attention_cpu(
q_extend,
k_extend,
v_extend,
o_extend,
k_buffer,
v_buffer,
req_to_token,
req_pool_indices,
seq_lens,
extend_seq_lens,
extend_start_loc,
max_len_extend,
sm_scale,
logit_cap,
is_cross_attn,
sliding_window_size,
encoder_lens,
sinks,
tree_mask,
)
def bmm_cpu(
out: torch.Tensor,
mat1: torch.Tensor,
@@ -3747,15 +3593,6 @@ def bmm_cpu(
torch.ops._C.bmm_cpu(out, mat1, mat2, is_vnni, scale)
def amx_mla_concat_and_cache(
kv_c_normed: torch.Tensor,
k_pe: torch.Tensor,
kv_cache: torch.Tensor,
slot_mapping: torch.Tensor,
) -> None:
torch.ops._C.concat_and_cache_mla_cpu(kv_c_normed, k_pe, kv_cache, slot_mapping)
if hasattr(torch.ops._C, "hc_pre_fused_cpu"):
@register_fake("_C::hc_pre_fused_cpu")
vllm/compilation/passes/fusion/matcher_utils.py
-84
View File
@@ -24,9 +24,6 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
kNvfp4Dynamic,
)
from vllm.model_executor.layers.rotary_embedding import RotaryEmbedding
from vllm.model_executor.layers.rotary_embedding.deepseek_scaling_rope import (
DeepseekScalingRotaryEmbedding,
)
from vllm.platforms import current_platform
ROTARY_OP = torch.ops._C.rotary_embedding.default
@@ -292,87 +289,6 @@ class MatcherRMSNormGated(MatcherCustomOp):
)
class MatcherDeepseekScalingRotaryEmbedding(MatcherCustomOp):
def __init__(
self,
is_neox: bool,
head_size: int,
num_heads: int,
num_kv_heads: int,
use_flashinfer: bool = False,
enabled: bool | None = None,
) -> None:
if enabled is None:
enabled = DeepseekScalingRotaryEmbedding.enabled()
super().__init__(enabled)
self.is_neox = is_neox
self.head_size = head_size
self.num_heads = num_heads
self.num_kv_heads = num_kv_heads
self.q_size = self.num_heads * self.head_size
self.kv_size = self.num_kv_heads * self.head_size
self.rotary_dim = head_size
self.use_flashinfer = use_flashinfer
def inputs(self) -> list[torch.Tensor]:
positions = self.empty_int64(5)
query = self.empty(5, self.num_heads, self.head_size)
key = self.empty(5, self.num_kv_heads, self.head_size)
cos_sin_cache = self.empty(4096, self.rotary_dim)
return [positions, query, key, cos_sin_cache]
def forward_custom(
self,
positions: torch.Tensor,
query: torch.Tensor,
key: torch.Tensor | None,
cos_sin_cache: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor | None]:
if self.use_flashinfer:
torch.ops.vllm.flashinfer_rotary_embedding(
positions,
query,
key,
self.head_size,
cos_sin_cache,
self.is_neox,
)
return query, key
result: tuple[torch.Tensor, torch.Tensor | None] = (
DeepseekScalingRotaryEmbedding.forward_static(
positions,
query,
key,
self.head_size,
self.rotary_dim,
cos_sin_cache,
self.is_neox,
)
)
return result
def forward_native(
self,
positions: torch.Tensor,
query: torch.Tensor,
key: torch.Tensor | None,
cos_sin_cache: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor | None]:
result: tuple[torch.Tensor, torch.Tensor | None] = (
DeepseekScalingRotaryEmbedding.forward_static(
positions,
query,
key,
self.head_size,
self.rotary_dim,
cos_sin_cache,
self.is_neox,
)
)
return result
class MatcherQuantFP8(MatcherCustomOp):
def __init__(
self,
vllm/compilation/passes/fx_utils.py
-7
View File
@@ -21,13 +21,6 @@ def find_auto_fn_maybe(nodes: Iterable[fx.Node], op: OpOverload) -> fx.Node | No
return None
# Returns the first auto_functionalized node with the given op
def find_auto_fn(nodes: Iterable[fx.Node], op: OpOverload) -> fx.Node:
node = find_auto_fn_maybe(nodes, op)
assert node is not None, f"Could not find {op} in nodes {nodes}"
return node
# An auto-functionalization-aware utility for finding nodes with a specific op
# Also handles op overload packets and finds all overloads
def find_op_nodes(
vllm/compilation/passes/vllm_inductor_pass.py
-21
View File
@@ -277,27 +277,6 @@ def fold_consecutive_reshapes(gm: fx.GraphModule) -> None:
gm.graph.erase_node(inp)
def remove_noop_reshapes(gm: fx.GraphModule) -> None:
"""Drop reshape ops whose output shape equals their input shape.
Companion to :func:`fold_consecutive_reshapes`. ``make_fx`` records a
reshape to the shape the input already has; the compiled graph has already
dropped it, so the pattern only matches once we drop it too.
"""
aten_reshape = torch.ops.aten.reshape.default
for node in list(gm.graph.nodes):
if not is_func(node, aten_reshape):
continue
inp = node.args[0]
if not isinstance(inp, fx.Node):
continue
val, inp_val = node.meta.get("val"), inp.meta.get("val")
if val is None or inp_val is None or val.shape != inp_val.shape:
continue
node.replace_all_uses_with(inp)
gm.graph.erase_node(node)
def _remove_noop_permutes(gm: fx.GraphModule) -> None:
for node in gm.graph.nodes:
if not is_func(node, torch.ops.aten.permute.default):
vllm/config/utils.py
-58
View File
@@ -430,31 +430,6 @@ class Range:
return self.__str__()
def handle_deprecated(
config: ConfigT,
old_name: str,
new_name_or_names: str | list[str],
removal_version: str,
) -> None:
old_val = getattr(config, old_name)
if old_val is None:
return
if isinstance(new_name_or_names, str):
new_names = [new_name_or_names]
else:
new_names = new_name_or_names
msg = (
f"{old_name} is deprecated and will be removed in {removal_version}. "
f"Use {', '.join(new_names)} instead."
)
logger.warning(msg)
for new_name in new_names:
setattr(config, new_name, old_val)
def get_from_deprecated_env_if_set(
env_name: str,
removal_version: str,
@@ -483,36 +458,3 @@ def get_from_deprecated_env_if_set(
return value
return None
def set_from_deprecated_env_if_set(
config: ConfigT,
env_name: str,
removal_version: str,
field_name: str,
to_bool: bool = False,
to_int: bool = False,
) -> None:
"""Set object field from deprecated environment variable with warning.
Args:
config: Config object to set the field on
env_name: Name of the deprecated environment variable
removal_version: Version when the env var will be removed
field_name: Name of the field to set
to_bool: Whether to convert the environment variable value to boolean
to_int: Whether to convert the environment variable value to integer
Returns:
None
"""
if to_bool and to_int:
raise ValueError("Cannot convert to both boolean and integer.")
env_value = get_from_deprecated_env_if_set(env_name, removal_version, field_name)
if env_value is not None:
field_value: str | bool | int = env_value
if to_bool:
field_value = env_value.lower() in ("1", "true")
elif to_int:
field_value = int(env_value)
setattr(config, field_name, field_value)
vllm/device_allocator/alloc_conf.py
-22
View File
@@ -88,28 +88,6 @@ def with_conf_flag(conf: str, key: str, enabled: bool) -> str:
return ",".join(fields)
def expandable_segments_enabled_from_env(
environ: "os._Environ[str] | dict[str, str] | None" = None,
) -> bool:
return conf_flag_enabled(alloc_conf_from_env(environ), EXPANDABLE_SEGMENTS)
def expandable_segments_enabled() -> bool | None:
"""Live ``expandable_segments`` state, or ``None`` if it cannot be read.
``None`` is deliberately distinct from ``False``: the environment can never
reflect a runtime write, so a caller verifying a write it just issued must
not read "cannot tell" as "the write did not take".
"""
try:
settings = torch.cuda.memory._snapshot().get("allocator_settings")
if settings is not None and EXPANDABLE_SEGMENTS in settings:
return bool(settings[EXPANDABLE_SEGMENTS])
except Exception: # noqa: BLE001
pass
return None
def set_alloc_conf(conf: str) -> None:
"""Write a complete allocator config string."""
setter = getattr(torch._C, "_accelerator_setAllocatorSettings", None)
vllm/distributed/ec_transfer/ec_connector/encoder_output_width.py
-17
View File
@@ -83,20 +83,3 @@ def measure_encoder_output_widths(
logger.info("EC producer encoder output widths: %s", reports[0])
ec_config.mm_encoder_output_widths = reports[0]
def get_encoder_output_width(vllm_config: "VllmConfig", modality: str) -> int:
"""Encoder output width for `modality`, as measured at producer startup.
Raises:
ValueError: No width was measured for `modality`, so a transfer
cannot be sized.
"""
assert vllm_config.ec_transfer_config is not None
widths = vllm_config.ec_transfer_config.mm_encoder_output_widths
if modality not in widths:
raise ValueError(
f"No encoder output width was measured for modality {modality!r}; "
f"measured: {sorted(widths)}"
)
return widths[modality]
vllm/distributed/ec_transfer/ec_connector/utils.py
+1 -42
View File
@@ -2,17 +2,14 @@
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""EC connector helper utilities."""
from typing import TYPE_CHECKING, Any
from typing import TYPE_CHECKING
import torch
from vllm.logger import init_logger
from vllm.utils.collection_utils import is_list_of
from vllm.v1.outputs import ECConnectorOutput, ModelRunnerOutput
if TYPE_CHECKING:
from vllm.config import ModelConfig
from vllm.multimodal.inputs import MultiModalFeatureSpec
logger = init_logger(__name__)
@@ -52,44 +49,6 @@ class PlaceholderMetadataResolver:
return fields
def collect_ec_item_metadata(
mm_features: "list[MultiModalFeatureSpec]",
resolver: PlaceholderMetadataResolver,
) -> dict[str, dict[str, Any]]:
"""Build one `ec_transfer_params` entry per feature for `request_finished()`.
Keyed by mm_hash, each entry carries a `metadata` dict with whatever
placeholder fields `resolver` says this model needs published for its
modality, so a consumer can skip media preprocessing. `item_indices`
identifies every occurrence in `mm_features`, including repeated hashes.
Audio token counts come from the placeholder; other metadata requires
`data` and falls back to raw media when it is unavailable on cache hits.
A connector that also has transfer coordinates to report
(e.g. NIXL peer_host/peer_port/size_bytes) merges those in alongside
`metadata`, not into it.
"""
items: dict[str, dict[str, Any]] = {}
for index, feature in enumerate(mm_features):
metadata: dict[str, Any] = {}
wanted = resolver.fields_for(feature.modality)
if "audio_num_tokens" in wanted:
metadata["audio_num_tokens"] = [feature.mm_position.get_num_embeds()]
if feature.data is not None:
for key, value in feature.data.get_data().items():
if key not in wanted:
continue
if isinstance(value, torch.Tensor):
metadata[key] = torch.atleast_1d(value).tolist()
elif is_list_of(value, (int, float)):
# Some metadata (e.g. Qwen3-VL video timestamps) is
# produced as a plain list rather than a tensor.
metadata[key] = value
entry = items.setdefault(feature.identifier, {"item_indices": []})
entry["metadata"] = metadata
entry["item_indices"].append(index)
return items
class ECOutputAggregator:
"""Merge every worker's EC connector output onto the single
ModelRunnerOutput that reaches the scheduler.
vllm/distributed/kv_transfer/kv_connector/utils.py
-98
View File
@@ -20,7 +20,6 @@ from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
from vllm.platforms import current_platform
from vllm.utils.torch_utils import async_tensor_h2d
from vllm.v1.attention.backend import AttentionBackend
from vllm.v1.kv_cache_interface import AttentionSpec, KVCacheConfig
from vllm.v1.outputs import KVConnectorOutput, ModelRunnerOutput
if TYPE_CHECKING:
@@ -34,16 +33,6 @@ EngineId = str
BlockIds = tuple[list[int], ...] | list[list[int]]
def clip_ssm_state_blocks(blocks: list[int], num_spec_blocks: int) -> list[int]:
"""Drop speculative scratch slots and keep only the running-state slot."""
if not blocks:
return blocks
# Keep at least the running-state slot, even before scratch allocation.
if num_scratch := min(num_spec_blocks, len(blocks) - 1):
blocks = blocks[:-num_scratch]
return blocks[-1:]
def get_kv_connector_cache_layout(vllm_config: VllmConfig | None = None):
# NOTE (NickLucche) When running disaggregated PD with NIXL, LBHNC layout is
# used for faster transfer.
@@ -241,57 +230,6 @@ def copy_kv_blocks(
copy_fn(src_tensor, dst_tensor, src_indices, dst_indices)
def kv_postprocess_blksize_on_receive(cache, indices, block_size_ratio):
"""Transforms the layout of received KV cache blocks to the local block_size.
(Only works for local blocksize > remote blocksize)
Example:
local blocksize = 16 tokens, remote blocksize = 4 tokens
local block[0] = remote block[0, 1, 2, 3]
remote is |h0-b0|h1-b0|h2-b0|h3-b0|h0-b1|h1-b1|h2-b1|h3-b1|...
local is |h0-b0..................|h1-b0..................|...
permute is to:
1. view => view remote as n_blocks * remote_shape(H,remoteN,D)
2. permute => (H, nblocks, remoteN, D)
3. flatten => (H, localN, D)
"""
blocks_to_update = cache.index_select(0, indices)
n_kv_heads, block_size, head_size = blocks_to_update.shape[1:]
remote_block_size = block_size // block_size_ratio
n_blocks = block_size_ratio
permuted_blocks = (
blocks_to_update.reshape(-1, n_blocks, n_kv_heads, remote_block_size, head_size)
.permute(0, 2, 1, 3, 4)
.flatten(2, 3)
)
cache.index_copy_(0, indices, permuted_blocks)
def kv_postprocess_blksize_and_layout_on_receive(cache, indices, block_size_ratio):
"""Transforms the layout of received KV cache to the local block_size and LBHNC.
(Only works for local blocksize >= remote blocksize)
prefill is LBHNC, smaller block_size
decode(local) is LBNHC, larger block_size
"""
blocks_to_update = cache.index_select(0, indices)
n_kv_heads, block_size, head_size = blocks_to_update.shape[1:]
remote_block_size = block_size // block_size_ratio
n_blocks = block_size_ratio
# View the received bytes in memory order as the remote head-major blocks.
permuted_blocks = (
blocks_to_update.transpose(1, 2)
.reshape(-1, n_blocks, n_kv_heads, remote_block_size, head_size)
.permute(0, 2, 1, 3, 4)
.flatten(2, 3)
)
cache.index_copy_(0, indices, permuted_blocks)
def yield_req_data(
scheduler_output,
) -> Iterator[tuple[str, tuple[list[int], ...] | None, bool]]:
@@ -352,42 +290,6 @@ def get_current_attn_backends(
]
def get_current_attn_backends_and_specs(
vllm_config: VllmConfig,
kv_cache_config: KVCacheConfig,
fallback: list[type[AttentionBackend]],
) -> tuple[list[type[AttentionBackend]], list[AttentionSpec | None]]:
"""Distinct (backend, spec) pairs of the transfer layers, else ``fallback``.
Compressed caches are skipped: their kernel pages never split transferred
blocks.
"""
pairs: dict[tuple[type[AttentionBackend], AttentionSpec | None], None] = {}
layer_type = cast(type[Any], AttentionLayerBase)
for group in kv_cache_config.transfer_groups:
specs = getattr(group.kv_cache_spec, "kv_cache_specs", {})
layers = get_layers_from_vllm_config(vllm_config, layer_type, group.layer_names)
for name, layer in layers.items():
spec = specs.get(name, group.kv_cache_spec)
if isinstance(spec, AttentionSpec) and spec.tokens_per_state > 1:
layout = vllm_config.cache_config.get_resolved_kv_cache_layout()
assert layout.is_block_outermost, "Compressed caches must be packed"
continue
if not isinstance(spec, AttentionSpec):
spec = None
pairs[layer.get_attn_backend(), spec] = None
if not pairs:
return fallback, [None] * len(fallback)
return [backend for backend, _ in pairs], [spec for _, spec in pairs]
def get_current_attn_backend(
vllm_config: VllmConfig, layer_names: list[str] | None = None
) -> type[AttentionBackend]:
"""Get the first attention backend for the given layers."""
return get_current_attn_backends(vllm_config, layer_names)[0]
# ---- Per-engine transfer info ----
vllm/distributed/kv_transfer/kv_connector/v1/offloading/canonical_mapping.py
-13
View File
@@ -17,7 +17,6 @@ import torch
from vllm.v1.kv_cache_interface import (
AttentionSpec,
KVCacheConfig,
KVCacheLayout,
KVCacheSpec,
MambaSpec,
MLAAttentionSpec,
@@ -32,18 +31,6 @@ if TYPE_CHECKING:
CANONICAL_FORMAT_VERSION = 1
def canonical_format_id(kv_cache_layout: str) -> str:
"""Identity of the canonical byte format, for namespacing persisted KV.
Canonical pages keep the worker's KV layout family, so the id couples the
format version with that family; consumers must match it exactly. The family
keeps its historical NHD/HND spelling so ids stay stable for KV persisted
before the layout enum existed."""
layout = KVCacheLayout[kv_cache_layout]
legacy = {KVCacheLayout.LBNHC: "nhd", KVCacheLayout.LBHNC: "hnd"}
family = legacy.get(layout, layout.name.lower())
return f"v{CANONICAL_FORMAT_VERSION}-{family}"
@dataclass(frozen=True)
class _RankContext:
"""Sharding parameters of one worker rank within the offload group."""
vllm/distributed/parallel_state.py
-14
View File
@@ -1607,20 +1607,6 @@ def get_etp_group() -> GroupCoordinator:
_ENGRAM_DP: GroupCoordinator | None = None
def get_engram_dp_group() -> GroupCoordinator | None:
"""Return the DP replicas that share one engram embedding table.
None when every replica holds a full (TP-sharded) copy, which is the case
for models without engram layers and for replicas that span nodes.
"""
return _ENGRAM_DP
def get_engram_dp_size() -> int:
"""Number of DP replicas one engram embedding table is sharded over."""
return _ENGRAM_DP.world_size if _ENGRAM_DP is not None else 1
_DCP: GroupCoordinator | None = None
vllm/entrypoints/launchers/cli_args.py
-7
View File
@@ -403,13 +403,6 @@ def validate_parsed_serve_args(args: argparse.Namespace):
validate_multi_port_external_lb_args(args)
def create_parser_for_docs() -> FlexibleArgumentParser:
parser_for_docs = FlexibleArgumentParser(
prog="-m vllm.entrypoints.launchers.api_server.entry"
)
return make_arg_parser(parser_for_docs)
def propagate_flash_late_interaction(args, engine_args) -> None:
"""Propagate `--disable-flash-late-interaction` into the engine config.
vllm/entrypoints/serve/engine/typing.py
+1 -11
View File
@@ -20,8 +20,7 @@ from vllm.entrypoints.serve.tokenize.protocol import (
TokenizeCompletionRequest,
TokenizeResponse,
)
from vllm.renderers import ChatParams, TokenizeParams
from vllm.renderers.chat_utils import ChatTemplateContentFormatOption
from vllm.renderers import TokenizeParams
class RendererRequest(Protocol):
@@ -29,15 +28,6 @@ class RendererRequest(Protocol):
raise NotImplementedError
class RendererChatRequest(RendererRequest, Protocol):
def build_chat_params(
self,
default_template: str | None,
default_template_content_format: ChatTemplateContentFormatOption,
) -> ChatParams:
raise NotImplementedError
CompletionLikeRequest: TypeAlias = (
CompletionRequest
| TokenizeCompletionRequest
vllm/entrypoints/serve/utils/api_utils.py
-13
View File
@@ -324,19 +324,6 @@ def _jsonify_arg_value(value: Any) -> Any:
return repr(value)
def jsonify_non_default_args(
args: Namespace | EngineArgs,
*,
exclude: set[str] | None = None,
) -> dict[str, Any]:
non_default_args = get_non_default_args(args)
if exclude is not None:
for key in exclude:
non_default_args.pop(key, None)
return {key: _jsonify_arg_value(value) for key, value in non_default_args.items()}
# Fields whose values must never be logged verbatim.
_SENSITIVE_ARG_FIELDS = frozenset({"api_key", "hf_token", "watermark_config"})
vllm/model_executor/kernels/attention/__init__.py
-2
View File
@@ -1,2 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
vllm/model_executor/kernels/linear/cute_dsl/_ll_fp32w_dotprod.py
-5
View File
@@ -450,8 +450,3 @@ class LLFp32WDotprod:
if const_expr(self.use_pdl):
cute.arch.griddepcontrol_launch_dependents()
def make_host_fp32w(
k_val: int, bs: int = 128, token_groups: int = 1, epb: int = 1, m: int = 1
):
return LLFp32WDotprod(m=m, k=k_val, bs=bs, token_groups=token_groups, epb=epb)
vllm/model_executor/layers/fused_moe/config.py
-18
View File
@@ -828,24 +828,6 @@ def nvfp4_moe_quant_config(
)
def mxfp4_moe_quant_config(
w1_scale: torch.Tensor,
w2_scale: torch.Tensor,
) -> FusedMoEQuantConfig:
"""Construct a quant config for MXFP4 x MXFP4 MoE.
MXFP4 uses block scaling only (E8M0 scales, 32-element groups), with no
separate alphas / global activation scales in this config.
"""
return FusedMoEQuantConfig.make(
"mxfp4",
w1_scale=w1_scale,
w2_scale=w2_scale,
per_act_token_quant=False,
per_out_ch_quant=False,
block_shape=None,
)
def nvfp4_w4a16_moe_quant_config(
g1_alphas: torch.Tensor,
g2_alphas: torch.Tensor,
vllm/model_executor/layers/fused_moe/experts/flashinfer_cutlass_moe.py
-26
View File
@@ -35,32 +35,6 @@ from vllm.utils.flashinfer import (
logger = init_logger(__name__)
def is_valid_flashinfer_cutlass_fused_moe(
hidden_states: torch.Tensor, w1: torch.Tensor, w2: torch.Tensor
) -> bool:
"""Check if the given problem size is supported by the FlashInfer CUTLASS MoE
kernel.
"""
if not has_flashinfer_cutlass_fused_moe():
logger.debug_once(
"FlashInferExperts disabled: flashinfer_cutlass_fused_moe not available."
)
return False
# Data type checks
if (
w1.dtype != torch.uint8
or w2.dtype != torch.uint8
or hidden_states.dtype not in [torch.float32, torch.float16, torch.bfloat16]
):
logger.debug_once(
"FlashInferExperts disabled: w1/w2 must be torch.uint8 "
f"(got w1={w1.dtype}, w2={w2.dtype}), hidden_states must be "
f"float32, float16, or bfloat16 (got {hidden_states.dtype})."
)
return False
return True
class FlashInferExperts(mk.FusedMoEExpertsModular):
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
if self.quant_config.use_nvfp4_w4a4:
vllm/model_executor/layers/fused_moe/flashinfer_moe_ep.py
-5
View File
@@ -1,7 +1,6 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from __future__ import annotations
import contextlib
import sys
@@ -103,10 +102,6 @@ _BACKEND_SPECS = {
assert frozenset(_BACKEND_SPECS) == FLASHINFER_MOE_EP_BACKENDS
def is_flashinfer_moe_ep_backend(moe_backend: str) -> bool:
return moe_backend in _BACKEND_SPECS
def flashinfer_moe_ep_backend_spec(
moe_backend: str,
) -> FlashInferMoeEpBackendSpec:
vllm/model_executor/layers/linear.py
+1 -90
View File
@@ -2,7 +2,7 @@
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from abc import abstractmethod
from collections.abc import Callable, Iterable
from collections.abc import Iterable
from typing import Any
import torch
@@ -1543,95 +1543,6 @@ class MinimaxM3QKVParallelLinearWithIndexer(QKVParallelLinear):
param_data.copy_(loaded_weight)
class KimiK3MergedQKVGateLinear(MergedColumnParallelLinear):
"""Kimi-K3 QKV-A projection fused with its output-gate projection.
NOTE: Kimi-K3-specific. This is tailored to its replicated Q/KV latent
projections and tensor-parallel output gate; it is not a general-purpose
linear layer. It lives here alongside ``MergedColumnParallelLinear``, whose
sharding and weight-loading machinery it reuses.
The per-rank shard layout is ``[q_a | kv_a | gate]``. The latent shards
stay replicated across tensor-parallel ranks while the gate shard is
tensor-parallel, matching the standalone projections they replace.
"""
def __init__(
self,
hidden_size: int,
q_lora_rank: int,
kv_lora_rank: int,
qk_rope_head_dim: int,
total_num_heads: int,
v_head_dim: int,
bias: bool = False,
quant_config: QuantizationConfig | None = None,
prefix: str = "",
) -> None:
tp_size = get_tensor_model_parallel_world_size()
# MergedColumnParallelLinear divides every logical shard by TP size.
# Inflate the replicated latent shards to retain their checkpoint widths.
output_sizes = [
q_lora_rank * tp_size,
(kv_lora_rank + qk_rope_head_dim) * tp_size,
total_num_heads * v_head_dim,
]
super().__init__(
input_size=hidden_size,
output_sizes=output_sizes,
bias=bias,
quant_config=quant_config,
prefix=prefix,
)
def _load_with_replicated_latents(
self,
loader: Callable[..., None],
param: Parameter,
loaded_weight: torch.Tensor,
loaded_shard_id: tuple[int, ...] | int | None,
) -> None:
is_replicated = isinstance(loaded_shard_id, int) and loaded_shard_id < 2
if not is_replicated:
loader(param, loaded_weight, loaded_shard_id)
return
# The base loader narrows every checkpoint shard using its TP rank.
# Replicated latent shards instead load the rank-zero slice on every rank.
tp_rank = self.tp_rank
param_tp_rank = getattr(param, "tp_rank", None)
self.tp_rank = 0
if param_tp_rank is not None:
param.tp_rank = 0
try:
loader(param, loaded_weight, loaded_shard_id)
finally:
self.tp_rank = tp_rank
if param_tp_rank is not None:
param.tp_rank = param_tp_rank
def weight_loader(
self,
param: Parameter,
loaded_weight: torch.Tensor,
loaded_shard_id: tuple[int, ...] | int | None = None,
) -> None:
self._load_with_replicated_latents(
super().weight_loader, param, loaded_weight, loaded_shard_id
)
def weight_loader_v2(
self,
param: BasevLLMParameter,
loaded_weight: torch.Tensor,
loaded_shard_id: tuple[int, ...] | int | None = None,
) -> None:
self._load_with_replicated_latents(
super().weight_loader_v2, param, loaded_weight, loaded_shard_id
)
# --8<-- [start:row_parallel_linear]
@PluggableLayer.register("row_parallel_linear")
class RowParallelLinear(LinearBase):
vllm/model_executor/layers/mamba/ops/cpu/causal_conv1d.py
-31
View File
@@ -1,7 +1,6 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from __future__ import annotations
import torch
import torch.nn.functional as F
@@ -115,33 +114,3 @@ def causal_conv1d_update_cpu(
pad_slot_id,
)
def causal_conv1d_update_torch(
x: torch.Tensor,
conv_state: torch.Tensor,
weight: torch.Tensor,
bias: torch.Tensor | None = None,
activation: str | None = None,
) -> torch.Tensor:
"""Pure PyTorch fallback for causal_conv1d_update.
Currently used as a fallback for Arm (aarch64) to leverage
oneDNN/ACL F.conv1d kernels for batched decoding.
"""
assert activation in {None, "silu", "swish"}
_, dim, seq_len = x.shape
state_len = conv_state.shape[-1]
x_new = torch.cat([conv_state, x], dim=-1).to(weight.dtype)
conv_state.copy_(x_new[:, :, -state_len:])
out = F.conv1d(
x_new,
weight.unsqueeze(1),
bias,
padding=0,
groups=dim,
)[:, :, -seq_len:]
if activation in ("silu", "swish"):
out = F.silu(out)
return out
vllm/model_executor/layers/mamba/ops/mamba_ssm.py
-151
View File
@@ -22,9 +22,6 @@ from vllm.triton_utils import HAS_TRITON, tl, triton
from vllm.utils.platform_utils import get_device_name_as_file_name
from vllm.v1.attention.backends.utils import NULL_BLOCK_ID
if current_platform.is_xpu():
from vllm._xpu_ops import xpu_ops
logger = init_logger(__name__)
TRITON3 = HAS_TRITON and (version.parse(triton.__version__) >= version.parse("3.0.0"))
@@ -699,154 +696,6 @@ def selective_state_update(
)
def selective_scan_fn(
u,
ssm_states,
delta,
A,
B,
C,
D=None,
z=None,
delta_bias=None,
delta_softplus=False,
query_start_loc=None,
cache_indices=None,
has_initial_state=None,
null_block_id=NULL_BLOCK_ID,
block_size=1024,
block_idx_first_scheduled_token=None,
block_idx_last_scheduled_token=None,
initial_state_idx=None,
cu_chunk_seqlen=None,
last_chunk_indices=None,
) -> torch.Tensor:
"""u: (dim, total_length) for varlen or (batch, dim, seqlen)
applies changes in place.
ssm_states: (batch, dim, dstate) or (batch, nheads, dim, dstate)
applies changes in place.
delta: (dim, total_length) for varlen or (batch, dim, seqlen)
A: (dim, dstate)
B: (ngroups, dstate, total_length) for varlen or
(batch,ngroups,dstate,seqlen)
C: (ngroups, dstate, total_length) for varlen or
(batch,ngroups,dstate,seqlen)
D: (dim,)
z: (dim, total_length) for varlen or (batch, dim, seqlen)
dt_bias: (dim,) or (dim)
query_start_loc: (batch + 1) int32
The cumulative sequence lengths of the sequences in
the batch, used to index into sequence. prepended with 0.
for example: query_start_loc = torch.Tensor([0,10,16,17]),
x.shape=(dim,17)
cache_indices: (batch) int32
A tensor with each cell is a correspondent
input and output ssm_state indices
- Without APC: (batch,) - single state index per batch item
- With APC: (batch, max_positions) - cache block indices for read/write
Each non-zero value indicates a cache block to load from and/or write to.
has_initial_state: (batch) bool
A tensor populated with ones and zeros,
indicate if the ssm_state at the corresponding index should be
used as initial state. Not providing argument assumes
there's no initial state
null_block_id: int
if cache_indices is passed, lets the kernel identify padding entries
that will not be processed,
for example: cache_indices = [null_block_id, 1 ,20 ,null_block_id]
in this case, the kernel will not process entries at indices 0 and 3
block_size: int
The block size to align the cached states to
block_idx_first_scheduled_token: (batch,), dtype int32
The pointer into cache_indices, where the first
cache block to be filled is located.
block_idx_last_scheduled_token: (batch,), dtype int32
The pointer into cache_indices, where the last cache block
to be filled is located.
initial_state_idx: (batch,), dtype int32
The pointer into cache_indices, where the cache block
containing the initial state is located.
Returns
-------
output: (dim, total_length) for varlen or (batch, dim, seqlen)
supports inplace replacement
"""
if u.stride(-1) != 1:
u = u.contiguous()
if delta.stride(-1) != 1:
delta = delta.contiguous()
if D is not None:
D = D.contiguous()
if B.stride(-1) != 1:
B = B.contiguous()
if C.stride(-1) != 1:
C = C.contiguous()
if z is not None and z.stride(-1) != 1:
z = z.contiguous()
if B.dim() == 3 and query_start_loc is None:
B = B.unsqueeze(1)
if B.dim() == 2 and query_start_loc is not None:
B = B.unsqueeze(0)
if C.dim() == 3 and query_start_loc is None:
C = C.unsqueeze(1)
if C.dim() == 2 and query_start_loc is not None:
C = C.unsqueeze(0)
if current_platform.is_xpu():
xpu_ops.selective_scan_fwd(
u,
delta,
A,
B,
C,
D,
z,
delta_bias,
delta_softplus,
query_start_loc,
cache_indices,
has_initial_state,
ssm_states,
null_block_id,
block_size,
block_idx_first_scheduled_token,
block_idx_last_scheduled_token,
initial_state_idx,
cu_chunk_seqlen,
last_chunk_indices,
)
else:
ops.selective_scan_fwd(
u,
delta,
A,
B,
C,
D,
z,
delta_bias,
delta_softplus,
query_start_loc,
cache_indices,
has_initial_state,
ssm_states,
null_block_id,
block_size,
block_idx_first_scheduled_token,
block_idx_last_scheduled_token,
initial_state_idx,
cu_chunk_seqlen,
last_chunk_indices,
)
if z is None:
return delta # output written inplace to delta
else:
return z # output written inplace to z
from vllm.platforms import current_platform # noqa: E402
if current_platform.is_cpu():
vllm/model_executor/layers/quantization/utils/fp8_utils.py
-19
View File
@@ -50,25 +50,6 @@ def is_fp8(x: torch.dtype | torch.Tensor) -> bool:
return x == torch.float8_e4m3fn or x == torch.float8_e4m3fnuz
def get_fp8_block_weight_scale(layer: torch.nn.Module) -> torch.Tensor | None:
"""Return the block-FP8 weight scale for supported quant methods."""
return getattr(layer, "weight_scale_inv", None)
def input_to_float8(
x: torch.Tensor, dtype: torch.dtype | None = None
) -> tuple[torch.Tensor, torch.Tensor]:
"""This function quantizes input values to float8 values "
"with tensor-wise quantization."""
dtype = current_platform.fp8_dtype() if dtype is None else dtype
finfo = torch.finfo(dtype)
min_val, max_val = x.aminmax()
amax = torch.maximum(min_val.abs(), max_val.abs()).clamp(min=1e-12)
scale = finfo.max / amax
x_scl_sat = (x * scale).clamp(min=finfo.min, max=finfo.max)
return x_scl_sat.to(dtype).contiguous(), scale.float().reciprocal()
@triton.jit
def _per_token_group_quant_fp8(
# Pointers to inputs and output
vllm/model_executor/layers/quantization/utils/marlin_utils_fp4.py
-110
View File
@@ -656,113 +656,3 @@ def prepare_moe_mxfp4_layer_for_marlin(
return w13, w2, w13_scale, w2_scale, w13_bias, w2_bias
def rand_marlin_weight_nvfp4_like(weight, group_size, input_dtype=None):
is_a_8bit = input_dtype is not None and input_dtype.itemsize == 1
assert not is_a_8bit, "NVFP4 weight + INT8/FP8 activation is not supported."
assert group_size > 0
size_n, size_k = weight.shape
scales = weight.view(size_n, -1, group_size).abs().max(-1)[0] / 6
global_scale = scales.max() / 448
scales = (scales / global_scale).to(torch.float8_e4m3fn)
fp4_weight = torch.randint(
0, 256, (size_n, size_k // 2), dtype=torch.uint8, device=weight.device
)
fp4_weight_part_1 = (fp4_weight & 0b10000000) | ((fp4_weight & 0b01110000) >> 2)
fp4_weight_part_1 = fp4_weight_part_1.view(torch.float8_e4m3fn)
fp4_weight_part_1 = fp4_weight_part_1.to(weight.dtype) * (2**6)
fp4_weight2 = fp4_weight << 4
fp4_weight_part_2 = (fp4_weight2 & 0b10000000) | ((fp4_weight2 & 0b01110000) >> 2)
fp4_weight_part_2 = fp4_weight_part_2.view(torch.float8_e4m3fn)
fp4_weight_part_2 = fp4_weight_part_2.to(weight.dtype) * (2**6)
weight_ref = torch.cat(
[fp4_weight_part_2.unsqueeze(2), fp4_weight_part_1.unsqueeze(2)], 2
).view(size_n, size_k)
weight_ref = (
weight_ref
* global_scale.to(weight.dtype)
* scales.repeat_interleave(group_size, 1).to(weight.dtype)
)
marlin_qweight = ops.gptq_marlin_repack(
b_q_weight=fp4_weight.view(torch.int32).T.contiguous(),
size_k=size_k,
size_n=size_n,
num_bits=4,
is_a_8bit=is_a_8bit,
)
marlin_scales = marlin_permute_scales(
s=scales.T.to(weight.dtype),
size_k=size_k,
size_n=size_n,
group_size=group_size,
is_a_8bit=is_a_8bit,
)
marlin_scales, scale_factor = nvfp4_marlin_process_scales(marlin_scales)
global_scale = nvfp4_marlin_process_global_scale(global_scale).to(torch.float32)
global_scale = global_scale / scale_factor
return weight_ref.T, marlin_qweight, marlin_scales, global_scale
def rand_marlin_weight_mxfp4_like(weight, group_size, input_dtype=None):
is_a_8bit = input_dtype is not None and input_dtype.itemsize == 1
if is_a_8bit:
assert input_dtype == torch.float8_e4m3fn, (
"MXFP4 weight + INT8 activation is not supported."
)
assert group_size > 0
size_n, size_k = weight.shape
scales = torch.randint(
110,
120,
(size_n, size_k // group_size),
dtype=torch.uint8,
device=weight.device,
)
scales = scales.view(torch.float8_e8m0fnu)
fp4_weight = torch.randint(
0, 256, (size_n, size_k // 2), dtype=torch.uint8, device=weight.device
)
fp4_weight_part_1 = (fp4_weight & 0b10000000) | ((fp4_weight & 0b01110000) >> 2)
fp4_weight_part_1 = fp4_weight_part_1.view(torch.float8_e4m3fn)
fp4_weight_part_1 = fp4_weight_part_1.to(weight.dtype) * (2**6)
fp4_weight2 = fp4_weight << 4
fp4_weight_part_2 = (fp4_weight2 & 0b10000000) | ((fp4_weight2 & 0b01110000) >> 2)
fp4_weight_part_2 = fp4_weight_part_2.view(torch.float8_e4m3fn)
fp4_weight_part_2 = fp4_weight_part_2.to(weight.dtype) * (2**6)
weight_ref = torch.cat(
[fp4_weight_part_2.unsqueeze(2), fp4_weight_part_1.unsqueeze(2)], 2
).view(size_n, size_k)
weight_ref = weight_ref * scales.repeat_interleave(group_size, 1).to(weight.dtype)
fp4_weight = fp4_weight.view(torch.int32).T.contiguous()
marlin_qweight = ops.gptq_marlin_repack(
b_q_weight=fp4_weight,
size_k=size_k,
size_n=size_n,
num_bits=4,
is_a_8bit=is_a_8bit,
)
marlin_scales = marlin_permute_scales(
s=scales.T.to(weight.dtype),
size_k=size_k,
size_n=size_n,
group_size=group_size,
is_a_8bit=is_a_8bit,
)
marlin_scales = mxfp4_marlin_process_scales(marlin_scales, input_dtype=input_dtype)
return weight_ref.T, marlin_qweight, marlin_scales.to(torch.float8_e8m0fnu)
vllm/model_executor/layers/quantization/utils/marlin_utils_fp8.py
-38
View File
@@ -561,41 +561,3 @@ def prepare_mxfp8_moe_layer_for_marlin(
return w13, w2, w13_scale, w2_scale
def marlin_quant_fp8_torch(weight, group_size, input_dtype=None):
is_a_8bit = input_dtype is not None and input_dtype.itemsize == 1
if is_a_8bit:
assert input_dtype == torch.float8_e4m3fn
size_n, size_k = weight.shape
if group_size != -1:
scales = weight.view(size_n, -1, group_size).abs().max(-1)[0] / 448
repeated_scales = scales.repeat_interleave(group_size, 1)
fp8_weight = (weight / repeated_scales).to(torch.float8_e4m3fn)
weight_ref = fp8_weight.to(weight.dtype) * repeated_scales
else:
scales = weight.view(size_n, 1, group_size).abs().max(-1)[0] / 448
repeated_scales = scales.repeat_interleave(size_k, 1)
fp8_weight = (weight / repeated_scales).to(torch.float8_e4m3fn)
weight_ref = fp8_weight.to(weight.dtype) * repeated_scales
packed_weight = pack_fp8_to_int32(fp8_weight, False).T.contiguous()
marlin_qweight = ops.gptq_marlin_repack(
b_q_weight=packed_weight,
size_k=size_k,
size_n=size_n,
num_bits=8,
is_a_8bit=is_a_8bit,
)
marlin_scales = marlin_permute_scales(
s=scales.T,
size_k=size_k,
size_n=size_n,
group_size=group_size,
is_a_8bit=is_a_8bit,
)
marlin_scales = fp8_fused_exponent_bias_into_scales(marlin_scales)
return weight_ref.T, marlin_qweight, marlin_scales
vllm/model_executor/layers/quantization/utils/ocp_mx_utils.py
-51
View File
@@ -1,7 +1,6 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from enum import Enum
from typing import Any
from vllm.logger import init_logger
from vllm.model_executor.layers.quantization.utils.quant_utils import (
@@ -47,56 +46,6 @@ OCP_MX_DTYPES = {
}
def ocp_mx_weight_dtype_and_rows(
weight_quant: dict[str, Any] | None,
) -> tuple[str, int] | None:
"""Normalize a Quark weight quant config to ``(mx_dtype, scale_block_rows)``.
Quark spells OCP MX weights two ways. The canonical one is 1-D per-group:
``qscheme="per_group"``, ``group_size=32``, ``scale_format="e8m0"``. MXFP8
checkpoints may instead use a 2-D per-block spelling
(``qscheme="per_block"``, ``block_size=[R, 32]``,
``scale_type="float8_e8m0fnu"``), which carries one scale per ``R`` weight
rows rather than one per row. ``scale_block_rows`` is that ``R``; it is 1
for the canonical spelling, where the two layouts coincide.
Returns ``None`` when the config is not an OCP MX weight quantization.
"""
if not isinstance(weight_quant, dict):
return None
dtype = weight_quant.get("dtype")
if not isinstance(dtype, str):
return None
mx_dtype = dtype.replace("fp", "mxfp")
if mx_dtype not in _WEIGHT_QUANT_KEY_MAP:
return None
qscheme = weight_quant.get("qscheme")
if qscheme == "per_group":
if weight_quant.get("group_size") != OCP_MX_BLOCK_SIZE:
return None
if weight_quant.get("scale_format") != "e8m0":
return None
return mx_dtype, 1
if qscheme == "per_block":
# Only the unpacked dtypes are accepted here: expanding a 2-D scale
# over a sub-byte packed weight has no checkpoint to validate against.
if mx_dtype not in OCP_MX_UNPACKED_DTYPES:
return None
block_size = list(weight_quant.get("block_size") or [])
if len(block_size) != 2 or block_size[1] != OCP_MX_BLOCK_SIZE:
return None
if weight_quant.get("scale_type") != "float8_e8m0fnu":
return None
if weight_quant.get("symmetric") is not True or weight_quant.get("is_dynamic"):
return None
return mx_dtype, block_size[0]
return None
class OCP_MX_Scheme(str, Enum):
w_mxfp4 = "w_mxfp4"
w_mxfp4_a_mxfp4 = "w_mxfp4_a_mxfp4"
vllm/model_executor/layers/quantization/utils/quant_utils.py
+1 -111
View File
@@ -5,7 +5,7 @@
from collections.abc import Callable, Mapping
from dataclasses import dataclass
from types import MappingProxyType
from typing import TYPE_CHECKING, ClassVar, Literal, NamedTuple
from typing import ClassVar, Literal, NamedTuple
import numpy
import torch
@@ -15,9 +15,6 @@ from vllm.distributed.parallel_state import get_ep_group, get_tp_group
from vllm.platforms import current_platform
from vllm.scalar_type import ScalarType, scalar_types
if TYPE_CHECKING:
from vllm.model_executor.layers.linear import LinearBase
FP8_DTYPE = current_platform.fp8_dtype()
FP4_DTYPE = torch.uint8
MXFP_SCALE_DTYPE = torch.uint8
@@ -508,63 +505,6 @@ def get_attribute_fallback(obj, attributes: list[str]):
raise AttributeError(f"'{obj}' has no recognized attributes: {attributes}.")
def get_and_maybe_dequant_weights(
layer: "LinearBase", out_dtype: torch.dtype = torch.float32
):
"""Return layer's unquantized weights in [out, in] layout."""
from vllm.model_executor.layers.linear import UnquantizedLinearMethod
from vllm.model_executor.layers.quantization.fp8 import Fp8LinearMethod
from vllm.model_executor.layers.quantization.online.fp8 import (
Fp8PerTensorOnlineLinearMethod,
)
weight = get_attribute_fallback(layer, ["weight", "qweight", "weight_packed"])
# Unquantized layer: just return base weights
if layer.quant_method is None or isinstance(
layer.quant_method, UnquantizedLinearMethod
):
return weight.to(out_dtype)
# Simple Fp8 case: rescale with tensor or block weight scales
if (
isinstance(
layer.quant_method, (Fp8LinearMethod, Fp8PerTensorOnlineLinearMethod)
)
and not layer.quant_method.use_marlin
# DeepGEMM transforms the scales using `transform_sf_into_required_layout` into
# a layout that is not compatible with `scaled_dequantize`.
and not layer.quant_method.use_deep_gemm
):
weight_scales = get_attribute_fallback(
layer, ["weight_scale", "weight_scale_inv"]
)
dequant_weights = scaled_dequantize(
weight,
weight_scales,
group_shape=layer.weight_block_size,
out_dtype=out_dtype,
)
# per-tensor scaling stores weights in [in, out] layout
if not layer.quant_method.block_quant:
dequant_weights = dequant_weights.T
return dequant_weights
# NOTE: Most generic base case
# - Call the layer with identity matrix which returns unquantized weights.
# - Must be used with extra care when dealing with static activation quantization:
# quantizing 1.0 may lead to over/underflows
# - Should only be used offline, since it's O(N^3)
assert hasattr(layer, "input_size_per_partition")
eye = torch.eye(
layer.input_size_per_partition,
dtype=out_dtype,
device=weight.device,
)
dequant_weights = layer.quant_method.apply(layer, eye, bias=None).to(out_dtype)
return dequant_weights.T
def pack_quantized_values_into_int32(
w_q: torch.Tensor, wtype: ScalarType, packed_dim: int = 0
):
@@ -789,25 +729,6 @@ SUPPORTED_GPTQ_QUANT_TYPES = [scalar_types.uint4b8, scalar_types.uint8b128]
SUPPORTED_GROUP_SIZES = [-1, 32, 64, 128]
def gptq_quantize_weights(
w: torch.Tensor,
quant_type: ScalarType,
group_size: int,
):
size_k, _ = w.shape
assert w.is_floating_point(), "w must be float"
assert quant_type in SUPPORTED_GPTQ_QUANT_TYPES, (
f"Unsupported gptq type = {quant_type}"
)
assert group_size in SUPPORTED_GROUP_SIZES + [size_k], (
f"Unsupported groupsize = {group_size}"
)
w_ref, w_q, w_s, _ = quantize_weights(w, quant_type, group_size)
return w_ref, w_q, w_s
def pack_rows(
q_w: torch.Tensor,
num_bits: int,
@@ -889,37 +810,6 @@ def unpack_cols(
return q_res
def gptq_pack(
q_w: torch.Tensor,
num_bits: int,
size_k: int,
size_n: int,
):
return pack_rows(q_w, num_bits, size_k, size_n)
def awq_pack(
q_w: torch.Tensor,
num_bits: int,
size_k: int,
size_n: int,
):
assert q_w.shape == (size_k, size_n)
# Interleave column dim (for the dequantize code) and pack it to int32
if num_bits == 4:
interleave = numpy.array([0, 2, 4, 6, 1, 3, 5, 7])
elif num_bits == 8:
interleave = numpy.array([0, 2, 1, 3])
else:
raise Exception("num_bits must be 4 or 8, got {}".format(num_bits))
q_w = q_w.reshape((-1, len(interleave)))[:, interleave].ravel()
q_w = q_w.reshape((-1, size_n)).contiguous()
return pack_cols(q_w, num_bits, size_k, size_n)
def convert_bf16_scales_to_fp8(
quant_fp8: Callable, scales: torch.Tensor
) -> tuple[torch.Tensor, torch.Tensor]:
vllm/model_executor/model_loader/tensorizer.py
-12
View File
@@ -5,7 +5,6 @@ import argparse
import contextlib
import contextvars
import dataclasses
import json
import os
import tempfile
import threading
@@ -75,17 +74,6 @@ def is_valid_deserialization_uri(uri: str | None) -> bool:
return False
def tensorizer_kwargs_arg(value):
loaded = json.loads(value)
if not isinstance(loaded, dict):
raise argparse.ArgumentTypeError(
f"Not deserializable to dict: {value}. serialization_kwargs and "
f"deserialization_kwargs must be "
f"deserializable from a JSON string to a dictionary. "
)
return loaded
class MetaTensorMode(TorchDispatchMode):
def __torch_dispatch__(self, func, types, args=(), kwargs=None):
kwargs = kwargs or {}
vllm/model_executor/model_loader/weight_utils.py
+1 -50
View File
@@ -13,7 +13,7 @@ import tempfile
import threading
import time
from collections.abc import Callable, Generator, Iterable
from contextlib import contextmanager, suppress
from contextlib import contextmanager
from pathlib import Path
from typing import IO, TYPE_CHECKING, Any
@@ -893,40 +893,6 @@ def np_cache_weights_iterator(
yield name, torch.from_numpy(param)
def drop_checkpoint_cache(
model_name_or_path: str,
revision: str | None = None,
cache_dir: str | None = None,
) -> None:
checkpoint_dir: Path | None = None
if os.path.isdir(model_name_or_path):
checkpoint_dir = Path(model_name_or_path)
else:
with suppress(OSError, ValueError):
checkpoint_dir = Path(
hf_api().snapshot_download(
model_name_or_path,
revision=revision,
cache_dir=cache_dir,
local_files_only=True,
)
)
if checkpoint_dir is None:
logger.info_once(
"No local checkpoint found for %s; skipping page-cache eviction",
model_name_or_path,
)
return
for path in checkpoint_dir.glob("*.safetensors"):
try:
with path.open("rb") as file:
os.posix_fadvise(file.fileno(), 0, 0, os.POSIX_FADV_DONTNEED)
except OSError as exc:
logger.warning(
"Could not release checkpoint page cache for %s: %s", path, exc
)
def _get_checkpoints_size_bytes(files: list[str]) -> int:
"""Return the total size of the checkpoint files in bytes."""
if not files:
@@ -1498,21 +1464,6 @@ def default_weight_loader(param: torch.Tensor, loaded_weight: torch.Tensor) -> N
raise
def row_parallel_weight_loader(
param: torch.Tensor, loaded_weight: torch.Tensor
) -> None:
"""Load weights that are row-parallelized."""
tp_rank = get_tensor_model_parallel_rank()
shard_dim = 0 if param.dim() != 1 else None
if shard_dim is not None:
shard_size = param.data.shape[shard_dim]
start_idx = tp_rank * shard_size
loaded_weight = loaded_weight.narrow(shard_dim, start_idx, shard_size)
return default_weight_loader(param, loaded_weight)
LoaderFunction = Callable[[torch.Tensor, torch.Tensor], None]
vllm/model_executor/models/adapters.py
-35
View File
@@ -11,11 +11,7 @@ import torch.nn as nn
from vllm.config import VllmConfig
from vllm.logger import init_logger
from vllm.model_executor.layers.activation import get_act_fn
from vllm.model_executor.models.config import VerifyAndUpdateConfig
from vllm.transformers_utils.config import (
try_get_dense_modules,
)
from vllm.transformers_utils.repo_utils import get_hf_file_bytes
from .interfaces import supports_multimodal
@@ -37,37 +33,6 @@ _GENERATE_SUFFIXES = [
]
def _load_st_projector(model_config: "ModelConfig") -> nn.Module | None:
"""Load Sentence-Transformers Dense projection layers."""
dense_modules = try_get_dense_modules(
model_config.model, revision=model_config.revision
)
if dense_modules is None:
return None
try:
layers = []
for layer_config in dense_modules:
folder = layer_config["folder"]
linear = nn.Linear(
layer_config["in_features"],
layer_config["out_features"],
bias=layer_config.get("bias", True),
dtype=model_config.head_dtype,
)
if not _load_dense_weights(linear, folder, model_config):
continue
layers.append(linear)
if act_name := layer_config.get("activation_function"):
layers.append(get_act_fn(act_name))
return nn.Sequential(*layers).to(dtype=model_config.head_dtype)
except Exception:
logger.exception("ST projector loading failed")
return None
def _load_dense_weights(
linear: nn.Linear, folder: str, model_config: "ModelConfig"
) -> bool:
vllm/model_executor/models/utils.py
-100
View File
@@ -601,26 +601,6 @@ def get_spec_layer_idx_from_weight_name(
return None
def skip_spec_layers(
weights: Iterable[tuple[str, torch.Tensor]], config: "PreTrainedConfig"
) -> Iterable[tuple[str, torch.Tensor]]:
"""Drop MTP spec-layer weights (loaded by the MTP head, not the base model).
Args:
weights: Iterable of `(name, tensor)` checkpoint pairs.
config: The model config, passed to `get_spec_layer_idx_from_weight_name`.
Yields:
`(name, tensor)` pairs whose weight is not an MTP-layer weight.
"""
return (
(name, w)
for name, w in weights
if get_spec_layer_idx_from_weight_name(config, name) is None
)
def init_vllm_registered_model(
vllm_config: VllmConfig,
*,
@@ -706,14 +686,6 @@ def _embedding_count_expression(embeddings: NestedTensors) -> str:
return " + ".join(_embedding_count_expression(inner) for inner in embeddings)
def split_list_into_ranges(lst: torch.Tensor, interval: int) -> list[list[int]]:
ranges: list[list[int]] = [[] for _ in range((max(lst) // interval) + 1)]
for num in lst:
index = num // interval
ranges[index].append(num)
return ranges
def _merge_multimodal_embeddings(
inputs_embeds: torch.Tensor,
multimodal_embeddings: NestedTensors,
@@ -941,17 +913,6 @@ def get_pp_missing_layer_names(model: torch.nn.Module) -> list[str]:
return missing_layer_names
def is_pp_missing_parameter(name: str, model: torch.nn.Module) -> bool:
"""Check if a parameter is missing in a pipeline parallel model."""
if isinstance(model, (StageMissingLayer, PPMissingLayer)):
return True
return any(
name.startswith(missing_layer_name)
for missing_layer_name in get_pp_missing_layer_names(model)
)
def make_empty_intermediate_tensors_factory(keys: list[str], hidden_size: int):
def make_empty_intermediate_tensors(
batch_size: int,
@@ -1047,32 +1008,6 @@ def cast_overflow_tensors(tensors: torch.Tensor, offset: float = 1000) -> torch.
return torch.clamp(tensors, min=-clamp_value, max=clamp_value)
def fast_topk(
values: torch.Tensor, topk: int, dim: int
) -> tuple[torch.Tensor, torch.Tensor]:
"""Optimized topk implementation that uses torch.max for k=1 case.
This function provides better performance for the common case of k=1
by using torch.max instead of the more general torch.topk.
Args:
values: Input tensor to find top-k values from
topk: Number of top values to return (k). Must be > 0.
dim: Dimension along which to compute topk
Returns:
Tuple of (values, indices) where values are the top-k values
and indices are their corresponding indices in the input tensor
"""
if topk == 1:
# Use max along the specified dimension to get both value and index
return torch.max(values, dim=dim, keepdim=True)
else:
# Use topk for efficiency with larger k values
return torch.topk(values, topk, dim=dim)
# Chunk x along the num_tokens axis for sequence parallelism
# NOTE: This is wrapped in a torch custom op to work around the following issue:
# The output tensor can have a sequence length 0 at small input sequence lengths
@@ -1142,20 +1077,6 @@ def process_eagle_weight(
model.has_own_embed_tokens = True
def get_layer_index(feature_layer_index: int, num_hidden_layers: int) -> int:
"""Given a signed vision feature layer, get the number of hidden layers
needed to leverage it.
Args:
feature_layer_index: Index of a required layer in the visual encoder.
num_hidden_layers: The total number of hidden layers in the visual encoder.
"""
if feature_layer_index < 0:
return num_hidden_layers + feature_layer_index + 1
return feature_layer_index
def scatter_output_slices(
output: torch.Tensor,
indices: list[int],
@@ -1171,24 +1092,3 @@ def scatter_output_slices(
dest[idx] = sliced.clone() if clone else sliced
offset += n_tok
def parse_diarized_timestamp(marker: str) -> float | None:
if (
not marker
or not marker.isascii()
or marker.count(".") > 1
or not marker.replace(".", "").isdigit()
):
return None
return float(marker)
def parse_diarized_speaker(speaker: str) -> str | None:
if (
len(speaker) < 2
or speaker[0] != "S"
or not speaker[1:].isascii()
or not speaker[1:].isdigit()
):
return None
return speaker
vllm/model_executor/models/vision.py
-139
View File
@@ -122,145 +122,6 @@ def _get_vision_feature_selector(
raise ValueError(f"Unexpected feature select strategy: {strategy!r}")
def get_num_selected_vision_tokens(
num_vision_tokens: int,
strategy: VisionFeatureSelectStrategy | str,
) -> int:
if callable(strategy):
dummy_features = torch.empty(1, num_vision_tokens, 64) # [B, L, D]
dummy_selected_features = strategy(dummy_features)
return dummy_selected_features.shape[1]
if strategy == "class":
return 1
if strategy == "default":
return num_vision_tokens - 1
if strategy == "full":
return num_vision_tokens
raise ValueError(f"Unexpected feature select strategy: {strategy!r}")
def resolve_visual_encoder_outputs(
encoder_outputs: torch.Tensor | list[torch.Tensor],
post_layer_norm: torch.nn.LayerNorm | None,
*,
select_layers: list[int] | None = None,
max_possible_layers: int | None = None,
last_hs_proc: Callable[[torch.Tensor], torch.Tensor] | None = None,
feature_select_strategy: VisionFeatureSelectStrategy | None = None,
) -> torch.Tensor:
"""Given the outputs a visual encoder module that may correspond to the
output of the last layer, or a list of hidden states to be stacked,
handle post normalization and resolve it into a single output tensor.
Args:
encoder_outputs: Output of encoder's last layer or all hidden states.
post_layer_norm: Post norm to apply to the output of the encoder.
select_layers: Optional layer indices to grab from the encoder
outputs; if provided, encoder outputs must be a list.
max_possible_layers: Total layers in the fully loaded visual encoder.
last_hs_proc: Optional callable to be applied to the last layer if it
is used, e.g., pooling head for Siglip. This is done prior to
feature selection and layer normalization. If select_layers are
provided, the output of last_hs_proc must be able to be
concatenated with the other select_layers along the last dimension.
feature_select_strategy: Defines how to select the hidden states
from each layer.
"""
if select_layers is None:
if not isinstance(encoder_outputs, torch.Tensor):
raise ValueError(
"Expected only a single encoder output when "
"`select_layers` is not provided"
)
# Preprocess the encoder outputs as needed, e.g., map head
# and layer norm for siglip, which runs before feature selection
if last_hs_proc is not None:
encoder_outputs = last_hs_proc(encoder_outputs)
if feature_select_strategy is not None:
select_features = _get_vision_feature_selector(feature_select_strategy)
encoder_outputs = select_features(encoder_outputs)
if post_layer_norm is not None:
return post_layer_norm(encoder_outputs)
return encoder_outputs
if max_possible_layers is None:
raise ValueError(
"`max_possible_layers` must be provided alongside `select_layers`"
)
# Get the hidden states corresponding to the layer indices.
# Negative values are relative to the full visual encoder,
# so offset them depending on how many layers were loaded.
# NOTE: this assumes that encoder_outputs is a list containing
# the inputs to the visual encoder, followed by the hidden states
# of each layer.
num_loaded_layers = len(encoder_outputs) - 1
offset = max_possible_layers - num_loaded_layers
hs_pool = [
encoder_outputs[layer_idx]
if layer_idx >= 0
else encoder_outputs[layer_idx + offset]
for layer_idx in select_layers
]
uses_last_layer = select_layers[-1] in (max_possible_layers - 1, -1)
if last_hs_proc is not None and uses_last_layer:
hs_pool[-1] = last_hs_proc(hs_pool[-1])
if feature_select_strategy is not None:
select_features = _get_vision_feature_selector(feature_select_strategy)
hs_pool = [select_features(hs) for hs in hs_pool]
# Apply post-norm on the final hidden state if we are using it
if post_layer_norm is not None and uses_last_layer:
hs_pool[-1] = post_layer_norm(hs_pool[-1])
return torch.cat(hs_pool, dim=-1)
def run_dp_sharded_vision_model(
image_input: torch.Tensor, vision_model: torch.nn.Module
) -> torch.Tensor:
"""Run a vision model with data parallelism (DP) sharding. The function
will shard the input image tensor on the first dimension and run the vision
model
Args:
image_input (torch.Tensor): Image input tensor.
vision_model (torch.nn.Module): Vision model.
Returns:
torch.Tensor: Output image embeddings
"""
num_chunks = image_input.shape[0]
mp_world_size = get_tensor_model_parallel_world_size()
num_chunks_per_rank = (num_chunks + mp_world_size - 1) // mp_world_size
num_padded_chunks = num_chunks_per_rank * mp_world_size - num_chunks
pad = (0,) * (2 * (image_input.dim() - 1)) + (0, num_padded_chunks)
image_input_padded = torch.nn.functional.pad(image_input, pad)
rank = get_tensor_model_parallel_rank()
image_input_per_rank = image_input_padded[
rank * num_chunks_per_rank : (rank + 1) * num_chunks_per_rank, ...
]
vision_embeddings = vision_model(image_input_per_rank)
# Ensure tensor is contiguous before all_gather
vision_embeddings = vision_embeddings.contiguous()
vision_embeddings = tensor_model_parallel_all_gather(vision_embeddings, dim=0)
vision_embeddings = vision_embeddings[:num_chunks, ...]
return vision_embeddings
def get_load_balance_assignment(
sizes: list[int],
num_gpus: int = 2,
vllm/model_executor/utils.py
-38
View File
@@ -144,44 +144,6 @@ def replace_parameter(
setattr(layer, param_name, new_param)
def get_packed_modules_mapping(model: torch.nn.Module) -> dict[str, list[str]]:
parent_map = getattr(model, "packed_modules_mapping", None)
parent_map = copy.deepcopy(parent_map) if parent_map is not None else {}
# don't infer mapping if the model has defined it explicitly.
if parent_map:
return parent_map
# We only check main components instead of whole model submodules
for child in model.children():
child_map = getattr(child, "packed_modules_mapping", None)
child_map = copy.deepcopy(child_map) if child_map is not None else {}
if any((k in parent_map and parent_map[k] != v) for k, v in child_map.items()):
raise ValueError(
f"Can't update {type(model).__name__}'s packed_modules_mapping "
f"safely because of conflicts from {type(child).__name__}."
)
else:
parent_map.update(child_map)
return parent_map
def get_moe_expert_mapping(
model: torch.nn.Module,
) -> list[tuple[str, str, int, str]]:
"""Get the expert mapping from a model.
It will be retrieved from the first module that has a `get_expert_mapping` method.
If the model manually implements `get_expert_mapping`, it will be used.
Otherwise, it will use the first RoutedExperts layer."""
for _, module in model.named_modules():
get_mapping = getattr(module, "get_expert_mapping", None)
if get_mapping is not None:
return get_mapping()
raise ValueError("No module in the model has a `get_expert_mapping` method.")
def maybe_disable_graph_partition(current_backend: str) -> dict[str, bool]:
if current_backend == "inductor" and is_torch_equal_or_newer("2.9.0.dev"):
return {"graph_partition": False}
vllm/model_executor/warmup/deep_gemm_warmup.py
-132
View File
@@ -40,9 +40,6 @@ from vllm.utils.math_utils import cdiv
from vllm.utils.platform_utils import num_compute_units
from vllm.v1.worker.workspace import current_workspace_manager
if TYPE_CHECKING:
from vllm.models.deepseek_v4.nvidia.model import DeepseekV4MoE
from vllm.models.kimi_k3.nvidia.model import KimiMoE
def _generate_optimal_warmup_m_values(
@@ -156,22 +153,6 @@ def _fused_moe_grouped_gemm_may_use_deep_gemm(module: torch.nn.Module) -> bool:
return isinstance(fused_experts, (DeepGemmExperts, TritonOrDeepGemmExperts))
# Model construction already imports the model modules. Avoid importing them here.
def _mega_moe_uses_mega_gate(module: torch.nn.Module) -> "TypeGuard[DeepseekV4MoE]":
# DeepSeek-V4 routes with bf16_mega_gate, Kimi K3 with its GateLinear.
dsv4 = sys.modules.get("vllm.models.deepseek_v4.nvidia.model")
return dsv4 is not None and isinstance(module, dsv4.DeepseekV4MoE)
def _mega_moe_may_use_deep_gemm(
module: torch.nn.Module,
) -> "TypeGuard[DeepseekV4MoE | KimiMoE]":
if _mega_moe_uses_mega_gate(module):
return module.use_native_mega_moe
kimi = sys.modules.get("vllm.models.kimi_k3.nvidia.model")
return kimi is not None and isinstance(module, kimi.KimiMoE) and module.use_mega_moe
FP8_GEMM_NT_WARMUP_CACHE: set[torch.Size] = set()
@@ -333,102 +314,6 @@ def _deepgemm_grouped_fp8_gemm_nt_contiguous_warmup(
GROUPED_FP8_GEMM_NT_CONTIGUOUS_WARMUP_CACHE.add(w.size())
FP8_FP4_MEGA_MOE_WARMUP_CACHE: set[tuple] = set()
def _get_fp8_fp4_mega_moe_warmup_key(module: "DeepseekV4MoE | KimiMoE") -> tuple:
gate = module.gate
return (
gate.output_size,
getattr(gate, "tid2eid", None) is not None,
gate.e_score_correction_bias is not None,
getattr(module, "image_sentinel_lo", 0),
module.experts.has_fused_shared_experts,
)
def _get_fp8_fp4_mega_moe_m_values(
module: "DeepseekV4MoE | KimiMoE", max_tokens: int
) -> list[int]:
"""Get the largest M per distinct fp8_fp4_mega_moe config and, where used,
bf16_mega_gate config."""
dg = _import_deep_gemm()
gate, experts = module.gate, module.experts
uses_mega_gate = _mega_moe_uses_mega_gate(module)
if module.use_sequence_parallel:
max_tokens = cdiv(max_tokens, module.tp_size)
m_values: dict[tuple, int] = {}
for num_tokens in range(1, max_tokens + 1):
gate_config = (
dg.get_bf16_mega_gate_config(
num_tokens, gate.input_size, gate.output_size, experts.top_k
)
if uses_mega_gate and hasattr(dg, "get_bf16_mega_gate_config")
else {}
)
block_m = experts.get_block_m(num_tokens)
# Keep the largest M: DeepSeek-V4 only calls bf16_mega_gate above a
# small token threshold, and a config's largest M is always above it.
m_values[(tuple(gate_config.items()), block_m)] = num_tokens
return sorted(m_values.values())
def _deepgemm_fp8_fp4_mega_moe_warmup(
module: "DeepseekV4MoE | KimiMoE",
max_tokens: int,
pbar: tqdm | None = None,
):
from vllm.models.deepseek_v4.nvidia.model import (
prepare_mega_gate_routing_metadata,
)
key = _get_fp8_fp4_mega_moe_warmup_key(module)
if key in FP8_FP4_MEGA_MOE_WARMUP_CACHE:
return
m_values = _get_fp8_fp4_mega_moe_m_values(module, max_tokens)
max_m = m_values[-1] # sorted
device = module.experts.get_symm_buffer().x.device
input_ids = torch.zeros(max_m, device=device, dtype=torch.int64)
hidden_states = torch.zeros(
(max_m, module.gate.input_size), device=device, dtype=torch.bfloat16
)
is_padding = torch.zeros(max_m + 1, device=device, dtype=torch.bool)
# With padding skip on, sequence parallel gives ranks > 0 a padding mask
# slice at any alignment.
sp_masked = module.use_sequence_parallel and envs.VLLM_MOE_SKIP_PADDING
offsets = (0, 1) if sp_masked else (0,)
# Run the whole layer so the gate sees the same optional inputs as serving.
for num_tokens in m_values:
inputs: tuple = (hidden_states[:num_tokens],)
if _mega_moe_uses_mega_gate(module):
ids = input_ids[:num_tokens]
inputs += (
ids,
prepare_mega_gate_routing_metadata(
ids,
has_hash_routing=module.gate.tid2eid is not None,
image_sentinel_base_id=module.image_sentinel_lo or None,
),
)
for offset in offsets:
forward_context = ForwardContext(
no_compile_layers={},
attn_metadata={},
slot_mapping={},
is_padding=is_padding[offset : offset + num_tokens],
)
with override_forward_context(forward_context):
module(*inputs)
if pbar is not None:
pbar.update(1)
FP8_FP4_MEGA_MOE_WARMUP_CACHE.add(key)
def deepgemm_fp8_gemm_nt_warmup(
model: torch.nn.Module, max_tokens: int, pbar: tqdm | None = None
):
@@ -454,21 +339,11 @@ def deepgemm_grouped_fp8_gemm_nt_contiguous_warmup(
)
@torch.inference_mode()
def deepgemm_fp8_fp4_mega_moe_warmup(
model: torch.nn.Module, max_tokens: int, pbar: tqdm | None = None
):
for m in model.modules():
if _mega_moe_may_use_deep_gemm(m):
_deepgemm_fp8_fp4_mega_moe_warmup(m, max_tokens, pbar=pbar)
def _count_warmup_iterations(model: torch.nn.Module, max_tokens: int) -> int:
seen_fp8_sizes: set[torch.Size] = set(FP8_GEMM_NT_WARMUP_CACHE)
seen_grouped_sizes: set[torch.Size] = set(
GROUPED_FP8_GEMM_NT_CONTIGUOUS_WARMUP_CACHE
)
seen_mega_moe_keys: set[tuple] = set(FP8_FP4_MEGA_MOE_WARMUP_CACHE)
total = 0
for m in model.modules():
@@ -489,11 +364,6 @@ def _count_warmup_iterations(model: torch.nn.Module, max_tokens: int) -> int:
if w2.size() not in seen_grouped_sizes:
total += n_values
seen_grouped_sizes.add(w2.size())
elif _mega_moe_may_use_deep_gemm(m):
key = _get_fp8_fp4_mega_moe_warmup_key(m)
if key not in seen_mega_moe_keys:
total += len(_get_fp8_fp4_mega_moe_m_values(m, max_tokens))
seen_mega_moe_keys.add(key)
return total
@@ -508,8 +378,6 @@ def deep_gemm_warmup(model: torch.nn.Module, max_tokens: int):
with tqdm(total=total, desc="DeepGEMM warmup") as pbar:
deepgemm_fp8_gemm_nt_warmup(model, max_tokens, pbar)
deepgemm_grouped_fp8_gemm_nt_contiguous_warmup(model, max_tokens, pbar)
deepgemm_fp8_fp4_mega_moe_warmup(model, max_tokens, pbar)
else:
deepgemm_fp8_gemm_nt_warmup(model, max_tokens, None)
deepgemm_grouped_fp8_gemm_nt_contiguous_warmup(model, max_tokens, None)
deepgemm_fp8_fp4_mega_moe_warmup(model, max_tokens, None)
vllm/model_executor/warmup/flashinfer_autotune_cache.py
-16
View File
@@ -3,9 +3,7 @@
"""FlashInfer autotune cache helpers."""
import hashlib
import os
import tempfile
from contextlib import suppress
from pathlib import Path
from typing import TYPE_CHECKING
@@ -142,17 +140,3 @@ def sync_flashinfer_autotune_cache(
if not AutoTuner.get().load_configs(f.name):
raise RuntimeError("FlashInfer autotune cache is incompatible")
def write_flashinfer_autotune_cache(cache_path: Path, contents: bytes) -> None:
cache_path.parent.mkdir(parents=True, exist_ok=True)
fd, tmp_path = tempfile.mkstemp(
dir=cache_path.parent, suffix=".tmp", prefix=f".{cache_path.name}."
)
try:
with os.fdopen(fd, "wb") as f:
f.write(contents)
os.replace(tmp_path, cache_path)
except BaseException:
with suppress(OSError):
os.unlink(tmp_path)
raise
vllm/model_executor/warmup/jit_warmup_triton_helper.py
-19
View File
@@ -736,25 +736,6 @@ class TritonPointerInputVariant:
return TritonWarmupTensor(dtype, aligned=self.is_aligned(name), shape=shape)
def _triton_constexpr_arg_names(
kernel: Callable[..., Any],
function_def: ast.FunctionDef,
arg_names: tuple[str, ...],
) -> frozenset[str]:
constexprs = getattr(kernel, "constexprs", None)
if constexprs is not None:
return frozenset(arg_names[index] for index in constexprs)
names: set[str] = set()
for arg in function_def.args.args + function_def.args.kwonlyargs:
if arg.annotation is None:
continue
annotation = get_ast_full_name(arg.annotation)
if annotation in ("tl.constexpr", "triton.language.constexpr", "constexpr"):
names.add(arg.arg)
return frozenset(names)
def _leftmost_name(node: ast.AST) -> str | None:
if isinstance(node, ast.Name):
return node.id
vllm/models/deepseek_v4/common/eplb_util.py
-32
View File
@@ -2,8 +2,6 @@
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Shared EPLB MoE registration helpers for DeepSeek V4 family models."""
from collections.abc import Iterable
from typing import Any
import torch.nn as nn
@@ -18,33 +16,3 @@ def dspark_draft_supports_eplb(draft_model_config: ModelConfig) -> bool:
"""
return getattr(draft_model_config.hf_config, "model_type", None) == "deepseek_v4"
def collect_moe_layers(
moe_model: Any,
layers: Iterable[nn.Module],
config: object,
*,
decoder_layer_type: type[nn.Module],
moe_type: type[nn.Module],
) -> None:
"""Populate ``MixtureOfExperts`` fields from decoder layers for EPLB.
Works for both target backbones and DSpark draft stacks. Draft models pass
their shorter ``model.layers`` list; pipeline-parallel placeholder layers
are skipped via ``decoder_layer_type``.
"""
moe_model.num_expert_groups = getattr(config, "n_group", 1)
moe_model.moe_layers = []
moe_model.moe_mlp_layers = []
example_moe = None
for layer in layers:
if not isinstance(layer, decoder_layer_type):
continue
if not isinstance(layer.ffn, moe_type):
continue
example_moe = layer.ffn
moe_model.moe_mlp_layers.append(layer.ffn)
moe_model.moe_layers.append(layer.ffn.experts)
moe_model.num_moe_layers = len(moe_model.moe_layers)
moe_model.extract_moe_parameters(example_moe)
vllm/multimodal/audio.py
-17
View File
@@ -34,23 +34,6 @@ except ImportError:
torchaudio = PlaceholderModule("torchaudio") # type: ignore[assignment]
# ============================================================
# Aligned with `librosa.get_duration` function
def get_audio_duration(*, y: npt.NDArray[np.floating], sr: float = 22050) -> float:
"""Get the duration of an audio array in seconds.
Args:
y: Audio time series. Can be 1D (samples,) or 2D (channels, samples).
sr: Sample rate of the audio in Hz.
Returns:
Duration of the audio in seconds.
"""
n_samples = y.shape[-1]
return float(n_samples) / sr
class ChannelReduction(str, Enum):
"""Method to reduce multi-channel audio to target channels."""
vllm/multimodal/inputs.py
-7
View File
@@ -333,13 +333,6 @@ A dictionary containing nested tensors which have been batched via
"""
def batched_tensors_equal(a: BatchedTensorInputs, b: BatchedTensorInputs) -> bool:
"""Equality check between
[`BatchedTensorInputs`][vllm.multimodal.inputs.BatchedTensorInputs] objects.
"""
return all(k in b and nested_tensors_equal(a[k], b[k]) for k in a)
@dataclass
class MultiModalFeatureSpec:
"""Represents a single multimodal input with its processed data and metadata.
vllm/multimodal/utils.py
-16
View File
@@ -410,19 +410,3 @@ def copy_mm_embedding_modality(
dst.modality = modality # type: ignore[attr-defined]
return dst
def get_mm_embedding_modalities(
multimodal_embeddings: Sequence["torch.Tensor"],
) -> list[str]:
"""Collect per-embedding modalities previously set on the tensors."""
modalities: list[str] = []
for i, emb in enumerate(multimodal_embeddings):
modality = getattr(emb, "modality", None)
if modality is None:
raise ValueError(
f"Missing modality on multimodal embedding at index {i}. "
"Encoder gather must set embed.modality before interleaved "
"audio-in-video merge."
)
modalities.append(modality)
return modalities
vllm/multimodal/video_decoders/opencv.py
+1 -31
View File
@@ -9,11 +9,7 @@ import numpy.typing as npt
from vllm.logger import init_logger
from vllm.utils.import_utils import PlaceholderModule
from .base import (
VideoSourceMetadata,
VideoTargetMetadata,
check_frame_pixel_limit,
)
from .base import VideoSourceMetadata
try:
import cv2
@@ -25,32 +21,6 @@ except ImportError:
logger = init_logger(__name__)
def decode_opencv(
loader_cls,
data: bytes,
target: VideoTargetMetadata,
sampling_kwargs: dict,
*,
frame_recovery: bool = False,
) -> tuple[npt.NDArray, VideoSourceMetadata, list[int], list[int]]:
cap = OpenCVVideoBackendMixin.open_video_capture(data)
check_frame_pixel_limit(
int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)),
int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)),
)
source = loader_cls._prepare_source(OpenCVVideoBackendMixin.get_video_metadata(cap))
frame_idx = loader_cls.compute_frames_index_to_sample(
source=source, target=target, **sampling_kwargs
)
frames, valid = loader_cls.read_frames(
cap,
frame_idx,
total_frames_num=source.total_frames_num,
frame_recovery=frame_recovery,
)
return frames, source, frame_idx, valid
class OpenCVVideoBackendMixin:
@staticmethod
def get_cv2_video_api():
vllm/multimodal/video_decoders/pynvvideocodec.py
-25
View File
@@ -14,7 +14,6 @@ from vllm.logger import init_logger
from vllm.utils.mem_constants import MiB_bytes
from .base import (
PYNVVIDEOCODEC_DEFAULT_HW_DECODERS,
VideoSourceMetadata,
VideoTargetMetadata,
check_frame_pixel_limit,
@@ -23,30 +22,6 @@ from .base import (
logger = init_logger(__name__)
def decode_pynvvideocodec(
loader_cls,
data: bytes,
target: VideoTargetMetadata,
sampling_kwargs: dict,
*,
hw_decoders: int = PYNVVIDEOCODEC_DEFAULT_HW_DECODERS,
max_width: int | None = None,
max_height: int | None = None,
) -> tuple[npt.NDArray, VideoSourceMetadata, list[int], list[int]]:
max_width, max_height = validate_pynvvideocodec_max_dimensions(
max_width, max_height
)
PyNvVideoCodecVideoBackendMixin._configure_decoder_slots(hw_decoders)
return PyNvVideoCodecVideoBackendMixin.decode_frames_pynvvideocodec(
loader_cls,
data,
target,
max_width=max_width,
max_height=max_height,
**sampling_kwargs,
)
class PyNvVideoCodecSourceMetadata(NamedTuple):
"""Metadata needed before GPU video decode."""
vllm/parser/utils.py
-51
View File
@@ -97,54 +97,3 @@ def scan_json_object_end(raw: str, start: int) -> int | None:
return i + 1
return None
def extract_json_member_object(raw: str, key: str) -> str | None:
"""Extract the raw text span of the top-level object-valued member ``key``
from a (possibly incomplete) ``{...}`` JSON wrapper.
Returns the verbatim substring (prefix-stable as the input grows, which
the engine's argument-delta diffing relies on), possibly an unterminated
object prefix; ``None`` when the value has not started. Raises
``ValueError`` when the member is present but its value is not a JSON
object. A verbatim substring (not re-serialized JSON) is required because
the engine diffs successive converter outputs and expects each to extend
the previous one; the scanner is string/escape-aware so a matching literal
inside another value cannot mislead it, and it recovers a partial span when
the wrapper is truncated (where ``json.loads`` would fail)."""
depth = 0
in_string = False
escape = False
string_start = -1
last_string: str | None = None
for i, ch in enumerate(raw):
if escape:
escape = False
continue
if in_string:
if ch == "\\":
escape = True
elif ch == '"':
in_string = False
if depth == 1:
last_string = raw[string_start + 1 : i]
continue
if ch == '"':
in_string = True
string_start = i
elif ch == ":" and depth == 1 and last_string == key:
value_start = i + 1
while value_start < len(raw) and raw[value_start] in _JSON_WS:
value_start += 1
if value_start >= len(raw):
return None
if raw[value_start] != "{":
raise ValueError(f"member {key!r} must be a JSON object")
value_end = scan_json_object_end(raw, value_start)
if value_end is None:
return raw[value_start:]
return raw[value_start:value_end]
elif ch in "{[":
depth += 1
elif ch in "}]":
depth -= 1
return None
vllm/platforms/rocm.py
-92
View File
@@ -5,7 +5,6 @@ import os
from functools import wraps
from typing import TYPE_CHECKING
import regex as re
import torch
import vllm.envs as envs
@@ -152,77 +151,6 @@ _ON_RDNA = _ON_GFX1X and not _ON_CDNA
_ON_RDNA4 = any(arch in _GCN_ARCH for arch in ["gfx1200", "gfx1201"])
def _capability_from_gcn_arch(gcn_arch: str) -> tuple[int, int] | None:
"""Parse (major, minor) from a GCN arch string, mirroring how
HIP derives hipDeviceProp_t.major / .minor.
Format: gfx<MAJOR><MINOR><STEPPING>
- 1-digit major (gfx9xx): "gfx" + M + m + stepping
- 2-digit major (gfx1xxx): "gfx" + MM + m + stepping
Examples:
gfx90a -> (9, 0) gfx942 -> (9, 4) gfx950 -> (9, 5)
gfx1100 -> (11, 0) gfx1101 -> (11, 0) gfx1200 -> (12, 0)
Returns None only when the string is not gfx-prefixed at all
(i.e. not a ROCm arch string). Raises on any string that looks
like a GCN arch but does not match a known layout.
"""
m = re.match(r"gfx(\d+)", gcn_arch)
if not m:
# Not a gfx string at all — caller should fall back to torch.cuda
return None
digits = m.group(1)
n = len(digits)
if n < 2:
raise ValueError(
f"GCN arch '{gcn_arch}' has too few digits ({n}) after 'gfx' "
f"to derive a (major, minor) capability. "
f"Please file a vLLM issue with your GPU model."
)
if n in (2, 3):
# 1-digit major: gfx9 family
# len 2: major + minor (e.g. gfx90 from gfx90a)
# len 3: major + minor + step (e.g. gfx942)
major = int(digits[0])
minor = int(digits[1])
elif n == 4:
# 2-digit major: gfx10xx, gfx11xx, gfx12xx
# major(2) + minor(1) + stepping(1)
major = int(digits[:2])
minor = int(digits[2])
elif n >= 5:
raise ValueError(
f"GCN arch '{gcn_arch}' has {n} digits after 'gfx', which "
f"exceeds the known 4-digit layout (MMms). Cannot determine "
f"major/minor split unambiguously. "
f"Please file a vLLM issue with your GPU model."
)
if major < 9:
raise ValueError(
f"Parsed unknown ROCm architecture from GCN arch '{gcn_arch}': "
f"major={major}, minor={minor}. "
f"Major version < 9 is not expected for any supported AMD GPU. "
f"Please file a vLLM issue with your GPU model."
)
if major > 12:
raise ValueError(
f"Parsed unknown ROCm architecture from GCN arch '{gcn_arch}': "
f"major={major}, minor={minor}. "
f"Major version > 12 is beyond currently known AMD generations. "
f"Please file a vLLM issue with your GPU model so support "
f"can be added."
)
return (major, minor)
def on_gfx1x() -> bool:
return _ON_GFX1X and not _ON_CDNA
@@ -235,14 +163,6 @@ def on_gfx1100() -> bool:
return _ON_GFX1100
def on_gfx1151() -> bool:
return _ON_GFX1151
def on_gfx115x() -> bool:
return _ON_GFX115X
def on_gfx12x() -> bool:
return _ON_GFX12X and not _ON_CDNA
@@ -259,10 +179,6 @@ def on_mi3xx() -> bool:
return _ON_MI3XX
def on_gfx9() -> bool:
return _ON_GFX9
def on_gfx90a() -> bool:
return _ON_GFX90A
@@ -275,14 +191,6 @@ def on_gfx950() -> bool:
return _ON_GFX950
def on_cdna() -> bool:
return _ON_CDNA
def on_rdna() -> bool:
return _ON_RDNA
def get_cdna_version() -> int:
if on_gfx90a():
return 2
vllm/tool_parsers/utils.py
+1 -95
View File
@@ -8,7 +8,7 @@ import math
import warnings
from collections.abc import Sequence
from dataclasses import dataclass
from json import JSONDecodeError, JSONDecoder
from json import JSONDecodeError
from typing import Any, TypeAlias
import partial_json_parser
@@ -20,11 +20,8 @@ from openai.types.responses import (
ToolChoiceFunction,
)
from openai.types.responses.tool import Tool as ResponsesTool
from partial_json_parser.core.options import Allow
from vllm.entrypoints.generate.base.protocol import (
DeltaFunctionCall,
DeltaToolCall,
FunctionCall,
ToolCall,
)
@@ -66,25 +63,6 @@ def custom_tool_as_function_dict(tool: CustomTool) -> dict[str, Any]:
}
def safe_literal_eval(text: str):
with warnings.catch_warnings():
warnings.simplefilter("ignore", SyntaxWarning)
return ast.literal_eval(text)
def partial_tag_overlap(text: str, tag: str) -> int:
"""Length of the longest prefix of *tag* that matches a suffix of *text*.
E.g. text ending in ``"<tool_"`` returns 6 when tag is ``"<tool_call>"``.
Returns 0 when there is no overlap.
"""
max_check = min(len(tag) - 1, len(text))
for k in range(max_check, 0, -1):
if text.endswith(tag[:k]):
return k
return 0
def find_common_prefix(s1: str, s2: str) -> str:
"""Finds a common prefix that is shared between two strings, if there is one.
Order of arguments is NOT important.
@@ -156,26 +134,6 @@ def extract_intermediate_diff(curr: str, old: str) -> str:
return diff
# partial_json_parser doesn't support extra data and
# JSONDecoder.raw_decode doesn't support partial JSON
def partial_json_loads(input_str: str, flags: Allow) -> tuple[Any, int]:
try:
return (partial_json_parser.loads(input_str, flags), len(input_str))
except JSONDecodeError as e:
if "Extra data" in e.msg:
dec = JSONDecoder()
return dec.raw_decode(input_str)
raise
def is_complete_json(input_str: str) -> bool:
try:
json.loads(input_str)
return True
except JSONDecodeError:
return False
def _is_json_finite(obj: Any) -> bool:
"""Whether *obj* can be serialized to valid JSON.
@@ -190,12 +148,6 @@ def _is_json_finite(obj: Any) -> bool:
return False
def consume_space(i: int, s: str) -> int:
while i < len(s) and s[i].isspace():
i += 1
return i
_NAMESPACE_TOOL_SEPARATOR = "__"
@@ -1237,49 +1189,3 @@ def coerce_to_schema_type(value: str, schema_type: str | list[str]) -> Any:
return value
return parsed
def compute_tool_delta(
previously_sent_args: str,
new_call: ToolCall,
index: int,
withheld_suffix: str,
) -> DeltaToolCall | None:
"""Compute the incremental delta between previously streamed arguments
and the current tool call state.
Returns:
A DeltaToolCall with only the new argument characters, or None
if there is no difference from what was previously sent.
"""
new_call_args = new_call.function.arguments
if withheld_suffix:
if not new_call_args.endswith(withheld_suffix):
msg = (
f"Tool call arguments '{new_call_args}' do not end with "
f"expected withheld suffix '{withheld_suffix}'"
)
logger.error(msg)
raise ValueError(msg)
new_call_args = new_call_args[: -len(withheld_suffix)]
if not previously_sent_args:
return DeltaToolCall(
id=new_call.id,
type="function",
index=index,
function=DeltaFunctionCall(
name=new_call.function.name,
arguments=new_call_args,
),
)
arg_diff = new_call_args[len(previously_sent_args) :]
return (
DeltaToolCall(
id=None,
index=index,
function=DeltaFunctionCall(arguments=arg_diff),
)
if arg_diff
else None
)
vllm/transformers_utils/chat_templates/registry.py
-15
View File
@@ -16,21 +16,6 @@ ChatTemplatePath: TypeAlias = Path | Callable[[str], Path | None]
_MODEL_TYPE_TO_CHAT_TEMPLATE_FALLBACK: dict[str, ChatTemplatePath] = {}
def register_chat_template_fallback_path(
model_type: str,
chat_template: ChatTemplatePath,
) -> None:
if model_type in _MODEL_TYPE_TO_CHAT_TEMPLATE_FALLBACK:
logger.warning(
"Model type %s already has a chat template registered. "
"It will be overwritten by the new chat template %s.",
model_type,
chat_template,
)
_MODEL_TYPE_TO_CHAT_TEMPLATE_FALLBACK[model_type] = chat_template
def get_chat_template_fallback_path(
model_type: str,
tokenizer_name_or_path: str,
vllm/transformers_utils/processor.py
-37
View File
@@ -418,18 +418,6 @@ def get_feature_extractor(
cached_get_feature_extractor = lru_cache(get_feature_extractor)
def cached_feature_extractor_from_config(
model_config: "ModelConfig",
**kwargs: Any,
):
return cached_get_feature_extractor(
model_config.model,
revision=model_config.revision,
trust_remote_code=model_config.trust_remote_code,
**_merge_mm_kwargs(model_config, AutoFeatureExtractor, **kwargs),
)
def get_image_processor(
processor_name: str,
*args: Any,
@@ -471,18 +459,6 @@ def get_image_processor(
cached_get_image_processor = lru_cache(get_image_processor)
def cached_image_processor_from_config(
model_config: "ModelConfig",
**kwargs: Any,
):
return cached_get_image_processor(
model_config.model,
revision=model_config.revision,
trust_remote_code=model_config.trust_remote_code,
**_merge_mm_kwargs(model_config, AutoImageProcessor, **kwargs),
)
def get_video_processor(
processor_name: str,
*args: Any,
@@ -523,16 +499,3 @@ def get_video_processor(
cached_get_video_processor = lru_cache(get_video_processor)
def cached_video_processor_from_config(
model_config: "ModelConfig",
processor_cls: type[_V] | None = None,
**kwargs: Any,
):
return cached_get_video_processor(
model_config.model,
revision=model_config.revision,
trust_remote_code=model_config.trust_remote_code,
processor_cls_overrides=processor_cls, # type: ignore[arg-type]
**_merge_mm_kwargs(model_config, AutoVideoProcessor, **kwargs),
)
vllm/triton_utils/force_first_config.py
-5
View File
@@ -10,11 +10,6 @@ logger = init_logger(__name__)
_installed: bool = False
def is_installed() -> bool:
"""Return whether the first-valid-config patch is currently installed."""
return _installed
def install() -> None:
"""Install the Autotuner.run replacement."""
global _installed
vllm/usage/usage_lib.py
-5
View File
@@ -42,11 +42,6 @@ _USAGE_ENV_VARS_TO_COLLECT = [
]
def set_runtime_usage_data(key: str, value: str | int | bool) -> None:
"""Set global usage data that will be sent with every usage heartbeat."""
_GLOBAL_RUNTIME_DATA[key] = value
def is_usage_stats_enabled():
"""Determine whether or not we can send usage stats to the server.
The logic is as follows:
vllm/utils/async_utils.py
+1 -25
View File
@@ -9,7 +9,7 @@ import asyncio
import contextlib
from asyncio import FIRST_COMPLETED, AbstractEventLoop, Future, Task
from collections.abc import AsyncGenerator, Awaitable, Callable
from concurrent.futures import Executor, ThreadPoolExecutor
from concurrent.futures import Executor
from functools import partial
from typing import TYPE_CHECKING, TypeVar
@@ -43,30 +43,6 @@ def make_async(
return _async_wrapper
def make_async_with_semaphore(
func: Callable[P, T],
executor: ThreadPoolExecutor,
) -> Callable[P, Awaitable[T]]:
"""Take a blocking function, and run it on in an executor thread.
This function prevents the blocking function from blocking the
asyncio event loop.
The code in this function needs to be thread safe.
The function is wrapped in a semaphore to limit the number of
concurrent executions making it easier to cancel tasks before they start.
"""
semaphore = asyncio.Semaphore(executor._max_workers)
async def _async_wrapper(*args: P.args, **kwargs: P.kwargs) -> T:
loop = asyncio.get_event_loop()
p_func = partial(func, *args, **kwargs)
async with semaphore:
return await loop.run_in_executor(executor, p_func)
return _async_wrapper
def run_in_loop(loop: AbstractEventLoop, function: Callable, *args):
if in_loop(loop):
function(*args)
vllm/utils/b12x.py
-9
View File
@@ -48,11 +48,6 @@ _B12X_SUBMODULES = {
}
def has_b12x() -> bool:
"""Return whether the B12X package is installed."""
return _HAS_B12X
def _get_submodule(module_name: str) -> ModuleType | None:
return _B12X_SUBMODULES.get(module_name)
@@ -77,10 +72,6 @@ def get_b12x_fused_moe() -> ModuleType | None:
return _get_submodule("b12x.moe.fused_moe")
def get_b12x_paged_attention() -> ModuleType | None:
return _get_submodule("b12x.attention.paged")
def b12x_warmup_token_counts(
*,
max_tokens: int,
vllm/utils/collection_utils.py
+1 -7
View File
@@ -7,7 +7,7 @@ This is similar in concept to the `collections` module.
import math
from collections import defaultdict
from collections.abc import Callable, Generator, Hashable, Iterable, Mapping, Sequence
from collections.abc import Callable, Hashable, Iterable, Mapping, Sequence
from typing import Generic, Literal, TypeVar, assert_never, overload
from typing_extensions import TypeIs
@@ -102,12 +102,6 @@ def common_prefix(items: Sequence[Sequence[T] | str]) -> Sequence[T] | str:
return shortest
def chunk_list(lst: list[T], chunk_size: int) -> Generator[list[T]]:
"""Yield successive chunk_size chunks from lst."""
for i in range(0, len(lst), chunk_size):
yield lst[i : i + chunk_size]
def flatten_2d_lists(lists: Iterable[Iterable[T]]) -> list[T]:
"""Flatten a list of lists to a single list."""
return [item for sublist in lists for item in sublist]
vllm/utils/cpu_resource_utils.py
-56
View File
@@ -9,7 +9,6 @@ import sys
from dataclasses import dataclass
from functools import cache
import psutil
import regex as re
from vllm.logger import init_logger
@@ -186,61 +185,6 @@ def parse_id_list(raw_str: str) -> list[int]:
return sorted(list(set(result)))
def get_memory_node_info(node_id: int = 0) -> MemoryNodeInfo:
if sys.platform == "darwin":
# MacOS has no memory node
return MemoryNodeInfo(
total_memory=psutil.virtual_memory().total,
available_memory=psutil.virtual_memory().available,
)
meminfo_path = f"/sys/devices/system/node/node{node_id}/meminfo"
if not os.path.exists(meminfo_path):
# Non-NUMA systems (e.g. many RISC-V boards) don't expose per-node
# meminfo. Fall back to system-wide numbers from psutil.
vm = psutil.virtual_memory()
return MemoryNodeInfo(
total_memory=vm.total,
available_memory=vm.available,
)
meminfo = {}
with open(meminfo_path) as f:
for line in f:
# Each line looks like: "Node 0 MemTotal: 97421888 kB"
parts = line.split()
key = parts[2].rstrip(":")
# convert to Bytes
value = int(parts[3]) * 1024
meminfo[key] = value
total_memory = meminfo["MemTotal"]
free_memory = meminfo["MemFree"]
active_file_memory = meminfo["Active(file)"]
inactive_file_memory = meminfo["Inactive(file)"]
reclaimable_memory = meminfo["SReclaimable"]
available_memory = (
free_memory + active_file_memory + inactive_file_memory + reclaimable_memory
)
# Honor cgroup memory limit (containers / k8s pods). NUMA meminfo
# reflects host-wide numbers; without this, gpu_memory_utilization
# would be applied to host RAM instead of the pod's limit. cgroup
# does not expose per-NUMA-node limits, so we just clamp the totals
# against the pod-wide limit here.
cgroup_limit = get_cgroup_memory_limit()
cgroup_usage = get_cgroup_memory_usage()
if cgroup_limit is not None and cgroup_limit < total_memory:
total_memory = cgroup_limit
cgroup_available = cgroup_limit - (cgroup_usage or 0)
available_memory = max(0, min(available_memory, cgroup_available))
return MemoryNodeInfo(
total_memory=total_memory,
available_memory=available_memory,
)
def get_allowed_cpu_list() -> list[LogicalCPUInfo]:
cpu_list = _get_cpu_list()
if sys.platform == "linux":
vllm/utils/deep_gemm.py
-18
View File
@@ -34,12 +34,6 @@ _DEEPGEMM_BLACKWELL_EXCLUDED_MODEL_TYPES: dict[str, tuple[int, ...]] = {
}
def get_paged_mqa_page_sizes() -> tuple[int, ...]:
"""KV page sizes (in cache entries) the paged-MQA logits kernels take (only
64 on SM120); larger storage blocks are split into one of these."""
return (64,) if current_platform.is_device_capability_family(120) else (32, 64)
def should_auto_disable_deep_gemm(model_type: str | None) -> bool:
"""Check if DeepGemm should be auto-disabled for this model on Blackwell.
@@ -799,18 +793,6 @@ def fp8_fp4_paged_mqa_logits(
)
def has_deep_gemm_sparse_mqa() -> bool:
"""Whether the installed DeepGEMM provides the sparse-indexer kernels
(``fp8_fp4_(paged_)sparse_mqa_logits``, added in DeepGEMM 2.8, SM100-only)."""
_lazy_init()
return (
_fp8_fp4_sparse_mqa_logits_impl is not None
and _fp8_fp4_paged_sparse_mqa_logits_impl is not None
and _get_sparse_mqa_logits_metadata_impl is not None
and _get_paged_sparse_mqa_logits_metadata_impl is not None
)
def get_sparse_mqa_logits_metadata(
cu_seqlen_ks: torch.Tensor,
cu_seqlen_ke: torch.Tensor,
vllm/utils/import_utils.py
-10
View File
@@ -552,11 +552,6 @@ def check_moonep_system_support() -> None:
)
def is_numba_available() -> bool:
"""Whether the optional `numba` package is available."""
return _has_module("numba")
def has_triton_kernels() -> bool:
"""Whether the optional `triton_kernels` package is available."""
is_available = _has_module("triton_kernels") or _has_module(
@@ -629,11 +624,6 @@ def has_helion() -> bool:
return _has_module("helion")
def has_aiter() -> bool:
"""Whether the optional `aiter` package is available."""
return _has_module("aiter")
def has_mori() -> bool:
"""Whether the optional `mori` package is available."""
return _has_module("mori")
vllm/utils/math_utils.py
-5
View File
@@ -19,11 +19,6 @@ def next_power_of_2(n: int) -> int:
return 1 if n < 1 else 1 << (n - 1).bit_length()
def prev_power_of_2(n: int) -> int:
"""The previous power of 2 (inclusive)."""
return 0 if n < 1 else 1 << (n.bit_length() - 1)
def round_up(x: int, y: int) -> int:
"""Round up x to the nearest multiple of y."""
return ((x + y - 1) // y) * y
vllm/utils/mem_utils.py
+1 -9
View File
@@ -14,19 +14,11 @@ import torch.types
from vllm.logger import init_logger
from vllm.platforms import current_platform
from .mem_constants import GiB_bytes, KiB_bytes, MiB_bytes
from .mem_constants import GiB_bytes, MiB_bytes
logger = init_logger(__name__)
def format_kib(b: int) -> str:
return f"{round(b / KiB_bytes, 2)}"
def format_mib(b: int) -> str:
return f"{round(b / MiB_bytes, 2)}"
def format_gib(b: int) -> str:
return f"{round(b / GiB_bytes, 2)}"
vllm/utils/multi_stream_utils.py
-141
View File
@@ -1,141 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from collections.abc import Callable
from enum import Enum
from typing import Any
import torch
class EventType(Enum):
Main = 0
Attention = 1
def maybe_execute_in_parallel(
fn0: Callable[[], Any],
fn1: Callable[[], Any],
event0: torch.cuda.Event,
event1: torch.cuda.Event,
aux_stream: torch.cuda.Stream | None = None,
) -> tuple[Any, Any]:
"""Run two functions potentially in parallel on separate CUDA streams.
When aux_stream is provided, fn0 runs on the current (default) stream and
fn1 runs on aux_stream, synchronized via CUDA events. When aux_stream is
None or a breakable CUDA graph capture is active, both functions execute
sequentially on the current stream.
This design follows TensorRT-LLM's maybe_execute_in_parallel pattern
(tensorrt_llm/_torch/modules/multi_stream_utils.py).
Args:
fn0: Callable for the default stream.
fn1: Callable for the auxiliary stream.
event0: CUDA event recorded before fn0 so aux_stream can wait.
event1: CUDA event recorded after fn1 so default stream can wait.
aux_stream: The second CUDA stream for fn1.
Multi-stream is disabled when aux_stream is None or a breakable
CUDA graph capture is active.
Returns:
Tuple of (fn0_result, fn1_result).
"""
if aux_stream is not None:
from vllm.compilation.breakable_cudagraph import BreakableCUDAGraphCapture
if BreakableCUDAGraphCapture.is_active():
aux_stream = None
if aux_stream is not None:
event0.record()
result0 = fn0()
with torch.cuda.stream(aux_stream):
event0.wait()
result1 = fn1()
event1.record()
event1.wait()
else:
result0 = fn0()
result1 = fn1()
return (result0, result1)
def execute_in_parallel(
default_fn: Callable[[], Any],
aux_fns: list[Callable[[], Any] | None],
start_event: torch.cuda.Event,
done_events: list[torch.cuda.Event],
aux_streams: list[torch.cuda.Stream] | None = None,
enable: bool = False,
default_first: bool = False,
) -> tuple[Any, list[Any]]:
"""Run default_fn on the current stream and aux_fns concurrently on
aux_streams.
Generalizes maybe_execute_in_parallel to N aux callables. Slots where
aux_fns[i] is None are skipped (no stream switch, no event record); their
corresponding entry in the returned aux_results list is None.
start_event fans out from the current stream to every launched aux stream;
done_events[i] is recorded after aux_fns[i] so the current stream joins
before returning. Falls back to sequential execution on the current stream
when aux_streams is None or enable is False; in that case default_fn runs
first, then aux_fns in order.
Args:
default_fn: Callable for the default (current) stream.
aux_fns: Per-aux callables; entries may be None to skip.
start_event: CUDA event recorded on the current stream before
default_fn so each launched aux stream can wait on it.
done_events: One CUDA event per aux slot, recorded after the
corresponding aux_fn. Length must match aux_fns.
aux_streams: Per-aux CUDA streams. Length must match aux_fns.
Multi-stream is disabled when None.
enable: Opt-in switch for the multi-stream path. Defaults to False,
so callers that pass aux_streams must also pass enable=True
(typically gated by an env var) to actually overlap. When False,
execution falls back to sequential on the current stream.
default_first: When True, enqueue the default chain before the aux
chains. The default chain is typically the longest, so it starts
first; the aux streams still only wait on start_event, so they
overlap the default chain. Defaults to False.
Returns:
Tuple of (default_result, aux_results) where aux_results[i] is the
result of aux_fns[i] (or None when skipped).
"""
aux_results: list[Any]
if aux_streams is None or not enable:
default_result = default_fn()
aux_results = [fn() if fn is not None else None for fn in aux_fns]
return default_result, aux_results
assert len(aux_fns) == len(aux_streams) == len(done_events), (
"aux_fns, aux_streams, and done_events must be the same length"
)
aux_results = [None] * len(aux_fns)
pending: list[torch.cuda.Event] = []
start_event.record()
if default_first:
default_result = default_fn()
for i, fn in enumerate(aux_fns):
if fn is None:
continue
with torch.cuda.stream(aux_streams[i]):
start_event.wait()
aux_results[i] = fn()
done_events[i].record()
pending.append(done_events[i])
if not default_first:
default_result = default_fn()
for ev in pending:
ev.wait()
return default_result, aux_results
vllm/utils/network_utils.py
-38
View File
@@ -127,25 +127,6 @@ def is_valid_ipv6_address(address: str) -> bool:
return False
def split_host_port(host_port: str) -> tuple[str, int]:
# ipv6
if host_port.startswith("["):
host, port = host_port.rsplit("]", 1)
host = host[1:]
port = port.split(":")[1]
return host, int(port)
else:
host, port = host_port.split(":")
return host, int(port)
def join_host_port(host: str, port: int) -> str:
if is_valid_ipv6_address(host):
return f"[{host}]:{port}"
else:
return f"{host}:{port}"
def get_distributed_init_method(ip: str, port: int) -> str:
return get_tcp_uri(ip, port)
@@ -305,25 +286,6 @@ def split_zmq_path(path: str) -> tuple[str, str, str]:
return scheme, host, port
def make_zmq_path(scheme: str, host: str, port: int | None = None) -> str:
"""Make a ZMQ path from its parts.
Args:
scheme: The ZMQ transport scheme (e.g. tcp, ipc, inproc).
host: The host - can be an IPv4 address, IPv6 address, or hostname.
port: Optional port number, only used for TCP sockets.
Returns:
A properly formatted ZMQ path string.
"""
if port is None:
return f"{scheme}://{host}"
if is_valid_ipv6_address(host):
return f"{scheme}://[{host}]:{port}"
return f"{scheme}://{host}:{port}"
# Adapted from: https://github.com/sgl-project/sglang/blob/v0.4.1/python/sglang/srt/utils.py#L783 # noqa: E501
# Apply the existing make_zmq_socket memory policy to every socket owner.
_ZMQ_LARGE_BUFFER_SIZE = 512 * 1024**2
vllm/utils/torch_utils.py
-131
View File
@@ -97,13 +97,6 @@ def kv_cache_uses_per_token_head_scales(kv_cache_dtype: str) -> bool:
return kv_cache_dtype.endswith("per_token_head")
def is_meta_module(module: torch.nn.Module) -> bool:
"""Return True if module contains any meta parameters or buffers."""
return any(p.is_meta for p in module.parameters()) or any(
b.is_meta for b in module.buffers()
)
def is_strictly_contiguous(t: torch.Tensor) -> bool:
"""Check if tensor is contiguous AND has no degenerate strides.
@@ -133,21 +126,6 @@ def is_strictly_contiguous(t: torch.Tensor) -> bool:
return True
def is_non_overlapping_and_dense(t: torch.Tensor) -> bool:
"""Check if the tensor's elements cover one gapless, non-overlapping byte
range, in any dimension order (i.e. a permuted view of a contiguous
tensor); ``is_contiguous()`` additionally requires row-major order.
"""
expected_stride = 1
for size, stride in sorted(zip(t.shape, t.stride()), key=lambda p: p[1]):
if size == 1:
continue
if stride != expected_stride:
return False
expected_stride *= size
return True
def canonicalize_singleton_dim_strides(t: torch.Tensor) -> torch.Tensor:
"""Fix degenerate strides on size=1 dimensions for CUDA TMA compatibility.
@@ -609,110 +587,6 @@ def nvfp4_split_data_scale(
return data, scale
def create_kv_caches_with_random_flash(
num_blocks: int,
block_size: int,
num_layers: int,
num_heads: int,
head_size: int,
cache_dtype: str | torch.dtype | None,
model_dtype: str | torch.dtype | None = None,
seed: int | None = None,
device: str | None = "cuda",
cache_layout: str | None = "NHD",
) -> tuple[list[torch.Tensor], list[torch.Tensor]]:
set_random_seed(seed)
dtype = get_kv_cache_torch_dtype(cache_dtype, model_dtype)
generic_kv_cache_shape = (num_blocks, 2, block_size, num_heads, head_size)
assert cache_layout in ("NHD", "HND")
stride_order = (0, 1, 2, 3, 4) if cache_layout == "NHD" else (0, 1, 3, 2, 4)
kv_cache_allocation_shape = tuple(generic_kv_cache_shape[i] for i in stride_order)
scale = head_size**-0.5
key_caches: list[torch.Tensor] = []
value_caches: list[torch.Tensor] = []
for _ in range(num_layers):
if isinstance(cache_dtype, str) and cache_dtype.startswith("nvfp4"):
# Full page dim: fp4 data + fp8 block scales per head.
# Per page layout: [K_data | K_scale | V_data | V_scale]
# Returns [:, 0] and [:, 1] like all other dtypes.
full_dim = nvfp4_kv_cache_full_dim(head_size)
nvfp4_shape = (num_blocks, 2, block_size, num_heads, full_dim)
nvfp4_phys = tuple(nvfp4_shape[i] for i in stride_order)
inv = [stride_order.index(i) for i in range(len(stride_order))]
key_value_cache = torch.randint(
0,
256,
nvfp4_phys,
dtype=dtype,
device=device,
).permute(*inv)
else:
key_value_cache = torch.empty(
size=kv_cache_allocation_shape, dtype=dtype, device=device
).permute(*stride_order)
if cache_dtype in ["auto", "half", "bfloat16", "float"]:
key_value_cache.uniform_(-scale, scale)
elif cache_dtype == "fp8":
_generate_random_fp8(key_value_cache, -scale, scale)
else:
raise ValueError(f"Does not support key cache of type {cache_dtype}")
key_caches.append(key_value_cache[:, 0])
value_caches.append(key_value_cache[:, 1])
return key_caches, value_caches
def create_kv_caches_with_random(
num_blocks: int,
block_size: int,
num_layers: int,
num_heads: int,
head_size: int,
cache_dtype: str | torch.dtype | None,
model_dtype: str | torch.dtype | None = None,
seed: int | None = None,
device: str | None = "cuda",
) -> tuple[list[torch.Tensor], list[torch.Tensor]]:
if cache_dtype == "fp8" and head_size % 16:
raise ValueError(
f"Does not support key cache of type fp8 with head_size {head_size}"
)
set_random_seed(seed)
dtype = get_kv_cache_torch_dtype(cache_dtype, model_dtype)
scale = head_size**-0.5
x = 16 // torch.tensor([], dtype=dtype).element_size()
key_cache_shape = (num_blocks, num_heads, head_size // x, block_size, x)
key_caches: list[torch.Tensor] = []
for _ in range(num_layers):
key_cache = torch.empty(size=key_cache_shape, dtype=dtype, device=device)
if cache_dtype in ["auto", "half", "bfloat16", "float"]:
key_cache.uniform_(-scale, scale)
elif cache_dtype == "fp8":
_generate_random_fp8(key_cache, -scale, scale)
else:
raise ValueError(f"Does not support key cache of type {cache_dtype}")
key_caches.append(key_cache)
value_cache_shape = (num_blocks, num_heads, head_size, block_size)
value_caches: list[torch.Tensor] = []
for _ in range(num_layers):
value_cache = torch.empty(size=value_cache_shape, dtype=dtype, device=device)
if cache_dtype in ["auto", "half", "bfloat16", "float"]:
value_cache.uniform_(-scale, scale)
elif cache_dtype == "fp8":
_generate_random_fp8(value_cache, -scale, scale)
else:
raise ValueError(f"Does not support value cache of type {cache_dtype}")
value_caches.append(value_cache)
return key_caches, value_caches
def async_tensor_h2d(
data: list | np.ndarray | torch.Tensor,
device: str | torch.device | None = None,
@@ -1056,11 +930,6 @@ def _encode_layer_name(layer_name: str) -> str | LayerName:
return LayerName(layer_name) if _USE_LAYERNAME else layer_name
# Supports XPU Graph with PyTorch versions >= 2.11.0.dev for XPU platform
def supports_xpu_graph() -> bool:
return is_torch_equal_or_newer("2.11.0.dev")
# create a library to hold the custom op
vllm_lib = Library("vllm", "FRAGMENT") # noqa
vllm/v1/attention/backends/linear_attn.py
-10
View File
@@ -94,16 +94,6 @@ class LinearAttentionMetadataBuilder(AttentionMetadataBuilder[LinearAttentionMet
return attn_metadata
class BailingLinearAttentionBackend(LinearAttentionBackend):
@staticmethod
def get_name() -> str:
return "BAILING_LINEAR_ATTN"
@staticmethod
def get_builder_cls() -> type["BailingLinearAttentionMetadataBuilder"]:
return BailingLinearAttentionMetadataBuilder
@dataclass
class BailingLinearAttentionMetadata(LinearAttentionMetadata):
state_indices_tensor_d: torch.Tensor | None = None
vllm/v1/attention/backends/mamba2_attn.py
-70
View File
@@ -1,6 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import itertools
from dataclasses import dataclass, replace
from typing import Any
@@ -11,7 +10,6 @@ from vllm.model_executor.layers.mamba.checkpoint import (
MambaPrefillCheckpointBuilder,
MambaPrefillCheckpointMetadata,
)
from vllm.utils.torch_utils import async_tensor_h2d
from vllm.v1.attention.backend import (
AttentionBackend,
CommonAttentionMetadata,
@@ -24,74 +22,6 @@ from vllm.v1.attention.backends.mamba_attn import (
from vllm.v1.kv_cache_interface import KVCacheSpec, MambaSpec
def compute_varlen_chunk_metadata(
query_start_loc: torch.Tensor,
chunk_size: int,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
"""Build chunk-aligned, variable-length metadata used by Mamba2 SSD kernels.
Given per-sequence cumulative token starts `query_start_loc` of shape [B+1]
and a physical `chunk_size`, returns three tensors on the same device:
- cu_chunk_seqlens: (nchunks+1,) int32 exclusive prefix-sum of
logical-chunk lengths (each logical chunk never crosses a sequence or
physical-chunk boundary).
- last_chunk_indices: (B,) int32 index of the last logical chunk
for each sequence (=-1 for empty sequences).
- seq_idx_chunks: (nchunks,) int32 sequence index for each logical
chunk in order.
This is intentionally lightweight and CPU-side; it mirrors the metadata
produced by the V1 Mamba2 meta-data builder and is exported so tests
(and other callers) can avoid duplicating the logic.
"""
assert query_start_loc.ndim == 1, "query_start_loc must be 1-D [B+1]"
assert int(query_start_loc[0].item()) == 0, "query_start_loc[0] must be 0"
device = query_start_loc.device
qsl64 = query_start_loc.to(torch.int64)
starts = qsl64[:-1].tolist()
ends = qsl64[1:].tolist()
total = int(qsl64[-1].item())
chunk_lens: list[int] = []
seq_idx_chunks: list[int] = []
last_chunk_indices: list[int] = [-1] * len(starts)
for b, (s, e) in enumerate(zip(starts, ends)):
if e <= s:
# empty sequence
continue
pos = s
while pos < e:
# split at both sequence boundaries and physical chunk boundaries
room = chunk_size - (pos % chunk_size)
take = min(room, e - pos)
chunk_lens.append(int(take))
seq_idx_chunks.append(b)
last_chunk_indices[b] = len(chunk_lens) - 1
pos += take
# Exclusive prefix sum over logical-chunk lengths
if chunk_lens:
cu_chunk_seqlens_list = [0] + list(itertools.accumulate(chunk_lens))
# Final boundary must equal total tokens (check on host to avoid a sync)
assert cu_chunk_seqlens_list[-1] == total
else:
cu_chunk_seqlens_list = [0]
cu_chunk_seqlens = async_tensor_h2d(
cu_chunk_seqlens_list, dtype=torch.int32, device=device
)
# last_chunk_indices is empty when there are no sequences (len(starts) == 0).
last_chunk_indices_t = async_tensor_h2d(
last_chunk_indices, dtype=torch.int32, device=device
)
seq_idx_chunks_t = async_tensor_h2d(
seq_idx_chunks, dtype=torch.int32, device=device
)
return cu_chunk_seqlens, last_chunk_indices_t, seq_idx_chunks_t
class Mamba2AttentionBackend(AttentionBackend):
@staticmethod
def get_name() -> str:
vllm/v1/attention/backends/utils.py
-55
View File
@@ -875,34 +875,6 @@ def split_decodes_and_prefills(
return (num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens)
def split_prefill_chunks(
seq_lens_cpu: torch.Tensor, workspace_size: int, request_offset: int = 0
) -> list[tuple[int, int]]:
"""Split the prefill requests into chunks such that the total sequence length
of each chunk is less than or equal to the workspace size.
Args:
seq_lens_cpu: The sequence lengths of the prefill requests on CPU.
workspace_size: The maximum workspace size (in tokens) per chunk.
request_offset: The offset to add to the request indices.
Returns:
A list of tuples of (reqs_start, reqs_end) representing chunk boundaries.
"""
chunk_bounds = []
i, n = 0, len(seq_lens_cpu)
assert torch.all(seq_lens_cpu <= workspace_size).item()
while i < n:
start, chunk_total = i, 0
while i < n and (chunk_total + (s := seq_lens_cpu[i].item())) <= workspace_size:
chunk_total += s
i += 1
chunk_bounds.append((start + request_offset, i + request_offset))
return chunk_bounds
def reorder_batch_to_split_decodes_and_prefills(
input_batch: "InputBatch",
scheduler_output: "SchedulerOutput",
@@ -983,33 +955,6 @@ def reorder_batch_to_split_decodes_and_prefills(
return True
def reshape_query_for_spec_decode(query: torch.Tensor, batch_size: int) -> torch.Tensor:
"""Reshapes the query tensor for the specified batch size, so that
it has shape (batch_size, seq_len, num_heads, head_dim).
"""
assert query.dim() == 3, f"query must be 3D, got {query.dim()}D"
total_tokens = query.shape[0]
num_heads = query.shape[1]
head_dim = query.shape[2]
assert total_tokens % batch_size == 0, (
f"{total_tokens=} is not divisible by {batch_size=}"
)
seq_len = total_tokens // batch_size
return query.view(batch_size, seq_len, num_heads, head_dim)
def reshape_attn_output_for_spec_decode(attn_output: torch.Tensor) -> torch.Tensor:
"""Reshapes the attention output tensor, so that
the batch_size and seq_len dimensions are combined.
"""
if attn_output.dim() == 3:
# Already in the correct shape
return attn_output
assert attn_output.dim() == 4, f"attn_output must be 4D, got {attn_output.dim()}D"
total_tokens = attn_output.shape[0] * attn_output.shape[1]
return attn_output.view(total_tokens, attn_output.shape[2], attn_output.shape[3])
@runtime_checkable
class KVSharingFastPrefillMetadata(Protocol):
logits_indices_padded: torch.Tensor | None = None
vllm/v1/attention/ops/dcp.py
-207
View File
@@ -2,10 +2,8 @@
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""MLA DCP collective selection and direct symmetric-memory implementations."""
from __future__ import annotations
import functools
from collections.abc import Callable
from dataclasses import dataclass
from typing import TYPE_CHECKING, Any, Protocol
@@ -14,7 +12,6 @@ import torch.distributed as dist
import vllm.envs as envs
from vllm.config import VllmConfig
from vllm.distributed import get_dcp_group
from vllm.logger import init_logger
from vllm.model_executor.warmup.jit_warmup import (
WarmupIntRange,
@@ -1445,208 +1442,4 @@ class DCPCombine(Protocol):
) -> torch.Tensor: ...
class MLADCPManager:
"""Select and own layer-level collective implementations for MLA DCP."""
_kv_gather: Callable[[torch.Tensor, torch.Tensor], object]
def __init__(
self,
vllm_config: VllmConfig,
device: torch.device,
num_heads: int,
query_head_dim: int,
output_head_dim: int,
query_dtype: torch.dtype,
output_dtype: torch.dtype,
padded_num_heads: int | None,
is_lse_base_on_e: bool,
use_pcp: bool,
) -> None:
parallel_config = vllm_config.parallel_config
self.group = get_dcp_group()
self.device = torch.device(device)
self.num_ubatches = max(parallel_config.num_ubatches, 1)
self.max_num_tokens = get_dcp_workspace_max_num_tokens(vllm_config)
self.use_a2a = parallel_config.dcp_comm_backend == "a2a"
self.padded_num_heads = padded_num_heads
self.combine = self._init_combine(
num_heads,
output_head_dim,
output_dtype,
is_lse_base_on_e,
use_pcp,
)
if vllm_config.kernel_config.enable_jit_warmup and not self.use_a2a:
_CORRECT_ATTN_CP_OUT_KERNEL.register_warmup(
vllm_config,
output_dtype=output_dtype,
num_heads=num_heads,
head_dim=output_head_dim,
is_base_e=is_lse_base_on_e,
)
self.query_gather = (
None
if use_pcp
else self._init_query_gather(
num_heads,
query_head_dim,
query_dtype,
)
)
def _init_combine(
self,
num_heads: int,
head_dim: int,
dtype: torch.dtype,
is_lse_base_on_e: bool,
use_pcp: bool,
) -> DCPCombine:
direct_workspace = None
if self.use_a2a:
direct_workspace = get_direct_dcp_a2a_workspace(
self.group,
self.device,
self.max_num_tokens,
num_heads,
head_dim,
dtype,
self.num_ubatches,
)
if direct_workspace is not None:
logger.info_once("Using direct symmetric-memory DCP A2A for MLA.")
return functools.partial(
self._direct_workspace_combine,
direct_workspace,
is_lse_base_on_e=is_lse_base_on_e,
)
combine_fn = (
dcp_a2a_lse_reduce
if self.use_a2a
else cp_lse_ag_out_ar
if use_pcp
else cp_lse_ag_out_rs
)
return functools.partial(
combine_fn,
cp_group=self.group,
is_lse_base_on_e=is_lse_base_on_e,
)
def _direct_workspace_combine(
self,
direct_workspace: DirectDCPA2AWorkspace,
partial_output: torch.Tensor,
partial_lse: torch.Tensor,
is_lse_base_on_e: bool,
seq_lens: torch.Tensor | None = None,
query_start_loc: torch.Tensor | None = None,
) -> torch.Tensor:
# Forced MQA path pass all batch tokens (including prefill) into combine,
# which may exceed the direct symmetric-memory workspace. Fall back to
# the nccl a2a combine for those cases.
if partial_output.shape[0] <= direct_workspace.max_num_tokens:
return direct_workspace.lse_reduce(
partial_output,
partial_lse,
is_lse_base_on_e,
seq_lens,
query_start_loc,
)
return dcp_a2a_lse_reduce(
partial_output,
partial_lse,
cp_group=self.group,
is_lse_base_on_e=is_lse_base_on_e,
seq_lens=seq_lens,
query_start_loc=query_start_loc,
)
def _init_query_gather(
self,
num_heads: int,
head_dim: int,
dtype: torch.dtype,
) -> Callable[[torch.Tensor], torch.Tensor]:
direct_workspace = get_direct_dcp_q_gather_workspace(
self.group,
self.device,
self.max_num_tokens,
num_heads,
head_dim,
dtype,
self.num_ubatches,
self.padded_num_heads,
)
if direct_workspace is not None:
logger.info_once("Using direct symmetric-memory DCP query gather for MLA.")
return functools.partial(
self._direct_workspace_query_gather,
direct_workspace,
)
return self._gather_query
def _gather_query(self, query: torch.Tensor) -> torch.Tensor:
query = self.group.all_gather(query, dim=1)
if self.padded_num_heads is not None:
query = reserve_query_head_storage(query, self.padded_num_heads)
return query
def _direct_workspace_query_gather(
self,
direct_workspace: DirectDCPQGatherWorkspace,
query: torch.Tensor,
) -> torch.Tensor:
# Forced MQA path can be taken for long sparse MLA prefills, whose batch size
# may exceed the direct symmetric-memory workspace. Fall back to allgather
# for those cases.
if query.shape[0] <= direct_workspace.max_num_tokens:
return direct_workspace.gather(query)
return self._gather_query(query)
def init_kv_gather(
self,
workspace: torch.Tensor,
max_gathered_tokens: int,
) -> None:
world_size = self.group.world_size
assert max_gathered_tokens > 0
assert max_gathered_tokens % world_size == 0
assert workspace.ndim == 2
assert workspace.is_contiguous()
assert workspace.shape[0] == (
max_gathered_tokens + max_gathered_tokens // world_size
)
assert workspace.shape[1] > 0
direct_workspace = get_direct_dcp_kv_gather_workspace(
self.group,
workspace.device,
max_gathered_tokens,
workspace.shape[1],
workspace.dtype,
self.num_ubatches,
)
if direct_workspace is not None:
logger.info_once(
"Using direct symmetric-memory DCP chunked-context KV gather for MLA."
)
self._kv_gather = direct_workspace.gather
else:
self._kv_gather = functools.partial(
torch.distributed.all_gather_into_tensor,
group=self.group.device_group,
)
def kv_gather(
self,
gathered_kv: torch.Tensor,
local_kv: torch.Tensor,
) -> object:
return self._kv_gather(gathered_kv, local_kv)
_CORRECT_ATTN_CP_OUT_KERNEL = CorrectAttnCPOutKernel()
vllm/v1/core/kv_cache_utils.py
-13
View File
@@ -130,19 +130,6 @@ def resolve_none_hash_seed(hash_fn: Callable[[Any], bytes]) -> str:
return DEFAULT_NONE_HASH_SEED
def get_none_hash_seed() -> str:
"""Return the seed NONE_HASH was derived from.
Components that must agree on NONE_HASH across processes (the P2P tier
advertises this during its connect handshake) read the resolved seed here
instead of re-deriving it, so they observe the random seed too. Falls back
to the deterministic seed before ``init_none_hash`` has run.
"""
if _NONE_HASH_SEED is None:
return DEFAULT_NONE_HASH_SEED
return _NONE_HASH_SEED
def init_none_hash(hash_fn: Callable[[Any], bytes]):
global NONE_HASH, _NONE_HASH_SEED
vllm/v1/hisparse/layout.py
-31
View File
@@ -14,7 +14,6 @@ from vllm.v1.hisparse.runtime import (
get_hisparse_host_block_stride,
use_shared_hisparse_host_pool,
)
from vllm.v1.hisparse.types import ACTIVE_TAIL_PAGES
from vllm.v1.kv_cache_interface import (
HiSparseHotSpec,
HiSparseResidentSpec,
@@ -375,33 +374,3 @@ def get_hisparse_kv_cache_config(
),
)
def get_hisparse_steady_state_concurrency(
vllm_config: VllmConfig, kv_cache_config: KVCacheConfig
) -> float:
"""Max concurrency at max_model_len once running requests read from host.
A request reading from host pins only its active tail of each resident
group, but admitting one still takes its full in-flight window, so the
bound is all-but-one requests at steady state plus one being admitted.
"""
admission_blocks = 0
steady_blocks = 0
host_blocks = 0
for group in kv_cache_config.kv_cache_groups:
spec = group.kv_cache_spec
required = cdiv(spec.max_memory_usage_bytes(vllm_config), spec.page_size_bytes)
if group.host_resident:
host_blocks += required
continue
admission_blocks += required
if isinstance(spec, HiSparseResidentSpec):
required = min(required, ACTIVE_TAIL_PAGES)
steady_blocks += required
assert kv_cache_config.hisparse_host_num_blocks is not None
num_blocks = kv_cache_config.num_blocks
if num_blocks < admission_blocks:
gpu_concurrency = num_blocks / admission_blocks
else:
gpu_concurrency = 1 + (num_blocks - admission_blocks) / steady_blocks
return min(gpu_concurrency, kv_cache_config.hisparse_host_num_blocks / host_blocks)
vllm/v1/hisparse/runtime.py
+2 -138
View File
@@ -2,16 +2,14 @@
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""NVIDIA HiSparse resident, host, and hot-cache data plane."""
from __future__ import annotations
import math
import mmap
import time
from collections.abc import Callable, Mapping, Sequence
from collections.abc import Callable
from dataclasses import dataclass
from typing import TYPE_CHECKING, Any, TypeAlias
import numpy as np
import psutil
import torch
@@ -22,9 +20,8 @@ from vllm.logger import init_logger
from vllm.platforms import current_platform
from vllm.utils.math_utils import cdiv, round_up
from vllm.utils.mem_utils import get_max_shared_memory_bytes
from vllm.utils.torch_utils import async_tensor_h2d, current_stream
from vllm.utils.torch_utils import current_stream
from vllm.v1.attention.backend import max_decode_query_len
from vllm.v1.hisparse.types import SparseKVResidencyUpdate
from vllm.v1.kv_offload.cpu.shared_offload_region import SharedOffloadRegion
from vllm.v1.simple_kv_offload.cuda_mem_ops import pin_tensor
@@ -33,7 +30,6 @@ logger = init_logger(__name__)
HOST_REGISTER_CHUNK_BYTES = 256 * 2**30
if TYPE_CHECKING:
from vllm.v1.attention.backends.mla.index_group import HiSparseMLAIndexGroup
from vllm.v1.kv_cache_interface import KVCacheConfig
# fp8_ds_mla KV row: 512 B quantized NoPE + 16 B scales + 128 B RoPE.
@@ -588,47 +584,6 @@ class HiSparsePrefillStagingPlan:
self.gpu_source_key = source_key
def build_hisparse_prefill_staging_plan(
block_table: torch.Tensor,
seq_lens: torch.Tensor,
block_size: int,
staging_block_capacity: int,
) -> HiSparsePrefillStagingPlan:
"""Build an asynchronous layer-independent host-cache staging remap."""
device = block_table.device
used = (seq_lens.to(torch.int64) + block_size - 1) // block_size
packed_offsets = torch.arange(staging_block_capacity, device=device)
row_ends = torch.cumsum(used, dim=0)
rows = torch.searchsorted(row_ends, packed_offsets, right=True)
valid = rows < block_table.shape[0]
rows = rows.clamp(max=block_table.shape[0] - 1)
row_starts = torch.cat((torch.zeros_like(row_ends[:1]), row_ends[:-1]))
columns = packed_offsets - row_starts[rows]
valid &= columns < block_table.shape[1]
flat_positions = rows * block_table.shape[1] + columns.clamp(
max=block_table.shape[1] - 1
)
referenced = torch.where(valid, block_table.flatten()[flat_positions], 0)
referenced_bt, row_ids = hisparse_prefill_staging_remap(referenced, block_size)
scratch_index = block_table.numel()
scatter_positions = torch.where(valid, flat_positions, scratch_index)
remapped = torch.where(valid, referenced_bt.flatten(), 0)
new_bt_storage = torch.zeros(scratch_index + 1, dtype=torch.int32, device=device)
new_bt_storage.scatter_(0, scatter_positions, remapped)
new_bt = new_bt_storage[:-1].view_as(block_table)
valid_rows = row_ids >= 0
dst_rows = torch.arange(row_ids.shape[1], dtype=torch.int32, device=device).view(
1, -1
)
return HiSparsePrefillStagingPlan(
block_table=new_bt,
row_ids=row_ids,
dst_rows=dst_rows,
miss_mask=valid_rows.to(torch.int32),
block_size=block_size,
)
def _has_hisparse_ops() -> bool:
if not hasattr(torch.ops, "_C_cache_ops"):
return False
@@ -1135,44 +1090,6 @@ class HiSparseRuntime:
return physical_topk_indices
def update_hisparse_residency(
residency: torch.Tensor,
updates: Mapping[str, SparseKVResidencyUpdate],
request_ids: Sequence[str],
request_state_indices: torch.Tensor,
) -> None:
"""Apply scheduled requests' residency changes.
``residency`` is ``[state rows, resident groups, pages]``: the GPU block of
each page, or block 0 when the page is read from the host.
"""
if not updates:
return
batch_rows = {request_id: row for row, request_id in enumerate(request_ids)}
num_groups = residency.shape[1]
rows: list[int] = []
pages: list[int] = []
block_ids: list[list[int]] = [[] for _ in range(num_groups)]
for request_id, update in updates.items():
rows.append(batch_rows[request_id])
pages.extend(update.pages)
for group_block_ids, update_block_ids in zip(
block_ids, update.block_ids, strict=True
):
group_block_ids.extend(update_block_ids)
host_values = np.empty((len(pages), 2 + num_groups), dtype=np.int32)
host_values[:, 0] = np.repeat(
rows, [len(update.pages) for update in updates.values()]
)
host_values[:, 1] = pages
host_values[:, 2:] = np.asarray(block_ids, dtype=np.int32).T
values = async_tensor_h2d(
host_values, device=request_state_indices.device, dtype=torch.int32
)
state_rows = request_state_indices[values[:, 0]]
residency[state_rows, :, values[:, 1]] = values[:, 2:]
class HiSparseCacheHandle:
"""Attention-facing handle for resident KV and sparse offload state."""
@@ -1314,56 +1231,3 @@ def initialize_hisparse_runtime_buffers(
for cache_handle in cache_handles:
cache_handle.runtime.request_state_indices = request_state_indices
def create_hisparse_cache_handle(
vllm_config: VllmConfig,
model_top_k: int,
*,
is_index_group_leader: bool,
row_width: int,
kv_dtype: torch.dtype,
index_group: HiSparseMLAIndexGroup | None = None,
device: torch.device | str | None = None,
) -> HiSparseCacheHandle | None:
config = ResolvedHiSparseConfig.from_vllm_config(vllm_config, model_top_k)
if config is None:
return None
max_num_reqs = vllm_config.scheduler_config.max_num_seqs
max_swap_rows = _get_max_swap_rows(vllm_config)
if device is None:
device = torch.device(
current_platform.device_type, torch.accelerator.current_device_index()
)
hisparse_group = (
None
if is_index_group_leader or index_group is None
else index_group.hisparse_group
)
runtime = HiSparseRuntime(
config=config,
max_num_reqs=max_num_reqs,
max_swap_rows=max_swap_rows,
row_width=row_width,
kv_dtype=kv_dtype,
device=device,
index_group=hisparse_group,
copy_stream=index_group.side_stream if index_group is not None else None,
logical_topk_ready=(
index_group.logical_topk_ready
if index_group is not None and index_group.has_indexer
else None
),
)
if is_index_group_leader and index_group is not None:
index_group.hisparse_group = runtime.index_group
logger.info_once(
"Enabled experimental HiSparse HMA hot cache: top_k=%d, "
"device_buffer_size=%d (%d LRU rows), max_num_seqs=%d.",
config.top_k,
config.device_buffer_size,
config.device_buffer_size,
max_num_reqs,
)
return HiSparseCacheHandle(runtime)
vllm/v1/pool/late_interaction.py
+1 -21
View File
@@ -6,7 +6,7 @@ from collections.abc import Sequence
import torch
from vllm.pooling_params import LateInteractionParams, PoolingParams
from vllm.pooling_params import PoolingParams
LATE_INTERACTION_MODE_CACHE_QUERY = "cache_query"
LATE_INTERACTION_MODE_SCORE_DOC = "score_doc"
@@ -36,26 +36,6 @@ def get_late_interaction_engine_index(
return zlib.crc32(query_key.encode("utf-8")) % num_engines
def build_late_interaction_query_params(
query_key: str,
query_uses: int,
) -> LateInteractionParams:
return LateInteractionParams(
mode=LATE_INTERACTION_MODE_CACHE_QUERY,
query_key=query_key,
query_uses=max(1, int(query_uses)),
)
def build_late_interaction_doc_params(
query_key: str,
) -> LateInteractionParams:
return LateInteractionParams(
mode=LATE_INTERACTION_MODE_SCORE_DOC,
query_key=query_key,
)
def compute_maxsim_score_batched(
q_embs: Sequence[torch.Tensor],
d_embs: Sequence[torch.Tensor],
vllm/v1/worker/ubatching.py
-11
View File
@@ -157,17 +157,6 @@ def dbo_current_ubatch_id() -> int:
return _THREAD_ID_TO_CONTEXT[threading.get_ident()]
def dbo_current_ubatch_token_offset() -> int:
"""Row offset of the calling ubatch inside the parent batch.
Returns 0 when no DBO context is active on this thread.
"""
if len(_THREAD_ID_TO_CONTEXT) == 0:
return 0
ctx = _CURRENT_CONTEXTS[_THREAD_ID_TO_CONTEXT[threading.get_ident()]]
return ctx.token_offset
def _register_ubatch_function(func):
def wrapper(*args, **kwargs):
if len(_THREAD_ID_TO_CONTEXT) > 0:
vllm/version.py
-6
View File
@@ -29,9 +29,3 @@ def _prev_minor_version_was(version_str):
assert isinstance(__version_tuple__[1], int)
return version_str == f"{__version_tuple__[0]}.{__version_tuple__[1] - 1}"
def _prev_minor_version():
"""For the purpose of testing, return a previous minor version number."""
# In dev tree, this will return "0.-1", but that will work fine"
assert isinstance(__version_tuple__[1], int)
return f"{__version_tuple__[0]}.{__version_tuple__[1] - 1}"
vllm/vllm_flash_attn/flash_attn_interface.py
-175
View File
@@ -582,178 +582,3 @@ def compile_flash_attn_varlen_func_from_specs(
compile_only=True,
)
def sparse_attn_func(
q,
k,
v,
block_count,
block_offset,
column_count,
column_index,
dropout_p=0.0,
softmax_scale=None,
causal=False,
softcap=0.0, # 0.0 means deactivated
alibi_slopes=None,
deterministic=False,
return_attn_probs=False,
*,
return_softmax_lse=False,
out=None,
):
"""Compute attention with vertical and slash sparsity patterns.
Most Arguments are the same with the flash_attn_func interface, except for 4 extra args:
block_count and block_offset for slash sparsity patterns, and
column_count and column_index for vertical sparsity patterns.
For more details please refer to Appendix C.4.2 of paper https://arxiv.org/abs/2407.02490.
Arguments:
q: (batch_size, seqlen, nheads, headdim)
k: (batch_size, seqlen, nheads_k, headdim)
v: (batch_size, seqlen, nheads_k, headdim)
block_count: (batch_size, nheads, cdiv(seqlen, BLOCK_M))
block_offset: (batch_size, nheads, cdiv(seqlen, BLOCK_M), NNZ_S)
column_count: (batch_size, nheads, cdiv(seqlen, BLOCK_M))
column_index: (batch_size, nheads, cdiv(seqlen, BLOCK_M), NNZ_V)
dropout_p: float. Dropout probability.
softmax_scale: float. The scaling of QK^T before applying softmax.
Default to 1 / sqrt(headdim).
causal: bool. Whether to apply causal attention mask (e.g., for auto-regressive modeling).
alibi_slopes: (nheads,) or (batch_size, nheads), fp32. A bias of
(-alibi_slope * |i + seqlen_k - seqlen_q - j|)
is added to the attention score of query i and key j.
deterministic: bool. Whether to use the deterministic implementation of the backward pass,
which is slightly slower and uses more memory. The forward pass is always deterministic.
return_attn_probs: bool. Whether to return the attention probabilities. This option is for
testing only. The returned probabilities are not guaranteed to be correct
(they might not have the right scaling).
softcap: float. Anything > 0 activates softcapping attention.
return_softmax_lse: bool. Whether to also return softmax_lse.
out: optional tensor to write the output into.
Return:
out: (batch_size, seqlen, nheads, headdim).
softmax_lse [optional, if return_softmax_lse=True]: (batch_size, nheads, seqlen). The
logsumexp of each row of the matrix QK^T * scaling (e.g., log of the softmax
normalization factor).
"""
if softmax_scale is None:
softmax_scale = q.shape[-1] ** (-0.5)
q, k, v = [maybe_contiguous(x) for x in (q, k, v)]
out, softmax_lse = torch.ops._vllm_fa2_C.fwd_sparse(
q,
k,
v,
block_count,
block_offset,
column_count,
column_index,
out,
alibi_slopes,
dropout_p,
softmax_scale,
causal,
softcap,
return_attn_probs and dropout_p > 0,
None,
)
return (out, softmax_lse) if return_softmax_lse else out
def sparse_attn_varlen_func(
q,
k,
v,
block_count,
block_offset,
column_count,
column_index,
cu_seqlens_q,
cu_seqlens_k,
max_seqlen_q,
max_seqlen_k,
dropout_p=0.0,
softmax_scale=None,
causal=False,
softcap=0.0, # 0.0 means deactivated
alibi_slopes=None,
deterministic=False,
return_attn_probs=False,
*,
return_softmax_lse=False,
out=None,
):
"""Compute attention with vertical and slash sparsity patterns.
Most Arguments are the same with the flash_attn_varlen_func interface, except for 4 extra args:
block_count and block_offset for slash sparsity patterns, and
column_count and column_index for vertical sparsity patterns.
For more details please refer to Appendix C.4.2 of paper https://arxiv.org/abs/2407.02490.
Arguments:
q: (total_q, nheads, headdim), where total_q = total number of query tokens in the batch.
k: (total_k, nheads_k, headdim), where total_k = total number of key tokens in the batch.
v: (total_k, nheads_k, headdim), where total_k = total number of key tokens in the batch.
block_count: (batch_size, nheads, cdiv(seqlen, BLOCK_M))
block_offset: (batch_size, nheads, cdiv(seqlen, BLOCK_M), NNZ_S)
column_count: (batch_size, nheads, cdiv(seqlen, BLOCK_M))
column_index: (batch_size, nheads, cdiv(seqlen, BLOCK_M), NNZ_V)
cu_seqlens_q: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths
of the sequences in the batch, used to index into q.
cu_seqlens_k: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths
of the sequences in the batch, used to index into kv.
max_seqlen_q: int. Maximum query sequence length in the batch.
max_seqlen_k: int. Maximum key sequence length in the batch.
dropout_p: float. Dropout probability.
softmax_scale: float. The scaling of QK^T before applying softmax.
Default to 1 / sqrt(headdim).
causal: bool. Whether to apply causal attention mask (e.g., for auto-regressive modeling).
softcap: float. Anything > 0 activates softcapping attention.
alibi_slopes: (nheads,) or (batch_size, nheads), fp32. A bias of
(-alibi_slope * |i + seqlen_k - seqlen_q - j|)
is added to the attention score of query i and key j.
deterministic: bool. Whether to use the deterministic implementation of the backward pass,
which is slightly slower and uses more memory. The forward pass is always deterministic.
return_attn_probs: bool. Whether to return the attention probabilities. This option is for
testing only. The returned probabilities are not guaranteed to be correct
(they might not have the right scaling).
return_softmax_lse: bool. Whether to also return softmax_lse.
out: optional tensor to write the output into.
Return:
out: (total, nheads, headdim).
softmax_lse [optional, if return_softmax_lse=True]: (nheads, total_q_seqlen). The
logsumexp of each row of the matrix QK^T * scaling (e.g., log of the softmax
normalization factor).
"""
if softmax_scale is None:
softmax_scale = q.shape[-1] ** (-0.5)
q, k, v = [maybe_contiguous(x) for x in (q, k, v)]
out, softmax_lse = torch.ops._vllm_fa2_C.varlen_fwd_sparse(
q,
k,
v,
block_count,
block_offset,
column_count,
column_index,
out,
cu_seqlens_q,
cu_seqlens_k,
None,
alibi_slopes,
max_seqlen_q,
max_seqlen_k,
dropout_p,
softmax_scale,
False,
causal,
softcap,
return_attn_probs and dropout_p > 0,
None,
)
return (out, softmax_lse) if return_softmax_lse else out