diff --git a/.buildkite/check-torch-abi.py b/.buildkite/check-torch-abi.py index ebf2fb4b27..b70e8ab6d0 100644 --- a/.buildkite/check-torch-abi.py +++ b/.buildkite/check-torch-abi.py @@ -15,6 +15,7 @@ # Shrink and remove over time. ALLOWED_UNSTABLE_LIBRARIES: tuple[str, ...] = ( "_C_fork.abi3.so", + "_flashkda_C.abi3.so", "vllm_flash_attn/_vllm_fa2_C.abi3.so", "vllm_flash_attn/_vllm_fa3_C.abi3.so", "third_party/deep_gemm/_C*.so", diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 6064ebd9e0..9005b81b5a 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -4,7 +4,7 @@ default_install_hook_types: default_stages: - pre-commit # Run locally - manual # Run in CI -exclude: '^(reference/|shellcheck-stable/|aphrodite/third_party/.*)' +exclude: '^(reference/|shellcheck-stable/|aphrodite/third_party/.*|aphrodite/models/kimi_k3/nvidia/ops/third_party/.*|aphrodite/models/kimi_k3/amd/ops/third_party/.*)' repos: - repo: https://github.com/astral-sh/ruff-pre-commit rev: v0.14.0 @@ -126,6 +126,21 @@ repos: language: system always_run: true pass_filenames: false + - id: check-patch-artifacts + name: Reject patch backup files + entry: bash + args: + - -c + - | + artifacts="$(git ls-files | grep -E '(\.orig|\.rej|\.bak|~)$' || true)" + if [[ -n "$artifacts" ]]; then + echo "Patch backup files must not be committed:" + echo "$artifacts" + exit 1 + fi + language: system + always_run: true + pass_filenames: false - id: enforce-import-regex-instead-of-re name: Enforce import regex as re entry: python tools/pre_commit/enforce_regex_import.py diff --git a/.sync/vllm-sha b/.sync/vllm-sha index 7be2e43b48..3c9542fda3 100644 --- a/.sync/vllm-sha +++ b/.sync/vllm-sha @@ -1 +1 @@ -8a7b3c299053efbca2669081ddf50a46b6a9d149 +9a4e5f95390fc759ada001444637fe2e96a29ad9 diff --git a/CMakeLists.txt b/CMakeLists.txt index 2fd8f0c968..782aecd67d 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -424,8 +424,11 @@ if(APHRODITE_GPU_LANG STREQUAL "CUDA" OR APHRODITE_GPU_LANG STREQUAL "HIP") "csrc/libtorch_stable/mamba/selective_scan_fwd.cu" "csrc/libtorch_stable/cache_kernels.cu" "csrc/libtorch_stable/cache_kernels_fused.cu" + "csrc/libtorch_stable/custom_all_gather_reduce_scatter.cu" + "csrc/libtorch_stable/custom_all_gather_reduce_scatter_ops.cpp" "csrc/libtorch_stable/custom_all_reduce.cu" - "csrc/libtorch_stable/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu") + "csrc/libtorch_stable/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu" + "csrc/libtorch_stable/fused_kimi_k3_mla_key_concat_kv_cache_kernel.cu") if(APHRODITE_GPU_LANG STREQUAL "CUDA" AND DEFINED CMAKE_CUDA_COMPILER_VERSION AND @@ -1139,6 +1142,23 @@ if(APHRODITE_GPU_LANG STREQUAL "CUDA" OR APHRODITE_GPU_LANG STREQUAL "HIP") set(MLA_ARCHS) endif() + if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0) + cuda_archs_loose_intersection(FUSED_KDA_DECODE_ARCHS + "9.0a;10.0f;12.0f" "${CUDA_ARCHS}") + endif() + if(FUSED_KDA_DECODE_ARCHS) + set(FUSED_KDA_DECODE_SRC + "csrc/libtorch_stable/kimi_k3/fused_kda_decode_kernel.cu") + set_gencode_flags_for_srcs( + SRCS "${FUSED_KDA_DECODE_SRC}" + CUDA_ARCHS "${FUSED_KDA_DECODE_ARCHS}") + set_property(SOURCE ${FUSED_KDA_DECODE_SRC} APPEND PROPERTY + COMPILE_OPTIONS "$<$:--use_fast_math>") + list(APPEND APHRODITE_STABLE_EXT_SRC "${FUSED_KDA_DECODE_SRC}") + message(STATUS + "Building fused KDA decode for archs: ${FUSED_KDA_DECODE_ARCHS}") + endif() + if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0) cuda_archs_loose_intersection(KIMI_K3_ATTN_RES_ARCHS "10.0f" "${CUDA_ARCHS}") @@ -1240,6 +1260,10 @@ if(APHRODITE_GPU_LANG STREQUAL "CUDA" OR APHRODITE_GPU_LANG STREQUAL "HIP") target_compile_definitions(_C_stable_libtorch PRIVATE APHRODITE_ENABLE_SM89_DSA=1) endif() + if(FUSED_KDA_DECODE_ARCHS) + target_compile_definitions(_C_stable_libtorch PRIVATE + APHRODITE_ENABLE_FUSED_KDA_DECODE=1) + endif() if(KIMI_K3_ATTN_RES_ARCHS) target_compile_definitions(_C_stable_libtorch PRIVATE APHRODITE_ENABLE_KIMI_K3_ATTN_RES=1) @@ -1543,6 +1567,7 @@ if (APHRODITE_GPU_LANG STREQUAL "CUDA") include(cmake/external_projects/deepgemm.cmake) include(cmake/external_projects/fmha_sm100.cmake) include(cmake/external_projects/flashmla.cmake) + include(cmake/external_projects/flashkda.cmake) include(cmake/external_projects/qutlass.cmake) include(cmake/external_projects/tml_fa4.cmake) diff --git a/aphrodite/_custom_ops.py b/aphrodite/_custom_ops.py index e653f533c1..89b112ffdc 100644 --- a/aphrodite/_custom_ops.py +++ b/aphrodite/_custom_ops.py @@ -2895,17 +2895,6 @@ def topk_softmax( e_score_correction_bias: torch.Tensor | None = None, is_padding: torch.Tensor | None = None, ) -> None: - if current_platform.is_xpu(): - # TODO: Remove after vllm-xpu-kernels supports is_padding. - torch.ops._moe_C.topk_softmax( - topk_weights, - topk_ids, - token_expert_indices, - gating_output, - renormalize, - e_score_correction_bias, - ) - return torch.ops._moe_C.topk_softmax( topk_weights, topk_ids, @@ -2951,20 +2940,6 @@ def topk_hash_softplus_sqrt( hash_indices_table: torch.Tensor | None = None, is_padding: torch.Tensor | None = None, ) -> None: - if current_platform.is_xpu(): - # TODO: Remove after vllm-xpu-kernels supports is_padding. - torch.ops._moe_C.topk_softplus_sqrt( - topk_weights, - topk_indices, - token_expert_indices, - gating_output, - renormalize, - routed_scaling_factor, - e_score_correction_bias, - input_tokens, - hash_indices_table, - ) - return torch.ops._moe_C.topk_softplus_sqrt( topk_weights, topk_indices, @@ -3232,6 +3207,52 @@ def fused_minimax_m3_qknorm_rope_kv_insert( ) +def fused_kda_decode( + x: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor | None, + conv_state: torch.Tensor, + raw_g: torch.Tensor, + raw_beta: torch.Tensor, + A_log: torch.Tensor, + dt_bias: torch.Tensor, + state_indices: torch.Tensor, + state: torch.Tensor, + out: torch.Tensor | None = None, + lower_bound: float | None = None, + output_gate: torch.Tensor | None = None, + norm_weight: torch.Tensor | None = None, + norm_eps: float = 1e-5, +) -> torch.Tensor: + if out is None: + out = torch.empty( + 1, + x.shape[0], + raw_g.shape[2], + raw_g.shape[3], + dtype=x.dtype, + device=x.device, + ) + torch.ops._C.fused_kda_decode( + x, + weight, + bias, + conv_state, + raw_g, + raw_beta, + A_log, + dt_bias, + state_indices, + state, + out, + lower_bound, + output_gate, + norm_weight, + norm_eps, + ) + return out + + def concat_and_cache_mla( kv_c: torch.Tensor, k_pe: torch.Tensor, @@ -3243,6 +3264,26 @@ def concat_and_cache_mla( torch.ops._C_cache_ops.concat_and_cache_mla(kv_c, k_pe, kv_cache, slot_mapping, kv_cache_dtype, scale) +def concat_and_cache_mla_grouped( + kv_c: torch.Tensor, + k_pe: torch.Tensor, + kv_cache_ptrs: torch.Tensor, + slot_mapping: torch.Tensor, + block_size: int, + block_stride: int, + entry_stride: int, +) -> None: + torch.ops._C_cache_ops.concat_and_cache_mla_grouped( + kv_c, + k_pe, + kv_cache_ptrs, + slot_mapping, + block_size, + block_stride, + entry_stride, + ) + + def kimi_k3_attn_res( prefix: torch.Tensor, delta: torch.Tensor, @@ -3600,6 +3641,57 @@ def all_reduce( torch.ops._C_custom_ar.all_reduce(fa, inp, out, reg_buffer, reg_buffer_sz_bytes) +def custom_all_gather( + fa: int, + inp: torch.Tensor, + out: torch.Tensor, + reg_buffer: int, + reg_buffer_sz_bytes: int, +) -> None: + torch.ops._C_custom_ar.custom_all_gather(fa, inp, out, reg_buffer, reg_buffer_sz_bytes) + + +def mnnvl_lamport_all_gather( + fa: int, + inp: torch.Tensor, + out: torch.Tensor, + local_buffer: int, + multicast_buffer: int, + epoch_buffer: int, + stage_sz_bytes: int, +) -> None: + torch.ops._C_custom_ar.mnnvl_lamport_all_gather( + fa, + inp, + out, + local_buffer, + multicast_buffer, + epoch_buffer, + stage_sz_bytes, + ) + + +def custom_reduce_scatter( + fa: int, + inp: torch.Tensor, + out: torch.Tensor, + reg_buffer: int, + reg_buffer_sz_bytes: int, +) -> None: + torch.ops._C_custom_ar.custom_reduce_scatter(fa, inp, out, reg_buffer, reg_buffer_sz_bytes) + + +def mnnvl_lamport_reduce_scatter( + fa: int, + inp: torch.Tensor, + out: torch.Tensor, + local_buffer: int, + epoch_buffer: int, + stage_sz_bytes: int, +) -> None: + torch.ops._C_custom_ar.mnnvl_lamport_reduce_scatter(fa, inp, out, local_buffer, epoch_buffer, stage_sz_bytes) + + def dispose(fa: int) -> None: torch.ops._C_custom_ar.dispose(fa) @@ -3698,18 +3790,18 @@ def dsv3_fused_a_gemm( output: torch.Tensor, mat_a: torch.Tensor, mat_b: torch.Tensor, + enable_pdl: bool = False, ) -> None: - """DeepSeek V3 fused A GEMM (SM 9.0+, bf16 only, 1-16 tokens). + """Low-latency fused-A-style GEMM (SM 9.0+, BF16, 1-16 tokens). - Computes output = mat_a @ mat_b.T where: - mat_a: [num_tokens, 7168] row-major bf16 (hidden states) - mat_b: [7168, 2112] column-major bf16 (weight transposed) - output: [num_tokens, 2112] row-major bf16 + Computes ``output = mat_a @ mat_b`` for the compiled Kimi K3 and + DeepSeek V3 projection shapes. ``mat_a`` and ``output`` are row-major; + ``mat_b`` is the column-major transposed weight. ``enable_pdl`` permits + programmatic dependent launch for callers that have validated it. - Optimized for the DeepSeek V2/V3 QKV A-projection at small batch sizes. - Requires SM 9.0+ (Hopper). + Requires SM 9.0+. """ - torch.ops._C.dsv3_fused_a_gemm(output, mat_a, mat_b) + torch.ops._C.dsv3_fused_a_gemm(output, mat_a, mat_b, enable_pdl) if hasattr(torch.ops._C, "weight_packed_linear"): @@ -4237,6 +4329,7 @@ def cpu_attn_get_scheduler_metadata( isa: str, enable_kv_split: bool, dynamic_causal: torch.Tensor | None = None, + kv_cache_dtype: str = "auto", ) -> torch.Tensor: scheduler_metadata = torch.ops._C.get_scheduler_metadata( num_reqs, @@ -4251,6 +4344,7 @@ def cpu_attn_get_scheduler_metadata( isa, enable_kv_split, dynamic_causal, + kv_cache_dtype, ) return scheduler_metadata @@ -4503,19 +4597,6 @@ def fusedQuantizeMx( raise ValueError(f"invalid method {method!r}, must be 'quest' or 'abs_max'") -if hasattr(torch.ops._qutlass_C, "fusedQuantizeNvAbsMax"): - - @register_fake("_qutlass_C::fusedQuantizeNvAbsMax") - def _fake_fused_quantize_nv_absmax( - a: torch.Tensor, - b: torch.Tensor, - xh_e2m1: torch.Tensor, - xh_e4m3: torch.Tensor, - global_scale: torch.Tensor, - ): - return xh_e2m1, xh_e4m3 - - def fusedQuantizeNv(a: torch.Tensor, b: torch.Tensor, global_scale: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: xh_e2m1 = torch.empty(*a.shape[:-1], a.size(-1) // 2, dtype=torch.uint8, device=a.device) @@ -4526,7 +4607,33 @@ def fusedQuantizeNv(a: torch.Tensor, b: torch.Tensor, global_scale: torch.Tensor padded_cols = n_col_blocks * 4 xh_e4m3 = torch.empty(padded_rows, padded_cols, dtype=torch.float8_e4m3fn, device=a.device) - return torch.ops._qutlass_C.fusedQuantizeNvAbsMax(a, b, xh_e2m1, xh_e4m3, global_scale) + safeFusedQuantizeNv(a, b, xh_e2m1, xh_e4m3, global_scale) + return xh_e2m1, xh_e4m3 + + +@torch.library.custom_op("aphrodite::safeFusedQuantizeNv", mutates_args=("xh_e2m1", "xh_e4m3")) +def safeFusedQuantizeNv( + a: torch.Tensor, + b: torch.Tensor, + xh_e2m1: torch.Tensor, + xh_e4m3: torch.Tensor, + global_scale: torch.Tensor, +) -> None: + """Call QUTLASS fused quantization with explicitly mutated output tensors.""" + torch.ops._qutlass_C.fusedQuantizeNvAbsMax(a, b, xh_e2m1, xh_e4m3, global_scale) + + +if hasattr(torch.ops._qutlass_C, "fusedQuantizeNv"): + + @register_fake("aphrodite::safeFusedQuantizeNv") + def _fake_fused_quantize_nv( + a: torch.Tensor, + b: torch.Tensor, + xh_e2m1: torch.Tensor, + xh_e4m3: torch.Tensor, + global_scale: torch.Tensor, + ) -> None: + return def hadacore_transform(x: torch.Tensor, inplace: bool = True) -> torch.Tensor: diff --git a/aphrodite/config/model.py b/aphrodite/config/model.py index fc8595d9cd..a8a4d6363e 100644 --- a/aphrodite/config/model.py +++ b/aphrodite/config/model.py @@ -81,7 +81,16 @@ RunnerOption = Literal["auto", "generate", "pooling", "draft"] ConvertType = Literal["none", "embed", "classify"] ConvertOption = Literal["auto", ConvertType] -TokenizerMode = Literal["auto", "hf", "slow", "mistral", "deepseek_v32", "deepseek_v4", "inkling"] +TokenizerMode = Literal[ + "auto", + "hf", + "slow", + "mistral", + "deepseek_v32", + "deepseek_v4", + "inkling", + "kimi_k3", +] ModelDType = Literal["auto", "half", "float16", "bfloat16", "float", "float32"] LogprobsMode = Literal["raw_logits", "raw_logprobs", "processed_logits", "processed_logprobs"] PROCESSED_LOGPROBS_MODES: tuple[LogprobsMode, ...] = ( @@ -133,6 +142,8 @@ class ModelConfig: - "mistral" will always use the tokenizer from `mistral_common`. - "deepseek_v32" will always use the tokenizer from `deepseek_v32`. - "deepseek_v4" will always use the tokenizer from `deepseek_v4`. + - "kimi_k3" will always use the "hf" tokenizer but render chat prompts + with Kimi K3's Python XTML encoding instead of a Jinja template. - Other custom values can be supported via plugins. To swap the Rust BPE backend that powers HF fast tokenizers for the @@ -375,7 +386,7 @@ class ModelConfig: interleave_mm_strings: InitVar[bool | None] = cast(InitVar[bool | None], None) skip_mm_profiling: InitVar[bool | None] = cast(InitVar[bool | None], None) video_pruning_rate: InitVar[float | None] = cast(InitVar[float | None], None) - video_pruning_method: InitVar[str | None] = None + video_pruning_method: InitVar[str | None] = cast(InitVar[str | None], None) mm_tensor_ipc: InitVar[MMTensorIPC | None] = cast(InitVar[MMTensorIPC | None], None) mm_ipc_gpu_memory_gb: InitVar[float | None] = cast(InitVar[float | None], None) @@ -615,6 +626,8 @@ def __post_init__( self.tokenizer_mode = "terratorch" elif arch == "MoonshotKimiaForCausalLM": self.tokenizer_mode = "kimi_audio" + elif arch == "KimiK3ForConditionalGeneration": + self.tokenizer_mode = "kimi_k3" elif arch == "DeepseekV32ForCausalLM": self.tokenizer_mode = "deepseek_v32" elif arch == "DeepseekV4ForCausalLM": diff --git a/aphrodite/config/speculative.py b/aphrodite/config/speculative.py index eed596d16c..24f8a12ffc 100644 --- a/aphrodite/config/speculative.py +++ b/aphrodite/config/speculative.py @@ -822,6 +822,15 @@ def __post_init__(self): else: raise NotImplementedError(f"Unsupported speculative method: '{self.method}'") + if self.method in ("eagle", "eagle3"): + # EAGLE drafts share the target's positional space; a + # draft checkpoint with a smaller max_position_embeddings + # than the target under-sizes its rotary cache (#48894). + SpeculativeConfig._maybe_override_draft_max_position_embeddings( + self.draft_model_config.hf_config, + self.target_model_config.max_model_len, + ) + # Replace hf_config for EAGLE draft_model if self.method in ("eagle", "eagle3", "dflash"): from aphrodite.transformers_utils.configs.eagle import EAGLEConfig @@ -992,6 +1001,34 @@ def _maybe_override_draft_max_model_len( ) return result + @staticmethod + def _maybe_override_draft_max_position_embeddings( + draft_hf_config: PretrainedConfig, + target_max_model_len: int, + ) -> None: + """Raise an EAGLE draft's max position embeddings to the target's. + + The proposer feeds the draft positions up to the target's + max_model_len, while max_position_embeddings sizes the draft's rotary + cache. A smaller checkpoint value makes that cache gather go out of + bounds (#48894). + + Args: + draft_hf_config: The draft model's HF config, mutated in place. + target_max_model_len: The target model's max_model_len. + """ + draft_max_position_embeddings = getattr(draft_hf_config, "max_position_embeddings", None) + if draft_max_position_embeddings is None or draft_max_position_embeddings >= target_max_model_len: + return + logger.info( + "Overriding draft model max_position_embeddings from %d to the " + "target model's max_model_len (%d); EAGLE drafts share the " + "target's positional space.", + draft_max_position_embeddings, + target_max_model_len, + ) + draft_hf_config.max_position_embeddings = target_max_model_len + @staticmethod def _verify_and_get_draft_tp( target_parallel_config: ParallelConfig, diff --git a/aphrodite/distributed/device_communicators/shm_broadcast.py b/aphrodite/distributed/device_communicators/shm_broadcast.py index fa3e08c1e7..8657b803db 100644 --- a/aphrodite/distributed/device_communicators/shm_broadcast.py +++ b/aphrodite/distributed/device_communicators/shm_broadcast.py @@ -1,6 +1,8 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import copyreg import functools +import io import os import pickle import shutil @@ -378,6 +380,66 @@ def get_metadata(self, current_idx: int): yield buf +def _rebuild_tensor(buf: Any, shape: tuple[int, ...], dtype_str: str) -> torch.Tensor: + """Rebuild a tensor from an out-of-band pickle buffer. + + Counterpart of `_reduce_tensor`. Note that pickle passes the original + buffer-providing object from `loads(buffers=...)` straight to this + function (no `PickleBuffer` wrapper on the receiving side), so `buf` is + a `zmq.Frame`, a `memoryview` of a shared-memory ring chunk, or `bytes` + if the buffer was serialized in-band. + """ + dtype = getattr(torch, dtype_str) + assert isinstance(dtype, torch.dtype) + if isinstance(buf, zmq.Frame): + # ZMQ frames own their message memory independently of any context, + # so the tensor can safely alias it with zero copies. The tensor's + # storage keeps the frame (and thus its bytes) alive via a strong + # reference for as long as the tensor is. + try: + return torch.frombuffer(buf, dtype=torch.uint8).view(dtype).view(shape) + except ValueError: + # Empty or read-only frame buffer; fall through to the copy path. + pass + # Shared-memory ring buffer chunks are reused by the writer once all + # readers have marked them read, so we must copy out of them. bytearray + # (vs bytes) keeps the resulting tensor writable, matching normal tensor + # semantics. + raw = bytearray(buf) + if not raw: + assert 0 in shape + return torch.empty(shape, dtype=dtype) + return torch.frombuffer(raw, dtype=torch.uint8).view(dtype).view(shape) + + +def _reduce_tensor(tensor: torch.Tensor): + """Reduce a CPU tensor to a `PickleBuffer` for out-of-band pickling. + + `torch.Tensor.__reduce_ex__` copies the tensor bytes into the pickle + byte stream via `torch.serialization` and never emits a `PickleBuffer`, + which defeats the out-of-band buffer handling in `MessageQueue.enqueue`. + This reducer instead exposes the tensor's memory directly, so large + tensors (e.g. `prompt_embeds` in `SchedulerOutput`) traverse the queue + without being copied into and back out of the pickled message. + """ + if tensor.device.type == "cpu" and tensor.layout == torch.strided and not tensor.requires_grad: + try: + # The uint8 view exposes the raw bytes via the buffer protocol, + # including for dtypes numpy doesn't recognize (bfloat16, fp8, ...). + # reshape(-1) first so that 0-dim tensors can be viewed as well. + raw = tensor.contiguous().reshape(-1).view(torch.uint8).numpy() + except RuntimeError: + # Exotic tensors (e.g. with the conjugate bit set) that don't + # support aliasing views; let torch handle them. + pass + else: + dtype_str = str(tensor.dtype).removeprefix("torch.") + return _rebuild_tensor, (PickleBuffer(raw), tuple(tensor.shape), dtype_str) + + # Fall back to torch's default (copying) reduction. + return tensor.__reduce_ex__(pickle.HIGHEST_PROTOCOL) + + @dataclass class Handle: local_reader_ranks: list[int] = field(default_factory=list) @@ -756,7 +818,23 @@ def oob_callback(buf: PickleBuffer) -> bool: total_bytes += len(raw_buf) + 4 return False - all_buffers[0] = pickle.dumps(obj, protocol=pickle.HIGHEST_PROTOCOL, buffer_callback=oob_callback) + # CPU tensors are routed through `_reduce_tensor` so that their + # bytes are emitted as out-of-band buffers instead of being + # copied into the pickle stream by torch's default reducer. + # Start from `copyreg.dispatch_table` to preserve globally + # registered reducers (e.g. `re.Pattern`); the per-pickler + # dispatch table would otherwise shadow them. + dispatch_table = dict(copyreg.dispatch_table) + dispatch_table[torch.Tensor] = _reduce_tensor + with io.BytesIO() as bio: + pickler = pickle.Pickler( + bio, + protocol=pickle.HIGHEST_PROTOCOL, + buffer_callback=oob_callback, + ) + pickler.dispatch_table = dispatch_table + pickler.dump(obj) + all_buffers[0] = bio.getvalue() if self.n_local_reader > 0: if total_bytes + len(all_buffers[0]) >= self.buffer.max_chunk_bytes: with self.acquire_write(timeout) as buf: diff --git a/aphrodite/distributed/ec_transfer/ec_connector/base.py b/aphrodite/distributed/ec_transfer/ec_connector/base.py index e16bc82079..6adb8f7cf6 100644 --- a/aphrodite/distributed/ec_transfer/ec_connector/base.py +++ b/aphrodite/distributed/ec_transfer/ec_connector/base.py @@ -263,3 +263,15 @@ def request_finished(self, request: "Request") -> tuple[bool, dict[str, Any] | N get_finished(). """ return False, None + + def has_pending_push_work(self) -> bool: + """Return True if the connector has push-mode work that requires + the engine main loop to keep stepping (e.g. for EPD, + Producer has push work when Xfer is in progress - Consumer + is reading it). + This mirrors exactly the KV Connector's has_pending_push_work(). + + Connectors that don't implement push-based EC transfer should + leave this as False. + """ + return False 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 034798decc..0cd7468407 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 @@ -977,7 +977,7 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]): caches_data = [] # With hybrid allocator, layers can share a kv cache tensor - seen_base_addresses = [] + seen_base_addresses: list[int] = [] # K and V are packed into the content dim, so each attention layer is a # single NIXL region whose block transfers as one unit. Mamba layers instead @@ -1016,11 +1016,17 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]): curr_tensor_size_bytes = num_blocks * physical_page_size base_addr = cache.data_ptr() + is_mla_region = isinstance(layer_spec, (MLAAttentionSpec, SlidingWindowMLASpec)) if base_addr in seen_base_addresses: # NOTE (NickLucche) HMA employs memory pooling to share tensors # across groups. This results in skipping all tensors but the ones # pointed to by group0. Also, generally we will have more blocks # per tensor but fewer regions. + # A shared tensor may back both SSM and attention layers (e.g. + # KDA+MLA in KimiLinear); the region's FA view is MLA whichever + # layer registered it first. + idx = seen_base_addresses.index(base_addr) + self._region_is_mla[idx] |= is_mla_region logger.debug("Skipping %s because it's already seen", layer_name) continue logger.debug("Registering layer %s with cache shape: %s", layer_name, cache.shape) @@ -1030,7 +1036,6 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]): self.block_len_per_layer.append(physical_page_size // self._physical_blocks_per_logical_kv_block) else: self.block_len_per_layer.append(physical_page_size) - is_mla_region = isinstance(layer_spec, (MLAAttentionSpec, SlidingWindowMLASpec)) self._region_is_mla.append(is_mla_region) if not is_mla_region: @@ -1607,7 +1612,22 @@ def _validate_remote_agent_handshake(self, nixl_agent_meta: NixlAgentMetadata, r # the per-rank KV head ratio rather than the raw tp_ratio, because GQA # replication caps per-rank heads at 1 when tp > total_kv_heads # (issue #45330). Mamba uses the ssm_sizes counterpart, so skip here. - if not self._has_mamba: + if self._has_mamba and self.use_mla: + # Hybrid MLA+SSM (e.g. KimiLinear's KDA+MLA): regions are + # kernel-granularity views of the mamba-unified page. The MLA + # per-token page is TP-independent, so the block lengths must + # match up to the kernel block size ratio even under + # heterogeneous TP (remote kernel blocks may be smaller). + # SSM geometry is validated via ssm_sizes/conv offsets instead. + assert self.block_len_per_layer == [ + block_len * block_size_ratio for block_len in nixl_agent_meta.block_lens + ], ( + "Hybrid MLA kernel-granularity block lengths must match " + f"between P and D (block_size_ratio={block_size_ratio}): " + f"local={self.block_len_per_layer}, " + f"remote={nixl_agent_meta.block_lens}." + ) + elif not self._has_mamba: assert len(self.block_len_per_layer) == len(nixl_agent_meta.block_lens), ( "Number of KV layers must match between prefill and decode" ) diff --git a/aphrodite/distributed/kv_transfer/kv_connector/v1/nixl/push_worker.py b/aphrodite/distributed/kv_transfer/kv_connector/v1/nixl/push_worker.py index b43ab23e1c..474d15a629 100644 --- a/aphrodite/distributed/kv_transfer/kv_connector/v1/nixl/push_worker.py +++ b/aphrodite/distributed/kv_transfer/kv_connector/v1/nixl/push_worker.py @@ -53,7 +53,10 @@ ReqMeta, TransferHandle, ) -from aphrodite.distributed.kv_transfer.kv_connector.v1.nixl.tp_mapping import ReadSpec +from aphrodite.distributed.kv_transfer.kv_connector.v1.nixl.tp_mapping import ( + ReadSpec, + _is_attention_spec, +) from aphrodite.distributed.kv_transfer.kv_connector.v1.nixl.utils import get_base_request_id from aphrodite.logger import init_logger @@ -498,37 +501,37 @@ def _xfer_blocks_for_req(self, req_id: str, meta: ReqMeta): local_block_ids = meta.local_physical_block_ids num_groups = len(local_block_ids) - if self.use_mla and tp_ratio < 0: - # MLA latent is replicated across D's TP ranks: the tp-mapping - # collapses to one rank (fine for reads), but push must WRITE every - # D rank or the rest decode stale KV; only the dst differs per rank. + # MLA latent is replicated across D's TP ranks: the tp-mapping + # collapses it to one rank (fine for reads), but push must WRITE every + # D rank or the rest decode stale KV. For hybrid MLA+SSM the sharded + # SSM state already targets every covered D rank, so only the + # attention groups need widening; pure MLA writes to all handshaked + # ranks (only the dst differs per rank). + replicate_attn = self.use_mla and tp_ratio < 0 + if replicate_attn and not self._has_mamba: assert len(plan.all_source_ranks) == 1 - mla_local_ids = [list(ids) for ids in local_block_ids] - mla_remote_ids = [list(ids) for ids in remote_block_ids] - read_specs = [ - ReadSpec( - remote_rank=rank, - local_block_ids=mla_local_ids, - remote_block_ids=mla_remote_ids, - ) - for rank in self.dst_xfer_side_handles[engine_id] - ] + write_ranks = sorted(self.dst_xfer_side_handles[engine_id]) else: - read_specs = [ - ReadSpec( - remote_rank=rank, - local_block_ids=[ - list(local_block_ids[g]) if rank in plan.source_ranks_per_group[g] else [] - for g in range(num_groups) - ], - remote_block_ids=[ - list(remote_block_ids[g]) if rank in plan.source_ranks_per_group[g] else [] - for g in range(num_groups) - ], - ) - for rank in plan.all_source_ranks + write_ranks = list(plan.all_source_ranks) + + def group_ids(block_ids: BlockIds, rank: int) -> BlockIds: + return [ + list(block_ids[g]) + if (replicate_attn and _is_attention_spec(self._group_spec_types[g])) + or rank in plan.source_ranks_per_group[g] + else [] + for g in range(num_groups) ] + read_specs = [ + ReadSpec( + remote_rank=rank, + local_block_ids=group_ids(local_block_ids, rank), + remote_block_ids=group_ids(remote_block_ids, rank), + ) + for rank in write_ranks + ] + handles: list[int] = [] for i, spec in enumerate(read_specs): remote_block_size = remote_info.remote_block_size @@ -540,7 +543,10 @@ def _xfer_blocks_for_req(self, req_id: str, meta: ReqMeta): remote_block_size, req_id, ) - if tp_ratio < 0 and not self.use_mla: + if tp_ratio < 0 and (not self.use_mla or len(plan.all_source_ranks) > 1): + # Multiple targets: write each rank its chunk of local memory. + # Hybrid MLA+SSM also lands here: its split handles replicate + # the attention descriptors and chunk only the SSM state. split_key = (tp_ratio, remote_block_size) local_xfer_side_handle = self.src_xfer_handles_by_tp_ratio[split_key][i] else: diff --git a/aphrodite/entrypoints/chat_utils.py b/aphrodite/entrypoints/chat_utils.py index dd68ccd4f8..c4aca9c6da 100644 --- a/aphrodite/entrypoints/chat_utils.py +++ b/aphrodite/entrypoints/chat_utils.py @@ -1878,13 +1878,19 @@ def get_history_tool_calls_cnt(conversation: list[ConversationMessage]): return idx -_KIMI_MODEL_TYPES = ("kimi_k2", "kimi_k25") +_KIMI_MODEL_TYPES = ("kimi_k2", "kimi_k25", "kimi_k3") def get_tool_call_id_type(model_config: ModelConfig) -> str: """Return the tool-call ID type for a given model configuration.""" hf_overrides = getattr(model_config, "hf_overrides", None) - if model_config.hf_text_config.model_type in _KIMI_MODEL_TYPES or ( + hf_config = getattr(model_config, "hf_config", None) + hf_text_config = getattr(model_config, "hf_text_config", None) + model_types = ( + getattr(hf_config, "model_type", None), + getattr(hf_text_config, "model_type", None), + ) + if any(model_type in _KIMI_MODEL_TYPES for model_type in model_types) or ( isinstance(hf_overrides, dict) and hf_overrides.get("model_type") in _KIMI_MODEL_TYPES ): return "kimi_k2" diff --git a/aphrodite/entrypoints/openai/api_server.py b/aphrodite/entrypoints/openai/api_server.py index c098b8ff2e..031e34b85e 100644 --- a/aphrodite/entrypoints/openai/api_server.py +++ b/aphrodite/entrypoints/openai/api_server.py @@ -28,7 +28,6 @@ from aphrodite.entrypoints.chat_utils import load_chat_template from aphrodite.entrypoints.launcher import serve_http from aphrodite.entrypoints.openai.cli_args import make_arg_parser, validate_parsed_serve_args -from aphrodite.entrypoints.openai.engine.protocol import GenerationError from aphrodite.entrypoints.openai.models.protocol import BaseModelPath from aphrodite.entrypoints.openai.models.serving import OpenAIServingModels from aphrodite.entrypoints.serve.elastic_ep.middleware import ScalingMiddleware @@ -42,20 +41,15 @@ ) from aphrodite.entrypoints.serve.utils.request_logger import RequestLogger from aphrodite.entrypoints.serve.utils.server_utils import ( - engine_error_handler, + aphrodite_error_handler, exception_handler, - generation_error_handler, get_uvicorn_log_config, http_exception_handler, lifespan, log_response, validation_exception_handler, ) -from aphrodite.exceptions import ( - APHRODITENotFoundError, - APHRODITEUnprocessableEntityError, - APHRODITEValidationError, -) +from aphrodite.exceptions import APHRODITEError from aphrodite.logger import init_logger from aphrodite.reasoning import ReasoningParserManager from aphrodite.renderers.online_derenderer import OnlineDerenderer @@ -67,7 +61,6 @@ from aphrodite.utils.argparse_utils import FlexibleArgumentParser from aphrodite.utils.network_utils import is_valid_ipv6_address from aphrodite.utils.system_utils import decorate_logs, set_ulimit -from aphrodite.v1.engine.exceptions import EngineDeadError, EngineGenerateError from aphrodite.version import __version__ as APHRODITE_VERSION prometheus_multiproc_dir: tempfile.TemporaryDirectory @@ -284,17 +277,8 @@ def build_app( app.exception_handler(HTTPException)(http_exception_handler) app.exception_handler(RequestValidationError)(validation_exception_handler) - app.exception_handler(EngineGenerateError)(engine_error_handler) - app.exception_handler(EngineDeadError)(engine_error_handler) - app.exception_handler(GenerationError)(generation_error_handler) - # Register specific exception types so they are handled by - # ExceptionMiddleware (inside the Prometheus middleware) rather than - # ServerErrorMiddleware (outside it). Without this, these exceptions - # propagate through Prometheus as unhandled and get recorded as 5xx - # even though they result in 4xx responses to the client. - app.exception_handler(APHRODITEValidationError)(exception_handler) - app.exception_handler(APHRODITEUnprocessableEntityError)(exception_handler) - app.exception_handler(APHRODITENotFoundError)(exception_handler) + app.exception_handler(APHRODITEError)(aphrodite_error_handler) + # Fallback handlers for raw exceptions not yet migrated to AphroditeError. app.exception_handler(ValueError)(exception_handler) app.exception_handler(TypeError)(exception_handler) app.exception_handler(OverflowError)(exception_handler) diff --git a/aphrodite/entrypoints/openai/chat_completion/protocol.py b/aphrodite/entrypoints/openai/chat_completion/protocol.py index 29151d7f00..21871fa651 100644 --- a/aphrodite/entrypoints/openai/chat_completion/protocol.py +++ b/aphrodite/entrypoints/openai/chat_completion/protocol.py @@ -612,6 +612,12 @@ def build_chat_params( ), media_io_kwargs=self.media_io_kwargs, return_assistant_tokens_mask=bool(self.return_assistant_tokens_mask), + # No-tools requests default to tool_choice="none" at the API + # layer. Collapse that default before rendering, so K3 emits a + # model-visible tool-choice instruction only for requests with a + # tools block. + tool_choice=self.tool_choice if self.tools else None, + response_format=self.response_format, ) def build_tok_params(self, model_config: ModelConfig) -> TokenizeParams: diff --git a/aphrodite/entrypoints/openai/engine/protocol.py b/aphrodite/entrypoints/openai/engine/protocol.py index 9e7901775e..8021ecba8d 100644 --- a/aphrodite/entrypoints/openai/engine/protocol.py +++ b/aphrodite/entrypoints/openai/engine/protocol.py @@ -19,7 +19,7 @@ from aphrodite.config.utils import replace from aphrodite.entrypoints.chat_utils import make_tool_call_id -from aphrodite.exceptions import APHRODITEValidationError +from aphrodite.exceptions import APHRODITEServerError, APHRODITEValidationError from aphrodite.logger import init_logger from aphrodite.sampling_params import StructuredOutputsParams from aphrodite.utils import random_uuid @@ -387,7 +387,7 @@ def _serialize(self, handler): return data -class GenerationError(Exception): +class GenerationError(APHRODITEServerError): """raised when finish_reason indicates internal server error (500)""" def __init__(self, message: str = "Internal server error"): diff --git a/aphrodite/entrypoints/openai/responses/protocol.py b/aphrodite/entrypoints/openai/responses/protocol.py index 7da5812b84..507250eca5 100644 --- a/aphrodite/entrypoints/openai/responses/protocol.py +++ b/aphrodite/entrypoints/openai/responses/protocol.py @@ -329,6 +329,7 @@ def build_chat_params( extra_kwargs, ), media_io_kwargs=self.media_io_kwargs, + tool_choice=self.tool_choice if self.tools else None, ) def build_tok_params(self, model_config: ModelConfig) -> TokenizeParams: diff --git a/aphrodite/entrypoints/serve/utils/error_response.py b/aphrodite/entrypoints/serve/utils/error_response.py index e2b988b3f9..a41810b690 100644 --- a/aphrodite/entrypoints/serve/utils/error_response.py +++ b/aphrodite/entrypoints/serve/utils/error_response.py @@ -26,7 +26,9 @@ def create_error_response( logger.debug("create_error_response called with %s: %s", type(exc).__name__, exc) from aphrodite.exceptions import ( + APHRODITEClientError, APHRODITENotFoundError, + APHRODITEServerError, APHRODITEUnprocessableEntityError, APHRODITEValidationError, ) @@ -43,8 +45,20 @@ def create_error_response( err_type = "NotFoundError" status_code = HTTPStatus.NOT_FOUND param = None + elif isinstance(exc, APHRODITEClientError): + err_type = "BadRequestError" + status_code = HTTPStatus.BAD_REQUEST + param = None + elif isinstance(exc, GenerationError): + err_type = "InternalServerError" + status_code = exc.status_code + param = None + elif isinstance(exc, APHRODITEServerError): + err_type = "InternalServerError" + status_code = HTTPStatus.INTERNAL_SERVER_ERROR + param = None + # Fallback for raw exceptions not yet migrated to AphroditeError. elif isinstance(exc, (ValueError, TypeError, OverflowError)): - # Common validation errors from user input err_type = "BadRequestError" status_code = HTTPStatus.BAD_REQUEST param = None @@ -52,10 +66,6 @@ def create_error_response( err_type = "NotImplementedError" status_code = HTTPStatus.NOT_IMPLEMENTED param = None - elif isinstance(exc, GenerationError): - err_type = "InternalServerError" - status_code = exc.status_code - param = None elif any(cls.__name__ == "TemplateError" for cls in type(exc).__mro__): # jinja2.TemplateError and its subclasses (avoid importing jinja2) err_type = "BadRequestError" diff --git a/aphrodite/entrypoints/serve/utils/server_utils.py b/aphrodite/entrypoints/serve/utils/server_utils.py index 9f1df0451c..e434e122fb 100644 --- a/aphrodite/entrypoints/serve/utils/server_utils.py +++ b/aphrodite/entrypoints/serve/utils/server_utils.py @@ -31,7 +31,7 @@ create_error_response, sanitize_message, ) -from aphrodite.exceptions import APHRODITEValidationError +from aphrodite.exceptions import APHRODITEError, APHRODITEValidationError from aphrodite.logger import init_logger from aphrodite.utils.gc_utils import freeze_gc_heap from aphrodite.v1.engine.exceptions import EngineDeadError, EngineGenerateError @@ -316,6 +316,15 @@ async def log_response(request: Request, call_next): return response +async def aphrodite_error_handler(req: Request, exc: APHRODITEError): + """Dispatch an Aphrodite-specific error to the appropriate handler.""" + if isinstance(exc, (EngineGenerateError, EngineDeadError)): + return await engine_error_handler(req, exc) + if isinstance(exc, GenerationError): + return await generation_error_handler(req, exc) + return await exception_handler(req, exc) + + async def engine_error_handler(req: Request, exc: EngineDeadError | EngineGenerateError): """ APHRODITE V1 AsyncLLM catches exceptions and returns diff --git a/aphrodite/envs.py b/aphrodite/envs.py index 5d2ebee86a..89cf5c0ab0 100755 --- a/aphrodite/envs.py +++ b/aphrodite/envs.py @@ -270,6 +270,7 @@ APHRODITE_NCCL_INCLUDE_PATH: str | None = None APHRODITE_GC_DEBUG: str = "" APHRODITE_DEBUG_WORKSPACE: bool = False + APHRODITE_ENABLE_K3_LATENT_MOE_TAIL_FUSION: bool = False APHRODITE_DISABLE_SHARED_EXPERTS_STREAM: bool = False APHRODITE_SHARED_EXPERTS_STREAM_TOKEN_THRESHOLD: int = 256 APHRODITE_MULTI_STREAM_GEMM_TOKEN_THRESHOLD: int = 1024 @@ -1665,6 +1666,11 @@ def _resolve_rust_cli_path() -> str | None: # Debug workspace allocations. # logging of workspace resize operations. "APHRODITE_DEBUG_WORKSPACE": lambda: bool(int(os.getenv("APHRODITE_DEBUG_WORKSPACE", "0"))), + # Enable the experimental Kimi K3 latent-MoE tail fusion. + # Currently supported only on SM100 with TP=8/16 and BF16. + "APHRODITE_ENABLE_K3_LATENT_MOE_TAIL_FUSION": lambda: bool( + int(os.getenv("APHRODITE_ENABLE_K3_LATENT_MOE_TAIL_FUSION", "0")) + ), # Disables parallel execution of shared_experts via separate cuda stream "APHRODITE_DISABLE_SHARED_EXPERTS_STREAM": lambda: bool( int(os.getenv("APHRODITE_DISABLE_SHARED_EXPERTS_STREAM", "0")) diff --git a/aphrodite/exceptions.py b/aphrodite/exceptions.py index ab495f5da0..50b2b6af83 100644 --- a/aphrodite/exceptions.py +++ b/aphrodite/exceptions.py @@ -6,7 +6,19 @@ from typing import Any -class AphroditeValidationError(ValueError): +class AphroditeError(Exception): + """Base class for all Aphrodite-specific errors.""" + + +class AphroditeClientError(AphroditeError): + """Base class for errors caused by the client request (4xx).""" + + +class AphroditeServerError(AphroditeError): + """Base class for errors caused by the server (5xx).""" + + +class AphroditeValidationError(AphroditeClientError): """Aphrodite-specific validation error for request validation failures. Args: @@ -36,7 +48,7 @@ def __str__(self): return f"{base} ({', '.join(extras)})" if extras else base -class AphroditeNotFoundError(Exception): +class AphroditeNotFoundError(AphroditeClientError): """Aphrodite-specific NotFoundError""" pass @@ -66,7 +78,7 @@ def __str__(self): return self.message -class AphroditeUnprocessableEntityError(ValueError): +class AphroditeUnprocessableEntityError(AphroditeClientError): """Aphrodite-specific error for unprocessable entity requests. Args: @@ -97,6 +109,9 @@ def __str__(self): # Backward compatibility with older upstream-derived imports. +APHRODITEError = AphroditeError +APHRODITEClientError = AphroditeClientError +APHRODITEServerError = AphroditeServerError APHRODITEValidationError = AphroditeValidationError APHRODITENotFoundError = AphroditeNotFoundError APHRODITEUnprocessableEntityError = AphroditeUnprocessableEntityError diff --git a/aphrodite/inputs/engine.py b/aphrodite/inputs/engine.py index fcdee6b6b0..743947b596 100644 --- a/aphrodite/inputs/engine.py +++ b/aphrodite/inputs/engine.py @@ -7,6 +7,8 @@ from typing_extensions import NotRequired, TypedDict, assert_never +from aphrodite.exceptions import APHRODITEValidationError + if TYPE_CHECKING: import torch @@ -284,7 +286,7 @@ class EncoderDecoderInput(TypedDict): def _validate_enc_input(enc_input: SingletonInput) -> EncoderInput: if enc_input["type"] == "embeds": - raise ValueError("Embedding inputs are not supported for encoder-decoder models") + raise APHRODITEValidationError("Embedding inputs are not supported for encoder-decoder models") if enc_input["type"] == "multimodal" and "encoder_prompt_token_ids" not in enc_input: raise RuntimeError("You should register an encoder-decoder multi-modal processor for encoder-decoder models.") @@ -294,7 +296,7 @@ def _validate_enc_input(enc_input: SingletonInput) -> EncoderInput: def _validate_dec_input(dec_input: SingletonInput) -> DecoderEngineInput: if dec_input["type"] == "embeds": - raise ValueError("Embedding inputs are not supported for encoder-decoder models") + raise APHRODITEValidationError("Embedding inputs are not supported for encoder-decoder models") return dec_input diff --git a/aphrodite/model_executor/kernels/linear/cute_dsl/_skinny_gemm.py b/aphrodite/model_executor/kernels/linear/cute_dsl/_skinny_gemm.py new file mode 100644 index 0000000000..51a69fe133 --- /dev/null +++ b/aphrodite/model_executor/kernels/linear/cute_dsl/_skinny_gemm.py @@ -0,0 +1,176 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from __future__ import annotations + +import cutlass +import cutlass.cute as cute +from cuda.bindings.driver import CUstream +from cutlass import const_expr + + +class CuteSkinnyGemm: + """Shape-dynamic low-latency GEMM for small token counts. + + Computes ``C[M, N] = A[M, K] @ B[N, K].T + residual`` with BF16 or FP16 + inputs, FP32 accumulators, and an output matching the input dtype. The + residual term is optional and is added before the output conversion. N and + K are runtime values. The tiny M dimension is fully unrolled alongside a + small set of tuning parameters. + """ + + def __init__( + self, + *, + element_type, + num_rows: int, + block_size: int, + outputs_per_block: int, + vector_width: int = 8, + k_unroll: int = 1, + has_residual: bool = False, + use_pdl: bool = False, + ) -> None: + if block_size % cute.arch.WARP_SIZE != 0: + raise ValueError("block_size must be a multiple of the warp size") + self.element_type = element_type + self.num_rows = num_rows + self.block_size = block_size + self.outputs_per_block = outputs_per_block + self.vector_width = vector_width + self.k_unroll = k_unroll + self.has_residual = has_residual + self.use_pdl = use_pdl + self.num_warps = block_size // cute.arch.WARP_SIZE + + @cute.jit + def __call__( + self, + gA: cute.Tensor, + gB: cute.Tensor, + gResidual: cute.Tensor, + gC: cute.Tensor, + stream: CUstream, + ) -> None: + n = cute.size(gB, mode=[0]) + k = cute.size(gA, mode=[1]) + copy_a = cute.make_copy_atom( + cute.nvgpu.CopyG2ROp(), + self.element_type, + num_bits_per_copy=self.vector_width * self.element_type.width, + load_cache_mode=cute.nvgpu.LoadCacheMode.ALWAYS, + ) + copy_b = cute.make_copy_atom( + cute.nvgpu.CopyG2ROp(), + self.element_type, + num_bits_per_copy=self.vector_width * self.element_type.width, + load_cache_mode=cute.nvgpu.LoadCacheMode.STREAMING, + ) + self.kernel(gA, gB, gResidual, gC, k, copy_a, copy_b).launch( + grid=[cute.ceil_div(n, self.outputs_per_block), 1, 1], + block=[self.block_size, 1, 1], + smem=self.num_rows * self.outputs_per_block * self.num_warps * 4, + stream=stream, + use_pdl=self.use_pdl, + min_blocks_per_mp=1, + ) + + @cute.kernel + def kernel( + self, + gA: cute.Tensor, + gB: cute.Tensor, + gResidual: cute.Tensor, + gC: cute.Tensor, + k_extent: cutlass.Int32, + copy_a: cute.CopyAtom, + copy_b: cute.CopyAtom, + ) -> None: + tidx, _, _ = cute.arch.thread_idx() + block_idx, _, _ = cute.arch.block_idx() + warp_idx = cute.arch.warp_idx() + + num_rows: cutlass.Constexpr = self.num_rows + outputs_per_block: cutlass.Constexpr = self.outputs_per_block + vector_width: cutlass.Constexpr = self.vector_width + block_size: cutlass.Constexpr = self.block_size + num_warps: cutlass.Constexpr = self.num_warps + + acc_layout = cute.make_layout((num_rows, outputs_per_block), stride=(outputs_per_block, 1)) + acc = cute.make_rmem_tensor(acc_layout, cutlass.Float32) + acc.fill(0.0) + + if const_expr(self.use_pdl): + cute.arch.griddepcontrol_wait() + + n_base = block_idx * outputs_per_block + k_tile_size: cutlass.Constexpr = block_size * vector_width + num_k_tiles = k_extent // k_tile_size + + gA_vec = cute.logical_divide(gA, (None, vector_width)) + gB_vec = cute.logical_divide(gB, (None, vector_width)) + # Layout after both divides is (M/N, K_TILE, K_LANE, K_VEC). + tA_all = cute.logical_divide(gA_vec, (None, (None, block_size))) + tB_all = cute.logical_divide(gB_vec, (None, (None, block_size))) + tA = tA_all[None, (None, (tidx, None))] + + a_regs = cute.make_rmem_tensor( + cute.make_layout((num_rows, vector_width), stride=(vector_width, 1)), + self.element_type, + ) + b_regs = cute.make_rmem_tensor( + cute.make_layout((outputs_per_block, vector_width), stride=(vector_width, 1)), + self.element_type, + ) + + for k_tile in cutlass.range(num_k_tiles, unroll=self.k_unroll): + for mi in cutlass.range_constexpr(num_rows): + cute.copy(copy_a, tA[mi, None, k_tile], a_regs[mi, None]) + + for ni in cutlass.range_constexpr(outputs_per_block): + n_idx = n_base + ni + tB = tB_all[n_idx, (None, (tidx, None))] + cute.copy(copy_b, tB[None, k_tile], b_regs[ni, None]) + + for vi in cutlass.range_constexpr(vector_width): + for mi in cutlass.range_constexpr(num_rows): + for ni in cutlass.range_constexpr(outputs_per_block): + acc[mi, ni] = acc[mi, ni] + a_regs[mi, vi].to(cutlass.Float32) * b_regs[ni, vi].to( + cutlass.Float32 + ) + + for mi in cutlass.range_constexpr(num_rows): + for ni in cutlass.range_constexpr(outputs_per_block): + acc[mi, ni] = cute.arch.warp_reduction_sum(acc[mi, ni]) + + smem_layout = cute.make_layout( + (num_rows, outputs_per_block, num_warps), + stride=(outputs_per_block * num_warps, num_warps, 1), + ) + smem = cutlass.utils.SmemAllocator() + partials = smem.allocate_tensor(cutlass.Float32, smem_layout, byte_alignment=16) + with cute.arch.elect_one(): + for mi in cutlass.range_constexpr(num_rows): + for ni in cutlass.range_constexpr(outputs_per_block): + partials[mi, ni, warp_idx] = acc[mi, ni] + + cute.arch.sync_threads() + if tidx == 0: + for mi in cutlass.range_constexpr(num_rows): + for ni in cutlass.range_constexpr(outputs_per_block): + n_idx = n_base + ni + total = ( + partials[mi, ni, None] + .load() + .reduce( + cute.ReductionOp.ADD, + init_val=cutlass.Float32(0.0), + reduction_profile=0, + ) + ) + if const_expr(self.has_residual): + total += gResidual[mi, n_idx].to(cutlass.Float32) + gC[mi, n_idx] = cutlass.Float32(total).to(self.element_type) + + if const_expr(self.use_pdl): + cute.arch.griddepcontrol_launch_dependents() diff --git a/aphrodite/model_executor/kernels/linear/cute_dsl/skinny_gemm.py b/aphrodite/model_executor/kernels/linear/cute_dsl/skinny_gemm.py new file mode 100644 index 0000000000..b8671d1ed4 --- /dev/null +++ b/aphrodite/model_executor/kernels/linear/cute_dsl/skinny_gemm.py @@ -0,0 +1,246 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from __future__ import annotations + +import logging +from collections.abc import Iterable +from dataclasses import dataclass +from functools import partial +from typing import Any + +import torch + +logger = logging.getLogger(__name__) + +_cutedsl_available: bool | None = None + + +@dataclass(frozen=True, slots=True) +class SkinnyGemmConfig: + num_rows: int + block_size: int + outputs_per_block: int + k_unroll: int = 1 + vector_width: int = 8 + + +class ShapeDynamicSkinnyGemm: + def __init__(self) -> None: + self._compiled: dict[tuple[torch.dtype, SkinnyGemmConfig, bool], Any] = {} + self._warmup_configs: set[tuple[torch.dtype, SkinnyGemmConfig, bool]] = set() + self._warmup_registered = False + + @staticmethod + def is_available() -> bool: + global _cutedsl_available + if _cutedsl_available is not None: + return _cutedsl_available + try: + import cutlass # noqa: F401 + import cutlass.cute # noqa: F401 + + _cutedsl_available = True + except ImportError: + _cutedsl_available = False + logger.info("cuteDSL is not available; skinny GEMM is disabled") + return _cutedsl_available + + @staticmethod + def _config(m: int, n: int, k: int) -> SkinnyGemmConfig: + num_rows = m + wide_block = 224 + if m == 1 and k >= 7168 and k % (wide_block * 8) == 0: + if n % 3 == 0: + k_unroll = 2 if n <= 2304 else 4 + return SkinnyGemmConfig(num_rows, wide_block, 3, k_unroll) + if 2304 < n < 4096 and n % 2 == 0: + return SkinnyGemmConfig(num_rows, wide_block, 2, k_unroll=4) + + if k <= 2048 or k % (128 * 8) != 0: + outputs_per_block = 2 if k <= 2048 else 4 + if n % outputs_per_block: + outputs_per_block = 1 + if k % (64 * 8) == 0: + return SkinnyGemmConfig(num_rows, 64, outputs_per_block, 2) + if k % (32 * 8) == 0: + return SkinnyGemmConfig(num_rows, 32, outputs_per_block, 2) + return SkinnyGemmConfig(num_rows, 32, outputs_per_block, 2, vector_width=4) + + block_size = 64 if 4096 <= n < 8192 else 128 + outputs_per_block = 1 if m == 1 and n <= 2304 else 2 + if n % outputs_per_block: + outputs_per_block = 1 + k_unroll = 2 if n <= 2304 or n >= 16384 else 1 + return SkinnyGemmConfig( + num_rows, + block_size, + outputs_per_block, + k_unroll=k_unroll, + ) + + @staticmethod + def _cutlass_dtype(dtype: torch.dtype): + from cutlass import BFloat16, Float16 + + return BFloat16 if dtype == torch.bfloat16 else Float16 + + @staticmethod + def _stream(): + from cuda.bindings.driver import CUstream + + from aphrodite.utils.torch_utils import current_stream + + return CUstream(current_stream().cuda_stream) + + @staticmethod + def _use_pdl() -> bool: + from aphrodite.platforms import current_platform + + return current_platform.is_arch_support_pdl() + + def _compile( + self, + dtype: torch.dtype, + config: SkinnyGemmConfig, + has_residual: bool, + ) -> None: + import cutlass.cute as cute + from quack.compile_utils import make_fake_tensor + + from ._skinny_gemm import CuteSkinnyGemm + + element_type = self._cutlass_dtype(dtype) + n = cute.sym_int(divisibility=config.outputs_per_block) + k = cute.sym_int(divisibility=config.block_size * config.vector_width) + a = make_fake_tensor( + element_type, + (config.num_rows, k), + divisibility=config.vector_width, + ) + b = make_fake_tensor( + element_type, + (n, k), + divisibility=config.vector_width, + ) + c = make_fake_tensor(element_type, (config.num_rows, n), divisibility=1) + residual = make_fake_tensor(element_type, (config.num_rows, n), divisibility=1) + kernel = CuteSkinnyGemm( + element_type=element_type, + num_rows=config.num_rows, + block_size=config.block_size, + outputs_per_block=config.outputs_per_block, + vector_width=config.vector_width, + k_unroll=config.k_unroll, + has_residual=has_residual, + use_pdl=self._use_pdl(), + ) + self._compiled[(dtype, config, has_residual)] = cute.compile( + kernel, + a, + b, + residual, + c, + self._stream(), + options="--enable-tvm-ffi --ptxas-options -maxrregcount=64", + ) + + def request_warmup( + self, + dtype: torch.dtype, + shapes: Iterable[tuple[int, int, int]], + ) -> None: + """Request compilation of ``(M, N, K)`` shapes before graph capture.""" + self.request_warmup_configs( + dtype, + (self._config(m, n, k) for m, n, k in shapes), + ) + + def request_warmup_configs( + self, + dtype: torch.dtype, + configs: Iterable[SkinnyGemmConfig], + *, + has_residual: bool = False, + ) -> None: + """Request compilation of explicit measured configs before capture.""" + self._warmup_configs.update((dtype, config, has_residual) for config in configs) + if self._warmup_registered: + return + from aphrodite.model_executor.warmup.cutedsl_warmup import ( + register_cutedsl_warmup_provider, + ) + + register_cutedsl_warmup_provider(self) + self._warmup_registered = True + + def get_cutedsl_warmup_compile_units(self): + from aphrodite.model_executor.warmup.cutedsl_warmup import CuTeDSLCompileUnit + + return tuple( + CuTeDSLCompileUnit( + name=("shape-dynamic skinny GEMM with residual" if has_residual else "shape-dynamic skinny GEMM"), + key=("shape-dynamic-skinny-gemm", dtype, config, has_residual), + compile=partial(self._compile, dtype, config, has_residual), + ) + for dtype, config, has_residual in sorted( + self._warmup_configs, + key=lambda item: ( + str(item[0]), + item[1].num_rows, + item[1].block_size, + item[1].outputs_per_block, + item[1].k_unroll, + item[1].vector_width, + item[2], + ), + ) + ) + + def __call__( + self, + a: torch.Tensor, + b: torch.Tensor, + config: SkinnyGemmConfig | None = None, + residual: torch.Tensor | None = None, + ) -> torch.Tensor: + if a.dim() != 2 or b.dim() != 2: + raise ValueError("a and b must be 2D tensors") + if a.dtype not in (torch.bfloat16, torch.float16) or b.dtype != a.dtype: + raise ValueError("a and b must have the same BF16 or FP16 dtype") + if not a.is_cuda or not b.is_cuda or a.device != b.device: + raise ValueError("a and b must be CUDA tensors on the same device") + if not a.is_contiguous() or not b.is_contiguous(): + raise ValueError("a and b must be contiguous") + if a.shape[1] != b.shape[1]: + raise ValueError("a and b must have matching K dimensions") + if not 1 <= a.shape[0] <= 16: + raise ValueError("shape-dynamic skinny GEMM requires 1 <= M <= 16") + if residual is not None: + if residual.dim() != 2 or residual.shape != (a.shape[0], b.shape[0]): + raise ValueError("residual must have shape (M, N)") + if residual.dtype != a.dtype: + raise ValueError("residual must have the same dtype as a and b") + if residual.device != a.device or not residual.is_cuda: + raise ValueError("residual must be on the same CUDA device") + if not residual.is_contiguous(): + raise ValueError("residual must be contiguous") + + config = config or self._config(a.shape[0], b.shape[0], a.shape[1]) + if config.num_rows != a.shape[0]: + raise ValueError("config num_rows must match M") + if b.shape[0] % config.outputs_per_block != 0: + raise ValueError("N must be divisible by outputs_per_block") + if a.shape[1] % (config.block_size * config.vector_width) != 0: + raise ValueError("K must be divisible by block_size * vector_width for this config") + has_residual = residual is not None + cache_key = (a.dtype, config, has_residual) + if cache_key not in self._compiled: + self._compile(a.dtype, config, has_residual) + output = torch.empty((a.shape[0], b.shape[0]), dtype=a.dtype, device=a.device) + residual_arg = output if residual is None else residual + self._compiled[cache_key](a, b, residual_arg, output, self._stream()) + return output + + +shape_dynamic_skinny_gemm = ShapeDynamicSkinnyGemm() diff --git a/aphrodite/model_executor/layers/activation.py b/aphrodite/model_executor/layers/activation.py index ba83f033b4..d2c88d07e0 100644 --- a/aphrodite/model_executor/layers/activation.py +++ b/aphrodite/model_executor/layers/activation.py @@ -153,6 +153,49 @@ def forward_xpu(self, x: torch.Tensor) -> torch.Tensor: return self.forward_cuda(x) +@CustomOp.register("situ_and_mul") +class SituAndMul(CustomOp): + """SituGLU activation used by Kimi models. + + Computes beta * tanh(gate / beta) * sigmoid(gate) * up. When + ``linear_beta`` is set, the up projection is also softly clipped with + linear_beta * tanh(up / linear_beta). + """ + + def __init__( + self, + beta: float = 1.0, + linear_beta: float | None = None, + *, + compile_native: bool = True, + ): + super().__init__(compile_native=compile_native) + self.beta = float(beta) + self.linear_beta = None if linear_beta is None else float(linear_beta) + if current_platform.is_cuda_alike(): + self.op = torch.ops._C.situ_and_mul + + def forward_native(self, x: torch.Tensor) -> torch.Tensor: + d = x.shape[-1] // 2 + gate = x[..., :d].float() + up = x[..., d:].float() + gate = self.beta * torch.tanh(gate / self.beta) * torch.sigmoid(gate) + if self.linear_beta is not None: + up = self.linear_beta * torch.tanh(up / self.linear_beta) + return (gate * up).to(x.dtype) + + def forward_cuda(self, x: torch.Tensor) -> torch.Tensor: + # Fused CUDA kernel: writes straight to `out`, no fp32 temporaries. + # linear_beta<=0 signals "unset" to the kernel (up passed through). + d = x.shape[-1] // 2 + out = torch.empty(x.shape[:-1] + (d,), dtype=x.dtype, device=x.device) + self.op(out, x, self.beta, -1.0 if self.linear_beta is None else self.linear_beta) + return out + + def forward_xpu(self, x: torch.Tensor) -> torch.Tensor: + return self.forward_native(x) + + @CustomOp.register("silu_and_mul_with_clamp") class SiluAndMulWithClamp(CustomOp): """SwiGLU activation with input clamping (used by some MoE shared experts). diff --git a/aphrodite/model_executor/layers/fused_moe/activation.py b/aphrodite/model_executor/layers/fused_moe/activation.py index 05afd0b857..87d3f9bb43 100644 --- a/aphrodite/model_executor/layers/fused_moe/activation.py +++ b/aphrodite/model_executor/layers/fused_moe/activation.py @@ -22,6 +22,7 @@ class MoEActivation(Enum): # expects the *packed* layout ([all gates; all ups]), as produced by a # MergedColumnParallelLinear gate_up_proj (e.g. MiniMax-M3). SWIGLUOAI = "swigluoai" + SITU = "situ" SWIGLUOAI_UNINTERLEAVE = "swigluoai_uninterleave" SWIGLUSTEP = "swiglustep" @@ -77,6 +78,7 @@ def from_str(cls, s: str) -> "MoEActivation": MoEActivation.SILU: "silu_and_mul", MoEActivation.GELU: "gelu_and_mul", MoEActivation.GELU_TANH: "gelu_tanh_and_mul", + MoEActivation.SITU: "situ_and_mul", MoEActivation.SWIGLUOAI: "swigluoai_and_mul", MoEActivation.SWIGLUOAI_UNINTERLEAVE: "silu_and_mul_with_clamp", MoEActivation.SWIGLUSTEP: "swiglustep_and_mul", @@ -133,6 +135,8 @@ def apply_moe_activation( beta: float = 0.0, topk_ids: torch.Tensor | None = None, expert_map: torch.Tensor | None = None, + activation_situ_beta: float | None = None, + activation_situ_linear_beta: float | None = None, ) -> torch.Tensor: """Apply MoE activation function. @@ -161,6 +165,21 @@ def apply_moe_activation( torch.ops._C.gelu_and_mul(output, input) elif activation == MoEActivation.GELU_TANH: torch.ops._C.gelu_tanh_and_mul(output, input) + elif activation == MoEActivation.SITU: + # Fused CUDA kernel: writes straight to `output`, no fp32 temporaries. + # (The pure-torch fallback below upcast both halves to fp32 and + # allocated ~8 temporaries per call, blowing up MoE memory.) + # Both betas come from FusedMoEConfig; a missing beta means the caller + # bypassed the config plumbing, so fail rather than silently use 1.0. + # linear_beta is genuinely optional: <= 0 signals "unset" to the kernel + # (up passed through), matching SituAndMul(linear_beta=None). + assert activation_situ_beta is not None, "SITU requires activation_situ_beta from FusedMoEConfig" + torch.ops._C.situ_and_mul( + output, + input, + activation_situ_beta, + -1.0 if activation_situ_linear_beta is None else activation_situ_linear_beta, + ) elif activation == MoEActivation.SWIGLUOAI: torch.ops._C.swigluoai_and_mul(output, input) elif activation == MoEActivation.SWIGLUOAI_UNINTERLEAVE: diff --git a/aphrodite/model_executor/layers/fused_moe/config.py b/aphrodite/model_executor/layers/fused_moe/config.py index 5e48b7e2f9..e4a8d85138 100644 --- a/aphrodite/model_executor/layers/fused_moe/config.py +++ b/aphrodite/model_executor/layers/fused_moe/config.py @@ -1288,6 +1288,10 @@ class FusedMoEConfig: swiglu_alpha: float | None = None swiglu_beta: float | None = None + # SituGLU parameters used by Kimi sit(u/v2) activations. + activation_situ_beta: float | None = None + activation_situ_linear_beta: float | None = None + max_capture_size: int = 0 # Set by __post_init__ diff --git a/aphrodite/model_executor/layers/fused_moe/experts/aiter_mxfp8_moe.py b/aphrodite/model_executor/layers/fused_moe/experts/aiter_mxfp8_moe.py index 92af44ffed..4584980167 100644 --- a/aphrodite/model_executor/layers/fused_moe/experts/aiter_mxfp8_moe.py +++ b/aphrodite/model_executor/layers/fused_moe/experts/aiter_mxfp8_moe.py @@ -6,10 +6,13 @@ ``convert_to_fp8_moe_kernel_format``. """ +import math + import torch import aphrodite.model_executor.layers.fused_moe.modular_kernel as mk from aphrodite.logger import init_logger +from aphrodite.model_executor.layers.fused_moe.activation import MoEActivation from aphrodite.model_executor.layers.fused_moe.experts.mxfp8_emulation_moe import ( Mxfp8TritonExpertsBase, ) @@ -17,6 +20,9 @@ logger = init_logger(__name__) +_AITER_SWIGLU_ALPHA = 1.702 +_AITER_SWIGLU_BETA = 1.0 + def is_aiter_mxfp8_moe_available() -> bool: """True when the FlyDSL MXFP8 MoE can run here: gfx950, the ``flydsl`` @@ -89,6 +95,24 @@ def is_supported_config(cls, moe_config, weight_key, activation_key, activation_ # reason when the device is fine but the flydsl package is missing. if is_supported and not is_aiter_mxfp8_moe_available(): return False, ("kernel requires the aiter flydsl package, which is not installed") + if is_supported and moe_config.activation != MoEActivation.SWIGLUOAI_UNINTERLEAVE: + return False, ( + "kernel hardcodes SwiGLU-OAI activation and requires " + f"activation={MoEActivation.SWIGLUOAI_UNINTERLEAVE.value}; " + f"got activation={moe_config.activation.value}" + ) + if is_supported and ( + moe_config.swiglu_alpha is None + or not math.isclose(float(moe_config.swiglu_alpha), _AITER_SWIGLU_ALPHA) + or moe_config.swiglu_beta is None + or not math.isclose(float(moe_config.swiglu_beta), _AITER_SWIGLU_BETA) + ): + return False, ( + "kernel hardcodes SwiGLU-OAI with " + f"alpha={_AITER_SWIGLU_ALPHA} and beta={_AITER_SWIGLU_BETA}; " + f"got swiglu_alpha={moe_config.swiglu_alpha} and " + f"swiglu_beta={moe_config.swiglu_beta}" + ) return is_supported, reason def apply( diff --git a/aphrodite/model_executor/layers/fused_moe/experts/deep_gemm_moe.py b/aphrodite/model_executor/layers/fused_moe/experts/deep_gemm_moe.py index 96a715221b..965bc5f8c2 100644 --- a/aphrodite/model_executor/layers/fused_moe/experts/deep_gemm_moe.py +++ b/aphrodite/model_executor/layers/fused_moe/experts/deep_gemm_moe.py @@ -263,7 +263,16 @@ def _act_mul_quant( # 3. fallback path for non-SiLU activations in non‑UE8M0 cases. act_out = torch.empty((M_sum, activation_out_dim), dtype=input.dtype, device=input.device) self.activation(activation, act_out, input) - return per_token_group_quant_fp8(act_out, block_k, column_major_scales=True, out_q=output) + if scale_fmt == DeepGemmQuantScaleFMT.UE8M0: + return per_token_group_quant_fp8_packed_for_deepgemm(act_out, block_k, use_ue8m0=True, out_q=output) + use_ue8m0 = scale_fmt == DeepGemmQuantScaleFMT.FLOAT32_CEIL_UE8M0 + return per_token_group_quant_fp8( + act_out, + block_k, + column_major_scales=True, + out_q=output, + use_ue8m0=use_ue8m0, + ) def apply( self, @@ -414,7 +423,14 @@ def _supports_quant_scheme( @staticmethod def _supports_activation(activation: MoEActivation) -> bool: - return activation in [MoEActivation.SILU, MoEActivation.SWIGLUSTEP] + # SILU has fused gate+mul+quant kernels; SWIGLUSTEP/SITU take the + # general path (activation applied via self.activation, which forwards + # the situ betas, then FP8 requant). + return activation in [ + MoEActivation.SILU, + MoEActivation.SWIGLUSTEP, + MoEActivation.SITU, + ] @staticmethod def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bool: @@ -455,16 +471,15 @@ def _act_mul_quant( M_sum, N = input.size() activation_out_dim = self.adjust_N_for_activation(N, activation) - if scale_fmt == DeepGemmQuantScaleFMT.UE8M0: - assert activation == MoEActivation.SILU - return fused_silu_mul_fp8_quant_packed( - input=input, - output_q=output, - group_size=block_k, - clamp_limit=self.gemm1_clamp_limit, - ) - if activation == MoEActivation.SILU: + # Fused gate+mul+quant kernels for the common SILU case. + if scale_fmt == DeepGemmQuantScaleFMT.UE8M0: + return fused_silu_mul_fp8_quant_packed( + input=input, + output_q=output, + group_size=block_k, + clamp_limit=self.gemm1_clamp_limit, + ) use_ue8m0 = scale_fmt == DeepGemmQuantScaleFMT.FLOAT32_CEIL_UE8M0 return silu_mul_per_token_group_quant_fp8_colmajor( input=input, diff --git a/aphrodite/model_executor/layers/fused_moe/experts/flashinfer_cutedsl_moe.py b/aphrodite/model_executor/layers/fused_moe/experts/flashinfer_cutedsl_moe.py index af54690846..8e2b95053f 100644 --- a/aphrodite/model_executor/layers/fused_moe/experts/flashinfer_cutedsl_moe.py +++ b/aphrodite/model_executor/layers/fused_moe/experts/flashinfer_cutedsl_moe.py @@ -13,6 +13,9 @@ from aphrodite.model_executor.layers.fused_moe.topk_weight_and_reduce import ( TopKWeightAndReduceNoOP, ) +from aphrodite.model_executor.layers.quantization.utils.flashinfer_utils import ( + activation_to_flashinfer_int, +) from aphrodite.model_executor.layers.quantization.utils.quant_utils import ( QuantKey, kNvfp4Dynamic, @@ -68,7 +71,7 @@ def _supports_current_device() -> bool: @staticmethod def _supports_no_act_and_mul() -> bool: - return False + return True @staticmethod def _supports_quant_scheme( @@ -82,7 +85,7 @@ def _supports_quant_scheme( @staticmethod def _supports_activation(activation: MoEActivation) -> bool: - return activation == MoEActivation.SILU + return activation in (MoEActivation.SILU, MoEActivation.RELU2_NO_MUL) @staticmethod def _supports_parallel_config( @@ -155,4 +158,5 @@ def apply( num_local_experts=self.local_num_experts, local_expert_offset=self.local_expert_offset, moe_output=output, + activation_type=activation_to_flashinfer_int(activation), ) 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 145a04f14f..fa070d4ad9 100644 --- a/aphrodite/model_executor/layers/fused_moe/experts/marlin_moe.py +++ b/aphrodite/model_executor/layers/fused_moe/experts/marlin_moe.py @@ -93,6 +93,8 @@ def _fused_marlin_moe( clamp_limit: float | None = None, gemm1_alpha: float = 1.0, gemm1_beta: float = 0.0, + activation_situ_beta: float | None = None, + activation_situ_linear_beta: float | None = None, ) -> torch.Tensor: assert hidden_states.ndim == 2 M, K = hidden_states.size() @@ -169,6 +171,8 @@ def _fused_marlin_moe( beta=gemm1_beta, topk_ids=topk_ids, expert_map=expert_map, + activation_situ_beta=activation_situ_beta, + activation_situ_linear_beta=activation_situ_linear_beta, ) if output is None: @@ -250,6 +254,8 @@ def fused_marlin_moe( clamp_limit: float | None = None, gemm1_alpha: float = 1.0, gemm1_beta: float = 0.0, + activation_situ_beta: float | None = None, + activation_situ_linear_beta: float | None = None, ) -> torch.Tensor: """ This function computes a Mixture of Experts (MoE) layer using two sets of @@ -351,6 +357,8 @@ def fused_marlin_moe( num_tokens_post_padded=num_tokens_post_padded, activation=activation, activation_func=activation_func, + activation_situ_beta=activation_situ_beta, + activation_situ_linear_beta=activation_situ_linear_beta, input_global_scale1=input_global_scale1, input_global_scale2=input_global_scale2, global_scale1=global_scale1, @@ -417,6 +425,9 @@ def batched_fused_marlin_moe( clamp_limit: float | None = None, gemm1_alpha: float = 1.0, gemm1_beta: float = 0.0, + activation_func: Callable[..., None] = apply_moe_activation, + activation_situ_beta: float | None = None, + activation_situ_linear_beta: float | None = None, ) -> torch.Tensor: """ This function massages the inputs so the batched hidden_states can be @@ -520,6 +531,9 @@ def batched_fused_marlin_moe( quant_type=quant_type, apply_router_weight_on_input=apply_router_weight_on_input, activation=activation, + activation_func=activation_func, + activation_situ_beta=activation_situ_beta, + activation_situ_linear_beta=activation_situ_linear_beta, expert_map=expert_map, block_size_m=block_size_m, sorted_token_ids=sorted_token_ids, @@ -631,6 +645,7 @@ def _supports_activation(activation: MoEActivation) -> bool: MoEActivation.SILU, MoEActivation.GELU, MoEActivation.GELU_TANH, + MoEActivation.SITU, MoEActivation.SWIGLUOAI, MoEActivation.SWIGLUOAI_UNINTERLEAVE, MoEActivation.SWIGLUSTEP, @@ -775,6 +790,8 @@ def apply( global_num_experts=global_num_experts, activation=activation, activation_func=self.activation, + activation_situ_beta=self.moe_config.activation_situ_beta, + activation_situ_linear_beta=(self.moe_config.activation_situ_linear_beta), moe_sum=self.moe_sum, expert_map=expert_map, output=output, @@ -815,6 +832,8 @@ def activation_with_lora( beta: float = 0.0, topk_ids: torch.Tensor | None = None, expert_map: torch.Tensor | None = None, + activation_situ_beta: float | None = None, + activation_situ_linear_beta: float | None = None, ) -> None: # act_input = intermediate_cache1 (M*topk, 2N for gated) # act_output = intermediate_cache2 (M*topk, N) @@ -853,6 +872,8 @@ def activation_with_lora( beta=beta, topk_ids=topk_ids, expert_map=expert_map, + activation_situ_beta=activation_situ_beta, + activation_situ_linear_beta=activation_situ_linear_beta, ) lora_state["cache2"] = act_output @@ -900,6 +921,8 @@ def moe_sum_with_lora( global_num_experts=global_num_experts, activation=activation, activation_func=activation_with_lora, + activation_situ_beta=self.moe_config.activation_situ_beta, + activation_situ_linear_beta=self.moe_config.activation_situ_linear_beta, moe_sum=moe_sum_with_lora, expert_map=expert_map, output=output, @@ -1003,6 +1026,39 @@ def apply( apply_router_weight_on_input: bool, ): assert expert_tokens_meta is not None, "Num valid tokens per batch is required" + + def activation_func( + act: MoEActivation, + act_output: torch.Tensor, + act_input: torch.Tensor, + *, + activation_situ_beta: float | None = None, + activation_situ_linear_beta: float | None = None, + **kwargs, + ) -> None: + if act != MoEActivation.SITU: + self.activation( + act, + act_output, + act_input, + activation_situ_beta=activation_situ_beta, + activation_situ_linear_beta=activation_situ_linear_beta, + **kwargs, + ) + return + + num_experts, max_num_tokens = hidden_states.shape[:2] + beta = activation_situ_beta + linear_beta = activation_situ_linear_beta + assert beta is not None, "SITU requires activation_situ_beta from FusedMoEConfig" + torch.ops._C.masked_situ_and_mul( + act_output.view(num_experts, max_num_tokens, -1), + act_input.view(num_experts, max_num_tokens, -1), + expert_tokens_meta.expert_num_tokens, + beta, + -1.0 if linear_beta is None else linear_beta, + ) + return batched_fused_marlin_moe( hidden_states=hidden_states, expert_num_tokens=expert_tokens_meta.expert_num_tokens, @@ -1033,4 +1089,7 @@ def apply( clamp_limit=self.gemm1_clamp_limit, gemm1_alpha=self.gemm1_alpha, gemm1_beta=self.gemm1_beta, + activation_func=activation_func, + activation_situ_beta=self.moe_config.activation_situ_beta, + activation_situ_linear_beta=self.moe_config.activation_situ_linear_beta, ) diff --git a/aphrodite/model_executor/layers/fused_moe/experts/mxfp8_emulation_moe.py b/aphrodite/model_executor/layers/fused_moe/experts/mxfp8_emulation_moe.py index 35b7ec8e81..d7f58a38d4 100644 --- a/aphrodite/model_executor/layers/fused_moe/experts/mxfp8_emulation_moe.py +++ b/aphrodite/model_executor/layers/fused_moe/experts/mxfp8_emulation_moe.py @@ -107,17 +107,13 @@ def activation( limit = self.quant_config.gemm1_clamp_limit if limit is None: raise ValueError("SWIGLUOAI_UNINTERLEAVE requires gemm1_clamp_limit") - alpha = self.quant_config.gemm1_alpha - alpha = 1.702 if alpha is None else float(alpha) - beta = self.quant_config.gemm1_beta - beta = 1.0 if beta is None else float(beta) apply_moe_activation( activation, output, input, clamp_limit=float(limit), - alpha=alpha, - beta=beta, + alpha=self.gemm1_alpha, + beta=self.gemm1_beta, ) return super().activation(activation, output, input) diff --git a/aphrodite/model_executor/layers/fused_moe/experts/mxfp8_native_moe.py b/aphrodite/model_executor/layers/fused_moe/experts/mxfp8_native_moe.py index 195f6e0900..e076fb6671 100644 --- a/aphrodite/model_executor/layers/fused_moe/experts/mxfp8_native_moe.py +++ b/aphrodite/model_executor/layers/fused_moe/experts/mxfp8_native_moe.py @@ -352,10 +352,8 @@ def apply( expert_tokens_meta: mk.ExpertTokensMetadata | None, apply_router_weight_on_input: bool, ): - alpha = self.quant_config.gemm1_alpha - alpha = 1.702 if alpha is None else float(alpha) - beta = self.quant_config.gemm1_beta - beta = 1.0 if beta is None else float(beta) + # `self.gemm1_alpha` and `self.gemm1_beta` are set by + # `TritonExperts.__init__`. limit = self.quant_config.gemm1_clamp_limit limit = None if limit is None else float(limit) out = fused_moe_mxfp8_native( @@ -366,8 +364,8 @@ def apply( self.w2_scale_val, topk_weights, topk_ids, - alpha=alpha, - beta=beta, + alpha=self.gemm1_alpha, + beta=self.gemm1_beta, limit=limit, global_num_experts=global_num_experts, expert_map=expert_map, diff --git a/aphrodite/model_executor/layers/fused_moe/experts/trtllm_mxfp4_moe.py b/aphrodite/model_executor/layers/fused_moe/experts/trtllm_mxfp4_moe.py index 903c266d78..3dc419509f 100644 --- a/aphrodite/model_executor/layers/fused_moe/experts/trtllm_mxfp4_moe.py +++ b/aphrodite/model_executor/layers/fused_moe/experts/trtllm_mxfp4_moe.py @@ -75,6 +75,33 @@ def __init__( else: self.gemm1_clamp_limit = None + # SITU (SituGLU) TRTLLM-Gen kernel computes + # left = alpha * tanh(x0 / alpha) * sigmoid(x0) # gate (x0) + # right = beta * tanh(x1 / beta) # up (x1) + # which matches vLLM's situ_and_mul with (beta, linear_beta), so map + # situ beta -> gatedActAlpha (gemm1_alpha) and situ linear_beta -> + # gatedActBeta (gemm1_beta). Both must be > 0. + if moe_config.activation == MoEActivation.SITU: + situ_beta = moe_config.activation_situ_beta + situ_linear_beta = moe_config.activation_situ_linear_beta + assert situ_beta is not None and situ_beta > 0, "SITU requires activation_situ_beta > 0" + assert situ_linear_beta is not None and situ_linear_beta > 0, ( + "TRTLLM SiTuGlu requires activation_situ_linear_beta > 0 (the private cubin has no up-passthrough path)" + ) + self.gemm1_alpha = torch.full( + (self.local_num_experts,), + float(situ_beta), + dtype=torch.float32, + device=device, + ) + self.gemm1_beta = torch.full( + (self.local_num_experts,), + float(situ_linear_beta), + dtype=torch.float32, + device=device, + ) + self.gemm1_clamp_limit = None + self.max_capture_size = moe_config.max_capture_size @staticmethod @@ -99,7 +126,19 @@ def _supports_quant_scheme( @staticmethod def _supports_activation(activation: MoEActivation) -> bool: - return activation in (MoEActivation.SWIGLUOAI, MoEActivation.SILU) + return activation in ( + MoEActivation.SWIGLUOAI, + MoEActivation.SILU, + MoEActivation.SITU, + ) + + @staticmethod + def _flashinfer_activation_type(activation: MoEActivation) -> int: + from flashinfer.fused_moe.core import ActivationType + + if activation == MoEActivation.SITU: + return ActivationType.Situ.value + return ActivationType.Swiglu.value @staticmethod def activation_format() -> mk.FusedMoEActivationFormat: @@ -136,6 +175,7 @@ def _supports_routing_method( activation_key: QuantKey | None, ) -> bool: return routing_method in [ + RoutingMethodType.DeepSeekV3, RoutingMethodType.Renormalize, RoutingMethodType.RenormalizeNaive, ] @@ -186,8 +226,8 @@ def apply( device=hidden_states.device, ) trtllm_fp4_block_scale_moe( - routing_logits=router_logits.to(torch.bfloat16), - routing_bias=None, + routing_logits=router_logits, + routing_bias=e_score_correction_bias, hidden_states=x_quant, hidden_states_scale=x_scale, gemm1_weights=w1, @@ -204,14 +244,15 @@ def apply( output2_scale_scalar=None, num_experts=global_num_experts, top_k=self.topk, - n_group=None, - topk_group=None, + n_group=(num_expert_group or 0), + topk_group=(topk_group or 0), intermediate_size=self.intermediate_size_per_partition, local_expert_offset=self.ep_rank * self.local_num_experts, local_num_experts=self.local_num_experts, - routed_scaling_factor=None, + routed_scaling_factor=routed_scaling_factor, routing_method_type=self.routing_method_type, do_finalize=True, + activation_type=self._flashinfer_activation_type(activation), tune_max_num_tokens=max(self.max_capture_size, 1), output=output, routing_replay_out=routing_replay_out, @@ -331,6 +372,7 @@ def apply( "routing_method_type": RoutingMethodType.Renormalize, "do_finalize": True, "enable_pdl": True, + "activation_type": self._flashinfer_activation_type(activation), "output": output, "tune_max_num_tokens": max(self.max_capture_size, 1), } diff --git a/aphrodite/model_executor/layers/fused_moe/experts/trtllm_nvfp4_moe.py b/aphrodite/model_executor/layers/fused_moe/experts/trtllm_nvfp4_moe.py index 1c87a1439a..1bc2986d21 100644 --- a/aphrodite/model_executor/layers/fused_moe/experts/trtllm_nvfp4_moe.py +++ b/aphrodite/model_executor/layers/fused_moe/experts/trtllm_nvfp4_moe.py @@ -106,6 +106,24 @@ def _per_expert(val: float | None) -> torch.Tensor | None: self.gemm1_alpha = None self.gemm1_beta = None + # SITU (Kimi SituGLU) TRTLLM-Gen kernel computes + # left=alpha*tanh(x0/alpha)*sigmoid(x0), right=beta*tanh(x1/beta), + # matching vLLM's situ_and_mul, so map situ beta -> gatedActAlpha + # (gemm1_alpha) and situ linear_beta -> gatedActBeta (gemm1_beta). + # These operate on the dequantized gate/up, so they are NOT folded by + # g1_alphas in process_weights_after_loading. + self.is_situ = moe_config.activation == MoEActivation.SITU + if self.is_situ: + situ_beta = moe_config.activation_situ_beta + situ_linear_beta = moe_config.activation_situ_linear_beta + assert situ_beta is not None and situ_beta > 0, "SITU requires activation_situ_beta > 0" + assert situ_linear_beta is not None and situ_linear_beta > 0, ( + "TRTLLM SiTuGlu requires activation_situ_linear_beta > 0 (the private cubin has no up-passthrough path)" + ) + self.gemm1_alpha = _per_expert(situ_beta) + self.gemm1_beta = _per_expert(situ_linear_beta) + self.gemm1_clamp_limit = None + logger.debug_once( "activation=%s, gemm1_alpha=%s, gemm1_beta=%s, gemm1_clamp_limit=%s", moe_config.activation, @@ -138,7 +156,10 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None: # (via the in-place mul above) and never changes again, so this is a # static, per-expert constant. Register on the layer so EPLB # rearranges it alongside the other expert tensors. - if self.gemm1_clamp_limit is not None: + # SITU alpha/beta act on the dequantized gate/up (tanh clamps), not the + # raw GEMM1 accumulator, so they are registered as-is without the + # g1_alphas fold used by the SwiGLU-OAI clamp/beta below. + if self.gemm1_clamp_limit is not None and not self.is_situ: gemm1_clamp_limit = self.gemm1_clamp_limit / self.quant_config.g1_alphas layer.register_parameter( "gemm1_clamp_limit", @@ -151,7 +172,7 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None: # raw. Register both on the layer so EPLB rearranges them with the # other per-expert tensors. if self.gemm1_beta is not None: - gemm1_beta = self.gemm1_beta / self.quant_config.g1_alphas + gemm1_beta = self.gemm1_beta if self.is_situ else self.gemm1_beta / self.quant_config.g1_alphas layer.register_parameter( "gemm1_beta", torch.nn.Parameter(gemm1_beta, requires_grad=False), @@ -189,13 +210,15 @@ def _supports_quant_scheme( @staticmethod def _supports_activation(activation: MoEActivation) -> bool: - """Supports SiLU, RELU^2 non-gated, GELU, and clamped SwiGLU-OAI.""" + """Supports SiLU, RELU^2 non-gated, GELU, clamped SwiGLU-OAI, and the + private SituGLU (SITU) cubin.""" return activation in [ MoEActivation.SILU, MoEActivation.RELU2_NO_MUL, MoEActivation.GELU, MoEActivation.GELU_TANH, MoEActivation.SWIGLUOAI_UNINTERLEAVE, + MoEActivation.SITU, ] @staticmethod diff --git a/aphrodite/model_executor/layers/fused_moe/fused_moe.py b/aphrodite/model_executor/layers/fused_moe/fused_moe.py index 132b3f666b..39fd1ac06f 100644 --- a/aphrodite/model_executor/layers/fused_moe/fused_moe.py +++ b/aphrodite/model_executor/layers/fused_moe/fused_moe.py @@ -28,9 +28,12 @@ from aphrodite.model_executor.layers.fused_moe.utils import ( enable_swap_ab, moe_kernel_quantize_input, + resolve_moe_use_td, + warn_if_moe_use_td_ineffective, ) from aphrodite.platforms import current_platform from aphrodite.triton_utils import tl, triton +from aphrodite.triton_utils.allocation import set_triton_allocator from aphrodite.utils.math_utils import next_power_of_2 from aphrodite.utils.platform_utils import get_device_name_as_file_name from aphrodite.utils.torch_utils import direct_register_custom_op @@ -327,6 +330,8 @@ def fused_moe_kernel( per_channel_quant: tl.constexpr, HAS_BIAS: tl.constexpr, SWAP_AB: tl.constexpr, + # Tensor-descriptor path for the A gather and B load in the K-loop. + USE_TD: tl.constexpr = False, ): """ Implements the fused computation for a Mixture of Experts (MOE) using @@ -416,7 +421,25 @@ def fused_moe_kernel( offs_bn = (pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N).to(tl.int64)) % N offs_k = tl.arange(0, BLOCK_SIZE_K) - if SWAP_AB: + # TD gather and the SWAP_AB accumulator layout are mutually exclusive. + tl.static_assert(not (USE_TD and SWAP_AB)) + if USE_TD: + # ``tt.descriptor_gather`` requires block_shape[0] == 1 and i32 idx. + m_td = num_valid_tokens // top_k + a_desc = tl.make_tensor_descriptor( + base=a_ptr, + shape=(m_td, K), + strides=(stride_am, stride_ak), + block_shape=(1, BLOCK_SIZE_K), + ) + b_desc = tl.make_tensor_descriptor( + base=b_ptr + off_experts * stride_be, + shape=(N, K), + strides=(stride_bn, stride_bk), + block_shape=(BLOCK_SIZE_N, BLOCK_SIZE_K), + ) + gather_idx = (offs_token // top_k).to(tl.int32) + elif SWAP_AB: a_ptrs = a_ptr + (offs_k[:, None] * stride_ak + offs_token[None, :] // top_k * stride_am) b_ptrs = b_ptr + off_experts * stride_be + (offs_bn[:, None] * stride_bn + offs_k[None, :] * stride_bk) else: @@ -460,18 +483,21 @@ def fused_moe_kernel( for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)): # Load the next block of A and B, generate a mask by checking the # K dimension. - if SWAP_AB: + if USE_TD: + a = a_desc.gather(gather_idx, k * BLOCK_SIZE_K) + b = b_desc.load([pid_n * BLOCK_SIZE_N, k * BLOCK_SIZE_K]).T + elif SWAP_AB: a_mask = (offs_k[:, None] < K - k * BLOCK_SIZE_K) & token_mask[None, :] b_mask = offs_k[None, :] < K - k * BLOCK_SIZE_K + a = tl.load(a_ptrs, mask=a_mask, other=0.0) + b = tl.load(b_ptrs, mask=b_mask, other=0.0) else: - a_mask = token_mask[:, None] & (offs_k[None, :] < K - k * BLOCK_SIZE_K) - b_mask = offs_k[:, None] < K - k * BLOCK_SIZE_K - a = tl.load( - a_ptrs, - mask=a_mask, - other=0.0, - ) - b = tl.load(b_ptrs, mask=b_mask, other=0.0) + a = tl.load( + a_ptrs, + mask=token_mask[:, None] & (offs_k[None, :] < K - k * BLOCK_SIZE_K), + other=0.0, + ) + b = tl.load(b_ptrs, mask=offs_k[:, None] < K - k * BLOCK_SIZE_K, other=0.0) # We accumulate along the K dimension. if use_int8_w8a16: accumulator = tl.dot(a, b.to(compute_type), acc=accumulator) @@ -496,9 +522,10 @@ def fused_moe_kernel( accumulator += tl.dot(a, b) else: accumulator += tl.dot(a, b) - # Advance the ptrs to the next K block. - a_ptrs += BLOCK_SIZE_K * stride_ak - b_ptrs += BLOCK_SIZE_K * stride_bk + if not USE_TD: + # Advance the ptrs to the next K block. + a_ptrs += BLOCK_SIZE_K * stride_ak + b_ptrs += BLOCK_SIZE_K * stride_bk if SWAP_AB: accumulator = tl.trans(accumulator, (1, 0)) @@ -722,6 +749,18 @@ def invoke_fused_moe_triton_kernel( else: SWAP_AB = False + # Quantized weights always carry a B_scale (see the asserts below); key off + # that rather than enumerating quant flags, which misses w8a16-fp8/nvfp4/etc. + is_quantized = B_scale is not None + warn_if_moe_use_td_ineffective("TRITON", is_quantized=is_quantized) + + # TD path is unvalidated under quantization; fall back to the pointer path. + use_td = resolve_moe_use_td() and not is_quantized + if use_td: + # The TD path builds a tensor descriptor inside the kernel, which + # requires a PyTorch-backed scratch allocator to be registered. + set_triton_allocator(A.device) + if use_fp8_w8a8 or use_int8_w8a8: assert B_scale is not None assert block_shape is None or triton.cdiv(B.size(-2), block_shape[0]) == B_scale.size(-2) @@ -753,6 +792,15 @@ def invoke_fused_moe_triton_kernel( BLOCK_SIZE_K = config.pop("BLOCK_SIZE_K") if block_shape is not None: BLOCK_SIZE_K = min(BLOCK_SIZE_K, min(block_shape[0], block_shape[1])) + if use_td and A.size(1) % BLOCK_SIZE_K != 0: + logger.warning_once( + "Disabling APHRODITE_TRITON_USE_TD for this MoE launch: K=%d is not " + "a multiple of BLOCK_SIZE_K=%d, which triggers a known Triton " + "tensor-descriptor + tl.dot miscompilation.", + A.size(1), + BLOCK_SIZE_K, + ) + use_td = False fused_moe_kernel[grid]( A, B, @@ -795,6 +843,7 @@ def invoke_fused_moe_triton_kernel( HAS_BIAS=HAS_BIAS, BLOCK_SIZE_K=BLOCK_SIZE_K, SWAP_AB=SWAP_AB, + USE_TD=use_td, **config, ) diff --git a/aphrodite/model_executor/layers/fused_moe/layer.py b/aphrodite/model_executor/layers/fused_moe/layer.py index ae566adb63..1f97a233e9 100644 --- a/aphrodite/model_executor/layers/fused_moe/layer.py +++ b/aphrodite/model_executor/layers/fused_moe/layer.py @@ -115,6 +115,8 @@ def FusedMoE( swiglu_limit: float | None = None, swiglu_alpha: float | None = None, swiglu_beta: float | None = None, + activation_situ_beta: float | None = None, + activation_situ_linear_beta: float | None = None, e_score_correction_bias: torch.Tensor | None = None, apply_router_weight_on_input: bool = False, activation: str = "silu", @@ -173,6 +175,8 @@ def FusedMoE( scoring_func: Scoring function for routing ("softmax" or others) routed_scaling_factor: Scaling factor applied to topk_weights or output swiglu_limit: SwiGLU activation limit + activation_situ_beta: SituGLU activation beta + activation_situ_linear_beta: SituGLU linear beta e_score_correction_bias: Expert score correction bias tensor apply_router_weight_on_input: Whether to apply router weights on input activation: Activation function name ("silu", "gelu", etc.) @@ -337,6 +341,8 @@ def FusedMoE( swiglu_limit=swiglu_limit, swiglu_alpha=swiglu_alpha, swiglu_beta=swiglu_beta, + activation_situ_beta=activation_situ_beta, + activation_situ_linear_beta=activation_situ_linear_beta, max_capture_size=aphrodite_config.compilation_config.max_cudagraph_capture_size, skip_final_all_reduce=skip_final_all_reduce, ) diff --git a/aphrodite/model_executor/layers/fused_moe/modular_kernel.py b/aphrodite/model_executor/layers/fused_moe/modular_kernel.py index 80c3877049..43d5949dfd 100644 --- a/aphrodite/model_executor/layers/fused_moe/modular_kernel.py +++ b/aphrodite/model_executor/layers/fused_moe/modular_kernel.py @@ -872,6 +872,8 @@ def activation( beta: float = 0.0, topk_ids: torch.Tensor | None = None, expert_map: torch.Tensor | None = None, + activation_situ_beta: float | None = None, + activation_situ_linear_beta: float | None = None, ) -> None: apply_moe_activation( activation, @@ -882,6 +884,14 @@ def activation( beta=beta, topk_ids=topk_ids, expert_map=expert_map, + activation_situ_beta=( + self.moe_config.activation_situ_beta if activation_situ_beta is None else activation_situ_beta + ), + activation_situ_linear_beta=( + self.moe_config.activation_situ_linear_beta + if activation_situ_linear_beta is None + else activation_situ_linear_beta + ), ) @abstractmethod @@ -1281,6 +1291,14 @@ def _fused_experts( activation, ) + use_output_alias = ( + output_alias is not None + and output_alias.shape == fused_out.shape + and output_alias.dtype == fused_out.dtype + and output_alias.device == fused_out.device + and output_alias.is_contiguous() + ) + # If caller's output buffer already matches fused_out shape/dtype, alias # to skip the redundant copy in TopKWeightAndReduceNoOP.apply downstream. # This eliminates ~94% of __amd_rocclr_copyBuffer events (Copy 2 of the @@ -1288,15 +1306,10 @@ def _fused_experts( if current_platform.is_rocm(): from aphrodite._aiter_ops import rocm_aiter_ops - if ( - rocm_aiter_ops.is_fused_moe_enabled() - and output_alias is not None - and output_alias.shape == fused_out.shape - and output_alias.dtype == fused_out.dtype - and output_alias.device == fused_out.device - and output_alias.is_contiguous() - ): + if use_output_alias and rocm_aiter_ops.is_fused_moe_enabled(): fused_out = output_alias + elif use_output_alias: + fused_out = output_alias self.fused_experts.apply( output=fused_out, diff --git a/aphrodite/model_executor/layers/fused_moe/oracle/mxfp8.py b/aphrodite/model_executor/layers/fused_moe/oracle/mxfp8.py index e91499d5ca..ca8a94786c 100644 --- a/aphrodite/model_executor/layers/fused_moe/oracle/mxfp8.py +++ b/aphrodite/model_executor/layers/fused_moe/oracle/mxfp8.py @@ -12,10 +12,10 @@ kMxfp8Dynamic, kMxfp8Static, ) -from aphrodite.platforms import current_platform logger = init_logger(__name__) +# Ordered by priority. _SUPPORTED_BACKENDS = ( Fp8MoeBackend.FLASHINFER_TRTLLM, Fp8MoeBackend.DEEPGEMM, @@ -26,6 +26,8 @@ # devices / no flydsl / EP it is skipped and native is used. Fp8MoeBackend.AITER_MXFP8, Fp8MoeBackend.HUMMING, + Fp8MoeBackend.TRITON_MXFP8, + Fp8MoeBackend.EMULATION, ) _BACKEND_NAME_MAP: dict[str, Fp8MoeBackend] = { @@ -61,15 +63,12 @@ def _mxfp8_backend_to_kernel_cls( return [AiterMxfp8Experts] if backend == Fp8MoeBackend.TRITON_MXFP8: - # Explicit ``--moe-backend triton``: the Triton mxfp8 path, i.e. - # dot_scaled on MX-capable HW (gfx950) and BF16 emulation otherwise. - # Mirrors the ROCm auto-fallback in ``_select_rocm_mxfp8_backend``. - if current_platform.supports_mx(): - from aphrodite.model_executor.layers.fused_moe.experts.mxfp8_native_moe import ( - Mxfp8NativeTritonExperts, - ) + from aphrodite.model_executor.layers.fused_moe.experts.mxfp8_native_moe import ( + Mxfp8NativeTritonExperts, + ) - return [Mxfp8NativeTritonExperts] + return [Mxfp8NativeTritonExperts] + if backend == Fp8MoeBackend.EMULATION: from aphrodite.model_executor.layers.fused_moe.experts.mxfp8_emulation_moe import ( Mxfp8EmulationTritonExperts, ) @@ -103,35 +102,6 @@ def _select_kernel_cls( raise ValueError(f"No supported MXFP8 expert class for {backend.value}: {last_reason}") -def _select_rocm_mxfp8_backend() -> tuple[Fp8MoeBackend, type[mk.FusedMoEExperts]]: - """ROCm fallback when no auto-selected MXFP8 backend is available. - - The aiter FlyDSL backend (``AITER_MXFP8``) is auto-picked earlier by - ``select_mxfp8_moe_backend`` via ``_SUPPORTED_BACKENDS`` when usable, or - explicitly via ``--moe-backend aiter``; this fallback handles the rest - (native dot_scaled on gfx950, else BF16 emulation). - """ - - if current_platform.supports_mx(): - from aphrodite.model_executor.layers.fused_moe.experts.mxfp8_native_moe import ( - Mxfp8NativeTritonExperts, - ) - - logger.info_once("Using native CDNA4 (gfx950) MXFP8 dot_scaled MoE backend.") - return Fp8MoeBackend.TRITON_MXFP8, Mxfp8NativeTritonExperts - - from aphrodite.model_executor.layers.fused_moe.experts.mxfp8_emulation_moe import ( - Mxfp8EmulationTritonExperts, - ) - - logger.info_once( - "No native MXFP8 MoE backend available on this device; " - "MXFP8 weights will be dequantized to BF16 once at load time and the " - "MoE will run in BF16 (no per-step dequant)." - ) - return Fp8MoeBackend.EMULATION, Mxfp8EmulationTritonExperts - - def select_mxfp8_moe_backend( config: FusedMoEConfig, ) -> tuple[Fp8MoeBackend, type[mk.FusedMoEExperts]]: @@ -165,8 +135,5 @@ def select_mxfp8_moe_backend( logger.info_once("Using '%s' MxFp8 MoE backend.", backend.value) return backend, experts_cls - # simplify the logic for rocm, refactor later when more backends are supported - if current_platform.is_rocm(): - return _select_rocm_mxfp8_backend() - + # TODO: add debug log with reason. raise ValueError("No MXFP8 MoE backends available.") diff --git a/aphrodite/model_executor/layers/fused_moe/oracle/unquantized.py b/aphrodite/model_executor/layers/fused_moe/oracle/unquantized.py index a358f703bf..61e7587320 100644 --- a/aphrodite/model_executor/layers/fused_moe/oracle/unquantized.py +++ b/aphrodite/model_executor/layers/fused_moe/oracle/unquantized.py @@ -332,6 +332,13 @@ def make_unquantized_moe_kernel( experts_cls: type[mk.FusedMoEExperts], routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None, ) -> mk.FusedMoEKernel: + from aphrodite.model_executor.layers.fused_moe.utils import ( + warn_if_moe_use_td_ineffective, + ) + + # Warn against the selected backend, not each probed candidate. + warn_if_moe_use_td_ineffective(backend.value, is_quantized=False) + # Create Prepare/Finalize is_monolithic = issubclass(experts_cls, mk.FusedMoEExpertsMonolithic) prepare_finalize = maybe_make_prepare_finalize( diff --git a/aphrodite/model_executor/layers/fused_moe/router/router_factory.py b/aphrodite/model_executor/layers/fused_moe/router/router_factory.py index 297418a436..ae08382363 100644 --- a/aphrodite/model_executor/layers/fused_moe/router/router_factory.py +++ b/aphrodite/model_executor/layers/fused_moe/router/router_factory.py @@ -9,6 +9,7 @@ from aphrodite.distributed.eplb.eplb_state import EplbLayerState from aphrodite.model_executor.layers.fused_moe.config import ( RoutingMethodType, + get_routing_method_type, ) from aphrodite.model_executor.layers.fused_moe.router.aiter_shared_routed_fused_moe_router import ( # noqa: E501 AiterSharedRoutedFusedMoERouter, @@ -67,7 +68,8 @@ def create_fused_moe_router( The selection logic follows this priority order: 1. RoutingSimulatorRouter - if APHRODITE_MOE_ROUTING_SIMULATION_STRATEGY env var is set 2. ZeroExpertRouter - if zero_expert_type is not None - 3. GroupedTopKRouter - if use_grouped_topk is True + 3. GroupedTopKRouter - if use_grouped_topk is True and the grouping is not + degenerate (at most one group, with topk_group <= 1) 4. CustomRoutingRouter - if custom_routing_function is not None 5. FusedTopKBiasRouter - if e_score_correction_bias is not None 6. AiterSharedRoutedFusedMoERouter - if num_fused_shared_experts > 0 @@ -136,30 +138,40 @@ def create_fused_moe_router( assert custom_routing_function is None if num_expert_group is None or topk_group is None: raise ValueError("num_expert_group and topk_group must be provided when use_grouped_topk is True") - grouped_topk_router = GroupedTopKRouter( - top_k=top_k, - global_num_experts=global_num_experts, - eplb_state=eplb_state, - num_expert_group=num_expert_group, - topk_group=topk_group, - renormalize=renormalize, - scoring_func=scoring_func, - routed_scaling_factor=routed_scaling_factor, - e_score_correction_bias=e_score_correction_bias, - num_fused_shared_experts=num_fused_shared_experts, - ) - if ( - grouped_topk_router.routing_method_type != RoutingMethodType.Unspecified - or num_expert_group > 1 - or topk_group > 1 - ): - return grouped_topk_router - - # If routing_method for GroupedTopKRouter is Unspecified and there is only - # one group, fallback to standard top-k routing - use_grouped_topk = False - num_expert_group = None - topk_group = None + + # For topk_group <= 1, grouped implementation is pure overhead. + degenerate_grouping = num_expert_group <= 1 and topk_group <= 1 + # FusedTopKRouter cannot apply routed_scaling_factor, FusedTopKBiasRouter can. + scaling_handled_downstream = routed_scaling_factor == 1.0 or e_score_correction_bias is not None + + # Degenerating must not change the advertised routing method, which drives + # kernel selection. num_expert_group only affects it for biased routing. + def advertised_routing_method(groups: int | None) -> RoutingMethodType: + return get_routing_method_type( + scoring_func=scoring_func, + top_k=top_k, + renormalize=renormalize, + num_expert_group=groups, + has_e_score_bias=e_score_correction_bias is not None, + routed_scaling_factor=routed_scaling_factor, + ) + + routing_method_preserved = advertised_routing_method(num_expert_group) == advertised_routing_method(None) + + if not (degenerate_grouping and scaling_handled_downstream and routing_method_preserved): + return GroupedTopKRouter( + top_k=top_k, + global_num_experts=global_num_experts, + eplb_state=eplb_state, + num_expert_group=num_expert_group, + topk_group=topk_group, + renormalize=renormalize, + scoring_func=scoring_func, + routed_scaling_factor=routed_scaling_factor, + e_score_correction_bias=e_score_correction_bias, + num_fused_shared_experts=num_fused_shared_experts, + ) + # Otherwise fall through to the non-grouped chain below. if custom_routing_function is not None: return CustomRoutingRouter( diff --git a/aphrodite/model_executor/layers/fused_moe/runner/latent_moe_runner.py b/aphrodite/model_executor/layers/fused_moe/runner/latent_moe_runner.py new file mode 100644 index 0000000000..afdf8ef4a1 --- /dev/null +++ b/aphrodite/model_executor/layers/fused_moe/runner/latent_moe_runner.py @@ -0,0 +1,222 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import torch + +import aphrodite.envs as envs +from aphrodite.config import get_current_aphrodite_config as get_current_aphrodite_config +from aphrodite.distributed import tensor_model_parallel_all_reduce +from aphrodite.logger import init_logger +from aphrodite.model_executor.layers.fused_allreduce_gemma_rms_norm import ( + _AR_RESIDUAL_RMS_NORM, + _can_use_flashinfer, + flashinfer_trtllm_fused_allreduce_norm, +) +from aphrodite.model_executor.layers.layernorm import RMSNorm +from aphrodite.utils.torch_utils import aux_stream, current_stream + +from .moe_runner import MoERunner, _unpack + +logger = init_logger(__name__) + + +class LatentMoERunner(MoERunner): + """MoE runner for latent MoE with a replicated routed up-projection. + + Fused path (tp>1, un-reduced combine output, shared expert, no SP): + concatenates the un-reduced latent partial (dim d) and the un-reduced + shared partial (dim D) into one contiguous buffer, all-reduces once, then + splits. The latent half is normed and up-projected locally (replicated + up-proj -> full hidden), and the shared add folds into the GEMM epilogue + (``torch.addmm``). One collective total, no post-reduction communication. + + Native path: the replicated up-proj produces the full hidden dim on every + rank, so the base runner combines routed + shared correctly at any TP size + (using two collectives instead of the fused path's one). + """ + + def __init__( + self, + *args, + enable_k3_latent_moe_tail_fusion: bool = False, + **kwargs, + ) -> None: + super().__init__(*args, **kwargs) + self.enable_k3_latent_moe_tail_fusion = enable_k3_latent_moe_tail_fusion + use_fused_path = self._use_fused_path() + if self.enable_k3_latent_moe_tail_fusion and use_fused_path and self.moe_config.tp_size not in (8, 16): + logger.warning_once( + "K3 latent-MoE tail fusion currently supports TP=8 and TP=16, " + "but TP=%d is configured. Falling back to the default path.", + self.moe_config.tp_size, + ) + self.enable_k3_latent_moe_tail_fusion = False + + if self.enable_k3_latent_moe_tail_fusion and use_fused_path: + aphrodite_config = get_current_aphrodite_config() + if aphrodite_config.parallel_config.use_ubatching: + raise ValueError("K3 latent-MoE tail fusion does not support DBO or ubatching.") + if aphrodite_config.model_config.enable_sleep_mode: + raise ValueError("K3 latent-MoE tail fusion does not support sleep mode.") + transform = self.routed_output_transform + assert transform is not None + norm = transform.norm + assert norm is not None + from aphrodite.models.kimi_k3.nvidia.ops.latent_moe_tail import ( + KimiK3LatentMoETailOp, + ) + + op = KimiK3LatentMoETailOp.initialize( + hidden_size=transform.up_proj.weight.shape[0], + latent_size=norm.weight.shape[0], + dtype=norm.weight.dtype, + device=norm.weight.device, + rms_eps=norm.variance_epsilon, + ) + self._k3_latent_moe_tail_op = op + + def _get_zero_residual(self, hidden_states: torch.Tensor) -> torch.Tensor: + """Read-only zero ``residual_in`` for the fused AR+RMSNorm kernel. + + flashinfer requires a residual buffer even when there is no residual to + add. + """ + numel = hidden_states.numel() + buf = getattr(self, "_zero_residual", None) + if buf is None or buf.dtype != hidden_states.dtype or buf.device != hidden_states.device or buf.numel() < numel: + buf = torch.zeros(numel, dtype=hidden_states.dtype, device=hidden_states.device) + self._zero_residual = buf + return buf[:numel].view_as(hidden_states) + + def _use_fused_path(self) -> bool: + # The fused path merges the latent and shared reductions into one + # all-reduce, so it needs actual TP parallelism, a shared expert (to + # concat), an un-reduced combine output, and no sequence parallelism. + return ( + self.moe_config.tp_size > 1 + and self._shared_experts is not None + and not self._fused_output_is_reduced + and not self.moe_config.is_sequence_parallel + ) + + def forward( + self, + hidden_states: torch.Tensor, + router_logits: torch.Tensor, + input_ids: torch.Tensor | None = None, + shared_experts_input: torch.Tensor | None = None, + ) -> torch.Tensor: + if self._use_fused_path(): + return self._fused_forward(hidden_states, router_logits, input_ids, shared_experts_input) + return super().forward(hidden_states, router_logits, input_ids, shared_experts_input) + + def _fused_forward( + self, + hidden_states: torch.Tensor, + router_logits: torch.Tensor, + input_ids: torch.Tensor | None = None, + shared_experts_input: torch.Tensor | None = None, + ) -> torch.Tensor: + # When the caller pre-applies the routed input transform outside the + # runner (e.g. to overlap it on a separate stream), it passes the + # already-transformed routed input as ``hidden_states`` and the original + # hidden states as ``shared_experts_input``; skip the transform then. + if shared_experts_input is None: + hidden_states, shared_experts_input = self.apply_routed_input_transform(hidden_states) + hidden_states, og_hidden_dim_pre_xform, og_hidden_dim_post_xform = self._maybe_pad_hidden_states( + shared_experts_input, + hidden_states, + ) + + result = self._forward_entry( + hidden_states, + router_logits, + shared_experts_input, + input_ids, + self._encode_layer_name(), + self.moe_config.hidden_dim_unpadded if self._quant_method.has_unpadded_output else 0, + ) + + shared_output, fused_output = _unpack(result) + assert shared_output is not None + + if og_hidden_dim_pre_xform is not None: + fused_output = fused_output[..., :og_hidden_dim_pre_xform] + + transform = self.routed_output_transform + assert transform is not None + + if self.enable_k3_latent_moe_tail_fusion: + op = self._k3_latent_moe_tail_op + if 0 < fused_output.shape[0] <= op.contract.max_num_tokens: + norm = transform.norm + assert norm is not None + result = op( + fused_output, + shared_output, + norm.weight, + transform.up_proj.weight, + ) + result = self._maybe_reduce_final_output(result, og_hidden_dim_post_xform, output_is_reduced=True) + return self._maybe_add_zero_expert_output(result) + + fused_latent = None + if transform.norm is not None: + fused_latent = self.allreduce_norm_latent_out(fused_output, transform.norm) + else: + fused_latent = tensor_model_parallel_all_reduce(fused_output) + + shared_expert_stream = ( + aux_stream() if shared_output.size(0) <= envs.APHRODITE_SHARED_EXPERTS_STREAM_TOKEN_THRESHOLD else None + ) + if shared_expert_stream is not None: + # overlap shared expert allreduce with latent up_proj + main = current_stream() + shared_output.record_stream(shared_expert_stream) + shared_expert_stream.wait_stream(main) + with torch.cuda.stream(shared_expert_stream): + shared_output = tensor_model_parallel_all_reduce(shared_output) + result = torch.mm(fused_latent, transform.up_proj.weight.t()) + main.wait_stream(shared_expert_stream) + else: + shared_output = tensor_model_parallel_all_reduce(shared_output) + result = torch.mm(fused_latent, transform.up_proj.weight.t()) + result.add_(shared_output) + + # Output is already fully reduced; this only strips padding. + result = self._maybe_reduce_final_output(result, og_hidden_dim_post_xform, output_is_reduced=True) + + return self._maybe_add_zero_expert_output(result) + + def allreduce_norm_latent_out( + self, + hidden_states: torch.Tensor, + norm: RMSNorm, + ) -> tuple[torch.Tensor, torch.Tensor]: + """All-reduce + add residual + (standard) RMSNorm, fused via flashinfer.""" + if self.moe_config.tp_size == 1: + return norm(hidden_states) + + if flashinfer_trtllm_fused_allreduce_norm is not None: + ok, max_token_num = _can_use_flashinfer(hidden_states, self.moe_config.tp_size) + if ok: + norm_out = torch.empty_like(hidden_states) + # With norm_out provided, the kernel writes the new residual + # (all_reduce(hidden_states) + residual) into the hidden_states + # buffer and the normalized result into norm_out. + flashinfer_trtllm_fused_allreduce_norm( + allreduce_in=hidden_states, + residual=self._get_zero_residual(hidden_states), + rms_gamma=norm.weight, + rms_eps=norm.variance_epsilon, + world_size=self.moe_config.tp_size, + weight_bias=0.0, + launch_with_pdl=True, + fp32_acc=True, + max_token_num=max_token_num, + pattern_code=_AR_RESIDUAL_RMS_NORM, + norm_out=norm_out, + ) + return norm_out + + reduced = tensor_model_parallel_all_reduce(hidden_states) + return norm(reduced) diff --git a/aphrodite/model_executor/layers/fused_moe/runner/moe_runner.py b/aphrodite/model_executor/layers/fused_moe/runner/moe_runner.py index ab012a71ce..f1eac5d206 100644 --- a/aphrodite/model_executor/layers/fused_moe/runner/moe_runner.py +++ b/aphrodite/model_executor/layers/fused_moe/runner/moe_runner.py @@ -341,7 +341,10 @@ def _maybe_fuse_gate_weights(self): def _quant_method(self) -> FusedMoEMethodBase: return self.routed_experts.quant_method - def apply_routed_input_transform(self, hidden_states: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor | None]: + def apply_routed_input_transform( + self, + hidden_states: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor | None]: """Apply transform for routed experts (e.g., latent projection). This is called by FusedMoE.forward_native. The original hidden_states @@ -405,6 +408,7 @@ def _fused_output_is_reduced(self) -> bool: def _maybe_reduce_shared_expert_output( self, shared_output: torch.Tensor | None, + fused_output_is_reduced: bool | None = None, ) -> torch.Tensor | None: """All-reduce shared expert output when the combine kernel already reduced fused output. @@ -414,14 +418,34 @@ def _maybe_reduce_shared_expert_output( * If we have SP (TP=N, DP=M, EP), there is a separate AG step handled in the model. """ - if shared_output is not None and not self.moe_config.is_sequence_parallel and self._fused_output_is_reduced: + if fused_output_is_reduced is None: + fused_output_is_reduced = self._fused_output_is_reduced + + if shared_output is not None and not self.moe_config.is_sequence_parallel and fused_output_is_reduced: shared_output = tensor_model_parallel_all_reduce(shared_output) return shared_output + def _maybe_reduce_routed_output_before_transform( + self, + fused_output: torch.Tensor, + fused_output_is_reduced: bool, + ) -> tuple[torch.Tensor, bool]: + """All-reduce latent routed output before its output transform.""" + if ( + self.routed_output_transform is not None + and not self.moe_config.is_sequence_parallel + and (self.moe_config.tp_size > 1 or self.moe_config.ep_size > 1) + and not fused_output_is_reduced + ): + fused_output = tensor_model_parallel_all_reduce(fused_output) + fused_output_is_reduced = True + return fused_output, fused_output_is_reduced + def _maybe_reduce_final_output( self, states: torch.Tensor, trunc_size: int | None, + output_is_reduced: bool | None = None, ) -> torch.Tensor: """All-reduce the combined output if needed. @@ -438,11 +462,14 @@ def _maybe_reduce_final_output( # We don't need to reduce the final output if: # - We are not running with TP or DP # - The MK already reduced the fused output itself. + if output_is_reduced is None: + output_is_reduced = self._fused_output_is_reduced + if ( not self.moe_config.is_sequence_parallel and not self.moe_config.skip_final_all_reduce and (self.moe_config.tp_size > 1 or self.moe_config.ep_size > 1) - and not self._fused_output_is_reduced + and not output_is_reduced ): states = tensor_model_parallel_all_reduce(states) @@ -612,6 +639,7 @@ def forward( hidden_states: torch.Tensor, router_logits: torch.Tensor, input_ids: torch.Tensor | None = None, + shared_experts_input: torch.Tensor | None = None, ) -> torch.Tensor: """Invoke the fused moe layer. @@ -635,7 +663,8 @@ def forward( # Apply transform for routed experts (e.g., latent projection # for latent MoE) - hidden_states, shared_experts_input = self.apply_routed_input_transform(hidden_states) + if shared_experts_input is None: + hidden_states, shared_experts_input = self.apply_routed_input_transform(hidden_states) # Record before `_maybe_pad_hidden_states` pads activations to match # `moe_config.hidden_dim`, e.g. after `align_trtllm_fp4_moe_hidden_dim_for_fi` @@ -671,9 +700,18 @@ def forward( if og_hidden_dim_pre_xform is not None: fused_output = fused_output[..., :og_hidden_dim_pre_xform] - # If combine kernel already reduced fused, reduce shared to match. + fused_output_is_reduced = self._fused_output_is_reduced + + # Latent routed output has to be reduced before output transform, + # because the transform may include non-linear normalization. + fused_output, fused_output_is_reduced = self._maybe_reduce_routed_output_before_transform( + fused_output, + fused_output_is_reduced, + ) + + # If routed output is already reduced, reduce shared to match. # See note above re: the two all-reduce points. - shared_output = self._maybe_reduce_shared_expert_output(shared_output) + shared_output = self._maybe_reduce_shared_expert_output(shared_output, fused_output_is_reduced) shared_output, fused_output = self._maybe_apply_routed_scale_to_output(shared_output, fused_output) @@ -685,7 +723,7 @@ def forward( else: result = fused_output - result = self._maybe_reduce_final_output(result, og_hidden_dim_post_xform) + result = self._maybe_reduce_final_output(result, og_hidden_dim_post_xform, fused_output_is_reduced) return self._maybe_add_zero_expert_output(result) diff --git a/aphrodite/model_executor/layers/fused_moe/runner/moe_runner_interface.py b/aphrodite/model_executor/layers/fused_moe/runner/moe_runner_interface.py index 3f3aa9177a..bc77fb3b56 100644 --- a/aphrodite/model_executor/layers/fused_moe/runner/moe_runner_interface.py +++ b/aphrodite/model_executor/layers/fused_moe/runner/moe_runner_interface.py @@ -36,6 +36,7 @@ def forward( hidden_states: torch.Tensor, router_logits: torch.Tensor, input_ids: torch.Tensor | None = None, + shared_experts_input: torch.Tensor | None = None, ) -> torch.Tensor: raise NotImplementedError diff --git a/aphrodite/model_executor/layers/fused_moe/utils.py b/aphrodite/model_executor/layers/fused_moe/utils.py index d8deb7ad2a..1c27dc52dd 100644 --- a/aphrodite/model_executor/layers/fused_moe/utils.py +++ b/aphrodite/model_executor/layers/fused_moe/utils.py @@ -7,7 +7,9 @@ import torch import torch.nn.functional as F +import aphrodite.envs as envs from aphrodite import _custom_ops as ops +from aphrodite.logger import init_logger from aphrodite.model_executor.layers.quantization.utils.fp8_utils import ( per_token_group_quant_fp8, ) @@ -39,6 +41,8 @@ if TYPE_CHECKING: from aphrodite.model_executor.layers.fused_moe.config import FusedMoEConfig +logger = init_logger(__name__) + @triton.jit def _count_expert_num_tokens( @@ -570,3 +574,52 @@ def swiglu_limit_func( @functools.lru_cache def enable_swap_ab(BLOCK_SIZE_M: int, BLOCK_SIZE_N: int) -> bool: return current_platform.is_device_capability(90) and BLOCK_SIZE_M < 64 and BLOCK_SIZE_N >= 64 + + +def moe_use_td_hw_supported() -> bool: + """Return whether the current device supports the fused MoE TD path.""" + if current_platform.is_xpu(): + return True + if current_platform.is_cuda(): + return current_platform.has_device_capability(100) + return False + + +def resolve_moe_use_td() -> bool: + """Resolve the tri-state ``APHRODITE_TRITON_USE_TD`` setting for MoE.""" + override = envs.APHRODITE_TRITON_USE_TD + if override is None: + return current_platform.is_xpu() + return override + + +_warned_moe_use_td_ineffective = False + + +def warn_if_moe_use_td_ineffective(active_backend: str, is_quantized: bool = False) -> None: + """Warn once when the tensor-descriptor override cannot take effect.""" + global _warned_moe_use_td_ineffective + if _warned_moe_use_td_ineffective: + return + if envs.APHRODITE_TRITON_USE_TD is None: + return + is_triton = active_backend.upper() == "TRITON" + if is_triton and not is_quantized: + return + if not is_triton: + reason = ( + f"the active MoE backend is {active_backend!r}; pass " + "`--moe-backend triton` to enable the tensor-descriptor path" + ) + else: + reason = ( + "the model uses quantized MoE weights; the TD path is " + "currently restricted to non-quantized weights and falls " + "back to the pointer path" + ) + logger.warning( + "APHRODITE_TRITON_USE_TD is set to %s but %s.", + envs.APHRODITE_TRITON_USE_TD, + reason, + ) + _warned_moe_use_td_ineffective = True diff --git a/aphrodite/model_executor/layers/linear.py b/aphrodite/model_executor/layers/linear.py index bcbd31bbe9..caa92f58c9 100644 --- a/aphrodite/model_executor/layers/linear.py +++ b/aphrodite/model_executor/layers/linear.py @@ -50,6 +50,7 @@ "UnquantizedLinearMethod", "CompressedTensorsLinearMethod", "CompressedTensorsLinearTransformMethod", + "QutlassNvFP4LinearMethod", "AutoAWQMarlinLinearMethod", "AutoAWQLinearMethod", "AutoGPTQLinearMethod", diff --git a/aphrodite/model_executor/layers/mamba/ops/causal_conv1d.py b/aphrodite/model_executor/layers/mamba/ops/causal_conv1d.py index 412744f7f1..fa3a170917 100644 --- a/aphrodite/model_executor/layers/mamba/ops/causal_conv1d.py +++ b/aphrodite/model_executor/layers/mamba/ops/causal_conv1d.py @@ -1009,6 +1009,7 @@ def causal_conv1d_update( block_idx_last_scheduled_token: torch.Tensor | None = None, initial_state_idx: torch.Tensor | None = None, validate_data=False, + out: torch.Tensor | None = None, ): """ x: Input tensor which can take the following shapes: @@ -1046,7 +1047,8 @@ def causal_conv1d_update( for example: conv_state_indices = [null_block_id, 1, 20, null_block_id] in this case, the kernel will not process entries at indices 0 and 3 - out: (batch, dim) or (batch, dim, seqlen) or (num_tokens, dim), same shape as `x` + out: optional output tensor with the same shape as `x`. When omitted, + the input is overwritten. """ if validate_data: assert null_block_id is not None @@ -1058,10 +1060,18 @@ def causal_conv1d_update( original_x_dtype = x.dtype x = x.to(conv_state.dtype) + if out is None: + out = x + else: + if out.shape != x.shape: + raise ValueError(f"`out` shape {tuple(out.shape)} must match `x` shape {tuple(x.shape)}.") + if out.dtype != original_x_dtype or out.device != x.device: + raise ValueError("`out` must have the same dtype and device as the input `x`.") unsqueeze = query_start_loc is None and x.dim() == 2 if unsqueeze: # make it (batch, dim, seqlen) with seqlen == 1 x = x.unsqueeze(-1) + out = out.unsqueeze(-1) if query_start_loc is None: batch, dim, seqlen = x.shape else: diff --git a/aphrodite/model_executor/layers/mamba/ops/gather_initial_states.py b/aphrodite/model_executor/layers/mamba/ops/gather_initial_states.py new file mode 100644 index 0000000000..0379955a32 --- /dev/null +++ b/aphrodite/model_executor/layers/mamba/ops/gather_initial_states.py @@ -0,0 +1,81 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import torch + +from aphrodite.platforms import current_platform +from aphrodite.triton_utils import tl, triton + + +@triton.jit +def _gather_initial_states_kernel( + state_ptr, + indices_ptr, + has_initial_state_ptr, + output_ptr, + stride_state_batch, + stride_indices, + stride_has_initial_state, + row_size: tl.constexpr, + BLOCK_SIZE: tl.constexpr, + launch_pdl: tl.constexpr, +): + block_idx = tl.program_id(0) + batch_idx = tl.program_id(1) + offsets = block_idx * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + mask = offsets < row_size + + if launch_pdl: + tl.extra.cuda.gdc_wait() + tl.extra.cuda.gdc_launch_dependents() + + has_initial_state = tl.load(has_initial_state_ptr + batch_idx * stride_has_initial_state).to(tl.int1) + state_idx = tl.load(indices_ptr + batch_idx * stride_indices).to(tl.int64) + state_idx = tl.where(has_initial_state, state_idx, 0) + values = tl.load( + state_ptr + state_idx * stride_state_batch + offsets, + mask=mask & has_initial_state, + other=0.0, + ) + tl.store(output_ptr + batch_idx * row_size + offsets, values, mask=mask) + + +def gather_initial_states( + state: torch.Tensor, + indices: torch.Tensor, + has_initial_state: torch.Tensor, +) -> torch.Tensor: + """Gather dense state rows, replacing uninitialized rows with zeros.""" + assert state.ndim >= 2 + assert state.is_cuda + assert indices.ndim == 1 and has_initial_state.ndim == 1 + assert indices.shape == has_initial_state.shape + assert indices.device == state.device + assert has_initial_state.device == state.device + assert indices.dtype in (torch.int32, torch.int64) + assert has_initial_state.dtype == torch.bool + + row_size = state[0].numel() + # Mamba pages may pad stride(0), but each state row remains dense. + assert state[0].is_contiguous() + output = torch.empty( + (indices.numel(), *state.shape[1:]), + dtype=state.dtype, + device=state.device, + ) + block_size = min(triton.next_power_of_2(row_size), 1024) + grid = (triton.cdiv(row_size, block_size), indices.numel()) + _gather_initial_states_kernel[grid]( + state, + indices, + has_initial_state, + output, + state.stride(0), + indices.stride(0), + has_initial_state.stride(0), + row_size=row_size, + BLOCK_SIZE=block_size, + num_warps=8, + launch_pdl=current_platform.is_arch_support_pdl(), + ) + return output diff --git a/aphrodite/model_executor/layers/quantization/compressed_tensors/transform/linear.py b/aphrodite/model_executor/layers/quantization/compressed_tensors/transform/linear.py index 66f532361e..baf47e0eb1 100644 --- a/aphrodite/model_executor/layers/quantization/compressed_tensors/transform/linear.py +++ b/aphrodite/model_executor/layers/quantization/compressed_tensors/transform/linear.py @@ -25,6 +25,7 @@ from aphrodite.model_executor.layers.quantization.compressed_tensors.transform.utils import ( # noqa: E501 TransformTuple, ) +from aphrodite.platforms import current_platform class CompressedTensorsLinearTransformMethod(LinearMethodBase): @@ -48,11 +49,10 @@ def from_schemes( assert input_tfms or output_tfms - if is_qutlass_fp4_scheme(quant_scheme, input_tfms): + if is_qutlass_fp4_scheme(quant_scheme, input_tfms) and current_platform.has_device_capability(100): return QutlassNvFP4LinearMethod(quant_method, input_tfms, output_tfms) # hadacore or dense gemm is selected by Transform module - return cls(quant_method, input_tfms, output_tfms) def __init__( diff --git a/aphrodite/model_executor/layers/quantization/compressed_tensors/transform/module.py b/aphrodite/model_executor/layers/quantization/compressed_tensors/transform/module.py index 1561547208..7822d78c26 100644 --- a/aphrodite/model_executor/layers/quantization/compressed_tensors/transform/module.py +++ b/aphrodite/model_executor/layers/quantization/compressed_tensors/transform/module.py @@ -32,7 +32,7 @@ class HadamardTransform(torch.nn.Module): transforms: dict[int, TransformTuple] # info parsed from transforms config weight: SharedWeightParameter # container for shared tensors - scales: dict[int, float] # hadamard scale, usually sqrt(matrix.size(0)) + scaled_data_ptrs: set[int] = set() def __init__( self, @@ -44,7 +44,6 @@ def __init__( ): super().__init__() self.transforms = transforms - self.scales = {} if get_tensor_model_parallel_world_size() > 1: raise NotImplementedError("Online transforms with tensor parallelism is not supported") @@ -60,26 +59,26 @@ def __init__( weight_size = self._get_weight_size(layer, scheme, args, input_size, output_size) data_key = self._get_data_key(scheme, weight_size) + # load up in model's default precision, rather than using scheme.precision self.weight.add_partition( part_index, data_key, size=(weight_size, weight_size), - dtype=scheme.precision, ) # validate that shared tensors and schemes are correct self._validate_input_transforms() def process_weights_after_loading(self): - for part_id in self.weight.partitions: - data = self.weight.partitions[part_id].data - + for part_id, partition in self.weight.partitions.items(): # required by torch.compile self.weight.process_weights_after_loading() - # precompute scale as a runtime multiply, not division - # do not fold into weight in order to utilize FWHT - self.scales[part_id] = 1 / math.sqrt(data.size(0)) + # Merge normalization scale directly into weight, must be done only once + data_ptr = partition.data.data_ptr() + if data_ptr not in HadamardTransform.scaled_data_ptrs: + partition.data.div_(math.sqrt(partition.data.size(0))) + HadamardTransform.scaled_data_ptrs.add(data_ptr) # FUTURE: avoid runtime transpose by processing weights # prior to apply @@ -105,16 +104,15 @@ def forward(self, value: Tensor, part_id: int = 0) -> Tensor: else: weight = self.weight.partitions[part_id] weight = weight if self.transforms[part_id].args.inverse else weight.T # linear := x(W.T) - scale = self.scales[part_id] if self.transforms[part_id].scheme.head_dim is not None: value = value.unflatten(-1, (-1, weight.size(0))) - value = dispatch_unquantized_gemm()(self, value.to(weight.dtype), weight, None).to(value.dtype) * scale + value = dispatch_unquantized_gemm()(self, value.to(weight.dtype), weight, None).to(value.dtype) value = value.flatten(-2, -1) return value - return dispatch_unquantized_gemm()(self, value.to(weight.dtype), weight, None).to(value.dtype) * scale + return dispatch_unquantized_gemm()(self, value.to(weight.dtype), weight, None).to(value.dtype) def _get_data_key(self, scheme: TransformScheme, weight_size: int) -> Hashable: return (id(scheme), weight_size) diff --git a/aphrodite/model_executor/layers/quantization/compressed_tensors/transform/schemes/linear_qutlass_nvfp4.py b/aphrodite/model_executor/layers/quantization/compressed_tensors/transform/schemes/linear_qutlass_nvfp4.py index ab98b965df..e4781e34ce 100644 --- a/aphrodite/model_executor/layers/quantization/compressed_tensors/transform/schemes/linear_qutlass_nvfp4.py +++ b/aphrodite/model_executor/layers/quantization/compressed_tensors/transform/schemes/linear_qutlass_nvfp4.py @@ -2,7 +2,13 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import torch +from torch.nn.parameter import Parameter +from aphrodite._custom_ops import fusedQuantizeNv +from aphrodite.model_executor.kernels.linear import ( + _LINEAR_BACKEND_KERNEL_MAP, + NvFp4LinearKernel, +) from aphrodite.model_executor.layers.quantization.compressed_tensors.compressed_tensors import ( # noqa: E501 CompressedTensorsScheme, CompressedTensorsW4A4Fp4, @@ -11,18 +17,29 @@ CompressedTensorsLinearTransformMethod, TransformTuple, ) +from aphrodite.model_executor.layers.quantization.qutlass_utils import to_blocked +from aphrodite.model_executor.layers.quantization.utils.nvfp4_utils import ( + slice_nvfp4_output, +) +from aphrodite.utils.flashinfer import ( + flashinfer_scaled_fp4_mm, +) __all__ = ["is_qutlass_fp4_scheme", "QutlassNvFP4LinearMethod"] +NVFP4_MAX = 6.0 + +# QUTLASS supports transform block sizes (16, 32, 64, 128) for NVFP4 +# https://github.com/IST-DASLab/qutlass/blob/v0.2.0/qutlass/csrc/bindings.cpp#L413-L414 def is_qutlass_fp4_scheme( quant_scheme: CompressedTensorsScheme | None, input_tfms: dict[int, TransformTuple], ) -> bool: return ( - isinstance(quant_scheme, (CompressedTensorsW4A4Fp4,)) - and len(input_tfms) == 1 - and input_tfms[0].scheme.head_dim == quant_scheme.group_size + isinstance(quant_scheme, CompressedTensorsW4A4Fp4) + and len(input_tfms) >= 1 + and all(input_tfm.scheme.head_dim in (16, 32, 64, 128) for input_tfm in input_tfms.values()) ) @@ -50,15 +67,81 @@ def create_weights( ) assert self.input_transform is not None - assert len(self.input_transform.weight) == 1 - assert self.input_transform.weight[0].size(0) == layer.scheme.group_size + assert len(self.input_transform.weight.partitions) >= 1 return ret + @staticmethod + def _get_flashinfer_gemm_backend(kernel: NvFp4LinearKernel) -> str: + """ + Given a kernel, find the string that is needed to be passed into + `flashinfer_scaled_fp4_mm`, using + aphrodite.model_executor.kernels.linear._LINEAR_BACKEND_KERNEL_MAP as source of truth + """ + kernel_type = type(kernel) + for key, kernels in _LINEAR_BACKEND_KERNEL_MAP.items(): + if not key.startswith("flashinfer_") or kernel_type not in kernels: + continue + backend = key.removeprefix("flashinfer_") + # flashinfer GEMM backend uses "cute-dsl", not "cutedsl" + return backend.replace("cutedsl", "cute-dsl") + raise ValueError(f"QutlassNvFP4 transform requires a FlashInfer kernel, got {kernel_type.__name__}") + + def process_weights_after_loading(self, layer): + super().process_weights_after_loading(layer) + + assert self.input_transform is not None + layer.hadamard_matrix = self.input_transform.weight.partitions[0].data + + # fusedQuantizeNv stores raw absmax as block scales (sf = absmax), + # while CT weights use sf = absmax * SFScaleVal / 6.0. The GEMM + # computes alpha * sum(fp4_a * sf_a * fp4_w * sf_w), so alpha must + # compensate: alpha = weight_global_scale / 6.0 + layer.fused_alpha = Parameter(layer.weight_global_scale / NVFP4_MAX, requires_grad=False) + + layer.fused_global_scale = Parameter( + torch.tensor( + [NVFP4_MAX], + dtype=torch.float32, + device=layer.weight_global_scale.device, + ), + requires_grad=False, + ) + + layer.flashinfer_gemm_backend = self._get_flashinfer_gemm_backend(layer.scheme.kernel) + def apply( self, layer: torch.nn.Module, x: torch.Tensor, bias: torch.Tensor | None = None, ) -> torch.Tensor: - raise NotImplementedError() + assert bias is None + output_size = layer.output_size_per_partition + output_shape = [*x.shape[:-1], output_size] + + x_flat = x.contiguous().flatten(end_dim=-2) + + x_fp4, x_scales = fusedQuantizeNv(x_flat, layer.hadamard_matrix, layer.fused_global_scale) + + x_scales_blocked = to_blocked(x_scales, backend="triton").view(x_scales.shape) + + out = flashinfer_scaled_fp4_mm( + x_fp4, + layer.weight, + x_scales_blocked, + layer.weight_scale, + layer.fused_alpha, + x.dtype, + backend=layer.flashinfer_gemm_backend, + ) + + out = slice_nvfp4_output(out, output_size) + + if self.output_transform is not None: + for part_id, (start, length) in enumerate(self.partition_ranges): + out[:, start : start + length] = self.output_transform( + out[:, start : start + length].clone(), part_id=part_id + ) + + return out.view(*output_shape) diff --git a/aphrodite/model_executor/layers/quantization/qutlass_utils.py b/aphrodite/model_executor/layers/quantization/qutlass_utils.py index 542c5baac8..408e53b222 100644 --- a/aphrodite/model_executor/layers/quantization/qutlass_utils.py +++ b/aphrodite/model_executor/layers/quantization/qutlass_utils.py @@ -84,6 +84,7 @@ def triton_scale_swizzle( ) +@torch.library.custom_op("aphrodite::triton_mx_block_rearrange", mutates_args=()) def triton_mx_block_rearrange(scale_tensor: torch.Tensor) -> torch.Tensor: """ Rearranges an E8M0 tensor scale from row-major format to @@ -140,6 +141,14 @@ def triton_mx_block_rearrange(scale_tensor: torch.Tensor) -> torch.Tensor: return out +@triton_mx_block_rearrange.register_fake +def _triton_mx_block_rearrange_fake(scale_tensor: torch.Tensor) -> torch.Tensor: + rows, cols = scale_tensor.shape + padded_rows = cdiv(rows, 128) * 128 + padded_cols = cdiv(cols, 4) * 4 + return scale_tensor.new_empty((padded_rows, padded_cols)) + + def to_blocked(input_matrix: torch.Tensor, backend: Literal["torch", "triton"] = "triton") -> torch.Tensor: """ Rearrange a large matrix by breaking it into blocks and applying @@ -153,7 +162,7 @@ def to_blocked(input_matrix: torch.Tensor, backend: Literal["torch", "triton"] = backend: "torch" (PyTorch path) or "triton" (Triton kernel) Returns: - Rearranged tensor of shape (32*cdiv(H,128), 16*cdiv(W,4)) + Rearranged flattened tensor of size (32*cdiv(H,128) * 16*cdiv(W,4)) """ if backend == "triton": return triton_mx_block_rearrange(input_matrix).flatten() 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 3fd3fba1dc..8dc06130b3 100644 --- a/aphrodite/model_executor/layers/quantization/utils/flashinfer_fp4_moe.py +++ b/aphrodite/model_executor/layers/quantization/utils/flashinfer_fp4_moe.py @@ -99,7 +99,8 @@ def prepare_nvfp4_moe_layer_for_flashinfer_cutedsl( """Prepare weights for the CuteDSL wrapper-based NvFP4 MoE backend. Converts weight scale factors to MMA layout expected by CuteDslMoEWrapper, - and interleaves w13 gate/linear rows. + and interleaves w13 gate/linear rows for gated activations. Non-gated + activations use a single w13 projection and keep its row order unchanged. """ from flashinfer.cute_dsl.utils import convert_sf_to_mma_layout @@ -108,13 +109,14 @@ def prepare_nvfp4_moe_layer_for_flashinfer_cutedsl( a13_scale = a13_scale.max().to(torch.float32).repeat(num_experts) a2_scale = a2_scale.max().to(torch.float32).repeat(num_experts) - half = w13.shape[1] // 2 - w13 = torch.cat([w13[:, half:], w13[:, :half]], dim=1) - w13_scale = torch.cat([w13_scale[:, half:], w13_scale[:, :half]], dim=1) + if layer.activation.is_gated: + half = w13.shape[1] // 2 + w13 = torch.cat([w13[:, half:], w13[:, :half]], dim=1) + w13_scale = torch.cat([w13_scale[:, half:], w13_scale[:, :half]], dim=1) - # Interleave up/gate rows for w13 weights and scales. - w13 = interleave_linear_and_gate(w13, group_size=64, dim=1) - w13_scale = interleave_linear_and_gate(w13_scale, group_size=64, dim=1) + # Interleave up/gate rows for w13 weights and scales. + w13 = interleave_linear_and_gate(w13, group_size=64, dim=1) + w13_scale = interleave_linear_and_gate(w13_scale, group_size=64, dim=1) # Convert w13 scale factors: linear → swizzled → MMA layout. w13_scale = swizzle_blockscale(w13_scale) diff --git a/aphrodite/model_executor/models/config.py b/aphrodite/model_executor/models/config.py index 5ba4246c62..d2c7e07f0e 100644 --- a/aphrodite/model_executor/models/config.py +++ b/aphrodite/model_executor/models/config.py @@ -712,6 +712,20 @@ def verify_and_update_config(aphrodite_config: "AphroditeConfig") -> None: ) +class Qwen3_5ForCausalLMConfig(Qwen3_5ForConditionalGenerationConfig): + @staticmethod + def verify_and_update_config(aphrodite_config: "AphroditeConfig") -> None: + Qwen3_5ForConditionalGenerationConfig.verify_and_update_config(aphrodite_config) + + # Text-only Qwen3.5 models use one-dimensional positions. Remove the + # M-RoPE fields inherited from the multimodal configuration. + hf_text_config = aphrodite_config.model_config.hf_text_config + rope_parameters = getattr(hf_text_config, "rope_parameters", None) + if rope_parameters is not None: + rope_parameters.pop("mrope_section", None) + rope_parameters.pop("mrope_interleaved", None) + + class ColQwen3_5Config(Qwen3_5ForConditionalGenerationConfig): """Apply the attention contract declared by a ColQwen3.5 checkpoint.""" @@ -825,7 +839,9 @@ def verify_and_update_config(aphrodite_config: "AphroditeConfig") -> None: "Qwen2ForRewardModel": Qwen2ForRewardModelConfig, "Qwen3ForSequenceClassification": Qwen3ForSequenceClassificationConfig, "Qwen3VLForSequenceClassification": Qwen3VLForSequenceClassificationConfig, + "Qwen3_5ForCausalLM": Qwen3_5ForCausalLMConfig, "Qwen3_5ForConditionalGeneration": Qwen3_5ForConditionalGenerationConfig, + "Qwen3_5MoeForCausalLM": Qwen3_5ForCausalLMConfig, "Qwen3_5MoeForConditionalGeneration": Qwen3_5ForConditionalGenerationConfig, "UnlimitedOCRForCausalLM": UnlimitedOCRForCausalLMConfig, "VoyageQwen3BidirectionalEmbedModel": VoyageQwen3BidirectionalEmbedModelConfig, diff --git a/aphrodite/model_executor/models/glm4_1v.py b/aphrodite/model_executor/models/glm4_1v.py index f982493f7e..7679ebeac8 100644 --- a/aphrodite/model_executor/models/glm4_1v.py +++ b/aphrodite/model_executor/models/glm4_1v.py @@ -1031,8 +1031,8 @@ def _get_vision_info( preprocessed_size = ImageSize(width=image_width, height=image_height) # NOTE: Frames are padded to be divisible by `temporal_patch_size` - # https://github.com/huggingface/transformers/blob/v4.48.3/src/transformers/models/qwen2_vl/image_processing_qwen2_vl.py#L294 - padded_num_frames = num_frames + num_frames % temporal_patch_size + # https://github.com/huggingface/transformers/blob/v5.13.0/src/transformers/models/qwen2_vl/video_processing_qwen2_vl.py#L249-L252 + padded_num_frames = num_frames + (-num_frames % temporal_patch_size) grid_t = max(padded_num_frames // temporal_patch_size, 1) grid_h = preprocessed_size.height // patch_size diff --git a/aphrodite/model_executor/models/kanana_v.py b/aphrodite/model_executor/models/kanana_v.py index 72377eff05..ed5f7ba086 100644 --- a/aphrodite/model_executor/models/kanana_v.py +++ b/aphrodite/model_executor/models/kanana_v.py @@ -403,8 +403,8 @@ def _get_vision_info( preprocessed_size = ImageSize(width=image_width, height=image_height) # NOTE: Frames are padded to be divisible by `temporal_patch_size` - # https://github.com/huggingface/transformers/blob/v4.48.3/src/transformers/models/qwen2_vl/image_processing_qwen2_vl.py#L294 - padded_num_frames = num_frames + num_frames % temporal_patch_size + # https://github.com/huggingface/transformers/blob/v5.13.0/src/transformers/models/qwen2_vl/video_processing_qwen2_vl.py#L249-L252 + padded_num_frames = num_frames + (-num_frames % temporal_patch_size) grid_t = max(padded_num_frames // temporal_patch_size, 1) grid_h = preprocessed_size.height // patch_size diff --git a/aphrodite/model_executor/models/keye.py b/aphrodite/model_executor/models/keye.py index bb82a8cde5..c2ef7f45c7 100644 --- a/aphrodite/model_executor/models/keye.py +++ b/aphrodite/model_executor/models/keye.py @@ -954,7 +954,7 @@ def _get_vision_info( else: preprocessed_size = ImageSize(width=image_width, height=image_height) - padded_num_frames = num_frames + num_frames % temporal_patch_size + padded_num_frames = num_frames + (-num_frames % temporal_patch_size) grid_t = max(padded_num_frames // temporal_patch_size, 1) grid_h = preprocessed_size.height // patch_size diff --git a/aphrodite/model_executor/models/llava_onevision2.py b/aphrodite/model_executor/models/llava_onevision2.py index de275fa19d..3bb868c0e4 100644 --- a/aphrodite/model_executor/models/llava_onevision2.py +++ b/aphrodite/model_executor/models/llava_onevision2.py @@ -1265,7 +1265,7 @@ def _get_vision_info( preprocessed = ImageSize(width=rw, height=rh) else: preprocessed = ImageSize(width=image_width, height=image_height) - padded_frames = num_frames + num_frames % temporal_patch_size + padded_frames = num_frames + (-num_frames % temporal_patch_size) grid_t = max(padded_frames // temporal_patch_size, 1) grid_h = preprocessed.height // patch_size grid_w = preprocessed.width // patch_size diff --git a/aphrodite/model_executor/models/mimo_v2_omni.py b/aphrodite/model_executor/models/mimo_v2_omni.py index bb56ad794f..e7c2312a73 100644 --- a/aphrodite/model_executor/models/mimo_v2_omni.py +++ b/aphrodite/model_executor/models/mimo_v2_omni.py @@ -682,7 +682,7 @@ def _get_vision_info( effective_frames = num_frames * tokens_per_second else: effective_frames = num_frames - padded_num_frames = effective_frames + effective_frames % temporal_patch_size + padded_num_frames = effective_frames + (-effective_frames % temporal_patch_size) grid_t = max(padded_num_frames // temporal_patch_size, 1) grid_h = preprocessed_size.height // patch_size grid_w = preprocessed_size.width // patch_size diff --git a/aphrodite/model_executor/models/qwen2_vl.py b/aphrodite/model_executor/models/qwen2_vl.py index 9c1a98748d..5274c715aa 100644 --- a/aphrodite/model_executor/models/qwen2_vl.py +++ b/aphrodite/model_executor/models/qwen2_vl.py @@ -870,8 +870,8 @@ def _get_vision_info( preprocessed_size = ImageSize(width=image_width, height=image_height) # NOTE: Frames are padded to be divisible by `temporal_patch_size` - # https://github.com/huggingface/transformers/blob/v4.48.3/src/transformers/models/qwen2_vl/image_processing_qwen2_vl.py#L294 - padded_num_frames = num_frames + num_frames % temporal_patch_size + # https://github.com/huggingface/transformers/blob/v5.13.0/src/transformers/models/qwen2_vl/video_processing_qwen2_vl.py#L249-L252 + padded_num_frames = num_frames + (-num_frames % temporal_patch_size) grid_t = max(padded_num_frames // temporal_patch_size, 1) grid_h = preprocessed_size.height // patch_size diff --git a/aphrodite/model_executor/models/qwen3_5.py b/aphrodite/model_executor/models/qwen3_5.py index 7ca266a8ca..c577620260 100644 --- a/aphrodite/model_executor/models/qwen3_5.py +++ b/aphrodite/model_executor/models/qwen3_5.py @@ -275,6 +275,7 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: class Qwen3_5ForCausalLMBase( nn.Module, HasInnerState, + IsHybrid, SupportsEagle3, SupportsLoRA, SupportsPP, @@ -348,6 +349,43 @@ def forward( return hidden_states + @classmethod + def get_mamba_state_dtype_from_config( + cls, + aphrodite_config: AphroditeConfig, + ) -> tuple[torch.dtype, torch.dtype]: + return MambaStateDtypeCalculator.gated_delta_net_state_dtype( + aphrodite_config.model_config.dtype, + aphrodite_config.cache_config.mamba_cache_dtype, + aphrodite_config.cache_config.mamba_ssm_cache_dtype, + ) + + @classmethod + def get_mamba_state_shape_from_config( + cls, aphrodite_config: AphroditeConfig + ) -> tuple[tuple[int, int], tuple[int, int]]: + parallel_config = aphrodite_config.parallel_config + hf_config = aphrodite_config.model_config.hf_text_config + tp_size = parallel_config.tensor_parallel_size + num_spec = ( + aphrodite_config.speculative_config.num_speculative_tokens if aphrodite_config.speculative_config else 0 + ) + return MambaStateShapeCalculator.gated_delta_net_state_shape( + tp_size, + hf_config.linear_num_key_heads, + hf_config.linear_num_value_heads, + hf_config.linear_key_head_dim, + hf_config.linear_value_head_dim, + hf_config.linear_conv_kernel_dim, + num_spec, + ) + + @classmethod + def get_mamba_state_copy_func( + cls, + ) -> tuple[MambaStateCopyFunc, MambaStateCopyFunc]: + return MambaStateCopyFuncCalculator.gated_delta_net_state_copy_func() + def compute_logits( self, hidden_states: torch.Tensor, diff --git a/aphrodite/model_executor/models/registry.py b/aphrodite/model_executor/models/registry.py index 60185f29fa..d85d165177 100644 --- a/aphrodite/model_executor/models/registry.py +++ b/aphrodite/model_executor/models/registry.py @@ -196,6 +196,8 @@ "Qwen2MoeForCausalLM": ("qwen2_moe", "Qwen2MoeForCausalLM"), "Qwen3ForCausalLM": ("qwen3", "Qwen3ForCausalLM"), "Qwen3MoeForCausalLM": ("qwen3_moe", "Qwen3MoeForCausalLM"), + "Qwen3_5ForCausalLM": ("qwen3_5", "Qwen3_5ForCausalLM"), + "Qwen3_5MoeForCausalLM": ("qwen3_5", "Qwen3_5MoeForCausalLM"), "RWForCausalLM": ("falcon", "FalconForCausalLM"), "SarvamMoEForCausalLM": ("sarvam", "SarvamMoEForCausalLM"), "SarvamMLAForCausalLM": ("sarvam", "SarvamMLAForCausalLM"), diff --git a/aphrodite/models/common/__init__.py b/aphrodite/models/common/__init__.py new file mode 100644 index 0000000000..208f01a7cb --- /dev/null +++ b/aphrodite/models/common/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project diff --git a/aphrodite/models/common/ops/__init__.py b/aphrodite/models/common/ops/__init__.py new file mode 100644 index 0000000000..0d0fe8de26 --- /dev/null +++ b/aphrodite/models/common/ops/__init__.py @@ -0,0 +1,9 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Ops shared across model implementations.""" + +from .fused_qk_rmsnorm import fused_q_kv_rmsnorm + +__all__ = [ + "fused_q_kv_rmsnorm", +] diff --git a/aphrodite/models/common/ops/fused_qk_rmsnorm.py b/aphrodite/models/common/ops/fused_qk_rmsnorm.py new file mode 100644 index 0000000000..47d8e10b49 --- /dev/null +++ b/aphrodite/models/common/ops/fused_qk_rmsnorm.py @@ -0,0 +1,101 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import torch + +from aphrodite.platforms import current_platform +from aphrodite.triton_utils import tl, triton + + +@triton.jit +def _fused_q_kv_rmsnorm_kernel( + q_ptr, + q_out_ptr, + q_weight_ptr, + q_in_stride, + q_out_stride, + kv_ptr, + kv_out_ptr, + kv_weight_ptr, + kv_in_stride, + kv_out_stride, + eps, + Q_SIZE: tl.constexpr, + KV_SIZE: tl.constexpr, + BLOCK_SIZE: tl.constexpr, + launch_pdl: tl.constexpr, +): + # num_tokens goes on grid-x (max 2**31 - 1); task goes on grid-y. + # CUDA's grid-y/z are capped at 65535, so putting num_tokens there crashes + # the launch at max-num-batched-tokens >= 65536 with "invalid argument". + # int64: q_in_stride can be ~24K (128 heads × 192) and overflows int32 + # past num_tokens ~87K under large chunked prefill. + token_idx = tl.program_id(0).to(tl.int64) + pid_task = tl.program_id(1) + + if pid_task == 0: + SIZE = Q_SIZE + row_in = q_ptr + token_idx * q_in_stride + weight_ptr = q_weight_ptr + row_out = q_out_ptr + token_idx * q_out_stride + else: + SIZE = KV_SIZE + row_in = kv_ptr + token_idx * kv_in_stride + weight_ptr = kv_weight_ptr + row_out = kv_out_ptr + token_idx * kv_out_stride + + if launch_pdl: + tl.extra.cuda.gdc_wait() + tl.extra.cuda.gdc_launch_dependents() + + # RMSNorm in fp32 throughout — matches csrc/layernorm_kernels.cu's + # `(scalar_t)(x * s_variance * w)` and DeepseekV4's compressor kernel, which + # keep x, rrms, and w all in fp32 and perform a single cast at store. + block = tl.arange(0, BLOCK_SIZE) + mask = block < SIZE + x = tl.load(row_in + block, mask=mask, other=0.0).to(tl.float32) + variance = tl.sum(x * x, axis=0) / SIZE + rrms = tl.rsqrt(variance + eps) + w = tl.load(weight_ptr + block, mask=mask, other=0.0).to(tl.float32) + y = x * rrms * w + tl.store(row_out + block, y.to(row_out.dtype.element_ty), mask=mask) + + +def fused_q_kv_rmsnorm( + qr: torch.Tensor, + kv: torch.Tensor, + q_weight: torch.Tensor, + kv_weight: torch.Tensor, + eps: float, +) -> tuple[torch.Tensor, torch.Tensor]: + assert qr.ndim == 2 and kv.ndim == 2 + assert qr.shape[0] == kv.shape[0], f"token dim mismatch: qr={qr.shape}, kv={kv.shape}" + assert qr.stride(-1) == 1 and kv.stride(-1) == 1 + assert q_weight.is_contiguous() and kv_weight.is_contiguous() + + q_size = qr.shape[1] + kv_size = kv.shape[1] + num_tokens = qr.shape[0] + qr_out = torch.empty_like(qr) + kv_out = torch.empty_like(kv) + if num_tokens == 0: + return qr_out, kv_out + + block_size = triton.next_power_of_2(max(q_size, kv_size)) + _fused_q_kv_rmsnorm_kernel[(num_tokens, 2)]( + qr, + qr_out, + q_weight, + qr.stride(0), + qr_out.stride(0), + kv, + kv_out, + kv_weight, + kv.stride(0), + kv_out.stride(0), + eps, + Q_SIZE=q_size, + KV_SIZE=kv_size, + BLOCK_SIZE=block_size, + launch_pdl=current_platform.is_arch_support_pdl(), + ) + return qr_out, kv_out diff --git a/aphrodite/models/inkling/nvidia/moe.py b/aphrodite/models/inkling/nvidia/moe.py index 41a1f26e62..5f83f58956 100644 --- a/aphrodite/models/inkling/nvidia/moe.py +++ b/aphrodite/models/inkling/nvidia/moe.py @@ -559,6 +559,9 @@ def load_expert_weight(self, name: str, weight: torch.Tensor) -> list[str]: elif vals.shape[1] != target_width: raise ValueError(f"cannot load {tuple(weight.shape)} into {tuple(param.shape)}") param.data[lids] = vals.reshape(len(gids), *param.shape[1:]).to(param.device) + elif key == "w2_weight_scale" and weight.shape[-1] == 1: + # Per-output-channel scales are replicated across TP ranks. + param.data[lids] = weight[gids].to(device=param.device, dtype=param.dtype) elif key.startswith("w13"): # Checkpoint w13 rows are interleaved [g0, u0, g1, u1, ...]; the # fused param layout is [w1(gate); w3(up)]. The TP-local rows form diff --git a/aphrodite/models/inkling/nvidia/ops/fa4_rel_attention.py b/aphrodite/models/inkling/nvidia/ops/fa4_rel_attention.py index a3fc2c0b7d..0801d16080 100644 --- a/aphrodite/models/inkling/nvidia/ops/fa4_rel_attention.py +++ b/aphrodite/models/inkling/nvidia/ops/fa4_rel_attention.py @@ -128,13 +128,20 @@ def inkling_fa4_rel_attention( cute_window = (None, None) if window_size == (-1, -1) else window_size rel_logits = rel_logits.contiguous() + flash_attn_varlen_func: Callable[..., Any] if _use_sheared_bias(): - from aphrodite.third_party.tml_fa4 import flash_attn_varlen_func + from aphrodite.third_party.tml_fa4 import ( + flash_attn_varlen_func as tml_flash_attn_varlen_func, + ) + flash_attn_varlen_func = tml_flash_attn_varlen_func bias_kwargs: dict[str, Any] = {"rel_bias": rel_logits} else: - from aphrodite.vllm_flash_attn.cute import flash_attn_varlen_func + from aphrodite.vllm_flash_attn.cute import ( + flash_attn_varlen_func as cute_flash_attn_varlen_func, + ) + flash_attn_varlen_func = cute_flash_attn_varlen_func bias_kwargs = { "score_mod": _get_score_mod(rel_extent), "aux_tensors": [rel_logits], diff --git a/aphrodite/models/kimi_k3/__init__.py b/aphrodite/models/kimi_k3/__init__.py index 208f01a7cb..2b382f03ae 100644 --- a/aphrodite/models/kimi_k3/__init__.py +++ b/aphrodite/models/kimi_k3/__init__.py @@ -1,2 +1,28 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Kimi K3 model — hardware-isolated entry point. + +The implementation lives under ``nvidia/`` and ``amd/``; this module picks the +right one for the current platform and re-exports the public classes used by +the model registry. (Mirrors ``aphrodite.models.minimax_m3``.) +""" + +from typing import TYPE_CHECKING + +from aphrodite.platforms import current_platform + +# The NVIDIA branch is the static default that type-checkers see; the ROCm +# branch overrides it at runtime (kept type-compatible via type: ignore). +if TYPE_CHECKING or not current_platform.is_rocm(): + from .nvidia.model import KimiK3ForConditionalGeneration, KimiLinearForCausalLM + from .nvidia.mtp import KimiK3MTP +else: + from .amd.linear import KimiLinearForCausalLM # type: ignore[assignment] + from .amd.model import KimiK3ForConditionalGeneration # type: ignore[assignment] + from .amd.mtp import KimiK3MTP # type: ignore[assignment] + +__all__ = [ + "KimiK3ForConditionalGeneration", + "KimiK3MTP", + "KimiLinearForCausalLM", +] diff --git a/aphrodite/models/kimi_k3/amd/linear.py b/aphrodite/models/kimi_k3/amd/linear.py new file mode 100644 index 0000000000..7a99b8873e --- /dev/null +++ b/aphrodite/models/kimi_k3/amd/linear.py @@ -0,0 +1,995 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from collections.abc import Iterable +from typing import Any + +import torch +from torch import nn + +from aphrodite.config import AphroditeConfig, CacheConfig +from aphrodite.distributed import ( + get_pp_group, + get_tensor_model_parallel_world_size, +) +from aphrodite.logger import init_logger +from aphrodite.model_executor.layers.activation import SiluAndMul, SituAndMul +from aphrodite.model_executor.layers.fused_moe import ( + FusedMoE, + fused_moe_make_expert_params_mapping, +) +from aphrodite.model_executor.layers.fused_moe.router.gate_linear import GateLinear +from aphrodite.model_executor.layers.layernorm import RMSNorm +from aphrodite.model_executor.layers.linear import ( + ColumnParallelLinear, + MergedColumnParallelLinear, + ReplicatedLinear, + RowParallelLinear, +) +from aphrodite.model_executor.layers.logits_processor import LogitsProcessor +from aphrodite.model_executor.layers.mamba.gdn.kimi_gdn_linear_attn import ( + KimiGatedDeltaNetAttention, +) +from aphrodite.model_executor.layers.mamba.mamba_utils import ( + MambaStateCopyFunc, + MambaStateCopyFuncCalculator, + MambaStateDtypeCalculator, + MambaStateShapeCalculator, +) +from aphrodite.model_executor.layers.mla import MLAModules, MultiHeadLatentAttentionWrapper +from aphrodite.model_executor.layers.quantization.base_config import QuantizationConfig +from aphrodite.model_executor.layers.vocab_parallel_embedding import ( + ParallelLMHead, + VocabParallelEmbedding, +) +from aphrodite.model_executor.model_loader.weight_utils import ( + default_weight_loader, + maybe_remap_kv_scale_name, +) +from aphrodite.model_executor.models.interfaces import ( + EagleModelMixin, + HasInnerState, + IsHybrid, + MixtureOfExperts, + SupportsPP, +) +from aphrodite.model_executor.models.utils import ( + AutoWeightsLoader, + PPMissingLayer, + get_spec_layer_idx_from_weight_name, + is_pp_missing_parameter, + make_layers, + maybe_prefix, +) +from aphrodite.models.kimi_k3.amd.ops.attn_res import attn_res +from aphrodite.sequence import IntermediateTensors +from aphrodite.transformers_utils.configs.kimi_linear import KimiLinearConfig +from aphrodite.utils.math_utils import cdiv + +logger = init_logger(__name__) + + +class KimiMLP(nn.Module): + def __init__( + self, + hidden_size: int, + intermediate_size: int, + hidden_act: str, + quant_config: QuantizationConfig | None = None, + reduce_results: bool = True, + prefix: str = "", + activation_situ_beta: float | None = None, + activation_situ_linear_beta: float | None = None, + ) -> None: + super().__init__() + + self.gate_up_proj = MergedColumnParallelLinear( + hidden_size, + [intermediate_size] * 2, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.gate_up_proj", + ) + self.down_proj = RowParallelLinear( + intermediate_size, + hidden_size, + bias=False, + quant_config=quant_config, + reduce_results=reduce_results, + prefix=f"{prefix}.down_proj", + ) + if hidden_act == "silu": + self.act_fn = SiluAndMul() + elif hidden_act == "situ": + self.act_fn = SituAndMul( + beta=activation_situ_beta or 1.0, + linear_beta=activation_situ_linear_beta, + ) + else: + raise ValueError(f"Unsupported activation: {hidden_act}. Only silu and situ are supported.") + + def forward(self, x): + gate_up, _ = self.gate_up_proj(x) + x = self.act_fn(gate_up) + x, _ = self.down_proj(x) + return x + + +class KimiRoutedOutputTransform(nn.Module): + def __init__( + self, + norm: RMSNorm | None, + up_proj: ReplicatedLinear, + ) -> None: + super().__init__() + self.norm = norm + self.up_proj = up_proj + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + if self.norm is not None: + hidden_states = self.norm(hidden_states) + hidden_states, _ = self.up_proj(hidden_states) + return hidden_states + + +def _apply_attn_res( + prefix_sum: torch.Tensor, + block_residual: torch.Tensor, + proj: ReplicatedLinear, + norm: RMSNorm, + num_valid_blocks: int, +) -> torch.Tensor: + if num_valid_blocks <= 0: + return prefix_sum + + return attn_res( + prefix_sum, + block_residual, + norm.weight, + proj.weight.squeeze(0), + num_valid_blocks, + norm.variance_epsilon, + ) + + +class KimiMoE(nn.Module): + def __init__( + self, + config: KimiLinearConfig, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + layer_idx: int = 0, + ): + super().__init__() + hidden_size = config.hidden_size + moe_intermediate_size = config.moe_intermediate_size + num_experts = config.num_experts + num_experts_per_token = config.num_experts_per_token + assert moe_intermediate_size is not None + assert num_experts is not None + assert num_experts_per_token is not None + moe_renormalize = config.moe_renormalize + routed_expert_hidden_size = config.routed_expert_hidden_size + self.use_latent_moe = routed_expert_hidden_size is not None + self.moe_hidden_size = routed_expert_hidden_size if routed_expert_hidden_size is not None else hidden_size + self.latent_moe_use_norm = config.latent_moe_use_norm + self.tp_size = get_tensor_model_parallel_world_size() + self.routed_scaling_factor = config.routed_scaling_factor + self.num_shared_experts = config.num_shared_experts + self.layer_idx = layer_idx + self.padded_moe_intermediate_size = moe_intermediate_size + min_moe_intermediate_per_partition = getattr(config, "min_moe_intermediate_per_partition", 256) + if self.tp_size > 1: + moe_intermediate_per_partition = moe_intermediate_size // self.tp_size + if moe_intermediate_per_partition < min_moe_intermediate_per_partition: + self.padded_moe_intermediate_size = min_moe_intermediate_per_partition * self.tp_size + activation_situ_beta = config.activation_situ_beta if config.hidden_act == "situ" else None + activation_situ_linear_beta = config.activation_situ_linear_beta if config.hidden_act == "situ" else None + + # Route with fp32 logits for numerically stable expert selection. + self.gate = GateLinear( + input_size=hidden_size, + output_size=num_experts, + bias=False, + out_dtype=torch.float32, + prefix=f"{prefix}.gate", + ) + + self.gate.e_score_correction_bias = nn.Parameter(torch.empty(num_experts)) + + self.shared_experts: KimiMLP | None + if self.num_shared_experts is not None: + shared_intermediate_size = moe_intermediate_size * self.num_shared_experts + self.shared_experts = KimiMLP( + hidden_size=config.hidden_size, + intermediate_size=shared_intermediate_size, + hidden_act=config.hidden_act, + quant_config=quant_config, + reduce_results=False, + prefix=f"{prefix}.shared_experts", + activation_situ_beta=activation_situ_beta, + activation_situ_linear_beta=activation_situ_linear_beta, + ) + else: + self.shared_experts = None + + self.routed_expert_down_proj: ReplicatedLinear | None + self.routed_expert_norm: RMSNorm | None + self.routed_expert_up_proj: ReplicatedLinear | None + self.routed_output_transform: KimiRoutedOutputTransform | None + if self.use_latent_moe: + self.routed_expert_down_proj = ReplicatedLinear( + hidden_size, + self.moe_hidden_size, + bias=False, + quant_config=None, + prefix=f"{prefix}.routed_expert_down_proj", + ) + self.routed_expert_norm = ( + RMSNorm(self.moe_hidden_size, eps=config.rms_norm_eps) if self.latent_moe_use_norm else None + ) + self.routed_expert_up_proj = ReplicatedLinear( + self.moe_hidden_size, + hidden_size, + bias=False, + quant_config=None, + prefix=f"{prefix}.routed_expert_up_proj", + ) + self.routed_output_transform = KimiRoutedOutputTransform( + self.routed_expert_norm, self.routed_expert_up_proj + ) + else: + self.routed_expert_down_proj = None + self.routed_expert_norm = None + self.routed_expert_up_proj = None + self.routed_output_transform = None + + self.experts = FusedMoE( + shared_experts=self.shared_experts, + num_experts=num_experts, + top_k=num_experts_per_token, + hidden_size=self.moe_hidden_size, + intermediate_size=self.padded_moe_intermediate_size, + activation=config.hidden_act, + activation_situ_beta=activation_situ_beta, + activation_situ_linear_beta=activation_situ_linear_beta, + renormalize=moe_renormalize, + quant_config=quant_config, + use_grouped_topk=config.use_grouped_topk, + num_expert_group=config.num_expert_group, + topk_group=config.topk_group, + prefix=f"{prefix}.experts", + scoring_func=config.moe_router_activation_func, + e_score_correction_bias=self.gate.e_score_correction_bias, + routed_scaling_factor=self.routed_scaling_factor, + routed_input_transform=self.routed_expert_down_proj, + routed_output_transform=self.routed_output_transform, + ) + if self.padded_moe_intermediate_size != moe_intermediate_size: + w13_weight = getattr(self.experts, "w13_weight", None) + if w13_weight is None: + w13_weight = self.experts.w13_weight_packed + w2_weight = getattr(self.experts, "w2_weight", None) + if w2_weight is None: + w2_weight = self.experts.w2_weight_packed + w13_weight.data.zero_() + w2_weight.data.zero_() + self.experts.moe_config.intermediate_size_per_partition_unpadded = moe_intermediate_size // self.tp_size + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + num_tokens, hidden_size = hidden_states.shape + hidden_states = hidden_states.view(-1, hidden_size) + router_logits, _ = self.gate(hidden_states) + final_hidden_states = self.experts(hidden_states=hidden_states, router_logits=router_logits) + return final_hidden_states.view(num_tokens, hidden_size) + + +class KimiMLAAttention(nn.Module): + """ + Main reference: DeepseekV2 vllm Implementation + """ + + def __init__( + self, + config: KimiLinearConfig, + hidden_size: int, + num_heads: int, + qk_nope_head_dim: int, + qk_rope_head_dim: int, + v_head_dim: int, + q_lora_rank: int | None, + kv_lora_rank: int, + use_nope: bool = False, + cache_config: CacheConfig | None = None, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + **kwargs, + ) -> None: + super().__init__() + self.hidden_size = hidden_size + self.qk_nope_head_dim = qk_nope_head_dim + self.qk_rope_head_dim = qk_rope_head_dim + self.qk_head_dim = qk_nope_head_dim + qk_rope_head_dim + self.v_head_dim = v_head_dim + self.q_lora_rank = q_lora_rank + self.kv_lora_rank = kv_lora_rank + self.num_heads = num_heads + tp_size = get_tensor_model_parallel_world_size() + self.num_local_heads = num_heads // tp_size + self.scaling = self.qk_head_dim**-0.5 + self.use_nope = use_nope + assert self.use_nope is True + assert num_heads % tp_size == 0 + if self.q_lora_rank is not None: + self.fused_qkv_a_proj = MergedColumnParallelLinear( + self.hidden_size, + [self.q_lora_rank, self.kv_lora_rank + self.qk_rope_head_dim], + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.fused_qkv_a_proj", + disable_tp=True, + ) + else: + self.kv_a_proj_with_mqa = ReplicatedLinear( + self.hidden_size, + self.kv_lora_rank + self.qk_rope_head_dim, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.kv_a_proj_with_mqa", + ) + if self.q_lora_rank is not None: + self.q_a_layernorm = RMSNorm( + self.q_lora_rank, + eps=config.rms_norm_eps, + ) + self.q_b_proj = ColumnParallelLinear( + self.q_lora_rank, + self.num_heads * self.qk_head_dim, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.q_b_proj", + ) + else: + self.q_proj = ColumnParallelLinear( + self.hidden_size, + self.num_heads * self.qk_head_dim, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.q_proj", + ) + self.kv_a_layernorm = RMSNorm( + self.kv_lora_rank, + eps=config.rms_norm_eps, + ) + self.kv_b_proj = ColumnParallelLinear( + self.kv_lora_rank, + self.num_heads * (self.qk_nope_head_dim + self.v_head_dim), + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.kv_b_proj", + ) + self.o_proj = RowParallelLinear( + self.num_heads * self.v_head_dim, + self.hidden_size, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.o_proj", + ) + + self.use_output_gate = config.mla_use_output_gate + if self.use_output_gate: + projection_size = self.num_heads * self.v_head_dim + self.g_proj = ColumnParallelLinear( + self.hidden_size, + projection_size, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.g_proj", + ) + + # TODO: Remove this mypy workaround once the K3 PR is fully merged. + mla_modules = MLAModules( # type: ignore[call-arg] + kv_a_layernorm=self.kv_a_layernorm, + kv_b_proj=self.kv_b_proj, + rotary_emb=None, + o_proj=self.o_proj, + fused_qkv_a_proj=self.fused_qkv_a_proj if self.q_lora_rank is not None else None, + kv_a_proj_with_mqa=self.kv_a_proj_with_mqa if self.q_lora_rank is None else None, + q_a_layernorm=self.q_a_layernorm if self.q_lora_rank is not None else None, + q_b_proj=self.q_b_proj if self.q_lora_rank is not None else None, + q_proj=self.q_proj if self.q_lora_rank is None else None, + indexer=None, + is_sparse=False, + topk_indices_buffer=None, + g_proj=getattr(self, "g_proj", None), + ) + self.mla_attn = MultiHeadLatentAttentionWrapper( + self.hidden_size, + self.num_local_heads, + self.scaling, + self.qk_nope_head_dim, + self.qk_rope_head_dim, + self.v_head_dim, + self.q_lora_rank, + self.kv_lora_rank, + mla_modules, + cache_config, + quant_config, + prefix, + ) + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + output: torch.Tensor, + ) -> None: + output[:] = self.mla_attn(positions, hidden_states) + + +class KimiDecoderLayer(nn.Module): + def __init__( + self, + config: KimiLinearConfig, + aphrodite_config: AphroditeConfig, + prefix: str = "", + ) -> None: + super().__init__() + self.hidden_size = config.hidden_size + self.layer_idx = int(prefix.rsplit(".", 1)[1]) + + self.is_moe = config.is_moe + layer_idx = self.layer_idx + model_config = aphrodite_config.model_config + cache_config = aphrodite_config.cache_config + quant_config = aphrodite_config.quant_config + + if config.is_kda_layer(layer_idx): + self.self_attn = KimiGatedDeltaNetAttention( + config, + aphrodite_config, + prefix=f"{prefix}.self_attn", + ) + else: + qk_nope_head_dim = config.qk_nope_head_dim + qk_rope_head_dim = config.qk_rope_head_dim + v_head_dim = config.v_head_dim + kv_lora_rank = config.kv_lora_rank + mla_use_nope = config.mla_use_nope + assert qk_nope_head_dim is not None + assert qk_rope_head_dim is not None + assert v_head_dim is not None + assert kv_lora_rank is not None + assert mla_use_nope is not None + self.self_attn = KimiMLAAttention( + layer_idx=layer_idx, + hidden_size=self.hidden_size, + num_heads=config.num_attention_heads, + quant_config=quant_config, + cache_config=cache_config, + model_config=model_config, + prefix=f"{prefix}.self_attn", + config=config, + qk_nope_head_dim=qk_nope_head_dim, + qk_rope_head_dim=qk_rope_head_dim, + v_head_dim=v_head_dim, + q_lora_rank=config.q_lora_rank, + kv_lora_rank=kv_lora_rank, + use_nope=mla_use_nope, + ) + + if ( + self.is_moe + and config.num_experts is not None + and layer_idx >= config.first_k_dense_replace + and layer_idx % config.moe_layer_freq == 0 + ): + self.block_sparse_moe = KimiMoE( + config=config, + quant_config=quant_config, + prefix=f"{prefix}.block_sparse_moe", + layer_idx=layer_idx, + ) + self.mlp = self.block_sparse_moe + else: + self.mlp = KimiMLP( + hidden_size=self.hidden_size, + intermediate_size=config.intermediate_size, + hidden_act=config.hidden_act, + quant_config=quant_config, + prefix=f"{prefix}.mlp", + activation_situ_beta=config.activation_situ_beta, + activation_situ_linear_beta=config.activation_situ_linear_beta, + ) + self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.post_attention_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + + attn_res_block_size = config.attn_res_block_size + self.use_attn_residuals = attn_res_block_size is not None + if attn_res_block_size is not None: + self.attn_res_block_size = attn_res_block_size + self.is_block_write_layer = layer_idx % self.attn_res_block_size == 0 + self.block_write_idx = layer_idx // self.attn_res_block_size + self.prev_valid_blocks = cdiv(layer_idx, self.attn_res_block_size) + self.self_attention_res_norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.mlp_res_norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.self_attention_res_proj = ReplicatedLinear( + config.hidden_size, + 1, + bias=False, + quant_config=None, + prefix=f"{prefix}.self_attention_res_proj", + ) + self.mlp_res_proj = ReplicatedLinear( + config.hidden_size, + 1, + bias=False, + quant_config=None, + prefix=f"{prefix}.mlp_res_proj", + ) + + def _run_self_attn( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + ) -> torch.Tensor: + attn_output = torch.empty_like(hidden_states) + self.self_attn( + hidden_states=hidden_states, + positions=positions, + output=attn_output, + ) + return attn_output + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + residual: torch.Tensor | None, + **kwargs, + ) -> tuple[torch.Tensor, torch.Tensor]: + if self.use_attn_residuals: + assert residual is not None + return self.forward_attn_residual(positions, hidden_states, residual) + + # Self Attention + if residual is None: + residual = hidden_states + hidden_states = self.input_layernorm(hidden_states) + else: + hidden_states, residual = self.input_layernorm(hidden_states, residual) + + hidden_states = self._run_self_attn(positions, hidden_states) + + # Fully Connected + hidden_states, residual = self.post_attention_layernorm(hidden_states, residual) + hidden_states = self.mlp(hidden_states) + return hidden_states, residual + + def forward_attn_residual( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + block_residual: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + prefix_sum = hidden_states + hidden_states = _apply_attn_res( + prefix_sum, + block_residual, + self.self_attention_res_proj, + self.self_attention_res_norm, + self.prev_valid_blocks, + ) + + if self.is_block_write_layer: + block_residual[:, self.block_write_idx, :].copy_(prefix_sum) + prefix_sum = None + + hidden_states = self.input_layernorm(hidden_states) + hidden_states = self._run_self_attn(positions, hidden_states) + + if prefix_sum is not None: + prefix_sum = prefix_sum + hidden_states + else: + prefix_sum = hidden_states + + mlp_valid_blocks = self.prev_valid_blocks + (1 if self.is_block_write_layer else 0) + hidden_states = _apply_attn_res( + prefix_sum, + block_residual, + self.mlp_res_proj, + self.mlp_res_norm, + mlp_valid_blocks, + ) + + hidden_states = self.post_attention_layernorm(hidden_states) + hidden_states = self.mlp(hidden_states) + prefix_sum = prefix_sum + hidden_states + return prefix_sum, block_residual + + +class KimiLinearModel(nn.Module, EagleModelMixin): + def __init__(self, *, aphrodite_config: AphroditeConfig, prefix: str = ""): + super().__init__() + + config = aphrodite_config.model_config.hf_text_config + self.config = config + + self.vocab_size = config.vocab_size + + if get_pp_group().is_first_rank: + self.embed_tokens = VocabParallelEmbedding( + config.vocab_size, + config.hidden_size, + prefix=f"{prefix}.embed_tokens", + ) + else: + self.embed_tokens = PPMissingLayer() + + def get_layer(prefix: str): + return KimiDecoderLayer( + config, + aphrodite_config, + prefix, + ) + + self.start_layer, self.end_layer, self.layers = make_layers( + config.num_hidden_layers, + get_layer, + prefix=f"{prefix}.layers", + ) + + if get_pp_group().is_last_rank: + self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + if config.attn_res_block_size is not None: + self.output_attn_res_norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.output_attn_res_proj = ReplicatedLinear( + config.hidden_size, + 1, + bias=False, + quant_config=None, + prefix=f"{prefix}.output_attn_res_proj", + ) + else: + self.norm = PPMissingLayer() + if config.attn_res_block_size is not None: + self.output_attn_res_norm = PPMissingLayer() + self.output_attn_res_proj = PPMissingLayer() + + world_size = get_tensor_model_parallel_world_size() + assert config.num_attention_heads % world_size == 0, "num_attention_heads must be divisible by world_size" + + def make_empty_intermediate_tensors( + self, + batch_size: int, + dtype: torch.dtype, + device: torch.device, + ) -> IntermediateTensors: + residual_shape: tuple[int, ...] = (batch_size, self.config.hidden_size) + if self.config.attn_res_block_size is not None: + residual_shape = ( + batch_size, + cdiv(self.start_layer, self.config.attn_res_block_size), + self.config.hidden_size, + ) + return IntermediateTensors( + { + "hidden_states": torch.zeros((batch_size, self.config.hidden_size), dtype=dtype, device=device), + "residual": torch.zeros(residual_shape, dtype=dtype, device=device), + } + ) + + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: + return self.embed_tokens(input_ids) + + def _maybe_add_hidden_state( + self, + aux_hidden_states: list[torch.Tensor], + layer_idx: int, + hidden_states: torch.Tensor, + residual: torch.Tensor | None, + ) -> list[torch.Tensor]: + if self.config.attn_res_block_size is not None: + # attn-res `residual` is a block-state bank, not an additive + # residual; None makes the mixin capture the prefix sum directly. + residual = None + return super()._maybe_add_hidden_state(aux_hidden_states, layer_idx, hidden_states, residual) + + def forward( + self, + input_ids: torch.Tensor | None, + positions: torch.Tensor, + intermediate_tensors: IntermediateTensors | None, + inputs_embeds: torch.Tensor | None = None, + **kwargs, + ) -> torch.Tensor | IntermediateTensors | tuple[torch.Tensor, list[torch.Tensor]]: + if get_pp_group().is_first_rank: + if inputs_embeds is not None: + hidden_states = inputs_embeds + else: + hidden_states = self.embed_input_ids(input_ids) + residual = None + else: + assert intermediate_tensors is not None + hidden_states = intermediate_tensors["hidden_states"] + residual = intermediate_tensors["residual"] + + aux_hidden_states = self._maybe_add_hidden_state([], self.start_layer, hidden_states, residual) + + if self.config.attn_res_block_size is None: + for layer_idx, layer in enumerate( + self.layers[self.start_layer : self.end_layer], + start=self.start_layer, + ): + hidden_states, residual = layer( + positions=positions, + hidden_states=hidden_states, + residual=residual, + ) + 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}) + + # NOTE: the final norm is applied in compute_logits instead of here, + # so the MTP draft model receives the pre-norm hidden states. + if residual is not None: + hidden_states = hidden_states + residual + if aux_hidden_states: + return hidden_states, aux_hidden_states + return hidden_states + + attn_res_block_num = cdiv(self.end_layer, self.config.attn_res_block_size) + block_residual = hidden_states.new_empty(hidden_states.size(0), attn_res_block_num, hidden_states.size(1)) + if residual is not None: + block_residual[:, : residual.size(1), :].copy_(residual) + residual = block_residual + + for layer_idx, layer in enumerate( + self.layers[self.start_layer : self.end_layer], + start=self.start_layer, + ): + hidden_states, residual = layer( + positions=positions, + hidden_states=hidden_states, + residual=residual, + ) + if (layer_idx + 1) in self.aux_hidden_state_layers: + # AMD attn-res layer already returns prefix_sum + MLP delta as + # hidden_states; the override drops the block bank in residual. + 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}) + + hidden_states = _apply_attn_res( + hidden_states, + residual, + self.output_attn_res_proj, + self.output_attn_res_norm, + attn_res_block_num, + ) + # NOTE: the final norm is applied in compute_logits instead of here, so + # the MTP draft model receives the pre-norm hidden states. + if aux_hidden_states: + return hidden_states, aux_hidden_states + return hidden_states + + def load_weights( + self, + weights: Iterable[tuple[str, torch.Tensor] | tuple[str, torch.Tensor, dict[str, Any]]], + ) -> set[str]: + kda_config = self.config.linear_attn_config + use_full_rank_gate = bool(kda_config and kda_config.get("use_full_rank_gate", False)) + beta_shard_id = 5 if use_full_rank_gate else 3 + stacked_params_mapping = [ + # (param_name, shard_name, shard_id) + (".in_proj_qkvgfab", ".q_proj", 0), + (".in_proj_qkvgfab", ".k_proj", 1), + (".in_proj_qkvgfab", ".v_proj", 2), + (".in_proj_qkvgfab", ".b_proj", beta_shard_id), + (".in_proj_qkvgfab", ".f_a_proj", 4), + (".conv1d", ".q_conv1d", 0), + (".conv1d", ".k_conv1d", 1), + (".conv1d", ".v_conv1d", 2), + (".gate_up_proj", ".gate_proj", 0), + (".gate_up_proj", ".up_proj", 1), + ] + if use_full_rank_gate: + stacked_params_mapping.append((".in_proj_qkvgfab", ".g_proj", 3)) + if getattr(self.config, "q_lora_rank", None) is not None: + stacked_params_mapping += [ + (".fused_qkv_a_proj", ".q_a_proj", 0), + (".fused_qkv_a_proj", ".kv_a_proj_with_mqa", 1), + ] + if self.config.is_moe: + # Params for weights, fp8 weight scales, fp8 activation scales + # (param_name, weight_name, expert_id, shard_id) + expert_params_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.num_experts, + ) + else: + expert_params_mapping = [] + params_dict = dict(self.named_parameters()) + # Under the MXFP4 quant interface the routed experts register unpacked + # params (``w13_weight``), while the compressed-tensors checkpoint names + # them ``.weight_packed``. Rebind so the expert mapping resolves; scales + # already share the ``.weight_scale`` suffix. + experts_unpacked = not any(n.endswith("w13_weight_packed") for n in params_dict) + loaded_params: set[str] = set() + for args in weights: + name, loaded_weight = args[0], args[1] + kwargs: dict[str, Any] = args[2] if len(args) > 2 else {} + if "rotary_emb.inv_freq" in name: + continue + if experts_unpacked and name.endswith(".weight_packed"): + name = name.replace(".weight_packed", ".weight") + + spec_layer = get_spec_layer_idx_from_weight_name(self.config, name) + if spec_layer is not None: + continue # skip spec decode layers for main model + if "rotary_emb.cos_cached" in name or "rotary_emb.sin_cached" in name: + # Models trained using ColossalAI may include these tensors in + # the checkpoint. Skip them. + continue + for param_name, weight_name, shard_id in stacked_params_mapping: + if weight_name not in name: + continue + # We have mlp.experts[0].gate_proj in the checkpoint. + # Since we handle the experts below in expert_params_mapping, + # we need to skip here BEFORE we update the name, otherwise + # name will be updated to mlp.experts[0].gate_up_proj, which + # will then be updated below in expert_params_mapping + # for mlp.experts[0].gate_gate_up_proj, which breaks load. + if ("mlp.experts." in name) and name not in params_dict: + continue + name_mapped = name.replace(weight_name, param_name) + # Packed projections are only present on compatible layers. + if name_mapped not in params_dict: + continue + name = name_mapped + # Skip loading extra bias for GPTQ models. + if name.endswith(".bias") and name not in params_dict: + continue + if is_pp_missing_parameter(name, self): + continue + param = params_dict[name] + weight_loader = param.weight_loader + weight_loader(param, loaded_weight, shard_id) + break + else: + for ( + expert_param_name, + expert_weight_name, + expert_id, + expert_shard_id, + ) in expert_params_mapping: + if expert_weight_name not in name: + continue + name = name.replace(expert_weight_name, expert_param_name) + if is_pp_missing_parameter(name, self): + continue + param = params_dict[name] + weight_loader = param.weight_loader + weight_loader( + param, + loaded_weight, + name, + expert_id=expert_id, + shard_id=expert_shard_id, + ) + break + else: + # Skip loading extra bias for GPTQ models. + if name.endswith(".bias") and name not in params_dict and not self.config.is_linear_attn: # noqa: E501 + continue + # Remapping the name of FP8 kv-scale. + remapped_name = maybe_remap_kv_scale_name(name, params_dict) + if remapped_name is None: + continue + name = remapped_name + if is_pp_missing_parameter(name, self): + continue + + param = params_dict[name] + weight_loader = getattr(param, "weight_loader", default_weight_loader) + weight_loader(param, loaded_weight, **kwargs) + loaded_params.add(name) + return loaded_params + + +class KimiLinearForCausalLM(nn.Module, HasInnerState, SupportsPP, MixtureOfExperts, IsHybrid): + def __init__(self, *, aphrodite_config: AphroditeConfig, prefix: str = ""): + super().__init__() + self.model_config = aphrodite_config.model_config + self.aphrodite_config = aphrodite_config + self.config = self.model_config.hf_config + quant_config = aphrodite_config.quant_config + self.quant_config = quant_config + self.model = KimiLinearModel(aphrodite_config=aphrodite_config, prefix=maybe_prefix(prefix, "model")) + if get_pp_group().is_last_rank: + self.lm_head = ParallelLMHead( + self.config.vocab_size, + self.config.hidden_size, + quant_config=quant_config, + prefix=maybe_prefix(prefix, "lm_head"), + ) + else: + self.lm_head = PPMissingLayer() + logit_scale = getattr(self.config, "logit_scale", 1.0) + self.logits_processor = LogitsProcessor(self.config.vocab_size, scale=logit_scale) + + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: + return self.model.embed_input_ids(input_ids) + + def make_empty_intermediate_tensors( + self, + batch_size: int, + dtype: torch.dtype, + device: torch.device, + ) -> IntermediateTensors: + return self.model.make_empty_intermediate_tensors(batch_size, dtype, device) + + def forward( + self, + input_ids: torch.Tensor | None, + positions: torch.Tensor, + intermediate_tensors: IntermediateTensors | None = None, + inputs_embeds: torch.Tensor | None = None, + **kwargs, + ) -> torch.Tensor | IntermediateTensors: + hidden_states = self.model(input_ids, positions, intermediate_tensors, inputs_embeds, **kwargs) + return hidden_states + + @classmethod + def get_mamba_state_dtype_from_config( + cls, + aphrodite_config: "AphroditeConfig", + ) -> tuple[torch.dtype, torch.dtype]: + return MambaStateDtypeCalculator.kda_state_dtype( + aphrodite_config.model_config.dtype, aphrodite_config.cache_config.mamba_cache_dtype + ) + + @classmethod + def get_mamba_state_shape_from_config( + cls, aphrodite_config: "AphroditeConfig" + ) -> tuple[tuple[int, int], tuple[int, int, int]]: + parallel_config = aphrodite_config.parallel_config + hf_config = aphrodite_config.model_config.hf_config + tp_size = parallel_config.tensor_parallel_size + num_spec = ( + aphrodite_config.speculative_config.num_speculative_tokens if aphrodite_config.speculative_config else 0 + ) + return MambaStateShapeCalculator.kda_state_shape( + tp_size, + hf_config.linear_attn_config["num_heads"], + hf_config.linear_attn_config["head_dim"], + conv_kernel_size=hf_config.linear_attn_config["short_conv_kernel_size"], + num_spec=num_spec, + ) + + @classmethod + def get_mamba_state_copy_func( + cls, + ) -> tuple[MambaStateCopyFunc, MambaStateCopyFunc]: + return MambaStateCopyFuncCalculator.kda_state_copy_func() + + def compute_logits( + self, + hidden_states: torch.Tensor, + ) -> torch.Tensor | None: + # The model's final norm is applied here (not at the end of forward) so + # that the pre-norm hidden states can be fed to the MTP draft model. + hidden_states = self.model.norm(hidden_states, None) + return self.logits_processor(self.lm_head, hidden_states) + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + loader = AutoWeightsLoader( + self, + skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None), + ) + return loader.load_weights(weights) diff --git a/aphrodite/models/kimi_k3/amd/model.py b/aphrodite/models/kimi_k3/amd/model.py new file mode 100644 index 0000000000..7f185c8406 --- /dev/null +++ b/aphrodite/models/kimi_k3/amd/model.py @@ -0,0 +1,224 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Kimi-K3 multimodal model implementation for vLLM.""" + +from collections.abc import Iterable +from typing import cast + +import torch +from torch import nn + +from aphrodite.config import AphroditeConfig +from aphrodite.model_executor.layers.quantization import QuantizationConfig +from aphrodite.model_executor.layers.quantization.compressed_tensors import ( + compressed_tensors, +) +from aphrodite.model_executor.models.interfaces import ( + HasInnerState, + IsHybrid, + SupportsEagle3, + SupportsMultiModal, + SupportsPP, + SupportsQuant, +) +from aphrodite.model_executor.models.kimi_k25 import KimiK25MediaPixelInputs +from aphrodite.model_executor.models.kimi_k25_vit import ( + KimiK25MultiModalProjector, + MoonViT3dPretrainedModel, + vision_tower_forward, +) +from aphrodite.model_executor.models.utils import ( + AutoWeightsLoader, + WeightsMapper, + init_vllm_registered_model, + maybe_prefix, +) +from aphrodite.model_executor.models.vision import is_vit_use_data_parallel +from aphrodite.multimodal import MULTIMODAL_REGISTRY +from aphrodite.multimodal.inputs import NestedTensors +from aphrodite.platforms import current_platform +from aphrodite.sequence import IntermediateTensors +from aphrodite.transformers_utils.configs.kimi_k3 import KimiK3Config + +from ..common.mm_preprocess import ( + KimiK3DummyInputsBuilder, + KimiK3MultiModalProcessor, + KimiK3ProcessingInfo, +) +from .linear import KimiLinearForCausalLM + + +@MULTIMODAL_REGISTRY.register_processor( + KimiK3MultiModalProcessor, + info=KimiK3ProcessingInfo, + dummy_inputs=KimiK3DummyInputsBuilder, +) +class KimiK3ForConditionalGeneration( + nn.Module, + SupportsMultiModal, + SupportsPP, + SupportsQuant, + SupportsEagle3, + HasInnerState, + IsHybrid, +): + """Kimi-K3 model with Kimi-K2.5 vision and KimiLinear text.""" + + supports_encoder_tp_data = True + + hf_to_vllm_mapper = WeightsMapper( + orig_to_new_prefix={ + "language_model.layers.": "language_model.model.layers.", + "mm_projector.proj.0": "mm_projector.linear_1", + "mm_projector.proj.2": "mm_projector.linear_2", + } + ) + + @classmethod + def get_placeholder_str(cls, modality: str, i: int) -> str | None: + if modality == "image": + return "<|kimi_image_placeholder|>" + raise ValueError(f"Unsupported modality: {modality}") + + def __init__( + self, + aphrodite_config: AphroditeConfig, + prefix: str = "", + ) -> None: + super().__init__() + model_config = aphrodite_config.model_config + config: KimiK3Config = model_config.hf_config + self.config = config + quant_config = aphrodite_config.quant_config + + multimodal_config = model_config.multimodal_config + assert multimodal_config is not None + self.use_data_parallel = is_vit_use_data_parallel(config.vision_config.num_attention_heads) + self.hidden_size = config.text_config.hidden_size + self.device = current_platform.current_device() + + with self._mark_tower_model(aphrodite_config, "image"): + self.vision_tower = MoonViT3dPretrainedModel( + config.vision_config, + quant_config=self._maybe_ignore_quant_config(quant_config), + prefix=maybe_prefix(prefix, "vision_tower"), + ) + if self._maybe_ignore_quant_config(quant_config) is not None: + self.vision_tower = self.vision_tower.to(device=self.device) + else: + self.vision_tower = self.vision_tower.to(device=self.device, dtype=model_config.dtype) + + self.mm_projector = KimiK25MultiModalProjector( + config=config.vision_config, + use_data_parallel=self.use_data_parallel, + quant_config=self._maybe_ignore_quant_config(quant_config), + prefix=maybe_prefix(prefix, "mm_projector"), + ) + self.mm_projector = self.mm_projector.to(device=self.device, dtype=model_config.dtype) + + self.quant_config = quant_config + with self._mark_language_model(aphrodite_config): + self.language_model = init_vllm_registered_model( + aphrodite_config=aphrodite_config, + hf_config=config.text_config, + prefix=maybe_prefix(prefix, "language_model"), + architectures=["KimiLinearForCausalLM"], + ) + self.make_empty_intermediate_tensors = ( # type: ignore[method-assign] + self.language_model.make_empty_intermediate_tensors + ) + self.media_placeholder: int = self.config.media_placeholder_token_id + + def _maybe_ignore_quant_config(self, quant_config: QuantizationConfig | None) -> QuantizationConfig | None: + if isinstance(quant_config, compressed_tensors.CompressedTensorsConfig): + return None + return quant_config + + def _parse_and_validate_media_input(self, **kwargs: object) -> KimiK25MediaPixelInputs | None: + pixel_values = kwargs.pop("pixel_values", None) + grid_thws = kwargs.pop("grid_thws", None) + if pixel_values is None: + return None + + if isinstance(pixel_values, list): + pixel_values = torch.cat(cast(list[torch.Tensor], pixel_values), dim=0) + if not isinstance(pixel_values, torch.Tensor): + raise TypeError(f"pixel_values must be a tensor or a list of tensors, got {type(pixel_values)}") + + if len(pixel_values.shape) == 5 or len(pixel_values.shape) == 3: + pixel_values = pixel_values.reshape(pixel_values.shape[0] * pixel_values.shape[1], *pixel_values.shape[2:]) + + target_dtype = next(self.vision_tower.parameters()).dtype + pixel_values = pixel_values.to(target_dtype) + assert isinstance(grid_thws, torch.Tensor), f"expect grid_thws to be a tensor, got {type(grid_thws)}" + grid_thws = grid_thws.reshape(-1, grid_thws.shape[-1]) + assert grid_thws.ndim == 2 and grid_thws.size(1) == 3, f"unexpected shape for grid_thws: {grid_thws.shape}" + + return KimiK25MediaPixelInputs( + type="pixel_values", + pixel_values=pixel_values, + grid_thws=grid_thws, + ) + + def _process_media_input(self, media_input: KimiK25MediaPixelInputs) -> list[torch.Tensor]: + media_features = vision_tower_forward( + self.vision_tower, + media_input["pixel_values"], + media_input["grid_thws"], + mm_projector=self.mm_projector, + use_data_parallel=self.use_data_parallel, + ) + return media_features + + def embed_multimodal(self, **kwargs: object) -> NestedTensors | None: + media_input = self._parse_and_validate_media_input(**kwargs) + if media_input is None: + return None + return self._process_media_input(media_input) + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + intermediate_tensors: IntermediateTensors | None = None, + inputs_embeds: torch.Tensor | None = None, + **kwargs: object, + ) -> IntermediateTensors: + if intermediate_tensors is not None: + inputs_embeds = None + hidden_states = self.language_model( + input_ids=input_ids, + positions=positions, + intermediate_tensors=intermediate_tensors, + inputs_embeds=inputs_embeds, + ) + return hidden_states + + def compute_logits(self, hidden_states: torch.Tensor, **kwargs) -> torch.Tensor: + return self.language_model.compute_logits(hidden_states) + + def copy_inputs_before_cuda_graphs(self, input_buffers, **kwargs): + return self.language_model.mamba_cache.copy_inputs_before_cuda_graphs(input_buffers, **kwargs) + + def get_seqlen_agnostic_capture_inputs(self, batch_size: int): + return self.language_model.mamba_cache.get_seqlen_agnostic_capture_inputs(batch_size) + + @classmethod + def get_mamba_state_dtype_from_config(cls, aphrodite_config: AphroditeConfig): + text_config = aphrodite_config.model_config.hf_config.text_config + temp_vllm_config = aphrodite_config.with_hf_config(text_config) + return KimiLinearForCausalLM.get_mamba_state_dtype_from_config(temp_vllm_config) + + @classmethod + def get_mamba_state_shape_from_config(cls, aphrodite_config: AphroditeConfig): + text_config = aphrodite_config.model_config.hf_config.text_config + temp_vllm_config = aphrodite_config.with_hf_config(text_config) + return KimiLinearForCausalLM.get_mamba_state_shape_from_config(temp_vllm_config) + + @classmethod + def get_mamba_state_copy_func(cls): + return KimiLinearForCausalLM.get_mamba_state_copy_func() + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]): + loader = AutoWeightsLoader(self) + return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) diff --git a/aphrodite/models/kimi_k3/amd/mtp.py b/aphrodite/models/kimi_k3/amd/mtp.py new file mode 100644 index 0000000000..b8e88072ff --- /dev/null +++ b/aphrodite/models/kimi_k3/amd/mtp.py @@ -0,0 +1,389 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Inference-only Kimi-K3 Multi-Token-Prediction (MTP) draft model.""" + +import copy +from collections.abc import Iterable + +import torch +import torch.nn as nn + +from aphrodite.config import AphroditeConfig +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.logits_processor import LogitsProcessor +from aphrodite.model_executor.layers.quantization import QuantizationConfig +from aphrodite.model_executor.layers.vocab_parallel_embedding import ( + ParallelLMHead, + VocabParallelEmbedding, +) +from aphrodite.model_executor.model_loader.weight_utils import ( + default_weight_loader, + maybe_remap_kv_scale_name, +) +from aphrodite.model_executor.models.utils import get_pp_missing_layer_names, maybe_prefix +from aphrodite.sequence import IntermediateTensors +from aphrodite.transformers_utils.configs.kimi_linear import KimiLinearConfig + +from ..common.mtp import fused_mtp_input +from .linear import KimiDecoderLayer, get_spec_layer_idx_from_weight_name + +logger = init_logger(__name__) + + +class SharedHead(nn.Module): + def __init__( + self, + config: KimiLinearConfig, + prefix: str, + quant_config: QuantizationConfig | None = None, + ) -> None: + super().__init__() + self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.head = ParallelLMHead( + config.vocab_size, + config.hidden_size, + quant_config=quant_config, + prefix=maybe_prefix(prefix, "head"), + ) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + return self.norm(hidden_states) + + +class KimiK3MultiTokenPredictorLayer(nn.Module): + def __init__( + self, + config: KimiLinearConfig, + aphrodite_config: AphroditeConfig, + prefix: str, + ) -> None: + super().__init__() + self.config = config + quant_config = aphrodite_config.quant_config + + self.enorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.hnorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.eh_proj = nn.Linear(config.hidden_size * 2, config.hidden_size, bias=False) + + self.shared_head = SharedHead(config=config, prefix=prefix, quant_config=quant_config) + # The MTP block is a standard KimiDecoderLayer, but it must NOT use the + # attn-residual (block-residual) scheme even when the base model does: + # the draft starts from the target's hidden state, without the main + # model's accumulated per-block residual tensor. We disable it by + # shallow-copying the config with ``attn_res_block_size=None``. + block_config = copy.copy(config) + block_config.attn_res_block_size = None + # NOTE: the prefix must end in the numeric spec-layer index so that + # KimiDecoderLayer can parse ``layer_idx`` and pick MLA (full attn). + self.mtp_block = KimiDecoderLayer(block_config, aphrodite_config, prefix=prefix) + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + previous_hidden_states: torch.Tensor, + inputs_embeds: torch.Tensor | None = None, + spec_step_index: int = 0, + ) -> tuple[torch.Tensor, torch.Tensor]: + assert inputs_embeds is not None + hidden_states = self.eh_proj( + fused_mtp_input( + positions, + inputs_embeds, + previous_hidden_states, + self.enorm.weight, + self.hnorm.weight, + self.enorm.variance_epsilon, + ) + ) + + hidden_states, residual = self.mtp_block( + positions=positions, + hidden_states=hidden_states, + residual=None, + ) + # Produce the normalized logits input and the pre-norm recurrent state + # in one fused add-RMSNorm launch. + logits_hidden_states, hidden_states = self.shared_head.norm(hidden_states, residual) + return logits_hidden_states, hidden_states + + +class KimiK3MultiTokenPredictor(nn.Module): + def __init__(self, *, aphrodite_config: AphroditeConfig, prefix: str = ""): + super().__init__() + config: KimiLinearConfig = aphrodite_config.model_config.hf_text_config + self.config = config + self.mtp_start_layer_idx = config.num_hidden_layers + self.num_mtp_layers = config.num_nextn_predict_layers + + self.layers = torch.nn.ModuleDict( + { + str(idx): KimiK3MultiTokenPredictorLayer(config, aphrodite_config, f"{prefix}.layers.{idx}") + for idx in range( + self.mtp_start_layer_idx, + self.mtp_start_layer_idx + self.num_mtp_layers, + ) + } + ) + self.embed_tokens = VocabParallelEmbedding( + config.vocab_size, + config.hidden_size, + prefix=maybe_prefix(prefix, "embed_tokens"), + ) + self.logits_processor = LogitsProcessor(config.vocab_size) + + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: + return self.embed_tokens(input_ids) + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + previous_hidden_states: torch.Tensor, + inputs_embeds: torch.Tensor | None = None, + spec_step_idx: int = 0, + ) -> tuple[torch.Tensor, torch.Tensor]: + if inputs_embeds is None: + inputs_embeds = self.embed_tokens(input_ids) + current_step_idx = spec_step_idx % self.num_mtp_layers + return self.layers[str(self.mtp_start_layer_idx + current_step_idx)]( + input_ids, + positions, + previous_hidden_states, + inputs_embeds, + current_step_idx, + ) + + def compute_logits( + self, + hidden_states: torch.Tensor, + spec_step_idx: int = 0, + ) -> torch.Tensor: + current_step_idx = spec_step_idx % self.num_mtp_layers + mtp_layer = self.layers[str(self.mtp_start_layer_idx + current_step_idx)] + logits = self.logits_processor(mtp_layer.shared_head.head, hidden_states) + return logits + + +class KimiK3MTP(nn.Module): + def __init__(self, *, aphrodite_config: AphroditeConfig, prefix: str = ""): + super().__init__() + self.config = aphrodite_config.model_config.hf_text_config + self.quant_config = aphrodite_config.quant_config + self.model = KimiK3MultiTokenPredictor(aphrodite_config=aphrodite_config, prefix=maybe_prefix(prefix, "model")) + + 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, + hidden_states: torch.Tensor, + intermediate_tensors: IntermediateTensors | None = None, + inputs_embeds: torch.Tensor | None = None, + spec_step_idx: int = 0, + ) -> tuple[torch.Tensor, torch.Tensor]: + return self.model( + input_ids, + positions, + hidden_states, + inputs_embeds, + spec_step_idx, + ) + + def compute_logits( + self, + hidden_states: torch.Tensor, + spec_step_idx: int = 0, + ) -> torch.Tensor | None: + return self.model.compute_logits(hidden_states, spec_step_idx) + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + # Mirror KimiLinearForCausalLM.load_weights naming: leading-dot shard + # names, q_lora-conditional fused QKV, and w1/w2/w3 expert weights. + kda_config = self.config.linear_attn_config + use_full_rank_gate = bool(kda_config and kda_config.get("use_full_rank_gate", False)) + beta_shard_id = 5 if use_full_rank_gate else 3 + stacked_params_mapping = [ + # (param_name, shard_name, shard_id) + (".in_proj_qkvgfab", ".q_proj", 0), + (".in_proj_qkvgfab", ".k_proj", 1), + (".in_proj_qkvgfab", ".v_proj", 2), + (".in_proj_qkvgfab", ".b_proj", beta_shard_id), + (".in_proj_qkvgfab", ".f_a_proj", 4), + (".conv1d", ".q_conv1d", 0), + (".conv1d", ".k_conv1d", 1), + (".conv1d", ".v_conv1d", 2), + (".gate_up_proj", ".gate_proj", 0), + (".gate_up_proj", ".up_proj", 1), + ] + if use_full_rank_gate: + stacked_params_mapping.append((".in_proj_qkvgfab", ".g_proj", 3)) + if getattr(self.config, "q_lora_rank", None) is not None: + stacked_params_mapping += [ + (".fused_qkv_a_proj", ".q_a_proj", 0), + (".fused_qkv_a_proj", ".kv_a_proj_with_mqa", 1), + ] + + expert_params_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.num_experts, + ) + if self.config.is_moe + else [] + ) + + pp_missing_layer_names = get_pp_missing_layer_names(self) + params_dict = dict(self.named_parameters()) + # Under the MXFP4 quant interface the routed experts register unpacked + # params (``w13_weight``), while the compressed-tensors checkpoint names + # them ``.weight_packed``. Rebind so the expert mapping resolves; scales + # already share the ``.weight_scale`` suffix. + experts_unpacked = not any(n.endswith("w13_weight_packed") for n in params_dict) + loaded_params: set[str] = set() + for name, loaded_weight in weights: + if "rotary_emb.inv_freq" in name: + continue + # The multimodal checkpoint prefixes text weights with + # ``language_model.``; strip it so names match this draft model's + # parameter paths (``model.layers.{i}.``). Non-text weights + # (vision_tower, mm_projector, ...) never match a spec layer below. + if name.startswith("language_model."): + name = name[len("language_model.") :] + if experts_unpacked and name.endswith(".weight_packed"): + name = name.replace(".weight_packed", ".weight") + spec_layer = get_spec_layer_idx_from_weight_name(self.config, name) + if spec_layer is None: + continue + name = self._rewrite_spec_layer_name(spec_layer, name) + + for param_name, weight_name, shard_id in stacked_params_mapping: + if weight_name not in name: + continue + # Routed experts (``.experts.{i}.w1/w2/w3``) are handled by the + # expert mapping below; skip them here. Shared experts + # (``.shared_experts.``) use gate/up_proj and fall through. + if ".experts." in name: + continue + name_mapped = name.replace(weight_name, param_name) + # Only take this mapping if the fused destination actually + # exists (e.g. QKV fusion is only present when q_lora is used). + if name_mapped not in params_dict: + continue + if name_mapped in pp_missing_layer_names: + continue + name = name_mapped + param = params_dict[name] + weight_loader = param.weight_loader + weight_loader(param, loaded_weight, shard_id) + break + else: + for ( + expert_param_name, + expert_weight_name, + expert_id, + expert_shard_id, + ) in expert_params_mapping: + if expert_weight_name not in name: + continue + name_mapped = name.replace(expert_weight_name, expert_param_name) + if name_mapped in pp_missing_layer_names: + continue + param = params_dict[name_mapped] + weight_loader = param.weight_loader + weight_loader( + param, + loaded_weight, + name_mapped, + shard_id=expert_shard_id, + expert_id=expert_id, + ) + name = name_mapped + break + else: + if name.endswith(".bias") and name not in params_dict: + continue + remapped_name = maybe_remap_kv_scale_name(name, params_dict) + if remapped_name is None: + continue + name = remapped_name + + # The embedding is shared across MTP layers; only the first + # spec layer carries the hoisted (non-".layers") copy. + if spec_layer != self.model.mtp_start_layer_idx and (".layers" not in name): + continue + if name in pp_missing_layer_names: + continue + # The base model uses an attn-residual scheme whose per-layer + # weights (self_attention_res_*, mlp_res_*) are not used by + # the draft block; such names have no matching parameter and + # are safely skipped. + 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) + loaded_params.add(name) + + # Validate that weights were loaded for each expected MTP layer. + loaded_layers: set[int] = set() + for param_name in loaded_params: + spec_layer = get_spec_layer_idx_from_weight_name(self.config, param_name) + if spec_layer is not None: + loaded_layers.add(spec_layer) + for layer_idx in range( + self.model.mtp_start_layer_idx, + self.model.mtp_start_layer_idx + self.model.num_mtp_layers, + ): + if layer_idx not in loaded_layers: + raise ValueError( + f"MTP speculative decoding layer {layer_idx} weights " + f"missing from checkpoint. The checkpoint may not include " + f"the MTP layer weights. Use a checkpoint that includes " + f"MTP layer weights, or disable speculative decoding." + ) + + return loaded_params + + def _rewrite_spec_layer_name(self, spec_layer: int, name: str) -> str: + """Rewrite a checkpoint weight name to this module's parameter path. + + Top-level MTP submodules (enorm/hnorm/eh_proj/shared_head) stay under + ``model.layers.{spec_layer}.*``; the shared ``embed_tokens`` is hoisted + to ``model.*``; everything else is a transformer-block weight and gets + ``.mtp_block`` inserted. + """ + spec_layer_weight_names = [ + "embed_tokens", + "enorm", + "hnorm", + "eh_proj", + "shared_head", + ] + shared_weight_names = ["embed_tokens"] + spec_layer_weight = False + shared_weight = False + for weight_name in spec_layer_weight_names: + if weight_name in name: + spec_layer_weight = True + if weight_name in shared_weight_names: + shared_weight = True + break + if not spec_layer_weight: + name = name.replace( + f"model.layers.{spec_layer}.", + f"model.layers.{spec_layer}.mtp_block.", + ) + elif shared_weight: + name = name.replace(f"model.layers.{spec_layer}.", "model.") + return name diff --git a/aphrodite/models/kimi_k3/amd/ops/third_party/__init__.py b/aphrodite/models/kimi_k3/amd/ops/third_party/__init__.py new file mode 100644 index 0000000000..208f01a7cb --- /dev/null +++ b/aphrodite/models/kimi_k3/amd/ops/third_party/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project diff --git a/aphrodite/models/kimi_k3/amd/ops/third_party/kda/__init__.py b/aphrodite/models/kimi_k3/amd/ops/third_party/kda/__init__.py new file mode 100644 index 0000000000..94a4268ba2 --- /dev/null +++ b/aphrodite/models/kimi_k3/amd/ops/third_party/kda/__init__.py @@ -0,0 +1,47 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# +# AMD/ROCm vendored copy of the Kimi-K3 KDA triton kernels. +# +# Provenance: mirror of aphrodite/models/kimi_k3/nvidia/ops/third_party/kda (tracker +# mke-tracker @ 7adebfcf9; FLA vendored per PRs #39/#86). Split per-vendor so +# AMD can carry gfx950-specific kernel changes without touching the NVIDIA copy. +# +# fla-org/flash-linear-attention#869 (the unmerged ROCm fixes our earlier +# amd_fla shim carried against FLA 0.5.0) is covered here by the *newer* vendored +# FLA rather than the literal patch: +# - transpose-state-layout workaround: N/A (kernels rewritten; no +# transpose_state_layout path remains), +# - AMD autotune configs: present (is_amd num_warps/num_stages branches), +# - OOB-mask correctness fix: present (all tl.load use mask=..., other=0). +# Validated on gfx950: no core-dump, gsm8k 94.1%. +# +# AMD-specific deltas vs the NVIDIA copy: NONE yet (byte-identical). Keep in sync +# with the NVIDIA copy on FLA updates; any divergence should be an intentional, +# documented gfx950-specific change (a #869-style AMD-only fix). + +from .chunk import ( + chunk_kda, + chunk_kda_fwd, + chunk_kda_with_fused_gate, + chunk_kda_with_fused_gate_fwd, + fused_kda_gate, + fused_kda_gate_chunk_cumsum, +) +from .fused_recurrent import ( + fused_recurrent_kda, + fused_recurrent_kda_fwd, + fused_recurrent_kda_packed_decode, +) + +__all__ = [ + "chunk_kda", + "chunk_kda_fwd", + "chunk_kda_with_fused_gate", + "chunk_kda_with_fused_gate_fwd", + "fused_kda_gate", + "fused_kda_gate_chunk_cumsum", + "fused_recurrent_kda", + "fused_recurrent_kda_fwd", + "fused_recurrent_kda_packed_decode", +] diff --git a/aphrodite/models/kimi_k3/amd/ops/third_party/kda/chunk.py b/aphrodite/models/kimi_k3/amd/ops/third_party/kda/chunk.py new file mode 100644 index 0000000000..9362af1238 --- /dev/null +++ b/aphrodite/models/kimi_k3/amd/ops/third_party/kda/chunk.py @@ -0,0 +1,935 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Songlin Yang, Yu Zhang, Zhiyuan Li +# +# This file contains code copied from the flash-linear-attention project. +# The original source code was licensed under the MIT license and included +# the following copyright notice: +# Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li +# ruff: noqa: E501 + + +import torch + +from aphrodite.third_party.flash_linear_attention.ops.chunk_delta_h import ( + chunk_gated_delta_rule_fwd_h, +) +from aphrodite.third_party.flash_linear_attention.ops.cumsum import chunk_local_cumsum +from aphrodite.third_party.flash_linear_attention.ops.index import prepare_chunk_indices +from aphrodite.third_party.flash_linear_attention.ops.l2norm import l2norm_fwd +from aphrodite.third_party.flash_linear_attention.ops.op import exp2, log +from aphrodite.third_party.flash_linear_attention.ops.utils import FLA_CHUNK_SIZE, is_amd +from aphrodite.triton_utils import tl, triton +from aphrodite.utils.math_utils import RCP_LN2, cdiv, next_power_of_2 + +from .chunk_intra import chunk_kda_fwd_intra + +BT_LIST_AUTOTUNE = [32, 64, 128] +NUM_WARPS_AUTOTUNE = [2, 4, 8, 16] if is_amd else [4, 8, 16, 32] + + +@triton.heuristics( + { + "STORE_QG": lambda args: args["qg"] is not None, + "STORE_KG": lambda args: args["kg"] is not None, + "IS_VARLEN": lambda args: args["cu_seqlens"] is not None, + } +) +@triton.autotune( + configs=[ + triton.Config({}, num_warps=num_warps, num_stages=num_stages) + for num_warps in [2, 4, 8] + for num_stages in [2, 3, 4] + ], + key=["H", "K", "V", "BT", "BK", "BV", "IS_VARLEN"], +) +@triton.jit(do_not_specialize=["T"]) +def recompute_w_u_fwd_kernel( + q, + k, + qg, + kg, + v, + beta, + w, + u, + A, + gk, + cu_seqlens, + chunk_indices, + T, + H: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BT: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + STORE_QG: tl.constexpr, + STORE_KG: tl.constexpr, + IS_VARLEN: tl.constexpr, + DOT_PRECISION: tl.constexpr, +): + i_t, i_bh = tl.program_id(0), tl.program_id(1) + i_b, i_h = i_bh // H, i_bh % H + if IS_VARLEN: + i_n, i_t = ( + tl.load(chunk_indices + i_t * 2).to(tl.int32), + tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32), + ) + bos, eos = ( + tl.load(cu_seqlens + i_n).to(tl.int32), + tl.load(cu_seqlens + i_n + 1).to(tl.int32), + ) + T = eos - bos + else: + bos, eos = i_b * T, i_b * T + T + p_b = tl.make_block_ptr(beta + bos * H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) + b_b = tl.load(p_b, boundary_check=(0,)).to(tl.float32) + + p_A = tl.make_block_ptr( + A + (bos * H + i_h) * BT, (T, BT), (H * BT, 1), (i_t * BT, 0), (BT, BT), (1, 0) + ) + b_A = tl.load(p_A, boundary_check=(0, 1)) + + for i_v in range(tl.cdiv(V, BV)): + p_v = tl.make_block_ptr( + v + (bos * H + i_h) * V, + (T, V), + (H * V, 1), + (i_t * BT, i_v * BV), + (BT, BV), + (1, 0), + ) + p_u = tl.make_block_ptr( + u + (bos * H + i_h) * V, + (T, V), + (H * V, 1), + (i_t * BT, i_v * BV), + (BT, BV), + (1, 0), + ) + b_v = tl.load(p_v, boundary_check=(0, 1)) + b_vb = (b_v * b_b[:, None]).to(b_v.dtype) + b_u = tl.dot(b_A, b_vb, input_precision=DOT_PRECISION) + tl.store(p_u, b_u.to(p_u.dtype.element_ty), boundary_check=(0, 1)) + + for i_k in range(tl.cdiv(K, BK)): + p_w = tl.make_block_ptr( + w + (bos * H + i_h) * K, + (T, K), + (H * K, 1), + (i_t * BT, i_k * BK), + (BT, BK), + (1, 0), + ) + p_k = tl.make_block_ptr( + k + (bos * H + i_h) * K, + (T, K), + (H * K, 1), + (i_t * BT, i_k * BK), + (BT, BK), + (1, 0), + ) + b_k = tl.load(p_k, boundary_check=(0, 1)) + b_kb = b_k * b_b[:, None] + + p_gk = tl.make_block_ptr( + gk + (bos * H + i_h) * K, + (T, K), + (H * K, 1), + (i_t * BT, i_k * BK), + (BT, BK), + (1, 0), + ) + b_gk = tl.load(p_gk, boundary_check=(0, 1)) + b_kb *= exp2(b_gk) + if STORE_QG: + p_q = tl.make_block_ptr( + q + (bos * H + i_h) * K, + (T, K), + (H * K, 1), + (i_t * BT, i_k * BK), + (BT, BK), + (1, 0), + ) + p_qg = tl.make_block_ptr( + qg + (bos * H + i_h) * K, + (T, K), + (H * K, 1), + (i_t * BT, i_k * BK), + (BT, BK), + (1, 0), + ) + b_q = tl.load(p_q, boundary_check=(0, 1)) + b_qg = b_q * exp2(b_gk) + tl.store(p_qg, b_qg.to(p_qg.dtype.element_ty), boundary_check=(0, 1)) + if STORE_KG: + last_idx = min(i_t * BT + BT, T) - 1 + + o_k = i_k * BK + tl.arange(0, BK) + m_k = o_k < K + b_gn = tl.load( + gk + ((bos + last_idx) * H + i_h) * K + o_k, mask=m_k, other=0.0 + ) + b_kg = b_k * exp2(b_gn - b_gk) + + p_kg = tl.make_block_ptr( + kg + (bos * H + i_h) * K, + (T, K), + (H * K, 1), + (i_t * BT, i_k * BK), + (BT, BK), + (1, 0), + ) + tl.store(p_kg, b_kg.to(p_kg.dtype.element_ty), boundary_check=(0, 1)) + + b_w = tl.dot(b_A, b_kb.to(b_k.dtype)) + tl.store(p_w, b_w.to(p_w.dtype.element_ty), boundary_check=(0, 1)) + + +def recompute_w_u_fwd( + k: torch.Tensor, + v: torch.Tensor, + beta: torch.Tensor, + A: torch.Tensor, + q: torch.Tensor | None = None, + gk: torch.Tensor | None = None, + cu_seqlens: torch.Tensor | None = None, + chunk_indices: torch.Tensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor]: + B, T, H, K, V = *k.shape, v.shape[-1] + BT = A.shape[-1] + BK = 64 + BV = 64 + + if chunk_indices is None and cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) + NT = cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) + + w = torch.empty_like(k) + u = torch.empty_like(v) + kg = torch.empty_like(k) if gk is not None else None + recompute_w_u_fwd_kernel[(NT, B * H)]( + q=q, + k=k, + qg=None, + kg=kg, + v=v, + beta=beta, + w=w, + u=u, + A=A, + gk=gk, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + T=T, + H=H, + K=K, + V=V, + BT=BT, + BK=BK, + BV=BV, + DOT_PRECISION="ieee", + ) + return w, u, None, kg + + +@triton.heuristics({"IS_VARLEN": lambda args: args["cu_seqlens"] is not None}) +@triton.autotune( + configs=[ + triton.Config({"BK": BK, "BV": BV}, num_warps=num_warps, num_stages=num_stages) + for BK in [32, 64] + for BV in [64, 128] + for num_warps in [2, 4, 8] + for num_stages in [2, 3, 4] + ], + key=["BT"], +) +@triton.jit(do_not_specialize=["T"]) +def chunk_gla_fwd_kernel_o( + q, + v, + g, + h, + o, + A, + cu_seqlens, + chunk_indices, + scale, + T, + H: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BT: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_v, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_b, i_h = i_bh // H, i_bh % H + if IS_VARLEN: + i_tg = i_t + i_n, i_t = ( + tl.load(chunk_indices + i_t * 2).to(tl.int32), + tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32), + ) + bos, eos = ( + tl.load(cu_seqlens + i_n).to(tl.int32), + tl.load(cu_seqlens + i_n + 1).to(tl.int32), + ) + T = eos - bos + NT = tl.cdiv(T, BT) + else: + NT = tl.cdiv(T, BT) + i_tg = i_b * NT + i_t + bos, eos = i_b * T, i_b * T + T + + m_s = tl.arange(0, BT)[:, None] >= tl.arange(0, BT)[None, :] + + b_o = tl.zeros([BT, BV], dtype=tl.float32) + for i_k in range(tl.cdiv(K, BK)): + p_q = tl.make_block_ptr( + q + (bos * H + i_h) * K, + (T, K), + (H * K, 1), + (i_t * BT, i_k * BK), + (BT, BK), + (1, 0), + ) + p_g = tl.make_block_ptr( + g + (bos * H + i_h) * K, + (T, K), + (H * K, 1), + (i_t * BT, i_k * BK), + (BT, BK), + (1, 0), + ) + p_h = tl.make_block_ptr( + h + (i_tg * H + i_h) * K * V, + (V, K), + (K, 1), + (i_v * BV, i_k * BK), + (BV, BK), + (1, 0), + ) + + # [BT, BK] + b_q = tl.load(p_q, boundary_check=(0, 1)) + b_q = (b_q * scale).to(b_q.dtype) + # [BT, BK] + b_g = tl.load(p_g, boundary_check=(0, 1)) + # [BT, BK] + b_qg = (b_q * exp2(b_g)).to(b_q.dtype) + # [BV, BK] + b_h = tl.load(p_h, boundary_check=(0, 1)) + # [BT, BV] + if i_k >= 0: + b_o += tl.dot(b_qg, tl.trans(b_h).to(b_qg.dtype)) + p_v = tl.make_block_ptr( + v + (bos * H + i_h) * V, + (T, V), + (H * V, 1), + (i_t * BT, i_v * BV), + (BT, BV), + (1, 0), + ) + p_o = tl.make_block_ptr( + o + (bos * H + i_h) * V, + (T, V), + (H * V, 1), + (i_t * BT, i_v * BV), + (BT, BV), + (1, 0), + ) + p_A = tl.make_block_ptr( + A + (bos * H + i_h) * BT, (T, BT), (H * BT, 1), (i_t * BT, 0), (BT, BT), (1, 0) + ) + # [BT, BV] + b_v = tl.load(p_v, boundary_check=(0, 1)) + # [BT, BT] + b_A = tl.load(p_A, boundary_check=(0, 1)) + b_A = tl.where(m_s, b_A, 0.0).to(b_v.dtype) + b_o += tl.dot(b_A, b_v, allow_tf32=False) + tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1)) + + +def chunk_gla_fwd_o_gk( + q: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + A: torch.Tensor, + h: torch.Tensor, + o: torch.Tensor, + scale: float, + cu_seqlens: torch.Tensor | None = None, + chunk_indices: torch.Tensor | None = None, + chunk_size: int = FLA_CHUNK_SIZE, +): + B, T, H, K, V = *q.shape, v.shape[-1] + BT = chunk_size + + if chunk_indices is None and cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size) + NT = cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) + + def grid(meta): + return (cdiv(V, meta["BV"]), NT, B * H) + + chunk_gla_fwd_kernel_o[grid]( + q=q, + v=v, + g=g, + h=h, + o=o, + A=A, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + scale=scale, + T=T, + H=H, + K=K, + V=V, + BT=BT, + ) + return o + + +@triton.heuristics( + { + "HAS_BIAS": lambda args: args["g_bias"] is not None, + "IS_VARLEN": lambda args: args["cu_seqlens"] is not None, + } +) +@triton.autotune( + configs=[ + triton.Config({"BS": BS}, num_warps=num_warps) + for BS in [32, 64] + for num_warps in [2, 4, 8] + ], + key=["H", "S", "BT", "IS_VARLEN"], +) +@triton.jit(do_not_specialize=["T"]) +def kda_gate_chunk_cumsum_vector_kernel( + s, + raw_beta, + A_log, + g_bias, + o, + beta_out, + cu_seqlens, + chunk_indices, + cumsum_scale, + lower_bound, + beta, + threshold, + T, + stride_beta_batch, + stride_beta_token, + stride_beta_head, + H: tl.constexpr, + S: tl.constexpr, + BT: tl.constexpr, + BS: tl.constexpr, + HAS_BIAS: tl.constexpr, + IS_VARLEN: tl.constexpr, + USE_LOWER_BOUND: tl.constexpr, +): + i_s, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_b, i_h = i_bh // H, i_bh % H + if IS_VARLEN: + i_n, i_t = ( + tl.load(chunk_indices + i_t * 2).to(tl.int32), + tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32), + ) + bos, eos = ( + tl.load(cu_seqlens + i_n).to(tl.int32), + tl.load(cu_seqlens + i_n + 1).to(tl.int32), + ) + T = eos - bos + else: + bos = i_b * T + + if i_s == 0: + o_beta_t = tl.arange(0, BT) + m_beta = i_t * BT + o_beta_t < T + if IS_VARLEN: + p_beta = ( + raw_beta + + (bos + i_t * BT + o_beta_t) * stride_beta_token + + i_h * stride_beta_head + ) + else: + p_beta = ( + raw_beta + + i_b * stride_beta_batch + + (i_t * BT + o_beta_t) * stride_beta_token + + i_h * stride_beta_head + ) + b_beta = tl.load(p_beta, mask=m_beta, other=0.0).to(tl.float32) + p_beta_out = beta_out + (bos + i_t * BT + o_beta_t) * H + i_h + tl.store(p_beta_out, tl.sigmoid(b_beta), mask=m_beta) + return + + i_s -= 1 + + p_s = tl.make_block_ptr( + s + (bos * H + i_h) * S, + (T, S), + (H * S, 1), + (i_t * BT, i_s * BS), + (BT, BS), + (1, 0), + ) + p_o = tl.make_block_ptr( + o + (bos * H + i_h) * S, + (T, S), + (H * S, 1), + (i_t * BT, i_s * BS), + (BT, BS), + (1, 0), + ) + + b_s = tl.load(p_s, boundary_check=(0, 1)).to(tl.float32) + if HAS_BIAS: + p_bias = tl.make_block_ptr( + g_bias + i_h * S, + (S,), + (1,), + (i_s * BS,), + (BS,), + (0,), + ) + b_bias = tl.load(p_bias, boundary_check=(0,)).to(tl.float32) + b_s += b_bias[None, :] + + b_a = tl.exp(tl.load(A_log + i_h).to(tl.float32)) + if USE_LOWER_BOUND: + b_gate = lower_bound * tl.sigmoid(b_a * b_s) + else: + b_g_scaled = b_s * beta + b_softplus = tl.where( + b_g_scaled > threshold, + b_s, + (1.0 / beta) * log(1.0 + tl.exp(b_g_scaled)), + ) + b_gate = -b_a * b_softplus + + # Boundary loads return zero, but bias and gate activation can make padded + # rows nonzero. Padding trails valid rows, so it only affects masked stores. + b_o = tl.cumsum(b_gate, axis=0) * cumsum_scale + tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1)) + + +def fused_kda_gate_chunk_cumsum( + raw_g: torch.Tensor, + raw_beta: torch.Tensor, + A_log: torch.Tensor, + g_bias: torch.Tensor | None = None, + beta: float = 1.0, + threshold: float = 20.0, + lower_bound: float | None = None, + cu_seqlens: torch.Tensor | None = None, + chunk_indices: torch.Tensor | None = None, + chunk_size: int = FLA_CHUNK_SIZE, + output_dtype: torch.dtype | None = torch.float, +) -> tuple[torch.Tensor, torch.Tensor]: + if cu_seqlens is not None: + assert raw_g.shape[0] == 1, ( + "Only batch size 1 is supported when cu_seqlens are provided" + ) + B, T, H, D = raw_g.shape + if raw_beta.shape != (B, T, H): + raise ValueError(f"Expected raw_beta shape {(B, T, H)}, got {raw_beta.shape}") + if chunk_indices is None and cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size) + NT = cdiv(T, chunk_size) if cu_seqlens is None else len(chunk_indices) + + A_log = A_log.reshape(-1) + if g_bias is not None: + g_bias = g_bias.reshape(-1) + y = torch.empty_like(raw_g, dtype=output_dtype or raw_g.dtype) + beta_out = torch.empty(raw_beta.shape, device=raw_beta.device, dtype=torch.float32) + + def grid(meta): + # For each (chunk, head), program 0 computes beta without extending a + # gate tile's critical path. The remaining programs cover the gate dim. + return (cdiv(meta["S"], meta["BS"]) + 1, NT, B * H) + + kda_gate_chunk_cumsum_vector_kernel[grid]( + s=raw_g, + raw_beta=raw_beta, + A_log=A_log, + g_bias=g_bias, + o=y, + beta_out=beta_out, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + # RCP_LN2 folds in the natural-log -> log2 conversion so downstream + # exp2-based kernels reproduce exp(g). Keep this in sync with the + # `use_exp2=True` path in `_chunk_kda_fwd_with_cumulative_g`. + cumsum_scale=RCP_LN2, + lower_bound=lower_bound or 0.0, + beta=beta, + threshold=threshold, + T=T, + stride_beta_batch=raw_beta.stride(0), + stride_beta_token=raw_beta.stride(1), + stride_beta_head=raw_beta.stride(2), + H=H, + S=D, + BT=chunk_size, + USE_LOWER_BOUND=lower_bound is not None, + ) + return y, beta_out + + +def _chunk_kda_fwd_with_cumulative_g( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + scale: float, + initial_state: torch.Tensor, + output_final_state: bool, + cu_seqlens: torch.Tensor | None = None, + chunk_indices: torch.Tensor | None = None, + chunk_size: int = FLA_CHUNK_SIZE, + safe_gate: bool = False, +): + # `g` must already be chunk-local cumulatively-summed AND scaled by + # RCP_LN2 (so the downstream exp2-based kernels reproduce exp(g)). + # Use `chunk_kda_fwd` or `chunk_kda_with_fused_gate_fwd` instead of + # calling this helper directly unless that invariant is upheld. + Aqk, A = chunk_kda_fwd_intra( + q=q, + k=k, + gk=g, + beta=beta, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + chunk_size=chunk_size, + safe_gate=safe_gate, + ) + w, u, _, kg = recompute_w_u_fwd( + k=k, + v=v, + beta=beta, + A=A, + gk=g, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + ) + del A + h, v_new, final_state = chunk_gated_delta_rule_fwd_h( + k=kg, + w=w, + u=u, + gk=g, + initial_state=initial_state, + output_final_state=output_final_state, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + use_exp2=True, + ) + del w, u, kg + o = chunk_gla_fwd_o_gk( + q=q, + v=v_new, + g=g, + A=Aqk, + h=h, + o=v, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + chunk_size=chunk_size, + ) + del Aqk, v_new, h + return o, final_state + + +def chunk_kda_fwd( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + scale: float, + initial_state: torch.Tensor, + output_final_state: bool, + cu_seqlens: torch.Tensor | None = None, +): + chunk_size = FLA_CHUNK_SIZE + chunk_indices = ( + prepare_chunk_indices(cu_seqlens, chunk_size) + if cu_seqlens is not None + else None + ) + g = chunk_local_cumsum( + g, + chunk_size=chunk_size, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + ) + # KDA evaluates cumulative gate decays with exp2. Convert from natural-log + # space so exp(x) is preserved as exp2(x / ln(2)). + g = g * RCP_LN2 + return _chunk_kda_fwd_with_cumulative_g( + q=q, + k=k, + v=v, + g=g, + beta=beta, + scale=scale, + initial_state=initial_state, + output_final_state=output_final_state, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + chunk_size=chunk_size, + ) + + +def chunk_kda_with_fused_gate_fwd( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + raw_g: torch.Tensor, + raw_beta: torch.Tensor, + A_log: torch.Tensor, + g_bias: torch.Tensor | None, + scale: float, + initial_state: torch.Tensor, + output_final_state: bool, + lower_bound: float | None = None, + cu_seqlens: torch.Tensor | None = None, +): + chunk_size = FLA_CHUNK_SIZE + chunk_indices = ( + prepare_chunk_indices(cu_seqlens, chunk_size) + if cu_seqlens is not None + else None + ) + g, beta = fused_kda_gate_chunk_cumsum( + raw_g, + raw_beta=raw_beta, + A_log=A_log, + g_bias=g_bias, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + chunk_size=chunk_size, + lower_bound=lower_bound, + ) + return _chunk_kda_fwd_with_cumulative_g( + q=q, + k=k, + v=v, + g=g, + beta=beta, + scale=scale, + initial_state=initial_state, + output_final_state=output_final_state, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + chunk_size=chunk_size, + safe_gate=lower_bound is not None, + ) + + +def chunk_kda( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + scale: float = None, + initial_state: torch.Tensor = None, + output_final_state: bool = False, + use_qk_l2norm_in_kernel: bool = False, + cu_seqlens: torch.Tensor | None = None, + **kwargs, +): + if scale is None: + scale = k.shape[-1] ** -0.5 + + if use_qk_l2norm_in_kernel: + q = l2norm_fwd(q.contiguous()) + k = l2norm_fwd(k.contiguous()) + + o, final_state = chunk_kda_fwd( + q=q, + k=k, + v=v.contiguous(), + g=g.contiguous(), + beta=beta.contiguous(), + scale=scale, + initial_state=initial_state.contiguous(), + output_final_state=output_final_state, + cu_seqlens=cu_seqlens, + ) + return o, final_state + + +def chunk_kda_with_fused_gate( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + raw_g: torch.Tensor, + raw_beta: torch.Tensor, + A_log: torch.Tensor, + g_bias: torch.Tensor | None, + scale: float | None = None, + initial_state: torch.Tensor | None = None, + output_final_state: bool = False, + lower_bound: float | None = None, + use_qk_l2norm_in_kernel: bool = False, + cu_seqlens: torch.Tensor | None = None, + **kwargs, +): + """Run chunk KDA from raw gate and beta projections.""" + if scale is None: + scale = k.shape[-1] ** -0.5 + + if use_qk_l2norm_in_kernel: + q = l2norm_fwd(q.contiguous()) + k = l2norm_fwd(k.contiguous()) + + o, final_state = chunk_kda_with_fused_gate_fwd( + q=q, + k=k, + v=v.contiguous(), + raw_g=raw_g.contiguous(), + raw_beta=raw_beta, + A_log=A_log, + g_bias=g_bias, + scale=scale, + initial_state=initial_state.contiguous() if initial_state is not None else None, + output_final_state=output_final_state, + lower_bound=lower_bound, + cu_seqlens=cu_seqlens, + ) + return o, final_state + + +@triton.autotune( + configs=[ + triton.Config({"BT": bt}, num_warps=nw, num_stages=ns) + for bt in BT_LIST_AUTOTUNE + for nw in NUM_WARPS_AUTOTUNE + for ns in [2, 3] + ], + key=["H", "D"], +) +@triton.jit +def kda_gate_fwd_kernel( + g, + A, + y, + g_bias, + lower_bound, + beta: tl.constexpr, + threshold: tl.constexpr, + T, + H, + D: tl.constexpr, + BT: tl.constexpr, + BD: tl.constexpr, + HAS_BIAS: tl.constexpr, + USE_LOWER_BOUND: tl.constexpr, +): + i_t, i_h = tl.program_id(0), tl.program_id(1) + n_t = i_t * BT + + b_a = tl.exp(tl.load(A + i_h).to(tl.float32)) + + stride_row = H * D + stride_col = 1 + + g_ptr = tl.make_block_ptr( + base=g + i_h * D, + shape=(T, D), + strides=(stride_row, stride_col), + offsets=(n_t, 0), + block_shape=(BT, BD), + order=(1, 0), + ) + + y_ptr = tl.make_block_ptr( + base=y + i_h * D, + shape=(T, D), + strides=(stride_row, stride_col), + offsets=(n_t, 0), + block_shape=(BT, BD), + order=(1, 0), + ) + + b_g = tl.load(g_ptr, boundary_check=(0, 1)).to(tl.float32) + + if HAS_BIAS: + n_d = tl.arange(0, BD) + bias_mask = n_d < D + b_bias = tl.load(g_bias + i_h * D + n_d, mask=bias_mask, other=0.0).to( + tl.float32 + ) + b_g = b_g + b_bias[None, :] + + if USE_LOWER_BOUND: + b_y = lower_bound * tl.sigmoid(b_a * b_g) + else: + g_scaled = b_g * beta + use_linear = g_scaled > threshold + sp = tl.where(use_linear, b_g, (1.0 / beta) * log(1.0 + tl.exp(g_scaled))) + b_y = -b_a * sp + + tl.store(y_ptr, b_y.to(y.dtype.element_ty), boundary_check=(0, 1)) + + +def fused_kda_gate( + g: torch.Tensor, + A: torch.Tensor, + head_k_dim: int, + g_bias: torch.Tensor | None = None, + beta: float = 1.0, + threshold: float = 20.0, + lower_bound: float | None = None, +) -> torch.Tensor: + """ + Forward pass for KDA gate: + input g: [..., H*D] + param A: [H] or [1, 1, H, 1] + beta: softplus beta parameter + threshold: softplus threshold parameter + return : [..., H, D] + """ + orig_shape = g.shape[:-1] + + g = g.view(-1, g.shape[-1]) + T = g.shape[0] + HD = g.shape[1] + H = A.numel() + assert H * head_k_dim == HD + + y = torch.empty_like(g, dtype=torch.float32) + + def grid(meta): + return (cdiv(T, meta["BT"]), H) + + kda_gate_fwd_kernel[grid]( + g, + A, + y, + g_bias, + lower_bound or 0.0, + beta, + threshold, + T, + H, + head_k_dim, + BD=next_power_of_2(head_k_dim), + HAS_BIAS=g_bias is not None, + USE_LOWER_BOUND=lower_bound is not None, + ) + + y = y.view(*orig_shape, H, head_k_dim) + return y diff --git a/aphrodite/models/kimi_k3/amd/ops/third_party/kda/chunk_intra.py b/aphrodite/models/kimi_k3/amd/ops/third_party/kda/chunk_intra.py new file mode 100644 index 0000000000..71735f96a6 --- /dev/null +++ b/aphrodite/models/kimi_k3/amd/ops/third_party/kda/chunk_intra.py @@ -0,0 +1,662 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Songlin Yang, Yu Zhang, Zhiyuan Li +# +# This file contains code copied from the flash-linear-attention project. +# The original source was licensed under the MIT license. +# Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li +# Forward-only adaptation of flash-linear-attention 0.5.0. +# ruff: noqa: E501 + +import torch + +from aphrodite.platforms import current_platform +from aphrodite.third_party.flash_linear_attention.ops.index import prepare_chunk_indices +from aphrodite.third_party.flash_linear_attention.ops.op import exp2, gather +from aphrodite.third_party.flash_linear_attention.ops.utils import is_gather_supported +from aphrodite.triton_utils import tl, triton + +from .chunk_intra_token_parallel import chunk_kda_fwd_intra_token_parallel + +################################################################################ +# Fused inter + solve_tril kernel: compute off-diagonal Akk and solve in one pass +################################################################################ + + +@triton.heuristics( + { + "IS_VARLEN": lambda args: args["cu_seqlens"] is not None, + } +) +@triton.autotune( + configs=[ + triton.Config({"BK": BK}, num_warps=num_warps) + for BK in [32, 64] + for num_warps in [1, 2, 4] + ], + key=["H", "HV", "K", "BC"], +) +@triton.jit(do_not_specialize=["T"]) +def chunk_kda_fwd_kernel_inter_solve_fused( + q, + k, + g, + beta, + Aqk, + Akkd, + Akk, + scale, + cu_seqlens, + chunk_indices, + T, + H: tl.constexpr, + HV: tl.constexpr, + K: tl.constexpr, + BT: tl.constexpr, + BC: tl.constexpr, + BK: tl.constexpr, + IS_VARLEN: tl.constexpr, + USE_SAFE_GATE: tl.constexpr, + SOLVE_TRIL_DOT_PRECISION: tl.constexpr, +): + """ + Fused kernel: compute inter-subchunk Akk + solve_tril in one pass. + Prerequisite: token_parallel has already computed diagonal Akk blocks in Akkd. + + This kernel: + 1. Computes off-diagonal Aqk blocks -> writes to global + 2. Computes off-diagonal Akk blocks -> keeps in registers + 3. Loads diagonal Akk blocks from Akkd (fp32) + 4. Does forward substitution on diagonals + 5. Computes merged Akk_inv + 6. Writes Akk_inv to Akk + """ + i_t, i_bh = tl.program_id(0), tl.program_id(1) + i_b, i_hv = i_bh // HV, i_bh % HV + i_h = i_hv // (HV // H) + + if IS_VARLEN: + i_n, i_t = ( + tl.load(chunk_indices + i_t * 2).to(tl.int32), + tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32), + ) + bos, eos = ( + tl.load(cu_seqlens + i_n).to(tl.int32), + tl.load(cu_seqlens + i_n + 1).to(tl.int32), + ) + T = eos - bos + else: + bos, eos = i_b * T, i_b * T + T + + if i_t * BT >= T: + return + + i_tc0 = i_t * BT + i_tc1 = i_t * BT + BC + i_tc2 = i_t * BT + 2 * BC + i_tc3 = i_t * BT + 3 * BC + + q += (bos * H + i_h) * K + k += (bos * H + i_h) * K + g += (bos * HV + i_hv) * K + Aqk += (bos * HV + i_hv) * BT + Akk += (bos * HV + i_hv) * BT + Akkd += (bos * HV + i_hv) * BC + + o_i = tl.arange(0, BC) + m_tc1 = (i_tc1 + o_i) < T + m_tc2 = (i_tc2 + o_i) < T + m_tc3 = (i_tc3 + o_i) < T + + b_Aqk10 = tl.zeros([BC, BC], dtype=tl.float32) + b_Akk10 = tl.zeros([BC, BC], dtype=tl.float32) + + b_Aqk20 = tl.zeros([BC, BC], dtype=tl.float32) + b_Akk20 = tl.zeros([BC, BC], dtype=tl.float32) + b_Aqk21 = tl.zeros([BC, BC], dtype=tl.float32) + b_Akk21 = tl.zeros([BC, BC], dtype=tl.float32) + + b_Aqk30 = tl.zeros([BC, BC], dtype=tl.float32) + b_Akk30 = tl.zeros([BC, BC], dtype=tl.float32) + b_Aqk31 = tl.zeros([BC, BC], dtype=tl.float32) + b_Akk31 = tl.zeros([BC, BC], dtype=tl.float32) + b_Aqk32 = tl.zeros([BC, BC], dtype=tl.float32) + b_Akk32 = tl.zeros([BC, BC], dtype=tl.float32) + + ################################################################################ + # off-diagonal blocks + ################################################################################ + for i_k in range(tl.cdiv(K, BK)): + o_k = i_k * BK + tl.arange(0, BK) + m_k = o_k < K + + p_k0 = tl.make_block_ptr( + k, (T, K), (H * K, 1), (i_tc0, i_k * BK), (BC, BK), (1, 0) + ) + p_g0 = tl.make_block_ptr( + g, (T, K), (HV * K, 1), (i_tc0, i_k * BK), (BC, BK), (1, 0) + ) + b_k0 = tl.load(p_k0, boundary_check=(0, 1)).to(tl.float32) + b_g0 = tl.load(p_g0, boundary_check=(0, 1)).to(tl.float32) + + if i_tc1 < T: + p_q1 = tl.make_block_ptr( + q, (T, K), (H * K, 1), (i_tc1, i_k * BK), (BC, BK), (1, 0) + ) + p_k1 = tl.make_block_ptr( + k, (T, K), (H * K, 1), (i_tc1, i_k * BK), (BC, BK), (1, 0) + ) + p_g1 = tl.make_block_ptr( + g, (T, K), (HV * K, 1), (i_tc1, i_k * BK), (BC, BK), (1, 0) + ) + # [BC, BK] + b_q1 = tl.load(p_q1, boundary_check=(0, 1)).to(tl.float32) + b_k1 = tl.load(p_k1, boundary_check=(0, 1)).to(tl.float32) + b_g1 = tl.load(p_g1, boundary_check=(0, 1)).to(tl.float32) + # [BK] + b_gn1 = tl.load(g + i_tc1 * HV * K + o_k, mask=m_k, other=0).to(tl.float32) + # [BC, BK] + b_gqn = tl.where(m_tc1[:, None], exp2(b_g1 - b_gn1[None, :]), 0) + # [BK, BC] + b_kgt = tl.trans(b_k0 * exp2(b_gn1[None, :] - b_g0)) + # [BC, BC] + b_Aqk10 += tl.dot(b_q1 * b_gqn, b_kgt) + b_Akk10 += tl.dot(b_k1 * b_gqn, b_kgt) + + if i_tc2 < T: + p_q2 = tl.make_block_ptr( + q, (T, K), (H * K, 1), (i_tc2, i_k * BK), (BC, BK), (1, 0) + ) + p_k2 = tl.make_block_ptr( + k, (T, K), (H * K, 1), (i_tc2, i_k * BK), (BC, BK), (1, 0) + ) + p_g2 = tl.make_block_ptr( + g, (T, K), (HV * K, 1), (i_tc2, i_k * BK), (BC, BK), (1, 0) + ) + # [BC, BK] + b_q2 = tl.load(p_q2, boundary_check=(0, 1)).to(tl.float32) + b_k2 = tl.load(p_k2, boundary_check=(0, 1)).to(tl.float32) + b_g2 = tl.load(p_g2, boundary_check=(0, 1)).to(tl.float32) + # [BK] + b_gn2 = tl.load(g + i_tc2 * HV * K + o_k, mask=m_k, other=0).to( + tl.float32 + ) + # [BC, BK] + b_gqn2 = tl.where(m_tc2[:, None], exp2(b_g2 - b_gn2[None, :]), 0) + b_qg2 = b_q2 * b_gqn2 + b_kg2 = b_k2 * b_gqn2 + # [BK, BC] + b_kgt = tl.trans(b_k0 * exp2(b_gn2[None, :] - b_g0)) + b_Aqk20 += tl.dot(b_qg2, b_kgt) + b_Akk20 += tl.dot(b_kg2, b_kgt) + # [BC, BC] + b_kgt = tl.trans(b_k1 * exp2(b_gn2[None, :] - b_g1)) + # [BC, BC] + b_Aqk21 += tl.dot(b_qg2, b_kgt) + b_Akk21 += tl.dot(b_kg2, b_kgt) + + if i_tc3 < T: + p_q3 = tl.make_block_ptr( + q, (T, K), (H * K, 1), (i_tc3, i_k * BK), (BC, BK), (1, 0) + ) + p_k3 = tl.make_block_ptr( + k, (T, K), (H * K, 1), (i_tc3, i_k * BK), (BC, BK), (1, 0) + ) + p_g3 = tl.make_block_ptr( + g, (T, K), (HV * K, 1), (i_tc3, i_k * BK), (BC, BK), (1, 0) + ) + # [BC, BK] + b_q3 = tl.load(p_q3, boundary_check=(0, 1)).to(tl.float32) + b_k3 = tl.load(p_k3, boundary_check=(0, 1)).to(tl.float32) + b_g3 = tl.load(p_g3, boundary_check=(0, 1)).to(tl.float32) + # [BK] + b_gn3 = tl.load(g + i_tc3 * HV * K + o_k, mask=m_k, other=0).to( + tl.float32 + ) + # [BC, BK] + b_gqn3 = tl.where(m_tc3[:, None], exp2(b_g3 - b_gn3[None, :]), 0) + b_qg3 = b_q3 * b_gqn3 + b_kg3 = b_k3 * b_gqn3 + # [BK, BC] + b_kgt = tl.trans(b_k0 * exp2(b_gn3[None, :] - b_g0)) + # [BC, BC] + b_Aqk30 += tl.dot(b_qg3, b_kgt) + b_Akk30 += tl.dot(b_kg3, b_kgt) + # [BK, BC] + b_kgt = tl.trans(b_k1 * exp2(b_gn3[None, :] - b_g1)) + # [BC, BC] + b_Aqk31 += tl.dot(b_qg3, b_kgt) + b_Akk31 += tl.dot(b_kg3, b_kgt) + # [BK, BC] + b_kgt = tl.trans(b_k2 * exp2(b_gn3[None, :] - b_g2)) + # [BC, BC] + b_Aqk32 += tl.dot(b_qg3, b_kgt) + b_Akk32 += tl.dot(b_kg3, b_kgt) + + ################################################################################ + # save off-diagonal Aqk blocks and prepare Akk + ################################################################################ + if i_tc1 < T: + p_Aqk10 = tl.make_block_ptr( + Aqk, (T, BT), (HV * BT, 1), (i_tc1, 0), (BC, BC), (1, 0) + ) + tl.store( + p_Aqk10, (b_Aqk10 * scale).to(Aqk.dtype.element_ty), boundary_check=(0, 1) + ) + + p_b1 = tl.make_block_ptr( + beta + bos * HV + i_hv, (T,), (HV,), (i_tc1,), (BC,), (0,) + ) + b_b1 = tl.load(p_b1, boundary_check=(0,)).to(tl.float32) + b_Akk10 = b_Akk10 * b_b1[:, None] + if i_tc2 < T: + p_Aqk20 = tl.make_block_ptr( + Aqk, (T, BT), (HV * BT, 1), (i_tc2, 0), (BC, BC), (1, 0) + ) + p_Aqk21 = tl.make_block_ptr( + Aqk, (T, BT), (HV * BT, 1), (i_tc2, BC), (BC, BC), (1, 0) + ) + tl.store( + p_Aqk20, (b_Aqk20 * scale).to(Aqk.dtype.element_ty), boundary_check=(0, 1) + ) + tl.store( + p_Aqk21, (b_Aqk21 * scale).to(Aqk.dtype.element_ty), boundary_check=(0, 1) + ) + + p_b2 = tl.make_block_ptr( + beta + bos * HV + i_hv, (T,), (HV,), (i_tc2,), (BC,), (0,) + ) + b_b2 = tl.load(p_b2, boundary_check=(0,)).to(tl.float32) + b_Akk20 = b_Akk20 * b_b2[:, None] + b_Akk21 = b_Akk21 * b_b2[:, None] + if i_tc3 < T: + p_Aqk30 = tl.make_block_ptr( + Aqk, (T, BT), (HV * BT, 1), (i_tc3, 0), (BC, BC), (1, 0) + ) + p_Aqk31 = tl.make_block_ptr( + Aqk, (T, BT), (HV * BT, 1), (i_tc3, BC), (BC, BC), (1, 0) + ) + p_Aqk32 = tl.make_block_ptr( + Aqk, (T, BT), (HV * BT, 1), (i_tc3, 2 * BC), (BC, BC), (1, 0) + ) + tl.store( + p_Aqk30, (b_Aqk30 * scale).to(Aqk.dtype.element_ty), boundary_check=(0, 1) + ) + tl.store( + p_Aqk31, (b_Aqk31 * scale).to(Aqk.dtype.element_ty), boundary_check=(0, 1) + ) + tl.store( + p_Aqk32, (b_Aqk32 * scale).to(Aqk.dtype.element_ty), boundary_check=(0, 1) + ) + + p_b3 = tl.make_block_ptr( + beta + bos * HV + i_hv, (T,), (HV,), (i_tc3,), (BC,), (0,) + ) + b_b3 = tl.load(p_b3, boundary_check=(0,)).to(tl.float32) + b_Akk30 = b_Akk30 * b_b3[:, None] + b_Akk31 = b_Akk31 * b_b3[:, None] + b_Akk32 = b_Akk32 * b_b3[:, None] + + p_Akk00 = tl.make_block_ptr( + Akkd, (T, BC), (HV * BC, 1), (i_tc0, 0), (BC, BC), (1, 0) + ) + p_Akk11 = tl.make_block_ptr( + Akkd, (T, BC), (HV * BC, 1), (i_tc1, 0), (BC, BC), (1, 0) + ) + p_Akk22 = tl.make_block_ptr( + Akkd, (T, BC), (HV * BC, 1), (i_tc2, 0), (BC, BC), (1, 0) + ) + p_Akk33 = tl.make_block_ptr( + Akkd, (T, BC), (HV * BC, 1), (i_tc3, 0), (BC, BC), (1, 0) + ) + b_Ai00 = tl.load(p_Akk00, boundary_check=(0, 1)).to(tl.float32) + b_Ai11 = tl.load(p_Akk11, boundary_check=(0, 1)).to(tl.float32) + b_Ai22 = tl.load(p_Akk22, boundary_check=(0, 1)).to(tl.float32) + b_Ai33 = tl.load(p_Akk33, boundary_check=(0, 1)).to(tl.float32) + + ################################################################################ + # forward substitution on diagonals + ################################################################################ + + if not USE_SAFE_GATE: + m_A = o_i[:, None] > o_i[None, :] + m_I = o_i[:, None] == o_i[None, :] + + b_Ai00 = -tl.where(m_A, b_Ai00, 0) + b_Ai11 = -tl.where(m_A, b_Ai11, 0) + b_Ai22 = -tl.where(m_A, b_Ai22, 0) + b_Ai33 = -tl.where(m_A, b_Ai33, 0) + + for i in range(2, min(BC, T - i_tc0)): + b_a00 = -tl.load(Akkd + (i_tc0 + i) * HV * BC + o_i) + b_a00 = tl.where(o_i < i, b_a00, 0.0) + b_a00 += tl.sum(b_a00[:, None] * b_Ai00, 0) + b_Ai00 = tl.where((o_i == i)[:, None], b_a00, b_Ai00) + for i in range(BC + 2, min(2 * BC, T - i_tc0)): + b_a11 = -tl.load(Akkd + (i_tc0 + i) * HV * BC + o_i) + b_a11 = tl.where(o_i < i - BC, b_a11, 0.0) + b_a11 += tl.sum(b_a11[:, None] * b_Ai11, 0) + b_Ai11 = tl.where((o_i == i - BC)[:, None], b_a11, b_Ai11) + for i in range(2 * BC + 2, min(3 * BC, T - i_tc0)): + b_a22 = -tl.load(Akkd + (i_tc0 + i) * HV * BC + o_i) + b_a22 = tl.where(o_i < i - 2 * BC, b_a22, 0.0) + b_a22 += tl.sum(b_a22[:, None] * b_Ai22, 0) + b_Ai22 = tl.where((o_i == i - 2 * BC)[:, None], b_a22, b_Ai22) + for i in range(3 * BC + 2, min(4 * BC, T - i_tc0)): + b_a33 = -tl.load(Akkd + (i_tc0 + i) * HV * BC + o_i) + b_a33 = tl.where(o_i < i - 3 * BC, b_a33, 0.0) + b_a33 += tl.sum(b_a33[:, None] * b_Ai33, 0) + b_Ai33 = tl.where((o_i == i - 3 * BC)[:, None], b_a33, b_Ai33) + + b_Ai00 += m_I + b_Ai11 += m_I + b_Ai22 += m_I + b_Ai33 += m_I + + ################################################################################ + # compute merged inverse using off-diagonals + ################################################################################ + + # we used tf32 to maintain matrix inverse's precision whenever possible. + b_Ai10 = -tl.dot( + tl.dot(b_Ai11, b_Akk10, input_precision=SOLVE_TRIL_DOT_PRECISION), + b_Ai00, + input_precision=SOLVE_TRIL_DOT_PRECISION, + ) + b_Ai21 = -tl.dot( + tl.dot(b_Ai22, b_Akk21, input_precision=SOLVE_TRIL_DOT_PRECISION), + b_Ai11, + input_precision=SOLVE_TRIL_DOT_PRECISION, + ) + b_Ai32 = -tl.dot( + tl.dot(b_Ai33, b_Akk32, input_precision=SOLVE_TRIL_DOT_PRECISION), + b_Ai22, + input_precision=SOLVE_TRIL_DOT_PRECISION, + ) + + b_Ai20 = -tl.dot( + b_Ai22, + tl.dot(b_Akk20, b_Ai00, input_precision=SOLVE_TRIL_DOT_PRECISION) + + tl.dot(b_Akk21, b_Ai10, input_precision=SOLVE_TRIL_DOT_PRECISION), + input_precision=SOLVE_TRIL_DOT_PRECISION, + ) + b_Ai31 = -tl.dot( + b_Ai33, + tl.dot(b_Akk31, b_Ai11, input_precision=SOLVE_TRIL_DOT_PRECISION) + + tl.dot(b_Akk32, b_Ai21, input_precision=SOLVE_TRIL_DOT_PRECISION), + input_precision=SOLVE_TRIL_DOT_PRECISION, + ) + b_Ai30 = -tl.dot( + b_Ai33, + tl.dot(b_Akk30, b_Ai00, input_precision=SOLVE_TRIL_DOT_PRECISION) + + tl.dot(b_Akk31, b_Ai10, input_precision=SOLVE_TRIL_DOT_PRECISION) + + tl.dot(b_Akk32, b_Ai20, input_precision=SOLVE_TRIL_DOT_PRECISION), + input_precision=SOLVE_TRIL_DOT_PRECISION, + ) + + ################################################################################ + # store full Akk_inv to Akk + ################################################################################ + + p_Akk00 = tl.make_block_ptr( + Akk, (T, BT), (HV * BT, 1), (i_tc0, 0), (BC, BC), (1, 0) + ) + p_Akk10 = tl.make_block_ptr( + Akk, (T, BT), (HV * BT, 1), (i_tc1, 0), (BC, BC), (1, 0) + ) + p_Akk11 = tl.make_block_ptr( + Akk, (T, BT), (HV * BT, 1), (i_tc1, BC), (BC, BC), (1, 0) + ) + p_Akk20 = tl.make_block_ptr( + Akk, (T, BT), (HV * BT, 1), (i_tc2, 0), (BC, BC), (1, 0) + ) + p_Akk21 = tl.make_block_ptr( + Akk, (T, BT), (HV * BT, 1), (i_tc2, BC), (BC, BC), (1, 0) + ) + p_Akk22 = tl.make_block_ptr( + Akk, (T, BT), (HV * BT, 1), (i_tc2, 2 * BC), (BC, BC), (1, 0) + ) + p_Akk30 = tl.make_block_ptr( + Akk, (T, BT), (HV * BT, 1), (i_tc3, 0), (BC, BC), (1, 0) + ) + p_Akk31 = tl.make_block_ptr( + Akk, (T, BT), (HV * BT, 1), (i_tc3, BC), (BC, BC), (1, 0) + ) + p_Akk32 = tl.make_block_ptr( + Akk, (T, BT), (HV * BT, 1), (i_tc3, 2 * BC), (BC, BC), (1, 0) + ) + p_Akk33 = tl.make_block_ptr( + Akk, (T, BT), (HV * BT, 1), (i_tc3, 3 * BC), (BC, BC), (1, 0) + ) + + tl.store(p_Akk00, b_Ai00.to(Akk.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_Akk10, b_Ai10.to(Akk.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_Akk11, b_Ai11.to(Akk.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_Akk20, b_Ai20.to(Akk.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_Akk21, b_Ai21.to(Akk.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_Akk22, b_Ai22.to(Akk.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_Akk30, b_Ai30.to(Akk.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_Akk31, b_Ai31.to(Akk.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_Akk32, b_Ai32.to(Akk.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_Akk33, b_Ai33.to(Akk.dtype.element_ty), boundary_check=(0, 1)) + + +@triton.heuristics( + { + "IS_VARLEN": lambda args: args["cu_seqlens"] is not None, + } +) +@triton.autotune( + configs=[ + triton.Config({}, num_warps=num_warps, num_stages=num_stages) + for num_warps in [1, 2, 4, 8] + for num_stages in [2, 3, 4] + ], + key=["BK", "NC", "BT", "HV"], +) +@triton.jit(do_not_specialize=["B", "T"]) +def chunk_kda_fwd_kernel_intra_sub_chunk( + q, + k, + g, + beta, + Aqk, + Akk, + scale, + cu_seqlens, + chunk_indices, + T, + H: tl.constexpr, + HV: tl.constexpr, + K: tl.constexpr, + BT: tl.constexpr, + BC: tl.constexpr, + BK: tl.constexpr, + IS_VARLEN: tl.constexpr, + USE_GATHER: tl.constexpr, +): + i_t, i_i, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_b, i_hv = i_bh // HV, i_bh % HV + i_h = i_hv // (HV // H) + + if IS_VARLEN: + i_n, i_t = ( + tl.load(chunk_indices + i_t * 2).to(tl.int32), + tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32), + ) + bos, eos = ( + tl.load(cu_seqlens + i_n).to(tl.int32), + tl.load(cu_seqlens + i_n + 1).to(tl.int32), + ) + T = eos - bos + else: + bos, eos = i_b * T, i_b * T + T + + i_ti = i_t * BT + i_i * BC + if i_ti >= T: + return + + o_c = i_ti + tl.arange(0, BC) + m_c = o_c < T + + q = q + (bos * H + i_h) * K + k = k + (bos * H + i_h) * K + g = g + (bos * HV + i_hv) * K + beta = beta + bos * HV + i_hv + Aqk = Aqk + (bos * HV + i_hv) * BT + Akk = Akk + (bos * HV + i_hv) * BC + + p_q = tl.make_block_ptr(q, (T, K), (H * K, 1), (i_ti, 0), (BC, BK), (1, 0)) + p_k = tl.make_block_ptr(k, (T, K), (H * K, 1), (i_ti, 0), (BC, BK), (1, 0)) + p_g = tl.make_block_ptr(g, (T, K), (HV * K, 1), (i_ti, 0), (BC, BK), (1, 0)) + + p_beta = tl.make_block_ptr(beta, (T,), (HV,), (i_ti,), (BC,), (0,)) + + b_q = tl.load(p_q, boundary_check=(0, 1)) + b_k = tl.load(p_k, boundary_check=(0, 1)) + b_g = tl.load(p_g, boundary_check=(0, 1)) + b_beta = tl.load(p_beta, boundary_check=(0,)).to(tl.float32) + + if USE_GATHER: + b_gn = gather( + b_g, tl.full([1, BK], min(BC // 2, T - i_ti - 1), dtype=tl.int16), axis=0 + ) + else: + # caculate offset + p_gn = g + (i_ti + min(BC // 2, T - i_ti - 1)) * HV * K + tl.arange(0, BK) + b_gn = tl.load(p_gn, mask=tl.arange(0, BK) < K, other=0.0) + b_gn = b_gn[None, :] + + # current block, keep numerical stability by subtracting the left boundary + # less than 85 to avoid overflow in exp2 + b_gm = (b_g - b_gn).to(tl.float32) + + b_gq = tl.where(m_c[:, None], exp2(b_gm), 0.0) + b_gk = tl.where(m_c[:, None], exp2(-b_gm), 0.0) + + b_kgt = tl.trans(b_k * b_gk) + + b_Aqk = tl.dot(b_q * b_gq, b_kgt) * scale + b_Akk = tl.dot(b_k * b_gq, b_kgt) * b_beta[:, None] + + o_i = tl.arange(0, BC) + m_Aqk = o_i[:, None] >= o_i[None, :] + m_Akk = o_i[:, None] > o_i[None, :] + m_I = o_i[:, None] == o_i[None, :] + + b_Aqk = tl.where(m_Aqk, b_Aqk, 0.0) + b_Akk = tl.where(m_Akk, b_Akk, 0.0) + + p_Aqk = tl.make_block_ptr( + Aqk, (T, BT), (HV * BT, 1), (i_ti, i_i * BC), (BC, BC), (1, 0) + ) + p_Akk = tl.make_block_ptr(Akk, (T, BC), (HV * BC, 1), (i_ti, 0), (BC, BC), (1, 0)) + tl.store(p_Aqk, b_Aqk.to(Aqk.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_Akk, b_Akk.to(Akk.dtype.element_ty), boundary_check=(0, 1)) + + tl.debug_barrier() + + ################################################################################ + # forward substitution + ################################################################################ + + b_Ai = -b_Akk + for i in range(2, min(BC, T - i_ti)): + b_a = -tl.load(Akk + (i_ti + i) * HV * BC + o_i) + b_a = tl.where(o_i < i, b_a, 0.0) + b_a += tl.sum(b_a[:, None] * b_Ai, 0) + b_Ai = tl.where((o_i == i)[:, None], b_a, b_Ai) + b_Ai += m_I + tl.store(p_Akk, b_Ai.to(Akk.dtype.element_ty), boundary_check=(0, 1)) + + +def chunk_kda_fwd_intra( + q: torch.Tensor, + k: torch.Tensor, + gk: torch.Tensor | None = None, + beta: torch.Tensor | None = None, + scale: float | None = None, + cu_seqlens: torch.LongTensor | None = None, + chunk_size: int = 64, + chunk_indices: torch.LongTensor | None = None, + safe_gate: bool = False, +): + B, T, H, K, HV = *k.shape, gk.shape[2] + BT = chunk_size + BC = 16 + if chunk_indices is None and cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) + NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) + NC = triton.cdiv(BT, BC) + + Aqk = torch.empty(B, T, HV, BT, device=k.device, dtype=k.dtype) + # Akk must be zero-initialized - kernel only writes lower triangular + Akk = torch.zeros(B, T, HV, BT, device=k.device, dtype=k.dtype) + # Separate fp32 buffer for diagonal 16x16 blocks (for precision in solve_tril) + Akkd = torch.empty(B, T, HV, BC, device=k.device, dtype=torch.float32) + + # Compute diagonal blocks into Akkd in fp32. + if safe_gate: + grid = (NT, NC, B * HV) + BK = triton.next_power_of_2(K) + chunk_kda_fwd_kernel_intra_sub_chunk[grid]( + q=q, + k=k, + g=gk, + beta=beta, + Aqk=Aqk, + Akk=Akkd, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + T=T, + H=H, + HV=HV, + K=K, + BT=BT, + BC=BC, + BK=BK, + USE_GATHER=is_gather_supported, + ) + else: + Aqk, Akkd = chunk_kda_fwd_intra_token_parallel( + q=q, + k=k, + gk=gk, + beta=beta, + Aqk=Aqk, + Akk=Akkd, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_size=BT, + sub_chunk_size=BC, + ) + + # Step 2: Fused inter + solve_tril (works for both fixed-len and varlen) + solve_tril_dot_precision = ( + "tf32" + if current_platform.is_cuda() and current_platform.has_device_capability(80) + else "ieee" + ) + grid = (NT, B * HV) + chunk_kda_fwd_kernel_inter_solve_fused[grid]( + q=q, + k=k, + g=gk, + beta=beta, + Aqk=Aqk, + Akkd=Akkd, + Akk=Akk, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + T=T, + H=H, + HV=HV, + K=K, + BT=BT, + BC=BC, + USE_SAFE_GATE=safe_gate, + SOLVE_TRIL_DOT_PRECISION=solve_tril_dot_precision, + ) + return Aqk, Akk diff --git a/aphrodite/models/kimi_k3/amd/ops/third_party/kda/chunk_intra_token_parallel.py b/aphrodite/models/kimi_k3/amd/ops/third_party/kda/chunk_intra_token_parallel.py new file mode 100644 index 0000000000..73c860ec45 --- /dev/null +++ b/aphrodite/models/kimi_k3/amd/ops/third_party/kda/chunk_intra_token_parallel.py @@ -0,0 +1,197 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Songlin Yang, Yu Zhang, Zhiyuan Li +# +# This file contains code copied from the flash-linear-attention project. +# The original source was licensed under the MIT license. +# Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li +# Forward-only adaptation of flash-linear-attention 0.5.0. +# ruff: noqa: E501 + +# Token-parallel implementation of KDA intra chunk kernel + +import torch + +from aphrodite.third_party.flash_linear_attention.ops.op import exp2 +from aphrodite.triton_utils import tl, triton + + +@triton.heuristics( + { + "IS_VARLEN": lambda args: args["cu_seqlens"] is not None, + } +) +@triton.autotune( + configs=[ + triton.Config({"BH": BH}, num_warps=num_warps) + for BH in [1, 2, 4, 8] + for num_warps in [1, 2, 4, 8] + ], + key=["K", "H", "HV"], +) +@triton.jit(do_not_specialize=["T", "N"]) +def chunk_kda_fwd_kernel_intra_token_parallel( + q, + k, + g, + beta, + Aqk, + Akk, + scale, + cu_seqlens, + N, + T, + H: tl.constexpr, + HV: tl.constexpr, + K: tl.constexpr, + BT: tl.constexpr, + BC: tl.constexpr, + BH: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_tg, i_hg = tl.program_id(0), tl.program_id(1) + + if IS_VARLEN: + i_n = 0 + left, right = 0, N + + # Unrolled binary search (max B=2^32) + # We can limit iterations based on expected max batch size if needed + # 20 iterations covers B=1M, usually enough + for _ in range(20): + if left < right: + mid = (left + right) // 2 + if i_tg < tl.load(cu_seqlens + mid + 1).to(tl.int32): + right = mid + else: + left = mid + 1 + i_n = left + + bos, eos = ( + tl.load(cu_seqlens + i_n).to(tl.int32), + tl.load(cu_seqlens + i_n + 1).to(tl.int32), + ) + T = eos - bos + i_t = i_tg - bos + else: + bos = (i_tg // T) * T + i_t = i_tg % T + + if i_t >= T: + return + + i_c = i_t // BT + i_s = (i_t % BT) // BC + i_tc = i_c * BT + i_ts = i_tc + i_s * BC + + G: tl.constexpr = HV // H + + q += bos * H * K + k += bos * H * K + g += bos * HV * K + Aqk += bos * HV * BT + Akk += bos * HV * BC + beta += bos * HV + + BK: tl.constexpr = triton.next_power_of_2(K) + o_hv = i_hg * BH + tl.arange(0, BH) + o_h = o_hv // G + o_k = tl.arange(0, BK) + m_hv = o_hv < HV + m_k = o_k < K + m_hk = m_hv[:, None] & m_k[None, :] + + # q/k: [B, T, H, K], manual load via mapped qk head index + p_qk = o_h[:, None] * K + o_k[None, :] + b_q = tl.load(q + i_t * H * K + p_qk, mask=m_hk, other=0).to(tl.float32) + b_k = tl.load(k + i_t * H * K + p_qk, mask=m_hk, other=0).to(tl.float32) + + # g: [B, T, HV, K], beta: [B, T, HV] + p_g = tl.make_block_ptr( + g + i_t * HV * K, (HV, K), (K, 1), (i_hg * BH, 0), (BH, BK), (1, 0) + ) + p_beta = tl.make_block_ptr(beta + i_t * HV, (HV,), (1,), (i_hg * BH,), (BH,), (0,)) + b_g = tl.load(p_g, boundary_check=(0, 1)).to(tl.float32) + b_beta = tl.load(p_beta, boundary_check=(0,)).to(tl.float32) + b_k *= b_beta[:, None] + + for j in range(i_ts, min(i_t + 1, min(T, i_ts + BC))): + b_kj = tl.load(k + j * H * K + p_qk, mask=m_hk, other=0).to(tl.float32) + p_gj = tl.make_block_ptr( + g + j * HV * K, (HV, K), (K, 1), (i_hg * BH, 0), (BH, BK), (1, 0) + ) + b_gj = tl.load(p_gj, boundary_check=(0, 1)).to(tl.float32) + + b_kgj = tl.where(m_k[None, :], b_kj * exp2(b_g - b_gj), 0.0) + b_Aqk = tl.sum(b_q * b_kgj, axis=1) * scale + b_Akk = tl.sum(b_k * b_kgj, axis=1) * tl.where(j < i_t, 1.0, 0.0) + + tl.store( + Aqk + i_t * HV * BT + o_hv * BT + j % BT, + b_Aqk.to(Aqk.dtype.element_ty), + mask=m_hv, + ) + tl.store( + Akk + i_t * HV * BC + o_hv * BC + j - i_ts, + b_Akk.to(Akk.dtype.element_ty), + mask=m_hv, + ) + + +def chunk_kda_fwd_intra_token_parallel( + q: torch.Tensor, + k: torch.Tensor, + gk: torch.Tensor, + beta: torch.Tensor, + Aqk: torch.Tensor, + Akk: torch.Tensor, + scale: float, + cu_seqlens: torch.LongTensor | None = None, + chunk_size: int = 64, + sub_chunk_size: int = 16, +) -> None: + """ + Token-parallel implementation: each token gets its own thread block. + Supports both fixed-length and variable-length sequences. + Reduces wasted computation on padding. + + Writes directly to Aqk and Akk tensors (in-place). + + Args: + q: [B, T, H, K] + k: [B, T, H, K] + gk: [B, T, HV, K] cumsum of gates (HV >= H for GVA) + beta: [B, T, HV] + Aqk: [B, T, HV, BT] output tensor to write to + Akk: [B, T, HV, BC] output tensor for diagonal blocks (fp32) + scale: attention scale + chunk_size: BT (default 64) + sub_chunk_size: BC (default 16) + """ + B, T, H, K, HV = *q.shape, gk.shape[2] + N = len(cu_seqlens) - 1 if cu_seqlens is not None else B + BT = chunk_size + BC = sub_chunk_size + + def grid(meta): + return (B * T, triton.cdiv(HV, meta["BH"])) + + chunk_kda_fwd_kernel_intra_token_parallel[grid]( + q=q, + k=k, + g=gk, + beta=beta, + Aqk=Aqk, + Akk=Akk, + scale=scale, + cu_seqlens=cu_seqlens, + N=N, + T=T, + H=H, + HV=HV, + K=K, + BT=BT, + BC=BC, + ) + return Aqk, Akk diff --git a/aphrodite/models/kimi_k3/amd/ops/third_party/kda/fused_recurrent.py b/aphrodite/models/kimi_k3/amd/ops/third_party/kda/fused_recurrent.py new file mode 100644 index 0000000000..560a06f4f2 --- /dev/null +++ b/aphrodite/models/kimi_k3/amd/ops/third_party/kda/fused_recurrent.py @@ -0,0 +1,621 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Songlin Yang, Yu Zhang, Zhiyuan Li +# +# This file contains code adapted from the flash-linear-attention project. +# The original source was licensed under the MIT license. +# Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li +# ruff: noqa: E501 + +import torch + +from aphrodite.third_party.flash_linear_attention.ops.op import exp, log +from aphrodite.triton_utils import tl, triton +from aphrodite.utils.math_utils import cdiv, next_power_of_2 + + +@triton.heuristics( + { + "HAS_DT_BIAS": lambda args: args["dt_bias"] is not None, + "USE_LOWER_BOUND": lambda args: args["lower_bound"] is not None, + } +) +@triton.jit +def _kda_gate_beta_fwd_kernel( + raw_g, + raw_beta, + A_log, + dt_bias, + gate, + beta_out, + lower_bound, + softplus_beta: tl.constexpr, + softplus_threshold: tl.constexpr, + T, + stride_g_token: tl.constexpr, + stride_beta_token: tl.constexpr, + H: tl.constexpr, + D: tl.constexpr, + BT: tl.constexpr, + BD: tl.constexpr, + HAS_DT_BIAS: tl.constexpr, + USE_LOWER_BOUND: tl.constexpr, +): + i_t, i_h = tl.program_id(0), tl.program_id(1) + o_t = i_t * BT + tl.arange(0, BT) + o_d = tl.arange(0, BD) + m_t = o_t < T + m_d = o_d < D + + p_g = raw_g + o_t[:, None] * stride_g_token + i_h * D + o_d[None, :] + b_g = tl.load(p_g, mask=m_t[:, None] & m_d[None, :], other=0.0).to(tl.float32) + if HAS_DT_BIAS: + b_bias = tl.load( + dt_bias + i_h * D + o_d, + mask=m_d, + other=0.0, + ).to(tl.float32) + b_g += b_bias[None, :] + + b_a = exp(tl.load(A_log + i_h).to(tl.float32)) + if USE_LOWER_BOUND: + b_gate = lower_bound * tl.sigmoid(b_a * b_g) + else: + b_scaled = b_g * softplus_beta + b_softplus = tl.where( + b_scaled > softplus_threshold, + b_g, + log(1.0 + tl.exp(b_scaled)) / softplus_beta, + ) + b_gate = -b_a * b_softplus + + p_gate = gate + (o_t[:, None] * H + i_h) * D + o_d[None, :] + tl.store( + p_gate, + b_gate, + mask=m_t[:, None] & m_d[None, :], + ) + + b_beta = tl.load( + raw_beta + o_t * stride_beta_token + i_h, + mask=m_t, + other=0.0, + ).to(tl.float32) + tl.store(beta_out + o_t * H + i_h, tl.sigmoid(b_beta), mask=m_t) + + +def _fused_kda_gate_beta( + raw_g: torch.Tensor, + raw_beta: torch.Tensor, + A_log: torch.Tensor, + dt_bias: torch.Tensor | None, + lower_bound: float | None, +) -> tuple[torch.Tensor, torch.Tensor]: + B, T, H, D = raw_g.shape + assert B == 1 + assert raw_beta.shape == (B, T, H) + assert raw_g.stride()[2:] == (D, 1) + assert raw_beta.stride(2) == 1 + gate = torch.empty((B, T, H, D), dtype=torch.float32, device=raw_g.device) + beta = torch.empty((B, T, H), dtype=torch.float32, device=raw_beta.device) + + BT = 16 + _kda_gate_beta_fwd_kernel[(cdiv(T, BT), H)]( + raw_g=raw_g, + raw_beta=raw_beta, + A_log=A_log, + dt_bias=dt_bias, + gate=gate, + beta_out=beta, + lower_bound=lower_bound, + softplus_beta=1.0, + softplus_threshold=20.0, + T=T, + stride_g_token=raw_g.stride(1), + stride_beta_token=raw_beta.stride(1), + H=H, + D=D, + BT=BT, + BD=next_power_of_2(D), + num_warps=4, + ) + return gate, beta + + +@triton.heuristics( + { + "IS_SPEC_DECODING": lambda args: args["num_accepted_tokens"] is not None, + "HAS_DT_BIAS": lambda args: args["dt_bias"] is not None, + "USE_LOWER_BOUND": lambda args: args["lower_bound"] is not None, + } +) +@triton.jit(do_not_specialize=["N", "T"]) +def fused_recurrent_kda_fwd_kernel( + q, + k, + v, + g, + beta, + A_log, + dt_bias, + out, + state, + cu_seqlens, + state_indices, + num_accepted_tokens, + lower_bound, + scale: tl.constexpr, + N: tl.int64, + T: tl.int64, + H: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + stride_qkv_token: tl.constexpr, + stride_g_token: tl.constexpr, + stride_beta_token: tl.constexpr, + stride_out_token: tl.constexpr, + stride_state_token: tl.constexpr, + stride_indices_seq: tl.constexpr, + IS_SPEC_DECODING: tl.constexpr, + USE_QK_L2NORM_IN_KERNEL: tl.constexpr, + USE_GATE_IN_KERNEL: tl.constexpr, + APPLY_BETA_SIGMOID: tl.constexpr, + HAS_DT_BIAS: tl.constexpr, + USE_LOWER_BOUND: tl.constexpr, + num_stages: tl.constexpr, +): + pid = tl.program_id(0) + i_v = pid % tl.cdiv(V, BV) + i_nh = pid // tl.cdiv(V, BV) + i_n, i_h = i_nh // H, i_nh % H + bos = tl.load(cu_seqlens + i_n).to(tl.int64) + eos = tl.load(cu_seqlens + i_n + 1).to(tl.int64) + sequence_length = eos - bos + if sequence_length == 0: + return + + o_k = tl.arange(0, BK) + o_v = i_v * BV + tl.arange(0, BV) + m_k = o_k < K + m_v = o_v < V + m_state = m_v[:, None] & m_k[None, :] + + if IS_SPEC_DECODING: + initial_token = tl.load(num_accepted_tokens + i_n).to(tl.int64) - 1 + else: + initial_token = 0 + state_index = tl.load(state_indices + i_n * stride_indices_seq + initial_token).to( + tl.int64 + ) + p_out = out + bos * stride_out_token + i_h * V + o_v + if state_index <= 0: + tl.store(p_out, tl.zeros([BV], dtype=tl.float32), mask=m_v) + return + + p_state = ( + state + + state_index * stride_state_token + + i_h * V * K + + o_v[:, None] * K + + o_k[None, :] + ) + b_state = tl.load(p_state, mask=m_state, other=0.0).to(tl.float32) + + p_q = q + bos * stride_qkv_token + i_h * K + o_k + p_k = k + bos * stride_qkv_token + i_h * K + o_k + p_v = v + bos * stride_qkv_token + i_h * V + o_v + p_g = g + bos * stride_g_token + i_h * K + o_k + p_beta = beta + bos * stride_beta_token + i_h + for i_t in tl.range(0, sequence_length, num_stages=num_stages): + b_q = tl.load(p_q, mask=m_k, other=0.0, eviction_policy="evict_last").to( + tl.float32 + ) + b_k = tl.load(p_k, mask=m_k, other=0.0, eviction_policy="evict_last").to( + tl.float32 + ) + b_v = tl.load(p_v, mask=m_v, other=0.0, eviction_policy="evict_first").to( + tl.float32 + ) + if USE_QK_L2NORM_IN_KERNEL: + b_q = b_q / tl.sqrt(tl.sum(b_q * b_q) + 1e-6) + b_k = b_k / tl.sqrt(tl.sum(b_k * b_k) + 1e-6) + b_q *= scale + + b_gate = tl.load( + p_g, + mask=m_k, + other=0.0, + eviction_policy="evict_last", + ).to(tl.float32) + if USE_GATE_IN_KERNEL: + if HAS_DT_BIAS: + b_bias = tl.load( + dt_bias + i_h * K + o_k, + mask=m_k, + other=0.0, + ).to(tl.float32) + b_gate += b_bias + b_a = exp(tl.load(A_log + i_h).to(tl.float32)) + if USE_LOWER_BOUND: + b_gate = lower_bound * tl.sigmoid(b_a * b_gate) + else: + b_softplus = tl.where( + b_gate > 20.0, + b_gate, + log(1.0 + tl.exp(b_gate)), + ) + b_gate = -b_a * b_softplus + + b_state *= exp(b_gate[None, :]) + b_v -= tl.sum(b_state * b_k[None, :], axis=1) + b_beta = tl.load(p_beta, eviction_policy="evict_last").to(tl.float32) + if APPLY_BETA_SIGMOID: + b_beta = tl.sigmoid(b_beta) + b_v *= b_beta + b_state += b_v[:, None] * b_k[None, :] + b_out = tl.sum(b_state * b_q[None, :], axis=1) + tl.store( + p_out, + b_out.to(p_out.dtype.element_ty), + mask=m_v, + eviction_policy="evict_first", + ) + + final_state_index = tl.load(state_indices + i_n * stride_indices_seq + i_t).to( + tl.int64 + ) + if final_state_index > 0: + p_final_state = ( + state + + final_state_index * stride_state_token + + i_h * V * K + + o_v[:, None] * K + + o_k[None, :] + ) + tl.store( + p_final_state, + b_state.to(p_final_state.dtype.element_ty), + mask=m_state, + ) + + p_q += stride_qkv_token + p_k += stride_qkv_token + p_v += stride_qkv_token + p_g += stride_g_token + p_beta += stride_beta_token + p_out += stride_out_token + + +def fused_recurrent_kda_fwd( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + scale: float | None = None, + initial_state: torch.Tensor | None = None, + inplace_final_state: bool = True, + cu_seqlens: torch.Tensor | None = None, + ssm_state_indices: torch.Tensor | None = None, + num_accepted_tokens: torch.Tensor | None = None, + use_qk_l2norm_in_kernel: bool = True, + A_log: torch.Tensor | None = None, + dt_bias: torch.Tensor | None = None, + lower_bound: float | None = None, + use_gate_in_kernel: bool = False, + use_beta_sigmoid_in_kernel: bool = False, + out: torch.Tensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor]: + """Launch recurrent KDA with dense inner dimensions and row strides.""" + B, T, H, K = q.shape + V = v.shape[-1] + assert B == 1 and k.shape == q.shape + assert v.shape == (B, T, H, V) and g.shape == (B, T, H, K) + assert beta.shape == (B, T, H) + assert initial_state is not None + assert cu_seqlens is not None + assert ssm_state_indices is not None + assert inplace_final_state + if out is None: + out = torch.empty_like(v) + assert out.shape == v.shape + assert initial_state.shape[1:] == (H, V, K) + assert ssm_state_indices.ndim in (1, 2) + + assert q.stride()[2:] == k.stride()[2:] == (K, 1) + assert v.stride()[2:] == out.stride()[2:] == (V, 1) + assert g.stride()[2:] == (K, 1) + assert beta.stride(2) == 1 + assert q.stride(1) == k.stride(1) == v.stride(1) + assert initial_state.stride()[1:] == (V * K, K, 1) + N = cu_seqlens.numel() - 1 + if ssm_state_indices.ndim == 1: + assert T == N + assert num_accepted_tokens is None + else: + assert ssm_state_indices.stride(1) == 1 + assert cu_seqlens.is_contiguous() + if use_gate_in_kernel: + assert A_log is not None and A_log.is_contiguous() + assert dt_bias is None or dt_bias.is_contiguous() + + if scale is None: + scale = K**-0.5 + + BV = 32 if use_gate_in_kernel else 8 + num_warps = 4 if use_gate_in_kernel else 1 + grid = (cdiv(V, BV) * N * H,) + fused_recurrent_kda_fwd_kernel[grid]( + q=q, + k=k, + v=v, + g=g, + beta=beta, + A_log=A_log, + dt_bias=dt_bias, + out=out, + state=initial_state, + cu_seqlens=cu_seqlens, + state_indices=ssm_state_indices, + num_accepted_tokens=num_accepted_tokens, + lower_bound=lower_bound, + scale=scale, + N=N, + T=T, + H=H, + K=K, + V=V, + BK=next_power_of_2(K), + BV=BV, + stride_qkv_token=q.stride(1), + stride_g_token=g.stride(1), + stride_beta_token=beta.stride(1), + stride_out_token=out.stride(1), + stride_state_token=initial_state.stride(0), + stride_indices_seq=ssm_state_indices.stride(0), + IS_SPEC_DECODING=num_accepted_tokens is not None, + USE_QK_L2NORM_IN_KERNEL=use_qk_l2norm_in_kernel, + USE_GATE_IN_KERNEL=use_gate_in_kernel, + APPLY_BETA_SIGMOID=use_beta_sigmoid_in_kernel, + num_warps=num_warps, + num_stages=2, + ) + return out, initial_state + + +def fused_recurrent_kda( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + raw_g: torch.Tensor, + raw_beta: torch.Tensor, + A_log: torch.Tensor, + dt_bias: torch.Tensor | None, + lower_bound: float | None, + initial_state: torch.Tensor, + cu_seqlens: torch.Tensor, + ssm_state_indices: torch.Tensor, + num_accepted_tokens: torch.Tensor | None = None, + out: torch.Tensor | None = None, + fuse_gate: bool | None = None, +) -> tuple[torch.Tensor, torch.Tensor]: + """Run recurrent KDA from raw gate and beta inputs. + + This vLLM wrapper applies the gate activation and beta sigmoid, selecting + whether to materialize them before launching the recurrent kernel. + """ + if fuse_gate is None: + # gfx950: always fuse the gate. + fuse_gate = True + + if fuse_gate: + gate = raw_g + beta = raw_beta + else: + gate, beta = _fused_kda_gate_beta( + raw_g, + raw_beta, + A_log, + dt_bias, + lower_bound, + ) + return fused_recurrent_kda_fwd( + q=q, + k=k, + v=v, + g=gate, + beta=beta, + scale=q.shape[-1] ** -0.5, + initial_state=initial_state, + inplace_final_state=True, + cu_seqlens=cu_seqlens, + ssm_state_indices=ssm_state_indices, + num_accepted_tokens=num_accepted_tokens, + use_qk_l2norm_in_kernel=True, + A_log=A_log if fuse_gate else None, + dt_bias=dt_bias if fuse_gate else None, + lower_bound=lower_bound if fuse_gate else None, + use_gate_in_kernel=fuse_gate, + use_beta_sigmoid_in_kernel=fuse_gate, + out=out, + ) + + +@triton.jit +def fused_recurrent_kda_packed_decode_kernel( + mixed_qkv, + raw_g, + raw_beta, + A_log, + dt_bias, + out, + state, + state_indices, + lower_bound, + scale: tl.constexpr, + stride_mixed_token: tl.constexpr, + stride_g_token: tl.constexpr, + stride_beta_token: tl.constexpr, + stride_state_token: tl.constexpr, + H: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + SOFTPLUS_THRESHOLD: tl.constexpr, + USE_LOWER_BOUND: tl.constexpr, +): + i_v, i_nh = tl.program_id(0), tl.program_id(1) + i_n, i_h = i_nh // H, i_nh % H + + o_k = tl.arange(0, BK) + o_v = i_v * BV + tl.arange(0, BV) + mask_k = o_k < K + mask_v = o_v < V + mask_state = mask_v[:, None] & mask_k[None, :] + + state_idx = tl.load(state_indices + i_n).to(tl.int64) + p_out = out + (i_n * H + i_h) * V + o_v + if state_idx <= 0: + tl.store(p_out, tl.zeros([BV], dtype=tl.float32), mask=mask_v) + return + + p_state = state + state_idx * stride_state_token + p_state += i_h * V * K + o_v[:, None] * K + o_k[None, :] + b_state = tl.load(p_state, mask=mask_state, other=0).to(tl.float32) + + # Q, K, and V occupy consecutive channel ranges, while the token stride + # may also include the output-gate projection that follows packed QKV. + p_mixed = mixed_qkv + i_n * stride_mixed_token + b_q = tl.load(p_mixed + i_h * K + o_k, mask=mask_k, other=0).to(tl.float32) + b_k = tl.load( + p_mixed + H * K + i_h * K + o_k, + mask=mask_k, + other=0, + ).to(tl.float32) + b_v = tl.load( + p_mixed + 2 * H * K + i_h * V + o_v, + mask=mask_v, + other=0, + ).to(tl.float32) + + b_q /= tl.sqrt(tl.sum(b_q * b_q) + 1e-6) + b_k /= tl.sqrt(tl.sum(b_k * b_k) + 1e-6) + b_q *= scale + + p_g = raw_g + i_n * stride_g_token + i_h * K + o_k + b_g = tl.load(p_g, mask=mask_k, other=0).to(tl.float32) + b_bias = tl.load(dt_bias + i_h * K + o_k, mask=mask_k, other=0).to(tl.float32) + b_a = exp(tl.load(A_log + i_h).to(tl.float32)) + b_g += b_bias + if USE_LOWER_BOUND: + b_gate = lower_bound * tl.sigmoid(b_a * b_g) + else: + b_softplus = tl.where( + b_g > SOFTPLUS_THRESHOLD, + b_g, + log(1.0 + tl.exp(b_g)), + ) + b_gate = -b_a * b_softplus + + b_state *= exp(b_gate[None, :]) + b_v -= tl.sum(b_state * b_k[None, :], axis=1) + b_beta = tl.sigmoid( + tl.load(raw_beta + i_n * stride_beta_token + i_h).to(tl.float32) + ) + b_v *= b_beta + b_state += b_v[:, None] * b_k[None, :] + b_out = tl.sum(b_state * b_q[None, :], axis=1) + + tl.store(p_out, b_out.to(p_out.dtype.element_ty), mask=mask_v) + tl.store(p_state, b_state.to(p_state.dtype.element_ty), mask=mask_state) + + +def fused_recurrent_kda_packed_decode( + mixed_qkv: torch.Tensor, + raw_g: torch.Tensor, + raw_beta: torch.Tensor, + A_log: torch.Tensor, + dt_bias: torch.Tensor, + lower_bound: float | None, + initial_state: torch.Tensor, + state_indices: torch.Tensor, + scale: float | None = None, +) -> tuple[torch.Tensor, torch.Tensor]: + """Run one-token KDA decode directly from packed post-conv QKV.""" + if mixed_qkv.ndim != 2 or mixed_qkv.stride(-1) != 1: + raise ValueError("`mixed_qkv` must be 2D and contiguous in its last dim.") + if raw_g.ndim != 4 or raw_g.shape[0] != 1: + raise ValueError("`raw_g` must have shape [1, B, H, K].") + if raw_beta.ndim != 3 or raw_beta.shape[0] != 1: + raise ValueError("`raw_beta` must have shape [1, B, H].") + if initial_state.ndim != 4: + raise ValueError("`initial_state` must have shape [cache, H, V, K].") + _, H, V, K = initial_state.shape + if raw_g.stride()[2:] != (K, 1): + raise ValueError("`raw_g` must be contiguous within each token.") + if raw_beta.stride(2) != 1: + raise ValueError("`raw_beta` heads must be contiguous.") + if initial_state.stride()[1:] != (V * K, K, 1): + raise ValueError("`initial_state` must be contiguous within each cache slot.") + if state_indices.ndim != 1 or state_indices.stride(0) != 1: + raise ValueError("`state_indices` must be contiguous and one-dimensional.") + if A_log.ndim != 1 or not A_log.is_contiguous(): + raise ValueError("`A_log` must be contiguous and one-dimensional.") + if not dt_bias.is_contiguous(): + raise ValueError("`dt_bias` must be contiguous.") + + device = mixed_qkv.device + if any( + x.device != device + for x in (raw_g, raw_beta, A_log, dt_bias, initial_state, state_indices) + ): + raise ValueError("All packed KDA inputs must be on the same device.") + + B = mixed_qkv.shape[0] + if raw_g.shape != (1, B, H, K): + raise ValueError(f"Unexpected raw gate shape {tuple(raw_g.shape)}.") + if raw_beta.shape != (1, B, H): + raise ValueError(f"Unexpected raw beta shape {tuple(raw_beta.shape)}.") + if mixed_qkv.shape[1] != 2 * H * K + H * V: + raise ValueError(f"Unexpected packed QKV shape {tuple(mixed_qkv.shape)}.") + if A_log.numel() != H or dt_bias.numel() != H * K: + raise ValueError("`A_log` or `dt_bias` has an incompatible shape.") + if state_indices.shape[0] != B: + raise ValueError("`state_indices` must contain one entry per token.") + + BK = next_power_of_2(K) + BV = min(next_power_of_2(V), 32) + if scale is None: + scale = K**-0.5 + + out = torch.empty((1, B, H, V), dtype=mixed_qkv.dtype, device=device) + grid = (cdiv(V, BV), B * H) + fused_recurrent_kda_packed_decode_kernel[grid]( + mixed_qkv=mixed_qkv, + raw_g=raw_g, + raw_beta=raw_beta, + A_log=A_log, + dt_bias=dt_bias, + out=out, + state=initial_state, + state_indices=state_indices, + lower_bound=lower_bound or 0.0, + scale=scale, + stride_mixed_token=mixed_qkv.stride(0), + stride_g_token=raw_g.stride(1), + stride_beta_token=raw_beta.stride(1), + stride_state_token=initial_state.stride(0), + H=H, + K=K, + V=V, + BK=BK, + BV=BV, + SOFTPLUS_THRESHOLD=20.0, + USE_LOWER_BOUND=lower_bound is not None, + num_warps=4, + num_stages=2, + ) + return out, initial_state diff --git a/aphrodite/models/kimi_k3/common/__init__.py b/aphrodite/models/kimi_k3/common/__init__.py new file mode 100644 index 0000000000..208f01a7cb --- /dev/null +++ b/aphrodite/models/kimi_k3/common/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project diff --git a/aphrodite/models/kimi_k3/common/mm_preprocess.py b/aphrodite/models/kimi_k3/common/mm_preprocess.py new file mode 100644 index 0000000000..04ce1a2c42 --- /dev/null +++ b/aphrodite/models/kimi_k3/common/mm_preprocess.py @@ -0,0 +1,320 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Shared Kimi-K3 multimodal preprocessing.""" + +import math +from collections.abc import Mapping, Sequence +from typing import Any, cast + +import torch +from transformers import BatchFeature + +from aphrodite.config.multimodal import BaseDummyOptions, ImageDummyOptions +from aphrodite.inputs import MultiModalDataDict +from aphrodite.logger import init_logger +from aphrodite.multimodal.inputs import ( + MultiModalFieldConfig, + MultiModalKwargsItems, +) +from aphrodite.multimodal.parse import ImageProcessorItems, ImageSize, MultiModalDataItems +from aphrodite.multimodal.processing import ( + BaseDummyInputsBuilder, + BaseMultiModalProcessor, + BaseProcessingInfo, + InputProcessingContext, + PromptReplacement, + PromptUpdate, + PromptUpdateDetails, +) +from aphrodite.transformers_utils.configs.kimi_k3 import KimiK3Config +from aphrodite.transformers_utils.processor import cached_get_image_processor +from aphrodite.transformers_utils.processors.kimi_k3 import KimiK3Processor + +logger = init_logger(__name__) + + +def navit_resize_image( + width: int, + height: int, + patch_size: int, + merge_kernel_size: int, + in_patch_limit: int, + patch_limit_on_one_side: int, + fixed_output_tokens: int | None, +): + # Apply the patch limits. + s1 = math.sqrt(in_patch_limit / (max(1.0, width // patch_size) * max(1.0, height // patch_size))) + s2 = patch_limit_on_one_side * patch_size / width + s3 = patch_limit_on_one_side * patch_size / height + scale = min(1.0, s1, s2, s3) + new_w, new_h = max(1, int(width * scale)), max(1, int(height * scale)) + new_w = min(new_w, patch_limit_on_one_side * patch_size) + new_h = min(new_h, patch_limit_on_one_side * patch_size) + + factor = merge_kernel_size * patch_size + + pad_height = (factor - new_h % factor) % factor + pad_width = (factor - new_w % factor) % factor + + if fixed_output_tokens is not None: + num_tokens = fixed_output_tokens + else: + # Calculate new dimensions after padding and patching + token_height = (new_h + pad_height) // factor + token_width = (new_w + pad_width) // factor + + assert token_height * merge_kernel_size <= patch_limit_on_one_side, ( + f"token_height {token_height} * merge_kernel_size {merge_kernel_size} > " + f"patch_limit_on_one_side {patch_limit_on_one_side}" + ) + assert token_width * merge_kernel_size <= patch_limit_on_one_side, ( + f"token_width {token_width} * merge_kernel_size {merge_kernel_size} > " + f"patch_limit_on_one_side {patch_limit_on_one_side}" + ) + + num_tokens = token_height * token_width + return { + "num_tokens": num_tokens, + "new_width": new_w, + "new_height": new_h, + "pad_width": pad_width, + "pad_height": pad_height, + "sampled_nframes": 1, + } + + +class KimiK3ProcessingInfo(BaseProcessingInfo): + """Processing information for the image-only Kimi-K3 model. + + K3 uses the standard ``image`` modality (unlike K2.5's unified + ``vision_chunk``), so it builds its own ``KimiK3Processor`` wrapper around + the checkpoint's image processor and resolves the ``<|media_pad|>`` token + id the same way K2.5 does. + """ + + def __init__(self, ctx: InputProcessingContext) -> None: + super().__init__(ctx) + + self.hf_config = hf_config = self.get_hf_config() + + tokenizer = self.get_tokenizer() + image_processor = cached_get_image_processor( + self.ctx.model_config.model, + revision=self.ctx.model_config.revision, + trust_remote_code=self.ctx.model_config.trust_remote_code, + ) + + # Resolve token ID from the tokenizer because transformers v5 + # may remap token IDs vs config.json. + config_token_id = hf_config.media_placeholder_token_id + resolved_token_id = tokenizer.convert_tokens_to_ids("<|media_pad|>") + unk_token_id = getattr(tokenizer, "unk_token_id", None) + is_valid_resolved = isinstance(resolved_token_id, int) and ( + unk_token_id is None or resolved_token_id != unk_token_id + ) + if is_valid_resolved and resolved_token_id != config_token_id: + logger.warning_once( + "Kimi-K3 config.media_placeholder_token_id (%d) disagrees " + "with tokenizer mapping for <|media_pad|> (%d). " + "Using tokenizer value.", + config_token_id, + resolved_token_id, + ) + media_token_id = resolved_token_id + # Patch config so downstream code also sees the correct ID. + hf_config.media_placeholder_token_id = resolved_token_id + else: + media_token_id = config_token_id + + self.media_token_id = media_token_id + self.media_token = tokenizer.decode(media_token_id) + + self.image_processor = image_processor + self.hf_processor = KimiK3Processor( + tokenizer=tokenizer, + image_processor=image_processor, + ) + self.media_tokens_calculator = image_processor.media_tokens_calculator + + def get_hf_processor(self, **kwargs: object) -> KimiK3Processor: + return self.hf_processor + + def get_hf_config(self) -> KimiK3Config: + return self.ctx.get_hf_config(KimiK3Config) + + def get_supported_mm_limits(self) -> Mapping[str, int | None]: + # None means unlimited + return {"image": None} + + @classmethod + def get_max_image_size( + cls, + patch_size: int, + merge_kernel_size: int, + in_patch_limit: int, + patch_limit_on_one_side: int, + fixed_output_tokens: int | None, + ) -> ImageSize: + max_side = patch_limit_on_one_side * patch_size + best_score = (-1, -1) + best_size = (max_side, max_side) + + for width_patches in range(patch_limit_on_one_side + 1): + width = min((width_patches + 1) * patch_size - 1, max_side) + for height_patches in range(width_patches, patch_limit_on_one_side + 1): + height = min((height_patches + 1) * patch_size - 1, max_side) + resize_config = navit_resize_image( + width, + height, + patch_size, + merge_kernel_size, + in_patch_limit, + patch_limit_on_one_side, + fixed_output_tokens, + ) + padded_width = resize_config["new_width"] + resize_config["pad_width"] + padded_height = resize_config["new_height"] + resize_config["pad_height"] + num_patches = padded_width // patch_size * (padded_height // patch_size) + score = (resize_config["num_tokens"], num_patches) + if score > best_score: + best_score = score + best_size = (width, height) + return ImageSize(width=best_size[0], height=best_size[1]) + + +class KimiK3DummyInputsBuilder(BaseDummyInputsBuilder[KimiK3ProcessingInfo]): + """Builds image-based dummy inputs for K3 profiling. + + The dummy text is made of ``<|kimi_image_placeholder|>`` tokens — exactly + the placeholder that K3's ``_get_prompt_updates`` expands — and the dummy + mm data is a plain list of PIL images under the ``image`` key. + """ + + def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str: + num_images = mm_counts.get("image", 0) + return self.info.get_hf_config().image_placeholder * num_images + + def get_dummy_mm_data( + self, + seq_len: int, + mm_counts: Mapping[str, int], + mm_options: Mapping[str, BaseDummyOptions] | None = None, + ) -> MultiModalDataDict: + media_proc_cfg = self.info.image_processor.media_proc_cfg + max_size = self.info.get_max_image_size( + media_proc_cfg["patch_size"], + media_proc_cfg["merge_kernel_size"], + media_proc_cfg["in_patch_limit"], + media_proc_cfg["patch_limit_on_one_side"], + media_proc_cfg["fixed_output_tokens"], + ) + num_images = mm_counts.get("image", 0) + image_overrides = cast( + ImageDummyOptions | None, + mm_options.get("image") if mm_options else None, + ) + return { + "image": self._get_dummy_images( + width=max_size.width, + height=max_size.height, + num_images=num_images, + overrides=image_overrides, + ) + } + + +class KimiK3MultiModalProcessor(BaseMultiModalProcessor[KimiK3ProcessingInfo]): + """Image-only multi-modal processor for Kimi-K3.""" + + def _call_hf_processor( + self, + prompt: str, + mm_data: Mapping[str, object], + mm_kwargs: Mapping[str, object], + tok_kwargs: Mapping[str, object], + ) -> BatchFeature: + # Override so the base always routes through the text+mm path + # (`KimiK3Processor.__call__`). Otherwise the mm-only fast path calls + # the checkpoint image processor directly with bare PIL images, but it + # requires `{"type": "image", "image": PIL}` media dicts that only our + # wrapper builds. + return super()._call_hf_processor(prompt, mm_data, mm_kwargs, tok_kwargs) + + def _hf_processor_applies_updates( + self, + prompt_text: str, + mm_items: MultiModalDataItems, + hf_processor_mm_kwargs: Mapping[str, object], + tokenization_kwargs: Mapping[str, object], + ) -> bool: + return False + + def _get_mm_fields_config( + self, + hf_inputs: BatchFeature, + hf_processor_mm_kwargs: Mapping[str, object], + ) -> Mapping[str, MultiModalFieldConfig]: + """Slice the flattened patch tensor back into per-image items. + + ``pixel_values`` holds all patches from every image concatenated; each + image's patch count is ``prod(grid_thws[i])``. ``grid_thws`` is one + ``[N_t, N_h, N_w]`` row per image. + """ + grid_thws = hf_inputs.get("grid_thws", torch.empty((0, 3))) + grid_sizes = grid_thws.prod(-1) + + return dict( + pixel_values=MultiModalFieldConfig.flat_from_sizes("image", grid_sizes), + grid_thws=MultiModalFieldConfig.batched("image", keep_on_cpu=True), + ) + + def _get_prompt_updates( + self, + mm_items: MultiModalDataItems, + hf_processor_mm_kwargs: Mapping[str, Any], + out_mm_kwargs: MultiModalKwargsItems, + ) -> Sequence[PromptUpdate]: + """Expand each K3 image placeholder into a resolution-aware update. + + K3's prompt carries a single ``<|kimi_image_placeholder|>`` token per + image. This replaces that token with + ``<|media_begin|>image {w}x{h}<|media_content|>{pads}<|media_end|>``, + embedding the per-image resolution in the prompt and marking only the + ``<|media_pad|>`` positions as embedding slots (the number of pads is + the feature size returned by ``media_tokens_calculator``). + """ + media_token_id = self.info.media_token_id + media_token = self.info.media_token + image_placeholder = self.info.get_hf_config().image_placeholder + + def get_replacement(item_idx: int) -> PromptUpdateDetails[str]: + images = mm_items.get_items("image", ImageProcessorItems) + image = images.get(item_idx) + if image is None: + raise ValueError(f"Missing image data at index {item_idx}") + + # The checkpoint image processor works on media dicts, so wrap the + # PIL image before asking it for the token count. + num_media_token = self.info.media_tokens_calculator({"type": "image", "image": image}) + pads = media_token * num_media_token + + # NOTE: `width`/`height` are the ORIGINAL upload dimensions, not the + # post-preprocess (smart-resized) ones. `image` comes from the + # untouched parsed `mm_items`; the checkpoint image processor + # (`KimiK3VisionProcessor.preprocess`) only produces new tensors via + # `image.resize(...)` and never mutates the stored PIL. This matches + # the reference HF processor (`KimiK3Processor.preprocess_medias`), + # which also builds the prompt from the original `img.size`. The + # resize is reflected only in the pad count above. + width, height = images.get_image_size(item_idx) + full = f"<|media_begin|>image {width}x{height}<|media_content|>{pads}<|media_end|>" + + return PromptUpdateDetails.select_token_id(full, media_token_id) + + return [ + PromptReplacement( + modality="image", + target=image_placeholder, + replacement=get_replacement, + ), + ] diff --git a/aphrodite/models/kimi_k3/common/mtp.py b/aphrodite/models/kimi_k3/common/mtp.py new file mode 100644 index 0000000000..e22612abf5 --- /dev/null +++ b/aphrodite/models/kimi_k3/common/mtp.py @@ -0,0 +1,94 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Fused Kimi-K3 MTP input preparation.""" + +import torch + +from aphrodite.triton_utils import tl, triton + + +@triton.jit +def _rms_norm(x, weight, eps, hidden_size: tl.constexpr): + x = x.to(tl.float32) + variance = tl.sum(x * x, axis=0) / hidden_size + return x * tl.rsqrt(variance + eps) * weight.to(tl.float32) + + +@triton.jit +def _fused_mtp_input_kernel( + positions_ptr, + inputs_embeds_ptr, + previous_hidden_states_ptr, + enorm_weight_ptr, + hnorm_weight_ptr, + output_ptr, + eps, + inputs_embeds_stride, + previous_hidden_states_stride, + output_stride, + hidden_size: tl.constexpr, + block_size: tl.constexpr, +): + token_idx = tl.program_id(0).to(tl.int64) + input_idx = tl.program_id(1) + offsets = tl.arange(0, block_size) + mask = offsets < hidden_size + + if input_idx == 0: + position = tl.load(positions_ptr + token_idx) + values = tl.load( + inputs_embeds_ptr + token_idx * inputs_embeds_stride + offsets, + mask=mask & (position != 0), + other=0.0, + ) + weight = tl.load(enorm_weight_ptr + offsets, mask=mask, other=0.0) + else: + values = tl.load( + previous_hidden_states_ptr + token_idx * previous_hidden_states_stride + offsets, + mask=mask, + other=0.0, + ) + weight = tl.load(hnorm_weight_ptr + offsets, mask=mask, other=0.0) + + output = _rms_norm(values, weight, eps, hidden_size) + tl.store( + output_ptr + token_idx * output_stride + input_idx * hidden_size + offsets, + output, + mask=mask, + ) + + +def fused_mtp_input( + positions: torch.Tensor, + inputs_embeds: torch.Tensor, + previous_hidden_states: torch.Tensor, + enorm_weight: torch.Tensor, + hnorm_weight: torch.Tensor, + eps: float, +) -> torch.Tensor: + """Mask and normalize both MTP inputs into the projection layout.""" + num_tokens, hidden_size = inputs_embeds.shape + output = torch.empty( + num_tokens, + 2 * hidden_size, + dtype=inputs_embeds.dtype, + device=inputs_embeds.device, + ) + if num_tokens == 0: + return output + + _fused_mtp_input_kernel[(num_tokens, 2)]( + positions, + inputs_embeds, + previous_hidden_states, + enorm_weight, + hnorm_weight, + output, + eps, + inputs_embeds.stride(0), + previous_hidden_states.stride(0), + output.stride(0), + hidden_size, + triton.next_power_of_2(hidden_size), + ) + return output diff --git a/aphrodite/models/kimi_k3/nvidia/dspark_mla.py b/aphrodite/models/kimi_k3/nvidia/dspark_mla.py new file mode 100644 index 0000000000..b54eb6bbe1 --- /dev/null +++ b/aphrodite/models/kimi_k3/nvidia/dspark_mla.py @@ -0,0 +1,495 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""K3 dense MLA draft model for DSpark speculative decoding.""" + +from collections.abc import Iterable + +import torch +import torch.nn as nn +import torch.nn.functional as F + +import aphrodite._custom_ops as ops +from aphrodite.config import AphroditeConfig +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.models.qwen3_dspark import DSparkMarkovHead +from aphrodite.model_executor.models.utils import ( + AutoWeightsLoader, + WeightsMapper, + get_draft_quant_config, + maybe_prefix, +) +from aphrodite.models.kimi_k3.nvidia.mla import MultiHeadLatentAttention +from aphrodite.models.kimi_k3.nvidia.model import KimiMLP +from aphrodite.utils.torch_utils import is_quantized_kv_cache + + +class ReplicatedDSparkMarkovHead(DSparkMarkovHead): + """DSpark Markov head with full weights on every TP rank.""" + + def __init__( + self, + vocab_size: int, + draft_vocab_size: int, + markov_rank: int, + prefix: str, + ) -> None: + # TODO: Remove this mypy workaround once the K3 PR is fully merged. + super().__init__( # type: ignore[call-arg] + vocab_size, + draft_vocab_size, + markov_rank, + prefix, + replicated=True, + ) + + +class K3DSparkDecoderLayer(nn.Module): + def __init__( + self, + *, + aphrodite_config: AphroditeConfig, + config, + layer_idx: int, + start_layer_id: int, + prefix: str, + ) -> None: + super().__init__() + quant_config = get_draft_quant_config(aphrodite_config) + self.self_attn = MultiHeadLatentAttention( + config=config, + hidden_size=config.hidden_size, + num_heads=config.num_attention_heads, + qk_nope_head_dim=config.qk_nope_head_dim, + qk_rope_head_dim=config.qk_rope_head_dim, + v_head_dim=config.v_head_dim, + q_lora_rank=config.q_lora_rank, + kv_lora_rank=config.kv_lora_rank, + cache_config=aphrodite_config.cache_config, + quant_config=quant_config, + prefix=maybe_prefix(prefix, f"layers.{start_layer_id + layer_idx}.self_attn"), + use_rope=True, + non_causal_multi_token_decode=True, + ) + self.mlp = KimiMLP( + hidden_size=config.hidden_size, + intermediate_size=config.intermediate_size, + hidden_act=config.hidden_act, + quant_config=quant_config, + prefix=maybe_prefix(prefix, f"layers.{start_layer_id + layer_idx}.mlp"), + ) + self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.post_attention_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + residual: torch.Tensor | None, + ) -> tuple[torch.Tensor, torch.Tensor]: + if residual is None: + residual = hidden_states + hidden_states = self.input_layernorm(hidden_states) + else: + hidden_states, residual = self.input_layernorm(hidden_states, residual) + + hidden_states = self.self_attn( + positions=positions, + hidden_states=hidden_states, + ) + hidden_states, residual = self.post_attention_layernorm(hidden_states, residual) + hidden_states = self.mlp(hidden_states) + return hidden_states, residual + + +class K3DSparkModel(nn.Module): + def __init__( + self, + *, + aphrodite_config: AphroditeConfig, + start_layer_id: int, + prefix: str, + ) -> None: + super().__init__() + assert aphrodite_config.speculative_config is not None + self.config = aphrodite_config.speculative_config.draft_model_config.hf_config + self.quant_config = get_draft_quant_config(aphrodite_config) + + # The frozen target embedding is aliased after the draft checkpoint loads. + self.embed_tokens: nn.Module | None = None + + self.context_proj = ReplicatedLinear( + self.config.target_hidden_size * self.config.num_target_layers, + self.config.hidden_size, + bias=False, + return_bias=False, + quant_config=self.quant_config, + prefix=maybe_prefix(prefix, "context_proj"), + ) + self.context_norm = RMSNorm(self.config.hidden_size, eps=self.config.rms_norm_eps) + + self.layers = nn.ModuleList( + [ + K3DSparkDecoderLayer( + aphrodite_config=aphrodite_config, + config=self.config, + layer_idx=layer_idx, + start_layer_id=start_layer_id, + prefix=prefix, + ) + for layer_idx in range(self.config.num_hidden_layers) + ] + ) + self.final_norm = RMSNorm(self.config.hidden_size, eps=self.config.rms_norm_eps) + self.markov_head = ReplicatedDSparkMarkovHead( + self.config.vocab_size, + self.config.draft_vocab_size, + self.config.markov_rank, + prefix=maybe_prefix(prefix, "markov_head"), + ) + self._context_kv_fusion_available: bool | None = None + self._max_num_context_tokens = aphrodite_config.scheduler_config.max_num_batched_tokens + + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: + assert self.embed_tokens is not None + return self.embed_tokens(input_ids) + + def combine_hidden_states(self, hidden_states: torch.Tensor) -> torch.Tensor: + return self.context_norm(self.context_proj(hidden_states)) + + @torch.inference_mode() + def precompute_and_store_context_kv( + self, + context_states: torch.Tensor, + context_positions: torch.Tensor, + context_slot_mapping: torch.Tensor | list[torch.Tensor | None] | None = None, + ) -> None: + """Project target-derived context into each draft layer's latent cache.""" + if self._context_kv_fusion_available is None: + self._build_fused_context_kv_buffers() + if self._context_kv_fusion_available: + self._precompute_fused_context_kv(context_states, context_positions, context_slot_mapping) + return + + # Quantized fallback. Directly invoking the projection modules preserves + # their quantization methods, at the cost of also computing unused Q rows. + for layer_idx, layer in enumerate(self.layers): + attn = layer.self_attn + assert attn.fused_qkv_a_proj is not None + assert attn.q_lora_rank is not None + assert attn.rotary_emb is not None + qkv_lora = attn.fused_qkv_a_proj(context_states)[0] + kv_lora = qkv_lora[..., attn.q_lora_rank :] + kv_c, k_pe = kv_lora.split([attn.kv_lora_rank, attn.qk_rope_head_dim], dim=-1) + kv_c = attn.kv_a_layernorm(kv_c) + k_pe = k_pe.unsqueeze(1) + # DeepSeek YaRN's FlashInfer path requires paired Q/K tensors. + # The vLLM CUDA op supports rotating one tensor in place and + # consumes the same (possibly scaled fp32) cos/sin cache. + rotary_emb = attn.rotary_emb + ops.rotary_embedding( + context_positions, + k_pe, + None, + rotary_emb.head_size, + rotary_emb.cos_sin_cache, + rotary_emb.is_neox_style, + ) + + slot_mapping = ( + context_slot_mapping[layer_idx] + if isinstance(context_slot_mapping, (list, tuple)) + else context_slot_mapping + ) + if slot_mapping is None: + continue + attn.impl.do_kv_cache_update( + kv_c, + k_pe, + attn.kv_cache, + slot_mapping, + attn.kv_cache_dtype, + attn._k_scale, + ) + + def _build_fused_context_kv_buffers(self) -> None: + """Build a cross-layer KV-only A projection after checkpoint loading.""" + if self.quant_config is not None: + self._context_kv_fusion_available = False + return + + attentions = [layer.self_attn for layer in self.layers] + if not attentions or any( + attn.fused_qkv_a_proj is None or not hasattr(attn.fused_qkv_a_proj, "weight") for attn in attentions + ): + self._context_kv_fusion_available = False + return + + attn0 = attentions[0] + assert attn0.q_lora_rank is not None + kv_width = attn0.kv_lora_rank + attn0.qk_rope_head_dim + kv_weights = [] + for attn in attentions: + assert attn.q_lora_rank is not None + assert ( + attn.q_lora_rank == attn0.q_lora_rank + and attn.kv_lora_rank == attn0.kv_lora_rank + and attn.qk_rope_head_dim == attn0.qk_rope_head_dim + and attn.kv_a_layernorm.variance_epsilon == attn0.kv_a_layernorm.variance_epsilon + ), "All MLA DSpark layers must share their latent KV geometry." + kv_weights.append(attn.fused_qkv_a_proj.weight.detach().narrow(0, attn.q_lora_rank, kv_width)) + + # [L * (kv_lora_rank + rope_dim), hidden_size]. The underlying fused + # A weights are replicated (`disable_tp=True`), so this is valid on + # every TP rank without communication. + self._fused_context_kv_weight = torch.cat(kv_weights, dim=0) + self._context_kv_norm_weights = torch.stack( + [attn.kv_a_layernorm.weight.detach() for attn in attentions], dim=0 + ).contiguous() + self._num_context_layers = len(attentions) + self._context_kv_width = kv_width + self._context_kv_lora_rank = attn0.kv_lora_rank + self._context_rope_dim = attn0.qk_rope_head_dim + self._context_rms_norm_eps = attn0.kv_a_layernorm.variance_epsilon + self._context_positions_repeated = torch.empty( + self._num_context_layers * self._max_num_context_tokens, + dtype=torch.int64, + device=self._fused_context_kv_weight.device, + ) + self._context_kv_fusion_available = True + + def _precompute_fused_context_kv( + self, + context_states: torch.Tensor, + context_positions: torch.Tensor, + context_slot_mapping: torch.Tensor | list[torch.Tensor | None] | None, + ) -> None: + num_ctx = context_states.shape[0] + num_layers = self._num_context_layers + + # One KV-only GEMM replaces five full Q+KV GEMMs. For K3 this projects + # 5*576 rows rather than 5*2112 rows (72.7% fewer A-projection FLOPs). + all_kv = F.linear(context_states, self._fused_context_kv_weight) + all_kv = all_kv.view(num_ctx, num_layers, self._context_kv_width) + all_kv_c = all_kv[..., : self._context_kv_lora_rank] + all_k_pe = all_kv[..., self._context_kv_lora_rank :] + + # Layer-major layout lets the 2-D RMSNorm weights select a distinct row + # for each draft layer in one grouped kernel. + all_kv_c = all_kv_c.permute(1, 0, 2).contiguous() + all_kv_c_normed = torch.empty_like(all_kv_c) + ops.rms_norm( + all_kv_c_normed, + all_kv_c, + self._context_kv_norm_weights, + self._context_rms_norm_eps, + ) + + all_k_pe = all_k_pe.permute(1, 0, 2).contiguous() + all_k_pe_flat = all_k_pe.view(num_layers * num_ctx, 1, self._context_rope_dim) + repeated_positions = self._context_positions_repeated[: num_layers * num_ctx] + repeated_positions.view(num_layers, num_ctx).copy_(context_positions) + # Keep the single-tensor context RoPE on vLLM's optimized CUDA op; + # DeepSeek YaRN's FlashInfer wrapper assumes a non-null key tensor. + rotary_emb = self.layers[0].self_attn.rotary_emb + assert rotary_emb is not None + ops.rotary_embedding( + repeated_positions, + all_k_pe_flat, + None, + rotary_emb.head_size, + rotary_emb.cos_sin_cache, + rotary_emb.is_neox_style, + ) + all_k_pe = all_k_pe_flat.view(num_layers, num_ctx, 1, self._context_rope_dim) + + if context_slot_mapping is None: + return + + cache_layers = [layer.self_attn for layer in self.layers] + if ( + not is_quantized_kv_cache(cache_layers[0].kv_cache_dtype) + and self._has_uniform_block_layout(cache_layers) + and (isinstance(context_slot_mapping, torch.Tensor) or all(s is not None for s in context_slot_mapping)) + ): + # Grouped context KV insert only supports unquantized (bf16) KV cache + # and assumes that all layers share the same block layout. + + if isinstance(context_slot_mapping, (list, tuple)): + per_layer_slot_mappings = [s for s in context_slot_mapping if s is not None] + if len({s.data_ptr() for s in per_layer_slot_mappings}) == 1: + # All rows alias to the same slot mapping. + slot_mapping = per_layer_slot_mappings[0].unsqueeze(0).expand(num_layers, -1) + else: + slot_mapping = torch.stack(per_layer_slot_mappings, dim=0) + else: + # Broadcast the single shared context_slot_mapping tensor. + slot_mapping = context_slot_mapping.unsqueeze(0).expand(num_layers, -1) + + ref_cache = cache_layers[0].kv_cache + ops.concat_and_cache_mla_grouped( + all_kv_c_normed, + all_k_pe.squeeze(2), + self._get_context_kv_cache_ptrs(cache_layers), + slot_mapping, + ref_cache.size(1), + ref_cache.stride(0), + ref_cache.stride(1), + ) + return + + for layer_idx, layer in enumerate(self.layers): + slot_mapping = ( + context_slot_mapping[layer_idx] + if isinstance(context_slot_mapping, (list, tuple)) + else context_slot_mapping + ) + if slot_mapping is None: + continue + attn = layer.self_attn + attn.impl.do_kv_cache_update( + all_kv_c_normed[layer_idx], + all_k_pe[layer_idx], + attn.kv_cache, + slot_mapping, + attn.kv_cache_dtype, + attn._k_scale, + ) + + def _has_uniform_block_layout( + self, + cache_layers: list[MultiHeadLatentAttention], + ) -> bool: + if not hasattr(self, "_layers_share_kv_block_layout"): + ref_cache = cache_layers[0].kv_cache + self._layers_share_kv_block_layout = all( + cl.kv_cache.size(1) == ref_cache.size(1) + and cl.kv_cache.stride(0) == ref_cache.stride(0) + and cl.kv_cache.stride(1) == ref_cache.stride(1) + for cl in cache_layers + ) + return self._layers_share_kv_block_layout + + def _get_context_kv_cache_ptrs( + self, + cache_layers: list[MultiHeadLatentAttention], + ) -> torch.Tensor: + # The per-layer KV cache base pointers are stable after allocation, so + # build the pointer array once and return it on every call. + if not hasattr(self, "_context_cache_ptrs"): + ref_cache = cache_layers[0].kv_cache + cache_ptrs = torch.tensor( + [cl.kv_cache.data_ptr() for cl in cache_layers], + dtype=torch.int64, + device=ref_cache.device, + ) + self._context_cache_ptrs = cache_ptrs + return self._context_cache_ptrs + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + inputs_embeds: torch.Tensor | None = None, + ) -> torch.Tensor: + if inputs_embeds is None: + inputs_embeds = self.embed_input_ids(input_ids) + + hidden_states = inputs_embeds + residual = None + for layer in self.layers: + hidden_states, residual = layer( + positions=positions, + hidden_states=hidden_states, + residual=residual, + ) + hidden_states, _ = self.final_norm(hidden_states, residual) + return hidden_states + + +class K3DSparkForCausalLM(nn.Module): + has_own_embed_tokens = False + has_own_lm_head = False + draft_id_to_target_id = None + checkpoint_skip_substrs = ("confidence_head", "embed_tokens", "lm_head") + + hf_to_vllm_mapper = WeightsMapper( + orig_to_new_prefix={"": "model."}, + orig_to_new_stacked={ + ".gate_proj": (".gate_up_proj", 0), + ".up_proj": (".gate_up_proj", 1), + ".q_a_proj": (".fused_qkv_a_proj", 0), + ".kv_a_proj_with_mqa": (".fused_qkv_a_proj", 1), + }, + ) + + 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 + target_layer_num = aphrodite_config.model_config.get_num_layers(aphrodite_config.parallel_config) + self.model = K3DSparkModel( + aphrodite_config=aphrodite_config, + start_layer_id=target_layer_num, + prefix=maybe_prefix(prefix, "model"), + ) + + # Assigned by load_dspark_model from the target. Keeping no placeholder + # avoids a transient full-vocabulary allocation for this 163k-vocab model. + self.lm_head: nn.Module | None = None + logit_scale = getattr(self.config, "logit_scale", 1.0) + self.logits_processor = LogitsProcessor(self.config.draft_vocab_size, scale=logit_scale) + + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: + return self.model.embed_input_ids(input_ids) + + def combine_hidden_states(self, hidden_states: torch.Tensor) -> torch.Tensor: + return self.model.combine_hidden_states(hidden_states) + + def get_draft_kv_cache_layer_names(self) -> list[str]: + return [layer.self_attn.layer_name for layer in self.model.layers] + + def precompute_and_store_context_kv( + self, + context_states: torch.Tensor, + context_positions: torch.Tensor, + context_slot_mapping: torch.Tensor | list[torch.Tensor | None] | None = None, + ) -> None: + self.model.precompute_and_store_context_kv(context_states, context_positions, context_slot_mapping) + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + inputs_embeds: torch.Tensor | None = None, + ) -> torch.Tensor: + return self.model(input_ids, positions, inputs_embeds) + + def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor: + assert self.lm_head is not None + return self.logits_processor(self.lm_head, hidden_states) + + def compute_draft_logits(self, hidden_states: torch.Tensor) -> torch.Tensor: + return self.compute_logits(hidden_states) + + def map_draft_to_target(self, draft_ids: torch.Tensor) -> torch.Tensor: + return draft_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) + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + # confidence_head is training-only. The frozen target embedding and LM + # head are shared after this draft-specific checkpoint is loaded. + loader = AutoWeightsLoader( + self, + skip_substrs=list(self.checkpoint_skip_substrs), + ) + loaded_weights = loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) + self.model._build_fused_context_kv_buffers() + return loaded_weights diff --git a/aphrodite/models/kimi_k3/nvidia/kda.py b/aphrodite/models/kimi_k3/nvidia/kda.py new file mode 100644 index 0000000000..27cd047674 --- /dev/null +++ b/aphrodite/models/kimi_k3/nvidia/kda.py @@ -0,0 +1,735 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from collections.abc import Callable + +import torch +from einops import rearrange +from torch import nn +from torch.nn.parameter import Parameter + +from aphrodite import _custom_ops as ops +from aphrodite.compilation.breakable_cudagraph import eager_break_during_capture +from aphrodite.config import AphroditeConfig +from aphrodite.distributed import divide, get_tensor_model_parallel_rank +from aphrodite.forward_context import get_forward_context +from aphrodite.logger import init_logger +from aphrodite.model_executor.layers.linear import ( + ColumnParallelLinear, + MergedColumnParallelLinear, + RowParallelLinear, +) +from aphrodite.model_executor.layers.mamba.gdn.base import GatedDeltaNetAttention +from aphrodite.model_executor.layers.mamba.mamba_utils import ( + MambaStateDtypeCalculator, + MambaStateShapeCalculator, + is_conv_state_dim_first, +) +from aphrodite.model_executor.layers.mamba.ops.causal_conv1d import ( + causal_conv1d_fn, + causal_conv1d_update, +) +from aphrodite.model_executor.layers.mamba.ops.gather_initial_states import ( + gather_initial_states, +) +from aphrodite.model_executor.model_loader.weight_utils import ( + default_weight_loader, + sharded_weight_loader, +) +from aphrodite.model_executor.parameter import BaseAphroditeParameter as BasevLLMParameter +from aphrodite.model_executor.utils import set_weight_attrs +from aphrodite.models.kimi_k3.nvidia.kda_metadata import ( + KimiK3KDAAttentionBackend, + KimiK3KDAMetadata, +) +from aphrodite.platforms import current_platform +from aphrodite.third_party.flash_linear_attention.ops.kda import FusedRMSNormGated +from aphrodite.transformers_utils.configs.kimi_linear import KimiLinearConfig +from aphrodite.v1.attention.backend import AttentionBackend + +logger = init_logger(__name__) + +_KDA_GATE_LOGBOUND_MIN = -5.0 + + +def a_log_weight_loader( + shard_axis: int, +) -> Callable[[torch.Tensor, torch.Tensor], None]: + """Load KDA A_log stored as either old 4D or current 1D weights.""" + + def loader(param: torch.Tensor, loaded_weight: torch.Tensor) -> None: + tp_rank = get_tensor_model_parallel_rank() + shard_size = param.data.shape[shard_axis] + start_idx = tp_rank * shard_size + + if loaded_weight.dim() == 4: + assert loaded_weight.shape[:2] == (1, 1), ( + f"Expected old A_log shape (1, 1, H, 1), got {loaded_weight.shape}" + ) + assert loaded_weight.shape[-1] == 1, f"Expected old A_log last dim to be 1, got {loaded_weight.shape}" + loaded_weight = loaded_weight.view(loaded_weight.shape[2]) + + loaded_weight = loaded_weight.narrow(shard_axis, start_idx, shard_size) + return default_weight_loader(param, loaded_weight) + + return loader + + +class _KimiGDNMergedColumnParallelLinear(MergedColumnParallelLinear): + """Merged projection with one output replicated across TP ranks.""" + + tp_rank: int + + def __init__( + self, + input_size: int, + output_sizes: list[int], + replicated_shard_id: int, + tp_size: int, + **kwargs, + ) -> None: + self.replicated_shard_id = replicated_shard_id + output_sizes = output_sizes.copy() + output_sizes[replicated_shard_id] *= tp_size + super().__init__(input_size, output_sizes, **kwargs) + + def weight_loader( + self, + param: Parameter, + loaded_weight: torch.Tensor, + loaded_shard_id: tuple[int, ...] | int | None = None, + ) -> None: + tp_rank = self.tp_rank + param_tp_rank = getattr(param, "tp_rank", None) + if loaded_shard_id == self.replicated_shard_id: + self.tp_rank = 0 + if param_tp_rank is not None: + param.tp_rank = 0 + try: + super().weight_loader(param, loaded_weight, loaded_shard_id) + finally: + self.tp_rank = tp_rank + if param_tp_rank is not None: + param.tp_rank = param_tp_rank + + def weight_loader_v2( + self, + param: BasevLLMParameter, + loaded_weight: torch.Tensor, + loaded_shard_id: tuple[int, ...] | int | None = None, + ) -> None: + tp_rank = self.tp_rank + param_tp_rank = getattr(param, "tp_rank", None) + if loaded_shard_id == self.replicated_shard_id: + self.tp_rank = 0 + if param_tp_rank is not None: + param.tp_rank = 0 + try: + super().weight_loader_v2(param, loaded_weight, loaded_shard_id) + finally: + self.tp_rank = tp_rank + if param_tp_rank is not None: + param.tp_rank = param_tp_rank + + +def is_fused_kda_decode_supported( + num_heads: int, + head_dim: int, + conv_width: int, + num_spec: int, + input_dtype: torch.dtype, + conv_state_dtype: torch.dtype, +) -> bool: + if ( + num_heads not in (12, 24, 48, 96) + or head_dim != 128 + or conv_width != 4 + or num_spec != 0 + or input_dtype != torch.bfloat16 + or conv_state_dtype != torch.bfloat16 + or is_conv_state_dim_first() + or not hasattr(torch.ops._C, "fused_kda_decode") + ): + return False + # SM90 is architecture-specific; SM10x and SM12x use family binaries. + return ( + current_platform.is_device_capability(90) + or current_platform.is_device_capability_family(100) + or current_platform.is_device_capability_family(120) + ) + + +def is_flashkda_supported( + head_dim: int, + dtype: torch.dtype, + lower_bound: float | None, +) -> bool: + capability = current_platform.get_device_capability() + return ( + capability is not None + and capability.major in (9, 10, 12) + and head_dim == 128 + and dtype == torch.bfloat16 + and lower_bound is not None + ) + + +def _flashkda_prefill( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + A_log: torch.Tensor, + dt_bias: torch.Tensor, + lower_bound: float, + initial_state: torch.Tensor, + cu_seqlens: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + import aphrodite._flashkda_C # noqa: F401 + + out = torch.empty(v.shape, dtype=v.dtype, device=v.device) + final_state = torch.empty_like(initial_state) + workspace = torch.empty( + torch.ops._flashkda_C.get_workspace_size( + q.shape[0] * q.shape[1], + q.shape[2], + cu_seqlens.numel() - 1, + ), + dtype=torch.uint8, + device=q.device, + ) + # FlashKDA hardcodes dense Q/K/V/G strides. Beta may be row-strided because + # FlashKDA materializes its transposed [H, T] layout internally. + # TODO: Teach FlashKDA to consume beta in [T, H] layout directly instead + # of transposing it to contiguous [H, T] storage internally. + torch.ops._flashkda_C.fwd( + q.contiguous(), + k.contiguous(), + v.contiguous(), + g.contiguous(), + beta, + q.shape[-1] ** -0.5, + out, + workspace, + A_log.contiguous(), + dt_bias.view(-1, q.shape[-1]).contiguous(), + lower_bound, + initial_state.contiguous(), + final_state, + cu_seqlens.contiguous(), + ) + return out, final_state + + +def resolve_kda_prefill_backend( + backend: str, + head_dim: int, + dtype: torch.dtype, + lower_bound: float | None, +) -> str: + if backend not in ("auto", "triton", "flashkda"): + raise ValueError(f"Unsupported KDA prefill backend: {backend}") + supported = is_flashkda_supported(head_dim, dtype, lower_bound) + if backend == "flashkda" and not supported: + raise RuntimeError("FlashKDA requires CUDA SM90/SM10x/SM12x, bfloat16, head_dim=128, and a bounded KDA gate.") + if supported and backend != "triton": + logger.info_once("Using FlashKDA KDA prefill backend.") + return "flashkda" + return "triton" + + +def _make_decode_conv1d_weight_loader( + dims: list[int], + tp_size: int, + tp_rank: int, + decode_conv1d_weight: torch.Tensor | None, +) -> Callable[..., None]: + sharded_dims = [dim // tp_size for dim in dims] + + def weight_loader( + param: torch.Tensor, + loaded_weight: torch.Tensor, + loaded_shard_id: int, + ) -> None: + if loaded_weight.dim() == 2: + loaded_weight = loaded_weight.unsqueeze(1) + shard_size = sharded_dims[loaded_shard_id] + source_start = tp_rank * shard_size + target_start = sum(sharded_dims[:loaded_shard_id]) + loaded_shard = loaded_weight[source_start : source_start + shard_size] + param.data[target_start : target_start + shard_size].copy_(loaded_shard) + if decode_conv1d_weight is not None and not param.is_meta: + decode_conv1d_weight[loaded_shard_id].copy_(loaded_shard.squeeze(1).transpose(0, 1)) + + return weight_loader + + +def _make_decode_norm_weight_loader( + decode_norm_weight: torch.Tensor, +) -> Callable[..., None]: + def weight_loader(param: torch.Tensor, loaded_weight: torch.Tensor) -> None: + default_weight_loader(param, loaded_weight) + if not param.is_meta: + decode_norm_weight.copy_(param.data) + + return weight_loader + + +class KimiK3DeltaAttention(GatedDeltaNetAttention): + def get_attn_backend(self) -> type[AttentionBackend]: + return KimiK3KDAAttentionBackend + + def get_state_dtype( + self, + ) -> tuple[torch.dtype, torch.dtype]: + if self.model_config is None or self.cache_config is None: + raise ValueError("model_config and cache_config must be set") + return MambaStateDtypeCalculator.kda_state_dtype(self.model_config.dtype, self.cache_config.mamba_cache_dtype) + + def get_state_shape( + self, + ) -> tuple[tuple[int, ...], tuple[int, ...]]: + return MambaStateShapeCalculator.kda_state_shape( + self.tp_size, + self.num_heads, + self.head_dim, + conv_kernel_size=self.conv_size, + num_spec=self.num_spec, + ) + + def __init__( + self, + config: KimiLinearConfig, + aphrodite_config: AphroditeConfig, + prefix: str = "", + ) -> None: + super().__init__(config, aphrodite_config, prefix) + + kda_config = config.linear_attn_config # type: ignore[attr-defined] + assert kda_config is not None, "linear_attn_config must be set" + self.head_dim = kda_config["head_dim"] + self.num_heads = kda_config["num_heads"] + assert self.num_heads % self.tp_size == 0 + self.local_num_heads = divide(self.num_heads, self.tp_size) + self.projection_size = self.head_dim * self.num_heads + self.local_projection_size = divide(self.projection_size, self.tp_size) + self.conv_size = kda_config["short_conv_kernel_size"] + assert kda_config.get("use_full_rank_gate", False), "KimiK3DeltaAttention requires a full-rank gate" + + # Keep f_a before the narrow beta shard, then pad each TP-local row + # to select the aligned BF16 GEMM path. + qkvg_output_sizes = [self.projection_size] * 4 + in_proj_output_sizes = qkvg_output_sizes + [ + self.head_dim, + self.num_heads, + ] + local_output_size = 4 * self.local_projection_size + self.head_dim + self.local_num_heads + self.in_proj_padding = -local_output_size % 16 + if self.in_proj_padding: + in_proj_output_sizes.append(self.in_proj_padding * self.tp_size) + self.in_proj_qkvgfab = _KimiGDNMergedColumnParallelLinear( + self.hidden_size, + in_proj_output_sizes, + replicated_shard_id=4, + tp_size=self.tp_size, + bias=False, + quant_config=self.quant_config, + prefix=f"{prefix}.in_proj_qkvgfab", + ) + if self.in_proj_padding: + self.in_proj_qkvgfab.weight.data[-self.in_proj_padding :].zero_() + + self.f_b_proj = ColumnParallelLinear( + self.head_dim, + self.projection_size, + bias=False, + quant_config=self.quant_config, + prefix=f"{prefix}.f_b_proj", + ) + self.dt_bias = nn.Parameter(torch.empty(self.local_projection_size, dtype=torch.float32)) + set_weight_attrs(self.dt_bias, {"weight_loader": sharded_weight_loader(0)}) + + # One packed parameter and cache let decode run a single conv update. + # Prefill slices them back into Q/K/V to obtain dense outputs cheaply. + self.conv1d = ColumnParallelLinear( + input_size=self.conv_size, + output_size=3 * self.projection_size, + bias=False, + params_dtype=torch.float32, + prefix=f"{prefix}.conv1d", + ) + self.conv1d.weight.data = self.conv1d.weight.data.unsqueeze(1) + # Keep a width-major copy for fused decode without changing the layout + # consumed by the prefill and fallback decode kernels. + conv_state_dtype, _ = self.get_state_dtype() + decode_conv1d_weight = None + if is_fused_kda_decode_supported( + self.local_num_heads, + self.head_dim, + self.conv_size, + self.num_spec, + aphrodite_config.model_config.dtype, + conv_state_dtype, + ): + logger.info_once("Fused KDA decode kernel (conv+KDA+norm) is enabled.") + decode_conv1d_weight = torch.empty( + 3, + self.conv_size, + self.local_projection_size, + dtype=self.conv1d.weight.dtype, + device=self.conv1d.weight.device, + ) + self.register_buffer("decode_conv1d_weight", decode_conv1d_weight, persistent=False) + delattr(self.conv1d.weight, "weight_loader") + set_weight_attrs( + self.conv1d.weight, + { + "weight_loader": _make_decode_conv1d_weight_loader( + [self.projection_size] * 3, + self.tp_size, + self.tp_rank, + decode_conv1d_weight, + ) + }, + ) + + self.A_log = nn.Parameter(torch.empty(self.local_num_heads, dtype=torch.float32)) + set_weight_attrs(self.A_log, {"weight_loader": a_log_weight_loader(0)}) + + self.gate_lower_bound: float | None = kda_config.get("gate_lower_bound", None) + if self.gate_lower_bound is not None: + assert _KDA_GATE_LOGBOUND_MIN <= self.gate_lower_bound < 0, ( + f"KDA gate lower bound must be in [{_KDA_GATE_LOGBOUND_MIN}, 0). Got {self.gate_lower_bound}." + ) + + additional_config = aphrodite_config.additional_config + backend = ( + additional_config.get("kda_prefill_backend", "auto") if isinstance(additional_config, dict) else "auto" + ) + self.kda_prefill_backend = resolve_kda_prefill_backend( + backend, + self.head_dim, + aphrodite_config.model_config.dtype, + self.gate_lower_bound, + ) + + self.o_norm = FusedRMSNormGated(self.head_dim, activation="sigmoid") + decode_norm_weight = None + if decode_conv1d_weight is not None: + decode_norm_weight = torch.empty( + self.head_dim, + dtype=torch.float32, + device=self.o_norm.weight.device, + ) + self.register_buffer("decode_norm_weight", decode_norm_weight, persistent=False) + if decode_norm_weight is not None: + # Upcast once while loading; direct BF16 norm weights slow the + # fully fused decode kernel. + if hasattr(self.o_norm.weight, "weight_loader"): + delattr(self.o_norm.weight, "weight_loader") + set_weight_attrs( + self.o_norm.weight, + {"weight_loader": _make_decode_norm_weight_loader(decode_norm_weight)}, + ) + self.o_proj = RowParallelLinear( + self.projection_size, + self.hidden_size, + bias=False, + quant_config=self.quant_config, + prefix=f"{prefix}.o_proj", + ) + + compilation_config = aphrodite_config.compilation_config + if prefix in compilation_config.static_forward_context: + raise ValueError(f"Duplicate layer name: {prefix}") + compilation_config.static_forward_context[prefix] = self + + def forward( + self, + hidden_states: torch.Tensor, + positions: torch.Tensor, + ) -> torch.Tensor: + num_tokens = hidden_states.size(0) + projected_qkvgfab = self.in_proj_qkvgfab(hidden_states)[0] + split_sizes = [ + 3 * self.local_projection_size, + self.local_projection_size, + self.head_dim, + self.local_num_heads, + ] + if self.in_proj_padding: + split_sizes.append(self.in_proj_padding) + projected = projected_qkvgfab.split(split_sizes, dim=-1) + mixed_qkv, g_proj_states, f_a, beta = projected[:4] + + g1 = self.f_b_proj(f_a)[0] + beta = beta.unsqueeze(0) + g1 = rearrange(g1, "n (h d) -> 1 n h d", d=self.head_dim) + g2 = rearrange(g_proj_states, "... (h d) -> ... h d", d=self.head_dim) + core_attn_out = torch.empty( + (1, num_tokens, self.local_num_heads, self.head_dim), + dtype=hidden_states.dtype, + device=hidden_states.device, + ) + self._forward( + mixed_qkv=mixed_qkv, + g1=g1, + g2=g2, + beta=beta, + core_attn_out=core_attn_out, + ) + core_attn_out = rearrange(core_attn_out, "1 n h d -> n (h d)") + return self.o_proj(core_attn_out)[0] + + @eager_break_during_capture + def _forward( + self, + mixed_qkv: torch.Tensor, + g1: torch.Tensor, + g2: torch.Tensor, + beta: torch.Tensor, + core_attn_out: torch.Tensor, + ) -> None: + forward_context = get_forward_context() + attn_metadata_raw = forward_context.attn_metadata + if attn_metadata_raw is None: + return + + from aphrodite.models.kimi_k3.nvidia.ops.third_party.kda import ( + chunk_kda_with_fused_gate, + fused_recurrent_kda, + fused_recurrent_kda_packed_decode, + ) + + assert isinstance(attn_metadata_raw, dict) + attn_metadata_narrowed = attn_metadata_raw[self.prefix] + assert isinstance(attn_metadata_narrowed, KimiK3KDAMetadata) + m = attn_metadata_narrowed + has_initial_state = m.has_initial_state + non_spec_query_start_loc = m.non_spec_query_start_loc + non_spec_state_indices_tensor = m.non_spec_state_indices_tensor + spec_token_indx = m.spec_token_indx + non_spec_token_indx = m.non_spec_token_indx + spec_state_indices_tensor = m.spec_state_indices_tensor + spec_query_start_loc = m.spec_query_start_loc + num_accepted_tokens = m.num_accepted_tokens + num_actual_tokens = m.num_actual_tokens + has_spec_decode = m.num_spec_decodes > 0 + mixed_qkv = mixed_qkv[:num_actual_tokens] + g1 = g1[:, :num_actual_tokens] + beta = beta[:, :num_actual_tokens] + + conv_state, recurrent_state = self.kv_cache + # The convolution kernels consume (..., dim, width - 1). + if not is_conv_state_dim_first(): + conv_state = conv_state.transpose(-1, -2) + + if ( + self.decode_conv1d_weight is not None + and self.decode_norm_weight is not None + and not has_spec_decode + and m.num_prefills == 0 + and m.num_decodes > 0 + ): + assert non_spec_state_indices_tensor is not None + ops.fused_kda_decode( + x=mixed_qkv, + weight=self.decode_conv1d_weight, + bias=self.conv1d.bias, + conv_state=conv_state, + raw_g=g1, + raw_beta=beta, + A_log=self.A_log, + dt_bias=self.dt_bias, + state_indices=non_spec_state_indices_tensor[:num_actual_tokens], + state=recurrent_state, + out=core_attn_out[:, :num_actual_tokens], + lower_bound=self.gate_lower_bound, + output_gate=g2[:num_actual_tokens], + norm_weight=self.decode_norm_weight, + norm_eps=self.o_norm.eps, + ) + return + + conv_weights = self.conv1d.weight.view(self.conv1d.weight.size(0), self.conv1d.weight.size(2)) + q_conv_weight, k_conv_weight, v_conv_weight = conv_weights.split(self.local_projection_size, dim=0) + q_conv_state, k_conv_state, v_conv_state = conv_state.split(self.local_projection_size, dim=-2) + + # Separate multi-query speculative tokens from prefill/plain decode. + if has_spec_decode: + if m.num_prefills == 0 and m.num_decodes == 0: + mixed_qkv_spec = mixed_qkv + g1_spec, beta_spec = g1, beta + mixed_qkv_ns = g1_ns = beta_ns = None + else: + assert spec_token_indx is not None + assert non_spec_token_indx is not None + mixed_qkv_spec = mixed_qkv.index_select(0, spec_token_indx) + g1_spec = g1.index_select(1, spec_token_indx) + beta_spec = beta.index_select(1, spec_token_indx) + mixed_qkv_ns = mixed_qkv.index_select(0, non_spec_token_indx) + g1_ns = g1.index_select(1, non_spec_token_indx) + beta_ns = beta.index_select(1, non_spec_token_indx) + else: + mixed_qkv_spec = g1_spec = beta_spec = None + mixed_qkv_ns, g1_ns, beta_ns = mixed_qkv, g1, beta + + # Spec-decode multi-query path. + core_attn_out_spec = None + if has_spec_decode: + assert spec_state_indices_tensor is not None + assert spec_query_start_loc is not None + spec_conv_indices = spec_state_indices_tensor[:, 0][: m.num_spec_decodes] + spec_max_query_len = spec_state_indices_tensor.size(-1) + spec_conv_out = torch.empty_like(mixed_qkv_spec) + mixed_qkv_spec = causal_conv1d_update( + mixed_qkv_spec, + conv_state, + conv_weights, + self.conv1d.bias, + activation="silu", + conv_state_indices=spec_conv_indices, + num_accepted_tokens=num_accepted_tokens, + query_start_loc=spec_query_start_loc, + max_query_len=spec_max_query_len, + validate_data=False, + out=spec_conv_out, + ) + q_spec, k_spec, v_spec = ( + rearrange(x, "n (h d) -> 1 n h d", d=self.head_dim) + for x in mixed_qkv_spec.split(self.local_projection_size, dim=-1) + ) + spec_cu_seqlens = spec_query_start_loc[: m.num_spec_decodes + 1] + spec_out = core_attn_out[:, : q_spec.shape[1]] if m.num_prefills == 0 and m.num_decodes == 0 else None + core_attn_out_spec, _ = fused_recurrent_kda( + q=q_spec, + k=k_spec, + v=v_spec, + raw_g=g1_spec, + raw_beta=beta_spec, + A_log=self.A_log, + dt_bias=self.dt_bias, + lower_bound=self.gate_lower_bound, + initial_state=recurrent_state, + cu_seqlens=spec_cu_seqlens, + ssm_state_indices=spec_state_indices_tensor, + num_accepted_tokens=num_accepted_tokens, + out=spec_out, + ) + + # Prefill or plain-decode path. + core_attn_out_non_spec = None + if mixed_qkv_ns is not None: + assert g1_ns is not None and beta_ns is not None + if m.num_prefills > 0: + q_ns, k_ns, v_ns = mixed_qkv_ns.split(self.local_projection_size, dim=-1) + + # Separate convolution calls accept row-strided packed inputs + # and produce dense Q/K/V without an additional V copy. + def _prefill_conv( + x: torch.Tensor, + state: torch.Tensor, + weight: torch.Tensor, + ) -> torch.Tensor: + return causal_conv1d_fn( + x.transpose(0, 1), + weight, + None, + activation="silu", + conv_states=state, + has_initial_state=has_initial_state, + cache_indices=non_spec_state_indices_tensor, + query_start_loc=non_spec_query_start_loc, + metadata=m, + ).transpose(0, 1) + + q_ns = _prefill_conv(q_ns, q_conv_state, q_conv_weight) + k_ns = _prefill_conv(k_ns, k_conv_state, k_conv_weight) + v_ns = _prefill_conv(v_ns, v_conv_state, v_conv_weight) + q_ns, k_ns, v_ns = (rearrange(x, "n (h d) -> 1 n h d", d=self.head_dim) for x in (q_ns, k_ns, v_ns)) + + assert non_spec_state_indices_tensor is not None + assert has_initial_state is not None + initial_state = gather_initial_states( + recurrent_state, + non_spec_state_indices_tensor, + has_initial_state, + ) + if self.kda_prefill_backend == "flashkda": + assert self.gate_lower_bound is not None + ( + core_attn_out_non_spec, + last_recurrent_state, + ) = _flashkda_prefill( + q=q_ns, + k=k_ns, + v=v_ns, + g=g1_ns, + beta=beta_ns, + A_log=self.A_log, + dt_bias=self.dt_bias, + lower_bound=self.gate_lower_bound, + initial_state=initial_state, + cu_seqlens=non_spec_query_start_loc, + ) + else: + ( + core_attn_out_non_spec, + last_recurrent_state, + ) = chunk_kda_with_fused_gate( + q=q_ns, + k=k_ns, + v=v_ns, + raw_g=g1_ns, + raw_beta=beta_ns, + A_log=self.A_log, + g_bias=self.dt_bias, + lower_bound=self.gate_lower_bound, + initial_state=initial_state, + output_final_state=True, + use_qk_l2norm_in_kernel=True, + cu_seqlens=non_spec_query_start_loc, + ) + recurrent_state[non_spec_state_indices_tensor] = last_recurrent_state + else: + # Pure non-speculative decode. + assert non_spec_state_indices_tensor is not None + decode_conv_indices = non_spec_state_indices_tensor[: mixed_qkv_ns.size(0)] + packed_conv_out = torch.empty_like(mixed_qkv_ns) + mixed_qkv_ns = causal_conv1d_update( + mixed_qkv_ns, + conv_state, + conv_weights, + self.conv1d.bias, + activation="silu", + conv_state_indices=decode_conv_indices, + validate_data=True, + out=packed_conv_out, + ) + ( + core_attn_out_non_spec, + _, + ) = fused_recurrent_kda_packed_decode( + mixed_qkv=mixed_qkv_ns, + raw_g=g1_ns, + raw_beta=beta_ns, + A_log=self.A_log, + dt_bias=self.dt_bias, + lower_bound=self.gate_lower_bound, + initial_state=recurrent_state, + state_indices=decode_conv_indices, + ) + + # Restore the scheduler's original token order for mixed batches. + if core_attn_out_spec is not None and core_attn_out_non_spec is not None: + core_attn_out.index_copy_(1, spec_token_indx, core_attn_out_spec) + core_attn_out.index_copy_(1, non_spec_token_indx, core_attn_out_non_spec) + elif core_attn_out_non_spec is not None: + # TODO: prefill and decode kernels write directly to core_attn_out + core_attn_out[0, :num_actual_tokens] = core_attn_out_non_spec[0, :num_actual_tokens] + else: + assert core_attn_out_spec is not None + # Triton normalizes in place, so this is a self-copy with no device + # work. Keep it for the out-of-place native implementation. + core_attn_out.copy_(self.o_norm(core_attn_out, g2)) diff --git a/aphrodite/models/kimi_k3/nvidia/kda_metadata.py b/aphrodite/models/kimi_k3/nvidia/kda_metadata.py new file mode 100644 index 0000000000..5f1a8d049c --- /dev/null +++ b/aphrodite/models/kimi_k3/nvidia/kda_metadata.py @@ -0,0 +1,474 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Kimi-K3 specialization of GDN attention metadata. + +The request classification and cudagraph staging intentionally mirror +``GDNAttentionMetadataBuilder``. Kimi-K3 builds the metadata required by its +prefill KDA kernel internally, so this builder omits the shared FLA chunk +metadata construction. +""" + +from dataclasses import dataclass +from functools import cache + +import torch + +from aphrodite.platforms import current_platform +from aphrodite.triton_utils import tl, triton +from aphrodite.utils.torch_utils import async_tensor_h2d +from aphrodite.v1.attention.backend import CommonAttentionMetadata +from aphrodite.v1.attention.backends.gdn_attn import ( + GDNAttentionBackend, + GDNAttentionMetadata, + GDNAttentionMetadataBuilder, +) +from aphrodite.v1.attention.backends.utils import ( + NULL_BLOCK_ID, + compute_causal_conv1d_metadata, + split_decodes_and_prefills, +) +from aphrodite.v1.kv_cache_interface import MambaSpec + + +@cache +def _metadata_launch_pdl() -> bool: + return current_platform.is_arch_support_pdl() + + +@triton.jit(do_not_specialize=["num_requests"]) +def _get_aligned_state_indices_kernel( + block_table_ptr, + seq_lens_ptr, + state_indices_ptr, + block_table_stride_0: tl.constexpr, + block_table_stride_1: tl.constexpr, + seq_lens_stride: tl.constexpr, + state_indices_stride_0: tl.constexpr, + state_indices_stride_1: tl.constexpr, + num_requests, + CACHE_BLOCK_SIZE: tl.constexpr, + NUM_STATE_SLOTS: tl.constexpr, + BLOCK_STATE_SLOTS: tl.constexpr, + BLOCK_ROWS: tl.constexpr, + launch_pdl: tl.constexpr, +): + if launch_pdl: + tl.extra.cuda.gdc_wait() + tl.extra.cuda.gdc_launch_dependents() + + rows = tl.program_id(0) * BLOCK_ROWS + tl.arange(0, BLOCK_ROWS) + valid_row = rows < num_requests + seq_lens = tl.load( + seq_lens_ptr + rows * seq_lens_stride, + mask=valid_row, + other=1, + ) + # Triton truncates signed division toward zero, unlike PyTorch floor + # division. Clamping makes both semantics equivalent for seq_lens <= 0. + first_state_slot = tl.maximum((seq_lens - 1) // CACHE_BLOCK_SIZE, 0) + + state_slots = tl.arange(0, BLOCK_STATE_SLOTS) + valid_state_slot = state_slots < NUM_STATE_SLOTS + state_indices = tl.load( + block_table_ptr + + rows[:, None] * block_table_stride_0 + + (first_state_slot[:, None] + state_slots[None, :]) * block_table_stride_1, + mask=valid_row[:, None] & valid_state_slot[None, :], + ) + tl.store( + state_indices_ptr + rows[:, None] * state_indices_stride_0 + state_slots[None, :] * state_indices_stride_1, + state_indices, + mask=valid_row[:, None] & valid_state_slot[None, :], + ) + + +def _mamba_get_block_table_tensor( + block_table: torch.Tensor, + seq_lens: torch.Tensor, + kv_cache_spec: MambaSpec, + mamba_cache_mode: str, +) -> torch.Tensor: + if mamba_cache_mode in ("all", "none"): + return block_table + + assert block_table.is_cuda and seq_lens.is_cuda + num_requests = block_table.shape[0] + num_state_slots = 1 + kv_cache_spec.num_speculative_blocks + state_indices = torch.empty( + (num_requests, num_state_slots), + dtype=block_table.dtype, + device=block_table.device, + ) + BLOCK_ROWS = 32 + grid = (triton.cdiv(num_requests, BLOCK_ROWS),) + _get_aligned_state_indices_kernel[grid]( + block_table, + seq_lens, + state_indices, + block_table.stride(0), + block_table.stride(1), + seq_lens.stride(0), + state_indices.stride(0), + state_indices.stride(1), + num_requests, + CACHE_BLOCK_SIZE=kv_cache_spec.block_size, + NUM_STATE_SLOTS=num_state_slots, + BLOCK_STATE_SLOTS=triton.next_power_of_2(num_state_slots), + BLOCK_ROWS=BLOCK_ROWS, + num_warps=1, + launch_pdl=_metadata_launch_pdl(), + ) + return state_indices + + +@triton.jit(do_not_specialize=["num_spec_decodes", "batch_size"]) +def _stage_spec_decode_metadata_kernel( + state_indices_ptr, + query_start_loc_ptr, + num_accepted_tokens_ptr, + staged_state_indices_ptr, + staged_query_start_loc_ptr, + staged_num_accepted_tokens_ptr, + state_indices_stride_0: tl.constexpr, + state_indices_stride_1: tl.constexpr, + staged_state_indices_stride_0: tl.constexpr, + staged_state_indices_stride_1: tl.constexpr, + num_spec_decodes, + batch_size, + NUM_STATE_SLOTS: tl.constexpr, + BLOCK_STATE_SLOTS: tl.constexpr, + NULL_STATE_ID: tl.constexpr, + BLOCK_ROWS: tl.constexpr, + launch_pdl: tl.constexpr, +): + if launch_pdl: + tl.extra.cuda.gdc_wait() + tl.extra.cuda.gdc_launch_dependents() + + rows = tl.program_id(0) * BLOCK_ROWS + tl.arange(0, BLOCK_ROWS) + real_request = rows < num_spec_decodes + + state_slots = tl.arange(0, BLOCK_STATE_SLOTS) + valid_state_slot = state_slots < NUM_STATE_SLOTS + state_indices = tl.load( + state_indices_ptr + rows[:, None] * state_indices_stride_0 + state_slots[None, :] * state_indices_stride_1, + mask=real_request[:, None] & valid_state_slot[None, :], + other=NULL_STATE_ID, + ) + tl.store( + staged_state_indices_ptr + + rows[:, None] * staged_state_indices_stride_0 + + state_slots[None, :] * staged_state_indices_stride_1, + state_indices, + mask=(rows < batch_size)[:, None] & valid_state_slot[None, :], + ) + + query_row = tl.minimum(rows, num_spec_decodes) + query_start_loc = tl.load( + query_start_loc_ptr + query_row, + mask=rows <= batch_size, + ) + tl.store( + staged_query_start_loc_ptr + rows, + query_start_loc, + mask=rows <= batch_size, + ) + + num_accepted_tokens = tl.load( + num_accepted_tokens_ptr + rows, + mask=real_request, + other=1, + ) + tl.store( + staged_num_accepted_tokens_ptr + rows, + num_accepted_tokens, + mask=rows < batch_size, + ) + + +def stage_spec_decode_metadata( + state_indices: torch.Tensor, + query_start_loc: torch.Tensor, + num_accepted_tokens: torch.Tensor, + staged_state_indices: torch.Tensor, + staged_query_start_loc: torch.Tensor, + staged_num_accepted_tokens: torch.Tensor, + *, + num_spec_decodes: int, +) -> None: + """Stage speculative-decode metadata into CUDA-graph buffers.""" + assert state_indices.is_cuda + assert state_indices.ndim == 2 + batch_size, num_state_slots = staged_state_indices.shape + BLOCK_ROWS = 32 + grid = (triton.cdiv(batch_size + 1, BLOCK_ROWS),) + _stage_spec_decode_metadata_kernel[grid]( + state_indices, + query_start_loc, + num_accepted_tokens, + staged_state_indices, + staged_query_start_loc, + staged_num_accepted_tokens, + state_indices.stride(0), + state_indices.stride(1), + staged_state_indices.stride(0), + staged_state_indices.stride(1), + num_spec_decodes, + batch_size, + NUM_STATE_SLOTS=num_state_slots, + BLOCK_STATE_SLOTS=triton.next_power_of_2(num_state_slots), + NULL_STATE_ID=NULL_BLOCK_ID, + BLOCK_ROWS=BLOCK_ROWS, + num_warps=1, + launch_pdl=_metadata_launch_pdl(), + ) + + +@dataclass +class KimiK3KDAMetadata(GDNAttentionMetadata): + pass + + +class KimiK3KDAMetadataBuilder(GDNAttentionMetadataBuilder): + def build( # type: ignore[override] + self, + common_prefix_len: int, + common_attn_metadata: CommonAttentionMetadata, + num_accepted_tokens: torch.Tensor | None = None, + num_decode_draft_tokens_cpu: torch.Tensor | None = None, + fast_build: bool = False, + ) -> KimiK3KDAMetadata: + m = common_attn_metadata + query_start_loc = m.query_start_loc + query_start_loc_cpu = m.query_start_loc_cpu + assert isinstance(self.kv_cache_spec, MambaSpec) + # Equivalent PyTorch "align" path: + # start = ((seq_lens - 1) // block_size).clamp_(min=0) + # offsets = torch.arange(1 + num_speculative_blocks, dtype=torch.int32) + # indices = (start[:, None] + offsets).to(torch.int64) + # block_table_tensor = torch.gather(block_table, 1, indices) + block_table_tensor = _mamba_get_block_table_tensor( + m.block_table_tensor, + m.seq_lens, + self.kv_cache_spec, + self.aphrodite_config.cache_config.mamba_cache_mode, + ) + + if not self.use_spec_decode or num_decode_draft_tokens_cpu is None: + spec_sequence_masks_cpu = None + num_spec_decodes = 0 + else: + spec_sequence_masks_cpu = num_decode_draft_tokens_cpu >= 0 + # A nonnegative entry identifies a spec request. If no draft token + # was scheduled, process the whole batch as non-spec instead. + if num_decode_draft_tokens_cpu[spec_sequence_masks_cpu].sum().item() == 0: + spec_sequence_masks_cpu = None + num_spec_decodes = 0 + else: + num_spec_decodes = spec_sequence_masks_cpu.sum().item() + + if num_spec_decodes == 0: + # The runner orders ordinary decodes before prefills. + num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens = split_decodes_and_prefills( + m, decode_threshold=1 + ) + num_spec_decode_tokens = 0 + spec_token_indx = None + non_spec_token_indx = None + spec_state_indices_tensor = None + non_spec_state_indices_tensor = block_table_tensor[:, 0] + spec_query_start_loc = None + non_spec_query_start_loc = query_start_loc + non_spec_query_start_loc_cpu = query_start_loc_cpu + num_accepted_tokens = None + else: + assert spec_sequence_masks_cpu is not None + assert num_accepted_tokens is not None + query_lens_cpu = query_start_loc_cpu.diff() + num_query_tokens = query_start_loc_cpu[-1].item() + + # Exclude zero-length cudagraph padding from request-indexed + # non-spec metadata. + active_non_spec_mask_cpu = (~spec_sequence_masks_cpu) & (query_lens_cpu > 0) + non_spec_query_lens_cpu = query_lens_cpu[active_non_spec_mask_cpu] + num_non_spec_requests = non_spec_query_lens_cpu.numel() + num_non_spec_tokens = non_spec_query_lens_cpu.sum().item() + + # Query length alone cannot distinguish a true decode from a + # one-token prefill chunk. Packed decode is safe only when every + # active non-spec request is a true, one-token decode. + assert m.is_prefilling is not None + assert m.is_prefilling.device.type == "cpu" + non_spec_is_prefilling = m.is_prefilling[active_non_spec_mask_cpu] + use_prefill = torch.any(non_spec_is_prefilling | (non_spec_query_lens_cpu != 1)).item() + if use_prefill: + num_prefills = num_non_spec_requests + num_prefill_tokens = num_non_spec_tokens + num_decodes = 0 + num_decode_tokens = 0 + else: + num_prefills = 0 + num_prefill_tokens = 0 + num_decodes = num_non_spec_requests + num_decode_tokens = num_non_spec_tokens + + num_spec_decode_tokens = num_query_tokens - num_non_spec_tokens + + if num_prefills == 0 and num_decodes == 0: + spec_token_indx = None + non_spec_token_indx = None + # Real requests precede trailing cudagraph padding. + spec_state_indices_tensor = block_table_tensor[:num_spec_decodes, : self.num_spec + 1] + non_spec_state_indices_tensor = None + # Padding trails real requests, so this prefix already contains + # the correct cumulative token counts. + spec_query_start_loc = query_start_loc[: num_spec_decodes + 1] + non_spec_query_start_loc = None + non_spec_query_start_loc_cpu = None + num_accepted_tokens = num_accepted_tokens[:num_spec_decodes] + else: + query_lens = query_start_loc.diff() + spec_sequence_masks_gpu = async_tensor_h2d(spec_sequence_masks_cpu, device=query_start_loc.device) + spec_token_masks = torch.repeat_interleave( + spec_sequence_masks_gpu, + query_lens, + output_size=num_query_tokens, + ) + # Stable partitioning preserves request-local token order in + # both subgroup tensors. + index = torch.argsort(spec_token_masks, stable=True) + num_non_spec_tokens = num_prefill_tokens + num_decode_tokens + non_spec_token_indx = index[:num_non_spec_tokens] + spec_token_indx = index[num_non_spec_tokens:] + + # Spec requests carry one state slot per speculative step; + # non-spec requests use only their current state slot. + spec_state_indices_tensor = block_table_tensor[spec_sequence_masks_cpu, : self.num_spec + 1] + non_spec_state_indices_tensor = block_table_tensor[active_non_spec_mask_cpu, 0] + + spec_query_lens = query_lens[spec_sequence_masks_cpu] + spec_query_start_loc = torch.zeros( + num_spec_decodes + 1, + dtype=torch.int32, + device=query_start_loc.device, + ) + torch.cumsum( + spec_query_lens, + dim=0, + out=spec_query_start_loc[1:], + ) + if num_prefills > 0: + non_spec_query_lens = query_lens[active_non_spec_mask_cpu] + non_spec_query_start_loc = torch.zeros( + non_spec_query_lens.size(0) + 1, + dtype=torch.int32, + device=query_start_loc.device, + ) + torch.cumsum( + non_spec_query_lens, + dim=0, + out=non_spec_query_start_loc[1:], + ) + non_spec_query_start_loc_cpu = torch.zeros( + non_spec_query_lens_cpu.size(0) + 1, + dtype=torch.int32, + ) + torch.cumsum( + non_spec_query_lens_cpu, + dim=0, + out=non_spec_query_start_loc_cpu[1:], + ) + else: + # Packed decode consumes one row per request and does not + # use cumulative sequence lengths. + non_spec_query_start_loc = None + non_spec_query_start_loc_cpu = None + + num_accepted_tokens = num_accepted_tokens[spec_sequence_masks_cpu] + + # Unlike the shared GDN layer, Kimi-K3's prefill KDA wrapper prepares + # its own chunk indices. Only causal-convolution metadata is needed here. + nums_dict, batch_ptr, token_chunk_offset_ptr = None, None, None + if num_prefills > 0: + has_initial_state = m.compute_num_computed_tokens() > 0 + if spec_sequence_masks_cpu is not None: + has_initial_state = has_initial_state[active_non_spec_mask_cpu] + assert non_spec_query_start_loc_cpu is not None + nums_dict, batch_ptr, token_chunk_offset_ptr = compute_causal_conv1d_metadata( + non_spec_query_start_loc_cpu, + device=query_start_loc.device, + ) + else: + has_initial_state = None + + # Prepare per-request tensors for cudagraph replay. num_actual_tokens + # may be token-padded, while state/query/acceptance metadata is indexed + # by request. + batch_size = m.num_reqs + if ( + self.use_full_cuda_graph + and num_spec_decodes > 0 + and num_prefills == 0 + and num_decodes == 0 + and num_spec_decodes <= self.decode_cudagraph_max_bs + and num_spec_decode_tokens <= self.decode_cudagraph_max_bs + ): + # Equivalent PyTorch staging: + # state[:N].copy_(state_src); state[N:].fill_(NULL_BLOCK_ID) + # qsl[:N + 1].copy_(qsl_src); qsl[N + 1:].fill_(qsl_src[-1]) + # accepted[:N].copy_(accepted_src); accepted[N:].fill_(1) + stage_spec_decode_metadata( + state_indices=spec_state_indices_tensor, + query_start_loc=spec_query_start_loc, + num_accepted_tokens=num_accepted_tokens, + staged_state_indices=self.spec_state_indices_tensor[:batch_size], + staged_query_start_loc=self.spec_query_start_loc[: batch_size + 1], + staged_num_accepted_tokens=self.num_accepted_tokens[:batch_size], + num_spec_decodes=num_spec_decodes, + ) + + spec_state_indices_tensor = self.spec_state_indices_tensor[:batch_size] + spec_query_start_loc = self.spec_query_start_loc[: batch_size + 1] + num_accepted_tokens = self.num_accepted_tokens[:batch_size] + + if ( + self.use_full_cuda_graph + and num_prefills == 0 + and num_spec_decodes == 0 + and num_decodes <= self.decode_cudagraph_max_bs + ): + self.non_spec_state_indices_tensor[:num_decodes].copy_(non_spec_state_indices_tensor, non_blocking=True) + self.non_spec_state_indices_tensor[num_decodes:batch_size].fill_(NULL_BLOCK_ID) + non_spec_state_indices_tensor = self.non_spec_state_indices_tensor[:batch_size] + + return KimiK3KDAMetadata( + num_prefills=num_prefills, + num_prefill_tokens=num_prefill_tokens, + num_decodes=num_decodes, + num_decode_tokens=num_decode_tokens, + num_spec_decodes=num_spec_decodes, + num_spec_decode_tokens=num_spec_decode_tokens, + num_actual_tokens=m.num_actual_tokens, + has_initial_state=has_initial_state, + spec_query_start_loc=spec_query_start_loc, + non_spec_query_start_loc=non_spec_query_start_loc, + spec_state_indices_tensor=spec_state_indices_tensor, + non_spec_state_indices_tensor=non_spec_state_indices_tensor, + spec_sequence_masks=None, + spec_token_indx=spec_token_indx, + non_spec_token_indx=non_spec_token_indx, + num_accepted_tokens=num_accepted_tokens, + nums_dict=nums_dict, + batch_ptr=batch_ptr, + token_chunk_offset_ptr=token_chunk_offset_ptr, + ) + + +class KimiK3KDAAttentionBackend(GDNAttentionBackend): + @staticmethod + def get_name() -> str: + return "KIMI_K3_KDA" + + @staticmethod + def get_builder_cls() -> type[KimiK3KDAMetadataBuilder]: + return KimiK3KDAMetadataBuilder diff --git a/aphrodite/models/kimi_k3/nvidia/low_latency_gemm.py b/aphrodite/models/kimi_k3/nvidia/low_latency_gemm.py new file mode 100644 index 0000000000..8c4308608b --- /dev/null +++ b/aphrodite/models/kimi_k3/nvidia/low_latency_gemm.py @@ -0,0 +1,487 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Kimi-K3 decode GEMM selection for unquantized BF16 on SM103. + +Dispatch is purely by local ``(N, K)`` shape and token count ``M`` — the module +name plays no role. Each measured shape maps to a :class:`ProjectionSpec` +holding the winning backend per token count. The static part of the decision is +resolved once per module at install time into a small ``{M: call}`` plan, so the +per-forward path is a single dict lookup. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Literal + +import torch +from torch import nn + +import aphrodite.envs as envs +from aphrodite import _custom_ops as ops +from aphrodite.model_executor.kernels.linear.cute_dsl.skinny_gemm import ( + SkinnyGemmConfig, + shape_dynamic_skinny_gemm, +) +from aphrodite.model_executor.layers.linear import LinearBase, UnquantizedLinearMethod +from aphrodite.model_executor.layers.vocab_parallel_embedding import ( + ParallelLMHead, + UnquantizedEmbeddingMethod, +) +from aphrodite.platforms import current_platform + +Backend = Literal["cute", "dsv3_fused_a"] +# A resolved per-token-count call: the backend plus its CuTe config (None for +# dsv3, which needs no config). +ResolvedCall = tuple[Backend, SkinnyGemmConfig | None] + + +@dataclass(frozen=True, slots=True) +class ProjectionSpec: + n: int + k: int + dsv3_tokens: frozenset[int] = frozenset() + cute_configs: tuple[tuple[int, SkinnyGemmConfig], ...] = () + residual_configs: tuple[tuple[int, SkinnyGemmConfig], ...] = () + name: str = "" # optional debug label; never used for dispatch + + def cute_config(self, num_tokens: int) -> SkinnyGemmConfig | None: + return dict(self.cute_configs).get(num_tokens) + + def residual_config(self, num_tokens: int) -> SkinnyGemmConfig | None: + return dict(self.residual_configs).get(num_tokens) + + +def _cute( + num_tokens: int, + block_size: int, + outputs_per_block: int, + k_unroll: int, + vector_width: int = 8, +) -> SkinnyGemmConfig: + return SkinnyGemmConfig( + num_tokens, + block_size, + outputs_per_block, + k_unroll, + vector_width, + ) + + +_M1_TO_16 = frozenset(range(1, 17)) +_M1 = frozenset({1}) + +# Keyed by local (N, K). Where two projections share a shape (only 1536x7168: +# shared_gate_up_proj and mla_g_proj) the entry is unified. +KIMI_K3_PROJECTIONS: dict[tuple[int, int], ProjectionSpec] = { + (1536, 128): ProjectionSpec(1536, 128, _M1_TO_16, name="f_b_proj"), + (3072, 128): ProjectionSpec(3072, 128, _M1_TO_16, name="f_b_proj"), + # 1536x7168 is shared by shared_gate_up_proj and mla_g_proj. dsv3 M1..16 is + # only crash-safe once the mla_g aux-stream/PDL capture fix lands (subtask + # task_7388aba1); the fallback if it cannot be fixed is dsv3_tokens=_M1. + (1536, 7168): ProjectionSpec(1536, 7168, _M1_TO_16, name="shared_gate_up_proj/mla_g_proj"), + (3072, 7168): ProjectionSpec( + 3072, + 7168, + cute_configs=( + (1, _cute(1, 224, 3, 4)), + (2, _cute(2, 128, 3, 2)), + (3, _cute(3, 128, 2, 1)), + (4, _cute(4, 64, 2, 2)), + (5, _cute(5, 128, 3, 1)), + ), + name="shared_gate_up_proj", + ), + (2112, 7168): ProjectionSpec(2112, 7168, _M1_TO_16, name="fused_qkv_a_proj"), + (2304, 1536): ProjectionSpec(2304, 1536, _M1_TO_16, name="q_b_proj"), + (4608, 1536): ProjectionSpec(4608, 1536, _M1_TO_16, name="q_b_proj"), + (3584, 7168): ProjectionSpec( + 3584, + 7168, + frozenset(range(2, 9)), + ((1, _cute(1, 224, 2, 4)),), + name="routed_expert_down_proj", + ), + (6288, 7168): ProjectionSpec( + 6288, + 7168, + cute_configs=( + (1, _cute(1, 224, 3, 4)), + (2, _cute(2, 64, 3, 2)), + (3, _cute(3, 32, 3, 4)), + (4, _cute(4, 128, 6, 1)), + ), + name="in_proj_qkvgfab", + ), + (12448, 7168): ProjectionSpec( + 12448, + 7168, + cute_configs=( + (1, _cute(1, 224, 4, 2)), + (2, _cute(2, 64, 4, 2)), + (3, _cute(3, 64, 2, 2)), + ), + name="in_proj_qkvgfab", + ), + (7168, 768): ProjectionSpec(7168, 768, _M1_TO_16, name="shared_down_proj"), + (7168, 1536): ProjectionSpec(7168, 1536, cute_configs=((1, _cute(1, 96, 4, 2)),), name="o_proj"), + (7168, 3072): ProjectionSpec( + 7168, + 3072, + cute_configs=( + (1, _cute(1, 96, 2, 4)), + (2, _cute(2, 32, 4, 4)), + ), + name="o_proj", + ), + (7168, 3584): ProjectionSpec( + 7168, + 3584, + cute_configs=( + (1, _cute(1, 224, 4, 2)), + (2, _cute(2, 64, 4, 2)), + ), + residual_configs=( + (1, _cute(1, 64, 4, 2)), + (2, _cute(2, 64, 7, 2)), + (3, _cute(3, 64, 2, 1)), + (4, _cute(4, 64, 2, 1)), + ), + name="routed_expert_up_proj", + ), + (7168, 4224): ProjectionSpec( + 7168, + 4224, + cute_configs=((1, _cute(1, 96, 4, 2, 4)),), + name="dense_down_proj", + ), + (7168, 8448): ProjectionSpec( + 7168, + 8448, + cute_configs=( + (1, _cute(1, 32, 4, 4)), + (2, _cute(2, 96, 4, 1)), + (3, _cute(3, 96, 4, 1)), + ), + name="dense_down_proj", + ), + (8448, 7168): ProjectionSpec( + 8448, + 7168, + cute_configs=( + (1, _cute(1, 224, 3, 4)), + (2, _cute(2, 32, 4, 4)), + ), + name="dense_gate_up_proj", + ), + (16896, 7168): ProjectionSpec( + 16896, + 7168, + cute_configs=( + (1, _cute(1, 224, 6, 4)), + (2, _cute(2, 32, 4, 4)), + ), + name="dense_gate_up_proj", + ), + (20480, 7168): ProjectionSpec( + 20480, + 7168, + cute_configs=( + (1, _cute(1, 224, 4, 2)), + (2, _cute(2, 64, 4, 2)), + (3, _cute(3, 64, 2, 2)), + (4, _cute(4, 64, 4, 1)), + ), + name="lm_head", + ), + (40960, 7168): ProjectionSpec( + 40960, + 7168, + cute_configs=( + (1, _cute(1, 128, 4, 2)), + (2, _cute(2, 64, 4, 2)), + (3, _cute(3, 64, 2, 2)), + (4, _cute(4, 64, 4, 1)), + ), + name="lm_head", + ), + # TP16. Measured on B300 over M=1..16 with the same >=5% threshold as the + # entries above. The replicated projections (2112x7168, 3584x7168, + # 7168x3584) keep their shapes at TP16 and reuse the entries above, and + # o_proj lands on 7168x768, which shared_down_proj already covers. + (3216, 7168): ProjectionSpec( + 3216, + 7168, + # Both gaps in this range are measured, not oversights: dsv3 is only + # 4% ahead at M6..M8, and at M16 cuBLAS switches to a faster kernel + # (11.42us vs dsv3's 11.83us) after trailing it by 6-8% at M9..M15. + frozenset(range(9, 16)), + cute_configs=( + (1, _cute(1, 224, 3, 4)), + (2, _cute(2, 128, 4, 2)), + (3, _cute(3, 128, 2, 1)), + (4, _cute(4, 64, 2, 2)), + (5, _cute(5, 128, 3, 1)), + ), + name="in_proj_qkvgfab", + ), + (768, 7168): ProjectionSpec( + 768, + 7168, + frozenset(range(5, 17)), + cute_configs=( + (1, _cute(1, 224, 2, 4)), + (2, _cute(2, 224, 2, 2)), + (3, _cute(3, 224, 2, 2)), + (4, _cute(4, 224, 2, 2)), + ), + name="mla_g_proj/shared_gate_up_proj", + ), + (1152, 1536): ProjectionSpec( + 1152, + 1536, + frozenset(range(2, 17)), + ((1, _cute(1, 192, 3, 4)),), + name="q_b_proj", + ), + (768, 128): ProjectionSpec(768, 128, _M1_TO_16, name="f_b_proj"), + # dsv3 drops under 5% from M9 on for this shape. + (7168, 384): ProjectionSpec(7168, 384, frozenset(range(1, 9)), name="shared_down_proj"), + (4224, 7168): ProjectionSpec( + 4224, + 7168, + frozenset(range(4, 9)), + cute_configs=( + (1, _cute(1, 224, 3, 4)), + (2, _cute(2, 128, 2, 1)), + (3, _cute(3, 64, 2, 2)), + ), + name="dense_gate_up_proj", + ), + (10240, 7168): ProjectionSpec( + 10240, + 7168, + cute_configs=( + (1, _cute(1, 224, 4, 2)), + (2, _cute(2, 32, 2, 4)), + (3, _cute(3, 64, 4, 1)), + (4, _cute(4, 64, 4, 1)), + ), + name="lm_head", + ), + # 7168x2112 (TP16 dense down_proj) has no entry on purpose: K=2112 divides + # none of the fused-A tile_k values, and the CuTe kernel is left with + # vector_width=2, which measured slower than cuBLAS. +} + + +def _backend_for(spec: ProjectionSpec, num_tokens: int, has_residual: bool) -> Backend | None: + if has_residual: + return "cute" if spec.residual_config(num_tokens) is not None else None + if spec.cute_config(num_tokens) is not None: + return "cute" + if num_tokens in spec.dsv3_tokens: + return "dsv3_fused_a" + return None + + +def select_kimi_k3_backend( + num_tokens: int, + n: int, + k: int, + *, + has_residual: bool = False, +) -> Backend | None: + """Backend for a local ``(N, K)`` at ``num_tokens``, or None to fall back.""" + spec = KIMI_K3_PROJECTIONS.get((n, k)) + return _backend_for(spec, num_tokens, has_residual) if spec is not None else None + + +def _build_plan(spec: ProjectionSpec) -> dict[int, ResolvedCall]: + plan: dict[int, ResolvedCall] = {} + for num_tokens in range(1, 17): + backend = _backend_for(spec, num_tokens, has_residual=False) + if backend == "cute": + plan[num_tokens] = ("cute", spec.cute_config(num_tokens)) + elif backend == "dsv3_fused_a": + plan[num_tokens] = ("dsv3_fused_a", None) + return plan + + +def _build_residual_plan(spec: ProjectionSpec) -> dict[int, SkinnyGemmConfig]: + return {num_tokens: config for num_tokens, config in spec.residual_configs} + + +def _is_sm103() -> bool: + return current_platform.is_device_capability((10, 3)) + + +def _is_packed_row_major(tensor: torch.Tensor) -> bool: + return tensor.dim() == 2 and tensor.stride() == (tensor.shape[1], 1) + + +def _runtime_ok(x: torch.Tensor, weight: torch.Tensor) -> bool: + return ( + _is_packed_row_major(x) + and _is_packed_row_major(weight) + and x.dtype == torch.bfloat16 + and weight.dtype == torch.bfloat16 + and x.is_cuda + and weight.is_cuda + and x.device == weight.device + and x.shape[1] == weight.shape[1] + ) + + +def _residual_ok(x: torch.Tensor, weight: torch.Tensor, residual: torch.Tensor) -> bool: + return ( + residual.dim() == 2 + and residual.dtype == torch.bfloat16 + and residual.is_cuda + and residual.device == x.device + and residual.is_contiguous() + and residual.shape == (x.shape[0], weight.shape[0]) + ) + + +def _run_plan(plan: dict[int, ResolvedCall], x: torch.Tensor, weight: torch.Tensor) -> torch.Tensor | None: + entry = plan.get(x.shape[0]) + if entry is None: + return None + backend, config = entry + if backend == "cute": + if not shape_dynamic_skinny_gemm.is_available(): + return None + return shape_dynamic_skinny_gemm(x, weight, config, None) + if not hasattr(torch.ops._C, "dsv3_fused_a_gemm"): + return None + output = torch.empty((x.shape[0], weight.shape[0]), dtype=x.dtype, device=x.device) + ops.dsv3_fused_a_gemm(output, x, weight.t(), enable_pdl=True) + return output + + +def _run_residual_plan( + residual_plan: dict[int, SkinnyGemmConfig], + x: torch.Tensor, + weight: torch.Tensor, + residual: torch.Tensor, +) -> torch.Tensor | None: + config = residual_plan.get(x.shape[0]) + if config is None or not shape_dynamic_skinny_gemm.is_available(): + return None + return shape_dynamic_skinny_gemm(x, weight, config, residual) + + +def try_low_latency_gemm( + x: torch.Tensor, + weight: torch.Tensor, + residual: torch.Tensor | None = None, +) -> torch.Tensor | None: + """Run the shape-selected low-latency kernel, or None to fall back. + + Resolves the plan from the shape table on each call; production installs a + precomputed plan (see :func:`enable_kimi_k3_low_latency_gemm`) and does not + use this path. + """ + if envs.APHRODITE_BATCH_INVARIANT or not _is_sm103() or not _runtime_ok(x, weight): + return None + spec = KIMI_K3_PROJECTIONS.get((weight.shape[0], weight.shape[1])) + if spec is None: + return None + if residual is None: + return _run_plan(_build_plan(spec), x, weight) + if not _residual_ok(x, weight, residual): + return None + return _run_residual_plan(_build_residual_plan(spec), x, weight, residual) + + +class _KimiK3LowLatencyApply: + """Mixin: try the precomputed plan, else defer to the base method.""" + + def __init__(self, plan: dict[int, ResolvedCall]) -> None: + self._plan = plan + + def apply( + self, + layer: nn.Module, + x: torch.Tensor, + bias: torch.Tensor | None = None, + ) -> torch.Tensor: + if bias is None and not envs.APHRODITE_BATCH_INVARIANT and _runtime_ok(x, layer.weight): + output = _run_plan(self._plan, x, layer.weight) + if output is not None: + return output + return super().apply(layer, x, bias) # type: ignore[misc] + + +class KimiK3LowLatencyLinearMethod(_KimiK3LowLatencyApply, UnquantizedLinearMethod): + def __init__( + self, + plan: dict[int, ResolvedCall], + residual_plan: dict[int, SkinnyGemmConfig], + ) -> None: + super().__init__(plan) + self._residual_plan = residual_plan + + def apply_with_residual( + self, + layer: nn.Module, + x: torch.Tensor, + residual: torch.Tensor, + ) -> torch.Tensor: + if ( + not envs.APHRODITE_BATCH_INVARIANT + and _runtime_ok(x, layer.weight) + and _residual_ok(x, layer.weight, residual) + ): + output = _run_residual_plan(self._residual_plan, x, layer.weight, residual) + if output is not None: + return output + return torch.addmm(residual, x, layer.weight.t()) + + +class KimiK3LowLatencyEmbeddingMethod(_KimiK3LowLatencyApply, UnquantizedEmbeddingMethod): + pass + + +def enable_kimi_k3_low_latency_gemm( + module: nn.Module, + dtype: torch.dtype, +) -> None: + """Install shape-selected low-latency GEMMs and register CuTe warmups. + + Modules are matched purely by type, an exactly-unquantized method, and a + local ``(N, K)`` present in :data:`KIMI_K3_PROJECTIONS`. + """ + if dtype != torch.bfloat16 or not _is_sm103(): + return + + warmup_configs: set[SkinnyGemmConfig] = set() + residual_warmup_configs: set[SkinnyGemmConfig] = set() + for child in module.modules(): + is_linear = isinstance(child, LinearBase) and type(child.quant_method) is UnquantizedLinearMethod + # ParallelLMHead is a VocabParallelEmbedding subclass; embed_tokens is + # the parent type, so isinstance already excludes it. + is_head = isinstance(child, ParallelLMHead) and type(child.quant_method) is UnquantizedEmbeddingMethod + if not (is_linear or is_head): + continue + weight = getattr(child, "weight", None) + if weight is None or weight.dim() != 2: + continue + spec = KIMI_K3_PROJECTIONS.get((weight.shape[0], weight.shape[1])) + if spec is None: + continue + if is_linear: + child.quant_method = KimiK3LowLatencyLinearMethod(_build_plan(spec), _build_residual_plan(spec)) + else: + child.quant_method = KimiK3LowLatencyEmbeddingMethod(_build_plan(spec)) + # Warm up only the configs measured for this module's local (N, K) so a + # TP8 deployment does not compile TP4 configs and vice versa. + warmup_configs.update(config for _, config in spec.cute_configs) + residual_warmup_configs.update(config for _, config in spec.residual_configs) + + if shape_dynamic_skinny_gemm.is_available(): + if warmup_configs: + shape_dynamic_skinny_gemm.request_warmup_configs(dtype, warmup_configs) + if residual_warmup_configs: + shape_dynamic_skinny_gemm.request_warmup_configs(dtype, residual_warmup_configs, has_residual=True) diff --git a/aphrodite/models/kimi_k3/nvidia/mla.py b/aphrodite/models/kimi_k3/nvidia/mla.py new file mode 100644 index 0000000000..c25e9c13f2 --- /dev/null +++ b/aphrodite/models/kimi_k3/nvidia/mla.py @@ -0,0 +1,725 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Clean Multi-head Latent Attention for Kimi-K3 (NVIDIA). + +This is a self-contained MLA layer that owns the full attention path: + + hidden_states + -> fused pre-attention ops (fused_qkv_a_proj / norms / q_b_proj) + -> explicit prefill / decode split + prefill: fused key-concat + cache-insert kernel -> run_prefill_new_tokens + (+ chunked-context merge); dispatched by cache dtype + (bf16 / plain fp8 / fp8_ds_mla) + decode : W_UK absorb (BMM1) -> fused q-concat + cache-insert kernel + -> impl.forward_mqa -> W_UV up-proj (MQA) + -> optional output gate + -> o_proj + +Unlike ``MultiHeadLatentAttentionWrapper`` (which delegates orchestration to +``MLAAttention.forward``), this class *is* the ``AttentionLayerBase``: it selects +the backend, builds the impl, registers itself in the forward context, owns the +KV cache, and absorbs ``kv_b_proj`` into ``W_UK_T`` / ``W_UV`` -- mirroring the +``DeepseekV4Attention`` structure. + +K3 specifics: optional rotary embedding (disabled for the target model's NoPE +layers, enabled for DSpark) and an optional sigmoid output gate (``g_proj``). + +Out of scope (extension points, not wired here): context parallelism (DCP/PCP), +sparse/indexer MLA, and the ROCm/aiter fp8/fp4 BMM fast paths. +""" + +import math +from typing import TYPE_CHECKING, cast + +import torch +from torch import nn + +from aphrodite.compilation.breakable_cudagraph import eager_break_during_capture +from aphrodite.config import ( + AphroditeConfig, + CacheConfig, +) +from aphrodite.config import ( + get_current_aphrodite_config as get_current_aphrodite_config, +) +from aphrodite.distributed import get_tensor_model_parallel_world_size +from aphrodite.forward_context import get_forward_context +from aphrodite.logger import init_logger +from aphrodite.model_executor.layers.attention.attention import ( + _init_kv_cache_quant, + set_default_quant_scales, + should_load_quant_weights, +) +from aphrodite.model_executor.layers.attention_layer_base import AttentionLayerBase +from aphrodite.model_executor.layers.layernorm import RMSNorm +from aphrodite.model_executor.layers.linear import ( + ColumnParallelLinear, + MergedColumnParallelLinear, + ReplicatedLinear, + RowParallelLinear, +) +from aphrodite.model_executor.layers.quantization import QuantizationConfig +from aphrodite.model_executor.layers.quantization.utils.quant_utils import ( + get_and_maybe_dequant_weights, +) +from aphrodite.model_executor.layers.rotary_embedding import RotaryEmbedding, get_rope +from aphrodite.model_executor.utils import replace_parameter +from aphrodite.models.common.ops import fused_q_kv_rmsnorm +from aphrodite.models.kimi_k3.nvidia.ops.fused_mla_key_concat_kv_cache import ( + fused_mla_decode_q_concat_kv_cache_insert, + fused_mla_key_concat_ds_mla_insert, + fused_mla_key_concat_kv_cache_insert, + fused_mla_qkv_quant_kv_cache_fp8_insert, +) +from aphrodite.platforms import current_platform +from aphrodite.transformers_utils.configs.kimi_linear import KimiLinearConfig +from aphrodite.utils.multi_stream_utils import maybe_execute_in_parallel +from aphrodite.utils.torch_utils import ( + is_quantized_kv_cache, + kv_cache_dtype_str_to_dtype, +) +from aphrodite.v1.attention.backend import ( + AttentionBackend, + AttentionType, + MLAAttentionImpl, +) +from aphrodite.v1.attention.backends.mla.prefill import get_mla_prefill_backend +from aphrodite.v1.attention.ops.merge_attn_states import merge_attn_states +from aphrodite.v1.attention.selector import get_attn_backend +from aphrodite.v1.kv_cache_interface import KVCacheSpec, MLAAttentionSpec, get_kv_quant_mode + +if TYPE_CHECKING: + from aphrodite.model_executor.layers.attention.mla_attention import MLACommonMetadata + +logger = init_logger(__name__) + +# Below this many tokens, overlap the g_proj GEMM on the aux stream with the +# attention front-end (the GEMM is small and launch-bound, so the overlap +# hides it); at or above it, run the gate on the main stream. +_GATE_MULTI_STREAM_TOKEN_THRESHOLD = 512 + + +@torch.compile(backend=current_platform.simple_compile_backend) +def _gate_sigmoid_mul(attn_out: torch.Tensor, gate: torch.Tensor) -> torch.Tensor: + """Apply the sigmoid output gate to a precomputed ``g_proj`` projection.""" + return attn_out * gate.sigmoid() + + +class MultiHeadLatentAttention(nn.Module, AttentionLayerBase): + """Kimi-K3 Multi-head Latent Attention with optional RoPE and output gate.""" + + def __init__( + self, + config: KimiLinearConfig, + hidden_size: int, + num_heads: int, + qk_nope_head_dim: int, + qk_rope_head_dim: int, + v_head_dim: int, + q_lora_rank: int | None, + kv_lora_rank: int, + use_output_gate: bool = False, + cache_config: CacheConfig | None = None, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + aux_stream: torch.cuda.Stream | None = None, + use_rope: bool = False, + non_causal_multi_token_decode: bool = False, + ) -> None: + super().__init__() + self.hidden_size = hidden_size + self.qk_nope_head_dim = qk_nope_head_dim + self.qk_rope_head_dim = qk_rope_head_dim + self.qk_head_dim = qk_nope_head_dim + qk_rope_head_dim + self.v_head_dim = v_head_dim + self.q_lora_rank = q_lora_rank + self.kv_lora_rank = kv_lora_rank + self.non_causal_multi_token_decode = non_causal_multi_token_decode + # Latent "head" seen by the attention kernel / KV cache. + self.head_size = kv_lora_rank + qk_rope_head_dim + self.scale = self.qk_head_dim**-0.5 + self.rms_norm_eps = config.rms_norm_eps + self.layer_name = prefix + + self.rotary_emb: RotaryEmbedding | None = None + if use_rope: + rope_parameters = dict(config.rope_parameters) + if rope_parameters["rope_type"] != "default": + rope_parameters["rope_type"] = ( + "deepseek_yarn" if rope_parameters.get("apply_yarn_scaling", True) else "deepseek_llama_scaling" + ) + self.rotary_emb = get_rope( + qk_rope_head_dim, + max_position=config.max_position_embeddings, + rope_parameters=rope_parameters, + is_neox_style=False, + dtype=torch.float32, + ) + if rope_parameters["rope_type"] == "deepseek_yarn": + mscale_all_dim = rope_parameters.get("mscale_all_dim", False) + scaling_factor = rope_parameters["factor"] + mscale = 1.0 if scaling_factor <= 1 else 0.1 * float(mscale_all_dim) * math.log(scaling_factor) + 1.0 + self.scale *= mscale * mscale + # The fused epilogues read the cos/sin table directly in fp32 and run + # the RoPE math in fp32, so there is no per-forward dtype cast (and no + # precision loss). deepseek_yarn builds cos_sin_cache in fp32 already; + # dtype=torch.float32 above forces it for the default rope too (the + # DSpark draft, which has no yarn scaling). + assert self.rotary_emb.cos_sin_cache.dtype == torch.float32, ( + f"K3 fused MLA RoPE requires an fp32 cos/sin cache; got {self.rotary_emb.cos_sin_cache.dtype}." + ) + + tp_size = get_tensor_model_parallel_world_size() + assert num_heads % tp_size == 0 + self.num_heads = num_heads + self.num_local_heads = num_heads // tp_size + + # ---- Pre-attention projections (fusable front-end) ---- + # Two query variants: a low-rank q-LoRA path (Kimi-K3) fused with the + # kv-down proj, or an uncompressed q path (Kimi-Linear, ``q_lora_rank`` + # None) with a standalone ``q_proj`` and separate ``kv_a_proj_with_mqa``. + if self.q_lora_rank is not None: + # Fused q-down + kv-down projection. Replicated (disable_tp) because + # the low-rank latents are shared across TP ranks; TP splitting + # happens at q_b_proj / kv_b_proj. Checkpoint weights ``q_a_proj`` + # and ``kv_a_proj_with_mqa`` map onto shards 0 and 1 respectively. + self.fused_qkv_a_proj = MergedColumnParallelLinear( + self.hidden_size, + [self.q_lora_rank, self.kv_lora_rank + self.qk_rope_head_dim], + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.fused_qkv_a_proj", + disable_tp=True, + ) + self.q_a_layernorm = RMSNorm(self.q_lora_rank, eps=config.rms_norm_eps) + self.q_b_proj = ColumnParallelLinear( + self.q_lora_rank, + self.num_heads * self.qk_head_dim, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.q_b_proj", + ) + else: + # Uncompressed query: full-rank q_proj (TP-split over heads) plus a + # replicated kv-down projection (shared latent across TP ranks). + self.q_proj = ColumnParallelLinear( + self.hidden_size, + self.num_heads * self.qk_head_dim, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.q_proj", + ) + self.kv_a_proj_with_mqa = ReplicatedLinear( + self.hidden_size, + self.kv_lora_rank + self.qk_rope_head_dim, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.kv_a_proj_with_mqa", + ) + self.kv_a_layernorm = RMSNorm(self.kv_lora_rank, eps=config.rms_norm_eps) + self.kv_b_proj = ColumnParallelLinear( + self.kv_lora_rank, + self.num_heads * (self.qk_nope_head_dim + self.v_head_dim), + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.kv_b_proj", + ) + + # ---- Post-attention projections ---- + self.use_output_gate = use_output_gate + self.g_proj = ( + ColumnParallelLinear( + self.hidden_size, + self.num_heads * self.v_head_dim, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.g_proj", + ) + if use_output_gate + else None + ) + # Aux stream (created at the model level, DeepseekV4 convention) for + # overlapping the g_proj GEMM with the attention front-end. None on + # ROCm/non-cuda -> maybe_execute_in_parallel falls back to sequential. + self.aux_stream = aux_stream + self._gate_events = ( + [torch.cuda.Event(), torch.cuda.Event()] + if self.g_proj is not None and current_platform.is_cuda_alike() + else None + ) + self.o_proj = RowParallelLinear( + self.num_heads * self.v_head_dim, + self.hidden_size, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.o_proj", + ) + + # ---- Attention backend / impl / KV cache ---- + self.quant_config = quant_config + if cache_config is not None: + self.kv_cache_dtype = cache_config.cache_dtype + else: + self.kv_cache_dtype = "auto" + + dtype = torch.get_default_dtype() + self.attn_backend = get_attn_backend( + self.head_size, + dtype, + self.kv_cache_dtype, + use_mla=True, + use_sparse=False, + num_heads=self.num_local_heads, + ) + _init_kv_cache_quant(self, quant_config, prefix) + # Unit (1.0) scale for the fused fp8 prefill path: q/k/v are cast + # unscaled to match forward_mha (the prefill flash path does not + # dequantize); only the cache uses _k_scale. + self.register_buffer("_one_scale", torch.ones(1, dtype=torch.float32), persistent=False) + + impl_cls = cast(type[MLAAttentionImpl], self.attn_backend.get_impl_cls()) + self.impl = impl_cls( # type: ignore[assignment] + num_heads=self.num_local_heads, + head_size=self.head_size, + scale=self.scale, + num_kv_heads=1, + alibi_slopes=None, + sliding_window=None, + kv_cache_dtype=self.kv_cache_dtype, + logits_soft_cap=None, + attn_type=AttentionType.DECODER, + kv_sharing_target_layer_name=None, + q_lora_rank=self.q_lora_rank, + kv_lora_rank=self.kv_lora_rank, + qk_nope_head_dim=self.qk_nope_head_dim, + qk_rope_head_dim=self.qk_rope_head_dim, + qk_head_dim=self.qk_head_dim, + v_head_dim=self.v_head_dim, + kv_b_proj=self.kv_b_proj, + indexer=None, + ) + self.q_pad_num_heads = getattr(self.impl, "q_pad_num_heads", None) + + aphrodite_config = get_current_aphrodite_config() + parallel_config = aphrodite_config.parallel_config + assert ( + parallel_config.decode_context_parallel_size <= 1 and parallel_config.prefill_context_parallel_size <= 1 + ), "Kimi-K3 MultiHeadLatentAttention does not support context parallelism." + self.prefill_backend = get_mla_prefill_backend(aphrodite_config)( + num_heads=self.num_local_heads, + scale=self.scale, + kv_lora_rank=self.kv_lora_rank, + qk_nope_head_dim=self.qk_nope_head_dim, + qk_rope_head_dim=self.qk_rope_head_dim, + v_head_dim=self.v_head_dim, + aphrodite_config=aphrodite_config, + ) + + compilation_config = aphrodite_config.compilation_config + if prefix in compilation_config.static_forward_context: + raise ValueError(f"Duplicate layer name: {prefix}") + compilation_config.static_forward_context[prefix] = self + self.kv_cache = torch.tensor([]) + + # ------------------------------------------------------------------ + # AttentionLayerBase interface + # ------------------------------------------------------------------ + def get_attn_backend(self) -> type[AttentionBackend]: + return self.attn_backend + + def get_kv_cache_spec(self, aphrodite_config: AphroditeConfig) -> KVCacheSpec: + kv_cache_dtype = kv_cache_dtype_str_to_dtype(self.kv_cache_dtype, aphrodite_config.model_config) + # TODO: Remove this mypy workaround once the K3 PR is fully merged. + return MLAAttentionSpec( # type: ignore[call-arg] + block_size=aphrodite_config.cache_config.block_size, + num_kv_heads=1, + head_size=self.head_size, + dtype=kv_cache_dtype, + cache_dtype_str=self.kv_cache_dtype, + kv_quant_mode=get_kv_quant_mode(self.kv_cache_dtype), + non_causal_multi_token_decode=self.non_causal_multi_token_decode, + ) + + def process_weights_after_loading(self, act_dtype: torch.dtype) -> None: + """Absorb ``kv_b_proj`` into decode-time ``W_UK_T`` / ``W_UV`` bmm weights. + + ``kv_b_proj`` produces ``[k_nope; v]`` per head from the ``kv_lora_rank`` + latent. For the MQA decode path we pre-split it so that queries are + projected into latent space by ``W_UK_T`` and the attention output is + projected back to ``v`` by ``W_UV`` -- avoiding materializing full K/V. + """ + kv_b_proj_weight = get_and_maybe_dequant_weights(self.kv_b_proj, out_dtype=act_dtype).T + assert kv_b_proj_weight.shape == ( + self.kv_lora_rank, + self.num_local_heads * (self.qk_nope_head_dim + self.v_head_dim), + ), f"{kv_b_proj_weight.shape=}" + kv_b_proj_weight = kv_b_proj_weight.view( + self.kv_lora_rank, + self.num_local_heads, + self.qk_nope_head_dim + self.v_head_dim, + ) + W_UK, W_UV = kv_b_proj_weight.split([self.qk_nope_head_dim, self.v_head_dim], dim=-1) + # (L, N, V) -> (N, L, V) + replace_parameter(self, "W_UV", W_UV.transpose(0, 1), prefer_copy=True) + # (L, N, P) -> (N, P, L) + replace_parameter(self, "W_UK_T", W_UK.permute(1, 2, 0), prefer_copy=True) + + quant_method = self.quant_config.get_quant_method(self, prefix=self.layer_name) if self.quant_config else None + if not should_load_quant_weights(quant_method): + set_default_quant_scales(self, register_buffer=False) + + # Precompute reciprocal scales once here (scales are final after load; + # K3 has no runtime calculate_kv_scales path) so the fp8 fused kernels + # in the decode/prefill hot path take a ready inverse instead of + # launching a per-step reciprocal kernel. + self.register_buffer("_q_scale_inv", self._q_scale.reciprocal().reshape(1), persistent=False) + self.register_buffer("_k_scale_inv", self._k_scale.reciprocal().reshape(1), persistent=False) + + def _v_up_proj(self, x: torch.Tensor, out: torch.Tensor) -> None: + """Project latent attention output back to ``v`` via ``W_UV`` (bmm).""" + # (B, N, L) -> (N, B, L) + x = x.view(-1, self.num_local_heads, self.kv_lora_rank).transpose(0, 1) + out = out.view(-1, self.num_local_heads, self.v_head_dim) + # (N, B, L) x (N, L, V) -> (N, B, V) written transposed into (B, N, V) + torch.bmm(x, self.W_UV, out=out.transpose(0, 1)) + + def _attn_read_kv_cache(self) -> torch.Tensor: + """Latent cache as seen by the attention read kernels (decode / context). + + A plain per-tensor fp8 cache is stored as ``uint8``; view it as fp8 so + the backend reads it as E4M3 rather than fp4/E2M1 -- the latter doubles + the perceived head dim (``head_size * 2``) and fails the kernel's + ``head_dim_k == head_dim_q`` check. Mirrors ``MLAAttention.forward``; + the fp8_ds_mla layout keeps its native uint8 view. + """ + cache = self.kv_cache + if is_quantized_kv_cache(self.kv_cache_dtype) and self.kv_cache_dtype != "fp8_ds_mla": + return cache.view(current_platform.fp8_dtype()) + return cache + + # ------------------------------------------------------------------ + # Forward + # ------------------------------------------------------------------ + def _forward_attn( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + ) -> torch.Tensor: + """Attention front-end: fused qkv-a proj -> norms -> q_b -> attention. + + Returns the pre-gate attention output ``[num_tokens, + num_local_heads * v_head_dim]``. On a profile/dummy run + it returns a zeroed buffer. + """ + if self.q_lora_rank is not None: + qkv_lora = self.fused_qkv_a_proj(hidden_states)[0] + q_c, kv_c, k_pe = qkv_lora.split([self.q_lora_rank, self.kv_lora_rank, self.qk_rope_head_dim], dim=-1) + q_c, kv_c_normed = fused_q_kv_rmsnorm( + q_c, + kv_c, + self.q_a_layernorm.weight.data, + self.kv_a_layernorm.weight.data, + self.rms_norm_eps, + ) + q = self.q_b_proj(q_c)[0].view(-1, self.num_local_heads, self.qk_head_dim) + else: + # Uncompressed query: project directly (no q-LoRA, no q norm) and + # normalize only the kv latent. + q = self.q_proj(hidden_states)[0].view(-1, self.num_local_heads, self.qk_head_dim) + kv_lora = self.kv_a_proj_with_mqa(hidden_states)[0] + kv_c, k_pe = kv_lora.split([self.kv_lora_rank, self.qk_rope_head_dim], dim=-1) + kv_c_normed = self.kv_a_layernorm(kv_c) + k_pe = k_pe.unsqueeze(1) + + attn_out = torch.empty( + (hidden_states.shape[0], self.num_local_heads * self.v_head_dim), + dtype=hidden_states.dtype, + device=hidden_states.device, + ) + self._attention(positions, q, kv_c_normed, k_pe, attn_out) + return attn_out + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + ) -> torch.Tensor: + # Both branches produce (attn_out, gate); they differ only in whether + # the g_proj GEMM is overlapped on the aux stream. + g_proj = self.g_proj + events = self._gate_events + if ( + g_proj is not None + and events is not None + and self.aux_stream is not None + and hidden_states.shape[0] < _GATE_MULTI_STREAM_TOKEN_THRESHOLD + ): + attn_out, gate = maybe_execute_in_parallel( + lambda: self._forward_attn(positions, hidden_states), + lambda: g_proj(hidden_states)[0], + events[0], + events[1], + self.aux_stream, + ) + else: + attn_out = self._forward_attn(positions, hidden_states) + gate = g_proj(hidden_states)[0] if g_proj is not None else None + + if gate is not None: + attn_out = _gate_sigmoid_mul(attn_out, gate) + + # ``o_proj`` (RowParallelLinear + out-of-place all-reduce) returns a + # fresh private tensor, so return it directly rather than copying into a + # caller buffer -- the previous ``output[:] = ...`` convention forced an + # extra [num_tokens, hidden] copy per layer. + return self.o_proj(attn_out)[0] + + @eager_break_during_capture + def _attention( + self, + positions: torch.Tensor, + q: torch.Tensor, + kv_c_normed: torch.Tensor, + k_pe: torch.Tensor, + attn_out: torch.Tensor, + ) -> None: + forward_context = get_forward_context() + attn_metadata_by_layer = forward_context.attn_metadata + if attn_metadata_by_layer is None: + attn_out.zero_() + return + assert isinstance(attn_metadata_by_layer, dict) + attn_metadata = cast("MLACommonMetadata", attn_metadata_by_layer[self.layer_name]) + + num_actual_toks = attn_metadata.num_actual_tokens + slot_mapping_by_layer = forward_context.slot_mapping + assert isinstance(slot_mapping_by_layer, dict) + slot_mapping = slot_mapping_by_layer[self.layer_name] + + q = q[:num_actual_toks] + kv_c_normed = kv_c_normed[:num_actual_toks] + k_pe = k_pe[:num_actual_toks] + positions = positions[:num_actual_toks] + attn_out = attn_out[:num_actual_toks] + + cos_sin_cache = None + rope_positions = None + if self.rotary_emb is not None: + # Pass the fp32 cos/sin table straight to the fused epilogue (it reads + # fp32 and does the RoPE math in fp32) -- no per-forward dtype cast. + cos_sin_cache = self.rotary_emb.cos_sin_cache + rope_positions = positions + + # Decode tokens are laid out first, prefill tokens after. The fused + # prefill covers every supported config (bf16 / plain-fp8 / + # fp8_ds_mla), so there is no dense-MHA (forward_mha) fallback. + num_mqa_tokens = attn_metadata.num_decode_tokens + num_mha_tokens = q.size(0) - num_mqa_tokens + + # Both the prefill and decode fused epilogues write their own cache + # slice, so there is no separate do_kv_cache_update. + + # ---- Prefill: fused key-concat + cache-insert + attention ---- + if num_mha_tokens > 0: + self._forward_prefill_fused( + q[num_mqa_tokens:], + kv_c_normed[num_mqa_tokens:], + k_pe[num_mqa_tokens:], + rope_positions[num_mqa_tokens:] if rope_positions is not None else None, + cos_sin_cache, + slot_mapping[num_mqa_tokens:num_actual_toks], + attn_metadata, + attn_out[num_mqa_tokens:], + ) + + # ---- Decode: latent multi-query attention ---- + if num_mqa_tokens > 0: + mqa_q_nope, mqa_q_pe = q[:num_mqa_tokens].split([self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1) + # BMM1: absorb q_nope into latent space. (N,B,P) x (N,P,L) -> (B,N,L) + ql_nope = torch.bmm(mqa_q_nope.transpose(0, 1), self.W_UK_T).transpose(0, 1) + # Fused: concat mqa_q = [ql_nope | q_pe] and insert the decode-token + # latent into the paged cache (one launch, right before forward_mqa). + mqa_q = self._decode_concat_cache( + ql_nope, + mqa_q_pe, + kv_c_normed[:num_mqa_tokens], + k_pe[:num_mqa_tokens], + rope_positions[:num_mqa_tokens] if rope_positions is not None else None, + cos_sin_cache, + slot_mapping[:num_mqa_tokens], + ) + latent_out, _lse = self.impl.forward_mqa( # type: ignore[attr-defined] + mqa_q, self._attn_read_kv_cache(), attn_metadata, self + ) + self._v_up_proj(latent_out, out=attn_out[:num_mqa_tokens]) + + def _decode_concat_cache( + self, + ql_nope: torch.Tensor, + q_pe: torch.Tensor, + kv_c_normed: torch.Tensor, + k_pe: torch.Tensor, + positions: torch.Tensor | None, + cos_sin_cache: torch.Tensor | None, + slot_mapping: torch.Tensor, + ) -> torch.Tensor: + """Fused decode query-concat + latent cache insert, dispatched by cache + dtype (same policy as prefill: fp8 cache -> fp8 query).""" + if self.kv_cache_dtype == "fp8_ds_mla": + cache = self.kv_cache + if cache.dtype != torch.uint8: + cache = cache.view(torch.uint8) + return fused_mla_decode_q_concat_kv_cache_insert( + ql_nope, + q_pe, + kv_c_normed, + k_pe, + cache, + slot_mapping, + ds_mla=True, + positions=positions, + cos_sin_cache=cos_sin_cache, + ) + if is_quantized_kv_cache(self.kv_cache_dtype): + assert self.impl.supports_quant_query_input, ( # type: ignore[attr-defined] + "Kimi-K3 fp8 KV cache decode requires a backend that accepts an fp8 (quantized) query input." + ) + cache = self.kv_cache + if cache.dtype != torch.float8_e4m3fn: + cache = cache.view(torch.float8_e4m3fn) + return fused_mla_decode_q_concat_kv_cache_insert( + ql_nope, + q_pe, + kv_c_normed, + k_pe, + cache, + slot_mapping, + q_scale_inv=self._q_scale_inv, + cache_scale_inv=self._k_scale_inv, + positions=positions, + cos_sin_cache=cos_sin_cache, + ) + return fused_mla_decode_q_concat_kv_cache_insert( + ql_nope, + q_pe, + kv_c_normed, + k_pe, + self.kv_cache, + slot_mapping, + positions=positions, + cos_sin_cache=cos_sin_cache, + ) + + def _forward_prefill_fused( + self, + q: torch.Tensor, + kv_c_normed: torch.Tensor, + k_pe: torch.Tensor, + positions: torch.Tensor | None, + cos_sin_cache: torch.Tensor | None, + slot_mapping: torch.Tensor, + attn_metadata, + out: torch.Tensor, + ) -> None: + """Prefill using the fused key-concat + cache-insert kernel. + + Replaces ``_concat_k_nope_k_pe`` and the prefill cache write with one + fused kernel launch, dispatched by cache dtype. The chunked context + gather + online-softmax merge are delegated to the impl. + + Supported configs (K3 fp8 policy): + - bf16 cache -> bf16 prefill query + - plain fp8 cache -> fp8 prefill query (unscaled q/k/v; cache _k_scale) + - fp8_ds_mla cache -> bf16 prefill query (656B per-tile self-scaled) + """ + prefill = attn_metadata.prefill + has_context = prefill.chunked_context is not None + fp8_prefill = prefill.q_data_type == current_platform.fp8_dtype() + + kv_nope = self.kv_b_proj(kv_c_normed)[0].view(-1, self.num_local_heads, self.qk_nope_head_dim + self.v_head_dim) + k_nope, v = kv_nope.split([self.qk_nope_head_dim, self.v_head_dim], dim=-1) + + if self.kv_cache_dtype == "fp8_ds_mla": + # fp8_ds_mla cache (656B, per-tile self-scaled); bf16 attention. + assert not fp8_prefill, ( + "Kimi-K3 fp8_ds_mla uses a bf16 prefill query; fp8 prefill query is not supported with fp8_ds_mla." + ) + kv_cache = self.kv_cache + if kv_cache.dtype != torch.uint8: + kv_cache = kv_cache.view(torch.uint8) + k = fused_mla_key_concat_ds_mla_insert( + q, + k_nope, + k_pe, + kv_c_normed, + kv_cache, + slot_mapping, + positions, + cos_sin_cache, + ) + elif is_quantized_kv_cache(self.kv_cache_dtype): + assert fp8_prefill, ( + "Kimi-K3 fp8 KV cache requires an fp8 prefill query; enable " + "--attention-config '{\"use_prefill_query_quantization\": true}'." + ) + # Plain per-tensor fp8: quant q/k/v (unscaled, matching forward_mha's + # unscaled `.to(fp8)`) and insert the fp8 latent (scaled by _k_scale). + kv_cache = self.kv_cache + if kv_cache.dtype != torch.float8_e4m3fn: + kv_cache = kv_cache.view(torch.float8_e4m3fn) + q, k, v = fused_mla_qkv_quant_kv_cache_fp8_insert( + q, + k_nope, + k_pe, + kv_c_normed, + v, + kv_cache, + slot_mapping, + self._one_scale, + self._one_scale, + self._one_scale, + self._k_scale_inv, + positions, + cos_sin_cache, + ) + else: + # Concat full K = [k_nope | k_pe] and insert [kv_c_normed | k_pe] + # into the paged cache for these prefill tokens, in one launch. + k = fused_mla_key_concat_kv_cache_insert( + q, + k_nope, + k_pe, + kv_c_normed, + self.kv_cache, + slot_mapping, + positions, + cos_sin_cache, + ) + + # When there is no chunked context, backends that honor `out` write the + # attention result straight into it, avoiding a slice+flatten+copy. + writes_out = not has_context and prefill.prefill_backend.supports_out() + output_prefill = prefill.prefill_backend.run_prefill_new_tokens( + q=q, + k=k, + v=v, + return_softmax_lse=has_context, + out=(out.view(-1, self.num_local_heads, self.v_head_dim) if writes_out else None), + ) + + if has_context: + context_output, context_lse = self.impl._compute_prefill_context( # type: ignore[attr-defined] + q, self._attn_read_kv_cache(), attn_metadata, self._k_scale + ) + suffix_output, suffix_lse = output_prefill + out = out.view(-1, self.num_local_heads, self.v_head_dim) + merge_attn_states( + output=out, + prefix_output=context_output[..., : self.v_head_dim], + prefix_lse=context_lse, + suffix_output=suffix_output[..., : self.v_head_dim], + suffix_lse=suffix_lse, + prefill_tokens_with_context=prefill.chunked_context.prefill_tokens_with_context, + ) + elif not writes_out: + out.copy_(output_prefill[..., : self.v_head_dim].flatten(start_dim=-2)) diff --git a/aphrodite/models/kimi_k3/nvidia/model.py b/aphrodite/models/kimi_k3/nvidia/model.py new file mode 100644 index 0000000000..5d802a48db --- /dev/null +++ b/aphrodite/models/kimi_k3/nvidia/model.py @@ -0,0 +1,1722 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Kimi-K3 multimodal model implementation for vLLM.""" + +import math +from collections.abc import Iterable +from typing import Any, cast + +import torch +from torch import nn + +import aphrodite.envs as envs +from aphrodite.config import AphroditeConfig +from aphrodite.distributed import ( + get_ep_group, + get_pp_group, + get_tensor_model_parallel_world_size, +) +from aphrodite.forward_context import get_forward_context, is_forward_context_available +from aphrodite.logger import init_logger +from aphrodite.model_executor.layers.activation import SiluAndMul, SituAndMul +from aphrodite.model_executor.layers.fused_moe import ( + FusedMoE, + fused_moe_make_expert_params_mapping, +) +from aphrodite.model_executor.layers.fused_moe.router.base_router import ( + eplb_map_to_physical_and_record, +) +from aphrodite.model_executor.layers.fused_moe.router.gate_linear import GateLinear +from aphrodite.model_executor.layers.fused_moe.router.grouped_topk_router import ( + fused_grouped_topk, +) +from aphrodite.model_executor.layers.fused_moe.runner.latent_moe_runner import ( + LatentMoERunner, +) +from aphrodite.model_executor.layers.layernorm import RMSNorm +from aphrodite.model_executor.layers.linear import ( + MergedColumnParallelLinear, + ReplicatedLinear, + RowParallelLinear, +) +from aphrodite.model_executor.layers.logits_processor import LogitsProcessor +from aphrodite.model_executor.layers.mamba.gdn.kimi_gdn_linear_attn import ( + KimiGatedDeltaNetAttention as KimiLinearGatedDeltaNetAttention, +) +from aphrodite.model_executor.layers.mamba.mamba_utils import ( + MambaStateCopyFunc, + MambaStateCopyFuncCalculator, + MambaStateDtypeCalculator, + MambaStateShapeCalculator, +) +from aphrodite.model_executor.layers.quantization import QuantizationConfig +from aphrodite.model_executor.layers.quantization.compressed_tensors import ( + compressed_tensors, +) +from aphrodite.model_executor.layers.vocab_parallel_embedding import ( + ParallelLMHead, + VocabParallelEmbedding, +) +from aphrodite.model_executor.model_loader.weight_utils import ( + default_weight_loader, + maybe_remap_kv_scale_name, +) +from aphrodite.model_executor.models.interfaces import ( + EagleModelMixin, + HasInnerState, + IsHybrid, + MixtureOfExperts, + SupportsEagle3, + SupportsEncoderCudaGraph, + SupportsMultiModal, + SupportsPP, + SupportsQuant, +) +from aphrodite.model_executor.models.kimi_k25 import KimiK25MediaPixelInputs +from aphrodite.model_executor.models.kimi_k25_vit import ( + KimiK25MultiModalProjector, + MoonViT3dPretrainedModel, + vision_tower_forward, +) +from aphrodite.model_executor.models.utils import ( + AutoWeightsLoader, + PPMissingLayer, + WeightsMapper, + init_vllm_registered_model, + is_pp_missing_parameter, + make_layers, + maybe_prefix, +) +from aphrodite.model_executor.models.vision import is_vit_use_data_parallel +from aphrodite.models.deepseek_v4.nvidia.model import DeepseekV4MegaMoEExperts +from aphrodite.models.deepseek_v4.nvidia.ops.prepare_megamoe import prepare_megamoe_inputs +from aphrodite.models.kimi_k3.nvidia.kda import KimiK3DeltaAttention +from aphrodite.models.kimi_k3.nvidia.low_latency_gemm import ( + enable_kimi_k3_low_latency_gemm, +) +from aphrodite.models.kimi_k3.nvidia.mla import MultiHeadLatentAttention +from aphrodite.models.kimi_k3.nvidia.ops import attn_res +from aphrodite.models.kimi_k3.nvidia.ops.sequence_parallel import ( + sp_all_gather, + sp_padding_mask, + sp_reduce_scatter, + sp_shard, +) +from aphrodite.multimodal import MULTIMODAL_REGISTRY +from aphrodite.multimodal.inputs import NestedTensors +from aphrodite.platforms import current_platform +from aphrodite.sequence import IntermediateTensors +from aphrodite.transformers_utils.configs.kimi_k3 import KimiK3Config +from aphrodite.transformers_utils.configs.kimi_linear import KimiLinearConfig +from aphrodite.utils.math_utils import cdiv +from aphrodite.utils.multi_stream_utils import maybe_execute_in_parallel +from aphrodite.utils.torch_utils import aux_stream +from aphrodite.v1.worker.ubatching import dbo_current_ubatch_id + +from ..common.mm_preprocess import ( + KimiK3DummyInputsBuilder, + KimiK3MultiModalProcessor, + KimiK3ProcessingInfo, +) + +logger = init_logger(__name__) + + +class KimiMLP(nn.Module): + def __init__( + self, + hidden_size: int, + intermediate_size: int, + hidden_act: str, + quant_config: QuantizationConfig | None = None, + reduce_results: bool = True, + use_sequence_parallel: bool = False, + prefix: str = "", + activation_situ_beta: float | None = None, + activation_situ_linear_beta: float | None = None, + ) -> None: + super().__init__() + + self.gate_up_proj = MergedColumnParallelLinear( + hidden_size, + [intermediate_size] * 2, + bias=False, + quant_config=quant_config, + disable_tp=use_sequence_parallel, + prefix=f"{prefix}.gate_up_proj", + ) + self.down_proj = RowParallelLinear( + intermediate_size, + hidden_size, + bias=False, + quant_config=quant_config, + reduce_results=reduce_results, + disable_tp=use_sequence_parallel, + prefix=f"{prefix}.down_proj", + ) + if hidden_act == "silu": + self.act_fn = SiluAndMul() + elif hidden_act == "situ": + self.act_fn = SituAndMul( + beta=activation_situ_beta or 1.0, + linear_beta=activation_situ_linear_beta, + ) + else: + raise ValueError(f"Unsupported activation: {hidden_act}. Only silu and situ are supported.") + + def forward(self, x): + gate_up, _ = self.gate_up_proj(x) + x = self.act_fn(gate_up) + x, _ = self.down_proj(x) + return x + + +class KimiRoutedOutputTransform(nn.Module): + def __init__( + self, + norm: RMSNorm | None, + up_proj: ReplicatedLinear, + ) -> None: + super().__init__() + self.norm = norm + self.up_proj = up_proj + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + if self.norm is not None: + hidden_states = self.norm(hidden_states) + hidden_states, _ = self.up_proj(hidden_states) + return hidden_states + + +class KimiK3MegaMoEExperts(DeepseekV4MegaMoEExperts): + """Kimi K3 adapter for the DeepGEMM MegaMoE kernel.""" + + _kimi_symm_buffer_cache: dict[tuple[object, ...], object] = {} + _synchronized_ep_groups: set[tuple[int, int]] = set() + _transformed_l1_weights: tuple[torch.Tensor, torch.Tensor] | None + _transformed_l2_weights: tuple[torch.Tensor, torch.Tensor] | None + w13_weight: nn.Parameter | None + w13_weight_scale: nn.Parameter | None + w2_weight: nn.Parameter | None + w2_weight_scale: nn.Parameter | None + + def __init__( + self, + *args, + activation: str, + activation_beta: float | None, + activation_linear_beta: float | None, + **kwargs, + ): + super().__init__(*args, **kwargs) + self.activation = activation + self.activation_beta = activation_beta + self.activation_linear_beta = activation_linear_beta + + def synchronize_first_launch(self) -> None: + ep_group = get_ep_group() + device = torch.accelerator.current_device_index() + key = (id(ep_group.cpu_group), device) + if key in self._synchronized_ep_groups: + return + torch.accelerator.synchronize() + torch.distributed.barrier(group=ep_group.cpu_group) + self._synchronized_ep_groups.add(key) + + def finalize_weights(self) -> None: + if self._transformed_l1_weights is not None: + return + + self._check_runtime_supported() + assert self.w13_weight is not None + assert self.w13_weight_scale is not None + assert self.w2_weight is not None + assert self.w2_weight_scale is not None + from aphrodite.utils.deep_gemm import _import_deep_gemm + + deep_gemm = _import_deep_gemm() + w13_scale = deep_gemm.transform_sf_into_required_layout( + self._ue8m0_uint8_to_float(self.w13_weight_scale.data).contiguous(), + 2 * self.intermediate_size, + self.hidden_size, + (1, 32), + self.num_local_experts, + ) + w2_scale = deep_gemm.transform_sf_into_required_layout( + self._ue8m0_uint8_to_float(self.w2_weight_scale.data).contiguous(), + self.hidden_size, + self.intermediate_size, + (1, 32), + self.num_local_experts, + ) + self._transformed_l1_weights, self._transformed_l2_weights = deep_gemm.transform_weights_for_mega_moe( + (self.w13_weight.data.view(torch.int8).contiguous(), w13_scale), + (self.w2_weight.data.view(torch.int8).contiguous(), w2_scale), + activation=self.activation, + ) + self.w13_weight = None + self.w13_weight_scale = None + self.w2_weight = None + self.w2_weight_scale = None + + def get_symm_buffer(self): + from aphrodite.utils.deep_gemm import _import_deep_gemm + + deep_gemm = _import_deep_gemm() + group = get_ep_group().device_group + device = torch.accelerator.current_device_index() + key = ( + id(group), + device, + self.num_experts, + self.max_num_tokens, + self.top_k, + self.hidden_size, + self.intermediate_size, + self.activation, + ) + symm_buffer = self._kimi_symm_buffer_cache.get(key) + if symm_buffer is None: + symm_buffer = deep_gemm.get_symm_buffer_for_mega_moe( + group, + self.num_experts, + self.max_num_tokens, + self.top_k, + self.hidden_size, + self.intermediate_size, + activation=self.activation, + ) + self._kimi_symm_buffer_cache[key] = symm_buffer + return symm_buffer + + def forward( + self, + hidden_states: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + *, + activation_clamp: float | None, + fast_math: bool = True, + ) -> torch.Tensor: + self.synchronize_first_launch() + if hidden_states.shape[0] > self.max_num_tokens: + raise ValueError( + f"Kimi K3 MegaMoE got {hidden_states.shape[0]} tokens, " + f"but its symmetric buffer supports {self.max_num_tokens}." + ) + y = torch.empty_like(hidden_states, dtype=torch.bfloat16) + from aphrodite.utils.deep_gemm import _import_deep_gemm + + deep_gemm = _import_deep_gemm() + symm_buffer = self.get_symm_buffer() + num_tokens = hidden_states.shape[0] + is_padding = None + if envs.APHRODITE_MOE_SKIP_PADDING and is_forward_context_available(): + is_padding = get_forward_context().is_padding + if is_padding is not None: + is_padding = is_padding[:num_tokens] + + eplb_state = self.eplb_state + if eplb_state.logical_to_physical_map is not None: + assert eplb_state.expert_load_view is not None + assert eplb_state.logical_replica_count is not None + assert eplb_state.should_record_tensor is not None + if is_padding is not None: + topk_ids = torch.where(is_padding.unsqueeze(1), -1, topk_ids) + topk_ids = eplb_map_to_physical_and_record( + topk_ids=topk_ids, + expert_load_view=eplb_state.expert_load_view, + logical_to_physical_map=eplb_state.logical_to_physical_map, + logical_replica_count=eplb_state.logical_replica_count, + record_enabled=eplb_state.should_record_tensor, + num_unpadded_tokens=eplb_state.num_unpadded_tokens_tensors[dbo_current_ubatch_id()] + if eplb_state.num_unpadded_tokens_tensors is not None + else None, + ) + + prepare_megamoe_inputs( + hidden_states, + topk_weights, + topk_ids, + symm_buffer.x[:num_tokens], + symm_buffer.x_sf[:num_tokens], + symm_buffer.topk_idx[:num_tokens], + symm_buffer.topk_weights[:num_tokens], + is_padding=is_padding, + ) + self.finalize_weights() + assert self._transformed_l1_weights is not None + assert self._transformed_l2_weights is not None + deep_gemm.fp8_fp4_mega_moe( + y, + self._transformed_l1_weights, + self._transformed_l2_weights, + symm_buffer, + activation_clamp=activation_clamp, + activation=self.activation, + activation_beta=self.activation_beta, + activation_linear_beta=self.activation_linear_beta, + fast_math=fast_math, + ) + return y + + +def make_kimi_k3_mega_moe_expert_params_mapping( + num_experts: int, +) -> list[tuple[str, str, int, str]]: + mapping = [] + for expert_id in range(num_experts): + for shard_id in ("w1", "w2", "w3"): + param_prefix = "w13" if shard_id in ("w1", "w3") else "w2" + for suffix in ("weight_packed", "weight_scale"): + param_suffix = "weight" if suffix == "weight_packed" else suffix + mapping.append( + ( + f"experts.{param_prefix}_{param_suffix}", + f"experts.{expert_id}.{shard_id}.{suffix}", + expert_id, + shard_id, + ) + ) + return mapping + + +class KimiMoE(nn.Module): + def __init__( + self, + config: KimiLinearConfig, + aphrodite_config: AphroditeConfig, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + layer_idx: int = 0, + use_sequence_parallel: bool = False, + ): + super().__init__() + hidden_size = config.hidden_size + moe_intermediate_size = config.moe_intermediate_size + num_experts = config.num_experts + num_experts_per_token = config.num_experts_per_token + assert moe_intermediate_size is not None + assert num_experts is not None + assert num_experts_per_token is not None + moe_renormalize = config.moe_renormalize + routed_expert_hidden_size = config.routed_expert_hidden_size + self.use_latent_moe = routed_expert_hidden_size is not None + self.moe_hidden_size = routed_expert_hidden_size if routed_expert_hidden_size is not None else hidden_size + self.latent_moe_use_norm = config.latent_moe_use_norm + self.tp_size = get_tensor_model_parallel_world_size() + self.routed_scaling_factor = config.routed_scaling_factor + self.moe_renormalize = moe_renormalize + self.use_grouped_topk = config.use_grouped_topk + self.num_expert_group = config.num_expert_group + self.topk_group = config.topk_group + self.moe_router_activation_func = config.moe_router_activation_func + self.num_shared_experts = config.num_shared_experts + self.layer_idx = layer_idx + self.use_mega_moe = aphrodite_config.kernel_config.moe_backend == "deep_gemm_mega_moe" + if self.use_mega_moe and not aphrodite_config.parallel_config.enable_expert_parallel: + raise NotImplementedError( + "Kimi K3 MegaMoE requires expert parallel. Enable it with --enable-expert-parallel." + ) + if self.use_mega_moe and config.hidden_act != "situ": + raise ValueError("Kimi K3 MegaMoE requires SITU activation.") + if self.use_mega_moe and not self.use_latent_moe: + raise ValueError("Kimi K3 MegaMoE requires latent MoE projections.") + if self.use_mega_moe and not self.use_grouped_topk: + raise ValueError("Kimi K3 MegaMoE requires grouped top-k routing.") + if self.use_mega_moe and (self.num_expert_group != 1 or self.topk_group != 1): + raise NotImplementedError("Kimi K3 MegaMoE currently requires one expert group.") + self.padded_moe_intermediate_size = moe_intermediate_size + min_moe_intermediate_per_partition = getattr(config, "min_moe_intermediate_per_partition", 256) + if self.tp_size > 1: + moe_intermediate_per_partition = moe_intermediate_size // self.tp_size + if moe_intermediate_per_partition < min_moe_intermediate_per_partition: + self.padded_moe_intermediate_size = min_moe_intermediate_per_partition * self.tp_size + activation_situ_beta = config.activation_situ_beta if config.hidden_act == "situ" else None + activation_situ_linear_beta = config.activation_situ_linear_beta if config.hidden_act == "situ" else None + + # Route with fp32 logits for numerically stable expert selection. + self.gate = GateLinear( + input_size=hidden_size, + output_size=num_experts, + bias=False, + out_dtype=torch.float32, + prefix=f"{prefix}.gate", + ) + + self.gate.e_score_correction_bias = nn.Parameter(torch.empty(num_experts, dtype=torch.float32)) + + self.shared_experts: KimiMLP | None + if self.num_shared_experts is not None: + shared_intermediate_size = moe_intermediate_size * self.num_shared_experts + self.shared_experts = KimiMLP( + hidden_size=config.hidden_size, + intermediate_size=shared_intermediate_size, + hidden_act=config.hidden_act, + quant_config=quant_config, + reduce_results=False, + use_sequence_parallel=use_sequence_parallel, + prefix=f"{prefix}.shared_experts", + activation_situ_beta=activation_situ_beta, + activation_situ_linear_beta=activation_situ_linear_beta, + ) + else: + self.shared_experts = None + + self.routed_expert_down_proj: ReplicatedLinear | None + self.routed_expert_norm: RMSNorm | None + self.routed_expert_up_proj: ReplicatedLinear | None + self.routed_output_transform: KimiRoutedOutputTransform | None + if self.use_latent_moe: + self.routed_expert_down_proj = ReplicatedLinear( + hidden_size, + self.moe_hidden_size, + bias=False, + quant_config=None, + prefix=f"{prefix}.routed_expert_down_proj", + ) + self.routed_expert_norm = ( + RMSNorm(self.moe_hidden_size, eps=config.rms_norm_eps) if self.latent_moe_use_norm else None + ) + # Replicated up-proj: the full weight lives on every rank and + # produces the full hidden dim locally. This lets LatentMoERunner + # fuse the latent and shared reductions into a single all-reduce + # (concat the two partials, reduce once), then run the up-proj and + # shared add locally with no further collective. + self.routed_expert_up_proj = ReplicatedLinear( + self.moe_hidden_size, + hidden_size, + bias=False, + quant_config=None, + prefix=f"{prefix}.routed_expert_up_proj", + ) + + self.routed_output_transform = KimiRoutedOutputTransform( + self.routed_expert_norm, self.routed_expert_up_proj + ) + # Auxiliary CUDA stream to overlap the router gate with the routed + # down projection on decode-sized batches (gated by + # APHRODITE_ROUTED_DOWN_PROJ_STREAM_TOKEN_THRESHOLD). + self._down_proj_stream: torch.cuda.Stream | None = aux_stream() + self._down_proj_events = (torch.cuda.Event(), torch.cuda.Event()) + else: + self.routed_expert_down_proj = None + self.routed_expert_norm = None + self.routed_expert_up_proj = None + self.routed_output_transform = None + + if self.use_mega_moe: + ep_group = get_ep_group() + ep_size = ep_group.world_size + ep_rank = ep_group.rank_in_group + if num_experts % ep_size != 0: + raise ValueError(f"Kimi K3 num_experts={num_experts} must be divisible by EP size {ep_size}.") + num_local_experts = num_experts // ep_size + self.experts = KimiK3MegaMoEExperts( + aphrodite_config, + num_experts=num_experts, + num_local_experts=num_local_experts, + experts_start_idx=ep_rank * num_local_experts, + top_k=num_experts_per_token, + hidden_size=self.moe_hidden_size, + intermediate_size=self.padded_moe_intermediate_size, + prefix=f"{prefix}.experts", + activation="situ", + activation_beta=activation_situ_beta, + activation_linear_beta=activation_situ_linear_beta, + ) + else: + enable_tail_fusion = envs.APHRODITE_ENABLE_K3_LATENT_MOE_TAIL_FUSION + self.experts = FusedMoE( + shared_experts=self.shared_experts, + num_experts=num_experts, + top_k=num_experts_per_token, + hidden_size=self.moe_hidden_size, + intermediate_size=self.padded_moe_intermediate_size, + activation=config.hidden_act, + activation_situ_beta=activation_situ_beta, + activation_situ_linear_beta=activation_situ_linear_beta, + renormalize=moe_renormalize, + quant_config=quant_config, + use_grouped_topk=config.use_grouped_topk, + num_expert_group=config.num_expert_group, + topk_group=config.topk_group, + prefix=f"{prefix}.experts", + scoring_func=config.moe_router_activation_func, + e_score_correction_bias=self.gate.e_score_correction_bias, + routed_scaling_factor=self.routed_scaling_factor, + # Down projection runs outside FusedMoE so it can overlap the + # router gate on the aux stream (see forward()); the original + # hidden states are passed to forward() as shared_experts_input + # so shared experts still see the untransformed input. + routed_input_transform=None, + routed_output_transform=self.routed_output_transform, + is_sequence_parallel=use_sequence_parallel, + runner_cls=LatentMoERunner if self.use_latent_moe else None, + runner_args=({"enable_k3_latent_moe_tail_fusion": enable_tail_fusion} if self.use_latent_moe else None), + ) + if self.padded_moe_intermediate_size != moe_intermediate_size: + w13_weight = getattr(self.experts, "w13_weight", None) + if w13_weight is None: + w13_weight = getattr(self.experts, "w13_weight_packed", None) + w2_weight = getattr(self.experts, "w2_weight", None) + if w2_weight is None: + w2_weight = getattr(self.experts, "w2_weight_packed", None) + if w13_weight is not None: + w13_weight.data.zero_() + if w2_weight is not None: + w2_weight.data.zero_() + self.experts.moe_config.intermediate_size_per_partition_unpadded = moe_intermediate_size // self.tp_size + + def _maybe_overlap_router_and_down_proj( + self, hidden_states: torch.Tensor + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]: + """Compute the routed-expert down projection alongside the router, + overlapping them on separate CUDA streams when latent MoE is enabled. + + The router gate and the down projection both read ``hidden_states``, so + the gate runs on the default stream and the down projection on the aux + stream, joined via ``maybe_execute_in_parallel``. For MegaMoE the + grouped top-k selection consumes only the gate logits, so it also runs + on the default stream and overlaps the down projection. + + Returns: + ``(routed_hidden_states, router_output, topk_ids)``. + ``routed_hidden_states`` is the down-projected latent (or the + original ``hidden_states`` when latent MoE is disabled). For MegaMoE + ``router_output`` holds the grouped top-k weights and ``topk_ids`` + the selected experts; otherwise ``router_output`` holds the raw gate + logits and ``topk_ids`` is ``None``. + """ + + def _router( + hidden_states: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor | None]: + router_logits, _ = self.gate(hidden_states) + if not self.use_mega_moe: + return router_logits, None + return fused_grouped_topk( + hidden_states=hidden_states, + gating_output=router_logits, + topk=self.experts.top_k, + renormalize=self.moe_renormalize, + e_score_correction_bias=self.gate.e_score_correction_bias.data, + num_expert_group=self.num_expert_group, + topk_group=self.topk_group, + scoring_func=self.moe_router_activation_func, + routed_scaling_factor=self.routed_scaling_factor, + ) + + down_proj = self.routed_expert_down_proj + if down_proj is None: + router_output, topk_ids = _router(hidden_states) + return hidden_states, router_output, topk_ids + num_tokens = hidden_states.shape[0] + (router_output, topk_ids), (routed_hidden_states, _) = maybe_execute_in_parallel( + lambda: _router(hidden_states), + lambda: down_proj(hidden_states), + self._down_proj_events[0], + self._down_proj_events[1], + self._down_proj_stream if num_tokens <= envs.APHRODITE_ROUTED_DOWN_PROJ_STREAM_TOKEN_THRESHOLD else None, + ) + return routed_hidden_states, router_output, topk_ids + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + num_tokens, hidden_size = hidden_states.shape + hidden_states = hidden_states.view(-1, hidden_size) + # Overlap the gate with the routed down projection; the returned hidden + # states are already down-projected. Keep the original ``hidden_states`` + # for the shared experts. + routed_hidden_states, router_output, topk_ids = self._maybe_overlap_router_and_down_proj(hidden_states) + if self.use_mega_moe: + assert self.routed_output_transform is not None + assert topk_ids is not None + final_hidden_states = self.experts( + routed_hidden_states, + router_output, + topk_ids, + activation_clamp=None, + ) + final_hidden_states = self.routed_output_transform(final_hidden_states) + if self.shared_experts is not None: + final_hidden_states = final_hidden_states + self.shared_experts(hidden_states) + else: + # Routed experts consume the down-projected latent; shared experts + # (inside FusedMoE) get the original hidden states via + # shared_experts_input. + final_hidden_states = self.experts( + hidden_states=routed_hidden_states, + router_logits=router_output, + shared_experts_input=hidden_states, + ) + return final_hidden_states.view(num_tokens, hidden_size) + + +class KimiDecoderLayer(nn.Module): + def __init__( + self, + config: KimiLinearConfig, + aphrodite_config: AphroditeConfig, + prefix: str = "", + aux_stream: torch.cuda.Stream | None = None, + ) -> None: + super().__init__() + self.hidden_size = config.hidden_size + self.layer_idx = int(prefix.rsplit(".", 1)[1]) + + self.is_moe = config.is_moe + layer_idx = self.layer_idx + cache_config = aphrodite_config.cache_config + quant_config = aphrodite_config.quant_config + parallel_config = aphrodite_config.parallel_config + self.is_moe_layer = ( + self.is_moe + and config.num_experts is not None + and layer_idx >= config.first_k_dense_replace + and layer_idx % config.moe_layer_freq == 0 + ) + + use_mega_moe = aphrodite_config.kernel_config.moe_backend == "deep_gemm_mega_moe" + self.use_sequence_parallel = ( + parallel_config.pipeline_parallel_size == 1 + and parallel_config.enable_expert_parallel + and parallel_config.tensor_parallel_size > 1 + and (use_mega_moe or parallel_config.data_parallel_size > 1) + ) + if config.is_kda_layer(layer_idx): + kda_config = config.linear_attn_config + assert kda_config is not None + # This class also serves standalone Kimi-Linear through the model + # registry. Only Kimi-K3's full-rank gate uses the private KDA path. + if kda_config.get("use_full_rank_gate", False): + self.self_attn = KimiK3DeltaAttention( + config, + aphrodite_config, + prefix=f"{prefix}.self_attn", + ) + self._self_attn_writes_output = False + else: + self.self_attn = KimiLinearGatedDeltaNetAttention( + config, + aphrodite_config, + prefix=f"{prefix}.self_attn", + ) + self._self_attn_writes_output = True + else: + qk_nope_head_dim = config.qk_nope_head_dim + qk_rope_head_dim = config.qk_rope_head_dim + v_head_dim = config.v_head_dim + kv_lora_rank = config.kv_lora_rank + mla_use_nope = config.mla_use_nope + assert qk_nope_head_dim is not None + assert qk_rope_head_dim is not None + assert v_head_dim is not None + assert kv_lora_rank is not None + assert mla_use_nope, "Kimi-K3 MLA (MultiHeadLatentAttention) is NoPE-only" + # q_lora_rank may be None (Kimi-Linear): the MLA layer then uses an + # uncompressed q_proj instead of the fused q-LoRA front-end. + self.self_attn = MultiHeadLatentAttention( + config=config, + hidden_size=self.hidden_size, + num_heads=config.num_attention_heads, + qk_nope_head_dim=qk_nope_head_dim, + qk_rope_head_dim=qk_rope_head_dim, + v_head_dim=v_head_dim, + q_lora_rank=config.q_lora_rank, + kv_lora_rank=kv_lora_rank, + use_output_gate=bool(config.mla_use_output_gate), + cache_config=cache_config, + quant_config=quant_config, + prefix=f"{prefix}.self_attn", + aux_stream=aux_stream, + ) + self._self_attn_writes_output = False + + if self.use_sequence_parallel: + self.self_attn.o_proj.reduce_results = False + + if self.is_moe_layer: + self.block_sparse_moe = KimiMoE( + config=config, + aphrodite_config=aphrodite_config, + quant_config=quant_config, + prefix=f"{prefix}.block_sparse_moe", + layer_idx=layer_idx, + use_sequence_parallel=self.use_sequence_parallel, + ) + self.mlp = self.block_sparse_moe + else: + self.mlp = KimiMLP( + hidden_size=self.hidden_size, + intermediate_size=config.intermediate_size, + hidden_act=config.hidden_act, + quant_config=quant_config, + prefix=f"{prefix}.mlp", + use_sequence_parallel=self.use_sequence_parallel, + activation_situ_beta=config.activation_situ_beta, + activation_situ_linear_beta=config.activation_situ_linear_beta, + ) + self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.post_attention_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + + attn_res_block_size = config.attn_res_block_size + self.use_attn_res = attn_res_block_size is not None + if self.use_attn_res: + assert attn_res_block_size is not None + self.attn_res_block_size = attn_res_block_size + self.is_block_write_layer = layer_idx % self.attn_res_block_size == 0 + self.block_write_idx = layer_idx // self.attn_res_block_size + self.prev_valid_blocks = cdiv(layer_idx, self.attn_res_block_size) + self.self_attention_res_norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.mlp_res_norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.self_attention_res_proj = ReplicatedLinear( + config.hidden_size, + 1, + bias=False, + quant_config=None, + prefix=f"{prefix}.self_attention_res_proj", + ) + self.mlp_res_proj = ReplicatedLinear( + config.hidden_size, + 1, + bias=False, + quant_config=None, + prefix=f"{prefix}.mlp_res_proj", + ) + + def _run_self_attn( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + ) -> torch.Tensor: + if self._self_attn_writes_output: + output = torch.empty_like(hidden_states) + self.self_attn( + hidden_states=hidden_states, + positions=positions, + output=output, + ) + return output + return self.self_attn( + hidden_states=hidden_states, + positions=positions, + ) + + def _pre_attn_norm( + self, + hidden_states: torch.Tensor | None, + residual: torch.Tensor | None, + prefix_sum: torch.Tensor | None, + ) -> tuple[torch.Tensor, torch.Tensor | None, torch.Tensor]: + if not self.use_attn_res: + assert hidden_states is not None + if residual is None: + residual = hidden_states + hidden_states = self.input_layernorm(hidden_states) + else: + hidden_states, residual = self.input_layernorm(hidden_states, residual) + return hidden_states, prefix_sum, residual + + assert prefix_sum is not None + assert residual is not None + hidden_states = attn_res( + prefix_sum, + hidden_states, + residual, + self.self_attention_res_norm.weight, + self.self_attention_res_proj.weight.squeeze(0), + self.input_layernorm.weight, + num_blocks=self.prev_valid_blocks, + block_write_idx=(self.block_write_idx if self.is_block_write_layer else -1), + eps=self.self_attention_res_norm.variance_epsilon, + output_norm_eps=self.input_layernorm.variance_epsilon, + ) + return hidden_states, prefix_sum, residual + + def _post_attn_norm( + self, + hidden_states: torch.Tensor, + residual: torch.Tensor, + prefix_sum: torch.Tensor | None, + ) -> tuple[torch.Tensor, torch.Tensor | None, torch.Tensor]: + if not self.use_attn_res: + hidden_states, residual = self.post_attention_layernorm(hidden_states, residual) + return hidden_states, prefix_sum, residual + + assert prefix_sum is not None + if self.is_block_write_layer: + prefix_sum = hidden_states + prefix_delta = None + else: + prefix_delta = hidden_states + mlp_valid_blocks = self.prev_valid_blocks + self.is_block_write_layer + hidden_states = attn_res( + prefix_sum, + prefix_delta, + residual, + self.mlp_res_norm.weight, + self.mlp_res_proj.weight.squeeze(0), + self.post_attention_layernorm.weight, + num_blocks=mlp_valid_blocks, + block_write_idx=-1, + eps=self.mlp_res_norm.variance_epsilon, + output_norm_eps=self.post_attention_layernorm.variance_epsilon, + ) + return hidden_states, prefix_sum, residual + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor | None, + residual: torch.Tensor | None, + prefix_sum: torch.Tensor | None = None, + **kwargs, + ) -> tuple[torch.Tensor, torch.Tensor | None, torch.Tensor]: + hidden_states, prefix_sum, residual = self._pre_attn_norm(hidden_states, residual, prefix_sum) + assert hidden_states is not None + + if self.use_sequence_parallel: + hidden_states = sp_all_gather(hidden_states) + # Remove SP padding before attention. + hidden_states = hidden_states[: positions.shape[0]] + + # Attention. + hidden_states = self._run_self_attn(positions, hidden_states) + + if self.use_sequence_parallel: + # Add SP padding if needed, and then perform reduce scatter. + hidden_states = sp_reduce_scatter(hidden_states) + + hidden_states, prefix_sum, residual = self._post_attn_norm(hidden_states, residual, prefix_sum) + + # MoE/MLP. + hidden_states = self.mlp(hidden_states) + return hidden_states, prefix_sum, residual + + +class KimiLinearModel(nn.Module, EagleModelMixin): + def __init__(self, *, aphrodite_config: AphroditeConfig, prefix: str = ""): + super().__init__() + + config = aphrodite_config.model_config.hf_text_config + self.config = config + self.attn_res_block_size: int | None = config.attn_res_block_size + self.use_attn_res = self.attn_res_block_size is not None + parallel_config = aphrodite_config.parallel_config + use_mega_moe = aphrodite_config.kernel_config.moe_backend == "deep_gemm_mega_moe" + self.use_sequence_parallel = ( + parallel_config.pipeline_parallel_size == 1 + and parallel_config.enable_expert_parallel + and parallel_config.tensor_parallel_size > 1 + and (use_mega_moe or parallel_config.data_parallel_size > 1) + ) + + self.vocab_size = config.vocab_size + + if get_pp_group().is_first_rank: + self.embed_tokens = VocabParallelEmbedding( + config.vocab_size, + config.hidden_size, + prefix=f"{prefix}.embed_tokens", + ) + else: + self.embed_tokens = PPMissingLayer() + + # Aux stream for overlapping the MLA g_proj output-gate GEMM with the + # attention front-end (DeepseekV4 convention: created at the model + # level and threaded into each attention layer). + aux_stream = torch.cuda.Stream() + + def get_layer(prefix: str): + return KimiDecoderLayer( + config, + aphrodite_config, + prefix, + aux_stream=aux_stream, + ) + + self.start_layer, self.end_layer, self.layers = make_layers( + config.num_hidden_layers, + get_layer, + prefix=f"{prefix}.layers", + ) + self.num_attn_res_blocks = ( + cdiv(self.end_layer, self.attn_res_block_size) if self.attn_res_block_size is not None else 0 + ) + + if get_pp_group().is_last_rank: + self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + if self.use_attn_res: + self.output_attn_res_norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.output_attn_res_proj = ReplicatedLinear( + config.hidden_size, + 1, + bias=False, + quant_config=None, + prefix=f"{prefix}.output_attn_res_proj", + ) + else: + self.norm = PPMissingLayer() + if self.use_attn_res: + self.output_attn_res_norm = PPMissingLayer() + self.output_attn_res_proj = PPMissingLayer() + + world_size = get_tensor_model_parallel_world_size() + assert config.num_attention_heads % world_size == 0, "num_attention_heads must be divisible by world_size" + + def make_empty_intermediate_tensors( + self, + batch_size: int, + dtype: torch.dtype, + device: torch.device, + ) -> IntermediateTensors: + residual_shape: tuple[int, ...] = (batch_size, self.config.hidden_size) + if self.use_attn_res: + assert self.attn_res_block_size is not None + residual_shape = ( + batch_size, + cdiv(self.start_layer, self.attn_res_block_size), + self.config.hidden_size, + ) + return IntermediateTensors( + { + "hidden_states": torch.zeros((batch_size, self.config.hidden_size), dtype=dtype, device=device), + "residual": torch.zeros(residual_shape, dtype=dtype, device=device), + } + ) + + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: + return self.embed_tokens(input_ids) + + def forward( + self, + input_ids: torch.Tensor | None, + positions: torch.Tensor, + intermediate_tensors: IntermediateTensors | None, + inputs_embeds: torch.Tensor | None = None, + **kwargs, + ) -> torch.Tensor | IntermediateTensors | tuple[torch.Tensor, list[torch.Tensor]]: + if get_pp_group().is_first_rank: + if inputs_embeds is not None: + hidden_states = inputs_embeds + else: + hidden_states = self.embed_input_ids(input_ids) + residual = None + else: + assert intermediate_tensors is not None + hidden_states = intermediate_tensors["hidden_states"] + residual = intermediate_tensors["residual"] + assert hidden_states is not None + + aux_hidden_states: list[torch.Tensor] = [] + if self.start_layer in self.aux_hidden_state_layers: + if self.use_attn_res or residual is None: + aux_hidden_states.append(hidden_states) + else: + aux_hidden_states.append(hidden_states + residual) + + full_num_tokens = positions.shape[0] + if self.use_sequence_parallel: + if envs.APHRODITE_MOE_SKIP_PADDING and is_forward_context_available(): + forward_context = get_forward_context() + forward_context.is_padding = sp_padding_mask(forward_context.is_padding, hidden_states) + hidden_states = sp_shard(hidden_states) + assert residual is None, "Currently, SP is not supported with PP" + + prefix_sum = None + if self.use_attn_res: + block_residual = hidden_states.new_empty( + hidden_states.size(0), + self.num_attn_res_blocks, + hidden_states.size(1), + ) + if residual is not None: + block_residual[:, : residual.size(1), :].copy_(residual) + prefix_sum = hidden_states + hidden_states = None + residual = block_residual + + for layer_idx, layer in enumerate( + self.layers[self.start_layer : self.end_layer], + start=self.start_layer, + ): + hidden_states, prefix_sum, residual = layer( + positions=positions, + hidden_states=hidden_states, + prefix_sum=prefix_sum, + residual=residual, + ) + if (layer_idx + 1) in self.aux_hidden_state_layers: + if self.use_attn_res: + assert prefix_sum is not None + aux_hidden_state = prefix_sum + hidden_states + else: + assert residual is not None + aux_hidden_state = hidden_states + residual + + if self.use_sequence_parallel: + # Gather SP-sharded aux hidden states. + # TODO: Optimize this. + aux_hidden_state = sp_all_gather(aux_hidden_state) + aux_hidden_state = aux_hidden_state[:full_num_tokens] + aux_hidden_states.append(aux_hidden_state) + + assert hidden_states is not None + assert residual is not None + if not get_pp_group().is_last_rank: + assert not self.use_sequence_parallel, "Currently, SP is not supported with PP" + if prefix_sum is not None: + hidden_states = hidden_states + prefix_sum + return IntermediateTensors({"hidden_states": hidden_states, "residual": residual}) + + if self.use_attn_res: + assert prefix_sum is not None + hidden_states = attn_res( + prefix_sum, + hidden_states, + residual, + self.output_attn_res_norm.weight, + self.output_attn_res_proj.weight.squeeze(0), + None, + num_blocks=self.num_attn_res_blocks, + block_write_idx=-1, + eps=self.output_attn_res_norm.variance_epsilon, + output_norm_eps=0.0, + ) + else: + hidden_states = hidden_states + residual + + if self.use_sequence_parallel: + # Gather SP-sharded hidden states. + hidden_states = sp_all_gather(hidden_states) + hidden_states = hidden_states[:full_num_tokens] + + # NOTE: the final norm is applied in compute_logits instead of here, so + # the MTP draft model receives the pre-norm hidden states. + if aux_hidden_states: + return hidden_states, aux_hidden_states + return hidden_states + + def load_weights( + self, + weights: Iterable[tuple[str, torch.Tensor] | tuple[str, torch.Tensor, dict[str, Any]]], + ) -> set[str]: + kda_config = self.config.linear_attn_config + use_full_rank_gate = bool(kda_config and kda_config.get("use_full_rank_gate", False)) + beta_shard_id = 5 if use_full_rank_gate else 3 + stacked_params_mapping = [ + # (param_name, shard_name, shard_id) + (".in_proj_qkvgfab", ".q_proj", 0), + (".in_proj_qkvgfab", ".k_proj", 1), + (".in_proj_qkvgfab", ".v_proj", 2), + (".in_proj_qkvgfab", ".b_proj", beta_shard_id), + (".in_proj_qkvgfab", ".f_a_proj", 4), + (".conv1d", ".q_conv1d", 0), + (".conv1d", ".k_conv1d", 1), + (".conv1d", ".v_conv1d", 2), + (".gate_up_proj", ".gate_proj", 0), + (".gate_up_proj", ".up_proj", 1), + ] + if use_full_rank_gate: + stacked_params_mapping.append((".in_proj_qkvgfab", ".g_proj", 3)) + if getattr(self.config, "q_lora_rank", None) is not None: + stacked_params_mapping += [ + (".fused_qkv_a_proj", ".q_a_proj", 0), + (".fused_qkv_a_proj", ".kv_a_proj_with_mqa", 1), + ] + use_mega_moe = any(module.use_mega_moe for module in self.modules() if isinstance(module, KimiMoE)) + if self.config.is_moe and use_mega_moe: + expert_params_mapping = make_kimi_k3_mega_moe_expert_params_mapping(self.config.num_experts) + elif self.config.is_moe: + # Params for weights, fp8 weight scales, fp8 activation scales + # (param_name, weight_name, expert_id, shard_id) + expert_params_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.num_experts, + ) + else: + expert_params_mapping = [] + params_dict = dict(self.named_parameters()) + # Under the MXFP4 quant interface the routed experts register unpacked + # params (``w13_weight``), while the compressed-tensors checkpoint names + # them ``.weight_packed``. Rebind so the expert mapping resolves; scales + # already share the ``.weight_scale`` suffix. + experts_unpacked = not use_mega_moe and not any(n.endswith("w13_weight_packed") for n in params_dict) + loaded_params: set[str] = set() + for args in weights: + name, loaded_weight = args[0], args[1] + kwargs: dict[str, Any] = args[2] if len(args) > 2 else {} + if "rotary_emb.inv_freq" in name: + continue + if experts_unpacked and name.endswith(".weight_packed"): + name = name.replace(".weight_packed", ".weight") + + spec_layer = get_spec_layer_idx_from_weight_name(self.config, name) + if spec_layer is not None: + continue # skip spec decode layers for main model + if "rotary_emb.cos_cached" in name or "rotary_emb.sin_cached" in name: + # Models trained using ColossalAI may include these tensors in + # the checkpoint. Skip them. + continue + for param_name, weight_name, shard_id in stacked_params_mapping: + if weight_name not in name: + continue + # We have mlp.experts[0].gate_proj in the checkpoint. + # Since we handle the experts below in expert_params_mapping, + # we need to skip here BEFORE we update the name, otherwise + # name will be updated to mlp.experts[0].gate_up_proj, which + # will then be updated below in expert_params_mapping + # for mlp.experts[0].gate_gate_up_proj, which breaks load. + if ("mlp.experts." in name) and name not in params_dict: + continue + name_mapped = name.replace(weight_name, param_name) + # Packed projections are only present on compatible layers. + if name_mapped not in params_dict: + continue + name = name_mapped + # Skip loading extra bias for GPTQ models. + if name.endswith(".bias") and name not in params_dict: + continue + if is_pp_missing_parameter(name, self): + continue + param = params_dict[name] + weight_loader = param.weight_loader + weight_loader(param, loaded_weight, shard_id) + break + else: + for ( + expert_param_name, + expert_weight_name, + expert_id, + expert_shard_id, + ) in expert_params_mapping: + if expert_weight_name not in name: + continue + name = name.replace(expert_weight_name, expert_param_name) + if is_pp_missing_parameter(name, self): + continue + param = params_dict[name] + weight_loader = param.weight_loader + weight_loader( + param, + loaded_weight, + name, + expert_id=expert_id, + shard_id=expert_shard_id, + ) + break + else: + # Skip loading extra bias for GPTQ models. + if name.endswith(".bias") and name not in params_dict and not self.config.is_linear_attn: # noqa: E501 + continue + # Remapping the name of FP8 kv-scale. + remapped_name = maybe_remap_kv_scale_name(name, params_dict) + if remapped_name is None: + continue + name = remapped_name + if is_pp_missing_parameter(name, self): + continue + + param = params_dict[name] + weight_loader = getattr(param, "weight_loader", default_weight_loader) + weight_loader(param, loaded_weight, **kwargs) + loaded_params.add(name) + return loaded_params + + def finalize_mega_moe_weights(self) -> None: + for module in self.modules(): + if isinstance(module, KimiMoE) and module.use_mega_moe: + module.experts.finalize_weights() + + +class KimiLinearForCausalLM(nn.Module, HasInnerState, SupportsPP, MixtureOfExperts, IsHybrid): + def __init__(self, *, aphrodite_config: AphroditeConfig, prefix: str = ""): + super().__init__() + self.model_config = aphrodite_config.model_config + self.aphrodite_config = aphrodite_config + self.config = self.model_config.hf_config + quant_config = aphrodite_config.quant_config + self.quant_config = quant_config + self.model = KimiLinearModel(aphrodite_config=aphrodite_config, prefix=maybe_prefix(prefix, "model")) + if get_pp_group().is_last_rank: + self.lm_head = ParallelLMHead( + self.config.vocab_size, + self.config.hidden_size, + quant_config=quant_config, + prefix=maybe_prefix(prefix, "lm_head"), + ) + else: + self.lm_head = PPMissingLayer() + enable_kimi_k3_low_latency_gemm(self, self.model_config.dtype) + logit_scale = getattr(self.config, "logit_scale", 1.0) + self.logits_processor = LogitsProcessor(self.config.vocab_size, scale=logit_scale) + + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: + return self.model.embed_input_ids(input_ids) + + def make_empty_intermediate_tensors( + self, + batch_size: int, + dtype: torch.dtype, + device: torch.device, + ) -> IntermediateTensors: + return self.model.make_empty_intermediate_tensors(batch_size, dtype, device) + + def forward( # type: ignore[override] + self, + input_ids: torch.Tensor | None, + positions: torch.Tensor, + intermediate_tensors: IntermediateTensors | None = None, + inputs_embeds: torch.Tensor | None = None, + **kwargs, + ) -> torch.Tensor | IntermediateTensors | tuple[torch.Tensor, list[torch.Tensor]]: + return self.model(input_ids, positions, intermediate_tensors, inputs_embeds, **kwargs) + + @classmethod + def get_mamba_state_dtype_from_config( + cls, + aphrodite_config: "AphroditeConfig", + ) -> tuple[torch.dtype, torch.dtype]: + return MambaStateDtypeCalculator.kda_state_dtype( + aphrodite_config.model_config.dtype, aphrodite_config.cache_config.mamba_cache_dtype + ) + + @classmethod + def get_mamba_state_shape_from_config( + cls, aphrodite_config: "AphroditeConfig" + ) -> tuple[tuple[int, int], tuple[int, int, int]]: + parallel_config = aphrodite_config.parallel_config + hf_config = aphrodite_config.model_config.hf_config + tp_size = parallel_config.tensor_parallel_size + num_spec = ( + aphrodite_config.speculative_config.num_speculative_tokens if aphrodite_config.speculative_config else 0 + ) + return MambaStateShapeCalculator.kda_state_shape( + tp_size, + hf_config.linear_attn_config["num_heads"], + hf_config.linear_attn_config["head_dim"], + conv_kernel_size=hf_config.linear_attn_config["short_conv_kernel_size"], + num_spec=num_spec, + ) + + @classmethod + def get_mamba_state_copy_func( + cls, + ) -> tuple[MambaStateCopyFunc, MambaStateCopyFunc]: + return MambaStateCopyFuncCalculator.kda_state_copy_func() + + def compute_logits( + self, + hidden_states: torch.Tensor, + ) -> torch.Tensor | None: + # The model's final norm is applied here (not at the end of forward) so + # that the pre-norm hidden states can be fed to the MTP draft model. + hidden_states = self.model.norm(hidden_states, None) + return self.logits_processor(self.lm_head, hidden_states) + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + loader = AutoWeightsLoader( + self, + skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None), + ) + loaded = loader.load_weights(weights) + self.model.finalize_mega_moe_weights() + # The fused MultiHeadLatentAttention's process_weights_after_loading + # (W_UK_T / W_UV absorption) is driven by the loader's generic post-load + # hook for any AttentionLayerBase, so no manual trigger is needed here. + return loaded + + +def get_spec_layer_idx_from_weight_name(config: KimiLinearConfig, weight_name: str) -> int | None: + if hasattr(config, "num_nextn_predict_layers") and (config.num_nextn_predict_layers > 0): + layer_idx = config.num_hidden_layers + for i in range(config.num_nextn_predict_layers): + # Match regardless of the surrounding prefix. The name may arrive as + # ``model.layers.{i}.``, a bare ``layers.{i}.`` (after AutoWeightsLoader + # has stripped the ``model.`` prefix in the main model), or with the + # multimodal ``language_model.model.layers.{i}.`` prefix. + if f"layers.{layer_idx + i}." in weight_name: + return layer_idx + i + return None + + +@MULTIMODAL_REGISTRY.register_processor( + KimiK3MultiModalProcessor, + info=KimiK3ProcessingInfo, + dummy_inputs=KimiK3DummyInputsBuilder, +) +class KimiK3ForConditionalGeneration( + nn.Module, + SupportsMultiModal, + SupportsEncoderCudaGraph, + SupportsPP, + SupportsQuant, + SupportsEagle3, + HasInnerState, + IsHybrid, +): + """Kimi-K3 model with Kimi-K2.5 vision and KimiLinear text.""" + + supports_encoder_tp_data = True + + hf_to_vllm_mapper = WeightsMapper( + orig_to_new_prefix={ + "language_model.layers.": "language_model.model.layers.", + "mm_projector.proj.0": "mm_projector.linear_1", + "mm_projector.proj.2": "mm_projector.linear_2", + } + ) + + @classmethod + def get_placeholder_str(cls, modality: str, i: int) -> str | None: + if modality == "image": + return "<|kimi_image_placeholder|>" + raise ValueError(f"Unsupported modality: {modality}") + + def __init__( + self, + aphrodite_config: AphroditeConfig, + prefix: str = "", + ) -> None: + super().__init__() + model_config = aphrodite_config.model_config + config: KimiK3Config = model_config.hf_config + self.config = config + self.model_config = model_config + quant_config = aphrodite_config.quant_config + + multimodal_config = model_config.multimodal_config + assert multimodal_config is not None + self.use_data_parallel = is_vit_use_data_parallel(config.vision_config.num_attention_heads) + self.hidden_size = config.text_config.hidden_size + self.device = current_platform.current_device() + + with self._mark_tower_model(aphrodite_config, "image"): + self.vision_tower = MoonViT3dPretrainedModel( + config.vision_config, + quant_config=self._maybe_ignore_quant_config(quant_config), + prefix=maybe_prefix(prefix, "vision_tower"), + ) + if self._maybe_ignore_quant_config(quant_config) is not None: + self.vision_tower = self.vision_tower.to(device=self.device) + else: + self.vision_tower = self.vision_tower.to(device=self.device, dtype=model_config.dtype) + + vision_attn = self.vision_tower.encoder.blocks[0].attn + if vision_attn.is_flash_attn_backend and vision_attn._fa_version == 4: + from aphrodite.models.kimi_k3.nvidia.ops.vision_fa4_warmup import ( + KimiK3VisionFA4WarmupConfig, + register_kimi_k3_vision_fa4_warmup, + ) + + merge_height, merge_width = config.vision_config.merge_kernel_size + mm_config = model_config.get_multimodal_config() + assert mm_config is not None + register_kimi_k3_vision_fa4_warmup( + KimiK3VisionFA4WarmupConfig( + num_heads=vision_attn.num_heads, + head_dim=vision_attn.head_size, + dtype=vision_attn.dtype, + max_batch_size=( + aphrodite_config.scheduler_config.max_num_seqs * mm_config.get_limit_per_prompt("image") + ), + max_seqlen=( + aphrodite_config.scheduler_config.max_num_encoder_input_tokens * merge_height * merge_width + ), + ) + ) + + self.mm_projector = KimiK25MultiModalProjector( + config=config.vision_config, + use_data_parallel=self.use_data_parallel, + quant_config=self._maybe_ignore_quant_config(quant_config), + prefix=maybe_prefix(prefix, "mm_projector"), + ) + self.mm_projector = self.mm_projector.to(device=self.device, dtype=model_config.dtype) + + self.quant_config = quant_config + with self._mark_language_model(aphrodite_config): + self.language_model = init_vllm_registered_model( + aphrodite_config=aphrodite_config, + hf_config=config.text_config, + prefix=maybe_prefix(prefix, "language_model"), + architectures=["KimiLinearForCausalLM"], + ) + self.make_empty_intermediate_tensors = ( # type: ignore[method-assign] + self.language_model.make_empty_intermediate_tensors + ) + self.media_placeholder: int = self.config.media_placeholder_token_id + + # -- SupportsEncoderCudaGraph protocol methods -- + + def get_encoder_cudagraph_config(self): + from aphrodite.v1.worker.encoder_cudagraph_defs import EncoderCudaGraphConfig + + return EncoderCudaGraphConfig( + modalities=["image"], + buffer_keys=[ + "pixel_values", + "pos_embeds", + "rope_freqs_cis", + "cu_seqlens", + "max_seqlen", + "sequence_lengths", + "merge_gather_idx", + ], + out_hidden_size=self.hidden_size, + ) + + def get_encoder_cudagraph_budget_range(self, aphrodite_config: AphroditeConfig) -> tuple[int, int]: + min_budget = 64 + max_budget = min( + aphrodite_config.scheduler_config.max_num_batched_tokens, + self.model_config.max_model_len, + ) + return min_budget, max_budget + + @staticmethod + def _get_grid_thws(mm_kwargs: dict[str, Any]) -> list[list[int]]: + grid_thws = mm_kwargs["grid_thws"] + if not isinstance(grid_thws, list): + grid_thws = grid_thws.tolist() + return grid_thws + + @staticmethod + def _get_pixel_values(mm_kwargs: dict[str, Any]) -> torch.Tensor: + pixel_values = mm_kwargs["pixel_values"] + if isinstance(pixel_values, list): + pixel_values = torch.cat(pixel_values) + if pixel_values.ndim in (3, 5): + pixel_values = pixel_values.reshape( + pixel_values.shape[0] * pixel_values.shape[1], + *pixel_values.shape[2:], + ) + return pixel_values + + def get_encoder_cudagraph_item_specs(self, mm_kwargs: dict[str, Any]): + from aphrodite.v1.worker.encoder_cudagraph_defs import EncoderItemSpec + + kh, kw = self.config.vision_config.merge_kernel_size + return [ + EncoderItemSpec( + input_size=t * h * w, + output_tokens=(h // kh) * (w // kw), + ) + for t, h, w in self._get_grid_thws(mm_kwargs) + ] + + def select_encoder_cudagraph_items(self, mm_kwargs: dict[str, Any], indices: list[int]) -> dict[str, Any]: + grid_thws = self._get_grid_thws(mm_kwargs) + pixel_values = self._get_pixel_values(mm_kwargs) + source_grid = mm_kwargs["grid_thws"] + + if not indices: + empty_grid = source_grid[:0] if isinstance(source_grid, torch.Tensor) else [] + return {"pixel_values": pixel_values[:0], "grid_thws": empty_grid} + + patch_counts = [t * h * w for t, h, w in grid_thws] + offsets = [0] + for count in patch_counts: + offsets.append(offsets[-1] + count) + selected_pixel_values = torch.cat([pixel_values[offsets[i] : offsets[i + 1]] for i in indices]) + grid_device = source_grid.device if isinstance(source_grid, torch.Tensor) else None + selected_grid = torch.tensor( + [grid_thws[i] for i in indices], + dtype=torch.long, + device=grid_device, + ) + return {"pixel_values": selected_pixel_values, "grid_thws": selected_grid} + + def prepare_encoder_cudagraph_capture_inputs( + self, + token_budget: int, + max_batch_size: int, + max_frames_per_batch: int, + device: torch.device, + dtype: torch.dtype, + path: str = "default", + ): + from aphrodite.v1.worker.encoder_cudagraph_defs import ( + EncoderCudaGraphCaptureInputs, + ) + + kh, kw = self.config.vision_config.merge_kernel_size + per_item_output = (token_budget + max_batch_size - 1) // max_batch_size + rope = self.vision_tower.encoder.rope_2d + max_output_width = rope.max_width // kw + max_output_height = rope.max_height // kh + output_width = min(math.ceil(math.sqrt(per_item_output)), max_output_width) + output_height = (per_item_output + output_width - 1) // output_width + if output_height > max_output_height: + output_height = max_output_height + output_width = (per_item_output + output_height - 1) // output_height + if output_width > max_output_width: + raise ValueError( + f"Encoder CUDA graph budget {token_budget} exceeds K3 RoPE capacity for max_batch_size={max_batch_size}" + ) + grid_thws = [[1, output_height * kh, output_width * kw] for _ in range(max_batch_size)] + + patch_size: int | tuple[int, int] = self.config.vision_config.patch_size + if isinstance(patch_size, int): + patch_size = (patch_size, patch_size) + total_patches = sum(t * h * w for t, h, w in grid_thws) + pixel_values = torch.randn( + total_patches, + 3, + patch_size[0], + patch_size[1], + device=device, + dtype=dtype, + ) + metadata = self.vision_tower.prepare_encoder_cudagraph_metadata( + grid_thws, + max_batch_size=max_batch_size, + max_seqlen_override=max( + token_budget * kh * kw, + max(t * h * w for t, h, w in grid_thws), + ), + device=device, + ) + return EncoderCudaGraphCaptureInputs(values=metadata | {"pixel_values": pixel_values}) + + def prepare_encoder_cudagraph_replay_buffers( + self, + mm_kwargs: dict[str, Any], + max_batch_size: int, + max_frames_per_batch: int, + path: str = "default", + ): + from aphrodite.v1.worker.encoder_cudagraph_defs import ( + EncoderCudaGraphReplayBuffers, + ) + + pixel_values = self._get_pixel_values(mm_kwargs) + metadata = self.vision_tower.prepare_encoder_cudagraph_metadata( + self._get_grid_thws(mm_kwargs), + max_batch_size=max_batch_size, + device=pixel_values.device, + ) + return EncoderCudaGraphReplayBuffers(values=metadata | {"pixel_values": pixel_values}) + + def _project_encoder_features(self, image_features: torch.Tensor) -> torch.Tensor: + projector_dtype = next(self.mm_projector.parameters()).dtype + if image_features.dtype != projector_dtype: + image_features = image_features.to(projector_dtype) + output = self.mm_projector(image_features) + return output.reshape(-1, output.shape[-1]) + + def encoder_cudagraph_forward( + self, + values: dict[str, torch.Tensor], + path: str = "default", + ) -> torch.Tensor: + pixel_values = values.pop("pixel_values") + image_features = self.vision_tower(pixel_values, None, encoder_metadata=values) + return self._project_encoder_features(image_features) + + def encoder_eager_forward( + self, + mm_kwargs: dict[str, Any], + path: str = "default", + ) -> torch.Tensor: + image_features = self.vision_tower( + self._get_pixel_values(mm_kwargs).to(next(self.vision_tower.parameters()).dtype), + self._get_grid_thws(mm_kwargs), + ) + return self._project_encoder_features(torch.cat(image_features)) + + def _maybe_ignore_quant_config(self, quant_config: QuantizationConfig | None) -> QuantizationConfig | None: + if isinstance(quant_config, compressed_tensors.CompressedTensorsConfig): + return None + return quant_config + + def _parse_and_validate_media_input(self, **kwargs: object) -> KimiK25MediaPixelInputs | None: + pixel_values = kwargs.pop("pixel_values", None) + grid_thws = kwargs.pop("grid_thws", None) + if pixel_values is None: + return None + + if isinstance(pixel_values, list): + pixel_values = torch.cat(cast(list[torch.Tensor], pixel_values), dim=0) + if not isinstance(pixel_values, torch.Tensor): + raise TypeError(f"pixel_values must be a tensor or a list of tensors, got {type(pixel_values)}") + + if len(pixel_values.shape) == 5 or len(pixel_values.shape) == 3: + pixel_values = pixel_values.reshape(pixel_values.shape[0] * pixel_values.shape[1], *pixel_values.shape[2:]) + + target_dtype = next(self.vision_tower.parameters()).dtype + pixel_values = pixel_values.to(target_dtype) + assert isinstance(grid_thws, torch.Tensor), f"expect grid_thws to be a tensor, got {type(grid_thws)}" + grid_thws = grid_thws.reshape(-1, grid_thws.shape[-1]) + assert grid_thws.ndim == 2 and grid_thws.size(1) == 3, f"unexpected shape for grid_thws: {grid_thws.shape}" + + return KimiK25MediaPixelInputs( + type="pixel_values", + pixel_values=pixel_values, + grid_thws=grid_thws, + ) + + def _process_media_input(self, media_input: KimiK25MediaPixelInputs) -> list[torch.Tensor]: + media_features = vision_tower_forward( + self.vision_tower, + media_input["pixel_values"], + media_input["grid_thws"], + mm_projector=self.mm_projector, + use_data_parallel=self.use_data_parallel, + ) + return media_features + + def embed_multimodal(self, **kwargs: object) -> NestedTensors | None: + media_input = self._parse_and_validate_media_input(**kwargs) + if media_input is None: + return None + return self._process_media_input(media_input) + + def forward( # type: ignore[override] + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + intermediate_tensors: IntermediateTensors | None = None, + inputs_embeds: torch.Tensor | None = None, + **kwargs: object, + ) -> torch.Tensor | IntermediateTensors | tuple[torch.Tensor, list[torch.Tensor]]: + if intermediate_tensors is not None: + inputs_embeds = None + return self.language_model( + input_ids=input_ids, + positions=positions, + intermediate_tensors=intermediate_tensors, + inputs_embeds=inputs_embeds, + ) + + def compute_logits(self, hidden_states: torch.Tensor, **kwargs) -> torch.Tensor: + return self.language_model.compute_logits(hidden_states) + + def copy_inputs_before_cuda_graphs(self, input_buffers, **kwargs): + return self.language_model.mamba_cache.copy_inputs_before_cuda_graphs(input_buffers, **kwargs) + + def get_seqlen_agnostic_capture_inputs(self, batch_size: int): + return self.language_model.mamba_cache.get_seqlen_agnostic_capture_inputs(batch_size) + + @classmethod + def get_mamba_state_dtype_from_config(cls, aphrodite_config: AphroditeConfig): + text_config = aphrodite_config.model_config.hf_config.text_config + temp_vllm_config = aphrodite_config.with_hf_config(text_config) + return KimiLinearForCausalLM.get_mamba_state_dtype_from_config(temp_vllm_config) + + @classmethod + def get_mamba_state_shape_from_config(cls, aphrodite_config: AphroditeConfig): + text_config = aphrodite_config.model_config.hf_config.text_config + temp_vllm_config = aphrodite_config.with_hf_config(text_config) + return KimiLinearForCausalLM.get_mamba_state_shape_from_config(temp_vllm_config) + + @classmethod + def get_mamba_state_copy_func(cls): + return KimiLinearForCausalLM.get_mamba_state_copy_func() + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]): + loader = AutoWeightsLoader(self) + return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) diff --git a/aphrodite/models/kimi_k3/nvidia/mtp.py b/aphrodite/models/kimi_k3/nvidia/mtp.py new file mode 100644 index 0000000000..47c7d17d28 --- /dev/null +++ b/aphrodite/models/kimi_k3/nvidia/mtp.py @@ -0,0 +1,417 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Inference-only Kimi-K3 Multi-Token-Prediction (MTP) draft model.""" + +import copy +from collections.abc import Iterable + +import torch +import torch.nn as nn + +import aphrodite.envs as envs +from aphrodite.config import AphroditeConfig +from aphrodite.forward_context import get_forward_context, is_forward_context_available +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.logits_processor import LogitsProcessor +from aphrodite.model_executor.layers.quantization import QuantizationConfig +from aphrodite.model_executor.layers.vocab_parallel_embedding import ( + ParallelLMHead, + VocabParallelEmbedding, +) +from aphrodite.model_executor.model_loader.weight_utils import ( + default_weight_loader, + maybe_remap_kv_scale_name, +) +from aphrodite.model_executor.models.utils import get_pp_missing_layer_names, maybe_prefix +from aphrodite.sequence import IntermediateTensors +from aphrodite.transformers_utils.configs.kimi_linear import KimiLinearConfig + +from ..common.mtp import fused_mtp_input +from .low_latency_gemm import enable_kimi_k3_low_latency_gemm +from .model import ( + KimiDecoderLayer, + KimiMoE, + get_spec_layer_idx_from_weight_name, + make_kimi_k3_mega_moe_expert_params_mapping, +) +from .ops.sequence_parallel import sp_all_gather, sp_padding_mask, sp_shard + +logger = init_logger(__name__) + + +class SharedHead(nn.Module): + def __init__( + self, + config: KimiLinearConfig, + prefix: str, + quant_config: QuantizationConfig | None = None, + ) -> None: + super().__init__() + self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.head = ParallelLMHead( + config.vocab_size, + config.hidden_size, + quant_config=quant_config, + prefix=maybe_prefix(prefix, "head"), + ) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + return self.norm(hidden_states) + + +class KimiK3MultiTokenPredictorLayer(nn.Module): + def __init__( + self, + config: KimiLinearConfig, + aphrodite_config: AphroditeConfig, + prefix: str, + ) -> None: + super().__init__() + self.config = config + quant_config = aphrodite_config.quant_config + + self.enorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.hnorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.eh_proj = nn.Linear(config.hidden_size * 2, config.hidden_size, bias=False) + + self.shared_head = SharedHead(config=config, prefix=prefix, quant_config=quant_config) + # The MTP block starts without the base model's AttnRes state. + block_config = copy.copy(config) + block_config.attn_res_block_size = None + # NOTE: the prefix must end in the numeric spec-layer index so that + # KimiDecoderLayer can parse ``layer_idx`` and pick MLA (full attn). + # Own aux stream for the MLA g_proj output-gate overlap (DeepseekV4 + # convention: the MTP block creates its own stream). + aux_stream = torch.cuda.Stream() + self.mtp_block = KimiDecoderLayer(block_config, aphrodite_config, prefix=prefix, aux_stream=aux_stream) + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + previous_hidden_states: torch.Tensor, + inputs_embeds: torch.Tensor | None = None, + spec_step_index: int = 0, + ) -> tuple[torch.Tensor, torch.Tensor]: + assert inputs_embeds is not None + hidden_states = self.eh_proj( + fused_mtp_input( + positions, + inputs_embeds, + previous_hidden_states, + self.enorm.weight, + self.hnorm.weight, + self.enorm.variance_epsilon, + ) + ) + if self.mtp_block.use_sequence_parallel: + if envs.APHRODITE_MOE_SKIP_PADDING and is_forward_context_available(): + forward_context = get_forward_context() + forward_context.is_padding = sp_padding_mask(forward_context.is_padding, hidden_states) + hidden_states = sp_shard(hidden_states) + + hidden_states, _, residual = self.mtp_block( + positions=positions, + hidden_states=hidden_states, + residual=None, + ) + if self.mtp_block.use_sequence_parallel: + assert residual is not None + hidden_states = hidden_states + residual + hidden_states = sp_all_gather(hidden_states)[: positions.shape[0]] + logits_hidden_states = self.shared_head.norm(hidden_states) + return logits_hidden_states, hidden_states + + # Produce the normalized logits input and the pre-norm recurrent state + # in one fused add-RMSNorm launch. + logits_hidden_states, hidden_states = self.shared_head.norm(hidden_states, residual) + return logits_hidden_states, hidden_states + + +class KimiK3MultiTokenPredictor(nn.Module): + def __init__(self, *, aphrodite_config: AphroditeConfig, prefix: str = ""): + super().__init__() + config: KimiLinearConfig = aphrodite_config.model_config.hf_text_config + self.config = config + self.mtp_start_layer_idx = config.num_hidden_layers + self.num_mtp_layers = config.num_nextn_predict_layers + + self.layers = torch.nn.ModuleDict( + { + str(idx): KimiK3MultiTokenPredictorLayer(config, aphrodite_config, f"{prefix}.layers.{idx}") + for idx in range( + self.mtp_start_layer_idx, + self.mtp_start_layer_idx + self.num_mtp_layers, + ) + } + ) + self.embed_tokens = VocabParallelEmbedding( + config.vocab_size, + config.hidden_size, + prefix=maybe_prefix(prefix, "embed_tokens"), + ) + self.logits_processor = LogitsProcessor(config.vocab_size) + + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: + return self.embed_tokens(input_ids) + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + previous_hidden_states: torch.Tensor, + inputs_embeds: torch.Tensor | None = None, + spec_step_idx: int = 0, + ) -> tuple[torch.Tensor, torch.Tensor]: + if inputs_embeds is None: + inputs_embeds = self.embed_tokens(input_ids) + current_step_idx = spec_step_idx % self.num_mtp_layers + return self.layers[str(self.mtp_start_layer_idx + current_step_idx)]( + input_ids, + positions, + previous_hidden_states, + inputs_embeds, + current_step_idx, + ) + + def compute_logits( + self, + hidden_states: torch.Tensor, + spec_step_idx: int = 0, + ) -> torch.Tensor: + current_step_idx = spec_step_idx % self.num_mtp_layers + mtp_layer = self.layers[str(self.mtp_start_layer_idx + current_step_idx)] + logits = self.logits_processor(mtp_layer.shared_head.head, hidden_states) + return logits + + +class KimiK3MTP(nn.Module): + def __init__(self, *, aphrodite_config: AphroditeConfig, prefix: str = ""): + super().__init__() + self.config = aphrodite_config.model_config.hf_text_config + self.quant_config = aphrodite_config.quant_config + self.model = KimiK3MultiTokenPredictor(aphrodite_config=aphrodite_config, prefix=maybe_prefix(prefix, "model")) + enable_kimi_k3_low_latency_gemm(self, aphrodite_config.model_config.dtype) + + 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, + hidden_states: torch.Tensor, + intermediate_tensors: IntermediateTensors | None = None, + inputs_embeds: torch.Tensor | None = None, + spec_step_idx: int = 0, + ) -> tuple[torch.Tensor, torch.Tensor]: + return self.model( + input_ids, + positions, + hidden_states, + inputs_embeds, + spec_step_idx, + ) + + def compute_logits( + self, + hidden_states: torch.Tensor, + spec_step_idx: int = 0, + ) -> torch.Tensor | None: + return self.model.compute_logits(hidden_states, spec_step_idx) + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + # Mirror KimiLinearForCausalLM.load_weights naming: leading-dot shard + # names, q_lora-conditional fused QKV, and w1/w2/w3 expert weights. + kda_config = self.config.linear_attn_config + use_full_rank_gate = bool(kda_config and kda_config.get("use_full_rank_gate", False)) + beta_shard_id = 5 if use_full_rank_gate else 3 + stacked_params_mapping = [ + # (param_name, shard_name, shard_id) + (".in_proj_qkvgfab", ".q_proj", 0), + (".in_proj_qkvgfab", ".k_proj", 1), + (".in_proj_qkvgfab", ".v_proj", 2), + (".in_proj_qkvgfab", ".b_proj", beta_shard_id), + (".in_proj_qkvgfab", ".f_a_proj", 4), + (".conv1d", ".q_conv1d", 0), + (".conv1d", ".k_conv1d", 1), + (".conv1d", ".v_conv1d", 2), + (".gate_up_proj", ".gate_proj", 0), + (".gate_up_proj", ".up_proj", 1), + ] + if use_full_rank_gate: + stacked_params_mapping.append((".in_proj_qkvgfab", ".g_proj", 3)) + if getattr(self.config, "q_lora_rank", None) is not None: + stacked_params_mapping += [ + (".fused_qkv_a_proj", ".q_a_proj", 0), + (".fused_qkv_a_proj", ".kv_a_proj_with_mqa", 1), + ] + + use_mega_moe = any(module.use_mega_moe for module in self.modules() if isinstance(module, KimiMoE)) + if self.config.is_moe and use_mega_moe: + expert_params_mapping = make_kimi_k3_mega_moe_expert_params_mapping(self.config.num_experts) + elif self.config.is_moe: + expert_params_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.num_experts, + ) + else: + expert_params_mapping = [] + + pp_missing_layer_names = get_pp_missing_layer_names(self) + params_dict = dict(self.named_parameters()) + # Under the MXFP4 quant interface the routed experts register unpacked + # params (``w13_weight``), while the compressed-tensors checkpoint names + # them ``.weight_packed``. Rebind so the expert mapping resolves; scales + # already share the ``.weight_scale`` suffix. + experts_unpacked = not use_mega_moe and not any(n.endswith("w13_weight_packed") for n in params_dict) + loaded_params: set[str] = set() + for name, loaded_weight in weights: + if "rotary_emb.inv_freq" in name: + continue + # The multimodal checkpoint prefixes text weights with + # ``language_model.``; strip it so names match this draft model's + # parameter paths (``model.layers.{i}.``). Non-text weights + # (vision_tower, mm_projector, ...) never match a spec layer below. + if name.startswith("language_model."): + name = name[len("language_model.") :] + if experts_unpacked and name.endswith(".weight_packed"): + name = name.replace(".weight_packed", ".weight") + spec_layer = get_spec_layer_idx_from_weight_name(self.config, name) + if spec_layer is None: + continue + name = self._rewrite_spec_layer_name(spec_layer, name) + + for param_name, weight_name, shard_id in stacked_params_mapping: + if weight_name not in name: + continue + # Routed experts (``.experts.{i}.w1/w2/w3``) are handled by the + # expert mapping below; skip them here. Shared experts + # (``.shared_experts.``) use gate/up_proj and fall through. + if ".experts." in name: + continue + name_mapped = name.replace(weight_name, param_name) + # Only take this mapping if the fused destination actually + # exists (e.g. QKV fusion is only present when q_lora is used). + if name_mapped not in params_dict: + continue + if name_mapped in pp_missing_layer_names: + continue + name = name_mapped + param = params_dict[name] + weight_loader = param.weight_loader + weight_loader(param, loaded_weight, shard_id) + break + else: + for ( + expert_param_name, + expert_weight_name, + expert_id, + expert_shard_id, + ) in expert_params_mapping: + if expert_weight_name not in name: + continue + name_mapped = name.replace(expert_weight_name, expert_param_name) + if name_mapped in pp_missing_layer_names: + continue + param = params_dict[name_mapped] + weight_loader = param.weight_loader + weight_loader( + param, + loaded_weight, + name_mapped, + shard_id=expert_shard_id, + expert_id=expert_id, + ) + name = name_mapped + break + else: + if name.endswith(".bias") and name not in params_dict: + continue + remapped_name = maybe_remap_kv_scale_name(name, params_dict) + if remapped_name is None: + continue + name = remapped_name + + # The embedding is shared across MTP layers; only the first + # spec layer carries the hoisted (non-".layers") copy. + if spec_layer != self.model.mtp_start_layer_idx and (".layers" not in name): + continue + if name in pp_missing_layer_names: + continue + # The base model uses an attn-residual scheme whose per-layer + # weights (self_attention_res_*, mlp_res_*) are not used by + # the draft block; such names have no matching parameter and + # are safely skipped. + 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) + loaded_params.add(name) + + # Validate that weights were loaded for each expected MTP layer. + loaded_layers: set[int] = set() + for param_name in loaded_params: + spec_layer = get_spec_layer_idx_from_weight_name(self.config, param_name) + if spec_layer is not None: + loaded_layers.add(spec_layer) + for layer_idx in range( + self.model.mtp_start_layer_idx, + self.model.mtp_start_layer_idx + self.model.num_mtp_layers, + ): + if layer_idx not in loaded_layers: + raise ValueError( + f"MTP speculative decoding layer {layer_idx} weights " + f"missing from checkpoint. The checkpoint may not include " + f"the MTP layer weights. Use a checkpoint that includes " + f"MTP layer weights, or disable speculative decoding." + ) + + if use_mega_moe: + for module in self.modules(): + if isinstance(module, KimiMoE) and module.use_mega_moe: + module.experts.finalize_weights() + + return loaded_params + + def _rewrite_spec_layer_name(self, spec_layer: int, name: str) -> str: + """Rewrite a checkpoint weight name to this module's parameter path. + + Top-level MTP submodules (enorm/hnorm/eh_proj/shared_head) stay under + ``model.layers.{spec_layer}.*``; the shared ``embed_tokens`` is hoisted + to ``model.*``; everything else is a transformer-block weight and gets + ``.mtp_block`` inserted. + """ + spec_layer_weight_names = [ + "embed_tokens", + "enorm", + "hnorm", + "eh_proj", + "shared_head", + ] + shared_weight_names = ["embed_tokens"] + spec_layer_weight = False + shared_weight = False + for weight_name in spec_layer_weight_names: + if weight_name in name: + spec_layer_weight = True + if weight_name in shared_weight_names: + shared_weight = True + break + if not spec_layer_weight: + name = name.replace( + f"model.layers.{spec_layer}.", + f"model.layers.{spec_layer}.mtp_block.", + ) + elif shared_weight: + name = name.replace(f"model.layers.{spec_layer}.", "model.") + return name diff --git a/aphrodite/models/kimi_k3/nvidia/ops/cute_dsl/__init__.py b/aphrodite/models/kimi_k3/nvidia/ops/cute_dsl/__init__.py new file mode 100644 index 0000000000..208f01a7cb --- /dev/null +++ b/aphrodite/models/kimi_k3/nvidia/ops/cute_dsl/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project diff --git a/aphrodite/models/kimi_k3/nvidia/ops/cute_dsl/latent_moe_tail/__init__.py b/aphrodite/models/kimi_k3/nvidia/ops/cute_dsl/latent_moe_tail/__init__.py new file mode 100644 index 0000000000..7eda574b83 --- /dev/null +++ b/aphrodite/models/kimi_k3/nvidia/ops/cute_dsl/latent_moe_tail/__init__.py @@ -0,0 +1,14 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""CuTe DSL kernels for KimiK3LatentMoETailOp.""" + +from .allreduce_rmsnorm_reduce_scatter_early_exit import CollectiveKernel +from .fused_add_multicast_gemm import AdaptiveUpProjectionKernel +from .lamport_copy import LamportCopyKernel + +__all__ = [ + "AdaptiveUpProjectionKernel", + "CollectiveKernel", + "LamportCopyKernel", +] diff --git a/aphrodite/models/kimi_k3/nvidia/ops/cute_dsl/latent_moe_tail/allreduce_rmsnorm_reduce_scatter_early_exit.py b/aphrodite/models/kimi_k3/nvidia/ops/cute_dsl/latent_moe_tail/allreduce_rmsnorm_reduce_scatter_early_exit.py new file mode 100644 index 0000000000..f9361cdd86 --- /dev/null +++ b/aphrodite/models/kimi_k3/nvidia/ops/cute_dsl/latent_moe_tail/allreduce_rmsnorm_reduce_scatter_early_exit.py @@ -0,0 +1,931 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# Specialized from FlashInfer's oneshotAllreduceFusionKernel. + +"""Routed AllReduce/RMSNorm with CTA-specialized ReduceScatter early exit.""" + +from __future__ import annotations + +from typing import Any + +import cuda.bindings.driver as cuda +import cutlass +import cutlass.cute as cute +import torch +import torch.distributed as dist +import torch.distributed._symmetric_memory as symm_mem +from cutlass import BFloat16, Float32, Int32, Int64, Uint32 + +from .primitives import ( + NUM_LAMPORT_BUFFERS, + PACKED_BYTES, + VEC_BF16, + bf16x8_to_packed_u32x4, + block_sum_specialized, + fragment_is_dirty, + load_global_u32x4, + load_volatile_u32, + map_shared_to_peer, + packed_u32x4_to_bf16x8, + red_async_release_gpu_add_u32, + sanitize_negative_zero, + store_global_u32x4, + store_lamport_sentinel_128, + store_shared_cluster_f32, + to_cute, + to_cute_dynamic_m, +) + + +def _mapping( + tp_size: int, + latent_dim: int, + hidden_dim: int, +) -> tuple[int, int, int, int]: + shard_dim = hidden_dim // tp_size + cluster_ctas = latent_dim // shard_dim + threads = shard_dim // VEC_BF16 + shared_roles = tp_size // cluster_ctas + return shard_dim, cluster_ctas, threads, shared_roles + + +def validate_shape(*, tp_size: int, latent_dim: int, hidden_dim: int) -> None: + """Validate constraints imposed by the fused collective mapping.""" + + if tp_size <= 0 or latent_dim <= 0 or hidden_dim <= 0: + raise ValueError("collective dimensions must be positive") + if hidden_dim % tp_size: + raise ValueError("hidden_dim must be divisible by tp_size") + shard, cluster, threads, _ = _mapping(tp_size, latent_dim, hidden_dim) + if shard % VEC_BF16: + raise ValueError("hidden_dim / tp_size must be divisible by 8") + if latent_dim % shard: + raise ValueError("latent_dim must be an integer multiple of shard_dim") + if cluster > 16 or cluster & (cluster - 1): + raise ValueError("collective cluster width must be a power of two <= 16") + if tp_size % cluster: + raise ValueError("tp_size must be divisible by collective cluster width") + if not 32 <= threads <= 1024: + raise ValueError("collective threads per CTA must be in [32, 1024]") + if threads < cluster: + raise ValueError("collective threads must cover every cluster CTA") + + +def _select_routed_schedule( + tp_size: int, + latent_dim: int, + hidden_dim: int, + max_m: int, +) -> tuple[int, int]: + """Match upstream MNNVL's one-cluster-per-token occupancy policy.""" + + _, cluster_ctas, threads, _ = _mapping(tp_size, latent_dim, hidden_dim) + sm_count = torch.cuda.get_device_properties(torch.accelerator.current_device_index()).multi_processor_count + while max_m * cluster_ctas > sm_count and cluster_ctas > 1 and threads <= 512: + cluster_ctas //= 2 + threads *= 2 + return threads, cluster_ctas + + +class AllReduceRMSNormWithReduceScatterEarlyExit: + """One routed role plus one ReduceScatter role per destination group.""" + + def __init__( + self, + *, + rank: int, + tp_size: int, + latent_dim: int, + hidden_dim: int, + max_m: int, + max_token_ctas: int, + fp32_internal: bool = False, + include_reduce_scatter: bool = True, + include_routed: bool = True, + ): + validate_shape( + tp_size=tp_size, + latent_dim=latent_dim, + hidden_dim=hidden_dim, + ) + if not 0 <= rank < tp_size: + raise ValueError(f"rank must be in [0,{tp_size}), got {rank}") + if not include_routed and not include_reduce_scatter: + raise ValueError("at least one collective role must be enabled") + self.rank = rank + self.tp_size = tp_size + self.latent_dim = latent_dim + self.hidden_dim = hidden_dim + ( + self.shard_dim, + mapped_cluster, + mapped_threads, + shared_roles, + ) = _mapping(tp_size, latent_dim, hidden_dim) + if include_reduce_scatter: + self.cluster_ctas = mapped_cluster + self.threads = mapped_threads + else: + self.threads, self.cluster_ctas = _select_routed_schedule(tp_size, latent_dim, hidden_dim, max_m) + # Diagnostic specialization: a one-role grid executes only the routed + # AllReduce/RMSNorm path. Keep this compile-time so the production + # fused path is unchanged when include_reduce_scatter=True. + if include_routed and include_reduce_scatter: + self.roles = 1 + shared_roles + elif include_routed: + self.roles = 1 + else: + self.roles = shared_roles + self.warps = (self.threads + 31) // 32 + self.last_warp_lanes = self.threads - (self.warps - 1) * 32 + self.last_warp_mask = (1 << self.last_warp_lanes) - 1 + # The upstream MNNVL oneshot protocol assigns one cluster to each + # token. Keep that exact ownership in the routed-only diagnostic; + # reusing a cluster for multiple token waves allows the Lamport + # generation metadata to change between waves. + self.token_ctas = min(max_m, max_token_ctas) if include_reduce_scatter else max_m + self.fp32_internal = fp32_internal + self.include_reduce_scatter = include_reduce_scatter + self.include_routed = include_routed + + @cute.jit + def __call__( + self, + latent_source: cute.Tensor, + gamma: cute.Tensor, + latent_output: cute.Tensor, + routed_workspace: cute.Tensor, + latent_flags: cute.Tensor, + latent_multicast_ptr: Int64, + shared_source: cute.Tensor, + shared_output: cute.Tensor, + shared_workspace: cute.Tensor, + shared_flags: cute.Tensor, + shared_peer_ptrs: cute.Tensor, + m: Int32, + epsilon: Float32, + stream: cuda.CUstream, + ): + self.kernel( + latent_source, + gamma, + latent_output, + routed_workspace, + latent_flags, + latent_multicast_ptr, + shared_source, + shared_output, + shared_workspace, + shared_flags, + shared_peer_ptrs, + m, + epsilon, + ).launch( + grid=(self.token_ctas, self.cluster_ctas, self.roles), + block=(self.threads, 1, 1), + cluster=(1, self.cluster_ctas, 1), + smem=(self.warps + self.cluster_ctas) * 4, + stream=stream, + use_pdl=True, + ) + + @cute.kernel + def kernel( + self, + latent_source: cute.Tensor, + gamma: cute.Tensor, + latent_output: cute.Tensor, + routed_workspace: cute.Tensor, + latent_flags: cute.Tensor, + latent_multicast_ptr: Int64, + shared_source: cute.Tensor, + shared_output: cute.Tensor, + shared_workspace: cute.Tensor, + shared_flags: cute.Tensor, + shared_peer_ptrs: cute.Tensor, + m: Int32, + epsilon: Float32, + ): + tidx, _, _ = cute.arch.thread_idx() + token_cta, cta_y, role = cute.arch.block_idx() + logical_role = role + if cutlass.const_expr(not self.include_routed): + # A shared-only grid starts at z=0, while _token_device reserves + # logical role 0 for the routed collective. + logical_role = role + Int32(1) + cluster_rank = cute.arch.make_warp_uniform(cute.arch.block_idx_in_cluster()) + + cute.arch.griddepcontrol_wait() + token = token_cta + while token < m: + self._token_device( + latent_source, + gamma, + latent_output, + routed_workspace, + latent_flags, + latent_multicast_ptr, + shared_source, + shared_output, + shared_workspace, + shared_flags, + shared_peer_ptrs, + m, + epsilon, + token, + token_cta, + cta_y, + logical_role, + cluster_rank, + tidx, + ) + token = token + self.token_ctas + + @cute.jit + def _token_device( + self, + latent_source: cute.Tensor, + gamma: cute.Tensor, + latent_output: cute.Tensor, + routed_workspace: cute.Tensor, + latent_flags: cute.Tensor, + latent_multicast_ptr: Int64, + shared_source: cute.Tensor, + shared_output: cute.Tensor, + shared_workspace: cute.Tensor, + shared_flags: cute.Tensor, + shared_peer_ptrs: cute.Tensor, + m: Int32, + epsilon: Float32, + token: Int32, + token_cta: Int32, + cta_y: Int32, + role: Int32, + cluster_rank: Int32, + tidx: Int32, + ): + if role == 0: + # ---------------- routed AllReduce + RMSNorm ---------------- + packed_idx = cluster_rank * self.threads + tidx + element_offset = Int64(token) * self.latent_dim + Int64(packed_idx) * VEC_BF16 + current_index = cute.arch.load((latent_flags.iterator + 0).llvm_ptr, Uint32) + dirty_index = cute.arch.load((latent_flags.iterator + 1).llvm_ptr, Uint32) + bytes_per_buffer = cute.arch.load((latent_flags.iterator + 2).llvm_ptr, Uint32) + dirty_num_stages = cute.arch.load((latent_flags.iterator + 3).llvm_ptr, Uint32) + bytes_to_clear = cute.arch.load((latent_flags.iterator + 4).llvm_ptr, Uint32) + current_elements = Int64(current_index) * (Int64(bytes_per_buffer) // Int64(2)) + dirty_elements = Int64(dirty_index) * (Int64(bytes_per_buffer) // Int64(2)) + + local_ptr = cute.make_ptr( + BFloat16, + (latent_source.iterator + element_offset).llvm_ptr, + cute.AddressSpace.gmem, + assumed_align=16, + ) + local_packed = sanitize_negative_zero(load_global_u32x4(local_ptr, volatile=False)) + multicast_offset = ( + Int64(current_index) * Int64(bytes_per_buffer) + + ((Int64(token) * self.tp_size + self.rank) * self.latent_dim + Int64(packed_idx) * VEC_BF16) * 2 + ) + store_global_u32x4( + latent_multicast_ptr + multicast_offset, + local_packed, + volatile=False, + ) + + cute.arch.cluster_arrive() + if cluster_rank == 0 and tidx < 32: + cute.arch.cluster_wait() + if tidx == 0: + red_async_release_gpu_add_u32(latent_flags.iterator + 8, Uint32(1)) + + global_tid = (Int64(token) * self.cluster_ctas + Int64(cta_y)) * self.threads + Int64(tidx) + total_threads = Int64(m) * self.cluster_ctas * self.threads + clear_fragments = (Int64(bytes_to_clear) + PACKED_BYTES - 1) // PACKED_BYTES + clear_idx = global_tid + if dirty_num_stages > Uint32(0): + while clear_idx < clear_fragments: + clear_ptr = cute.make_ptr( + BFloat16, + (routed_workspace.iterator + dirty_elements + clear_idx * VEC_BF16).llvm_ptr, + cute.AddressSpace.gmem, + assumed_align=16, + ) + store_lamport_sentinel_128(clear_ptr) + clear_idx = clear_idx + total_threads + + rank_words = cute.make_rmem_tensor(cute.make_layout((self.tp_size, 4), stride=(4, 1)), Uint32) + for word in cutlass.range_constexpr(4): + rank_words[self.rank, word] = local_packed[word] + valid = False + while not valid: + valid = True + for source_rank in cutlass.range_constexpr(self.tp_size): + if cutlass.const_expr(source_rank != self.rank): + remote_element = current_elements + ( + (Int64(token) * self.tp_size + source_rank) * self.latent_dim + Int64(packed_idx) * VEC_BF16 + ) + remote_ptr = cute.make_ptr( + BFloat16, + (routed_workspace.iterator + remote_element).llvm_ptr, + cute.AddressSpace.gmem, + assumed_align=16, + ) + remote = load_global_u32x4(remote_ptr, volatile=True) + for word in cutlass.range_constexpr(4): + rank_words[source_rank, word] = remote[word] + valid = valid & (not fragment_is_dirty(remote)) + + accum = cute.make_rmem_tensor(cute.make_layout((VEC_BF16,)), Float32) + for element in cutlass.range_constexpr(VEC_BF16): + accum[element] = Float32(0.0) + for source_rank in cutlass.range_constexpr(self.tp_size): + values = packed_u32x4_to_bf16x8(rank_words[source_rank, None].load()).to(Float32) + for element in cutlass.range_constexpr(VEC_BF16): + accum[element] = accum[element] + values[element] + + # Preserve the original early PDL point before RMSNorm. + cute.arch.griddepcontrol_launch_dependents() + + if cutlass.const_expr(self.fp32_internal): + # High-precision fused mode: retain the rank reduction in + # FP32 through the RMS square and row reduction. + norm_input = accum.load() + norm_square = norm_input * norm_input + else: + norm_input_bf16 = accum.load().to(BFloat16) + norm_input = norm_input_bf16.to(Float32) + # Upstream-compatible mode: FlashInfer evaluates BF16 * + # BF16 first, then promotes the rounded square to FP32. + norm_square = (norm_input_bf16 * norm_input_bf16).to(Float32) + thread_sum = norm_square.reduce( + cute.ReductionOp.ADD, + init_val=Float32(0.0), + reduction_profile=0, + ) + smem = cutlass.utils.SmemAllocator() + warp_sums = smem.allocate_tensor(Float32, cute.make_layout((self.warps,)), byte_alignment=4) + cluster_sums = smem.allocate_tensor( + Float32, + cute.make_layout((self.cluster_ctas,)), + byte_alignment=4, + ) + block_sum = block_sum_specialized( + thread_sum, + warp_sums, + tidx, + self.warps, + self.last_warp_lanes, + self.last_warp_mask, + ) + if tidx < self.cluster_ctas: + local_slot = cluster_sums.iterator + cluster_rank + remote_slot = map_shared_to_peer(local_slot, Int32(tidx)) + store_shared_cluster_f32(remote_slot, block_sum) + cute.arch.cluster_arrive() + cute.arch.cluster_wait() + + full_sum = Float32(0.0) + for peer in cutlass.range_constexpr(self.cluster_ctas): + full_sum = full_sum + cluster_sums[peer] + inv_rms = cute.math.rsqrt(full_sum / Float32(self.latent_dim) + epsilon, fastmath=True) + gamma_ptr = cute.make_ptr( + BFloat16, + (gamma.iterator + Int64(packed_idx) * VEC_BF16).llvm_ptr, + cute.AddressSpace.gmem, + assumed_align=16, + ) + gamma_values = packed_u32x4_to_bf16x8(load_global_u32x4(gamma_ptr, volatile=False)) + result = (norm_input * inv_rms * gamma_values.to(Float32)).to(BFloat16) + store_global_u32x4( + Int64((latent_output.iterator + element_offset).toint()), + bf16x8_to_packed_u32x4(result), + volatile=False, + ) + + # The x=0 CTA rotates only after reaching its final token wave. + # Waiting for all M arrivals then guarantees every token-wave CTA + # loaded the current generation before the metadata is advanced. + if token_cta == 0 and token + self.token_ctas >= m and cta_y == 0 and tidx == 0: + access_counter = latent_flags.iterator + 8 + arrived = load_volatile_u32(access_counter) + while arrived < Uint32(m): + arrived = load_volatile_u32(access_counter) + next_index = (current_index + Uint32(1)) % Uint32(NUM_LAMPORT_BUFFERS) + actual_bytes = Uint32(m) * Uint32(self.tp_size * self.latent_dim * 2) + cute.arch.store((latent_flags.iterator + 0).llvm_ptr, next_index) + cute.arch.store((latent_flags.iterator + 1).llvm_ptr, current_index) + cute.arch.store( + (latent_flags.iterator + 2).llvm_ptr, + bytes_per_buffer, + ) + cute.arch.store((latent_flags.iterator + 3).llvm_ptr, Uint32(1)) + cute.arch.store((latent_flags.iterator + 4).llvm_ptr, actual_bytes) + for index in cutlass.range_constexpr(5, 8): + cute.arch.store( + (latent_flags.iterator + index).llvm_ptr, + Uint32(0), + ) + cute.arch.store(access_counter.llvm_ptr, Uint32(0)) + + else: + # ---------------- shared ReduceScatter ---------------- + shared_group = role - 1 + destination = shared_group * self.cluster_ctas + cluster_rank + current_index = cute.arch.load((shared_flags.iterator + 0).llvm_ptr, Uint32) + dirty_index = cute.arch.load((shared_flags.iterator + 1).llvm_ptr, Uint32) + bytes_per_buffer = cute.arch.load((shared_flags.iterator + 2).llvm_ptr, Uint32) + dirty_num_stages = cute.arch.load((shared_flags.iterator + 3).llvm_ptr, Uint32) + bytes_to_clear = cute.arch.load((shared_flags.iterator + 4).llvm_ptr, Uint32) + current_elements = Int64(current_index) * (Int64(bytes_per_buffer) // Int64(2)) + dirty_elements = Int64(dirty_index) * (Int64(bytes_per_buffer) // Int64(2)) + + source_element = ( + Int64(token) * self.hidden_dim + Int64(destination) * self.shard_dim + Int64(tidx) * VEC_BF16 + ) + source_ptr = cute.make_ptr( + BFloat16, + (shared_source.iterator + source_element).llvm_ptr, + cute.AddressSpace.gmem, + assumed_align=16, + ) + local_packed = sanitize_negative_zero(load_global_u32x4(source_ptr, volatile=False)) + peer_base = cute.arch.load( + (shared_peer_ptrs.iterator + destination).llvm_ptr, + Int64, + ) + destination_element = current_elements + ( + (Int64(token) * self.tp_size + self.rank) * self.shard_dim + Int64(tidx) * VEC_BF16 + ) + store_global_u32x4( + peer_base + destination_element * 2, + local_packed, + volatile=False, + ) + + # One arrival per shared destination group and token. + cute.arch.cluster_arrive() + if cluster_rank == 0 and tidx < 32: + cute.arch.cluster_wait() + if tidx == 0: + red_async_release_gpu_add_u32(shared_flags.iterator + 8, Uint32(1)) + + global_tid = (Int64(token) * self.tp_size + Int64(destination)) * self.threads + Int64(tidx) + total_threads = Int64(m) * self.tp_size * self.threads + clear_fragments = (Int64(bytes_to_clear) + PACKED_BYTES - 1) // PACKED_BYTES + clear_idx = global_tid + if dirty_num_stages > Uint32(0): + while clear_idx < clear_fragments: + clear_ptr = cute.make_ptr( + BFloat16, + (shared_workspace.iterator + dirty_elements + clear_idx * VEC_BF16).llvm_ptr, + cute.AddressSpace.gmem, + assumed_align=16, + ) + store_lamport_sentinel_128(clear_ptr) + clear_idx = clear_idx + total_threads + + if destination == self.rank: + rank_words = cute.make_rmem_tensor(cute.make_layout((self.tp_size, 4), stride=(4, 1)), Uint32) + valid = False + while not valid: + valid = True + for source_rank in cutlass.range_constexpr(self.tp_size): + remote_element = current_elements + ( + (Int64(token) * self.tp_size + source_rank) * self.shard_dim + Int64(tidx) * VEC_BF16 + ) + remote_ptr = cute.make_ptr( + BFloat16, + (shared_workspace.iterator + remote_element).llvm_ptr, + cute.AddressSpace.gmem, + assumed_align=16, + ) + remote = load_global_u32x4(remote_ptr, volatile=True) + for word in cutlass.range_constexpr(4): + rank_words[source_rank, word] = remote[word] + valid = valid & (not fragment_is_dirty(remote)) + + accum = cute.make_rmem_tensor(cute.make_layout((VEC_BF16,)), Float32) + for element in cutlass.range_constexpr(VEC_BF16): + accum[element] = Float32(0.0) + for source_rank in cutlass.range_constexpr(self.tp_size): + values = packed_u32x4_to_bf16x8(rank_words[source_rank, None].load()).to(Float32) + for element in cutlass.range_constexpr(VEC_BF16): + accum[element] = accum[element] + values[element] + result = accum.load().to(BFloat16) + output_element = Int64(token) * self.hidden_dim + self.rank * self.shard_dim + Int64(tidx) * VEC_BF16 + store_global_u32x4( + Int64((shared_output.iterator + output_element).toint()), + bf16x8_to_packed_u32x4(result), + volatile=False, + ) + + if destination == self.rank: + cute.arch.barrier() + + cute.arch.griddepcontrol_launch_dependents() + + if ( + token_cta == 0 + and token + self.token_ctas >= m + and shared_group == 0 + and cluster_rank == 0 + and tidx == 0 + ): + access_counter = shared_flags.iterator + 8 + arrived = load_volatile_u32(access_counter) + target = Uint32(m) * Uint32(self.tp_size // self.cluster_ctas) + while arrived < target: + arrived = load_volatile_u32(access_counter) + next_index = (current_index + Uint32(1)) % Uint32(NUM_LAMPORT_BUFFERS) + actual_bytes = Uint32(m) * Uint32(self.tp_size * self.shard_dim * 2) + cute.arch.store((shared_flags.iterator + 0).llvm_ptr, next_index) + cute.arch.store((shared_flags.iterator + 1).llvm_ptr, current_index) + cute.arch.store( + (shared_flags.iterator + 2).llvm_ptr, + bytes_per_buffer, + ) + cute.arch.store((shared_flags.iterator + 3).llvm_ptr, Uint32(1)) + cute.arch.store((shared_flags.iterator + 4).llvm_ptr, actual_bytes) + for index in cutlass.range_constexpr(5, 8): + cute.arch.store( + (shared_flags.iterator + index).llvm_ptr, + Uint32(0), + ) + cute.arch.store(access_counter.llvm_ptr, Uint32(0)) + + +_COMPILED: dict[tuple[object, ...], Any] = {} + + +def _routed_workspace_cute(workspace: torch.Tensor): + return to_cute(workspace.view(torch.bfloat16), 16) + + +def _compile_key( + rank: int, + tp_size: int, + latent_dim: int, + hidden_dim: int, + max_m: int, + max_token_ctas: int, + fp32_internal: bool, + include_reduce_scatter: bool, + include_routed: bool, +): + return ( + torch.accelerator.current_device_index(), + rank, + tp_size, + latent_dim, + hidden_dim, + max_m, + max_token_ctas, + fp32_internal, + include_reduce_scatter, + include_routed, + ) + + +def _runtime_args( + latent_source: torch.Tensor, + gamma: torch.Tensor, + latent_output: torch.Tensor, + routed_workspace: torch.Tensor, + routed_flags: torch.Tensor, + routed_multicast_ptr: int, + shared_source: torch.Tensor, + shared_output: torch.Tensor, + shared_workspace: torch.Tensor, + shared_flags: torch.Tensor, + shared_peer_ptrs: torch.Tensor, + rms_eps: float, +): + return ( + to_cute_dynamic_m(latent_source, mode=0, assumed_align=16), + to_cute(gamma, 16), + to_cute(latent_output, 16), + _routed_workspace_cute(routed_workspace), + to_cute(routed_flags, 16), + Int64(routed_multicast_ptr), + to_cute_dynamic_m(shared_source, mode=0, assumed_align=16), + to_cute(shared_output, 16), + to_cute(shared_workspace, 16), + to_cute(shared_flags, 16), + to_cute(shared_peer_ptrs, 16), + Int32(latent_source.shape[0]), + Float32(rms_eps), + cuda.CUstream(torch.cuda.current_stream(latent_source.device).cuda_stream), + ) + + +def compile_kernel( + *, + rank: int, + tp_size: int, + latent_dim: int, + hidden_dim: int, + max_m: int, + max_token_ctas: int, + latent_output: torch.Tensor, + routed_workspace: torch.Tensor, + routed_flags: torch.Tensor, + routed_multicast_ptr: int, + shared_output: torch.Tensor, + shared_workspace: torch.Tensor, + shared_flags: torch.Tensor, + shared_peer_ptrs: torch.Tensor, + rms_eps: float, + fp32_internal: bool, + include_reduce_scatter: bool = True, + include_routed: bool = True, +) -> None: + """Compile the rank/M specialization without retaining caller tensors.""" + + key = _compile_key( + rank, + tp_size, + latent_dim, + hidden_dim, + max_m, + max_token_ctas, + fp32_internal, + include_reduce_scatter, + include_routed, + ) + if key in _COMPILED: + return + device = latent_output.device + latent = torch.empty((max_m, latent_dim), dtype=torch.bfloat16, device=device) + gamma = torch.empty((latent_dim,), dtype=torch.bfloat16, device=device) + shared = torch.empty((max_m, hidden_dim), dtype=torch.bfloat16, device=device) + kernel = AllReduceRMSNormWithReduceScatterEarlyExit( + rank=rank, + tp_size=tp_size, + latent_dim=latent_dim, + hidden_dim=hidden_dim, + max_m=max_m, + max_token_ctas=max_token_ctas, + fp32_internal=fp32_internal, + include_reduce_scatter=include_reduce_scatter, + include_routed=include_routed, + ) + _COMPILED[key] = cute.compile( + kernel, + *_runtime_args( + latent, + gamma, + latent_output, + routed_workspace, + routed_flags, + routed_multicast_ptr, + shared, + shared_output, + shared_workspace, + shared_flags, + shared_peer_ptrs, + rms_eps, + ), + ) + + +def launch( + latent_source: torch.Tensor, + gamma: torch.Tensor, + latent_output: torch.Tensor, + routed_workspace: torch.Tensor, + routed_flags: torch.Tensor, + routed_multicast_ptr: int, + shared_source: torch.Tensor, + shared_output: torch.Tensor, + shared_workspace: torch.Tensor, + shared_flags: torch.Tensor, + shared_peer_ptrs: torch.Tensor, + rms_eps: float, + *, + rank: int, + tp_size: int, + latent_dim: int, + hidden_dim: int, + max_m: int, + max_token_ctas: int, + fp32_internal: bool, + include_reduce_scatter: bool = True, + include_routed: bool = True, +) -> None: + compile_kernel( + rank=rank, + tp_size=tp_size, + latent_dim=latent_dim, + hidden_dim=hidden_dim, + max_m=max_m, + max_token_ctas=max_token_ctas, + latent_output=latent_output, + routed_workspace=routed_workspace, + routed_flags=routed_flags, + routed_multicast_ptr=routed_multicast_ptr, + shared_output=shared_output, + shared_workspace=shared_workspace, + shared_flags=shared_flags, + shared_peer_ptrs=shared_peer_ptrs, + rms_eps=rms_eps, + fp32_internal=fp32_internal, + include_reduce_scatter=include_reduce_scatter, + include_routed=include_routed, + ) + _COMPILED[ + _compile_key( + rank, + tp_size, + latent_dim, + hidden_dim, + max_m, + max_token_ctas, + fp32_internal, + include_reduce_scatter, + include_routed, + ) + ]( + *_runtime_args( + latent_source, + gamma, + latent_output, + routed_workspace, + routed_flags, + routed_multicast_ptr, + shared_source, + shared_output, + shared_workspace, + shared_flags, + shared_peer_ptrs, + rms_eps, + ) + ) + + +class CollectiveKernel: + """Own and launch the routed AllReduce/RMSNorm plus shared ReduceScatter.""" + + def __init__( + self, + *, + group: dist.ProcessGroup, + rank: int, + tp_size: int, + latent_dim: int, + hidden_dim: int, + max_m: int, + max_token_ctas: int, + rms_eps: float, + fp32_internal: bool, + ) -> None: + validate_shape( + tp_size=tp_size, + latent_dim=latent_dim, + hidden_dim=hidden_dim, + ) + self.rank = rank + self.tp_size = tp_size + self.latent_dim = latent_dim + self.hidden_dim = hidden_dim + self.shard_dim = hidden_dim // tp_size + self.max_m = max_m + self.max_token_ctas = max_token_ctas + self.rms_eps = float(rms_eps) + self.fp32_internal = fp32_internal + device = torch.device("cuda", torch.accelerator.current_device_index()) + + bytes_per_routed_buffer = max_m * tp_size * latent_dim * 2 + routed_bytes = NUM_LAMPORT_BUFFERS * bytes_per_routed_buffer + self._routed_workspace = symm_mem.empty( + routed_bytes // 4, + dtype=torch.float32, + device=device, + ) + self._routed_symm_mem = symm_mem.rendezvous(self._routed_workspace, group) + self._routed_workspace.fill_(-0.0) + actual_bytes_per_buffer = self._routed_symm_mem.buffer_size // NUM_LAMPORT_BUFFERS // 16 * 16 + if actual_bytes_per_buffer < bytes_per_routed_buffer: + raise RuntimeError("routed symmetric workspace is too small") + self._routed_flags = torch.tensor( + [0, 2, actual_bytes_per_buffer, 0, 0, 0, 0, 0, 0], + dtype=torch.uint32, + device=device, + ) + routed_multicast_ptr = self._routed_symm_mem.multicast_ptr + if routed_multicast_ptr is None or routed_multicast_ptr == 0: + raise RuntimeError("routed NVLS multicast mapping is unavailable") + self._routed_multicast_ptr = int(routed_multicast_ptr) + + self._latent_output = torch.empty((max_m, latent_dim), dtype=torch.bfloat16, device=device) + self._shared_output = torch.empty((max_m, hidden_dim), dtype=torch.bfloat16, device=device) + shard_start = rank * self.shard_dim + shard_end = shard_start + self.shard_dim + self._shared_shard = self._shared_output[:, shard_start:shard_end] + + self._shared_workspace = symm_mem.empty( + (NUM_LAMPORT_BUFFERS, max_m, tp_size, self.shard_dim), + dtype=torch.bfloat16, + device=device, + ) + self._shared_symm_mem = symm_mem.rendezvous(self._shared_workspace, group) + self._shared_workspace.view(torch.int32).fill_(-0x80000000) + self._shared_flags = torch.zeros(12, dtype=torch.int32, device=device) + self._shared_flags[1] = 1 + self._shared_flags[2] = max_m * tp_size * self.shard_dim * 2 + peer_ptrs = [ + self._shared_symm_mem.get_buffer( + peer, + self._shared_workspace.shape, + torch.bfloat16, + ).data_ptr() + for peer in range(tp_size) + ] + if any(pointer == 0 for pointer in peer_ptrs): + raise RuntimeError("shared LSA peer mapping is unavailable") + self._shared_peer_ptrs = torch.tensor(peer_ptrs, dtype=torch.int64, device=device) + + torch.accelerator.synchronize(device) + dist.barrier(group=group, device_ids=[device.index]) + for owner in range(tp_size): + if rank == owner: + compile_kernel( + rank=rank, + tp_size=tp_size, + latent_dim=latent_dim, + hidden_dim=hidden_dim, + max_m=max_m, + max_token_ctas=max_token_ctas, + latent_output=self._latent_output, + routed_workspace=self._routed_workspace, + routed_flags=self._routed_flags, + routed_multicast_ptr=self._routed_multicast_ptr, + shared_output=self._shared_output, + shared_workspace=self._shared_workspace, + shared_flags=self._shared_flags, + shared_peer_ptrs=self._shared_peer_ptrs, + rms_eps=self.rms_eps, + fp32_internal=fp32_internal, + ) + dist.barrier(group=group, device_ids=[device.index]) + + def __call__( + self, + latent_source: torch.Tensor, + shared_source: torch.Tensor, + gamma: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + if latent_source.ndim != 2 or shared_source.ndim != 2: + raise ValueError("latent_source and shared_source must be rank-2") + m = latent_source.shape[0] + device = self._routed_workspace.device + expected = ( + (latent_source, (m, self.latent_dim), "latent_source"), + (shared_source, (m, self.hidden_dim), "shared_source"), + (gamma, (self.latent_dim,), "gamma"), + ) + for tensor, shape, name in expected: + if ( + tensor.shape != shape + or tensor.dtype != torch.bfloat16 + or tensor.device != device + or not tensor.is_contiguous() + ): + raise ValueError(f"{name} must be contiguous CUDA BF16 {list(shape)}") + if not 1 <= m <= self.max_m: + raise ValueError(f"runtime M={m} must be in [1, {self.max_m}]") + + with torch.accelerator.device_index(device.index): + launch( + latent_source, + gamma, + self._latent_output, + self._routed_workspace, + self._routed_flags, + self._routed_multicast_ptr, + shared_source, + self._shared_output, + self._shared_workspace, + self._shared_flags, + self._shared_peer_ptrs, + self.rms_eps, + rank=self.rank, + tp_size=self.tp_size, + latent_dim=self.latent_dim, + hidden_dim=self.hidden_dim, + max_m=self.max_m, + max_token_ctas=self.max_token_ctas, + fp32_internal=self.fp32_internal, + ) + return ( + self._latent_output[:m], + self._shared_shard, + ) + + @property + def latent_output(self) -> torch.Tensor: + return self._latent_output + + @property + def shared_output(self) -> torch.Tensor: + return self._shared_output diff --git a/aphrodite/models/kimi_k3/nvidia/ops/cute_dsl/latent_moe_tail/fused_add_multicast_gemm.py b/aphrodite/models/kimi_k3/nvidia/ops/cute_dsl/latent_moe_tail/fused_add_multicast_gemm.py new file mode 100644 index 0000000000..b12d6ad01c --- /dev/null +++ b/aphrodite/models/kimi_k3/nvidia/ops/cute_dsl/latent_moe_tail/fused_add_multicast_gemm.py @@ -0,0 +1,1199 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +# Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause + +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: + +# 1. Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. + +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. + +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. + +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +"""Blackwell GEMM with a fused shared-shard add and multicast epilogue. + +Modified from the original CUTLASS CuTe DSL SM100 persistent GEMM tutorial. +The PDL wait before A loading orders the up-projection and shared-expert +inputs after the producer collective. +""" + +import math +from typing import Any + +import cuda.bindings.driver as cuda +import cutlass +import cutlass.cute as cute +import cutlass.pipeline as pipeline +import cutlass.utils as utils +import torch +import torch.distributed as dist +import torch.distributed._symmetric_memory as symm_mem +from cutlass.cute.nvgpu import cpasync, tcgen05 +from cutlass.cute.nvgpu.common import CacheEvictionPriority +from cutlass.cute.runtime import from_dlpack +from cutlass.pipeline import pipeline_init_arrive, pipeline_init_wait + +from .fused_add_multicast_skinny_gemm import ( + FusedAddMulticastSkinnyGemmKernel, +) +from .primitives import CUDAGraphCompatibleWrapper + + +def _as_cute(tensor: torch.Tensor, *, dynamic_m: bool = False): + converted = from_dlpack(CUDAGraphCompatibleWrapper(tensor.detach()), assumed_align=16) + if dynamic_m: + converted = converted.mark_compact_shape_dynamic( + mode=1, + stride_order=tensor.dim_order(), + ) + return converted + + +def validate_configuration( + *, + latent_dim: int, + shard_dim: int, + mma_tiler_mn: tuple[int, int], + cluster_shape_mn: tuple[int, int], + b_prime_stages: int, +) -> None: + """Validate constraints imposed by this BF16 SM100 GEMM.""" + + if latent_dim <= 0 or shard_dim <= 0: + raise ValueError("GEMM K and N must be positive") + if latent_dim % 8 or shard_dim % 8: + raise ValueError("GEMM K and N must be divisible by 8 BF16 values") + if mma_tiler_mn[0] not in (64, 128): + raise ValueError("MMA M tile must be 64 or 128") + if mma_tiler_mn[1] not in range(32, 257, 32): + raise ValueError("MMA N tile must be a multiple of 32 in [32, 256]") + if ( + len(cluster_shape_mn) != 2 + or any(value <= 0 or value & (value - 1) for value in cluster_shape_mn) + or math.prod(cluster_shape_mn) > 16 + ): + raise ValueError("cluster dimensions must be powers of two with product <= 16") + if not 0 <= b_prime_stages <= math.ceil(latent_dim / 128): + raise ValueError("b_prime_stages exceeds the GEMM K-tile count") + + +@cute.jit +def _epilogue_tma_store_add_shared( + gemm_kernel, + epi_tidx: cutlass.Int32, + warp_idx: cutlass.Int32, + tma_atom_c: cute.CopyAtom, + tCtAcc_base: cute.Tensor, + sC: cute.Tensor, + tCgC_base: cute.Tensor, + tCgShared_base: cute.Tensor, + tCcC_base: cute.Tensor, + mC_mnl: cute.Tensor, + mShared_mnl: cute.Tensor, + epi_tile: cute.Tile, + num_tiles_executed: cutlass.Int32, + mma_tile_coord_mnl, + acc_consumer_state: pipeline.PipelineState, + acc_pipeline: pipeline.PipelineAsync, + c_pipeline: pipeline.PipelineTmaStore, +) -> pipeline.PipelineState: + """BF16 GEMM rounding + BF16 shared shard addition, then swizzled S2G TMA.""" + sm100 = utils.gemm.sm100 + tCgC = sm100.transform_partitioned_tensor_layout(tCgC_base) + tCgShared = sm100.transform_partitioned_tensor_layout(tCgShared_base) + tCcC = sm100.transform_partitioned_tensor_layout(tCcC_base) + tCtAcc = sm100.transform_partitioned_tensor_layout(tCtAcc_base) + + tiled_copy_t2r, tTR_tAcc_base, tTR_rAcc = sm100.epilogue_tmem_copy_and_partition( + gemm_kernel, + epi_tidx, + tCtAcc, + tCgC, + epi_tile, + False, + ) + tTR_rC = cute.make_rmem_tensor(tTR_rAcc.shape, gemm_kernel.c_dtype) + tTR_rShared = cute.make_rmem_tensor(tTR_rAcc.shape, gemm_kernel.c_dtype) + tiled_copy_r2s, tRS_rC, tRS_sC = sm100.epilogue_smem_copy_and_partition( + gemm_kernel, tiled_copy_t2r, tTR_rC, epi_tidx, sC + ) + + tCgC_epi = cute.flat_divide(tCgC, epi_tile) + bSG_sC, bSG_gC_partitioned = cpasync.tma_partition( + tma_atom_c, + 0, + cute.make_layout(1), + cute.group_modes(sC, 0, 2), + cute.group_modes(tCgC_epi, 0, 2), + ) + epilog_sync_barrier = pipeline.NamedBarrier( + barrier_id=gemm_kernel.epilog_sync_bar_id, + num_threads=32 * len(gemm_kernel.epilogue_warp_id), + ) + + bSG_gC = bSG_gC_partitioned[(None, None, None, *mma_tile_coord_mnl)] + tTR_tAcc = tTR_tAcc_base[(None, None, None, None, None, acc_consumer_state.index)] + thr_copy_t2r = tiled_copy_t2r.get_slice(epi_tidx) + tTR_gShared_partitioned = thr_copy_t2r.partition_D(cute.flat_divide(tCgShared, epi_tile)) + tTR_cC_partitioned = thr_copy_t2r.partition_D(cute.flat_divide(tCcC, epi_tile)) + tTR_gShared = tTR_gShared_partitioned[(None, None, None, None, None, *mma_tile_coord_mnl)] + tTR_cC = tTR_cC_partitioned[(None, None, None, None, None, *mma_tile_coord_mnl)] + + exemplar = tTR_gShared_partitioned[(None, None, None, 0, 0, 0, 0, 0)] + mcl_r = cute.max_common_layout(tTR_rShared.layout, exemplar.layout) + shared_copy_bits = min( + exemplar.iterator.alignment * 8, + cute.size(mcl_r) * gemm_kernel.c_dtype.width, + 128, + ) + shared_g2r_atom = cute.make_copy_atom( + cute.nvgpu.CopyG2ROp(), + gemm_kernel.c_dtype, + num_bits_per_copy=shared_copy_bits, + l1c_evict_priority=CacheEvictionPriority.NO_ALLOCATE, + ) + + tTR_tAcc = cute.group_modes(tTR_tAcc, 3, cute.rank(tTR_tAcc)) + tTR_gShared = cute.group_modes(tTR_gShared, 3, cute.rank(tTR_gShared)) + tTR_cC = cute.group_modes(tTR_cC, 3, cute.rank(tTR_cC)) + bSG_gC = cute.group_modes(bSG_gC, 1, cute.rank(bSG_gC)) + + subtile_cnt = cute.size(tTR_tAcc.shape, mode=[3]) + previous_subtile_count = num_tiles_executed * subtile_cnt + cute.arch.griddepcontrol_wait() + for subtile_idx in range(subtile_cnt): + tTR_gShared_subtile = tTR_gShared[(None, None, None, subtile_idx)] + tTR_cC_subtile = tTR_cC[(None, None, None, subtile_idx)] + pred_shape = (1, *tTR_cC_subtile.shape[1:]) + pred = cute.make_rmem_tensor(pred_shape, cutlass.Boolean) + for m_idx in range(tTR_cC_subtile.shape[1]): + for n_idx in range(tTR_cC_subtile.shape[2]): + pred[(0, m_idx, n_idx)] = cute.elem_less(tTR_cC_subtile[(0, m_idx, n_idx)], mC_mnl.shape) + tTR_rShared.store(cute.zeros_like(tTR_rShared, dtype=gemm_kernel.c_dtype)) + cute.copy( + shared_g2r_atom, + tTR_gShared_subtile, + tTR_rShared, + pred=pred, + ) + + # Load the shared addend before waiting for the accumulator. + if subtile_idx == 0: + acc_pipeline.consumer_wait(acc_consumer_state) + tTR_tAcc_mn = tTR_tAcc[(None, None, None, subtile_idx)] + cute.copy(tiled_copy_t2r, tTR_tAcc_mn, tTR_rAcc) + + gemm_vec = tiled_copy_r2s.retile(tTR_rAcc).load().to(gemm_kernel.c_dtype) + shared_vec = tiled_copy_r2s.retile(tTR_rShared).load() + fused_vec = (gemm_vec.to(cutlass.Float32) + shared_vec.to(cutlass.Float32)).to(gemm_kernel.c_dtype) + # The symmetric output is an in-band Lamport mailbox whose empty + # marker contains BF16 -0. Normalize either signed zero to +0 so a + # legitimate result can never be mistaken for an unwritten fragment. + fused_vec = cute.where( + fused_vec == cute.zeros_like(fused_vec), + cute.zeros_like(fused_vec), + fused_vec, + ) + tRS_rC.store(fused_vec) + + c_buffer = (previous_subtile_count + subtile_idx) % gemm_kernel.num_c_stage + cute.copy(tiled_copy_r2s, tRS_rC, tRS_sC[(None, None, None, c_buffer)]) + cute.arch.fence_proxy("async.shared", space="cta") + epilog_sync_barrier.arrive_and_wait() + if warp_idx == gemm_kernel.epilogue_warp_id[0]: + cute.copy( + tma_atom_c, + bSG_sC[(None, c_buffer)], + bSG_gC[(None, subtile_idx)], + ) + c_pipeline.producer_commit() + c_pipeline.producer_acquire() + epilog_sync_barrier.arrive_and_wait() + + epilog_sync_barrier.arrive_and_wait() + with cute.arch.elect_one(): + acc_pipeline.consumer_release(acc_consumer_state) + acc_consumer_state.advance() + return acc_consumer_state + + +def _compute_stages( + tiled_mma: cute.TiledMma, + mma_tiler_mnk: tuple[int, int, int], + a_dtype, + b_dtype, + c_dtype, + smem_capacity: int, + c_smem_layout, +) -> tuple[int, int, int]: + """Choose accumulator, mainloop, and epilogue stage counts.""" + num_acc_stage = 2 + num_c_stage = 2 + a_smem_layout_stage_one = utils.sm100.make_smem_layout_a(tiled_mma, mma_tiler_mnk, a_dtype, 1) + b_smem_layout_staged_one = utils.sm100.make_smem_layout_b(tiled_mma, mma_tiler_mnk, b_dtype, 1) + + ab_bytes_per_stage = cute.size_in_bytes(a_dtype, a_smem_layout_stage_one) + cute.size_in_bytes( + b_dtype, b_smem_layout_staged_one + ) + mbar_helpers_bytes = 1024 + + c_bytes_per_stage = cute.size_in_bytes(c_dtype, c_smem_layout) + c_bytes = c_bytes_per_stage * num_c_stage + num_ab_stage = (smem_capacity - (mbar_helpers_bytes + c_bytes)) // ab_bytes_per_stage + num_c_stage += ( + smem_capacity - ab_bytes_per_stage * num_ab_stage - (mbar_helpers_bytes + c_bytes) + ) // c_bytes_per_stage + return num_acc_stage, num_ab_stage, num_c_stage + + +class FusedAddMulticastGemm: + """Persistent Blackwell GEMM with a shared-add epilogue. + + B priming may overlap the producer collective. The PDL wait before A + loading orders both inputs before the shared-add epilogue. + """ + + def __init__( + self, + mma_tiler_mn: tuple[int, int], + cluster_shape_mn: tuple[int, int], + b_prime_stages: int = 2, + ): + self.acc_dtype = cutlass.Float32 + self.cluster_shape_mn = cluster_shape_mn + self.mma_tiler = (*mma_tiler_mn, 1) + # B primes the combined A+B pipeline before the PDL wait. + self.b_prime_stages = b_prime_stages + self.cta_group = tcgen05.CtaGroup.ONE + self.epilogue_warp_id = (0, 1, 2, 3) + self.mma_warp_id = 4 + self.tma_warp_id = 5 + self.threads_per_cta = 32 * len((self.mma_warp_id, self.tma_warp_id, *self.epilogue_warp_id)) + self.epilog_sync_bar_id = 1 + self.tmem_alloc_sync_bar_id = 2 + + def _create_tiled_mma(self): + return utils.sm100.make_trivial_tiled_mma( + self.a_dtype, + self.a_major_mode, + self.b_major_mode, + self.acc_dtype, + self.cta_group, + self.mma_tiler[:2], + ) + + def _setup_attributes(self): + """Derive layouts and stage counts from the compiled tensor shapes.""" + tiled_mma = self._create_tiled_mma() + + mma_inst_shape_k = cute.size(tiled_mma.shape_mnk, mode=[2]) + mma_inst_tile_k = 4 + self.mma_tiler = ( + self.mma_tiler[0], + self.mma_tiler[1], + mma_inst_shape_k * mma_inst_tile_k, + ) + self.cta_tile_shape_mnk = ( + self.mma_tiler[0] // cute.size(tiled_mma.thr_id.shape), + self.mma_tiler[1], + self.mma_tiler[2], + ) + + self.cluster_layout_vmnk = cute.tiled_divide( + cute.make_layout((*self.cluster_shape_mn, 1)), + (tiled_mma.thr_id.shape,), + ) + + self.num_mcast_ctas_a = cute.size(self.cluster_layout_vmnk.shape[2]) + self.num_mcast_ctas_b = cute.size(self.cluster_layout_vmnk.shape[1]) + self.is_a_mcast = self.num_mcast_ctas_a > 1 + self.is_b_mcast = self.num_mcast_ctas_b > 1 + + self.epi_tile = utils.sm100.compute_epilogue_tile_shape( + self.cta_tile_shape_mnk, + False, + self.c_layout, + self.c_dtype, + ) + c_smem_layout = utils.sm100.make_smem_layout_epi(self.c_dtype, self.c_layout, self.epi_tile, 1) + + self.num_acc_stage, self.num_ab_stage, self.num_c_stage = _compute_stages( + tiled_mma, + self.mma_tiler, + self.a_dtype, + self.b_dtype, + self.c_dtype, + utils.get_smem_capacity_in_bytes(), + c_smem_layout, + ) + + self.a_smem_layout_staged = utils.sm100.make_smem_layout_a( + tiled_mma, self.mma_tiler, self.a_dtype, self.num_ab_stage + ) + self.b_smem_layout_staged = utils.sm100.make_smem_layout_b( + tiled_mma, self.mma_tiler, self.b_dtype, self.num_ab_stage + ) + + self.c_smem_layout_staged = utils.sm100.make_smem_layout_epi( + self.c_dtype, self.c_layout, self.epi_tile, self.num_c_stage + ) + + self.num_tmem_alloc_cols = self._compute_num_tmem_alloc_cols( + tiled_mma, self.mma_tiler, self.num_acc_stage, "sm_100" + ) + + @cute.jit + def __call__( + self, + a: cute.Tensor, + b: cute.Tensor, + c: cute.Tensor, + shared_shard: cute.Tensor, + c_multicast_i64: cutlass.Int64, + max_active_clusters: cutlass.Constexpr, + stream: cuda.CUstream, + ): + """Launch the persistent GEMM.""" + # Preserve C's logical strided layout but point the TMA descriptor at + # this rank's shard inside the LSA multicast mapping. One TMA store is + # therefore replicated into the same shard on all eight ranks. + c = cute.make_tensor( + cute.make_ptr( + c.element_type, + c_multicast_i64, + cute.AddressSpace.gmem, + assumed_align=16, + ), + c.layout, + ) + + self.a_dtype: type[cutlass.Numeric] = a.element_type + self.b_dtype: type[cutlass.Numeric] = b.element_type + self.c_dtype: type[cutlass.Numeric] = c.element_type + self.a_major_mode = utils.LayoutEnum.from_tensor(a).mma_major_mode() + self.b_major_mode = utils.LayoutEnum.from_tensor(b).mma_major_mode() + self.c_layout = utils.LayoutEnum.from_tensor(c) + + if cutlass.const_expr(self.a_dtype != self.b_dtype): + raise TypeError(f"Type must match: {self.a_dtype} != {self.b_dtype}") + + tiled_mma = self._create_tiled_mma() + + self._setup_attributes() + if cutlass.const_expr(self.b_prime_stages > self.num_ab_stage): + raise ValueError("b_prime_stages exceeds the compiled A/B pipeline stage count") + + atom_thr_size = cute.size(tiled_mma.thr_id.shape) + + a_op = utils.sm100.cluster_shape_to_tma_atom_A(self.cluster_shape_mn, tiled_mma.thr_id) + a_smem_layout = cute.slice_(self.a_smem_layout_staged, (None, None, None, 0)) + tma_atom_a, tma_tensor_a = cute.nvgpu.make_tiled_tma_atom_A( + a_op, + a, + a_smem_layout, + self.mma_tiler, + tiled_mma, + self.cluster_layout_vmnk.shape, + internal_type=(cutlass.TFloat32 if a.element_type is cutlass.Float32 else None), + ) + + b_op = utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, tiled_mma.thr_id) + b_smem_layout = cute.slice_(self.b_smem_layout_staged, (None, None, None, 0)) + tma_atom_b, tma_tensor_b = cute.nvgpu.make_tiled_tma_atom_B( + b_op, + b, + b_smem_layout, + self.mma_tiler, + tiled_mma, + self.cluster_layout_vmnk.shape, + internal_type=(cutlass.TFloat32 if b.element_type is cutlass.Float32 else None), + ) + + a_copy_size = cute.size_in_bytes(self.a_dtype, a_smem_layout) + b_copy_size = cute.size_in_bytes(self.b_dtype, b_smem_layout) + self.num_tma_load_bytes = (a_copy_size + b_copy_size) * atom_thr_size + + epi_smem_layout = cute.select(self.c_smem_layout_staged, mode=[0, 1]) + tma_atom_c, tma_tensor_c = cpasync.make_tiled_tma_atom( + cpasync.CopyBulkTensorTileS2GOp(), c, epi_smem_layout, self.epi_tile + ) + + tile_sched_params, grid = self._compute_grid( + c, self.cta_tile_shape_mnk, self.cluster_shape_mn, max_active_clusters + ) + self.kernel( + tiled_mma, + tma_atom_a, + tma_tensor_a, + tma_atom_b, + tma_tensor_b, + tma_atom_c, + tma_tensor_c, + self.cluster_layout_vmnk, + self.a_smem_layout_staged, + self.b_smem_layout_staged, + self.c_smem_layout_staged, + self.epi_tile, + tile_sched_params, + shared_shard, + ).launch( + grid=grid, + block=[self.threads_per_cta, 1, 1], + cluster=(*self.cluster_shape_mn, 1), + stream=stream, + use_pdl=True, + ) + return + + @cute.kernel + def kernel( + self, + tiled_mma: cute.TiledMma, + tma_atom_a: cute.CopyAtom, + mA_mkl: cute.Tensor, + tma_atom_b: cute.CopyAtom, + mB_nkl: cute.Tensor, + tma_atom_c: cute.CopyAtom, + mC_mnl: cute.Tensor, + cluster_layout_vmnk: cute.Layout, + a_smem_layout_staged: cute.ComposedLayout, + b_smem_layout_staged: cute.ComposedLayout, + c_smem_layout_staged: cute.Layout | cute.ComposedLayout, + epi_tile: cute.Tile, + tile_sched_params: utils.PersistentTileSchedulerParams, + shared_shard: cute.Tensor, + ): + self._gemm_device( + tiled_mma, + tma_atom_a, + mA_mkl, + tma_atom_b, + mB_nkl, + tma_atom_c, + mC_mnl, + cluster_layout_vmnk, + a_smem_layout_staged, + b_smem_layout_staged, + c_smem_layout_staged, + epi_tile, + tile_sched_params, + shared_shard, + ) + + @cute.jit + def _gemm_device( + self, + tiled_mma: cute.TiledMma, + tma_atom_a: cute.CopyAtom, + mA_mkl: cute.Tensor, + tma_atom_b: cute.CopyAtom, + mB_nkl: cute.Tensor, + tma_atom_c: cute.CopyAtom, + mC_mnl: cute.Tensor, + cluster_layout_vmnk: cute.Layout, + a_smem_layout_staged: cute.ComposedLayout, + b_smem_layout_staged: cute.ComposedLayout, + c_smem_layout_staged: cute.Layout | cute.ComposedLayout, + epi_tile: cute.Tile, + tile_sched_params: utils.PersistentTileSchedulerParams, + shared_shard: cute.Tensor, + ): + warp_idx = cute.arch.warp_idx() + warp_idx = cute.arch.make_warp_uniform(warp_idx) + + if warp_idx == self.tma_warp_id: + cpasync.prefetch_descriptor(tma_atom_a) + cpasync.prefetch_descriptor(tma_atom_b) + cpasync.prefetch_descriptor(tma_atom_c) + + bidx, bidy, bidz = cute.arch.block_idx() + mma_tile_coord_v = bidx % cute.size(tiled_mma.thr_id.shape) + is_leader_cta = mma_tile_coord_v == 0 + cta_rank_in_cluster = cute.arch.make_warp_uniform(cute.arch.block_idx_in_cluster()) + block_in_cluster_coord_vmnk = cluster_layout_vmnk.get_flat_coord(cta_rank_in_cluster) + tidx, _, _ = cute.arch.thread_idx() + + @cute.struct + class SharedStorage: + ab_full_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage * 2] + acc_full_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage * 2] + tmem_dealloc_mbar_ptr: cutlass.Int64 + tmem_holding_buf: cutlass.Int32 + + smem = utils.SmemAllocator() + storage = smem.allocate(SharedStorage) + + ab_pipeline_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread) + num_tma_producer = self.num_mcast_ctas_a + self.num_mcast_ctas_b - 1 + ab_pipeline_consumer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread, num_tma_producer) + ab_pipeline = pipeline.PipelineTmaUmma.create( + barrier_storage=storage.ab_full_mbar_ptr.data_ptr(), + num_stages=self.num_ab_stage, + producer_group=ab_pipeline_producer_group, + consumer_group=ab_pipeline_consumer_group, + tx_count=self.num_tma_load_bytes, + cta_layout_vmnk=cluster_layout_vmnk, + defer_sync=True, + ) + ab_producer, ab_consumer = ab_pipeline.make_participants() + + acc_pipeline_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread) + num_acc_consumer_threads = len(self.epilogue_warp_id) + acc_pipeline_consumer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread, num_acc_consumer_threads) + acc_pipeline = pipeline.PipelineUmmaAsync.create( + barrier_storage=storage.acc_full_mbar_ptr.data_ptr(), + num_stages=self.num_acc_stage, + producer_group=acc_pipeline_producer_group, + consumer_group=acc_pipeline_consumer_group, + cta_layout_vmnk=cluster_layout_vmnk, + defer_sync=True, + ) + + tmem_alloc_barrier = pipeline.NamedBarrier( + barrier_id=self.tmem_alloc_sync_bar_id, + num_threads=32 * len((self.mma_warp_id, *self.epilogue_warp_id)), + ) + tmem = utils.TmemAllocator( + storage.tmem_holding_buf, + barrier_for_retrieve=tmem_alloc_barrier, + allocator_warp_id=self.epilogue_warp_id[0], + is_two_cta=False, + two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar_ptr, + ) + + pipeline_init_arrive(cluster_shape_mn=cluster_layout_vmnk, is_relaxed=True) + + sA = smem.allocate_tensor( + element_type=self.a_dtype, + layout=a_smem_layout_staged.outer, + byte_alignment=128, + swizzle=a_smem_layout_staged.inner, + ) + sB = smem.allocate_tensor( + element_type=self.b_dtype, + layout=b_smem_layout_staged.outer, + byte_alignment=128, + swizzle=b_smem_layout_staged.inner, + ) + + a_full_mcast_mask = None + b_full_mcast_mask = None + if cutlass.const_expr(self.is_a_mcast or self.is_b_mcast): + a_full_mcast_mask = cpasync.create_tma_multicast_mask( + cluster_layout_vmnk, block_in_cluster_coord_vmnk, mcast_mode=2 + ) + b_full_mcast_mask = cpasync.create_tma_multicast_mask( + cluster_layout_vmnk, block_in_cluster_coord_vmnk, mcast_mode=1 + ) + + gA_mkl = cute.local_tile(mA_mkl, cute.slice_(self.mma_tiler, (None, 0, None)), (None, None, None)) + gB_nkl = cute.local_tile(mB_nkl, cute.slice_(self.mma_tiler, (0, None, None)), (None, None, None)) + gC_mnl = cute.local_tile(mC_mnl, cute.slice_(self.mma_tiler, (None, None, 0)), (None, None, None)) + # Shared shard is physically [M, shard_dim]. Give it the same logical MNL + # view as C so its epilogue partition is coordinate-identical. + mShared_mnl = cute.make_tensor( + shared_shard.iterator, + cute.append(shared_shard.layout, cute.make_layout((1,), stride=(0,))), + ) + gShared_mnl = cute.local_tile( + mShared_mnl, + cute.slice_(self.mma_tiler, (None, None, 0)), + (None, None, None), + ) + k_tile_cnt = cute.size(gA_mkl, mode=[3]) + + thr_mma = tiled_mma.get_slice(mma_tile_coord_v) + tCgA = thr_mma.partition_A(gA_mkl) + tCgB = thr_mma.partition_B(gB_nkl) + tCgC = thr_mma.partition_C(gC_mnl) + tCgShared = thr_mma.partition_C(gShared_mnl) + + # Predicate the partial M tile when M is smaller than the MMA tile. + idC = cute.make_identity_tensor(mC_mnl.shape) + cC_mnl = cute.local_tile(idC, cute.slice_(self.mma_tiler, (None, None, 0)), (None, None, None)) + tCcC = thr_mma.partition_C(cC_mnl) + + a_cta_layout = cute.make_layout(cute.slice_(cluster_layout_vmnk, (0, 0, None, 0)).shape) + tAsA, tAgA = cpasync.tma_partition( + tma_atom_a, + block_in_cluster_coord_vmnk[2], + a_cta_layout, + cute.group_modes(sA, 0, 3), + cute.group_modes(tCgA, 0, 3), + ) + b_cta_layout = cute.make_layout(cute.slice_(cluster_layout_vmnk, (0, None, 0, 0)).shape) + tBsB, tBgB = cpasync.tma_partition( + tma_atom_b, + block_in_cluster_coord_vmnk[1], + b_cta_layout, + cute.group_modes(sB, 0, 3), + cute.group_modes(tCgB, 0, 3), + ) + + tCrA = tiled_mma.make_fragment_A(sA) + tCrB = tiled_mma.make_fragment_B(sB) + acc_shape = tiled_mma.partition_shape_C(self.mma_tiler[:2]) + tCtAcc_fake = tiled_mma.make_fragment_C(cute.append(acc_shape, self.num_acc_stage)) + + pipeline_init_wait(cluster_shape_mn=cluster_layout_vmnk) + + gemm_grid_z = cute.arch.grid_dim()[2] + tile_sched = utils.StaticPersistentTileScheduler.create( + tile_sched_params, + cute.arch.block_idx(), + ( + cute.arch.grid_dim()[0], + cute.arch.grid_dim()[1], + gemm_grid_z, + ), + ) + work_tile = tile_sched.initial_work_tile_info() + + if warp_idx == self.tma_warp_id: + while work_tile.is_valid_tile: + cur_tile_coord = work_tile.tile_idx + mma_tile_coord_mnl = ( + cur_tile_coord[0] // cute.size(tiled_mma.thr_id.shape), + cur_tile_coord[1], + cur_tile_coord[2], + ) + + tAgA_slice = tAgA[(None, mma_tile_coord_mnl[0], None, mma_tile_coord_mnl[2])] + tBgB_slice = tBgB[(None, mma_tile_coord_mnl[1], None, mma_tile_coord_mnl[2])] + + # Prime a short prefix of the existing combined A+B ring with + # B only. Its barrier still expects A+B bytes, so the MMA + # consumer cannot observe a half-filled stage. + ab_producer.reset() + peek_ab_empty_status = ab_producer.try_acquire() + + for k_tile in cutlass.range(0, self.b_prime_stages, 1, unroll=1): + handle = ab_producer.acquire_and_advance(peek_ab_empty_status) + cute.copy( + tma_atom_b, + tBgB_slice[(None, handle.count)], + tBsB[(None, handle.index)], + tma_bar_ptr=handle.barrier, + mcast_mask=b_full_mcast_mask, + ) + peek_ab_empty_status = cutlass.Boolean(1) + if handle.count + 1 < self.b_prime_stages: + peek_ab_empty_status = ab_producer.try_acquire() + + cute.arch.griddepcontrol_wait() + + # Supply A to the very same stages after the producer AR has + # programmatically released this dependent kernel. + a_fill_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Producer, self.num_ab_stage) + for k_tile in cutlass.range(0, self.b_prime_stages, 1, unroll=1): + a_barrier = ab_pipeline.producer_get_barrier(a_fill_state) + cute.copy( + tma_atom_a, + tAgA_slice[(None, k_tile)], + tAsA[(None, a_fill_state.index)], + tma_bar_ptr=a_barrier, + mcast_mask=a_full_mcast_mask, + ) + a_fill_state.advance() + + peek_ab_empty_status = ab_producer.try_acquire() + + for k_tile in cutlass.range(self.b_prime_stages, k_tile_cnt, 1, unroll=1): + handle = ab_producer.acquire_and_advance(peek_ab_empty_status) + + cute.copy( + tma_atom_a, + tAgA_slice[(None, handle.count)], + tAsA[(None, handle.index)], + tma_bar_ptr=handle.barrier, + mcast_mask=a_full_mcast_mask, + ) + cute.copy( + tma_atom_b, + tBgB_slice[(None, handle.count)], + tBsB[(None, handle.index)], + tma_bar_ptr=handle.barrier, + mcast_mask=b_full_mcast_mask, + ) + + peek_ab_empty_status = cutlass.Boolean(1) + if handle.count + 1 < k_tile_cnt: + peek_ab_empty_status = ab_producer.try_acquire() + + tile_sched.advance_to_next_work() + work_tile = tile_sched.get_current_work() + + ab_producer.tail() + + if warp_idx == self.mma_warp_id: + tmem.wait_for_alloc() + tmem_ptr = tmem.retrieve_ptr(self.acc_dtype) + tCtAcc_base = cute.make_tensor(tmem_ptr, tCtAcc_fake.layout) + + acc_producer_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Producer, self.num_acc_stage) + + while work_tile.is_valid_tile: + cur_tile_coord = work_tile.tile_idx + mma_tile_coord_mnl = ( + cur_tile_coord[0] // cute.size(tiled_mma.thr_id.shape), + cur_tile_coord[1], + cur_tile_coord[2], + ) + + tCtAcc = tCtAcc_base[(None, None, None, acc_producer_state.index)] + + ab_consumer.reset() + peek_ab_full_status = cutlass.Boolean(1) + if is_leader_cta: + peek_ab_full_status = ab_consumer.try_wait() + + if is_leader_cta: + acc_pipeline.producer_acquire(acc_producer_state) + + tiled_mma.set(tcgen05.Field.ACCUMULATE, False) + + for k_tile in range(k_tile_cnt): + if is_leader_cta: + handle = ab_consumer.wait_and_advance(peek_ab_full_status) + + num_kblocks = cute.size(tCrA, mode=[2]) + for kblk_idx in cutlass.range(num_kblocks, unroll_full=True): + kblk_crd = (None, None, kblk_idx, handle.index) + + cute.gemm( + tiled_mma, + tCtAcc, + tCrA[kblk_crd], + tCrB[kblk_crd], + tCtAcc, + ) + tiled_mma.set(tcgen05.Field.ACCUMULATE, True) + + handle.release() + + peek_ab_full_status = cutlass.Boolean(1) + if handle.count + 1 < k_tile_cnt: + peek_ab_full_status = ab_consumer.try_wait() + + if is_leader_cta: + acc_pipeline.producer_commit(acc_producer_state) + acc_producer_state.advance() + + tile_sched.advance_to_next_work() + work_tile = tile_sched.get_current_work() + + acc_pipeline.producer_tail(acc_producer_state) + + sC = smem.allocate_tensor( + element_type=self.c_dtype, + layout=c_smem_layout_staged.outer, + byte_alignment=128, + swizzle=c_smem_layout_staged.inner, + ) + + if warp_idx < self.mma_warp_id: + tmem.allocate(self.num_tmem_alloc_cols) + + tmem.wait_for_alloc() + tmem_ptr = tmem.retrieve_ptr(self.acc_dtype) + tCtAcc_base = cute.make_tensor(tmem_ptr, tCtAcc_fake.layout) + + acc_consumer_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, self.num_acc_stage) + c_producer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, + 32 * len(self.epilogue_warp_id), + ) + c_pipeline = pipeline.PipelineTmaStore.create(num_stages=self.num_c_stage, producer_group=c_producer_group) + while work_tile.is_valid_tile: + cur_tile_coord = work_tile.tile_idx + mma_tile_coord_mnl = ( + cur_tile_coord[0] // cute.size(tiled_mma.thr_id.shape), + cur_tile_coord[1], + cur_tile_coord[2], + ) + tile_sched.advance_to_next_work() + work_tile = tile_sched.get_current_work() + + num_tiles_executed = tile_sched.num_tiles_executed + acc_consumer_state = _epilogue_tma_store_add_shared( + self, + tidx, + warp_idx, + tma_atom_c, + tCtAcc_base, + sC, + tCgC, + tCgShared, + tCcC, + mC_mnl, + mShared_mnl, + epi_tile, + num_tiles_executed, + mma_tile_coord_mnl, + acc_consumer_state, + acc_pipeline, + c_pipeline, + ) + + c_pipeline.producer_tail() + + tmem.relinquish_alloc_permit() + tmem.free(tmem_ptr) + + # Allow the Lamport copy to become resident before this grid + # fully retires. Its griddepcontrol.wait still enforces complete + # producer-grid ordering before mailbox inspection. + cute.arch.griddepcontrol_launch_dependents() + + @staticmethod + def _compute_grid( + c: cute.Tensor, + cta_tile_shape_mnk: tuple[int, int, int], + cluster_shape_mn: tuple[int, int], + max_active_clusters: cutlass.Constexpr, + ) -> tuple[utils.PersistentTileSchedulerParams, tuple[int, int, int]]: + """Build the static persistent schedule.""" + c_shape = cute.slice_(cta_tile_shape_mnk, (None, None, 0)) + gc = cute.zipped_divide(c, tiler=c_shape) + num_ctas_mnl = gc[(0, (None, None, None))].shape + cluster_shape_mnl = (*cluster_shape_mn, 1) + + tile_sched_params = utils.PersistentTileSchedulerParams(num_ctas_mnl, cluster_shape_mnl) + grid = utils.StaticPersistentTileScheduler.get_grid_shape(tile_sched_params, max_active_clusters) + + return tile_sched_params, grid + + @staticmethod + def _compute_num_tmem_alloc_cols( + tiled_mma: cute.TiledMma, + mma_tiler: tuple[int, int, int], + num_acc_stage: int, + arch: str, + ) -> int: + """Return the required tensor-memory column count.""" + acc_shape = tiled_mma.partition_shape_C(mma_tiler[:2]) + tCtAcc_fake = tiled_mma.make_fragment_C(cute.append(acc_shape, num_acc_stage)) + num_tmem_alloc_cols = utils.get_num_tmem_alloc_cols(tCtAcc_fake, arch=arch) + + return num_tmem_alloc_cols + + +@cute.jit +def launch_kernel( + gemm_op: cutlass.Constexpr, + a: cute.Tensor, # (l, m, k) + b: cute.Tensor, # (l, n, k) + c: cute.Tensor, # (l, m, n) + shared_shard: cute.Tensor, # (m, shard_dim), private TP shard + rows: cutlass.Int64, + c_multicast_i64: cutlass.Int64, + full_hidden_dim: cutlass.Constexpr, + shard_dim: cutlass.Constexpr, + max_active_clusters: cutlass.Constexpr, + stream: cuda.CUstream, +): + """Launch the fused-add multicast GEMM using PyTorch BMM tensor order.""" + # C is passed as the fixed-capacity [1,max_m,H] symmetric mailbox. Only M + # is runtime-variable; construct this rank's logical [1,M,S] view here so + # H and S remain compile-time constants and no dynamic strided host view is + # needed on every call. __call__ later replaces the local base pointer with + # the already rank-offset multicast address. + c = cute.make_tensor( + c.iterator, + cute.make_layout( + (1, rows, shard_dim), + stride=(0, full_hidden_dim, 1), + ), + ) + # (l,m,k) -> (m,k,l) + a = cute.make_tensor(a.iterator, cute.select(a.layout, mode=[1, 2, 0])) + # (l,n,k) -> (n,k,l) + b = cute.make_tensor(b.iterator, cute.select(b.layout, mode=[1, 2, 0])) + # (l,m,n) -> (m,n,l) + c = cute.make_tensor(c.iterator, cute.select(c.layout, mode=[1, 2, 0])) + + gemm_op( + a, + b, + c, + shared_shard, + c_multicast_i64, + max_active_clusters, + stream, + ) + + +_COMPILED: dict[tuple[object, ...], object] = {} + + +def compile_kernel( + mnkl: tuple[int, int, int, int], + a: cute.Tensor, + b: cute.Tensor, + c: cute.Tensor, + shared_shard: cute.Tensor, + full_hidden_dim: int, + shard_dim: int, + mma_tiler_mn: tuple[int, int] = (64, 32), + cluster_shape_mn: tuple[int, int] = (1, 8), + max_active_clusters: cutlass.Constexpr = None, + b_prime_stages: int = 2, +): + key = ( + torch.accelerator.current_device_index(), + mnkl, + full_hidden_dim, + shard_dim, + mma_tiler_mn, + cluster_shape_mn, + max_active_clusters, + b_prime_stages, + ) + if key in _COMPILED: + return _COMPILED[key] + + from cutlass.cute.runtime import make_fake_stream, make_fake_tensor + + gemm = FusedAddMulticastGemm( + mma_tiler_mn, + cluster_shape_mn, + b_prime_stages, + ) + validate_configuration( + latent_dim=mnkl[2], + shard_dim=mnkl[1], + mma_tiler_mn=mma_tiler_mn, + cluster_shape_mn=cluster_shape_mn, + b_prime_stages=b_prime_stages, + ) + if any(tensor.element_type is not cutlass.BFloat16 for tensor in (a, b, c)): + raise ValueError("up-projection tensors must be BF16") + + # The producer writes [M, full_hidden_dim]. GEMM receives a rank-local + # [M, shard_dim] view: contiguous within a row, with full_hidden_dim as + # its leading dimension. + shared_shard_compile = make_fake_tensor( + shared_shard.element_type, + (mnkl[0], shard_dim), + stride=(full_hidden_dim, 1), + assumed_align=16, + ) + stream = make_fake_stream() + compiled = cute.compile( + launch_kernel, + gemm, + a, + b, + c, + shared_shard_compile, + cutlass.Int64(mnkl[0]), + cutlass.Int64(0), + full_hidden_dim, + shard_dim, + max_active_clusters, + stream, + ) + _COMPILED[key] = compiled + return compiled + + +class AdaptiveUpProjectionKernel: + """Dispatch static-M Skinny or dynamic-M WGMMA into one mailbox.""" + + def __init__( + self, + *, + group: dist.ProcessGroup, + rank: int, + tp_size: int, + latent_dim: int, + hidden_dim: int, + max_m: int, + skinny_max_m: int, + mma_tiler_mn: tuple[int, int], + cluster_shape_mn: tuple[int, int], + b_prime_stages: int, + ) -> None: + if hidden_dim % tp_size: + raise ValueError("hidden_dim must be divisible by TP size") + if not 0 <= skinny_max_m <= min(8, max_m): + raise ValueError("skinny_max_m must be in [0, min(8, max_m)]") + self.rank = rank + self.tp_size = tp_size + self.latent_dim = latent_dim + self.hidden_dim = hidden_dim + self.shard_dim = hidden_dim // tp_size + self.max_m = max_m + self.skinny_max_m = skinny_max_m + self.mma_tiler_mn = mma_tiler_mn + self.cluster_shape_mn = cluster_shape_mn + self.b_prime_stages = b_prime_stages + device = torch.device("cuda", torch.accelerator.current_device_index()) + self._device = device + self._dynamic: Any | None = None + self._skinny_by_m: dict[int, FusedAddMulticastSkinnyGemmKernel] = {} + validate_configuration( + latent_dim=latent_dim, + shard_dim=self.shard_dim, + mma_tiler_mn=mma_tiler_mn, + cluster_shape_mn=cluster_shape_mn, + b_prime_stages=b_prime_stages, + ) + if skinny_max_m and self.latent_dim % (224 * 8): + raise ValueError("Skinny up-projection requires latent_dim divisible by 1792.") + + self._mailbox = symm_mem.empty( + (1, max_m, hidden_dim), + dtype=torch.bfloat16, + device=device, + ) + self._mailbox_symm_mem = symm_mem.rendezvous(self._mailbox, group) + self._mailbox.view(torch.int32).fill_(-0x80000000) + multicast_ptr = self._mailbox_symm_mem.multicast_ptr + if multicast_ptr is None or multicast_ptr == 0: + raise RuntimeError("mailbox NVLS multicast mapping is unavailable") + self._mailbox_multicast_ptr = int(multicast_ptr) + + cluster_size = math.prod(cluster_shape_mn) + self._max_active_clusters = utils.HardwareInfo().get_max_active_clusters(cluster_size) + self._mailbox_c = _as_cute(self._mailbox) + + def compile_dynamic(self) -> None: + if self._dynamic is not None: + return + device = self._device + with torch.accelerator.device_index(device.index): + compile_latent = torch.empty( + (1, self.max_m, self.latent_dim), + dtype=torch.bfloat16, + device=device, + ) + compile_weight = torch.empty( + (self.shard_dim, self.latent_dim), + dtype=torch.bfloat16, + device=device, + ) + compile_shared = torch.empty( + (self.max_m, self.hidden_dim), + dtype=torch.bfloat16, + device=device, + )[ + :, + self.rank * self.shard_dim : (self.rank + 1) * self.shard_dim, + ] + compile_latent_c = _as_cute(compile_latent, dynamic_m=True) + compile_weight_c = _as_cute(compile_weight.unsqueeze(0)) + compile_shared_c = _as_cute(compile_shared) + + self._dynamic = compile_kernel( + (self.max_m, self.shard_dim, self.latent_dim, 1), + compile_latent_c, + compile_weight_c, + self._mailbox_c, + compile_shared_c, + self.hidden_dim, + self.shard_dim, + self.mma_tiler_mn, + self.cluster_shape_mn, + self._max_active_clusters, + self.b_prime_stages, + ) + + def compile_skinny(self, m: int) -> None: + if not 1 <= m <= self.skinny_max_m: + raise ValueError(f"Skinny up-projection requires M in [1, {self.skinny_max_m}].") + if m in self._skinny_by_m: + return + with torch.accelerator.device_index(self._device.index): + self._skinny_by_m[m] = FusedAddMulticastSkinnyGemmKernel( + rank=self.rank, + tp_size=self.tp_size, + latent_dim=self.latent_dim, + hidden_dim=self.hidden_dim, + num_rows=m, + ) + + def ensure_compiled(self, m: int) -> None: + if not 1 <= m <= self.max_m: + raise ValueError(f"runtime M={m} must be in [1, {self.max_m}]") + if m <= self.skinny_max_m: + self.compile_skinny(m) + else: + self.compile_dynamic() + + def __call__( + self, + latent: torch.Tensor, + weight: torch.Tensor, + shared_shard: torch.Tensor, + ) -> torch.Tensor: + if latent.ndim != 2: + raise ValueError("latent must be rank-2") + m = latent.shape[0] + device = self._mailbox.device + expected = ( + (latent, (m, self.latent_dim), "latent"), + ( + weight, + (self.shard_dim, self.latent_dim), + "weight", + ), + ( + shared_shard, + (self.max_m, self.shard_dim), + "shared_shard", + ), + ) + for tensor, shape, name in expected: + if tensor.shape != shape or tensor.dtype != torch.bfloat16 or tensor.device != device: + raise ValueError(f"{name} must be CUDA torch.bfloat16 {list(shape)}") + if not latent.is_contiguous() or not weight.is_contiguous() or shared_shard.stride() != (self.hidden_dim, 1): + raise ValueError("up-projection inputs have unsupported strides") + if not 1 <= m <= self.max_m: + raise ValueError(f"runtime M={m} must be in [1, {self.max_m}]") + + if m <= self.skinny_max_m: + skinny = self._skinny_by_m.get(m) + if skinny is None: + raise RuntimeError(f"Skinny up-projection M={m} was not compiled before launch.") + return skinny( + latent, + weight, + shared_shard, + self._mailbox, + self._mailbox_multicast_ptr, + ) + + if self._dynamic is None: + raise RuntimeError("Dynamic up-projection was not compiled before launch.") + with torch.accelerator.device_index(device.index): + stream = cuda.CUstream(torch.cuda.current_stream(device).cuda_stream) + self._dynamic( + _as_cute(latent.unsqueeze(0), dynamic_m=True), + _as_cute(weight.unsqueeze(0)), + self._mailbox_c, + _as_cute(shared_shard), + cutlass.Int64(m), + cutlass.Int64(self._mailbox_multicast_ptr + self.rank * self.shard_dim * 2), + stream, + ) + return self._mailbox diff --git a/aphrodite/models/kimi_k3/nvidia/ops/cute_dsl/latent_moe_tail/fused_add_multicast_skinny_gemm.py b/aphrodite/models/kimi_k3/nvidia/ops/cute_dsl/latent_moe_tail/fused_add_multicast_skinny_gemm.py new file mode 100644 index 0000000000..68f88c1a2d --- /dev/null +++ b/aphrodite/models/kimi_k3/nvidia/ops/cute_dsl/latent_moe_tail/fused_add_multicast_skinny_gemm.py @@ -0,0 +1,424 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Static-M SIMT skinny GEMM with a shared-add multicast epilogue.""" + +from __future__ import annotations + +from dataclasses import dataclass + +import cuda.bindings.driver as cuda +import cutlass +import cutlass.cute as cute +import torch +from cutlass import BFloat16, Float32, Int64, const_expr +from cutlass.cute.runtime import from_dlpack + +from .primitives import ( + CUDAGraphCompatibleWrapper, + bf16x2_to_u32, + bf16x4_to_packed_u32x2, + bf16x8_to_packed_u32x4, + sanitize_negative_zero, + sanitize_negative_zero_u32, + sanitize_negative_zero_u32x2, + store_global_u32, + store_global_u32x2, + store_global_u32x4, +) + + +@dataclass(frozen=True) +class SkinnyConfig: + block_size: int = 224 + outputs_per_block: int = 2 + k_unroll: int = 2 + vector_width: int = 8 + prefetch_b_before_pdl: bool = True + + +def config_for_m(num_rows: int, shard_dim: int = 896) -> SkinnyConfig: + if shard_dim == 448: + if num_rows >= 6: + return SkinnyConfig( + block_size=224, + outputs_per_block=2, + k_unroll=1, + ) + outputs_per_block = 2 if num_rows <= 3 else 4 + return SkinnyConfig( + block_size=448, + outputs_per_block=outputs_per_block, + k_unroll=1, + ) + if num_rows == 1: + return SkinnyConfig(outputs_per_block=8, k_unroll=1) + return SkinnyConfig(outputs_per_block=4) + + +def _as_cute(tensor: torch.Tensor): + return from_dlpack( + CUDAGraphCompatibleWrapper(tensor.detach()), + assumed_align=16, + ) + + +class FusedAddMulticastSkinnyGemm: + """SIMT GEMM adapted from the existing Skinny GEMM.""" + + def __init__( + self, + *, + num_rows: int, + hidden_dim: int, + config: SkinnyConfig, + ) -> None: + if config.block_size % 32: + raise ValueError("skinny block_size must be a multiple of 32") + self.num_rows = num_rows + self.hidden_dim = hidden_dim + self.block_size = config.block_size + self.outputs_per_block = config.outputs_per_block + self.k_unroll = config.k_unroll + self.vector_width = config.vector_width + self.prefetch_b_before_pdl = config.prefetch_b_before_pdl + self.num_warps = config.block_size // 32 + + @cute.jit + def __call__( + self, + gA: cute.Tensor, + gB: cute.Tensor, + gShared: cute.Tensor, + output_multicast_ptr: Int64, + stream: cuda.CUstream, + ) -> None: + n = cute.size(gB, mode=[0]) + k = cute.size(gA, mode=[1]) + copy_a = cute.make_copy_atom( + cute.nvgpu.CopyG2ROp(), + BFloat16, + num_bits_per_copy=self.vector_width * BFloat16.width, + load_cache_mode=cute.nvgpu.LoadCacheMode.ALWAYS, + ) + copy_b = cute.make_copy_atom( + cute.nvgpu.CopyG2ROp(), + BFloat16, + num_bits_per_copy=self.vector_width * BFloat16.width, + load_cache_mode=cute.nvgpu.LoadCacheMode.STREAMING, + ) + self.kernel( + gA, + gB, + gShared, + output_multicast_ptr, + k, + copy_a, + copy_b, + ).launch( + grid=[cute.ceil_div(n, self.outputs_per_block), 1, 1], + block=[self.block_size, 1, 1], + smem=(self.num_rows * self.outputs_per_block * self.num_warps * 4), + stream=stream, + use_pdl=True, + min_blocks_per_mp=1, + ) + + @cute.kernel + def kernel( + self, + gA: cute.Tensor, + gB: cute.Tensor, + gShared: cute.Tensor, + output_multicast_ptr: Int64, + k_extent: cutlass.Int32, + copy_a: cute.CopyAtom, + copy_b: cute.CopyAtom, + ) -> None: + tidx, _, _ = cute.arch.thread_idx() + block_idx, _, _ = cute.arch.block_idx() + warp_idx = cute.arch.warp_idx() + + outputs_per_block: cutlass.Constexpr = self.outputs_per_block + vector_width: cutlass.Constexpr = self.vector_width + block_size: cutlass.Constexpr = self.block_size + num_warps: cutlass.Constexpr = self.num_warps + num_rows: cutlass.Constexpr = self.num_rows + + acc = cute.make_rmem_tensor( + cute.make_layout( + (num_rows, outputs_per_block), + stride=(outputs_per_block, 1), + ), + Float32, + ) + acc.fill(0.0) + + n_base = block_idx * outputs_per_block + k_tile_size: cutlass.Constexpr = block_size * vector_width + num_k_tiles = k_extent // k_tile_size + gA_vec = cute.logical_divide(gA, (None, vector_width)) + gB_vec = cute.logical_divide(gB, (None, vector_width)) + tA_all = cute.logical_divide(gA_vec, (None, (None, block_size))) + tB_all = cute.logical_divide(gB_vec, (None, (None, block_size))) + tA = tA_all[None, (None, (tidx, None))] + + a_regs = cute.make_rmem_tensor( + cute.make_layout( + (num_rows, vector_width), + stride=(vector_width, 1), + ), + BFloat16, + ) + b_regs = cute.make_rmem_tensor( + cute.make_layout( + (outputs_per_block, vector_width), + stride=(vector_width, 1), + ), + BFloat16, + ) + + if const_expr(self.prefetch_b_before_pdl): + for ni in cutlass.range_constexpr(outputs_per_block): + tB = tB_all[n_base + ni, (None, (tidx, None))] + cute.copy(copy_b, tB[None, 0], b_regs[ni, None]) + + cute.arch.griddepcontrol_wait() + + for mi in cutlass.range_constexpr(num_rows): + cute.copy(copy_a, tA[mi, None, 0], a_regs[mi, None]) + if const_expr(not self.prefetch_b_before_pdl): + for ni in cutlass.range_constexpr(outputs_per_block): + tB = tB_all[n_base + ni, (None, (tidx, None))] + cute.copy(copy_b, tB[None, 0], b_regs[ni, None]) + for vi in cutlass.range_constexpr(vector_width): + for mi in cutlass.range_constexpr(num_rows): + for ni in cutlass.range_constexpr(outputs_per_block): + acc[mi, ni] = acc[mi, ni] + a_regs[mi, vi].to(Float32) * b_regs[ni, vi].to(Float32) + + for k_tile in cutlass.range(1, num_k_tiles, unroll=self.k_unroll): + for mi in cutlass.range_constexpr(num_rows): + cute.copy( + copy_a, + tA[mi, None, k_tile], + a_regs[mi, None], + ) + for ni in cutlass.range_constexpr(outputs_per_block): + tB = tB_all[n_base + ni, (None, (tidx, None))] + cute.copy(copy_b, tB[None, k_tile], b_regs[ni, None]) + for vi in cutlass.range_constexpr(vector_width): + for mi in cutlass.range_constexpr(num_rows): + for ni in cutlass.range_constexpr(outputs_per_block): + acc[mi, ni] = acc[mi, ni] + a_regs[mi, vi].to(Float32) * b_regs[ni, vi].to(Float32) + + for mi in cutlass.range_constexpr(num_rows): + for ni in cutlass.range_constexpr(outputs_per_block): + acc[mi, ni] = cute.arch.warp_reduction_sum(acc[mi, ni]) + + smem_layout = cute.make_layout( + (num_rows, outputs_per_block, num_warps), + stride=(outputs_per_block * num_warps, num_warps, 1), + ) + smem = cutlass.utils.SmemAllocator() + partials = smem.allocate_tensor( + Float32, + smem_layout, + byte_alignment=16, + ) + with cute.arch.elect_one(): + for mi in cutlass.range_constexpr(num_rows): + for ni in cutlass.range_constexpr(outputs_per_block): + partials[mi, ni, warp_idx] = acc[mi, ni] + + cute.arch.sync_threads() + if tidx == 0: + fused = cute.make_rmem_tensor( + cute.make_layout((outputs_per_block,)), + BFloat16, + ) + for mi in cutlass.range_constexpr(num_rows): + for ni in cutlass.range_constexpr(outputs_per_block): + total = ( + partials[mi, ni, None] + .load() + .reduce( + cute.ReductionOp.ADD, + init_val=Float32(0.0), + reduction_profile=0, + ) + ) + gemm_value = Float32(total).to(BFloat16) + fused[ni] = (gemm_value.to(Float32) + gShared[mi, n_base + ni].to(Float32)).to(BFloat16) + output_offset = Int64((mi * self.hidden_dim + n_base) * 2) + if const_expr(outputs_per_block == 2): + packed = sanitize_negative_zero_u32(bf16x2_to_u32(fused.load())) + store_global_u32( + output_multicast_ptr + output_offset, + packed, + ) + elif const_expr(outputs_per_block == 4): + packed = sanitize_negative_zero_u32x2(bf16x4_to_packed_u32x2(fused.load())) + store_global_u32x2( + output_multicast_ptr + output_offset, + packed, + ) + else: + packed = sanitize_negative_zero(bf16x8_to_packed_u32x4(fused.load())) + store_global_u32x4( + output_multicast_ptr + output_offset, + packed, + ) + + cute.arch.griddepcontrol_launch_dependents() + + +_COMPILED: dict[tuple[object, ...], object] = {} + + +def compile_kernel( + *, + num_rows: int, + latent_dim: int, + hidden_dim: int, + shard_dim: int, + config: SkinnyConfig, +): + key = ( + torch.accelerator.current_device_index(), + num_rows, + latent_dim, + hidden_dim, + shard_dim, + config, + ) + if key in _COMPILED: + return _COMPILED[key] + + from cutlass.cute.runtime import make_fake_stream, make_fake_tensor + + if shard_dim % config.outputs_per_block: + raise ValueError("shard_dim must be divisible by outputs_per_block") + if latent_dim % (config.block_size * config.vector_width): + raise ValueError("latent_dim must be divisible by block_size * vector_width") + + a = make_fake_tensor( + BFloat16, + (num_rows, latent_dim), + stride=(latent_dim, 1), + assumed_align=16, + ) + b = make_fake_tensor( + BFloat16, + (shard_dim, latent_dim), + stride=(latent_dim, 1), + assumed_align=16, + ) + shared = make_fake_tensor( + BFloat16, + (num_rows, shard_dim), + stride=(hidden_dim, 1), + assumed_align=16, + ) + compiled = cute.compile( + FusedAddMulticastSkinnyGemm( + num_rows=num_rows, + hidden_dim=hidden_dim, + config=config, + ), + a, + b, + shared, + Int64(0), + make_fake_stream(), + options="--ptxas-options -maxrregcount=128", + ) + _COMPILED[key] = compiled + return compiled + + +class FusedAddMulticastSkinnyGemmKernel: + """Buffer-free compiled launcher for one static M.""" + + def __init__( + self, + *, + rank: int, + tp_size: int, + latent_dim: int, + hidden_dim: int, + num_rows: int, + ) -> None: + if not 1 <= num_rows <= 8: + raise ValueError("skinny backend requires static M in [1, 8]") + if hidden_dim % tp_size: + raise ValueError("hidden_dim must be divisible by TP size") + self.rank = rank + self.latent_dim = latent_dim + self.hidden_dim = hidden_dim + self.shard_dim = hidden_dim // tp_size + self.num_rows = num_rows + self._skinny = compile_kernel( + num_rows=num_rows, + latent_dim=latent_dim, + hidden_dim=hidden_dim, + shard_dim=self.shard_dim, + config=config_for_m(num_rows, self.shard_dim), + ) + + def __call__( + self, + latent: torch.Tensor, + weight: torch.Tensor, + shared_shard: torch.Tensor, + mailbox: torch.Tensor, + mailbox_multicast_ptr: int, + ) -> torch.Tensor: + device = mailbox.device + expected = ( + ( + latent, + (self.num_rows, self.latent_dim), + torch.bfloat16, + "latent", + ), + ( + weight, + (self.shard_dim, self.latent_dim), + torch.bfloat16, + "weight", + ), + ( + shared_shard, + (mailbox.shape[1], self.shard_dim), + torch.bfloat16, + "shared_shard", + ), + ( + mailbox, + (1, mailbox.shape[1], self.hidden_dim), + torch.bfloat16, + "mailbox", + ), + ) + for tensor, shape, dtype, name in expected: + if tensor.shape != shape or tensor.dtype != dtype or tensor.device != device: + raise ValueError(f"{name} must be CUDA {dtype} {list(shape)}") + if ( + not latent.is_contiguous() + or not weight.is_contiguous() + or shared_shard.stride() != (self.hidden_dim, 1) + or not mailbox.is_contiguous() + ): + raise ValueError("skinny up-projection inputs have unsupported strides") + if mailbox.shape[1] < self.num_rows: + raise ValueError("mailbox capacity is smaller than runtime M") + + with torch.accelerator.device_index(device.index): + self._skinny( + _as_cute(latent), + _as_cute(weight), + _as_cute(shared_shard[: self.num_rows]), + Int64(mailbox_multicast_ptr + self.rank * self.shard_dim * 2), + cuda.CUstream(torch.cuda.current_stream(device).cuda_stream), + ) + return mailbox diff --git a/aphrodite/models/kimi_k3/nvidia/ops/cute_dsl/latent_moe_tail/lamport_copy.py b/aphrodite/models/kimi_k3/nvidia/ops/cute_dsl/latent_moe_tail/lamport_copy.py new file mode 100644 index 0000000000..e61c3b9ccc --- /dev/null +++ b/aphrodite/models/kimi_k3/nvidia/ops/cute_dsl/latent_moe_tail/lamport_copy.py @@ -0,0 +1,220 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Lamport mailbox-to-local copy and reset.""" + +from __future__ import annotations + +import functools + +import cuda.bindings.driver as cuda +import cutlass +import cutlass.cute as cute +import torch +from cutlass.cute.runtime import make_fake_compact_tensor, make_fake_stream + +from .primitives import ( + VEC_BF16, + fragment_is_dirty, + load_global_u32x4, + store_global_u32x4, + store_lamport_sentinel_128, + to_cute, + to_cute_dynamic_m, +) + + +class LamportCopy: + """Consume the local physical copy of an NVLS-multicast mailbox.""" + + def __init__(self, hidden_dim: int, ctas: int, threads: int): + self.hidden_dim = hidden_dim + self.ctas = ctas + self.threads = threads + + @cute.jit + def __call__( + self, + symmetric_mailbox: cute.Tensor, + local_output: cute.Tensor, + m: cutlass.Int32, + stream: cuda.CUstream, + ): + self.kernel(symmetric_mailbox, local_output, m).launch( + grid=(self.ctas, 1, 1), + block=(self.threads, 1, 1), + stream=stream, + use_pdl=True, + ) + + @cute.kernel + def kernel( + self, + symmetric_mailbox: cute.Tensor, + local_output: cute.Tensor, + m: cutlass.Int32, + ): + # The CTA may be scheduled early, but mailbox inspection must not pass + # the producer GEMM's programmatic completion point. + cute.arch.griddepcontrol_wait() + + tidx, _, _ = cute.arch.thread_idx() + block, _, _ = cute.arch.block_idx() + thread = cutlass.Int64(block * self.threads + tidx) + stride = cutlass.Int64(self.ctas * self.threads) + fragments = cutlass.Int64(m) * cutlass.Int64(self.hidden_dim // VEC_BF16) + + fragment = thread + while fragment < fragments: + element = fragment * VEC_BF16 + source = cute.make_ptr( + cutlass.BFloat16, + (symmetric_mailbox.iterator + element).llvm_ptr, + cute.AddressSpace.gmem, + assumed_align=16, + ) + packed = load_global_u32x4(source, volatile=True) + while fragment_is_dirty(packed): + packed = load_global_u32x4(source, volatile=True) + + destination = cutlass.Int64((local_output.iterator + element).toint()) + store_global_u32x4(destination, packed, volatile=False) + fragment = fragment + stride + + # The returned ordinary tensor is complete. A same-stream successor + # may overlap the mailbox cleanup below. + cute.arch.griddepcontrol_launch_dependents() + + fragment = thread + while fragment < fragments: + element = fragment * VEC_BF16 + source = cute.make_ptr( + cutlass.BFloat16, + (symmetric_mailbox.iterator + element).llvm_ptr, + cute.AddressSpace.gmem, + assumed_align=16, + ) + store_lamport_sentinel_128(source) + fragment = fragment + stride + + +@functools.cache +def compile_kernel( + hidden_dim: int, + max_m: int, + ctas: int, + threads: int, + device_index: int, +): + if hidden_dim <= 0 or hidden_dim % VEC_BF16: + raise ValueError("hidden_dim must be a positive multiple of 8") + if max_m <= 0: + raise ValueError("max_m must be positive") + if ctas <= 0: + raise ValueError("Lamport copy ctas must be positive") + if not 1 <= threads <= 1024: + raise ValueError("Lamport copy threads must be in [1, 1024]") + with torch.accelerator.device_index(device_index): + mailbox = make_fake_compact_tensor(cutlass.BFloat16, (max_m * hidden_dim,), assumed_align=16) + output = make_fake_compact_tensor( + cutlass.BFloat16, + (cute.sym_int32(divisibility=VEC_BF16),), + assumed_align=16, + ) + return cute.compile( + LamportCopy(hidden_dim, ctas, threads), + mailbox, + output, + cutlass.Int32(max_m), + make_fake_stream(), + ) + + +def launch( + symmetric_mailbox: torch.Tensor, + local_output: torch.Tensor, + *, + m: int, + hidden_dim: int, + max_m: int, + ctas: int, + threads: int, +) -> None: + if not 1 <= m <= max_m: + raise ValueError(f"runtime M={m} must be in [1, {max_m}]") + if ( + symmetric_mailbox.ndim != 3 + or symmetric_mailbox.dtype != torch.bfloat16 + or not symmetric_mailbox.is_cuda + or not symmetric_mailbox.is_contiguous() + or symmetric_mailbox.shape[0] != 1 + or symmetric_mailbox.shape[2] != hidden_dim + ): + raise ValueError(f"symmetric_mailbox must be contiguous CUDA BF16 [1,M,{hidden_dim}]") + if symmetric_mailbox.shape[1] != max_m: + raise ValueError("symmetric mailbox must retain its full max_m capacity") + if ( + local_output.shape != (1, m, hidden_dim) + or local_output.dtype != torch.bfloat16 + or local_output.device != symmetric_mailbox.device + or not local_output.is_contiguous() + ): + raise ValueError("local_output must be contiguous CUDA BF16 [1,M,H]") + + device_index = symmetric_mailbox.device.index + stream = cuda.CUstream(torch.cuda.current_stream(symmetric_mailbox.device).cuda_stream) + compile_kernel(hidden_dim, max_m, ctas, threads, device_index)( + to_cute(symmetric_mailbox.flatten(), 16), + to_cute_dynamic_m( + local_output.flatten(), + mode=0, + assumed_align=16, + ), + cutlass.Int32(m), + stream, + ) + + +class LamportCopyKernel: + """Copy a borrowed symmetric mailbox into a fresh local tensor.""" + + def __init__( + self, + *, + hidden_dim: int, + max_m: int, + ctas: int, + threads: int, + ) -> None: + self.hidden_dim = hidden_dim + self.max_m = max_m + self.ctas = ctas + self.threads = threads + compile_kernel( + hidden_dim, + max_m, + ctas, + threads, + torch.accelerator.current_device_index(), + ) + + def __call__(self, symmetric_mailbox: torch.Tensor, *, m: int) -> torch.Tensor: + if not symmetric_mailbox.is_cuda: + raise ValueError("symmetric_mailbox must be a CUDA tensor") + device = symmetric_mailbox.device + with torch.accelerator.device_index(device.index): + output = torch.empty( + (1, m, self.hidden_dim), + dtype=torch.bfloat16, + device=device, + ) + launch( + symmetric_mailbox, + output, + m=m, + hidden_dim=self.hidden_dim, + max_m=self.max_m, + ctas=self.ctas, + threads=self.threads, + ) + return output diff --git a/aphrodite/models/kimi_k3/nvidia/ops/cute_dsl/latent_moe_tail/primitives.py b/aphrodite/models/kimi_k3/nvidia/ops/cute_dsl/latent_moe_tail/primitives.py new file mode 100644 index 0000000000..072d23a2b1 --- /dev/null +++ b/aphrodite/models/kimi_k3/nvidia/ops/cute_dsl/latent_moe_tail/primitives.py @@ -0,0 +1,431 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Shared CuTe DSL primitives; this module does not define a CUDA kernel.""" + +from __future__ import annotations + +import cutlass +import cutlass.cute as cute +import torch +from cutlass import BFloat16, Float32, Int32, Int64, Uint16, Uint32 +from cutlass._mlir import ir +from cutlass._mlir.dialects import llvm, vector +from cutlass.cute.runtime import from_dlpack +from cutlass.cutlass_dsl import T, dsl_user_op + +VEC_BF16 = 8 +PACKED_BYTES = 16 +NUM_LAMPORT_BUFFERS = 3 +NEG_ZERO_F32_BITS = 0x80000000 +NEG_ZERO_BF16_BITS = 0x8000 + + +class CUDAGraphCompatibleWrapper: + """DLPack view that does not synchronize with the producer stream.""" + + def __init__(self, tensor: torch.Tensor): + self.tensor = tensor + + def __dlpack__(self, stream=None): + return self.tensor.__dlpack__(stream=-1) + + def __dlpack_device__(self): + return self.tensor.__dlpack_device__() + + +def to_cute(tensor: torch.Tensor, assumed_align: int = 16) -> cute.Tensor: + return from_dlpack(CUDAGraphCompatibleWrapper(tensor.detach()), assumed_align=assumed_align) + + +def to_cute_dynamic_m( + tensor: torch.Tensor, + *, + mode: int, + assumed_align: int = 16, +) -> cute.Tensor: + """Expose exactly one compact tensor mode as a runtime shape. + + Model dimensions remain part of the compiled tensor type. Only the token + mode is symbolic, so changing M within an Op's capacity reuses the same + compiled kernel. + """ + + return to_cute(tensor, assumed_align).mark_compact_shape_dynamic( + mode=mode, + stride_order=tensor.dim_order(), + ) + + +@dsl_user_op +def load_global_u32x4( + pointer: cute.Pointer, + *, + volatile: cutlass.Constexpr[bool] = False, + loc=None, + ip=None, +): + """Load one 128-bit fragment as four u32 registers. + + The volatile form is the Lamport polling load. Marking the asm + side-effecting prevents loop-invariant motion and common-subexpression + elimination across polling iterations. + """ + + address = pointer.toint(loc=loc, ip=ip) + opcode = "ld.volatile.global.v4.u32" if volatile else "ld.global.v4.u32" + out = llvm.inline_asm( + llvm.StructType.get_literal([T.i32()] * 4), + [address.ir_value(loc=loc, ip=ip)], + f"{opcode} {{$0, $1, $2, $3}}, [$4];", + "=r,=r,=r,=r,l", + has_side_effects=volatile, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + packed = vector.from_elements( + ir.VectorType.get([4], T.i32(), loc=loc), + [llvm.extractvalue(T.i32(), out, [i], loc=loc, ip=ip) for i in range(4)], + loc=loc, + ip=ip, + ) + return cute.TensorSSA(packed, 4, Uint32) + + +@dsl_user_op +def store_global_u32x4( + address: Int64, + packed, + *, + volatile: cutlass.Constexpr[bool] = False, + loc=None, + ip=None, +) -> None: + """Store four packed words to an ordinary or NVLS multicast global VA.""" + + words = [packed[i].ir_value(loc=loc, ip=ip) for i in range(4)] + opcode = "st.volatile.global.v4.u32" if volatile else "st.global.v4.u32" + llvm.inline_asm( + None, + [address.ir_value(loc=loc, ip=ip), *words], + f"{opcode} [$0], {{$1, $2, $3, $4}};", + "l,r,r,r,r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def store_global_u32x2( + address: Int64, + packed, + *, + loc=None, + ip=None, +) -> None: + words = [packed[i].ir_value(loc=loc, ip=ip) for i in range(2)] + llvm.inline_asm( + None, + [address.ir_value(loc=loc, ip=ip), *words], + "st.global.v2.u32 [$0], {$1, $2};", + "l,r,r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def store_global_u32( + address: Int64, + word: Uint32, + *, + loc=None, + ip=None, +) -> None: + llvm.inline_asm( + None, + [ + address.ir_value(loc=loc, ip=ip), + word.ir_value(loc=loc, ip=ip), + ], + "st.global.u32 [$0], $1;", + "l,r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def store_lamport_sentinel_128(pointer: cute.Pointer, *, loc=None, ip=None) -> None: + """Reset one Lamport fragment to four FP32 negative-zero bit patterns.""" + + address = pointer.toint(loc=loc, ip=ip) + value = Uint32(NEG_ZERO_F32_BITS).ir_value(loc=loc, ip=ip) + llvm.inline_asm( + None, + [address.ir_value(loc=loc, ip=ip), value, value, value, value], + "st.global.v4.u32 [$0], {$1, $2, $3, $4};", + "l,r,r,r,r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def red_async_release_gpu_add_u32(pointer: cute.Pointer, value: Uint32, *, loc=None, ip=None) -> None: + """The exact SM100 arrival primitive used by upstream LamportFlags.""" + + address = pointer.toint(loc=loc, ip=ip) + llvm.inline_asm( + None, + [ + address.ir_value(loc=loc, ip=ip), + value.ir_value(loc=loc, ip=ip), + ], + "red.async.release.global.gpu.add.u32 [$0], $1;", + "l,r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def load_volatile_u32(pointer: cute.Pointer, *, loc=None, ip=None) -> Uint32: + address = pointer.toint(loc=loc, ip=ip) + return Uint32( + llvm.inline_asm( + T.i32(), + [address.ir_value(loc=loc, ip=ip)], + "ld.volatile.global.u32 $0, [$1];", + "=r,l", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def map_shared_to_peer( + smem_ptr: cute.Pointer, + peer_rank: Int32, + *, + loc=None, + ip=None, +) -> Int32: + """Map a local shared-memory slot to the same slot in a peer CTA.""" + + smem_address = smem_ptr.toint(loc=loc, ip=ip).ir_value() + return Int32( + llvm.inline_asm( + T.i32(), + [smem_address, peer_rank.ir_value(loc=loc, ip=ip)], + "mapa.shared::cluster.u32 $0, $1, $2;", + "=r,r,r", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def store_shared_cluster_f32( + remote_address: Int32, + value: Float32, + *, + loc=None, + ip=None, +) -> None: + llvm.inline_asm( + None, + [ + remote_address.ir_value(loc=loc, ip=ip), + value.ir_value(loc=loc, ip=ip), + ], + "st.shared::cluster.f32 [$0], $1;", + "r,f", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def packed_u32x4_to_bf16x8(packed, *, loc=None, ip=None): + target_type = ir.VectorType.get([VEC_BF16], BFloat16.mlir_type, loc=loc) + values = llvm.bitcast( + target_type, + packed.ir_value(loc=loc, ip=ip), + loc=loc, + ip=ip, + ) + return cute.TensorSSA(values, VEC_BF16, BFloat16) + + +@dsl_user_op +def bf16x8_to_packed_u32x4(values, *, loc=None, ip=None): + target_type = ir.VectorType.get([4], T.i32(), loc=loc) + packed = llvm.bitcast( + target_type, + values.ir_value(loc=loc, ip=ip), + loc=loc, + ip=ip, + ) + return cute.TensorSSA(packed, 4, Uint32) + + +@dsl_user_op +def bf16x4_to_packed_u32x2(values, *, loc=None, ip=None): + target_type = ir.VectorType.get([2], T.i32(), loc=loc) + packed = llvm.bitcast( + target_type, + values.ir_value(loc=loc, ip=ip), + loc=loc, + ip=ip, + ) + return cute.TensorSSA(packed, 2, Uint32) + + +@dsl_user_op +def bf16x2_to_u32(values, *, loc=None, ip=None): + packed = llvm.bitcast( + T.i32(), + values.ir_value(loc=loc, ip=ip), + loc=loc, + ip=ip, + ) + return Uint32(packed) + + +@cute.jit +def sanitize_negative_zero_u32(word): + low = Uint16(word & Uint32(0xFFFF)) + high = Uint16(word >> Uint32(16)) + if low == Uint16(NEG_ZERO_BF16_BITS): + word = word & Uint32(0xFFFF0000) + if high == Uint16(NEG_ZERO_BF16_BITS): + word = word & Uint32(0x0000FFFF) + return word + + +@cute.jit +def sanitize_negative_zero_u32x2(packed): + result = cute.make_rmem_tensor(cute.make_layout((2,)), Uint32) + for i in cutlass.range_constexpr(2): + result[i] = sanitize_negative_zero_u32(packed[i]) + return result.load() + + +@cute.jit +def sanitize_negative_zero(packed): + """Turn real BF16 -0 into +0 so it cannot equal the empty sentinel.""" + + result = cute.make_rmem_tensor(cute.make_layout((4,)), Uint32) + for i in cutlass.range_constexpr(4): + word = packed[i] + low = Uint16(word & Uint32(0xFFFF)) + high = Uint16(word >> Uint32(16)) + if low == Uint16(NEG_ZERO_BF16_BITS): + word = word & Uint32(0xFFFF0000) + if high == Uint16(NEG_ZERO_BF16_BITS): + word = word & Uint32(0x0000FFFF) + result[i] = word + return result.load() + + +@cute.jit +def fragment_is_dirty(packed): + """Bit-exact upstream sentinel check: one comparison per 32-bit word.""" + + dirty = packed[0] == Uint32(NEG_ZERO_F32_BITS) + for i in cutlass.range_constexpr(1, 4): + dirty = dirty | (packed[i] == Uint32(NEG_ZERO_F32_BITS)) + return dirty + + +@cute.jit +def warp_sum_specialized( + value: Float32, + warp_idx: Int32, + lane: Int32, + warps: cutlass.Constexpr[int], + last_warp_lanes: cutlass.Constexpr[int], + last_warp_mask: cutlass.Constexpr[int], +) -> Float32: + """Warp sum supporting a compile-time partial final warp.""" + + if warp_idx == Int32(warps - 1) and cutlass.const_expr(last_warp_lanes < 32): + for offset in cutlass.range_constexpr(16, 0, -1): + # range_constexpr does not provide powers-of-two stepping. + if cutlass.const_expr(offset in (16, 8, 4, 2, 1)): + other = cute.arch.shuffle_sync_bfly( + value, + offset=offset, + mask=last_warp_mask, + mask_and_clamp=31, + ) + if (lane ^ Int32(offset)) < Int32(last_warp_lanes): + value = value + other + else: + for offset in cutlass.range_constexpr(16, 0, -1): + if cutlass.const_expr(offset in (16, 8, 4, 2, 1)): + value = value + cute.arch.shuffle_sync_bfly( + value, + offset=offset, + mask=-1, + mask_and_clamp=31, + ) + return value + + +@cute.jit +def block_sum_specialized( + value: Float32, + warp_sums: cute.Tensor, + tidx: Int32, + warps: cutlass.Constexpr[int], + last_warp_lanes: cutlass.Constexpr[int], + last_warp_mask: cutlass.Constexpr[int], +) -> Float32: + """Upstream-equivalent FP32 block reduction.""" + + lane = cute.arch.lane_idx() + warp_idx = cute.arch.warp_idx() + value = warp_sum_specialized(value, warp_idx, lane, warps, last_warp_lanes, last_warp_mask) + if lane == 0: + warp_sums[warp_idx] = value + cute.arch.barrier() + + block_sum = Float32(0.0) + if warp_idx == 0: + if lane < Int32(warps): + block_sum = warp_sums[lane] + block_sum = cute.arch.warp_reduction_sum(block_sum) + if lane == 0: + warp_sums[0] = block_sum + cute.arch.barrier() + return warp_sums[0] diff --git a/aphrodite/models/kimi_k3/nvidia/ops/fused_mla_key_concat_kv_cache.py b/aphrodite/models/kimi_k3/nvidia/ops/fused_mla_key_concat_kv_cache.py new file mode 100644 index 0000000000..c55a46f29f --- /dev/null +++ b/aphrodite/models/kimi_k3/nvidia/ops/fused_mla_key_concat_kv_cache.py @@ -0,0 +1,234 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Fused MLA prefill and decode epilogues for Kimi-K3. + +Thin wrappers over the CUDA ops in +``csrc/libtorch_stable/fused_kimi_k3_mla_key_concat_kv_cache_kernel.cu``, which +mirror ``fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_{bf16,fp8}_insert``. + +- ``fused_mla_key_concat_kv_cache_insert`` (bf16): optionally apply RoPE, + concat the full per-head key ``[k_nope | k_pe]`` into ``k_out``, and insert + the latent ``[kv_c_normed | k_pe]`` into the paged cache. +- ``fused_mla_qkv_quant_kv_cache_fp8_insert`` (fp8): additionally quantize + ``q``/``k``/``v`` to E4M3 with ``q_scale`` / ``k_scale`` / ``v_scale`` (the + cache shares ``k_scale``, as in ``concat_and_cache_mla``). + +The optional ``positions`` / ``cos_sin_cache`` pair enables GPT-J-style RoPE +inside the epilogue. Omitting both keeps the K3 NoPE fast path. The kernels use +Programmatic Dependent Launch to overlap the tail of the producing GEMMs on +sm_90+. +""" + +import torch + + +def fused_mla_key_concat_kv_cache_insert( + q: torch.Tensor, # [Tp, H, qk_head_dim], RoPE is applied in place + k_nope: torch.Tensor, # [Tp, H, qk_nope_head_dim] + k_pe: torch.Tensor, # [Tp, qk_rope_head_dim] or [Tp, 1, qk_rope_head_dim] + kv_c_normed: torch.Tensor, # [Tp, kv_lora_rank] + kv_cache: torch.Tensor, # [num_blocks, block_size, kv_lora_rank + rope] + slot_mapping: torch.Tensor, # [Tp] int64 + positions: torch.Tensor | None = None, # [Tp] int64 + cos_sin_cache: torch.Tensor | None = None, # [max_position, rope] +) -> torch.Tensor: + """Apply optional RoPE, concat K, and insert the paged latent (bf16). + + Returns the full key ``[Tp, H, qk_nope_head_dim + qk_rope_head_dim]``; + optionally rotates ``q`` and writes ``kv_cache`` in place. + """ + k_pe = k_pe.reshape(k_pe.shape[0], -1) + tp, num_heads, qk_nope_head_dim = k_nope.shape + qk_head_dim = qk_nope_head_dim + k_pe.shape[1] + k_out = torch.empty((tp, num_heads, qk_head_dim), dtype=k_nope.dtype, device=k_nope.device) + if tp == 0: + return k_out + torch.ops._C.fused_kimi_k3_mla_key_concat_kv_cache_insert( + q, + k_nope, + k_pe, + kv_c_normed, + k_out, + kv_cache, + slot_mapping, + kv_cache.shape[1], + positions, + cos_sin_cache, + ) + return k_out + + +def fused_mla_key_concat_ds_mla_insert( + q: torch.Tensor, # [Tp, H, qk_head_dim], RoPE is applied in place + k_nope: torch.Tensor, # [Tp, H, qk_nope_head_dim] + k_pe: torch.Tensor, # [Tp, qk_rope_head_dim] or [Tp, 1, qk_rope_head_dim] + kv_c_normed: torch.Tensor, # [Tp, kv_lora_rank] + kv_cache: torch.Tensor, # [num_blocks, block_size, 656] uint8 (fp8_ds_mla) + slot_mapping: torch.Tensor, # [Tp] int64 + positions: torch.Tensor | None = None, # [Tp] int64 + cos_sin_cache: torch.Tensor | None = None, # [max_position, rope] +) -> torch.Tensor: + """Concat full K (bf16) and insert the latent in the fp8_ds_mla layout. + + The cache uses DeepSeek's 656-byte block-scaled layout (NoPE in 4 tiles of + 128 with per-tile dynamic fp8 scales, RoPE as bf16) -- self-scaling, so no + scale argument. Returns the bf16 full key; optionally rotates ``q`` and + writes ``kv_cache`` in place. + """ + k_pe = k_pe.reshape(k_pe.shape[0], -1) + tp, num_heads, qk_nope_head_dim = k_nope.shape + qk_head_dim = qk_nope_head_dim + k_pe.shape[1] + k_out = torch.empty((tp, num_heads, qk_head_dim), dtype=k_nope.dtype, device=k_nope.device) + if tp == 0: + return k_out + torch.ops._C.fused_kimi_k3_mla_key_concat_ds_mla_insert( + q, + k_nope, + k_pe, + kv_c_normed, + k_out, + kv_cache, + slot_mapping, + kv_cache.shape[1], + positions, + cos_sin_cache, + ) + return k_out + + +def fused_mla_qkv_quant_kv_cache_fp8_insert( + q: torch.Tensor, # [Tp, H, qk_head_dim] + k_nope: torch.Tensor, # [Tp, H, qk_nope_head_dim] + k_pe: torch.Tensor, # [Tp, qk_rope_head_dim] or [Tp, 1, qk_rope_head_dim] + kv_c_normed: torch.Tensor, # [Tp, kv_lora_rank] + v: torch.Tensor, # [Tp, H, v_head_dim] + kv_cache: torch.Tensor, # [num_blocks, block_size, kv_lora_rank + rope] fp8 + slot_mapping: torch.Tensor, # [Tp] int64 + q_scale_inv: torch.Tensor, # scalar fp32, 1 / q scale (attention query) + k_scale_inv: torch.Tensor, # scalar fp32, 1 / k scale (attention key) + v_scale_inv: torch.Tensor, # scalar fp32, 1 / v scale (attention value) + cache_scale_inv: torch.Tensor, # scalar fp32, 1 / kv scale (cache latent) + positions: torch.Tensor | None = None, # [Tp] int64 + cos_sin_cache: torch.Tensor | None = None, # [max_position, rope] +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Quantize q/k/v to fp8 and insert the fp8 latent into the paged cache. + + The attention key ``k_fp8`` and the cache latent use *separate* scales + (``k_scale_inv`` vs ``cache_scale_inv``): the cache must be quantized with + ``_k_scale`` (read back by decode / context), while the prefill attention + q/k/v currently stay unscaled (the prefill flash path does not dequantize). + + Returns ``(q_fp8, k_fp8, v_fp8)``; writes the fp8 ``kv_cache`` in place. + """ + k_pe = k_pe.reshape(k_pe.shape[0], -1) + tp, num_heads, _ = q.shape + qk_head_dim = q.shape[2] + v_head_dim = v.shape[2] + fp8 = torch.float8_e4m3fn + q_fp8 = torch.empty((tp, num_heads, qk_head_dim), dtype=fp8, device=q.device) + k_fp8 = torch.empty((tp, num_heads, qk_head_dim), dtype=fp8, device=q.device) + v_fp8 = torch.empty((tp, num_heads, v_head_dim), dtype=fp8, device=q.device) + if tp == 0: + return q_fp8, k_fp8, v_fp8 + torch.ops._C.fused_kimi_k3_mla_qkv_quant_kv_cache_fp8_insert( + q, + k_nope, + k_pe, + kv_c_normed, + v, + q_fp8, + k_fp8, + v_fp8, + kv_cache, + slot_mapping, + q_scale_inv, + k_scale_inv, + v_scale_inv, + cache_scale_inv, + kv_cache.shape[1], + positions, + cos_sin_cache, + ) + return q_fp8, k_fp8, v_fp8 + + +def fused_mla_decode_q_concat_kv_cache_insert( + ql_nope: torch.Tensor, # [B, H, kv_lora_rank] (BMM1 output, absorbed q) + q_pe: torch.Tensor, # [B, H, qk_rope_head_dim] + kv_c_normed: torch.Tensor, # [B, kv_lora_rank] + k_pe: torch.Tensor, # [B, qk_rope_head_dim] or [B, 1, qk_rope_head_dim] + kv_cache: torch.Tensor, # [num_blocks, block_size, entry] + slot_mapping: torch.Tensor, # [B] int64 + *, + ds_mla: bool = False, + q_scale_inv: torch.Tensor | None = None, # scalar fp32, 1 / q scale + cache_scale_inv: torch.Tensor | None = None, # scalar fp32, 1 / kv scale + positions: torch.Tensor | None = None, # [B] int64 + cos_sin_cache: torch.Tensor | None = None, # [max_position, rope] +) -> torch.Tensor: + """Concat the latent decode query ``mqa_q = [ql_nope | q_pe]`` and insert the + latent ``[kv_c_normed | k_pe]`` into the paged cache, in one launch (runs + right before ``forward_mqa``). + + Dispatched by cache format: + - bf16 -> bf16 mqa_q, bf16 cache + - plain fp8 -> fp8 mqa_q (q_scale_inv), fp8 cache (cache_scale_inv) + - fp8_ds_mla -> bf16 mqa_q, 656B block-scaled cache + + Returns ``mqa_q`` of shape ``[B, H, kv_lora_rank + qk_rope_head_dim]``; + writes ``kv_cache`` in place. + """ + k_pe = k_pe.reshape(k_pe.shape[0], -1) + b, num_heads, kv_lora_rank = ql_nope.shape + entry = kv_lora_rank + q_pe.shape[-1] + fp8_q = q_scale_inv is not None + out_dtype = torch.float8_e4m3fn if fp8_q else ql_nope.dtype + mqa_q = torch.empty((b, num_heads, entry), dtype=out_dtype, device=ql_nope.device) + if b == 0: + return mqa_q + + if ds_mla: + cache = kv_cache if kv_cache.dtype == torch.uint8 else kv_cache.view(torch.uint8) + torch.ops._C.fused_kimi_k3_mla_decode_q_concat_ds_mla_insert( + ql_nope, + q_pe, + kv_c_normed, + k_pe, + mqa_q, + cache, + slot_mapping, + cache.shape[1], + positions, + cos_sin_cache, + ) + elif fp8_q: + assert cache_scale_inv is not None, "fp8 decode requires cache_scale_inv" + cache = kv_cache if kv_cache.dtype == torch.float8_e4m3fn else kv_cache.view(torch.float8_e4m3fn) + torch.ops._C.fused_kimi_k3_mla_decode_q_concat_kv_cache_fp8_insert( + ql_nope, + q_pe, + kv_c_normed, + k_pe, + mqa_q, + cache, + slot_mapping, + q_scale_inv, + cache_scale_inv, + cache.shape[1], + positions, + cos_sin_cache, + ) + else: + torch.ops._C.fused_kimi_k3_mla_decode_q_concat_kv_cache_insert( + ql_nope, + q_pe, + kv_c_normed, + k_pe, + mqa_q, + kv_cache, + slot_mapping, + kv_cache.shape[1], + positions, + cos_sin_cache, + ) + return mqa_q diff --git a/aphrodite/models/kimi_k3/nvidia/ops/latent_moe_tail.py b/aphrodite/models/kimi_k3/nvidia/ops/latent_moe_tail.py new file mode 100644 index 0000000000..955429175f --- /dev/null +++ b/aphrodite/models/kimi_k3/nvidia/ops/latent_moe_tail.py @@ -0,0 +1,245 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from dataclasses import dataclass +from functools import partial +from typing import ClassVar + +import torch +import torch.distributed as dist + +from aphrodite.distributed import get_tp_group +from aphrodite.model_executor.warmup.cutedsl_warmup import ( + CuTeDSLCompileUnit, + register_cutedsl_warmup_provider, +) + +_MAX_NUM_TOKENS = 16 +_SKINNY_MAX_NUM_TOKENS = 5 +_MMA_TILER_MN = (64, 32) +_GEMM_CLUSTER_MN = (1, 8) +_B_PRIME_STAGES = 2 +_COLLECTIVE_TOKEN_CTAS = 8 +_LAMPORT_COPY_CTAS = 32 +_LAMPORT_COPY_THREADS = 224 +_SUPPORTED_TP_SIZES = (8, 16) + + +@dataclass(frozen=True) +class KimiK3LatentMoETailContract: + tp_group_id: int + tp_size: int + device: torch.device + dtype: torch.dtype + hidden_size: int + latent_size: int + max_num_tokens: int + rms_eps: float + + +class KimiK3LatentMoETailOp: + """Process-wide cached K3 latent-MoE tail implementation.""" + + _instances: ClassVar[dict[KimiK3LatentMoETailContract, "KimiK3LatentMoETailOp"]] = {} + + @classmethod + def _contract_and_group( + cls, + *, + hidden_size: int, + latent_size: int, + dtype: torch.dtype, + device: torch.device, + rms_eps: float, + ) -> tuple[KimiK3LatentMoETailContract, dist.ProcessGroup]: + tp = get_tp_group() + group = tp.device_group + device = torch.device(device) + tp_device = torch.device(tp.device) + if device != tp_device: + raise ValueError(f"Input device {device} does not match TP device {tp_device}.") + return ( + KimiK3LatentMoETailContract( + tp_group_id=id(group), + tp_size=dist.get_world_size(group), + device=device, + dtype=dtype, + hidden_size=hidden_size, + latent_size=latent_size, + max_num_tokens=_MAX_NUM_TOKENS, + rms_eps=float(rms_eps), + ), + group, + ) + + @classmethod + def initialize( + cls, + *, + hidden_size: int, + latent_size: int, + dtype: torch.dtype, + device: torch.device, + rms_eps: float, + ) -> "KimiK3LatentMoETailOp": + contract, group = cls._contract_and_group( + hidden_size=hidden_size, + latent_size=latent_size, + dtype=dtype, + device=device, + rms_eps=rms_eps, + ) + op = cls._instances.get(contract) + if op is None: + op = cls(contract, group) + cls._instances[contract] = op + return op + + def __init__( + self, + contract: KimiK3LatentMoETailContract, + group: dist.ProcessGroup, + ) -> None: + if contract.tp_size not in _SUPPORTED_TP_SIZES: + raise ValueError(f"K3 latent-MoE tail fusion requires TP in {_SUPPORTED_TP_SIZES}, got {contract.tp_size}.") + if contract.device.type != "cuda": + raise ValueError("K3 latent-MoE tail fusion requires a CUDA device.") + if torch.cuda.get_device_capability(contract.device)[0] != 10: + raise ValueError("K3 latent-MoE tail fusion requires SM100.") + if contract.dtype != torch.bfloat16: + raise ValueError("K3 latent-MoE tail fusion requires bfloat16.") + if (contract.hidden_size, contract.latent_size) != (7168, 3584): + raise ValueError("K3 latent-MoE tail fusion requires hidden_size=7168 and latent_size=3584.") + + self.contract = contract + self.rank = dist.get_rank(group) + from .cute_dsl.latent_moe_tail import ( + AdaptiveUpProjectionKernel, + CollectiveKernel, + LamportCopyKernel, + ) + + with torch.accelerator.device_index(contract.device.index): + self._collective = CollectiveKernel( + group=group, + rank=self.rank, + tp_size=contract.tp_size, + latent_dim=contract.latent_size, + hidden_dim=contract.hidden_size, + max_m=contract.max_num_tokens, + max_token_ctas=_COLLECTIVE_TOKEN_CTAS, + rms_eps=contract.rms_eps, + fp32_internal=False, + ) + self._up_projection = AdaptiveUpProjectionKernel( + group=group, + rank=self.rank, + tp_size=contract.tp_size, + latent_dim=contract.latent_size, + hidden_dim=contract.hidden_size, + max_m=contract.max_num_tokens, + skinny_max_m=_SKINNY_MAX_NUM_TOKENS, + mma_tiler_mn=_MMA_TILER_MN, + cluster_shape_mn=_GEMM_CLUSTER_MN, + b_prime_stages=_B_PRIME_STAGES, + ) + self._lamport_copy = LamportCopyKernel( + hidden_dim=contract.hidden_size, + max_m=contract.max_num_tokens, + ctas=_LAMPORT_COPY_CTAS, + threads=_LAMPORT_COPY_THREADS, + ) + register_cutedsl_warmup_provider(self) + + def get_cutedsl_warmup_compile_units(self) -> tuple[CuTeDSLCompileUnit, ...]: + contract = self.contract + skinny_units = tuple( + CuTeDSLCompileUnit( + name="K3 latent MoE tail Skinny up-projection", + key=( + "k3-latent-moe-tail-skinny-up-projection", + contract, + m, + ), + compile=partial(self._up_projection.compile_skinny, m), + ) + for m in range(1, self._up_projection.skinny_max_m + 1) + ) + return skinny_units + ( + CuTeDSLCompileUnit( + name="K3 latent MoE tail dynamic up-projection", + key=( + "k3-latent-moe-tail-dynamic-up-projection", + contract, + _MMA_TILER_MN, + _GEMM_CLUSTER_MN, + _B_PRIME_STAGES, + ), + compile=self._up_projection.compile_dynamic, + ), + ) + + def __call__( + self, + routed_output: torch.Tensor, + shared_output: torch.Tensor, + rms_weight: torch.Tensor, + up_weight: torch.Tensor, + ) -> torch.Tensor: + self._validate_inputs( + routed_output, + shared_output, + rms_weight, + up_weight, + ) + self._up_projection.ensure_compiled(routed_output.shape[0]) + latent, shared_shard = self._collective( + routed_output, + shared_output, + rms_weight, + ) + local_hidden_size = self.contract.hidden_size // self.contract.tp_size + local_up_weight = up_weight.narrow( + 0, + self.rank * local_hidden_size, + local_hidden_size, + ) + mailbox = self._up_projection( + latent, + local_up_weight, + shared_shard, + ) + return self._lamport_copy( + mailbox, + m=routed_output.shape[0], + ).squeeze(0) + + def _validate_inputs( + self, + routed_output: torch.Tensor, + shared_output: torch.Tensor, + rms_weight: torch.Tensor, + up_weight: torch.Tensor, + ) -> None: + contract = self.contract + if routed_output.ndim != 2: + raise ValueError("routed_output must be a 2D tensor.") + num_tokens = routed_output.shape[0] + if routed_output.shape != (num_tokens, contract.latent_size): + raise ValueError(f"routed_output must have shape [M, {contract.latent_size}].") + if shared_output.shape != (num_tokens, contract.hidden_size): + raise ValueError(f"shared_output must have shape [M, {contract.hidden_size}].") + if rms_weight.shape != (contract.latent_size,): + raise ValueError(f"rms_weight must have shape [{contract.latent_size}].") + if up_weight.shape != (contract.hidden_size, contract.latent_size): + raise ValueError(f"up_weight must have shape [{contract.hidden_size}, {contract.latent_size}].") + if not 1 <= num_tokens <= contract.max_num_tokens: + raise ValueError(f"K3 latent-MoE tail fusion requires between 1 and {contract.max_num_tokens} tokens.") + + tensors = (routed_output, shared_output, rms_weight, up_weight) + if any(tensor.device != contract.device for tensor in tensors): + raise ValueError("All inputs must be on the contract device.") + if any(tensor.dtype != contract.dtype for tensor in tensors): + raise ValueError("All inputs must use the contract dtype.") + if any(not tensor.is_contiguous() for tensor in tensors): + raise ValueError("All inputs must be contiguous.") diff --git a/aphrodite/models/kimi_k3/nvidia/ops/sequence_parallel.py b/aphrodite/models/kimi_k3/nvidia/ops/sequence_parallel.py new file mode 100644 index 0000000000..481621dbc2 --- /dev/null +++ b/aphrodite/models/kimi_k3/nvidia/ops/sequence_parallel.py @@ -0,0 +1,73 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import torch + +from aphrodite.distributed import ( + get_tensor_model_parallel_rank, + get_tensor_model_parallel_world_size, + get_tp_group, + tensor_model_parallel_all_gather, + tensor_model_parallel_reduce_scatter, +) + + +def _custom_collective( + name: str, + x: torch.Tensor, +) -> torch.Tensor | None: + device_communicator = get_tp_group().device_communicator + if device_communicator is None: + return None + collective = getattr(device_communicator, name, None) + if collective is None: + return None + return collective(x) + + +def sp_all_gather(x: torch.Tensor) -> torch.Tensor: + output = _custom_collective("custom_all_gather", x) + if output is not None: + return output + return tensor_model_parallel_all_gather(x, 0) + + +def sp_reduce_scatter(x: torch.Tensor) -> torch.Tensor: + assert x.ndim == 2 + tp_size = get_tensor_model_parallel_world_size() + sp_pad = (-x.shape[0]) % tp_size + if sp_pad > 0: + x = torch.nn.functional.pad(x, (0, 0, 0, sp_pad)) + output = _custom_collective("custom_reduce_scatter", x) + if output is not None: + return output + return tensor_model_parallel_reduce_scatter(x, 0) + + +def sp_shard(x: torch.Tensor) -> torch.Tensor: + assert x.ndim == 2 + tp_size = get_tensor_model_parallel_world_size() + tp_rank = get_tensor_model_parallel_rank() + sp_pad = (-x.shape[0]) % tp_size + if sp_pad > 0: + x = torch.nn.functional.pad(x, (0, 0, 0, sp_pad)) + chunk = x.shape[0] // tp_size + return x[tp_rank * chunk : (tp_rank + 1) * chunk] + + +def sp_padding_mask( + is_padding: torch.Tensor | None, + hidden_states: torch.Tensor, +) -> torch.Tensor: + num_tokens = hidden_states.shape[0] + if is_padding is None: + is_padding = hidden_states.new_zeros(num_tokens, dtype=torch.bool) + assert is_padding.shape[0] == num_tokens + + tp_size = get_tensor_model_parallel_world_size() + sp_pad = (-num_tokens) % tp_size + if sp_pad > 0: + is_padding = torch.nn.functional.pad(is_padding, (0, sp_pad), value=True) + chunk = is_padding.shape[0] // tp_size + tp_rank = get_tensor_model_parallel_rank() + return is_padding[tp_rank * chunk : (tp_rank + 1) * chunk] diff --git a/aphrodite/models/kimi_k3/nvidia/ops/third_party/__init__.py b/aphrodite/models/kimi_k3/nvidia/ops/third_party/__init__.py new file mode 100644 index 0000000000..208f01a7cb --- /dev/null +++ b/aphrodite/models/kimi_k3/nvidia/ops/third_party/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project diff --git a/aphrodite/models/kimi_k3/nvidia/ops/third_party/kda/__init__.py b/aphrodite/models/kimi_k3/nvidia/ops/third_party/kda/__init__.py new file mode 100644 index 0000000000..92936016ab --- /dev/null +++ b/aphrodite/models/kimi_k3/nvidia/ops/third_party/kda/__init__.py @@ -0,0 +1,28 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from .chunk import ( + chunk_kda, + chunk_kda_fwd, + chunk_kda_with_fused_gate, + chunk_kda_with_fused_gate_fwd, + fused_kda_gate, + fused_kda_gate_chunk_cumsum, +) +from .fused_recurrent import ( + fused_recurrent_kda, + fused_recurrent_kda_fwd, + fused_recurrent_kda_packed_decode, +) + +__all__ = [ + "chunk_kda", + "chunk_kda_fwd", + "chunk_kda_with_fused_gate", + "chunk_kda_with_fused_gate_fwd", + "fused_kda_gate", + "fused_kda_gate_chunk_cumsum", + "fused_recurrent_kda", + "fused_recurrent_kda_fwd", + "fused_recurrent_kda_packed_decode", +] diff --git a/aphrodite/models/kimi_k3/nvidia/ops/third_party/kda/chunk.py b/aphrodite/models/kimi_k3/nvidia/ops/third_party/kda/chunk.py new file mode 100644 index 0000000000..5a7860bcf4 --- /dev/null +++ b/aphrodite/models/kimi_k3/nvidia/ops/third_party/kda/chunk.py @@ -0,0 +1,938 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Songlin Yang, Yu Zhang, Zhiyuan Li +# +# This file contains code copied from the flash-linear-attention project. +# The original source code was licensed under the MIT license and included +# the following copyright notice: +# Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li +# ruff: noqa: E501 + + +import torch + +from aphrodite.third_party.flash_linear_attention.ops.chunk_delta_h import ( + chunk_gated_delta_rule_fwd_h, +) +from aphrodite.third_party.flash_linear_attention.ops.cumsum import chunk_local_cumsum +from aphrodite.third_party.flash_linear_attention.ops.index import prepare_chunk_indices +from aphrodite.third_party.flash_linear_attention.ops.l2norm import l2norm_fwd +from aphrodite.third_party.flash_linear_attention.ops.op import exp2, log +from aphrodite.third_party.flash_linear_attention.ops.utils import FLA_CHUNK_SIZE, is_amd +from aphrodite.triton_utils import tl, triton +from aphrodite.utils.math_utils import RCP_LN2, cdiv, next_power_of_2 + +from .chunk_intra import chunk_kda_fwd_intra + +BT_LIST_AUTOTUNE = [32, 64, 128] +NUM_WARPS_AUTOTUNE = [2, 4, 8, 16] if is_amd else [4, 8, 16, 32] + + +@triton.heuristics( + { + "STORE_QG": lambda args: args["qg"] is not None, + "STORE_KG": lambda args: args["kg"] is not None, + "IS_VARLEN": lambda args: args["cu_seqlens"] is not None, + } +) +@triton.autotune( + configs=[ + triton.Config({}, num_warps=num_warps, num_stages=num_stages) + for num_warps in [2, 4, 8] + for num_stages in [2, 3, 4] + ], + key=["H", "K", "V", "BT", "BK", "BV", "IS_VARLEN"], +) +@triton.jit(do_not_specialize=["T"]) +def recompute_w_u_fwd_kernel( + q, + k, + qg, + kg, + v, + beta, + w, + u, + A, + gk, + cu_seqlens, + chunk_indices, + T, + H: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BT: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + STORE_QG: tl.constexpr, + STORE_KG: tl.constexpr, + IS_VARLEN: tl.constexpr, + DOT_PRECISION: tl.constexpr, +): + i_t, i_bh = tl.program_id(0), tl.program_id(1) + i_b, i_h = i_bh // H, i_bh % H + if IS_VARLEN: + i_n, i_t = ( + tl.load(chunk_indices + i_t * 2).to(tl.int32), + tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32), + ) + bos, eos = ( + tl.load(cu_seqlens + i_n).to(tl.int32), + tl.load(cu_seqlens + i_n + 1).to(tl.int32), + ) + T = eos - bos + else: + bos, eos = i_b * T, i_b * T + T + p_b = tl.make_block_ptr(beta + bos * H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) + b_b = tl.load(p_b, boundary_check=(0,)).to(tl.float32) + + p_A = tl.make_block_ptr( + A + (bos * H + i_h) * BT, (T, BT), (H * BT, 1), (i_t * BT, 0), (BT, BT), (1, 0) + ) + b_A = tl.load(p_A, boundary_check=(0, 1)) + + for i_v in range(tl.cdiv(V, BV)): + p_v = tl.make_block_ptr( + v + (bos * H + i_h) * V, + (T, V), + (H * V, 1), + (i_t * BT, i_v * BV), + (BT, BV), + (1, 0), + ) + p_u = tl.make_block_ptr( + u + (bos * H + i_h) * V, + (T, V), + (H * V, 1), + (i_t * BT, i_v * BV), + (BT, BV), + (1, 0), + ) + b_v = tl.load(p_v, boundary_check=(0, 1)) + b_vb = (b_v * b_b[:, None]).to(b_v.dtype) + b_u = tl.dot(b_A, b_vb, input_precision=DOT_PRECISION) + tl.store(p_u, b_u.to(p_u.dtype.element_ty), boundary_check=(0, 1)) + + for i_k in range(tl.cdiv(K, BK)): + p_w = tl.make_block_ptr( + w + (bos * H + i_h) * K, + (T, K), + (H * K, 1), + (i_t * BT, i_k * BK), + (BT, BK), + (1, 0), + ) + p_k = tl.make_block_ptr( + k + (bos * H + i_h) * K, + (T, K), + (H * K, 1), + (i_t * BT, i_k * BK), + (BT, BK), + (1, 0), + ) + b_k = tl.load(p_k, boundary_check=(0, 1)) + b_kb = b_k * b_b[:, None] + + p_gk = tl.make_block_ptr( + gk + (bos * H + i_h) * K, + (T, K), + (H * K, 1), + (i_t * BT, i_k * BK), + (BT, BK), + (1, 0), + ) + b_gk = tl.load(p_gk, boundary_check=(0, 1)) + b_kb *= exp2(b_gk) + if STORE_QG: + p_q = tl.make_block_ptr( + q + (bos * H + i_h) * K, + (T, K), + (H * K, 1), + (i_t * BT, i_k * BK), + (BT, BK), + (1, 0), + ) + p_qg = tl.make_block_ptr( + qg + (bos * H + i_h) * K, + (T, K), + (H * K, 1), + (i_t * BT, i_k * BK), + (BT, BK), + (1, 0), + ) + b_q = tl.load(p_q, boundary_check=(0, 1)) + b_qg = b_q * exp2(b_gk) + tl.store(p_qg, b_qg.to(p_qg.dtype.element_ty), boundary_check=(0, 1)) + if STORE_KG: + last_idx = min(i_t * BT + BT, T) - 1 + + o_k = i_k * BK + tl.arange(0, BK) + m_k = o_k < K + b_gn = tl.load( + gk + ((bos + last_idx) * H + i_h) * K + o_k, mask=m_k, other=0.0 + ) + b_kg = b_k * exp2(b_gn - b_gk) + + p_kg = tl.make_block_ptr( + kg + (bos * H + i_h) * K, + (T, K), + (H * K, 1), + (i_t * BT, i_k * BK), + (BT, BK), + (1, 0), + ) + tl.store(p_kg, b_kg.to(p_kg.dtype.element_ty), boundary_check=(0, 1)) + + b_w = tl.dot(b_A, b_kb.to(b_k.dtype)) + tl.store(p_w, b_w.to(p_w.dtype.element_ty), boundary_check=(0, 1)) + + +def recompute_w_u_fwd( + k: torch.Tensor, + v: torch.Tensor, + beta: torch.Tensor, + A: torch.Tensor, + q: torch.Tensor | None = None, + gk: torch.Tensor | None = None, + cu_seqlens: torch.Tensor | None = None, + chunk_indices: torch.Tensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor]: + B, T, H, K, V = *k.shape, v.shape[-1] + BT = A.shape[-1] + BK = 64 + BV = 64 + + if chunk_indices is None and cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) + NT = cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) + + w = torch.empty_like(k) + u = torch.empty_like(v) + kg = torch.empty_like(k) if gk is not None else None + recompute_w_u_fwd_kernel[(NT, B * H)]( + q=q, + k=k, + qg=None, + kg=kg, + v=v, + beta=beta, + w=w, + u=u, + A=A, + gk=gk, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + T=T, + H=H, + K=K, + V=V, + BT=BT, + BK=BK, + BV=BV, + DOT_PRECISION="ieee", + ) + return w, u, None, kg + + +@triton.heuristics({"IS_VARLEN": lambda args: args["cu_seqlens"] is not None}) +@triton.autotune( + configs=[ + triton.Config({"BK": BK, "BV": BV}, num_warps=num_warps, num_stages=num_stages) + for BK in [32, 64] + for BV in [64, 128] + for num_warps in [2, 4, 8] + for num_stages in [2, 3, 4] + ], + key=["BT"], +) +@triton.jit(do_not_specialize=["T"]) +def chunk_gla_fwd_kernel_o( + q, + v, + g, + h, + o, + A, + cu_seqlens, + chunk_indices, + scale, + T, + H: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BT: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_v, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_b, i_h = i_bh // H, i_bh % H + if IS_VARLEN: + i_tg = i_t + i_n, i_t = ( + tl.load(chunk_indices + i_t * 2).to(tl.int32), + tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32), + ) + bos, eos = ( + tl.load(cu_seqlens + i_n).to(tl.int32), + tl.load(cu_seqlens + i_n + 1).to(tl.int32), + ) + T = eos - bos + NT = tl.cdiv(T, BT) + else: + NT = tl.cdiv(T, BT) + i_tg = i_b * NT + i_t + bos, eos = i_b * T, i_b * T + T + + m_s = tl.arange(0, BT)[:, None] >= tl.arange(0, BT)[None, :] + + b_o = tl.zeros([BT, BV], dtype=tl.float32) + for i_k in range(tl.cdiv(K, BK)): + p_q = tl.make_block_ptr( + q + (bos * H + i_h) * K, + (T, K), + (H * K, 1), + (i_t * BT, i_k * BK), + (BT, BK), + (1, 0), + ) + p_g = tl.make_block_ptr( + g + (bos * H + i_h) * K, + (T, K), + (H * K, 1), + (i_t * BT, i_k * BK), + (BT, BK), + (1, 0), + ) + p_h = tl.make_block_ptr( + h + (i_tg * H + i_h) * K * V, + (V, K), + (K, 1), + (i_v * BV, i_k * BK), + (BV, BK), + (1, 0), + ) + + # [BT, BK] + b_q = tl.load(p_q, boundary_check=(0, 1)) + b_q = (b_q * scale).to(b_q.dtype) + # [BT, BK] + b_g = tl.load(p_g, boundary_check=(0, 1)) + # [BT, BK] + b_qg = (b_q * exp2(b_g)).to(b_q.dtype) + # [BV, BK] + b_h = tl.load(p_h, boundary_check=(0, 1)) + # [BT, BV] + if i_k >= 0: + b_o += tl.dot(b_qg, tl.trans(b_h).to(b_qg.dtype)) + p_v = tl.make_block_ptr( + v + (bos * H + i_h) * V, + (T, V), + (H * V, 1), + (i_t * BT, i_v * BV), + (BT, BV), + (1, 0), + ) + p_o = tl.make_block_ptr( + o + (bos * H + i_h) * V, + (T, V), + (H * V, 1), + (i_t * BT, i_v * BV), + (BT, BV), + (1, 0), + ) + p_A = tl.make_block_ptr( + A + (bos * H + i_h) * BT, (T, BT), (H * BT, 1), (i_t * BT, 0), (BT, BT), (1, 0) + ) + # [BT, BV] + b_v = tl.load(p_v, boundary_check=(0, 1)) + # [BT, BT] + b_A = tl.load(p_A, boundary_check=(0, 1)) + b_A = tl.where(m_s, b_A, 0.0).to(b_v.dtype) + b_o += tl.dot(b_A, b_v, allow_tf32=False) + tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1)) + + +def chunk_gla_fwd_o_gk( + q: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + A: torch.Tensor, + h: torch.Tensor, + o: torch.Tensor, + scale: float, + cu_seqlens: torch.Tensor | None = None, + chunk_indices: torch.Tensor | None = None, + chunk_size: int = FLA_CHUNK_SIZE, +): + B, T, H, K, V = *q.shape, v.shape[-1] + BT = chunk_size + + if chunk_indices is None and cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size) + NT = cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) + + def grid(meta): + return (cdiv(V, meta["BV"]), NT, B * H) + + chunk_gla_fwd_kernel_o[grid]( + q=q, + v=v, + g=g, + h=h, + o=o, + A=A, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + scale=scale, + T=T, + H=H, + K=K, + V=V, + BT=BT, + ) + return o + + +@triton.heuristics( + { + "HAS_BIAS": lambda args: args["g_bias"] is not None, + "IS_VARLEN": lambda args: args["cu_seqlens"] is not None, + } +) +@triton.autotune( + configs=[ + triton.Config({"BS": BS}, num_warps=num_warps) + for BS in [32, 64] + for num_warps in [2, 4, 8] + ], + key=["H", "S", "BT", "IS_VARLEN"], +) +@triton.jit(do_not_specialize=["T"]) +def kda_gate_chunk_cumsum_vector_kernel( + s, + raw_beta, + A_log, + g_bias, + o, + beta_out, + cu_seqlens, + chunk_indices, + cumsum_scale, + lower_bound, + beta, + threshold, + T, + stride_beta_batch, + stride_beta_token, + stride_beta_head, + H: tl.constexpr, + S: tl.constexpr, + BT: tl.constexpr, + BS: tl.constexpr, + HAS_BIAS: tl.constexpr, + IS_VARLEN: tl.constexpr, + USE_LOWER_BOUND: tl.constexpr, +): + i_s, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_b, i_h = i_bh // H, i_bh % H + if IS_VARLEN: + i_n, i_t = ( + tl.load(chunk_indices + i_t * 2).to(tl.int32), + tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32), + ) + bos, eos = ( + tl.load(cu_seqlens + i_n).to(tl.int32), + tl.load(cu_seqlens + i_n + 1).to(tl.int32), + ) + T = eos - bos + else: + bos = i_b * T + + if i_s == 0: + o_beta_t = tl.arange(0, BT) + m_beta = i_t * BT + o_beta_t < T + if IS_VARLEN: + p_beta = ( + raw_beta + + (bos + i_t * BT + o_beta_t) * stride_beta_token + + i_h * stride_beta_head + ) + else: + p_beta = ( + raw_beta + + i_b * stride_beta_batch + + (i_t * BT + o_beta_t) * stride_beta_token + + i_h * stride_beta_head + ) + b_beta = tl.load(p_beta, mask=m_beta, other=0.0).to(tl.float32) + p_beta_out = beta_out + (bos + i_t * BT + o_beta_t) * H + i_h + tl.store(p_beta_out, tl.sigmoid(b_beta), mask=m_beta) + return + + i_s -= 1 + + p_s = tl.make_block_ptr( + s + (bos * H + i_h) * S, + (T, S), + (H * S, 1), + (i_t * BT, i_s * BS), + (BT, BS), + (1, 0), + ) + p_o = tl.make_block_ptr( + o + (bos * H + i_h) * S, + (T, S), + (H * S, 1), + (i_t * BT, i_s * BS), + (BT, BS), + (1, 0), + ) + + b_s = tl.load(p_s, boundary_check=(0, 1)).to(tl.float32) + if HAS_BIAS: + p_bias = tl.make_block_ptr( + g_bias + i_h * S, + (S,), + (1,), + (i_s * BS,), + (BS,), + (0,), + ) + b_bias = tl.load(p_bias, boundary_check=(0,)).to(tl.float32) + b_s += b_bias[None, :] + + b_a = tl.exp(tl.load(A_log + i_h).to(tl.float32)) + if USE_LOWER_BOUND: + b_gate = lower_bound * tl.sigmoid(b_a * b_s) + else: + b_g_scaled = b_s * beta + b_softplus = tl.where( + b_g_scaled > threshold, + b_s, + (1.0 / beta) * log(1.0 + tl.exp(b_g_scaled)), + ) + b_gate = -b_a * b_softplus + + # Boundary loads return zero, but bias and gate activation can make padded + # rows nonzero. Padding trails valid rows, so it only affects masked stores. + b_o = tl.cumsum(b_gate, axis=0) * cumsum_scale + tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1)) + + +def fused_kda_gate_chunk_cumsum( + raw_g: torch.Tensor, + raw_beta: torch.Tensor, + A_log: torch.Tensor, + g_bias: torch.Tensor | None = None, + beta: float = 1.0, + threshold: float = 20.0, + lower_bound: float | None = None, + cu_seqlens: torch.Tensor | None = None, + chunk_indices: torch.Tensor | None = None, + chunk_size: int = FLA_CHUNK_SIZE, + output_dtype: torch.dtype | None = torch.float, +) -> tuple[torch.Tensor, torch.Tensor]: + if cu_seqlens is not None: + assert raw_g.shape[0] == 1, ( + "Only batch size 1 is supported when cu_seqlens are provided" + ) + B, T, H, D = raw_g.shape + if raw_beta.shape != (B, T, H): + raise ValueError( + f"Expected raw_beta shape {(B, T, H)}, got {raw_beta.shape}" + ) + if chunk_indices is None and cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size) + NT = cdiv(T, chunk_size) if cu_seqlens is None else len(chunk_indices) + + A_log = A_log.reshape(-1) + if g_bias is not None: + g_bias = g_bias.reshape(-1) + y = torch.empty_like(raw_g, dtype=output_dtype or raw_g.dtype) + beta_out = torch.empty(raw_beta.shape, device=raw_beta.device, dtype=torch.float32) + + def grid(meta): + # For each (chunk, head), program 0 computes beta without extending a + # gate tile's critical path. The remaining programs cover the gate dim. + return (cdiv(meta["S"], meta["BS"]) + 1, NT, B * H) + + kda_gate_chunk_cumsum_vector_kernel[grid]( + s=raw_g, + raw_beta=raw_beta, + A_log=A_log, + g_bias=g_bias, + o=y, + beta_out=beta_out, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + # RCP_LN2 folds in the natural-log -> log2 conversion so downstream + # exp2-based kernels reproduce exp(g). Keep this in sync with the + # `use_exp2=True` path in `_chunk_kda_fwd_with_cumulative_g`. + cumsum_scale=RCP_LN2, + lower_bound=lower_bound or 0.0, + beta=beta, + threshold=threshold, + T=T, + stride_beta_batch=raw_beta.stride(0), + stride_beta_token=raw_beta.stride(1), + stride_beta_head=raw_beta.stride(2), + H=H, + S=D, + BT=chunk_size, + USE_LOWER_BOUND=lower_bound is not None, + ) + return y, beta_out + + +def _chunk_kda_fwd_with_cumulative_g( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + scale: float, + initial_state: torch.Tensor, + output_final_state: bool, + cu_seqlens: torch.Tensor | None = None, + chunk_indices: torch.Tensor | None = None, + chunk_size: int = FLA_CHUNK_SIZE, + safe_gate: bool = False, +): + # `g` must already be chunk-local cumulatively-summed AND scaled by + # RCP_LN2 (so the downstream exp2-based kernels reproduce exp(g)). + # Use `chunk_kda_fwd` or `chunk_kda_with_fused_gate_fwd` instead of + # calling this helper directly unless that invariant is upheld. + Aqk, A = chunk_kda_fwd_intra( + q=q, + k=k, + gk=g, + beta=beta, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + chunk_size=chunk_size, + safe_gate=safe_gate, + ) + w, u, _, kg = recompute_w_u_fwd( + k=k, + v=v, + beta=beta, + A=A, + gk=g, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + ) + del A + h, v_new, final_state = chunk_gated_delta_rule_fwd_h( + k=kg, + w=w, + u=u, + gk=g, + initial_state=initial_state, + output_final_state=output_final_state, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + use_exp2=True, + ) + del w, u, kg + o = chunk_gla_fwd_o_gk( + q=q, + v=v_new, + g=g, + A=Aqk, + h=h, + o=v, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + chunk_size=chunk_size, + ) + del Aqk, v_new, h + return o, final_state + + +def chunk_kda_fwd( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + scale: float, + initial_state: torch.Tensor, + output_final_state: bool, + cu_seqlens: torch.Tensor | None = None, +): + chunk_size = FLA_CHUNK_SIZE + chunk_indices = ( + prepare_chunk_indices(cu_seqlens, chunk_size) + if cu_seqlens is not None + else None + ) + g = chunk_local_cumsum( + g, + chunk_size=chunk_size, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + ) + # KDA evaluates cumulative gate decays with exp2. Convert from natural-log + # space so exp(x) is preserved as exp2(x / ln(2)). + g = g * RCP_LN2 + return _chunk_kda_fwd_with_cumulative_g( + q=q, + k=k, + v=v, + g=g, + beta=beta, + scale=scale, + initial_state=initial_state, + output_final_state=output_final_state, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + chunk_size=chunk_size, + ) + + +def chunk_kda_with_fused_gate_fwd( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + raw_g: torch.Tensor, + raw_beta: torch.Tensor, + A_log: torch.Tensor, + g_bias: torch.Tensor | None, + scale: float, + initial_state: torch.Tensor, + output_final_state: bool, + lower_bound: float | None = None, + cu_seqlens: torch.Tensor | None = None, +): + chunk_size = FLA_CHUNK_SIZE + chunk_indices = ( + prepare_chunk_indices(cu_seqlens, chunk_size) + if cu_seqlens is not None + else None + ) + g, beta = fused_kda_gate_chunk_cumsum( + raw_g, + raw_beta=raw_beta, + A_log=A_log, + g_bias=g_bias, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + chunk_size=chunk_size, + lower_bound=lower_bound, + ) + return _chunk_kda_fwd_with_cumulative_g( + q=q, + k=k, + v=v, + g=g, + beta=beta, + scale=scale, + initial_state=initial_state, + output_final_state=output_final_state, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + chunk_size=chunk_size, + safe_gate=lower_bound is not None, + ) + + +def chunk_kda( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + scale: float = None, + initial_state: torch.Tensor = None, + output_final_state: bool = False, + use_qk_l2norm_in_kernel: bool = False, + cu_seqlens: torch.Tensor | None = None, + **kwargs, +): + if scale is None: + scale = k.shape[-1] ** -0.5 + + if use_qk_l2norm_in_kernel: + q = l2norm_fwd(q.contiguous()) + k = l2norm_fwd(k.contiguous()) + + o, final_state = chunk_kda_fwd( + q=q, + k=k, + v=v.contiguous(), + g=g.contiguous(), + beta=beta.contiguous(), + scale=scale, + initial_state=initial_state.contiguous(), + output_final_state=output_final_state, + cu_seqlens=cu_seqlens, + ) + return o, final_state + + +def chunk_kda_with_fused_gate( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + raw_g: torch.Tensor, + raw_beta: torch.Tensor, + A_log: torch.Tensor, + g_bias: torch.Tensor | None, + scale: float | None = None, + initial_state: torch.Tensor | None = None, + output_final_state: bool = False, + lower_bound: float | None = None, + use_qk_l2norm_in_kernel: bool = False, + cu_seqlens: torch.Tensor | None = None, + **kwargs, +): + """Run chunk KDA from raw gate and beta projections.""" + if scale is None: + scale = k.shape[-1] ** -0.5 + + if use_qk_l2norm_in_kernel: + q = l2norm_fwd(q.contiguous()) + k = l2norm_fwd(k.contiguous()) + + o, final_state = chunk_kda_with_fused_gate_fwd( + q=q, + k=k, + v=v.contiguous(), + raw_g=raw_g.contiguous(), + raw_beta=raw_beta, + A_log=A_log, + g_bias=g_bias, + scale=scale, + initial_state=initial_state.contiguous() if initial_state is not None else None, + output_final_state=output_final_state, + lower_bound=lower_bound, + cu_seqlens=cu_seqlens, + ) + return o, final_state + + +@triton.autotune( + configs=[ + triton.Config({"BT": bt}, num_warps=nw, num_stages=ns) + for bt in BT_LIST_AUTOTUNE + for nw in NUM_WARPS_AUTOTUNE + for ns in [2, 3] + ], + key=["H", "D"], +) +@triton.jit +def kda_gate_fwd_kernel( + g, + A, + y, + g_bias, + lower_bound, + beta: tl.constexpr, + threshold: tl.constexpr, + T, + H, + D: tl.constexpr, + BT: tl.constexpr, + BD: tl.constexpr, + HAS_BIAS: tl.constexpr, + USE_LOWER_BOUND: tl.constexpr, +): + i_t, i_h = tl.program_id(0), tl.program_id(1) + n_t = i_t * BT + + b_a = tl.exp(tl.load(A + i_h).to(tl.float32)) + + stride_row = H * D + stride_col = 1 + + g_ptr = tl.make_block_ptr( + base=g + i_h * D, + shape=(T, D), + strides=(stride_row, stride_col), + offsets=(n_t, 0), + block_shape=(BT, BD), + order=(1, 0), + ) + + y_ptr = tl.make_block_ptr( + base=y + i_h * D, + shape=(T, D), + strides=(stride_row, stride_col), + offsets=(n_t, 0), + block_shape=(BT, BD), + order=(1, 0), + ) + + b_g = tl.load(g_ptr, boundary_check=(0, 1)).to(tl.float32) + + if HAS_BIAS: + n_d = tl.arange(0, BD) + bias_mask = n_d < D + b_bias = tl.load(g_bias + i_h * D + n_d, mask=bias_mask, other=0.0).to( + tl.float32 + ) + b_g = b_g + b_bias[None, :] + + if USE_LOWER_BOUND: + b_y = lower_bound * tl.sigmoid(b_a * b_g) + else: + g_scaled = b_g * beta + use_linear = g_scaled > threshold + sp = tl.where(use_linear, b_g, (1.0 / beta) * log(1.0 + tl.exp(g_scaled))) + b_y = -b_a * sp + + tl.store(y_ptr, b_y.to(y.dtype.element_ty), boundary_check=(0, 1)) + + +def fused_kda_gate( + g: torch.Tensor, + A: torch.Tensor, + head_k_dim: int, + g_bias: torch.Tensor | None = None, + beta: float = 1.0, + threshold: float = 20.0, + lower_bound: float | None = None, +) -> torch.Tensor: + """ + Forward pass for KDA gate: + input g: [..., H*D] + param A: [H] or [1, 1, H, 1] + beta: softplus beta parameter + threshold: softplus threshold parameter + return : [..., H, D] + """ + orig_shape = g.shape[:-1] + + g = g.view(-1, g.shape[-1]) + T = g.shape[0] + HD = g.shape[1] + H = A.numel() + assert H * head_k_dim == HD + assert g.stride() == (HD, 1) + + y = torch.empty_like(g, dtype=torch.float32) + + def grid(meta): + return (cdiv(T, meta["BT"]), H) + + kda_gate_fwd_kernel[grid]( + g, + A, + y, + g_bias, + lower_bound or 0.0, + beta, + threshold, + T, + H, + head_k_dim, + BD=next_power_of_2(head_k_dim), + HAS_BIAS=g_bias is not None, + USE_LOWER_BOUND=lower_bound is not None, + ) + + y = y.view(*orig_shape, H, head_k_dim) + return y diff --git a/aphrodite/models/kimi_k3/nvidia/ops/third_party/kda/chunk_intra.py b/aphrodite/models/kimi_k3/nvidia/ops/third_party/kda/chunk_intra.py new file mode 100644 index 0000000000..efd114a5bf --- /dev/null +++ b/aphrodite/models/kimi_k3/nvidia/ops/third_party/kda/chunk_intra.py @@ -0,0 +1,559 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Songlin Yang, Yu Zhang, Zhiyuan Li +# +# This file contains code copied from the flash-linear-attention project. +# The original source was licensed under the MIT license. +# Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li +# Forward-only adaptation of flash-linear-attention 0.5.0. +# ruff: noqa: E501 + +import torch + +from aphrodite.platforms import current_platform +from aphrodite.third_party.flash_linear_attention.ops.index import prepare_chunk_indices +from aphrodite.third_party.flash_linear_attention.ops.op import exp2, gather +from aphrodite.third_party.flash_linear_attention.ops.utils import is_gather_supported +from aphrodite.triton_utils import tl, triton + +from .chunk_intra_token_parallel import chunk_kda_fwd_intra_token_parallel + +################################################################################ +# Fused inter + solve_tril kernel: compute off-diagonal Akk and solve in one pass +################################################################################ + + +@triton.heuristics({ + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, +}) +@triton.autotune( + configs=[ + triton.Config({'BK': BK}, num_warps=num_warps) + for BK in [32, 64] + for num_warps in [1, 2, 4] + ], + key=["H", "HV", "K", "BC"], +) +@triton.jit(do_not_specialize=['T']) +def chunk_kda_fwd_kernel_inter_solve_fused( + q, + k, + g, + beta, + Aqk, + Akkd, + Akk, + scale, + cu_seqlens, + chunk_indices, + T, + H: tl.constexpr, + HV: tl.constexpr, + K: tl.constexpr, + BT: tl.constexpr, + BC: tl.constexpr, + BK: tl.constexpr, + IS_VARLEN: tl.constexpr, + USE_SAFE_GATE: tl.constexpr, + SOLVE_TRIL_DOT_PRECISION: tl.constexpr, +): + """ + Fused kernel: compute inter-subchunk Akk + solve_tril in one pass. + Prerequisite: token_parallel has already computed diagonal Akk blocks in Akkd. + + This kernel: + 1. Computes off-diagonal Aqk blocks -> writes to global + 2. Computes off-diagonal Akk blocks -> keeps in registers + 3. Loads diagonal Akk blocks from Akkd (fp32) + 4. Does forward substitution on diagonals + 5. Computes merged Akk_inv + 6. Writes Akk_inv to Akk + """ + i_t, i_bh = tl.program_id(0), tl.program_id(1) + i_b, i_hv = i_bh // HV, i_bh % HV + i_h = i_hv // (HV // H) + + if IS_VARLEN: + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + T = eos - bos + else: + bos, eos = i_b * T, i_b * T + T + + if i_t * BT >= T: + return + + i_tc0 = i_t * BT + i_tc1 = i_t * BT + BC + i_tc2 = i_t * BT + 2 * BC + i_tc3 = i_t * BT + 3 * BC + + q += (bos * H + i_h) * K + k += (bos * H + i_h) * K + g += (bos * HV + i_hv) * K + Aqk += (bos * HV + i_hv) * BT + Akk += (bos * HV + i_hv) * BT + Akkd += (bos * HV + i_hv) * BC + + o_i = tl.arange(0, BC) + m_tc1 = (i_tc1 + o_i) < T + m_tc2 = (i_tc2 + o_i) < T + m_tc3 = (i_tc3 + o_i) < T + + b_Aqk10 = tl.zeros([BC, BC], dtype=tl.float32) + b_Akk10 = tl.zeros([BC, BC], dtype=tl.float32) + + b_Aqk20 = tl.zeros([BC, BC], dtype=tl.float32) + b_Akk20 = tl.zeros([BC, BC], dtype=tl.float32) + b_Aqk21 = tl.zeros([BC, BC], dtype=tl.float32) + b_Akk21 = tl.zeros([BC, BC], dtype=tl.float32) + + b_Aqk30 = tl.zeros([BC, BC], dtype=tl.float32) + b_Akk30 = tl.zeros([BC, BC], dtype=tl.float32) + b_Aqk31 = tl.zeros([BC, BC], dtype=tl.float32) + b_Akk31 = tl.zeros([BC, BC], dtype=tl.float32) + b_Aqk32 = tl.zeros([BC, BC], dtype=tl.float32) + b_Akk32 = tl.zeros([BC, BC], dtype=tl.float32) + + ################################################################################ + # off-diagonal blocks + ################################################################################ + for i_k in range(tl.cdiv(K, BK)): + o_k = i_k * BK + tl.arange(0, BK) + m_k = o_k < K + + p_k0 = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_tc0, i_k * BK), (BC, BK), (1, 0)) + p_g0 = tl.make_block_ptr(g, (T, K), (HV*K, 1), (i_tc0, i_k * BK), (BC, BK), (1, 0)) + b_k0 = tl.load(p_k0, boundary_check=(0, 1)).to(tl.float32) + b_g0 = tl.load(p_g0, boundary_check=(0, 1)).to(tl.float32) + + if i_tc1 < T: + p_q1 = tl.make_block_ptr(q, (T, K), (H*K, 1), (i_tc1, i_k * BK), (BC, BK), (1, 0)) + p_k1 = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_tc1, i_k * BK), (BC, BK), (1, 0)) + p_g1 = tl.make_block_ptr(g, (T, K), (HV*K, 1), (i_tc1, i_k * BK), (BC, BK), (1, 0)) + # [BC, BK] + b_q1 = tl.load(p_q1, boundary_check=(0, 1)).to(tl.float32) + b_k1 = tl.load(p_k1, boundary_check=(0, 1)).to(tl.float32) + b_g1 = tl.load(p_g1, boundary_check=(0, 1)).to(tl.float32) + # [BK] + b_gn1 = tl.load(g + i_tc1 * HV*K + o_k, mask=m_k, other=0).to(tl.float32) + # [BC, BK] + b_gqn = tl.where(m_tc1[:, None], exp2(b_g1 - b_gn1[None, :]), 0) + # [BK, BC] + b_kgt = tl.trans(b_k0 * exp2(b_gn1[None, :] - b_g0)) + # [BC, BC] + b_Aqk10 += tl.dot(b_q1 * b_gqn, b_kgt) + b_Akk10 += tl.dot(b_k1 * b_gqn, b_kgt) + + if i_tc2 < T: + p_q2 = tl.make_block_ptr(q, (T, K), (H*K, 1), (i_tc2, i_k * BK), (BC, BK), (1, 0)) + p_k2 = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_tc2, i_k * BK), (BC, BK), (1, 0)) + p_g2 = tl.make_block_ptr(g, (T, K), (HV*K, 1), (i_tc2, i_k * BK), (BC, BK), (1, 0)) + # [BC, BK] + b_q2 = tl.load(p_q2, boundary_check=(0, 1)).to(tl.float32) + b_k2 = tl.load(p_k2, boundary_check=(0, 1)).to(tl.float32) + b_g2 = tl.load(p_g2, boundary_check=(0, 1)).to(tl.float32) + # [BK] + b_gn2 = tl.load(g + i_tc2 * HV*K + o_k, mask=m_k, other=0).to(tl.float32) + # [BC, BK] + b_gqn2 = tl.where(m_tc2[:, None], exp2(b_g2 - b_gn2[None, :]), 0) + b_qg2 = b_q2 * b_gqn2 + b_kg2 = b_k2 * b_gqn2 + # [BK, BC] + b_kgt = tl.trans(b_k0 * exp2(b_gn2[None, :] - b_g0)) + b_Aqk20 += tl.dot(b_qg2, b_kgt) + b_Akk20 += tl.dot(b_kg2, b_kgt) + # [BC, BC] + b_kgt = tl.trans(b_k1 * exp2(b_gn2[None, :] - b_g1)) + # [BC, BC] + b_Aqk21 += tl.dot(b_qg2, b_kgt) + b_Akk21 += tl.dot(b_kg2, b_kgt) + + if i_tc3 < T: + p_q3 = tl.make_block_ptr(q, (T, K), (H*K, 1), (i_tc3, i_k * BK), (BC, BK), (1, 0)) + p_k3 = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_tc3, i_k * BK), (BC, BK), (1, 0)) + p_g3 = tl.make_block_ptr(g, (T, K), (HV*K, 1), (i_tc3, i_k * BK), (BC, BK), (1, 0)) + # [BC, BK] + b_q3 = tl.load(p_q3, boundary_check=(0, 1)).to(tl.float32) + b_k3 = tl.load(p_k3, boundary_check=(0, 1)).to(tl.float32) + b_g3 = tl.load(p_g3, boundary_check=(0, 1)).to(tl.float32) + # [BK] + b_gn3 = tl.load(g + i_tc3 * HV*K + o_k, mask=m_k, other=0).to(tl.float32) + # [BC, BK] + b_gqn3 = tl.where(m_tc3[:, None], exp2(b_g3 - b_gn3[None, :]), 0) + b_qg3 = b_q3 * b_gqn3 + b_kg3 = b_k3 * b_gqn3 + # [BK, BC] + b_kgt = tl.trans(b_k0 * exp2(b_gn3[None, :] - b_g0)) + # [BC, BC] + b_Aqk30 += tl.dot(b_qg3, b_kgt) + b_Akk30 += tl.dot(b_kg3, b_kgt) + # [BK, BC] + b_kgt = tl.trans(b_k1 * exp2(b_gn3[None, :] - b_g1)) + # [BC, BC] + b_Aqk31 += tl.dot(b_qg3, b_kgt) + b_Akk31 += tl.dot(b_kg3, b_kgt) + # [BK, BC] + b_kgt = tl.trans(b_k2 * exp2(b_gn3[None, :] - b_g2)) + # [BC, BC] + b_Aqk32 += tl.dot(b_qg3, b_kgt) + b_Akk32 += tl.dot(b_kg3, b_kgt) + + ################################################################################ + # save off-diagonal Aqk blocks and prepare Akk + ################################################################################ + if i_tc1 < T: + p_Aqk10 = tl.make_block_ptr(Aqk, (T, BT), (HV*BT, 1), (i_tc1, 0), (BC, BC), (1, 0)) + tl.store(p_Aqk10, (b_Aqk10 * scale).to(Aqk.dtype.element_ty), boundary_check=(0, 1)) + + p_b1 = tl.make_block_ptr(beta + bos * HV + i_hv, (T,), (HV,), (i_tc1,), (BC,), (0,)) + b_b1 = tl.load(p_b1, boundary_check=(0,)).to(tl.float32) + b_Akk10 = b_Akk10 * b_b1[:, None] + if i_tc2 < T: + p_Aqk20 = tl.make_block_ptr(Aqk, (T, BT), (HV*BT, 1), (i_tc2, 0), (BC, BC), (1, 0)) + p_Aqk21 = tl.make_block_ptr(Aqk, (T, BT), (HV*BT, 1), (i_tc2, BC), (BC, BC), (1, 0)) + tl.store(p_Aqk20, (b_Aqk20 * scale).to(Aqk.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_Aqk21, (b_Aqk21 * scale).to(Aqk.dtype.element_ty), boundary_check=(0, 1)) + + p_b2 = tl.make_block_ptr(beta + bos * HV + i_hv, (T,), (HV,), (i_tc2,), (BC,), (0,)) + b_b2 = tl.load(p_b2, boundary_check=(0,)).to(tl.float32) + b_Akk20 = b_Akk20 * b_b2[:, None] + b_Akk21 = b_Akk21 * b_b2[:, None] + if i_tc3 < T: + p_Aqk30 = tl.make_block_ptr(Aqk, (T, BT), (HV*BT, 1), (i_tc3, 0), (BC, BC), (1, 0)) + p_Aqk31 = tl.make_block_ptr(Aqk, (T, BT), (HV*BT, 1), (i_tc3, BC), (BC, BC), (1, 0)) + p_Aqk32 = tl.make_block_ptr(Aqk, (T, BT), (HV*BT, 1), (i_tc3, 2*BC), (BC, BC), (1, 0)) + tl.store(p_Aqk30, (b_Aqk30 * scale).to(Aqk.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_Aqk31, (b_Aqk31 * scale).to(Aqk.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_Aqk32, (b_Aqk32 * scale).to(Aqk.dtype.element_ty), boundary_check=(0, 1)) + + p_b3 = tl.make_block_ptr(beta + bos * HV + i_hv, (T,), (HV,), (i_tc3,), (BC,), (0,)) + b_b3 = tl.load(p_b3, boundary_check=(0,)).to(tl.float32) + b_Akk30 = b_Akk30 * b_b3[:, None] + b_Akk31 = b_Akk31 * b_b3[:, None] + b_Akk32 = b_Akk32 * b_b3[:, None] + + p_Akk00 = tl.make_block_ptr(Akkd, (T, BC), (HV*BC, 1), (i_tc0, 0), (BC, BC), (1, 0)) + p_Akk11 = tl.make_block_ptr(Akkd, (T, BC), (HV*BC, 1), (i_tc1, 0), (BC, BC), (1, 0)) + p_Akk22 = tl.make_block_ptr(Akkd, (T, BC), (HV*BC, 1), (i_tc2, 0), (BC, BC), (1, 0)) + p_Akk33 = tl.make_block_ptr(Akkd, (T, BC), (HV*BC, 1), (i_tc3, 0), (BC, BC), (1, 0)) + b_Ai00 = tl.load(p_Akk00, boundary_check=(0, 1)).to(tl.float32) + b_Ai11 = tl.load(p_Akk11, boundary_check=(0, 1)).to(tl.float32) + b_Ai22 = tl.load(p_Akk22, boundary_check=(0, 1)).to(tl.float32) + b_Ai33 = tl.load(p_Akk33, boundary_check=(0, 1)).to(tl.float32) + + ################################################################################ + # forward substitution on diagonals + ################################################################################ + + if not USE_SAFE_GATE: + m_A = o_i[:, None] > o_i[None, :] + m_I = o_i[:, None] == o_i[None, :] + + b_Ai00 = -tl.where(m_A, b_Ai00, 0) + b_Ai11 = -tl.where(m_A, b_Ai11, 0) + b_Ai22 = -tl.where(m_A, b_Ai22, 0) + b_Ai33 = -tl.where(m_A, b_Ai33, 0) + + for i in range(2, min(BC, T - i_tc0)): + b_a00 = -tl.load(Akkd + (i_tc0 + i) * HV*BC + o_i) + b_a00 = tl.where(o_i < i, b_a00, 0.) + b_a00 += tl.sum(b_a00[:, None] * b_Ai00, 0) + b_Ai00 = tl.where((o_i == i)[:, None], b_a00, b_Ai00) + for i in range(BC + 2, min(2*BC, T - i_tc0)): + b_a11 = -tl.load(Akkd + (i_tc0 + i) * HV*BC + o_i) + b_a11 = tl.where(o_i < i - BC, b_a11, 0.) + b_a11 += tl.sum(b_a11[:, None] * b_Ai11, 0) + b_Ai11 = tl.where((o_i == i - BC)[:, None], b_a11, b_Ai11) + for i in range(2*BC + 2, min(3*BC, T - i_tc0)): + b_a22 = -tl.load(Akkd + (i_tc0 + i) * HV*BC + o_i) + b_a22 = tl.where(o_i < i - 2*BC, b_a22, 0.) + b_a22 += tl.sum(b_a22[:, None] * b_Ai22, 0) + b_Ai22 = tl.where((o_i == i - 2*BC)[:, None], b_a22, b_Ai22) + for i in range(3*BC + 2, min(4*BC, T - i_tc0)): + b_a33 = -tl.load(Akkd + (i_tc0 + i) * HV*BC + o_i) + b_a33 = tl.where(o_i < i - 3*BC, b_a33, 0.) + b_a33 += tl.sum(b_a33[:, None] * b_Ai33, 0) + b_Ai33 = tl.where((o_i == i - 3*BC)[:, None], b_a33, b_Ai33) + + b_Ai00 += m_I + b_Ai11 += m_I + b_Ai22 += m_I + b_Ai33 += m_I + + ################################################################################ + # compute merged inverse using off-diagonals + ################################################################################ + + # we used tf32 to maintain matrix inverse's precision whenever possible. + b_Ai10 = -tl.dot( + tl.dot(b_Ai11, b_Akk10, input_precision=SOLVE_TRIL_DOT_PRECISION), + b_Ai00, + input_precision=SOLVE_TRIL_DOT_PRECISION + ) + b_Ai21 = -tl.dot( + tl.dot(b_Ai22, b_Akk21, input_precision=SOLVE_TRIL_DOT_PRECISION), + b_Ai11, + input_precision=SOLVE_TRIL_DOT_PRECISION + ) + b_Ai32 = -tl.dot( + tl.dot(b_Ai33, b_Akk32, input_precision=SOLVE_TRIL_DOT_PRECISION), + b_Ai22, + input_precision=SOLVE_TRIL_DOT_PRECISION + ) + + b_Ai20 = -tl.dot( + b_Ai22, + tl.dot(b_Akk20, b_Ai00, input_precision=SOLVE_TRIL_DOT_PRECISION) + + tl.dot(b_Akk21, b_Ai10, input_precision=SOLVE_TRIL_DOT_PRECISION), + input_precision=SOLVE_TRIL_DOT_PRECISION + ) + b_Ai31 = -tl.dot( + b_Ai33, + tl.dot(b_Akk31, b_Ai11, input_precision=SOLVE_TRIL_DOT_PRECISION) + + tl.dot(b_Akk32, b_Ai21, input_precision=SOLVE_TRIL_DOT_PRECISION), + input_precision=SOLVE_TRIL_DOT_PRECISION + ) + b_Ai30 = -tl.dot( + b_Ai33, + tl.dot(b_Akk30, b_Ai00, input_precision=SOLVE_TRIL_DOT_PRECISION) + + tl.dot(b_Akk31, b_Ai10, input_precision=SOLVE_TRIL_DOT_PRECISION) + + tl.dot(b_Akk32, b_Ai20, input_precision=SOLVE_TRIL_DOT_PRECISION), + input_precision=SOLVE_TRIL_DOT_PRECISION + ) + + ################################################################################ + # store full Akk_inv to Akk + ################################################################################ + + p_Akk00 = tl.make_block_ptr(Akk, (T, BT), (HV*BT, 1), (i_tc0, 0), (BC, BC), (1, 0)) + p_Akk10 = tl.make_block_ptr(Akk, (T, BT), (HV*BT, 1), (i_tc1, 0), (BC, BC), (1, 0)) + p_Akk11 = tl.make_block_ptr(Akk, (T, BT), (HV*BT, 1), (i_tc1, BC), (BC, BC), (1, 0)) + p_Akk20 = tl.make_block_ptr(Akk, (T, BT), (HV*BT, 1), (i_tc2, 0), (BC, BC), (1, 0)) + p_Akk21 = tl.make_block_ptr(Akk, (T, BT), (HV*BT, 1), (i_tc2, BC), (BC, BC), (1, 0)) + p_Akk22 = tl.make_block_ptr(Akk, (T, BT), (HV*BT, 1), (i_tc2, 2*BC), (BC, BC), (1, 0)) + p_Akk30 = tl.make_block_ptr(Akk, (T, BT), (HV*BT, 1), (i_tc3, 0), (BC, BC), (1, 0)) + p_Akk31 = tl.make_block_ptr(Akk, (T, BT), (HV*BT, 1), (i_tc3, BC), (BC, BC), (1, 0)) + p_Akk32 = tl.make_block_ptr(Akk, (T, BT), (HV*BT, 1), (i_tc3, 2*BC), (BC, BC), (1, 0)) + p_Akk33 = tl.make_block_ptr(Akk, (T, BT), (HV*BT, 1), (i_tc3, 3*BC), (BC, BC), (1, 0)) + + tl.store(p_Akk00, b_Ai00.to(Akk.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_Akk10, b_Ai10.to(Akk.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_Akk11, b_Ai11.to(Akk.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_Akk20, b_Ai20.to(Akk.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_Akk21, b_Ai21.to(Akk.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_Akk22, b_Ai22.to(Akk.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_Akk30, b_Ai30.to(Akk.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_Akk31, b_Ai31.to(Akk.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_Akk32, b_Ai32.to(Akk.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_Akk33, b_Ai33.to(Akk.dtype.element_ty), boundary_check=(0, 1)) + + +@triton.heuristics({ + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, +}) +@triton.autotune( + configs=[ + triton.Config({}, num_warps=num_warps, num_stages=num_stages) + for num_warps in [1, 2, 4, 8] + for num_stages in [2, 3, 4] + ], + key=['BK', 'NC', 'BT', 'HV'], +) +@triton.jit(do_not_specialize=['B', 'T']) +def chunk_kda_fwd_kernel_intra_sub_chunk( + q, + k, + g, + beta, + Aqk, + Akk, + scale, + cu_seqlens, + chunk_indices, + T, + H: tl.constexpr, + HV: tl.constexpr, + K: tl.constexpr, + BT: tl.constexpr, + BC: tl.constexpr, + BK: tl.constexpr, + IS_VARLEN: tl.constexpr, + USE_GATHER: tl.constexpr, +): + i_t, i_i, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_b, i_hv = i_bh // HV, i_bh % HV + i_h = i_hv // (HV // H) + + if IS_VARLEN: + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + T = eos - bos + else: + bos, eos = i_b * T, i_b * T + T + + i_ti = i_t * BT + i_i * BC + if i_ti >= T: + return + + o_c = i_ti + tl.arange(0, BC) + m_c = o_c < T + + q = q + (bos * H + i_h) * K + k = k + (bos * H + i_h) * K + g = g + (bos * HV + i_hv) * K + beta = beta + bos * HV + i_hv + Aqk = Aqk + (bos * HV + i_hv) * BT + Akk = Akk + (bos * HV + i_hv) * BC + + p_q = tl.make_block_ptr(q, (T, K), (H*K, 1), (i_ti, 0), (BC, BK), (1, 0)) + p_k = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_ti, 0), (BC, BK), (1, 0)) + p_g = tl.make_block_ptr(g, (T, K), (HV*K, 1), (i_ti, 0), (BC, BK), (1, 0)) + + p_beta = tl.make_block_ptr(beta, (T,), (HV,), (i_ti,), (BC,), (0,)) + + b_q = tl.load(p_q, boundary_check=(0, 1)) + b_k = tl.load(p_k, boundary_check=(0, 1)) + b_g = tl.load(p_g, boundary_check=(0, 1)) + b_beta = tl.load(p_beta, boundary_check=(0,)).to(tl.float32) + + if USE_GATHER: + b_gn = gather(b_g, tl.full([1, BK], min(BC//2, T - i_ti - 1), dtype=tl.int16), axis=0) + else: + # caculate offset + p_gn = g + (i_ti + min(BC // 2, T - i_ti - 1)) * HV*K + tl.arange(0, BK) + b_gn = tl.load(p_gn, mask=tl.arange(0, BK) < K, other=0.0) + b_gn = b_gn[None, :] + + # current block, keep numerical stability by subtracting the left boundary + # less than 85 to avoid overflow in exp2 + b_gm = (b_g - b_gn).to(tl.float32) + + b_gq = tl.where(m_c[:, None], exp2(b_gm), 0.) + b_gk = tl.where(m_c[:, None], exp2(-b_gm), 0.) + + b_kgt = tl.trans(b_k * b_gk) + + b_Aqk = tl.dot(b_q * b_gq, b_kgt) * scale + b_Akk = tl.dot(b_k * b_gq, b_kgt) * b_beta[:, None] + + o_i = tl.arange(0, BC) + m_Aqk = o_i[:, None] >= o_i[None, :] + m_Akk = o_i[:, None] > o_i[None, :] + m_I = o_i[:, None] == o_i[None, :] + + b_Aqk = tl.where(m_Aqk, b_Aqk, 0.0) + b_Akk = tl.where(m_Akk, b_Akk, 0.0) + + p_Aqk = tl.make_block_ptr(Aqk, (T, BT), (HV*BT, 1), (i_ti, i_i * BC), (BC, BC), (1, 0)) + p_Akk = tl.make_block_ptr(Akk, (T, BC), (HV*BC, 1), (i_ti, 0), (BC, BC), (1, 0)) + tl.store(p_Aqk, b_Aqk.to(Aqk.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_Akk, b_Akk.to(Akk.dtype.element_ty), boundary_check=(0, 1)) + + tl.debug_barrier() + + ################################################################################ + # forward substitution + ################################################################################ + + b_Ai = -b_Akk + for i in range(2, min(BC, T - i_ti)): + b_a = -tl.load(Akk + (i_ti + i) * HV*BC + o_i) + b_a = tl.where(o_i < i, b_a, 0.) + b_a += tl.sum(b_a[:, None] * b_Ai, 0) + b_Ai = tl.where((o_i == i)[:, None], b_a, b_Ai) + b_Ai += m_I + tl.store(p_Akk, b_Ai.to(Akk.dtype.element_ty), boundary_check=(0, 1)) + + +def chunk_kda_fwd_intra( + q: torch.Tensor, + k: torch.Tensor, + gk: torch.Tensor | None = None, + beta: torch.Tensor | None = None, + scale: float | None = None, + cu_seqlens: torch.LongTensor | None = None, + chunk_size: int = 64, + chunk_indices: torch.LongTensor | None = None, + safe_gate: bool = False, +): + B, T, H, K, HV = *k.shape, gk.shape[2] + BT = chunk_size + BC = 16 + if chunk_indices is None and cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) + NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) + NC = triton.cdiv(BT, BC) + + Aqk = torch.empty(B, T, HV, BT, device=k.device, dtype=k.dtype) + # Akk must be zero-initialized - kernel only writes lower triangular + Akk = torch.zeros(B, T, HV, BT, device=k.device, dtype=k.dtype) + # Separate fp32 buffer for diagonal 16x16 blocks (for precision in solve_tril) + Akkd = torch.empty(B, T, HV, BC, device=k.device, dtype=torch.float32) + + # Compute diagonal blocks into Akkd in fp32. + if safe_gate: + grid = (NT, NC, B * HV) + BK = triton.next_power_of_2(K) + chunk_kda_fwd_kernel_intra_sub_chunk[grid]( + q=q, + k=k, + g=gk, + beta=beta, + Aqk=Aqk, + Akk=Akkd, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + T=T, + H=H, + HV=HV, + K=K, + BT=BT, + BC=BC, + BK=BK, + USE_GATHER=is_gather_supported, + ) + else: + Aqk, Akkd = chunk_kda_fwd_intra_token_parallel( + q=q, + k=k, + gk=gk, + beta=beta, + Aqk=Aqk, + Akk=Akkd, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_size=BT, + sub_chunk_size=BC, + ) + + # Step 2: Fused inter + solve_tril (works for both fixed-len and varlen) + solve_tril_dot_precision = ( + "tf32" + if current_platform.is_cuda() + and current_platform.has_device_capability(80) + else "ieee" + ) + grid = (NT, B * HV) + chunk_kda_fwd_kernel_inter_solve_fused[grid]( + q=q, + k=k, + g=gk, + beta=beta, + Aqk=Aqk, + Akkd=Akkd, + Akk=Akk, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + T=T, + H=H, + HV=HV, + K=K, + BT=BT, + BC=BC, + USE_SAFE_GATE=safe_gate, + SOLVE_TRIL_DOT_PRECISION=solve_tril_dot_precision, + ) + return Aqk, Akk diff --git a/aphrodite/models/kimi_k3/nvidia/ops/third_party/kda/chunk_intra_token_parallel.py b/aphrodite/models/kimi_k3/nvidia/ops/third_party/kda/chunk_intra_token_parallel.py new file mode 100644 index 0000000000..37c8bb34e9 --- /dev/null +++ b/aphrodite/models/kimi_k3/nvidia/ops/third_party/kda/chunk_intra_token_parallel.py @@ -0,0 +1,178 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Songlin Yang, Yu Zhang, Zhiyuan Li +# +# This file contains code copied from the flash-linear-attention project. +# The original source was licensed under the MIT license. +# Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li +# Forward-only adaptation of flash-linear-attention 0.5.0. +# ruff: noqa: E501 + +# Token-parallel implementation of KDA intra chunk kernel + +import torch + +from aphrodite.third_party.flash_linear_attention.ops.op import exp2 +from aphrodite.triton_utils import tl, triton + + +@triton.heuristics({ + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, +}) +@triton.autotune( + configs=[ + triton.Config({'BH': BH}, num_warps=num_warps) + for BH in [1, 2, 4, 8] + for num_warps in [1, 2, 4, 8] + ], + key=["K", "H", "HV"], +) +@triton.jit(do_not_specialize=['T', 'N']) +def chunk_kda_fwd_kernel_intra_token_parallel( + q, + k, + g, + beta, + Aqk, + Akk, + scale, + cu_seqlens, + N, + T, + H: tl.constexpr, + HV: tl.constexpr, + K: tl.constexpr, + BT: tl.constexpr, + BC: tl.constexpr, + BH: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_tg, i_hg = tl.program_id(0), tl.program_id(1) + + if IS_VARLEN: + i_n = 0 + left, right = 0, N + + # Unrolled binary search (max B=2^32) + # We can limit iterations based on expected max batch size if needed + # 20 iterations covers B=1M, usually enough + for _ in range(20): + if left < right: + mid = (left + right) // 2 + if i_tg < tl.load(cu_seqlens + mid + 1).to(tl.int32): + right = mid + else: + left = mid + 1 + i_n = left + + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + T = eos - bos + i_t = i_tg - bos + else: + bos = (i_tg // T) * T + i_t = i_tg % T + + if i_t >= T: + return + + i_c = i_t // BT + i_s = (i_t % BT) // BC + i_tc = i_c * BT + i_ts = i_tc + i_s * BC + + G: tl.constexpr = HV // H + + q += bos * H*K + k += bos * H*K + g += bos * HV*K + Aqk += bos * HV*BT + Akk += bos * HV*BC + beta += bos * HV + + BK: tl.constexpr = triton.next_power_of_2(K) + o_hv = i_hg * BH + tl.arange(0, BH) + o_h = o_hv // G + o_k = tl.arange(0, BK) + m_hv = o_hv < HV + m_k = o_k < K + m_hk = m_hv[:, None] & m_k[None, :] + + # q/k: [B, T, H, K], manual load via mapped qk head index + p_qk = o_h[:, None] * K + o_k[None, :] + b_q = tl.load(q + i_t * H * K + p_qk, mask=m_hk, other=0).to(tl.float32) + b_k = tl.load(k + i_t * H * K + p_qk, mask=m_hk, other=0).to(tl.float32) + + # g: [B, T, HV, K], beta: [B, T, HV] + p_g = tl.make_block_ptr(g + i_t * HV * K, (HV, K), (K, 1), (i_hg * BH, 0), (BH, BK), (1, 0)) + p_beta = tl.make_block_ptr(beta + i_t * HV, (HV,), (1,), (i_hg * BH,), (BH,), (0,)) + b_g = tl.load(p_g, boundary_check=(0, 1)).to(tl.float32) + b_beta = tl.load(p_beta, boundary_check=(0,)).to(tl.float32) + b_k *= b_beta[:, None] + + for j in range(i_ts, min(i_t + 1, min(T, i_ts + BC))): + b_kj = tl.load(k + j * H * K + p_qk, mask=m_hk, other=0).to(tl.float32) + p_gj = tl.make_block_ptr(g + j * HV * K, (HV, K), (K, 1), (i_hg * BH, 0), (BH, BK), (1, 0)) + b_gj = tl.load(p_gj, boundary_check=(0, 1)).to(tl.float32) + + b_kgj = tl.where(m_k[None, :], b_kj * exp2(b_g - b_gj), 0.0) + b_Aqk = tl.sum(b_q * b_kgj, axis=1) * scale + b_Akk = tl.sum(b_k * b_kgj, axis=1) * tl.where(j < i_t, 1.0, 0.0) + + tl.store(Aqk + i_t * HV * BT + o_hv * BT + j % BT, b_Aqk.to(Aqk.dtype.element_ty), mask=m_hv) + tl.store(Akk + i_t * HV * BC + o_hv * BC + j - i_ts, b_Akk.to(Akk.dtype.element_ty), mask=m_hv) + + +def chunk_kda_fwd_intra_token_parallel( + q: torch.Tensor, + k: torch.Tensor, + gk: torch.Tensor, + beta: torch.Tensor, + Aqk: torch.Tensor, + Akk: torch.Tensor, + scale: float, + cu_seqlens: torch.LongTensor | None = None, + chunk_size: int = 64, + sub_chunk_size: int = 16, +) -> None: + """ + Token-parallel implementation: each token gets its own thread block. + Supports both fixed-length and variable-length sequences. + Reduces wasted computation on padding. + + Writes directly to Aqk and Akk tensors (in-place). + + Args: + q: [B, T, H, K] + k: [B, T, H, K] + gk: [B, T, HV, K] cumsum of gates (HV >= H for GVA) + beta: [B, T, HV] + Aqk: [B, T, HV, BT] output tensor to write to + Akk: [B, T, HV, BC] output tensor for diagonal blocks (fp32) + scale: attention scale + chunk_size: BT (default 64) + sub_chunk_size: BC (default 16) + """ + B, T, H, K, HV = *q.shape, gk.shape[2] + N = len(cu_seqlens) - 1 if cu_seqlens is not None else B + BT = chunk_size + BC = sub_chunk_size + + def grid(meta): return (B * T, triton.cdiv(HV, meta['BH'])) + chunk_kda_fwd_kernel_intra_token_parallel[grid]( + q=q, + k=k, + g=gk, + beta=beta, + Aqk=Aqk, + Akk=Akk, + scale=scale, + cu_seqlens=cu_seqlens, + N=N, + T=T, + H=H, + HV=HV, + K=K, + BT=BT, + BC=BC, + ) + return Aqk, Akk diff --git a/aphrodite/models/kimi_k3/nvidia/ops/third_party/kda/fused_recurrent.py b/aphrodite/models/kimi_k3/nvidia/ops/third_party/kda/fused_recurrent.py new file mode 100644 index 0000000000..4cc220066d --- /dev/null +++ b/aphrodite/models/kimi_k3/nvidia/ops/third_party/kda/fused_recurrent.py @@ -0,0 +1,671 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Songlin Yang, Yu Zhang, Zhiyuan Li +# +# This file contains code adapted from the flash-linear-attention project. +# The original source was licensed under the MIT license. +# Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li +# ruff: noqa: E501 + +import torch + +from aphrodite.platforms import current_platform +from aphrodite.third_party.flash_linear_attention.ops.op import exp, log +from aphrodite.triton_utils import tl, triton +from aphrodite.utils.math_utils import cdiv, next_power_of_2 + + +@triton.heuristics( + { + "HAS_DT_BIAS": lambda args: args["dt_bias"] is not None, + "USE_LOWER_BOUND": lambda args: args["lower_bound"] is not None, + } +) +@triton.jit +def _kda_gate_beta_fwd_kernel( + raw_g, + raw_beta, + A_log, + dt_bias, + gate, + beta_out, + lower_bound, + softplus_beta: tl.constexpr, + softplus_threshold: tl.constexpr, + T, + stride_g_token: tl.constexpr, + stride_beta_token: tl.constexpr, + H: tl.constexpr, + D: tl.constexpr, + BT: tl.constexpr, + BD: tl.constexpr, + HAS_DT_BIAS: tl.constexpr, + USE_LOWER_BOUND: tl.constexpr, + launch_pdl: tl.constexpr, +): + if launch_pdl: + tl.extra.cuda.gdc_wait() + tl.extra.cuda.gdc_launch_dependents() + + i_t, i_h = tl.program_id(0), tl.program_id(1) + o_t = i_t * BT + tl.arange(0, BT) + o_d = tl.arange(0, BD) + m_t = o_t < T + m_d = o_d < D + + p_g = raw_g + o_t[:, None] * stride_g_token + i_h * D + o_d[None, :] + b_g = tl.load(p_g, mask=m_t[:, None] & m_d[None, :], other=0.0).to(tl.float32) + if HAS_DT_BIAS: + b_bias = tl.load( + dt_bias + i_h * D + o_d, + mask=m_d, + other=0.0, + ).to(tl.float32) + b_g += b_bias[None, :] + + b_a = exp(tl.load(A_log + i_h).to(tl.float32)) + if USE_LOWER_BOUND: + b_gate = lower_bound * tl.sigmoid(b_a * b_g) + else: + b_scaled = b_g * softplus_beta + b_softplus = tl.where( + b_scaled > softplus_threshold, + b_g, + log(1.0 + tl.exp(b_scaled)) / softplus_beta, + ) + b_gate = -b_a * b_softplus + + p_gate = gate + (o_t[:, None] * H + i_h) * D + o_d[None, :] + tl.store( + p_gate, + b_gate, + mask=m_t[:, None] & m_d[None, :], + ) + + b_beta = tl.load( + raw_beta + o_t * stride_beta_token + i_h, + mask=m_t, + other=0.0, + ).to(tl.float32) + tl.store(beta_out + o_t * H + i_h, tl.sigmoid(b_beta), mask=m_t) + + +def _fused_kda_gate_beta( + raw_g: torch.Tensor, + raw_beta: torch.Tensor, + A_log: torch.Tensor, + dt_bias: torch.Tensor | None, + lower_bound: float | None, +) -> tuple[torch.Tensor, torch.Tensor]: + B, T, H, D = raw_g.shape + assert B == 1 + assert raw_beta.shape == (B, T, H) + assert raw_g.stride()[2:] == (D, 1) + assert raw_beta.stride(2) == 1 + gate = torch.empty((B, T, H, D), dtype=torch.float32, device=raw_g.device) + beta = torch.empty((B, T, H), dtype=torch.float32, device=raw_beta.device) + + BT = 16 + _kda_gate_beta_fwd_kernel[(cdiv(T, BT), H)]( + raw_g=raw_g, + raw_beta=raw_beta, + A_log=A_log, + dt_bias=dt_bias, + gate=gate, + beta_out=beta, + lower_bound=lower_bound, + softplus_beta=1.0, + softplus_threshold=20.0, + T=T, + stride_g_token=raw_g.stride(1), + stride_beta_token=raw_beta.stride(1), + H=H, + D=D, + BT=BT, + BD=next_power_of_2(D), + num_warps=4, + launch_pdl=current_platform.is_arch_support_pdl(), + ) + return gate, beta + + +@triton.heuristics( + { + "IS_SPEC_DECODING": lambda args: args["num_accepted_tokens"] is not None, + "HAS_DT_BIAS": lambda args: args["dt_bias"] is not None, + "USE_LOWER_BOUND": lambda args: args["lower_bound"] is not None, + } +) +@triton.jit(do_not_specialize=["N", "T", "stride_beta_token"]) +def fused_recurrent_kda_fwd_kernel( + q, + k, + v, + g, + beta, + A_log, + dt_bias, + out, + state, + cu_seqlens, + state_indices, + num_accepted_tokens, + lower_bound, + scale: tl.constexpr, + N: tl.int64, + T: tl.int64, + H: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + stride_qkv_token: tl.constexpr, + stride_g_token: tl.constexpr, + stride_beta_token, + stride_out_token: tl.constexpr, + stride_state_token: tl.constexpr, + stride_indices_seq: tl.constexpr, + IS_SPEC_DECODING: tl.constexpr, + USE_QK_L2NORM_IN_KERNEL: tl.constexpr, + USE_GATE_IN_KERNEL: tl.constexpr, + APPLY_BETA_SIGMOID: tl.constexpr, + HAS_DT_BIAS: tl.constexpr, + USE_LOWER_BOUND: tl.constexpr, + num_stages: tl.constexpr, + launch_pdl: tl.constexpr, +): + if launch_pdl: + tl.extra.cuda.gdc_wait() + + pid = tl.program_id(0) + i_v = pid % tl.cdiv(V, BV) + i_nh = pid // tl.cdiv(V, BV) + i_n, i_h = i_nh // H, i_nh % H + bos = tl.load(cu_seqlens + i_n).to(tl.int64) + eos = tl.load(cu_seqlens + i_n + 1).to(tl.int64) + sequence_length = eos - bos + if sequence_length == 0: + return + + o_k = tl.arange(0, BK) + o_v = i_v * BV + tl.arange(0, BV) + m_k = o_k < K + m_v = o_v < V + m_state = m_v[:, None] & m_k[None, :] + + if IS_SPEC_DECODING: + initial_token = tl.load(num_accepted_tokens + i_n).to(tl.int64) - 1 + else: + initial_token = 0 + state_index = tl.load(state_indices + i_n * stride_indices_seq + initial_token).to( + tl.int64 + ) + p_out = out + bos * stride_out_token + i_h * V + o_v + if state_index <= 0: + tl.store(p_out, tl.zeros([BV], dtype=tl.float32), mask=m_v) + return + + p_state = ( + state + + state_index * stride_state_token + + i_h * V * K + + o_v[:, None] * K + + o_k[None, :] + ) + b_state = tl.load(p_state, mask=m_state, other=0.0).to(tl.float32) + + p_q = q + bos * stride_qkv_token + i_h * K + o_k + p_k = k + bos * stride_qkv_token + i_h * K + o_k + p_v = v + bos * stride_qkv_token + i_h * V + o_v + p_g = g + bos * stride_g_token + i_h * K + o_k + p_beta = beta + bos * stride_beta_token + i_h + for i_t in tl.range(0, sequence_length, num_stages=num_stages): + b_q = tl.load(p_q, mask=m_k, other=0.0, eviction_policy="evict_last").to( + tl.float32 + ) + b_k = tl.load(p_k, mask=m_k, other=0.0, eviction_policy="evict_last").to( + tl.float32 + ) + b_v = tl.load(p_v, mask=m_v, other=0.0, eviction_policy="evict_first").to( + tl.float32 + ) + if USE_QK_L2NORM_IN_KERNEL: + b_q = b_q / tl.sqrt(tl.sum(b_q * b_q) + 1e-6) + b_k = b_k / tl.sqrt(tl.sum(b_k * b_k) + 1e-6) + b_q *= scale + + b_gate = tl.load( + p_g, + mask=m_k, + other=0.0, + eviction_policy="evict_last", + ).to(tl.float32) + if USE_GATE_IN_KERNEL: + if HAS_DT_BIAS: + b_bias = tl.load( + dt_bias + i_h * K + o_k, + mask=m_k, + other=0.0, + ).to(tl.float32) + b_gate += b_bias + b_a = exp(tl.load(A_log + i_h).to(tl.float32)) + if USE_LOWER_BOUND: + b_gate = lower_bound * tl.sigmoid(b_a * b_gate) + else: + b_softplus = tl.where( + b_gate > 20.0, + b_gate, + log(1.0 + tl.exp(b_gate)), + ) + b_gate = -b_a * b_softplus + + b_state *= exp(b_gate[None, :]) + b_v -= tl.sum(b_state * b_k[None, :], axis=1) + b_beta = tl.load(p_beta, eviction_policy="evict_last").to(tl.float32) + if APPLY_BETA_SIGMOID: + b_beta = tl.sigmoid(b_beta) + b_v *= b_beta + b_state += b_v[:, None] * b_k[None, :] + b_out = tl.sum(b_state * b_q[None, :], axis=1) + tl.store( + p_out, + b_out.to(p_out.dtype.element_ty), + mask=m_v, + eviction_policy="evict_first", + ) + + final_state_index = tl.load(state_indices + i_n * stride_indices_seq + i_t).to( + tl.int64 + ) + if final_state_index > 0: + p_final_state = ( + state + + final_state_index * stride_state_token + + i_h * V * K + + o_v[:, None] * K + + o_k[None, :] + ) + tl.store( + p_final_state, + b_state.to(p_final_state.dtype.element_ty), + mask=m_state, + ) + + p_q += stride_qkv_token + p_k += stride_qkv_token + p_v += stride_qkv_token + p_g += stride_g_token + p_beta += stride_beta_token + p_out += stride_out_token + + if launch_pdl: + tl.extra.cuda.gdc_launch_dependents() + + +# Consumed by kimi_k3_triton_warmup.py during kernel_warmup(). +def get_fused_recurrent_kda_fwd_warmup_profiles( + num_heads: int, +) -> tuple[int, ...]: + """Return representative sequence counts for gated launch variants.""" + # The region above 192 head-sequences reuses the second launch variant. + return ( + 1, + 48 // num_heads + 1, + 96 // num_heads + 1, + ) + + +def fused_recurrent_kda_fwd( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + scale: float | None = None, + initial_state: torch.Tensor | None = None, + inplace_final_state: bool = True, + cu_seqlens: torch.Tensor | None = None, + ssm_state_indices: torch.Tensor | None = None, + num_accepted_tokens: torch.Tensor | None = None, + use_qk_l2norm_in_kernel: bool = True, + A_log: torch.Tensor | None = None, + dt_bias: torch.Tensor | None = None, + lower_bound: float | None = None, + use_gate_in_kernel: bool = False, + use_beta_sigmoid_in_kernel: bool = False, + out: torch.Tensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor]: + """Launch recurrent KDA with dense inner dimensions and row strides.""" + B, T, H, K = q.shape + V = v.shape[-1] + assert B == 1 and k.shape == q.shape + assert v.shape == (B, T, H, V) and g.shape == (B, T, H, K) + assert beta.shape == (B, T, H) + assert initial_state is not None + assert cu_seqlens is not None + assert ssm_state_indices is not None + assert inplace_final_state + if out is None: + out = torch.empty_like(v) + assert out.shape == v.shape + assert initial_state.shape[1:] == (H, V, K) + assert ssm_state_indices.ndim in (1, 2) + + assert q.stride()[2:] == k.stride()[2:] == (K, 1) + assert v.stride()[2:] == out.stride()[2:] == (V, 1) + assert g.stride()[2:] == (K, 1) + assert beta.stride(2) == 1 + assert q.stride(1) == k.stride(1) == v.stride(1) + assert initial_state.stride()[1:] == (V * K, K, 1) + N = cu_seqlens.numel() - 1 + if ssm_state_indices.ndim == 1: + assert T == N + assert num_accepted_tokens is None + else: + assert ssm_state_indices.stride(1) == 1 + assert cu_seqlens.is_contiguous() + if use_gate_in_kernel: + assert A_log is not None and A_log.is_contiguous() + assert dt_bias is None or dt_bias.is_contiguous() + + if scale is None: + scale = K**-0.5 + + if use_gate_in_kernel: + # Tuned on GB300 for Kimi-K3 shapes. Keep the warmup profiles above in + # sync with these boundaries. + head_sequences = H * N + if head_sequences <= 48: + BV, num_stages = 4, 4 + elif head_sequences <= 96: + BV, num_stages = 8, 3 + elif head_sequences <= 192: + BV, num_stages = 16, 3 + else: + BV, num_stages = 8, 3 + num_warps = 1 + else: + BV, num_warps, num_stages = 8, 1, 2 + grid = (cdiv(V, BV) * N * H,) + fused_recurrent_kda_fwd_kernel[grid]( + q=q, + k=k, + v=v, + g=g, + beta=beta, + A_log=A_log, + dt_bias=dt_bias, + out=out, + state=initial_state, + cu_seqlens=cu_seqlens, + state_indices=ssm_state_indices, + num_accepted_tokens=num_accepted_tokens, + lower_bound=lower_bound, + scale=scale, + N=N, + T=T, + H=H, + K=K, + V=V, + BK=next_power_of_2(K), + BV=BV, + stride_qkv_token=q.stride(1), + stride_g_token=g.stride(1), + stride_beta_token=beta.stride(1), + stride_out_token=out.stride(1), + stride_state_token=initial_state.stride(0), + stride_indices_seq=ssm_state_indices.stride(0), + IS_SPEC_DECODING=num_accepted_tokens is not None, + USE_QK_L2NORM_IN_KERNEL=use_qk_l2norm_in_kernel, + USE_GATE_IN_KERNEL=use_gate_in_kernel, + APPLY_BETA_SIGMOID=use_beta_sigmoid_in_kernel, + num_warps=num_warps, + num_stages=num_stages, + launch_pdl=current_platform.is_arch_support_pdl(), + ) + return out, initial_state + + +def fused_recurrent_kda( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + raw_g: torch.Tensor, + raw_beta: torch.Tensor, + A_log: torch.Tensor, + dt_bias: torch.Tensor | None, + lower_bound: float | None, + initial_state: torch.Tensor, + cu_seqlens: torch.Tensor, + ssm_state_indices: torch.Tensor, + num_accepted_tokens: torch.Tensor | None = None, + out: torch.Tensor | None = None, + fuse_gate: bool | None = None, +) -> tuple[torch.Tensor, torch.Tensor]: + """Run recurrent KDA from raw gate and beta inputs. + + This vLLM wrapper applies the gate activation and beta sigmoid, selecting + whether to materialize them before launching the recurrent kernel. + """ + if fuse_gate is None: + fuse_gate = True + + if fuse_gate: + gate = raw_g + beta = raw_beta + else: + gate, beta = _fused_kda_gate_beta( + raw_g, + raw_beta, + A_log, + dt_bias, + lower_bound, + ) + return fused_recurrent_kda_fwd( + q=q, + k=k, + v=v, + g=gate, + beta=beta, + scale=q.shape[-1] ** -0.5, + initial_state=initial_state, + inplace_final_state=True, + cu_seqlens=cu_seqlens, + ssm_state_indices=ssm_state_indices, + num_accepted_tokens=num_accepted_tokens, + use_qk_l2norm_in_kernel=True, + A_log=A_log if fuse_gate else None, + dt_bias=dt_bias if fuse_gate else None, + lower_bound=lower_bound if fuse_gate else None, + use_gate_in_kernel=fuse_gate, + use_beta_sigmoid_in_kernel=fuse_gate, + out=out, + ) + + +@triton.jit( + do_not_specialize=["stride_beta_token", "stride_state_indices"] +) +def fused_recurrent_kda_packed_decode_kernel( + mixed_qkv, + raw_g, + raw_beta, + A_log, + dt_bias, + out, + state, + state_indices, + lower_bound, + scale: tl.constexpr, + stride_mixed_token: tl.constexpr, + stride_g_token: tl.constexpr, + stride_beta_token, + stride_state_token: tl.constexpr, + stride_state_indices, + H: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + SOFTPLUS_THRESHOLD: tl.constexpr, + USE_LOWER_BOUND: tl.constexpr, + launch_pdl: tl.constexpr, +): + i_v, i_nh = tl.program_id(0), tl.program_id(1) + i_n, i_h = i_nh // H, i_nh % H + + o_k = tl.arange(0, BK) + o_v = i_v * BV + tl.arange(0, BV) + mask_k = o_k < K + mask_v = o_v < V + mask_state = mask_v[:, None] & mask_k[None, :] + + if launch_pdl: + tl.extra.cuda.gdc_wait() + tl.extra.cuda.gdc_launch_dependents() + + state_idx = tl.load(state_indices + i_n * stride_state_indices).to(tl.int64) + p_out = out + (i_n * H + i_h) * V + o_v + if state_idx <= 0: + tl.store(p_out, tl.zeros([BV], dtype=tl.float32), mask=mask_v) + return + + p_state = state + state_idx * stride_state_token + p_state += i_h * V * K + o_v[:, None] * K + o_k[None, :] + b_state = tl.load(p_state, mask=mask_state, other=0).to(tl.float32) + + # Q, K, and V occupy consecutive channel ranges, while the token stride + # may also include the output-gate projection that follows packed QKV. + p_mixed = mixed_qkv + i_n * stride_mixed_token + b_q = tl.load(p_mixed + i_h * K + o_k, mask=mask_k, other=0).to(tl.float32) + b_k = tl.load( + p_mixed + H * K + i_h * K + o_k, + mask=mask_k, + other=0, + ).to(tl.float32) + b_v = tl.load( + p_mixed + 2 * H * K + i_h * V + o_v, + mask=mask_v, + other=0, + ).to(tl.float32) + + b_q /= tl.sqrt(tl.sum(b_q * b_q) + 1e-6) + b_k /= tl.sqrt(tl.sum(b_k * b_k) + 1e-6) + b_q *= scale + + p_g = raw_g + i_n * stride_g_token + i_h * K + o_k + b_g = tl.load(p_g, mask=mask_k, other=0).to(tl.float32) + b_bias = tl.load(dt_bias + i_h * K + o_k, mask=mask_k, other=0).to(tl.float32) + b_a = exp(tl.load(A_log + i_h).to(tl.float32)) + b_g += b_bias + if USE_LOWER_BOUND: + b_gate = lower_bound * tl.sigmoid(b_a * b_g) + else: + b_softplus = tl.where( + b_g > SOFTPLUS_THRESHOLD, + b_g, + log(1.0 + tl.exp(b_g)), + ) + b_gate = -b_a * b_softplus + + b_state *= exp(b_gate[None, :]) + b_v -= tl.sum(b_state * b_k[None, :], axis=1) + b_beta = tl.sigmoid( + tl.load(raw_beta + i_n * stride_beta_token + i_h).to(tl.float32) + ) + b_v *= b_beta + b_state += b_v[:, None] * b_k[None, :] + b_out = tl.sum(b_state * b_q[None, :], axis=1) + + tl.store(p_out, b_out.to(p_out.dtype.element_ty), mask=mask_v) + tl.store(p_state, b_state.to(p_state.dtype.element_ty), mask=mask_state) + + +def fused_recurrent_kda_packed_decode( + mixed_qkv: torch.Tensor, + raw_g: torch.Tensor, + raw_beta: torch.Tensor, + A_log: torch.Tensor, + dt_bias: torch.Tensor, + lower_bound: float | None, + initial_state: torch.Tensor, + state_indices: torch.Tensor, + scale: float | None = None, +) -> tuple[torch.Tensor, torch.Tensor]: + """Run one-token KDA decode directly from packed post-conv QKV.""" + if mixed_qkv.ndim != 2 or mixed_qkv.stride(-1) != 1: + raise ValueError("`mixed_qkv` must be 2D and contiguous in its last dim.") + if raw_g.ndim != 4 or raw_g.shape[0] != 1: + raise ValueError("`raw_g` must have shape [1, B, H, K].") + if raw_beta.ndim != 3 or raw_beta.shape[0] != 1: + raise ValueError("`raw_beta` must have shape [1, B, H].") + if initial_state.ndim != 4: + raise ValueError("`initial_state` must have shape [cache, H, V, K].") + _, H, V, K = initial_state.shape + if raw_g.stride()[2:] != (K, 1): + raise ValueError("`raw_g` must be contiguous within each token.") + if raw_beta.stride(2) != 1: + raise ValueError("`raw_beta` heads must be contiguous.") + if initial_state.stride()[1:] != (V * K, K, 1): + raise ValueError("`initial_state` must be contiguous within each cache slot.") + if state_indices.ndim != 1: + raise ValueError("`state_indices` must be one-dimensional.") + if A_log.ndim != 1 or not A_log.is_contiguous(): + raise ValueError("`A_log` must be contiguous and one-dimensional.") + if not dt_bias.is_contiguous(): + raise ValueError("`dt_bias` must be contiguous.") + + device = mixed_qkv.device + if any( + x.device != device + for x in (raw_g, raw_beta, A_log, dt_bias, initial_state, state_indices) + ): + raise ValueError("All packed KDA inputs must be on the same device.") + + B = mixed_qkv.shape[0] + if raw_g.shape != (1, B, H, K): + raise ValueError(f"Unexpected raw gate shape {tuple(raw_g.shape)}.") + if raw_beta.shape != (1, B, H): + raise ValueError(f"Unexpected raw beta shape {tuple(raw_beta.shape)}.") + if mixed_qkv.shape[1] != 2 * H * K + H * V: + raise ValueError(f"Unexpected packed QKV shape {tuple(mixed_qkv.shape)}.") + if A_log.numel() != H or dt_bias.numel() != H * K: + raise ValueError("`A_log` or `dt_bias` has an incompatible shape.") + if state_indices.shape[0] != B: + raise ValueError("`state_indices` must contain one entry per token.") + + BK = next_power_of_2(K) + BV = min(next_power_of_2(V), 32) + if scale is None: + scale = K**-0.5 + + out = torch.empty((1, B, H, V), dtype=mixed_qkv.dtype, device=device) + grid = (cdiv(V, BV), B * H) + fused_recurrent_kda_packed_decode_kernel[grid]( + mixed_qkv=mixed_qkv, + raw_g=raw_g, + raw_beta=raw_beta, + A_log=A_log, + dt_bias=dt_bias, + out=out, + state=initial_state, + state_indices=state_indices, + lower_bound=lower_bound or 0.0, + scale=scale, + stride_mixed_token=mixed_qkv.stride(0), + stride_g_token=raw_g.stride(1), + stride_beta_token=raw_beta.stride(1), + stride_state_token=initial_state.stride(0), + stride_state_indices=state_indices.stride(0), + H=H, + K=K, + V=V, + BK=BK, + BV=BV, + SOFTPLUS_THRESHOLD=20.0, + USE_LOWER_BOUND=lower_bound is not None, + num_warps=4, + num_stages=2, + launch_pdl=current_platform.is_arch_support_pdl(), + ) + return out, initial_state diff --git a/aphrodite/models/kimi_k3/nvidia/ops/vision_fa4_warmup.py b/aphrodite/models/kimi_k3/nvidia/ops/vision_fa4_warmup.py new file mode 100644 index 0000000000..00bebf9689 --- /dev/null +++ b/aphrodite/models/kimi_k3/nvidia/ops/vision_fa4_warmup.py @@ -0,0 +1,191 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Startup compilation of Kimi-K3 vision FA4 kernels.""" + +from __future__ import annotations + +import math +from collections.abc import Iterator +from dataclasses import dataclass +from functools import partial + +import torch + +from aphrodite.model_executor.warmup.cutedsl_warmup import ( + CuTeDSLCompileUnit, + register_cutedsl_warmup_provider, +) +from aphrodite.platforms import current_platform + +_FA4_TILE_SIZE = 128 +_FA4_MAX_SPLITS = 128 + + +@dataclass(frozen=True) +class KimiK3VisionFA4WarmupConfig: + num_heads: int + head_dim: int + dtype: torch.dtype + max_batch_size: int + max_seqlen: int + + +@dataclass(frozen=True) +class _FA4WarmupProbe: + batch_size: int + max_seqlen: int + dispatch_key: tuple[int, bool, int | None] + + +def _combine_log_max_splits(head_dim: int, num_splits: int) -> int: + k_block_size = 64 if head_dim <= 64 else 128 + tile_m = 8 if k_block_size == 128 else 16 + min_log_max_splits = 5 if tile_m == 8 else 4 + return max(math.ceil(math.log2(num_splits)), min_log_max_splits) + + +def _get_dispatch( + config: KimiK3VisionFA4WarmupConfig, + *, + batch_size: int, + max_seqlen: int, + num_sms: int, +) -> tuple[int, bool, int | None]: + # These are the shape-dependent fields in FA4's forward and combine + # compile keys for noncausal varlen MHA. Compile units still pass + # num_splits=0 so the production selector makes the final decision. + q_stage = 2 if max_seqlen > _FA4_TILE_SIZE else 1 + effective_q_tile = q_stage * _FA4_TILE_SIZE + num_m_blocks = math.ceil(max_seqlen / effective_q_tile) + num_n_blocks = math.ceil(max_seqlen / _FA4_TILE_SIZE) + + if num_n_blocks <= 4: + num_splits = 1 + else: + total_m_blocks = batch_size * config.num_heads * num_m_blocks + num_splits = min( + num_sms // total_m_blocks, + _FA4_MAX_SPLITS, + num_n_blocks, + ) + + is_split_kv = num_splits > 1 + combine_key = _combine_log_max_splits(config.head_dim, num_splits) if is_split_kv else None + return q_stage, is_split_kv, combine_key + + +def _get_warmup_probes( + config: KimiK3VisionFA4WarmupConfig, + *, + num_sms: int, +) -> tuple[_FA4WarmupProbe, ...]: + if config.max_batch_size <= 0 or config.max_seqlen <= 0: + return () + + probes: dict[tuple[int, bool, int | None], _FA4WarmupProbe] = {} + + def add_probe(batch_size: int, max_seqlen: int) -> None: + dispatch_key = _get_dispatch( + config, + batch_size=batch_size, + max_seqlen=max_seqlen, + num_sms=num_sms, + ) + if dispatch_key not in probes: + probes[dispatch_key] = _FA4WarmupProbe( + batch_size=batch_size, + max_seqlen=max_seqlen, + dispatch_key=dispatch_key, + ) + + add_probe(batch_size=1, max_seqlen=1) + if config.max_seqlen > _FA4_TILE_SIZE: + add_probe(batch_size=1, max_seqlen=_FA4_TILE_SIZE + 1) + + max_n_blocks = math.ceil(config.max_seqlen / _FA4_TILE_SIZE) + for num_n_blocks in range(5, max_n_blocks + 1): + max_seqlen = min( + (num_n_blocks - 1) * _FA4_TILE_SIZE + 1, + config.max_seqlen, + ) + num_m_blocks = math.ceil(max_seqlen / (2 * _FA4_TILE_SIZE)) + max_split_batch_size = min( + config.max_batch_size, + num_sms // (2 * config.num_heads * num_m_blocks), + ) + if max_split_batch_size == 0: + break + for batch_size in range(1, max_split_batch_size + 1): + add_probe(batch_size, max_seqlen) + + return tuple(probes.values()) + + +def _iter_compile_units( + config: KimiK3VisionFA4WarmupConfig, +) -> Iterator[CuTeDSLCompileUnit]: + device_id = current_platform.current_device() + num_sms = current_platform.num_compute_units(device_id) + for probe in _get_warmup_probes(config, num_sms=num_sms): + yield CuTeDSLCompileUnit( + name="kimi_k3_vision_fa4", + key=("kimi_k3_vision_fa4", config, probe.dispatch_key), + compile=partial(_compile, config, probe), + ) + + +def _compile( + config: KimiK3VisionFA4WarmupConfig, + probe: _FA4WarmupProbe, +) -> None: + from aphrodite.vllm_flash_attn import flash_attn_varlen_func + + device = current_platform.current_device() + total_tokens = probe.batch_size * probe.max_seqlen + qkv = torch.empty( + 3, + total_tokens, + config.num_heads, + config.head_dim, + device=device, + dtype=config.dtype, + ) + cu_seqlens = torch.arange( + 0, + total_tokens + 1, + probe.max_seqlen, + device=qkv.device, + dtype=torch.int32, + ) + flash_attn_varlen_func( + qkv[0], + qkv[1], + qkv[2], + cu_seqlens_q=cu_seqlens, + cu_seqlens_k=cu_seqlens, + max_seqlen_q=probe.max_seqlen, + max_seqlen_k=probe.max_seqlen, + dropout_p=0.0, + causal=False, + softmax_scale=config.head_dim**-0.5, + fa_version=4, + num_splits=0, + ) + + +class _WarmupProvider: + def __init__(self) -> None: + self.configs: set[KimiK3VisionFA4WarmupConfig] = set() + + def get_cutedsl_warmup_compile_units(self) -> tuple[CuTeDSLCompileUnit, ...]: + return tuple(unit for config in self.configs for unit in _iter_compile_units(config)) + + +_PROVIDER = _WarmupProvider() + + +def register_kimi_k3_vision_fa4_warmup( + config: KimiK3VisionFA4WarmupConfig, +) -> None: + _PROVIDER.configs.add(config) + register_cutedsl_warmup_provider(_PROVIDER) diff --git a/aphrodite/models/minimax_m3/common/mm_preprocess.py b/aphrodite/models/minimax_m3/common/mm_preprocess.py index 553d98aa80..8d42c1ec21 100644 --- a/aphrodite/models/minimax_m3/common/mm_preprocess.py +++ b/aphrodite/models/minimax_m3/common/mm_preprocess.py @@ -3,8 +3,9 @@ import math from collections.abc import Mapping, Sequence -from typing import cast +from typing import Any, Literal, cast +import numpy.typing as npt import torch from transformers import BatchFeature from transformers.video_utils import VideoMetadata @@ -439,10 +440,39 @@ def get_video_replacement(item_idx: int): ] -# TODO(Isotr0py): Tie with MinimaxVideoProcessor -# after https://github.com/vllm-project/vllm/pull/44126 -@VIDEO_LOADER_REGISTRY.register("minimax_m3_vl") +@VIDEO_LOADER_REGISTRY.register( + name="minimax_m3_vl", + video_processor="MiniMaxM3VLVideoProcessor", +) class MiniMaxM3VideoBackend(VideoBackend): + @classmethod + def load_bytes( + cls, + data: bytes, + num_frames: int = -1, + fps: int = 1, + max_duration: int = 300, + frame_recovery: bool = False, + *, + backend: Literal[ + "opencv", + "pyav", + "torchcodec", + "pynvvideocodec", + "deepstream", + ] = "opencv", + **kwargs, + ) -> tuple[npt.NDArray, dict[str, Any]]: + return super().load_bytes( + data, + num_frames=num_frames, + fps=fps, + max_duration=max_duration, + frame_recovery=frame_recovery, + backend=backend, + **kwargs, + ) + @classmethod def compute_frames_index_to_sample( cls, @@ -474,7 +504,9 @@ def compute_frames_index_to_sample( indices.append(target_frame) prev_kept_ts = target_frame / video_fps - last_frame_idx = total_frames - 1 + # Because HF sample_frames includes the last frame, + # we will use HF as the standard. + last_frame_idx = total_frames last_ts = last_frame_idx / video_fps if indices and indices[-1] != last_frame_idx and last_ts - prev_kept_ts > eps: indices.append(last_frame_idx) diff --git a/aphrodite/parser/kimi_k3.py b/aphrodite/parser/kimi_k3.py new file mode 100644 index 0000000000..79390de80a --- /dev/null +++ b/aphrodite/parser/kimi_k3.py @@ -0,0 +1,136 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the Aphrodite project + +from __future__ import annotations + +from collections.abc import Sequence +from typing import TYPE_CHECKING + +from aphrodite.entrypoints.openai.engine.protocol import ( + DeltaMessage, + FunctionCall, +) +from aphrodite.parser.abstract_parser import DelegatingParser +from aphrodite.reasoning.kimi_k3_reasoning_parser import KimiK3ReasoningParser + +if TYPE_CHECKING: + from aphrodite.entrypoints.openai.chat_completion.protocol import ( + ChatCompletionRequest, + ) + from aphrodite.entrypoints.openai.responses.protocol import ResponsesRequest + + +class KimiK3Parser(DelegatingParser): + """Compose the Kimi K3 reasoning and tool parsers for XTML output.""" + + # TODO: Switch Kimi K3 to the parser engine once its XTML reasoning/tool + # path is covered there. + def _extract_tool_calls( + self, + content: str | None, + request: ChatCompletionRequest | ResponsesRequest, + enable_auto_tools: bool = False, + ) -> tuple[list[FunctionCall] | None, str | None]: + if self._tool_parser is None or not enable_auto_tools: + return super()._extract_tool_calls(content, request, enable_auto_tools) + + tool_call_info = self.extract_tool_calls(content or "", request=request) + if request.tool_choice == "none": + return [], tool_call_info.content + if not tool_call_info.tools_called: + return None, tool_call_info.content + + tool_calls = [ + FunctionCall( + id=tool_call.id, + name=tool_call.function.name, + arguments=tool_call.function.arguments, + ) + for tool_call in tool_call_info.tool_calls + ] + parsed_content = tool_call_info.content + if parsed_content and parsed_content.strip() == "": + parsed_content = None + return tool_calls, parsed_content + + def _extract_tool_calls_streaming( + self, + previous_text: str, + current_text: str, + delta_text: str, + previous_token_ids: Sequence[int], + current_token_ids: Sequence[int], + delta_token_ids: Sequence[int], + request: ChatCompletionRequest | ResponsesRequest, + tool_call_idx: int | None = None, + tool_call_id_type: str = "random", + function_name_returned: bool = False, + ) -> tuple[DeltaMessage | None, bool]: + if request.tool_choice != "none": + return super()._extract_tool_calls_streaming( + previous_text, + current_text, + delta_text, + previous_token_ids, + current_token_ids, + delta_token_ids, + request, + tool_call_idx=tool_call_idx, + tool_call_id_type=tool_call_id_type, + function_name_returned=function_name_returned, + ) + + delta_message = self.extract_tool_calls_streaming( + previous_text, + current_text, + delta_text, + previous_token_ids, + current_token_ids, + delta_token_ids, + request, + ) + if delta_message is not None: + delta_message.tool_calls = [] + return delta_message, False + + def parse_delta( + self, + delta_text: str, + delta_token_ids: list[int], + request: ChatCompletionRequest | ResponsesRequest, + prompt_token_ids: list[int] | None = None, + *, + finished: bool, + ) -> DeltaMessage | None: + state = self._stream_state + previous_content = state.previous_text if state.reasoning_ended else "" + delta_message = super().parse_delta( + delta_text, + delta_token_ids, + request, + prompt_token_ids, + finished=finished, + ) + + if ( + self._tool_parser is not None + or not isinstance(self._reasoning_parser, KimiK3ReasoningParser) + or not state.reasoning_ended + or delta_message is None + ): + return delta_message + + stripped = self._reasoning_parser.strip_content_streaming( + previous_text=previous_content, + current_text=state.previous_text, + ) + delta_message.content = stripped.content if stripped is not None else None + if ( + delta_message.role is None + and delta_message.content is None + and delta_message.reasoning is None + and not delta_message.tool_calls + ): + return None + return delta_message diff --git a/aphrodite/parser/parser_manager.py b/aphrodite/parser/parser_manager.py index 640f7ddd3b..b0a400c5ec 100644 --- a/aphrodite/parser/parser_manager.py +++ b/aphrodite/parser/parser_manager.py @@ -117,6 +117,18 @@ def get_parser( MistralParser.tool_parser_cls = tool_parser_cls return MistralParser + if reasoning_parser_name == "kimi_k3" or tool_parser_name == "kimi_k3": + from aphrodite.parser.kimi_k3 import KimiK3Parser + + r_cls = reasoning_parser_cls + t_cls = tool_parser_cls + + class _KimiK3Parser(KimiK3Parser): + reasoning_parser_cls = r_cls + tool_parser_cls = t_cls + + return _KimiK3Parser + from aphrodite.parser.abstract_parser import DelegatingParser r_cls = reasoning_parser_cls diff --git a/aphrodite/pooling_params.py b/aphrodite/pooling_params.py index b4f00a1424..f763f249c8 100644 --- a/aphrodite/pooling_params.py +++ b/aphrodite/pooling_params.py @@ -7,6 +7,7 @@ import msgspec from aphrodite.config import ModelConfig, PoolerConfig +from aphrodite.exceptions import APHRODITEValidationError from aphrodite.logger import init_logger from aphrodite.sampling_params import RequestOutputKind from aphrodite.tasks import PoolingTask, check_removed_pooling_task @@ -145,7 +146,7 @@ def _verify_step_pooling( invalid_parameters.append(k) if invalid_parameters: - raise ValueError( + raise APHRODITEValidationError( f"Task {self.task} only supports {valid_parameters} " f"parameters, does not support " f"{invalid_parameters} parameters" @@ -170,20 +171,20 @@ def _set_default_parameters(self, model_config: ModelConfig): valid_range = f"[1, {embedding_size}]" dimensions_in_range = 1 <= dimensions <= embedding_size if not model_config.is_matryoshka: - raise ValueError( + raise APHRODITEValidationError( f"Model {model_name!r} does not support Matryoshka " f"embeddings; dimensions must be unset " f"(received dimensions={dimensions})." ) if not dimensions_in_range: - raise ValueError( + raise APHRODITEValidationError( f"Model {model_name!r} only supports dimensions in range {valid_range}, got {dimensions}." ) mds = model_config.matryoshka_dimensions if mds is not None and dimensions not in mds: - raise ValueError( + raise APHRODITEValidationError( f"Model {model_name!r} only supports Matryoshka dimensions {str(mds)}, got {dimensions}." ) @@ -206,7 +207,7 @@ def _verify_valid_parameters(self): invalid_parameters.append(k) if invalid_parameters: - raise ValueError( + raise APHRODITEValidationError( f"Task {self.task!r} only supports {valid_parameters} " f"parameters, does not support " f"{invalid_parameters} parameters" @@ -229,4 +230,4 @@ def __repr__(self) -> str: def __post_init__(self) -> None: check_removed_pooling_task(self.task) if self.output_kind != RequestOutputKind.FINAL_ONLY: - raise ValueError(f"For pooling output_kind has to be FINAL_ONLY, got {self.output_kind!r}") + raise APHRODITEValidationError(f"For pooling output_kind has to be FINAL_ONLY, got {self.output_kind!r}") diff --git a/aphrodite/reasoning/__init__.py b/aphrodite/reasoning/__init__.py index 186b84b03b..0b7cf7b164 100644 --- a/aphrodite/reasoning/__init__.py +++ b/aphrodite/reasoning/__init__.py @@ -84,6 +84,10 @@ "kimi_k2_reasoning_parser", "KimiK2ReasoningParser", ), + "kimi_k3": ( + "kimi_k3_reasoning_parser", + "KimiK3ReasoningParser", + ), "mimo": ( "qwen3_engine_reasoning_parser", "Qwen3ParserReasoningAdapter", diff --git a/aphrodite/reasoning/kimi_k3_reasoning_parser.py b/aphrodite/reasoning/kimi_k3_reasoning_parser.py new file mode 100644 index 0000000000..f96e8356fd --- /dev/null +++ b/aphrodite/reasoning/kimi_k3_reasoning_parser.py @@ -0,0 +1,357 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the Aphrodite project +"""Reasoning parser for the Kimi K3 (XTML) chat format. + +This strips the ``think`` channel out of generated text and hands the remainder +(``response`` + ``tools`` channels) downstream. Kimi K3 wraps the thinking +channel as an XTML element built from special tokens:: + + <|open|>think<|sep|> <|close|>think<|sep|> + +Two subtleties drive the implementation: + * Unlike Kimi-K2 (a single ```` token), each K3 marker is a 3-token + sequence, so the token-id helpers search for the marker *subsequence* + rather than a single id. + * In thinking mode the serving layer may feed ``<|open|>think<|sep|>`` as the + generation prefix, so the model's output can begin *inside* the think + channel with no open marker. The text paths therefore treat a missing open + marker as "reasoning starts at offset 0". + +When thinking is disabled (``chat_template_kwargs={"thinking": False}`` or +``{"enable_thinking": False}``, i.e. instruct mode) the parser returns every +delta as normal content; there is simply no think channel to extract. +""" + +from collections.abc import Sequence +from typing import TYPE_CHECKING + +import regex as re +from transformers import PreTrainedTokenizerBase + +from aphrodite.entrypoints.openai.engine.protocol import DeltaMessage +from aphrodite.reasoning import ReasoningParser + +if TYPE_CHECKING: + from aphrodite.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest + from aphrodite.entrypoints.openai.responses.protocol import ResponsesRequest + + +def _subseq_index(haystack: Sequence[int], needle: Sequence[int]) -> int: + """Return start index of the last occurrence of needle in haystack, or -1.""" + n = len(needle) + if n == 0: + return -1 + for i in range(len(haystack) - n, -1, -1): + if list(haystack[i : i + n]) == list(needle): + return i + return -1 + + +class KimiK3ReasoningParser(ReasoningParser): + """Reasoning parser for the Kimi K3 (XTML) think channel.""" + + def __init__(self, tokenizer: PreTrainedTokenizerBase, *args, **kwargs): + super().__init__(tokenizer) + + if not self.model_tokenizer: + raise ValueError( + "The model tokenizer must be passed to the ReasoningParser constructor during construction." + ) + + # thinking can be disabled via chat_template_kwargs -> identity fallthrough + chat_kwargs = kwargs.get("chat_template_kwargs", {}) or {} + thinking = chat_kwargs.get("thinking", None) + if thinking is None: + thinking = chat_kwargs.get("enable_thinking", True) + self._thinking_enabled = bool(thinking) + + # XTML markers as literal strings (skip_special_tokens=False at serve time) + self._think_open = "<|open|>think<|sep|>" + self._think_close = "<|close|>think<|sep|>" + + # Content-channel markers that must be stripped from the final output. + # The tool parser handles these when active, but when tool_choice is + # "none" (the default for requests without tools) the tool parser is + # bypassed and the reasoning parser must strip them itself. + self._response_open = "<|open|>response<|sep|>" + self._response_close = "<|close|>response<|sep|>" + self._message_close = "<|close|>message<|sep|>" + + # Tolerant matchers: defense-in-depth against Aphrodite's added-token spacing + # ("<|open|> think <|sep|>"). The `\s*` is a no-op on clean input, so the + # normal path stays byte-exact; it only helps if some serving path leaves + # spaces_between_special_tokens on. See adjust_request for the real fix. + open_marker = r"<\|open\|>" + close_marker = r"<\|close\|>" + sep_marker = r"<\|sep\|>" + self._think_open_re = re.compile(open_marker + r"\s*think\s*" + sep_marker) + self._think_close_re = re.compile(close_marker + r"\s*think\s*" + sep_marker) + self._response_open_re = re.compile(open_marker + r"\s*response\s*" + sep_marker) + self._response_close_re = re.compile(close_marker + r"\s*response\s*" + sep_marker) + self._message_close_re = re.compile(close_marker + r"\s*message\s*" + sep_marker) + + # marker token-id subsequences (each marker is 3 tokens) + self._think_open_ids = tokenizer.encode(self._think_open, add_special_tokens=False) + self._think_close_ids = tokenizer.encode(self._think_close, add_special_tokens=False) + self._last_streaming_delta_token_ids: tuple[int, ...] | None = None + self._last_streaming_content_token_ids: list[int] | None = None + + @property + def reasoning_start_str(self) -> str | None: + return self._think_open + + @property + def reasoning_end_str(self) -> str | None: + return self._think_close + + def adjust_request( + self, + request: "ChatCompletionRequest | ResponsesRequest", + ) -> "ChatCompletionRequest | ResponsesRequest": + request.skip_special_tokens = False + if hasattr(request, "spaces_between_special_tokens"): + request.spaces_between_special_tokens = False + return request + + def is_reasoning_end(self, input_ids: Sequence[int]) -> bool: + if not self._thinking_enabled: + return True + # Reasoning has ended only if the *most recent* think block is closed: + # the last close marker must come after the last open marker. A plain + # "a close marker exists anywhere" check false-positives in multi-turn / + # agent continuations, where the chat template keeps a prior turn's + # think channel (with its <|close|>think<|sep|>) in the prompt while the + # current turn is still reasoning (its <|open|>think<|sep|> is the newest + # marker). See _subseq_index (returns the last occurrence). + last_close = _subseq_index(input_ids, self._think_close_ids) + last_open = _subseq_index(input_ids, self._think_open_ids) + if last_open == -1: + # No open marker in scope (e.g. the open was consumed as the + # generation prefix): a close marker means reasoning has ended. + return last_close != -1 + return last_close > last_open + + def _extract_content_ids(self, input_ids: list[int]) -> list[int]: + if not self._thinking_enabled: + return input_ids + idx = _subseq_index(input_ids, self._think_close_ids) + if idx == -1: + return [] # still reasoning + return input_ids[idx + len(self._think_close_ids) :] + + def extract_content_ids(self, input_ids: list[int]) -> list[int]: + cached_delta_ids = self._last_streaming_delta_token_ids + cached_content_ids = self._last_streaming_content_token_ids + self._last_streaming_delta_token_ids = None + self._last_streaming_content_token_ids = None + if cached_delta_ids == tuple(input_ids) and cached_content_ids is not None: + return cached_content_ids + return self._extract_content_ids(input_ids) + + def _strip_content_wrapper(self, text: str) -> str: + """Strip ``<|open|>response<|sep|>…<|close|>response<|sep|>`` wrapper and + ``<|close|>message<|sep|>`` from *text*. + + When Kimi K3 tool parsing is active it needs the raw XTML + ``response`` + ``tools`` channels. Otherwise the reasoning parser cleans + up the response wrapper itself so API users do not see XTML markers. + """ + # Try structured unwrap first: extract body of response channel + m_ro = self._response_open_re.search(text) + m_rc = self._response_close_re.search(text, m_ro.end() if m_ro else 0) + if m_ro is not None and m_rc is not None: + text = text[m_ro.end() : m_rc.start()] + elif m_ro is not None: + # response opened but not closed (truncated) + text = text[m_ro.end() :] + else: + # No response wrapper — strip stray markers if present + text = self._response_open_re.sub("", text) + text = self._response_close_re.sub("", text) + text = self._message_close_re.sub("", text) + return text + + @staticmethod + def _should_preserve_tool_channels( + request: "ChatCompletionRequest | ResponsesRequest", + ) -> bool: + return bool(getattr(request, "tools", None)) and (getattr(request, "tool_choice", None) != "none") + + def _content_after_reasoning( + self, + text: str, + request: "ChatCompletionRequest | ResponsesRequest", + ) -> str | None: + if self._should_preserve_tool_channels(request): + return text or None + return self._strip_content_wrapper(text) or None + + def extract_reasoning( + self, model_output: str, request: "ChatCompletionRequest | ResponsesRequest" + ) -> tuple[str | None, str | None]: + """Split full text into ``(reasoning, rest)`` for the non-streaming path. + + Handles three shapes: + * no think channel at all -> ``(None, model_output)`` (all content) + * open marker present -> reasoning starts after ``<|open|>think<|sep|>`` + * open marker absent but a close marker exists (gen-prefix consumed + the open) -> reasoning starts at offset 0 + ``rest`` is whatever follows the close marker, fed on to the tool parser. + """ + if not self._thinking_enabled: + return None, self._content_after_reasoning(model_output, request) + + m_open = self._think_open_re.search(model_output) + # reasoning content begins right after think-open (or at start if the + # open marker was already consumed as a generation prefix) + content_start = m_open.end() if m_open is not None else 0 + # if there is no think channel at all, everything is content + if m_open is None and self._think_close_re.search(model_output) is None: + return None, self._content_after_reasoning(model_output, request) + + m_close = self._think_close_re.search(model_output, content_start) + if m_close is not None: + reasoning = model_output[content_start : m_close.start()] + rest = model_output[m_close.end() :] + return (reasoning or None, self._content_after_reasoning(rest, request)) + # think not closed -> still reasoning, no content yet + return (model_output[content_start:] or None, None) + + def _reasoning_text_ready_to_emit(self, text: str) -> str: + """Return the reasoning prefix that is safe to stream now. + + Work from accumulated text, not from the current delta alone. That turns + split-open and split-close handling into the same prefix-diff problem. + + Example 1: split think-open marker. + chunks: ``<|open|>`` / ``think`` / ``<|sep|>reasoning`` + current text after chunk 1: ``<|open|>`` -> emit ``""`` + current text after chunk 2: ``<|open|>think`` -> emit ``""`` + current text after chunk 3: ``<|open|>think<|sep|>reasoning`` + -> emit ``reasoning`` + + Example 2: split think-close marker. + chunks: ``reasoning`` / ``<|close|>`` / ``think<|sep|>...`` + after ``<|close|>``, the suffix is only a partial close marker, so the + sendable reasoning is still ``reasoning`` and the delta is empty. Once + ``think<|sep|>`` arrives, the close branch hands the following response + or tools text to the downstream parser. + """ + m_open = self._think_open_re.search(text) + if m_open is not None: + text = text[m_open.end() :] + overlap = 0 + for marker in (self._think_open, self._think_close): + max_check = min(len(marker) - 1, len(text)) + for n in range(max_check, 0, -1): + if text.endswith(marker[:n]): + overlap = max(overlap, n) + break + return text[:-overlap] if overlap else text + + def _content_ready_to_emit(self, text: str) -> str: + """Return the content prefix that is safe to stream now. + + Mirrors ``_reasoning_text_ready_to_emit`` but for the post-reasoning + content phase. Strips the ``<|open|>response<|sep|>`` prefix, holds back + any partial marker suffix, and removes complete + ``<|close|>response<|sep|>`` / ``<|close|>message<|sep|>`` markers. + """ + # Strip response-open prefix + m_open = self._response_open_re.search(text) + if m_open is not None: + text = text[m_open.end() :] + + # Remove complete close/message markers + text = self._response_close_re.sub("", text) + text = self._message_close_re.sub("", text) + + # Hold back partial markers at the end + overlap = 0 + for marker in ( + self._response_open, + self._response_close, + self._message_close, + ): + max_check = min(len(marker) - 1, len(text)) + for n in range(max_check, 0, -1): + if text.endswith(marker[:n]): + overlap = max(overlap, n) + break + return text[:-overlap] if overlap else text + + def strip_content_streaming( + self, + previous_text: str, + current_text: str, + ) -> DeltaMessage | None: + """Strip XTML content wrappers from streaming deltas after reasoning. + + Called by KimiK3Parser when no tool parser is configured, so the + reasoning parser handles ``<|open|>response<|sep|>`` / + ``<|close|>response<|sep|>`` / ``<|close|>message<|sep|>`` stripping itself. + + Works from accumulated text (``previous_text`` / ``current_text`` + already contain only post-reasoning content). + """ + current_safe = self._content_ready_to_emit(current_text) + previous_safe = self._content_ready_to_emit(previous_text) + if current_safe.startswith(previous_safe): + delta = current_safe[len(previous_safe) :] + else: + delta = current_safe + return DeltaMessage(content=delta) if delta else None + + def extract_reasoning_streaming( + self, + previous_text: str, + current_text: str, + delta_text: str, + previous_token_ids: Sequence[int], + current_token_ids: Sequence[int], + delta_token_ids: Sequence[int], + ) -> DeltaMessage | None: + self._last_streaming_delta_token_ids = None + self._last_streaming_content_token_ids = None + if not self._thinking_enabled: + return DeltaMessage(content=delta_text) + + # reasoning already ended -> downstream content + if self._think_close_re.search(previous_text): + return DeltaMessage(content=delta_text) + + # the close marker completes within this delta's accumulated text: + # split the buffer at the close marker into reasoning vs trailing content. + m_close = self._think_close_re.search(current_text) + if m_close is not None: + self._last_streaming_delta_token_ids = tuple(delta_token_ids) + self._last_streaming_content_token_ids = self._extract_content_ids(list(current_token_ids)) + m_open = self._think_open_re.search(current_text) + r_start = m_open.end() if m_open is not None else 0 + reasoning = current_text[r_start : m_close.start()] + already_sent = self._reasoning_text_ready_to_emit(previous_text) + if reasoning.startswith(already_sent): + reasoning_delta = reasoning[len(already_sent) :] + else: + reasoning_delta = reasoning + content = current_text[m_close.end() :] + return DeltaMessage( + reasoning=reasoning_delta or None, + content=content or None, + ) + + current_reasoning = self._reasoning_text_ready_to_emit(current_text) + previous_reasoning = self._reasoning_text_ready_to_emit(previous_text) + if current_reasoning.startswith(previous_reasoning): + reasoning_delta = current_reasoning[len(previous_reasoning) :] + else: + reasoning_delta = current_reasoning + if not reasoning_delta: + return None + return DeltaMessage(reasoning=reasoning_delta) + + # Backward-compatible aliases for existing unit tests and downstream users + # that still call the pre-split method names. + extract_reasoning_content = extract_reasoning + extract_reasoning_content_streaming = extract_reasoning_streaming diff --git a/aphrodite/renderers/kimi_k3.py b/aphrodite/renderers/kimi_k3.py new file mode 100644 index 0000000000..86ae3d23b5 --- /dev/null +++ b/aphrodite/renderers/kimi_k3.py @@ -0,0 +1,213 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the Aphrodite project +from typing import Any, cast + +from aphrodite.config import AphroditeConfig +from aphrodite.entrypoints.chat_utils import ( + ChatCompletionMessageParam, + ConversationMessage, + parse_chat_messages, + parse_chat_messages_async, +) +from aphrodite.exceptions import AphroditeValidationError +from aphrodite.multimodal.media.connector import merge_media_io_kwargs +from aphrodite.tokenizers.hf import HfTokenizer +from aphrodite.utils.async_utils import make_async + +from .base import BaseRenderer +from .inputs import DictPrompt +from .inputs.preprocess import parse_dec_only_prompt +from .params import ChatParams + +# Keep the original image mode (including the alpha channel) for K3 instead of +# flattening images onto a background color. Server-level (--media-io-kwargs) +# and request-level media_io_kwargs still take precedence over this default. +_K3_MEDIA_IO_DEFAULTS: dict[str, dict[str, Any]] = {"image": {"image_mode": None}} +_K3_THINKING_EFFORTS = ("low", "high", "max") + + +def _merge_k3_media_io_kwargs( + media_io_kwargs: dict[str, dict[str, Any]] | None, +) -> dict[str, dict[str, Any]] | None: + return merge_media_io_kwargs(_K3_MEDIA_IO_DEFAULTS, media_io_kwargs) + + +def _dump_k3_template_value(value: Any) -> Any: + model_dump = getattr(value, "model_dump", None) + if callable(model_dump): + return model_dump(mode="json", exclude_none=True) + + to_dict = getattr(value, "to_dict", None) + if callable(to_dict): + return to_dict() + + return value + + +def _apply_k3_thinking_kwargs(kwargs: dict[str, Any]) -> None: + if (enable_thinking := kwargs.pop("enable_thinking", None)) is not None: + kwargs.setdefault("thinking", enable_thinking) + + reasoning_effort = kwargs.pop("reasoning_effort", None) + if reasoning_effort == "none": + kwargs.setdefault("thinking", False) + elif reasoning_effort is not None: + kwargs.setdefault("thinking_effort", reasoning_effort) + + thinking_effort = kwargs.get("thinking_effort") + if thinking_effort is not None and thinking_effort not in _K3_THINKING_EFFORTS: + supported = ", ".join(_K3_THINKING_EFFORTS) + raise AphroditeValidationError( + f"Kimi K3 supports thinking_effort values: {supported}", + parameter="thinking_effort", + value=thinking_effort, + ) + + +def _normalize_k3_tool_messages( + conversation: list[ConversationMessage], +) -> list[dict[str, Any]]: + """Reorder tool-result messages to match assistant tool_call order. + + Supports matching by ``tool_call_id`` or by the synthetic + ``"{tool}:{zero_based_index}"`` alias. When any tool message in a + block cannot be resolved, the whole block is left in its original + order (graceful fallback). + + Returns a new list; caller-owned message dicts are not mutated. + """ + normalized: list[dict[str, Any]] = [] + i = 0 + + while i < len(conversation): + message = conversation[i] + normalized.append(dict(message)) + i += 1 + + if message.get("role") != "assistant": + continue + + tool_calls = message.get("tool_calls") + if not tool_calls: + continue + + # Build the lookup table from the assistant's tool_calls. + targets_by_id: dict[str, tuple[int, str]] = {} + for position, tool_call in enumerate(tool_calls): + function = tool_call.get("function") + if not isinstance(function, dict): + continue + name = function.get("name") + if not isinstance(name, str) or not name: + continue + + call_target = (position, name) + aliases = [f"{name}:{position}"] + if tool_call_id := tool_call.get("id"): + aliases.insert(0, str(tool_call_id)) + for alias in aliases: + targets_by_id.setdefault(alias, call_target) + + # Collect consecutive tool messages. + block_start = i + while i < len(conversation) and conversation[i].get("role") == "tool": + i += 1 + if i == block_start: + continue + + # Try to resolve every tool message in the block. + resolved: list[tuple[int, int, dict[str, Any]]] = [] + for order, tool_message in enumerate(conversation[block_start:i]): + tool_call_id = tool_message.get("tool_call_id") + resolved_target = targets_by_id.get(str(tool_call_id)) if tool_call_id is not None else None + if resolved_target is None: + normalized.extend(dict(item) for item in conversation[block_start:i]) + break + + position, name = resolved_target + enriched = dict(tool_message) + enriched["tool"] = name + enriched["index"] = position + 1 + resolved.append((position, order, enriched)) + else: + resolved.sort(key=lambda item: (item[0], item[1])) + normalized.extend(item for _, _, item in resolved) + + return normalized + + +class KimiK3Renderer(BaseRenderer[HfTokenizer]): + """Render chat prompts with Kimi K3's Python XTML encoding. + + K3 ships no Jinja chat template; its tokenizer renders messages through + ``encoding_k3`` instead. We tokenize eagerly so the structural markers keep + their special-token ids while user- and tool-supplied text stays ordinary. + """ + + def __init__(self, config: AphroditeConfig, tokenizer: HfTokenizer | None) -> None: + super().__init__(config, tokenizer) + + self._apply_chat_template_async = make_async(self._apply_chat_template, executor=self._executor) + + def _apply_chat_template( + self, + conversation: list[dict[str, Any]], + params: ChatParams, + ) -> list[int]: + # Tokenize eagerly: K3 encodes structural markers as special tokens and + # user/tool text as ordinary tokens, so we cannot defer to a plain + # re-tokenization of the rendered string downstream. + kwargs = params.get_apply_chat_template_kwargs() + _apply_k3_thinking_kwargs(kwargs) + if params.tool_choice not in (None, "auto"): + kwargs["tool_choice"] = _dump_k3_template_value(params.tool_choice) + if params.response_format is not None: + kwargs["response_format"] = _dump_k3_template_value(params.response_format) + kwargs["tokenize"] = True + return self.get_tokenizer().apply_chat_template(conversation, **kwargs) + + def render_messages( + self, + messages: list[ChatCompletionMessageParam], + params: ChatParams, + ) -> tuple[list[ConversationMessage], DictPrompt]: + conversation, mm_data, mm_uuids = parse_chat_messages( + messages, + self.model_config, + content_format="string", + media_io_kwargs=_merge_k3_media_io_kwargs(params.media_io_kwargs), + mm_processor_kwargs=params.mm_processor_kwargs, + ) + + rendered_conversation = _normalize_k3_tool_messages(conversation) + prompt = parse_dec_only_prompt(self._apply_chat_template(rendered_conversation, params)) + if mm_data is not None: + prompt["multi_modal_data"] = mm_data + if mm_uuids is not None: + prompt["multi_modal_uuids"] = mm_uuids + + return cast(list[ConversationMessage], rendered_conversation), prompt + + async def render_messages_async( + self, + messages: list[ChatCompletionMessageParam], + params: ChatParams, + ) -> tuple[list[ConversationMessage], DictPrompt]: + conversation, mm_data, mm_uuids = await parse_chat_messages_async( + messages, + self.model_config, + content_format="string", + media_io_kwargs=_merge_k3_media_io_kwargs(params.media_io_kwargs), + mm_processor_kwargs=params.mm_processor_kwargs, + ) + + rendered_conversation = _normalize_k3_tool_messages(conversation) + token_ids = await self._apply_chat_template_async(rendered_conversation, params) + prompt = parse_dec_only_prompt(token_ids) + if mm_data is not None: + prompt["multi_modal_data"] = mm_data + if mm_uuids is not None: + prompt["multi_modal_uuids"] = mm_uuids + + return cast(list[ConversationMessage], rendered_conversation), prompt diff --git a/aphrodite/renderers/online_renderer.py b/aphrodite/renderers/online_renderer.py index e961ebbd7d..d6faac3c1c 100644 --- a/aphrodite/renderers/online_renderer.py +++ b/aphrodite/renderers/online_renderer.py @@ -47,6 +47,19 @@ logger = init_logger(__name__) +def _reused_prompt_token_ids(request: Any) -> list[int] | None: + """Pop prompt token ids forwarded for decode-side reuse, if any. + + Disaggregated serving carries the prefill stage's ids in + ``kv_transfer_params`` so the decode stage can skip re-tokenizing. Removing + the key keeps the id list out of the engine's sampling metadata. + """ + kv = getattr(request, "kv_transfer_params", None) + if not isinstance(kv, dict): + return None + return kv.pop("prompt_token_ids", None) or None + + class OnlineRenderer: def __init__( self, @@ -98,6 +111,12 @@ async def render_chat( Called directly by render_chat_request and delegated to by OpenAIServingChat.render_chat_request after its engine-aware checks. + + Decode-side token reuse (ids forwarded in ``kv_transfer_params``) is + handled deeper, in ``preprocess_chat`` / ``_make_request_with_harmony``, + so it skips only templating and tokenization while tool-choice + validation and ``adjust_request`` still run and the output is + detokenized (text-out). """ tokenizer = self.renderer.tokenizer @@ -172,6 +191,13 @@ def _make_request_with_harmony( should_include_tools: bool = True, ): """Build Harmony (GPT-OSS) messages and engine prompt from a chat request.""" + reuse_ids = _reused_prompt_token_ids(request) + if reuse_ids: + # Decode-side token reuse: feed the forwarded ids straight to the + # engine. Harmony has no adjust_request hook to preserve. + engine_input = tokens_input(reuse_ids, cache_salt=request.cache_salt) + return [], [engine_input] + messages: list[OpenAIMessage] = [] # because of issues with pydantic we need to potentially @@ -337,15 +363,24 @@ async def preprocess_chat( default_mm_processor_kwargs=getattr(request, "mm_processor_kwargs", None), ) - (conversation,), (engine_input,) = await renderer.render_chat_async( - [messages], - chat_params, - tok_params, - prompt_extras={ - k: v for k in ("mm_processor_kwargs", "cache_salt") if (v := getattr(request, k, None)) is not None - }, - skip_mm_cache=skip_mm_cache, - ) + reuse_ids = _reused_prompt_token_ids(request) + if reuse_ids: + # Decode-side token reuse: feed the forwarded ids straight to the + # engine, skipping templating and tokenization. ``messages`` are not + # tokenized, so conversation is empty. The adjust_request tail below + # still runs. + conversation: list[ConversationMessage] = [] + engine_input = tokens_input(reuse_ids, cache_salt=getattr(request, "cache_salt", None)) + else: + (conversation,), (engine_input,) = await renderer.render_chat_async( + [messages], + chat_params, + tok_params, + prompt_extras={ + k: v for k in ("mm_processor_kwargs", "cache_salt") if (v := getattr(request, k, None)) is not None + }, + skip_mm_cache=skip_mm_cache, + ) # tool parsing is done only if a tool_parser has been set and if # tool_choice is not "none" (if tool_choice is "none" but a tool_parser diff --git a/aphrodite/renderers/params.py b/aphrodite/renderers/params.py index dacf9426fa..d7415e314c 100644 --- a/aphrodite/renderers/params.py +++ b/aphrodite/renderers/params.py @@ -88,6 +88,12 @@ class ChatParams: return_assistant_tokens_mask: bool = False """Request a per-token assistant mask from apply_chat_template.""" + tool_choice: Any | None = None + """Request-level tool choice for renderers that need API metadata.""" + + response_format: Any | None = None + """Request-level response format for renderers that need API metadata.""" + def with_defaults( self, default_chat_template_kwargs: dict[str, Any] | None = None, @@ -113,6 +119,8 @@ def with_defaults( self.mm_processor_kwargs, ), return_assistant_tokens_mask=self.return_assistant_tokens_mask, + tool_choice=self.tool_choice, + response_format=self.response_format, ) def get_apply_chat_template_kwargs(self) -> dict[str, Any]: diff --git a/aphrodite/renderers/registry.py b/aphrodite/renderers/registry.py index 96efb4411b..a130cb76b9 100644 --- a/aphrodite/renderers/registry.py +++ b/aphrodite/renderers/registry.py @@ -24,6 +24,7 @@ "deepseek_v4": ("deepseek_v4", "DeepseekV4Renderer"), "hf": ("hf", "HfRenderer"), "kimi_audio": ("hf", "HfRenderer"), + "kimi_k3": ("kimi_k3", "KimiK3Renderer"), "mistral": ("mistral", "MistralRenderer"), "terratorch": ("terratorch", "TerratorchRenderer"), "inkling": ("inkling", "InklingRenderer"), diff --git a/aphrodite/sampling_params.py b/aphrodite/sampling_params.py index 7d3f2a6d7c..371495fd4c 100644 --- a/aphrodite/sampling_params.py +++ b/aphrodite/sampling_params.py @@ -127,12 +127,12 @@ def __post_init__(self): ] ) if count > 1: - raise ValueError( + raise APHRODITEValidationError( "You can only use one kind of structured outputs constraint " f"but multiple are specified: {self.__dict__}" ) if count < 1: - raise ValueError( + raise APHRODITEValidationError( f"You must use one kind of structured outputs constraint but none are specified: {self.__dict__}" ) @@ -188,13 +188,13 @@ class RepetitionDetectionParams: def __post_init__(self): if self.max_pattern_size < 0 or self.min_pattern_size < 0 or self.min_pattern_size > self.max_pattern_size: - raise ValueError( + raise APHRODITEValidationError( "max_pattern_size, min_pattern_size must be >=0, " "with min_pattern_size <= max_pattern_size. " "Set both to 0 to disable repetitive pattern detection." ) if self.max_pattern_size > 0 and self.min_count < 2: - raise ValueError( + raise APHRODITEValidationError( "min_count must be >= 2 to detect repetitive patterns " "in engine output. If you do not wish to detect repetitive " "patterns, set max_pattern_size to 0." @@ -643,12 +643,12 @@ def __post_init__(self) -> None: def _verify_args(self) -> None: if not isinstance(self.n, int): - raise ValueError(f"n must be an int, but is of type {type(self.n)}") + raise APHRODITEValidationError(f"n must be an int, but is of type {type(self.n)}") if self.n < 1: - raise ValueError(f"n must be at least 1, got {self.n}.") + raise APHRODITEValidationError(f"n must be at least 1, got {self.n}.") max_n = envs.APHRODITE_MAX_N_SEQUENCES if self.n > max_n: - raise ValueError( + raise APHRODITEValidationError( f"n must be at most {max_n}, got {self.n}. " "To increase this limit, set the APHRODITE_MAX_N_SEQUENCES " "environment variable." @@ -687,9 +687,9 @@ def _verify_args(self) -> None: ) # quietly accept -1 as disabled, but prefer 0 if self.top_k < -1: - raise ValueError(f"top_k must be 0 (disable), or at least 1, got {self.top_k}.") + raise APHRODITEValidationError(f"top_k must be 0 (disable), or at least 1, got {self.top_k}.") if not isinstance(self.top_k, int): - raise TypeError(f"top_k must be an integer, got {type(self.top_k).__name__}") + raise APHRODITEValidationError(f"top_k must be an integer, got {type(self.top_k).__name__}") if self.no_repeat_ngram_size < 0: raise ValueError(f"no_repeat_ngram_size must be non-negative, got {self.no_repeat_ngram_size}.") if self.dynatemp_min < 0.0 or self.dynatemp_max < 0.0: @@ -699,7 +699,7 @@ def _verify_args(self) -> None: if self.top_a < 0.0: raise ValueError(f"top_a must be non-negative, got {self.top_a}.") if not 0.0 <= self.min_p <= 1.0: - raise ValueError(f"min_p must be in [0, 1], got {self.min_p}.") + raise APHRODITEValidationError(f"min_p must be in [0, 1], got {self.min_p}.") if not 0.0 < self.tfs <= 1.0: raise ValueError(f"tfs must be in (0, 1], got {self.tfs}.") if self.eta_cutoff < 0.0: @@ -715,9 +715,9 @@ def _verify_args(self) -> None: value=self.max_tokens, ) if self.min_tokens < 0: - raise ValueError(f"min_tokens must be greater than or equal to 0, got {self.min_tokens}.") + raise APHRODITEValidationError(f"min_tokens must be greater than or equal to 0, got {self.min_tokens}.") if self.max_tokens is not None and self.min_tokens > self.max_tokens: - raise ValueError( + raise APHRODITEValidationError( f"min_tokens must be less than or equal to max_tokens={self.max_tokens}, got {self.min_tokens}." ) if self.stream_interval is not None and self.stream_interval < 1: @@ -740,12 +740,12 @@ def _verify_args(self) -> None: ) assert isinstance(self.stop_token_ids, list) if not all(isinstance(st_id, int) for st_id in self.stop_token_ids): - raise ValueError(f"stop_token_ids must contain only integers, got {self.stop_token_ids}.") + raise APHRODITEValidationError(f"stop_token_ids must contain only integers, got {self.stop_token_ids}.") assert isinstance(self.stop, list) if any(not stop_str for stop_str in self.stop): - raise ValueError("stop cannot contain an empty string.") + raise APHRODITEValidationError("stop cannot contain an empty string.") if self.stop and not self.detokenize: - raise ValueError( + raise APHRODITEValidationError( "stop strings are only supported when detokenize is True. Set detokenize=True to use stop." ) if self.xtc_threshold < 0.0: @@ -772,11 +772,11 @@ def _verify_args(self) -> None: raise ValueError("dry_early_exit_match_len must be non-negative.") assert isinstance(self.bad_words, list) if any(not bad_word for bad_word in self.bad_words): - raise ValueError(f"bad_words cannot contain an empty string. Got bad_words={self.bad_words}") + raise APHRODITEValidationError(f"bad_words cannot contain an empty string. Got bad_words={self.bad_words}") def _verify_greedy_sampling(self) -> None: if self.n > 1: - raise ValueError(f"n must be 1 when using greedy sampling, got {self.n}.") + raise APHRODITEValidationError(f"n must be 1 when using greedy sampling, got {self.n}.") def update_from_generation_config( self, diff --git a/aphrodite/third_party/flash_linear_attention/ops/fused_norm_gate.py b/aphrodite/third_party/flash_linear_attention/ops/fused_norm_gate.py new file mode 100644 index 0000000000..21d8718d2d --- /dev/null +++ b/aphrodite/third_party/flash_linear_attention/ops/fused_norm_gate.py @@ -0,0 +1,412 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Songlin Yang, Yu Zhang +# +# This file contains code copied from the flash-linear-attention project. +# The original source was licensed under the MIT license. +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang +# ruff: noqa: E501 + +import torch +import torch.nn as nn + +from aphrodite.model_executor.custom_op import CustomOp +from aphrodite.platforms import current_platform +from aphrodite.triton_utils import tl, triton +from aphrodite.utils.math_utils import cdiv, next_power_of_2 + +@triton.heuristics( + { + "STORE_RESIDUAL_OUT": lambda args: args["residual_out"] is not None, + "HAS_RESIDUAL": lambda args: args["residual"] is not None, + "HAS_WEIGHT": lambda args: args["w"] is not None, + "HAS_BIAS": lambda args: args["b"] is not None, + } +) +@triton.jit +def layer_norm_gated_fwd_kernel( + x, # pointer to the input + g, # pointer to the gate + y, # pointer to the output + w, # pointer to the weights + b, # pointer to the biases + residual, # pointer to the residual + residual_out, # pointer to the residual + mean, # pointer to the mean + rstd, # pointer to the 1/std + eps, # epsilon to avoid division by zero + T, # number of rows in x + H: tl.constexpr, # number of heads + g_stride_n, + D: tl.constexpr, # number of columns in x + BT: tl.constexpr, + BD: tl.constexpr, + ACTIVATION: tl.constexpr, + IS_RMS_NORM: tl.constexpr, + STORE_RESIDUAL_OUT: tl.constexpr, + HAS_RESIDUAL: tl.constexpr, + HAS_WEIGHT: tl.constexpr, + HAS_BIAS: tl.constexpr, + launch_pdl: tl.constexpr, +): + i_t = tl.program_id(0) + + o_d = tl.arange(0, BD) + m_d = o_d < D + + if launch_pdl: + tl.extra.cuda.gdc_wait() + tl.extra.cuda.gdc_launch_dependents() + + p_x = tl.make_block_ptr(x, (T, D), (D, 1), (i_t * BT, 0), (BT, BD), (1, 0)) + b_x = tl.load(p_x, boundary_check=(0, 1)).to(tl.float32) + if HAS_RESIDUAL: + p_res = tl.make_block_ptr( + residual, (T, D), (D, 1), (i_t * BT, 0), (BT, BD), (1, 0) + ) + b_x += tl.load(p_res, boundary_check=(0, 1)).to(tl.float32) + if STORE_RESIDUAL_OUT: + p_res_out = tl.make_block_ptr( + residual_out, (T, D), (D, 1), (i_t * BT, 0), (BT, BD), (1, 0) + ) + tl.store(p_res_out, b_x.to(p_res_out.dtype.element_ty), boundary_check=(0, 1)) + if not IS_RMS_NORM: + b_mean = tl.sum(b_x, axis=1) / D + p_mean = tl.make_block_ptr(mean, (T,), (1,), (i_t * BT,), (BT,), (0,)) + tl.store(p_mean, b_mean.to(p_mean.dtype.element_ty), boundary_check=(0,)) + b_xbar = tl.where(m_d[None, :], b_x - b_mean[:, None], 0.0) + b_var = tl.sum(b_xbar * b_xbar, axis=1) / D + else: + b_xbar = tl.where(m_d[None, :], b_x, 0.0) + b_var = tl.sum(b_xbar * b_xbar, axis=1) / D + b_rstd = 1 / tl.sqrt(b_var + eps) + + p_rstd = tl.make_block_ptr(rstd, (T,), (1,), (i_t * BT,), (BT,), (0,)) + tl.store(p_rstd, b_rstd.to(p_rstd.dtype.element_ty), boundary_check=(0,)) + + if HAS_WEIGHT: + b_w = tl.load(w + o_d, mask=m_d).to(tl.float32) + if HAS_BIAS: + b_b = tl.load(b + o_d, mask=m_d).to(tl.float32) + b_x_hat = ( + (b_x - b_mean[:, None]) * b_rstd[:, None] + if not IS_RMS_NORM + else b_x * b_rstd[:, None] + ) + b_y = b_x_hat * b_w[None, :] if HAS_WEIGHT else b_x_hat + if HAS_BIAS: + b_y = b_y + b_b[None, :] + + # swish/sigmoid output gate + o_t = i_t * BT + tl.arange(0, BT) + o_g = (o_t // H) * g_stride_n + (o_t % H) * D + b_g = tl.load( + g + o_g[:, None] + o_d[None, :], + mask=(o_t[:, None] < T) & m_d[None, :], + other=0.0, + ).to(tl.float32) + if ACTIVATION == "swish" or ACTIVATION == "silu": + b_y = b_y * b_g * tl.sigmoid(b_g) + elif ACTIVATION == "sigmoid": + b_y = b_y * tl.sigmoid(b_g) + + # Write output + p_y = tl.make_block_ptr(y, (T, D), (D, 1), (i_t * BT, 0), (BT, BD), (1, 0)) + tl.store(p_y, b_y.to(p_y.dtype.element_ty), boundary_check=(0, 1)) + + +@triton.heuristics( + { + "STORE_RESIDUAL_OUT": lambda args: args["residual_out"] is not None, + "HAS_RESIDUAL": lambda args: args["residual"] is not None, + "HAS_WEIGHT": lambda args: args["w"] is not None, + "HAS_BIAS": lambda args: args["b"] is not None, + } +) +@triton.jit +def layer_norm_gated_fwd_kernel1( + x, # pointer to the input + g, # pointer to the gate + y, # pointer to the output + w, # pointer to the weights + b, # pointer to the biases + residual, # pointer to the residual + residual_out, # pointer to the residual + mean, # pointer to the mean + rstd, # pointer to the 1/std + eps, # epsilon to avoid division by zero + D: tl.constexpr, # number of columns in x + BD: tl.constexpr, + ACTIVATION: tl.constexpr, + IS_RMS_NORM: tl.constexpr, + STORE_RESIDUAL_OUT: tl.constexpr, + HAS_RESIDUAL: tl.constexpr, + HAS_WEIGHT: tl.constexpr, + HAS_BIAS: tl.constexpr, + launch_pdl: tl.constexpr, +): + i_t = tl.program_id(0) + x += i_t * D + y += i_t * D + g += i_t * D + if HAS_RESIDUAL: + residual += i_t * D + if STORE_RESIDUAL_OUT: + residual_out += i_t * D + + if launch_pdl: + tl.extra.cuda.gdc_wait() + tl.extra.cuda.gdc_launch_dependents() + + o_d = tl.arange(0, BD) + m_d = o_d < D + b_x = tl.load(x + o_d, mask=m_d, other=0.0).to(tl.float32) + if HAS_RESIDUAL: + b_x += tl.load(residual + o_d, mask=m_d, other=0.0).to(tl.float32) + if STORE_RESIDUAL_OUT: + tl.store(residual_out + o_d, b_x, mask=m_d) + if not IS_RMS_NORM: + b_mean = tl.sum(b_x, axis=0) / D + tl.store(mean + i_t, b_mean) + b_xbar = tl.where(m_d, b_x - b_mean, 0.0) + b_var = tl.sum(b_xbar * b_xbar, axis=0) / D + else: + b_xbar = tl.where(m_d, b_x, 0.0) + b_var = tl.sum(b_xbar * b_xbar, axis=0) / D + b_rstd = 1 / tl.sqrt(b_var + eps) + tl.store(rstd + i_t, b_rstd) + + if HAS_WEIGHT: + b_w = tl.load(w + o_d, mask=m_d).to(tl.float32) + if HAS_BIAS: + b_b = tl.load(b + o_d, mask=m_d).to(tl.float32) + b_x_hat = (b_x - b_mean) * b_rstd if not IS_RMS_NORM else b_x * b_rstd + b_y = b_x_hat * b_w if HAS_WEIGHT else b_x_hat + if HAS_BIAS: + b_y = b_y + b_b + + # swish/sigmoid output gate + b_g = tl.load(g + o_d, mask=m_d, other=0.0).to(tl.float32) + if ACTIVATION == "swish" or ACTIVATION == "silu": + b_y = b_y * b_g * tl.sigmoid(b_g) + elif ACTIVATION == "sigmoid": + b_y = b_y * tl.sigmoid(b_g) + + # Write output + tl.store(y + o_d, b_y, mask=m_d) + + +def layer_norm_gated_fwd( + x: torch.Tensor, + g: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor, + activation: str = "swish", + eps: float = 1e-5, + residual: torch.Tensor = None, + out_dtype: torch.dtype = None, + residual_dtype: torch.dtype = None, + is_rms_norm: bool = False, + H: int = 1, + g_stride_n: int | None = None, +): + if residual is not None: + residual_dtype = residual.dtype + T, D = x.shape + if g_stride_n is None: + g_stride_n = D + assert T % H == 0 + if residual is not None: + assert residual.shape == (T, D) + if weight is not None: + assert weight.shape == (D,) + if bias is not None: + assert bias.shape == (D,) + # allocate output + y = x if out_dtype is None else torch.empty_like(x, dtype=out_dtype) + if residual is not None or ( + residual_dtype is not None and residual_dtype != x.dtype + ): + residual_out = torch.empty(T, D, device=x.device, dtype=residual_dtype) + else: + residual_out = None + mean = ( + torch.empty((T,), dtype=torch.float, device=x.device) + if not is_rms_norm + else None + ) + rstd = torch.empty((T,), dtype=torch.float, device=x.device) + # Less than 64KB per feature: enqueue fused kernel + MAX_FUSED_SIZE = 65536 // x.element_size() + BD = min(MAX_FUSED_SIZE, next_power_of_2(D)) + if D > BD: + raise RuntimeError("This layer norm doesn't support feature dim >= 64KB.") + if D <= 512: + BT = 16 + layer_norm_gated_fwd_kernel[(cdiv(T, BT),)]( + x=x, + g=g, + y=y, + w=weight, + b=bias, + residual=residual, + residual_out=residual_out, + mean=mean, + rstd=rstd, + eps=eps, + T=T, + H=H, + g_stride_n=g_stride_n, + D=D, + BD=BD, + BT=BT, + ACTIVATION=activation, + IS_RMS_NORM=is_rms_norm, + num_warps=8, + launch_pdl=current_platform.is_arch_support_pdl(), + ) + else: + layer_norm_gated_fwd_kernel1[(T,)]( + x=x, + g=g, + y=y, + w=weight, + b=bias, + residual=residual, + residual_out=residual_out, + mean=mean, + rstd=rstd, + eps=eps, + D=D, + BD=BD, + ACTIVATION=activation, + IS_RMS_NORM=is_rms_norm, + num_warps=4, + launch_pdl=current_platform.is_arch_support_pdl(), + ) + # residual_out is None if residual is None and residual_dtype == input_dtype + return y, mean, rstd, residual_out if residual_out is not None else x + + +def rms_norm_gated( + x: torch.Tensor, + g: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor, + activation: str = "swish", + residual: torch.Tensor | None = None, + prenorm: bool = False, + residual_in_fp32: bool = False, + eps: float = 1e-6, +): + x_shape_og = x.shape + # reshape input data into 2D tensor + x = x.contiguous().reshape(-1, x.shape[-1]) + D = x.shape[-1] + # The tiled kernel supports row-strided gates; kernel1 does not. + if D <= 512: + H = 1 if g.ndim == 2 else g.shape[-2] + g = g.view(-1, H, D) + g_stride_n = g.stride(0) + else: + g = g.contiguous() + H = 1 + g_stride_n = D + if residual is not None: + assert residual.shape == x_shape_og + residual = residual.contiguous().reshape(-1, residual.shape[-1]) + residual_dtype = ( + residual.dtype + if residual is not None + else (torch.float if residual_in_fp32 else None) + ) + y, _, _, residual_out = layer_norm_gated_fwd( + x=x, + g=g, + weight=weight, + bias=bias, + activation=activation, + eps=eps, + residual=residual, + residual_dtype=residual_dtype, + is_rms_norm=True, + H=H, + g_stride_n=g_stride_n, + ) + y = y.reshape(x_shape_og) + return y if not prenorm else (y, residual_out.reshape(x_shape_og)) + + +@CustomOp.register("fused_rms_norm_gated") +class FusedRMSNormGated(CustomOp): + def __init__( + self, + hidden_size: int, + elementwise_affine: bool = True, + eps: float = 1e-5, + activation: str = "swish", + device: torch.device | None = None, + dtype: torch.dtype | None = None, + ) -> None: + factory_kwargs = {"device": device, "dtype": dtype} + super().__init__() + + self.hidden_size = hidden_size + self.elementwise_affine = elementwise_affine + self.eps = eps + self.activation = activation + + if self.activation not in ["swish", "silu", "sigmoid"]: + raise ValueError(f"Unsupported activation: {self.activation}") + + if elementwise_affine: + self.weight = nn.Parameter(torch.empty(hidden_size, **factory_kwargs)) + else: + self.register_parameter("weight", None) + self.register_parameter("bias", None) + + def forward_native( + self, + x: torch.Tensor, + g: torch.Tensor, + residual: torch.Tensor | None = None, + prenorm: bool = False, + residual_in_fp32: bool = False, + ) -> torch.Tensor: + """Decomposed PyTorch ops for torch.compile/inductor fusion.""" + # TODO(https://github.com/vllm-project/aphrodite/issues/36175): implement + # native residual/prenorm path and unify with RMSNormGated. + # For now, fall back to the triton kernel. + if residual is not None or prenorm: + return self.forward_cuda(x, g, residual, prenorm, residual_in_fp32) + x_float = x.float() + variance = x_float.pow(2).mean(dim=-1, keepdim=True) + x_normed = x_float * torch.rsqrt(variance + self.eps) + if self.weight is not None: + x_normed = x_normed * self.weight.float() + g_float = g.float() + if self.activation in ("swish", "silu"): + out = x_normed * g_float * torch.sigmoid(g_float) + else: # sigmoid + out = x_normed * torch.sigmoid(g_float) + return out.to(x.dtype) + + def forward_cuda( + self, + x: torch.Tensor, + g: torch.Tensor, + residual: torch.Tensor | None = None, + prenorm: bool = False, + residual_in_fp32: bool = False, + ) -> torch.Tensor: + return rms_norm_gated( + x, + g, + self.weight, + self.bias, + self.activation, + residual=residual, + eps=self.eps, + prenorm=prenorm, + residual_in_fp32=residual_in_fp32, + ) diff --git a/aphrodite/third_party/flash_linear_attention/ops/kda.py b/aphrodite/third_party/flash_linear_attention/ops/kda.py index 062d073b00..8fe500aba7 100644 --- a/aphrodite/third_party/flash_linear_attention/ops/kda.py +++ b/aphrodite/third_party/flash_linear_attention/ops/kda.py @@ -166,6 +166,8 @@ def layer_norm_gated_fwd_kernel( rstd, # pointer to the 1/std eps, # epsilon to avoid division by zero T, # number of rows in x + H: tl.constexpr, # number of heads + g_stride_n: tl.constexpr, D: tl.constexpr, # number of columns in x BT: tl.constexpr, BD: tl.constexpr, @@ -213,8 +215,13 @@ def layer_norm_gated_fwd_kernel( b_y = b_y + b_b[None, :] # swish/sigmoid output gate - p_g = tl.make_block_ptr(g, (T, D), (D, 1), (i_t * BT, 0), (BT, BD), (1, 0)) - b_g = tl.load(p_g, boundary_check=(0, 1)).to(tl.float32) + o_t = i_t * BT + tl.arange(0, BT) + o_g = (o_t // H) * g_stride_n + (o_t % H) * D + b_g = tl.load( + g + o_g[:, None] + o_d[None, :], + mask=(o_t[:, None] < T) & m_d[None, :], + other=0.0, + ).to(tl.float32) if ACTIVATION == "swish" or ACTIVATION == "silu": b_y = b_y * b_g * tl.sigmoid(b_g) elif ACTIVATION == "sigmoid": @@ -312,10 +319,15 @@ def layer_norm_gated_fwd( out_dtype: torch.dtype = None, residual_dtype: torch.dtype = None, is_rms_norm: bool = False, + H: int = 1, + g_stride_n: int | None = None, ): if residual is not None: residual_dtype = residual.dtype T, D = x.shape + if g_stride_n is None: + g_stride_n = D + assert T % H == 0 if residual is not None: assert residual.shape == (T, D) if weight is not None: @@ -335,10 +347,8 @@ def layer_norm_gated_fwd( BD = min(MAX_FUSED_SIZE, next_power_of_2(D)) if D > BD: raise RuntimeError("This layer norm doesn't support feature dim >= 64KB.") - # heuristics for number of warps - if D <= 512: - BT = 32 + BT = 16 layer_norm_gated_fwd_kernel[(cdiv(T, BT),)]( x=x, g=g, @@ -351,12 +361,14 @@ def layer_norm_gated_fwd( rstd=rstd, eps=eps, T=T, + H=H, + g_stride_n=g_stride_n, D=D, BD=BD, BT=BT, ACTIVATION=activation, IS_RMS_NORM=is_rms_norm, - num_warps=4, + num_warps=8, ) else: layer_norm_gated_fwd_kernel1[(T,)]( @@ -394,7 +406,16 @@ def rms_norm_gated( x_shape_og = x.shape # reshape input data into 2D tensor x = x.contiguous().reshape(-1, x.shape[-1]) - g = g.contiguous().reshape(-1, g.shape[-1]) + D = x.shape[-1] + # The tiled kernel supports row-strided gates; kernel1 does not. + if D <= 512: + H = 1 if g.ndim == 2 else g.shape[-2] + g = g.view(-1, H, D) + g_stride_n = g.stride(0) + else: + g = g.contiguous() + H = 1 + g_stride_n = D if residual is not None: assert residual.shape == x_shape_og residual = residual.contiguous().reshape(-1, residual.shape[-1]) @@ -409,6 +430,8 @@ def rms_norm_gated( residual=residual, residual_dtype=residual_dtype, is_rms_norm=True, + H=H, + g_stride_n=g_stride_n, ) y = y.reshape(x_shape_og) return y if not prenorm else (y, residual_out.reshape(x_shape_og)) @@ -1143,6 +1166,7 @@ def kda_gate_cumsum_fwd_kernel( cu_seqlens, chunk_indices, cumsum_scale, + lower_bound, beta, threshold, T, @@ -1152,6 +1176,7 @@ def kda_gate_cumsum_fwd_kernel( BD: tl.constexpr, HAS_BIAS: tl.constexpr, IS_VARLEN: tl.constexpr, + USE_LOWER_BOUND: tl.constexpr, ): i_d, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) i_b, i_h = i_bh // H, i_bh % H @@ -1191,14 +1216,17 @@ def kda_gate_cumsum_fwd_kernel( b_bias = tl.load(g_bias + i_h * D + o_d, mask=o_d < D, other=0.0).to(tl.float32) b_g = b_g + b_bias[None, :] - b_a = -tl.exp(tl.load(A + i_h).to(tl.float32)) - b_g_scaled = b_g * beta - b_softplus = tl.where( - b_g_scaled > threshold, - b_g, - (1.0 / beta) * log(1.0 + tl.exp(b_g_scaled)), - ) - b_gate = b_a * b_softplus + b_a = tl.exp(tl.load(A + i_h).to(tl.float32)) + if USE_LOWER_BOUND: + b_gate = lower_bound * tl.sigmoid(b_a * b_g) + else: + b_g_scaled = b_g * beta + b_softplus = tl.where( + b_g_scaled > threshold, + b_g, + (1.0 / beta) * log(1.0 + tl.exp(b_g_scaled)), + ) + b_gate = -b_a * b_softplus # Out-of-bounds rows (load returns 0, but softplus/bias can still make # b_gate non-zero) participate in the dot product. They only contribute to @@ -1216,6 +1244,7 @@ def fused_kda_gate_chunk_cumsum( g_bias: torch.Tensor | None = None, beta: float = 1.0, threshold: float = 20.0, + lower_bound: float | None = None, cu_seqlens: torch.Tensor | None = None, chunk_indices: torch.Tensor | None = None, chunk_size: int = FLA_CHUNK_SIZE, @@ -1247,16 +1276,17 @@ def grid(meta): # exp2-based kernels reproduce exp(g). Keep this in sync with the # `use_exp2=True` path in `_chunk_kda_fwd_with_cumulative_g`. cumsum_scale=RCP_LN2, + lower_bound=lower_bound or 0.0, beta=beta, threshold=threshold, T=T, H=H, D=D, BT=chunk_size, + USE_LOWER_BOUND=lower_bound is not None, ) return y - def _chunk_kda_fwd_with_cumulative_g( q: torch.Tensor, k: torch.Tensor, @@ -1374,6 +1404,7 @@ def chunk_kda_with_fused_gate_fwd( scale: float, initial_state: torch.Tensor, output_final_state: bool, + lower_bound: float | None = None, cu_seqlens: torch.Tensor | None = None, ): chunk_size = FLA_CHUNK_SIZE @@ -1385,6 +1416,7 @@ def chunk_kda_with_fused_gate_fwd( cu_seqlens=cu_seqlens, chunk_indices=chunk_indices, chunk_size=chunk_size, + lower_bound=lower_bound, ) return _chunk_kda_fwd_with_cumulative_g( q=q, @@ -1446,6 +1478,7 @@ def chunk_kda_with_fused_gate( scale: float | None = None, initial_state: torch.Tensor | None = None, output_final_state: bool = False, + lower_bound: float | None = None, use_qk_l2norm_in_kernel: bool = False, cu_seqlens: torch.Tensor | None = None, **kwargs, @@ -1469,6 +1502,7 @@ def chunk_kda_with_fused_gate( scale=scale, initial_state=initial_state.contiguous() if initial_state is not None else None, output_final_state=output_final_state, + lower_bound=lower_bound, cu_seqlens=cu_seqlens, ) return o, final_state @@ -1489,6 +1523,7 @@ def kda_gate_fwd_kernel( A, y, g_bias, + lower_bound, beta: tl.constexpr, threshold: tl.constexpr, T, @@ -1497,12 +1532,12 @@ def kda_gate_fwd_kernel( BT: tl.constexpr, BD: tl.constexpr, HAS_BIAS: tl.constexpr, + USE_LOWER_BOUND: tl.constexpr, ): i_t, i_h = tl.program_id(0), tl.program_id(1) n_t = i_t * BT - b_a = tl.load(A + i_h).to(tl.float32) - b_a = -tl.exp(b_a) + b_a = tl.exp(tl.load(A + i_h).to(tl.float32)) stride_row = H * D stride_col = 1 @@ -1533,13 +1568,13 @@ def kda_gate_fwd_kernel( b_bias = tl.load(g_bias + i_h * D + n_d, mask=bias_mask, other=0.0).to(tl.float32) b_g = b_g + b_bias[None, :] - # softplus(x, beta) = (1/beta) * log(1 + exp(beta * x)) - # When beta * x > threshold, use linear approximation x - # Use threshold to switch to linear when beta*x > threshold - g_scaled = b_g * beta - use_linear = g_scaled > threshold - sp = tl.where(use_linear, b_g, (1.0 / beta) * log(1.0 + tl.exp(g_scaled))) - b_y = b_a * sp + if USE_LOWER_BOUND: + b_y = lower_bound * tl.sigmoid(b_a * b_g) + else: + g_scaled = b_g * beta + use_linear = g_scaled > threshold + sp = tl.where(use_linear, b_g, (1.0 / beta) * log(1.0 + tl.exp(g_scaled))) + b_y = -b_a * sp tl.store(y_ptr, b_y.to(y.dtype.element_ty), boundary_check=(0, 1)) @@ -1551,6 +1586,7 @@ def fused_kda_gate( g_bias: torch.Tensor | None = None, beta: float = 1.0, threshold: float = 20.0, + lower_bound: float | None = None, ) -> torch.Tensor: """ Forward pass for KDA gate: @@ -1578,6 +1614,7 @@ def grid(meta): A, y, g_bias, + lower_bound or 0.0, beta, threshold, T, @@ -1585,6 +1622,7 @@ def grid(meta): head_k_dim, BD=next_power_of_2(head_k_dim), HAS_BIAS=g_bias is not None, + USE_LOWER_BOUND=lower_bound is not None, ) y = y.view(*orig_shape, H, head_k_dim) diff --git a/aphrodite/tokenizers/registry.py b/aphrodite/tokenizers/registry.py index 980865c8b0..5fd7b3b0ee 100644 --- a/aphrodite/tokenizers/registry.py +++ b/aphrodite/tokenizers/registry.py @@ -44,6 +44,7 @@ "deepseek_v4": ("deepseek_v4", "DeepseekV4Tokenizer"), "hf": ("hf", "CachedHfTokenizer"), "kimi_audio": ("kimi_audio", "KimiAudioTokenizer"), + "kimi_k3": ("hf", "CachedHfTokenizer"), "mistral": ("mistral", "MistralTokenizer"), # Inkling uses the plain HF tokenizer for token operations; the "inkling" # mode exists to select the InklingRenderer, which renders chat to diff --git a/aphrodite/tool_parsers/__init__.py b/aphrodite/tool_parsers/__init__.py index 7b503626dc..1c8c272a44 100644 --- a/aphrodite/tool_parsers/__init__.py +++ b/aphrodite/tool_parsers/__init__.py @@ -102,6 +102,10 @@ "kimi_k2_tool_parser", "KimiK2ToolParser", ), + "kimi_k3": ( + "kimi_k3_tool_parser", + "KimiK3ToolParser", + ), "llama3_json": ( "llama_tool_parser", "Llama3JsonToolParser", diff --git a/aphrodite/tool_parsers/kimi_k3_tool_parser.py b/aphrodite/tool_parsers/kimi_k3_tool_parser.py new file mode 100644 index 0000000000..98dac83889 --- /dev/null +++ b/aphrodite/tool_parsers/kimi_k3_tool_parser.py @@ -0,0 +1,379 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the Aphrodite project +"""Tool-call parser for the Kimi K3 (XTML) chat format. + +This turns the generated XTML ``response`` and ``tools`` channels back into +OpenAI-compatible ``content`` and ``tool_calls``. + +K3 assistant tool calls live in a nested ``tools`` channel:: + + <|open|>tools<|sep|> + <|open|>call tool="python" index="1"<|sep|> + <|open|>argument key="code" type="string"<|sep|>print(1)<|close|>argument<|sep|> + <|open|>argument key="opts" type="object"<|sep|>{"a":1}<|close|>argument<|sep|> + <|close|>call<|sep|> + <|close|>tools<|sep|> + +where ``<|open|>``, ``<|close|>``, ``<|sep|>`` are dedicated special tokens. The plain +reply lives in a sibling ``response`` channel which is unwrapped into content. +``response`` always precedes ``tools`` (see the protocol channel grammar). + +Argument decoding mirrors the template's type tagging (inverse encoding): + * type="string" -> value is the RAW text, no unescaping (template emits it raw) + * other types -> value is JSON-decoded (number/boolean/null/object/array) + +Attribute values (``tool=``, ``index=``, ``key=``, ``type=``) ARE escaped on the +encode side (``&`` -> ``&``, ``"`` -> ``"``), so ``_attrs`` reverses +them on decode (``"`` before ``&`` -- reverse of encode order). + +Known limitation: because string argument and response bodies are emitted raw, +a value that literally contains ``<|close|>argument<|sep|>`` or +``<|close|>response<|sep|>`` is indistinguishable from a real closing marker. +""" + +import json +from collections.abc import Sequence + +import regex as re +from openai.types.responses import ToolChoiceFunction + +from aphrodite.entrypoints.openai.chat_completion.protocol import ( + ChatCompletionNamedToolChoiceParam, + ChatCompletionRequest, +) +from aphrodite.entrypoints.openai.engine.protocol import ( + DeltaFunctionCall, + DeltaMessage, + DeltaToolCall, + ExtractedToolCallInformation, + FunctionCall, + ToolCall, +) +from aphrodite.entrypoints.openai.responses.protocol import ResponsesRequest +from aphrodite.exceptions import AphroditeValidationError +from aphrodite.logger import init_logger +from aphrodite.tokenizers import TokenizerLike +from aphrodite.tool_parsers.abstract_tool_parser import Tool, ToolParser + +logger = init_logger(__name__) + +_O, _C, _S = r"<\|open\|>", r"<\|close\|>", r"<\|sep\|>" +_TEXT_UNTIL_SEP = r"(?:(?!" + _S + r").)*?" + + +def _partial_tag_overlap(text: str, tag: str) -> int: + max_len = min(len(text), len(tag) - 1) + for n in range(max_len, 0, -1): + if text.endswith(tag[:n]): + return n + return 0 + + +class KimiK3ToolParser(ToolParser): + supports_required_and_named = False + # Enables the Aphrodite-side XTML structural tag builder + # (``get_kimi_k3_structural_tag`` in ``structural_tag_registry``). With + # ``APHRODITE_ENFORCE_STRICT_TOOL_CALLING`` on (default), ``_apply_structural_tag`` + # constrains generation to K3's ``<|open|>tools<|sep|>`` channel for + # ``required`` (and ``auto`` when a tool sets ``strict``) instead of the + # generic JSON guided decoding, which conflicts with the XTML format. + structural_tag_model = "kimi_k3" + + def __init__(self, tokenizer: TokenizerLike, tools: list[Tool] | None = None): + super().__init__(tokenizer, tools) + + self.tools_open = "<|open|>tools<|sep|>" + self.tools_close = "<|close|>tools<|sep|>" + self.response_open = "<|open|>response<|sep|>" + self.response_close = "<|close|>response<|sep|>" + + # Regexes operate on detokenized text. The XTML markers reach us as the + # literal strings <|open|>/<|close|>/<|sep|>. adjust_request keeps them + # from being stripped and suppresses spacing between adjacent token + # pieces. As DEFENSE-IN-DEPTH we still tolerate optional whitespace + # WITHIN a marker (e.g. "<|open|> tools <|sep|>") in case some serving + # path leaves Aphrodite's added-token spacing on -- this `\s*` is a no-op on + # clean input, so the normal byte-exact path is unaffected. We do NOT + # strip body whitespace, so clean content stays byte-exact. + # All bodies use non-greedy (.*?) so each block stops at its own first + # closing marker -- see the module-level KNOWN LIMITATION about + # literal-marker values. + self._tools_open_re = re.compile(_O + r"\s*tools\s*" + _S) + self._tools_close_re = re.compile(_C + r"\s*tools\s*" + _S) + self._response_open_re = re.compile(_O + r"\s*response\s*" + _S) + self._response_close_re = re.compile(_C + r"\s*response\s*" + _S) + self._message_close_re = re.compile(_C + r"\s*message\s*" + _S) + self._call_re = re.compile( + _O + r"\s*call\s+(?P" + _TEXT_UNTIL_SEP + r")" + _S + r"(?P.*?)" + _C + r"\s*call\s*" + _S, + re.DOTALL, + ) + self._arg_re = re.compile( + _O + + r"\s*argument\s+(?P" + + _TEXT_UNTIL_SEP + + r")" + + _S + + r"(?P.*?)" + + _C + + r"\s*argument\s*" + + _S, + re.DOTALL, + ) + # attr segment: key="value" (value already escaped on the encode side) + self._attr_re = re.compile(r'(?P\w+)="(?P[^"]*)"') + self._response_re = re.compile( + _O + r"\s*response\s*" + _S + r"(?P.*?)" + _C + r"\s*response\s*" + _S, + re.DOTALL, + ) + + # streaming state + self._sent_content_idx = 0 + self._sent_tool_call_count = 0 + + if not self.model_tokenizer: + raise ValueError("The model tokenizer must be passed to the ToolParser constructor during construction.") + + def adjust_request( + self, request: ChatCompletionRequest | ResponsesRequest + ) -> ChatCompletionRequest | ResponsesRequest: + named = isinstance( + request.tool_choice, + (ChatCompletionNamedToolChoiceParam, ToolChoiceFunction), + ) + structured_outputs = getattr(request, "structured_outputs", None) + has_structural_tag = structured_outputs is not None and structured_outputs.structural_tag is not None + if named and not has_structural_tag: + # Without the XTML structural tag there is no way to force the + # named call (the generic JSON guided-decoding path conflicts + # with the XTML channel format). + raise AphroditeValidationError( + "Named tool choice for Kimi K3 requires strict tool calling " + "(APHRODITE_ENFORCE_STRICT_TOOL_CALLING) so the XTML structural " + "tag can force the call. Otherwise use `tool_choice` set to " + '"auto", "required", or "none".', + parameter="tool_choice", + value=request.tool_choice, + ) + + if request.tools and (request.tool_choice == "required" or named): + # K3 emits tool calls in XTML. When strict tool calling is enabled, + # DelegatingParser._apply_structural_tag has already attached the K3 + # XTML structural tag (structural_tag_model = "kimi_k3"). We return + # early to skip the generic parent path, which would otherwise attach + # JSON guided decoding when strict calling is off -- that JSON + # constraint conflicts with the <|open|>tools<|sep|> channel. + request.skip_special_tokens = False + if hasattr(request, "spaces_between_special_tokens"): + request.spaces_between_special_tokens = False + return request + + request = super().adjust_request(request) + # The XTML markers (<|open|>/<|close|>/<|sep|>, + # <|open|>response<|sep|> ...) must + # reach this parser as CONTIGUOUS literal text. Two request flags govern + # that, and Aphrodite's detokenizer couples them: + # effective_spaces_between = + # skip_special_tokens OR spaces_between_special_tokens + # (aphrodite/v1/engine/detokenizer.py). + # The detokenizer treats these control tokens as separate sub-texts and, + # when the effective flag is True, joins them with spaces -> + # "<|open|> response <|sep|>", which the regexes below would NOT match. + # Forcing BOTH flags off is the only way to suppress that spacing. We + # set both unconditionally: the response channel is unwrapped from these + # markers even with no tools. + request.skip_special_tokens = False + if hasattr(request, "spaces_between_special_tokens"): + request.spaces_between_special_tokens = False + return request + + def _attrs(self, s: str) -> dict[str, str]: + return {m["k"]: m["v"].replace(""", '"').replace("&", "&") for m in self._attr_re.finditer(s)} + + def _decode_call(self, attrs: str, body: str) -> ToolCall | None: + """Decode one ``call`` block into a :class:`ToolCall`. + + ``attrs`` is the text between ``<|open|>call `` and ``<|sep|>`` (carries + ``tool=`` / ``index=``); ``body`` is the sequence of ``argument`` blocks. + Each argument is re-typed per its ``type=`` tag: strings pass through + raw, everything else is JSON-decoded (falling back to raw text if the + JSON is malformed, so a partial stream never raises). Returns ``None`` + when no tool name is present (an empty/garbage block is dropped). + """ + call_attrs = self._attrs(attrs) + tool_name = call_attrs.get("tool", "") + tool_index = call_attrs.get("index", "") + arguments: dict = {} + for arg_match in self._arg_re.finditer(body): + arg_attrs = self._attrs(arg_match["attrs"]) + key = arg_attrs.get("key", "") + arg_type = arg_attrs.get("type", "string") + raw_value = arg_match["val"] + if arg_type == "string": + arguments[key] = raw_value + else: + try: + arguments[key] = json.loads(raw_value) + except json.JSONDecodeError: + arguments[key] = raw_value + if not tool_name: + return None + tool_call_id = tool_name + if tool_index: + try: + tool_call_id = f"{tool_name}:{int(tool_index) - 1}" + except ValueError: + tool_call_id = f"{tool_name}:{tool_index}" + # id uses the API-side zero-based call ordinal; XTML's message index + # stays one-based when rendering tool result messages. + return ToolCall( + id=tool_call_id, + type="function", + function=FunctionCall( + name=tool_name, + arguments=json.dumps(arguments, ensure_ascii=False), + ), + ) + + def _strip_response_content(self, text: str) -> str | None: + """Strip XTML response/message markers from generated response text. + + In chat serving, ``<|open|>response<|sep|>`` is often part of the prompt + generation prefix, so the model output may only contain the body plus + ``<|close|>response<|sep|>``. Handle both that consumed-prefix shape and a + complete ``<|open|>response<|sep|>... <|close|>response<|sep|>`` wrapper. + """ + m_open = self._response_open_re.search(text) + if m_open is not None: + m_close = self._response_close_re.search(text, m_open.end()) + if m_close is not None: + text = text[m_open.end() : m_close.start()] + else: + text = text[m_open.end() :] + else: + text = self._response_close_re.sub("", text) + text = self._message_close_re.sub("", text) + return text or None + + def _content(self, model_output: str, before: str) -> str | None: + # prefer the unwrapped response channel; else the text before the tools + m = self._response_re.search(model_output) + if m is not None: + return m["c"] or None + return self._strip_response_content(before) + + def _extract_response_content(self, current_text: str) -> str | None: + # Streaming response text is computed from the accumulated text. This is + # what keeps split markers from leaking: + # <|open|> / response / <|sep|>Hi -> emit only "Hi" after open closes + # Hi<|open|> / tools / <|sep|>... -> emit "Hi", hold the tools marker + # Hi<|close|> / response / <|sep|>... -> emit "Hi", hold the close marker + # Tool calls are even simpler: they are not emitted until a full + # <|close|>call<|sep|> is present in current_text. + m_open = self._response_open_re.search(current_text) + # In the normal chat path, the response-open marker may be consumed as + # the generation prefix. Then generated text starts directly with the + # response body or with <|close|>response<|sep|> before a tools channel. + body_start = m_open.end() if m_open is not None else 0 + m_tools = self._tools_open_re.search(current_text, body_start) + m_rclose = self._response_close_re.search(current_text, body_start) + tools_start = m_tools.start() if m_tools else -1 + response_end = m_rclose.start() if m_rclose else -1 + + candidates = [i for i in (tools_start, response_end) if i != -1] + if candidates: + sendable_idx = min(candidates) + else: + overlap = max( + _partial_tag_overlap(current_text, self.response_open), + _partial_tag_overlap(current_text, self.response_close), + _partial_tag_overlap(current_text, self.tools_open), + ) + sendable_idx = len(current_text) - overlap + + if sendable_idx <= body_start: + return None + if self._sent_content_idx < body_start: + self._sent_content_idx = body_start + if sendable_idx <= self._sent_content_idx: + return None + + content = current_text[self._sent_content_idx : sendable_idx] + self._sent_content_idx = sendable_idx + return content or None + + def extract_tool_calls(self, model_output: str, request: ChatCompletionRequest) -> ExtractedToolCallInformation: + m_open = self._tools_open_re.search(model_output) + if m_open is None: + # no tools channel -> content is the response channel (unwrapped) + return ExtractedToolCallInformation( + tools_called=False, + tool_calls=[], + content=self._content(model_output, model_output), + ) + try: + before = model_output[: m_open.start()] + start = m_open.end() + m_close = self._tools_close_re.search(model_output, start) + section = model_output[start:] if m_close is None else model_output[start : m_close.start()] + + tool_calls = [ + tc + for m in self._call_re.finditer(section) + if (tc := self._decode_call(m["attrs"], m["body"])) is not None + ] + if not tool_calls: + return ExtractedToolCallInformation( + tools_called=False, + tool_calls=[], + content=self._content(model_output, before), + ) + return ExtractedToolCallInformation( + tools_called=True, + tool_calls=tool_calls, + content=self._content(model_output, before), + ) + except Exception: + logger.exception("Error extracting K3 tool calls.") + return ExtractedToolCallInformation(tools_called=False, tool_calls=[], content=model_output) + + def extract_tool_calls_streaming( + self, + previous_text: str, + current_text: str, + delta_text: str, + previous_token_ids: Sequence[int], + current_token_ids: Sequence[int], + delta_token_ids: Sequence[int], + request: ChatCompletionRequest, + ) -> DeltaMessage | None: + # Conservative streaming: stream unwrapped response-channel text, then + # buffer tool calls and emit each call once its block closes. + content = self._extract_response_content(current_text) + + # tools channel is open: parse fully-closed calls we have not emitted yet + m_tools = self._tools_open_re.search(current_text) + if m_tools is None: + return DeltaMessage(content=content) if content else None + + section = current_text[m_tools.end() :] + calls = [ + tc for m in self._call_re.finditer(section) if (tc := self._decode_call(m["attrs"], m["body"])) is not None + ] + if len(calls) <= self._sent_tool_call_count: + return DeltaMessage(content=content) if content else None + new = calls[self._sent_tool_call_count :] + + deltas = [ + DeltaToolCall( + index=self._sent_tool_call_count + i, + id=tc.id, + type="function", + function=DeltaFunctionCall(name=tc.function.name, arguments=tc.function.arguments).model_dump( + exclude_none=True + ), + ) + for i, tc in enumerate(new) + ] + self._sent_tool_call_count = len(calls) + return DeltaMessage(content=content, tool_calls=deltas) diff --git a/aphrodite/tool_parsers/structural_tag_registry.py b/aphrodite/tool_parsers/structural_tag_registry.py index cc1e7d5b7f..2f518d3200 100644 --- a/aphrodite/tool_parsers/structural_tag_registry.py +++ b/aphrodite/tool_parsers/structural_tag_registry.py @@ -19,7 +19,12 @@ AnyTextFormat, ConstStringFormat, JSONSchemaFormat, + OptionalFormat, + OrFormat, + PlusFormat, + RegexFormat, SequenceFormat, + StarFormat, TagFormat, TagsWithSeparatorFormat, TriggeredTagsFormat, @@ -63,7 +68,7 @@ "deepseek_v4", } ) -APHRODITE_BUILTIN_STRUCTURAL_TAG_MODELS = frozenset({"hermes"}) +APHRODITE_BUILTIN_STRUCTURAL_TAG_MODELS = frozenset({"hermes", "kimi_k3"}) SUPPORTED_STRUCTURAL_TAG_MODELS = XGRAMMAR_BUILTIN_STRUCTURAL_TAG_MODELS | APHRODITE_BUILTIN_STRUCTURAL_TAG_MODELS _APHRODITE_STRUCTURAL_TAG_REGISTRY: dict[str, StructuralTagBuilder] = {} @@ -327,3 +332,249 @@ def get_minimax_structural_tag( ) return StructuralTag(format=suffix_tag) + + +# --------------------------------------------------------------------------- +# Kimi K3 (XTML channel format) +# --------------------------------------------------------------------------- +# K3 assistant output after the reasoning gate (``<|close|>think<|sep|>``): +# <|open|>response<|sep|> <|close|>response<|sep|> +# [ <|open|>tools<|sep|> +# <|open|>call tool="NAME" index="1"<|sep|> +# <|open|>argument key="K" type="TYPE"<|sep|>VALUE<|close|>argument<|sep|> +# <|close|>call<|sep|> ... +# <|close|>tools<|sep|> ] +# See ``encoding_k3.py`` (_render_assistant_segments / _open_tag / _attr) and +# ``kimi_k3_tool_parser.py`` for the exact byte-level encoding this mirrors. +_K3_OPEN = "<|open|>" +_K3_CLOSE = "<|close|>" +_K3_SEP = "<|sep|>" +_K3_RESPONSE_OPEN = f"{_K3_OPEN}response{_K3_SEP}" +_K3_RESPONSE_CLOSE = f"{_K3_CLOSE}response{_K3_SEP}" +_K3_TOOLS_OPEN = f"{_K3_OPEN}tools{_K3_SEP}" +_K3_TOOLS_CLOSE = f"{_K3_CLOSE}tools{_K3_SEP}" +_K3_CALL_CLOSE = f"{_K3_CLOSE}call{_K3_SEP}" +_K3_ARG_CLOSE = f"{_K3_CLOSE}argument{_K3_SEP}" +# The model closes the assistant turn with <|close|>message<|sep|> right before +# the end-of-message token. It is generated (not part of the prompt prefix), so +# the tag must permit it or the FSM would mask the model's natural terminator. +_K3_MESSAGE_CLOSE = f"{_K3_CLOSE}message{_K3_SEP}" + +# JSON-schema type -> K3 XTML ``type=`` attribute value. Mirrors +# ``encoding_k3._xtml_type`` (integer collapses onto number). +_K3_JSON_TO_XTML_TYPE = { + "string": "string", + "integer": "number", + "number": "number", + "boolean": "boolean", + "null": "null", + "object": "object", + "array": "array", +} + + +def _k3_escape_attr(value: str) -> str: + """Mirror ``encoding_k3._escape_attr_value`` (``&`` then ``"``).""" + return str(value).replace("&", "&").replace('"', """) + + +_K3_STRING_ATOM = r"(?:[^<]|<[^|])" +"""One raw-string character: anything but the ambiguous "<|" marker prefix. +Allows '<' inside values (e.g. HTML snippets); a value *ending* in '<' or +containing a literal "<|" is not expressible and falls back to AnyText via +the pattern checks below never matching those cases at build time (schemas +cannot know values, so the only build-time effect is the length bound).""" + + +def _k3_bounded_string_regex(prop: dict[str, Any]) -> str | None: + """Length/pattern constraint for the raw string channel, if expressible. + + The XTML string channel emits values raw (not JSON-quoted), so + JSONSchemaFormat cannot enforce string constraints there; unconstrained + AnyText lets maxLength/pattern violations through (observed on the walle + verifier: over-long junk strings pass the grammar and fail validation). + xgrammar's regex engine has no lookahead, so the close marker is kept + unambiguous by excluding the "<|" prefix from value characters. + + Returns a regex for the value, or None to keep permissive AnyText. + """ + max_len = prop.get("maxLength") + min_len = prop.get("minLength", 0) + if not isinstance(max_len, int) or max_len < 0 or max_len > 4096: + return None + if not isinstance(min_len, int) or min_len < 0 or min_len > max_len: + min_len = 0 + return _K3_STRING_ATOM + f"{{{min_len},{max_len}}}" + + +def _k3_argument_tag( + key: str, + schema: dict[str, Any], + root_defs: dict[str, Any] | None = None, +) -> TagFormat | None: + """Build one ``argument`` XTML tag for property ``key``. + + ``string`` values are emitted raw (bounded by the close marker); every other + JSON type is emitted as JSON and validated against the property schema. A + property whose type is a union / missing is left permissive (any XTML type, + raw value) so a valid call is never rejected. + + ``root_defs`` carries the tool parameters' root-level ``$defs`` / + ``definitions``: slicing a property out of the parameters document orphans + its ``#/$defs/...`` references, so those tables must be re-attached to keep + the embedded schema self-contained. + """ + prop = schema if isinstance(schema, dict) else {} + json_type = prop.get("type") + xtml_type = _K3_JSON_TO_XTML_TYPE.get(json_type) if isinstance(json_type, str) else None + if xtml_type is None: + # Unknown / union type: constrain the key but keep the value permissive. + return None + begin = f'{_K3_OPEN}argument key="{_k3_escape_attr(key)}" type="{xtml_type}"{_K3_SEP}' + if xtml_type == "string": + # Raw string channel: JSONSchemaFormat can't apply (values are not + # JSON-quoted), but an enum/const of strings is a finite set that can + # be enforced exactly with const-string alternation. Enum semantics + # are exclusive, so this never over-rejects. Fall back to permissive + # AnyText for open-ended strings or non-representable enums. + enum_values = prop.get("enum") + if enum_values is None and isinstance(prop.get("const"), str): + enum_values = [prop["const"]] + if ( + isinstance(enum_values, list) + and enum_values + and len(enum_values) <= 256 + and all(isinstance(v, str) for v in enum_values) + and not any("<|" in v for v in enum_values) + ): + branches = [ConstStringFormat(value=v) for v in enum_values] + content: Any = branches[0] if len(branches) == 1 else OrFormat(elements=branches) + elif (bounded := _k3_bounded_string_regex(prop)) is not None: + content = RegexFormat(pattern=bounded) + else: + content = AnyTextFormat(excludes=[_K3_CLOSE]) + else: + embedded = prop + if root_defs: + embedded = dict(prop) + for defs_key, defs_value in root_defs.items(): + embedded.setdefault(defs_key, defs_value) + content = JSONSchemaFormat(json_schema=embedded) + return TagFormat(begin=begin, content=content, end=_K3_ARG_CLOSE) + + +def _k3_permissive_argument_tag() -> TagFormat: + """A key/type-agnostic ``argument`` tag: any attributes, raw value. + + Used as a fallback so tools with union/loose schemas still get the XTML + skeleton constrained without over-rejecting the value. + """ + return TagFormat( + begin=_K3_OPEN + "argument ", + content=SequenceFormat( + elements=[ + RegexFormat(pattern=r"[^<]*" + _K3_SEP.replace("|", r"\|")), + AnyTextFormat(excludes=[_K3_CLOSE]), + ] + ), + end=_K3_ARG_CLOSE, + ) + + +def _k3_arguments_block(parameters: dict[str, Any] | bool) -> Any: + """Build ``argument`` tags for a tool's parameter schema. + + Require at least one tag when the root schema declares required properties. + Otherwise, keep accepting zero-or-more tags. Arguments remain order-agnostic + and non-unique. + """ + if not isinstance(parameters, dict): + return StarFormat(content=_k3_permissive_argument_tag()) + props = parameters.get("properties") + if not isinstance(props, dict) or not props: + # No declared properties: allow any argument blocks (or none). + return StarFormat(content=_k3_permissive_argument_tag()) + root_defs = { + defs_key: parameters[defs_key] + for defs_key in ("$defs", "definitions") + if isinstance(parameters.get(defs_key), dict) + } + tags: list[TagFormat] = [] + for key, prop in props.items(): + tag = _k3_argument_tag(key, prop, root_defs) + tags.append(tag if tag is not None else _k3_permissive_argument_tag()) + inner = tags[0] if len(tags) == 1 else OrFormat(elements=list(tags)) + required = parameters.get("required") + if isinstance(required, list) and required: + return PlusFormat(content=inner) + return StarFormat(content=inner) + + +def _k3_call_tag(tool: FunctionToolParam) -> TagFormat: + """One ``call`` tag: ``<|open|>call tool="N" index=""<|sep|> args``.""" + function = tool.function + parameters = _get_function_parameters(function) + begin = f'{_K3_OPEN}call tool="{_k3_escape_attr(function.name)}" index="' + return TagFormat( + begin=begin, + content=SequenceFormat( + elements=[ + RegexFormat(pattern=r"[0-9]+"), + ConstStringFormat(value=f'"{_K3_SEP}'), + _k3_arguments_block(parameters), + ] + ), + end=_K3_CALL_CLOSE, + ) + + +def _k3_response_prefix() -> list[Any]: + """The response channel that always precedes the tools channel. + + ``response`` is generated in thinking mode (prefix ends at + ``<|open|>think<|sep|>``) but is part of the generation prefix in + non-thinking mode, so its open marker is optional. The body is bounded by + the response close marker. + """ + return [ + OptionalFormat(content=ConstStringFormat(value=_K3_RESPONSE_OPEN)), + TagFormat(begin="", content=AnyTextFormat(), end=_K3_RESPONSE_CLOSE), + ] + + +def _k3_tools_channel(tools: list[FunctionToolParam]) -> TagFormat: + return TagFormat( + begin=_K3_TOOLS_OPEN, + content=TagsWithSeparatorFormat( + tags=[_k3_call_tag(tool) for tool in tools], + separator="", + at_least_one=True, + ), + end=_K3_TOOLS_CLOSE, + ) + + +@register_aphrodite_structural_tag("kimi_k3") +def get_kimi_k3_structural_tag( + tools: list[FunctionToolParam], + builtin_tools: list[BuiltinToolParam], + tool_choice: SimplifiedToolChoice, + reasoning: bool, +) -> StructuralTag: + del builtin_tools, reasoning + + trailer = OptionalFormat(content=ConstStringFormat(value=_K3_MESSAGE_CLOSE)) + + if not tools: + return StructuralTag(format=SequenceFormat(elements=[*_k3_response_prefix(), trailer])) + + if tool_choice == "auto": + tools_part: Any = OptionalFormat(content=_k3_tools_channel(tools)) + elif tool_choice == "forced": + # K3 rejects named tool choice upstream; treat defensively as a single + # mandatory call of the first tool. + tools_part = _k3_tools_channel(tools[:1]) + else: # required + tools_part = _k3_tools_channel(tools) + + return StructuralTag(format=SequenceFormat(elements=[*_k3_response_prefix(), tools_part, trailer])) diff --git a/aphrodite/transformers_utils/config.py b/aphrodite/transformers_utils/config.py index 5d6fa86ae6..a062dbe50a 100644 --- a/aphrodite/transformers_utils/config.py +++ b/aphrodite/transformers_utils/config.py @@ -124,7 +124,9 @@ def __getitem__(self, key): qwen3_asr="Qwen3ASRConfig", qwen3_next="Qwen3NextConfig", qwen3_5="Qwen3_5Config", + qwen3_5_text="Qwen3_5TextConfig", qwen3_5_moe="Qwen3_5MoeConfig", + qwen3_5_moe_text="Qwen3_5MoeTextConfig", laguna="LagunaConfig", lfm2_moe="Lfm2MoeConfig", **{"unlimited-ocr": "UnlimitedOCRConfig"}, diff --git a/aphrodite/transformers_utils/configs/kimi_k3.py b/aphrodite/transformers_utils/configs/kimi_k3.py new file mode 100644 index 0000000000..5b2266d6dd --- /dev/null +++ b/aphrodite/transformers_utils/configs/kimi_k3.py @@ -0,0 +1,136 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Kimi-K3 multimodal configuration.""" + +from transformers.configuration_utils import PretrainedConfig + +from aphrodite.logger import init_logger +from aphrodite.transformers_utils.configs.kimi_linear import KimiLinearConfig + +logger = init_logger(__name__) + + +class KimiK3VisionConfig(PretrainedConfig): + model_type = "kimi_k3_vision" + + def __init__( + self, + patch_size: int = 14, + init_pos_emb_height: int = 64, + init_pos_emb_width: int = 64, + init_pos_emb_time: int = 4, + pos_emb_type: str = "divided_fixed", + vt_num_attention_heads: int = 12, + vt_num_hidden_layers: int = 27, + vt_hidden_size: int = 1024, + vt_intermediate_size: int = 4096, + merge_kernel_size: tuple[int, int] = (2, 2), + video_attn_type: str = "spatial_temporal", + merge_type: str = "sd2_tpool", + _attn_implementation: str = "flash_attention_2", + mm_projector_type: str = "patchmergerv2", + mm_hidden_size: int | None = None, + projector_hidden_act: str = "gelu", + projector_ln_eps: float = 1e-5, + qkv_hidden_size: int = 1536, + norm_type: str = "rmsnorm", + attn_bias: bool = False, + patch_embed_proj_bias: bool = False, + mlp_type: str = "mlp2", + linear_bias: bool = False, + activation_func: str = "gelu_pytorch_tanh", + pos_emb_interpolation_mode: str = "bilinear", + text_hidden_size: int = 2304, + **kwargs, + ): + super().__init__(**kwargs) + + self.patch_size = patch_size + self.init_pos_emb_height = init_pos_emb_height + self.init_pos_emb_width = init_pos_emb_width + self.init_pos_emb_time = init_pos_emb_time + self.pos_emb_type = pos_emb_type + self.vt_num_attention_heads = vt_num_attention_heads + self.vt_num_hidden_layers = vt_num_hidden_layers + self.vt_hidden_size = vt_hidden_size + self.vt_intermediate_size = vt_intermediate_size + self.merge_kernel_size = tuple(merge_kernel_size) + self.video_attn_type = video_attn_type + self.merge_type = merge_type + self._attn_implementation = _attn_implementation + + self.mm_projector_type = mm_projector_type + self.mm_hidden_size = mm_hidden_size if mm_hidden_size is not None else vt_hidden_size + self.projector_hidden_act = projector_hidden_act + self.projector_ln_eps = projector_ln_eps + self.text_hidden_size = text_hidden_size + + self.qkv_hidden_size = qkv_hidden_size + self.norm_type = norm_type + self.attn_bias = attn_bias + self.patch_embed_proj_bias = patch_embed_proj_bias + self.mlp_type = mlp_type + self.linear_bias = linear_bias + self.activation_func = activation_func + self.pos_emb_interpolation_mode = pos_emb_interpolation_mode + + # Aliases consumed by the Kimi-K2.5 vision implementation. + self.num_attention_heads = vt_num_attention_heads + self.num_hidden_layers = vt_num_hidden_layers + self.hidden_size = vt_hidden_size + self.intermediate_size = vt_intermediate_size + + +class KimiK3Config(PretrainedConfig): + model_type = "kimi_k3" + + def __init__( + self, + text_config: dict | KimiLinearConfig | None = None, + vision_config: dict | KimiK3VisionConfig | None = None, + ignore_index: int = -100, + media_placeholder_token_id: int = 163605, + pad_token_id: int = 0, + image_placeholder: str = "<|kimi_image_placeholder|>", + **kwargs, + ): + if text_config is None: + self.text_config = KimiLinearConfig() + elif isinstance(text_config, dict): + self.text_config = KimiLinearConfig(**text_config) + else: + self.text_config = text_config + + if vision_config is None: + self.vision_config = KimiK3VisionConfig() + elif isinstance(vision_config, dict): + self.vision_config = KimiK3VisionConfig(**vision_config) + else: + self.vision_config = vision_config + + # K3's vision projector output must match the text model's hidden size. + # Override any value provided in vision_config for safety. + if self.vision_config.text_hidden_size != self.text_config.hidden_size: + logger.info( + "Overriding vision_config.text_hidden_size from %s to %s to match text_config.hidden_size", + self.vision_config.text_hidden_size, + self.text_config.hidden_size, + ) + self.vision_config.text_hidden_size = self.text_config.hidden_size + + self.ignore_index = ignore_index + self.media_placeholder_token_id = media_placeholder_token_id + self.image_placeholder = image_placeholder + + if getattr(self.text_config, "quantization_config", None) is not None: + self.quantization_config = self.text_config.quantization_config + + super().__init__(pad_token_id=pad_token_id, **kwargs) + + @property + def hidden_size(self) -> int: + return self.text_config.hidden_size + + @property + def vocab_size(self) -> int: + return self.text_config.vocab_size diff --git a/aphrodite/transformers_utils/processors/kimi_k3.py b/aphrodite/transformers_utils/processors/kimi_k3.py new file mode 100644 index 0000000000..9a05b2a83c --- /dev/null +++ b/aphrodite/transformers_utils/processors/kimi_k3.py @@ -0,0 +1,61 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from transformers import BaseImageProcessor, BatchFeature, TensorType +from transformers.processing_utils import ProcessorMixin + +from aphrodite.tokenizers.hf import HfTokenizer + + +class KimiK3Processor(ProcessorMixin): + """HF-style processor wrapper for the image-only Kimi-K3 model. + + K3 exposes the standard ``image`` modality, so vLLM calls this processor + with ``images=[PIL, ...]``. The underlying checkpoint image processor + (``KimiK3VisionProcessor``) works on ``{"type": "image", "image": PIL}`` + media dicts, so this wrapper adapts bare PIL images into that shape before + delegating to ``preprocess``. + + Text is only tokenized here; the single ``<|kimi_image_placeholder|>`` + token per image is expanded into the resolution-aware media block by the + model's ``_get_prompt_updates`` on the vLLM side. + """ + + attributes = ["image_processor", "tokenizer"] + + def __init__( + self, + image_processor: BaseImageProcessor, + tokenizer: HfTokenizer, + ) -> None: + self.image_processor = image_processor + self.tokenizer = tokenizer + + def __call__( + self, + text: str | list[str] | None = None, + images: object | list[object] | None = None, + return_tensors: str | TensorType | None = None, + **kwargs, + ) -> BatchFeature: + if images is not None: + if not isinstance(images, list): + images = [images] + medias = [{"type": "image", "image": image} for image in images] + mm_inputs = self.image_processor.preprocess( + medias, + return_tensors=return_tensors, + ) + else: + mm_inputs = {} + + if text is not None: + if not isinstance(text, list): + text = [text] + text_inputs = self.tokenizer(text) + else: + text_inputs = {} + + return BatchFeature( + data={**text_inputs, **mm_inputs}, + tensor_type=return_tensors, + ) diff --git a/aphrodite/v1/attention/backends/cpu_attn.py b/aphrodite/v1/attention/backends/cpu_attn.py index c2314d8792..ad4371f7b6 100644 --- a/aphrodite/v1/attention/backends/cpu_attn.py +++ b/aphrodite/v1/attention/backends/cpu_attn.py @@ -153,12 +153,12 @@ def __init__( if self.window_size is None: self.window_size = -1 self.block_size = aphrodite_config.cache_config.block_size - kv_cache_dtype_str = aphrodite_config.cache_config.cache_dtype + self.kv_cache_dtype = aphrodite_config.cache_config.cache_dtype self.isa = _get_attn_isa( self.dtype, self.block_size, self.head_dim, - kv_cache_dtype_str, + self.kv_cache_dtype, ) self.is_cross_attention = isinstance(kv_cache_spec, CrossAttentionSpec) self.is_encoder_only_attention = isinstance(kv_cache_spec, EncoderOnlyAttentionSpec) @@ -224,6 +224,7 @@ def build( isa=self.isa, enable_kv_split=envs.APHRODITE_CPU_ATTN_SPLIT_KV, dynamic_causal=dynamic_casual, + kv_cache_dtype=self.kv_cache_dtype, ) attn_metadata = CPUAttentionMetadata( diff --git a/aphrodite/v1/attention/ops/rocm_aiter_mla_sparse.py b/aphrodite/v1/attention/ops/rocm_aiter_mla_sparse.py index f7b3bd423b..356ac6b6b5 100644 --- a/aphrodite/v1/attention/ops/rocm_aiter_mla_sparse.py +++ b/aphrodite/v1/attention/ops/rocm_aiter_mla_sparse.py @@ -412,6 +412,7 @@ def rocm_fp8_paged_mqa_logits( KVBlockSize=block_size, WavePerEU=2, ) + out_logits.nan_to_num_(float("-inf")) return out_logits deepgemm_fp8_paged_mqa_logits_stage1 = aiter_paged_mqa_logits_module.deepgemm_fp8_paged_mqa_logits_stage1 batch_size, next_n, heads, _ = q_fp8.shape diff --git a/aphrodite/v1/attention/ops/triton_merge_attn_states.py b/aphrodite/v1/attention/ops/triton_merge_attn_states.py index 3fc8c66774..0965e7b482 100644 --- a/aphrodite/v1/attention/ops/triton_merge_attn_states.py +++ b/aphrodite/v1/attention/ops/triton_merge_attn_states.py @@ -135,6 +135,15 @@ def merge_attn_states( # attention backend. prefix_head_stride = prefix_output.stride(1) output_head_stride = output.stride(1) + # lse tensors are [NUM_HEADS, NUM_TOKENS] but may be non-contiguous views + # (e.g. a transpose of a backend's [NUM_TOKENS, NUM_HEADS] output), so index + # them by their actual strides rather than assuming a contiguous layout. + prefix_lse_head_stride = prefix_lse.stride(0) + prefix_lse_token_stride = prefix_lse.stride(1) + suffix_lse_head_stride = suffix_lse.stride(0) + suffix_lse_token_stride = suffix_lse.stride(1) + output_lse_head_stride = output_lse.stride(0) if output_lse is not None else 0 + output_lse_token_stride = output_lse.stride(1) if output_lse is not None else 0 # If prefill_tokens_with_context is None, all tokens should use prefix context if prefill_tokens_with_context is None: @@ -150,6 +159,12 @@ def merge_attn_states( suffix_lse, prefix_head_stride, output_head_stride, + prefix_lse_head_stride, + prefix_lse_token_stride, + suffix_lse_head_stride, + suffix_lse_token_stride, + output_lse_head_stride, + output_lse_token_stride, output_scale, prefill_tokens_with_context, head_size, @@ -169,6 +184,12 @@ def merge_attn_states_kernel( suffix_lse, # [NUM_HEADS, NUM_TOKENS] prefix_head_stride, output_head_stride, + prefix_lse_head_stride, + prefix_lse_token_stride, + suffix_lse_head_stride, + suffix_lse_token_stride, + output_lse_head_stride, + output_lse_token_stride, output_scale, # scale tensor or None prefill_tokens_with_context, HEAD_SIZE: tl.constexpr, @@ -179,7 +200,6 @@ def merge_attn_states_kernel( FP8_MAX: tl.constexpr = float8_info.max, ): token_idx = tl.program_id(0) - num_tokens = tl.num_programs(0) head_idx = tl.program_id(1) num_heads = tl.num_programs(1) @@ -191,9 +211,12 @@ def merge_attn_states_kernel( # For tokens without context (token_idx >= prefill_tokens_with_context), # directly copy from suffix_output if not prefix_mask: - s_lse = tl.load(suffix_lse + head_idx * num_tokens + token_idx) + s_lse = tl.load(suffix_lse + head_idx * suffix_lse_head_stride + token_idx * suffix_lse_token_stride) if OUTPUT_LSE: - tl.store(output_lse + head_idx * num_tokens + token_idx, s_lse) + tl.store( + output_lse + head_idx * output_lse_head_stride + token_idx * output_lse_token_stride, + s_lse, + ) s_out = tl.load( suffix_output + token_idx * num_heads * prefix_head_stride + head_idx * prefix_head_stride + head_arange, @@ -214,8 +237,8 @@ def merge_attn_states_kernel( # For tokens with context (token_idx < prefill_tokens_with_context), # perform normal merge operation - p_lse = tl.load(prefix_lse + head_idx * num_tokens + token_idx) - s_lse = tl.load(suffix_lse + head_idx * num_tokens + token_idx) + p_lse = tl.load(prefix_lse + head_idx * prefix_lse_head_stride + token_idx * prefix_lse_token_stride) + s_lse = tl.load(suffix_lse + head_idx * suffix_lse_head_stride + token_idx * suffix_lse_token_stride) # FA2 and FA3 have different behavior for when the sum-exp is 0, this namely # arises with 0 len seqlens. FA3 returns -inf here while FA2 returns inf. @@ -236,7 +259,10 @@ def merge_attn_states_kernel( if OUTPUT_LSE: out_lse = tl.log(out_se) + max_lse out_lse = tl.where(max_lse == float("-inf"), float("-inf"), out_lse) - tl.store(output_lse + head_idx * num_tokens + token_idx, out_lse) + tl.store( + output_lse + head_idx * output_lse_head_stride + token_idx * output_lse_token_stride, + out_lse, + ) p_out = tl.load( prefix_output + token_idx * num_heads * prefix_head_stride + head_idx * prefix_head_stride + head_arange, diff --git a/aphrodite/v1/core/sched/scheduler.py b/aphrodite/v1/core/sched/scheduler.py index a2e0c8d33e..74f6252324 100644 --- a/aphrodite/v1/core/sched/scheduler.py +++ b/aphrodite/v1/core/sched/scheduler.py @@ -2120,6 +2120,7 @@ def has_requests(self) -> bool: self.has_unfinished_requests() or self.has_finished_requests() or (self.connector is not None and self.connector.has_pending_push_work()) + or (self.ec_connector is not None and self.ec_connector.has_pending_push_work()) ) def reset_prefix_cache(self, reset_running_requests: bool = False, reset_connector: bool = False) -> bool: diff --git a/aphrodite/v1/engine/async_llm.py b/aphrodite/v1/engine/async_llm.py index 289790d598..fcf7298138 100644 --- a/aphrodite/v1/engine/async_llm.py +++ b/aphrodite/v1/engine/async_llm.py @@ -21,6 +21,7 @@ from aphrodite.engine.arg_utils import AsyncEngineArgs from aphrodite.engine.protocol import EngineClient, StreamingInput from aphrodite.entrypoints.serve.elastic_ep.middleware import set_scaling_elastic_ep +from aphrodite.exceptions import APHRODITEClientError, APHRODITEValidationError from aphrodite.inputs import EngineInput, PromptType from aphrodite.logger import init_logger from aphrodite.lora.request import LoRARequest @@ -302,7 +303,7 @@ async def add_request( is_pooling = isinstance(params, PoolingParams) if self.aphrodite_config.cache_config.kv_sharing_fast_prefill and not is_pooling and params.prompt_logprobs: - raise ValueError( + raise APHRODITEValidationError( "--kv-sharing-fast-prefill produces incorrect logprobs for " "prompt tokens, please disable it when the requests need " "prompt logprobs" @@ -467,7 +468,7 @@ async def handle_inputs(): ) req.external_req_id = request_id if req.prompt_embeds is not None: - raise ValueError("prompt_embeds not supported for streaming inputs") + raise APHRODITEValidationError("prompt_embeds not supported for streaming inputs") prompt_text, _, _ = extract_prompt_components(self.model_config, input_chunk.prompt) await self._add_request(req, prompt_text, None, 0, queue) except (asyncio.CancelledError, GeneratorExit): @@ -499,7 +500,7 @@ def _validate_streaming_input_sampling_params( or params.output_kind == RequestOutputKind.FINAL_ONLY or params.stop ): - raise ValueError( + raise APHRODITEValidationError( "Input streaming not currently supported " "for pooling models, n > 1, request_kind = FINAL_ONLY " "or with stop strings." @@ -588,7 +589,7 @@ async def generate( raise # Request validation error. - except ValueError as e: + except APHRODITEClientError as e: if self.log_requests: logger.info("Request %s failed (bad request): %s.", request_id, e) raise @@ -841,7 +842,7 @@ async def encode( raise # Request validation error. - except ValueError: + except APHRODITEClientError: if self.log_requests: logger.info("Request %s failed (bad request).", request_id) raise diff --git a/aphrodite/v1/engine/exceptions.py b/aphrodite/v1/engine/exceptions.py index d9f79a019e..6c318241cd 100644 --- a/aphrodite/v1/engine/exceptions.py +++ b/aphrodite/v1/engine/exceptions.py @@ -1,12 +1,15 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -class EngineGenerateError(Exception): +from aphrodite.exceptions import APHRODITEServerError + + +class EngineGenerateError(APHRODITEServerError): """Raised when a AsyncLLM.generate() fails. Recoverable.""" pass -class EngineDeadError(Exception): +class EngineDeadError(APHRODITEServerError): """Raised when the EngineCore dies. Unrecoverable.""" def __init__(self, *args, suppress_context: bool = False, **kwargs): diff --git a/aphrodite/v1/engine/input_processor.py b/aphrodite/v1/engine/input_processor.py index 6c748a789f..a22ffb50d9 100644 --- a/aphrodite/v1/engine/input_processor.py +++ b/aphrodite/v1/engine/input_processor.py @@ -7,6 +7,7 @@ import aphrodite.envs as envs from aphrodite.config import AphroditeConfig +from aphrodite.exceptions import APHRODITEValidationError from aphrodite.inputs import ( EngineInput, PromptType, @@ -86,7 +87,7 @@ def _validate_params( if isinstance(params, SamplingParams): supported_generation_tasks = [task for task in supported_tasks if task in GENERATION_TASKS] if not supported_generation_tasks: - raise ValueError("This model does not support generation") + raise APHRODITEValidationError("This model does not support generation") params.verify( self.model_config, @@ -97,13 +98,13 @@ def _validate_params( if params.thinking_token_budget is not None: if self.aphrodite_config.reasoning_config is None or not self.aphrodite_config.reasoning_config.enabled: - raise ValueError( + raise APHRODITEValidationError( "thinking_token_budget is set but reasoning_config is " "not configured. Please set --reasoning-parser " "and/or --reasoning-config to use thinking_token_budget." ) if self.use_v2_model_runner: - raise ValueError( + raise APHRODITEValidationError( "thinking_token_budget is not yet supported by the V2 " "model runner. Run Aphrodite with APHRODITE_USE_V2_MODEL_RUNNER=0 " "to use thinking_token_budget." @@ -111,7 +112,7 @@ def _validate_params( elif isinstance(params, PoolingParams): supported_pooling_tasks = [task for task in supported_tasks if task in POOLING_TASKS] if not supported_pooling_tasks: - raise ValueError("This model does not support pooling") + raise APHRODITEValidationError("This model does not support pooling") if params.task is None: if "token_embed" in supported_pooling_tasks: @@ -122,7 +123,9 @@ def _validate_params( params.task = "plugin" if params.task not in supported_pooling_tasks: - raise ValueError(f"Unsupported task: {params.task!r} Supported tasks: {supported_pooling_tasks}") + raise APHRODITEValidationError( + f"Unsupported task: {params.task!r} Supported tasks: {supported_pooling_tasks}" + ) params.verify(self.model_config) else: @@ -134,7 +137,7 @@ def _validate_lora(self, lora_request: LoRARequest | None) -> None: # LoRA request passed in while LoRA is not enabled if not self.lora_config: - raise ValueError(f"Got lora_request {lora_request} but LoRA is not enabled!") + raise APHRODITEValidationError(f"Got lora_request {lora_request} but LoRA is not enabled!") if self.tokenizer is not None: logger.warning_once( @@ -240,7 +243,7 @@ def process_inputs( dp_local_size = parallel_config.data_parallel_size_local num_ranks = dp_local_size if parallel_config.local_engines_only else dp_size if data_parallel_rank is not None and not (0 <= data_parallel_rank < num_ranks): - raise ValueError(f"data_parallel_rank {data_parallel_rank} is out of range [0, {num_ranks}).") + raise APHRODITEValidationError(f"data_parallel_rank {data_parallel_rank} is out of range [0, {num_ranks}).") if isinstance(prompt, dict) and "type" in prompt: if tokenization_kwargs: @@ -365,7 +368,7 @@ def _validate_prompt_len( return if prompt_len == 0 and prompt_type == "decoder": - raise ValueError(f"The {prompt_type} prompt cannot be empty") + raise APHRODITEValidationError(f"The {prompt_type} prompt cannot be empty") model_config = self.model_config max_prompt_len = model_config.max_model_len if prompt_type == "decoder" else self.mm_encoder_cache_size @@ -380,7 +383,7 @@ def _validate_prompt_len( else: suggestion = "Make sure that `max_model_len` is no smaller than the number of text tokens." - raise ValueError( + raise APHRODITEValidationError( f"The {prompt_type} prompt (length {prompt_len}) is " f"longer than the maximum model length of {max_prompt_len}. " f"{suggestion}" @@ -390,7 +393,7 @@ def _validate_prompt_len( "Make sure that `max_model_len` is no smaller than the " "number of text tokens (prompt + requested output tokens)." ) - raise ValueError( + raise APHRODITEValidationError( f"The {prompt_type} prompt (length {prompt_len}) plus the number of " f"requested output tokens (at least 1) is longer than the maximum " f"model length of {max_prompt_len}. {suggestion}" @@ -416,7 +419,7 @@ def _validate_model_input( for mm_position in mm_positions: num_embeds = mm_position.get_num_embeds() if num_embeds > self.mm_encoder_cache_size: - raise ValueError( + raise APHRODITEValidationError( f"The {prompt_type} prompt contains a(n) {modality} item " f"with {num_embeds} embedding tokens, which exceeds the " f"pre-allocated encoder cache size " @@ -440,7 +443,7 @@ def _validate_model_input( # truly out-of-vocabulary. model_vocab_size = model_config.get_vocab_size() if max_input_id > max(tokenizer.max_token_id, model_vocab_size - 1): - raise ValueError(f"Token id {max_input_id} is out of vocabulary") + raise APHRODITEValidationError(f"Token id {max_input_id} is out of vocabulary") def _validate_model_inputs( self, diff --git a/aphrodite/v1/kv_offload/base.py b/aphrodite/v1/kv_offload/base.py index bcff9ecbc5..c1207dd4eb 100644 --- a/aphrodite/v1/kv_offload/base.py +++ b/aphrodite/v1/kv_offload/base.py @@ -13,8 +13,6 @@ import numpy as np import torch -from aphrodite.logger import init_logger - if TYPE_CHECKING: from aphrodite.distributed.kv_transfer.kv_connector.v1.offloading.metrics import ( OffloadingConnectorStats, @@ -27,8 +25,6 @@ # Use the helper functions below to construct / decompose keys. OffloadKey = NewType("OffloadKey", bytes) -logger = init_logger(__name__) - def make_offload_key(block_hash: bytes, group_idx: int) -> OffloadKey: """Pack a block hash and group index into an `OffloadKey`.""" @@ -532,10 +528,6 @@ def build_metric_definitions(cls, extra_config: dict[str, Any]) -> dict[str, "Of return {} def __init__(self, config: OffloadingConfig): - logger.warning( - "Initializing OffloadingSpec. This API is experimental and " - "subject to change in the future as we iterate the design." - ) self.config = config self.extra_config = config.extra_config self.replicated_layout: bool = False diff --git a/aphrodite/v1/kv_offload/cpu/manager.py b/aphrodite/v1/kv_offload/cpu/manager.py index 492c0962a3..f15d67938f 100644 --- a/aphrodite/v1/kv_offload/cpu/manager.py +++ b/aphrodite/v1/kv_offload/cpu/manager.py @@ -2,7 +2,6 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from collections import OrderedDict from collections.abc import Collection, Iterable -from typing import Literal from typing_extensions import override @@ -24,19 +23,15 @@ CPULoadStoreSpec, CPUOffloadingMetrics, ) -from aphrodite.v1.kv_offload.cpu.policies.arc import ARCCachePolicy from aphrodite.v1.kv_offload.cpu.policies.base import BlockStatus, CachePolicy -from aphrodite.v1.kv_offload.cpu.policies.lru import LRUCachePolicy - -_CACHE_POLICIES: dict[str, type[CachePolicy]] = { - "lru": LRUCachePolicy, - "arc": ARCCachePolicy, -} +from aphrodite.v1.kv_offload.cpu.policies.factory import CachePolicyFactory class CPUOffloadingManager(OffloadingManager): """ - An OffloadingManager with a pluggable CachePolicy (LRU or ARC). + An OffloadingManager with a pluggable CachePolicy, resolved by name via + CachePolicyFactory (built in: "lru", "arc"; external policies can either + register their own or be loaded out-of-tree via cache_policy_module_path). The manager owns all shared logic: ref-counting, event emission, block pool management, and the prepare_store/complete_store skeletons. @@ -47,7 +42,8 @@ class CPUOffloadingManager(OffloadingManager): def __init__( self, num_blocks: int, - cache_policy: Literal["lru", "arc"] = "lru", + cache_policy: str = "lru", + cache_policy_module_path: str | None = None, enable_events: bool = False, store_threshold: int = 1, max_tracker_size: int = 64_000, @@ -57,9 +53,7 @@ def __init__( self._num_allocated_blocks: int = 0 self._free_list: list[int] = [] self.events: list[OffloadingEvent] | None = [] if enable_events else None - policy_cls = _CACHE_POLICIES.get(cache_policy) - if policy_cls is None: - raise ValueError(f"Unknown cache policy: {cache_policy!r}. Supported: {list(_CACHE_POLICIES)}") + policy_cls = CachePolicyFactory.get_cache_policy_cls(cache_policy, cache_policy_module_path) self._policy: CachePolicy = policy_cls(cache_capacity=num_blocks) # Track the number of blocks in the cache that are evictable. i.e. ref_cnt 0. self._num_evictable_cache_blocks: int = 0 diff --git a/aphrodite/v1/kv_offload/cpu/policies/arc.py b/aphrodite/v1/kv_offload/cpu/policies/arc.py index d7b0488ac4..ae3b2d361b 100644 --- a/aphrodite/v1/kv_offload/cpu/policies/arc.py +++ b/aphrodite/v1/kv_offload/cpu/policies/arc.py @@ -48,7 +48,7 @@ class ARCCachePolicy(CachePolicy): """ def __init__(self, cache_capacity: int): - self.cache_capacity: int = cache_capacity + super().__init__(cache_capacity) self.target_t1_size: float = 0.0 self.t1: OrderedDict[OffloadKey, BlockStatus] = OrderedDict() self.t2: OrderedDict[OffloadKey, BlockStatus] = OrderedDict() diff --git a/aphrodite/v1/kv_offload/cpu/policies/base.py b/aphrodite/v1/kv_offload/cpu/policies/base.py index 3d5fdc059d..c37d5b1bdc 100644 --- a/aphrodite/v1/kv_offload/cpu/policies/base.py +++ b/aphrodite/v1/kv_offload/cpu/policies/base.py @@ -41,8 +41,8 @@ class CachePolicy(ABC): and eviction, so they cannot be separated cleanly. """ - @abstractmethod - def __init__(self, cache_capacity: int) -> None: ... + def __init__(self, cache_capacity: int) -> None: + self.cache_capacity = cache_capacity @abstractmethod def get(self, key: OffloadKey) -> BlockStatus | None: diff --git a/aphrodite/v1/kv_offload/cpu/policies/factory.py b/aphrodite/v1/kv_offload/cpu/policies/factory.py new file mode 100644 index 0000000000..4105456ac2 --- /dev/null +++ b/aphrodite/v1/kv_offload/cpu/policies/factory.py @@ -0,0 +1,58 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import importlib +from collections.abc import Callable + +from aphrodite.logger import init_logger +from aphrodite.v1.kv_offload.cpu.policies.base import CachePolicy + +logger = init_logger(__name__) + + +class CachePolicyFactory: + """Registry for CachePolicy implementations, resolved by name. + + Built-in policies are pre-registered below. External policies can either + register a friendly short name up front or pass a module path at lookup + time without modifying Aphrodite. + """ + + _registry: dict[str, Callable[[], type[CachePolicy]]] = {} + + @classmethod + def register_cache_policy(cls, name: str, module_path: str, class_name: str) -> None: + """Register a cache policy with a lazy-loading module and class name.""" + if name in cls._registry: + raise ValueError(f"Cache policy '{name}' is already registered.") + + def loader() -> type[CachePolicy]: + module = importlib.import_module(module_path) + return getattr(module, class_name) + + cls._registry[name] = loader + + @classmethod + def get_cache_policy_cls(cls, name: str, module_path: str | None = None) -> type[CachePolicy]: + """Get a registered or out-of-tree cache policy class by name.""" + if name in cls._registry: + return cls._registry[name]() + if module_path is None: + raise ValueError( + f"Unknown cache policy: {name!r}. Supported: {list(cls._registry)}. " + "For an out-of-tree policy, also set cache_policy_module_path." + ) + logger.warning( + "Loading out-of-tree cache policy '%s' from '%s'. This API is " + "experimental and subject to change in the future as we " + "iterate the design.", + name, + module_path, + ) + module = importlib.import_module(module_path) + policy_cls = getattr(module, name) + assert issubclass(policy_cls, CachePolicy) + return policy_cls + + +CachePolicyFactory.register_cache_policy("lru", "aphrodite.v1.kv_offload.cpu.policies.lru", "LRUCachePolicy") +CachePolicyFactory.register_cache_policy("arc", "aphrodite.v1.kv_offload.cpu.policies.arc", "ARCCachePolicy") diff --git a/aphrodite/v1/kv_offload/cpu/policies/lru.py b/aphrodite/v1/kv_offload/cpu/policies/lru.py index 4f2d93cddf..aaf2815169 100644 --- a/aphrodite/v1/kv_offload/cpu/policies/lru.py +++ b/aphrodite/v1/kv_offload/cpu/policies/lru.py @@ -19,6 +19,7 @@ class LRUCachePolicy(CachePolicy): """ def __init__(self, cache_capacity: int): + super().__init__(cache_capacity) # Blocks with ref_cnt 0 (not participating in any loads/stores) ordered in LRU self.evictable_blocks: OrderedDict[OffloadKey, None] = OrderedDict() self.blocks: dict[OffloadKey, BlockStatus] = {} diff --git a/aphrodite/v1/kv_offload/cpu/spec.py b/aphrodite/v1/kv_offload/cpu/spec.py index c64f7ce34a..a59a7883d1 100644 --- a/aphrodite/v1/kv_offload/cpu/spec.py +++ b/aphrodite/v1/kv_offload/cpu/spec.py @@ -114,6 +114,7 @@ def __init__(self, config: OffloadingConfig): self._worker: CPUOffloadingWorker | None = None self.eviction_policy: str = self.extra_config.get("eviction_policy", "lru") + self.cache_policy_module_path: str | None = self.extra_config.get("cache_policy_module_path") @override def get_manager(self) -> OffloadingManager: @@ -128,7 +129,8 @@ def get_manager(self) -> OffloadingManager: self._manager = CPUOffloadingManager( num_blocks=self.num_blocks, - cache_policy=self.eviction_policy, # type: ignore[arg-type] + cache_policy=self.eviction_policy, + cache_policy_module_path=self.cache_policy_module_path, enable_events=self.kv_events_config.enable_kv_cache_events, store_threshold=store_threshold, max_tracker_size=max_tracker_size, diff --git a/aphrodite/v1/kv_offload/factory.py b/aphrodite/v1/kv_offload/factory.py index 60dad49560..d9b43b7d8e 100644 --- a/aphrodite/v1/kv_offload/factory.py +++ b/aphrodite/v1/kv_offload/factory.py @@ -35,6 +35,13 @@ def get_spec_cls(cls, extra_config: Mapping[str, Any]) -> type[OffloadingSpec]: spec_module_path = extra_config.get("spec_module_path") if spec_module_path is None: raise ValueError(f"Unsupported spec type: {spec_name}") + logger.warning( + "Loading out-of-tree offloading spec '%s' from '%s'. This " + "API is experimental and subject to change in the future " + "as we iterate the design.", + spec_name, + spec_module_path, + ) spec_module = importlib.import_module(spec_module_path) spec_cls = getattr(spec_module, spec_name) assert issubclass(spec_cls, OffloadingSpec) diff --git a/aphrodite/v1/kv_offload/tiering/manager.py b/aphrodite/v1/kv_offload/tiering/manager.py index 09c3ac8b93..e97f01694f 100644 --- a/aphrodite/v1/kv_offload/tiering/manager.py +++ b/aphrodite/v1/kv_offload/tiering/manager.py @@ -93,11 +93,13 @@ def __init__( num_blocks: int, mmap_region: SharedOffloadRegion, cache_policy: str = "lru", + cache_policy_module_path: str | None = None, enable_events: bool = False, ): super().__init__( num_blocks=num_blocks, - cache_policy=cache_policy, # type: ignore[arg-type] + cache_policy=cache_policy, + cache_policy_module_path=cache_policy_module_path, enable_events=enable_events, ) self._mmap_region = mmap_region diff --git a/aphrodite/v1/kv_offload/tiering/spec.py b/aphrodite/v1/kv_offload/tiering/spec.py index a42e2886a1..333d7904cd 100644 --- a/aphrodite/v1/kv_offload/tiering/spec.py +++ b/aphrodite/v1/kv_offload/tiering/spec.py @@ -13,8 +13,13 @@ - blocks_per_chunk: (optional) Chunk size in GPU blocks (default: 1). Must be greater than 0. Use this instead of block_size when KV cache groups have different block sizes. - - eviction_policy: (optional) Primary tier eviction policy: "lru" or - "arc" (default: "lru") + - eviction_policy: (optional) Primary tier eviction policy: built-in "lru"/ + "arc", or the name of a policy registered via CachePolicyFactory, or an + out-of-tree CachePolicy class name paired with cache_policy_module_path + (default: "lru") + - cache_policy_module_path: (optional) Python import path to load + eviction_policy from when it names an out-of-tree CachePolicy not + registered via CachePolicyFactory - secondary_tiers: (optional) List of secondary tier configurations Each secondary tier config is a dict with: - type: (required) Type of secondary tier (e.g., "example", "storage", "network") @@ -176,7 +181,8 @@ def get_manager(self) -> OffloadingManager: # Create primary tier (CPU-based) primary_tier = CPUPrimaryTierOffloadingManager( num_blocks=self.num_blocks, - cache_policy=self.eviction_policy, # type: ignore[arg-type] + cache_policy=self.eviction_policy, + cache_policy_module_path=self.cache_policy_module_path, enable_events=self.kv_events_config.enable_kv_cache_events, mmap_region=scheduler_mmap, ) diff --git a/aphrodite/v1/sample/logits_processor/__init__.py b/aphrodite/v1/sample/logits_processor/__init__.py index 7cd36b2ed6..9b3e10db15 100644 --- a/aphrodite/v1/sample/logits_processor/__init__.py +++ b/aphrodite/v1/sample/logits_processor/__init__.py @@ -10,6 +10,7 @@ import torch +from aphrodite.exceptions import APHRODITEValidationError from aphrodite.logger import init_logger from aphrodite.logits_process import LogitsProcessor as RequestLogitsProcessor from aphrodite.sampling_params import SamplingParams @@ -211,7 +212,10 @@ def validate_logits_processors_parameters( ): logits_processors = tuple(logits_processors) if logits_processors is not None else None for logits_procs in cached_load_custom_logitsprocs(logits_processors): - logits_procs.validate_params(sampling_params) + try: + logits_procs.validate_params(sampling_params) + except ValueError as e: + raise APHRODITEValidationError(str(e)) from e class AdapterLogitsProcessor(LogitsProcessor): diff --git a/aphrodite/v1/sample/logits_processor/interface.py b/aphrodite/v1/sample/logits_processor/interface.py index ba4f58b67e..9037b8ec25 100644 --- a/aphrodite/v1/sample/logits_processor/interface.py +++ b/aphrodite/v1/sample/logits_processor/interface.py @@ -62,7 +62,10 @@ class LogitsProcessor(ABC): def validate_params(cls, sampling_params: SamplingParams): """Validate sampling params for this logits processor. - Raise ValueError for invalid ones. + Raise ``APHRODITEValidationError`` (preferred) / ``ValueError`` (backward + compatible) for invalid params. Bare ``ValueError`` is converted to + ``APHRODITEValidationError`` at the engine boundary so online serving + returns HTTP 400. """ return None diff --git a/aphrodite/v1/spec_decode/dflash.py b/aphrodite/v1/spec_decode/dflash.py index 4b99ec4395..924c669083 100644 --- a/aphrodite/v1/spec_decode/dflash.py +++ b/aphrodite/v1/spec_decode/dflash.py @@ -38,8 +38,12 @@ def __init__( # Only next_token_ids and mask tokens are query tokens, all other context is K/V self.max_query_tokens = self.max_batch_size * (1 + self.num_speculative_tokens) + self.max_padded_query_tokens = max( + self.max_query_tokens, + aphrodite_config.compilation_config.max_cudagraph_capture_size or 0, + ) # Positions covers both context states + query states - self.max_positions = self.max_num_tokens + self.max_query_tokens + self.max_positions = self.max_num_tokens + self.max_padded_query_tokens # Separate context buffers to keep query buffer addresses stable for CUDA graphs self._context_slot_mapping_buffer = torch.zeros( @@ -48,7 +52,7 @@ def __init__( device=device, ) self._slot_mapping_buffer = torch.zeros( - self.max_query_tokens, + self.max_padded_query_tokens, dtype=torch.int64, device=device, ) @@ -58,7 +62,7 @@ def __init__( device=device, ) self.positions = torch.zeros( - self.max_query_tokens, + self.max_padded_query_tokens, dtype=torch.int64, device=device, ) diff --git a/aphrodite/v1/spec_decode/llm_base_proposer.py b/aphrodite/v1/spec_decode/llm_base_proposer.py index dbd2f0f98d..5d369ca838 100644 --- a/aphrodite/v1/spec_decode/llm_base_proposer.py +++ b/aphrodite/v1/spec_decode/llm_base_proposer.py @@ -1605,7 +1605,8 @@ def initialize_attn_backend( attention_groups: dict[tuple[str, str], AttentionGroup] = {} if kv_cache_spec is not None: - for layer_name in self._draft_attn_layer_names: + # Iterate deterministically because draft layer names are a set. + for layer_name in sorted(self._draft_attn_layer_names): attn_backend = all_attn_layers[layer_name].get_attn_backend() backend_key = attn_backend.full_cls_name() if backend_key not in attention_groups: @@ -1634,7 +1635,11 @@ def initialize_attn_backend( attention_groups[backend_key].layer_names.append(layer_name) self.draft_attn_groups = list(attention_groups.values()) - self.block_size = self.draft_attn_groups[0].get_metadata_builder().kv_cache_spec.block_size + if kernel_block_sizes is not None and 0 <= self.kv_cache_gid < len(kernel_block_sizes): + # Slot mappings use the kernel-granularity block table. + self.block_size = kernel_block_sizes[self.kv_cache_gid] + else: + self.block_size = self.draft_attn_groups[0].get_metadata_builder().kv_cache_spec.block_size logger.debug("Using block size %d for drafting layers", self.block_size) def _determine_batch_execution_and_padding( diff --git a/aphrodite/v1/worker/block_table.py b/aphrodite/v1/worker/block_table.py index 4914859029..770b033e7d 100644 --- a/aphrodite/v1/worker/block_table.py +++ b/aphrodite/v1/worker/block_table.py @@ -125,6 +125,11 @@ def move_row(self, src: int, tgt: int) -> None: block_table_np = self.block_table.np block_table_np[tgt, :num_blocks] = block_table_np[src, :num_blocks] self.num_blocks_per_row[tgt] = num_blocks + # Clear the vacated source row: dummy-run batches dereference stale + # rows as mamba state slots and write state in place there, possibly + # after the blocks have been freed and reallocated. + block_table_np[src, :num_blocks] = 0 + self.num_blocks_per_row[src] = 0 def swap_row(self, src: int, tgt: int) -> None: src_tgt, tgt_src = [src, tgt], [tgt, src] diff --git a/aphrodite/v1/worker/gpu/block_table.py b/aphrodite/v1/worker/gpu/block_table.py index 4a4e2f8c8e..66ded8224a 100644 --- a/aphrodite/v1/worker/gpu/block_table.py +++ b/aphrodite/v1/worker/gpu/block_table.py @@ -147,7 +147,11 @@ def get_dummy_block_tables(self, num_reqs: int) -> tuple[torch.Tensor, ...]: # Therefore, this method must return the persistent tensor # with the same memory address as that used during the model's forward pass, # rather than allocating a new tensor. - return tuple(block_table[:num_reqs] for block_table in self.input_block_tables) + # + # Zero the rows so dummy runs write mamba state to the reserved null + # block rather than through the previous real step's (stale) block + # ids, which may point at blocks since freed and reallocated. + return tuple(block_table[:num_reqs].zero_() for block_table in self.input_block_tables) def compute_slot_mappings( self, diff --git a/aphrodite/v1/worker/gpu/model_runner.py b/aphrodite/v1/worker/gpu/model_runner.py index bcd5a6f907..cbb65983a1 100644 --- a/aphrodite/v1/worker/gpu/model_runner.py +++ b/aphrodite/v1/worker/gpu/model_runner.py @@ -1223,6 +1223,10 @@ def execute_model( slot_mappings, self.attn_groups, self.kv_cache_config, + # FULL replay reads capture-time metadata buffers. Re-stage them + # from the zeroed dummy block tables instead of retaining state + # indices from the previous real batch. + for_capture=dummy_run and batch_desc.cg_mode == CUDAGraphMode.FULL, ) input_ids = input_batch.input_ids diff --git a/aphrodite/v1/worker/gpu_model_runner.py b/aphrodite/v1/worker/gpu_model_runner.py index ddf9a6e9a6..ddbd96282c 100644 --- a/aphrodite/v1/worker/gpu_model_runner.py +++ b/aphrodite/v1/worker/gpu_model_runner.py @@ -897,9 +897,16 @@ def init_fp8_kv_scales(self) -> None: return kv_caches = getattr(self, "kv_caches", []) - for cache_tensor in kv_caches: - if cache_tensor is not None: - cache_tensor.zero_() + for cache_entry in kv_caches: + if cache_entry is None: + continue + # Hybrid models (Mamba, DeltaNet) store per-layer state as a + # list of tensors rather than a single tensor. + if isinstance(cache_entry, list): + for t in cache_entry: + t.zero_() + else: + cache_entry.zero_() k_attr_names = ("_k_scale", "k_scale") v_attr_names = ("_v_scale", "v_scale") @@ -5539,7 +5546,10 @@ def _dummy_run( num_reqs=num_reqs_padded, max_query_len=max_query_len, ubatch_slices=(ubatch_slices_padded if pad_attn else ubatch_slices), - for_cudagraph_capture=is_graph_capturing, + # FULL replay reads capture-time metadata buffers. Re-stage them + # from the zeroed dummy block tables instead of retaining state + # indices from the previous real batch. + for_cudagraph_capture=(is_graph_capturing or cudagraph_runtime_mode == CUDAGraphMode.FULL), slot_mappings=slot_mappings_by_group, use_spec_decode=self.speculative_config is not None, ) diff --git a/benchmarks/kernels/benchmark_k3_cutedsl_residual.py b/benchmarks/kernels/benchmark_k3_cutedsl_residual.py new file mode 100644 index 0000000000..c18fad4fba --- /dev/null +++ b/benchmarks/kernels/benchmark_k3_cutedsl_residual.py @@ -0,0 +1,340 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Benchmark the Kimi-K3 latent MoE addmm against CuTe residual GEMM. + +The benchmark covers ``BF16[M, 3584] @ BF16[7168, 3584].T + BF16[M, 7168]`` +with FP32 accumulation and BF16 output. Both backends execute through CUDA +Graph replay. Weights and residuals rotate across buffers exceeding L2 so the +comparison models the full latent MoE projection-and-add path. +""" + +from __future__ import annotations + +import argparse +import dataclasses +import importlib.util +import json +import math +import statistics +from collections.abc import Callable, Sequence +from pathlib import Path +from typing import Any + +import cutlass +import cutlass.cute as cute +import torch +from cuda.bindings import driver as cuda +from cuda.bindings.driver import CUstream +from quack.compile_utils import make_fake_tensor + +N = 7168 +K = 3584 + + +@dataclasses.dataclass(frozen=True, slots=True) +class Config: + block_size: int + outputs_per_block: int + k_unroll: int + vector_width: int = 8 + + +def parse_config(value: str) -> Config: + try: + parts = [int(part) for part in value.split(",")] + except ValueError as error: + raise argparse.ArgumentTypeError("config must be BLOCK,OUTPUTS,K_UNROLL[,VECTOR_WIDTH]") from error + if len(parts) == 3: + return Config(*parts) + if len(parts) == 4: + return Config(*parts) + raise argparse.ArgumentTypeError("config must be BLOCK,OUTPUTS,K_UNROLL[,VECTOR_WIDTH]") + + +def production_residual_config(m: int) -> Config | None: + """The measured Latent-MoE residual config for M, from the K3 table.""" + from aphrodite.models.kimi_k3.nvidia.low_latency_gemm import KIMI_K3_PROJECTIONS + + spec = KIMI_K3_PROJECTIONS.get((N, K)) + config = spec.residual_config(m) if spec is not None else None + if config is None: + return None + return Config( + config.block_size, + config.outputs_per_block, + config.k_unroll, + config.vector_width, + ) + + +def candidate_configs(mode: str, selected: Config | None, m: int) -> list[Config]: + if mode == "selected": + if selected is not None: + return [selected] + # No explicit --config: fall back to the production table for this M. + config = production_residual_config(m) + return [config] if config is not None else [] + if mode == "baseline": + return [Config(224, 4, 2)] + return [ + Config(block_size, outputs_per_block, k_unroll, vector_width) + for vector_width in (4, 8) + for block_size in (32, 64, 128, 224, 448) + if block_size % 32 == 0 and K % (block_size * vector_width) == 0 + for outputs_per_block in (1, 2, 4, 7, 8) + if N % outputs_per_block == 0 + for k_unroll in (1, 2, 4) + ] + + +def load_kernel_class(path: Path): + spec = importlib.util.spec_from_file_location("cute_skinny_device", path) + if spec is None or spec.loader is None: + raise RuntimeError(f"cannot load CuTe kernel from {path}") + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module.CuteSkinnyGemm + + +def stream() -> CUstream: + return CUstream(torch.cuda.current_stream().cuda_stream) + + +def compile_kernel(kernel_class, m: int, config: Config, max_registers: int): + element_type = cutlass.BFloat16 + n = cute.sym_int(divisibility=config.outputs_per_block) + k = cute.sym_int(divisibility=config.block_size * config.vector_width) + a = make_fake_tensor(element_type, (m, k), divisibility=config.vector_width) + b = make_fake_tensor(element_type, (n, k), divisibility=config.vector_width) + residual = make_fake_tensor(element_type, (m, n), divisibility=1) + c = make_fake_tensor(element_type, (m, n), divisibility=1) + kernel = kernel_class( + element_type=element_type, + num_rows=m, + block_size=config.block_size, + outputs_per_block=config.outputs_per_block, + vector_width=config.vector_width, + k_unroll=config.k_unroll, + has_residual=True, + use_pdl=True, + ) + return cute.compile( + kernel, + a, + b, + residual, + c, + stream(), + options=( + f"--enable-tvm-ffi --keep-cubin --ptxas-options -maxrregcount={max_registers} --ptxas-options -lineinfo" + ), + ) + + +def resource_usage(compiled) -> dict[str, Any]: + executor = getattr(compiled, "_default_executor", None) + context = getattr(executor, "exec_context", None) + functions = getattr(context, "kernel_functions", None) + if not functions: + return {"resource_metrics_available": False} + + def attribute(name, function) -> int: + error, value = cuda.cuFuncGetAttribute(name, function) + if error != cuda.CUresult.CUDA_SUCCESS: + raise RuntimeError(f"cuFuncGetAttribute failed with {error}") + return int(value) + + registers = [attribute(cuda.CUfunction_attribute.CU_FUNC_ATTRIBUTE_NUM_REGS, function) for function in functions] + local_bytes = [ + attribute( + cuda.CUfunction_attribute.CU_FUNC_ATTRIBUTE_LOCAL_SIZE_BYTES, + function, + ) + for function in functions + ] + return { + "resource_metrics_available": True, + "registers_per_thread": max(registers, default=0), + "spill_bytes": max(local_bytes, default=0), + } + + +def rotating_buffer_count(m: int, multiplier: float, limit: int) -> int: + properties = torch.cuda.get_device_properties(0) + bytes_per_pair = (N * K + m * N) * 2 + target = math.ceil(multiplier * properties.L2_cache_size) + return max(2, min(limit, math.ceil(target / bytes_per_pair))) + + +def graph_samples( + launch: Callable[[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor], None], + activation: torch.Tensor, + weights: Sequence[torch.Tensor], + residuals: Sequence[torch.Tensor], + repeats: int, + replays: int, +) -> tuple[list[float], list[torch.Tensor]]: + outputs = [torch.empty_like(residual) for residual in residuals] + for weight, residual, output in zip(weights, residuals, outputs): + launch(activation, weight, residual, output) + torch.accelerator.synchronize() + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + for weight, residual, output in zip(weights, residuals, outputs): + launch(activation, weight, residual, output) + for _ in range(20): + graph.replay() + torch.accelerator.synchronize() + + samples = [] + for _ in range(repeats): + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + start.record() + for _ in range(replays): + graph.replay() + end.record() + end.synchronize() + samples.append(start.elapsed_time(end) * 1000.0 / (replays * len(weights))) + return samples, outputs + + +def summarize(samples: Sequence[float]) -> dict[str, Any]: + ordered = sorted(samples) + + def percentile(fraction: float) -> float: + position = fraction * (len(ordered) - 1) + lower = math.floor(position) + upper = math.ceil(position) + if lower == upper: + return ordered[lower] + weight = position - lower + return ordered[lower] * (1.0 - weight) + ordered[upper] * weight + + mean = statistics.mean(samples) + return { + "median_us": statistics.median(samples), + "p10_us": percentile(0.1), + "p90_us": percentile(0.9), + "mean_us": mean, + "cv_pct": statistics.pstdev(samples) / mean * 100.0, + "samples_us": list(samples), + } + + +def correctness( + output: torch.Tensor, + activation: torch.Tensor, + weight: torch.Tensor, + residual: torch.Tensor, +) -> dict[str, Any]: + actual = output.float() + reference = activation.float() @ weight.float().t() + residual.float() + error = (actual - reference).abs() + scaled_error = error / (reference.abs() + 1.0) + cosine = torch.nn.functional.cosine_similarity(actual.flatten(), reference.flatten(), dim=0).item() + return { + "valid": cosine > 0.999, + "cosine": cosine, + "max_abs_error": error.max().item(), + "max_scaled_error": scaled_error.max().item(), + } + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--kernel", type=Path, required=True) + parser.add_argument("--output", type=Path, required=True) + parser.add_argument("--mode", choices=("baseline", "sweep", "selected"), default="baseline") + parser.add_argument("--config", type=parse_config) + parser.add_argument("--m", type=int, action="append") + parser.add_argument("--config-shard", type=int, default=0) + parser.add_argument("--num-config-shards", type=int, default=1) + parser.add_argument("--repeats", type=int, default=21) + parser.add_argument("--replays", type=int, default=200) + parser.add_argument("--cache-multiplier", type=float, default=3.0) + parser.add_argument("--max-buffers", type=int, default=32) + parser.add_argument("--max-registers", type=int, default=64) + args = parser.parse_args() + + token_counts = args.m or list(range(1, 17)) + if any(not 1 <= m <= 16 for m in token_counts): + raise ValueError("expected 1 <= M <= 16") + if not 0 <= args.config_shard < args.num_config_shards: + raise ValueError("config shard must be in [0, num_config_shards)") + torch.accelerator.set_device_index(0) + if torch.cuda.get_device_capability() != (10, 3): + raise RuntimeError("this benchmark requires SM103") + + kernel_class = load_kernel_class(args.kernel) + properties = torch.cuda.get_device_properties(0) + metadata = { + "device": properties.name, + "compute_capability": list(torch.cuda.get_device_capability()), + "torch_version": torch.__version__, + "cuda_version": torch.version.cuda, + } + args.output.parent.mkdir(parents=True, exist_ok=True) + with args.output.open("w", encoding="utf-8") as output_file: + for m in token_counts: + configs = candidate_configs(args.mode, args.config, m) + torch.manual_seed(20260722 + m) + count = rotating_buffer_count(m, args.cache_multiplier, args.max_buffers) + activation = torch.randn((m, K), device="cuda", dtype=torch.bfloat16) + weights = [torch.randn((N, K), device="cuda", dtype=torch.bfloat16) for _ in range(count)] + residuals = [torch.randn((m, N), device="cuda", dtype=torch.bfloat16) for _ in range(count)] + candidates: list[tuple[str, Config | None]] = [("cublas_addmm", None)] + candidates.extend( + ("cute_residual", config) + for index, config in enumerate(configs) + if index % args.num_config_shards == args.config_shard + ) + for backend, config in candidates: + row: dict[str, Any] = { + "m": m, + "n": N, + "k": K, + "backend": backend, + "mode": args.mode, + "config": dataclasses.asdict(config) if config else {}, + "num_buffers": count, + "cache_multiplier": args.cache_multiplier, + **metadata, + } + try: + if backend == "cublas_addmm": + launch = lambda a, b, residual, c: torch.addmm(residual, a, b.t(), out=c) + else: + if config is None: + raise AssertionError("missing CuTe config") + compiled = compile_kernel(kernel_class, m, config, args.max_registers) + launch = lambda a, b, residual, c, fn=compiled: fn(a, b, residual, c, stream()) + row.update(resource_usage(compiled)) + samples, outputs = graph_samples( + launch, + activation, + weights, + residuals, + args.repeats, + args.replays, + ) + row.update(correctness(outputs[0], activation, weights[0], residuals[0])) + row.update(summarize(samples)) + except Exception as error: # noqa: BLE001 + row.update( + { + "valid": False, + "error": f"{type(error).__name__}: {error}", + } + ) + output_file.write(json.dumps(row, sort_keys=True) + "\n") + output_file.flush() + print(json.dumps(row, sort_keys=True), flush=True) + + del activation, weights, residuals + torch.accelerator.empty_cache() + + +if __name__ == "__main__": + main() diff --git a/benchmarks/kernels/benchmark_kimi_k3_latent_moe_tail.py b/benchmarks/kernels/benchmark_kimi_k3_latent_moe_tail.py new file mode 100644 index 0000000000..b590d34081 --- /dev/null +++ b/benchmarks/kernels/benchmark_kimi_k3_latent_moe_tail.py @@ -0,0 +1,787 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Benchmark the Kimi K3 latent-MoE tail and its up-projection kernels. + +The ``up-projection`` subcommand isolates the TP-local dynamic and static-M +skinny GEMMs. It rotates weights through a working set larger than L2 to model +successive model layers. + +The ``whole-tail`` subcommand measures the distributed operator. Its reference +path includes two AllReduces, RMSNorm, the replicated up-projection, and the +final add. CUDA-event samples report the slowest rank so cross-rank skew is +included. + +Examples: + +.. code-block:: console + + .venv/bin/python \ + benchmarks/kernels/benchmark_kimi_k3_latent_moe_tail.py up-projection + + torchrun --nproc-per-node=8 \ + benchmarks/kernels/benchmark_kimi_k3_latent_moe_tail.py whole-tail + +For multi-node runs, launch one ``torchrun`` agent per node and use a shared +rendezvous endpoint. +""" + +from __future__ import annotations + +import argparse +import json +import math +import os +import statistics +from collections.abc import Callable, Sequence +from dataclasses import asdict +from pathlib import Path +from typing import Any + +import cutlass +import cutlass.utils as utils +import torch +import torch.distributed as dist +import torch.nn.functional as F +from cuda.bindings import driver as cuda + +from aphrodite.distributed import get_tp_group +from aphrodite.distributed.parallel_state import ( + init_distributed_environment, + initialize_model_parallel, + set_custom_all_reduce, +) +from aphrodite.model_executor.warmup.cutedsl_warmup import cutedsl_warmup +from aphrodite.models.kimi_k3.nvidia.ops import latent_moe_tail +from aphrodite.models.kimi_k3.nvidia.ops.cute_dsl.latent_moe_tail import ( + fused_add_multicast_gemm, + fused_add_multicast_skinny_gemm, +) + +HIDDEN_SIZE = 7168 +LATENT_SIZE = 3584 +RMS_EPS = 0.1 +MAX_NUM_TOKENS = 16 +MMA_TILER_MN = (64, 32) +CLUSTER_SHAPE_MN = (1, 8) +B_PRIME_STAGES = 2 + + +def parse_up_projection_config( + value: str, +) -> fused_add_multicast_skinny_gemm.SkinnyConfig: + try: + values = [int(part) for part in value.split(",")] + except ValueError as error: + raise argparse.ArgumentTypeError("config must be BLOCK,OUTPUTS,K_UNROLL[,VECTOR_WIDTH[,PREFETCH_B]]") from error + if len(values) in (3, 4): + return fused_add_multicast_skinny_gemm.SkinnyConfig(*values) + if len(values) == 5 and values[4] in (0, 1): + return fused_add_multicast_skinny_gemm.SkinnyConfig( + *values[:4], + prefetch_b_before_pdl=bool(values[4]), + ) + raise argparse.ArgumentTypeError( + "config must be BLOCK,OUTPUTS,K_UNROLL[,VECTOR_WIDTH[,PREFETCH_B]], where PREFETCH_B is 0 or 1" + ) + + +def parse_tail_skinny_config( + value: str, +) -> tuple[int, fused_add_multicast_skinny_gemm.SkinnyConfig]: + try: + values = [int(part) for part in value.split(",")] + except ValueError as error: + raise argparse.ArgumentTypeError( + "config must be M,BLOCK,OUTPUTS,K_UNROLL[,VECTOR_WIDTH[,PREFETCH_B]]" + ) from error + if len(values) == 4: + num_tokens, *config = values + return num_tokens, fused_add_multicast_skinny_gemm.SkinnyConfig(*config) + if len(values) == 5: + num_tokens, *config = values + return num_tokens, fused_add_multicast_skinny_gemm.SkinnyConfig(*config) + if len(values) == 6 and values[5] in (0, 1): + num_tokens, block, outputs, unroll, vector_width, prefetch = values + return num_tokens, fused_add_multicast_skinny_gemm.SkinnyConfig( + block, + outputs, + unroll, + vector_width, + bool(prefetch), + ) + raise argparse.ArgumentTypeError( + "config must be M,BLOCK,OUTPUTS,K_UNROLL[,VECTOR_WIDTH[,PREFETCH_B]], where PREFETCH_B is 0 or 1" + ) + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + subparsers = parser.add_subparsers(dest="scope", required=True) + + up_projection = subparsers.add_parser( + "up-projection", + help="Benchmark the isolated TP-local up-projection kernels.", + ) + up_projection.add_argument( + "--backend", + choices=("dynamic", "skinny", "both"), + default="both", + ) + up_projection.add_argument("--tp-size", type=int, default=16) + up_projection.add_argument( + "--num-tokens", + type=int, + nargs="+", + default=[*range(1, 9), 16], + ) + up_projection.add_argument( + "--skinny-config", + type=parse_up_projection_config, + action="append", + help="Benchmark a static-M config for every selected token count.", + ) + up_projection.add_argument("--cache-multiplier", type=float, default=2.0) + up_projection.add_argument("--max-weights", type=int, default=64) + up_projection.add_argument("--warmup-replays", type=int, default=10) + up_projection.add_argument("--samples", type=int, default=31) + up_projection.add_argument("--output", type=Path) + + whole_tail = subparsers.add_parser( + "whole-tail", + help="Benchmark the distributed latent-MoE tail operator.", + ) + whole_tail.add_argument( + "--backend", + choices=("reference", "fused", "both"), + default="both", + ) + whole_tail.add_argument( + "--num-tokens", + type=int, + nargs="+", + default=[1, 5, 8, 16], + ) + whole_tail.add_argument("--warmup-replays", type=int, default=20) + whole_tail.add_argument("--samples", type=int, default=51) + whole_tail.add_argument( + "--skinny-max-num-tokens", + type=int, + nargs="+", + help="Override the fused operator's static-M cutoff; use 0 for dynamic-only.", + ) + whole_tail.add_argument( + "--skinny-config", + type=parse_tail_skinny_config, + action="append", + help="Override one static-M config for tuning.", + ) + whole_tail.add_argument("--output", type=Path) + return parser.parse_args() + + +def percentile(samples: Sequence[float], fraction: float) -> float: + ordered = sorted(samples) + position = fraction * (len(ordered) - 1) + lower = math.floor(position) + upper = math.ceil(position) + if lower == upper: + return ordered[lower] + upper_weight = position - lower + return ordered[lower] * (1.0 - upper_weight) + ordered[upper] * upper_weight + + +def summarize(samples_us: Sequence[float]) -> dict[str, Any]: + mean_us = statistics.mean(samples_us) + return { + "median_us": statistics.median(samples_us), + "p10_us": percentile(samples_us, 0.1), + "p90_us": percentile(samples_us, 0.9), + "mean_us": mean_us, + "cv_pct": statistics.pstdev(samples_us) / mean_us * 100.0, + "samples_us": list(samples_us), + } + + +def rotating_weight_count( + shard_size: int, + cache_multiplier: float, + limit: int, +) -> int: + properties = torch.cuda.get_device_properties(torch.accelerator.current_device_index()) + weight_bytes = shard_size * LATENT_SIZE * 2 + target_bytes = math.ceil(properties.L2_cache_size * cache_multiplier) + return max(2, min(limit, math.ceil(target_bytes / weight_bytes))) + + +def capture_up_projection_graph( + launches: Sequence[Callable[[], None]], +) -> torch.cuda.CUDAGraph: + for launch in launches: + launch() + torch.accelerator.synchronize() + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + for launch in launches: + launch() + torch.accelerator.synchronize() + return graph + + +def benchmark_up_projection_graph( + graph: torch.cuda.CUDAGraph, + *, + operations_per_replay: int, + warmup_replays: int, + samples: int, +) -> dict[str, Any]: + for _ in range(warmup_replays): + graph.replay() + torch.accelerator.synchronize() + + samples_us = [] + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + for _ in range(samples): + start.record() + graph.replay() + end.record() + end.synchronize() + samples_us.append(start.elapsed_time(end) * 1000.0 / operations_per_replay) + return summarize(samples_us) + + +class DynamicKernel: + def __init__( + self, + shard_size: int, + mailbox: torch.Tensor, + shared_shard: torch.Tensor, + ) -> None: + self.shard_size = shard_size + self.mailbox = mailbox + self.mailbox_c = fused_add_multicast_gemm._as_cute(mailbox) + compile_latent = torch.empty( + (1, MAX_NUM_TOKENS, LATENT_SIZE), + dtype=torch.bfloat16, + device=mailbox.device, + ) + compile_weight = torch.empty( + (1, shard_size, LATENT_SIZE), + dtype=torch.bfloat16, + device=mailbox.device, + ) + cluster_size = math.prod(CLUSTER_SHAPE_MN) + max_active_clusters = utils.HardwareInfo().get_max_active_clusters(cluster_size) + self.compiled = fused_add_multicast_gemm.compile_kernel( + (MAX_NUM_TOKENS, shard_size, LATENT_SIZE, 1), + fused_add_multicast_gemm._as_cute( + compile_latent, + dynamic_m=True, + ), + fused_add_multicast_gemm._as_cute(compile_weight), + self.mailbox_c, + fused_add_multicast_gemm._as_cute(shared_shard), + HIDDEN_SIZE, + shard_size, + MMA_TILER_MN, + CLUSTER_SHAPE_MN, + max_active_clusters, + B_PRIME_STAGES, + ) + + def launch( + self, + latent: torch.Tensor, + weight: torch.Tensor, + shared_shard: torch.Tensor, + ) -> None: + stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) + self.compiled( + fused_add_multicast_gemm._as_cute( + latent.unsqueeze(0), + dynamic_m=True, + ), + fused_add_multicast_gemm._as_cute(weight.unsqueeze(0)), + self.mailbox_c, + fused_add_multicast_gemm._as_cute(shared_shard), + cutlass.Int64(latent.shape[0]), + cutlass.Int64(self.mailbox.data_ptr()), + stream, + ) + + +class SkinnyKernel: + def __init__( + self, + num_tokens: int, + shard_size: int, + config: fused_add_multicast_skinny_gemm.SkinnyConfig, + ) -> None: + self.compiled = fused_add_multicast_skinny_gemm.compile_kernel( + num_rows=num_tokens, + latent_dim=LATENT_SIZE, + hidden_dim=HIDDEN_SIZE, + shard_dim=shard_size, + config=config, + ) + + def launch( + self, + latent: torch.Tensor, + weight: torch.Tensor, + shared_shard: torch.Tensor, + mailbox: torch.Tensor, + ) -> None: + self.compiled( + fused_add_multicast_skinny_gemm._as_cute(latent), + fused_add_multicast_skinny_gemm._as_cute(weight), + fused_add_multicast_skinny_gemm._as_cute(shared_shard), + cutlass.Int64(mailbox.data_ptr()), + cuda.CUstream(torch.cuda.current_stream().cuda_stream), + ) + + +def check_up_projection_output( + actual: torch.Tensor, + latent: torch.Tensor, + weight: torch.Tensor, + shared_shard: torch.Tensor, +) -> None: + gemm = F.linear(latent.float(), weight.float()).to(torch.bfloat16) + expected = (gemm.float() + shared_shard.float()).to(torch.bfloat16) + torch.testing.assert_close(actual, expected, atol=8e-2, rtol=3e-2) + + +def make_up_projection_launches( + launch: Callable[[torch.Tensor, torch.Tensor, torch.Tensor], None], + latent: torch.Tensor, + weights: Sequence[torch.Tensor], + shared_shard: torch.Tensor, +) -> list[Callable[[], None]]: + return [lambda weight=weight: launch(latent, weight, shared_shard) for weight in weights] + + +def benchmark_up_projection(args: argparse.Namespace) -> None: + if args.tp_size <= 0 or HIDDEN_SIZE % args.tp_size: + raise ValueError("TP size must be positive and divide the hidden size") + if any(not 1 <= num_tokens <= MAX_NUM_TOKENS for num_tokens in args.num_tokens): + raise ValueError("--num-tokens values must be in [1, 16]") + if args.cache_multiplier <= 0 or args.max_weights <= 0: + raise ValueError("cache multiplier and max weights must be positive") + if args.warmup_replays < 0 or args.samples <= 0: + raise ValueError("warmup replays must be nonnegative and samples positive") + + torch.accelerator.set_device_index(0) + device = torch.device("cuda", 0) + if torch.cuda.get_device_capability(device)[0] != 10: + raise RuntimeError("Kimi K3 latent-MoE tail requires SM100") + + shard_size = HIDDEN_SIZE // args.tp_size + weight_count = rotating_weight_count( + shard_size, + args.cache_multiplier, + args.max_weights, + ) + torch.manual_seed(20260726) + weights = [ + torch.randn( + (shard_size, LATENT_SIZE), + dtype=torch.bfloat16, + device=device, + ) + / LATENT_SIZE**0.5 + for _ in range(weight_count) + ] + mailbox = torch.empty( + (1, MAX_NUM_TOKENS, HIDDEN_SIZE), + dtype=torch.bfloat16, + device=device, + ) + shared = torch.randn( + (MAX_NUM_TOKENS, HIDDEN_SIZE), + dtype=torch.bfloat16, + device=device, + ) + shared_shard = shared[:, :shard_size] + use_dynamic = args.backend in ("dynamic", "both") + use_skinny = args.backend in ("skinny", "both") + dynamic_kernel = DynamicKernel(shard_size, mailbox, shared_shard) if use_dynamic else None + + results = [] + for num_tokens in args.num_tokens: + latent = torch.randn( + (num_tokens, LATENT_SIZE), + dtype=torch.bfloat16, + device=device, + ) + result: dict[str, Any] = {"num_tokens": num_tokens} + if dynamic_kernel is not None: + launches = make_up_projection_launches( + dynamic_kernel.launch, + latent, + weights, + shared_shard, + ) + graph = capture_up_projection_graph(launches) + result["dynamic"] = benchmark_up_projection_graph( + graph, + operations_per_replay=len(launches), + warmup_replays=args.warmup_replays, + samples=args.samples, + ) + check_up_projection_output( + mailbox[0, :num_tokens, :shard_size], + latent, + weights[-1], + shared_shard[:num_tokens], + ) + if use_skinny: + configs = args.skinny_config or [ + fused_add_multicast_skinny_gemm.config_for_m( + num_tokens, + shard_size, + ) + ] + skinny_results = [] + for config in configs: + skinny_kernel = SkinnyKernel(num_tokens, shard_size, config) + + def launch_skinny( + latent: torch.Tensor, + weight: torch.Tensor, + shared_shard: torch.Tensor, + *, + skinny_kernel: SkinnyKernel = skinny_kernel, + num_tokens: int = num_tokens, + ) -> None: + skinny_kernel.launch( + latent, + weight, + shared_shard[:num_tokens], + mailbox, + ) + + launches = make_up_projection_launches( + launch_skinny, + latent, + weights, + shared_shard, + ) + graph = capture_up_projection_graph(launches) + timing = benchmark_up_projection_graph( + graph, + operations_per_replay=len(launches), + warmup_replays=args.warmup_replays, + samples=args.samples, + ) + check_up_projection_output( + mailbox[0, :num_tokens, :shard_size], + latent, + weights[-1], + shared_shard[:num_tokens], + ) + skinny_results.append( + { + "config": asdict(config), + **timing, + } + ) + result["skinny"] = skinny_results + results.append(result) + + properties = torch.cuda.get_device_properties(device) + report = { + "scope": "up-projection", + "device": properties.name, + "compute_capability": list(torch.cuda.get_device_capability(device)), + "tp_size": args.tp_size, + "shard_size": shard_size, + "weight_count": weight_count, + "cache_multiplier": args.cache_multiplier, + "warmup_replays": args.warmup_replays, + "samples": args.samples, + "results": results, + } + rendered = json.dumps(report, indent=2) + print(rendered, flush=True) + if args.output is not None: + args.output.parent.mkdir(parents=True, exist_ok=True) + args.output.write_text(rendered + "\n", encoding="utf-8") + + +def capture_tail_graph( + operation: Callable[[], torch.Tensor], + cpu_group: dist.ProcessGroup, +) -> tuple[torch.cuda.CUDAGraph, torch.Tensor]: + for _ in range(3): + dist.barrier(group=cpu_group) + output = operation() + torch.accelerator.synchronize() + + dist.barrier(group=cpu_group) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + output = operation() + torch.accelerator.synchronize() + return graph, output + + +def benchmark_tail_graph( + graph: torch.cuda.CUDAGraph, + *, + warmup_replays: int, + samples: int, + device_group: dist.ProcessGroup, + cpu_group: dist.ProcessGroup, +) -> dict[str, Any]: + for _ in range(warmup_replays): + graph.replay() + torch.accelerator.synchronize() + + dist.barrier(group=cpu_group) + starts = [torch.cuda.Event(enable_timing=True) for _ in range(samples + 1)] + ends = [torch.cuda.Event(enable_timing=True) for _ in range(samples + 1)] + for start, end in zip(starts, ends): + start.record() + graph.replay() + end.record() + torch.accelerator.synchronize() + + samples_us = torch.tensor( + [start.elapsed_time(end) * 1000.0 for start, end in zip(starts, ends)], + dtype=torch.float64, + device=torch.accelerator.current_device_index(), + ) + dist.all_reduce(samples_us, op=dist.ReduceOp.MAX, group=device_group) + return summarize(samples_us[1:].tolist()) + + +def make_inputs( + num_tokens: int, + rank: int, + device: torch.device, +) -> tuple[torch.Tensor, torch.Tensor]: + torch.manual_seed(20260726 + 100 * num_tokens + rank) + routed = torch.randn( + (num_tokens, LATENT_SIZE), + dtype=torch.bfloat16, + device=device, + ).mul_(0.01) + shared = torch.randn( + (num_tokens, HIDDEN_SIZE), + dtype=torch.bfloat16, + device=device, + ) + return routed, shared + + +def make_reference( + routed: torch.Tensor, + shared: torch.Tensor, + rms_weight: torch.Tensor, + up_weight: torch.Tensor, + device_group: dist.ProcessGroup, +) -> Callable[[], torch.Tensor]: + routed_workspace = torch.empty_like(routed) + shared_workspace = torch.empty_like(shared) + + def reference() -> torch.Tensor: + routed_workspace.copy_(routed) + dist.all_reduce(routed_workspace, group=device_group) + normalized = F.rms_norm( + routed_workspace, + (LATENT_SIZE,), + rms_weight, + RMS_EPS, + ) + projected = F.linear(normalized, up_weight) + shared_workspace.copy_(shared) + dist.all_reduce(shared_workspace, group=device_group) + return projected.add(shared_workspace) + + return reference + + +def check_fused_output( + fused_output: torch.Tensor, + reference: Callable[[], torch.Tensor], + cpu_group: dist.ProcessGroup, +) -> None: + dist.barrier(group=cpu_group) + expected = reference() + torch.testing.assert_close(fused_output, expected, atol=8e-2, rtol=3e-2) + + +def benchmark_whole_tail(args: argparse.Namespace) -> None: + if any(not 1 <= num_tokens <= 16 for num_tokens in args.num_tokens): + raise ValueError("--num-tokens values must be in [1, 16]") + if args.warmup_replays < 0 or args.samples <= 0: + raise ValueError("warmup replays must be nonnegative and samples positive") + if args.skinny_max_num_tokens is not None and any(not 0 <= cutoff <= 8 for cutoff in args.skinny_max_num_tokens): + raise ValueError("--skinny-max-num-tokens must be in [0, 8]") + skinny_configs = dict(args.skinny_config or ()) + if len(skinny_configs) != len(args.skinny_config or ()): + raise ValueError("--skinny-config must not repeat an M value") + if any(not 1 <= num_tokens <= 8 for num_tokens in skinny_configs): + raise ValueError("--skinny-config M values must be in [1, 8]") + if not {"RANK", "WORLD_SIZE", "LOCAL_RANK"} <= os.environ.keys(): + raise RuntimeError("launch this benchmark with torchrun") + + rank = int(os.environ["RANK"]) + world_size = int(os.environ["WORLD_SIZE"]) + local_rank = int(os.environ["LOCAL_RANK"]) + device = torch.device("cuda", local_rank) + torch.accelerator.set_device_index(device) + init_distributed_environment() + if world_size > 8: + set_custom_all_reduce(False) + initialize_model_parallel(tensor_model_parallel_size=world_size) + device_group = get_tp_group().device_group + cpu_group = dist.new_group(backend="gloo") + + if torch.cuda.get_device_capability(device)[0] != 10: + raise RuntimeError("Kimi K3 latent-MoE tail requires SM100") + + torch.manual_seed(20260726) + rms_weight = 1 + 0.1 * torch.randn( + LATENT_SIZE, + dtype=torch.bfloat16, + device=device, + ) + up_weight = ( + torch.randn( + (HIDDEN_SIZE, LATENT_SIZE), + dtype=torch.bfloat16, + device=device, + ) + / LATENT_SIZE**0.5 + ) + + use_reference = args.backend in ("reference", "both") + use_fused = args.backend in ("fused", "both") + fused_ops = [] + if use_fused: + production_config_for_m = fused_add_multicast_skinny_gemm.config_for_m + + def config_for_m( + num_rows: int, + shard_dim: int = 896, + ) -> fused_add_multicast_skinny_gemm.SkinnyConfig: + config = skinny_configs.get(num_rows) + if config is not None: + return config + return production_config_for_m(num_rows, shard_dim) + + fused_add_multicast_skinny_gemm.config_for_m = config_for_m + cutoffs = args.skinny_max_num_tokens or [latent_moe_tail._SKINNY_MAX_NUM_TOKENS] + for cutoff in cutoffs: + latent_moe_tail._SKINNY_MAX_NUM_TOKENS = cutoff + latent_moe_tail.KimiK3LatentMoETailOp._instances.clear() + fused_ops.append( + ( + cutoff, + latent_moe_tail.KimiK3LatentMoETailOp.initialize( + hidden_size=HIDDEN_SIZE, + latent_size=LATENT_SIZE, + dtype=torch.bfloat16, + device=device, + rms_eps=RMS_EPS, + ), + ) + ) + cutedsl_warmup() + + results = [] + for num_tokens in args.num_tokens: + routed, shared = make_inputs(num_tokens, rank, device) + reference = make_reference( + routed, + shared, + rms_weight, + up_weight, + device_group, + ) + result: dict[str, Any] = {"num_tokens": num_tokens} + if use_reference: + reference_graph, _ = capture_tail_graph(reference, cpu_group) + result["reference"] = benchmark_tail_graph( + reference_graph, + warmup_replays=args.warmup_replays, + samples=args.samples, + device_group=device_group, + cpu_group=cpu_group, + ) + for cutoff, fused_op in fused_ops: + + def fused( + routed: torch.Tensor = routed, + shared: torch.Tensor = shared, + fused_op: latent_moe_tail.KimiK3LatentMoETailOp = fused_op, + ) -> torch.Tensor: + return fused_op(routed, shared, rms_weight, up_weight) + + fused_graph, fused_output = capture_tail_graph(fused, cpu_group) + fused_key = "fused" if len(fused_ops) == 1 else f"fused_skinny_max_{cutoff}" + result[fused_key] = benchmark_tail_graph( + fused_graph, + warmup_replays=args.warmup_replays, + samples=args.samples, + device_group=device_group, + cpu_group=cpu_group, + ) + check_fused_output(fused_output, reference, cpu_group) + if "reference" in result: + speedup = result["reference"]["median_us"] / result[fused_key]["median_us"] + if len(fused_ops) == 1: + result["speedup"] = speedup + else: + result[f"{fused_key}_speedup"] = speedup + results.append(result) + + properties = torch.cuda.get_device_properties(device) + report = { + "scope": "whole-tail", + "device": properties.name, + "compute_capability": list(torch.cuda.get_device_capability(device)), + "world_size": world_size, + "torch_version": torch.__version__, + "cuda_version": torch.version.cuda, + "warmup_replays": args.warmup_replays, + "samples": args.samples, + "skinny_max_num_tokens": [cutoff for cutoff, _ in fused_ops], + "skinny_configs": {str(num_tokens): asdict(config) for num_tokens, config in skinny_configs.items()}, + "timing_scope": { + "reference": ( + "two input copies, two AllReduces, RMSNorm, full replicated up-projection GEMM, and final add" + ), + "fused": ( + "routed AllReduce/RMSNorm plus shared ReduceScatter, sharded up-projection/multicast, and Lamport copy" + ), + }, + "results": results, + } + if rank == 0: + rendered = json.dumps(report, indent=2) + print(rendered, flush=True) + if args.output is not None: + args.output.parent.mkdir(parents=True, exist_ok=True) + args.output.write_text(rendered + "\n", encoding="utf-8") + + dist.barrier(group=cpu_group) + + +def main() -> None: + args = parse_args() + if args.scope == "up-projection": + benchmark_up_projection(args) + return + + from aphrodite.config import AphroditeConfig, set_current_aphrodite_config + + with set_current_aphrodite_config(AphroditeConfig()): + benchmark_whole_tail(args) + + +if __name__ == "__main__": + main() diff --git a/benchmarks/kernels/benchmark_kimi_k3_sp_collectives.py b/benchmarks/kernels/benchmark_kimi_k3_sp_collectives.py new file mode 100644 index 0000000000..ed35efa2c2 --- /dev/null +++ b/benchmarks/kernels/benchmark_kimi_k3_sp_collectives.py @@ -0,0 +1,239 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import argparse +import json +import os +import statistics +from collections.abc import Callable + +import torch +import torch.distributed as dist + +import aphrodite._custom_ops as ops +from aphrodite.distributed.device_communicators.custom_all_reduce import CustomAllreduce + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser() + parser.add_argument("--tokens", type=int, nargs="+", default=[8, 32, 128, 1024]) + parser.add_argument("--hidden-size", type=int, default=7168) + parser.add_argument("--graph-repeats", type=int, default=20) + parser.add_argument("--warmup-replays", type=int, default=5) + parser.add_argument("--samples", type=int, default=15) + return parser.parse_args() + + +def capture_graph(op: Callable[[], None], repeats: int) -> torch.cuda.CUDAGraph: + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + for _ in range(3): + op() + stream.synchronize() + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, stream=stream): + for _ in range(repeats): + op() + torch.cuda.current_stream().wait_stream(stream) + return graph + + +def max_rank_graph_time( + graph: torch.cuda.CUDAGraph, + repeats: int, + warmup_replays: int, + samples: int, + device_group: dist.ProcessGroup, + cpu_group: dist.ProcessGroup, +) -> float: + for _ in range(warmup_replays): + graph.replay() + torch.accelerator.synchronize() + + timings = [] + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + for _ in range(samples): + dist.barrier(group=cpu_group) + start.record() + graph.replay() + end.record() + end.synchronize() + elapsed = torch.tensor( + start.elapsed_time(end) / repeats, + dtype=torch.float64, + device=torch.accelerator.current_device_index(), + ) + dist.all_reduce(elapsed, op=dist.ReduceOp.MAX, group=device_group) + timings.append(elapsed.item()) + return statistics.median(timings) + + +def check_outputs( + comm: CustomAllreduce, + local: torch.Tensor, + reduce_input: torch.Tensor, + device_group: dist.ProcessGroup, +) -> None: + expected_gather = torch.empty( + (local.shape[0] * dist.get_world_size(), local.shape[1]), + dtype=local.dtype, + device=local.device, + ) + dist.all_gather_into_tensor(expected_gather, local, group=device_group) + gathered = comm.custom_all_gather(local) + assert gathered is not None + torch.testing.assert_close(gathered, expected_gather) + + expected_scatter = torch.empty_like(local) + dist.reduce_scatter_tensor( + expected_scatter, + reduce_input.clone(), + group=device_group, + ) + scattered = comm.custom_reduce_scatter(reduce_input) + assert scattered is not None + torch.testing.assert_close(scattered, expected_scatter) + + +def benchmark_shape( + comm: CustomAllreduce, + global_tokens: int, + hidden_size: int, + graph_repeats: int, + warmup_replays: int, + samples: int, + device_group: dist.ProcessGroup, + cpu_group: dist.ProcessGroup, +) -> dict[str, float | int]: + world_size = dist.get_world_size() + rank = dist.get_rank() + padded_tokens = (global_tokens + world_size - 1) // world_size * world_size + local_tokens = padded_tokens // world_size + local = torch.full( + (local_tokens, hidden_size), + rank + 1, + dtype=torch.bfloat16, + device=torch.accelerator.current_device_index(), + ) + reduce_input = torch.full( + (padded_tokens, hidden_size), + rank + 1, + dtype=torch.bfloat16, + device=local.device, + ) + check_outputs(comm, local, reduce_input, device_group) + + custom_gather_out = torch.empty( + (padded_tokens, hidden_size), + dtype=local.dtype, + device=local.device, + ) + custom_scatter_out = torch.empty_like(local) + nccl_gather_out = torch.empty_like(custom_gather_out) + nccl_scatter_out = torch.empty_like(local) + + def custom_ag() -> None: + ops.mnnvl_lamport_all_gather( + comm._ptr, + local, + custom_gather_out, + comm.mnnvl_lamport_ag_local_ptr, + comm.mnnvl_lamport_ag_multicast_ptr, + comm.mnnvl_lamport_ag_epoch_ptr, + comm.mnnvl_buffer_size, + ) + + def custom_rs() -> None: + ops.mnnvl_lamport_reduce_scatter( + comm._ptr, + reduce_input, + custom_scatter_out, + comm.mnnvl_lamport_rs_local_ptr, + comm.mnnvl_lamport_rs_epoch_ptr, + comm.mnnvl_buffer_size, + ) + + def nccl_ag() -> None: + dist.all_gather_into_tensor(nccl_gather_out, local, group=device_group) + + def nccl_rs() -> None: + dist.reduce_scatter_tensor( + nccl_scatter_out, + reduce_input, + group=device_group, + ) + + graphs = { + "custom_ag_us": capture_graph(custom_ag, graph_repeats), + "nccl_ag_us": capture_graph(nccl_ag, graph_repeats), + "custom_rs_us": capture_graph(custom_rs, graph_repeats), + "nccl_rs_us": capture_graph(nccl_rs, graph_repeats), + } + times = { + name: max_rank_graph_time( + graph, + graph_repeats, + warmup_replays, + samples, + device_group, + cpu_group, + ) + * 1000 + for name, graph in graphs.items() + } + torch.testing.assert_close(custom_gather_out, nccl_gather_out) + torch.testing.assert_close(custom_scatter_out, nccl_scatter_out) + return { + "global_tokens": global_tokens, + "padded_tokens": padded_tokens, + "local_bytes": local.nbytes, + "full_bytes": reduce_input.nbytes, + **times, + "ag_speedup": times["nccl_ag_us"] / times["custom_ag_us"], + "rs_speedup": times["nccl_rs_us"] / times["custom_rs_us"], + } + + +def main() -> None: + args = parse_args() + local_rank = int(os.environ["LOCAL_RANK"]) + torch.accelerator.set_device_index(local_rank) + dist.init_process_group("nccl") + device_group = dist.group.WORLD + cpu_group = dist.new_group(backend="gloo") + + comm = CustomAllreduce( + group=cpu_group, + device=torch.device("cuda", local_rank), + ) + assert not comm.disabled + assert comm.world_size == 16 + assert comm.mnnvl_only + assert comm.mnnvl_multicast_ptr + + results = [ + benchmark_shape( + comm, + tokens, + args.hidden_size, + args.graph_repeats, + args.warmup_replays, + args.samples, + device_group, + cpu_group, + ) + for tokens in args.tokens + ] + if dist.get_rank() == 0: + print(json.dumps(results, indent=2), flush=True) + + comm.close() + dist.destroy_process_group(cpu_group) + dist.destroy_process_group() + + +if __name__ == "__main__": + main() diff --git a/cmake/external_projects/flashkda.cmake b/cmake/external_projects/flashkda.cmake new file mode 100644 index 0000000000..1e8cdfe455 --- /dev/null +++ b/cmake/external_projects/flashkda.cmake @@ -0,0 +1,74 @@ +include(FetchContent) + +if(DEFINED ENV{FLASH_KDA_SRC_DIR}) + set(FLASH_KDA_SRC_DIR $ENV{FLASH_KDA_SRC_DIR}) +endif() + +if(FLASH_KDA_SRC_DIR) + FetchContent_Declare( + flashkda + SOURCE_DIR ${FLASH_KDA_SRC_DIR} + ) +else() + FetchContent_Declare( + flashkda + GIT_REPOSITORY https://github.com/vllm-project/FlashKDA.git + GIT_TAG a3e42bbbece3bb38f7c426b880315294a336e82f + GIT_PROGRESS TRUE + GIT_SUBMODULES cutlass + ) +endif() + +FetchContent_MakeAvailable(flashkda) +message(STATUS "FlashKDA is available at ${flashkda_SOURCE_DIR}") + +set(FLASH_KDA_SUPPORT_ARCHS) +if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.0) + list(APPEND FLASH_KDA_SUPPORT_ARCHS "9.0a") +endif() +if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0) + list(APPEND FLASH_KDA_SUPPORT_ARCHS "10.0f" "12.0f") +elseif(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.9) + list(APPEND FLASH_KDA_SUPPORT_ARCHS "10.0a" "10.3a" "12.0a") +endif() + +cuda_archs_loose_intersection( + FLASH_KDA_ARCHS "${FLASH_KDA_SUPPORT_ARCHS}" "${CUDA_ARCHS}") + +if(FLASH_KDA_ARCHS) + message(STATUS "FlashKDA CUDA architectures: ${FLASH_KDA_ARCHS}") + + set(FLASH_KDA_SOURCES + csrc/flashkda_registration.cpp + ${flashkda_SOURCE_DIR}/csrc/flash_kda.cpp + ${flashkda_SOURCE_DIR}/csrc/smxx/fwd_launch.cu) + set(FLASH_KDA_INCLUDES + ${flashkda_SOURCE_DIR}/csrc + ${flashkda_SOURCE_DIR}/cutlass/include + ${flashkda_SOURCE_DIR}/cutlass/examples/common + ${flashkda_SOURCE_DIR}/cutlass/tools/util/include) + + set_gencode_flags_for_srcs( + SRCS "${FLASH_KDA_SOURCES}" + CUDA_ARCHS "${FLASH_KDA_ARCHS}") + + define_extension_target( + _flashkda_C + DESTINATION vllm + LANGUAGE ${APHRODITE_GPU_LANG} + SOURCES ${FLASH_KDA_SOURCES} + COMPILE_FLAGS ${APHRODITE_GPU_FLAGS} + ARCHITECTURES ${APHRODITE_GPU_ARCHES} + INCLUDE_DIRECTORIES ${FLASH_KDA_INCLUDES} + USE_SABI 3 + WITH_SOABI) + + target_compile_options(_flashkda_C PRIVATE + $<$:-UPy_LIMITED_API --expt-relaxed-constexpr --expt-extended-lambda --use_fast_math -O3> + $<$:-UPy_LIMITED_API>) +else() + message(STATUS + "FlashKDA will not compile: CUDA >=12.0 and a supported architecture " + "(SM90, SM10x, or SM12x) are required") + add_custom_target(_flashkda_C) +endif() diff --git a/csrc/cpu/cpu_attn.cpp b/csrc/cpu/cpu_attn.cpp index 518928976d..6843e19436 100644 --- a/csrc/cpu/cpu_attn.cpp +++ b/csrc/cpu/cpu_attn.cpp @@ -30,7 +30,8 @@ torch::Tensor get_scheduler_metadata( const torch::Tensor& query_start_loc, const bool causal, const int64_t window_size, const std::string& isa_hint, const bool enable_kv_split, - const std::optional& dynamic_causal) { + const std::optional& dynamic_causal, + const std::string& kv_cache_dtype) { cpu_attention::ISA isa; if (isa_hint == "amx") { isa = cpu_attention::ISA::AMX; @@ -65,9 +66,11 @@ torch::Tensor get_scheduler_metadata( input.dynamic_causal = dynamic_causal.has_value() ? dynamic_causal->data_ptr() : nullptr; + const int64_t kv_cache_idx = + static_cast(parse_fp8_kv_dtype(kv_cache_dtype)); APHRODITE_DISPATCH_FLOATING_TYPES(dtype, "get_scheduler_metadata", [&]() { - CPU_ATTN_DISPATCH(head_dim, isa, 0, [&]() { - input.elem_size = sizeof(scalar_t); + CPU_ATTN_DISPATCH(head_dim, isa, kv_cache_idx, [&]() { + input.elem_size = sizeof(attn_impl::kv_cache_t); input.q_buffer_elem_size = sizeof(attn_impl::q_buffer_t); input.logits_buffer_elem_size = sizeof(attn_impl::logits_buffer_t); input.output_buffer_elem_size = diff --git a/csrc/cpu/cpu_types_vxe.hpp b/csrc/cpu/cpu_types_vxe.hpp index 641fceeffe..6f151a5247 100644 --- a/csrc/cpu/cpu_types_vxe.hpp +++ b/csrc/cpu/cpu_types_vxe.hpp @@ -270,7 +270,7 @@ struct FP32Vec4 : public Vec { explicit FP32Vec4(__vector float data) : reg(data) {} - explicit FP32Vec4(const FP32Vec4& data) : reg(data.reg) {} + FP32Vec4(const FP32Vec4& data) : reg(data.reg) {} }; struct FP32Vec8 : public Vec { @@ -299,7 +299,7 @@ struct FP32Vec8 : public Vec { explicit FP32Vec8(f32x4x2_t data) : reg(data) {} - explicit FP32Vec8(const FP32Vec8& data) { + FP32Vec8(const FP32Vec8& data) { reg.val[0] = data.reg.val[0]; reg.val[1] = data.reg.val[1]; } @@ -644,7 +644,7 @@ struct FP32Vec16 : public Vec { explicit FP32Vec16(f32x4x4_t data) : reg(data) {} - explicit FP32Vec16(const FP32Vec16& data) { + FP32Vec16(const FP32Vec16& data) { reg.val[0] = data.reg.val[0]; reg.val[1] = data.reg.val[1]; reg.val[2] = data.reg.val[2]; @@ -1198,4 +1198,4 @@ FORCE_INLINE void prefetch(const void* addr) { }; // namespace vec_op -#endif \ No newline at end of file +#endif diff --git a/csrc/cpu/torch_bindings.cpp b/csrc/cpu/torch_bindings.cpp index 80b9409d99..ee8e37110e 100644 --- a/csrc/cpu/torch_bindings.cpp +++ b/csrc/cpu/torch_bindings.cpp @@ -163,7 +163,8 @@ torch::Tensor get_scheduler_metadata( const torch::Tensor& query_start_loc, const bool casual, const int64_t window_size, const std::string& isa_hint, const bool enable_kv_split, - const std::optional& dynamic_causal); + const std::optional& dynamic_causal, + const std::string& kv_cache_dtype); void cpu_attn_reshape_and_cache(const torch::Tensor& key, const torch::Tensor& value, @@ -577,7 +578,8 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) { "get_scheduler_metadata(int num_req, int num_heads_q, int num_heads_kv, " "int head_dim, Tensor seq_lens, ScalarType dtype, Tensor " "query_start_loc, bool casual, int window_size, str isa_hint, bool " - "enable_kv_split, Tensor? dynamic_causal) -> Tensor", + "enable_kv_split, Tensor? dynamic_causal, " + "str kv_cache_dtype=\"auto\") -> Tensor", &get_scheduler_metadata); ops.def( "cpu_attn_reshape_and_cache(Tensor key, Tensor value, Tensor(a2!) " diff --git a/csrc/custom_all_gather_reduce_scatter.cuh b/csrc/custom_all_gather_reduce_scatter.cuh new file mode 100644 index 0000000000..0f60a31b39 --- /dev/null +++ b/csrc/custom_all_gather_reduce_scatter.cuh @@ -0,0 +1,327 @@ +#pragma once + +#include "custom_collective_common.cuh" + +namespace aphrodite { + +constexpr int kMnnvlLamportAgThreads = 128; +constexpr int kMnnvlLamportRsThreads = 256; +constexpr int kMnnvlLamportConcurrentPollMaxPacks = 8192; + +using CopyPack = array_t; + +template +__global__ void __launch_bounds__(512, 1) + cross_device_all_gather(RankData* _dp, RankSignals sg, Signal* self_sg, + CopyPack* __restrict__ result, int rank, + int size_per_rank) { + auto dp = *_dp; + int tid = blockIdx.x * blockDim.x + threadIdx.x; + int stride = gridDim.x * blockDim.x; + barrier_at_start(sg, self_sg, rank); +#pragma unroll + for (int src_rank = 0; src_rank < ngpus; ++src_rank) { + auto src = reinterpret_cast(dp.ptrs[src_rank]); + auto dst = result + src_rank * size_per_rank; + for (int idx = tid; idx < size_per_rank; idx += stride) { + dst[idx] = src[idx]; + } + } + barrier_at_end(sg, self_sg, rank); +} + +template +__global__ void __launch_bounds__(512, 1) + cross_device_reduce_scatter(RankData* _dp, RankSignals sg, Signal* self_sg, + T* __restrict__ result, int rank, + int size_per_rank) { + using P = typename packed_t::P; + using A = typename packed_t::A; + auto dp = *_dp; + auto offset = rank * size_per_rank; + barrier_at_start(sg, self_sg, rank); + for (int idx = blockIdx.x * blockDim.x + threadIdx.x; idx < size_per_rank; + idx += gridDim.x * blockDim.x) { + reinterpret_cast(result)[idx] = + packed_reduce((const P**)&dp.ptrs[0], offset + idx); + } + barrier_at_end(sg, self_sg, rank); +} + +template +union LamportPack { + P packed; + uint32_t words[sizeof(P) / sizeof(uint32_t)]; +}; + +template +DINLINE LamportPack

load_lamport_pack(const P* ptr) { + static_assert(sizeof(P) == 16); + LamportPack

value; +#if !defined(USE_ROCM) + asm volatile("ld.volatile.global.v4.u32 {%0, %1, %2, %3}, [%4];" + : "=r"(value.words[0]), "=r"(value.words[1]), + "=r"(value.words[2]), "=r"(value.words[3]) + : "l"(ptr) + : "memory"); +#else + const volatile uint32_t* src = + reinterpret_cast(ptr); + #pragma unroll + for (int i = 0; i < sizeof(P) / sizeof(uint32_t); ++i) { + value.words[i] = src[i]; + } +#endif + return value; +} + +template +DINLINE bool is_lamport_dirty(const LamportPack

& value) { +#pragma unroll + for (int i = 0; i < sizeof(P) / sizeof(uint32_t); ++i) { + if (value.words[i] == 0x80000000U) return true; + } + return false; +} + +template +DINLINE P lamport_sentinel() { + LamportPack

value; +#pragma unroll + for (int i = 0; i < sizeof(P) / sizeof(uint32_t); ++i) { + value.words[i] = 0x80000000U; + } + return value.packed; +} + +template +DINLINE P sanitize_lamport_payload(P packed) { + LamportPack

value{.packed = packed}; +#pragma unroll + for (int i = 0; i < sizeof(P) / sizeof(uint32_t); ++i) { + if (value.words[i] == 0x80000000U) value.words[i] = 0; + } + return value.packed; +} + +template +DINLINE P wait_lamport_payload(const P* ptr) { + auto value = load_lamport_pack(ptr); + while (is_lamport_dirty(value)) value = load_lamport_pack(ptr); + return value.packed; +} + +template +DINLINE void wait_lamport_payloads(const P* base, int rank, int rank_stride, + P local_value, P (&values)[ngpus]) { + bool ready[ngpus]; +#pragma unroll + for (int src_rank = 0; src_rank < ngpus; ++src_rank) { + ready[src_rank] = src_rank == rank; + if (src_rank == rank) values[src_rank] = local_value; + } + + int remaining = ngpus - 1; + while (remaining != 0) { +#pragma unroll + for (int src_rank = 0; src_rank < ngpus; ++src_rank) { + if (!ready[src_rank]) { + auto value = load_lamport_pack(base + src_rank * rank_stride); + if (!is_lamport_dirty(value)) { + values[src_rank] = value.packed; + ready[src_rank] = true; + --remaining; + } + } + } + } +} + +template +DINLINE P reduce_lamport_payloads(const P* current_local, const P* packed_input, + int rank, int size_per_rank, int idx) { + P source_zero = + rank == 0 ? packed_input[idx] : wait_lamport_payload(current_local + idx); + A tmp = upcast(source_zero); +#pragma unroll + for (int src_rank = 1; src_rank < ngpus; ++src_rank) { + P value = src_rank == rank + ? packed_input[rank * size_per_rank + idx] + : wait_lamport_payload(current_local + + src_rank * size_per_rank + idx); + packed_assign_add(tmp, upcast(value)); + } + return sanitize_lamport_payload(downcast

(tmp)); +} + +DINLINE void lamport_cta_arrive(uint32_t* counter) { +#if !defined(USE_ROCM) + if (threadIdx.x < 32) { + asm volatile("barrier.cta.sync 1, %0;" : : "r"(blockDim.x) : "memory"); + if (threadIdx.x == 0) { + #if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 1000 + asm volatile("red.async.release.global.gpu.add.u32 [%0], 1;" + : + : "l"(counter) + : "memory"); + #elif defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 700 + asm volatile("red.release.global.gpu.add.u32 [%0], 1;" + : + : "l"(counter) + : "memory"); + #else + atomicAdd(counter, 1); + #endif + } + } else { + asm volatile("barrier.cta.arrive 1, %0;" : : "r"(blockDim.x) : "memory"); + } +#else + __syncthreads(); + if (threadIdx.x == 0) atomicAdd(counter, 1); +#endif +} + +template +__global__ void __launch_bounds__(kMnnvlLamportAgThreads, 1) + mnnvl_lamport_all_gather(RankData* _dp, const T* __restrict__ input, + T* __restrict__ result, + T* __restrict__ multicast_buffer, + uint32_t* __restrict__ epochs, int rank, + int size_per_rank, int stage_size) { + using P = typename packed_t::P; +#if !defined(USE_ROCM) && CUDA_VERSION >= 12000 && defined(__CUDA_ARCH__) && \ + (__CUDA_ARCH__ >= 900) + cudaGridDependencySynchronize(); +#endif + auto dp = *_dp; + int tid = blockIdx.x * blockDim.x + threadIdx.x; + int stride = gridDim.x * blockDim.x; + uint32_t epoch = epochs[0]; + int current_stage = epoch % 3; + int dirty_stage = (epoch + 1) % 3; + int dirty_size = epochs[2 + dirty_stage]; + auto local_buffer = reinterpret_cast(const_cast(dp.ptrs[rank])); + auto current_local = local_buffer + current_stage * stage_size; + auto dirty_local = local_buffer + dirty_stage * stage_size; + auto current_multicast = + reinterpret_cast(multicast_buffer) + current_stage * stage_size; + auto packed_input = reinterpret_cast(input); + auto packed_result = reinterpret_cast(result); + + int total_size = size_per_rank * ngpus; + P local_value; + if (tid < size_per_rank) { + local_value = packed_input[tid]; + current_multicast[rank * size_per_rank + tid] = + sanitize_lamport_payload(local_value); + } +#if !defined(USE_ROCM) && CUDA_VERSION >= 12000 && defined(__CUDA_ARCH__) && \ + (__CUDA_ARCH__ >= 900) + cudaTriggerProgrammaticLaunchCompletion(); +#endif + + lamport_cta_arrive(&epochs[1]); + + for (int idx = tid; idx < dirty_size; idx += stride) { + dirty_local[idx] = lamport_sentinel

(); + } + + if (tid < size_per_rank) { +#pragma unroll + for (int src_rank = 0; src_rank < ngpus; ++src_rank) { + int output_idx = src_rank * size_per_rank + tid; + P value = src_rank == rank + ? local_value + : wait_lamport_payload(current_local + output_idx); + packed_result[output_idx] = value; + } + } + + if (tid == 0) { + while (*reinterpret_cast(&epochs[1]) < gridDim.x); + epochs[2 + current_stage] = total_size; + epochs[0] = epoch + 1; + epochs[1] = 0; + } +} + +template +__global__ void __launch_bounds__(kMnnvlLamportRsThreads, 1) + mnnvl_lamport_reduce_scatter_kernel(RankData* _dp, + const T* __restrict__ input, + T* __restrict__ result, + uint32_t* __restrict__ epochs, int rank, + int size_per_rank, int stage_size) { + using P = typename packed_t::P; + using A = typename packed_t::A; +#if !defined(USE_ROCM) && CUDA_VERSION >= 12000 && defined(__CUDA_ARCH__) && \ + (__CUDA_ARCH__ >= 900) + cudaGridDependencySynchronize(); +#endif + auto dp = *_dp; + int dst_rank = blockIdx.x % ngpus; + int tile = blockIdx.x / ngpus; + int idx = tile * blockDim.x + threadIdx.x; + int tid = blockIdx.x * blockDim.x + threadIdx.x; + int stride = gridDim.x * blockDim.x; + uint32_t epoch = epochs[0]; + int current_stage = epoch % 3; + int dirty_stage = (epoch + 1) % 3; + int dirty_size = epochs[2 + dirty_stage]; + auto local_buffer = reinterpret_cast(const_cast(dp.ptrs[rank])); + auto current_local = local_buffer + current_stage * stage_size; + auto dirty_local = local_buffer + dirty_stage * stage_size; + auto packed_input = reinterpret_cast(input); + + if (idx < size_per_rank && dst_rank != rank) { + auto dst = reinterpret_cast(const_cast(dp.ptrs[dst_rank])) + + current_stage * stage_size + rank * size_per_rank; + auto src = packed_input + dst_rank * size_per_rank; + dst[idx] = sanitize_lamport_payload(src[idx]); + } +#if !defined(USE_ROCM) && CUDA_VERSION >= 12000 && defined(__CUDA_ARCH__) && \ + (__CUDA_ARCH__ >= 900) + cudaTriggerProgrammaticLaunchCompletion(); +#endif + + lamport_cta_arrive(&epochs[1]); + + for (int idx = tid; idx < dirty_size; idx += stride) { + dirty_local[idx] = lamport_sentinel

(); + } + + if (idx < size_per_rank && dst_rank == rank) { + if constexpr (ngpus == 4) { + if (size_per_rank > kMnnvlLamportConcurrentPollMaxPacks) { + reinterpret_cast(result)[idx] = + reduce_lamport_payloads(current_local, packed_input, + rank, size_per_rank, idx); + } else { + P values[ngpus]; + wait_lamport_payloads( + current_local + idx, rank, size_per_rank, + packed_input[rank * size_per_rank + idx], values); + A tmp = upcast(values[0]); +#pragma unroll + for (int src_rank = 1; src_rank < ngpus; ++src_rank) { + packed_assign_add(tmp, upcast(values[src_rank])); + } + reinterpret_cast(result)[idx] = + sanitize_lamport_payload(downcast

(tmp)); + } + } else { + reinterpret_cast(result)[idx] = reduce_lamport_payloads( + current_local, packed_input, rank, size_per_rank, idx); + } + } + + if (tid == 0) { + while (*reinterpret_cast(&epochs[1]) < gridDim.x); + epochs[2 + current_stage] = size_per_rank * ngpus; + epochs[0] = epoch + 1; + epochs[1] = 0; + } +} + +} // namespace aphrodite diff --git a/csrc/custom_all_reduce.cuh b/csrc/custom_all_reduce.cuh index a08b25820e..76198ef83d 100644 --- a/csrc/custom_all_reduce.cuh +++ b/csrc/custom_all_reduce.cuh @@ -1,304 +1,8 @@ #pragma once -#include -#include -#include -#include - -#if defined(USE_ROCM) -typedef __hip_bfloat16 nv_bfloat16; -#endif - -#include -#include -#include -#include -#include -#include -#include -#include - -namespace aphrodite { -#define CUDACHECK(cmd) \ - do { \ - cudaError_t e = cmd; \ - if (e != cudaSuccess) { \ - printf("Failed: Cuda error %s:%d '%s'\n", __FILE__, __LINE__, \ - cudaGetErrorString(e)); \ - exit(EXIT_FAILURE); \ - } \ - } while (0) - -// Maximal number of blocks in allreduce kernel. -constexpr int kMaxBlocks = 36; - -// Default number of blocks in allreduce kernel. -#ifndef USE_ROCM -const int defaultBlockLimit = 36; -CUpointer_attribute rangeStartAddrAttr = CU_POINTER_ATTRIBUTE_RANGE_START_ADDR; -#else -const int defaultBlockLimit = 16; -hipPointer_attribute rangeStartAddrAttr = - HIP_POINTER_ATTRIBUTE_RANGE_START_ADDR; -#endif - -// Counter may overflow, but it's fine since unsigned int overflow is -// well-defined behavior. -using FlagType = uint32_t; - -// Two sets of peer counters are needed for two syncs: starting and ending an -// operation. The reason is that it's possible for peer GPU block to arrive at -// the second sync point while the current GPU block haven't passed the first -// sync point. Thus, peer GPU may write counter+1 while current GPU is busy -// waiting for counter. We use alternating counter array to avoid this -// possibility. -struct Signal { - alignas(128) FlagType start[kMaxBlocks][8]; - alignas(128) FlagType end[kMaxBlocks][8]; - alignas(128) FlagType _flag[kMaxBlocks]; // incremental flags for each rank -}; - -struct __align__(16) RankData { - const void* ptrs[8]; -}; - -struct __align__(16) RankSignals { - Signal* signals[8]; -}; - -// like std::array, but aligned -template -struct __align__(alignof(T) * sz) array_t { - T data[sz]; - using type = T; - static constexpr int size = sz; -}; - -// use packed type to maximize memory efficiency -// goal: generate ld.128 and st.128 instructions -template -struct packed_t { - // the (P)acked type for load/store - using P = array_t; - // the (A)ccumulator type for reduction - using A = array_t; -}; - -#define DINLINE __device__ __forceinline__ - -// scalar cast functions -DINLINE float upcast_s(half val) { return __half2float(val); } - -template -DINLINE T downcast_s(float val); -template <> -DINLINE half downcast_s(float val) { - return __float2half(val); -} - -// scalar add functions -// for some reason when compiling with Pytorch, the + operator for half and -// bfloat is disabled so we call the intrinsics directly -DINLINE half& assign_add(half& a, half b) { - a = __hadd(a, b); - return a; -} -DINLINE float& assign_add(float& a, float b) { return a += b; } - -#if (__CUDA_ARCH__ >= 800 || !defined(__CUDA_ARCH__)) -DINLINE float upcast_s(nv_bfloat16 val) { return __bfloat162float(val); } -template <> -DINLINE nv_bfloat16 downcast_s(float val) { - return __float2bfloat16(val); -} -DINLINE nv_bfloat16& assign_add(nv_bfloat16& a, nv_bfloat16 b) { - a = __hadd(a, b); - return a; -} -#endif - -template -DINLINE array_t& packed_assign_add(array_t& a, array_t b) { -#pragma unroll - for (int i = 0; i < N; i++) { - assign_add(a.data[i], b.data[i]); - } - return a; -} - -template -DINLINE array_t upcast(array_t val) { - if constexpr (std::is_same::value) { - return val; - } else { - array_t out; -#pragma unroll - for (int i = 0; i < N; i++) { - out.data[i] = upcast_s(val.data[i]); - } - return out; - } -} - -template -DINLINE O downcast(array_t val) { - if constexpr (std::is_same::value) { - return val; - } else { - O out; -#pragma unroll - for (int i = 0; i < O::size; i++) { - out.data[i] = downcast_s(val.data[i]); - } - return out; - } -} - -#if !defined(USE_ROCM) - -static DINLINE void st_flag_release(FlagType* flag_addr, FlagType flag) { - #if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 700 - asm volatile("st.release.sys.global.u32 [%1], %0;" ::"r"(flag), - "l"(flag_addr)); - #else - asm volatile("membar.sys; st.volatile.global.u32 [%1], %0;" ::"r"(flag), - "l"(flag_addr)); - #endif -} - -static DINLINE FlagType ld_flag_acquire(FlagType* flag_addr) { - FlagType flag; - #if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 700 - asm volatile("ld.acquire.sys.global.u32 %0, [%1];" - : "=r"(flag) - : "l"(flag_addr)); - #else - asm volatile("ld.volatile.global.u32 %0, [%1]; membar.gl;" - : "=r"(flag) - : "l"(flag_addr)); - #endif - return flag; -} - -static DINLINE void st_flag_volatile(FlagType* flag_addr, FlagType flag) { - asm volatile("st.volatile.global.u32 [%1], %0;" ::"r"(flag), "l"(flag_addr)); -} - -static DINLINE FlagType ld_flag_volatile(FlagType* flag_addr) { - FlagType flag; - asm volatile("ld.volatile.global.u32 %0, [%1];" - : "=r"(flag) - : "l"(flag_addr)); - return flag; -} - -// This function is meant to be used as the first synchronization in the all -// reduce kernel. Thus, it doesn't need to make any visibility guarantees for -// prior memory accesses. Note: volatile writes will not be reordered against -// other volatile writes. -template -DINLINE void barrier_at_start(const RankSignals& sg, Signal* self_sg, - int rank) { - uint32_t flag = self_sg->_flag[blockIdx.x] + 1; - if (threadIdx.x < ngpus) { - auto peer_counter_ptr = &sg.signals[threadIdx.x]->start[blockIdx.x][rank]; - auto self_counter_ptr = &self_sg->start[blockIdx.x][threadIdx.x]; - // Write the expected counter value to peer and wait for correct value - // from peer. - st_flag_volatile(peer_counter_ptr, flag); - while (ld_flag_volatile(self_counter_ptr) != flag); - } - __syncthreads(); - // use one thread to update flag - if (threadIdx.x == 0) self_sg->_flag[blockIdx.x] = flag; -} - -// This function is meant to be used as the second or the final -// synchronization barrier in the all reduce kernel. If it's the final -// synchronization barrier, we don't need to make any visibility guarantees -// for prior memory accesses. -template -DINLINE void barrier_at_end(const RankSignals& sg, Signal* self_sg, int rank) { - __syncthreads(); - uint32_t flag = self_sg->_flag[blockIdx.x] + 1; - if (threadIdx.x < ngpus) { - auto peer_counter_ptr = &sg.signals[threadIdx.x]->end[blockIdx.x][rank]; - auto self_counter_ptr = &self_sg->end[blockIdx.x][threadIdx.x]; - // Write the expected counter value to peer and wait for correct value from - // peer. - if constexpr (!final_sync) { - st_flag_release(peer_counter_ptr, flag); - while (ld_flag_acquire(self_counter_ptr) != flag); - } else { - st_flag_volatile(peer_counter_ptr, flag); - while (ld_flag_volatile(self_counter_ptr) != flag); - } - } - if constexpr (!final_sync) __syncthreads(); - - // use one thread to update flag - if (threadIdx.x == 0) self_sg->_flag[blockIdx.x] = flag; -} - -#else - -template -DINLINE void barrier_at_start(const RankSignals& sg, Signal* self_sg, - int rank) { - uint32_t flag = self_sg->_flag[blockIdx.x] + 1; - if (threadIdx.x < ngpus) { - // simultaneously write to the corresponding flag of all ranks. - // Latency = 1 p2p write - __scoped_atomic_store_n(&sg.signals[threadIdx.x]->start[blockIdx.x][rank], - flag, __ATOMIC_RELAXED, __MEMORY_SCOPE_SYSTEM); - // wait until we got true from all ranks - while (__scoped_atomic_load_n(&self_sg->start[blockIdx.x][threadIdx.x], - __ATOMIC_RELAXED, - __MEMORY_SCOPE_DEVICE) < flag); - } - __syncthreads(); - // use one thread to update flag - if (threadIdx.x == 0) self_sg->_flag[blockIdx.x] = flag; -} - -template -DINLINE void barrier_at_end(const RankSignals& sg, Signal* self_sg, int rank) { - __syncthreads(); - uint32_t flag = self_sg->_flag[blockIdx.x] + 1; - if (threadIdx.x < ngpus) { - // simultaneously write to the corresponding flag of all ranks. - // Latency = 1 p2p write - __scoped_atomic_store_n(&sg.signals[threadIdx.x]->end[blockIdx.x][rank], - flag, - final_sync ? __ATOMIC_RELAXED : __ATOMIC_RELEASE, - __MEMORY_SCOPE_SYSTEM); - // wait until we got true from all ranks - while ( - __scoped_atomic_load_n(&self_sg->end[blockIdx.x][threadIdx.x], - final_sync ? __ATOMIC_RELAXED : __ATOMIC_ACQUIRE, - __MEMORY_SCOPE_DEVICE) < flag); - } - if constexpr (!final_sync) __syncthreads(); - // use one thread to update flag - if (threadIdx.x == 0) self_sg->_flag[blockIdx.x] = flag; -} - -#endif - -template -DINLINE P packed_reduce(const P* ptrs[], int idx) { - A tmp = upcast(ptrs[0][idx]); -#pragma unroll - for (int i = 1; i < ngpus; i++) { - packed_assign_add(tmp, upcast(ptrs[i][idx])); - } - return downcast

(tmp); -} - -template -__global__ void __launch_bounds__(512, 1) - cross_device_reduce_1stage(RankData* _dp, RankSignals sg, Signal* self_sg, - T* __restrict__ result, int rank, int size) { +#include "custom_collective_common.cuh" +cross_device_reduce_1stage(RankData* _dp, RankSignals sg, Signal* self_sg, + T* __restrict__ result, int rank, int size) { using P = typename packed_t::P; using A = typename packed_t::A; // note: we don't reorder the address so the accumulation order is the same @@ -617,6 +321,21 @@ class CustomAllreduce { #undef KL } + void allgather(cudaStream_t stream, void* input, void* output, int size_bytes, + int threads = 512, int block_limit = defaultBlockLimit); + template + void mnnvl_lamport_allgather(cudaStream_t stream, T* input, T* output, + void* local_buffer, void* multicast_buffer, + uint32_t* epochs, int size_bytes, + int stage_size_bytes); + template + void reduce_scatter(cudaStream_t stream, T* input, T* output, int size, + int threads = 512, int block_limit = defaultBlockLimit); + template + void mnnvl_lamport_reduce_scatter(cudaStream_t stream, T* input, T* output, + void* local_buffer, uint32_t* epochs, + int size, int stage_size_bytes); + ~CustomAllreduce() { for (auto [_, ptr] : ipc_handles_) { CUDACHECK(cudaIpcCloseMemHandle(ptr)); diff --git a/csrc/custom_collective_common.cuh b/csrc/custom_collective_common.cuh new file mode 100644 index 0000000000..d66afc0856 --- /dev/null +++ b/csrc/custom_collective_common.cuh @@ -0,0 +1,332 @@ +#pragma once + +#include +#include +#include +#include + +#if defined(USE_ROCM) +typedef __hip_bfloat16 nv_bfloat16; +#endif + +#include +#include +#include +#include +#include +#include +#include +#include + +namespace aphrodite { +constexpr int kMaxCustomCollectiveRanks = 16; + +#define CUDACHECK(cmd) \ + do { \ + cudaError_t e = cmd; \ + if (e != cudaSuccess) { \ + printf("Failed: Cuda error %s:%d '%s'\n", __FILE__, __LINE__, \ + cudaGetErrorString(e)); \ + exit(EXIT_FAILURE); \ + } \ + } while (0) + +// Maximal number of blocks in allreduce kernel. +constexpr int kMaxBlocks = 36; + +// Default number of blocks in allreduce kernel. +#ifndef USE_ROCM +inline constexpr int defaultBlockLimit = 36; +inline CUpointer_attribute rangeStartAddrAttr = + CU_POINTER_ATTRIBUTE_RANGE_START_ADDR; +#else +inline constexpr int defaultBlockLimit = 16; +inline hipPointer_attribute rangeStartAddrAttr = + HIP_POINTER_ATTRIBUTE_RANGE_START_ADDR; +#endif + +// Counter may overflow, but it's fine since unsigned int overflow is +// well-defined behavior. +using FlagType = uint32_t; + +// Two sets of peer counters are needed for two syncs: starting and ending an +// operation. The reason is that it's possible for peer GPU block to arrive at +// the second sync point while the current GPU block haven't passed the first +// sync point. Thus, peer GPU may write counter+1 while current GPU is busy +// waiting for counter. We use alternating counter array to avoid this +// possibility. +struct Signal { + alignas(128) FlagType start[kMaxBlocks][kMaxCustomCollectiveRanks]; + alignas(128) FlagType end[kMaxBlocks][kMaxCustomCollectiveRanks]; + alignas(128) FlagType _flag[kMaxBlocks]; // incremental flags for each rank +}; + +struct __align__(16) RankData { + const void* ptrs[kMaxCustomCollectiveRanks]; +}; + +struct __align__(16) RankSignals { + Signal* signals[kMaxCustomCollectiveRanks]; +}; + +// like std::array, but aligned +template +struct __align__(alignof(T) * sz) array_t { + T data[sz]; + using type = T; + static constexpr int size = sz; +}; + +// use packed type to maximize memory efficiency +// goal: generate ld.128 and st.128 instructions +template +struct packed_t { + // the (P)acked type for load/store + using P = array_t; + // the (A)ccumulator type for reduction + using A = array_t; +}; + +#define DINLINE __device__ __forceinline__ + +// scalar cast functions +DINLINE float upcast_s(half val) { return __half2float(val); } + +template +DINLINE T downcast_s(float val); +template <> +DINLINE half downcast_s(float val) { + return __float2half(val); +} + +// scalar add functions +// for some reason when compiling with Pytorch, the + operator for half and +// bfloat is disabled so we call the intrinsics directly +DINLINE half& assign_add(half& a, half b) { + a = __hadd(a, b); + return a; +} +DINLINE float& assign_add(float& a, float b) { return a += b; } + +#if (__CUDA_ARCH__ >= 800 || !defined(__CUDA_ARCH__)) +DINLINE float upcast_s(nv_bfloat16 val) { return __bfloat162float(val); } +template <> +DINLINE nv_bfloat16 downcast_s(float val) { + return __float2bfloat16(val); +} +DINLINE nv_bfloat16& assign_add(nv_bfloat16& a, nv_bfloat16 b) { + a = __hadd(a, b); + return a; +} +#endif + +template +DINLINE array_t& packed_assign_add(array_t& a, array_t b) { +#pragma unroll + for (int i = 0; i < N; i++) { + assign_add(a.data[i], b.data[i]); + } + return a; +} + +template +DINLINE array_t upcast(array_t val) { + if constexpr (std::is_same::value) { + return val; + } else { + array_t out; +#pragma unroll + for (int i = 0; i < N; i++) { + out.data[i] = upcast_s(val.data[i]); + } + return out; + } +} + +template +DINLINE O downcast(array_t val) { + if constexpr (std::is_same::value) { + return val; + } else { + O out; +#pragma unroll + for (int i = 0; i < O::size; i++) { + out.data[i] = downcast_s(val.data[i]); + } + return out; + } +} + +#if !defined(USE_ROCM) + +static DINLINE void st_flag_release(FlagType* flag_addr, FlagType flag) { + #if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 700 + asm volatile("st.release.sys.global.u32 [%1], %0;" ::"r"(flag), + "l"(flag_addr)); + #else + asm volatile("membar.sys; st.volatile.global.u32 [%1], %0;" ::"r"(flag), + "l"(flag_addr)); + #endif +} + +static DINLINE FlagType ld_flag_acquire(FlagType* flag_addr) { + FlagType flag; + #if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 700 + asm volatile("ld.acquire.sys.global.u32 %0, [%1];" + : "=r"(flag) + : "l"(flag_addr)); + #else + asm volatile("ld.volatile.global.u32 %0, [%1]; membar.gl;" + : "=r"(flag) + : "l"(flag_addr)); + #endif + return flag; +} + +static DINLINE void st_flag_volatile(FlagType* flag_addr, FlagType flag) { + asm volatile("st.volatile.global.u32 [%1], %0;" ::"r"(flag), "l"(flag_addr)); +} + +static DINLINE FlagType ld_flag_volatile(FlagType* flag_addr) { + FlagType flag; + asm volatile("ld.volatile.global.u32 %0, [%1];" + : "=r"(flag) + : "l"(flag_addr)); + return flag; +} + +// This function is meant to be used as the first synchronization in the all +// reduce kernel. Thus, it doesn't need to make any visibility guarantees for +// prior memory accesses. Note: volatile writes will not be reordered against +// other volatile writes. +template +DINLINE void barrier_at_start(const RankSignals& sg, Signal* self_sg, + int rank) { + uint32_t flag = self_sg->_flag[blockIdx.x] + 1; + if (threadIdx.x < ngpus) { + auto peer_counter_ptr = &sg.signals[threadIdx.x]->start[blockIdx.x][rank]; + auto self_counter_ptr = &self_sg->start[blockIdx.x][threadIdx.x]; + // Write the expected counter value to peer and wait for correct value + // from peer. + st_flag_volatile(peer_counter_ptr, flag); + while (ld_flag_volatile(self_counter_ptr) != flag); + } + __syncthreads(); + // use one thread to update flag + if (threadIdx.x == 0) self_sg->_flag[blockIdx.x] = flag; +} + +template +DINLINE void barrier_at_start_release(const RankSignals& sg, Signal* self_sg, + int rank) { + __syncthreads(); + uint32_t flag = self_sg->_flag[blockIdx.x] + 1; + if (threadIdx.x < ngpus) { + auto peer_counter_ptr = &sg.signals[threadIdx.x]->start[blockIdx.x][rank]; + auto self_counter_ptr = &self_sg->start[blockIdx.x][threadIdx.x]; + st_flag_release(peer_counter_ptr, flag); + while (ld_flag_acquire(self_counter_ptr) != flag); + } + __syncthreads(); + if (threadIdx.x == 0) self_sg->_flag[blockIdx.x] = flag; +} + +// This function is meant to be used as the second or the final +// synchronization barrier in the all reduce kernel. If it's the final +// synchronization barrier, we don't need to make any visibility guarantees +// for prior memory accesses. +template +DINLINE void barrier_at_end(const RankSignals& sg, Signal* self_sg, int rank) { + __syncthreads(); + uint32_t flag = self_sg->_flag[blockIdx.x] + 1; + if (threadIdx.x < ngpus) { + auto peer_counter_ptr = &sg.signals[threadIdx.x]->end[blockIdx.x][rank]; + auto self_counter_ptr = &self_sg->end[blockIdx.x][threadIdx.x]; + // Write the expected counter value to peer and wait for correct value from + // peer. + if constexpr (!final_sync) { + st_flag_release(peer_counter_ptr, flag); + while (ld_flag_acquire(self_counter_ptr) != flag); + } else { + st_flag_volatile(peer_counter_ptr, flag); + while (ld_flag_volatile(self_counter_ptr) != flag); + } + } + if constexpr (!final_sync) __syncthreads(); + + // use one thread to update flag + if (threadIdx.x == 0) self_sg->_flag[blockIdx.x] = flag; +} + +#else + +template +DINLINE void barrier_at_start(const RankSignals& sg, Signal* self_sg, + int rank) { + uint32_t flag = self_sg->_flag[blockIdx.x] + 1; + if (threadIdx.x < ngpus) { + // simultaneously write to the corresponding flag of all ranks. + // Latency = 1 p2p write + __scoped_atomic_store_n(&sg.signals[threadIdx.x]->start[blockIdx.x][rank], + flag, __ATOMIC_RELAXED, __MEMORY_SCOPE_SYSTEM); + // wait until we got true from all ranks + while (__scoped_atomic_load_n(&self_sg->start[blockIdx.x][threadIdx.x], + __ATOMIC_RELAXED, + __MEMORY_SCOPE_DEVICE) < flag); + } + __syncthreads(); + // use one thread to update flag + if (threadIdx.x == 0) self_sg->_flag[blockIdx.x] = flag; +} + +template +DINLINE void barrier_at_start_release(const RankSignals& sg, Signal* self_sg, + int rank) { + __syncthreads(); + uint32_t flag = self_sg->_flag[blockIdx.x] + 1; + if (threadIdx.x < ngpus) { + __scoped_atomic_store_n(&sg.signals[threadIdx.x]->start[blockIdx.x][rank], + flag, __ATOMIC_RELEASE, __MEMORY_SCOPE_SYSTEM); + while (__scoped_atomic_load_n(&self_sg->start[blockIdx.x][threadIdx.x], + __ATOMIC_ACQUIRE, + __MEMORY_SCOPE_DEVICE) < flag); + } + __syncthreads(); + if (threadIdx.x == 0) self_sg->_flag[blockIdx.x] = flag; +} + +template +DINLINE void barrier_at_end(const RankSignals& sg, Signal* self_sg, int rank) { + __syncthreads(); + uint32_t flag = self_sg->_flag[blockIdx.x] + 1; + if (threadIdx.x < ngpus) { + // simultaneously write to the corresponding flag of all ranks. + // Latency = 1 p2p write + __scoped_atomic_store_n(&sg.signals[threadIdx.x]->end[blockIdx.x][rank], + flag, + final_sync ? __ATOMIC_RELAXED : __ATOMIC_RELEASE, + __MEMORY_SCOPE_SYSTEM); + // wait until we got true from all ranks + while ( + __scoped_atomic_load_n(&self_sg->end[blockIdx.x][threadIdx.x], + final_sync ? __ATOMIC_RELAXED : __ATOMIC_ACQUIRE, + __MEMORY_SCOPE_DEVICE) < flag); + } + if constexpr (!final_sync) __syncthreads(); + // use one thread to update flag + if (threadIdx.x == 0) self_sg->_flag[blockIdx.x] = flag; +} + +#endif + +template +DINLINE P packed_reduce(const P* ptrs[], int idx) { + A tmp = upcast(ptrs[0][idx]); +#pragma unroll + for (int i = 1; i < ngpus; i++) { + packed_assign_add(tmp, upcast(ptrs[i][idx])); + } + return downcast

(tmp); +} + +} // namespace aphrodite diff --git a/csrc/flashkda_registration.cpp b/csrc/flashkda_registration.cpp new file mode 100644 index 0000000000..9058119e11 --- /dev/null +++ b/csrc/flashkda_registration.cpp @@ -0,0 +1,17 @@ +#include "core/registration.h" +#include "flash_kda.h" + +TORCH_LIBRARY(_flashkda_C, m) { + m.def("get_workspace_size(int T_total, int H, int N=1) -> int", + &get_workspace_size); + m.def( + "fwd(Tensor q, Tensor k, Tensor v, Tensor g, Tensor beta, float scale, " + "Tensor(a!) out, Tensor workspace, Tensor A_log, Tensor dt_bias, " + "float lower_bound, " + "Tensor? initial_state=None, Tensor(b!)? final_state=None, " + "Tensor? cu_seqlens=None) -> ()"); +} + +TORCH_LIBRARY_IMPL(_flashkda_C, CUDA, m) { m.impl("fwd", &fwd); } + +REGISTER_EXTENSION(_flashkda_C) diff --git a/csrc/libtorch_stable/activation_kernels.cu b/csrc/libtorch_stable/activation_kernels.cu index 6e3e6860e6..4f4cde78df 100644 --- a/csrc/libtorch_stable/activation_kernels.cu +++ b/csrc/libtorch_stable/activation_kernels.cu @@ -477,6 +477,66 @@ __global__ void swigluoai_and_mul_kernel( } } +// SITU (Kimi SituGLU) gated activation. Non-interleaved layout: +// input = [gate(d), up(d)] per token. +// gate_out = beta * tanh(gate / beta) * sigmoid(gate) +// up_out = (linear_beta > 0) ? linear_beta * tanh(up / linear_beta) : up +// out = gate_out * up_out +// Compute is done in fp32 and written straight to `out` -- no intermediate +// tensors and no full-tensor fp32 upcast (the pure-torch forward_native +// allocated ~8 fp32 temporaries per call, which blows up MoE profiling). +template +__global__ void situ_and_mul_kernel( + scalar_t* __restrict__ out, // [..., d] + const scalar_t* __restrict__ input, // [..., 2, d] + const int d, const float beta, const float linear_beta) { + const int64_t row = blockIdx.x; + const scalar_t* gate_ptr = input + row * 2 * d; + const scalar_t* up_ptr = gate_ptr + d; + scalar_t* out_ptr = out + row * d; + const bool clamp_up = linear_beta > 0.0f; + const float inv_beta = 1.0f / beta; + const float inv_linear_beta = clamp_up ? 1.0f / linear_beta : 0.0f; + for (int64_t idx = threadIdx.x; idx < d; idx += blockDim.x) { + const float g = (float)APHRODITE_LDG(&gate_ptr[idx]); + const float u = (float)APHRODITE_LDG(&up_ptr[idx]); + const float gate_out = beta * tanhf(g * inv_beta) / (1.0f + expf(-g)); + const float up_out = + clamp_up ? linear_beta * tanhf(u * inv_linear_beta) : u; + out_ptr[idx] = (scalar_t)(gate_out * up_out); + } +} + +template +__global__ void masked_situ_and_mul_kernel( + scalar_t* __restrict__ out, const scalar_t* __restrict__ input, + const int* __restrict__ expert_num_tokens, const int max_num_tokens, + const int d, const float beta, const float linear_beta) { + const int expert = blockIdx.y; + const int num_tokens = expert_num_tokens[expert]; + const int idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx >= d || num_tokens == 0) { + return; + } + + const bool clamp_up = linear_beta > 0.0f; + const float inv_beta = 1.0f / beta; + const float inv_linear_beta = clamp_up ? 1.0f / linear_beta : 0.0f; + const int64_t expert_row = static_cast(expert) * max_num_tokens; + for (int token = 0; token < num_tokens; ++token) { + const int64_t row = expert_row + token; + const scalar_t* gate_ptr = input + row * 2 * d; + const scalar_t* up_ptr = gate_ptr + d; + scalar_t* out_ptr = out + row * d; + const float g = (float)APHRODITE_LDG(&gate_ptr[idx]); + const float u = (float)APHRODITE_LDG(&up_ptr[idx]); + const float gate_out = beta * tanhf(g * inv_beta) / (1.0f + expf(-g)); + const float up_out = + clamp_up ? linear_beta * tanhf(u * inv_linear_beta) : u; + out_ptr[idx] = (scalar_t)(gate_out * up_out); + } +} + } // namespace aphrodite #define LAUNCH_ACTIVATION_GATE_KERNEL_WITH_PARAM(KERNEL, PACKED_KERNEL, PARAM) \ @@ -569,6 +629,56 @@ void swigluoai_and_mul(torch::stable::Tensor& out, // [..., d] double alpha, double limit) { LAUNCH_SIGLUOAI_AND_MUL(aphrodite::swigluoai_and_mul, alpha, limit); } + +// Kimi SITU gated activation. `linear_beta <= 0` means "unset" (up passed +// through), matching SituAndMul(linear_beta=None) on the Python side. +void situ_and_mul(torch::stable::Tensor& out, // [..., d] + torch::stable::Tensor& input, // [..., 2 * d] + double beta, double linear_beta) { + int d = input.size(-1) / 2; + int64_t num_tokens = input.numel() / input.size(-1); + if (num_tokens == 0) { + return; + } + dim3 grid(num_tokens); + dim3 block(std::min(d, 1024)); + const torch::stable::accelerator::DeviceGuard device_guard( + input.get_device_index()); + const cudaStream_t stream = get_current_cuda_stream(); + APHRODITE_STABLE_DISPATCH_FLOATING_TYPES( + input.scalar_type(), "situ_and_mul_kernel", [&] { + aphrodite::situ_and_mul_kernel<<>>( + out.mutable_data_ptr(), input.const_data_ptr(), + d, (float)beta, (float)linear_beta); + }); +} + +void masked_situ_and_mul(torch::stable::Tensor& out, // [E, T, d] + torch::stable::Tensor& input, // [E, T, 2 * d] + const torch::stable::Tensor& expert_num_tokens, + double beta, double linear_beta) { + int num_experts = input.size(0); + int max_num_tokens = input.size(1); + int d = input.size(2) / 2; + if (num_experts == 0 || max_num_tokens == 0) { + return; + } + constexpr int block_size = 256; + dim3 grid((d + block_size - 1) / block_size, num_experts); + dim3 block(block_size); + const torch::stable::accelerator::DeviceGuard device_guard( + input.get_device_index()); + const cudaStream_t stream = get_current_cuda_stream(); + APHRODITE_STABLE_DISPATCH_FLOATING_TYPES( + input.scalar_type(), "masked_situ_and_mul_kernel", [&] { + aphrodite::masked_situ_and_mul_kernel + <<>>( + out.mutable_data_ptr(), + input.const_data_ptr(), + expert_num_tokens.const_data_ptr(), max_num_tokens, d, + (float)beta, (float)linear_beta); + }); +} namespace aphrodite { // Element-wise activation kernel template. diff --git a/csrc/libtorch_stable/attention/merge_attn_states.cu b/csrc/libtorch_stable/attention/merge_attn_states.cu index 8eb6801bdd..d648cc3cfe 100644 --- a/csrc/libtorch_stable/attention/merge_attn_states.cu +++ b/csrc/libtorch_stable/attention/merge_attn_states.cu @@ -21,7 +21,10 @@ __global__ void merge_attn_states_kernel( const float* prefix_lse, const scalar_t* suffix_output, const float* suffix_lse, const uint num_tokens, const uint num_heads, const uint head_size, const uint prefix_head_stride, - const uint output_head_stride, const uint prefix_num_tokens, + const uint output_head_stride, const uint prefix_lse_head_stride, + const uint prefix_lse_token_stride, const uint suffix_lse_head_stride, + const uint suffix_lse_token_stride, const uint output_lse_head_stride, + const uint output_lse_token_stride, const uint prefix_num_tokens, const float* output_scale) { // Inputs always load 128-bit packs (pack_size elements of scalar_t). // Outputs store pack_size elements of output_t, which is smaller for FP8. @@ -84,15 +87,19 @@ __global__ void merge_attn_states_kernel( } } if (output_lse != nullptr && pack_idx == 0) { - float s_lse = suffix_lse[head_idx * num_tokens + token_idx]; - output_lse[head_idx * num_tokens + token_idx] = s_lse; + float s_lse = suffix_lse[head_idx * suffix_lse_head_stride + + token_idx * suffix_lse_token_stride]; + output_lse[head_idx * output_lse_head_stride + + token_idx * output_lse_token_stride] = s_lse; } return; } // For tokens within prefix range, merge prefix and suffix - float p_lse = prefix_lse[head_idx * num_tokens + token_idx]; - float s_lse = suffix_lse[head_idx * num_tokens + token_idx]; + float p_lse = prefix_lse[head_idx * prefix_lse_head_stride + + token_idx * prefix_lse_token_stride]; + float s_lse = suffix_lse[head_idx * suffix_lse_head_stride + + token_idx * suffix_lse_token_stride]; p_lse = std::isinf(p_lse) ? -std::numeric_limits::infinity() : p_lse; s_lse = std::isinf(s_lse) ? -std::numeric_limits::infinity() : s_lse; @@ -132,7 +139,8 @@ __global__ void merge_attn_states_kernel( } // We only need to write to output_lse once per head. if (output_lse != nullptr && pack_idx == 0) { - output_lse[head_idx * num_tokens + token_idx] = max_lse; + output_lse[head_idx * output_lse_head_stride + + token_idx * output_lse_token_stride] = max_lse; } return; } @@ -187,7 +195,8 @@ __global__ void merge_attn_states_kernel( // We only need to write to output_lse once per head. if (output_lse != nullptr && pack_idx == 0) { float out_lse = logf(out_se) + max_lse; - output_lse[head_idx * num_tokens + token_idx] = out_lse; + output_lse[head_idx * output_lse_head_stride + + token_idx * output_lse_token_stride] = out_lse; } } @@ -221,6 +230,9 @@ __global__ void merge_attn_states_kernel( reinterpret_cast(suffix_output.data_ptr()), \ reinterpret_cast(suffix_lse.data_ptr()), num_tokens, \ num_heads, head_size, prefix_head_stride, output_head_stride, \ + prefix_lse_head_stride, prefix_lse_token_stride, \ + suffix_lse_head_stride, suffix_lse_token_stride, \ + output_lse_head_stride, output_lse_token_stride, \ prefix_num_tokens, output_scale_ptr); \ } @@ -259,6 +271,19 @@ void merge_attn_states_launcher( const uint head_size = output.size(2); const uint prefix_head_stride = prefix_output.stride(1); const uint output_head_stride = output.stride(1); + // lse tensors are [NUM_HEADS, NUM_TOKENS] but may be non-contiguous views + // (e.g. a transpose of a backend's [NUM_TOKENS, NUM_HEADS] output), so index + // them by their actual strides rather than assuming a contiguous layout. + const uint prefix_lse_head_stride = prefix_lse.stride(0); + const uint prefix_lse_token_stride = prefix_lse.stride(1); + const uint suffix_lse_head_stride = suffix_lse.stride(0); + const uint suffix_lse_token_stride = suffix_lse.stride(1); + uint output_lse_head_stride = 0; + uint output_lse_token_stride = 0; + if (output_lse.has_value()) { + output_lse_head_stride = output_lse.value().stride(0); + output_lse_token_stride = output_lse.value().stride(1); + } // Thread mapping is based on input BF16 pack_size const uint pack_size = 16 / sizeof(scalar_t); STD_TORCH_CHECK(head_size % pack_size == 0, diff --git a/csrc/libtorch_stable/cache_kernels.cu b/csrc/libtorch_stable/cache_kernels.cu index 7ade78f56f..694b6c4780 100644 --- a/csrc/libtorch_stable/cache_kernels.cu +++ b/csrc/libtorch_stable/cache_kernels.cu @@ -443,6 +443,55 @@ __global__ void concat_and_cache_mla_kernel( copy(k_pe, kv_cache, k_pe_stride, block_stride, pe_dim, kv_lora_rank); } +// Grouped variant of concat_and_cache_mla: inserts the context K/V for every +// draft layer in a single launch. Grid is (num_tokens, num_layers); each layer +// reads its own cache base pointer from kv_cache_ptrs (same pointer-array +// pattern as copy_blocks_kernel). bf16 only, so it is a raw 16-bit copy with no +// scaling or quantization; scalar_t is uint16_t for portability. +template +__global__ void concat_and_cache_mla_grouped_kernel( + const scalar_t* __restrict__ kv_c, // [num_layers, num_tokens, + // kv_lora_rank] + const scalar_t* __restrict__ k_pe, // [num_layers, num_tokens, pe_dim] + const int64_t* __restrict__ kv_cache_ptrs, // [num_layers] + const int64_t* __restrict__ slot_mapping, // [num_layers, num_tokens] + const int64_t kv_c_layer_stride, const int64_t kv_c_token_stride, + const int64_t k_pe_layer_stride, const int64_t k_pe_token_stride, + const int64_t slot_layer_stride, const int64_t block_stride, + const int64_t entry_stride, const int kv_lora_rank, const int pe_dim, + const int block_size) { + const int64_t token_idx = blockIdx.x; + const int64_t layer_idx = blockIdx.y; + const int64_t slot_idx = + slot_mapping[layer_idx * slot_layer_stride + token_idx]; + // NOTE: slot_idx can be -1 if the token is padded + if (slot_idx < 0) { + return; + } + const int64_t block_idx = slot_idx / block_size; + const int64_t block_offset = slot_idx % block_size; + + scalar_t* __restrict__ kv_cache = + reinterpret_cast(kv_cache_ptrs[layer_idx]); + const scalar_t* __restrict__ kv_c_layer = + kv_c + layer_idx * kv_c_layer_stride; + const scalar_t* __restrict__ k_pe_layer = + k_pe + layer_idx * k_pe_layer_stride; + + auto copy = [&](const scalar_t* __restrict__ src, int64_t src_token_stride, + int size, int offset) { + for (int i = threadIdx.x; i < size; i += blockDim.x) { + const int64_t src_idx = token_idx * src_token_stride + i; + const int64_t dst_idx = + block_idx * block_stride + block_offset * entry_stride + i + offset; + kv_cache[dst_idx] = src[src_idx]; + } + }; + + copy(kv_c_layer, kv_c_token_stride, kv_lora_rank, 0); + copy(k_pe_layer, k_pe_token_stride, pe_dim, kv_lora_rank); +} + template __global__ void concat_and_cache_ds_mla_kernel( const scalar_t* __restrict__ kv_c, // [num_tokens, kv_lora_rank] @@ -903,6 +952,53 @@ void concat_and_cache_mla( } } +void concat_and_cache_mla_grouped( + torch::stable::Tensor& kv_c, // [num_layers, num_tokens, kv_lora_rank] + torch::stable::Tensor& k_pe, // [num_layers, num_tokens, pe_dim] + torch::stable::Tensor& kv_cache_ptrs, // [num_layers] int64, on device + torch::stable::Tensor& slot_mapping, // [num_layers, num_tokens] int64 + int64_t block_size, int64_t block_stride, int64_t entry_stride) { + int num_layers = kv_c.size(0); + int num_tokens = kv_c.size(1); + int kv_lora_rank = kv_c.size(2); + int pe_dim = k_pe.size(2); + + STD_TORCH_CHECK( + kv_c.scalar_type() == torch::headeronly::ScalarType::BFloat16 && + k_pe.scalar_type() == torch::headeronly::ScalarType::BFloat16, + "concat_and_cache_mla_grouped only supports a bf16 KV cache; got kv_c=", + kv_c.scalar_type(), ", k_pe=", k_pe.scalar_type()); + STD_TORCH_CHECK( + kv_cache_ptrs.scalar_type() == torch::headeronly::ScalarType::Long, + "kv_cache_ptrs must be int64"); + + if (num_tokens == 0 || num_layers == 0) { + return; + } + + const int64_t kv_c_layer_stride = kv_c.stride(0); + const int64_t kv_c_token_stride = kv_c.stride(1); + const int64_t k_pe_layer_stride = k_pe.stride(0); + const int64_t k_pe_token_stride = k_pe.stride(1); + const int64_t slot_layer_stride = slot_mapping.stride(0); + + const torch::stable::accelerator::DeviceGuard device_guard( + kv_c.get_device_index()); + const cudaStream_t stream = get_current_cuda_stream(); + + dim3 grid(num_tokens, num_layers); + dim3 block(std::min(kv_lora_rank, 512)); + aphrodite::concat_and_cache_mla_grouped_kernel + <<>>( + reinterpret_cast(kv_c.data_ptr()), + reinterpret_cast(k_pe.data_ptr()), + kv_cache_ptrs.const_data_ptr(), + slot_mapping.const_data_ptr(), kv_c_layer_stride, + kv_c_token_stride, k_pe_layer_stride, k_pe_token_stride, + slot_layer_stride, block_stride, entry_stride, kv_lora_rank, pe_dim, + block_size); +} + namespace aphrodite { template diff --git a/csrc/libtorch_stable/custom_all_gather_reduce_scatter.cu b/csrc/libtorch_stable/custom_all_gather_reduce_scatter.cu new file mode 100644 index 0000000000..b906549a0b --- /dev/null +++ b/csrc/libtorch_stable/custom_all_gather_reduce_scatter.cu @@ -0,0 +1,362 @@ +#include "torch_utils.h" + +#include +#include +#include +#include + +#include "custom_all_reduce.cuh" +#include "custom_all_gather_reduce_scatter.cuh" + +namespace aphrodite { + +void CustomAllreduce::allgather(cudaStream_t stream, void* input, void* output, + int size_bytes, int threads, int block_limit) { + if (size_bytes % sizeof(CopyPack) != 0) + throw std::runtime_error( + "custom allgather requires input byte size to be a multiple of " + + std::to_string(sizeof(CopyPack))); + + auto ptrs = buffers_.at(input); + int size_per_rank = size_bytes / sizeof(CopyPack); + int total_size = size_per_rank * world_size_; + int blocks = std::min(block_limit, (total_size + threads - 1) / threads); + +#define AG_CASE(ngpus) \ + case ngpus: \ + cross_device_all_gather<<>>( \ + ptrs, sg_, self_sg_, reinterpret_cast(output), rank_, \ + size_per_rank); \ + break; + + switch (world_size_) { + AG_CASE(2) + AG_CASE(4) + AG_CASE(6) + AG_CASE(8) + default: + throw std::runtime_error( + "custom allgather only supports num gpus in (2,4,6,8)"); + } +#undef AG_CASE +} + +template +void CustomAllreduce::mnnvl_lamport_allgather(cudaStream_t stream, T* input, + T* output, void* local_buffer, + void* multicast_buffer, + uint32_t* epochs, int size_bytes, + int stage_size_bytes) { + if (size_bytes % sizeof(typename packed_t::P) != 0 || + stage_size_bytes % sizeof(typename packed_t::P) != 0) + throw std::runtime_error( + "MNNVL Lamport allgather requires 16-byte aligned sizes"); + + auto ptrs = buffers_.at(local_buffer); + int size_per_rank = size_bytes / sizeof(typename packed_t::P); + int stage_size = stage_size_bytes / sizeof(typename packed_t::P); + int blocks = + (size_per_rank + kMnnvlLamportAgThreads - 1) / kMnnvlLamportAgThreads; + +#if !defined(USE_ROCM) && CUDA_VERSION >= 12000 + cudaLaunchAttribute attributes[1]{}; + attributes[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attributes[0].val.programmaticStreamSerializationAllowed = 1; + cudaLaunchConfig_t config{.gridDim = dim3(blocks), + .blockDim = dim3(kMnnvlLamportAgThreads), + .dynamicSmemBytes = 0, + .stream = stream, + .attrs = attributes, + .numAttrs = 1}; + #define MNNVL_LAMPORT_AG_LAUNCH(ngpus) \ + CUDACHECK(cudaLaunchKernelEx(&config, &mnnvl_lamport_all_gather, \ + ptrs, input, output, \ + reinterpret_cast(multicast_buffer), \ + epochs, rank_, size_per_rank, stage_size)) +#else + #define MNNVL_LAMPORT_AG_LAUNCH(ngpus) \ + mnnvl_lamport_all_gather \ + <<>>( \ + ptrs, input, output, reinterpret_cast(multicast_buffer), \ + epochs, rank_, size_per_rank, stage_size) +#endif + +#define MNNVL_LAMPORT_AG_CASE(ngpus) \ + case ngpus: \ + MNNVL_LAMPORT_AG_LAUNCH(ngpus); \ + break; + + switch (world_size_) { + MNNVL_LAMPORT_AG_CASE(2) + MNNVL_LAMPORT_AG_CASE(4) + MNNVL_LAMPORT_AG_CASE(6) + MNNVL_LAMPORT_AG_CASE(8) + MNNVL_LAMPORT_AG_CASE(16) + default: + throw std::runtime_error( + "MNNVL Lamport allgather only supports num gpus in (2,4,6,8,16)"); + } +#undef MNNVL_LAMPORT_AG_CASE +#undef MNNVL_LAMPORT_AG_LAUNCH +} + +template +void CustomAllreduce::reduce_scatter(cudaStream_t stream, T* input, T* output, + int size, int threads, int block_limit) { + auto packed_size = packed_t::P::size; + if (size % (packed_size * world_size_) != 0) + throw std::runtime_error( + "custom reduce-scatter requires each output shard byte size to be " + "a multiple of 16"); + + auto ptrs = buffers_.at(input); + int size_per_rank = size / packed_size / world_size_; + int blocks = std::min(block_limit, (size_per_rank + threads - 1) / threads); + +#define RS_CASE(ngpus) \ + case ngpus: \ + cross_device_reduce_scatter<<>>( \ + ptrs, sg_, self_sg_, output, rank_, size_per_rank); \ + break; + + switch (world_size_) { + RS_CASE(2) + RS_CASE(4) + RS_CASE(6) + RS_CASE(8) + default: + throw std::runtime_error( + "custom reduce-scatter only supports num gpus in (2,4,6,8)"); + } +#undef RS_CASE +} + +template +void CustomAllreduce::mnnvl_lamport_reduce_scatter(cudaStream_t stream, + T* input, T* output, + void* local_buffer, + uint32_t* epochs, int size, + int stage_size_bytes) { + auto packed_size = packed_t::P::size; + if (size % (packed_size * world_size_) != 0 || + stage_size_bytes % sizeof(typename packed_t::P) != 0) + throw std::runtime_error( + "MNNVL Lamport reduce-scatter requires 16-byte aligned sizes"); + + auto ptrs = buffers_.at(local_buffer); + int size_per_rank = size / packed_size / world_size_; + int stage_size = stage_size_bytes / sizeof(typename packed_t::P); + int blocks_per_rank = + (size_per_rank + kMnnvlLamportRsThreads - 1) / kMnnvlLamportRsThreads; + int blocks = blocks_per_rank * world_size_; + +#if !defined(USE_ROCM) && CUDA_VERSION >= 12000 + cudaLaunchAttribute attributes[1]{}; + attributes[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attributes[0].val.programmaticStreamSerializationAllowed = 1; + cudaLaunchConfig_t config{.gridDim = dim3(blocks), + .blockDim = dim3(kMnnvlLamportRsThreads), + .dynamicSmemBytes = 0, + .stream = stream, + .attrs = attributes, + .numAttrs = 1}; + #define MNNVL_LAMPORT_RS_LAUNCH(ngpus) \ + CUDACHECK(cudaLaunchKernelEx( \ + &config, &mnnvl_lamport_reduce_scatter_kernel, ptrs, input, \ + output, epochs, rank_, size_per_rank, stage_size)) +#else + #define MNNVL_LAMPORT_RS_LAUNCH(ngpus) \ + mnnvl_lamport_reduce_scatter_kernel \ + <<>>( \ + ptrs, input, output, epochs, rank_, size_per_rank, stage_size) +#endif + +#define MNNVL_LAMPORT_RS_CASE(ngpus) \ + case ngpus: \ + MNNVL_LAMPORT_RS_LAUNCH(ngpus); \ + break; + + switch (world_size_) { + MNNVL_LAMPORT_RS_CASE(2) + MNNVL_LAMPORT_RS_CASE(4) + MNNVL_LAMPORT_RS_CASE(6) + MNNVL_LAMPORT_RS_CASE(8) + MNNVL_LAMPORT_RS_CASE(16) + default: + throw std::runtime_error( + "MNNVL Lamport reduce-scatter only supports num gpus in " + "(2,4,6,8,16)"); + } +#undef MNNVL_LAMPORT_RS_CASE +#undef MNNVL_LAMPORT_RS_LAUNCH +} + +} // namespace aphrodite + +using fptr_t = int64_t; +static_assert(sizeof(void*) == sizeof(fptr_t)); + +bool _is_weak_contiguous(torch::stable::Tensor& t); + +void custom_all_gather(fptr_t _fa, torch::stable::Tensor& inp, + torch::stable::Tensor& out, fptr_t _reg_buffer, + int64_t reg_buffer_sz_bytes) { + auto fa = reinterpret_cast(_fa); + const torch::stable::accelerator::DeviceGuard device_guard( + inp.get_device_index()); + const cudaStream_t stream = get_current_cuda_stream(inp.get_device_index()); + + STD_TORCH_CHECK((inp.scalar_type()) == (out.scalar_type())); + STD_TORCH_CHECK((inp.numel() * fa->world_size_) == (out.numel())); + STD_TORCH_CHECK(_is_weak_contiguous(out)); + STD_TORCH_CHECK(_is_weak_contiguous(inp)); + auto input_size = inp.numel() * inp.element_size(); + auto reg_buffer = reinterpret_cast(_reg_buffer); + STD_TORCH_CHECK(reg_buffer != nullptr); + STD_TORCH_CHECK((input_size) <= (reg_buffer_sz_bytes)); + STD_CUDA_CHECK(cudaMemcpyAsync(reg_buffer, inp.const_data_ptr(), input_size, + cudaMemcpyDeviceToDevice, stream)); + fa->allgather(stream, reg_buffer, out.mutable_data_ptr(), input_size); +} + +void mnnvl_lamport_all_gather(fptr_t _fa, torch::stable::Tensor& inp, + torch::stable::Tensor& out, fptr_t _local_buffer, + fptr_t _multicast_buffer, fptr_t _epoch_buffer, + int64_t stage_sz_bytes) { + auto fa = reinterpret_cast(_fa); + const torch::stable::accelerator::DeviceGuard device_guard( + inp.get_device_index()); + const cudaStream_t stream = get_current_cuda_stream(inp.get_device_index()); + + STD_TORCH_CHECK((inp.scalar_type()) == (out.scalar_type())); + STD_TORCH_CHECK((inp.numel() * fa->world_size_) == (out.numel())); + STD_TORCH_CHECK(_is_weak_contiguous(out)); + STD_TORCH_CHECK(_is_weak_contiguous(inp)); + auto input_size = inp.numel() * inp.element_size(); + STD_TORCH_CHECK((input_size * fa->world_size_) <= stage_sz_bytes); + auto local_buffer = reinterpret_cast(_local_buffer); + auto multicast_buffer = reinterpret_cast(_multicast_buffer); + auto epochs = reinterpret_cast(_epoch_buffer); + switch (out.scalar_type()) { + case torch::headeronly::ScalarType::Float: { + fa->mnnvl_lamport_allgather( + stream, reinterpret_cast(inp.mutable_data_ptr()), + reinterpret_cast(out.mutable_data_ptr()), local_buffer, + multicast_buffer, epochs, input_size, stage_sz_bytes); + break; + } + case torch::headeronly::ScalarType::Half: { + fa->mnnvl_lamport_allgather( + stream, reinterpret_cast(inp.mutable_data_ptr()), + reinterpret_cast(out.mutable_data_ptr()), local_buffer, + multicast_buffer, epochs, input_size, stage_sz_bytes); + break; + } +#if (__CUDA_ARCH__ >= 800 || !defined(__CUDA_ARCH__)) + case torch::headeronly::ScalarType::BFloat16: { + fa->mnnvl_lamport_allgather( + stream, reinterpret_cast(inp.mutable_data_ptr()), + reinterpret_cast(out.mutable_data_ptr()), local_buffer, + multicast_buffer, epochs, input_size, stage_sz_bytes); + break; + } +#endif + default: + throw std::runtime_error( + "MNNVL Lamport allgather only supports float32, float16 and " + "bfloat16"); + } +} + +void custom_reduce_scatter(fptr_t _fa, torch::stable::Tensor& inp, + torch::stable::Tensor& out, fptr_t _reg_buffer, + int64_t reg_buffer_sz_bytes) { + auto fa = reinterpret_cast(_fa); + const torch::stable::accelerator::DeviceGuard device_guard( + inp.get_device_index()); + const cudaStream_t stream = get_current_cuda_stream(inp.get_device_index()); + + STD_TORCH_CHECK((inp.scalar_type()) == (out.scalar_type())); + STD_TORCH_CHECK((out.numel() * fa->world_size_) == (inp.numel())); + STD_TORCH_CHECK(_is_weak_contiguous(out)); + STD_TORCH_CHECK(_is_weak_contiguous(inp)); + auto input_size = inp.numel() * inp.element_size(); + auto reg_buffer = reinterpret_cast(_reg_buffer); + STD_TORCH_CHECK(reg_buffer != nullptr); + STD_TORCH_CHECK((input_size) <= (reg_buffer_sz_bytes)); + STD_CUDA_CHECK(cudaMemcpyAsync(reg_buffer, inp.const_data_ptr(), input_size, + cudaMemcpyDeviceToDevice, stream)); + switch (out.scalar_type()) { + case torch::headeronly::ScalarType::Float: { + fa->reduce_scatter( + stream, reinterpret_cast(reg_buffer), + reinterpret_cast(out.mutable_data_ptr()), inp.numel()); + break; + } + case torch::headeronly::ScalarType::Half: { + fa->reduce_scatter(stream, reinterpret_cast(reg_buffer), + reinterpret_cast(out.mutable_data_ptr()), + inp.numel()); + break; + } +#if (__CUDA_ARCH__ >= 800 || !defined(__CUDA_ARCH__)) + case torch::headeronly::ScalarType::BFloat16: { + fa->reduce_scatter( + stream, reinterpret_cast(reg_buffer), + reinterpret_cast(out.mutable_data_ptr()), inp.numel()); + break; + } +#endif + default: + throw std::runtime_error( + "custom reduce-scatter only supports float32, float16 and bfloat16"); + } +} + +void mnnvl_lamport_reduce_scatter(fptr_t _fa, torch::stable::Tensor& inp, + torch::stable::Tensor& out, + fptr_t _local_buffer, fptr_t _epoch_buffer, + int64_t stage_sz_bytes) { + auto fa = reinterpret_cast(_fa); + const torch::stable::accelerator::DeviceGuard device_guard( + inp.get_device_index()); + const cudaStream_t stream = get_current_cuda_stream(inp.get_device_index()); + + STD_TORCH_CHECK((inp.scalar_type()) == (out.scalar_type())); + STD_TORCH_CHECK((out.numel() * fa->world_size_) == (inp.numel())); + STD_TORCH_CHECK(_is_weak_contiguous(out)); + STD_TORCH_CHECK(_is_weak_contiguous(inp)); + auto input_size = inp.numel() * inp.element_size(); + STD_TORCH_CHECK(input_size <= stage_sz_bytes); + auto local_buffer = reinterpret_cast(_local_buffer); + auto epochs = reinterpret_cast(_epoch_buffer); + switch (out.scalar_type()) { + case torch::headeronly::ScalarType::Float: { + fa->mnnvl_lamport_reduce_scatter( + stream, reinterpret_cast(inp.mutable_data_ptr()), + reinterpret_cast(out.mutable_data_ptr()), local_buffer, + epochs, inp.numel(), stage_sz_bytes); + break; + } + case torch::headeronly::ScalarType::Half: { + fa->mnnvl_lamport_reduce_scatter( + stream, reinterpret_cast(inp.mutable_data_ptr()), + reinterpret_cast(out.mutable_data_ptr()), local_buffer, epochs, + inp.numel(), stage_sz_bytes); + break; + } +#if (__CUDA_ARCH__ >= 800 || !defined(__CUDA_ARCH__)) + case torch::headeronly::ScalarType::BFloat16: { + fa->mnnvl_lamport_reduce_scatter( + stream, reinterpret_cast(inp.mutable_data_ptr()), + reinterpret_cast(out.mutable_data_ptr()), local_buffer, + epochs, inp.numel(), stage_sz_bytes); + break; + } +#endif + default: + throw std::runtime_error( + "MNNVL Lamport reduce-scatter only supports float32, float16 and " + "bfloat16"); + } +} diff --git a/csrc/libtorch_stable/custom_all_gather_reduce_scatter_ops.cpp b/csrc/libtorch_stable/custom_all_gather_reduce_scatter_ops.cpp new file mode 100644 index 0000000000..198b78d56d --- /dev/null +++ b/csrc/libtorch_stable/custom_all_gather_reduce_scatter_ops.cpp @@ -0,0 +1,29 @@ +#include "ops.h" +#include "core/registration.h" + +#include + +STABLE_TORCH_LIBRARY_FRAGMENT(_C_custom_ar, custom_ag_rs) { + custom_ag_rs.def( + "custom_all_gather(int fa, Tensor inp, Tensor! out, int reg_buffer, " + "int reg_buffer_sz_bytes) -> ()"); + custom_ag_rs.def( + "mnnvl_lamport_all_gather(int fa, Tensor inp, Tensor! out, int " + "local_buffer, int multicast_buffer, int epoch_buffer, int " + "stage_sz_bytes) -> ()"); + custom_ag_rs.def( + "custom_reduce_scatter(int fa, Tensor inp, Tensor! out, int reg_buffer, " + "int reg_buffer_sz_bytes) -> ()"); + custom_ag_rs.def( + "mnnvl_lamport_reduce_scatter(int fa, Tensor inp, Tensor! out, int " + "local_buffer, int epoch_buffer, int stage_sz_bytes) -> ()"); +} + +STABLE_TORCH_LIBRARY_IMPL(_C_custom_ar, CUDA, custom_ag_rs) { + custom_ag_rs.impl("custom_all_gather", TORCH_BOX(&custom_all_gather)); + custom_ag_rs.impl("mnnvl_lamport_all_gather", + TORCH_BOX(&mnnvl_lamport_all_gather)); + custom_ag_rs.impl("custom_reduce_scatter", TORCH_BOX(&custom_reduce_scatter)); + custom_ag_rs.impl("mnnvl_lamport_reduce_scatter", + TORCH_BOX(&mnnvl_lamport_reduce_scatter)); +} diff --git a/csrc/libtorch_stable/custom_all_reduce.cu b/csrc/libtorch_stable/custom_all_reduce.cu index 6d43964edd..047841d9c4 100644 --- a/csrc/libtorch_stable/custom_all_reduce.cu +++ b/csrc/libtorch_stable/custom_all_reduce.cu @@ -18,14 +18,14 @@ fptr_t init_custom_ar(const std::vector& fake_ipc_ptrs, torch::stable::Tensor& rank_data, int64_t rank, bool fully_connected) { int world_size = fake_ipc_ptrs.size(); - if (world_size > 8) - throw std::invalid_argument("world size > 8 is not supported"); + if (world_size > aphrodite::kMaxCustomCollectiveRanks) + throw std::invalid_argument("world size > 16 is not supported"); if (world_size % 2 != 0) throw std::invalid_argument("Odd num gpus is not supported for now"); if (rank < 0 || rank >= world_size) throw std::invalid_argument("invalid rank passed in"); - aphrodite::Signal* ipc_ptrs[8]; + aphrodite::Signal* ipc_ptrs[aphrodite::kMaxCustomCollectiveRanks]; for (int i = 0; i < world_size; i++) { ipc_ptrs[i] = reinterpret_cast(fake_ipc_ptrs[i]); } @@ -124,7 +124,7 @@ int64_t meta_size() { return sizeof(aphrodite::Signal); } void register_buffer(fptr_t _fa, const std::vector& fake_ipc_ptrs) { auto fa = reinterpret_cast(_fa); STD_TORCH_CHECK(fake_ipc_ptrs.size() == fa->world_size_); - void* ipc_ptrs[8]; + void* ipc_ptrs[aphrodite::kMaxCustomCollectiveRanks]; for (int i = 0; i < fake_ipc_ptrs.size(); i++) { ipc_ptrs[i] = reinterpret_cast(fake_ipc_ptrs[i]); } diff --git a/csrc/libtorch_stable/dsv3_fused_a_gemm.cu b/csrc/libtorch_stable/dsv3_fused_a_gemm.cu index 585004c047..bac9869a02 100644 --- a/csrc/libtorch_stable/dsv3_fused_a_gemm.cu +++ b/csrc/libtorch_stable/dsv3_fused_a_gemm.cu @@ -647,17 +647,17 @@ __global__ __launch_bounds__(256, 1) void fused_a_gemm_kernel( #endif } -template +template void invokeFusedAGemm(T* output, T const* mat_a, T const* mat_b, int num_tokens, - cudaStream_t const stream) { - constexpr int gemm_m = kHdOut; // 2112 - int const gemm_n = num_tokens; // 1-16 - constexpr int gemm_k = kHdIn; // 7168 + cudaStream_t const stream, bool enable_pdl) { + constexpr int gemm_m = kHdOut; + int const gemm_n = num_tokens; + constexpr int gemm_k = kHdIn; constexpr int batch_size = 1; std::swap(mat_a, mat_b); constexpr int tile_m = 16; - constexpr int tile_n = kTileN; // 8 or 16 - constexpr int tile_k = std::max(256, 1024 / tile_n); // 256 + constexpr int tile_n = kTileN; + constexpr int tile_k = kTileK; constexpr int max_stage_cnt = 1024 * 192 / ((tile_m + tile_n) * tile_k * sizeof(bf16_t)); constexpr int k_iter_cnt = gemm_k / tile_k; @@ -679,7 +679,8 @@ void invokeFusedAGemm(T* output, T const* mat_a, T const* mat_b, int num_tokens, config.stream = stream; cudaLaunchAttribute attrs[1]; attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; - attrs[0].val.programmaticStreamSerializationAllowed = getEnvEnablePDL(); + attrs[0].val.programmaticStreamSerializationAllowed = + enable_pdl || getEnvEnablePDL(); config.numAttrs = 1; config.attrs = attrs; if (smem_bytes >= (48 * 1024)) { @@ -694,36 +695,50 @@ void invokeFusedAGemm(T* output, T const* mat_a, T const* mat_b, int num_tokens, output, mat_a, mat_b, gemm_n); } -template void invokeFusedAGemm<__nv_bfloat16, 7168, 2112, 8>( - __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, int num_tokens, - cudaStream_t); - -template void invokeFusedAGemm<__nv_bfloat16, 7168, 2112, 16>( - __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, int num_tokens, - cudaStream_t); +template +void invokeFusedAGemmForTokens(T* output, T const* mat_a, T const* mat_b, + int num_tokens, cudaStream_t const stream, + bool enable_pdl) { + if (num_tokens <= 8) { + invokeFusedAGemm( + output, mat_a, mat_b, num_tokens, stream, enable_pdl); + } else { + invokeFusedAGemm( + output, mat_a, mat_b, num_tokens, stream, enable_pdl); + } +} void dsv3_fused_a_gemm(torch::stable::Tensor& output, torch::stable::Tensor const& mat_a, - torch::stable::Tensor const& mat_b) { + torch::stable::Tensor const& mat_b, bool enable_pdl) { STD_TORCH_CHECK(mat_a.dim() == 2 && mat_b.dim() == 2 && output.dim() == 2); int const num_tokens = mat_a.size(0); int const hd_in = mat_a.size(1); int const hd_out = mat_b.size(1); - constexpr int kHdIn = 7168; - constexpr int kHdOut = 2112; STD_TORCH_CHECK(num_tokens >= 1 && num_tokens <= 16, "required 1 <= mat_a.shape[0] <= 16"); - STD_TORCH_CHECK(hd_in == kHdIn, "required mat_a.shape[1] == 7168"); - STD_TORCH_CHECK(hd_out == kHdOut, "required mat_b.shape[1] == 2112"); STD_TORCH_CHECK(output.size(0) == num_tokens, "required output.shape[0] == mat_a.shape[0]"); STD_TORCH_CHECK(output.size(1) == hd_out, "required output.shape[1] == mat_b.shape[1]"); + STD_TORCH_CHECK(mat_b.size(0) == hd_in, + "required mat_b.shape[0] == mat_a.shape[1]"); - STD_TORCH_CHECK(mat_a.stride(1) == 1, "mat_a must be a row major tensor"); - STD_TORCH_CHECK(output.stride(1) == 1, "output must be a row major tensor"); - STD_TORCH_CHECK(mat_b.stride(0) == 1, "mat_b must be a column major tensor"); + STD_TORCH_CHECK(mat_a.get_device_index() == mat_b.get_device_index() && + mat_a.get_device_index() == output.get_device_index(), + "mat_a, mat_b, and output must be on the same device"); + + // The kernels index global memory with raw pointers and packed strides, so + // reject any padded or transposed view rather than reading out of bounds. + STD_TORCH_CHECK( + mat_a.stride(0) == hd_in && mat_a.stride(1) == 1, + "mat_a must be a packed row-major [num_tokens, hd_in] tensor"); + STD_TORCH_CHECK( + output.stride(0) == hd_out && output.stride(1) == 1, + "output must be a packed row-major [num_tokens, hd_out] tensor"); + STD_TORCH_CHECK(mat_b.stride(0) == 1 && mat_b.stride(1) == hd_in, + "mat_b must be a packed column-major [hd_in, hd_out] tensor"); STD_TORCH_CHECK( mat_a.scalar_type() == torch::headeronly::ScalarType::BFloat16 && @@ -738,19 +753,85 @@ void dsv3_fused_a_gemm(torch::stable::Tensor& output, STD_TORCH_CHECK(getSMVersion() >= 90, "required CUDA ARCH >= SM_90"); auto stream = get_current_cuda_stream(mat_a.get_device_index()); - if (num_tokens <= 8) { - invokeFusedAGemm<__nv_bfloat16, kHdIn, kHdOut, 8>( - reinterpret_cast<__nv_bfloat16*>(output.mutable_data_ptr()), - reinterpret_cast<__nv_bfloat16 const*>(mat_a.data_ptr()), - reinterpret_cast<__nv_bfloat16 const*>(mat_b.data_ptr()), num_tokens, - stream); - } else { - invokeFusedAGemm<__nv_bfloat16, kHdIn, kHdOut, 16>( - reinterpret_cast<__nv_bfloat16*>(output.mutable_data_ptr()), - reinterpret_cast<__nv_bfloat16 const*>(mat_a.data_ptr()), - reinterpret_cast<__nv_bfloat16 const*>(mat_b.data_ptr()), num_tokens, - stream); + auto* output_ptr = + reinterpret_cast<__nv_bfloat16*>(output.mutable_data_ptr()); + auto const* mat_a_ptr = + reinterpret_cast<__nv_bfloat16 const*>(mat_a.data_ptr()); + auto const* mat_b_ptr = + reinterpret_cast<__nv_bfloat16 const*>(mat_b.data_ptr()); + +#define DISPATCH_DSV3_SHAPE(HD_IN, HD_OUT) \ + if (hd_in == HD_IN && hd_out == HD_OUT) { \ + invokeFusedAGemmForTokens<__nv_bfloat16, HD_IN, HD_OUT>( \ + output_ptr, mat_a_ptr, mat_b_ptr, num_tokens, stream, enable_pdl); \ + return; \ } + + // Shapes the Kimi-K3 selector routes to dsv3_fused_a (see the dsv3 winners + // in KIMI_K3_PROJECTIONS) plus the DeepSeek V2/V3 QKV A-projection. + DISPATCH_DSV3_SHAPE(7168, 1536) + DISPATCH_DSV3_SHAPE(7168, 2112) + DISPATCH_DSV3_SHAPE(1536, 2304) + DISPATCH_DSV3_SHAPE(1536, 4608) + DISPATCH_DSV3_SHAPE(7168, 3584) + DISPATCH_DSV3_SHAPE(768, 7168) + // TP16 dsv3 winners, as (hd_in=K, hd_out=N). TP16 dense down_proj is absent + // because hd_in=2112 is not a multiple of any supported tile_k. + DISPATCH_DSV3_SHAPE(1536, 1152) + DISPATCH_DSV3_SHAPE(7168, 768) + DISPATCH_DSV3_SHAPE(7168, 3216) + DISPATCH_DSV3_SHAPE(7168, 4224) + +#ifdef APHRODITE_K3_BENCH_SHAPES + // The selector routes these shapes to CuTe or the default GEMM, so they are + // never reached in production. They are compiled only for offline + // DSV3-vs-CuTe benchmarking. + DISPATCH_DSV3_SHAPE(7168, 6288) + DISPATCH_DSV3_SHAPE(1536, 7168) + DISPATCH_DSV3_SHAPE(3584, 7168) + DISPATCH_DSV3_SHAPE(7168, 8448) + DISPATCH_DSV3_SHAPE(7168, 20480) + DISPATCH_DSV3_SHAPE(7168, 3072) + DISPATCH_DSV3_SHAPE(7168, 12448) + DISPATCH_DSV3_SHAPE(3072, 7168) + DISPATCH_DSV3_SHAPE(8448, 7168) + DISPATCH_DSV3_SHAPE(7168, 16896) + DISPATCH_DSV3_SHAPE(7168, 40960) +#endif + +#undef DISPATCH_DSV3_SHAPE + + if (hd_in == 128 && hd_out == 1536) { + invokeFusedAGemmForTokens<__nv_bfloat16, 128, 1536, 128>( + output_ptr, mat_a_ptr, mat_b_ptr, num_tokens, stream, enable_pdl); + return; + } + if (hd_in == 128 && hd_out == 3072) { + invokeFusedAGemmForTokens<__nv_bfloat16, 128, 3072, 128>( + output_ptr, mat_a_ptr, mat_b_ptr, num_tokens, stream, enable_pdl); + return; + } + // TP16 KDA f_b_proj and shared_expert down_proj. Neither hd_in is a multiple + // of 256, so both need the 128 tile_k. + if (hd_in == 128 && hd_out == 768) { + invokeFusedAGemmForTokens<__nv_bfloat16, 128, 768, 128>( + output_ptr, mat_a_ptr, mat_b_ptr, num_tokens, stream, enable_pdl); + return; + } + if (hd_in == 384 && hd_out == 7168) { + invokeFusedAGemmForTokens<__nv_bfloat16, 384, 7168, 128>( + output_ptr, mat_a_ptr, mat_b_ptr, num_tokens, stream, enable_pdl); + return; + } +#ifdef APHRODITE_K3_BENCH_SHAPES + if (hd_in == 4224 && hd_out == 7168) { + invokeFusedAGemmForTokens<__nv_bfloat16, 4224, 7168, 128>( + output_ptr, mat_a_ptr, mat_b_ptr, num_tokens, stream, enable_pdl); + return; + } +#endif + + STD_TORCH_CHECK(false, "unsupported DSV3 fused-A GEMM shape"); } STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, m) { diff --git a/csrc/libtorch_stable/fused_kimi_k3_mla_key_concat_kv_cache_kernel.cu b/csrc/libtorch_stable/fused_kimi_k3_mla_key_concat_kv_cache_kernel.cu new file mode 100644 index 0000000000..a88160fb81 --- /dev/null +++ b/csrc/libtorch_stable/fused_kimi_k3_mla_key_concat_kv_cache_kernel.cu @@ -0,0 +1,1237 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * SPDX-FileCopyrightText: Copyright contributors to the vLLM project + * + * Fused Kimi-K3 MLA prefill + decode epilogues with optional RoPE. + * + * Prefill: runs after the q_b_proj / kv_b_proj GEMMs, one launch per token + * slice. Decode: runs after BMM1 (q_nope x W_UK) right before forward_mqa, + * concatenating mqa_q = [ql_nope | q_pe] and inserting the latent cache + * (fused_kimi_k3_mla_decode_q_concat_kv_cache_{,fp8_,ds_mla_}insert). + * + * Prefill variants: + * + * bf16 (fused_kimi_k3_mla_key_concat_kv_cache_insert): + * - optional in-place q RoPE: rotate q[t, h, 128:192] + * - full key concat: k_out[t, h] = [k_nope[t, h] | k_pe[t]] (per head) + * - latent cache insert: cache[slot(t)] = [kv_c_normed[t] | k_pe[t]] + * (v is used as-is in bf16, so it is not touched here.) + * + * fp8 (fused_kimi_k3_mla_qkv_quant_kv_cache_fp8_insert): + * - q_fp8[t, h] = quant(q[t, h], q_scale) + * - k_fp8[t, h] = quant([k_nope[t, h] | k_pe[t]], k_scale) + * - v_fp8[t, h] = quant(v[t, h], v_scale) + * - cache[slot(t)] = quant([kv_c_normed[t] | k_pe[t]], k_scale) + * matching MLA's _q_scale / _k_scale / _v_scale (the cache latent uses a + * separate cache scale). Per-tensor E4M3. + * + * fp8_ds_mla (fused_kimi_k3_mla_key_concat_ds_mla_insert): + * - full key concat (bf16), and cache insert in DeepSeek's 656-byte + * block-scaled layout (NoPE fp8 in 4 tiles of 128 with per-tile dynamic + * scales, RoPE bf16), bit-compatible with concat_and_cache_ds_mla_kernel. + * + * Both use Programmatic Dependent Launch (PDL) to overlap the tail of the + * producing GEMMs on sm_90+, and are structured after + * `fusedDeepseekV4FullCacheKernel`: one grid, one warp per (token, slot) with + * `slotsPerToken = num_heads + 1`. Slots [0, H) do the per-head work; the extra + * slot H does the per-token cache insert. All dims are multiples of 8, so bf16 + * copies move one uint4 (8 elems) per step and fp8 stores pack 8 elems into a + * uint2. + */ + +#include "torch_utils.h" + +#include +#include +#include +#include +#include +#include + +#include "cuda_compat.h" +#include "dispatch_utils.h" +#include "type_convert.cuh" + +#ifndef USE_ROCM + #include + #include "../quantization/w8a8/fp8/nvidia/quant_utils.cuh" +#else + #include + #include "../quantization/w8a8/fp8/amd/quant_utils.cuh" +#endif +#include +#include +#include + +#ifdef USE_ROCM +__device__ __forceinline__ uint8_t rocm_cvt_float_to_fp8_e4m3(float val) { + #if defined(__gfx942__) + __hip_fp8_e4m3_fnuz fp8_val(val); + #else + __hip_fp8_e4m3 fp8_val(val); + #endif + return reinterpret_cast(fp8_val); +} +#endif + +namespace aphrodite { +namespace kimi_k3_fused_ops { + +namespace { +inline int getSMVersion() { + auto* props = get_device_prop(); + return props->major * 10 + props->minor; +} +} // namespace + +// ──────────────────────────────────────────────────────────────────────────── +// Constants (Kimi-K3 MLA) +// ──────────────────────────────────────────────────────────────────────────── +constexpr int kKvLoraRank = 512; // L +constexpr int kQkNopeHeadDim = 128; // P +constexpr int kQkRopeHeadDim = 64; // R +constexpr int kQkHeadDim = kQkNopeHeadDim + kQkRopeHeadDim; // 192 +constexpr int kVHeadDim = 128; // V +constexpr int kCacheEntry = kKvLoraRank + kQkRopeHeadDim; // 576 +constexpr int kVecElems = 8; // 8 bf16 == one uint4 load / one uint2 fp8 store + +#if defined(USE_ROCM) && defined(__gfx942__) +constexpr float kFp8Max = 224.0f; +#else +constexpr float kFp8Max = 448.0f; +#endif +// Divisor for fp8_ds_mla per-tile dynamic scales (matches cache_kernels.cu). +// fp8_ds_mla 656B entry: [0,512) NoPE fp8 (4 tiles of 128), [512,528) 4 fp32 +// tile scales, [528,656) RoPE 64 bf16. +constexpr float kFp8ScaleDivisor = kFp8Max; + +// Copy 8 source elements (one uint4 of bf16/fp16) to `dst`. FP8=false stores a +// uint4 (bf16); FP8=true decodes to fp32, scales by `scale_inv`, saturates to +// ±kFp8Max and packs into a uint2 of E4M3. +template +__device__ __forceinline__ void copyChunk8(void* dst, const scalar_t* src, + float scale_inv, + const float* cos_sin = nullptr, + int rope_elem_base = 0) { + uint4 const v = *reinterpret_cast(src); + if constexpr (FP8 || APPLY_ROPE) { + using Converter = aphrodite::_typeConvert; + auto const* p = + reinterpret_cast(&v); + float f[kVecElems]; +#pragma unroll + for (int i = 0; i < 4; i++) { + float2 x = Converter::convert(p[i]); + f[2 * i] = x.x; + f[2 * i + 1] = x.y; + } + if constexpr (APPLY_ROPE) { +#pragma unroll + for (int i = 0; i < kVecElems / 2; i++) { + int const pair_idx = rope_elem_base / 2 + i; + float const cos = static_cast(cos_sin[pair_idx]); + float const sin = + static_cast(cos_sin[pair_idx + kQkRopeHeadDim / 2]); + float const x = f[2 * i]; + float const y = f[2 * i + 1]; + f[2 * i] = x * cos - y * sin; + f[2 * i + 1] = x * sin + y * cos; + } + } + if constexpr (!FP8) { + uint4 out; + auto* o = reinterpret_cast(&out); +#pragma unroll + for (int i = 0; i < kVecElems / 2; i++) { + o[i] = Converter::convert(make_float2(f[2 * i], f[2 * i + 1])); + } + *reinterpret_cast(dst) = out; + return; + } +#ifndef USE_ROCM + uint2 out; + auto* o2 = reinterpret_cast<__nv_fp8x2_storage_t*>(&out); + #pragma unroll + for (int i = 0; i < 4; i++) { + float2 s = make_float2(f[2 * i] * scale_inv, f[2 * i + 1] * scale_inv); + s.x = fminf(fmaxf(s.x, -kFp8Max), kFp8Max); + s.y = fminf(fmaxf(s.y, -kFp8Max), kFp8Max); + o2[i] = __nv_cvt_float2_to_fp8x2(s, __NV_SATFINITE, __NV_E4M3); + } + *reinterpret_cast(dst) = out; +#else + uint8_t out[kVecElems]; + #pragma unroll + for (int i = 0; i < kVecElems; i++) { + float s = fminf(fmaxf(f[i] * scale_inv, -kFp8Max), kFp8Max); + out[i] = rocm_cvt_float_to_fp8_e4m3(s); + } + *reinterpret_cast(dst) = *reinterpret_cast(out); +#endif + } else { + *reinterpret_cast(dst) = v; + } +} + +// Concat + store one head's full key: dst[e] = [k_nope | k_pe], e in [0, 192). +// FP8 dst is byte-addressed; bf16 dst is scalar_t-addressed (dst_elem_size). +template +__device__ __forceinline__ void writeFullKey(void* dst, const scalar_t* k_nope, + const scalar_t* k_pe, int laneId, + int dst_elem_size, float scale_inv, + const float* cos_sin = nullptr) { + auto* d = reinterpret_cast(dst); + for (int e = laneId * kVecElems; e < kQkHeadDim; e += 32 * kVecElems) { + if (e < kQkNopeHeadDim) { + copyChunk8(d + e * dst_elem_size, k_nope + e, scale_inv); + } else { + int const rope_e = e - kQkNopeHeadDim; + copyChunk8( + d + e * dst_elem_size, k_pe + rope_e, scale_inv, cos_sin, rope_e); + } + } +} + +// Store a prefill query, rotating only q[..., 128:192]. For bf16 dst may alias +// q (in-place); fp8 writes the quantized query directly to its output. +template +__device__ __forceinline__ void writePrefillQuery( + void* dst, const scalar_t* q, int laneId, int dst_elem_size, + float scale_inv, const float* cos_sin = nullptr) { + auto* d = reinterpret_cast(dst); + if constexpr (FP8) { + for (int e = laneId * kVecElems; e < kQkHeadDim; e += 32 * kVecElems) { + if (e < kQkNopeHeadDim) { + copyChunk8(d + e * dst_elem_size, q + e, scale_inv); + } else { + int const rope_e = e - kQkNopeHeadDim; + copyChunk8(d + e * dst_elem_size, q + e, + scale_inv, cos_sin, rope_e); + } + } + } else if constexpr (APPLY_ROPE) { + for (int e = laneId * kVecElems; e < kQkRopeHeadDim; e += 32 * kVecElems) { + copyChunk8( + d + (kQkNopeHeadDim + e) * dst_elem_size, q + kQkNopeHeadDim + e, + scale_inv, cos_sin, e); + } + } +} + +// Concat + store a 576-wide latent: dst[e] = [a512 | b64], e in [0, 576). Used +// for the decode query mqa_q = [ql_nope | q_pe] and the plain latent cache +// entry [kv_c | k_pe]. FP8 packs to E4M3 (dst_elem_size 1); bf16 stores uint4. +template +__device__ __forceinline__ void writeLatent576(void* dst, const scalar_t* a512, + const scalar_t* b64, int laneId, + int dst_elem_size, + float scale_inv, + const float* cos_sin = nullptr) { + auto* d = reinterpret_cast(dst); + for (int e = laneId * kVecElems; e < kCacheEntry; e += 32 * kVecElems) { + if (e < kKvLoraRank) { + copyChunk8(d + e * dst_elem_size, a512 + e, scale_inv); + } else { + int const rope_e = e - kKvLoraRank; + copyChunk8(d + e * dst_elem_size, b64 + rope_e, + scale_inv, cos_sin, rope_e); + } + } +} + +// Write [kv_c | k_pe] into the fp8_ds_mla 656B entry using one warp: NoPE 512 +// as fp8 in 4 tiles of 128 (per-tile dynamic absmax scale, 4 fp32 scales at +// [512,528)), RoPE 64 as bf16 at [528,656). Bit-compatible with +// concat_and_cache_ds_mla_kernel. +template +__device__ __forceinline__ void writeDsMlaCache( + uint8_t* row, const scalar_t* kvc, const scalar_t* pe, int laneId, + const float* cos_sin = nullptr) { + constexpr int kElemsPerLane = kKvLoraRank / 32; // 16 + int const tile = laneId >> 3; // 8 lanes per tile + scalar_t vals[kElemsPerLane]; + *reinterpret_cast(vals) = + *reinterpret_cast(kvc + laneId * kElemsPerLane); + *reinterpret_cast(vals + 8) = + *reinterpret_cast(kvc + laneId * kElemsPerLane + 8); + + float max_abs = 0.0f; +#pragma unroll + for (int i = 0; i < kElemsPerLane; i++) { + max_abs = fmaxf(max_abs, fabsf(static_cast(vals[i]))); + } +#pragma unroll + for (int offset = 4; offset > 0; offset /= 2) { + max_abs = fmaxf(max_abs, APHRODITE_SHFL_XOR_SYNC_WIDTH(max_abs, offset, 8)); + } + float const tile_scale = fmaxf(max_abs / kFp8ScaleDivisor, FLT_MIN); + if ((laneId & 7) == 0) { + reinterpret_cast(row)[kKvLoraRank / 4 + tile] = tile_scale; + } + uint8_t res[kElemsPerLane]; +#pragma unroll + for (int i = 0; i < kElemsPerLane; i++) { + res[i] = + fp8::scaled_convert( + vals[i], tile_scale); + } + *reinterpret_cast(row + laneId * kElemsPerLane) = + *reinterpret_cast(res); + scalar_t* row16 = reinterpret_cast(row); + scalar_t* rope_dst = row16 + kKvLoraRank / 2 + 8 + laneId * 2; + if constexpr (APPLY_ROPE) { + using Converter = aphrodite::_typeConvert; + using packed_t = typename Converter::packed_hip_type; + packed_t const src = *reinterpret_cast(pe + laneId * 2); + float2 const xy = Converter::convert(src); + float const cos = static_cast(cos_sin[laneId]); + float const sin = static_cast(cos_sin[laneId + kQkRopeHeadDim / 2]); + *reinterpret_cast(rope_dst) = Converter::convert( + make_float2(xy.x * cos - xy.y * sin, xy.x * sin + xy.y * cos)); + } else { + *reinterpret_cast(rope_dst) = + *reinterpret_cast(pe + laneId * 2); + } +} + +// ──────────────────────────────────────────────────────────────────────────── +// bf16 variant: optional q RoPE + full-key concat + latent cache insert +// ──────────────────────────────────────────────────────────────────────────── +template +__global__ void fusedKimiK3MLAKeyConcatKVCacheInsertKernel( + scalar_t* __restrict__ q, int64_t const q_tok_stride, + int64_t const q_head_stride, const scalar_t* __restrict__ k_nope, + int64_t const kn_tok_stride, int64_t const kn_head_stride, + const scalar_t* __restrict__ k_pe, int64_t const k_pe_tok_stride, + const scalar_t* __restrict__ kv_c, int64_t const kv_c_tok_stride, + scalar_t* __restrict__ k_out, int64_t const ko_tok_stride, + int64_t const ko_head_stride, scalar_t* __restrict__ k_cache, + int64_t const cache_block_stride, int64_t const cache_token_stride, + const int64_t* __restrict__ slot_mapping, + const int64_t* __restrict__ position_ids, + const float* __restrict__ cos_sin_cache, int const num_tokens, + int const num_heads, int const cache_block_size) { + int const warpsPerBlock = blockDim.x / 32; + int const laneId = threadIdx.x % 32; + int const globalWarpIdx = blockIdx.x * warpsPerBlock + threadIdx.x / 32; + int const slotsPerToken = num_heads + 1; + int const tokenIdx = globalWarpIdx / slotsPerToken; + int const slotIdx = globalWarpIdx % slotsPerToken; + if (tokenIdx >= num_tokens) return; + +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + cudaGridDependencySynchronize(); +#endif + + const float* rope_cache = nullptr; + if constexpr (APPLY_ROPE) { + rope_cache = cos_sin_cache + position_ids[tokenIdx] * kQkRopeHeadDim; + } + + if (slotIdx < num_heads) { + scalar_t* qh = q + tokenIdx * q_tok_stride + slotIdx * q_head_stride; + writePrefillQuery( + qh, qh, laneId, sizeof(scalar_t), 1.0f, rope_cache); + writeFullKey( + k_out + tokenIdx * ko_tok_stride + slotIdx * ko_head_stride, + k_nope + tokenIdx * kn_tok_stride + slotIdx * kn_head_stride, + k_pe + tokenIdx * k_pe_tok_stride, laneId, sizeof(scalar_t), 1.0f, + rope_cache); + } else { + int64_t const slot_id = slot_mapping[tokenIdx]; + if (slot_id >= 0) { + scalar_t* row = k_cache + + (slot_id / cache_block_size) * cache_block_stride + + (slot_id % cache_block_size) * cache_token_stride; + writeLatent576( + row, kv_c + tokenIdx * kv_c_tok_stride, + k_pe + tokenIdx * k_pe_tok_stride, laneId, sizeof(scalar_t), 1.0f, + rope_cache); + } + } + +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + cudaTriggerProgrammaticLaunchCompletion(); +#endif +} + +// ──────────────────────────────────────────────────────────────────────────── +// fp8 variant: quant q / k / v + latent cache insert +// ──────────────────────────────────────────────────────────────────────────── +template +__global__ void fusedKimiK3MLAQKVQuantKVCacheFp8Kernel( + const scalar_t* __restrict__ q, int64_t const q_tok_stride, + int64_t const q_head_stride, const scalar_t* __restrict__ k_nope, + int64_t const kn_tok_stride, int64_t const kn_head_stride, + const scalar_t* __restrict__ k_pe, int64_t const k_pe_tok_stride, + const scalar_t* __restrict__ kv_c, int64_t const kv_c_tok_stride, + const scalar_t* __restrict__ v, int64_t const v_tok_stride, + int64_t const v_head_stride, uint8_t* __restrict__ q_fp8, + int64_t const qo_tok_stride, int64_t const qo_head_stride, + uint8_t* __restrict__ k_fp8, int64_t const ko_tok_stride, + int64_t const ko_head_stride, uint8_t* __restrict__ v_fp8, + int64_t const vo_tok_stride, int64_t const vo_head_stride, + uint8_t* __restrict__ k_cache, int64_t const cache_block_stride, + int64_t const cache_token_stride, const int64_t* __restrict__ slot_mapping, + const float* __restrict__ q_scale_inv, + const float* __restrict__ k_scale_inv, + const float* __restrict__ v_scale_inv, + const float* __restrict__ cache_scale_inv, int const num_tokens, + int const num_heads, int const cache_block_size, + const int64_t* __restrict__ position_ids, + const float* __restrict__ cos_sin_cache) { + int const warpsPerBlock = blockDim.x / 32; + int const laneId = threadIdx.x % 32; + int const globalWarpIdx = blockIdx.x * warpsPerBlock + threadIdx.x / 32; + int const slotsPerToken = num_heads + 1; + int const tokenIdx = globalWarpIdx / slotsPerToken; + int const slotIdx = globalWarpIdx % slotsPerToken; + if (tokenIdx >= num_tokens) return; + +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + cudaGridDependencySynchronize(); +#endif + + const float* rope_cache = nullptr; + if constexpr (APPLY_ROPE) { + rope_cache = cos_sin_cache + position_ids[tokenIdx] * kQkRopeHeadDim; + } + + if (slotIdx < num_heads) { + int const h = slotIdx; + // q_fp8[t, h] = quant(q[t, h], q_scale) + float const qsi = __ldg(q_scale_inv); + const scalar_t* qh = q + tokenIdx * q_tok_stride + h * q_head_stride; + uint8_t* qo = q_fp8 + tokenIdx * qo_tok_stride + h * qo_head_stride; + writePrefillQuery(qo, qh, laneId, 1, qsi, + rope_cache); + // k_fp8[t, h] = quant([k_nope | k_pe], k_scale) + writeFullKey( + k_fp8 + tokenIdx * ko_tok_stride + h * ko_head_stride, + k_nope + tokenIdx * kn_tok_stride + h * kn_head_stride, + k_pe + tokenIdx * k_pe_tok_stride, laneId, 1, __ldg(k_scale_inv), + rope_cache); + // v_fp8[t, h] = quant(v[t, h], v_scale) + float const vsi = __ldg(v_scale_inv); + const scalar_t* vh = v + tokenIdx * v_tok_stride + h * v_head_stride; + uint8_t* vo = v_fp8 + tokenIdx * vo_tok_stride + h * vo_head_stride; + for (int e = laneId * kVecElems; e < kVHeadDim; e += 32 * kVecElems) { + copyChunk8(vo + e, vh + e, vsi); + } + } else { + int64_t const slot_id = slot_mapping[tokenIdx]; + if (slot_id >= 0) { + // The cache latent uses _k_scale (read back by decode / context); the + // attention key (k_fp8 above) uses its own k_scale. + float const ksi = __ldg(cache_scale_inv); + uint8_t* row = k_cache + + (slot_id / cache_block_size) * cache_block_stride + + (slot_id % cache_block_size) * cache_token_stride; + writeLatent576( + row, kv_c + tokenIdx * kv_c_tok_stride, + k_pe + tokenIdx * k_pe_tok_stride, laneId, 1, ksi, rope_cache); + } + } + +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + cudaTriggerProgrammaticLaunchCompletion(); +#endif +} + +// ──────────────────────────────────────────────────────────────────────────── +// ds_mla variant: concat full key (bf16) + fp8_ds_mla latent cache insert +// +// Cache entry (656 bytes), matching concat_and_cache_ds_mla_kernel: +// [0, 512) NoPE 512 vals as fp8, 4 tiles of 128, each dynamically scaled +// [512, 528) 4 fp32 per-tile scales +// [528, 656) RoPE 64 vals as bf16 (unquantized) +// The cache slot uses one warp: lane L quantizes NoPE elems [L*16, L*16+16) +// (tile = L>>3, absmax-reduced within its 8-lane tile group), then all 32 lanes +// write 2 RoPE bf16 each. +// ──────────────────────────────────────────────────────────────────────────── +template +__global__ void fusedKimiK3MLAKeyConcatDsMlaInsertKernel( + scalar_t* __restrict__ q, int64_t const q_tok_stride, + int64_t const q_head_stride, const scalar_t* __restrict__ k_nope, + int64_t const kn_tok_stride, int64_t const kn_head_stride, + const scalar_t* __restrict__ k_pe, int64_t const k_pe_tok_stride, + const scalar_t* __restrict__ kv_c, int64_t const kv_c_tok_stride, + scalar_t* __restrict__ k_out, int64_t const ko_tok_stride, + int64_t const ko_head_stride, uint8_t* __restrict__ k_cache, + int64_t const cache_block_stride, int64_t const cache_token_stride, + const int64_t* __restrict__ slot_mapping, int const num_tokens, + int const num_heads, int const cache_block_size, + const int64_t* __restrict__ position_ids, + const float* __restrict__ cos_sin_cache) { + int const warpsPerBlock = blockDim.x / 32; + int const laneId = threadIdx.x % 32; + int const globalWarpIdx = blockIdx.x * warpsPerBlock + threadIdx.x / 32; + int const slotsPerToken = num_heads + 1; + int const tokenIdx = globalWarpIdx / slotsPerToken; + int const slotIdx = globalWarpIdx % slotsPerToken; + if (tokenIdx >= num_tokens) return; + +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + cudaGridDependencySynchronize(); +#endif + + const float* rope_cache = nullptr; + if constexpr (APPLY_ROPE) { + rope_cache = cos_sin_cache + position_ids[tokenIdx] * kQkRopeHeadDim; + } + + if (slotIdx < num_heads) { + scalar_t* qh = q + tokenIdx * q_tok_stride + slotIdx * q_head_stride; + writePrefillQuery( + qh, qh, laneId, sizeof(scalar_t), 1.0f, rope_cache); + // Full key (bf16): k_out[t, h] = [k_nope[t, h] | k_pe[t]]. + writeFullKey( + k_out + tokenIdx * ko_tok_stride + slotIdx * ko_head_stride, + k_nope + tokenIdx * kn_tok_stride + slotIdx * kn_head_stride, + k_pe + tokenIdx * k_pe_tok_stride, laneId, sizeof(scalar_t), 1.0f, + rope_cache); + } else { + int64_t const slot_id = slot_mapping[tokenIdx]; + if (slot_id >= 0) { + uint8_t* row = k_cache + + (slot_id / cache_block_size) * cache_block_stride + + (slot_id % cache_block_size) * cache_token_stride; + writeDsMlaCache( + row, kv_c + tokenIdx * kv_c_tok_stride, + k_pe + tokenIdx * k_pe_tok_stride, laneId, rope_cache); + } + } + +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + cudaTriggerProgrammaticLaunchCompletion(); +#endif +} + +// ──────────────────────────────────────────────────────────────────────────── +// Decode epilogue: concat mqa_q = [ql_nope | q_pe] (576) + latent cache insert, +// run right before forward_mqa. Q_FP8 quantizes mqa_q; KV_FP8 quantizes the +// plain per-tensor cache. (ds_mla cache uses the separate kernel below.) +// ──────────────────────────────────────────────────────────────────────────── +template +__global__ void fusedKimiK3MLADecodeQConcatKVCacheKernel( + const scalar_t* __restrict__ ql_nope, int64_t const qn_tok_stride, + int64_t const qn_head_stride, const scalar_t* __restrict__ q_pe, + int64_t const qpe_tok_stride, int64_t const qpe_head_stride, + const scalar_t* __restrict__ kv_c, int64_t const kv_c_tok_stride, + const scalar_t* __restrict__ k_pe, int64_t const k_pe_tok_stride, + void* __restrict__ mqa_q, int64_t const mq_tok_stride, + int64_t const mq_head_stride, void* __restrict__ k_cache, + int64_t const cache_block_stride, int64_t const cache_token_stride, + const int64_t* __restrict__ slot_mapping, + const float* __restrict__ q_scale_inv, + const float* __restrict__ cache_scale_inv, int const num_tokens, + int const num_heads, int const cache_block_size, + const int64_t* __restrict__ position_ids, + const float* __restrict__ cos_sin_cache) { + constexpr int kMqElem = Q_FP8 ? 1 : sizeof(scalar_t); + constexpr int kCacheElem = KV_FP8 ? 1 : sizeof(scalar_t); + int const warpsPerBlock = blockDim.x / 32; + int const laneId = threadIdx.x % 32; + int const globalWarpIdx = blockIdx.x * warpsPerBlock + threadIdx.x / 32; + int const slotsPerToken = num_heads + 1; + int const tokenIdx = globalWarpIdx / slotsPerToken; + int const slotIdx = globalWarpIdx % slotsPerToken; + if (tokenIdx >= num_tokens) return; + +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + cudaGridDependencySynchronize(); +#endif + + const float* rope_cache = nullptr; + if constexpr (APPLY_ROPE) { + rope_cache = cos_sin_cache + position_ids[tokenIdx] * kQkRopeHeadDim; + } + + if (slotIdx < num_heads) { + float const qsi = Q_FP8 ? __ldg(q_scale_inv) : 1.0f; + writeLatent576( + reinterpret_cast(mqa_q) + + (tokenIdx * mq_tok_stride + slotIdx * mq_head_stride) * kMqElem, + ql_nope + tokenIdx * qn_tok_stride + slotIdx * qn_head_stride, + q_pe + tokenIdx * qpe_tok_stride + slotIdx * qpe_head_stride, laneId, + kMqElem, qsi, rope_cache); + } else { + int64_t const slot_id = slot_mapping[tokenIdx]; + if (slot_id >= 0) { + float const ksi = KV_FP8 ? __ldg(cache_scale_inv) : 1.0f; + writeLatent576( + reinterpret_cast(k_cache) + + (slot_id / cache_block_size * cache_block_stride + + slot_id % cache_block_size * cache_token_stride) * + kCacheElem, + kv_c + tokenIdx * kv_c_tok_stride, k_pe + tokenIdx * k_pe_tok_stride, + laneId, kCacheElem, ksi, rope_cache); + } + } + +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + cudaTriggerProgrammaticLaunchCompletion(); +#endif +} + +// Decode epilogue for fp8_ds_mla: concat mqa_q (bf16) + ds_mla cache insert. +template +__global__ void fusedKimiK3MLADecodeQConcatDsMlaKernel( + const scalar_t* __restrict__ ql_nope, int64_t const qn_tok_stride, + int64_t const qn_head_stride, const scalar_t* __restrict__ q_pe, + int64_t const qpe_tok_stride, int64_t const qpe_head_stride, + const scalar_t* __restrict__ kv_c, int64_t const kv_c_tok_stride, + const scalar_t* __restrict__ k_pe, int64_t const k_pe_tok_stride, + scalar_t* __restrict__ mqa_q, int64_t const mq_tok_stride, + int64_t const mq_head_stride, uint8_t* __restrict__ k_cache, + int64_t const cache_block_stride, int64_t const cache_token_stride, + const int64_t* __restrict__ slot_mapping, int const num_tokens, + int const num_heads, int const cache_block_size, + const int64_t* __restrict__ position_ids, + const float* __restrict__ cos_sin_cache) { + int const warpsPerBlock = blockDim.x / 32; + int const laneId = threadIdx.x % 32; + int const globalWarpIdx = blockIdx.x * warpsPerBlock + threadIdx.x / 32; + int const slotsPerToken = num_heads + 1; + int const tokenIdx = globalWarpIdx / slotsPerToken; + int const slotIdx = globalWarpIdx % slotsPerToken; + if (tokenIdx >= num_tokens) return; + +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + cudaGridDependencySynchronize(); +#endif + + const float* rope_cache = nullptr; + if constexpr (APPLY_ROPE) { + rope_cache = cos_sin_cache + position_ids[tokenIdx] * kQkRopeHeadDim; + } + + if (slotIdx < num_heads) { + writeLatent576( + mqa_q + tokenIdx * mq_tok_stride + slotIdx * mq_head_stride, + ql_nope + tokenIdx * qn_tok_stride + slotIdx * qn_head_stride, + q_pe + tokenIdx * qpe_tok_stride + slotIdx * qpe_head_stride, laneId, + sizeof(scalar_t), 1.0f, rope_cache); + } else { + int64_t const slot_id = slot_mapping[tokenIdx]; + if (slot_id >= 0) { + uint8_t* row = k_cache + + (slot_id / cache_block_size) * cache_block_stride + + (slot_id % cache_block_size) * cache_token_stride; + writeDsMlaCache( + row, kv_c + tokenIdx * kv_c_tok_stride, + k_pe + tokenIdx * k_pe_tok_stride, laneId, rope_cache); + } + } + +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + cudaTriggerProgrammaticLaunchCompletion(); +#endif +} + +// PDL-aware launch of a (token, num_heads + 1)-warp grid. +template +static void launchPdl(KernelT kernel, int num_tokens, int num_heads, + cudaStream_t stream, Args... args) { + constexpr int kBlockSize = 256; + constexpr int kWarpsPerBlock = kBlockSize / 32; + int64_t const total_warps = + static_cast(num_tokens) * (num_heads + 1); + int const grid = + static_cast((total_warps + kWarpsPerBlock - 1) / kWarpsPerBlock); +#ifndef USE_ROCM + static int const sm_version = getSMVersion(); + cudaLaunchConfig_t config; + config.gridDim = dim3(grid); + config.blockDim = dim3(kBlockSize); + config.dynamicSmemBytes = 0; + config.stream = stream; + cudaLaunchAttribute attrs[1]; + attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attrs[0].val.programmaticStreamSerializationAllowed = 1; + config.attrs = attrs; + config.numAttrs = (sm_version >= 90) ? 1 : 0; + cudaLaunchKernelEx(&config, kernel, args...); +#else + // clang-format off + // hipify's CUDA->HIP regex catastrophically backtracks on "> > >"; keep the + // launch closer as ">>>". clang-format would otherwise re-split it (it does + // not parse the CUDA launch syntax in this libtorch_stable file). + kernel<<>>(args...); + // clang-format on +#endif +} + +} // namespace kimi_k3_fused_ops +} // namespace aphrodite + +// ──────────────────────────────────────────────────────────────────────────── +// Torch op wrappers +// ──────────────────────────────────────────────────────────────────────────── +namespace { +bool check_rope_inputs( + std::optional const& position_ids, + std::optional const& cos_sin_cache, + torch::stable::Tensor const& /*input*/, int64_t num_tokens) { + using torch::headeronly::ScalarType; + STD_TORCH_CHECK(position_ids.has_value() == cos_sin_cache.has_value(), + "position_ids and cos_sin_cache must be provided together"); + if (!position_ids.has_value()) return false; + + auto const& positions = position_ids.value(); + auto const& rope_cache = cos_sin_cache.value(); + STD_TORCH_CHECK(positions.device().is_cuda() && positions.dim() == 1 && + positions.scalar_type() == ScalarType::Long && + positions.size(0) == num_tokens, + "position_ids must be int64 CUDA with shape [num_tokens]"); + STD_TORCH_CHECK(rope_cache.device().is_cuda() && rope_cache.dim() == 2 && + rope_cache.size(1) == 64 && rope_cache.stride(1) == 1 && + rope_cache.scalar_type() == ScalarType::Float, + "cos_sin_cache must have shape [max_position, 64], unit " + "last-dim stride, and be fp32 (RoPE math runs in fp32)"); + return true; +} +} // namespace + +void fused_kimi_k3_mla_key_concat_kv_cache_insert( + torch::stable::Tensor& q, // [Tp, H, 192] + torch::stable::Tensor const& k_nope, // [Tp, H, 128] + torch::stable::Tensor const& k_pe, // [Tp, 64] + torch::stable::Tensor const& kv_c_normed, // [Tp, 512] + torch::stable::Tensor& k_out, // [Tp, H, 192] bf16, written + torch::stable::Tensor& k_cache, // [nblk, bs, 576] bf16, written + torch::stable::Tensor const& slot_mapping, // [Tp] int64 + int64_t cache_block_size, std::optional position_ids, + std::optional cos_sin_cache) { + using torch::headeronly::ScalarType; + namespace kk3 = aphrodite::kimi_k3_fused_ops; + STD_TORCH_CHECK( + k_nope.device().is_cuda() && k_nope.dim() == 3 && k_nope.size(2) == 128, + "k_nope shape [Tp, H, 128] CUDA"); + STD_TORCH_CHECK(q.device().is_cuda() && q.dim() == 3 && q.size(2) == 192, + "q shape [Tp, H, 192] CUDA"); + // k_pe is a strided view of the fused QKV-LoRA GEMM output; the kernel takes + // its row stride and reads the 64 cols contiguously, so unit last-dim stride + // (not full contiguity) is all that is required. + STD_TORCH_CHECK(k_pe.device().is_cuda() && k_pe.dim() == 2 && + k_pe.stride(1) == 1 && k_pe.size(1) == 64, + "k_pe shape [Tp, 64], unit last-dim stride, CUDA"); + STD_TORCH_CHECK(kv_c_normed.device().is_cuda() && + kv_c_normed.is_contiguous() && kv_c_normed.dim() == 2 && + kv_c_normed.size(1) == 512, + "kv_c_normed shape [Tp, 512] contiguous CUDA"); + STD_TORCH_CHECK(k_out.device().is_cuda() && k_out.is_contiguous() && + k_out.dim() == 3 && k_out.size(2) == 192, + "k_out shape [Tp, H, 192] contiguous CUDA"); + STD_TORCH_CHECK(k_cache.device().is_cuda() && k_cache.dim() == 3 && + k_cache.size(1) == cache_block_size && + k_cache.size(2) == 576 && k_cache.stride(2) == 1, + "k_cache shape [nblk, block_size, 576] contiguous CUDA"); + STD_TORCH_CHECK(slot_mapping.device().is_cuda() && + slot_mapping.scalar_type() == ScalarType::Long, + "slot_mapping must be int64 CUDA"); + ScalarType const dt = k_nope.scalar_type(); + STD_TORCH_CHECK(q.scalar_type() == dt && k_pe.scalar_type() == dt && + kv_c_normed.scalar_type() == dt && + k_out.scalar_type() == dt && k_cache.scalar_type() == dt, + "all tensors must share k_nope's (bf16/fp16) dtype"); + + int const num_tokens = static_cast(k_nope.size(0)); + int const num_heads = static_cast(k_nope.size(1)); + STD_TORCH_CHECK(static_cast(k_out.size(1)) == num_heads, + "k_out head count must match k_nope"); + STD_TORCH_CHECK(q.size(0) == num_tokens && q.size(1) == num_heads, + "q token/head dimensions must match k_nope"); + bool const apply_rope = + check_rope_inputs(position_ids, cos_sin_cache, q, num_tokens); + if (num_tokens == 0) return; + + const torch::stable::accelerator::DeviceGuard device_guard( + k_nope.get_device_index()); + const cudaStream_t stream = + get_current_cuda_stream(k_nope.get_device_index()); + + APHRODITE_STABLE_DISPATCH_HALF_TYPES( + dt, "fused_kimi_k3_mla_key_concat_kv_cache_insert", [&] { + auto launch = [&](auto kernel) { + kk3::launchPdl( + kernel, num_tokens, num_heads, stream, + reinterpret_cast(q.mutable_data_ptr()), q.stride(0), + q.stride(1), + reinterpret_cast(k_nope.const_data_ptr()), + k_nope.stride(0), k_nope.stride(1), + reinterpret_cast(k_pe.const_data_ptr()), + k_pe.stride(0), + reinterpret_cast(kv_c_normed.const_data_ptr()), + kv_c_normed.stride(0), + reinterpret_cast(k_out.mutable_data_ptr()), + k_out.stride(0), k_out.stride(1), + reinterpret_cast(k_cache.mutable_data_ptr()), + k_cache.stride(0), k_cache.stride(1), + slot_mapping.const_data_ptr(), + apply_rope ? position_ids.value().const_data_ptr() + : nullptr, + apply_rope ? reinterpret_cast( + cos_sin_cache.value().const_data_ptr()) + : nullptr, + num_tokens, num_heads, static_cast(cache_block_size)); + }; + if (apply_rope) { + launch( + kk3::fusedKimiK3MLAKeyConcatKVCacheInsertKernel); + } else { + launch( + kk3::fusedKimiK3MLAKeyConcatKVCacheInsertKernel); + } + }); +} + +void fused_kimi_k3_mla_key_concat_ds_mla_insert( + torch::stable::Tensor& q, // [Tp, H, 192] + torch::stable::Tensor const& k_nope, // [Tp, H, 128] bf16 + torch::stable::Tensor const& k_pe, // [Tp, 64] bf16 + torch::stable::Tensor const& kv_c_normed, // [Tp, 512] bf16 + torch::stable::Tensor& k_out, // [Tp, H, 192] bf16, written + torch::stable::Tensor& k_cache, // [nblk, bs, 656] uint8, written + torch::stable::Tensor const& slot_mapping, // [Tp] int64 + int64_t cache_block_size, std::optional position_ids, + std::optional cos_sin_cache) { + using torch::headeronly::ScalarType; + namespace kk3 = aphrodite::kimi_k3_fused_ops; + ScalarType const dt = k_nope.scalar_type(); + STD_TORCH_CHECK( + k_nope.device().is_cuda() && k_nope.dim() == 3 && k_nope.size(2) == 128, + "k_nope shape [Tp, H, 128] CUDA"); + STD_TORCH_CHECK(q.device().is_cuda() && q.scalar_type() == dt && + q.dim() == 3 && q.size(2) == 192, + "q shape [Tp, H, 192]"); + STD_TORCH_CHECK(k_pe.device().is_cuda() && k_pe.dim() == 2 && + k_pe.stride(1) == 1 && k_pe.scalar_type() == dt && + k_pe.size(1) == 64, + "k_pe shape [Tp, 64], unit last-dim stride"); + STD_TORCH_CHECK(kv_c_normed.device().is_cuda() && + kv_c_normed.is_contiguous() && + kv_c_normed.scalar_type() == dt && + kv_c_normed.dim() == 2 && kv_c_normed.size(1) == 512, + "kv_c_normed shape [Tp, 512] contiguous"); + STD_TORCH_CHECK(k_out.device().is_cuda() && k_out.is_contiguous() && + k_out.scalar_type() == dt && k_out.dim() == 3 && + k_out.size(2) == 192, + "k_out shape [Tp, H, 192] contiguous"); + // fp8_ds_mla entry is 656 bytes stored as uint8. + STD_TORCH_CHECK( + k_cache.device().is_cuda() && k_cache.scalar_type() == ScalarType::Byte && + k_cache.dim() == 3 && k_cache.size(1) == cache_block_size && + k_cache.size(2) == 656 && k_cache.stride(2) == 1, + "k_cache shape [nblk, block_size, 656] uint8 contiguous"); + STD_TORCH_CHECK(slot_mapping.device().is_cuda() && + slot_mapping.scalar_type() == ScalarType::Long, + "slot_mapping must be int64 CUDA"); + + int const num_tokens = static_cast(k_nope.size(0)); + int const num_heads = static_cast(k_nope.size(1)); + STD_TORCH_CHECK(static_cast(k_out.size(1)) == num_heads, + "k_out head count must match k_nope"); + STD_TORCH_CHECK(q.size(0) == num_tokens && q.size(1) == num_heads, + "q token/head dimensions must match k_nope"); + bool const apply_rope = + check_rope_inputs(position_ids, cos_sin_cache, q, num_tokens); + if (num_tokens == 0) return; + + const torch::stable::accelerator::DeviceGuard device_guard( + k_nope.get_device_index()); + const cudaStream_t stream = + get_current_cuda_stream(k_nope.get_device_index()); + + APHRODITE_STABLE_DISPATCH_HALF_TYPES( + dt, "fused_kimi_k3_mla_key_concat_ds_mla_insert", [&] { + auto launch = [&](auto kernel) { + kk3::launchPdl( + kernel, num_tokens, num_heads, stream, + reinterpret_cast(q.mutable_data_ptr()), q.stride(0), + q.stride(1), + reinterpret_cast(k_nope.const_data_ptr()), + k_nope.stride(0), k_nope.stride(1), + reinterpret_cast(k_pe.const_data_ptr()), + k_pe.stride(0), + reinterpret_cast(kv_c_normed.const_data_ptr()), + kv_c_normed.stride(0), + reinterpret_cast(k_out.mutable_data_ptr()), + k_out.stride(0), k_out.stride(1), + reinterpret_cast(k_cache.mutable_data_ptr()), + k_cache.stride(0), k_cache.stride(1), + slot_mapping.const_data_ptr(), num_tokens, num_heads, + static_cast(cache_block_size), + apply_rope ? position_ids.value().const_data_ptr() + : nullptr, + apply_rope ? reinterpret_cast( + cos_sin_cache.value().const_data_ptr()) + : nullptr); + }; + if (apply_rope) { + launch(kk3::fusedKimiK3MLAKeyConcatDsMlaInsertKernel); + } else { + launch( + kk3::fusedKimiK3MLAKeyConcatDsMlaInsertKernel); + } + }); +} + +void fused_kimi_k3_mla_qkv_quant_kv_cache_fp8_insert( + torch::stable::Tensor const& q, // [Tp, H, 192] bf16 + torch::stable::Tensor const& k_nope, // [Tp, H, 128] bf16 + torch::stable::Tensor const& k_pe, // [Tp, 64] bf16 + torch::stable::Tensor const& kv_c_normed, // [Tp, 512] bf16 + torch::stable::Tensor const& v, // [Tp, H, 128] bf16 + torch::stable::Tensor& q_fp8, // [Tp, H, 192] fp8, written + torch::stable::Tensor& k_fp8, // [Tp, H, 192] fp8, written + torch::stable::Tensor& v_fp8, // [Tp, H, 128] fp8, written + torch::stable::Tensor& k_cache, // [nblk, bs, 576] fp8, written + torch::stable::Tensor const& slot_mapping, // [Tp] int64 + torch::stable::Tensor const& q_scale_inv, // scalar fp32 (1 / q scale) + torch::stable::Tensor const& k_scale_inv, // scalar fp32 (1 / k scale) + torch::stable::Tensor const& v_scale_inv, // scalar fp32 (1 / v scale) + torch::stable::Tensor const& cache_scale_inv, // scalar fp32 (1 / kv scale) + int64_t cache_block_size, std::optional position_ids, + std::optional cos_sin_cache) { + using torch::headeronly::ScalarType; + namespace kk3 = aphrodite::kimi_k3_fused_ops; + ScalarType const dt = k_nope.scalar_type(); + auto check_in = [&](torch::stable::Tensor const& t, int d2, char const* n) { + STD_TORCH_CHECK(t.device().is_cuda() && t.scalar_type() == dt && + t.dim() == 3 && t.size(2) == d2, + n); + }; + check_in(q, 192, "q shape [Tp, H, 192]"); + check_in(k_nope, 128, "k_nope shape [Tp, H, 128]"); + check_in(v, 128, "v shape [Tp, H, 128]"); + STD_TORCH_CHECK(k_pe.device().is_cuda() && k_pe.dim() == 2 && + k_pe.stride(1) == 1 && k_pe.scalar_type() == dt && + k_pe.size(1) == 64, + "k_pe shape [Tp, 64], unit last-dim stride"); + STD_TORCH_CHECK(kv_c_normed.device().is_cuda() && + kv_c_normed.is_contiguous() && + kv_c_normed.scalar_type() == dt && + kv_c_normed.dim() == 2 && kv_c_normed.size(1) == 512, + "kv_c_normed shape [Tp, 512] contiguous"); + auto check_out = [&](torch::stable::Tensor const& t, int d2, char const* n) { + STD_TORCH_CHECK(t.device().is_cuda() && t.is_contiguous() && + t.scalar_type() == ScalarType::Float8_e4m3fn && + t.dim() == 3 && t.size(2) == d2, + n); + }; + check_out(q_fp8, 192, "q_fp8 shape [Tp, H, 192] fp8 contiguous"); + check_out(k_fp8, 192, "k_fp8 shape [Tp, H, 192] fp8 contiguous"); + check_out(v_fp8, 128, "v_fp8 shape [Tp, H, 128] fp8 contiguous"); + STD_TORCH_CHECK(k_cache.device().is_cuda() && k_cache.dim() == 3 && + k_cache.size(1) == cache_block_size && + k_cache.size(2) == 576 && k_cache.stride(2) == 1 && + k_cache.scalar_type() == ScalarType::Float8_e4m3fn, + "k_cache shape [nblk, block_size, 576] fp8 contiguous"); + STD_TORCH_CHECK(slot_mapping.device().is_cuda() && + slot_mapping.scalar_type() == ScalarType::Long, + "slot_mapping must be int64 CUDA"); + auto check_scale = [&](torch::stable::Tensor const& s, char const* n) { + STD_TORCH_CHECK(s.device().is_cuda() && + s.scalar_type() == ScalarType::Float && s.size(0) == 1, + n); + }; + check_scale(q_scale_inv, "q_scale_inv must be scalar float32 CUDA"); + check_scale(k_scale_inv, "k_scale_inv must be scalar float32 CUDA"); + check_scale(v_scale_inv, "v_scale_inv must be scalar float32 CUDA"); + check_scale(cache_scale_inv, "cache_scale_inv must be scalar float32 CUDA"); + + int const num_tokens = static_cast(k_nope.size(0)); + int const num_heads = static_cast(k_nope.size(1)); + bool const apply_rope = + check_rope_inputs(position_ids, cos_sin_cache, q, num_tokens); + if (num_tokens == 0) return; + + const torch::stable::accelerator::DeviceGuard device_guard( + k_nope.get_device_index()); + const cudaStream_t stream = + get_current_cuda_stream(k_nope.get_device_index()); + + APHRODITE_STABLE_DISPATCH_HALF_TYPES( + dt, "fused_kimi_k3_mla_qkv_quant_kv_cache_fp8_insert", [&] { + auto launch = [&](auto kernel) { + kk3::launchPdl( + kernel, num_tokens, num_heads, stream, + reinterpret_cast(q.const_data_ptr()), + q.stride(0), q.stride(1), + reinterpret_cast(k_nope.const_data_ptr()), + k_nope.stride(0), k_nope.stride(1), + reinterpret_cast(k_pe.const_data_ptr()), + k_pe.stride(0), + reinterpret_cast(kv_c_normed.const_data_ptr()), + kv_c_normed.stride(0), + reinterpret_cast(v.const_data_ptr()), + v.stride(0), v.stride(1), + reinterpret_cast(q_fp8.mutable_data_ptr()), + q_fp8.stride(0), q_fp8.stride(1), + reinterpret_cast(k_fp8.mutable_data_ptr()), + k_fp8.stride(0), k_fp8.stride(1), + reinterpret_cast(v_fp8.mutable_data_ptr()), + v_fp8.stride(0), v_fp8.stride(1), + reinterpret_cast(k_cache.mutable_data_ptr()), + k_cache.stride(0), k_cache.stride(1), + slot_mapping.const_data_ptr(), + q_scale_inv.const_data_ptr(), + k_scale_inv.const_data_ptr(), + v_scale_inv.const_data_ptr(), + cache_scale_inv.const_data_ptr(), num_tokens, num_heads, + static_cast(cache_block_size), + apply_rope ? position_ids.value().const_data_ptr() + : nullptr, + apply_rope ? reinterpret_cast( + cos_sin_cache.value().const_data_ptr()) + : nullptr); + }; + if (apply_rope) { + launch(kk3::fusedKimiK3MLAQKVQuantKVCacheFp8Kernel); + } else { + launch(kk3::fusedKimiK3MLAQKVQuantKVCacheFp8Kernel); + } + }); +} + +// ──────────────────────────────────────────────────────────────────────────── +// Decode epilogue torch ops +// ──────────────────────────────────────────────────────────────────────────── +namespace { +// Shared shape checks for the decode ops. Verifies the query/latent inputs and +// slot_mapping; caller checks mqa_q / k_cache dtypes for its variant. +void check_decode_inputs(torch::stable::Tensor const& ql_nope, + torch::stable::Tensor const& q_pe, + torch::stable::Tensor const& kv_c_normed, + torch::stable::Tensor const& k_pe, + torch::stable::Tensor const& mqa_q, + torch::stable::Tensor const& slot_mapping) { + using torch::headeronly::ScalarType; + auto const dt = ql_nope.scalar_type(); + STD_TORCH_CHECK(ql_nope.device().is_cuda() && ql_nope.dim() == 3 && + ql_nope.size(2) == 512, + "ql_nope shape [B, H, 512] CUDA"); + STD_TORCH_CHECK(q_pe.device().is_cuda() && q_pe.scalar_type() == dt && + q_pe.dim() == 3 && q_pe.size(2) == 64, + "q_pe shape [B, H, 64]"); + STD_TORCH_CHECK(kv_c_normed.device().is_cuda() && + kv_c_normed.is_contiguous() && + kv_c_normed.scalar_type() == dt && + kv_c_normed.dim() == 2 && kv_c_normed.size(1) == 512, + "kv_c_normed shape [B, 512] contiguous"); + STD_TORCH_CHECK(k_pe.device().is_cuda() && k_pe.dim() == 2 && + k_pe.stride(1) == 1 && k_pe.scalar_type() == dt && + k_pe.size(1) == 64, + "k_pe shape [B, 64], unit last-dim stride"); + STD_TORCH_CHECK(mqa_q.device().is_cuda() && mqa_q.is_contiguous() && + mqa_q.dim() == 3 && mqa_q.size(2) == 576, + "mqa_q shape [B, H, 576] contiguous"); + STD_TORCH_CHECK(slot_mapping.device().is_cuda() && + slot_mapping.scalar_type() == ScalarType::Long, + "slot_mapping must be int64 CUDA"); +} +} // namespace + +void fused_kimi_k3_mla_decode_q_concat_kv_cache_insert( + torch::stable::Tensor const& ql_nope, // [B, H, 512] bf16 + torch::stable::Tensor const& q_pe, // [B, H, 64] bf16 + torch::stable::Tensor const& kv_c_normed, // [B, 512] bf16 + torch::stable::Tensor const& k_pe, // [B, 64] bf16 + torch::stable::Tensor& mqa_q, // [B, H, 576] bf16, written + torch::stable::Tensor& k_cache, // [nblk, bs, 576] bf16, written + torch::stable::Tensor const& slot_mapping, // [B] int64 + int64_t cache_block_size, std::optional position_ids, + std::optional cos_sin_cache) { + using torch::headeronly::ScalarType; + namespace kk3 = aphrodite::kimi_k3_fused_ops; + ScalarType const dt = ql_nope.scalar_type(); + check_decode_inputs(ql_nope, q_pe, kv_c_normed, k_pe, mqa_q, slot_mapping); + STD_TORCH_CHECK(mqa_q.scalar_type() == dt && k_cache.scalar_type() == dt, + "mqa_q / k_cache must match ql_nope dtype (bf16)"); + STD_TORCH_CHECK(k_cache.device().is_cuda() && k_cache.dim() == 3 && + k_cache.size(1) == cache_block_size && + k_cache.size(2) == 576 && k_cache.stride(2) == 1, + "k_cache shape [nblk, block_size, 576] contiguous"); + + int const num_tokens = static_cast(ql_nope.size(0)); + int const num_heads = static_cast(ql_nope.size(1)); + bool const apply_rope = + check_rope_inputs(position_ids, cos_sin_cache, q_pe, num_tokens); + if (num_tokens == 0) return; + const torch::stable::accelerator::DeviceGuard device_guard( + ql_nope.get_device_index()); + const cudaStream_t stream = + get_current_cuda_stream(ql_nope.get_device_index()); + + APHRODITE_STABLE_DISPATCH_HALF_TYPES( + dt, "fused_kimi_k3_mla_decode_q_concat_kv_cache_insert", [&] { + auto launch = [&](auto kernel) { + kk3::launchPdl( + kernel, num_tokens, num_heads, stream, + reinterpret_cast(ql_nope.const_data_ptr()), + ql_nope.stride(0), ql_nope.stride(1), + reinterpret_cast(q_pe.const_data_ptr()), + q_pe.stride(0), q_pe.stride(1), + reinterpret_cast(kv_c_normed.const_data_ptr()), + kv_c_normed.stride(0), + reinterpret_cast(k_pe.const_data_ptr()), + k_pe.stride(0), mqa_q.mutable_data_ptr(), mqa_q.stride(0), + mqa_q.stride(1), k_cache.mutable_data_ptr(), k_cache.stride(0), + k_cache.stride(1), slot_mapping.const_data_ptr(), + nullptr, nullptr, num_tokens, num_heads, + static_cast(cache_block_size), + apply_rope ? position_ids.value().const_data_ptr() + : nullptr, + apply_rope ? reinterpret_cast( + cos_sin_cache.value().const_data_ptr()) + : nullptr); + }; + if (apply_rope) { + launch(kk3::fusedKimiK3MLADecodeQConcatKVCacheKernel); + } else { + launch(kk3::fusedKimiK3MLADecodeQConcatKVCacheKernel); + } + }); +} + +void fused_kimi_k3_mla_decode_q_concat_kv_cache_fp8_insert( + torch::stable::Tensor const& ql_nope, // [B, H, 512] bf16 + torch::stable::Tensor const& q_pe, // [B, H, 64] bf16 + torch::stable::Tensor const& kv_c_normed, // [B, 512] bf16 + torch::stable::Tensor const& k_pe, // [B, 64] bf16 + torch::stable::Tensor& mqa_q, // [B, H, 576] fp8, written + torch::stable::Tensor& k_cache, // [nblk, bs, 576] fp8, written + torch::stable::Tensor const& slot_mapping, // [B] int64 + torch::stable::Tensor const& q_scale_inv, // scalar fp32 (1 / q scale) + torch::stable::Tensor const& cache_scale_inv, // scalar fp32 (1 / kv scale) + int64_t cache_block_size, std::optional position_ids, + std::optional cos_sin_cache) { + using torch::headeronly::ScalarType; + namespace kk3 = aphrodite::kimi_k3_fused_ops; + ScalarType const dt = ql_nope.scalar_type(); + check_decode_inputs(ql_nope, q_pe, kv_c_normed, k_pe, mqa_q, slot_mapping); + STD_TORCH_CHECK(mqa_q.scalar_type() == ScalarType::Float8_e4m3fn, + "mqa_q must be float8_e4m3fn"); + STD_TORCH_CHECK(k_cache.device().is_cuda() && k_cache.dim() == 3 && + k_cache.size(1) == cache_block_size && + k_cache.size(2) == 576 && k_cache.stride(2) == 1 && + k_cache.scalar_type() == ScalarType::Float8_e4m3fn, + "k_cache shape [nblk, block_size, 576] fp8 contiguous"); + auto check_scale = [&](torch::stable::Tensor const& s, char const* n) { + STD_TORCH_CHECK(s.device().is_cuda() && + s.scalar_type() == ScalarType::Float && s.size(0) == 1, + n); + }; + check_scale(q_scale_inv, "q_scale_inv must be scalar float32 CUDA"); + check_scale(cache_scale_inv, "cache_scale_inv must be scalar float32 CUDA"); + + int const num_tokens = static_cast(ql_nope.size(0)); + int const num_heads = static_cast(ql_nope.size(1)); + bool const apply_rope = + check_rope_inputs(position_ids, cos_sin_cache, q_pe, num_tokens); + if (num_tokens == 0) return; + const torch::stable::accelerator::DeviceGuard device_guard( + ql_nope.get_device_index()); + const cudaStream_t stream = + get_current_cuda_stream(ql_nope.get_device_index()); + + APHRODITE_STABLE_DISPATCH_HALF_TYPES( + dt, "fused_kimi_k3_mla_decode_q_concat_kv_cache_fp8_insert", [&] { + auto launch = [&](auto kernel) { + kk3::launchPdl( + kernel, num_tokens, num_heads, stream, + reinterpret_cast(ql_nope.const_data_ptr()), + ql_nope.stride(0), ql_nope.stride(1), + reinterpret_cast(q_pe.const_data_ptr()), + q_pe.stride(0), q_pe.stride(1), + reinterpret_cast(kv_c_normed.const_data_ptr()), + kv_c_normed.stride(0), + reinterpret_cast(k_pe.const_data_ptr()), + k_pe.stride(0), mqa_q.mutable_data_ptr(), mqa_q.stride(0), + mqa_q.stride(1), k_cache.mutable_data_ptr(), k_cache.stride(0), + k_cache.stride(1), slot_mapping.const_data_ptr(), + q_scale_inv.const_data_ptr(), + cache_scale_inv.const_data_ptr(), num_tokens, num_heads, + static_cast(cache_block_size), + apply_rope ? position_ids.value().const_data_ptr() + : nullptr, + apply_rope ? reinterpret_cast( + cos_sin_cache.value().const_data_ptr()) + : nullptr); + }; + if (apply_rope) { + launch(kk3::fusedKimiK3MLADecodeQConcatKVCacheKernel); + } else { + launch(kk3::fusedKimiK3MLADecodeQConcatKVCacheKernel); + } + }); +} + +void fused_kimi_k3_mla_decode_q_concat_ds_mla_insert( + torch::stable::Tensor const& ql_nope, // [B, H, 512] bf16 + torch::stable::Tensor const& q_pe, // [B, H, 64] bf16 + torch::stable::Tensor const& kv_c_normed, // [B, 512] bf16 + torch::stable::Tensor const& k_pe, // [B, 64] bf16 + torch::stable::Tensor& mqa_q, // [B, H, 576] bf16, written + torch::stable::Tensor& k_cache, // [nblk, bs, 656] uint8, written + torch::stable::Tensor const& slot_mapping, // [B] int64 + int64_t cache_block_size, std::optional position_ids, + std::optional cos_sin_cache) { + using torch::headeronly::ScalarType; + namespace kk3 = aphrodite::kimi_k3_fused_ops; + ScalarType const dt = ql_nope.scalar_type(); + check_decode_inputs(ql_nope, q_pe, kv_c_normed, k_pe, mqa_q, slot_mapping); + STD_TORCH_CHECK(mqa_q.scalar_type() == dt, "mqa_q must be bf16 for ds_mla"); + STD_TORCH_CHECK( + k_cache.device().is_cuda() && k_cache.scalar_type() == ScalarType::Byte && + k_cache.dim() == 3 && k_cache.size(1) == cache_block_size && + k_cache.size(2) == 656 && k_cache.stride(2) == 1, + "k_cache shape [nblk, block_size, 656] uint8 contiguous"); + + int const num_tokens = static_cast(ql_nope.size(0)); + int const num_heads = static_cast(ql_nope.size(1)); + bool const apply_rope = + check_rope_inputs(position_ids, cos_sin_cache, q_pe, num_tokens); + if (num_tokens == 0) return; + const torch::stable::accelerator::DeviceGuard device_guard( + ql_nope.get_device_index()); + const cudaStream_t stream = + get_current_cuda_stream(ql_nope.get_device_index()); + + APHRODITE_STABLE_DISPATCH_HALF_TYPES( + dt, "fused_kimi_k3_mla_decode_q_concat_ds_mla_insert", [&] { + auto launch = [&](auto kernel) { + kk3::launchPdl( + kernel, num_tokens, num_heads, stream, + reinterpret_cast(ql_nope.const_data_ptr()), + ql_nope.stride(0), ql_nope.stride(1), + reinterpret_cast(q_pe.const_data_ptr()), + q_pe.stride(0), q_pe.stride(1), + reinterpret_cast(kv_c_normed.const_data_ptr()), + kv_c_normed.stride(0), + reinterpret_cast(k_pe.const_data_ptr()), + k_pe.stride(0), + reinterpret_cast(mqa_q.mutable_data_ptr()), + mqa_q.stride(0), mqa_q.stride(1), + reinterpret_cast(k_cache.mutable_data_ptr()), + k_cache.stride(0), k_cache.stride(1), + slot_mapping.const_data_ptr(), num_tokens, num_heads, + static_cast(cache_block_size), + apply_rope ? position_ids.value().const_data_ptr() + : nullptr, + apply_rope ? reinterpret_cast( + cos_sin_cache.value().const_data_ptr()) + : nullptr); + }; + if (apply_rope) { + launch(kk3::fusedKimiK3MLADecodeQConcatDsMlaKernel); + } else { + launch(kk3::fusedKimiK3MLADecodeQConcatDsMlaKernel); + } + }); +} diff --git a/csrc/libtorch_stable/kimi_k3/fused_kda_decode_kernel.cu b/csrc/libtorch_stable/kimi_k3/fused_kda_decode_kernel.cu new file mode 100644 index 0000000000..0badeb9b77 --- /dev/null +++ b/csrc/libtorch_stable/kimi_k3/fused_kda_decode_kernel.cu @@ -0,0 +1,1130 @@ +/* + * Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + */ + +#include +#include +#include +#include +#include + +#include "../torch_utils.h" +#include "../../cuda_compat.h" + +namespace { + +constexpr int kDimK = 128; +constexpr int kDimV = 128; +constexpr int kKernelWidth = 4; +constexpr int kConvStateWidth = kKernelWidth - 1; +constexpr int kThreads = 256; +constexpr int kWarps = kThreads / 32; +constexpr int kChunkV = 32; +constexpr int kNumChunks = kDimV / kChunkV; +constexpr int kRowsPerWarp = kChunkV / kWarps; + +struct KdaDecodeStrides { + int64_t x_row; + int64_t beta_row; + int64_t onorm_row; + int64_t conv_slot; + int64_t state_slot; +}; + +__device__ __forceinline__ float bf16_load(const __nv_bfloat16* ptr, + int64_t idx) { + return __bfloat162float(ptr[idx]); +} + +__device__ __forceinline__ float bf16_load(const float* ptr, int64_t idx) { + return ptr[idx]; +} + +template +__device__ __forceinline__ float conv_weight_load(const float* ptr, int channel, + int width) { + return ptr[width * kChannels + channel]; +} + +__device__ __forceinline__ __nv_bfloat16 bf16_store(float value) { + return __float2bfloat16(value); +} + +template +__device__ __forceinline__ void store_state_float4(float* ptr, float4 value) { + if constexpr (kUseCacheGlobalStore) { + __stcg(reinterpret_cast(ptr), value); + } else { + *reinterpret_cast(ptr) = value; + } +} + +__device__ __forceinline__ float sigmoid_fast(float x) { + return 1.0f / (1.0f + __expf(-x)); +} + +__device__ __forceinline__ float silu_fast(float x) { + return x * sigmoid_fast(x); +} + +__device__ __forceinline__ float softplus_fast(float x) { + return x > 20.0f ? x : log1pf(__expf(x)); +} + +__device__ __forceinline__ float warp_reduce_sum(float value) { +#pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) { + value += __shfl_xor_sync(0xffffffffu, value, offset); + } + return value; +} + +__device__ __forceinline__ void cp_async_cg_16b(float* smem_ptr, + const float* gmem_ptr) { + uint32_t smem_addr = + static_cast(__cvta_generic_to_shared(smem_ptr)); + asm volatile("cp.async.cg.shared.global [%0], [%1], 16;\n" + : + : "r"(smem_addr), "l"(gmem_ptr)); +} + +__device__ __forceinline__ void cp_async_commit() { + asm volatile("cp.async.commit_group;\n" ::); +} + +__device__ __forceinline__ void cp_async_wait_all() { + asm volatile("cp.async.wait_all;\n" ::: "memory"); +} + +__device__ __forceinline__ void cp_async_wait_group_1() { + asm volatile("cp.async.wait_group 1;\n" ::: "memory"); +} + +template +__device__ __forceinline__ void cp_async_state_chunk_for(float* s_state, + const float* state, + int slot, int i_hv, + int HV, int chunk) { + constexpr int kFloat4PerChunk = kChunkV * kDimK / 4; + const int tid = threadIdx.x; + const int stage = chunk & 1; + const int v_base = chunk * kChunkV; + for (int linear4 = tid; linear4 < kFloat4PerChunk; linear4 += kCopyThreads) { + const int elem = linear4 * 4; + const int row = elem / kDimK; + const int k = elem - row * kDimK; + float* dst = s_state + (stage * kChunkV + row) * kDimK + k; + const float* src = + state + ((slot * HV + i_hv) * kDimV + v_base + row) * kDimK + k; + cp_async_cg_16b(dst, src); + } + cp_async_commit(); +} + +__device__ __forceinline__ void cp_async_state_chunk(float* s_state, + const float* state, + int slot, int i_hv, int HV, + int chunk) { + cp_async_state_chunk_for(s_state, state, slot, i_hv, HV, chunk); +} + +__device__ __forceinline__ float block_reduce_sum(float value, float* scratch) { + const int tid = threadIdx.x; + const int lane = tid & 31; + const int warp = tid >> 5; + + float warp_total = warp_reduce_sum(value); + if (lane == 0) { + scratch[warp] = warp_total; + } + __syncthreads(); + + float block_total = 0.0f; + if (warp == 0) { + block_total = lane < kWarps ? scratch[lane] : 0.0f; + block_total = warp_reduce_sum(block_total); + if (lane == 0) { + scratch[0] = block_total; + } + } + __syncthreads(); + return scratch[0]; +} + +struct Sum2 { + float x; + float y; +}; + +__device__ __forceinline__ Sum2 warp_reduce_sum_pair(float x, float y) { +#pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) { + x += __shfl_xor_sync(0xffffffffu, x, offset); + y += __shfl_xor_sync(0xffffffffu, y, offset); + } + return {x, y}; +} + +template +__device__ __forceinline__ Sum2 block_reduce_sum2_for(float x, float y, + float* scratch) { + const int lane = threadIdx.x & 31; + const int warp = threadIdx.x >> 5; + + const float warp_x = warp_reduce_sum(x); + const float warp_y = warp_reduce_sum(y); + if (lane == 0) { + scratch[warp] = warp_x; + scratch[kReduceWarps + warp] = warp_y; + } + __syncthreads(); + + float block_x = 0.0f; + float block_y = 0.0f; + if (warp == 0) { + block_x = lane < kReduceWarps ? scratch[lane] : 0.0f; + block_y = lane < kReduceWarps ? scratch[kReduceWarps + lane] : 0.0f; + block_x = warp_reduce_sum(block_x); + block_y = warp_reduce_sum(block_y); + if (lane == 0) { + scratch[0] = block_x; + scratch[1] = block_y; + } + } + __syncthreads(); + return {scratch[0], scratch[1]}; +} + +__device__ __forceinline__ Sum2 block_reduce_sum2(float x, float y, + float* scratch) { + return block_reduce_sum2_for(x, y, scratch); +} + +template +__device__ __forceinline__ float block_reduce_sum_active_for(float value, + float* scratch) { + const int lane = threadIdx.x & 31; + const int warp = threadIdx.x >> 5; + + float warp_total = 0.0f; + if (warp < kReduceWarps) { + warp_total = warp_reduce_sum(value); + } + if (lane == 0 && warp < kReduceWarps) { + scratch[warp] = warp_total; + } + __syncthreads(); + + float block_total = 0.0f; + if (warp == 0) { + block_total = lane < kReduceWarps ? scratch[lane] : 0.0f; + block_total = warp_reduce_sum(block_total); + if (lane == 0) { + scratch[0] = block_total; + } + } + __syncthreads(); + return scratch[0]; +} + +template +__device__ __forceinline__ Sum2 block_reduce_sum2_active_for(float x, float y, + float* scratch) { + const int lane = threadIdx.x & 31; + const int warp = threadIdx.x >> 5; + + float warp_x = 0.0f; + float warp_y = 0.0f; + if (warp < kReduceWarps) { + warp_x = warp_reduce_sum(x); + warp_y = warp_reduce_sum(y); + } + if (lane == 0 && warp < kReduceWarps) { + scratch[warp] = warp_x; + scratch[kReduceWarps + warp] = warp_y; + } + __syncthreads(); + + float block_x = 0.0f; + float block_y = 0.0f; + if (warp == 0) { + block_x = lane < kReduceWarps ? scratch[lane] : 0.0f; + block_y = lane < kReduceWarps ? scratch[kReduceWarps + lane] : 0.0f; + block_x = warp_reduce_sum(block_x); + block_y = warp_reduce_sum(block_y); + if (lane == 0) { + scratch[0] = block_x; + scratch[1] = block_y; + } + } + __syncthreads(); + return {scratch[0], scratch[1]}; +} + +template +__global__ +__launch_bounds__(kThreads, 2) void kda_decode_fusion_many_heads_kernel( + const __nv_bfloat16* __restrict__ x_q, + const __nv_bfloat16* __restrict__ x_k, + const __nv_bfloat16* __restrict__ x_v, const float* __restrict__ w_q_t, + const float* __restrict__ w_k_t, const float* __restrict__ w_v_t, + const float* __restrict__ bias_q, const float* __restrict__ bias_k, + const float* __restrict__ bias_v, __nv_bfloat16* __restrict__ cs_q, + __nv_bfloat16* __restrict__ cs_k, __nv_bfloat16* __restrict__ cs_v, + const float* __restrict__ a_log, const __nv_bfloat16* __restrict__ g, + const float* __restrict__ dt_bias, const __nv_bfloat16* __restrict__ beta, + const __nv_bfloat16* __restrict__ onorm_g, + const float* __restrict__ onorm_weight, + const int* __restrict__ ssm_state_indices, + const int* __restrict__ cu_seqlens, float* __restrict__ state, + __nv_bfloat16* __restrict__ out, int B, int H, int HV, float lower_bound, + float scale, float onorm_eps, KdaDecodeStrides strides) { + const int tid = threadIdx.x; + const int lane = tid & 31; + const int warp = tid >> 5; + int i_n; + int i_hv; + int i_h; + int bos; + int slot; +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900 + cudaGridDependencySynchronize(); +#endif + if constexpr (kUseStaticDecodeLayout) { + if constexpr (kUseHeadGrid) { + i_n = blockIdx.x; + i_hv = blockIdx.y; + } else { + const int nhv = blockIdx.x; + i_n = nhv / kFixedValueHeads; + i_hv = nhv - i_n * kFixedValueHeads; + } + i_h = i_hv; + bos = i_n; + slot = ssm_state_indices == nullptr ? i_n : ssm_state_indices[i_n]; + } else { + const int nhv = blockIdx.x; + i_n = nhv / HV; + i_hv = nhv - i_n * HV; + const int hv_per_h = HV / H; + i_h = i_hv / hv_per_h; + + bos = cu_seqlens == nullptr ? i_n : cu_seqlens[i_n]; + const int eos = cu_seqlens == nullptr ? i_n + 1 : cu_seqlens[i_n + 1]; + if (eos <= bos) { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900 + cudaTriggerProgrammaticLaunchCompletion(); +#endif + return; + } + slot = ssm_state_indices == nullptr ? i_n : ssm_state_indices[i_n]; + } + + constexpr int kLocalDim = kFixedHeads * kDimK; + constexpr int kPackedDim = 3 * kLocalDim; + const int hk_off = i_h * kDimK; + const int hv_off = i_hv * kDimV; + constexpr int hv_count = kFixedValueHeads; + float* const state_for_slot = state + slot * strides.state_slot; + const int64_t conv_slot_offset = slot * strides.conv_slot; + __nv_bfloat16* const cs_q_for_slot = cs_q + conv_slot_offset; + __nv_bfloat16* const cs_k_for_slot = cs_k + conv_slot_offset; + __nv_bfloat16* const cs_v_for_slot = cs_v + conv_slot_offset; + + __shared__ float s_state[2][kChunkV][kDimK]; + __shared__ float s_q[kDimK]; + __shared__ float s_k[kDimK]; + __shared__ float s_decay[kDimK]; + __shared__ float s_v[kDimV]; + __shared__ float s_o[kDimV]; + __shared__ float s_reduce[kThreads]; + __shared__ float s_beta; + float pre_onorm_gate = 0.0f; + float pre_onorm_weight = 0.0f; + + cp_async_state_chunk(&s_state[0][0][0], state_for_slot, 0, i_hv, hv_count, 0); + + if constexpr (kUpdateConvState) { + if (tid < kDimK) { + const int k = tid; + const int hk = hk_off + k; + const int64_t xq_idx = bos * strides.x_row + i_h * kDimK + k; + const float exp_a = + __shfl_sync(0xffffffffu, lane == 0 ? __expf(a_log[i_h]) : 0.0f, 0); + + float q_acc = bias_q == nullptr ? 0.0f : bf16_load(bias_q, hk); + float k_acc = bias_k == nullptr ? 0.0f : bf16_load(bias_k, hk); + __nv_bfloat16 q_shift0 = __float2bfloat16(0.0f); + __nv_bfloat16 q_shift1 = __float2bfloat16(0.0f); + __nv_bfloat16 k_shift0 = __float2bfloat16(0.0f); + __nv_bfloat16 k_shift1 = __float2bfloat16(0.0f); +#pragma unroll + for (int w = 0; w < kConvStateWidth; ++w) { + const __nv_bfloat16 q_state = cs_q_for_slot[hk + w * kPackedDim]; + const __nv_bfloat16 k_state = cs_k_for_slot[hk + w * kPackedDim]; + q_acc += __bfloat162float(q_state) * + conv_weight_load(w_q_t, hk, w); + k_acc += __bfloat162float(k_state) * + conv_weight_load(w_k_t, hk, w); + if (w == 1) { + q_shift0 = q_state; + k_shift0 = k_state; + } else if (w == 2) { + q_shift1 = q_state; + k_shift1 = k_state; + } + } + const __nv_bfloat16 q_new = x_q[xq_idx]; + const __nv_bfloat16 k_new = x_k[xq_idx]; + q_acc += __bfloat162float(q_new) * + conv_weight_load(w_q_t, hk, kKernelWidth - 1); + k_acc += __bfloat162float(k_new) * + conv_weight_load(w_k_t, hk, kKernelWidth - 1); + + cs_q_for_slot[hk] = q_shift0; + cs_q_for_slot[hk + kPackedDim] = q_shift1; + cs_q_for_slot[hk + 2 * kPackedDim] = q_new; + cs_k_for_slot[hk] = k_shift0; + cs_k_for_slot[hk + kPackedDim] = k_shift1; + cs_k_for_slot[hk + 2 * kPackedDim] = k_new; + + s_q[k] = silu_fast(q_acc); + s_k[k] = silu_fast(k_acc); + + const int64_t gate_idx = bos * kLocalDim + i_hv * kDimK + k; + const float g_raw = bf16_load(g, gate_idx) + dt_bias[hk]; + if constexpr (kUseLowerBound) { + s_decay[k] = __expf(lower_bound * sigmoid_fast(exp_a * g_raw)); + } else { + s_decay[k] = __expf(-exp_a * softplus_fast(g_raw)); + } + } + } else { + if (tid < kDimK) { + const int k = tid; + const int hk = hk_off + k; + const float exp_a = + __shfl_sync(0xffffffffu, lane == 0 ? __expf(a_log[i_h]) : 0.0f, 0); + + float q_acc = bias_q == nullptr ? 0.0f : bf16_load(bias_q, hk); + float k_acc = bias_k == nullptr ? 0.0f : bf16_load(bias_k, hk); +#pragma unroll + for (int w = 0; w < kConvStateWidth; ++w) { + const int cs_idx = hk + w * kPackedDim; + q_acc += bf16_load(cs_q_for_slot, cs_idx) * + conv_weight_load(w_q_t, hk, w); + k_acc += bf16_load(cs_k_for_slot, cs_idx) * + conv_weight_load(w_k_t, hk, w); + } + q_acc += bf16_load(x_q, bos * strides.x_row + i_h * kDimK + k) * + conv_weight_load(w_q_t, hk, kKernelWidth - 1); + k_acc += bf16_load(x_k, bos * strides.x_row + i_h * kDimK + k) * + conv_weight_load(w_k_t, hk, kKernelWidth - 1); + + s_q[k] = silu_fast(q_acc); + s_k[k] = silu_fast(k_acc); + + const int64_t gate_idx = bos * kLocalDim + i_hv * kDimK + k; + const float g_raw = bf16_load(g, gate_idx) + dt_bias[hk]; + if constexpr (kUseLowerBound) { + s_decay[k] = __expf(lower_bound * sigmoid_fast(exp_a * g_raw)); + } else { + s_decay[k] = __expf(-exp_a * softplus_fast(g_raw)); + } + } + } + + if constexpr (kUpdateConvState) { + if (tid < kDimV) { + const int v = tid; + const int hvv = hv_off + v; + const int64_t xv_idx = bos * strides.x_row + i_hv * kDimV + v; + + float v_acc = bias_v == nullptr ? 0.0f : bf16_load(bias_v, hvv); + __nv_bfloat16 v_shift0 = __float2bfloat16(0.0f); + __nv_bfloat16 v_shift1 = __float2bfloat16(0.0f); +#pragma unroll + for (int w = 0; w < kConvStateWidth; ++w) { + const __nv_bfloat16 v_state = cs_v_for_slot[hvv + w * kPackedDim]; + v_acc += __bfloat162float(v_state) * + conv_weight_load(w_v_t, hvv, w); + if (w == 1) { + v_shift0 = v_state; + } else if (w == 2) { + v_shift1 = v_state; + } + } + const __nv_bfloat16 v_new = x_v[xv_idx]; + v_acc += __bfloat162float(v_new) * + conv_weight_load(w_v_t, hvv, kKernelWidth - 1); + cs_v_for_slot[hvv] = v_shift0; + cs_v_for_slot[hvv + kPackedDim] = v_shift1; + cs_v_for_slot[hvv + 2 * kPackedDim] = v_new; + s_v[v] = silu_fast(v_acc); + + if constexpr (kApplyOnorm && kPreloadOnormParams) { + const int64_t gate_idx = i_n * strides.onorm_row + i_hv * kDimV + v; + pre_onorm_gate = sigmoid_fast(bf16_load(onorm_g, gate_idx)); + pre_onorm_weight = onorm_weight[v]; + } + } + } else { + if (tid < kDimV) { + const int v = tid; + const int hvv = hv_off + v; + + float v_acc = bias_v == nullptr ? 0.0f : bf16_load(bias_v, hvv); +#pragma unroll + for (int w = 0; w < kConvStateWidth; ++w) { + const int cs_idx = hvv + w * kPackedDim; + v_acc += bf16_load(cs_v_for_slot, cs_idx) * + conv_weight_load(w_v_t, hvv, w); + } + v_acc += bf16_load(x_v, bos * strides.x_row + i_hv * kDimV + v) * + conv_weight_load(w_v_t, hvv, kKernelWidth - 1); + s_v[v] = silu_fast(v_acc); + + if constexpr (kApplyOnorm && kPreloadOnormParams) { + const int64_t gate_idx = i_n * strides.onorm_row + i_hv * kDimV + v; + pre_onorm_gate = sigmoid_fast(bf16_load(onorm_g, gate_idx)); + pre_onorm_weight = onorm_weight[v]; + } + } + } + + if (tid == 0) { + const float beta_raw = bf16_load(beta, bos * strides.beta_row + i_hv); + if constexpr (kApplyBetaSigmoid) { + s_beta = sigmoid_fast(beta_raw); + } else { + s_beta = beta_raw; + } + } + __syncthreads(); + + if constexpr (kPrefetchNextStateChunk && kNumChunks > 1) { + cp_async_state_chunk(&s_state[0][0][0], state_for_slot, 0, i_hv, hv_count, + 1); + } + + const float q_sq = tid < kDimK ? s_q[tid] * s_q[tid] : 0.0f; + const float k_sq = tid < kDimK ? s_k[tid] * s_k[tid] : 0.0f; + Sum2 qk_sum; + if constexpr (kUseActiveQkReduction) { + qk_sum = block_reduce_sum2_active_for(q_sq, k_sq, s_reduce); + } else { + qk_sum = block_reduce_sum2(q_sq, k_sq, s_reduce); + } + if (tid < kDimK) { + s_q[tid] *= rsqrtf(qk_sum.x + 1.0e-6f) * scale; + s_k[tid] *= rsqrtf(qk_sum.y + 1.0e-6f); + } + __syncthreads(); + + const int k_base = lane * 4; + const float4 q4 = *reinterpret_cast(s_q + k_base); + const float4 k4 = *reinterpret_cast(s_k + k_base); + const float4 decay4 = *reinterpret_cast(s_decay + k_base); + float r_q[4] = {q4.x, q4.y, q4.z, q4.w}; + float r_k[4] = {k4.x, k4.y, k4.z, k4.w}; + float r_decay[4] = {decay4.x, decay4.y, decay4.z, decay4.w}; + float o_sumsq = 0.0f; + +#pragma unroll + for (int chunk = 0; chunk < kNumChunks; ++chunk) { + if constexpr (kPrefetchNextStateChunk && kNumChunks > 1) { + if (chunk + 1 < kNumChunks) { + cp_async_wait_group_1(); + } else { + cp_async_wait_all(); + } + } else { + cp_async_wait_all(); + } + if constexpr (!kSkipWarpSync) { + __syncwarp(); + } + + if constexpr (!kPrefetchNextStateChunk) { + if (chunk + 1 < kNumChunks) { + cp_async_state_chunk(&s_state[0][0][0], state_for_slot, 0, i_hv, + hv_count, chunk + 1); + } + } + +#pragma unroll + for (int row = 0; row < kRowsPerWarp; row += 2) { + const int v_row_a = warp + row * kWarps; + const int v_row_b = warp + (row + 1) * kWarps; + const int v0 = chunk * kChunkV + v_row_a; + const int v1 = chunk * kChunkV + v_row_b; + float h_a_vals[4]; + float h_b_vals[4]; + float dot_hk_a = 0.0f; + float dot_hk_b = 0.0f; + + const float4 raw_h_a = *reinterpret_cast( + &s_state[chunk & 1][v_row_a][k_base]); + const float4 raw_h_b = *reinterpret_cast( + &s_state[chunk & 1][v_row_b][k_base]); + h_a_vals[0] = raw_h_a.x * r_decay[0]; + h_a_vals[1] = raw_h_a.y * r_decay[1]; + h_a_vals[2] = raw_h_a.z * r_decay[2]; + h_a_vals[3] = raw_h_a.w * r_decay[3]; + h_b_vals[0] = raw_h_b.x * r_decay[0]; + h_b_vals[1] = raw_h_b.y * r_decay[1]; + h_b_vals[2] = raw_h_b.z * r_decay[2]; + h_b_vals[3] = raw_h_b.w * r_decay[3]; + dot_hk_a = h_a_vals[0] * r_k[0] + h_a_vals[1] * r_k[1] + + h_a_vals[2] * r_k[2] + h_a_vals[3] * r_k[3]; + dot_hk_b = h_b_vals[0] * r_k[0] + h_b_vals[1] * r_k[1] + + h_b_vals[2] * r_k[2] + h_b_vals[3] * r_k[3]; + + const Sum2 dot_hk = warp_reduce_sum_pair(dot_hk_a, dot_hk_b); + const float v_new0 = (s_v[v0] - dot_hk.x) * s_beta; + const float v_new1 = (s_v[v1] - dot_hk.y) * s_beta; + + float dot_hq_a = 0.0f; + float dot_hq_b = 0.0f; + const int state_idx_a = (i_hv * kDimV + v0) * kDimK + k_base; + const int state_idx_b = (i_hv * kDimV + v1) * kDimK + k_base; + const float h_a_0 = h_a_vals[0] + r_k[0] * v_new0; + const float h_a_1 = h_a_vals[1] + r_k[1] * v_new0; + const float h_a_2 = h_a_vals[2] + r_k[2] * v_new0; + const float h_a_3 = h_a_vals[3] + r_k[3] * v_new0; + const float h_b_0 = h_b_vals[0] + r_k[0] * v_new1; + const float h_b_1 = h_b_vals[1] + r_k[1] * v_new1; + const float h_b_2 = h_b_vals[2] + r_k[2] * v_new1; + const float h_b_3 = h_b_vals[3] + r_k[3] * v_new1; + if constexpr (kComputeOutputBeforeStore) { + dot_hq_a = + h_a_0 * r_q[0] + h_a_1 * r_q[1] + h_a_2 * r_q[2] + h_a_3 * r_q[3]; + dot_hq_b = + h_b_0 * r_q[0] + h_b_1 * r_q[1] + h_b_2 * r_q[2] + h_b_3 * r_q[3]; + store_state_float4( + state_for_slot + state_idx_a, + make_float4(h_a_0, h_a_1, h_a_2, h_a_3)); + store_state_float4( + state_for_slot + state_idx_b, + make_float4(h_b_0, h_b_1, h_b_2, h_b_3)); + } else { + store_state_float4( + state_for_slot + state_idx_a, + make_float4(h_a_0, h_a_1, h_a_2, h_a_3)); + store_state_float4( + state_for_slot + state_idx_b, + make_float4(h_b_0, h_b_1, h_b_2, h_b_3)); + dot_hq_a = + h_a_0 * r_q[0] + h_a_1 * r_q[1] + h_a_2 * r_q[2] + h_a_3 * r_q[3]; + dot_hq_b = + h_b_0 * r_q[0] + h_b_1 * r_q[1] + h_b_2 * r_q[2] + h_b_3 * r_q[3]; + } + + const Sum2 dot_hq = warp_reduce_sum_pair(dot_hq_a, dot_hq_b); + if (lane == 0) { + s_o[v0] = dot_hq.x; + s_o[v1] = dot_hq.y; + if constexpr (kApplyOnorm && kAccumulateOnormSumsq) { + o_sumsq += dot_hq.x * dot_hq.x + dot_hq.y * dot_hq.y; + } + } + } + + if constexpr (kPrefetchNextStateChunk) { + if (chunk + 2 < kNumChunks) { + cp_async_state_chunk(&s_state[0][0][0], state_for_slot, 0, i_hv, + hv_count, chunk + 2); + } + } + } + __syncthreads(); + +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900 + cudaTriggerProgrammaticLaunchCompletion(); +#endif + + if constexpr (kApplyOnorm) { + if constexpr (kAccumulateOnormSumsq) { + if (lane == 0) { + s_reduce[warp] = o_sumsq; + } + __syncthreads(); + + float total_sumsq = 0.0f; + if (warp == 0) { + total_sumsq = lane < kWarps ? s_reduce[lane] : 0.0f; + total_sumsq = warp_reduce_sum(total_sumsq); + if (lane == 0) { + s_reduce[0] = total_sumsq; + } + } + __syncthreads(); + + if (tid < kDimV) { + const int64_t out_idx = i_n * kLocalDim + i_hv * kDimV + tid; + const float raw_o = s_o[tid]; + const float rstd = + rsqrtf(s_reduce[0] / static_cast(kDimV) + onorm_eps); + float gate; + float weight; + if constexpr (kPreloadOnormParams) { + gate = pre_onorm_gate; + weight = pre_onorm_weight; + } else { + const int64_t gate_idx = i_n * strides.onorm_row + i_hv * kDimV + tid; + gate = sigmoid_fast(bf16_load(onorm_g, gate_idx)); + weight = onorm_weight[tid]; + } + const float y = raw_o * rstd * weight * gate; + out[out_idx] = bf16_store(y); + } + } else { + const float raw_o = tid < kDimV ? s_o[tid] : 0.0f; + const float o_sq = raw_o * raw_o; + float sumsq; + if constexpr (kUseActiveOnormReduction || kUseActiveQkReduction) { + sumsq = block_reduce_sum_active_for(o_sq, s_reduce); + } else { + sumsq = block_reduce_sum(o_sq, s_reduce); + } + + if (tid < kDimV) { + const int64_t out_idx = i_n * kLocalDim + i_hv * kDimV + tid; + const float rstd = + rsqrtf(sumsq / static_cast(kDimV) + onorm_eps); + float gate; + float weight; + if constexpr (kPreloadOnormParams) { + gate = pre_onorm_gate; + weight = pre_onorm_weight; + } else { + const int64_t gate_idx = i_n * strides.onorm_row + i_hv * kDimV + tid; + gate = sigmoid_fast(bf16_load(onorm_g, gate_idx)); + weight = onorm_weight[tid]; + } + const float y = raw_o * rstd * weight * gate; + out[out_idx] = bf16_store(y); + } + } + } else { + if (tid < kDimV) { + const int64_t out_idx = i_n * kLocalDim + i_hv * kDimV + tid; + out[out_idx] = bf16_store(s_o[tid]); + } + } +} + +template +void launch_kda_decode_many_heads_raw( + const void* x_q, const void* x_k, const void* x_v, const void* w_q_t, + const void* w_k_t, const void* w_v_t, const void* bias_q, + const void* bias_k, const void* bias_v, void* cs_q, void* cs_k, void* cs_v, + const float* a_log, const void* g, const float* dt_bias, const void* beta, + const void* onorm_g, const float* onorm_weight, + const int* ssm_state_indices, const int* cu_seqlens, float* state, + void* out, int B, int H, int HV, float lower_bound, float scale, + float onorm_eps, KdaDecodeStrides strides, cudaStream_t stream) { + auto kernel = &kda_decode_fusion_many_heads_kernel< + kApplyOnorm, true, kHeads, kHeads, true, false, false, false, false, + false, true, true, true, kUpdateConvState, kUseLowerBound, + kApplyBetaSigmoid>; + cudaLaunchConfig_t config{}; + config.gridDim = dim3(B, kHeads); + config.blockDim = dim3(kThreads); + config.dynamicSmemBytes = 0; + config.stream = stream; + cudaLaunchAttribute attrs[1]; + attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attrs[0].val.programmaticStreamSerializationAllowed = 1; + config.attrs = attrs; + config.numAttrs = 1; + cudaLaunchKernelEx(&config, kernel, + reinterpret_cast(x_q), + reinterpret_cast(x_k), + reinterpret_cast(x_v), + reinterpret_cast(w_q_t), + reinterpret_cast(w_k_t), + reinterpret_cast(w_v_t), + reinterpret_cast(bias_q), + reinterpret_cast(bias_k), + reinterpret_cast(bias_v), + reinterpret_cast<__nv_bfloat16*>(cs_q), + reinterpret_cast<__nv_bfloat16*>(cs_k), + reinterpret_cast<__nv_bfloat16*>(cs_v), a_log, + reinterpret_cast(g), dt_bias, + reinterpret_cast(beta), + reinterpret_cast(onorm_g), + onorm_weight, ssm_state_indices, cu_seqlens, state, + reinterpret_cast<__nv_bfloat16*>(out), B, H, HV, + lower_bound, scale, onorm_eps, strides); +} + +template +void launch_kda_decode_many_heads_selected( + const void* x_q, const void* x_k, const void* x_v, const void* w_q_t, + const void* w_k_t, const void* w_v_t, const void* bias_q, + const void* bias_k, const void* bias_v, void* cs_q, void* cs_k, void* cs_v, + const float* a_log, const void* g, const float* dt_bias, const void* beta, + const void* onorm_g, const float* onorm_weight, + const int* ssm_state_indices, const int* cu_seqlens, float* state, + void* out, int B, int H, int HV, bool update_conv_cache, float lower_bound, + float scale, float onorm_eps, KdaDecodeStrides strides, + cudaStream_t stream) { +#define LAUNCH_KDA_DECODE(NUM_HEADS) \ + do { \ + if (update_conv_cache) { \ + launch_kda_decode_many_heads_raw( \ + x_q, x_k, x_v, w_q_t, w_k_t, w_v_t, bias_q, bias_k, bias_v, cs_q, \ + cs_k, cs_v, a_log, g, dt_bias, beta, onorm_g, onorm_weight, \ + ssm_state_indices, cu_seqlens, state, out, B, H, HV, lower_bound, \ + scale, onorm_eps, strides, stream); \ + } else { \ + launch_kda_decode_many_heads_raw( \ + x_q, x_k, x_v, w_q_t, w_k_t, w_v_t, bias_q, bias_k, bias_v, cs_q, \ + cs_k, cs_v, a_log, g, dt_bias, beta, onorm_g, onorm_weight, \ + ssm_state_indices, cu_seqlens, state, out, B, H, HV, lower_bound, \ + scale, onorm_eps, strides, stream); \ + } \ + } while (false) + + switch (H) { + case 12: + LAUNCH_KDA_DECODE(12); + break; + case 24: + LAUNCH_KDA_DECODE(24); + break; + case 48: + LAUNCH_KDA_DECODE(48); + break; + case 96: + LAUNCH_KDA_DECODE(96); + break; + default: + STD_TORCH_CHECK(false, "Unsupported number of heads: ", H); + } + +#undef LAUNCH_KDA_DECODE +} + +struct KdaDecodeLaunchParams { + const void* x_q; + const void* x_k; + const void* x_v; + const void* w_q_t; + const void* w_k_t; + const void* w_v_t; + const void* bias_q; + const void* bias_k; + const void* bias_v; + void* cs_q; + void* cs_k; + void* cs_v; + const float* a_log; + const void* g; + const float* dt_bias; + const void* beta; + const void* onorm_g; + const float* onorm_weight; + const int* ssm_state_indices; + const int* cu_seqlens; + float* state; + void* out; + int B; + int H; + int HV; + bool update_conv_cache; + float lower_bound; + float scale; + float onorm_eps; + KdaDecodeStrides strides; + cudaStream_t stream; +}; + +template +void launch_kda_decode_selected_backend(const KdaDecodeLaunchParams& p) { + launch_kda_decode_many_heads_selected( + p.x_q, p.x_k, p.x_v, p.w_q_t, p.w_k_t, p.w_v_t, p.bias_q, p.bias_k, + p.bias_v, p.cs_q, p.cs_k, p.cs_v, p.a_log, p.g, p.dt_bias, p.beta, + p.onorm_g, p.onorm_weight, p.ssm_state_indices, p.cu_seqlens, p.state, + p.out, p.B, p.H, p.HV, p.update_conv_cache, p.lower_bound, p.scale, + p.onorm_eps, p.strides, p.stream); +} + +template +void dispatch_kda_decode_beta(const KdaDecodeLaunchParams& p, + bool apply_beta_sigmoid) { + if (apply_beta_sigmoid) { + launch_kda_decode_selected_backend(p); + } else { + launch_kda_decode_selected_backend(p); + } +} + +template +void dispatch_kda_decode_decay(const KdaDecodeLaunchParams& p, + bool use_lower_bound, bool apply_beta_sigmoid) { + if (use_lower_bound) { + dispatch_kda_decode_beta(p, apply_beta_sigmoid); + } else { + dispatch_kda_decode_beta(p, apply_beta_sigmoid); + } +} + +void dispatch_kda_decode_features(const KdaDecodeLaunchParams& p, + bool apply_onorm, bool use_lower_bound, + bool apply_beta_sigmoid) { + if (apply_onorm) { + dispatch_kda_decode_decay(p, use_lower_bound, apply_beta_sigmoid); + } else { + dispatch_kda_decode_decay(p, use_lower_bound, apply_beta_sigmoid); + } +} + +} // namespace + +extern "C" void launch_kda_decode_many_heads_cuda( + const void* x_q, const void* x_k, const void* x_v, const void* w_q_t, + const void* w_k_t, const void* w_v_t, const void* bias_q, + const void* bias_k, const void* bias_v, void* cs_q, void* cs_k, void* cs_v, + const float* a_log, const void* g, const float* dt_bias, const void* beta, + const void* onorm_g, const float* onorm_weight, + const int* ssm_state_indices, const int* cu_seqlens, float* state, + void* out, int B, int H, int HV, bool apply_onorm, bool update_conv_cache, + bool use_lower_bound, bool apply_beta_sigmoid, float lower_bound, + float scale, float onorm_eps, const int64_t* raw_strides, + cudaStream_t stream) { + const KdaDecodeStrides strides{raw_strides[0], raw_strides[1], raw_strides[2], + raw_strides[3], raw_strides[4]}; + const KdaDecodeLaunchParams params{x_q, + x_k, + x_v, + w_q_t, + w_k_t, + w_v_t, + bias_q, + bias_k, + bias_v, + cs_q, + cs_k, + cs_v, + a_log, + g, + dt_bias, + beta, + onorm_g, + onorm_weight, + ssm_state_indices, + cu_seqlens, + state, + out, + B, + H, + HV, + update_conv_cache, + lower_bound, + scale, + onorm_eps, + strides, + stream}; + dispatch_kda_decode_features(params, apply_onorm, use_lower_bound, + apply_beta_sigmoid); +} + +void fused_kda_decode( + torch::stable::Tensor const& x, torch::stable::Tensor const& weight, + std::optional bias, + torch::stable::Tensor& conv_state, torch::stable::Tensor const& raw_g, + torch::stable::Tensor const& raw_beta, torch::stable::Tensor const& a_log, + torch::stable::Tensor const& dt_bias, + torch::stable::Tensor const& state_indices, torch::stable::Tensor& state, + torch::stable::Tensor& out, std::optional lower_bound, + std::optional output_gate, + std::optional norm_weight, double norm_eps) { + using torch::headeronly::ScalarType; + constexpr int kHeadDim = 128; + constexpr int kConvWidth = 4; + + STD_TORCH_CHECK(x.is_cuda() && x.scalar_type() == ScalarType::BFloat16, + "x must be a CUDA bfloat16 tensor"); + STD_TORCH_CHECK(weight.is_cuda() && weight.scalar_type() == ScalarType::Float, + "weight must be a CUDA float32 tensor"); + STD_TORCH_CHECK( + conv_state.is_cuda() && conv_state.scalar_type() == ScalarType::BFloat16, + "conv_state must be a CUDA bfloat16 tensor"); + STD_TORCH_CHECK( + raw_g.is_cuda() && raw_g.scalar_type() == ScalarType::BFloat16, + "raw_g must be a CUDA bfloat16 tensor"); + STD_TORCH_CHECK( + raw_beta.is_cuda() && raw_beta.scalar_type() == ScalarType::BFloat16, + "raw_beta must be a CUDA bfloat16 tensor"); + STD_TORCH_CHECK(a_log.is_cuda() && a_log.scalar_type() == ScalarType::Float, + "A_log must be a CUDA float32 tensor"); + STD_TORCH_CHECK( + dt_bias.is_cuda() && dt_bias.scalar_type() == ScalarType::Float, + "dt_bias must be a CUDA float32 tensor"); + STD_TORCH_CHECK(state.is_cuda() && state.scalar_type() == ScalarType::Float, + "state must be a CUDA float32 tensor"); + STD_TORCH_CHECK(out.is_cuda() && out.scalar_type() == ScalarType::BFloat16, + "out must be a CUDA bfloat16 tensor"); + STD_TORCH_CHECK( + state_indices.is_cuda() && state_indices.scalar_type() == ScalarType::Int, + "state_indices must be a CUDA int32 tensor"); + + STD_TORCH_CHECK(x.dim() == 2, "x must have shape [B, 3 * H * 128]"); + int const batch_size = static_cast(x.size(0)); + int64_t const qkv_width = x.size(1); + STD_TORCH_CHECK(qkv_width % (3 * kHeadDim) == 0, + "x must have shape [B, 3 * H * 128]"); + int64_t const num_heads = qkv_width / (3 * kHeadDim); + STD_TORCH_CHECK( + num_heads == 12 || num_heads == 24 || num_heads == 48 || num_heads == 96, + "H must be 12, 24, 48, or 96, got ", num_heads); + STD_TORCH_CHECK(batch_size > 0, + "KDA decode fusion requires at least one row"); + int const dim = num_heads * kHeadDim; + + STD_TORCH_CHECK(weight.dim() == 3 && weight.is_contiguous() && + weight.size(0) == 3 && weight.size(1) == kConvWidth && + weight.size(2) == dim, + "weight must have shape [3, 4, H * 128]"); + STD_TORCH_CHECK(conv_state.dim() == 3 && conv_state.size(1) == 3 * dim && + conv_state.size(2) == kConvWidth - 1, + "conv_state must have shape [slots, 3 * H * 128, 3]"); + STD_TORCH_CHECK(raw_g.dim() == 4 && raw_g.size(0) == 1 && + raw_g.size(1) == batch_size && + raw_g.size(2) == num_heads && raw_g.size(3) == kHeadDim, + "raw_g must have shape [1, B, H, 128]"); + STD_TORCH_CHECK(raw_beta.dim() == 3 && raw_beta.size(0) == 1 && + raw_beta.size(1) == batch_size && + raw_beta.size(2) == num_heads, + "raw_beta must have shape [1, B, H]"); + STD_TORCH_CHECK(a_log.is_contiguous() && a_log.numel() == num_heads, + "A_log must be contiguous with H elements"); + STD_TORCH_CHECK(dt_bias.is_contiguous() && dt_bias.numel() == dim, + "dt_bias must be contiguous with H * 128 elements"); + STD_TORCH_CHECK( + state_indices.is_contiguous() && state_indices.numel() == batch_size, + "state_indices must be contiguous with B elements"); + STD_TORCH_CHECK(state.dim() == 4 && state.size(1) == num_heads && + state.size(2) == kHeadDim && state.size(3) == kHeadDim, + "state must have shape [slots, H, 128, 128]"); + STD_TORCH_CHECK(out.dim() == 4 && out.size(0) == 1 && + out.size(1) == batch_size && out.size(2) == num_heads && + out.size(3) == kHeadDim, + "out must have shape [1, B, H, 128]"); + STD_TORCH_CHECK(x.stride(1) == 1, + "x must be contiguous in its channel dimension"); + STD_TORCH_CHECK(conv_state.stride(0) >= 3 * dim * (kConvWidth - 1) && + conv_state.stride(1) == 1 && + conv_state.stride(2) == 3 * dim, + "conv_state must use the SD cache layout"); + STD_TORCH_CHECK(state.stride(0) >= num_heads * kHeadDim * kHeadDim && + state.stride(1) == kHeadDim * kHeadDim && + state.stride(2) == kHeadDim && state.stride(3) == 1, + "state must have contiguous [H, 128, 128] slot contents"); + STD_TORCH_CHECK(raw_g.is_contiguous(), "raw_g must be contiguous"); + STD_TORCH_CHECK(raw_beta.stride(2) == 1, + "raw_beta must be contiguous in its head dimension"); + STD_TORCH_CHECK(out.is_contiguous(), "out must be contiguous"); + + bool const apply_onorm = output_gate.has_value(); + STD_TORCH_CHECK(apply_onorm == norm_weight.has_value(), + "output_gate and norm_weight must be provided together"); + void const* output_gate_ptr = nullptr; + float const* norm_weight_ptr = nullptr; + int64_t output_gate_row_stride = 0; + if (apply_onorm) { + STD_TORCH_CHECK(output_gate->is_cuda() && + output_gate->scalar_type() == ScalarType::BFloat16, + "output_gate must be a CUDA bfloat16 tensor"); + bool const gate_is_3d = + output_gate->dim() == 3 && output_gate->size(0) == batch_size && + output_gate->size(1) == num_heads && output_gate->size(2) == kHeadDim; + bool const gate_is_4d = + output_gate->dim() == 4 && output_gate->size(0) == 1 && + output_gate->size(1) == batch_size && + output_gate->size(2) == num_heads && output_gate->size(3) == kHeadDim; + STD_TORCH_CHECK(gate_is_3d || gate_is_4d, + "output_gate must have shape [B, H, 128] or " + "[1, B, H, 128]"); + int const row_dim = gate_is_3d ? 0 : 1; + STD_TORCH_CHECK(output_gate->stride(output_gate->dim() - 1) == 1, + "output_gate must be contiguous in its last dimension"); + STD_TORCH_CHECK(output_gate->stride(row_dim + 1) == kHeadDim, + "output_gate must have contiguous head rows"); + STD_TORCH_CHECK(norm_weight->is_cuda() && + norm_weight->scalar_type() == ScalarType::Float, + "norm_weight must be a CUDA float32 tensor"); + STD_TORCH_CHECK( + norm_weight->is_contiguous() && norm_weight->numel() == kHeadDim, + "norm_weight must be contiguous with 128 elements"); + STD_TORCH_CHECK(norm_eps >= 0.0, "norm_eps must be non-negative"); + output_gate_ptr = output_gate->data_ptr(); + norm_weight_ptr = static_cast(norm_weight->data_ptr()); + output_gate_row_stride = output_gate->stride(row_dim); + } + + void const* bias_ptr = nullptr; + if (bias.has_value()) { + STD_TORCH_CHECK(bias->is_cuda() && bias->scalar_type() == ScalarType::Float, + "bias must be a CUDA float32 tensor"); + STD_TORCH_CHECK(bias->is_contiguous() && bias->numel() == 3 * dim, + "bias must be contiguous with 3 * H * 128 elements"); + bias_ptr = bias->data_ptr(); + } + + auto const* x_ptr = static_cast(x.data_ptr()); + auto const* weight_ptr = static_cast(weight.data_ptr()); + auto* conv_ptr = static_cast(conv_state.data_ptr()); + auto const* bias_bytes = static_cast(bias_ptr); + int64_t const segment_bytes = dim * sizeof(__nv_bfloat16); + int64_t const weight_segment_bytes = + dim * kConvWidth * static_cast(sizeof(float)); + int64_t const conv_segment_bytes = + dim * conv_state.stride(1) * sizeof(__nv_bfloat16); + int64_t const bias_segment_bytes = dim * sizeof(float); + std::array const strides{ + x.stride(0), raw_beta.stride(1), output_gate_row_stride, + conv_state.stride(0), state.stride(0), + }; + bool const use_lower_bound = lower_bound.has_value(); + float const lower_bound_value = + use_lower_bound ? static_cast(*lower_bound) : 0.0f; + + torch::stable::accelerator::DeviceGuard const device_guard( + x.get_device_index()); + cudaStream_t const stream = get_current_cuda_stream(x.get_device_index()); + launch_kda_decode_many_heads_cuda( + x_ptr, x_ptr + segment_bytes, x_ptr + 2 * segment_bytes, weight_ptr, + weight_ptr + weight_segment_bytes, weight_ptr + 2 * weight_segment_bytes, + bias_bytes, + bias_bytes == nullptr ? nullptr : bias_bytes + bias_segment_bytes, + bias_bytes == nullptr ? nullptr : bias_bytes + 2 * bias_segment_bytes, + conv_ptr, conv_ptr + conv_segment_bytes, + conv_ptr + 2 * conv_segment_bytes, + static_cast(a_log.data_ptr()), raw_g.data_ptr(), + static_cast(dt_bias.data_ptr()), raw_beta.data_ptr(), + output_gate_ptr, norm_weight_ptr, + static_cast(state_indices.data_ptr()), nullptr, + static_cast(state.data_ptr()), out.data_ptr(), batch_size, + num_heads, num_heads, apply_onorm, true, use_lower_bound, true, + lower_bound_value, 0.08838834764831845f, static_cast(norm_eps), + strides.data(), stream); + cudaError_t const error = cudaGetLastError(); + STD_TORCH_CHECK( + error == cudaSuccess, + "Kimi K3 KDA decode kernel launch failed: ", cudaGetErrorString(error)); +} diff --git a/csrc/libtorch_stable/moe/grouped_topk_kernels.cu b/csrc/libtorch_stable/moe/grouped_topk_kernels.cu index dadfbc3a5c..5d394f5555 100644 --- a/csrc/libtorch_stable/moe/grouped_topk_kernels.cu +++ b/csrc/libtorch_stable/moe/grouped_topk_kernels.cu @@ -25,6 +25,7 @@ #include "libtorch_stable/torch_utils.h" #include +#include #include #include #include @@ -449,7 +450,8 @@ enum ScoringFunc { SCORING_SIGMOID = 1 // apply sigmoid }; -// Efficient sigmoid approximation from TensorRT-LLM +// Adapted from +// https://github.com/NVIDIA/TensorRT-LLM/blob/v1.3.0rc2/cpp/tensorrt_llm/kernels/noAuxTcKernels.cu __device__ inline float sigmoid_accurate(float x) { return 0.5f * tanhf(0.5f * x) + 0.5f; } @@ -891,6 +893,434 @@ __global__ void grouped_topk_fused_small_expert_count_kernel( #endif } +// Adapted from +// https://github.com/flashinfer-ai/flashinfer/blob/06400d062a2d51564bbe781f6f811d0b75ca593e/include/flashinfer/trtllm/fused_moe/RoutingKernelTopK.cuh +namespace single_group_topk { +namespace detail { + +static constexpr int BlockDim = 256; +static constexpr uint32_t FullWarpMask = 0xffffffffU; +static constexpr float InvalidScore = -INFINITY; + +// TopK-only tuning: use wider workers and keep these tiers on the block path. +template +static constexpr bool UseTunedBlockPath = + MaxNumTopExperts == 16 && (MaxNumExperts == 896 || MaxNumExperts == 1024); + +template +__device__ __forceinline__ void preprocess_score(T input, BiasT correction_bias, + float& unbiased_score, + float& selection_score) { + unbiased_score = 0.0F; + selection_score = InvalidScore; + float const input_float = cuda_cast(input); + float const bias = cuda_cast(correction_bias); + if (!is_finite(input_float) || !is_finite(bias)) { + return; + } + + float const unbiased = apply_scoring(input_float); + float const biased = unbiased + bias; + if constexpr (SF == SCORING_NONE) { + if (!is_finite(biased)) { + return; + } + } + unbiased_score = unbiased; + selection_score = biased == 0.0F ? 0.0F : biased; +} + +template +__device__ __forceinline__ void write_outputs( + cg::thread_block_tile const& warp, float lane_selection_score, + float lane_unbiased, int32_t lane_expert, int32_t lane, int32_t token, + int32_t topk, float* topk_values, IdxT* topk_indices, bool renormalize, + float routed_scaling_factor) { + bool const finite_selection = + lane < topk && lane_selection_score != InvalidScore; + lane_unbiased = finite_selection ? lane_unbiased : 0.0F; + unsigned const finite_mask = __ballot_sync(FullWarpMask, finite_selection); + float const sum = cg::reduce(warp, lane_unbiased, cg::plus{}); + + if (lane < topk) { + float output = 0.0F; + if (finite_mask == 0) { + if (renormalize) { + output = 1.0F / static_cast(topk); + } + } else if (finite_selection) { + float scale = routed_scaling_factor; + if (renormalize) { + scale /= sum + 1e-20F; + } + output = lane_unbiased * scale; + } + + int64_t const output_index = int64_t{token} * topk + lane; + topk_values[output_index] = output; + topk_indices[output_index] = static_cast(lane_expert); + } +} + +template +__global__ void __launch_bounds__(BlockDim) + single_group_topk_block_kernel(T const* scores, float* topk_values, + IdxT* topk_indices, BiasT const* bias, + int64_t num_experts, int64_t topk, + bool renormalize, + float routed_scaling_factor, + bool enable_pdl) { + static constexpr int NumChunks = (MaxNumExperts + WARP_SIZE - 1) / WARP_SIZE; + static constexpr int WorkerValuesPerLane = + UseTunedBlockPath ? 8 : 4; + static constexpr int ExpertsPerWorkerWarp = WorkerValuesPerLane * WARP_SIZE; + using LaneOwnedRange = + reduce_topk::HighExpertLaneOwnedTopKRange; + static constexpr int NumWorkerWarps = + (MaxNumExperts + ExpertsPerWorkerWarp - 1) / ExpertsPerWorkerWarp; + static constexpr int NumIntermediate = NumWorkerWarps * MaxNumTopExperts; + static constexpr int MergeValuesPerLane = + (NumIntermediate + WARP_SIZE - 1) / WARP_SIZE; + static constexpr bool LaneOwnedResourcesFit = + NumWorkerWarps <= BlockDim / WARP_SIZE && MergeValuesPerLane <= 64; + static constexpr bool UseHierarchicalLaneTopK = + LaneOwnedRange::kEnabled && LaneOwnedResourcesFit; + + static_assert(NumChunks <= 64); + static_assert(MaxNumTopExperts <= WARP_SIZE); + +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + if (enable_pdl) { + cudaGridDependencySynchronize(); + } +#endif + + __shared__ float __attribute((aligned(128))) biased_scores[MaxNumExperts]; + __shared__ float __attribute((aligned(128))) unbiased_scores[MaxNumExperts]; + + int32_t const token = static_cast(blockIdx.x); + int32_t const lane = static_cast(threadIdx.x) % WARP_SIZE; + int32_t const warp_id = static_cast(threadIdx.x) / WARP_SIZE; + int32_t const num_experts_i32 = static_cast(num_experts); + int32_t const topk_i32 = static_cast(topk); + T const* token_scores = scores + int64_t{token} * num_experts; + + for (int32_t expert = static_cast(threadIdx.x); + expert < num_experts_i32; expert += BlockDim) { + preprocess_score(token_scores[expert], bias[expert], + unbiased_scores[expert], + biased_scores[expert]); + } + __syncthreads(); + + auto warp = cg::tiled_partition(cg::this_thread_block()); + + if constexpr (UseHierarchicalLaneTopK) { + __shared__ float + __attribute((aligned(128))) intermediate_scores[NumIntermediate]; + __shared__ int32_t + __attribute((aligned(128))) intermediate_indices[NumIntermediate]; + + if (warp_id < NumWorkerWarps) { + float local_scores[WorkerValuesPerLane]; + int32_t local_indices[WorkerValuesPerLane]; +#pragma unroll + for (int index = 0; index < WorkerValuesPerLane; ++index) { + int32_t const expert = + warp_id * ExpertsPerWorkerWarp + index * WARP_SIZE + lane; + local_scores[index] = + expert < num_experts_i32 ? biased_scores[expert] : InvalidScore; + local_indices[index] = expert; + } + + float lane_score; + int32_t lane_expert; + reduce_topk::reduceTopKForLane( + warp, lane_score, lane_expert, local_scores, local_indices, + InvalidScore, lane); + if (lane < MaxNumTopExperts) { + int32_t const intermediate = warp_id * MaxNumTopExperts + lane; + bool const active = lane < topk_i32; + intermediate_scores[intermediate] = active ? lane_score : InvalidScore; + intermediate_indices[intermediate] = + active ? lane_expert : MaxNumExperts; + } + } + __syncthreads(); + + if (warp_id != 0) { + return; + } + + float merge_scores[MergeValuesPerLane]; + int32_t merge_indices[MergeValuesPerLane]; +#pragma unroll + for (int index = 0; index < MergeValuesPerLane; ++index) { + int32_t const intermediate = index * WARP_SIZE + lane; + bool const active = intermediate < NumIntermediate; + merge_scores[index] = + active ? intermediate_scores[intermediate] : InvalidScore; + merge_indices[index] = + active ? intermediate_indices[intermediate] : MaxNumExperts; + } + + float lane_score; + int32_t lane_expert; + reduce_topk::reduceTopKForLane( + warp, lane_score, lane_expert, merge_scores, merge_indices, + InvalidScore, lane); + float const lane_unbiased = + lane < topk_i32 && lane_expert >= 0 && lane_expert < num_experts_i32 + ? unbiased_scores[lane_expert] + : 0.0F; + write_outputs(warp, lane_score, lane_unbiased, lane_expert, lane, token, + topk_i32, topk_values, topk_indices, renormalize, + routed_scaling_factor); + } else { + if (warp_id != 0) { + return; + } + + float local_scores[NumChunks]; + int32_t local_indices[NumChunks]; +#pragma unroll + for (int index = 0; index < NumChunks; ++index) { + int32_t const expert = index * WARP_SIZE + lane; + local_scores[index] = + expert < num_experts_i32 ? biased_scores[expert] : InvalidScore; + local_indices[index] = expert; + } + + float top_scores[MaxNumTopExperts]; + int32_t top_experts[MaxNumTopExperts]; + reduce_topk::reduceTopK(warp, top_scores, top_experts, local_scores, + local_indices, InvalidScore, topk_i32); + float const lane_score = lane < topk_i32 ? top_scores[lane] : InvalidScore; + int32_t const lane_expert = lane < topk_i32 ? top_experts[lane] : -1; + float const lane_unbiased = + lane < topk_i32 && lane_expert >= 0 && lane_expert < num_experts_i32 + ? unbiased_scores[lane_expert] + : 0.0F; + write_outputs(warp, lane_score, lane_unbiased, lane_expert, lane, token, + topk_i32, topk_values, topk_indices, renormalize, + routed_scaling_factor); + } + +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + if (enable_pdl) { + cudaTriggerProgrammaticLaunchCompletion(); + } +#endif +} + +template +struct WarpTopKLaunchConfig { + static constexpr int DefaultBlockDim = + MaxNumExperts <= 1024 ? MaxNumExperts : 1024; + static constexpr int BlockDim = DefaultBlockDim > 256 ? 256 : DefaultBlockDim; + static constexpr int NumWarps = BlockDim / WARP_SIZE; + static constexpr int MaxBlockScale = + (DefaultBlockDim + BlockDim - 1) / BlockDim; + static constexpr int MaxBlocks = 1024 * MaxBlockScale; + + static_assert(BlockDim % WARP_SIZE == 0); + + static uint32_t grid_dim(int64_t num_tokens) { + int64_t const token_blocks = (num_tokens + NumWarps - 1) / NumWarps; + int64_t const selected = + token_blocks < MaxBlocks ? token_blocks : MaxBlocks; + return static_cast(selected > 0 ? selected : 1); + } +}; + +template +__global__ void __launch_bounds__(WarpTopKLaunchConfig::BlockDim) + single_group_topk_warp_kernel(T const* scores, float* topk_values, + IdxT* topk_indices, BiasT const* bias, + int64_t num_tokens, int64_t num_experts, + int64_t topk, bool renormalize, + float routed_scaling_factor, + bool enable_pdl) { + static constexpr int NumChunks = (MaxNumExperts + WARP_SIZE - 1) / WARP_SIZE; + static constexpr int WarpBlockDim = + WarpTopKLaunchConfig::BlockDim; + using LaneOwnedRange = + reduce_topk::HighExpertLaneOwnedTopKRange; + + static_assert(NumChunks <= 64); + static_assert(MaxNumTopExperts <= WARP_SIZE); + +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + if (enable_pdl) { + cudaGridDependencySynchronize(); + } +#endif + + int32_t const lane = static_cast(threadIdx.x) % WARP_SIZE; + int32_t const warp_id = static_cast(threadIdx.x) / WARP_SIZE; + int32_t const global_warp = + static_cast(blockIdx.x) * WarpBlockDim / WARP_SIZE + warp_id; + int32_t const global_warp_stride = + static_cast(gridDim.x) * WarpBlockDim / WARP_SIZE; + int32_t const num_experts_i32 = static_cast(num_experts); + int32_t const topk_i32 = static_cast(topk); + auto warp = cg::tiled_partition(cg::this_thread_block()); + + for (int32_t token = global_warp; token < num_tokens; + token += global_warp_stride) { + T const* token_scores = scores + int64_t{token} * num_experts; + float local_scores[NumChunks]; + int32_t local_indices[NumChunks]; +#pragma unroll + for (int index = 0; index < NumChunks; ++index) { + int32_t const expert = index * WARP_SIZE + lane; + float unbiased; + float selection; + if (expert < num_experts_i32) { + preprocess_score(token_scores[expert], bias[expert], + unbiased, selection); + } else { + selection = InvalidScore; + } + local_scores[index] = selection; + local_indices[index] = expert; + } + + float lane_score; + int32_t lane_expert; + if constexpr (LaneOwnedRange::kEnabled) { + reduce_topk::reduceTopKForLane( + warp, lane_score, lane_expert, local_scores, local_indices, + InvalidScore, lane); + } else { + float top_scores[MaxNumTopExperts]; + int32_t top_experts[MaxNumTopExperts]; + reduce_topk::reduceTopK(warp, top_scores, top_experts, local_scores, + local_indices, InvalidScore, topk_i32); + lane_score = lane < topk_i32 ? top_scores[lane] : InvalidScore; + lane_expert = lane < topk_i32 ? top_experts[lane] : -1; + } + + float lane_unbiased = 0.0F; + if (lane < topk_i32 && lane_expert >= 0 && lane_expert < num_experts_i32) { + lane_unbiased = lane_score - cuda_cast(bias[lane_expert]); + } + write_outputs(warp, lane_score, lane_unbiased, lane_expert, lane, token, + topk_i32, topk_values, topk_indices, renormalize, + routed_scaling_factor); + } + +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + if (enable_pdl) { + cudaTriggerProgrammaticLaunchCompletion(); + } +#endif +} + +template +struct Tier { + static constexpr int kExperts = Experts; + static constexpr int kTopK = TopK; +}; + +template +struct TierList {}; + +using SigmoidBiasTiers = + TierList, Tier<256, 8>, Tier<384, 8>, Tier<512, 8>, + Tier<512, 22>, Tier<768, 16>, Tier<896, 16>, Tier<1024, 16>>; + +using PrecomputedSoftmaxBiasTiers = + TierList, Tier<128, 8>, Tier<160, 8>, Tier<256, 8>, + Tier<256, 16>, Tier<512, 8>, Tier<512, 16>, Tier<512, 22>, + Tier<512, 32>, Tier<576, 8>, Tier<768, 16>, Tier<896, 16>, + Tier<1024, 16>>; + +template +void launch(T* scores, float* topk_values, IdxT* topk_indices, + BiasT const* bias, int64_t num_tokens, int64_t num_experts, + int64_t topk, bool renormalize, double routed_scaling_factor, + bool enable_pdl, cudaLaunchConfig_t& config) { + config.dynamicSmemBytes = 0; + bool const use_block_kernel = + UseTunedBlockPath || + MaxNumExperts > 1024 || num_experts >= 1024 || + (num_experts >= 256 && num_tokens <= 1024); + if (use_block_kernel) { + config.gridDim = static_cast(num_tokens); + config.blockDim = BlockDim; + cudaLaunchKernelEx( + &config, + &single_group_topk_block_kernel, + scores, topk_values, topk_indices, bias, num_experts, topk, renormalize, + static_cast(routed_scaling_factor), enable_pdl); + } else { + using WarpConfig = WarpTopKLaunchConfig; + config.gridDim = WarpConfig::grid_dim(num_tokens); + config.blockDim = WarpConfig::BlockDim; + cudaLaunchKernelEx( + &config, + &single_group_topk_warp_kernel, + scores, topk_values, topk_indices, bias, num_tokens, num_experts, topk, + renormalize, static_cast(routed_scaling_factor), enable_pdl); + } +} + +template +bool dispatch(TierList<>*, T*, float*, IdxT*, BiasT const*, int64_t, int64_t, + int64_t, bool, double, bool, cudaLaunchConfig_t&) { + return false; +} + +template +bool dispatch(TierList*, T* scores, float* topk_values, + IdxT* topk_indices, BiasT const* bias, int64_t num_tokens, + int64_t num_experts, int64_t topk, bool renormalize, + double routed_scaling_factor, bool enable_pdl, + cudaLaunchConfig_t& config) { + if (num_experts <= First::kExperts && topk <= First::kTopK) { + launch( + scores, topk_values, topk_indices, bias, num_tokens, num_experts, topk, + renormalize, routed_scaling_factor, enable_pdl, config); + return true; + } + return dispatch( + static_cast*>(nullptr), scores, topk_values, + topk_indices, bias, num_tokens, num_experts, topk, renormalize, + routed_scaling_factor, enable_pdl, config); +} + +} // namespace detail + +template +bool invoke(T* scores, float* topk_values, IdxT* topk_indices, + BiasT const* bias, int64_t num_tokens, int64_t num_experts, + int64_t topk, bool renormalize, double routed_scaling_factor, + bool enable_pdl, cudaLaunchConfig_t& config) { + static_assert(SF == SCORING_NONE || SF == SCORING_SIGMOID); + if constexpr (SF == SCORING_SIGMOID) { + return detail::dispatch( + static_cast(nullptr), scores, topk_values, + topk_indices, bias, num_tokens, num_experts, topk, renormalize, + routed_scaling_factor, enable_pdl, config); + } else { + return detail::dispatch( + static_cast(nullptr), scores, + topk_values, topk_indices, bias, num_tokens, num_experts, topk, + renormalize, routed_scaling_factor, enable_pdl, config); + } +} + +} // namespace single_group_topk + template void invokeNoAuxTc(T* scores, float* topk_values, IdxT* topk_indices, BiasT const* bias, int64_t const num_tokens, @@ -906,6 +1336,12 @@ void invokeNoAuxTc(T* scores, float* topk_values, IdxT* topk_indices, attrs[0].val.programmaticStreamSerializationAllowed = enable_pdl; config.numAttrs = 1; config.attrs = attrs; + if (n_group == 1 && topk_group == 1 && + single_group_topk::invoke( + scores, topk_values, topk_indices, bias, num_tokens, num_experts, + topk, renormalize, routed_scaling_factor, enable_pdl, config)) { + return; + } // Check if we can use the optimized // grouped_topk_fused_small_expert_count_kernel diff --git a/csrc/libtorch_stable/moe/moeTopKFuncs.cuh b/csrc/libtorch_stable/moe/moeTopKFuncs.cuh index 81a216bd6f..bacb63b3cf 100644 --- a/csrc/libtorch_stable/moe/moeTopKFuncs.cuh +++ b/csrc/libtorch_stable/moe/moeTopKFuncs.cuh @@ -1,6 +1,7 @@ /* * Adapted from * https://github.com/NVIDIA/TensorRT-LLM/blob/v1.3.0rc2/cpp/tensorrt_llm/kernels/moeTopKFuncs.cuh + * https://github.com/flashinfer-ai/flashinfer/blob/06400d062a2d51564bbe781f6f811d0b75ca593e/include/flashinfer/trtllm/fused_moe/RoutingKernelTopK.cuh * Copyright (c) 2026, The vLLM team. * SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION. All rights * reserved. SPDX-License-Identifier: Apache-2.0 @@ -23,6 +24,9 @@ #include #include +#include +#include + namespace aphrodite { namespace moe { namespace reduce_topk { @@ -38,11 +42,10 @@ struct TopKRedType { "Top K reduction only implemented for int, float, float16 and bfloat16"); using TypeCmp = std::conditional_t; - using IdxT = std::conditional_t; static constexpr int kMoveBits = (sizeof(T) == 4) ? 32 : 16; static constexpr int kMaxIdx = 65535; - TypeCmp compValIdx; + TypeCmp compVal; static __host__ __device__ inline TypeCmp makeCmpVal(T val, int32_t idx = 0) { auto valueBits = cub::Traits::TwiddleIn( @@ -69,69 +72,175 @@ struct TopKRedType { __host__ __device__ TopKRedType() = default; __host__ __device__ TopKRedType(T val, int32_t idx) - : compValIdx(makeCmpVal(val, idx)) {} + : compVal(makeCmpVal(val, idx)) {} - __host__ __device__ operator TypeCmp() const noexcept { return compValIdx; } + __host__ __device__ operator TypeCmp() const noexcept { return compVal; } __device__ inline TypeCmp reduce( cg::thread_block_tile const& warp) { - return cg::reduce(warp, compValIdx, cg::greater{}); +#ifdef __CUDA_ARCH__ + static constexpr bool kHAS_FAST_REDUX = (__CUDA_ARCH__ / 100) >= 10; +#else + static constexpr bool kHAS_FAST_REDUX = false; +#endif + if constexpr (!kHAS_FAST_REDUX) { + return cg::reduce(warp, compVal, cg::greater{}); + } else if constexpr (sizeof(TypeCmp) == 8) { + uint32_t hi = static_cast(compVal >> 32); + uint32_t lo = static_cast(compVal & 0xffffffffu); + uint32_t maxHi; + asm volatile("redux.sync.max.u32 %0, %1, 0xffffffff;\n" + : "=r"(maxHi) + : "r"(hi)); + uint32_t loContrib = hi == maxHi ? lo : 0u; + uint32_t maxLo; + asm volatile("redux.sync.max.u32 %0, %1, 0xffffffff;\n" + : "=r"(maxLo) + : "r"(loContrib)); + return (static_cast(maxHi) << 32) | static_cast(maxLo); + } else { + TypeCmp result; + asm volatile("redux.sync.max.u32 %0, %1, 0xffffffff;\n" + : "=r"(result) + : "r"(compVal)); + return result; + } } }; -//////////////////////////////////////////////////////////////////////////////////////////////////// - -template -struct TopKIdx { - // by default, empty +template +struct IsPowerOf2 { + static constexpr bool value = N > 0 && (N & (N - 1)) == 0; }; -template -struct TopKIdx { - static constexpr int K = K_; - int32_t val[K]; +template +struct NextPow2 { + private: + static constexpr unsigned u = static_cast(N - 1); + static constexpr unsigned s1 = u | (u >> 1); + static constexpr unsigned s2 = s1 | (s1 >> 2); + static constexpr unsigned s3 = s2 | (s2 >> 4); + static constexpr unsigned s4 = s3 | (s3 >> 8); + static constexpr unsigned s5 = s4 | (s4 >> 16); + + public: + static constexpr int value = N <= 1 ? 1 : static_cast(s5 + 1); }; -//////////////////////////////////////////////////////////////////////////////////////////////////// +template +__device__ __forceinline__ void topkCompareSwap(T* a) { + if constexpr (A < Size && B < Size) { + if (a[A] < a[B]) { + T tmp = a[A]; + a[A] = a[B]; + a[B] = tmp; + } + } else { + (void)a; + } +} -#define TOPK_SWAP(I, J) \ - { \ - auto pairMin = min(topK[I].compValIdx, topK[J].compValIdx); \ - auto pairMax = max(topK[I].compValIdx, topK[J].compValIdx); \ - topK[I].compValIdx = pairMax; \ - topK[J].compValIdx = pairMin; \ +template +__device__ __forceinline__ void topkMergePairs(T* a) { + if constexpr (I + Step < End) { + topkCompareSwap(a); + topkMergePairs(a); + } else { + (void)a; + } +} + +template +__device__ __forceinline__ void topkOEM(T* a) { + constexpr int M = R * 2; + if constexpr (M < N) { + topkOEM(a); + topkOEM(a); + topkMergePairs(a); + } else if constexpr (R < N) { + topkCompareSwap(a); + } else { + (void)a; + } +} + +template +__device__ __forceinline__ void topkSortBatcher(T* a) { + if constexpr (N > 1) { + constexpr int Half = N / 2; + topkSortBatcher(a); + topkSortBatcher(a); + topkOEM(a); + } else { + (void)a; } +} template -struct Sort; +struct Sort { + static_assert(N > 0 && N <= 64, "Sort only supports N in range [1, 64]"); + + static __device__ void run(RedType* topK) { + if constexpr (IsPowerOf2::value) { +#pragma unroll + for (int k = 2; k <= N; k *= 2) { +#pragma unroll + for (int j = k / 2; j > 0; j /= 2) { +#pragma unroll + for (int i = 0; i < N; ++i) { + int ixj = i ^ j; + if (ixj > i) { + if ((i & k) == 0) { + if (topK[i].compVal < topK[ixj].compVal) { + auto tmp = topK[i].compVal; + topK[i].compVal = topK[ixj].compVal; + topK[ixj].compVal = tmp; + } + } else { + if (topK[i].compVal > topK[ixj].compVal) { + auto tmp = topK[i].compVal; + topK[i].compVal = topK[ixj].compVal; + topK[ixj].compVal = tmp; + } + } + } + } + } + } + } else { + constexpr int P = NextPow2::value; + topkSortBatcher<0, P, N, RedType>(topK); + } + } +}; template struct Sort<1, RedType> { - static __device__ void run(RedType* topK) {} + static __device__ void run(RedType*) {} }; template struct Sort<2, RedType> { - static __device__ void run(RedType* topK) { TOPK_SWAP(0, 1); } + static __device__ void run(RedType* topK) { topkCompareSwap<0, 1, 2>(topK); } }; template struct Sort<3, RedType> { static __device__ void run(RedType* topK) { - TOPK_SWAP(0, 1); - TOPK_SWAP(1, 2); - TOPK_SWAP(0, 1); + topkCompareSwap<0, 1, 3>(topK); + topkCompareSwap<1, 2, 3>(topK); + topkCompareSwap<0, 1, 3>(topK); } }; template struct Sort<4, RedType> { static __device__ void run(RedType* topK) { - TOPK_SWAP(0, 2); - TOPK_SWAP(1, 3); - TOPK_SWAP(0, 1); - TOPK_SWAP(2, 3); - TOPK_SWAP(1, 2); + topkCompareSwap<0, 2, 4>(topK); + topkCompareSwap<1, 3, 4>(topK); + topkCompareSwap<0, 1, 4>(topK); + topkCompareSwap<2, 3, 4>(topK); + topkCompareSwap<1, 2, 4>(topK); } }; @@ -147,24 +256,23 @@ __forceinline__ __device__ void reduceTopK( typename RedType::TypeCmp packedMax{}; #pragma unroll for (int kk = 0; kk < actualK; ++kk) { - topK = - kk > 0 && packedMax == topK.compValIdx ? RedType{minValue, idx} : topK; - // get the next largest value + topK = kk > 0 && packedMax == topK.compVal ? RedType{minValue, idx} : topK; packedMax = topK.reduce(warp); RedType::unpack(out[kk], outIdx[kk], packedMax); } }; -template -__device__ void reduceTopKFunc(cg::thread_block_tile const& warp, - Type (&out)[K], int32_t (&outIdx)[K], - Type (&value)[N], int32_t (&idx)[N], - Type minValue, int actualK = K) { +template +__forceinline__ __device__ void reduceTopK( + cg::thread_block_tile const& warp, Type (&out)[K], + int32_t (&outIdx)[K], Type (&value)[N], int32_t (&idx)[N], + Type const minValue, int actualK = K) { static_assert(K > 0, "Top K must have K > 0"); - static_assert(K < kWARP_SIZE, "Top K must have K < kWARP_SIZE"); + static_assert(K <= kWARP_SIZE, "Top K must have K <= kWARP_SIZE"); static_assert(N > 0, "Top K must have N > 0"); - static_assert(N < 5, - "Only support candidates number less than or equal to 128"); + static_assert(N <= 64, + "Only support candidates number less than or equal to " + "64*32=2048"); using RedType = TopKRedType; RedType topK[N]; #pragma unroll @@ -172,85 +280,88 @@ __device__ void reduceTopKFunc(cg::thread_block_tile const& warp, topK[nn] = RedType{value[nn], idx[nn]}; } - if constexpr (!IsSorted) { - Sort::run(topK); - } + Sort::run(topK); + typename RedType::TypeCmp packedMax{}; -#pragma unroll for (int kk = 0; kk < actualK; ++kk) { - bool update = kk > 0 && packedMax == topK[0].compValIdx; + bool update = kk > 0 && packedMax == topK[0].compVal; #pragma unroll for (int nn = 0; nn < N; ++nn) { topK[nn] = update && nn == N - 1 ? RedType{minValue, idx[nn]} : update ? topK[nn + 1] : topK[nn]; } - // get the next largest value packedMax = topK[0].reduce(warp); RedType::unpack(out[kk], outIdx[kk], packedMax); } }; +template +struct LaneOwnedTopKRange { + static_assert(MinExperts > 0 && MinExperts <= MaxExperts); + static_assert(MinTopExperts > 0 && MinTopExperts <= MaxTopExperts); + static constexpr bool kEnabled = + NumExperts >= MinExperts && NumExperts <= MaxExperts && + NumTopExperts >= MinTopExperts && NumTopExperts <= MaxTopExperts; +}; + +static constexpr int kHIGH_EXPERT_LANE_OWNED_TOPK_MIN_EXPERTS = 512; +static constexpr int kHIGH_EXPERT_LANE_OWNED_TOPK_MAX_EXPERTS = 1024; +static constexpr int kHIGH_EXPERT_LANE_OWNED_TOPK_MIN_TOP_EXPERTS = 9; +static constexpr int kHIGH_EXPERT_LANE_OWNED_TOPK_MAX_TOP_EXPERTS = 16; + +template +using HighExpertLaneOwnedTopKRange = + LaneOwnedTopKRange; + template -__forceinline__ __device__ void reduceTopK( - cg::thread_block_tile const& warp, Type (&out)[K], - int32_t (&outIdx)[K], Type (&value)[N], int32_t (&idx)[N], - Type const minValue, int actualK = K) { +__forceinline__ __device__ void reduceTopKForLane( + cg::thread_block_tile const& warp, Type& out, int32_t& outIdx, + Type (&value)[N], int32_t (&idx)[N], Type const minValue, int32_t laneIdx) { static_assert(K > 0, "Top K must have K > 0"); - static_assert(K < kWARP_SIZE, "Top K must have K < kWARP_SIZE"); + static_assert(K <= kWARP_SIZE, "Top K must have K <= kWARP_SIZE"); static_assert(N > 0, "Top K must have N > 0"); - static_assert( - N <= 16, - "Only support candidates number less than or equal to 16*32=512"); - static_assert(N <= 4 || N % 4 == 0, - "Only support candidates number is a multiple of 4*32=128 or " - "less than or equal to 4"); + static_assert(N <= 64, + "Only support candidates number less than or equal to " + "64*32=2048"); using RedType = TopKRedType; + RedType topK[N]; +#pragma unroll + for (int nn = 0; nn < N; ++nn) { + topK[nn] = RedType{value[nn], idx[nn]}; + } - if constexpr (N <= 4) { - reduceTopKFunc(warp, out, outIdx, value, idx, minValue, - actualK); - } else { - constexpr int numLoops = N / 4; - constexpr int numResults = (numLoops * K - 1) / kWARP_SIZE + 1; - - Type topKBufferValue[numResults]; - int32_t topKBufferIdx[numResults]; - int32_t laneIdx = threadIdx.x % kWARP_SIZE; + Sort::run(topK); - for (int ii = 0; ii < numResults; ++ii) { - topKBufferValue[ii] = minValue; - topKBufferIdx[ii] = ii * kWARP_SIZE - 1; + typename RedType::TypeCmp packedMax{}; + typename RedType::TypeCmp lanePacked{}; +#pragma unroll + for (int kk = 0; kk < K; ++kk) { + bool update = kk > 0 && packedMax == topK[0].compVal; +#pragma unroll + for (int nn = 0; nn < N; ++nn) { + topK[nn] = update && nn == N - 1 ? RedType{minValue, idx[nn]} + : update ? topK[nn + 1] + : topK[nn]; } - for (int loop = 0; loop < numLoops; ++loop) { - int start = loop * 4; - Type topKValue[K]; - int32_t topKIdx[K]; - Type inValue[4]; - int32_t inIdx[4]; - for (int i = 0; i < 4; ++i) { - inValue[i] = value[start + i]; - inIdx[i] = idx[start + i]; - } - reduceTopKFunc(warp, topKValue, topKIdx, inValue, inIdx, - minValue, actualK); - int inOffset = laneIdx % K; - if (laneIdx >= loop * K && laneIdx < (loop + 1) * K) { - topKBufferValue[0] = topKValue[inOffset]; - topKBufferIdx[0] = topKIdx[inOffset]; - } - if (loop == numLoops - 1 && (laneIdx < (numLoops * K - kWARP_SIZE))) { - topKBufferValue[1] = topKValue[inOffset]; - topKBufferIdx[1] = topKIdx[inOffset]; - } + packedMax = topK[0].reduce(warp); + if (laneIdx == kk) { + lanePacked = packedMax; } - - reduceTopKFunc(warp, out, outIdx, topKBufferValue, - topKBufferIdx, minValue, actualK); } -}; -#undef TOPK_SWAP + if (laneIdx < K) { + RedType::unpack(out, outIdx, lanePacked); + } else { + out = minValue; + outIdx = -1; + } +} } // namespace reduce_topk } // namespace moe diff --git a/csrc/libtorch_stable/ops.h b/csrc/libtorch_stable/ops.h index c2bc9c4c66..29d58083c9 100644 --- a/csrc/libtorch_stable/ops.h +++ b/csrc/libtorch_stable/ops.h @@ -288,6 +288,61 @@ void fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_bf16_insert( torch::stable::Tensor const& cos_sin_cache, double eps, int64_t cache_block_size); +void fused_kimi_k3_mla_key_concat_kv_cache_insert( + torch::stable::Tensor& q, torch::stable::Tensor const& k_nope, + torch::stable::Tensor const& k_pe, torch::stable::Tensor const& kv_c_normed, + torch::stable::Tensor& k_out, torch::stable::Tensor& k_cache, + torch::stable::Tensor const& slot_mapping, int64_t cache_block_size, + std::optional position_ids, + std::optional cos_sin_cache); + +void fused_kimi_k3_mla_key_concat_ds_mla_insert( + torch::stable::Tensor& q, torch::stable::Tensor const& k_nope, + torch::stable::Tensor const& k_pe, torch::stable::Tensor const& kv_c_normed, + torch::stable::Tensor& k_out, torch::stable::Tensor& k_cache, + torch::stable::Tensor const& slot_mapping, int64_t cache_block_size, + std::optional position_ids, + std::optional cos_sin_cache); + +void fused_kimi_k3_mla_qkv_quant_kv_cache_fp8_insert( + torch::stable::Tensor const& q, torch::stable::Tensor const& k_nope, + torch::stable::Tensor const& k_pe, torch::stable::Tensor const& kv_c_normed, + torch::stable::Tensor const& v, torch::stable::Tensor& q_fp8, + torch::stable::Tensor& k_fp8, torch::stable::Tensor& v_fp8, + torch::stable::Tensor& k_cache, torch::stable::Tensor const& slot_mapping, + torch::stable::Tensor const& q_scale_inv, + torch::stable::Tensor const& k_scale_inv, + torch::stable::Tensor const& v_scale_inv, + torch::stable::Tensor const& cache_scale_inv, int64_t cache_block_size, + std::optional position_ids, + std::optional cos_sin_cache); + +void fused_kimi_k3_mla_decode_q_concat_kv_cache_insert( + torch::stable::Tensor const& ql_nope, torch::stable::Tensor const& q_pe, + torch::stable::Tensor const& kv_c_normed, torch::stable::Tensor const& k_pe, + torch::stable::Tensor& mqa_q, torch::stable::Tensor& k_cache, + torch::stable::Tensor const& slot_mapping, int64_t cache_block_size, + std::optional position_ids, + std::optional cos_sin_cache); + +void fused_kimi_k3_mla_decode_q_concat_kv_cache_fp8_insert( + torch::stable::Tensor const& ql_nope, torch::stable::Tensor const& q_pe, + torch::stable::Tensor const& kv_c_normed, torch::stable::Tensor const& k_pe, + torch::stable::Tensor& mqa_q, torch::stable::Tensor& k_cache, + torch::stable::Tensor const& slot_mapping, + torch::stable::Tensor const& q_scale_inv, + torch::stable::Tensor const& cache_scale_inv, int64_t cache_block_size, + std::optional position_ids, + std::optional cos_sin_cache); + +void fused_kimi_k3_mla_decode_q_concat_ds_mla_insert( + torch::stable::Tensor const& ql_nope, torch::stable::Tensor const& q_pe, + torch::stable::Tensor const& kv_c_normed, torch::stable::Tensor const& k_pe, + torch::stable::Tensor& mqa_q, torch::stable::Tensor& k_cache, + torch::stable::Tensor const& slot_mapping, int64_t cache_block_size, + std::optional position_ids, + std::optional cos_sin_cache); + void fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_fp8_insert( torch::stable::Tensor const& q, torch::stable::Tensor const& kv, torch::stable::Tensor& q_fp8, torch::stable::Tensor& k_cache, @@ -327,6 +382,19 @@ void fused_minimax_m3_qknorm_rope_kv_insert( std::optional index_q_out, const std::string& kv_cache_dtype, bool skip_index_branch); +#ifdef APHRODITE_ENABLE_FUSED_KDA_DECODE +void fused_kda_decode( + torch::stable::Tensor const& x, torch::stable::Tensor const& weight, + std::optional bias, + torch::stable::Tensor& conv_state, torch::stable::Tensor const& raw_g, + torch::stable::Tensor const& raw_beta, torch::stable::Tensor const& a_log, + torch::stable::Tensor const& dt_bias, + torch::stable::Tensor const& state_indices, torch::stable::Tensor& state, + torch::stable::Tensor& out, std::optional lower_bound, + std::optional output_gate, + std::optional norm_weight, double norm_eps); +#endif + #ifdef APHRODITE_ENABLE_KIMI_K3_ATTN_RES void kimi_k3_attn_res(torch::stable::Tensor& prefix, torch::stable::Tensor const& delta, @@ -427,6 +495,20 @@ fptr_t init_custom_ar(const std::vector& fake_ipc_ptrs, void all_reduce(fptr_t _fa, torch::stable::Tensor& inp, torch::stable::Tensor& out, fptr_t reg_buffer, int64_t reg_buffer_sz_bytes); +void custom_all_gather(fptr_t _fa, torch::stable::Tensor& inp, + torch::stable::Tensor& out, fptr_t reg_buffer, + int64_t reg_buffer_sz_bytes); +void mnnvl_lamport_all_gather(fptr_t _fa, torch::stable::Tensor& inp, + torch::stable::Tensor& out, fptr_t local_buffer, + fptr_t multicast_buffer, fptr_t epoch_buffer, + int64_t stage_sz_bytes); +void custom_reduce_scatter(fptr_t _fa, torch::stable::Tensor& inp, + torch::stable::Tensor& out, fptr_t reg_buffer, + int64_t reg_buffer_sz_bytes); +void mnnvl_lamport_reduce_scatter(fptr_t _fa, torch::stable::Tensor& inp, + torch::stable::Tensor& out, + fptr_t local_buffer, fptr_t epoch_buffer, + int64_t stage_sz_bytes); void dispose(fptr_t _fa); int64_t meta_size(); void register_buffer(fptr_t _fa, const std::vector& fake_ipc_ptrs); @@ -465,6 +547,12 @@ void fatrelu_and_mul(torch::stable::Tensor& out, torch::stable::Tensor& input, double threshold); void swigluoai_and_mul(torch::stable::Tensor& out, torch::stable::Tensor& input, double alpha = 1.702, double limit = 7.0); +void situ_and_mul(torch::stable::Tensor& out, torch::stable::Tensor& input, + double beta = 1.0, double linear_beta = -1.0); +void masked_situ_and_mul(torch::stable::Tensor& out, + torch::stable::Tensor& input, + const torch::stable::Tensor& expert_num_tokens, + double beta = 1.0, double linear_beta = -1.0); void gelu_new(torch::stable::Tensor& out, torch::stable::Tensor& input); void gelu_fast(torch::stable::Tensor& out, torch::stable::Tensor& input); void gelu_quick(torch::stable::Tensor& out, torch::stable::Tensor& input); @@ -540,6 +628,13 @@ void concat_and_cache_mla(torch::stable::Tensor& kv_c, const std::string& kv_cache_dtype, torch::stable::Tensor& scale); +void concat_and_cache_mla_grouped(torch::stable::Tensor& kv_c, + torch::stable::Tensor& k_pe, + torch::stable::Tensor& kv_cache_ptrs, + torch::stable::Tensor& slot_mapping, + int64_t block_size, int64_t block_stride, + int64_t entry_stride); + // NOTE: k_pe and kv_c order is flipped compared to concat_and_cache_mla void concat_and_cache_mla_rope_fused( torch::stable::Tensor& positions, torch::stable::Tensor& q_pe, diff --git a/csrc/libtorch_stable/torch_bindings.cpp b/csrc/libtorch_stable/torch_bindings.cpp index 0a5ef9c11d..aae1a1a94d 100644 --- a/csrc/libtorch_stable/torch_bindings.cpp +++ b/csrc/libtorch_stable/torch_bindings.cpp @@ -377,7 +377,8 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) { // DeepSeek V3 fused A GEMM (SM 9.0+, bf16 only, 1-16 tokens). // conditionally compiled so impl registration is in source file ops.def( - "dsv3_fused_a_gemm(Tensor! output, Tensor mat_a, Tensor mat_b) -> ()"); + "dsv3_fused_a_gemm(Tensor! output, Tensor mat_a, Tensor mat_b, " + "bool enable_pdl=False) -> ()"); // BF16/FP32 x FP32 -> FP32 router GEMM for H=3072, E=256, M<=32 (SM90+). // conditionally compiled so impl registration is in source file @@ -500,6 +501,48 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) { "Tensor fp8_scale, Tensor q_fp8_scale_inv, float eps, " "int cache_block_size) -> ()"); + // Kimi-K3 MLA epilogues: optional RoPE followed by concat/cache insertion. + ops.def( + "fused_kimi_k3_mla_key_concat_kv_cache_insert(" + "Tensor! q, Tensor k_nope, Tensor k_pe, Tensor kv_c_normed, " + "Tensor! k_out, Tensor! k_cache, Tensor slot_mapping, " + "int cache_block_size, Tensor? position_ids=None, " + "Tensor? cos_sin_cache=None) -> ()"); + ops.def( + "fused_kimi_k3_mla_key_concat_ds_mla_insert(" + "Tensor! q, Tensor k_nope, Tensor k_pe, Tensor kv_c_normed, " + "Tensor! k_out, Tensor! k_cache, Tensor slot_mapping, " + "int cache_block_size, Tensor? position_ids=None, " + "Tensor? cos_sin_cache=None) -> ()"); + ops.def( + "fused_kimi_k3_mla_qkv_quant_kv_cache_fp8_insert(" + "Tensor q, Tensor k_nope, Tensor k_pe, Tensor kv_c_normed, Tensor v, " + "Tensor! q_fp8, Tensor! k_fp8, Tensor! v_fp8, Tensor! k_cache, " + "Tensor slot_mapping, Tensor q_scale_inv, Tensor k_scale_inv, " + "Tensor v_scale_inv, Tensor cache_scale_inv, int cache_block_size, " + "Tensor? position_ids=None, Tensor? cos_sin_cache=None) -> ()"); + + // Kimi-K3 MLA decode epilogue: concat mqa_q = [ql_nope | q_pe] and insert the + // latent [kv_c_normed | k_pe] into the paged cache (bf16 / fp8 / fp8_ds_mla). + ops.def( + "fused_kimi_k3_mla_decode_q_concat_kv_cache_insert(" + "Tensor ql_nope, Tensor q_pe, Tensor kv_c_normed, Tensor k_pe, " + "Tensor! mqa_q, Tensor! k_cache, Tensor slot_mapping, " + "int cache_block_size, Tensor? position_ids=None, " + "Tensor? cos_sin_cache=None) -> ()"); + ops.def( + "fused_kimi_k3_mla_decode_q_concat_kv_cache_fp8_insert(" + "Tensor ql_nope, Tensor q_pe, Tensor kv_c_normed, Tensor k_pe, " + "Tensor! mqa_q, Tensor! k_cache, Tensor slot_mapping, " + "Tensor q_scale_inv, Tensor cache_scale_inv, int cache_block_size, " + "Tensor? position_ids=None, Tensor? cos_sin_cache=None) -> ()"); + ops.def( + "fused_kimi_k3_mla_decode_q_concat_ds_mla_insert(" + "Tensor ql_nope, Tensor q_pe, Tensor kv_c_normed, Tensor k_pe, " + "Tensor! mqa_q, Tensor! k_cache, Tensor slot_mapping, " + "int cache_block_size, Tensor? position_ids=None, " + "Tensor? cos_sin_cache=None) -> ()"); + #ifndef USE_ROCM ops.def( "minimax_allreduce_rms_qk(" @@ -521,6 +564,16 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) { "int block_size, Tensor!? q_out, Tensor!? index_q_out, " "str kv_cache_dtype, bool skip_index_branch=False) -> ()"); +#ifdef APHRODITE_ENABLE_FUSED_KDA_DECODE + ops.def( + "fused_kda_decode(" + "Tensor x, Tensor weight, Tensor? bias, Tensor! conv_state, " + "Tensor raw_g, Tensor raw_beta, Tensor A_log, Tensor dt_bias, " + "Tensor state_indices, Tensor! state, Tensor! out, " + "float? lower_bound=None, Tensor? output_gate=None, " + "Tensor? norm_weight=None, float norm_eps=1e-5) -> ()"); +#endif + #ifdef APHRODITE_ENABLE_KIMI_K3_ATTN_RES ops.def( "kimi_k3_attn_res(" @@ -613,6 +666,14 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) { "limit=7.0) " "-> ()"); + // SituGLU implementation used in Kimi models. + ops.def( + "situ_and_mul(Tensor! out, Tensor input, float beta=1.0, float " + "linear_beta=-1.0) -> ()"); + ops.def( + "masked_situ_and_mul(Tensor! out, Tensor input, Tensor " + "expert_num_tokens, float beta=1.0, float linear_beta=-1.0) -> ()"); + // GELU implementation used in GPT-2. ops.def("gelu_new(Tensor! out, Tensor input) -> ()"); @@ -771,11 +832,27 @@ STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) { ops.impl( "fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_fp8_insert", TORCH_BOX(&fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_fp8_insert)); + ops.impl("fused_kimi_k3_mla_key_concat_kv_cache_insert", + TORCH_BOX(&fused_kimi_k3_mla_key_concat_kv_cache_insert)); + ops.impl("fused_kimi_k3_mla_key_concat_ds_mla_insert", + TORCH_BOX(&fused_kimi_k3_mla_key_concat_ds_mla_insert)); + ops.impl("fused_kimi_k3_mla_qkv_quant_kv_cache_fp8_insert", + TORCH_BOX(&fused_kimi_k3_mla_qkv_quant_kv_cache_fp8_insert)); + ops.impl("fused_kimi_k3_mla_decode_q_concat_kv_cache_insert", + TORCH_BOX(&fused_kimi_k3_mla_decode_q_concat_kv_cache_insert)); + ops.impl("fused_kimi_k3_mla_decode_q_concat_kv_cache_fp8_insert", + TORCH_BOX(&fused_kimi_k3_mla_decode_q_concat_kv_cache_fp8_insert)); + ops.impl("fused_kimi_k3_mla_decode_q_concat_ds_mla_insert", + TORCH_BOX(&fused_kimi_k3_mla_decode_q_concat_ds_mla_insert)); #ifndef USE_ROCM ops.impl("minimax_allreduce_rms_qk", TORCH_BOX(&minimax_allreduce_rms_qk)); #endif ops.impl("fused_minimax_m3_qknorm_rope_kv_insert", TORCH_BOX(&fused_minimax_m3_qknorm_rope_kv_insert)); +#ifdef APHRODITE_ENABLE_FUSED_KDA_DECODE + ops.impl("fused_kda_decode", TORCH_BOX(&fused_kda_decode)); +#endif + #ifdef APHRODITE_ENABLE_KIMI_K3_ATTN_RES ops.impl("kimi_k3_attn_res", TORCH_BOX(&kimi_k3_attn_res)); #endif @@ -801,6 +878,8 @@ STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) { ops.impl("gelu_tanh_and_mul", TORCH_BOX(&gelu_tanh_and_mul)); ops.impl("fatrelu_and_mul", TORCH_BOX(&fatrelu_and_mul)); ops.impl("swigluoai_and_mul", TORCH_BOX(&swigluoai_and_mul)); + ops.impl("situ_and_mul", TORCH_BOX(&situ_and_mul)); + ops.impl("masked_situ_and_mul", TORCH_BOX(&masked_situ_and_mul)); ops.impl("gelu_new", TORCH_BOX(&gelu_new)); ops.impl("gelu_fast", TORCH_BOX(&gelu_fast)); ops.impl("gelu_quick", TORCH_BOX(&gelu_quick)); @@ -899,6 +978,15 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C_cache_ops, ops) { " str kv_cache_dtype," " Tensor scale) -> ()"); + // Grouped concat_and_cache_mla across all layers (bf16 only). Each + // layer's cache base pointer is read from kv_cache_ptrs. + ops.def( + "concat_and_cache_mla_grouped(Tensor kv_c, Tensor k_pe," + " Tensor kv_cache_ptrs," + " Tensor slot_mapping," + " int block_size, int block_stride," + " int entry_stride) -> ()"); + // Rotate Q and K, then write to kv cache for MLA ops.def( "concat_and_cache_mla_rope_fused(" @@ -997,6 +1085,8 @@ STABLE_TORCH_LIBRARY_IMPL(_C_cache_ops, CUDA, ops) { ops.impl("reshape_and_cache", TORCH_BOX(&reshape_and_cache)); ops.impl("reshape_and_cache_flash", TORCH_BOX(&reshape_and_cache_flash)); ops.impl("concat_and_cache_mla", TORCH_BOX(&concat_and_cache_mla)); + ops.impl("concat_and_cache_mla_grouped", + TORCH_BOX(&concat_and_cache_mla_grouped)); ops.impl("concat_and_cache_mla_rope_fused", TORCH_BOX(&concat_and_cache_mla_rope_fused)); ops.impl("convert_fp8", TORCH_BOX(&convert_fp8)); diff --git a/csrc/quantization/w8a8/fp8/nvidia/quant_utils.cuh b/csrc/quantization/w8a8/fp8/nvidia/quant_utils.cuh index 590622eb12..8dde19da25 100644 --- a/csrc/quantization/w8a8/fp8/nvidia/quant_utils.cuh +++ b/csrc/quantization/w8a8/fp8/nvidia/quant_utils.cuh @@ -13,10 +13,18 @@ namespace aphrodite { namespace fp8 { #ifdef ENABLE_FP8 +// Unspecialized conversions are a compile error: the old passthrough +// (`return x;`) silently skipped fp8 encoding for any (Tout, Tin) pair +// without a specialization below (e.g. the torch stable-ABI scalar types), +// corrupting quantized data with no runtime signal. +template +inline constexpr bool _no_conversion_specialization = false; + template __inline__ __device__ Tout vec_conversion( const Tin& x, const __nv_fp8_interpretation_t fp8_type = __NV_E4M3) { - return x; + static_assert(_no_conversion_specialization, + "no vec_conversion specialization for this (Tout, Tin) pair"); } // float -> c10::Float8_e4m3fn @@ -301,7 +309,9 @@ __inline__ __device__ bf16_8_t vec_conversion( template __inline__ __device__ Tout scaled_vec_conversion( const Tin& x, const float scale, const __nv_fp8_interpretation_t fp8_type) { - return x; + static_assert( + _no_conversion_specialization, + "no scaled_vec_conversion specialization for this (Tout, Tin) pair"); } // fp8 -> half @@ -492,6 +502,25 @@ __inline__ __device__ uint8_t scaled_vec_conversion( __builtin_unreachable(); // Suppress missing return statement warning } +// torch stable-ABI (headeronly) scalar types delegate to the CUDA-native +// conversions, so libtorch_stable kernels dispatched on c10::BFloat16 / +// c10::Half quantize correctly without manual casts. +template <> +__inline__ __device__ uint8_t scaled_vec_conversion( + const c10::BFloat16& a, const float scale, + const __nv_fp8_interpretation_t fp8_type) { + return scaled_vec_conversion( + reinterpret_cast(a), scale, fp8_type); +} + +template <> +__inline__ __device__ uint8_t scaled_vec_conversion( + const c10::Half& a, const float scale, + const __nv_fp8_interpretation_t fp8_type) { + return scaled_vec_conversion( + reinterpret_cast(a), scale, fp8_type); +} + // float -> fp8 template <> __inline__ __device__ uint8_t scaled_vec_conversion( diff --git a/csrc/spinloop.cpp b/csrc/spinloop.cpp index c29e48a5f0..3285b9a3de 100644 --- a/csrc/spinloop.cpp +++ b/csrc/spinloop.cpp @@ -7,7 +7,7 @@ extern "C" { #if defined(__i386__) || defined(__x86_64__) #include - #include + #include #endif #if defined(CLOCK_MONOTONIC_RAW) diff --git a/docker/Dockerfile.xpu b/docker/Dockerfile.xpu index 8fbe4c3b5a..951fe14256 100644 --- a/docker/Dockerfile.xpu +++ b/docker/Dockerfile.xpu @@ -5,9 +5,6 @@ WORKDIR /workspace/ ARG PYTHON_VERSION=3.12 ARG PIP_EXTRA_INDEX_URL="https://download.pytorch.org/whl/xpu" -RUN wget -O- https://apt.repos.intel.com/intel-gpg-keys/GPG-PUB-KEY-INTEL-SW-PRODUCTS.PUB | gpg --dearmor | tee /usr/share/keyrings/oneapi-archive-keyring.gpg > /dev/null && \ - echo "deb [signed-by=/usr/share/keyrings/oneapi-archive-keyring.gpg] https://apt.repos.intel.com/oneapi all main" | tee /etc/apt/sources.list.d/oneAPI.list - RUN apt clean && apt-get update -y && \ apt-get install -y --no-install-recommends --fix-missing \ curl \ @@ -26,9 +23,6 @@ RUN apt clean && apt-get update -y && \ python3.12-dev \ python3-pip -RUN apt update && apt upgrade -y && \ - apt install -y intel-oneapi-compiler-dpcpp-cpp-2025.3 - # Install UMD RUN mkdir neo && \ cd neo && \ @@ -50,16 +44,6 @@ RUN curl -LsSf https://astral.sh/uv/install.sh | sh RUN uv venv --python ${PYTHON_VERSION} --seed ${VIRTUAL_ENV} ENV PATH="$VIRTUAL_ENV/bin:$PATH" -# This oneccl contains the BMG support which is not the case for default version of oneapi 2025.3. -ARG ONECCL_INSTALLER="intel-oneccl-2021.15.9.14_offline.sh" -RUN wget "https://github.com/uxlfoundation/oneCCL/releases/download/2021.15.9/${ONECCL_INSTALLER}" && \ - bash "${ONECCL_INSTALLER}" -a --silent --eula accept && \ - rm "${ONECCL_INSTALLER}" && \ - echo "source /opt/intel/oneapi/setvars.sh --force" >> /root/.bashrc && \ - echo "source /opt/intel/oneapi/ccl/2021.15/env/vars.sh --force" >> /root/.bashrc -RUN rm -f /opt/intel/oneapi/ccl/latest && \ - ln -s /opt/intel/oneapi/ccl/2021.15 /opt/intel/oneapi/ccl/latest - SHELL ["bash", "-c"] CMD ["bash", "-c", "source /root/.bashrc && exec bash"] @@ -80,14 +64,10 @@ RUN --mount=type=cache,target=/root/.cache/uv \ uv pip install --upgrade pip && \ uv pip install -r requirements/xpu.txt && \ uv pip install grpcio-tools protobuf nanobind && \ - source /opt/intel/oneapi/setvars.sh --force && \ - source /opt/intel/oneapi/ccl/2021.15/env/vars.sh --force && \ export CMAKE_PREFIX_PATH="$(python3 -c 'import site; print(site.getsitepackages()[0])'):${CMAKE_PREFIX_PATH}" && \ uv pip install --no-build-isolation -r /workspace/aphrodite/requirements/test/xpu.txt - - -ENV LD_LIBRARY_PATH="$LD_LIBRARY_PATH:/usr/local/lib/" +ENV LD_LIBRARY_PATH=/opt/venv/lib:/usr/local/lib COPY . . ARG GIT_REPO_CHECK=0 @@ -164,10 +144,6 @@ RUN --mount=type=cache,target=/root/.cache/uv \ # FIX triton RUN --mount=type=cache,target=/root/.cache/uv \ uv pip uninstall triton triton-xpu && \ - uv pip install triton-xpu==3.7.0 - -# remove torch bundled oneccl to avoid conflicts -RUN --mount=type=cache,target=/root/.cache/uv \ - uv pip uninstall oneccl oneccl-devel + uv pip install triton-xpu==3.7.2 ENTRYPOINT ["aphrodite", "serve"] diff --git a/pyproject.toml b/pyproject.toml index bfe4f81ce5..693e2c6c77 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -137,9 +137,11 @@ extend-exclude = ["tests/models/fixtures/*", "tests/prompts/*", "tests/tokenizer "examples/pooling/token_embed/*", "tests/models/language/pooling/*", "aphrodite/third_party/*", "aphrodite/entrypoints/serve/instrumentator/static/*", "tests/entrypoints/openai/speech_to_text/test_transcription_validation.py", "tests/entrypoints/speech_to_text/transcription/test_transcription_validation.py", - "rust/src/chat/src/renderer/deepseek_v32/fixtures/*", + "rust/src/chat/src/renderer/deepseek_v32/fixtures/*", "rust/src/parser/**", "docs/governance/process.md", "docs/assets/contributing/aphrodite_bench_serve_timeline.html", - "tests/v1/engine/test_fast_incdec_prefix_err.py", ".git/*"] + "tests/v1/engine/test_fast_incdec_prefix_err.py", ".git/*", + "rust/src/text/src/output/decoded.rs", + "rust/src/tokenizer/src/incremental.rs"] ignore-hidden = false [tool.typos.default] diff --git a/requirements/test/xpu.txt b/requirements/test/xpu.txt index 60ea4f5189..30fdee5025 100644 --- a/requirements/test/xpu.txt +++ b/requirements/test/xpu.txt @@ -140,7 +140,7 @@ docopt==0.6.2 # via num2words docstring-parser==0.18.0 # via anthropic -dpcpp-cpp-rt==2025.3.2 +dpcpp-cpp-rt==2026.0.0 # via # onemkl-sycl-blas # onemkl-sycl-dft @@ -253,27 +253,27 @@ ijson==3.5.0 # via -r requirements/test/../common.txt imageio==2.37.3 # via scikit-image -impi-rt==2021.17.2 +impi-rt==2021.18.0 # via # oneccl # torch iniconfig==2.3.0 # via pytest -intel-cmplr-lib-rt==2025.3.2 +intel-cmplr-lib-rt==2026.0.0 # via # intel-sycl-rt # torch -intel-cmplr-lib-ur==2025.3.2 +intel-cmplr-lib-ur==2026.0.0 # via # intel-openmp # intel-sycl-rt # torch -intel-cmplr-lic-rt==2025.3.2 +intel-cmplr-lic-rt==2026.0.0 # via # intel-opencl-rt # intel-sycl-rt # torch -intel-opencl-rt==2025.3.2 +intel-opencl-rt==2026.0.0 # via # dpcpp-cpp-rt # onemkl-sycl-blas @@ -282,14 +282,14 @@ intel-opencl-rt==2025.3.2 # onemkl-sycl-rng # onemkl-sycl-sparse # torch -intel-openmp==2025.3.2 +intel-openmp==2026.0.0 # via # dpcpp-cpp-rt # mkl # torch -intel-pti==0.16.0 +intel-pti==0.17.0 # via torch -intel-sycl-rt==2025.3.2 +intel-sycl-rt==2026.0.0 # via # dpcpp-cpp-rt # oneccl @@ -378,7 +378,7 @@ mistral-common==1.11.5 # -c requirements/common.txt # -r requirements/test/../common.txt # -r requirements/test/xpu.in -mkl==2025.3.1 +mkl==2026.0.0 # via # onemkl-sycl-blas # onemkl-sycl-dft @@ -453,28 +453,28 @@ numpy==2.2.6 # torchvision # transformers # xgrammar -oneccl==2021.17.2 +oneccl==2022.0.0 # via # oneccl-devel # torch -oneccl-devel==2021.17.2 +oneccl-devel==2022.0.0 # via torch -onemkl-license==2025.3.1 +onemkl-license==2026.0.0 # via # mkl # torch -onemkl-sycl-blas==2025.3.1 +onemkl-sycl-blas==2026.0.0 # via # onemkl-sycl-lapack # onemkl-sycl-sparse # torch -onemkl-sycl-dft==2025.3.1 +onemkl-sycl-dft==2026.0.0 # via torch -onemkl-sycl-lapack==2025.3.1 +onemkl-sycl-lapack==2026.0.0 # via torch -onemkl-sycl-rng==2025.3.1 +onemkl-sycl-rng==2026.0.0 # via torch -onemkl-sycl-sparse==2025.3.1 +onemkl-sycl-sparse==2026.0.0 # via torch openai==2.44.0 # via @@ -719,6 +719,8 @@ pyyaml==6.0.3 # timm # transformers # uvicorn +pyzes==0.1.1 + # via torch pyzmq==27.1.0 # via # -c requirements/common.txt @@ -871,14 +873,14 @@ tabledata==1.3.4 # via pytablewriter tabulate==0.10.0 # via sacrebleu -tbb==2022.3.1 +tbb==2023.0.0 # via # intel-opencl-rt # mkl # torch tblib==3.1.0 # via -r requirements/test/xpu.in -tcmlib==1.4.1 +tcmlib==1.5.0 # via # tbb # torch @@ -910,7 +912,7 @@ tokenizers==0.22.2 # -c requirements/common.txt # -r requirements/test/../common.txt # transformers -torch==2.12.0+xpu +torch==2.13.0+xpu # via # -c requirements/xpu.txt # accelerate @@ -920,7 +922,7 @@ torch==2.12.0+xpu # timm # torchvision # xgrammar -torchvision==0.27.0+xpu +torchvision==0.28.0+xpu # via timm tqdm==4.67.3 # via @@ -946,7 +948,7 @@ transformers==5.14.1 # xgrammar triton==3.7.1 # via xgrammar -triton-xpu==3.7.1 +triton-xpu==3.7.2 # via torch typepy==1.3.4 # via @@ -1001,7 +1003,7 @@ typing-inspection==0.4.2 # mcp # pydantic # pydantic-settings -umf==1.0.3 +umf==1.1.0 # via # intel-cmplr-lib-ur # torch diff --git a/requirements/xpu.txt b/requirements/xpu.txt index 4d9ab239d2..491ff65020 100644 --- a/requirements/xpu.txt +++ b/requirements/xpu.txt @@ -12,10 +12,10 @@ jinja2>=3.1.6 datasets # for benchmark scripts numba == 0.65.0 # Required for N-gram speculative decoding --extra-index-url=https://download.pytorch.org/whl/xpu -torch==2.12.0 +torch==2.13.0 torchaudio torchvision torchcodec >= 0.14 # Required for the torchcodec video decoding backend -auto_round_lib==0.14.1 -aphrodite_xpu_kernels @ https://github.com/vllm-project/vllm-xpu-kernels/releases/download/v0.1.11.1/vllm_xpu_kernels-0.1.11.1-cp38-abi3-manylinux_2_28_x86_64.whl +auto_round_lib==0.14.2 +aphrodite_xpu_kernels @ https://github.com/vllm-project/vllm-xpu-kernels/releases/download/v0.1.12/vllm_xpu_kernels-0.1.12-cp38-abi3-manylinux_2_28_x86_64.whl diff --git a/rust/Cargo.lock b/rust/Cargo.lock index c03f287ed4..66ac60ed55 100644 --- a/rust/Cargo.lock +++ b/rust/Cargo.lock @@ -131,6 +131,7 @@ name = "aphrodite-bench" version = "0.1.0" dependencies = [ "anyhow", + "aphrodite-tracing", "base64 0.22.1", "bytes", "chrono", @@ -156,7 +157,6 @@ dependencies = [ "tokio", "tokio-stream", "tracing", - "tracing-subscriber", "url", "uuid", ] @@ -218,6 +218,7 @@ dependencies = [ "aphrodite-engine-core-client", "aphrodite-managed-engine", "aphrodite-server", + "aphrodite-tracing", "clap", "educe", "expect-test", @@ -228,11 +229,9 @@ dependencies = [ "serde_json", "serde_with", "thiserror-ext", - "time", "tokio", "tokio-util", "tracing", - "tracing-subscriber", "uuid", ] @@ -469,6 +468,15 @@ dependencies = [ "tracing", ] +[[package]] +name = "aphrodite-tracing" +version = "0.1.0" +dependencies = [ + "time", + "tracing", + "tracing-subscriber", +] + [[package]] name = "aphrodite-tool-parser-py" version = "0.1.0" @@ -2574,7 +2582,7 @@ checksum = "11d3d7f243d5c5a8b9bb5d6dd2b1602c0cb0b9db1621bafc7ed66e35ff9fe092" [[package]] name = "llm-multimodal" version = "1.7.1" -source = "git+https://github.com/smg-project/llm-multimodal?rev=5390032d6dc8a3e6fdc83acd320260367eb4b9b5#5390032d6dc8a3e6fdc83acd320260367eb4b9b5" +source = "git+https://github.com/smg-project/llm-multimodal?rev=15adba5e025d8636ba4a334fb379b1371f6196a1#15adba5e025d8636ba4a334fb379b1371f6196a1" dependencies = [ "anyhow", "base64 0.22.1", diff --git a/rust/Cargo.toml b/rust/Cargo.toml index d0b3ada027..307e3a5541 100644 --- a/rust/Cargo.toml +++ b/rust/Cargo.toml @@ -13,6 +13,7 @@ members = [ "src/server", "src/text", "src/tokenizer", + "src/tracing", ] resolver = "3" @@ -58,7 +59,7 @@ indexmap = "2.13.0" indicatif = "0.18.4" itertools = "0.14.0" libc = "0.2.177" -llm-multimodal = { git = "https://github.com/smg-project/llm-multimodal", rev = "5390032d6dc8a3e6fdc83acd320260367eb4b9b5", default-features = false, features = ["native-tls"] } +llm-multimodal = { git = "https://github.com/smg-project/llm-multimodal", rev = "15adba5e025d8636ba4a334fb379b1371f6196a1", default-features = false, features = ["native-tls"] } 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"] } @@ -143,6 +144,7 @@ aphrodite-parser = { path = "src/parser" } aphrodite-server = { path = "src/server" } aphrodite-text = { path = "src/text" } aphrodite-tokenizer = { path = "src/tokenizer" } +aphrodite-tracing = { path = "src/tracing" } winnow = { version = "1.0.2", features = ["simd"] } xgrammar-structural-tag = "0.2.0" zeromq = { version = "0.6.0", default-features = false, features = [ diff --git a/rust/src/bench/Cargo.toml b/rust/src/bench/Cargo.toml index 5e134a8e8a..cbb3d72c79 100644 --- a/rust/src/bench/Cargo.toml +++ b/rust/src/bench/Cargo.toml @@ -32,9 +32,9 @@ tokenizers.workspace = true tokio.workspace = true tokio-stream.workspace = true tracing.workspace = true -tracing-subscriber.workspace = true url.workspace = true uuid.workspace = true +aphrodite-tracing.workspace = true [lints] workspace = true diff --git a/rust/src/bench/src/main.rs b/rust/src/bench/src/main.rs index 271761600e..5c16d1bd17 100644 --- a/rust/src/bench/src/main.rs +++ b/rust/src/bench/src/main.rs @@ -19,18 +19,8 @@ struct Cli { args: aphrodite_bench::BenchServeArgs, } -// TODO: unify the tracing subscriber used by different binaries. -fn init_tracing() { - let filter = tracing_subscriber::EnvFilter::try_from_default_env() - .unwrap_or_else(|_| tracing_subscriber::EnvFilter::new("info")); - let _ = tracing_subscriber::fmt() - .with_env_filter(filter) - .with_writer(std::io::stderr) - .try_init(); -} - fn main() -> anyhow::Result<()> { - init_tracing(); + aphrodite_tracing::init_tracing("Bench"); let cli = Cli::parse(); diff --git a/rust/src/chat/src/backend/hf.rs b/rust/src/chat/src/backend/hf.rs index 54476fed62..6b7b061ede 100644 --- a/rust/src/chat/src/backend/hf.rs +++ b/rust/src/chat/src/backend/hf.rs @@ -20,7 +20,7 @@ use crate::output::{ use crate::renderer::hf::{HfChatRenderer, MultimodalRenderInfo}; use crate::renderer::{ DeepSeekV4ChatRenderer, DeepSeekV32ChatRenderer, DynChatRenderer, HarmonyChatRenderer, - InklingChatRenderer, + InklingChatRenderer, KimiK3ChatRenderer, }; use crate::request::ChatRequest; use crate::{DynChatOutputProcessor, RendererSelection}; @@ -57,6 +57,7 @@ impl HfChatBackend { processor_config: files.processor_config_path.as_deref(), }, tokenizer.clone(), + options.limit_mm_per_prompt.clone(), )? }; let multimodal_render_info = resolve_multimodal_render_info(multimodal_model_info.as_ref()); @@ -73,6 +74,7 @@ impl HfChatBackend { RendererSelection::DeepSeekV4 => Arc::new(DeepSeekV4ChatRenderer::new()), RendererSelection::Harmony => Arc::new(HarmonyChatRenderer::new()?), RendererSelection::Inkling => Arc::new(InklingChatRenderer::new(tokenizer.clone())?), + RendererSelection::KimiK3 => Arc::new(KimiK3ChatRenderer::new(tokenizer.clone())), }; info!( @@ -230,6 +232,7 @@ mod tests { chat_template_content_format: Default::default(), chat_template: None, default_chat_template_kwargs: HashMap::new(), + limit_mm_per_prompt: HashMap::new(), }, test_tokenizer(), ) diff --git a/rust/src/chat/src/backend/mod.rs b/rust/src/chat/src/backend/mod.rs index 298c4429e3..da366c82e4 100644 --- a/rust/src/chat/src/backend/mod.rs +++ b/rust/src/chat/src/backend/mod.rs @@ -8,7 +8,7 @@ use serde_json::Value; use aphrodite_text::{DynTextBackend, TextBackend}; use crate::error::Result; -use crate::multimodal::MultimodalModelInfo; +use crate::multimodal::{MmLimitPerPrompt, MultimodalModelInfo}; use crate::output::DynChatOutputProcessor; use crate::renderer::DynChatRenderer; use crate::request::ChatRequest; @@ -74,6 +74,9 @@ pub struct LoadModelBackendsOptions { /// Optional server-default keyword arguments merged into every /// chat-template render before request-level `chat_template_kwargs`. pub default_chat_template_kwargs: HashMap, + /// Maximum number of input items allowed per prompt for each modality. + /// Unspecified modalities are unlimited. + pub limit_mm_per_prompt: MmLimitPerPrompt, } /// Shared backends loaded from a model id. diff --git a/rust/src/chat/src/error.rs b/rust/src/chat/src/error.rs index f552983726..b1e72f7ba3 100644 --- a/rust/src/chat/src/error.rs +++ b/rust/src/chat/src/error.rs @@ -23,6 +23,8 @@ pub enum Error { UnsupportedMultimodalContent(&'static str), #[error("`{modality}` input is not supported by this model")] UnsupportedModality { modality: String }, + #[error("At most {limit} {modality}(s) may be provided in one prompt.")] + MmLimitExceeded { modality: String, limit: usize }, #[error("multimodal preprocessing error: {0}")] Multimodal(#[message] String), #[error("{kind} parsing is not available for model `{model_id}`")] @@ -87,7 +89,8 @@ impl Error { Self::Text(error) => error.is_request_validation_error(), Self::UnsupportedMultimodalRenderer | Self::UnsupportedMultimodalContent(_) - | Self::UnsupportedModality { .. } => true, + | Self::UnsupportedModality { .. } + | Self::MmLimitExceeded { .. } => true, _ => false, } diff --git a/rust/src/chat/src/lib.rs b/rust/src/chat/src/lib.rs index 35ef6bf555..5a10d0312f 100644 --- a/rust/src/chat/src/lib.rs +++ b/rust/src/chat/src/lib.rs @@ -33,7 +33,8 @@ pub use parser::tool::{ToolParser, ToolParserError, ToolParserFactory}; pub use renderer::hf::ChatTemplateContentFormatOption; pub use renderer::{ ChatRenderer, DeepSeekV4ChatRenderer, DeepSeekV32ChatRenderer, DynChatRenderer, - HarmonyChatRenderer, InklingChatRenderer, RenderedPrompt, RendererSelection, + HarmonyChatRenderer, InklingChatRenderer, KimiK3ChatRenderer, RenderedPrompt, + RendererSelection, }; pub use request::{ ChatContent, ChatContentPart, ChatMessage, ChatOptions, ChatRequest, ChatRole, ChatTool, @@ -353,6 +354,12 @@ mod tests { .unwrap(); } + #[test] + fn validate_parser_overrides_accepts_explicit_kimi_k3() { + let selection = ParserSelection::Explicit("kimi_k3".to_string()); + validate_parser_overrides(&selection, &selection).unwrap(); + } + #[test] fn validate_parser_overrides_accepts_auto_and_none() { validate_parser_overrides(&ParserSelection::Auto, &ParserSelection::None).unwrap(); @@ -366,7 +373,7 @@ mod tests { ) .unwrap_err(); - expect_test::expect!["tool parser `definitely_missing_tool_parser` is not registered (choose from: deepseek_v3, deepseek_v31, deepseek_v32, deepseek_v4, gemma4, glm45, glm47, granite4, hermes, hy_v3, inkling, internlm, kimi_k2, llama3_json, llama4_json, minimax_m2, minimax_m3, mistral, phi4_mini_json, qwen3_coder, qwen3_xml, seed_oss)"].assert_eq(&error.to_report_string()); + expect_test::expect!["tool parser `definitely_missing_tool_parser` is not registered (choose from: deepseek_v3, deepseek_v31, deepseek_v32, deepseek_v4, gemma4, glm45, glm47, granite4, hermes, hy_v3, inkling, internlm, kimi_k2, kimi_k3, llama3_json, llama4_json, minimax_m2, minimax_m3, mistral, phi4_mini_json, qwen3_coder, qwen3_xml, seed_oss)"].assert_eq(&error.to_report_string()); } #[test] @@ -377,6 +384,6 @@ mod tests { ) .unwrap_err(); - expect_test::expect!["reasoning parser `definitely_missing_reasoning_parser` is not registered (choose from: cohere_cmd, deepseek_r1, deepseek_v3, deepseek_v4, gemma4, glm45, inkling, kimi, kimi_k2, minimax_m2, minimax_m3, nemotron_v3, qwen3, seed_oss, step3, step3p5)"].assert_eq(&error.to_report_string()); + expect_test::expect!["reasoning parser `definitely_missing_reasoning_parser` is not registered (choose from: cohere_cmd, deepseek_r1, deepseek_v3, deepseek_v4, gemma4, glm45, inkling, kimi, kimi_k2, kimi_k3, minimax_m2, minimax_m3, nemotron_v3, qwen3, seed_oss, step3, step3p5)"].assert_eq(&error.to_report_string()); } } diff --git a/rust/src/chat/src/multimodal.rs b/rust/src/chat/src/multimodal.rs index bd3cac01d3..0a03ec8717 100644 --- a/rust/src/chat/src/multimodal.rs +++ b/rust/src/chat/src/multimodal.rs @@ -11,7 +11,7 @@ //! Raw media stays above `aphrodite-text`; this module lowers it into token IDs and //! opaque tensor payloads before the request is handed to text generation. -use std::collections::HashSet; +use std::collections::{BTreeMap, HashMap, HashSet}; use std::fs; use std::path::Path; use std::sync::{Arc, LazyLock}; @@ -24,6 +24,7 @@ use llm_multimodal::{ PromptReplacement, Tokenizer as TokenResolver, TrackedMedia, VideoClip, VisionPreProcessor, VisionProcessorRegistry, }; +use serde::{Deserialize, Serialize}; use thiserror_ext::AsReport as _; use tracing::warn; use aphrodite_engine_core_client::protocol::dtype::ModelDtype; @@ -52,6 +53,71 @@ pub struct MultimodalModelInfo { video: Option, audio: Option, media_connector: Arc, + /// Maximum number of input items allowed per prompt for each modality. + limit_mm_per_prompt: MmLimitPerPrompt, +} + +/// Per-modality item-count limits configured by `--limit-mm-per-prompt`. +/// +/// Modalities absent from the map are unlimited. +pub type MmLimitPerPrompt = HashMap; + +/// Modalities that `--limit-mm-per-prompt` can be keyed by. +/// +/// Closed on purpose: these are exactly the keys Python accepts, per +/// `MultiModalDummyOptionsBuiltins` in `vllm/config/multimodal.py`. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum MmLimitModality { + Image, + Audio, + Video, +} + +impl MmLimitModality { + /// The wire name, matching Python's modality strings. + pub fn as_str(self) -> &'static str { + match self { + Self::Image => "image", + Self::Audio => "audio", + Self::Video => "video", + } + } +} + +/// One modality's limit, in either of the two shapes Python accepts. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde( + untagged, + expecting = "an item count, or an object with an optional `count` field" +)] +pub enum MmLimitSpec { + /// Legacy form: `"image": 16` + Count(usize), + + /// Configurable form: + /// `"video": {"count": 1, "num_frames": 32}` + Options { + /// Absent means unlimited, matching an absent modality. + #[serde(default, skip_serializing_if = "Option::is_none")] + count: Option, + + /// Preserve Python-owned options for forwarding. Never interpreted + /// here: they size the engine's dummy-profiling encoder cache, which + /// has no Rust counterpart. + #[serde(flatten)] + extra: BTreeMap, + }, +} + +impl MmLimitSpec { + /// The configured item count, or `None` when this modality is unlimited. + pub fn count(&self) -> Option { + match self { + Self::Count(count) => Some(*count), + Self::Options { count, .. } => *count, + } + } } /// Model metadata and tokenizer access shared by all multimodal specs. @@ -289,6 +355,7 @@ impl MultimodalModelInfo { model_type: Option, files: MultimodalConfigFiles<'_>, tokenizer: DynTokenizer, + limit_mm_per_prompt: MmLimitPerPrompt, ) -> Result> { let config = match files.config { Some(path) => { @@ -321,7 +388,12 @@ impl MultimodalModelInfo { tokenizer: TokenizerResolver(tokenizer), }; - Self::from_loaded(context, preprocessor_config, video_preprocessor_config) + Self::from_loaded( + context, + preprocessor_config, + video_preprocessor_config, + limit_mm_per_prompt, + ) } /// Resolve multimodal support from an assembled context and parsed @@ -330,6 +402,7 @@ impl MultimodalModelInfo { context: MultimodalModelContext, preprocessor_config: PreProcessorConfig, video_preprocessor_config: PreProcessorConfig, + limit_mm_per_prompt: MmLimitPerPrompt, ) -> Result> { let (image, video) = Self::resolve_vision_lanes( &context, @@ -358,6 +431,7 @@ impl MultimodalModelInfo { video, audio, media_connector, + limit_mm_per_prompt, })) } @@ -572,7 +646,55 @@ fn input_audio_data_url(data: &str, format: Option<&str>) -> Result { Ok(format!("data:{mime_type};base64,{data}")) } +/// The modality a content part counts against, or `None` for plain text. +/// +/// Embedding inputs share their base modality's budget rather than getting one +/// of their own, matching Python's `modality.replace("_embeds", "")` in +/// `vllm/entrypoints/chat_utils.py`. +fn media_part_limit_modality(part: &MediaContentPart) -> Option { + match part { + MediaContentPart::Text { .. } => None, + MediaContentPart::ImageUrl { .. } + | MediaContentPart::ImageData { .. } + | MediaContentPart::ImageEmbeds { .. } => Some(MmLimitModality::Image), + MediaContentPart::AudioUrl { .. } | MediaContentPart::AudioData { .. } => { + Some(MmLimitModality::Audio) + } + MediaContentPart::VideoUrl { .. } | MediaContentPart::VideoData { .. } => { + Some(MmLimitModality::Video) + } + } +} + impl MultimodalModelInfo { + /// Reject requests exceeding `--limit-mm-per-prompt`'s configured + /// per-modality item count, before any fetch/decode work is spent on them. + /// + /// Modalities without a configured count are unlimited. + fn validate_mm_limits(&self, media_parts: &[MediaContentPart]) -> Result<()> { + let mut counts: HashMap = HashMap::new(); + for part in media_parts { + if let Some(modality) = media_part_limit_modality(part) { + *counts.entry(modality).or_default() += 1; + } + } + + for (modality, count) in counts { + let Some(limit) = self.limit_mm_per_prompt.get(&modality).and_then(MmLimitSpec::count) + else { + continue; + }; + if count > limit { + return Err(Error::MmLimitExceeded { + modality: modality.as_str().to_string(), + limit, + }); + } + } + + Ok(()) + } + /// Run media fetch, per-modality preprocessing, prompt expansion, and /// feature build. /// @@ -590,8 +712,7 @@ impl MultimodalModelInfo { return Ok(Vec::new()); } - // TODO: enforce per-modality item-count limits, aligned with the - // engine's `--limit-mm-per-prompt` semantics. + self.validate_mm_limits(&media_parts)?; let fetched = self.fetch_media(media_parts).await?; let mut prepared = Vec::new(); @@ -759,10 +880,11 @@ mod tests { .with_regular_token("<|video_pad|>", QWEN3_VIDEO_PAD_ID) } - fn test_info( + fn test_info_with_limits( model_type: &str, config: serde_json::Value, tokenizer: TestTokenizer, + limit_mm_per_prompt: MmLimitPerPrompt, ) -> MultimodalModelInfo { let context = MultimodalModelContext { model_id: format!("{model_type}-test"), @@ -775,11 +897,20 @@ mod tests { context, PreProcessorConfig::default(), PreProcessorConfig::default(), + limit_mm_per_prompt, ) .unwrap() .unwrap_or_else(|| panic!("{model_type} multimodal support should resolve")) } + fn test_info( + model_type: &str, + config: serde_json::Value, + tokenizer: TestTokenizer, + ) -> MultimodalModelInfo { + test_info_with_limits(model_type, config, tokenizer, HashMap::new()) + } + fn llama4_info() -> MultimodalModelInfo { let config = serde_json::json!({ "model_type": "llama4", @@ -789,16 +920,19 @@ mod tests { test_info("llama4", config, llama4_tokenizer()) } - pub(super) fn qwen3_vl_info() -> MultimodalModelInfo { - let config = serde_json::json!({ + fn qwen3_vl_config() -> serde_json::Value { + 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()) + }) + } + + pub(super) fn qwen3_vl_info() -> MultimodalModelInfo { + test_info("qwen3_vl", qwen3_vl_config(), qwen3_vl_tokenizer()) } #[test] @@ -895,4 +1029,123 @@ mod tests { ); assert!(input_audio_data_url("AAAA", Some("flac")).is_err()); } + + fn image_url_part() -> MediaContentPart { + MediaContentPart::ImageUrl { + url: "https://example.com/image.png".to_string(), + detail: None, + uuid: None, + } + } + + fn qwen3_vl_info_with_limits(limit_mm_per_prompt: MmLimitPerPrompt) -> MultimodalModelInfo { + test_info_with_limits( + "qwen3_vl", + qwen3_vl_config(), + qwen3_vl_tokenizer(), + limit_mm_per_prompt, + ) + } + + #[test] + fn validate_mm_limits_ignores_text_parts() { + let info = qwen3_vl_info(); + let parts = vec![ + MediaContentPart::Text { + text: "hello".to_string(), + }, + MediaContentPart::Text { + text: "world".to_string(), + }, + ]; + assert!(info.validate_mm_limits(&parts).is_ok()); + } + + #[test] + fn validate_mm_limits_leaves_unconfigured_modalities_unlimited() { + let info = qwen3_vl_info(); + let parts: Vec<_> = std::iter::repeat_with(image_url_part).take(1_000).collect(); + + assert!(info.validate_mm_limits(&parts).is_ok()); + } + + #[test] + fn validate_mm_limits_enforces_configured_limit_at_the_boundary() { + let info = qwen3_vl_info_with_limits(HashMap::from([( + MmLimitModality::Image, + MmLimitSpec::Count(1), + )])); + assert!(info.validate_mm_limits(&[image_url_part()]).is_ok()); + + let error = info.validate_mm_limits(&[image_url_part(), image_url_part()]).unwrap_err(); + assert_eq!( + error.to_report_string(), + "At most 1 image(s) may be provided in one prompt." + ); + // Confirms the HTTP-mapping bug found during implementation stays fixed: + // this must map to 400, not 500. + assert!(error.is_request_validation_error()); + } + + #[test] + fn validate_mm_limits_counts_image_embeds_against_the_image_limit() { + let info = qwen3_vl_info_with_limits(HashMap::from([( + MmLimitModality::Image, + MmLimitSpec::Count(1), + )])); + let image_embeds_part = MediaContentPart::ImageEmbeds { + payload: serde_json::Value::String("AAAA".to_string()), + uuid: None, + }; + + let error = info.validate_mm_limits(&[image_url_part(), image_embeds_part]).unwrap_err(); + assert_eq!( + error.to_report_string(), + "At most 1 image(s) may be provided in one prompt." + ); + } + + /// An options object without a `count` carries only profiling keys, which + /// say nothing about how many items are allowed. + #[test] + fn validate_mm_limits_treats_a_count_less_options_object_as_unlimited() { + let info = qwen3_vl_info_with_limits(HashMap::from([( + MmLimitModality::Image, + MmLimitSpec::Options { + count: None, + extra: BTreeMap::from([("width".to_string(), serde_json::json!(512))]), + }, + )])); + + assert!(info.validate_mm_limits(&[image_url_part(), image_url_part()]).is_ok()); + } + + fn parse_limits(json: &str) -> MmLimitPerPrompt { + serde_json::from_str(json).expect("limit map should parse") + } + + #[test] + fn limit_map_parses_both_shapes_python_accepts() { + let limits = parse_limits(r#"{"image": 16, "video": {"count": 1, "num_frames": 32}}"#); + + assert_eq!(limits[&MmLimitModality::Image].count(), Some(16)); + assert_eq!(limits[&MmLimitModality::Video].count(), Some(1)); + assert_eq!(limits.get(&MmLimitModality::Audio), None); + } + + #[test] + fn limit_map_rejects_keys_python_does_not_accept() { + assert!(serde_json::from_str::(r#"{"image_embeds": 1}"#).is_err()); + } + + /// Managed mode forwards this map back to Python as JSON, where + /// `BaseDummyOptions.count` is a non-optional `int` under `extra="forbid"`. + /// Emitting `"count": null` would make the engine subprocess fail to start. + #[test] + fn limit_map_round_trips_without_emitting_a_null_count() { + let source = r#"{"video":{"num_frames":32}}"#; + let encoded = serde_json::to_string(&parse_limits(source)).expect("map should serialize"); + + assert_eq!(encoded, source); + } } diff --git a/rust/src/chat/src/multimodal/video.rs b/rust/src/chat/src/multimodal/video.rs index 7f8958cdac..5eee545af8 100644 --- a/rust/src/chat/src/multimodal/video.rs +++ b/rust/src/chat/src/multimodal/video.rs @@ -207,6 +207,7 @@ mod tests { Some("qwen3_vl".to_string()), files, Arc::new(qwen3_vl_tokenizer()), + std::collections::HashMap::new(), ) }; diff --git a/rust/src/chat/src/parser/reasoning/mod.rs b/rust/src/chat/src/parser/reasoning/mod.rs index 8f5004b60b..2246602a8b 100644 --- a/rust/src/chat/src/parser/reasoning/mod.rs +++ b/rust/src/chat/src/parser/reasoning/mod.rs @@ -27,6 +27,7 @@ pub mod names { pub const GLM45: &str = "glm45"; pub const KIMI: &str = "kimi"; pub const KIMI_K2: &str = "kimi_k2"; + pub const KIMI_K3: &str = "kimi_k3"; pub const MINIMAX_M2: &str = "minimax_m2"; pub const MINIMAX_M3: &str = "minimax_m3"; pub const NEMOTRON_V3: &str = "nemotron_v3"; @@ -68,6 +69,7 @@ impl ReasoningParserFactory { .register_parser::(names::GLM45) .register_parser::(names::KIMI) .register_parser::(names::KIMI_K2) + .register_unified_dummy(names::KIMI_K3) .register_parser::(names::MINIMAX_M2) .register_parser::(names::MINIMAX_M3) .register_parser::(names::NEMOTRON_V3) diff --git a/rust/src/chat/src/parser/tool/mod.rs b/rust/src/chat/src/parser/tool/mod.rs index d7da6db294..492977ee18 100644 --- a/rust/src/chat/src/parser/tool/mod.rs +++ b/rust/src/chat/src/parser/tool/mod.rs @@ -33,6 +33,7 @@ pub mod names { // also routes to `Internlm2ToolParser` despite the version-agnostic name. pub const INTERNLM: &str = "internlm"; pub const KIMI_K2: &str = "kimi_k2"; + pub const KIMI_K3: &str = "kimi_k3"; pub const LLAMA3_JSON: &str = "llama3_json"; pub const LLAMA4_JSON: &str = "llama4_json"; pub const MINIMAX_M2: &str = "minimax_m2"; @@ -78,6 +79,7 @@ impl ToolParserFactory { .register_parser::(names::HY_V3) .register_parser::(names::INTERNLM) .register_parser::(names::KIMI_K2) + .register_unified_dummy(names::KIMI_K3) .register_parser::(names::LLAMA3_JSON) .register_parser::(names::LLAMA4_JSON) .register_parser::(names::MINIMAX_M2) diff --git a/rust/src/chat/src/parser/unified.rs b/rust/src/chat/src/parser/unified.rs index 9cb9a1c47e..3c9a131904 100644 --- a/rust/src/chat/src/parser/unified.rs +++ b/rust/src/chat/src/parser/unified.rs @@ -5,7 +5,9 @@ use std::sync::LazyLock; -pub use aphrodite_parser::unified::{Gemma4UnifiedParser, InklingUnifiedParser, UnifiedParser}; +pub use aphrodite_parser::unified::{ + Gemma4UnifiedParser, InklingUnifiedParser, KimiK3UnifiedParser, UnifiedParser, +}; use aphrodite_tokenizer::DynTokenizer; use crate::parser::ParserFactory; @@ -15,6 +17,7 @@ use crate::request::ChatTool; pub mod names { pub const GEMMA4: &str = "gemma4"; pub const INKLING: &str = "inkling"; + pub const KIMI_K3: &str = "kimi_k3"; } /// Constructor signature for one registered unified parser implementation. @@ -39,11 +42,14 @@ impl UnifiedParserFactory { factory.register_parser::(names::GEMMA4); factory.register_parser::(names::INKLING); + factory.register_parser::(names::KIMI_K3); factory .register_pattern("gemma-4", names::GEMMA4) .register_pattern("gemma4", names::GEMMA4) - .register_pattern("inkling", names::INKLING); + .register_pattern("inkling", names::INKLING) + .register_pattern("kimi-k3", names::KIMI_K3) + .register_pattern("kimi_k3", names::KIMI_K3); factory } @@ -121,4 +127,20 @@ mod tests { ); factory.create(names::INKLING, &[], Arc::new(inkling_tokenizer())).unwrap(); } + + #[test] + fn factory_registers_kimi_k3() { + let factory = UnifiedParserFactory::new(); + let tokenizer = TestTokenizer::new() + .with_regular_token("<|open|>", 1001) + .with_regular_token("<|close|>", 1002) + .with_regular_token("<|sep|>", 1003); + + assert!(factory.contains(names::KIMI_K3)); + assert_eq!( + factory.resolve_name_for_model("moonshotai/Kimi-K3"), + Some(names::KIMI_K3) + ); + factory.create(names::KIMI_K3, &[], Arc::new(tokenizer)).unwrap(); + } } diff --git a/rust/src/chat/src/renderer/deepseek_v32/tests.rs b/rust/src/chat/src/renderer/deepseek_v32/tests.rs index 465bcf88d1..2fa4f26724 100644 --- a/rust/src/chat/src/renderer/deepseek_v32/tests.rs +++ b/rust/src/chat/src/renderer/deepseek_v32/tests.rs @@ -59,7 +59,7 @@ fn fixture_request(input_name: &str) -> ChatRequest { fn deepseek_fixture_options() -> FixtureRequestOptions { FixtureRequestOptions { - enable_thinking: true, + enable_thinking: Some(true), no_generation_prompt_when_last_assistant: true, } } diff --git a/rust/src/chat/src/renderer/deepseek_v4/tests.rs b/rust/src/chat/src/renderer/deepseek_v4/tests.rs index 73380cef26..9068d460ee 100644 --- a/rust/src/chat/src/renderer/deepseek_v4/tests.rs +++ b/rust/src/chat/src/renderer/deepseek_v4/tests.rs @@ -27,7 +27,7 @@ fn fixture_request(input_name: &str) -> ChatRequest { fn deepseek_fixture_options() -> FixtureRequestOptions { FixtureRequestOptions { - enable_thinking: true, + enable_thinking: Some(true), no_generation_prompt_when_last_assistant: true, } } diff --git a/rust/src/chat/src/renderer/harmony/tests.rs b/rust/src/chat/src/renderer/harmony/tests.rs index 301d9707e7..4849c93223 100644 --- a/rust/src/chat/src/renderer/harmony/tests.rs +++ b/rust/src/chat/src/renderer/harmony/tests.rs @@ -22,7 +22,7 @@ fn fixture_request(input_name: &str) -> ChatRequest { fixture_chat_request( &fixture_path(input_name), FixtureRequestOptions { - enable_thinking: false, + enable_thinking: Some(false), no_generation_prompt_when_last_assistant: false, }, ) diff --git a/rust/src/chat/src/renderer/inkling/tests.rs b/rust/src/chat/src/renderer/inkling/tests.rs index e232863ba6..4faec70700 100644 --- a/rust/src/chat/src/renderer/inkling/tests.rs +++ b/rust/src/chat/src/renderer/inkling/tests.rs @@ -105,7 +105,7 @@ fn fixture_request(name: &str) -> ChatRequest { fn inkling_fixture_options() -> FixtureRequestOptions { FixtureRequestOptions { - enable_thinking: false, + enable_thinking: Some(false), no_generation_prompt_when_last_assistant: false, } } diff --git a/rust/src/chat/src/renderer/kimi_k3/encoding.rs b/rust/src/chat/src/renderer/kimi_k3/encoding.rs new file mode 100644 index 0000000000..49e74e99fb --- /dev/null +++ b/rust/src/chat/src/renderer/kimi_k3/encoding.rs @@ -0,0 +1,608 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +//! Kimi K3 XTML prompt renderer. +//! +//! Port of Moonshot remote-code `encoding_k3.py::build_chat_segments()`. + +use std::collections::HashMap; + +use serde_json::{Map, Value, json}; +use aphrodite_tokenizer::Tokenizer; + +use crate::error::{Error, Result}; +use crate::request::{ + ChatContent, ChatContentPart, ChatMessage, ChatRequest, ChatTool, ChatToolChoice, +}; +use crate::{AssistantContentBlock, AssistantToolCall}; + +pub(super) const OPEN: &str = "<|open|>"; +pub(super) const CLOSE: &str = "<|close|>"; +pub(super) const SEP: &str = "<|sep|>"; +pub(super) const END_OF_MSG: &str = "<|end_of_msg|>"; +pub(super) const IMAGE_PLACEHOLDER: &str = "<|media_pad|>"; + +const DEFAULT_THINKING_EFFORT: &str = "max"; +const VALID_THINKING_EFFORTS: &[&str] = &["low", "high", "max"]; + +/// K3 prompt encoder preserving Python's per-segment tokenization boundaries. +pub(super) struct K3TokenWriter<'a> { + tokenizer: &'a dyn Tokenizer, + token_ids: Vec, +} + +impl<'a> K3TokenWriter<'a> { + pub(super) fn new(tokenizer: &'a dyn Tokenizer) -> Self { + Self { + tokenizer, + token_ids: Vec::new(), + } + } + + /// Encode one trusted segment with normal added-token recognition. + pub(super) fn control(&mut self, text: &str) -> Result<()> { + if !text.is_empty() { + self.token_ids.extend(self.tokenizer.encode(text, false)?); + } + Ok(()) + } + + /// Encode one literal segment while bypassing every added-token matcher. + pub(super) fn ordinary(&mut self, text: &str) -> Result<()> { + if !text.is_empty() { + self.token_ids.extend(self.tokenizer.encode_ordinary(text)?); + } + Ok(()) + } + + pub(super) fn finish(self) -> Vec { + self.token_ids + } +} + +/// Render and tokenize one chat request using K3's segment-aware contract. +pub(super) fn render_request(request: &ChatRequest, tokenizer: &dyn Tokenizer) -> Result> { + let thinking = thinking_enabled(request)?; + let thinking_effort = thinking.then(|| thinking_effort(request)).transpose()?; + let tools = request_tools(request); + let mut out = K3TokenWriter::new(tokenizer); + + if !tools.is_empty() { + write_tool_declare(&mut out, tools, false)?; + } + + if let Some(effort) = thinking_effort { + // Preserve the checkpoint's literal guidance text: it still names + // `medium`, although the validator above no longer accepts it. + write_internal_system( + &mut out, + "thinking-effort", + &format!( + "`thinking_effort` guides on how much to think in your \ + thinking channel (not including the response channel), \ + supported values include `low`, `medium`, `high`, and `max`.\n\ + Now the system is invoked with `thinking_effort={effort}`." + ), + )?; + } + + // Track prior assistant tool-call ids for tool-result reordering / naming. + let mut tool_call_id_index: HashMap = HashMap::new(); + let mut pending_tool_run: Vec<(usize, ChatMessage)> = Vec::new(); + + let flush_tool_run = |out: &mut K3TokenWriter<'_>, + run: &mut Vec<(usize, ChatMessage)>, + id_index: &HashMap| + -> Result<()> { + if run.is_empty() { + return Ok(()); + } + + let mut resolved = Vec::with_capacity(run.len()); + let mut unresolved = false; + for (offset, message) in run.drain(..) { + let ChatMessage::ToolResponse { + content, + tool_call_id, + } = message + else { + unreachable!("pending tool run only holds tool responses"); + }; + match id_index.get(&tool_call_id) { + Some(&(position, ref name)) => { + resolved.push((position, offset, content, Some(name.clone()))); + } + None => { + unresolved = true; + resolved.push((usize::MAX, offset, content, None)); + } + } + } + + if unresolved { + // Preserve original order when the run cannot be fully matched. + resolved.sort_by_key(|item| item.1); + } else { + resolved.sort_by_key(|item| (item.0, item.1)); + } + + for (xtml_index, (_, _, content, name)) in resolved.into_iter().enumerate() { + let tool_name = name.as_deref().ok_or_else(|| { + Error::ChatTemplate( + "Kimi K3 tool messages need a resolvable tool name: \ + carry a matching tool_call_id against a preceding \ + assistant tool_call" + .to_string(), + ) + })?; + write_tool_message(out, tool_name, xtml_index + 1, &content)?; + } + Ok(()) + }; + + for (message_index, message) in request.messages.iter().enumerate() { + match message { + ChatMessage::ToolResponse { .. } => { + pending_tool_run.push((message_index, message.clone())); + continue; + } + _ => { + flush_tool_run(&mut out, &mut pending_tool_run, &tool_call_id_index)?; + } + } + + match message { + // Python: role=system with a `tools` field renders a dynamic + // tool-declare (`## New Tools Available`). Map that to Developer + // messages that carry tools (OpenAI "developer" / system-tools). + ChatMessage::Developer { + content, + tools: Some(local_tools), + } if !local_tools.is_empty() => { + write_tool_declare(&mut out, local_tools, true)?; + if !content_is_empty(content) { + write_role_message(&mut out, "system", None, content)?; + } + } + ChatMessage::System { content } | ChatMessage::Developer { content, .. } => { + write_role_message(&mut out, "system", None, content)?; + } + ChatMessage::User { content } => { + write_role_message(&mut out, "user", None, content)?; + } + ChatMessage::Assistant { content } => { + tool_call_id_index.clear(); + let mut call_position = 0usize; + for block in content { + if let AssistantContentBlock::ToolCall(call) = block { + call_position += 1; + if !call.id.is_empty() { + tool_call_id_index + .entry(call.id.clone()) + .or_insert((call_position, call.name.clone())); + } + } + } + + write_assistant_message(&mut out, content, thinking)?; + } + ChatMessage::ToolResponse { .. } => unreachable!("handled above"), + } + } + flush_tool_run(&mut out, &mut pending_tool_run, &tool_call_id_index)?; + + match &request.tool_choice { + ChatToolChoice::Required => { + write_internal_system( + &mut out, + "tool-choice", + "The system is invoked with `tool_choice=required`.\n\ + You MUST call tools in the next message.", + )?; + } + // Emit only when tools are present: Rust defaults tool_choice to None + // for tool-free requests, which must not inject a tool-choice message. + ChatToolChoice::None if !request.tools.is_empty() => { + write_internal_system( + &mut out, + "tool-choice", + "The system is invoked with `tool_choice=none`.\n\ + You MUST NOT call any tools in the next message.", + )?; + } + ChatToolChoice::None | ChatToolChoice::Auto | ChatToolChoice::Function { .. } => {} + } + + write_response_format(&mut out, request)?; + + if request.chat_options.add_generation_prompt() { + write_open_tag(&mut out, "message", &[("role", "assistant")])?; + write_open_tag(&mut out, if thinking { "think" } else { "response" }, &[])?; + } + + Ok(out.finish()) +} + +fn request_tools(request: &ChatRequest) -> &[ChatTool] { + // Declare tools whenever the request carries them. K3 tool-declare is + // independent of tool_choice; tool_choice only injects control messages. + request.tools.as_slice() +} + +fn thinking_enabled(request: &ChatRequest) -> Result { + if let Some(thinking) = request.parse_template_bool("thinking")? { + return Ok(thinking); + } + if let Some(enable_thinking) = request.parse_template_bool("enable_thinking")? { + return Ok(enable_thinking); + } + Ok(request + .chat_options + .reasoning_effort + .map(|effort| effort != crate::request::ReasoningEffort::None) + .unwrap_or(true)) +} + +fn thinking_effort(request: &ChatRequest) -> Result { + let effort = if let Some(value) = request.chat_options.template_kwargs.get("thinking_effort") { + value.as_str().ok_or_else(|| { + Error::ChatTemplate(format!( + "template kwarg `thinking_effort` must be a string, got {value}" + )) + })? + } else if let Some(effort) = request.chat_options.reasoning_effort { + effort.as_str() + } else if let Some(value) = request.chat_options.template_kwargs.get("reasoning_effort") { + value.as_str().ok_or_else(|| { + Error::ChatTemplate(format!( + "template kwarg `reasoning_effort` must be a string, got {value}" + )) + })? + } else { + DEFAULT_THINKING_EFFORT + }; + + if !VALID_THINKING_EFFORTS.contains(&effort) { + return Err(Error::ChatTemplate(format!( + "unsupported thinking_effort={effort:?}; supported values are `low`, `high`, and `max`" + ))); + } + Ok(effort.to_string()) +} + +fn content_is_empty(content: &ChatContent) -> bool { + match content { + ChatContent::Text(text) => text.is_empty(), + ChatContent::Parts(parts) => parts.iter().all(|part| match part { + ChatContentPart::Text { text } => text.is_empty(), + ChatContentPart::ImageUrl { .. } + | ChatContentPart::VideoUrl { .. } + | ChatContentPart::InputAudio { .. } + | ChatContentPart::AudioUrl { .. } => false, + }), + } +} + +fn write_tool_declare( + out: &mut K3TokenWriter<'_>, + tools: &[ChatTool], + dynamic: bool, +) -> Result<()> { + let mut specs = Vec::with_capacity(tools.len()); + for tool in tools { + let mut function = Map::new(); + function.insert( + "description".to_string(), + Value::String(tool.description.clone().unwrap_or_default()), + ); + function.insert("name".to_string(), Value::String(tool.name.clone())); + function.insert("parameters".to_string(), sort_json(&tool.parameters)); + specs.push(json!({ + "function": Value::Object(function), + "type": "function", + })); + } + let payload = compact_json(&sort_json(&Value::Array(specs)))?; + + let body = if dynamic { + format!( + "## New Tools Available\n\ + The system dynamically extends the toolset via lazy-loading.\n\ + You have access to all existing and extended tools.\n\ + Here are the specs for the extended tools.\n\n\ + ```json\n\ + {payload}\n\ + ```" + ) + } else { + format!( + "# Tools\n\ + Here are the available tools, described in JSONSchema.\n\n\ + ```json\n\ + {payload}\n\ + ```" + ) + }; + + write_internal_system(out, "tool-declare", &body) +} + +fn write_internal_system( + out: &mut K3TokenWriter<'_>, + message_type: &str, + body: &str, +) -> Result<()> { + write_open_tag( + out, + "message", + &[("role", "system"), ("type", message_type)], + )?; + out.ordinary(body.trim())?; + write_close_tag(out, "message")?; + out.control(END_OF_MSG) +} + +fn write_role_message( + out: &mut K3TokenWriter<'_>, + role: &str, + name: Option<&str>, + content: &ChatContent, +) -> Result<()> { + let mut attrs = vec![("role", role.to_string())]; + if let Some(name) = name.filter(|name| !name.is_empty()) { + attrs.push(("name", name.to_string())); + } + let attr_refs: Vec<(&str, &str)> = attrs.iter().map(|(k, v)| (*k, v.as_str())).collect(); + write_open_tag(out, "message", &attr_refs)?; + write_content(out, content)?; + write_close_tag(out, "message")?; + out.control(END_OF_MSG) +} + +fn write_tool_message( + out: &mut K3TokenWriter<'_>, + tool_name: &str, + index: usize, + content: &ChatContent, +) -> Result<()> { + let index_str = index.to_string(); + write_open_tag( + out, + "message", + &[("role", "tool"), ("tool", tool_name), ("index", &index_str)], + )?; + write_content(out, content)?; + write_close_tag(out, "message")?; + out.control(END_OF_MSG) +} + +fn write_assistant_message( + out: &mut K3TokenWriter<'_>, + content: &[AssistantContentBlock], + thinking: bool, +) -> Result<()> { + write_open_tag(out, "message", &[("role", "assistant")])?; + + let mut reasoning = String::new(); + let mut response = String::new(); + let mut tool_calls = Vec::new(); + for block in content { + match block { + AssistantContentBlock::Reasoning { text } => reasoning.push_str(text), + AssistantContentBlock::Text { text } => response.push_str(text), + AssistantContentBlock::ToolCall(call) => tool_calls.push(call), + } + } + + // The think channel is structural: in thinking mode every assistant + // message carries open/close tags even when there is no reasoning content. + // In non-thinking mode the channel is dropped entirely. + if thinking { + write_open_tag(out, "think", &[])?; + if !reasoning.trim().is_empty() { + out.ordinary(&reasoning)?; + } + write_close_tag(out, "think")?; + } + + write_open_tag(out, "response", &[])?; + out.ordinary(&response)?; + write_close_tag(out, "response")?; + + if !tool_calls.is_empty() { + write_open_tag(out, "tools", &[])?; + for (index, tool_call) in tool_calls.into_iter().enumerate() { + write_assistant_tool_call(out, tool_call, index + 1)?; + } + write_close_tag(out, "tools")?; + } + + write_close_tag(out, "message")?; + out.control(END_OF_MSG) +} + +fn write_assistant_tool_call( + out: &mut K3TokenWriter<'_>, + tool_call: &AssistantToolCall, + index: usize, +) -> Result<()> { + let index_str = index.to_string(); + write_open_tag( + out, + "call", + &[ + ("tool", tool_call.name.as_str()), + ("index", index_str.as_str()), + ], + )?; + + let (args, json_block) = normalize_tool_arguments(&tool_call.arguments)?; + if let Some(raw) = json_block { + write_open_tag(out, "json", &[("type", "object")])?; + out.ordinary(&raw)?; + write_close_tag(out, "json")?; + } else { + for (key, value) in args { + let typ = xtml_type(&value); + write_open_tag(out, "argument", &[("key", key.as_str()), ("type", typ)])?; + out.ordinary(&xtml_value(&value))?; + write_close_tag(out, "argument")?; + } + } + + write_close_tag(out, "call") +} + +fn write_content(out: &mut K3TokenWriter<'_>, content: &ChatContent) -> Result<()> { + match content { + ChatContent::Text(text) => write_text_with_images(out, text), + ChatContent::Parts(parts) => { + for part in parts { + match part { + ChatContentPart::Text { text } => write_text_with_images(out, text)?, + ChatContentPart::ImageUrl { .. } => out.control(IMAGE_PLACEHOLDER)?, + ChatContentPart::VideoUrl { .. } => { + return Err(Error::UnsupportedMultimodalContent("video_url")); + } + ChatContentPart::InputAudio { .. } => { + return Err(Error::UnsupportedMultimodalContent("input_audio")); + } + ChatContentPart::AudioUrl { .. } => { + return Err(Error::UnsupportedMultimodalContent("audio_url")); + } + } + } + Ok(()) + } + } +} + +fn write_text_with_images(out: &mut K3TokenWriter<'_>, text: &str) -> Result<()> { + // Placeholder expansion is left as the literal K3 image token; multimodal + // preprocessing can replace it once image prompts are known. + out.ordinary(text) +} + +fn write_response_format(out: &mut K3TokenWriter<'_>, request: &ChatRequest) -> Result<()> { + let Some(rf) = request.chat_options.response_format.as_ref() else { + return Ok(()); + }; + + let rf_type = rf.get("type").and_then(Value::as_str).or_else(|| rf.as_str()).unwrap_or(""); + + match rf_type { + "json_object" => { + write_internal_system( + out, + "response-format", + "The system is invoked with `response_format=json_object`.\n\ + Your response must be raw JSON data without markdown code \ + blocks (```json) or any additional formatting.", + )?; + } + "json_schema" => { + let schema = extract_response_schema(rf); + let schema_json = compact_json(&sort_json(&schema.unwrap_or(Value::Null)))?; + write_internal_system( + out, + "response-format", + &format!( + "The system is invoked with `response_format=json_schema`.\n\ + Your response must be raw JSON data without markdown code \ + blocks (```json) or any additional formatting.\n\ + The JSON data must match the following schema:\n\ + ```json\n\ + {schema_json}\n\ + ```" + ), + )?; + } + _ => {} + } + Ok(()) +} + +fn extract_response_schema(response_format: &Value) -> Option { + let json_schema = response_format.get("json_schema")?; + if let Some(schema) = json_schema.get("schema") { + return Some(schema.clone()); + } + if let Some(schema) = json_schema.get("json_schema") { + return Some(schema.clone()); + } + Some(json_schema.clone()) +} + +fn normalize_tool_arguments(arguments: &str) -> Result<(Map, Option)> { + let trimmed = arguments.trim(); + if trimmed.is_empty() { + return Ok((Map::new(), None)); + } + match serde_json::from_str::(trimmed) { + Ok(Value::Object(map)) => Ok((map, None)), + Ok(_) => Err(Error::ChatTemplate( + "Kimi K3 tool call arguments must be a JSON object".to_string(), + )), + Err(_) => Ok((Map::new(), Some(arguments.to_string()))), + } +} + +fn xtml_type(value: &Value) -> &'static str { + match value { + Value::Bool(_) => "boolean", + Value::Null => "null", + Value::Number(_) => "number", + Value::String(_) => "string", + Value::Object(_) => "object", + Value::Array(_) => "array", + } +} + +fn xtml_value(value: &Value) -> String { + match value { + Value::String(text) => text.clone(), + other => compact_json(other).unwrap_or_else(|_| other.to_string()), + } +} + +fn write_open_tag(out: &mut K3TokenWriter<'_>, tag: &str, attrs: &[(&str, &str)]) -> Result<()> { + out.control(OPEN)?; + out.ordinary(tag)?; + for (key, value) in attrs { + out.ordinary(&format!(" {key}"))?; + out.ordinary("=\"")?; + out.ordinary(&escape_attr_value(value))?; + out.ordinary("\"")?; + } + out.control(SEP) +} + +fn write_close_tag(out: &mut K3TokenWriter<'_>, tag: &str) -> Result<()> { + out.control(CLOSE)?; + out.ordinary(tag)?; + out.control(SEP) +} + +fn escape_attr_value(value: &str) -> String { + value.replace('&', "&").replace('"', """) +} + +fn compact_json(value: &Value) -> Result { + serde_json::to_string(value).map_err(|error| Error::ChatTemplate(error.to_string())) +} + +fn sort_json(value: &Value) -> Value { + match value { + Value::Array(items) => Value::Array(items.iter().map(sort_json).collect()), + Value::Object(map) => { + let mut sorted = Map::new(); + let mut keys = map.keys().collect::>(); + keys.sort(); + for key in keys { + sorted.insert(key.clone(), sort_json(&map[key])); + } + Value::Object(sorted) + } + _ => value.clone(), + } +} diff --git a/rust/src/chat/src/renderer/kimi_k3/fixtures/controls_thinking_off_input.json b/rust/src/chat/src/renderer/kimi_k3/fixtures/controls_thinking_off_input.json new file mode 100644 index 0000000000..569201f3b9 --- /dev/null +++ b/rust/src/chat/src/renderer/kimi_k3/fixtures/controls_thinking_off_input.json @@ -0,0 +1,15 @@ +{ + "messages": [ + { + "role": "user", + "content": "json pls" + } + ], + "add_generation_prompt": true, + "response_format": { + "type": "json_object" + }, + "template_kwargs": { + "thinking": false + } +} diff --git a/rust/src/chat/src/renderer/kimi_k3/fixtures/controls_thinking_off_output.txt b/rust/src/chat/src/renderer/kimi_k3/fixtures/controls_thinking_off_output.txt new file mode 100644 index 0000000000..57ce4a6876 --- /dev/null +++ b/rust/src/chat/src/renderer/kimi_k3/fixtures/controls_thinking_off_output.txt @@ -0,0 +1,2 @@ +<|open|>message role="user"<|sep|>json pls<|close|>message<|sep|><|end_of_msg|><|open|>message role="system" type="response-format"<|sep|>The system is invoked with `response_format=json_object`. +Your response must be raw JSON data without markdown code blocks (```json) or any additional formatting.<|close|>message<|sep|><|end_of_msg|><|open|>message role="assistant"<|sep|><|open|>response<|sep|> \ No newline at end of file diff --git a/rust/src/chat/src/renderer/kimi_k3/fixtures/dynamic_system_tool_declare_input.json b/rust/src/chat/src/renderer/kimi_k3/fixtures/dynamic_system_tool_declare_input.json new file mode 100644 index 0000000000..61ecebc268 --- /dev/null +++ b/rust/src/chat/src/renderer/kimi_k3/fixtures/dynamic_system_tool_declare_input.json @@ -0,0 +1,62 @@ +{ + "tools": [ + { + "type": "function", + "function": { + "name": "get_weather", + "description": "weather", + "parameters": { + "type": "object", + "properties": { + "city": { + "type": "string" + }, + "days": { + "type": "number" + } + }, + "required": [ + "city" + ] + } + } + } + ], + "messages": [ + { + "role": "user", + "content": "hi" + }, + { + "role": "developer", + "tools": [ + { + "type": "function", + "function": { + "name": "calc", + "description": "calculator", + "parameters": { + "type": "object", + "properties": { + "x": { + "type": "number" + } + }, + "required": [ + "x" + ] + } + } + } + ] + }, + { + "role": "user", + "content": "use calc" + } + ], + "add_generation_prompt": true, + "template_kwargs": { + "thinking": true + } +} diff --git a/rust/src/chat/src/renderer/kimi_k3/fixtures/dynamic_system_tool_declare_output.txt b/rust/src/chat/src/renderer/kimi_k3/fixtures/dynamic_system_tool_declare_output.txt new file mode 100644 index 0000000000..edec2bebf2 --- /dev/null +++ b/rust/src/chat/src/renderer/kimi_k3/fixtures/dynamic_system_tool_declare_output.txt @@ -0,0 +1,14 @@ +<|open|>message role="system" type="tool-declare"<|sep|># Tools +Here are the available tools, described in JSONSchema. + +```json +[{"function":{"description":"weather","name":"get_weather","parameters":{"properties":{"city":{"type":"string"},"days":{"type":"number"}},"required":["city"],"type":"object"}},"type":"function"}] +```<|close|>message<|sep|><|end_of_msg|><|open|>message role="system" type="thinking-effort"<|sep|>`thinking_effort` guides on how much to think in your thinking channel (not including the response channel), supported values include `low`, `medium`, `high`, and `max`. +Now the system is invoked with `thinking_effort=max`.<|close|>message<|sep|><|end_of_msg|><|open|>message role="user"<|sep|>hi<|close|>message<|sep|><|end_of_msg|><|open|>message role="system" type="tool-declare"<|sep|>## New Tools Available +The system dynamically extends the toolset via lazy-loading. +You have access to all existing and extended tools. +Here are the specs for the extended tools. + +```json +[{"function":{"description":"calculator","name":"calc","parameters":{"properties":{"x":{"type":"number"}},"required":["x"],"type":"object"}},"type":"function"}] +```<|close|>message<|sep|><|end_of_msg|><|open|>message role="user"<|sep|>use calc<|close|>message<|sep|><|end_of_msg|><|open|>message role="assistant"<|sep|><|open|>think<|sep|> \ No newline at end of file diff --git a/rust/src/chat/src/renderer/kimi_k3/fixtures/history_preserve_and_image_input.json b/rust/src/chat/src/renderer/kimi_k3/fixtures/history_preserve_and_image_input.json new file mode 100644 index 0000000000..b7cf40c9b8 --- /dev/null +++ b/rust/src/chat/src/renderer/kimi_k3/fixtures/history_preserve_and_image_input.json @@ -0,0 +1,31 @@ +{ + "messages": [ + { + "role": "user", + "content": [ + { + "type": "text", + "text": "see this" + }, + { + "type": "image_url", + "image_url": "https://example.com/a.png" + } + ] + }, + { + "role": "assistant", + "content": "old answer", + "reasoning_content": "old reasoning" + }, + { + "role": "user", + "content": "continue" + } + ], + "add_generation_prompt": true, + "template_kwargs": { + "thinking": true, + "thinking_effort": "high" + } +} diff --git a/rust/src/chat/src/renderer/kimi_k3/fixtures/history_preserve_and_image_output.txt b/rust/src/chat/src/renderer/kimi_k3/fixtures/history_preserve_and_image_output.txt new file mode 100644 index 0000000000..03582f1180 --- /dev/null +++ b/rust/src/chat/src/renderer/kimi_k3/fixtures/history_preserve_and_image_output.txt @@ -0,0 +1,2 @@ +<|open|>message role="system" type="thinking-effort"<|sep|>`thinking_effort` guides on how much to think in your thinking channel (not including the response channel), supported values include `low`, `medium`, `high`, and `max`. +Now the system is invoked with `thinking_effort=high`.<|close|>message<|sep|><|end_of_msg|><|open|>message role="user"<|sep|>see this<|media_pad|><|close|>message<|sep|><|end_of_msg|><|open|>message role="assistant"<|sep|><|open|>think<|sep|>old reasoning<|close|>think<|sep|><|open|>response<|sep|>old answer<|close|>response<|sep|><|close|>message<|sep|><|end_of_msg|><|open|>message role="user"<|sep|>continue<|close|>message<|sep|><|end_of_msg|><|open|>message role="assistant"<|sep|><|open|>think<|sep|> \ No newline at end of file diff --git a/rust/src/chat/src/renderer/kimi_k3/fixtures/tools_history_and_required_input.json b/rust/src/chat/src/renderer/kimi_k3/fixtures/tools_history_and_required_input.json new file mode 100644 index 0000000000..3a64f8dd0b --- /dev/null +++ b/rust/src/chat/src/renderer/kimi_k3/fixtures/tools_history_and_required_input.json @@ -0,0 +1,87 @@ +{ + "tools": [ + { + "type": "function", + "function": { + "name": "get_weather", + "description": "weather", + "parameters": { + "type": "object", + "properties": { + "city": { + "type": "string" + }, + "days": { + "type": "number" + } + }, + "required": [ + "city" + ] + } + } + }, + { + "type": "function", + "function": { + "name": "calc", + "description": "calculator", + "parameters": { + "type": "object", + "properties": { + "x": { + "type": "number" + } + }, + "required": [ + "x" + ] + } + } + } + ], + "messages": [ + { + "role": "user", + "content": "Hangzhou weather and calc" + }, + { + "role": "assistant", + "content": "I'll check.", + "reasoning_content": "Need tools.", + "tool_calls": [ + { + "id": "get_weather:0", + "type": "function", + "function": { + "name": "get_weather", + "arguments": "{\"city\":\"Hangzhou\",\"days\":1}" + } + }, + { + "id": "calc:1", + "type": "function", + "function": { + "name": "calc", + "arguments": "{\"x\":2}" + } + } + ] + }, + { + "role": "tool", + "tool_call_id": "calc:1", + "content": "2" + }, + { + "role": "tool", + "tool_call_id": "get_weather:0", + "content": "{\"city\":\"Hangzhou\",\"condition\":\"rain\"}" + } + ], + "add_generation_prompt": true, + "tool_choice": "required", + "template_kwargs": { + "thinking": true + } +} diff --git a/rust/src/chat/src/renderer/kimi_k3/fixtures/tools_history_and_required_output.txt b/rust/src/chat/src/renderer/kimi_k3/fixtures/tools_history_and_required_output.txt new file mode 100644 index 0000000000..1c24810072 --- /dev/null +++ b/rust/src/chat/src/renderer/kimi_k3/fixtures/tools_history_and_required_output.txt @@ -0,0 +1,8 @@ +<|open|>message role="system" type="tool-declare"<|sep|># Tools +Here are the available tools, described in JSONSchema. + +```json +[{"function":{"description":"weather","name":"get_weather","parameters":{"properties":{"city":{"type":"string"},"days":{"type":"number"}},"required":["city"],"type":"object"}},"type":"function"},{"function":{"description":"calculator","name":"calc","parameters":{"properties":{"x":{"type":"number"}},"required":["x"],"type":"object"}},"type":"function"}] +```<|close|>message<|sep|><|end_of_msg|><|open|>message role="system" type="thinking-effort"<|sep|>`thinking_effort` guides on how much to think in your thinking channel (not including the response channel), supported values include `low`, `medium`, `high`, and `max`. +Now the system is invoked with `thinking_effort=max`.<|close|>message<|sep|><|end_of_msg|><|open|>message role="user"<|sep|>Hangzhou weather and calc<|close|>message<|sep|><|end_of_msg|><|open|>message role="assistant"<|sep|><|open|>think<|sep|>Need tools.<|close|>think<|sep|><|open|>response<|sep|>I'll check.<|close|>response<|sep|><|open|>tools<|sep|><|open|>call tool="get_weather" index="1"<|sep|><|open|>argument key="city" type="string"<|sep|>Hangzhou<|close|>argument<|sep|><|open|>argument key="days" type="number"<|sep|>1<|close|>argument<|sep|><|close|>call<|sep|><|open|>call tool="calc" index="2"<|sep|><|open|>argument key="x" type="number"<|sep|>2<|close|>argument<|sep|><|close|>call<|sep|><|close|>tools<|sep|><|close|>message<|sep|><|end_of_msg|><|open|>message role="tool" tool="get_weather" index="1"<|sep|>{"city":"Hangzhou","condition":"rain"}<|close|>message<|sep|><|end_of_msg|><|open|>message role="tool" tool="calc" index="2"<|sep|>2<|close|>message<|sep|><|end_of_msg|><|open|>message role="system" type="tool-choice"<|sep|>The system is invoked with `tool_choice=required`. +You MUST call tools in the next message.<|close|>message<|sep|><|end_of_msg|><|open|>message role="assistant"<|sep|><|open|>think<|sep|> \ No newline at end of file diff --git a/rust/src/chat/src/renderer/kimi_k3/mod.rs b/rust/src/chat/src/renderer/kimi_k3/mod.rs new file mode 100644 index 0000000000..8399e81522 --- /dev/null +++ b/rust/src/chat/src/renderer/kimi_k3/mod.rs @@ -0,0 +1,39 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +//! Native Kimi K3 XTML chat renderer. + +mod encoding; +#[cfg(test)] +mod tests; + +use aphrodite_text::Prompt; +use aphrodite_text::tokenizer::DynTokenizer; + +use super::{ChatRenderer, RenderedPrompt, request_template_kwargs}; +use crate::Result; +use crate::request::ChatRequest; + +/// Dedicated Kimi K3 XTML renderer. +#[derive(Clone)] +pub struct KimiK3ChatRenderer { + tokenizer: DynTokenizer, +} + +impl KimiK3ChatRenderer { + /// Create a Kimi K3 renderer. + pub fn new(tokenizer: DynTokenizer) -> Self { + Self { tokenizer } + } +} + +impl ChatRenderer for KimiK3ChatRenderer { + fn render(&self, request: &ChatRequest) -> Result { + request.validate()?; + + Ok(RenderedPrompt { + prompt: Prompt::TokenIds(encoding::render_request(request, self.tokenizer.as_ref())?), + effective_template_kwargs: request_template_kwargs(request), + }) + } +} diff --git a/rust/src/chat/src/renderer/kimi_k3/tests.rs b/rust/src/chat/src/renderer/kimi_k3/tests.rs new file mode 100644 index 0000000000..2df479accb --- /dev/null +++ b/rust/src/chat/src/renderer/kimi_k3/tests.rs @@ -0,0 +1,267 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +//! Golden fixtures generated from HF remote-code `encoding_k3.py`. + +use std::path::PathBuf; +use std::sync::Arc; + +use expect_test::{expect, expect_file}; +use serde_json::json; +use aphrodite_text::Prompt; +use aphrodite_text::tokenizer::DynTokenizer; +use aphrodite_tokenizer::Tokenizer; +use aphrodite_tokenizer::test_utils::TestTokenizer; + +use super::KimiK3ChatRenderer; +use crate::AssistantContentBlock; +use crate::ChatRenderer; +use crate::renderer::kimi_k3::encoding::{CLOSE, END_OF_MSG, IMAGE_PLACEHOLDER, OPEN, SEP}; +use crate::renderer::test_utils::{FixtureRequestOptions, fixture_chat_request}; +use crate::request::{ChatContentPart, ChatMessage, GenerationPromptMode, ReasoningEffort}; + +const OPEN_ID: u32 = 256; +const CLOSE_ID: u32 = 257; +const SEP_ID: u32 = 258; +const END_OF_MSG_ID: u32 = 259; +const MEDIA_ID: u32 = 260; + +fn test_tokenizer() -> TestTokenizer { + TestTokenizer::new() + .with_special_token(OPEN, OPEN_ID) + .with_special_token(CLOSE, CLOSE_ID) + .with_special_token(SEP, SEP_ID) + .with_special_token(END_OF_MSG, END_OF_MSG_ID) + .with_special_token(IMAGE_PLACEHOLDER, MEDIA_ID) +} + +fn render_token_ids(request: &crate::request::ChatRequest, tokenizer: DynTokenizer) -> Vec { + let prompt = KimiK3ChatRenderer::new(tokenizer).render(request).unwrap().prompt; + let Prompt::TokenIds(token_ids) = prompt else { + panic!("kimi k3 renderer should return token IDs") + }; + token_ids +} + +fn render_request(request: &crate::request::ChatRequest) -> String { + let tokenizer: DynTokenizer = Arc::new(test_tokenizer()); + let token_ids = render_token_ids(request, tokenizer.clone()); + tokenizer.decode(&token_ids, false).unwrap() +} + +fn fixture_path(name: &str) -> PathBuf { + PathBuf::from(env!("CARGO_MANIFEST_DIR")) + .join("src/renderer/kimi_k3/fixtures") + .join(name) +} + +fn kimi_k3_fixture_options() -> FixtureRequestOptions { + FixtureRequestOptions { + // Fixture JSON owns thinking via `template_kwargs`. + enable_thinking: None, + no_generation_prompt_when_last_assistant: false, + } +} + +fn assert_golden(name: &str) { + let input_name = format!("{name}_input.json"); + let request = fixture_chat_request(&fixture_path(&input_name), kimi_k3_fixture_options()); + let rendered = render_request(&request); + expect_file![format!("fixtures/{name}_output.txt")].assert_eq(&rendered); +} + +#[test] +fn golden_history_preserve_and_image() { + assert_golden("history_preserve_and_image"); +} + +#[test] +fn golden_tools_history_and_required() { + assert_golden("tools_history_and_required"); +} + +#[test] +fn golden_controls_thinking_off() { + assert_golden("controls_thinking_off"); +} + +#[test] +fn golden_dynamic_system_tool_declare() { + assert_golden("dynamic_system_tool_declare"); +} + +#[test] +fn token_writer_protects_literal_control_and_media_markers() { + let tokenizer = Arc::new(test_tokenizer()); + let user_text = format!("literal {OPEN} and {}", super::encoding::IMAGE_PLACEHOLDER); + let mut request = crate::request::ChatRequest::for_test(); + request.messages = vec![ChatMessage::user(vec![ + ChatContentPart::text(user_text), + ChatContentPart::image_url("data:image/png;base64,test"), + ])]; + request + .chat_options + .template_kwargs + .insert("thinking".to_string(), json!(false)); + request.chat_options.generation_prompt_mode = GenerationPromptMode::NoGenerationPrompt; + + let token_ids = render_token_ids(&request, tokenizer.clone()); + + assert_eq!( + token_ids.iter().filter(|&&token_id| token_id == OPEN_ID).count(), + 1 + ); + assert_eq!( + token_ids.iter().filter(|&&token_id| token_id == MEDIA_ID).count(), + 1 + ); + + let flattened = tokenizer.decode(&token_ids, false).unwrap(); + let flattened_ids = tokenizer.encode(&flattened, false).unwrap(); + assert_eq!( + flattened_ids.iter().filter(|&&token_id| token_id == OPEN_ID).count(), + 2 + ); + assert_eq!( + flattened_ids.iter().filter(|&&token_id| token_id == MEDIA_ID).count(), + 2 + ); +} + +#[test] +fn thinking_history_renders_empty_think_channel() { + let mut request = crate::request::ChatRequest::for_test(); + request.messages = vec![ + ChatMessage::user("question"), + ChatMessage::assistant_text("answer"), + ChatMessage::user("follow-up"), + ]; + request.chat_options.generation_prompt_mode = GenerationPromptMode::NoGenerationPrompt; + + let rendered = render_request(&request); + + assert!(rendered.contains( + "<|open|>message role=\"assistant\"<|sep|>\ + <|open|>think<|sep|><|close|>think<|sep|>\ + <|open|>response<|sep|>answer<|close|>response<|sep|>" + )); +} + +#[test] +fn non_thinking_history_omits_reasoning_channel() { + let mut request = crate::request::ChatRequest::for_test(); + request.messages = vec![ + ChatMessage::user("question"), + ChatMessage::assistant_blocks(vec![ + AssistantContentBlock::Reasoning { + text: "hidden reasoning".to_string(), + }, + AssistantContentBlock::Text { + text: "answer".to_string(), + }, + ]), + ChatMessage::user("follow-up"), + ]; + request + .chat_options + .template_kwargs + .insert("thinking".to_string(), json!(false)); + request.chat_options.generation_prompt_mode = GenerationPromptMode::NoGenerationPrompt; + + let rendered = render_request(&request); + + assert!(!rendered.contains("hidden reasoning")); + assert!(!rendered.contains("<|open|>think<|sep|>")); + assert!(rendered.contains( + "<|open|>message role=\"assistant\"<|sep|>\ + <|open|>response<|sep|>answer<|close|>response<|sep|>" + )); +} + +#[test] +fn defaults_thinking_effort_to_max() { + let rendered = render_request(&crate::request::ChatRequest::for_test()); + + expect![[r#"<|open|>message role="system" type="thinking-effort"<|sep|>`thinking_effort` guides on how much to think in your thinking channel (not including the response channel), supported values include `low`, `medium`, `high`, and `max`. +Now the system is invoked with `thinking_effort=max`.<|close|>message<|sep|><|end_of_msg|><|open|>message role="user"<|sep|>test<|close|>message<|sep|><|end_of_msg|><|open|>message role="assistant"<|sep|><|open|>think<|sep|>"#]] + .assert_eq(&rendered); +} + +#[test] +fn translates_standard_thinking_kwargs() { + let mut request = crate::request::ChatRequest::for_test(); + request + .chat_options + .template_kwargs + .insert("enable_thinking".to_string(), json!(true)); + request + .chat_options + .template_kwargs + .insert("reasoning_effort".to_string(), json!("high")); + + let rendered = render_request(&request); + + assert!(rendered.contains("thinking_effort=high")); + assert!(rendered.ends_with("<|open|>think<|sep|>")); +} + +#[test] +fn native_k3_kwargs_take_precedence() { + let mut request = crate::request::ChatRequest::for_test(); + request.chat_options.template_kwargs.extend([ + ("thinking".to_string(), json!(true)), + ("enable_thinking".to_string(), json!(false)), + ("thinking_effort".to_string(), json!("low")), + ("reasoning_effort".to_string(), json!("high")), + ]); + + let rendered = render_request(&request); + + assert!(rendered.contains("thinking_effort=low")); + assert!(!rendered.contains("thinking_effort=high")); +} + +#[test] +fn standard_none_disables_thinking() { + let mut request = crate::request::ChatRequest::for_test(); + request.chat_options.template_kwargs.extend([ + ("enable_thinking".to_string(), json!(false)), + ("reasoning_effort".to_string(), json!("none")), + ]); + + let rendered = render_request(&request); + + assert!(!rendered.contains("type=\"thinking-effort\"")); + assert!(rendered.ends_with("<|open|>response<|sep|>")); +} + +#[test] +fn typed_none_disables_thinking() { + let mut request = crate::request::ChatRequest::for_test(); + request.chat_options.reasoning_effort = Some(ReasoningEffort::None); + + let rendered = render_request(&request); + + assert!(!rendered.contains("type=\"thinking-effort\"")); + assert!(rendered.ends_with("<|open|>response<|sep|>")); +} + +#[test] +fn rejects_removed_medium_thinking_effort() { + let mut request = crate::request::ChatRequest::for_test(); + request + .chat_options + .template_kwargs + .insert("thinking_effort".to_string(), json!("medium")); + + let error = KimiK3ChatRenderer::new(Arc::new(test_tokenizer())) + .render(&request) + .unwrap_err(); + + expect![[r#" + ChatTemplate( + "unsupported thinking_effort=\"medium\"; supported values are `low`, `high`, and `max`", + ) + "#]] + .assert_debug_eq(&error); +} diff --git a/rust/src/chat/src/renderer/mod.rs b/rust/src/chat/src/renderer/mod.rs index fb9a15625e..5a177fe5b6 100644 --- a/rust/src/chat/src/renderer/mod.rs +++ b/rust/src/chat/src/renderer/mod.rs @@ -15,6 +15,7 @@ pub mod deepseek_v4; pub mod harmony; pub mod hf; mod inkling; +mod kimi_k3; mod selection; #[cfg(test)] mod test_utils; @@ -23,6 +24,7 @@ pub use deepseek_v4::DeepSeekV4ChatRenderer; pub use deepseek_v32::DeepSeekV32ChatRenderer; pub use harmony::HarmonyChatRenderer; pub use inkling::InklingChatRenderer; +pub use kimi_k3::KimiK3ChatRenderer; pub use selection::RendererSelection; /// Rendered chat prompt submitted to the text backend. diff --git a/rust/src/chat/src/renderer/selection.rs b/rust/src/chat/src/renderer/selection.rs index 711c0d87dd..5b6642c4c0 100644 --- a/rust/src/chat/src/renderer/selection.rs +++ b/rust/src/chat/src/renderer/selection.rs @@ -26,6 +26,8 @@ pub enum RendererSelection { Harmony, /// Force the Inkling native token renderer. Inkling, + /// Force the Kimi K3 XTML renderer. + KimiK3, } impl RendererSelection { @@ -37,6 +39,7 @@ impl RendererSelection { pub const HF_LITERAL: &str = "hf"; pub const INKLING_LITERAL: &str = "inkling"; pub const INKLING_MODEL_TYPE: &str = "inkling_mm_model"; + pub const KIMI_K3_LITERAL: &str = "kimi_k3"; /// Resolve the renderer selection using the given model type string, if /// it's `Auto`. @@ -47,6 +50,7 @@ impl RendererSelection { Self::DEEPSEEK_V4_LITERAL => Self::DeepSeekV4, Self::GPT_OSS_MODEL_TYPE => Self::Harmony, Self::INKLING_MODEL_TYPE => Self::Inkling, + Self::KIMI_K3_LITERAL => Self::KimiK3, _ => Self::Hf, }, selection => selection, @@ -70,6 +74,8 @@ impl FromStr for RendererSelection { Ok(Self::Harmony) } else if value.eq_ignore_ascii_case(Self::INKLING_LITERAL) { Ok(Self::Inkling) + } else if value.eq_ignore_ascii_case(Self::KIMI_K3_LITERAL) { + Ok(Self::KimiK3) } else { Err(format!( "unknown renderer `{value}` (expected one of: {})", @@ -88,6 +94,7 @@ impl fmt::Display for RendererSelection { Self::DeepSeekV4 => f.write_str(Self::DEEPSEEK_V4_LITERAL), Self::Harmony => f.write_str(Self::HARMONY_LITERAL), Self::Inkling => f.write_str(Self::INKLING_LITERAL), + Self::KimiK3 => f.write_str(Self::KIMI_K3_LITERAL), } } } @@ -114,7 +121,7 @@ mod tests { fn renderer_selection_expected_error_message() { let err = RendererSelection::from_str("unknown").unwrap_err(); expect_test::expect![ - "unknown renderer `unknown` (expected one of: auto, hf, deepseek_v32, deepseek_v4, harmony, inkling)" + "unknown renderer `unknown` (expected one of: auto, hf, deepseek_v32, deepseek_v4, harmony, inkling, kimi_k3)" ] .assert_eq(&err); } diff --git a/rust/src/chat/src/renderer/test_utils.rs b/rust/src/chat/src/renderer/test_utils.rs index 08eeca8176..0a4edc8137 100644 --- a/rust/src/chat/src/renderer/test_utils.rs +++ b/rust/src/chat/src/renderer/test_utils.rs @@ -1,6 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright contributors to the vLLM project +use std::collections::HashMap; use std::fs; use std::path::Path; @@ -16,8 +17,9 @@ use crate::request::{ /// Options for constructing a [`ChatRequest`] from a fixture file. #[derive(Debug, Clone, Copy)] pub(crate) struct FixtureRequestOptions { - /// Whether to set the template kwarg `[enable_]thinking=true`. - pub enable_thinking: bool, + /// Optional thinking toggle applied only when the fixture does not already + /// set `thinking` / `enable_thinking` in [`FixtureRequest::template_kwargs`]. + pub enable_thinking: Option, /// Whether fixtures ending in an assistant message should omit the /// trailing generation prompt. pub no_generation_prompt_when_last_assistant: bool, @@ -46,6 +48,15 @@ pub(crate) struct FixtureRequest { messages: Vec, add_generation_prompt: Option, reasoning_effort: Option, + /// Standard response format passed to model-specific renderers. + #[serde(default)] + response_format: Option, + /// Extra chat-template kwargs (thinking, preserve_thinking, …). + #[serde(default)] + template_kwargs: HashMap, + /// When omitted, defaults to `auto` if tools are present, otherwise `none`. + #[serde(default)] + tool_choice: Option, } impl FixtureFile { @@ -57,6 +68,9 @@ impl FixtureFile { messages, add_generation_prompt: None, reasoning_effort: None, + response_format: None, + template_kwargs: HashMap::new(), + tool_choice: None, }, } } @@ -69,6 +83,7 @@ pub(crate) enum FixtureMessage { content: FixtureContent, }, Developer { + #[serde(default)] content: FixtureContent, #[serde(default)] tools: Vec, @@ -98,6 +113,12 @@ pub(crate) enum FixtureContent { Parts(Vec), } +impl Default for FixtureContent { + fn default() -> Self { + Self::Text(String::new()) + } +} + #[derive(Debug, Deserialize)] #[serde(tag = "type", rename_all = "snake_case")] pub(crate) enum FixtureContentPart { @@ -136,6 +157,13 @@ struct FixtureToolCallFunction { impl FixtureRequest { fn into_chat_request(self, options: FixtureRequestOptions) -> ChatRequest { + let tools = to_chat_tools(&self.tools); + let tool_choice = self.tool_choice.unwrap_or(if tools.is_empty() { + ChatToolChoice::None + } else { + ChatToolChoice::Auto + }); + let mut request = ChatRequest { request_id: "renderer-fixture".to_string(), messages: self @@ -144,12 +172,8 @@ impl FixtureRequest { .enumerate() .map(|(index, message)| fixture_message_to_chat_message(index, message)) .collect(), - tools: to_chat_tools(&self.tools), - tool_choice: if self.tools.is_empty() { - ChatToolChoice::None - } else { - ChatToolChoice::Auto - }, + tools, + tool_choice, ..ChatRequest::for_test() }; @@ -162,9 +186,15 @@ impl FixtureRequest { request.chat_options.generation_prompt_mode = GenerationPromptMode::NoGenerationPrompt; } request.chat_options.reasoning_effort = self.reasoning_effort; - if options.enable_thinking { - for key in ["thinking", "enable_thinking"] { - request.chat_options.template_kwargs.insert(key.to_string(), Value::Bool(true)); + request.chat_options.response_format = self.response_format; + request.chat_options.template_kwargs.extend(self.template_kwargs); + + // Options supply a default thinking toggle only when the fixture did not. + if let Some(thinking) = options.enable_thinking { + let kwargs = &mut request.chat_options.template_kwargs; + if !kwargs.contains_key("thinking") && !kwargs.contains_key("enable_thinking") { + kwargs.insert("thinking".to_string(), Value::Bool(thinking)); + kwargs.insert("enable_thinking".to_string(), Value::Bool(thinking)); } } diff --git a/rust/src/chat/src/request.rs b/rust/src/chat/src/request.rs index 3aa412190e..42c1944d09 100644 --- a/rust/src/chat/src/request.rs +++ b/rust/src/chat/src/request.rs @@ -397,6 +397,10 @@ pub struct ChatOptions { /// Effort level exposed to chat templates for reasoning models. pub reasoning_effort: Option, + /// Standard response format available to model-specific renderers. + #[serde(default)] + pub response_format: Option, + /// Additional keyword arguments exposed to the chat template. pub template_kwargs: HashMap, } @@ -407,6 +411,7 @@ impl Default for ChatOptions { generation_prompt_mode: GenerationPromptMode::StartNewAssistant, chat_template: None, reasoning_effort: None, + response_format: None, template_kwargs: HashMap::new(), } } diff --git a/rust/src/chat/tests/roundtrip.rs b/rust/src/chat/tests/roundtrip.rs index 46cea6bd09..07b26f716f 100644 --- a/rust/src/chat/tests/roundtrip.rs +++ b/rust/src/chat/tests/roundtrip.rs @@ -212,6 +212,23 @@ impl RoundtripCase { } } + /// Kimi K3 XTML tool/reasoning channels (native renderer + unified parser). + /// + /// Needs HF tokenizer files under `HF_HOME` (`tiktoken.model` + + /// `tokenizer_config.json`). Weights are not required for this text-level + /// roundtrip. + fn kimi_k3() -> Self { + Self { + model_id: "moonshotai/Kimi-K3", + assistant_stop_suffix: "<|end_of_msg|>", + tool_call_parser: ParserSelection::Auto, + reasoning_parser: ParserSelection::Auto, + thinking_behavior: ThinkingBehavior::Toggleable { default: true }, + json_fmt: compact_json_fmt(), + sort_json_keys: false, + } + } + /// SeedOSS with `` / `` reasoning tags. fn seed_oss() -> Self { Self { @@ -255,7 +272,7 @@ impl RoundtripCase { fn gpt_oss() -> Self { Self { model_id: "openai/gpt-oss-20b", - assistant_stop_suffix: "", // not applicable for token-id cases + assistant_stop_suffix: "", tool_call_parser: ParserSelection::Auto, reasoning_parser: ParserSelection::Auto, thinking_behavior: ThinkingBehavior::Always { value: true }, @@ -268,7 +285,7 @@ impl RoundtripCase { fn inkling() -> Self { Self { model_id: "thinkingmachines/Inkling", - assistant_stop_suffix: "", + assistant_stop_suffix: "<|content_model_end_sampling|>", tool_call_parser: ParserSelection::Auto, reasoning_parser: ParserSelection::Auto, thinking_behavior: ThinkingBehavior::Always { value: true }, @@ -312,6 +329,8 @@ roundtrip_tests! { nemotron_v3 => [reasoning_and_content], gemma4 => [tool_call_mix], // Gemma4 strips reasoning in history if there's no tool call kimi_k25 => [tool_call_mix], // Kimi K2.5 strips reasoning in history + // K3 drops plain-assistant reasoning in history; tool-call turns keep it. + kimi_k3 => [tool_call_mix], gpt_oss => [tool_call_mix], // Harmony strips reasoning in history if there's no tool call inkling => [reasoning_and_content, tool_call_mix], } @@ -670,14 +689,23 @@ fn decoded_completion_stream( .collect() } Prompt::TokenIds(token_ids) => { - ensure!( - assistant_stop_suffix.is_empty(), - "token-id roundtrip cases do not support text stop suffixes" - ); + let body = if assistant_stop_suffix.is_empty() { + token_ids.as_slice() + } else { + let stop_token_ids = tokenizer + .encode(assistant_stop_suffix, false) + .context("failed to encode token-id completion stop suffix")?; + token_ids.strip_suffix(stop_token_ids.as_slice()).with_context(|| { + format!( + "token-id completion did not end with {:?}: {:?}", + assistant_stop_suffix, token_ids + ) + })? + }; incremental_decode_chunks( tokenizer, &prompt_token_ids, - token_ids, + body, TOKEN_COMPLETION_CHUNK_TOKENS, )? } diff --git a/rust/src/cmd/Cargo.toml b/rust/src/cmd/Cargo.toml index 7fbdcce063..02270cf7a5 100644 --- a/rust/src/cmd/Cargo.toml +++ b/rust/src/cmd/Cargo.toml @@ -23,17 +23,16 @@ serde.workspace = true serde_json.workspace = true serde_with.workspace = true thiserror-ext.workspace = true -time.workspace = true tokio = { workspace = true, features = ["signal"] } tokio-util.workspace = true tracing.workspace = true -tracing-subscriber.workspace = true uuid.workspace = true aphrodite-bench.workspace = true aphrodite-chat.workspace = true aphrodite-engine-core-client.workspace = true aphrodite-managed-engine.workspace = true aphrodite-server.workspace = true +aphrodite-tracing.workspace = true [dev-dependencies] expect-test.workspace = true diff --git a/rust/src/cmd/src/cli.rs b/rust/src/cmd/src/cli.rs index b6253affd5..222cf3fe50 100644 --- a/rust/src/cmd/src/cli.rs +++ b/rust/src/cmd/src/cli.rs @@ -23,6 +23,7 @@ use serde_with::{DefaultOnNull, OneOrMany, serde_as}; use thiserror_ext::AsReport as _; use uuid::Uuid; use aphrodite_chat::ReasoningParserFactory; +use aphrodite_chat::multimodal::MmLimitPerPrompt; use aphrodite_engine_core_client::TransportMode; use aphrodite_managed_engine::ManagedEngineConfig; use aphrodite_managed_engine::cli::{ManagedEngineArgs, repartition_managed_engine_args}; @@ -186,6 +187,17 @@ pub struct SharedRuntimeArgs { #[serde(default)] pub default_chat_template_kwargs: Option>, + /// The maximum number of input items allowed per prompt for each + /// modality, as a JSON object (e.g. `{"image": 16, "video": 2}`). + /// + /// Also accepts the engine's configurable form + /// (e.g. `{"video": {"count": 1, "num_frames": 32}}`); the extra + /// profiling options are forwarded to the engine untouched. + /// Unspecified modalities are unlimited. + #[arg(long, value_parser = parse_json::, value_name = "JSON", default_value = "{}")] + #[serde(default)] + pub limit_mm_per_prompt: MmLimitPerPrompt, + /// The format to render message content within a chat template. /// /// * "auto" detects the format from the template @@ -348,6 +360,18 @@ impl SharedRuntimeArgs { .expect("profiler config serialization should not fail") } + /// Return the per-modality limits as JSON for managed Python engine + /// forwarding, or `None` when nothing is configured. + /// + /// Round-tripping the parsed map rather than the raw argument keeps the + /// engine's own profiling options (`num_frames`, `width`, ...) intact. + pub fn limit_mm_per_prompt_json(&self) -> Option { + (!self.limit_mm_per_prompt.is_empty()).then(|| { + serde_json::to_string(&self.limit_mm_per_prompt) + .expect("limit-mm-per-prompt serialization should not fail") + }) + } + /// Apply fallback logic for API key configuration from env variables. fn apply_env_api_key_fallback(&mut self) { if self.api_key.is_empty() @@ -399,6 +423,7 @@ impl SharedRuntimeArgs { language_model_only: self.language_model_only, chat_template: self.chat_template, default_chat_template_kwargs: self.default_chat_template_kwargs, + limit_mm_per_prompt: self.limit_mm_per_prompt, chat_template_content_format: self.chat_template_content_format, max_logprobs: self.max_logprobs, api_server_options, @@ -451,6 +476,7 @@ impl SharedRuntimeArgs { language_model_only: self.language_model_only, chat_template: self.chat_template, default_chat_template_kwargs: self.default_chat_template_kwargs, + limit_mm_per_prompt: self.limit_mm_per_prompt, chat_template_content_format: self.chat_template_content_format, max_logprobs: self.max_logprobs, api_server_options, @@ -666,6 +692,7 @@ impl ServeArgs { self.runtime.disable_log_stats, self.runtime.shutdown_timeout, handshake_port, + self.runtime.limit_mm_per_prompt_json(), ) } } diff --git a/rust/src/cmd/src/cli/tests.rs b/rust/src/cmd/src/cli/tests.rs index 363d158707..d9ed32fe0a 100644 --- a/rust/src/cmd/src/cli/tests.rs +++ b/rust/src/cmd/src/cli/tests.rs @@ -64,6 +64,7 @@ fn serve_args_forward_python_flags_with_separator() { http_timeout_keep_alive: None, chat_template: None, default_chat_template_kwargs: None, + limit_mm_per_prompt: {}, chat_template_content_format: Auto, enable_log_requests: false, enable_prompt_tokens_details: false, @@ -666,7 +667,7 @@ fn serve_args_reject_unknown_renderer_value() { .unwrap_err(); expect![[r#" - error: invalid value 'definitely_missing' for '--tokenizer-mode ': unknown renderer `definitely_missing` (expected one of: auto, hf, deepseek_v32, deepseek_v4, harmony, inkling) + error: invalid value 'definitely_missing' for '--tokenizer-mode ': unknown renderer `definitely_missing` (expected one of: auto, hf, deepseek_v32, deepseek_v4, harmony, inkling, kimi_k3) For more information, try '--help'. "#]] @@ -762,6 +763,7 @@ fn frontend_args_accept_json() { http_timeout_keep_alive: None, chat_template: None, default_chat_template_kwargs: None, + limit_mm_per_prompt: {}, chat_template_content_format: Auto, enable_log_requests: false, enable_prompt_tokens_details: false, @@ -1117,6 +1119,24 @@ fn frontend_args_json_rejects_malformed_json() { "#]].assert_eq(&error.to_string()); } +#[test] +fn serve_args_reject_unsupported_modality_in_limit_mm_per_prompt() { + let error = Cli::try_parse_from([ + "vllm-rs", + "serve", + "Qwen/Qwen3-0.6B", + "--limit-mm-per-prompt", + r#"{"unsupported_modality": 1}"#, + ]) + .unwrap_err(); + + expect![[r#" + error: invalid value '{"unsupported_modality": 1}' for '--limit-mm-per-prompt ': invalid JSON object: unknown variant `unsupported_modality`, expected one of `image`, `audio`, `video` at line 1 column 23 + + For more information, try '--help'. + "#]].assert_eq(&error.to_string()); +} + #[test] fn serve_args_reject_flags_before_model() { let error = Cli::try_parse_from(["aphrodite-rs", "serve", "--python", "python3", "Qwen/Qwen3-0.6B"]) @@ -1331,6 +1351,7 @@ fn serve_args_accept_handshake_aliases() { http_timeout_keep_alive: None, chat_template: None, default_chat_template_kwargs: None, + limit_mm_per_prompt: {}, chat_template_content_format: Auto, enable_log_requests: false, enable_prompt_tokens_details: false, @@ -1474,6 +1495,7 @@ fn serve_frontend_config_uses_dp_address_as_advertised_host() { language_model_only: false, chat_template: None, default_chat_template_kwargs: None, + limit_mm_per_prompt: {}, chat_template_content_format: Auto, max_logprobs: None, api_server_options: ApiServerOptions { @@ -1558,6 +1580,7 @@ fn serve_frontend_config_keeps_tcp_transport_for_non_local_only_topology() { language_model_only: false, chat_template: None, default_chat_template_kwargs: None, + limit_mm_per_prompt: {}, chat_template_content_format: Auto, max_logprobs: None, api_server_options: ApiServerOptions { @@ -1660,6 +1683,7 @@ fn frontend_config_uses_external_coordinator_when_coordinator_address_is_present language_model_only: false, chat_template: None, default_chat_template_kwargs: None, + limit_mm_per_prompt: {}, chat_template_content_format: Auto, max_logprobs: None, api_server_options: ApiServerOptions { diff --git a/rust/src/cmd/src/cli/unsupported.rs b/rust/src/cmd/src/cli/unsupported.rs index 7a3ebac50b..ea4b178b12 100644 --- a/rust/src/cmd/src/cli/unsupported.rs +++ b/rust/src/cmd/src/cli/unsupported.rs @@ -299,11 +299,6 @@ pub struct EngineUnsupportedArgs { )] pub kv_sharing_fast_prefill: Option, - /// The maximum number of input items and options allowed per - /// prompt for each modality. - #[arg(long)] - pub limit_mm_per_prompt: Option, - /// Additional args passed to process media inputs, keyed by modalities. #[arg(long)] pub media_io_kwargs: Option, diff --git a/rust/src/cmd/src/main.rs b/rust/src/cmd/src/main.rs index c5ee8e0dc2..bfffa9123a 100644 --- a/rust/src/cmd/src/main.rs +++ b/rust/src/cmd/src/main.rs @@ -2,7 +2,6 @@ // SPDX-FileCopyrightText: Copyright contributors to the vLLM project mod cli; -mod logging; use std::env; use std::ffi::OsStr; @@ -89,7 +88,7 @@ fn main() -> Result<()> { "serve" | "frontend" => "RustFrontend", _ => "Rust", }; - logging::init_tracing(process_label); + aphrodite_tracing::init_tracing(process_label); let cli = Cli::parse(); diff --git a/rust/src/managed-engine/src/cli.rs b/rust/src/managed-engine/src/cli.rs index ef1da2cce5..05dfebd3e4 100644 --- a/rust/src/managed-engine/src/cli.rs +++ b/rust/src/managed-engine/src/cli.rs @@ -91,6 +91,7 @@ impl ManagedEngineArgs { disable_log_stats: bool, shutdown_timeout: u64, handshake_port: u16, + limit_mm_per_prompt: Option, ) -> ManagedEngineConfig { let mut python_args = self.python_args; // Manually forward some args to the Python engine. @@ -126,6 +127,10 @@ impl ManagedEngineArgs { python_args.push("--data-parallel-size-local".to_string()); python_args.push(data_parallel_size_local.to_string()); } + if let Some(limit_mm_per_prompt) = limit_mm_per_prompt { + python_args.push("--limit-mm-per-prompt".to_string()); + python_args.push(limit_mm_per_prompt); + } ManagedEngineConfig { python: self.python, diff --git a/rust/src/parser/src/unified/kimi_k3.rs b/rust/src/parser/src/unified/kimi_k3.rs new file mode 100644 index 0000000000..98cf1373d3 --- /dev/null +++ b/rust/src/parser/src/unified/kimi_k3.rs @@ -0,0 +1,1149 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +//! Unified parser for the Kimi K3 (XTML) chat format. +//! +//! Original Python implementations: +//! - `vllm/reasoning/kimi_k3_reasoning_parser.py` +//! - `vllm/tool_parsers/kimi_k3_tool_parser.py` +//! +//! K3 wraps one assistant message into XTML channels built from the dedicated +//! special tokens `<|open|>`, `<|close|>`, and `<|sep|>`: +//! +//! ```text +//! <|open|>think<|sep|>reasoning<|close|>think<|sep|> +//! <|open|>response<|sep|>visible answer<|close|>response<|sep|> +//! <|open|>tools<|sep|> +//! <|open|>call tool="get_weather" index="1"<|sep|> +//! <|open|>argument key="city" type="string"<|sep|>Hangzhou<|close|>argument<|sep|> +//! <|close|>call<|sep|> +//! <|close|>tools<|sep|> +//! <|close|>message<|sep|> +//! ``` +//! +//! In chat serving the generation prompt ends with `<|open|>think<|sep|>` +//! (thinking) or `<|open|>response<|sep|>` (instruct), so the model output +//! starts *inside* that channel without re-emitting the open tag; +//! [`UnifiedParser::initialize`] detects this from the prompt token IDs. +//! +//! Argument decoding mirrors the renderer's type tagging (inverse encoding): +//! `type="string"` values pass the raw text through, other types are +//! JSON-decoded, and a raw `json` block is passed through unmodified. Attribute +//! values reverse the renderer escaping (`"` before `&`). +//! +//! Known limitation (shared with the Python parser): string argument and +//! response bodies are emitted raw, so a value that literally contains +//! `<|close|>argument<|sep|>` or `<|close|>response<|sep|>` is +//! indistinguishable from a real closing marker. + +mod structural_tag; + +pub use structural_tag::KimiK3StructuralTagBuilder; + +use serde_json::{Map, Value}; +use aphrodite_tokenizer::DynTokenizer; +use winnow::ascii::{multispace0 as ws0, multispace1 as ws1}; +use winnow::combinator::{alt, delimited, eof, preceded, repeat, seq, terminated}; +use winnow::error::{ContextError, ErrMode, ModalResult, StrContext}; +use winnow::prelude::*; +use winnow::stream::Partial; +use winnow::token::{literal, rest, take_till, take_until, take_while}; + +use self::structural_tag::KIMI_K3_STRUCTURAL_TAG_BUILDER; +use super::{Result, UnifiedParser, UnifiedParserOutput, token_id}; +use crate::tool::{StructuralTagBuilder, Tool, ToolCallDelta}; +use crate::unified::parsing_failed; +use crate::utils::{MarkerScanState, parse_buffered_event, safe_text_len_mul, take_until_marker}; + +const OPEN: &str = "<|open|>"; +const SEP: &str = "<|sep|>"; +const END_OF_MSG: &str = "<|end_of_msg|>"; + +const THINK_OPEN: &str = "<|open|>think<|sep|>"; +const THINK_CLOSE: &str = "<|close|>think<|sep|>"; +const RESPONSE_OPEN: &str = "<|open|>response<|sep|>"; +const RESPONSE_CLOSE: &str = "<|close|>response<|sep|>"; +const TOOLS_OPEN: &str = "<|open|>tools<|sep|>"; +const TOOLS_CLOSE: &str = "<|close|>tools<|sep|>"; +const MESSAGE_CLOSE: &str = "<|close|>message<|sep|>"; +const CALL_OPEN: &str = "<|open|>call"; +const CALL_CLOSE: &str = "<|close|>call<|sep|>"; +const ARG_OPEN: &str = "<|open|>argument"; +const ARG_CLOSE: &str = "<|close|>argument<|sep|>"; +const JSON_OPEN: &str = "<|open|>json"; +const JSON_CLOSE: &str = "<|close|>json<|sep|>"; + +const IDLE_MARKERS: &[&str] = &[ + THINK_OPEN, + RESPONSE_OPEN, + TOOLS_OPEN, + MESSAGE_CLOSE, + END_OF_MSG, +]; +const REASONING_MARKERS: &[&str] = &[THINK_CLOSE, END_OF_MSG]; +const RESPONSE_MARKERS: &[&str] = &[RESPONSE_CLOSE, TOOLS_OPEN, MESSAGE_CLOSE, END_OF_MSG]; +const EPILOGUE_MARKERS: &[&str] = &[TOOLS_OPEN, MESSAGE_CLOSE, END_OF_MSG]; +const TOOLS_MARKERS: &[&str] = &[CALL_OPEN, TOOLS_CLOSE, MESSAGE_CLOSE, END_OF_MSG]; + +/// Channel tags are a couple of text tokens; longer `<|open|>…<|sep|>` spans in +/// the prompt tail (attribute-bearing message opens, message bodies) never name +/// a generation channel. +const MAX_PREFILL_TAG_TOKENS: usize = 8; + +type KimiK3Input<'i> = Partial<&'i str>; + +#[derive(Debug, Clone, PartialEq, Eq)] +enum KimiK3Event { + Text { + len: usize, + }, + Reasoning { + len: usize, + }, + /// Structural noise consumed without emitting anything. + Skip, + ThinkOpen, + ThinkClose, + ResponseOpen, + ResponseClose, + ToolsOpen, + ToolsClose, + /// The assistant message closed; everything after it is ignored. + MessageEnd, + CallOpen { + name: String, + index: Option, + }, + CallComplete { + arguments: String, + }, +} + +#[derive(Debug, Clone, Default, PartialEq, Eq)] +enum KimiK3Mode { + /// Before any channel opens: channel opens are expected, but raw text + /// falls through as visible text so marker-free output still streams. + #[default] + Idle, + /// Inside the `think` channel. + Reasoning, + /// Inside the `response` channel. + Response, + /// After the `response` (or `tools`) channel closed: only a `tools` + /// channel or the message close may follow, and noise is never content. + Epilogue, + /// Inside the `tools` channel, between `call` blocks. + Tools, + /// Inside one `call` block, buffering its body until the close marker. + Call { + name: String, + index: Option, + scan: MarkerScanState, + }, + /// After the message closed: ignore the rest (EOS leakage guard). + Done, +} + +/// Unified parser for Kimi K3 XTML think / response / tools channels. +pub struct KimiK3UnifiedParser { + buffer: String, + mode: KimiK3Mode, + /// Parser-provided tool-call IDs (`{tool}:{zero_based_index}`) by tool + /// index; its length is also the count of emitted calls. + call_ids: Vec, + tokenizer: DynTokenizer, + open_token_id: u32, + sep_token_id: u32, +} + +impl KimiK3UnifiedParser { + /// Create a Kimi K3 parser. + pub fn new(_tools: &[Tool], tokenizer: DynTokenizer) -> Result { + let open_token_id = token_id(tokenizer.as_ref(), OPEN)?; + let sep_token_id = token_id(tokenizer.as_ref(), SEP)?; + + Ok(Self { + buffer: String::new(), + mode: KimiK3Mode::default(), + call_ids: Vec::new(), + tokenizer, + open_token_id, + sep_token_id, + }) + } + + /// Detect the prefilled generation channel from the prompt tail. + /// + /// `add_generation_prompt` ends the prompt with `<|open|>think<|sep|>` or + /// `<|open|>response<|sep|>`, so generation starts inside that channel and + /// never re-emits the open tag. Locate the last `<|open|>…<|sep|>` pair and + /// decode the short tag between the two structural tokens. + fn initialize_mode(&mut self, prompt_token_ids: &[u32]) { + self.mode = KimiK3Mode::Idle; + + let Some(sep_pos) = prompt_token_ids.iter().rposition(|&id| id == self.sep_token_id) else { + return; + }; + let Some(open_pos) = + prompt_token_ids[..sep_pos].iter().rposition(|&id| id == self.open_token_id) + else { + return; + }; + let tag_ids = &prompt_token_ids[open_pos + 1..sep_pos]; + if tag_ids.is_empty() || tag_ids.len() > MAX_PREFILL_TAG_TOKENS { + return; + } + let Ok(tag) = self.tokenizer.decode(tag_ids, /* skip_special_tokens */ false) else { + return; + }; + self.mode = match tag.trim() { + "think" => KimiK3Mode::Reasoning, + "response" => KimiK3Mode::Response, + _ => KimiK3Mode::Idle, + }; + } + + fn apply_event(&mut self, event: KimiK3Event, output: &mut UnifiedParserOutput) -> Result<()> { + match event { + KimiK3Event::Text { len } => output.push_text(self.buffer[..len].to_string()), + KimiK3Event::Reasoning { len } => { + output.push_reasoning(self.buffer[..len].to_string()); + } + KimiK3Event::Skip => {} + KimiK3Event::ThinkOpen => self.mode = KimiK3Mode::Reasoning, + KimiK3Event::ThinkClose => self.mode = KimiK3Mode::Idle, + KimiK3Event::ResponseOpen => self.mode = KimiK3Mode::Response, + KimiK3Event::ResponseClose => self.mode = KimiK3Mode::Epilogue, + KimiK3Event::ToolsOpen => self.mode = KimiK3Mode::Tools, + KimiK3Event::ToolsClose => self.mode = KimiK3Mode::Epilogue, + KimiK3Event::MessageEnd => self.mode = KimiK3Mode::Done, + KimiK3Event::CallOpen { name, index } => { + self.mode = KimiK3Mode::Call { + name, + index, + scan: MarkerScanState::default(), + }; + } + KimiK3Event::CallComplete { arguments } => { + let mode = std::mem::replace(&mut self.mode, KimiK3Mode::Tools); + let KimiK3Mode::Call { name, index, .. } = mode else { + return Err(parsing_failed!( + "Kimi K3 call completion without an active tool call" + )); + }; + // An empty/garbage call block without a tool name is dropped, + // matching the Python parser. + if name.is_empty() { + return Ok(()); + } + + let tool_index = self.call_ids.len(); + self.call_ids.push(tool_call_id_for(&name, index.as_deref())); + output.push_call(ToolCallDelta { + tool_index, + name: Some(name), + arguments, + }); + } + } + Ok(()) + } + + fn reset_state(&mut self) -> String { + self.mode = KimiK3Mode::Idle; + self.call_ids.clear(); + std::mem::take(&mut self.buffer) + } +} + +impl UnifiedParser for KimiK3UnifiedParser { + fn create(tools: &[Tool], tokenizer: DynTokenizer) -> Result> + where + Self: Sized + 'static, + { + Self::new(tools, tokenizer).map(|parser| Box::new(parser) as Box) + } + + fn initialize(&mut self, prompt_token_ids: &[u32]) -> Result<()> { + self.buffer.clear(); + self.call_ids.clear(); + self.initialize_mode(prompt_token_ids); + Ok(()) + } + + fn preserve_special_tokens(&self) -> bool { + true + } + + fn structural_tag_builder(&self) -> Option<&dyn StructuralTagBuilder> { + Some(&KIMI_K3_STRUCTURAL_TAG_BUILDER) + } + + fn tool_call_id(&self, tool_index: usize) -> Option<&str> { + self.call_ids.get(tool_index).map(String::as_str) + } + + fn parse_into(&mut self, chunk: &str, output: &mut UnifiedParserOutput) -> Result<()> { + self.buffer.push_str(chunk); + + while let Some((event, consumed_len)) = parse_buffered_event(&self.buffer, |input| { + parse_next_kimi_k3_event(input, &mut self.mode) + })? { + self.apply_event(event, output)?; + self.buffer.drain(..consumed_len); + } + + Ok(()) + } + + fn finish(&mut self) -> Result { + let mut output = UnifiedParserOutput::default(); + + match &self.mode { + KimiK3Mode::Idle | KimiK3Mode::Response => { + output.push_text(std::mem::take(&mut self.buffer)); + } + KimiK3Mode::Reasoning => output.push_reasoning(std::mem::take(&mut self.buffer)), + KimiK3Mode::Epilogue | KimiK3Mode::Done => self.buffer.clear(), + // A tools channel truncated between complete calls loses only its + // closing markers; keep the calls already emitted. + KimiK3Mode::Tools if self.buffer.is_empty() => {} + KimiK3Mode::Tools | KimiK3Mode::Call { .. } => { + return Err(parsing_failed!("incomplete Kimi K3 tool call")); + } + } + + // Keep call_ids so tool_call_id() stays available after the stream ends. + self.mode = KimiK3Mode::Idle; + Ok(output) + } + + fn reset(&mut self) -> String { + self.reset_state() + } +} + +/// Build the API-side tool-call ID from the XTML one-based `index` attribute. +/// +/// The ID uses the zero-based call ordinal; XTML's message index stays +/// one-based when rendering tool result messages. +fn tool_call_id_for(name: &str, index: Option<&str>) -> String { + match index { + None => name.to_string(), + Some(raw) => match raw.parse::() { + Ok(one_based) => format!("{name}:{}", one_based - 1), + Err(_) => format!("{name}:{raw}"), + }, + } +} + +/// Parse one Kimi K3 event from buffered streaming input. +fn parse_next_kimi_k3_event( + input: &mut KimiK3Input<'_>, + mode: &mut KimiK3Mode, +) -> ModalResult { + match mode { + KimiK3Mode::Idle => parse_idle_event(input), + KimiK3Mode::Reasoning => parse_reasoning_event(input), + KimiK3Mode::Response => parse_response_event(input), + KimiK3Mode::Epilogue => parse_epilogue_event(input), + KimiK3Mode::Tools => parse_tools_event(input), + KimiK3Mode::Call { scan, .. } => call_body_event(input, scan), + KimiK3Mode::Done => parse_done_event(input), + } +} + +/// Parse an event while waiting for the next channel open. +fn parse_idle_event(input: &mut KimiK3Input<'_>) -> ModalResult { + alt(( + literal(THINK_OPEN).value(KimiK3Event::ThinkOpen), + literal(RESPONSE_OPEN).value(KimiK3Event::ResponseOpen), + literal(TOOLS_OPEN).value(KimiK3Event::ToolsOpen), + message_end_event, + safe_idle_text_event, + )) + .parse_next(input) +} + +/// Parse an event inside the `think` channel. +fn parse_reasoning_event(input: &mut KimiK3Input<'_>) -> ModalResult { + alt(( + literal(THINK_CLOSE).value(KimiK3Event::ThinkClose), + // `<|end_of_msg|>` can reach the parser under `ignore_eos` or + // `include_stop_str_in_output`; never leak it into reasoning. + literal(END_OF_MSG).value(KimiK3Event::MessageEnd), + safe_reasoning_event, + )) + .parse_next(input) +} + +/// Parse an event inside the `response` channel. +fn parse_response_event(input: &mut KimiK3Input<'_>) -> ModalResult { + alt(( + literal(RESPONSE_CLOSE).value(KimiK3Event::ResponseClose), + // The response body also implicitly ends at a `tools` channel. + literal(TOOLS_OPEN).value(KimiK3Event::ToolsOpen), + message_end_event, + safe_response_text_event, + )) + .parse_next(input) +} + +/// Parse an event after the response channel closed. +fn parse_epilogue_event(input: &mut KimiK3Input<'_>) -> ModalResult { + alt(( + literal(TOOLS_OPEN).value(KimiK3Event::ToolsOpen), + message_end_event, + skip_epilogue_noise_event, + )) + .parse_next(input) +} + +/// Parse an event inside the `tools` channel, between `call` blocks. +fn parse_tools_event(input: &mut KimiK3Input<'_>) -> ModalResult { + alt(( + call_open_event, + literal(TOOLS_CLOSE).value(KimiK3Event::ToolsClose), + // Defensive: an unterminated tools channel still ends with the message. + message_end_event, + skip_tools_noise_event, + )) + .parse_next(input) +} + +/// Parse a message close or end-of-message marker. +fn message_end_event(input: &mut KimiK3Input<'_>) -> ModalResult { + alt((literal(MESSAGE_CLOSE), literal(END_OF_MSG))) + .value(KimiK3Event::MessageEnd) + .parse_next(input) +} + +/// Ignore everything after the assistant message closed. +fn parse_done_event(input: &mut KimiK3Input<'_>) -> ModalResult { + rest.value(KimiK3Event::Skip).parse_next(input) +} + +/// Parse safe text while waiting for the next channel marker. +fn safe_idle_text_event(input: &mut KimiK3Input<'_>) -> ModalResult { + safe_text_len_mul(input, IDLE_MARKERS).map(|len| KimiK3Event::Text { len }) +} + +/// Parse safe reasoning before the think close marker. +fn safe_reasoning_event(input: &mut KimiK3Input<'_>) -> ModalResult { + safe_text_len_mul(input, REASONING_MARKERS).map(|len| KimiK3Event::Reasoning { len }) +} + +/// Parse safe response text before the next channel marker. +fn safe_response_text_event(input: &mut KimiK3Input<'_>) -> ModalResult { + safe_text_len_mul(input, RESPONSE_MARKERS).map(|len| KimiK3Event::Text { len }) +} + +/// Skip non-content noise after the response channel closed. +fn skip_epilogue_noise_event(input: &mut KimiK3Input<'_>) -> ModalResult { + safe_text_len_mul(input, EPILOGUE_MARKERS).map(|_| KimiK3Event::Skip) +} + +/// Skip non-content noise between `call` blocks. +fn skip_tools_noise_event(input: &mut KimiK3Input<'_>) -> ModalResult { + safe_text_len_mul(input, TOOLS_MARKERS).map(|_| KimiK3Event::Skip) +} + +/// Parse a `call` open tag into its tool name and one-based index. +fn call_open_event(input: &mut KimiK3Input<'_>) -> ModalResult { + let (attrs,) = seq!( + _: literal(CALL_OPEN), + take_until(0.., SEP), + _: literal(SEP), + ) + .parse_next(input)?; + let attrs = parse_tag_attrs(attrs)?; + + Ok(KimiK3Event::CallOpen { + name: attr_value(&attrs, "tool").unwrap_or_default().to_string(), + index: attr_value(&attrs, "index") + .filter(|index| !index.is_empty()) + .map(str::to_string), + }) +} + +/// Parse the buffered call body through `<|close|>call<|sep|>` into a +/// completed call. +fn call_body_event( + input: &mut KimiK3Input<'_>, + scan: &mut MarkerScanState, +) -> ModalResult { + let (body,) = seq!( + take_until_marker(CALL_CLOSE, scan), + _: literal(CALL_CLOSE), + ) + .parse_next(input)?; + let arguments = parse_call_arguments(body)?; + + Ok(KimiK3Event::CallComplete { arguments }) +} + +/// Parse a complete `call` body into the OpenAI-style arguments JSON string. +/// +/// The body is either one raw `json` block (passed through unmodified) or a +/// sequence of typed `argument` blocks (converted per their `type` tags). +fn parse_call_arguments(body: &str) -> ModalResult { + let mut input = body; + terminated( + delimited(ws0, alt((json_block_arguments, typed_arguments)), ws0), + eof, + ) + .parse_next(&mut input) + .map_err(|_| xtml_error("Kimi K3 call body")) +} + +/// Parse one raw `json` argument block, passing its body through unmodified. +fn json_block_arguments(input: &mut &str) -> ModalResult { + let (raw,) = seq!( + _: literal(JSON_OPEN), + _: take_until(0.., SEP), // attrs (`type="object"`), unused on decode + _: literal(SEP), + take_until(0.., JSON_CLOSE), + _: literal(JSON_CLOSE), + ) + .parse_next(input)?; + Ok(raw.to_string()) +} + +/// Parse typed `argument` blocks into a serialized JSON object, preserving +/// argument order. +fn typed_arguments(input: &mut &str) -> ModalResult { + let pairs: Vec<(String, Value)> = + repeat(0.., terminated(argument_block, ws0)).parse_next(input)?; + let arguments = pairs.into_iter().collect::>(); + serde_json::to_string(&arguments).map_err(|_| xtml_error("Kimi K3 arguments")) +} + +/// Parse one typed `argument` block into its key/value pair. +fn argument_block(input: &mut &str) -> ModalResult<(String, Value)> { + let (attrs, raw) = seq!( + _: literal(ARG_OPEN), + take_until(0.., SEP), + _: literal(SEP), + take_until(0.., ARG_CLOSE), + _: literal(ARG_CLOSE), + ) + .parse_next(input)?; + let attrs = parse_tag_attrs(attrs)?; + + let key = attr_value(&attrs, "key").unwrap_or_default().to_string(); + let arg_type = attr_value(&attrs, "type").unwrap_or("string"); + Ok((key, decode_argument_value(arg_type, raw))) +} + +/// Decode one typed argument value per its XTML `type` tag. +/// +/// `string` values pass the raw text through (the renderer emits them +/// unescaped); other types are JSON-decoded, falling back to the raw text on +/// malformed payloads so one quirky value does not fail the whole call. +fn decode_argument_value(arg_type: &str, raw: &str) -> Value { + if arg_type == "string" { + return Value::String(raw.to_string()); + } + serde_json::from_str(raw).unwrap_or_else(|_| Value::String(raw.to_string())) +} + +/// Parse a complete XTML attribute string like ` tool="get_weather" index="1"`. +fn parse_tag_attrs(attrs: &str) -> ModalResult> { + let mut input = attrs; + terminated(repeat(0.., preceded(ws1, tag_attr)), (ws0, eof)) + .parse_next(&mut input) + .map_err(|_| xtml_error("XTML tag attributes")) +} + +/// Parse one XTML `key="value"` attribute pair. +fn tag_attr(input: &mut &str) -> ModalResult<(String, String)> { + seq!( + take_while(1.., |char: char| char.is_alphanumeric() || char == '_').map(str::to_string), + _: literal("=\""), + take_till(0.., '"').map(unescape_attr_value), + _: literal("\""), + ) + .parse_next(input) +} + +/// Reverse XTML attribute escaping: `"` first, then `&` (the inverse +/// of the encode order). +fn unescape_attr_value(value: &str) -> String { + value.replace(""", "\"").replace("&", "&") +} + +/// Look up one parsed attribute value by key. +fn attr_value<'a>(attrs: &'a [(String, String)], key: &str) -> Option<&'a str> { + attrs.iter().find(|(name, _)| name == key).map(|(_, value)| value.as_str()) +} + +/// Build a cut error for determinably malformed XTML structure. +fn xtml_error(label: &'static str) -> ErrMode { + let mut error = ContextError::new(); + error.push(StrContext::Label(label)); + ErrMode::Cut(error) +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use serde_json::{Value, json}; + use thiserror_ext::AsReport; + use aphrodite_tokenizer::Tokenizer as _; + use aphrodite_tokenizer::test_utils::TestTokenizer; + + use super::{ + END_OF_MSG, KimiK3UnifiedParser, OPEN, RESPONSE_CLOSE, RESPONSE_OPEN, SEP, THINK_CLOSE, + THINK_OPEN, TOOLS_CLOSE, TOOLS_OPEN, + }; + use crate::tool::ToolCallDelta; + use crate::unified::{ + UnifiedParser, UnifiedParserError, UnifiedParserEvent, UnifiedParserOutput, + }; + + const OPEN_ID: u32 = 256; + const CLOSE_ID: u32 = 257; + const SEP_ID: u32 = 258; + const END_OF_MSG_ID: u32 = 259; + + fn tokenizer() -> TestTokenizer { + TestTokenizer::new() + .with_special_token(OPEN, OPEN_ID) + .with_special_token("<|close|>", CLOSE_ID) + .with_special_token(SEP, SEP_ID) + .with_special_token(END_OF_MSG, END_OF_MSG_ID) + } + + trait UnifiedParserTestExt { + fn parse_chunk(&mut self, chunk: &str) -> super::Result; + fn parse_complete(&mut self, text: &str) -> super::Result; + } + + impl UnifiedParserTestExt for KimiK3UnifiedParser { + fn parse_chunk(&mut self, chunk: &str) -> super::Result { + let mut output = UnifiedParserOutput::default(); + self.parse_into(chunk, &mut output)?; + Ok(output) + } + + fn parse_complete(&mut self, text: &str) -> super::Result { + let mut output = self.parse_chunk(text)?; + output.append(self.finish()?); + Ok(output) + } + } + + trait UnifiedOutputTestExt { + fn normal_text(&self) -> String; + fn reasoning_text(&self) -> String; + fn calls(&self) -> Vec; + } + + impl UnifiedOutputTestExt for UnifiedParserOutput { + fn normal_text(&self) -> String { + self.events + .iter() + .filter_map(|event| match event { + UnifiedParserEvent::Text(text) => Some(text.as_str()), + _ => None, + }) + .collect() + } + + fn reasoning_text(&self) -> String { + self.events + .iter() + .filter_map(|event| match event { + UnifiedParserEvent::Reasoning(text) => Some(text.as_str()), + _ => None, + }) + .collect() + } + + fn calls(&self) -> Vec { + self.events + .iter() + .filter_map(|event| match event { + UnifiedParserEvent::ToolCall(call) => Some(call.clone()), + _ => None, + }) + .collect() + } + } + + fn test_parser() -> KimiK3UnifiedParser { + KimiK3UnifiedParser::new(&[], Arc::new(tokenizer())).unwrap() + } + + fn collect_stream(parser: &mut KimiK3UnifiedParser, chunks: &[&str]) -> UnifiedParserOutput { + let mut output = UnifiedParserOutput::default(); + for chunk in chunks { + output.append(parser.parse_chunk(chunk).unwrap()); + } + output.append(parser.finish().unwrap()); + output + } + + /// Split `text` into small chunks to stress marker-split handling. + fn char_chunks(text: &str, size: usize) -> Vec { + let chars: Vec = text.chars().collect(); + chars.chunks(size).map(|chunk| chunk.iter().collect()).collect() + } + + fn arg(key: &str, arg_type: &str, value: &str) -> String { + format!( + "{OPEN}argument key=\"{key}\" type=\"{arg_type}\"{SEP}{value}<|close|>argument{SEP}" + ) + } + + fn call(attrs: &str, body: &str) -> String { + format!("{OPEN}call {attrs}{SEP}{body}<|close|>call{SEP}") + } + + fn thinking_output(reasoning: &str, response: &str, tools_body: &str) -> String { + let mut output = format!("{THINK_OPEN}{reasoning}{THINK_CLOSE}"); + output.push_str(&format!("{RESPONSE_OPEN}{response}{RESPONSE_CLOSE}")); + if !tools_body.is_empty() { + output.push_str(&format!("{TOOLS_OPEN}{tools_body}{TOOLS_CLOSE}")); + } + output.push_str("<|close|>message<|sep|>"); + output + } + + fn first_call(output: &UnifiedParserOutput) -> ToolCallDelta { + output.calls().first().expect("expected one tool call").clone() + } + + #[test] + fn kimi_k3_create_requires_structural_tokens() { + let error = match KimiK3UnifiedParser::new(&[], Arc::new(TestTokenizer::new())) { + Ok(_) => panic!("expected missing token error"), + Err(error) => error, + }; + + assert!(matches!( + error, + UnifiedParserError::MissingToken { token } if token == OPEN + )); + } + + #[test] + fn kimi_k3_parses_reasoning_response_and_typed_tool_call() { + let body = [ + arg("city", "string", "Hangzhou"), + arg("days", "number", "1.5"), + arg("detailed", "boolean", "true"), + arg("filters", "object", r#"{"kind":"rain"}"#), + arg("hours", "array", "[8,20]"), + ] + .concat(); + let text = thinking_output( + "Need the weather tool.", + "I'll check.", + &call("tool=\"get_weather\" index=\"1\"", &body), + ); + + let mut parser = test_parser(); + let output = parser.parse_complete(&text).unwrap(); + + assert_eq!(output.reasoning_text(), "Need the weather tool."); + assert_eq!(output.normal_text(), "I'll check."); + let call = first_call(&output); + assert_eq!(call.tool_index, 0); + assert_eq!(call.name.as_deref(), Some("get_weather")); + assert_eq!( + serde_json::from_str::(&call.arguments).unwrap(), + json!({ + "city": "Hangzhou", + "days": 1.5, + "detailed": true, + "filters": { "kind": "rain" }, + "hours": [8, 20], + }) + ); + assert_eq!(parser.tool_call_id(0), Some("get_weather:0")); + } + + #[test] + fn kimi_k3_arguments_preserve_order_and_number_formatting() { + let body = [ + arg("y", "number", "1.0"), + arg("x", "number", "2"), + arg("items", "array", r#"["left","right"]"#), + ] + .concat(); + let text = thinking_output("t", "", &call("tool=\"add\" index=\"1\"", &body)); + + let mut parser = test_parser(); + let output = parser.parse_complete(&text).unwrap(); + + assert_eq!( + first_call(&output).arguments, + r#"{"y":1.0,"x":2,"items":["left","right"]}"# + ); + } + + #[test] + fn kimi_k3_streaming_splits_markers_across_chunks() { + let text = thinking_output( + "step by step", + "the answer", + &call("tool=\"calc\" index=\"1\"", &arg("x", "number", "42")), + ); + + for size in [1, 3, 7] { + let chunks = char_chunks(&text, size); + let chunk_refs: Vec<&str> = chunks.iter().map(String::as_str).collect(); + let output = collect_stream(&mut test_parser(), &chunk_refs); + + assert_eq!(output.reasoning_text(), "step by step", "chunk size {size}"); + assert_eq!(output.normal_text(), "the answer", "chunk size {size}"); + assert_eq!(first_call(&output).name.as_deref(), Some("calc")); + assert_eq!(first_call(&output).arguments, r#"{"x":42}"#); + } + } + + #[test] + fn kimi_k3_streaming_emits_text_incrementally() { + let mut parser = test_parser(); + let prompt = tokenizer().encode("<|open|>response<|sep|>", false).unwrap(); + parser.initialize(&prompt).unwrap(); + + let first = parser.parse_chunk("Hel").unwrap(); + assert_eq!(first.normal_text(), "Hel"); + + let second = parser.parse_chunk("lo<|close|>resp").unwrap(); + assert_eq!(second.normal_text(), "lo"); + + let mut output = parser.parse_chunk("onse<|sep|>").unwrap(); + output.append(parser.finish().unwrap()); + assert_eq!(output.normal_text(), ""); + } + + #[test] + fn kimi_k3_initialize_think_prefill_starts_in_reasoning() { + let mut parser = test_parser(); + let prompt = tokenizer() + .encode( + "<|open|>message role=\"assistant\"<|sep|><|open|>think<|sep|>", + false, + ) + .unwrap(); + parser.initialize(&prompt).unwrap(); + + let output = parser + .parse_complete(&format!( + "reasoning{THINK_CLOSE}{RESPONSE_OPEN}answer{RESPONSE_CLOSE}<|close|>message{SEP}" + )) + .unwrap(); + + assert_eq!(output.reasoning_text(), "reasoning"); + assert_eq!(output.normal_text(), "answer"); + } + + #[test] + fn kimi_k3_initialize_response_prefill_starts_in_response() { + let mut parser = test_parser(); + let prompt = tokenizer() + .encode( + "<|open|>message role=\"assistant\"<|sep|><|open|>response<|sep|>", + false, + ) + .unwrap(); + parser.initialize(&prompt).unwrap(); + + let output = parser + .parse_complete(&format!("answer{RESPONSE_CLOSE}<|close|>message{SEP}")) + .unwrap(); + + assert_eq!(output.normal_text(), "answer"); + assert!(output.reasoning_text().is_empty()); + } + + #[test] + fn kimi_k3_initialize_message_open_prefill_starts_idle() { + let mut parser = test_parser(); + let prompt = + tokenizer().encode("<|open|>message role=\"assistant\"<|sep|>", false).unwrap(); + parser.initialize(&prompt).unwrap(); + + let output = parser + .parse_complete(&format!("{THINK_OPEN}reason{THINK_CLOSE}{RESPONSE_OPEN}hi")) + .unwrap(); + + assert_eq!(output.reasoning_text(), "reason"); + assert_eq!(output.normal_text(), "hi"); + } + + #[test] + fn kimi_k3_plain_text_falls_through_as_text() { + let output = collect_stream(&mut test_parser(), &["plain ", "answer"]); + + assert_eq!(output.normal_text(), "plain answer"); + assert!(output.reasoning_text().is_empty()); + assert!(output.calls().is_empty()); + } + + #[test] + fn kimi_k3_tool_call_waits_for_close_marker() { + let mut parser = test_parser(); + let mut output = UnifiedParserOutput::default(); + + let argument = arg("x", "number", "1"); + for chunk in [ + TOOLS_OPEN, + "<|open|>call tool=\"calc\" index=\"1\"<|sep|>", + argument.as_str(), + ] { + output.append(parser.parse_chunk(chunk).unwrap()); + assert!(output.calls().is_empty()); + } + + output.append(parser.parse_chunk("<|close|>call<|sep|>").unwrap()); + + assert_eq!(first_call(&output).name.as_deref(), Some("calc")); + assert_eq!(first_call(&output).arguments, r#"{"x":1}"#); + } + + #[test] + fn kimi_k3_parses_multiple_tool_calls() { + let tools_body = format!( + "{}{}", + call( + "tool=\"get_weather\" index=\"1\"", + &arg("city", "string", "SF") + ), + call("tool=\"get_time\" index=\"2\"", ""), + ); + let text = thinking_output("t", "r", &tools_body); + + let mut parser = test_parser(); + let output = parser.parse_complete(&text).unwrap(); + + let calls = output.calls(); + assert_eq!(calls.len(), 2); + assert_eq!(calls[0].tool_index, 0); + assert_eq!(calls[0].name.as_deref(), Some("get_weather")); + assert_eq!(calls[0].arguments, r#"{"city":"SF"}"#); + assert_eq!(calls[1].tool_index, 1); + assert_eq!(calls[1].name.as_deref(), Some("get_time")); + assert_eq!(calls[1].arguments, "{}"); + assert_eq!(parser.tool_call_id(0), Some("get_weather:0")); + assert_eq!(parser.tool_call_id(1), Some("get_time:1")); + } + + #[test] + fn kimi_k3_json_block_arguments_pass_through_raw() { + // Spacing and key order must survive unmodified: raw `json` blocks are + // not validated or normalized. + let raw = r#"{"b": 1, "a": [2 , 3]}"#; + let body = format!("{OPEN}json type=\"object\"{SEP}{raw}<|close|>json{SEP}"); + let text = thinking_output("t", "", &call("tool=\"run\" index=\"1\"", &body)); + + let mut parser = test_parser(); + let output = parser.parse_complete(&text).unwrap(); + + assert_eq!(first_call(&output).arguments, raw); + } + + #[test] + fn kimi_k3_string_argument_passes_raw_text_through() { + let value = "line one\nline two {\"not\": \"json\"} & "; + let text = thinking_output( + "t", + "", + &call( + "tool=\"write\" index=\"1\"", + &arg("content", "string", value), + ), + ); + + let mut parser = test_parser(); + let output = parser.parse_complete(&text).unwrap(); + + assert_eq!( + serde_json::from_str::(&first_call(&output).arguments).unwrap(), + json!({ "content": value }) + ); + } + + #[test] + fn kimi_k3_malformed_typed_argument_falls_back_to_raw_text() { + let text = thinking_output( + "t", + "", + &call( + "tool=\"calc\" index=\"1\"", + &arg("x", "number", "not a number"), + ), + ); + + let mut parser = test_parser(); + let output = parser.parse_complete(&text).unwrap(); + + assert_eq!(first_call(&output).arguments, r#"{"x":"not a number"}"#); + } + + #[test] + fn kimi_k3_attribute_values_are_unescaped() { + let text = thinking_output( + "t", + "", + &call( + "tool=\"a"b&c\" index=\"1\"", + &arg("key", "string", "value"), + ), + ); + + let mut parser = test_parser(); + let output = parser.parse_complete(&text).unwrap(); + + assert_eq!(first_call(&output).name.as_deref(), Some("a\"b&c")); + } + + #[test] + fn kimi_k3_call_without_tool_name_is_dropped() { + let tools_body = format!( + "{}{}", + call("index=\"1\"", &arg("x", "number", "1")), + call("tool=\"real\" index=\"2\"", ""), + ); + let text = thinking_output("t", "", &tools_body); + + let mut parser = test_parser(); + let output = parser.parse_complete(&text).unwrap(); + + let calls = output.calls(); + assert_eq!(calls.len(), 1); + assert_eq!(calls[0].tool_index, 0); + assert_eq!(calls[0].name.as_deref(), Some("real")); + assert_eq!(parser.tool_call_id(0), Some("real:1")); + } + + #[test] + fn kimi_k3_tool_call_id_follows_index_attribute() { + let tools_body = format!( + "{}{}{}", + call("tool=\"first\" index=\"3\"", ""), + call("tool=\"second\"", ""), + call("tool=\"third\" index=\"x\"", ""), + ); + let text = thinking_output("t", "", &tools_body); + + let mut parser = test_parser(); + parser.parse_complete(&text).unwrap(); + + assert_eq!(parser.tool_call_id(0), Some("first:2")); + assert_eq!(parser.tool_call_id(1), Some("second")); + assert_eq!(parser.tool_call_id(2), Some("third:x")); + } + + #[test] + fn kimi_k3_ignores_output_after_message_close() { + let mut parser = test_parser(); + let output = parser + .parse_complete(&format!( + "{RESPONSE_OPEN}answer{RESPONSE_CLOSE}<|close|>message{SEP}junk{END_OF_MSG}" + )) + .unwrap(); + + assert_eq!(output.normal_text(), "answer"); + } + + #[test] + fn kimi_k3_epilogue_noise_is_not_content() { + let text = format!( + "{RESPONSE_OPEN}answer{RESPONSE_CLOSE}\n{TOOLS_OPEN}{}{TOOLS_CLOSE}\n<|close|>message{SEP}", + call("tool=\"calc\" index=\"1\"", ""), + ); + + let mut parser = test_parser(); + let output = parser.parse_complete(&text).unwrap(); + + assert_eq!(output.normal_text(), "answer"); + assert_eq!(output.calls().len(), 1); + } + + #[test] + fn kimi_k3_finish_flushes_unclosed_reasoning() { + let mut parser = test_parser(); + let mut output = parser.parse_chunk(&format!("{THINK_OPEN}still thinking")).unwrap(); + output.append(parser.finish().unwrap()); + + assert_eq!(output.reasoning_text(), "still thinking"); + assert!(output.normal_text().is_empty()); + } + + #[test] + fn kimi_k3_finish_flushes_partial_marker_as_text() { + let mut parser = test_parser(); + let mut output = parser.parse_chunk("answer<|clo").unwrap(); + output.append(parser.finish().unwrap()); + + assert_eq!(output.normal_text(), "answer<|clo"); + } + + #[test] + fn kimi_k3_finish_fails_mid_tool_call() { + let mut parser = test_parser(); + parser + .parse_chunk(&format!( + "{TOOLS_OPEN}<|open|>call tool=\"calc\" index=\"1\"<|sep|>{}", + arg("x", "number", "1") + )) + .unwrap(); + + let error = parser.finish().unwrap_err(); + + assert!(error.to_report_string().contains("incomplete Kimi K3 tool call")); + } + + #[test] + fn kimi_k3_finish_after_truncated_tools_keeps_complete_calls() { + let mut parser = test_parser(); + let mut output = parser + .parse_chunk(&format!( + "{TOOLS_OPEN}{}", + call("tool=\"calc\" index=\"1\"", &arg("x", "number", "1")) + )) + .unwrap(); + output.append(parser.finish().unwrap()); + + assert_eq!(first_call(&output).name.as_deref(), Some("calc")); + } + + #[test] + fn kimi_k3_malformed_call_attributes_fail_fast() { + let mut parser = test_parser(); + let error = parser + .parse_chunk(&format!("{TOOLS_OPEN}<|open|>call garbage attrs<|sep|>")) + .unwrap_err(); + + assert!(error.to_report_string().contains("XTML tag attributes")); + } + + #[test] + fn kimi_k3_empty_response_channel_emits_nothing() { + let text = thinking_output("t", "", &call("tool=\"calc\" index=\"1\"", "")); + + let mut parser = test_parser(); + let output = parser.parse_complete(&text).unwrap(); + + assert!(output.normal_text().is_empty()); + assert_eq!(output.reasoning_text(), "t"); + assert_eq!(output.calls().len(), 1); + } + + #[test] + fn kimi_k3_reset_returns_buffered_text() { + let mut parser = test_parser(); + let prompt = tokenizer().encode("<|open|>response<|sep|>", false).unwrap(); + parser.initialize(&prompt).unwrap(); + parser.parse_chunk("answer<|close|>resp").unwrap(); + + let raw = parser.reset(); + + assert_eq!(raw, "<|close|>resp"); + } +} diff --git a/rust/src/parser/src/unified/kimi_k3/structural_tag.rs b/rust/src/parser/src/unified/kimi_k3/structural_tag.rs new file mode 100644 index 0000000000..d3515ad2dc --- /dev/null +++ b/rust/src/parser/src/unified/kimi_k3/structural_tag.rs @@ -0,0 +1,503 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +//! Structural-tag grammar for Kimi K3 XTML tool calls. + +use serde_json::{Map, Value}; +use xgrammar_structural_tag::Result; +use xgrammar_structural_tag::builders::{ + StructuralTagBuilder, StructuralTagContext, StructuralTagOptions, +}; +use xgrammar_structural_tag::format::{Format, JsonSchemaFormat, StructuralTag, TagFormat}; +use xgrammar_structural_tag::tool::{BuilderToolChoice, FunctionToolParam, function_parameters}; + +use super::{ + ARG_CLOSE, CALL_CLOSE, END_OF_MSG, JSON_CLOSE, JSON_OPEN, MESSAGE_CLOSE, OPEN, RESPONSE_CLOSE, + RESPONSE_OPEN, SEP, THINK_CLOSE, TOOLS_CLOSE, TOOLS_OPEN, +}; + +pub(super) static KIMI_K3_STRUCTURAL_TAG_BUILDER: KimiK3StructuralTagBuilder = + KimiK3StructuralTagBuilder; + +const XTML_TYPES: &[&str] = &["string", "number", "boolean", "null", "object", "array"]; + +/// Kimi K3 XTML structural-tag builder. +#[derive(Debug, Clone, Copy, Default)] +pub struct KimiK3StructuralTagBuilder; + +impl StructuralTagBuilder for KimiK3StructuralTagBuilder { + fn build(&self, ctx: StructuralTagContext<'_>) -> Result { + let mut elements = response_prefix(ctx.options.reasoning); + + let tools = match ctx.tool_choice { + // Serving lowering filters empty tools before calling the builder, + // while direct `build_structural_tag(..., auto, ...)` calls do not. + BuilderToolChoice::Auto if ctx.function_tools.is_empty() => None, + BuilderToolChoice::Auto => Some(Format::optional(tools_channel( + ctx.function_tools, + ctx.options, + ))), + BuilderToolChoice::Forced | BuilderToolChoice::Required => { + Some(tools_channel(ctx.function_tools, ctx.options)) + } + }; + if let Some(tools) = tools { + elements.push(tools); + } + elements.push(Format::optional(Format::const_string(MESSAGE_CLOSE))); + + Ok(StructuralTag::new(Format::sequence(elements))) + } +} + +fn response_prefix(reasoning: bool) -> Vec { + let mut elements = Vec::new(); + if reasoning { + elements.push(Format::tag( + "", + Format::any_text_excluding(&[THINK_CLOSE, END_OF_MSG]), + THINK_CLOSE, + )); + elements.push(Format::const_string(RESPONSE_OPEN)); + } else { + elements.push(Format::optional(Format::const_string(RESPONSE_OPEN))); + } + elements.push(Format::tag( + "", + Format::any_text_excluding(&[RESPONSE_CLOSE, TOOLS_OPEN, MESSAGE_CLOSE, END_OF_MSG]), + RESPONSE_CLOSE, + )); + elements +} + +fn tools_channel(tools: &[FunctionToolParam], options: StructuralTagOptions) -> Format { + let calls = tools.iter().map(|tool| call_tag(tool, options)).collect(); + Format::tag( + TOOLS_OPEN, + Format::tags_with_separator(calls, "", true, false), + TOOLS_CLOSE, + ) +} + +fn call_tag(tool: &FunctionToolParam, options: StructuralTagOptions) -> TagFormat { + let parameters = function_parameters(&tool.function); + let call_body = Format::or(vec![ + typed_arguments(¶meters, options), + raw_json_arguments(¶meters, options), + ]); + + TagFormat::new( + format!( + "{OPEN}call tool=\"{}\" index=\"", + escape_attr_value(&tool.function.name) + ), + Format::sequence(vec![ + Format::regex("[1-9][0-9]*"), + Format::const_string(format!("\"{SEP}")), + call_body, + ]), + CALL_CLOSE, + ) +} + +fn typed_arguments(parameters: &Value, options: StructuralTagOptions) -> Format { + let Some(schema) = parameters.as_object() else { + return if parameters == &Value::Bool(false) { + Format::const_string("") + } else { + Format::star(permissive_argument()) + }; + }; + let Some(properties) = schema.get("properties").and_then(Value::as_object) else { + return Format::star(permissive_argument()); + }; + if properties.is_empty() { + return Format::star(permissive_argument()); + } + + let root_defs = root_definitions(schema); + let arguments = properties + .iter() + .flat_map(|(key, schema)| argument_tags(key, schema, &root_defs, options)) + .map(Format::Tag) + .collect::>(); + let arguments = match arguments.as_slice() { + [argument] => argument.clone(), + _ => Format::or(arguments), + }; + // Keep typed arguments order-agnostic and non-unique, but do not allow an + // empty call when the root schema declares required properties. + if schema + .get("required") + .and_then(Value::as_array) + .is_some_and(|required| !required.is_empty()) + { + Format::plus(arguments) + } else { + Format::star(arguments) + } +} + +fn argument_tags( + key: &str, + schema: &Value, + root_defs: &Map, + options: StructuralTagOptions, +) -> Vec { + let types = schema_types(schema); + types + .into_iter() + .map(|xtml_type| { + let content = if xtml_type == "string" { + string_argument_content(schema) + } else { + json_schema( + attach_root_definitions(&narrow_schema_type(schema, xtml_type), root_defs), + options, + ) + }; + TagFormat::new( + format!( + "{OPEN}argument key=\"{}\" type=\"{xtml_type}\"{SEP}", + escape_attr_value(key) + ), + content, + ARG_CLOSE, + ) + }) + .collect() +} + +fn schema_types(schema: &Value) -> Vec<&'static str> { + let Some(schema) = schema.as_object() else { + return XTML_TYPES.to_vec(); + }; + let mut types = Vec::new(); + match schema.get("type") { + Some(Value::String(value)) => push_schema_type(&mut types, value), + Some(Value::Array(values)) => { + for value in values.iter().filter_map(Value::as_str) { + push_schema_type(&mut types, value); + } + } + _ => {} + } + if types.is_empty() + && let Some(value) = schema.get("const") + { + push_value_type(&mut types, value); + } + if types.is_empty() + && let Some(values) = schema.get("enum").and_then(Value::as_array) + { + for value in values { + push_value_type(&mut types, value); + } + } + if types.is_empty() { + XTML_TYPES.to_vec() + } else { + types + } +} + +fn push_value_type(types: &mut Vec<&'static str>, value: &Value) { + let xtml_type = match value { + Value::String(_) => "string", + Value::Number(_) => "number", + Value::Bool(_) => "boolean", + Value::Null => "null", + Value::Object(_) => "object", + Value::Array(_) => "array", + }; + if !types.contains(&xtml_type) { + types.push(xtml_type); + } +} + +fn narrow_schema_type(schema: &Value, xtml_type: &str) -> Value { + let Some(mut schema) = schema.as_object().cloned() else { + return schema.clone(); + }; + let json_type = if xtml_type == "number" && explicitly_integer_only(&schema) { + "integer" + } else { + xtml_type + }; + schema.insert("type".to_string(), Value::String(json_type.to_string())); + Value::Object(schema) +} + +fn explicitly_integer_only(schema: &Map) -> bool { + match schema.get("type") { + Some(Value::String(value)) => value == "integer", + Some(Value::Array(values)) => { + let values = values.iter().filter_map(Value::as_str).collect::>(); + values.contains(&"integer") && !values.contains(&"number") + } + _ => false, + } +} + +fn push_schema_type(types: &mut Vec<&'static str>, json_type: &str) { + let xtml_type = match json_type { + "string" => Some("string"), + "integer" | "number" => Some("number"), + "boolean" => Some("boolean"), + "null" => Some("null"), + "object" => Some("object"), + "array" => Some("array"), + _ => None, + }; + if let Some(xtml_type) = xtml_type + && !types.contains(&xtml_type) + { + types.push(xtml_type); + } +} + +fn string_argument_content(schema: &Value) -> Format { + let Some(schema) = schema.as_object() else { + return Format::any_text_excluding(&[ARG_CLOSE, CALL_CLOSE]); + }; + let values = schema + .get("enum") + .and_then(Value::as_array) + .cloned() + .or_else(|| schema.get("const").cloned().map(|value| vec![value])); + let Some(values) = values else { + return Format::any_text_excluding(&[ARG_CLOSE, CALL_CLOSE]); + }; + if values.is_empty() + || values.len() > 256 + || values + .iter() + .any(|value| value.as_str().is_none_or(|value| value.contains("<|"))) + { + return Format::any_text_excluding(&[ARG_CLOSE, CALL_CLOSE]); + } + let values = values.iter().filter_map(Value::as_str).collect::>(); + match values.as_slice() { + [value] => Format::const_string(*value), + _ => Format::or(values.into_iter().map(Format::const_string).collect()), + } +} + +fn raw_json_arguments(parameters: &Value, options: StructuralTagOptions) -> Format { + Format::tag( + format!("{JSON_OPEN} type=\"object\"{SEP}"), + json_schema(parameters.clone(), options), + JSON_CLOSE, + ) +} + +fn permissive_argument() -> Format { + let key = Format::regex(r#"(?:[^<\"&]|&(?:amp|quot);|<[^|])*"#); + let alternatives = XTML_TYPES + .iter() + .map(|xtml_type| { + Format::sequence(vec![ + key.clone(), + Format::const_string(format!("\" type=\"{xtml_type}\"{SEP}")), + if *xtml_type == "string" { + Format::any_text_excluding(&[ARG_CLOSE, CALL_CLOSE]) + } else { + Format::json_schema(Value::Bool(true)) + }, + ]) + }) + .collect(); + Format::tag( + format!("{OPEN}argument key=\""), + Format::or(alternatives), + ARG_CLOSE, + ) +} + +fn json_schema(schema: Value, options: StructuralTagOptions) -> Format { + Format::JsonSchema( + JsonSchemaFormat::new(schema) + .with_any_order(options.any_order) + .with_max_whitespace_cnt(options.max_whitespace_cnt), + ) +} + +fn root_definitions(schema: &Map) -> Map { + ["$defs", "definitions"] + .into_iter() + .filter_map(|key| schema.get(key).map(|value| (key.to_string(), value.clone()))) + .collect() +} + +fn attach_root_definitions(schema: &Value, root_defs: &Map) -> Value { + let Some(mut schema) = schema.as_object().cloned() else { + return schema.clone(); + }; + for (key, value) in root_defs { + schema.entry(key.clone()).or_insert_with(|| value.clone()); + } + Value::Object(schema) +} + +fn escape_attr_value(value: &str) -> String { + value.replace('&', "&").replace('"', """) +} + +#[cfg(test)] +mod tests { + use expect_test::expect; + use serde_json::json; + use xgrammar_structural_tag::builders::StructuralTagOptions; + use xgrammar_structural_tag::{ + FunctionDefinition, FunctionToolParam, ToolChoice, ToolParam, build_structural_tag, + }; + + use super::KimiK3StructuralTagBuilder; + + fn tool(name: &str, parameters: serde_json::Value) -> ToolParam { + ToolParam::Function(FunctionToolParam::new( + FunctionDefinition::new(name).with_parameters(parameters), + )) + } + + #[test] + fn required_structural_tag_matches_xtml_channels() { + let tools = vec![tool( + "get_weather", + json!({ + "$defs": { + "place": { "type": "object", "properties": { "city": { "type": "string" } } } + }, + "type": "object", + "properties": { + "unit": { "type": "string", "enum": ["celsius", "fahrenheit"] }, + "place": { "$ref": "#/$defs/place", "type": "object" } + }, + "required": ["place"] + }), + )]; + let tag = build_structural_tag( + KimiK3StructuralTagBuilder, + &tools, + ToolChoice::required(), + StructuralTagOptions::default().with_reasoning(false), + ) + .unwrap(); + + expect![[r##"{"type":"structural_tag","format":{"type":"sequence","elements":[{"type":"optional","content":{"type":"const_string","value":"<|open|>response<|sep|>"}},{"type":"tag","begin":"","content":{"type":"any_text","excludes":["<|close|>response<|sep|>","<|open|>tools<|sep|>","<|close|>message<|sep|>","<|end_of_msg|>"]},"end":"<|close|>response<|sep|>"},{"type":"tag","begin":"<|open|>tools<|sep|>","content":{"type":"tags_with_separator","tags":[{"begin":"<|open|>call tool=\"get_weather\" index=\"","content":{"type":"sequence","elements":[{"type":"regex","pattern":"[1-9][0-9]*"},{"type":"const_string","value":"\"<|sep|>"},{"type":"or","elements":[{"type":"plus","content":{"type":"or","elements":[{"type":"tag","begin":"<|open|>argument key=\"unit\" type=\"string\"<|sep|>","content":{"type":"or","elements":[{"type":"const_string","value":"celsius"},{"type":"const_string","value":"fahrenheit"}]},"end":"<|close|>argument<|sep|>"},{"type":"tag","begin":"<|open|>argument key=\"place\" type=\"object\"<|sep|>","content":{"type":"json_schema","json_schema":{"$ref":"#/$defs/place","type":"object","$defs":{"place":{"type":"object","properties":{"city":{"type":"string"}}}}},"style":"json","any_order":false,"max_whitespace_cnt":null},"end":"<|close|>argument<|sep|>"}]}},{"type":"tag","begin":"<|open|>json type=\"object\"<|sep|>","content":{"type":"json_schema","json_schema":{"$defs":{"place":{"type":"object","properties":{"city":{"type":"string"}}}},"type":"object","properties":{"unit":{"type":"string","enum":["celsius","fahrenheit"]},"place":{"$ref":"#/$defs/place","type":"object"}},"required":["place"]},"style":"json","any_order":false,"max_whitespace_cnt":null},"end":"<|close|>json<|sep|>"}]}]},"end":"<|close|>call<|sep|>"}],"separator":"","at_least_one":true,"stop_after_first":false},"end":"<|close|>tools<|sep|>"},{"type":"optional","content":{"type":"const_string","value":"<|close|>message<|sep|>"}}]}}"##]].assert_eq(&tag.to_json_string().unwrap()); + } + + #[test] + fn typed_arguments_require_one_tag_only_for_nonempty_required() { + let required = super::typed_arguments( + &json!({ + "type": "object", + "properties": { "query": { "type": "string" } }, + "required": ["query"] + }), + StructuralTagOptions::default(), + ); + let optional = super::typed_arguments( + &json!({ + "type": "object", + "properties": { "query": { "type": "string" } } + }), + StructuralTagOptions::default(), + ); + let empty_required = super::typed_arguments( + &json!({ + "type": "object", + "properties": { "query": { "type": "string" } }, + "required": [] + }), + StructuralTagOptions::default(), + ); + + assert_eq!(serde_json::to_value(required).unwrap()["type"], "plus"); + assert_eq!(serde_json::to_value(optional).unwrap()["type"], "star"); + assert_eq!( + serde_json::to_value(empty_required).unwrap()["type"], + "star" + ); + } + + #[test] + fn reasoning_grammar_starts_inside_prefilled_think_channel() { + let tools = vec![tool("ping", json!({ "type": "object", "properties": {} }))]; + let tag = build_structural_tag( + KimiK3StructuralTagBuilder, + &tools, + ToolChoice::auto(), + StructuralTagOptions::default().with_reasoning(true), + ) + .unwrap(); + let value = serde_json::to_value(tag).unwrap(); + + assert_eq!( + value["format"]["elements"][0]["end"], + "<|close|>think<|sep|>" + ); + assert_eq!( + value["format"]["elements"][1]["value"], + "<|open|>response<|sep|>" + ); + assert_eq!(value["format"]["elements"][3]["type"], "optional"); + } + + #[test] + fn forced_choice_keeps_only_the_named_tool() { + let tools = vec![ + tool("search", json!({ "type": "object" })), + tool("lookup", json!({ "type": "object" })), + ]; + let tag = build_structural_tag( + KimiK3StructuralTagBuilder, + &tools, + ToolChoice::function("lookup"), + StructuralTagOptions::default().with_reasoning(false), + ) + .unwrap() + .to_json_string() + .unwrap(); + + assert!(tag.contains("lookup")); + assert!(!tag.contains("search")); + } + + #[test] + fn union_argument_content_matches_its_xtml_type() { + let tools = vec![tool( + "set_count", + json!({ + "type": "object", + "properties": { + "count": { "type": ["integer", "null"] } + } + }), + )]; + let tag = build_structural_tag( + KimiK3StructuralTagBuilder, + &tools, + ToolChoice::required(), + StructuralTagOptions::default(), + ) + .unwrap() + .to_json_string() + .unwrap(); + + assert!(tag.contains(r#"type=\"number\""#)); + assert!(tag.contains(r#""json_schema":{"type":"integer"}"#), "{tag}"); + assert!(tag.contains(r#"type=\"null\""#)); + assert!(tag.contains(r#""json_schema":{"type":"null"}"#), "{tag}"); + } + + #[test] + fn unsafe_string_enum_falls_back_as_a_whole() { + let format = super::string_argument_content(&json!({ + "type": "string", + "enum": ["safe", "<|unsafe"] + })); + + assert_eq!(serde_json::to_value(format).unwrap()["type"], "any_text"); + } +} diff --git a/rust/src/parser/src/unified/mod.rs b/rust/src/parser/src/unified/mod.rs index 1836cb7397..092d5c852c 100644 --- a/rust/src/parser/src/unified/mod.rs +++ b/rust/src/parser/src/unified/mod.rs @@ -6,10 +6,12 @@ mod combined; mod gemma4; mod inkling; +mod kimi_k3; pub use combined::CombinedParser; pub use gemma4::Gemma4UnifiedParser; pub use inkling::InklingUnifiedParser; +pub use kimi_k3::{KimiK3StructuralTagBuilder, KimiK3UnifiedParser}; use aphrodite_tokenizer::DynTokenizer; use thiserror::Error; use thiserror_ext::Macro; diff --git a/rust/src/server/examples/external_engine_openai_qwen.rs b/rust/src/server/examples/external_engine_openai_qwen.rs index 734133d316..79ba7be4d1 100644 --- a/rust/src/server/examples/external_engine_openai_qwen.rs +++ b/rust/src/server/examples/external_engine_openai_qwen.rs @@ -1,6 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright contributors to the vLLM project +use std::collections::HashMap; use std::time::Duration; use anyhow::{Context, Result, bail}; @@ -70,6 +71,7 @@ async fn main() -> Result<()> { language_model_only: false, chat_template: None, default_chat_template_kwargs: None, + limit_mm_per_prompt: HashMap::new(), chat_template_content_format: ChatTemplateContentFormatOption::Auto, max_logprobs: None, api_server_options: ApiServerOptions::default(), diff --git a/rust/src/server/src/config.rs b/rust/src/server/src/config.rs index 5dd16a85d5..464c24d5de 100644 --- a/rust/src/server/src/config.rs +++ b/rust/src/server/src/config.rs @@ -10,6 +10,7 @@ use axum::http::{HeaderName, HeaderValue, Method}; use educe::Educe; use serde::Serialize; use serde_json::Value; +use aphrodite_chat::multimodal::MmLimitPerPrompt; use aphrodite_chat::{ChatTemplateContentFormatOption, ParserSelection, RendererSelection}; use aphrodite_engine_core_client::{CoordinatorMode as EngineCoreCoordinatorMode, TransportMode}; @@ -184,6 +185,9 @@ pub struct Config { pub chat_template: Option, /// Server-default keyword arguments merged into every chat-template render. pub default_chat_template_kwargs: Option>, + /// Maximum number of input items allowed per prompt for each modality. + /// Unspecified modalities are unlimited. + pub limit_mm_per_prompt: MmLimitPerPrompt, /// How to serialize `message.content` for chat-template rendering. pub chat_template_content_format: ChatTemplateContentFormatOption, /// Optional maximum number of top log probabilities accepted by the diff --git a/rust/src/server/src/lib.rs b/rust/src/server/src/lib.rs index b879008c7f..f21d1931f1 100644 --- a/rust/src/server/src/lib.rs +++ b/rust/src/server/src/lib.rs @@ -101,6 +101,7 @@ async fn build_state(config: &Config) -> Result> { .default_chat_template_kwargs .clone() .unwrap_or_default(), + limit_mm_per_prompt: config.limit_mm_per_prompt.clone(), }, ) .await diff --git a/rust/src/server/src/routes/openai/chat_completions/convert.rs b/rust/src/server/src/routes/openai/chat_completions/convert.rs index fa216ba822..7576a3ef05 100644 --- a/rust/src/server/src/routes/openai/chat_completions/convert.rs +++ b/rust/src/server/src/routes/openai/chat_completions/convert.rs @@ -89,6 +89,15 @@ pub(super) fn prepare_chat_request( )?; let template_kwargs = request.chat_template_kwargs.unwrap_or_default(); + let response_format = + request.response_format.as_ref().map(serde_json::to_value).transpose().map_err( + |error| { + ApiError::invalid_request( + format!("failed to serialize response_format: {error}"), + Some("response_format"), + ) + }, + )?; let include_usage = (request.stream_options.as_ref()) .and_then(|options| options.include_usage) @@ -150,6 +159,7 @@ pub(super) fn prepare_chat_request( generation_prompt_mode, chat_template: request.chat_template, reasoning_effort: request.reasoning_effort, + response_format, template_kwargs, }, tools: convert_tools(request.tools)?, @@ -400,6 +410,7 @@ fn convert_tool_choice(tool_choice: Option<&ToolChoice>) -> Result aphrodite_ch } fn qwen_multimodal_model_info() -> aphrodite_chat::multimodal::MultimodalModelInfo { + qwen_multimodal_model_info_with_limits(std::collections::HashMap::new()) +} + +fn qwen_multimodal_model_info_with_limits( + limit_mm_per_prompt: aphrodite_chat::multimodal::MmLimitPerPrompt, +) -> aphrodite_chat::multimodal::MultimodalModelInfo { let config_path = std::env::temp_dir().join(format!( "aphrodite-server-qwen-config-{}.json", uuid::Uuid::new_v4() @@ -591,6 +597,7 @@ fn qwen_multimodal_model_info() -> aphrodite_chat::multimodal::MultimodalModelIn ..Default::default() }, Arc::new(fake_chat_tokenizer()), + limit_mm_per_prompt, ) .expect("load multimodal info") .expect("qwen multimodal info is registered"); @@ -2305,6 +2312,74 @@ async fn non_stream_chat_image_url_reaches_engine_mm_features() { assert_eq!(json["choices"][0]["message"]["content"], "hi"); } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +#[serial] +async fn non_stream_chat_rejects_when_image_count_exceeds_limit_mm_per_prompt() { + // The request is rejected by `--limit-mm-per-prompt` validation before + // ever reaching the engine, so the mock engine task is never awaited. + let (chat, _engine_task) = test_models_with_engine_outputs_and_backend( + b"engine-openai-mm-limit", + default_stream_output_specs(), + Arc::new(FakeChatBackend::with_multimodal_model_info( + qwen_multimodal_model_info_with_limits(std::collections::HashMap::from([( + aphrodite_chat::multimodal::MmLimitModality::Image, + aphrodite_chat::multimodal::MmLimitSpec::Count(1), + )])), + )), + ) + .await; + let app = build_router(Arc::new(AppState::new( + vec!["Qwen/Qwen1.5-0.5B-Chat".to_string()], + chat, + ))); + + let response = app + .clone() + .call( + Request::builder() + .method("POST") + .uri("/v1/chat/completions") + .header("content-type", "application/json") + .body(Body::from( + json!({ + "model": "Qwen/Qwen1.5-0.5B-Chat", + "stream": false, + "messages": [{ + "role": "user", + "content": [ + {"type": "text", "text": "describe "}, + { + "type": "image_url", + "image_url": { + "url": "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mNk+A8AAQUBAScY42YAAAAASUVORK5CYII=" + } + }, + { + "type": "image_url", + "image_url": { + "url": "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mNk+A8AAQUBAScY42YAAAAASUVORK5CYII=" + } + } + ] + }] + }) + .to_string(), + )) + .expect("build request"), + ) + .await + .expect("call app"); + + assert_eq!(response.status(), StatusCode::BAD_REQUEST); + let body = to_bytes(response.into_body(), usize::MAX).await.expect("read body"); + let json: serde_json::Value = serde_json::from_slice(&body).expect("decode json"); + assert_eq!(json["error"]["type"], "invalid_request_error"); + assert_eq!( + json["error"]["message"], + "At most 1 image(s) may be provided in one prompt." + ); +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] #[serial] async fn non_stream_chat_includes_logprobs_and_prompt_logprobs() { diff --git a/rust/src/server/src/routes/tokenize/types.rs b/rust/src/server/src/routes/tokenize/types.rs index 085f977ec8..0b86a71cb3 100644 --- a/rust/src/server/src/routes/tokenize/types.rs +++ b/rust/src/server/src/routes/tokenize/types.rs @@ -82,6 +82,7 @@ impl TokenizeChatRequest { generation_prompt_mode, chat_template: self.chat_template, reasoning_effort: None, + response_format: None, template_kwargs: self.chat_template_kwargs.unwrap_or_default(), }, tools: convert_tools(self.tools)?, diff --git a/rust/src/tokenizer/src/tiktoken.rs b/rust/src/tokenizer/src/tiktoken.rs index 7d087534d3..6d9612c093 100644 --- a/rust/src/tokenizer/src/tiktoken.rs +++ b/rust/src/tokenizer/src/tiktoken.rs @@ -532,7 +532,7 @@ fn detect_bpe_pattern(config: &TiktokenModelConfig) -> &'static str { let model_type = config.effective_model_type(); match model_type { - Some("kimi" | "kimi_k2" | "kimi_k25" | "deepseek_v3") => KIMI_PATTERN, + Some("kimi" | "kimi_k2" | "kimi_k25" | "kimi_k3" | "deepseek_v3") => KIMI_PATTERN, _ => CL100K_BASE_PATTERN, } } @@ -817,6 +817,7 @@ mod tests { #[test] fn tiktoken_detects_kimi_pattern_from_model_type() { let kimi = config_json!({ "model_type": "kimi_k25" }); + let kimi_k3 = config_json!({ "model_type": "kimi_k3" }); let baseten_kimi = config_json!({ "model_type": "deepseek_v3" }); let nested_kimi = config_json!({ "model_type": "composite_wrapper", @@ -830,6 +831,7 @@ mod tests { let missing = config_json!({ "text_config": {} }); assert_eq!(detect_bpe_pattern(&kimi), KIMI_PATTERN); + assert_eq!(detect_bpe_pattern(&kimi_k3), KIMI_PATTERN); assert_eq!(detect_bpe_pattern(&baseten_kimi), KIMI_PATTERN); assert_eq!(detect_bpe_pattern(&nested_kimi), CL100K_BASE_PATTERN); assert_eq!(detect_bpe_pattern(&generic), CL100K_BASE_PATTERN); diff --git a/rust/src/tracing/Cargo.toml b/rust/src/tracing/Cargo.toml new file mode 100644 index 0000000000..f1555ca5f7 --- /dev/null +++ b/rust/src/tracing/Cargo.toml @@ -0,0 +1,14 @@ +[package] +name = "aphrodite-tracing" +version.workspace = true +edition.workspace = true +description = "Shared tracing subscriber and log formatting for Aphrodite Rust binaries" +license.workspace = true + +[dependencies] +time.workspace = true +tracing.workspace = true +tracing-subscriber.workspace = true + +[lints] +workspace = true diff --git a/rust/src/cmd/src/logging.rs b/rust/src/tracing/src/lib.rs similarity index 98% rename from rust/src/cmd/src/logging.rs rename to rust/src/tracing/src/lib.rs index 09ec92313d..38d82165ab 100644 --- a/rust/src/cmd/src/logging.rs +++ b/rust/src/tracing/src/lib.rs @@ -1,6 +1,8 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright contributors to the vLLM project +//! Shared tracing subscriber and log formatting for Aphrodite Rust binaries. + use std::{env, fmt, process}; use time::UtcOffset; @@ -26,8 +28,8 @@ const RESET: &str = "\x1b[0m"; const APHRODITE_TIME_FORMAT: &[time::format_description::FormatItem<'static>] = format_description!("[month]-[day] [hour]:[minute]:[second]"); -/// Install the process-wide Aphrodite-style tracing subscriber for the CLI binary. -pub(crate) fn init_tracing(process_label: &str) { +/// Install the process-wide Aphrodite-style tracing subscriber. +pub fn init_tracing(process_label: &str) { let filter = build_targets_filter( env::var("APHRODITE_LOGGING_LEVEL").ok().as_deref(), env::var("RUST_LOG").ok().as_deref(), diff --git a/setup.py b/setup.py index b78fa66f04..405d764ddc 100644 --- a/setup.py +++ b/setup.py @@ -638,7 +638,7 @@ def _read_requirements(filename: str) -> list[str]: # copying the relevant .py files from the source repository. ext_modules.append(CMakeExtension(name="aphrodite.triton_kernels", optional=True)) -if sys.version_info >= (3, 11): +if not _is_xpu() and sys.version_info >= (3, 11): ext_modules.append(CMakeExtension(name="aphrodite.spinloop")) ext_modules.append(CMakeExtension(name="aphrodite.fs_io_C")) @@ -659,6 +659,8 @@ def _read_requirements(filename: str) -> list[str]: # not targeting a hopper system ext_modules.append(CMakeExtension(name="aphrodite._flashmla_C", optional=True)) ext_modules.append(CMakeExtension(name="aphrodite._flashmla_extension_C", optional=True)) + if CUDA_HOME and get_nvcc_cuda_version() >= Version("12.0"): + ext_modules.append(CMakeExtension(name="aphrodite._flashkda_C", optional=True)) if CUDA_HOME and get_nvcc_cuda_version() >= Version("12.3"): # DeepGEMM requires CUDA 12.3+ (SM90/SM100) # Optional since it won't build on unsupported architectures diff --git a/tests/compile/test_dynamic_shapes_compilation.py b/tests/compile/test_dynamic_shapes_compilation.py index a330cd2667..c75ba87ef4 100644 --- a/tests/compile/test_dynamic_shapes_compilation.py +++ b/tests/compile/test_dynamic_shapes_compilation.py @@ -1,14 +1,13 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -import gc import tempfile from contextlib import contextmanager import pytest import torch -from aphrodite import LLM, SamplingParams +from aphrodite import SamplingParams from aphrodite.compilation.decorators import support_torch_compile from aphrodite.config import AphroditeConfig, CompilationConfig, set_current_aphrodite_config from aphrodite.config.compilation import ( @@ -19,7 +18,6 @@ from aphrodite.forward_context import set_forward_context from aphrodite.utils.torch_utils import is_torch_equal_or_newer from tests.models.utils import check_logprobs_close -from tests.utils import wait_for_rocm_memory_to_settle def get_test_models(): @@ -49,6 +47,7 @@ def get_test_models(): @pytest.mark.skipif(not is_torch_equal_or_newer("2.10.0"), reason="requires torch 2.10") def test_dynamic_shapes_compilation( monkeypatch, + aphrodite_runner, model_name, shapes_type, use_aot_compile, @@ -77,9 +76,13 @@ def test_dynamic_shapes_compilation( print(f"Testing {shapes_type.name} dynamic shapes...") - # Initialize the model with specific dynamic shapes configuration - model = LLM( - model=model_name, + sampling_params = SamplingParams(max_tokens=5, temperature=0, logprobs=10) + test_prompts = [prompt, "The capital of France is"] + + # AphroditeRunner shuts down the engine core on exit, so the eager model + # below never races a lingering compiled engine for GPU memory. + with aphrodite_runner( + model_name, compilation_config={ "mode": CompilationMode.APHRODITE_COMPILE, "dynamic_shapes_config": { @@ -88,33 +91,25 @@ def test_dynamic_shapes_compilation( }, }, max_model_len=1024, - ) - - sampling_params = SamplingParams(max_tokens=5, temperature=0, logprobs=10) - test_prompts = [prompt, "The capital of France is"] - - compiled_outputs = [] - for p in test_prompts: - output = model.generate(p, sampling_params)[0].outputs[0] - assert len(output.text.strip()) > 0, "Compiled model produced empty output" - compiled_outputs.append((output.token_ids, output.text, output.logprobs)) - - del model - gc.collect() - torch.accelerator.empty_cache() - torch.accelerator.synchronize() - wait_for_rocm_memory_to_settle() - - eager_model = LLM(model=model_name, enforce_eager=True, max_model_len=1024) - eager_outputs = [] - for p in test_prompts: - output = eager_model.generate(p, sampling_params)[0].outputs[0] - assert len(output.text.strip()) > 0, "Eager model produced empty output" - eager_outputs.append((output.token_ids, output.text, output.logprobs)) - del eager_model - gc.collect() - torch.accelerator.empty_cache() - torch.accelerator.synchronize() + enable_chunked_prefill=None, + ) as aphrodite_model: + compiled_outputs = [] + for p in test_prompts: + output = aphrodite_model.llm.generate(p, sampling_params)[0].outputs[0] + assert len(output.text.strip()) > 0, "Compiled model produced empty output" + compiled_outputs.append((output.token_ids, output.text, output.logprobs)) + + with aphrodite_runner( + model_name, + enforce_eager=True, + max_model_len=1024, + enable_chunked_prefill=None, + ) as aphrodite_model: + eager_outputs = [] + for p in test_prompts: + output = aphrodite_model.llm.generate(p, sampling_params)[0].outputs[0] + assert len(output.text.strip()) > 0, "Eager model produced empty output" + eager_outputs.append((output.token_ids, output.text, output.logprobs)) check_logprobs_close( outputs_0_lst=eager_outputs, @@ -228,44 +223,35 @@ def test(model_class, input1, input2, is_01_specialization=False): @pytest.mark.skipif(not is_torch_equal_or_newer("2.10.0"), reason="requires torch 2.10") -def test_piecewise_backend_empty_sym_shape_indices(): +def test_piecewise_backend_empty_sym_shape_indices(aphrodite_runner): """Test that PiecewiseBackend handles empty sym_shape_indices correctly. When all inputs have static shapes (no torch.SymInt), sym_shape_indices will be empty. The fix in PiecewiseBackend.__call__ handles this case by using the first compiled range_entry. """ - gc.collect() - torch.accelerator.empty_cache() - torch.accelerator.synchronize() - # Use small max_model_len and max_num_batched_tokens to encourage # static shape compilation with empty sym_shape_indices - llm = LLM( - model="Qwen/Qwen3-0.6B", + with aphrodite_runner( + "Qwen/Qwen3-0.6B", max_model_len=512, max_num_batched_tokens=1, + enable_chunked_prefill=None, compilation_config={ "mode": CompilationMode.APHRODITE_COMPILE, "dynamic_shapes_config": { "type": DynamicShapesType.BACKED.value, }, }, - ) - - sampling_params = SamplingParams(temperature=0, top_p=0.95, max_tokens=10) - - # Generate with static shape inputs - output = llm.generate("Hello, my name is", sampling_params=sampling_params) - result = output[0].outputs[0].text - assert len(result) > 0, "Should generate non-empty output" - - # Generate again to verify compilation works with empty sym_shape_indices - output = llm.generate("The capital of France is", sampling_params=sampling_params) - result = output[0].outputs[0].text - assert len(result) > 0, "Should generate non-empty output on second run" - - del llm - gc.collect() - torch.accelerator.empty_cache() - torch.accelerator.synchronize() + ) as aphrodite_model: + sampling_params = SamplingParams(temperature=0, top_p=0.95, max_tokens=10) + + # Generate with static shape inputs + output = aphrodite_model.llm.generate("Hello, my name is", sampling_params=sampling_params) + result = output[0].outputs[0].text + assert len(result) > 0, "Should generate non-empty output" + + # Generate again to verify compilation works with empty sym_shape_indices + output = aphrodite_model.llm.generate("The capital of France is", sampling_params=sampling_params) + result = output[0].outputs[0].text + assert len(result) > 0, "Should generate non-empty output on second run" diff --git a/tests/config/test_speculative_draft_max_position_embeddings.py b/tests/config/test_speculative_draft_max_position_embeddings.py new file mode 100644 index 0000000000..c37f515455 --- /dev/null +++ b/tests/config/test_speculative_draft_max_position_embeddings.py @@ -0,0 +1,95 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import logging + +import pytest +from transformers import PretrainedConfig + +from aphrodite.config.model import ModelConfig +from aphrodite.config.parallel import ParallelConfig +from aphrodite.config.speculative import SpeculativeConfig + +EAGLE3_DRAFT = "yuhuili/EAGLE3-LLaMA3.1-Instruct-8B" +LLAMA3_TARGET = "unsloth/Meta-Llama-3.1-8B-Instruct" +AR_MODEL = "JackFram/llama-68m" + +_LOGGER = "aphrodite.config.speculative" +_OVERRIDE_MSG = "Overriding draft model max_position_embeddings" + + +@pytest.fixture +def aphrodite_caplog(caplog: pytest.LogCaptureFixture, monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr(logging.getLogger("aphrodite"), "propagate", True) + with caplog.at_level(logging.INFO, logger=_LOGGER): + yield caplog + + +def _override_logged(caplog: pytest.LogCaptureFixture) -> bool: + return any(_OVERRIDE_MSG in record.getMessage() for record in caplog.records) + + +@pytest.mark.cpu_test +def test_override_raises_smaller_value( + aphrodite_caplog: pytest.LogCaptureFixture, +): + hf_config = PretrainedConfig(max_position_embeddings=2048) + SpeculativeConfig._maybe_override_draft_max_position_embeddings(hf_config, target_max_model_len=8192) + assert hf_config.max_position_embeddings == 8192 + assert _override_logged(aphrodite_caplog) + + +@pytest.mark.cpu_test +def test_override_keeps_sufficient_value( + aphrodite_caplog: pytest.LogCaptureFixture, +): + hf_config = PretrainedConfig(max_position_embeddings=8192) + SpeculativeConfig._maybe_override_draft_max_position_embeddings(hf_config, target_max_model_len=8192) + assert hf_config.max_position_embeddings == 8192 + assert not _override_logged(aphrodite_caplog) + + +@pytest.mark.cpu_test +def test_override_ignores_missing_attribute( + aphrodite_caplog: pytest.LogCaptureFixture, +): + hf_config = PretrainedConfig() + hf_config.__dict__.pop("max_position_embeddings", None) + SpeculativeConfig._maybe_override_draft_max_position_embeddings(hf_config, target_max_model_len=8192) + assert not hasattr(hf_config, "max_position_embeddings") + assert not _override_logged(aphrodite_caplog) + + +@pytest.mark.cpu_test +@pytest.mark.parametrize("method", ["eagle", "eagle3"]) +def test_eagle_draft_inherits_target_max_model_len(method: str, aphrodite_caplog: pytest.LogCaptureFixture): + target_model_config = ModelConfig(LLAMA3_TARGET) + assert target_model_config.max_model_len > 2048 + speculative_config = SpeculativeConfig( + target_model_config=target_model_config, + target_parallel_config=ParallelConfig(), + model=EAGLE3_DRAFT, + method=method, + num_speculative_tokens=3, + ) + draft_hf_config = speculative_config.draft_model_config.hf_config + assert draft_hf_config.max_position_embeddings == target_model_config.max_model_len + assert _override_logged(aphrodite_caplog) + + +@pytest.mark.cpu_test +def test_independent_draft_model_keeps_its_own_limit( + aphrodite_caplog: pytest.LogCaptureFixture, +): + target_model_config = ModelConfig(AR_MODEL, hf_overrides={"max_position_embeddings": 8192}) + assert target_model_config.max_model_len == 8192 + speculative_config = SpeculativeConfig( + target_model_config=target_model_config, + target_parallel_config=ParallelConfig(), + model=AR_MODEL, + method="draft_model", + num_speculative_tokens=3, + ) + draft_hf_config = speculative_config.draft_model_config.hf_config + assert draft_hf_config.max_position_embeddings == 2048 + assert not _override_logged(aphrodite_caplog) diff --git a/tests/distributed/test_custom_all_reduce.py b/tests/distributed/test_custom_all_reduce.py index 504e0b1cc6..27a84cfb66 100644 --- a/tests/distributed/test_custom_all_reduce.py +++ b/tests/distributed/test_custom_all_reduce.py @@ -2,6 +2,8 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import random +from types import SimpleNamespace +from unittest.mock import Mock import pytest import ray @@ -9,6 +11,8 @@ import torch.distributed as dist from aphrodite.distributed.communication_op import tensor_model_parallel_all_reduce # noqa +from aphrodite.distributed.device_communicators import custom_all_reduce +from aphrodite.distributed.device_communicators.custom_all_reduce import CustomAllreduce from aphrodite.distributed.parallel_state import get_tp_group, graph_capture from ..utils import ( @@ -23,6 +27,55 @@ test_sizes[i] -= v % 8 +def test_sp16_dispatches_only_to_mnnvl_lamport( + monkeypatch: pytest.MonkeyPatch, +): + """SP16 uses the MNNVL Lamport kernels and rejects same-host dispatch.""" + comm = CustomAllreduce.__new__(CustomAllreduce) + comm.disabled = False + comm.world_size = 16 + comm.fully_connected = False + comm.mnnvl_only = True + comm._IS_CAPTURING = False + comm.max_mnnvl_all_gather_size = 2 * 1024 * 1024 + comm.max_mnnvl_reduce_scatter_size = 16 * 1024 * 1024 + comm.mnnvl_multicast_ptr = 1 + comm.mnnvl_lamport_ag_local_ptr = 1 + comm.mnnvl_lamport_ag_multicast_ptr = 1 + comm.mnnvl_lamport_ag_epoch_ptr = 1 + comm.mnnvl_lamport_rs_local_ptr = 1 + comm.mnnvl_lamport_rs_epoch_ptr = 1 + comm.mnnvl_buffer_size = 32 * 1024 * 1024 + comm._ptr = 0 + + lamport_all_gather = Mock() + lamport_reduce_scatter = Mock() + monkeypatch.setattr(custom_all_reduce.current_platform, "is_cuda", lambda: True) + monkeypatch.setattr( + custom_all_reduce, + "ops", + SimpleNamespace( + mnnvl_lamport_all_gather=lamport_all_gather, + mnnvl_lamport_reduce_scatter=lamport_reduce_scatter, + ), + ) + + gathered = comm.custom_all_gather(torch.empty((8, 8), dtype=torch.bfloat16)) + scattered = comm.custom_reduce_scatter(torch.empty((16, 8), dtype=torch.bfloat16)) + + assert gathered is not None + assert scattered is not None + lamport_all_gather.assert_called_once() + lamport_reduce_scatter.assert_called_once() + assert not comm.should_custom_ar(torch.empty(8, dtype=torch.bfloat16)) + assert not comm.should_custom_all_gather(torch.empty((8, 8), dtype=torch.int32)) + assert not comm.should_custom_all_gather(torch.empty((131073, 8), dtype=torch.bfloat16)) + + comm.mnnvl_only = False + assert not comm.should_custom_all_gather(torch.empty((8, 8), dtype=torch.bfloat16)) + assert not comm.should_custom_reduce_scatter(torch.empty((16, 8), dtype=torch.bfloat16)) + + @ray.remote(num_gpus=1, max_calls=1) def graph_allreduce( monkeypatch: pytest.MonkeyPatch, @@ -80,6 +133,28 @@ def graph_allreduce( torch.testing.assert_close(out1, inp1) torch.testing.assert_close(out2, inp2) + fa = get_tp_group().device_communicator.ca_comm + tp_rank = rank % tp_size + with graph_capture(device=device) as graph_capture_context: + local = torch.full((512, 4096), tp_rank + 1, dtype=torch.bfloat16, device=device) + reduce_input = torch.full( + (512 * tp_size, 4096), + tp_rank + 1, + dtype=torch.bfloat16, + device=device, + ) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, stream=graph_capture_context.stream): + gathered = fa.custom_all_gather(local) + scattered = fa.custom_reduce_scatter(reduce_input) + graph.replay() + assert gathered is not None + assert scattered is not None + expected_gather = torch.cat([torch.full_like(local, peer_rank + 1) for peer_rank in range(tp_size)]) + expected_scatter = torch.full_like(local, tp_size * (tp_size + 1) // 2) + torch.testing.assert_close(gathered, expected_gather) + torch.testing.assert_close(scattered, expected_scatter) + @ray.remote(num_gpus=1, max_calls=1) def eager_allreduce( @@ -110,6 +185,23 @@ def eager_allreduce( out = fa.all_reduce(out, registered=False) torch.testing.assert_close(out, inp * (tp_size**num_communication)) + group = get_tp_group().device_group + tp_rank = rank % tp_size + for dtype in [torch.float32, torch.float16, torch.bfloat16]: + local = torch.full((64, 4096), tp_rank + 1, dtype=dtype, device=device) + expected_gather = torch.empty((64 * tp_size, 4096), dtype=dtype, device=device) + dist.all_gather_into_tensor(expected_gather, local, group=group) + gathered = fa.custom_all_gather(local) + assert gathered is not None + torch.testing.assert_close(gathered, expected_gather) + + reduce_input = torch.full((64 * tp_size, 4096), tp_rank + 1, dtype=dtype, device=device) + expected_scatter = torch.empty((64, 4096), dtype=dtype, device=device) + dist.reduce_scatter_tensor(expected_scatter, reduce_input.clone(), group=group) + scattered = fa.custom_reduce_scatter(reduce_input) + assert scattered is not None + torch.testing.assert_close(scattered, expected_scatter) + inp = torch.ones(sz * 4, dtype=torch.bfloat16, device=device) out = inp for _ in range(num_communication): @@ -130,3 +222,14 @@ def test_custom_allreduce( if world_size > torch.accelerator.device_count(): pytest.skip("Not enough GPUs to run the test.") multi_process_parallel(monkeypatch, tp_size, pipeline_parallel_size, test_target) + + +@pytest.mark.parametrize("test_target", [eager_allreduce, graph_allreduce]) +def test_custom_collectives_world_size_four( + monkeypatch: pytest.MonkeyPatch, + test_target, +): + """Exercise the four-rank kernel specialization used by Kimi SP.""" + if torch.accelerator.device_count() < 4: + pytest.skip("Not enough GPUs to run the test.") + multi_process_parallel(monkeypatch, 4, 1, test_target) diff --git a/tests/distributed/test_shm_broadcast.py b/tests/distributed/test_shm_broadcast.py index 0e5baece29..aa54d6d695 100644 --- a/tests/distributed/test_shm_broadcast.py +++ b/tests/distributed/test_shm_broadcast.py @@ -1,6 +1,8 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import io +import pickle import random import threading import time @@ -10,12 +12,15 @@ import multiprocess as mp import numpy as np import pytest +import torch import torch.distributed as dist from aphrodite.distributed.device_communicators import shm_broadcast from aphrodite.distributed.device_communicators.shm_broadcast import ( MessageQueue, ShmRingBuffer, + _rebuild_tensor, + _reduce_tensor, check_shm_free_space, ) from aphrodite.distributed.utils import StatelessProcessGroup @@ -340,6 +345,144 @@ def test_message_queue_busy_to_idle(): distributed_run(worker_fn_test_busy_to_idle, 4) +@worker_fn_wrapper +def worker_fn_tensor_broadcast(): + rank = dist.get_rank() + writer_rank = 0 + message_queue = MessageQueue.create_from_process_group(dist.group.WORLD, 8 * 1024 * 1024, 4, writer_rank) + + # Both ranks construct the identical reference payload. + torch.manual_seed(42) + payload = { + # 2MiB: rides the shm ring as an out-of-band buffer (the receiving + # side must copy out of the reusable ring chunk). + "mid": torch.randn(1024, 512), + # 16MiB > max_chunk_bytes: overflows to the zmq socket (the + # receiving side aliases the zmq.Frame zero-copy). + "big": torch.randn(4096, 2048, dtype=torch.bfloat16), + "nested": ["plain", 123, {"inner": torch.arange(5)}], + } + + if rank == writer_rank: + with mock.patch( + "aphrodite.distributed.device_communicators.shm_broadcast._reduce_tensor", + wraps=_reduce_tensor, + ) as wrapped_reduce: + message_queue.enqueue(payload) + assert wrapped_reduce.call_count == 3 + # Cycle the ring (max_chunks=4) several times over so that aliased + # ring chunks would be overwritten. + for i in range(16): + message_queue.enqueue({"junk": torch.full((1024, 512), float(i))}) + else: + received = message_queue.dequeue(timeout=30) + for key in ("mid", "big"): + assert torch.equal(received[key], payload[key]), key + assert received[key].dtype == payload[key].dtype, key + assert torch.equal(received["nested"][2]["inner"], torch.arange(5)) + + snapshot = received["mid"].clone() + for i in range(16): + junk = message_queue.dequeue(timeout=30) + assert torch.equal(junk["junk"], torch.full((1024, 512), float(i))) + # Tensors received via the shm ring must not alias chunk memory + # that the writer has reused for subsequent messages. + assert torch.equal(received["mid"], snapshot) + # Rebuilt tensors must be writable, like regular tensors. + received["mid"] += 1.0 + received["big"][0, 0] = 1.0 + + dist.barrier() + print(f"tensor broadcast passed the test! Rank {rank}") + + +def test_tensor_broadcast(): + distributed_run(worker_fn_tensor_broadcast, 2) + + +def _dumps_oob(obj) -> tuple[bytes, list]: + """Pickle `obj` like `MessageQueue.enqueue`: tensor dispatch + OOB buffers.""" + buffers = [] + + def callback(buf: pickle.PickleBuffer) -> bool: + raw = buf.raw() + if raw.nbytes < 1024 * 1024: + return True + buffers.append(raw) + return False + + bio = io.BytesIO() + pickler = pickle.Pickler(bio, protocol=pickle.HIGHEST_PROTOCOL, buffer_callback=callback) + pickler.dispatch_table = {torch.Tensor: _reduce_tensor} + pickler.dump(obj) + return bio.getvalue(), buffers + + +@pytest.mark.parametrize( + "case", + [ + "small", + "mid", + "bf16", + "fp8", + "empty", + "scalar", + "noncontig", + "requires_grad", + "conj", + "param", + ], +) +def test_tensor_pickle_roundtrip(case: str): + tensor = { + # Inlined in-band (< 1MiB) and out-of-band (>= 1MiB) buffers. + "small": lambda: torch.randn(100, 10), + "mid": lambda: torch.randn(1024, 512), + # Dtypes numpy doesn't recognize. + "bf16": lambda: torch.randn(512, 512, dtype=torch.bfloat16), + "fp8": lambda: torch.randn(32, 32).to(torch.float8_e4m3fn), + # Shape edge cases. + "empty": lambda: torch.empty(0, 8), + "scalar": lambda: torch.tensor(3.14), + "noncontig": lambda: torch.randn(64, 64).t(), + # These fall back to torch's default reducer. + "requires_grad": lambda: torch.randn(8, 8, requires_grad=True), + "conj": lambda: torch.randn(4, dtype=torch.complex64).conj(), + "param": lambda: torch.nn.Parameter(torch.randn(4), requires_grad=False), + }[case]() + + data, buffers = _dumps_oob({"tensor": tensor, "meta": list(range(10))}) + received = pickle.loads(data, buffers=buffers)["tensor"] + + assert received.shape == tensor.shape + assert received.dtype == tensor.dtype + if tensor.dtype == torch.float8_e4m3fn: + assert torch.equal(received.view(torch.uint8), tensor.view(torch.uint8)) + else: + assert torch.equal(received, tensor) + assert received.requires_grad == tensor.requires_grad + assert isinstance(received, type(tensor)) + if tensor.numel() and not tensor.requires_grad: + # Rebuilt tensors must be writable, like regular tensors. + received.view(-1)[0] = 1.0 + + +@pytest.mark.parametrize("case", ["cuda", "requires_grad", "conj"]) +def test_reduce_tensor_fallback(case: str): + """Tensors without safe aliases must fall back to torch's default reduction.""" + if case == "cuda": + if not torch.cuda.is_available(): + pytest.skip("requires CUDA") + tensor = torch.randn(4, device="cuda") + elif case == "requires_grad": + tensor = torch.randn(8, requires_grad=True) + else: + tensor = torch.randn(4, dtype=torch.complex64).conj() + + reduced = _reduce_tensor(tensor) + assert reduced[0] is not _rebuild_tensor + + @pytest.mark.parametrize("should_warn", [False, True]) def test_reader_timeout_caps_indefinite_waits(should_warn): with ( diff --git a/tests/entrypoints/multimodal/openai/chat_completion/test_audio.py b/tests/entrypoints/multimodal/openai/chat_completion/test_audio.py index 262a49ba39..c420b55230 100644 --- a/tests/entrypoints/multimodal/openai/chat_completion/test_audio.py +++ b/tests/entrypoints/multimodal/openai/chat_completion/test_audio.py @@ -9,6 +9,7 @@ from aphrodite.assets.audio import AudioAsset from aphrodite.multimodal.utils import encode_audio_base64, encode_audio_url, fetch_audio +from aphrodite.platforms import current_platform from tests.utils import RemoteOpenAIServer MODEL_NAME = "fixie-ai/ultravox-v0_5-llama-3_2-1b" @@ -18,6 +19,10 @@ ] MAXIMUM_AUDIOS = 2 +# Disable prefix caching on ROCm to reduce non-determinism in +# streaming-vs-non-streaming comparisons. +_ROCM_ARGS = ["--no-enable-prefix-caching"] if current_platform.is_rocm() else [] + @pytest.fixture(scope="module") def server(): @@ -32,6 +37,7 @@ def server(): "--trust-remote-code", "--limit-mm-per-prompt", json.dumps({"audio": MAXIMUM_AUDIOS}), + *_ROCM_ARGS, ] with RemoteOpenAIServer(MODEL_NAME, args) as remote_server: diff --git a/tests/entrypoints/openai/chat_completion/test_chat_completion.py b/tests/entrypoints/openai/chat_completion/test_chat_completion.py index d8be15e2e2..e17791549f 100644 --- a/tests/entrypoints/openai/chat_completion/test_chat_completion.py +++ b/tests/entrypoints/openai/chat_completion/test_chat_completion.py @@ -150,3 +150,92 @@ async def test_empty_grammar(client: openai.AsyncOpenAI, model_name: str) -> Non ], extra_body={"structured_outputs": {"grammar": ""}}, ) + + +# Decode-side token reuse for disaggregated serving. The router forwards the +# prefill stage's prompt token ids in kv_transfer_params so the decode stage +# skips re-tokenizing. + +TOKEN_IN_MESSAGES = [{"role": "user", "content": "Hello, how are you today?"}] +DECODE_MESSAGES = [{"role": "user", "content": "unrelated decode-side text"}] + + +@pytest.mark.asyncio +async def test_kv_transfer_prompt_token_ids_round_trip(client: openai.AsyncOpenAI): + """Ids forwarded in kv_transfer_params are used verbatim, skipping tokenize. + + The decode request carries different messages, so a response whose + prompt_token_ids match the forwarded ids proves the ids were used rather + than the request's own messages. Generated text is not compared across + requests because Aphrodite greedy decoding is not bitwise-reproducible. + """ + baseline = await client.chat.completions.create( + model=MODEL_NAME, + messages=TOKEN_IN_MESSAGES, + max_completion_tokens=16, + temperature=0, + extra_body={"return_token_ids": True}, + ) + reused_ids = baseline.prompt_token_ids + assert reused_ids + + decode = await client.chat.completions.create( + model=MODEL_NAME, + messages=DECODE_MESSAGES, + max_completion_tokens=16, + temperature=0, + extra_body={ + "kv_transfer_params": {"prompt_token_ids": reused_ids}, + "return_token_ids": True, + }, + ) + + # The engine saw the forwarded ids, not the decode request's own messages. + assert decode.prompt_token_ids == reused_ids + # text-out: reuse still yields a detokenized message. + assert decode.choices[0].message.content + + +@pytest.mark.asyncio +async def test_kv_transfer_prompt_token_ids_streaming(client: openai.AsyncOpenAI): + """Decode-side token reuse streams chat-formatted text-out.""" + baseline = await client.chat.completions.create( + model=MODEL_NAME, + messages=TOKEN_IN_MESSAGES, + max_completion_tokens=16, + temperature=0, + extra_body={"return_token_ids": True}, + ) + reused_ids = baseline.prompt_token_ids + assert reused_ids + + stream = await client.chat.completions.create( + model=MODEL_NAME, + messages=DECODE_MESSAGES, + max_completion_tokens=16, + temperature=0, + stream=True, + extra_body={ + "kv_transfer_params": {"prompt_token_ids": reused_ids}, + "return_token_ids": True, + }, + ) + + content = "" + delta_token_ids: list[int] = [] + first_chunk = True + async for chunk in stream: + if first_chunk: + # prompt_token_ids arrives once, on the first chunk. + assert chunk.prompt_token_ids == reused_ids + first_chunk = False + if not chunk.choices: + continue + if chunk.choices[0].delta.content: + content += chunk.choices[0].delta.content + if tids := getattr(chunk.choices[0], "token_ids", None): + delta_token_ids.extend(tids) + + # streamed text-out, reconstructed from deltas, with generated token ids. + assert content + assert delta_token_ids diff --git a/tests/entrypoints/openai/chat_completion/test_serving_chat.py b/tests/entrypoints/openai/chat_completion/test_serving_chat.py index d785c448d6..2842e4374b 100644 --- a/tests/entrypoints/openai/chat_completion/test_serving_chat.py +++ b/tests/entrypoints/openai/chat_completion/test_serving_chat.py @@ -2231,3 +2231,34 @@ async def result_generator(): assert reasons[0] == "tool_calls", ( f"Choice {choice_idx}: expected finish_reason='tool_calls', got '{reasons[0]}'" ) + + +def test_make_request_with_harmony_reuses_kv_transfer_prompt_token_ids(): + """The Harmony reuse branch honors ids forwarded in kv_transfer_params. + + A GPT-OSS server is impractical to stand up here, so this exercises the + branch directly on a harmony-configured renderer. + """ + engine = MockEngine() + engine.model_config.hf_config = MockHFConfig(model_type="gpt_oss") + models = OpenAIServingModels(engine, BASE_MODEL_PATHS) + online_renderer = _build_online_renderer(engine, models.registry) + assert online_renderer.use_harmony + + request = ChatCompletionRequest( + model=MODEL_NAME, + messages=[{"role": "user", "content": "hi"}], + kv_transfer_params={ + "prompt_token_ids": [10, 20, 30], + "do_remote_prefill": True, + }, + ) + conversation, engine_inputs = online_renderer._make_request_with_harmony(request) + + assert conversation == [] + assert len(engine_inputs) == 1 + engine_input = engine_inputs[0] + assert engine_input["type"] == "token" + assert engine_input["prompt_token_ids"] == [10, 20, 30] + # The reuse key is consumed and other kv_transfer_params are preserved. + assert request.kv_transfer_params == {"do_remote_prefill": True} diff --git a/tests/evals/gpt_oss/configs/gpt-oss-20b-xpu-baseline.yaml b/tests/evals/gpt_oss/configs/gpt-oss-20b-xpu-baseline.yaml new file mode 100644 index 0000000000..78a583888b --- /dev/null +++ b/tests/evals/gpt_oss/configs/gpt-oss-20b-xpu-baseline.yaml @@ -0,0 +1,5 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +model_name: openai/gpt-oss-20b +metric_threshold: 0.568 +reasoning_effort: low diff --git a/tests/evals/gpt_oss/configs/gpt-oss-20b-xpu-triton-attn.yaml b/tests/evals/gpt_oss/configs/gpt-oss-20b-xpu-triton-attn.yaml new file mode 100644 index 0000000000..e711ffb331 --- /dev/null +++ b/tests/evals/gpt_oss/configs/gpt-oss-20b-xpu-triton-attn.yaml @@ -0,0 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +model_name: openai/gpt-oss-20b +metric_threshold: 0.568 +reasoning_effort: low +server_args: "--attention-backend TRITON_ATTN" diff --git a/tests/evals/gpt_oss/configs/models-xpu.txt b/tests/evals/gpt_oss/configs/models-xpu.txt new file mode 100644 index 0000000000..a9de9266b1 --- /dev/null +++ b/tests/evals/gpt_oss/configs/models-xpu.txt @@ -0,0 +1,3 @@ +# Intel XPU model configurations for GPQA evaluation +gpt-oss-20b-xpu-baseline.yaml +gpt-oss-20b-xpu-triton-attn.yaml diff --git a/tests/kernels/attention/test_cpu_attn.py b/tests/kernels/attention/test_cpu_attn.py index a61e7317f1..e63f14724b 100644 --- a/tests/kernels/attention/test_cpu_attn.py +++ b/tests/kernels/attention/test_cpu_attn.py @@ -498,6 +498,7 @@ def varlen_with_paged_kv( isa=isa, enable_kv_split=False, dynamic_causal=dynamic_causal_tensor, + kv_cache_dtype=kv_cache_dtype, ) out_without_split = torch.empty_like(query) @@ -533,6 +534,7 @@ def varlen_with_paged_kv( isa=isa, enable_kv_split=True, dynamic_causal=dynamic_causal_tensor, + kv_cache_dtype=kv_cache_dtype, ) out_with_split = torch.empty_like(query) @@ -765,6 +767,24 @@ def test_varlen_with_paged_kv_normal_amx( ) +@pytest.mark.skipif(not torch.cpu._is_amx_tile_supported(), reason="no AMX support.") +def test_varlen_with_paged_kv_fp8_large_prefill_amx() -> None: + varlen_with_paged_kv( + seq_lens=[(1024, 1024)] * 4, + num_heads=(16, 2), + head_size=256, + sliding_window=None, + dtype=torch.bfloat16, + block_size=2176, + soft_cap=None, + num_blocks=4, + use_alibi=False, + use_sink=False, + isa="amx", + kv_cache_dtype="fp8_e4m3", + ) + + @pytest.mark.parametrize("seq_lens", SEQ_LENS) @pytest.mark.parametrize("num_heads", NUM_HEADS) @pytest.mark.parametrize("head_size", HEAD_SIZES_VEC16) diff --git a/tests/kernels/attention/test_kimi_k3_mla_fused_epilogue.py b/tests/kernels/attention/test_kimi_k3_mla_fused_epilogue.py new file mode 100644 index 0000000000..34139fc2f7 --- /dev/null +++ b/tests/kernels/attention/test_kimi_k3_mla_fused_epilogue.py @@ -0,0 +1,201 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""RoPE equivalence tests for the fused Kimi-K3 MLA epilogues.""" + +import pytest +import torch + +from aphrodite.models.kimi_k3.nvidia.ops.fused_mla_key_concat_kv_cache import ( + fused_mla_decode_q_concat_kv_cache_insert, + fused_mla_key_concat_ds_mla_insert, + fused_mla_key_concat_kv_cache_insert, + fused_mla_qkv_quant_kv_cache_fp8_insert, +) +from aphrodite.platforms import current_platform + +pytestmark = pytest.mark.skipif(not current_platform.is_cuda(), reason="Kimi-K3 fused MLA requires CUDA") + +_DTYPE = torch.bfloat16 +_NUM_TOKENS = 3 +_NUM_HEADS = 4 +_BLOCK_SIZE = 8 +_POSITIONS = (1, 7, 13) +_SLOTS = (0, 3, 9) + + +def _randn(*shape: int) -> torch.Tensor: + return torch.randn(*shape, device="cuda", dtype=_DTYPE) * 0.2 + + +def _rope_cache(max_position: int = 32) -> torch.Tensor: + inv_freq = 1.0 / (50000 ** (torch.arange(0, 64, 2, dtype=torch.float32, device="cuda") / 64)) + positions = torch.arange(max_position, dtype=torch.float32, device="cuda") + freqs = torch.outer(positions, inv_freq) + # The fused epilogue reads the cos/sin table in fp32 (RoPE math runs in fp32). + return torch.cat((freqs.cos(), freqs.sin()), dim=-1) + + +def _apply_gptj_rope(x: torch.Tensor, positions: torch.Tensor, cos_sin_cache: torch.Tensor) -> torch.Tensor: + cos, sin = cos_sin_cache.index_select(0, positions).chunk(2, dim=-1) + for _ in range(x.ndim - 2): + cos = cos.unsqueeze(1) + sin = sin.unsqueeze(1) + x1 = x[..., ::2].float() + x2 = x[..., 1::2].float() + out1 = x1 * cos.float() - x2 * sin.float() + out2 = x2 * cos.float() + x1 * sin.float() + return torch.stack((out1, out2), dim=-1).flatten(-2).to(x.dtype) + + +def _cache_rows(cache: torch.Tensor, slots: torch.Tensor) -> torch.Tensor: + return cache.reshape(-1, cache.shape[-1]).index_select(0, slots) + + +def _assert_fp8_close(actual: torch.Tensor, expected: torch.Tensor) -> None: + torch.testing.assert_close( + actual.float(), + expected.to(torch.float8_e4m3fn).float(), + atol=0.03125, + rtol=0.15, + ) + + +@pytest.mark.parametrize("cache_kind", ["bf16", "fp8", "fp8_ds_mla"]) +@torch.inference_mode() +def test_prefill_epilogue_fuses_gptj_rope(cache_kind: str) -> None: + torch.manual_seed(0) + positions = torch.tensor(_POSITIONS, device="cuda", dtype=torch.int64) + slots = torch.tensor(_SLOTS, device="cuda", dtype=torch.int64) + cos_sin_cache = _rope_cache() + q = _randn(_NUM_TOKENS, _NUM_HEADS, 192) + k_nope = _randn(_NUM_TOKENS, _NUM_HEADS, 128) + k_pe = _randn(_NUM_TOKENS, 64) + kv_c = _randn(_NUM_TOKENS, 512) + v = _randn(_NUM_TOKENS, _NUM_HEADS, 128) + + q_expected = q.clone() + q_expected[..., 128:] = _apply_gptj_rope(q_expected[..., 128:], positions, cos_sin_cache) + k_pe_expected = _apply_gptj_rope(k_pe, positions, cos_sin_cache) + k_expected = torch.cat((k_nope, k_pe_expected[:, None, :].expand(-1, _NUM_HEADS, -1)), dim=-1) + cache_expected = torch.cat((kv_c, k_pe_expected), dim=-1) + + if cache_kind == "bf16": + cache = torch.zeros(2, _BLOCK_SIZE, 576, device="cuda", dtype=_DTYPE) + q_actual = q.clone() + k_actual = fused_mla_key_concat_kv_cache_insert( + q_actual, + k_nope, + k_pe, + kv_c, + cache, + slots, + positions, + cos_sin_cache, + ) + torch.testing.assert_close(q_actual, q_expected) + torch.testing.assert_close(k_actual, k_expected) + torch.testing.assert_close(_cache_rows(cache, slots), cache_expected) + elif cache_kind == "fp8": + cache = torch.zeros(2, _BLOCK_SIZE, 576, device="cuda", dtype=torch.float8_e4m3fn) + one = torch.ones(1, device="cuda", dtype=torch.float32) + q_actual, k_actual, v_actual = fused_mla_qkv_quant_kv_cache_fp8_insert( + q, + k_nope, + k_pe, + kv_c, + v, + cache, + slots, + one, + one, + one, + one, + positions, + cos_sin_cache, + ) + _assert_fp8_close(q_actual, q_expected) + _assert_fp8_close(k_actual, k_expected) + _assert_fp8_close(v_actual, v) + _assert_fp8_close(_cache_rows(cache, slots), cache_expected) + else: + cache = torch.zeros(2, _BLOCK_SIZE, 656, device="cuda", dtype=torch.uint8) + q_actual = q.clone() + k_actual = fused_mla_key_concat_ds_mla_insert( + q_actual, + k_nope, + k_pe, + kv_c, + cache, + slots, + positions, + cos_sin_cache, + ) + rope_cache = _cache_rows(cache, slots)[:, 528:656].view(_DTYPE) + torch.testing.assert_close(q_actual, q_expected) + torch.testing.assert_close(k_actual, k_expected) + torch.testing.assert_close(rope_cache, k_pe_expected) + + +@pytest.mark.parametrize("cache_kind", ["bf16", "fp8", "fp8_ds_mla"]) +@torch.inference_mode() +def test_decode_epilogue_fuses_gptj_rope(cache_kind: str) -> None: + torch.manual_seed(1) + positions = torch.tensor(_POSITIONS, device="cuda", dtype=torch.int64) + slots = torch.tensor(_SLOTS, device="cuda", dtype=torch.int64) + cos_sin_cache = _rope_cache() + ql_nope = _randn(_NUM_TOKENS, _NUM_HEADS, 512) + q_pe = _randn(_NUM_TOKENS, _NUM_HEADS, 64) + kv_c = _randn(_NUM_TOKENS, 512) + k_pe = _randn(_NUM_TOKENS, 64) + + q_pe_expected = _apply_gptj_rope(q_pe, positions, cos_sin_cache) + k_pe_expected = _apply_gptj_rope(k_pe, positions, cos_sin_cache) + q_expected = torch.cat((ql_nope, q_pe_expected), dim=-1) + cache_expected = torch.cat((kv_c, k_pe_expected), dim=-1) + + kwargs = {"positions": positions, "cos_sin_cache": cos_sin_cache} + if cache_kind == "bf16": + cache = torch.zeros(2, _BLOCK_SIZE, 576, device="cuda", dtype=_DTYPE) + q_actual = fused_mla_decode_q_concat_kv_cache_insert(ql_nope, q_pe, kv_c, k_pe, cache, slots, **kwargs) + torch.testing.assert_close(q_actual, q_expected) + torch.testing.assert_close(_cache_rows(cache, slots), cache_expected) + elif cache_kind == "fp8": + cache = torch.zeros(2, _BLOCK_SIZE, 576, device="cuda", dtype=torch.float8_e4m3fn) + one = torch.ones(1, device="cuda", dtype=torch.float32) + q_actual = fused_mla_decode_q_concat_kv_cache_insert( + ql_nope, + q_pe, + kv_c, + k_pe, + cache, + slots, + q_scale_inv=one, + cache_scale_inv=one, + **kwargs, + ) + _assert_fp8_close(q_actual, q_expected) + _assert_fp8_close(_cache_rows(cache, slots), cache_expected) + else: + cache = torch.zeros(2, _BLOCK_SIZE, 656, device="cuda", dtype=torch.uint8) + q_actual = fused_mla_decode_q_concat_kv_cache_insert( + ql_nope, q_pe, kv_c, k_pe, cache, slots, ds_mla=True, **kwargs + ) + rope_cache = _cache_rows(cache, slots)[:, 528:656].view(_DTYPE) + torch.testing.assert_close(q_actual, q_expected) + torch.testing.assert_close(rope_cache, k_pe_expected) + + +@torch.inference_mode() +def test_decode_epilogue_preserves_nope_path() -> None: + torch.manual_seed(2) + slots = torch.tensor(_SLOTS, device="cuda", dtype=torch.int64) + ql_nope = _randn(_NUM_TOKENS, _NUM_HEADS, 512) + q_pe = _randn(_NUM_TOKENS, _NUM_HEADS, 64) + kv_c = _randn(_NUM_TOKENS, 512) + k_pe = _randn(_NUM_TOKENS, 64) + cache = torch.zeros(2, _BLOCK_SIZE, 576, device="cuda", dtype=_DTYPE) + + q_actual = fused_mla_decode_q_concat_kv_cache_insert(ql_nope, q_pe, kv_c, k_pe, cache, slots) + + torch.testing.assert_close(q_actual, torch.cat((ql_nope, q_pe), dim=-1)) + torch.testing.assert_close(_cache_rows(cache, slots), torch.cat((kv_c, k_pe), dim=-1)) diff --git a/tests/kernels/attention/test_rocm_triton_attn_dsv4.py b/tests/kernels/attention/test_rocm_triton_attn_dsv4.py index 2966846100..9a571cb09e 100644 --- a/tests/kernels/attention/test_rocm_triton_attn_dsv4.py +++ b/tests/kernels/attention/test_rocm_triton_attn_dsv4.py @@ -1,6 +1,8 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from types import SimpleNamespace + import pytest import torch @@ -175,6 +177,63 @@ def _ragged_from_rows(rows: list[list[int]], device: torch.device) -> tuple[torc ) +@torch.inference_mode() +def test_paged_mqa_logits_do_not_contain_nan(monkeypatch) -> None: + from aphrodite._aiter_ops import rocm_aiter_ops + from aphrodite.v1.attention.ops import rocm_aiter_mla_sparse as mod + + device = torch.device("cuda") + + class FakeWorkspaceManager: + def get_simultaneous(self, *shapes_and_dtypes): + return [torch.empty(shape, dtype=dtype, device=device) for shape, dtype in shapes_and_dtypes] + + def fake_paged_mqa_logits( + q_fp8, + kv_cache_fp8, + weights, + out_logits, + context_lens, + block_tables, + max_seq_len, + **kwargs, + ): + del ( + q_fp8, + kv_cache_fp8, + weights, + context_lens, + block_tables, + max_seq_len, + kwargs, + ) + out_logits.fill_(float("nan")) + + monkeypatch.setattr(mod, "_ON_GFX942", False) + monkeypatch.setattr(mod, "_ON_GFX950", True) + monkeypatch.setattr(rocm_aiter_ops, "is_enabled", lambda: True) + monkeypatch.setattr( + mod, + "paged_mqa_logits_module", + lambda: SimpleNamespace(deepgemm_fp8_paged_mqa_logits=fake_paged_mqa_logits), + ) + monkeypatch.setattr(mod, "current_workspace_manager", lambda: FakeWorkspaceManager()) + + q_fp8 = torch.empty((1, 1, 1, 1), dtype=torch.uint8, device=device) + kv_cache_fp8 = torch.empty((1, 1, 1, 5), dtype=torch.uint8, device=device) + logits = mod.rocm_fp8_paged_mqa_logits( + q_fp8, + kv_cache_fp8, + torch.empty((1, 1), dtype=torch.float32, device=device), + torch.ones(1, dtype=torch.int32, device=device), + torch.zeros((1, 1), dtype=torch.int32, device=device), + torch.empty(0, dtype=torch.int32, device=device), + 1, + ) + + assert not torch.isnan(logits).any() + + @torch.inference_mode() def test_compute_global_topk_ragged_indices_and_indptr() -> None: from aphrodite.models.deepseek_v4.amd.rocm import ( diff --git a/tests/kernels/core/test_activation.py b/tests/kernels/core/test_activation.py index eab3d942a6..6c2595885d 100644 --- a/tests/kernels/core/test_activation.py +++ b/tests/kernels/core/test_activation.py @@ -185,6 +185,44 @@ def test_silu_and_mul_with_clamp( opcheck(torch.ops._C.silu_and_mul_with_clamp, (out_buf, x, swiglu_limit)) +@pytest.mark.parametrize("linear_beta", [-1.0, 2.0]) +@pytest.mark.parametrize("dtype", [torch.half, torch.bfloat16]) +@torch.inference_mode() +def test_masked_situ_and_mul( + default_vllm_config, + linear_beta: float, + dtype: torch.dtype, +) -> None: + """Masked SITU computes valid expert rows and preserves padded zeros.""" + device = CUDA_DEVICES[0] + num_experts, max_num_tokens, d = 4, 7, 512 + beta = 1.5 + input = torch.randn(num_experts, max_num_tokens, 2 * d, dtype=dtype, device=device) + expert_num_tokens = torch.tensor([0, 1, 4, 7], dtype=torch.int32, device=device) + output = torch.zeros(num_experts, max_num_tokens, d, dtype=dtype, device=device) + + torch.ops._C.masked_situ_and_mul(output, input, expert_num_tokens, beta, linear_beta) + + gate, up = input.float().chunk(2, dim=-1) + expected = beta * torch.tanh(gate / beta) * torch.sigmoid(gate) + if linear_beta > 0: + up = linear_beta * torch.tanh(up / linear_beta) + expected = (expected * up).to(dtype) + for expert, num_tokens in enumerate(expert_num_tokens.cpu().tolist()): + torch.testing.assert_close( + output[expert, :num_tokens], + expected[expert, :num_tokens], + atol=get_default_atol(output), + rtol=get_default_rtol(output), + ) + assert torch.count_nonzero(output[expert, num_tokens:]) == 0 + + opcheck( + torch.ops._C.masked_situ_and_mul, + (output, input, expert_num_tokens, beta, linear_beta), + ) + + @pytest.mark.parametrize( "activation", [ diff --git a/tests/kernels/core/test_fused_q_kv_rmsnorm.py b/tests/kernels/core/test_fused_q_kv_rmsnorm.py index 840d8f816a..8bdf5c5a51 100644 --- a/tests/kernels/core/test_fused_q_kv_rmsnorm.py +++ b/tests/kernels/core/test_fused_q_kv_rmsnorm.py @@ -13,7 +13,7 @@ import pytest import torch -from aphrodite.models.deepseek_v4.common.ops import fused_q_kv_rmsnorm +from aphrodite.models.common.ops import fused_q_kv_rmsnorm from aphrodite.platforms import current_platform pytestmark = pytest.mark.skipif( diff --git a/tests/kernels/core/test_fused_rms_norm_gated.py b/tests/kernels/core/test_fused_rms_norm_gated.py index 2fdf17afa5..5f6e1b05b5 100644 --- a/tests/kernels/core/test_fused_rms_norm_gated.py +++ b/tests/kernels/core/test_fused_rms_norm_gated.py @@ -7,7 +7,9 @@ import pytest import torch -from aphrodite.third_party.flash_linear_attention.ops.kda import FusedRMSNormGated +from aphrodite.third_party.flash_linear_attention.ops.fused_norm_gate import ( + FusedRMSNormGated, +) from aphrodite.utils.torch_utils import set_random_seed DTYPES = [torch.bfloat16] diff --git a/tests/kernels/moe/test_grouped_topk.py b/tests/kernels/moe/test_grouped_topk.py index 778037d624..78edbe82a1 100644 --- a/tests/kernels/moe/test_grouped_topk.py +++ b/tests/kernels/moe/test_grouped_topk.py @@ -95,3 +95,241 @@ def test_grouped_topk( torch.testing.assert_close(baseline_topk_weights, test_topk_weights, atol=2e-2, rtol=0) torch.testing.assert_close(baseline_topk_ids, test_topk_ids, atol=0, rtol=0) + + +def _run_single_group_topk( + logits: torch.Tensor, + bias: torch.Tensor, + topk: int, + *, + scoring_func: str, + renormalize: bool, + routed_scaling_factor: float = 1.0, +) -> tuple[torch.Tensor, torch.Tensor]: + return fused_grouped_topk( + hidden_states=torch.empty((logits.shape[0], 0), dtype=logits.dtype, device=logits.device), + gating_output=logits, + topk=topk, + renormalize=renormalize, + e_score_correction_bias=bias, + num_expert_group=1, + topk_group=1, + scoring_func=scoring_func, + routed_scaling_factor=routed_scaling_factor, + ) + + +def _single_group_reference( + logits: torch.Tensor, + bias: torch.Tensor, + topk: int, + *, + scoring_func: str, + renormalize: bool, + routed_scaling_factor: float = 1.0, +) -> tuple[torch.Tensor, torch.Tensor]: + if scoring_func == "sigmoid": + scores = 0.5 * torch.tanh(0.5 * logits.float()) + 0.5 + else: + scores = torch.softmax(logits, dim=-1).float() + indices = torch.argsort(scores + bias.float(), dim=-1, descending=True, stable=True)[:, :topk] + values = scores.gather(1, indices) + if renormalize: + values /= values.sum(dim=-1, keepdim=True) + 1e-20 + values *= routed_scaling_factor + return values, indices.to(torch.int32) + + +@pytest.mark.skipif(not current_platform.is_cuda(), reason="This test is skipped on non-CUDA platform.") +def test_grouped_topk_single_group_large_batch(): + set_random_seed(0) + logits = torch.randn((1536, 896), dtype=torch.bfloat16, device="cuda") + bias = torch.randn((896,), dtype=torch.float32, device="cuda") + + expected_values, expected_ids = _single_group_reference(logits, bias, 16, scoring_func="sigmoid", renormalize=True) + actual_values, actual_ids = _run_single_group_topk(logits, bias, 16, scoring_func="sigmoid", renormalize=True) + + torch.testing.assert_close(actual_ids, expected_ids) + torch.testing.assert_close(actual_values, expected_values, atol=2e-5, rtol=0) + + +@pytest.mark.skipif(not current_platform.is_cuda(), reason="This test is skipped on non-CUDA platform.") +@pytest.mark.parametrize( + "num_experts,topk,input_dtype,bias_dtype", + [ + (512, 9, torch.bfloat16, torch.float32), + (512, 16, torch.float16, torch.float16), + (513, 9, torch.float32, torch.bfloat16), + (513, 16, torch.bfloat16, torch.float32), + (895, 9, torch.float16, torch.bfloat16), + (896, 16, torch.float32, torch.float16), + (897, 9, torch.bfloat16, torch.bfloat16), + (897, 16, torch.float16, torch.float32), + (1024, 9, torch.float32, torch.bfloat16), + (1024, 16, torch.bfloat16, torch.float16), + ], +) +@pytest.mark.parametrize( + "scoring_func,renormalize,routed_scaling_factor", + [ + ("sigmoid", True, 1.0), + ("sigmoid", False, 2.5), + ("softmax", True, 2.5), + ("softmax", False, 1.0), + ], +) +def test_grouped_topk_single_group_tiers( + num_experts: int, + topk: int, + input_dtype: torch.dtype, + bias_dtype: torch.dtype, + scoring_func: str, + renormalize: bool, + routed_scaling_factor: float, +): + set_random_seed(7) + logits = torch.randn((17, num_experts), dtype=input_dtype, device="cuda") + bias = torch.randn((num_experts,), dtype=bias_dtype, device="cuda") + + expected_values, expected_ids = _single_group_reference( + logits, + bias, + topk, + scoring_func=scoring_func, + renormalize=renormalize, + routed_scaling_factor=routed_scaling_factor, + ) + actual_values, actual_ids = _run_single_group_topk( + logits, + bias, + topk, + scoring_func=scoring_func, + renormalize=renormalize, + routed_scaling_factor=routed_scaling_factor, + ) + + torch.testing.assert_close(actual_ids, expected_ids) + torch.testing.assert_close(actual_values, expected_values, atol=2e-5, rtol=0) + + +@pytest.mark.skipif(not current_platform.is_cuda(), reason="This test is skipped on non-CUDA platform.") +@pytest.mark.parametrize( + "num_experts,topk,scoring_func", + [ + (128, 8, "sigmoid"), + (129, 8, "sigmoid"), + (257, 8, "sigmoid"), + (385, 8, "sigmoid"), + (512, 9, "sigmoid"), + (513, 9, "sigmoid"), + (769, 9, "sigmoid"), + (897, 16, "sigmoid"), + (1024, 16, "sigmoid"), + (128, 4, "softmax"), + (128, 5, "softmax"), + (129, 8, "softmax"), + (161, 8, "softmax"), + (256, 9, "softmax"), + (257, 8, "softmax"), + (512, 9, "softmax"), + (512, 17, "softmax"), + (512, 23, "softmax"), + (513, 8, "softmax"), + (577, 9, "softmax"), + (769, 9, "softmax"), + (897, 9, "softmax"), + (1024, 16, "softmax"), + ], +) +def test_grouped_topk_single_group_capacity_tiers( + num_experts: int, + topk: int, + scoring_func: str, +): + set_random_seed(11) + logits = torch.randn((3, num_experts), dtype=torch.bfloat16, device="cuda") + bias = torch.randn((num_experts,), dtype=torch.float32, device="cuda") + expected_values, expected_ids = _single_group_reference( + logits, + bias, + topk, + scoring_func=scoring_func, + renormalize=True, + routed_scaling_factor=2.5, + ) + actual_values, actual_ids = _run_single_group_topk( + logits, + bias, + topk, + scoring_func=scoring_func, + renormalize=True, + routed_scaling_factor=2.5, + ) + + torch.testing.assert_close(actual_ids, expected_ids) + torch.testing.assert_close(actual_values, expected_values, atol=2e-5, rtol=0) + + +@pytest.mark.skipif(not current_platform.is_cuda(), reason="This test is skipped on non-CUDA platform.") +@pytest.mark.parametrize("num_experts", [512, 896, 1024]) +def test_grouped_topk_single_group_stable_ties(num_experts: int): + logits = torch.zeros((1, num_experts), dtype=torch.bfloat16, device="cuda") + bias = torch.zeros((num_experts,), dtype=torch.float32, device="cuda") + + actual_values, actual_ids = _run_single_group_topk( + logits, + bias, + 16, + scoring_func="sigmoid", + renormalize=True, + routed_scaling_factor=2.5, + ) + + expected_ids = torch.arange(16, dtype=torch.int32, device="cuda")[None] + expected_values = torch.full((1, 16), 2.5 / 16, dtype=torch.float32, device="cuda") + torch.testing.assert_close(actual_ids, expected_ids) + torch.testing.assert_close(actual_values, expected_values, atol=2e-5, rtol=0) + + +@pytest.mark.skipif(not current_platform.is_cuda(), reason="This test is skipped on non-CUDA platform.") +@pytest.mark.parametrize("num_experts", [512, 896, 1024]) +@pytest.mark.parametrize("num_finite", [0, 15]) +@pytest.mark.parametrize("renormalize", [False, True]) +def test_grouped_topk_single_group_nonfinite_scores(num_experts: int, num_finite: int, renormalize: bool): + logits = torch.full((1, num_experts), float("nan"), dtype=torch.bfloat16, device="cuda") + if num_finite: + logits[0, :num_finite] = torch.arange(num_finite, dtype=torch.bfloat16, device="cuda") + logits[0, num_finite] = torch.inf + logits[0, num_finite + 1] = -torch.inf + bias = torch.zeros((num_experts,), dtype=torch.float32, device="cuda") + + actual_values, actual_ids = _run_single_group_topk( + logits, + bias, + 16, + scoring_func="sigmoid", + renormalize=renormalize, + routed_scaling_factor=2.5, + ) + + if num_finite == 0: + expected_ids = torch.arange(16, dtype=torch.int32, device="cuda")[None] + if renormalize: + expected_values = torch.full((1, 16), 1 / 16, dtype=torch.float32, device="cuda") + else: + expected_values = torch.zeros((1, 16), dtype=torch.float32, device="cuda") + else: + expected_ids = torch.cat( + ( + torch.arange(num_finite - 1, -1, -1, dtype=torch.int32, device="cuda"), + torch.tensor([num_finite], dtype=torch.int32, device="cuda"), + ) + )[None] + finite_values = logits[0, :num_finite].float().sigmoid().flip(0) + if renormalize: + finite_values /= finite_values.sum() + finite_values *= 2.5 + expected_values = torch.cat((finite_values, torch.zeros(1, dtype=torch.float32, device="cuda")))[None] + + torch.testing.assert_close(actual_ids, expected_ids) + torch.testing.assert_close(actual_values, expected_values, atol=2e-5, rtol=0) diff --git a/tests/kernels/moe/test_moe.py b/tests/kernels/moe/test_moe.py index a35ccc158e..aa2ed11326 100644 --- a/tests/kernels/moe/test_moe.py +++ b/tests/kernels/moe/test_moe.py @@ -30,6 +30,9 @@ batched_fused_marlin_moe, fused_marlin_moe, ) +from aphrodite.model_executor.layers.fused_moe.utils import ( + moe_use_td_hw_supported, +) from aphrodite.model_executor.layers.quantization.utils.marlin_utils import ( marlin_permute_bias, ) @@ -47,6 +50,7 @@ from aphrodite.model_executor.layers.quantization.utils.quant_utils import quantize_weights from aphrodite.platforms import current_platform from aphrodite.scalar_type import ScalarType, scalar_types +from aphrodite.triton_utils import tl from aphrodite.utils.math_utils import next_power_of_2 from aphrodite.utils.torch_utils import set_random_seed from tests.kernels.moe.utils import ( @@ -56,6 +60,8 @@ ) from tests.kernels.utils import opcheck, stack_and_dev, torch_experts, torch_moe +DEVICE_TYPE = current_platform.device_type + def iterative_moe( hidden_states: torch.Tensor, @@ -289,6 +295,7 @@ def run_moe_test( @pytest.mark.parametrize("ep_size", EP_SIZE) @pytest.mark.parametrize("dtype", [torch.bfloat16]) @pytest.mark.parametrize("padding", [True, False]) +@pytest.mark.parametrize("use_td", [False, True]) def test_fused_moe( m: int, n: int, @@ -298,9 +305,19 @@ def test_fused_moe( ep_size: int, dtype: torch.dtype, padding: bool, + use_td: bool, monkeypatch, workspace_init, ): + if use_td and not hasattr(tl, "make_tensor_descriptor"): + pytest.skip("Triton < 3.6 lacks tl.make_tensor_descriptor") + if use_td and not moe_use_td_hw_supported(): + pytest.skip( + "tensor_descriptor.gather requires XPU or NVIDIA Blackwell " + "(sm100+); lowers to tile::gather4, which is unsupported on " + "Hopper (sm90) and earlier" + ) + monkeypatch.setenv("APHRODITE_TRITON_USE_TD", "1" if use_td else "0") set_random_seed(7) # @@ -311,17 +328,17 @@ def test_fused_moe( # Setup test data # - a = torch.randn((m, k), device="cuda", dtype=dtype) / 10 - w1 = torch.randn((e, 2 * n, k), device="cuda", dtype=dtype) / 10 - w2 = torch.randn((e, k, n), device="cuda", dtype=dtype) / 10 + a = torch.randn((m, k), device=DEVICE_TYPE, dtype=dtype) / 10 + w1 = torch.randn((e, 2 * n, k), device=DEVICE_TYPE, dtype=dtype) / 10 + w2 = torch.randn((e, k, n), device=DEVICE_TYPE, dtype=dtype) / 10 - score = torch.randn((m, e), device="cuda", dtype=dtype) + score = torch.randn((m, e), device=DEVICE_TYPE, dtype=dtype) if ep_size > 1: local_e = e // ep_size - e_ids = torch.randint(0, e, (local_e,), device="cuda", dtype=torch.int32) - e_map = torch.full((e,), -1, device="cuda", dtype=torch.int32) - e_map[e_ids] = torch.arange(local_e, device="cuda", dtype=torch.int32) + e_ids = torch.randint(0, e, (local_e,), device=DEVICE_TYPE, dtype=torch.int32) + e_map = torch.full((e,), -1, device=DEVICE_TYPE, dtype=torch.int32) + e_map[e_ids] = torch.arange(local_e, device=DEVICE_TYPE, dtype=torch.int32) w1 = w1[e_ids] w2 = w2[e_ids] else: diff --git a/tests/kernels/moe/test_moe_align_block_size.py b/tests/kernels/moe/test_moe_align_block_size.py index 248a799bdb..82e44216e7 100644 --- a/tests/kernels/moe/test_moe_align_block_size.py +++ b/tests/kernels/moe/test_moe_align_block_size.py @@ -234,6 +234,7 @@ def test_moe_align_block_size_with_expert_map( experts = torch.randperm(num_experts, device="cuda")[:topk] for k in range(topk): topk_ids[i, k] = experts[k] if (experts[k] in local_experts) or not mask_inactive_experts else -1 + topk_ids[0, 0] = -1 actual_sorted_ids, actual_expert_ids, actual_num_tokens = moe_align_block_size( topk_ids=topk_ids, diff --git a/tests/kernels/moe/test_mxfp8_aiter_backend_selection.py b/tests/kernels/moe/test_mxfp8_aiter_backend_selection.py index 191657daf0..da1d61575b 100644 --- a/tests/kernels/moe/test_mxfp8_aiter_backend_selection.py +++ b/tests/kernels/moe/test_mxfp8_aiter_backend_selection.py @@ -18,9 +18,17 @@ if not current_platform.is_rocm(): pytest.skip("This test can only run on ROCm.", allow_module_level=True) +from aphrodite.model_executor.layers.fused_moe.activation import ( # noqa: E402 + MoEActivation, +) from aphrodite.model_executor.layers.fused_moe.experts.aiter_mxfp8_moe import ( # noqa: E402 + _AITER_SWIGLU_ALPHA, + _AITER_SWIGLU_BETA, AiterMxfp8Experts, ) +from aphrodite.model_executor.layers.fused_moe.experts.mxfp8_emulation_moe import ( # noqa: E402 + Mxfp8EmulationTritonExperts, +) from aphrodite.model_executor.layers.fused_moe.modular_kernel import ( # noqa: E402 FusedMoEActivationFormat, ) @@ -32,6 +40,7 @@ _SUPPORTED_BACKENDS, _mxfp8_backend_to_kernel_cls, _select_kernel_cls, + select_mxfp8_moe_backend, ) from aphrodite.model_executor.layers.quantization.utils.quant_utils import ( # noqa: E402 kMxfp8Dynamic, @@ -43,7 +52,15 @@ def _config(ep_size: int = 1): - cfg = make_dummy_moe_config(num_experts=128, experts_per_token=4, hidden_dim=6144) + # AiterMxfp8Experts hardcodes SwiGLU-OAI: match its required activation and + # alpha/beta so is_supported_config doesn't reject the config on those grounds. + cfg = make_dummy_moe_config( + num_experts=128, + experts_per_token=4, + hidden_dim=6144, + activation=MoEActivation.SWIGLUOAI_UNINTERLEAVE, + ) + cfg = dataclasses.replace(cfg, swiglu_alpha=_AITER_SWIGLU_ALPHA, swiglu_beta=_AITER_SWIGLU_BETA) if ep_size != 1: cfg = dataclasses.replace( cfg, @@ -72,12 +89,6 @@ def test_aiter_mxfp8_registered(): assert _mxfp8_backend_to_kernel_cls(Fp8MoeBackend.AITER_MXFP8) == [AiterMxfp8Experts] -def test_triton_selectable(): - assert _BACKEND_NAME_MAP["triton"] is Fp8MoeBackend.TRITON_MXFP8 - # Not auto-selected (only reachable explicitly), so FlyDSL still wins auto. - assert Fp8MoeBackend.TRITON_MXFP8 not in _SUPPORTED_BACKENDS - - @pytest.mark.parametrize("ep_size", [1, 2]) def test_ep_supported(ep_size): """FlyDSL accepts both TP and EP: apply() forwards expert_map as expert_mask.""" @@ -118,3 +129,24 @@ def test_explicit_moe_backend_aiter(): pytest.raises(ValueError, match="flydsl package"), ): _select_kernel_cls(Fp8MoeBackend.AITER_MXFP8, _config(1)) + + +def test_gfx950_picks_aiter(): + """Auto-select on real ROCm hardware with flydsl usable -> FlyDSL wins.""" + # NOTE: Fp8MoeBackend.AITER_MXFP8 does not require APHRODITE_ROCM_USE_AITER=1 + with ( + patch(f"{_AITER_MOD}.current_platform.supports_mx", return_value=True), + _flydsl_installed(True), + ): + backend, experts_cls = select_mxfp8_moe_backend(_config()) + assert backend is Fp8MoeBackend.AITER_MXFP8 + assert experts_cls is AiterMxfp8Experts + + +def test_gfx942_picks_emulation(): + """flydsl unusable (e.g. gfx942, no FlyDSL support) -> native Triton + dot_scaled backend wins instead.""" + with patch(f"{_AITER_MOD}.current_platform.supports_mx", return_value=False): + backend, experts_cls = select_mxfp8_moe_backend(_config()) + assert backend is Fp8MoeBackend.EMULATION + assert experts_cls is Mxfp8EmulationTritonExperts diff --git a/tests/kernels/moe/test_routing.py b/tests/kernels/moe/test_routing.py index 9923ea6348..7bf493ef9b 100644 --- a/tests/kernels/moe/test_routing.py +++ b/tests/kernels/moe/test_routing.py @@ -8,9 +8,19 @@ from aphrodite._aiter_ops import rocm_aiter_ops from aphrodite.distributed.eplb.eplb_state import EplbLayerState +from aphrodite.model_executor.layers.fused_moe.config import RoutingMethodType from aphrodite.model_executor.layers.fused_moe.router.base_router import ( eplb_map_to_physical_and_record, ) +from aphrodite.model_executor.layers.fused_moe.router.fused_topk_bias_router import ( + FusedTopKBiasRouter, +) +from aphrodite.model_executor.layers.fused_moe.router.fused_topk_router import ( + FusedTopKRouter, +) +from aphrodite.model_executor.layers.fused_moe.router.grouped_topk_router import ( + GroupedTopKRouter, +) from aphrodite.model_executor.layers.fused_moe.router.router_factory import ( create_fused_moe_router, ) @@ -36,6 +46,112 @@ def _is_aiter_capable() -> bool: NUM_EXPERTS = [8, 16, 64] +def test_degenerate_grouped_config_uses_standard_topk() -> None: + router = create_fused_moe_router( + top_k=4, + global_num_experts=128, + use_grouped_topk=True, + num_expert_group=1, + topk_group=1, + scoring_func="softmax", + renormalize=True, + ) + + assert isinstance(router, FusedTopKRouter) + hidden_states, router_logits = make_test_data(32, 256, 128) + + topk_weights, topk_ids = router.select_experts(hidden_states, router_logits) + baseline_weights, baseline_ids = baseline_fused_topk( + router_logits, + top_k=4, + renormalize=True, + ) + + assert_routing_results_close( + topk_weights, + topk_ids, + baseline_weights, + baseline_ids, + ) + + +def test_multiple_expert_groups_use_grouped_topk() -> None: + router = create_fused_moe_router( + top_k=4, + global_num_experts=128, + use_grouped_topk=True, + num_expert_group=8, + topk_group=4, + scoring_func="softmax", + renormalize=True, + ) + + assert isinstance(router, GroupedTopKRouter) + + +def test_degenerate_grouped_config_with_bias_uses_topk_bias() -> None: + router = create_fused_moe_router( + top_k=4, + global_num_experts=128, + use_grouped_topk=True, + num_expert_group=1, + topk_group=1, + scoring_func="softmax", + renormalize=True, + e_score_correction_bias=torch.empty(128), + ) + + assert isinstance(router, FusedTopKBiasRouter) + + +def test_degenerate_grouped_config_with_bias_keeps_routed_scale() -> None: + router = create_fused_moe_router( + top_k=4, + global_num_experts=128, + use_grouped_topk=True, + num_expert_group=1, + topk_group=1, + scoring_func="softmax", + renormalize=True, + routed_scaling_factor=1.1, + e_score_correction_bias=torch.empty(128), + ) + + assert isinstance(router, FusedTopKBiasRouter) + assert router.routed_scaling_factor == 1.1 + + +def test_degenerate_deepseek_v3_routing_stays_grouped() -> None: + router = create_fused_moe_router( + top_k=4, + global_num_experts=128, + use_grouped_topk=True, + num_expert_group=1, + topk_group=1, + scoring_func="sigmoid", + renormalize=True, + e_score_correction_bias=torch.empty(128), + ) + + assert isinstance(router, GroupedTopKRouter) + assert router.routing_method_type == RoutingMethodType.DeepSeekV3 + + +def test_single_expert_group_with_non_unit_scale_uses_grouped_topk() -> None: + router = create_fused_moe_router( + top_k=4, + global_num_experts=128, + use_grouped_topk=True, + num_expert_group=1, + topk_group=1, + scoring_func="softmax", + renormalize=True, + routed_scaling_factor=1.1, + ) + + assert isinstance(router, GroupedTopKRouter) + + def setup_eplb_state(enable_eplb: bool, global_num_experts: int) -> EplbLayerState | None: if not enable_eplb: return None diff --git a/tests/kernels/test_bf16_skinny_gemm.py b/tests/kernels/test_bf16_skinny_gemm.py new file mode 100644 index 0000000000..744f8ded6d --- /dev/null +++ b/tests/kernels/test_bf16_skinny_gemm.py @@ -0,0 +1,607 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Tests for the Kimi-K3 SM103 decode GEMM selector (shape-only dispatch).""" + +from pathlib import Path +from types import SimpleNamespace +from typing import Any + +import pytest +import regex as re +import torch +from torch import nn + +from aphrodite.model_executor.kernels.linear.cute_dsl.skinny_gemm import SkinnyGemmConfig +from aphrodite.models.kimi_k3.nvidia import low_latency_gemm as k3_gemm +from aphrodite.models.kimi_k3.nvidia.low_latency_gemm import KIMI_K3_PROJECTIONS + +# Keyed by local (N, K): (cute token counts, dsv3 token counts). 1536x7168 is +# the unified shared_gate_up_proj/mla_g_proj entry (dsv3 M1..16). +EXPECTED_SELECTIONS = { + (1536, 128): (set(), set(range(1, 17))), + (3072, 128): (set(), set(range(1, 17))), + (1536, 7168): (set(), set(range(1, 17))), + (3072, 7168): (set(range(1, 6)), set()), + (2112, 7168): (set(), set(range(1, 17))), + (2304, 1536): (set(), set(range(1, 17))), + (4608, 1536): (set(), set(range(1, 17))), + (3584, 7168): ({1}, set(range(2, 9))), + (6288, 7168): (set(range(1, 5)), set()), + (12448, 7168): (set(range(1, 4)), set()), + (7168, 768): (set(), set(range(1, 17))), + (7168, 1536): ({1}, set()), + (7168, 3072): ({1, 2}, set()), + (7168, 3584): ({1, 2}, set()), + (7168, 4224): ({1}, set()), + (7168, 8448): (set(range(1, 4)), set()), + (8448, 7168): ({1, 2}, set()), + (16896, 7168): ({1, 2}, set()), + (20480, 7168): (set(range(1, 5)), set()), + (40960, 7168): (set(range(1, 5)), set()), + # TP16. + (3216, 7168): (set(range(1, 6)), set(range(9, 16))), + (768, 7168): (set(range(1, 5)), set(range(5, 17))), + (1152, 1536): ({1}, set(range(2, 17))), + (768, 128): (set(), set(range(1, 17))), + (7168, 384): (set(), set(range(1, 9))), + (4224, 7168): (set(range(1, 4)), set(range(4, 9))), + (10240, 7168): (set(range(1, 5)), set()), +} + +CUTE_CASES = [ + (spec.n, spec.k, num_tokens) for spec in k3_gemm.KIMI_K3_PROJECTIONS.values() for num_tokens, _ in spec.cute_configs +] + +RESIDUAL_CUTE_CASES = [ + (spec.n, spec.k, num_tokens) + for spec in k3_gemm.KIMI_K3_PROJECTIONS.values() + for num_tokens, _ in spec.residual_configs +] + +EXPECTED_CUTE_CONFIGS = { + (3072, 7168, 1): (224, 3, 4, 8), + (3072, 7168, 2): (128, 3, 2, 8), + (3072, 7168, 3): (128, 2, 1, 8), + (3072, 7168, 4): (64, 2, 2, 8), + (3072, 7168, 5): (128, 3, 1, 8), + (3584, 7168, 1): (224, 2, 4, 8), + (6288, 7168, 1): (224, 3, 4, 8), + (6288, 7168, 2): (64, 3, 2, 8), + (6288, 7168, 3): (32, 3, 4, 8), + (6288, 7168, 4): (128, 6, 1, 8), + (12448, 7168, 1): (224, 4, 2, 8), + (12448, 7168, 2): (64, 4, 2, 8), + (12448, 7168, 3): (64, 2, 2, 8), + (7168, 1536, 1): (96, 4, 2, 8), + (7168, 3072, 1): (96, 2, 4, 8), + (7168, 3072, 2): (32, 4, 4, 8), + (7168, 3584, 1): (224, 4, 2, 8), + (7168, 3584, 2): (64, 4, 2, 8), + (7168, 4224, 1): (96, 4, 2, 4), + (7168, 8448, 1): (32, 4, 4, 8), + (7168, 8448, 2): (96, 4, 1, 8), + (7168, 8448, 3): (96, 4, 1, 8), + (8448, 7168, 1): (224, 3, 4, 8), + (8448, 7168, 2): (32, 4, 4, 8), + (16896, 7168, 1): (224, 6, 4, 8), + (16896, 7168, 2): (32, 4, 4, 8), + (20480, 7168, 1): (224, 4, 2, 8), + (20480, 7168, 2): (64, 4, 2, 8), + (20480, 7168, 3): (64, 2, 2, 8), + (20480, 7168, 4): (64, 4, 1, 8), + (40960, 7168, 1): (128, 4, 2, 8), + (40960, 7168, 2): (64, 4, 2, 8), + (40960, 7168, 3): (64, 2, 2, 8), + (40960, 7168, 4): (64, 4, 1, 8), + # TP16. + (3216, 7168, 1): (224, 3, 4, 8), + (3216, 7168, 2): (128, 4, 2, 8), + (3216, 7168, 3): (128, 2, 1, 8), + (3216, 7168, 4): (64, 2, 2, 8), + (3216, 7168, 5): (128, 3, 1, 8), + (768, 7168, 1): (224, 2, 4, 8), + (768, 7168, 2): (224, 2, 2, 8), + (768, 7168, 3): (224, 2, 2, 8), + (768, 7168, 4): (224, 2, 2, 8), + (1152, 1536, 1): (192, 3, 4, 8), + (4224, 7168, 1): (224, 3, 4, 8), + (4224, 7168, 2): (128, 2, 1, 8), + (4224, 7168, 3): (64, 2, 2, 8), + (10240, 7168, 1): (224, 4, 2, 8), + (10240, 7168, 2): (32, 2, 4, 8), + (10240, 7168, 3): (64, 4, 1, 8), + (10240, 7168, 4): (64, 4, 1, 8), +} + +EXPECTED_RESIDUAL_CUTE_CONFIGS = { + (7168, 3584, 1): (64, 4, 2, 8), + (7168, 3584, 2): (64, 7, 2, 8), + (7168, 3584, 3): (64, 2, 1, 8), + (7168, 3584, 4): (64, 2, 1, 8), +} + + +def _config_tuple(config) -> tuple[int, int, int, int]: + return ( + config.block_size, + config.outputs_per_block, + config.k_unroll, + config.vector_width, + ) + + +def test_table_is_keyed_by_shape() -> None: + for (n, k), spec in k3_gemm.KIMI_K3_PROJECTIONS.items(): + assert (spec.n, spec.k) == (n, k) + + +def test_every_dsv3_routed_shape_is_instantiated() -> None: + """dsv3_fused_a_gemm specializes on (K, N); an unlisted shape raises. + + The table routes by shape while the kernel is built per shape, so a missing + instantiation only shows up at the token counts that route to dsv3. Checking + it here needs no GPU, which is the point -- a GPU-only check is exactly what + let (3216, 7168) ship without its DISPATCH_DSV3_SHAPE(7168, 3216). + """ + source = (Path(__file__).resolve().parents[2] / "csrc" / "libtorch_stable" / "dsv3_fused_a_gemm.cu").read_text( + encoding="utf-8" + ) + # Benchmark-only shapes live behind APHRODITE_K3_BENCH_SHAPES and are not built + # by default, so they must not count as available. + production_macros = source.split("#ifdef APHRODITE_K3_BENCH_SHAPES")[0] + explicit = source.split("#undef DISPATCH_DSV3_SHAPE")[1].split("#ifdef APHRODITE_K3_BENCH_SHAPES")[0] + compiled = { + (int(hd_in), int(hd_out)) + for hd_in, hd_out in re.findall(r"DISPATCH_DSV3_SHAPE\((\d+),\s*(\d+)\)", production_macros) + } | {(int(hd_in), int(hd_out)) for hd_in, hd_out in re.findall(r"hd_in == (\d+) && hd_out == (\d+)", explicit)} + assert compiled, "failed to parse the dispatch list" + + missing = sorted( + (spec.n, spec.k) + for spec in KIMI_K3_PROJECTIONS.values() + if spec.dsv3_tokens and (spec.k, spec.n) not in compiled + ) + assert not missing, f"routed to dsv3 with no instantiation: {missing}; add DISPATCH_DSV3_SHAPE(K, N) for each" + + +def test_packed_row_major_rejects_single_row_slice() -> None: + packed = torch.empty(1, 128) + sliced = torch.empty(1, 144)[:, :128] + + assert packed.is_contiguous() + assert sliced.is_contiguous() + assert k3_gemm._is_packed_row_major(packed) + assert not k3_gemm._is_packed_row_major(sliced) + + +def test_cute_configs_match_measured_table() -> None: + actual = { + (spec.n, spec.k, num_tokens): _config_tuple(config) + for spec in k3_gemm.KIMI_K3_PROJECTIONS.values() + for num_tokens, config in spec.cute_configs + } + assert actual == EXPECTED_CUTE_CONFIGS + + +def test_residual_cute_configs_match_measured_table() -> None: + actual = { + (spec.n, spec.k, num_tokens): _config_tuple(config) + for spec in k3_gemm.KIMI_K3_PROJECTIONS.values() + for num_tokens, config in spec.residual_configs + } + assert actual == EXPECTED_RESIDUAL_CUTE_CONFIGS + + +@pytest.mark.parametrize("key", EXPECTED_SELECTIONS) +def test_sm103_selector_table(key: tuple[int, int]) -> None: + n, k = key + cute_tokens, dsv3_tokens = EXPECTED_SELECTIONS[key] + for num_tokens in range(1, 17): + backend = k3_gemm.select_kimi_k3_backend(num_tokens, n, k) + if num_tokens in cute_tokens: + assert backend == "cute" + elif num_tokens in dsv3_tokens: + assert backend == "dsv3_fused_a" + else: + assert backend is None + + +@pytest.mark.parametrize("key", EXPECTED_SELECTIONS) +def test_selector_requires_supported_shape_and_tokens(key: tuple[int, int]) -> None: + n, k = key + assert k3_gemm.select_kimi_k3_backend(0, n, k) is None + assert k3_gemm.select_kimi_k3_backend(17, n, k) is None + assert k3_gemm.select_kimi_k3_backend(1, n + 1, k) is None + assert k3_gemm.select_kimi_k3_backend(1, n, k + 1) is None + + +def test_unlisted_shape_and_unselected_tokens_fall_back() -> None: + # Shape absent from the table. + assert k3_gemm.select_kimi_k3_backend(1, 1000, 1000) is None + # o_proj (7168,1536) is CuTe M1 only; M2+ falls back. + assert k3_gemm.select_kimi_k3_backend(2, 7168, 1536) is None + + +@pytest.mark.parametrize("num_tokens", range(1, 17)) +def test_sm103_residual_selector_table(num_tokens: int) -> None: + backend = k3_gemm.select_kimi_k3_backend(num_tokens, 7168, 3584, has_residual=True) + assert backend == ("cute" if num_tokens <= 4 else None) + + +def test_build_plan_matches_selector() -> None: + for spec in k3_gemm.KIMI_K3_PROJECTIONS.values(): + plan = k3_gemm._build_plan(spec) + for num_tokens in range(1, 17): + backend = k3_gemm.select_kimi_k3_backend(num_tokens, spec.n, spec.k) + if backend is None: + assert num_tokens not in plan + else: + assert plan[num_tokens][0] == backend + + +def test_installation_is_shape_specific_and_unquantized( + monkeypatch: pytest.MonkeyPatch, +) -> None: + class FakeLinear(nn.Module): + def __init__(self, quant_method: object, n: int, k: int) -> None: + super().__init__() + self.quant_method = quant_method + self.weight = torch.empty(n, k) + + class FakeHead(nn.Module): + def __init__(self, n: int, k: int) -> None: + super().__init__() + self.quant_method = k3_gemm.UnquantizedEmbeddingMethod() + self.weight = torch.empty(n, k) + + root = nn.Module() + # dsv3-only shape (no cute warmup contribution). + root.dsv3_only = FakeLinear(k3_gemm.UnquantizedLinearMethod(), 2304, 1536) + # quantized: must be left untouched. + quantized_method = object() + root.quantized = FakeLinear(quantized_method, 6288, 7168) + # cute shape. + root.cute = FakeLinear(k3_gemm.UnquantizedLinearMethod(), 6288, 7168) + # cute + residual shape. + root.residual = FakeLinear(k3_gemm.UnquantizedLinearMethod(), 7168, 3584) + # shape absent from the table: must be left untouched. + root.unlisted = FakeLinear(k3_gemm.UnquantizedLinearMethod(), 1234, 5678) + root.lm_head = FakeHead(20480, 7168) + + monkeypatch.setattr(k3_gemm, "LinearBase", FakeLinear) + monkeypatch.setattr(k3_gemm, "ParallelLMHead", FakeHead) + monkeypatch.setattr(k3_gemm, "_is_sm103", lambda: True) + warmup_configs: set[SkinnyGemmConfig] = set() + residual_warmup_configs: set[SkinnyGemmConfig] = set() + monkeypatch.setattr(k3_gemm.shape_dynamic_skinny_gemm, "is_available", lambda: True) + + def request_warmup_configs(dtype, configs, *, has_residual=False): + target = residual_warmup_configs if has_residual else warmup_configs + target.update(configs) + + monkeypatch.setattr( + k3_gemm.shape_dynamic_skinny_gemm, + "request_warmup_configs", + request_warmup_configs, + ) + + k3_gemm.enable_kimi_k3_low_latency_gemm(root, torch.bfloat16) + + assert isinstance(root.dsv3_only.quant_method, k3_gemm.KimiK3LowLatencyLinearMethod) + assert isinstance(root.cute.quant_method, k3_gemm.KimiK3LowLatencyLinearMethod) + assert isinstance(root.residual.quant_method, k3_gemm.KimiK3LowLatencyLinearMethod) + assert root.quantized.quant_method is quantized_method + assert type(root.unlisted.quant_method) is k3_gemm.UnquantizedLinearMethod + assert isinstance(root.lm_head.quant_method, k3_gemm.KimiK3LowLatencyEmbeddingMethod) + # Warmup covers only the installed modules' local (N, K). + assert warmup_configs == { + config + for key in ((6288, 7168), (7168, 3584), (20480, 7168)) + for _, config in k3_gemm.KIMI_K3_PROJECTIONS[key].cute_configs + } + assert residual_warmup_configs == { + config for _, config in k3_gemm.KIMI_K3_PROJECTIONS[(7168, 3584)].residual_configs + } + + +@pytest.mark.parametrize( + "dtype,platform_enabled", + [(torch.float16, True), (torch.bfloat16, False)], +) +def test_installation_requires_bf16_sm103( + monkeypatch: pytest.MonkeyPatch, + dtype: torch.dtype, + platform_enabled: bool, +) -> None: + class FakeLinear(nn.Module): + def __init__(self) -> None: + super().__init__() + self.quant_method = k3_gemm.UnquantizedLinearMethod() + self.weight = torch.empty(2304, 1536) + + root = nn.Module() + root.projection = FakeLinear() + monkeypatch.setattr(k3_gemm, "LinearBase", FakeLinear) + monkeypatch.setattr(k3_gemm, "_is_sm103", lambda: platform_enabled) + + k3_gemm.enable_kimi_k3_low_latency_gemm(root, dtype) + + assert type(root.projection.quant_method) is k3_gemm.UnquantizedLinearMethod + + +def _require_sm103_and_dsv3() -> None: + if not torch.cuda.is_available() or torch.cuda.get_device_capability() != (10, 3): + pytest.skip("Kimi-K3 production selection requires SM103") + if not hasattr(torch.ops._C, "dsv3_fused_a_gemm"): + pytest.skip("dsv3_fused_a_gemm was not built") + + +def _require_sm103_and_cute() -> None: + if not torch.cuda.is_available() or torch.cuda.get_device_capability() != (10, 3): + pytest.skip("Kimi-K3 production selection requires SM103") + if not k3_gemm.shape_dynamic_skinny_gemm.is_available(): + pytest.skip("CuTe DSL is not available") + + +@pytest.mark.parametrize("n,k,num_tokens", CUTE_CASES) +def test_cute_selected_shapes(n: int, k: int, num_tokens: int) -> None: + _require_sm103_and_cute() + torch.manual_seed(42) + x = torch.randn(num_tokens, k, dtype=torch.bfloat16, device="cuda") + weight = torch.randn(n, k, dtype=torch.bfloat16, device="cuda") + + output = k3_gemm.try_low_latency_gemm(x, weight) + + assert output is not None + reference = torch.nn.functional.linear(x, weight) + cosine = torch.nn.functional.cosine_similarity(output.float().flatten(), reference.float().flatten(), dim=0).item() + assert cosine > 0.999 + + +def _dsv3_probe_tokens(tokens: frozenset[int]) -> set[int]: + """Extremes, plus both sides of the kernel's num_tokens<=8 tile_n branch.""" + if not tokens: + return set() + return {min(tokens), max(tokens)} | ({8, 9} & set(tokens)) + + +# Derived from the table rather than hand-listed, so a shape routed to dsv3 +# cannot be added without being exercised here. +DSV3_CASES = sorted( + (num_tokens, spec.n, spec.k) + for spec in KIMI_K3_PROJECTIONS.values() + for num_tokens in _dsv3_probe_tokens(spec.dsv3_tokens) +) + + +@pytest.mark.parametrize("num_tokens,n,k", DSV3_CASES) +def test_dsv3_selected_shapes(num_tokens: int, n: int, k: int) -> None: + _require_sm103_and_dsv3() + spec = k3_gemm.KIMI_K3_PROJECTIONS[(n, k)] + assert num_tokens in spec.dsv3_tokens + torch.manual_seed(42) + x = torch.randn(num_tokens, k, dtype=torch.bfloat16, device="cuda") + weight = torch.randn(n, k, dtype=torch.bfloat16, device="cuda") + + output = k3_gemm.try_low_latency_gemm(x, weight) + + assert output is not None + reference = torch.nn.functional.linear(x, weight) + cosine = torch.nn.functional.cosine_similarity(output.float().flatten(), reference.float().flatten(), dim=0).item() + assert cosine > 0.999 + + +def test_nonpacked_single_token_dsv3_falls_back() -> None: + _require_sm103_and_dsv3() + n, k = 1536, 128 + storage = torch.randn(1, k + 16, dtype=torch.bfloat16, device="cuda") + x = storage[:, :k] + weight = torch.randn(n, k, dtype=torch.bfloat16, device="cuda") + spec = k3_gemm.KIMI_K3_PROJECTIONS[(n, k)] + method = k3_gemm.KimiK3LowLatencyLinearMethod(k3_gemm._build_plan(spec), k3_gemm._build_residual_plan(spec)) + + assert x.is_contiguous() + assert x.stride() == (k + 16, 1) + assert not k3_gemm._runtime_ok(x, weight) # strict guard rejects the view + output = method.apply(SimpleNamespace(weight=weight), x) + + reference = torch.nn.functional.linear(x, weight) + torch.testing.assert_close(output, reference) + + +def test_selected_kernels_cuda_graph_capture() -> None: + _require_sm103_and_cute() + _require_sm103_and_dsv3() + cute_spec = k3_gemm.KIMI_K3_PROJECTIONS[(6288, 7168)] + dsv3_spec = k3_gemm.KIMI_K3_PROJECTIONS[(1536, 128)] + cute_x = torch.randn(1, cute_spec.k, dtype=torch.bfloat16, device="cuda") + cute_weight = torch.randn(cute_spec.n, cute_spec.k, dtype=torch.bfloat16, device="cuda") + dsv3_x = torch.randn(1, dsv3_spec.k, dtype=torch.bfloat16, device="cuda") + dsv3_weight = torch.randn(dsv3_spec.n, dsv3_spec.k, dtype=torch.bfloat16, device="cuda") + k3_gemm.try_low_latency_gemm(cute_x, cute_weight) + k3_gemm.try_low_latency_gemm(dsv3_x, dsv3_weight) + torch.accelerator.synchronize() + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + cute_output = k3_gemm.try_low_latency_gemm(cute_x, cute_weight) + dsv3_output = k3_gemm.try_low_latency_gemm(dsv3_x, dsv3_weight) + graph.replay() + torch.accelerator.synchronize() + + assert cute_output is not None + assert dsv3_output is not None + for output, activation, weight in ( + (cute_output, cute_x, cute_weight), + (dsv3_output, dsv3_x, dsv3_weight), + ): + reference = torch.nn.functional.linear(activation, weight) + cosine = torch.nn.functional.cosine_similarity( + output.float().flatten(), reference.float().flatten(), dim=0 + ).item() + assert cosine > 0.999 + + +@pytest.mark.parametrize("num_tokens", [1, 8, 9, 16]) +def test_dsv3_cuda_graph_capture_tile_branches(num_tokens: int) -> None: + """Capture DSV3 across the num_tokens<=8 vs >8 tile_n branch.""" + _require_sm103_and_dsv3() + spec = k3_gemm.KIMI_K3_PROJECTIONS[(1536, 128)] + x = torch.randn(num_tokens, spec.k, dtype=torch.bfloat16, device="cuda") + weight = torch.randn(spec.n, spec.k, dtype=torch.bfloat16, device="cuda") + k3_gemm.try_low_latency_gemm(x, weight) + torch.accelerator.synchronize() + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + output = k3_gemm.try_low_latency_gemm(x, weight) + graph.replay() + torch.accelerator.synchronize() + + assert output is not None + reference = torch.nn.functional.linear(x, weight) + cosine = torch.nn.functional.cosine_similarity(output.float().flatten(), reference.float().flatten(), dim=0).item() + assert cosine > 0.999 + + +@pytest.mark.parametrize("n,k,num_tokens", RESIDUAL_CUTE_CASES) +def test_cute_residual_epilogue(n: int, k: int, num_tokens: int) -> None: + _require_sm103_and_cute() + torch.manual_seed(42 + num_tokens) + x = torch.randn(num_tokens, k, dtype=torch.bfloat16, device="cuda") + weight = torch.randn(n, k, dtype=torch.bfloat16, device="cuda") + residual = torch.randn(num_tokens, n, dtype=torch.bfloat16, device="cuda") + spec = k3_gemm.KIMI_K3_PROJECTIONS[(n, k)] + config = spec.residual_config(num_tokens) + assert config is not None + + output = k3_gemm.shape_dynamic_skinny_gemm(x, weight, config, residual) + + reference = x.float() @ weight.float().t() + residual.float() + cosine = torch.nn.functional.cosine_similarity(output.float().flatten(), reference.flatten(), dim=0).item() + assert cosine > 0.999 + + +@pytest.mark.parametrize("num_tokens", range(1, 17)) +def test_cute_residual_epilogue_all_supported_token_counts(num_tokens: int) -> None: + _require_sm103_and_cute() + from aphrodite.model_executor.kernels.linear.cute_dsl.skinny_gemm import ( + ShapeDynamicSkinnyGemm, + ) + + n, k = 64, 512 + x = torch.randn(num_tokens, k, dtype=torch.bfloat16, device="cuda") + weight = torch.randn(n, k, dtype=torch.bfloat16, device="cuda") + residual = torch.randn(num_tokens, n, dtype=torch.bfloat16, device="cuda") + config = ShapeDynamicSkinnyGemm._config(num_tokens, n, k) + + output = k3_gemm.shape_dynamic_skinny_gemm(x, weight, config, residual) + + reference = x.float() @ weight.float().t() + residual.float() + torch.testing.assert_close(output.float(), reference, rtol=2e-2, atol=2e-1) + + +@pytest.mark.parametrize("num_tokens", range(1, 5)) +def test_cute_residual_epilogue_cuda_graph_capture(num_tokens: int) -> None: + _require_sm103_and_cute() + spec = k3_gemm.KIMI_K3_PROJECTIONS[(7168, 3584)] + config = spec.residual_config(num_tokens) + assert config is not None + x = torch.randn(num_tokens, spec.k, dtype=torch.bfloat16, device="cuda") + weight = torch.randn(spec.n, spec.k, dtype=torch.bfloat16, device="cuda") + residual = torch.randn(num_tokens, spec.n, dtype=torch.bfloat16, device="cuda") + k3_gemm.shape_dynamic_skinny_gemm(x, weight, config, residual) + torch.accelerator.synchronize() + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + output = k3_gemm.shape_dynamic_skinny_gemm(x, weight, config, residual) + graph.replay() + torch.accelerator.synchronize() + + reference = x.float() @ weight.float().t() + residual.float() + cosine = torch.nn.functional.cosine_similarity(output.float().flatten(), reference.flatten(), dim=0).item() + assert cosine > 0.999 + + +class _SkinnyGemmSpy: + """Wraps the skinny-GEMM singleton to record whether CuTe was invoked.""" + + def __init__(self, real: Any) -> None: + self._real = real + self.calls: list[int] = [] + + def __call__(self, a, b, config=None, residual=None): + self.calls.append(a.shape[0]) + return self._real(a, b, config, residual) + + def is_available(self) -> bool: + return self._real.is_available() + + +@pytest.mark.parametrize("num_tokens", [1, 2, 3, 4]) +def test_latent_moe_production_layout_residual( + monkeypatch: pytest.MonkeyPatch, + num_tokens: int, +) -> None: + """The real Latent-MoE residual is a non-packed slice of a cat buffer. + + The strict packed-row-major guard rejects such a slice at every token count + (a size-1 leading dim reads as contiguous but its stride is not packed), so + the CuTe residual epilogue never fires for this production layout and the + method falls back to addmm. Output is correct regardless of the path. + """ + _require_sm103_and_cute() + latent_dim, shared_dim = 3584, 7168 # routed_expert_up_proj K, N + torch.manual_seed(7 + num_tokens) + buf = torch.randn(num_tokens, latent_dim + shared_dim, dtype=torch.bfloat16, device="cuda") + latent = buf[:, :latent_dim] # non-contiguous view (row stride = full width) + residual = buf[:, latent_dim:] # non-contiguous view + weight = torch.randn(shared_dim, latent_dim, dtype=torch.bfloat16, device="cuda") + + spec = k3_gemm.KIMI_K3_PROJECTIONS[(shared_dim, latent_dim)] + method = k3_gemm.KimiK3LowLatencyLinearMethod(k3_gemm._build_plan(spec), k3_gemm._build_residual_plan(spec)) + spy = _SkinnyGemmSpy(k3_gemm.shape_dynamic_skinny_gemm) + monkeypatch.setattr(k3_gemm, "shape_dynamic_skinny_gemm", spy) + + layer = SimpleNamespace(weight=weight) + output = method.apply_with_residual(layer, latent, residual) + + reference = latent.float() @ weight.float().t() + residual.float() + cosine = torch.nn.functional.cosine_similarity(output.float().flatten(), reference.flatten(), dim=0).item() + assert cosine > 0.999 # correct regardless of the path taken + assert not spy.calls, "non-packed buf-slice residual must fall back to addmm at every M" + + +def test_residual_dispatch_falls_back_to_addmm( + monkeypatch: pytest.MonkeyPatch, +) -> None: + fallback = torch.randn(2, 3) + residual = torch.randn(2, 3) + x = torch.randn(2, 4) + weight = torch.randn(3, 4) + monkeypatch.setattr(torch, "addmm", lambda *args: fallback) + # CPU tensors fail the runtime check, forcing the addmm fallback. + method = k3_gemm.KimiK3LowLatencyLinearMethod({}, {}) + + output = method.apply_with_residual(SimpleNamespace(weight=weight), x, residual) + + assert output is fallback + + +def test_fallback_preserves_default_method(monkeypatch: pytest.MonkeyPatch) -> None: + fallback = torch.empty(2, 8) + monkeypatch.setattr( + k3_gemm.UnquantizedLinearMethod, + "apply", + lambda *args: fallback, + ) + # 1-D input fails the runtime check, forcing the base-method fallback. + method = k3_gemm.KimiK3LowLatencyLinearMethod({}, {}) + + output = method.apply( + SimpleNamespace(weight=torch.empty(0)), + torch.empty(0), + ) + + assert output is fallback diff --git a/tests/lora/test_qwenvl.py b/tests/lora/test_qwenvl.py index 96a2afb89c..66c50e6369 100644 --- a/tests/lora/test_qwenvl.py +++ b/tests/lora/test_qwenvl.py @@ -179,6 +179,18 @@ def run_beam_search_test( QWEN3VL_MODEL_PATH = "Qwen/Qwen3-VL-4B-Instruct" +def _enable_deterministic_lora_shrink(monkeypatch: pytest.MonkeyPatch) -> None: + # These tests assert exact greedy outputs. Force the Triton LoRA shrink + # kernel to use SPLIT_K=1 so it stores the complete reduction directly + # instead of accumulating split-K partial results with atomic_add. This + # targets reduction determinism, not full batch invariance. + monkeypatch.setenv("APHRODITE_BATCH_INVARIANT", "1") + # The kernel configuration reads APHRODITE_BATCH_INVARIANT at import time. + # Spawn the engine process so it observes this setting even if the LoRA + # Triton utilities were already imported during test collection. + monkeypatch.setenv("APHRODITE_WORKER_MULTIPROC_METHOD", "spawn") + + def test_qwen2vl_lora(qwen2vl_lora_files): """Test Qwen 2.0 VL model with LoRA""" config = TestConfig(model_path=QWEN2VL_MODEL_PATH, lora_path=qwen2vl_lora_files) @@ -235,7 +247,12 @@ def test_qwen25vl_vision_lora(qwen25vl_vision_lora_files): ) -def test_qwen3vl_vision_lora(qwen3vl_vision_lora_files): +def test_qwen3vl_vision_lora( + qwen3vl_vision_lora_files, + monkeypatch: pytest.MonkeyPatch, +): + _enable_deterministic_lora_shrink(monkeypatch) + config = TestConfig( model_path=QWEN3VL_MODEL_PATH, lora_path=qwen3vl_vision_lora_files, @@ -258,6 +275,7 @@ def test_qwen2vl_multiple_lora_types( qwen2vl_language_lora_files, qwen2vl_vision_tower_connector_lora_files, qwen2vl_vision_tower_lora_files, + monkeypatch: pytest.MonkeyPatch, ): """ Test multiple LoRA adapter types (language, vision tower + connector, @@ -268,6 +286,8 @@ def test_qwen2vl_multiple_lora_types( the multimodal encoder cache correctly manages state transitions between language-only and vision-enabled LoRA adapters. """ + _enable_deterministic_lora_shrink(monkeypatch) + config = TestConfig( model_path=QWEN2VL_MODEL_PATH, # We'll override the lora_path for each specific test, but need to provide diff --git a/tests/models/inkling/test_moe_weight_layout.py b/tests/models/inkling/test_moe_weight_layout.py index a12c3f2f5f..ec5c03eadf 100644 --- a/tests/models/inkling/test_moe_weight_layout.py +++ b/tests/models/inkling/test_moe_weight_layout.py @@ -173,6 +173,28 @@ def test_moe_loads_compressed_tensors_global_scale(projection: str, scale_kind: assert loaded == [f"experts.routed_experts.{projection}_{scale_kind}_global_scale"] +@pytest.mark.parametrize(("projection", "checkpoint_rows"), [("w13", 8), ("w2", 4)]) +def test_moe_loads_channelwise_scale_for_tp(projection: str, checkpoint_rows: int) -> None: + param = torch.nn.Parameter(torch.empty(2, 4, 1)) + experts = SimpleNamespace( + **{f"{projection}_weight_scale": param}, + moe_config=SimpleNamespace(moe_parallel_config=SimpleNamespace(tp_rank=1)), + ) + layer = SimpleNamespace( + experts=SimpleNamespace(routed_experts=experts), + _local_expert_slots=lambda: {0: 0, 2: 1}, + ) + checkpoint_scale = torch.arange(3 * checkpoint_rows).reshape(3, checkpoint_rows, 1) + + loaded = moe.InklingMoE.load_expert_weight(layer, f"experts.{projection}_weight_scale", checkpoint_scale) + + expected = checkpoint_scale[[0, 2]] + if projection == "w13": + expected = expected[:, 4:].reshape(2, 2, 2, 1).transpose(1, 2).flatten(1, 2) + torch.testing.assert_close(param, expected.float()) + assert loaded == [f"experts.routed_experts.{projection}_weight_scale"] + + def test_sink_down_projection_is_packed_during_load(monkeypatch) -> None: monkeypatch.setattr(moe, "get_tensor_model_parallel_world_size", lambda: 2) monkeypatch.setattr(moe, "get_tensor_model_parallel_rank", lambda: 1) diff --git a/tests/models/kimi_k3/test_attn_res.py b/tests/models/kimi_k3/test_attn_res.py index 092657e258..0a49b797a1 100644 --- a/tests/models/kimi_k3/test_attn_res.py +++ b/tests/models/kimi_k3/test_attn_res.py @@ -5,6 +5,7 @@ import torch import torch.nn.functional as F +from aphrodite.models.kimi_k3.common.mtp import fused_mtp_input from aphrodite.models.kimi_k3.nvidia.ops import attn_res from aphrodite.platforms import current_platform @@ -169,3 +170,33 @@ def test_attn_res_without_output_norm(): ) torch.testing.assert_close(actual, expected, atol=8e-2, rtol=3e-2) + + +@pytest.mark.parametrize("num_tokens", [0, 1, 17]) +def test_fused_mtp_input(num_tokens: int): + positions = torch.arange(num_tokens, device="cuda") + inputs_embeds = _randn_with_row_padding(num_tokens, HIDDEN_SIZE, padding=7) + previous_hidden_states = _randn_with_row_padding(num_tokens, HIDDEN_SIZE, padding=11) + enorm_weight = torch.randn(HIDDEN_SIZE, device="cuda", dtype=torch.bfloat16) + hnorm_weight = torch.randn(HIDDEN_SIZE, device="cuda", dtype=torch.bfloat16) + + masked_inputs_embeds = torch.where(positions.unsqueeze(-1) == 0, 0, inputs_embeds) + expected = torch.cat( + ( + F.rms_norm(masked_inputs_embeds, (HIDDEN_SIZE,), enorm_weight, EPS), + F.rms_norm(previous_hidden_states, (HIDDEN_SIZE,), hnorm_weight, EPS), + ), + dim=-1, + ) + actual = fused_mtp_input( + positions, + inputs_embeds, + previous_hidden_states, + enorm_weight, + hnorm_weight, + EPS, + ) + + torch.testing.assert_close(actual, expected, atol=2e-2, rtol=2e-2) + assert actual.shape == (num_tokens, 2 * HIDDEN_SIZE) + assert actual.is_contiguous() diff --git a/tests/models/kimi_k3/test_eagle3.py b/tests/models/kimi_k3/test_eagle3.py new file mode 100644 index 0000000000..ba972d485d --- /dev/null +++ b/tests/models/kimi_k3/test_eagle3.py @@ -0,0 +1,130 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from types import SimpleNamespace +from unittest.mock import Mock + +import torch + +from aphrodite.model_executor.models.interfaces import supports_eagle3 +from aphrodite.models.kimi_k3.nvidia import model as kimi_model +from aphrodite.models.kimi_k3.nvidia.model import ( + KimiK3ForConditionalGeneration, + KimiLinearModel, +) + + +def _make_kimi_linear_model() -> KimiLinearModel: + model = object.__new__(KimiLinearModel) + object.__setattr__(model, "aux_hidden_state_layers", (2,)) + object.__setattr__(model, "use_sequence_parallel", False) + return model + + +def test_kimi_k3_advertises_eagle3_support(): + assert supports_eagle3(KimiK3ForConditionalGeneration) + + +def test_kimi_k3_uses_shared_eagle3_layer_configuration(): + target = object.__new__(KimiK3ForConditionalGeneration) + torch.nn.Module.__init__(target) + model = _make_kimi_linear_model() + object.__setattr__(model, "layers", [None] * 93) + language_model = SimpleNamespace( + embed_input_ids=lambda _: None, + model=model, + ) + object.__setattr__(target, "language_model", language_model) + object.__setattr__(target, "_language_model_names", ["language_model"]) + + target.set_aux_hidden_state_layers((2, 46, 90)) + + assert model.aux_hidden_state_layers == (2, 46, 90) + assert target.get_eagle3_default_aux_hidden_state_layers() == ( + 2, + 46, + 90, + ) + + +def test_kimi_linear_forward_extracts_standard_aux_hidden_states(monkeypatch): + model = _make_kimi_linear_model() + initial_hidden_states = torch.tensor([[1.0, 2.0]]) + layer_hidden_states = torch.tensor([[3.0, 4.0]]) + layer_residual = torch.tensor([[5.0, 6.0]]) + + object.__setattr__(model, "start_layer", 0) + object.__setattr__(model, "end_layer", 1) + object.__setattr__( + model, + "layers", + [Mock(return_value=(layer_hidden_states, None, layer_residual))], + ) + object.__setattr__(model, "aux_hidden_state_layers", (0, 1)) + object.__setattr__(model, "use_attn_res", False) + monkeypatch.setattr( + kimi_model, + "get_pp_group", + lambda: SimpleNamespace(is_first_rank=True, is_last_rank=True), + ) + + output, aux_hidden_states = model.forward( + input_ids=None, + positions=torch.tensor([0]), + intermediate_tensors=None, + inputs_embeds=initial_hidden_states, + ) + + expected_layer_output = layer_hidden_states + layer_residual + torch.testing.assert_close(output, expected_layer_output) + torch.testing.assert_close(aux_hidden_states[0], initial_hidden_states) + torch.testing.assert_close(aux_hidden_states[1], expected_layer_output) + + +def test_kimi_linear_forward_extracts_attn_res_aux_hidden_states(monkeypatch): + model = _make_kimi_linear_model() + initial_hidden_states = torch.tensor([[1.0, 2.0]]) + layer_hidden_states = torch.tensor([[3.0, 4.0]]) + prefix_sum = torch.tensor([[5.0, 6.0]]) + block_residual = torch.tensor([[[7.0, 8.0]]]) + final_hidden_states = torch.tensor([[9.0, 10.0]]) + + object.__setattr__(model, "start_layer", 0) + object.__setattr__(model, "end_layer", 1) + object.__setattr__( + model, + "layers", + [Mock(return_value=(layer_hidden_states, prefix_sum, block_residual))], + ) + object.__setattr__(model, "aux_hidden_state_layers", (0, 1)) + object.__setattr__(model, "use_attn_res", True) + object.__setattr__(model, "num_attn_res_blocks", 1) + object.__setattr__( + model, + "output_attn_res_norm", + SimpleNamespace(weight=torch.ones(2), variance_epsilon=1e-5), + ) + object.__setattr__( + model, + "output_attn_res_proj", + SimpleNamespace(weight=torch.ones(1, 2)), + ) + monkeypatch.setattr( + kimi_model, + "get_pp_group", + lambda: SimpleNamespace(is_first_rank=True, is_last_rank=True), + ) + final_attn_res = Mock(return_value=final_hidden_states) + monkeypatch.setattr(kimi_model, "attn_res", final_attn_res) + + output, aux_hidden_states = model.forward( + input_ids=None, + positions=torch.tensor([0]), + intermediate_tensors=None, + inputs_embeds=initial_hidden_states, + ) + + torch.testing.assert_close(output, final_hidden_states) + torch.testing.assert_close(aux_hidden_states[0], initial_hidden_states) + torch.testing.assert_close(aux_hidden_states[1], prefix_sum + layer_hidden_states) + assert final_attn_res.call_args.args[2] is block_residual diff --git a/tests/models/kimi_k3/test_kda.py b/tests/models/kimi_k3/test_kda.py new file mode 100644 index 0000000000..d0860ed14a --- /dev/null +++ b/tests/models/kimi_k3/test_kda.py @@ -0,0 +1,736 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Precision tests for vllm's chunk_kda Triton operator. + +Compares chunk_kda against a naive recurrent reference (float32). +Uses torch.rand for q/k/v to match FLA's test pattern. +""" + +import pytest +import torch +import torch.nn.functional as F + +from aphrodite import _custom_ops as ops +from aphrodite.model_executor.layers.mamba.ops.causal_conv1d import causal_conv1d_update +from aphrodite.model_executor.layers.mamba.ops.gather_initial_states import ( + gather_initial_states, +) +from aphrodite.models.kimi_k3.nvidia.kda import ( + is_flashkda_supported, + is_fused_kda_decode_supported, +) +from aphrodite.models.kimi_k3.nvidia.ops.third_party.kda import ( + chunk_kda, + chunk_kda_with_fused_gate, + fused_kda_gate, + fused_recurrent_kda, + fused_recurrent_kda_fwd, + fused_recurrent_kda_packed_decode, +) +from aphrodite.third_party.flash_linear_attention.ops.l2norm import l2norm_fwd + +DEVICE = "cuda" + + +@torch.inference_mode() +def test_gather_initial_states_correctness(): + row_size = 8 * 128 * 128 + storage = torch.randn(5, row_size + 256, dtype=torch.float32, device=DEVICE) + state = storage[:, :row_size].view(5, 8, 128, 128) + assert not state.is_contiguous() + assert state[0].is_contiguous() + indices = torch.tensor([4, 1, 3], dtype=torch.int32, device=DEVICE) + has_initial_state = torch.tensor([True, False, True], device=DEVICE) + + expected = state[indices].clone() + expected[~has_initial_state] = 0 + + torch.testing.assert_close( + gather_initial_states(state, indices, has_initial_state), + expected, + ) + + +def naive_recurrent_kda( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + scale: float | None = None, + initial_state: torch.Tensor | None = None, + output_final_state: bool = False, +) -> tuple[torch.Tensor, torch.Tensor | None]: + """Naive recurrent KDA reference, ported from FLA's naive.py.""" + dtype = v.dtype + B, T, H, K = q.shape + V = v.shape[-1] + if scale is None: + scale = K**-0.5 + + q, k, v, g, beta = (x.to(torch.float) for x in [q, k, v, g, beta]) + q = q * scale + + S = k.new_zeros(B, H, K, V).to(q) + if initial_state is not None: + S += initial_state + o = torch.zeros_like(v) + for i in range(T): + q_i, k_i, v_i, g_i, b_i = q[:, i], k[:, i], v[:, i], g[:, i], beta[:, i] + S = S * g_i[..., None].exp() + S = S + torch.einsum( + "bhk,bhv->bhkv", + b_i[..., None] * k_i, + v_i - (k_i[..., None] * S).sum(-2), + ) + o[:, i] = torch.einsum("bhk,bhkv->bhv", q_i, S) + if not output_final_state: + S = None + return o.to(dtype), S + + +def assert_close( + name: str, + ref: torch.Tensor, + tri: torch.Tensor, + ratio: float, + err_atol: float = 1e-6, +): + """RMSE-based relative error comparison.""" + abs_err = (ref.detach() - tri.detach()).flatten().abs().max().item() + rmse_diff = (ref.detach() - tri.detach()).flatten().square().mean().sqrt().item() + rmse_base = ref.detach().flatten().square().mean().sqrt().item() + rel_err = rmse_diff / (rmse_base + 1e-8) + print(f"{name:>4} | abs={abs_err:.6f} | rmse={rel_err:.6f} | thr={ratio}") + if abs_err <= err_atol: + return + assert not torch.isnan(ref).any(), f"{name}: NaN detected in ref" + assert not torch.isnan(tri).any(), f"{name}: NaN detected in tri" + assert rel_err < ratio, f"{name}: max abs err {abs_err:.6f}, rmse ratio {rel_err:.6f} >= {ratio}" + + +@pytest.mark.parametrize( + ("H", "D", "cu_seqlens", "dtype"), + [ + pytest.param( + *test, + id="H{}-D{}-cu{}-{}".format(*test), + ) + for test in [ + (32, 128, [0, 64], torch.float16), + (32, 128, [0, 1024], torch.float16), + (32, 128, [0, 15], torch.float16), + (32, 128, [0, 256, 512, 768, 1024], torch.float16), + (32, 128, [0, 15, 100, 300, 1200], torch.float16), + (64, 128, [0, 256, 500, 1000], torch.float16), + (32, 128, [0, 8192], torch.float16), + (32, 128, [0, 256, 500, 1000], torch.bfloat16), + ] + ], +) +@torch.inference_mode() +def test_chunk_kda( + H: int, + D: int, + cu_seqlens: list[int], + dtype: torch.dtype, +): + T = cu_seqlens[-1] + torch.manual_seed(42) + B = 1 + cu_seqlens_t = torch.LongTensor(cu_seqlens).to(DEVICE) + N = len(cu_seqlens) - 1 + + q = torch.rand(B, T, H, D, dtype=dtype, device=DEVICE) + k = torch.rand(B, T, H, D, dtype=dtype, device=DEVICE) + v = torch.rand(B, T, H, D, dtype=dtype, device=DEVICE) + g = F.logsigmoid(torch.randn(B, T, H, D, dtype=torch.float32, device=DEVICE)).to(dtype) + beta = torch.rand(B, T, H, dtype=dtype, device=DEVICE).sigmoid() + h0 = torch.randn(N, H, D, D, dtype=torch.float32, device=DEVICE) + + # Naive reference with l2norm_fwd (same kernel as chunk_kda) + ref_outputs = [] + ref_states = [] + for i in range(N): + s, e = cu_seqlens[i], cu_seqlens[i + 1] + q_i = l2norm_fwd(q[:, s:e].contiguous()) + k_i = l2norm_fwd(k[:, s:e].contiguous()) + o_i, ht_i = naive_recurrent_kda( + q_i, + k_i, + v[:, s:e], + g[:, s:e], + beta[:, s:e], + initial_state=h0[i], + output_final_state=True, + ) + ref_outputs.append(o_i) + ref_states.append(ht_i) + ref_o = torch.cat(ref_outputs, dim=1) + ref_ht = torch.cat(ref_states, dim=0) + + # h0 transposed to (V, K) layout for the kernel; naive uses (K, V) + tri_o, tri_ht = chunk_kda( + q=q.clone(), + k=k.clone(), + v=v.clone(), + g=g.clone(), + beta=beta.clone(), + initial_state=h0.transpose(-1, -2).contiguous().clone(), + output_final_state=True, + cu_seqlens=cu_seqlens_t, + use_qk_l2norm_in_kernel=True, + ) + + assert not torch.isnan(tri_o).any(), "Triton output o contains NaN" + assert not torch.isnan(tri_ht).any(), "Triton output ht contains NaN" + assert_close("o", ref_o, tri_o, 0.005) + assert_close("ht", ref_ht, tri_ht.transpose(-1, -2).contiguous(), 0.005) + + +@pytest.mark.parametrize( + ("cu_seqlens", "dtype", "lower_bound"), + [ + ([0, 64], torch.float16, None), + ([0, 15, 100, 300], torch.bfloat16, None), + ([0, 15, 100, 300], torch.bfloat16, -3.0), + ], +) +@torch.inference_mode() +def test_chunk_kda_fused_gate_cumsum_matches_unfused( + cu_seqlens: list[int], + dtype: torch.dtype, + lower_bound: float | None, +): + H, D = 8, 64 + T = cu_seqlens[-1] + N = len(cu_seqlens) - 1 + torch.manual_seed(123) + + cu_seqlens_t = torch.tensor(cu_seqlens, dtype=torch.int32, device=DEVICE) + q = torch.randn(1, T, H, D, dtype=dtype, device=DEVICE) + k = torch.randn(1, T, H, D, dtype=dtype, device=DEVICE) + v = torch.randn(1, T, H, D, dtype=dtype, device=DEVICE) + raw_g = torch.randn(1, T, H, D, dtype=dtype, device=DEVICE) + beta_storage = torch.randn(1, T, 2 * H + 3, dtype=dtype, device=DEVICE) + raw_beta = beta_storage[..., 1 : 2 * H + 1 : 2] + beta = raw_beta.float().sigmoid() + A_log = (torch.randn(H, dtype=torch.float32, device=DEVICE) * 0.5).contiguous() + dt_bias = (torch.randn(H * D, dtype=torch.float32, device=DEVICE) * 0.1).contiguous() + h0 = torch.randn(N, H, D, D, dtype=torch.float32, device=DEVICE) + initial_state = h0.transpose(-1, -2).contiguous() + + gate = fused_kda_gate( + raw_g.reshape(T, H * D), + A_log, + D, + g_bias=dt_bias, + lower_bound=lower_bound, + ) + if lower_bound is not None: + expected_gate = lower_bound * torch.sigmoid( + A_log.exp()[None, :, None] * (raw_g.float().view(T, H, D) + dt_bias.view(H, D)) + ) + torch.testing.assert_close(gate, expected_gate) + gate = gate.unsqueeze(0) + old_o, old_ht = chunk_kda( + q=q.clone(), + k=k.clone(), + v=v.clone(), + g=gate, + beta=beta, + initial_state=initial_state.clone(), + output_final_state=True, + cu_seqlens=cu_seqlens_t, + use_qk_l2norm_in_kernel=True, + ) + new_o, new_ht = chunk_kda_with_fused_gate( + q=q.clone(), + k=k.clone(), + v=v.clone(), + raw_g=raw_g, + raw_beta=raw_beta, + A_log=A_log, + g_bias=dt_bias, + lower_bound=lower_bound, + initial_state=initial_state.clone(), + output_final_state=True, + cu_seqlens=cu_seqlens_t, + use_qk_l2norm_in_kernel=True, + ) + + assert_close("o", old_o, new_o, 1e-3, err_atol=1e-3) + assert_close("ht", old_ht, new_ht, 1e-3, err_atol=1e-3) + + +@pytest.mark.parametrize("num_seqs", [1, 8, 32]) +@pytest.mark.parametrize("lower_bound", [-5.0, None]) +@pytest.mark.parametrize("state_indices_stride", [1, 8]) +@torch.inference_mode() +def test_packed_kda_decode_correctness( + num_seqs: int, + lower_bound: float | None, + state_indices_stride: int, +): + H, D = 8, 128 + torch.manual_seed(321) + + packed_storage = torch.randn( + num_seqs, + 3 * H * D + 1, + dtype=torch.bfloat16, + device=DEVICE, + ) + mixed_qkv = packed_storage[:, : 3 * H * D] + assert mixed_qkv.stride(0) == 3 * H * D + 1 + q, k, v = (x.contiguous().view(1, num_seqs, H, D) for x in mixed_qkv.split(H * D, dim=-1)) + raw_g = torch.randn( + 1, + num_seqs, + H, + D, + dtype=torch.bfloat16, + device=DEVICE, + ) + raw_beta = torch.randn( + 1, + num_seqs, + H, + dtype=torch.bfloat16, + device=DEVICE, + ) + beta = raw_beta.float().sigmoid() + A_log = torch.randn(H, dtype=torch.float32, device=DEVICE) * 0.5 + dt_bias = torch.randn(H, D, dtype=torch.float32, device=DEVICE) * 0.1 + state_storage = torch.randn( + num_seqs + 1, + H * D * D + 17, + dtype=torch.float32, + device=DEVICE, + ) + state = state_storage[:, : H * D * D].view(num_seqs + 1, H, D, D) + assert not state.is_contiguous() + assert state.stride()[1:] == (D * D, D, 1) + state_indices_storage = torch.zeros( + num_seqs, + state_indices_stride, + dtype=torch.int32, + device=DEVICE, + ) + state_indices = state_indices_storage[:, 0] + state_indices.copy_( + torch.arange( + 1, + num_seqs + 1, + dtype=torch.int32, + device=DEVICE, + ) + ) + gate = fused_kda_gate( + raw_g.reshape(num_seqs, H * D), + A_log, + D, + g_bias=dt_bias, + lower_bound=lower_bound, + ).unsqueeze(0) + dense_state = state.clone() + dense_out, _ = fused_recurrent_kda_fwd( + q=q, + k=k, + v=v, + g=gate, + beta=beta, + scale=D**-0.5, + initial_state=dense_state, + inplace_final_state=True, + cu_seqlens=torch.arange( + num_seqs + 1, + dtype=torch.int32, + device=DEVICE, + ), + ssm_state_indices=state_indices, + use_qk_l2norm_in_kernel=True, + ) + packed_state = state + packed_out, _ = fused_recurrent_kda_packed_decode( + mixed_qkv=mixed_qkv, + raw_g=raw_g, + raw_beta=raw_beta, + A_log=A_log, + dt_bias=dt_bias, + lower_bound=lower_bound, + initial_state=packed_state, + state_indices=state_indices, + ) + + assert_close("o", dense_out, packed_out, 1e-3, err_atol=1e-3) + assert_close("ht", dense_state, packed_state, 1e-3, err_atol=1e-3) + + +@pytest.mark.parametrize( + ("H", "fuse_gate"), + [(12, True), (12, False), (12, None), (96, None)], +) +@pytest.mark.parametrize("lower_bound", [-5.0, None]) +@torch.inference_mode() +def test_kda_spec_decode_correctness( + H: int, + fuse_gate: bool | None, + lower_bound: float | None, +): + num_seqs, query_len, D = 3, 3, 128 + T = num_seqs * query_len + torch.manual_seed(1234) + + qkv_storage = torch.randn( + 1, + T, + 3 * H * D + 7, + dtype=torch.bfloat16, + device=DEVICE, + ) + packed_qkv = qkv_storage[..., : 3 * H * D] + q, k, v = (x.view(1, T, H, D) for x in packed_qkv.split(H * D, dim=-1)) + gate_storage = torch.randn( + 1, + T, + H * D + 5, + dtype=torch.bfloat16, + device=DEVICE, + ) + raw_g = gate_storage[..., : H * D].view(1, T, H, D) + beta_storage = torch.randn( + 1, + T, + H + 1, + dtype=torch.bfloat16, + device=DEVICE, + ) + raw_beta = beta_storage[..., :H] + A_log = 0.5 * torch.randn(H, dtype=torch.float32, device=DEVICE) + dt_bias = 0.1 * torch.randn(H, D, dtype=torch.float32, device=DEVICE) + cu_seqlens = torch.arange( + 0, + T + 1, + query_len, + dtype=torch.int32, + device=DEVICE, + ) + state_indices = torch.arange( + 1, + T + 1, + dtype=torch.int32, + device=DEVICE, + ).view(num_seqs, query_len) + num_accepted_tokens = torch.tensor( + [1, 2, 3], + dtype=torch.int32, + device=DEVICE, + ) + state_storage = 0.01 * torch.randn( + T + 1, + H * D * D + 17, + dtype=torch.float32, + device=DEVICE, + ) + state = state_storage[:, : H * D * D].view(T + 1, H, D, D) + output_storage = torch.full( + (1, T, H * D + 11), + torch.nan, + dtype=torch.bfloat16, + device=DEVICE, + ) + output = output_storage[..., : H * D].view(1, T, H, D) + + gate = fused_kda_gate( + raw_g.contiguous().view(T, H * D), + A_log, + D, + g_bias=dt_bias, + lower_bound=lower_bound, + ).unsqueeze(0) + beta = raw_beta.float().sigmoid() + q_norm = l2norm_fwd(q.contiguous()) + k_norm = l2norm_fwd(k.contiguous()) + expected_state = state.clone() + expected_outputs = [] + for seq, accepted in enumerate(num_accepted_tokens.tolist()): + recurrent_state = expected_state[state_indices[seq, accepted - 1]].transpose(-1, -2) + start = seq * query_len + for token in range(query_len): + token_slice = slice(start + token, start + token + 1) + token_output, recurrent_state = naive_recurrent_kda( + q_norm[:, token_slice], + k_norm[:, token_slice], + v[:, token_slice], + gate[:, token_slice], + beta[:, token_slice], + initial_state=recurrent_state, + output_final_state=True, + ) + assert recurrent_state is not None + expected_outputs.append(token_output) + expected_state[state_indices[seq, token]] = recurrent_state.transpose(-1, -2) + expected = torch.cat(expected_outputs, dim=1) + + actual_state = state.clone() + actual, _ = fused_recurrent_kda( + q=q, + k=k, + v=v, + raw_g=raw_g, + raw_beta=raw_beta, + A_log=A_log, + dt_bias=dt_bias, + lower_bound=lower_bound, + initial_state=actual_state, + cu_seqlens=cu_seqlens, + ssm_state_indices=state_indices, + num_accepted_tokens=num_accepted_tokens, + out=output, + fuse_gate=fuse_gate, + ) + + assert actual.data_ptr() == output.data_ptr() + assert_close("o", expected, actual, 1e-3, err_atol=1e-3) + used_states = state_indices.flatten().long() + assert_close( + "ht", + expected_state[used_states], + actual_state[used_states], + 3e-3, + err_atol=3e-3, + ) + assert torch.isnan(output_storage[..., H * D :]).all() + + +@pytest.mark.parametrize( + ("num_heads", "num_seqs", "lower_bound", "fuse_output_norm"), + [ + (12, 1, -5.0, True), + (12, 4, None, False), + (24, 4, None, False), + (48, 1, -5.0, True), + (96, 1, -5.0, True), + ], +) +@torch.inference_mode() +def test_fused_kda_decode_correctness( + num_heads: int, + num_seqs: int, + lower_bound: float | None, + fuse_output_norm: bool, +): + D, W = 128, 4 + if not is_fused_kda_decode_supported( + num_heads, + D, + W, + num_spec=0, + input_dtype=torch.bfloat16, + conv_state_dtype=torch.bfloat16, + ): + pytest.skip("Fused KDA decode is not supported on this platform") + torch.manual_seed(967 + num_heads + num_seqs) + dim = num_heads * D + slots = num_seqs + 2 + packed_x_storage = torch.randn(num_seqs, 3 * dim + 17, dtype=torch.bfloat16, device=DEVICE) + packed_x = packed_x_storage[:, : 3 * dim] + weight = 0.1 * torch.randn(3 * dim, W, dtype=torch.float32, device=DEVICE) + conv_seed = 0.1 * torch.randn( + slots, + W - 1, + 3 * dim, + dtype=torch.bfloat16, + device=DEVICE, + ).transpose(1, 2) + raw_g = torch.randn( + 1, + num_seqs, + num_heads, + D, + dtype=torch.bfloat16, + device=DEVICE, + ) + raw_beta_storage = torch.randn( + 1, + num_seqs, + num_heads + 1, + dtype=torch.bfloat16, + device=DEVICE, + ) + raw_beta = raw_beta_storage[:, :, :num_heads] + output_gate_storage = torch.randn( + num_seqs, + dim + 7, + dtype=torch.bfloat16, + device=DEVICE, + ) + output_gate = output_gate_storage[:, :dim].view(num_seqs, num_heads, D) + norm_weight = torch.randn(D, dtype=torch.float32, device=DEVICE) + norm_eps = 1e-5 + A_log = 0.5 * torch.randn(num_heads, dtype=torch.float32, device=DEVICE) + dt_bias = 0.1 * torch.randn(dim, dtype=torch.float32, device=DEVICE) + state_indices = torch.arange( + num_seqs, + 0, + -1, + dtype=torch.int32, + device=DEVICE, + ) + state_seed = 0.01 * torch.randn( + slots, + num_heads, + D, + D, + dtype=torch.float32, + device=DEVICE, + ) + + conv_ref = conv_seed.clone() + state_ref = state_seed.clone() + mixed_qkv = causal_conv1d_update( + packed_x, + conv_ref, + weight, + activation="silu", + conv_state_indices=state_indices, + validate_data=True, + out=torch.empty_like(packed_x), + ) + expected, _ = fused_recurrent_kda_packed_decode( + mixed_qkv=mixed_qkv, + raw_g=raw_g, + raw_beta=raw_beta, + A_log=A_log, + dt_bias=dt_bias, + lower_bound=lower_bound, + initial_state=state_ref, + state_indices=state_indices, + ) + if fuse_output_norm: + expected_float = expected.float() + expected = ( + expected_float + * torch.rsqrt(expected_float.square().mean(dim=-1, keepdim=True) + norm_eps) + * norm_weight + * output_gate.float().sigmoid().unsqueeze(0) + ).to(expected.dtype) + + conv_slot_elements = 3 * dim * (W - 1) + state_slot_elements = num_heads * D * D + conv_slot_bytes = conv_slot_elements * torch.bfloat16.itemsize + page_bytes = conv_slot_bytes + state_slot_elements * torch.float32.itemsize + cache_storage = torch.empty(slots * page_bytes, dtype=torch.uint8, device=DEVICE) + conv_actual = torch.as_strided( + cache_storage.view(torch.bfloat16), + size=(slots, 3 * dim, W - 1), + stride=(page_bytes // torch.bfloat16.itemsize, 1, 3 * dim), + ) + state_actual = torch.as_strided( + cache_storage.view(torch.float32), + size=(slots, num_heads, D, D), + stride=(page_bytes // torch.float32.itemsize, D * D, D, 1), + storage_offset=conv_slot_bytes // torch.float32.itemsize, + ) + conv_actual.copy_(conv_seed) + state_actual.copy_(state_seed) + fused_weight = weight.reshape(3, dim, W).transpose(1, 2).contiguous() + actual = ops.fused_kda_decode( + x=packed_x, + weight=fused_weight, + bias=None, + conv_state=conv_actual, + raw_g=raw_g, + raw_beta=raw_beta, + A_log=A_log, + dt_bias=dt_bias, + state_indices=state_indices, + state=state_actual, + lower_bound=lower_bound, + output_gate=output_gate if fuse_output_norm else None, + norm_weight=norm_weight if fuse_output_norm else None, + norm_eps=norm_eps, + ) + + torch.testing.assert_close(actual, expected, atol=3e-2, rtol=3e-2) + torch.testing.assert_close(conv_actual, conv_ref, atol=0, rtol=0) + torch.testing.assert_close(state_actual, state_ref, atol=3e-2, rtol=3e-2) + + +def test_fused_kda_decode_rejects_speculative_conv_state(): + assert not is_fused_kda_decode_supported( + num_heads=12, + head_dim=128, + conv_width=4, + num_spec=2, + input_dtype=torch.bfloat16, + conv_state_dtype=torch.bfloat16, + ) + + +@torch.inference_mode() +def test_flashkda_correctness(): + if not is_flashkda_supported(128, torch.bfloat16, -3.0): + pytest.skip("FlashKDA is not supported on this platform") + + import aphrodite._flashkda_C # noqa: F401 + + B, T, H, D = 1, 48, 2, 128 + torch.manual_seed(11) + q, k, v, raw_g = [torch.randn(B, T, H, D, dtype=torch.bfloat16, device=DEVICE) for _ in range(4)] + beta_logits = torch.randn(B, T, H, dtype=torch.bfloat16, device=DEVICE) + A_log = torch.randn(H, dtype=torch.float32, device=DEVICE) * 0.5 + dt_bias = torch.randn(H, D, dtype=torch.float32, device=DEVICE) * 0.1 + initial_state = torch.randn(2, H, D, D, dtype=torch.float32, device=DEVICE) + cu_seqlens = torch.tensor([0, 17, T], dtype=torch.int32, device=DEVICE) + lower_bound = -3.0 + + gate = lower_bound * torch.sigmoid(A_log.exp()[None, None, :, None] * (raw_g.float() + dt_bias[None, None, :, :])) + beta = beta_logits.float().sigmoid() + q_norm = l2norm_fwd(q.contiguous()) + k_norm = l2norm_fwd(k.contiguous()) + + expected_outputs = [] + expected_states = [] + for i, (start, end) in enumerate(zip(cu_seqlens[:-1].tolist(), cu_seqlens[1:].tolist())): + output, final_state = naive_recurrent_kda( + q_norm[:, start:end], + k_norm[:, start:end], + v[:, start:end], + gate[:, start:end], + beta[:, start:end], + initial_state=initial_state[i].transpose(-1, -2), + output_final_state=True, + ) + expected_outputs.append(output) + expected_states.append(final_state) + expected_out = torch.cat(expected_outputs, dim=1) + expected_state = torch.cat(expected_states).transpose(-1, -2).contiguous() + + actual_out = torch.empty_like(v) + actual_state = torch.empty_like(initial_state) + workspace = torch.empty( + torch.ops._flashkda_C.get_workspace_size(T, H, cu_seqlens.numel() - 1), + dtype=torch.uint8, + device=DEVICE, + ) + torch.ops._flashkda_C.fwd( + q, + k, + v, + raw_g, + beta_logits, + D**-0.5, + actual_out, + workspace, + A_log, + dt_bias, + lower_bound, + initial_state, + actual_state, + cu_seqlens, + ) + + assert_close("o", expected_out, actual_out, 0.01) + assert_close("ht", expected_state, actual_state, 0.01) diff --git a/tests/models/kimi_k3/test_kda_metadata.py b/tests/models/kimi_k3/test_kda_metadata.py new file mode 100644 index 0000000000..4c70ea48d5 --- /dev/null +++ b/tests/models/kimi_k3/test_kda_metadata.py @@ -0,0 +1,390 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from dataclasses import fields + +import pytest +import torch + +from aphrodite.config import SpeculativeConfig +from aphrodite.config.compilation import CUDAGraphMode +from aphrodite.models.kimi_k3.nvidia.kda_metadata import ( + KimiK3KDAAttentionBackend, + KimiK3KDAMetadata, + KimiK3KDAMetadataBuilder, + _mamba_get_block_table_tensor, + stage_spec_decode_metadata, +) +from aphrodite.v1.attention.backend import AttentionMetadataBuilder +from aphrodite.v1.attention.backends.gdn_attn import ( + GDNAttentionBackend, + GDNAttentionMetadata, + GDNAttentionMetadataBuilder, +) +from aphrodite.v1.attention.backends.utils import ( + NULL_BLOCK_ID, + mamba_get_block_table_tensor, +) +from aphrodite.v1.kv_cache_interface import MambaSpec +from tests.v1.attention.utils import ( + BatchSpec, + create_aphrodite_config, + create_common_attn_metadata, +) + +BLOCK_SIZE = 16 +DEVICE = torch.device("cpu") +PRUNED_METADATA_FIELDS = { + "chunk_indices", + "chunk_offsets", + "prefill_query_start_loc", + "prefill_state_indices", + "prefill_has_initial_state", + "spec_sequence_masks", +} + + +def _assert_matches_shared_gdn(reference, actual: KimiK3KDAMetadata): + for field in fields(KimiK3KDAMetadata): + actual_value = getattr(actual, field.name) + expected_value = getattr(reference, field.name) + if field.name in PRUNED_METADATA_FIELDS: + assert actual_value is None + continue + if ( + field.name in {"spec_token_indx", "non_spec_token_indx"} + and actual.num_spec_decodes > 0 + and actual.num_prefills == 0 + and actual.num_decodes == 0 + ): + assert actual_value is None + continue + if isinstance(actual_value, torch.Tensor): + torch.testing.assert_close(actual_value, expected_value) + elif field.name == "nums_dict": + assert (actual_value is None) == (expected_value is None) + if actual_value is not None: + assert actual_value[8]["tot"] == expected_value[8]["tot"] + torch.testing.assert_close(actual_value[8]["nums"], expected_value[8]["nums"]) + else: + assert actual_value == expected_value + + +def _make_builder( + builder_cls: type[AttentionMetadataBuilder], + num_speculative_tokens: int, + full_cuda_graph: bool, + device: torch.device = DEVICE, + mamba_cache_mode: str = "none", +) -> AttentionMetadataBuilder: + aphrodite_config = create_aphrodite_config( + model_name="Qwen/Qwen3.5-0.8B", + block_size=BLOCK_SIZE, + ) + if num_speculative_tokens: + aphrodite_config.speculative_config = SpeculativeConfig( + method="ngram", + num_speculative_tokens=num_speculative_tokens, + ) + aphrodite_config.compilation_config.cudagraph_mode = ( + CUDAGraphMode.FULL_AND_PIECEWISE if full_cuda_graph else CUDAGraphMode.NONE + ) + aphrodite_config.cache_config.mamba_cache_mode = mamba_cache_mode + return builder_cls( + kv_cache_spec=MambaSpec( + block_size=BLOCK_SIZE, + shapes=((16, 64),), + dtypes=(torch.float16,), + num_speculative_blocks=num_speculative_tokens, + ), + layer_names=["layer.0"], + aphrodite_config=aphrodite_config, + device=device, + ) + + +@pytest.mark.parametrize( + ( + "batch", + "num_decode_draft_tokens", + "num_speculative_tokens", + "full_cuda_graph", + "is_prefilling", + ), + [ + pytest.param( + BatchSpec(seq_lens=[50, 30], query_lens=[3, 3]), + [2, 2], + 2, + False, + [False, False], + id="pure-spec-decode", + ), + pytest.param( + BatchSpec(seq_lens=[100, 65, 20], query_lens=[50, 1, 3]), + [-1, -1, 2], + 2, + False, + [True, False, False], + id="mixed-prefill-and-spec-decode", + ), + pytest.param( + BatchSpec(seq_lens=[40, 30], query_lens=[1, 1]), + None, + 0, + False, + [False, False], + id="regular-decode", + ), + pytest.param( + BatchSpec(seq_lens=[40, 30], query_lens=[1, 1]), + [0, 0], + 2, + False, + [False, False], + id="no-scheduled-draft-tokens", + ), + ], +) +def test_kimi_k3_kda_metadata_matches_shared_gdn( + batch: BatchSpec, + num_decode_draft_tokens: list[int] | None, + num_speculative_tokens: int, + full_cuda_graph: bool, + is_prefilling: list[bool], +): + kwargs: dict[str, torch.Tensor] = {} + if num_decode_draft_tokens is not None: + kwargs = { + "num_decode_draft_tokens_cpu": torch.tensor(num_decode_draft_tokens, dtype=torch.int32), + "num_accepted_tokens": torch.ones(batch.batch_size, dtype=torch.int32, device=DEVICE), + } + + common_attn_metadata = create_common_attn_metadata(batch, BLOCK_SIZE, DEVICE).replace( + is_prefilling=torch.tensor(is_prefilling, dtype=torch.bool) + ) + reference = _make_builder( + GDNAttentionMetadataBuilder, + num_speculative_tokens, + full_cuda_graph, + ).build( + 0, + common_attn_metadata, + **kwargs, + ) + actual = _make_builder( + KimiK3KDAMetadataBuilder, + num_speculative_tokens, + full_cuda_graph, + ).build(0, common_attn_metadata, **kwargs) + + assert isinstance(actual, KimiK3KDAMetadata) + _assert_matches_shared_gdn(reference, actual) + + +def test_mixed_regular_and_spec_decode_uses_packed_decode_metadata(): + batch = BatchSpec(seq_lens=[100, 65, 20], query_lens=[1, 1, 3]) + common_attn_metadata = create_common_attn_metadata(batch, BLOCK_SIZE, DEVICE).replace( + is_prefilling=torch.tensor([False, False, False]) + ) + actual = _make_builder( + KimiK3KDAMetadataBuilder, + num_speculative_tokens=2, + full_cuda_graph=False, + ).build( + 0, + common_attn_metadata, + num_decode_draft_tokens_cpu=torch.tensor([-1, -1, 2], dtype=torch.int32), + num_accepted_tokens=torch.ones(3, dtype=torch.int32, device=DEVICE), + ) + + # The K3 layer dispatches the non-spec subgroup to packed decode whenever + # it contains no prefill request. + assert actual.num_decodes == 2 + assert actual.num_decode_tokens == 2 + assert actual.num_prefills == 0 + assert actual.num_prefill_tokens == 0 + assert actual.has_initial_state is None + assert actual.nums_dict is None + assert actual.non_spec_query_start_loc is None + torch.testing.assert_close(actual.non_spec_token_indx, torch.tensor([0, 1])) + torch.testing.assert_close(actual.spec_token_indx, torch.tensor([2, 3, 4])) + torch.testing.assert_close( + actual.spec_query_start_loc, + torch.tensor([0, 3], dtype=torch.int32), + ) + + +def test_mixed_regular_and_spec_decode_excludes_request_padding(): + batch = BatchSpec(seq_lens=[16, 65, 20], query_lens=[0, 1, 3]) + common_attn_metadata = create_common_attn_metadata(batch, BLOCK_SIZE, DEVICE).replace( + is_prefilling=torch.tensor([False, False, False]) + ) + actual = _make_builder( + KimiK3KDAMetadataBuilder, + num_speculative_tokens=2, + full_cuda_graph=False, + ).build( + 0, + common_attn_metadata, + num_decode_draft_tokens_cpu=torch.tensor([-1, -1, 2], dtype=torch.int32), + num_accepted_tokens=torch.ones(3, dtype=torch.int32, device=DEVICE), + ) + + assert actual.num_decodes == 1 + assert actual.non_spec_state_indices_tensor is not None + assert actual.non_spec_state_indices_tensor.shape == (1,) + torch.testing.assert_close(actual.non_spec_token_indx, torch.tensor([0])) + torch.testing.assert_close(actual.spec_token_indx, torch.tensor([1, 2, 3])) + + +@pytest.mark.parametrize( + ("seq_len", "expected_has_initial_state"), + [ + pytest.param(1, False, id="first-token-prefill"), + pytest.param(65, True, id="final-one-token-prefill-chunk"), + ], +) +def test_mixed_one_token_prefill_and_spec_decode_uses_prefill_metadata( + seq_len: int, + expected_has_initial_state: bool, +): + batch = BatchSpec(seq_lens=[seq_len, 20], query_lens=[1, 3]) + common_attn_metadata = create_common_attn_metadata(batch, BLOCK_SIZE, DEVICE).replace( + is_prefilling=torch.tensor([True, False]) + ) + actual = _make_builder( + KimiK3KDAMetadataBuilder, + num_speculative_tokens=2, + full_cuda_graph=False, + ).build( + 0, + common_attn_metadata, + num_decode_draft_tokens_cpu=torch.tensor([-1, 2], dtype=torch.int32), + num_accepted_tokens=torch.ones(2, dtype=torch.int32, device=DEVICE), + ) + + assert actual.num_prefills == 1 + assert actual.num_prefill_tokens == 1 + assert actual.num_decodes == 0 + assert actual.num_decode_tokens == 0 + assert actual.has_initial_state is not None + assert actual.has_initial_state.tolist() == [expected_has_initial_state] + assert actual.non_spec_query_start_loc is not None + torch.testing.assert_close( + actual.non_spec_query_start_loc, + torch.tensor([0, 1], dtype=torch.int32), + ) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +def test_kimi_k3_kda_cudagraph_capture_matches_shared_gdn(): + device = torch.device("cuda") + batch = BatchSpec(seq_lens=[50, 30], query_lens=[3, 3]) + common_attn_metadata = create_common_attn_metadata(batch, BLOCK_SIZE, device).replace( + is_prefilling=torch.tensor([False, False]) + ) + reference = _make_builder( + GDNAttentionMetadataBuilder, + num_speculative_tokens=2, + full_cuda_graph=True, + device=device, + ).build_for_cudagraph_capture(common_attn_metadata) + actual = _make_builder( + KimiK3KDAMetadataBuilder, + num_speculative_tokens=2, + full_cuda_graph=True, + device=device, + ).build_for_cudagraph_capture(common_attn_metadata) + + assert isinstance(actual, KimiK3KDAMetadata) + _assert_matches_shared_gdn(reference, actual) + + +def test_kimi_k3_kda_backend_uses_private_metadata_builder(): + assert KimiK3KDAAttentionBackend.get_builder_cls() is KimiK3KDAMetadataBuilder + assert KimiK3KDAAttentionBackend.is_ssm() + assert issubclass(KimiK3KDAAttentionBackend, GDNAttentionBackend) + assert issubclass(KimiK3KDAMetadata, GDNAttentionMetadata) + assert issubclass(KimiK3KDAMetadataBuilder, GDNAttentionMetadataBuilder) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +def test_stage_spec_decode_metadata_matches_pytorch(): + device = torch.device("cuda") + num_spec_decodes = 33 + batch_size = 65 + num_state_slots = 3 + state_indices = torch.arange( + num_spec_decodes * 32, + dtype=torch.int32, + device=device, + ).reshape(num_spec_decodes, 32)[:, :num_state_slots] + query_start_loc = torch.arange(num_spec_decodes + 1, dtype=torch.int32, device=device) * num_state_slots + num_accepted_tokens = torch.arange(num_spec_decodes, dtype=torch.int32, device=device) % num_state_slots + 1 + + staged_state_indices = torch.empty((batch_size, num_state_slots), dtype=torch.int32, device=device) + staged_query_start_loc = torch.empty(batch_size + 1, dtype=torch.int32, device=device) + staged_num_accepted_tokens = torch.empty(batch_size, dtype=torch.int32, device=device) + stage_spec_decode_metadata( + state_indices, + query_start_loc, + num_accepted_tokens, + staged_state_indices, + staged_query_start_loc, + staged_num_accepted_tokens, + num_spec_decodes=num_spec_decodes, + ) + + expected_state_indices = torch.full_like(staged_state_indices, NULL_BLOCK_ID) + expected_state_indices[:num_spec_decodes] = state_indices + expected_query_start_loc = torch.full( + (batch_size + 1,), + query_start_loc[-1], + dtype=torch.int32, + device=device, + ) + expected_query_start_loc[: num_spec_decodes + 1] = query_start_loc + expected_num_accepted_tokens = torch.ones(batch_size, dtype=torch.int32, device=device) + expected_num_accepted_tokens[:num_spec_decodes] = num_accepted_tokens + + torch.testing.assert_close(staged_state_indices, expected_state_indices) + torch.testing.assert_close(staged_query_start_loc, expected_query_start_loc) + torch.testing.assert_close(staged_num_accepted_tokens, expected_num_accepted_tokens) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +def test_aligned_block_table_matches_shared_gdn(): + device = torch.device("cuda") + seq_lens = torch.tensor( + [0, 1, 15, 16, 17, 31, 32, 33, 511, 512, 513], + dtype=torch.int32, + device=device, + ).repeat(6)[:65] + block_table_storage = torch.arange( + seq_lens.numel() * 128, + dtype=torch.int32, + device=device, + ).reshape(seq_lens.numel(), 128) + block_table = block_table_storage[:, ::2] + kv_cache_spec = MambaSpec( + block_size=BLOCK_SIZE, + shapes=((16, 64),), + dtypes=(torch.float16,), + num_speculative_blocks=2, + ) + + expected = mamba_get_block_table_tensor( + block_table, + seq_lens, + kv_cache_spec, + "align", + ) + actual = _mamba_get_block_table_tensor( + block_table, + seq_lens, + kv_cache_spec, + "align", + ) + + torch.testing.assert_close(actual, expected) diff --git a/tests/models/kimi_k3/test_latent_moe_tail.py b/tests/models/kimi_k3/test_latent_moe_tail.py new file mode 100644 index 0000000000..246c738ca2 --- /dev/null +++ b/tests/models/kimi_k3/test_latent_moe_tail.py @@ -0,0 +1,145 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import pytest +import ray +import torch +import torch.distributed as dist +import torch.nn.functional as F + +from aphrodite.distributed import get_tp_group +from aphrodite.model_executor.warmup.cutedsl_warmup import cutedsl_warmup +from aphrodite.models.kimi_k3.nvidia.ops.latent_moe_tail import KimiK3LatentMoETailOp +from aphrodite.platforms import current_platform +from tests.utils import ( + init_test_distributed_environment, + multi_gpu_test, + multi_process_parallel, +) + +HIDDEN_SIZE = 7168 +LATENT_SIZE = 3584 +EPS = 0.1 + + +@ray.remote(num_gpus=1, max_calls=1) +def _test_latent_moe_tail_worker( + monkeypatch: pytest.MonkeyPatch, + tp_size: int, + pp_size: int, + rank: int, + distributed_init_port: str, +) -> None: + monkeypatch.delenv("CUDA_VISIBLE_DEVICES", raising=False) + device = torch.device(f"cuda:{rank}") + torch.accelerator.set_device_index(device) + init_test_distributed_environment( + tp_size, + pp_size, + rank, + distributed_init_port, + ) + + torch.manual_seed(0) + rms_weight = 1 + 0.1 * torch.randn( + LATENT_SIZE, + device=device, + dtype=torch.bfloat16, + ) + up_weight = ( + torch.randn( + HIDDEN_SIZE, + LATENT_SIZE, + device=device, + dtype=torch.bfloat16, + ) + / LATENT_SIZE**0.5 + ) + + group = get_tp_group().device_group + op = KimiK3LatentMoETailOp.initialize( + hidden_size=HIDDEN_SIZE, + latent_size=LATENT_SIZE, + dtype=torch.bfloat16, + device=device, + rms_eps=EPS, + ) + cutedsl_warmup() + + for iteration, num_tokens in enumerate((1, 5, 8, 16, 5)): + torch.manual_seed(100 * iteration + rank + 1) + routed_output = torch.randn( + num_tokens, + LATENT_SIZE, + device=device, + dtype=torch.bfloat16, + ).mul_(0.01) + shared_output = torch.randn( + num_tokens, + HIDDEN_SIZE, + device=device, + dtype=torch.bfloat16, + ) + + routed_reference = routed_output.clone() + shared_reference = shared_output.clone() + dist.all_reduce(routed_reference, group=group) + dist.all_reduce(shared_reference, group=group) + expected = F.linear( + F.rms_norm( + routed_reference, + (LATENT_SIZE,), + rms_weight, + EPS, + ), + up_weight, + ) + expected.add_(shared_reference) + + actual = op( + routed_output, + shared_output, + rms_weight, + up_weight, + ) + torch.testing.assert_close(actual, expected, atol=8e-2, rtol=3e-2) + assert actual.is_contiguous() + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + graph_output = op( + routed_output, + shared_output, + rms_weight, + up_weight, + ) + graph.replay() + torch.testing.assert_close(graph_output, expected, atol=8e-2, rtol=3e-2) + + +def _run_latent_moe_tail_test( + monkeypatch: pytest.MonkeyPatch, + tp_size: int, +) -> None: + if not current_platform.is_device_capability_family(100): + pytest.skip("K3 latent-MoE tail fusion requires SM100") + multi_process_parallel( + monkeypatch, + tp_size, + 1, + _test_latent_moe_tail_worker, + ) + + +@multi_gpu_test(num_gpus=8) +def test_latent_moe_tail_tp8_matches_native_path( + monkeypatch: pytest.MonkeyPatch, +) -> None: + _run_latent_moe_tail_test(monkeypatch, 8) + + +@multi_gpu_test(num_gpus=16) +def test_latent_moe_tail_tp16_matches_native_path( + monkeypatch: pytest.MonkeyPatch, +) -> None: + _run_latent_moe_tail_test(monkeypatch, 16) diff --git a/tests/models/kimi_k3/test_sequence_parallel.py b/tests/models/kimi_k3/test_sequence_parallel.py new file mode 100644 index 0000000000..860f62feb0 --- /dev/null +++ b/tests/models/kimi_k3/test_sequence_parallel.py @@ -0,0 +1,343 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from types import MethodType, SimpleNamespace +from unittest.mock import Mock + +import pytest +import torch +from torch import nn + +from aphrodite.config import ParallelConfig +from aphrodite.models.kimi_k3.nvidia import model as kimi_model +from aphrodite.models.kimi_k3.nvidia import mtp as kimi_mtp +from aphrodite.models.kimi_k3.nvidia.ops import sequence_parallel as sp_ops +from aphrodite.platforms import current_platform + + +class _IdentityNorm(nn.Module): + def __init__(self, hidden_size: int = 2) -> None: + super().__init__() + self.weight = nn.Parameter(torch.ones(hidden_size), requires_grad=False) + self.variance_epsilon = 1e-5 + + def forward( + self, + hidden_states: torch.Tensor, + residual: torch.Tensor | None = None, + ): + if residual is None: + return hidden_states + return hidden_states, residual + + +class _RecordingMoE(nn.Module): + def __init__(self) -> None: + super().__init__() + self.num_tokens = 0 + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + self.num_tokens = hidden_states.shape[0] + return hidden_states + + +class _Projection(nn.Module): + def __init__(self, hidden_size: int = 2) -> None: + super().__init__() + self.weight = nn.Parameter( + torch.ones(1, hidden_size), + requires_grad=False, + ) + + +class _SequenceParallelMTPBlock: + use_sequence_parallel = True + + def __call__( + self, + *, + positions: torch.Tensor, + hidden_states: torch.Tensor, + residual: torch.Tensor | None, + ): + assert residual is None + return hidden_states * 2, None, hidden_states * 3 + + +def _mock_sequence_parallel_collectives(monkeypatch): + monkeypatch.setattr( + kimi_model, + "sp_reduce_scatter", + lambda tensor: tensor.chunk(2, dim=0)[0], + ) + monkeypatch.setattr( + kimi_model, + "sp_shard", + lambda tensor: torch.nn.functional.pad(tensor, (0, 0, 0, 1))[:2], + ) + monkeypatch.setattr( + kimi_model, + "sp_all_gather", + lambda tensor: torch.cat([tensor, tensor], dim=0), + ) + + +@pytest.mark.parametrize( + ("num_tokens", "is_padding", "tp_rank", "expected"), + [ + (1, None, 0, [False]), + (1, None, 1, [True]), + (5, None, 2, [False, True]), + (5, None, 3, [True, True]), + (5, [False, True, False, False, False], 0, [False, True]), + ], +) +def test_sp_padding_mask_marks_added_rows( + monkeypatch, + num_tokens: int, + is_padding: list[bool] | None, + tp_rank: int, + expected: list[bool], +): + monkeypatch.setattr(sp_ops, "get_tensor_model_parallel_world_size", lambda: 4) + monkeypatch.setattr(sp_ops, "get_tensor_model_parallel_rank", lambda: tp_rank) + + hidden_states = torch.empty(num_tokens, 2) + padding = torch.tensor(is_padding) if is_padding is not None else None + actual = sp_ops.sp_padding_mask(padding, hidden_states) + + torch.testing.assert_close(actual, torch.tensor(expected)) + + +@pytest.mark.parametrize( + ("data_parallel_size", "expected"), + [ + (1, False), + (2, True), + ], +) +def test_moe_sequence_parallel_requires_data_parallel( + monkeypatch, + data_parallel_size: int, + expected: bool, +): + monkeypatch.setattr(current_platform, "device_count", lambda: 2) + parallel_config = ParallelConfig( + tensor_parallel_size=2, + data_parallel_size=data_parallel_size, + enable_expert_parallel=True, + all2all_backend="allgather_reducescatter", + ) + + assert parallel_config.use_sequence_parallel_moe is expected + + +def test_kimi_decoder_layer_keeps_moe_states_sequence_sharded(monkeypatch): + layer = object.__new__(kimi_model.KimiDecoderLayer) + nn.Module.__init__(layer) + layer.use_attn_res = False + layer.use_sequence_parallel = True + layer.input_layernorm = _IdentityNorm() + layer.post_attention_layernorm = _IdentityNorm() + layer.mlp = _RecordingMoE() + layer._run_self_attn = MethodType( + lambda self, positions, hidden_states: hidden_states, + layer, + ) + + _mock_sequence_parallel_collectives(monkeypatch) + + positions = torch.arange(3) + full_hidden_states = torch.arange(6, dtype=torch.float32).view(3, 2) + hidden_states = kimi_model.sp_shard(full_hidden_states) + hidden_states, prefix_sum, residual = layer( + positions=positions, + hidden_states=hidden_states, + residual=None, + ) + + assert prefix_sum is None + assert hidden_states.shape == residual.shape == (2, 2) + assert layer.mlp.num_tokens == 2 + + hidden_states, prefix_sum, residual = layer( + positions=positions, + hidden_states=hidden_states, + residual=residual, + ) + + assert prefix_sum is None + assert hidden_states.shape == residual.shape == (2, 2) + assert layer.mlp.num_tokens == 2 + + +def test_kimi_attn_residual_states_stay_sequence_sharded(monkeypatch): + layer = object.__new__(kimi_model.KimiDecoderLayer) + nn.Module.__init__(layer) + layer.use_attn_res = True + layer.use_sequence_parallel = True + layer.prev_valid_blocks = 0 + layer.block_write_idx = 0 + layer.is_block_write_layer = False + layer.input_layernorm = _IdentityNorm() + layer.post_attention_layernorm = _IdentityNorm() + layer.self_attention_res_norm = _IdentityNorm() + layer.mlp_res_norm = _IdentityNorm() + layer.self_attention_res_proj = _Projection() + layer.mlp_res_proj = _Projection() + layer.mlp = _RecordingMoE() + layer._run_self_attn = MethodType( + lambda self, positions, hidden_states: hidden_states, + layer, + ) + + _mock_sequence_parallel_collectives(monkeypatch) + monkeypatch.setattr( + kimi_model, + "attn_res", + lambda prefix_sum, hidden_states, *args, **kwargs: ( + prefix_sum if hidden_states is None else prefix_sum + hidden_states + ), + ) + + prefix_sum = kimi_model.sp_shard(torch.arange(6, dtype=torch.float32).view(3, 2)) + block_residual = torch.zeros(2, 1, 2) + hidden_states, prefix_sum, block_residual = layer( + positions=torch.arange(3), + hidden_states=None, + prefix_sum=prefix_sum, + residual=block_residual, + ) + + assert hidden_states.shape == prefix_sum.shape == (2, 2) + assert block_residual.shape == (2, 1, 2) + assert layer.mlp.num_tokens == 2 + + +def test_kimi_mtp_restores_sequence_parallel_output(monkeypatch): + layer = object.__new__(kimi_mtp.KimiK3MultiTokenPredictorLayer) + nn.Module.__init__(layer) + layer.enorm = _IdentityNorm() + layer.hnorm = _IdentityNorm() + layer.eh_proj = nn.Identity() + object.__setattr__(layer, "mtp_block", _SequenceParallelMTPBlock()) + + final_norm = Mock(side_effect=lambda hidden_states: hidden_states + 1) + object.__setattr__( + layer, + "shared_head", + SimpleNamespace(norm=final_norm), + ) + + monkeypatch.setattr( + kimi_mtp, + "fused_mtp_input", + lambda positions, inputs_embeds, *args: inputs_embeds, + ) + monkeypatch.setattr( + kimi_mtp, + "sp_shard", + lambda tensor: torch.nn.functional.pad(tensor, (0, 0, 0, 1))[:2], + ) + monkeypatch.setattr( + kimi_mtp, + "sp_all_gather", + lambda tensor: torch.cat([tensor, tensor], dim=0), + ) + + inputs_embeds = torch.arange(6, dtype=torch.float32).view(3, 2) + logits_hidden_states, hidden_states = layer( + input_ids=torch.zeros(3, dtype=torch.long), + positions=torch.arange(3), + previous_hidden_states=torch.zeros_like(inputs_embeds), + inputs_embeds=inputs_embeds, + ) + + sharded_states = torch.nn.functional.pad(inputs_embeds, (0, 0, 0, 1))[:2] + expected_hidden_states = torch.cat( + [sharded_states * 5, sharded_states * 5], + dim=0, + )[:3] + torch.testing.assert_close(hidden_states, expected_hidden_states) + torch.testing.assert_close(logits_hidden_states, expected_hidden_states + 1) + final_norm.assert_called_once() + torch.testing.assert_close(final_norm.call_args.args[0], expected_hidden_states) + + +def test_sp_all_gather_uses_custom_kernel(monkeypatch): + hidden_states = torch.arange(4, dtype=torch.float32).view(2, 2) + expected = torch.cat([hidden_states, hidden_states]) + custom_all_gather = Mock(return_value=expected) + device_communicator = SimpleNamespace( + custom_all_gather=custom_all_gather, + ) + monkeypatch.setattr( + sp_ops, + "get_tp_group", + lambda: SimpleNamespace(device_communicator=device_communicator), + ) + fallback = Mock(side_effect=AssertionError("unexpected fallback")) + monkeypatch.setattr(sp_ops, "tensor_model_parallel_all_gather", fallback) + + output = sp_ops.sp_all_gather(hidden_states) + + torch.testing.assert_close(output, expected) + custom_all_gather.assert_called_once_with(hidden_states) + fallback.assert_not_called() + + +def test_sp_reduce_scatter_uses_custom_kernel_after_padding(monkeypatch): + hidden_states = torch.arange(6, dtype=torch.float32).view(3, 2) + expected = torch.arange(4, dtype=torch.float32).view(2, 2) + custom_reduce_scatter = Mock(return_value=expected) + device_communicator = SimpleNamespace( + custom_reduce_scatter=custom_reduce_scatter, + ) + monkeypatch.setattr( + sp_ops, + "get_tp_group", + lambda: SimpleNamespace(device_communicator=device_communicator), + ) + monkeypatch.setattr( + sp_ops, + "get_tensor_model_parallel_world_size", + lambda: 2, + ) + fallback = Mock(side_effect=AssertionError("unexpected fallback")) + monkeypatch.setattr(sp_ops, "tensor_model_parallel_reduce_scatter", fallback) + + output = sp_ops.sp_reduce_scatter(hidden_states) + + torch.testing.assert_close(output, expected) + padded = custom_reduce_scatter.call_args.args[0] + assert padded.shape == (4, 2) + torch.testing.assert_close(padded[:3], hidden_states) + torch.testing.assert_close(padded[3], torch.zeros(2)) + fallback.assert_not_called() + + +def test_sp_collectives_fall_back_without_custom_kernel(monkeypatch): + hidden_states = torch.arange(4, dtype=torch.float32).view(2, 2) + monkeypatch.setattr( + sp_ops, + "get_tp_group", + lambda: SimpleNamespace(device_communicator=None), + ) + monkeypatch.setattr( + sp_ops, + "get_tensor_model_parallel_world_size", + lambda: 2, + ) + all_gather = Mock(return_value=hidden_states) + reduce_scatter = Mock(return_value=hidden_states) + monkeypatch.setattr(sp_ops, "tensor_model_parallel_all_gather", all_gather) + monkeypatch.setattr( + sp_ops, + "tensor_model_parallel_reduce_scatter", + reduce_scatter, + ) + + torch.testing.assert_close(sp_ops.sp_all_gather(hidden_states), hidden_states) + torch.testing.assert_close(sp_ops.sp_reduce_scatter(hidden_states), hidden_states) + all_gather.assert_called_once_with(hidden_states, 0) + reduce_scatter.assert_called_once_with(hidden_states, 0) diff --git a/tests/models/multimodal/generation/test_common.py b/tests/models/multimodal/generation/test_common.py index 859643ae1a..d06151fc4a 100644 --- a/tests/models/multimodal/generation/test_common.py +++ b/tests/models/multimodal/generation/test_common.py @@ -1207,7 +1207,7 @@ def test_custom_inputs_models( create_new_process_for_each_test=True, ), ) -@create_new_process_for_each_test() +@create_new_process_for_each_test("spawn") def test_single_image_models_heavy( tmp_path: PosixPath, model_type: str, diff --git a/tests/models/multimodal/processing/test_glm4_1v.py b/tests/models/multimodal/processing/test_glm4_1v.py index e202c4f465..e43e2cba95 100644 --- a/tests/models/multimodal/processing/test_glm4_1v.py +++ b/tests/models/multimodal/processing/test_glm4_1v.py @@ -61,6 +61,31 @@ def test_encoder_cudagraph_uses_model_video_frame_limit(): assert Glm4vForConditionalGeneration.get_max_frames_per_video(model) == 600 +@pytest.mark.parametrize( + ("temporal_patch_size", "expected_grid_t"), + [(2, 9), (4, 5), (8, 3)], +) +def test_vision_info_rounds_up_temporal_frames( + temporal_patch_size: int, + expected_grid_t: int, +): + info = Mock(spec=Glm4vProcessingInfo) + vision_config = info.get_hf_config.return_value.vision_config + vision_config.patch_size = 14 + vision_config.spatial_merge_size = 2 + vision_config.temporal_patch_size = temporal_patch_size + + _, num_vision_tokens = Glm4vProcessingInfo._get_vision_info( + info, + image_width=28, + image_height=28, + num_frames=17, + do_resize=False, + ) + + assert num_vision_tokens == expected_grid_t + + @pytest.mark.parametrize("model_id", ["zai-org/GLM-4.1V-9B-Thinking"]) @pytest.mark.parametrize("expected_toks_per_frame", [299]) @pytest.mark.parametrize( diff --git a/tests/models/quantization/test_mxfp8.py b/tests/models/quantization/test_mxfp8.py index 51f009bed9..80b66d66d0 100644 --- a/tests/models/quantization/test_mxfp8.py +++ b/tests/models/quantization/test_mxfp8.py @@ -17,7 +17,9 @@ """ import pytest +import torch +from aphrodite.platforms import current_platform from tests.quantization.utils import is_quant_method_supported from ..utils import check_logprobs_close @@ -77,6 +79,170 @@ def test_mxfp8_logprobs( ) +@pytest.mark.skipif( + not is_quant_method_supported("mxfp8"), + reason="mxfp8 is not supported on this GPU type (requires sm_100+).", +) +@pytest.mark.skipif( + not current_platform.is_rocm(), + reason="AITER MXFP8 MoE backend is ROCm-only.", +) +@pytest.mark.quant_model +def test_mxfp8_aiter_requires_swigluoai_activation( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from aphrodite.model_executor.layers.fused_moe.activation import MoEActivation + from aphrodite.model_executor.layers.fused_moe.config import ( + FusedMoEConfig, + FusedMoEParallelConfig, + RoutingMethodType, + ) + from aphrodite.model_executor.layers.fused_moe.experts import aiter_mxfp8_moe + from aphrodite.model_executor.layers.fused_moe.oracle.mxfp8 import ( + select_mxfp8_moe_backend, + ) + + monkeypatch.setattr( + aiter_mxfp8_moe.AiterMxfp8Experts, + "_supports_current_device", + staticmethod(lambda: True), + ) + monkeypatch.setattr( + aiter_mxfp8_moe, + "is_aiter_mxfp8_moe_available", + lambda: True, + ) + + config = FusedMoEConfig( + num_experts=8, + experts_per_token=2, + hidden_dim=256, + intermediate_size=256, + num_local_experts=8, + num_logical_experts=8, + moe_parallel_config=FusedMoEParallelConfig.make_no_parallel(), + activation=MoEActivation.SILU, + in_dtype=torch.bfloat16, + device="cuda", + routing_method=RoutingMethodType.Renormalize, + moe_backend="aiter", + ) + + with pytest.raises(ValueError, match="requires activation=swigluoai_uninterleave"): + select_mxfp8_moe_backend(config) + + +@pytest.mark.skipif( + not is_quant_method_supported("mxfp8"), + reason="mxfp8 is not supported on this GPU type (requires sm_100+).", +) +@pytest.mark.skipif( + not current_platform.is_rocm(), + reason="AITER MXFP8 MoE backend is ROCm-only.", +) +@pytest.mark.quant_model +def test_mxfp8_aiter_requires_swigluoai_params( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from aphrodite.model_executor.layers.fused_moe.activation import MoEActivation + from aphrodite.model_executor.layers.fused_moe.config import ( + FusedMoEConfig, + FusedMoEParallelConfig, + RoutingMethodType, + ) + from aphrodite.model_executor.layers.fused_moe.experts import aiter_mxfp8_moe + from aphrodite.model_executor.layers.fused_moe.oracle.mxfp8 import ( + select_mxfp8_moe_backend, + ) + + monkeypatch.setattr( + aiter_mxfp8_moe.AiterMxfp8Experts, + "_supports_current_device", + staticmethod(lambda: True), + ) + monkeypatch.setattr( + aiter_mxfp8_moe, + "is_aiter_mxfp8_moe_available", + lambda: True, + ) + + config = FusedMoEConfig( + num_experts=8, + experts_per_token=2, + hidden_dim=256, + intermediate_size=256, + num_local_experts=8, + num_logical_experts=8, + moe_parallel_config=FusedMoEParallelConfig.make_no_parallel(), + activation=MoEActivation.SWIGLUOAI_UNINTERLEAVE, + in_dtype=torch.bfloat16, + device="cuda", + routing_method=RoutingMethodType.Renormalize, + moe_backend="aiter", + ) + + with pytest.raises(ValueError, match="hardcodes SwiGLU-OAI"): + select_mxfp8_moe_backend(config) + + +@pytest.mark.skipif( + not is_quant_method_supported("mxfp8"), + reason="mxfp8 is not supported on this GPU type (requires sm_100+).", +) +@pytest.mark.skipif( + not current_platform.is_rocm(), + reason="AITER MXFP8 MoE backend is ROCm-only.", +) +@pytest.mark.quant_model +def test_mxfp8_aiter_accepts_swigluoai_params( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from aphrodite.model_executor.layers.fused_moe.activation import MoEActivation + from aphrodite.model_executor.layers.fused_moe.config import ( + FusedMoEConfig, + FusedMoEParallelConfig, + RoutingMethodType, + ) + from aphrodite.model_executor.layers.fused_moe.experts import aiter_mxfp8_moe + from aphrodite.model_executor.layers.fused_moe.oracle.fp8 import Fp8MoeBackend + from aphrodite.model_executor.layers.fused_moe.oracle.mxfp8 import ( + select_mxfp8_moe_backend, + ) + + monkeypatch.setattr( + aiter_mxfp8_moe.AiterMxfp8Experts, + "_supports_current_device", + staticmethod(lambda: True), + ) + monkeypatch.setattr( + aiter_mxfp8_moe, + "is_aiter_mxfp8_moe_available", + lambda: True, + ) + + config = FusedMoEConfig( + num_experts=8, + experts_per_token=2, + hidden_dim=256, + intermediate_size=256, + num_local_experts=8, + num_logical_experts=8, + moe_parallel_config=FusedMoEParallelConfig.make_no_parallel(), + activation=MoEActivation.SWIGLUOAI_UNINTERLEAVE, + in_dtype=torch.bfloat16, + device="cuda", + routing_method=RoutingMethodType.Renormalize, + moe_backend="aiter", + swiglu_alpha=aiter_mxfp8_moe._AITER_SWIGLU_ALPHA, + swiglu_beta=aiter_mxfp8_moe._AITER_SWIGLU_BETA, + ) + + backend, experts_cls = select_mxfp8_moe_backend(config) + + assert backend == Fp8MoeBackend.AITER_MXFP8 + assert experts_cls is aiter_mxfp8_moe.AiterMxfp8Experts + + @pytest.mark.skipif( not is_quant_method_supported("mxfp8"), reason="mxfp8 is not supported on this GPU type (requires sm_100+).", diff --git a/tests/models/registry.py b/tests/models/registry.py index 7f6767d451..5287076a4e 100644 --- a/tests/models/registry.py +++ b/tests/models/registry.py @@ -432,6 +432,8 @@ def check_available_online( "Qwen2MoeForCausalLM": _HfExamplesInfo("Qwen/Qwen1.5-MoE-A2.7B-Chat"), "Qwen3ForCausalLM": _HfExamplesInfo("Qwen/Qwen3-8B"), "Qwen3MoeForCausalLM": _HfExamplesInfo("Qwen/Qwen3-30B-A3B"), + "Qwen3_5ForCausalLM": _HfExamplesInfo("codecho/Qwen3.5-0.8B-text-only"), + "Qwen3_5MoeForCausalLM": _HfExamplesInfo("codecho/Qwen3.5-35B-A3B-text-only"), "MellumForCausalLM": _HfExamplesInfo("JetBrains/Mellum2-12B-A2.5B-Base"), "Qwen3NextForCausalLM": _HfExamplesInfo( "Qwen/Qwen3-Next-80B-A3B-Instruct", diff --git a/tests/multimodal/test_video.py b/tests/multimodal/test_video.py index badd113b1b..678c5c3f85 100644 --- a/tests/multimodal/test_video.py +++ b/tests/multimodal/test_video.py @@ -14,6 +14,7 @@ from transformers.video_utils import VideoMetadata from aphrodite.assets.base import get_aphrodite_public_assets +from aphrodite.models.minimax_m3.common.mm_preprocess import MiniMaxM3VideoBackend from aphrodite.multimodal.video import ( PYNVVIDEOCODEC_DECODER_CACHE_SIZE, PYNVVIDEOCODEC_VIDEO_BACKEND, @@ -379,6 +380,12 @@ def test_cosmos3_edge_uses_qwen3_vl_video_backend(): {"fps": 2}, id="qwen2_5_vl", ), + pytest.param( + "MiniMaxAI/MiniMax-M3", + MiniMaxM3VideoBackend, + None, + id="minimax_m3_vl", + ), ], ) def test_video_processor_from_model_repo( diff --git a/tests/quantization/test_compressed_tensors.py b/tests/quantization/test_compressed_tensors.py index 2e1775e0b7..0967ab1e81 100644 --- a/tests/quantization/test_compressed_tensors.py +++ b/tests/quantization/test_compressed_tensors.py @@ -444,6 +444,11 @@ def check_model(model): "Flat is better than nested.\nSparse is better than dense.", 150.0, ), + ( + "nm-testing/Llama-3.2-1B-Instruct-quipv16-nvfp4", + "Flat is better than nested.\nSparse is better than dense.", + 150.0, + ), ], ) def test_compressed_tensors_transforms_perplexity(aphrodite_runner, model, prompt, exp_perplexity): diff --git a/tests/reasoning/test_kimi_k3_reasoning_parser.py b/tests/reasoning/test_kimi_k3_reasoning_parser.py new file mode 100644 index 0000000000..d8a9784fcf --- /dev/null +++ b/tests/reasoning/test_kimi_k3_reasoning_parser.py @@ -0,0 +1,244 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the Aphrodite project + +import pytest + +from aphrodite.entrypoints.openai.chat_completion.protocol import ( + ChatCompletionRequest, +) +from aphrodite.entrypoints.openai.engine.protocol import DeltaMessage +from aphrodite.parser.kimi_k3 import KimiK3Parser +from aphrodite.parser.parser_manager import ParserManager +from aphrodite.reasoning.kimi_k3_reasoning_parser import KimiK3ReasoningParser + +pytestmark = pytest.mark.skip_global_cleanup + +OPEN = "<|open|>" +CLOSE = "<|close|>" +SEP = "<|sep|>" +THINK_OPEN = f"{OPEN}think{SEP}" +THINK_CLOSE = f"{CLOSE}think{SEP}" +RESPONSE_OPEN = f"{OPEN}response{SEP}" + + +class DummyTokenizer: + def get_vocab(self) -> dict[str, int]: + return {} + + def encode(self, text: str, add_special_tokens: bool = False) -> list[int]: + if text == THINK_OPEN: + return [1, 2, 3] + if text == THINK_CLOSE: + return [4, 2, 3] + return [ord(ch) for ch in text] + + +class ReasoningOnlyParser(KimiK3Parser): + reasoning_parser_cls = KimiK3ReasoningParser + + +def test_parser_manager_selects_kimi_k3_parser_for_reasoning_only(): + parser_cls = ParserManager.get_parser(reasoning_parser_name="kimi_k3") + + assert parser_cls is not None + assert issubclass(parser_cls, KimiK3Parser) + assert parser_cls.reasoning_parser_cls is KimiK3ReasoningParser + assert parser_cls.tool_parser_cls is None + + +def test_parser_selection_thinking_disabled(): + parser = KimiK3ReasoningParser(DummyTokenizer(), chat_template_kwargs={"thinking": False}) + + assert parser._thinking_enabled is False + + +def test_extract_reasoning_with_xtml_tags(): + parser = KimiK3ReasoningParser(DummyTokenizer()) + request = ChatCompletionRequest(model="test-model", messages=[]) + + reasoning, content = parser.extract_reasoning_content( + f"{THINK_OPEN}step{THINK_CLOSE}{RESPONSE_OPEN}answer", + request, + ) + + assert reasoning == "step" + assert content == "answer" + + +def test_extract_reasoning_with_generation_prefix_consumed(): + parser = KimiK3ReasoningParser(DummyTokenizer()) + request = ChatCompletionRequest(model="test-model", messages=[]) + + reasoning, content = parser.extract_reasoning_content( + f"step{THINK_CLOSE}{RESPONSE_OPEN}answer", + request, + ) + + assert reasoning == "step" + assert content == "answer" + + +def test_delegating_parser_strips_response_wrapper_without_tool_parser(): + parser = ReasoningOnlyParser(DummyTokenizer()) + request = ChatCompletionRequest(model="test-model", messages=[]) + + reasoning, content, tool_calls = parser.parse( + f"{THINK_OPEN}step{THINK_CLOSE}{RESPONSE_OPEN}answer", + request, + ) + + assert reasoning == "step" + assert content == "answer" + assert tool_calls == [] + + +def test_is_reasoning_end_uses_full_input_ids(): + parser = KimiK3ReasoningParser(DummyTokenizer()) + + assert not parser.is_reasoning_end([4, 2]) + assert parser.is_reasoning_end([4, 2, 3]) + + +def test_is_reasoning_end_ignores_stale_close_from_prior_turn(): + # DummyTokenizer: THINK_OPEN -> [1, 2, 3], THINK_CLOSE -> [4, 2, 3]. + # Multi-turn / agent continuation: a prior turn's think channel (its close + # marker) is kept in the prompt, then the current turn opens a new think + # block that has not closed yet. Reasoning must read as NOT ended, otherwise + # the structured-output gate constrains the current turn's reasoning. + parser = KimiK3ReasoningParser(DummyTokenizer()) + + stale_close = [4, 2, 3] + new_open = [1, 2, 3] + # prior close, then current-turn open still unclosed -> not ended + assert not parser.is_reasoning_end([*stale_close, *new_open]) + # ...then the current turn emits its own close -> ended + assert parser.is_reasoning_end([*stale_close, *new_open, *stale_close]) + # open with no close yet -> not ended + assert not parser.is_reasoning_end([*new_open]) + + +def test_streaming_split_open_marker_is_held_back(): + parser = KimiK3ReasoningParser(DummyTokenizer()) + + first = parser.extract_reasoning_content_streaming( + previous_text="", + current_text=OPEN, + delta_text=OPEN, + previous_token_ids=[], + current_token_ids=[1], + delta_token_ids=[1], + ) + second = parser.extract_reasoning_content_streaming( + previous_text=OPEN, + current_text=f"{OPEN}think", + delta_text="think", + previous_token_ids=[1], + current_token_ids=[1, 2], + delta_token_ids=[2], + ) + third = parser.extract_reasoning_content_streaming( + previous_text=f"{OPEN}think", + current_text=THINK_OPEN + "step", + delta_text=f"{SEP}step", + previous_token_ids=[1, 2], + current_token_ids=[1, 2, 3, 9], + delta_token_ids=[3, 9], + ) + + assert first is None + assert second is None + assert isinstance(third, DeltaMessage) + assert third.reasoning == "step" + + +def test_streaming_split_close_marker_hands_content_downstream(): + parser = KimiK3ReasoningParser(DummyTokenizer()) + + previous_text = f"{THINK_OPEN}step" + partial_close = parser.extract_reasoning_content_streaming( + previous_text=previous_text, + current_text=previous_text + CLOSE, + delta_text=CLOSE, + previous_token_ids=[1, 2, 3, 9], + current_token_ids=[1, 2, 3, 9, 4], + delta_token_ids=[4], + ) + closed = parser.extract_reasoning_content_streaming( + previous_text=previous_text + CLOSE, + current_text=previous_text + f"{THINK_CLOSE}{RESPONSE_OPEN}answer", + delta_text=f"think{SEP}{RESPONSE_OPEN}answer", + previous_token_ids=[1, 2, 3, 9, 4], + current_token_ids=[1, 2, 3, 9, 4, 2, 3, 10], + delta_token_ids=[2, 3, 10], + ) + + assert partial_close is None + assert isinstance(closed, DeltaMessage) + assert closed.reasoning is None + assert closed.content == f"{RESPONSE_OPEN}answer" + assert parser.extract_content_ids([2, 3, 10]) == [10] + + +def test_thinking_disabled_streams_content(): + parser = KimiK3ReasoningParser(DummyTokenizer(), chat_template_kwargs={"enable_thinking": False}) + + delta = parser.extract_reasoning_content_streaming( + previous_text="", + current_text=f"{RESPONSE_OPEN}answer", + delta_text=f"{RESPONSE_OPEN}answer", + previous_token_ids=[], + current_token_ids=[1], + delta_token_ids=[1], + ) + + assert isinstance(delta, DeltaMessage) + assert delta.content == f"{RESPONSE_OPEN}answer" + assert delta.reasoning is None + + +def test_delegating_parser_thinking_false_streams_response_content(): + parser = ReasoningOnlyParser(DummyTokenizer(), chat_template_kwargs={"thinking": False}) + request = ChatCompletionRequest( + model="test-model", + messages=[], + chat_template_kwargs={"thinking": False}, + ) + + first = parser.parse_delta( + delta_text="OK", + delta_token_ids=[10], + request=request, + prompt_token_ids=[1], + finished=False, + ) + partial_close = parser.parse_delta( + delta_text=CLOSE, + delta_token_ids=[2], + request=request, + prompt_token_ids=[1], + finished=False, + ) + closed = parser.parse_delta( + delta_text=f"response{SEP}", + delta_token_ids=[3, 4], + request=request, + prompt_token_ids=[1], + finished=False, + ) + + assert first is not None + assert first.content == "OK" + assert first.reasoning is None + assert partial_close is None + assert closed is None + + +def test_adjust_request_keeps_xtml_markers_contiguous(): + parser = KimiK3ReasoningParser(DummyTokenizer()) + request = ChatCompletionRequest(model="test-model", messages=[]) + + adjusted = parser.adjust_request(request) + + assert adjusted.skip_special_tokens is False + if hasattr(adjusted, "spaces_between_special_tokens"): + assert adjusted.spaces_between_special_tokens is False diff --git a/tests/renderers/test_kimi_k3.py b/tests/renderers/test_kimi_k3.py new file mode 100644 index 0000000000..82dd2a47ea --- /dev/null +++ b/tests/renderers/test_kimi_k3.py @@ -0,0 +1,337 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the Aphrodite project +from dataclasses import dataclass, field +from typing import Any + +import pytest + +from aphrodite.exceptions import AphroditeValidationError +from aphrodite.renderers import ChatParams +from aphrodite.renderers.kimi_k3 import KimiK3Renderer, _merge_k3_media_io_kwargs +from aphrodite.renderers.registry import RENDERER_REGISTRY +from aphrodite.tokenizers.registry import TokenizerRegistry + + +class StubTokenizer: + """Stands in for the model's TikTokenTokenizer. + + Records the kwargs it is called with and returns fixed token ids, so tests + can assert how the renderer drives ``apply_chat_template`` without the real + (git-LFS) tiktoken vocabulary. + """ + + def __init__(self, token_ids: list[int]) -> None: + self.token_ids = token_ids + self.calls: list[dict[str, Any]] = [] + self.conversations: list[list[dict[str, Any]]] = [] + + def apply_chat_template(self, conversation, **kwargs) -> list[int]: + self.conversations.append(conversation) + self.calls.append(kwargs) + return list(self.token_ids) + + +@dataclass +class MockHFConfig: + model_type: str = "kimi_k3" + + +@dataclass +class MockModelConfig: + runner_type: str = "generate" + is_multimodal_model: bool = False + multimodal_config: Any = None + hf_config: MockHFConfig = field(default_factory=MockHFConfig) + allowed_local_media_path: str = "" + allowed_media_domains: Any = None + enable_prompt_embeds: bool = False + renderer_num_workers: int = 1 + + +@dataclass +class MockParallelConfig: + _api_process_rank: int = 0 + + +@dataclass +class MockVllmConfig: + model_config: MockModelConfig + parallel_config: MockParallelConfig + + +def _make_renderer(tokenizer: StubTokenizer) -> KimiK3Renderer: + config = MockVllmConfig(MockModelConfig(), MockParallelConfig()) + return KimiK3Renderer(config, tokenizer) + + +def test_kimi_k3_registered(): + assert RENDERER_REGISTRY.load_renderer_cls("kimi_k3").__name__ == "KimiK3Renderer" + assert TokenizerRegistry.load_tokenizer_cls("kimi_k3").__name__ == "CachedHfTokenizer" + + +def test_k3_media_io_defaults_preserve_original_mode(): + # Default: K3 keeps the original image mode (no background flattening). + assert _merge_k3_media_io_kwargs(None) == {"image": {"image_mode": None}} + + # Server-/request-level values take precedence over the K3 default. + assert _merge_k3_media_io_kwargs({"image": {"image_mode": "RGB"}}) == {"image": {"image_mode": "RGB"}} + + # Unrelated image kwargs are merged with the default. + assert _merge_k3_media_io_kwargs({"image": {"rgba_background_color": (0, 0, 0)}}) == { + "image": {"image_mode": None, "rgba_background_color": (0, 0, 0)} + } + + +def test_apply_chat_template_forces_tokenize_and_pins_return_dict(): + tokenizer = StubTokenizer([7, 8, 9]) + renderer = _make_renderer(tokenizer) + tools = [{"type": "function", "function": {"name": "search"}}] + params = ChatParams(chat_template_kwargs={"tools": tools, "tokenize": False, "thinking": True}) + + token_ids = renderer._apply_chat_template([{"role": "user", "content": "hi"}], params) + + assert token_ids == [7, 8, 9] + kwargs = tokenizer.calls[-1] + # tokenize is forced on even though the request asked for False, so K3 keeps + # the special-vs-ordinary token distinction instead of re-tokenizing a string. + assert kwargs["tokenize"] is True + # return_dict is pinned False so we always get a flat list of ids. + assert kwargs["return_dict"] is False + assert kwargs["tools"] == tools + assert kwargs["thinking"] is True + + +def test_apply_chat_template_translates_standard_thinking_kwargs(): + # Standard enable_thinking/reasoning_effort kwargs must be translated + # to K3's native thinking/thinking_effort. + tokenizer = StubTokenizer([7, 8, 9]) + renderer = _make_renderer(tokenizer) + params = ChatParams(chat_template_kwargs={"enable_thinking": False, "reasoning_effort": "none"}) + + renderer._apply_chat_template([{"role": "user", "content": "hi"}], params) + + kwargs = tokenizer.calls[-1] + assert kwargs["thinking"] is False + assert "thinking_effort" not in kwargs + assert "enable_thinking" not in kwargs + assert "reasoning_effort" not in kwargs + + +@pytest.mark.parametrize("reasoning_effort", ["low", "high", "max"]) +def test_apply_chat_template_translates_supported_reasoning_effort( + reasoning_effort: str, +): + tokenizer = StubTokenizer([7, 8, 9]) + renderer = _make_renderer(tokenizer) + params = ChatParams(chat_template_kwargs={"reasoning_effort": reasoning_effort}) + + renderer._apply_chat_template([{"role": "user", "content": "hi"}], params) + + kwargs = tokenizer.calls[-1] + assert kwargs["thinking_effort"] == reasoning_effort + assert "reasoning_effort" not in kwargs + + +@pytest.mark.parametrize("reasoning_effort", ["minimal", "medium", "xhigh"]) +def test_apply_chat_template_rejects_unsupported_reasoning_effort( + reasoning_effort: str, +): + tokenizer = StubTokenizer([7, 8, 9]) + renderer = _make_renderer(tokenizer) + params = ChatParams(chat_template_kwargs={"reasoning_effort": reasoning_effort}) + + with pytest.raises(AphroditeValidationError, match="thinking_effort") as exc_info: + renderer._apply_chat_template([{"role": "user", "content": "hi"}], params) + + assert exc_info.value.parameter == "thinking_effort" + assert exc_info.value.value == reasoning_effort + assert tokenizer.calls == [] + + +def test_apply_chat_template_validates_canonical_native_thinking_effort(): + tokenizer = StubTokenizer([7, 8, 9]) + renderer = _make_renderer(tokenizer) + params = ChatParams( + chat_template_kwargs={ + "thinking_effort": "low", + "reasoning_effort": "medium", + } + ) + + renderer._apply_chat_template([{"role": "user", "content": "hi"}], params) + + assert tokenizer.calls[-1]["thinking_effort"] == "low" + + +@pytest.mark.parametrize("thinking_effort", ["none", "minimal", "medium", "xhigh"]) +def test_apply_chat_template_rejects_unsupported_native_thinking_effort( + thinking_effort: str, +): + tokenizer = StubTokenizer([7, 8, 9]) + renderer = _make_renderer(tokenizer) + params = ChatParams(chat_template_kwargs={"thinking_effort": thinking_effort}) + + with pytest.raises(AphroditeValidationError, match="thinking_effort") as exc_info: + renderer._apply_chat_template([{"role": "user", "content": "hi"}], params) + + assert exc_info.value.parameter == "thinking_effort" + assert exc_info.value.value == thinking_effort + assert tokenizer.calls == [] + + +def test_apply_chat_template_native_k3_kwargs_take_precedence(): + tokenizer = StubTokenizer([7, 8, 9]) + renderer = _make_renderer(tokenizer) + params = ChatParams( + chat_template_kwargs={ + "thinking": True, + "enable_thinking": False, + "thinking_effort": "low", + "reasoning_effort": "high", + } + ) + + renderer._apply_chat_template([{"role": "user", "content": "hi"}], params) + + kwargs = tokenizer.calls[-1] + assert kwargs["thinking"] is True + assert kwargs["thinking_effort"] == "low" + + +def test_apply_chat_template_adds_k3_api_metadata(): + tokenizer = StubTokenizer([7, 8, 9]) + renderer = _make_renderer(tokenizer) + response_format = {"type": "json_object"} + params = ChatParams( + tool_choice="required", + response_format=response_format, + ) + + renderer._apply_chat_template([{"role": "user", "content": "hi"}], params) + + kwargs = tokenizer.calls[-1] + assert kwargs["tool_choice"] == "required" + assert kwargs["response_format"] == response_format + + +def test_apply_chat_template_auto_tool_choice_keeps_template_kwarg(): + tokenizer = StubTokenizer([7, 8, 9]) + renderer = _make_renderer(tokenizer) + params = ChatParams( + chat_template_kwargs={"tool_choice": "required"}, + tool_choice="auto", + ) + + renderer._apply_chat_template([{"role": "user", "content": "hi"}], params) + + assert tokenizer.calls[-1]["tool_choice"] == "required" + + +def test_apply_chat_template_omits_tool_choice_without_tools(): + tokenizer = StubTokenizer([7, 8, 9]) + renderer = _make_renderer(tokenizer) + + renderer._apply_chat_template([{"role": "user", "content": "hi"}], ChatParams(tool_choice=None)) + + assert "tool_choice" not in tokenizer.calls[-1] + + +def test_render_messages_returns_token_prompt(): + renderer = _make_renderer(StubTokenizer([1, 2, 3])) + + conversation, prompt = renderer.render_messages([{"role": "user", "content": "hi"}], ChatParams()) + + assert prompt == {"prompt_token_ids": [1, 2, 3]} + assert "multi_modal_data" not in prompt + assert conversation[0]["role"] == "user" + + +def test_render_messages_derives_private_xtml_tool_attrs(): + tokenizer = StubTokenizer([1, 2, 3]) + renderer = _make_renderer(tokenizer) + + conversation, _ = renderer.render_messages( + [ + { + "role": "assistant", + "content": "", + "tool_calls": [ + { + "id": "lookup:0", + "type": "function", + "function": {"name": "lookup", "arguments": "{}"}, + }, + { + "id": "lookup:1", + "type": "function", + "function": {"name": "lookup", "arguments": "{}"}, + }, + ], + }, + { + "role": "tool", + "tool_call_id": "lookup:1", + "tool": "client-supplied-name", + "index": 99, + "content": "second", + }, + { + "role": "tool", + "tool_call_id": "lookup:0", + "content": "first", + }, + ], + ChatParams(), + ) + + assert [message["content"] for message in conversation[1:]] == [ + "first", + "second", + ] + assert conversation[1]["tool"] == "lookup" + assert conversation[1]["index"] == 1 + assert conversation[2]["tool"] == "lookup" + assert conversation[2]["index"] == 2 + assert tokenizer.conversations[-1] == conversation + + +def test_render_messages_ignores_client_supplied_xtml_tool_attrs(): + tokenizer = StubTokenizer([1, 2, 3]) + renderer = _make_renderer(tokenizer) + + conversation, _ = renderer.render_messages( + [ + { + "role": "assistant", + "content": "", + "tool_calls": [ + { + "id": "lookup:0", + "type": "function", + "function": {"name": "lookup", "arguments": "{}"}, + } + ], + }, + { + "role": "tool", + "tool_call_id": "unknown", + "tool": "lookup", + "index": 1, + "content": "result", + }, + ], + ChatParams(), + ) + + assert "tool" not in conversation[1] + assert "index" not in conversation[1] + + +@pytest.mark.asyncio +async def test_render_messages_async_returns_token_prompt(): + renderer = _make_renderer(StubTokenizer([4, 5])) + + conversation, prompt = await renderer.render_messages_async([{"role": "user", "content": "hi"}], ChatParams()) + + assert prompt == {"prompt_token_ids": [4, 5]} + assert conversation[0]["role"] == "user" diff --git a/tests/tool_parsers/test_kimi_k3_named_tool_choice.py b/tests/tool_parsers/test_kimi_k3_named_tool_choice.py new file mode 100644 index 0000000000..42f815dc99 --- /dev/null +++ b/tests/tool_parsers/test_kimi_k3_named_tool_choice.py @@ -0,0 +1,60 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the Aphrodite project +"""Named tool choice for Kimi K3: allowed when the XTML structural tag is +attached (strict tool calling), rejected otherwise.""" + +import pytest + +from aphrodite.entrypoints.openai.chat_completion.protocol import ( + ChatCompletionRequest, +) +from aphrodite.exceptions import AphroditeValidationError +from aphrodite.sampling_params import StructuredOutputsParams + + +class _DummyTokenizer: + def get_vocab(self): + return {} + + def encode(self, text, add_special_tokens=False): + return [ord(c) for c in text] + + +def _request(with_tag: bool) -> ChatCompletionRequest: + req = ChatCompletionRequest( + model="k3", + messages=[{"role": "user", "content": "hi"}], + tools=[ + { + "type": "function", + "function": { + "name": "get_weather", + "parameters": { + "type": "object", + "properties": {"city": {"type": "string"}}, + "required": ["city"], + }, + }, + } + ], + tool_choice={"type": "function", "function": {"name": "get_weather"}}, + ) + if with_tag: + req.structured_outputs = StructuredOutputsParams(structural_tag='{"type": "structural_tag", "format": {}}') + return req + + +def _parser(): + from aphrodite.tool_parsers.kimi_k3_tool_parser import KimiK3ToolParser + + return KimiK3ToolParser(_DummyTokenizer()) + + +def test_named_choice_allowed_with_structural_tag(): + req = _parser().adjust_request(_request(with_tag=True)) + assert req.skip_special_tokens is False + + +def test_named_choice_rejected_without_structural_tag(): + with pytest.raises(AphroditeValidationError): + _parser().adjust_request(_request(with_tag=False)) diff --git a/tests/tool_parsers/test_structural_tag_registry.py b/tests/tool_parsers/test_structural_tag_registry.py index 36fe15e6e9..f2a5ce09ab 100644 --- a/tests/tool_parsers/test_structural_tag_registry.py +++ b/tests/tool_parsers/test_structural_tag_registry.py @@ -5,7 +5,8 @@ from unittest.mock import MagicMock import pytest -from xgrammar import StructuralTag +from xgrammar import Grammar, StructuralTag +from xgrammar.testing import _is_grammar_accept_string from aphrodite.entrypoints.openai.chat_completion.protocol import ( ChatCompletionNamedFunction, @@ -26,6 +27,7 @@ from aphrodite.tool_parsers.glm47_moe_tool_parser import Glm47MoeModelToolParser from aphrodite.tool_parsers.hermes_tool_parser import Hermes2ProToolParser from aphrodite.tool_parsers.kimi_k2_tool_parser import KimiK2ToolParser +from aphrodite.tool_parsers.kimi_k3_tool_parser import KimiK3ToolParser from aphrodite.tool_parsers.llama_tool_parser import Llama3JsonToolParser from aphrodite.tool_parsers.minimax_m2_tool_parser import MinimaxM2ToolParser from aphrodite.tool_parsers.qwen3_engine_tool_parser import Qwen3EngineToolParser @@ -167,6 +169,183 @@ def test_hermes_required_tool_calls_use_empty_separator(): assert tag.format.separator == "" +# --------------------------------------------------------------------------- +# Kimi K3 (XTML channel format) structural tag +# --------------------------------------------------------------------------- +_K3_RESPONSE_OPEN = "<|open|>response<|sep|>" +_K3_RESPONSE_CLOSE = "<|close|>response<|sep|>" +_K3_TOOLS_OPEN = "<|open|>tools<|sep|>" +_K3_TOOLS_CLOSE = "<|close|>tools<|sep|>" +_K3_CALL_CLOSE = "<|close|>call<|sep|>" +_K3_ARG_CLOSE = "<|close|>argument<|sep|>" +_K3_MESSAGE_CLOSE = "<|close|>message<|sep|>" + + +def _k3_tools_by_name() -> list[ChatCompletionToolsParam]: + return [ + ChatCompletionToolsParam( + type="function", + function={ + "name": "get_weather", + "parameters": { + "type": "object", + "properties": { + "city": {"type": "string"}, + "days": {"type": "integer"}, + }, + "required": ["city"], + }, + }, + ), + ChatCompletionToolsParam( + type="function", + function={ + "name": "run_command", + "parameters": { + "type": "object", + "properties": {"command": {"type": "string"}}, + "required": ["command"], + }, + }, + ), + ] + + +def _k3_arg(key: str, typ: str, val: str) -> str: + return f'<|open|>argument key="{key}" type="{typ}"<|sep|>{val}{_K3_ARG_CLOSE}' + + +def _k3_call(name: str, args: str, idx: int = 1) -> str: + return f'<|open|>call tool="{name}" index="{idx}"<|sep|>{args}{_K3_CALL_CLOSE}' + + +def _k3_response(content: str = "") -> str: + return f"{_K3_RESPONSE_OPEN}{content}{_K3_RESPONSE_CLOSE}" + + +def _k3_tools(*calls: str) -> str: + return f"{_K3_TOOLS_OPEN}{''.join(calls)}{_K3_TOOLS_CLOSE}" + + +def _k3_grammar(tool_choice, tools=None): + tag = get_model_structural_tag( + model="kimi_k3", + tools=tools if tools is not None else _k3_tools_by_name(), + tool_choice=tool_choice, + reasoning=False, + ) + assert isinstance(tag, StructuralTag) + return Grammar.from_structural_tag(tag) + + +def test_kimi_k3_registered_as_aphrodite_builtin(): + assert "kimi_k3" in APHRODITE_BUILTIN_STRUCTURAL_TAG_MODELS + assert KimiK3ToolParser.structural_tag_model == "kimi_k3" + + +def test_kimi_k3_auto_without_strict_is_unconstrained(): + # auto + no strict tool => no structural tag (matches the strict gate). + tag = get_model_structural_tag( + model="kimi_k3", + tools=_k3_tools_by_name(), + tool_choice="auto", + reasoning=False, + ) + assert tag is None + + +@pytest.mark.parametrize( + "body", + [ + # single required arg + _k3_response() + _k3_tools(_k3_call("get_weather", _k3_arg("city", "string", "Paris"))), + # response content + two args (string + number) + _k3_response("Checking.") + + _k3_tools( + _k3_call( + "get_weather", + _k3_arg("city", "string", "Paris") + _k3_arg("days", "number", "3"), + ) + ), + # args in reverse order (parser is order-agnostic) + _k3_response() + + _k3_tools( + _k3_call( + "get_weather", + _k3_arg("days", "number", "3") + _k3_arg("city", "string", "Paris"), + ) + ), + # two calls, second tool + _k3_response() + + _k3_tools( + _k3_call("get_weather", _k3_arg("city", "string", "Paris"), 1), + _k3_call("run_command", _k3_arg("command", "string", "ls -la"), 2), + ), + # string value with regex metacharacters / spaces + _k3_response() + _k3_tools(_k3_call("run_command", _k3_arg("command", "string", "grep -E 'a|b{2,}' x.py"))), + # trailing message-close marker (model's natural turn terminator) + _k3_response() + _k3_tools(_k3_call("get_weather", _k3_arg("city", "string", "Paris"))) + _K3_MESSAGE_CLOSE, + # non-thinking mode: response-open is the prompt prefix, so it is absent + _K3_RESPONSE_CLOSE + _k3_tools(_k3_call("get_weather", _k3_arg("city", "string", "Paris"))), + ], +) +def test_kimi_k3_required_accepts_valid_tool_calls(body: str): + assert _is_grammar_accept_string(_k3_grammar("required"), body) + + +@pytest.mark.parametrize( + "body", + [ + # unknown tool name + _k3_response() + _k3_tools(_k3_call("get_temperature", _k3_arg("city", "string", "x"))), + # number arg given a non-numeric JSON value + _k3_response() + _k3_tools(_k3_call("get_weather", _k3_arg("days", "number", "abc"))), + # undeclared argument key + _k3_response() + + _k3_tools( + _k3_call( + "get_weather", + _k3_arg("city", "string", "Paris") + _k3_arg("zzz", "string", "x"), + ) + ), + # required schema but no argument tags + _k3_response() + _k3_tools(_k3_call("get_weather", "")), + # missing tools close marker + _k3_response() + _K3_TOOLS_OPEN + _k3_call("get_weather", _k3_arg("city", "string", "Paris")), + # required but no tool call + _k3_response("hello"), + ], +) +def test_kimi_k3_required_rejects_invalid(body: str): + assert not _is_grammar_accept_string(_k3_grammar("required"), body) + + +def test_kimi_k3_schema_without_required_accepts_empty_call(): + tools = [ + ChatCompletionToolsParam( + type="function", + function={ + "name": "get_weather", + "parameters": { + "type": "object", + "properties": {"city": {"type": "string"}}, + }, + }, + ) + ] + grammar = _k3_grammar("required", tools=tools) + body = _k3_response() + _k3_tools(_k3_call("get_weather", "")) + + assert _is_grammar_accept_string(grammar, body) + + +def test_kimi_k3_auto_strict_allows_response_only(sample_tools_strict): + # With a strict tool the tag is built; the tools channel is optional so a + # plain response (no tool call) is still valid. + grammar = _k3_grammar("auto", tools=sample_tools_strict) + assert _is_grammar_accept_string(grammar, _k3_response("Just answering.")) + + @pytest.mark.parametrize("model", sorted(XGRAMMAR_BUILTIN_STRUCTURAL_TAG_MODELS)) def test_get_model_structural_tag_supports_named_tool_choice( model: str, @@ -336,3 +515,143 @@ def test_get_function_parameters_relaxes_function_strict_false(): ) assert _get_function_parameters(function) is True + + +def _k3_tools_with_root_defs() -> list[ChatCompletionToolsParam]: + return [ + ChatCompletionToolsParam( + type="function", + function={ + "name": "make_config", + "parameters": { + "type": "object", + "properties": { + "config": { + "type": "object", + "properties": { + "build": {"$ref": "#/$defs/build"}, + "index": {"type": "string"}, + }, + "required": ["index"], + "additionalProperties": False, + }, + }, + "required": ["config"], + "$defs": { + "build": { + "type": "object", + "properties": {"outDir": {"type": "string"}}, + "additionalProperties": False, + } + }, + }, + }, + ) + ] + + +def test_kimi_k3_property_ref_to_root_defs_compiles_and_accepts(): + # Root-level $defs referenced from inside a property schema (the walle + # TestReferences shape). Slicing the property out of the parameters + # document orphans "#/$defs/..." unless the builder re-attaches $defs; + # before the fix Grammar.from_structural_tag raised on the dangling ref. + grammar = _k3_grammar("required", tools=_k3_tools_with_root_defs()) + + body = _k3_response() + _k3_tools( + _k3_call( + "make_config", + _k3_arg( + "config", + "object", + '{"build": {"outDir": "dist"}, "index": "a.html"}', + ), + ) + ) + assert _is_grammar_accept_string(grammar, body) + + +def _k3_tools_with_string_enum() -> list[ChatCompletionToolsParam]: + return [ + ChatCompletionToolsParam( + type="function", + function={ + "name": "set_unit", + "parameters": { + "type": "object", + "properties": { + "unit": { + "type": "string", + "enum": ["celsius", " fahrenheit", "\tkelvin"], + }, + }, + "required": ["unit"], + }, + }, + ) + ] + + +@pytest.mark.parametrize("value", ["celsius", " fahrenheit", "\tkelvin"]) +def test_kimi_k3_string_enum_accepts_exact_values(value: str): + # Raw string channel with enum: constrained to the exact enum values, + # including leading-whitespace variants the model otherwise flubs. + grammar = _k3_grammar("required", tools=_k3_tools_with_string_enum()) + body = _k3_response() + _k3_tools(_k3_call("set_unit", _k3_arg("unit", "string", value))) + assert _is_grammar_accept_string(grammar, body) + + +@pytest.mark.parametrize("value", ["kelvin", "Celsius", "celsius ", ""]) +def test_kimi_k3_string_enum_rejects_non_members(value: str): + grammar = _k3_grammar("required", tools=_k3_tools_with_string_enum()) + body = _k3_response() + _k3_tools(_k3_call("set_unit", _k3_arg("unit", "string", value))) + assert not _is_grammar_accept_string(grammar, body) + + +def _k3_tools_with_maxlen() -> list[ChatCompletionToolsParam]: + return [ + ChatCompletionToolsParam( + type="function", + function={ + "name": "set_note", + "parameters": { + "type": "object", + "properties": { + "note": {"type": "string", "maxLength": 8, "minLength": 2}, + }, + "required": ["note"], + }, + }, + ) + ] + + +def test_kimi_k3_string_maxlength_bounds_raw_channel(): + # Raw string channel with maxLength/minLength: enforced via a bounded + # regex that keeps the "<|" marker prefix unambiguous but still allows + # a bare '<' inside values. + grammar = _k3_grammar("required", tools=_k3_tools_with_maxlen()) + + def body(val: str) -> str: + return _k3_response() + _k3_tools(_k3_call("set_note", _k3_arg("note", "string", val))) + + assert _is_grammar_accept_string(grammar, body("ab")) + assert _is_grammar_accept_string(grammar, body("a dict[str, int]: + return {} + + def encode(self, text: str, add_special_tokens: bool = False) -> list[int]: + if text == THINK_OPEN: + return [1, 2, 3] + if text == THINK_CLOSE: + return [4, 2, 3] + return [ord(ch) for ch in text] + + +class KimiK3DelegatingParser(KimiK3Parser): + reasoning_parser_cls = KimiK3ReasoningParser + tool_parser_cls = KimiK3ToolParser + + +def test_parser_manager_selects_kimi_k3_parser(): + parser_cls = ParserManager.get_parser( + tool_parser_name="kimi_k3", + reasoning_parser_name="kimi_k3", + enable_auto_tools=True, + ) + + assert parser_cls is not None + assert issubclass(parser_cls, KimiK3Parser) + assert parser_cls.reasoning_parser_cls is KimiK3ReasoningParser + assert parser_cls.tool_parser_cls is KimiK3ToolParser + + +def _request() -> ChatCompletionRequest: + return ChatCompletionRequest( + model="test-model", + messages=[], + tools=[ + { + "type": "function", + "function": { + "name": "calc", + "parameters": {"type": "object", "properties": {}}, + }, + } + ], + tool_choice="auto", + ) + + +def _named_request() -> ChatCompletionRequest: + return ChatCompletionRequest( + model="test-model", + messages=[], + tools=[ + { + "type": "function", + "function": { + "name": "calc", + "parameters": {"type": "object", "properties": {}}, + }, + } + ], + tool_choice={"type": "function", "function": {"name": "calc"}}, + ) + + +def _responses_request(*, tool_choice="auto") -> ResponsesRequest: + return ResponsesRequest.model_validate( + { + "model": "test-model", + "input": "Call the calc tool.", + "tools": [ + { + "type": "function", + "name": "calc", + "parameters": {"type": "object", "properties": {}}, + } + ], + "tool_choice": tool_choice, + } + ) + + +def _arg(key: str, typ: str, value: str) -> str: + return f'{OPEN}argument key="{key}" type="{typ}"{SEP}{value}{CLOSE}argument{SEP}' + + +def _call(tool: str, index: int, *args: str) -> str: + body = "".join(args) + return f'{OPEN}call tool="{tool}" index="{index}"{SEP}{body}{CLOSE}call{SEP}' + + +def _response(content: str) -> str: + return f"{OPEN}response{SEP}{content}{RESPONSE_CLOSE}" + + +def _tools(*calls: str) -> str: + return f"{OPEN}tools{SEP}{''.join(calls)}{CLOSE}tools{SEP}" + + +def test_extract_tool_calls_with_response_and_typed_arguments(): + parser = KimiK3ToolParser(DummyTokenizer()) + + output = _response("answer") + _tools( + _call( + "calc", + 1, + _arg("x", "number", "1"), + _arg("flag", "boolean", "true"), + _arg("text", "string", "raw"), + ) + ) + extracted = parser.extract_tool_calls(output, _request()) + + assert extracted.tools_called is True + assert extracted.content == "answer" + assert len(extracted.tool_calls) == 1 + tool_call = extracted.tool_calls[0] + assert tool_call.id == "calc:0" + assert tool_call.function.name == "calc" + assert json.loads(tool_call.function.arguments) == { + "x": 1, + "flag": True, + "text": "raw", + } + + +def test_delegating_parser_preserves_tool_calls_after_reasoning(): + parser = KimiK3DelegatingParser(DummyTokenizer()) + output = f"{THINK_OPEN}step{THINK_CLOSE}" + _response("answer") + _tools(_call("calc", 1, _arg("x", "number", "1"))) + + reasoning, content, tool_calls = parser.parse( + output, + _request(), + enable_auto_tools=True, + ) + + assert reasoning == "step" + assert content == "answer" + assert tool_calls is not None + assert len(tool_calls) == 1 + assert tool_calls[0].id == "calc:0" + assert tool_calls[0].name == "calc" + assert json.loads(tool_calls[0].arguments) == {"x": 1} + + +def test_delegating_parser_required_tool_choice_uses_xtml_parser(): + parser = KimiK3DelegatingParser(DummyTokenizer()) + request = _request().model_copy(update={"tool_choice": "required"}) + output = f"{THINK_OPEN}step{THINK_CLOSE}" + _response("") + _tools(_call("calc", 1, _arg("x", "number", "1"))) + + reasoning, content, tool_calls = parser.parse( + output, + request, + enable_auto_tools=True, + ) + + assert reasoning == "step" + assert content is None + assert tool_calls is not None + assert len(tool_calls) == 1 + assert tool_calls[0].name == "calc" + assert json.loads(tool_calls[0].arguments) == {"x": 1} + + +def test_delegating_parser_named_tool_choice_uses_xtml_parser(): + parser = KimiK3DelegatingParser(DummyTokenizer()) + output = f"{THINK_OPEN}step{THINK_CLOSE}" + _response("") + _tools(_call("calc", 1, _arg("x", "number", "1"))) + + reasoning, content, tool_calls = parser.parse( + output, + _named_request(), + enable_auto_tools=True, + ) + + assert reasoning == "step" + assert content is None + assert tool_calls is not None + assert len(tool_calls) == 1 + assert tool_calls[0].name == "calc" + assert json.loads(tool_calls[0].arguments) == {"x": 1} + + +def test_delegating_parser_auto_no_call_strips_consumed_response_prefix(): + parser = KimiK3DelegatingParser(DummyTokenizer(), chat_template_kwargs={"thinking": False}) + request = _request().model_copy(update={"chat_template_kwargs": {"thinking": False}}) + + reasoning, content, tool_calls = parser.parse( + f"answer{RESPONSE_CLOSE}", + request, + enable_auto_tools=True, + ) + + assert reasoning is None + assert content == "answer" + assert tool_calls is None + + +def test_delegating_parser_required_call_strips_consumed_response_prefix(): + parser = KimiK3DelegatingParser(DummyTokenizer(), chat_template_kwargs={"thinking": False}) + request = _request().model_copy( + update={ + "tool_choice": "required", + "chat_template_kwargs": {"thinking": False}, + } + ) + output = RESPONSE_CLOSE + _tools(_call("calc", 1, _arg("x", "number", "1"))) + + reasoning, content, tool_calls = parser.parse( + output, + request, + enable_auto_tools=True, + ) + + assert reasoning is None + assert content is None + assert tool_calls is not None + assert len(tool_calls) == 1 + assert tool_calls[0].name == "calc" + assert json.loads(tool_calls[0].arguments) == {"x": 1} + + +def test_delegating_parser_truncated_tools_do_not_leak_xtml(): + parser = KimiK3DelegatingParser(DummyTokenizer(), chat_template_kwargs={"thinking": False}) + request = _request().model_copy( + update={ + "tool_choice": "required", + "chat_template_kwargs": {"thinking": False}, + } + ) + + reasoning, content, tool_calls = parser.parse( + (f'{RESPONSE_CLOSE}{OPEN}tools{SEP}{OPEN}call tool="calc" index="1"'), + request, + enable_auto_tools=True, + ) + + assert reasoning is None + assert content is None + assert tool_calls is None + + +def test_extract_tool_calls_unescapes_attributes(): + parser = KimiK3ToolParser(DummyTokenizer()) + + output = _tools(_call("a&b"c", 1, _arg("k&q", "string", "v"))) + extracted = parser.extract_tool_calls(output, _request()) + + assert extracted.tools_called is True + assert extracted.tool_calls[0].function.name == 'a&b"c' + assert json.loads(extracted.tool_calls[0].function.arguments) == {"k&q": "v"} + + +def test_extract_tool_calls_allows_less_than_in_attributes(): + parser = KimiK3ToolParser(DummyTokenizer()) + + output = _tools(_call("calc None: + if current_platform.is_rocm(): + request.getfixturevalue("disable_aphrodite_compile_cache") + + @pytest.mark.parametrize( "speculative_config", [ @@ -167,21 +173,26 @@ def reset_torch_dynamo(): }, ], ) +@pytest.mark.usefixtures("disable_aphrodite_compile_cache_on_rocm") @single_gpu_only @large_gpu_mark(min_gb=20) def test_ngram_and_suffix_correctness( speculative_config: dict, model_name: str, + aphrodite_runner, ): - spec_llm = LLM( - model=model_name, + with aphrodite_runner( + model_name, + # Keep LLM defaults; AphroditeRunner only provides lifecycle cleanup here. + trust_remote_code=False, + enable_chunked_prefill=None, speculative_config=speculative_config, max_model_len=4096, - ) - evaluate_llm_for_gsm8k(spec_llm) - del spec_llm - torch.accelerator.empty_cache() - cleanup_dist_env_and_memory() + # Preserve LLM's default compilation/cudagraph configuration. Without + # this, AphroditeRunner injects its reduced test-only capture sizes. + compilation_config=CompilationConfig(), + ) as runner: + evaluate_llm_for_gsm8k(runner.llm) @pytest.mark.parametrize("async_scheduling", [True], ids=["async"]) diff --git a/tests/v1/ec_connector/integration/run_epd_correctness_test.sh b/tests/v1/ec_connector/integration/run_epd_correctness_test.sh index 6e3e6bd04b..9540bdfbf0 100644 --- a/tests/v1/ec_connector/integration/run_epd_correctness_test.sh +++ b/tests/v1/ec_connector/integration/run_epd_correctness_test.sh @@ -21,12 +21,19 @@ GIT_ROOT="${GIT_ROOT:-$(cd -- "${SCRIPT_DIR}/../../../.." && pwd -P)}" # Model to test MODEL="${MODEL:-Qwen/Qwen2.5-VL-3B-Instruct}" +MAX_MODEL_LEN="${MAX_MODEL_LEN:-10240}" +GPU_MEMORY_UTILIZATION="${GPU_MEMORY_UTILIZATION:-0.7}" +MAX_NUM_SEQS="${MAX_NUM_SEQS:-128}" # Set 1 to use multimodal prompts; else to use text-only USE_MM_PROMPTS="${USE_MM_PROMPTS:-1}" -MM_FLAG="" -if [ "$USE_MM_PROMPTS" = "1" ]; then - MM_FLAG="--use_mm_prompts" +USE_TWO_IMAGE_PROMPT="${USE_TWO_IMAGE_PROMPT:-1}" +TEST_FLAGS=() +if [[ "$USE_MM_PROMPTS" == "1" ]]; then + TEST_FLAGS+=(--use_mm_prompts) +fi +if [[ "$USE_TWO_IMAGE_PROMPT" != "1" ]]; then + TEST_FLAGS+=(--skip_two_image_prompt) fi # GPU configuration @@ -36,6 +43,16 @@ GPU_D="${GPU_D:-2}" GPU_SINGLE="${GPU_SINGLE:-$GPU_P}" GPU_PD="${GPU_PD:-$GPU_P}" +# Device platform and affinity environment variable +DEVICE_PLATFORM="${DEVICE_PLATFORM:-cuda}" +if [[ -z "${DEVICE_AFFINITY_ENV:-}" ]]; then + if [[ "${DEVICE_PLATFORM,,}" == "xpu" ]]; then + DEVICE_AFFINITY_ENV="ZE_AFFINITY_MASK" + else + DEVICE_AFFINITY_ENV="CUDA_VISIBLE_DEVICES" + fi +fi + # Port ENCODE_PORT="${ENCODE_PORT:-19534}" PREFILL_PORT="${PREFILL_PORT:-19535}" @@ -87,11 +104,12 @@ run_baseline() { # Start baseline instance echo "Starting baseline instance on GPU $GPU_SINGLE, port $PORT" - CUDA_VISIBLE_DEVICES="$GPU_SINGLE" aphrodite serve "$MODEL" \ + env "$DEVICE_AFFINITY_ENV=$GPU_SINGLE" aphrodite serve "$MODEL" \ --port "$PORT" \ + --max-model-len "$MAX_MODEL_LEN" \ --enforce-eager \ - --gpu-memory-utilization 0.7 \ - --max-num-seqs 128 \ + --gpu-memory-utilization 0.9 \ + --max-num-seqs "$MAX_NUM_SEQS" \ --allowed-local-media-path "${GIT_ROOT}"/tests/v1/ec_connector/integration \ > "$LOG_PATH"/baseline.log 2>&1 & @@ -112,7 +130,7 @@ run_baseline() { --model_name "$MODEL" \ --mode baseline \ --baseline_file "$BASELINE_FILE" \ - $MM_FLAG + "${TEST_FLAGS[@]}" # Cleanup baseline echo "Stopping baseline instance..." @@ -139,14 +157,15 @@ run_epd_1e_1pd() { # Start encoder instance echo "Starting encoder instance on GPU $GPU_E, port $ENCODE_PORT" - CUDA_VISIBLE_DEVICES="$GPU_E" aphrodite serve "$MODEL" \ + env "$DEVICE_AFFINITY_ENV=$GPU_E" aphrodite serve "$MODEL" \ --port "$ENCODE_PORT" \ + --max-model-len "$MAX_MODEL_LEN" \ --enforce-eager \ --gpu-memory-utilization 0.01 \ --enable-request-id-headers \ --no-enable-prefix-caching \ --max-num-batched-tokens 114688 \ - --max-num-seqs 128 \ + --max-num-seqs "$MAX_NUM_SEQS" \ --allowed-local-media-path "${GIT_ROOT}"/tests/v1/ec_connector/integration \ --ec-transfer-config '{ "ec_connector": "ECExampleConnector", @@ -160,12 +179,13 @@ run_epd_1e_1pd() { # Start prefill+decode instance echo "Starting PD instance on GPU $GPU_PD, port $PREFILL_DECODE_PORT" - CUDA_VISIBLE_DEVICES="$GPU_PD" aphrodite serve "$MODEL" \ + env "$DEVICE_AFFINITY_ENV=$GPU_PD" aphrodite serve "$MODEL" \ --port "$PREFILL_DECODE_PORT" \ + --max-model-len "$MAX_MODEL_LEN" \ --enforce-eager \ - --gpu-memory-utilization 0.7 \ + --gpu-memory-utilization "$GPU_MEMORY_UTILIZATION" \ --enable-request-id-headers \ - --max-num-seqs 128 \ + --max-num-seqs "$MAX_NUM_SEQS" \ --allowed-local-media-path "${GIT_ROOT}"/tests/v1/ec_connector/integration \ --ec-transfer-config '{ "ec_connector": "ECExampleConnector", @@ -212,7 +232,7 @@ run_epd_1e_1pd() { --model_name "$MODEL" \ --mode disagg \ --baseline_file "$BASELINE_FILE" \ - $MM_FLAG + "${TEST_FLAGS[@]}" # Cleanup echo "✓✓ 1E+1PD Correctness Test finished" @@ -242,14 +262,15 @@ run_baseline_1p_1d() { # Start prefill instance echo "Starting prefill instance on GPU $GPU_P, port $PREFILL_PORT" - CUDA_VISIBLE_DEVICES="$GPU_P" \ + env "$DEVICE_AFFINITY_ENV=$GPU_P" \ APHRODITE_NIXL_SIDE_CHANNEL_PORT=5559 \ aphrodite serve "$MODEL" \ --port "$PREFILL_PORT" \ + --max-model-len "$MAX_MODEL_LEN" \ --enforce-eager \ - --gpu-memory-utilization 0.7 \ + --gpu-memory-utilization "$GPU_MEMORY_UTILIZATION" \ --enable-request-id-headers \ - --max-num-seqs 128 \ + --max-num-seqs "$MAX_NUM_SEQS" \ --allowed-local-media-path "${GIT_ROOT}"/tests/v1/ec_connector/integration \ --kv-transfer-config '{ "kv_connector": "NixlConnector", @@ -260,14 +281,15 @@ run_baseline_1p_1d() { # Start decode instance echo "Starting decode instance on GPU $GPU_D, port $DECODE_PORT" - CUDA_VISIBLE_DEVICES="$GPU_D" \ + env "$DEVICE_AFFINITY_ENV=$GPU_D" \ APHRODITE_NIXL_SIDE_CHANNEL_PORT=6000 \ aphrodite serve "$MODEL" \ --port "$DECODE_PORT" \ + --max-model-len "$MAX_MODEL_LEN" \ --enforce-eager \ - --gpu-memory-utilization 0.7 \ + --gpu-memory-utilization "$GPU_MEMORY_UTILIZATION" \ --enable-request-id-headers \ - --max-num-seqs 128 \ + --max-num-seqs "$MAX_NUM_SEQS" \ --allowed-local-media-path "${GIT_ROOT}"/tests/v1/ec_connector/integration \ --kv-transfer-config '{ "kv_connector": "NixlConnector", @@ -309,7 +331,7 @@ run_baseline_1p_1d() { --model_name "$MODEL" \ --mode baseline_pd \ --baseline_file "$BASELINE_PD_FILE" \ - $MM_FLAG + "${TEST_FLAGS[@]}" # Cleanup echo "Stopping PD (1P+1D) instances..." @@ -339,14 +361,15 @@ run_epd_1e_1p_1d() { # Start encoder instance echo "Starting encoder instance on GPU $GPU_E, port $ENCODE_PORT" - CUDA_VISIBLE_DEVICES="$GPU_E" aphrodite serve "$MODEL" \ + env "$DEVICE_AFFINITY_ENV=$GPU_E" aphrodite serve "$MODEL" \ --port "$ENCODE_PORT" \ + --max-model-len "$MAX_MODEL_LEN" \ --enforce-eager \ --gpu-memory-utilization 0.01 \ --enable-request-id-headers \ --no-enable-prefix-caching \ --max-num-batched-tokens 114688 \ - --max-num-seqs 128 \ + --max-num-seqs "$MAX_NUM_SEQS" \ --allowed-local-media-path "${GIT_ROOT}"/tests/v1/ec_connector/integration \ --ec-transfer-config '{ "ec_connector": "ECExampleConnector", @@ -360,14 +383,15 @@ run_epd_1e_1p_1d() { # Start prefill instance echo "Starting prefill instance on GPU $GPU_P, port $PREFILL_PORT" - CUDA_VISIBLE_DEVICES="$GPU_P" \ + env "$DEVICE_AFFINITY_ENV=$GPU_P" \ APHRODITE_NIXL_SIDE_CHANNEL_PORT=5559 \ aphrodite serve "$MODEL" \ --port "$PREFILL_PORT" \ + --max-model-len "$MAX_MODEL_LEN" \ --enforce-eager \ - --gpu-memory-utilization 0.7 \ + --gpu-memory-utilization "$GPU_MEMORY_UTILIZATION" \ --enable-request-id-headers \ - --max-num-seqs 128 \ + --max-num-seqs "$MAX_NUM_SEQS" \ --allowed-local-media-path "${GIT_ROOT}"/tests/v1/ec_connector/integration \ --ec-transfer-config '{ "ec_connector": "ECExampleConnector", @@ -385,14 +409,15 @@ run_epd_1e_1p_1d() { # Start decode instance echo "Starting decode instance on GPU $GPU_D, port $DECODE_PORT" - CUDA_VISIBLE_DEVICES="$GPU_D" \ + env "$DEVICE_AFFINITY_ENV=$GPU_D" \ APHRODITE_NIXL_SIDE_CHANNEL_PORT=6000 \ aphrodite serve "$MODEL" \ --port "$DECODE_PORT" \ + --max-model-len "$MAX_MODEL_LEN" \ --enforce-eager \ - --gpu-memory-utilization 0.7 \ + --gpu-memory-utilization "$GPU_MEMORY_UTILIZATION" \ --enable-request-id-headers \ - --max-num-seqs 128 \ + --max-num-seqs "$MAX_NUM_SEQS" \ --allowed-local-media-path "${GIT_ROOT}"/tests/v1/ec_connector/integration \ --kv-transfer-config '{ "kv_connector": "NixlConnector", @@ -438,7 +463,7 @@ run_epd_1e_1p_1d() { --model_name "$MODEL" \ --mode disagg \ --baseline_file "$BASELINE_PD_FILE" \ - $MM_FLAG + "${TEST_FLAGS[@]}" # Cleanup echo "✓✓ 1E+1P+1D Correctness Test finished" @@ -465,7 +490,7 @@ run_epd_1e_1pd # Step 3: Test baseline 1P + 1D run_baseline_1p_1d -# Step 4: Test 1E + 1P + 1D +# # Step 4: Test 1E + 1P + 1D run_epd_1e_1p_1d # Cleanup output file diff --git a/tests/v1/ec_connector/integration/test_epd_correctness.py b/tests/v1/ec_connector/integration/test_epd_correctness.py index 3b8a74182e..11527316b0 100644 --- a/tests/v1/ec_connector/integration/test_epd_correctness.py +++ b/tests/v1/ec_connector/integration/test_epd_correctness.py @@ -185,6 +185,12 @@ def main(): help="Use multimodal prompts (default: use text-only for quick testing)", ) + parser.add_argument( + "--skip_two_image_prompt", + action="store_true", + help="Skip the two-image multimodal prompt", + ) + args = parser.parse_args() print(f"Service URL: {args.service_url}") @@ -214,7 +220,7 @@ def main(): # Select prompts to use if args.use_mm_prompts: - test_prompts = SAMPLE_PROMPTS_MM + test_prompts = SAMPLE_PROMPTS_MM[:1] if args.skip_two_image_prompt else SAMPLE_PROMPTS_MM print("Using multimodal prompts") else: test_prompts = SAMPLE_PROMPTS_TEXT diff --git a/tests/v1/kv_connector/nixl_integration/run_accuracy_test.sh b/tests/v1/kv_connector/nixl_integration/run_accuracy_test.sh index 638e25beae..7fc26ddac8 100755 --- a/tests/v1/kv_connector/nixl_integration/run_accuracy_test.sh +++ b/tests/v1/kv_connector/nixl_integration/run_accuracy_test.sh @@ -227,9 +227,12 @@ run_tests_for_model() { # Calculate side channel port SIDE_CHANNEL_PORT=$((5659 + i * $DECODER_TP_SIZE)) INTERNAL_PORT=$((DECODER_INTERNAL_PORT_BASE + i * INTERNAL_PORT_STRIDE)) - DECODER_INTERNAL_PORT_ENV= + # For non-DP mode, set APHRODITE_PORT to pin the internal port; + # For DP mode, set APHRODITE_DP_MASTER_PORT instead to avoid race condition. if [[ -z "${DP_EP:-}" ]]; then DECODER_INTERNAL_PORT_ENV="APHRODITE_PORT=$INTERNAL_PORT" + else + DECODER_INTERNAL_PORT_ENV="APHRODITE_DP_MASTER_PORT=$INTERNAL_PORT" fi echo "Starting decode instance $i on GPU $GPU_ID, port $PORT" diff --git a/tests/v1/kv_connector/unit/test_nixl_connector_hma.py b/tests/v1/kv_connector/unit/test_nixl_connector_hma.py index 56a280461b..32b3dbb56f 100644 --- a/tests/v1/kv_connector/unit/test_nixl_connector_hma.py +++ b/tests/v1/kv_connector/unit/test_nixl_connector_hma.py @@ -1347,3 +1347,102 @@ def test_logical_to_kernel_block_ids_with_remote_ratio( remote_physical_per_logical, ) assert list(result) == expected_kernel_block_ids, f"Expected {expected_kernel_block_ids}, got {result}" + + +def _make_hybrid_mla_kv_cache_config(num_blocks: int = 4): + """Build a KimiLinear-shaped config with one MLA and two KDA groups.""" + from aphrodite.v1.attention.backends.registry import MambaAttentionBackendEnum + from aphrodite.v1.kv_cache_interface import ( + KVCacheConfig, + KVCacheGroupSpec, + KVCacheTensor, + MambaSpec, + MLAAttentionSpec, + ) + + mla_spec = MLAAttentionSpec(block_size=12, num_kv_heads=1, head_size=6, dtype=torch.float16) + unified_page = mla_spec.page_size_bytes + kda_spec = MambaSpec( + block_size=12, + shapes=((8, 3), (1, 4, 4)), + dtypes=(torch.float16, torch.float32), + page_size_padded=unified_page, + mamba_type=MambaAttentionBackendEnum.GDN_ATTN, + ) + assert kda_spec.page_size_bytes == unified_page + return KVCacheConfig( + num_blocks=num_blocks, + kv_cache_tensors=[ + KVCacheTensor( + size=num_blocks * unified_page, + shared_by=[f"mla.{i}", f"kda_a.{i}", f"kda_b.{i}"], + ) + for i in range(2) + ], + kv_cache_groups=[ + KVCacheGroupSpec(["mla.0", "mla.1"], mla_spec), + KVCacheGroupSpec(["kda_a.0", "kda_a.1"], kda_spec), + KVCacheGroupSpec(["kda_b.0", "kda_b.1"], kda_spec), + ], + ) + + +@pytest.mark.cpu_test +def test_push_write_hybrid_mla_replicates_attention(): + """Hybrid MLA+SSM push replicates attention while SSM stays sharded.""" + import threading + from collections import defaultdict + + from aphrodite.distributed.kv_transfer.kv_connector.v1.nixl.push_worker import ( + NixlPushConnectorWorker, + ) + from aphrodite.distributed.kv_transfer.kv_connector.v1.nixl.tp_mapping import ( + TPMapping, + ) + from aphrodite.v1.kv_cache_interface import MambaSpec, MLAAttentionSpec + + worker = object.__new__(NixlPushConnectorWorker) + worker.shutdown = lambda: None + worker.use_mla = True + worker._has_mamba = True + worker._group_spec_types = (MLAAttentionSpec, MambaSpec) + worker.transfer_topo = MagicMock() + worker.transfer_topo.tp_ratio.return_value = -2 + remote_info = MagicMock() + remote_info.remote_physical_blocks_per_logical = 1 + remote_info.remote_block_size = 4 + worker.transfer_topo.get_engine_info.return_value = remote_info + + engine_id = "remote-engine" + worker.tp_mappings = { + engine_id: TPMapping( + source_ranks_per_group=((0,), (0, 1)), + all_source_ranks=(0, 1), + rank_to_attention_slot={0: 0, 1: 0}, + rank_offset_factor=0, + ) + } + worker.dst_xfer_side_handles = {engine_id: {0: 100, 1: 101}} + worker.src_xfer_handles_by_tp_ratio = {(-2, 4): [200, 201]} + worker.src_xfer_handles_by_block_size = {4: 300} + worker._sending_transfers = defaultdict(list) + worker._sending_transfers_lock = threading.Lock() + worker.kv_cache_config = _make_hybrid_mla_kv_cache_config() + worker._xfer_blocks = MagicMock(return_value=1) + + meta = MagicMock() + meta.remote.engine_id = engine_id + meta.remote.block_ids = [[7, 8], [3]] + meta.local_physical_block_ids = [[1, 2], [5]] + + worker._xfer_blocks_for_req("req-1", meta) + + calls = worker._xfer_blocks.call_args_list + assert len(calls) == 2 + for call, rank, local_handle, remote_handle in zip(calls, (0, 1), (200, 201), (100, 101)): + spec = call.kwargs["read_spec"] + assert spec.remote_rank == rank + assert spec.local_block_ids == [[1, 2], [5]] + assert spec.remote_block_ids == [[7, 8], [3]] + assert call.kwargs["local_xfer_side_handle"] == local_handle + assert call.kwargs["remote_xfer_side_handle"] == remote_handle diff --git a/tests/v1/kv_connector/unit/test_nixl_desc_geometry.py b/tests/v1/kv_connector/unit/test_nixl_desc_geometry.py index 4110dd91c2..5a2ffb4ef2 100644 --- a/tests/v1/kv_connector/unit/test_nixl_desc_geometry.py +++ b/tests/v1/kv_connector/unit/test_nixl_desc_geometry.py @@ -597,3 +597,20 @@ def test_mla_hybrid_large_ppl_geometry(num_tokens): num_tokens=num_tokens, tp_size=8, ) + + +@pytest.mark.cpu_test +def test_mismatched_mla_kernel_page_rejected_for_mla_hybrid(): + """Reject MLA hybrid handshakes with mismatched per-token pages.""" + worker = _make_mla_hybrid_worker(local_block_size=12, kernel_block_size=4, num_logical_blocks=8) + meta_r = _make_remote_meta( + worker, + remote_block_size=8, + remote_kernel_block_size=4, + remote_num_logical=12, + remote_ssm_sizes=(24, 32), + ) + # Equal kernel block sizes (ratio 1), but a half-sized per-token page. + meta_r.block_lens = [x // 2 for x in worker.block_len_per_layer] + with pytest.raises((AssertionError, RuntimeError)): + worker.add_remote_agent(meta_r, remote_tp_rank=0, remote_tp_size=2) diff --git a/tests/v1/kv_connector/unit/test_nixl_push_connector.py b/tests/v1/kv_connector/unit/test_nixl_push_connector.py index c72953ca4b..a42fbe9243 100644 --- a/tests/v1/kv_connector/unit/test_nixl_push_connector.py +++ b/tests/v1/kv_connector/unit/test_nixl_push_connector.py @@ -43,6 +43,7 @@ from aphrodite.distributed.kv_transfer.kv_connector.v1.nixl.utils import ( get_base_request_id, ) +from aphrodite.v1.kv_cache_interface import FullAttentionSpec from aphrodite.v1.outputs import KVConnectorOutput from .utils import make_nixl_push_scheduler @@ -323,6 +324,9 @@ def fresh(cls) -> _StubWriterWorker: w.engine_id = "test-decode-engine" w._remote_agents = {} w._physical_blocks_per_logical_kv_block = 1 + # Single non-hybrid attention group, matching the stub block id lists. + w._has_mamba = False + w._group_spec_types = (FullAttentionSpec,) # Track _do_start_push_kv invocations. calls: list[tuple[str, Any, dict[str, Any]]] = [] diff --git a/tests/v1/kv_offload/cpu/__init__.py b/tests/v1/kv_offload/cpu/__init__.py new file mode 100644 index 0000000000..8b13789179 --- /dev/null +++ b/tests/v1/kv_offload/cpu/__init__.py @@ -0,0 +1 @@ + diff --git a/tests/v1/kv_offload/cpu/policies/__init__.py b/tests/v1/kv_offload/cpu/policies/__init__.py new file mode 100644 index 0000000000..8b13789179 --- /dev/null +++ b/tests/v1/kv_offload/cpu/policies/__init__.py @@ -0,0 +1 @@ + diff --git a/tests/v1/kv_offload/cpu/policies/test_factory.py b/tests/v1/kv_offload/cpu/policies/test_factory.py new file mode 100644 index 0000000000..1cb8318d63 --- /dev/null +++ b/tests/v1/kv_offload/cpu/policies/test_factory.py @@ -0,0 +1,92 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from collections.abc import Iterable + +import pytest + +from aphrodite.v1.kv_offload.base import OffloadKey, ReqContext +from aphrodite.v1.kv_offload.cpu.manager import CPUOffloadingManager +from aphrodite.v1.kv_offload.cpu.policies.arc import ARCCachePolicy +from aphrodite.v1.kv_offload.cpu.policies.base import BlockStatus, CachePolicy +from aphrodite.v1.kv_offload.cpu.policies.factory import CachePolicyFactory +from aphrodite.v1.kv_offload.cpu.policies.lru import LRUCachePolicy + + +class _DummyCachePolicy(CachePolicy): + def __init__(self, cache_capacity: int) -> None: + super().__init__(cache_capacity) + + def get(self, key: OffloadKey) -> BlockStatus | None: + return None + + def insert(self, key: OffloadKey, block: BlockStatus) -> None: + pass + + def remove(self, key: OffloadKey) -> None: + pass + + def touch(self, keys: Iterable[OffloadKey], req_context: ReqContext) -> None: + pass + + def evict(self, n: int, protected: set[OffloadKey]) -> list[tuple[OffloadKey, BlockStatus]] | None: + return None + + def clear(self) -> None: + pass + + +@pytest.fixture(autouse=True) +def restore_cache_policy_registry(): + original = dict(CachePolicyFactory._registry) + yield + CachePolicyFactory._registry = original + + +class TestCachePolicyFactory: + def test_pre_registered_policies_can_be_imported(self): + for name in CachePolicyFactory._registry: + policy_cls = CachePolicyFactory._registry[name]() + assert issubclass(policy_cls, CachePolicy) + + def test_lru_and_arc_registered(self): + assert CachePolicyFactory.get_cache_policy_cls("lru") is LRUCachePolicy + assert CachePolicyFactory.get_cache_policy_cls("arc") is ARCCachePolicy + + def test_register_and_resolve_custom_policy(self): + CachePolicyFactory.register_cache_policy( + "dummy", + "tests.v1.kv_offload.cpu.policies.test_factory", + "_DummyCachePolicy", + ) + policy_cls = CachePolicyFactory.get_cache_policy_cls("dummy") + assert policy_cls is _DummyCachePolicy + + manager = CPUOffloadingManager(num_blocks=4, cache_policy="dummy") + assert isinstance(manager._policy, _DummyCachePolicy) + + def test_unregistered_policy_raises(self): + with pytest.raises(ValueError, match="Unknown cache policy"): + CachePolicyFactory.get_cache_policy_cls("nonexistent") + + def test_duplicate_registration_raises(self): + with pytest.raises(ValueError, match="is already registered"): + CachePolicyFactory.register_cache_policy("lru", "some.module", "SomeClass") + + def test_dynamic_load_via_cache_policy_module_path(self): + policy_cls = CachePolicyFactory.get_cache_policy_cls( + "_DummyCachePolicy", + "tests.v1.kv_offload.cpu.policies.test_factory", + ) + assert policy_cls is _DummyCachePolicy + + def test_manager_resolves_policy_via_module_path(self): + manager = CPUOffloadingManager( + num_blocks=4, + cache_policy="_DummyCachePolicy", + cache_policy_module_path="tests.v1.kv_offload.cpu.policies.test_factory", + ) + assert isinstance(manager._policy, _DummyCachePolicy) + + def test_unregistered_policy_without_module_path_raises(self): + with pytest.raises(ValueError, match="Unknown cache policy"): + CachePolicyFactory.get_cache_policy_cls("nonexistent", None) diff --git a/tests/v1/kv_offload/cpu/test_manager.py b/tests/v1/kv_offload/cpu/test_manager.py index d8640a78bc..7d00676d12 100644 --- a/tests/v1/kv_offload/cpu/test_manager.py +++ b/tests/v1/kv_offload/cpu/test_manager.py @@ -35,6 +35,7 @@ def make_req_context(req_id: str = "", kv_transfer_params: dict | None = None) - def make_cpu_manager( num_blocks: int = 4, cache_policy: str = "lru", + cache_policy_module_path: str | None = None, enable_events: bool = False, store_threshold: int = 0, max_tracker_size: int = 64_000, @@ -42,6 +43,7 @@ def make_cpu_manager( return CPUOffloadingManager( num_blocks=num_blocks, cache_policy=cache_policy, + cache_policy_module_path=cache_policy_module_path, enable_events=enable_events, store_threshold=store_threshold, max_tracker_size=max_tracker_size, diff --git a/tests/v1/spec_decode/test_llm_base_proposer.py b/tests/v1/spec_decode/test_llm_base_proposer.py new file mode 100644 index 0000000000..cd2079e294 --- /dev/null +++ b/tests/v1/spec_decode/test_llm_base_proposer.py @@ -0,0 +1,97 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from types import SimpleNamespace + +import pytest + +import aphrodite.v1.spec_decode.llm_base_proposer as llm_base_proposer +from aphrodite.v1.spec_decode.eagle import EagleProposer + +SCHEDULER_BLOCK_SIZE = 256 +KERNEL_BLOCK_SIZE = 64 + + +class _FakeAttentionGroup: + def __init__(self, backend, layer_names, kv_cache_spec, kv_cache_group_id): + self.backend = backend + self.layer_names = list(layer_names) + self.kv_cache_spec = kv_cache_spec + self.kv_cache_group_id = kv_cache_group_id + self.kernel_block_size = None + + def create_metadata_builders(self, aphrodite_config, device, kernel_block_size=None): + self.kernel_block_size = kernel_block_size + + def get_metadata_builder(self): + return SimpleNamespace(kv_cache_spec=self.kv_cache_spec) + + +def _make_proposer(monkeypatch: pytest.MonkeyPatch, layer_names: set[str]) -> EagleProposer: + fake_layers = {} + for name in layer_names: + backend = SimpleNamespace(full_cls_name=lambda: "FakeBackend") + fake_layers[name] = SimpleNamespace(get_attn_backend=lambda backend=backend: backend) + monkeypatch.setattr( + llm_base_proposer, + "get_layers_from_aphrodite_config", + lambda *args, **kwargs: fake_layers, + ) + monkeypatch.setattr(llm_base_proposer, "AttentionGroup", _FakeAttentionGroup) + + proposer = EagleProposer.__new__(EagleProposer) + proposer.aphrodite_config = None + proposer.device = None + proposer._draft_attn_layer_names = set(layer_names) + proposer.kv_cache_gid = -1 + proposer.draft_attn_groups = [] + proposer.block_size = -1 + return proposer + + +def _make_kv_cache_config(layer_names: set[str]) -> SimpleNamespace: + spec = SimpleNamespace(block_size=SCHEDULER_BLOCK_SIZE) + group = SimpleNamespace(layer_names=list(layer_names), kv_cache_spec=spec) + return SimpleNamespace(kv_cache_groups=[group]) + + +def test_block_size_uses_kernel_block_size(monkeypatch: pytest.MonkeyPatch): + layer_names = {"draft.0.self_attn.attn"} + proposer = _make_proposer(monkeypatch, layer_names) + + proposer.initialize_attn_backend( + _make_kv_cache_config(layer_names), + kernel_block_sizes=[KERNEL_BLOCK_SIZE], + ) + + assert proposer.block_size == KERNEL_BLOCK_SIZE + assert proposer.block_size != SCHEDULER_BLOCK_SIZE + assert proposer.draft_attn_groups[0].kernel_block_size == KERNEL_BLOCK_SIZE + + +def test_block_size_falls_back_to_kv_cache_spec( + monkeypatch: pytest.MonkeyPatch, +): + layer_names = {"draft.0.self_attn.attn"} + proposer = _make_proposer(monkeypatch, layer_names) + + proposer.initialize_attn_backend(_make_kv_cache_config(layer_names), kernel_block_sizes=None) + + assert proposer.block_size == SCHEDULER_BLOCK_SIZE + + +def test_draft_layer_iteration_is_deterministic( + monkeypatch: pytest.MonkeyPatch, +): + layer_names = {"draft.c.attn", "draft.a.attn", "draft.b.attn"} + expected_order = sorted(layer_names) + + for insertion_order in (expected_order, expected_order[::-1]): + proposer = _make_proposer(monkeypatch, set(insertion_order)) + proposer.initialize_attn_backend( + _make_kv_cache_config(set(insertion_order)), + kernel_block_sizes=[KERNEL_BLOCK_SIZE], + ) + assert len(proposer.draft_attn_groups) == 1 + assert proposer.draft_attn_groups[0].layer_names == expected_order + assert proposer.block_size == KERNEL_BLOCK_SIZE diff --git a/tests/v1/worker/test_gpu_block_table.py b/tests/v1/worker/test_gpu_block_table.py index 9ce059b37d..1a5a5067cb 100644 --- a/tests/v1/worker/test_gpu_block_table.py +++ b/tests/v1/worker/test_gpu_block_table.py @@ -130,3 +130,66 @@ def test_block_tables_apply_staged_writes_single_group(): block_tables.block_tables[0].gpu[0, :2], torch.tensor([1, 2], dtype=torch.int32, device=device), ) + + +def test_v1_block_table_move_row_clears_vacated_row(): + """condense() moves the last row into a freed slot; the vacated row must + not keep stale block ids. Padded dummy-run batches dereference stale rows + as mamba state slots (bypassing the NULL_BLOCK_ID fill of real decode + padding) and write state in place there — corrupting the blocks' new + owner once they are reallocated, e.g. to an in-flight NIXL load.""" + from aphrodite.v1.worker.block_table import BlockTable + + block_table = BlockTable( + block_size=16, + max_num_reqs=4, + max_num_blocks_per_req=8, + max_num_batched_tokens=64, + pin_memory=False, + device=torch.device("cuda"), + kernel_block_size=16, + cp_kv_cache_interleave_size=1, + ) + block_table.add_row([7, 8, 9], row_idx=0) + block_table.add_row([4, 5], row_idx=1) + + block_table.move_row(1, 0) + + assert block_table.block_table.np[0, :2].tolist() == [4, 5] + assert block_table.num_blocks_per_row[0] == 2 + # The vacated source row routes to the reserved null block. + assert block_table.num_blocks_per_row[1] == 0 + assert (block_table.block_table.np[1] == 0).all() + + +def test_get_dummy_block_tables_returns_zeroed_rows(): + """Dummy runs bypass the gather, so the persistent input_block_tables + hold the previous real step's rows. Mamba/GDN metadata routes in-place + state writes through block_table[:, 0] (dummy slot mappings are + PAD-filled, state indices are not), so stale rows would direct dummy + state writes at freed — possibly reallocated — blocks. + get_dummy_block_tables must hand out zeroed (null block) rows while + preserving the persistent storage address for CUDA graphs.""" + device = torch.device("cuda") + block_tables = BlockTables( + block_sizes=[16], + max_num_reqs=4, + max_num_batched_tokens=64, + max_num_blocks_per_group=[8], + device=device, + kernel_block_sizes=[16], + ) + # Simulate a real step: stage a request's blocks and gather them into + # the persistent input block tables. + block_tables.append_block_ids(req_index=0, new_block_ids=([1, 2],), overwrite=True) + block_tables.apply_staged_writes() + idx_mapping = torch.zeros(1, dtype=torch.int32, device=device) + block_tables.gather_block_tables(idx_mapping, num_reqs_padded=1) + torch.accelerator.synchronize() + assert block_tables.input_block_tables[0][0, 0].item() == 1 + + dummy = block_tables.get_dummy_block_tables(num_reqs=1) + torch.accelerator.synchronize() + assert (dummy[0] == 0).all() + # CUDA graph invariant: same persistent tensor, not a fresh allocation. + assert dummy[0].data_ptr() == block_tables.input_block_tables[0].data_ptr() diff --git a/tests/v1/worker/test_gpu_model_runner.py b/tests/v1/worker/test_gpu_model_runner.py index 31c3fef19f..62ab8109d7 100644 --- a/tests/v1/worker/test_gpu_model_runner.py +++ b/tests/v1/worker/test_gpu_model_runner.py @@ -1609,3 +1609,57 @@ def test_mamba_cache_raises_when_max_num_seqs_exceeds_blocks(): with pytest.raises(ValueError, match="max_num_seqs"): runner.initialize_kv_cache(kv_cache_config) + + +class TestInitFp8KvScalesHybridModels: + """Verify init_fp8_kv_scales handles heterogeneous kv_caches entries. + + Hybrid models (Mamba, DeltaNet) store per-layer state as a list of tensors + rather than a single tensor. init_fp8_kv_scales must iterate both forms. + """ + + @staticmethod + def _make_runner_stub(kv_caches): + runner = Mock(spec=GPUModelRunner) + runner.cache_config = SimpleNamespace(cache_dtype="fp8_e4m3") + runner.kv_caches = kv_caches + runner.compilation_config = SimpleNamespace(static_forward_context={}) + runner.init_fp8_kv_scales = GPUModelRunner.init_fp8_kv_scales.__get__(runner, GPUModelRunner) + return runner + + def test_zeroes_both_tensor_and_list_entries(self): + single_tensor = torch.ones(4, 8) + list_tensors = [torch.ones(2, 4), torch.ones(3, 6)] + + runner = self._make_runner_stub([single_tensor, list_tensors]) + runner.init_fp8_kv_scales() + + assert (single_tensor == 0).all() + assert all((t == 0).all() for t in list_tensors) + + def test_skips_none_entries(self): + tensor = torch.ones(4, 8) + runner = self._make_runner_stub([None, tensor, None]) + runner.init_fp8_kv_scales() + + assert (tensor == 0).all() + + def test_noop_when_kv_cache_not_quantized(self): + tensor = torch.ones(4, 8) + runner = self._make_runner_stub([tensor]) + runner.cache_config.cache_dtype = "auto" + runner.init_fp8_kv_scales() + + assert (tensor == 1).all() + + def test_mixed_none_tensor_and_list(self): + t1 = torch.ones(2, 2) + t2 = torch.ones(3, 3) + list_entry = [torch.ones(1, 1), torch.ones(1, 1)] + + runner = self._make_runner_stub([None, t1, list_entry, None, t2]) + runner.init_fp8_kv_scales() + + assert (t1 == 0).all() + assert (t2 == 0).all() + assert all((t == 0).all() for t in list_entry) diff --git a/tools/pre_commit/check_pickle_imports.py b/tools/pre_commit/check_pickle_imports.py index 64c1f63937..4921b739cd 100644 --- a/tools/pre_commit/check_pickle_imports.py +++ b/tools/pre_commit/check_pickle_imports.py @@ -53,6 +53,7 @@ "aphrodite/diffusion/runtime/models/registry.py", # tests exercising serialization paths "tests/compile/test_aot_compile.py", + "tests/distributed/test_shm_broadcast.py", "tests/distributed/test_weight_transfer.py", "tests/multimodal/media/test_base.py", "tests/tokenizers_/test_hf.py",