diff --git a/.sync/vllm-sha b/.sync/vllm-sha index 8828e4021a..a78f98fe82 100644 --- a/.sync/vllm-sha +++ b/.sync/vllm-sha @@ -1 +1 @@ -ea0fa34f4992c09fc5aa4ae9a12e67e03125289c +83762b77b07e97c77f986e4bf5a9474952e47bb3 diff --git a/CMakeLists.txt b/CMakeLists.txt index 9c03d2dbfc..d384f83ac5 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -70,6 +70,15 @@ endif() # set(TORCH_SUPPORTED_VERSION_CUDA "2.13.0") set(TORCH_SUPPORTED_VERSION_ROCM "2.13.0") +# TORCH_NIGHTLY=1 builds run against unpinned nightly wheels, so the supported- +# version check would always warn. Only treat it as a nightly build when the +# value is exactly "1" (the bootstrap exports TORCH_NIGHTLY=0 by default, which +# must NOT suppress the warning for normal builds). +if (DEFINED ENV{TORCH_NIGHTLY} AND "$ENV{TORCH_NIGHTLY}" STREQUAL "1") + set(TORCH_NIGHTLY_BUILD TRUE) +else() + set(TORCH_NIGHTLY_BUILD FALSE) +endif() # # Try to find python package with an executable that exactly matches @@ -177,7 +186,7 @@ endif() if (NOT HIP_FOUND AND NOT PYTORCH_FOUND_HIP AND CUDA_FOUND) set(APHRODITE_GPU_LANG "CUDA") - if (NOT Torch_VERSION VERSION_EQUAL ${TORCH_SUPPORTED_VERSION_CUDA}) + if (NOT TORCH_NIGHTLY_BUILD AND NOT Torch_VERSION VERSION_EQUAL ${TORCH_SUPPORTED_VERSION_CUDA}) message(WARNING "Pytorch version ${TORCH_SUPPORTED_VERSION_CUDA} " "expected for CUDA build, saw ${Torch_VERSION} instead.") endif() @@ -190,7 +199,7 @@ elseif(HIP_FOUND OR PYTORCH_FOUND_HIP) enable_language(HIP) # ROCm 5.X and 6.X - if (ROCM_VERSION_DEV_MAJOR GREATER_EQUAL 5 AND + if (NOT TORCH_NIGHTLY_BUILD AND ROCM_VERSION_DEV_MAJOR GREATER_EQUAL 5 AND Torch_VERSION VERSION_LESS ${TORCH_SUPPORTED_VERSION_ROCM}) message(WARNING "Pytorch version >= ${TORCH_SUPPORTED_VERSION_ROCM} " "expected for ROCm build, saw ${Torch_VERSION} instead.") @@ -383,6 +392,7 @@ if(APHRODITE_GPU_LANG STREQUAL "CUDA" OR APHRODITE_GPU_LANG STREQUAL "HIP") "csrc/libtorch_stable/cuda_view.cu" "csrc/libtorch_stable/cuda_utils_kernels.cu" "csrc/libtorch_stable/activation_kernels.cu" + "csrc/libtorch_stable/ngram_embedding_kernels.cu" "csrc/libtorch_stable/quantization/activation_kernels.cu" "csrc/libtorch_stable/quantization/w8a8/int8/scaled_quant.cu" "csrc/libtorch_stable/quantization/w8a8/fp8/common.cu" diff --git a/aphrodite/_custom_ops.py b/aphrodite/_custom_ops.py index 4074370b48..0b167c14f7 100644 --- a/aphrodite/_custom_ops.py +++ b/aphrodite/_custom_ops.py @@ -483,6 +483,37 @@ def rms_norm( torch.ops._C.rms_norm(out, input, weight, epsilon) +# LongCat n-gram embedding index kernel (see csrc/.../ngram_embedding_kernels.cu). +def ngram_compute_n_gram_ids( + ne_n: int, + ne_k: int, + ne_weights: torch.Tensor, + ne_mods: torch.Tensor, + exclusive_ne_embedder_size_sums: torch.Tensor, + exclusive_req_len_sums: torch.Tensor, + ne_token_table: torch.Tensor, + row_indices: torch.Tensor, + column_starts: torch.Tensor, + n_gram_ids: torch.Tensor, +) -> None: + """Compute concatenated (offset) n-gram ids for a ragged prefill batch. + + Writes ``n_gram_ids`` of shape ``[token_num, (ne_n-1)*ne_k]``. + """ + torch.ops._C.ngram_compute_n_gram_ids( + ne_n, + ne_k, + ne_weights, + ne_mods, + exclusive_ne_embedder_size_sums, + exclusive_req_len_sums, + ne_token_table, + row_indices, + column_starts, + n_gram_ids, + ) + + def fused_add_rms_norm( input: torch.Tensor, residual: torch.Tensor, @@ -2757,18 +2788,6 @@ def topk_sigmoid( e_score_correction_bias: torch.Tensor | None = None, routed_scaling_factor: float = 1.0, ) -> None: - if current_platform.is_xpu(): - # xpu doesn't support routed_scaling_factor currently, will revert - # in next vllm-xpu-kernels bumpup - torch.ops._moe_C.topk_sigmoid( - topk_weights, - topk_ids, - token_expert_indices, - gating_output, - renormalize, - e_score_correction_bias, - ) - return torch.ops._moe_C.topk_sigmoid( topk_weights, topk_ids, diff --git a/aphrodite/config/aphrodite.py b/aphrodite/config/aphrodite.py index 8f05efe952..c83c2c2b6e 100644 --- a/aphrodite/config/aphrodite.py +++ b/aphrodite/config/aphrodite.py @@ -70,6 +70,7 @@ "DeepseekV2ForCausalLM", "Qwen2MoeForCausalLM", "GraniteMoeForCausalLM", + "LongcatFlashNgramForCausalLM", } ) diff --git a/aphrodite/config/cache.py b/aphrodite/config/cache.py index 9773fb9c93..baa021e94a 100644 --- a/aphrodite/config/cache.py +++ b/aphrodite/config/cache.py @@ -53,17 +53,12 @@ class CacheConfig: """Whether block_size was explicitly provided. Derived automatically.""" user_specified_mamba_block_size: bool = field(default=False, init=False) """Whether mamba_block_size was explicitly provided. Derived automatically.""" - hash_block_size: int | None = Field(default=None, gt=0) - """Block size (in tokens) used for computing Request's block_hashes. + prefix_match_unit: int | None = Field(default=None, gt=0) + """The finest token boundary (in tokens) a prefix-cache hit can land on. - This can be set to a finer granularity than the physical KV cache block - sizes (e.g. 8) as long as every KV cache group's `block_size` is divisible - by it. This enables prefix-caching keys to be computed at the finest common - granularity and then merged for larger physical block sizes. - - This config is not static default. If left unspecified, Aphrodite will choose a - default based on the resolved KV cache groups (typically the smallest KV - cache block size when there are multiple groups). + Prefix-cache keys are computed every `prefix_match_unit` tokens. It can + be set finer than the physical KV cache block sizes as long as every KV + cache group's `block_size` is divisible by it. """ gpu_memory_utilization: float = Field(default=0.92, gt=0, le=1) """The fraction of GPU memory to be used for the model executor, which can @@ -209,7 +204,7 @@ def compute_hash(self) -> str: "enable_prefix_caching", "prefix_caching_hash_algo", # Prefix-caching implementation detail (doesn't affect compiled graph). - "hash_block_size", + "prefix_match_unit", "mamba_page_size_padded", "skip_page_size_padded", "user_specified_block_size", diff --git a/aphrodite/config/speculative.py b/aphrodite/config/speculative.py index b9595f75c3..c91e4b2186 100644 --- a/aphrodite/config/speculative.py +++ b/aphrodite/config/speculative.py @@ -479,7 +479,7 @@ def hf_config_override(hf_config: PretrainedConfig) -> PretrainedConfig: "architectures": ["Qwen3_5MoeMTP" if is_moe else "Qwen3_5MTP"], } ) - if hf_config.model_type == "longcat_flash": + if hf_config.model_type in ("longcat_flash", "longcat_flash_ngram"): hf_config.model_type = "longcat_flash_mtp" n_predict = getattr(hf_config, "num_nextn_predict_layers", 1) hf_config.update({"n_predict": n_predict, "architectures": ["LongCatFlashMTPModel"]}) @@ -842,6 +842,28 @@ def __post_init__(self): if self.num_speculative_tokens is None: raise ValueError("A speculative model was provided, but `num_speculative_tokens` was not provided") + if self.method == "dspark": + # DSpark is a semi-autoregressive *block* drafter. A + # speculative length smaller than the checkpoint's block + # feeds the block / Markov-head machinery an unsupported + # layout and yields incorrect (garbled) output rather than + # merely lower acceptance. Require num_speculative_tokens to + # be at least the block size (e.g. 5 or 7 for DeepSeek-V4). + dspark_block_size = getattr( + self.draft_model_config.hf_config, + "dspark_block_size", + None, + ) + if dspark_block_size is not None and self.num_speculative_tokens < dspark_block_size: + raise ValueError( + "DSpark requires num_speculative_tokens >= " + f"dspark_block_size ({dspark_block_size}); got " + f"{self.num_speculative_tokens}. Smaller values " + "produce incorrect output. Use " + f"num_speculative_tokens={dspark_block_size} or " + "larger (e.g. 7)." + ) + self.draft_tensor_parallel_size = SpeculativeConfig._verify_and_get_draft_tp( self.target_parallel_config, self.draft_tensor_parallel_size, diff --git a/aphrodite/distributed/kv_events.py b/aphrodite/distributed/kv_events.py index 1a287715cc..e454ce6dc3 100644 --- a/aphrodite/distributed/kv_events.py +++ b/aphrodite/distributed/kv_events.py @@ -44,6 +44,8 @@ class KVCacheEvent( MEDIUM_GPU = "GPU" MEDIUM_CPU = "CPU" +MEDIUM_FS = "FS" +MEDIUM_OBJ = "OBJ" class BlockStored(KVCacheEvent): diff --git a/aphrodite/distributed/kv_transfer/kv_connector/utils.py b/aphrodite/distributed/kv_transfer/kv_connector/utils.py index 9ba73f536c..90b1b2a231 100644 --- a/aphrodite/distributed/kv_transfer/kv_connector/utils.py +++ b/aphrodite/distributed/kv_transfer/kv_connector/utils.py @@ -10,7 +10,12 @@ import torch -from aphrodite.config import AphroditeConfig, get_current_aphrodite_config, get_layers_from_aphrodite_config +from aphrodite.config import ( + AphroditeConfig, + get_current_aphrodite_config, + get_layers_from_aphrodite_config, + set_current_aphrodite_config, +) from aphrodite.distributed.kv_transfer.kv_connector.factory import KVConnectorFactory from aphrodite.logger import init_logger from aphrodite.model_executor.layers.attention_layer_base import AttentionLayerBase @@ -326,14 +331,15 @@ def get_current_attn_backends( logger.debug("No layers found in the Aphrodite config. Falling back to default attention backend.") from aphrodite.v1.attention.selector import get_attn_backend - return [ - get_attn_backend( - head_size=aphrodite_config.model_config.get_head_size(), - dtype=aphrodite_config.model_config.dtype, - kv_cache_dtype=aphrodite_config.cache_config.cache_dtype, - use_mla=aphrodite_config.model_config.use_mla, - ) - ] + with set_current_aphrodite_config(aphrodite_config): + return [ + get_attn_backend( + head_size=aphrodite_config.model_config.get_head_size(), + dtype=aphrodite_config.model_config.dtype, + kv_cache_dtype=aphrodite_config.cache_config.cache_dtype, + use_mla=aphrodite_config.model_config.use_mla, + ) + ] def get_current_attn_backend( @@ -396,8 +402,6 @@ def __post_init__(self): self._engines: dict[tuple[EngineId, int], EngineTransferInfo] = {} - # Figure out whether the first dimension of the cache is K/V - # or num_blocks. attn_backend = self.attn_backends[0] if not self.is_mamba: _MOCK_BLOCK_SIZE = 16 @@ -408,10 +412,16 @@ def __post_init__(self): head_size=1, ) logger.debug("Test kv_cache_shape: %s", kv_cache_shape) - # Non-MLA backends caches have 5 dims [num_blocks, 2, H,N,D], - # we just mock num_blocks to 1 for the dimension check below. - # Hybrid SSM models assume a single blocks_first layout - self._is_kv_layout_blocks_first = self.is_mamba or (len(kv_cache_shape) == 5 and kv_cache_shape[0] == 1) + assert kv_cache_shape[0] == 1, ( + "KV cache layout must be blocks-first; expected mocked " + f"num_blocks=1 in leading dim, got shape {kv_cache_shape}." + ) + if not self.is_mla: + assert len(kv_cache_shape) == 4, ( + "Attention KV cache layout must be standardized as " + "[num_blocks, num_kv_heads, block_size, content_size], " + f"got shape {kv_cache_shape}." + ) self._cross_layers_blocks = False if self.tensor_shape is not None: @@ -465,27 +475,17 @@ def unregister_remote_engine(self, remote_engine_id: EngineId) -> None: # Layout properties # ============================================================ - @property - def is_kv_layout_blocks_first(self) -> bool: - return self._is_kv_layout_blocks_first - @property def cross_layers_blocks(self) -> bool: return self._cross_layers_blocks @property def virtually_split_kv_in_blocks(self) -> bool: - # Whether to logically split each block into K and V halves. - # Applies when K/V are interleaved within each block (blocks-first), - # but NOT when cross-layer blocks are used — cross-layer blocks have - # per-layer K/V interleaving (L0_K, L0_V, L1_K, L1_V, ...) so a - # simple half-split does not separate K from V. - return self._is_kv_layout_blocks_first and not self._cross_layers_blocks - - @property - def split_k_and_v(self) -> bool: - # Whether to register regions for K and V separately (when present). - return not (self._cross_layers_blocks or self.is_mla or self.is_kv_layout_blocks_first) + # Whether to logically split each block into two separately-indexable + # sub-regions. With K and V packed into the content dim, an attention + # block transfers as a single unit. Only Mamba still needs this, to + # index its two state regions separately. + return self.is_mamba and not self._cross_layers_blocks # ============================================================ # Common methods @@ -580,8 +580,9 @@ def get_transfer_cache_regions( # Swap [2<>num_blocks] dims for hybrid SSM layout. cache = cache.transpose(0, 1) - # Regular case: backends like FA register K/V in separate regions - return cache if self.split_k_and_v else [cache] + # K and V are packed into one tensor (content dim), so each layer + # registers as a single region. + return [cache] def describe(self, remote_engine_id: EngineId, remote_pp_rank: int = 0) -> str: """One-line summary of transfer config for logging.""" diff --git a/aphrodite/distributed/kv_transfer/kv_connector/v1/mooncake/store/coordinator.py b/aphrodite/distributed/kv_transfer/kv_connector/v1/mooncake/store/coordinator.py index 697bafe9a5..1249d3a802 100644 --- a/aphrodite/distributed/kv_transfer/kv_connector/v1/mooncake/store/coordinator.py +++ b/aphrodite/distributed/kv_transfer/kv_connector/v1/mooncake/store/coordinator.py @@ -29,10 +29,15 @@ class ExternalCachedBlockPool: """Duck-typed BlockPool backed by a ``(group_id, hash)`` exists set.""" - def __init__(self, exists: set[tuple[int, bytes]] | None = None) -> None: + def __init__( + self, + hash_block_size: int, + exists: set[tuple[int, bytes]] | None = None, + ) -> None: # ``exists=None`` is used on the recv side where hit_length is already # determined and we just want each spec's manager to apply its own mask. self._exists = exists + self.hash_block_size = hash_block_size self.null_block = KVCacheBlock(block_id=0) # Dummy ID 1 for present block for duck-typing. self._present_block = KVCacheBlock(block_id=1) @@ -155,7 +160,7 @@ def load_mask( masks, _ = self.find_longest_cache_hit( block_hashes, token_len, - ExternalCachedBlockPool(), + ExternalCachedBlockPool(self.hash_block_size), apply_eagle=False, ) return masks @@ -253,9 +258,8 @@ def _find_hit_blocks( eagle_indices = self.eagle_attn_group_indices if apply_eagle else set() if len(self.attention_groups) == 1: spec, group_ids, manager_cls = self.attention_groups[0] - hashes = self.block_hashes_for_spec(block_hashes, spec) - hit_blocks = manager_cls.find_longest_cache_hit( - block_hashes=hashes, # type: ignore[arg-type] + hit_blocks, hit_length = manager_cls.find_longest_cache_hit( + block_hashes=block_hashes, # type: ignore[arg-type] max_length=max_length, kv_cache_group_ids=group_ids, block_pool=cast(BlockPool, cached_block_pool), @@ -267,11 +271,12 @@ def _find_hit_blocks( blocks_by_group: list[list[KVCacheBlock]] = [[] for _ in range(num_groups)] for gid, blks in zip(group_ids, hit_blocks, strict=True): blocks_by_group[gid] = blks - return tuple(blocks_by_group), len(hit_blocks[0]) * spec.block_size + return tuple(blocks_by_group), hit_length num_groups = len(self.kv_cache_groups) hit_length = max_length hit_blocks_by_group: list[list[KVCacheBlock] | None] = [None] * num_groups + hit_length_by_group: list[int] = [0] * num_groups is_simple_hybrid = len(self.attention_groups) == 2 and isinstance( self.attention_groups[0][0], FullAttentionSpec @@ -282,18 +287,18 @@ def _find_hit_blocks( curr_hit_length = hit_length for idx, (spec, group_ids, manager_cls) in enumerate(self.attention_groups): - cached = hit_blocks_by_group[group_ids[0]] + first_group_id = group_ids[0] + cached = hit_blocks_by_group[first_group_id] if isinstance(spec, FullAttentionSpec) and cached is not None: - curr_hit_length = curr_hit_length // spec.block_size * spec.block_size + curr_hit_length = min(curr_hit_length, hit_length_by_group[first_group_id]) continue drop_eagle_block = idx in eagle_indices and idx not in eagle_verified _max_length = curr_hit_length if drop_eagle_block: _max_length = min(curr_hit_length + spec.block_size, max_length) - hashes = self.block_hashes_for_spec(block_hashes, spec) - hit_blocks = manager_cls.find_longest_cache_hit( - block_hashes=hashes, # type: ignore[arg-type] + hit_blocks, _new_hit_length = manager_cls.find_longest_cache_hit( + block_hashes=block_hashes, # type: ignore[arg-type] max_length=_max_length, kv_cache_group_ids=group_ids, block_pool=cast(BlockPool, cached_block_pool), @@ -301,7 +306,6 @@ def _find_hit_blocks( drop_eagle_block=drop_eagle_block, alignment_tokens=self.lcm_block_size, ) - _new_hit_length = len(hit_blocks[0]) * spec.block_size if drop_eagle_block: eagle_verified.add(idx) elif _new_hit_length < curr_hit_length: @@ -309,6 +313,7 @@ def _find_hit_blocks( curr_hit_length = _new_hit_length for gid, blocks in zip(group_ids, hit_blocks, strict=True): hit_blocks_by_group[gid] = blocks + hit_length_by_group[gid] = _new_hit_length if curr_hit_length >= hit_length: break @@ -325,6 +330,7 @@ def _find_hit_blocks( full_blks = hit_blocks_by_group[gid] assert full_blks is not None del full_blks[num_blocks:] + hit_length_by_group[gid] = hit_length return ( tuple(blks if blks is not None else [] for blks in hit_blocks_by_group), diff --git a/aphrodite/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py b/aphrodite/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py index c4cae2f6af..dde53ba0b1 100644 --- a/aphrodite/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py +++ b/aphrodite/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py @@ -1434,7 +1434,9 @@ def lookup(self, token_len: int, block_hashes: Sequence[BlockHash]) -> int: } _masks, hit_length = self.coord.find_longest_cache_hit( - block_hashes, token_len, ExternalCachedBlockPool(exists_set) + block_hashes, + token_len, + ExternalCachedBlockPool(self.hash_block_size, exists_set), ) return hit_length diff --git a/aphrodite/distributed/kv_transfer/kv_connector/v1/moriio/moriio_layout.py b/aphrodite/distributed/kv_transfer/kv_connector/v1/moriio/moriio_layout.py index 835f4b0d73..d021512f73 100644 --- a/aphrodite/distributed/kv_transfer/kv_connector/v1/moriio/moriio_layout.py +++ b/aphrodite/distributed/kv_transfer/kv_connector/v1/moriio/moriio_layout.py @@ -7,6 +7,7 @@ import torch from aphrodite.v1.kv_cache_interface import ( + AttentionSpec, KVCacheConfig, KVCacheSpec, MLAAttentionSpec, @@ -49,6 +50,11 @@ def is_mla_cache_layer(layer_to_spec: Mapping[str, KVCacheSpec], layer_name: str return isinstance(spec, (MLAAttentionSpec, SlidingWindowMLASpec)) +def _content_packed_dim(spec: AttentionSpec) -> int: + head_size_v = getattr(spec, "head_size_v", spec.head_size) + return spec.head_size + head_size_v + + def _spec_dim_matches(value: int, expected: int | None) -> bool: return expected is None or value == expected @@ -203,6 +209,30 @@ def get_layer_transfer_geometry( split_kv_regions=False, ) + if ( + not is_mla_cache + and isinstance(spec, AttentionSpec) + and len(shape) == 4 + and shape[1] == spec.num_kv_heads + and shape[2] == spec.block_size + and shape[3] == _content_packed_dim(spec) + ): + num_blocks, num_kv_heads, block_size, packed_dim = shape + slot_size_bytes = num_kv_heads * packed_dim * element_size + block_len = block_size * slot_size_bytes + return LayerTransferGeometry( + num_blocks=num_blocks, + block_size=block_size, + block_len=block_len, + slot_size_bytes=slot_size_bytes, + block_stride=stride[0], + local_kv_stride=None, + remote_kv_stride=None, + transfers_per_block=1, + regions_per_block=1, + split_kv_regions=False, + ) + cache_kind = "MLA" if is_mla_cache else "K/V" raise ValueError(f"Unsupported MoRIIO {cache_kind} cache shape for layer {layer_name}: {tuple(shape)}") diff --git a/aphrodite/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py b/aphrodite/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py index f965702abe..63046aa25d 100644 --- a/aphrodite/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py +++ b/aphrodite/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py @@ -200,12 +200,10 @@ def _fa_desc_replicated(self, num_fa_descs: int) -> list[bool]: if n_regions == 0 or self.num_regions == 0: return [False] * num_fa_descs nblk = num_fa_descs // self.num_regions - virtually_split = self.transfer_topo.virtually_split_kv_in_blocks flags: list[bool] = [] for i in range(n_regions): replicated = self._is_region_replicated(i) - num_streams = 1 if replicated or not virtually_split else 2 - flags.extend([replicated] * (num_streams * nblk)) + flags.extend([replicated] * nblk) assert len(flags) == num_fa_descs, f"FA desc flags {len(flags)} != num_fa_descs {num_fa_descs}" return flags @@ -630,27 +628,26 @@ def initialize_host_xfer_buffer(self, kv_caches: dict[str, torch.Tensor]) -> Non NOT directly supported by NIXL (e.g., tpu) """ xfer_buffers: dict[str, torch.Tensor] = {} - inv_order = [0, 1, 3, 2, 4] try: for layer_name, kv_cache in kv_caches.items(): kv_shape = kv_cache.shape kv_dtype = kv_cache.dtype permute_shape = False - if ( - self.kv_cache_layout == "NHD" - and self.aphrodite_config.kv_transfer_config is not None - and self.aphrodite_config.kv_transfer_config.enable_permute_local_kv - ): - logger.info_once( - "'enable_permute_local_kv' flag is enabled while " - "device KV Layout is NHD. Init host buffer with" - " HND to better support Decode/Prefill TP_ratio > 1." - ) - # Since NHD will not support Decode/Prefill TP_ratio > 1, - # we can leverage host_buffer for permute - self.host_buffer_kv_cache_layout = "HND" - kv_shape = tuple(kv_shape[i] for i in inv_order) if not self.use_mla else kv_shape - permute_shape = not self.use_mla + inv_order = (0, 2, 1, 3) + if not self.use_mla: + assert kv_cache.ndim == 4 + + if self.kv_cache_layout == "NHD": + if self.kv_transfer_config.enable_permute_local_kv: + logger.info_once( + "'enable_permute_local_kv' flag is enabled while " + "device KV Layout is NHD. Init host buffer with" + " HND to better support Decode/Prefill TP_ratio > 1." + ) + self.host_buffer_kv_cache_layout = "HND" + else: + kv_shape = tuple(kv_shape[i] for i in inv_order) + permute_shape = True xfer_buffers[layer_name] = torch.empty(kv_shape, dtype=kv_dtype, device="cpu") if permute_shape: @@ -1027,9 +1024,7 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]): f"backend={self.backend_name}, " "all_backends=" f"{[backend.get_name() for backend in self.attn_backends]}, " - f"kv_cache_layout={self.kv_cache_layout}, " - "blocks_first=" - f"{self.transfer_topo.is_kv_layout_blocks_first}" + f"kv_cache_layout={self.kv_cache_layout}" ) # Need to make sure the device ID is non-negative for NIXL, @@ -1052,19 +1047,6 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]): regions_per_layer = self.num_regions // num_local_layers self._remote_region_offset = regions_per_layer * start_layer - if self.transfer_topo.virtually_split_kv_in_blocks: - # NOTE (NickLucche) When FlashInfer is used, memory is registered - # with joint KV for each block. This minimizes the overhead in - # registerMem allowing faster descs queries. In order to be able to - # split on kv_heads dim as required by heterogeneous TP, one must - # be able to index K/V separately. Hence we double the number - # of 'virtual' regions here and halve `block_len` below. - # Similarly for Mamba layers, we register SSM+Conv as a single region and - # then duplicate it logically to be able to index SSM/Conv separately. - # Exception: key-only REPLICATE regions (MLA) have no V half, so - # they contribute a single desc stream and are not doubled. - self.num_regions = sum(1 if self._is_region_replicated(i) else 2 for i in range(len(self._region_is_mla))) - # Total local FA descriptors (boundary between FA and mamba descs). self.num_descs = self.num_regions * self.num_blocks @@ -1212,16 +1194,6 @@ def _build_fa_local( addr = base_addr + block_offset result.append((addr, kv_block_len, self.device_id)) - if self.transfer_topo.virtually_split_kv_in_blocks and not self._is_region_replicated(i): - # Separate and interleave K/V regions to maintain the same - # descs ordering. This is needed for selecting contiguous heads - # when split across TP ranks. (Skipped for key-only REPLICATE.) - second_split = self.get_backend_aware_kv_block_len(layer_idx=i, first_split=False, mamba_view=False) - for block_id in range(num_blocks): - block_offset = block_id * page_stride - addr = base_addr + block_offset - v_addr = addr + kv_block_len - result.append((v_addr, second_split, self.device_id)) return result def _build_fa_remote( @@ -1261,17 +1233,6 @@ def _build_fa_remote( addr = base_addr + block_offset + rank_offset result.append((addr, local_block_len, nixl_agent_meta.device_id)) - emits_v = self.transfer_topo.virtually_split_kv_in_blocks and not replicated - if emits_v: - # With FlashInfer index V separately to allow head splitting. - second_split = self.get_backend_aware_kv_block_len(layer_idx=i, first_split=False, mamba_view=False) - second_split = second_split // num_reads - for block_id in range(num_blocks): - block_offset = block_id * page_size - addr = base_addr + block_offset + rank_offset - # Hop over the first split of remote page, K, to read V. - v_addr = addr + nixl_agent_meta.block_lens[i] // 2 - result.append((v_addr, second_split, nixl_agent_meta.device_id)) return result def register_local_xfer_handler( @@ -1697,20 +1658,16 @@ def post_process_device_kv_on_receive( block_size_ratio, ) - split_k_and_v = self.transfer_topo.split_k_and_v - for block_ids in block_ids_list: indices = torch.tensor(block_ids, device=self.device_type, dtype=torch.long) - for _, cache_or_caches in self.device_kv_caches.items(): - cache_list = cache_or_caches if split_k_and_v else [cache_or_caches] - for cache in cache_list: - if self.enable_permute_local_kv and block_size_ratio > 1: - kv_postprocess_blksize_and_layout_on_receive(cache, indices, block_size_ratio) - elif self.enable_permute_local_kv: - kv_postprocess_layout_on_receive(cache, indices) - else: - kv_postprocess_blksize_on_receive(cache, indices, block_size_ratio) + for cache in self.device_kv_caches.values(): + if self.enable_permute_local_kv and block_size_ratio > 1: + kv_postprocess_blksize_and_layout_on_receive(cache, indices, block_size_ratio) + elif self.enable_permute_local_kv: + kv_postprocess_layout_on_receive(cache, indices) + else: + kv_postprocess_blksize_on_receive(cache, indices, block_size_ratio) def post_process_device_kv_on_receive_heterogeneous_attn(self, block_ids: list[int]): """ @@ -2139,12 +2096,10 @@ def get_backend_aware_kv_block_len(self, layer_idx: int, first_split: bool = Tru |1st_split-2nd_split| |1st_split-2nd_split | """ assert self.transfer_topo is not None - virtually_split = self.transfer_topo.virtually_split_kv_in_blocks - if virtually_split and mamba_view: + if self.transfer_topo.virtually_split_kv_in_blocks and mamba_view: block_len = self._mamba_ssm_size[not first_split] else: - half_block = virtually_split and not self._is_region_replicated(layer_idx) - block_len = self.block_len_per_layer[layer_idx] // (2 if half_block else 1) + block_len = self.block_len_per_layer[layer_idx] return block_len def get_kv_connector_stats(self) -> KVConnectorStats | None: diff --git a/aphrodite/distributed/weight_transfer/base.py b/aphrodite/distributed/weight_transfer/base.py index 5032e398eb..e5827047e2 100644 --- a/aphrodite/distributed/weight_transfer/base.py +++ b/aphrodite/distributed/weight_transfer/base.py @@ -72,6 +72,7 @@ class WeightTransferEngine(ABC, Generic[TInitInfo, TUpdateInfo]): # Subclasses should override these class attributes init_info_cls: type[TInitInfo] update_info_cls: type[TUpdateInfo] + supports_draft_weight_update: bool = True def __init__( self, @@ -95,6 +96,22 @@ def __init__( self.model_config = aphrodite_config.model_config self.device = device self.model = model + self._default_model_config = self.model_config + self._default_model = model + + def set_weight_update_target( + self, + model: torch.nn.Module, + model_config: Any, + ) -> None: + """Set the model that will receive the active weight update.""" + self.model = model + self.model_config = model_config + + def reset_weight_update_target(self) -> None: + """Restore weight updates to the engine's default target model.""" + self.model = self._default_model + self.model_config = self._default_model_config def parse_init_info(self, init_dict: dict[str, Any]) -> TInitInfo: """ diff --git a/aphrodite/distributed/weight_transfer/sparse_nccl_engine.py b/aphrodite/distributed/weight_transfer/sparse_nccl_engine.py index 66bb59cf69..9251459a5d 100644 --- a/aphrodite/distributed/weight_transfer/sparse_nccl_engine.py +++ b/aphrodite/distributed/weight_transfer/sparse_nccl_engine.py @@ -98,6 +98,7 @@ class SparseNCCLWeightTransferEngine( init_info_cls = NCCLWeightTransferInitInfo update_info_cls = SparseNCCLWeightTransferUpdateInfo + supports_draft_weight_update = False def __init__( self, diff --git a/aphrodite/engine/arg_utils.py b/aphrodite/engine/arg_utils.py index abbdd41bb5..5ceb134ad4 100644 --- a/aphrodite/engine/arg_utils.py +++ b/aphrodite/engine/arg_utils.py @@ -629,6 +629,7 @@ class EngineArgs: mamba_cache_dtype: MambaDType = CacheConfig.mamba_cache_dtype mamba_ssm_cache_dtype: MambaDType = CacheConfig.mamba_ssm_cache_dtype mamba_block_size: int | None = get_field(CacheConfig, "mamba_block_size") + prefix_match_unit: int | None = get_field(CacheConfig, "prefix_match_unit") mamba_cache_mode: MambaCacheMode = CacheConfig.mamba_cache_mode mamba_backend: MambaBackendEnum = MambaBackendEnum.TRITON @@ -1044,6 +1045,7 @@ def add_cli_args(parser: FlexibleArgumentParser) -> FlexibleArgumentParser: cache_group.add_argument("--mamba-cache-dtype", **cache_kwargs["mamba_cache_dtype"]) cache_group.add_argument("--mamba-ssm-cache-dtype", **cache_kwargs["mamba_ssm_cache_dtype"]) cache_group.add_argument("--mamba-block-size", **cache_kwargs["mamba_block_size"]) + cache_group.add_argument("--prefix-match-unit", **cache_kwargs["prefix_match_unit"]) cache_group.add_argument("--mamba-cache-mode", **cache_kwargs["mamba_cache_mode"]) cache_group.add_argument("--kv-offloading-size", **cache_kwargs["kv_offloading_size"]) cache_group.add_argument("--kv-offloading-backend", **cache_kwargs["kv_offloading_backend"]) @@ -1641,6 +1643,7 @@ def create_engine_config( mamba_cache_dtype=self.mamba_cache_dtype, mamba_ssm_cache_dtype=self.mamba_ssm_cache_dtype, mamba_block_size=self.mamba_block_size, + prefix_match_unit=self.prefix_match_unit, mamba_cache_mode=self.mamba_cache_mode, kv_offloading_size=self.kv_offloading_size, kv_offloading_backend=self.kv_offloading_backend, diff --git a/aphrodite/engine/protocol.py b/aphrodite/engine/protocol.py index f94fe36b0e..432571aaeb 100644 --- a/aphrodite/engine/protocol.py +++ b/aphrodite/engine/protocol.py @@ -239,6 +239,10 @@ async def start_weight_update(self) -> None: """Start a new weight update.""" raise NotImplementedError + async def start_draft_weight_update(self) -> None: + """Start a new weight update targeting the speculative draft model.""" + raise NotImplementedError + async def update_weights(self, request: WeightTransferUpdateRequest) -> None: """Batched weight update for RL training.""" raise NotImplementedError diff --git a/aphrodite/entrypoints/chat_utils.py b/aphrodite/entrypoints/chat_utils.py index ff489f6491..e6b783f265 100644 --- a/aphrodite/entrypoints/chat_utils.py +++ b/aphrodite/entrypoints/chat_utils.py @@ -775,10 +775,17 @@ async def resolve_items( if not self._items_by_modality: return None, None - resolved_items_by_modality = { - modality: await asyncio.gather(*(item() for item in items)) - for modality, items in self._items_by_modality.items() - } + resolved_items_by_modality: dict[str, list[Any]] = {} + for modality, items in self._items_by_modality.items(): + results = await asyncio.gather(*(item() for item in items), return_exceptions=True) + for result in results: + if isinstance(result, BaseException): + # Gathering with return_exceptions=True lets every task in + # this modality finish (or itself fail) before we raise, + # instead of abandoning still-in-flight fetches (real + # network/thread-pool work) the moment the first one fails. + raise result + resolved_items_by_modality[modality] = results mm_processor = self.mm_processor if self._model_config.is_multimodal_model else None return _resolve_items( diff --git a/aphrodite/entrypoints/llm.py b/aphrodite/entrypoints/llm.py index 88e629b868..a82813c33f 100644 --- a/aphrodite/entrypoints/llm.py +++ b/aphrodite/entrypoints/llm.py @@ -862,6 +862,10 @@ def finish_weight_update(self) -> None: """Finish the current weight update.""" self.llm_engine.collective_rpc("finish_weight_update") + def start_draft_weight_update(self) -> None: + """Start a new weight update targeting the speculative draft model.""" + self.llm_engine.collective_rpc("start_draft_weight_update") + def __repr__(self) -> str: """Return a transformers-style hierarchical view of the model.""" # Cache the result to avoid repeated collective_rpc calls diff --git a/aphrodite/entrypoints/openai/chat_completion/serving.py b/aphrodite/entrypoints/openai/chat_completion/serving.py index 1cafaa1421..bcc9c95f71 100644 --- a/aphrodite/entrypoints/openai/chat_completion/serving.py +++ b/aphrodite/entrypoints/openai/chat_completion/serving.py @@ -564,14 +564,8 @@ async def chat_completion_stream_generator( prompt_token_ids=res.prompt_token_ids, finished=output.finish_reason is not None, ) - if delta_message is not None: - if delta_message.tool_calls: - tools_streamed[i] = True - - if delta_message.reasoning and not request.include_reasoning: - delta_message.reasoning = None - if not (delta_message.content or delta_message.tool_calls): - delta_message = None + if delta_message is not None and delta_message.tool_calls: + tools_streamed[i] = True # handle streaming just a content delta (no parsers) else: @@ -586,11 +580,19 @@ async def chat_completion_stream_generator( # "control token" for tool calls or the parser otherwise # wasn't ready to send a token, then # get the next token without streaming a chunk + # When reasoning is hidden, suppress per-token + # metadata (logprobs, token_ids) on every chunk to + # prevent leaking reasoning tokens through decoded + # token text in logprob entries or raw token IDs. + hide_stream_metadata = not request.include_reasoning and parser is not None + if hide_stream_metadata: + logprobs = None + if delta_message is None: # NOTE: If return_token_ids is enabled, we still need to # send a chunk with token_ids even if delta_message is None # to ensure all tokens are included in the response - if output.finish_reason is None and not request.return_token_ids: + if output.finish_reason is None and (not request.return_token_ids or hide_stream_metadata): continue delta_message = DeltaMessage() @@ -622,6 +624,8 @@ async def chat_completion_stream_generator( delta=True, ) + include_token_ids = request.return_token_ids and not hide_stream_metadata + if output.finish_reason is None: # Send token-by-token response for each request.n choice_data = ChatCompletionResponseStreamChoice( @@ -629,7 +633,7 @@ async def chat_completion_stream_generator( delta=delta_message, logprobs=logprobs, finish_reason=None, - token_ids=(as_list(output.token_ids) if request.return_token_ids else None), + token_ids=(as_list(output.token_ids) if include_token_ids else None), ) # if the model is finished generating @@ -653,7 +657,7 @@ async def chat_completion_stream_generator( logprobs=logprobs, finish_reason=finish_reason_, stop_reason=output.stop_reason, - token_ids=(as_list(output.token_ids) if request.return_token_ids else None), + token_ids=(as_list(output.token_ids) if include_token_ids else None), ) finish_reason_sent[i] = True @@ -817,12 +821,15 @@ async def chat_completion_full_generator( enable_auto_tools=self.enable_auto_tools, model_output_token_ids=token_ids, ) + suppress_metadata = not request.include_reasoning if not request.include_reasoning: reasoning = None + logprobs = None else: reasoning = None content = output.text tool_calls = [] + suppress_metadata = False auto_tools_called = False is_named_tool_choice = ( @@ -906,7 +913,7 @@ async def chat_completion_full_generator( if output.finish_reason else "stop", stop_reason=output.stop_reason, - token_ids=(as_list(output.token_ids) if request.return_token_ids else None), + token_ids=(as_list(output.token_ids) if request.return_token_ids and not suppress_metadata else None), routed_experts=routed_experts_b64, ) choice_data = maybe_filter_parallel_tool_calls(choice_data, request) diff --git a/aphrodite/entrypoints/openai/dp_supervisor.py b/aphrodite/entrypoints/openai/dp_supervisor.py index 7f9f8ddebb..e724609410 100644 --- a/aphrodite/entrypoints/openai/dp_supervisor.py +++ b/aphrodite/entrypoints/openai/dp_supervisor.py @@ -248,14 +248,45 @@ def is_ready(self) -> bool: async def run(self) -> None: loop = asyncio.get_running_loop() + decorate_logs("DPSupervisor") # K8s sends SIGTERM for shutdown - begin graceful termination. for sig in (signal.SIGTERM, signal.SIGINT): loop.add_signal_handler(sig, partial(self._handle_signal, sig)) - # Launch DPSupervisor Server. + supervisor_server: uvicorn.Server | None = None + supervisor_server_task: asyncio.Task[None] | None = None + try: + # Launch Aphrodite DP Servers and begin monitoring them. + self._start_children() + monitor_task = asyncio.create_task(self._monitor_children(), name="dp-monitor") + + # Only start the DPSupervisor server once the Aphrodite DP Servers + # are ready. This avoids the supervisor /health endpoint returning + # 503 to external load balancer probes while the engines initialize. + await self._wait_until_ready(monitor_task) + if self.is_ready and not monitor_task.done(): + supervisor_server, supervisor_server_task = await self._start_server() + + await monitor_task + finally: + self._is_ready = False + await self._shutdown_children() + + # Shutdown the DP Supervisor server. + if supervisor_server is not None and supervisor_server_task is not None: + supervisor_server.should_exit = True + await supervisor_server_task + + async def _start_server(self) -> tuple[uvicorn.Server, asyncio.Task[None]]: + """ + Launch the DPSupervisor HTTP server. + + Called only after the Aphrodite DP Servers are ready so that /health + does not return 503 to external probes while the engines are + initializing. + """ app = _build_dp_supervisor_app(self) - decorate_logs("DPSupervisor") host = self.args.host or "0.0.0.0" config = uvicorn.Config( app, @@ -282,18 +313,20 @@ async def run(self) -> None: raise RuntimeError("DPSupervisor exited before startup.") await asyncio.sleep(0.05) logger.info("Started DPSupervisor on %s:%d", host, self.supervisor_port) + return supervisor_server, supervisor_server_task - # Launch and Monitor Aphrodite Server Processes. - try: - self._start_children() - await self._monitor_children() - finally: - self._is_ready = False - await self._shutdown_children() + async def _wait_until_ready(self, monitor_task: asyncio.Task[None]) -> None: + """ + Block until the Aphrodite DP Servers are ready or shutdown is triggered. - # Shutdown the DP Supervisor server. - supervisor_server.should_exit = True - await supervisor_server_task + Returns early if monitoring stops (e.g. a DP Server dies during + startup), in which case the supervisor server is never started. + """ + logger.info("Waiting for Aphrodite DP Servers to become ready.") + while not self._is_ready and not self._shutdown_event.is_set(): + if monitor_task.done(): + return + await asyncio.sleep(0.05) def _handle_signal(self, signum: int) -> None: """ diff --git a/aphrodite/entrypoints/openai/responses/context.py b/aphrodite/entrypoints/openai/responses/context.py index d63068f407..460bafbd24 100644 --- a/aphrodite/entrypoints/openai/responses/context.py +++ b/aphrodite/entrypoints/openai/responses/context.py @@ -305,7 +305,9 @@ def __init__( self.finish_reason: str | None = None self.enable_auto_tools = enable_auto_tools - self.response_parser = response_parser + self.response_parser = response_parser or ( + parser_cls(tokenizer, request.tools) if parser_cls is not None else None + ) self.parser_cls = parser_cls self.request = request @@ -339,6 +341,8 @@ def append_output(self, output: RequestOutput) -> None: enable_auto_tools=self.enable_auto_tools, model_output_token_ids=completion.token_ids, ) + if not self.request.include_reasoning: + reasoning = None self.response_messages.extend( build_response_output_items( reasoning=reasoning, diff --git a/aphrodite/entrypoints/openai/responses/protocol.py b/aphrodite/entrypoints/openai/responses/protocol.py index 7f2c526066..cb12b369d1 100644 --- a/aphrodite/entrypoints/openai/responses/protocol.py +++ b/aphrodite/entrypoints/openai/responses/protocol.py @@ -159,6 +159,15 @@ class ResponsesRequest(OpenAIBaseModel): previous_response_id: str | None = None prompt: ResponsePrompt | None = None reasoning: Reasoning | None = None + include_reasoning: bool = Field( + default=True, + description=( + "Whether to include reasoning content in the response. " + "When false, reasoning tokens are still generated but " + "excluded from the output. This reduces network traffic " + "without affecting model inference." + ), + ) service_tier: Literal["auto", "default", "flex", "scale", "priority"] = "auto" store: bool | None = True stream: bool | None = False diff --git a/aphrodite/entrypoints/openai/responses/serving.py b/aphrodite/entrypoints/openai/responses/serving.py index a3f0d66fdf..d2b294eeeb 100644 --- a/aphrodite/entrypoints/openai/responses/serving.py +++ b/aphrodite/entrypoints/openai/responses/serving.py @@ -997,6 +997,9 @@ def _make_response_output_items( enable_auto_tools=self.enable_auto_tools, model_output_token_ids=final_output.token_ids, ) + if not request.include_reasoning: + reasoning = None + logprobs = None return build_response_output_items( reasoning=reasoning, content=content, @@ -1248,12 +1251,15 @@ async def _process_simple_streaming_events( _increment_sequence_number_and_return: Callable[[StreamingResponsesResponse], StreamingResponsesResponse], ) -> AsyncGenerator[StreamingResponsesResponse, None]: processor = SimpleStreamingEventProcessor(tools=request.tools) + hide_stream_metadata = not request.include_reasoning and self.parser is not None def _get_logprobs( output: CompletionOutput, ) -> list[response_text_delta_event.Logprob]: if not request.is_include_output_logprobs(): return [] + if hide_stream_metadata: + return [] return self._create_stream_response_logprobs( token_ids=output.token_ids, logprobs=output.logprobs, diff --git a/aphrodite/entrypoints/serve/dev/rlhf/api_router.py b/aphrodite/entrypoints/serve/dev/rlhf/api_router.py index c9cf807b73..8fff3c7913 100644 --- a/aphrodite/entrypoints/serve/dev/rlhf/api_router.py +++ b/aphrodite/entrypoints/serve/dev/rlhf/api_router.py @@ -91,6 +91,51 @@ async def resume_generation(raw_request: Request) -> JSONResponse: ) +@router.post("/abort_requests") +async def abort_requests(raw_request: Request) -> JSONResponse: + """Abort in-flight requests without pausing the scheduler. + + Empty/missing ``request_ids`` aborts all in-flight requests. + """ + + engine = engine_client(raw_request) + + try: + body = await raw_request.json() + except json.JSONDecodeError as e: + raise HTTPException(status_code=400, detail="Invalid JSON format") from e # noqa: B904 + + request_ids = body.get("request_ids") + + try: + if request_ids: + # Body ids are external (user-supplied) request ids. + await engine.abort(request_ids) + else: + # The dev RL server runs AsyncLLM; abort everything it is tracking. + # request_states is keyed by internal ids; parent_requests holds + # parallel-sampling parents. Abort both as internal ids. + from aphrodite.v1.engine.async_llm import AsyncLLM + + assert isinstance(engine, AsyncLLM) + op = engine.output_processor + request_ids = [ + *op.request_states.keys(), + *op.parent_requests.keys(), + ] + await engine.abort(request_ids, internal=True) + return JSONResponse( + content={"status": "aborted", "aborted": len(request_ids)}, + status_code=HTTPStatus.OK.value, + ) + except Exception as err: # pragma: no cover - defensive + logger.exception("Failed to abort requests") + return JSONResponse( + content={"error": f"Failed to abort requests: {err}"}, + status_code=HTTPStatus.INTERNAL_SERVER_ERROR.value, + ) + + @router.get("/is_paused") async def is_paused(raw_request: Request) -> JSONResponse: """Return the current pause status.""" @@ -131,6 +176,12 @@ async def start_weight_update(raw_request: Request): return JSONResponse(content={"message": "Weight update started"}) +@router.post("/start_draft_weight_update") +async def start_draft_weight_update(raw_request: Request): + await engine_client(raw_request).start_draft_weight_update() + return JSONResponse(content={"message": "Draft weight update started"}) + + @router.post("/update_weights") async def update_weights(raw_request: Request): try: diff --git a/aphrodite/envs.py b/aphrodite/envs.py index 8ee86052fe..78a16d60ce 100755 --- a/aphrodite/envs.py +++ b/aphrodite/envs.py @@ -197,6 +197,7 @@ APHRODITE_BLOCKSCALE_FP8_GEMM_FLASHINFER: bool = True APHRODITE_USE_FLASHINFER_MOE_INT4: bool = False APHRODITE_FLASHINFER_AUTOTUNE_CACHE_DIR: str | None = None + APHRODITE_FLASHINFER_AUTOTUNE_SKIP_OPS: list[str] | None = None APHRODITE_FLASHINFER_ALLREDUCE_BACKEND: Literal["auto", "trtllm", "mnnvl"] = "auto" APHRODITE_FLASHINFER_WORKSPACE_BUFFER_SIZE: int = 394 * 1024 * 1024 APHRODITE_XGRAMMAR_CACHE_MB: int = 0 @@ -1404,6 +1405,14 @@ def _resolve_rust_frontend_path() -> str | None: "MOONCAKE_REQUESTER_LOCAL_HOSTNAME": lambda: os.getenv("MOONCAKE_REQUESTER_LOCAL_HOSTNAME"), # Override the directory for the FlashInfer autotune config cache. "APHRODITE_FLASHINFER_AUTOTUNE_CACHE_DIR": lambda: os.getenv("APHRODITE_FLASHINFER_AUTOTUNE_CACHE_DIR", None), + # Comma-separated FlashInfer op names to exclude from autotuning, using + # the heuristic fallback tactic instead. Unset: skip "fp4_gemm" when the + # CuTe-DSL NVFP4 linear kernel is selected. Empty: skip nothing. + "APHRODITE_FLASHINFER_AUTOTUNE_SKIP_OPS": lambda: ( + None + if "APHRODITE_FLASHINFER_AUTOTUNE_SKIP_OPS" not in os.environ + else [v.strip() for v in os.environ["APHRODITE_FLASHINFER_AUTOTUNE_SKIP_OPS"].split(",") if v.strip()] + ), # Flashinfer fused allreduce backend. "APHRODITE_FLASHINFER_ALLREDUCE_BACKEND": env_with_choices( "APHRODITE_FLASHINFER_ALLREDUCE_BACKEND", @@ -1869,6 +1878,7 @@ def compile_factors() -> dict[str, object]: "APHRODITE_DEBUG_LOG_API_SERVER_RESPONSE", "APHRODITE_TUNED_CONFIG_FOLDER", "APHRODITE_FLASHINFER_AUTOTUNE_CACHE_DIR", + "APHRODITE_FLASHINFER_AUTOTUNE_SKIP_OPS", "APHRODITE_ENGINE_ITERATION_TIMEOUT_S", "APHRODITE_HTTP_TIMEOUT_KEEP_ALIVE", "APHRODITE_EXECUTE_MODEL_TIMEOUT_SECONDS", diff --git a/aphrodite/lora/layers/utils.py b/aphrodite/lora/layers/utils.py index de00c3422e..9873f2e029 100644 --- a/aphrodite/lora/layers/utils.py +++ b/aphrodite/lora/layers/utils.py @@ -57,6 +57,9 @@ def _get_lora_device(base_layer: nn.Module) -> torch.device: # GPTQ/AWQ elif hasattr(base_layer, "qweight"): return base_layer.qweight.device + # INC WNA16 (AutoRound) + elif hasattr(base_layer, "ark_linear"): + return base_layer.ark_linear.qweight.device # MoE layer elif hasattr(base_layer, "w2_weight"): return base_layer.w2_weight.device diff --git a/aphrodite/model_executor/layers/fused_moe/experts/marlin_moe.py b/aphrodite/model_executor/layers/fused_moe/experts/marlin_moe.py index 9af65fa95d..91254167f4 100644 --- a/aphrodite/model_executor/layers/fused_moe/experts/marlin_moe.py +++ b/aphrodite/model_executor/layers/fused_moe/experts/marlin_moe.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Fused MoE utilities for GPTQ.""" +import math from collections.abc import Callable import torch @@ -306,6 +307,12 @@ def fused_marlin_moe( assert num_bits in [4, 8] assert topk_weights.dtype == torch.float32 + if global_num_experts == -1: + global_num_experts = E + else: + # Set M to estimated valid tokens per rank. + M = math.ceil(M * E / global_num_experts) + # M block size selection logic # TODO: tune this further for specific models for block_size_m in [8, 16, 32, 48, 64]: @@ -315,8 +322,6 @@ def fused_marlin_moe( if input_dtype is not None and input_dtype.itemsize == 1: block_size_m = max(block_size_m, 16) - if global_num_experts == -1: - global_num_experts = E sorted_token_ids, expert_ids, num_tokens_post_padded = moe_align_block_size( topk_ids, block_size_m, diff --git a/aphrodite/model_executor/layers/fused_moe/experts/xpu_moe.py b/aphrodite/model_executor/layers/fused_moe/experts/xpu_moe.py index 0eed6e7154..ca3ae7a02b 100644 --- a/aphrodite/model_executor/layers/fused_moe/experts/xpu_moe.py +++ b/aphrodite/model_executor/layers/fused_moe/experts/xpu_moe.py @@ -62,11 +62,6 @@ def __init__( max_num_tokens, num_dispatchers, ) - self.is_fp8 = False - self.is_int4 = False - self.is_mxfp4 = False - self.is_block_fp8 = False - self.is_mxfp8 = False self.gemm1_clamp_limit = quant_config.gemm1_clamp_limit self.fused_moe_impl: XpuFusedMoe | None = None @@ -149,17 +144,11 @@ def apply( expert_tokens_meta: mk.ExpertTokensMetadata | None, apply_router_weight_on_input: bool, ): - # The kernel takes is_fp8/is_int4/is_mxfp4 as independent booleans. - # In this hierarchy each subclass flips exactly one to True; assert - # the invariant so a future subclass that sets two doesn't silently - # miscompute (kernel-side priority is undocumented). - assert sum([self.is_fp8, self.is_int4, self.is_mxfp4]) <= 1, ( - "XPUExperts: at most one of is_fp8, is_int4, is_mxfp4 may be True; " - f"got is_fp8={self.is_fp8}, is_int4={self.is_int4}, " - f"is_mxfp4={self.is_mxfp4}." - ) if self.fused_moe_impl is None: topk = topk_ids.size(-1) + if self.quant_config is not None and self.quant_config.weight_quant_dtype == "mxfp4": + w1 = w1.view(torch.float4_e2m1fn_x2) + w2 = w2.view(torch.float4_e2m1fn_x2) self.fused_moe_impl = XpuFusedMoe( w13=w1, w13_scales=self.w1_scale, @@ -172,11 +161,6 @@ def apply( num_experts=self.moe_config.num_local_experts, ep_rank=self.moe_config.ep_rank, ep_size=self.moe_config.ep_size, - is_fp8=self.is_fp8, - is_int4=self.is_int4, - is_mxfp4=self.is_mxfp4, - is_mxfp8=self.is_mxfp8, - is_block_fp8=self.is_block_fp8, gemm1_clamp_limit=self.gemm1_clamp_limit, ) assert self.fused_moe_impl is not None @@ -202,7 +186,6 @@ def __init__( max_num_tokens, num_dispatchers, ) - self.is_fp8 = True @staticmethod def _supports_quant_scheme( @@ -231,7 +214,6 @@ def __init__( num_dispatchers, ) assert quant_config.quant_dtype == "mxfp8" - self.is_mxfp8 = True @staticmethod def _supports_quant_scheme( @@ -259,7 +241,6 @@ def __init__( max_num_tokens, num_dispatchers, ) - self.is_block_fp8 = True @staticmethod def _supports_quant_scheme( @@ -298,7 +279,6 @@ def __init__( max_num_tokens, num_dispatchers, ) - self.is_int4 = True @staticmethod def _supports_quant_scheme( @@ -325,7 +305,6 @@ def __init__( max_num_tokens, num_dispatchers, ) - self.is_mxfp4 = True @staticmethod def _supports_quant_scheme( diff --git a/aphrodite/model_executor/layers/fused_moe/router/gate_linear.py b/aphrodite/model_executor/layers/fused_moe/router/gate_linear.py index 5d87e5ea5e..3357a2719a 100644 --- a/aphrodite/model_executor/layers/fused_moe/router/gate_linear.py +++ b/aphrodite/model_executor/layers/fused_moe/router/gate_linear.py @@ -15,8 +15,8 @@ class GateLinear(ReplicatedLinear): """MoE gate linear layer with multi-tier GEMM dispatch: 1. DSV3 specialized kernel (SM90+, M<=16, H=7168 E=256/384, H=6144 E=256) - 2. fp32 specialized kernel (SM90+, bf16/fp32 in, fp32 out, - M<=32, H=3072, E=256) + 2. fp32 specialized kernel (SM90+, bf16/fp32 in, fp32 out, M<=32, + (H, E) in {(3072, 256), (6144, 128), (6144, 256)}) 3. cuBLAS bf16×bf16→fp32 (SM90+ + bf16 weight + fp32 out_dtype) 4. F.linear via ReplicatedLinear (ultimate fallback) @@ -36,7 +36,8 @@ class GateLinear(ReplicatedLinear): # (hidden_size, num_experts) pairs with an instantiated fp32 kernel: # (3072, 256) -> MiniMax-M2/M2.5, (6144, 128) -> MiniMax-M3 - FP32_SUPPORTED_SHAPES = {(3072, 256), (6144, 128)} + # (6144, 256) -> GLM-5.2 + FP32_SUPPORTED_SHAPES = {(3072, 256), (6144, 128), (6144, 256)} FP32_MAX_TOKENS = 32 def __init__( @@ -127,7 +128,7 @@ def forward(self, x: torch.Tensor) -> torch.Tensor | tuple[torch.Tensor, Paramet ) return output, None - # Tier 2: fp32 specialized kernel (H=3072, E=256, M<=32) + # Tier 2: fp32 specialized kernel (supported H/E pairs, M<=32) # Dispatch is wrapped in a custom op so that torch.compile/CUDA-graph # capture does not freeze the runtime num_tokens branch. if self.allow_fp32_router_gemm and x.dtype in ( diff --git a/aphrodite/model_executor/layers/hpc/rope_norm.py b/aphrodite/model_executor/layers/hpc/rope_norm.py index 84100c4b67..7d113f6a9d 100644 --- a/aphrodite/model_executor/layers/hpc/rope_norm.py +++ b/aphrodite/model_executor/layers/hpc/rope_norm.py @@ -61,15 +61,15 @@ def hpc_rope_norm_forward( return attn_layer = forward_context.no_compile_layers[layer_name] - # bind_kv_cache stores the per-layer KV cache as a single 5D tensor - # (num_blocks, 2, block_size, num_kv_heads, head_size), so use it directly. + # bind_kv_cache stores the per-layer KV cache as a packed 4D tensor + # (num_blocks, num_kv_heads, block_size, 2 * head_size). kv_cache = attn_layer.kv_cache if kv_cache.numel() == 0: output.zero_() return - assert kv_cache.dim() == 5, f"Expected kv_cache to have 5 dims, got {tuple(kv_cache.shape)}" + assert kv_cache.dim() == 4, f"Expected kv_cache to have 4 dims, got {tuple(kv_cache.shape)}" rope_norm = _hpc_rope_norm_instances[layer_name] rope_norm._forward_impl(qkv, kv_cache, attn_metadata, attn_layer, output) @@ -300,6 +300,7 @@ def _forward_impl( # rope_norm_store_kv_fp8 kernel can write quantized K/V in-place. if self.use_fp8: kv_cache = kv_cache.view(torch.float8_e4m3fn) + key_cache, value_cache = kv_cache.transpose(1, 2).split(self.head_dim, dim=-1) # Per-tensor K/V scales (shape [1]) used by the FP8 kernel. k_scale = attn_layer._k_scale.reshape(1) @@ -319,8 +320,8 @@ def _forward_impl( if self.use_fp8: _, q_scale, split_k_flag = hpc.rope_norm_store_kv_fp8( - key_cache=kv_cache[:, 0], - value_cache=kv_cache[:, 1], + key_cache=key_cache, + value_cache=value_cache, qkv=qkv_prefill, cos_sin=self.cos_sin_cache, num_seqlen_per_req=seq_lens_prefill, @@ -339,8 +340,8 @@ def _forward_impl( attn_metadata.hpc_prefill_q_scale = q_scale else: hpc.rope_norm_store_kv( - kv_cache[:, 0], - kv_cache[:, 1], + key_cache, + value_cache, qkv_prefill, self.cos_sin_cache, seq_lens_prefill, @@ -363,8 +364,8 @@ def _forward_impl( if self.use_fp8: _, q_scale, split_k_flag = hpc.rope_norm_store_kv_fp8( - key_cache=kv_cache[:, 0], - value_cache=kv_cache[:, 1], + key_cache=key_cache, + value_cache=value_cache, qkv=qkv_decode, cos_sin=self.cos_sin_cache, num_seqlen_per_req=num_seq_kvcache, @@ -385,8 +386,8 @@ def _forward_impl( attn_metadata.hpc_split_k_flag = split_k_flag else: hpc.rope_norm_store_kv( - kv_cache[:, 0], - kv_cache[:, 1], + key_cache, + value_cache, qkv_decode, self.cos_sin_cache, num_seq_kvcache, diff --git a/aphrodite/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py b/aphrodite/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py index 8672010a6f..39273533f8 100644 --- a/aphrodite/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py +++ b/aphrodite/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py @@ -431,6 +431,7 @@ def __init__( aphrodite_config: AphroditeConfig, prefix: str = "", gqa_interleaved_layout=False, + reduce_results: bool = True, ) -> None: super().__init__(config, aphrodite_config, prefix) @@ -539,6 +540,7 @@ def __init__( self.hidden_size, bias=False, input_is_parallel=True, + reduce_results=reduce_results, quant_config=self.quant_config, prefix=f"{prefix}.out_proj", ) @@ -820,17 +822,14 @@ def rearrange_mixed_qkv(self, mixed_qkv): def forward( self, hidden_states: torch.Tensor, - output: torch.Tensor, - ): - self._forward_method(hidden_states, output) + ) -> torch.Tensor: + return self._forward_method(hidden_states) def _output_projection( self, core_attn_out: torch.Tensor, z: torch.Tensor, - output: torch.Tensor, - num_tokens: int, - ): + ) -> torch.Tensor: """Part 3: RMSNormGated + output linear projection. The RMSNormGated + quant sequence is eligible for fusion @@ -842,13 +841,13 @@ def _output_projection( core_attn_out = self.norm(core_attn_out, z) core_attn_out = core_attn_out.reshape(z_shape_og) core_attn_out = core_attn_out.flatten(-2) # ... h d -> ... (h d) - output[:num_tokens], _ = self.out_proj(core_attn_out) + output, _ = self.out_proj(core_attn_out) + return output def forward_hip( self, hidden_states: torch.Tensor, - output: torch.Tensor, - ): + ) -> torch.Tensor: """ROCm forward using AITER Triton fused projection+attention when available, otherwise falling back to the generic CUDA path.""" if GDN_AITER_TRITON_AVAILABLE: @@ -877,15 +876,14 @@ def forward_hip( use_aiter=True, ) - self._output_projection(core_attn_out, z, output, num_tokens) + return self._output_projection(core_attn_out, z) else: - self.forward_cuda(hidden_states, output) + return self.forward_cuda(hidden_states) def forward_cuda( self, hidden_states: torch.Tensor, - output: torch.Tensor, - ): + ) -> torch.Tensor: """ Forward pass with three parts: 1. Input projection @@ -936,13 +934,12 @@ def forward_cuda( # ============================================================ # Part 3: Output Projection # ============================================================ - self._output_projection(core_attn_out, z, output, num_tokens) + return self._output_projection(core_attn_out, z) def forward_xpu( self, hidden_states: torch.Tensor, - output: torch.Tensor, - ): + ) -> torch.Tensor: """ Forward pass with three parts: 1. Input projection @@ -985,13 +982,13 @@ def forward_xpu( core_attn_out = self.norm(core_attn_out, z) core_attn_out = core_attn_out.reshape(z_shape_og) core_attn_out = core_attn_out.flatten(-2) # ... h d -> ... (h d) - output[:num_tokens], _ = self.out_proj(core_attn_out) + out, _ = self.out_proj(core_attn_out) + return out def forward_cpu( self, hidden_states: torch.Tensor, - output: torch.Tensor, - ): + ) -> torch.Tensor: assert not hasattr(self, "in_proj_qkv"), "lora isn't supported on CPU." mixed_qkvz, _ = self.in_proj_qkvz(hidden_states) @@ -1031,7 +1028,8 @@ def forward_cpu( core_attn_out = self.norm(core_attn_out, z) core_attn_out = core_attn_out.reshape(z_shape_og) core_attn_out = core_attn_out.flatten(-2) # ... h d -> ... (h d) - output[:num_tokens], _ = self.out_proj(core_attn_out) + out, _ = self.out_proj(core_attn_out) + return out def _warmup_prefill_kernels(self, qkv_or_qkvz: torch.Tensor, v_dim: int) -> None: """Warm up GDN prefill kernels during V1 profiling. diff --git a/aphrodite/model_executor/layers/quantization/utils/flashinfer_fp4_moe.py b/aphrodite/model_executor/layers/quantization/utils/flashinfer_fp4_moe.py index 5c8fdd1ec2..7eb3b7148f 100644 --- a/aphrodite/model_executor/layers/quantization/utils/flashinfer_fp4_moe.py +++ b/aphrodite/model_executor/layers/quantization/utils/flashinfer_fp4_moe.py @@ -23,25 +23,41 @@ logger = init_logger(__name__) - __all__ = [ "reorder_w1w3_to_w3w1", ] def reorder_w1w3_to_w3w1(weight: torch.Tensor, scale: torch.Tensor, dim: int = -2) -> tuple[torch.Tensor, torch.Tensor]: - """Re-order the concatenated `[w1, w3]` tensors to `[w3, w1]`""" + """Re-order concatenated `[w1, w3]` tensors to `[w3, w1]` in-place. + + `weight` and `scale` must be contiguous; they remain contiguous on return. + """ + assert weight.is_contiguous(), "weight must be contiguous" + assert scale.is_contiguous(), "scale must be contiguous" size = weight.size(dim) assert size % 2 == 0, f"Expected even size in dim {dim}, got {size}" half = size // 2 + d = dim % weight.dim() - w1, w3 = weight.split(half, dim=dim) - s1, s3 = scale.split(half, dim=dim) - - return ( - torch.cat([w3, w1], dim=dim).contiguous(), - torch.cat([s3, s1], dim=dim).contiguous(), + # 64 MB transient cap + bytes_per_row = max( + weight.numel() // size * weight.element_size(), + scale.numel() // size * scale.element_size(), ) + chunk = max(1, min(half, (64 << 20) // max(bytes_per_row, 1))) + + fa, fb = [slice(None)] * weight.dim(), [slice(None)] * weight.dim() + for off in range(0, half, chunk): + end = min(off + chunk, half) + fa[d], fb[d] = slice(off, end), slice(half + off, half + end) + a, b = tuple(fa), tuple(fb) + for t in (weight, scale): + tmp = t[b].clone() + t[b] = t[a] + t[a] = tmp + + return weight, scale def interleave_linear_and_gate( diff --git a/aphrodite/model_executor/layers/quantization/utils/marlin_utils_fp4.py b/aphrodite/model_executor/layers/quantization/utils/marlin_utils_fp4.py index baee1faae7..eefb38dffd 100644 --- a/aphrodite/model_executor/layers/quantization/utils/marlin_utils_fp4.py +++ b/aphrodite/model_executor/layers/quantization/utils/marlin_utils_fp4.py @@ -290,6 +290,37 @@ def prepare_fp4_layer_for_marlin(layer: torch.nn.Module, input_dtype: torch.dtyp return +def _repack_marlin_experts( + weight: torch.Tensor, + size_n: int, + size_k: int, + perm: torch.Tensor, + is_a_8bit: bool, +) -> torch.Tensor: + """Repack each expert to marlin format into a preallocated output.""" + num_experts = weight.shape[0] + out: torch.Tensor | None = None + for i in range(num_experts): + qweight = weight[i].view(torch.int32).T.contiguous() + marlin_qweight = ops.gptq_marlin_repack( + b_q_weight=qweight, + perm=perm, + size_k=size_k, + size_n=size_n, + num_bits=4, + is_a_8bit=is_a_8bit, + ) + if out is None: + out = torch.empty( + (num_experts, *marlin_qweight.shape), + dtype=marlin_qweight.dtype, + device=marlin_qweight.device, + ) + out[i] = marlin_qweight + assert out is not None + return out + + def prepare_nvfp4_moe_layer_for_marlin( layer: RoutedExperts, w13: torch.Tensor, @@ -353,7 +384,6 @@ def pad_w2(x: torch.Tensor, packing: int) -> torch.Tensor: # WEIGHT # Repack weights to marlin format def repack_weight(weight: torch.Tensor, name: str) -> torch.Tensor: - tensor_list = [] if "w13" in name: size_n, size_k = N * num_shards, K assert weight.shape == (E, size_n, size_k // 2) @@ -365,20 +395,7 @@ def repack_weight(weight: torch.Tensor, name: str) -> torch.Tensor: weight = pad_w2(weight, packing=2) size_k = padded_N - for i in range(E): - qweight = weight[i].view(torch.int32).T.contiguous() - - marlin_qweight = ops.gptq_marlin_repack( - b_q_weight=qweight, - perm=perm, - size_k=size_k, - size_n=size_n, - num_bits=4, - is_a_8bit=is_a_8bit, - ) - tensor_list.append(marlin_qweight) - - return torch.cat([x.unsqueeze(0) for x in tensor_list], 0) + return _repack_marlin_experts(weight, size_n, size_k, perm, is_a_8bit) w13 = repack_weight(w13, "w13") w2 = repack_weight(w2, "w2") @@ -450,7 +467,6 @@ def prepare_moe_fp4_layer_for_marlin(layer: torch.nn.Module, input_dtype: torch. # Repack weights to marlin format for name in ["w13_weight", "w2_weight"]: weight = getattr(layer, name) - tensor_list = [] if "w13" in name: size_n, size_k = n * 2, k else: @@ -458,20 +474,7 @@ def prepare_moe_fp4_layer_for_marlin(layer: torch.nn.Module, input_dtype: torch. assert weight.shape == (e, size_n, size_k // 2) - for i in range(e): - qweight = weight[i].view(torch.int32).T.contiguous() - - marlin_qweight = ops.gptq_marlin_repack( - b_q_weight=qweight, - perm=perm, - size_k=size_k, - size_n=size_n, - num_bits=4, - is_a_8bit=is_a_8bit, - ) - tensor_list.append(marlin_qweight) - - weight = torch.cat([x.unsqueeze(0) for x in tensor_list], 0) + weight = _repack_marlin_experts(weight, size_n, size_k, perm, is_a_8bit) weight = torch.nn.Parameter(weight, requires_grad=False) setattr(layer, name, weight) @@ -587,7 +590,6 @@ def prepare_moe_mxfp4_layer_for_marlin( # WEIGHT: Repack weights to marlin format def repack_weight(weight: torch.Tensor, name: str) -> torch.Tensor: - tensor_list = [] if "w13" in name: size_n, size_k = n * 2, k else: @@ -595,18 +597,7 @@ def repack_weight(weight: torch.Tensor, name: str) -> torch.Tensor: assert weight.shape == (e, size_n, size_k // 2) - for i in range(e): - qweight = weight[i].view(torch.int32).T.contiguous() - marlin_qweight = ops.gptq_marlin_repack( - b_q_weight=qweight, - perm=perm, - size_k=size_k, - size_n=size_n, - num_bits=4, - is_a_8bit=is_a_8bit, - ) - tensor_list.append(marlin_qweight) - return torch.cat([x.unsqueeze(0) for x in tensor_list], 0) + return _repack_marlin_experts(weight, size_n, size_k, perm, is_a_8bit) w13 = repack_weight(w13, "w13") w2 = repack_weight(w2, "w2") diff --git a/aphrodite/model_executor/model_loader/weight_utils.py b/aphrodite/model_executor/model_loader/weight_utils.py index 9f3aa23632..71bada256c 100644 --- a/aphrodite/model_executor/model_loader/weight_utils.py +++ b/aphrodite/model_executor/model_loader/weight_utils.py @@ -568,6 +568,12 @@ def filter_duplicate_safetensors_files(hf_weights_files: list[str], hf_folder: s weight_files_in_index = set() for weight_name in weight_map: weight_files_in_index.add(os.path.join(hf_folder, weight_map[weight_name])) + # Check if files referenced in model.safetensors.index.json actually exist. + # Raise error if any file is missing. + hf_weights_files_set = set(hf_weights_files) + missing_files = weight_files_in_index - hf_weights_files_set + if missing_files: + raise FileNotFoundError(f"Weight files referenced in index but missing: {missing_files}") # Filter out any fields that are not found in the index file. hf_weights_files = [f for f in hf_weights_files if f in weight_files_in_index] return hf_weights_files diff --git a/aphrodite/model_executor/models/config.py b/aphrodite/model_executor/models/config.py index bf81f5b1e3..87014a9887 100644 --- a/aphrodite/model_executor/models/config.py +++ b/aphrodite/model_executor/models/config.py @@ -466,7 +466,7 @@ def verify_and_update_model_config(model_config: "ModelConfig") -> None: "last": "LAST", } - pooling_type = pooling_type_map.get(hf_config.pooling, None) + pooling_type = pooling_type_map.get(hf_config.pooling) if pooling_type is None: raise ValueError(f"pool_type {hf_config.pooling!r} not supported") @@ -753,6 +753,23 @@ def verify_and_update_model_config(model_config: "ModelConfig") -> None: model_config.hf_config.embedding_size = model_config.hf_config.num_labels +class LongcatFlashNgramForCausalLMConfig(VerifyAndUpdateConfig): + @staticmethod + def verify_and_update_config(aphrodite_config: "AphroditeConfig") -> None: + # LongCat-Flash-Lite's zero-expert MoE trips a data-dependent assert + # under torch.compile, and its n-gram inputs_embeds are only wired for + # FULL cudagraph capture (PIECEWISE prefill drops them). Default to + # no-compile + FULL cudagraph (prefill runs eager) unless the user + # configured compilation explicitly. + from aphrodite.config.compilation import CompilationMode, CUDAGraphMode + + compilation_config = aphrodite_config.compilation_config + if compilation_config.mode is None: + compilation_config.mode = CompilationMode.NONE + if compilation_config.cudagraph_mode is None: + compilation_config.cudagraph_mode = CUDAGraphMode.FULL + + MODELS_CONFIG_MAP: dict[str, type[VerifyAndUpdateConfig]] = { "ColBERTJinaRobertaModel": JinaRobertaModelConfig, "ColQwen3_5": ColQwen3_5Config, @@ -766,6 +783,7 @@ def verify_and_update_model_config(model_config: "ModelConfig") -> None: "Gemma4ForConditionalGeneration": Gemma4Config, "Gemma4UnifiedForConditionalGeneration": Gemma4Config, "GptOssForCausalLM": GptOssForCausalLMConfig, + "LongcatFlashNgramForCausalLM": LongcatFlashNgramForCausalLMConfig, "GteModel": SnowflakeGteNewModelConfig, "GteNewForSequenceClassification": GteNewModelConfig, "GteNewModel": GteNewModelConfig, diff --git a/aphrodite/model_executor/models/granitemoehybrid.py b/aphrodite/model_executor/models/granitemoehybrid.py index 18a4b20a73..aca64f5047 100644 --- a/aphrodite/model_executor/models/granitemoehybrid.py +++ b/aphrodite/model_executor/models/granitemoehybrid.py @@ -308,8 +308,12 @@ def forward( ALL_DECODER_LAYER_TYPES = { + # Transformers < 5.13.0 "attention": GraniteMoeHybridAttentionDecoderLayer, "mamba": GraniteMoeHybridMambaDecoderLayer, + # Transformers >= 5.13.0 + "full_attention": GraniteMoeHybridAttentionDecoderLayer, + "linear_attention": GraniteMoeHybridMambaDecoderLayer, } diff --git a/aphrodite/model_executor/models/hunyuan_vision.py b/aphrodite/model_executor/models/hunyuan_vision.py index 14690eb2f5..e93497d416 100644 --- a/aphrodite/model_executor/models/hunyuan_vision.py +++ b/aphrodite/model_executor/models/hunyuan_vision.py @@ -73,6 +73,7 @@ BaseProcessingInfo, PromptReplacement, PromptUpdate, + PromptUpdateDetails, ) from aphrodite.sequence import IntermediateTensors from aphrodite.transformers_utils.configs.hunyuan_vl import ( @@ -721,8 +722,10 @@ def _get_prompt_updates( hf_processor = self.info.get_hf_processor(**hf_processor_mm_kwargs) image_processor = self.info.get_image_processor(**hf_processor_mm_kwargs) - placeholder = { + token_ids = { "image": hf_processor.image_token_id, + "image_start": hf_processor.image_start_token_id, + "image_end": hf_processor.image_end_token_id, } merge_size = image_processor.merge_size @@ -734,12 +737,15 @@ def get_replacement_hunyuan_vl(item_idx: int, modality: str): _, grid_h, grid_w = grid_thw num_tokens = (int(grid_h) // merge_size) * (int(grid_w) // merge_size + 1) + 2 - return [placeholder[modality]] * num_tokens + tokens = ( + [token_ids[f"{modality}_start"]] + [token_ids[modality]] * num_tokens + [token_ids[f"{modality}_end"]] + ) + return PromptUpdateDetails.select_token_id(tokens, token_ids[modality]) return [ PromptReplacement( modality=modality, - target=[placeholder[modality]], + target=[token_ids[modality]], replacement=partial(get_replacement_hunyuan_vl, modality=modality), ) for modality in ("image",) diff --git a/aphrodite/model_executor/models/longcat_flash.py b/aphrodite/model_executor/models/longcat_flash.py index 0ce8b25f75..fa7dec9a24 100644 --- a/aphrodite/model_executor/models/longcat_flash.py +++ b/aphrodite/model_executor/models/longcat_flash.py @@ -306,7 +306,7 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: hidden_states = hidden_states.view(-1, hidden_dim) # Align to FusedMoE padded hidden size to avoid dim mismatch - padded_hidden = self.experts.hidden_size + padded_hidden = self.experts.moe_config.hidden_dim if hidden_dim < padded_hidden: hidden_states_padded = torch.nn.functional.pad( hidden_states, @@ -647,10 +647,15 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: ) self_attn.w_kc = w_kc.transpose(1, 2).contiguous().transpose(1, 2) self_attn.w_vc = w_vc.contiguous().transpose(1, 2) - if self.config.mla_scale_q_lora: + # Guard against compounding on incremental load_weights calls: + # the in-place ``*=`` would otherwise re-apply the MLA LoRA + # scaling to the layernorm weights on each pass. + if self.config.mla_scale_q_lora and not getattr(self_attn, "_mla_q_lora_scaled", False): self_attn.q_a_layernorm.weight.data *= (self.config.hidden_size / self.config.q_lora_rank) ** 0.5 - if self.config.mla_scale_kv_lora: + self_attn._mla_q_lora_scaled = True + if self.config.mla_scale_kv_lora and not getattr(self_attn, "_mla_kv_lora_scaled", False): self_attn.kv_a_layernorm.weight.data *= (self.config.hidden_size / self.config.kv_lora_rank) ** 0.5 + self_attn._mla_kv_lora_scaled = True return loaded_params diff --git a/aphrodite/model_executor/models/longcat_flash_mtp.py b/aphrodite/model_executor/models/longcat_flash_mtp.py index 3803e87467..4b9a8ea599 100644 --- a/aphrodite/model_executor/models/longcat_flash_mtp.py +++ b/aphrodite/model_executor/models/longcat_flash_mtp.py @@ -122,8 +122,10 @@ def forward( class LongCatFlashMTP(nn.Module): def __init__(self, *, aphrodite_config: AphroditeConfig, prefix: str = ""): super().__init__() - # LongCat MTP without MoE layers - aphrodite_config.model_config.hf_config.n_routed_experts = None + # LongCat MTP has no MoE layers: clear n_routed_experts so the predictor + # builds a dense MLP. object.__setattr__ bypasses the ngram remote + # config's strict validation (it rejects setting the int field to None). + object.__setattr__(aphrodite_config.model_config.hf_config, "n_routed_experts", None) self.config = FlashConfig(**aphrodite_config.model_config.hf_config.__dict__) self.quant_config = ( None if "mtp" in getattr(self.config, "disable_quant_module", []) else aphrodite_config.quant_config @@ -272,10 +274,14 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: ) self_attn.w_kc = w_kc.transpose(1, 2).contiguous().transpose(1, 2) self_attn.w_vc = w_vc.contiguous().transpose(1, 2) - if self.config.mla_scale_q_lora: + # Guard against compounding on incremental load_weights calls (the + # in-place *= would otherwise double-apply the LoRA scaling). + if self.config.mla_scale_q_lora and not getattr(self_attn, "_mla_q_lora_scaled", False): self_attn.q_a_layernorm.weight.data *= (self.config.hidden_size / self.config.q_lora_rank) ** 0.5 - if self.config.mla_scale_kv_lora: + self_attn._mla_q_lora_scaled = True + if self.config.mla_scale_kv_lora and not getattr(self_attn, "_mla_kv_lora_scaled", False): self_attn.kv_a_layernorm.weight.data *= (self.config.hidden_size / self.config.kv_lora_rank) ** 0.5 + self_attn._mla_kv_lora_scaled = True return loaded_params def _rewrite_spec_layer_name(self, spec_layer: int, name: str, new_to_old_names_mapping: dict) -> str: diff --git a/aphrodite/model_executor/models/longcat_flash_ngram.py b/aphrodite/model_executor/models/longcat_flash_ngram.py new file mode 100644 index 0000000000..f2ad873cc2 --- /dev/null +++ b/aphrodite/model_executor/models/longcat_flash_ngram.py @@ -0,0 +1,370 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Inference-only LongCat-Flash-Lite (n-gram embedding) model. + +``LongcatFlashNgramForCausalLM`` is LongCat-Flash (MLA dual-attention + +zero-expert MoE + YaRN) plus an n-gram embedding input layer: each position's +embedding fuses the token embedding with hashed embeddings of the preceding +``n`` tokens. That per-request token history is isolated in a Model-Runner-V2 +:class:`LongcatNgramModelState` (mirroring ``DiffusionGemmaModelState``), so +``get_model_state_cls`` makes the model MRV2-only. +""" + +from collections.abc import Iterable +from typing import Any + +import torch +from torch import nn + +from aphrodite import _custom_ops as ops +from aphrodite.config import AphroditeConfig +from aphrodite.distributed import get_pp_group +from aphrodite.model_executor.layers.logits_processor import LogitsProcessor +from aphrodite.model_executor.layers.vocab_parallel_embedding import ( + ParallelLMHead, + VocabParallelEmbedding, +) +from aphrodite.v1.core.sched.output import NewRequestData +from aphrodite.v1.worker.gpu.input_batch import InputBatch +from aphrodite.v1.worker.gpu.model_states.default import DefaultModelState +from aphrodite.v1.worker.gpu.states import RequestState + +from .interfaces import SupportsLoRA, SupportsPP +from .longcat_flash import FlashConfig, FlashModel +from .utils import AutoWeightsLoader, PPMissingLayer, maybe_prefix + + +def uses_ngram_embedding(config: FlashConfig) -> bool: + return getattr(config, "ngram_vocab_size_ratio", None) is not None + + +def _config_dtype(config: FlashConfig) -> torch.dtype: + dt = getattr(config, "torch_dtype", None) or getattr(config, "dtype", None) + if isinstance(dt, torch.dtype): + return dt + return getattr(torch, str(dt), None) or torch.bfloat16 + + +class NgramEmbedding(nn.Module): + """Token embedding fused with hashed n-gram embeddings. + + TP-sharded: the ``k*(n-1)`` per-embedder tables are concatenated into one + :class:`VocabParallelEmbedding` (``oe_embedder``) with per-embedder offsets, + and the projections are stacked into one ``oe_projection`` applied with a + single ``bmm``. Hashing math is ported from the HF reference. + """ + + def __init__(self, config: FlashConfig, base_embeddings: nn.Module) -> None: + super().__init__() + self.config = config + self.word_embeddings = base_embeddings + + self.m = config.ngram_vocab_size_ratio * config.vocab_size + self.k = config.emb_split_num + self.n = config.emb_neighbor_num + self.pad_id = config.pad_token_id + self.eos_token_id = config.eos_token_id + self._dtype = _config_dtype(config) + + self._init_ngram_embeddings() + + def _init_ngram_embeddings(self) -> None: + self.num_embedders = self.k * (self.n - 1) + oe_dim = self.config.hidden_size // self.num_embedders + self.oe_dim = oe_dim + + # Exclusive prefix sums of per-embedder table sizes; each embedder's + # local id is offset into the single concatenated table. + sizes = [int(self.m + i * 2 + 1) for i in range(self.num_embedders)] + offsets = [0] + for s in sizes: + offsets.append(offsets[-1] + s) + self._offsets = offsets # len num_embedders + 1 + self._sizes = sizes + + self.oe_embedder = VocabParallelEmbedding(offsets[-1], oe_dim, params_dtype=self._dtype) + # Stacked projections: oe_projection[i] = post_projs[i].weight.T + self.oe_projection = nn.Parameter( + torch.empty(self.num_embedders, oe_dim, self.config.hidden_size, dtype=self._dtype), + requires_grad=False, + ) + + # Precomputed tables for the CUDA n-gram id kernel (ngram_embedding + # _kernels.cu): ne_weights[i][j][delta] = vocab^delta mod ne_mods[i][j], + # ne_mods[i][j] = m + 2*(i*k+j) + 1. Registered as non-persistent buffers + # so they follow the module to the device (not part of the checkpoint). + vocab = self.config.vocab_size + ne_weights = torch.zeros(self.n - 1, self.k, self.n, dtype=torch.int32) + ne_mods = torch.zeros(self.n - 1, self.k, dtype=torch.int32) + for i in range(self.n - 1): + for j in range(self.k): + mod = int(self.m + 2 * (i * self.k + j) + 1) + ne_mods[i, j] = mod + for delta in range(self.n): + ne_weights[i, j, delta] = pow(vocab, delta, mod) + self.register_buffer("ne_weights", ne_weights, persistent=False) + self.register_buffer("ne_mods", ne_mods, persistent=False) + self.register_buffer( + "exclusive_sizes", + torch.tensor(offsets, dtype=torch.int32), + persistent=False, + ) + + def load_weight(self, weight_name: str, loaded_weight: torch.Tensor) -> str: + """Split a per-embedder checkpoint weight into the sharded layout. + + Returns the destination parameter's qualified name (relative to the + enclosing model) so the caller can mark it loaded for completeness + checks. + """ + if "ngram_embeddings.embedders." in weight_name: + index = int(weight_name.split("ngram_embeddings.embedders.")[1].split(".")[0]) + lo, hi = self._offsets[index], self._offsets[index + 1] + assert hi - lo == loaded_weight.shape[0], f"{hi - lo=} {loaded_weight.shape[0]=}" + shard = self.oe_embedder.shard_indices + tp_start, tp_end = shard.org_vocab_start_index, shard.org_vocab_end_index + load_start, load_end = max(lo, tp_start), min(hi, tp_end) + if load_start < load_end: + self.oe_embedder.weight.data[load_start - tp_start : load_end - tp_start] = loaded_weight[ + load_start - lo : load_end - lo + ] + return "ngram_embeddings.oe_embedder.weight" + elif "ngram_embeddings.post_projs." in weight_name: + index = int(weight_name.split("ngram_embeddings.post_projs.")[1].split(".")[0]) + self.oe_projection.data[index].copy_(loaded_weight.t()) + return "ngram_embeddings.oe_projection" + else: + raise AssertionError(f"Unexpected ngram weight: {weight_name}") + + def embed_batched(self, input_ids: torch.Tensor, oe_ids: torch.Tensor) -> torch.Tensor: + """Fused n-gram embedding for a flat batch given precomputed ids. + + Args: + input_ids: ``[num_tokens]`` current token per position. + oe_ids: ``[num_tokens, num_embedders]`` global (offset) n-gram ids, + as produced by the ``ngram_compute_n_gram_ids`` kernel. + Returns: ``[num_tokens, hidden]``. + """ + word = self.word_embeddings(input_ids) # [N, H] + flat = oe_ids.permute(1, 0).contiguous() # [num_embedders, N] + oe = self.oe_embedder(flat) # [num_embedders, N, oe_dim] + proj = torch.bmm(oe, self.oe_projection) # [num_embedders, N, H] + all_h = torch.cat([word.unsqueeze(0), proj], dim=0) # [ne+1, N, H] + return all_h.mean(dim=0) # [N, H] + + +class FlashNgramModel(FlashModel): + """FlashModel whose input embedding is an :class:`NgramEmbedding`.""" + + def __init__(self, *, aphrodite_config: AphroditeConfig, prefix: str = "") -> None: + # Each FlashDecoderLayer is a *dual* layer (2 attentions), so the number + # of decoder layers is ``num_layers``. The ngram HF config sets + # ``num_hidden_layers`` to a multiple of that (attention-module count), + # which FlashModel would otherwise build as too many (dead) layers. + hf = aphrodite_config.model_config.hf_config + num_layers = getattr(hf, "num_layers", None) + if num_layers is not None and hf.num_hidden_layers != num_layers: + hf.num_hidden_layers = num_layers + super().__init__(aphrodite_config=aphrodite_config, prefix=prefix) + if get_pp_group().is_first_rank and uses_ngram_embedding(self.config): + self.ngram_embeddings = NgramEmbedding(self.config, self.embed_tokens) + else: + self.ngram_embeddings = None + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + # Names arrive with the ``model.`` prefix already stripped (routed here + # by AutoWeightsLoader). Split the concatenated/sharded ngram tables and + # stacked projections; delegate everything else to FlashModel. + loaded: set[str] = set() + rest: list[tuple[str, torch.Tensor]] = [] + for name, w in weights: + if self.ngram_embeddings is not None and ( + "ngram_embeddings.embedders." in name or "ngram_embeddings.post_projs." in name + ): + loaded.add(self.ngram_embeddings.load_weight(name, w)) + else: + rest.append((name, w)) + loaded |= super().load_weights(rest) + return loaded + + +class LongcatFlashNgramForCausalLM(nn.Module, SupportsLoRA, SupportsPP): + """LongCat-Flash-Lite for causal LM (MRV2-only, n-gram embedding).""" + + packed_modules_mapping = { + "qkv_proj": ["q_proj", "k_proj", "v_proj"], + "gate_up_proj": ["gate_proj", "up_proj"], + } + + def __init__(self, *, aphrodite_config: AphroditeConfig, prefix: str = "") -> None: + super().__init__() + if not aphrodite_config.use_v2_model_runner: + raise NotImplementedError( + "LongcatFlashNgramForCausalLM (LongCat-Flash-Lite) requires the " + "V2 model runner for its n-gram embedding state; it is selected " + "automatically unless APHRODITE_USE_V2_MODEL_RUNNER=0 is set." + ) + config = FlashConfig(**aphrodite_config.model_config.hf_config.__dict__) + config.intermediate_size = getattr(config, "ffn_hidden_size", config.intermediate_size) + self.config = config + self.quant_config = aphrodite_config.quant_config + + self.model = FlashNgramModel(aphrodite_config=aphrodite_config, prefix=maybe_prefix(prefix, "model")) + if get_pp_group().is_last_rank: + self.lm_head = ParallelLMHead( + config.vocab_size, + config.hidden_size, + quant_config=self.quant_config, + prefix=maybe_prefix(prefix, "lm_head"), + ) + else: + self.lm_head = PPMissingLayer() + self.logits_processor = LogitsProcessor(config.vocab_size) + self.make_empty_intermediate_tensors = self.model.make_empty_intermediate_tensors + + @staticmethod + def get_model_state_cls() -> type["LongcatNgramModelState"]: + return LongcatNgramModelState + + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: + return self.model.embed_input_ids(input_ids) + + def forward( + self, + input_ids: torch.Tensor | None, + positions: torch.Tensor, + intermediate_tensors=None, + inputs_embeds: torch.Tensor | None = None, + ): + # inputs_embeds is produced by LongcatNgramModelState.prepare_inputs. + return self.model(input_ids, positions, intermediate_tensors, inputs_embeds) + + def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor | None: + return self.logits_processor(self.lm_head, hidden_states) + + def get_expert_mapping(self): + return self.model.get_expert_mapping() + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + # AutoWeightsLoader routes ``model.*`` to FlashNgramModel.load_weights + # (which handles the ngram split) and ``lm_head.*`` to the head. MTP + # weights are not part of this model. + loader = AutoWeightsLoader(self, skip_prefixes=["model.mtp."]) + return loader.load_weights(weights) + + +class LongcatNgramModelState(DefaultModelState): + """Per-request n-gram token history for LongCat-Flash-Lite. + + Maintains a small CPU-side per-slot context (last ``n-1`` processed tokens) + and a persistent ``inputs_embeds`` buffer. ``prepare_inputs`` computes the + fused n-gram embedding per request into the buffer, handed to the model + forward as ``inputs_embeds``. + """ + + def __init__(self, aphrodite_config, model, encoder_cache, device) -> None: + super().__init__(aphrodite_config, model, encoder_cache, device) + config = model.config + self.ngram = model.model.ngram_embeddings + self.n = int(config.emb_neighbor_num) + self.ctx_len = self.n - 1 + self.eos_id = int(config.eos_token_id) + + # Per-slot left-context: last ``n-1`` processed tokens, EOS negated. A + # negative entry (incl. the -1 fill) marks a context boundary that stops + # the n-gram walk (matches the kernel's EOS break / fresh-request start). + self.token_context = torch.full((self.max_num_reqs, self.ctx_len), -1, dtype=torch.int32, device=device) + + self._inputs_embeds_buf = torch.zeros( + self.max_num_tokens, + config.hidden_size, + dtype=self.dtype, + device=device, + ) + + def _neg_eos(self, toks: list[int]) -> list[int]: + return [-t if t == self.eos_id else t for t in toks] + + def add_request(self, req_index: int, new_req_data: NewRequestData) -> None: + super().add_request(req_index, new_req_data) # rope positions + # Fresh request -> no left-context (-1 fill). On resume, seed from the + # already-processed token tail. Use prefill_token_ids (full processed + # sequence incl. generated tokens on v2 resume), like DefaultModelState; + # the prompt alone would be too short when resuming after decode. + ctx = [-1] * self.ctx_len + ncomp = new_req_data.num_computed_tokens + toks_src = new_req_data.prefill_token_ids or new_req_data.prompt_token_ids + if ncomp > 0 and toks_src is not None: + lo = max(0, ncomp - self.ctx_len) + toks = self._neg_eos(list(toks_src[lo:ncomp])) + ctx[self.ctx_len - len(toks) :] = toks + self.token_context[req_index] = torch.tensor(ctx, dtype=torch.int32, device=self.token_context.device) + + def prepare_inputs(self, input_batch: InputBatch, req_states: RequestState) -> dict[str, Any]: + model_inputs = super().prepare_inputs(input_batch, req_states) # positions + num_tokens = input_batch.num_tokens + num_padded = input_batch.num_tokens_after_padding + input_ids = input_batch.input_ids[:num_tokens] + embeds = self._inputs_embeds_buf[:num_padded] + + oe_ids = self._compute_oe_ids(input_batch) + embeds[:num_tokens].copy_(self.ngram.embed_batched(input_ids, oe_ids)) + model_inputs["inputs_embeds"] = embeds + return model_inputs + + def prepare_dummy_inputs(self, num_reqs: int, num_tokens: int) -> dict[str, Any]: + # FULL cudagraph replay reads only the captured buffers, so capture must + # reference the same persistent ``inputs_embeds`` buffer prepare_inputs + # re-fills (the base class wires this for multimodal models only). + model_inputs = super().prepare_dummy_inputs(num_reqs, num_tokens) # positions + model_inputs["inputs_embeds"] = self._inputs_embeds_buf[:num_tokens] + return model_inputs + + def _compute_oe_ids(self, input_batch: InputBatch) -> torch.Tensor: + """Batched global n-gram ids ``[num_tokens, num_embedders]``. + + Assembles an ephemeral per-request token table (``[n-1] context ++ + current tokens``, EOS-negated) and runs the ``ngram_compute_n_gram_ids`` + CUDA kernel for the whole batch, then rolls each slot's context forward. + """ + device = self.token_context.device + num_tokens = input_batch.num_tokens + num_reqs = input_batch.num_reqs + ctx_len = self.ctx_len + idx_mapping = input_batch.idx_mapping[:num_reqs].long() + qsl = input_batch.query_start_loc[: num_reqs + 1].to(torch.int32) + cur = input_batch.input_ids[:num_tokens].to(torch.int32) + + cur_neg = torch.where(cur == self.eos_id, -cur, cur) + req_lens = qsl[1:] - qsl[:-1] + max_len = int(req_lens.max().item()) + width = ctx_len + max_len + + # table[r] = [context(n-1) | current tokens | pad(-1)] + table = torch.full((num_reqs, width), -1, dtype=torch.int32, device=device) + table[:, :ctx_len] = self.token_context[idx_mapping] + tok_req = torch.repeat_interleave(torch.arange(num_reqs, device=device), req_lens.long()) + col = ctx_len + (torch.arange(num_tokens, device=device) - qsl[:-1].long()[tok_req]) + table[tok_req, col] = cur_neg + + column_starts = torch.full((num_reqs,), ctx_len, dtype=torch.int32, device=device) + row_indices = torch.arange(num_reqs, dtype=torch.int64, device=device) + n_gram_ids = torch.empty(num_tokens, self.ngram.num_embedders, dtype=torch.int32, device=device) + ops.ngram_compute_n_gram_ids( + self.n, + self.ngram.k, + self.ngram.ne_weights, + self.ngram.ne_mods, + self.ngram.exclusive_sizes, + qsl, + table, + row_indices, + column_starts, + n_gram_ids, + ) + + # Roll context: new context = last n-1 of [context | current] per slot. + gather = req_lens.long().unsqueeze(1) + torch.arange(ctx_len, device=device).unsqueeze(0) + rows = torch.arange(num_reqs, device=device).unsqueeze(1) + self.token_context[idx_mapping] = table[rows, gather] + return n_gram_ids.long() diff --git a/aphrodite/model_executor/models/mistral_large_3.py b/aphrodite/model_executor/models/mistral_large_3.py index 2a4e5b4c80..d4a72ee894 100644 --- a/aphrodite/model_executor/models/mistral_large_3.py +++ b/aphrodite/model_executor/models/mistral_large_3.py @@ -2,60 +2,67 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from collections.abc import Iterable -import regex as re +import regex import torch from aphrodite.model_executor.models.deepseek_v2 import DeepseekV3ForCausalLM +from aphrodite.model_executor.models.utils import AutoWeightsLoader, WeightsMapper class MistralLarge3ForCausalLM(DeepseekV3ForCausalLM): - # fmt: off - remapping = { - r"layers\.(\d+)\.attention_norm\.weight": r"model.layers.\1.input_layernorm.weight", # noqa: E501 - r"layers\.(\d+)\.attention\.wq_a\.(\w+)": r"model.layers.\1.self_attn.q_a_proj.\2", # noqa: E501 - r"layers\.(\d+)\.attention\.q_a_norm\.weight": r"model.layers.\1.self_attn.q_a_layernorm.weight", # noqa: E501 - r"layers\.(\d+)\.attention\.wq_b\.(\w+)": r"model.layers.\1.self_attn.q_b_proj.\2", # noqa: E501 - r"layers\.(\d+)\.attention\.wkv_a_with_mqa\.(\w+)": r"model.layers.\1.self_attn.kv_a_proj_with_mqa.\2", # noqa: E501 - r"layers\.(\d+)\.attention\.kv_a_norm\.weight": r"model.layers.\1.self_attn.kv_a_layernorm.weight", # noqa: E501 - r"layers\.(\d+)\.attention\.wkv_b\.(\w+)": r"model.layers.\1.self_attn.kv_b_proj.\2", # noqa: E501 - r"layers\.(\d+)\.attention\.wo\.(\w+)": r"model.layers.\1.self_attn.o_proj.\2", # noqa: E501 - r"layers\.(\d+)\.ffn_norm\.weight": r"model.layers.\1.post_attention_layernorm.weight", # noqa: E501 - r"layers\.(\d+)\.feed_forward\.w1\.(\w+)": r"model.layers.\1.mlp.gate_proj.\2", # noqa: E501 - r"layers\.(\d+)\.feed_forward\.w2\.(\w+)": r"model.layers.\1.mlp.down_proj.\2", # noqa: E501 - r"layers\.(\d+)\.feed_forward\.w3\.(\w+)": r"model.layers.\1.mlp.up_proj.\2", # noqa: E501 - r"layers\.(\d+)\.gate\.weight": r"model.layers.\1.mlp.gate.weight", # noqa: E501 - r"layers\.(\d+)\.shared_experts\.w1\.(\w+)": r"model.layers.\1.mlp.shared_experts.gate_proj.\2", # noqa: E501 - r"layers\.(\d+)\.shared_experts\.w2\.(\w+)": r"model.layers.\1.mlp.shared_experts.down_proj.\2", # noqa: E501 - r"layers\.(\d+)\.shared_experts\.w3\.(\w+)": r"model.layers.\1.mlp.shared_experts.up_proj.\2", # noqa: E501 - r"layers\.(\d+)\.experts\.(\d+)\.w1\.(\w+)": r"model.layers.\1.mlp.experts.\2.gate_proj.\3", # noqa: E501 - r"layers\.(\d+)\.experts\.(\d+)\.w2\.(\w+)": r"model.layers.\1.mlp.experts.\2.down_proj.\3", # noqa: E501 - r"layers\.(\d+)\.experts\.(\d+)\.w3\.(\w+)": r"model.layers.\1.mlp.experts.\2.up_proj.\3", # noqa: E501 - r"norm\.weight": "model.norm.weight", # noqa: E501 - r"tok_embeddings\.weight": "model.embed_tokens.weight", # noqa: E501 - r"output\.weight": "lm_head.weight", # noqa: E501 - } - # fmt: on + # WeightsMapper applies all matching patterns sequentially (no break on first + # match). This is safe here because every pattern is anchored at both ends + # (\A...\Z) and after substitution the resulting key always starts with + # "model." or "lm_head.", so no later pattern can accidentally match again. + hf_to_aphrodite_mapper = WeightsMapper( + orig_to_new_regex={ # noqa: B950 + regex.compile(r"\Alayers\.(\d+)\.attention_norm\.weight\Z"): (r"model.layers.\1.input_layernorm.weight"), + regex.compile(r"\Alayers\.(\d+)\.attention\.wq_a\.(\w+)\Z"): (r"model.layers.\1.self_attn.q_a_proj.\2"), + regex.compile(r"\Alayers\.(\d+)\.attention\.q_a_norm\.weight\Z"): ( + r"model.layers.\1.self_attn.q_a_layernorm.weight" + ), + regex.compile(r"\Alayers\.(\d+)\.attention\.wq_b\.(\w+)\Z"): (r"model.layers.\1.self_attn.q_b_proj.\2"), + regex.compile(r"\Alayers\.(\d+)\.attention\.wkv_a_with_mqa\.(\w+)\Z"): ( + r"model.layers.\1.self_attn.kv_a_proj_with_mqa.\2" + ), + regex.compile(r"\Alayers\.(\d+)\.attention\.kv_a_norm\.weight\Z"): ( + r"model.layers.\1.self_attn.kv_a_layernorm.weight" + ), + regex.compile(r"\Alayers\.(\d+)\.attention\.wkv_b\.(\w+)\Z"): (r"model.layers.\1.self_attn.kv_b_proj.\2"), + regex.compile(r"\Alayers\.(\d+)\.attention\.wo\.(\w+)\Z"): (r"model.layers.\1.self_attn.o_proj.\2"), + regex.compile(r"\Alayers\.(\d+)\.ffn_norm\.weight\Z"): (r"model.layers.\1.post_attention_layernorm.weight"), + regex.compile(r"\Alayers\.(\d+)\.feed_forward\.w1\.(\w+)\Z"): (r"model.layers.\1.mlp.gate_proj.\2"), + regex.compile(r"\Alayers\.(\d+)\.feed_forward\.w2\.(\w+)\Z"): (r"model.layers.\1.mlp.down_proj.\2"), + regex.compile(r"\Alayers\.(\d+)\.feed_forward\.w3\.(\w+)\Z"): (r"model.layers.\1.mlp.up_proj.\2"), + regex.compile(r"\Alayers\.(\d+)\.gate\.weight\Z"): (r"model.layers.\1.mlp.gate.weight"), + regex.compile(r"\Alayers\.(\d+)\.shared_experts\.w1\.(\w+)\Z"): ( + r"model.layers.\1.mlp.shared_experts.gate_proj.\2" + ), + regex.compile(r"\Alayers\.(\d+)\.shared_experts\.w2\.(\w+)\Z"): ( + r"model.layers.\1.mlp.shared_experts.down_proj.\2" + ), + regex.compile(r"\Alayers\.(\d+)\.shared_experts\.w3\.(\w+)\Z"): ( + r"model.layers.\1.mlp.shared_experts.up_proj.\2" + ), + regex.compile(r"\Alayers\.(\d+)\.experts\.(\d+)\.w1\.(\w+)\Z"): ( + r"model.layers.\1.mlp.experts.\2.gate_proj.\3" + ), + regex.compile(r"\Alayers\.(\d+)\.experts\.(\d+)\.w2\.(\w+)\Z"): ( + r"model.layers.\1.mlp.experts.\2.down_proj.\3" + ), + regex.compile(r"\Alayers\.(\d+)\.experts\.(\d+)\.w3\.(\w+)\Z"): ( + r"model.layers.\1.mlp.experts.\2.up_proj.\3" + ), + regex.compile(r"\Anorm\.weight\Z"): "model.norm.weight", + regex.compile(r"\Atok_embeddings\.weight\Z"): "model.embed_tokens.weight", + regex.compile(r"\Aoutput\.weight\Z"): "lm_head.weight", + }, + orig_to_new_suffix={ + ".qscale_act": ".input_scale", + ".qscale_weight": ".weight_scale", + }, + ) def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - return super().load_weights(map(self._remap_mistral_to_ds, weights)) - - def _remap_mistral_to_ds(self, weight: tuple[str, torch.Tensor]) -> tuple[str, torch.Tensor]: - """Remap Mistral parameters to DeepseekV2 parameters.""" - name, loaded_weight = weight - - for k, v in self.remapping.items(): - match = re.fullmatch(k, name) - if match: - name = re.sub(k, v, name) - break - else: - raise ValueError(f"Cannot remap {name}") - - # Remapping scale names. We could do this in the regex above but it - # would triple the number of lines for most layers. - if name.endswith(".qscale_act"): - name = re.sub(r"\.qscale_act$", ".input_scale", name) - elif name.endswith(".qscale_weight"): - name = re.sub(r"\.qscale_weight$", ".weight_scale", name) - - return name, loaded_weight + loader = AutoWeightsLoader(self) + return loader.load_weights(weights, mapper=self.hf_to_aphrodite_mapper) diff --git a/aphrodite/model_executor/models/mistral_large_3_eagle.py b/aphrodite/model_executor/models/mistral_large_3_eagle.py index 9219d2a7ab..438db61853 100644 --- a/aphrodite/model_executor/models/mistral_large_3_eagle.py +++ b/aphrodite/model_executor/models/mistral_large_3_eagle.py @@ -5,6 +5,7 @@ from collections.abc import Iterable from functools import partial +import regex import torch import torch.nn as nn @@ -22,7 +23,7 @@ from aphrodite.model_executor.models.mistral_large_3 import MistralLarge3ForCausalLM from .interfaces import SupportsMultiModal -from .utils import make_empty_intermediate_tensors_factory, maybe_prefix +from .utils import WeightsMapper, make_empty_intermediate_tensors_factory, maybe_prefix logger = init_logger(__name__) @@ -99,11 +100,13 @@ def forward( class EagleMistralLarge3ForCausalLM(MistralLarge3ForCausalLM): - remapping = MistralLarge3ForCausalLM.remapping | { - r"eagle_linear\.weight": r"model.fc.weight", - r"eagle_linear\.qscale_act": r"model.fc.input_scale", - r"eagle_linear\.qscale_weight": r"model.fc.weight_scale", - } + hf_to_aphrodite_mapper = MistralLarge3ForCausalLM.hf_to_aphrodite_mapper | WeightsMapper( + orig_to_new_regex={ + regex.compile(r"\Aeagle_linear\.weight\Z"): r"model.fc.weight", + regex.compile(r"\Aeagle_linear\.qscale_act\Z"): r"model.fc.input_scale", + regex.compile(r"\Aeagle_linear\.qscale_weight\Z"): r"model.fc.weight_scale", + }, + ) def __init__(self, *, aphrodite_config: AphroditeConfig, prefix: str = ""): target_layer_num = aphrodite_config.model_config.get_num_layers(aphrodite_config.parallel_config) diff --git a/aphrodite/model_executor/models/modernbert.py b/aphrodite/model_executor/models/modernbert.py index 3de1b685d9..69ce41f598 100644 --- a/aphrodite/model_executor/models/modernbert.py +++ b/aphrodite/model_executor/models/modernbert.py @@ -220,7 +220,13 @@ def forward( @support_torch_compile @default_pooling_type(seq_pooling_type="CLS") class ModernBertModel(nn.Module): - hf_to_aphrodite_mapper = WeightsMapper(orig_to_new_prefix={"layers.": "encoder_layer.layers."}) + hf_to_aphrodite_mapper = WeightsMapper( + orig_to_new_prefix={ + "model.layers.": "encoder_layer.layers.", + "layers.": "encoder_layer.layers.", + "model.": "", + } + ) def __init__( self, @@ -244,6 +250,8 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: for name, loaded_weight in weights: if name.endswith(".bias") and name not in params_dict: continue + if name not in params_dict: + continue param = params_dict[name] weight_loader = getattr(param, "weight_loader", default_weight_loader) weight_loader(param, loaded_weight) diff --git a/aphrodite/model_executor/models/olmo3.py b/aphrodite/model_executor/models/olmo3.py index f2bd1027d0..f993659db7 100644 --- a/aphrodite/model_executor/models/olmo3.py +++ b/aphrodite/model_executor/models/olmo3.py @@ -135,11 +135,9 @@ def __init__(self, *, aphrodite_config: AphroditeConfig, prefix: str = ""): ) # Rotary embeddings. Rope scaling is only applied on full attention layers. - if sliding_window is None: - rope_parameters = self.config.rope_parameters - else: - rope_theta = self.config.rope_parameters["rope_theta"] - rope_parameters = {"rope_type": "default", "rope_theta": rope_theta} + rope_parameters = self.config.rope_parameters + attn_type = "full_attention" if sliding_window is None else "sliding_attention" + rope_parameters = rope_parameters.get(attn_type, rope_parameters) self.rotary_emb = get_rope( self.head_dim, max_position=self.max_position_embeddings, diff --git a/aphrodite/model_executor/models/qwen3_5.py b/aphrodite/model_executor/models/qwen3_5.py index 6b598813c1..1bfbca99c2 100644 --- a/aphrodite/model_executor/models/qwen3_5.py +++ b/aphrodite/model_executor/models/qwen3_5.py @@ -122,10 +122,15 @@ def __init__( config = aphrodite_config.model_config.hf_text_config model_config = aphrodite_config.model_config cache_config = aphrodite_config.cache_config + parallel_config = aphrodite_config.parallel_config quant_config = aphrodite_config.quant_config self.layer_type = layer_type self.layer_idx = extract_layer_index(prefix) + is_moe_layer = config.model_type == "qwen3_5_moe_text" + self.use_attn_reduce_scatter_for_moe = ( + parallel_config.use_sequence_parallel_moe and parallel_config.pipeline_parallel_size == 1 and is_moe_layer + ) if self.layer_type == "linear_attention": self.linear_attn = QwenGatedDeltaNetAttention( @@ -133,6 +138,7 @@ def __init__( aphrodite_config=aphrodite_config, prefix=f"{prefix}.linear_attn", gqa_interleaved_layout=False, + reduce_results=not self.use_attn_reduce_scatter_for_moe, ) elif self.layer_type == "full_attention": self.self_attn = Qwen3NextAttention( @@ -140,6 +146,7 @@ def __init__( model_config=model_config, cache_config=cache_config, quant_config=quant_config, + reduce_results=not self.use_attn_reduce_scatter_for_moe, prefix=f"{prefix}.self_attn", ) else: diff --git a/aphrodite/model_executor/models/qwen3_dflash.py b/aphrodite/model_executor/models/qwen3_dflash.py index 1b2ad56cf4..9194ab29a0 100644 --- a/aphrodite/model_executor/models/qwen3_dflash.py +++ b/aphrodite/model_executor/models/qwen3_dflash.py @@ -729,6 +729,14 @@ def combine_hidden_states( needs_squeeze = hidden_states.dim() == 1 if needs_squeeze: hidden_states = hidden_states.unsqueeze(0) + expected = self.model.fc.input_size + if hidden_states.shape[-1] != expected: + raise ValueError( + f"DFlash drafter expects {expected} concatenated aux hidden " + f"features but received {hidden_states.shape[-1]}. This usually " + "means the draft model's target_layer_ids reference layers that " + "do not exist in the target model (incompatible draft/target pair)." + ) result = self.model.fc(hidden_states) if needs_squeeze: result = result.squeeze(0) diff --git a/aphrodite/model_executor/models/qwen3_next.py b/aphrodite/model_executor/models/qwen3_next.py index 8035a487cf..9e1e2f0cf5 100644 --- a/aphrodite/model_executor/models/qwen3_next.py +++ b/aphrodite/model_executor/models/qwen3_next.py @@ -16,6 +16,7 @@ get_pp_group, get_tensor_model_parallel_world_size, tensor_model_parallel_all_gather, + tensor_model_parallel_reduce_scatter, ) from aphrodite.logger import init_logger from aphrodite.model_executor.layers.attention import Attention @@ -187,13 +188,17 @@ def __init__(self, aphrodite_config: AphroditeConfig, prefix: str = ""): shared_expert_gate=self.shared_expert_gate if self.shared_expert is None else None, ) - def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + def forward( + self, + hidden_states: torch.Tensor, + already_sequence_parallel: bool = False, + ) -> torch.Tensor: # NOTE: hidden_states can have either 1D or 2D shape. orig_shape = hidden_states.shape num_tokens, hidden_dim = hidden_states.shape hidden_states = hidden_states.view(-1, hidden_dim) - if self.is_sequence_parallel: + if self.is_sequence_parallel and not already_sequence_parallel: hidden_states = sequence_parallel_chunk(hidden_states) if self.experts.is_internal_router: @@ -204,7 +209,7 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: router_logits, _ = self.gate(hidden_states) final_hidden_states = self.experts(hidden_states=hidden_states, router_logits=router_logits) - if self.is_sequence_parallel: + if self.is_sequence_parallel and not already_sequence_parallel: final_hidden_states = tensor_model_parallel_all_gather(final_hidden_states, 0) final_hidden_states = final_hidden_states[:num_tokens] @@ -218,6 +223,7 @@ def __init__( model_config: ModelConfig | None = None, cache_config: CacheConfig | None = None, quant_config: QuantizationConfig | None = None, + reduce_results: bool = True, prefix: str = "", ) -> None: super().__init__() @@ -258,6 +264,7 @@ def __init__( self.total_num_heads * self.head_dim, config.hidden_size, bias=False, + reduce_results=reduce_results, quant_config=quant_config, prefix=f"{prefix}.o_proj", ) @@ -356,15 +363,15 @@ def _project_qkv_gate( def forward( self, positions: torch.Tensor, - output: torch.Tensor, hidden_states: torch.Tensor, - ): + ) -> torch.Tensor: qkv, _ = self.qkv_proj(hidden_states) q, k, v, gate = self._project_qkv_gate(qkv, positions) attn_output = self.attn(q, k, v) if gate is not None: attn_output = attn_output * torch.sigmoid(gate) - output[:], _ = self.o_proj(attn_output) + output, _ = self.o_proj(attn_output) + return output class Qwen3NextDecoderLayer(nn.Module): @@ -380,16 +387,26 @@ def __init__( model_config = aphrodite_config.model_config cache_config = aphrodite_config.cache_config quant_config = aphrodite_config.quant_config + parallel_config = aphrodite_config.parallel_config self.layer_type = layer_type self.layer_idx = extract_layer_index(prefix) + mlp_only_layers = [] if not hasattr(config, "mlp_only_layers") else config.mlp_only_layers + is_moe_layer = (self.layer_idx not in mlp_only_layers) and ( + config.num_experts > 0 and (self.layer_idx + 1) % config.decoder_sparse_step == 0 + ) + self.use_attn_reduce_scatter_for_moe = ( + parallel_config.use_sequence_parallel_moe and parallel_config.pipeline_parallel_size == 1 and is_moe_layer + ) + if self.layer_type == "linear_attention": self.linear_attn = QwenGatedDeltaNetAttention( config, aphrodite_config=aphrodite_config, prefix=f"{prefix}.linear_attn", gqa_interleaved_layout=True, + reduce_results=not self.use_attn_reduce_scatter_for_moe, ) elif self.layer_type == "full_attention": self.self_attn = Qwen3NextAttention( @@ -397,15 +414,13 @@ def __init__( model_config=model_config, cache_config=cache_config, quant_config=quant_config, + reduce_results=not self.use_attn_reduce_scatter_for_moe, prefix=f"{prefix}.self_attn", ) else: raise ValueError(f"Invalid layer_type {self.layer_type}") - mlp_only_layers = [] if not hasattr(config, "mlp_only_layers") else config.mlp_only_layers - if (self.layer_idx not in mlp_only_layers) and ( - config.num_experts > 0 and (self.layer_idx + 1) % config.decoder_sparse_step == 0 - ): + if is_moe_layer: self.mlp = Qwen3NextSparseMoeBlock( aphrodite_config=aphrodite_config, prefix=f"{prefix}.mlp", @@ -446,27 +461,30 @@ def forward( positions: torch.Tensor = None, **kwargs: object, ): + full_num_tokens = positions.shape[-1] + input_is_sequence_parallel = ( + self.use_attn_reduce_scatter_for_moe and residual is not None and hidden_states.shape[0] != full_num_tokens + ) + if residual is None: residual = hidden_states hidden_states = self.input_layernorm(hidden_states) else: hidden_states, residual = self.input_layernorm(hidden_states, residual) - self_attention_output = torch.empty_like(hidden_states) + if input_is_sequence_parallel: + hidden_states = tensor_model_parallel_all_gather(hidden_states, 0) + hidden_states = hidden_states[:full_num_tokens] + if self.layer_type == "linear_attention": - self.linear_attn( - hidden_states=hidden_states, - output=self_attention_output, - ) + hidden_states = self.linear_attn(hidden_states=hidden_states) elif self.layer_type == "full_attention": - self.self_attn( + hidden_states = self.self_attn( hidden_states=hidden_states, - output=self_attention_output, positions=positions, ) else: raise ValueError("Invalid layer_type") - hidden_states = self_attention_output if self.layer_scale: if len(hidden_states.shape) == 2: @@ -474,9 +492,23 @@ def forward( else: hidden_states = hidden_states * (self.attn_layer_scale.to(hidden_states.dtype) + 1) + if self.use_attn_reduce_scatter_for_moe: + tp_world_size = get_tensor_model_parallel_world_size() + sp_pad = (-hidden_states.shape[0]) % tp_world_size + hidden_states = torch.nn.functional.pad(hidden_states, (0, 0, 0, sp_pad)) + hidden_states = tensor_model_parallel_reduce_scatter(hidden_states, 0) + if not input_is_sequence_parallel: + residual = sequence_parallel_chunk(residual) + # Fully Connected hidden_states, residual = self.post_attention_layernorm(hidden_states, residual) - hidden_states = self.mlp(hidden_states) + if self.use_attn_reduce_scatter_for_moe: + hidden_states = self.mlp( + hidden_states, + already_sequence_parallel=True, + ) + else: + hidden_states = self.mlp(hidden_states) if self.layer_scale: if len(hidden_states.shape) == 2: @@ -490,6 +522,24 @@ def forward( return hidden_states, residual +def _all_gather_hidden_and_residual( + hidden_states: torch.Tensor, + residual: torch.Tensor | None, + full_num_tokens: int, + hidden_size: int, +) -> tuple[torch.Tensor, torch.Tensor | None]: + if residual is None: + hidden_states = tensor_model_parallel_all_gather(hidden_states, 0) + hidden_states = hidden_states[:full_num_tokens] + return hidden_states, None + + combined_states = torch.cat([hidden_states, residual], dim=-1) + combined_states = tensor_model_parallel_all_gather(combined_states, 0) + combined_states = combined_states[:full_num_tokens] + hidden_states, residual = combined_states.split([hidden_size, hidden_size], dim=-1) + return hidden_states, residual + + @support_torch_compile class Qwen3NextModel(nn.Module, EagleModelMixin): hf_to_aphrodite_mapper = WeightsMapper( @@ -542,6 +592,8 @@ def get_layer(prefix: str): else: self.norm = PPMissingLayer() + self.aux_hidden_state_layers: tuple[int, ...] = () + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: return self.embed_tokens(input_ids) @@ -563,20 +615,42 @@ def forward( hidden_states = intermediate_tensors["hidden_states"] residual = intermediate_tensors["residual"] + full_num_tokens = positions.shape[-1] aux_hidden_states = self._maybe_add_hidden_state([], 0, hidden_states, residual) for layer_idx, layer in enumerate( islice(self.layers, self.start_layer, self.end_layer), start=self.start_layer, ): + if hidden_states.shape[0] != full_num_tokens and not layer.use_attn_reduce_scatter_for_moe: + hidden_states, residual = _all_gather_hidden_and_residual( + hidden_states, + residual, + full_num_tokens, + self.config.hidden_size, + ) hidden_states, residual = layer( positions=positions, hidden_states=hidden_states, residual=residual, ) + if (layer_idx + 1) in self.aux_hidden_state_layers and hidden_states.shape[0] != full_num_tokens: + hidden_states, residual = _all_gather_hidden_and_residual( + hidden_states, + residual, + full_num_tokens, + self.config.hidden_size, + ) self._maybe_add_hidden_state(aux_hidden_states, layer_idx + 1, hidden_states, residual) if not get_pp_group().is_last_rank: return IntermediateTensors({"hidden_states": hidden_states, "residual": residual}) + if hidden_states.shape[0] != full_num_tokens: + hidden_states, residual = _all_gather_hidden_and_residual( + hidden_states, + residual, + full_num_tokens, + self.config.hidden_size, + ) hidden_states, _ = self.norm(hidden_states, residual) if aux_hidden_states: return hidden_states, aux_hidden_states diff --git a/aphrodite/model_executor/models/registry.py b/aphrodite/model_executor/models/registry.py index e6d58e6532..d894099cb1 100644 --- a/aphrodite/model_executor/models/registry.py +++ b/aphrodite/model_executor/models/registry.py @@ -146,6 +146,10 @@ # For decapoda-research/llama-* "LLaMAForCausalLM": ("llama", "LlamaForCausalLM"), "LongcatFlashForCausalLM": ("longcat_flash", "LongcatFlashForCausalLM"), + "LongcatFlashNgramForCausalLM": ( + "longcat_flash_ngram", + "LongcatFlashNgramForCausalLM", + ), "MambaForCausalLM": ("mamba", "MambaForCausalLM"), "Mamba2ForCausalLM": ("mamba2", "Mamba2ForCausalLM"), "MellumForCausalLM": ("mellum", "MellumForCausalLM"), diff --git a/aphrodite/model_executor/warmup/deep_gemm_warmup.py b/aphrodite/model_executor/warmup/deep_gemm_warmup.py index 8659e85d82..68bb2e2e36 100644 --- a/aphrodite/model_executor/warmup/deep_gemm_warmup.py +++ b/aphrodite/model_executor/warmup/deep_gemm_warmup.py @@ -128,9 +128,6 @@ def _fp8_linear_may_use_deep_gemm(module: torch.nn.Module) -> bool: Return True if the input module/layer could be processed with DeepGEMM. """ - # FIXME: this logic is brittle and incorrect - since we - # could use DeepGEMM with for than just Fp8LinearMethod - block_size = get_mk_alignment_for_contiguous_layout()[0] if not ( isinstance(module, LinearBase) and isinstance(module.quant_method, Fp8LinearMethod) @@ -146,6 +143,8 @@ def _fp8_linear_may_use_deep_gemm(module: torch.nn.Module) -> bool: ): return False + block_size = get_mk_alignment_for_contiguous_layout()[0] + w, _, block_sizes = _extract_data_from_linear_base_module(module) return ( block_sizes == get_mk_alignment_for_contiguous_layout() diff --git a/aphrodite/model_executor/warmup/kernel_warmup.py b/aphrodite/model_executor/warmup/kernel_warmup.py index 3a3d1bd1b2..3ec3b52d75 100644 --- a/aphrodite/model_executor/warmup/kernel_warmup.py +++ b/aphrodite/model_executor/warmup/kernel_warmup.py @@ -120,6 +120,25 @@ def _is_flashinfer_backend(backend): cutedsl_warmup() +def _flashinfer_autotune_skip_ops(runner: "GPUModelRunner") -> set[str] | None: + if envs.APHRODITE_FLASHINFER_AUTOTUNE_SKIP_OPS is not None: + return set(envs.APHRODITE_FLASHINFER_AUTOTUNE_SKIP_OPS) or None + + from aphrodite.model_executor.kernels.linear import ( + FlashInferCuteDslNvFp4LinearKernel, + ) + + for module in runner.get_model().modules(): + for holder_name in ("quant_method", "scheme"): + kernel = getattr(getattr(module, holder_name, None), "kernel", None) + # CuTe-DSL mm_fp4 tuning JIT-compiles every tactic and its + # fallback is already the heuristic; all mm_fp4 backends share + # the "fp4_gemm" op name, so skip only when cute-dsl is selected. + if isinstance(kernel, FlashInferCuteDslNvFp4LinearKernel): + return {"fp4_gemm"} + return None + + def flashinfer_autotune(runner: "GPUModelRunner") -> None: """ Autotune FlashInfer operations. @@ -136,6 +155,15 @@ def flashinfer_autotune(runner: "GPUModelRunner") -> None: import aphrodite.utils.flashinfer as fi_utils from aphrodite.distributed.parallel_state import get_world_group + autotune_kwargs: dict[str, object] = {} + skip_ops = _flashinfer_autotune_skip_ops(runner) + if skip_ops: + logger.info( + "Skipping FlashInfer autotuning for ops %s", + sorted(skip_ops), + ) + autotune_kwargs["skip_ops"] = skip_ops + use_persistent_cache = True # When distributed, tune on every rank so the collectives stay synchronized. @@ -143,7 +171,7 @@ def flashinfer_autotune(runner: "GPUModelRunner") -> None: use_persistent_cache = False if not use_persistent_cache: - with torch.inference_mode(), fi_utils.autotune(): + with torch.inference_mode(), fi_utils.autotune(**autotune_kwargs): runner._dummy_run( num_tokens=runner.scheduler_config.max_num_batched_tokens, skip_eplb=True, @@ -171,7 +199,7 @@ def flashinfer_autotune(runner: "GPUModelRunner") -> None: with torch.inference_mode(): if is_leader: - with fi_utils.autotune(tune_mode=True, cache=str(cache_path)): + with fi_utils.autotune(tune_mode=True, cache=str(cache_path), **autotune_kwargs): runner._dummy_run(**dummy_run_kwargs) else: runner._dummy_run(**dummy_run_kwargs) diff --git a/aphrodite/models/deepseek_v4/__init__.py b/aphrodite/models/deepseek_v4/__init__.py index 892f212301..bfac20b2ba 100644 --- a/aphrodite/models/deepseek_v4/__init__.py +++ b/aphrodite/models/deepseek_v4/__init__.py @@ -7,6 +7,8 @@ classes used by the model registry and quantization config lookup. """ +from typing import TYPE_CHECKING + from aphrodite.platforms import current_platform from .quant_config import DeepseekV4FP8Config @@ -14,17 +16,21 @@ # Pick the per-platform implementation. The NVIDIA branch is the static # default that mypy sees; the ROCm/XPU branches override at runtime and are # kept type-compatible via ``# type: ignore[assignment]``. -if current_platform.is_rocm(): +if TYPE_CHECKING: + from .nvidia.dspark import DSparkDeepseekV4ForCausalLM + from .nvidia.model import DeepseekV4ForCausalLM + from .nvidia.mtp import DeepSeekV4MTP +elif current_platform.is_rocm(): + from .amd.dspark import ( # type: ignore[assignment] + DSparkDeepseekV4ForCausalLM, + ) from .amd.model import DeepseekV4ForCausalLM from .amd.mtp import DeepSeekV4MTP - - # DSpark is NVIDIA-only for now. - DSparkDeepseekV4ForCausalLM = None # type: ignore[assignment] elif current_platform.is_xpu(): from .xpu.model import DeepseekV4ForCausalLM # type: ignore[assignment] from .xpu.mtp import DeepSeekV4MTP # type: ignore[assignment] - DSparkDeepseekV4ForCausalLM = None # type: ignore[assignment] + DSparkDeepseekV4ForCausalLM = None # type: ignore[assignment, misc] else: from .nvidia.dspark import ( # type: ignore[assignment] DSparkDeepseekV4ForCausalLM, diff --git a/aphrodite/models/deepseek_v4/amd/dspark.py b/aphrodite/models/deepseek_v4/amd/dspark.py new file mode 100644 index 0000000000..e44218c55b --- /dev/null +++ b/aphrodite/models/deepseek_v4/amd/dspark.py @@ -0,0 +1,459 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""DSpark draft model for DeepSeek-V4 on ROCm/AMD (gfx950). + +ROCm port of ``nvidia/dspark.py``. Follows the same nvidia->amd recipe used for +``amd/mtp.py``: + + * import ``DeepseekV4DecoderLayer`` from the AMD ``.model`` (aiter/triton + attention + MHC CustomOp path) instead of the nvidia one; + * route the MHC head through the ``HCHeadOp`` CustomOp dispatcher (aiter / + tilelang / triton / torch) instead of calling the tilelang kernels directly, + and gate the trailing ``mhc_post`` on ``use_fused_mhc`` (False on the aiter + path, where the decoder layer already applies hc_post in-layer); + * drop the mega-MoE weight path (``make_deepseek_v4_expert_params_mapping`` / + ``use_mega_moe`` / ``finalize_mega_moe_weights`` do not exist in amd/model.py). + +Everything else - the semi-autoregressive drafting hooks, the Markov head, the +sliding-window context-KV insert, and the checkpoint ``mtp.*`` weight remap - is +pure torch / Triton and shared with the nvidia implementation unchanged. +""" + +from collections.abc import Iterable + +import regex as re +import torch +import torch.nn as nn + +from aphrodite.config import AphroditeConfig, get_current_aphrodite_config +from aphrodite.distributed import ( + get_tensor_model_parallel_rank, + get_tensor_model_parallel_world_size, +) +from aphrodite.logger import init_logger +from aphrodite.model_executor.layers.fused_moe import ( + fused_moe_make_expert_params_mapping, +) +from aphrodite.model_executor.layers.layernorm import RMSNorm +from aphrodite.model_executor.layers.linear import ReplicatedLinear +from aphrodite.model_executor.layers.logits_processor import LogitsProcessor +from aphrodite.model_executor.layers.mhc import HCHeadOp +from aphrodite.model_executor.layers.vocab_parallel_embedding import ( + ParallelLMHead, + VocabParallelEmbedding, +) +from aphrodite.model_executor.model_loader.weight_utils import default_weight_loader +from aphrodite.model_executor.models.qwen3_dspark import ( + DSparkMarkovHead, +) +from aphrodite.model_executor.models.utils import maybe_prefix + +from .model import DeepseekV4DecoderLayer + +logger = init_logger(__name__) + +# MoE expert scale suffix differs by expert dtype (mirrors deepseek_v4 loaders): +# fp4 experts register ``.weight_scale``; block-fp8 experts ``.weight_scale_inv``. +_EXPERT_SCALE_RE = re.compile(r"\.experts\.\d+\.w[123]\.scale$") + + +class DSparkDeepseekV4Model(nn.Module): + def __init__(self, *, aphrodite_config: AphroditeConfig, prefix: str = "") -> None: + super().__init__() + assert aphrodite_config.speculative_config is not None + config = aphrodite_config.speculative_config.draft_model_config.hf_config + self.config = config + self.hidden_size = config.hidden_size + self.hc_mult = config.hc_mult + self.hc_eps = config.hc_eps + self.rms_norm_eps = config.rms_norm_eps + self.num_hidden_layers = config.num_hidden_layers + self.target_layer_ids = tuple(config.dspark_target_layer_ids) + + self.num_dspark_layers = getattr(config, "n_mtp_layers", None) or 3 + + # Shared with the target (aliased by the speculator's loading utility). + self.embed_tokens = VocabParallelEmbedding( + config.vocab_size, + config.hidden_size, + prefix=maybe_prefix(prefix, "embed_tokens"), + ) + + self.main_proj = ReplicatedLinear( + config.hidden_size * len(self.target_layer_ids), + config.hidden_size, + bias=False, + return_bias=False, + quant_config=aphrodite_config.quant_config, + prefix=maybe_prefix(prefix, "main_proj"), + ) + self.main_norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + + current_aphrodite_config = get_current_aphrodite_config() + self.layers = nn.ModuleList( + [ + DeepseekV4DecoderLayer( + current_aphrodite_config, + prefix=maybe_prefix(prefix, f"layers.{self.num_hidden_layers + i}"), + ) + for i in range(self.num_dspark_layers) + ] + ) + + # Heads: final norm + hc_head, and the Markov head + # Loaded from the "final" MTP layer weights (mtp.*) in the target checkpoint + self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + hc_dim = self.hc_mult * config.hidden_size + self.hc_head_fn = nn.Parameter( + torch.empty(self.hc_mult, hc_dim, dtype=torch.float32), + requires_grad=False, + ) + self.hc_head_base = nn.Parameter(torch.empty(self.hc_mult, dtype=torch.float32), requires_grad=False) + self.hc_head_scale = nn.Parameter(torch.empty(1, dtype=torch.float32), requires_grad=False) + draft_vocab_size = getattr(config, "draft_vocab_size", None) or config.vocab_size + self.markov_head = DSparkMarkovHead( + config.vocab_size, + draft_vocab_size, + config.dspark_markov_rank, + prefix=maybe_prefix(prefix, "markov_head"), + ) + + # MHC head CustomOp dispatcher (aiter / tilelang / triton / torch), + # replacing the direct nvidia tilelang kernel call. + self.hc_head_op = HCHeadOp() + + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: + return self.embed_tokens(input_ids) + + def combine_hidden_states(self, aux_hidden_states: torch.Tensor) -> torch.Tensor: + """main_x = main_norm(main_proj(concat of target aux hidden states)). + + ``aux_hidden_states`` is [T, hidden_size * len(target_layer_ids)]. + """ + return self.main_norm(self.main_proj(aux_hidden_states)) + + @torch.inference_mode() + def precompute_and_store_context_kv( + self, + main_x: torch.Tensor, + context_positions: torch.Tensor, + context_slot_mappings: list[torch.Tensor | None] | None = None, + ) -> None: + """Insert the sliding-window context KV for every draft layer. + + Mirrors the reference DSparkAttention: each layer derives its context KV + from the SAME projected target hidden ``main_x``, via that layer's own + ``wkv`` + ``kv_norm`` + RoPE + quant, then writes it at the + layer's context slots. + + ``context_slot_mappings`` is a per-layer list (each entry is the context + slot mapping for that layer's kv-cache group, since the hybrid manager may + place draft layers in different groups). ``None`` (or a ``None`` entry) + runs the projection to reserve workspace but writes nothing (profiling). + """ + for i, layer in enumerate(self.layers): + slot_mapping = None if context_slot_mappings is None else context_slot_mappings[i] + attn = layer.attn + # Optimized DSV4 MLA path: wkv part of the fused wq_a|wkv projection + # (q_lora part discarded), then RoPE/quant/insert via the fused op. + qr_kv, _ = attn.fused_wqa_wkv(main_x) + kv = qr_kv[..., attn.q_lora_rank :] + kv = attn.kv_norm(kv) + if slot_mapping is None: + continue + _insert_context_kv(attn, kv, context_positions, slot_mapping) + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + inputs_embeds: torch.Tensor | None = None, + ) -> torch.Tensor: + if inputs_embeds is None: + inputs_embeds = self.embed_input_ids(input_ids) + # Expand to hc_mult copies for hyper-connections ([T, H] -> [T, hc, H]). + hidden_states = inputs_embeds.unsqueeze(-2).repeat(1, self.hc_mult, 1) + + residual = post_mix = res_mix = None + for layer in self.layers: + hidden_states, residual, post_mix, res_mix = layer( + hidden_states, + positions, + input_ids, + post_mix, + res_mix, + residual, + ) + if self.layers[-1].use_fused_mhc: + hidden_states = self.layers[-1].hc_post(hidden_states, residual, post_mix, res_mix) + # hc_head reduces the hc copies; return the PRE-norm head hidden + hidden_states = self.hc_head_op( + hidden_states, + self.hc_head_fn, + self.hc_head_scale, + self.hc_head_base, + self.rms_norm_eps, + self.hc_eps, + ) + return hidden_states + + +def _insert_context_kv( + attn: nn.Module, + kv: torch.Tensor, + positions: torch.Tensor, + slot_mapping: torch.Tensor, +) -> None: + """RoPE + quant + paged-cache insert of (already kv_norm'd) context KV. + + Reuses the DSV4 fused insert ops (which also process a query; we pass a dummy + query and discard it, since context tokens have no query). Mirrors + ``DeepseekV4Attention._fused_qnorm_rope_kv_insert``. + """ + swa_cache = attn.swa_cache_layer.kv_cache + block_size = attn.swa_cache_layer.block_size + cos_sin_cache = attn.rotary_emb.cos_sin_cache + cache_dtype = swa_cache.dtype + n_ctx = kv.shape[0] + dummy_q = torch.zeros( + (n_ctx, attn.n_local_heads, attn.head_dim), + dtype=kv.dtype, + device=kv.device, + ) + if cache_dtype == torch.uint8: + # fp8_ds_mla UE8M0 paged layout + swa_2d = swa_cache.view(swa_cache.shape[0], -1) + torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert( + dummy_q, + kv, + swa_2d, + slot_mapping, + positions, + cos_sin_cache, + attn.padded_heads, + attn.eps, + block_size, + ) + elif cache_dtype == torch.bfloat16: + swa_3d = swa_cache.view(-1, block_size, attn.head_dim) + torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_bf16_insert( + dummy_q, + kv, + swa_3d, + slot_mapping, + positions, + cos_sin_cache, + attn.eps, + block_size, + ) + else: # per-tensor fp8 (torch.float8_e4m3fn) + # TODO(ben): double-check if this is being dispatched correctly for FI backend + swa_3d = swa_cache.view(-1, block_size, attn.head_dim) + dummy_q_fp8 = torch.zeros_like(dummy_q, dtype=torch.float8_e4m3fn) + torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_fp8_insert( + dummy_q, + kv, + dummy_q_fp8, + swa_3d, + slot_mapping, + positions, + cos_sin_cache, + attn._flashinfer_fp8_kv_scale, + attn._flashinfer_fp8_q_scale_inv, + attn.eps, + block_size, + ) + + +class DSparkDeepseekV4ForCausalLM(nn.Module): + # Draft weights ship in the target checkpoint (mtp.*) without embed/head, so + # load_dspark_model always aliases the target's. + has_own_embed_tokens = False + has_own_lm_head = False + # Full-vocab draft: draft ids are target ids, no remapping needed. + draft_id_to_target_id = None + + def __init__(self, *, aphrodite_config: AphroditeConfig, prefix: str = "") -> None: + super().__init__() + assert aphrodite_config.speculative_config is not None + self.draft_model_config = aphrodite_config.speculative_config.draft_model_config + self.config = self.draft_model_config.hf_config + self.model = DSparkDeepseekV4Model(aphrodite_config=aphrodite_config, prefix=maybe_prefix(prefix, "model")) + # Shared with the target (aliased by the speculator's load utility). + self.lm_head = ParallelLMHead( + self.config.vocab_size, + self.config.hidden_size, + prefix=maybe_prefix(prefix, "lm_head"), + ) + self.logits_processor = LogitsProcessor(self.config.vocab_size) + + # --- Hooks used by the speculator ------------------------------------- + + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: + return self.model.embed_input_ids(input_ids) + + def combine_hidden_states(self, aux_hidden_states: torch.Tensor) -> torch.Tensor: + return self.model.combine_hidden_states(aux_hidden_states) + + def get_draft_kv_cache_layer_names(self) -> list[str]: + # DSV4 MLA path: each draft layer's sliding-window cache is a separate + # layer, named by its prefix. + return [layer.attn.swa_cache_layer.prefix for layer in self.model.layers] + + def precompute_and_store_context_kv( + self, + context_states: torch.Tensor, + context_positions: torch.Tensor, + context_slot_mappings: list[torch.Tensor | None] | None = None, + ) -> None: + self.model.precompute_and_store_context_kv(context_states, context_positions, context_slot_mappings) + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + inputs_embeds: torch.Tensor | None = None, + ) -> torch.Tensor: + # Returns the pre-norm hc_head hidden ([T, hidden_size]). + return self.model(input_ids, positions, inputs_embeds) + + def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor: + """Base logits U_k = lm_head(norm(head_hidden)).""" + return self.logits_processor(self.lm_head, self.model.norm(hidden_states)) + + def compute_draft_logits(self, hidden_states: torch.Tensor) -> torch.Tensor: + # Full-vocab draft: base logits, no d2t scatter. + return self.compute_logits(hidden_states) + + def map_draft_to_target(self, draft_ids: torch.Tensor) -> torch.Tensor: + return draft_ids # full-vocab: draft ids are target ids + + def markov_embed(self, token_ids: torch.Tensor) -> torch.Tensor: + return self.model.markov_head.embed(token_ids) + + def markov_bias(self, markov_embed: torch.Tensor) -> torch.Tensor: + return self.model.markov_head.bias(markov_embed, self.logits_processor) + + # --- Weight loading ---------------------------------------------------- + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + """Load the ``mtp.{0,1,2}.*`` draft weights from the target checkpoint. + + Non-mtp weights (embed/head/main layers) belong to the target model and + are skipped here. ``embed_tokens``/``lm_head`` are aliased from the target. + """ + expert_mapping = fused_moe_make_expert_params_mapping( + self, + ckpt_gate_proj_name="w1", + ckpt_down_proj_name="w2", + ckpt_up_proj_name="w3", + num_experts=self.config.n_routed_experts, + ) + expert_scale_suffix = ( + ".weight_scale" if getattr(self.config, "expert_dtype", "fp4") == "fp4" else ".weight_scale_inv" + ) + + # (param_name, ckpt_shard_name, shard_id) for non-expert stacked params. + stacked_params_mapping = [ + ("gate_up_proj", "w1", 0), + ("gate_up_proj", "w3", 1), + ("attn.fused_wqa_wkv", "attn.wq_a", 0), + ("attn.fused_wqa_wkv", "attn.wkv", 1), + ] + + params_dict = dict(self.named_parameters()) + loaded_params: set[str] = set() + + tp_size = get_tensor_model_parallel_world_size() + tp_rank = get_tensor_model_parallel_rank() + n_local_head = self.config.num_attention_heads // tp_size + head_start = n_local_head * tp_rank + head_end = n_local_head * (tp_rank + 1) + + for name, loaded_weight in weights: + mapped = self._remap_dspark_name(name) + if mapped is None: + continue + name = mapped + + # ``.scale`` -> per-method scale suffix. + if name.endswith(".scale"): + suffix = expert_scale_suffix if _EXPERT_SCALE_RE.search(name) else ".weight_scale_inv" + name = name.removesuffix(".scale") + suffix + + # E8M0 expert scales: keep raw exponent bytes. + if ".experts." in name: + if "weight_scale" in name and loaded_weight.dtype == torch.float8_e8m0fnu: + loaded_weight = loaded_weight.view(torch.uint8) + for param_name, weight_name, expert_id, shard_id in expert_mapping: + if weight_name not in name: + continue + name_mapped = name.replace(weight_name, param_name) + param = params_dict[name_mapped] + success = param.weight_loader( + param, + loaded_weight, + name_mapped, + shard_id=shard_id, + expert_id=expert_id, + return_success=True, + ) + if success: + loaded_params.add(name_mapped) + break + continue + + # Stacked rules only apply to decoder-layer weights. Head-stack params + # (main_proj/norm/hc_head/markov_head) load directly — otherwise e.g. + # "markov_w1" would collide with the "w1" shard rule. + is_layer_param = name.startswith("model.layers.") + for param_name, weight_name, stacked_shard_id in stacked_params_mapping: + if not is_layer_param or weight_name not in name: + continue + name = name.replace(weight_name, param_name) + param = params_dict[name] + param.weight_loader(param, loaded_weight, stacked_shard_id) + loaded_params.add(name) + break + else: + if "attn_sink" in name: + narrow = loaded_weight[head_start:head_end] + params_dict[name][: narrow.shape[0]].copy_(narrow) + loaded_params.add(name) + continue + if ".shared_experts.w2" in name: + name = name.replace(".shared_experts.w2", ".shared_experts.down_proj") + if name.endswith(".ffn.gate.bias"): + name = name.replace(".ffn.gate.bias", ".ffn.gate.e_score_correction_bias") + param = params_dict[name] + weight_loader = getattr(param, "weight_loader", default_weight_loader) + weight_loader(param, loaded_weight) + loaded_params.add(name) + logger.info_once("DSpark draft model loaded: %d params", len(loaded_params)) + return loaded_params + + def _remap_dspark_name(self, name: str) -> str | None: + """Map a checkpoint ``mtp.{i}.*`` name to this model's parameter path. + + Returns None for non-mtp weights (owned by the target model). + """ + m = re.match(r"mtp\.(\d+)\.(.*)", name) + if m is None: + return None + stage = int(m.group(1)) + rest = m.group(2) + # The confidence head is not wired into inference yet; drop its weights. + if rest.startswith("confidence_head."): + return None + # Head-stack params live at model level (mtp.last), context combiner at + # model level (mtp.0); everything else is a per-layer decoder block. + head_prefixes = ( + "norm.", + "hc_head_fn", + "hc_head_base", + "hc_head_scale", + "markov_head.", + ) + if rest.startswith(("main_proj.", "main_norm.")) or rest.startswith(head_prefixes): + return f"model.{rest}" + return f"model.layers.{stage}.{rest}" diff --git a/aphrodite/models/deepseek_v4/amd/rocm.py b/aphrodite/models/deepseek_v4/amd/rocm.py index 244a46f70e..559a77333d 100644 --- a/aphrodite/models/deepseek_v4/amd/rocm.py +++ b/aphrodite/models/deepseek_v4/amd/rocm.py @@ -267,149 +267,6 @@ def compute_global_topk_ragged_indices_and_indptr( return global_topk_ragged, topk_indptr, topk_lens -@triton.jit -def _compute_combined_lens_kernel( - combined_lens_ptr, - query_start_loc_ptr, - seq_lens_ptr, - TOP_K: tl.constexpr, - COMPRESS_RATIO: tl.constexpr, - WINDOW_SIZE: tl.constexpr, -): - batch_idx = tl.program_id(0) - worker_id = tl.program_id(1) - num_workers = tl.num_programs(1) - - base = tl.load(query_start_loc_ptr) - query_start = tl.load(query_start_loc_ptr + batch_idx) - base - query_end = tl.load(query_start_loc_ptr + batch_idx + 1) - base - query_len = query_end - query_start - seq_len = tl.load(seq_lens_ptr + batch_idx) - start_pos = seq_len - query_len - - for token_idx in range(query_start + worker_id, query_end, num_workers): - token_idx_in_query = token_idx - query_start - pos = start_pos + token_idx_in_query - topk_len = tl.minimum((pos + 1) // COMPRESS_RATIO, TOP_K) - swa_len = tl.minimum(pos + 1, WINDOW_SIZE) - tl.store(combined_lens_ptr + token_idx, topk_len + swa_len) - - -@triton.jit -def _combine_topk_swa_indices_ragged_kernel( - combined_ragged_ptr, - combined_indptr_ptr, - topk_indices_ptr, - topk_indices_stride, - query_start_loc_ptr, - seq_lens_ptr, - gather_lens_ptr, - M, - N, - topk_width, - TOP_K: tl.constexpr, - COMPRESS_RATIO: tl.constexpr, - WINDOW_SIZE: tl.constexpr, - BLOCK_SIZE: tl.constexpr, -): - batch_idx = tl.program_id(0) - worker_id = tl.program_id(1) - block_idx = tl.program_id(2) - num_workers = tl.num_programs(1) - - base = tl.load(query_start_loc_ptr) - query_start = tl.load(query_start_loc_ptr + batch_idx) - base - query_end = tl.load(query_start_loc_ptr + batch_idx + 1) - base - query_len = query_end - query_start - seq_len = tl.load(seq_lens_ptr + batch_idx) - gather_len = tl.load(gather_lens_ptr + batch_idx) - start_pos = seq_len - query_len - gather_start = seq_len - gather_len - - for token_idx in range(query_start + worker_id, query_end, num_workers): - token_idx_in_query = token_idx - query_start - pos = start_pos + token_idx_in_query - topk_len = tl.minimum((pos + 1) // COMPRESS_RATIO, TOP_K) - swa_len = tl.minimum(pos + 1, WINDOW_SIZE) - combined_len = topk_len + swa_len - - offset = block_idx * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) - if block_idx * BLOCK_SIZE < combined_len: - out_start = tl.load(combined_indptr_ptr + token_idx) - topk_mask = (offset < topk_len) & (offset < topk_width) - topk_vals = tl.load( - topk_indices_ptr + token_idx * topk_indices_stride + offset, - mask=topk_mask, - other=-1, - ) - tl.store( - combined_ragged_ptr + out_start + offset, - topk_vals + M * batch_idx, - mask=topk_mask, - ) - - swa_offset = offset - topk_len - swa_mask = (offset >= topk_len) & (swa_offset < swa_len) - tl.store( - combined_ragged_ptr + out_start + offset, - M * batch_idx + N + swa_offset + pos - swa_len + 1 - gather_start, - mask=swa_mask, - ) - - -def combine_topk_swa_indices_ragged( - topk_indices: torch.Tensor, - query_start_loc: torch.Tensor, - seq_lens: torch.Tensor, - gather_lens: torch.Tensor, - window_size: int, - compress_ratio: int, - topk: int, - M: int, - N: int, -) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - topk_indices = topk_indices.reshape(topk_indices.shape[0], -1).contiguous() - num_tokens = topk_indices.shape[0] - num_reqs = seq_lens.shape[0] - combined_lens = torch.empty(num_tokens, dtype=torch.int32, device=topk_indices.device) - - num_workers = 128 - _compute_combined_lens_kernel[(num_reqs, num_workers)]( - combined_lens, - query_start_loc, - seq_lens, - TOP_K=topk, - COMPRESS_RATIO=compress_ratio, - WINDOW_SIZE=window_size, - ) - - combined_indptr = _build_indptr_from_lengths(combined_lens) - combined_ragged = torch.empty( - num_tokens * (topk + window_size), - dtype=torch.int32, - device=topk_indices.device, - ) - if combined_ragged.numel() > 0: - block = 128 - _combine_topk_swa_indices_ragged_kernel[(num_reqs, num_workers, triton.cdiv(topk + window_size, block))]( - combined_ragged, - combined_indptr, - topk_indices, - topk_indices.stride(0), - query_start_loc, - seq_lens, - gather_lens, - M, - N, - topk_indices.shape[-1], - TOP_K=topk, - COMPRESS_RATIO=compress_ratio, - WINDOW_SIZE=window_size, - BLOCK_SIZE=block, - ) - return combined_ragged, combined_indptr, combined_lens - - def _copy_ragged_to_graph_buffers( ragged_indices: torch.Tensor, ragged_indptr: torch.Tensor, diff --git a/aphrodite/models/minimax_m3/common/ops/sparse_attn.py b/aphrodite/models/minimax_m3/common/ops/sparse_attn.py index c99569498d..64eab096c1 100644 --- a/aphrodite/models/minimax_m3/common/ops/sparse_attn.py +++ b/aphrodite/models/minimax_m3/common/ops/sparse_attn.py @@ -49,7 +49,7 @@ @triton.jit(do_not_specialize_on_alignment=["seq_lens", "prefix_lens"]) def _gqa_sparse_fwd_kernel( q_ptr, # [total_q, num_heads, head_dim] - kv_cache_ptr, # main cache: [num_blocks, 2, 128, num_kv_heads, head_dim] + kv_cache_ptr, # main cache: [num_blocks, num_kv_heads, 128, 2*head_dim] t_ptr, # topk_idx: [num_kv_heads, total_q, topk] o_ptr, # [total_q, num_heads, head_dim] block_table_ptr, # [num_reqs, max_blocks] @@ -67,9 +67,8 @@ def _gqa_sparse_fwd_kernel( stride_qh, stride_qd, stride_kv_blk, - stride_kv_kv, - stride_kv_pos, stride_kv_h, + stride_kv_pos, stride_kv_d, stride_th, stride_tn, @@ -139,9 +138,8 @@ def _gqa_sparse_fwd_kernel( k = tl.load( kv_cache_ptr + page * stride_kv_blk - + 0 * stride_kv_kv - + off_n[None, :] * stride_kv_pos + pid_kh * stride_kv_h + + off_n[None, :] * stride_kv_pos + off_d[:, None] * stride_kv_d, mask=d_mask[:, None] & pos_mask[None, :], other=0.0, @@ -161,10 +159,9 @@ def _gqa_sparse_fwd_kernel( v = tl.load( kv_cache_ptr + page * stride_kv_blk - + 1 * stride_kv_kv - + off_n[:, None] * stride_kv_pos + pid_kh * stride_kv_h - + off_d[None, :] * stride_kv_d, + + off_n[:, None] * stride_kv_pos + + (head_dim + off_d[None, :]) * stride_kv_d, mask=pos_mask[:, None] & d_mask[None, :], other=0.0, ) @@ -203,7 +200,7 @@ def _gqa_sparse_fwd_kernel( @triton.jit(do_not_specialize=["decode_query_len"]) def _gqa_sparse_decode_kernel( q_ptr, # [total_q, num_heads, head_dim] - kv_cache_ptr, # main cache: [num_blocks, 2, 128, num_kv_heads, head_dim] + kv_cache_ptr, # main cache: [num_blocks, num_kv_heads, 128, 2*head_dim] t_ptr, # topk_idx: [num_kv_heads, total_q, topk] o_ptr, # partial out: [NUM_TOPK_CHUNKS, total_q, num_heads, head_dim] lse_ptr, # partial lse (log2): [NUM_TOPK_CHUNKS, total_q, num_heads] @@ -219,9 +216,8 @@ def _gqa_sparse_decode_kernel( stride_qh, stride_qd, stride_kv_blk, - stride_kv_kv, - stride_kv_pos, stride_kv_h, + stride_kv_pos, stride_kv_d, stride_th, stride_tn, @@ -297,9 +293,8 @@ def _gqa_sparse_decode_kernel( k = tl.load( kv_cache_ptr + page * stride_kv_blk - + 0 * stride_kv_kv - + off_n[None, :] * stride_kv_pos + pid_kh * stride_kv_h + + off_n[None, :] * stride_kv_pos + off_d[:, None] * stride_kv_d, mask=d_mask[:, None] & pos_mask[None, :], other=0.0, @@ -316,10 +311,9 @@ def _gqa_sparse_decode_kernel( v = tl.load( kv_cache_ptr + page * stride_kv_blk - + 1 * stride_kv_kv - + off_n[:, None] * stride_kv_pos + pid_kh * stride_kv_h - + off_d[None, :] * stride_kv_d, + + off_n[:, None] * stride_kv_pos + + (head_dim + off_d[None, :]) * stride_kv_d, mask=pos_mask[:, None] & d_mask[None, :], other=0.0, ) @@ -411,7 +405,7 @@ def _merge_topk_attn_out_kernel( @torch.no_grad() def minimax_m3_sparse_attn( q: torch.Tensor, # [total_q, num_heads, head_dim] - kv_cache: torch.Tensor, # [num_blocks, 2, 128, num_kv_heads, head_dim] + kv_cache: torch.Tensor, # [num_blocks, num_kv_heads, 128, 2*head_dim] topk_idx: torch.Tensor, # [num_kv_heads, total_q, topk] block_table: torch.Tensor, # [batch, max_blocks] cu_seqlens_q: torch.Tensor, # [batch+1] int32 @@ -452,7 +446,6 @@ def minimax_m3_sparse_attn( kv_cache.stride(1), kv_cache.stride(2), kv_cache.stride(3), - kv_cache.stride(4), topk_idx.stride(0), topk_idx.stride(1), topk_idx.stride(2), @@ -469,7 +462,7 @@ def minimax_m3_sparse_attn( @torch.no_grad() def minimax_m3_sparse_attn_decode( q: torch.Tensor, # [total_q, num_heads, head_dim] - kv_cache: torch.Tensor, # [num_blocks, 2, 128, num_kv_heads, head_dim] + kv_cache: torch.Tensor, # [num_blocks, num_kv_heads, 128, 2*head_dim] topk_idx: torch.Tensor, # [num_kv_heads, total_q, topk] block_table: torch.Tensor, # [num_reqs, max_blocks] seq_lens: torch.Tensor, # [num_reqs] int32 @@ -518,7 +511,6 @@ def minimax_m3_sparse_attn_decode( kv_cache.stride(1), kv_cache.stride(2), kv_cache.stride(3), - kv_cache.stride(4), topk_idx.stride(0), topk_idx.stride(1), topk_idx.stride(2), diff --git a/aphrodite/models/minimax_m3/common/sparse_attention.py b/aphrodite/models/minimax_m3/common/sparse_attention.py index 0305dea44b..bc66e05d19 100644 --- a/aphrodite/models/minimax_m3/common/sparse_attention.py +++ b/aphrodite/models/minimax_m3/common/sparse_attention.py @@ -104,20 +104,25 @@ def get_kv_cache_shape( head_size: int, cache_dtype_str: str = "auto", ) -> tuple[int, ...]: - return (num_blocks, 2, block_size, num_kv_heads, head_size) + # K and V are packed into the content dim: logical (B, H, N, 2*hs). + return (num_blocks, num_kv_heads, block_size, 2 * head_size) @staticmethod def get_kv_cache_stride_order( include_num_layers_dimension: bool = False, ) -> tuple[int, ...]: - # Permutation from get_kv_cache_shape to the actual memory layout. + # `stride_order` indicates the permutation that gets us from + # `get_kv_cache_shape` (logical (B, H, N, 2*hs)) to the actual memory + # layout we want. if include_num_layers_dimension: raise NotImplementedError # no cross-layer KV blocks in M3 cache_layout = get_kv_cache_layout() if cache_layout == "NHD": - stride_order = (0, 1, 2, 3, 4) + # (num_blocks, block_size, num_kv_heads, 2*head_size) + stride_order = (0, 2, 1, 3) elif cache_layout == "HND": - stride_order = (0, 1, 3, 2, 4) + # (num_blocks, num_kv_heads, block_size, 2*head_size) + stride_order = (0, 1, 2, 3) else: raise ValueError(f"Unknown cache layout format {cache_layout}.") return stride_order diff --git a/aphrodite/models/minimax_m3/nvidia/sparse_attention_msa.py b/aphrodite/models/minimax_m3/nvidia/sparse_attention_msa.py index 7f8ab0c789..dccd233547 100644 --- a/aphrodite/models/minimax_m3/nvidia/sparse_attention_msa.py +++ b/aphrodite/models/minimax_m3/nvidia/sparse_attention_msa.py @@ -77,8 +77,7 @@ def forward( # strided view directly (topK stays innermost-contiguous). prefill_topk = topk[nd:num_tokens].transpose(0, 1) qp = q[nd:] - k_cache = kv_cache[:, 0].transpose(1, 2) - v_cache = kv_cache[:, 1].transpose(1, 2) + k_cache, v_cache = kv_cache.split(self.head_size, dim=-1) k2q_row_ptr, k2q_q_indices, schedule = build_k2q_csr( prefill_topk, p.cu_seqlens_q, diff --git a/aphrodite/multimodal/video.py b/aphrodite/multimodal/video.py index 3370a1d4fd..5b26d8ebbf 100644 --- a/aphrodite/multimodal/video.py +++ b/aphrodite/multimodal/video.py @@ -206,6 +206,7 @@ def create_hf_metadata( VIDEO_LOADER_REGISTRY = VideoLoaderRegistry() PYNVVIDEOCODEC_VIDEO_BACKEND: Literal["pynvvideocodec"] = "pynvvideocodec" +DEEPSTREAM_VIDEO_BACKEND: Literal["deepstream"] = "deepstream" # Fixed upper bound reserved for persistent PyNvVideoCodec decoder surfaces. PYNVVIDEOCODEC_DECODER_GPU_MEMORY_BYTES = 128 * MiB_bytes PYNVVIDEOCODEC_DECODER_CACHE_SIZE = 2 @@ -801,6 +802,106 @@ def decode_frames_pynvvideocodec( return frames, source, frame_idx, valid_frame_indices +class DeepStreamVideoBackendMixin: + """NVIDIA DeepStream (NVDEC) GPU-decode codec utilities. + + Decoding runs on a shared pool of daemon threads inside one CUDA + context (see the ``nvidia-deepstream-videodecode-cu13`` package). The + container bytes are pushed into an ``appsrc`` GStreamer pipeline, so no + local file path is required. HTTP and base64 sources decode identically + to local files. + + Like the OpenCV/PyAV mixins, this provides only the codec layer. + Frame *selection* lives in the loader's + ``compute_frames_index_to_sample`` and arrives here as an explicit + list of frame indices. + """ + + # Process-wide lazy decode pool, shared across all DeepStream backends. + _pool: ClassVar[Any] = None + _pool_lock: ClassVar[Any] = None + + @classmethod + def _get_pool(cls, pool_size: int | None = None): + """Lazy-initialize the shared decode pool on first use. + + ``pool_size`` (number of decode worker threads) comes from + ``--media-io-kwargs`` (``{"video": {"pool_size": N}}``); when unset it + defaults to the existing ``APHRODITE_MEDIA_LOADING_THREAD_COUNT`` so no + DeepStream-specific env var is needed. The pool is a process-wide + singleton, so the first decode's value wins. + """ + if DeepStreamVideoBackendMixin._pool is not None: + return DeepStreamVideoBackendMixin._pool + if DeepStreamVideoBackendMixin._pool_lock is None: + DeepStreamVideoBackendMixin._pool_lock = threading.Lock() + with DeepStreamVideoBackendMixin._pool_lock: + if DeepStreamVideoBackendMixin._pool is not None: + return DeepStreamVideoBackendMixin._pool + + from nvidia.deepstream_videodecode import DecodePool + + if pool_size is None: + pool_size = int(os.environ.get("APHRODITE_MEDIA_LOADING_THREAD_COUNT", 8)) + pool_size = max(1, min(int(pool_size), 16)) + logger.info( + "[DeepStream] initializing decode pool with %d workers", + pool_size, + ) + DeepStreamVideoBackendMixin._pool = DecodePool(num_workers=pool_size) + return DeepStreamVideoBackendMixin._pool + + @classmethod + def decode_indices( + cls, + data: bytes, + frame_indices: list[int], + codec: str = "", + pool_size: int | None = None, + timeout_sec: float = 120.0, + ) -> tuple[npt.NDArray, list[int]]: + """Decode the requested frame indices from raw container bytes. + + The whole stream is decoded; the pool keeps exactly the frames whose + decode-order index is in ``frame_indices`` (1:1, frame-exact) and + sends EOS once the last one is matched. + + ``codec`` (e.g. ``"h264"``/``"hevc"``) lets the pool keep its NVDEC + session warm across same-codec streams and rebuild only on a codec + change. Frames are returned as a CPU NHWC uint8 array so the + upstream multimodal parser sees the same shape as the other + backends. + """ + if not frame_indices: + raise ValueError("DeepStream backend received no frame indices") + + result = cls._get_pool(pool_size).decode( + data, + target_indices=frame_indices, + codec=codec, + max_frames=len(frame_indices), + timeout_sec=timeout_sec, + ) + if result.error: + raise ValueError(f"DeepStream decode failed: {result.error}") + if result.frames is None or result.n_kept == 0: + raise ValueError("DeepStream decode produced no frames") + + valid = frame_indices[: result.n_kept] + # GPU -> CPU NHWC uint8 at the codec boundary (one PCIe copy); keeps + # the array shape identical to the OpenCV/PyAV backends. Copy into + # pinned host memory so the D2H transfer uses the fast pinned path. + gpu = result.frames + if gpu.is_cuda: + host = torch.empty(gpu.shape, dtype=gpu.dtype, pin_memory=True) + host.copy_(gpu, non_blocking=True) + torch.cuda.current_stream().synchronize() + arr = host.numpy() + else: + arr = gpu.numpy() + return arr, valid + + @VIDEO_LOADER_REGISTRY.register("opencv") class VideoBackend( VideoLoader, @@ -808,14 +909,15 @@ class VideoBackend( PyAVVideoBackendMixin, TorchCodecVideoBackendMixin, PyNvVideoCodecVideoBackendMixin, + DeepStreamVideoBackendMixin, ): """Uniform-sampling video backend. Samples ``num_frames`` uniformly across the video (or one frame every ``1/fps`` seconds, whichever produces fewer frames). The decoding codec is selected via the ``backend`` kwarg (``"opencv"``, ``"pyav"``, - ``"torchcodec"``, or ``"pynvvideocodec"``), which can be passed through - ``--media-io-kwargs``. Defaults to ``"opencv"``. + ``"torchcodec"``, ``"pynvvideocodec"``, or ``"deepstream"``), which can be + passed through ``--media-io-kwargs``. Defaults to ``"opencv"``. """ _sampling_suffix: ClassVar[str] = "" @@ -858,7 +960,7 @@ def load_bytes( max_duration: int = 300, frame_recovery: bool = False, *, - backend: Literal["opencv", "pyav", "torchcodec", "pynvvideocodec"] = "opencv", + backend: Literal["opencv", "pyav", "torchcodec", "pynvvideocodec", "deepstream"] = "opencv", num_ffmpeg_threads: int = 0, seek_mode: Literal["exact", "approximate"] = "exact", **kwargs, @@ -874,7 +976,7 @@ def load_bytes( frame_recovery: Enable forward-scan recovery for failed frames. Only honored by the OpenCV codec. backend: Decoding codec — ``"opencv"``, ``"pyav"``, - ``"torchcodec"``, or ``"pynvvideocodec"``. + ``"torchcodec"``, ``"pynvvideocodec"`` or ``"deepstream"``. num_ffmpeg_threads: Number of FFmpeg decoding threads, only used by TorchCodec. ``0`` uses the FFmpeg default. seek_mode: Seek mode for the TorchCodec decoder, only used by @@ -932,10 +1034,36 @@ def load_bytes( target, **kwargs, ) + elif backend == DEEPSTREAM_VIDEO_BACKEND: + if frame_recovery: + raise ValueError(f"frame_recovery is not supported for `{DEEPSTREAM_VIDEO_BACKEND}` backend") + # Decode-pool size comes from media-io-kwargs (no env var); the + # pool is a process-wide singleton so the first decode's value + # wins. Pop it so it isn't forwarded to the frame sampler. + pool_size = kwargs.pop("pool_size", None) + from nvidia.deepstream_videodecode import probe_metadata + + total_frames, original_fps, duration, width, height, codec = probe_metadata(data) + _check_frame_pixel_limit(width, height) + source = cls._prepare_source( + VideoSourceMetadata( + total_frames_num=total_frames, + original_fps=original_fps, + duration=duration, + ) + ) + frame_idx = cls.compute_frames_index_to_sample(source=source, target=target, **kwargs) + frames, valid = cls.decode_indices( + data, + frame_idx, + codec=codec, + pool_size=pool_size, + ) else: raise ValueError( f"Unknown video codec backend {backend!r}; " - "valid options: 'opencv', 'pyav', 'torchcodec', 'pynvvideocodec'." + "valid options: 'opencv', 'pyav', 'torchcodec', " + "'pynvvideocodec' and 'deepstream'." ) if len(valid) < len(frame_idx): @@ -987,6 +1115,32 @@ def load_bytes( ) +@VIDEO_LOADER_REGISTRY.register(DEEPSTREAM_VIDEO_BACKEND, requires_gpu=True) +class DeepStreamVideoBackend(VideoBackend): + """Hardware-accelerated video backend using NVIDIA DeepStream.""" + + @classmethod + def load_bytes( + cls, + data: bytes, + num_frames: int = -1, + fps: int = -1, + max_duration: int = 300, + frame_recovery: bool = False, + **kwargs, + ) -> tuple[npt.NDArray, dict[str, Any]]: + kwargs.pop("backend", None) + return super().load_bytes( + data, + num_frames=num_frames, + fps=fps, + max_duration=max_duration, + frame_recovery=frame_recovery, + backend=DEEPSTREAM_VIDEO_BACKEND, + **kwargs, + ) + + @VIDEO_LOADER_REGISTRY.register( "qwen3_vl", video_processor="Qwen3VLVideoProcessor", @@ -1022,7 +1176,7 @@ def load_bytes( max_duration: int = 300, frame_recovery: bool = False, *, - backend: Literal["opencv", "pyav", "torchcodec", "pynvvideocodec"] = "opencv", + backend: Literal["opencv", "pyav", "torchcodec", "pynvvideocodec", "deepstream"] = "opencv", **kwargs, ) -> tuple[npt.NDArray, dict[str, Any]]: return super().load_bytes( @@ -1098,7 +1252,7 @@ def load_bytes( max_duration: int = 300, frame_recovery: bool = False, *, - backend: Literal["opencv", "pyav", "torchcodec", "pynvvideocodec"] = "opencv", + backend: Literal["opencv", "pyav", "torchcodec", "pynvvideocodec", "deepstream"] = "opencv", **kwargs, ) -> tuple[npt.NDArray, dict[str, Any]]: return super().load_bytes( @@ -1180,7 +1334,7 @@ def load_bytes( max_duration: int = 300, frame_recovery: bool = False, *, - backend: Literal["opencv", "pyav", "torchcodec", "pynvvideocodec"] = "opencv", + backend: Literal["opencv", "pyav", "torchcodec", "pynvvideocodec", "deepstream"] = "opencv", **kwargs, ) -> tuple[npt.NDArray, dict[str, Any]]: return super().load_bytes( @@ -1301,7 +1455,7 @@ def load_bytes( max_duration: int = 300, frame_recovery: bool = False, *, - backend: Literal["opencv", "pyav", "torchcodec", "pynvvideocodec"] = "opencv", + backend: Literal["opencv", "pyav", "torchcodec", "pynvvideocodec", "deepstream"] = "opencv", **kwargs, ) -> tuple[npt.NDArray, dict[str, Any]]: return super().load_bytes( @@ -1393,7 +1547,7 @@ def load_bytes( max_duration: int = 300, frame_recovery: bool = False, *, - backend: Literal["opencv", "pyav", "torchcodec", "pynvvideocodec"] = "opencv", + backend: Literal["opencv", "pyav", "torchcodec", "pynvvideocodec", "deepstream"] = "opencv", **kwargs, ) -> tuple[npt.NDArray, dict[str, Any]]: frames, metadata = super().load_bytes( @@ -1696,7 +1850,7 @@ def load_bytes( max_duration: int = 300, frame_recovery: bool = False, *, - backend: Literal["opencv", "pyav", "torchcodec", "pynvvideocodec"] = "opencv", + backend: Literal["opencv", "pyav", "torchcodec", "pynvvideocodec", "deepstream"] = "opencv", **kwargs, ) -> tuple[npt.NDArray, dict[str, Any]]: frames, metadata = super().load_bytes( diff --git a/aphrodite/parser/abstract_parser.py b/aphrodite/parser/abstract_parser.py index 98da5632bb..5f1cef71cd 100644 --- a/aphrodite/parser/abstract_parser.py +++ b/aphrodite/parser/abstract_parser.py @@ -431,7 +431,7 @@ def _extract_tool_calls( tool_calls = list[FunctionCall]() if is_named_tool_choice and supports_required_and_named: - if content is None: + if content is None or (isinstance(content, str) and not content.strip()): return [], None function_name = self._get_function_name(request) tool_calls.append( @@ -478,6 +478,13 @@ def _extract_tool_calls( content = None else: # No tool calls. + # For required/named tool choice (when falling back to auto + # parsing), if content is empty or whitespace-only, return + # empty list with None content. + if (is_required_tool_choice or is_named_tool_choice) and ( + content is None or (isinstance(content, str) and not content.strip()) + ): + return [], None return None, content return tool_calls, content @@ -867,6 +874,15 @@ def parse_delta( delta_message = self.finalize_generation(delta_message, request, state) delta_message = self._flush_engine_parsers(delta_message) + # Suppress reasoning deltas if not requested. + if delta_message and not request.include_reasoning: + delta_message.reasoning = None + + # If only reasoning was in the message (no content, no tool_calls), + # skip emitting entirely. + if not delta_message.content and not delta_message.tool_calls: + delta_message = None + return delta_message def _flush_engine_parsers(self, delta_message: DeltaMessage | None) -> DeltaMessage | None: diff --git a/aphrodite/parser/engine/parser_engine.py b/aphrodite/parser/engine/parser_engine.py index 4d0bae9bd1..138115655b 100644 --- a/aphrodite/parser/engine/parser_engine.py +++ b/aphrodite/parser/engine/parser_engine.py @@ -429,7 +429,15 @@ def parse_delta( if finished: events.extend(self._engine.finish()) result = self._events_to_delta(events, finished=finished) - return self._strip_trailing_reasoning(result) + result = self._strip_trailing_reasoning(result) + + # Suppress reasoning deltas if not requested. + if result and not request.include_reasoning: + result.reasoning = None + if not result.content and not result.tool_calls: + result = None + + return result def _strip_trailing_reasoning( self, diff --git a/aphrodite/parser/harmony.py b/aphrodite/parser/harmony.py index 3b296b28ff..ddbaec3848 100644 --- a/aphrodite/parser/harmony.py +++ b/aphrodite/parser/harmony.py @@ -270,6 +270,16 @@ def parse_delta( delta_message.reasoning = combined_reasoning if tool_messages: delta_message.tool_calls = tool_messages + + # Suppress reasoning deltas if not requested. + if delta_message and not request.include_reasoning: + delta_message.reasoning = None + + # If only reasoning was in the message (no content, no tool_calls), + # skip emitting entirely. + if not delta_message.content and not delta_message.tool_calls: + return None + return delta_message def process_chunk(self, token_ids: Sequence[int]) -> ChunkResult: diff --git a/aphrodite/platforms/rocm.py b/aphrodite/platforms/rocm.py index b965a0da46..45e6b45898 100644 --- a/aphrodite/platforms/rocm.py +++ b/aphrodite/platforms/rocm.py @@ -427,8 +427,8 @@ def _get_backend_priorities( ] backends = [] - # ROCM_ATTN uses (2, num_blocks, ...) KV cache layout which is - # incompatible with KV connectors that require blocks-first layout. + # Keep ROCM_ATTN disabled for KV connectors until connector transfer + # semantics are validated for its asymmetric native K/V cache views. if not use_kv_connector: backends.append(AttentionBackendEnum.ROCM_ATTN) if rocm_aiter_ops.is_mha_enabled(): @@ -511,6 +511,19 @@ def get_valid_backends( attn_selector_config.use_sparse, attn_selector_config.use_kv_connector, ) + from aphrodite.config import get_current_aphrodite_config_or_none + + aphrodite_config = get_current_aphrodite_config_or_none() + is_encoder_decoder = ( + getattr(getattr(aphrodite_config, "model_config", None), "attn_type", None) == "encoder_decoder" + ) + # ROCM_ATTN still uses a legacy attention layout (KV is the outer + # dimension) that is incompatible with the encoder backend layouts. The + # encoder and decoder need the layouts to match. This is currently + # enforced implicitly. + # TODO: Make this explicit in the selector in a future PR. + if is_encoder_decoder and AttentionBackendEnum.ROCM_ATTN in backend_priorities: + backend_priorities.remove(AttentionBackendEnum.ROCM_ATTN) for priority, backend in enumerate(backend_priorities): try: backend_class = backend.get_class() diff --git a/aphrodite/transformers_utils/model_arch_config_convertor.py b/aphrodite/transformers_utils/model_arch_config_convertor.py index 7d5714e5d6..369eada3c7 100644 --- a/aphrodite/transformers_utils/model_arch_config_convertor.py +++ b/aphrodite/transformers_utils/model_arch_config_convertor.py @@ -256,6 +256,7 @@ def is_deepseek_mla(self) -> bool: "kimi_k2", "kimi_linear", "longcat_flash", + "longcat_flash_ngram", "pangu_ultra_moe", "pangu_ultra_moe_mtp", "bailing_hybrid", diff --git a/aphrodite/transformers_utils/processors/mimo_v2_omni.py b/aphrodite/transformers_utils/processors/mimo_v2_omni.py index e028efb6b1..52c03b212b 100644 --- a/aphrodite/transformers_utils/processors/mimo_v2_omni.py +++ b/aphrodite/transformers_utils/processors/mimo_v2_omni.py @@ -7,33 +7,21 @@ """ import contextlib -import copy -import io import logging import math from collections import OrderedDict from concurrent.futures import ThreadPoolExecutor, as_completed from dataclasses import dataclass, field -from io import BytesIO from typing import Any, Literal import numpy as np import regex as re -import requests import torch import torch.nn.functional as F from PIL import Image from transformers import BatchFeature, TensorType from transformers.processing_utils import ProcessorMixin -try: - from torchcodec.decoders import AudioDecoder - - _HAS_TORCHCODEC = True -except ImportError: - AudioDecoder = None - _HAS_TORCHCODEC = False - try: import torchaudio from torchaudio.transforms import MelSpectrogram as _MelSpectrogram @@ -62,7 +50,7 @@ @dataclass class ImageInput: - # PIL.Image | str (path/url/base64) | bytes | torch.Tensor (C,H,W) + # PIL.Image | torch.Tensor (C,H,W) image: Any max_pixels: int | None = None min_pixels: int | None = None @@ -87,7 +75,7 @@ class VideoInput: @dataclass class AudioInput: - # str (path/url/base64) | bytes | tuple[waveform_1D, sr] + # tuple[waveform_1D, sr] # | np.ndarray | torch.Tensor (T,n_vq) audio: Any @@ -166,14 +154,6 @@ def _smart_resize(h: int, w: int, factor: int, min_px: int, max_px: int) -> tupl return int(h_bar), int(w_bar) -def _to_rgb(img: Image.Image) -> Image.Image: - if img.mode == "RGBA": - bg = Image.new("RGB", img.size, (255, 255, 255)) - bg.paste(img, mask=img.split()[3]) - return bg - return img.convert("RGB") - - def _standardize(images: torch.Tensor) -> torch.Tensor: key = str(images.device) if key not in _mean_std_cache: @@ -222,27 +202,6 @@ def _transform_single( return _standardize(out).squeeze(0), w_bar, h_bar -def _fetch_image(src: Any) -> Image.Image: - if isinstance(src, Image.Image): - return _to_rgb(src) - if isinstance(src, bytes): - return _to_rgb(copy.deepcopy(Image.open(BytesIO(src)))) - if isinstance(src, str): - if src.startswith(("http://", "https://")): - r = requests.get(src, timeout=30) - r.raise_for_status() - return _to_rgb(copy.deepcopy(Image.open(BytesIO(r.content)))) - if src.startswith("file://"): - return _to_rgb(Image.open(src[7:])) - if src.startswith("data:image"): - import pybase64 as _b64 - - _, b64 = src.split("base64,", 1) - return _to_rgb(copy.deepcopy(Image.open(BytesIO(_b64.b64decode(b64))))) - return _to_rgb(Image.open(src)) - raise ValueError(f"Unrecognized image source: {type(src)}") - - # --------------------------------------------------------------------------- # Core processor # --------------------------------------------------------------------------- @@ -253,6 +212,10 @@ class MiMoVLProcessor: Handles image/video/audio preprocessing and token sequence construction. Ported from SGLang's MiMoVLProcessor. + + Media strings (URL / path / data:) are resolved upstream in Aphrodite's + media pipeline; this processor accepts only already-decoded inputs + (``PIL.Image`` for images, ``(waveform, sr)`` tuples for audio). """ def __init__( @@ -424,30 +387,14 @@ def _resolve_vid_kw(self, vid: VideoInput) -> dict: return kw def preprocess_audio(self, audio: Any) -> tuple[torch.Tensor, int]: - """Decode audio bytes/path/tuple → (mel_spec (T, n_mels), token_len).""" - if isinstance(audio, tuple): - waveform, original_sr = audio - else: - if AudioDecoder is None: - raise RuntimeError("torchcodec is required for audio. Install with: pip install torchcodec") - if isinstance(audio, bytes): - file_obj: Any = io.BytesIO(audio) - elif isinstance(audio, str): - if audio.startswith("data:"): - import pybase64 as _b64 - - file_obj = io.BytesIO(_b64.b64decode(audio.split(",")[1])) - elif audio.startswith(("http://", "https://")): - r = requests.get(audio, timeout=30) - r.raise_for_status() - file_obj = io.BytesIO(r.content) - else: - file_obj = audio - else: - raise ValueError(f"Unsupported audio source type: {type(audio)}") - samples = AudioDecoder(file_obj).get_all_samples() - waveform = samples.data - original_sr = samples.sample_rate + """Convert a pre-loaded ``(waveform, sr)`` tuple into mel features.""" + if not isinstance(audio, tuple): + raise ValueError( + f"Unsupported audio source type: {type(audio)}. Audio must be a " + "pre-decoded (waveform, sample_rate) tuple; URL/path/bytes " + "resolution is handled upstream in Aphrodite's media pipeline." + ) + waveform, original_sr = audio if original_sr != self.audio_sampling_rate: if original_sr not in self._resamplers: @@ -474,8 +421,6 @@ def preprocess_audio(self, audio: Any) -> tuple[torch.Tensor, int]: def process_image(self, image: ImageInput) -> torch.Tensor: kw = self._resolve_img_kw(image) src = image.image - if isinstance(src, (str, bytes)): - src = _fetch_image(src) tensor, _, _ = _transform_single( src, factor=self.patch_size * self.merge_size, @@ -538,7 +483,7 @@ def process_audio(self, audio: AudioInput) -> Any: src = audio.audio if isinstance(src, np.ndarray): src = (torch.from_numpy(src).float(), self.audio_sampling_rate) - if isinstance(src, (str, bytes, tuple)): + if isinstance(src, tuple): return self.preprocess_audio(src) # Pre-tokenized tensor (T, n_vq) assert isinstance(src, torch.Tensor) and src.ndim == 2 @@ -813,7 +758,7 @@ class MiMoOmniProcessor(ProcessorMixin): """HuggingFace-compatible ProcessorMixin wrapper for MiMo-Omni. Accepts PIL images, pre-decoded video tuples (frames_TCHW, timestamps_T), - and audio (file path / bytes / (waveform, sr) tuple / numpy array). + and audio ((waveform, sr) tuple / numpy array). """ attributes = ["tokenizer"] @@ -923,8 +868,12 @@ def __init__( ) @classmethod - def from_hf_config(cls, tokenizer: Any, hf_config: Any) -> "MiMoOmniProcessor": - """Convenience factory: instantiate directly from an HF model config object.""" + def from_hf_config( + cls, + tokenizer: Any, + hf_config: Any, + ) -> "MiMoOmniProcessor": + """Instantiate directly from an HF model config object.""" vc = hf_config.vision_config if isinstance(vc, dict): patch_size = vc.get("patch_size", 14) @@ -1031,8 +980,8 @@ def __call__( images: PIL.Image or list[PIL.Image]. videos: list of ``(frames_TCHW: torch.Tensor, timestamps_T: torch.Tensor)`` tuples (pre-decoded). - audio: list of ``str`` (path/url/base64), ``bytes``, - ``(waveform_1D, sample_rate)`` tuples, or ``np.ndarray``. + audio: list of ``(waveform_1D, sample_rate)`` tuples or + ``np.ndarray``. return_tensors: Passed to :class:`BatchFeature`. Returns: diff --git a/aphrodite/utils/torch_utils.py b/aphrodite/utils/torch_utils.py index c62d8322ec..c840985f3a 100644 --- a/aphrodite/utils/torch_utils.py +++ b/aphrodite/utils/torch_utils.py @@ -400,26 +400,24 @@ def nvfp4_kv_cache_full_dim(head_size: int) -> int: return head_size // 2 + head_size // 16 -def _nvfp4_split_data_scale( +def nvfp4_split_data_scale( kv_side: torch.Tensor, ) -> tuple[torch.Tensor, torch.Tensor]: - """Split a single NVFP4 KV-side buffer into data and scale views. + """Split one side (K or V) of an NVFP4 KV cache into data and scale. - The input is a 4D tensor for one KV side (K or V) whose last - dimension is ``full_dim = data_dim + scale_dim``. The physical - layout within each side is [data | scale], both packed contiguously. + The input is a 4D uint8 tensor whose last dimension is + ``full_dim = data_dim + scale_dim``. The physical layout within each + side is ``[data | scale]``, both packed contiguously. + + The caller is responsible for slicing K and V from the combined cache + first (e.g. ``kv_cache.split(num_kv_heads, dim=1)``). Args: - kv_side: 4D uint8 tensor with shape - ``(num_pages, dim_1, dim_2, full_dim)``. - May be in any permutation order (NHD or HND). + kv_side: 4D uint8 tensor ``(B, H, N, full_dim)``. Returns: - ``(data, scale)`` where - ``data`` is a uint8 view with shape - ``(num_pages, dim_1, dim_2, data_dim)``. - ``scale`` is a float8_e4m3fn view with shape - ``(num_pages, dim_1, dim_2, scale_dim)``. + ``(data, scale)`` where *data* is uint8 and *scale* is + float8_e4m3fn, both views of the same storage. """ num_pages = kv_side.shape[0] dim_1, dim_2 = kv_side.shape[1], kv_side.shape[2] @@ -452,38 +450,6 @@ def _nvfp4_split_data_scale( return data, scale -def nvfp4_kv_cache_split_views(kv_cache: torch.Tensor) -> tuple[tuple, tuple]: - """Split an NVFP4 KV cache tensor into data and scale views. - - Accepts either a 5D tensor ``(num_pages, 2, dim_2, dim_3, full_dim)`` - or a 4D single-side tensor ``(num_pages, dim_2, dim_3, full_dim)``. - - Per-page layout: [K_data | K_scale | V_data | V_scale]. - Each KV side is self-contained (data followed by its scale), so the - 5D case simply splits each side independently. - - The returned views are in the same dim order as the input (NHD or - HND), so callers get views matching whichever order they passed in. - - Args: - kv_cache: 5D or 4D uint8 tensor where the last dimension is - ``full_dim = data_dim + scale_dim = 9 * head_size / 16``. - - Returns: - For 5D input: - ``(k_data, v_data), (k_scale, v_scale)`` - For 4D input (single KV side): - ``(data,), (scale,)`` - """ - if kv_cache.dim() == 4: - data, scale = _nvfp4_split_data_scale(kv_cache) - return (data,), (scale,) - - k_data, k_scale = _nvfp4_split_data_scale(kv_cache[:, 0]) - v_data, v_scale = _nvfp4_split_data_scale(kv_cache[:, 1]) - return (k_data, v_data), (k_scale, v_scale) - - def create_kv_caches_with_random_flash( num_blocks: int, block_size: int, diff --git a/aphrodite/v1/attention/backend.py b/aphrodite/v1/attention/backend.py index 7598f7d614..cf46e872a9 100644 --- a/aphrodite/v1/attention/backend.py +++ b/aphrodite/v1/attention/backend.py @@ -122,11 +122,11 @@ def get_kv_cache_stride_order( ) -> tuple[int, ...]: """ Get the physical (memory layout) ordering of the kv cache dimensions. - e.g. if the KV cache shape is - [2, num_blocks, block_size, num_heads, head_size], - and get_kv_cache_stride_order returns (1, 3, 0, 2, 4) then the physical + Standard attention backends pack K and V into the content dim, giving + the logical shape [num_blocks, num_heads, block_size, 2 * head_size]. + e.g. if get_kv_cache_stride_order returns (0, 2, 1, 3) then the physical ordering of dimensions is - [num_blocks, num_heads, 2, block_size, head_size]. + [num_blocks, block_size, num_heads, 2 * head_size]. If this function is unimplemented / raises NotImplementedError, the physical layout of the KV cache will match the logical shape. @@ -135,9 +135,9 @@ def get_kv_cache_stride_order( include_num_layers_dimension: if True, includes an additional num_layers dimension, which is assumed to be prepended to the logical KV cache shape. - With the above example, a return value (2, 4, 0, 1, 3, 5) + With the above example, a return value (1, 0, 3, 2, 4) corresponds to - [num_blocks, num_heads, num_layers, 2, block_size, head_size]. + [num_blocks, num_layers, block_size, num_heads, 2 * head_size]. If an additional dimension is NOT included in the returned tuple, the physical layout will not include a layers dimension. diff --git a/aphrodite/v1/attention/backends/flash_attn.py b/aphrodite/v1/attention/backends/flash_attn.py index 86ff7ee98b..b563f08901 100755 --- a/aphrodite/v1/attention/backends/flash_attn.py +++ b/aphrodite/v1/attention/backends/flash_attn.py @@ -129,25 +129,29 @@ def get_kv_cache_shape( ) -> tuple[int, ...]: if block_size % 16 != 0: raise ValueError("Block size must be a multiple of 16.") - return (num_blocks, 2, block_size, num_kv_heads, head_size) + # K and V are packed into the content dim: logical (B, H, N, 2*D). + return (num_blocks, num_kv_heads, block_size, 2 * head_size) @staticmethod def get_kv_cache_stride_order( include_num_layers_dimension: bool = False, ) -> tuple[int, ...]: - # `stride_order` indicates the permutation that gets - # us from `get_kv_cache_shape` to the actual memory layout we want. + # `stride_order` indicates the permutation that gets us from + # `get_kv_cache_shape` (logical (B, H, N, 2*D)) to the actual memory + # layout we want. cache_layout = get_kv_cache_layout() if cache_layout == "NHD" and include_num_layers_dimension: - # (num_blocks, num_layers, 2, block_size, num_kv_heads, head_size) - return (1, 0, 2, 3, 4, 5) + # (num_blocks, num_layers, block_size, num_kv_heads, 2*head_size) + return (1, 0, 3, 2, 4) elif cache_layout == "NHD": - stride_order = (0, 1, 2, 3, 4) + # (num_blocks, block_size, num_kv_heads, 2*head_size) + stride_order = (0, 2, 1, 3) elif cache_layout == "HND" and include_num_layers_dimension: - # (num_blocks, num_kv_heads, num_layers, 2, block_size, head_size) - return (1, 4, 0, 2, 3, 5) + # (num_blocks, num_kv_heads, num_layers, block_size, 2*head_size) + return (1, 2, 0, 3, 4) elif cache_layout == "HND": - stride_order = (0, 1, 3, 2, 4) + # (num_blocks, num_kv_heads, block_size, 2*head_size) + stride_order = (0, 1, 2, 3) else: raise ValueError(f"Unknown cache layout format {cache_layout}.") return stride_order @@ -725,7 +729,7 @@ def forward( key: shape = [num_tokens, num_kv_heads, head_size] value: shape = [num_tokens, num_kv_heads, head_size] kv_cache: shape = - [num_blocks, 2, block_size, num_kv_heads, head_size] + [num_blocks, num_kv_heads, block_size, 2 * head_size] attn_metadata: Metadata for attention. Returns: shape = [num_tokens, num_heads * head_size] @@ -768,8 +772,8 @@ def forward( layer, ) - # For decoder and cross-attention, use KV cache as before - key_cache, value_cache = kv_cache.unbind(1) + # (B, H, N, 2*D) -> ((B, N, H, D), (B, N, H, D)) + key_cache, value_cache = kv_cache.transpose(1, 2).split(self.head_size, dim=-1) # Fix degenerate strides on size-1 dims (e.g. num_kv_heads=1 with TP). # FA3/4 on H100+ uses TMA, which requires ≥16-byte stride alignment. # See aphrodite.utils.torch_utils.canonicalize_singleton_dim_strides. @@ -962,7 +966,8 @@ def do_kv_cache_update( # Scatter write into the KV cache using slot_mapping indices. # No TMA kernel is invoked here, so stride canonicalization is not needed. - key_cache, value_cache = kv_cache.unbind(1) + # (B, H, N, 2*D) -> ((B, N, H, D), (B, N, H, D)) + key_cache, value_cache = kv_cache.transpose(1, 2).split(self.head_size, dim=-1) # Reshape the input keys and values and store them in the cache. # Skip this if sharing KV cache with an earlier attention layer. diff --git a/aphrodite/v1/attention/backends/flash_attn_diffkv.py b/aphrodite/v1/attention/backends/flash_attn_diffkv.py index 8ce5f1b438..95cbe8e6f0 100644 --- a/aphrodite/v1/attention/backends/flash_attn_diffkv.py +++ b/aphrodite/v1/attention/backends/flash_attn_diffkv.py @@ -83,10 +83,12 @@ def get_kv_cache_shape( ) -> tuple[int, ...]: if block_size % 16 != 0: raise ValueError("Block size must be a multiple of 16.") + # Logical (blocks-first, head-major) layout: K and V (with their + # different head sizes) packed in the content dim. return ( num_blocks, - block_size, num_kv_heads, + block_size, head_size + FlashAttentionDiffKVBackend.head_size_v, ) @@ -94,21 +96,22 @@ def get_kv_cache_shape( def get_kv_cache_stride_order( include_num_layers_dimension: bool = False, ) -> tuple[int, ...]: - # `stride_order` indicates the permutation that gets - # us from `get_kv_cache_shape` to the actual memory layout we want. + # `stride_order` indicates the permutation that gets us from + # `get_kv_cache_shape` (logical (B, H, N, C_k+C_v)) to the actual + # memory layout we want. cache_layout = get_kv_cache_layout() if cache_layout == "NHD" and include_num_layers_dimension: - # (num_blocks, num_layers, block_size, - # num_kv_heads, head_size + head_size_v) - return (1, 0, 2, 3, 4) + # (num_blocks, num_layers, block_size, num_kv_heads, C_k+C_v) + return (1, 0, 3, 2, 4) elif cache_layout == "NHD": - stride_order = (0, 1, 2, 3) + # (num_blocks, block_size, num_kv_heads, C_k+C_v) + stride_order = (0, 2, 1, 3) elif cache_layout == "HND" and include_num_layers_dimension: - # (num_blocks, num_kv_heads, num_layers, - # block_size, head_size + head_size_v) - return (1, 3, 0, 2, 4) + # (num_blocks, num_kv_heads, num_layers, block_size, C_k+C_v) + return (1, 2, 0, 3, 4) elif cache_layout == "HND": - stride_order = (0, 2, 1, 3) + # (num_blocks, num_kv_heads, block_size, C_k+C_v) + stride_order = (0, 1, 2, 3) else: raise ValueError(f"Unknown cache layout format {cache_layout}.") return stride_order @@ -142,21 +145,14 @@ def do_kv_cache_update( # we use direct Q, K, V tensors without caching return - # Unlike standard FlashAttn which splits kv_cache via unbind(0), # DiffKV packs K and V into a single tensor along the last dim: # kv_cache shape: [num_blocks, block_size, num_kv_heads, # head_size_k + head_size_v] - # The triton kernel handles this combined layout directly. - # - # NOTE(woosuk): key and value are padded while slot_mapping is - # not padded. However, we don't need to do key[:num_actual_tokens] - # and value[:num_actual_tokens] because the reshape_and_cache_flash - # op uses the slot_mapping's shape to determine the number of - # actual tokens. + # (B, H, N, C) -> (B, N, H, C) for kernel compatibility. triton_reshape_and_cache_flash_diffkv( key, value, - kv_cache, + kv_cache.transpose(1, 2), slot_mapping, self.kv_cache_dtype, layer._k_scale, @@ -225,10 +221,8 @@ def forward( layer, ) - # For decoder and cross-attention, use KV cache as before - # Different head_size for K and V - key_cache = kv_cache[..., : self.head_size] - value_cache = kv_cache[..., self.head_size :] + # (B, H, N, 2*hs) -> ((B, N, H, hs), (B, N, H, hs)) + key_cache, value_cache = kv_cache.transpose(1, 2).split(self.head_size, dim=-1) # Fix degenerate strides on size-1 dims (e.g. num_kv_heads=1 with TP). # FA3/4 on H100+ uses TMA, which requires ≥16-byte stride alignment. # See aphrodite.utils.torch_utils.canonicalize_singleton_dim_strides. diff --git a/aphrodite/v1/attention/backends/flashinfer.py b/aphrodite/v1/attention/backends/flashinfer.py index 30dd2132d3..755142537e 100755 --- a/aphrodite/v1/attention/backends/flashinfer.py +++ b/aphrodite/v1/attention/backends/flashinfer.py @@ -51,7 +51,7 @@ is_quantized_kv_cache, is_strictly_contiguous, nvfp4_kv_cache_full_dim, - nvfp4_kv_cache_split_views, + nvfp4_split_data_scale, ) from aphrodite.v1.attention.backend import ( AttentionBackend, @@ -111,9 +111,12 @@ def _trtllm_prefill_attn_kvfp8_dequant( src_stride_page, src_stride_kv, src_stride_head, + src_stride_block, + src_stride_head_size, DST_K_CACHE_STRIDE: tl.constexpr, DST_KV_CACHE_STRIDE: tl.constexpr, HEAD_STRIDE: tl.constexpr, + HEAD_SIZE: tl.constexpr, NUM_KV_HEADS: tl.constexpr, ): batch_idx = tl.program_id(0).to(tl.int64) @@ -129,27 +132,40 @@ def _trtllm_prefill_attn_kvfp8_dequant( v_scale_val = tl.load(v_scale_ptr) mock_page_idx = batch_idx * block_table_stride + mock_block_table_idx + 1 - head_offsets = tl.arange(0, HEAD_STRIDE) + logical_offsets = tl.arange(0, HEAD_STRIDE) + block_offsets = logical_offsets // HEAD_SIZE + head_size_offsets = logical_offsets % HEAD_SIZE for h in range(NUM_KV_HEADS): h_off = tl.cast(h, tl.int64) # Read K from source (supports non-contiguous page/kv/head strides) - src_k = orig_page_num * src_stride_page + h_off * src_stride_head + head_offsets + src_k = ( + orig_page_num * src_stride_page + + h_off * src_stride_head + + block_offsets * src_stride_block + + head_size_offsets * src_stride_head_size + ) fp8_k = tl.load(kv_cache_ptr + src_k) dequant_k = (fp8_k.to(tl.float32) * k_scale_val).to(dequant_dtype) # Write K to contiguous mock cache - dst_k = mock_page_idx * DST_KV_CACHE_STRIDE + h * HEAD_STRIDE + head_offsets + dst_k = mock_page_idx * DST_KV_CACHE_STRIDE + h * HEAD_STRIDE + logical_offsets tl.store(mock_kv_cache_ptr + dst_k, dequant_k) # Read V from source (offset by src_stride_kv for the V half) - src_v = orig_page_num * src_stride_page + src_stride_kv + h_off * src_stride_head + head_offsets + src_v = ( + orig_page_num * src_stride_page + + src_stride_kv + + h_off * src_stride_head + + block_offsets * src_stride_block + + head_size_offsets * src_stride_head_size + ) fp8_v = tl.load(kv_cache_ptr + src_v) dequant_v = (fp8_v.to(tl.float32) * v_scale_val).to(dequant_dtype) # Write V to contiguous mock cache - dst_v = mock_page_idx * DST_KV_CACHE_STRIDE + DST_K_CACHE_STRIDE + h * HEAD_STRIDE + head_offsets + dst_v = mock_page_idx * DST_KV_CACHE_STRIDE + DST_K_CACHE_STRIDE + h * HEAD_STRIDE + logical_offsets tl.store(mock_kv_cache_ptr + dst_v, dequant_v) @@ -171,9 +187,6 @@ def trtllm_prefill_attn_kvfp8_dequant( kv_cache_stride = k_cache_stride * s[1] strides = kv_cache.stride() - assert strides[3] == head_size and strides[4] == 1, ( - f"For kv cache layouts, (block_size, head_size) dimensions must be contiguous, got strides {strides}" - ) new_s = (batch_size * num_of_page_per_token + 1, s[1], s[2], s[3], s[4]) # mock kv cache contains just the pages needed by this prefill @@ -196,9 +209,12 @@ def trtllm_prefill_attn_kvfp8_dequant( strides[0], strides[1], strides[2], + strides[3], + strides[4], k_cache_stride, kv_cache_stride, head_stride, + head_size, num_kv_heads, ) return mock_kv_cache, mock_block_table @@ -273,7 +289,7 @@ def run( self, layer: torch.nn.Module, prefill_query: torch.Tensor, - kv_cache_permute: torch.Tensor, + kv_cache_tuple: tuple[torch.Tensor, torch.Tensor], key: torch.Tensor, value: torch.Tensor, out: torch.Tensor, @@ -281,7 +297,7 @@ def run( prefill_query_across_dcp = get_dcp_group().all_gather(prefill_query.contiguous(), dim=1) output_context_tmp, lse_context_tmp = self._context.run( prefill_query_across_dcp, - kv_cache_permute, + kv_cache_tuple, k_scale=layer._k_scale_float, v_scale=layer._v_scale_float, return_lse=True, @@ -371,28 +387,31 @@ def get_kv_cache_shape( cache_dtype_str: str = "auto", ) -> tuple[int, ...]: if cache_dtype_str == "nvfp4": - # Packed layout: fp4 data + fp8 block scales in last dim - last_dim = nvfp4_kv_cache_full_dim(head_size) - return (num_blocks, 2, block_size, num_kv_heads, last_dim) - return (num_blocks, 2, block_size, num_kv_heads, head_size) + full_dim = nvfp4_kv_cache_full_dim(head_size) + return (num_blocks, 2 * num_kv_heads, block_size, full_dim) + # Pack K and V in the content dim (B, H, N, 2*hs). + return (num_blocks, num_kv_heads, block_size, 2 * head_size) @staticmethod def get_kv_cache_stride_order( include_num_layers_dimension: bool = False, ) -> tuple[int, ...]: # `stride_order` indicates the permutation that gets us from - # `get_kv_cache_shape` to the actual memory layout we want. + # `get_kv_cache_shape` (logical (B, H, N, 2*hs)) to the actual memory + # layout we want. cache_layout = get_kv_cache_layout() if cache_layout == "NHD" and include_num_layers_dimension: - # (num_blocks, num_layers, 2, block_size, num_kv_heads, head_size) - return (1, 0, 2, 3, 4, 5) + # (num_blocks, num_layers, block_size, num_kv_heads, 2*head_size) + return (1, 0, 3, 2, 4) elif cache_layout == "NHD": - stride_order = (0, 1, 2, 3, 4) + # (num_blocks, block_size, num_kv_heads, 2*head_size) + stride_order = (0, 2, 1, 3) elif cache_layout == "HND" and include_num_layers_dimension: - # (num_blocks, 2, num_kv_heads, num_layers, block_size, head_size) - return (1, 2, 4, 0, 3, 5) + # (num_blocks, num_kv_heads, num_layers, block_size, 2*head_size) + return (1, 2, 0, 3, 4) elif cache_layout == "HND": - stride_order = (0, 1, 3, 2, 4) + # (num_blocks, num_kv_heads, block_size, 2*head_size) + stride_order = (0, 1, 2, 3) else: raise ValueError(f"Unknown cache layout format {cache_layout}.") return stride_order @@ -1580,7 +1599,16 @@ def forward( if attn_metadata.use_cascade: # Cascade attention (rare case). assert attn_metadata.cascade_wrapper is not None - output.copy_(attn_metadata.cascade_wrapper.run(query, kv_cache)) + stride_order = FlashInferBackend.get_kv_cache_stride_order() + if self.is_kvcache_nvfp4: + kv_cache_views = tuple( + cache.permute(*stride_order) for cache in kv_cache.split(self.num_kv_heads, dim=1) + ) + else: + kv_perm = kv_cache.permute(*stride_order) + kv_cache_views = kv_perm.split(self.head_size, dim=-1) + kv_tuple = tuple(canonicalize_singleton_dim_strides(cache) for cache in kv_cache_views) + output.copy_(attn_metadata.cascade_wrapper.run(query, kv_tuple)) return output # When using spec decoding, num_decodes can be < num_decode_tokens @@ -1605,12 +1633,23 @@ def forward( ) kv_cache_permute = fixed - # For NVFP4, the kv_cache last dim is full_dim (data + scale packed). - # Split into correctly-strided data and scale views. + # Split K/V as zero-copy views. NVFP4 stores K/V as separate head + # groups; other dtypes pack K/V in the content dim. + hs = self.head_size nvfp4_kv_data = None nvfp4_kv_block_scales = None if self.is_kvcache_nvfp4: - nvfp4_kv_data, nvfp4_kv_block_scales = nvfp4_kv_cache_split_views(kv_cache_permute) + k_cache, v_cache = kv_cache.split(self.num_kv_heads, dim=1) + kv_cache_tuple = ( + canonicalize_singleton_dim_strides(k_cache.permute(*stride_order)), + canonicalize_singleton_dim_strides(v_cache.permute(*stride_order)), + ) + k_data, k_sf = nvfp4_split_data_scale(kv_cache_tuple[0]) + v_data, v_sf = nvfp4_split_data_scale(kv_cache_tuple[1]) + nvfp4_kv_data = (k_data, v_data) + nvfp4_kv_block_scales = (k_sf, v_sf) + else: + kv_cache_tuple = kv_cache_permute.split(hs, dim=-1) use_dcp = self.dcp_world_size > 1 @@ -1645,7 +1684,7 @@ def forward( prefill_wrapper.run( layer, prefill_query, - kv_cache_permute, + kv_cache_tuple, key[num_decode_tokens:], value[num_decode_tokens:], out=output[num_decode_tokens:], @@ -1658,7 +1697,9 @@ def forward( assert prefill_wrapper._causal == attn_metadata.causal if self.is_kvcache_nvfp4: - kv_cache_permute = nvfp4_kv_data + kv_cache_for_fi = nvfp4_kv_data + else: + kv_cache_for_fi = kv_cache_tuple kv_cache_sf = nvfp4_kv_block_scales if self.is_kvcache_nvfp4 else None # NVFP4 trtllm kernel only supports FP8 output. @@ -1672,7 +1713,7 @@ def forward( prefill_wrapper.run( prefill_query, - kv_cache_permute, + kv_cache_for_fi, q_scale=layer._q_scale_float, k_scale=layer._k_scale_float, v_scale=layer._v_scale_float, @@ -1736,22 +1777,27 @@ def forward( # TRTLLM prefill attention does not support BF16 Q # and fp8 kv cache. So to enable prefill attention # with fp8 kv cache, we can construct a mock block - # and mock kv cache with BF16 KV involved in the prefill - # + # and mock kv cache with BF16 KV involved in the prefill. kv_cache_permute = canonicalize_singleton_dim_strides(kv_cache_permute) kv_strides = kv_cache_permute.stride() assert kv_strides[-1] == 1 and kv_strides[-2] == kv_cache_permute.shape[-1], ( f"KV cache inner dims (block_size, head_size) must be contiguous, got strides {kv_strides}" ) + # fp8 uses (B, H, N, 2*hs); reshape to (B, 2, H, N, hs) + # for the dequant kernel. The dequant kernel handles the + # interleaved K/V block stride. + B_kv, H_kv, N_kv = kv_cache_permute.shape[:3] + kv_cache_5d = kv_cache_permute.view(B_kv, H_kv, N_kv, 2, hs) + kv_cache_5d = kv_cache_5d.permute(0, 3, 1, 2, 4) mock_kv_cache, mock_block_table = trtllm_prefill_attn_kvfp8_dequant( - kv_cache_permute, + kv_cache_5d, block_tables_prefill, layer._k_scale, layer._v_scale, attn_metadata.q_data_type_prefill, ) else: - mock_kv_cache = kv_cache_permute + mock_kv_cache = kv_cache_tuple mock_block_table = block_tables_prefill trtllm_batch_context_with_kv_cache( @@ -1799,7 +1845,9 @@ def forward( assert decode_wrapper._sm_scale == self.scale if self.is_kvcache_nvfp4: - kv_cache_permute = nvfp4_kv_data + kv_cache_for_fi = nvfp4_kv_data + else: + kv_cache_for_fi = kv_cache_tuple kv_cache_sf = nvfp4_kv_block_scales if self.is_kvcache_nvfp4 else None # NVFP4 kernel only supports FP8 output. @@ -1820,7 +1868,7 @@ def forward( ) decode_wrapper.run( decode_query, - kv_cache_permute, + kv_cache_for_fi, q_scale=layer._q_scale_float, k_scale=layer._k_scale_float, v_scale=layer._v_scale_float, @@ -1837,7 +1885,7 @@ def forward( else: decode_wrapper.run( decode_query, - kv_cache_permute, + kv_cache_for_fi, q_scale=layer._q_scale_float, k_scale=layer._k_scale_float, v_scale=layer._v_scale_float, @@ -1913,7 +1961,7 @@ def forward( trtllm_batch_decode_with_kv_cache( query=decode_query, - kv_cache=(nvfp4_kv_data if self.is_kvcache_nvfp4 else kv_cache_permute), + kv_cache=(nvfp4_kv_data if self.is_kvcache_nvfp4 else kv_cache_tuple), workspace_buffer=workspace_buffer, block_tables=block_tables_decode, seq_lens=seq_lens_decode, @@ -1950,8 +1998,14 @@ def do_kv_cache_update( # and value[:num_actual_tokens] because the reshape_and_cache_flash # op uses the slot_mapping's shape to determine the number of # actual tokens. - k_cache = kv_cache[:, 0] - v_cache = kv_cache[:, 1] + if self.is_kvcache_nvfp4: + # (B, 2*H, N, full_dim) -> ((B, N, H, full_dim), + # (B, N, H, full_dim)); + # K heads first, then V heads. + k_cache, v_cache = kv_cache.transpose(1, 2).split(self.num_kv_heads, dim=-2) + else: + # (B, H, N, 2*hs) -> ((B, N, H, hs), (B, N, H, hs)) + k_cache, v_cache = kv_cache.transpose(1, 2).split(self.head_size, dim=-1) torch.ops._C_cache_ops.reshape_and_cache_flash( key, value, diff --git a/aphrodite/v1/attention/backends/flex_attention.py b/aphrodite/v1/attention/backends/flex_attention.py index 172e2ae913..9fd9d86901 100644 --- a/aphrodite/v1/attention/backends/flex_attention.py +++ b/aphrodite/v1/attention/backends/flex_attention.py @@ -124,15 +124,16 @@ def get_kv_cache_shape( head_size: int, cache_dtype_str: str = "auto", ) -> tuple[int, ...]: - return (num_blocks, 2, block_size, num_kv_heads, head_size) + # K and V are packed into the content dim: logical (B, H, N, 2*hs). + return (num_blocks, num_kv_heads, block_size, 2 * head_size) @staticmethod def get_kv_cache_stride_order( include_num_layers_dimension: bool = False, ) -> tuple[int, ...]: if include_num_layers_dimension: - return (1, 0, 3, 2, 4, 5) - return (0, 2, 1, 3, 4) + return (1, 0, 3, 2, 4) + return (0, 2, 1, 3) @staticmethod def get_builder_cls() -> type["FlexAttentionMetadataBuilder"]: @@ -1139,7 +1140,8 @@ def do_kv_cache_update( if self.attn_type == AttentionType.ENCODER_ONLY: return - key_cache, value_cache = kv_cache.unbind(1) + # (B, H, N, 2*hs) -> ((B, N, H, hs), (B, N, H, hs)) + key_cache, value_cache = kv_cache.transpose(1, 2).split(self.head_size, dim=-1) torch.ops._C_cache_ops.reshape_and_cache_flash( key, value, @@ -1170,7 +1172,7 @@ def forward( key: shape = [num_tokens, num_kv_heads, head_size] value: shape = [num_tokens, num_kv_heads, head_size] kv_cache: shape = - [num_blocks, 2, block_size, num_kv_heads, head_size] + [num_blocks, num_kv_heads, block_size, 2 * head_size] attn_metadata: Metadata for attention. Returns: shape = [num_tokens, num_heads * head_size] @@ -1234,9 +1236,11 @@ def forward( else: assert self.attn_type == AttentionType.DECODER - key_cache, value_cache = kv_cache.unbind(1) + kv_cache = kv_cache.transpose(1, 2) + hs = self.head_size + key_cache, value_cache = kv_cache.split(hs, dim=-1) - # Flatten (num_blocks, block_size) into a single token dim + # Flatten (num_blocks, block_size) into a single token dim. key_cache = key_cache.view(-1, self.num_kv_heads, self.head_size) value_cache = value_cache.view(-1, self.num_kv_heads, self.head_size) query, key_tensor, value_tensor = map( diff --git a/aphrodite/v1/attention/backends/hpc_attn.py b/aphrodite/v1/attention/backends/hpc_attn.py index 3eee1fdec7..ef80090c62 100644 --- a/aphrodite/v1/attention/backends/hpc_attn.py +++ b/aphrodite/v1/attention/backends/hpc_attn.py @@ -286,15 +286,15 @@ def get_kv_cache_shape( head_size: int, cache_dtype_str: str = "auto", ) -> tuple[int, ...]: - return (num_blocks, 2, block_size, num_kv_heads, head_size) + return (num_blocks, num_kv_heads, block_size, 2 * head_size) @staticmethod def get_kv_cache_stride_order( include_num_layers_dimension: bool = False, ) -> tuple[int, ...]: if include_num_layers_dimension: - return (1, 0, 2, 3, 4, 5) - return (0, 1, 2, 3, 4) + return (1, 0, 3, 2, 4) + return (0, 2, 1, 3) @classmethod def get_supported_head_sizes(cls) -> list[int]: @@ -431,11 +431,12 @@ def forward( # Write KV cache if not already done by HpcRopeNorm. if self.kv_sharing_target_layer_name is None and not hpc_kv_written: + key_cache, value_cache = kv_cache.transpose(1, 2).split(self.head_size, dim=-1) torch.ops._C_cache_ops.reshape_and_cache_flash( key, value, - kv_cache[:, 0], - kv_cache[:, 1], + key_cache, + value_cache, attn_metadata.slot_mapping, self.kv_cache_dtype, layer._k_scale, @@ -445,6 +446,7 @@ def forward( if self.use_fp8: torch_dtype = _get_fp8_dtype_for_kv_cache(self.kv_cache_dtype) kv_cache = kv_cache.view(torch_dtype) + key_cache, value_cache = kv_cache.transpose(1, 2).split(self.head_size, dim=-1) if self.use_fp8: if not hpc_kv_written: @@ -476,8 +478,8 @@ def forward( if self.use_fp8: hpc.attention_with_kvcache_prefill_fp8( q_prefill, - kv_cache[:, 0], - kv_cache[:, 1], + key_cache, + value_cache, hpc_prefill_q_scale, k_scale, v_scale, @@ -491,8 +493,8 @@ def forward( else: hpc.attention_with_kvcache_prefill_bf16( q_prefill, - kv_cache[:, 0], - kv_cache[:, 1], + key_cache, + value_cache, cu_seqlens_prefill, block_table_prefill, seq_lens_prefill, @@ -513,8 +515,8 @@ def forward( if self.use_fp8: hpc.attention_decode_fp8( q_decode, - kv_cache[:, 0], - kv_cache[:, 1], + key_cache, + value_cache, block_table_decode, num_seq_kvcache, hpc_decode_q_scale, @@ -534,8 +536,8 @@ def forward( else: hpc.attention_decode_bf16( q_decode, - kv_cache[:, 0], - kv_cache[:, 1], + key_cache, + value_cache, block_table_decode, num_seq_kvcache, mtp=mtp, diff --git a/aphrodite/v1/attention/backends/rocm_aiter_fa.py b/aphrodite/v1/attention/backends/rocm_aiter_fa.py index b3bb1d30b5..0e4c0c79d1 100644 --- a/aphrodite/v1/attention/backends/rocm_aiter_fa.py +++ b/aphrodite/v1/attention/backends/rocm_aiter_fa.py @@ -700,7 +700,8 @@ def get_kv_cache_shape( ) -> tuple[int, ...]: if block_size % 16 != 0: raise ValueError("Block size must be a multiple of 16.") - return (num_blocks, 2, block_size, num_kv_heads, head_size) + # K and V are packed into the content dim: logical (B, H, N, 2*hs). + return (num_blocks, num_kv_heads, block_size, 2 * head_size) @classmethod def supports_compute_capability(cls, capability: DeviceCapability) -> bool: @@ -988,7 +989,8 @@ def forward( # Whenever making a change in this method, please benchmark the # performance to make sure it does not introduce any overhead. num_actual_tokens = attn_metadata.num_actual_tokens - key_cache, value_cache = kv_cache.unbind(1) + # (B, H, N, 2*hs) -> ((B, N, H, hs), (B, N, H, hs)) + key_cache, value_cache = kv_cache.transpose(1, 2).split(self.head_size, dim=-1) if is_quantized_kv_cache(self.kv_cache_dtype): key_cache = key_cache.view(current_platform.fp8_dtype()) @@ -1279,7 +1281,8 @@ def do_kv_cache_update( kv_cache: torch.Tensor, slot_mapping: torch.Tensor, ): - key_cache, value_cache = kv_cache.unbind(1) + # (B, H, N, 2*hs) -> ((B, N, H, hs), (B, N, H, hs)) + key_cache, value_cache = kv_cache.transpose(1, 2).split(self.head_size, dim=-1) # key and value may be None in the case of cross attention. They are # calculated once based on the output from the encoder and then cached @@ -1340,7 +1343,8 @@ def do_rope_and_kv_cache_update( kv_cache: torch.Tensor, layer_slot_mapping: torch.Tensor, ): - key_cache, value_cache = kv_cache.unbind(1) + # (B, H, N, 2*hs) -> ((B, N, H, hs), (B, N, H, hs)) + key_cache, value_cache = kv_cache.transpose(1, 2).split(self.head_size, dim=-1) flash_layout = True is_fp8_kv_cache = is_quantized_kv_cache(self.kv_cache_dtype) diff --git a/aphrodite/v1/attention/backends/rocm_aiter_unified_attn.py b/aphrodite/v1/attention/backends/rocm_aiter_unified_attn.py index 51566d1c8c..14903f8a25 100644 --- a/aphrodite/v1/attention/backends/rocm_aiter_unified_attn.py +++ b/aphrodite/v1/attention/backends/rocm_aiter_unified_attn.py @@ -87,7 +87,8 @@ def get_kv_cache_shape( ) -> tuple[int, ...]: if block_size % 16 != 0: raise ValueError("Block size must be a multiple of 16.") - return (num_blocks, 2, block_size, num_kv_heads, head_size) + # K and V are packed into the content dim: logical (B, H, N, 2*hs). + return (num_blocks, num_kv_heads, block_size, 2 * head_size) @staticmethod def use_cascade_attention(*args, **kwargs) -> bool: @@ -146,10 +147,8 @@ def __init__( self.supports_quant_query_input = True def _split_kv_cache(self, kv_cache: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: - # Blocks-first ``(num_blocks, 2, ...)``. The model runner normalizes any - # shared decoder/cross-attention allocation to this layout, so no - # per-backend restriding is needed here. - return kv_cache.unbind(1) + # (B, H, N, 2*hs) -> ((B, N, H, hs), (B, N, H, hs)) + return kv_cache.transpose(1, 2).split(self.head_size, dim=-1) def forward( self, diff --git a/aphrodite/v1/attention/backends/triton_attn.py b/aphrodite/v1/attention/backends/triton_attn.py index 818d13f42c..2c5313f585 100644 --- a/aphrodite/v1/attention/backends/triton_attn.py +++ b/aphrodite/v1/attention/backends/triton_attn.py @@ -314,6 +314,7 @@ def get_kv_cache_shape( ) -> tuple[int, ...]: if block_size % 16 != 0: raise ValueError("Block size must be a multiple of 16.") + # K and V are packed into the content dim: logical (B, H, N, 2*hs). if kv_cache_uses_per_token_head_scales(cache_dtype_str): # Pad the head dim by sizeof(float32)/sizeof(cache_dtype) so the # per-(token, head) scale fits inline after the quantized data; @@ -331,26 +332,30 @@ def get_kv_cache_shape( data_head_size = head_size // 2 else: data_head_size = head_size - return (num_blocks, 2, block_size, num_kv_heads, data_head_size + scale_pad) - return (num_blocks, 2, block_size, num_kv_heads, head_size) + padded_hs = data_head_size + scale_pad + return (num_blocks, num_kv_heads, block_size, 2 * padded_hs) + return (num_blocks, num_kv_heads, block_size, 2 * head_size) @staticmethod def get_kv_cache_stride_order( include_num_layers_dimension: bool = False, ) -> tuple[int, ...]: - # `stride_order` indicates the permutation that gets - # us from `get_kv_cache_shape` to the actual memory layout we want. + # `stride_order` indicates the permutation that gets us from + # `get_kv_cache_shape` (logical (B, H, N, 2*hs)) to the actual memory + # layout we want. cache_layout = get_kv_cache_layout() if cache_layout == "NHD" and include_num_layers_dimension: - # (num_blocks, num_layers, 2, block_size, num_kv_heads, head_size) - return (1, 0, 2, 3, 4, 5) + # (num_blocks, num_layers, block_size, num_kv_heads, 2*head_size) + return (1, 0, 3, 2, 4) elif cache_layout == "NHD": - stride_order = (0, 1, 2, 3, 4) + # (num_blocks, block_size, num_kv_heads, 2*head_size) + stride_order = (0, 2, 1, 3) elif cache_layout == "HND" and include_num_layers_dimension: - # (num_blocks, num_kv_heads, num_layers, 2, block_size, head_size) - return (1, 4, 0, 2, 3, 5) + # (num_blocks, num_kv_heads, num_layers, block_size, 2*head_size) + return (1, 2, 0, 3, 4) elif cache_layout == "HND": - stride_order = (0, 1, 3, 2, 4) + # (num_blocks, num_kv_heads, block_size, 2*head_size) + stride_order = (0, 1, 2, 3) else: raise ValueError(f"Unknown cache layout: {cache_layout}") return stride_order @@ -400,12 +405,16 @@ class TritonAttentionImpl(AttentionImpl): _v_scale_cache: torch.Tensor | None = None def _ensure_scale_caches(self, kv_cache: torch.Tensor) -> None: - """Extract per-head scale views from the padded head dimension. + """Extract per-head scale views from the padded content dimension. - The KV cache shape is ``(num_blocks, 2, block_size, nkv, hs+pad)`` - where ``pad = sizeof(float32) / sizeof(cache_dtype)``. The last - ``pad`` elements of each head hold one float32 scale. We create - strided float32 views over those bytes. + The KV cache is packed as logical shape + ``(num_blocks, nkv, block_size, 2 * (hs + pad))`` where + ``pad = sizeof(float32) / sizeof(cache_dtype)``. The content dim holds + ``[K(hs) | K_scale(pad) | V(hs) | V_scale(pad)]`` per (head, slot); the + last ``pad`` elements of each half hold one float32 scale. We create + strided float32 views over those bytes. ``kv_cache`` must be the + packed logical tensor (call before any transpose), but may have HND or + NHD physical strides. Scale shape: ``(num_blocks, block_size, num_kv_heads)`` """ @@ -413,39 +422,43 @@ def _ensure_scale_caches(self, kv_cache: torch.Tensor) -> None: return from aphrodite.utils.torch_utils import get_dtype_size - num_blocks, _, block_size, nkv, padded_hs = kv_cache.shape + num_blocks, nkv, block_size, content = kv_cache.shape dtype_sz = kv_cache.element_size() scale_pad = get_dtype_size(torch.float32) // dtype_sz # e.g. 4 + padded_hs = content // 2 hs = padded_hs - scale_pad raw = kv_cache.untyped_storage() base_f32 = torch.tensor([], dtype=torch.float32, device=kv_cache.device).set_(raw) - # In the raw bytes, each (block, kv_half, slot, head) occupies - # padded_hs * dtype_sz bytes. The scale float32 sits at byte - # offset hs * dtype_sz within that region. - kv_half_bytes = block_size * nkv * padded_hs * dtype_sz - full_block_f32 = 2 * kv_half_bytes // 4 # stride between blocks - slot_f32 = nkv * padded_hs * dtype_sz // 4 # stride between slots - head_f32 = padded_hs * dtype_sz // 4 # stride between heads - scale_off_f32 = hs * dtype_sz // 4 # offset to scale within head + def to_f32_units(elements: int) -> int: + nbytes = elements * dtype_sz + assert nbytes % 4 == 0 + return nbytes // 4 - # K scales: kv_half=0 + strides = kv_cache.stride() + block_f32 = to_f32_units(strides[0]) + head_f32 = to_f32_units(strides[1]) + slot_f32 = to_f32_units(strides[2]) + base_off_f32 = to_f32_units(kv_cache.storage_offset()) + k_scale_off_f32 = base_off_f32 + to_f32_units(hs) + v_scale_off_f32 = base_off_f32 + to_f32_units(padded_hs + hs) + + # K scales (first content half) self._k_scale_cache = torch.as_strided( base_f32, size=(num_blocks, block_size, nkv), - stride=(full_block_f32, slot_f32, head_f32), - storage_offset=scale_off_f32, + stride=(block_f32, slot_f32, head_f32), + storage_offset=k_scale_off_f32, ) self._k_scale_cache.fill_(1.0) - # V scales: kv_half=1, offset by kv_half_bytes - v_base_f32 = kv_half_bytes // 4 + # V scales (second content half) self._v_scale_cache = torch.as_strided( base_f32, size=(num_blocks, block_size, nkv), - stride=(full_block_f32, slot_f32, head_f32), - storage_offset=v_base_f32 + scale_off_f32, + stride=(block_f32, slot_f32, head_f32), + storage_offset=v_scale_off_f32, ) self._v_scale_cache.fill_(1.0) @@ -558,7 +571,7 @@ def forward( key: shape = [num_tokens, num_kv_heads, head_size] value: shape = [num_tokens, num_kv_heads, head_size] kv_cache: shape = - [num_blocks, 2, block_size, num_kv_heads, head_size] + [num_blocks, num_kv_heads, block_size, 2 * head_size] attn_metadata: Metadata for attention. Returns: shape = [num_tokens, num_heads * head_size] @@ -598,6 +611,7 @@ def forward( layer, ) + # KV cache arrives in logical (B, H, N, 2*hs) order. # Per-token-head quantized KV cache: handled by the core unified # kernel, which dequantizes per-(token, head) inline via constexpr # branches (INT8 / FP8) and dispatches to the packed INT4 kernel. @@ -608,7 +622,9 @@ def forward( q_descale = k_descale = v_descale = None # FP8 per-tensor / auto path (original flow). else: - key_cache, value_cache = kv_cache.unbind(1) + kv_cache = kv_cache.transpose(1, 2) + hs = self.head_size + key_cache, value_cache = kv_cache.split(hs, dim=-1) if is_quantized_kv_cache(self.kv_cache_dtype) and key_cache.dtype != self.fp8_dtype: key_cache = key_cache.view(self.fp8_dtype) value_cache = value_cache.view(self.fp8_dtype) @@ -682,7 +698,8 @@ def forward( def _pth_key_value_caches(self, kv_cache: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: """Per-token-head K/V cache views (ensures scale caches; FP8 retyped).""" self._ensure_scale_caches(kv_cache) - key_cache, value_cache = kv_cache.unbind(1) + padded_hs = kv_cache.shape[-1] // 2 + key_cache, value_cache = kv_cache.transpose(1, 2).split(padded_hs, dim=-1) if self._kv_quant_mode == KVQuantMode.FP8_PER_TOKEN_HEAD: key_cache = key_cache.view(self.fp8_dtype) value_cache = value_cache.view(self.fp8_dtype) @@ -747,13 +764,9 @@ def do_kv_cache_update( return # Reshape the input keys and values and store them in the cache. if self._is_per_token_head_quant: - self._ensure_scale_caches(kv_cache) - key_cache, value_cache = kv_cache.unbind(1) + key_cache, value_cache = self._pth_key_value_caches(kv_cache) k_scale_cache = self._k_scale_cache v_scale_cache = self._v_scale_cache - if self._kv_quant_mode == KVQuantMode.FP8_PER_TOKEN_HEAD: - key_cache = key_cache.view(self.fp8_dtype) - value_cache = value_cache.view(self.fp8_dtype) triton_reshape_and_cache_flash_per_token_head_quant( key, value, @@ -766,7 +779,8 @@ def do_kv_cache_update( ) return # For decoder and cross-attention, use KV cache as before. - key_cache, value_cache = kv_cache.unbind(1) + # (B, H, N, 2*hs) -> ((B, N, H, hs), (B, N, H, hs)) + key_cache, value_cache = kv_cache.transpose(1, 2).split(self.head_size, dim=-1) if is_quantized_kv_cache(self.kv_cache_dtype): key_cache = key_cache.view(self.fp8_dtype) value_cache = value_cache.view(self.fp8_dtype) @@ -798,7 +812,8 @@ def do_rope_and_kv_cache_update( kv_cache: torch.Tensor, layer_slot_mapping: torch.Tensor, ): - key_cache, value_cache = kv_cache.unbind(1) + # (B, H, N, 2*hs) -> ((B, N, H, hs), (B, N, H, hs)) + key_cache, value_cache = kv_cache.transpose(1, 2).split(self.head_size, dim=-1) flash_layout = True is_fp8_kv_cache = is_quantized_kv_cache(self.kv_cache_dtype) diff --git a/aphrodite/v1/attention/backends/triton_attn_diffkv.py b/aphrodite/v1/attention/backends/triton_attn_diffkv.py index b0a372ff72..8fece8694e 100644 --- a/aphrodite/v1/attention/backends/triton_attn_diffkv.py +++ b/aphrodite/v1/attention/backends/triton_attn_diffkv.py @@ -2,12 +2,9 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Triton attention backend with different K/V head dimensions (DiffKV). -The KV cache layout is identical to ``FlashAttentionDiffKVBackend`` — K -and V are packed along the last dim: - - [num_blocks, block_size, num_kv_heads, head_size_qk + head_size_v] - -so existing helpers (``triton_reshape_and_cache_flash_diffkv``) are reused. +The KV cache layout is identical to ``FlashAttentionDiffKVBackend``: K and V +are packed along the last dim in the logical shape +``[num_blocks, num_kv_heads, block_size, head_size_qk + head_size_v]``. """ from typing import ClassVar @@ -106,10 +103,12 @@ def get_kv_cache_shape( ) -> tuple[int, ...]: if block_size % 16 != 0: raise ValueError("Block size must be a multiple of 16.") + # Logical (blocks-first, head-major) layout: K and V (with their + # different head sizes) packed in the content dim. return ( num_blocks, - block_size, num_kv_heads, + block_size, head_size + TritonAttentionDiffKVBackend.head_size_v, ) @@ -117,19 +116,22 @@ def get_kv_cache_shape( def get_kv_cache_stride_order( include_num_layers_dimension: bool = False, ) -> tuple[int, ...]: + # `stride_order` indicates the permutation that gets us from + # `get_kv_cache_shape` (logical (B, H, N, C_k+C_v)) to the actual + # memory layout we want. cache_layout = get_kv_cache_layout() if cache_layout == "NHD" and include_num_layers_dimension: - # (num_blocks, num_layers, block_size, - # num_kv_heads, head_size + head_size_v) - return (1, 0, 2, 3, 4) + # (num_blocks, num_layers, block_size, num_kv_heads, C_k+C_v) + return (1, 0, 3, 2, 4) elif cache_layout == "NHD": - return (0, 1, 2, 3) + # (num_blocks, block_size, num_kv_heads, C_k+C_v) + return (0, 2, 1, 3) elif cache_layout == "HND" and include_num_layers_dimension: - # (num_blocks, num_kv_heads, num_layers, - # block_size, head_size + head_size_v) - return (1, 3, 0, 2, 4) + # (num_blocks, num_kv_heads, num_layers, block_size, C_k+C_v) + return (1, 2, 0, 3, 4) elif cache_layout == "HND": - return (0, 2, 1, 3) + # (num_blocks, num_kv_heads, block_size, C_k+C_v) + return (0, 1, 2, 3) else: raise ValueError(f"Unknown cache layout format {cache_layout}.") @@ -169,13 +171,12 @@ def do_kv_cache_update( kv_cache: torch.Tensor, slot_mapping: torch.Tensor, ) -> None: - # Cache is packed [..., head_size_qk + head_size_v]; the diffkv - # reshape kernel writes K to [..., :head_size_qk] and V to - # [..., head_size_qk:hqk+hv]. + # Cache is logical (B, H, N, C); the diffkv reshape kernel expects + # (B, N, H, C). triton_reshape_and_cache_flash_diffkv( key, value, - kv_cache, + kv_cache.transpose(1, 2), slot_mapping, self.kv_cache_dtype, layer._k_scale, @@ -204,7 +205,7 @@ def forward( query: [num_tokens, num_heads, head_size_qk] key: [num_tokens, num_kv_heads, head_size_qk] value: [num_tokens, num_kv_heads, head_size_v] - kv_cache: [num_blocks, block_size, num_kv_heads, + kv_cache: [num_blocks, num_kv_heads, block_size, head_size_qk + head_size_v] output: [num_tokens, num_heads, head_size_v] """ @@ -220,8 +221,8 @@ def forward( head_size_qk = self.head_size head_size_v = TritonAttentionDiffKVBackend.head_size_v - # Slice the packed cache into K / V views. Strides on dims 0/1/2 - # match the original cache; dim 3 stays contiguous (stride 1). + # Triton DiffKV kernels consume (B, N, H, D) cache views. + kv_cache = kv_cache.transpose(1, 2) key_cache = kv_cache[..., :head_size_qk] value_cache = kv_cache[..., head_size_qk : head_size_qk + head_size_v] diff --git a/aphrodite/v1/attention/backends/turboquant_attn.py b/aphrodite/v1/attention/backends/turboquant_attn.py index a4d37fb84a..357adc4438 100644 --- a/aphrodite/v1/attention/backends/turboquant_attn.py +++ b/aphrodite/v1/attention/backends/turboquant_attn.py @@ -141,9 +141,10 @@ def get_kv_cache_shape( Standard attention backends use (2, num_blocks, block_size, num_kv_heads, head_dim) with a leading 2 to separate K and V. TurboQuant packs K+V - into a single interleaved slot per head per position, so the cache is: + into a single interleaved slot per head per position. The logical + (blocks-first, head-major) shape is: - (num_blocks, block_size, num_kv_heads, slot_size_aligned) + (num_blocks, num_kv_heads, block_size, slot_size_aligned) Each slot = [key_packed | value_packed | padding]. This is safe because TQ has its own get_kv_cache_shape override and @@ -159,7 +160,7 @@ def get_kv_cache_shape( ) tq_config = TurboQuantConfig.from_cache_dtype(cache_dtype_str, head_size) - return (num_blocks, block_size, num_kv_heads, tq_config.slot_size_aligned) + return (num_blocks, num_kv_heads, block_size, tq_config.slot_size_aligned) @classmethod def supports_kv_cache_dtype(cls, kv_cache_dtype: CacheDType | None) -> bool: @@ -410,6 +411,8 @@ def do_kv_cache_update( k = key[:N].view(N, self.num_kv_heads, self.head_size) v = value[:N].view(N, self.num_kv_heads, self.head_size) + # (B, H, N, C) -> (B, N, H, C) for TQ kernels + kv_cache = kv_cache.transpose(1, 2) self._store_kv(k, v, kv_cache, slot_mapping, layer) def forward( @@ -437,6 +440,9 @@ def forward( if attn_metadata is None: return output.fill_(0) + # (B, H, N, C) -> (B, N, H, C) for TQ kernels + kv_cache = kv_cache.transpose(1, 2) + # Slice to actual tokens N = attn_metadata.num_actual_tokens if N <= 0: diff --git a/aphrodite/v1/attention/ops/triton_reshape_and_cache_flash.py b/aphrodite/v1/attention/ops/triton_reshape_and_cache_flash.py index 9525d357fb..43959f1679 100644 --- a/aphrodite/v1/attention/ops/triton_reshape_and_cache_flash.py +++ b/aphrodite/v1/attention/ops/triton_reshape_and_cache_flash.py @@ -436,6 +436,7 @@ def reshape_and_cache_kernel_flash_diffkv( value_stride: tl.int64, block_stride: tl.int64, page_stride: tl.int64, + head_stride: tl.int64, num_heads: tl.constexpr, head_size_k: tl.constexpr, head_size_v: tl.constexpr, @@ -460,7 +461,7 @@ def reshape_and_cache_kernel_flash_diffkv( src_key_idx = token_idx * key_stride + tile_i * head_size_k src_value_idx = token_idx * value_stride + tile_i * head_size_v - tgt_idx = block_idx * block_stride + block_offset * page_stride + tile_i * (head_size_k + head_size_v) + tgt_idx = block_idx * block_stride + block_offset * page_stride + tile_i * head_stride # [TILE_SIZE] key_load = tl.load(key_ptr + src_key_idx + tile_offs, mask=tile_offs < head_size_k) @@ -499,7 +500,7 @@ def reshape_and_cache_kernel_flash_diffkv( def triton_reshape_and_cache_flash_diffkv( key: torch.Tensor, # [num_tokens, num_heads, head_size] value: torch.Tensor, # [num_tokens, num_heads, head_size_v] - # [num_blocks, block_size, num_heads, head_size + head_size_v] + # Strided [num_blocks, block_size, num_heads, head_size + head_size_v]. kv_cache: torch.Tensor, slot_mapping: torch.Tensor, # [num_tokens] kv_cache_dtype: str, # "auto", "fp8" @@ -515,6 +516,7 @@ def triton_reshape_and_cache_flash_diffkv( v_stride = value.stride()[0] block_stride = kv_cache.stride()[0] page_stride = kv_cache.stride()[1] + head_stride = kv_cache.stride()[2] kv_cache_torch_dtype = current_platform.fp8_dtype() if is_quantized_kv_cache(kv_cache_dtype) else kv_cache.dtype @@ -552,6 +554,7 @@ def triton_reshape_and_cache_flash_diffkv( value_stride=v_stride, block_stride=block_stride, page_stride=page_stride, + head_stride=head_stride, num_heads=num_heads, head_size_k=head_size_k, head_size_v=head_size_v, diff --git a/aphrodite/v1/attention/ops/triton_turboquant_store.py b/aphrodite/v1/attention/ops/triton_turboquant_store.py index a5cd523294..90073540bd 100644 --- a/aphrodite/v1/attention/ops/triton_turboquant_store.py +++ b/aphrodite/v1/attention/ops/triton_turboquant_store.py @@ -176,7 +176,7 @@ def _tq_fused_store_fp8( # ── FP8 KEY: cast to FP8 in-kernel and store ───────────────── d_offs = tl.arange(0, BLOCK_D) d_mask = d_offs < D - k_vals = tl.load(Key_ptr + base + d_offs, mask=d_mask, other=0.0) + k_vals = tl.load(Key_ptr + base + d_offs, mask=d_mask, other=0.0).to(tl.float32) k_fp8 = k_vals.to(tl.float8e4b15) if FP8_E4B15 else k_vals.to(tl.float8e4nv) k_bytes = k_fp8.to(tl.uint8, bitcast=True) tl.store(KV_cache_ptr + slot_base + d_offs, k_bytes, mask=d_mask) @@ -371,7 +371,7 @@ def triton_turboquant_store( _tq_fused_store_fp8[grid]( k_flat, v_flat, - kv_cache.view(-1), + kv_cache, slot_mapping, stride_cache_block=stride_block, stride_cache_pos=stride_pos, @@ -407,7 +407,7 @@ def triton_turboquant_store( norms.squeeze(1), v_flat, midpoints, - kv_cache.view(-1), + kv_cache, slot_mapping, stride_cache_block=stride_block, stride_cache_pos=stride_pos, diff --git a/aphrodite/v1/core/block_pool.py b/aphrodite/v1/core/block_pool.py index 3b353662ce..79276ea541 100644 --- a/aphrodite/v1/core/block_pool.py +++ b/aphrodite/v1/core/block_pool.py @@ -14,8 +14,6 @@ from aphrodite.v1.core.kv_cache_metrics import KVCacheMetricsCollector from aphrodite.v1.core.kv_cache_utils import ( BlockHash, - BlockHashList, - BlockHashListWithBlockSize, BlockHashWithGroupId, ExternalBlockHash, FreeKVCacheBlockQueue, @@ -25,6 +23,7 @@ get_group_id, make_block_hash_with_group_id, maybe_convert_block_hash, + resolve_block_hashes, ) from aphrodite.v1.request import Request @@ -251,14 +250,7 @@ def cache_full_blocks( return new_full_blocks = blocks[num_cached_blocks:num_full_blocks] assert block_mask is None or len(block_mask) == len(new_full_blocks) - if block_size == self.hash_block_size: - # Common case. - block_hashes: BlockHashList = request.block_hashes - else: - # block_size is a multiple of hash_block_size. This happens when - # different KV cache groups have different block sizes. - assert block_size % self.hash_block_size == 0 - block_hashes = BlockHashListWithBlockSize(request.block_hashes, self.hash_block_size, block_size) + block_hashes = resolve_block_hashes(request.block_hashes, self.hash_block_size, block_size) assert len(block_hashes) >= num_full_blocks new_block_hashes = block_hashes[num_cached_blocks:] @@ -315,19 +307,79 @@ def cache_full_blocks( extra_keys_list.append(extra_keys) self.kv_event_queue.append( - BlockStored( + self._build_block_stored_event( + request, block_hashes=new_hashes, parent_block_hash=parent_block_hash, - token_ids=request.all_token_ids[start_token_idx:end_token_idx], + start_token_idx=start_token_idx, + end_token_idx=end_token_idx, block_size=block_size, - lora_id=request.lora_request.adapter_id if request.lora_request else None, - medium=MEDIUM_GPU, - lora_name=request.lora_request.name if request.lora_request else None, - extra_keys=extra_keys_list if extra_keys_list else None, - group_idx=kv_cache_group_id, + kv_cache_group_id=kv_cache_group_id, + extra_keys_list=extra_keys_list, ) ) + def _build_block_stored_event( + self, + request: Request, + block_hashes: list[ExternalBlockHash] | None, + parent_block_hash: ExternalBlockHash | None, + start_token_idx: int, + end_token_idx: int, + block_size: int, + kv_cache_group_id: int, + extra_keys_list: list[tuple[Any, ...] | None], + ) -> BlockStored: + """Build a ``BlockStored`` KV event for ``request``.""" + return BlockStored( + block_hashes=block_hashes, + parent_block_hash=parent_block_hash, + token_ids=request.all_token_ids[start_token_idx:end_token_idx], + block_size=block_size, + lora_id=request.lora_request.adapter_id if request.lora_request else None, + medium=MEDIUM_GPU, + lora_name=request.lora_request.name if request.lora_request else None, + extra_keys=extra_keys_list if extra_keys_list else None, + group_idx=kv_cache_group_id, + ) + + def emit_cached_block_events( + self, + request: Request, + num_cached_blocks: int, + block_size: int, + kv_cache_group_id: int, + ) -> None: + """Generate ``BlockStored`` events for prefix-cache-reused blocks.""" + if not self.enable_kv_cache_events or num_cached_blocks == 0: + return + + block_hashes = resolve_block_hashes(request.block_hashes, self.hash_block_size, block_size) + assert len(block_hashes) >= num_cached_blocks + + cached_hashes: list[ExternalBlockHash] = [] + extra_keys_list: list[tuple[Any, ...] | None] = [] + curr_mm_idx = 0 + for i in range(num_cached_blocks): + block_start = i * block_size + block_end = block_start + block_size + cached_hashes.append(maybe_convert_block_hash(block_hashes[i])) + extra_keys, curr_mm_idx = generate_block_hash_extra_keys(request, block_start, block_end, curr_mm_idx) + extra_keys_list.append(extra_keys) + + self.kv_event_queue.append( + self._build_block_stored_event( + request, + block_hashes=cached_hashes, + parent_block_hash=None, + start_token_idx=0, + end_token_idx=num_cached_blocks * block_size, + block_size=block_size, + kv_cache_group_id=kv_cache_group_id, + extra_keys_list=extra_keys_list, + ) + ) + def cache_partial_block( self, request: Request, @@ -487,6 +539,18 @@ def _insert_block_hash( self.cached_block_hashes_by_block.setdefault(block.block_id, set()).add(block_hash_with_group_id) self.cached_block_hash_to_block.insert(block_hash_with_group_id, block) + def move_block_hashes( + self, + src_block: KVCacheBlock, + dst_block: KVCacheBlock, + ) -> None: + """Re-point ``src_block``'s prefix-cache entries to ``dst_block``.""" + assert dst_block.block_hash is None + assert dst_block.block_id not in self.cached_block_hashes_by_block + num_tokens = src_block.block_hash_num_tokens + for block_hash in self._remove_cached_block_hashes(src_block): + self._insert_block_hash(block_hash, dst_block, num_tokens=num_tokens) + def get_new_blocks(self, num_blocks: int) -> list[KVCacheBlock]: """Get new blocks from the free block pool. diff --git a/aphrodite/v1/core/kv_cache_coordinator.py b/aphrodite/v1/core/kv_cache_coordinator.py index 2296818296..093cbfbaa0 100644 --- a/aphrodite/v1/core/kv_cache_coordinator.py +++ b/aphrodite/v1/core/kv_cache_coordinator.py @@ -5,12 +5,11 @@ from typing import NamedTuple from aphrodite import envs +from aphrodite.utils.math_utils import cdiv from aphrodite.v1.core.block_pool import BlockPool from aphrodite.v1.core.kv_cache_metrics import KVCacheMetricsCollector from aphrodite.v1.core.kv_cache_utils import ( BlockHash, - BlockHashList, - BlockHashListWithBlockSize, KVCacheBlock, ) from aphrodite.v1.core.single_type_kv_cache_manager import ( @@ -127,6 +126,7 @@ def get_num_blocks_to_allocate( new_computed_blocks: tuple[Sequence[KVCacheBlock], ...], num_encoder_tokens: int, total_computed_tokens: int, + num_local_computed_tokens: int, num_tokens_main_model: int, apply_admission_cap: bool = False, ) -> int: @@ -163,6 +163,7 @@ def get_num_blocks_to_allocate( num_encoder_tokens, [], 0, + 0, num_encoder_tokens, apply_admission_cap=apply_admission_cap, ) @@ -172,6 +173,7 @@ def get_num_blocks_to_allocate( num_tokens, new_computed_blocks[i], total_computed_tokens, + num_local_computed_tokens, num_tokens_main_model, apply_admission_cap=apply_admission_cap, ) @@ -460,7 +462,7 @@ def find_longest_cache_hit( block_hashes: list[BlockHash], max_cache_hit_length: int, ) -> tuple[tuple[list[KVCacheBlock], ...], int]: - hit_blocks = self.single_type_managers[0].find_longest_cache_hit( + hit_blocks, hit_length = self.single_type_managers[0].find_longest_cache_hit( block_hashes=block_hashes, max_length=max_cache_hit_length, kv_cache_group_ids=[0], @@ -471,7 +473,7 @@ def find_longest_cache_hit( dcp_world_size=self.dcp_world_size, pcp_world_size=self.pcp_world_size, ) - return hit_blocks, len(hit_blocks[0]) * self.block_size + return hit_blocks, hit_length class SpecGroup(NamedTuple): @@ -532,8 +534,20 @@ def __init__( ) assert dcp_world_size == 1, "DCP not support hybrid attn now." assert pcp_world_size == 1, "PCP not support hybrid attn now." + self.dcp_world_size = dcp_world_size + self.pcp_world_size = pcp_world_size + self.enable_partial_hash_hits = dcp_world_size == 1 and any( + isinstance(g.kv_cache_spec, MambaSpec) + and g.kv_cache_spec.mamba_cache_mode == "align" + and g.kv_cache_spec.block_size > hash_block_size + for g in kv_cache_config.kv_cache_groups + ) self.verify_and_split_kv_cache_groups() + @property + def _cache_hit_alignment_tokens(self) -> int: + return self.hash_block_size if self.enable_partial_hash_hits else self.scheduler_block_size + def verify_and_split_kv_cache_groups(self) -> None: """ Groups KV cache groups by their spec type for efficient batch processing @@ -569,12 +583,10 @@ def verify_and_split_kv_cache_groups(self) -> None: self.single_type_managers[gid].use_eagle = True def cache_blocks(self, request: Request, num_computed_tokens: int) -> None: - # Cache hits in this coordinator are always a multiple of - # ``scheduler_block_size`` tokens (see ``find_longest_cache_hit``). - # Within an aligned region, SWA groups may only consult a subset of blocks - # per ``scheduler_block_size``-segment so the unused blocks also stay - # out of the prefix-cache hash map. - aligned_num_computed_tokens = num_computed_tokens // self.scheduler_block_size * self.scheduler_block_size + if self.enable_partial_hash_hits: + aligned_num_computed_tokens = num_computed_tokens + else: + aligned_num_computed_tokens = num_computed_tokens // self.scheduler_block_size * self.scheduler_block_size for manager in self.single_type_managers: num_tokens_to_cache = aligned_num_computed_tokens # EAGLE groups match one block past each aligned boundary and drop @@ -617,15 +629,11 @@ def find_longest_cache_hit( - The number of tokens of the longest cache hit. """ - def _get_block_hashes(kv_cache_spec: KVCacheSpec) -> BlockHashList: - if kv_cache_spec.block_size == self.hash_block_size: - return block_hashes - return BlockHashListWithBlockSize(block_hashes, self.hash_block_size, kv_cache_spec.block_size) - num_groups = len(self.kv_cache_config.kv_cache_groups) hit_length = max_cache_hit_length longest_hit_length = 0 hit_blocks_by_group: list[list[KVCacheBlock] | None] = [None] * num_groups + hit_length_by_group: list[int] = [0] * num_groups # Simple hybrid (1 full attn + 1 other): one iteration suffices. # Full attn is always first if it exists. @@ -642,30 +650,38 @@ def _get_block_hashes(kv_cache_spec: KVCacheSpec) -> BlockHashList: curr_hit_length = hit_length for idx, (spec, group_ids, manager_cls, use_eagle) in enumerate(self.attention_groups): - cached_blocks = hit_blocks_by_group[group_ids[0]] + first_group_id = group_ids[0] + group_block_size = self.single_type_managers[first_group_id].block_size + cached_blocks = hit_blocks_by_group[first_group_id] if isinstance(spec, FullAttentionSpec) and cached_blocks is not None: # Full attention is downward-closed: we only need to look # up cached blocks once; on subsequent iterations just trim # to the (reduced) current hit length. - curr_hit_length = curr_hit_length // spec.block_size * spec.block_size + curr_hit_length = min(curr_hit_length, hit_length_by_group[first_group_id]) continue drop_eagle_block = use_eagle and idx not in eagle_verified _max_length = curr_hit_length - if drop_eagle_block: - # Eagle needs to match one more block and then pop the last. - _max_length = min(curr_hit_length + spec.block_size, max_cache_hit_length) - hit_blocks = manager_cls.find_longest_cache_hit( - block_hashes=_get_block_hashes(spec), + if drop_eagle_block and not isinstance(spec, MambaSpec): + eagle_margin = ( + self.hash_block_size + if self.enable_partial_hash_hits + and manager_cls.supports_fine_grained_hash_lookup + and group_block_size > self.hash_block_size + else group_block_size + ) + _max_length = min(curr_hit_length + eagle_margin, max_cache_hit_length) + hit_blocks, _new_hit_length = manager_cls.find_longest_cache_hit( + block_hashes=block_hashes, max_length=_max_length, kv_cache_group_ids=group_ids, block_pool=self.block_pool, kv_cache_spec=spec, drop_eagle_block=drop_eagle_block, - alignment_tokens=self.scheduler_block_size, + alignment_tokens=self._cache_hit_alignment_tokens, + dcp_world_size=(self.dcp_world_size if isinstance(spec, FullAttentionSpec) else 1), ) - _new_hit_length = len(hit_blocks[0]) * spec.block_size if drop_eagle_block: eagle_verified.add(idx) elif _new_hit_length < curr_hit_length: @@ -674,6 +690,7 @@ def _get_block_hashes(kv_cache_spec: KVCacheSpec) -> BlockHashList: curr_hit_length = _new_hit_length for group_id, blocks in zip(group_ids, hit_blocks): hit_blocks_by_group[group_id] = blocks + hit_length_by_group[group_id] = _new_hit_length longest_hit_length = max(longest_hit_length, curr_hit_length) @@ -686,10 +703,12 @@ def _get_block_hashes(kv_cache_spec: KVCacheSpec) -> BlockHashList: # Truncate full attention blocks to final hit_length (if present) first_group = self.attention_groups[0] if isinstance(first_group.spec, FullAttentionSpec): - num_blocks = hit_length // first_group.spec.block_size + group_block_size = self.single_type_managers[first_group.group_ids[0]].block_size + num_blocks = cdiv(hit_length, group_block_size) for group_id in first_group.group_ids: if (blks := hit_blocks_by_group[group_id]) is not None: del blks[num_blocks:] + hit_length_by_group[group_id] = hit_length # Uncached shared prefix detection: If any attn. group cached a longer prefix # than the current prefix, it is an uncached common prefix across requests: @@ -707,26 +726,20 @@ def find_longest_cache_hit_per_group( (blocks_per_group, hit_lengths_per_group) """ - def _get_block_hashes(kv_cache_spec: KVCacheSpec) -> BlockHashList: - if kv_cache_spec.block_size == self.hash_block_size: - return block_hashes - return BlockHashListWithBlockSize(block_hashes, self.hash_block_size, kv_cache_spec.block_size) - num_groups = len(self.kv_cache_config.kv_cache_groups) hit_blocks: list[list[KVCacheBlock]] = [[] for _ in range(num_groups)] hit_lengths: list[int] = [0] * num_groups for spec, group_ids, manager_cls, use_eagle in self.attention_groups: - blocks = manager_cls.find_longest_cache_hit( - block_hashes=_get_block_hashes(spec), + blocks, group_hit = manager_cls.find_longest_cache_hit( + block_hashes=block_hashes, max_length=max_cache_hit_length, kv_cache_group_ids=group_ids, block_pool=self.block_pool, kv_cache_spec=spec, drop_eagle_block=use_eagle, - alignment_tokens=self.scheduler_block_size, + alignment_tokens=self._cache_hit_alignment_tokens, ) - group_hit = len(blocks[0]) * spec.block_size for gid, blks in zip(group_ids, blocks): hit_blocks[gid] = blks hit_lengths[gid] = group_hit diff --git a/aphrodite/v1/core/kv_cache_manager.py b/aphrodite/v1/core/kv_cache_manager.py index 2e2500de90..1720e56b17 100644 --- a/aphrodite/v1/core/kv_cache_manager.py +++ b/aphrodite/v1/core/kv_cache_manager.py @@ -11,7 +11,7 @@ from aphrodite.utils.math_utils import cdiv from aphrodite.v1.core.kv_cache_coordinator import get_kv_cache_coordinator from aphrodite.v1.core.kv_cache_metrics import KVCacheMetricsCollector -from aphrodite.v1.core.kv_cache_utils import KVCacheBlock +from aphrodite.v1.core.kv_cache_utils import KVCacheBlock, KVCacheBlockCopy from aphrodite.v1.kv_cache_interface import ( AttentionSpec, CrossAttentionSpec, @@ -127,6 +127,7 @@ def __init__( max_in_flight_tokens = max_model_len self.enable_caching = enable_caching + self.enable_kv_cache_events = enable_kv_cache_events self.use_eagle = use_eagle self.log_stats = log_stats self.metrics_collector = metrics_collector @@ -222,6 +223,23 @@ def get_computed_blocks(self, request: Request) -> tuple[KVCacheBlocks, int]: request.block_hashes, max_cache_hit_length ) + if ( + num_new_computed_tokens > 0 + and self.enable_kv_cache_events + and getattr(request, "kv_cache_report_mode", "incremental") == "full" + ): + for group_idx, group_blocks in enumerate(computed_blocks): + if len(group_blocks) == 0: + continue + group = self.kv_cache_config.kv_cache_groups[group_idx] + block_size = group.kv_cache_spec.block_size + self.block_pool.emit_cached_block_events( + request, + len(group_blocks), + block_size, + group_idx, + ) + if self.log_stats: assert self.prefix_cache_stats is not None self.prefix_cache_stats.record( @@ -365,6 +383,7 @@ def allocate_slots( new_computed_blocks=new_computed_block_list, num_encoder_tokens=num_encoder_tokens, total_computed_tokens=total_computed_tokens, + num_local_computed_tokens=num_local_computed_tokens, num_tokens_main_model=full_num_tokens, apply_admission_cap=True, ) @@ -396,6 +415,7 @@ def allocate_slots( new_computed_blocks=new_computed_block_list, num_encoder_tokens=num_encoder_tokens, total_computed_tokens=num_local_computed_tokens + num_external_computed_tokens, + num_local_computed_tokens=num_local_computed_tokens, num_tokens_main_model=num_tokens_main_model, ) @@ -610,6 +630,23 @@ def take_new_block_ids(self) -> list[int]: ids.extend(mgr.take_new_block_ids()) return ids + def take_kv_cache_block_copies( + self, + ) -> tuple[list[KVCacheBlockCopy], list[KVCacheBlock]]: + """Drain pending copies and return their retained endpoints.""" + pending_copies: list[tuple[KVCacheBlock, KVCacheBlock]] = [] + for mgr in self.coordinator.single_type_managers: + pending_copies.extend(mgr.take_pending_cow_copies()) + copies = [ + KVCacheBlockCopy( + src_block_id=source_block.block_id, + dst_block_id=cow_block.block_id, + ) + for source_block, cow_block in pending_copies + ] + retained_blocks = [block for pair in pending_copies for block in pair] + return copies, retained_blocks + def new_step_starts(self) -> None: - """Called when a new step is started.""" + """Notify the coordinator that a new step is starting.""" self.coordinator.new_step_starts() diff --git a/aphrodite/v1/core/kv_cache_utils.py b/aphrodite/v1/core/kv_cache_utils.py index 30461a6e22..a5c000fe4a 100644 --- a/aphrodite/v1/core/kv_cache_utils.py +++ b/aphrodite/v1/core/kv_cache_utils.py @@ -10,7 +10,7 @@ from collections.abc import Callable, Iterable, Iterator, Sequence from dataclasses import dataclass, replace from functools import partial -from typing import Any, NewType, TypeAlias, cast, overload +from typing import Any, NamedTuple, NewType, TypeAlias, cast, overload from aphrodite import envs from aphrodite.config import AphroditeConfig @@ -174,6 +174,11 @@ def __repr__(self) -> str: ) +class KVCacheBlockCopy(NamedTuple): + src_block_id: int + dst_block_id: int + + class FreeKVCacheBlockQueue: """This class organizes a list of KVCacheBlock objects to a doubly linked list of free blocks. We implement this class instead of using Python @@ -600,7 +605,7 @@ def resolve_kv_cache_block_sizes( group's block size — context parallelism is not supported here. - ``hash_block_size`` is the granularity at which ``Request.block_hashes`` is computed. Single group: equals scheduler block size. Multiple groups: - ``cache_config.hash_block_size`` override if set, else the GCD of group + ``cache_config.prefix_match_unit`` override if set, else the GCD of group block sizes; every group's block size must be divisible by it. Returns the scheduler block size (i.e. disables finer hashing) if block hashing is inactive or a mamba group's block size diverges from the cache @@ -639,12 +644,12 @@ def resolve_kv_cache_block_sizes( ): return scheduler_block_size, scheduler_block_size - requested = cache_config.hash_block_size + requested = cache_config.prefix_match_unit hash_block_size = requested if requested is not None else math.gcd(*group_block_sizes) if any(bs % hash_block_size != 0 for bs in group_block_sizes): raise ValueError( - f"Invalid hash_block_size={hash_block_size}; all KV cache group " - f"block sizes must be divisible by hash_block_size. " + f"Invalid prefix_match_unit={hash_block_size}; all KV cache group " + f"block sizes must be divisible by prefix_match_unit. " f"Got group block sizes={group_block_sizes}." ) return scheduler_block_size, hash_block_size @@ -2075,3 +2080,33 @@ def _get_value_at(self, idx: int) -> BlockHash: BlockHashList = list[BlockHash] | BlockHashListWithBlockSize + + +def resolve_block_hashes( + block_hashes: BlockHashList, + hash_block_size: int, + block_size: int, + *, + supports_fine_grained_hash_lookup: bool = False, + alignment_tokens: int | None = None, +) -> BlockHashList: + """Resolve the block-hash view at ``block_size``. + + When ``block_size`` equals ``hash_block_size``, reuse the precomputed block + hashes directly; otherwise view them at ``block_size`` granularity. + Fine-grained lookup keeps the original hashes for partial cache hits. + """ + if block_size == hash_block_size: + return block_hashes + if isinstance(block_hashes, BlockHashListWithBlockSize): + assert block_hashes.scale_factor == block_size // hash_block_size + return block_hashes + if ( + supports_fine_grained_hash_lookup + and alignment_tokens is not None + and alignment_tokens < block_size + and block_size % alignment_tokens == 0 + ): + return block_hashes + assert block_size % hash_block_size == 0 + return BlockHashListWithBlockSize(block_hashes, hash_block_size, block_size) diff --git a/aphrodite/v1/core/sched/output.py b/aphrodite/v1/core/sched/output.py index fc1be74f7d..e55f475581 100644 --- a/aphrodite/v1/core/sched/output.py +++ b/aphrodite/v1/core/sched/output.py @@ -16,10 +16,12 @@ from aphrodite.multimodal.inputs import MultiModalFeatureSpec from aphrodite.pooling_params import PoolingParams from aphrodite.sampling_params import SamplingParams + from aphrodite.v1.core.kv_cache_utils import KVCacheBlockCopy from aphrodite.v1.request import Request else: ECConnectorMetadata = object KVConnectorMetadata = object + KVCacheBlockCopy = object LoRARequest = object MultiModalFeatureSpec = object PoolingParams = object @@ -230,6 +232,9 @@ class SchedulerOutput: # preventing stale NaN/data from corrupting attention or SSM computation. new_block_ids_to_zero: list[int] | None = None + # CoW copies to apply after zeroing new blocks and before forward. + kv_cache_block_copies: list[KVCacheBlockCopy] | None = None + # Dynamic speculative decoding: optimal K chosen by scheduler. # Number of spec tokens to schedule for the next step. num_spec_tokens_to_schedule: int = 0 diff --git a/aphrodite/v1/core/sched/scheduler.py b/aphrodite/v1/core/sched/scheduler.py index 39bd639c93..d0e0ec54b4 100644 --- a/aphrodite/v1/core/sched/scheduler.py +++ b/aphrodite/v1/core/sched/scheduler.py @@ -245,6 +245,7 @@ def __init__( # Create the KV cache manager. if hash_block_size is None: hash_block_size = block_size + self.hash_block_size = hash_block_size self.kv_cache_manager = KVCacheManager( kv_cache_config=kv_cache_config, max_model_len=self.max_model_len, @@ -279,6 +280,7 @@ def __init__( self.has_mamba_layers = kv_cache_config.has_mamba_layers self.needs_kv_cache_zeroing = kv_cache_config.needs_kv_cache_zeroing self.need_mamba_block_aligned_split = self.has_mamba_layers and self.cache_config.mamba_cache_mode == "align" + self.mamba_partial_cache_hit = self.need_mamba_block_aligned_split and self.hash_block_size < self.block_size # Counts of non-empty steps scheduled / processed. update_from_output # is called once per scheduled step in FIFO order, so these stay in sync. @@ -344,15 +346,27 @@ def _mamba_block_aligned_split( if self.requires_eagle_cache_drop: last_cache_position = max(last_cache_position - block_size, 0) num_computed_tokens_after_sched = num_computed_tokens + num_new_tokens - if num_computed_tokens_after_sched < last_cache_position: - # align to block_size - num_new_tokens = num_new_tokens // block_size * block_size + next_boundary = (num_computed_tokens // block_size + 1) * block_size + if ( + num_computed_tokens % block_size != 0 + and next_boundary <= last_cache_position + and num_computed_tokens_after_sched > next_boundary + ): + num_new_tokens = next_boundary - num_computed_tokens + elif num_computed_tokens_after_sched < last_cache_position: + aligned_end = num_computed_tokens_after_sched // block_size * block_size + num_new_tokens = max(aligned_end - num_computed_tokens, 0) elif num_computed_tokens < last_cache_position < num_computed_tokens_after_sched: # force to cache the last chunk num_new_tokens = last_cache_position - num_computed_tokens - else: - # prefill the last few tokens - pass + elif self.mamba_partial_cache_hit: + tail_boundary = request.num_prompt_tokens // self.hash_block_size * self.hash_block_size + if ( + num_computed_tokens < tail_boundary < num_computed_tokens_after_sched + and tail_boundary < request.num_prompt_tokens + and tail_boundary > last_cache_position + ): + num_new_tokens = tail_boundary - num_computed_tokens # Marconi cache admission optimization: # cache common prefixes by scheduling num_new_tokens = common prefix length @@ -997,6 +1011,10 @@ def schedule(self, throttle_prefills: bool = False) -> SchedulerOutput: # does not grow unbounded; only kv-cache zeroing consumes them. new_attn_block_ids = self.kv_cache_manager.take_new_block_ids() new_block_ids_to_zero = (new_attn_block_ids or None) if self.needs_kv_cache_zeroing else None + kv_cache_block_copies, cow_retained_blocks = self.kv_cache_manager.take_kv_cache_block_copies() + if kv_cache_block_copies: + self._free_cow_retained_blocks(cow_retained_blocks, self.sched_step_seq + 1) + pending_kv_cache_block_copies = kv_cache_block_copies or None # Dynamic speculative decoding: compute optimal K num_spec_tokens_to_schedule = self.num_spec_tokens @@ -1019,6 +1037,7 @@ def schedule(self, throttle_prefills: bool = False) -> SchedulerOutput: finished_req_ids=self.finished_req_ids, free_encoder_mm_hashes=self.encoder_cache_manager.get_freed_mm_hashes(), new_block_ids_to_zero=new_block_ids_to_zero, + kv_cache_block_copies=pending_kv_cache_block_copies, num_spec_tokens_to_schedule=num_spec_tokens_to_schedule, ) @@ -1970,11 +1989,18 @@ def _free_request_blocks(self, request: Request): if blocks: self.deferred_frees.append((self.sched_step_seq, blocks)) + def _free_cow_retained_blocks(self, blocks: list[KVCacheBlock], fence_seq: int) -> None: + """Release CoW copy retentions once their copy step is processed.""" + if not self.defer_block_free or fence_seq <= self.processed_step_seq: + self.kv_cache_manager.block_pool.free_blocks(blocks) + return + self.deferred_frees.append((fence_seq, blocks[::-1])) + def _drain_deferred_frees(self): """Return deferred blocks whose fence step has completed. - Entries are appended with monotonically non-decreasing fences, so - stop at the first one that is still pending. + Fences are appended in near-monotonic order; stop at the first pending + one, and any satisfied entry behind it is freed later. """ while self.deferred_frees: fence, _ = self.deferred_frees[0] @@ -2200,6 +2226,7 @@ def _request_remaining_blocks(self, request: Request) -> int: new_computed_blocks=self.kv_cache_manager.empty_kv_cache_blocks.blocks, num_encoder_tokens=0, total_computed_tokens=request.num_computed_tokens, + num_local_computed_tokens=request.num_computed_tokens, num_tokens_main_model=full_num_tokens, apply_admission_cap=True, ) diff --git a/aphrodite/v1/core/single_type_kv_cache_manager.py b/aphrodite/v1/core/single_type_kv_cache_manager.py index 0bc5ef3ddb..a9198cead6 100644 --- a/aphrodite/v1/core/single_type_kv_cache_manager.py +++ b/aphrodite/v1/core/single_type_kv_cache_manager.py @@ -4,13 +4,16 @@ from abc import ABC, abstractmethod from collections import defaultdict from collections.abc import Sequence +from typing import ClassVar from aphrodite.utils.math_utils import cdiv from aphrodite.v1.core.block_pool import BlockPool from aphrodite.v1.core.kv_cache_utils import ( BlockHashList, + BlockHashListWithBlockSize, BlockHashWithGroupId, KVCacheBlock, + resolve_block_hashes, ) from aphrodite.v1.kv_cache_interface import ( ChunkedLocalAttentionSpec, @@ -36,6 +39,8 @@ class SingleTypeKVCacheManager(ABC): logic of one specific type of attention layer. """ + supports_fine_grained_hash_lookup: ClassVar[bool] = False + def __init__( self, kv_cache_spec: KVCacheSpec, @@ -106,16 +111,29 @@ def __init__( # determining the attention groups. self.use_eagle = False + # Partial-hit copy-on-write bookkeeping. Populated only by fine-grained + # managers (full attention, mamba "align"); harmlessly empty elsewhere. + self._partial_hit_reqs: dict[str, tuple[int, KVCacheBlock]] = {} + self._pending_cow_copies: list[tuple[KVCacheBlock, KVCacheBlock]] = [] + @classmethod def _get_num_evictable_blocks(cls, blocks: Sequence[KVCacheBlock]): return sum(blk.ref_cnt == 0 and not blk.is_null for blk in blocks) + def _has_partial_local_hit( + self, + new_computed_blocks: Sequence[KVCacheBlock], + num_local_computed_tokens: int, + ) -> bool: + return len(new_computed_blocks) > 0 and num_local_computed_tokens % self.block_size != 0 + def get_num_blocks_to_allocate( self, request_id: str, num_tokens: int, new_computed_blocks: Sequence[KVCacheBlock], total_computed_tokens: int, + num_local_computed_tokens: int, num_tokens_main_model: int, apply_admission_cap: bool = False, ) -> int: @@ -185,6 +203,8 @@ def get_num_blocks_to_allocate( # ref_cnt == 0), it will be removed from the free queue when touched by # the allocated request, so we must count it in the free-capacity check. num_evictable_blocks = self._get_num_evictable_blocks(new_computed_blocks[num_skipped_new_computed_blocks:]) + if self._has_partial_local_hit(new_computed_blocks, num_local_computed_tokens): + num_new_blocks += 1 return num_new_blocks + num_evictable_blocks def add_local_computed_blocks( @@ -234,6 +254,10 @@ def add_local_computed_blocks( # them so cache_blocks() will not try to re-cache blocks that already # have a block_hash set. self.num_cached_block[request_id] = len(req_blocks) + if self._has_partial_local_hit(new_computed_blocks, num_local_computed_tokens): + block_idx = num_local_computed_tokens // self.block_size + self._partial_hit_reqs[request_id] = (block_idx, new_computed_blocks[-1]) + self.num_cached_block[request_id] = block_idx def allocate_external_computed_blocks( self, @@ -287,17 +311,26 @@ def allocate_new_blocks(self, request_id: str, num_tokens: int, num_tokens_main_ Returns: The new allocated blocks. """ + cow_blocks: list[KVCacheBlock] = [] + if request_id in self._partial_hit_reqs: + block_idx, source_block = self._partial_hit_reqs.pop(request_id) + cow_block = self.block_pool.get_new_blocks(1)[0] + self._apply_cow(request_id, block_idx, source_block, cow_block) + if self._record_new_block_ids: + self.new_block_ids.append(cow_block.block_id) + cow_blocks.append(cow_block) + req_blocks = self.req_to_blocks[request_id] num_required_blocks = cdiv(num_tokens, self.block_size) num_new_blocks = num_required_blocks - len(req_blocks) if num_new_blocks <= 0: - return [] + return cow_blocks else: new_blocks = self.block_pool.get_new_blocks(num_new_blocks) req_blocks.extend(new_blocks) if self._record_new_block_ids: self.new_block_ids.extend(b.block_id for b in new_blocks) - return new_blocks + return cow_blocks + new_blocks def take_new_block_ids(self) -> list[int]: """Drain and return block IDs allocated since the last call.""" @@ -305,6 +338,26 @@ def take_new_block_ids(self) -> list[int]: self.new_block_ids = [] return ids + def take_pending_cow_copies(self) -> list[tuple[KVCacheBlock, KVCacheBlock]]: + pending_copies = self._pending_cow_copies + self._pending_cow_copies = [] + return pending_copies + + def _apply_cow( + self, + request_id: str, + block_idx: int, + source_block: KVCacheBlock, + cow_block: KVCacheBlock, + ) -> None: + req_blocks = self.req_to_blocks[request_id] + assert block_idx < len(req_blocks) + assert req_blocks[block_idx] is source_block + assert not source_block.is_null and source_block.ref_cnt > 0 + req_blocks[block_idx] = cow_block + self._pending_cow_copies.append((source_block, cow_block)) + cow_block.ref_cnt += 1 + def cache_blocks( self, request: Request, @@ -386,6 +439,7 @@ def pop_blocks_for_free(self, request_id: str) -> list[KVCacheBlock]: # Default to [] in case a request is freed (aborted) before alloc. req_blocks = self.req_to_blocks.pop(request_id, []) self.num_cached_block.pop(request_id, None) + self._partial_hit_reqs.pop(request_id, None) return req_blocks def free(self, request_id: str) -> None: @@ -427,7 +481,7 @@ def find_longest_cache_hit( alignment_tokens: int, dcp_world_size: int = 1, pcp_world_size: int = 1, - ) -> tuple[list[KVCacheBlock], ...]: + ) -> tuple[tuple[list[KVCacheBlock], ...], int]: """ Get the longest cache hit prefix of the blocks that is not longer than `max_length`. The prefix should be a common prefix hit for all the @@ -551,6 +605,8 @@ def new_step_starts(self) -> None: class FullAttentionManager(SingleTypeKVCacheManager): + supports_fine_grained_hash_lookup: ClassVar[bool] = True + @classmethod def find_longest_cache_hit( cls, @@ -563,35 +619,93 @@ def find_longest_cache_hit( alignment_tokens: int, dcp_world_size: int = 1, pcp_world_size: int = 1, - ) -> tuple[list[KVCacheBlock], ...]: + ) -> tuple[tuple[list[KVCacheBlock], ...], int]: assert isinstance(kv_cache_spec, FullAttentionSpec | ChunkedLocalAttentionSpec), ( "FullAttentionManager can only be used for full attention and chunked local attention groups" ) - computed_blocks: tuple[list[KVCacheBlock], ...] = tuple([] for _ in range(len(kv_cache_group_ids))) block_size = kv_cache_spec.block_size if dcp_world_size * pcp_world_size > 1: block_size *= dcp_world_size * pcp_world_size - max_num_blocks = max_length // block_size - for block_hash in itertools.islice(block_hashes, max_num_blocks): - # block_hashes is a chain of block hashes. If a block hash is not - # in the cached_block_hash_to_id, the following block hashes are - # not computed yet for sure. - if cached_block := block_pool.get_cached_block(block_hash, kv_cache_group_ids): - for computed, cached in zip(computed_blocks, cached_block): + block_hashes = resolve_block_hashes( + block_hashes, + block_pool.hash_block_size, + block_size, + supports_fine_grained_hash_lookup=cls.supports_fine_grained_hash_lookup, + alignment_tokens=alignment_tokens, + ) + + fine_grained = alignment_tokens < block_size and block_size % alignment_tokens == 0 + if fine_grained: + assert isinstance(block_hashes, list) + full_block_hashes: BlockHashList = BlockHashListWithBlockSize(block_hashes, alignment_tokens, block_size) + else: + full_block_hashes = block_hashes + + computed_blocks: tuple[list[KVCacheBlock], ...] = tuple([] for _ in range(len(kv_cache_group_ids))) + for block_hash in itertools.islice(full_block_hashes, max_length // block_size): + cached_block = block_pool.get_cached_block(block_hash, kv_cache_group_ids) + if not cached_block: + break + for computed, cached in zip(computed_blocks, cached_block): + computed.append(cached) + hit_length = len(computed_blocks[0]) * block_size + + if fine_grained: + assert isinstance(block_hashes, list) + scale_factor = block_size // alignment_tokens + first_partial_idx = len(computed_blocks[0]) * scale_factor + max_partial_idx = min( + first_partial_idx + scale_factor - 1, + max_length // alignment_tokens, + len(block_hashes), + ) + for fine_idx in range(max_partial_idx - 1, first_partial_idx - 1, -1): + cached_tail = block_pool.get_cached_block(block_hashes[fine_idx], kv_cache_group_ids) + if not cached_tail: + continue + for computed, cached in zip(computed_blocks, cached_tail): computed.append(cached) - else: + hit_length = (fine_idx + 1) * alignment_tokens break - if drop_eagle_block and computed_blocks[0]: - # Need to drop the last matched block if eagle is enabled. - for computed in computed_blocks: - computed.pop() - while ( - block_size != alignment_tokens # Faster for common case. - and len(computed_blocks[0]) * block_size % alignment_tokens != 0 - ): - for computed in computed_blocks: - computed.pop() - return computed_blocks + + if drop_eagle_block and hit_length > 0: + hit_length -= min(alignment_tokens, block_size) + hit_length -= hit_length % alignment_tokens + num_blocks = cdiv(hit_length, block_size) + for computed in computed_blocks: + del computed[num_blocks:] + return computed_blocks, hit_length + + def cache_blocks( + self, + request: Request, + num_tokens: int, + retention_interval: int | None = None, + ) -> None: + super().cache_blocks(request, num_tokens, retention_interval=retention_interval) + hash_block_size = self.block_pool.hash_block_size + if self.block_size != hash_block_size: + self._cache_partial_tail_block(request, num_tokens) + + def _cache_partial_tail_block(self, request: Request, num_tokens: int) -> None: + hash_block_size = self.block_pool.hash_block_size + boundary_tokens = request.num_prompt_tokens // hash_block_size * hash_block_size + if boundary_tokens == 0 or boundary_tokens > num_tokens: + return + if boundary_tokens % self.block_size == 0: + return + + blocks = self.req_to_blocks[request.request_id] + block_idx = boundary_tokens // self.block_size + if block_idx >= len(blocks): + return + self.block_pool.cache_partial_block( + request=request, + block=blocks[block_idx], + num_tokens=boundary_tokens, + kv_cache_group_id=self.kv_cache_group_id, + block_size=self.block_size, + ) def get_num_common_prefix_blocks(self, running_request_id: str) -> int: blocks = self.req_to_blocks[running_request_id] @@ -674,12 +788,22 @@ def find_longest_cache_hit( alignment_tokens: int, dcp_world_size: int = 1, pcp_world_size: int = 1, - ) -> tuple[list[KVCacheBlock], ...]: + ) -> tuple[tuple[list[KVCacheBlock], ...], int]: assert isinstance(kv_cache_spec, SlidingWindowSpec), ( "SlidingWindowManager can only be used for sliding window groups" ) assert dcp_world_size == 1, "DCP not support sliding window attn now." assert pcp_world_size == 1, "PCP not support sliding window attn now." + assert alignment_tokens % kv_cache_spec.block_size == 0, ( + "SlidingWindowManager does not support fine-grained (partial) cache hits" + ) + block_hashes = resolve_block_hashes( + block_hashes, + block_pool.hash_block_size, + kv_cache_spec.block_size, + supports_fine_grained_hash_lookup=cls.supports_fine_grained_hash_lookup, + alignment_tokens=alignment_tokens, + ) # The number of contiguous blocks needed for a prefix cache hit. sliding_window_contiguous_blocks = cls._contiguous_blocks_for_hit( @@ -692,7 +816,9 @@ def find_longest_cache_hit( # sliding_window_contiguous_blocks), # which is good for low cache hit rate scenarios. max_num_blocks = max_length // kv_cache_spec.block_size - computed_blocks = tuple([block_pool.null_block] * max_num_blocks for _ in range(len(kv_cache_group_ids))) + computed_blocks: tuple[list[KVCacheBlock], ...] = tuple( + [block_pool.null_block] * max_num_blocks for _ in range(len(kv_cache_group_ids)) + ) block_size = kv_cache_spec.block_size num_contiguous_blocks = 0 match_found = False @@ -739,7 +865,8 @@ def find_longest_cache_hit( while block_size != alignment_tokens and len(computed_blocks[0]) * block_size % alignment_tokens != 0: for computed in computed_blocks: computed.pop() - return computed_blocks + hit_length = len(computed_blocks[0]) * block_size + return computed_blocks, hit_length @classmethod def reachable_block_mask( @@ -860,7 +987,7 @@ def find_longest_cache_hit( alignment_tokens: int, dcp_world_size: int = 1, pcp_world_size: int = 1, - ) -> tuple[list[KVCacheBlock], ...]: + ) -> tuple[tuple[list[KVCacheBlock], ...], int]: """ For chunked local attention, we need to find the longest cache hit prefix of the blocks that is not longer than `max_length`. The prefix @@ -905,6 +1032,13 @@ def find_longest_cache_hit( assert kv_cache_spec.block_size == alignment_tokens, ( "KV cache groups with different block sizes are not compatible with chunked local attention now" ) + block_hashes = resolve_block_hashes( + block_hashes, + block_pool.hash_block_size, + kv_cache_spec.block_size, + supports_fine_grained_hash_lookup=cls.supports_fine_grained_hash_lookup, + alignment_tokens=alignment_tokens, + ) max_num_blocks = max_length // kv_cache_spec.block_size if max_length > 0: local_attention_start_idx = ( @@ -927,7 +1061,8 @@ def find_longest_cache_hit( computed.append(cached) else: break - return computed_blocks + hit_length = len(computed_blocks[0]) * kv_cache_spec.block_size + return computed_blocks, hit_length def get_num_skipped_tokens(self, num_computed_tokens: int) -> int: """ @@ -981,11 +1116,15 @@ def get_num_common_prefix_blocks(self, running_request_id: str) -> int: class MambaManager(SingleTypeKVCacheManager): + supports_fine_grained_hash_lookup: ClassVar[bool] = True + def __init__(self, kv_cache_spec: MambaSpec, block_pool: BlockPool, **kwargs) -> None: super().__init__(kv_cache_spec, block_pool, **kwargs) - self.cached_blocks_this_step: set[BlockHashWithGroupId] = set() + # Mamba align mode uses the spec block size for its recurrent state. + self.block_size = kv_cache_spec.block_size self.mamba_cache_mode = kv_cache_spec.mamba_cache_mode self.num_speculative_blocks: int = kv_cache_spec.num_speculative_blocks + self.cached_blocks_this_step: set[BlockHashWithGroupId] = set() if self.mamba_cache_mode == "align": # Mapping from request ID to the index of the block # allocated in the previous step @@ -1005,13 +1144,38 @@ def find_longest_cache_hit( alignment_tokens: int, dcp_world_size: int = 1, pcp_world_size: int = 1, - ) -> tuple[list[KVCacheBlock], ...]: + ) -> tuple[tuple[list[KVCacheBlock], ...], int]: assert isinstance(kv_cache_spec, MambaSpec), "MambaManager can only be used for mamba groups" assert dcp_world_size == 1, "DCP not support mamba now." assert pcp_world_size == 1, "PCP not support mamba now." + block_hashes = resolve_block_hashes( + block_hashes, + block_pool.hash_block_size, + kv_cache_spec.block_size, + supports_fine_grained_hash_lookup=cls.supports_fine_grained_hash_lookup, + alignment_tokens=alignment_tokens, + ) computed_blocks: tuple[list[KVCacheBlock], ...] = tuple([] for _ in range(len(kv_cache_group_ids))) + hit_length = 0 block_size = kv_cache_spec.block_size + if alignment_tokens < block_size and block_size % alignment_tokens == 0: + assert isinstance(block_hashes, list) + hash_block_size = alignment_tokens + scale_factor = block_size // hash_block_size + max_num_partial_units = min(max_length // hash_block_size, len(block_hashes)) + for fine_idx in range(max_num_partial_units - 1, -1, -1): + num_tokens = (fine_idx + 1) * hash_block_size + block_hash = block_hashes[fine_idx] + if cached_block := block_pool.get_cached_block(block_hash, kv_cache_group_ids): + block_idx = fine_idx // scale_factor + for computed, cached in zip(computed_blocks, cached_block): + computed.extend([block_pool.null_block] * block_idx) + computed.append(cached) + hit_length = num_tokens + break + return computed_blocks, hit_length + max_num_blocks = max_length // block_size # Search from right to left and early stop when a match is found. for i in range(max_num_blocks - 1, -1, -1): @@ -1031,9 +1195,10 @@ def find_longest_cache_hit( # so we insert dummy blocks at the beginning: computed.extend([block_pool.null_block] * i) computed.append(cached) + hit_length = (i + 1) * block_size break # we just need the last match - early stopping - return computed_blocks + return computed_blocks, hit_length @classmethod def reachable_block_mask( @@ -1126,6 +1291,7 @@ def get_num_blocks_to_allocate( num_tokens: int, new_computed_blocks: Sequence[KVCacheBlock], total_computed_tokens: int, + num_local_computed_tokens: int, num_tokens_main_model: int, apply_admission_cap: bool = False, ) -> int: @@ -1146,6 +1312,7 @@ def get_num_blocks_to_allocate( num_tokens, new_computed_blocks, total_computed_tokens, + num_local_computed_tokens, num_tokens_main_model, apply_admission_cap=apply_admission_cap, ) @@ -1161,15 +1328,19 @@ def get_num_blocks_to_allocate( # num_tokens can include draft tokens that will later be rejected. num_required_blocks = cdiv(num_tokens, self.block_size) + self.num_speculative_blocks num_new_blocks = num_required_blocks - len(new_computed_blocks) - len(self.req_to_blocks[request_id]) + has_partial_hit = ( + self._has_partial_local_hit(new_computed_blocks, num_local_computed_tokens) + or request_id in self._partial_hit_reqs + ) + if has_partial_hit: + num_new_blocks = max(num_new_blocks, 0) + 1 if num_new_blocks > 0: if request_id in self._allocated_block_reqs: # Old request. Needs at most 1 more blocks as we can reuse the # speculative blocks in previous step. - num_new_blocks = 1 + num_new_blocks = 1 + int(has_partial_hit) else: - # First prefill. Allocate 1 block for running state and the - # speculative blocks. - num_new_blocks = 1 + self.num_speculative_blocks + num_new_blocks = 1 + self.num_speculative_blocks + int(has_partial_hit) num_evictable_computed_blocks = self._get_num_evictable_blocks(new_computed_blocks) return num_new_blocks + num_evictable_computed_blocks @@ -1193,9 +1364,11 @@ def allocate_new_blocks(self, request_id: str, num_tokens: int, num_tokens_main_ # NOTE(tdouble): this is an over-estimate of how many blocks we need because # num_tokens can include draft tokens that will later be rejected. num_required_blocks = cdiv(num_tokens, self.block_size) + self.num_speculative_blocks + partial_hit = self._partial_hit_reqs.get(request_id) + has_partial_hit = partial_hit is not None # `num_required_blocks` might be less than `len(req_blocks)` if blocks are # over-allocated at last round. - if num_required_blocks <= len(req_blocks): + if num_required_blocks <= len(req_blocks) and not has_partial_hit: return [] else: prev_block_len = len(req_blocks) @@ -1224,14 +1397,33 @@ def allocate_new_blocks(self, request_id: str, num_tokens: int, num_tokens_main_ else: break num_new_blocks = num_required_blocks - len(req_blocks) + if has_partial_hit: + num_new_blocks = max(num_new_blocks, 0) + 1 if blocks_allocated: - assert num_new_blocks <= 1 + assert num_new_blocks <= 1 + int(has_partial_hit) else: - assert num_new_blocks <= self.num_speculative_blocks + 1 + assert num_new_blocks <= self.num_speculative_blocks + 1 + int(has_partial_hit) new_blocks = self.block_pool.get_new_blocks(num_new_blocks) + returned_blocks = req_blocks[prev_block_len:] + if partial_hit is not None: + block_idx, source_block = partial_hit + cow_block = new_blocks[0] + new_blocks = new_blocks[1:] + if blocks_allocated: + assert req_blocks[block_idx] is source_block + self.block_pool.move_block_hashes(source_block, cow_block) + self._pending_cow_copies.append((source_block, cow_block)) + source_block.ref_cnt += 1 + if cow_block.block_hash is not None: + self.cached_blocks_this_step.add(cow_block.block_hash) + else: + self._apply_cow(request_id, block_idx, source_block, cow_block) + returned_blocks = [cow_block] + returned_blocks req_blocks.extend(new_blocks) self._allocated_block_reqs.add(request_id) - return req_blocks[prev_block_len:] + self._partial_hit_reqs.pop(request_id, None) + returned_blocks.extend(new_blocks) + return returned_blocks def pop_blocks_for_free(self, request_id: str) -> list[KVCacheBlock]: if self.mamba_cache_mode == "align": @@ -1256,6 +1448,10 @@ def cache_blocks( num_cached_blocks_before = self.num_cached_block.get(request.request_id, 0) super().cache_blocks(request, num_tokens, retention_interval=retention_interval) num_cached_blocks_after = self.num_cached_block.get(request.request_id, 0) + if self.mamba_cache_mode == "align": + partial_hash = self._cache_partial_tail_block(request, num_tokens) + if partial_hash is not None: + self.cached_blocks_this_step.add(partial_hash) if num_cached_blocks_after > num_cached_blocks_before: for block in self.req_to_blocks[request.request_id][num_cached_blocks_before:num_cached_blocks_after]: # Skip null blocks (align-mode skipped states) and blocks that @@ -1269,6 +1465,42 @@ def cache_blocks( def new_step_starts(self) -> None: self.cached_blocks_this_step.clear() + def _cache_partial_tail_block( + self, + request: Request, + num_tokens: int, + ) -> BlockHashWithGroupId | None: + hash_block_size = self.block_pool.hash_block_size + if self.block_size == hash_block_size: + return None + if num_tokens % self.block_size == 0: + return None + if num_tokens % hash_block_size != 0: + return None + latest_prompt_hash_boundary = request.num_prompt_tokens // hash_block_size * hash_block_size + if num_tokens != latest_prompt_hash_boundary: + return None + + block_idx = num_tokens // self.block_size + blocks = self.req_to_blocks[request.request_id] + if block_idx >= len(blocks): + return None + source_block = blocks[block_idx] + if source_block.is_null: + return None + + partial_hash = self.block_pool.cache_partial_block( + request=request, + block=source_block, + num_tokens=num_tokens, + kv_cache_group_id=self.kv_cache_group_id, + block_size=self.block_size, + ) + if partial_hash is not None: + self._partial_hit_reqs[request.request_id] = (block_idx, source_block) + self.num_cached_block[request.request_id] = block_idx + return partial_hash + class CrossAttentionManager(SingleTypeKVCacheManager): """Manager for cross-attention KV cache in encoder-decoder models.""" @@ -1320,7 +1552,7 @@ def find_longest_cache_hit( alignment_tokens: int, dcp_world_size: int = 1, pcp_world_size: int = 1, - ) -> tuple[list[KVCacheBlock], ...]: + ) -> tuple[tuple[list[KVCacheBlock], ...], int]: assert isinstance(kv_cache_spec, CrossAttentionSpec), ( "CrossAttentionManager can only be used for cross-attention groups" ) diff --git a/aphrodite/v1/engine/async_llm.py b/aphrodite/v1/engine/async_llm.py index 98b0329929..637ba6dff8 100644 --- a/aphrodite/v1/engine/async_llm.py +++ b/aphrodite/v1/engine/async_llm.py @@ -1038,6 +1038,10 @@ async def start_weight_update(self) -> None: """Start a new weight update.""" await self.collective_rpc("start_weight_update") + async def start_draft_weight_update(self) -> None: + """Start a new weight update targeting the speculative draft model.""" + await self.collective_rpc("start_draft_weight_update") + async def update_weights(self, request: WeightTransferUpdateRequest) -> None: """ Batched weight update for RL training. diff --git a/aphrodite/v1/kv_offload/tiering/fs/manager.py b/aphrodite/v1/kv_offload/tiering/fs/manager.py index 594cdc04d4..a35129a2fa 100644 --- a/aphrodite/v1/kv_offload/tiering/fs/manager.py +++ b/aphrodite/v1/kv_offload/tiering/fs/manager.py @@ -19,7 +19,7 @@ import json import os from collections.abc import Iterable -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, ClassVar try: from aphrodite.fs_io_C import batch_lookup as batch_lookup_C @@ -30,11 +30,18 @@ from typing_extensions import override +from aphrodite.distributed.kv_events import MEDIUM_FS from aphrodite.logger import init_logger -from aphrodite.v1.kv_offload.base import LookupResult, OffloadKey, ReqContext +from aphrodite.v1.kv_offload.base import ( + LookupResult, + OffloadingEvent, + OffloadKey, + ReqContext, +) from aphrodite.v1.kv_offload.file_mapper import FileMapper from aphrodite.v1.kv_offload.tiering.async_lookup import AsyncLookupManager from aphrodite.v1.kv_offload.tiering.base import ( + JobId, JobMetadata, JobResult, RequestOffloadingContext, @@ -90,6 +97,8 @@ class FileSystemTierManager(SecondaryTierManager): content. """ + medium: ClassVar[str] = MEDIUM_FS + def __init__( self, offloading_spec: "OffloadingSpec", @@ -98,6 +107,7 @@ def __init__( root_dir: str, n_read_threads: int = 16, n_write_threads: int = 16, + enable_kv_events: bool = False, ): """ Args: @@ -108,9 +118,26 @@ def __init__( root_dir: Root directory for block files. n_read_threads: Number of read-priority I/O threads. n_write_threads: Number of write-priority I/O threads. + enable_kv_events: Emit BlockStored KV events for blocks + successfully stored to this tier. Effective only when KV + cache events are enabled globally (kv_events_config). """ super().__init__(offloading_spec, primary_kv_view, tier_type) + self.events: list[OffloadingEvent] | None = None + if enable_kv_events: + if offloading_spec.kv_events_config.enable_kv_cache_events: + self.events = [] + else: + logger.warning( + "enable_kv_events is set on secondary tier '%s' but KV " + "cache events are disabled globally; the tier will not " + "emit events.", + tier_type, + ) + # Keys of in-flight store jobs, tracked only when events are enabled. + self._store_job_keys: dict[JobId, list[OffloadKey]] = {} + # Extract block size from primary view assert primary_kv_view.strides is not None, "primary_kv_view.strides cannot be None" self._block_size: int = primary_kv_view.strides[0] @@ -151,6 +178,8 @@ def lookup(self, key: OffloadKey, req_context: ReqContext) -> LookupResult: @override def submit_store(self, job_metadata: JobMetadata) -> None: + if self.events is not None: + self._store_job_keys[job_metadata.job_id] = list(job_metadata.keys) tasks = ( functools.partial( store_block, @@ -182,7 +211,20 @@ def get_finished_jobs(self) -> Iterable[JobResult]: """ Collect completed jobs from the finished-jobs queue. """ - return (JobResult(job_id=job_id, success=success) for job_id, success in self._pool.get_finished()) + results = [] + for job_id, success in self._pool.get_finished(): + if self.events is not None: + keys = self._store_job_keys.pop(job_id, None) + if success and keys: + self.events.append(OffloadingEvent(keys=keys, medium=self.medium, removed=False)) + results.append(JobResult(job_id=job_id, success=success)) + return results + + @override + def take_events(self) -> Iterable[OffloadingEvent]: + if self.events is not None: + yield from self.events + self.events.clear() @override def drain_jobs(self) -> None: diff --git a/aphrodite/v1/kv_offload/tiering/obj/manager.py b/aphrodite/v1/kv_offload/tiering/obj/manager.py index 4bbc44e49d..3d59cbbf07 100644 --- a/aphrodite/v1/kv_offload/tiering/obj/manager.py +++ b/aphrodite/v1/kv_offload/tiering/obj/manager.py @@ -5,15 +5,22 @@ import ctypes import time from collections.abc import Iterable -from typing import TYPE_CHECKING, NamedTuple +from typing import TYPE_CHECKING, ClassVar, NamedTuple +from aphrodite.distributed.kv_events import MEDIUM_OBJ from aphrodite.distributed.nixl_utils import NixlWrapper as nixl_agent from aphrodite.distributed.nixl_utils import nixl_agent_config from aphrodite.logger import init_logger -from aphrodite.v1.kv_offload.base import LookupResult, OffloadKey, ReqContext +from aphrodite.v1.kv_offload.base import ( + LookupResult, + OffloadingEvent, + OffloadKey, + ReqContext, +) from aphrodite.v1.kv_offload.file_mapper import FileMapper from aphrodite.v1.kv_offload.tiering.async_lookup import AsyncLookupManager from aphrodite.v1.kv_offload.tiering.base import ( + JobId, JobMetadata, JobResult, RequestOffloadingContext, @@ -88,6 +95,8 @@ class ObjectStoreSecondaryTierManager(SecondaryTierManager): primary tier. Object keys are formed as ``{prefix}/{hash_shard}/{hash}.bin``. """ + medium: ClassVar[str] = MEDIUM_OBJ + def __init__( self, offloading_spec: "OffloadingSpec", @@ -96,8 +105,36 @@ def __init__( store_config: dict, prefix: str = "", io_threads: int = 4, + enable_kv_events: bool = False, ): + """ + Args: + offloading_spec: Offloading configuration. + primary_kv_view: Memoryview of the primary tier's CPU KV cache. + tier_type: Tier type identifier, set by SecondaryTierFactory. + store_config: Object store connection parameters (see ObjStoreConfig). + prefix: Key prefix prepended to all object keys. + io_threads: Number of NIXL I/O threads. + enable_kv_events: Emit BlockStored KV events for blocks + successfully stored to this tier. Effective only when KV + cache events are enabled globally (kv_events_config). + """ super().__init__(offloading_spec, primary_kv_view, tier_type) + + self.events: list[OffloadingEvent] | None = None + if enable_kv_events: + if offloading_spec.kv_events_config.enable_kv_cache_events: + self.events = [] + else: + logger.warning( + "enable_kv_events is set on secondary tier '%s' but KV " + "cache events are disabled globally; the tier will not " + "emit events.", + tier_type, + ) + # Keys of in-flight store jobs, tracked only when events are enabled. + self._store_job_keys: dict[JobId, list[OffloadKey]] = {} + agent_config = nixl_agent_config(backends=[]) self._agent = nixl_agent("ObjAgent", agent_config) obj_config = ObjStoreConfig(**store_config) @@ -217,6 +254,8 @@ def lookup(self, key: OffloadKey, req_context: ReqContext) -> LookupResult: return LookupResult.HIT if result else LookupResult.MISS def submit_store(self, job_metadata: JobMetadata) -> None: + if self.events is not None: + self._store_job_keys[job_metadata.job_id] = list(job_metadata.keys) obj_keys = (self._file_mapper.get_file_name(k) for k in job_metadata.keys) self._submit_transfer(job_metadata.job_id, job_metadata.block_ids, obj_keys, NIXL_WRITE) @@ -261,8 +300,18 @@ def get_finished_jobs(self) -> Iterable[JobResult]: self._poll_active_transfers() results = self._pending_results self._pending_results = [] + if self.events is not None: + for result in results: + keys = self._store_job_keys.pop(result.job_id, None) + if result.success and keys: + self.events.append(OffloadingEvent(keys=keys, medium=self.medium, removed=False)) return results + def take_events(self) -> Iterable[OffloadingEvent]: + if self.events is not None: + yield from self.events + self.events.clear() + def drain_jobs(self) -> None: """Block until every submitted transfer has completed or failed. diff --git a/aphrodite/v1/request.py b/aphrodite/v1/request.py index 6a10a64b0b..6754a054f3 100644 --- a/aphrodite/v1/request.py +++ b/aphrodite/v1/request.py @@ -96,6 +96,7 @@ def __init__( # P/D: Connector-specific KV transfer parameters. self.kv_transfer_params: dict[str, Any] | None = None + self.kv_cache_report_mode = "incremental" if pooling_params is not None: # Pooling models. @@ -109,6 +110,7 @@ def __init__( if sampling_params.extra_args is not None: self.kv_transfer_params = sampling_params.extra_args.get("kv_transfer_params") + self.kv_cache_report_mode = sampling_params.extra_args.get("kv_cache_report_mode", "incremental") else: raise ValueError("sampling_params and pooling_params can't both be unset") diff --git a/aphrodite/v1/sample/thinking_budget_state.py b/aphrodite/v1/sample/thinking_budget_state.py index eb3bd8c0d6..7751647b21 100644 --- a/aphrodite/v1/sample/thinking_budget_state.py +++ b/aphrodite/v1/sample/thinking_budget_state.py @@ -7,7 +7,7 @@ import torch from aphrodite.platforms import current_platform -from aphrodite.utils.torch_utils import PIN_MEMORY, async_tensor_h2d +from aphrodite.utils.torch_utils import async_tensor_h2d from aphrodite.v1.sample.logits_processor.interface import ( BatchUpdate, MoveDirectionality, @@ -22,10 +22,11 @@ def maybe_create_thinking_budget_state_holder( max_num_seqs: int, num_spec_tokens: int, device: torch.device, + is_pin_memory: bool, ) -> "ThinkingBudgetStateHolder | None": if reasoning_config is None: return None - return ThinkingBudgetStateHolder(reasoning_config, max_num_seqs, num_spec_tokens, device, PIN_MEMORY) + return ThinkingBudgetStateHolder(reasoning_config, max_num_seqs, num_spec_tokens, device, is_pin_memory) class ThinkingBudgetStateHolder: @@ -169,17 +170,6 @@ def _find_last_sequence_index(target_list: list[int], token_ids: list[int]) -> i return i return -1 - @staticmethod - def _find_last_sequence_index_from(target_list: list[int], token_ids: list[int], search_start: int) -> int: - """Last occurrence of ``token_ids`` at or after ``search_start``.""" - if not token_ids: - return -1 - lo = max(0, search_start) - for i in range(len(target_list) - len(token_ids), lo - 1, -1): - if target_list[i : i + len(token_ids)] == token_ids: - return i - return -1 - def _init_state_entry(self, prompt_tok_ids: list[int] | None, thinking_token_budget: int) -> dict[str, Any]: if prompt_tok_ids is None: last_start = -1 @@ -225,8 +215,6 @@ def _init_state_entry(self, prompt_tok_ids: list[int] | None, thinking_token_bud "force_index": [], "start_thinking": start_thinking, "end_thinking": -1, - "start_search_pos": 0, - "end_search_pos": 0, "in_spec_mode": False, "bonus_token_forced": False, "continue_thinking": continue_thinking, @@ -242,41 +230,36 @@ def _update_think_state(self, state: dict[str, Any]) -> None: state["force_index"] = [] return - output_tok_ids = state.get("output_tok_ids", []) if state["start_thinking"] == -1: - seq_len = len(self.think_start_token_ids) scan_offset = state.get("scan_offset", 0) - start_thinking = self._find_last_sequence_index_from( - output_tok_ids, - self.think_start_token_ids, - max(scan_offset, state["start_search_pos"] - (seq_len - 1)), - ) - if start_thinking >= 0 and scan_offset > 0: - # Re-entry after a forced end: budget was already exhausted - # in a prior block, so immediately force-close this one. - # scan_offset > 0 is only set after forced-end completion - # (never after natural end), so this won't block legitimate - # re-entries where budget remains. - state["start_thinking"] = start_thinking - state["in_think"] = False - state["in_end"] = True - state["end_count"] = 0 - state["force_index"] = [0] - return + output_slice = state.get("output_tok_ids", [])[scan_offset:] + start_thinking = self._find_last_sequence_index(output_slice, self.think_start_token_ids) + if start_thinking >= 0: + start_thinking += scan_offset state["start_thinking"] = start_thinking - if start_thinking == -1: - state["start_search_pos"] = len(output_tok_ids) if state["end_thinking"] == -1: - seq_len = len(self.think_end_token_ids) scan_offset = state.get("scan_offset", 0) - end_thinking = self._find_last_sequence_index_from( - output_tok_ids, - self.think_end_token_ids, - max(scan_offset, state["end_search_pos"] - (seq_len - 1)), - ) + output_slice = state.get("output_tok_ids", [])[scan_offset:] + end_thinking = self._find_last_sequence_index(output_slice, self.think_end_token_ids) + if end_thinking >= 0: + end_thinking += scan_offset state["end_thinking"] = end_thinking - if end_thinking == -1: - state["end_search_pos"] = len(output_tok_ids) + + if ( + not state.get("in_end", False) + and state["start_thinking"] >= 0 + and state["end_thinking"] >= 0 + and state["end_thinking"] > state["start_thinking"] + and not state.get("continue_thinking", False) + ): + state["in_think"] = False + state["think_count"] = 0 + state["continue_thinking"] = False + state["start_thinking"] = -1 + state["end_thinking"] = -1 + state["scan_offset"] = len(state.get("output_tok_ids", [])) + state["check_count_down"] = state["thinking_token_budget"] + return if state["start_thinking"] == -1: return @@ -296,7 +279,17 @@ def _update_think_state(self, state: dict[str, Any]) -> None: predicted_countdown = current_step_countdown - len(state["spec_token_ids"]) - 1 # We only proceed further if we have counted down the thinking budget # to 0 or less and when we are in the "in think" mode. - if not state.get("in_end", False) and predicted_countdown >= 0 and state["start_thinking"] > -1: + # Exception: when continue_thinking=True and a natural is + # detected (end_thinking != -1), fall through to handle the exit - + # even if the budget hasn't expired yet. For continue_thinking=False, + # the early natural-end detection block above already handles it. + natural_end_with_continue = state.get("continue_thinking", False) and state["end_thinking"] != -1 + if ( + not state.get("in_end", False) + and predicted_countdown >= 0 + and state["start_thinking"] > -1 + and not natural_end_with_continue + ): state["check_count_down"] = current_step_countdown state["prev_output_length"] = len(state.get("output_tok_ids", [])) return @@ -363,6 +356,10 @@ def _update_think_state(self, state: dict[str, Any]) -> None: # Case: ......... - exiting think mode state["in_think"] = False state["think_count"] = 0 + state["continue_thinking"] = False + state["start_thinking"] = -1 + state["end_thinking"] = -1 + state["scan_offset"] = len(state.get("output_tok_ids", [])) elif absolute_start_pos >= 0 and not state["continue_thinking"]: # Found think start - entering think mode @@ -374,6 +371,10 @@ def _update_think_state(self, state: dict[str, Any]) -> None: # Found think end - exiting think mode state["in_think"] = False state["think_count"] = 0 + state["continue_thinking"] = False + state["start_thinking"] = -1 + state["end_thinking"] = -1 + state["scan_offset"] = len(state.get("output_tok_ids", [])) elif state["in_think"]: # Continue thinking mode, increment count by new tokens diff --git a/aphrodite/v1/worker/cp_utils.py b/aphrodite/v1/worker/cp_utils.py index fec81dcb88..c30c8fe96a 100644 --- a/aphrodite/v1/worker/cp_utils.py +++ b/aphrodite/v1/worker/cp_utils.py @@ -2,6 +2,8 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from typing import TYPE_CHECKING, Any, cast +import torch + from aphrodite.config import AphroditeConfig, get_layers_from_aphrodite_config from aphrodite.distributed import get_dcp_group, get_pcp_group @@ -55,3 +57,7 @@ def get_total_cp_world_size(): # DCP might not be initialized in testing dcp_world_size = 1 return dcp_world_size * pcp_world_size + + +def should_skip_dcp_context_attention(context_kv_lens: torch.Tensor) -> bool: + return bool(context_kv_lens.max().item() == 0) diff --git a/aphrodite/v1/worker/gpu/attn_utils.py b/aphrodite/v1/worker/gpu/attn_utils.py index cce7cd94b0..0caf50b538 100644 --- a/aphrodite/v1/worker/gpu/attn_utils.py +++ b/aphrodite/v1/worker/gpu/attn_utils.py @@ -204,7 +204,9 @@ def _reshape_attention_kv_cache( assert inv_order[0] == 0 page_bytes = prod(kv_cache_shape[1:]) * get_dtype_size(dtype) kv_cache = ( - kv_raw_tensor.view(-1, block_stride)[:, offset : offset + page_bytes].view(dtype).view(kv_cache_shape) + kv_raw_tensor.view(-1, block_stride)[:, offset : offset + page_bytes] + .view(dtype) + .view(permuted_kv_cache_shape) ) elif kv_cache_spec.page_size_padded is not None: # Use a strided view to skip the padding between physical pages. @@ -224,7 +226,7 @@ def _reshape_attention_kv_cache( page_stride = kv_cache_spec.page_size_bytes // dtype_size num_blocks_dim = inv_order[0] - strides = list(torch.empty(permuted_kv_cache_shape).stride()) + strides = list(torch.empty(permuted_kv_cache_shape, device="meta").stride()) strides[num_blocks_dim] = page_stride kv_cache = torch.as_strided( @@ -496,7 +498,12 @@ def init_kv_cache( shared_kv_cache_layers=shared_kv_cache_layers, kv_cache_config=kv_cache_config, ) - bind_kv_cache(kv_caches, forward_context, runner_kv_caches) + # Dual-attention models (e.g. LongCat-Flash) put two Attention modules per + # decoder layer, so a layer name carries two integers (layer + module index). + num_attn_module = ( + 2 if aphrodite_config.model_config.hf_config.model_type in ("longcat_flash", "longcat_flash_ngram") else 1 + ) + bind_kv_cache(kv_caches, forward_context, runner_kv_caches, num_attn_module) return kv_caches diff --git a/aphrodite/v1/worker/gpu/model_runner.py b/aphrodite/v1/worker/gpu/model_runner.py index 957c6f1ad3..d57074bfcb 100644 --- a/aphrodite/v1/worker/gpu/model_runner.py +++ b/aphrodite/v1/worker/gpu/model_runner.py @@ -112,7 +112,7 @@ from aphrodite.v1.worker.gpu.states import RequestState from aphrodite.v1.worker.gpu.structured_outputs import StructuredOutputsWorker from aphrodite.v1.worker.lora_model_runner_mixin import LoRAModelRunnerMixin -from aphrodite.v1.worker.utils import KVBlockZeroer +from aphrodite.v1.worker.utils import KVBlockZeroer, copy_kv_cache_blocks_inplace logger = init_logger(__name__) @@ -355,6 +355,12 @@ def load_model(self, load_dummy_weights: bool = False, *args, **kwargs) -> None: def get_model(self) -> nn.Module: return self.model + def get_draft_model(self) -> nn.Module | None: + speculator = self.speculator + if not isinstance(speculator, DraftModelSpeculator): + return None + return speculator.model + def reload_weights(self, *args, **kwargs) -> None: # TODO(Wentao): Use full version instead of import when fully migrated to v2 from aphrodite.v1.worker.gpu_model_runner import GPUModelRunner as GPUModelRunnerV1 @@ -795,6 +801,12 @@ def update_requests(self, scheduler_output: SchedulerOutput) -> None: if scheduler_output.new_block_ids_to_zero: assert self.kv_block_zeroer is not None self.kv_block_zeroer.zero_block_ids(scheduler_output.new_block_ids_to_zero) + if scheduler_output.kv_cache_block_copies: + copy_kv_cache_blocks_inplace( + self.kv_caches, + self.kv_cache_config.num_blocks, + scheduler_output.kv_cache_block_copies, + ) def prepare_inputs(self, scheduler_output: SchedulerOutput, batch_desc: BatchExecutionDescriptor) -> InputBatch: num_tokens = scheduler_output.total_num_scheduled_tokens diff --git a/aphrodite/v1/worker/gpu_input_batch.py b/aphrodite/v1/worker/gpu_input_batch.py index c51e89b961..acef2b7213 100644 --- a/aphrodite/v1/worker/gpu_input_batch.py +++ b/aphrodite/v1/worker/gpu_input_batch.py @@ -116,6 +116,7 @@ def __init__( max_num_reqs, num_spec_tokens, device, + PIN_MEMORY, ) self.thinking_token_budget_reqs: set[str] = set() self.is_pooling_model = is_pooling_model diff --git a/aphrodite/v1/worker/gpu_model_runner.py b/aphrodite/v1/worker/gpu_model_runner.py index e29d1b9254..b8a76c7781 100644 --- a/aphrodite/v1/worker/gpu_model_runner.py +++ b/aphrodite/v1/worker/gpu_model_runner.py @@ -225,6 +225,7 @@ KVBlockZeroer, add_kv_sharing_layers_to_kv_cache_groups, bind_kv_cache, + copy_kv_cache_blocks_inplace, prepare_kernel_block_sizes, sanity_check_mm_encoder_outputs, ) @@ -1084,6 +1085,12 @@ def _update_states(self, scheduler_output: "SchedulerOutput") -> Callable | None # stale NaN/data from corrupting attention or SSM computation. if scheduler_output.new_block_ids_to_zero: self._zero_block_ids(scheduler_output.new_block_ids_to_zero) + if scheduler_output.kv_cache_block_copies: + copy_kv_cache_blocks_inplace( + self.kv_caches, + self.kv_cache_config.num_blocks, + scheduler_output.kv_cache_block_copies, + ) # Free the cached encoder outputs. for mm_hash in scheduler_output.free_encoder_mm_hashes: @@ -3007,6 +3014,15 @@ def get_model(self) -> nn.Module: return self.model.unwrap() return self.model + def get_draft_model(self) -> nn.Module | None: + drafter = getattr(self, "drafter", None) + if drafter is None: + return None + model = getattr(drafter, "model", None) + if isinstance(model, (CUDAGraphWrapper, UBatchWrapper, BreakableCUDAGraphWrapper)): + return cast(nn.Module, model.unwrap()) + return cast(nn.Module | None, model) + def get_supported_generation_tasks(self) -> list[GenerationTask]: model = self.get_model() supported_tasks = list[GenerationTask]() diff --git a/aphrodite/v1/worker/gpu_worker.py b/aphrodite/v1/worker/gpu_worker.py index ec0033aca5..74fd2eb58a 100644 --- a/aphrodite/v1/worker/gpu_worker.py +++ b/aphrodite/v1/worker/gpu_worker.py @@ -154,6 +154,7 @@ def __init__( # Buffers saved before sleep self._sleep_saved_buffers: dict[str, torch.Tensor] = {} + self._sleep_rebuild_draft_metadata_buffers = False # Weight transfer engine is created in `load_model` once the model # is available, since the engine needs a reference to the model. @@ -194,6 +195,9 @@ def sleep(self, level: int = 1) -> None: if level == 2: model = self.model_runner.model self._sleep_saved_buffers = {name: buffer.cpu().clone() for name, buffer in model.named_buffers()} + draft = self.get_draft_model() + inner = getattr(draft, "model", None) if draft is not None else None + self._sleep_rebuild_draft_metadata_buffers = inner is not None and hasattr(inner, "_build_fused_kv_buffers") self._get_sleep_mode_backend().suspend(level) @@ -225,6 +229,14 @@ def wake_up(self, tags: list[str] | None = None) -> None: buffer.data.copy_(self._sleep_saved_buffers[name].data) self._sleep_saved_buffers = {} + if self._sleep_rebuild_draft_metadata_buffers: + draft = self.get_draft_model() + if draft is not None: + inner = getattr(draft, "model", None) + if inner is not None and hasattr(inner, "_build_fused_kv_buffers"): + inner._build_fused_kv_buffers() + self._sleep_rebuild_draft_metadata_buffers = False + if tags is None or "kv_cache" in tags: self.model_runner.post_kv_cache_wake_up() @@ -454,11 +466,14 @@ def determine_available_memory(self) -> int: profile_torch_peak = torch.accelerator.memory_stats(self.device).get("allocated_bytes.all.peak", 0) # Profile CUDA graph memory if graphs will be captured. - # Skip on ROCm/HIP/XPU as graph pool handles and get_memory_info - # behave differently and can produce incorrect/negative estimates. + # ROCm is included: #44825 moved the profiler to + # torch.accelerator.get_memory_info (reliable on ROCm, as used by + # the AMD-CI mem tests), and graph_pool_handle resolves to the same + # torch.cuda handle the live capture path already uses on ROCm. + # XPU stays excluded (see #39977). cudagraph_memory_estimate = 0 if ( - current_platform.is_cuda() + current_platform.is_cuda_alike() and self.aphrodite_config.compilation_config.cudagraph_mode != CUDAGraphMode.NONE ): cudagraph_memory_estimate = self.model_runner.profile_cudagraph_memory() @@ -469,8 +484,7 @@ def determine_available_memory(self) -> int: profile_result.non_torch_increase + profile_result.torch_peak_increase + profile_result.weights_memory ) - # On ROCm, cudagraph_memory_estimate is always 0 so this is a no-op. - # On CUDA, respect the opt-in flag as originally designed. + # Respect the opt-in flag as originally designed. cudagraph_memory_estimate_applied = ( cudagraph_memory_estimate if envs.APHRODITE_MEMORY_PROFILER_ESTIMATE_CUDAGRAPHS else 0 ) @@ -686,11 +700,6 @@ def initialize_from_config(self, kv_cache_config: KVCacheConfig) -> None: @instrument(span_name="Warmup (GPU)") def compile_or_warm_up_model(self) -> CompilationTimes: warmup_sizes: list[int] = [] - cg_capture_sizes: list[int] = [] - - if self.aphrodite_config.compilation_config.cudagraph_mode != CUDAGraphMode.NONE: - cg_sizes = self.aphrodite_config.compilation_config.cudagraph_capture_sizes - cg_capture_sizes = [] if cg_sizes is None else cg_sizes if self.aphrodite_config.compilation_config.mode == CompilationMode.APHRODITE_COMPILE: # warm up sizes that are not in cudagraph capture sizes, @@ -698,8 +707,11 @@ def compile_or_warm_up_model(self) -> CompilationTimes: # e.g. for the max-num-batched token size in chunked prefill. compile_sizes = self.aphrodite_config.compilation_config.compile_sizes warmup_sizes = compile_sizes.copy() if compile_sizes is not None else [] # type: ignore[assignment] + cg_capture_sizes: list[int] = [] if self.aphrodite_config.compilation_config.cudagraph_mode != CUDAGraphMode.NONE: + cg_sizes = self.aphrodite_config.compilation_config.cudagraph_capture_sizes + cg_capture_sizes = [] if cg_sizes is None else cg_sizes warmup_sizes = [x for x in warmup_sizes if x not in cg_capture_sizes] compile_ranges = self.aphrodite_config.compilation_config.get_compile_ranges() @@ -712,20 +724,6 @@ def compile_or_warm_up_model(self) -> CompilationTimes: if not any(x in compile_range for x in all_sizes): warmup_sizes.append(compile_range.end) - # TODO(LucasWilkinson, akaratza): Remove when MRV1 is deprecated - if ( - current_platform.is_rocm() - and not self.use_v2_model_runner - and self.aphrodite_config.compilation_config.cudagraph_mode != CUDAGraphMode.NONE - and get_pp_group().is_last_rank - ): - max_num_reqs = min( - self.scheduler_config.max_num_seqs, - self.scheduler_config.max_num_batched_tokens, - ) - if max_num_reqs not in cg_capture_sizes and max_num_reqs not in warmup_sizes: - warmup_sizes.append(max_num_reqs) - # We skip EPLB here since we don't want to record dummy metrics for size in sorted(warmup_sizes, reverse=True): logger.info("Compile and warming up model for size %d", size) @@ -876,6 +874,22 @@ def reset_encoder_cache(self) -> None: def get_model(self) -> nn.Module: return self.model_runner.get_model() + def get_draft_model(self) -> nn.Module | None: + return self.model_runner.get_draft_model() + + def _set_draft_weight_update_target(self) -> None: + assert self.weight_transfer_engine is not None + + draft_model = self.get_draft_model() + if draft_model is None: + raise RuntimeError("Draft model weight update requested, but no draft model is configured.") + + speculative_config = self.speculative_config + if speculative_config is None or speculative_config.draft_model_config is None: + raise RuntimeError("Draft model weight update requested, but no draft model config is configured.") + + self.weight_transfer_engine.set_weight_update_target(draft_model, speculative_config.draft_model_config) + def get_supported_tasks(self) -> tuple[SupportedTask, ...]: return self.model_runner.get_supported_tasks() @@ -1110,15 +1124,36 @@ def start_weight_update(self) -> None: the configured weight transfer engine. The worker only tracks that a session is active. """ + self._start_weight_update() + + def start_draft_weight_update(self) -> None: + """ + Like start_weight_update, but retargets the engine at the speculative + draft model for this session. + """ + self._start_weight_update(is_draft=True) + + def _start_weight_update(self, is_draft: bool = False) -> None: self._check_weight_transfer_engine() assert self.weight_transfer_engine is not None + if is_draft and not self.weight_transfer_engine.supports_draft_weight_update: + raise RuntimeError( + f"{type(self.weight_transfer_engine).__name__} does not support draft model weight updates." + ) + if self._weight_update_active: raise RuntimeError( "start_weight_update called while a weight update is already active. Call finish_weight_update first." ) - self.weight_transfer_engine.start_weight_update() + try: + if is_draft: + self._set_draft_weight_update_target() + self.weight_transfer_engine.start_weight_update() + except BaseException: + self.weight_transfer_engine.reset_weight_update_target() + raise self._weight_update_active = True def update_weights(self, update_info: dict) -> None: @@ -1127,6 +1162,8 @@ def update_weights(self, update_info: dict) -> None: start_weight_update must be called before update_weights and finish_weight_update must be called after all chunks have been sent. + Every chunk loads into whichever model the session's start_weight_update + / start_draft_weight_update call selected. Args: update_info: Dictionary containing backend-specific update info @@ -1141,6 +1178,7 @@ def update_weights(self, update_info: dict) -> None: self.weight_transfer_engine.update_weights(update_info) except BaseException: self._weight_update_active = False + self.weight_transfer_engine.reset_weight_update_target() raise def finish_weight_update(self) -> None: @@ -1152,6 +1190,7 @@ def finish_weight_update(self) -> None: raise RuntimeError("finish_weight_update called without a matching start_weight_update.") self.weight_transfer_engine.finish_weight_update() + self.weight_transfer_engine.reset_weight_update_target() self._weight_update_active = False def shutdown(self) -> None: diff --git a/aphrodite/v1/worker/utils.py b/aphrodite/v1/worker/utils.py index 07b9fc3848..cfc4c4acf2 100644 --- a/aphrodite/v1/worker/utils.py +++ b/aphrodite/v1/worker/utils.py @@ -2,11 +2,12 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import math from collections import defaultdict -from collections.abc import Iterable +from collections.abc import Iterable, Sequence from dataclasses import dataclass, field from itertools import product as iprod from typing import Any +import numpy as np import torch from aphrodite.config import AphroditeConfig, CacheConfig @@ -18,11 +19,13 @@ from aphrodite.triton_utils import tl, triton from aphrodite.utils.math_utils import largest_power_of_2_divisor from aphrodite.utils.mem_utils import MemorySnapshot, format_gib +from aphrodite.utils.torch_utils import async_tensor_h2d from aphrodite.v1.attention.backend import ( AttentionBackend, AttentionMetadataBuilder, MultipleOf, ) +from aphrodite.v1.core.kv_cache_utils import KVCacheBlockCopy from aphrodite.v1.kv_cache_interface import ( AttentionSpec, EncoderOnlyAttentionSpec, @@ -510,6 +513,41 @@ def bind_kv_cache( forward_context[layer_name].kv_cache = kv_cache +def copy_kv_cache_blocks_inplace( + kv_caches: Iterable[torch.Tensor | list[torch.Tensor]], + num_blocks: int, + kv_cache_block_copies: Sequence[KVCacheBlockCopy], +) -> None: + if not kv_cache_block_copies: + return + + storage_tensors: list[torch.Tensor] = [] + seen_storage: set[int] = set() + for entry in kv_caches: + tensors = entry if isinstance(entry, (list, tuple)) else (entry,) + for tensor in tensors: + ptr = tensor.untyped_storage().data_ptr() + if ptr in seen_storage: + continue + seen_storage.add(ptr) + storage_tensors.append(tensor) + + if not storage_tensors: + return + device = storage_tensors[0].device + indices_np = np.array(kv_cache_block_copies, dtype=np.int64) + indices = async_tensor_h2d(indices_np, device=device) + src_indices, dst_indices = indices.unbind(dim=1) + + for tensor in storage_tensors: + assert tensor.device == device + blocks = torch.empty(0, dtype=torch.uint8, device=device) + blocks.set_(tensor.untyped_storage()) + assert blocks.numel() % num_blocks == 0 + blocks = blocks.view(num_blocks, -1) + blocks[dst_indices] = blocks[src_indices] + + def is_residual_scattered_for_sp(aphrodite_config: AphroditeConfig, num_input_tokens: int) -> bool: """Check if the residual tensor is scattered for sequence parallelism. diff --git a/build_and_upload_docker.sh b/build_and_upload_docker.sh index 144e8f9e15..6a2d23ea27 100755 --- a/build_and_upload_docker.sh +++ b/build_and_upload_docker.sh @@ -7,7 +7,7 @@ MAX_JOBS=64 NVCC_THREADS=64 CUDA_VERSION=13.0.0 PYTHON_VERSION=3.12 -TORCH_CUDA_ARCH_LIST="7.0 7.5 8.0 8.9 9.0 10.0 12.0" +TORCH_CUDA_ARCH_LIST="7.0 7.5 8.0 8.9 9.0 10.0 11.0 12.0" while [[ "$#" -gt 0 ]]; do case $1 in diff --git a/build_wheel.sh b/build_wheel.sh index ca1a149f69..8d2a1a6f06 100755 --- a/build_wheel.sh +++ b/build_wheel.sh @@ -4,7 +4,7 @@ CUDA_VERSION=${CUDA_VERSION:-13.0.0} PYTHON_VERSION=${PYTHON_VERSION:-3.12} MAX_JOBS=${MAX_JOBS:-} NVCC_THREADS=${NVCC_THREADS:-} -TORCH_CUDA_ARCH_LIST=${TORCH_CUDA_ARCH_LIST:-"7.0 7.5 8.0 8.9 9.0 10.0 12.0"} +TORCH_CUDA_ARCH_LIST=${TORCH_CUDA_ARCH_LIST:-"7.0 7.5 8.0 8.9 9.0 10.0 11.0 12.0"} DOCKER_BUILDKIT=1 docker build . -f docker/Dockerfile --target build --tag alpindale/aphrodite-build \ --build-arg CUDA_VERSION=$CUDA_VERSION \ diff --git a/cmake/external_projects/flashmla.cmake b/cmake/external_projects/flashmla.cmake index 518cd1fcf0..caff230ace 100644 --- a/cmake/external_projects/flashmla.cmake +++ b/cmake/external_projects/flashmla.cmake @@ -19,7 +19,7 @@ else() FetchContent_Declare( flashmla GIT_REPOSITORY https://github.com/vllm-project/FlashMLA - GIT_TAG a6ec2ba7bd0a7dff98b3f4d3e6b52b159c48d78b + GIT_TAG b70aff3d110a2b1a037e62eac295166b5143643a GIT_PROGRESS TRUE CONFIGURE_COMMAND "" BUILD_COMMAND "" @@ -183,4 +183,3 @@ else() add_custom_target(_flashmla_C) add_custom_target(_flashmla_extension_C) endif() - diff --git a/cmake/external_projects/vllm_flash_attn.cmake b/cmake/external_projects/vllm_flash_attn.cmake index e8c3979034..80a0a83b68 100644 --- a/cmake/external_projects/vllm_flash_attn.cmake +++ b/cmake/external_projects/vllm_flash_attn.cmake @@ -39,7 +39,7 @@ else() FetchContent_Declare( aphrodite-flash-attn GIT_REPOSITORY https://github.com/vllm-project/flash-attention.git - GIT_TAG b3964b1d8b95d8e8447435668ab169a2700bab65 + GIT_TAG bb9a72e7dde0dc614ffc663e052cd6a19ce73a42 GIT_PROGRESS TRUE # Don't share the aphrodite-flash-attn build between build types BINARY_DIR ${CMAKE_BINARY_DIR}/aphrodite-flash-attn diff --git a/csrc/libtorch_stable/fp32_router_gemm.cu b/csrc/libtorch_stable/fp32_router_gemm.cu index 64393fad61..3207082d57 100644 --- a/csrc/libtorch_stable/fp32_router_gemm.cu +++ b/csrc/libtorch_stable/fp32_router_gemm.cu @@ -1,13 +1,18 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright contributors to the vLLM project // -// Router GEMM: activation(T) x weight(fp32) -> fp32, H=3072, E=256, M<=32. +// Router GEMM: activation(T) x weight(fp32) -> fp32, M<=32, for the +// supported (E, H) pairs listed at the bottom of this file. // Supports bf16 or fp32 activation; weight is always fp32. // Adapted from dsv3_router_gemm_float_out.cu. +// (E=256, H=6144) bf16 uses a B300-tuned wide-block geometry; see +// invokeFp32RouterGemm. #include #include +#include + // --------------------------------------------------------------------------- // Load helpers // --------------------------------------------------------------------------- @@ -73,94 +78,113 @@ __device__ __forceinline__ void load_activation<__nv_bfloat16, 8>( // InputT : type of activation (float or __nv_bfloat16) // Weight is always fp32; output is always fp32. // VPT = 16 / sizeof(InputT): 4 for fp32, 8 for bf16 -template -__global__ __launch_bounds__(128, 1) void fp32_router_gemm_kernel( - float* out, InputT const* mat_a, float const* mat_b) { +// Each block computes kEPB expert columns; wider blocks / kEPB > 1 are +// selected per (shape, M) in invokeFp32RouterGemm (B300-tuned, see below). +// kTGroups > 1 splits the tokens across groups of kBlockSize threads within +// the block: all groups scan the same weight K-slices (group 0 misses to +// DRAM, later groups hit L1) so weight traffic stays 1x, while per-thread +// accumulator registers drop by kTGroups (at M=16 the 32 fp32 accumulators +// push the kernel to 128 regs/thread and 1 block/SM). +template +__global__ __launch_bounds__( + kBlockSize* kTGroups, 1) void fp32_router_gemm_kernel(float* out, + InputT const* mat_a, + float const* mat_b) { constexpr int VPT = 16 / sizeof(InputT); constexpr int k_elems_per_k_iteration = VPT * kBlockSize; constexpr int k_iterations = kHiddenDim / k_elems_per_k_iteration; + static_assert(kHiddenDim % k_elems_per_k_iteration == 0); + static_assert(kNumTokens % kTGroups == 0); constexpr int kWarpSize = 32; - constexpr int kNumWarps = kBlockSize / kWarpSize; + constexpr int kNumWarps = kBlockSize / kWarpSize; // per token group + constexpr int kMG = kNumTokens / kTGroups; // tokens per group - int const n_idx = blockIdx.x; - int const tid = threadIdx.x; + int const e_base = blockIdx.x * kEPB; + int const tid = threadIdx.x % kBlockSize; + int const m0 = (threadIdx.x / kBlockSize) * kMG; int const warpId = tid / kWarpSize; int const laneId = tid % kWarpSize; - float acc[kNumTokens] = {}; - __shared__ float sm_reduction[kNumTokens][kNumWarps]; - - float const* b_col = mat_b + n_idx * kHiddenDim; - - int k_bases[k_iterations]; -#pragma unroll - for (int ki = 0; ki < k_iterations; ki++) { - k_bases[ki] = ki * k_elems_per_k_iteration + tid * VPT; - } + float acc[kMG][kEPB] = {}; + __shared__ float sm_reduction[kNumTokens][kEPB][kNumWarps]; #if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) cudaGridDependencySynchronize(); + // Fire the PDL trigger right after our own wait instead of at kernel end: + // a gridsync-ing consumer is unaffected (its wait always targets full grid + // completion), while a consumer that reads none of our outputs (e.g. the + // NVFP4 activation quant, which reads the same hidden_states) can launch + // now and fully overlap this kernel's body. + cudaTriggerProgrammaticLaunchCompletion(); #endif +#pragma unroll for (int ki = 0; ki < k_iterations; ki++) { - int const k_base = k_bases[ki]; + int const k_base = ki * k_elems_per_k_iteration + tid * VPT; - float b_float[VPT]; - load_weight(b_col + k_base, b_float); + float b_float[kEPB][VPT]; +#pragma unroll + for (int e = 0; e < kEPB; e++) { + load_weight(mat_b + (e_base + e) * kHiddenDim + k_base, b_float[e]); + } #pragma unroll - for (int m_idx = 0; m_idx < kNumTokens; m_idx++) { + for (int m_idx = 0; m_idx < kMG; m_idx++) { float a_float[VPT]; - load_activation(mat_a + m_idx * kHiddenDim + k_base, - a_float); + load_activation( + mat_a + (size_t)(m0 + m_idx) * kHiddenDim + k_base, a_float); +#pragma unroll + for (int e = 0; e < kEPB; e++) { #pragma unroll - for (int k = 0; k < VPT; k++) { - acc[m_idx] += a_float[k] * b_float[k]; + for (int k = 0; k < VPT; k++) { + acc[m_idx][e] += a_float[k] * b_float[e][k]; + } } } } // Warp-level butterfly reduction #pragma unroll - for (int m = 0; m < kNumTokens; m++) { - float sum = acc[m]; - sum += __shfl_xor_sync(0xffffffff, sum, 16); - sum += __shfl_xor_sync(0xffffffff, sum, 8); - sum += __shfl_xor_sync(0xffffffff, sum, 4); - sum += __shfl_xor_sync(0xffffffff, sum, 2); - sum += __shfl_xor_sync(0xffffffff, sum, 1); - if (laneId == 0) sm_reduction[m][warpId] = sum; + for (int m = 0; m < kMG; m++) { +#pragma unroll + for (int e = 0; e < kEPB; e++) { + float sum = acc[m][e]; + sum += __shfl_xor_sync(0xffffffff, sum, 16); + sum += __shfl_xor_sync(0xffffffff, sum, 8); + sum += __shfl_xor_sync(0xffffffff, sum, 4); + sum += __shfl_xor_sync(0xffffffff, sum, 2); + sum += __shfl_xor_sync(0xffffffff, sum, 1); + if (laneId == 0) sm_reduction[m0 + m][e][warpId] = sum; + } } __syncthreads(); - if (tid == 0) { -#pragma unroll - for (int m = 0; m < kNumTokens; m++) { - float final_sum = 0.0f; + // Parallel finalize: one thread per (m, e) output. + for (int idx = threadIdx.x; idx < kNumTokens * kEPB; + idx += kBlockSize * kTGroups) { + int const m = idx / kEPB; + int const e = idx % kEPB; + float final_sum = 0.0f; #pragma unroll - for (int w = 0; w < kNumWarps; w++) final_sum += sm_reduction[m][w]; - out[m * kNumExperts + n_idx] = final_sum; - } + for (int w = 0; w < kNumWarps; w++) final_sum += sm_reduction[m][e][w]; + out[m * kNumExperts + e_base + e] = final_sum; } - -#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) - cudaTriggerProgrammaticLaunchCompletion(); -#endif } // --------------------------------------------------------------------------- // Launcher // --------------------------------------------------------------------------- -template -void invokeFp32RouterGemm(float* output, InputT const* mat_a, - float const* mat_b, cudaStream_t stream) { - constexpr int kBlockSize = 128; +template +static void launchFp32RouterGemm(float* output, InputT const* mat_a, + float const* mat_b, cudaStream_t stream) { + static_assert(kNumExperts % kEPB == 0); cudaLaunchConfig_t config; - config.gridDim = kNumExperts; - config.blockDim = kBlockSize; + config.gridDim = kNumExperts / kEPB; + config.blockDim = kBlockSize * kTGroups; config.dynamicSmemBytes = 0; config.stream = stream; cudaLaunchAttribute attrs[1]; @@ -168,15 +192,112 @@ void invokeFp32RouterGemm(float* output, InputT const* mat_a, attrs[0].val.programmaticStreamSerializationAllowed = 1; config.numAttrs = 1; config.attrs = attrs; - cudaLaunchKernelEx(&config, - fp32_router_gemm_kernel, - output, mat_a, mat_b); + cudaLaunchKernelEx( + &config, + fp32_router_gemm_kernel, + output, mat_a, mat_b); +} + +static bool isBlackwellFamily() { + static int sm = []() { + int dev = 0, major = 0, minor = 0; + cudaGetDevice(&dev); + cudaDeviceGetAttribute(&major, cudaDevAttrComputeCapabilityMajor, dev); + cudaDeviceGetAttribute(&minor, cudaDevAttrComputeCapabilityMinor, dev); + return major * 10 + minor; + }(); + return sm >= 100; +} + +template +void invokeFp32RouterGemm(float* output, InputT const* mat_a, + float const* mat_b, cudaStream_t stream) { + // Geometry tuned on B300 per supported shape, bf16 activation, under a + // production-fidelity harness (CUDA-graph replay, per-layer cold weights). + // GLM-5.2 (E=256, H=6144): + // M <= 4 : BS=768, EPB=1 (2.7us vs cast+cuBLAS 8.1us at M=1) + // M in [5, 15] + // or odd : BS=384, EPB=2 (crossover vs BS=768 measured in (4, 8)) + // M >= 16, even : BS=192, EPB=2, 2 token groups (M=16 4.79us vs 5.04, + // M=24 5.71 vs 6.38, M=32 6.79 vs 7.72; M=12 loses at + // 0.97x, so the boundary is 16). + // Only enabled on the Blackwell family where it was validated; Hopper and + // other shapes / fp32 activation keep the legacy geometry. + if constexpr (std::is_same_v && kNumExperts == 256 && + kHiddenDim == 6144) { + if (!isBlackwellFamily()) { + launchFp32RouterGemm( + output, mat_a, mat_b, stream); + return; + } + if constexpr (kNumTokens <= 4) { + launchFp32RouterGemm( + output, mat_a, mat_b, stream); + } else if constexpr (kNumTokens >= 16 && kNumTokens % 2 == 0) { + launchFp32RouterGemm(output, mat_a, mat_b, stream); + } else { + launchFp32RouterGemm( + output, mat_a, mat_b, stream); + } + } else if constexpr (std::is_same_v && + kNumExperts == 128 && kHiddenDim == 6144) { + // MiniMax-M3. Legacy 128/1 only fills 128 blocks and pays the same + // accumulator register cliffs; B300 sweep: + // even M in [6, 10] : BS=384, EPB=1, 2 token groups (1.26-1.43x) + // even M >= 12 : BS=192, EPB=1, 2 token groups (1.59-1.66x at + // M >= 18; re-measured on B300+B200: 192 also wins + // M=12/14 by 5-11%% on both, ties 384 at 16) + // M <= 5 / odd : BS=384, EPB=1 (1.03-1.19x) + if (!isBlackwellFamily()) { + launchFp32RouterGemm( + output, mat_a, mat_b, stream); + return; + } + if constexpr (kNumTokens >= 12 && kNumTokens % 2 == 0) { + launchFp32RouterGemm(output, mat_a, mat_b, stream); + } else if constexpr (kNumTokens >= 6 && kNumTokens % 2 == 0) { + launchFp32RouterGemm(output, mat_a, mat_b, stream); + } else { + launchFp32RouterGemm( + output, mat_a, mat_b, stream); + } + } else if constexpr (std::is_same_v && + kNumExperts == 256 && kHiddenDim == 3072) { + // MiniMax-M2/M2.5. The 3.1MB weight is latency-floor bound at small M + // (legacy already optimal); token groups win only at even M >= 8 + // (1.05-1.17x). EPB crossover measured between 12 and 16. + if (!isBlackwellFamily()) { + launchFp32RouterGemm( + output, mat_a, mat_b, stream); + return; + } + if constexpr (kNumTokens >= 14 && kNumTokens % 2 == 0) { + // M=14 originally measured 0.91x and stayed on legacy; two fresh + // sweeps (B300 dev1 + B200) both put 192/2/tg2 ahead by 3.5-4%%. + launchFp32RouterGemm(output, mat_a, mat_b, stream); + } else if constexpr (kNumTokens >= 8 && kNumTokens <= 12 && + kNumTokens % 2 == 0) { + launchFp32RouterGemm(output, mat_a, mat_b, stream); + } else { + launchFp32RouterGemm( + output, mat_a, mat_b, stream); + } + } else { + launchFp32RouterGemm( + output, mat_a, mat_b, stream); + } } // --------------------------------------------------------------------------- // Explicit instantiations: M=1..32, for both input types, for the supported -// (E, H) pairs: (256, 3072) [MiniMax-M2/M2.5] and (128, 6144) [MiniMax-M3]. +// (E, H) pairs: (256, 3072) [MiniMax-M2/M2.5], (128, 6144) [MiniMax-M3] +// and (256, 6144) [GLM-5.2]. // --------------------------------------------------------------------------- #define INSTANTIATE(T, M, E, H) \ @@ -221,6 +342,8 @@ INSTANTIATE_ALL(float, 256, 3072) INSTANTIATE_ALL(__nv_bfloat16, 256, 3072) INSTANTIATE_ALL(float, 128, 6144) INSTANTIATE_ALL(__nv_bfloat16, 128, 6144) +INSTANTIATE_ALL(float, 256, 6144) +INSTANTIATE_ALL(__nv_bfloat16, 256, 6144) #undef INSTANTIATE_ALL #undef INSTANTIATE diff --git a/csrc/libtorch_stable/fp32_router_gemm_entry.cu b/csrc/libtorch_stable/fp32_router_gemm_entry.cu index fc09193c86..8643e33cbe 100644 --- a/csrc/libtorch_stable/fp32_router_gemm_entry.cu +++ b/csrc/libtorch_stable/fp32_router_gemm_entry.cu @@ -25,10 +25,12 @@ inline int getSMVersion() { static constexpr int FP32_MAX_TOKENS = 32; // Supported (hidden_dim, num_experts) pairs (must match the instantiations in -// fp32_router_gemm.cu): (3072, 256) for MiniMax-M2/M2.5, (6144, 128) for M3. +// fp32_router_gemm.cu): (3072, 256) for MiniMax-M2/M2.5, (6144, 128) for M3, +// (6144, 256) for GLM-5.2. static inline bool fp32_router_gemm_supported(int hidden_dim, int num_experts) { return (hidden_dim == 3072 && num_experts == 256) || - (hidden_dim == 6144 && num_experts == 128); + (hidden_dim == 6144 && num_experts == 128) || + (hidden_dim == 6144 && num_experts == 256); } // Forward declarations — 4 template params must match fp32_router_gemm.cu @@ -77,6 +79,9 @@ void dispatchFp32RouterGemm(int num_experts, int hidden_dim, int num_tokens, } else if (num_experts == 128 && hidden_dim == 6144) { Fp32LoopUnroller::unroll( num_tokens, output, mat_a, mat_b, stream); + } else if (num_experts == 256 && hidden_dim == 6144) { + Fp32LoopUnroller::unroll( + num_tokens, output, mat_a, mat_b, stream); } else { throw std::invalid_argument( "fp32_router_gemm: unsupported (hidden_dim, num_experts) pair"); @@ -111,7 +116,7 @@ void fp32_router_gemm( STD_TORCH_CHECK( fp32_router_gemm_supported(hidden_dim, num_experts), "fp32_router_gemm: supported (hidden_dim, num_experts) pairs are " - "(3072, 256) and (6144, 128)"); + "(3072, 256), (6144, 128) and (6144, 256)"); STD_TORCH_CHECK(num_tokens <= FP32_MAX_TOKENS, "fp32_router_gemm: num_tokens must be in [0, 32]"); STD_TORCH_CHECK( diff --git a/csrc/libtorch_stable/fused_minimax_m3_qknorm_rope_kv_insert_kernel.cu b/csrc/libtorch_stable/fused_minimax_m3_qknorm_rope_kv_insert_kernel.cu index 74ef917a2f..9ecc0a282d 100644 --- a/csrc/libtorch_stable/fused_minimax_m3_qknorm_rope_kv_insert_kernel.cu +++ b/csrc/libtorch_stable/fused_minimax_m3_qknorm_rope_kv_insert_kernel.cu @@ -288,16 +288,15 @@ __global__ void fusedMiniMaxM3QNormRopeKVInsertKernel( int64_t const* __restrict__ positions, // [N] i64 int64_t const* __restrict__ slot_mapping, // main K/V slots or nullptr int64_t const* __restrict__ index_slot_mapping, // index K slots/nullptr - cache_t* __restrict__ kv_cache, // [nb,2,bs,nkv,128] or nullptr + cache_t* __restrict__ kv_cache, // [nb,nkv,bs,2*128] or nullptr out_idx_t* __restrict__ index_cache, // [nb*bs, 128]; scalar_t or e4m3 byte float const eps, int const rotary_dim, int const num_tokens, int const nq, int const nkv, int const niq, int const block_size, - // kv_cache strides (in elements) for logical shape [nb, 2, bs, nkv, 128]. - // The head_dim (last) dim is always innermost-contiguous (stride 1), so the - // NHD/HND layout choice is fully captured by these four strides: NHD keeps - // s_token < s_head, HND swaps them. dim_base addresses head_dim directly. - int64_t const kv_s_block, int64_t const kv_s_kv, int64_t const kv_s_token, - int64_t const kv_s_head) { + // kv_cache strides (in elements) for logical shape [nb, nkv, bs, 2*128]. + // The content (last) dim is always innermost-contiguous (stride 1), so the + // NHD/HND layout choice is captured by the head/token strides. + int64_t const kv_s_block, int64_t const kv_s_head, int64_t const kv_s_token, + int64_t const kv_s_dim) { #if (!defined(__CUDA_ARCH__) || __CUDA_ARCH__ < 800) && !defined(USE_ROCM) // _typeConvert is unavailable on pre-Ampere; the M3 kernel only // runs with bf16/fp16 inputs in practice. Discard the bf16 body there. @@ -438,16 +437,16 @@ __global__ void fusedMiniMaxM3QNormRopeKVInsertKernel( storeElems(index_cache + sm * kHeadDim + dim_base, elems); } } else if (isK || isV) { - // kv_cache logical shape [num_blocks, 2, block_size, nkv, head_dim]. + // kv_cache logical shape [num_blocks, nkv, block_size, 2*head_dim]. // Paging is logical (block = sm/block_size, token = sm%block_size); // the physical NHD/HND layout is honoured via the passed strides. int64_t const b = sm / block_size; int64_t const t = sm % block_size; int const kv = isK ? 0 : 1; - int64_t const off = - b * kv_s_block + kv * kv_s_kv + t * kv_s_token + head * kv_s_head; - storeCacheElems(kv_cache + off + dim_base, - elems); + int64_t const off = b * kv_s_block + head * kv_s_head + + t * kv_s_token + + (kv * kHeadDim + dim_base) * kv_s_dim; + storeCacheElems(kv_cache + off, elems); } } } @@ -474,8 +473,8 @@ void launchFusedMiniMaxM3( int64_t const* index_slot_mapping, cache_t* kv_cache, void* index_cache, float const eps, int const rotary_dim, int const num_tokens, int const nq, int const nkv, int const niq, int const block_size, - int64_t const kv_s_block, int64_t const kv_s_kv, int64_t const kv_s_token, - int64_t const kv_s_head, bool const has_index, bool const insert_kv, + int64_t const kv_s_block, int64_t const kv_s_head, int64_t const kv_s_token, + int64_t const kv_s_dim, bool const has_index, bool const insert_kv, bool const fp8_idx, cudaStream_t stream) { // Index outputs are scalar_t (bf16) or e4m3 bytes (uint8_t); reinterpret the // void* pointers per instantiation in the LAUNCH macro. @@ -517,20 +516,20 @@ void launchFusedMiniMaxM3( iq_norm_w, ik_norm_w, cos_sin_cache, positions, slot_mapping, \ index_slot_mapping, kv_cache, reinterpret_cast(index_cache), \ eps, rotary_dim, num_tokens, nq, nkv, niq, block_size, kv_s_block, \ - kv_s_kv, kv_s_token, kv_s_head) + kv_s_head, kv_s_token, kv_s_dim) #else // ROCm: standard kernel launch syntax (no PDL/stream serialization). // clang-format off #define LAUNCH(IS_SPARSE, INSERT, FP8, OUT_T) \ - fusedMiniMaxM3QNormRopeKVInsertKernel \ <<>>( \ qkv, q_out, reinterpret_cast(index_q_out), q_norm_w, \ k_norm_w, iq_norm_w, ik_norm_w, cos_sin_cache, positions, \ slot_mapping, index_slot_mapping, kv_cache, \ reinterpret_cast(index_cache), eps, rotary_dim, \ - num_tokens, nq, nkv, niq, block_size, kv_s_block, kv_s_kv, \ - kv_s_token, kv_s_head) + num_tokens, nq, nkv, niq, block_size, kv_s_block, kv_s_head, \ + kv_s_token, kv_s_dim) // clang-format on #endif @@ -585,8 +584,8 @@ void launchFusedMiniMaxM3( ? reinterpret_cast(index_cache->data_ptr()) \ : nullptr, \ static_cast(eps), static_cast(rotary_dim), num_tokens, nq, \ - nkv, niq, static_cast(block_size), kv_s_block, kv_s_kv, kv_s_token, \ - kv_s_head, has_index, insert_kv, fp8_idx, stream) + nkv, niq, static_cast(block_size), kv_s_block, kv_s_head, \ + kv_s_token, kv_s_dim, has_index, insert_kv, fp8_idx, stream) // ──────────────────────────────────────────────────────────────────────────── // Torch op wrapper @@ -603,7 +602,7 @@ void fused_minimax_m3_qknorm_rope_kv_insert( int64_t num_index_heads, // niq; 0 => dense std::optional slot_mapping, // [N] i64 std::optional index_slot_mapping, // [N] i64 - std::optional kv_cache, // [nb,2,bs,nkv,128] + std::optional kv_cache, // [nb,nkv,bs,2*128] std::optional index_cache, // [nb,bs,128] int64_t block_size, std::optional q_out, // [N, nq*128] contiguous @@ -674,11 +673,11 @@ void fused_minimax_m3_qknorm_rope_kv_insert( index_k_norm_weight->numel() == kHeadDim, "index norm weights must have 128 elements"); } - // kv_cache strides (logical shape [nb, 2, bs, nkv, head_dim]). Read straight + // kv_cache strides (logical shape [nb, nkv, bs, 2*head_dim]). Read straight // off the tensor so the kernel honours whatever physical layout the attention - // backend allocated (NHD: stride order (0,1,2,3,4); HND: (0,1,3,2,4)). No new + // backend allocated (NHD: stride order (0,2,1,3); HND: (0,1,2,3)). No new // op argument is needed -- the strides ride along with the tensor itself. - int64_t kv_s_block = 0, kv_s_kv = 0, kv_s_token = 0, kv_s_head = 0; + int64_t kv_s_block = 0, kv_s_head = 0, kv_s_token = 0, kv_s_dim = 0; torch::stable::Tensor const* effective_index_slot_mapping = nullptr; if (insert_kv) { STD_TORCH_CHECK( @@ -708,13 +707,13 @@ void fused_minimax_m3_qknorm_rope_kv_insert( index_cache->scalar_type() == torch::headeronly::ScalarType::Float8_e4m3fn), "insert mode requires index_cache matching qkv dtype or fp8 e4m3"); - STD_TORCH_CHECK(kv_cache->dim() == 5 && kv_cache->stride(4) == 1, - "kv_cache must be [nb,2,bs,nkv,head_dim] with contiguous " - "head_dim (stride(4)==1)"); + STD_TORCH_CHECK(kv_cache->dim() == 4 && kv_cache->stride(3) == 1, + "kv_cache must be [nb,nkv,bs,2*head_dim] with contiguous " + "content dim (stride(3)==1)"); kv_s_block = kv_cache->stride(0); - kv_s_kv = kv_cache->stride(1); + kv_s_head = kv_cache->stride(1); kv_s_token = kv_cache->stride(2); - kv_s_head = kv_cache->stride(3); + kv_s_dim = kv_cache->stride(3); effective_index_slot_mapping = index_slot_mapping.has_value() ? &index_slot_mapping.value() : &slot_mapping.value(); diff --git a/csrc/libtorch_stable/moe/moe_align_sum_kernels.cu b/csrc/libtorch_stable/moe/moe_align_sum_kernels.cu index d94d644365..448a3f2e03 100644 --- a/csrc/libtorch_stable/moe/moe_align_sum_kernels.cu +++ b/csrc/libtorch_stable/moe/moe_align_sum_kernels.cu @@ -82,6 +82,21 @@ __global__ void batched_moe_align_block_size_kernel( } } // namespace batched_moe_align_block_size +template +__device__ __forceinline__ int get_local_expert_id( + size_t idx, const scalar_t* __restrict__ topk_ids, + int32_t* __restrict__ expert_map, int32_t num_experts, + bool has_expert_map) { + int expert_id = topk_ids[idx]; + if (expert_id >= num_experts || expert_id < 0) { + return -1; + } + if (has_expert_map) { + expert_id = expert_map[expert_id]; + } + return expert_id; +} + template __device__ void _moe_align_block_size( const scalar_t* __restrict__ topk_ids, @@ -126,20 +141,15 @@ __device__ void _moe_align_block_size( const size_t stride = blockDim.x; for (size_t i = tid; i < numel; i += stride) { - int expert_id = topk_ids[i]; - if (expert_id >= num_experts) { - continue; + if (int expert_id = get_local_expert_id(i, topk_ids, expert_map, + num_experts, has_expert_map); + expert_id != -1) { + int warp_idx = expert_id / experts_per_warp; + int expert_offset = expert_id % experts_per_warp; + int mask = token_mask == nullptr ? 1 : token_mask[i / topk_num]; + atomicAdd(&shared_counts[warp_idx * experts_per_warp + expert_offset], + mask); } - if (has_expert_map) { - expert_id = expert_map[expert_id]; - // filter invalid experts - if (expert_id == -1) continue; - } - int warp_idx = expert_id / experts_per_warp; - int expert_offset = expert_id % experts_per_warp; - int mask = token_mask == nullptr ? 1 : token_mask[i / topk_num]; - atomicAdd(&shared_counts[warp_idx * experts_per_warp + expert_offset], - mask); } __syncthreads(); @@ -227,14 +237,12 @@ __device__ void _moe_align_block_size_small_batch_expert( } for (size_t i = tid; i < numel; i += stride) { - int32_t expert_id = topk_ids[i]; - if (has_expert_map) { - expert_id = expert_map[expert_id]; - // filter invalid expert - if (expert_id == -1) continue; + if (int expert_id = get_local_expert_id(i, topk_ids, expert_map, + num_experts, has_expert_map); + expert_id != -1) { + int mask = token_mask == nullptr ? 1 : token_mask[i / topk_num]; + tokens_cnts[(tid + 1) * num_experts + expert_id] += mask; } - int mask = token_mask == nullptr ? 1 : token_mask[i / topk_num]; - tokens_cnts[(tid + 1) * num_experts + expert_id] += mask; } __syncthreads(); @@ -276,18 +284,16 @@ __device__ void _moe_align_block_size_small_batch_expert( } for (size_t i = tid; i < numel; i += stride) { - int32_t expert_id = topk_ids[i]; - if (has_expert_map) { - expert_id = expert_map[expert_id]; - // filter invalid expert - if (expert_id == -1) continue; - } - int32_t rank_post_pad = - tokens_cnts[tid * num_experts + expert_id] + cumsum[expert_id]; - - if (token_mask == nullptr || token_mask[i / topk_num]) { - sorted_token_ids[sorted_token_ids_offset + rank_post_pad] = i; - ++tokens_cnts[tid * num_experts + expert_id]; + if (int expert_id = get_local_expert_id(i, topk_ids, expert_map, + num_experts, has_expert_map); + expert_id != -1) { + int32_t rank_post_pad = + tokens_cnts[tid * num_experts + expert_id] + cumsum[expert_id]; + + if (token_mask == nullptr || token_mask[i / topk_num]) { + sorted_token_ids[sorted_token_ids_offset + rank_post_pad] = i; + ++tokens_cnts[tid * num_experts + expert_id]; + } } } } @@ -303,22 +309,15 @@ __device__ void _count_and_sort_expert_tokens( const size_t stride = blockDim.x * gridDim.y; for (size_t i = tid; i < numel; i += stride) { - int32_t expert_id = topk_ids[i]; - if (expert_id >= num_experts) { - continue; - } - - if (has_expert_map) { - expert_id = expert_map[expert_id]; - // filter invalid experts - if (expert_id == -1) continue; - } - - if (token_mask == nullptr || token_mask[i / topk_num]) { - int32_t rank_post_pad = atomicAdd( - &cumsum_buffer[(model_offset * (num_experts + 1)) + expert_id], 1); - sorted_token_ids[max_num_tokens_padded * model_offset + rank_post_pad] = - i; + if (int expert_id = get_local_expert_id(i, topk_ids, expert_map, + num_experts, has_expert_map); + expert_id != -1) { + if (token_mask == nullptr || token_mask[i / topk_num]) { + int32_t rank_post_pad = atomicAdd( + &cumsum_buffer[(model_offset * (num_experts + 1)) + expert_id], 1); + sorted_token_ids[max_num_tokens_padded * model_offset + rank_post_pad] = + i; + } } } } diff --git a/csrc/libtorch_stable/ngram_embedding_kernels.cu b/csrc/libtorch_stable/ngram_embedding_kernels.cu new file mode 100644 index 0000000000..a89a2d1c99 --- /dev/null +++ b/csrc/libtorch_stable/ngram_embedding_kernels.cu @@ -0,0 +1,96 @@ +// N-gram embedding index kernel for LongCat-Flash (n-gram embedding variant). +// +// Adapted from SGLang: +// https://github.com/sgl-project/sglang/blob/main/python/sglang/jit_kernel/csrc/ngram_embedding.cuh +// +// For each position, computes the hashed n-gram embedding ids that index the +// concatenated embedder table. Integer tensors are int32 except ``row_indices`` +// (int64); the token table is ``[max_running_reqs, max_context_len]`` int32, +// where a negative entry marks an ignored token (e.g. an EOS boundary). + +#include "torch_utils.h" + +#include "ops.h" + +#include + +namespace aphrodite::ngram_embedding { + +constexpr int kBlockThreads = 256; + +__global__ void ComputeNGramIdsKernel( + int batch_size, int ne_n, int ne_k, + int* ne_weights, // [ne_n-1, ne_k, ne_n] + int* ne_mods, // [ne_n-1, ne_k] + int* exclusive_ne_embedder_size_sums, // [(ne_n-1)*ne_k + 1] + int* exclusive_req_len_sums, // [batch_size + 1] + int* ne_token_table, // [max_running_reqs, max_context_len] + int max_context_len, + const int64_t* __restrict__ row_indices, // [batch_size] + int* column_starts, // [batch_size] + int* n_gram_ids // [token_num, (ne_n-1)*ne_k] +) { + const int req_id = blockIdx.x % batch_size; + const int config_id = (blockIdx.x - req_id) / batch_size; + // n and k are offset from their physical meaning: n = real_n - 2, k = real_k + // - 1 (they index into ne_weights / ne_mods). + const int k = config_id % ne_k; + const int n = (config_id - config_id % ne_k) / ne_k; + const int ne_weight_base_idx = n * ne_k * ne_n + k * ne_n; + const int ne_mod = ne_mods[n * ne_k + k]; + for (int i = exclusive_req_len_sums[req_id] + threadIdx.x; + i < exclusive_req_len_sums[req_id + 1]; i += blockDim.x) { + uint64_t n_gram_id = 0; + const int64_t current_token_offset = i - exclusive_req_len_sums[req_id]; + const int64_t req_token_table_index = + row_indices[req_id] * static_cast(max_context_len); + const int64_t current_token_table_index = + req_token_table_index + column_starts[req_id] + current_token_offset; + for (int j = 0; j < n + 2; j++) { + if (current_token_table_index - j < req_token_table_index) { + break; // outside this request's range + } + if (ne_token_table[current_token_table_index - j] < 0) { + break; // ignored token + } + const uint64_t term = + (uint64_t)ne_token_table[current_token_table_index - j] * + (uint64_t)ne_weights[ne_weight_base_idx + j]; + n_gram_id += term % ne_mod; + } + n_gram_id %= ne_mod; + n_gram_id += exclusive_ne_embedder_size_sums[n * ne_k + k]; + n_gram_ids[i * (ne_n - 1) * ne_k + n * ne_k + k] = (int)(n_gram_id); + } +} + +} // namespace aphrodite::ngram_embedding + +void ngram_compute_n_gram_ids( + int64_t ne_n, int64_t ne_k, torch::stable::Tensor& ne_weights, + torch::stable::Tensor& ne_mods, + torch::stable::Tensor& exclusive_ne_embedder_size_sums, + torch::stable::Tensor& exclusive_req_len_sums, + torch::stable::Tensor& ne_token_table, torch::stable::Tensor& row_indices, + torch::stable::Tensor& column_starts, torch::stable::Tensor& n_gram_ids) { + const int batch_size = static_cast(exclusive_req_len_sums.size(0) - 1); + const int max_context_len = static_cast(ne_token_table.size(1)); + const int num_configs = (static_cast(ne_n) - 1) * static_cast(ne_k); + const int grid_size = num_configs * batch_size; + if (grid_size <= 0) return; + + const torch::stable::accelerator::DeviceGuard device_guard( + ne_weights.get_device_index()); + const cudaStream_t stream = get_current_cuda_stream(); + aphrodite::ngram_embedding::ComputeNGramIdsKernel<<< + grid_size, aphrodite::ngram_embedding::kBlockThreads, 0, stream>>>( + batch_size, static_cast(ne_n), static_cast(ne_k), + ne_weights.mutable_data_ptr(), + ne_mods.mutable_data_ptr(), + exclusive_ne_embedder_size_sums.mutable_data_ptr(), + exclusive_req_len_sums.mutable_data_ptr(), + ne_token_table.mutable_data_ptr(), max_context_len, + row_indices.const_data_ptr(), + column_starts.mutable_data_ptr(), + n_gram_ids.mutable_data_ptr()); +} diff --git a/csrc/libtorch_stable/ops.h b/csrc/libtorch_stable/ops.h index 56a1173f09..cd087d474e 100644 --- a/csrc/libtorch_stable/ops.h +++ b/csrc/libtorch_stable/ops.h @@ -586,3 +586,12 @@ void cp_gather_indexer_k_quant_cache( // quant_block_size * 4] const torch::stable::Tensor& block_table, // [batch_size, num_blocks] const torch::stable::Tensor& cu_seq_lens); // [batch_size + 1] + +// LongCat n-gram embedding index kernel (see ngram_embedding_kernels.cu). +void ngram_compute_n_gram_ids( + int64_t ne_n, int64_t ne_k, torch::stable::Tensor& ne_weights, + torch::stable::Tensor& ne_mods, + torch::stable::Tensor& exclusive_ne_embedder_size_sums, + torch::stable::Tensor& exclusive_req_len_sums, + torch::stable::Tensor& ne_token_table, torch::stable::Tensor& row_indices, + torch::stable::Tensor& column_starts, torch::stable::Tensor& n_gram_ids); diff --git a/csrc/libtorch_stable/torch_bindings.cpp b/csrc/libtorch_stable/torch_bindings.cpp index d675a0a6b0..581310f2db 100644 --- a/csrc/libtorch_stable/torch_bindings.cpp +++ b/csrc/libtorch_stable/torch_bindings.cpp @@ -661,9 +661,22 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) { "Tensor? initial_state_idx," "Tensor? cu_chunk_seqlen," "Tensor? last_chunk_indices) -> ()"); + + // LongCat n-gram embedding index kernel. All tensor args are marked mutable + // to match the (non-const) stable-Tensor& C++ signature; only ne_token_table + // and n_gram_ids are actually written in place. + ops.def( + "ngram_compute_n_gram_ids(int ne_n, int ne_k, Tensor(a!) ne_weights, " + "Tensor(b!) ne_mods, Tensor(c!) exclusive_ne_embedder_size_sums, " + "Tensor(d!) exclusive_req_len_sums, Tensor(e!) ne_token_table, " + "Tensor(f!) row_indices, Tensor(g!) column_starts, " + "Tensor(h!) n_gram_ids) -> ()"); } STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) { + // LongCat n-gram embedding index kernel. + ops.impl("ngram_compute_n_gram_ids", TORCH_BOX(&ngram_compute_n_gram_ids)); + // Per-token group quantization ops.impl("per_token_group_fp8_quant", TORCH_BOX(&per_token_group_quant_fp8)); ops.impl("per_token_group_fp8_quant_packed", diff --git a/docker/Dockerfile b/docker/Dockerfile index 52dac54966..b32ac2d09c 100644 --- a/docker/Dockerfile +++ b/docker/Dockerfile @@ -219,7 +219,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \ # Explicitly set the list to avoid issues with torch 2.2 # See https://github.com/pytorch/pytorch/pull/123243 # From versions.json: .torch.cuda_arch_list -ARG torch_cuda_arch_list='7.5 8.0 8.6 8.9 9.0 10.0 12.0+PTX' +ARG torch_cuda_arch_list='7.5 8.0 8.6 8.9 9.0 10.0 11.0 12.0+PTX' ENV TORCH_CUDA_ARCH_LIST=${torch_cuda_arch_list} #################### BUILD BASE IMAGE #################### @@ -655,7 +655,7 @@ ARG PIP_EXTRA_INDEX_URL UV_EXTRA_INDEX_URL ENV UV_HTTP_TIMEOUT=500 # install kv_connectors if requested -ARG torch_cuda_arch_list='7.5 8.0 8.6 8.9 9.0 10.0 12.0+PTX' +ARG torch_cuda_arch_list='7.5 8.0 8.6 8.9 9.0 10.0 11.0 12.0+PTX' ENV TORCH_CUDA_ARCH_LIST=${torch_cuda_arch_list} RUN --mount=type=cache,target=/root/.cache/uv \ --mount=type=bind,source=requirements/kv_connectors.txt,target=/tmp/kv_connectors.txt,ro \ diff --git a/docker/export_wheels.sh b/docker/export_wheels.sh index 0cc5ccc2b0..c9ad1fa350 100755 --- a/docker/export_wheels.sh +++ b/docker/export_wheels.sh @@ -20,7 +20,7 @@ set -e CUDA_VERSION="${CUDA_VERSION:-13.0.2}" TARGETPLATFORM="${TARGETPLATFORM:-linux/amd64}" PYTHON_VERSION="${PYTHON_VERSION:-3.12}" -TORCH_CUDA_ARCH_LIST="${TORCH_CUDA_ARCH_LIST:-7.0 7.5 8.0 8.9 9.0 10.0 12.0}" +TORCH_CUDA_ARCH_LIST="${TORCH_CUDA_ARCH_LIST:-7.0 7.5 8.0 8.9 9.0 10.0 11.0 12.0}" MAX_JOBS="${MAX_JOBS:-2}" NVCC_THREADS="${NVCC_THREADS:-8}" APHRODITE_VERSION_OVERRIDE="${APHRODITE_VERSION_OVERRIDE:-}" diff --git a/docker/versions.json b/docker/versions.json index f690365b4e..da76f8a5a4 100644 --- a/docker/versions.json +++ b/docker/versions.json @@ -41,7 +41,7 @@ "default": "0" }, "TORCH_CUDA_ARCH_LIST": { - "default": "7.5 8.0 8.6 8.9 9.0 10.0 12.0+PTX" + "default": "7.5 8.0 8.6 8.9 9.0 10.0 11.0 12.0+PTX" }, "DEEPEP_COMMIT_HASH": { "default": "73b6ea4" diff --git a/docs/src/content/docs/usage/models.md b/docs/src/content/docs/usage/models.md index e11ca7053b..b66403fe96 100644 --- a/docs/src/content/docs/usage/models.md +++ b/docs/src/content/docs/usage/models.md @@ -84,6 +84,7 @@ On ROCm platforms, Mistral and Mixtral are capped to 4096 max context length due | `Blip2ForConditionalGeneration` | Image | `Salesforce/blip2-opt-6.7b` | | `ChameleonForConditionalGeneration` | Image | `facebook/chameleon-7b` | | `ChatGLMModel` | Image | `THUDM/chatglm3-6b` | +| `Cosmos3ForConditionalGeneration` | Image, Video | `nvidia/Cosmos3-Nano`, `nvidia/Cosmos3-Super` | | `InternVLChatModel` | Image | `OpenGVLab/InternVL2-8B` | | `LlavaForConditionalGeneration` | Image | `llava-hf/llava-v1.5-7b-hf` | | `LlavaNextForConditionalGeneration` | Image | `llava-hf/llava-v1.6-mistral-7b-hf` | diff --git a/docs/src/content/docs/usage/vlm.md b/docs/src/content/docs/usage/vlm.md index 30afb588ce..944868c538 100644 --- a/docs/src/content/docs/usage/vlm.md +++ b/docs/src/content/docs/usage/vlm.md @@ -125,11 +125,12 @@ export APHRODITE_IMAGE_FETCH_TIMEOUT= ### Video Decoding Backend Aphrodite decodes video bytes into frames using a selectable decoding backend. -Three FFmpeg-backed software backends are supported: +The following decoding backends are supported: - `opencv` (default): OpenCV-based decoder. - `pyav`: PyAV decoder. - `torchcodec`: TorchCodec (PyTorch-native) decoder. +- `deepstream`: NVIDIA DeepStream (NVDEC) GPU decoder. `torchcodec` lets you choose which FFmpeg version is used, while `opencv` and `pyav` rely on whichever FFmpeg build they were linked against. @@ -154,6 +155,55 @@ aphrodite run Qwen/Qwen3-VL-30B-A3B-Instruct \ --media-io-kwargs '{"video": {"backend": "torchcodec", "seek_mode": "approximate", "num_ffmpeg_threads": 4}}' ``` +#### GPU Video Decoding with DeepStream (NVDEC) + +By default Aphrodite decodes video on the CPU. On NVIDIA GPUs you can instead +decode directly on the hardware video engine (NVDEC) with the DeepStream +backend, which keeps decoding off the CPU and can significantly increase video +throughput. + +Install the backend (Linux x86-64 only): + +```bash +pip install aphrodite-engine[deepstream] +``` + +The pip wheel bundles the DeepStream libraries but still relies on a few system +packages that pip cannot install. On Ubuntu: + +```bash +apt-get install -y \ + gstreamer1.0-tools gstreamer1.0-plugins-base gstreamer1.0-plugins-good \ + gstreamer1.0-plugins-bad gstreamer1.0-libav \ + python3-gi python3-gst-1.0 libv4l-0 cuda-libraries-13-0 +``` + +Select the backend either with an environment variable: + +```bash +export APHRODITE_VIDEO_LOADER_BACKEND=deepstream +aphrodite run Qwen/Qwen3-VL-30B-A3B-Instruct +``` + +or per request via `--media-io-kwargs`: + +```bash +aphrodite run Qwen/Qwen3-VL-30B-A3B-Instruct \ + --media-io-kwargs '{"video": {"backend": "deepstream"}}' +``` + +DeepStream-specific parameters: + +- `pool_size`: Number of GPU decode workers in the process-wide decode pool + (clamped to `[1, 16]`). When unset it defaults to + `APHRODITE_MEDIA_LOADING_THREAD_COUNT` (default `8`). The pool is a + singleton, so the first request's value wins. + +```bash +aphrodite run Qwen/Qwen3-VL-30B-A3B-Instruct \ + --media-io-kwargs '{"video": {"backend": "deepstream", "pool_size": 12}}' +``` + :::tip There is no need to format the prompt in the API request since it'll be handled by the server. ::: diff --git a/requirements/test/cpu.txt b/requirements/test/cpu.txt index a35e821a5a..1b7ce2f10f 100644 --- a/requirements/test/cpu.txt +++ b/requirements/test/cpu.txt @@ -962,7 +962,7 @@ s3transfer==0.10.3 # via boto3 sacrebleu==2.4.3 # via lm-eval -safetensors==0.7.0 +safetensors==0.8.0 # via # -r requirements/test/../common.txt # accelerate @@ -1159,7 +1159,7 @@ tqdm==4.67.3 # segmentation-models-pytorch # sentence-transformers # transformers -transformers==5.10.4 +transformers==5.13.1 # via # -r requirements/test/../common.txt # -r requirements/test/cuda.in diff --git a/requirements/test/cuda.in b/requirements/test/cuda.in index 4cd1bd96b3..a287f5d329 100644 --- a/requirements/test/cuda.in +++ b/requirements/test/cuda.in @@ -39,7 +39,7 @@ open_clip_torch==2.32.0 # Required for nemotron_vl test, Nemotron Parse in test_ datamodel_code_generator # required for minicpm3 test lm-eval[api]>=0.4.12 # required for model evaluation test mteb[bm25s]>=2, <3 # required for mteb test -transformers==5.10.4 +transformers==5.13.1 tokenizers==0.22.2 schemathesis>=4.0.0 # Required for openai schema test. # quantization diff --git a/requirements/test/cuda.txt b/requirements/test/cuda.txt index 97d6cdebf8..98c2aee87e 100644 --- a/requirements/test/cuda.txt +++ b/requirements/test/cuda.txt @@ -1055,7 +1055,7 @@ s3transfer==0.10.3 # via boto3 sacrebleu==2.4.3 # via lm-eval -safetensors==0.7.0 +safetensors==0.8.0 # via # -c requirements/common.txt # -r requirements/test/../common.txt @@ -1263,7 +1263,7 @@ tqdm==4.67.3 # segmentation-models-pytorch # sentence-transformers # transformers -transformers==5.10.4 +transformers==5.13.1 # via # -c requirements/common.txt # -r requirements/test/../common.txt diff --git a/requirements/test/nightly-torch.txt b/requirements/test/nightly-torch.txt index 826473db1b..dfefa8239c 100644 --- a/requirements/test/nightly-torch.txt +++ b/requirements/test/nightly-torch.txt @@ -29,7 +29,7 @@ opencv-python-headless >= 4.13.0 # required for video test datamodel_code_generator # required for minicpm3 test lm-eval[api]>=0.4.12 # required for model evaluation test mteb[bm25s]>=2, <3 # required for mteb test -transformers==5.10.4 +transformers==5.13.1 tokenizers==0.22.2 schemathesis>=4.0.0 # Required for openai schema test. # quantization diff --git a/requirements/test/rocm.in b/requirements/test/rocm.in index b1c9a473e2..c5f8f85f2f 100644 --- a/requirements/test/rocm.in +++ b/requirements/test/rocm.in @@ -35,7 +35,7 @@ open_clip_torch==2.32.0 # Required for nemotron_vl test, Nemotron Parse in test_ datamodel_code_generator # required for minicpm3 test lm-eval[api]>=0.4.12 # required for model evaluation test mteb[bm25s]>=2, <3 # required for mteb test -transformers==5.10.4 +transformers==5.13.1 tokenizers==0.22.2 schemathesis>=4.0.0 # Required for openai schema test # quantization diff --git a/requirements/test/rocm.txt b/requirements/test/rocm.txt index 52a0ad88c9..1b4484c3a2 100644 --- a/requirements/test/rocm.txt +++ b/requirements/test/rocm.txt @@ -1035,7 +1035,7 @@ s3transfer==0.16.0 # via boto3 sacrebleu==2.6.0 # via lm-eval -safetensors==0.7.0 +safetensors==0.8.0 # via # -c requirements/common.txt # -r requirements/test/../common.txt @@ -1218,7 +1218,7 @@ tqdm==4.67.3 # sentence-transformers # tilelang # transformers -transformers==5.10.4 +transformers==5.13.1 # via # -c requirements/common.txt # -r requirements/test/../common.txt diff --git a/requirements/test/xpu.in b/requirements/test/xpu.in index 98ebda7de4..56d1374e06 100644 --- a/requirements/test/xpu.in +++ b/requirements/test/xpu.in @@ -17,7 +17,7 @@ accelerate arctic-inference lm_eval[api]>=0.4.12 modelscope<1.38 -transformers==5.10.4 +transformers==5.13.1 # --- Audio Processing --- librosa diff --git a/requirements/test/xpu.txt b/requirements/test/xpu.txt index 6335fc90cf..cd74568411 100644 --- a/requirements/test/xpu.txt +++ b/requirements/test/xpu.txt @@ -779,7 +779,7 @@ rpds-py==0.30.0 # referencing sacrebleu==2.6.0 # via lm-eval -safetensors==0.7.0 +safetensors==0.8.0 # via # -c requirements/common.txt # -r requirements/test/../common.txt @@ -940,7 +940,7 @@ tqdm==4.67.3 # pqdm # sentence-transformers # transformers -transformers==5.10.4 +transformers==5.13.1 # via # -c requirements/common.txt # -r requirements/test/../common.txt diff --git a/requirements/xpu.txt b/requirements/xpu.txt index d595c905d0..4490a38da3 100644 --- a/requirements/xpu.txt +++ b/requirements/xpu.txt @@ -18,4 +18,4 @@ torchvision torchcodec >= 0.14 # Required for the torchcodec video decoding backend auto_round_lib>=0.14.0 -aphrodite_xpu_kernels @ https://github.com/vllm-project/vllm-xpu-kernels/releases/download/v0.1.10.1/vllm_xpu_kernels-0.1.10.1-cp38-abi3-manylinux_2_28_x86_64.whl +aphrodite_xpu_kernels @ https://github.com/vllm-project/vllm-xpu-kernels/releases/download/v0.1.11/vllm_xpu_kernels-0.1.11-cp38-abi3-manylinux_2_28_x86_64.whl diff --git a/rust/Cargo.lock b/rust/Cargo.lock index d950956d82..b5ea0b8c6f 100644 --- a/rust/Cargo.lock +++ b/rust/Cargo.lock @@ -2483,7 +2483,7 @@ checksum = "11d3d7f243d5c5a8b9bb5d6dd2b1602c0cb0b9db1621bafc7ed66e35ff9fe092" [[package]] name = "llm-multimodal" version = "1.7.1" -source = "git+https://github.com/smg-project/llm-multimodal?rev=7d74582aeaf0e4086a44964382655d22f1af0686#7d74582aeaf0e4086a44964382655d22f1af0686" +source = "git+https://github.com/smg-project/llm-multimodal?rev=c8a29dcc755139fdc26185f400ea48c6d6d48273#c8a29dcc755139fdc26185f400ea48c6d6d48273" dependencies = [ "anyhow", "base64 0.22.1", @@ -2774,8 +2774,8 @@ dependencies = [ "portable-atomic", "portable-atomic-util", "rawpointer", + "serde", ] - [[package]] name = "ndarray" version = "0.17.2" diff --git a/rust/Cargo.toml b/rust/Cargo.toml index 7868165792..08352acefb 100644 --- a/rust/Cargo.toml +++ b/rust/Cargo.toml @@ -53,12 +53,12 @@ hyper-util = { version = "0.1.20", features = [ indexmap = "2.13.0" itertools = "0.14.0" libc = "0.2.177" -llm-multimodal = { git = "https://github.com/smg-project/llm-multimodal", rev = "7d74582aeaf0e4086a44964382655d22f1af0686" } +llm-multimodal = { git = "https://github.com/smg-project/llm-multimodal", rev = "c8a29dcc755139fdc26185f400ea48c6d6d48273" } mimalloc = "0.1.52" minijinja = { version = "2.0", features = ["unstable_machinery", "json", "builtins", "loader", "loop_controls", "preserve_order"] } minijinja-contrib = { version = "2.0", features = ["pycompat"] } native-tls-vendored = { package = "native-tls", version = "0.2.18", features = ["vendored"] } -ndarray = { version = "0.16.1", features = ["serde"] } +ndarray = { version = "0.17", features = ["serde"] } openai-harmony = { package = "oss-harmony", git = "https://github.com/oss-harmony/harmony", tag = "v0.0.11", default-features = false } openai-protocol = "1.6.0" openssl = "0.10" diff --git a/rust/src/chat/Cargo.toml b/rust/src/chat/Cargo.toml index fcca07346b..0425042dba 100644 --- a/rust/src/chat/Cargo.toml +++ b/rust/src/chat/Cargo.toml @@ -42,6 +42,7 @@ anyhow.workspace = true bytes.workspace = true clap.workspace = true expect-test.workspace = true +ndarray.workspace = true paste.workspace = true rmp-serde.workspace = true serial_test.workspace = true diff --git a/rust/src/chat/src/backend/hf.rs b/rust/src/chat/src/backend/hf.rs index 1c3aebf111..b94a446f70 100644 --- a/rust/src/chat/src/backend/hf.rs +++ b/rust/src/chat/src/backend/hf.rs @@ -10,7 +10,7 @@ use crate::backend::{ NewChatOutputProcessorOptions, }; use crate::error::Result; -use crate::multimodal::MultimodalModelInfo; +use crate::multimodal::{MultimodalConfigFiles, MultimodalModelInfo}; use crate::output::{ DefaultChatOutputProcessor, HarmonyChatOutputProcessor, validate_harmony_parser_overrides, }; @@ -46,8 +46,12 @@ impl HfChatBackend { MultimodalModelInfo::from_paths( model_id.clone(), (!model_type.is_empty()).then_some(model_type.to_string()), - files.config_path.as_deref(), - files.preprocessor_config_path.as_deref(), + MultimodalConfigFiles { + config: files.config_path.as_deref(), + preprocessor_config: files.preprocessor_config_path.as_deref(), + video_preprocessor_config: files.video_preprocessor_config_path.as_deref(), + processor_config: files.processor_config_path.as_deref(), + }, tokenizer.clone(), )? }; @@ -139,8 +143,11 @@ pub(super) async fn load_model_backends( fn resolve_multimodal_render_info( info: Option<&MultimodalModelInfo>, ) -> Option { + use llm_multimodal::Modality; + info.map(|info| MultimodalRenderInfo { - placeholder_token: info.placeholder_token().to_string(), + image_token: info.placeholder_token(Modality::Image).map(str::to_string), + video_token: info.placeholder_token(Modality::Video).map(str::to_string), }) } @@ -192,6 +199,8 @@ mod tests { tokenizer_config_path: Some(tokenizer_config_path), generation_config_path: None, preprocessor_config_path: None, + video_preprocessor_config_path: None, + processor_config_path: None, chat_template_path: None, config_path: Some(config_path), } diff --git a/rust/src/chat/src/error.rs b/rust/src/chat/src/error.rs index 3ab1fa47a5..6e2662b1b7 100644 --- a/rust/src/chat/src/error.rs +++ b/rust/src/chat/src/error.rs @@ -1,5 +1,5 @@ use thiserror::Error; -use thiserror_ext::Macro; +use thiserror_ext::{AsReport as _, Macro}; type BoxedError = Box; @@ -18,6 +18,8 @@ pub enum Error { UnsupportedMultimodalRenderer, #[error("unsupported multimodal content: {0}")] UnsupportedMultimodalContent(&'static str), + #[error("`{modality}` input is not supported by this model")] + UnsupportedModality { modality: String }, #[error("multimodal preprocessing error: {0}")] Multimodal(#[message] String), #[error("{kind} parsing is not available for model `{model_id}`")] @@ -80,11 +82,39 @@ impl Error { match self { Self::PromptTooLong { .. } => true, Self::Text(error) => error.is_request_validation_error(), + Self::UnsupportedMultimodalRenderer + | Self::UnsupportedMultimodalContent(_) + | Self::UnsupportedModality { .. } => true, + _ => false, } } } +impl From for Error { + fn from(error: llm_multimodal::MediaConnectorError) -> Self { + Self::Multimodal(error.to_report_string()) + } +} + +impl From for Error { + fn from(error: llm_multimodal::MultiModalError) -> Self { + Self::Multimodal(error.to_report_string()) + } +} + +impl From for Error { + fn from(error: llm_multimodal::TransformError) -> Self { + Self::Multimodal(error.to_report_string()) + } +} + +impl From for Error { + fn from(error: llm_multimodal::registry::ModelRegistryError) -> Self { + Self::Multimodal(error.to_report_string()) + } +} + /// Format the available-parser suffix used in user-facing error messages. fn available_parser_hint(available_names: &[String]) -> String { if available_names.is_empty() { diff --git a/rust/src/chat/src/multimodal.rs b/rust/src/chat/src/multimodal.rs index 5e25f45d68..70014099c0 100644 --- a/rust/src/chat/src/multimodal.rs +++ b/rust/src/chat/src/multimodal.rs @@ -1,8 +1,8 @@ -//! Chat-layer multimodal image preparation. +//! Chat-layer multimodal media preparation. //! -//! This module owns the narrow image-only multimodal path for chat requests: -//! it extracts image parts from structured chat messages, fetches and -//! preprocesses them through `llm-multimodal`, expands rendered prompt +//! This module owns the multimodal path for chat requests: it extracts media +//! parts from structured chat messages, fetches and preprocesses them through +//! `llm-multimodal` one modality at a time, expands rendered prompt //! placeholders after tokenization, and builds the engine-facing //! `MmFeatures` payload. //! @@ -16,19 +16,15 @@ use std::sync::{Arc, LazyLock}; use itertools::izip; use llm_multimodal::{ - AsyncMultiModalTracker, FieldLayout, MediaConnector, MediaConnectorConfig, MediaContentPart, - Modality, ModelMetadata, ModelProcessorSpec, ModelRegistry, PreProcessorConfig, - PreprocessedEncoderInputs as PreprocessedImages, PromptReplacement, Tokenizer as TokenResolver, - TrackedMedia, VisionPreProcessor as ImagePreProcessor, - VisionProcessorRegistry as ImageProcessorRegistry, + AsyncMultiModalTracker, FieldLayout, ImageFrame, MediaConnector, MediaConnectorConfig, + MediaContentPart, Modality, ModelMetadata, ModelProcessorSpec, ModelRegistry, + PreProcessorConfig, PreprocessedEncoderInputs, PromptReplacement, Tokenizer as TokenResolver, + TrackedMedia, VideoClip, VisionPreProcessor, VisionProcessorRegistry, }; +use thiserror_ext::AsReport as _; use tracing::warn; use aphrodite_engine_core_client::protocol::dtype::ModelDtype; -use aphrodite_engine_core_client::protocol::multimodal::{ - MmBatchedField, MmFeatureSpec, MmFeatures, MmField, MmFieldElem, MmFlatField, MmKwargsItem, - MmSharedField, MmSlice, PlaceholderRange, SliceSpec, -}; -use aphrodite_engine_core_client::protocol::tensor::WireTensor; +use aphrodite_engine_core_client::protocol::multimodal::{MmFeatureSpec, MmFeatures, MmKwargsItem}; use aphrodite_text::Prompt; use aphrodite_text::tokenizer::{DynTokenizer, Tokenizer}; @@ -36,14 +32,20 @@ use crate::error::{Error, Result, bail_multimodal, multimodal}; use crate::renderer::RenderedPrompt; use crate::request::{ChatContent, ChatContentPart, ChatMessage, ChatRequest}; +mod expand; +mod image; mod tensor; +mod video; + +use self::expand::expand_prompt_token_ids; /// Resolved multimodal support for one loaded model. #[derive(Clone)] pub struct MultimodalModelInfo { context: MultimodalModelContext, spec: ResolvedMultimodalSpec, - image_processor: ResolvedImageProcessor, + image: Option, + video: Option, media_connector: Arc, } @@ -75,91 +77,171 @@ impl MultimodalModelContext { REGISTRY.lookup(&self.metadata()) } - /// Resolve a static image preprocessor for one loaded model. - fn resolve_image_processor(&self) -> Option<&'static dyn ImagePreProcessor> { - static REGISTRY: LazyLock = - LazyLock::new(ImageProcessorRegistry::with_defaults); + /// Resolve a static vision preprocessor for one loaded model. + /// + /// The vision preprocessor serves both the image and video modalities. + fn resolve_vision_processor(&self) -> Option<&'static dyn VisionPreProcessor> { + static REGISTRY: LazyLock = + LazyLock::new(VisionProcessorRegistry::with_defaults); REGISTRY.find(&self.model_id, self.model_type.as_deref()) } } -/// Static model-specific prompt and tensor-layout behavior. +/// Static model-specific tensor-layout behavior shared across modalities. #[derive(Clone)] struct ResolvedMultimodalSpec { raw: &'static dyn ModelProcessorSpec, - placeholder_token: String, - placeholder_marker_token_id: u32, - placeholder_embed_token_id: u32, field_layouts: HashMap, keep_on_cpu_keys: HashSet, } impl ResolvedMultimodalSpec { - fn new(raw: &'static dyn ModelProcessorSpec, context: &MultimodalModelContext) -> Result { - let metadata = context.metadata(); - let placeholder_token = - raw.placeholder_token(&metadata).map_err(|error| multimodal!("{error}"))?; - // This is the rendered prompt marker, so resolve it from the token - // string itself. Do not use `ModelProcessorSpec::placeholder_token_id()`: - // for some specs that ID is the replacement vision/patch token, - // not necessarily the token ID of `placeholder_token`. - let placeholder_marker_token_id = - context.tokenizer().token_to_id(&placeholder_token).ok_or_else(|| { - multimodal!( - "placeholder token `{placeholder_token}` is not in the tokenizer vocabulary" - ) - })?; - let placeholder_embed_token_id = - raw.placeholder_token_id(&metadata).map_err(|error| multimodal!("{error}"))? as u32; - - Ok(Self { + fn new(raw: &'static dyn ModelProcessorSpec) -> Self { + Self { raw, - placeholder_token, - placeholder_marker_token_id, - placeholder_embed_token_id, field_layouts: raw.field_layouts(), keep_on_cpu_keys: raw.keep_on_cpu_keys().into_iter().collect(), - }) + } } - fn prompt_replacements( + fn prompt_replacements_for( &self, context: &MultimodalModelContext, - preprocessed: &PreprocessedImages, + preprocessed: &PreprocessedEncoderInputs, + modality: Modality, ) -> Result> { - self.raw - .prompt_replacements(&context.metadata(), preprocessed) - .map_err(|error| multimodal!("{error}")) + Ok(self.raw.prompt_replacements_for(&context.metadata(), preprocessed, modality)?) } } -/// Static image preprocessor plus its loaded config. +/// Resolved placeholder tokens for one modality. #[derive(Clone)] -struct ResolvedImageProcessor { - raw: &'static dyn ImagePreProcessor, +struct ResolvedPlaceholder { + token: String, + /// The token ID emitted for `token` in the rendered prompt. + marker_token_id: u32, + /// The model-declared embed token ID marked in `is_embed` masks. + embed_token_id: u32, +} + +impl ResolvedPlaceholder { + fn resolve( + raw: &'static dyn ModelProcessorSpec, + context: &MultimodalModelContext, + modality: Modality, + ) -> Result { + let metadata = context.metadata(); + let token = raw.placeholder_token_for(&metadata, modality)?; + // This is the rendered prompt marker, so resolve it from the token + // string itself. Do not use `ModelProcessorSpec::placeholder_token_id_for()`: + // for some specs that ID is the replacement vision/patch token, + // not necessarily the token ID of the placeholder token. + let marker_token_id = context.tokenizer().token_to_id(&token).ok_or_else(|| { + multimodal!("placeholder token `{token}` is not in the tokenizer vocabulary") + })?; + let embed_token_id = raw.placeholder_token_id_for(&metadata, modality)? as u32; + + Ok(Self { + token, + marker_token_id, + embed_token_id, + }) + } +} + +/// Static per-modality vision preprocessor plus its loaded config and +/// resolved placeholder tokens. +#[derive(Clone)] +struct ModalitySupport { + placeholder: ResolvedPlaceholder, + processor: &'static dyn VisionPreProcessor, config: PreProcessorConfig, } -/// Request-scoped fetched media, kept together with tracker UUID metadata. -struct FetchedImageMedia { - frames: Vec>, - uuids: Vec>, +/// Model-repo config file locations consumed by multimodal support. +#[derive(Debug, Default, Clone, Copy)] +pub struct MultimodalConfigFiles<'a> { + pub config: Option<&'a Path>, + pub preprocessor_config: Option<&'a Path>, + /// Video-specific preprocessor config (`video_preprocessor_config.json`). + pub video_preprocessor_config: Option<&'a Path>, + /// Combined processor config (`processor_config.json`), whose modality + /// sections are fallback preprocessor config sources. + pub processor_config: Option<&'a Path>, +} + +/// Load a modality's dedicated preprocessor config, falling back to its section +/// in the combined processor config. +fn load_preprocessor_config( + dedicated_path: Option<&Path>, + dedicated_name: &str, + processor_config_path: Option<&Path>, + processor_section: &str, +) -> Result> { + if let Some(path) = dedicated_path { + let text = fs::read_to_string(path) + .map_err(|error| multimodal!("failed to read {dedicated_name}: {error}"))?; + let config = PreProcessorConfig::from_json(&text) + .map_err(|error| multimodal!("failed to parse {dedicated_name}: {error}"))?; + return Ok(Some(config)); + } + + let Some(path) = processor_config_path else { + return Ok(None); + }; + let text = fs::read_to_string(path) + .map_err(|error| multimodal!("failed to read processor_config.json: {error}"))?; + let value: serde_json::Value = serde_json::from_str(&text) + .map_err(|error| multimodal!("failed to parse processor_config.json: {error}"))?; + let Some(processor) = value.get(processor_section) else { + return Ok(None); + }; + let config = PreProcessorConfig::from_value(processor.clone()).map_err(|error| { + multimodal!("failed to parse {processor_section} from processor_config.json: {error}") + })?; + Ok(Some(config)) +} + +/// Request-scoped fetched media, split per modality with tracker UUID +/// metadata preserved in request order. +struct FetchedMedia { + images: Vec>, + image_uuids: Vec>, + videos: Vec>, + video_uuids: Vec>, +} + +/// One modality's preprocessed output, ready for the shared expansion and +/// feature-assembly tail. +struct PreparedMedia { + modality: Modality, + placeholder: ResolvedPlaceholder, + /// One replacement per media item, in request order. + replacements: Vec, + /// One entry per media item, aligned with `replacements`. + items: Vec, +} + +/// One media item's complete engine kwargs plus identity metadata. +struct PreparedItem { + data: MmKwargsItem, + hash: String, + uuid: Option, } impl MultimodalModelInfo { /// Load and resolve multimodal support from model files. /// - /// Returns `Ok(Some(_))` only when both the model spec and image processor - /// are registered. File read/parse failures are real errors; unsupported - /// model families are logged and returned as `Ok(None)`. + /// Returns `Ok(Some(_))` only when the model spec is registered and at + /// least one modality resolves. File read/parse failures are real errors; + /// unsupported model families are logged and returned as `Ok(None)`. pub fn from_paths( model_id: String, model_type: Option, - config_path: Option<&Path>, - preprocessor_config_path: Option<&Path>, + files: MultimodalConfigFiles<'_>, tokenizer: DynTokenizer, ) -> Result> { - let config = match config_path { + let config = match files.config { Some(path) => { let text = fs::read_to_string(path) .map_err(|error| multimodal!("failed to read config.json: {error}"))?; @@ -168,17 +250,20 @@ impl MultimodalModelInfo { } None => serde_json::Value::Object(Default::default()), }; - let preprocessor_config = match preprocessor_config_path { - Some(path) => { - let text = fs::read_to_string(path).map_err(|error| { - multimodal!("failed to read preprocessor_config.json: {error}") - })?; - PreProcessorConfig::from_json(&text).map_err(|error| { - multimodal!("failed to parse preprocessor_config.json: {error}") - })? - } - None => PreProcessorConfig::default(), - }; + let image_preprocessor_config = load_preprocessor_config( + files.preprocessor_config, + "preprocessor_config.json", + files.processor_config, + "image_processor", + )? + .unwrap_or_default(); + let video_preprocessor_config = load_preprocessor_config( + files.video_preprocessor_config, + "video_preprocessor_config.json", + files.processor_config, + "video_processor", + )? + .unwrap_or_else(|| image_preprocessor_config.clone()); let context = MultimodalModelContext { model_id, @@ -187,7 +272,21 @@ impl MultimodalModelInfo { tokenizer: TokenizerResolver(tokenizer), }; - let Some(spec) = context.resolve_model_spec() else { + Self::from_loaded( + context, + image_preprocessor_config, + video_preprocessor_config, + ) + } + + /// Resolve multimodal support from an assembled context and parsed + /// preprocessor configs. + fn from_loaded( + context: MultimodalModelContext, + image_preprocessor_config: PreProcessorConfig, + video_preprocessor_config: PreProcessorConfig, + ) -> Result> { + let Some(raw_spec) = context.resolve_model_spec() else { warn!( model_id = context.model_id, model_type = context.model_type, @@ -195,47 +294,99 @@ impl MultimodalModelInfo { ); return Ok(None); }; - let spec = ResolvedMultimodalSpec::new(spec, &context)?; - let Some(image_processor) = context.resolve_image_processor() else { + let Some(processor) = context.resolve_vision_processor() else { warn!( model_id = context.model_id, model_type = context.model_type, - "image processor is not registered; disabling multimodal support for this model" + "vision processor is not registered; disabling multimodal support for this model" ); return Ok(None); }; - let media_connector = Arc::new( - MediaConnector::new(reqwest::Client::new(), MediaConnectorConfig::default()) - .map_err(|error| multimodal!("{error}"))?, - ); + // Warn and disable the modality if the placeholder resolution fails. + let resolve_placeholder = + |modality: Modality| match ResolvedPlaceholder::resolve(raw_spec, &context, modality) { + Ok(placeholder) => Some(placeholder), + Err(error) => { + warn!( + model_id = context.model_id, + %modality, + error = %error.as_report(), + "placeholder tokens did not resolve; disabling this modality for this model" + ); + None + } + }; + + let image = resolve_placeholder(Modality::Image).map(|placeholder| ModalitySupport { + placeholder, + processor, + config: image_preprocessor_config, + }); + + let video = resolve_placeholder(Modality::Video).and_then(|placeholder| { + // Placeholder expansion attributes markers to modalities by token + // ID, so a marker shared with the image modality is ambiguous. + let image_marker = image.as_ref().map(|image| image.placeholder.marker_token_id); + if image_marker == Some(placeholder.marker_token_id) { + warn!( + model_id = context.model_id, + token = placeholder.token, + "video placeholder token collides with the image placeholder; disabling video support for this model" + ); + None + } else { + Some(ModalitySupport { + placeholder, + processor, + config: video_preprocessor_config, + }) + } + }); + + if image.is_none() && video.is_none() { + warn!( + model_id = context.model_id, + model_type = context.model_type, + "no multimodal modality resolved; disabling multimodal support for this model" + ); + return Ok(None); + } + + let media_connector = Arc::new(MediaConnector::new( + reqwest::Client::new(), + MediaConnectorConfig::default(), + )?); Ok(Some(Self { context, - spec, - image_processor: ResolvedImageProcessor { - raw: image_processor, - config: preprocessor_config, - }, + spec: ResolvedMultimodalSpec::new(raw_spec), + image, + video, media_connector, })) } - /// Return the template-visible placeholder token for this model. + /// Return the template-visible placeholder token for one modality, when + /// this model supports it. /// - /// The HF renderer uses this token while flattening image content in string - /// content format. - pub fn placeholder_token(&self) -> &str { - &self.spec.placeholder_token + /// The HF renderer uses these tokens while flattening media content in + /// string content format. + pub fn placeholder_token(&self, modality: Modality) -> Option<&str> { + match modality { + Modality::Image => self.image.as_ref()?.placeholder.token.as_str().into(), + Modality::Video => self.video.as_ref()?.placeholder.token.as_str().into(), + _ => None, + } } } /// Finalize a rendered chat prompt into text-generation input. /// /// Text-only requests pass through unchanged as `Prompt::Text`. Multimodal -/// requests are tokenized in chat, their image placeholders are expanded, and -/// preprocessed image features are attached for engine-core transport. +/// requests are tokenized in chat, their media placeholders are expanded, and +/// preprocessed media features are attached for engine-core transport. pub(crate) async fn finalize_rendered_prompt( request: &ChatRequest, rendered: RenderedPrompt, @@ -260,7 +411,7 @@ pub(crate) async fn finalize_rendered_prompt( Ok((Prompt::TokenIds(prompt_token_ids), Some(prepared))) } -/// Extract image media parts from chat messages in message/content order. +/// Extract media parts from chat messages in message/content order. /// /// Assistant history is skipped because generated assistant blocks are already /// represented as text for prompt rendering in this crate. @@ -289,6 +440,12 @@ fn extract_media_parts(request: &ChatRequest) -> Result> { detail: *detail, uuid: uuid.clone(), }), + ChatContentPart::VideoUrl { video_url, uuid } => { + all_parts.push(MediaContentPart::VideoUrl { + url: video_url.clone(), + uuid: uuid.clone(), + }) + } } } } @@ -296,8 +453,8 @@ fn extract_media_parts(request: &ChatRequest) -> Result> { } impl MultimodalModelInfo { - /// Run media fetch, image preprocessing, prompt expansion, and feature - /// build. + /// Run media fetch, per-modality preprocessing, prompt expansion, and + /// feature build. /// /// `prompt_token_ids` is mutated in place because placeholder expansion /// changes both the final prompt and the offsets recorded in @@ -313,12 +470,47 @@ impl MultimodalModelInfo { } let media_parts_len = media_parts.len(); - let fetched = self.fetch_images(media_parts).await?; - let preprocessed = self.preprocess_images(&fetched.frames).await?; - let replacements = self.spec.prompt_replacements(&self.context, &preprocessed)?; - let ranges = self.expand_prompt_tokens(prompt_token_ids, replacements)?; + // TODO: enforce per-modality item-count limits, aligned with the + // engine's `--limit-mm-per-prompt` semantics. + let fetched = self.fetch_media(media_parts).await?; + + let mut prepared = Vec::new(); + if !fetched.images.is_empty() { + prepared + .push(self.prepare_images(fetched.images, fetched.image_uuids, model_dtype).await?); + } + if !fetched.videos.is_empty() { + prepared + .push(self.prepare_videos(fetched.videos, fetched.video_uuids, model_dtype).await?); + } + + let mut ranges = expand_prompt_token_ids(prompt_token_ids, &prepared)?; + + let mut features = Vec::with_capacity(media_parts_len); + for media in prepared { + let media_ranges = ranges.remove(&media.modality).unwrap_or_default(); + if media_ranges.len() != media.items.len() { + bail_multimodal!( + "number of expanded `{}` placeholders {} does not match number of media items {}", + media.modality, + media_ranges.len(), + media.items.len() + ); + } + for (item, range) in izip!(media.items, media_ranges) { + features.push(MmFeatureSpec { + data: Some(item.data), + modality: media.modality.to_string(), + identifier: item.uuid.unwrap_or_else(|| item.hash.clone()), + mm_position: range, + mm_hash: Some(item.hash), + }); + } + } + // Mirror the Python frontend (`argsort_mm_positions`): features are + // ordered by their placeholder position in the prompt. + features.sort_by_key(|feature| feature.mm_position.offset); - let features = self.build_features(preprocessed, fetched, ranges, model_dtype)?; if features.len() != media_parts_len { bail_multimodal!( "number of built multimodal features {} does not match number of media parts {}", @@ -329,219 +521,51 @@ impl MultimodalModelInfo { Ok(features) } - /// Fetch all image parts and preserve their request-order UUID metadata. - async fn fetch_images(&self, media_parts: Vec) -> Result { + /// Fetch all media parts and split them per modality, preserving their + /// request-order UUID metadata. + async fn fetch_media(&self, media_parts: Vec) -> Result { let mut tracker = AsyncMultiModalTracker::new(Arc::clone(&self.media_connector)); for part in media_parts { - tracker.push_part(part).map_err(|error| multimodal!("{error}"))?; + tracker.push_part(part)?; } - let tracker_output = tracker.finalize().await.map_err(|error| multimodal!("{error}"))?; - let images = tracker_output.data.get(&Modality::Image).cloned().unwrap_or_default(); - let uuids = tracker_output.uuids.get(&Modality::Image).cloned().unwrap_or_default(); + let mut tracker_output = tracker.finalize().await?; - let frames = images + let images = tracker_output + .data + .remove(&Modality::Image) + .unwrap_or_default() .into_iter() .map(|media| match media { TrackedMedia::Image(frame) => Ok(frame), - _ => Err(Error::UnsupportedMultimodalContent("non-image")), + _ => Err(multimodal!( + "tracker returned non-image media for the image modality" + )), }) .collect::>>()?; + let image_uuids = tracker_output.uuids.remove(&Modality::Image).unwrap_or_default(); - Ok(FetchedImageMedia { frames, uuids }) - } + let videos = tracker_output + .data + .remove(&Modality::Video) + .unwrap_or_default() + .into_iter() + .map(|media| match media { + TrackedMedia::Video(clip) => Ok(clip), + _ => Err(multimodal!( + "tracker returned non-video media for the video modality" + )), + }) + .collect::>>()?; + let video_uuids = tracker_output.uuids.remove(&Modality::Video).unwrap_or_default(); - /// Preprocess fetched image frames with the model's resolved image - /// processor. - /// - /// The processor work is CPU-heavy relative to request wiring, so it runs - /// in a blocking task and returns owned tensors ready for wire - /// conversion. - async fn preprocess_images( - &self, - image_frames: &[Arc], - ) -> Result { - let config = self.image_processor.config.clone(); - let processor = self.image_processor.raw; - let images = image_frames.iter().map(|frame| frame.data().clone()).collect::>(); - - // TODO: is it still necessary given that we've already in a dedicated runtime? - tokio::task::spawn_blocking(move || { - processor.preprocess(&images, &config).map_err(|error| multimodal!("{error}")) + Ok(FetchedMedia { + images, + image_uuids, + videos, + video_uuids, }) - .await - .map_err(|error| multimodal!("image preprocessing task failed: {error}"))? - } - - /// Replace rendered placeholder markers with model-specific replacement - /// tokens. - /// - /// Replacements are consumed in order, matching the original media-part - /// order. The returned ranges point into the already-expanded prompt. - fn expand_prompt_tokens( - &self, - prompt_token_ids: &mut Vec, - replacements: Vec, - ) -> Result> { - expand_prompt_token_ids( - prompt_token_ids, - replacements, - self.spec.placeholder_marker_token_id, - self.spec.placeholder_embed_token_id, - &self.spec.placeholder_token, - ) } - - /// Convert preprocessed image tensors into engine-core multimodal features. - /// - /// One `MmFeatureSpec` is produced per image. Tensor fields are - /// sliced according to the model spec's field layout declarations. - fn build_features( - &self, - preprocessed: PreprocessedImages, - images: FetchedImageMedia, - ranges: Vec, - model_dtype: ModelDtype, - ) -> Result { - let len = images.frames.len(); - let tensors = tensor::collect_tensors(preprocessed, model_dtype)?; - - let mut features = Vec::with_capacity(images.frames.len()); - for (index, (frame, uuid, range)) in izip!(images.frames, images.uuids, ranges).enumerate() - { - let mut data = MmKwargsItem::new(); - for (key, tensor) in &tensors { - let keep_on_cpu = self.spec.keep_on_cpu_keys.contains(key); - let (value, field) = match self.spec.field_layouts.get(key) { - Some(FieldLayout::Batched) => ( - tensor.batched_value_at(index)?, - MmField::Batched(MmBatchedField { keep_on_cpu }), - ), - Some(FieldLayout::Flat { sizes_key }) => { - let sizes = tensors.get(sizes_key).ok_or_else(|| { - multimodal!("flat tensor sizes key `{sizes_key}` is missing") - })?; - let (start, end) = tensor::flat_range_for_index(sizes, sizes_key, index)?; - ( - tensor.flat_value_range(start, end)?, - MmField::Flat(MmFlatField { - slices: vec![MmSlice::Slice(SliceSpec { - start: Some(0), - stop: Some((end - start) as isize), - step: None, - })], - dim: 0, - keep_on_cpu, - }), - ) - } - None => ( - tensor.clone(), - MmField::Shared(MmSharedField { - batch_size: len, - keep_on_cpu, - }), - ), - }; - - data.insert( - key.clone(), - MmFieldElem { - data: Some(value.try_into()?), - field, - }, - ); - } - - let hash = frame.hash.clone(); - features.push(MmFeatureSpec { - data: Some(data), - modality: "image".to_string(), - identifier: uuid.unwrap_or_else(|| hash.clone()), - mm_position: range, - mm_hash: Some(hash), - }); - } - - Ok(features) - } -} - -fn expand_prompt_token_ids( - prompt_token_ids: &mut Vec, - replacements: Vec, - placeholder_marker_token_id: u32, - placeholder_embed_token_id: u32, - placeholder_token: &str, -) -> Result> { - if replacements.is_empty() { - return Ok(Vec::new()); - } - - let replacement_growth = replacements.iter().fold(0usize, |total, replacement| { - total.saturating_add(replacement.tokens.len().saturating_sub(1)) - }); - let mut expanded = - Vec::with_capacity(prompt_token_ids.len().saturating_add(replacement_growth)); - let mut ranges = Vec::with_capacity(replacements.len()); - let mut cursor = 0usize; - - for replacement in replacements { - if replacement.modality != Modality::Image { - bail_multimodal!( - "unsupported prompt replacement modality `{}`", - replacement.modality - ); - } - - let offset = find_next_token(prompt_token_ids, placeholder_marker_token_id, cursor) - .ok_or_else(|| { - multimodal!( - "placeholder token `{placeholder_token}` was not found in tokenized prompt" - ) - })?; - - if replacement.tokens.is_empty() { - bail_multimodal!("placeholder token `{placeholder_token}` expanded to no tokens"); - } - - let replacement_len = replacement.tokens.len(); - let is_embed = { - let mask = replacement - .tokens - .iter() - .map(|&token| token as u32 == placeholder_embed_token_id) - .collect::>(); - WireTensor::from_bool(vec![replacement_len], mask).map_err(Error::Multimodal)? - }; - - expanded.extend_from_slice(&prompt_token_ids[cursor..offset]); - let expanded_offset = expanded.len(); - expanded.extend(replacement.tokens.into_iter().map(|token| token as u32)); - ranges.push(PlaceholderRange { - offset: expanded_offset, - length: replacement_len, - is_embed: Some(is_embed), - }); - cursor = offset + 1; - } - - expanded.extend_from_slice(&prompt_token_ids[cursor..]); - *prompt_token_ids = expanded; - - Ok(ranges) -} - -/// Find `needle` in `haystack`, starting at `start`. -/// -/// This is intentionally order-preserving rather than a global replace: each -/// image consumes the next placeholder occurrence. -fn find_next_token(haystack: &[u32], needle: u32, start: usize) -> Option { - haystack - .get(start..)? - .iter() - .position(|token| *token == needle) - .map(|offset| start + offset) } /// Adapter from the frontend tokenizer trait to `llm-multimodal`. @@ -566,18 +590,19 @@ impl TokenResolver for TokenizerResolver { mod tests { use std::sync::Arc; - use llm_multimodal::TokenId; - use aphrodite_engine_core_client::protocol::tensor::WireArrayData; use aphrodite_tokenizer::test_utils::TestTokenizer; use super::*; - const LLAMA4_IMAGE_START_ID: u32 = 200088; - const LLAMA4_IMAGE_END_ID: u32 = 200089; - const LLAMA4_IMAGE_ID: u32 = 200090; - const LLAMA4_PATCH_ID: u32 = 200092; - const LLAMA4_TILE_X_SEPARATOR_ID: u32 = 200093; - const LLAMA4_TILE_Y_SEPARATOR_ID: u32 = 200094; + pub(super) const LLAMA4_IMAGE_START_ID: u32 = 200088; + pub(super) const LLAMA4_IMAGE_END_ID: u32 = 200089; + pub(super) const LLAMA4_IMAGE_ID: u32 = 200090; + pub(super) const LLAMA4_PATCH_ID: u32 = 200092; + pub(super) const LLAMA4_TILE_X_SEPARATOR_ID: u32 = 200093; + pub(super) const LLAMA4_TILE_Y_SEPARATOR_ID: u32 = 200094; + + pub(super) const QWEN3_IMAGE_PAD_ID: u32 = 151655; + pub(super) const QWEN3_VIDEO_PAD_ID: u32 = 151656; fn llama4_tokenizer() -> TestTokenizer { TestTokenizer::new() @@ -589,33 +614,31 @@ mod tests { .with_regular_token("<|tile_y_separator|>", LLAMA4_TILE_Y_SEPARATOR_ID) } - fn test_info(model_type: &str, config: serde_json::Value) -> MultimodalModelInfo { + pub(super) fn qwen3_vl_tokenizer() -> TestTokenizer { + TestTokenizer::new() + .with_regular_token("<|image_pad|>", QWEN3_IMAGE_PAD_ID) + .with_regular_token("<|video_pad|>", QWEN3_VIDEO_PAD_ID) + } + + fn test_info( + model_type: &str, + config: serde_json::Value, + tokenizer: TestTokenizer, + ) -> MultimodalModelInfo { let context = MultimodalModelContext { model_id: format!("{model_type}-test"), model_type: Some(model_type.to_string()), config, - tokenizer: TokenizerResolver(Arc::new(llama4_tokenizer())), + tokenizer: TokenizerResolver(Arc::new(tokenizer)), }; - let spec = context - .resolve_model_spec() - .unwrap_or_else(|| panic!("{model_type} spec should match")); - let spec = ResolvedMultimodalSpec::new(spec, &context).unwrap(); - let raw_image_processor = context - .resolve_image_processor() - .unwrap_or_else(|| panic!("{model_type} image processor should match")); - let media_connector = Arc::new( - MediaConnector::new(reqwest::Client::new(), MediaConnectorConfig::default()).unwrap(), - ); - MultimodalModelInfo { + MultimodalModelInfo::from_loaded( context, - spec, - image_processor: ResolvedImageProcessor { - raw: raw_image_processor, - config: PreProcessorConfig::default(), - }, - media_connector, - } + PreProcessorConfig::default(), + PreProcessorConfig::default(), + ) + .unwrap() + .unwrap_or_else(|| panic!("{model_type} multimodal support should resolve")) } fn llama4_info() -> MultimodalModelInfo { @@ -624,173 +647,96 @@ mod tests { "image_token_index": LLAMA4_PATCH_ID, "vision_config": {"image_size": 336, "patch_size": 14} }); - test_info("llama4", config) + test_info("llama4", config, llama4_tokenizer()) } - fn llama4_single_tile_replacement() -> PromptReplacement { - PromptReplacement::sequence( - Modality::Image, - "<|image|>", - vec![ - LLAMA4_IMAGE_START_ID as TokenId, - LLAMA4_IMAGE_ID as TokenId, - LLAMA4_PATCH_ID as TokenId, - LLAMA4_PATCH_ID as TokenId, - LLAMA4_IMAGE_END_ID as TokenId, - ], - ) + pub(super) fn qwen3_vl_info() -> MultimodalModelInfo { + let config = serde_json::json!({ + "model_type": "qwen3_vl", + "image_token_id": QWEN3_IMAGE_PAD_ID, + "video_token_id": QWEN3_VIDEO_PAD_ID, + "vision_start_token_id": 151652, + "vision_end_token_id": 151653, + "vision_config": {"patch_size": 16} + }); + test_info("qwen3_vl", config, qwen3_vl_tokenizer()) } - fn llama4_multi_tile_replacement() -> PromptReplacement { - PromptReplacement::sequence( - Modality::Image, - "<|image|>", - vec![ - LLAMA4_IMAGE_START_ID as TokenId, - LLAMA4_PATCH_ID as TokenId, - LLAMA4_TILE_X_SEPARATOR_ID as TokenId, - LLAMA4_PATCH_ID as TokenId, - LLAMA4_TILE_Y_SEPARATOR_ID as TokenId, - LLAMA4_IMAGE_ID as TokenId, - LLAMA4_PATCH_ID as TokenId, - LLAMA4_IMAGE_END_ID as TokenId, - ], + #[test] + fn from_paths_resolves_image_config_from_processor_config() { + let dir = tempfile::tempdir().unwrap(); + let config_path = dir.path().join("config.json"); + std::fs::write( + &config_path, + serde_json::json!({ + "model_type": "qwen3_vl", + "image_token_id": QWEN3_IMAGE_PAD_ID, + }) + .to_string(), ) - } + .unwrap(); + let processor_config_path = dir.path().join("processor_config.json"); + std::fs::write( + &processor_config_path, + r#"{"image_processor":{"size":{"shortest_edge":64}}}"#, + ) + .unwrap(); + + let info = MultimodalModelInfo::from_paths( + "qwen3-vl-test".to_string(), + Some("qwen3_vl".to_string()), + MultimodalConfigFiles { + config: Some(&config_path), + processor_config: Some(&processor_config_path), + ..Default::default() + }, + Arc::new(qwen3_vl_tokenizer()), + ) + .unwrap() + .unwrap(); - fn assert_bool_mask(range: &PlaceholderRange, expected: &[bool]) { - let tensor = range.is_embed.as_ref().expect("is_embed mask"); - assert_eq!(tensor.dtype, "bool"); - assert_eq!(tensor.shape, vec![expected.len()]); - assert_eq!( - tensor.data, - WireArrayData::RawView(expected.iter().map(|value| u8::from(*value)).collect()) - ); + assert_eq!(info.image.unwrap().config.get_shortest_edge(), Some(64)); } #[test] - fn expand_prompt_tokens_marks_only_llama4_patch_tokens_as_embed() { - let info = llama4_info(); - let mut prompt_token_ids = vec![1, LLAMA4_IMAGE_ID, 2]; - let replacements = vec![llama4_multi_tile_replacement()]; - - let ranges = info.expand_prompt_tokens(&mut prompt_token_ids, replacements).unwrap(); + fn qwen3_vl_resolves_image_and_video_support() { + let info = qwen3_vl_info(); assert_eq!( - prompt_token_ids, - vec![ - 1, - LLAMA4_IMAGE_START_ID, - LLAMA4_PATCH_ID, - LLAMA4_TILE_X_SEPARATOR_ID, - LLAMA4_PATCH_ID, - LLAMA4_TILE_Y_SEPARATOR_ID, - LLAMA4_IMAGE_ID, - LLAMA4_PATCH_ID, - LLAMA4_IMAGE_END_ID, - 2, - ] + info.placeholder_token(Modality::Image), + Some("<|image_pad|>") ); - assert_eq!(ranges[0].offset, 1); - assert_eq!(ranges[0].length, 8); - assert_bool_mask( - &ranges[0], - &[false, true, false, true, false, false, true, false], + assert_eq!( + info.placeholder_token(Modality::Video), + Some("<|video_pad|>") + ); + assert_ne!( + info.image.as_ref().unwrap().placeholder.marker_token_id, + info.video.as_ref().unwrap().placeholder.marker_token_id, ); } #[test] - fn expand_prompt_tokens_errors_when_placeholder_missing() { - let info = llama4_info(); - let mut prompt_token_ids = vec![1, 2, 3]; - let replacements = vec![llama4_single_tile_replacement()]; - - let error = info.expand_prompt_tokens(&mut prompt_token_ids, replacements).unwrap_err(); - - assert!(matches!(error, Error::Multimodal(message) if message.contains("not found"))); - } - - #[test] - fn expand_prompt_tokens_ignores_empty_replacements() { - let info = llama4_info(); - let mut prompt_token_ids = vec![1, LLAMA4_IMAGE_ID, 2]; - let original_prompt_token_ids = prompt_token_ids.clone(); - - let ranges = info.expand_prompt_tokens(&mut prompt_token_ids, Vec::new()).unwrap(); - - assert!(ranges.is_empty()); - assert_eq!(prompt_token_ids, original_prompt_token_ids); - } - - #[test] - fn expand_prompt_tokens_leaves_prompt_unchanged_when_later_placeholder_missing() { - let info = llama4_info(); - let mut prompt_token_ids = vec![1, LLAMA4_IMAGE_ID, 2]; - let original_prompt_token_ids = prompt_token_ids.clone(); - let replacements = vec![ - llama4_single_tile_replacement(), - llama4_single_tile_replacement(), - ]; - - let error = info.expand_prompt_tokens(&mut prompt_token_ids, replacements).unwrap_err(); - - assert!(matches!(error, Error::Multimodal(message) if message.contains("not found"))); - assert_eq!(prompt_token_ids, original_prompt_token_ids); - } + fn qwen3_vl_without_video_token_id_disables_video_support_only() { + let config = serde_json::json!({ + "model_type": "qwen3_vl", + "image_token_id": QWEN3_IMAGE_PAD_ID, + "vision_config": {"patch_size": 16} + }); + let info = test_info("qwen3_vl", config, qwen3_vl_tokenizer()); - #[test] - fn expand_prompt_tokens_errors_when_replacement_is_empty() { - let info = llama4_info(); - let mut prompt_token_ids = vec![1, LLAMA4_IMAGE_ID, 2]; - let original_prompt_token_ids = prompt_token_ids.clone(); - let replacements = vec![PromptReplacement::sequence( - Modality::Image, - "<|image|>", - Vec::new(), - )]; - - let error = info.expand_prompt_tokens(&mut prompt_token_ids, replacements).unwrap_err(); - - assert!( - matches!(error, Error::Multimodal(message) if message.contains("expanded to no tokens")) + assert_eq!( + info.placeholder_token(Modality::Image), + Some("<|image_pad|>") ); - assert_eq!(prompt_token_ids, original_prompt_token_ids); + assert_eq!(info.placeholder_token(Modality::Video), None); } #[test] - fn expand_prompt_tokens_skips_llama4_image_marker_inside_replacement() { + fn llama4_resolves_image_support_only() { let info = llama4_info(); - let mut prompt_token_ids = vec![1, LLAMA4_IMAGE_ID, 2, LLAMA4_IMAGE_ID, 3]; - let replacements = vec![ - llama4_single_tile_replacement(), - llama4_single_tile_replacement(), - ]; - - let ranges = info.expand_prompt_tokens(&mut prompt_token_ids, replacements).unwrap(); - assert_eq!( - prompt_token_ids, - vec![ - 1, - LLAMA4_IMAGE_START_ID, - LLAMA4_IMAGE_ID, - LLAMA4_PATCH_ID, - LLAMA4_PATCH_ID, - LLAMA4_IMAGE_END_ID, - 2, - LLAMA4_IMAGE_START_ID, - LLAMA4_IMAGE_ID, - LLAMA4_PATCH_ID, - LLAMA4_PATCH_ID, - LLAMA4_IMAGE_END_ID, - 3, - ] - ); - assert_eq!(ranges[0].offset, 1); - assert_eq!(ranges[0].length, 5); - assert_bool_mask(&ranges[0], &[false, false, true, true, false]); - assert_eq!(ranges[1].offset, 7); - assert_eq!(ranges[1].length, 5); - assert_bool_mask(&ranges[1], &[false, false, true, true, false]); + assert_eq!(info.placeholder_token(Modality::Image), Some("<|image|>")); + assert_eq!(info.placeholder_token(Modality::Video), None); } } diff --git a/rust/src/chat/src/multimodal/expand.rs b/rust/src/chat/src/multimodal/expand.rs new file mode 100644 index 0000000000..3ea20db3b9 --- /dev/null +++ b/rust/src/chat/src/multimodal/expand.rs @@ -0,0 +1,446 @@ +//! Prompt placeholder expansion shared across modalities. + +use std::collections::{HashMap, VecDeque}; + +use llm_multimodal::{Modality, PromptReplacement}; +use aphrodite_engine_core_client::protocol::multimodal::PlaceholderRange; +use aphrodite_engine_core_client::protocol::tensor::WireTensor; + +use super::PreparedMedia; +use crate::error::{Error, Result, bail_multimodal}; + +/// One modality's queue of pending placeholder replacements for prompt +/// expansion. +struct ExpansionLane<'a> { + modality: Modality, + marker_token_id: u32, + embed_token_id: u32, + placeholder_token: String, + replacements: VecDeque<&'a PromptReplacement>, +} + +impl<'a> ExpansionLane<'a> { + fn from_prepared(media: &'a PreparedMedia) -> Option { + if media.replacements.is_empty() { + return None; + } + + Some(Self { + modality: media.modality, + marker_token_id: media.placeholder.marker_token_id, + embed_token_id: media.placeholder.embed_token_id, + placeholder_token: media.placeholder.token.clone(), + replacements: media.replacements.iter().collect(), + }) + } +} + +/// Replace rendered placeholder markers with model-specific replacement +/// tokens across all modalities in one left-to-right pass. +/// +/// Each prepared modality consumes its own marker occurrences in order, +/// matching the original media-part order within that modality; markers of +/// different modalities may interleave freely. +/// +/// The returned ranges point into the already-expanded prompt, grouped per +/// modality in item order. +pub(super) fn expand_prompt_token_ids( + prompt_token_ids: &mut Vec, + prepared: &[PreparedMedia], +) -> Result>> { + let mut lanes = prepared.iter().filter_map(ExpansionLane::from_prepared).collect::>(); + if lanes.is_empty() { + return Ok(HashMap::new()); + } + + let replacement_growth = lanes + .iter() + .flat_map(|lane| lane.replacements.iter()) + .fold(0usize, |total, replacement| { + total.saturating_add(replacement.tokens.len().saturating_sub(1)) + }); + let expanded_len = prompt_token_ids.len().saturating_add(replacement_growth); + + let mut expanded = Vec::with_capacity(expanded_len); + let mut ranges = HashMap::>::new(); + + for &token in prompt_token_ids.iter() { + let lane = lanes + .iter_mut() + .find(|lane| lane.marker_token_id == token && !lane.replacements.is_empty()); + let Some(lane) = lane else { + expanded.push(token); + continue; + }; + + let replacement = lane.replacements.pop_front().expect("lane queue is non-empty"); + debug_assert_eq!(replacement.modality, lane.modality); + if replacement.tokens.is_empty() { + bail_multimodal!( + "placeholder token `{}` expanded to no tokens", + lane.placeholder_token + ); + } + + let replacement_len = replacement.tokens.len(); + let is_embed = { + let mask = replacement + .tokens + .iter() + .map(|&token| token as u32 == lane.embed_token_id) + .collect::>(); + WireTensor::from_bool(vec![replacement_len], mask).map_err(Error::Multimodal)? + }; + + let expanded_offset = expanded.len(); + expanded.extend(replacement.tokens.iter().map(|&token| token as u32)); + ranges.entry(lane.modality).or_default().push(PlaceholderRange { + offset: expanded_offset, + length: replacement_len, + is_embed: Some(is_embed), + }); + } + + for lane in &lanes { + if !lane.replacements.is_empty() { + bail_multimodal!( + "placeholder token `{}` was not found in tokenized prompt for {} remaining `{}` item(s)", + lane.placeholder_token, + lane.replacements.len(), + lane.modality + ); + } + } + + *prompt_token_ids = expanded; + + Ok(ranges) +} + +#[cfg(test)] +mod tests { + use llm_multimodal::TokenId; + use aphrodite_engine_core_client::protocol::tensor::WireArrayData; + + use super::super::tests::{ + LLAMA4_IMAGE_END_ID, LLAMA4_IMAGE_ID, LLAMA4_IMAGE_START_ID, LLAMA4_PATCH_ID, + LLAMA4_TILE_X_SEPARATOR_ID, LLAMA4_TILE_Y_SEPARATOR_ID, QWEN3_IMAGE_PAD_ID, + QWEN3_VIDEO_PAD_ID, + }; + use super::super::{PreparedMedia, ResolvedPlaceholder}; + use super::*; + + /// Build prepared media directly from placeholder token IDs. + fn prepared_media( + modality: Modality, + placeholder_token: &str, + marker_token_id: u32, + embed_token_id: u32, + replacements: Vec, + ) -> PreparedMedia { + PreparedMedia { + modality, + placeholder: ResolvedPlaceholder { + token: placeholder_token.to_string(), + marker_token_id, + embed_token_id, + }, + replacements, + items: Vec::new(), + } + } + + /// Llama4 image prepared media: the `<|image|>` marker expands to + /// sequences whose embed positions are the `<|patch|>` tokens. + fn llama4_prepared(replacements: Vec) -> PreparedMedia { + prepared_media( + Modality::Image, + "<|image|>", + LLAMA4_IMAGE_ID, + LLAMA4_PATCH_ID, + replacements, + ) + } + + fn qwen3_image_prepared(replacements: Vec) -> PreparedMedia { + prepared_media( + Modality::Image, + "<|image_pad|>", + QWEN3_IMAGE_PAD_ID, + QWEN3_IMAGE_PAD_ID, + replacements, + ) + } + + fn qwen3_video_prepared(replacements: Vec) -> PreparedMedia { + prepared_media( + Modality::Video, + "<|video_pad|>", + QWEN3_VIDEO_PAD_ID, + QWEN3_VIDEO_PAD_ID, + replacements, + ) + } + + fn llama4_single_tile_replacement() -> PromptReplacement { + PromptReplacement::sequence( + Modality::Image, + "<|image|>", + vec![ + LLAMA4_IMAGE_START_ID as TokenId, + LLAMA4_IMAGE_ID as TokenId, + LLAMA4_PATCH_ID as TokenId, + LLAMA4_PATCH_ID as TokenId, + LLAMA4_IMAGE_END_ID as TokenId, + ], + ) + } + + fn llama4_multi_tile_replacement() -> PromptReplacement { + PromptReplacement::sequence( + Modality::Image, + "<|image|>", + vec![ + LLAMA4_IMAGE_START_ID as TokenId, + LLAMA4_PATCH_ID as TokenId, + LLAMA4_TILE_X_SEPARATOR_ID as TokenId, + LLAMA4_PATCH_ID as TokenId, + LLAMA4_TILE_Y_SEPARATOR_ID as TokenId, + LLAMA4_IMAGE_ID as TokenId, + LLAMA4_PATCH_ID as TokenId, + LLAMA4_IMAGE_END_ID as TokenId, + ], + ) + } + + fn assert_bool_mask(range: &PlaceholderRange, expected: &[bool]) { + let tensor = range.is_embed.as_ref().expect("is_embed mask"); + assert_eq!(tensor.dtype, "bool"); + assert_eq!(tensor.shape, vec![expected.len()]); + assert_eq!( + tensor.data, + WireArrayData::RawView(expected.iter().map(|value| u8::from(*value)).collect()) + ); + } + + #[test] + fn expand_prompt_tokens_marks_only_llama4_patch_tokens_as_embed() { + let mut prompt_token_ids = vec![1, LLAMA4_IMAGE_ID, 2]; + let prepared = vec![llama4_prepared(vec![llama4_multi_tile_replacement()])]; + + let ranges = expand_prompt_token_ids(&mut prompt_token_ids, &prepared).unwrap(); + let ranges = &ranges[&Modality::Image]; + + assert_eq!( + prompt_token_ids, + vec![ + 1, + LLAMA4_IMAGE_START_ID, + LLAMA4_PATCH_ID, + LLAMA4_TILE_X_SEPARATOR_ID, + LLAMA4_PATCH_ID, + LLAMA4_TILE_Y_SEPARATOR_ID, + LLAMA4_IMAGE_ID, + LLAMA4_PATCH_ID, + LLAMA4_IMAGE_END_ID, + 2, + ] + ); + assert_eq!(ranges[0].offset, 1); + assert_eq!(ranges[0].length, 8); + assert_bool_mask( + &ranges[0], + &[false, true, false, true, false, false, true, false], + ); + } + + #[test] + fn expand_prompt_tokens_errors_when_placeholder_missing() { + let mut prompt_token_ids = vec![1, 2, 3]; + let prepared = vec![llama4_prepared(vec![llama4_single_tile_replacement()])]; + + let error = expand_prompt_token_ids(&mut prompt_token_ids, &prepared).unwrap_err(); + + assert!(matches!(error, Error::Multimodal(message) if message.contains("not found"))); + } + + #[test] + fn expand_prompt_tokens_ignores_empty_replacements() { + let mut prompt_token_ids = vec![1, LLAMA4_IMAGE_ID, 2]; + let original_prompt_token_ids = prompt_token_ids.clone(); + let prepared = vec![llama4_prepared(Vec::new())]; + + let ranges = expand_prompt_token_ids(&mut prompt_token_ids, &prepared).unwrap(); + + assert!(ranges.is_empty()); + assert_eq!(prompt_token_ids, original_prompt_token_ids); + } + + #[test] + fn expand_prompt_tokens_leaves_prompt_unchanged_when_later_placeholder_missing() { + let mut prompt_token_ids = vec![1, LLAMA4_IMAGE_ID, 2]; + let original_prompt_token_ids = prompt_token_ids.clone(); + let prepared = vec![llama4_prepared(vec![ + llama4_single_tile_replacement(), + llama4_single_tile_replacement(), + ])]; + + let error = expand_prompt_token_ids(&mut prompt_token_ids, &prepared).unwrap_err(); + + assert!(matches!(error, Error::Multimodal(message) if message.contains("not found"))); + assert_eq!(prompt_token_ids, original_prompt_token_ids); + } + + #[test] + fn expand_prompt_tokens_errors_when_replacement_is_empty() { + let mut prompt_token_ids = vec![1, LLAMA4_IMAGE_ID, 2]; + let original_prompt_token_ids = prompt_token_ids.clone(); + let prepared = vec![llama4_prepared(vec![PromptReplacement::sequence( + Modality::Image, + "<|image|>", + Vec::new(), + )])]; + + let error = expand_prompt_token_ids(&mut prompt_token_ids, &prepared).unwrap_err(); + + assert!( + matches!(error, Error::Multimodal(message) if message.contains("expanded to no tokens")) + ); + assert_eq!(prompt_token_ids, original_prompt_token_ids); + } + + #[test] + fn expand_prompt_tokens_skips_llama4_image_marker_inside_replacement() { + let mut prompt_token_ids = vec![1, LLAMA4_IMAGE_ID, 2, LLAMA4_IMAGE_ID, 3]; + let prepared = vec![llama4_prepared(vec![ + llama4_single_tile_replacement(), + llama4_single_tile_replacement(), + ])]; + + let ranges = expand_prompt_token_ids(&mut prompt_token_ids, &prepared).unwrap(); + let ranges = &ranges[&Modality::Image]; + + assert_eq!( + prompt_token_ids, + vec![ + 1, + LLAMA4_IMAGE_START_ID, + LLAMA4_IMAGE_ID, + LLAMA4_PATCH_ID, + LLAMA4_PATCH_ID, + LLAMA4_IMAGE_END_ID, + 2, + LLAMA4_IMAGE_START_ID, + LLAMA4_IMAGE_ID, + LLAMA4_PATCH_ID, + LLAMA4_PATCH_ID, + LLAMA4_IMAGE_END_ID, + 3, + ] + ); + assert_eq!(ranges[0].offset, 1); + assert_eq!(ranges[0].length, 5); + assert_bool_mask(&ranges[0], &[false, false, true, true, false]); + assert_eq!(ranges[1].offset, 7); + assert_eq!(ranges[1].length, 5); + assert_bool_mask(&ranges[1], &[false, false, true, true, false]); + } + + #[test] + fn expand_prompt_tokens_interleaves_image_and_video_prepared_media() { + let mut prompt_token_ids = vec![ + 1, + QWEN3_IMAGE_PAD_ID, + 2, + QWEN3_VIDEO_PAD_ID, + 3, + QWEN3_IMAGE_PAD_ID, + 4, + ]; + let prepared = vec![ + qwen3_image_prepared(vec![ + PromptReplacement::repeated( + Modality::Image, + "<|image_pad|>", + QWEN3_IMAGE_PAD_ID as TokenId, + 2, + ), + PromptReplacement::repeated( + Modality::Image, + "<|image_pad|>", + QWEN3_IMAGE_PAD_ID as TokenId, + 3, + ), + ]), + qwen3_video_prepared(vec![PromptReplacement::repeated( + Modality::Video, + "<|video_pad|>", + QWEN3_VIDEO_PAD_ID as TokenId, + 4, + )]), + ]; + + let ranges = expand_prompt_token_ids(&mut prompt_token_ids, &prepared).unwrap(); + + assert_eq!( + prompt_token_ids, + vec![ + 1, + QWEN3_IMAGE_PAD_ID, + QWEN3_IMAGE_PAD_ID, + 2, + QWEN3_VIDEO_PAD_ID, + QWEN3_VIDEO_PAD_ID, + QWEN3_VIDEO_PAD_ID, + QWEN3_VIDEO_PAD_ID, + 3, + QWEN3_IMAGE_PAD_ID, + QWEN3_IMAGE_PAD_ID, + QWEN3_IMAGE_PAD_ID, + 4, + ] + ); + + let image_ranges = &ranges[&Modality::Image]; + assert_eq!(image_ranges[0].offset, 1); + assert_eq!(image_ranges[0].length, 2); + assert_bool_mask(&image_ranges[0], &[true, true]); + assert_eq!(image_ranges[1].offset, 9); + assert_eq!(image_ranges[1].length, 3); + assert_bool_mask(&image_ranges[1], &[true, true, true]); + + let video_ranges = &ranges[&Modality::Video]; + assert_eq!(video_ranges[0].offset, 4); + assert_eq!(video_ranges[0].length, 4); + assert_bool_mask(&video_ranges[0], &[true, true, true, true]); + } + + #[test] + fn expand_prompt_tokens_error_names_modality_with_leftover_replacements() { + let mut prompt_token_ids = vec![1, QWEN3_IMAGE_PAD_ID, 2]; + let original_prompt_token_ids = prompt_token_ids.clone(); + let prepared = vec![ + qwen3_image_prepared(vec![PromptReplacement::repeated( + Modality::Image, + "<|image_pad|>", + QWEN3_IMAGE_PAD_ID as TokenId, + 2, + )]), + qwen3_video_prepared(vec![PromptReplacement::repeated( + Modality::Video, + "<|video_pad|>", + QWEN3_VIDEO_PAD_ID as TokenId, + 4, + )]), + ]; + + let error = expand_prompt_token_ids(&mut prompt_token_ids, &prepared).unwrap_err(); + + assert!(matches!( + error, + Error::Multimodal(message) + if message.contains("<|video_pad|>") && message.contains("`video`") + )); + assert_eq!(prompt_token_ids, original_prompt_token_ids); + } +} diff --git a/rust/src/chat/src/multimodal/image.rs b/rust/src/chat/src/multimodal/image.rs new file mode 100644 index 0000000000..53741b6534 --- /dev/null +++ b/rust/src/chat/src/multimodal/image.rs @@ -0,0 +1,141 @@ +//! Image-modality preparation: batch preprocessing and per-item feature +//! build. + +use std::sync::Arc; + +use itertools::izip; +use llm_multimodal::{FieldLayout, ImageFrame, Modality, PreprocessedEncoderInputs}; +use aphrodite_engine_core_client::protocol::dtype::ModelDtype; +use aphrodite_engine_core_client::protocol::multimodal::{ + MmBatchedField, MmField, MmFieldElem, MmFlatField, MmKwargsItem, MmSharedField, MmSlice, + SliceSpec, +}; + +use super::{ModalitySupport, MultimodalModelInfo, PreparedItem, PreparedMedia, tensor}; +use crate::error::{Error, Result, bail_multimodal, multimodal}; + +impl MultimodalModelInfo { + /// Preprocess all fetched image frames as one batch and build per-item + /// features. + pub(super) async fn prepare_images( + &self, + frames: Vec>, + uuids: Vec>, + model_dtype: ModelDtype, + ) -> Result { + let support = self.image.as_ref().ok_or_else(|| Error::UnsupportedModality { + modality: Modality::Image.to_string(), + })?; + let preprocessed = self.preprocess_images(support, &frames).await?; + let replacements = + self.spec + .prompt_replacements_for(&self.context, &preprocessed, Modality::Image)?; + if replacements.len() != frames.len() { + bail_multimodal!( + "number of image prompt replacements {} does not match number of images {}", + replacements.len(), + frames.len() + ); + } + let items = self.build_image_items(preprocessed, &frames, uuids, model_dtype)?; + + Ok(PreparedMedia { + modality: Modality::Image, + placeholder: support.placeholder.clone(), + replacements, + items, + }) + } + + /// Preprocess fetched image frames with the model's resolved vision + /// processor. + /// + /// The processor work is CPU-heavy relative to request wiring, so it runs + /// in a blocking task and returns owned tensors ready for wire + /// conversion. + async fn preprocess_images( + &self, + support: &ModalitySupport, + image_frames: &[Arc], + ) -> Result { + let config = support.config.clone(); + let processor = support.processor; + let images = image_frames.iter().map(|frame| frame.data().clone()).collect::>(); + + // TODO: is it still necessary given that we've already in a dedicated runtime? + tokio::task::spawn_blocking(move || Ok(processor.preprocess(&images, &config)?)) + .await + .map_err(|error| multimodal!("image preprocessing task failed: {error}"))? + } + + /// Convert one batch of preprocessed image tensors into per-item engine + /// kwargs. + /// + /// Tensor fields are sliced per item according to the model spec's field + /// layout declarations. + fn build_image_items( + &self, + preprocessed: PreprocessedEncoderInputs, + frames: &[Arc], + uuids: Vec>, + model_dtype: ModelDtype, + ) -> Result> { + let len = frames.len(); + let tensors = tensor::collect_tensors(preprocessed, "pixel_values", model_dtype)?; + + let mut items = Vec::with_capacity(len); + for (index, (frame, uuid)) in izip!(frames, uuids).enumerate() { + let mut data = MmKwargsItem::new(); + for (key, tensor) in &tensors { + let keep_on_cpu = self.spec.keep_on_cpu_keys.contains(key); + let (value, field) = match self.spec.field_layouts.get(key) { + Some(FieldLayout::Batched) => ( + tensor.batched_value_at(index)?, + MmField::Batched(MmBatchedField { keep_on_cpu }), + ), + Some(FieldLayout::Flat { sizes_key }) => { + let sizes = tensors.get(sizes_key).ok_or_else(|| { + multimodal!("flat tensor sizes key `{sizes_key}` is missing") + })?; + let (start, end) = tensor::flat_range_for_index(sizes, sizes_key, index)?; + ( + tensor.flat_value_range(start, end)?, + MmField::Flat(MmFlatField { + slices: vec![MmSlice::Slice(SliceSpec { + start: Some(0), + stop: Some((end - start) as isize), + step: None, + })], + dim: 0, + keep_on_cpu, + }), + ) + } + None => ( + tensor.clone(), + MmField::Shared(MmSharedField { + batch_size: len, + keep_on_cpu, + }), + ), + }; + + data.insert( + key.clone(), + MmFieldElem { + data: Some(value.try_into()?), + field, + }, + ); + } + + items.push(PreparedItem { + data, + hash: frame.hash.clone(), + uuid, + }); + } + + Ok(items) + } +} diff --git a/rust/src/chat/src/multimodal/tensor.rs b/rust/src/chat/src/multimodal/tensor.rs index 6c9bab8ad1..dbed8e1512 100644 --- a/rust/src/chat/src/multimodal/tensor.rs +++ b/rust/src/chat/src/multimodal/tensor.rs @@ -1,7 +1,7 @@ use std::collections::HashMap; use half::{bf16, f16}; -use llm_multimodal::{ModelSpecificValue, PreprocessedEncoderInputs as PreprocessedImages}; +use llm_multimodal::{ModelSpecificValue, PreprocessedEncoderInputs}; use aphrodite_engine_core_client::protocol::dtype::ModelDtype; use aphrodite_engine_core_client::protocol::multimodal::MmKwargValue as ProtocolKwargValue; use aphrodite_engine_core_client::protocol::tensor::{ShapeExt as _, WireTensor}; @@ -25,25 +25,31 @@ pub(super) enum KwargValue { Passthrough(ProtocolKwargValue), } -/// Collect `pixel_values` and model-specific outputs into one tensor map. +/// Collect the primary encoder input and model-specific outputs into one +/// tensor map. +/// +/// `primary_key` names the encoder-input tensor as the model's forward kwargs +/// expect it (e.g. `pixel_values` for images, `pixel_values_videos` for +/// videos). pub(super) fn collect_tensors( - preprocessed: PreprocessedImages, + preprocessed: PreprocessedEncoderInputs, + primary_key: &str, float_dtype: ModelDtype, ) -> Result> { - let PreprocessedImages { + let PreprocessedEncoderInputs { encoder_input, model_specific, .. } = preprocessed; - let pixel_values = { + let primary_value = { let shape = encoder_input.shape().to_vec(); let data = encoder_input.into_iter().collect(); KwargValue::from_f32_tensor(data, shape, float_dtype)? }; let mut tensors = HashMap::new(); - tensors.insert("pixel_values".to_string(), pixel_values); + tensors.insert(primary_key.to_string(), primary_value); for (key, value) in model_specific { tensors.insert(key, KwargValue::from_model_specific(value, float_dtype)?); } @@ -124,10 +130,22 @@ impl TryFrom for ProtocolKwargValue { } impl KwargValue { - /// Extract one image from a batched tensor field. + /// First-axis length for tensor values; `None` for passthrough kwargs. + pub(super) fn first_dim(&self) -> Option { + match self { + Self::F32Tensor { shape, .. } + | Self::F16Tensor { shape, .. } + | Self::Bf16Tensor { shape, .. } + | Self::I64Tensor { shape, .. } + | Self::U32Tensor { shape, .. } => shape.first().copied(), + Self::Passthrough(_) => None, + } + } + + /// Extract one media item from a batched tensor field. /// - /// Batched fields use their first axis as image index and drop that axis in - /// the per-feature value, matching Aphrodite's batched-field semantics. + /// Batched fields use their first axis as media-item index and drop that + /// axis in the per-feature value, matching Aphrodite's batched-field semantics. pub(super) fn batched_value_at(&self, index: usize) -> Result { match self { Self::F32Tensor { data, shape } => { @@ -154,9 +172,9 @@ impl KwargValue { } } - /// Extract one image's variable-length range from a flat tensor field. + /// Extract one media item's variable-length range from a flat tensor field. /// - /// Flat fields keep the first axis as the sliced length for this image. + /// Flat fields keep the first axis as the sliced length for this item. pub(super) fn flat_value_range(&self, start: usize, end: usize) -> Result { match self { Self::F32Tensor { data, shape } => { @@ -184,10 +202,10 @@ impl KwargValue { } } -/// Compute the first-axis range for one image in a flat tensor. +/// Compute the first-axis range for one media item in a flat tensor. /// /// `sizes_key` names a companion tensor whose entries are cumulative slice -/// sizes per image. +/// sizes per media item. pub(super) fn flat_range_for_index( sizes: &KwargValue, sizes_key: &str, @@ -195,7 +213,7 @@ pub(super) fn flat_range_for_index( ) -> Result<(usize, usize)> { let sizes = tensor_as_usize_vec(sizes)?; let size = *sizes.get(index).ok_or_else(|| { - multimodal!("flat tensor sizes key `{sizes_key}` has no entry for image {index}") + multimodal!("flat tensor sizes key `{sizes_key}` has no entry for media item {index}") })?; let start = sizes[..index].iter().sum::(); Ok((start, start + size)) diff --git a/rust/src/chat/src/multimodal/video.rs b/rust/src/chat/src/multimodal/video.rs new file mode 100644 index 0000000000..f33c4675b3 --- /dev/null +++ b/rust/src/chat/src/multimodal/video.rs @@ -0,0 +1,316 @@ +//! Video-modality preparation: per-clip preprocessing, config resolution, +//! and per-item feature build. + +use std::sync::Arc; + +use itertools::izip; +use llm_multimodal::{FieldLayout, Modality, PreprocessedEncoderInputs, VideoClip}; +use thiserror_ext::AsReport as _; +use tracing::warn; +use aphrodite_engine_core_client::protocol::dtype::ModelDtype; +use aphrodite_engine_core_client::protocol::multimodal::{ + MmBatchedField, MmField, MmFieldElem, MmFlatField, MmKwargsItem, MmSharedField, MmSlice, + SliceSpec, +}; + +use super::{ModalitySupport, MultimodalModelInfo, PreparedItem, PreparedMedia, tensor}; +use crate::error::{Error, Result, bail_multimodal, multimodal}; + +/// Forward-kwargs name of the primary video encoder input. +/// +/// Video-capable vLLM models read `pixel_values_videos` alongside +/// `video_grid_thw`, mirroring the HF processor output naming. +const VIDEO_PRIMARY_KEY: &str = "pixel_values_videos"; + +impl MultimodalModelInfo { + /// Preprocess fetched video clips one at a time and build per-item + /// features. + /// + /// Unlike images, each clip runs through the preprocessor independently + /// (a batch of one), so its tensors are complete per item and need no + /// cross-item slicing. + pub(super) async fn prepare_videos( + &self, + clips: Vec>, + uuids: Vec>, + model_dtype: ModelDtype, + ) -> Result { + let support = self.video.as_ref().ok_or_else(|| Error::UnsupportedModality { + modality: Modality::Video.to_string(), + })?; + let mut replacements = Vec::with_capacity(clips.len()); + let mut items = Vec::with_capacity(clips.len()); + + for (clip, uuid) in izip!(&clips, uuids) { + let preprocessed = self.preprocess_video_clip(support, Arc::clone(clip)).await?; + let mut clip_replacements = + self.spec + .prompt_replacements_for(&self.context, &preprocessed, Modality::Video)?; + if clip_replacements.len() != 1 { + bail_multimodal!( + "expected exactly one prompt replacement per video clip, got {}", + clip_replacements.len() + ); + } + replacements.push(clip_replacements.pop().unwrap()); + items.push(self.build_video_item( + preprocessed, + clip.hash.clone(), + uuid, + model_dtype, + )?); + } + + Ok(PreparedMedia { + modality: Modality::Video, + placeholder: support.placeholder.clone(), + replacements, + items, + }) + } + + /// Preprocess one decoded video clip with the model's resolved vision + /// processor. + async fn preprocess_video_clip( + &self, + support: &ModalitySupport, + clip: Arc, + ) -> Result { + let config = support.config.clone(); + let processor = support.processor; + + tokio::task::spawn_blocking(move || { + // Prefer the borrowed-RGB fast path, which avoids materializing a + // `DynamicImage` per sampled frame after media decode. + if let Some(rgb_video) = clip.rgb_video() { + match rgb_video.frame_refs() { + Ok(frame_refs) => match processor.preprocess_video_rgb(&frame_refs, &config) { + Ok(preprocessed) => return Ok(preprocessed), + Err(error) => warn!( + error = %error.as_report(), + "RGB video preprocessing fast path failed; falling back to materialized frames" + ), + }, + Err(error) => warn!( + error, + "RGB video frame refs are invalid; falling back to materialized frames" + ), + } + } + + let frames = clip.materialized_frames().map_err(|error| multimodal!("{error}"))?; + Ok(processor.preprocess_video(&frames, &config)?) + }) + .await + .map_err(|error| multimodal!("video preprocessing task failed: {error}"))? + } + + /// Convert one preprocessed video clip into engine kwargs. + /// + /// The clip is a batch of one, so no per-item slicing is required: the + /// primary tensor ships as a full-range flat field (the engine re-batches + /// flat fields by concatenating along the declared dim, matching vLLM's + /// `flat_from_sizes` treatment of video patches), and batched metadata + /// tensors drop their singleton batch axis. + fn build_video_item( + &self, + preprocessed: PreprocessedEncoderInputs, + hash: String, + uuid: Option, + model_dtype: ModelDtype, + ) -> Result { + let tensors = tensor::collect_tensors(preprocessed, VIDEO_PRIMARY_KEY, model_dtype)?; + + let mut data = MmKwargsItem::new(); + for (key, tensor) in tensors { + let keep_on_cpu = self.spec.keep_on_cpu_keys.contains(&key); + let (value, field) = if key == VIDEO_PRIMARY_KEY { + let len = tensor + .first_dim() + .ok_or_else(|| multimodal!("video encoder input `{key}` is not a tensor"))?; + ( + tensor, + MmField::Flat(MmFlatField { + slices: vec![MmSlice::Slice(SliceSpec { + start: Some(0), + stop: Some(len as isize), + step: None, + })], + dim: 0, + keep_on_cpu, + }), + ) + } else if matches!( + self.spec.field_layouts.get(&key), + Some(FieldLayout::Batched) + ) { + ( + tensor.batched_value_at(0)?, + MmField::Batched(MmBatchedField { keep_on_cpu }), + ) + } else { + ( + tensor, + MmField::Shared(MmSharedField { + batch_size: 1, + keep_on_cpu, + }), + ) + }; + + data.insert( + key, + MmFieldElem { + data: Some(value.try_into()?), + field, + }, + ); + } + + Ok(PreparedItem { data, hash, uuid }) + } +} + +#[cfg(test)] +mod tests { + use std::collections::HashMap; + use std::sync::Arc; + + use llm_multimodal::ModelSpecificValue; + use ndarray::ArrayD; + use aphrodite_engine_core_client::protocol::multimodal::MmKwargValue; + + use super::super::tests::{ + QWEN3_IMAGE_PAD_ID, QWEN3_VIDEO_PAD_ID, qwen3_vl_info, qwen3_vl_tokenizer, + }; + use super::super::{MultimodalConfigFiles, MultimodalModelInfo}; + use super::*; + + #[test] + fn from_paths_resolves_video_config_from_dedicated_file_or_processor_config() { + let dir = tempfile::tempdir().unwrap(); + let config_path = dir.path().join("config.json"); + std::fs::write( + &config_path, + serde_json::json!({ + "model_type": "qwen3_vl", + "image_token_id": QWEN3_IMAGE_PAD_ID, + "video_token_id": QWEN3_VIDEO_PAD_ID, + }) + .to_string(), + ) + .unwrap(); + + let info_for = |files: MultimodalConfigFiles<'_>| { + MultimodalModelInfo::from_paths( + "qwen3-vl-test".to_string(), + Some("qwen3_vl".to_string()), + files, + Arc::new(qwen3_vl_tokenizer()), + ) + }; + + // Dedicated video preprocessor config file. + let video_config_path = dir.path().join("video_preprocessor_config.json"); + std::fs::write(&video_config_path, r#"{"size":{"shortest_edge":128}}"#).unwrap(); + let info = info_for(MultimodalConfigFiles { + config: Some(&config_path), + video_preprocessor_config: Some(&video_config_path), + ..Default::default() + }) + .unwrap() + .unwrap(); + assert!(info.video.is_some()); + + // `video_processor` section of the combined processor config. + let processor_config_path = dir.path().join("processor_config.json"); + std::fs::write( + &processor_config_path, + r#"{"video_processor":{"size":{"shortest_edge":128}}}"#, + ) + .unwrap(); + let info = info_for(MultimodalConfigFiles { + config: Some(&config_path), + processor_config: Some(&processor_config_path), + ..Default::default() + }) + .unwrap() + .unwrap(); + assert!(info.video.is_some()); + + // Neither source: video support still resolves on the image config. + let info = info_for(MultimodalConfigFiles { + config: Some(&config_path), + ..Default::default() + }) + .unwrap() + .unwrap(); + assert!(info.video.is_some()); + + // Malformed dedicated file is a real error, not a silent fallback. + std::fs::write(&video_config_path, r#"{"size""#).unwrap(); + let error = match info_for(MultimodalConfigFiles { + config: Some(&config_path), + video_preprocessor_config: Some(&video_config_path), + ..Default::default() + }) { + Err(error) => error, + Ok(_) => panic!("malformed video preprocessor config should fail"), + }; + assert!(matches!( + error, + Error::Multimodal(message) + if message.contains("failed to parse video_preprocessor_config.json") + )); + } + + #[test] + fn build_video_item_names_primary_tensor_and_layouts() { + let info = qwen3_vl_info(); + // One clip flattened to 6 patches with 4 features each. + let preprocessed = PreprocessedEncoderInputs { + encoder_input: ArrayD::zeros(vec![6, 4]), + feature_token_counts: vec![6], + item_sizes: vec![(32, 32)], + model_specific: HashMap::from([ + ( + "video_grid_thw".to_string(), + ModelSpecificValue::int_2d(vec![1, 2, 3], 1, 3), + ), + ( + "patches_per_video".to_string(), + ModelSpecificValue::int_1d(vec![6]), + ), + ]), + }; + + let item = info + .build_video_item( + preprocessed, + "".to_string(), + None, + ModelDtype::Float32, + ) + .unwrap(); + + let primary = &item.data[VIDEO_PRIMARY_KEY]; + assert!(matches!( + &primary.field, + MmField::Flat(MmFlatField { slices, dim: 0, .. }) + if matches!( + slices.as_slice(), + [MmSlice::Slice(SliceSpec { start: Some(0), stop: Some(6), step: None })] + ) + )); + + // Batched metadata drops its singleton batch axis per item. + let grid = &item.data["video_grid_thw"]; + assert!(matches!(&grid.field, MmField::Batched(_))); + let MmKwargValue::Tensor(grid_tensor) = grid.data.as_ref().unwrap() else { + panic!("expected tensor value for video_grid_thw"); + }; + assert_eq!(grid_tensor.shape, vec![3]); + + assert_eq!(item.hash, ""); + } +} diff --git a/rust/src/chat/src/renderer/hf/mod.rs b/rust/src/chat/src/renderer/hf/mod.rs index 3e252ac0a1..ab694ec543 100644 --- a/rust/src/chat/src/renderer/hf/mod.rs +++ b/rust/src/chat/src/renderer/hf/mod.rs @@ -31,9 +31,14 @@ pub use template::{load_chat_template, resolve_chat_template}; pub use self::format::ChatTemplateContentFormatOption; -#[derive(Debug, Clone)] +/// Template-visible placeholder tokens per supported modality. +/// +/// A `None` token means the loaded model does not support that modality, and +/// content parts of that modality are rejected during rendering. +#[derive(Debug, Clone, Default)] pub struct MultimodalRenderInfo { - pub placeholder_token: String, + pub image_token: Option, + pub video_token: Option, } /// Hugging Face chat-template renderer backed by the local Jinja chat-template @@ -254,6 +259,7 @@ enum TemplateContent { enum TemplateContentPart { Text { text: String }, Image, + Video, } #[derive(Debug, Serialize)] @@ -417,9 +423,17 @@ fn to_template_openai_content( } // All multimodal contents are normalized to `{ "type": }`. ChatContentPart::ImageUrl { .. } => { - multimodal.ok_or(Error::UnsupportedMultimodalContent("image_url"))?; + multimodal + .and_then(|multimodal| multimodal.image_token.as_ref()) + .ok_or(Error::UnsupportedMultimodalContent("image_url"))?; Ok(TemplateContentPart::Image) } + ChatContentPart::VideoUrl { .. } => { + multimodal + .and_then(|multimodal| multimodal.video_token.as_ref()) + .ok_or(Error::UnsupportedMultimodalContent("video_url"))?; + Ok(TemplateContentPart::Video) + } }) .collect(), } @@ -437,9 +451,16 @@ fn to_template_string_content( match part { ChatContentPart::Text { text } => out.push_str(text), ChatContentPart::ImageUrl { .. } => { - let multimodal = - multimodal.ok_or(Error::UnsupportedMultimodalContent("image_url"))?; - out.push_str(&multimodal.placeholder_token); + let image_token = multimodal + .and_then(|multimodal| multimodal.image_token.as_ref()) + .ok_or(Error::UnsupportedMultimodalContent("image_url"))?; + out.push_str(image_token); + } + ChatContentPart::VideoUrl { .. } => { + let video_token = multimodal + .and_then(|multimodal| multimodal.video_token.as_ref()) + .ok_or(Error::UnsupportedMultimodalContent("video_url"))?; + out.push_str(video_token); } } } @@ -468,7 +489,7 @@ fn append_continue_final_message_tag(message: &mut TemplateMessage) -> Result parts.iter_mut().rev().find_map(|part| match part { TemplateContentPart::Text { text } => Some(text), - TemplateContentPart::Image => None, + TemplateContentPart::Image | TemplateContentPart::Video => None, }), }; let text = text.ok_or_else(|| { @@ -577,7 +598,8 @@ mod tests { ) -> Result { HfChatRenderer::new(Some(template.to_string()), HashMap::new(), content_format)? .with_multimodal(Some(MultimodalRenderInfo { - placeholder_token: "".to_string(), + image_token: Some("".to_string()), + video_token: Some("