From fca1511e7691cac53b9ce615ef8db1d5f251d0d6 Mon Sep 17 00:00:00 2001 From: AlpinDale Date: Wed, 29 Jul 2026 10:07:22 +0000 Subject: [PATCH 01/41] [sync] [Bugfix][Multimodal] Fix video temporal padding estimates (#49030) Upstream-vLLM: 1db989bbf14b9896c1306fd722ced9a9b1238465 Co-authored-by: labAxiaoming <34019940+labAxiaoming@users.noreply.github.com> --- .sync/vllm-sha | 2 +- aphrodite/model_executor/models/glm4_1v.py | 4 +-- aphrodite/model_executor/models/kanana_v.py | 4 +-- aphrodite/model_executor/models/keye.py | 2 +- .../model_executor/models/llava_onevision2.py | 2 +- .../model_executor/models/mimo_v2_omni.py | 2 +- aphrodite/model_executor/models/qwen2_vl.py | 4 +-- .../multimodal/processing/test_glm4_1v.py | 25 +++++++++++++++++++ 8 files changed, 35 insertions(+), 10 deletions(-) diff --git a/.sync/vllm-sha b/.sync/vllm-sha index 7be2e43b48..d88337440f 100644 --- a/.sync/vllm-sha +++ b/.sync/vllm-sha @@ -1 +1 @@ -8a7b3c299053efbca2669081ddf50a46b6a9d149 +1db989bbf14b9896c1306fd722ced9a9b1238465 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/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( From c43a4b5e0e0ac4c8b70d7306e669744e038cba79 Mon Sep 17 00:00:00 2001 From: AlpinDale Date: Wed, 29 Jul 2026 10:10:10 +0000 Subject: [PATCH 02/41] [sync] [ROCm][Bugfix] Sanitize AITER paged-MQA logits before sparse top-k for DeepSeek-V4 (#49714) Upstream-vLLM: 6c7e679f048dc6123caecc3985150766e455ff22 Co-authored-by: Shanshan Shen <467638484@qq.com> --- .sync/vllm-sha | 2 +- .../v1/attention/ops/rocm_aiter_mla_sparse.py | 1 + .../attention/test_rocm_triton_attn_dsv4.py | 59 +++++++++++++++++++ 3 files changed, 61 insertions(+), 1 deletion(-) diff --git a/.sync/vllm-sha b/.sync/vllm-sha index d88337440f..0b7996aac2 100644 --- a/.sync/vllm-sha +++ b/.sync/vllm-sha @@ -1 +1 @@ -1db989bbf14b9896c1306fd722ced9a9b1238465 +6c7e679f048dc6123caecc3985150766e455ff22 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/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 ( From 8fd92b195c58391874da74a4828e9f86ddaef990 Mon Sep 17 00:00:00 2001 From: AlpinDale Date: Wed, 29 Jul 2026 10:11:29 +0000 Subject: [PATCH 03/41] [sync] [BugFix] Fix clang spinloop mwaitx include (#45532) Upstream-vLLM: 118bcde449c923ff20fe6336a009e2cef2ba5423 Co-authored-by: johnnyychiu --- .sync/vllm-sha | 2 +- csrc/spinloop.cpp | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/.sync/vllm-sha b/.sync/vllm-sha index 0b7996aac2..66689e759a 100644 --- a/.sync/vllm-sha +++ b/.sync/vllm-sha @@ -1 +1 @@ -6c7e679f048dc6123caecc3985150766e455ff22 +118bcde449c923ff20fe6336a009e2cef2ba5423 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) From 6c3fd7228ad243011ffbd4eb907f2b5cd0371a59 Mon Sep 17 00:00:00 2001 From: AlpinDale Date: Wed, 29 Jul 2026 10:13:54 +0000 Subject: [PATCH 04/41] [sync] [Rust Frontend] Extract shared tracing setup logic into `vllm-tracing` (#50129) Upstream-vLLM: d552a68645b60fe194ed488e9c3df211e4d2fa6f Co-authored-by: Bugen Zhao --- .sync/vllm-sha | 2 +- rust/Cargo.lock | 14 +++++++++++--- rust/Cargo.toml | 2 ++ rust/src/bench/Cargo.toml | 2 +- rust/src/bench/src/main.rs | 12 +----------- rust/src/cmd/Cargo.toml | 3 +-- rust/src/cmd/src/main.rs | 3 +-- rust/src/tracing/Cargo.toml | 14 ++++++++++++++ .../src/{cmd/src/logging.rs => tracing/src/lib.rs} | 6 ++++-- 9 files changed, 36 insertions(+), 22 deletions(-) create mode 100644 rust/src/tracing/Cargo.toml rename rust/src/{cmd/src/logging.rs => tracing/src/lib.rs} (98%) diff --git a/.sync/vllm-sha b/.sync/vllm-sha index 66689e759a..043768888b 100644 --- a/.sync/vllm-sha +++ b/.sync/vllm-sha @@ -1 +1 @@ -118bcde449c923ff20fe6336a009e2cef2ba5423 +d552a68645b60fe194ed488e9c3df211e4d2fa6f diff --git a/rust/Cargo.lock b/rust/Cargo.lock index 4ee1b4be81..88af7ed147 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", ] @@ -217,6 +217,7 @@ dependencies = [ "aphrodite-engine-core-client", "aphrodite-managed-engine", "aphrodite-server", + "aphrodite-tracing", "clap", "educe", "expect-test", @@ -227,11 +228,9 @@ dependencies = [ "serde_json", "serde_with", "thiserror-ext", - "time", "tokio", "tokio-util", "tracing", - "tracing-subscriber", "uuid", ] @@ -468,6 +467,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" diff --git a/rust/Cargo.toml b/rust/Cargo.toml index d0b3ada027..0218a2c909 100644 --- a/rust/Cargo.toml +++ b/rust/Cargo.toml @@ -13,6 +13,7 @@ members = [ "src/server", "src/text", "src/tokenizer", + "src/tracing", ] resolver = "3" @@ -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/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/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/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(), From 907ea785df8b1e013e1991888eca2ed6cb356f82 Mon Sep 17 00:00:00 2001 From: AlpinDale Date: Wed, 29 Jul 2026 10:15:20 +0000 Subject: [PATCH 05/41] [sync] [Bugfix][Spec Decode] Size DFlash query buffers for cudagraph-padded batches (#50065) Upstream-vLLM: 0b6aa3c47ce69a1f3c8a19cafe3b9dc2871d1f6b Co-authored-by: Siddhant Bharti <42143349+siddhant-bharti@users.noreply.github.com> --- .sync/vllm-sha | 2 +- aphrodite/v1/spec_decode/dflash.py | 10 +++++++--- 2 files changed, 8 insertions(+), 4 deletions(-) diff --git a/.sync/vllm-sha b/.sync/vllm-sha index 043768888b..58f7ba5c41 100644 --- a/.sync/vllm-sha +++ b/.sync/vllm-sha @@ -1 +1 @@ -d552a68645b60fe194ed488e9c3df211e4d2fa6f +0b6aa3c47ce69a1f3c8a19cafe3b9dc2871d1f6b 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, ) From d10f05ab758310ed46e61b03ab168af380637fb1 Mon Sep 17 00:00:00 2001 From: AlpinDale Date: Wed, 29 Jul 2026 10:18:28 +0000 Subject: [PATCH 06/41] [sync] [Model] Add Kimi K3 support: Rust frontend [1/2] (#50104) Upstream-vLLM: 2899dca8432d40632987b0ec24253a8fe6df2710 Co-authored-by: Bugen Zhao --- .sync/vllm-sha | 2 +- pyproject.toml | 6 +- rust/Cargo.lock | 2 +- rust/Cargo.toml | 2 +- rust/src/chat/src/backend/hf.rs | 3 +- rust/src/chat/src/lib.rs | 13 +- rust/src/chat/src/parser/reasoning/mod.rs | 2 + rust/src/chat/src/parser/tool/mod.rs | 2 + rust/src/chat/src/parser/unified.rs | 26 +- .../chat/src/renderer/deepseek_v32/tests.rs | 2 +- .../chat/src/renderer/deepseek_v4/tests.rs | 2 +- rust/src/chat/src/renderer/harmony/tests.rs | 2 +- rust/src/chat/src/renderer/inkling/tests.rs | 2 +- .../src/chat/src/renderer/kimi_k3/encoding.rs | 608 +++++++++ .../fixtures/controls_thinking_off_input.json | 15 + .../fixtures/controls_thinking_off_output.txt | 2 + .../dynamic_system_tool_declare_input.json | 62 + .../dynamic_system_tool_declare_output.txt | 14 + .../history_preserve_and_image_input.json | 31 + .../history_preserve_and_image_output.txt | 2 + .../tools_history_and_required_input.json | 87 ++ .../tools_history_and_required_output.txt | 8 + rust/src/chat/src/renderer/kimi_k3/mod.rs | 39 + rust/src/chat/src/renderer/kimi_k3/tests.rs | 267 ++++ rust/src/chat/src/renderer/mod.rs | 2 + rust/src/chat/src/renderer/selection.rs | 9 +- rust/src/chat/src/renderer/test_utils.rs | 52 +- rust/src/chat/src/request.rs | 5 + rust/src/chat/tests/roundtrip.rs | 42 +- rust/src/cmd/src/cli/tests.rs | 2 +- rust/src/parser/src/unified/kimi_k3.rs | 1149 +++++++++++++++++ .../src/unified/kimi_k3/structural_tag.rs | 503 ++++++++ rust/src/parser/src/unified/mod.rs | 2 + .../routes/openai/chat_completions/convert.rs | 73 +- rust/src/server/src/routes/tokenize/types.rs | 1 + rust/src/tokenizer/src/tiktoken.rs | 4 +- 36 files changed, 3006 insertions(+), 39 deletions(-) create mode 100644 rust/src/chat/src/renderer/kimi_k3/encoding.rs create mode 100644 rust/src/chat/src/renderer/kimi_k3/fixtures/controls_thinking_off_input.json create mode 100644 rust/src/chat/src/renderer/kimi_k3/fixtures/controls_thinking_off_output.txt create mode 100644 rust/src/chat/src/renderer/kimi_k3/fixtures/dynamic_system_tool_declare_input.json create mode 100644 rust/src/chat/src/renderer/kimi_k3/fixtures/dynamic_system_tool_declare_output.txt create mode 100644 rust/src/chat/src/renderer/kimi_k3/fixtures/history_preserve_and_image_input.json create mode 100644 rust/src/chat/src/renderer/kimi_k3/fixtures/history_preserve_and_image_output.txt create mode 100644 rust/src/chat/src/renderer/kimi_k3/fixtures/tools_history_and_required_input.json create mode 100644 rust/src/chat/src/renderer/kimi_k3/fixtures/tools_history_and_required_output.txt create mode 100644 rust/src/chat/src/renderer/kimi_k3/mod.rs create mode 100644 rust/src/chat/src/renderer/kimi_k3/tests.rs create mode 100644 rust/src/parser/src/unified/kimi_k3.rs create mode 100644 rust/src/parser/src/unified/kimi_k3/structural_tag.rs diff --git a/.sync/vllm-sha b/.sync/vllm-sha index 58f7ba5c41..91a044c68a 100644 --- a/.sync/vllm-sha +++ b/.sync/vllm-sha @@ -1 +1 @@ -0b6aa3c47ce69a1f3c8a19cafe3b9dc2871d1f6b +2899dca8432d40632987b0ec24253a8fe6df2710 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/rust/Cargo.lock b/rust/Cargo.lock index 88af7ed147..e7685bf91c 100644 --- a/rust/Cargo.lock +++ b/rust/Cargo.lock @@ -2530,7 +2530,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 0218a2c909..307e3a5541 100644 --- a/rust/Cargo.toml +++ b/rust/Cargo.toml @@ -59,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"] } diff --git a/rust/src/chat/src/backend/hf.rs b/rust/src/chat/src/backend/hf.rs index 54476fed62..a23ecd69fa 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}; @@ -73,6 +73,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!( 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/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/src/cli/tests.rs b/rust/src/cmd/src/cli/tests.rs index 363d158707..c5caf11670 100644 --- a/rust/src/cmd/src/cli/tests.rs +++ b/rust/src/cmd/src/cli/tests.rs @@ -666,7 +666,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'. "#]] 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/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 &'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); From 0c91ec24a30af11232b7f3191daab4deac170180 Mon Sep 17 00:00:00 2001 From: AlpinDale Date: Wed, 29 Jul 2026 10:21:55 +0000 Subject: [PATCH 07/41] [sync] perf: dispatch non-grouped bias-less topk routing methods to fused path (#49618) Upstream-vLLM: bb3b61f2fd2333ab165ebaba13f133db4210b9f2 Co-authored-by: Julien Debache --- .sync/vllm-sha | 2 +- .../layers/fused_moe/router/router_factory.py | 62 ++++++---- tests/kernels/moe/test_routing.py | 116 ++++++++++++++++++ 3 files changed, 154 insertions(+), 26 deletions(-) diff --git a/.sync/vllm-sha b/.sync/vllm-sha index 91a044c68a..14c6d52832 100644 --- a/.sync/vllm-sha +++ b/.sync/vllm-sha @@ -1 +1 @@ -2899dca8432d40632987b0ec24253a8fe6df2710 +bb3b61f2fd2333ab165ebaba13f133db4210b9f2 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/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 From 8c70aefd7cbf9c92bec53675383b6bcb89d67870 Mon Sep 17 00:00:00 2001 From: AlpinDale Date: Wed, 29 Jul 2026 10:25:30 +0000 Subject: [PATCH 08/41] [sync] [Perf] Zero-copy torch.Tensor pickling in shm_broadcast MessageQueue (#48442) Upstream-vLLM: a07fac758f2cae58f6ebad29da726e9432012c13 Co-authored-by: Ruinan Ma <97484148+mrn3088@users.noreply.github.com> --- .sync/vllm-sha | 2 +- .../device_communicators/shm_broadcast.py | 80 +++++++++- tests/distributed/test_shm_broadcast.py | 143 ++++++++++++++++++ tools/pre_commit/check_pickle_imports.py | 1 + 4 files changed, 224 insertions(+), 2 deletions(-) diff --git a/.sync/vllm-sha b/.sync/vllm-sha index 14c6d52832..c722ba7297 100644 --- a/.sync/vllm-sha +++ b/.sync/vllm-sha @@ -1 +1 @@ -bb3b61f2fd2333ab165ebaba13f133db4210b9f2 +a07fac758f2cae58f6ebad29da726e9432012c13 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/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/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", From f7ce99ff899580dbfbca1ced174026105539a743 Mon Sep 17 00:00:00 2001 From: AlpinDale Date: Wed, 29 Jul 2026 10:27:59 +0000 Subject: [PATCH 09/41] [sync] [Test] Make EPD correctness tests configurable for XPU (#50110) Upstream-vLLM: e7f6a39db8d1f05ea923b5f976141db838b11c44 Co-authored-by: liuzhenwei --- .sync/vllm-sha | 2 +- .../integration/run_epd_correctness_test.sh | 85 ++++++++++++------- .../integration/test_epd_correctness.py | 8 +- 3 files changed, 63 insertions(+), 32 deletions(-) diff --git a/.sync/vllm-sha b/.sync/vllm-sha index c722ba7297..3687ba3e92 100644 --- a/.sync/vllm-sha +++ b/.sync/vllm-sha @@ -1 +1 @@ -a07fac758f2cae58f6ebad29da726e9432012c13 +e7f6a39db8d1f05ea923b5f976141db838b11c44 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 From 4878509f84edf354b5b165e20d1faebb0f1facc5 Mon Sep 17 00:00:00 2001 From: AlpinDale Date: Wed, 29 Jul 2026 10:32:10 +0000 Subject: [PATCH 10/41] [sync] [MXFP8][ROCm] Fix MXFP8 MoE backend selection (#49747) Upstream-vLLM: 5369f7b7b89e48d396fa98a9d33f3fba654af8ee Co-authored-by: fxmarty-amd --- .sync/vllm-sha | 2 +- .../fused_moe/experts/aiter_mxfp8_moe.py | 24 +++ .../fused_moe/experts/mxfp8_emulation_moe.py | 8 +- .../fused_moe/experts/mxfp8_native_moe.py | 10 +- .../layers/fused_moe/oracle/mxfp8.py | 51 +----- .../moe/test_mxfp8_aiter_backend_selection.py | 42 ++++- tests/models/quantization/test_mxfp8.py | 166 ++++++++++++++++++ 7 files changed, 241 insertions(+), 62 deletions(-) diff --git a/.sync/vllm-sha b/.sync/vllm-sha index 3687ba3e92..56eb87c6f9 100644 --- a/.sync/vllm-sha +++ b/.sync/vllm-sha @@ -1 +1 @@ -e7f6a39db8d1f05ea923b5f976141db838b11c44 +5369f7b7b89e48d396fa98a9d33f3fba654af8ee 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/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/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/tests/kernels/moe/test_mxfp8_aiter_backend_selection.py b/tests/kernels/moe/test_mxfp8_aiter_backend_selection.py index 191657daf0..2f7ccf1241 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_native_moe import ( # noqa: E402 + Mxfp8NativeTritonExperts, +) 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,20 @@ 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.""" + with _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_triton(): + """flydsl unusable (e.g. gfx942, no FlyDSL support) -> native Triton + dot_scaled backend wins instead.""" + with _flydsl_installed(False): + backend, experts_cls = select_mxfp8_moe_backend(_config()) + assert backend is Fp8MoeBackend.TRITON_MXFP8 + assert experts_cls is Mxfp8NativeTritonExperts 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+).", From 5a81c9267caf662d1444e3a26260ac0cd71c1cf6 Mon Sep 17 00:00:00 2001 From: AlpinDale Date: Wed, 29 Jul 2026 10:33:53 +0000 Subject: [PATCH 11/41] [sync] [ROCm][CI] Stabilize ROCm audio streaming test (#50163) Upstream-vLLM: 176256b9628d2d8d22db41d7a72c15540f438026 Co-authored-by: Andreas Karatzas --- .sync/vllm-sha | 2 +- .../multimodal/openai/chat_completion/test_audio.py | 6 ++++++ 2 files changed, 7 insertions(+), 1 deletion(-) diff --git a/.sync/vllm-sha b/.sync/vllm-sha index 56eb87c6f9..fb1bebbe98 100644 --- a/.sync/vllm-sha +++ b/.sync/vllm-sha @@ -1 +1 @@ -5369f7b7b89e48d396fa98a9d33f3fba654af8ee +176256b9628d2d8d22db41d7a72c15540f438026 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: From a28ddfbd2ea7490b509e617124e1a4853887296b Mon Sep 17 00:00:00 2001 From: AlpinDale Date: Wed, 29 Jul 2026 10:35:11 +0000 Subject: [PATCH 12/41] [sync] [CI][NIXL] Fix flaky DP+EP test port conflict (#50171) Upstream-vLLM: fe65aa6a97981c815800daa4bf49d8356b8c4988 Co-authored-by: Divakar Verma <137818590+divakar-amd@users.noreply.github.com> --- .sync/vllm-sha | 2 +- tests/v1/kv_connector/nixl_integration/run_accuracy_test.sh | 5 ++++- 2 files changed, 5 insertions(+), 2 deletions(-) diff --git a/.sync/vllm-sha b/.sync/vllm-sha index fb1bebbe98..2decf6ca64 100644 --- a/.sync/vllm-sha +++ b/.sync/vllm-sha @@ -1 +1 @@ -176256b9628d2d8d22db41d7a72c15540f438026 +fe65aa6a97981c815800daa4bf49d8356b8c4988 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" From 874dc326fa1e0aa10b38934901bf744f616bdc37 Mon Sep 17 00:00:00 2001 From: AlpinDale Date: Wed, 29 Jul 2026 10:37:48 +0000 Subject: [PATCH 13/41] [sync] [Bugfix] Fix /wake_up crash on hybrid models (Mamba/DeltaNet) (#41602) Upstream-vLLM: 56f31af62afe6553369b14af225bfafb788e99a0 Co-authored-by: Kevin Glynn --- .sync/vllm-sha | 2 +- aphrodite/v1/worker/gpu_model_runner.py | 13 ++++-- tests/v1/worker/test_gpu_model_runner.py | 54 ++++++++++++++++++++++++ 3 files changed, 65 insertions(+), 4 deletions(-) diff --git a/.sync/vllm-sha b/.sync/vllm-sha index 2decf6ca64..6de30644a6 100644 --- a/.sync/vllm-sha +++ b/.sync/vllm-sha @@ -1 +1 @@ -fe65aa6a97981c815800daa4bf49d8356b8c4988 +56f31af62afe6553369b14af225bfafb788e99a0 diff --git a/aphrodite/v1/worker/gpu_model_runner.py b/aphrodite/v1/worker/gpu_model_runner.py index f2698c149c..098044bb2f 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") 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) From ebb8534b6061020fa27f99952503dba7d669b13e Mon Sep 17 00:00:00 2001 From: AlpinDale Date: Wed, 29 Jul 2026 10:40:57 +0000 Subject: [PATCH 14/41] [sync] [BugFix] Stop dummy runs from writing mamba state through stale block-table rows (#49757) Upstream-vLLM: 6fbbcf2151c34cd9313c223b8cbf19afa77f2ea6 Co-authored-by: Nick Hill --- .sync/vllm-sha | 2 +- aphrodite/v1/worker/block_table.py | 5 ++ aphrodite/v1/worker/gpu/block_table.py | 6 ++- aphrodite/v1/worker/gpu/model_runner.py | 4 ++ aphrodite/v1/worker/gpu_model_runner.py | 5 +- tests/v1/worker/test_gpu_block_table.py | 63 +++++++++++++++++++++++++ 6 files changed, 82 insertions(+), 3 deletions(-) diff --git a/.sync/vllm-sha b/.sync/vllm-sha index 6de30644a6..88dc9d23dc 100644 --- a/.sync/vllm-sha +++ b/.sync/vllm-sha @@ -1 +1 @@ -56f31af62afe6553369b14af225bfafb788e99a0 +6fbbcf2151c34cd9313c223b8cbf19afa77f2ea6 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 098044bb2f..2b151fa916 100644 --- a/aphrodite/v1/worker/gpu_model_runner.py +++ b/aphrodite/v1/worker/gpu_model_runner.py @@ -5546,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/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() From 820c72d950af3a79097a37a6eb7771b03ef4ef38 Mon Sep 17 00:00:00 2001 From: AlpinDale Date: Wed, 29 Jul 2026 10:44:30 +0000 Subject: [PATCH 15/41] [sync] [Model] Add Inkling compressed-tensors dynamic FP8 support (#48876) Upstream-vLLM: 17a74b745b459eb72ea9eea2945bc83b1c17e0a0 Co-authored-by: Krishna Teja Chitty-Venkata <44275589+krishnateja95@users.noreply.github.com> --- .sync/vllm-sha | 2 +- aphrodite/models/inkling/nvidia/moe.py | 3 +++ .../models/inkling/test_moe_weight_layout.py | 22 +++++++++++++++++++ 3 files changed, 26 insertions(+), 1 deletion(-) diff --git a/.sync/vllm-sha b/.sync/vllm-sha index 88dc9d23dc..0822593344 100644 --- a/.sync/vllm-sha +++ b/.sync/vllm-sha @@ -1 +1 @@ -6fbbcf2151c34cd9313c223b8cbf19afa77f2ea6 +17a74b745b459eb72ea9eea2945bc83b1c17e0a0 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/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) From 0663b5a0340e824705eaad1d30cdf5d356aa241d Mon Sep 17 00:00:00 2001 From: AlpinDale Date: Wed, 29 Jul 2026 10:45:58 +0000 Subject: [PATCH 16/41] [sync] [ROCm][CI] Stabilize ngram and suffix correctness test (#50190) Upstream-vLLM: 7398a30d79758559704d4f28173b3594b2ec1345 Co-authored-by: Andreas Karatzas --- .sync/vllm-sha | 2 +- tests/v1/e2e/spec_decode/test_spec_decode.py | 27 ++++++++++++++------ 2 files changed, 20 insertions(+), 9 deletions(-) diff --git a/.sync/vllm-sha b/.sync/vllm-sha index 0822593344..cc6037fb8d 100644 --- a/.sync/vllm-sha +++ b/.sync/vllm-sha @@ -1 +1 @@ -17a74b745b459eb72ea9eea2945bc83b1c17e0a0 +7398a30d79758559704d4f28173b3594b2ec1345 diff --git a/tests/v1/e2e/spec_decode/test_spec_decode.py b/tests/v1/e2e/spec_decode/test_spec_decode.py index ff5d1dfbcb..a834a2d9ed 100644 --- a/tests/v1/e2e/spec_decode/test_spec_decode.py +++ b/tests/v1/e2e/spec_decode/test_spec_decode.py @@ -14,7 +14,7 @@ from aphrodite.assets.base import APHRODITE_S3_BUCKET_URL from aphrodite.assets.image import VLM_IMAGES_DIR from aphrodite.benchmarks.datasets import InstructCoderDataset -from aphrodite.config import AphroditeConfig, replace +from aphrodite.config import AphroditeConfig, CompilationConfig, replace from aphrodite.distributed import cleanup_dist_env_and_memory from aphrodite.engine.arg_utils import EngineArgs from aphrodite.platforms import current_platform @@ -152,6 +152,12 @@ def reset_torch_dynamo(): torch._dynamo.reset() +@pytest.fixture +def disable_aphrodite_compile_cache_on_rocm(request: pytest.FixtureRequest) -> 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"]) From a4aa34b645e81af2b1681f460dc7a7402f89a265 Mon Sep 17 00:00:00 2001 From: AlpinDale Date: Wed, 29 Jul 2026 10:49:19 +0000 Subject: [PATCH 17/41] [sync] [CompressedTensors] FP4 Qutlass Integration (#43229) Upstream-vLLM: 30c2718eaafc230927f0d8ac47439d040b46a7c3 Co-authored-by: Kyle Sayers --- .sync/vllm-sha | 2 +- aphrodite/_custom_ops.py | 41 +- aphrodite/_custom_ops.py.orig | 4576 +++++++++++++++++ aphrodite/model_executor/layers/linear.py | 1 + .../compressed_tensors/transform/linear.py | 4 +- .../compressed_tensors/transform/module.py | 22 +- .../transform/module.py.orig | 157 + .../transform/schemes/linear_qutlass_nvfp4.py | 95 +- .../layers/quantization/qutlass_utils.py | 11 +- .../layers/quantization/qutlass_utils.py.orig | 178 + tests/quantization/test_compressed_tensors.py | 5 + .../test_compressed_tensors.py.orig | 919 ++++ 12 files changed, 5975 insertions(+), 36 deletions(-) create mode 100644 aphrodite/_custom_ops.py.orig create mode 100644 aphrodite/model_executor/layers/quantization/compressed_tensors/transform/module.py.orig create mode 100644 aphrodite/model_executor/layers/quantization/qutlass_utils.py.orig create mode 100644 tests/quantization/test_compressed_tensors.py.orig diff --git a/.sync/vllm-sha b/.sync/vllm-sha index cc6037fb8d..c57e38d54e 100644 --- a/.sync/vllm-sha +++ b/.sync/vllm-sha @@ -1 +1 @@ -7398a30d79758559704d4f28173b3594b2ec1345 +30c2718eaafc230927f0d8ac47439d040b46a7c3 diff --git a/aphrodite/_custom_ops.py b/aphrodite/_custom_ops.py index e653f533c1..c58512051e 100644 --- a/aphrodite/_custom_ops.py +++ b/aphrodite/_custom_ops.py @@ -4503,19 +4503,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 +4513,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/_custom_ops.py.orig b/aphrodite/_custom_ops.py.orig new file mode 100644 index 0000000000..e653f533c1 --- /dev/null +++ b/aphrodite/_custom_ops.py.orig @@ -0,0 +1,4576 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from enum import IntEnum +from functools import cache +from typing import TYPE_CHECKING, Literal + +import torch + +import aphrodite.envs as envs +from aphrodite.logger import init_logger +from aphrodite.platforms import current_platform +from aphrodite.scalar_type import ScalarType +from aphrodite.utils.flashinfer import ( + flashinfer_quant_nvfp4_8x4_sf_layout, +) +from aphrodite.utils.math_utils import cdiv + +logger = init_logger(__name__) + +current_platform.import_kernels() + +if TYPE_CHECKING: + + def register_fake(fn): + return lambda name: fn +else: + try: + from torch.library import register_fake + except ImportError: + from torch.library import impl_abstract as register_fake + + +# scaled_fp4_quant functional + out variant for torch.compile buffer management + + +def create_fp4_scale_tensor( + m: int, + n: int, + device: torch.device, + is_sf_swizzled_layout: bool, +) -> torch.Tensor: + """ + Allocate the output scale tensor for scaled_fp4_quant. + + When is_sf_swizzled_layout=True, we use rounded values to store the + swizzled scales. Due to the requirement of the Tensor Core, the minimum + tile is 128x4 for the scales. So, we first pad the scales to multiples + of 128 (rows) and 4 (cols). Then, the scales (in float8_e4m3fn) are + packed into an int32 for every 4 values. More: + https://docs.nvidia.com/cuda/parallel-thread-execution/ + #tcgen05-mma-scale-factor-b-layout-4x + """ + from aphrodite.utils.math_utils import round_up + + block_size = 16 + if is_sf_swizzled_layout: + rounded_m = round_up(m, 128) + scale_n = n // block_size + rounded_n = round_up(scale_n, 4) + # Must be zero-initialized: the swizzled scale buffer is padded to + # (round_up(m, 128), round_up(scale_n, 4) // 4) but the NVFP4 quant + # kernel does not write every padded element that the downstream + # NVFP4 GEMM reads. torch.empty leaves those padded scale factors + # uninitialized, which corrupts dequantization and causes a severe + # Blackwell NVFP4 decode throughput/output-length regression. + return torch.zeros((rounded_m, rounded_n // 4), device=device, dtype=torch.int32) + else: + return torch.empty((m, n // block_size), device=device, dtype=torch.uint8) + + +def create_fp4_output_tensors( + m: int, + n: int, + device: torch.device, + is_sf_swizzled_layout: bool, + padded_n: int | None = None, +) -> tuple[torch.Tensor, torch.Tensor]: + """ + Allocate both output tensors for scaled_fp4_quant: + (quantized_output, output_scale). + + Must match the C++ scaled_fp4_quant_func allocation exactly when + ``padded_n`` is ``None``. When ``padded_n`` is provided, allocate a larger + packed-FP4 output/scale buffer so the quantization kernel can write + CUTLASS-compatible K padding directly + """ + physical_n = padded_n if padded_n is not None else n + output = torch.empty((m, physical_n // 2), device=device, dtype=torch.uint8) + output_scale = create_fp4_scale_tensor(m, physical_n, device, is_sf_swizzled_layout) + return output, output_scale + + +if hasattr(torch.ops, "_C") and hasattr(torch.ops._C, "scaled_fp4_quant"): + + @register_fake("_C::scaled_fp4_quant") + def _scaled_fp4_quant_fake( + input: torch.Tensor, + input_scale: torch.Tensor, + is_sf_swizzled_layout: bool, + ) -> tuple[torch.Tensor, torch.Tensor]: + n = input.shape[-1] + m = input.numel() // n + return create_fp4_output_tensors(m, n, input.device, is_sf_swizzled_layout) + + @register_fake("_C::scaled_fp4_quant.out") + def _scaled_fp4_quant_out_fake( + input: torch.Tensor, + input_scale: torch.Tensor, + is_sf_swizzled_layout: bool, + *, + output: torch.Tensor, + output_scale: torch.Tensor, + ) -> None: + return None + + +def exl3_gemm( + a: torch.Tensor, + b: torch.Tensor, + c: torch.Tensor, + suh: torch.Tensor | None, + a_had: torch.Tensor | None, + svh: torch.Tensor | None, + force_shape_idx: int = -1, + mcg: bool = False, + mul1: bool = False, + force_num_sms: int = 0, +) -> None: + torch.ops._C.exl3_gemm( + a, + b, + c, + suh, + a_had, + svh, + force_shape_idx, + mcg, + mul1, + force_num_sms, + ) + + +def exl3_mgemm( + a: torch.Tensor, + b: torch.Tensor, + c: torch.Tensor, + suh: torch.Tensor, + a_had: torch.Tensor, + svh: torch.Tensor, + indices: torch.Tensor | None, + weights: torch.Tensor | None, + k: int, + force_shape_idx: int = -1, + mcg: bool = False, + mul1: bool = False, + min_index: int = -1, + max_index: int = -1, + force_num_sms: int = 0, +) -> None: + torch.ops._C.exl3_mgemm( + a, + b, + c, + suh, + a_had, + svh, + indices, + weights, + k, + force_shape_idx, + mcg, + mul1, + min_index, + max_index, + force_num_sms, + ) + + +def exl3_reconstruct( + unpacked: torch.Tensor, + packed: torch.Tensor, + k: int, + mcg: bool = False, + mul1: bool = False, +) -> None: + torch.ops._C.exl3_reconstruct(unpacked, packed, k, mcg, mul1) + + +def exl3_had_r_128( + input: torch.Tensor, + output: torch.Tensor, + pre_scale: torch.Tensor | None, + post_scale: torch.Tensor | None, + scale: float = 1.0, +) -> None: + torch.ops._C.exl3_had_r_128(input, output, pre_scale, post_scale, scale) + + +def exl3_hgemm(a: torch.Tensor, b: torch.Tensor, c: torch.Tensor) -> None: + torch.ops._C.exl3_hgemm(a, b, c) + + +def exl3_moe( + hidden_state: torch.Tensor, + output_state: torch.Tensor, + expert_count: torch.Tensor, + token_sorted: torch.Tensor, + weight_sorted: torch.Tensor, + temp_state_g: torch.Tensor, + temp_state_u: torch.Tensor, + temp_intermediate_g: torch.Tensor, + temp_intermediate_u: torch.Tensor, + act_function: int, + k_gate: int, + k_up: int, + k_down: int, + gate_ptrs_trellis: torch.Tensor, + gate_ptrs_suh: torch.Tensor, + gate_ptrs_svh: torch.Tensor, + up_ptrs_trellis: torch.Tensor, + up_ptrs_suh: torch.Tensor, + up_ptrs_svh: torch.Tensor, + down_ptrs_trellis: torch.Tensor, + down_ptrs_suh: torch.Tensor, + down_ptrs_svh: torch.Tensor, + gate_mcg: bool, + gate_mul1: bool, + up_mcg: bool, + up_mul1: bool, + down_mcg: bool, + down_mul1: bool, + act_limit: float = 0.0, +) -> None: + torch.ops._C.exl3_moe( + hidden_state, + output_state, + expert_count, + token_sorted, + weight_sorted, + temp_state_g, + temp_state_u, + temp_intermediate_g, + temp_intermediate_u, + act_function, + k_gate, + k_up, + k_down, + gate_ptrs_trellis, + gate_ptrs_suh, + gate_ptrs_svh, + up_ptrs_trellis, + up_ptrs_suh, + up_ptrs_svh, + down_ptrs_trellis, + down_ptrs_suh, + down_ptrs_svh, + gate_mcg, + gate_mul1, + up_mcg, + up_mul1, + down_mcg, + down_mul1, + act_limit, + ) + + +if hasattr(torch.ops, "_C") and hasattr(torch.ops._C, "exl3_gemm"): + + @register_fake("_C::exl3_gemm") + def _exl3_gemm_fake( + a: torch.Tensor, + b: torch.Tensor, + c: torch.Tensor, + suh: torch.Tensor | None, + a_had: torch.Tensor | None, + svh: torch.Tensor | None, + force_shape_idx: int, + mcg: bool, + mul1: bool, + force_num_sms: int, + ) -> None: + return None + + @register_fake("_C::exl3_mgemm") + def _exl3_mgemm_fake( + a: torch.Tensor, + b: torch.Tensor, + c: torch.Tensor, + suh: torch.Tensor, + a_had: torch.Tensor, + svh: torch.Tensor, + indices: torch.Tensor | None, + weights: torch.Tensor | None, + k: int, + force_shape_idx: int, + mcg: bool, + mul1: bool, + min_index: int, + max_index: int, + force_num_sms: int, + ) -> None: + return None + + @register_fake("_C::exl3_reconstruct") + def _exl3_reconstruct_fake( + unpacked: torch.Tensor, + packed: torch.Tensor, + k: int, + mcg: bool, + mul1: bool, + ) -> None: + return None + + @register_fake("_C::exl3_had_r_128") + def _exl3_had_r_128_fake( + input: torch.Tensor, + output: torch.Tensor, + pre_scale: torch.Tensor | None, + post_scale: torch.Tensor | None, + scale: float, + ) -> None: + return None + + @register_fake("_C::exl3_hgemm") + def _exl3_hgemm_fake( + a: torch.Tensor, + b: torch.Tensor, + c: torch.Tensor, + ) -> None: + return None + + if hasattr(torch.ops._C, "exl3_moe"): + + @register_fake("_C::exl3_moe") + def _exl3_moe_fake( + hidden_state: torch.Tensor, + output_state: torch.Tensor, + expert_count: torch.Tensor, + token_sorted: torch.Tensor, + weight_sorted: torch.Tensor, + temp_state_g: torch.Tensor, + temp_state_u: torch.Tensor, + temp_intermediate_g: torch.Tensor, + temp_intermediate_u: torch.Tensor, + act_function: int, + k_gate: int, + k_up: int, + k_down: int, + gate_ptrs_trellis: torch.Tensor, + gate_ptrs_suh: torch.Tensor, + gate_ptrs_svh: torch.Tensor, + up_ptrs_trellis: torch.Tensor, + up_ptrs_suh: torch.Tensor, + up_ptrs_svh: torch.Tensor, + down_ptrs_trellis: torch.Tensor, + down_ptrs_suh: torch.Tensor, + down_ptrs_svh: torch.Tensor, + gate_mcg: bool, + gate_mul1: bool, + up_mcg: bool, + up_mul1: bool, + down_mcg: bool, + down_mul1: bool, + act_limit: float, + ) -> None: + return None + + +def paged_attention_rocm( + out: torch.Tensor, + exp_sum: torch.Tensor, + max_logits: torch.Tensor, + tmp_out: torch.Tensor, + query: torch.Tensor, + key_cache: torch.Tensor, + value_cache: torch.Tensor, + num_kv_heads: int, + scale: float, + block_tables: torch.Tensor, + seq_lens: torch.Tensor, + query_start_loc: torch.Tensor | None, + block_size: int, + max_seq_len: int, + alibi_slopes: torch.Tensor | None, + kv_cache_dtype: str, + k_scale: torch.Tensor, + v_scale: torch.Tensor, + fp8_out_scale: torch.Tensor | None = None, + mfma_type: str = "fp8" if envs.APHRODITE_ROCM_FP8_MFMA_PAGE_ATTN else "f16", +) -> None: + torch.ops._rocm_C.paged_attention( + out, + exp_sum, + max_logits, + tmp_out, + query, + key_cache, + value_cache, + num_kv_heads, + scale, + block_tables, + seq_lens, + query_start_loc, + block_size, + max_seq_len, + alibi_slopes, + kv_cache_dtype, + k_scale, + v_scale, + fp8_out_scale, + mfma_type, + ) + + +def mla_decode_kvcache_cpu( + out: torch.Tensor, + query: torch.Tensor, + kv_cache: torch.Tensor, + scale: float, + block_tables: torch.Tensor, + seq_lens: torch.Tensor, +) -> None: + torch.ops._C.mla_decode_kvcache(out, query, kv_cache, scale, block_tables, seq_lens) + + +# merge attn states ops +def merge_attn_states( + output: torch.Tensor, + prefix_output: torch.Tensor, + prefix_lse: torch.Tensor, + suffix_output: torch.Tensor, + suffix_lse: torch.Tensor, + output_lse: torch.Tensor | None = None, + prefill_tokens_with_context: int | None = None, + output_scale: torch.Tensor | None = None, +) -> None: + torch.ops._C.merge_attn_states( + output, + output_lse, + prefix_output, + prefix_lse, + suffix_output, + suffix_lse, + prefill_tokens_with_context, + output_scale, + ) + + +# pos encoding ops +def rotary_embedding( + positions: torch.Tensor, + query: torch.Tensor, + key: torch.Tensor | None, + head_size: int, + cos_sin_cache: torch.Tensor, + is_neox: bool, + rope_dim_offset: int = 0, + inverse: bool = False, +) -> None: + if rope_dim_offset == 0 and not inverse: + torch.ops._C.rotary_embedding(positions, query, key, head_size, cos_sin_cache, is_neox) + else: + torch.ops._C.rotary_embedding( + positions, + query, + key, + head_size, + cos_sin_cache, + is_neox, + rope_dim_offset, + inverse, + ) + + +# layer norm ops +def rms_norm( + out: torch.Tensor, + input: torch.Tensor, + weight: torch.Tensor | None, + epsilon: float, +) -> None: + torch.ops._C.rms_norm(out, input, weight, epsilon) + + +# LongCat n-gram embedding index kernel (see csrc/.../ngram_embedding_kernels.cu). +def ngram_compute_n_gram_ids( + ne_n: int, + ne_k: int, + ne_weights: torch.Tensor, + ne_mods: torch.Tensor, + exclusive_ne_embedder_size_sums: torch.Tensor, + exclusive_req_len_sums: torch.Tensor, + ne_token_table: torch.Tensor, + row_indices: torch.Tensor, + column_starts: torch.Tensor, + n_gram_ids: torch.Tensor, +) -> None: + """Compute concatenated (offset) n-gram ids for a ragged prefill batch. + + Writes ``n_gram_ids`` of shape ``[token_num, (ne_n-1)*ne_k]``. + """ + torch.ops._C.ngram_compute_n_gram_ids( + ne_n, + ne_k, + ne_weights, + ne_mods, + exclusive_ne_embedder_size_sums, + exclusive_req_len_sums, + ne_token_table, + row_indices, + column_starts, + n_gram_ids, + ) + + +def fused_add_rms_norm( + input: torch.Tensor, + residual: torch.Tensor, + weight: torch.Tensor | None, + epsilon: float, +) -> None: + # Note: this func is batch invariant + torch.ops._C.fused_add_rms_norm(input, residual, weight, epsilon) + + +def fused_qk_norm_rope( + qkv: torch.Tensor, + num_heads_q: int, + num_heads_k: int, + num_heads_v: int, + head_dim: int, + eps: float, + q_weight: torch.Tensor, + k_weight: torch.Tensor, + cos_sin_cache: torch.Tensor, + is_neox: bool, + position_ids: torch.Tensor, + forced_token_heads_per_warp: int = -1, +) -> None: + torch.ops._C.fused_qk_norm_rope( + qkv, + num_heads_q, + num_heads_k, + num_heads_v, + head_dim, + eps, + q_weight, + k_weight, + cos_sin_cache, + is_neox, + position_ids, + forced_token_heads_per_warp, + ) + + +def apply_repetition_penalties_torch( + logits: torch.Tensor, + prompt_mask: torch.Tensor, + output_mask: torch.Tensor, + repetition_penalties: torch.Tensor, +) -> None: + repetition_penalties = repetition_penalties.unsqueeze(dim=1).repeat(1, logits.size(1)) + # If token appears in prompt or output, apply, otherwise use 1.0 for no-op. + penalties = torch.where(prompt_mask | output_mask, repetition_penalties, 1.0) + # If logits are positive, divide by penalty, otherwise multiply by penalty. + scaling = torch.where(logits > 0, 1.0 / penalties, penalties) + logits *= scaling + + +def apply_repetition_penalties_cuda( + logits: torch.Tensor, + prompt_mask: torch.Tensor, + output_mask: torch.Tensor, + repetition_penalties: torch.Tensor, +) -> None: + torch.ops._C.apply_repetition_penalties_(logits, prompt_mask, output_mask, repetition_penalties) + + +def apply_repetition_penalties( + logits: torch.Tensor, + prompt_mask: torch.Tensor, + output_mask: torch.Tensor, + repetition_penalties: torch.Tensor, +) -> None: + """Apply repetition penalties to logits in-place. + + Args: + logits: The logits tensor of shape [num_seqs, vocab_size]. + prompt_mask: A boolean tensor indicating which tokens appear in the prompt. + output_mask: A boolean tensor indicating which tokens appear in the output. + repetition_penalties: The repetition penalties of shape (num_seqs, ). + """ + if logits.is_cuda and logits.is_contiguous(): + apply_repetition_penalties_cuda(logits, prompt_mask, output_mask, repetition_penalties) + else: + apply_repetition_penalties_torch(logits, prompt_mask, output_mask, repetition_penalties) + + +# fused quant layer norm ops +def rms_norm_dynamic_per_token_quant( + input: torch.Tensor, + weight: torch.Tensor, + epsilon: float, + quant_dtype: torch.dtype, + scale_ub: torch.Tensor | None = None, + residual: torch.Tensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor]: + output = torch.empty(input.shape, dtype=quant_dtype, device=input.device) + scales = torch.empty((input.numel() // input.shape[-1], 1), device=input.device, dtype=torch.float32) + + torch.ops._C.rms_norm_dynamic_per_token_quant(output, input, weight, scales, epsilon, scale_ub, residual) + return output, scales + + +# fused quant layer norm ops blocked +def rms_norm_per_block_quant( + input: torch.Tensor, + weight: torch.Tensor, + epsilon: float, + quant_dtype: torch.dtype, + group_size: list[int], + scale_ub: torch.Tensor | None = None, + residual: torch.Tensor | None = None, + is_scale_transposed: bool = False, + tma_alignment: int = 0, +) -> tuple[torch.Tensor, torch.Tensor]: + assert len(group_size) == 2 + output = torch.empty(input.shape, dtype=quant_dtype, device=input.device) + if is_scale_transposed: + if tma_alignment == 0: + scales = torch.empty( + (input.shape[-1] // group_size[1], input.numel() // input.shape[-1]), + device=input.device, + dtype=torch.float32, + ).transpose(0, 1) + else: + m = input.shape[-2] + sf_k = input.shape[-1] // group_size[1] + tma_aligned_m = (m + tma_alignment - 1) // tma_alignment * tma_alignment + shape = input.shape[:-2] + (m, sf_k) + stride = (1, tma_aligned_m) if input.dim() == 2 else (tma_aligned_m * sf_k, 1, tma_aligned_m) + scales = torch.empty_strided(shape, stride, device=input.device, dtype=torch.float32) + else: + scales = torch.empty( + (input.numel() // input.shape[-1], input.shape[-1] // group_size[1]), + device=input.device, + dtype=torch.float32, + ) + + assert tma_alignment in [0, 4], "Expected TMA alignment 0 or 4, but got " + str(tma_alignment) + + torch.ops._C.rms_norm_per_block_quant( + output, + input, + weight, + scales, + epsilon, + scale_ub, + residual, + group_size[1], + is_scale_transposed, + ) + return output, scales + + +# fused silu_and_mul + block quant +def silu_and_mul_per_block_quant( + input: torch.Tensor, + group_size: int, # Changed from list[int] + quant_dtype: torch.dtype, + scale_ub: torch.Tensor | None = None, + is_scale_transposed: bool = False, +) -> tuple[torch.Tensor, torch.Tensor]: + assert input.ndim == 2, f"input must be 2D [batch, hidden*2], got {input.shape}" + assert input.shape[-1] % 2 == 0, f"input last dim must be even (gate||up layout), got {input.shape[-1]}" + + # Output is half the width of input (after silu_and_mul) + num_tokens = input.shape[0] + hidden_size = input.shape[-1] // 2 # Divide by 2 because input is [gate || up] + + # Allocate output tensor (FP8 or INT8) + output = torch.empty((num_tokens, hidden_size), device=input.device, dtype=quant_dtype) + + # Allocate scales tensor + num_groups = hidden_size // group_size # Directly use group_size + if is_scale_transposed: + scales = torch.empty( + (num_groups, num_tokens), + device=input.device, + dtype=torch.float32, + ).t() + else: + scales = torch.empty( + (num_tokens, num_groups), + device=input.device, + dtype=torch.float32, + ) + + # Call the C++ kernel + torch.ops._C.silu_and_mul_per_block_quant( + output, + input, + scales, + group_size, # Pass directly as int + scale_ub, + is_scale_transposed, + ) + + return output, scales + + +def silu_mul(out: torch.Tensor, gate: torch.Tensor, up: torch.Tensor) -> None: + torch.ops._C.silu_mul(out, gate, up) + + +def make_gate_up_indices( + out: torch.Tensor, + indices: torch.Tensor, + offset: int, +) -> None: + torch.ops._C.make_gate_up_indices(out, indices, offset) + + +# quantization ops +# awq +def awq_dequantize( + qweight: torch.Tensor, + scales: torch.Tensor, + zeros: torch.Tensor, + split_k_iters: int, + thx: int, + thy: int, +) -> torch.Tensor: + if envs.APHRODITE_USE_TRITON_AWQ: + from aphrodite.model_executor.layers.quantization.awq_triton import ( + awq_dequantize_triton, + ) + + return awq_dequantize_triton(qweight, scales, zeros) + return torch.ops._C.awq_dequantize(qweight, scales, zeros, split_k_iters, thx, thy) + + +if hasattr(torch.ops._C, "awq_dequantize"): + + @register_fake("_C::awq_dequantize") + def _awq_dequantize_fake( + qweight: torch.Tensor, + scales: torch.Tensor, + zeros: torch.Tensor, + split_k_iters: torch.SymInt, + thx: int, + thy: int, + ) -> torch.Tensor: + in_c = qweight.size(0) + qout_c = qweight.size(1) + out_c = qout_c * 8 + return torch.empty((in_c, out_c), dtype=scales.dtype, device=scales.device) + + +def awq_gemm( + input: torch.Tensor, + qweight: torch.Tensor, + scales: torch.Tensor, + qzeros: torch.Tensor, + split_k_iters: int, +) -> torch.Tensor: + if envs.APHRODITE_USE_TRITON_AWQ: + from aphrodite.model_executor.layers.quantization.awq_triton import awq_gemm_triton + + return awq_gemm_triton(input, qweight, scales, qzeros, split_k_iters) + return torch.ops._C.awq_gemm(input, qweight, scales, qzeros, split_k_iters) + + +if hasattr(torch.ops._C, "awq_gemm"): + + @register_fake("_C::awq_gemm") + def _awq_gemm_fake( + input: torch.Tensor, + qweight: torch.Tensor, + scales: torch.Tensor, + qzeros: torch.Tensor, + split_k_iters: torch.SymInt, + ) -> torch.Tensor: + num_in_feats = input.size(0) + return torch.empty( + (split_k_iters, num_in_feats, qweight.size(1) * 8), + dtype=input.dtype, + device=input.device, + ).sum(0) + + +# gptq +def gptq_gemm( + a: torch.Tensor, + b_q_weight: torch.Tensor, + b_gptq_qzeros: torch.Tensor, + b_gptq_scales: torch.Tensor, + b_g_idx: torch.Tensor, + use_exllama: bool, + use_v2_format: bool, + bit: int, +) -> torch.Tensor: + return torch.ops._C.gptq_gemm( + a, + b_q_weight, + b_gptq_qzeros, + b_gptq_scales, + b_g_idx, + use_exllama, + use_v2_format, + bit, + ) + + +if hasattr(torch.ops._C, "gptq_gemm"): + + @register_fake("_C::gptq_gemm") + def _gptq_gemm_fake( + a: torch.Tensor, + b_q_weight: torch.Tensor, + b_gptq_qzeros: torch.Tensor, + b_gptq_scales: torch.Tensor, + b_g_idx: torch.Tensor, + use_exllama: bool, + use_v2_format: bool, + bit: int, + ) -> torch.Tensor: + return torch.empty((a.size(0), b_q_weight.size(1)), dtype=a.dtype, device=a.device) + + +def gptq_shuffle(q_weight: torch.Tensor, q_perm: torch.Tensor, bit: int) -> None: + torch.ops._C.gptq_shuffle(q_weight, q_perm, bit) + + +def gptq_gemm_rdna3( + a: torch.Tensor, + b_q_weight: torch.Tensor, + b_qzeros: torch.Tensor, + b_scales: torch.Tensor, + b_g_idx: torch.Tensor, + use_v2_format: bool, +) -> torch.Tensor: + return torch.ops._rocm_C.gptq_gemm_rdna3(a, b_q_weight, b_qzeros, b_scales, b_g_idx, use_v2_format) + + +if hasattr(torch.ops, "_rocm_C") and hasattr(torch.ops._rocm_C, "gptq_gemm_rdna3"): + + @register_fake("_rocm_C::gptq_gemm_rdna3") + def _gptq_gemm_rdna3_fake( + a: torch.Tensor, + b_q_weight: torch.Tensor, + b_qzeros: torch.Tensor, + b_scales: torch.Tensor, + b_g_idx: torch.Tensor, + use_v2_format: bool, + ) -> torch.Tensor: + return torch.empty((a.size(0), b_q_weight.size(1)), dtype=a.dtype, device=a.device) + + +if hasattr(torch.ops, "_rocm_C") and hasattr(torch.ops._rocm_C, "gptq_gemm_rdna3_wmma"): + + @register_fake("_rocm_C::gptq_gemm_rdna3_wmma") + def _gptq_gemm_rdna3_wmma_fake( + a: torch.Tensor, + b_q_weight: torch.Tensor, + b_qzeros: torch.Tensor, + b_scales: torch.Tensor, + b_g_idx: torch.Tensor, + use_v2_format: bool, + ) -> torch.Tensor: + return torch.empty((a.size(0), b_q_weight.size(1)), dtype=a.dtype, device=a.device) + + +def moe_gptq_gemm_rdna3( + a: torch.Tensor, + c: torch.Tensor, + b_q_weight: torch.Tensor, + b_scales: torch.Tensor, + b_qzeros: torch.Tensor, + topk_weights: torch.Tensor, + sorted_token_ids: torch.Tensor, + expert_ids: torch.Tensor, + num_tokens_post_padded: torch.Tensor, + top_k: int, + block_size_m: int, + mul_topk_weight: bool, + output_topk: int = 0, +) -> None: + torch.ops._rocm_C.moe_gptq_gemm_rdna3( + a, + c, + b_q_weight, + b_scales, + b_qzeros, + topk_weights, + sorted_token_ids, + expert_ids, + num_tokens_post_padded, + top_k, + block_size_m, + mul_topk_weight, + output_topk, + ) + + +if hasattr(torch.ops, "_rocm_C") and hasattr(torch.ops._rocm_C, "moe_gptq_gemm_rdna3"): + + @register_fake("_rocm_C::moe_gptq_gemm_rdna3") + def _moe_gptq_gemm_rdna3_fake( + a: torch.Tensor, + c: torch.Tensor, + b_q_weight: torch.Tensor, + b_scales: torch.Tensor, + b_qzeros: torch.Tensor, + topk_weights: torch.Tensor, + sorted_token_ids: torch.Tensor, + expert_ids: torch.Tensor, + num_tokens_post_padded: torch.Tensor, + top_k: int, + block_size_m: int, + mul_topk_weight: bool, + output_topk: int = 0, + ) -> None: + return + + +if hasattr(torch.ops._C, "allspark_w8a16_gemm"): + + @register_fake("_C::allspark_w8a16_gemm") + def _allspark_w8a16_gemm_fake( + a: torch.Tensor, + b_qweight: torch.Tensor, + b_scales: torch.Tensor, + b_qzeros: torch.Tensor | None, + n: torch.SymInt, + group_size: torch.SymInt, + sm_count: torch.SymInt, + sm_version: torch.SymInt, + CUBLAS_M_THRESHOLD: torch.SymInt, + has_zp: bool, + n32k16_reorder: bool, + ) -> torch.Tensor: + m = a.size(0) + return torch.empty((m, n), device=a.device, dtype=a.dtype) + + +# cutlass +def cutlass_scaled_mm_supports_fp4(cuda_device_capability: int) -> bool: + return torch.ops._C.cutlass_scaled_mm_supports_fp4(cuda_device_capability) + + +def mxfp4_experts_quant_supported(cuda_device_capability: int) -> bool: + try: + return torch.ops._C.mxfp4_experts_quant_supported(cuda_device_capability) + except AttributeError: + # Return False on builds where the CUDA helper is not available. + return False + + +def cutlass_scaled_fp4_mm( + a: torch.Tensor, + b: torch.Tensor, + block_scale_a: torch.Tensor, + block_scale_b: torch.Tensor, + alpha: torch.Tensor, + out_dtype: torch.dtype, +) -> torch.Tensor: + assert a.ndim == 2 and b.ndim == 2 + m, n = a.shape[0], b.shape[0] + out = torch.empty((m, n), dtype=out_dtype, device=a.device) + torch.ops._C.cutlass_scaled_fp4_mm(out, a, b, block_scale_a, block_scale_b, alpha) + return out + + +def cutlass_scaled_mm_supports_fp8(cuda_device_capability: int) -> bool: + return torch.ops._C.cutlass_scaled_mm_supports_fp8(cuda_device_capability) + + +def cutlass_scaled_mm_supports_block_fp8(cuda_device_capability: int) -> bool: + return torch.ops._C.cutlass_scaled_mm_supports_block_fp8(cuda_device_capability) + + +def cutlass_scaled_mm( + a: torch.Tensor, + b: torch.Tensor, + scale_a: torch.Tensor, + scale_b: torch.Tensor, + out_dtype: torch.dtype, + bias: torch.Tensor | None = None, +) -> torch.Tensor: + """ + `cutlass_scaled_mm` implements a fused version of + `output = torch.mm((scale_a * a), (scale_b * b)).to(out_dtype)` + where scale_a * a and scale_b * b are implemented using numpy-style + broadcasting. + + In order to support blockwise scaling like found in DeepSeek V3 we also + support extended "group" broadcast rules. We extend the numpy-style + broadcasting rules with the following rule: + "if the extent of a dimension in the source shape is between 1 and + corresponding extent in the target shape we repeat each element along + that dimension src_shape[dim] // target_shape[dim] times consecutively" + example if we have: + a = [[1, 2], and target_shape = (2, 4) + [3, 4]] + then we would expand a to: + a = [[1, 1, 2, 2], + [3, 3, 4, 4]] + currently we only support the case: + scale_a.shape * [1, 128] == a.shape + scale_b.shape * [128, 128] == b.shape + """ + assert out_dtype is torch.bfloat16 or out_dtype is torch.float16 + assert bias is None or bias.numel() == b.shape[1] and bias.dtype == out_dtype + + # Massage the input to be 2D + target_shape = (*a.shape[:-1], b.shape[1]) + a = a.view(-1, a.shape[-1]) + + cutlass_compatible_b = b.shape[0] % 16 == 0 and b.shape[1] % 16 == 0 + if current_platform.is_rocm() or not cutlass_compatible_b: + from aphrodite.model_executor.layers.quantization.compressed_tensors.triton_scaled_mm import ( # noqa + triton_scaled_mm, + ) + + out = triton_scaled_mm(a, b, scale_a, scale_b, out_dtype, bias) + else: + out = torch.empty((a.shape[0], b.shape[1]), dtype=out_dtype, device=a.device) + torch.ops._C.cutlass_scaled_mm(out, a, b, scale_a, scale_b, bias) + + return out.view(*target_shape) + + +def cutlass_scaled_mm_azp( + a: torch.Tensor, + b: torch.Tensor, + scale_a: torch.Tensor, + scale_b: torch.Tensor, + out_dtype: torch.dtype, + azp_adj: torch.Tensor, + azp: torch.Tensor | None = None, + bias: torch.Tensor | None = None, +) -> torch.Tensor: + """ + Args: + azp_adj: In the per-tensor case, this should include the azp. + Always per-channel. + azp: Only set in the per-token case. Per-token if set. + """ + assert b.shape[0] % 16 == 0 and b.shape[1] % 16 == 0 + assert out_dtype is torch.bfloat16 or out_dtype is torch.float16 + assert bias is None or bias.numel() == b.shape[1] and bias.dtype == out_dtype + + # Massage the input to be 2D + target_shape = (*a.shape[:-1], b.shape[1]) + a = a.view(-1, a.shape[-1]) + assert azp is None or azp.numel() == a.shape[0] + + out = torch.empty((a.shape[0], b.shape[1]), dtype=out_dtype, device=a.device) + torch.ops._C.cutlass_scaled_mm_azp(out, a, b, scale_a, scale_b, azp_adj, azp, bias) + return out.view(*target_shape) + + +def cutlass_group_gemm_supported(cuda_device_capability: int) -> bool: + if cuda_device_capability < 90 or cuda_device_capability >= 110: + return False + try: + return torch.ops._C.cutlass_group_gemm_supported(cuda_device_capability) + except AttributeError: + # Return False on non-CUDA platforms where it is not available + return False + + +def get_cutlass_moe_mm_data( + topk_ids: torch.Tensor, + expert_offsets: torch.Tensor, + problem_sizes1: torch.Tensor, + problem_sizes2: torch.Tensor, + input_permutation: torch.Tensor, + output_permutation: torch.Tensor, + num_experts: int, + n: int, + k: int, + blockscale_offsets: torch.Tensor | None = None, + is_gated: bool = True, +): + """ + Prepare data necessary to perform CUTLASS grouped matrix multiplications + used in CUTLASS-based fused MoE. + + The function takes in topk_ids (token-expert mapping) and uses it to + compute: + - expert_offsets: Indices that mark at which token index each expert begins + its computation after the input is sorted with + input_permutation. The number of tokens computed with + expert E is expert_offsets[E + 1] - expert_offsets[E] + - problem_sizes1, problem_sizes2: MxNxK sizes of each expert's + multiplication in two grouped MMs used in + the fused MoE operation. + - input_permutation: Permutation that must be used to shuffle the input + before executing the MMs. + - output_permutation: Permutation that must be used to shuffle the output + after executing the MMs. + - blockscale_offsets: Optional argument passed for fp4 moe. Indices that + mark at which block scale index each expert begins + its computation. The number of block scale rows + computed with expert E is blockscale_offsets[E + 1] - + blockscale_offsets[E] + - is_gated: Whether the activation is gated (gate + up). When True, the + first GEMM N dimension is 2*n; when False, it is n. + """ + return torch.ops._C.get_cutlass_moe_mm_data( + topk_ids, + expert_offsets, + problem_sizes1, + problem_sizes2, + input_permutation, + output_permutation, + num_experts, + n, + k, + blockscale_offsets, + is_gated, + ) + + +def get_cutlass_moe_mm_problem_sizes_from_expert_offsets( + expert_first_token_offset: torch.Tensor, + problem_sizes1: torch.Tensor, + problem_sizes2: torch.Tensor, + n: int, + k: int, + swap_ab: bool, +): + """Compute per-expert (M, N, K) problem sizes from expert_first_token_offset""" + return torch.ops._C.get_cutlass_moe_mm_problem_sizes_from_expert_offsets( + expert_first_token_offset, + problem_sizes1, + problem_sizes2, + n, + k, + swap_ab, + ) + + +def shuffle_rows(input_tensor: torch.Tensor, dst2src_map: torch.Tensor): + """ + Shuffle and expand the input tensor according to the dst2src_map and store the result in output_tensor. + This is used in MoE to permute the input tensor before performing grouped matrix multiplications. + """ + num_tokens_permuted = dst2src_map.shape[0] + output_tensor = torch.empty( + (num_tokens_permuted, input_tensor.shape[1]), + device=input_tensor.device, + dtype=input_tensor.dtype, + ) + torch.ops._moe_C.shuffle_rows(input_tensor, dst2src_map, output_tensor) + return output_tensor + + +def get_cutlass_batched_moe_mm_data( + expert_offsets: torch.Tensor, + problem_sizes1: torch.Tensor, + problem_sizes2: torch.Tensor, + expert_num_tokens: torch.Tensor, + num_local_experts: int, + padded_m: int, + n: int, + k: int, +): + """ + Prepare data necessary to perform CUTLASS grouped matrix multiplications + used in CUTLASS-based fused MoE. + + The function takes in expert_num_tokens (token count per expert) and + non_zero_expert_idxs (consecutive indices of experts with non-zero token + counts) and uses them to compute: + - expert_offsets: Indices that mark at which token index each expert begins + its computation. + - problem_sizes1, problem_sizes2: MxNxK sizes of each expert's + multiplication in two grouped MMs used in + the fused MoE operation. + """ + return torch.ops._C.get_cutlass_batched_moe_mm_data( + expert_offsets, + problem_sizes1, + problem_sizes2, + expert_num_tokens, + num_local_experts, + padded_m, + n, + k, + ) + + +def cutlass_moe_mm( + out_tensors: torch.Tensor, + a_tensors: torch.Tensor, + b_tensors: torch.Tensor, + a_scales: torch.Tensor, + b_scales: torch.Tensor, + expert_offsets: torch.Tensor, + problem_sizes: torch.Tensor, + a_strides: torch.Tensor, + b_strides: torch.Tensor, + c_strides: torch.Tensor, + per_act_token: bool, + per_out_ch: bool, +): + """ + A single grouped matrix multiplication used in CUTLASS-based fused MoE. + The function executes fp8-quantized OUT = AB matrix multiplication. + + - expert_offsets: Indices that mark at which token index each expert begins + its computation. The number of tokens computed with + expert E is expert_offsets[E + 1] - expert_offsets[E] + - problem_sizes: MxNxK sizes of each expert's multiplication in two grouped + MMs used in the fused MoE operation. + - a/b/c_strides: The data strides passed to grouped matrix multiplication. + """ + return torch.ops._C.cutlass_moe_mm( + out_tensors, + a_tensors, + b_tensors, + a_scales, + b_scales, + expert_offsets, + problem_sizes, + a_strides, + b_strides, + c_strides, + per_act_token, + per_out_ch, + ) + + +def cutlass_fp4_moe_mm( + out_tensors: torch.Tensor, + a_tensors: torch.Tensor, + b_tensors: torch.Tensor, + a_scales: torch.Tensor, + b_scales: torch.Tensor, + alphas: torch.Tensor, + problem_sizes: torch.Tensor, + expert_offsets: torch.Tensor, + sf_offsets: torch.Tensor, +): + """ + An FP4 Blockscaled Group Gemm that takes in a_tensors, b_tensors and runs + the gemms for each combination based on the specified problem sizes. + + This is used as the MoE gemm during NVFP4 Quantized FusedMoE forward. + - a/b_tensors: the NVFP4 a_ptrs and b_ptrs tensors which are quantized + input and expert weights. + - a_/b_scales: The blockscales in FP8-E4M3 precision + - expert_offsets/sf_offsets: Indices that mark at which token index + each expert begins its computation. The number of tokens + computed with expert E is expert_offsets[E + 1] - + expert_offsets[E] And the sf_size per expert is + sf_offset[E+1] - sf_offset[E] + - problem_sizes: MxNxK sizes of each expert's multiplication in two grouped + MMs used in the fused MoE operation. + """ + return torch.ops._C.cutlass_fp4_group_mm( + out_tensors, + a_tensors, + b_tensors, + a_scales, + b_scales, + alphas, + problem_sizes, + expert_offsets, + sf_offsets, + ) + + +def cutlass_mxfp4_moe_mm( + out_tensors: torch.Tensor, + a_tensors: torch.Tensor, + b_tensors: torch.Tensor, + a_scales: torch.Tensor, + b_scales: torch.Tensor, + problem_sizes: torch.Tensor, + expert_offsets: torch.Tensor, + sf_offsets: torch.Tensor, +): + """ + An MXFP4 Blockscaled Group Gemm for MoE (MXFP4 x MXFP4). + + Uses mx_float4_t types with E8M0 scale factors and 32-element blocks. + - a/b_tensors: MXFP4 packed activations/weights (uint8, 2 E2M1 per byte) + - a_/b_scales: E8M0 blockscales (uint8, stored in swizzled layout) + - Epilogue uses scalar alpha=1, beta=0 inside the CUDA op (no global scales). + - expert_offsets/sf_offsets: expert boundary indices + - problem_sizes: (num_experts, 3) with (M, N, K) per expert + """ + return torch.ops._C.cutlass_mxfp4_group_mm( + out_tensors, + a_tensors, + b_tensors, + a_scales, + b_scales, + problem_sizes, + expert_offsets, + sf_offsets, + ) + + +# gptq_marlin +def gptq_marlin_repack( + b_q_weight: torch.Tensor, + perm: torch.Tensor, + size_k: int, + size_n: int, + num_bits: int, + is_a_8bit: bool = False, +) -> torch.Tensor: + return torch.ops._C.gptq_marlin_repack(b_q_weight, perm, size_k, size_n, num_bits, is_a_8bit) + + +if hasattr(torch.ops._C, "gptq_marlin_repack"): + + @register_fake("_C::gptq_marlin_repack") + def _gptq_marlin_repack_fake( + b_q_weight: torch.Tensor, + perm: torch.Tensor, + size_k: torch.SymInt, + size_n: torch.SymInt, + num_bits: int, + is_a_8bit: bool = False, + ) -> torch.Tensor: + pack_factor = 32 // num_bits + marlin_tile_size = 16 + return torch.empty( + (size_k // marlin_tile_size, size_n * marlin_tile_size // pack_factor), + dtype=b_q_weight.dtype, + device=b_q_weight.device, + ) + + +# awq_marlin +def awq_marlin_repack( + b_q_weight: torch.Tensor, + size_k: int, + size_n: int, + num_bits: int, + is_a_8bit: bool = False, +) -> torch.Tensor: + return torch.ops._C.awq_marlin_repack(b_q_weight, size_k, size_n, num_bits, is_a_8bit) + + +if hasattr(torch.ops._C, "awq_marlin_repack"): + + @register_fake("_C::awq_marlin_repack") + def _awq_marlin_repack_fake( + b_q_weight: torch.Tensor, + size_k: torch.SymInt, + size_n: torch.SymInt, + num_bits: int, + is_a_8bit: bool = False, + ) -> torch.Tensor: + pack_factor = 32 // num_bits + marlin_tile_size = 16 + return torch.empty( + (size_k // marlin_tile_size, size_n * marlin_tile_size // pack_factor), + dtype=b_q_weight.dtype, + device=b_q_weight.device, + ) + + +def gptq_marlin_moe_repack( + b_q_weight: torch.Tensor, + perm: torch.Tensor, + size_k: int, + size_n: int, + num_bits: int, + is_a_8bit: bool = False, +) -> torch.Tensor: + num_experts = b_q_weight.shape[0] + assert size_k % 16 == 0 + output = torch.empty( + (num_experts, size_k // 16, size_n * (num_bits // 2)), + device=b_q_weight.device, + dtype=b_q_weight.dtype, + ) + for e in range(num_experts): + output[e] = torch.ops._C.gptq_marlin_repack(b_q_weight[e], perm[e], size_k, size_n, num_bits, is_a_8bit) + return output + + +def awq_marlin_moe_repack( + b_q_weight: torch.Tensor, + perm: torch.Tensor, + size_k: int, + size_n: int, + num_bits: int, + is_a_8bit: bool = False, +) -> torch.Tensor: + num_experts = b_q_weight.shape[0] + assert size_k % 16 == 0 + output = torch.empty( + (num_experts, size_k // 16, size_n * (num_bits // 2)), + device=b_q_weight.device, + dtype=b_q_weight.dtype, + ) + for e in range(num_experts): + output[e] = torch.ops._C.awq_marlin_repack(b_q_weight[e], size_k, size_n, num_bits, is_a_8bit) + return output + + +def marlin_int4_fp8_preprocess( + qweight: torch.Tensor, + qzeros_or_none: torch.Tensor | None = None, + inplace: bool = False, +): + return torch.ops._C.marlin_int4_fp8_preprocess(qweight, qzeros_or_none, inplace) + + +def marlin_gemm( + a: torch.Tensor, + c: torch.Tensor | None, + b_q_weight: torch.Tensor, + b_bias: torch.Tensor | None, + b_scales: torch.Tensor, + a_scales: torch.Tensor | None, + global_scale: torch.Tensor | None, + b_zeros: torch.Tensor | None, + g_idx: torch.Tensor | None, + perm: torch.Tensor | None, + workspace: torch.Tensor, + b_q_type: ScalarType, + size_m: int, + size_n: int, + size_k: int, + is_k_full: bool = True, + use_atomic_add: bool = False, + use_fp32_reduce: bool = False, + is_zp_float: bool = False, +) -> torch.Tensor: + return torch.ops._C.marlin_gemm( + a, + c, + b_q_weight, + b_bias, + b_scales, + a_scales, + global_scale, + b_zeros, + g_idx, + perm, + workspace, + b_q_type.id, + size_m, + size_n, + size_k, + is_k_full, + use_atomic_add, + use_fp32_reduce, + is_zp_float, + ) + + +if hasattr(torch.ops._C, "marlin_gemm"): + + @register_fake("_C::marlin_gemm") + def _marlin_gemm_fake( + a: torch.Tensor, + c: torch.Tensor | None, + b_q_weight: torch.Tensor, + b_bias: torch.Tensor | None, + b_scales: torch.Tensor, + a_scales: torch.Tensor | None, + global_scale: torch.Tensor | None, + b_zeros: torch.Tensor | None, + g_idx: torch.Tensor | None, + perm: torch.Tensor | None, + workspace: torch.Tensor, + b_q_type_id: int, + size_m: torch.SymInt, + size_n: torch.SymInt, + size_k: torch.SymInt, + is_k_full: bool = True, + use_atomic_add: bool = False, + use_fp32_reduce: bool = False, + is_zp_float: bool = False, + ) -> torch.Tensor: + dtype = a.dtype + if dtype not in [torch.half, torch.bfloat16]: + dtype = b_scales.dtype + return torch.empty((size_m, size_n), device=a.device, dtype=dtype) + + +# machete +def machete_supported_schedules( + a_type: torch.dtype, + b_type: ScalarType, + group_scales_type: torch.dtype | None, + group_zeros_type: torch.dtype | None = None, + channel_scales_type: torch.dtype | None = None, + token_scales_type: torch.dtype | None = None, + out_type: torch.dtype | None = None, +) -> list[str]: + return torch.ops._C.machete_supported_schedules( + a_type, + b_type.id, + group_scales_type, + group_zeros_type, + channel_scales_type, + token_scales_type, + out_type, + ) + + +def machete_mm( + a: torch.Tensor, + # b_q Should be the tensor returned by machete_prepack_B + b_q: torch.Tensor, + b_type: ScalarType, + out_type: torch.dtype | None = None, + b_group_scales: torch.Tensor | None = None, + b_group_zeros: torch.Tensor | None = None, + b_group_size: int | None = None, + b_channel_scales: torch.Tensor | None = None, + a_token_scales: torch.Tensor | None = None, + schedule: str | None = None, +) -> torch.Tensor: + return torch.ops._C.machete_mm( + a, + b_q, + b_type.id, + out_type, + b_group_scales, + b_group_zeros, + b_group_size, + b_channel_scales, + a_token_scales, + schedule, + ) + + +if hasattr(torch.ops._C, "machete_mm"): + + @register_fake("_C::machete_mm") + def machete_mm_fake( + a: torch.Tensor, + # b_q Should be the tensor returned by machete_prepack_B + b_q: torch.Tensor, + b_type: ScalarType, + out_type: torch.dtype | None = None, + b_group_scales: torch.Tensor | None = None, + b_group_zeros: torch.Tensor | None = None, + b_group_size: int | None = None, + b_channel_scales: torch.Tensor | None = None, + a_token_scales: torch.Tensor | None = None, + schedule: str | None = None, + ) -> torch.Tensor: + m = a.size(0) + n = b_q.size(1) + return torch.empty((m, n), device=a.device, dtype=a.dtype) + + +def machete_prepack_B( + b_q_weight: torch.Tensor, + a_type: torch.dtype, + b_type: ScalarType, + group_scales_type: torch.dtype | None, +) -> torch.Tensor: + return torch.ops._C.machete_prepack_B(b_q_weight, a_type, b_type.id, group_scales_type) + + +if hasattr(torch.ops._C, "machete_prepack_B"): + + @register_fake("_C::machete_prepack_B") + def machete_prepack_B_fake( + b_q_weight: torch.Tensor, + a_type: torch.dtype, + b_type: ScalarType, + group_scales_type: torch.dtype | None, + ) -> torch.Tensor: + return torch.empty_like(b_q_weight, memory_format=torch.contiguous_format) + + +# Swordfish (Blackwell sm100/sm110 w4a16) +def swordfish_prepack_B( + b_q_weight: torch.Tensor, + size_k: int, + size_n: int, + num_bits: int = 4, + perm: torch.Tensor | None = None, +) -> torch.Tensor: + """Pack a GPTQ int weight (int32 [K*bits/32, N]) into Swordfish ABI v1 + (int32 [NB, KB, 512*bits/4] block-linear). perm applies the act_order + row sort during the repack.""" + return torch.ops._C.swordfish_prepack_B(b_q_weight, perm, size_k, size_n, num_bits) + + +if hasattr(torch.ops._C, "swordfish_prepack_B"): + + @register_fake("_C::swordfish_prepack_B") + def swordfish_prepack_B_fake( + b_q_weight: torch.Tensor, + perm: torch.Tensor | None, + size_k: int, + size_n: int, + num_bits: int, + ) -> torch.Tensor: + return torch.empty( + (size_n // 64, size_k // 64, 128 * num_bits), + dtype=torch.int32, + device=b_q_weight.device, + ) + + +def swordfish_mm( + a: torch.Tensor, + b_packed: torch.Tensor, + group_scales: torch.Tensor, + group_size: int, + size_k: int, + size_n: int, + group_zps: torch.Tensor | None = None, + num_bits: int = 4, + perm: torch.Tensor | None = None, +) -> torch.Tensor: + """w4a16/w8a16 GEMM: a [M, K] fp16/bf16 times a Swordfish ABI v1 packed + weight with per-group scales [groups, N]. group_zps holds prescaled + (8 - zp) * scale rows for zero-point checkpoints (AWQ/HQQ, 4-bit only). + perm is the act_order column sort; the op permutes the activations for + the fused paths and folds the sort into the dense tier's weight scatter.""" + return torch.ops._C.swordfish_mm(a, b_packed, group_scales, group_zps, perm, num_bits, group_size, size_k, size_n) + + +if hasattr(torch.ops._C, "swordfish_mm"): + + @register_fake("_C::swordfish_mm") + def swordfish_mm_fake( + a: torch.Tensor, + b_packed: torch.Tensor, + group_scales: torch.Tensor, + group_zps: torch.Tensor | None, + perm: torch.Tensor | None, + num_bits: int, + group_size: int, + size_k: int, + size_n: int, + ) -> torch.Tensor: + return torch.empty((a.shape[0], size_n), dtype=a.dtype, device=a.device) + + +def swordfish_dequant_dense( + b_packed: torch.Tensor, + group_scales: torch.Tensor, + group_zps: torch.Tensor | None, + num_bits: int, + group_size: int, + size_k: int, + size_n: int, + transpose: bool, +) -> torch.Tensor: + """Dequantize Swordfish-packed weights to dense fp16/bf16, [K, N] or + out-major [N, K], with an optional leading expert dimension.""" + return torch.ops._C.swordfish_dequant_dense( + b_packed, + group_scales, + group_zps, + num_bits, + group_size, + size_k, + size_n, + transpose, + ) + + +if hasattr(torch.ops._C, "swordfish_dequant_dense"): + + @register_fake("_C::swordfish_dequant_dense") + def swordfish_dequant_dense_fake( + b_packed: torch.Tensor, + group_scales: torch.Tensor, + group_zps: torch.Tensor | None, + num_bits: int, + group_size: int, + size_k: int, + size_n: int, + transpose: bool, + ) -> torch.Tensor: + shape: tuple[int, ...] = (size_n, size_k) if transpose else (size_k, size_n) + if b_packed.dim() == 4: + shape = (b_packed.shape[0],) + shape + return torch.empty(shape, dtype=group_scales.dtype, device=b_packed.device) + + +def swordfish_moe_mm( + a: torch.Tensor, + b_packed: torch.Tensor, + group_scales: torch.Tensor, + sorted_token_ids: torch.Tensor, + expert_ids: torch.Tensor, + num_tokens_post_padded: torch.Tensor, + topk_weights: torch.Tensor | None, + moe_block_size: int, + top_k: int, + mul_topk_weights: bool, + num_bits: int, + group_size: int, + size_k: int, + size_n: int, +) -> torch.Tensor: + """Fused-MoE GEMM over per-expert Swordfish ABI v1 weights + (int32 [E, NB, KB, words]) with token blocks prepared by + moe_align_block_size(moe_block_size in {16, 64}). The 64-token blocks + run the CTA-wide weight-staging kernel for batched shapes. Returns + [M * top_k, size_n].""" + return torch.ops._C.swordfish_moe_mm( + a, + b_packed, + group_scales, + sorted_token_ids, + expert_ids, + num_tokens_post_padded, + topk_weights, + moe_block_size, + top_k, + mul_topk_weights, + num_bits, + group_size, + size_k, + size_n, + ) + + +if hasattr(torch.ops._C, "swordfish_moe_mm"): + + @register_fake("_C::swordfish_moe_mm") + def swordfish_moe_mm_fake( + a: torch.Tensor, + b_packed: torch.Tensor, + group_scales: torch.Tensor, + sorted_token_ids: torch.Tensor, + expert_ids: torch.Tensor, + num_tokens_post_padded: torch.Tensor, + topk_weights: torch.Tensor | None, + moe_block_size: int, + top_k: int, + mul_topk_weights: bool, + num_bits: int, + group_size: int, + size_k: int, + size_n: int, + ) -> torch.Tensor: + return torch.empty((a.shape[0] * top_k, size_n), dtype=a.dtype, device=a.device) + + +def swordfish_prefill_mm( + a: torch.Tensor, + b_packed: torch.Tensor, + group_scales: torch.Tensor, + group_size: int, + size_k: int, + size_n: int, + group_zps: torch.Tensor | None = None, + num_bits: int = 4, +) -> torch.Tensor: + """w4a16 prefill GEMM (sm100 tcgen05 mixed-input mainloop fork): a [M, K] + bf16 times a Swordfish ABI v1 packed weight (int32 [NB, KB, 512]) with + bf16 per-group scales [groups, N]. v1: group_size 128, K/N % 128 == 0.""" + return torch.ops._C.swordfish_prefill_mm(a, b_packed, group_scales, group_zps, num_bits, group_size, size_k, size_n) + + +if hasattr(torch.ops._C, "swordfish_prefill_mm"): + + @register_fake("_C::swordfish_prefill_mm") + def swordfish_prefill_mm_fake( + a: torch.Tensor, + b_packed: torch.Tensor, + group_scales: torch.Tensor, + group_zps: torch.Tensor | None, + num_bits: int, + group_size: int, + size_k: int, + size_n: int, + ) -> torch.Tensor: + return torch.empty((a.shape[0], size_n), dtype=a.dtype, device=a.device) + + +# CUTLASS W4A8 +def cutlass_w4a8_mm( + a: torch.Tensor, + # b_q Should be the tensor returned by cutlass_encode_and_reorder_int4b + b_q: torch.Tensor, + b_group_scales: torch.Tensor, + b_group_size: int, + b_channel_scales: torch.Tensor, + a_token_scales: torch.Tensor, + out_type: torch.dtype | None = None, + maybe_schedule: str | None = None, +) -> torch.Tensor: + return torch.ops._C.cutlass_w4a8_mm( + a, + b_q, + b_group_scales, + b_group_size, + b_channel_scales, + a_token_scales, + out_type, + maybe_schedule, + ) + + +if hasattr(torch.ops._C, "cutlass_w4a8_mm"): + + @register_fake("_C::cutlass_w4a8_mm") + def cutlass_w4a8_mm_fake( + a: torch.Tensor, + # b_q Should be the tensor returned by cutlass_encode_and_reorder_int4b + b_q: torch.Tensor, + b_group_scales: torch.Tensor, + b_group_size: int, + b_channel_scales: torch.Tensor, + a_token_scales: torch.Tensor, + out_type: torch.dtype | None = None, + maybe_schedule: str | None = None, + ) -> torch.Tensor: + m = a.size(0) + n = b_q.size(1) + out_dtype = out_type if out_type is not None else torch.bfloat16 + return torch.empty((m, n), device=a.device, dtype=out_dtype) + + +def cutlass_pack_scale_fp8(scales: torch.Tensor) -> torch.Tensor: + return torch.ops._C.cutlass_pack_scale_fp8(scales) + + +if hasattr(torch.ops._C, "cutlass_pack_scale_fp8"): + + @register_fake("_C::cutlass_pack_scale_fp8") + def cutlass_pack_scale_fp8_fake(scales: torch.Tensor) -> torch.Tensor: + return torch.empty_like(scales, memory_format=torch.contiguous_format) + + +def cutlass_encode_and_reorder_int4b(b: torch.Tensor) -> torch.Tensor: + return torch.ops._C.cutlass_encode_and_reorder_int4b(b) + + +if hasattr(torch.ops._C, "cutlass_encode_and_reorder_int4b"): + + @register_fake("_C::cutlass_encode_and_reorder_int4b") + def cutlass_encode_and_reorder_int4b_fake(b: torch.Tensor) -> torch.Tensor: + return torch.empty_like(b, memory_format=torch.contiguous_format) + + +def cutlass_w4a8_moe_mm( + out_tensors: torch.Tensor, + a_tensors: torch.Tensor, + b_tensors: torch.Tensor, + a_scales: torch.Tensor, + b_scales: torch.Tensor, + b_group_scales: torch.Tensor, + b_group_size: int, + expert_offsets: torch.Tensor, + problem_sizes: torch.Tensor, + a_strides: torch.Tensor, + b_strides: torch.Tensor, + c_strides: torch.Tensor, + group_scale_strides: torch.Tensor, + maybe_schedule: str | None = None, +): + """ + Executes the CUTLASS-based fused-MoE grouped matrix multiplication for the + W4A8 quantization scheme. Uses group-wise quantization (INT4 -> FP8) + and both per-channel + per-token scaling in the epilogue. + + Args: + out_tensors: + Output buffer for all experts (updated in-place). + a_tensors: + FP8 (E4M3FN) activations for all experts. + b_tensors: + INT4-packed weight matrix for all experts, packed to INT32 + a_scales: + Per-token FP8 activation scales, applied in the epilogue. + b_scales: + Per-channel FP8 weight scales for each expert, applied in the epilogue. + b_group_scales: + FP8 scale values for group-wise INT4 weight blocks. + b_group_size: + Number of elements grouped under each entry of b_group_scales. + expert_offsets: + Cumulative token offsets + problem_sizes: + Per-expert (M, N, K) GEMM sizes used by the grouped GEMM launcher. + a/b/c/group_scale_strides: + Strides describing the memory layout of the input tensors. + maybe_schedule: + Optional override to choose a specific kernel or epilogue schedule. + + Returns: + out_tensors updated in-place with the dequantized INT4xFP8 grouped GEMM result. + """ + return torch.ops._C.cutlass_w4a8_moe_mm( + out_tensors, + a_tensors, + b_tensors, + a_scales, + b_scales, + b_group_scales, + b_group_size, + expert_offsets, + problem_sizes, + a_strides, + b_strides, + c_strides, + group_scale_strides, + maybe_schedule, + ) + + +def cutlass_encode_and_reorder_int4b_grouped( + b_tensors: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + return torch.ops._C.cutlass_encode_and_reorder_int4b_grouped(b_tensors) + + +if hasattr(torch.ops._C, "cutlass_encode_and_reorder_int4b_grouped"): + + @register_fake("_C::cutlass_encode_and_reorder_int4b_grouped") + def cutlass_encode_and_reorder_int4b_grouped_fake(b: torch.Tensor) -> torch.Tensor: + return torch.empty_like(b, memory_format=torch.contiguous_format) + + +def permute_cols(a: torch.Tensor, perm: torch.Tensor) -> torch.Tensor: + return torch.ops._C.permute_cols(a, perm) + + +if hasattr(torch.ops._C, "permute_cols"): + + @register_fake("_C::permute_cols") + def _permute_cols_fake(a: torch.Tensor, perm: torch.Tensor) -> torch.Tensor: + return torch.empty_like(a) + + +# fp4 +def scaled_fp4_quant( + input: torch.Tensor, + input_global_scale: torch.Tensor, + is_sf_swizzled_layout: bool = True, + backend: str = "none", + padded_n: int | None = None, +) -> tuple[torch.Tensor, torch.Tensor]: + """ + Quantize input tensor to FP4 and return quantized tensor and scale. + + This function quantizes the last dimension of the given tensor `input`. For + every 16 consecutive elements, a single dynamically computed scaling factor + is shared. This scaling factor is quantized using the `input_global_scale` + and is stored in a swizzled layout (see + https://docs.nvidia.com/cuda/parallel-thread-execution/#tcgen05-mma-scale-factor-b-layout-4x). + + Args: + input: The input tensor to be quantized to FP4 + input_global_scale: A scalar scaling factor for the entire tensor. + is_sf_swizzled_layout: Whether to store the scaling factors in the + swizzled layout (default `True`). + backend: Quantization kernel backend to dispatch to. For `"trtllm"` + backends the 8x4 scale-factor layout is selected for small + batches (m <= 32) instead of the 128x4 layout. + padded_n: Optional padded K dimension. When provided, the quantized + output and scale tensors are allocated for ``padded_n`` + + Returns: + tuple[torch.Tensor, torch.Tensor]: The output tensor in FP4 but every + two values are packed into a uint8 and float8_e4m3 scaling factors + in the sizzled layout. + """ + assert not current_platform.is_rocm() + assert input.ndim >= 1, f"input.ndim needs to be >= 1, but got {input.ndim}." + other_dims = 1 if input.ndim == 1 else -1 + input = input.reshape(other_dims, input.shape[-1]) + m, n = input.shape + block_size = 16 + + assert n % block_size == 0, f"last dim has to be multiple of 16, but got {n}." + assert input.dtype in (torch.float16, torch.bfloat16), ( + f"input.dtype needs to be fp16 or bf16 but got {input.dtype}." + ) + if padded_n is not None: + assert padded_n >= n, f"padded_n must be >= n, got padded_n={padded_n}, n={n}." + assert padded_n % block_size == 0, f"padded_n has to be a multiple of {block_size}, but got {padded_n}." + + use_8x4_sf_layout = True if "trtllm" in backend and m <= 32 else False # noqa: SIM210 + if use_8x4_sf_layout and padded_n is not None and padded_n != n: + # TODO: support this case + raise ValueError("padded_n is not supported with TRTLLM 8x4 scale layout.") + if use_8x4_sf_layout: + output, output_scale = flashinfer_quant_nvfp4_8x4_sf_layout(input, input_global_scale) + else: + # Pre-allocate and call .out variant (same behavior as old in-place API) + output, output_scale = create_fp4_output_tensors( + m, + n, + input.device, + is_sf_swizzled_layout, + padded_n=padded_n, + ) + torch.ops._C.scaled_fp4_quant.out( + input, + input_global_scale, + is_sf_swizzled_layout, + output=output, + output_scale=output_scale, + ) + + output_scale = output_scale.view(torch.float8_e4m3fn) + return output, output_scale + + +def scaled_fp4_experts_quant( + input_tensor: torch.Tensor, + input_global_scale: torch.Tensor, + expert_offsets: torch.Tensor, + blockscale_offsets: torch.Tensor, + topk: int, +) -> tuple[torch.Tensor, torch.Tensor]: + """ + Quantize input tensor to NVFP4 and return quantized tensor and scale, for + packed MoE Inputs. + Args: + input_tensor: The input tensor to be quantized to NVFP4 + input_global_scale: A scalar scaling factor for the entire tensor. + expert_offsets: The expert offsets tensor + blockscale_offsets: The blockscale offsets tensor + Outputs: + output: The quantized tensor in NVFP4 + output_scales: The blockscale tensor in FP8-E4M3 + """ + assert not current_platform.is_rocm() + assert input_tensor.ndim == 2, f"input.ndim needs to be == 2, but got {input_tensor.ndim}." + + # Control the maximum number of tokens per expert supported by the + # NVFP4 MoE Expert Quantization. This is used to prevent the kernel + # from running out of memory. This value can also be increased to support + # larger models. + MAX_TOKENS_PER_EXPERT = envs.APHRODITE_MAX_TOKENS_PER_EXPERT_FP4_MOE + m_numtopk, k = input_tensor.shape + + assert m_numtopk <= MAX_TOKENS_PER_EXPERT * topk, ( + f"m_numtopk must be less than MAX_TOKENS_PER_EXPERT(" + f"{MAX_TOKENS_PER_EXPERT})" + f" for cutlass_moe_fp4, observed m_numtopk = {m_numtopk}. Use" + f" APHRODITE_MAX_TOKENS_PER_EXPERT_FP4_MOE to set this value." + ) + scales_k = k // 16 + padded_k = (scales_k + (4 - 1)) // 4 + + # output is uint8 and packed fp4 values + output = torch.empty(m_numtopk, k // 2, device=input_tensor.device, dtype=torch.uint8) + output_scales = torch.empty( + MAX_TOKENS_PER_EXPERT * topk, + padded_k, + dtype=torch.int32, + device=input_tensor.device, + ) + torch.ops._C.scaled_fp4_experts_quant( + output, + output_scales, + input_tensor, + input_global_scale, + expert_offsets, + blockscale_offsets, + ) + output_scales = output_scales.view(torch.float8_e4m3fn) + return output, output_scales + + +def silu_and_mul_scaled_fp4_experts_quant( + input_tensor: torch.Tensor, + input_global_scale: torch.Tensor, + expert_offsets: torch.Tensor, + blockscale_offsets: torch.Tensor, + topk: int, +) -> tuple[torch.Tensor, torch.Tensor]: + """ + Fused SiLU+Mul+NVFP4 quantization for MoE intermediate activations. + + Args: + input_tensor: The input tensor with gate || up layout [m_topk, k*2] + input_global_scale: A per-expert scaling factor [n_experts] + expert_offsets: The expert offsets tensor [n_experts+1] + blockscale_offsets: The blockscale offsets tensor [n_experts+1] + topk: Number of top-k experts selected + Outputs: + output: The quantized tensor in NVFP4 [m_topk, k/2] + output_scales: The blockscale tensor in FP8-E4M3 + """ + assert not current_platform.is_rocm() + assert input_tensor.ndim == 2, f"input.ndim needs to be == 2, but got {input_tensor.ndim}." + + # Control the maximum number of tokens per expert supported by the + # NVFP4 MoE Expert Quantization. This is used to prevent the kernel + # from running out of memory. This value can also be increased to support + # larger models. + MAX_TOKENS_PER_EXPERT = envs.APHRODITE_MAX_TOKENS_PER_EXPERT_FP4_MOE + m_numtopk, k_times_2 = input_tensor.shape + assert k_times_2 % 2 == 0, "input width must be even (gate || up layout)" + k = k_times_2 // 2 + + assert m_numtopk <= MAX_TOKENS_PER_EXPERT * topk, ( + f"m_numtopk must be less than MAX_TOKENS_PER_EXPERT(" + f"{MAX_TOKENS_PER_EXPERT})" + f" for cutlass_moe_fp4, observed m_numtopk = {m_numtopk}. Use" + f" APHRODITE_MAX_TOKENS_PER_EXPERT_FP4_MOE to set this value." + ) + scales_k = k // 16 + padded_k = (scales_k + (4 - 1)) // 4 + + # output is uint8 and packed fp4 values + output = torch.empty(m_numtopk, k // 2, device=input_tensor.device, dtype=torch.uint8) + output_scales = torch.empty( + MAX_TOKENS_PER_EXPERT * topk, + padded_k, + dtype=torch.int32, + device=input_tensor.device, + ) + torch.ops._C.silu_and_mul_scaled_fp4_experts_quant( + output, + output_scales, + input_tensor, + input_global_scale, + expert_offsets, + blockscale_offsets, + ) + output_scales = output_scales.view(torch.float8_e4m3fn) + return output, output_scales + + +def mxfp4_experts_quant( + input_tensor: torch.Tensor, + expert_offsets: torch.Tensor, + blockscale_offsets: torch.Tensor, + n_experts: int, + topk: int, +) -> tuple[torch.Tensor, torch.Tensor]: + """ + Quantize input tensor to MXFP4 for packed MoE inputs. + Uses 32-element blocks with E8M0 (power-of-two) scale factors. + MXFP4 has no global scale - only block-level E8M0 scale factors. + + Args: + input_tensor: [m_topk, k] BF16/FP16 activations + expert_offsets: [n_experts+1] token boundaries per expert + blockscale_offsets: [n_experts+1] SF row boundaries per expert + n_experts: number of experts + topk: number of top-k experts + Returns: + output: [m_topk, k//2] packed E2M1 values (uint8) + output_scales: E8M0 blockscales in swizzled layout (uint8 view) + """ + assert not current_platform.is_rocm() + assert input_tensor.ndim == 2 + + MAX_TOKENS_PER_EXPERT = envs.APHRODITE_MAX_TOKENS_PER_EXPERT_FP4_MOE + m_numtopk, k = input_tensor.shape + + assert m_numtopk <= MAX_TOKENS_PER_EXPERT * topk, ( + f"m_numtopk must be less than MAX_TOKENS_PER_EXPERT(" + f"{MAX_TOKENS_PER_EXPERT})" + f" for cutlass_moe_mxfp4, observed m_numtopk = {m_numtopk}. Use" + f" APHRODITE_MAX_TOKENS_PER_EXPERT_FP4_MOE to set this value." + ) + scales_k = k // 32 + padded_k = (scales_k + (4 - 1)) // 4 + + output = torch.empty(m_numtopk, k // 2, device=input_tensor.device, dtype=torch.uint8) + output_scales = torch.empty( + MAX_TOKENS_PER_EXPERT * topk, + padded_k, + dtype=torch.int32, + device=input_tensor.device, + ) + torch.ops._C.mxfp4_experts_quant( + output, + output_scales, + input_tensor, + expert_offsets, + blockscale_offsets, + n_experts, + ) + # E8M0 SFs are stored as uint8 + output_scales = output_scales.view(torch.uint8) + return output, output_scales + + +def silu_and_mul_mxfp4_experts_quant( + input_tensor: torch.Tensor, + expert_offsets: torch.Tensor, + blockscale_offsets: torch.Tensor, + n_experts: int, + topk: int, +) -> tuple[torch.Tensor, torch.Tensor]: + """ + Fused SiLU+Mul+MXFP4 quantization for MoE intermediate activations. + MXFP4 has no global scale - only block-level E8M0 scale factors. + """ + assert not current_platform.is_rocm() + assert input_tensor.ndim == 2 + + MAX_TOKENS_PER_EXPERT = envs.APHRODITE_MAX_TOKENS_PER_EXPERT_FP4_MOE + m_numtopk, k_times_2 = input_tensor.shape + assert k_times_2 % 2 == 0, "input width must be even (gate || up layout)" + k = k_times_2 // 2 + + assert m_numtopk <= MAX_TOKENS_PER_EXPERT * topk + scales_k = k // 32 + padded_k = (scales_k + (4 - 1)) // 4 + + output = torch.empty(m_numtopk, k // 2, device=input_tensor.device, dtype=torch.uint8) + output_scales = torch.empty( + MAX_TOKENS_PER_EXPERT * topk, + padded_k, + dtype=torch.int32, + device=input_tensor.device, + ) + torch.ops._C.silu_and_mul_mxfp4_experts_quant( + output, + output_scales, + input_tensor, + expert_offsets, + blockscale_offsets, + n_experts, + ) + output_scales = output_scales.view(torch.uint8) + return output, output_scales + + +# fp8 +def scaled_fp8_quant( + input: torch.Tensor, + scale: torch.Tensor | None = None, + num_token_padding: int | None = None, + scale_ub: torch.Tensor | None = None, + use_per_token_if_dynamic: bool = False, + output: torch.Tensor | None = None, + group_shape: tuple[int, int] | None = None, +) -> tuple[torch.Tensor, torch.Tensor]: + """ + Quantize input tensor to FP8 and return quantized tensor and scale. + + This function supports both static and dynamic quantization: If you + provide the scale, it will use static scaling and if you omit it, + the scale will be determined dynamically. The function also allows + optional padding of the output tensors for downstream kernels that + will benefit from padding. + + Args: + input: The input tensor to be quantized to FP8 (must be 2D: [M, N]) + scale: Optional scaling factor for the FP8 quantization. Supports: + - 0D or [1]: per-tensor scaling + - 1D: requires explicit group_shape to disambiguate per-channel + vs per-token (use (-1, 1) for per-channel, (1, -1) for per-token) + - 2D [M/group_m, N/group_n]: group scaling (e.g. [M, N/128] for + DeepSeek-style (1,128) groups, or [M/128, N/128] for (128,128)) + scale_ub: Optional upper bound for scaling factor in dynamic + per token case + num_token_padding: If specified, pad the first dimension + of the output to at least this value. + use_per_token_if_dynamic: Whether to do per_tensor or per_token + in the dynamic quantization case. + group_shape: Optional tuple (group_m, group_n) specifying the group + shape for static quantization. Use -1 for "full extent" (e.g., + (-1, -1) for per-tensor, (-1, 1) for per-channel, etc.) + Required for 1D scales; optional for 2D scales. + + Returns: + tuple[torch.Tensor, torch.Tensor]: The output tensor in FP8 and + scaling factor. + """ + # This code assumes batch_dim and num_tokens are flattened + assert input.ndim == 2 + shape: tuple[int, int] | torch.Size = input.shape + # For ROCm on MI300, the output fp8 dtype is torch.float_e3m3fnuz + out_dtype: torch.dtype = current_platform.fp8_dtype() + if num_token_padding: + shape = (max(num_token_padding, input.shape[0]), shape[1]) + if output is None: + output = torch.empty(shape, device=input.device, dtype=out_dtype) + else: + assert num_token_padding is None, "padding not supported if output passed in" + assert output.dtype == out_dtype + + if scale is None: + if use_per_token_if_dynamic: + scale = torch.empty((shape[0], 1), device=input.device, dtype=torch.float32) + torch.ops._C.dynamic_per_token_scaled_fp8_quant(output, input, scale, scale_ub) + else: + scale = torch.empty(1, device=input.device, dtype=torch.float32) + torch.ops._C.dynamic_scaled_fp8_quant(output, input, scale) + else: + torch.ops._C.static_scaled_fp8_quant(output, input, scale, group_shape) + + return output, scale + + +# gptq allspark +def allspark_repack_weight( + qweight: torch.Tensor, + scale: torch.Tensor, + zero_point: torch.Tensor | None = None, + has_zp: bool = False, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """ + Rearrange qweight, scale, and zero_point(if asymmetric) to n32k16 format + for Ampere W8A16 Fused Gemm kernel + + Args: + qweight: uint8 weight tensor, original k x n format. + scale: fp16/bf16 weight scale tensor, 1 x n format. + zero_point: fp16/bf16 weight zero_point tensor, 1 x n format. + Must be provided for asymmetric quantization. + has_zp: if use symmetric quantization, has_zp = False. + if use asymmetric quantization, has_zp = True. + + Returns: + tuple[torch.Tensor, torch.Tensor, torch.Tensor | None] : + rearranged weight, scale, and optionally zero_point. + """ + K = qweight.shape[0] + N = qweight.shape[1] + N_32align = (N + 32 - 1) // 32 * 32 + + qweight_reorder = torch.empty((N_32align, K), device=qweight.device, dtype=qweight.dtype) + scale_reorder = torch.empty((1, N_32align), device=scale.device, dtype=scale.dtype) + zero_point_reorder = None + if has_zp: + assert zero_point is not None, "zero_point must be provided for asymmetric quantization." + zero_point_reorder = torch.empty((1, N_32align), device=zero_point.device, dtype=zero_point.dtype) + + torch.ops._C.rearrange_kn_weight_as_n32k16_order( + qweight, + scale, + zero_point, + has_zp, + qweight_reorder, + scale_reorder, + zero_point_reorder, + K, + N, + N_32align, + ) + + return qweight_reorder, scale_reorder, zero_point_reorder + + +def allspark_w8a16_gemm( + a: torch.Tensor, + b_qweight: torch.Tensor, + b_scales: torch.Tensor, + b_qzeros: torch.Tensor | None, + n: int, + group_size: int, + sm_count: int, + sm_version: int, + CUBLAS_M_THRESHOLD: int, + has_zp: bool, + n32k16_reorder: bool, +) -> torch.Tensor: + return torch.ops._C.allspark_w8a16_gemm( + a, + b_qweight, + b_scales, + b_qzeros, + n, + group_size, + sm_count, + sm_version, + CUBLAS_M_THRESHOLD, + has_zp, + n32k16_reorder, + ) + + +# int8 +def scaled_int8_quant( + input: torch.Tensor, + scale: torch.Tensor | None = None, + azp: torch.Tensor | None = None, + symmetric: bool = True, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]: + """ + Quantize the input tensor to int8 and return the quantized tensor and scale, and maybe azp. + + Args: + input: The input tensor to be quantized to int8. + scale: Optional scaling factor for the int8 quantization. + When not provided, we invoke dynamic-per-token quantization. + azp: Optional zero-point for the int8 quantization. + Must be provided for asymmetric quantization if `scale` is provided. + symmetric: Whether to use symmetric quantization (scale only, azp ignored). + + Returns: + tuple[torch.Tensor, torch.Tensor, torch.Tensor | None] : Output int8 tensor, scales, and optionally azp. + """ + output = torch.empty_like(input, dtype=torch.int8) + if scale is not None: + # static-per-tensor quantization. + assert symmetric == (azp is None), "azp must only be provided for asymmetric quantization." + torch.ops._C.static_scaled_int8_quant(output, input, scale, azp) + return output, scale, azp + + # dynamic-per-token quantization. + input_scales = torch.empty((input.numel() // input.shape[-1], 1), device=input.device, dtype=torch.float32) + input_azp = None if symmetric else torch.empty_like(input_scales, dtype=torch.int32) + torch.ops._C.dynamic_scaled_int8_quant(output, input.contiguous(), input_scales, input_azp) + return output, input_scales, input_azp + + +# gguf +def ggml_dequantize(W: torch.Tensor, quant_type: int, m: int, n: int, dtype: torch.dtype | None) -> torch.Tensor: + return torch.ops._C.ggml_dequantize(W, quant_type, m, n, dtype) + + +def ggml_mul_mat_vec_a8( + W: torch.Tensor, + X: torch.Tensor, + quant_type: int, + row: int, +) -> torch.Tensor: + return torch.ops._C.ggml_mul_mat_vec_a8(W, X, quant_type, row) + + +def ggml_mul_mat_a8( + W: torch.Tensor, + X: torch.Tensor, + quant_type: int, + row: int, +) -> torch.Tensor: + return torch.ops._C.ggml_mul_mat_a8(W, X, quant_type, row) + + +def ggml_moe_a8( + X: torch.Tensor, + W: torch.Tensor, + sorted_token_ids: torch.Tensor, + expert_ids: torch.Tensor, + num_tokens_post_padded: torch.Tensor, + quant_type: int, + row: int, + top_k: int, + tokens: int, +) -> torch.Tensor: + return torch.ops._C.ggml_moe_a8( + X, + W, + sorted_token_ids, + expert_ids, + num_tokens_post_padded, + quant_type, + row, + top_k, + tokens, + ) + + +def ggml_moe_a8_vec( + X: torch.Tensor, + W: torch.Tensor, + topk_ids: torch.Tensor, + top_k: int, + quant_type: int, + row: torch.SymInt, + tokens: torch.SymInt, +) -> torch.Tensor: + return torch.ops._C.ggml_moe_a8_vec(X, W, topk_ids, top_k, quant_type, row, tokens) + + +def ggml_moe_get_block_size(quant_type: int) -> int: + return torch.ops._C.ggml_moe_get_block_size(quant_type) + + +if hasattr(torch.ops._C, "ggml_dequantize"): + + @register_fake("_C::ggml_dequantize") + def _ggml_dequantize_fake( + W: torch.Tensor, + quant_type: int, + m: torch.SymInt, + n: torch.SymInt, + dtype: torch.dtype | None = None, + ) -> torch.Tensor: + return torch.empty((m, n), dtype=torch.float16, device=W.device) + + @register_fake("_C::ggml_mul_mat_vec_a8") + def _ggml_mul_mat_vec_a8_fake( + W: torch.Tensor, + X: torch.Tensor, + quant_type: int, + row: torch.SymInt, + ) -> torch.Tensor: + return torch.empty((X.shape[0], row), dtype=X.dtype, device=W.device) + + @register_fake("_C::ggml_mul_mat_a8") + def _ggml_mul_mat_a8_fake( + W: torch.Tensor, + X: torch.Tensor, + quant_type: int, + row: torch.SymInt, + ) -> torch.Tensor: + batch = X.size(0) + return torch.empty((batch, row), dtype=X.dtype, device=W.device) + + @register_fake("_C::ggml_moe_a8") + def _ggml_moe_a8_fake( + X: torch.Tensor, + W: torch.Tensor, + sorted_token_ids: torch.Tensor, + expert_ids: torch.Tensor, + num_tokens_post_padded: torch.Tensor, + quant_type: int, + row: torch.SymInt, + top_k: torch.SymInt, + tokens: torch.SymInt, + ) -> torch.Tensor: + tokens = X.size(0) + return torch.empty((tokens * top_k, row), dtype=torch.float16, device=W.device) + + +if hasattr(torch.ops._C, "ggml_moe_a8_vec"): + + @register_fake("_C::ggml_moe_a8_vec") + def _ggml_moe_a8_vec_fake( + X: torch.Tensor, + W: torch.Tensor, + topk_ids: torch.Tensor, + top_k: int, + quant_type: int, + row: torch.SymInt, + tokens: torch.SymInt, + ) -> torch.Tensor: + tokens = X.size(0) + return torch.empty((tokens * top_k, row), dtype=X.dtype, device=W.device) + + +# mamba +def selective_scan_fwd( + u: torch.Tensor, + delta: torch.Tensor, + A: torch.Tensor, + B: torch.Tensor, + C: torch.Tensor, + D_: torch.Tensor | None, + z_: torch.Tensor | None, + delta_bias_: torch.Tensor | None, + delta_softplus: bool, + query_start_loc: torch.Tensor | None, + cache_indices: torch.Tensor | None, + has_initial_state: torch.Tensor | None, + ssm_states: torch.Tensor, + null_block_id: int, + block_size: int = 1024, + block_idx_first_scheduled_token: torch.Tensor | None = None, + block_idx_last_scheduled_token: torch.Tensor | None = None, + initial_state_idx: torch.Tensor | None = None, + cu_chunk_seqlen: torch.Tensor | None = None, + last_chunk_indices: torch.Tensor | None = None, +): + torch.ops._C.selective_scan_fwd( + u, + delta, + A, + B, + C, + D_, + z_, + delta_bias_, + delta_softplus, + query_start_loc, + cache_indices, + has_initial_state, + ssm_states, + null_block_id, + block_size, + block_idx_first_scheduled_token, + block_idx_last_scheduled_token, + initial_state_idx, + cu_chunk_seqlen, + last_chunk_indices, + ) + + +def causal_conv1d_update_cpu_vec( + x: torch.Tensor, + conv_state: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor | None = None, + activation: str | None = None, + conv_state_indices: torch.Tensor | None = None, + query_start_loc: torch.Tensor | None = None, + pad_slot_id: int = 0, +) -> torch.Tensor: + return torch.ops._C.causal_conv1d_update_cpu_vec( + x, + conv_state, + weight, + bias, + activation, + conv_state_indices, + query_start_loc, + pad_slot_id, + ) + + +def selective_state_update_cpu( + state: torch.Tensor, + x: torch.Tensor, + dt: torch.Tensor, + A: torch.Tensor, + B: torch.Tensor, + C: torch.Tensor, + D: torch.Tensor | None, + z: torch.Tensor | None, + dt_bias: torch.Tensor | None, + dt_softplus: bool, + state_batch_indices: torch.Tensor | None, + dst_state_batch_indices: torch.Tensor | None, + null_block_id: int, + out: torch.Tensor, + num_accepted_tokens: torch.Tensor | None, + cu_seqlens: torch.Tensor | None, +) -> None: + torch.ops._C.selective_state_update_cpu( + state, + x, + dt, + A, + B, + C, + D, + z, + dt_bias, + dt_softplus, + state_batch_indices, + dst_state_batch_indices, + null_block_id, + out, + num_accepted_tokens, + cu_seqlens, + ) + + +def mamba_chunk_scan_fwd_cpu( + out: torch.Tensor, + final_states: torch.Tensor, + x: torch.Tensor, + dt: torch.Tensor, + A: torch.Tensor, + B: torch.Tensor, + C: torch.Tensor, + D: torch.Tensor | None, + z: torch.Tensor | None, + cu_seqlens: torch.Tensor, +) -> None: + """Prefill SSM scan kernel. out and final_states are written in-place.""" + torch.ops._C.mamba_chunk_scan_fwd_cpu( + out, + final_states, + x, + dt, + A, + B, + C, + D, + z, + cu_seqlens, + ) + + +# ROCm skinny gemms +def LLMM1(a: torch.Tensor, b: torch.Tensor, rows_per_block: int) -> torch.Tensor: + return torch.ops._rocm_C.LLMM1(a, b, rows_per_block) + + +def wvSplitK(a: torch.Tensor, b: torch.Tensor, cu_count: int, bias: torch.Tensor = None) -> torch.Tensor: + return torch.ops._rocm_C.wvSplitK(a, b, bias, cu_count) + + +def wvSplitK_int4_g( + weight: torch.Tensor, + activation: torch.Tensor, + scale: torch.Tensor, + cu_count: int, + group_size: int, + zero_points: torch.Tensor | None = None, + bias: torch.Tensor | None = None, +) -> torch.Tensor: + # The kernel is weight-major: weight is the packed int4 operand + # (in_a, [out_features, K/2]) and activation is in_b ([num_tokens, K]). + return torch.ops._rocm_C.wvSplitK_int4_g(weight, activation, scale, zero_points, bias, cu_count, group_size) + + +if hasattr(torch.ops, "_rocm_C") and hasattr(torch.ops._rocm_C, "wvSplitK_int4_g"): + + @register_fake("_rocm_C::wvSplitK_int4_g") + def _wvSplitK_int4_g_fake( + in_a: torch.Tensor, + in_b: torch.Tensor, + in_scale: torch.Tensor, + in_zero_points: torch.Tensor | None, + in_bias: torch.Tensor | None, + CuCount: int, + group_size: int, + ) -> torch.Tensor: + return torch.empty((in_b.size(0), in_a.size(0)), dtype=in_b.dtype, device=in_b.device) + + +def wvSplitKrc(a: torch.Tensor, b: torch.Tensor, cu_count: int, bias: torch.Tensor = None) -> torch.Tensor: + return torch.ops._rocm_C.wvSplitKrc(a, b, bias, cu_count) + + +def wvSplitKQ( + a: torch.Tensor, + b: torch.Tensor, + out_dtype: torch.dtype, + scale_a: torch.Tensor, + scale_b: torch.Tensor, + cu_count: int, + bias: torch.Tensor = None, +) -> torch.Tensor: + out = torch.empty((b.shape[0], a.shape[0]), dtype=out_dtype, device=b.device) + torch.ops._rocm_C.wvSplitKQ(a, b, bias, out, scale_a, scale_b, cu_count) + return out + + +# moe +def moe_sum( + input: torch.Tensor, + output: torch.Tensor, + topk_ids: torch.Tensor | None = None, + expert_map: torch.Tensor | None = None, +): + torch.ops._moe_C.moe_sum(input, output, topk_ids, expert_map) + + +def moe_align_block_size( + topk_ids: torch.Tensor, + num_experts: int, + block_size: int, + sorted_token_ids: torch.Tensor, + experts_ids: torch.Tensor, + num_tokens_post_pad: torch.Tensor, + expert_map: torch.Tensor | None = None, +) -> None: + torch.ops._moe_C.moe_align_block_size( + topk_ids, + num_experts, + block_size, + sorted_token_ids, + experts_ids, + num_tokens_post_pad, + expert_map, + ) + + +def batched_moe_align_block_size( + max_tokens_per_batch: int, + block_size: int, + expert_num_tokens: torch.Tensor, + sorted_ids: torch.Tensor, + expert_ids: torch.Tensor, + num_tokens_post_pad: torch.Tensor, +) -> None: + torch.ops._moe_C.batched_moe_align_block_size( + max_tokens_per_batch, + block_size, + expert_num_tokens, + sorted_ids, + expert_ids, + num_tokens_post_pad, + ) + + +def moe_lora_align_block_size( + topk_ids: torch.Tensor, + token_lora_mapping: torch.Tensor, + num_experts: int, + block_size: int, + max_loras: int, + max_num_tokens_padded: int, + max_num_m_blocks: int, + sorted_token_ids: torch.Tensor, + experts_ids: torch.Tensor, + num_tokens_post_pad: torch.Tensor, + adapter_enabled: torch.Tensor, + lora_ids: torch.Tensor, + expert_map: torch.Tensor | None = None, +) -> None: + torch.ops._moe_C.moe_lora_align_block_size( + topk_ids, + token_lora_mapping, + num_experts, + block_size, + max_loras, + max_num_tokens_padded, + max_num_m_blocks, + sorted_token_ids, + experts_ids, + num_tokens_post_pad, + adapter_enabled, + lora_ids, + expert_map, + ) + + +def moe_wna16_gemm( + input: torch.Tensor, + output: torch.Tensor, + b_qweight: torch.Tensor, + b_scales: torch.Tensor, + b_qzeros: torch.Tensor | None, + topk_weights: torch.Tensor | None, + sorted_token_ids: torch.Tensor, + experts_ids: torch.Tensor, + num_tokens_post_pad: torch.Tensor, + top_k: int, + BLOCK_SIZE_M: int, + BLOCK_SIZE_N: int, + BLOCK_SIZE_K: int, + bit: int, +) -> torch.Tensor: + if not current_platform.is_cuda(): + raise NotImplementedError("The optimized moe_wna16_gemm kernel is only available on CUDA platforms") + torch.ops._moe_C.moe_wna16_gemm( + input, + output, + b_qweight, + b_scales, + b_qzeros, + topk_weights, + sorted_token_ids, + experts_ids, + num_tokens_post_pad, + top_k, + BLOCK_SIZE_M, + BLOCK_SIZE_N, + BLOCK_SIZE_K, + bit, + ) + + +def dsv3_router_gemm( + hidden_states: torch.Tensor, + router_weight: torch.Tensor, + output_dtype: torch.dtype, +) -> torch.Tensor: + output = torch.empty( + hidden_states.shape[0], + router_weight.shape[0], + device=hidden_states.device, + dtype=output_dtype, + ) + torch.ops._moe_C.dsv3_router_gemm(output, hidden_states, router_weight) + return output + + +def fp32_router_gemm( + hidden_states: torch.Tensor, + router_weight: torch.Tensor, +) -> torch.Tensor: + output = torch.empty( + hidden_states.shape[0], + router_weight.shape[0], + device=hidden_states.device, + dtype=torch.float32, + ) + torch.ops._C.fp32_router_gemm(output, hidden_states, router_weight) + return output + + +if hasattr(torch.ops, "_C") and hasattr(torch.ops._C, "fp32_router_gemm"): + + @register_fake("_C::fp32_router_gemm") + def fp32_router_gemm_fake( + output: torch.Tensor, + mat_a: torch.Tensor, + mat_b: torch.Tensor, + ) -> None: + return + + +def topk_softmax( + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + token_expert_indices: torch.Tensor, + gating_output: torch.Tensor, + renormalize: bool = False, + 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, + token_expert_indices, + gating_output, + renormalize, + e_score_correction_bias, + is_padding, + ) + + +def topk_sigmoid( + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + token_expert_indices: torch.Tensor, + gating_output: torch.Tensor, + renormalize: bool = False, + e_score_correction_bias: torch.Tensor | None = None, + routed_scaling_factor: float = 1.0, + is_padding: torch.Tensor | None = None, +) -> None: + torch.ops._moe_C.topk_sigmoid( + topk_weights, + topk_ids, + token_expert_indices, + gating_output, + renormalize, + e_score_correction_bias, + routed_scaling_factor, + is_padding, + ) + + +def topk_hash_softplus_sqrt( + topk_weights: torch.Tensor, + topk_indices: torch.Tensor, + token_expert_indices: torch.Tensor, + gating_output: torch.Tensor, + renormalize: bool = False, + routed_scaling_factor: float = 1.0, + e_score_correction_bias: torch.Tensor | None = None, + input_tokens: torch.Tensor | None = None, + 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, + token_expert_indices, + gating_output, + renormalize, + routed_scaling_factor, + e_score_correction_bias, + input_tokens, + hash_indices_table, + is_padding, + ) + + +def grouped_topk( + scores: torch.Tensor, + num_expert_group: int, + topk_group: int, + topk: int, + renormalize: bool, + routed_scaling_factor: float, + bias: torch.Tensor, + scoring_func: int = 0, +): + """ + Perform grouped top-k routing for mixture of experts. + + Args: + scores: Raw inputs (logits if scoring_func=1, scores if scoring_func=0) + num_expert_group: Number of expert groups + topk_group: Number of groups to select + topk: Number of experts to select per token + renormalize: Whether to renormalize the output weights + routed_scaling_factor: Scaling factor for routing weights + bias: Bias tensor (e_score_correction_bias). Always fused in kernel. + scoring_func: 0=none (no activation), 1=sigmoid + """ + if not current_platform.is_cuda(): + raise NotImplementedError("The fused grouped_topk kernel is only available on CUDA platforms") + return torch.ops._moe_C.grouped_topk( + scores, + num_expert_group, + topk_group, + topk, + renormalize, + routed_scaling_factor, + bias, + scoring_func, + ) + + +def moe_wna16_marlin_gemm( + input: torch.Tensor, + output: torch.Tensor | None, + b_qweight: torch.Tensor, + b_bias: torch.Tensor | None, + b_scales: torch.Tensor, + a_scales: torch.Tensor | None, + global_scale: torch.Tensor | None, + b_qzeros: torch.Tensor | None, + g_idx: torch.Tensor | None, + perm: torch.Tensor | None, + workspace: torch.Tensor, + sorted_token_ids: torch.Tensor, + expert_ids: torch.Tensor, + num_tokens_past_padded: torch.Tensor, + topk_weights: torch.Tensor, + moe_block_size: int, + top_k: int, + mul_topk_weights: bool, + b_q_type: ScalarType, + size_m: int, + size_n: int, + size_k: int, + is_k_full: bool, + use_atomic_add: bool, + use_fp32_reduce: bool, + is_zp_float: bool, + thread_k: int = -1, + thread_n: int = -1, + blocks_per_sm: int = -1, +) -> torch.Tensor: + return torch.ops._moe_C.moe_wna16_marlin_gemm( + input, + output, + b_qweight, + b_bias, + b_scales, + a_scales, + global_scale, + b_qzeros, + g_idx, + perm, + workspace, + sorted_token_ids, + expert_ids, + num_tokens_past_padded, + topk_weights, + moe_block_size, + top_k, + mul_topk_weights, + b_q_type.id, + size_m, + size_n, + size_k, + is_k_full, + use_atomic_add, + use_fp32_reduce, + is_zp_float, + thread_k, + thread_n, + blocks_per_sm, + ) + + +if hasattr(torch.ops, "_moe_C") and hasattr(torch.ops._moe_C, "moe_wna16_marlin_gemm"): + + @register_fake("_moe_C::moe_wna16_marlin_gemm") + def moe_wna16_marlin_gemm_fake( + input: torch.Tensor, + output: torch.Tensor | None, + b_qweight: torch.Tensor, + b_bias: torch.Tensor | None, + b_scales: torch.Tensor, + a_scales: torch.Tensor | None, + global_scale: torch.Tensor | None, + b_qzeros: torch.Tensor | None, + g_idx: torch.Tensor | None, + perm: torch.Tensor | None, + workspace: torch.Tensor, + sorted_token_ids: torch.Tensor, + expert_ids: torch.Tensor, + num_tokens_past_padded: torch.Tensor, + topk_weights: torch.Tensor, + moe_block_size: int, + top_k: int, + mul_topk_weights: bool, + b_q_type: ScalarType, + size_m: int, + size_n: int, + size_k: int, + is_k_full: bool, + use_atomic_add: bool, + use_fp32_reduce: bool, + is_zp_float: bool, + ): + return torch.empty((size_m * top_k, size_n), dtype=input.dtype, device=input.device) + + +def reshape_and_cache( + key: torch.Tensor, + value: torch.Tensor, + key_cache: torch.Tensor, + value_cache: torch.Tensor, + slot_mapping: torch.Tensor, + kv_cache_dtype: str, + k_scale: torch.Tensor, + v_scale: torch.Tensor, +) -> None: + torch.ops._C_cache_ops.reshape_and_cache( + key, + value, + key_cache, + value_cache, + slot_mapping, + kv_cache_dtype, + k_scale, + v_scale, + ) + + +def reshape_and_cache_flash( + key: torch.Tensor, + value: torch.Tensor, + key_cache: torch.Tensor, + value_cache: torch.Tensor, + slot_mapping: torch.Tensor, + kv_cache_dtype: str, + k_scale: torch.Tensor, + v_scale: torch.Tensor, +) -> None: + torch.ops._C_cache_ops.reshape_and_cache_flash( + key, + value, + key_cache, + value_cache, + slot_mapping, + kv_cache_dtype, + k_scale, + v_scale, + ) + + +def fused_minimax_m3_qknorm_rope_kv_insert( + qkv: torch.Tensor, + q_norm_weight: torch.Tensor, + k_norm_weight: torch.Tensor, + cos_sin_cache: torch.Tensor, + positions: torch.Tensor, + num_heads: int, + num_kv_heads: int, + rotary_dim: int, + eps: float, + index_q_norm_weight: torch.Tensor | None = None, + index_k_norm_weight: torch.Tensor | None = None, + num_index_heads: int = 0, + slot_mapping: torch.Tensor | None = None, + index_slot_mapping: torch.Tensor | None = None, + kv_cache: torch.Tensor | None = None, + index_cache: torch.Tensor | None = None, + block_size: int = 0, + q_out: torch.Tensor | None = None, + index_q_out: torch.Tensor | None = None, + kv_cache_dtype: str = "auto", + skip_index_branch: bool = False, +) -> None: + """Fused MiniMax-M3 attention pre-processing (in-place). + + Applies Gemma RMSNorm + partial NeoX RoPE to ``qkv`` in place. ``qkv`` is a + single fused tensor: + + - dense layer (``num_index_heads == 0``): ``[q | k | v]``; + - sparse layer (``num_index_heads > 0``): ``[q | k | v | index_q | + index_k]`` — the index branch is read straight out of ``qkv``. + + When ``kv_cache`` is given (sparse serving), also scatter-inserts the + normed/roped k & v into the paged KV cache by ``slot_mapping`` and the + index key into ``index_cache`` by ``index_slot_mapping``. ``kv_cache_dtype`` + selects the cache storage/conversion path. If + ``index_slot_mapping`` is omitted, ``slot_mapping`` is used for both caches. + + If ``q_out`` / ``index_q_out`` (contiguous ``[N, nq*128]`` / ``[N, + niq*128]``) are given, the normed/roped q / index_q are written there + instead of in place — folding the de-interleave into this kernel's store so + callers skip a separate ``.contiguous()`` copy before the SM100 sparse + attention's flat TMA descriptor. + + When ``skip_index_branch`` is true, sparse rows still keep their packed + ``[index_q | index_k]`` tail, but the kernel only processes the main q/k/v + branches and main KV cache. This is used by MiniMax-M3 index-topk reuse + layers that consume top-k block ids selected by an earlier sparse layer. + """ + torch.ops._C.fused_minimax_m3_qknorm_rope_kv_insert( + qkv, + q_norm_weight, + k_norm_weight, + cos_sin_cache, + positions, + num_heads, + num_kv_heads, + rotary_dim, + eps, + index_q_norm_weight, + index_k_norm_weight, + num_index_heads, + slot_mapping, + index_slot_mapping, + kv_cache, + index_cache, + block_size, + q_out, + index_q_out, + kv_cache_dtype, + skip_index_branch, + ) + + +def concat_and_cache_mla( + kv_c: torch.Tensor, + k_pe: torch.Tensor, + kv_cache: torch.Tensor, + slot_mapping: torch.Tensor, + kv_cache_dtype: str, + scale: torch.Tensor, +) -> None: + torch.ops._C_cache_ops.concat_and_cache_mla(kv_c, k_pe, kv_cache, slot_mapping, kv_cache_dtype, scale) + + +def kimi_k3_attn_res( + prefix: torch.Tensor, + delta: torch.Tensor, + blocks: torch.Tensor, + norm_weight: torch.Tensor, + qk_weight: torch.Tensor, + output_norm_weight: torch.Tensor, + num_blocks: int, + eps: float, + output_norm_eps: float, +) -> torch.Tensor: + output = torch.empty_like(prefix) + torch.ops._C.kimi_k3_attn_res( + prefix, + delta, + blocks, + norm_weight, + qk_weight, + output_norm_weight, + output, + num_blocks, + eps, + output_norm_eps, + ) + return output + + +def concat_and_cache_mla_rope_fused( + positions: torch.Tensor, + q_pe: torch.Tensor, + k_pe: torch.Tensor, + kv_c: torch.Tensor, + cos_sin_cache: torch.Tensor, + is_neox: bool, + slot_mapping: torch.Tensor, + kv_cache: torch.Tensor, + kv_cache_dtype: str, + kv_cache_scale: torch.Tensor, +) -> None: + torch.ops._C_cache_ops.concat_and_cache_mla_rope_fused( + positions, + q_pe, + k_pe, + kv_c, + cos_sin_cache, + is_neox, + slot_mapping, + kv_cache, + kv_cache_dtype, + kv_cache_scale, + ) + + +def swap_blocks( + src: torch.Tensor, + dst: torch.Tensor, + block_size_in_bytes: int, + block_mapping: torch.Tensor, +) -> None: + """ + Copy specific blocks from one tensor to another. + + This method assumes each of the two input tensors is composed of + consecutive contiguous blocks, of size block_size_in_bytes. + i.e. the memory layout for each tensor is: + [block0] [block1] ... [block N] + + block_mapping determines the subset of blocks to copy of the source tensor, + and their matching destination block number on the destination tensor. + block_mapping is expected to be a tensor of shape (num_blocks_to_copy, 2) + where each block_mapping[i] represents a single copy operation, copying + block #block_mapping[i][0] from the source tensor + to block #block_mapping[i][1] on the destination tensor. + block_mapping should have dtype int64. + + The source and the destination tensors can be either on cpu or gpu, + but not both on cpu. + the block mapping tensor must on cpu. + """ + torch.ops._C_cache_ops.swap_blocks(src, dst, block_size_in_bytes, block_mapping) + + +def swap_blocks_batch( + src_ptrs: torch.Tensor, + dst_ptrs: torch.Tensor, + sizes: torch.Tensor, + is_src_access_order_any: bool = False, +) -> None: + """ + Batch version of swap_blocks: submit all copies in a single driver call. + + Each entry specifies a raw pointer copy: src_ptrs[i] -> dst_ptrs[i] + of sizes[i] bytes. All three tensors must be CPU tensors with the + platform-appropriate pointer dtype: int64 on CUDA/ROCm (required by + cache_kernels.cu) and uint64 on XPU (required by the XPU DMA engine). + On CUDA 12.8+ this uses cuMemcpyBatchAsync for minimal submission + overhead; on older CUDA it falls back to a loop of cudaMemcpyAsync. + + is_src_access_order_any: if True, pass CU_MEMCPY_SRC_ACCESS_ORDER_ANY to + cuMemcpyBatchAsync, letting the DMA engine prefetch source bytes + out of stream order. Only safe when no GPU stream is concurrently + writing to the source. Defaults to False (STREAM ordering), which + is always safe. + """ + if current_platform.is_xpu(): + torch.ops._C_cache_ops.swap_blocks_batch(src_ptrs, dst_ptrs, sizes) + else: + torch.ops._C_cache_ops.swap_blocks_batch(src_ptrs, dst_ptrs, sizes, is_src_access_order_any) + + +def convert_fp8(output: torch.Tensor, input: torch.Tensor, scale: float = 1.0, kv_dtype: str = "fp8") -> None: + torch.ops._C_cache_ops.convert_fp8(output, input, scale, kv_dtype) + + +def gather_and_maybe_dequant_cache( + src_cache: torch.Tensor, + dst: torch.Tensor, + block_table: torch.Tensor, + cu_seq_lens: torch.Tensor, + token_to_seq: torch.Tensor, + num_tokens: int, + kv_cache_dtype: str, + scale: torch.Tensor, + seq_starts: torch.Tensor | None = None, +) -> None: + torch.ops._C_cache_ops.gather_and_maybe_dequant_cache( + src_cache, + dst, + block_table, + cu_seq_lens, + token_to_seq, + num_tokens, + kv_cache_dtype, + scale, + seq_starts, + ) + + +def cp_gather_cache( + src_cache: torch.Tensor, + dst: torch.Tensor, + block_table: torch.Tensor, + cu_seq_lens: torch.Tensor, + batch_size: int, + seq_starts: torch.Tensor | None = None, +) -> None: + torch.ops._C_cache_ops.cp_gather_cache(src_cache, dst, block_table, cu_seq_lens, batch_size, seq_starts) + + +def cp_gather_and_upconvert_fp8_kv_cache( + src_cache: torch.Tensor, + dst: torch.Tensor, + block_table: torch.Tensor, + workspace_starts: torch.Tensor, + batch_size: int, + seq_starts: torch.Tensor | None = None, +) -> None: + """Gather and upconvert FP8 KV cache to BF16 workspace. + + Args: + src_cache: FP8 KV cache [num_blocks, block_size, 656] + dst: BF16 output workspace [total_tokens, 576] + block_table: Block indices [num_reqs, max_blocks] + workspace_starts: Workspace start offsets [num_reqs] + batch_size: Number of requests + seq_starts: Optional source sequence offsets [num_reqs] + """ + torch.ops._C_cache_ops.cp_gather_and_upconvert_fp8_kv_cache( + src_cache, dst, block_table, workspace_starts, batch_size, seq_starts + ) + + +def concat_mla_q( + ql_nope: torch.Tensor, + q_pe: torch.Tensor, + q_out: torch.Tensor, +) -> None: + """Concatenate query nope and rope for MLA/DSA attention. + + Args: + ql_nope: Query nope component [num_tokens, num_heads, nope_dim] + q_pe: Query rope component [num_tokens, num_heads, rope_dim] + q_out: Output tensor [num_tokens, num_heads, nope_dim + rope_dim] + """ + torch.ops._C_cache_ops.concat_mla_q(ql_nope, q_pe, q_out) + + +def indexer_k_quant_and_cache( + k: torch.Tensor, + kv_cache: torch.Tensor, + slot_mapping: torch.Tensor, + quant_block_size: int, + kv_cache_dtype: str, +) -> None: + torch.ops._C_cache_ops.indexer_k_quant_and_cache(k, kv_cache, slot_mapping, quant_block_size, kv_cache_dtype) + + +def top_k_per_row_prefill( + logits: torch.Tensor, + cu_seqlen_ks: torch.Tensor, + cu_seqlen_ke: torch.Tensor, + raw_topk_indices: torch.Tensor, + num_rows: int, + stride0: int, + stride1: int, + topk_tokens: int, +) -> None: + torch.ops._C.top_k_per_row_prefill( + logits, + cu_seqlen_ks, + cu_seqlen_ke, + raw_topk_indices, + num_rows, + stride0, + stride1, + topk_tokens, + ) + + +def top_k_per_row_decode( + logits: torch.Tensor, + next_n: int, + seq_lens: torch.Tensor, + raw_topk_indices: torch.Tensor, + num_rows: int, + stride0: int, + stride1: int, + topk_tokens: int, +) -> None: + torch.ops._C.top_k_per_row_decode( + logits, + next_n, + seq_lens, + raw_topk_indices, + num_rows, + stride0, + stride1, + topk_tokens, + ) + + +@cache +def supports_sm89_dsa() -> bool: + """Whether the sm89 DeepSeek sparse attention ops were compiled into _C.""" + return hasattr(torch.ops._C, "sm89_sparse_mla_fwd") + + +def sm89_fp8_paged_mqa_logits( + q: torch.Tensor, + kv_cache_pool: torch.Tensor, + weights: torch.Tensor, + seq_lens: torch.Tensor, + block_table: torch.Tensor, + sched: torch.Tensor, + logits: torch.Tensor, + clean_logits: bool, +) -> None: + """Paged fp8 MQA indexer logits for decode on sm89 (block size 64). + + q is [B, next_n, 64, 128] fp8 e4m3 viewed as uint8; kv_cache_pool [pool_pages, 8448] + uint8 (per page, 8192B fp8 keys + 256B fp32 per-token scales); weights [B*next_n, 64] + fp32; sched [(P+1), 2] int32 built by sm89_paged_mqa_logits_metadata; logits + [B*next_n, max_model_len] fp32, written for columns [0, seq_len) per row. With + clean_logits, the tail columns of each row's last partial page are set to -inf. + """ + torch.ops._C.sm89_fp8_paged_mqa_logits( + q, kv_cache_pool, weights, seq_lens, block_table, sched, logits, clean_logits + ) + + +def sm89_paged_mqa_logits_metadata(seq_lens: torch.Tensor, sched: torch.Tensor, next_n: int) -> None: + """Build the [(P+1), 2] int32 (request, page) partition table for sm89_fp8_paged_mqa_logits.""" + torch.ops._C.sm89_paged_mqa_logits_metadata(seq_lens, sched, next_n) + + +def sm89_fp8_mqa_logits( + q: torch.Tensor, + kv: torch.Tensor, + kv_scales: torch.Tensor, + weights: torch.Tensor, + cu_seqlen_ks: torch.Tensor, + cu_seqlen_ke: torch.Tensor, + logits: torch.Tensor, +) -> None: + """Prefill/ragged fp8 MQA indexer logits on sm89 over a contiguous gathered kv buffer. + + q is [M, 64, 128] fp8 e4m3 viewed as uint8; kv [N, 128] fp8 e4m3 viewed as uint8; + kv_scales [N] fp32; weights [M, 64] fp32; cu_seqlen_ks/cu_seqlen_ke [M] int32 + per-row key windows (clamped to [0, N]); logits [M, N] fp32, written only in the + [ks, ke) window of each row (the caller owns the fill value outside the window). + """ + torch.ops._C.sm89_fp8_mqa_logits(q, kv, kv_scales, weights, cu_seqlen_ks, cu_seqlen_ke, logits) + + +def sm89_sparse_mla_fwd( + q: torch.Tensor, + kv_cache_pool: torch.Tensor, + indices: torch.Tensor, + out: torch.Tensor, + lse: torch.Tensor, + sm_scale: float, + topk_lens: torch.Tensor | None = None, +) -> None: + """Sparse MLA forward on sm89, gathering directly from an fp8_ds_mla pool by slot. + + q is [T, h, 576] bf16; kv_cache_pool [S, 656] uint8 (fp8_ds_mla rows); indices + [T, topk] int32 physical slots, -1 padded; out [T, h, 512] bf16; lse [T, h] fp32. + Rows whose indices are all -1 produce zero output and -inf LSE. + + topk_lens is an optional [T] int32 of per-token valid counts. It requires index + rows to be front-compacted (valid slots first, then -1); the kernel then stops + after the last valid slot instead of scanning all topk columns, with output + identical to topk_lens=None. DCP passes its per-rank counts here + (~topk/dcp_world_size valid). + """ + torch.ops._C.sm89_sparse_mla_fwd(q, kv_cache_pool, indices, out, lse, sm_scale, topk_lens) + + +def cp_gather_indexer_k_quant_cache( + kv_cache: torch.Tensor, + dst_k: torch.Tensor, + dst_scale: torch.Tensor, + block_table: torch.Tensor, + cu_seq_lens: torch.Tensor, +) -> None: + torch.ops._C_cache_ops.cp_gather_indexer_k_quant_cache(kv_cache, dst_k, dst_scale, block_table, cu_seq_lens) + + +def get_device_attribute(attribute: int, device: int) -> int: + return torch.ops._C_cuda_utils.get_device_attribute(attribute, device) + + +def get_max_shared_memory_per_block_device_attribute(device: int) -> int: + # ruff: noqa: E501 + return torch.ops._C_cuda_utils.get_max_shared_memory_per_block_device_attribute(device) + + +# custom ar +def init_custom_ar( + ipc_tensors: list[torch.Tensor], + rank_data: torch.Tensor, + rank: int, + fully_connected: bool, +) -> int: + return torch.ops._C_custom_ar.init_custom_ar(ipc_tensors, rank_data, rank, fully_connected) + + +def all_reduce( + fa: int, + inp: torch.Tensor, + out: torch.Tensor, + reg_buffer: int, + reg_buffer_sz_bytes: int, +) -> None: + torch.ops._C_custom_ar.all_reduce(fa, inp, out, reg_buffer, reg_buffer_sz_bytes) + + +def dispose(fa: int) -> None: + torch.ops._C_custom_ar.dispose(fa) + + +def meta_size() -> int: + return torch.ops._C_custom_ar.meta_size() + + +def register_buffer(fa: int, ipc_tensors: list[int]) -> None: + return torch.ops._C_custom_ar.register_buffer(fa, ipc_tensors) + + +def get_graph_buffer_ipc_meta(fa: int) -> tuple[list[int], list[int]]: + return torch.ops._C_custom_ar.get_graph_buffer_ipc_meta(fa) + + +def register_graph_buffers(fa: int, handles: list[list[int]], offsets: list[list[int]]) -> None: + torch.ops._C_custom_ar.register_graph_buffers(fa, handles, offsets) + + +def allocate_shared_buffer_and_handle(size: int) -> tuple[int, torch.Tensor]: + return torch.ops._C_custom_ar.allocate_shared_buffer_and_handle(size) + + +def open_mem_handle(mem_handle: torch.Tensor): + return torch.ops._C_custom_ar.open_mem_handle(mem_handle) + + +def free_shared_buffer(ptr: int) -> None: + torch.ops._C_custom_ar.free_shared_buffer(ptr) + + +# quick all reduce +def init_custom_qr(rank: int, world_size: int, qr_max_size: int | None = None) -> int: + return torch.ops._C_custom_ar.init_custom_qr(rank, world_size, qr_max_size) + + +def qr_destroy(fa: int) -> None: + torch.ops._C_custom_ar.qr_destroy(fa) + + +def qr_all_reduce( + fa: int, + inp: torch.Tensor, + out: torch.Tensor, + quant_level: int, + cast_bf2half: bool = False, +) -> None: + torch.ops._C_custom_ar.qr_all_reduce(fa, inp, out, quant_level, cast_bf2half) + + +def qr_get_handle(fa: int) -> torch.Tensor: + return torch.ops._C_custom_ar.qr_get_handle(fa) + + +def qr_open_handles(fa: int, handles: list[torch.Tensor]) -> None: + return torch.ops._C_custom_ar.qr_open_handles(fa, handles) + + +def qr_max_size() -> int: + return torch.ops._C_custom_ar.qr_max_size() + + +def sm100_cutlass_mla_decode( + out: torch.Tensor, + lse: torch.Tensor, + q_nope: torch.Tensor, + q_pe: torch.Tensor, + kv_c_and_k_pe_cache: torch.Tensor, + seq_lens: torch.Tensor, + page_table: torch.Tensor, + workspace: torch.Tensor, + scale: float, + num_kv_splits: int, +) -> torch.Tensor: + torch.ops._C.sm100_cutlass_mla_decode( + out, + lse, + q_nope, + q_pe, + kv_c_and_k_pe_cache, + seq_lens, + page_table, + workspace, + scale, + num_kv_splits, + ) + return out + + +def sm100_cutlass_mla_get_workspace_size(max_seq_len: int, num_batches: int, sm_count: int, num_kv_splits: int) -> int: + return torch.ops._C.sm100_cutlass_mla_get_workspace_size(max_seq_len, num_batches, sm_count, num_kv_splits) + + +def dsv3_fused_a_gemm( + output: torch.Tensor, + mat_a: torch.Tensor, + mat_b: torch.Tensor, +) -> None: + """DeepSeek V3 fused A GEMM (SM 9.0+, bf16 only, 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 + + Optimized for the DeepSeek V2/V3 QKV A-projection at small batch sizes. + Requires SM 9.0+ (Hopper). + """ + torch.ops._C.dsv3_fused_a_gemm(output, mat_a, mat_b) + + +if hasattr(torch.ops._C, "weight_packed_linear"): + + @register_fake("_C::weight_packed_linear") + def weight_packed_linear_fake( + mat1: torch.Tensor, + mat2: torch.Tensor, + bias: torch.Tensor | None, + is_vnni: bool, + ) -> torch.Tensor: + return torch.empty((mat1.size(0), mat2.size(0)), dtype=mat1.dtype, device=mat2.device) + + +class CPUQuantMethod(IntEnum): + UNQUANT = 0 + INT8_W8A8 = 1 + FP8_W8A16 = 2 + INT4_W4A8 = 3 + MXFP4 = 4 + + +if hasattr(torch.ops._C, "fused_experts_cpu"): + + @register_fake("_C::fused_experts_cpu") + def fused_experts_cpu_fake( + hidden_states: torch.Tensor, + w1: torch.Tensor, + w2: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + inplace: bool, + moe_comp_method: CPUQuantMethod, + w1_scale: torch.Tensor | None, + w2_scale: torch.Tensor | None, + w1_zero: torch.Tensor | None, + w2_zero: torch.Tensor | None, + block_size: list[int] | None, + w1_bias: torch.Tensor | None, + w2_bias: torch.Tensor | None, + alpha: float | None, + limit: float | None, + is_vnni: bool, + ) -> torch.Tensor: + return torch.empty_like(hidden_states) + + +if hasattr(torch.ops._C, "dynamic_4bit_int_moe"): + + @register_fake("_C::dynamic_4bit_int_moe") + def dynamic_4bit_int_moe_fake( + x: torch.Tensor, + topk_ids: torch.Tensor, + topk_weights: torch.Tensor, + w13_packed: torch.Tensor, + w2_packed: torch.Tensor, + hidden_size: int, + intermediate_size: int, + group_size: int, + apply_router_weight_on_input: bool, + activation_kind: int, + ) -> torch.Tensor: + return x.new_empty((x.size(0), hidden_size)) + + +def fused_experts_cpu( + hidden_states: torch.Tensor, + w1: torch.Tensor, + w2: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + inplace: bool, + moe_comp_method: CPUQuantMethod, + w1_scale: torch.Tensor | None, + w2_scale: torch.Tensor | None, + w1_zero: torch.Tensor | None, + w2_zero: torch.Tensor | None, + block_size: list[int] | None, + w1_bias: torch.Tensor | None = None, + w2_bias: torch.Tensor | None = None, + alpha: float | None = None, + limit: float | None = None, + is_vnni: bool = True, +) -> torch.Tensor: + return torch.ops._C.fused_experts_cpu( + hidden_states, + w1, + w2, + topk_weights, + topk_ids, + inplace, + moe_comp_method, + w1_scale, + w2_scale, + w1_zero, + w2_zero, + block_size, + w1_bias, + w2_bias, + alpha, + limit, + is_vnni, + ) + + +if hasattr(torch.ops._C, "int8_scaled_mm_with_quant"): + + @register_fake("_C::int8_scaled_mm_with_quant") + def int8_scaled_mm_with_quant_fake( + mat1: torch.Tensor, + mat2: torch.Tensor, + scales2: torch.Tensor, + bias: torch.Tensor | None, + out_dtype: torch.dtype, + is_vnni: bool, + ) -> torch.Tensor: + M = mat1.size(0) + N = mat2.size(0) + return torch.empty((M, N), dtype=out_dtype) + + +class CPUQuantAlgo(IntEnum): + AWQ = 0 + GPTQ = 1 + + +if hasattr(torch.ops._C, "convert_weight_packed_scale_zp"): + + @register_fake("_C::convert_weight_packed_scale_zp") + def convert_weight_packed_scale_zp_fake( + qweight: torch.Tensor, + qzeros: torch.Tensor, + scales: torch.Tensor, + quant_method_4bit: CPUQuantAlgo, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + return ( + torch.empty_like(qweight), + torch.empty_like(qzeros), + torch.empty_like(scales), + ) + + +def convert_weight_packed_scale_zp( + qweight: torch.Tensor, + qzeros: torch.Tensor, + scales: torch.Tensor, + quant_method_4bit: CPUQuantAlgo, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + return torch.ops._C.convert_weight_packed_scale_zp( + qweight, + qzeros, + scales, + quant_method_4bit, + ) + + +if hasattr(torch.ops._C, "int4_scaled_mm_cpu"): + + @register_fake("_C::int4_scaled_mm_cpu") + def int4_scaled_mm_cpu_fake( + x: torch.Tensor, + w: torch.Tensor, + w_zeros: torch.Tensor, + w_scales: torch.Tensor, + bias: torch.Tensor | None, + ) -> torch.Tensor: + N = w_scales.size(0) * w_scales.size(-1) + return torch.empty((x.size(0), N), dtype=x.dtype, device=x.device) + + +def int4_scaled_mm_cpu( + x: torch.Tensor, + w: torch.Tensor, + w_zeros: torch.Tensor, + w_scales: torch.Tensor, + bias: torch.Tensor | None, +) -> torch.Tensor: + x_shape = x.shape + x_2d = x.reshape(-1, x_shape[-1]) if len(x_shape) > 2 else x + + out = torch.ops._C.int4_scaled_mm_cpu( + x_2d, + w, + w_zeros, + w_scales, + bias, + ) + out = out.reshape(x_shape[:-1] + (out.size(-1),)) if len(x_shape) > 2 else out + return out + + +if hasattr(torch.ops._C, "fp8_scaled_mm_cpu"): + + @register_fake("_C::fp8_scaled_mm_cpu") + def fp8_scaled_mm_cpu_fake( + mat1: torch.Tensor, + mat2: torch.Tensor, + scales2: torch.Tensor, + block_size: list[int], + bias: torch.Tensor | None, + out_dtype: torch.dtype, + is_vnni: bool, + ) -> torch.Tensor: + M = mat1.size(0) + N = mat2.size(0) + return torch.empty((M, N), dtype=out_dtype, device=mat1.device) + + +_supports_cpu_fp8_w8a16 = bool(hasattr(torch.ops._C, "fp8_scaled_mm_cpu")) + + +def fp8_scaled_mm_cpu( + mat1: torch.Tensor, + mat2: torch.Tensor, + scales2: torch.Tensor, + block_size: list[int], + bias: torch.Tensor | None, + out_dtype: torch.dtype, + is_vnni: bool, +) -> torch.Tensor: + return torch.ops._C.fp8_scaled_mm_cpu(mat1, mat2, scales2, block_size, bias, out_dtype, is_vnni) + + +def chunk_gated_delta_rule_cpu( + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + initial_state: torch.Tensor, + output_final_state: bool, + cu_seqlens: torch.Tensor, + head_first: bool, + use_qk_l2norm_in_kernel: bool, + eps: float = 1e-5, +) -> tuple[torch.Tensor, torch.Tensor]: + return torch.ops._C.chunk_gated_delta_rule_cpu( + query, + key, + value, + g, + beta, + initial_state, + output_final_state, + cu_seqlens, + head_first, + use_qk_l2norm_in_kernel, + eps, + ) + + +def fused_sigmoid_gating_delta_rule_update_cpu( + A_log: torch.Tensor, + dt_bias: torch.Tensor, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + a: torch.Tensor, + b: torch.Tensor, + initial_state_source: torch.Tensor, + initial_state_indices: torch.Tensor, + cu_seqlens: torch.Tensor, + use_qk_l2norm_in_kernel: bool, + softplus_beta: float = 1.0, + softplus_threshold: float = 20.0, +) -> torch.Tensor: + return torch.ops._C.fused_sigmoid_gating_delta_rule_update_cpu( + A_log, + dt_bias, + q, + k, + v, + a, + b, + initial_state_source, + initial_state_indices, + cu_seqlens, + use_qk_l2norm_in_kernel, + softplus_beta, + softplus_threshold, + ) + + +def fused_sigmoid_gating_delta_rule_update_spec_cpu( + A_log: torch.Tensor, + dt_bias: torch.Tensor, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + a: torch.Tensor, + b: torch.Tensor, + initial_state_source: torch.Tensor, + spec_state_indices: torch.Tensor, + num_accepted_tokens: torch.Tensor, + cu_seqlens: torch.Tensor, + use_qk_l2norm_in_kernel: bool, + softplus_beta: float = 1.0, + softplus_threshold: float = 20.0, +) -> torch.Tensor: + return torch.ops._C.fused_sigmoid_gating_delta_rule_update_spec_cpu( + A_log, + dt_bias, + q, + k, + v, + a, + b, + initial_state_source, + spec_state_indices, + num_accepted_tokens, + cu_seqlens, + use_qk_l2norm_in_kernel, + softplus_beta, + softplus_threshold, + ) + + +def fused_gdn_gating_cpu( + A_log: torch.Tensor, + a: torch.Tensor, + b: torch.Tensor, + dt_bias: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + return torch.ops._C.fused_gdn_gating_cpu( + A_log, + a, + b, + dt_bias, + ) + + +def causal_conv1d_weight_pack( + weight: torch.Tensor, +) -> torch.Tensor: + return torch.ops._C.causal_conv1d_weight_pack( + weight, + ) + + +def causal_conv1d_fwd_cpu( + x: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor | None, + conv_states: torch.Tensor | None, + query_start_loc: torch.Tensor | None, + cache_indices: torch.Tensor | None, + has_initial_state: torch.Tensor | None, + silu_activation: bool, + is_vnni: bool, +) -> torch.Tensor: + return torch.ops._C.causal_conv1d_fwd_cpu( + x, + weight, + bias, + conv_states, + query_start_loc, + cache_indices, + has_initial_state, + silu_activation, + -1, + is_vnni, + ) + + +def causal_conv1d_update_cpu( + x: torch.Tensor, + conv_states: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor | None, + silu_activation: bool, + conv_state_indices: torch.Tensor | None, + is_vnni: bool, + num_accepted_tokens: torch.Tensor | None = None, +) -> torch.Tensor: + return torch.ops._C.causal_conv1d_update_cpu( + x, + conv_states, + weight, + bias, + silu_activation, + num_accepted_tokens, + conv_state_indices, + -1, + is_vnni, + ) + + +class CPUDNNLGEMMHandler: + def __init__(self) -> None: + self.handler_tensor: torch.Tensor | None = None + self.n = -1 + self.k = -1 + self.dtor = torch.ops._C.release_dnnl_matmul_handler + + def __del__(self): + if self.handler_tensor is not None: + self.dtor(self.handler_tensor.item()) + + +_supports_onednn = bool(hasattr(torch.ops._C, "create_onednn_mm_handler")) + + +def is_onednn_acl_supported(): + return torch.ops._C.is_onednn_acl_supported() + + +def create_onednn_mm( + weight: torch.Tensor, # [K, N] + primitive_cache_size: int = 128, +) -> CPUDNNLGEMMHandler: + handler = CPUDNNLGEMMHandler() + handler.k, handler.n = weight.size() + # store the handler pointer in a tensor it doesn't get inlined + handler.handler_tensor = torch.tensor( + torch.ops._C.create_onednn_mm_handler(weight, primitive_cache_size), + dtype=torch.int64, + ) + return handler + + +def onednn_mm( + dnnl_handler: CPUDNNLGEMMHandler, + x: torch.Tensor, + bias: torch.Tensor | None, +) -> torch.Tensor: + output = torch.empty((*x.shape[0:-1], dnnl_handler.n), dtype=x.dtype) + torch.ops._C.onednn_mm(output, x.reshape(-1, dnnl_handler.k), bias, dnnl_handler.handler_tensor) + + return output + + +def create_onednn_scaled_mm( + weight: torch.Tensor, # [K, N] + weight_scales: torch.Tensor, + output_type: torch.dtype, + dynamic_quant: bool, + use_azp: bool, + primitive_cache_size: int = 128, +) -> CPUDNNLGEMMHandler: + handler = CPUDNNLGEMMHandler() + handler.k, handler.n = weight.size() + # store the handler pointer in a tensor so it doesn't get inlined + handler.handler_tensor = torch.tensor( + torch.ops._C.create_onednn_scaled_mm_handler( + weight, + weight_scales, + output_type, + dynamic_quant, + use_azp, + primitive_cache_size, + ), + dtype=torch.int64, + ) + return handler + + +def onednn_scaled_int8_quant( + input: torch.Tensor, + scale: torch.Tensor | None = None, + azp: torch.Tensor | None = None, + symmetric: bool = True, +): + """ + Quantize the input tensor to int8 and return the quantized tensor and scale, and maybe azp. + + Args: + input: The input tensor to be quantized to int8. + scale: Optional scaling factor for the int8 quantization. + When not provided, we invoke dynamic-per-token quantization. + azp: Optional zero-point for the int8 quantization. + Must be provided for asymmetric quantization if `scale` is provided. + symmetric: Whether to use symmetric quantization (scale only, azp ignored). + + Returns: + tuple[torch.Tensor, torch.Tensor, torch.Tensor | None] : Output int8 tensor, scales, and optionally azp. + """ + output = torch.empty_like(input, dtype=torch.int8) + token_num = input.numel() // input.shape[-1] + input = input.view((token_num, input.shape[-1])) + if scale is not None: + # static-per-tensor quantization. + assert symmetric == (azp is None), "azp must only be provided for asymmetric quantization." + torch.ops._C.static_scaled_int8_quant(output, input, scale, azp) + return output, scale, azp + + # dynamic-per-token quantization. + input_scales = torch.empty((token_num, 1), device=input.device, dtype=torch.float32) + input_azp = None if symmetric else torch.empty_like(input_scales, dtype=torch.int32) + torch.ops._C.dynamic_scaled_int8_quant(output, input, input_scales, input_azp) + return output, input_scales, input_azp + + +def onednn_scaled_mm( + dnnl_handler: CPUDNNLGEMMHandler, + x: torch.Tensor, + output: torch.Tensor, + input_scale: torch.Tensor | None, + input_zp: torch.Tensor | None, + input_zp_adj: torch.Tensor | None, + bias: torch.Tensor | None, +) -> torch.Tensor: + torch.ops._C.onednn_scaled_mm( + output, + x, + input_scale, + input_zp, + input_zp_adj, + bias, + dnnl_handler.handler_tensor, + ) + + return output + + +def cpu_attn_get_scheduler_metadata( + num_reqs: int, + num_heads: int, + num_kv_heads: int, + head_dim: int, + seq_lens: torch.Tensor, + dtype: torch.dtype, + query_start_loc: torch.Tensor, + causal: bool, + sliding_window_size: int, + isa: str, + enable_kv_split: bool, + dynamic_causal: torch.Tensor | None = None, +) -> torch.Tensor: + scheduler_metadata = torch.ops._C.get_scheduler_metadata( + num_reqs, + num_heads, + num_kv_heads, + head_dim, + seq_lens, + dtype, + query_start_loc, + causal, + sliding_window_size, + isa, + enable_kv_split, + dynamic_causal, + ) + return scheduler_metadata + + +def cpu_attn_reshape_and_cache( + key: torch.Tensor, + value: torch.Tensor, + key_cache: torch.Tensor, + value_cache: torch.Tensor, + slot_mapping: torch.Tensor, + isa: str, + k_scale: float = 1.0, + v_scale: float = 1.0, + kv_cache_dtype: str = "auto", +) -> None: + torch.ops._C.cpu_attn_reshape_and_cache( + key, + value, + key_cache, + value_cache, + slot_mapping, + isa, + k_scale, + v_scale, + kv_cache_dtype, + ) + + +def cpu_attention_with_kv_cache( + query: torch.Tensor, + key_cache: torch.Tensor, + value_cache: torch.Tensor, + output: torch.Tensor, + query_start_loc: torch.Tensor, + seq_lens: torch.Tensor, + scale: float, + causal: bool, + alibi_slopes: torch.Tensor | None, + sliding_window: int, + block_table: torch.Tensor, + softcap: float, + scheduler_metadata: torch.Tensor, + s_aux: torch.Tensor | None, + dynamic_causal: torch.Tensor | None = None, + k_scale: float = 1.0, + v_scale: float = 1.0, + kv_cache_dtype: str = "auto", +) -> None: + torch.ops._C.cpu_attention_with_kv_cache( + query, + key_cache, + value_cache, + output, + query_start_loc, + seq_lens, + scale, + causal, + alibi_slopes, + sliding_window, + block_table, + softcap, + scheduler_metadata, + s_aux, + dynamic_causal, + k_scale, + v_scale, + kv_cache_dtype, + ) + + +def cpu_gemm_wna16( + input: torch.Tensor, + q_weight: torch.Tensor, + scales: torch.Tensor, + zeros: torch.Tensor | None, + g_idx: torch.Tensor | None, + bias: torch.Tensor | None, + pack_factor: int, + isa_hint: str, +) -> torch.Tensor: + output = torch.empty((input.size(0), scales.size(1)), dtype=input.dtype) + torch.ops._C.cpu_gemm_wna16( + input, + q_weight, + output, + scales, + zeros, + g_idx, + bias, + pack_factor, + isa_hint, + ) + return output + + +def cpu_activation_lut_bf16(input: torch.Tensor, activation: str) -> torch.Tensor: + out = torch.empty_like(input) + torch.ops._C.activation_lut_bf16(out, input, activation) + return out + + +def cpu_prepack_moe_weight( + weight: torch.Tensor, + isa: str, +) -> torch.Tensor: + output = torch.empty_like(weight) + torch.ops._C.prepack_moe_weight(weight, output, isa) + return output + + +def cpu_prepack_moe_weight_int8( + weight: torch.Tensor, + isa: str, +) -> torch.Tensor: + output = torch.empty_like(weight) + torch.ops._C.prepack_moe_weight_int8(weight, output, isa) + return output + + +def cpu_fused_moe( + input: torch.Tensor, + w13: torch.Tensor, + w2: torch.Tensor, + w13_bias: torch.Tensor | None, + w2_bias: torch.Tensor | None, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + act: str, + isa: str, + skip_weighted: bool = False, +) -> torch.Tensor: + output = torch.empty_like(input) + torch.ops._C.cpu_fused_moe( + output, + input, + w13, + w2, + w13_bias, + w2_bias, + topk_weights, + topk_ids, + skip_weighted, + act, + isa, + ) + return output + + +def cpu_fused_moe_int8( + input: torch.Tensor, + w13: torch.Tensor, + w2: torch.Tensor, + w13_scale: torch.Tensor, + w2_scale: torch.Tensor, + w13_bias: torch.Tensor | None, + w2_bias: torch.Tensor | None, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + act: str, + isa: str, + skip_weighted: bool = False, +) -> torch.Tensor: + output = torch.empty_like(input) + torch.ops._C.cpu_fused_moe_int8( + output, + input, + w13, + w2, + w13_scale, + w2_scale, + w13_bias, + w2_bias, + topk_weights, + topk_ids, + skip_weighted, + act, + isa, + ) + return output + + +if hasattr(torch.ops._qutlass_C, "matmul_mxf4_bf16_tn"): + + @register_fake("_qutlass_C::matmul_mxf4_bf16_tn") + def _fake_matmul_mxf4_bf16_tn( + a: torch.Tensor, + b: torch.Tensor, + a_sf: torch.Tensor, + b_sf: torch.Tensor, + alpha: torch.Tensor, + ): + return a.new_empty(*a.shape[:-1], b.shape[0], dtype=torch.bfloat16) + + +def matmul_mxf4_bf16_tn( + a: torch.Tensor, + b: torch.Tensor, + a_sf: torch.Tensor, + b_sf: torch.Tensor, + alpha: torch.Tensor, +) -> torch.Tensor: + return torch.ops._qutlass_C.matmul_mxf4_bf16_tn(a, b, a_sf, b_sf, alpha) + + +if hasattr(torch.ops._qutlass_C, "fusedQuantizeMxQuest"): + + @register_fake("_qutlass_C::fusedQuantizeMxQuest") + def _fake_fused_quantize_mx_quest(a: torch.Tensor, b: torch.Tensor, xh_e2m1: torch.Tensor, xh_e8m0: torch.Tensor): + return xh_e2m1, xh_e8m0 + + +if hasattr(torch.ops._qutlass_C, "fusedQuantizeMxAbsMax"): + + @register_fake("_qutlass_C::fusedQuantizeMxAbsMax") + def _fake_fused_quantize_mx_absmax(a: torch.Tensor, b: torch.Tensor, xh_e2m1: torch.Tensor, xh_e8m0: torch.Tensor): + return xh_e2m1, xh_e8m0 + + +def fusedQuantizeMx( + a: torch.Tensor, b: torch.Tensor, *, method: Literal["quest", "abs_max"] = "quest" +) -> tuple[torch.Tensor, torch.Tensor]: + if a.dim() == 0: + raise ValueError("`a` must have at least 1 dimension.") + if a.size(-1) % 32 != 0: + raise ValueError(f"last dim of `a` must be divisible by 32, got {a.size(-1)}.") + if b.device != a.device: + raise ValueError("`a` and `b` must be on the same device.") + + xh_e2m1 = torch.empty(*a.shape[:-1], a.size(-1) // 2, dtype=torch.uint8, device=a.device) + + rows, cols = a.numel() // a.size(-1), a.size(-1) // 32 + n_row_blocks = cdiv(rows, 128) + n_col_blocks = cdiv(cols, 4) + padded_rows = n_row_blocks * 128 + padded_cols = n_col_blocks * 4 + + xh_e8m0 = torch.empty(padded_rows, padded_cols, dtype=torch.float8_e8m0fnu, device=a.device) + + if not hasattr(torch.ops, "_qutlass_C"): + raise RuntimeError( + "The `_qutlass_C` extension is not loaded. " + "Make sure your custom op library is imported before calling fusedQuantizeMx." + ) + + if method == "quest": + return torch.ops._qutlass_C.fusedQuantizeMxQuest(a, b, xh_e2m1, xh_e8m0) + elif method == "abs_max": + return torch.ops._qutlass_C.fusedQuantizeMxAbsMax(a, b, xh_e2m1, xh_e8m0) + else: + 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) + + rows, cols = a.numel() // a.size(-1), a.size(-1) // 16 + n_row_blocks = cdiv(rows, 128) + n_col_blocks = cdiv(cols, 4) + padded_rows = n_row_blocks * 128 + 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) + + +def hadacore_transform(x: torch.Tensor, inplace: bool = True) -> torch.Tensor: + """ + Perform Hadamard transforms using [Hadacore](https://arxiv.org/abs/2412.08832) + kernels. Note that these kernels exploit the recursive properties of + Sylvester Hadamards, and therefore do not require transform weight data + + Note that sylvester hadamard transforms are also symmetric, which means that + this function is also applies the (transpose <=> inverse) transform. + + Args: + x: value to be transformed inplace + inplace: modify value in place + + Returns: + value after transformation + """ + return torch.ops._C.hadacore_transform(x, inplace) + + +if hasattr(torch.ops._C, "hadacore_transform"): + + @register_fake("_C::hadacore_transform") + def _hadacore_transform_fake(x: torch.Tensor, inplace: bool) -> torch.Tensor: + return torch.empty_like(x) if not inplace else x + + +if hasattr(torch.ops._C, "minimax_allreduce_rms_qk"): + + @register_fake("_C::minimax_allreduce_rms_qk") + def _minimax_allreduce_rms_qk_fake( + qkv: torch.Tensor, + norm_weight_q: torch.Tensor, + norm_weight_k: torch.Tensor, + workspace: torch.Tensor, + q_size: int, + kv_size: int, + rank: int, + nranks: int, + eps: float, + ) -> tuple[torch.Tensor, torch.Tensor]: + token_num = qkv.shape[0] + return ( + torch.empty([token_num, q_size], dtype=qkv.dtype, device=qkv.device), + torch.empty([token_num, kv_size], dtype=qkv.dtype, device=qkv.device), + ) 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/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/module.py.orig b/aphrodite/model_executor/layers/quantization/compressed_tensors/transform/module.py.orig new file mode 100644 index 0000000000..1561547208 --- /dev/null +++ b/aphrodite/model_executor/layers/quantization/compressed_tensors/transform/module.py.orig @@ -0,0 +1,157 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import math +from collections.abc import Callable, Hashable + +import torch +from compressed_tensors.transform import ( + TransformArgs, + TransformLocation, + TransformScheme, +) +from torch import Tensor + +import aphrodite._custom_ops as ops +from aphrodite.distributed.parallel_state import get_tensor_model_parallel_world_size +from aphrodite.model_executor.layers.linear import LinearBase +from aphrodite.model_executor.layers.quantization.compressed_tensors.transform.utils import ( # noqa: E501 + TransformTuple, +) +from aphrodite.model_executor.layers.utils import dispatch_unquantized_gemm +from aphrodite.model_executor.layers.vocab_parallel_embedding import VocabParallelEmbedding +from aphrodite.model_executor.parameter import SharedWeightParameter + + +class HadamardTransform(torch.nn.Module): + """ + Class which handles weight loading, postprocessing, and application of + transforms. Meant to be used with `CompressedTensorsLinearTransformMethod` + and attention transforms method (not implemented yet) + """ + + 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)) + + def __init__( + self, + transforms: dict[int, TransformTuple], + layer: torch.nn.Module, + weight_loader: Callable, + input_size_per_partition: int, + output_partition_sizes: list[int], + ): + 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") + + # Similar to row/col parallel params, but tensors are separate + # to allow for loading with shared memory + self.weight = SharedWeightParameter(weight_loader=weight_loader) + + # create shared partition data for each partition of the original weight + input_size = input_size_per_partition + for part_index, (_scheme_name, scheme, args) in self.transforms.items(): + output_size = output_partition_sizes[part_index] + weight_size = self._get_weight_size(layer, scheme, args, input_size, output_size) + + data_key = self._get_data_key(scheme, weight_size) + 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 + + # 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)) + + # FUTURE: avoid runtime transpose by processing weights + # prior to apply + + def forward(self, value: Tensor, part_id: int = 0) -> Tensor: + if part_id not in self.weight.partitions: + return value + + # use hadacore if possible + if self.transforms[part_id].scheme.type == "hadamard": + if self.transforms[part_id].scheme.head_dim is not None: + weight_size = self.transforms[part_id].scheme.head_dim + value = value.unflatten(-1, (-1, weight_size)) + value = ops.hadacore_transform(value) + value = value.flatten(-2, -1) + + return value + + # sylvester transforms are symmetric, inv => transpose => original + return ops.hadacore_transform(value) + + # fall back to dense + 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 = value.flatten(-2, -1) + + return value + + return dispatch_unquantized_gemm()(self, value.to(weight.dtype), weight, None).to(value.dtype) * scale + + def _get_data_key(self, scheme: TransformScheme, weight_size: int) -> Hashable: + return (id(scheme), weight_size) + + def _get_weight_size( + self, + layer: torch.nn.Module, + scheme: TransformScheme, + args: TransformArgs, + input_size: int, + output_size: int, + ) -> int: + if scheme.head_dim is not None: + return scheme.head_dim + + if isinstance(layer, LinearBase): + if args.location == TransformLocation.INPUT: + return input_size + + elif args.location == TransformLocation.OUTPUT: + return output_size + + elif isinstance(layer, VocabParallelEmbedding): + if args.location == TransformLocation.INPUT: + return output_size + + elif args.location == TransformLocation.OUTPUT: + return input_size + + raise ValueError() + + def _validate_input_transforms(self): + assert len(self.transforms) > 0 + location = list(self.transforms.values())[0].args.location + + if location == TransformLocation.INPUT: + first_data = self.weight.partitions[0].data + for partition in self.weight.partitions.values(): + if partition.data.data_ptr() != first_data.data_ptr(): + raise ValueError("") 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/qutlass_utils.py.orig b/aphrodite/model_executor/layers/quantization/qutlass_utils.py.orig new file mode 100644 index 0000000000..542c5baac8 --- /dev/null +++ b/aphrodite/model_executor/layers/quantization/qutlass_utils.py.orig @@ -0,0 +1,178 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# +# Modified by Roberto L. Castro (Roberto.LopezCastro@ist.ac.at). +# +# Copied from https://github.com/pytorch/ao/tree/main/torchao/prototype/mx_formats +# +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. + +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. + +from typing import Literal + +import torch +from torch.library import wrap_triton + +from aphrodite.triton_utils import tl, triton +from aphrodite.utils.math_utils import cdiv + + +@triton.jit +def triton_scale_swizzle( + scale_ptr: torch.Tensor, + scale_rows: int, + scale_cols: int, + output_ptr: torch.Tensor, + input_row_stride: int, + output_block_stride: int, + BLOCK_ROWS: tl.constexpr, + BLOCK_COLS: tl.constexpr, +): + """ + Rearranges tensor data from row-major to block-scaled swizzle format. + + Args: + scale_ptr: Pointer to the input scale tensor + scale_rows: Number of rows in the scale tensor + scale_cols: Number of columns in the scale tensor + output_ptr: Pointer to the output tensor + input_row_stride: Stride between rows in the input tensor + output_block_stride: Stride between blocks in the output tensor + BLOCK_ROWS: Number of rows in a tile (compile-time constant) + BLOCK_COLS: Number of columns in a tile (compile-time constant) + """ + pid_row = tl.program_id(0) + pid_col = tl.program_id(1) + + rows = tl.arange(0, BLOCK_ROWS)[:, None] + cols = tl.arange(0, BLOCK_COLS)[None, :] + + # Calculate starting row and column for this tile + start_row = pid_row * BLOCK_ROWS + start_col = pid_col * BLOCK_COLS + global_rows = start_row + rows + global_cols = start_col + cols + + mask = (global_rows < scale_rows) & (global_cols < scale_cols) + + input_scales = tl.load( + scale_ptr + global_rows * input_row_stride + global_cols, + mask=mask, + other=0.0, + ) + + r_div_32 = rows // 32 + r_mod_32 = rows % 32 + + # 2) Rearrange to (32, 4, 4) then to final (32, 16) coordinates + dest_indices = r_mod_32 * 16 + r_div_32 * 4 + cols + + # Flatten + dest_indices_flat = tl.reshape(dest_indices, (BLOCK_ROWS * BLOCK_COLS)) + scales_flat = tl.reshape(input_scales, (BLOCK_ROWS * BLOCK_COLS)) + + # Calculate block offset using provided output block stride + LOCAL_NUMEL = BLOCK_ROWS * BLOCK_COLS + block_offset = pid_col * LOCAL_NUMEL + (pid_row * output_block_stride) + + tl.store( + output_ptr + block_offset + dest_indices_flat, + scales_flat, + ) + + +def triton_mx_block_rearrange(scale_tensor: torch.Tensor) -> torch.Tensor: + """ + Rearranges an E8M0 tensor scale from row-major format to + block-scaled swizzle format. + + This format is suitable for Tmem as described in NVIDIA documentation: + https://docs.nvidia.com/cuda/cublas/index.html#d-block-scaling-factors-layout + + Args: + scale_tensor: Input tensor in row-major format with 8-bit elements + + Returns: + Rearranged tensor in block-scaled swizzle format + """ + assert scale_tensor.element_size() == 1, "Expected element size to be 1 byte (8 bits)" + assert scale_tensor.is_contiguous(), "Input tensor must be contiguous" + + rows, cols = scale_tensor.shape + + # Calculate blocks needed + n_row_blocks = triton.cdiv(rows, 128) + n_col_blocks = triton.cdiv(cols, 4) + padded_rows = n_row_blocks * 128 + padded_cols = n_col_blocks * 4 + + out = scale_tensor.new_empty((padded_rows, padded_cols)) + + # Input stride (for row-major format) + input_row_stride = cols + + # We probably want handle multiple blocks per tile but + # for now keep it simple + BLOCK_ROWS, BLOCK_COLS = 128, 4 + + # Output block stride for the rearranged format + output_block_stride = BLOCK_ROWS * BLOCK_COLS * (padded_cols // BLOCK_COLS) + + grid = lambda META: ( + triton.cdiv(padded_rows, BLOCK_ROWS), + triton.cdiv(padded_cols, BLOCK_COLS), + ) + + wrap_triton(triton_scale_swizzle)[grid]( + scale_tensor.view(torch.uint8), + rows, + cols, + out.view(torch.uint8), + input_row_stride, + output_block_stride, + BLOCK_ROWS=BLOCK_ROWS, + BLOCK_COLS=BLOCK_COLS, + ) + + return out + + +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 + the rearrangement pattern. + + See: + https://docs.nvidia.com/cuda/cublas/index.html#d-block-scaling-factors-layout + + Args: + input_matrix: Input tensor of shape (H, W) + backend: "torch" (PyTorch path) or "triton" (Triton kernel) + + Returns: + Rearranged tensor of shape (32*cdiv(H,128), 16*cdiv(W,4)) + """ + if backend == "triton": + return triton_mx_block_rearrange(input_matrix).flatten() + elif backend != "torch": + raise ValueError(f'backend must be "torch" or "triton", got {backend!r}') + + rows, cols = input_matrix.shape + n_row_blocks = cdiv(rows, 128) + n_col_blocks = cdiv(cols, 4) + + # Calculate the padded shape + padded_rows = n_row_blocks * 128 + padded_cols = n_col_blocks * 4 + + padded = input_matrix + assert (rows, cols) == (padded_rows, padded_cols) + + # Rearrange the blocks + blocks = padded.view(n_row_blocks, 128, n_col_blocks, 4).permute(0, 2, 1, 3) + rearranged = blocks.reshape(-1, 4, 32, 4).transpose(1, 2).reshape(-1, 32, 16) + + return rearranged.flatten() 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/quantization/test_compressed_tensors.py.orig b/tests/quantization/test_compressed_tensors.py.orig new file mode 100644 index 0000000000..2e1775e0b7 --- /dev/null +++ b/tests/quantization/test_compressed_tensors.py.orig @@ -0,0 +1,919 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Test model set-up and weight loading for llmcompressor-quantized models. + +Run `pytest tests/quantization/test_compressed_tensors.py`. +""" + +from contextlib import contextmanager +from unittest.mock import Mock + +import pytest +import torch +from compressed_tensors.quantization import ( + ActivationOrdering, + QuantizationArgs, + QuantizationStrategy, + QuantizationType, +) + +from aphrodite.model_executor.kernels.linear import ( + Fp8BlockScaledMMLinearKernel, +) +from aphrodite.model_executor.layers.fused_moe import UnquantizedFusedMoEMethod +from aphrodite.model_executor.layers.quantization.compressed_tensors.compressed_tensors import ( # noqa: E501 + CompressedTensorsConfig, + CompressedTensorsLinearMethod, + CompressedTensorsW4A4Fp4, + CompressedTensorsW4A4Mxfp4, + CompressedTensorsW4A8Fp8, + CompressedTensorsW8A8Fp8, + CompressedTensorsW8A8Int8, + CompressedTensorsW8A8Mxfp8, + CompressedTensorsW8A16Fp8, + CompressedTensorsWNA8O8Int, + CompressedTensorsWNA16, +) +from aphrodite.model_executor.layers.quantization.compressed_tensors.utils import ( + find_matched_target, +) +from aphrodite.model_executor.layers.quantization.input_quant_fp8 import QuantFP8 +from aphrodite.model_executor.layers.vocab_parallel_embedding import ParallelLMHead +from aphrodite.platforms import current_platform +from aphrodite.v1.attention.backends.fa_utils import get_flash_attn_version +from tests.models.utils import check_logprobs_close + +# AITER only supports per-channel-per-channel INT8 gemm +# and per-tensor-per-tensor INT8 GEMM. +# It does not support mix precision MM and mix quantization scheme. +ROCM_AITER_SUPPORTED_INT8_MODEL = [ + "neuralmagic/Llama-3.2-1B-quantized.w8a8", + "nm-testing/tinyllama-oneshot-w8a8-channel-dynamic-token-v2", +] + +# TritonInt8ScaledMMLinearKernel only supports symmetric quantization. +ROCM_TRITON_SCALED_MM_SUPPORTED_INT8_MODEL = [ + "nm-testing/tinyllama-oneshot-w8w8-test-static-shape-change", + "nm-testing/tinyllama-oneshot-w8-channel-a8-tensor", + "neuralmagic/Llama-3.2-1B-quantized.w8a8", + "nm-testing/tinyllama-oneshot-w8a8-dynamic-token-v2", + "nm-testing/tinyllama-oneshot-w8a8-channel-dynamic-token-v2", +] + + +@pytest.fixture(scope="function", autouse=True) +def enable_pickle(monkeypatch): + """`LLM.apply_model` requires pickling a function.""" + monkeypatch.setenv("APHRODITE_ALLOW_INSECURE_SERIALIZATION", "1") + + +@pytest.mark.parametrize( + "model_args", + [ + ( + "nm-testing/tinyllama-oneshot-w8w8-test-static-shape-change", + "tensor", + QuantizationType.INT, + 2560, + True, + ), + ( + "nm-testing/asym-w8w8-int8-static-per-tensor-tiny-llama", + "tensor", + QuantizationType.INT, + 2560, + False, + ), + ], +) +def test_compressed_tensors_w8a8_static_setup(aphrodite_runner, model_args): + model_path, strategy, quant_type, shape_0, is_symmetric = model_args + + if current_platform.is_rocm() and model_path not in ROCM_TRITON_SCALED_MM_SUPPORTED_INT8_MODEL: + pytest.skip(f"Skip model {model_path} as it is not supported on ROCm.") + + with aphrodite_runner(model_path, enforce_eager=True) as llm: + + def check_model(model): + layer = model.model.layers[0] + + qkv_proj = layer.self_attn.qkv_proj + o_proj = layer.self_attn.o_proj + gate_up_proj = layer.mlp.gate_up_proj + down_proj = layer.mlp.down_proj + + # assert zp for symmetric and asymmetric cases + def zp_valid(zp: torch.Tensor | None): + if is_symmetric: + return zp is None + + return zp is not None and zp.dtype is torch.int32 + + assert zp_valid(qkv_proj.input_zero_point) + assert zp_valid(o_proj.input_zero_point) + assert zp_valid(gate_up_proj.input_zero_point) + assert zp_valid(down_proj.input_zero_point) + + assert isinstance(qkv_proj.quant_method, CompressedTensorsLinearMethod) + assert isinstance(o_proj.quant_method, CompressedTensorsLinearMethod) + assert isinstance(gate_up_proj.quant_method, CompressedTensorsLinearMethod) + assert isinstance(down_proj.quant_method, CompressedTensorsLinearMethod) + assert isinstance(qkv_proj.scheme, CompressedTensorsW8A8Int8) + + assert qkv_proj.scheme.strategy == strategy + assert qkv_proj.scheme.is_static_input_scheme + expected_type = torch.int8 + + assert qkv_proj.weight.dtype is expected_type + assert o_proj.weight.dtype is expected_type + assert gate_up_proj.weight.dtype is expected_type + + if qkv_proj.scheme.strategy == "tensor": + # Make sure it is a channelwise buffer + # After running process_weights_after_loading + assert len(qkv_proj.weight_scale.shape) == 2 + assert qkv_proj.weight_scale.shape[0] == shape_0 + assert qkv_proj.weight_scale.shape[1] == 1 + assert qkv_proj.weight_scale.dtype is torch.float32 + assert qkv_proj.input_scale.dtype is torch.float32 + + llm.apply_model(check_model) + + output = llm.generate_greedy(["Hello my name is"], max_tokens=4) + assert output + + +@pytest.mark.parametrize( + "model_path", + [ + "neuralmagic/Llama-3.2-1B-quantized.w8a8", + ], +) +@pytest.mark.parametrize("max_tokens", [4]) +@pytest.mark.parametrize("num_logprobs", [10]) +@pytest.mark.parametrize("use_aiter", [True, False] if current_platform.is_rocm() else [False]) +def test_compressed_tensors_w8a8_logprobs( + hf_runner, + aphrodite_runner, + example_prompts, + model_path, + max_tokens, + num_logprobs, + use_aiter, + monkeypatch, +): + if current_platform.is_rocm() and model_path not in ROCM_TRITON_SCALED_MM_SUPPORTED_INT8_MODEL: + pytest.skip(f"Skip model {model_path} as it is not supported on ROCm.") + + if use_aiter: + if model_path not in ROCM_AITER_SUPPORTED_INT8_MODEL: + pytest.skip(f"Skip model {model_path} as it is not support by aiter.") + # this will enable APHRODITE_ROCM_USE_AITER_LINEAR + monkeypatch.setenv("APHRODITE_ROCM_USE_AITER", "1") + + dtype = "bfloat16" + + # skip language translation prompt for the static per tensor models + if model_path in ( + "nm-testing/Meta-Llama-3-8B-Instruct-W8A8-Static-Per-Tensor-Sym", + "nm-testing/Meta-Llama-3-8B-Instruct-W8A8-Static-Per-Tensor-Asym", + ): + example_prompts = example_prompts[0:-1] + + with hf_runner(model_path, dtype=dtype) as hf_model: + hf_outputs = hf_model.generate_greedy_logprobs_limit(example_prompts, max_tokens, num_logprobs) + + with aphrodite_runner(model_path, dtype=dtype, enforce_eager=True) as aphrodite_model: + aphrodite_outputs = aphrodite_model.generate_greedy_logprobs(example_prompts, max_tokens, num_logprobs) + + check_logprobs_close( + outputs_0_lst=hf_outputs, + outputs_1_lst=aphrodite_outputs, + name_0="hf", + name_1="aphrodite", + ) + + if current_platform.is_rocm(): + torch.accelerator.synchronize() + + +def test_compressed_tensors_no_enforce_eager(aphrodite_runner): + model_path = "nm-testing/tinyllama-oneshot-w8w8-test-static-shape-change" + with aphrodite_runner(model_path) as llm: + output = llm.generate_greedy("Hello my name is", max_tokens=4) + assert output + + +@pytest.mark.parametrize( + "model_args", + [ + ("nm-testing/tinyllama-oneshot-w8a8-dynamic-token-v2", "tensor"), + ( + "nm-testing/tinyllama-oneshot-w8a8-channel-dynamic-token-v2", + "channel", + ), + ], +) +@pytest.mark.parametrize("use_aiter", [True, False] if current_platform.is_rocm() else [False]) +def test_compressed_tensors_w8a8_dynamic_per_token( + aphrodite_runner, + model_args, + use_aiter, + monkeypatch, +): + model_path, strategy = model_args + + if current_platform.is_rocm() and model_path not in ROCM_TRITON_SCALED_MM_SUPPORTED_INT8_MODEL: + pytest.skip(f"Skip model {model_path} as it is not supported on ROCm.") + + if use_aiter: + if model_path not in ROCM_AITER_SUPPORTED_INT8_MODEL: + pytest.skip(f"Skip model {model_path} as it is not support by aiter.") + # this will enable APHRODITE_ROCM_USE_AITER_LINEAR + monkeypatch.setenv("APHRODITE_ROCM_USE_AITER", "1") + + with aphrodite_runner(model_path, enforce_eager=True, dtype=torch.float16) as llm: + + def check_model(model): + layer = model.model.layers[0] + + qkv_proj = layer.self_attn.qkv_proj + + assert isinstance(qkv_proj.quant_method, CompressedTensorsLinearMethod) + assert isinstance(qkv_proj.scheme, CompressedTensorsW8A8Int8) + assert not qkv_proj.scheme.is_static_input_scheme + assert qkv_proj.scheme.strategy == strategy + assert qkv_proj.weight.dtype is torch.int8 + + llm.apply_model(check_model) + + output = llm.generate_greedy(["Hello my name is"], max_tokens=4) + assert output + + +@pytest.mark.parametrize( + "wNa16_args", + [ + ( + "nm-testing/tinyllama-oneshot-w4a16-channel-v2", + "channel", + None, + 8, + True, + False, + ), + ( + "nm-testing/TinyLlama-1.1B-Chat-v1.0-W4A16-G128-Asym-Updated-ActOrder", + "group", + 128, + 8, + False, + True, + ), + ], +) +@pytest.mark.skipif(not current_platform.is_cuda(), reason="The tests are skipped on non-CUDA platform.") +def test_compressed_tensors_wNa16(aphrodite_runner, wNa16_args): + model, strategy, group, pack_factor, symmetric, has_g_idx = wNa16_args + with aphrodite_runner(model, enforce_eager=True) as llm: + + def check_model(model): + layer = model.model.layers[0] + + qkv_proj = layer.self_attn.qkv_proj + assert isinstance(qkv_proj.quant_method, CompressedTensorsLinearMethod) + assert isinstance(qkv_proj.scheme, CompressedTensorsWNA16) + + assert qkv_proj.scheme.strategy == strategy + assert qkv_proj.scheme.group_size == (-1 if group is None else group) + + assert qkv_proj.scheme.pack_factor == pack_factor + assert qkv_proj.scheme.symmetric == symmetric + assert qkv_proj.scheme.has_g_idx == has_g_idx + + llm.apply_model(check_model) + + output = llm.generate_greedy("Hello my name is", max_tokens=4) + assert output + + +def test_compressed_tensors_fp8(aphrodite_runner): + model_path = "nm-testing/Meta-Llama-3-8B-FP8-compressed-tensors-test" + with aphrodite_runner(model_path, enforce_eager=True) as llm: + + def check_model(model): + layer = model.model.layers[0] + + qkv_proj = layer.self_attn.qkv_proj + + assert isinstance(qkv_proj.quant_method, CompressedTensorsLinearMethod) + assert isinstance( + qkv_proj.scheme, + (CompressedTensorsW8A8Fp8, CompressedTensorsW8A16Fp8), + ) + + assert qkv_proj.input_scale.dtype is torch.float32 + + if isinstance(qkv_proj.scheme, CompressedTensorsW8A8Fp8): + assert len(qkv_proj.input_scale.shape) == 0 + assert qkv_proj.weight.dtype is current_platform.fp8_dtype() + assert qkv_proj.weight_scale.dtype is torch.float32 + assert len(qkv_proj.weight_scale.shape) == 0 + + llm.apply_model(check_model) + + output = llm.generate_greedy("Hello my name is", max_tokens=4) + assert output + + +@pytest.mark.skipif(not current_platform.is_cuda(), reason="This test is skipped on non-CUDA platform.") +def test_compressed_tensors_kv_cache_fp8_per_tensor(aphrodite_runner): + model_path = "nm-testing/TinyLlama-1.1B-Chat-v1.0-kvcache-fp8-tensor" + with aphrodite_runner(model_path) as llm: + output = llm.generate_greedy("Hello world!", max_tokens=4) + assert output + + +@pytest.mark.skipif(not current_platform.is_cuda(), reason="This test is skipped on non-CUDA platform.") +def test_compressed_tensors_kv_cache_fp8_per_attn_head(aphrodite_runner): + model_path = "nm-testing/TinyLlama-1.1B-Chat-v1.0-kvcache-fp8-attn_head" + try: + fa_version = get_flash_attn_version() + except Exception: + pytest.skip("This test requires FlashAttention backend.") + if fa_version is None or fa_version < 3: + pytest.skip("This test requires FlashAttention version >= 3.") + + with aphrodite_runner(model_path, attention_config={"backend": "FLASH_ATTN"}) as llm: + output = llm.generate_greedy("Hello world!", max_tokens=4) + assert output + + +@contextmanager +def _nvfp4_marlin_error_context(model, capfd): + is_rocm_and_unsupported = model == "nm-testing/TinyLlama-1.1B-Chat-v1.0-NVFP4A16" and current_platform.is_rocm() + + if is_rocm_and_unsupported: + expected_error = ( + "ValueError: Forced NVFP4 kernel MarlinNvFp4LinearKernel is not supported: Marlin FP4 not available" + ) + with pytest.raises(RuntimeError, match="Engine core initialization failed"): + yield + + captured = capfd.readouterr() + assert expected_error in captured.out + captured.err + else: + yield + + +@pytest.mark.parametrize( + "args", + [ + ("nm-testing/TinyLlama-1.1B-Chat-v1.0-NVFP4A16", True), + ("nm-testing/TinyLlama-1.1B-Chat-v1.0-NVFP4", False), + ], +) +def test_compressed_tensors_nvfp4(aphrodite_runner, args, capfd): + model, use_a16 = args + with ( + _nvfp4_marlin_error_context(model, capfd), + aphrodite_runner(model, enforce_eager=True) as llm, + ): + + def check_model(model): + layer = model.model.layers[0] + + qkv_proj = layer.self_attn.qkv_proj + assert isinstance(qkv_proj.quant_method, CompressedTensorsLinearMethod) + assert isinstance(qkv_proj.scheme, CompressedTensorsW4A4Fp4) + assert qkv_proj.scheme.use_a16 == use_a16 + assert qkv_proj.scheme.group_size == 16 + + llm.apply_model(check_model) + output = llm.generate_greedy(["Hello my name is"], max_tokens=4) + print(output) + assert output + + +@pytest.mark.skipif( + not current_platform.is_cuda() or not current_platform.has_device_capability(90), + reason="W4A8 FP8 is not yet supported on this GPU type.", +) +@pytest.mark.parametrize( + "args", + [("czhu-cohere/TinyLlama-1.1B-Chat-v1.0-W4A8-e2e", CompressedTensorsW4A8Fp8)], +) +def test_compressed_tensors_w4a8_fp8(aphrodite_runner, args): + model, scheme = args + with aphrodite_runner(model, enforce_eager=True) as llm: + + def check_model(model): + layer = model.model.layers[0] + + qkv_proj = layer.self_attn.qkv_proj + o_proj = layer.self_attn.o_proj + gate_up_proj = layer.mlp.gate_up_proj + down_proj = layer.mlp.down_proj + + for proj in (qkv_proj, o_proj, gate_up_proj, down_proj): + assert isinstance(proj.quant_method, CompressedTensorsLinearMethod) + assert isinstance(proj.scheme, scheme) + + assert proj.weight_packed.dtype is torch.int32 + assert proj.weight_scale.dtype is torch.float8_e4m3fn + assert proj.weight_chan_scale.dtype is torch.float32 + assert proj.scheme.group_size == 128 + + llm.apply_model(check_model) + output = llm.generate_greedy("Hello my name is", max_tokens=4) + print(output) + assert output + + +@pytest.mark.skipif(not current_platform.is_cuda(), reason="This test is skipped on non-CUDA platform.") +@pytest.mark.parametrize( + "model,prompt,exp_perplexity", + [ + ( + "nm-testing/Llama-3.2-1B-Instruct-spinquantR1R2R4-w4a16", + "Flat is better than nested.\nSparse is better than dense.", + 150.0, + ), + ( + "nm-testing/Llama-3.2-1B-Instruct-quip-w4a16", + "Flat is better than nested.\nSparse is better than dense.", + 150.0, + ), + ], +) +def test_compressed_tensors_transforms_perplexity(aphrodite_runner, model, prompt, exp_perplexity): + with aphrodite_runner(model, enforce_eager=True) as llm: + perplexity = llm.generate_prompt_perplexity([prompt])[0] + print(perplexity) + assert perplexity <= exp_perplexity + + +def test_compressed_tensors_fp8_block_enabled(aphrodite_runner): + model_path = "RedHatAI/Qwen3-0.6B-FP8-BLOCK" + with aphrodite_runner(model_path, enforce_eager=True) as llm: + fp8_dtype = current_platform.fp8_dtype() + + def check_model(model): + layer = model.model.layers[0] + + qkv_proj = layer.self_attn.qkv_proj + assert isinstance(qkv_proj.quant_method, CompressedTensorsLinearMethod) + assert isinstance(qkv_proj.scheme, CompressedTensorsW8A8Fp8) + assert isinstance(qkv_proj.scheme.fp8_linear, Fp8BlockScaledMMLinearKernel) + + assert qkv_proj.weight.dtype is fp8_dtype + assert qkv_proj.weight_scale.dtype is torch.float32 + assert len(qkv_proj.weight.shape) == 2 + assert len(qkv_proj.weight_scale.shape) == 2 + + input_quant_op = qkv_proj.scheme.fp8_linear.quant_fp8 + assert isinstance(input_quant_op, QuantFP8) + assert input_quant_op._forward_method in ( + input_quant_op.forward_cuda, + input_quant_op.forward_hip, + input_quant_op.forward_xpu, + ) + + llm.apply_model(check_model) + + output = llm.generate_greedy("Hello my name is", max_tokens=4) + assert output + + +@pytest.mark.skipif( + not current_platform.is_cuda(), + reason="This test is not for non-CUDA platforms", +) +def test_compressed_tensors_moe_ignore_with_model(aphrodite_runner): + """ + Integration test for MoE layer ignore functionality with a real model. + + This test would verify that when loading a compressed-tensors quantized + MoE model where some MoE layers are in the ignore list, those layers + use UnquantizedFusedMoEMethod while non-ignored layers use the + quantized method. + + Expected model structure: + - Compressed-tensors quantized MoE model (e.g., Mixtral-based) + - Config with ignore list containing specific MoE layers + - Multiple MoE layers where some are quantized and some are not + """ + + # model_path = "nm-testing/tinysmokeqwen3moe-W4A16-first-only" # CT 12.3 + model_path = "nm-testing/tinysmokeqwen3moe-W4A16-first-only-CTstable" # CT 12.2 + + with aphrodite_runner(model_path, enforce_eager=True) as llm: + + def check_model(model): + from aphrodite.model_executor.layers.fused_moe import MoERunner + from aphrodite.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe import ( # noqa: E501 + CompressedTensorsMoEMethod, + ) + + # Check layer 0 MoE (should be quantized) + layer_quantized = model.model.layers[0].mlp.experts + assert isinstance(layer_quantized, MoERunner) + assert isinstance(layer_quantized._quant_method, CompressedTensorsMoEMethod) + + # Check layer 10 MoE (should be unquantized + ignored) + layer_unquantized = model.model.layers[3].mlp.experts + assert isinstance(layer_unquantized, MoERunner) + assert isinstance(layer_unquantized._quant_method, UnquantizedFusedMoEMethod) + + llm.apply_model(check_model) + + # Verify the model can generate output + output = llm.generate_greedy("Hello, my name is", max_tokens=4) + assert output + + +def test_w4a16_moe_torch_compile(aphrodite_runner): + """Regression test: MoE quant_config must be initialized inside the + moe_forward custom op, not just in forward_native which is compiled by + Dynamo (attribute mutations are not replayed at runtime). + + Without the fix in _moe_forward/_moe_forward_shared, this hits: + AssertionError: Hidden size mismatch 2048 != 1024 + because use_int4_w4a16 is False (moe_quant_config stays None). + """ + model_path = "nm-testing/tinysmokeqwen3moe-W4A16-first-only-CTstable" + + with aphrodite_runner( + model_path, + enforce_eager=False, + max_model_len=256, + compilation_config={ + "cudagraph_mode": "NONE", + }, + ) as llm: + output = llm.generate_greedy("Hi", max_tokens=1) + assert output + + +def _make_ct_config(*, target: str = "Linear") -> CompressedTensorsConfig: + """Build a minimal CompressedTensorsConfig with INT8 channel quant.""" + weight_quant = QuantizationArgs( + num_bits=8, + type=QuantizationType.INT, + strategy=QuantizationStrategy.CHANNEL, + symmetric=True, + dynamic=False, + ) + return CompressedTensorsConfig( + target_scheme_map={ + target: { + "weights": weight_quant, + "input_activations": None, + "format": "pack-quantized", + } + }, + ignore=[], + quant_format="pack-quantized", + ) + + +def test_get_quant_method_returns_linear_method_for_parallel_lm_head(): + """ParallelLMHead whose name matches a target must get a quantised method.""" + config = _make_ct_config(target="re:.*lm_head") + mock_lm_head = Mock(spec=ParallelLMHead) + mock_lm_head.__class__ = ParallelLMHead + + method = config.get_quant_method(mock_lm_head, prefix="model.lm_head") + + assert isinstance(method, CompressedTensorsLinearMethod), ( + f"Expected CompressedTensorsLinearMethod, got {type(method).__name__}" + ) + + +def test_get_quant_method_returns_none_for_ignored_parallel_lm_head(): + """ParallelLMHead on the ignore list should be left unquantized (None).""" + config = _make_ct_config(target="re:.*lm_head") + config.ignore = ["re:.*lm_head"] + mock_lm_head = Mock(spec=ParallelLMHead) + mock_lm_head.__class__ = ParallelLMHead + + method = config.get_quant_method(mock_lm_head, prefix="model.lm_head") + + assert method is None, f"Expected None for ignored ParallelLMHead, got {type(method).__name__}" + + +def test_get_quant_method_returns_none_for_unmatched_parallel_lm_head(): + """ParallelLMHead with target='Linear' (typical real model) must not crash. + + Most compressed-tensors models only target 'Linear'. ParallelLMHead does + not match that target, so get_quant_method should return None (unquantized) + instead of raising ValueError. + """ + config = _make_ct_config(target="Linear") + mock_lm_head = Mock(spec=ParallelLMHead) + mock_lm_head.__class__ = ParallelLMHead + + method = config.get_quant_method(mock_lm_head, prefix="model.lm_head") + + assert method is None, f"Expected None for unmatched ParallelLMHead, got {type(method).__name__}" + + +def test_find_matched_target_returns_none_on_no_match(): + result = find_matched_target( + layer_name="model.layers.0.self_attn.qkv_proj", + module=Mock(spec=torch.nn.Linear), + targets=["no_match_target"], + ) + assert result is None + + +def test_get_scheme_dict_returns_none_on_no_match(): + config = _make_ct_config(target="matched_layer") + result = config.get_scheme_dict( + layer=Mock(spec=torch.nn.Linear), + layer_name="model.layers.0.unmatched_layer", + ) + assert result is None + + +# Test constants for activation quantization +_STATIC_SYM_INT8_ACT = QuantizationArgs( + num_bits=8, + type=QuantizationType.INT, + strategy=QuantizationStrategy.TENSOR.value, + symmetric=True, + dynamic=False, +) + +_STATIC_ASYM_INT8_ACT = QuantizationArgs( + num_bits=8, + type=QuantizationType.INT, + strategy=QuantizationStrategy.TENSOR.value, + symmetric=False, + dynamic=False, +) + +_DYNAMIC_INT8_ACT = QuantizationArgs( + num_bits=8, + type=QuantizationType.INT, + strategy=QuantizationStrategy.TOKEN.value, + symmetric=True, + dynamic=True, +) + + +@pytest.mark.parametrize( + "weight_bits,weight_strategy,input_act,output_act,format,expected_scheme", + [ + # W8A8 int-quantized -> W8A8Int8 (regression test for #46389) + pytest.param( + 8, + QuantizationStrategy.CHANNEL.value, + _STATIC_SYM_INT8_ACT, + None, + "int-quantized", + CompressedTensorsW8A8Int8, + id="w8a8_channel_static_sym", + ), + pytest.param( + 8, + QuantizationStrategy.CHANNEL.value, + _STATIC_ASYM_INT8_ACT, + None, + "int-quantized", + CompressedTensorsW8A8Int8, + id="w8a8_channel_static_asym", + ), + pytest.param( + 8, + QuantizationStrategy.TENSOR.value, + _STATIC_SYM_INT8_ACT, + None, + "int-quantized", + CompressedTensorsW8A8Int8, + id="w8a8_tensor_static", + ), + pytest.param( + 8, + QuantizationStrategy.CHANNEL.value, + _DYNAMIC_INT8_ACT, + None, + "int-quantized", + CompressedTensorsW8A8Int8, + id="w8a8_channel_dynamic", + ), + # W8A8O8 int-quantized -> WNA8O8Int (both input and output) + pytest.param( + 8, + QuantizationStrategy.CHANNEL.value, + _STATIC_SYM_INT8_ACT, + _STATIC_SYM_INT8_ACT, + "int-quantized", + CompressedTensorsWNA8O8Int, + id="w8a8o8_channel", + ), + pytest.param( + 4, + QuantizationStrategy.GROUP.value, + _STATIC_SYM_INT8_ACT, + _STATIC_SYM_INT8_ACT, + "int-quantized", + CompressedTensorsWNA8O8Int, + id="w4a8o8_group", + ), + # Weight-only pack-quantized -> WNA16 + pytest.param( + 8, + QuantizationStrategy.CHANNEL.value, + None, + None, + "pack-quantized", + CompressedTensorsWNA16, + id="w8_pack", + ), + pytest.param( + 4, + QuantizationStrategy.GROUP.value, + None, + None, + "pack-quantized", + CompressedTensorsWNA16, + id="w4_pack", + ), + pytest.param( + 2, + QuantizationStrategy.GROUP.value, + None, + None, + "pack-quantized", + CompressedTensorsWNA16, + id="w2_pack", + ), + pytest.param( + 3, + QuantizationStrategy.GROUP.value, + None, + None, + "pack-quantized", + CompressedTensorsWNA16, + id="w3_pack", + ), + pytest.param( + 5, + QuantizationStrategy.GROUP.value, + None, + None, + "pack-quantized", + CompressedTensorsWNA16, + id="w5_pack", + ), + pytest.param( + 6, + QuantizationStrategy.GROUP.value, + None, + None, + "pack-quantized", + CompressedTensorsWNA16, + id="w6_pack", + ), + pytest.param( + 7, + QuantizationStrategy.GROUP.value, + None, + None, + "pack-quantized", + CompressedTensorsWNA16, + id="w7_pack", + ), + ], +) +def test_scheme_selection(weight_bits, weight_strategy, input_act, output_act, format, expected_scheme): + """Test that _get_scheme_from_parts selects the correct scheme. + + This parametrized test verifies scheme selection for various combinations + of weight bits, quantization strategies, input/output activations, and + compression formats. + + Key regression test: W8A8 int-quantized models with channel-wise weights + should use W8A8Int8 (true int8 gemm), not WNA8O8Int (fake-quant). + WNA8O8Int should only match when BOTH input and output activations are + present. + """ + weight_quant = QuantizationArgs( + num_bits=weight_bits, + type=QuantizationType.INT, + strategy=weight_strategy, + symmetric=True, + dynamic=False, + group_size=128 if weight_strategy == QuantizationStrategy.GROUP.value else None, + ) + + config = CompressedTensorsConfig( + target_scheme_map={}, + ignore=[], + quant_format=format, + ) + + scheme = config._get_scheme_from_parts( + weight_quant=weight_quant, + input_quant=input_act, + output_quant=output_act, + format=format, + ) + + assert isinstance(scheme, expected_scheme), ( + f"Expected {expected_scheme.__name__} for " + f"W{weight_bits} {weight_strategy} + " + f"input_act={input_act} + output_act={output_act} + " + f"format={format}, got {type(scheme).__name__}" + ) + + +@pytest.mark.skipif( + not current_platform.is_cuda() or not current_platform.has_device_capability(75), + reason="MXFP8 requires Turing (sm_75+) or newer.", +) +def test_compressed_tensors_mxfp8_moe_setup(aphrodite_runner): + """Verify MXFP8 scheme, dtypes, and generation for a MoE model.""" + model_path = "AliEdalati97/Qwen3-30B-A3B-MXFP8" + with aphrodite_runner( + model_path, + enforce_eager=True, + load_format="dummy", + hf_overrides={"num_hidden_layers": 4}, + ) as llm: + + def check_model(model): + from aphrodite.model_executor.layers.fused_moe import MoERunner + from aphrodite.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe.compressed_tensors_moe_w8a8_mxfp8 import ( # noqa: E501 + CompressedTensorsW8A8Mxfp8MoEMethod, + ) + + layer = model.model.layers[0] + + qkv = layer.self_attn.qkv_proj + assert isinstance(qkv.quant_method, CompressedTensorsLinearMethod) + assert isinstance(qkv.scheme, CompressedTensorsW8A8Mxfp8) + + experts = layer.mlp.experts + assert isinstance(experts, MoERunner) + assert isinstance(experts._quant_method, CompressedTensorsW8A8Mxfp8MoEMethod) + + llm.apply_model(check_model) + output = llm.generate_greedy("Hello my name is", max_tokens=4) + assert output + + +@pytest.mark.parametrize( + "actorder,group_size,part,full,expected", + [ + # actorder="group" with real grouping: must load full-K w2 scales and, + # when sharded (part != full), report is_k_full=False. + (ActivationOrdering.GROUP, 32, 64, 128, (True, 128, False)), + # actorder="group" but unsharded (part == full): full scales, k_full. + (ActivationOrdering.GROUP, 32, 128, 128, (True, 128, True)), + # actorder="group" with channel-wise (group_size == -1): no full load. + (ActivationOrdering.GROUP, -1, 64, 128, (False, 64, False)), + # "static"/"weight" reorder at quant time -> shard normally + k_full. + # Regression: static actorder under TP must keep is_k_full=True so the + # Marlin kernel never gets the invalid (group_size=16, is_k_full=0). + ("static", 32, 64, 128, (False, 64, True)), + ("weight", 32, 64, 128, (False, 64, True)), + (None, 32, 64, 128, (False, 64, True)), + ], +) +def test_wna16_marlin_moe_w2_scale_sharding(actorder, group_size, part, full, expected): + from aphrodite.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe.compressed_tensors_moe_wna16_marlin import ( # noqa: E501 + CompressedTensorsWNA16MarlinMoEMethod, + ) + + result = CompressedTensorsWNA16MarlinMoEMethod._w2_scale_sharding(actorder, group_size, part, full) + assert result == expected + + +@pytest.mark.skipif( + not current_platform.is_cuda() or not current_platform.has_device_capability(80), + reason="MXFP4 requires ampere or newer", +) +def test_compressed_tensors_mxfp4(aphrodite_runner): + model_path = "nm-testing/TinyLlama-1.1B-Chat-v1.0-MXFP4" + with aphrodite_runner(model_path, enforce_eager=True) as llm: + + def check_model(model): + layer = model.model.layers[0] + + qkv_proj = layer.self_attn.qkv_proj + o_proj = layer.self_attn.o_proj + gate_up_proj = layer.mlp.gate_up_proj + down_proj = layer.mlp.down_proj + + for proj in (qkv_proj, o_proj, gate_up_proj, down_proj): + assert isinstance(proj.quant_method, CompressedTensorsLinearMethod) + assert isinstance(proj.scheme, CompressedTensorsW4A4Mxfp4) + + # Verify group size + assert proj.scheme.group_size == 32 + + llm.apply_model(check_model) + output = llm.generate_greedy("Hello my name is", max_tokens=4) + assert output From 7036fc3cfec9ef6fe135a9c90041f38a24c8adf8 Mon Sep 17 00:00:00 2001 From: AlpinDale Date: Wed, 29 Jul 2026 10:51:31 +0000 Subject: [PATCH 18/41] [sync] Integrate CuTeDSL MoE for ReLU2 NVFP4 (#49580) Upstream-vLLM: 32a423ac0aad67f94f93e97f73338f484b55faec Co-authored-by: danielafrimi <45691845+danielafrimi@users.noreply.github.com> --- .sync/vllm-sha | 2 +- .../fused_moe/experts/flashinfer_cutedsl_moe.py | 8 ++++++-- .../quantization/utils/flashinfer_fp4_moe.py | 16 +++++++++------- 3 files changed, 16 insertions(+), 10 deletions(-) diff --git a/.sync/vllm-sha b/.sync/vllm-sha index c57e38d54e..27b3705b29 100644 --- a/.sync/vllm-sha +++ b/.sync/vllm-sha @@ -1 +1 @@ -30c2718eaafc230927f0d8ac47439d040b46a7c3 +32a423ac0aad67f94f93e97f73338f484b55faec 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/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) From 64ced6344b90ea0f3736fe9a75ced5b1665141ea Mon Sep 17 00:00:00 2001 From: AlpinDale Date: Wed, 29 Jul 2026 10:56:15 +0000 Subject: [PATCH 19/41] [sync] [KV Connector] Support NIXL P/D for hybrid MLA+SSM models (#49762) Upstream-vLLM: f37f03db4af69e4969242767734db3d9175055f7 Co-authored-by: Nick Hill --- .sync/vllm-sha | 2 +- .../kv_connector/v1/nixl/base_worker.py | 26 ++++- .../kv_connector/v1/nixl/push_worker.py | 64 ++++++------ .../unit/test_nixl_connector_hma.py | 99 +++++++++++++++++++ .../unit/test_nixl_desc_geometry.py | 17 ++++ .../unit/test_nixl_push_connector.py | 4 + 6 files changed, 179 insertions(+), 33 deletions(-) diff --git a/.sync/vllm-sha b/.sync/vllm-sha index 27b3705b29..16b8ea1455 100644 --- a/.sync/vllm-sha +++ b/.sync/vllm-sha @@ -1 +1 @@ -32a423ac0aad67f94f93e97f73338f484b55faec +f37f03db4af69e4969242767734db3d9175055f7 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/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]]] = [] From 7674166e616f4543c7936d855f7f45ba646e102c Mon Sep 17 00:00:00 2001 From: AlpinDale Date: Wed, 29 Jul 2026 11:02:35 +0000 Subject: [PATCH 20/41] [sync] [Frontend][Core] Standardize request error handling with VLLMError hierarchy (#49665) Upstream-vLLM: 58f96593971b6287cdc8867024a6dfb9985c545c Co-authored-by: Zach Zhu --- .sync/vllm-sha | 2 +- aphrodite/entrypoints/openai/api_server.py | 24 +++---------- .../entrypoints/openai/engine/protocol.py | 4 +-- .../entrypoints/serve/utils/error_response.py | 20 ++++++++--- .../entrypoints/serve/utils/server_utils.py | 11 +++++- aphrodite/exceptions.py | 21 ++++++++++-- aphrodite/inputs/engine.py | 6 ++-- aphrodite/pooling_params.py | 13 +++---- aphrodite/sampling_params.py | 34 +++++++++---------- aphrodite/v1/engine/async_llm.py | 11 +++--- aphrodite/v1/engine/exceptions.py | 7 ++-- aphrodite/v1/engine/input_processor.py | 27 ++++++++------- .../v1/sample/logits_processor/__init__.py | 6 +++- .../v1/sample/logits_processor/interface.py | 5 ++- 14 files changed, 113 insertions(+), 78 deletions(-) diff --git a/.sync/vllm-sha b/.sync/vllm-sha index 16b8ea1455..da8e336ab0 100644 --- a/.sync/vllm-sha +++ b/.sync/vllm-sha @@ -1 +1 @@ -f37f03db4af69e4969242767734db3d9175055f7 +58f96593971b6287cdc8867024a6dfb9985c545c 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/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/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/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/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/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/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/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 From c430ed053c804d9a9e3f4bdc7e3465d568539108 Mon Sep 17 00:00:00 2001 From: AlpinDale Date: Wed, 29 Jul 2026 11:04:25 +0000 Subject: [PATCH 21/41] [sync] [CI][ROCm] Stabilize Qwen2-VL LoRA test (#50161) Upstream-vLLM: 0bb548b60e88b74304008817b4908b65ec8b2393 Co-authored-by: Andreas Karatzas --- .sync/vllm-sha | 2 +- tests/lora/test_qwenvl.py | 22 +++++++++++++++++++++- 2 files changed, 22 insertions(+), 2 deletions(-) diff --git a/.sync/vllm-sha b/.sync/vllm-sha index da8e336ab0..513f34de28 100644 --- a/.sync/vllm-sha +++ b/.sync/vllm-sha @@ -1 +1 @@ -58f96593971b6287cdc8867024a6dfb9985c545c +0bb548b60e88b74304008817b4908b65ec8b2393 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 From 4c18769ef7473ac43e16b12fd44b2dac20ea7aba Mon Sep 17 00:00:00 2001 From: AlpinDale Date: Wed, 29 Jul 2026 11:09:40 +0000 Subject: [PATCH 22/41] [sync] Add CachePolicyFactory for pluggable/external eviction policies (#49114) Upstream-vLLM: dc1be79031d948d7a18c37600881e45ca708d913 Co-authored-by: Philip Pesic <116986036+philippesic@users.noreply.github.com> --- .sync/vllm-sha | 2 +- aphrodite/v1/kv_offload/base.py | 8 -- aphrodite/v1/kv_offload/cpu/manager.py | 20 ++-- aphrodite/v1/kv_offload/cpu/policies/arc.py | 2 +- aphrodite/v1/kv_offload/cpu/policies/base.py | 4 +- .../v1/kv_offload/cpu/policies/factory.py | 58 ++++++++++++ aphrodite/v1/kv_offload/cpu/policies/lru.py | 1 + aphrodite/v1/kv_offload/cpu/spec.py | 4 +- aphrodite/v1/kv_offload/factory.py | 7 ++ aphrodite/v1/kv_offload/tiering/manager.py | 4 +- aphrodite/v1/kv_offload/tiering/spec.py | 12 ++- tests/v1/kv_offload/cpu/__init__.py | 1 + tests/v1/kv_offload/cpu/policies/__init__.py | 1 + .../kv_offload/cpu/policies/test_factory.py | 94 +++++++++++++++++++ tests/v1/kv_offload/cpu/test_manager.py | 2 + 15 files changed, 190 insertions(+), 30 deletions(-) create mode 100644 aphrodite/v1/kv_offload/cpu/policies/factory.py create mode 100644 tests/v1/kv_offload/cpu/__init__.py create mode 100644 tests/v1/kv_offload/cpu/policies/__init__.py create mode 100644 tests/v1/kv_offload/cpu/policies/test_factory.py diff --git a/.sync/vllm-sha b/.sync/vllm-sha index 513f34de28..2406205fa8 100644 --- a/.sync/vllm-sha +++ b/.sync/vllm-sha @@ -1 +1 @@ -0bb548b60e88b74304008817b4908b65ec8b2393 +dc1be79031d948d7a18c37600881e45ca708d913 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/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..ace77c3b73 --- /dev/null +++ b/tests/v1/kv_offload/cpu/policies/test_factory.py @@ -0,0 +1,94 @@ +# 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, From 8dabf8930d57dd6ee5dae8bdc3b7a767e5655092 Mon Sep 17 00:00:00 2001 From: AlpinDale Date: Wed, 29 Jul 2026 11:12:16 +0000 Subject: [PATCH 23/41] [sync] [CPU] Fix s390x builds and update torch version in dockerfile (#50144) Upstream-vLLM: db7a79cbf74ec0b21e009faa7f1b563f1f88aeb2 Co-authored-by: Rehan Khan --- .sync/vllm-sha | 2 +- csrc/cpu/cpu_types_vxe.hpp | 8 ++++---- tests/v1/kv_offload/cpu/policies/test_factory.py | 4 +--- 3 files changed, 6 insertions(+), 8 deletions(-) diff --git a/.sync/vllm-sha b/.sync/vllm-sha index 2406205fa8..f625ed0d9e 100644 --- a/.sync/vllm-sha +++ b/.sync/vllm-sha @@ -1 +1 @@ -dc1be79031d948d7a18c37600881e45ca708d913 +db7a79cbf74ec0b21e009faa7f1b563f1f88aeb2 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/tests/v1/kv_offload/cpu/policies/test_factory.py b/tests/v1/kv_offload/cpu/policies/test_factory.py index ace77c3b73..1cb8318d63 100644 --- a/tests/v1/kv_offload/cpu/policies/test_factory.py +++ b/tests/v1/kv_offload/cpu/policies/test_factory.py @@ -28,9 +28,7 @@ def remove(self, key: OffloadKey) -> None: def touch(self, keys: Iterable[OffloadKey], req_context: ReqContext) -> None: pass - def evict( - self, n: int, protected: set[OffloadKey] - ) -> list[tuple[OffloadKey, BlockStatus]] | None: + def evict(self, n: int, protected: set[OffloadKey]) -> list[tuple[OffloadKey, BlockStatus]] | None: return None def clear(self) -> None: From 1d6f627d9ed23856d108eff13a781a1d24b62028 Mon Sep 17 00:00:00 2001 From: AlpinDale Date: Wed, 29 Jul 2026 11:14:33 +0000 Subject: [PATCH 24/41] [sync] [Test] dynamic_shapes_compilation (#49974) Upstream-vLLM: 6f91edf96d3f3272945809c04702380053bff4de Co-authored-by: Jared Wen --- .sync/vllm-sha | 2 +- .../test_dynamic_shapes_compilation.py | 102 ++++++++---------- 2 files changed, 45 insertions(+), 59 deletions(-) diff --git a/.sync/vllm-sha b/.sync/vllm-sha index f625ed0d9e..da016db90c 100644 --- a/.sync/vllm-sha +++ b/.sync/vllm-sha @@ -1 +1 @@ -db7a79cbf74ec0b21e009faa7f1b563f1f88aeb2 +6f91edf96d3f3272945809c04702380053bff4de 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" From 0c136e3d616c953c4243758b11530eb88d4259d6 Mon Sep 17 00:00:00 2001 From: AlpinDale Date: Wed, 29 Jul 2026 11:18:05 +0000 Subject: [PATCH 25/41] [sync] fused_moe: add VLLM_TRITON_USE_TD tensor-descriptor path (#42436) Upstream-vLLM: 6f00a1ae3bd4b86168667bce673998218f461c0f Co-authored-by: Artur Fierka --- .sync/vllm-sha | 2 +- .../layers/fused_moe/fused_moe.py | 75 +++++++++++++++---- .../layers/fused_moe/oracle/unquantized.py | 7 ++ .../model_executor/layers/fused_moe/utils.py | 53 +++++++++++++ tests/kernels/moe/test_moe.py | 31 ++++++-- 5 files changed, 147 insertions(+), 21 deletions(-) diff --git a/.sync/vllm-sha b/.sync/vllm-sha index da016db90c..8204a3f27f 100644 --- a/.sync/vllm-sha +++ b/.sync/vllm-sha @@ -1 +1 @@ -6f91edf96d3f3272945809c04702380053bff4de +6f00a1ae3bd4b86168667bce673998218f461c0f 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/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/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/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: From 02dbc4549c25d11222f930b25c084e8e682cbb14 Mon Sep 17 00:00:00 2001 From: AlpinDale Date: Wed, 29 Jul 2026 11:26:04 +0000 Subject: [PATCH 26/41] [sync] [Model] Add Kimi K3 support: model files and kernels [1/N] (#50089) Upstream-vLLM: 7c6729b769597541d327f666273cd972b8a5318d Co-authored-by: Jiangyun Zhu --- .buildkite/check-torch-abi.py | 1 + .buildkite/check-torch-abi.py.orig | 94 + .pre-commit-config.yaml | 2 +- .pre-commit-config.yaml.orig | 186 ++ .sync/vllm-sha | 2 +- CMakeLists.txt | 27 +- CMakeLists.txt.orig | 1551 +++++++++++++ aphrodite/_custom_ops.py | 133 +- aphrodite/_custom_ops.py.orig | 41 +- aphrodite/envs.py | 6 + aphrodite/envs.py.orig | 1996 +++++++++++++++++ .../kernels/linear/cute_dsl/_skinny_gemm.py | 176 ++ .../kernels/linear/cute_dsl/skinny_gemm.py | 246 ++ aphrodite/model_executor/layers/activation.py | 43 + .../model_executor/layers/activation.py.orig | 806 +++++++ .../layers/fused_moe/activation.py | 19 + .../layers/fused_moe/activation.py.orig | 188 ++ .../model_executor/layers/fused_moe/config.py | 4 + .../layers/fused_moe/config.py.orig | 1414 ++++++++++++ .../layers/fused_moe/experts/deep_gemm_moe.py | 37 +- .../fused_moe/experts/deep_gemm_moe.py.orig | 580 +++++ .../layers/fused_moe/experts/marlin_moe.py | 59 + .../fused_moe/experts/marlin_moe.py.orig | 1036 +++++++++ .../fused_moe/experts/trtllm_mxfp4_moe.py | 54 +- .../experts/trtllm_mxfp4_moe.py.orig | 342 +++ .../fused_moe/experts/trtllm_nvfp4_moe.py | 29 +- .../experts/trtllm_nvfp4_moe.py.orig | 526 +++++ .../model_executor/layers/fused_moe/layer.py | 6 + .../layers/fused_moe/layer.py.orig | 422 ++++ .../layers/fused_moe/modular_kernel.py | 29 +- .../layers/fused_moe/modular_kernel.py.orig | 1678 ++++++++++++++ .../fused_moe/runner/latent_moe_runner.py | 222 ++ .../layers/fused_moe/runner/moe_runner.py | 52 +- .../fused_moe/runner/moe_runner.py.orig | 951 ++++++++ .../fused_moe/runner/moe_runner_interface.py | 1 + .../layers/mamba/ops/causal_conv1d.py | 12 +- .../layers/mamba/ops/causal_conv1d.py.orig | 1178 ++++++++++ .../layers/mamba/ops/gather_initial_states.py | 81 + aphrodite/models/common/__init__.py | 2 + aphrodite/models/common/ops/__init__.py | 9 + .../models/common/ops/fused_qk_rmsnorm.py | 103 + .../inkling/nvidia/ops/fa4_rel_attention.py | 11 +- .../nvidia/ops/fa4_rel_attention.py.orig | 161 ++ aphrodite/models/kimi_k3/__init__.py | 26 + aphrodite/models/kimi_k3/amd/linear.py | 993 ++++++++ aphrodite/models/kimi_k3/amd/model.py | 224 ++ aphrodite/models/kimi_k3/amd/mtp.py | 389 ++++ .../kimi_k3/amd/ops/third_party/__init__.py | 2 + .../amd/ops/third_party/kda/__init__.py | 47 + .../kimi_k3/amd/ops/third_party/kda/chunk.py | 935 ++++++++ .../amd/ops/third_party/kda/chunk_intra.py | 662 ++++++ .../kda/chunk_intra_token_parallel.py | 197 ++ .../ops/third_party/kda/fused_recurrent.py | 621 +++++ aphrodite/models/kimi_k3/common/__init__.py | 2 + .../models/kimi_k3/common/mm_preprocess.py | 330 +++ aphrodite/models/kimi_k3/common/mtp.py | 96 + aphrodite/models/kimi_k3/nvidia/dspark_mla.py | 495 ++++ aphrodite/models/kimi_k3/nvidia/kda.py | 735 ++++++ .../models/kimi_k3/nvidia/kda_metadata.py | 474 ++++ .../models/kimi_k3/nvidia/low_latency_gemm.py | 487 ++++ aphrodite/models/kimi_k3/nvidia/mla.py | 719 ++++++ aphrodite/models/kimi_k3/nvidia/model.py | 1720 ++++++++++++++ aphrodite/models/kimi_k3/nvidia/mtp.py | 417 ++++ .../kimi_k3/nvidia/ops/cute_dsl/__init__.py | 2 + .../ops/cute_dsl/latent_moe_tail/__init__.py | 14 + ...educe_rmsnorm_reduce_scatter_early_exit.py | 1012 +++++++++ .../fused_add_multicast_gemm.py | 1291 +++++++++++ .../fused_add_multicast_skinny_gemm.py | 438 ++++ .../cute_dsl/latent_moe_tail/lamport_copy.py | 226 ++ .../cute_dsl/latent_moe_tail/primitives.py | 437 ++++ .../ops/fused_mla_key_concat_kv_cache.py | 234 ++ .../kimi_k3/nvidia/ops/latent_moe_tail.py | 245 ++ .../kimi_k3/nvidia/ops/sequence_parallel.py | 73 + .../nvidia/ops/third_party/__init__.py | 2 + .../nvidia/ops/third_party/kda/__init__.py | 28 + .../nvidia/ops/third_party/kda/chunk.py | 938 ++++++++ .../nvidia/ops/third_party/kda/chunk_intra.py | 559 +++++ .../kda/chunk_intra_token_parallel.py | 178 ++ .../ops/third_party/kda/fused_recurrent.py | 671 ++++++ .../kimi_k3/nvidia/ops/vision_fa4_warmup.py | 191 ++ .../ops/fused_norm_gate.py | 412 ++++ .../flash_linear_attention/ops/kda.py | 88 +- .../flash_linear_attention/ops/kda.py.orig | 1591 +++++++++++++ .../transformers_utils/configs/kimi_k3.py | 136 ++ .../transformers_utils/processors/kimi_k3.py | 61 + .../attention/ops/triton_merge_attn_states.py | 38 +- .../ops/triton_merge_attn_states.py.orig | 267 +++ .../kernels/benchmark_k3_cutedsl_residual.py | 340 +++ .../benchmark_kimi_k3_latent_moe_tail.py | 787 +++++++ .../benchmark_kimi_k3_sp_collectives.py | 239 ++ cmake/external_projects/flashkda.cmake | 74 + csrc/custom_all_gather_reduce_scatter.cuh | 327 +++ csrc/custom_all_reduce.cuh | 317 +-- csrc/custom_all_reduce.cuh.orig | 633 ++++++ csrc/custom_collective_common.cuh | 332 +++ csrc/flashkda_registration.cpp | 17 + csrc/libtorch_stable/activation_kernels.cu | 110 + .../activation_kernels.cu.orig | 724 ++++++ .../attention/merge_attn_states.cu | 39 +- csrc/libtorch_stable/cache_kernels.cu | 96 + csrc/libtorch_stable/cache_kernels.cu.orig | 1669 ++++++++++++++ .../custom_all_gather_reduce_scatter.cu | 362 +++ .../custom_all_gather_reduce_scatter_ops.cpp | 29 + csrc/libtorch_stable/custom_all_reduce.cu | 8 +- .../libtorch_stable/custom_all_reduce.cu.orig | 202 ++ csrc/libtorch_stable/dsv3_fused_a_gemm.cu | 151 +- ..._kimi_k3_mla_key_concat_kv_cache_kernel.cu | 1237 ++++++++++ .../kimi_k3/fused_kda_decode_kernel.cu | 1130 ++++++++++ .../moe/grouped_topk_kernels.cu | 438 +++- .../moe/grouped_topk_kernels.cu.orig | 1124 ++++++++++ csrc/libtorch_stable/moe/moeTopKFuncs.cuh | 309 ++- .../libtorch_stable/moe/moeTopKFuncs.cuh.orig | 257 +++ csrc/libtorch_stable/ops.h | 95 + csrc/libtorch_stable/ops.h.orig | 621 +++++ csrc/libtorch_stable/torch_bindings.cpp | 92 +- csrc/libtorch_stable/torch_bindings.cpp.orig | 1014 +++++++++ .../w8a8/fp8/nvidia/quant_utils.cuh | 33 +- setup.py | 2 + setup.py.orig | 796 +++++++ tests/distributed/test_custom_all_reduce.py | 103 + .../test_kimi_k3_mla_fused_epilogue.py | 201 ++ tests/kernels/core/test_activation.py | 38 + tests/kernels/core/test_activation.py.orig | 222 ++ tests/kernels/core/test_fused_q_kv_rmsnorm.py | 2 +- .../kernels/core/test_fused_rms_norm_gated.py | 4 +- tests/kernels/moe/test_grouped_topk.py | 238 ++ tests/kernels/moe/test_grouped_topk.py.orig | 316 +++ .../kernels/moe/test_moe_align_block_size.py | 1 + .../moe/test_moe_align_block_size.py.orig | 366 +++ tests/kernels/test_bf16_skinny_gemm.py | 607 +++++ tests/models/kimi_k3/test_attn_res.py | 31 + tests/models/kimi_k3/test_attn_res.py.orig | 171 ++ tests/models/kimi_k3/test_eagle3.py | 130 ++ tests/models/kimi_k3/test_kda.py | 736 ++++++ tests/models/kimi_k3/test_kda_metadata.py | 390 ++++ tests/models/kimi_k3/test_latent_moe_tail.py | 145 ++ .../models/kimi_k3/test_sequence_parallel.py | 343 +++ .../multimodal/generation/test_common.py | 2 +- .../multimodal/generation/test_common.py.orig | 1355 +++++++++++ 139 files changed, 51679 insertions(+), 545 deletions(-) create mode 100644 .buildkite/check-torch-abi.py.orig create mode 100644 .pre-commit-config.yaml.orig create mode 100644 CMakeLists.txt.orig create mode 100755 aphrodite/envs.py.orig create mode 100644 aphrodite/model_executor/kernels/linear/cute_dsl/_skinny_gemm.py create mode 100644 aphrodite/model_executor/kernels/linear/cute_dsl/skinny_gemm.py create mode 100644 aphrodite/model_executor/layers/activation.py.orig create mode 100644 aphrodite/model_executor/layers/fused_moe/activation.py.orig create mode 100644 aphrodite/model_executor/layers/fused_moe/config.py.orig create mode 100644 aphrodite/model_executor/layers/fused_moe/experts/deep_gemm_moe.py.orig create mode 100644 aphrodite/model_executor/layers/fused_moe/experts/marlin_moe.py.orig create mode 100644 aphrodite/model_executor/layers/fused_moe/experts/trtllm_mxfp4_moe.py.orig create mode 100644 aphrodite/model_executor/layers/fused_moe/experts/trtllm_nvfp4_moe.py.orig create mode 100644 aphrodite/model_executor/layers/fused_moe/layer.py.orig create mode 100644 aphrodite/model_executor/layers/fused_moe/modular_kernel.py.orig create mode 100644 aphrodite/model_executor/layers/fused_moe/runner/latent_moe_runner.py create mode 100644 aphrodite/model_executor/layers/fused_moe/runner/moe_runner.py.orig create mode 100644 aphrodite/model_executor/layers/mamba/ops/causal_conv1d.py.orig create mode 100644 aphrodite/model_executor/layers/mamba/ops/gather_initial_states.py create mode 100644 aphrodite/models/common/__init__.py create mode 100644 aphrodite/models/common/ops/__init__.py create mode 100644 aphrodite/models/common/ops/fused_qk_rmsnorm.py create mode 100644 aphrodite/models/inkling/nvidia/ops/fa4_rel_attention.py.orig create mode 100644 aphrodite/models/kimi_k3/amd/linear.py create mode 100644 aphrodite/models/kimi_k3/amd/model.py create mode 100644 aphrodite/models/kimi_k3/amd/mtp.py create mode 100644 aphrodite/models/kimi_k3/amd/ops/third_party/__init__.py create mode 100644 aphrodite/models/kimi_k3/amd/ops/third_party/kda/__init__.py create mode 100644 aphrodite/models/kimi_k3/amd/ops/third_party/kda/chunk.py create mode 100644 aphrodite/models/kimi_k3/amd/ops/third_party/kda/chunk_intra.py create mode 100644 aphrodite/models/kimi_k3/amd/ops/third_party/kda/chunk_intra_token_parallel.py create mode 100644 aphrodite/models/kimi_k3/amd/ops/third_party/kda/fused_recurrent.py create mode 100644 aphrodite/models/kimi_k3/common/__init__.py create mode 100644 aphrodite/models/kimi_k3/common/mm_preprocess.py create mode 100644 aphrodite/models/kimi_k3/common/mtp.py create mode 100644 aphrodite/models/kimi_k3/nvidia/dspark_mla.py create mode 100644 aphrodite/models/kimi_k3/nvidia/kda.py create mode 100644 aphrodite/models/kimi_k3/nvidia/kda_metadata.py create mode 100644 aphrodite/models/kimi_k3/nvidia/low_latency_gemm.py create mode 100644 aphrodite/models/kimi_k3/nvidia/mla.py create mode 100644 aphrodite/models/kimi_k3/nvidia/model.py create mode 100644 aphrodite/models/kimi_k3/nvidia/mtp.py create mode 100644 aphrodite/models/kimi_k3/nvidia/ops/cute_dsl/__init__.py create mode 100644 aphrodite/models/kimi_k3/nvidia/ops/cute_dsl/latent_moe_tail/__init__.py create mode 100644 aphrodite/models/kimi_k3/nvidia/ops/cute_dsl/latent_moe_tail/allreduce_rmsnorm_reduce_scatter_early_exit.py create mode 100644 aphrodite/models/kimi_k3/nvidia/ops/cute_dsl/latent_moe_tail/fused_add_multicast_gemm.py create mode 100644 aphrodite/models/kimi_k3/nvidia/ops/cute_dsl/latent_moe_tail/fused_add_multicast_skinny_gemm.py create mode 100644 aphrodite/models/kimi_k3/nvidia/ops/cute_dsl/latent_moe_tail/lamport_copy.py create mode 100644 aphrodite/models/kimi_k3/nvidia/ops/cute_dsl/latent_moe_tail/primitives.py create mode 100644 aphrodite/models/kimi_k3/nvidia/ops/fused_mla_key_concat_kv_cache.py create mode 100644 aphrodite/models/kimi_k3/nvidia/ops/latent_moe_tail.py create mode 100644 aphrodite/models/kimi_k3/nvidia/ops/sequence_parallel.py create mode 100644 aphrodite/models/kimi_k3/nvidia/ops/third_party/__init__.py create mode 100644 aphrodite/models/kimi_k3/nvidia/ops/third_party/kda/__init__.py create mode 100644 aphrodite/models/kimi_k3/nvidia/ops/third_party/kda/chunk.py create mode 100644 aphrodite/models/kimi_k3/nvidia/ops/third_party/kda/chunk_intra.py create mode 100644 aphrodite/models/kimi_k3/nvidia/ops/third_party/kda/chunk_intra_token_parallel.py create mode 100644 aphrodite/models/kimi_k3/nvidia/ops/third_party/kda/fused_recurrent.py create mode 100644 aphrodite/models/kimi_k3/nvidia/ops/vision_fa4_warmup.py create mode 100644 aphrodite/third_party/flash_linear_attention/ops/fused_norm_gate.py create mode 100644 aphrodite/third_party/flash_linear_attention/ops/kda.py.orig create mode 100644 aphrodite/transformers_utils/configs/kimi_k3.py create mode 100644 aphrodite/transformers_utils/processors/kimi_k3.py create mode 100644 aphrodite/v1/attention/ops/triton_merge_attn_states.py.orig create mode 100644 benchmarks/kernels/benchmark_k3_cutedsl_residual.py create mode 100644 benchmarks/kernels/benchmark_kimi_k3_latent_moe_tail.py create mode 100644 benchmarks/kernels/benchmark_kimi_k3_sp_collectives.py create mode 100644 cmake/external_projects/flashkda.cmake create mode 100644 csrc/custom_all_gather_reduce_scatter.cuh create mode 100644 csrc/custom_all_reduce.cuh.orig create mode 100644 csrc/custom_collective_common.cuh create mode 100644 csrc/flashkda_registration.cpp create mode 100644 csrc/libtorch_stable/activation_kernels.cu.orig create mode 100644 csrc/libtorch_stable/cache_kernels.cu.orig create mode 100644 csrc/libtorch_stable/custom_all_gather_reduce_scatter.cu create mode 100644 csrc/libtorch_stable/custom_all_gather_reduce_scatter_ops.cpp create mode 100644 csrc/libtorch_stable/custom_all_reduce.cu.orig create mode 100644 csrc/libtorch_stable/fused_kimi_k3_mla_key_concat_kv_cache_kernel.cu create mode 100644 csrc/libtorch_stable/kimi_k3/fused_kda_decode_kernel.cu create mode 100644 csrc/libtorch_stable/moe/grouped_topk_kernels.cu.orig create mode 100644 csrc/libtorch_stable/moe/moeTopKFuncs.cuh.orig create mode 100644 csrc/libtorch_stable/ops.h.orig create mode 100644 csrc/libtorch_stable/torch_bindings.cpp.orig create mode 100644 setup.py.orig create mode 100644 tests/kernels/attention/test_kimi_k3_mla_fused_epilogue.py create mode 100644 tests/kernels/core/test_activation.py.orig create mode 100644 tests/kernels/moe/test_grouped_topk.py.orig create mode 100644 tests/kernels/moe/test_moe_align_block_size.py.orig create mode 100644 tests/kernels/test_bf16_skinny_gemm.py create mode 100644 tests/models/kimi_k3/test_attn_res.py.orig create mode 100644 tests/models/kimi_k3/test_eagle3.py create mode 100644 tests/models/kimi_k3/test_kda.py create mode 100644 tests/models/kimi_k3/test_kda_metadata.py create mode 100644 tests/models/kimi_k3/test_latent_moe_tail.py create mode 100644 tests/models/kimi_k3/test_sequence_parallel.py create mode 100644 tests/models/multimodal/generation/test_common.py.orig 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/.buildkite/check-torch-abi.py.orig b/.buildkite/check-torch-abi.py.orig new file mode 100644 index 0000000000..ebf2fb4b27 --- /dev/null +++ b/.buildkite/check-torch-abi.py.orig @@ -0,0 +1,94 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the Aphrodite Engine project + +"""Audit Aphrodite compiled libraries for PyTorch stable ABI compliance.""" + +import fnmatch +import sys +from pathlib import Path + +from torch_abi_audit import inspect_package +from torch_abi_audit.report import ExtensionReport, PackageReport + +# Temporary allowlist of extensions not yet on the stable ABI. +# Shrink and remove over time. +ALLOWED_UNSTABLE_LIBRARIES: tuple[str, ...] = ( + "_C_fork.abi3.so", + "vllm_flash_attn/_vllm_fa2_C.abi3.so", + "vllm_flash_attn/_vllm_fa3_C.abi3.so", + "third_party/deep_gemm/_C*.so", +) + + +def _relative_path(lib: ExtensionReport, package_root: Path) -> str: + try: + return lib.path.relative_to(package_root).as_posix() + except ValueError: + return lib.path.name + + +def _is_torch_unstable(lib: ExtensionReport) -> bool: + return lib.error is None and lib.torch.uses_torch and not lib.torch.stable + + +def _matches_allowlist(rel_path: str, patterns: tuple[str, ...]) -> bool: + return any(fnmatch.fnmatch(rel_path, pattern) for pattern in patterns) + + +def _iter_libs(report: PackageReport) -> tuple[ExtensionReport, ...]: + return (*report.extensions, *report.bundled_libs) + + +def _collect_unstable(report: PackageReport) -> list[str]: + return sorted(_relative_path(lib, report.root) for lib in _iter_libs(report) if _is_torch_unstable(lib)) + + +def _find_stale_allowlist_entries(report: PackageReport, patterns: tuple[str, ...]) -> list[str]: + """Allowlist patterns that match a built library which is no longer unstable.""" + stale: list[str] = [] + for pattern in patterns: + for lib in _iter_libs(report): + if lib.error is not None: + continue + if not fnmatch.fnmatch(_relative_path(lib, report.root), pattern): + continue + if not _is_torch_unstable(lib): + stale.append(pattern) + break + return stale + + +def check_torch_abi( + package: str = "aphrodite", + patterns: tuple[str, ...] = ALLOWED_UNSTABLE_LIBRARIES, +) -> int: + report = inspect_package(package) + if report.error: + print(f"error: failed to inspect {package!r}: {report.error}", file=sys.stderr) + return 2 + + unstable = _collect_unstable(report) + unexpected = [rel_path for rel_path in unstable if not _matches_allowlist(rel_path, patterns)] + stale = _find_stale_allowlist_entries(report, patterns) + + if unexpected or stale: + if unexpected: + print( + f"Not allowed: torch-unstable libraries outside ALLOWED_UNSTABLE_LIBRARIES: {', '.join(unexpected)}", + file=sys.stderr, + ) + if stale: + print( + f"Not allowed: stale ALLOWED_UNSTABLE_LIBRARIES entries: {', '.join(stale)}", + file=sys.stderr, + ) + return 1 + + print("Torch stable ABI check passed.") + return 0 + + +if __name__ == "__main__": + print(">>> Auditing Aphrodite extension modules for PyTorch stable ABI compliance") + sys.exit(check_torch_abi()) diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 6064ebd9e0..9a51e319b9 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 diff --git a/.pre-commit-config.yaml.orig b/.pre-commit-config.yaml.orig new file mode 100644 index 0000000000..6064ebd9e0 --- /dev/null +++ b/.pre-commit-config.yaml.orig @@ -0,0 +1,186 @@ +default_install_hook_types: + - pre-commit + - commit-msg +default_stages: + - pre-commit # Run locally + - manual # Run in CI +exclude: '^(reference/|shellcheck-stable/|aphrodite/third_party/.*)' +repos: +- repo: https://github.com/astral-sh/ruff-pre-commit + rev: v0.14.0 + hooks: + - id: ruff-check + args: [--output-format, github, --fix] + - id: ruff-format +- repo: https://github.com/crate-ci/typos + rev: v1.43.5 + hooks: + - id: typos + exclude: '(kernels|csrc)/xqa/cubin/.*' + args: [--force-exclude] + require_serial: true +- repo: https://github.com/pre-commit/mirrors-clang-format + rev: v21.1.2 + hooks: + - id: clang-format + exclude: '(kernels|csrc)/(moe/topk_softmax_kernels.cu|quantization/gguf/(ggml-common.h|dequantize.cuh|vecdotq.cuh|mmq.cuh|mmvq.cuh)|quantization/exl3/exllamav3_ext/quant/exl3_kernel_map_samples.cuh|flash_mla|hadamard|backup|vmm|xqa)|aphrodite/third_party/.*' + types_or: [c++, cuda] + args: [--style=file, --verbose] + require_serial: true +- repo: https://github.com/DavidAnson/markdownlint-cli2 + rev: v0.21.0 + hooks: + - id: markdownlint-cli2 + language_version: lts + args: [--fix] + exclude: (^|/)CLAUDE\.md$ +- repo: https://github.com/rhysd/actionlint + rev: v1.7.7 + hooks: + - id: actionlint +- repo: https://github.com/astral-sh/uv-pre-commit + rev: 0.11.1 + hooks: + - id: pip-compile + args: [ + requirements/test/cuda.in, + -c, requirements/cuda.txt, + -o, requirements/test/cuda.txt, + --index-strategy, unsafe-best-match, + --torch-backend, cu130, + --python-platform, x86_64-manylinux_2_28, + --python-version, "3.12", + ] + files: ^requirements/(common|cuda|test/cuda)\.(in|txt)$ +- repo: local + hooks: + - id: mypy-local + name: Run mypy locally for lowest supported Python version + entry: python tools/pre_commit/mypy.py 0 "3.10" + stages: [pre-commit] # Don't run in CI + <<: &mypy_common + language: python + types_or: [python, pyi] + require_serial: true + additional_dependencies: ["mypy[faster-cache]==1.19.1", regex, types-cachetools, types-setuptools, types-PyYAML, types-requests, types-torch, pydantic] + - id: mypy-3.10 + name: Run mypy for Python 3.10 + entry: python tools/pre_commit/mypy.py 1 "3.10" + <<: *mypy_common + stages: [manual] + - id: mypy-3.11 + name: Run mypy for Python 3.11 + entry: python tools/pre_commit/mypy.py 1 "3.11" + <<: *mypy_common + stages: [manual] + - id: mypy-3.12 + name: Run mypy for Python 3.12 + entry: python tools/pre_commit/mypy.py 1 "3.12" + <<: *mypy_common + stages: [manual] + - id: mypy-3.13 + name: Run mypy for Python 3.13 + entry: python tools/pre_commit/mypy.py 1 "3.13" + <<: *mypy_common + stages: [manual] + - id: shellcheck + name: Lint shell scripts + entry: tools/pre_commit/shellcheck.sh + language: script + types: [shell] + files: ^tools/pre_commit/.*\.sh$ + - id: png-lint + name: Lint PNG exports from excalidraw + entry: tools/pre_commit/png-lint.sh + language: script + types: [png] + - id: signoff-commit + name: Sign-off Commit + entry: bash + args: + - -c + - | + if ! grep -q "^Signed-off-by: $(git config user.name) <$(git config user.email)>" "$(git rev-parse --git-path COMMIT_EDITMSG)"; then + printf "\nSigned-off-by: $(git config user.name) <$(git config user.email)>\n" >> "$(git rev-parse --git-path COMMIT_EDITMSG)" + fi + language: system + verbose: true + stages: [commit-msg] + - id: check-spdx-header + name: Check SPDX headers + entry: python tools/pre_commit/check_spdx_header.py + language: python + types_or: [python, rust, proto] + exclude: ^(reference/|aphrodite/vllm_flash_attn/|aphrodite/third_party/|tests/|examples/|benchmarks/) + - id: check-root-lazy-imports + name: Check root lazy imports + entry: python tools/pre_commit/check_init_lazy_imports.py + language: python + types: [python] + - id: check-filenames + name: Check for spaces in all filenames + entry: bash + args: + - -c + - 'git ls-files | grep " " && echo "Filenames should not contain spaces!" && exit 1 || exit 0' + 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 + language: python + types: [python] + pass_filenames: false + additional_dependencies: [regex] + - id: forbid-direct-triton-import + name: "Forbid direct 'import triton'" + entry: python tools/pre_commit/check_triton_import.py + language: python + types: [python] + pass_filenames: false + additional_dependencies: [regex] + - id: check-pickle-imports + name: Prevent new pickle/cloudpickle imports + entry: python tools/pre_commit/check_pickle_imports.py + language: python + types: [python] + additional_dependencies: [regex] + exclude: ^benchmarks/ + - id: check-huggingface-imports + name: Prevent direct Hugging Face Hub repository API imports + entry: python tools/pre_commit/check_huggingface_imports.py + language: python + types: [python] + additional_dependencies: [regex] + - id: check-torch-cuda-call + name: "Prevent new 'torch.cuda' APIs call" + entry: python tools/pre_commit/check_torch_cuda.py + language: python + types: [python] + additional_dependencies: [regex] + - id: validate-config + name: Validate configuration has default values and that each field has a docstring + entry: python tools/pre_commit/validate_config.py + language: python + additional_dependencies: [regex] + - id: validate-docker-versions + name: Validate docker/versions.json matches Dockerfile + entry: python tools/generate_versions_json.py --check + language: python + files: ^docker/(Dockerfile|versions\.json)$ + pass_filenames: false + additional_dependencies: [dockerfile-parse] + - id: check-boolean-context-manager + name: Check for boolean ops in with-statements + entry: python tools/pre_commit/check_boolean_context_manager.py + language: python + types: [python] + # Keep `suggestion` last + - id: suggestion + name: Suggestion + entry: bash -c 'echo "To bypass all the pre-commit hooks, add --no-verify to git commit. To skip a specific hook, prefix the commit command with SKIP=."' + language: system + verbose: true + pass_filenames: false + # Insert new entries above the `suggestion` entry diff --git a/.sync/vllm-sha b/.sync/vllm-sha index 8204a3f27f..1cf8d9f73e 100644 --- a/.sync/vllm-sha +++ b/.sync/vllm-sha @@ -1 +1 @@ -6f00a1ae3bd4b86168667bce673998218f461c0f +7c6729b769597541d327f666273cd972b8a5318d 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/CMakeLists.txt.orig b/CMakeLists.txt.orig new file mode 100644 index 0000000000..2fd8f0c968 --- /dev/null +++ b/CMakeLists.txt.orig @@ -0,0 +1,1551 @@ +cmake_minimum_required(VERSION 3.26) + +# When building directly using CMake, make sure you run the install step +# (it places the .so files in the correct location). +# +# Example: +# mkdir build && cd build +# cmake -G Ninja -DAPHRODITE_PYTHON_EXECUTABLE=`which python3` -DCMAKE_INSTALL_PREFIX=.. .. +# cmake --build . --target install +# +# If you want to only build one target, make sure to install it manually: +# cmake --build . --target _C +# cmake --install . --component _C +project(aphrodite_extensions LANGUAGES CXX) + +set(CMAKE_CXX_STANDARD 20) +set(CMAKE_CXX_STANDARD_REQUIRED ON) +set(CMAKE_CUDA_STANDARD 20) +set(CMAKE_CUDA_STANDARD_REQUIRED ON) +set(CMAKE_HIP_STANDARD 20) +set(CMAKE_HIP_STANDARD_REQUIRED ON) + +# PyTorch headers require C++20; GCC < 11.3 has incomplete C++20 support. +if(CMAKE_CXX_COMPILER_ID STREQUAL "GNU" AND CMAKE_CXX_COMPILER_VERSION VERSION_LESS "11.3") + message(FATAL_ERROR + "GCC >= 11.3 is required to build Aphrodite (found ${CMAKE_CXX_COMPILER_VERSION}). " + "PyTorch's C++20 headers require a compiler with full C++20 support. " + "See: https://github.com/pytorch/pytorch/pull/167929") +endif() + + +# CUDA by default, can be overridden by using -DAPHRODITE_TARGET_DEVICE=... (used by setup.py) +set(APHRODITE_TARGET_DEVICE "cuda" CACHE STRING "Target device backend for Aphrodite") +message(STATUS "Build type: ${CMAKE_BUILD_TYPE}") +message(STATUS "Target device: ${APHRODITE_TARGET_DEVICE}") + +include(${CMAKE_CURRENT_LIST_DIR}/cmake/utils.cmake) + +# Suppress potential warnings about unused manually-specified variables +set(ignoreMe "${APHRODITE_PYTHON_PATH}") + +# Prevent installation of dependencies (cutlass) by default. +install(CODE "set(CMAKE_INSTALL_LOCAL_ONLY TRUE)" ALL_COMPONENTS) + +# +# Supported python versions. These versions will be searched in order, the +# first match will be selected. These should be kept in sync with setup.py. +# +set(PYTHON_SUPPORTED_VERSIONS "3.10" "3.11" "3.12" "3.13" "3.14") + +# Supported AMD GPU architectures. +set(HIP_SUPPORTED_ARCHS "gfx906;gfx908;gfx90a;gfx942;gfx950;gfx1030;gfx1100;gfx1101;gfx1102;gfx1103;gfx1150;gfx1151;gfx1152;gfx1153;gfx1200;gfx1201") + +# ROCm installation prefix. Default to /opt/rocm but allow override via +# -DROCM_PATH=/your/rocm/path when invoking cmake. +if(NOT DEFINED ROCM_PATH) + set(ROCM_PATH "/opt/rocm" CACHE PATH "ROCm installation prefix") +else() + set(ROCM_PATH ${ROCM_PATH} CACHE PATH "ROCm installation prefix" FORCE) +endif() +# +# Supported/expected torch versions for CUDA/ROCm. +# +# Currently, having an incorrect pytorch version results in a warning +# rather than an error. +# +# Note: the CUDA torch version is derived from pyproject.toml and various +# requirements.txt files and should be kept consistent. The ROCm torch +# versions are derived from docker/Dockerfile.rocm +# +set(TORCH_SUPPORTED_VERSION_CUDA "2.13.0") +set(TORCH_SUPPORTED_VERSION_ROCM "2.13.0") +# TORCH_NIGHTLY=1 builds run against unpinned nightly wheels, so the supported- +# version check would always warn. Only treat it as a nightly build when the +# value is exactly "1" (the bootstrap exports TORCH_NIGHTLY=0 by default, which +# must NOT suppress the warning for normal builds). +if (DEFINED ENV{TORCH_NIGHTLY} AND "$ENV{TORCH_NIGHTLY}" STREQUAL "1") + set(TORCH_NIGHTLY_BUILD TRUE) +else() + set(TORCH_NIGHTLY_BUILD FALSE) +endif() + +# +# Try to find python package with an executable that exactly matches +# `APHRODITE_PYTHON_EXECUTABLE` and is one of the supported versions. +# +if (APHRODITE_PYTHON_EXECUTABLE) + find_python_from_executable(${APHRODITE_PYTHON_EXECUTABLE} "${PYTHON_SUPPORTED_VERSIONS}") +else() + message(FATAL_ERROR + "Please set APHRODITE_PYTHON_EXECUTABLE to the path of the desired python version" + " before running cmake configure.") +endif() + +# +# Update cmake's `CMAKE_PREFIX_PATH` with torch location. +# +append_cmake_prefix_path("torch" "torch.utils.cmake_prefix_path") + +# Ensure the 'nvcc' command is in the PATH +find_program(NVCC_EXECUTABLE nvcc) +if (CUDA_FOUND AND NOT NVCC_EXECUTABLE) + message(FATAL_ERROR "nvcc not found") +endif() + +# +# Import torch cmake configuration. +# Torch also imports CUDA (and partially HIP) languages with some customizations, +# so there is no need to do this explicitly with check_language/enable_language, +# etc. +# +find_package(Torch REQUIRED) + +# Supported NVIDIA architectures. +# This check must happen after find_package(Torch) because that's when CMAKE_CUDA_COMPILER_VERSION gets defined +if(DEFINED CMAKE_CUDA_COMPILER_VERSION AND + CMAKE_CUDA_COMPILER_VERSION VERSION_GREATER_EQUAL 13.4) + # Rubin (10.7) can run SM100 family code, but CUDA 13.4 also supports + # targeting it directly. + set(CUDA_SUPPORTED_ARCHS "7.5;8.0;8.6;8.7;8.9;9.0;10.0;10.7;11.0;12.0") +elseif(DEFINED CMAKE_CUDA_COMPILER_VERSION AND + CMAKE_CUDA_COMPILER_VERSION VERSION_GREATER_EQUAL 13.0) + # starting from CUDA 12.9 and Blackwell (10.0), we use family-specific targets (10.0f, 12.0f, etc) + # to support the whole generation without specifying all sub-architectures + # see: https://developer.nvidia.com/blog/nvidia-blackwell-and-nvidia-cuda-12-9-introduce-family-specific-architecture-features/ + set(CUDA_SUPPORTED_ARCHS "7.5;8.0;8.6;8.7;8.9;9.0;10.0;11.0;12.0") +elseif(DEFINED CMAKE_CUDA_COMPILER_VERSION AND + CMAKE_CUDA_COMPILER_VERSION VERSION_GREATER_EQUAL 12.8) + set(CUDA_SUPPORTED_ARCHS "7.5;8.0;8.6;8.7;8.9;9.0;10.0;10.1;10.3;12.0;12.1") +else() + set(CUDA_SUPPORTED_ARCHS "7.0;7.5;8.0;8.6;8.7;8.9;9.0") +endif() + +# +# spinloop extension (pure CXX; must stay above the non-CUDA device branch so +# CPU builds define the target before the early return) +# This extension requires SABI 3.11 since it relies on Py_buffer support. Loading +# failure is handled gracefully on Aphrodite side for lower Python versions. +# Skip the target entirely on Python < 3.11 so the build doesn't break. +# +if(Python_VERSION VERSION_GREATER_EQUAL "3.11") + set(APHRODITE_SPINLOOP_EXT_SRC "csrc/spinloop.cpp") + set(SPINLOOP_COMPILE_FLAGS "") + if(CMAKE_SYSTEM_PROCESSOR MATCHES "x86_64|amd64") + list(APPEND SPINLOOP_COMPILE_FLAGS "-mmwaitx") + endif() + define_extension_target( + spinloop + DESTINATION aphrodite + LANGUAGE CXX + SOURCES ${APHRODITE_SPINLOOP_EXT_SRC} + COMPILE_FLAGS ${SPINLOOP_COMPILE_FLAGS} + USE_SABI 3.11 + WITH_SOABI) +endif() + +# +# fs_io extension (pure CXX; must stay above the non-CUDA device branch +# so CPU builds define the target before the early return). +# GIL-releasing filesystem helpers for FileSystemTierManager. +# +if(Python_VERSION VERSION_GREATER_EQUAL "3.11") + define_extension_target( + fs_io_C + DESTINATION aphrodite + LANGUAGE CXX + SOURCES csrc/fs_io.cpp + USE_SABI 3.11 + WITH_SOABI) +endif() + +# +# Forward the non-CUDA device extensions to external CMake scripts. +# +if (NOT APHRODITE_TARGET_DEVICE STREQUAL "cuda" AND + NOT APHRODITE_TARGET_DEVICE STREQUAL "rocm") + if (APHRODITE_TARGET_DEVICE STREQUAL "cpu") + include(${CMAKE_CURRENT_LIST_DIR}/cmake/cpu_extension.cmake) + elseif (APHRODITE_TARGET_DEVICE STREQUAL "metal") + include(${CMAKE_CURRENT_LIST_DIR}/cmake/metal_extension.cmake) + else() + return() + endif() + return() +endif() + +# +# Set up GPU language and check the torch version and warn if it isn't +# what is expected. +# +if (NOT HIP_FOUND AND NOT PYTORCH_FOUND_HIP AND CUDA_FOUND) + set(APHRODITE_GPU_LANG "CUDA") + + if (NOT TORCH_NIGHTLY_BUILD AND NOT Torch_VERSION VERSION_EQUAL ${TORCH_SUPPORTED_VERSION_CUDA}) + message(WARNING "Pytorch version ${TORCH_SUPPORTED_VERSION_CUDA} " + "expected for CUDA build, saw ${Torch_VERSION} instead.") + endif() +elseif(HIP_FOUND OR PYTORCH_FOUND_HIP) + set(APHRODITE_GPU_LANG "HIP") + + # Importing torch recognizes and sets up some HIP/ROCm configuration but does + # not let cmake recognize .hip files. In order to get cmake to understand the + # .hip extension automatically, HIP must be enabled explicitly. + enable_language(HIP) + + # ROCm 5.X and 6.X + if (NOT TORCH_NIGHTLY_BUILD AND ROCM_VERSION_DEV_MAJOR GREATER_EQUAL 5 AND + Torch_VERSION VERSION_LESS ${TORCH_SUPPORTED_VERSION_ROCM}) + message(WARNING "Pytorch version >= ${TORCH_SUPPORTED_VERSION_ROCM} " + "expected for ROCm build, saw ${Torch_VERSION} instead.") + endif() +else() + message(FATAL_ERROR "Can't find CUDA or HIP installation.") +endif() + + +if(APHRODITE_GPU_LANG STREQUAL "CUDA") + # + # For cuda we want to be able to control which architectures we compile for on + # a per-file basis in order to cut down on compile time. So here we extract + # the set of architectures we want to compile for and remove the from the + # CMAKE_CUDA_FLAGS so that they are not applied globally. + # + # `+PTX` in TORCH_CUDA_ARCH_LIST is not preserved here. If a kernel really + # needs PTX, add `+PTX` to that kernel's component-specific arch list below. + # + clear_cuda_arches(CUDA_ARCH_FLAGS) + extract_unique_cuda_archs_ascending(CUDA_ARCHS "${CUDA_ARCH_FLAGS}") + message(STATUS "CUDA target architectures: ${CUDA_ARCHS}") + # Filter the target architectures by the supported supported archs + # since for some files we will build for all CUDA_ARCHS. + cuda_archs_loose_intersection(CUDA_ARCHS + "${CUDA_SUPPORTED_ARCHS}" "${CUDA_ARCHS}") + message(STATUS "CUDA supported target architectures: ${CUDA_ARCHS}") + if(NOT CUDA_ARCHS) + message(FATAL_ERROR + "No supported CUDA architectures; the build would produce a binary " + "with no usable kernels. Detected gencode flags: ${CUDA_ARCH_FLAGS}; " + "supported: ${CUDA_SUPPORTED_ARCHS}. " + "Set TORCH_CUDA_ARCH_LIST for your GPU (e.g. 12.0).") + endif() +else() + # + # For other GPU targets override the GPU architectures detected by cmake/torch + # and filter them by the supported versions for the current language. + # The final set of arches is stored in `APHRODITE_GPU_ARCHES`. + # + override_gpu_arches(APHRODITE_GPU_ARCHES + ${APHRODITE_GPU_LANG} + "${${APHRODITE_GPU_LANG}_SUPPORTED_ARCHS}") +endif() + +# +# Query torch for additional GPU compilation flags for the given +# `APHRODITE_GPU_LANG`. +# The final set of arches is stored in `APHRODITE_GPU_FLAGS`. +# +get_torch_gpu_compiler_flags(APHRODITE_GPU_FLAGS ${APHRODITE_GPU_LANG}) + +# +# Set nvcc parallelism. +# +if(NVCC_THREADS AND APHRODITE_GPU_LANG STREQUAL "CUDA") + list(APPEND APHRODITE_GPU_FLAGS "--threads=${NVCC_THREADS}") +endif() + +# +# Set compression mode for CUDA >=13.x. +# +if(APHRODITE_GPU_LANG STREQUAL "CUDA" AND + DEFINED CMAKE_CUDA_COMPILER_VERSION AND + CMAKE_CUDA_COMPILER_VERSION VERSION_GREATER_EQUAL 13.0) + list(APPEND APHRODITE_GPU_FLAGS "--compress-mode=size") +endif() + +# +# Set CUDA include flags for CXX compiler. +# +if(APHRODITE_GPU_LANG STREQUAL "CUDA") + set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -I${CUDA_TOOLKIT_ROOT_DIR}/include") + if(CUDA_VERSION VERSION_GREATER_EQUAL 13.0) + set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -I${CUDA_TOOLKIT_ROOT_DIR}/include/cccl") + endif() +endif() + +# +# Use FetchContent for C++ dependencies that are compiled as part of Aphrodite's build process. +# setup.py will override FETCHCONTENT_BASE_DIR to play nicely with sccache. +# Each dependency that produces build artifacts should override its BINARY_DIR to avoid +# conflicts between build types. It should instead be set to ${CMAKE_BINARY_DIR}/. +# +include(FetchContent) +file(MAKE_DIRECTORY ${FETCHCONTENT_BASE_DIR}) # Ensure the directory exists +message(STATUS "FetchContent base directory: ${FETCHCONTENT_BASE_DIR}") + +if(APHRODITE_GPU_LANG STREQUAL "HIP") + # + # Overriding the default -O set up by cmake, adding ggdb3 for the most verbose devug info + # + set(CMAKE_${APHRODITE_GPU_LANG}_FLAGS_DEBUG "${CMAKE_${APHRODITE_GPU_LANG}_FLAGS_DEBUG} -O0 -ggdb3") + set(CMAKE_CXX_FLAGS_DEBUG "${CMAKE_CXX_FLAGS_DEBUG} -O0 -ggdb3") + + # + # Certain HIP functions are marked as [[nodiscard]], yet aphrodite ignores the result which generates + # a lot of warnings that always mask real issues. Suppressing until this is properly addressed. + # + set(CMAKE_${APHRODITE_GPU_LANG}_FLAGS "${CMAKE_${APHRODITE_GPU_LANG}_FLAGS} -Wno-unused-result -Wno-unused-value") + set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wno-unused-result -Wno-unused-value") + + # When using LTO then *.cpp files must be compiled with same compiler as used linker + # So if HIP uses clang linker we also must use it + # Otherwise symbols will be missing from .so + if (CMAKE_CXX_FLAGS MATCHES "\-flto") + if(NOT CMAKE_CXX_COMPILER_ID STREQUAL CMAKE_HIP_COMPILER_ID) + message(FATAL_ERROR "LTO is enabled for ROCm build, but the C++ compiler (${CMAKE_CXX_COMPILER_ID}) and HIP compiler (${CMAKE_HIP_COMPILER_ID}) are different which is not supported. " + "Please ensure they are same by setting CXX=${CMAKE_HIP_COMPILER} environment variable. Or alternatively disable LTO.") + endif() + endif() +endif() + +# +# Define other extension targets +# + +# +# cumem_allocator extension +# + +set(APHRODITE_CUMEM_EXT_SRC + "csrc/cumem_allocator.cpp") + +set_gencode_flags_for_srcs( + SRCS "${APHRODITE_CUMEM_EXT_SRC}" + CUDA_ARCHS "${CUDA_ARCHS}") + +if(APHRODITE_GPU_LANG STREQUAL "CUDA" OR APHRODITE_GPU_LANG STREQUAL "HIP") + message(STATUS "Enabling cumem allocator extension.") + if(APHRODITE_GPU_LANG STREQUAL "CUDA") + # link against cuda driver library + list(APPEND CUMEM_LIBS CUDA::cuda_driver) + else() + # link against rocm driver library. Prefer an absolute path to + # libamdhip64.so inside ${ROCM_PATH}/lib if available, otherwise fall + # back to linking by name "amdhip64". + find_library(AMDHIP64_LIB + NAMES amdhip64 libamdhip64.so + PATHS ${ROCM_PATH}/lib + NO_DEFAULT_PATH) + if(AMDHIP64_LIB) + message(STATUS "Found libamdhip64 at ${AMDHIP64_LIB}") + list(APPEND CUMEM_LIBS ${AMDHIP64_LIB}) + else() + message(WARNING "libamdhip64 not found in ${ROCM_PATH}/lib; falling back to linking 'amdhip64' by name") + list(APPEND CUMEM_LIBS amdhip64) + endif() + endif() + define_extension_target( + cumem_allocator + DESTINATION aphrodite + LANGUAGE CXX + SOURCES ${APHRODITE_CUMEM_EXT_SRC} + LIBRARIES ${CUMEM_LIBS} + USE_SABI 3.8 + WITH_SOABI) +endif() + +# +# Legacy _C extension (ROCm only — CUDA ops migrated to _C_stable_libtorch) +# + +if(APHRODITE_GPU_LANG STREQUAL "HIP") + set(APHRODITE_EXT_SRC + "csrc/torch_bindings.cpp" + "csrc/custom_quickreduce.cu") + + message(STATUS "Enabling C extension.") + define_extension_target( + _C + DESTINATION aphrodite + LANGUAGE ${APHRODITE_GPU_LANG} + SOURCES ${APHRODITE_EXT_SRC} + COMPILE_FLAGS ${APHRODITE_GPU_FLAGS} + ARCHITECTURES ${APHRODITE_GPU_ARCHES} + INCLUDE_DIRECTORIES ${CUTLASS_INCLUDE_DIR} + INCLUDE_DIRECTORIES ${CUTLASS_TOOLS_UTIL_INCLUDE_DIR} + USE_SABI 3 + WITH_SOABI) + + # If CUTLASS is compiled on NVCC >= 12.5, it by default uses + # cudaGetDriverEntryPointByVersion as a wrapper to avoid directly calling the + # driver API. This causes problems when linking with earlier versions of CUDA. + # Setting this variable sidesteps the issue by calling the driver directly. + target_compile_definitions(_C PRIVATE CUTLASS_ENABLE_DIRECT_CUDA_DRIVER_CALL=1) +endif() # _C HIP endif + +if(APHRODITE_GPU_LANG STREQUAL "CUDA" OR APHRODITE_GPU_LANG STREQUAL "HIP") + # + # _C_stable_libtorch extension (ops registered via STABLE_TORCH_LIBRARY) + # + set(APHRODITE_STABLE_EXT_SRC + "csrc/libtorch_stable/torch_bindings.cpp" + "csrc/cpu/dry.cpp" + "csrc/libtorch_stable/cuda_view.cu" + "csrc/libtorch_stable/cuda_utils_kernels.cu" + "csrc/libtorch_stable/activation_kernels.cu" + "csrc/libtorch_stable/ngram_embedding_kernels.cu" + "csrc/libtorch_stable/quantization/activation_kernels.cu" + "csrc/libtorch_stable/quantization/w8a8/int8/scaled_quant.cu" + "csrc/libtorch_stable/quantization/w8a8/fp8/common.cu" + "csrc/libtorch_stable/quantization/w8a8/fp8/per_token_group_quant.cu" + "csrc/libtorch_stable/quantization/w8a8/int8/per_token_group_quant.cu" + "csrc/libtorch_stable/permute_cols.cu" + "csrc/libtorch_stable/quantization/gptq/q_gemm.cu" + "csrc/libtorch_stable/pos_encoding_kernels.cu" + "csrc/libtorch_stable/fused_qknorm_rope_kernel.cu" + "csrc/libtorch_stable/fused_minimax_m3_qknorm_rope_kv_insert_kernel.cu" + "csrc/libtorch_stable/layernorm_kernels.cu" + "csrc/libtorch_stable/layernorm_quant_kernels.cu" + "csrc/libtorch_stable/quantization/fused_kernels/fused_layernorm_dynamic_per_token_quant.cu" + "csrc/libtorch_stable/quantization/fused_kernels/fused_silu_mul_block_quant.cu" + "csrc/libtorch_stable/attention/merge_attn_states.cu" + "csrc/libtorch_stable/sampler.cu" + "csrc/libtorch_stable/topk.cu" + "csrc/libtorch_stable/mamba/selective_scan_fwd.cu" + "csrc/libtorch_stable/cache_kernels.cu" + "csrc/libtorch_stable/cache_kernels_fused.cu" + "csrc/libtorch_stable/custom_all_reduce.cu" + "csrc/libtorch_stable/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu") + + if(APHRODITE_GPU_LANG STREQUAL "CUDA" AND + DEFINED CMAKE_CUDA_COMPILER_VERSION AND + CMAKE_CUDA_COMPILER_VERSION VERSION_GREATER_EQUAL 12.0) + + if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0) + cuda_archs_loose_intersection(COOPERATIVE_TOPK_ARCHS + "9.0a;10.0f;10.1f;10.3f;10.7f;11.0f;12.0f;12.1f" "${CUDA_ARCHS}") + else() + cuda_archs_loose_intersection(COOPERATIVE_TOPK_ARCHS + "9.0a;10.0a;10.1a;10.3a;12.0a;12.1a" "${CUDA_ARCHS}") + endif() + + if(COOPERATIVE_TOPK_ARCHS) + list(APPEND APHRODITE_GPU_FLAGS "-DAPHRODITE_ENABLE_COOPERATIVE_TOPK=1") + + endif() + endif() + + if(APHRODITE_GPU_LANG STREQUAL "CUDA") + SET(CUTLASS_ENABLE_HEADERS_ONLY ON CACHE BOOL "Enable only the header library") + + # Set CUTLASS_REVISION. Used for FetchContent. Also fixes some bogus messages when building. + set(CUTLASS_REVISION "v4.4.2") + + # Use the specified CUTLASS source directory for compilation if APHRODITE_CUTLASS_SRC_DIR is provided + if (DEFINED ENV{APHRODITE_CUTLASS_SRC_DIR}) + set(APHRODITE_CUTLASS_SRC_DIR $ENV{APHRODITE_CUTLASS_SRC_DIR}) + endif() + + if(APHRODITE_CUTLASS_SRC_DIR) + if(NOT IS_ABSOLUTE APHRODITE_CUTLASS_SRC_DIR) + get_filename_component(APHRODITE_CUTLASS_SRC_DIR "${APHRODITE_CUTLASS_SRC_DIR}" ABSOLUTE) + endif() + message(STATUS "The APHRODITE_CUTLASS_SRC_DIR is set, using ${APHRODITE_CUTLASS_SRC_DIR} for compilation") + FetchContent_Declare(cutlass SOURCE_DIR ${APHRODITE_CUTLASS_SRC_DIR}) + else() + FetchContent_Declare( + cutlass + GIT_REPOSITORY https://github.com/nvidia/cutlass.git + # Please keep this in sync with CUTLASS_REVISION line above. + GIT_TAG ${CUTLASS_REVISION} + GIT_PROGRESS TRUE + + # Speed up CUTLASS download by retrieving only the specified GIT_TAG instead of the history. + # Important: If GIT_SHALLOW is enabled then GIT_TAG works only with branch names and tags. + # So if the GIT_TAG above is updated to a commit hash, GIT_SHALLOW must be set to FALSE + GIT_SHALLOW TRUE + ) + endif() + FetchContent_MakeAvailable(cutlass) + + list(APPEND APHRODITE_STABLE_EXT_SRC + "csrc/libtorch_stable/cutlass_extensions/common.cpp" + "csrc/libtorch_stable/quantization/w8a8/cutlass/scaled_mm_entry.cu" + "csrc/libtorch_stable/quantization/fp4/nvfp4_quant_entry.cu" + "csrc/libtorch_stable/quantization/fp4/nvfp4_scaled_mm_entry.cu" + "csrc/libtorch_stable/quantization/awq/gemm_kernels.cu" + "csrc/libtorch_stable/minimax_reduce_rms_kernel.cu") + + # + # Machete kernels + # + # The machete kernels only work on hopper and require CUDA 12.0 or later. + # Only build Machete kernels if we are building for something compatible with sm90a + cuda_archs_loose_intersection(MACHETE_ARCHS "9.0a" "${CUDA_ARCHS}") + if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.0 AND MACHETE_ARCHS) + # + # For the Machete kernels we automatically generate sources for various + # preselected input type pairs and schedules. + # Generate sources: + set(MACHETE_GEN_SCRIPT + ${CMAKE_CURRENT_SOURCE_DIR}/csrc/libtorch_stable/quantization/machete/generate.py) + set(MACHETE_GEN_SUPPORT + ${CMAKE_CURRENT_SOURCE_DIR}/csrc/cutlass_extensions/aphrodite_cutlass_library_extension.py) + set(MACHETE_GEN_SENTINEL + ${CMAKE_CURRENT_SOURCE_DIR}/csrc/libtorch_stable/quantization/machete/generated/.complete) + file(MD5 ${MACHETE_GEN_SCRIPT} MACHETE_GEN_SCRIPT_HASH) + file(MD5 ${MACHETE_GEN_SUPPORT} MACHETE_GEN_SUPPORT_HASH) + string(APPEND MACHETE_GEN_SCRIPT_HASH "-${MACHETE_GEN_SUPPORT_HASH}") + + message(STATUS "Machete generation script hash: ${MACHETE_GEN_SCRIPT_HASH}") + message(STATUS "Last run machete generate script hash: $CACHE{MACHETE_GEN_SCRIPT_HASH}") + + if (NOT DEFINED CACHE{MACHETE_GEN_SCRIPT_HASH} + OR NOT $CACHE{MACHETE_GEN_SCRIPT_HASH} STREQUAL ${MACHETE_GEN_SCRIPT_HASH} + OR NOT EXISTS ${MACHETE_GEN_SENTINEL}) + execute_process( + COMMAND ${CMAKE_COMMAND} -E env + PYTHONPATH=${CMAKE_CURRENT_SOURCE_DIR}/csrc/cutlass_extensions/:${CUTLASS_DIR}/python/:${APHRODITE_PYTHON_PATH}:$ENV{PYTHONPATH} + ${Python_EXECUTABLE} ${MACHETE_GEN_SCRIPT} + RESULT_VARIABLE machete_generation_result + OUTPUT_VARIABLE machete_generation_output + OUTPUT_FILE ${CMAKE_CURRENT_BINARY_DIR}/machete_generation.log + ERROR_FILE ${CMAKE_CURRENT_BINARY_DIR}/machete_generation.log + ) + + if (NOT machete_generation_result EQUAL 0) + message(FATAL_ERROR "Machete generation failed." + " Result: \"${machete_generation_result}\"" + "\nCheck the log for details: " + "${CMAKE_CURRENT_BINARY_DIR}/machete_generation.log") + else() + set(MACHETE_GEN_SCRIPT_HASH ${MACHETE_GEN_SCRIPT_HASH} + CACHE STRING "Last run machete generate script hash" FORCE) + message(STATUS "Machete generation completed successfully.") + endif() + else() + message(STATUS "Machete generation script has not changed, skipping generation.") + endif() + + # Add machete generated sources + file(GLOB MACHETE_GEN_SOURCES "csrc/libtorch_stable/quantization/machete/generated/*.cu") + list(APPEND APHRODITE_STABLE_EXT_SRC ${MACHETE_GEN_SOURCES}) + + # forward compatible + set_gencode_flags_for_srcs( + SRCS "${MACHETE_GEN_SOURCES}" + CUDA_ARCHS "${MACHETE_ARCHS}") + + list(APPEND APHRODITE_STABLE_EXT_SRC + csrc/libtorch_stable/quantization/machete/machete_pytorch.cu) + message(STATUS "Building Machete kernels for archs: ${MACHETE_ARCHS}") + else() + if (NOT ${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.0 + AND MACHETE_ARCHS) + message(STATUS "Not building Machete kernels as CUDA Compiler version is " + "not >= 12.0, we recommend upgrading to CUDA 12.0 or " + "later if you intend on running w4a16 quantized models on " + "Hopper.") + else() + message(STATUS "Not building Machete kernels as no compatible archs " + "found in CUDA target architectures") + endif() + endif() + + set_gencode_flags_for_srcs( + SRCS "${APHRODITE_STABLE_EXT_SRC}" + CUDA_ARCHS "${CUDA_ARCHS}") + + # + # Swordfish kernels: w4a16 GEMM for the Blackwell sm100 family + # (datacenter sm100 + Thor sm110). Family (f) variants need CUDA >= 13.0; + # older toolchains fall back to arch-exact sm100. + # + if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0) + cuda_archs_loose_intersection(SWORDFISH_ARCHS "10.0f;11.0f" "${CUDA_ARCHS}") + else() + cuda_archs_loose_intersection(SWORDFISH_ARCHS "10.0a" "${CUDA_ARCHS}") + endif() + if(SWORDFISH_ARCHS) + set(SWORDFISH_SRCS + "csrc/libtorch_stable/quantization/swordfish/swordfish_prepack.cu" + "csrc/libtorch_stable/quantization/swordfish/swordfish_mm.cu" + "csrc/libtorch_stable/quantization/swordfish/swordfish_moe.cu" + "csrc/libtorch_stable/quantization/swordfish/swordfish_dense_tier.cu" + "csrc/libtorch_stable/quantization/swordfish/swordfish_prefill.cu" + "csrc/libtorch_stable/quantization/swordfish/swordfish_prefill_f16.cu") + set_gencode_flags_for_srcs( + SRCS "${SWORDFISH_SRCS}" + CUDA_ARCHS "${SWORDFISH_ARCHS}") + list(APPEND APHRODITE_STABLE_EXT_SRC "${SWORDFISH_SRCS}") + message(STATUS "Building Swordfish kernels for archs: ${SWORDFISH_ARCHS}") + else() + message(STATUS "Not building Swordfish kernels as no compatible archs " + "found in CUDA target architectures") + endif() + + if(COOPERATIVE_TOPK_ARCHS) + list(APPEND APHRODITE_STABLE_EXT_SRC + "csrc/libtorch_stable/cooperative_topk.cu") + set_gencode_flags_for_srcs( + SRCS "csrc/libtorch_stable/cooperative_topk.cu" + CUDA_ARCHS "${COOPERATIVE_TOPK_ARCHS}") + endif() + + # sm89 DeepSeek sparse attention (DSA) kernels; fp8 e4m3 mma codegen for + # sm_89 requires CUDA 12.4+. + if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.4) + cuda_archs_loose_intersection(DSA_SM89_ARCHS "8.9+PTX" "${CUDA_ARCHS}") + endif() + if(DSA_SM89_ARCHS) + set(SM89_DSA_SRC + "csrc/libtorch_stable/attention/sm89_dsa/sm89_fp8_paged_mqa_logits.cu" + "csrc/libtorch_stable/attention/sm89_dsa/sm89_fp8_mqa_logits.cu" + "csrc/libtorch_stable/attention/sm89_dsa/sm89_sparse_mla_fwd.cu") + list(APPEND APHRODITE_STABLE_EXT_SRC ${SM89_DSA_SRC}) + set_gencode_flags_for_srcs( + SRCS "${SM89_DSA_SRC}" + CUDA_ARCHS "${DSA_SM89_ARCHS}") + message(STATUS "Building sm89 DSA kernels for archs: ${DSA_SM89_ARCHS}") + endif() + + # Only build Marlin kernels if we are building for at least some compatible archs. + # Keep building Marlin for 9.0 as there are some group sizes and shapes that + # are not supported by Machete yet. + + # marlin arches for fp16 output + # Family-conditional 12.0f (one cubin for SM12x family) requires CUDA >= 13.0; + # fall back to architecture-specific 12.0a;12.1a on CUDA < 13.0 (e.g. 12.8). + if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0) + cuda_archs_loose_intersection(MARLIN_ARCHS "8.0+PTX;12.0f" "${CUDA_ARCHS}") + else() + cuda_archs_loose_intersection(MARLIN_ARCHS "8.0+PTX;12.0a;12.1a" "${CUDA_ARCHS}") + endif() + # marlin has limited support for turing + cuda_archs_loose_intersection(MARLIN_SM75_ARCHS "7.5" "${CUDA_ARCHS}") + # marlin arches for bf16 output (we need 9.0 for bf16 atomicAdd PTX) + if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0) + cuda_archs_loose_intersection(MARLIN_BF16_ARCHS "8.0+PTX;9.0+PTX;12.0f" "${CUDA_ARCHS}") + else() + cuda_archs_loose_intersection(MARLIN_BF16_ARCHS "8.0+PTX;9.0+PTX;12.0a;12.1a" "${CUDA_ARCHS}") + endif() + # marlin arches for fp8 input + # - sm80 doesn't support fp8 computation + # - sm90 and sm100 don't support QMMA.16832.F32.E4M3.E4M3 SAAS instruction + # so we only enable fp8 computation for SM89 (e.g. RTX 40x0) and 12.0 (e.g. RTX 50x0) + if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0) + cuda_archs_loose_intersection(MARLIN_FP8_ARCHS "8.9;12.0f" "${CUDA_ARCHS}") + else() + cuda_archs_loose_intersection(MARLIN_FP8_ARCHS "8.9;12.0a;12.1a" "${CUDA_ARCHS}") + endif() + # marlin arches for other files + cuda_archs_loose_intersection(MARLIN_OTHER_ARCHS "7.5;8.0+PTX" "${CUDA_ARCHS}") + + if (MARLIN_OTHER_ARCHS) + + # + # For the Marlin kernels we automatically generate sources for various + # preselected input type pairs and schedules. + # Generate sources: + set(MARLIN_GEN_SCRIPT + ${CMAKE_CURRENT_SOURCE_DIR}/csrc/libtorch_stable/quantization/marlin/generate_kernels.py) + file(MD5 ${MARLIN_GEN_SCRIPT} MARLIN_GEN_SCRIPT_HASH) + list(JOIN CUDA_ARCHS "," CUDA_ARCHS_STR) + set(MARLIN_GEN_SCRIPT_HASH_AND_ARCH "${MARLIN_GEN_SCRIPT_HASH}(ARCH:${CUDA_ARCHS_STR})") + + message(STATUS "Marlin generation script hash: ${MARLIN_GEN_SCRIPT_HASH_AND_ARCH}") + message(STATUS "Last run Marlin generate script hash: $CACHE{MARLIN_GEN_SCRIPT_HASH_AND_ARCH}") + + if (NOT DEFINED CACHE{MARLIN_GEN_SCRIPT_HASH_AND_ARCH} + OR NOT $CACHE{MARLIN_GEN_SCRIPT_HASH_AND_ARCH} STREQUAL ${MARLIN_GEN_SCRIPT_HASH_AND_ARCH} + OR NOT EXISTS + ${CMAKE_CURRENT_SOURCE_DIR}/csrc/libtorch_stable/quantization/marlin/kernel_selector.h) + execute_process( + COMMAND ${CMAKE_COMMAND} -E env + PYTHONPATH=$ENV{PYTHONPATH} + ${Python_EXECUTABLE} ${MARLIN_GEN_SCRIPT} ${CUDA_ARCHS_STR} + RESULT_VARIABLE marlin_generation_result + OUTPUT_VARIABLE marlin_generation_result + OUTPUT_FILE ${CMAKE_CURRENT_BINARY_DIR}/marlin_generation.log + ERROR_FILE ${CMAKE_CURRENT_BINARY_DIR}/marlin_generation.log + ) + + if (NOT marlin_generation_result EQUAL 0) + message(FATAL_ERROR "Marlin generation failed." + " Result: \"${marlin_generation_result}\"" + "\nCheck the log for details: " + "${CMAKE_CURRENT_BINARY_DIR}/marlin_generation.log") + else() + set(MARLIN_GEN_SCRIPT_HASH_AND_ARCH ${MARLIN_GEN_SCRIPT_HASH_AND_ARCH} + CACHE STRING "Last run Marlin generate script hash and arch" FORCE) + message(STATUS "Marlin generation completed successfully.") + endif() + else() + message(STATUS "Marlin generation script has not changed, skipping generation.") + endif() + + if (MARLIN_ARCHS) + file(GLOB MARLIN_TEMPLATE_KERNEL_SRC "csrc/libtorch_stable/quantization/marlin/sm80_kernel_*_float16.cu") + set_gencode_flags_for_srcs( + SRCS "${MARLIN_TEMPLATE_KERNEL_SRC}" + CUDA_ARCHS "${MARLIN_ARCHS}") + if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8) + set_source_files_properties(${MARLIN_TEMPLATE_KERNEL_SRC} + PROPERTIES COMPILE_FLAGS "-static-global-template-stub=false") + endif() + list(APPEND APHRODITE_STABLE_EXT_SRC ${MARLIN_TEMPLATE_KERNEL_SRC}) + + file(GLOB MARLIN_TEMPLATE_BF16_KERNEL_SRC "csrc/libtorch_stable/quantization/marlin/sm80_kernel_*_bfloat16.cu") + set_gencode_flags_for_srcs( + SRCS "${MARLIN_TEMPLATE_BF16_KERNEL_SRC}" + CUDA_ARCHS "${MARLIN_BF16_ARCHS}") + if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8) + set_source_files_properties(${MARLIN_TEMPLATE_BF16_KERNEL_SRC} + PROPERTIES COMPILE_FLAGS "-static-global-template-stub=false") + endif() + list(APPEND APHRODITE_STABLE_EXT_SRC ${MARLIN_TEMPLATE_BF16_KERNEL_SRC}) + endif() + + if (MARLIN_SM75_ARCHS) + file(GLOB MARLIN_TEMPLATE_SM75_KERNEL_SRC "csrc/libtorch_stable/quantization/marlin/sm75_kernel_*.cu") + set_gencode_flags_for_srcs( + SRCS "${MARLIN_TEMPLATE_SM75_KERNEL_SRC}" + CUDA_ARCHS "${MARLIN_SM75_ARCHS}") + if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8) + set_source_files_properties(${MARLIN_TEMPLATE_SM75_KERNEL_SRC} + PROPERTIES COMPILE_FLAGS "-static-global-template-stub=false") + endif() + list(APPEND APHRODITE_STABLE_EXT_SRC ${MARLIN_TEMPLATE_SM75_KERNEL_SRC}) + endif() + + if (MARLIN_FP8_ARCHS) + file(GLOB MARLIN_TEMPLATE_FP8_KERNEL_SRC "csrc/libtorch_stable/quantization/marlin/sm89_kernel_*.cu") + set_gencode_flags_for_srcs( + SRCS "${MARLIN_TEMPLATE_FP8_KERNEL_SRC}" + CUDA_ARCHS "${MARLIN_FP8_ARCHS}") + if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8) + set_source_files_properties(${MARLIN_TEMPLATE_FP8_KERNEL_SRC} + PROPERTIES COMPILE_FLAGS "-static-global-template-stub=false") + endif() + list(APPEND APHRODITE_STABLE_EXT_SRC ${MARLIN_TEMPLATE_FP8_KERNEL_SRC}) + endif() + + set(MARLIN_SRCS + "csrc/libtorch_stable/quantization/marlin/marlin.cu" + "csrc/libtorch_stable/quantization/marlin/marlin_int4_fp8_preprocess.cu" + "csrc/libtorch_stable/quantization/marlin/gptq_marlin_repack.cu" + "csrc/libtorch_stable/quantization/marlin/awq_marlin_repack.cu") + set_gencode_flags_for_srcs( + SRCS "${MARLIN_SRCS}" + CUDA_ARCHS "${MARLIN_OTHER_ARCHS}") + if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8) + set_source_files_properties(${MARLIN_SRCS} + PROPERTIES COMPILE_FLAGS "-static-global-template-stub=false") + endif() + list(APPEND APHRODITE_STABLE_EXT_SRC "${MARLIN_SRCS}") + + message(STATUS "Building Marlin kernels for archs: ${MARLIN_OTHER_ARCHS}") + else() + message(STATUS "Not building Marlin kernels as no compatible archs found" + " in CUDA target architectures") + endif() + + # DeepSeek V3 fused A GEMM kernel (requires SM 9.0+, Hopper and later) + if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0) + cuda_archs_loose_intersection(DSV3_FUSED_A_GEMM_ARCHS "9.0a;10.0f;10.7f;11.0f;12.0f" "${CUDA_ARCHS}") + else() + cuda_archs_loose_intersection(DSV3_FUSED_A_GEMM_ARCHS "9.0a;10.0a;10.1a;10.3a;12.0a;12.1a" "${CUDA_ARCHS}") + endif() + if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.0 AND DSV3_FUSED_A_GEMM_ARCHS) + set(DSV3_FUSED_A_GEMM_SRCS "csrc/libtorch_stable/dsv3_fused_a_gemm.cu") + set_gencode_flags_for_srcs( + SRCS "${DSV3_FUSED_A_GEMM_SRCS}" + CUDA_ARCHS "${DSV3_FUSED_A_GEMM_ARCHS}") + list(APPEND APHRODITE_STABLE_EXT_SRC "${DSV3_FUSED_A_GEMM_SRCS}") + message(STATUS "Building dsv3_fused_a_gemm for archs: ${DSV3_FUSED_A_GEMM_ARCHS}") + else() + message(STATUS "Not building dsv3_fused_a_gemm as no compatible archs found " + "in CUDA target architectures.") + endif() + + # FP32 router GEMM (H=3072, E=256, M<=32). Requires SM90+ and CUDA >= 12.0. + cuda_archs_sm90plus(FP32_ROUTER_GEMM_ARCHS "${CUDA_ARCHS}") + if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.0 AND FP32_ROUTER_GEMM_ARCHS) + set(FP32_ROUTER_GEMM_SRCS + "csrc/libtorch_stable/fp32_router_gemm_entry.cu" + "csrc/libtorch_stable/fp32_router_gemm.cu") + set_gencode_flags_for_srcs( + SRCS "${FP32_ROUTER_GEMM_SRCS}" + CUDA_ARCHS "${FP32_ROUTER_GEMM_ARCHS}") + list(APPEND APHRODITE_STABLE_EXT_SRC "${FP32_ROUTER_GEMM_SRCS}") + message(STATUS "Building fp32_router_gemm for archs: ${FP32_ROUTER_GEMM_ARCHS}") + else() + message(STATUS "Not building fp32_router_gemm as no compatible archs found " + "(requires SM90+ and CUDA >= 12.0).") + endif() + + # Only build AllSpark kernels if we are building for at least some compatible archs. + # 8.0+PTX JIT-forwards onto Blackwell like marlin; the kernel is plain + # mma.sync + cp.async and its large-M path is dequant + cuBLAS, which owns + # channelwise w8a16 prefill on sm100/sm110. + cuda_archs_loose_intersection(ALLSPARK_ARCHS "8.0+PTX" "${CUDA_ARCHS}") + if (ALLSPARK_ARCHS) + set(ALLSPARK_SRCS + "csrc/libtorch_stable/quantization/gptq_allspark/allspark_repack.cu" + "csrc/libtorch_stable/quantization/gptq_allspark/allspark_qgemm_w8a16.cu") + set_gencode_flags_for_srcs( + SRCS "${ALLSPARK_SRCS}" + CUDA_ARCHS "${ALLSPARK_ARCHS}") + list(APPEND APHRODITE_STABLE_EXT_SRC "${ALLSPARK_SRCS}") + message(STATUS "Building AllSpark kernels for archs: ${ALLSPARK_ARCHS}") + else() + message(STATUS "Not building AllSpark kernels as no compatible archs found" + " in CUDA target architectures") + endif() + + # + # CUTLASS scaled_mm kernels (moved from _C to _C_stable_libtorch) + # + set(SCALED_MM_3X_ARCHS) + # The cutlass_scaled_mm kernels for Hopper (c3x, i.e. CUTLASS 3.x) require + # CUDA 12.0 or later + cuda_archs_loose_intersection(SCALED_MM_ARCHS "9.0a;" "${CUDA_ARCHS}") + if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.0 AND SCALED_MM_ARCHS) + set(SCALED_MM_SM90_SRCS + "csrc/libtorch_stable/quantization/w8a8/cutlass/scaled_mm_c3x_sm90.cu" + "csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_sm90_fp8.cu" + "csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_sm90_int8.cu" + "csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_azp_sm90_int8.cu" + "csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm90_fp8.cu") + set_gencode_flags_for_srcs( + SRCS "${SCALED_MM_SM90_SRCS}" + CUDA_ARCHS "${SCALED_MM_ARCHS}") + list(APPEND APHRODITE_STABLE_EXT_SRC "${SCALED_MM_SM90_SRCS}") + list(APPEND APHRODITE_GPU_FLAGS "-DENABLE_SCALED_MM_SM90=1") + # Let scaled_mm_c2x know it doesn't need to build these arches + list(APPEND SCALED_MM_3X_ARCHS "${SCALED_MM_ARCHS}") + message(STATUS "Building scaled_mm_c3x_sm90 for archs: ${SCALED_MM_ARCHS}") + else() + if (NOT ${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.0 AND SCALED_MM_ARCHS) + message(STATUS "Not building scaled_mm_c3x_sm90 as CUDA Compiler version is " + "not >= 12.0, we recommend upgrading to CUDA 12.0 or " + "later if you intend on running FP8 quantized models on " + "Hopper.") + else() + message(STATUS "Not building scaled_mm_c3x_sm90 as no compatible archs found " + "in CUDA target architectures") + endif() + endif() + + + # The cutlass_scaled_mm kernels for Blackwell SM12x (c3x, i.e. CUTLASS 3.x) require + # CUDA 12.8 or later + if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0) + cuda_archs_loose_intersection(SCALED_MM_ARCHS "12.0f" "${CUDA_ARCHS}") + else() + cuda_archs_loose_intersection(SCALED_MM_ARCHS "12.0a;12.1a" "${CUDA_ARCHS}") + endif() + if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND SCALED_MM_ARCHS) + set(SCALED_MM_SM120_SRCS + "csrc/libtorch_stable/quantization/w8a8/cutlass/scaled_mm_c3x_sm120.cu" + "csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_sm120_fp8.cu" + "csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm120_fp8.cu" + ) + set_gencode_flags_for_srcs( + SRCS "${SCALED_MM_SM120_SRCS}" + CUDA_ARCHS "${SCALED_MM_ARCHS}") + list(APPEND APHRODITE_STABLE_EXT_SRC "${SCALED_MM_SM120_SRCS}") + list(APPEND APHRODITE_GPU_FLAGS "-DENABLE_SCALED_MM_SM120=1") + # Let scaled_mm_c2x know it doesn't need to build these arches + list(APPEND SCALED_MM_3X_ARCHS "${SCALED_MM_ARCHS}") + message(STATUS "Building scaled_mm_c3x_sm120 for archs: ${SCALED_MM_ARCHS}") + else() + if (NOT ${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND SCALED_MM_ARCHS) + message(STATUS "Not building scaled_mm_c3x_sm120 as CUDA Compiler version is " + "not >= 12.8, we recommend upgrading to CUDA 12.8 or " + "later if you intend on running FP8 quantized models on " + "Blackwell.") + else() + message(STATUS "Not building scaled_mm_c3x_120 as no compatible archs found " + "in CUDA target architectures") + endif() + endif() + + + # The cutlass_scaled_mm kernels for Blackwell SM100 (c3x, i.e. CUTLASS 3.x) + # require CUDA 12.8 or later + if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0) + cuda_archs_loose_intersection(SCALED_MM_ARCHS "10.0f;10.7f;11.0f" "${CUDA_ARCHS}") + else() + cuda_archs_loose_intersection(SCALED_MM_ARCHS "10.0a;10.1a;10.3a" "${CUDA_ARCHS}") + endif() + if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND SCALED_MM_ARCHS) + set(SCALED_MM_SM100_SRCS + "csrc/libtorch_stable/quantization/w8a8/cutlass/scaled_mm_c3x_sm100.cu" + "csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_sm100_fp8.cu" + "csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm100_fp8.cu" + ) + set_gencode_flags_for_srcs( + SRCS "${SCALED_MM_SM100_SRCS}" + CUDA_ARCHS "${SCALED_MM_ARCHS}") + list(APPEND APHRODITE_STABLE_EXT_SRC "${SCALED_MM_SM100_SRCS}") + list(APPEND APHRODITE_GPU_FLAGS "-DENABLE_SCALED_MM_SM100=1") + # Let scaled_mm_c2x know it doesn't need to build these arches + list(APPEND SCALED_MM_3X_ARCHS "${SCALED_MM_ARCHS}") + message(STATUS "Building scaled_mm_c3x_sm100 for archs: ${SCALED_MM_ARCHS}") + else() + if (NOT ${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND SCALED_MM_ARCHS) + message(STATUS "Not building scaled_mm_c3x_sm100 as CUDA Compiler version is " + "not >= 12.8, we recommend upgrading to CUDA 12.8 or " + "later if you intend on running FP8 quantized models on " + "Blackwell.") + else() + message(STATUS "Not building scaled_mm_c3x_100 as no compatible archs found " + "in CUDA target architectures") + endif() + endif() + + # + # For the cutlass_scaled_mm kernels we want to build the c2x (CUTLASS 2.x) + # kernels for the remaining archs that are not already built for 3x. + # (Build 8.9 for FP8) + cuda_archs_loose_intersection(SCALED_MM_2X_ARCHS + "7.5;8.0;8.7;8.9+PTX" "${CUDA_ARCHS}") + # subtract out the archs that are already built for 3x + list(REMOVE_ITEM SCALED_MM_2X_ARCHS ${SCALED_MM_3X_ARCHS}) + if (SCALED_MM_2X_ARCHS) + set(SCALED_MM_C2X_SRCS "csrc/libtorch_stable/quantization/w8a8/cutlass/scaled_mm_c2x.cu") + set_gencode_flags_for_srcs( + SRCS "${SCALED_MM_C2X_SRCS}" + CUDA_ARCHS "${SCALED_MM_2X_ARCHS}") + list(APPEND APHRODITE_STABLE_EXT_SRC "${SCALED_MM_C2X_SRCS}") + list(APPEND APHRODITE_GPU_FLAGS "-DENABLE_SCALED_MM_C2X=1") + message(STATUS "Building scaled_mm_c2x for archs: ${SCALED_MM_2X_ARCHS}") + else() + if (SCALED_MM_3X_ARCHS) + message(STATUS "Not building scaled_mm_c2x as all archs are already built" + " for and covered by scaled_mm_c3x") + else() + message(STATUS "Not building scaled_mm_c2x as no compatible archs found " + "in CUDA target architectures") + endif() + endif() + + # + # CUTLASS MoE kernels (moved from _C to _C_stable_libtorch) + # + + # The MoE kernel cutlass_moe_mm requires CUDA 12.3 or later (and ONLY works + # on Hopper). get_cutlass_(batched_)moe_mm_data should only be compiled + # if it's possible to compile MoE kernels that use its output. + cuda_archs_loose_intersection(SCALED_MM_ARCHS "9.0a" "${CUDA_ARCHS}") + if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.3 AND SCALED_MM_ARCHS) + set(CUTLASS_MOE_SM90_SRCS "csrc/libtorch_stable/quantization/w8a8/cutlass/moe/grouped_mm_c3x_sm90.cu") + set_gencode_flags_for_srcs( + SRCS "${CUTLASS_MOE_SM90_SRCS}" + CUDA_ARCHS "${SCALED_MM_ARCHS}") + list(APPEND APHRODITE_STABLE_EXT_SRC "${CUTLASS_MOE_SM90_SRCS}") + list(APPEND APHRODITE_GPU_FLAGS "-DENABLE_CUTLASS_MOE_SM90=1") + message(STATUS "Building grouped_mm_c3x for archs: ${SCALED_MM_ARCHS}") + else() + if (NOT ${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.3 AND SCALED_MM_ARCHS) + message(STATUS "Not building grouped_mm_c3x kernels as CUDA Compiler version is " + "not >= 12.3, we recommend upgrading to CUDA 12.3 or later " + "if you intend on running FP8 quantized MoE models on Hopper.") + else() + message(STATUS "Not building grouped_mm_c3x as no compatible archs found " + "in CUDA target architectures.") + endif() + endif() + + if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0) + cuda_archs_loose_intersection(SCALED_MM_ARCHS "10.0f;10.7f;11.0f" "${CUDA_ARCHS}") + else() + cuda_archs_loose_intersection(SCALED_MM_ARCHS "10.0a;10.1a;10.3a" "${CUDA_ARCHS}") + endif() + if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND SCALED_MM_ARCHS) + set(CUTLASS_MOE_SM100_SRCS "csrc/libtorch_stable/quantization/w8a8/cutlass/moe/grouped_mm_c3x_sm100.cu") + set_gencode_flags_for_srcs( + SRCS "${CUTLASS_MOE_SM100_SRCS}" + CUDA_ARCHS "${SCALED_MM_ARCHS}") + list(APPEND APHRODITE_STABLE_EXT_SRC "${CUTLASS_MOE_SM100_SRCS}") + list(APPEND APHRODITE_GPU_FLAGS "-DENABLE_CUTLASS_MOE_SM100=1") + message(STATUS "Building grouped_mm_c3x for archs: ${SCALED_MM_ARCHS}") + else() + if (NOT ${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND SCALED_MM_ARCHS) + message(STATUS "Not building grouped_mm_c3x kernels as CUDA Compiler version is " + "not >= 12.8, we recommend upgrading to CUDA 12.8 or later " + "if you intend on running FP8 quantized MoE models on Blackwell.") + else() + message(STATUS "Not building grouped_mm_c3x as no compatible archs found " + "in CUDA target architectures.") + endif() + endif() + + # moe_data.cu is used by all CUTLASS MoE kernels. + if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0) + cuda_archs_loose_intersection(CUTLASS_MOE_DATA_ARCHS "9.0a;10.0f;10.7f;11.0f;12.0f" "${CUDA_ARCHS}") + else() + cuda_archs_loose_intersection(CUTLASS_MOE_DATA_ARCHS "9.0a;10.0a;10.1a;10.3a;12.0a;12.1a" "${CUDA_ARCHS}") + endif() + if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.3 AND CUTLASS_MOE_DATA_ARCHS) + set(CUTLASS_MOE_DATA_SRCS "csrc/libtorch_stable/quantization/w8a8/cutlass/moe/moe_data.cu") + set_gencode_flags_for_srcs( + SRCS "${CUTLASS_MOE_DATA_SRCS}" + CUDA_ARCHS "${CUTLASS_MOE_DATA_ARCHS}") + list(APPEND APHRODITE_STABLE_EXT_SRC "${CUTLASS_MOE_DATA_SRCS}") + message(STATUS "Building moe_data for archs: ${CUTLASS_MOE_DATA_ARCHS}") + else() + if (NOT ${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.3 AND CUTLASS_MOE_DATA_ARCHS) + message(STATUS "Not building moe_data as CUDA Compiler version is " + "not >= 12.3, we recommend upgrading to CUDA 12.3 or later " + "if you intend on running FP8 quantized MoE models on Hopper or Blackwell.") + else() + message(STATUS "Not building moe_data as no compatible archs found " + "in CUDA target architectures.") + endif() + endif() + + # + # FP4/NVFP4 kernels (moved from _C to _C_stable_libtorch) + # + + # SM12x FP4 kernels. These share some generic NVFP4 quantization entry + # sources with the SM10x/11x block below; set_gencode_flags_for_srcs appends + # per-source flags, so shared files accumulate both SM12x and SM10x/11x + # gencodes when both families are requested. + if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0) + cuda_archs_loose_intersection(FP4_SM120_ARCHS "12.0f" "${CUDA_ARCHS}") + else() + cuda_archs_loose_intersection(FP4_SM120_ARCHS "12.0a;12.1a" "${CUDA_ARCHS}") + endif() + if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND FP4_SM120_ARCHS) + set(FP4_SM120_SRCS + "csrc/libtorch_stable/quantization/fp4/nvfp4_quant_kernels.cu" + "csrc/libtorch_stable/quantization/fp4/activation_nvfp4_quant_fusion_kernels.cu" + "csrc/libtorch_stable/quantization/fp4/nvfp4_experts_quant.cu" + "csrc/libtorch_stable/quantization/fp4/nvfp4_scaled_mm_sm120_kernels.cu" + "csrc/libtorch_stable/quantization/fp4/nvfp4_blockwise_moe_kernel.cu" + "csrc/libtorch_stable/nvfp4_kv_cache_kernels.cu") + set_gencode_flags_for_srcs( + SRCS "${FP4_SM120_SRCS}" + CUDA_ARCHS "${FP4_SM120_ARCHS}") + list(APPEND APHRODITE_STABLE_EXT_SRC "${FP4_SM120_SRCS}") + list(APPEND APHRODITE_GPU_FLAGS "-DENABLE_NVFP4_SM120=1") + list(APPEND APHRODITE_GPU_FLAGS "-DENABLE_CUTLASS_MOE_SM120=1") + message(STATUS "Building SM12x NVFP4 for archs: ${FP4_SM120_ARCHS}") + else() + message(STATUS "Not building SM12x NVFP4 as no compatible archs were found.") + endif() + + # SM10x/11x FP4 kernels. MXFP4 experts quantization is currently compiled + # only in this block; SM12x has separate NVFP4 matmul/MoE kernels above. + if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0) + cuda_archs_loose_intersection(FP4_SM100_ARCHS "10.0f;10.7f;11.0f" "${CUDA_ARCHS}") + else() + cuda_archs_loose_intersection(FP4_SM100_ARCHS "10.0a;10.1a;10.3a" "${CUDA_ARCHS}") + endif() + if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND FP4_SM100_ARCHS) + set(FP4_SM100_SRCS + "csrc/libtorch_stable/quantization/fp4/nvfp4_quant_kernels.cu" + "csrc/libtorch_stable/quantization/fp4/activation_nvfp4_quant_fusion_kernels.cu" + "csrc/libtorch_stable/quantization/fp4/nvfp4_experts_quant.cu" + "csrc/libtorch_stable/quantization/fp4/nvfp4_scaled_mm_kernels.cu" + "csrc/libtorch_stable/quantization/fp4/nvfp4_blockwise_moe_kernel.cu" + "csrc/libtorch_stable/quantization/fp4/mxfp4_experts_quant.cu" + "csrc/libtorch_stable/quantization/fp4/mxfp4_blockwise_moe_kernel.cu" + "csrc/libtorch_stable/nvfp4_kv_cache_kernels.cu") + if(NOT ${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.9) + message(STATUS + "Building mxfp4_experts_quant unsupported stubs because CUDA compiler version is not >= 12.9 (found ${CMAKE_CUDA_COMPILER_VERSION}).") + endif() + set_gencode_flags_for_srcs( + SRCS "${FP4_SM100_SRCS}" + CUDA_ARCHS "${FP4_SM100_ARCHS}") + list(APPEND APHRODITE_STABLE_EXT_SRC "${FP4_SM100_SRCS}") + list(APPEND APHRODITE_GPU_FLAGS "-DENABLE_NVFP4_SM100=1") + list(APPEND APHRODITE_GPU_FLAGS "-DENABLE_CUTLASS_MOE_SM100=1") + message(STATUS "Building SM10x/11x NVFP4/MXFP4 for archs: ${FP4_SM100_ARCHS}") + else() + message(STATUS "Not building SM10x/11x NVFP4/MXFP4 as no compatible archs were found.") + endif() + + # + # W4A8 kernels (moved from _C to _C_stable_libtorch) + # + + # Only build W4A8 kernels if we are building for something compatible with sm90a + cuda_archs_loose_intersection(W4A8_ARCHS "9.0a" "${CUDA_ARCHS}") + if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.0 AND W4A8_ARCHS) + set(W4A8_SRCS + "csrc/libtorch_stable/quantization/cutlass_w4a8/w4a8_mm_entry.cu" + "csrc/libtorch_stable/quantization/cutlass_w4a8/w4a8_grouped_mm_entry.cu" + "csrc/libtorch_stable/quantization/cutlass_w4a8/w4a8_utils.cu" + ) + + set_gencode_flags_for_srcs( + SRCS "${W4A8_SRCS}" + CUDA_ARCHS "${W4A8_ARCHS}") + + list(APPEND APHRODITE_STABLE_EXT_SRC "${W4A8_SRCS}") + + message(STATUS "Building W4A8 kernels for archs: ${W4A8_ARCHS}") + else() + if (NOT ${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.0 + AND W4A8_ARCHS) + message(STATUS "Not building W4A8 kernels as CUDA Compiler version is " + "not >= 12.0, we recommend upgrading to CUDA 12.0 or " + "later if you intend on running w4a16 quantized models on " + "Hopper.") + else() + message(STATUS "Not building W4A8 kernels as no compatible archs " + "found in CUDA target architectures") + endif() + endif() + + # CUTLASS MLA Archs and flags. + # Runtime dispatch is gated in + # aphrodite/v1/attention/backends/mla/cutlass_mla.py. + if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0) + cuda_archs_loose_intersection(MLA_ARCHS "10.0f;10.7f;11.0f" "${CUDA_ARCHS}") + else() + cuda_archs_loose_intersection(MLA_ARCHS "10.0a;10.1a;10.3a" "${CUDA_ARCHS}") + endif() + if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND MLA_ARCHS) + set(CUTLASS_MLA_SRCS + "csrc/libtorch_stable/attention/mla/sm100_cutlass_mla_kernel.cu") + set_gencode_flags_for_srcs( + SRCS "${CUTLASS_MLA_SRCS}" + CUDA_ARCHS "${MLA_ARCHS}") + list(APPEND APHRODITE_STABLE_EXT_SRC "${CUTLASS_MLA_SRCS}") + list(APPEND APHRODITE_GPU_FLAGS "-DENABLE_CUTLASS_MLA=1") + # Add MLA-specific include directories only to MLA source files + set_source_files_properties(${CUTLASS_MLA_SRCS} + PROPERTIES INCLUDE_DIRECTORIES "${CUTLASS_DIR}/examples/77_blackwell_fmha;${CUTLASS_DIR}/examples/common") + message(STATUS "Building CUTLASS MLA for archs: ${MLA_ARCHS}") + else() + message(STATUS "Not building CUTLASS MLA as no compatible archs were found.") + # clear MLA_ARCHS + set(MLA_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}") + endif() + if(KIMI_K3_ATTN_RES_ARCHS) + set(KIMI_K3_ATTN_RES_SRC + "csrc/libtorch_stable/kimi_k3/attn_res_kernel.cu") + set_gencode_flags_for_srcs( + SRCS "${KIMI_K3_ATTN_RES_SRC}" + CUDA_ARCHS "${KIMI_K3_ATTN_RES_ARCHS}") + set_property(SOURCE ${KIMI_K3_ATTN_RES_SRC} APPEND PROPERTY + COMPILE_OPTIONS + "$<$:--expt-relaxed-constexpr;--expt-extended-lambda;--use_fast_math>") + list(APPEND APHRODITE_STABLE_EXT_SRC "${KIMI_K3_ATTN_RES_SRC}") + message(STATUS + "Building Kimi K3 AttnRes for archs: ${KIMI_K3_ATTN_RES_ARCHS}") + endif() + + # Hadacore kernels + cuda_archs_loose_intersection(HADACORE_ARCHS "8.0+PTX;9.0+PTX" "${CUDA_ARCHS}") + if(HADACORE_ARCHS) + set(HADACORE_SRCS "csrc/libtorch_stable/quantization/hadamard/hadacore/hadamard_transform_cuda.cu") + set_gencode_flags_for_srcs( + SRCS "${HADACORE_SRCS}" + CUDA_ARCHS "${HADACORE_ARCHS}") + list(APPEND APHRODITE_STABLE_EXT_SRC "${HADACORE_SRCS}") + message(STATUS "Building hadacore") + endif() + + # if CUDA endif + endif() + + if(APHRODITE_GPU_LANG STREQUAL "CUDA") + set(APHRODITE_EXL3_SRC + "csrc/libtorch_stable/exl3_torch_bindings.cpp" + "csrc/quantization/exl3/exllamav3_ext/graph.cu" + "csrc/quantization/exl3/exllamav3_ext/hgemm.cu" + "csrc/quantization/exl3/exllamav3_ext/quant/exl3_devctx.cu" + "csrc/quantization/exl3/exllamav3_ext/quant/exl3_gemm.cu" + "csrc/quantization/exl3/exllamav3_ext/quant/exl3_gemv.cu" + "csrc/quantization/exl3/exllamav3_ext/quant/exl3_kernel_map.cu" + "csrc/quantization/exl3/exllamav3_ext/quant/exl3_moe.cu" + "csrc/quantization/exl3/exllamav3_ext/quant/hadamard.cu" + "csrc/quantization/exl3/exllamav3_ext/quant/reconstruct.cu" + "csrc/quantization/exl3/exllamav3_ext/quant/util.cu" + "csrc/quantization/exl3/exllamav3_ext/quant/comp_units/exl3_comp_unit_1.cu" + "csrc/quantization/exl3/exllamav3_ext/quant/comp_units/exl3_comp_unit_2.cu" + "csrc/quantization/exl3/exllamav3_ext/quant/comp_units/exl3_comp_unit_3.cu" + "csrc/quantization/exl3/exllamav3_ext/quant/comp_units/exl3_comp_unit_4.cu" + "csrc/quantization/exl3/exllamav3_ext/quant/comp_units/exl3_comp_unit_5.cu" + "csrc/quantization/exl3/exllamav3_ext/quant/comp_units/exl3_comp_unit_6.cu" + "csrc/quantization/exl3/exllamav3_ext/quant/comp_units/exl3_comp_unit_7.cu" + "csrc/quantization/exl3/exllamav3_ext/quant/comp_units/exl3_comp_unit_8.cu") + + # EXL3 requires Ampere+ instructions. Preserve all configured SM80+ + # architectures, including native Thor SM110 cubins. + set(EXL3_ARCHS) + foreach(_arch IN LISTS CUDA_ARCHS) + string(REGEX MATCH "^[0-9]+\\.[0-9]+" _arch_num "${_arch}") + if(_arch_num AND NOT _arch_num VERSION_LESS "8.0") + list(APPEND EXL3_ARCHS "${_arch}") + endif() + endforeach() + message(STATUS "EXL3_ARCHS (SM80+ subset of CUDA_ARCHS): ${EXL3_ARCHS}") + set_gencode_flags_for_srcs( + SRCS "${APHRODITE_EXL3_SRC}" + CUDA_ARCHS "${EXL3_ARCHS}") + list(APPEND APHRODITE_STABLE_EXT_SRC "${APHRODITE_EXL3_SRC}") + endif() + + message(STATUS "Enabling C_stable extension.") + list(REMOVE_DUPLICATES APHRODITE_STABLE_EXT_SRC) + define_extension_target( + _C_stable_libtorch + DESTINATION aphrodite + LANGUAGE ${APHRODITE_GPU_LANG} + SOURCES ${APHRODITE_STABLE_EXT_SRC} + COMPILE_FLAGS ${APHRODITE_GPU_FLAGS} + ARCHITECTURES ${APHRODITE_GPU_ARCHES} + INCLUDE_DIRECTORIES ${CUTLASS_INCLUDE_DIR} ${CUTLASS_TOOLS_UTIL_INCLUDE_DIR} + USE_SABI 3 + WITH_SOABI) + + # Set TORCH_TARGET_VERSION for stable ABI compatibility. + # This ensures we only use C-shim APIs available in PyTorch 2.11. + # _C_stable_libtorch is abi compatible with PyTorch >= TORCH_TARGET_VERSION + # which is currently set to 2.11. + target_compile_definitions(_C_stable_libtorch PRIVATE + TORCH_TARGET_VERSION=0x020B000000000000ULL) + + # Needed to use cuda/hip APIs from C-shim + if(APHRODITE_GPU_LANG STREQUAL "CUDA") + target_compile_definitions(_C_stable_libtorch PRIVATE USE_CUDA) + if(COOPERATIVE_TOPK_ARCHS) + target_compile_definitions(_C_stable_libtorch PRIVATE + APHRODITE_ENABLE_COOPERATIVE_TOPK=1) + endif() + if(DSA_SM89_ARCHS) + target_compile_definitions(_C_stable_libtorch PRIVATE + APHRODITE_ENABLE_SM89_DSA=1) + endif() + if(KIMI_K3_ATTN_RES_ARCHS) + target_compile_definitions(_C_stable_libtorch PRIVATE + APHRODITE_ENABLE_KIMI_K3_ATTN_RES=1) + endif() + # Needed by CUTLASS kernels + target_compile_definitions(_C_stable_libtorch PRIVATE + CUTLASS_ENABLE_DIRECT_CUDA_DRIVER_CALL=1) + elseif(APHRODITE_GPU_LANG STREQUAL "HIP") + target_compile_definitions(_C_stable_libtorch PRIVATE USE_ROCM) + endif() + + # On ROCm, _C_stable_libtorch calls raw HIP APIs (e.g. hipGetDevice in + # get_device_prop()) which must resolve to the same libamdhip64.so that + # PyTorch uses. When PyTorch bundles its own copy (pip/conda wheels), + # the raw HIP calls would otherwise resolve to the system ROCm copy, + # initializing a second HIP runtime that corrupts device state (wrong + # device on DeviceGuard, core dumps on multi-GPU tests). + # + # If PyTorch doesn't bundle libamdhip64 (built from source against system + # ROCm), there is only one copy in the process and no action is needed — + # the HIP compiler already links the system libamdhip64 automatically. + if(APHRODITE_GPU_LANG STREQUAL "HIP") + find_library(_STABLE_TORCH_AMDHIP64 amdhip64 + PATHS "${TORCH_INSTALL_PREFIX}/lib" NO_DEFAULT_PATH) + if(_STABLE_TORCH_AMDHIP64) + message(STATUS "Found PyTorch-bundled libamdhip64 at ${_STABLE_TORCH_AMDHIP64}") + target_link_libraries(_C_stable_libtorch PRIVATE ${_STABLE_TORCH_AMDHIP64}) + endif() + endif() +endif() + +# +# _moe_C_stable_libtorch extension +# + +set(APHRODITE_MOE_EXT_SRC + "csrc/libtorch_stable/moe/torch_bindings.cpp" + "csrc/libtorch_stable/moe/moe_align_sum_kernels.cu" + "csrc/libtorch_stable/moe/topk_softmax_kernels.cu" + "csrc/libtorch_stable/moe/topk_softplus_sqrt_kernels.cu") + +if(APHRODITE_GPU_LANG STREQUAL "CUDA") + list(APPEND APHRODITE_MOE_EXT_SRC + "csrc/libtorch_stable/moe/moe_wna16.cu" + "csrc/libtorch_stable/moe/grouped_topk_kernels.cu") +endif() + +if(APHRODITE_GPU_LANG STREQUAL "CUDA") + set(MOE_PERMUTE_SRC + "csrc/libtorch_stable/moe/permute_unpermute_kernels/moe_permute_unpermute_kernel.cu" + "csrc/libtorch_stable/moe/moe_permute_unpermute_op.cu") + + list(APPEND APHRODITE_MOE_EXT_SRC "${MOE_PERMUTE_SRC}") +endif() + +set_gencode_flags_for_srcs( + SRCS "${APHRODITE_MOE_EXT_SRC}" + CUDA_ARCHS "${CUDA_ARCHS}") + +if(APHRODITE_GPU_LANG STREQUAL "CUDA") + set(APHRODITE_MOE_WNA16_SRC + "csrc/libtorch_stable/moe/moe_wna16.cu") + + set_gencode_flags_for_srcs( + SRCS "${APHRODITE_MOE_WNA16_SRC}" + CUDA_ARCHS "${CUDA_ARCHS}") + + list(APPEND APHRODITE_MOE_EXT_SRC "${APHRODITE_MOE_WNA16_SRC}") + # moe marlin arches + # note that we always set `use_atomic_add=False` for moe marlin now, + # so we don't need 9.0 for bf16 atomicAdd PTX + if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0) + cuda_archs_loose_intersection(MARLIN_MOE_ARCHS "8.0+PTX;12.0f" "${CUDA_ARCHS}") + else() + cuda_archs_loose_intersection(MARLIN_MOE_ARCHS "8.0+PTX;12.0a;12.1a" "${CUDA_ARCHS}") + endif() + # moe marlin has limited support for turing + cuda_archs_loose_intersection(MARLIN_MOE_SM75_ARCHS "7.5" "${CUDA_ARCHS}") + # moe marlin arches for fp8 input + # - sm80 doesn't support fp8 computation + # - sm90 and sm100 don't support QMMA.16832.F32.E4M3.E4M3 SAAS instruction + # so we only enable fp8 computation for SM89 (e.g. RTX 40x0) and 12.0 (e.g. RTX 50x0) + if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0) + cuda_archs_loose_intersection(MARLIN_MOE_FP8_ARCHS "8.9;12.0f" "${CUDA_ARCHS}") + else() + cuda_archs_loose_intersection(MARLIN_MOE_FP8_ARCHS "8.9;12.0a;12.1a" "${CUDA_ARCHS}") + endif() + # moe marlin arches for other files + cuda_archs_loose_intersection(MARLIN_MOE_OTHER_ARCHS "7.5;8.0+PTX" "${CUDA_ARCHS}") + if (MARLIN_MOE_OTHER_ARCHS) + + # + # For the Marlin MOE kernels we automatically generate sources for various + # preselected input type pairs and schedules. + # Generate sources: + set(MOE_MARLIN_GEN_SCRIPT + ${CMAKE_CURRENT_SOURCE_DIR}/csrc/libtorch_stable/moe/marlin_moe_wna16/generate_kernels.py) + file(MD5 ${MOE_MARLIN_GEN_SCRIPT} MOE_MARLIN_GEN_SCRIPT_HASH) + list(JOIN CUDA_ARCHS "," CUDA_ARCHS_STR) + set(MOE_MARLIN_GEN_SCRIPT_HASH_AND_ARCH "${MOE_MARLIN_GEN_SCRIPT_HASH}(ARCH:${CUDA_ARCHS_STR})") + + message(STATUS "Marlin MOE generation script hash with arch: ${MOE_MARLIN_GEN_SCRIPT_HASH_AND_ARCH}") + message(STATUS "Last run Marlin MOE generate script hash with arch: $CACHE{MOE_MARLIN_GEN_SCRIPT_HASH_AND_ARCH}") + + if (NOT DEFINED CACHE{MOE_MARLIN_GEN_SCRIPT_HASH_AND_ARCH} + OR NOT $CACHE{MOE_MARLIN_GEN_SCRIPT_HASH_AND_ARCH} STREQUAL ${MOE_MARLIN_GEN_SCRIPT_HASH_AND_ARCH} + OR NOT EXISTS + ${CMAKE_CURRENT_SOURCE_DIR}/csrc/libtorch_stable/moe/marlin_moe_wna16/kernel_selector.h) + execute_process( + COMMAND ${CMAKE_COMMAND} -E env + PYTHONPATH=$ENV{PYTHONPATH} + ${Python_EXECUTABLE} ${MOE_MARLIN_GEN_SCRIPT} ${CUDA_ARCHS_STR} + RESULT_VARIABLE moe_marlin_generation_result + OUTPUT_VARIABLE moe_marlin_generation_output + OUTPUT_FILE ${CMAKE_CURRENT_BINARY_DIR}/moe_marlin_generation.log + ERROR_FILE ${CMAKE_CURRENT_BINARY_DIR}/moe_marlin_generation.log + ) + + if (NOT moe_marlin_generation_result EQUAL 0) + message(FATAL_ERROR "Marlin MOE generation failed." + " Result: \"${moe_marlin_generation_result}\"" + "\nCheck the log for details: " + "${CMAKE_CURRENT_BINARY_DIR}/moe_marlin_generation.log") + else() + set(MOE_MARLIN_GEN_SCRIPT_HASH_AND_ARCH ${MOE_MARLIN_GEN_SCRIPT_HASH_AND_ARCH} + CACHE STRING "Last run Marlin MOE generate script hash" FORCE) + message(STATUS "Marlin MOE generation completed successfully.") + endif() + else() + message(STATUS "Marlin MOE generation script has not changed, skipping generation.") + endif() + + if (MARLIN_MOE_ARCHS) + file(GLOB MARLIN_MOE_SRC "csrc/libtorch_stable/moe/marlin_moe_wna16/sm80_kernel_*.cu") + set_gencode_flags_for_srcs( + SRCS "${MARLIN_MOE_SRC}" + CUDA_ARCHS "${MARLIN_MOE_ARCHS}") + if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8) + set_source_files_properties(${MARLIN_MOE_SRC} + PROPERTIES COMPILE_FLAGS "-static-global-template-stub=false") + endif() + list(APPEND APHRODITE_MOE_EXT_SRC ${MARLIN_MOE_SRC}) + endif() + + if (MARLIN_MOE_SM75_ARCHS) + file(GLOB MARLIN_MOE_SM75_SRC "csrc/libtorch_stable/moe/marlin_moe_wna16/sm75_kernel_*.cu") + set_gencode_flags_for_srcs( + SRCS "${MARLIN_MOE_SM75_SRC}" + CUDA_ARCHS "${MARLIN_MOE_SM75_ARCHS}") + if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8) + set_source_files_properties(${MARLIN_MOE_SM75_SRC} + PROPERTIES COMPILE_FLAGS "-static-global-template-stub=false") + endif() + list(APPEND APHRODITE_MOE_EXT_SRC ${MARLIN_MOE_SM75_SRC}) + endif() + + if (MARLIN_MOE_FP8_ARCHS) + file(GLOB MARLIN_MOE_FP8_SRC "csrc/libtorch_stable/moe/marlin_moe_wna16/sm89_kernel_*.cu") + set_gencode_flags_for_srcs( + SRCS "${MARLIN_MOE_FP8_SRC}" + CUDA_ARCHS "${MARLIN_MOE_FP8_ARCHS}") + if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8) + set_source_files_properties(${MARLIN_MOE_FP8_SRC} + PROPERTIES COMPILE_FLAGS "-static-global-template-stub=false") + endif() + list(APPEND APHRODITE_MOE_EXT_SRC ${MARLIN_MOE_FP8_SRC}) + endif() + + set(MARLIN_MOE_OTHER_SRC "csrc/libtorch_stable/moe/marlin_moe_wna16/ops.cu") + set_gencode_flags_for_srcs( + SRCS "${MARLIN_MOE_OTHER_SRC}" + CUDA_ARCHS "${MARLIN_MOE_OTHER_ARCHS}") + if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8) + set_source_files_properties(${MARLIN_MOE_OTHER_SRC} + PROPERTIES COMPILE_FLAGS "-static-global-template-stub=false") + endif() + list(APPEND APHRODITE_MOE_EXT_SRC "${MARLIN_MOE_OTHER_SRC}") + + message(STATUS "Building Marlin MOE kernels for archs: ${MARLIN_MOE_OTHER_ARCHS}") + else() + message(STATUS "Not building Marlin MOE kernels as no compatible archs found" + " in CUDA target architectures") + endif() + + # DeepSeek V3 router GEMM kernel requires SM90+ and CUDA >= 12.0. + # (fp32_router_gemm has been migrated to _C_stable_libtorch above.) + cuda_archs_sm90plus(SM90PLUS_ROUTER_GEMM_ARCHS "${CUDA_ARCHS}") + if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.0 AND SM90PLUS_ROUTER_GEMM_ARCHS) + set(DSV3_ROUTER_GEMM_SRC + "csrc/libtorch_stable/moe/dsv3_router_gemm_entry.cu" + "csrc/libtorch_stable/moe/dsv3_router_gemm_float_out.cu" + "csrc/libtorch_stable/moe/dsv3_router_gemm_bf16_out.cu") + set_gencode_flags_for_srcs( + SRCS "${DSV3_ROUTER_GEMM_SRC}" + CUDA_ARCHS "${SM90PLUS_ROUTER_GEMM_ARCHS}") + list(APPEND APHRODITE_MOE_EXT_SRC "${DSV3_ROUTER_GEMM_SRC}") + + message(STATUS "Building DSV3 router GEMM kernels for archs: ${SM90PLUS_ROUTER_GEMM_ARCHS}") + else() + message(STATUS "Not building DSV3 router GEMM kernels as no compatible archs found" + " (requires SM90+ and CUDA >= 12.0)") + endif() +endif() + +message(STATUS "Enabling MoE C_stable extension.") +define_extension_target( + _moe_C_stable_libtorch + DESTINATION aphrodite + LANGUAGE ${APHRODITE_GPU_LANG} + SOURCES ${APHRODITE_MOE_EXT_SRC} + COMPILE_FLAGS ${APHRODITE_GPU_FLAGS} + ARCHITECTURES ${APHRODITE_GPU_ARCHES} + INCLUDE_DIRECTORIES ${CUTLASS_INCLUDE_DIR} + INCLUDE_DIRECTORIES ${CUTLASS_TOOLS_UTIL_INCLUDE_DIR} + USE_SABI 3 + WITH_SOABI) + +# Set TORCH_TARGET_VERSION for stable ABI compatibility. +# This ensures we only use C-shim APIs available in PyTorch 2.11. +# _moe_C_stable_libtorch is abi compatible with PyTorch >= TORCH_TARGET_VERSION +# which is currently set to 2.11. +target_compile_definitions(_moe_C_stable_libtorch PRIVATE + TORCH_TARGET_VERSION=0x020B000000000000ULL) + +# Needed to use cuda/hip APIs from C-shim +if(APHRODITE_GPU_LANG STREQUAL "CUDA") + target_compile_definitions(_moe_C_stable_libtorch PRIVATE USE_CUDA) + # Needed by CUTLASS kernels + target_compile_definitions(_moe_C_stable_libtorch PRIVATE + CUTLASS_ENABLE_DIRECT_CUDA_DRIVER_CALL=1) +elseif(APHRODITE_GPU_LANG STREQUAL "HIP") + target_compile_definitions(_moe_C_stable_libtorch PRIVATE USE_ROCM) +endif() + +# On ROCm, _moe_C_stable_libtorch calls raw HIP APIs (e.g. hipGetDevice in +# get_device_prop()) which must resolve to the same libamdhip64.so that +# PyTorch uses. When PyTorch bundles its own copy (pip/conda wheels), +# the raw HIP calls would otherwise resolve to the system ROCm copy, +# initializing a second HIP runtime that corrupts device state (wrong +# device on DeviceGuard, core dumps on multi-GPU tests). +# +# If PyTorch doesn't bundle libamdhip64 (built from source against system +# ROCm), there is only one copy in the process and no action is needed — +# the HIP compiler already links the system libamdhip64 automatically. +if(APHRODITE_GPU_LANG STREQUAL "HIP") + find_library(_MOE_STABLE_TORCH_AMDHIP64 amdhip64 + PATHS "${TORCH_INSTALL_PREFIX}/lib" NO_DEFAULT_PATH) + if(_MOE_STABLE_TORCH_AMDHIP64) + message(STATUS "Found PyTorch-bundled libamdhip64 for _moe_C_stable_libtorch at ${_MOE_STABLE_TORCH_AMDHIP64}") + target_link_libraries(_moe_C_stable_libtorch PRIVATE ${_MOE_STABLE_TORCH_AMDHIP64}) + endif() +endif() + +if(APHRODITE_GPU_LANG STREQUAL "HIP") + # + # _rocm_C extension + # + set(APHRODITE_ROCM_EXT_SRC + "csrc/rocm/torch_bindings.cpp" + "csrc/rocm/skinny_gemms.cu" + "csrc/rocm/skinny_gemms_int4.cu" + "csrc/rocm/attention.cu") + + set(APHRODITE_ROCM_HAS_GFX1100 OFF) + if(APHRODITE_GPU_ARCHES MATCHES "gfx1100") + set(APHRODITE_ROCM_HAS_GFX1100 ON) + list(APPEND APHRODITE_ROCM_EXT_SRC + "csrc/rocm/q_gemm_rdna3.cu" + "csrc/rocm/q_gemm_rdna3_wmma.cu" + "csrc/rocm/moe_q_gemm_rdna3.cu") + endif() + + define_extension_target( + _rocm_C + DESTINATION aphrodite + LANGUAGE ${APHRODITE_GPU_LANG} + SOURCES ${APHRODITE_ROCM_EXT_SRC} + COMPILE_FLAGS ${APHRODITE_GPU_FLAGS} + ARCHITECTURES ${APHRODITE_GPU_ARCHES} + USE_SABI 3 + WITH_SOABI) + + if(APHRODITE_ROCM_HAS_GFX1100) + target_compile_definitions(_rocm_C PRIVATE APHRODITE_ROCM_GFX1100) + endif() +endif() + +# Must run after the last HIP `define_extension_target` so every extension +# has registered its sources. +if (APHRODITE_GPU_LANG STREQUAL "HIP") + aphrodite_finalize_hipify_target() +endif() + +# For CUDA and HIP builds also build the triton_kernels external package. +if(APHRODITE_GPU_LANG STREQUAL "CUDA" OR APHRODITE_GPU_LANG STREQUAL "HIP") + include(cmake/external_projects/triton_kernels.cmake) +endif() + +# For CUDA we also build and ship some external projects. +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/qutlass.cmake) + include(cmake/external_projects/tml_fa4.cmake) + + # aphrodite-flash-attn should be last as it overwrites some CMake functions + include(cmake/external_projects/vllm_flash_attn.cmake) +endif () diff --git a/aphrodite/_custom_ops.py b/aphrodite/_custom_ops.py index c58512051e..682b291a00 100644 --- a/aphrodite/_custom_ops.py +++ b/aphrodite/_custom_ops.py @@ -3232,6 +3232,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 +3289,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 +3666,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 +3815,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"): diff --git a/aphrodite/_custom_ops.py.orig b/aphrodite/_custom_ops.py.orig index e653f533c1..c58512051e 100644 --- a/aphrodite/_custom_ops.py.orig +++ b/aphrodite/_custom_ops.py.orig @@ -4503,19 +4503,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 +4513,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/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/envs.py.orig b/aphrodite/envs.py.orig new file mode 100755 index 0000000000..5d2ebee86a --- /dev/null +++ b/aphrodite/envs.py.orig @@ -0,0 +1,1996 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import functools +import json +import logging +import os +import sys +import tempfile +import uuid +from collections.abc import Callable +from typing import TYPE_CHECKING, Any, Literal + +if TYPE_CHECKING: + APHRODITE_HOST_IP: str = "" + APHRODITE_PORT: int | None = None + APHRODITE_RPC_BASE_PATH: str = tempfile.gettempdir() + APHRODITE_USE_MODELSCOPE: bool = False + APHRODITE_USE_FASTOKENS: bool = False + APHRODITE_RINGBUFFER_WARNING_INTERVAL: int = 60 + APHRODITE_NCCL_SO_PATH: str | None = None + LD_LIBRARY_PATH: str | None = None + APHRODITE_ROCM_SLEEP_MEM_CHUNK_SIZE: int = 256 + LOCAL_RANK: int = 0 + CUDA_VISIBLE_DEVICES: str | None = None + APHRODITE_ENGINE_ITERATION_TIMEOUT_S: int = 60 + APHRODITE_ENGINE_READY_TIMEOUT_S: int = 600 + APHRODITE_API_KEY: str | None = None + APHRODITE_DEBUG_LOG_API_SERVER_RESPONSE: bool = False + S3_ACCESS_KEY_ID: str | None = None + S3_SECRET_ACCESS_KEY: str | None = None + S3_ENDPOINT_URL: str | None = None + APHRODITE_MODEL_REDIRECT_PATH: str | None = None + APHRODITE_CACHE_ROOT: str = os.path.expanduser("~/.cache/aphrodite") + APHRODITE_CONFIG_ROOT: str = os.path.expanduser("~/.config/aphrodite") + APHRODITE_USAGE_STATS_SERVER: str = "" + APHRODITE_NO_USAGE_STATS: bool = False + APHRODITE_DO_NOT_TRACK: bool = False + APHRODITE_USAGE_SOURCE: str = "production" + APHRODITE_CONFIGURE_LOGGING: bool = True + APHRODITE_LOGGING_LEVEL: str = "INFO" + APHRODITE_LOGGING_PREFIX: str = "" + APHRODITE_LOGGING_STREAM: str = "ext://sys.stdout" + APHRODITE_LOGGING_CONFIG_PATH: str | None = None + APHRODITE_LOGGING_COLOR: str = "auto" + APHRODITE_LOGGING_VERBOSE: bool = False + NO_COLOR: bool = False + APHRODITE_LOG_STATS_INTERVAL: float = 10.0 + APHRODITE_REQUEST_LEVEL_METRICS: bool = True + APHRODITE_TRACE_FUNCTION: int = 0 + APHRODITE_USE_FLASHINFER_SAMPLER: bool = True + APHRODITE_PP_LAYER_PARTITION: str | None = None + APHRODITE_CPU_KVCACHE_SPACE: int | None = 0 + APHRODITE_CPU_OMP_THREADS_BIND: str = "auto" + APHRODITE_CPU_NUM_OF_RESERVED_CPU: int | None = None + APHRODITE_CPU_SGL_KERNEL: bool = False + APHRODITE_CPU_ATTN_SPLIT_KV: bool = True + APHRODITE_ZENTORCH_WEIGHT_PREPACK: bool = True + APHRODITE_CPU_INT4_W4A8: bool = True + APHRODITE_XLA_CACHE_PATH: str = os.path.join(APHRODITE_CACHE_ROOT, "xla_cache") + APHRODITE_XLA_CHECK_RECOMPILATION: bool = False + APHRODITE_SPARSE_INDEXER_MAX_LOGITS_MB: int = 512 + APHRODITE_USE_RAY_COMPILED_DAG_CHANNEL_TYPE: Literal["auto", "nccl", "shm"] = "auto" + APHRODITE_USE_RAY_COMPILED_DAG_OVERLAP_COMM: bool = False + APHRODITE_USE_RAY_WRAPPED_PP_COMM: bool = True + APHRODITE_USE_RAY_V2_EXECUTOR_BACKEND: bool = False + APHRODITE_DISTRIBUTED_USE_SPLIT_GROUP: bool = False + APHRODITE_XLA_USE_SPMD: bool = False + APHRODITE_WORKER_MULTIPROC_METHOD: Literal["fork", "spawn"] = "fork" + APHRODITE_ASSETS_CACHE: str = os.path.join(APHRODITE_CACHE_ROOT, "assets") + APHRODITE_ASSETS_CACHE_MODEL_CLEAN: bool = False + APHRODITE_IMAGE_FETCH_TIMEOUT: int = 5 + APHRODITE_VIDEO_FETCH_TIMEOUT: int = 30 + APHRODITE_AUDIO_FETCH_TIMEOUT: int = 10 + APHRODITE_MEDIA_CACHE: str = "" + APHRODITE_MEDIA_CACHE_MAX_SIZE_MB: int = 5120 + APHRODITE_MEDIA_CACHE_TTL_HOURS: float = 24 + APHRODITE_MEDIA_FETCH_MAX_RETRIES: int = 3 + APHRODITE_MEDIA_URL_ALLOW_REDIRECTS: bool = True + APHRODITE_MEDIA_LOADING_THREAD_COUNT: int = 8 + APHRODITE_MAX_AUDIO_CLIP_FILESIZE_MB: int = 25 + APHRODITE_MAX_AUDIO_DECODE_DURATION_S: int = 600 + APHRODITE_MAX_AUDIO_PREPROCESS_WORKERS: int = max(1, min(os.cpu_count() or 1, 2)) + APHRODITE_MAX_IMAGE_PIXELS: int = 178_956_970 + APHRODITE_VIDEO_LOADER_BACKEND: str = "opencv" + APHRODITE_MEDIA_CONNECTOR: str = "http" + APHRODITE_MM_HASHER_ALGORITHM: str = "blake3" + APHRODITE_TARGET_DEVICE: str = "cuda" + APHRODITE_USE_PRECOMPILED: bool = False + APHRODITE_PRECOMPILED_WHEEL_LOCATION: str | None = None + APHRODITE_PRECOMPILED_WHEEL_COMMIT: str | None = None + APHRODITE_SKIP_PRECOMPILED_VERSION_SUFFIX: bool = False + APHRODITE_MAIN_CUDA_VERSION: str = "13.0" + APHRODITE_FLOAT32_MATMUL_PRECISION: Literal["highest", "high", "medium"] = "highest" + APHRODITE_BATCH_INVARIANT: bool = False + APHRODITE_TRITON_USE_TD: bool | None = None + # Deprecated alias of APHRODITE_TRITON_USE_TD (removed in v0.25). + APHRODITE_TRITON_ATTN_USE_TD: bool | None = None + APHRODITE_GPU_SYNC_CHECK: Literal["warn", "error"] | None = None + MAX_JOBS: str | None = None + NVCC_THREADS: str | None = None + APHRODITE_DOCKER_BUILD_CONTEXT: bool = False + APHRODITE_BUILD_COMMIT: str = "unknown" + APHRODITE_BUILD_PIPELINE: str = "local" + APHRODITE_BUILD_URL: str = "" + APHRODITE_IMAGE_TAG: str = "" + APHRODITE_KEEP_ALIVE_ON_ENGINE_DEATH: bool = False + CMAKE_BUILD_TYPE: Literal["Debug", "Release", "RelWithDebInfo"] | None = None + VERBOSE: bool = False + APHRODITE_ALLOW_LONG_MAX_MODEL_LEN: bool = False + APHRODITE_HTTP_TIMEOUT_KEEP_ALIVE: int = 5 # seconds + APHRODITE_MAX_N_SEQUENCES: int = 16384 + APHRODITE_MAX_COMPLETION_PROMPTS: int = 1024 + APHRODITE_PLUGINS: list[str] | None = None + APHRODITE_LORA_RESOLVER_CACHE_DIR: str | None = None + APHRODITE_LORA_RESOLVER_HF_REPO_LIST: str | None = None + APHRODITE_USE_AOT_COMPILE: bool = False + APHRODITE_USE_BYTECODE_HOOK: bool = True + APHRODITE_FORCE_AOT_LOAD: bool = False + APHRODITE_USE_MEGA_AOT_ARTIFACT: bool = False + APHRODITE_USE_TRITON_AWQ: bool = False + APHRODITE_FASTSAFETENSORS_QUEUE_SIZE: int = 0 + APHRODITE_TRITON_FORCE_FIRST_CONFIG: bool = False + APHRODITE_ALLOW_RUNTIME_LORA_UPDATING: bool = False + APHRODITE_SKIP_P2P_CHECK: bool = False + APHRODITE_DISABLED_KERNELS: list[str] = [] + APHRODITE_ENABLE_FLA_PACKED_RECURRENT_DECODE: bool = True + APHRODITE_DISABLE_PYNCCL: bool = False + APHRODITE_USE_OINK_OPS: bool = False + APHRODITE_MXFP8_EMULATION_DEQUANT_AT_LOAD: bool = True + APHRODITE_ROCM_USE_AITER: bool = False + APHRODITE_ROCM_USE_AITER_CUSTOM_AR: bool = True + APHRODITE_ROCM_USE_AITER_LINEAR: bool = True + APHRODITE_ROCM_USE_AITER_LINEAR_HIPBMM: bool = False + APHRODITE_ROCM_USE_AITER_MOE: bool = True + APHRODITE_ROCM_AITER_MOE_DISPATCH_POLICY: int = 0 + APHRODITE_ROCM_USE_AITER_RMSNORM: bool = True + APHRODITE_ROCM_USE_AITER_MLA: bool = True + APHRODITE_ROCM_USE_AITER_MHA: bool = True + APHRODITE_ROCM_USE_AITER_FP4_ASM_GEMM: bool = False + APHRODITE_ROCM_USE_AITER_TRITON_ROPE: bool = False + APHRODITE_ROCM_USE_AITER_FP8BMM: bool = True + APHRODITE_ROCM_USE_AITER_FP4BMM: bool = True + APHRODITE_ROCM_USE_AITER_UNIFIED_ATTENTION: bool = False + APHRODITE_ROCM_USE_AITER_FUSION_SHARED_EXPERTS: bool = False + APHRODITE_ROCM_USE_AITER_TRITON_GEMM: bool = True + APHRODITE_ROCM_USE_SKINNY_GEMM: bool = True + APHRODITE_ROCM_FP8_PADDING: bool = True + APHRODITE_ROCM_MOE_PADDING: bool = True + APHRODITE_ROCM_SHUFFLE_KV_CACHE_LAYOUT: bool = False + APHRODITE_ENABLE_V1_MULTIPROCESSING: bool = True + APHRODITE_LOG_BATCHSIZE_INTERVAL: float = -1 + APHRODITE_DISABLE_COMPILE_CACHE: bool = False + APHRODITE_USE_LAYERNAME: bool = True + Q_SCALE_CONSTANT: int = 200 + K_SCALE_CONSTANT: int = 200 + V_SCALE_CONSTANT: int = 100 + APHRODITE_USE_RUST_FRONTEND: bool = False + APHRODITE_USE_RUST_BENCH: bool = False + APHRODITE_RUST_FRONTEND_PATH: str | None = "auto" + APHRODITE_SERVER_DEV_MODE: bool = False + APHRODITE_V1_OUTPUT_PROC_CHUNK_SIZE: int = 128 + APHRODITE_MLA_DISABLE: bool = False + APHRODITE_RAY_PER_WORKER_GPUS: float = 1.0 + APHRODITE_RAY_BUNDLE_INDICES: str = "" + APHRODITE_CUDART_SO_PATH: str | None = None + APHRODITE_DP_RANK: int = 0 + APHRODITE_DP_RANK_LOCAL: int = -1 + APHRODITE_DP_SIZE: int = 1 + APHRODITE_USE_STANDALONE_COMPILE: bool = True + APHRODITE_ENABLE_PREGRAD_PASSES: bool = True + APHRODITE_USE_BREAKABLE_CUDAGRAPH: bool = False + APHRODITE_DP_MASTER_IP: str = "" + APHRODITE_DP_MASTER_PORT: int = 0 + APHRODITE_RANDOMIZE_DP_DUMMY_INPUTS: bool = False + APHRODITE_RAY_DP_PACK_STRATEGY: Literal["strict", "fill", "span"] = "strict" + APHRODITE_RAY_DP_PLACEMENT_NODE_IPS: str = "" + APHRODITE_RAY_EXTRA_ENV_VAR_PREFIXES_TO_COPY: str = "" + APHRODITE_RAY_EXTRA_ENV_VARS_TO_COPY: str = "" + APHRODITE_MARLIN_USE_ATOMIC_ADD: bool = False + APHRODITE_MARLIN_INPUT_DTYPE: Literal["int8", "fp8"] | None = None + APHRODITE_HUMMING_ONLINE_QUANT_CONFIG: dict[str, Any] | None = None + APHRODITE_HUMMING_INPUT_QUANT_CONFIG: dict[str, Any] | None = None + APHRODITE_HUMMING_USE_F16_ACCUM: bool = False + APHRODITE_HUMMING_MOE_GEMM_TYPE: Literal["indexed", "grouped", "auto"] | None = None + APHRODITE_DEEPEPLL_NVFP4_DISPATCH: bool = False + APHRODITE_V1_USE_OUTLINES_CACHE: bool = False + APHRODITE_TPU_USING_PATHWAYS: bool = False + APHRODITE_USE_DEEP_GEMM: bool = True + APHRODITE_MOE_USE_DEEP_GEMM: bool = True + APHRODITE_USE_DEEP_GEMM_E8M0: bool = True + APHRODITE_USE_DEEP_GEMM_TMA_ALIGNED_SCALES: bool = True + APHRODITE_DCP_Q_REPLICATE: bool = False + APHRODITE_DEEP_GEMM_WARMUP: Literal[ + "skip", + "full", + "relax", + ] = "relax" + APHRODITE_USE_FUSED_MOE_GROUPED_TOPK: bool = True + APHRODITE_MOE_SKIP_PADDING: bool = True + APHRODITE_BLOCKSCALE_FP8_GEMM_FLASHINFER: bool = True + APHRODITE_USE_FLASHINFER_MOE_INT4: bool = False + APHRODITE_FLASHINFER_AUTOTUNE_CACHE_DIR: str | None = None + APHRODITE_FLASHINFER_AUTOTUNE_SKIP_OPS: list[str] | None = None + APHRODITE_FLASHINFER_ALLREDUCE_BACKEND: Literal["auto", "trtllm", "mnnvl"] = "auto" + APHRODITE_FLASHINFER_WORKSPACE_BUFFER_SIZE: int = 394 * 1024 * 1024 + APHRODITE_XGRAMMAR_CACHE_MB: int = 0 + APHRODITE_REGEX_COMPILATION_TIMEOUT_S: int = 5 + APHRODITE_MSGPACK_ZERO_COPY_THRESHOLD: int = 256 + APHRODITE_ALLOW_INSECURE_SERIALIZATION: bool = False + APHRODITE_DISABLE_REQUEST_ID_RANDOMIZATION: bool = False + APHRODITE_NIXL_SIDE_CHANNEL_HOST: str = "localhost" + APHRODITE_NIXL_SIDE_CHANNEL_PORT: int = 5600 + APHRODITE_P2P_SIDE_CHANNEL_HOST: str = "localhost" + APHRODITE_P2P_SIDE_CHANNEL_PORT: int = 5710 + APHRODITE_EC_SIDE_CHANNEL_HOST: str = "localhost" + APHRODITE_EC_SIDE_CHANNEL_PORT: int = 5601 + APHRODITE_MOONCAKE_BOOTSTRAP_PORT: int = 8998 + APHRODITE_MOONCAKE_STORE_TIER_LOG: bool = False + APHRODITE_MOONCAKE_LOAD_RECV_THREADS: int = 1 + APHRODITE_MOONCAKE_DISK_STAGING_USABLE_RATIO: float = 0.9 + MOONCAKE_PREFERRED_SEGMENT: str | None = None + MOONCAKE_REQUESTER_LOCAL_HOSTNAME: str | None = None + APHRODITE_MAX_TOKENS_PER_EXPERT_FP4_MOE: int = 163840 + APHRODITE_TOOL_PARSE_REGEX_TIMEOUT_SECONDS: int = 1 + APHRODITE_ENFORCE_STRICT_TOOL_CALLING: bool = True + APHRODITE_MQ_MAX_CHUNK_BYTES_MB: int = 16 + APHRODITE_EXECUTE_MODEL_TIMEOUT_SECONDS: int = 300 + APHRODITE_WORKER_SHUTDOWN_TIMEOUT_SECONDS: int = 5 + APHRODITE_KV_CACHE_LAYOUT: Literal["NHD", "HND"] | None = None + APHRODITE_SSM_CONV_STATE_LAYOUT: Literal["SD", "DS"] | None = None + APHRODITE_COMPUTE_NANS_IN_LOGITS: bool = False + APHRODITE_ROCM_QUICK_REDUCE_QUANTIZATION: Literal["FP", "INT8", "INT6", "INT4", "INT3", "NONE"] = "NONE" + APHRODITE_ROCM_QUICK_REDUCE_CAST_BF16_TO_FP16: bool = True + APHRODITE_ROCM_QUICK_REDUCE_MAX_SIZE_BYTES_MB: int | None = None + APHRODITE_ROCM_QUICK_REDUCE_MIN_SIZE_BYTES_MB: int | None = None + APHRODITE_ROCM_QUICK_REDUCE_QUANTIZATION_MIN_SIZE_KB: int | None = None + APHRODITE_MOONCAKE_ABORT_REQUEST_TIMEOUT: int = 480 + APHRODITE_ENABLE_CUDAGRAPH_GC: bool = False + APHRODITE_LOOPBACK_IP: str = "" + APHRODITE_ALLOW_CHUNKED_LOCAL_ATTN_WITH_HYBRID_KV_CACHE: bool = True + APHRODITE_ENABLE_RESPONSES_API_STORE: bool = False + APHRODITE_HAS_FLASHINFER_CUBIN: bool = False + APHRODITE_ROCM_FP8_MFMA_PAGE_ATTN: bool = False + APHRODITE_ALLREDUCE_USE_SYMM_MEM: bool = True + APHRODITE_ALLREDUCE_USE_FLASHINFER: bool = False + APHRODITE_TUNED_CONFIG_FOLDER: str | None = None + APHRODITE_ENABLE_STARTUP_PLAN: bool = False + APHRODITE_GPT_OSS_SYSTEM_TOOL_MCP_LABELS: set[str] = set() + APHRODITE_USE_EXPERIMENTAL_PARSER_CONTEXT: bool = False + APHRODITE_GPT_OSS_HARMONY_SYSTEM_INSTRUCTIONS: bool = False + APHRODITE_SYSTEM_START_DATE: str | None = None + APHRODITE_TOOL_JSON_ERROR_AUTOMATIC_RETRY: bool = False + APHRODITE_CUSTOM_SCOPES_FOR_PROFILING: bool = False + APHRODITE_NVTX_SCOPES_FOR_PROFILING: bool = False + APHRODITE_KV_EVENTS_USE_INT_BLOCK_HASHES: bool = True + APHRODITE_OBJECT_STORAGE_SHM_BUFFER_NAME: str = "APHRODITE_OBJECT_STORAGE_SHM_BUFFER" + APHRODITE_DEEPEP_BUFFER_SIZE_MB: int = 1024 + APHRODITE_DEEPEP_HIGH_THROUGHPUT_FORCE_INTRA_NODE: bool = False + APHRODITE_DEEPEP_LOW_LATENCY_USE_MNNVL: bool = False + APHRODITE_DEEPEP_V2_ALLOW_HYBRID_MODE: bool = True + APHRODITE_DEEPEP_V2_PREFER_OVERLAP: bool = False + APHRODITE_DEEPEP_V2_ALLOW_MULTIPLE_REDUCTION: bool = False + APHRODITE_DBO_COMM_SMS: int = 20 + APHRODITE_PATTERN_MATCH_DEBUG: str | None = None + APHRODITE_DEBUG_DUMP_PATH: str | None = None + APHRODITE_ENABLE_INDUCTOR_MAX_AUTOTUNE: bool = True + APHRODITE_ENABLE_INDUCTOR_COORDINATE_DESCENT_TUNING: bool = True + APHRODITE_USE_NCCL_SYMM_MEM: bool = False + APHRODITE_NCCL_INCLUDE_PATH: str | None = None + APHRODITE_GC_DEBUG: str = "" + APHRODITE_DEBUG_WORKSPACE: 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 + APHRODITE_COMPILE_CACHE_SAVE_FORMAT: Literal["binary", "unpacked"] = "binary" + APHRODITE_USE_V2_MODEL_RUNNER: bool | None = None + APHRODITE_LOG_MODEL_INSPECTION: bool = False + APHRODITE_DEBUG_MFU_METRICS: bool = False + APHRODITE_WEIGHT_OFFLOADING_DISABLE_PIN_MEMORY: bool = False + APHRODITE_WEIGHT_OFFLOADING_DISABLE_UVA: bool = False + APHRODITE_WSL2_ENABLE_PIN_MEMORY: bool = False + APHRODITE_DISABLE_LOG_LOGO: bool = False + APHRODITE_DISABLE_SM89_DSA: bool = False + APHRODITE_LORA_DISABLE_PDL: bool = False + APHRODITE_ENABLE_CUDA_COMPATIBILITY: bool = False + APHRODITE_CUDA_COMPATIBILITY_PATH: str | None = None + APHRODITE_SKIP_MODEL_NAME_VALIDATION: bool = False + """If set, Aphrodite will skip model name validation in API requests. + This allows any model name to be accepted in the 'model' field of requests, + making the server model-name agnostic. Useful for proxy/gateway scenarios.""" + APHRODITE_LOG_ROUTES: Literal["off", "compact", "full"] = "compact" + APHRODITE_ELASTIC_EP_SCALE_UP_LAUNCH: bool = False + APHRODITE_ELASTIC_EP_DRAIN_REQUESTS: bool = False + APHRODITE_MEMORY_PROFILER_ESTIMATE_CUDAGRAPHS: bool = True + APHRODITE_NIXL_EP_MAX_NUM_RANKS: int = 32 + APHRODITE_XPU_ENABLE_XPU_GRAPH: bool = False + APHRODITE_XPU_USE_SAMPLER_KERNEL: bool = True + APHRODITE_LORA_ENABLE_DUAL_STREAM: bool = False + APHRODITE_GPU_NIC_PCIE_MAPPING: str = "" + APHRODITE_NIC_SELECTION_VARS: str = "" + APHRODITE_PREFIX_CACHE_RETENTION_INTERVAL: int | None = None + + +def get_default_cache_root(): + return os.getenv( + "XDG_CACHE_HOME", + os.path.join(os.path.expanduser("~"), ".cache"), + ) + + +def get_default_config_root(): + return os.getenv( + "XDG_CONFIG_HOME", + os.path.join(os.path.expanduser("~"), ".config"), + ) + + +def maybe_convert_int(value: str | None) -> int | None: + if value is None: + return None + return int(value) + + +def maybe_convert_bool(value: str | None) -> bool | None: + if value is None: + return None + return bool(int(value)) + + +def maybe_convert_json_str_or_file(value: str | None) -> dict[str, Any] | None: + if value is None: + return None + if os.path.exists(value): + with open(value) as f: + return json.load(f) + return json.loads(value) + + +def disable_compile_cache() -> bool: + return bool(int(os.getenv("APHRODITE_DISABLE_COMPILE_CACHE", "0"))) + + +def use_aot_compile() -> bool: + from aphrodite.utils.torch_utils import is_torch_equal_or_newer + + default_value = "1" if is_torch_equal_or_newer("2.10.0") and not disable_compile_cache() else "0" + + return os.environ.get("APHRODITE_USE_AOT_COMPILE", default_value) == "1" + + +def use_mega_aot_artifact(): + from aphrodite.utils.torch_utils import is_torch_equal_or_newer + + default_value = "1" if is_torch_equal_or_newer("2.12.0.dev") and use_aot_compile() else "0" + + return os.environ.get("APHRODITE_USE_MEGA_AOT_ARTIFACT", default_value) == "1" + + +def env_with_choices( + env_name: str, + default: str | None, + choices: list[str] | Callable[[], list[str]], + case_sensitive: bool = True, +) -> Callable[[], str | None]: + """ + Create a lambda that validates environment variable against allowed choices + + Args: + env_name: Name of the environment variable + default: Default value if not set (can be None) + choices: List of valid string options or callable that returns list + case_sensitive: Whether validation should be case sensitive + + Returns: + Lambda function for environment_variables dict + """ + + def _get_validated_env() -> str | None: + value = os.getenv(env_name) + if value is None: + return default + + # Resolve choices if it's a callable (for lazy loading) + actual_choices = choices() if callable(choices) else choices + + if not case_sensitive: + check_value = value.lower() + check_choices = [choice.lower() for choice in actual_choices] + else: + check_value = value + check_choices = actual_choices + + if check_value not in check_choices: + raise ValueError(f"Invalid value '{value}' for {env_name}. Valid options: {actual_choices}.") + + return value + + return _get_validated_env + + +def env_list_with_choices( + env_name: str, + default: list[str], + choices: list[str] | Callable[[], list[str]], + case_sensitive: bool = True, +) -> Callable[[], list[str]]: + """ + Create a lambda that validates environment variable + containing comma-separated values against allowed choices + + Args: + env_name: Name of the environment variable + default: Default list of values if not set + choices: List of valid string options or callable that returns list + case_sensitive: Whether validation should be case sensitive + + Returns: + Lambda function for environment_variables + dict that returns list of strings + """ + + def _get_validated_env_list() -> list[str]: + value = os.getenv(env_name) + if value is None: + return default + + # Split comma-separated values and strip whitespace + values = [v.strip() for v in value.split(",") if v.strip()] + + if not values: + return default + + # Resolve choices if it's a callable (for lazy loading) + actual_choices = choices() if callable(choices) else choices + + # Validate each value + for val in values: + if not case_sensitive: + check_value = val.lower() + check_choices = [choice.lower() for choice in actual_choices] + else: + check_value = val + check_choices = actual_choices + + if check_value not in check_choices: + raise ValueError(f"Invalid value '{val}' in {env_name}. Valid options: {actual_choices}.") + + return values + + return _get_validated_env_list + + +def env_set_with_choices( + env_name: str, + default: list[str], + choices: list[str] | Callable[[], list[str]], + case_sensitive: bool = True, +) -> Callable[[], set[str]]: + """ + Creates a lambda which that validates environment variable + containing comma-separated values against allowed choices which + returns choices as a set. + """ + + def _get_validated_env_set() -> set[str]: + return set(env_list_with_choices(env_name, default, choices, case_sensitive)()) + + return _get_validated_env_set + + +def get_aphrodite_port() -> int | None: + """Get the port from APHRODITE_PORT environment variable. + + Returns: + The port number as an integer if APHRODITE_PORT is set, None otherwise. + + Raises: + ValueError: If APHRODITE_PORT is a URI, suggest k8s service discovery issue. + """ + if "APHRODITE_PORT" not in os.environ: + return None + + port = os.getenv("APHRODITE_PORT", "0") + + try: + return int(port) + except ValueError as err: + from urllib3.util import parse_url + + parsed = parse_url(port) + if parsed.scheme: + raise ValueError( + f"APHRODITE_PORT '{port}' appears to be a URI. " + "This may be caused by a Kubernetes service discovery issue," + "check the warning in: https://docs.aphrodite.ai/en/latest/configuration/env_vars.html" + ) from None + raise ValueError(f"APHRODITE_PORT '{port}' must be a valid integer") from err + + +def get_env_or_set_default( + env_name: str, + default_factory: Callable[[], str], +) -> Callable[[], str]: + """ + Create a lambda that returns an environment variable value if set, + or generates and sets a default value using the provided factory function. + """ + + def _get_or_set_default() -> str: + value = os.getenv(env_name) + if value is not None: + return value + + default_value = default_factory() + os.environ[env_name] = default_value + return default_value + + return _get_or_set_default + + +# The start-* and end* here are used by the documentation generator +# to extract the used env vars. + +# --8<-- [start:env-vars-definition] + +logger = logging.getLogger(__name__) + + +def _deprecated_triton_attn_use_td() -> None: + """Warn that APHRODITE_TRITON_ATTN_USE_TD was renamed. + + The old name is ignored; APHRODITE_TRITON_USE_TD is the supported variable. + """ + if "APHRODITE_TRITON_ATTN_USE_TD" in os.environ: + logger.warning( + "APHRODITE_TRITON_ATTN_USE_TD is deprecated and will be removed in " + "v0.25. Use APHRODITE_TRITON_USE_TD instead." + ) + return None + + +def _resolve_rust_cli_path() -> str | None: + """Resolve the aphrodite-rs binary path. + + Returns None unless APHRODITE_USE_RUST_FRONTEND or + APHRODITE_USE_RUST_BENCH is enabled. + When enabled, resolves APHRODITE_RUST_FRONTEND_PATH ("auto" by default) + to the actual binary path. + """ + use_rust = bool(int(os.environ.get("APHRODITE_USE_RUST_FRONTEND", "0"))) or bool( + int(os.environ.get("APHRODITE_USE_RUST_BENCH", "0")) + ) + raw = os.environ.get("APHRODITE_RUST_FRONTEND_PATH", "auto") + + if not use_rust: + if os.environ.get("APHRODITE_RUST_FRONTEND_PATH") is not None: + logger.warning( + "APHRODITE_RUST_FRONTEND_PATH is set without enabling " + "APHRODITE_USE_RUST_FRONTEND or APHRODITE_USE_RUST_BENCH. " + "Set one of them to 1 to use the aphrodite-rs binary." + ) + return None + + if raw.lower() in ("auto", "1", "true"): + pkg_dir = os.path.dirname(os.path.abspath(__file__)) + candidate = os.path.join(pkg_dir, "aphrodite-rs") + if os.path.isfile(candidate) and os.access(candidate, os.X_OK): + return candidate + + raise FileNotFoundError( + "APHRODITE_RUST_FRONTEND_PATH=auto but the aphrodite-rs binary was " + f"not found at {candidate}. " + "Build with setuptools-rust or set the path explicitly." + ) + return raw + + +environment_variables: dict[str, Callable[[], Any]] = { + # ================== Installation Time Env Vars ================== + # Target device of Aphrodite, supporting [cuda (by default), + # rocm, cpu] + "APHRODITE_TARGET_DEVICE": lambda: os.getenv("APHRODITE_TARGET_DEVICE", "cuda").lower(), + # Build-time only: skip all native builds (CMake/CUDA and Rust); binaries must + # be downloaded from the matching main-branch wheel. + "APHRODITE_USE_PRECOMPILED": lambda: os.getenv("APHRODITE_USE_PRECOMPILED", "").lower() in ("1", "true", "yes"), + # Exact local path or URL to use instead of resolving a per-commit wheel. + "APHRODITE_PRECOMPILED_WHEEL_LOCATION": lambda: os.getenv("APHRODITE_PRECOMPILED_WHEEL_LOCATION"), + # Full main-branch commit SHA, or "nightly", used instead of the checkout's + # merge-base with main. + "APHRODITE_PRECOMPILED_WHEEL_COMMIT": lambda: os.getenv("APHRODITE_PRECOMPILED_WHEEL_COMMIT"), + # Preserve the source-derived version in packaging contexts that consume + # precompiled artifacts internally. + "APHRODITE_SKIP_PRECOMPILED_VERSION_SUFFIX": lambda: ( + os.getenv("APHRODITE_SKIP_PRECOMPILED_VERSION_SUFFIX", "").lower() in ("1", "true", "yes") + ), + # Main CUDA version of Aphrodite. This follows PyTorch but can be overridden. + "APHRODITE_MAIN_CUDA_VERSION": lambda: os.getenv("APHRODITE_MAIN_CUDA_VERSION", "").lower() or "13.0", + # Controls PyTorch float32 matmul precision mode within Aphrodite workers. + # Valid options mirror torch.set_float32_matmul_precision + "APHRODITE_FLOAT32_MATMUL_PRECISION": env_with_choices( + "APHRODITE_FLOAT32_MATMUL_PRECISION", + "highest", + ["highest", "high", "medium"], + case_sensitive=False, + ), + # Enable batch-invariant mode: deterministic results regardless of + # batch composition. Requires NVIDIA GPU with compute capability >= 9.0. + "APHRODITE_BATCH_INVARIANT": lambda: bool(int(os.getenv("APHRODITE_BATCH_INVARIANT", "0"))), + # Use tensor descriptors for Q/K/V loads and output stores in the + # Triton unified-attention kernel. Enables HW 2D block reads on + # Intel XPU; the non-TD branch is dead-code-eliminated at Triton + # compile time so other platforms see no overhead. Tri-state override: + # unset (default) lets the `triton_attn` backend auto-select per + # platform (currently auto-enabled on XPU only); ``1`` forces TD on; + # ``0`` forces TD off. Useful for A/B benchmarking the TD path. + "APHRODITE_TRITON_USE_TD": lambda: {"1": True, "0": False}.get(os.getenv("APHRODITE_TRITON_USE_TD", "").strip()), + # If set, enable PyTorch's GPU<->CPU synchronization debug mode around + # the worker's `execute_model` and `sample_tokens` calls. Valid values + # are "warn" (print a warning on each sync) or "error" (raise on sync). + # Unset disables the check. See `torch.cuda.set_sync_debug_mode`. + "APHRODITE_GPU_SYNC_CHECK": env_with_choices("APHRODITE_GPU_SYNC_CHECK", None, ["warn", "error"]), + # Deprecated: renamed to APHRODITE_TRITON_USE_TD. Kept registered so it + # does not trip the unknown-env-var check; warns on use and is otherwise + # ignored. + "APHRODITE_TRITON_ATTN_USE_TD": lambda: _deprecated_triton_attn_use_td(), + # Maximum number of compilation jobs to run in parallel. + # By default this is the number of CPUs + "MAX_JOBS": lambda: os.getenv("MAX_JOBS", None), + # Number of threads to use for nvcc + # By default this is 1. + # If set, `MAX_JOBS` will be reduced to avoid oversubscribing the CPU. + "NVCC_THREADS": lambda: os.getenv("NVCC_THREADS", None), + # Used to mark that setup.py is running in a Docker build context, + # in order to force the use of precompiled binaries. + "APHRODITE_DOCKER_BUILD_CONTEXT": lambda: ( + os.environ.get("APHRODITE_DOCKER_BUILD_CONTEXT", "").strip().lower() in ("1", "true") + ), + # Build provenance metadata embedded in official aphrodite-openai images. + # Set via Docker ENV at image build time; informational only. + "APHRODITE_BUILD_COMMIT": lambda: os.environ.get("APHRODITE_BUILD_COMMIT", "unknown"), + "APHRODITE_BUILD_PIPELINE": lambda: os.environ.get("APHRODITE_BUILD_PIPELINE", "local"), + "APHRODITE_BUILD_URL": lambda: os.environ.get("APHRODITE_BUILD_URL", ""), + "APHRODITE_IMAGE_TAG": lambda: os.environ.get("APHRODITE_IMAGE_TAG", ""), + # CMake build type + # If not set, defaults to "Debug" or "RelWithDebInfo" + # Available options: "Debug", "Release", "RelWithDebInfo" + "CMAKE_BUILD_TYPE": env_with_choices("CMAKE_BUILD_TYPE", None, ["Debug", "Release", "RelWithDebInfo"]), + # If set, aphrodite will print verbose logs during installation + "VERBOSE": lambda: bool(int(os.getenv("VERBOSE", "0"))), + # Root directory for Aphrodite configuration files + # Defaults to `~/.config/aphrodite` unless `XDG_CONFIG_HOME` is set + # Note that this not only affects how aphrodite finds its configuration files + # during runtime, but also affects how aphrodite installs its configuration + # files during **installation**. + "APHRODITE_CONFIG_ROOT": lambda: os.path.expanduser( + os.getenv( + "APHRODITE_CONFIG_ROOT", + os.path.join(get_default_config_root(), "aphrodite"), + ) + ), + # ================== Runtime Env Vars ================== + # Root directory for Aphrodite cache files + # Defaults to `~/.cache/aphrodite` unless `XDG_CACHE_HOME` is set + "APHRODITE_CACHE_ROOT": lambda: os.path.expanduser( + os.getenv( + "APHRODITE_CACHE_ROOT", + os.path.join(get_default_cache_root(), "aphrodite"), + ) + ), + # used in distributed environment to determine the ip address + # of the current node, when the node has multiple network interfaces. + # If you are using multi-node inference, you should set this differently + # on each node. + "APHRODITE_HOST_IP": lambda: os.getenv("APHRODITE_HOST_IP", ""), + # used in distributed environment to manually set the communication port + # Note: if APHRODITE_PORT is set, and some code asks for multiple ports, the + # APHRODITE_PORT will be used as the first port, and the rest will be generated + # by incrementing the APHRODITE_PORT value. + "APHRODITE_PORT": get_aphrodite_port, + # path used for ipc when the frontend api server is running in + # multi-processing mode to communicate with the backend engine process. + "APHRODITE_RPC_BASE_PATH": lambda: os.getenv("APHRODITE_RPC_BASE_PATH", tempfile.gettempdir()), + # If true, will load models from ModelScope instead of Hugging Face Hub. + # note that the value is true or false, not numbers + "APHRODITE_USE_MODELSCOPE": lambda: os.environ.get("APHRODITE_USE_MODELSCOPE", "False").lower() == "true", + # If true, replace the Rust BPE backend that powers HF fast tokenizers + # with the `fastokens` (https://github.com/crusoecloud/fastokens) shim. + # Available in Aphrodite v0.23.0 and later. If your installed Aphrodite + # version does not recognize this environment variable, upgrade Aphrodite + # before enabling the override. Applies to any tokenizer mode that loads an + # HF fast tokenizer (`hf`, `deepseek_v32`, `deepseek_v4`, …). The + # `fastokens` Python package must be installed. + "APHRODITE_USE_FASTOKENS": lambda: bool(int(os.getenv("APHRODITE_USE_FASTOKENS", "0"))), + # Interval in seconds to log a warning message when the ring buffer is full + "APHRODITE_RINGBUFFER_WARNING_INTERVAL": lambda: int(os.environ.get("APHRODITE_RINGBUFFER_WARNING_INTERVAL", "60")), + # path to cudatoolkit home directory, under which should be bin, include, + # and lib directories. + "CUDA_HOME": lambda: os.environ.get("CUDA_HOME", None), + # Path to the NCCL library file. It is needed because nccl>=2.19 brought + # by PyTorch contains a bug: https://github.com/NVIDIA/nccl/issues/1234 + "APHRODITE_NCCL_SO_PATH": lambda: os.environ.get("APHRODITE_NCCL_SO_PATH", None), + # when `APHRODITE_NCCL_SO_PATH` is not set, aphrodite will try to find the nccl + # library file in the locations specified by `LD_LIBRARY_PATH` + "LD_LIBRARY_PATH": lambda: os.environ.get("LD_LIBRARY_PATH", None), + # flag to control the chunk size (in MB) for sleeping memory allocations under ROCm + "APHRODITE_ROCM_SLEEP_MEM_CHUNK_SIZE": lambda: int(os.environ.get("APHRODITE_ROCM_SLEEP_MEM_CHUNK_SIZE", "256")), + # Feature flag to enable/disable Inductor standalone compile. + # In torch <= 2.7 we ignore this flag; in torch >= 2.9 this is + # enabled by default. + "APHRODITE_USE_STANDALONE_COMPILE": lambda: os.environ.get("APHRODITE_USE_STANDALONE_COMPILE", "1") == "1", + # Inductor's pre-grad passes don't do anything for Aphrodite. + # The pre-grad passes get run even on cache-hit and negatively impact + # aphrodite cold compile times by O(1s) + # Can remove this after the following issue gets fixed + # TODO(luka): maybe_inplace requires this + # https://github.com/pytorch/pytorch/issues/174502 + "APHRODITE_ENABLE_PREGRAD_PASSES": lambda: os.environ.get("APHRODITE_ENABLE_PREGRAD_PASSES", "1") == "1", + # Experimental: breakable cudagraph does not rely on torch.compile + "APHRODITE_USE_BREAKABLE_CUDAGRAPH": lambda: os.environ.get("APHRODITE_USE_BREAKABLE_CUDAGRAPH", "0") == "1", + # Debug pattern matching inside custom passes. + # Should be set to the fx.Node name (e.g. 'getitem_34' or 'scaled_mm_3'). + "APHRODITE_PATTERN_MATCH_DEBUG": lambda: os.environ.get("APHRODITE_PATTERN_MATCH_DEBUG", None), + # Dump fx graphs to the given directory. + # It will override CompilationConfig.debug_dump_path if set. + "APHRODITE_DEBUG_DUMP_PATH": lambda: os.environ.get("APHRODITE_DEBUG_DUMP_PATH", None), + # Feature flag to enable/disable AOT compilation. This will ensure + # compilation is done in warmup phase and the compilation will be + # reused in subsequent calls. + "APHRODITE_USE_AOT_COMPILE": use_aot_compile, + # Feature flag to enable/disable bytecode in + # TorchCompileWithNoGuardsWrapper. + "APHRODITE_USE_BYTECODE_HOOK": lambda: bool(int(os.environ.get("APHRODITE_USE_BYTECODE_HOOK", "1"))), + # Force aphrodite to always load AOT compiled models from disk. Failure + # to load will result in a hard error when this is enabled. + # Will be ignored when APHRODITE_USE_AOT_COMPILE is disabled. + "APHRODITE_FORCE_AOT_LOAD": lambda: os.environ.get("APHRODITE_FORCE_AOT_LOAD", "0") == "1", + # Enable loading compiled models directly from cached standalone compile artifacts + # without re-splitting graph modules. This reduces overhead during model + # loading by using reconstruct_serializable_fn_from_mega_artifact. + "APHRODITE_USE_MEGA_AOT_ARTIFACT": use_mega_aot_artifact, + # local rank of the process in the distributed setting, used to determine + # the GPU device id + "LOCAL_RANK": lambda: int(os.environ.get("LOCAL_RANK", "0")), + # used to control the visible devices in the distributed setting + "CUDA_VISIBLE_DEVICES": lambda: os.environ.get("CUDA_VISIBLE_DEVICES", None), + # timeout for each iteration in the engine + "APHRODITE_ENGINE_ITERATION_TIMEOUT_S": lambda: int(os.environ.get("APHRODITE_ENGINE_ITERATION_TIMEOUT_S", "60")), + # Timeout in seconds for waiting for engine cores to become ready + # during startup. Default is 600 seconds (10 minutes). + "APHRODITE_ENGINE_READY_TIMEOUT_S": lambda: int(os.environ.get("APHRODITE_ENGINE_READY_TIMEOUT_S", "600")), + # API key for Aphrodite API server + "APHRODITE_API_KEY": lambda: os.environ.get("APHRODITE_API_KEY", None), + # Whether to log responses from API Server for debugging + "APHRODITE_DEBUG_LOG_API_SERVER_RESPONSE": lambda: ( + os.environ.get("APHRODITE_DEBUG_LOG_API_SERVER_RESPONSE", "False").lower() == "true" + ), + # S3 access information, used for tensorizer to load model from S3 + "S3_ACCESS_KEY_ID": lambda: os.environ.get("S3_ACCESS_KEY_ID", None), + "S3_SECRET_ACCESS_KEY": lambda: os.environ.get("S3_SECRET_ACCESS_KEY", None), + "S3_ENDPOINT_URL": lambda: os.environ.get("S3_ENDPOINT_URL", None), + # Usage stats collection + "APHRODITE_USAGE_STATS_SERVER": lambda: os.environ.get("APHRODITE_USAGE_STATS_SERVER", ""), + "APHRODITE_NO_USAGE_STATS": lambda: os.environ.get("APHRODITE_NO_USAGE_STATS", "0") == "1", + "APHRODITE_DO_NOT_TRACK": lambda: ( + (os.environ.get("APHRODITE_DO_NOT_TRACK", None) or os.environ.get("DO_NOT_TRACK", None) or "0") == "1" + ), + "APHRODITE_USAGE_SOURCE": lambda: os.environ.get("APHRODITE_USAGE_SOURCE", "production"), + # Logging configuration + # If set to 0, aphrodite will not configure logging + # If set to 1, aphrodite will configure logging using the default configuration + # or the configuration file specified by APHRODITE_LOGGING_CONFIG_PATH + "APHRODITE_CONFIGURE_LOGGING": lambda: bool(int(os.getenv("APHRODITE_CONFIGURE_LOGGING", "1"))), + "APHRODITE_LOGGING_CONFIG_PATH": lambda: os.getenv("APHRODITE_LOGGING_CONFIG_PATH"), + # this is used for configuring the default logging level + "APHRODITE_LOGGING_LEVEL": lambda: os.getenv("APHRODITE_LOGGING_LEVEL", "INFO").upper(), + # this is used for configuring the default logging stream + "APHRODITE_LOGGING_STREAM": lambda: os.getenv("APHRODITE_LOGGING_STREAM", "ext://sys.stdout"), + # if set, APHRODITE_LOGGING_PREFIX will be prepended to all log messages + "APHRODITE_LOGGING_PREFIX": lambda: os.getenv("APHRODITE_LOGGING_PREFIX", ""), + # Controls colored logging output. Options: "auto" (default, colors when terminal), + # "1" (always use colors), "0" (never use colors) + "APHRODITE_LOGGING_COLOR": lambda: os.getenv("APHRODITE_LOGGING_COLOR", "auto"), + # Toggle the more verbose Aphrodite logging format (file/line info) instead + # of the compact info format. + "APHRODITE_LOGGING_VERBOSE": lambda: os.environ.get("APHRODITE_LOGGING_VERBOSE", "0").lower() in ("1", "true"), + # Standard unix flag for disabling ANSI color codes + "NO_COLOR": lambda: os.getenv("NO_COLOR", "0") != "0", + # If set, aphrodite will log stats at this interval in seconds + # If not set, aphrodite will log stats every 10 seconds. + "APHRODITE_LOG_STATS_INTERVAL": lambda: ( + val if (val := float(os.getenv("APHRODITE_LOG_STATS_INTERVAL", "10."))) > 0.0 else 10.0 + ), + # If set to 1, aphrodite will collect and emit per-request level metrics. + "APHRODITE_REQUEST_LEVEL_METRICS": lambda: bool(int(os.getenv("APHRODITE_REQUEST_LEVEL_METRICS", "1"))), + # Trace function calls + # If set to 1, aphrodite will trace function calls + # Useful for debugging + "APHRODITE_TRACE_FUNCTION": lambda: int(os.getenv("APHRODITE_TRACE_FUNCTION", "0")), + # Whether to use the FlashInfer top-k / top-p sampler on CUDA. Enabled + # by default when the hardware supports it — set to 0 to opt out + # explicitly, which forces the PyTorch-native (Triton for bs>=8) path. + "APHRODITE_USE_FLASHINFER_SAMPLER": lambda: ( + bool(int(os.environ["APHRODITE_USE_FLASHINFER_SAMPLER"])) + if "APHRODITE_USE_FLASHINFER_SAMPLER" in os.environ + else True + ), + # Pipeline stage partition strategy + "APHRODITE_PP_LAYER_PARTITION": lambda: os.getenv("APHRODITE_PP_LAYER_PARTITION", None), + # (CPU backend only) CPU key-value cache space. + # default is None and will be set as 4 GB + "APHRODITE_CPU_KVCACHE_SPACE": lambda: ( + int(os.getenv("APHRODITE_CPU_KVCACHE_SPACE", "0")) if "APHRODITE_CPU_KVCACHE_SPACE" in os.environ else None + ), + # (CPU backend only) CPU core ids bound by OpenMP threads, e.g., "0-31", + # "0,1,2", "0-31,33". CPU cores of different ranks are separated by '|'. + "APHRODITE_CPU_OMP_THREADS_BIND": lambda: os.getenv("APHRODITE_CPU_OMP_THREADS_BIND", "auto"), + # (CPU backend only) CPU cores not used by OMP threads . + # Those CPU cores will not be used by OMP threads of a rank. + "APHRODITE_CPU_NUM_OF_RESERVED_CPU": lambda: ( + int(os.getenv("APHRODITE_CPU_NUM_OF_RESERVED_CPU", "0")) + if "APHRODITE_CPU_NUM_OF_RESERVED_CPU" in os.environ + else None + ), + # (CPU backend only) whether to use SGL kernels, optimized for small batch. + "APHRODITE_CPU_SGL_KERNEL": lambda: bool(int(os.getenv("APHRODITE_CPU_SGL_KERNEL", "0"))), + # (CPU backend only) whether to enable attention spilt KV. + "APHRODITE_CPU_ATTN_SPLIT_KV": lambda: bool(int(os.getenv("APHRODITE_CPU_ATTN_SPLIT_KV", "1"))), + # (Zen CPU backend) eagerly prepack weights into ZenDNN blocked layout + # at model load time. Eliminates per-inference layout conversion overhead. + "APHRODITE_ZENTORCH_WEIGHT_PREPACK": lambda: bool(int(os.getenv("APHRODITE_ZENTORCH_WEIGHT_PREPACK", "1"))), + # (CPU backend only) whether to use SGLang INT4 W4A8 kernels for AWQ. + "APHRODITE_CPU_INT4_W4A8": lambda: bool(int(os.getenv("APHRODITE_CPU_INT4_W4A8", "1"))), + # If the env var is set, Ray Compiled Graph uses the specified + # channel type to communicate between workers belonging to + # different pipeline-parallel stages. + # Available options: + # - "auto": use the default channel type + # - "nccl": use NCCL for communication + # - "shm": use shared memory and gRPC for communication + "APHRODITE_USE_RAY_COMPILED_DAG_CHANNEL_TYPE": env_with_choices( + "APHRODITE_USE_RAY_COMPILED_DAG_CHANNEL_TYPE", "auto", ["auto", "nccl", "shm"] + ), + # If the env var is set, it enables GPU communication overlap + # (experimental feature) in Ray's Compiled Graph. + "APHRODITE_USE_RAY_COMPILED_DAG_OVERLAP_COMM": lambda: bool( + int(os.getenv("APHRODITE_USE_RAY_COMPILED_DAG_OVERLAP_COMM", "0")) + ), + # If the env var is set, it uses a Ray Communicator wrapping + # Aphrodite's pipeline parallelism communicator to interact with Ray's + # Compiled Graph. Otherwise, it uses Ray's NCCL communicator. + "APHRODITE_USE_RAY_WRAPPED_PP_COMM": lambda: bool(int(os.getenv("APHRODITE_USE_RAY_WRAPPED_PP_COMM", "1"))), + # When True and distributed_executor_backend="ray", use RayExecutorV2 + # (MQ-based) instead of RayDistributedExecutor (compiled-graph backend). + "APHRODITE_USE_RAY_V2_EXECUTOR_BACKEND": lambda: bool(int(os.getenv("APHRODITE_USE_RAY_V2_EXECUTOR_BACKEND", "1"))), + # When True, GroupCoordinator constructs its CPU/device subgroups via + # ``torch.distributed.split_group(backend=...)`` + # and ``init_distributed_environment`` initializes the default PG with + # mixed ``cpu:gloo,cuda:nccl`` backend + eager ``device_id`` binding. + "APHRODITE_DISTRIBUTED_USE_SPLIT_GROUP": lambda: bool(int(os.getenv("APHRODITE_DISTRIBUTED_USE_SPLIT_GROUP", "0"))), + # Use dedicated multiprocess context for workers. + # Both spawn and fork work + "APHRODITE_WORKER_MULTIPROC_METHOD": env_with_choices( + "APHRODITE_WORKER_MULTIPROC_METHOD", "fork", ["spawn", "fork"] + ), + # Path to the cache for storing downloaded assets + "APHRODITE_ASSETS_CACHE": lambda: os.path.expanduser( + os.getenv( + "APHRODITE_ASSETS_CACHE", + os.path.join(get_default_cache_root(), "aphrodite", "assets"), + ) + ), + # If the env var is set, we will clean model file in + # this path $APHRODITE_ASSETS_CACHE/model_streamer/$model_name + "APHRODITE_ASSETS_CACHE_MODEL_CLEAN": lambda: bool(int(os.getenv("APHRODITE_ASSETS_CACHE_MODEL_CLEAN", "0"))), + # Timeout for fetching images when serving multimodal models + # Default is 5 seconds + "APHRODITE_IMAGE_FETCH_TIMEOUT": lambda: int(os.getenv("APHRODITE_IMAGE_FETCH_TIMEOUT", "5")), + # Timeout for fetching videos when serving multimodal models + # Default is 30 seconds + "APHRODITE_VIDEO_FETCH_TIMEOUT": lambda: int(os.getenv("APHRODITE_VIDEO_FETCH_TIMEOUT", "30")), + # Timeout for fetching audio when serving multimodal models + # Default is 10 seconds + "APHRODITE_AUDIO_FETCH_TIMEOUT": lambda: int(os.getenv("APHRODITE_AUDIO_FETCH_TIMEOUT", "10")), + # Directory for caching media downloads (images, video, audio fetched + # from URLs during inference). Empty string disables caching. + "APHRODITE_MEDIA_CACHE": lambda: os.getenv("APHRODITE_MEDIA_CACHE", ""), + # Maximum cache size in MB. When exceeded, least-recently-used entries + # are evicted. Default is 5120 (5 GB). + "APHRODITE_MEDIA_CACHE_MAX_SIZE_MB": lambda: int(os.getenv("APHRODITE_MEDIA_CACHE_MAX_SIZE_MB", "5120")), + # Time-to-live in hours for cached media files. Entries older than this + # are evicted regardless of cache size. Default is 24 hours. + "APHRODITE_MEDIA_CACHE_TTL_HOURS": lambda: float(os.getenv("APHRODITE_MEDIA_CACHE_TTL_HOURS", "24")), + # Maximum number of retries for fetching media (images, audio, video) + # from URLs. Each retry quadruples the timeout. Default is 3. + "APHRODITE_MEDIA_FETCH_MAX_RETRIES": lambda: int(os.getenv("APHRODITE_MEDIA_FETCH_MAX_RETRIES", "3")), + # Whether to allow HTTP redirects when fetching from media URLs. + # Default to True + "APHRODITE_MEDIA_URL_ALLOW_REDIRECTS": lambda: bool(int(os.getenv("APHRODITE_MEDIA_URL_ALLOW_REDIRECTS", "1"))), + # Max number of workers for the thread pool handling + # media bytes loading. Set to 1 to disable parallel processing. + # Default is 8 + "APHRODITE_MEDIA_LOADING_THREAD_COUNT": lambda: int(os.getenv("APHRODITE_MEDIA_LOADING_THREAD_COUNT", "8")), + # Maximum filesize in MB for a single audio file when processing + # speech-to-text requests. Files larger than this will be rejected. + # Default is 25 MB + "APHRODITE_MAX_AUDIO_CLIP_FILESIZE_MB": lambda: int(os.getenv("APHRODITE_MAX_AUDIO_CLIP_FILESIZE_MB", "25")), + # Maximum decoded audio duration in seconds. Compressed audio files + # (e.g. OPUS at very low bitrate) can expand into gigabytes of float32 + # PCM. This limit is enforced *during* decoding so the memory is never + # allocated. Default is 600s (10 minutes). + "APHRODITE_MAX_AUDIO_DECODE_DURATION_S": lambda: int(os.getenv("APHRODITE_MAX_AUDIO_DECODE_DURATION_S", "600")), + # Maximum number of worker threads used for STT preprocessing. The default + # intentionally caps at 2 because that performed best in profiling. + # https://github.com/vllm-project/vllm/pull/44612#issuecomment-4662757781 + "APHRODITE_MAX_AUDIO_PREPROCESS_WORKERS": lambda: int( + os.getenv( + "APHRODITE_MAX_AUDIO_PREPROCESS_WORKERS", + str(max(1, min(os.cpu_count() or 1, 2))), + ) + ), + # Maximum decoded image size in pixels. Small compressed images can + # expand into gigabytes of raster memory. This limit is enforced before + # decoding so the memory is never allocated. Default matches PIL's + # built-in 2x decompression-bomb threshold (~179M pixels, ~680 MB RGB). + "APHRODITE_MAX_IMAGE_PIXELS": lambda: int(os.getenv("APHRODITE_MAX_IMAGE_PIXELS", "178956970")), + # Backend for Video IO — selects the frame-sampling algorithm. + # - "opencv": uniform sampling. + # - "opencv_dynamic": duration-aware dynamic sampling. + # - "identity": returns raw video bytes for model processor to handle. + # + # Custom backend implementations can be registered + # via `@VIDEO_LOADER_REGISTRY.register("my_custom_video_loader")` and + # imported at runtime. + # If a non-existing backend is used, an AssertionError will be thrown. + "APHRODITE_VIDEO_LOADER_BACKEND": lambda: os.getenv("APHRODITE_VIDEO_LOADER_BACKEND", "opencv"), + # Media connector implementation. + # - "http": Default connector that supports fetching media via HTTP. + # + # Custom implementations can be registered + # via `@MEDIA_CONNECTOR_REGISTRY.register("my_custom_media_connector")` and + # imported at runtime. + # If a non-existing backend is used, an AssertionError will be thrown. + "APHRODITE_MEDIA_CONNECTOR": lambda: os.getenv("APHRODITE_MEDIA_CONNECTOR", "http"), + # Hash algorithm for multimodal content hashing. + # - "blake3": Default, fast cryptographic hash (not FIPS 140-3 compliant) + # - "sha256": FIPS 140-3 compliant, widely supported + # - "sha512": FIPS 140-3 compliant, faster on 64-bit systems + # Use sha256 or sha512 for FIPS compliance in government/enterprise deployments + "APHRODITE_MM_HASHER_ALGORITHM": env_with_choices( + "APHRODITE_MM_HASHER_ALGORITHM", + "blake3", + ["blake3", "sha256", "sha512"], + case_sensitive=False, + ), + # Path to the XLA persistent cache directory. + # Only used for XLA devices such as TPUs. + "APHRODITE_XLA_CACHE_PATH": lambda: os.path.expanduser( + os.getenv( + "APHRODITE_XLA_CACHE_PATH", + os.path.join(get_default_cache_root(), "aphrodite", "xla_cache"), + ) + ), + # If set, assert on XLA recompilation after each execution step. + "APHRODITE_XLA_CHECK_RECOMPILATION": lambda: bool(int(os.getenv("APHRODITE_XLA_CHECK_RECOMPILATION", "0"))), + # Enable SPMD mode for TPU backend. + "APHRODITE_XLA_USE_SPMD": lambda: bool(int(os.getenv("APHRODITE_XLA_USE_SPMD", "0"))), + # Maximum size (in MB) for logits tensor in sparse MLA indexer prefill chunks. + # Bounds the [M, N] float32 logits tensor to prevent CUDA OOM. + # Default: 512 MB + "APHRODITE_SPARSE_INDEXER_MAX_LOGITS_MB": lambda: int(os.getenv("APHRODITE_SPARSE_INDEXER_MAX_LOGITS_MB", "512")), + # If set, the OpenAI API server will stay alive even after the underlying + # AsyncLLMEngine errors and stops serving requests + "APHRODITE_KEEP_ALIVE_ON_ENGINE_DEATH": lambda: bool(int(os.getenv("APHRODITE_KEEP_ALIVE_ON_ENGINE_DEATH", "0"))), + # If the env var APHRODITE_ALLOW_LONG_MAX_MODEL_LEN is set, it allows + # the user to specify a max sequence length greater than + # the max length derived from the model's config.json. + # To enable this, set APHRODITE_ALLOW_LONG_MAX_MODEL_LEN=1. + "APHRODITE_ALLOW_LONG_MAX_MODEL_LEN": lambda: ( + os.environ.get("APHRODITE_ALLOW_LONG_MAX_MODEL_LEN", "0").strip().lower() in ("1", "true") + ), + # If set, forces FP8 Marlin to be used for FP8 quantization regardless + # of the hardware support for FP8 compute. + "APHRODITE_TEST_FORCE_FP8_MARLIN": lambda: ( + os.environ.get("APHRODITE_TEST_FORCE_FP8_MARLIN", "0").strip().lower() in ("1", "true") + ), + "APHRODITE_TEST_FORCE_LOAD_FORMAT": lambda: os.getenv("APHRODITE_TEST_FORCE_LOAD_FORMAT", "dummy"), + # Queue size for fastsafetensors ParallelLoader pipelined weight + # loading. Peak load-time VRAM is roughly + # model_weights + (1 + queue_size) * shard_size. + # Default 0 preserves the non-pipelined memory footprint so this + # change does not shrink the loadable-model envelope. Set to 1 + # (or higher) to overlap producing the next shard's device buffer + # with the consumer copying the current shard into model params, + # at the cost of `queue_size` extra shard-sized buffers resident + # at peak during loading. + "APHRODITE_FASTSAFETENSORS_QUEUE_SIZE": lambda: int(os.getenv("APHRODITE_FASTSAFETENSORS_QUEUE_SIZE", "0")), + # Timeout in seconds for keeping HTTP connections alive in API server + "APHRODITE_HTTP_TIMEOUT_KEEP_ALIVE": lambda: int(os.environ.get("APHRODITE_HTTP_TIMEOUT_KEEP_ALIVE", "5")), + # Maximum allowed value for the `n` sampling parameter (number of output + # sequences per request). Limits resource consumption to prevent + # denial-of-service via excessively large fan-out. Default: 16384. + "APHRODITE_MAX_N_SEQUENCES": lambda: int(os.environ.get("APHRODITE_MAX_N_SEQUENCES", "16384")), + # Maximum number of prompts allowed in a single /v1/completions request + # when the prompt field is a list. Prevents unbounded fan-out of engine + # requests from a single API call. Default: 1024. + "APHRODITE_MAX_COMPLETION_PROMPTS": lambda: int(os.environ.get("APHRODITE_MAX_COMPLETION_PROMPTS", "1024")), + # a list of plugin names to load, separated by commas. + # if this is not set, it means all plugins will be loaded + # if this is set to an empty string, no plugins will be loaded + "APHRODITE_PLUGINS": lambda: ( + None if "APHRODITE_PLUGINS" not in os.environ else os.environ["APHRODITE_PLUGINS"].split(",") + ), + # Retain local sliding-window KV checkpoints for prefix caching. + # Unset (default) preserves the dense local checkpointing behavior. `0` + # retains only the latest completed prompt boundary. Positive values retain + # checkpoints at the specified interval boundaries (rounded up to the + # prefix-cache alignment). + # Applies to sliding-window attention for now but not yet Mamba/linear attention. + "APHRODITE_PREFIX_CACHE_RETENTION_INTERVAL": lambda: ( + int(os.environ["APHRODITE_PREFIX_CACHE_RETENTION_INTERVAL"]) + if "APHRODITE_PREFIX_CACHE_RETENTION_INTERVAL" in os.environ + else None + ), + # a local directory to look in for unrecognized LoRA adapters. + # only works if plugins are enabled and + # APHRODITE_ALLOW_RUNTIME_LORA_UPDATING is enabled. + "APHRODITE_LORA_RESOLVER_CACHE_DIR": lambda: os.getenv("APHRODITE_LORA_RESOLVER_CACHE_DIR", None), + # A remote HF repo(s) containing one or more LoRA adapters, which + # may be downloaded and leveraged as needed. Only works if plugins + # are enabled and APHRODITE_ALLOW_RUNTIME_LORA_UPDATING is enabled. + # Values should be comma separated. + "APHRODITE_LORA_RESOLVER_HF_REPO_LIST": lambda: os.getenv("APHRODITE_LORA_RESOLVER_HF_REPO_LIST", None), + # If set, Aphrodite will use Triton implementations of AWQ. + "APHRODITE_USE_TRITON_AWQ": lambda: bool(int(os.getenv("APHRODITE_USE_TRITON_AWQ", "0"))), + # If set, monkey-patch triton.runtime.autotuner.Autotuner.run to skip + # benchmarking and select the first valid config (walking past invalid + # ones). Used to eliminate autotuning variability when measuring kernel + # performance and applied before running any kernel. + "APHRODITE_TRITON_FORCE_FIRST_CONFIG": lambda: ( + os.environ.get("APHRODITE_TRITON_FORCE_FIRST_CONFIG", "0").strip().lower() in ("1", "true") + ), + # If set, allow loading or unloading lora adapters in runtime, + "APHRODITE_ALLOW_RUNTIME_LORA_UPDATING": lambda: ( + os.environ.get("APHRODITE_ALLOW_RUNTIME_LORA_UPDATING", "0").strip().lower() in ("1", "true") + ), + # We assume drivers can report p2p status correctly. + # If the program hangs when using custom allreduce, + # potantially caused by a bug in the driver (535 series), + # if might be helpful to set APHRODITE_SKIP_P2P_CHECK=0 + # so that Aphrodite can verify if p2p is actually working. + # See https://github.com/vllm-project/vllm/blob/a9b15c606fea67a072416ea0ea115261a2756058/aphrodite/distributed/device_communicators/custom_all_reduce_utils.py#L101-L108 for details. # noqa + "APHRODITE_SKIP_P2P_CHECK": lambda: os.getenv("APHRODITE_SKIP_P2P_CHECK", "1") == "1", + # List of quantization kernels that should be disabled, used for testing + # and performance comparisons. Currently only affects MPLinearKernel + # selection + # (kernels: MacheteLinearKernel, MarlinLinearKernel, ExllamaLinearKernel) + "APHRODITE_DISABLED_KERNELS": lambda: ( + [] if "APHRODITE_DISABLED_KERNELS" not in os.environ else os.environ["APHRODITE_DISABLED_KERNELS"].split(",") + ), + "APHRODITE_ENABLE_FLA_PACKED_RECURRENT_DECODE": lambda: bool( + int(os.getenv("APHRODITE_ENABLE_FLA_PACKED_RECURRENT_DECODE", "1")) + ), + # Disable pynccl (using torch.distributed instead) + "APHRODITE_DISABLE_PYNCCL": lambda: os.getenv("APHRODITE_DISABLE_PYNCCL", "False").lower() in ("true", "1"), + # Optional: enable external Oink custom ops (e.g., Blackwell RMSNorm). + # Disabled by default. + "APHRODITE_USE_OINK_OPS": lambda: os.getenv("APHRODITE_USE_OINK_OPS", "False").lower() in ("true", "1"), + # Disable aiter ops unless specifically enabled. + # Acts as a parent switch to enable the rest of the other operations. + # On hardware without a native MXFP8 kernel (e.g. ROCm gfx942 / MI300), the + # MXFP8 emulation path dequantizes weights MXFP8->BF16 once at load time and + # runs as a BF16 checkpoint (no per-step dequant). Set to 0 to fall back to + # per-step dequant: keeps the 1-byte MXFP8 weights (~half the weight memory) + # at the cost of dequantizing every forward step (much slower). Default on. + "APHRODITE_MXFP8_EMULATION_DEQUANT_AT_LOAD": lambda: ( + os.getenv("APHRODITE_MXFP8_EMULATION_DEQUANT_AT_LOAD", "True").lower() in ("true", "1") + ), + "APHRODITE_ROCM_USE_AITER": lambda: os.getenv("APHRODITE_ROCM_USE_AITER", "False").lower() in ("true", "1"), + # Use AITER's CustomAllreduce as the custom-allreduce backend inside + # Aphrodite's CudaCommunicator on ROCm. + "APHRODITE_ROCM_USE_AITER_CUSTOM_AR": lambda: ( + os.getenv("APHRODITE_ROCM_USE_AITER_CUSTOM_AR", "True").lower() in ("true", "1") + ), + # use aiter linear op if aiter ops are enabled + # The following list of related ops + # - scaled_mm (per-tensor / rowwise) + # - use aiter tuned gemms for unquantized gemms + "APHRODITE_ROCM_USE_AITER_LINEAR": lambda: ( + os.getenv("APHRODITE_ROCM_USE_AITER_LINEAR", "True").lower() in ("true", "1") + ), + "APHRODITE_ROCM_USE_AITER_LINEAR_HIPBMM": lambda: ( + os.getenv("APHRODITE_ROCM_USE_AITER_LINEAR_HIPBMM", "False").lower() in ("true", "1") + ), + # Whether to use aiter moe ops. + # By default is enabled. + "APHRODITE_ROCM_USE_AITER_MOE": lambda: os.getenv("APHRODITE_ROCM_USE_AITER_MOE", "True").lower() in ("true", "1"), + # MoE sorting dispatch policy for AITER fused MoE kernels. + # 0 = auto (default): single-pass for small batches, multi-pass + # for large batches + # 1 = always single-pass: one kernel launch, no workspace, + # may be preferred for low-concurrency decode workloads + # 2 = always multi-pass: can be faster for MoE-heavy models + # (e.g., +2-5% on Qwen3-Next, +1.5% on DeepSeek-V3 at TP4, + # see PR #39177 for benchmarks) + "APHRODITE_ROCM_AITER_MOE_DISPATCH_POLICY": lambda: int(os.getenv("APHRODITE_ROCM_AITER_MOE_DISPATCH_POLICY", "0")), + # use aiter rms norm op if aiter ops are enabled. + "APHRODITE_ROCM_USE_AITER_RMSNORM": lambda: ( + os.getenv("APHRODITE_ROCM_USE_AITER_RMSNORM", "True").lower() in ("true", "1") + ), + # Whether to use aiter mla ops. + # By default is enabled. + "APHRODITE_ROCM_USE_AITER_MLA": lambda: os.getenv("APHRODITE_ROCM_USE_AITER_MLA", "True").lower() in ("true", "1"), + # Whether to use aiter mha ops. + # By default is enabled. + "APHRODITE_ROCM_USE_AITER_MHA": lambda: os.getenv("APHRODITE_ROCM_USE_AITER_MHA", "True").lower() in ("true", "1"), + # Whether to use aiter fp4 gemm asm. + # By default is disabled. + "APHRODITE_ROCM_USE_AITER_FP4_ASM_GEMM": lambda: ( + os.getenv("APHRODITE_ROCM_USE_AITER_FP4_ASM_GEMM", "False").lower() in ("true", "1") + ), + # Whether to use aiter rope. + # By default is disabled. + "APHRODITE_ROCM_USE_AITER_TRITON_ROPE": lambda: ( + os.getenv("APHRODITE_ROCM_USE_AITER_TRITON_ROPE", "False").lower() in ("true", "1") + ), + # Whether to use aiter triton fp8 bmm kernel + # By default is enabled. + "APHRODITE_ROCM_USE_AITER_FP8BMM": lambda: ( + os.getenv("APHRODITE_ROCM_USE_AITER_FP8BMM", "True").lower() in ("true", "1") + ), + # Whether to use aiter triton fp4 bmm kernel + # By default is enabled. + "APHRODITE_ROCM_USE_AITER_FP4BMM": lambda: ( + os.getenv("APHRODITE_ROCM_USE_AITER_FP4BMM", "True").lower() in ("true", "1") + ), + # Use AITER triton unified attention for V1 attention + "APHRODITE_ROCM_USE_AITER_UNIFIED_ATTENTION": lambda: ( + os.getenv("APHRODITE_ROCM_USE_AITER_UNIFIED_ATTENTION", "False").lower() in ("true", "1") + ), + # Whether to use aiter fusion shared experts ops. + # By default is disabled. + "APHRODITE_ROCM_USE_AITER_FUSION_SHARED_EXPERTS": lambda: ( + os.getenv("APHRODITE_ROCM_USE_AITER_FUSION_SHARED_EXPERTS", "False").lower() in ("true", "1") + ), + # Whether to use aiter triton kernels for gemm ops. + # By default is enabled. + "APHRODITE_ROCM_USE_AITER_TRITON_GEMM": lambda: ( + os.getenv("APHRODITE_ROCM_USE_AITER_TRITON_GEMM", "True").lower() in ("true", "1") + ), + # use rocm skinny gemms + "APHRODITE_ROCM_USE_SKINNY_GEMM": lambda: ( + os.getenv("APHRODITE_ROCM_USE_SKINNY_GEMM", "True").lower() in ("true", "1") + ), + # Pad the fp8 weights to 256 bytes for ROCm + "APHRODITE_ROCM_FP8_PADDING": lambda: bool(int(os.getenv("APHRODITE_ROCM_FP8_PADDING", "1"))), + # Pad the weights for the moe kernel + "APHRODITE_ROCM_MOE_PADDING": lambda: bool(int(os.getenv("APHRODITE_ROCM_MOE_PADDING", "1"))), + # Whether to use the shuffled kv cache layout + "APHRODITE_ROCM_SHUFFLE_KV_CACHE_LAYOUT": lambda: ( + os.getenv("APHRODITE_ROCM_SHUFFLE_KV_CACHE_LAYOUT", "False").lower() in ("true", "1") + ), + # Custom quick allreduce kernel for MI3* cards + # Choice of quantization level: FP, INT8, INT6, INT4, INT3 or NONE + # Recommended for large models to get allreduce + "APHRODITE_ROCM_QUICK_REDUCE_QUANTIZATION": env_with_choices( + "APHRODITE_ROCM_QUICK_REDUCE_QUANTIZATION", + "NONE", + ["FP", "INT8", "INT6", "INT4", "INT3", "NONE"], + ), + # Custom quick allreduce kernel for MI3* cards + # Due to the lack of the bfloat16 asm instruction, bfloat16 + # kernels are slower than fp16, + # If environment variable is set to 1, the input is converted to fp16 + "APHRODITE_ROCM_QUICK_REDUCE_CAST_BF16_TO_FP16": lambda: ( + os.getenv("APHRODITE_ROCM_QUICK_REDUCE_CAST_BF16_TO_FP16", "True").lower() in ("true", "1") + ), + # Custom quick allreduce kernel for MI3* cards. + # Controls the maximum allowed number of data bytes(MB) for custom quick + # allreduce communication. + # Default: 2048 MB. + # Data exceeding this size will use either custom allreduce or RCCL + # communication. + "APHRODITE_ROCM_QUICK_REDUCE_MAX_SIZE_BYTES_MB": lambda: maybe_convert_int( + os.environ.get("APHRODITE_ROCM_QUICK_REDUCE_MAX_SIZE_BYTES_MB", None) + ), + # Custom quick allreduce kernel for MI3* cards. + # Controls the minimum allowed number of data bytes(MB) required to use + # custom quick allreduce communication. + # If unset, use the built-in threshold table. + "APHRODITE_ROCM_QUICK_REDUCE_MIN_SIZE_BYTES_MB": lambda: maybe_convert_int( + os.environ.get("APHRODITE_ROCM_QUICK_REDUCE_MIN_SIZE_BYTES_MB", None) + ), + # Controls the minimum tensor size (KB, where 1 KB = 1024 bytes) required + # to use the configured QuickReduce codec. Smaller tensors use FP + # QuickReduce. This does not affect QuickReduce eligibility. + "APHRODITE_ROCM_QUICK_REDUCE_QUANTIZATION_MIN_SIZE_KB": lambda: maybe_convert_int( + os.environ.get("APHRODITE_ROCM_QUICK_REDUCE_QUANTIZATION_MIN_SIZE_KB", None) + ), + # Divisor for dynamic query scale factor calculation for FP8 KV Cache + "Q_SCALE_CONSTANT": lambda: int(os.getenv("Q_SCALE_CONSTANT", "200")), + # Divisor for dynamic key scale factor calculation for FP8 KV Cache + "K_SCALE_CONSTANT": lambda: int(os.getenv("K_SCALE_CONSTANT", "200")), + # Divisor for dynamic value scale factor calculation for FP8 KV Cache + "V_SCALE_CONSTANT": lambda: int(os.getenv("V_SCALE_CONSTANT", "100")), + # If set, enable multiprocessing in LLM for the V1 code path. + "APHRODITE_ENABLE_V1_MULTIPROCESSING": lambda: bool(int(os.getenv("APHRODITE_ENABLE_V1_MULTIPROCESSING", "1"))), + "APHRODITE_LOG_BATCHSIZE_INTERVAL": lambda: float(os.getenv("APHRODITE_LOG_BATCHSIZE_INTERVAL", "-1")), + "APHRODITE_DISABLE_COMPILE_CACHE": disable_compile_cache, + # If set to "0", disable LayerName opaque type for layer_name + # parameters in custom ops. Defaults to enabled on torch >= 2.11. + "APHRODITE_USE_LAYERNAME": lambda: bool(int(os.getenv("APHRODITE_USE_LAYERNAME", "1"))), + # If set, use the Rust frontend binary instead of the Python API server + # process(es). + "APHRODITE_USE_RUST_FRONTEND": lambda: bool(int(os.getenv("APHRODITE_USE_RUST_FRONTEND", "0"))), + # If set, use the packaged Rust client for `aphrodite bench serve`. + "APHRODITE_USE_RUST_BENCH": lambda: bool(int(os.getenv("APHRODITE_USE_RUST_BENCH", "0"))), + # Path to the aphrodite-rs binary. Defaults to "auto" which discovers + # the binary installed with the aphrodite package. Used when + # APHRODITE_USE_RUST_FRONTEND=1 or APHRODITE_USE_RUST_BENCH=1. + "APHRODITE_RUST_FRONTEND_PATH": lambda: _resolve_rust_cli_path(), + # If set, aphrodite will run in development mode, which will enable + # some additional endpoints for developing and debugging, + # e.g. `/reset_prefix_cache` + "APHRODITE_SERVER_DEV_MODE": lambda: bool(int(os.getenv("APHRODITE_SERVER_DEV_MODE", "0"))), + # Controls the maximum number of requests to handle in a + # single asyncio task when processing per-token outputs in the + # V1 AsyncLLM interface. It is applicable when handling a high + # concurrency of streaming requests. + # Setting this too high can result in a higher variance of + # inter-message latencies. Setting it too low can negatively impact + # TTFT and overall throughput. + "APHRODITE_V1_OUTPUT_PROC_CHUNK_SIZE": lambda: int(os.getenv("APHRODITE_V1_OUTPUT_PROC_CHUNK_SIZE", "128")), + # If set, Aphrodite will disable the MLA attention optimizations. + "APHRODITE_MLA_DISABLE": lambda: bool(int(os.getenv("APHRODITE_MLA_DISABLE", "0"))), + # If set, Aphrodite will pick up the provided Flash Attention MLA + # Number of GPUs per worker in Ray, if it is set to be a fraction, + # it allows ray to schedule multiple actors on a single GPU, + # so that users can colocate other actors on the same GPUs as Aphrodite. + "APHRODITE_RAY_PER_WORKER_GPUS": lambda: float(os.getenv("APHRODITE_RAY_PER_WORKER_GPUS", "1.0")), + # Bundle indices for Ray, if it is set, it can control precisely + # which indices are used for the Ray bundle, for every worker. + # Format: comma-separated list of integers, e.g. "0,1,2,3" + "APHRODITE_RAY_BUNDLE_INDICES": lambda: os.getenv("APHRODITE_RAY_BUNDLE_INDICES", ""), + # In some system, find_loaded_library() may not work. So we allow users to + # specify the path through environment variable APHRODITE_CUDART_SO_PATH. + "APHRODITE_CUDART_SO_PATH": lambda: os.getenv("APHRODITE_CUDART_SO_PATH", None), + # Rank of the process in the data parallel setting + "APHRODITE_DP_RANK": lambda: int(os.getenv("APHRODITE_DP_RANK", "0")), + # Rank of the process in the data parallel setting. + # Defaults to APHRODITE_DP_RANK when not set. + "APHRODITE_DP_RANK_LOCAL": lambda: int( + os.getenv("APHRODITE_DP_RANK_LOCAL", sys.modules[__name__].APHRODITE_DP_RANK) + ), + # World size of the data parallel setting + "APHRODITE_DP_SIZE": lambda: int(os.getenv("APHRODITE_DP_SIZE", "1")), + # IP address of the master node in the data parallel setting + "APHRODITE_DP_MASTER_IP": lambda: os.getenv("APHRODITE_DP_MASTER_IP", "127.0.0.1"), + # Port of the master node in the data parallel setting + "APHRODITE_DP_MASTER_PORT": lambda: int(os.getenv("APHRODITE_DP_MASTER_PORT", "0")), + # Randomize inputs during dummy runs when using Data Parallel + "APHRODITE_RANDOMIZE_DP_DUMMY_INPUTS": lambda: os.environ.get("APHRODITE_RANDOMIZE_DP_DUMMY_INPUTS", "0") == "1", + # Strategy to pack the data parallel ranks for Ray. + # Available options: + # - "fill": + # for DP master node, allocate exactly data-parallel-size-local DP ranks, + # for non-master nodes, allocate as many DP ranks as can fit; + # - "strict": + # allocate exactly data-parallel-size-local DP ranks to each picked node; + # - "span": + # Should be used only when a single DP rank requires multiple nodes. + # allocate one DP rank over as many nodes as required for set world_size; + # This environment variable is ignored if data-parallel-backend is not Ray. + "APHRODITE_RAY_DP_PACK_STRATEGY": lambda: os.getenv("APHRODITE_RAY_DP_PACK_STRATEGY", "strict"), + # Optional comma-separated list of node IPs that Ray data-parallel + # placement groups may use. When set, create_dp_placement_groups only + # considers these nodes (the DP master node is always included). + # This environment variable is ignored if data-parallel-backend is not Ray. + "APHRODITE_RAY_DP_PLACEMENT_NODE_IPS": lambda: os.getenv("APHRODITE_RAY_DP_PLACEMENT_NODE_IPS", ""), + # Comma-separated *additional* prefixes of env vars to copy from the + # driver to Ray workers. These are merged with the built-in defaults + # defined in ``aphrodite.ray.ray_env`` (APHRODITE_, etc.). Example: "MYLIB_,OTHER_" + "APHRODITE_RAY_EXTRA_ENV_VAR_PREFIXES_TO_COPY": lambda: os.getenv( + "APHRODITE_RAY_EXTRA_ENV_VAR_PREFIXES_TO_COPY", "" + ), + # Comma-separated *additional* individual env var names to copy from + # the driver to Ray workers. Merged with the built-in defaults + # defined in ``aphrodite.ray.ray_env`` (PYTHONHASHSEED). + # Example: "MY_SECRET,MY_FLAG" + "APHRODITE_RAY_EXTRA_ENV_VARS_TO_COPY": lambda: os.getenv("APHRODITE_RAY_EXTRA_ENV_VARS_TO_COPY", ""), + # Use model_redirect to redirect the model name to a local folder. + # `model_redirect` can be a json file mapping the model between + # repo_id and local folder: + # {"meta-llama/Llama-3.2-1B": "/tmp/Llama-3.2-1B"} + # or a space separated values table file: + # meta-llama/Llama-3.2-1B /tmp/Llama-3.2-1B + "APHRODITE_MODEL_REDIRECT_PATH": lambda: os.environ.get("APHRODITE_MODEL_REDIRECT_PATH", None), + # Whether to use atomicAdd reduce in gptq/awq marlin kernel. + "APHRODITE_MARLIN_USE_ATOMIC_ADD": lambda: os.environ.get("APHRODITE_MARLIN_USE_ATOMIC_ADD", "0") == "1", + # The activation dtype for marlin kernel + "APHRODITE_MARLIN_INPUT_DTYPE": env_with_choices("APHRODITE_MARLIN_INPUT_DTYPE", None, ["int8", "fp8"]), + # The online quantization dtype for humming kernel + "APHRODITE_HUMMING_ONLINE_QUANT_CONFIG": lambda: maybe_convert_json_str_or_file( + os.environ.get("APHRODITE_HUMMING_ONLINE_QUANT_CONFIG", None) + ), + # The activation dtype config for humming kernel + "APHRODITE_HUMMING_INPUT_QUANT_CONFIG": lambda: maybe_convert_json_str_or_file( + os.environ.get("APHRODITE_HUMMING_INPUT_QUANT_CONFIG", None) + ), + # Whether to use fp16 accumulator mma + "APHRODITE_HUMMING_USE_F16_ACCUM": lambda: maybe_convert_bool( + os.environ.get("APHRODITE_HUMMING_USE_F16_ACCUM", "0") + ), + # Whether to use indexed gemm for humming moe + # if 1, force use indexed gemm + # if 0, force use grouped gemm + # if None, choose better gemm type automatically + "APHRODITE_HUMMING_MOE_GEMM_TYPE": lambda: os.environ.get("APHRODITE_HUMMING_MOE_GEMM_TYPE", None), + # Whether to use DeepEPLL kernels for NVFP4 quantization and dispatch method + # only supported on Blackwell GPUs and with + # https://github.com/deepseek-ai/DeepEP/pull/341 + "APHRODITE_DEEPEPLL_NVFP4_DISPATCH": lambda: bool(int(os.getenv("APHRODITE_DEEPEPLL_NVFP4_DISPATCH", "0"))), + # Whether to turn on the outlines cache for V1 + # This cache is unbounded and on disk, so it's not safe to use in + # an environment with potentially malicious users. + "APHRODITE_V1_USE_OUTLINES_CACHE": lambda: os.environ.get("APHRODITE_V1_USE_OUTLINES_CACHE", "0") == "1", + # Whether using Pathways + "APHRODITE_TPU_USING_PATHWAYS": lambda: bool("proxy" in os.getenv("JAX_PLATFORMS", "").lower()), + # Allow use of DeepGemm kernels for fused moe ops. + "APHRODITE_USE_DEEP_GEMM": lambda: bool(int(os.getenv("APHRODITE_USE_DEEP_GEMM", "1"))), + # Allow use of DeepGemm specifically for MoE fused ops (overrides only MoE). + "APHRODITE_MOE_USE_DEEP_GEMM": lambda: bool(int(os.getenv("APHRODITE_MOE_USE_DEEP_GEMM", "1"))), + # Whether to use E8M0 scaling when DeepGEMM is used on Blackwell GPUs. + "APHRODITE_USE_DEEP_GEMM_E8M0": lambda: bool(int(os.getenv("APHRODITE_USE_DEEP_GEMM_E8M0", "1"))), + # Whether to create TMA-aligned scale tensor when DeepGEMM is used. + "APHRODITE_USE_DEEP_GEMM_TMA_ALIGNED_SCALES": lambda: bool( + int(os.getenv("APHRODITE_USE_DEEP_GEMM_TMA_ALIGNED_SCALES", "1")) + ), + # Opt-in MLA DCP query replication: skip the decode query all-gather. + "APHRODITE_DCP_Q_REPLICATE": lambda: bool(int(os.getenv("APHRODITE_DCP_Q_REPLICATE", "0"))), + # DeepGemm JITs the kernels on-demand. The warmup attempts to make DeepGemm + # JIT all the required kernels before model execution so there is no + # JIT'ing in the hot-path. However, this warmup increases the engine + # startup time by a couple of minutes. + # Available options: + # - "skip" : Skip warmup. + # - "full" : Warmup deepgemm by running all possible gemm shapes the + # engine could encounter. + # - "relax" : Select gemm shapes to run based on some heuristics. The + # heuristic aims to have the same effect as running all possible gemm + # shapes, but provides no guarantees. + "APHRODITE_DEEP_GEMM_WARMUP": env_with_choices( + "APHRODITE_DEEP_GEMM_WARMUP", + "relax", + [ + "skip", + "full", + "relax", + ], + ), + # Whether to use fused grouped_topk used for MoE expert selection. + "APHRODITE_USE_FUSED_MOE_GROUPED_TOPK": lambda: bool(int(os.getenv("APHRODITE_USE_FUSED_MOE_GROUPED_TOPK", "1"))), + # Skip cudagraph/DP padding tokens in the MoE path by forcing their expert + # ids to -1 so the dispatch and experts drop them. Requires a MoE kernel that + # treats topk_id == -1 as a skip sentinel + "APHRODITE_MOE_SKIP_PADDING": lambda: bool(int(os.getenv("APHRODITE_MOE_SKIP_PADDING", "1"))), + # Allow use of FlashInfer FP8 block-scale GEMM for linear layers. + # This uses TensorRT-LLM kernels and requires SM90+ (Hopper). + "APHRODITE_BLOCKSCALE_FP8_GEMM_FLASHINFER": lambda: bool( + int(os.getenv("APHRODITE_BLOCKSCALE_FP8_GEMM_FLASHINFER", "1")) + ), + # Allow use of FlashInfer MxInt4 MoE kernels for fused moe ops. + "APHRODITE_USE_FLASHINFER_MOE_INT4": lambda: bool(int(os.getenv("APHRODITE_USE_FLASHINFER_MOE_INT4", "0"))), + # Control the cache sized used by the xgrammar compiler. The default + # of 512 MB should be enough for roughly 1000 JSON schemas. + # It can be changed with this variable if needed for some reason. + "APHRODITE_XGRAMMAR_CACHE_MB": lambda: int(os.getenv("APHRODITE_XGRAMMAR_CACHE_MB", "512")), + # Maximum time in seconds allowed for regex compilation in structured + # output backends (xgrammar, outlines). Prevents ReDoS attacks where + # adversarial patterns cause exponential DFA state-space explosion. + # Set to 0 to disable the timeout (not recommended in production). + "APHRODITE_REGEX_COMPILATION_TIMEOUT_S": lambda: int(os.getenv("APHRODITE_REGEX_COMPILATION_TIMEOUT_S", "5")), + # Control the threshold for msgspec to use 'zero copy' for + # serialization/deserialization of tensors. Tensors below + # this limit will be encoded into the msgpack buffer, and + # tensors above will instead be sent via a separate message. + # While the sending side still actually copies the tensor + # in all cases, on the receiving side, tensors above this + # limit will actually be zero-copy decoded. The unit is bytes. + "APHRODITE_MSGPACK_ZERO_COPY_THRESHOLD": lambda: int(os.getenv("APHRODITE_MSGPACK_ZERO_COPY_THRESHOLD", "256")), + # If set, allow insecure serialization using pickle. + # This is useful for environments where it is deemed safe to use the + # insecure method and it is needed for some reason. + "APHRODITE_ALLOW_INSECURE_SERIALIZATION": lambda: bool( + int(os.getenv("APHRODITE_ALLOW_INSECURE_SERIALIZATION", "0")) + ), + # Temporary: skip adding random suffix to internal request IDs. May be + # needed for KV connectors that match request IDs across instances. + "APHRODITE_DISABLE_REQUEST_ID_RANDOMIZATION": lambda: bool( + int(os.getenv("APHRODITE_DISABLE_REQUEST_ID_RANDOMIZATION", "0")) + ), + # IP address used for NIXL handshake between remote agents. + "APHRODITE_NIXL_SIDE_CHANNEL_HOST": lambda: os.getenv("APHRODITE_NIXL_SIDE_CHANNEL_HOST", "localhost"), + # Port used for NIXL handshake between remote agents. + "APHRODITE_NIXL_SIDE_CHANNEL_PORT": lambda: int(os.getenv("APHRODITE_NIXL_SIDE_CHANNEL_PORT", "5600")), + # Address the P2P KV-offload control socket binds to. Defaults to + # ``localhost`` (loopback only); must be set to the node IP for + # cross-host P2P so remote peers can reach the socket. + "APHRODITE_P2P_SIDE_CHANNEL_HOST": lambda: os.getenv("APHRODITE_P2P_SIDE_CHANNEL_HOST", "localhost"), + # Port the P2P KV-offload control socket binds to. + "APHRODITE_P2P_SIDE_CHANNEL_PORT": lambda: int(os.getenv("APHRODITE_P2P_SIDE_CHANNEL_PORT", "5710")), + # IP address used for the EC connector's ZMQ side channel + # (producer ROUTER bind, consumer DEALER dial). + "APHRODITE_EC_SIDE_CHANNEL_HOST": lambda: os.getenv("APHRODITE_EC_SIDE_CHANNEL_HOST", "localhost"), + # Port for the EC connector's ZMQ side channel; advertised to peers + # via `ec_transfer_params.peer_port` on the producer's response. + "APHRODITE_EC_SIDE_CHANNEL_PORT": lambda: int(os.getenv("APHRODITE_EC_SIDE_CHANNEL_PORT", "5601")), + # Port used for Mooncake handshake between remote agents. + "APHRODITE_MOONCAKE_BOOTSTRAP_PORT": lambda: int(os.getenv("APHRODITE_MOONCAKE_BOOTSTRAP_PORT", "8998")), + # Log per-batch memory/disk tier breakdown on external GETs. + "APHRODITE_MOONCAKE_STORE_TIER_LOG": lambda: ( + os.getenv("APHRODITE_MOONCAKE_STORE_TIER_LOG", "False").lower() in ("true", "1") + ), + # Number of parallel KV-load receive threads per worker rank. Lets the + # per-request control overhead (Python prep + master key lookup) of one + # request overlap with the RDMA transfer of another, keeping the transfer + # engine's queue pairs busy. Helps when that overhead is significant or + # per-request batches are too small to saturate the link on their own. + "APHRODITE_MOONCAKE_LOAD_RECV_THREADS": lambda: int(os.getenv("APHRODITE_MOONCAKE_LOAD_RECV_THREADS", "1")), + # Fraction of the owner's DirectIO staging buffer to fill per GET batch. + "APHRODITE_MOONCAKE_DISK_STAGING_USABLE_RATIO": lambda: float( + os.getenv("APHRODITE_MOONCAKE_DISK_STAGING_USABLE_RATIO", "0.9") + ), + # Pin this rank to a specific owner segment ("host:port"). + "MOONCAKE_PREFERRED_SEGMENT": lambda: os.getenv("MOONCAKE_PREFERRED_SEGMENT"), + # Override the hostname the rank registers as a Mooncake requester. + "MOONCAKE_REQUESTER_LOCAL_HOSTNAME": lambda: os.getenv("MOONCAKE_REQUESTER_LOCAL_HOSTNAME"), + # Override the directory for the FlashInfer autotune config cache. + "APHRODITE_FLASHINFER_AUTOTUNE_CACHE_DIR": lambda: os.getenv("APHRODITE_FLASHINFER_AUTOTUNE_CACHE_DIR", None), + # Comma-separated FlashInfer op names to exclude from autotuning, using + # the heuristic fallback tactic instead. Unset: skip "fp4_gemm" when the + # CuTe-DSL NVFP4 linear kernel is selected. Empty: skip nothing. + "APHRODITE_FLASHINFER_AUTOTUNE_SKIP_OPS": lambda: ( + None + if "APHRODITE_FLASHINFER_AUTOTUNE_SKIP_OPS" not in os.environ + else [v.strip() for v in os.environ["APHRODITE_FLASHINFER_AUTOTUNE_SKIP_OPS"].split(",") if v.strip()] + ), + # Flashinfer fused allreduce backend. + "APHRODITE_FLASHINFER_ALLREDUCE_BACKEND": env_with_choices( + "APHRODITE_FLASHINFER_ALLREDUCE_BACKEND", + "auto", + ["auto", "trtllm", "mnnvl"], + ), + # Control the workspace buffer size for the FlashInfer backend. + "APHRODITE_FLASHINFER_WORKSPACE_BUFFER_SIZE": lambda: int( + os.getenv("APHRODITE_FLASHINFER_WORKSPACE_BUFFER_SIZE", str(394 * 1024 * 1024)) + ), + # Control the maximum number of tokens per expert supported by the + # NVFP4 MoE CUTLASS Kernel. This value is used to create a buffer for + # the blockscale tensor of activations NVFP4 Quantization. + # This is used to prevent the kernel from running out of memory. + "APHRODITE_MAX_TOKENS_PER_EXPERT_FP4_MOE": lambda: int( + os.getenv("APHRODITE_MAX_TOKENS_PER_EXPERT_FP4_MOE", "163840") + ), + # MoE routing strategy selector. + # See `RoutingSimulator.get_available_strategies()` # for available + # strategies. + # Custom routing strategies can be registered by + # RoutingSimulator.register_strategy() + # Note: custom strategies may not produce correct model outputs + "APHRODITE_MOE_ROUTING_SIMULATION_STRATEGY": lambda: os.environ.get( + "APHRODITE_MOE_ROUTING_SIMULATION_STRATEGY", "" + ).lower(), + # Regex timeout for use by the Aphrodite tool parsing plugins. + "APHRODITE_TOOL_PARSE_REGEX_TIMEOUT_SECONDS": lambda: int( + os.getenv("APHRODITE_TOOL_PARSE_REGEX_TIMEOUT_SECONDS", "1") + ), + # Enforce function parameter schemas in structural-tag based tool calling. + "APHRODITE_ENFORCE_STRICT_TOOL_CALLING": lambda: ( + os.getenv("APHRODITE_ENFORCE_STRICT_TOOL_CALLING", "True").lower() in ("true", "1") + ), + # Control the max chunk bytes (in MB) for the rpc message queue. + # Object larger than this threshold will be broadcast to worker + # processes via zmq. + "APHRODITE_MQ_MAX_CHUNK_BYTES_MB": lambda: int(os.getenv("APHRODITE_MQ_MAX_CHUNK_BYTES_MB", "16")), + # Timeout in seconds for execute_model RPC calls in multiprocessing + # executor (only applies when TP > 1). + "APHRODITE_EXECUTE_MODEL_TIMEOUT_SECONDS": lambda: int(os.getenv("APHRODITE_EXECUTE_MODEL_TIMEOUT_SECONDS", "300")), + # Timeout in seconds for engine and worker process shutdown + "APHRODITE_WORKER_SHUTDOWN_TIMEOUT_SECONDS": lambda: int( + os.getenv("APHRODITE_WORKER_SHUTDOWN_TIMEOUT_SECONDS", "5") + ), + # KV Cache layout used throughout aphrodite. + # Some common values are: + # - NHD + # - HND + # Where N=num_blocks, H=num_heads and D=head_size. The default value will + # leave the layout choice to the backend. Mind that backends may only + # implement and support a subset of all possible layouts. + "APHRODITE_KV_CACHE_LAYOUT": env_with_choices("APHRODITE_KV_CACHE_LAYOUT", None, ["NHD", "HND"]), + # SSM conv state layout used for Mamba models. + # - SD: (state_len, dim) — dim contiguous (default) + # - DS: (dim, state_len) — TP-sharded dim on dim1, + # consistent with SSM temporal state and HND KV cache layout. + "APHRODITE_SSM_CONV_STATE_LAYOUT": env_with_choices("APHRODITE_SSM_CONV_STATE_LAYOUT", None, ["SD", "DS"]), + # Enable checking whether the generated logits contain NaNs, + # indicating corrupted output. Useful for debugging low level bugs + # or bad hardware but it may add compute overhead. + "APHRODITE_COMPUTE_NANS_IN_LOGITS": lambda: bool(int(os.getenv("APHRODITE_COMPUTE_NANS_IN_LOGITS", "0"))), + # Timeout (in seconds) for MooncakeConnector in PD disaggregated setup. + "APHRODITE_MOONCAKE_ABORT_REQUEST_TIMEOUT": lambda: int( + os.getenv("APHRODITE_MOONCAKE_ABORT_REQUEST_TIMEOUT", "480") + ), + # If set, it means we pre-downloaded cubin files and flashinfer will + # read the cubin files directly. + "APHRODITE_HAS_FLASHINFER_CUBIN": lambda: bool(int(os.getenv("APHRODITE_HAS_FLASHINFER_CUBIN", "0"))), + # Controls garbage collection during CUDA graph capture. + # If set to 0 (default), enables GC freezing to speed up capture time. + # If set to 1, allows GC to run during capture. + "APHRODITE_ENABLE_CUDAGRAPH_GC": lambda: bool(int(os.getenv("APHRODITE_ENABLE_CUDAGRAPH_GC", "0"))), + # Used to force set up loopback IP + "APHRODITE_LOOPBACK_IP": lambda: os.getenv("APHRODITE_LOOPBACK_IP", ""), + # Used to set the process name prefix for Aphrodite processes. + # This is useful for debugging and monitoring purposes. + # The default value is "APHRODITE". + "APHRODITE_PROCESS_NAME_PREFIX": lambda: os.getenv("APHRODITE_PROCESS_NAME_PREFIX", "APHRODITE"), + # Allow chunked local attention with hybrid kv cache manager. + # Currently using the Hybrid KV cache manager with chunked local attention + # in the Llama4 models (the only models currently using chunked local attn) + # causes a latency regression. For this reason, we disable it by default. + # This flag is used to allow users to enable it if they want to (to save on + # kv-cache memory usage and enable longer contexts) + # TODO(lucas): Remove this flag once latency regression is resolved. + "APHRODITE_ALLOW_CHUNKED_LOCAL_ATTN_WITH_HYBRID_KV_CACHE": lambda: bool( + int(os.getenv("APHRODITE_ALLOW_CHUNKED_LOCAL_ATTN_WITH_HYBRID_KV_CACHE", "1")) + ), + # Enables support for the "store" option in the OpenAI Responses API. + # When set to 1, Aphrodite's OpenAI server will retain the input and output + # messages for those requests in memory. By default, this is disabled (0), + # and the "store" option is ignored. + # NOTE/WARNING: + # 1. Messages are kept in memory only (not persisted to disk) and will be + # lost when the Aphrodite server shuts down. + # 2. Enabling this option will cause a memory leak, as stored messages are + # never removed from memory until the server terminates. + "APHRODITE_ENABLE_RESPONSES_API_STORE": lambda: bool(int(os.getenv("APHRODITE_ENABLE_RESPONSES_API_STORE", "0"))), + # If set, use the fp8 mfma in rocm paged attention. + "APHRODITE_ROCM_FP8_MFMA_PAGE_ATTN": lambda: bool(int(os.getenv("APHRODITE_ROCM_FP8_MFMA_PAGE_ATTN", "0"))), + # Whether to use pytorch symmetric memory for allreduce + "APHRODITE_ALLREDUCE_USE_SYMM_MEM": lambda: bool(int(os.getenv("APHRODITE_ALLREDUCE_USE_SYMM_MEM", "1"))), + # Whether to use FlashInfer allreduce + "APHRODITE_ALLREDUCE_USE_FLASHINFER": lambda: bool(int(os.getenv("APHRODITE_ALLREDUCE_USE_FLASHINFER", "0"))), + # Experimental: use this to enable MCP tool calling for non harmony models + "APHRODITE_USE_EXPERIMENTAL_PARSER_CONTEXT": lambda: bool( + int(os.getenv("APHRODITE_USE_EXPERIMENTAL_PARSER_CONTEXT", "0")) + ), + # User override folder for tuned Triton-kernel configs. Shared by MoE, + # Mamba SSU, and LoRA. Filenames are distinct so one folder can hold all. + # Each component first checks this folder, then the configs shipped with + # Aphrodite (if any). If no JSON matches, it uses a hard-coded heuristic. + "APHRODITE_TUNED_CONFIG_FOLDER": lambda: os.getenv("APHRODITE_TUNED_CONFIG_FOLDER", None), + # Opt-in persistence of the startup plan. When enabled, each worker + # saves the memory-profiling result (the suggested --kv-cache-memory value + # and the free-memory baseline) under APHRODITE_CACHE_ROOT/startup_plan/, + # keyed by a hardware+config fingerprint, and later boots auto-apply it + # -- skipping memory profiling -- when the fingerprint matches and + # current free memory >= the recorded baseline. + # See aphrodite/v1/worker/startup_plan.py. + "APHRODITE_ENABLE_STARTUP_PLAN": lambda: bool(int(os.getenv("APHRODITE_ENABLE_STARTUP_PLAN", "0"))), + # Valid values are container,code_interpreter,web_search_preview + # ex APHRODITE_GPT_OSS_SYSTEM_TOOL_MCP_LABELS=container,code_interpreter + # If the server_label of your mcp tool is not in this list it will + # be completely ignored. + "APHRODITE_GPT_OSS_SYSTEM_TOOL_MCP_LABELS": env_set_with_choices( + "APHRODITE_GPT_OSS_SYSTEM_TOOL_MCP_LABELS", + default=[], + choices=["container", "code_interpreter", "web_search_preview"], + ), + # Allows harmony instructions to be injected on system messages + "APHRODITE_GPT_OSS_HARMONY_SYSTEM_INSTRUCTIONS": lambda: bool( + int(os.getenv("APHRODITE_GPT_OSS_HARMONY_SYSTEM_INSTRUCTIONS", "0")) + ), + # Pin the conversation start date injected into the Harmony system + # message. When unset the current date is used, which introduces + # non-determinism (different tokens -> different model behaviour at + # temperature=0). Set to an ISO date string, e.g. "2023-09-12", + # for reproducible inference or testing. + "APHRODITE_SYSTEM_START_DATE": lambda: os.getenv("APHRODITE_SYSTEM_START_DATE", None), + # Enable automatic retry when tool call JSON parsing fails + # If enabled, returns an error message to the model to retry + # If disabled (default), raises an exception and fails the request + "APHRODITE_TOOL_JSON_ERROR_AUTOMATIC_RETRY": lambda: bool( + int(os.getenv("APHRODITE_TOOL_JSON_ERROR_AUTOMATIC_RETRY", "0")) + ), + # Add optional custom scopes for profiling, disable to avoid overheads + "APHRODITE_CUSTOM_SCOPES_FOR_PROFILING": lambda: bool(int(os.getenv("APHRODITE_CUSTOM_SCOPES_FOR_PROFILING", "0"))), + # Add optional nvtx scopes for profiling, disable to avoid overheads + "APHRODITE_NVTX_SCOPES_FOR_PROFILING": lambda: bool(int(os.getenv("APHRODITE_NVTX_SCOPES_FOR_PROFILING", "0"))), + # Represent block hashes in KV cache events as 64-bit integers instead of + # raw bytes. Defaults to True for backward compatibility. + "APHRODITE_KV_EVENTS_USE_INT_BLOCK_HASHES": lambda: bool( + int(os.getenv("APHRODITE_KV_EVENTS_USE_INT_BLOCK_HASHES", "1")) + ), + # Name of the shared memory buffer used for object storage. + # Only effective when mm_config.mm_processor_cache_type == "shm". + # Automatically generates a unique UUID-based name per process tree + # if not explicitly set. + "APHRODITE_OBJECT_STORAGE_SHM_BUFFER_NAME": get_env_or_set_default( + "APHRODITE_OBJECT_STORAGE_SHM_BUFFER_NAME", + lambda: f"APHRODITE_OBJECT_STORAGE_SHM_BUFFER_{uuid.uuid4().hex}", + ), + # The size in MB of the buffers (NVL and RDMA) used by DeepEP + "APHRODITE_DEEPEP_BUFFER_SIZE_MB": lambda: int(os.getenv("APHRODITE_DEEPEP_BUFFER_SIZE_MB", "1024")), + # Force DeepEP to use intranode kernel for inter-node communication in + # high throughput mode. This is useful archive higher prefill throughput + # on system supports multi-node nvlink (e.g GB200). + "APHRODITE_DEEPEP_HIGH_THROUGHPUT_FORCE_INTRA_NODE": lambda: bool( + int(os.getenv("APHRODITE_DEEPEP_HIGH_THROUGHPUT_FORCE_INTRA_NODE", "0")) + ), + # Allow DeepEP to use MNNVL (multi-node nvlink) for internode_ll kernel, + # turn this for better latency on GB200 like system + "APHRODITE_DEEPEP_LOW_LATENCY_USE_MNNVL": lambda: bool( + int(os.getenv("APHRODITE_DEEPEP_LOW_LATENCY_USE_MNNVL", "0")) + ), + # DeepEP v2: enable two-tier NVLink+RDMA hybrid mode + "APHRODITE_DEEPEP_V2_ALLOW_HYBRID_MODE": lambda: bool(int(os.getenv("APHRODITE_DEEPEP_V2_ALLOW_HYBRID_MODE", "0"))), + # DeepEP v2: use fewer SMs at slight throughput cost + "APHRODITE_DEEPEP_V2_PREFER_OVERLAP": lambda: bool(int(os.getenv("APHRODITE_DEEPEP_V2_PREFER_OVERLAP", "0"))), + # DeepEP v2: trade precision for transfer size in combine + "APHRODITE_DEEPEP_V2_ALLOW_MULTIPLE_REDUCTION": lambda: bool( + int(os.getenv("APHRODITE_DEEPEP_V2_ALLOW_MULTIPLE_REDUCTION", "0")) + ), + # The number of SMs/CUs to allocate for communication kernels when + # running DBO; the rest will be allocated to compute. + # Default: 20 on CUDA (SMs), 64 on ROCm (CUs). + "APHRODITE_DBO_COMM_SMS": lambda: int( + os.getenv( + "APHRODITE_DBO_COMM_SMS", + "64" + if hasattr(__import__("torch").version, "hip") and __import__("torch").version.hip is not None + else "20", + ) + ), + # Enable max_autotune & coordinate_descent_tuning in inductor_config + # to compile static shapes passed from compile_sizes in compilation_config + # If set to 1, enable max_autotune; By default, this is enabled (1) + "APHRODITE_ENABLE_INDUCTOR_MAX_AUTOTUNE": lambda: bool( + int(os.getenv("APHRODITE_ENABLE_INDUCTOR_MAX_AUTOTUNE", "1")) + ), + # If set to 1, enable coordinate_descent_tuning; + # By default, this is enabled (1) + "APHRODITE_ENABLE_INDUCTOR_COORDINATE_DESCENT_TUNING": lambda: bool( + int(os.getenv("APHRODITE_ENABLE_INDUCTOR_COORDINATE_DESCENT_TUNING", "1")) + ), + # Flag to enable NCCL symmetric memory allocation and registration + "APHRODITE_USE_NCCL_SYMM_MEM": lambda: bool(int(os.getenv("APHRODITE_USE_NCCL_SYMM_MEM", "0"))), + # NCCL header path + "APHRODITE_NCCL_INCLUDE_PATH": lambda: os.environ.get("APHRODITE_NCCL_INCLUDE_PATH", None), + # GC debug config + # - APHRODITE_GC_DEBUG=0: disable GC debugger + # - APHRODITE_GC_DEBUG=1: enable GC debugger with gc.collect elpased times + # - APHRODITE_GC_DEBUG='{"top_objects":5}': enable GC debugger with + # top 5 collected objects + "APHRODITE_GC_DEBUG": lambda: os.getenv("APHRODITE_GC_DEBUG", ""), + # Debug workspace allocations. + # logging of workspace resize operations. + "APHRODITE_DEBUG_WORKSPACE": lambda: bool(int(os.getenv("APHRODITE_DEBUG_WORKSPACE", "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")) + ), + # Limits when we run shared_experts in a separate stream. + # We found out that for large batch sizes, the separate stream + # execution is not beneficial (most likely because of the input clone) + # TODO(alexm-redhat): Tune to be more dynamic based on GPU type + "APHRODITE_SHARED_EXPERTS_STREAM_TOKEN_THRESHOLD": lambda: int( + int(os.getenv("APHRODITE_SHARED_EXPERTS_STREAM_TOKEN_THRESHOLD", 256)) + ), + # Token-count cutoff for multi-stream overlap of the attention input + # GEMM with auxiliary GEMMs (e.g. fused_wqa_wkv overlapped with indexer + # weights / kv-score projections in DeepSeek-V4). At or below this many + # tokens the FP8 main GEMM has idle SMs to share with the bf16 aux GEMMs + # and overlap is a 5-45% win; above it the FP8 GEMM saturates the device + # and the cross-stream sync becomes pure overhead. Set to 0 to disable + # the multi-stream path entirely. See #PR 41526 for the empirical result + # for the default value of 1024 tokens. + "APHRODITE_MULTI_STREAM_GEMM_TOKEN_THRESHOLD": lambda: int( + os.getenv("APHRODITE_MULTI_STREAM_GEMM_TOKEN_THRESHOLD", "1024") + ), + # Format for saving torch.compile cache artifacts + # - "binary": saves as binary file + # Safe for multiple aphrodite serve processes accessing the same torch compile cache. + # - "unpacked": saves as directory structure (for inspection/debugging) + # NOT multiprocess safe - race conditions may occur with multiple processes. + # Allows viewing and setting breakpoints in Inductor's code output files. + "APHRODITE_COMPILE_CACHE_SAVE_FORMAT": env_with_choices( + "APHRODITE_COMPILE_CACHE_SAVE_FORMAT", "binary", ["binary", "unpacked"] + ), + # Flag to control the v2 model runner. If unset, use config defaults. + "APHRODITE_USE_V2_MODEL_RUNNER": lambda: maybe_convert_bool(os.getenv("APHRODITE_USE_V2_MODEL_RUNNER", None)), + # Log model inspection after loading. + # If enabled, logs a transformers-style hierarchical view of the model + # with quantization methods and attention backends. + "APHRODITE_LOG_MODEL_INSPECTION": lambda: bool(int(os.getenv("APHRODITE_LOG_MODEL_INSPECTION", "0"))), + # Debug logging for --enable-mfu-metrics + "APHRODITE_DEBUG_MFU_METRICS": lambda: bool(int(os.getenv("APHRODITE_DEBUG_MFU_METRICS", "0"))), + # Disable using pytorch's pin memory for CPU offloading. + "APHRODITE_WEIGHT_OFFLOADING_DISABLE_PIN_MEMORY": lambda: bool( + int(os.getenv("APHRODITE_WEIGHT_OFFLOADING_DISABLE_PIN_MEMORY", "0")) + ), + # Disable using UVA (Unified Virtual Addressing) for CPU offloading. + "APHRODITE_WEIGHT_OFFLOADING_DISABLE_UVA": lambda: bool( + int(os.getenv("APHRODITE_WEIGHT_OFFLOADING_DISABLE_UVA", "0")) + ), + # On WSL2 with a compatible kernel (>= 4.19.121), pinned memory is + # supported but disabled by default due to a small performance regression. + # Set to 1 when pinned memory or UVA is required (e.g. CPU offloading + # or v2 model runner). + "APHRODITE_WSL2_ENABLE_PIN_MEMORY": lambda: bool(int(os.getenv("APHRODITE_WSL2_ENABLE_PIN_MEMORY", "0"))), + # Disable logging of Aphrodite logo at server startup time. + "APHRODITE_DISABLE_LOG_LOGO": lambda: bool(int(os.getenv("APHRODITE_DISABLE_LOG_LOGO", "0"))), + # Kill switch for the sm89 (Ada, e.g. RTX 4090) DeepSeek sparse attention + # kernels. Skips the SM89_MLA_SPARSE backend and the sm89 indexer logits + # paths even if the kernels are compiled in. + "APHRODITE_DISABLE_SM89_DSA": lambda: bool(int(os.getenv("APHRODITE_DISABLE_SM89_DSA", "0"))), + # Disable PDL for LoRA, as enabling PDL with LoRA on SM100 causes + # Triton compilation to fail. + "APHRODITE_LORA_DISABLE_PDL": lambda: bool(int(os.getenv("APHRODITE_LORA_DISABLE_PDL", "0"))), + # Enable CUDA compatibility mode for datacenter GPUs with older + # driver versions than the CUDA toolkit major version of Aphrodite. + "APHRODITE_ENABLE_CUDA_COMPATIBILITY": lambda: ( + os.environ.get("APHRODITE_ENABLE_CUDA_COMPATIBILITY", "0").strip().lower() in ("1", "true") + ), + # Path to the CUDA compatibility libraries when CUDA compatibility is enabled. + "APHRODITE_CUDA_COMPATIBILITY_PATH": lambda: os.environ.get("APHRODITE_CUDA_COMPATIBILITY_PATH", None), + # Skip model name validation in OpenAI API requests. + # When set to 1, any model name will be accepted in the 'model' field + # of API requests. This is useful for proxy/gateway scenarios where + # the actual model is served but different names may be used in requests. + "APHRODITE_SKIP_MODEL_NAME_VALIDATION": lambda: ( + os.getenv("APHRODITE_SKIP_MODEL_NAME_VALIDATION", "0").strip().lower() in ("1", "true") + ), + # Controls route logging at server startup time. + # - "off": disable route logging + # - "compact": group routes by API family + # - "full": log every route individually + "APHRODITE_LOG_ROUTES": env_with_choices("APHRODITE_LOG_ROUTES", "compact", ["off", "compact", "full"]), + # Whether it is a scale up launch engine for elastic EP, + # Should only be set by EngineCoreClient. + "APHRODITE_ELASTIC_EP_SCALE_UP_LAUNCH": lambda: bool(int(os.getenv("APHRODITE_ELASTIC_EP_SCALE_UP_LAUNCH", "0"))), + # Whether to wait for all requests to drain before sending the + # scaling command in elastic EP. + "APHRODITE_ELASTIC_EP_DRAIN_REQUESTS": lambda: bool(int(os.getenv("APHRODITE_ELASTIC_EP_DRAIN_REQUESTS", "0"))), + # If set to 1, enable CUDA graph memory estimation during memory profiling. + # This profiles CUDA graph memory usage to provide more accurate KV cache + # memory allocation. Enabled by default as of v0.21.0 + "APHRODITE_MEMORY_PROFILER_ESTIMATE_CUDAGRAPHS": lambda: bool( + int(os.getenv("APHRODITE_MEMORY_PROFILER_ESTIMATE_CUDAGRAPHS", "1")) + ), + # NIXL EP environment variables + "APHRODITE_NIXL_EP_MAX_NUM_RANKS": lambda: int(os.getenv("APHRODITE_NIXL_EP_MAX_NUM_RANKS", "32")), + # Whether enable XPU graph on Intel GPU + "APHRODITE_XPU_ENABLE_XPU_GRAPH": lambda: bool(int(os.getenv("APHRODITE_XPU_ENABLE_XPU_GRAPH", "0"))), + # whether use xpu specific sample kernel + "APHRODITE_XPU_USE_SAMPLER_KERNEL": lambda: bool(int(os.getenv("APHRODITE_XPU_USE_SAMPLER_KERNEL", "1"))), + # Enable simple KV offload. + "APHRODITE_USE_SIMPLE_KV_OFFLOAD": lambda: bool(int(os.getenv("APHRODITE_USE_SIMPLE_KV_OFFLOAD", "0"))), + # Whether to enable dual cuda streams for LoRA computation + # (used by both BaseLinearLayerWithLoRA and FusedMoEWithLoRA to + # overlap the base layer compute with the LoRA fast path). + "APHRODITE_LORA_ENABLE_DUAL_STREAM": lambda: bool(int(os.getenv("APHRODITE_LORA_ENABLE_DUAL_STREAM", "0"))), + # If set to 1, use Python spinloop extension to poll in a more efficient + # way when using the mp backend. + "APHRODITE_USE_SPINLOOP_EXT": lambda: bool(int(os.getenv("APHRODITE_USE_SPINLOOP_EXT", "0"))), + # Comma-separated GPU_BDF=NIC_BDF pairs for RDMA NIC selection. + # Must be set together with APHRODITE_NIC_SELECTION_VARS. + "APHRODITE_GPU_NIC_PCIE_MAPPING": lambda: os.getenv("APHRODITE_GPU_NIC_PCIE_MAPPING", ""), + # Comma-separated list of env vars to set from the GPU-NIC mapping. + # Each entry is VAR_NAME or VAR_NAME: (suffix appended to + # RDMA device name). Must be set together with APHRODITE_GPU_NIC_PCIE_MAPPING. + "APHRODITE_NIC_SELECTION_VARS": lambda: os.getenv("APHRODITE_NIC_SELECTION_VARS", ""), +} + + +# --8<-- [end:env-vars-definition] + + +def __getattr__(name: str): + """ + Gets environment variables lazily. + + NOTE: After enable_envs_cache() invocation (which triggered after service + initialization), all environment variables will be cached. + """ + if name in environment_variables: + return environment_variables[name]() + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") + + +def _is_envs_cache_enabled() -> bool: + """Checked if __getattr__ is wrapped with functools.cache""" + global __getattr__ + return hasattr(__getattr__, "cache_clear") + + +def enable_envs_cache() -> None: + """ + Enables caching of environment variables. This is useful for performance + reasons, as it avoids the need to re-evaluate environment variables on + every call. + + NOTE: Currently, it's invoked after service initialization to reduce + runtime overhead. This also means that environment variables should NOT + be updated after the service is initialized. + """ + if _is_envs_cache_enabled(): + # Avoid wrapping functools.cache multiple times + return + # Tag __getattr__ with functools.cache + global __getattr__ + __getattr__ = functools.cache(__getattr__) + + # Cache all environment variables + for key in environment_variables: + __getattr__(key) + + +def disable_envs_cache() -> None: + """ + Resets the environment variables cache. It could be used to isolate environments + between unit tests. + """ + global __getattr__ + # If __getattr__ is wrapped by functions.cache, unwrap the caching layer. + if _is_envs_cache_enabled(): + assert hasattr(__getattr__, "__wrapped__") + __getattr__ = __getattr__.__wrapped__ + + +def __dir__(): + return list(environment_variables.keys()) + + +def is_set(name: str): + """Check if an environment variable is explicitly set.""" + if name in environment_variables: + return name in os.environ + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") + + +def validate_environ(hard_fail: bool) -> None: + for env in os.environ: + if env.startswith("APHRODITE_") and env not in environment_variables: + if hard_fail: + raise ValueError(f"Unknown Aphrodite environment variable detected: {env}") + else: + logger.warning("Unknown Aphrodite environment variable detected: %s", env) + + +def compile_factors() -> dict[str, object]: + """Return env vars used for torch.compile cache keys. + + Start with every known Aphrodite env var; drop entries in `ignored_factors`; + hash everything else. This keeps the cache key aligned across workers.""" + + ignored_factors: set[str] = { + "MAX_JOBS", + "APHRODITE_RPC_BASE_PATH", + "APHRODITE_USE_MODELSCOPE", + "APHRODITE_RINGBUFFER_WARNING_INTERVAL", + "APHRODITE_DEBUG_DUMP_PATH", + "APHRODITE_PORT", + "APHRODITE_CACHE_ROOT", + # Runtime memory-plan persistence; does not affect compiled graphs. + "APHRODITE_ENABLE_STARTUP_PLAN", + # Location-only derived paths: where a cache/config directory lives + # cannot affect compiled artifacts, and hashing them means relocating + # HOME or the XDG roots silently invalidates every compile cache + # (APHRODITE_CACHE_ROOT above and + # APHRODITE_FLASHINFER_AUTOTUNE_CACHE_DIR below are already ignored for + # the same reason). + "APHRODITE_XLA_CACHE_PATH", + "APHRODITE_CONFIG_ROOT", + "LD_LIBRARY_PATH", + "APHRODITE_SERVER_DEV_MODE", + "APHRODITE_DP_MASTER_IP", + "APHRODITE_DP_MASTER_PORT", + "APHRODITE_NIXL_SIDE_CHANNEL_HOST", + "APHRODITE_P2P_SIDE_CHANNEL_HOST", + "APHRODITE_P2P_SIDE_CHANNEL_PORT", + "APHRODITE_RANDOMIZE_DP_DUMMY_INPUTS", + "APHRODITE_MODEL_REDIRECT_PATH", + "APHRODITE_HOST_IP", + "APHRODITE_FORCE_AOT_LOAD", + "S3_ACCESS_KEY_ID", + "S3_SECRET_ACCESS_KEY", + "S3_ENDPOINT_URL", + "APHRODITE_USAGE_STATS_SERVER", + "APHRODITE_NO_USAGE_STATS", + "APHRODITE_DO_NOT_TRACK", + "APHRODITE_LOGGING_LEVEL", + "APHRODITE_LOGGING_PREFIX", + "APHRODITE_LOGGING_STREAM", + "APHRODITE_LOGGING_CONFIG_PATH", + "APHRODITE_LOGGING_COLOR", + "APHRODITE_LOGGING_VERBOSE", + "APHRODITE_LOG_STATS_INTERVAL", + "APHRODITE_REQUEST_LEVEL_METRICS", + "APHRODITE_LOG_ROUTES", + "APHRODITE_DEBUG_LOG_API_SERVER_RESPONSE", + "APHRODITE_TUNED_CONFIG_FOLDER", + "APHRODITE_FLASHINFER_AUTOTUNE_CACHE_DIR", + "APHRODITE_FLASHINFER_AUTOTUNE_SKIP_OPS", + "APHRODITE_ENGINE_ITERATION_TIMEOUT_S", + "APHRODITE_HTTP_TIMEOUT_KEEP_ALIVE", + "APHRODITE_EXECUTE_MODEL_TIMEOUT_SECONDS", + "APHRODITE_WORKER_SHUTDOWN_TIMEOUT_SECONDS", + "APHRODITE_KEEP_ALIVE_ON_ENGINE_DEATH", + "APHRODITE_IMAGE_FETCH_TIMEOUT", + "APHRODITE_VIDEO_FETCH_TIMEOUT", + "APHRODITE_AUDIO_FETCH_TIMEOUT", + "APHRODITE_MEDIA_CACHE", + "APHRODITE_MEDIA_CACHE_MAX_SIZE_MB", + "APHRODITE_MEDIA_CACHE_TTL_HOURS", + "APHRODITE_MEDIA_FETCH_MAX_RETRIES", + "APHRODITE_MEDIA_URL_ALLOW_REDIRECTS", + "APHRODITE_MEDIA_LOADING_THREAD_COUNT", + "APHRODITE_MAX_AUDIO_CLIP_FILESIZE_MB", + "APHRODITE_MAX_AUDIO_DECODE_DURATION_S", + "APHRODITE_MAX_AUDIO_PREPROCESS_WORKERS", + "APHRODITE_MAX_IMAGE_PIXELS", + "APHRODITE_VIDEO_LOADER_BACKEND", + "APHRODITE_MEDIA_CONNECTOR", + "APHRODITE_OBJECT_STORAGE_SHM_BUFFER_NAME", + "APHRODITE_ASSETS_CACHE", + "APHRODITE_ASSETS_CACHE_MODEL_CLEAN", + "APHRODITE_WORKER_MULTIPROC_METHOD", + "APHRODITE_ENABLE_V1_MULTIPROCESSING", + "APHRODITE_V1_OUTPUT_PROC_CHUNK_SIZE", + "APHRODITE_CPU_KVCACHE_SPACE", + "APHRODITE_CPU_MOE_PREPACK", + "APHRODITE_ZENTORCH_WEIGHT_PREPACK", + "APHRODITE_TEST_FORCE_LOAD_FORMAT", + "APHRODITE_ENABLE_CUDA_COMPATIBILITY", + "APHRODITE_CUDA_COMPATIBILITY_PATH", + "APHRODITE_SKIP_MODEL_NAME_VALIDATION", + "LOCAL_RANK", + "CUDA_VISIBLE_DEVICES", + "NO_COLOR", + } + + from aphrodite.config.utils import normalize_value + + factors: dict[str, object] = {} + for factor, getter in environment_variables.items(): + if factor in ignored_factors: + continue + + try: + raw = getter() + except Exception as exc: # pragma: no cover - defensive logging + logger.warning( + "Skipping environment variable %s while hashing compile factors: %s", + factor, + exc, + ) + continue + + factors[factor] = normalize_value(raw) + + ray_noset_env_vars = [ + # Refer to + # https://github.com/ray-project/ray/blob/c584b1ea97b00793d1def71eaf81537d70efba42/python/ray/_private/accelerators/nvidia_gpu.py#L11 + # https://github.com/ray-project/ray/blob/c584b1ea97b00793d1def71eaf81537d70efba42/python/ray/_private/accelerators/amd_gpu.py#L11 + # https://github.com/ray-project/ray/blob/b97d21dab233c2bd8ed7db749a82a1e594222b5c/python/ray/_private/accelerators/amd_gpu.py#L10 + # https://github.com/ray-project/ray/blob/c584b1ea97b00793d1def71eaf81537d70efba42/python/ray/_private/accelerators/npu.py#L12 + # https://github.com/ray-project/ray/blob/c584b1ea97b00793d1def71eaf81537d70efba42/python/ray/_private/accelerators/hpu.py#L12 + # https://github.com/ray-project/ray/blob/c584b1ea97b00793d1def71eaf81537d70efba42/python/ray/_private/accelerators/neuron.py#L14 + # https://github.com/ray-project/ray/blob/c584b1ea97b00793d1def71eaf81537d70efba42/python/ray/_private/accelerators/tpu.py#L38 + # https://github.com/ray-project/ray/blob/c584b1ea97b00793d1def71eaf81537d70efba42/python/ray/_private/accelerators/intel_gpu.py#L10 + # https://github.com/ray-project/ray/blob/c584b1ea97b00793d1def71eaf81537d70efba42/python/ray/_private/accelerators/rbln.py#L10 + "RAY_EXPERIMENTAL_NOSET_CUDA_VISIBLE_DEVICES", + "RAY_EXPERIMENTAL_NOSET_ROCR_VISIBLE_DEVICES", + "RAY_EXPERIMENTAL_NOSET_HIP_VISIBLE_DEVICES", + "RAY_EXPERIMENTAL_NOSET_ASCEND_RT_VISIBLE_DEVICES", + "RAY_EXPERIMENTAL_NOSET_HABANA_VISIBLE_MODULES", + "RAY_EXPERIMENTAL_NOSET_NEURON_RT_VISIBLE_CORES", + "RAY_EXPERIMENTAL_NOSET_TPU_VISIBLE_CHIPS", + "RAY_EXPERIMENTAL_NOSET_ONEAPI_DEVICE_SELECTOR", + "RAY_EXPERIMENTAL_NOSET_RBLN_RT_VISIBLE_DEVICES", + ] + + for var in ray_noset_env_vars: + factors[var] = normalize_value(os.getenv(var)) + + return factors diff --git a/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/activation.py.orig b/aphrodite/model_executor/layers/activation.py.orig new file mode 100644 index 0000000000..ba83f033b4 --- /dev/null +++ b/aphrodite/model_executor/layers/activation.py.orig @@ -0,0 +1,806 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Custom activation functions.""" + +import math + +import torch +import torch.nn as nn +import torch.nn.functional as F + +from aphrodite.distributed import ( + divide, + get_tensor_model_parallel_rank, + get_tensor_model_parallel_world_size, +) +from aphrodite.logger import init_logger +from aphrodite.model_executor.custom_op import CustomOp +from aphrodite.model_executor.utils import set_weight_attrs +from aphrodite.platforms import CpuArchEnum, current_platform +from aphrodite.triton_utils import tl, triton +from aphrodite.utils.collection_utils import LazyDict + +logger = init_logger(__name__) + + +@triton.jit +def _swiglustep_and_mul_kernel( + o_ptr, + o_stride, + x_ptr, + x_stride, + limit: tl.constexpr, + d: tl.constexpr, + BLOCK_SIZE: tl.constexpr, +) -> None: + i = tl.program_id(axis=0).to(tl.int64) + j = tl.program_id(axis=1) + o_row_ptr = o_ptr + o_stride * i + x_row_ptr = x_ptr + x_stride * i + offsets = j * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + mask = offsets < d + + gate = tl.load(x_row_ptr + offsets, mask=mask).to(tl.float32) + up = tl.load(x_row_ptr + offsets + d, mask=mask).to(tl.float32) + + gate_silu = tl.sigmoid(gate) * gate + gate_clamped = tl.minimum(gate_silu, limit) + up_clamped = tl.minimum(tl.maximum(up, -limit), limit) + + result = gate_clamped * up_clamped + result = result.to(x_ptr.dtype.element_ty) + tl.store(o_row_ptr + offsets, result, mask=mask) + + +def swiglustep_and_mul_triton(output: torch.Tensor, input: torch.Tensor, limit: float = 7.0): + b, n = input.shape + assert input.ndim == 2 + assert n % 2 == 0 + d = n // 2 + + def grid(meta): + return (b, triton.cdiv(d, meta["BLOCK_SIZE"])) + + _swiglustep_and_mul_kernel[grid]( + output, + output.stride(0), + input, + input.stride(0), + limit=limit, + d=d, + BLOCK_SIZE=1024, + ) + + +# --8<-- [start:fatrelu_and_mul] +@CustomOp.register("fatrelu_and_mul") +class FatreluAndMul(CustomOp): + """An activation function for FATReLU. + + The function computes x -> FATReLU(x[:d]) * x[d:] where + d = x.shape[-1] // 2. + This is used in openbmb/MiniCPM-S-1B-sft. + + Shapes: + x: (num_tokens, 2 * d) or (batch_size, seq_len, 2 * d) + return: (num_tokens, d) or (batch_size, seq_len, d) + """ + + # --8<-- [end:fatrelu_and_mul] + + def __init__(self, threshold: float = 0.0): + super().__init__() + self.threshold = threshold + if current_platform.is_cuda_alike(): + self.op = torch.ops._C.fatrelu_and_mul + elif current_platform.is_cpu(): + self._forward_method = self.forward_native + + def forward_native(self, x: torch.Tensor) -> torch.Tensor: + d = x.shape[-1] // 2 + x1 = x[..., :d] + x2 = x[..., d:] + x1 = F.threshold(x1, self.threshold, 0.0) + return x1 * x2 + + def forward_cuda(self, x: torch.Tensor) -> torch.Tensor: + d = x.shape[-1] // 2 + output_shape = x.shape[:-1] + (d,) + out = torch.empty(output_shape, dtype=x.dtype, device=x.device) + self.op(out, x, self.threshold) + return out + + +# --8<-- [start:silu_and_mul] +@CustomOp.register("silu_and_mul") +class SiluAndMul(CustomOp): + """An activation function for SwiGLU. + + The function computes x -> silu(x[:d]) * x[d:] where d = x.shape[-1] // 2. + + Shapes: + x: (num_tokens, 2 * d) or (batch_size, seq_len, 2 * d) + return: (num_tokens, d) or (batch_size, seq_len, d) + """ + + # --8<-- [end:silu_and_mul] + + def __init__(self, *, compile_native: bool = True, enforce_enable: bool = False): + super().__init__( + enforce_enable=enforce_enable, + compile_native=compile_native, + ) + if current_platform.is_cuda_alike() or current_platform.is_cpu() or current_platform.is_xpu(): + self.op = torch.ops._C.silu_and_mul + + @staticmethod + def forward_native(x: torch.Tensor) -> torch.Tensor: + """PyTorch-native implementation equivalent to forward().""" + d = x.shape[-1] // 2 + return F.silu(x[..., :d]) * x[..., d:] + + def forward_cuda(self, x: torch.Tensor) -> torch.Tensor: + d = x.shape[-1] // 2 + output_shape = x.shape[:-1] + (d,) + out = torch.empty(output_shape, dtype=x.dtype, device=x.device) + self.op(out, x) + return out + + def forward_cpu(self, x: torch.Tensor) -> torch.Tensor: + return self.forward_cuda(x) + + def forward_xpu(self, x: torch.Tensor) -> torch.Tensor: + return self.forward_cuda(x) + + +@CustomOp.register("silu_and_mul_with_clamp") +class SiluAndMulWithClamp(CustomOp): + """SwiGLU activation with input clamping (used by some MoE shared experts). + + Computes: + gate = clamp(x[..., :d], max=swiglu_limit) + up = clamp(x[..., d:], min=-swiglu_limit, max=swiglu_limit) + out = gate * sigmoid(alpha * gate) * (up + beta) + where d = x.shape[-1] // 2. The defaults alpha=1.0, beta=0.0 reduce this to + ``silu(gate) * up``; SwiGLU-OAI style models pass alpha (sigmoid scale) and + beta=1.0 (up bias). + + Shapes: + x: (num_tokens, 2 * d) or (batch_size, seq_len, 2 * d) + return: (num_tokens, d) or (batch_size, seq_len, d) + """ + + def __init__( + self, + swiglu_limit: float, + alpha: float = 1.0, + beta: float = 0.0, + *, + compile_native: bool = True, + ): + super().__init__(compile_native=compile_native) + self.swiglu_limit = float(swiglu_limit) + self.alpha = float(alpha) + self.beta = float(beta) + if current_platform.is_rocm() or current_platform.is_xpu(): + self._forward_method = self.forward_native + elif current_platform.is_cuda_alike(): + self.op = torch.ops._C.silu_and_mul_with_clamp + elif current_platform.is_cpu(): + self._forward_method = self.forward_native + + def forward_native(self, x: torch.Tensor) -> torch.Tensor: + d = x.shape[-1] // 2 + gate = torch.clamp(x[..., :d], max=self.swiglu_limit) + up = torch.clamp(x[..., d:], min=-self.swiglu_limit, max=self.swiglu_limit) + return gate * torch.sigmoid(self.alpha * gate) * (up + self.beta) + + def forward_cuda(self, x: torch.Tensor) -> torch.Tensor: + d = x.shape[-1] // 2 + output_shape = x.shape[:-1] + (d,) + out = torch.empty(output_shape, dtype=x.dtype, device=x.device) + self.op(out, x, self.swiglu_limit, self.alpha, self.beta) + return out + + def forward_xpu(self, x: torch.Tensor) -> torch.Tensor: + return self.forward_native(x) + + def extra_repr(self) -> str: + return f"swiglu_limit={self.swiglu_limit!r}, alpha={self.alpha!r}, beta={self.beta!r}" + + +# --8<-- [start:mul_and_silu] +@CustomOp.register("mul_and_silu") +class MulAndSilu(CustomOp): + """An activation function for SwiGLU. + + The function computes x -> x[:d] * silu(x[d:]) where d = x.shape[-1] // 2. + + Shapes: + x: (num_tokens, 2 * d) or (batch_size, seq_len, 2 * d) + return: (num_tokens, d) or (batch_size, seq_len, d) + """ + + # --8<-- [end:mul_and_silu] + + def __init__(self): + super().__init__() + if current_platform.is_cuda_alike() or current_platform.is_xpu(): + self.op = torch.ops._C.mul_and_silu + elif current_platform.is_cpu(): + self._forward_method = self.forward_native + + def forward_native(self, x: torch.Tensor) -> torch.Tensor: + """PyTorch-native implementation equivalent to forward().""" + d = x.shape[-1] // 2 + return x[..., :d] * F.silu(x[..., d:]) + + def forward_cuda(self, x: torch.Tensor) -> torch.Tensor: + d = x.shape[-1] // 2 + output_shape = x.shape[:-1] + (d,) + out = torch.empty(output_shape, dtype=x.dtype, device=x.device) + self.op(out, x) + return out + + def forward_xpu(self, x: torch.Tensor) -> torch.Tensor: + return self.forward_cuda(x) + + +# --8<-- [start:gelu_and_mul_sparse] +@CustomOp.register("gelu_and_mul_sparse") +class GeluAndMulSparse(CustomOp): + """An activation function for GeluAndMulSparse. + This activation function is used in Gemma3n. It computes: + up_proj = self.up_proj(x) + gate_proj = self.gate_proj(x) + gate_proj = self._gaussian_topk(gate_proj) # sparsity + activations = self.act_fn(gate_proj) # gelu + down_proj = self.down_proj(activations * up_proj) + Shapes: + x: (num_tokens, 2 * d) or (batch_size, seq_len, 2 * d) + return: (num_tokens, d) or (batch_size, seq_len, d) + """ + + # --8<-- [end:gelu_and_mul_sparse] + + def __init__(self, activation_sparsity: float, approximate: str = "none"): + super().__init__() + # Gelu. + self.approximate = approximate + if approximate not in ("none", "tanh"): + raise ValueError(f"Unknown approximate mode: {approximate}") + if current_platform.is_rocm() and approximate == "tanh": + # TODO:[ROCm] PyTorch native GELU with tanh is unstable with torch.compile + logger.warning_once( + "[ROCm] Pytorch's native GELU with tanh approximation is currently " + "unstable and produces garbage. Fallback to 'none' approximation." + ) + self.approximate = "none" + + # Sparsity. + if activation_sparsity == 0.0: + raise ValueError("activation_sparsity is 0.0. Please use GeluAndMul.") + target_sparsity_tensor = torch.tensor(activation_sparsity, dtype=torch.float32) + normal_dist = torch.distributions.normal.Normal(0, 1) + self.std_multiplier = normal_dist.icdf(target_sparsity_tensor) + + def _gaussian_topk(self, x: torch.Tensor) -> torch.Tensor: + """Get % sparse percentile of the Gaussian distribution.""" + # NOTE(rob): for TP>1, we could all-gather to get the means/std. + # But we do not do this because in expectation they are the same + # and in practice the eval scores are good without gathering. + mean = torch.mean(x, dim=-1, keepdim=True) + std = torch.std(x, dim=-1, keepdim=True, unbiased=False) + cutoff_x = mean + std * self.std_multiplier + return nn.functional.relu(x - cutoff_x) + + def forward_native(self, x: torch.Tensor) -> torch.Tensor: + """PyTorch-native implementation equivalent to forward().""" + d = x.shape[-1] // 2 + out = self._gaussian_topk(x[..., :d]) + out = F.gelu(out, approximate=self.approximate) + return out * x[..., d:] + + def forward_cuda(self, x: torch.Tensor) -> torch.Tensor: + return self.forward_native(x) + + +# --8<-- [start:gelu] +@CustomOp.register("gelu") +class GELU(CustomOp): + # --8<-- [end:gelu] + + def __init__(self): + super().__init__() + if ( + current_platform.is_cpu() + and current_platform.get_cpu_architecture() == CpuArchEnum.ARM + and hasattr(torch.ops._C, "activation_lut_bf16") + ): + self.op = torch.ops._C.activation_lut_bf16 + else: + self.op = None + + def forward_native(self, x: torch.Tensor) -> torch.Tensor: + return F.gelu(x, approximate="none") + + def forward_cpu(self, x: torch.Tensor) -> torch.Tensor: + if self.op and x.dtype == torch.bfloat16 and x.is_contiguous(): + out = torch.empty_like(x) + self.op(out, x, "gelu") + return out + return self.forward_native(x) + + def forward_cuda(self, x: torch.Tensor) -> torch.Tensor: + return self.forward_native(x) + + +# --8<-- [start:gelu_tanh] +@CustomOp.register("gelu_tanh") +class GELUTanh(CustomOp): + # --8<-- [end:gelu_tanh] + + def __init__(self): + super().__init__() + if ( + current_platform.is_cpu() + and current_platform.get_cpu_architecture() == CpuArchEnum.ARM + and hasattr(torch.ops._C, "gelu_tanh") + ): + self.op = torch.ops._C.gelu_tanh + else: + self.op = None + + def forward_native(self, x: torch.Tensor) -> torch.Tensor: + return F.gelu(x, approximate="tanh") + + def forward_cpu(self, x: torch.Tensor) -> torch.Tensor: + if self.op: + out = torch.empty_like(x) + self.op(out, x) + return out + return self.forward_native(x) + + def forward_cuda(self, x: torch.Tensor) -> torch.Tensor: + return self.forward_native(x) + + +# --8<-- [start:gelu_and_mul] +@CustomOp.register("gelu_and_mul") +class GeluAndMul(CustomOp): + """An activation function for GeGLU. + + The function computes x -> GELU(x[:d]) * x[d:] where d = x.shape[-1] // 2. + + Shapes: + x: (batch_size, seq_len, 2 * d) or (num_tokens, 2 * d) + return: (batch_size, seq_len, d) or (num_tokens, d) + """ + + # --8<-- [end:gelu_and_mul] + + def __init__(self, approximate: str = "none"): + super().__init__() + self.approximate = approximate + if approximate not in ("none", "tanh"): + raise ValueError(f"Unknown approximate mode: {approximate}") + if current_platform.is_cuda_alike() or current_platform.is_cpu() or current_platform.is_xpu(): + if approximate == "none": + self.op = torch.ops._C.gelu_and_mul + elif approximate == "tanh": + self.op = torch.ops._C.gelu_tanh_and_mul + if current_platform.is_rocm() and approximate == "tanh": + logger.warning_once( + "[ROCm] PyTorch's native GELU with tanh approximation is unstable " + "with torch.compile. For native implementation, fallback to 'none' " + "approximation. The custom kernel implementation is unaffected." + ) + + def forward_native(self, x: torch.Tensor) -> torch.Tensor: + """PyTorch-native implementation equivalent to forward().""" + # TODO: [ROCm] PyTorch's native GELU with tanh is unstable with torch.compile + approximate = self.approximate + if current_platform.is_rocm() and approximate == "tanh": + approximate = "none" + d = x.shape[-1] // 2 + return F.gelu(x[..., :d], approximate=approximate) * x[..., d:] + + def forward_cuda(self, x: torch.Tensor) -> torch.Tensor: + d = x.shape[-1] // 2 + output_shape = x.shape[:-1] + (d,) + out = torch.empty(output_shape, dtype=x.dtype, device=x.device) + self.op(out, x) + return out + + def forward_cpu(self, x: torch.Tensor) -> torch.Tensor: + if self.op: + return self.forward_cuda(x) + return self.native(x) + + def forward_xpu(self, x: torch.Tensor) -> torch.Tensor: + return self.forward_cuda(x) + + def extra_repr(self) -> str: + return f"approximate={repr(self.approximate)}" + + +# --8<-- [start:swigluoai_and_mul] +@CustomOp.register("swigluoai_and_mul") +class SwigluOAIAndMul(CustomOp): + # https://github.com/huggingface/transformers/blob/v4.55.0/src/transformers/models/gpt_oss/modeling_gpt_oss.py#L106-L110 + # --8<-- [end:swigluoai_and_mul] + + def __init__(self, alpha: float = 1.702, limit: float = 7.0): + super().__init__() + self.alpha = alpha + self.limit = limit + + def forward_native(self, x: torch.Tensor) -> torch.Tensor: + """PyTorch-native implementation equivalent to forward().""" + + gate, up = x[..., ::2], x[..., 1::2] + gate = gate.clamp(min=None, max=self.limit) + up = up.clamp(min=-self.limit, max=self.limit) + glu = gate * torch.sigmoid(gate * self.alpha) + gated_output = (up + 1) * glu + return gated_output + + def forward_cuda(self, x: torch.Tensor) -> torch.Tensor: + d = x.shape[-1] // 2 + output_shape = x.shape[:-1] + (d,) + out = torch.empty(output_shape, dtype=x.dtype, device=x.device) + torch.ops._C.swigluoai_and_mul(out, x, self.alpha, self.limit) + return out + + def extra_repr(self) -> str: + return f"alpha={repr(self.alpha)}, limit={repr(self.limit)}" + + +# --8<-- [start:swiglustep_and_mul] +@CustomOp.register("swiglustep_and_mul") +class SwigluStepAndMul(CustomOp): + """An activation function for SwiGLU with clamping. + + Computes x -> silu(x[:d]).clamp(max=limit) * x[d:].clamp(-limit, limit) + where d = x.shape[-1] // 2. + + Shapes: + x: (num_tokens, 2 * d) or (batch_size, seq_len, 2 * d) + return: (num_tokens, d) or (batch_size, seq_len, d) + """ + + def __init__(self, limit: float = 7.0): + super().__init__() + if limit is None: + raise ValueError("SwigluStepAndMul requires limit to be set.") + self.limit = limit + + def forward_native(self, x: torch.Tensor) -> torch.Tensor: + """PyTorch-native implementation equivalent to forward().""" + gate, up = x.chunk(2, dim=-1) + gate = F.silu(gate) + gate = gate.clamp(max=self.limit) + up = up.clamp(min=-self.limit, max=self.limit) + return gate * up + + def forward_cuda(self, x: torch.Tensor) -> torch.Tensor: + d = x.shape[-1] // 2 + output_shape = x.shape[:-1] + (d,) + out = torch.empty(output_shape, dtype=x.dtype, device=x.device) + swiglustep_and_mul_triton(out, x, self.limit) + return out + + def extra_repr(self) -> str: + return f"limit={repr(self.limit)}" + + +# --8<-- [start:gelu_new] +@CustomOp.register("gelu_new") +class NewGELU(CustomOp): + # --8<-- [end:gelu_new] + + def __init__(self): + super().__init__() + if current_platform.is_cuda_alike() or current_platform.is_cpu() or current_platform.is_xpu(): + self.op = torch.ops._C.gelu_new + + def forward_native(self, x: torch.Tensor) -> torch.Tensor: + """PyTorch-native implementation equivalent to forward().""" + c = math.sqrt(2.0 / math.pi) + return 0.5 * x * (1.0 + torch.tanh(c * (x + 0.044715 * torch.pow(x, 3.0)))) + + def forward_cuda(self, x: torch.Tensor) -> torch.Tensor: + out = torch.empty_like(x) + self.op(out, x) + return out + + def forward_xpu(self, x: torch.Tensor) -> torch.Tensor: + return self.forward_cuda(x) + + +# --8<-- [start:gelu_fast] +@CustomOp.register("gelu_fast") +class FastGELU(CustomOp): + # --8<-- [end:gelu_fast] + + def __init__(self): + super().__init__() + if current_platform.is_cuda_alike() or current_platform.is_cpu() or current_platform.is_xpu(): + self.op = torch.ops._C.gelu_fast + + def forward_native(self, x: torch.Tensor) -> torch.Tensor: + """PyTorch-native implementation equivalent to forward().""" + return 0.5 * x * (1.0 + torch.tanh(x * 0.7978845608 * (1.0 + 0.044715 * x * x))) + + def forward_cuda(self, x: torch.Tensor) -> torch.Tensor: + out = torch.empty_like(x) + self.op(out, x) + return out + + def forward_xpu(self, x: torch.Tensor) -> torch.Tensor: + return self.forward_cuda(x) + + +# --8<-- [start:quick_gelu] +@CustomOp.register("quick_gelu") +class QuickGELU(CustomOp): + # https://github.com/huggingface/transformers/blob/main/src/transformers/activations.py#L90 + # --8<-- [end:quick_gelu] + + def __init__(self): + super().__init__() + if current_platform.is_cuda_alike() or current_platform.is_cpu() or current_platform.is_xpu(): + self.op = torch.ops._C.gelu_quick + + def forward_native(self, x: torch.Tensor) -> torch.Tensor: + """PyTorch-native implementation equivalent to forward().""" + return x * torch.sigmoid(1.702 * x) + + def forward_cuda(self, x: torch.Tensor) -> torch.Tensor: + out = torch.empty_like(x) + self.op(out, x) + return out + + def forward_xpu(self, x: torch.Tensor) -> torch.Tensor: + return self.forward_cuda(x) + + +# --8<-- [start:relu2] +@CustomOp.register("relu2") +class ReLUSquaredActivation(CustomOp): + """ + Applies the relu^2 activation introduced in https://arxiv.org/abs/2109.08668v2 + """ + + # --8<-- [end:relu2] + + def __init__(self): + super().__init__() + if current_platform.is_cuda_alike(): + self.op = torch.ops._C.relu_squared + + def forward_native(self, x: torch.Tensor) -> torch.Tensor: + """PyTorch-native implementation equivalent to forward().""" + return torch.square(F.relu(x)) + + def forward_cuda(self, x: torch.Tensor) -> torch.Tensor: + out = torch.empty_like(x) + self.op(out, x) + return out + + +# --8<-- [start:xielu] +@CustomOp.register("xielu") +class XIELU(CustomOp): + """ + Applies the xIELU activation function introduced in https://arxiv.org/abs/2411.13010 + If the user has installed the nickjbrowning/XIELU, we import xIELU CUDA + Otherwise, we emit a single warning and use xIELU Python + """ + + # --8<-- [end:xielu] + + def __init__( + self, + alpha_p_init: float = 0.8, + alpha_n_init: float = 0.8, + beta: float = 0.5, + eps: float = -1e-6, + dtype: torch.dtype = torch.bfloat16, + with_vector_loads: bool = False, + ): + super().__init__() + self.alpha_p = nn.Parameter(torch.log(torch.exp(torch.tensor(alpha_p_init, dtype=dtype)) - 1).unsqueeze(0)) + self.alpha_n = nn.Parameter( + torch.log(torch.exp(torch.tensor(alpha_n_init - beta, dtype=dtype)) - 1).unsqueeze(0) + ) + self.register_buffer("beta", torch.tensor(beta, dtype=dtype)) + self.register_buffer("eps", torch.tensor(eps, dtype=dtype)) + self.with_vector_loads = with_vector_loads + # Temporary until xIELU CUDA fully implemented + self._beta_scalar = float(self.beta.detach().cpu().float().item()) + self._eps_scalar = float(self.eps.detach().cpu().float().item()) + + self._xielu_cuda_obj = None + try: + import xielu.ops # noqa: F401 + + self._xielu_cuda_obj = torch.classes.xielu.XIELU() + msg = "Using experimental xIELU CUDA." + try: + from torch._dynamo import allow_in_graph + + self._xielu_cuda_fn = allow_in_graph(self._xielu_cuda) + msg += " Enabled torch._dynamo for xIELU CUDA." + except Exception as err: + msg += f" Could not enable torch._dynamo for xIELU ({err}) - this may result in slower performance." + self._xielu_cuda_fn = self._xielu_cuda + logger.warning_once(msg) + except Exception as err: + logger.warning_once( + "CUDA-fused xIELU not available (%s) –" + " falling back to a Python version.\n" + "For CUDA xIELU (experimental), `pip install git+https://github.com/nickjbrowning/XIELU`", + str(err), + ) + + def _xielu_python(self, x: torch.Tensor) -> torch.Tensor: + alpha_p = nn.functional.softplus(self.alpha_p) + alpha_n = self.beta + nn.functional.softplus(self.alpha_n) + return torch.where( + x > 0, + alpha_p * x * x + self.beta * x, + (torch.expm1(torch.min(x, self.eps)) - x) * alpha_n + self.beta * x, + ) + + def _xielu_cuda(self, x: torch.Tensor) -> torch.Tensor: + """Firewall function to prevent torch.compile from seeing .item()""" + assert self._xielu_cuda_obj is not None, "XIELU CUDA object must not be None" + original_shape = x.shape + # CUDA kernel expects 3D tensors, reshape if needed + while x.dim() < 3: + x = x.unsqueeze(0) + if x.dim() > 3: + x = x.view(-1, 1, x.size(-1)) + if original_shape != x.shape: + logger.warning_once( + "Warning: xIELU input tensor expects 3 dimensions but got (shape: %s). Reshaping to (shape: %s).", + original_shape, + x.shape, + ) + result = self._xielu_cuda_obj.forward( + x, + self.alpha_p, + self.alpha_n, + # Temporary until xIELU CUDA fully implemented -> + # self.{beta,eps}.item() + self._beta_scalar, + self._eps_scalar, + self.with_vector_loads, + ) + return result.view(original_shape) + + def forward_native(self, input: torch.Tensor) -> torch.Tensor: + if self._xielu_cuda_obj is not None and input.is_cuda: + if not torch._dynamo.is_compiling(): + return self._xielu_cuda_fn(input) + else: + logger.warning_once("torch._dynamo is compiling, using Python version of xIELU.") + return self._xielu_python(input) + + def forward_cuda(self, input: torch.Tensor) -> torch.Tensor: + return self.forward_native(input) + + +class ScaledActivation(nn.Module): + """An activation function with post-scale parameters. + + This is used for some quantization methods like AWQ. + """ + + def __init__( + self, + act_module: nn.Module, + intermediate_size: int, + input_is_parallel: bool = True, + params_dtype: torch.dtype | None = None, + ): + super().__init__() + self.act = act_module + self.input_is_parallel = input_is_parallel + if input_is_parallel: + tp_size = get_tensor_model_parallel_world_size() + intermediate_size_per_partition = divide(intermediate_size, tp_size) + else: + intermediate_size_per_partition = intermediate_size + if params_dtype is None: + params_dtype = torch.get_default_dtype() + self.scales = nn.Parameter(torch.empty(intermediate_size_per_partition, dtype=params_dtype)) + set_weight_attrs(self.scales, {"weight_loader": self.weight_loader}) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.act(x) / self.scales + + def weight_loader(self, param: nn.Parameter, loaded_weight: torch.Tensor): + param_data = param.data + if self.input_is_parallel: + tp_rank = get_tensor_model_parallel_rank() + shard_size = param_data.shape[0] + start_idx = tp_rank * shard_size + loaded_weight = loaded_weight.narrow(0, start_idx, shard_size) + assert param_data.shape == loaded_weight.shape + param_data.copy_(loaded_weight) + + +_ACTIVATION_REGISTRY = LazyDict( + { + "gelu": lambda: GELU(), + "gelu_fast": lambda: FastGELU(), + "gelu_new": lambda: NewGELU(), + "gelu_pytorch_tanh": lambda: _get_gelu_pytorch_tanh(), + "relu": lambda: nn.ReLU(), + "relu2": lambda: ReLUSquaredActivation(), + "silu": lambda: nn.SiLU(), + "swish": lambda: nn.SiLU(), + "quick_gelu": lambda: QuickGELU(), + "tanh": lambda: nn.Tanh(), + "sigmoid": lambda: nn.Sigmoid(), + "xielu": lambda: XIELU(), + } +) + + +def _get_gelu_pytorch_tanh() -> nn.Module: + """Get PyTorch GELU with tanh approximation, with ROCm fallback + and fast GELU for ARM.""" + if current_platform.is_rocm(): + # TODO:[ROCm] PyTorch native GELU with tanh is unstable with torch.compile + logger.warning_once( + "[ROCm] PyTorch's native GELU with tanh approximation is unstable. " + "Falling back to GELU(approximate='none')." + ) + return nn.GELU(approximate="none") + if current_platform.is_cpu() and current_platform.get_cpu_architecture() == CpuArchEnum.ARM: + return GELUTanh() + return nn.GELU(approximate="tanh") + + +def get_act_fn(act_fn_name: str) -> nn.Module: + """Get an activation function by name.""" + act_fn_name = act_fn_name.lower() + + if act_fn_name.startswith("torch.nn.modules."): + activation_name = act_fn_name.split(".")[-1] + if activation_name == "identity": + return nn.Identity() + act_fn_name = activation_name + + if act_fn_name not in _ACTIVATION_REGISTRY: + raise ValueError(f"Activation function {act_fn_name!r} is not supported.") + + return _ACTIVATION_REGISTRY[act_fn_name] + + +_ACTIVATION_AND_MUL_REGISTRY: LazyDict[nn.Module] = LazyDict( + { + "gelu": lambda: GeluAndMul(), + "gelu_pytorch_tanh": lambda: GeluAndMul(approximate="tanh"), + "silu": lambda: SiluAndMul(), + "swish": lambda: SiluAndMul(), + "geglu": lambda: GeluAndMul(), + "swigluoai": lambda: SwigluOAIAndMul(), + } +) + + +def get_act_and_mul_fn(act_fn_name: str, *, compile_native: bool = True) -> nn.Module: + """Get an activation-and-mul (i.e. SiluAndMul) function by name.""" + act_fn_name = act_fn_name.lower() + + if not compile_native and act_fn_name in ("silu", "swish"): + return SiluAndMul(compile_native=False) + + try: + return _ACTIVATION_AND_MUL_REGISTRY[act_fn_name] + except KeyError: + raise ValueError(f"Activation function {act_fn_name!r} is not supported.") from None 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/activation.py.orig b/aphrodite/model_executor/layers/fused_moe/activation.py.orig new file mode 100644 index 0000000000..05afd0b857 --- /dev/null +++ b/aphrodite/model_executor/layers/fused_moe/activation.py.orig @@ -0,0 +1,188 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""MoE activation function enum and utilities.""" + +from enum import Enum + +import torch +import torch.nn.functional as F + + +class MoEActivation(Enum): + """Activation functions for MoE layers.""" + + # Gated activations (gate * activation(up)) expect input of shape [..., 2*d] + # and produce output of shape [..., d] + SILU = "silu" + GELU = "gelu" + GELU_TANH = "gelu_tanh" + RELU2 = "relu2" + # SWIGLUOAI expects gate/up *interleaved* in w13 ([gate0, up0, gate1, ...]), + # as in gpt-oss checkpoints. SWIGLUOAI_UNINTERLEAVE has identical math but + # expects the *packed* layout ([all gates; all ups]), as produced by a + # MergedColumnParallelLinear gate_up_proj (e.g. MiniMax-M3). + SWIGLUOAI = "swigluoai" + SWIGLUOAI_UNINTERLEAVE = "swigluoai_uninterleave" + SWIGLUSTEP = "swiglustep" + + # Non-gated activations (no mul with gate) expect input of shape [..., d] + # and produce output of shape [..., d]. + # NOTE: Non-gated activations require the "_no_mul" suffix to be present. + SILU_NO_MUL = "silu_no_mul" + GELU_NO_MUL = "gelu_no_mul" + GELU_TANH_NO_MUL = "gelu_tanh_no_mul" + RELU2_NO_MUL = "relu2_no_mul" + + @property + def is_gated(self) -> bool: + """Returns True if activation expects gate*activation(up) pattern. + + Gated activations expect input tensor with 2x the output size, + where the first half is the gate and second half is the up projection. + """ + return not self.value.endswith("_no_mul") + + @property + def custom_op_name(self) -> str: + """Maps to the CustomOp name of activations + in aphrodite/model_executor/layers/activation.py.""" + return _CUSTOM_OP_NAMES[self] + + def without_mul(self) -> "MoEActivation": + """Get the non-gated variant of this activation. + + For activations that have a _no_mul variant, returns that variant. + For activations without a _no_mul variant (or already _no_mul), + returns self. + """ + return _WITHOUT_MUL.get(self, self) + + @classmethod + def from_str(cls, s: str) -> "MoEActivation": + """Parse from string for backward compatibility.""" + s = _STR_ALIASES.get(s, s) + for member in cls: + if member.value == s: + return member + valid = [m.value for m in cls] + raise ValueError(f"Unknown MoE activation: {s!r}. Valid activations: {valid}") + + +# Module-level lookup tables used by MoEActivation functions. +_STR_ALIASES: dict[str, str] = { + "gelu_pytorch_tanh": "gelu_tanh", +} + +_CUSTOM_OP_NAMES: dict[MoEActivation, str] = { + MoEActivation.SILU: "silu_and_mul", + MoEActivation.GELU: "gelu_and_mul", + MoEActivation.GELU_TANH: "gelu_tanh_and_mul", + MoEActivation.SWIGLUOAI: "swigluoai_and_mul", + MoEActivation.SWIGLUOAI_UNINTERLEAVE: "silu_and_mul_with_clamp", + MoEActivation.SWIGLUSTEP: "swiglustep_and_mul", + MoEActivation.RELU2: "relu2", + MoEActivation.SILU_NO_MUL: "silu_and_mul", + MoEActivation.GELU_NO_MUL: "gelu_and_mul", + MoEActivation.GELU_TANH_NO_MUL: "gelu_tanh_and_mul", + MoEActivation.RELU2_NO_MUL: "relu2", +} + +_WITHOUT_MUL: dict[MoEActivation, MoEActivation] = { + MoEActivation.SILU: MoEActivation.SILU_NO_MUL, + MoEActivation.GELU: MoEActivation.GELU_NO_MUL, + MoEActivation.GELU_TANH: MoEActivation.GELU_TANH_NO_MUL, + MoEActivation.RELU2: MoEActivation.RELU2_NO_MUL, +} + + +def activation_without_mul(activation: str) -> str: + """Get the non-gated variant of an activation function. + + Args: + activation: The activation function name (e.g., "silu", "gelu") + + Returns: + The non-gated activation name (e.g., "silu_no_mul", "gelu_no_mul") + """ + return MoEActivation.from_str(activation).without_mul().value + + +def silu_and_mul_with_clamp( + output: torch.Tensor, + input: torch.Tensor, + clamp_limit: float, + topk_ids: torch.Tensor | None = None, + expert_map: torch.Tensor | None = None, +) -> None: + if topk_ids is not None and expert_map is not None: + from aphrodite.model_executor.layers.fused_moe.utils import swiglu_limit_func + + swiglu_limit_func(output, input, clamp_limit, topk_ids, expert_map) + else: + # Fused silu(clamp(gate)) * clamp(up); equivalent to swiglu_limit_func. + torch.ops._C.silu_and_mul_with_clamp(output, input, clamp_limit, 1.0, 0.0) + + +def apply_moe_activation( + activation: MoEActivation, + output: torch.Tensor, + input: torch.Tensor, + *, + clamp_limit: float | None = None, + alpha: float = 1.0, + beta: float = 0.0, + topk_ids: torch.Tensor | None = None, + expert_map: torch.Tensor | None = None, +) -> torch.Tensor: + """Apply MoE activation function. + + ``clamp_limit``/``alpha``/``beta`` (from the quant config) drive the clamped + SwiGLU kernels: ``SILU`` + ``clamp_limit`` and ``SWIGLUOAI_UNINTERLEAVE`` both + map to ``silu_and_mul_with_clamp``. Other activations ignore them. + """ + assert input.dim() == 2, "Input must be 2D" + assert output.dim() == 2, "Output must be 2D" + if activation.is_gated: + assert output.size(-1) * 2 == input.size(-1), ( + f"{activation.value} expects 2x ratio: {output.size(-1) * 2} vs {input.size(-1)}" + ) + else: + assert output.size(-1) == input.size(-1), ( + f"{activation.value} expects equal sizes: {output.size(-1)} vs {input.size(-1)}" + ) + + # Activations with gated multiplication (gate × activation(up)) + if activation == MoEActivation.SILU: + if clamp_limit is not None: + silu_and_mul_with_clamp(output, input, clamp_limit, topk_ids, expert_map) + else: + torch.ops._C.silu_and_mul(output, input) + elif activation == MoEActivation.GELU: + 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.SWIGLUOAI: + torch.ops._C.swigluoai_and_mul(output, input) + elif activation == MoEActivation.SWIGLUOAI_UNINTERLEAVE: + # SwiGLU-OAI on packed w13 (gate = first half, up = second half). + assert clamp_limit is not None, "SWIGLUOAI_UNINTERLEAVE requires clamp_limit" + torch.ops._C.silu_and_mul_with_clamp(output, input, clamp_limit, alpha, beta) + elif activation == MoEActivation.SWIGLUSTEP: + from aphrodite.model_executor.layers.activation import swiglustep_and_mul_triton + + swiglustep_and_mul_triton(output, input) + + # Activations without gated multiplication + elif activation == MoEActivation.SILU_NO_MUL: + output.copy_(F.silu(input)) + elif activation == MoEActivation.GELU_NO_MUL: + output.copy_(F.gelu(input)) + elif activation == MoEActivation.GELU_TANH_NO_MUL: + output.copy_(F.gelu(input, approximate="tanh")) + elif activation == MoEActivation.RELU2_NO_MUL: + F.relu(input, inplace=True) + torch.square(input, out=output) + else: + raise ValueError(f"Unsupported FusedMoe activation: {activation}") + + return output 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/config.py.orig b/aphrodite/model_executor/layers/fused_moe/config.py.orig new file mode 100644 index 0000000000..5e48b7e2f9 --- /dev/null +++ b/aphrodite/model_executor/layers/fused_moe/config.py.orig @@ -0,0 +1,1414 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from dataclasses import dataclass +from enum import IntEnum +from typing import Union + +import torch + +from aphrodite.config import ParallelConfig, SchedulerConfig +from aphrodite.config.kernel import MoEBackend +from aphrodite.distributed import get_dp_group, get_pcp_group, get_tensor_model_parallel_rank +from aphrodite.logger import init_logger +from aphrodite.model_executor.layers.fused_moe.activation import MoEActivation +from aphrodite.model_executor.layers.quantization.utils.ocp_mx_utils import ( + OCP_MX_DTYPES, + OCP_MX_Scheme, +) +from aphrodite.model_executor.layers.quantization.utils.quant_utils import GroupShape +from aphrodite.platforms import current_platform +from aphrodite.utils.import_utils import has_triton_kernels +from aphrodite.utils.math_utils import cdiv + +logger = init_logger(__name__) + +if has_triton_kernels(): + try: + from triton_kernels.matmul_ogs import PrecisionConfig + except (ImportError, AttributeError) as e: + logger.error( + "Failed to import Triton kernels. Please make sure your triton version is compatible. Error: %s", + e, + ) + + +def _get_config_dtype_str( + dtype: torch.dtype, + use_fp8_w8a8: bool = False, + use_fp8_w8a16: bool = False, + use_int8_w8a16: bool = False, + use_int4_w4a16: bool = False, + ocp_mx_scheme: str | None = None, +) -> str | None: + """ + Return a string used to construct the filename that contains the + tuning info for a particular quantization scheme. See + try_get_optimal_moe_config in fused_moe.py. + """ + if use_fp8_w8a8: + return "fp8_w8a8" + elif use_fp8_w8a16: + return "fp8_w8a16" + elif use_int8_w8a16: + return "int8_w8a16" + elif use_int4_w4a16: + return "int4_w4a16" + elif ocp_mx_scheme is not None: + # The output of this function is passed to `try_get_optimal_moe_config`, + # and as we only simulate OCP MX execution in fused_moe for now, + # we will NOT look for `*,dtype=w_mxfp4_a_mxfp4.json` for now. + return None + elif dtype == torch.float: + # avoiding cases where kernel fails when float32 MoE + # use fp16/bfloat16 configs + return "float32" + return None + + +def _quant_flags_to_group_shape( + quant_dtype: torch.dtype | str | None, + per_act_token_quant: bool, + per_out_ch_quant: bool, + block_shape: list[int] | None, +) -> tuple[GroupShape | None, GroupShape | None]: + """ + Convert MoE quantization flags into more generic GroupShapes. + """ + a_shape: GroupShape | None + w_shape: GroupShape | None + if block_shape is not None: + assert not per_act_token_quant + assert not per_out_ch_quant + # TODO(bnell): this is not quite right for activations since first + # dim should be 1. + a_shape = GroupShape(row=block_shape[0], col=block_shape[1]) + w_shape = GroupShape(row=block_shape[0], col=block_shape[1]) + else: + w_shape = None + a_shape = None if quant_dtype is None else GroupShape.PER_TENSOR + + if per_act_token_quant: + a_shape = GroupShape.PER_TOKEN + + if per_out_ch_quant: + w_shape = GroupShape.PER_TOKEN + + return a_shape, w_shape + + +# The type of method in top-K routing +# Please keep this in sync with the counterpart defined in https://github.com/flashinfer-ai/flashinfer/blob/main/include/flashinfer/trtllm/fused_moe/runner.h +class RoutingMethodType(IntEnum): + # Default: Softmax -> TopK + Default = (0,) + # Renormalize: TopK -> Softmax + Renormalize = (1,) + # DeepSeekV3: Sigmoid -> RoutingBiasAdd -> Top2 in group -> Top4 groups + # -> Top8 experts from the Top4 groups + DeepSeekV3 = (2,) + # Llama4: Top1 -> Sigmoid + Llama4 = (3,) + # RenormalizeNaive: Softmax -> TopK -> Renormalize + RenormalizeNaive = (4,) + # TopK: TopK (no softmax) + TopK = (5,) + # SigmoidRenorm: Sigmoid -> TopK -> Renormalize (divide by sum of top-K) + SigmoidRenorm = (6,) + # MiniMax2: Sigmoid + Bias -> TopK -> ScaledSumNormalize + # (routeScale=1.0, epsilon=1e-20) + MiniMax2 = (7,) + # Sigmoid: Sigmoid -> TopK (no renormalization) + Sigmoid = (8,) + # Unspecified + Unspecified = (9,) + # other routing types (not passed to FlashInfer kernels) + # Deepseek V4 -> sqrtsoftplus + Bias + Normalize + DeepseekV4 = (100,) + Custom = (101,) + Simulated = (102,) + + +def get_routing_method_type( + scoring_func: str, + top_k: int, + renormalize: bool, + num_expert_group: int | None, + has_e_score_bias: bool, + routed_scaling_factor: float | None = 1.0, +) -> RoutingMethodType: + if scoring_func == "sqrtsoftplus": + # DeepSeek V4 uses sqrtsoftplus routing with optional routing bias + # and top-k renormalization. + if renormalize: + return RoutingMethodType.DeepseekV4 + else: + return RoutingMethodType.Unspecified + + if has_e_score_bias: + if scoring_func == "sigmoid": + if not renormalize: + return RoutingMethodType.Unspecified + if (num_expert_group or 0) > 0: + return RoutingMethodType.DeepSeekV3 + if routed_scaling_factor in (None, 1.0): + return RoutingMethodType.MiniMax2 + return RoutingMethodType.Unspecified + else: + return RoutingMethodType.Unspecified + + if scoring_func == "sigmoid": + if renormalize: + return RoutingMethodType.SigmoidRenorm + return RoutingMethodType.Sigmoid + + if scoring_func == "softmax": + if renormalize: + return RoutingMethodType.RenormalizeNaive + else: + return RoutingMethodType.Default + + return RoutingMethodType.Unspecified + + +@dataclass +class FusedMoEQuantDesc: + """ + A quantization descriptor for fused MoE ops. This class can describe + either activations or weights. + """ + + # The quantized type of this parameters. None means unquantized or + # already quantized. + # TODO (bnell): use scalar_type instead of Union. + dtype: torch.dtype | str | None = None + + # A field that describes the quantization group shape, from quant_utils.py. + # * (-1, -1) for per-tensor quantization + # * (1, -1) for per-row quantization + # * (-1, 1) for per-column quantization + # * (128, 128) for 128x128 deepseek style block quantization + # * (1, 128) for deepseek style activation quantization + # (i.e. per-token-per-group) + shape: GroupShape | None = None + + # Quantization scales. + # TODO(bnell): maybe put PrecisionConfigs in subclass of QuantDesc? + scale: Union[torch.Tensor, "PrecisionConfig", None] = None + + # Quantization alphas or gscales, used for nvfp4 types. + # W4A8 FP8: used for per-channel scales + # TODO(bnell): put some of these in subclasses + alpha_or_gscale: torch.Tensor | None = None + + # Zero points for int4/int8 types + zp: torch.Tensor | None = None + + # Biases for GPT triton MoE + bias: torch.Tensor | None = None + + +# TODO(bnell): have subclasses for specific moe methods? +# e.g. for specific arguments bias, precision, etc. +@dataclass +class FusedMoEQuantConfig: + """ + The FusedMoEQuantConfig contains all the quantization parameters for + a single FusedMoEMethodBase operation. It consists of four + FusedMoEQuantDescs, one for each activation and set of weights. + + Each FusedMoEMethodBase must implement a get_fused_moe_quant_config + method to construct a FusedMoEQuantConfig for use with that class. + + FusedMoEQuant configs are only used for modular kernels, fused_experts + (from fused_moe.py), cutlass_moe_fp[48], rocm_aiter_fused_experts and + triton_kernel_moe_forward. Other MoE methods can ignore the + FusedMoEQuantConfig (for now) and hardcode it to None. + + There are currently some restrictions on what can be expressed: + - Most MoE ops only support similar quantization strategies for + each parameter, e.g. both weights must have the same GroupShape + and both activations must share the same GroupShape. One exception to + this is the cutlass moe which allows per channel quantization on the + outputs. Note: this restrictions are not always rigorously checked. + - Not all fused MoE functions support all the parameters, e.g. zero points, + global scales, alphas and biases are not universally supported. + - Fully general GroupShapes are not allowed. Activations only support + per token, per tensor or K-blocked. + - Weights are not required to have a GroupShape since they have already + been quantized. + + Other notes: + - PrecisionConfigs are specific to GPT OSS Triton. + - As a follow up it would probably make sense to subclass FusedMoEQuantDesc + or FusedMoEQuantConfig for particular FusedMoEMethodBase subclasses + so that only the required quantization parameters are used/stored. + """ + + # TODO(bnell) make sure a1_scales/a2_scales don't interfere with chunking + _a1: FusedMoEQuantDesc + _a2: FusedMoEQuantDesc + _w1: FusedMoEQuantDesc + _w2: FusedMoEQuantDesc + is_scale_swizzled: bool = True + + # MXFP4-specific TRTLLM parameters for SwiGLU activation clamping. + # These correspond to gemm1_alpha, gemm1_beta, gemm1_clamp_limit + # in TrtLlmMxfp4ExpertsBase. + gemm1_alpha: float | None = None + gemm1_beta: float | None = None + gemm1_clamp_limit: float | None = None + + mx_alignment: int = 0 + + def __post_init__(self): + assert not self.per_act_token_quant or self.block_shape is None, "illegal quantization" + + # + # Convenience accessors for various properties. + # + + @property + def quant_dtype(self) -> torch.dtype | str | None: + return self._a1.dtype + + @property + def weight_quant_dtype(self) -> torch.dtype | str | None: + return self._w1.dtype + + @property + def is_quantized(self) -> bool: + return self.quant_dtype is not None + + @property + def is_per_act_token(self) -> bool: + return self._a1.shape == GroupShape.PER_TOKEN + + @property + def per_act_token_quant(self) -> bool: + return self._a1.shape == GroupShape.PER_TOKEN + + @property + def per_out_ch_quant(self) -> bool: + return self._w1.shape == GroupShape.PER_TOKEN + + @property + def is_per_tensor(self) -> bool: + return self._a1.shape == GroupShape.PER_TENSOR + + @property + def block_shape(self) -> list[int] | None: + if ( + self._a1.shape is not None + and self._a1.shape != GroupShape.PER_TENSOR + and self._a1.shape != GroupShape.PER_TOKEN + ): + return [self._a1.shape.row, self._a1.shape.col] + else: + return None + + @property + def is_block_quantized(self) -> bool: + return self.block_shape is not None + + @property + def a1_scale(self) -> torch.Tensor | None: + assert self._a1.scale is None or isinstance(self._a1.scale, torch.Tensor) + return self._a1.scale + + @property + def a1_gscale(self) -> torch.Tensor | None: + return self._a1.alpha_or_gscale + + @property + def a2_scale(self) -> torch.Tensor | None: + assert self._a2.scale is None or isinstance(self._a2.scale, torch.Tensor) + return self._a2.scale + + @property + def a2_gscale(self) -> torch.Tensor | None: + return self._a2.alpha_or_gscale + + @property + def w1_scale(self) -> torch.Tensor | None: + assert self._w1.scale is None or isinstance(self._w1.scale, torch.Tensor) + return self._w1.scale + + @property + def w1_zp(self) -> torch.Tensor | None: + return self._w1.zp + + @property + def w1_bias(self) -> torch.Tensor | None: + return self._w1.bias + + @property + def w1_precision(self) -> "PrecisionConfig | None": + assert self._w1.scale is None or isinstance(self._w1.scale, PrecisionConfig) + return self._w1.scale + + @property + def g1_alphas(self) -> torch.Tensor | None: + return self._w1.alpha_or_gscale + + @property + def w2_scale(self) -> torch.Tensor | None: + assert self._w2.scale is None or isinstance(self._w2.scale, torch.Tensor) + return self._w2.scale + + @property + def w2_zp(self) -> torch.Tensor | None: + return self._w2.zp + + @property + def w2_bias(self) -> torch.Tensor | None: + return self._w2.bias + + @property + def w2_precision(self) -> "PrecisionConfig | None": + assert self._w2.scale is None or isinstance(self._w2.scale, PrecisionConfig) + return self._w2.scale + + @property + def g2_alphas(self) -> torch.Tensor | None: + return self._w2.alpha_or_gscale + + @property + def use_fp8_w8a8(self) -> bool: + return self.quant_dtype == current_platform.fp8_dtype() + + @property + def use_int8_w8a8(self) -> bool: + return self.quant_dtype == torch.int8 + + @property + def use_int8_w8a16(self) -> bool: + return self._a1.dtype is None and self._w1.dtype == torch.int8 + + @property + def use_fp8_w8a16(self) -> bool: + return self._a1.dtype is None and self._w1.dtype == current_platform.fp8_dtype() + + @property + def use_int4_w4a16(self) -> bool: + return self._a1.dtype is None and self._w1.dtype == "int4" + + @property + def use_nvfp4_w4a16(self) -> bool: + return self._a1.dtype is None and self._w1.dtype == "nvfp4" + + @property + def ocp_mx_scheme(self) -> str | None: + if not hasattr(self, "_ocp_mx_scheme"): + if (self._a1.dtype is not None and not isinstance(self._a1.dtype, str)) or ( + self._w1.dtype is not None and not isinstance(self._w1.dtype, str) + ): + self._ocp_mx_scheme = None + else: + ocp_mx_scheme = OCP_MX_Scheme.from_quant_dtype(self._a1.dtype, self._w1.dtype) + + if ocp_mx_scheme is not None: + ocp_mx_scheme = ocp_mx_scheme.value + + self._ocp_mx_scheme = ocp_mx_scheme + + return self._ocp_mx_scheme + + @property + def use_mxfp4_w4a16(self) -> bool: + return self._a1.dtype is None and self._w1.dtype == "mxfp4" + + @property + def use_mxfp4_w4a4(self) -> bool: + return self._a1.dtype == "mxfp4" and self._w1.dtype == "mxfp4" + + @property + def use_nvfp4_w4a4(self) -> bool: + return self.quant_dtype == "nvfp4" + + @property + def use_mxfp4_w4a8(self) -> bool: + return self._a1.dtype == "fp8" and self._w1.dtype == "mxfp4" + + def config_name(self, dtype: torch.dtype) -> str | None: + """ + Return a string used to construct the filename that contains the + tuning info for a particular quantization scheme. See + try_get_optimal_moe_config in fused_moe.py. + """ + return _get_config_dtype_str( + use_fp8_w8a8=self.use_fp8_w8a8, + use_fp8_w8a16=self.use_fp8_w8a16, + use_int8_w8a16=self.use_int8_w8a16, + use_int4_w4a16=self.use_int4_w4a16, + ocp_mx_scheme=self.ocp_mx_scheme, + dtype=dtype, + ) + + def scale_shape( + self, + max_tokens: int, + hidden_dim: int, + ) -> tuple[int, int] | None: + """ + Construct the proper activation scale shape for this + config. + """ + if self.is_quantized: + if self.is_block_quantized: + assert self.block_shape is not None + _, block_k = self.block_shape + k_tiles = cdiv(hidden_dim, block_k) + return (max_tokens, k_tiles) + elif self.is_per_act_token: + return (max_tokens, 1) + else: + return (1, 1) + else: + return None + + def batched_scale_shape( + self, + num_experts: int, + max_tokens: int, + hidden_dim: int, + ) -> tuple[int, int, int] | None: + """ + Construct the proper activation batched scale shape for this + config, e.g. (num experts, *scale_shape). + """ + if self.is_quantized: + scale_shape = self.scale_shape(max_tokens, hidden_dim) + assert scale_shape is not None + return (num_experts, *scale_shape) + else: + return None + + @staticmethod + def make( + quant_dtype: torch.dtype | str | None = None, + per_act_token_quant: bool = False, + per_out_ch_quant: bool = False, + block_shape: list[int] | None = None, + w1_scale: Union[torch.Tensor, "PrecisionConfig", None] = None, + w2_scale: Union[torch.Tensor, "PrecisionConfig", None] = None, + a1_scale: torch.Tensor | None = None, + a2_scale: torch.Tensor | None = None, + g1_alphas: torch.Tensor | None = None, + g2_alphas: torch.Tensor | None = None, + a1_gscale: torch.Tensor | None = None, + a2_gscale: torch.Tensor | None = None, + w1_bias: torch.Tensor | None = None, + w2_bias: torch.Tensor | None = None, + w1_zp: torch.Tensor | None = None, + w2_zp: torch.Tensor | None = None, + weight_dtype: torch.dtype | str | None = None, + is_scale_swizzled: bool = True, + gemm1_alpha: float | None = None, + gemm1_beta: float | None = None, + gemm1_clamp_limit: float | None = None, + ) -> "FusedMoEQuantConfig": + """ + General builder function for a FusedMoEQuantConfig. + - quant_dtype: Optional quantization type. None if activations are + unquantized or quantized prior to calling. Note: "nvfp4", "mxfp4", + "mxfp6_e3m2", "mxfp6_e2m3" are the only valid string values + for quant_dtype. + - per_act_token_quant: Activations have per token quantization. + - per_out_ch_quant: Outputs have per channel quantization. (only + for cutlass). + - block_shape: Optional block size for block-wise quantization. + Incompatible with per_act_token and per_out_ch quant. + - w1_scale: Optional scale to be used for w1. + - w2_scale: Optional scale to be used for w2. + - a1_scale: Optional scale to be used for a1. + - a2_scale: Optional scale to be used for a2. + - g1_alphas: Optional global quantization scales for w1 (for nvfp4). + Optional per-channel scales for w1 (for W4A8 FP8). + Optional dq scale i.e. w_scale * a_scale (for W8A8 fp8). + - g2_alphas: Optional global quantization scales for w2 (for nvfp4). + Optional per-channel scales for w2 (for W4A8 FP8). + Optional dq scale i.e. w_scale * a_scale (for W8A8 fp8). + - a1_gscale: Optional global quantization scales for a1 (1.0 /a2_scale). + - a2_gscale: Optional global quantization scales for a2 (1.0 /a2_scale). + + - w1_bias: Optional biases for w1 (GPT OSS Triton). + - w2_bias: Optional biases for w1 (GPT OSS Triton). + - w1_zp: Optional w1 zero points for int4/int8 quantization. + - w2_zp: Optional w2 zero points for int4/int8 quantization. + - is_scale_swizzled: Whether the activation scale-factor layout is + swizzled. Pass through to the underlying quantization kernel for + dtypes that distinguish layouts (nvfp4, mxfp8). Defaults to True. + - gemm1_alpha: Optional MXFP4 TRTLLM SwiGLU alpha parameter. + - gemm1_beta: Optional MXFP4 TRTLLM SwiGLU beta parameter. + - gemm1_clamp_limit: Optional MXFP4 TRTLLM SwiGLU clamp limit. + """ + assert not isinstance(quant_dtype, str) or quant_dtype in { + "nvfp4", + "mxfp4", + "mxfp6_e3m2", + "mxfp6_e2m3", + "mxfp8", + } + assert not isinstance(weight_dtype, str) or weight_dtype in { + "nvfp4", + "mxfp4", + "mxfp6_e3m2", + "mxfp6_e2m3", + "int4", + "mxfp8", + } + + if weight_dtype is None: + weight_dtype = quant_dtype + + a_shape, w_shape = _quant_flags_to_group_shape(quant_dtype, per_act_token_quant, per_out_ch_quant, block_shape) + quant_config = FusedMoEQuantConfig( + _a1=FusedMoEQuantDesc(quant_dtype, a_shape, a1_scale, a1_gscale), + _a2=FusedMoEQuantDesc(quant_dtype, a_shape, a2_scale, a2_gscale), + _w1=FusedMoEQuantDesc(weight_dtype, w_shape, w1_scale, g1_alphas, w1_zp, w1_bias), + _w2=FusedMoEQuantDesc(weight_dtype, w_shape, w2_scale, g2_alphas, w2_zp, w2_bias), + is_scale_swizzled=is_scale_swizzled, + gemm1_alpha=gemm1_alpha, + gemm1_beta=gemm1_beta, + gemm1_clamp_limit=gemm1_clamp_limit, + ) + assert quant_config.per_act_token_quant == per_act_token_quant + assert quant_config.per_out_ch_quant == per_out_ch_quant + assert quant_config.block_shape == block_shape + return quant_config + + +def fp8_w8a8_moe_quant_config( + w1_scale: torch.Tensor, + w2_scale: torch.Tensor, + a1_scale: torch.Tensor | None = None, + a2_scale: torch.Tensor | None = None, + w1_bias: torch.Tensor | None = None, + w2_bias: torch.Tensor | None = None, + per_act_token_quant: bool = False, + per_out_ch_quant: bool = False, + block_shape: list[int] | None = None, + a1_gscale: torch.Tensor | None = None, + a2_gscale: torch.Tensor | None = None, + g1_alphas: torch.Tensor | None = None, + g2_alphas: torch.Tensor | None = None, + gemm1_alpha: float | None = None, + gemm1_beta: float | None = None, + gemm1_clamp_limit: float | None = None, +) -> FusedMoEQuantConfig: + """ + Construct a quant config for fp8 activations and fp8 weights. + """ + return FusedMoEQuantConfig.make( + current_platform.fp8_dtype(), + w1_scale=w1_scale, + g1_alphas=g1_alphas, + w2_scale=w2_scale, + g2_alphas=g2_alphas, + w1_bias=w1_bias, + w2_bias=w2_bias, + a1_scale=a1_scale, + a1_gscale=a1_gscale, + a2_scale=a2_scale, + a2_gscale=a2_gscale, + per_act_token_quant=per_act_token_quant, + per_out_ch_quant=per_out_ch_quant, + block_shape=block_shape, + gemm1_alpha=gemm1_alpha, + gemm1_beta=gemm1_beta, + gemm1_clamp_limit=gemm1_clamp_limit, + ) + + +def int8_w8a8_moe_quant_config( + w1_scale: torch.Tensor, + w2_scale: torch.Tensor, + a1_scale: torch.Tensor | None, + a2_scale: torch.Tensor | None, + w1_bias: torch.Tensor | None = None, + w2_bias: torch.Tensor | None = None, + per_act_token_quant: bool = False, +) -> FusedMoEQuantConfig: + """ + Construct a quant config for int8 activations and int8 weights. + """ + return FusedMoEQuantConfig.make( + torch.int8, + w1_scale=w1_scale, + w2_scale=w2_scale, + a1_scale=a1_scale, + a2_scale=a2_scale, + w1_bias=w1_bias, + w2_bias=w2_bias, + per_act_token_quant=per_act_token_quant, + per_out_ch_quant=False, + block_shape=None, + ) + + +def gptq_marlin_moe_quant_config( + w1_scale: torch.Tensor, + w2_scale: torch.Tensor, + weight_bits: int, + group_size: int, + w1_zp: torch.Tensor | None = None, + w2_zp: torch.Tensor | None = None, + w1_bias: torch.Tensor | None = None, + w2_bias: torch.Tensor | None = None, +): + """ + Construct a quant config for gptq marlin quantization. + """ + from aphrodite.model_executor.layers.quantization.utils.quant_utils import GroupShape + + w_shape = None if group_size == -1 else GroupShape(row=1, col=group_size) + + # Activations are NOT quantized for GPTQ (fp16/bf16) + a_shape = w_shape # Same as weight shape for alignment + + # Determine weight dtype + if weight_bits == 4: + weight_dtype = "int4" + elif weight_bits == 8: + weight_dtype = torch.int8 + else: + raise ValueError(f"Unsupported weight_bits: {weight_bits}") + + return FusedMoEQuantConfig( + _a1=FusedMoEQuantDesc(dtype=None, shape=a_shape), + _a2=FusedMoEQuantDesc(dtype=None, shape=a_shape), + _w1=FusedMoEQuantDesc(weight_dtype, w_shape, w1_scale, None, w1_zp, w1_bias), + _w2=FusedMoEQuantDesc(weight_dtype, w_shape, w2_scale, None, w2_zp, w2_bias), + ) + + +def mxfp4_w4a16_moe_quant_config( + w1_scale: Union[torch.Tensor, "PrecisionConfig"], + w2_scale: Union[torch.Tensor, "PrecisionConfig"], + w1_bias: torch.Tensor | None = None, + w2_bias: torch.Tensor | None = None, + gemm1_alpha: float | None = None, + gemm1_beta: float | None = None, + gemm1_clamp_limit: float | None = None, +) -> FusedMoEQuantConfig: + """ + Construct a quant config for unquantized activations and mxfp4 weights. + """ + return FusedMoEQuantConfig( + _a1=FusedMoEQuantDesc(), + _a2=FusedMoEQuantDesc(), + _w1=FusedMoEQuantDesc("mxfp4", None, w1_scale, None, None, w1_bias), + _w2=FusedMoEQuantDesc("mxfp4", None, w2_scale, None, None, w2_bias), + gemm1_alpha=gemm1_alpha, + gemm1_beta=gemm1_beta, + gemm1_clamp_limit=gemm1_clamp_limit, + ) + + +def mxfp4_mxfp8_moe_quant_config( + w1_scale: Union[torch.Tensor, "PrecisionConfig"], + w2_scale: Union[torch.Tensor, "PrecisionConfig"], + a1_scale: torch.Tensor | None = None, + a2_scale: torch.Tensor | None = None, + w1_bias: torch.Tensor | None = None, + w2_bias: torch.Tensor | None = None, + block_shape: list[int] | None = None, + gemm1_alpha: float | None = None, + gemm1_beta: float | None = None, + gemm1_clamp_limit: float | None = None, + mx_alignment: int = 0, + is_scale_swizzled: bool = True, +) -> FusedMoEQuantConfig: + """ + Construct a quant config for mxfp4 activations and mxfp4 weights. + """ + return FusedMoEQuantConfig( + _a1=FusedMoEQuantDesc("mxfp8"), + _a2=FusedMoEQuantDesc("mxfp8"), + _w1=FusedMoEQuantDesc("mxfp4", None, w1_scale, None, None, w1_bias), + _w2=FusedMoEQuantDesc("mxfp4", None, w2_scale, None, None, w2_bias), + gemm1_alpha=gemm1_alpha, + gemm1_beta=gemm1_beta, + gemm1_clamp_limit=gemm1_clamp_limit, + mx_alignment=mx_alignment, + is_scale_swizzled=is_scale_swizzled, + ) + + +def mxfp4_w4a8_moe_quant_config( + w1_scale: Union[torch.Tensor, "PrecisionConfig"], + w2_scale: Union[torch.Tensor, "PrecisionConfig"], + a1_scale: torch.Tensor | None = None, + a2_scale: torch.Tensor | None = None, + w1_bias: torch.Tensor | None = None, + w2_bias: torch.Tensor | None = None, + block_shape: list[int] | None = None, + gemm1_clamp_limit: float | None = None, +) -> FusedMoEQuantConfig: + """ + Construct a quant config for fp8 activations and mxfp4 weights. + """ + return FusedMoEQuantConfig( + _a1=FusedMoEQuantDesc("fp8", None, a1_scale, None, None, None), + _a2=FusedMoEQuantDesc("fp8", None, a2_scale, None, None, None), + _w1=FusedMoEQuantDesc("mxfp4", None, w1_scale, None, None, w1_bias), + _w2=FusedMoEQuantDesc("mxfp4", None, w2_scale, None, None, w2_bias), + gemm1_clamp_limit=gemm1_clamp_limit, + ) + + +def ocp_mx_moe_quant_config( + quant_dtype: str, + w1_scale: Union[torch.Tensor, "PrecisionConfig"], + w2_scale: Union[torch.Tensor, "PrecisionConfig"], + weight_dtype: str | None = None, + a1_scale: torch.Tensor | None = None, + a2_scale: torch.Tensor | None = None, + w1_bias: torch.Tensor | None = None, + w2_bias: torch.Tensor | None = None, + block_shape: list[int] | None = None, + gemm1_alpha: float | None = None, + gemm1_beta: float | None = None, + gemm1_clamp_limit: float | None = None, +) -> FusedMoEQuantConfig: + """ + Construct a quant config for mxfp4 activations and mxfp4 weights. + """ + assert quant_dtype in OCP_MX_DTYPES + return FusedMoEQuantConfig.make( + quant_dtype=quant_dtype, + weight_dtype=weight_dtype, + w1_scale=w1_scale, + w2_scale=w2_scale, + a1_scale=a1_scale, + a2_scale=a2_scale, + w1_bias=w1_bias, + w2_bias=w2_bias, + per_act_token_quant=False, + per_out_ch_quant=False, + block_shape=block_shape, + gemm1_alpha=gemm1_alpha, + gemm1_beta=gemm1_beta, + gemm1_clamp_limit=gemm1_clamp_limit, + ) + + +def nvfp4_moe_quant_config( + g1_alphas: torch.Tensor, + g2_alphas: torch.Tensor, + a1_gscale: torch.Tensor, + a2_gscale: torch.Tensor, + w1_scale: torch.Tensor, + w2_scale: torch.Tensor, + w1_bias: torch.Tensor | None = None, + w2_bias: torch.Tensor | None = None, + is_scale_swizzled: bool = True, + gemm1_clamp_limit: float | None = None, +) -> FusedMoEQuantConfig: + """ + Construct a quant config for mxfp4 activations and nvp4 weights. + """ + return FusedMoEQuantConfig.make( + "nvfp4", + w1_scale=w1_scale, + w2_scale=w2_scale, + w1_bias=w1_bias, + w2_bias=w2_bias, + a1_gscale=a1_gscale, + a2_gscale=a2_gscale, + g1_alphas=g1_alphas, + g2_alphas=g2_alphas, + per_act_token_quant=False, + per_out_ch_quant=False, + block_shape=None, + is_scale_swizzled=is_scale_swizzled, + gemm1_clamp_limit=gemm1_clamp_limit, + ) + + +def mxfp4_moe_quant_config( + w1_scale: torch.Tensor, + w2_scale: torch.Tensor, +) -> FusedMoEQuantConfig: + """ + Construct a quant config for MXFP4 x MXFP4 MoE. + MXFP4 uses block scaling only (E8M0 scales, 32-element groups), with no + separate alphas / global activation scales in this config. + """ + return FusedMoEQuantConfig.make( + "mxfp4", + w1_scale=w1_scale, + w2_scale=w2_scale, + per_act_token_quant=False, + per_out_ch_quant=False, + block_shape=None, + ) + + +def nvfp4_w4a16_moe_quant_config( + g1_alphas: torch.Tensor, + g2_alphas: torch.Tensor, + w1_scale: torch.Tensor, + w2_scale: torch.Tensor, + gemm1_clamp_limit: float | None = None, +) -> FusedMoEQuantConfig: + """ + Construct a quant config for 16-but activations and nvp4 weights. + """ + return FusedMoEQuantConfig.make( + quant_dtype=None, + w1_scale=w1_scale, + w2_scale=w2_scale, + g1_alphas=g1_alphas, + g2_alphas=g2_alphas, + weight_dtype="nvfp4", + gemm1_clamp_limit=gemm1_clamp_limit, + ) + + +def int4_w4a16_moe_quant_config( + w1_scale: torch.Tensor, + w2_scale: torch.Tensor, + w1_zp: torch.Tensor | None = None, + w2_zp: torch.Tensor | None = None, + w1_bias: torch.Tensor | None = None, + w2_bias: torch.Tensor | None = None, + block_shape: list[int] | None = None, + a1_gscale: torch.Tensor | None = None, + a2_gscale: torch.Tensor | None = None, + gemm1_clamp_limit: float | None = None, + gemm1_alpha: float | None = None, + gemm1_beta: float | None = None, +) -> FusedMoEQuantConfig: + """ + Construct a quant config for 16-bit float activations and int4 weights. + """ + group_shape = GroupShape(*block_shape) if block_shape is not None else None + return FusedMoEQuantConfig( + _a1=FusedMoEQuantDesc(shape=group_shape, alpha_or_gscale=a1_gscale), + _a2=FusedMoEQuantDesc(shape=group_shape, alpha_or_gscale=a2_gscale), + _w1=FusedMoEQuantDesc("int4", group_shape, w1_scale, None, w1_zp, w1_bias), + _w2=FusedMoEQuantDesc("int4", group_shape, w2_scale, None, w2_zp, w2_bias), + gemm1_clamp_limit=gemm1_clamp_limit, + gemm1_alpha=gemm1_alpha, + gemm1_beta=gemm1_beta, + ) + + +def fp8_w8a16_moe_quant_config( + w1_scale: torch.Tensor, + w2_scale: torch.Tensor, + w1_bias: torch.Tensor | None = None, + w2_bias: torch.Tensor | None = None, + block_shape: list[int] | None = None, + gemm1_alpha: float | None = None, + gemm1_beta: float | None = None, + gemm1_clamp_limit: float | None = None, +) -> FusedMoEQuantConfig: + """ + Construct a quant config for 16-bit float activations and fp8 weights. + """ + group_shape = GroupShape(*block_shape) if block_shape is not None else None + fp8_dtype = current_platform.fp8_dtype() + return FusedMoEQuantConfig( + _a1=FusedMoEQuantDesc(), + _a2=FusedMoEQuantDesc(), + _w1=FusedMoEQuantDesc( + fp8_dtype, + group_shape, + w1_scale, + None, + None, + w1_bias, + ), + _w2=FusedMoEQuantDesc( + fp8_dtype, + group_shape, + w2_scale, + None, + None, + w2_bias, + ), + gemm1_alpha=gemm1_alpha, + gemm1_beta=gemm1_beta, + gemm1_clamp_limit=gemm1_clamp_limit, + ) + + +def int8_w8a16_moe_quant_config( + w1_scale: torch.Tensor, + w2_scale: torch.Tensor, + w1_zp: torch.Tensor | None = None, + w2_zp: torch.Tensor | None = None, + w1_bias: torch.Tensor | None = None, + w2_bias: torch.Tensor | None = None, + block_shape: list[int] | None = None, + a1_gscale: torch.Tensor | None = None, + a2_gscale: torch.Tensor | None = None, + gemm1_clamp_limit: float | None = None, + gemm1_alpha: float | None = None, + gemm1_beta: float | None = None, +) -> FusedMoEQuantConfig: + """ + Construct a quant config for 16-bit float activations and int8 weights. + """ + group_shape = GroupShape(*block_shape) if block_shape is not None else None + return FusedMoEQuantConfig( + _a1=FusedMoEQuantDesc(shape=group_shape, alpha_or_gscale=a1_gscale), + _a2=FusedMoEQuantDesc(shape=group_shape, alpha_or_gscale=a2_gscale), + _w1=FusedMoEQuantDesc(torch.int8, group_shape, w1_scale, None, w1_zp, w1_bias), + _w2=FusedMoEQuantDesc(torch.int8, group_shape, w2_scale, None, w2_zp, w2_bias), + gemm1_clamp_limit=gemm1_clamp_limit, + gemm1_alpha=gemm1_alpha, + gemm1_beta=gemm1_beta, + ) + + +def int4_w4afp8_moe_quant_config( + w1_scale: torch.Tensor, + w2_scale: torch.Tensor, + g1_alphas: torch.Tensor, + g2_alphas: torch.Tensor, + per_act_token_quant: bool = False, + per_out_ch_quant: bool = False, + block_shape: list[int] | None = None, +) -> FusedMoEQuantConfig: + """ + Construct a quant config for fp8 activations and int4 weights. + """ + return FusedMoEQuantConfig.make( + torch.float8_e4m3fn, # quant dtype for activations + w1_scale=w1_scale, + w2_scale=w2_scale, + g1_alphas=g1_alphas, + g2_alphas=g2_alphas, + per_act_token_quant=per_act_token_quant, + per_out_ch_quant=per_out_ch_quant, + block_shape=block_shape, + weight_dtype="int4", # weight dtype for weights + ) + + +def biased_moe_quant_config( + w1_bias: torch.Tensor | None, + w2_bias: torch.Tensor | None, + gemm1_alpha: float | None = None, + gemm1_beta: float | None = None, + gemm1_clamp_limit: float | None = None, +) -> FusedMoEQuantConfig: + """ + Construct a quant config for unquantized activations with biases. + + gemm1_alpha/gemm1_beta/gemm1_clamp_limit carry the SwiGLU gate params + through to the fused activation kernel (e.g. swigluoai_uninterleave). + """ + return FusedMoEQuantConfig( + _a1=FusedMoEQuantDesc(), + _a2=FusedMoEQuantDesc(), + _w1=FusedMoEQuantDesc(bias=w1_bias), + _w2=FusedMoEQuantDesc(bias=w2_bias), + gemm1_alpha=gemm1_alpha, + gemm1_beta=gemm1_beta, + gemm1_clamp_limit=gemm1_clamp_limit, + ) + + +# A FusedMoEQuantConfig constant for an unquantized MoE op. +FUSED_MOE_UNQUANTIZED_CONFIG: FusedMoEQuantConfig = FusedMoEQuantConfig.make() + + +@dataclass +class FusedMoEParallelConfig: + tp_size: int + pcp_size: int + dp_size: int + ep_size: int + tp_rank: int + pcp_rank: int + dp_rank: int + ep_rank: int + sp_size: int + + use_ep: bool # whether to use EP or not + all2all_backend: str # all2all backend for MoE communication + enable_eplb: bool # whether to enable expert load balancing + + @property + def is_sequence_parallel(self) -> bool: + return self.sp_size > 1 + + @property + def use_all2all_kernels(self): + return self.use_ep and (self.dp_size > 1 or self.pcp_size > 1 or self.is_sequence_parallel) + + @property + def use_deepep_ht_kernels(self): + return self.use_all2all_kernels and self.all2all_backend == "deepep_high_throughput" + + @property + def use_deepep_ll_kernels(self): + return self.use_all2all_kernels and self.all2all_backend == "deepep_low_latency" + + @property + def use_fi_nvl_two_sided_kernels(self): + return self.use_all2all_kernels and ( + self.all2all_backend == "flashinfer_all2allv" or self.all2all_backend == "flashinfer_nvlink_two_sided" + ) + + @property + def use_fi_nvl_one_sided_kernels(self): + return self.use_all2all_kernels and self.all2all_backend == "flashinfer_nvlink_one_sided" + + @property + def use_batched_activation_format(self): + return self.use_deepep_ll_kernels or self.use_nixl_ep_kernels + + @property + def needs_round_robin_routing_tables(self): + return self.use_deepep_ll_kernels or self.use_nixl_ep_kernels + + @property + def use_ag_rs_all2all_kernels(self): + return self.use_all2all_kernels and self.all2all_backend == "allgather_reducescatter" + + @property + def use_mori_kernels(self): + return self.use_all2all_kernels and self.all2all_backend in ( + "mori_high_throughput", + "mori_low_latency", + ) + + @property + def use_nixl_ep_kernels(self): + return self.use_all2all_kernels and self.all2all_backend == "nixl_ep" + + @property + def use_deepep_v2_kernels(self): + return self.use_all2all_kernels and self.all2all_backend == "deepep_v2" + + @staticmethod + def flatten_tp_across_dp_and_pcp( + tp_size: int, dp_size: int, dp_rank: int, pcp_size: int, pcp_rank: int + ) -> tuple[int, int]: + tp_rank = 0 if tp_size == 1 else get_tensor_model_parallel_rank() + # There are actually dp_size * pcp_size * tp_size devices. + # Update tp_size and tp_rank so we shard across all devices. + flatten_tp_size = dp_size * pcp_size * tp_size + flatten_tp_rank = dp_rank * pcp_size * tp_size + pcp_rank * tp_size + tp_rank + return flatten_tp_size, flatten_tp_rank + + @staticmethod + def make( + tp_size_: int, + pcp_size_: int, + dp_size_: int, + sp_size_: int, + aphrodite_parallel_config: ParallelConfig, + ) -> "FusedMoEParallelConfig": + """ + Determine MoE parallel configuration. Based on the input `tp_size_`, + `dp_size_` and aphrodite's parallel config, determine what + level's of parallelism to use in the fused moe layer. + + Args: + tp_size_ (int): `tp_size` passed into the FusedMoE constructor. + pcp_size_ (int): `pcp_size` passed into the FusedMoE constructor. + dp_size_ (int): `dp_size` passed into the FusedMoE constructor. + aphrodite_parallel_config (ParallelConfig): Aphrodite's parallel config + object which contains the `enable_expert_parallel` flag. + + Examples: + When there is no parallelism requested, + i.e. `tp_size_` = `pcp_size_` = `dp_size_` = 1, we simply return the sizes + unaltered and the ranks set to 0. + + Expert Parallelism is considered only when either `dp_size_`, `pcp_size_` or + `tp_size_` is non trivial. + + Note that PCP serves the same function as DP here. + + When TP = 2, DP(PCP) = 1 and EP = False, the configuration on different + devices: + + - device 0 : TP = {2, 0} DP = {1, 0} EP = {1, 0} // + legend : {size, rank} + - device 1 : TP = {2, 1} DP = {1, 0} EP = {1, 0} + - Comment : Tensors are sharded across 2 devices. + + When TP = 1, DP(PCP) = 2 and EP = False, the configuration on different + devices: + + - device 0 : TP = {2, 0} DP = {2, 0} EP = {1, 0} + - device 1 : TP = {2, 1} DP = {2, 1} EP = {1, 0} + - Comment: There are 2 engine instances and the tensors are sharded + across 2 decvices. + + When TP = 2, DP(PCP) = 2 and EP = False, the configuration on different + devices: + + - device 0: TP = {4, 0} DP = {2, 0} EP = {1, 0} + - device 1: TP = {4, 1} DP = {2, 0} EP = {1, 0} + - device 2: TP = {4, 2} DP = {2, 1} EP = {1, 0} + - device 3: TP = {4, 3} DP = {2, 1} EP = {1, 0} + - Comment: There are 2 engine instances and the tensors are sharded + across 4 devices. + + When, TP = 2, DP(PCP) = 1 and EP = True, the configuration on different + devices: + + - device 0: TP = {1, 0} DP = {1, 0} EP = {2, 0} + - device 1: TP = {1, 0} DP = {1, 0} EP = {2, 1} + - Comment: The experts are split between the 2 devices. + + When, TP = 1, DP(PCP) = 2 and EP = True, the configuration on different + devices: + + - device 0: TP = {1, 0} DP = {2, 0} EP = {2, 0} + - device 1: TP = {1, 0} DP = {2, 1} EP = {2, 1} + - Comment: There are 2 engine instances and the experts are split + between the 2 devices. + + When TP = 2, DP(PCP) = 2 and EP = True, the configuration on different + devices: + + - device 0: TP = {1, 0} DP = {2, 0} EP = {4, 0} + - device 1: TP = {1, 0} DP = {2, 0} EP = {4, 1} + - device 2: TP = {1, 0} DP = {2, 1} EP = {4, 2} + - device 3: TP = {1, 0} DP = {2, 1} EP = {4, 3} + - Comment: There are 2 engine instances and the experts are split + between the 4 devices. + """ + use_ep = dp_size_ * pcp_size_ * tp_size_ > 1 and aphrodite_parallel_config.enable_expert_parallel + + dp_size = dp_size_ + dp_rank = get_dp_group().rank_in_group if dp_size > 1 else 0 + pcp_size = pcp_size_ + pcp_rank = get_pcp_group().rank_in_group if pcp_size > 1 else 0 + tp_size, tp_rank = FusedMoEParallelConfig.flatten_tp_across_dp_and_pcp( + tp_size_, dp_size_, dp_rank, pcp_size_, pcp_rank + ) + + if not use_ep: + return FusedMoEParallelConfig( + tp_size=tp_size, + tp_rank=tp_rank, + pcp_size=pcp_size, + pcp_rank=pcp_rank, + dp_size=dp_size, + dp_rank=dp_rank, + ep_size=1, + ep_rank=0, + sp_size=sp_size_, + use_ep=False, + all2all_backend=aphrodite_parallel_config.all2all_backend, + enable_eplb=aphrodite_parallel_config.enable_eplb, + ) + # DP + EP / TP + EP / DP + TP + EP + assert use_ep + # In EP, each device owns a set of experts fully. There is no tensor + # parallel update tp_size, tp_rank, ep_size and ep_rank to reflect that. + ep_size = tp_size + ep_rank = tp_rank + return FusedMoEParallelConfig( + tp_size=1, + tp_rank=0, + pcp_size=pcp_size, + pcp_rank=pcp_rank, + dp_size=dp_size, + dp_rank=dp_rank, + ep_size=ep_size, + ep_rank=ep_rank, + sp_size=sp_size_, + use_ep=True, + all2all_backend=aphrodite_parallel_config.all2all_backend, + enable_eplb=aphrodite_parallel_config.enable_eplb, + ) + + @classmethod + def make_no_parallel(cls) -> "FusedMoEParallelConfig": + """For usage in CI/CD and testing.""" + return FusedMoEParallelConfig( + tp_size=1, + tp_rank=0, + pcp_size=1, + pcp_rank=0, + dp_size=1, + dp_rank=0, + ep_size=1, + ep_rank=0, + sp_size=1, + use_ep=False, + all2all_backend="allgather_reducescatter", + enable_eplb=False, + ) + + +# Adapted from pplx-kernels tests/all_to_all_utils.py +@dataclass +class FusedMoEConfig: + num_experts: int + experts_per_token: int + hidden_dim: int + intermediate_size: int + num_local_experts: int + num_logical_experts: int + activation: MoEActivation + device: torch.device | str + routing_method: RoutingMethodType + moe_parallel_config: FusedMoEParallelConfig + + # The activation type. + in_dtype: torch.dtype + + # Defaults to in_dtype if not specified. + router_logits_dtype: torch.dtype | None = None + + # Defaults to hidden_dim if not specified. + hidden_dim_unpadded: int | None = None + # Defaults to intermediate_size_per_partition if not specified. + intermediate_size_per_partition_unpadded: int | None = None + # Model specific override + intermediate_pad: int | None = None + + moe_backend: MoEBackend = "auto" + max_num_tokens: int = SchedulerConfig.DEFAULT_MAX_NUM_BATCHED_TOKENS_FOR_BATCHED_DP + has_bias: bool = False + is_lora_enabled: bool = False + + # When True, the MoE skips its final cross-rank all-reduce (and the separate + # shared-expert reduce), returning the partial per-rank sum. The caller is + # then responsible for the reduction (e.g. fusing it into the next RMSNorm). + # Only honored on the non-reduced (late-AR) TP path. Default False. + skip_final_all_reduce: bool = False + + # SwiGLU clamp limit. When set, backends that do not implement the clamp + # are filtered out by `FusedMoEExperts.is_supported_config` so the oracle + # cannot silently select one and drop the clamp. + swiglu_limit: float | None = None + swiglu_alpha: float | None = None + swiglu_beta: float | None = None + + max_capture_size: int = 0 + + # Set by __post_init__ + intermediate_size_per_partition: int = -1 + rocm_aiter_fmoe_enabled: bool = False + aiter_fmoe_shared_expert_enabled: bool = False + + def __post_init__(self): + from aphrodite._aiter_ops import rocm_aiter_ops + + tp_size = self.moe_parallel_config.tp_size + assert self.intermediate_size % tp_size == 0 + self.intermediate_size_per_partition = self.intermediate_size // tp_size + + if self.dp_size > 1: + logger.debug_once("Using FusedMoEConfig::max_num_tokens=%d", self.max_num_tokens) + + assert self.max_num_tokens > 0 + + if self.router_logits_dtype is None: + self.router_logits_dtype = self.in_dtype + + if self.hidden_dim_unpadded is None: + self.hidden_dim_unpadded = self.hidden_dim + if self.intermediate_size_per_partition_unpadded is None: + self.intermediate_size_per_partition_unpadded = self.intermediate_size_per_partition + + if self.is_act_and_mul: + self.rocm_aiter_fmoe_enabled = rocm_aiter_ops.is_fused_moe_enabled() + self.aiter_fmoe_shared_expert_enabled = rocm_aiter_ops.is_fusion_moe_shared_experts_enabled() + + if self.use_mori_kernels: + assert self.rocm_aiter_fmoe_enabled, "Mori needs to be used with aiter fused_moe for now." + assert not self.aiter_fmoe_shared_expert_enabled, ( + "Mori does not support fusion shared expert now. " + "Turn it off by setting APHRODITE_ROCM_USE_AITER_FUSION_SHARED_EXPERTS=0" + ) + + if not self.is_act_and_mul and not (current_platform.is_cuda_alike() or current_platform.is_xpu()): + raise NotImplementedError("is_act_and_mul=False is supported only for CUDA, XPU and ROCm for now") + + @property + def is_act_and_mul(self) -> bool: + return self.activation.is_gated + + @property + def tp_size(self): + return self.moe_parallel_config.tp_size + + @property + def dp_size(self): + return self.moe_parallel_config.dp_size + + @property + def pcp_size(self): + return self.moe_parallel_config.pcp_size + + @property + def ep_size(self): + return self.moe_parallel_config.ep_size + + @property + def sp_size(self): + return self.moe_parallel_config.sp_size + + @property + def is_sequence_parallel(self): + return self.moe_parallel_config.is_sequence_parallel + + @property + def tp_rank(self): + return self.moe_parallel_config.tp_rank + + @property + def dp_rank(self): + return self.moe_parallel_config.dp_rank + + @property + def pcp_rank(self): + return self.moe_parallel_config.pcp_rank + + @property + def ep_rank(self): + return self.moe_parallel_config.ep_rank + + @property + def use_ep(self): + return self.moe_parallel_config.use_ep + + @property + def use_deepep_ht_kernels(self): + return self.moe_parallel_config.use_deepep_ht_kernels + + @property + def use_deepep_ll_kernels(self): + return self.moe_parallel_config.use_deepep_ll_kernels + + @property + def use_mori_kernels(self): + return self.moe_parallel_config.use_mori_kernels + + @property + def use_fi_nvl_two_sided_kernels(self): + return self.moe_parallel_config.use_fi_nvl_two_sided_kernels + + @property + def use_fi_nvl_one_sided_kernels(self): + return self.moe_parallel_config.use_fi_nvl_one_sided_kernels + + @property + def use_ag_rs_all2all_kernels(self): + return self.moe_parallel_config.use_ag_rs_all2all_kernels + + @property + def use_nixl_ep_kernels(self): + return self.moe_parallel_config.use_nixl_ep_kernels + + @property + def use_deepep_v2_kernels(self): + return self.moe_parallel_config.use_deepep_v2_kernels + + @property + def needs_round_robin_routing_tables(self): + return self.moe_parallel_config.needs_round_robin_routing_tables 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/deep_gemm_moe.py.orig b/aphrodite/model_executor/layers/fused_moe/experts/deep_gemm_moe.py.orig new file mode 100644 index 0000000000..60995b557e --- /dev/null +++ b/aphrodite/model_executor/layers/fused_moe/experts/deep_gemm_moe.py.orig @@ -0,0 +1,580 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +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.config import ( + FusedMoEConfig, + FusedMoEParallelConfig, + FusedMoEQuantConfig, +) +from aphrodite.model_executor.layers.fused_moe.deep_gemm_utils import ( + compute_aligned_M_and_alignment, + deepgemm_moe_permute, + deepgemm_unpermute_and_reduce, +) +from aphrodite.model_executor.layers.fused_moe.topk_weight_and_reduce import ( + TopKWeightAndReduceNoOP, +) +from aphrodite.model_executor.layers.fused_moe.utils import _resize_cache +from aphrodite.model_executor.layers.quantization.utils.fp8_utils import ( + per_token_group_quant_fp8, + per_token_group_quant_fp8_packed_for_deepgemm, + silu_mul_per_token_group_quant_fp8_colmajor, +) +from aphrodite.model_executor.layers.quantization.utils.fp8_utils import ( + silu_mul_quant_fp8_packed_triton as fused_silu_mul_fp8_quant_packed, +) +from aphrodite.model_executor.layers.quantization.utils.quant_utils import ( + QuantKey, + kFp8Dynamic128Sym, + kFp8Static128BlockSym, + kMxfp4Static, + kMxfp8Dynamic, + kMxfp8Static, +) +from aphrodite.platforms import current_platform +from aphrodite.utils.deep_gemm import ( + DeepGemmQuantScaleFMT, + get_mk_alignment_for_contiguous_layout, + is_deep_gemm_supported, + m_grouped_fp8_fp4_gemm_nt_contiguous, + m_grouped_fp8_gemm_nt_contiguous, + mk_alignment_scope, +) +from aphrodite.utils.import_utils import has_deep_gemm + +logger = init_logger(__name__) + + +def _valid_deep_gemm_shape(M: int, N: int, K: int) -> bool: + align = get_mk_alignment_for_contiguous_layout()[0] + return align <= M and N % align == 0 and K % align == 0 + + +def _valid_deep_gemm(hidden_states: torch.Tensor, w1: torch.Tensor, w2: torch.Tensor) -> bool: + """ + Check if the given problem size is supported by the DeepGemm grouped + gemm kernel. All of M, N, K and the quantization block_shape must be + aligned by `dg.get_m_alignment_for_contiguous_layout()`. + """ + if not has_deep_gemm(): + logger.debug_once("DeepGemm disabled: deep_gemm not available.") + return False + + M = hidden_states.size(0) + _, K, N = w2.size() + + align = get_mk_alignment_for_contiguous_layout()[0] + + if not _valid_deep_gemm_shape(M, N, K): + logger.debug_once( + "DeepGemm disabled due to unaligned problem size. " + "M: %s, N: %s, K: %s. M should >= %s " + "and N and K must be multiples of %s. " + "This is not an error and we will fall back to triton.", + M, + N, + K, + align, + align, + ) + return False + elif N <= 512: + logger.debug_once( + "DeepGemm disabled for N <= 512. M: %s, N: %s, K: %s. " + "This means we will fallback to triton " + "for this specific shape for further speed up.", + M, + N, + K, + ) + return False + + if w1.dtype != torch.float8_e4m3fn or w2.dtype != torch.float8_e4m3fn: + logger.debug_once( + "DeepGemm disabled: invalid weight dtype(s). w1.dtype: %s, w2.dtype: %s", + w1.dtype, + w2.dtype, + ) + return False + + if not hidden_states.is_contiguous() or not w1.is_contiguous() or not w2.is_contiguous(): + logger.debug_once( + "DeepGemm disabled: weights or activations not contiguous. " + "hidden_states.is_contiguous(): %s, w1.is_contiguous(): %s, " + "w2.is_contiguous(): %s", + hidden_states.is_contiguous(), + w1.is_contiguous(), + w2.is_contiguous(), + ) + return False + + return True + + +class DeepGemmExperts(mk.FusedMoEExpertsModular): + """DeepGemm-based fused MoE expert implementation.""" + + def __init__(self, moe_config: FusedMoEConfig, quant_config: FusedMoEQuantConfig): + super().__init__(moe_config=moe_config, quant_config=quant_config) + # MXFP8: FP8 e4m3 values + UE8M0 1x32 block scales (Blackwell). Reuses + # the same grouped GEMM (aliased to fp8_fp4) with recipe (1, 32). + self.mxfp8 = quant_config.block_shape == [1, 32] + if self.mxfp8: + assert quant_config.quant_dtype == "mxfp8" + else: + assert quant_config.block_shape == get_mk_alignment_for_contiguous_layout() + assert quant_config.quant_dtype == torch.float8_e4m3fn + assert not quant_config.per_act_token_quant + assert not quant_config.per_out_ch_quant + + self.gemm1_clamp_limit = quant_config.gemm1_clamp_limit + # Gated-activation params: silu == swigluoai with alpha=1, beta=0. + # FP8 (silu) configs leave these None, reproducing plain silu. + self.gemm1_alpha = quant_config.gemm1_alpha if quant_config.gemm1_alpha is not None else 1.0 + self.gemm1_beta = quant_config.gemm1_beta if quant_config.gemm1_beta is not None else 0.0 + + @staticmethod + def activation_format() -> mk.FusedMoEActivationFormat: + return mk.FusedMoEActivationFormat.Standard + + @staticmethod + def _supports_current_device() -> bool: + return is_deep_gemm_supported() + + @staticmethod + def _supports_no_act_and_mul() -> bool: + return False + + @staticmethod + def _supports_quant_scheme( + weight_key: QuantKey | None, + activation_key: QuantKey | None, + ) -> bool: + if (weight_key, activation_key) == (kFp8Static128BlockSym, kFp8Dynamic128Sym): + return True + # MXFP8 1x32 uses the fp8_fp4 grouped GEMM with recipe (1, 32) — only + # available on Blackwell (SM100). + if (weight_key, activation_key) == (kMxfp8Static, kMxfp8Dynamic): + return current_platform.is_device_capability_family(100) + return False + + @staticmethod + def _supports_activation(activation: MoEActivation) -> bool: + # silu/swigluoai go through the fused alpha/beta kernel; swiglustep + # uses the unfused activation path. The fused kernel reads packed w13 + # (gate = first half, up = second half), so it implements the + # *uninterleaved* SwiGLU-OAI variant. + return activation in [ + MoEActivation.SILU, + MoEActivation.SWIGLUSTEP, + MoEActivation.SWIGLUOAI_UNINTERLEAVE, + ] + + @staticmethod + def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bool: + # NOTE(rob): discovered an IMA with this combination. Needs investigation. + return not ( + moe_parallel_config.use_fi_nvl_two_sided_kernels or moe_parallel_config.use_fi_nvl_one_sided_kernels + ) + + def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce: + return TopKWeightAndReduceNoOP() + + def workspace_shapes( + self, + M: int, + N: int, + K: int, + topk: int, + global_num_experts: int, + local_num_experts: int, + expert_tokens_meta: mk.ExpertTokensMetadata | None, + activation: MoEActivation, + ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]: + assert self.block_shape is not None + # Use the contiguous-layout M alignment (matches apply()); block_shape[0] + # is the quant block (1 for MXFP8) and would under-size the workspace. + block_m = get_mk_alignment_for_contiguous_layout()[0] + M_sum, align_used = compute_aligned_M_and_alignment(M, topk, local_num_experts, block_m, expert_tokens_meta) + assert M_sum % align_used == 0 + + activation_out_dim = self.adjust_N_for_activation(N, activation) + workspace1 = (M_sum, max(activation_out_dim, K)) + workspace2 = (M_sum, max(N, K)) + output = (M, K) + return (workspace1, workspace2, output) + + def _act_mul_quant( + self, input: torch.Tensor, output: torch.Tensor, activation: MoEActivation + ) -> tuple[torch.Tensor, torch.Tensor]: + assert self.block_shape is not None + block_k = self.block_shape[1] + scale_fmt = DeepGemmQuantScaleFMT.from_oracle() + + M_sum, N = input.size() + activation_out_dim = self.adjust_N_for_activation(N, activation) + + # silu and swigluoai are both expressible by the fused gated kernel via + # (alpha, beta): silu uses alpha=1, beta=0; swigluoai uses config values. + # The fused kernel reads packed w13, hence SWIGLUOAI_UNINTERLEAVE. + fused_gated = activation in ( + MoEActivation.SILU, + MoEActivation.SWIGLUOAI_UNINTERLEAVE, + ) + + # 1. DeepGemm UE8M0: fused gate+mul+clamp+quant+pack + if scale_fmt == DeepGemmQuantScaleFMT.UE8M0: + if fused_gated: + return fused_silu_mul_fp8_quant_packed( + input=input, + output_q=output, + group_size=block_k, + clamp_limit=self.gemm1_clamp_limit, + alpha=self.gemm1_alpha, + beta=self.gemm1_beta, + ) + act_out = torch.empty((M_sum, activation_out_dim), dtype=input.dtype, device=input.device) + self.activation(activation, act_out, input) + a2q, a2q_scale = per_token_group_quant_fp8_packed_for_deepgemm( + act_out, + block_k, + out_q=output, + ) + return a2q, a2q_scale + + # 2. Hopper / non‑E8M0: prefer the fused gate+mul+quant kernel + if fused_gated: + use_ue8m0 = scale_fmt == DeepGemmQuantScaleFMT.FLOAT32_CEIL_UE8M0 + return silu_mul_per_token_group_quant_fp8_colmajor( + input=input, + output=output, + use_ue8m0=use_ue8m0, + clamp_limit=self.gemm1_clamp_limit, + group_size=block_k, + alpha=self.gemm1_alpha, + beta=self.gemm1_beta, + ) + + # 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) + + def apply( + self, + output: torch.Tensor, + hidden_states: torch.Tensor, + w1: torch.Tensor, + w2: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + activation: MoEActivation, + global_num_experts: int, + expert_map: torch.Tensor | None, + a1q_scale: torch.Tensor | None, + a2_scale: torch.Tensor | None, + workspace13: torch.Tensor, + workspace2: torch.Tensor, + expert_tokens_meta: mk.ExpertTokensMetadata | None, + apply_router_weight_on_input: bool, + ): + assert a1q_scale is not None + assert a2_scale is None + assert self.block_shape is not None + assert self.w1_scale is not None + assert self.w2_scale is not None + + a1q = hidden_states + _, N, K = w1.size() + + local_num_experts = w1.size(0) + if global_num_experts == -1: + global_num_experts = local_num_experts + + assert w2.size(1) == K + + M_sum, _ = compute_aligned_M_and_alignment( + M=topk_ids.size(0), + num_topk=topk_ids.size(1), + local_num_experts=local_num_experts, + alignment=get_mk_alignment_for_contiguous_layout()[0], + expert_tokens_meta=expert_tokens_meta, + ) + + a1q_perm = _resize_cache(workspace13.view(dtype=torch.float8_e4m3fn), (M_sum, K)) + a1q, a1q_scale, expert_ids, inv_perm, align_used = deepgemm_moe_permute( + aq=a1q, + aq_scale=a1q_scale, + topk_ids=topk_ids, + local_num_experts=local_num_experts, + expert_map=expert_map, + expert_tokens_meta=expert_tokens_meta, + aq_out=a1q_perm, + # MXFP8 uses a 32-element activation-scale group (block_shape[1]); + # FP8-block keeps the default (128) alignment. + block_size=self.block_shape[1] if self.mxfp8 else None, + ) + assert a1q.size(0) == M_sum + + # MXFP8 (1x32) drives the fp8_fp4-aliased grouped GEMM with recipe + # (1, 32); the FP8 block path keeps the default (128) recipe. + gemm_kwargs = {"recipe_a": (1, self.block_shape[1]), "recipe_b": (1, self.block_shape[1])} if self.mxfp8 else {} + + # Cap DG's BLOCK_M heuristic at the workspace's per-expert alignment; + # otherwise the scheduler can pick the wrong expert id from m_indices + # under cudagraph replay. + with mk_alignment_scope(align_used): + mm1_out = _resize_cache(workspace2, (M_sum, N)) + m_grouped_fp8_gemm_nt_contiguous( + (a1q, a1q_scale), + (w1, self.w1_scale), + mm1_out, + expert_ids, + **gemm_kwargs, + ) + + activation_out_dim = self.adjust_N_for_activation(N, activation) + quant_out = _resize_cache(workspace13.view(dtype=torch.float8_e4m3fn), (M_sum, activation_out_dim)) + a2q, a2q_scale = self._act_mul_quant(input=mm1_out.view(-1, N), output=quant_out, activation=activation) + + mm2_out = _resize_cache(workspace2, (M_sum, K)) + m_grouped_fp8_gemm_nt_contiguous( + (a2q, a2q_scale), + (w2, self.w2_scale), + mm2_out, + expert_ids, + **gemm_kwargs, + ) + + if apply_router_weight_on_input: + topk_weights = torch.ones_like(topk_weights) + + deepgemm_unpermute_and_reduce( + a=mm2_out, + topk_ids=topk_ids, + topk_weights=topk_weights, + inv_perm=inv_perm, + expert_map=expert_map, + output=output, + ) + + +class DeepGemmFP4Experts(mk.FusedMoEExpertsModular): + """DeepGemm-based fused MoE expert implementation for FP4 weights. + + Uses m_grouped_fp8_fp4_gemm_nt_contiguous with FP8 activations and + MXFP4 (FP4 E2M1 packed as uint8) weights. Requires Blackwell-family + GPUs (SM100 datacenter or SM120 consumer). + """ + + # FP8 activation block size (hardcoded since mxfp4_w4a8 quant config + # does not set a block_shape on the activation descriptor). + _ACT_BLOCK_K = 128 + # FP4 weight block size + _WEIGHT_BLOCK_K = 32 + + def __init__(self, moe_config: FusedMoEConfig, quant_config: FusedMoEQuantConfig): + super().__init__(moe_config=moe_config, quant_config=quant_config) + assert quant_config.weight_quant_dtype == "mxfp4" + assert not quant_config.per_act_token_quant + assert not quant_config.per_out_ch_quant + + self.gemm1_clamp_limit = quant_config.gemm1_clamp_limit + + @staticmethod + def activation_format() -> mk.FusedMoEActivationFormat: + return mk.FusedMoEActivationFormat.Standard + + @staticmethod + def _supports_current_device() -> bool: + from aphrodite.platforms import current_platform + + return is_deep_gemm_supported() and ( + current_platform.is_device_capability_family(100) or current_platform.is_device_capability_family(120) + ) + + @staticmethod + def _supports_no_act_and_mul() -> bool: + return False + + @staticmethod + def _supports_quant_scheme( + weight_key: QuantKey | None, + activation_key: QuantKey | None, + ) -> bool: + SUPPORTED_W_A = [ + (kMxfp4Static, kFp8Dynamic128Sym), + ] + return (weight_key, activation_key) in SUPPORTED_W_A + + @staticmethod + def _supports_activation(activation: MoEActivation) -> bool: + # 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: + return not ( + moe_parallel_config.use_fi_nvl_two_sided_kernels or moe_parallel_config.use_fi_nvl_one_sided_kernels + ) + + def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce: + return TopKWeightAndReduceNoOP() + + def workspace_shapes( + self, + M: int, + N: int, + K: int, + topk: int, + global_num_experts: int, + local_num_experts: int, + expert_tokens_meta: mk.ExpertTokensMetadata | None, + activation: MoEActivation, + ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]: + block_m = get_mk_alignment_for_contiguous_layout()[0] + M_sum, align_used = compute_aligned_M_and_alignment(M, topk, local_num_experts, block_m, expert_tokens_meta) + assert M_sum % align_used == 0 + + activation_out_dim = self.adjust_N_for_activation(N, activation) + workspace1 = (M_sum, max(activation_out_dim, K)) + workspace2 = (M_sum, max(N, K)) + output = (M, K) + return (workspace1, workspace2, output) + + def _act_mul_quant( + self, input: torch.Tensor, output: torch.Tensor, activation: MoEActivation + ) -> tuple[torch.Tensor, torch.Tensor]: + block_k = self._ACT_BLOCK_K + scale_fmt = DeepGemmQuantScaleFMT.from_oracle() + + M_sum, N = input.size() + activation_out_dim = self.adjust_N_for_activation(N, activation) + + 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, + output=output, + use_ue8m0=use_ue8m0, + clamp_limit=self.gemm1_clamp_limit, + ) + + 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) + + def apply( + self, + output: torch.Tensor, + hidden_states: torch.Tensor, + w1: torch.Tensor, + w2: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + activation: MoEActivation, + global_num_experts: int, + expert_map: torch.Tensor | None, + a1q_scale: torch.Tensor | None, + a2_scale: torch.Tensor | None, + workspace13: torch.Tensor, + workspace2: torch.Tensor, + expert_tokens_meta: mk.ExpertTokensMetadata | None, + apply_router_weight_on_input: bool, + ): + assert a1q_scale is not None + assert a2_scale is None + assert self.w1_scale is not None + assert self.w2_scale is not None + + a1q = hidden_states + _, N, _ = w1.size() + # K comes from activations (full hidden dim), not from w1 which is + # packed FP4 (E, N, K//2). + K = a1q.size(1) + + local_num_experts = w1.size(0) + if global_num_experts == -1: + global_num_experts = local_num_experts + + M_sum, _ = compute_aligned_M_and_alignment( + M=topk_ids.size(0), + num_topk=topk_ids.size(1), + local_num_experts=local_num_experts, + alignment=get_mk_alignment_for_contiguous_layout()[0], + expert_tokens_meta=expert_tokens_meta, + ) + + a1q_perm = _resize_cache(workspace13.view(dtype=torch.float8_e4m3fn), (M_sum, K)) + a1q, a1q_scale, expert_ids, inv_perm, align_used = deepgemm_moe_permute( + aq=a1q, + aq_scale=a1q_scale, + topk_ids=topk_ids, + local_num_experts=local_num_experts, + expert_map=expert_map, + expert_tokens_meta=expert_tokens_meta, + aq_out=a1q_perm, + ) + assert a1q.size(0) == M_sum + + # Cap DG's BLOCK_M heuristic at the workspace's per-expert alignment; + # see DeepGemmExperts.apply for rationale. + with mk_alignment_scope(align_used): + # FC1: FP8 activations x FP4 weights + # DeepGEMM 2.4.2 requires FP4-packed weights as int8 (kPackedFP4). + mm1_out = _resize_cache(workspace2, (M_sum, N)) + m_grouped_fp8_fp4_gemm_nt_contiguous( + (a1q, a1q_scale), + (w1.view(torch.int8), self.w1_scale), + mm1_out, + expert_ids, + recipe_a=(1, self._ACT_BLOCK_K), + recipe_b=(1, self._WEIGHT_BLOCK_K), + ) + + # SwiGLU activation + FP8 requant + activation_out_dim = self.adjust_N_for_activation(N, activation) + quant_out = _resize_cache(workspace13.view(dtype=torch.float8_e4m3fn), (M_sum, activation_out_dim)) + a2q, a2q_scale = self._act_mul_quant(input=mm1_out.view(-1, N), output=quant_out, activation=activation) + + # FC2: FP8 activations x FP4 weights + mm2_out = _resize_cache(workspace2, (M_sum, K)) + m_grouped_fp8_fp4_gemm_nt_contiguous( + (a2q, a2q_scale), + (w2.view(torch.int8), self.w2_scale), + mm2_out, + expert_ids, + recipe_a=(1, self._ACT_BLOCK_K), + recipe_b=(1, self._WEIGHT_BLOCK_K), + ) + + if apply_router_weight_on_input: + topk_weights = torch.ones_like(topk_weights) + + deepgemm_unpermute_and_reduce( + a=mm2_out, + topk_ids=topk_ids, + topk_weights=topk_weights, + inv_perm=inv_perm, + expert_map=expert_map, + output=output, + ) 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/marlin_moe.py.orig b/aphrodite/model_executor/layers/fused_moe/experts/marlin_moe.py.orig new file mode 100644 index 0000000000..145a04f14f --- /dev/null +++ b/aphrodite/model_executor/layers/fused_moe/experts/marlin_moe.py.orig @@ -0,0 +1,1036 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Fused MoE utilities for GPTQ.""" + +import math +from collections.abc import Callable + +import torch + +import aphrodite._custom_ops as ops +import aphrodite.model_executor.layers.fused_moe.modular_kernel as mk +from aphrodite.model_executor.layers.fused_moe.activation import ( + MoEActivation, + apply_moe_activation, +) +from aphrodite.model_executor.layers.fused_moe.config import ( + FusedMoEConfig, + FusedMoEParallelConfig, + FusedMoEQuantConfig, +) +from aphrodite.model_executor.layers.fused_moe.experts.lora_experts_mixin import ( + LoRAExpertsMixin, +) +from aphrodite.model_executor.layers.fused_moe.moe_align_block_size import ( + batched_moe_align_block_size, + moe_align_block_size, +) +from aphrodite.model_executor.layers.fused_moe.topk_weight_and_reduce import ( + TopKWeightAndReduceDelegate, + TopKWeightAndReduceNoOP, +) +from aphrodite.model_executor.layers.fused_moe.utils import _resize_cache +from aphrodite.model_executor.layers.quantization.utils.marlin_utils import ( + get_marlin_input_dtype, + marlin_make_workspace_new, + marlin_moe_intermediate_size, + marlin_quant_input, +) +from aphrodite.model_executor.layers.quantization.utils.quant_utils import ( + QuantKey, + kFp8Static128BlockSym, + kFp8StaticChannelSym, + kFp8StaticTensorSym, + kInt4Static, + kInt4Static32, + kInt4Static32Asym, + kInt4StaticAsym, + kInt8Static, + kMxfp4Static, + kMxfp8Static, + kNvfp4Static, +) +from aphrodite.platforms import current_platform +from aphrodite.scalar_type import ScalarType, scalar_types + + +def _fused_marlin_moe( + hidden_states: torch.Tensor, + w1: torch.Tensor, + w2: torch.Tensor, + bias1: torch.Tensor | None, + bias2: torch.Tensor | None, + w1_scale: torch.Tensor, + w2_scale: torch.Tensor, + topk_weights: torch.Tensor, + num_topk: int, + quant_type: ScalarType, + apply_router_weight_on_input: bool, + expert_map: torch.Tensor | None, + block_size_m: int, + sorted_token_ids: torch.Tensor, + expert_ids: torch.Tensor, + num_tokens_post_padded: torch.Tensor, + activation: MoEActivation = MoEActivation.SILU, + activation_func: Callable[..., None] = apply_moe_activation, + topk_ids: torch.Tensor | None = None, + input_global_scale1: torch.Tensor | None = None, + input_global_scale2: torch.Tensor | None = None, + global_scale1: torch.Tensor | None = None, + global_scale2: torch.Tensor | None = None, + g_idx1: torch.Tensor | None = None, + g_idx2: torch.Tensor | None = None, + sort_indices1: torch.Tensor | None = None, + sort_indices2: torch.Tensor | None = None, + w1_zeros: torch.Tensor | None = None, + w2_zeros: torch.Tensor | None = None, + workspace: torch.Tensor | None = None, + intermediate_cache13: torch.Tensor | None = None, + intermediate_cache2: torch.Tensor | None = None, + output: torch.Tensor | None = None, + input_dtype: torch.dtype | None = None, + is_k_full: bool = True, + clamp_limit: float | None = None, + gemm1_alpha: float = 1.0, + gemm1_beta: float = 0.0, +) -> torch.Tensor: + assert hidden_states.ndim == 2 + M, K = hidden_states.size() + N = marlin_moe_intermediate_size(w1, w2) + w13_num_shards = 2 if activation.is_gated else 1 + if workspace is None: + workspace = marlin_make_workspace_new(hidden_states.device, 4) + + if intermediate_cache13 is None: + intermediate_cache13 = torch.empty( + (M * num_topk * max(w13_num_shards * N, K),), + device=hidden_states.device, + dtype=hidden_states.dtype, + ) + + if intermediate_cache2 is None: + intermediate_cache2 = torch.empty( + (M * num_topk, N), + device=hidden_states.device, + dtype=hidden_states.dtype, + ) + + intermediate_cache1 = _resize_cache(intermediate_cache13, (M * num_topk, w13_num_shards * N)) + + intermediate_cache3 = _resize_cache(intermediate_cache13, (M * num_topk, K)) + + intermediate_cache2 = _resize_cache(intermediate_cache2, (M * num_topk, N)) + + a_scales1 = None + gate_up_input = hidden_states + if input_dtype == torch.int8: + gate_up_input, a_scales1 = marlin_quant_input(hidden_states, input_dtype) + if input_global_scale1 is not None: + a_scales1 = a_scales1 * input_global_scale1 + elif input_dtype == torch.float8_e4m3fn: + gate_up_input, a_scales1 = marlin_quant_input(hidden_states, input_dtype) + + intermediate_cache1 = ops.moe_wna16_marlin_gemm( + gate_up_input, + intermediate_cache1, + w1, + bias1, + w1_scale, + a_scales1, + global_scale1, + w1_zeros, + g_idx1, + sort_indices1, + workspace, + sorted_token_ids, + expert_ids, + num_tokens_post_padded, + topk_weights, + moe_block_size=block_size_m, + top_k=num_topk, + mul_topk_weights=apply_router_weight_on_input, + b_q_type=quant_type, + size_m=M, + size_n=w13_num_shards * N, + size_k=K, + is_k_full=is_k_full, + use_atomic_add=False, + use_fp32_reduce=True, + is_zp_float=False, + ) + # apply_moe_activation fuses the clamp/gate params: SILU + clamp_limit and + # SWIGLUOAI_UNINTERLEAVE both map to the silu_and_mul_with_clamp kernel. + activation_func( + activation, + intermediate_cache2, + intermediate_cache1.view(-1, w13_num_shards * N), + clamp_limit=clamp_limit, + alpha=gemm1_alpha, + beta=gemm1_beta, + topk_ids=topk_ids, + expert_map=expert_map, + ) + + if output is None: + output = intermediate_cache3 + + a_scales2 = None + if input_dtype == torch.int8: + intermediate_cache2, a_scales2 = marlin_quant_input(intermediate_cache2, input_dtype) + if input_global_scale2 is not None: + a_scales2 = a_scales2 * input_global_scale2 + elif input_dtype == torch.float8_e4m3fn: + intermediate_cache2, a_scales2 = marlin_quant_input(intermediate_cache2, input_dtype) + + output = ops.moe_wna16_marlin_gemm( + intermediate_cache2, + output, + w2, + bias2, + w2_scale, + a_scales2, + global_scale2, + w2_zeros, + g_idx2, + sort_indices2, + workspace, + sorted_token_ids, + expert_ids, + num_tokens_post_padded, + topk_weights, + moe_block_size=block_size_m, + top_k=1, + mul_topk_weights=not apply_router_weight_on_input, + b_q_type=quant_type, + size_m=M * num_topk, + size_n=K, + size_k=N, + is_k_full=is_k_full, + use_atomic_add=False, + use_fp32_reduce=True, + is_zp_float=False, + ) + + return output + + +def fused_marlin_moe( + hidden_states: torch.Tensor, + w1: torch.Tensor, + w2: torch.Tensor, + bias1: torch.Tensor | None, + bias2: torch.Tensor | None, + w1_scale: torch.Tensor, + w2_scale: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + quant_type_id: int, + apply_router_weight_on_input: bool = False, + global_num_experts: int = -1, + activation: MoEActivation = MoEActivation.SILU, + activation_func: Callable[..., None] = apply_moe_activation, + moe_sum: Callable[..., torch.Tensor | None] | None = None, + expert_map: torch.Tensor | None = None, + input_global_scale1: torch.Tensor | None = None, + input_global_scale2: torch.Tensor | None = None, + global_scale1: torch.Tensor | None = None, + global_scale2: torch.Tensor | None = None, + g_idx1: torch.Tensor | None = None, + g_idx2: torch.Tensor | None = None, + sort_indices1: torch.Tensor | None = None, + sort_indices2: torch.Tensor | None = None, + w1_zeros: torch.Tensor | None = None, + w2_zeros: torch.Tensor | None = None, + workspace: torch.Tensor | None = None, + intermediate_cache13: torch.Tensor | None = None, + intermediate_cache2: torch.Tensor | None = None, + is_k_full: bool = True, + output: torch.Tensor | None = None, + input_dtype: torch.dtype | None = None, + clamp_limit: float | None = None, + gemm1_alpha: float = 1.0, + gemm1_beta: float = 0.0, +) -> torch.Tensor: + """ + This function computes a Mixture of Experts (MoE) layer using two sets of + weights, w1 and w2, and top-k gating mechanism. + + Parameters: + - hidden_states (torch.Tensor): The input tensor to the MoE layer. + - w1 (torch.Tensor): The first set of expert weights. + - w2 (torch.Tensor): The second set of expert weights. + - w1_scale (torch.Tensor): Scale to be used for w1. + - w2_scale (torch.Tensor): Scale to be used for w2. + - g_idx1 (torch.Tensor|None): The first set of act_order indices. + - g_idx2 (torch.Tensor|None): The second set of act_order indices. + - sort_indices1 (torch.Tensor|None): The first act_order input + permutation. + - sort_indices2 (torch.Tensor|None): The second act_order input + permutation. + - topk_weights (torch.Tensor): Top-k weights. + - topk_ids (torch.Tensor): Indices of topk-k elements. + - w1_zeros (torch.Tensor|None): Optional zero points to be used for w1. + - w2_zeros (torch.Tensor|None): Optional zero points to be used for w2. + - num_bits (bool): The number of bits in expert weights quantization. + + Returns: + - torch.Tensor: The output tensor after applying the MoE layer. + """ + + quant_type = ScalarType.from_id(quant_type_id) + assert quant_type in [ + scalar_types.uint4, + scalar_types.uint8b128, + scalar_types.uint4b8, + scalar_types.float8_e4m3fn, + scalar_types.float4_e2m1f, + ] + + bit4_scalar_types = [ + scalar_types.uint4, + scalar_types.uint4b8, + scalar_types.float4_e2m1f, + ] + num_bits = 4 if quant_type in bit4_scalar_types else 8 + + M, K = hidden_states.size() + E = w1.size(0) + topk = topk_ids.size(1) + + # Check constraints. + assert w1.size(1) * 16 == K, "Hidden size mismatch w1" + assert w2.size(2) // (num_bits // 2) == K, "Hidden size mismatch w2" + assert hidden_states.is_contiguous(), "Hidden_states must be contiguous" + assert w1.is_contiguous(), "Expert weights1 must be contiguous" + assert w2.is_contiguous(), "Expert weights2 must be contiguous" + assert hidden_states.dtype in [torch.float16, torch.bfloat16] + assert num_bits in [4, 8] + assert topk_weights.dtype == torch.float32 + + if global_num_experts == -1: + global_num_experts = E + else: + # Set M to estimated valid tokens per rank. + M = math.ceil(M * E / global_num_experts) + + # M block size selection logic + # TODO: tune this further for specific models + for block_size_m in [8, 16, 32, 48, 64]: + if M * topk / E / block_size_m < 0.9: + break + + if input_dtype is not None and input_dtype.itemsize == 1: + block_size_m = max(block_size_m, 16) + + sorted_token_ids, expert_ids, num_tokens_post_padded = moe_align_block_size( + topk_ids, + block_size_m, + global_num_experts, + expert_map, + ignore_invalid_experts=True, + ) + + assert activation is not None + moe_output = _fused_marlin_moe( + hidden_states=hidden_states, + w1=w1, + w2=w2, + bias1=bias1, + bias2=bias2, + w1_scale=w1_scale, + w2_scale=w2_scale, + topk_weights=topk_weights, + topk_ids=topk_ids, + num_topk=topk, + quant_type=quant_type, + apply_router_weight_on_input=apply_router_weight_on_input, + expert_map=expert_map, + block_size_m=block_size_m, + sorted_token_ids=sorted_token_ids, + expert_ids=expert_ids, + num_tokens_post_padded=num_tokens_post_padded, + activation=activation, + activation_func=activation_func, + input_global_scale1=input_global_scale1, + input_global_scale2=input_global_scale2, + global_scale1=global_scale1, + global_scale2=global_scale2, + g_idx1=g_idx1, + g_idx2=g_idx2, + sort_indices1=sort_indices1, + sort_indices2=sort_indices2, + w1_zeros=w1_zeros, + w2_zeros=w2_zeros, + workspace=workspace, + intermediate_cache13=intermediate_cache13, + intermediate_cache2=intermediate_cache2, + output=None, + input_dtype=input_dtype, + is_k_full=is_k_full, + clamp_limit=clamp_limit, + gemm1_alpha=gemm1_alpha, + gemm1_beta=gemm1_beta, + ).view(-1, topk, K) + + if output is None: + output = torch.empty_like(hidden_states) + + if moe_sum is None: + if expert_map is not None: + ops.moe_sum(moe_output, output, topk_ids, expert_map) + return output + return torch.sum(moe_output.view(-1, topk, K), dim=1, out=output) + else: + return moe_sum(moe_output, output, topk_ids, expert_map) + + +def batched_fused_marlin_moe( + hidden_states: torch.Tensor, + expert_num_tokens: torch.Tensor, + w1: torch.Tensor, + w2: torch.Tensor, + bias1: torch.Tensor | None, + bias2: torch.Tensor | None, + w1_scale: torch.Tensor, + w2_scale: torch.Tensor, + quant_type_id: int, + apply_router_weight_on_input: bool = False, + global_num_experts: int = -1, + activation: MoEActivation = MoEActivation.SILU, + expert_map: torch.Tensor | None = None, + input_global_scale1: torch.Tensor | None = None, + input_global_scale2: torch.Tensor | None = None, + global_scale1: torch.Tensor | None = None, + global_scale2: torch.Tensor | None = None, + g_idx1: torch.Tensor | None = None, + g_idx2: torch.Tensor | None = None, + sort_indices1: torch.Tensor | None = None, + sort_indices2: torch.Tensor | None = None, + w1_zeros: torch.Tensor | None = None, + w2_zeros: torch.Tensor | None = None, + workspace: torch.Tensor | None = None, + intermediate_cache13: torch.Tensor | None = None, + intermediate_cache2: torch.Tensor | None = None, + is_k_full: bool = True, + output: torch.Tensor | None = None, + input_dtype: torch.dtype | None = None, + clamp_limit: float | None = None, + gemm1_alpha: float = 1.0, + gemm1_beta: float = 0.0, +) -> torch.Tensor: + """ + This function massages the inputs so the batched hidden_states can be + presented as a 2D contiguous tensor that could be used with + _fused_marlin_moe. + + Note that both batched_fused_marlin_moe and fused_marlin_moe ultimately + use `ops.moe_wna16_marlin_gemm` for the gemm operation and + `ops.moe_mna16_marlin_gemm` supports only 2D contiguous hidden_states. + Note that the moe_align_block_size function indicates, + - What rows of the A matrix (hidden_states) to access during the + matmul, via sorted_ids output. + - What expert_id to use for each block matmul, via expert_ids output. + + In the batched version, the tokens are already grouped/batched by experts + they subscribe to. Due to this, we can represent the batched hidden_states + tensor of shape [B, MAX_TOKENS_PER_BATCH, K] as a 2D tensor of shape, + [B * MAX_TOKENS_PER_BATCH, K]. We may treat this a 2D contiguous tensor + with topk=1 as each token (row in the tensor) subscribes to exactly one + expert_id (which is the batch_id). With the expert_num_tokens tensor, that + indicates how many tokens are actually valid in each batch, the + batched_moe_align_block_size function constructs the sorted_ids and + expert_ids tensors, so only relevant/valid rows of A (hidden_states) + are accessed and are processed with the correct expert_ids. + """ + + assert hidden_states.ndim == 3, ( + f"hidden states must be batched. e.g. [B, MAX_TOKENS, K].But got {hidden_states.size()}" + ) + + quant_type = ScalarType.from_id(quant_type_id) + assert quant_type in [ + scalar_types.uint4, + scalar_types.uint8b128, + scalar_types.uint4b8, + scalar_types.float8_e4m3fn, + scalar_types.float4_e2m1f, + ] + + bit4_scalar_types = [ + scalar_types.uint4, + scalar_types.uint4b8, + scalar_types.float4_e2m1f, + ] + num_bits = 4 if quant_type in bit4_scalar_types else 8 + + B, BATCH_TOKENS_MAX, K = hidden_states.size() + M = hidden_states.view(-1, K).size(0) + E = w1.size(0) + + # Check constraints. + assert hidden_states.is_contiguous(), "Hidden_states must be contiguous" + assert hidden_states.dtype in [torch.float16, torch.bfloat16] + assert expert_num_tokens.size(0) == E + assert B == E, ( + "Batch must be as big as number of experts as the tokensare sorted into the batch/expert they belong to" + ) + assert w1.size(1) * 16 == K, "Hidden size mismatch w1" + assert w2.size(2) // (num_bits // 2) == K, "Hidden size mismatch w2" + assert w1.is_contiguous(), "Expert weights1 must be contiguous" + assert w2.is_contiguous(), "Expert weights2 must be contiguous" + assert num_bits in [4, 8] + + # Technically, the tokens are already separated by their expert ids. + # Hidden-States can just be squeezed to have just 2 dimensions, + # [B * MAX_TOKENS, K] and top_k can be interpreted as just 1. + topk = 1 + + # TODO(varun) : Choose a decent block size like in fused_marlin_moe + # Tune block_size_m based on expert capacity to reduce padding overhead. + block_size_m = 64 + for b_m in [8, 16, 32, 48, 64]: + if BATCH_TOKENS_MAX / b_m < 0.9: + block_size_m = b_m + break + + if input_dtype is not None and input_dtype.itemsize == 1: + block_size_m = max(block_size_m, 16) + + sorted_token_ids, expert_ids, num_tokens_post_padded = batched_moe_align_block_size( + max_tokens_per_batch=BATCH_TOKENS_MAX, + block_size=block_size_m, + expert_num_tokens=expert_num_tokens, + ) + + # TODO (varun): This can be avoided by plumbing the marlin kernel to + # ignore topk_weights when topk_weights_ptr is a nullptr. + topk_weights = torch.ones((M, topk), device=hidden_states.device, dtype=torch.float32) + + assert activation is not None + output = _fused_marlin_moe( + hidden_states=hidden_states.view(-1, K), + w1=w1, + w2=w2, + bias1=bias1, + bias2=bias2, + w1_scale=w1_scale, + w2_scale=w2_scale, + topk_weights=topk_weights, + num_topk=topk, + quant_type=quant_type, + apply_router_weight_on_input=apply_router_weight_on_input, + activation=activation, + expert_map=expert_map, + block_size_m=block_size_m, + sorted_token_ids=sorted_token_ids, + expert_ids=expert_ids, + num_tokens_post_padded=num_tokens_post_padded, + input_global_scale1=input_global_scale1, + input_global_scale2=input_global_scale2, + global_scale1=global_scale1, + global_scale2=global_scale2, + g_idx1=g_idx1, + g_idx2=g_idx2, + sort_indices1=sort_indices1, + sort_indices2=sort_indices2, + w1_zeros=w1_zeros, + w2_zeros=w2_zeros, + workspace=workspace, + intermediate_cache13=intermediate_cache13, + intermediate_cache2=intermediate_cache2, + output=output.view(-1, K) if output is not None else output, + input_dtype=input_dtype, + is_k_full=is_k_full, + clamp_limit=clamp_limit, + gemm1_alpha=gemm1_alpha, + gemm1_beta=gemm1_beta, + ) + + output = output.view(B, BATCH_TOKENS_MAX, K) + + return output + + +class MarlinExpertsBase(mk.FusedMoEExpertsModular): + def __init__( + self, + moe_config: FusedMoEConfig, + quant_config: FusedMoEQuantConfig, + max_num_tokens: int | None = None, + num_dispatchers: int | None = None, + w13_g_idx: torch.Tensor | None = None, + w2_g_idx: torch.Tensor | None = None, + w13_g_idx_sort_indices: torch.Tensor | None = None, + w2_g_idx_sort_indices: torch.Tensor | None = None, + is_k_full: bool = True, + ): + # TODO (varun) : Enable activation quantization + assert ( + quant_config.use_mxfp4_w4a16 + or quant_config.use_nvfp4_w4a16 + or quant_config.use_int4_w4a16 + or quant_config.use_int8_w8a16 + or quant_config.use_fp8_w8a16 + ), "Supports only {mxfp,nvfp,int}4_w4a16, int8_w8a16 or fp8_w8a16" + self.w13_g_idx = w13_g_idx + self.w2_g_idx = w2_g_idx + self.w13_g_idx_sort_indices = w13_g_idx_sort_indices + self.w2_g_idx_sort_indices = w2_g_idx_sort_indices + self.is_k_full = is_k_full + self.input_dtype = get_marlin_input_dtype() + self.gemm1_clamp_limit = quant_config.gemm1_clamp_limit + # Gated-activation params (used by SWIGLUOAI_UNINTERLEAVE on packed w13). + # silu == swigluoai with alpha=1, beta=0; configs that don't set these + # (plain silu) fall back to the silu identity. + self.gemm1_alpha = quant_config.gemm1_alpha if quant_config.gemm1_alpha is not None else 1.0 + self.gemm1_beta = quant_config.gemm1_beta if quant_config.gemm1_beta is not None else 0.0 + + super().__init__( + moe_config=moe_config, + quant_config=quant_config, + max_num_tokens=max_num_tokens, + num_dispatchers=num_dispatchers, + ) + + @staticmethod + def _supports_current_device() -> bool: + p = current_platform + return p.is_cuda() and p.has_device_capability((7, 5)) + + @staticmethod + def _supports_no_act_and_mul() -> bool: + return True + + @staticmethod + def _supports_quant_scheme( + weight_key: QuantKey | None, + activation_key: QuantKey | None, + ) -> bool: + # TODO(rob): add int4, int8 as integrations + # are migrated to use the oracle one-by-one. + SUPPORTED_W = [ + kFp8Static128BlockSym, + kFp8StaticChannelSym, + kFp8StaticTensorSym, + kMxfp4Static, + kMxfp8Static, + kNvfp4Static, + kInt4Static, + kInt8Static, + kInt4Static32, + kInt4StaticAsym, + kInt4Static32Asym, + ] + return weight_key in SUPPORTED_W + + @staticmethod + def _supports_activation(activation: MoEActivation) -> bool: + # Marlin uses apply_moe_activation() callback for activation, + # so any activation supported there can be used here. + return activation in [ + MoEActivation.SILU, + MoEActivation.GELU, + MoEActivation.GELU_TANH, + MoEActivation.SWIGLUOAI, + MoEActivation.SWIGLUOAI_UNINTERLEAVE, + MoEActivation.SWIGLUSTEP, + MoEActivation.SILU_NO_MUL, + MoEActivation.GELU_NO_MUL, + MoEActivation.GELU_TANH_NO_MUL, + MoEActivation.RELU2_NO_MUL, + ] + + @staticmethod + def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bool: + return not ( + moe_parallel_config.use_fi_nvl_two_sided_kernels or moe_parallel_config.use_fi_nvl_one_sided_kernels + ) + + @property + def quant_type_id(self) -> int: + if self.quant_config.use_int4_w4a16: + if self.w1_zp is not None or self.w2_zp is not None: + return scalar_types.uint4.id + return scalar_types.uint4b8.id + elif self.quant_config.use_int8_w8a16: + return scalar_types.uint8b128.id + elif self.quant_config.use_mxfp4_w4a16 or self.quant_config.use_nvfp4_w4a16: + return scalar_types.float4_e2m1f.id + elif self.quant_config.use_fp8_w8a16 and current_platform.fp8_dtype() == torch.float8_e4m3fn: + return scalar_types.float8_e4m3fn.id + else: + raise NotImplementedError("Unsupported quantization type.") + + def moe_problem_size( + self, + a1: torch.Tensor, + w1: torch.Tensor, + w2: torch.Tensor, + topk_ids: torch.Tensor, + ) -> tuple[int, int, int, int, int]: + assert w1.dim() == 3 and w2.dim() == 3 + + E = w1.size(0) + K = a1.size(-1) + N = marlin_moe_intermediate_size(w1, w2) + + if a1.dim() == 2: + # Make sure we are using the correct a1 (pre-permute). + assert topk_ids.size(0) == a1.size(0), f"{topk_ids.size(0)} != {a1.size(0)}" + M = a1.size(0) + else: + assert a1.dim() == 3 + assert a1.size(0) == E, f"{a1.size(0)} == {E}" + M = a1.size(1) # This is max_num_tokens + + assert topk_ids.dim() == 2 + topk = topk_ids.size(1) + + return E, M, N, K, topk + + +class MarlinExperts(LoRAExpertsMixin, MarlinExpertsBase): + """Marlin-based fused MoE expert implementation.""" + + def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce: + return TopKWeightAndReduceNoOP() + + @staticmethod + def activation_format() -> mk.FusedMoEActivationFormat: + return mk.FusedMoEActivationFormat.Standard + + def workspace_shapes( + self, + M: int, + N: int, + K: int, + topk: int, + global_num_experts: int, + local_num_experts: int, + expert_tokens_meta: mk.ExpertTokensMetadata | None, + activation: MoEActivation, + ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]: + # Modular Kernel provisions output buffer from workspace1. However in + # the fused_marlin_moe() function, the final torch.sum(), is defined + # essentially as, + # `torch.sum(workspace1, dim=1, out=output)` + # Having overlapping input and output tensors for torch.sum seems + # error prone and depends on how the torch.sum is implemented. + # For this reason we swap let the output buffer provision from + # workspace2. + + # Workspace/IntermediateCache allocation matching fused_marlin_moe() + # workspace1 = (M * topk * max(2 * N, K),) + # workspace2 = (M * topk, N) + + # Workspace/IntermediateCache allocation accounting for output buffer + # provisioning + workspace1 = (M * topk, max(N, K)) + workspace2 = (M * topk * max(2 * N, K),) + output = (M, K) + + return (workspace1, workspace2, output) + + def apply( + self, + output: torch.Tensor, + hidden_states: torch.Tensor, + w1: torch.Tensor, + w2: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + activation: MoEActivation, + global_num_experts: int, + expert_map: torch.Tensor | None, + a1q_scale: torch.Tensor | None, + a2_scale: torch.Tensor | None, + workspace13: torch.Tensor, + workspace2: torch.Tensor, + expert_tokens_meta: mk.ExpertTokensMetadata | None, + apply_router_weight_on_input: bool, + ): + assert self.w1_scale is not None + assert self.w2_scale is not None + + ctx = self._lora_context + if ctx is None: + fused_marlin_moe( + hidden_states=hidden_states, + w1=w1, + w2=w2, + bias1=self.w1_bias, + bias2=self.w2_bias, + w1_scale=self.w1_scale, + w2_scale=self.w2_scale, + topk_weights=topk_weights, + topk_ids=topk_ids, + global_scale1=self.g1_alphas, + global_scale2=self.g2_alphas, + input_global_scale1=self.a1_gscale, + input_global_scale2=self.a2_gscale, + w1_zeros=self.w1_zp, + w2_zeros=self.w2_zp, + quant_type_id=self.quant_type_id, + apply_router_weight_on_input=apply_router_weight_on_input, + global_num_experts=global_num_experts, + activation=activation, + activation_func=self.activation, + moe_sum=self.moe_sum, + expert_map=expert_map, + output=output, + # Workspaces are swapped in workspace_shapes() to account for proper + # output buffer allocation. Please refer to workspace_shapes(). + intermediate_cache13=workspace2, + intermediate_cache2=workspace13, + g_idx1=self.w13_g_idx, + g_idx2=self.w2_g_idx, + sort_indices1=self.w13_g_idx_sort_indices, + sort_indices2=self.w2_g_idx_sort_indices, + is_k_full=self.is_k_full, + input_dtype=self.input_dtype, + clamp_limit=self.gemm1_clamp_limit, + gemm1_alpha=self.gemm1_alpha, + gemm1_beta=self.gemm1_beta, + ) + return + + # LoRA path: wrap activation_func and moe_sum to inject LoRA at the + # two natural injection points. + # + # Marlin uses moe_align_block_size (same as TritonExperts) so + # intermediate_cache1 is indexed by flat (token, expert) pair index, + # which is compatible with add_lora_fused_moe's scatter mechanism. + + M = hidden_states.size(0) + top_k_num = topk_ids.size(1) + lora_state: dict = {} + + def activation_with_lora( + act_enum: MoEActivation, + act_output: torch.Tensor, + act_input: torch.Tensor, + *, + clamp_limit: float | None = None, + alpha: float = 1.0, + beta: float = 0.0, + topk_ids: torch.Tensor | None = None, + expert_map: torch.Tensor | None = None, + ) -> None: + # act_input = intermediate_cache1 (M*topk, 2N for gated) + # act_output = intermediate_cache2 (M*topk, N) + + ( + sorted_token_ids_lora, + expert_ids_lora, + num_tokens_post_padded_lora, + token_lora_mapping, + ) = self.apply_w13_lora( + ctx, + y=act_input, + x=hidden_states, + topk_ids=topk_ids, + topk_weights=topk_weights, + expert_map=expert_map, + w1=w1, + w2=w2, + num_tokens=M, + top_k_num=top_k_num, + ) + lora_state.update( + { + "sorted": sorted_token_ids_lora, + "eids": expert_ids_lora, + "npad": num_tokens_post_padded_lora, + "tlm": token_lora_mapping, + } + ) + self.activation( + act_enum, + act_output, + act_input, + clamp_limit=clamp_limit, + alpha=alpha, + beta=beta, + topk_ids=topk_ids, + expert_map=expert_map, + ) + lora_state["cache2"] = act_output + + def moe_sum_with_lora( + moe_out: torch.Tensor, + out: torch.Tensor, + topk_ids: torch.Tensor, + expert_map: torch.Tensor | None, + ) -> None: + # moe_out shape: (M, topk, K) + self.apply_w2_lora( + ctx, + y=moe_out, + x=lora_state["cache2"], + topk_weights=topk_weights, + sorted_token_ids_lora=lora_state["sorted"], + expert_ids_lora=lora_state["eids"], + num_tokens_post_padded_lora=lora_state["npad"], + token_lora_mapping=lora_state["tlm"], + num_tokens=M, + w1=w1, + w2=w2, + top_k_num=top_k_num, + ) + self.moe_sum(moe_out, out, topk_ids, expert_map) + + return fused_marlin_moe( + hidden_states=hidden_states, + w1=w1, + w2=w2, + bias1=self.w1_bias, + bias2=self.w2_bias, + w1_scale=self.w1_scale, + w2_scale=self.w2_scale, + topk_weights=topk_weights, + topk_ids=topk_ids, + global_scale1=self.g1_alphas, + global_scale2=self.g2_alphas, + input_global_scale1=self.a1_gscale, + input_global_scale2=self.a2_gscale, + w1_zeros=self.w1_zp, + w2_zeros=self.w2_zp, + quant_type_id=self.quant_type_id, + apply_router_weight_on_input=apply_router_weight_on_input, + global_num_experts=global_num_experts, + activation=activation, + activation_func=activation_with_lora, + moe_sum=moe_sum_with_lora, + expert_map=expert_map, + output=output, + intermediate_cache13=workspace2, + intermediate_cache2=workspace13, + g_idx1=self.w13_g_idx, + g_idx2=self.w2_g_idx, + sort_indices1=self.w13_g_idx_sort_indices, + sort_indices2=self.w2_g_idx_sort_indices, + is_k_full=self.is_k_full, + input_dtype=self.input_dtype, + clamp_limit=self.gemm1_clamp_limit, + gemm1_alpha=self.gemm1_alpha, + gemm1_beta=self.gemm1_beta, + ) + + def moe_sum( + self, + input: torch.Tensor, + output: torch.Tensor, + topk_ids: torch.Tensor, + expert_map: torch.Tensor | None, + ) -> None: + if expert_map is not None: + ops.moe_sum(input, output, topk_ids, expert_map) + else: + ops.moe_sum(input, output) + + +class BatchedMarlinExperts(MarlinExpertsBase): + """Batched Marlin-based fused MoE expert implementation.""" + + def __init__( + self, + moe_config: FusedMoEConfig, + quant_config: FusedMoEQuantConfig, + max_num_tokens: int, + num_dispatchers: int, + w13_g_idx: torch.Tensor | None = None, + w2_g_idx: torch.Tensor | None = None, + w13_g_idx_sort_indices: torch.Tensor | None = None, + w2_g_idx_sort_indices: torch.Tensor | None = None, + is_k_full: bool = True, + ): + super().__init__( + moe_config=moe_config, + quant_config=quant_config, + max_num_tokens=max_num_tokens, + num_dispatchers=num_dispatchers, + w13_g_idx=w13_g_idx, + w2_g_idx=w2_g_idx, + w13_g_idx_sort_indices=w13_g_idx_sort_indices, + w2_g_idx_sort_indices=w2_g_idx_sort_indices, + is_k_full=is_k_full, + ) + + def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce: + return TopKWeightAndReduceDelegate() + + @staticmethod + def activation_format() -> mk.FusedMoEActivationFormat: + return mk.FusedMoEActivationFormat.BatchedExperts + + def workspace_shapes( + self, + M: int, + N: int, + K: int, + topk: int, + global_num_experts: int, + local_num_experts: int, + expert_tokens_meta: mk.ExpertTokensMetadata | None, + activation: MoEActivation, + ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]: + assert self.num_dispatchers is not None + assert self.max_num_tokens is not None + num_dispatchers = self.num_dispatchers + num_experts = local_num_experts + max_num_tokens = self.max_num_tokens + workspace13 = (num_experts * max_num_tokens * num_dispatchers, max(K, N * 2)) + workspace2 = (num_experts * max_num_tokens * num_dispatchers, N) + output = (num_experts, max_num_tokens * num_dispatchers, K) + return (workspace13, workspace2, output) + + def apply( + self, + output: torch.Tensor, + hidden_states: torch.Tensor, + w1: torch.Tensor, + w2: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + activation: MoEActivation, + global_num_experts: int, + expert_map: torch.Tensor | None, + a1q_scale: torch.Tensor | None, + a2_scale: torch.Tensor | None, + workspace13: torch.Tensor, + workspace2: torch.Tensor, + expert_tokens_meta: mk.ExpertTokensMetadata | None, + apply_router_weight_on_input: bool, + ): + assert expert_tokens_meta is not None, "Num valid tokens per batch is required" + return batched_fused_marlin_moe( + hidden_states=hidden_states, + expert_num_tokens=expert_tokens_meta.expert_num_tokens, + w1=w1, + w2=w2, + bias1=self.w1_bias, + bias2=self.w2_bias, + w1_scale=self.w1_scale, + w2_scale=self.w2_scale, + quant_type_id=self.quant_type_id, + apply_router_weight_on_input=apply_router_weight_on_input, + global_num_experts=global_num_experts, + activation=activation, + expert_map=expert_map, + input_global_scale1=self.a1_gscale, + input_global_scale2=self.a2_gscale, + output=output, + intermediate_cache13=workspace13, + intermediate_cache2=workspace2, + g_idx1=self.w13_g_idx, + g_idx2=self.w2_g_idx, + sort_indices1=self.w13_g_idx_sort_indices, + sort_indices2=self.w2_g_idx_sort_indices, + w1_zeros=self.w1_zp, + w2_zeros=self.w2_zp, + input_dtype=self.input_dtype, + is_k_full=self.is_k_full, + clamp_limit=self.gemm1_clamp_limit, + gemm1_alpha=self.gemm1_alpha, + gemm1_beta=self.gemm1_beta, + ) 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_mxfp4_moe.py.orig b/aphrodite/model_executor/layers/fused_moe/experts/trtllm_mxfp4_moe.py.orig new file mode 100644 index 0000000000..903c266d78 --- /dev/null +++ b/aphrodite/model_executor/layers/fused_moe/experts/trtllm_mxfp4_moe.py.orig @@ -0,0 +1,342 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import torch + +import aphrodite.model_executor.layers.fused_moe.modular_kernel as mk +from aphrodite.model_executor.layers.fused_moe.activation import MoEActivation +from aphrodite.model_executor.layers.fused_moe.config import ( + FusedMoEConfig, + FusedMoEParallelConfig, + FusedMoEQuantConfig, + RoutingMethodType, +) +from aphrodite.model_executor.layers.fused_moe.topk_weight_and_reduce import ( + TopKWeightAndReduceNoOP, +) +from aphrodite.model_executor.layers.fused_moe.utils import trtllm_moe_pack_topk_ids_weights +from aphrodite.model_executor.layers.quantization.utils.quant_utils import ( + QuantKey, + kMxfp4Static, + kMxfp8Dynamic, +) +from aphrodite.platforms import current_platform +from aphrodite.utils.flashinfer import has_flashinfer + + +class TrtLlmMxfp4ExpertsBase: + """ + MXFP4 TRTLLM-Gen MoE kernels. Shared base for modular and monolithic. + """ + + def __init__( + self, + moe_config: FusedMoEConfig, + quant_config: FusedMoEQuantConfig, + **kwargs, + ): + self.moe_config = moe_config + self.quant_config = quant_config + + self.routing_method_type = moe_config.routing_method + self.topk = moe_config.experts_per_token + self.intermediate_size_per_partition = moe_config.intermediate_size_per_partition + self.hidden_dim = moe_config.hidden_dim + self.hidden_dim_unpadded = moe_config.hidden_dim_unpadded or moe_config.hidden_dim + self.local_num_experts = moe_config.num_local_experts + self.ep_rank = moe_config.moe_parallel_config.ep_rank + + # MXFP4-specific TRTLLM parameters from quant_config + device = torch.accelerator.current_device_index() + if quant_config.gemm1_alpha is not None: + self.gemm1_alpha = torch.tensor( + [quant_config.gemm1_alpha] * self.local_num_experts, + dtype=torch.float32, + device=device, + ) + else: + self.gemm1_alpha = None + + if quant_config.gemm1_beta is not None: + self.gemm1_beta = torch.tensor( + [quant_config.gemm1_beta] * self.local_num_experts, + dtype=torch.float32, + device=device, + ) + else: + self.gemm1_beta = None + + if quant_config.gemm1_clamp_limit is not None: + self.gemm1_clamp_limit = torch.tensor( + [quant_config.gemm1_clamp_limit] * self.local_num_experts, + dtype=torch.float32, + device=device, + ) + else: + self.gemm1_clamp_limit = None + + self.max_capture_size = moe_config.max_capture_size + + @staticmethod + def _supports_current_device() -> bool: + p = current_platform + return p.is_cuda() and p.is_device_capability_family(100) and has_flashinfer() + + @staticmethod + def _supports_no_act_and_mul() -> bool: + return False + + @staticmethod + def _supports_quant_scheme( + weight_key: QuantKey | None, + activation_key: QuantKey | None, + ) -> bool: + SUPPORTED_W_A = [ + (kMxfp4Static, None), + (kMxfp4Static, kMxfp8Dynamic), + ] + return (weight_key, activation_key) in SUPPORTED_W_A + + @staticmethod + def _supports_activation(activation: MoEActivation) -> bool: + return activation in (MoEActivation.SWIGLUOAI, MoEActivation.SILU) + + @staticmethod + def activation_format() -> mk.FusedMoEActivationFormat: + return mk.FusedMoEActivationFormat.Standard + + @property + def expects_unquantized_inputs(self) -> bool: + return False + + +class TrtLlmMxfp4ExpertsMonolithic(TrtLlmMxfp4ExpertsBase, mk.FusedMoEExpertsMonolithic): + """ + Monolithic version of the MXFP4 TRTLLM kernel (router + experts). + Wraps flashinfer.trtllm_fp4_block_scale_moe(). + """ + + def supports_routing_replay_capture(self) -> bool: + return True + + @staticmethod + def _supports_parallel_config( + moe_parallel_config: FusedMoEParallelConfig, + ) -> bool: + return ( + not moe_parallel_config.use_all2all_kernels + and not moe_parallel_config.enable_eplb + and moe_parallel_config.dp_size <= 1 + ) + + @staticmethod + def _supports_routing_method( + routing_method: RoutingMethodType, + weight_key: QuantKey | None, + activation_key: QuantKey | None, + ) -> bool: + return routing_method in [ + RoutingMethodType.Renormalize, + RoutingMethodType.RenormalizeNaive, + ] + + @staticmethod + def _supports_router_logits_dtype( + router_logits_dtype: torch.dtype | None, + routing_method: RoutingMethodType, + ) -> bool: + # Kernel converts to bfloat16 internally + return True + + def apply( + self, + hidden_states: torch.Tensor, + w1: torch.Tensor, + w2: torch.Tensor, + router_logits: torch.Tensor, + activation: MoEActivation, + global_num_experts: int, + expert_map: torch.Tensor | None, + a1q_scale: torch.Tensor | None, + apply_router_weight_on_input: bool, + # grouped topk + fused topk bias parameters + num_expert_group: int | None = None, + e_score_correction_bias: torch.Tensor | None = None, + routed_scaling_factor: float | None = None, + topk_group: int | None = None, + ) -> torch.Tensor: + from flashinfer import trtllm_fp4_block_scale_moe + + if a1q_scale is not None: + x_quant = hidden_states + x_scale = a1q_scale.view(torch.float8_e4m3fn) + else: + assert hidden_states.dtype == torch.bfloat16 + x_quant = hidden_states + x_scale = None + output = torch.empty( + *hidden_states.shape[:-1], + self.hidden_dim_unpadded, + dtype=torch.bfloat16, + device=hidden_states.device, + ) + + routing_replay_out = self._maybe_make_routing_replay_buffer( + num_tokens=hidden_states.shape[0], + device=hidden_states.device, + ) + trtllm_fp4_block_scale_moe( + routing_logits=router_logits.to(torch.bfloat16), + routing_bias=None, + hidden_states=x_quant, + hidden_states_scale=x_scale, + gemm1_weights=w1, + gemm1_weights_scale=self.w1_scale, + gemm1_bias=self.w1_bias, + gemm1_alpha=self.gemm1_alpha, + gemm1_beta=self.gemm1_beta, + gemm1_clamp_limit=self.gemm1_clamp_limit, + gemm2_weights=w2, + gemm2_weights_scale=self.w2_scale, + gemm2_bias=self.w2_bias, + output1_scale_scalar=None, + output1_scale_gate_scalar=None, + output2_scale_scalar=None, + num_experts=global_num_experts, + top_k=self.topk, + n_group=None, + topk_group=None, + 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, + routing_method_type=self.routing_method_type, + do_finalize=True, + tune_max_num_tokens=max(self.max_capture_size, 1), + output=output, + routing_replay_out=routing_replay_out, + ) + self._maybe_dispatch_routing_replay(routing_replay_out, num_tokens=hidden_states.shape[0]) + + return output + + +class TrtLlmMxfp4ExpertsModular(TrtLlmMxfp4ExpertsBase, mk.FusedMoEExpertsModular): + """ + Modular version of the MXFP4 TRTLLM kernel (just the experts). + Wraps flashinfer.trtllm_fp4_block_scale_routed_moe(). + Moved from trtllm_moe.py. + """ + + @staticmethod + def _supports_parallel_config( + moe_parallel_config: FusedMoEParallelConfig, + ) -> bool: + return True + + @staticmethod + def _supports_routing_method( + routing_method: RoutingMethodType, + weight_key: QuantKey | None, + activation_key: QuantKey | None, + ) -> bool: + # Modular kernel handles only the expert computation; + # routing is done externally, so accept any routing method. + return True + + def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce: + return TopKWeightAndReduceNoOP() + + def workspace_shapes( + self, + M: int, + N: int, + K: int, + topk: int, + global_num_experts: int, + local_num_experts: int, + expert_tokens_meta: mk.ExpertTokensMetadata | None, + activation: MoEActivation, + ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]: + # The workspaces for this implementation are managed by flashinfer. + workspace1 = (0,) + workspace2 = (0,) + output = (M, self.hidden_dim_unpadded) + return (workspace1, workspace2, output) + + def apply( + self, + output: torch.Tensor, + hidden_states: torch.Tensor, + w1: torch.Tensor, + w2: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + activation: MoEActivation, + global_num_experts: int, + expert_map: torch.Tensor | None, + a1q_scale: torch.Tensor | None, + a2_scale: torch.Tensor | None, + workspace13: torch.Tensor, + workspace2: torch.Tensor, + expert_tokens_meta: mk.ExpertTokensMetadata | None, + apply_router_weight_on_input: bool, + ): + topk = topk_ids.size(-1) + local_num_experts = w1.size(0) + intermediate_size = self.intermediate_size_per_partition + local_expert_offset = self.moe_config.ep_rank * local_num_experts + + if a1q_scale is not None: + x_quant = hidden_states + x_scale = a1q_scale.view(torch.float8_e4m3fn) + else: + assert hidden_states.dtype == torch.bfloat16 + x_quant = hidden_states + x_scale = None + + # Pack topk ids and weights into format expected by the kernel. + packed_tensor = trtllm_moe_pack_topk_ids_weights(topk_ids, topk_weights) + + assert self.w1_scale is not None + assert self.w2_scale is not None + kwargs = { + "topk_ids": packed_tensor, + "routing_bias": None, + "hidden_states": x_quant, + "hidden_states_scale": x_scale, + "gemm1_weights": w1, + "gemm1_weights_scale": self.w1_scale, + "gemm1_bias": self.w1_bias, + "gemm1_alpha": self.gemm1_alpha, + "gemm1_beta": self.gemm1_beta, + "gemm1_clamp_limit": self.gemm1_clamp_limit, + "gemm2_weights": w2, + "gemm2_weights_scale": self.w2_scale, + "gemm2_bias": self.w2_bias, + "output1_scale_scalar": None, + "output1_scale_gate_scalar": None, + "output2_scale_scalar": None, + "num_experts": global_num_experts, + "top_k": topk, + "n_group": None, + "topk_group": None, + "intermediate_size": intermediate_size, + "local_expert_offset": local_expert_offset, + "local_num_experts": local_num_experts, + "routed_scaling_factor": None, + # Modular kernel receives pre-routed tokens, so routing + # is already done. Use Renormalize as a safe default that + # the TRTLLM C++ kernel supports. + "routing_method_type": RoutingMethodType.Renormalize, + "do_finalize": True, + "enable_pdl": True, + "output": output, + "tune_max_num_tokens": max(self.max_capture_size, 1), + } + + from flashinfer import trtllm_fp4_block_scale_routed_moe + + trtllm_fp4_block_scale_routed_moe(**kwargs) + + return output 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/experts/trtllm_nvfp4_moe.py.orig b/aphrodite/model_executor/layers/fused_moe/experts/trtllm_nvfp4_moe.py.orig new file mode 100644 index 0000000000..1c87a1439a --- /dev/null +++ b/aphrodite/model_executor/layers/fused_moe/experts/trtllm_nvfp4_moe.py.orig @@ -0,0 +1,526 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + + +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.config import ( + FusedMoEConfig, + FusedMoEParallelConfig, + FusedMoEQuantConfig, + RoutingMethodType, +) +from aphrodite.model_executor.layers.fused_moe.topk_weight_and_reduce import ( + TopKWeightAndReduceNoOP, +) +from aphrodite.model_executor.layers.fused_moe.utils import ( + fi_moe_largest_bucket, + trtllm_moe_pack_topk_ids_weights, +) +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, + kNvfp4Static, +) +from aphrodite.platforms import current_platform +from aphrodite.utils.flashinfer import has_flashinfer_trtllm_fused_moe + +logger = init_logger(__name__) + +# Base scale for per-token NVFP4 activation quant; the kernel folds the +# per-token global scale (from the activation amax) on top of it. +_PER_TOKEN_BASE_GLOBAL_SCALE = 1.0 / (448.0 * 6.0) + + +class TrtLlmNvFp4ExpertsBase: + """ + NvFp4 TRTLLM-Gen MoE kernels. Supports modular and monolithic interface. + """ + + def __init__( + self, + moe_config: FusedMoEConfig, + quant_config: FusedMoEQuantConfig, + per_token_activation: bool = False, + ): + self.moe_config = moe_config + self.quant_config = quant_config + # Quantize the input here (deferred from prepare) to capture a per-token + # global scale, instead of a static one. + self.per_token_activation = per_token_activation + + self.routing_method_type = self.moe_config.routing_method + self.topk = moe_config.experts_per_token + self.intermediate_size_per_partition = moe_config.intermediate_size_per_partition + self.hidden_dim = moe_config.hidden_dim + self.hidden_dim_unpadded = moe_config.hidden_dim_unpadded or moe_config.hidden_dim + self.local_num_experts = moe_config.num_local_experts + self.ep_rank = moe_config.moe_parallel_config.ep_rank + + assert self.quant_config.g1_alphas is not None + assert self.quant_config.a2_gscale is not None + if moe_config.is_act_and_mul: + # g1_alpha_s = a13_scale * w13_scale_2 + # a2_gscale = (1 / a2_scale) + # g1_scale_c = a13_scale * w13_scale_2 / a2_scale + self.g1_scale_c = self.quant_config.g1_alphas * self.quant_config.a2_gscale + else: + self.g1_scale_c = self.quant_config.a2_gscale.clone() + + # Fall back to moe_config.swiglu_* when quant_config doesn't carry them + # (ModelOpt NVFP4 checkpoints store these on moe_config, not quant_config). + device = torch.accelerator.current_device_index() + + def _per_expert(val: float | None) -> torch.Tensor | None: + if val is None: + return None + return torch.full( + (self.local_num_experts,), + float(val), + dtype=torch.float32, + device=device, + ) + + clamp = quant_config.gemm1_clamp_limit + if clamp is None: + clamp = getattr(moe_config, "swiglu_limit", None) + alpha = quant_config.gemm1_alpha + if alpha is None: + alpha = getattr(moe_config, "swiglu_alpha", None) + beta = quant_config.gemm1_beta + if beta is None: + beta = getattr(moe_config, "swiglu_beta", None) + + if moe_config.is_act_and_mul: + self.gemm1_clamp_limit = _per_expert(clamp) + self.gemm1_alpha = _per_expert(alpha) + self.gemm1_beta = _per_expert(beta) + else: + self.gemm1_clamp_limit = None + self.gemm1_alpha = None + self.gemm1_beta = None + + logger.debug_once( + "activation=%s, gemm1_alpha=%s, gemm1_beta=%s, gemm1_clamp_limit=%s", + moe_config.activation, + alpha, + beta, + clamp, + ) + + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + layer.w13_weight_scale_2.data.mul_(layer.w13_input_scale) + layer.w2_weight_scale_2.data.mul_(layer.w2_input_scale) + # Recompute g1_scale_c since g1_alphas was just fused in-place. + # Register as a layer parameter so EPLB rearranges it alongside + # other expert weights. + assert self.quant_config.g1_alphas is not None + assert self.quant_config.a2_gscale is not None + if self.moe_config.is_act_and_mul: + g1_scale_c = self.quant_config.g1_alphas * self.quant_config.a2_gscale + else: + g1_scale_c = self.quant_config.a2_gscale.clone() + layer.register_parameter( + "g1_scale_c", + torch.nn.Parameter(g1_scale_c, requires_grad=False), + ) + self.g1_scale_c = layer.g1_scale_c + + # Pre-fold the per-expert g1_alphas (= output1_scale_gate_scalar) + # division so the TRTLLM kernel receives the raw-GEMM-space clamp + # directly. g1_alphas is set once here in process_weights_after_loading + # (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: + gemm1_clamp_limit = self.gemm1_clamp_limit / self.quant_config.g1_alphas + layer.register_parameter( + "gemm1_clamp_limit", + torch.nn.Parameter(gemm1_clamp_limit, requires_grad=False), + ) + self.gemm1_clamp_limit = layer.gemm1_clamp_limit + + # beta shifts the raw GEMM1 accumulator, so fold by g1_alphas like the + # clamp limit. alpha is applied to the dequantized gate, so it stays + # 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 + layer.register_parameter( + "gemm1_beta", + torch.nn.Parameter(gemm1_beta, requires_grad=False), + ) + self.gemm1_beta = layer.gemm1_beta + + if self.gemm1_alpha is not None: + layer.register_parameter( + "gemm1_alpha", + torch.nn.Parameter(self.gemm1_alpha, requires_grad=False), + ) + self.gemm1_alpha = layer.gemm1_alpha + + @staticmethod + def _supports_current_device() -> bool: + """Supports only Blackwell-family GPUs.""" + p = current_platform + return p.is_cuda() and p.is_device_capability_family(100) and has_flashinfer_trtllm_fused_moe() + + @staticmethod + def _supports_no_act_and_mul() -> bool: + """Supports non-gated MoE (i.e. Nemotron-Nano).""" + return True + + @staticmethod + def _supports_quant_scheme( + weight_key: QuantKey | None, + activation_key: QuantKey | None, + ) -> bool: + """Supports Nvfp4 quantization.""" + SUPPORTED_W_A = [ + (kNvfp4Static, kNvfp4Dynamic), + ] + return (weight_key, activation_key) in SUPPORTED_W_A + + @staticmethod + def _supports_activation(activation: MoEActivation) -> bool: + """Supports SiLU, RELU^2 non-gated, GELU, and clamped SwiGLU-OAI.""" + return activation in [ + MoEActivation.SILU, + MoEActivation.RELU2_NO_MUL, + MoEActivation.GELU, + MoEActivation.GELU_TANH, + MoEActivation.SWIGLUOAI_UNINTERLEAVE, + ] + + @staticmethod + def _supports_shape(hidden_dim: int) -> bool: + # Weights are zero-padded to 256-alignment at load time and the MoE + # runner pads activations via _maybe_pad_hidden_states, so any + # hidden_dim is accepted. + # NOTE: non-256-aligned dims will trigger a warning log and may + # cause performance degradation due to activation slicing. + return True + + @staticmethod + def activation_format() -> mk.FusedMoEActivationFormat: + return mk.FusedMoEActivationFormat.Standard + + @property + def expects_unquantized_inputs(self) -> bool: + return self.per_token_activation + + def _quantize_per_token_input( + self, + hidden_states: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """NVFP4-quantize activations with a per-token global scale. + + Returns ``(packed_fp4, block_scale, per_token_scale)``. + """ + from flashinfer import SfLayout, nvfp4_quantize + + hs_fp4, hs_block_scale, per_token_scale = nvfp4_quantize( + hidden_states, + _PER_TOKEN_BASE_GLOBAL_SCALE, + sfLayout=SfLayout.layout_linear, + per_token_activation=True, + ) + return hs_fp4, hs_block_scale, per_token_scale + + def _get_chunk_size(self) -> int: + MAX_GRID_Y = 65535 + MAX_TILE_TOKENS_DIM = 128 + + def _calc_max_supported_tokens(top_k: int, num_experts: int) -> int: + """Calculates the max number of supported tokens, so the CUDA grid.Y limit + won't be reached. + Based on getMaxNumCtasInBatchDim function in flashinfer's TRTLLM MoE runner: + https://github.com/flashinfer-ai/flashinfer/blob/719ee23fd82cb220d51ad118ca60198718f6c9d1/include/flashinfer/trtllm/fused_moe/runner.h#L97 + Which given numTokens, topK, numExperts, tileTokensDim calculates maxNumCtas + which is used as the CUDA grid.Y dimension, which we want to + be <= MAX_GRID_Y. Solving for numTokens gives the formula below. + """ + return (num_experts + (MAX_GRID_Y - num_experts + 1) * MAX_TILE_TOKENS_DIM - 1) // top_k + + # Using 305k or more causes IMA error in the kernel, so limit to 300k. + return min(300000, _calc_max_supported_tokens(self.topk, self.moe_config.num_experts)) + + +class TrtLlmNvFp4ExpertsModular(TrtLlmNvFp4ExpertsBase, mk.FusedMoEExpertsModular): + """ + Modular version of the implementation (just the experts). + """ + + @staticmethod + def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bool: + """The modular implementation supports all parallel configs.""" + return True + + def workspace_shapes( + self, + M: int, + N: int, + K: int, + topk: int, + global_num_experts: int, + local_num_experts: int, + expert_tokens_meta: mk.ExpertTokensMetadata | None, + activation: MoEActivation, + ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]: + if self.per_token_activation: + # Deferred input quant leaves K unpacked here, breaking the + # workspace assumptions below. Per-token NVFP4 is only supported on + # the monolithic (non-EP) path for now. + raise NotImplementedError( + "NVFP4 per-token activation is only supported on the monolithic (non-EP) FlashInfer TRTLLM MoE path." + ) + + # The workspaces for this implementation are managed by flashinfer. + workspace1 = (0,) + workspace2 = (0,) + + # Hidden states are Nvfp4, packed into int8 dtype, so we + # need to multiply K by 2 to get the output shape right. + assert self.hidden_dim == K * 2 + output = (M, self.hidden_dim) + + return (workspace1, workspace2, output) + + def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce: + return TopKWeightAndReduceNoOP() + + def _invoke_kernel( + self, + output: torch.Tensor, + hidden_states: torch.Tensor, + w1: torch.Tensor, + w2: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + activation: MoEActivation, + global_num_experts: int, + a1q_scale: torch.Tensor, + ): + import flashinfer + + assert self.quant_config.w1_scale is not None + assert self.quant_config.w2_scale is not None + + # Per-token: input is unquantized, quantize it here. Otherwise it was + # already quantized in prepare() with the static global scale. + if self.per_token_activation: + hidden_states, block_scale, per_token_scale = self._quantize_per_token_input(hidden_states) + else: + block_scale, per_token_scale = a1q_scale, None + + # Pack topk ids and weights into format expected by the kernel. + packed_tensor = trtllm_moe_pack_topk_ids_weights(topk_ids, topk_weights) + output1_scale_gate_scalar = self.quant_config.g1_alphas + + # Invoke kernel. + flashinfer.fused_moe.trtllm_fp4_block_scale_routed_moe( + topk_ids=packed_tensor, + routing_bias=None, + hidden_states=hidden_states, + hidden_states_scale=block_scale.view(torch.float8_e4m3fn).reshape(*hidden_states.shape[:-1], -1), + gemm1_weights=w1, + gemm1_weights_scale=self.quant_config.w1_scale.view(torch.float8_e4m3fn), + gemm1_bias=None, + gemm1_alpha=self.gemm1_alpha, + gemm1_beta=self.gemm1_beta, + gemm1_clamp_limit=self.gemm1_clamp_limit, + gemm2_weights=w2, + gemm2_weights_scale=self.quant_config.w2_scale.view(torch.float8_e4m3fn), + gemm2_bias=None, + output1_scale_scalar=self.g1_scale_c, + output1_scale_gate_scalar=output1_scale_gate_scalar, + output2_scale_scalar=self.quant_config.g2_alphas, + num_experts=global_num_experts, + top_k=self.topk, + n_group=0, + topk_group=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, + routing_method_type=1, # not used + do_finalize=True, + activation_type=activation_to_flashinfer_int(activation), + per_token_scale=per_token_scale, + output=output, + tune_max_num_tokens=min(fi_moe_largest_bucket(self.moe_config), self._get_chunk_size()), + ) + + def apply( + self, + output: torch.Tensor, + hidden_states: torch.Tensor, + w1: torch.Tensor, + w2: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + activation: MoEActivation, + global_num_experts: int, + expert_map: torch.Tensor | None, + a1q_scale: torch.Tensor | None, + a2_scale: torch.Tensor | None, + workspace13: torch.Tensor, + workspace2: torch.Tensor, + expert_tokens_meta: mk.ExpertTokensMetadata | None, + apply_router_weight_on_input: bool, + ): + assert self._supports_activation(activation) + # Per-token defers input quant to _invoke_kernel, so a1q_scale is None. + assert a1q_scale is not None or self.per_token_activation + + M = hidden_states.shape[0] + chunk_size = self._get_chunk_size() + + if chunk_size >= M: + self._invoke_kernel( + output, + hidden_states, + w1, + w2, + topk_weights, + topk_ids, + activation, + global_num_experts, + a1q_scale, + ) + else: + for start in range(0, M, chunk_size): + end = min(start + chunk_size, M) + self._invoke_kernel( + output[start:end], + hidden_states[start:end], + w1, + w2, + topk_weights[start:end], + topk_ids[start:end], + activation, + global_num_experts, + None if a1q_scale is None else a1q_scale[start:end], + ) + + +class TrtLlmNvFp4ExpertsMonolithic(TrtLlmNvFp4ExpertsBase, mk.FusedMoEExpertsMonolithic): + """ + Monolithic version of the kernel (router + experts). + """ + + def supports_routing_replay_capture(self) -> bool: + return True + + @staticmethod + def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bool: + """The modular implementation should be used for the Dp/Ep or EPLB case.""" + return not moe_parallel_config.use_all2all_kernels and not moe_parallel_config.enable_eplb + + @staticmethod + def _supports_routing_method( + routing_method_type: RoutingMethodType, + weight_key: QuantKey | None, + activation_key: QuantKey | None, + ) -> bool: + # NOTE(rob): this is a conservative list. + return routing_method_type in [ + RoutingMethodType.DeepSeekV3, + RoutingMethodType.Renormalize, + RoutingMethodType.RenormalizeNaive, + RoutingMethodType.Llama4, + RoutingMethodType.SigmoidRenorm, + RoutingMethodType.Sigmoid, + RoutingMethodType.MiniMax2, + RoutingMethodType.Simulated, + ] + + @staticmethod + def _supports_router_logits_dtype( + router_logits_dtype: torch.dtype | None, + routing_method: RoutingMethodType, + ) -> bool: + return router_logits_dtype in [torch.bfloat16, torch.float32] + + def apply( + self, + hidden_states: torch.Tensor, + w1: torch.Tensor, + w2: torch.Tensor, + router_logits: torch.Tensor, + activation: MoEActivation, + global_num_experts: int, + expert_map: torch.Tensor | None, + a1q_scale: torch.Tensor | None, + apply_router_weight_on_input: bool, + # grouped topk + fused topk bias parameters + num_expert_group: int | None = None, + e_score_correction_bias: torch.Tensor | None = None, + routed_scaling_factor: float | None = None, + topk_group: int | None = None, + ) -> torch.Tensor: + import flashinfer + + assert self._supports_activation(activation) + assert a1q_scale is not None or self.per_token_activation + assert self.quant_config.w1_scale is not None + assert self.quant_config.w2_scale is not None + assert (apply_router_weight_on_input and self.routing_method_type == RoutingMethodType.Llama4) or ( + not apply_router_weight_on_input and self.routing_method_type != RoutingMethodType.Llama4 + ) + + # Per-token: input is unquantized, quantize it here (see modular apply). + if self.per_token_activation: + hidden_states, block_scale, per_token_scale = self._quantize_per_token_input(hidden_states) + else: + block_scale, per_token_scale = a1q_scale, None + + output1_scale_gate_scalar = self.quant_config.g1_alphas + + routing_replay_out = self._maybe_make_routing_replay_buffer( + num_tokens=hidden_states.shape[0], + device=hidden_states.device, + ) + # Invoke kernel. + # NOTE: Activation padding and output + # truncation are handled by the MoE runner's + result = flashinfer.fused_moe.trtllm_fp4_block_scale_moe( + routing_logits=router_logits, + routing_bias=e_score_correction_bias, + hidden_states=hidden_states, + hidden_states_scale=block_scale.view(torch.float8_e4m3fn).reshape(*hidden_states.shape[:-1], -1), + gemm1_weights=w1, + gemm1_weights_scale=self.quant_config.w1_scale.view(torch.float8_e4m3fn), + gemm1_bias=None, + gemm1_alpha=self.gemm1_alpha, + gemm1_beta=self.gemm1_beta, + gemm1_clamp_limit=self.gemm1_clamp_limit, + gemm2_weights=w2, + gemm2_weights_scale=self.quant_config.w2_scale.view(torch.float8_e4m3fn), + gemm2_bias=None, + output1_scale_scalar=self.g1_scale_c, + output1_scale_gate_scalar=output1_scale_gate_scalar, + output2_scale_scalar=self.quant_config.g2_alphas, + num_experts=global_num_experts, + top_k=self.topk, + 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=routed_scaling_factor, + routing_method_type=self.routing_method_type, + do_finalize=True, + activation_type=activation_to_flashinfer_int(activation), + per_token_scale=per_token_scale, + tune_max_num_tokens=fi_moe_largest_bucket(self.moe_config), + routing_replay_out=routing_replay_out, + )[0] + self._maybe_dispatch_routing_replay(routing_replay_out, num_tokens=hidden_states.shape[0]) + return result 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/layer.py.orig b/aphrodite/model_executor/layers/fused_moe/layer.py.orig new file mode 100644 index 0000000000..ae566adb63 --- /dev/null +++ b/aphrodite/model_executor/layers/fused_moe/layer.py.orig @@ -0,0 +1,422 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from collections.abc import Callable +from typing import Any + +import torch + +import aphrodite.envs as envs +from aphrodite._aiter_ops import rocm_aiter_ops +from aphrodite.config import ParallelConfig, get_current_aphrodite_config +from aphrodite.distributed import ( + get_dp_group, + get_pcp_group, + get_tensor_model_parallel_world_size, +) +from aphrodite.distributed.eplb.eplb_state import EplbLayerState +from aphrodite.logger import init_logger +from aphrodite.model_executor.layers.fused_moe.activation import MoEActivation +from aphrodite.model_executor.layers.fused_moe.config import ( + FusedMoEConfig, + FusedMoEParallelConfig, +) +from aphrodite.model_executor.layers.fused_moe.expert_map_manager import ( + ExpertMapManager, +) +from aphrodite.model_executor.layers.fused_moe.routed_experts import RoutedExperts +from aphrodite.model_executor.layers.fused_moe.router.fused_moe_router import ( + FusedMoERouter, +) +from aphrodite.model_executor.layers.fused_moe.router.router_factory import ( + create_fused_moe_router, +) +from aphrodite.model_executor.layers.fused_moe.runner.moe_runner import ( + MoERunner, +) +from aphrodite.model_executor.layers.quantization.base_config import ( + QuantizationConfig, +) + +logger = init_logger(__name__) + + +def make_parallel_config( + tp_size: int | None, + dp_size: int | None, + pcp_size: int | None, + is_sequence_parallel: bool, + parallel_config: ParallelConfig, +) -> FusedMoEParallelConfig: + tp_size_ = tp_size if tp_size is not None else get_tensor_model_parallel_world_size() + dp_size_ = dp_size if dp_size is not None else get_dp_group().world_size + pcp_size_ = pcp_size if pcp_size is not None else get_pcp_group().world_size + sp_size = tp_size_ if is_sequence_parallel else 1 + + moe_parallel_config = FusedMoEParallelConfig.make( + tp_size_=tp_size_, + pcp_size_=pcp_size_, + dp_size_=dp_size_, + sp_size_=sp_size, + aphrodite_parallel_config=parallel_config, + ) + + assert moe_parallel_config.is_sequence_parallel == is_sequence_parallel + + logger.debug("FusedMoEParallelConfig = %s", str(moe_parallel_config)) + + return moe_parallel_config + + +def determine_expert_counts( + num_experts: int, + num_redundant_experts: int, + n_shared_experts: int | None, + is_act_and_mul: bool, +) -> tuple[int, int, int]: + global_num_experts = num_experts + num_redundant_experts + logical_num_experts = num_experts + # Shared-expert fusion: append the shared expert(s) as routed-expert slots + # so they run in the same grouped GEMM. Gated by + # APHRODITE_ROCM_USE_AITER_FUSION_SHARED_EXPERTS: either the native aiter fused-MoE + # path (env + master switch, via is_fusion_moe_shared_experts_enabled) or the + # backend-neutral router-append path (env alone, independent of the master + # switch; e.g. the MM3 triton/flydsl mxfp8 MoE). Gated activations only. + fuse_shared_enabled = ( + rocm_aiter_ops.is_fusion_moe_shared_experts_enabled() or envs.APHRODITE_ROCM_USE_AITER_FUSION_SHARED_EXPERTS + ) and is_act_and_mul + + num_fused_shared_experts = n_shared_experts if n_shared_experts is not None and fuse_shared_enabled else 0 + + return global_num_experts, logical_num_experts, num_fused_shared_experts + + +# TODO: rename this +def FusedMoE( + num_experts: int, # Global number of experts + top_k: int, + hidden_size: int, + intermediate_size: int, + intermediate_pad: int | None = None, + params_dtype: torch.dtype | None = None, + renormalize: bool = True, + use_grouped_topk: bool = False, + num_expert_group: int | None = None, + topk_group: int | None = None, + quant_config: QuantizationConfig | None = None, + tp_size: int | None = None, + dp_size: int | None = None, + pcp_size: int | None = None, + prefix: str = "", + custom_routing_function: Callable | None = None, + router: FusedMoERouter | None = None, + scoring_func: str = "softmax", + routed_scaling_factor: float = 1.0, + swiglu_limit: float | None = None, + swiglu_alpha: float | None = None, + swiglu_beta: float | None = None, + e_score_correction_bias: torch.Tensor | None = None, + apply_router_weight_on_input: bool = False, + activation: str = "silu", + enable_eplb: bool = False, + num_redundant_experts: int = 0, + has_bias: bool = False, + is_sequence_parallel: bool = False, + reduce_results: bool = True, + ckpt_names: tuple[str, str, str] = ("gate_proj", "down_proj", "up_proj"), + n_shared_experts: int | None = None, + router_logits_dtype: torch.dtype | None = None, + gate: torch.nn.Module | None = None, + shared_experts: torch.nn.Module | None = None, + shared_expert_gate: torch.nn.Module | None = None, + routed_input_transform: torch.nn.Module | None = None, + routed_output_transform: torch.nn.Module | None = None, + apply_routed_scale_to_output: bool = False, + zero_expert_type: str | None = None, + hash_indices_table: torch.Tensor | None = None, + runner_cls: type[MoERunner] | None = None, + runner_args: dict[str, Any] | None = None, + routed_experts_cls: type[RoutedExperts] | None = None, + routed_experts_args: dict[str, Any] | None = None, +) -> MoERunner: + """Factory function for creating MoE execution pipeline. + + Creates and configures a complete MoE execution pipeline including: + - Router (for token-to-expert assignment) + - RoutedExperts (containing expert weight parameters) + - MoERunner (orchestrates the complete forward pass) + + The experts contain both MergedColumnParallel weights (gate_up_proj/w13) + and RowParallelLinear weights (down_proj/w2). + + Note: Mixtral uses w1, w2, and w3 for gate, up, and down_proj. We + copy that naming convention here and handle any remapping in the + load_weights function in each model implementation. + + Args: + num_experts: Number of experts in the model (global count) + top_k: Number of experts selected for each token + hidden_size: Input hidden state size of the transformer + intermediate_size: Intermediate size of the experts + params_dtype: Data type for the parameters + renormalize: Whether to renormalize the logits in the router + use_grouped_topk: Whether to use grouped top-k routing + num_expert_group: Number of expert groups for grouped top-k + topk_group: Top-k value per group for grouped top-k + quant_config: Quantization configuration + tp_size: Tensor parallelism size (None = use global default) + dp_size: Data parallelism size (None = use global default) + pcp_size: Pipeline context parallelism size (None = use global default) + prefix: Layer name prefix for weight loading + custom_routing_function: Custom routing function override + router: Pre-configured router instance (None = create default) + 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 + 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.) + enable_eplb: Whether to enable expert parallelism load balancer + num_redundant_experts: Number of redundant experts for EPLB + has_bias: Whether expert layers have bias terms + is_sequence_parallel: Whether sequence parallelism is enabled + reduce_results: Whether to all-reduce the final output. Setting this + to False (to fuse the all-reduce downstream) is only honored on + the late-AR path. + ckpt_names: Checkpoint parameter name tuple (gate_proj, down_proj, + up_proj) used for weight loading + n_shared_experts: Number of shared experts to fuse into the routed + grouped GEMM (ROCm; requires aiter FSE or the router-append path) + router_logits_dtype: Data type for router logits buffers + gate: Pre-configured gate module + shared_experts: Pre-configured shared experts module + shared_expert_gate: Pre-configured shared expert gate module + routed_input_transform: Input transformation module + routed_output_transform: Output transformation module + apply_routed_scale_to_output: Whether to apply routed_scaling_factor to + output instead of topk_weights + zero_expert_type: Type of zero expert handling + hash_indices_table: Hash table for expert indices + runner_cls: Custom MoERunner class (None = use default MoERunner) + runner_args: Additional arguments for runner constructor + routed_experts_cls: Custom RoutedExperts class (None = use default) + routed_experts_args: Additional arguments for routed_experts constructor + + Returns: + MoERunner: Configured MoE execution pipeline ready for forward passes + """ + aphrodite_config = get_current_aphrodite_config() + + layer_name = prefix + + moe_activation = MoEActivation.from_str(activation) + is_act_and_mul = moe_activation.is_gated + + moe_parallel_config = make_parallel_config( + tp_size=tp_size, + dp_size=dp_size, + pcp_size=pcp_size, + is_sequence_parallel=is_sequence_parallel, + parallel_config=aphrodite_config.parallel_config, + ) + + # Resolve the deferred all-reduce request against the parallel config. + skip_final_all_reduce = ( + not reduce_results + and not moe_parallel_config.use_all2all_kernels + and not moe_parallel_config.is_sequence_parallel + and zero_expert_type is None + ) + + global_num_experts, logical_num_experts, num_fused_shared_experts = determine_expert_counts( + num_experts, + num_redundant_experts, + n_shared_experts, + is_act_and_mul, + ) + + # Initialize EPLB manager (or None?) + eplb_state: EplbLayerState | None = None + if enable_eplb: + use_ep = moe_parallel_config.use_ep + ep_size = moe_parallel_config.ep_size + if use_ep and global_num_experts % ep_size != 0: + raise ValueError( + f"EPLB currently only supports even distribution of " + f"experts across ranks. Got {global_num_experts} experts " + f"and {ep_size} EP ranks." + ) + eplb_state = EplbLayerState() + else: + assert num_redundant_experts == 0, "Redundant experts are only supported with EPLB." + + max_num_batched_tokens = aphrodite_config.scheduler_config.max_num_batched_tokens + + # Create ExpertMapManager to handle expert mapping and placement for EP. + # See ExpertMapManager for a detailed description of what it does and when + # it is required. + expert_map_manager = ExpertMapManager( + max_num_batched_tokens=max_num_batched_tokens, + top_k=top_k, + global_num_experts=global_num_experts, + num_redundant_experts=num_redundant_experts, + num_expert_group=num_expert_group, + moe_parallel_config=moe_parallel_config, + placement_strategy=aphrodite_config.parallel_config.expert_placement_strategy, + enable_eplb=eplb_state is not None, + num_fused_shared_experts=num_fused_shared_experts, + rocm_aiter_enabled=rocm_aiter_ops.is_fused_moe_enabled() and is_act_and_mul, + ) + + # TODO(bnell): we should not have to create a router if the kernel is + # monolithic. + if router is None: + router = create_fused_moe_router( + top_k=top_k, + global_num_experts=global_num_experts, + eplb_state=eplb_state, + renormalize=renormalize, + use_grouped_topk=use_grouped_topk, + num_expert_group=num_expert_group, + topk_group=topk_group, + custom_routing_function=custom_routing_function, + scoring_func=scoring_func, + # When apply_routed_scale_to_output is True, we set the scaling factor + # to 1.0 so it ends up being a nop. Applying the scale will be handled + # by the runner in this case. + # The member variable must be set in the same way as the router since + # some quantization methods can access it. + routed_scaling_factor=routed_scaling_factor if not apply_routed_scale_to_output else 1.0, + e_score_correction_bias=e_score_correction_bias, + num_fused_shared_experts=num_fused_shared_experts, + # Fused shared-expert slot weight. With apply_routed_scale_to_output + # the runner scales the combined output by routed_scaling_factor, so + # the shared slot weight must be 1/routed_scaling_factor for its net + # contribution to be 1.0 (matching the un-scaled separate-MLP add). + shared_expert_weight=( + (1.0 / routed_scaling_factor) + if (apply_routed_scale_to_output and num_fused_shared_experts > 0 and routed_scaling_factor) + else 1.0 + ), + zero_expert_type=zero_expert_type, + num_logical_experts=logical_num_experts, + hash_indices_table=hash_indices_table, + ) + + if params_dtype is None: + params_dtype = torch.get_default_dtype() + + # FIXME (varun): We should have a better way of inferring the activation + # datatype. This works for now as the tensor datatype entering the MoE + # operation is typically unquantized (i.e. float16/bfloat16). + if aphrodite_config.model_config is not None: + moe_in_dtype = aphrodite_config.model_config.dtype + else: + # TODO (bnell): This is a hack to get test_mixtral_moe to work + # since model_config is not set in the pytest test. + moe_in_dtype = params_dtype + + moe_config = FusedMoEConfig( + num_experts=global_num_experts, + experts_per_token=top_k, + hidden_dim=hidden_size, + intermediate_size=intermediate_size, + intermediate_pad=intermediate_pad, + num_local_experts=expert_map_manager.local_num_experts, + num_logical_experts=logical_num_experts, + moe_parallel_config=moe_parallel_config, + in_dtype=moe_in_dtype, + moe_backend=aphrodite_config.kernel_config.moe_backend, + router_logits_dtype=router_logits_dtype, + max_num_tokens=max_num_batched_tokens, + has_bias=has_bias, + is_lora_enabled=aphrodite_config.lora_config is not None, + activation=moe_activation, + device=aphrodite_config.device_config.device, + routing_method=router.routing_method_type, # Not ideal + swiglu_limit=swiglu_limit, + swiglu_alpha=swiglu_alpha, + swiglu_beta=swiglu_beta, + max_capture_size=aphrodite_config.compilation_config.max_cudagraph_capture_size, + skip_final_all_reduce=skip_final_all_reduce, + ) + + logger.debug("FusedMoEConfig = %s", moe_config) + + # Create RoutedExperts instance BEFORE create_weights() + # This will hold all expert weight parameters + if routed_experts_cls is None: + routed_experts_cls = RoutedExperts + + assert params_dtype is not None + routed_experts = routed_experts_cls( + layer_name, + params_dtype, + moe_config, + quant_config, + expert_map_manager=expert_map_manager, + ckpt_gate_proj_name=ckpt_names[0], + ckpt_down_proj_name=ckpt_names[1], + ckpt_up_proj_name=ckpt_names[2], + # Extra params that are needed by quant_methods, pass along for now + # Prefer getting these from other sources, e.g. moe_config or + # router object + renormalize=renormalize, + use_grouped_topk=use_grouped_topk, + num_expert_group=num_expert_group, + topk_group=topk_group, + custom_routing_function=custom_routing_function, + scoring_func=scoring_func, + routed_scaling_factor=routed_scaling_factor if not apply_routed_scale_to_output else 1.0, + swiglu_limit=swiglu_limit, + swiglu_alpha=swiglu_alpha, + swiglu_beta=swiglu_beta, + # TODO get from router? needs to be truncated? + e_score_correction_bias=e_score_correction_bias, + apply_router_weight_on_input=apply_router_weight_on_input, + **routed_experts_args if routed_experts_args is not None else {}, + ) + + if runner_cls is None: + runner_cls = MoERunner + + runner = runner_cls( + layer_name=layer_name, + moe_config=moe_config, + router=router, + routed_experts=routed_experts, + enable_dbo=aphrodite_config.parallel_config.enable_dbo, + gate=gate, + shared_expert_gate=shared_expert_gate, + shared_experts=shared_experts, + routed_input_transform=routed_input_transform, + routed_output_transform=routed_output_transform, + # When apply_routed_scale_to_output is True, we allow + # the scaling factor to be passed to the runner, otherwise + # we pass 1.0 so it ends up being a nop. + routed_scaling_factor=routed_scaling_factor if apply_routed_scale_to_output else 1.0, + **runner_args if runner_args is not None else {}, + ) + + return runner + + +def fused_moe_make_expert_params_mapping( + model: torch.nn.Module, + ckpt_gate_proj_name: str, + ckpt_down_proj_name: str, + ckpt_up_proj_name: str, + num_experts: int, + num_redundant_experts: int = 0, + routed_experts_prefix: str = "routed_experts", +) -> list[tuple[str, str, int, str]]: + """Delegate to EPLB manager.""" + return RoutedExperts.make_expert_params_mapping( + model, + ckpt_gate_proj_name, + ckpt_down_proj_name, + ckpt_up_proj_name, + num_experts, + num_redundant_experts, + routed_experts_prefix, + ) 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/modular_kernel.py.orig b/aphrodite/model_executor/layers/fused_moe/modular_kernel.py.orig new file mode 100644 index 0000000000..80c3877049 --- /dev/null +++ b/aphrodite/model_executor/layers/fused_moe/modular_kernel.py.orig @@ -0,0 +1,1678 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from abc import ABC, abstractmethod +from collections.abc import Callable +from dataclasses import dataclass +from enum import Enum +from math import prod +from typing import final + +import torch + +import aphrodite.envs as envs +from aphrodite.logger import init_logger +from aphrodite.model_executor.layers.fused_moe.activation import ( + MoEActivation, + apply_moe_activation, +) +from aphrodite.model_executor.layers.fused_moe.config import ( + FusedMoEConfig, + FusedMoEParallelConfig, + FusedMoEQuantConfig, + RoutingMethodType, +) +from aphrodite.model_executor.layers.fused_moe.runner.shared_experts import ( + SharedExperts, + SharedExpertsOrder, +) +from aphrodite.model_executor.layers.fused_moe.utils import ( + _resize_cache, +) +from aphrodite.model_executor.layers.quantization.utils.quant_utils import ( + QuantKey, +) +from aphrodite.platforms import current_platform +from aphrodite.v1.worker.ubatching import ( + dbo_enabled, + dbo_maybe_run_recv_hook, + dbo_register_recv_hook, + dbo_yield, +) +from aphrodite.v1.worker.workspace import current_workspace_manager + +logger = init_logger(__name__) + +# +# This file defines a set of base classes used to make MoE kernels more modular. +# The goal is to be able to utilize different communication mechanisms with +# any fused MoE kernel without needing to have combinatoric implementations. +# +# The fused moe kernels are broken down into the following components: +# +# [Router] → [Quantize-Dispatch] → [Permute-Experts-Unpermute] → [Combine] +# +# Each component will be independent of (but may inform) the others except for +# [Quantize-Dispatch] and `[Combine] (see below). The components can then be +# mixed and matched with so that DP+EP can be supported easily for multiple +# MoE kernel implementations. +# +# The following main classes are defined: +# * FusedMoEPrepareAndFinalizeModular - an abstract base class for preparation of MoE +# inputs (e.g. quantization, distribution) and finalization of Moe outputs. +# The prepare method must take care of any needed quantization and the +# finalize method, informed by the FusedMoEExpertsModular method, +# may apply weights and/or do the final reduction of the output. +# * FusedMoEExpertsModular - an abstract base class for the main fused +# MoE operation, i.e matmul + act_mul + optionally quant + matmul. +# Some FusedMoEExpertsModular implementations may choose to do +# the weight application and/or reduction. The class communicates this +# to [Finalize] via a TopKWeightAndReduce object. +# * FusedMoEModularKernel - an interface class that combines a +# FusedMoEPrepareAndFinalizeModular and a FusedMoEExpertsModular to +# provide the standard fused MoE kernel interface. +# * TopKWeightAndReduce - A TopKWeightAndReduce implementation chosen +# by the FusedMoEExpertsModular implementation that is passed +# on to [Finalize]. +# +# [Quantize-Prepare] and [Finalize] functionality are bundled into a single +# class `FusedMoEPrepareAndFinalizeModular` since they could use collective +# communication mechanisms that need to be consistent. +# + + +class FusedMoEActivationFormat(Enum): + """ + The standard activation format (num_tokens, hidden dim). + """ + + Standard = ("standard",) + """ + The batched experts format (num experts, max tokens per expert, hidden dim) + """ + BatchedExperts = ("batched_experts",) + + +@dataclass +class ExpertTokensMetadata: + """ + Metadata regarding expert-token routing. + """ + + expert_num_tokens: torch.Tensor + expert_num_tokens_cpu: torch.Tensor | None + + @staticmethod + def make_from_list(expert_num_tokens_list: list[int], device: str) -> "ExpertTokensMetadata": + expert_num_tokens_cpu = torch.tensor(expert_num_tokens_list, device="cpu", dtype=torch.int32) + return ExpertTokensMetadata( + expert_num_tokens=expert_num_tokens_cpu.to(device, non_blocking=True), + expert_num_tokens_cpu=expert_num_tokens_cpu, + ) + + +class TopKWeightAndReduce(ABC): + """ + An abstract base class for weight application and reduction implementations. + """ + + @abstractmethod + def apply( + self, + output: torch.Tensor | None, + fused_expert_output: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + apply_router_weight_on_input: bool, + ) -> torch.Tensor: + """ + Apply topk_weights to the fused_experts_outputs and/or reduce. + If an output tensor is not passed, it will be created in the + function. + """ + raise NotImplementedError + + +# +# PrepareResultType is a tuple of: +# - quantized + dispatched a. +# - quantized + dispatched a1_scales. +# - Optional ExpertTokensMetadata containing gpu/cpu tensors +# as big as the number of local experts with the information about the +# number of tokens assigned to each local expert. +# - Optional dispatched expert topk IDs +# - Optional dispatched expert topk weight +# +# See `prepare` method below. +# +PrepareResultType = tuple[ + torch.Tensor, + torch.Tensor | None, + ExpertTokensMetadata | None, + torch.Tensor | None, + torch.Tensor | None, +] + +# +# PrepareResultType is a tuple of: +# - quantized + dispatched a. +# - quantized + dispatched a1_scales. +# - dispatched router logits. +# +# See `prepare_monolithic` method below. +# +PrepareMonolithicResultType = tuple[ + torch.Tensor, + torch.Tensor | None, + torch.Tensor, +] + +ReceiverType = Callable[[], PrepareResultType] + +################################################################################ +# Prepare/Finalize +################################################################################ + + +class FusedMoEPrepareAndFinalize(ABC): + """ + An abstract base class for the [Quantize-Prepare] and [Finalize] steps + described above. + + There are two variants of this class: + * FusedMoEPrepareAndFinalizeModular - this operates on topk ids and weights + * FusedMoEPrepareAndFinalizeMonolithic - the operates on router_logits + """ + + def post_init_setup(self, fused_experts: "FusedMoEExperts"): + """ + Initialize FusedMoEPrepareAndFinalizeModular settings that depend on + FusedMoEExpertsModular experts object. + The FusedMoEPrepareAndFinalizeModular implementations that have such + dependencies may choose to override this function. + """ + return + + @property + @abstractmethod + def activation_format(self) -> FusedMoEActivationFormat: + """ + A property indicating the output format of the activations for the + 'prepare' method. + """ + raise NotImplementedError + + @abstractmethod + def topk_indices_dtype(self) -> torch.dtype | None: + """ + The PrepareFinalize All2All implementations generally constrain the + dtype of the topk_ids they support. This function returns the + required topk indices dtype so it can be respected. + Return None if there are no such restrictions. + """ + raise NotImplementedError + + @abstractmethod + def max_num_tokens_per_rank(self) -> int | None: + """ + Some PrepareFinalize All2All implementations are batched. Meaning, + they can process only as set of tokens at a time. This + function returns the batch size i.e the maximum number of tokens + the implementation can process at a time. + Return None if there are no such restrictions. + """ + raise NotImplementedError + + @abstractmethod + def num_dispatchers(self) -> int: + raise NotImplementedError + + @abstractmethod + def output_is_reduced(self) -> bool: + """ + Indicates whether or not the output of finalize is reduced across all + ranks. + """ + raise NotImplementedError + + def supports_async(self) -> bool: + """ + Indicates whether or not this class implements prepare_async and + finalize_async. + """ + return False + + def on_commit(self) -> None: + """ + Runs after this prepare/finalize has been committed to the active + MoE kernel. + """ + return + + +# TODO: pass FusedMoEParallelConfig in as ctor parameter? +class FusedMoEPrepareAndFinalizeModular(FusedMoEPrepareAndFinalize): + """ + An abstract base class for the [Quantize-Prepare] and [Finalize] steps + described above for the Modular case. + """ + + @abstractmethod + def prepare( + self, + a1: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + num_experts: int, + expert_map: torch.Tensor | None, + apply_router_weight_on_input: bool, + quant_config: FusedMoEQuantConfig, + defer_input_quant: bool, + ) -> PrepareResultType: + """ + Perform any quantization (and/or) dispatching needed for this kernel. + - a1: The (unquantized) input to the MoE layer. + - topk_ids: The topk ids. + - topk_weights: The topk weights. + - num_experts: The total number of experts in the global expert space. + - expert_map: A tensor mapping expert indices from the global expert + space to the local expert space of the expert parallel shard. + - apply_router_weight_on_input: When True, apply the weights to the + activations, before quantization + dispatching. + - quant_config: Quantization info provided by the fused experts. + - defer_input_quant: Runtime parameter indicating whether or not to + defer input quantization to the FusedMoEExpertsModular + in cases where the compute kernel expects unquantized inputs + + Returns a tuple of: + - quantized + dispatched a. + - Optional quantized + dispatched a1_scales. + - Optional ExpertTokensMetadata containing gpu/cpu tensors + as big as the number of local experts with the information about the + number of tokens assigned to each local expert. + - Optional dispatched expert topk IDs + - Optional dispatched expert topk weight + """ + raise NotImplementedError + + def prepare_async( + self, + a1: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + num_experts: int, + expert_map: torch.Tensor | None, + apply_router_weight_on_input: bool, + quant_config: FusedMoEQuantConfig, + defer_input_quant: bool, + ) -> tuple[Callable, ReceiverType] | ReceiverType: + """ + Perform any quantization (and/or) dispatching needed for this kernel + but do not wait for results from other workers. + - a1: The (unquantized) input to the MoE layer. + - a1_scale: Optional scales for a1 + - a2_scale: Optional scales for the second MoE gemm. Required to make + sure the quantization is consistent for both gemms. + - topk_ids: The topk ids. + - topk_weights: The topk weights. + - num_experts: The total number of experts in the global expert space. + - expert_map: A tensor mapping expert indices from the global expert + space to the local expert space of the expert parallel shard. + - apply_router_weight_on_input: When True, apply the weights to the + activations, before quantization + dispatching. + - defer_input_quant: Runtime parameter indicating whether or not to + defer input quantization to the FusedMoEExpertsModular + in cases where the compute kernel expects unquantized inputs + + Returns a callback or a hook callback pair that when invoked waits for + results from other workers and has the same return signature as + `prepare`, if a hook is returned this is more lightweight check that + the recv is complete without doing extra work (used by DBO, will be + refactored in the very near future) + + e.g. + + ret = obj.prepare_async(...) + + if isinstance(ret, tuple): + hook, receiver = ret + hook() + + if hook is not None: + a, a_scales, expert_meta, topk_ids, topk_weights = receiver() + + is equivalent to: + + a, a_scales, expert_meta, topk_ids, topk_weights = obj.prepare(...) + """ + raise NotImplementedError + + @abstractmethod + def finalize( + self, + output: torch.Tensor, + fused_expert_output: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + apply_router_weight_on_input: bool, + weight_and_reduce_impl: TopKWeightAndReduce, + ) -> None: + """ + Perform any combine plus apply weights and perform a reduction on the + fused experts output. + - output: The output tensor, written in place. Must be (M, K) shape. + - fused_expert_output: The unweighted, unreduced output of the fused + experts, it will have (M, topk, K) shape. + - topk_weights: The weights to be applied to the fused_experts_output. + - topk_ids: The topk_ids. + - apply_router_weight_on_input: When False, apply the weights to + fused_expert_output. + - weight_and_reduce_impl: An optional TopKWeightAndReduce + implementation. + """ + raise NotImplementedError + + def finalize_async( + self, + output: torch.Tensor, + fused_expert_output: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + apply_router_weight_on_input: bool, + weight_and_reduce_impl: TopKWeightAndReduce, + ) -> tuple[Callable, Callable] | Callable: + """ + Perform any combine plus apply weights and perform a reduction on the + fused experts output but do not wait for results from other workers. + - output: The output tensor, written in place. Must be (M, K) shape. + - fused_expert_output: The unweighted, unreduced output of the fused + experts, it will have (M, topk, K) shape. + - topk_weights: The weights to be applied to the fused_experts_output. + - topk_ids: The topk_ids. + - apply_router_weight_on_input: When False, apply the weights to + fused_expert_output. + - weight_and_reduce_impl: An optional TopKWeightAndReduce + implementation. + + Returns a callback or a hook callback pair that when invoked waits for + results from other workers and has the same return signature as + `finalize`, if a hook is returned this is more lightweight check that + the recv is complete without doing extra work (used by DBO, will be + refactored in the very near future) + + ret = obj.finalize_async(output, ...) + ... output not valid yet ... + if isinstance(ret, tuple): + hook, receiver = ret + hook() + receiver() + ... output valid here ... + + is equivalent to: + + obj.finalize(output, ...) + """ + raise NotImplementedError + + +class FusedMoEPrepareAndFinalizeMonolithic(FusedMoEPrepareAndFinalize): + """ + An abstract base class for the [Quantize-Prepare] and [Finalize] steps + described above for the monolithic case. + """ + + @abstractmethod + def prepare( + self, + a1: torch.Tensor, + router_logits: torch.Tensor, + quant_config: FusedMoEQuantConfig, + defer_input_quant: bool = False, + ) -> PrepareMonolithicResultType: + """ + Optional method for subclasses compatible with monolithic + FusedMoEExpertsModular kernels. + + Perform any quantization (and/or) dispatching needed for this kernel. + - a1: The (unquantized) input to the MoE layer. + - quant_config: Quantization info provided by the fused experts. + - defer_input_quant: Runtime parameter indicating whether or not to + defer input quantization to the FusedMoEExpertsModular + + Returns a tuple of: + - quantized + dispatched a. + - Optional quantized + dispatched a1_scales. + """ + raise NotImplementedError + + @abstractmethod + def finalize(self, fused_expert_output: torch.Tensor) -> torch.Tensor: + """ + Optional method for subclasses compatible with monolithic + FusedMoEExpertsModular kernels. + + Perform any combine plus apply weights and perform a reduction on the + fused experts output. + - fused_expert_output: The unweighted, unreduced output of the fused + experts, it will have (M, topk, K) shape. + """ + raise NotImplementedError + + +################################################################################ +# Experts +################################################################################ + + +# TODO: add supported activations method (return string) +class FusedMoEExperts(ABC): + def __init__( + self, + moe_config: FusedMoEConfig, + quant_config: FusedMoEQuantConfig, + max_num_tokens: int | None = None, + num_dispatchers: int | None = None, + ): + """ + moe_config: MoE layer configuration. + quant_config: Quantization parameters for this experts instance. + """ + if self.activation_format() == FusedMoEActivationFormat.Standard and ( + max_num_tokens is not None or num_dispatchers is not None + ): + raise ValueError( + "max_num_tokens and num_dispatchers should only be set for BatchedExperts activation format." + ) + elif self.activation_format() == FusedMoEActivationFormat.BatchedExperts and ( + max_num_tokens is None or num_dispatchers is None + ): + raise ValueError("max_num_tokens and num_dispatchers must be set for BatchedExperts activation format.") + + self.moe_config = moe_config + self.quant_config = quant_config + self.max_num_tokens = max_num_tokens + self.num_dispatchers = num_dispatchers + + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: # noqa: B027 + pass + + @staticmethod + def is_monolithic() -> bool: + raise NotImplementedError("Implemented by subclasses.") + + @property + def expects_unquantized_inputs(self) -> bool: + """ + Whether or not the PrepareFinalize should defer input quantization + in the prepare step. If True, then the Experts kernel will + execute the input quantization itself. + + Sample subclasses that override are AITER and FlashInfer CUTLASS. + """ + return False + + @staticmethod + @abstractmethod + def activation_format() -> FusedMoEActivationFormat: + """ + A property which is a tuple of the input and output activation formats + for the 'apply' method. + """ + raise NotImplementedError + + # + # Various helpers for registering support for various features. + # Used by the oracle to select a particular kernel for a deployment. + # + + @staticmethod + def is_supported_config( + cls: type["FusedMoEExperts"], + moe_config: FusedMoEConfig, + weight_key: QuantKey | None, + activation_key: QuantKey | None, + activation_format: FusedMoEActivationFormat, + ) -> tuple[bool, str | None]: + def _make_reason(reason: str) -> str: + return f"kernel does not support {reason}" + + if not cls._supports_current_device(): + return False, _make_reason(f"current device {current_platform.device_name}") + elif not (moe_config.is_act_and_mul or cls._supports_no_act_and_mul()): + return False, _make_reason("no act_and_mul MLP layer") + elif not cls._supports_activation(moe_config.activation): + return False, _make_reason(f"{moe_config.activation} activation") + elif not cls._supports_quant_scheme(weight_key, activation_key): + return False, _make_reason(f"quantization scheme {weight_key}x{activation_key}") + elif not cls._supports_parallel_config(moe_config.moe_parallel_config): + return False, _make_reason(f"parallel config {moe_config.moe_parallel_config}") + elif not cls._supports_routing_method(moe_config.routing_method, weight_key, activation_key): + return False, _make_reason(f"routing method {moe_config.routing_method}") + elif not cls._supports_router_logits_dtype( + moe_config.router_logits_dtype, + moe_config.routing_method, + ): + return False, _make_reason(f"router logits dtype {moe_config.router_logits_dtype}") + elif not cls._supports_shape(moe_config.hidden_dim): + return False, _make_reason(f"{moe_config.hidden_dim} hidden dim is not supported") + elif activation_format != cls.activation_format(): + return False, _make_reason(f"{activation_format.value} activation format") + elif envs.APHRODITE_BATCH_INVARIANT and not cls._supports_batch_invariance(): + return False, _make_reason("batch invariance") + elif moe_config.is_lora_enabled and not cls.supports_lora(): + return False, _make_reason("LoRA") + return True, None + + @staticmethod + @abstractmethod + def _supports_current_device() -> bool: + """ + Whether the kernel supports the current device type + (compute cability and current platform). + """ + raise NotImplementedError + + @staticmethod + @abstractmethod + def _supports_no_act_and_mul() -> bool: + """ + Whether the kernel supports act_and_mul=False, i.e. + non-gated MoE models like Nemotron-Nano. + """ + raise NotImplementedError + + @staticmethod + @abstractmethod + def _supports_quant_scheme( + weight_key: QuantKey | None, + activation_key: QuantKey | None, + ) -> bool: + raise NotImplementedError + + @staticmethod + @abstractmethod + def _supports_activation(activation: MoEActivation) -> bool: + """ + Whether the kernel supports a particular act function. + """ + raise NotImplementedError + + @staticmethod + @abstractmethod + def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bool: + """ + Whether the kernel supports deployment in particular parallel config. + + Can be overridden if a kernel does not support EP, SP or some other + configuration. + """ + raise NotImplementedError + + @staticmethod + def _supports_routing_method( + routing_method: RoutingMethodType, + weight_key: QuantKey | None, + activation_key: QuantKey | None, + ) -> bool: + """ + Whether the kernel supports a routing method (e.g. GroupedTopK). + + Can be overridden by monolithic kernels that execute the router + in addition to the experts if certain routers are not supported. + """ + return True + + @staticmethod + def _supports_router_logits_dtype( + router_logits_dtype: torch.dtype | None, + routing_method: RoutingMethodType, + ) -> bool: + """ + Whether a kernel supports a particular dtype for router logits input. + + Can be overridden by monolithic kernels that execute the router + in addition to the experts if certain dtypes are not supported. + """ + return True + + @staticmethod + def _supports_shape(hidden_dim: int) -> bool: + """ + Whether a kernel supports a particular shape. Can be overridden if a kernel + has specific shape requirements. + """ + return True + + @staticmethod + def _supports_batch_invariance() -> bool: + """ + Whether the kernel supports batch invariance, i.e. the output does not + depend on the order of the tokens in the input batch. This is useful + for determining if the kernel can used with APHRODITE_BATCH_INVARIANT=1. + """ + return False + + # + # Various helpers for accessing quantization parameters from the + # quant_config. + # + + @property + def quant_dtype(self) -> torch.dtype | str | None: + return self.quant_config.quant_dtype + + @property + def weight_quant_dtype(self) -> torch.dtype | str | None: + return self.quant_config.weight_quant_dtype + + @property + def block_shape(self) -> list[int] | None: + return self.quant_config.block_shape + + @property + def per_act_token_quant(self) -> bool: + return self.quant_config.per_act_token_quant + + @property + def per_out_ch_quant(self) -> bool: + return self.quant_config.per_out_ch_quant + + @property + def a1_scale(self) -> torch.Tensor | None: + return self.quant_config.a1_scale + + @property + def a2_scale(self) -> torch.Tensor | None: + return self.quant_config.a2_scale + + @property + def a1_gscale(self) -> torch.Tensor | None: + return self.quant_config.a1_gscale + + @property + def a2_gscale(self) -> torch.Tensor | None: + return self.quant_config.a2_gscale + + @property + def w1_scale(self) -> torch.Tensor | None: + return self.quant_config.w1_scale + + @property + def w2_scale(self) -> torch.Tensor | None: + return self.quant_config.w2_scale + + @property + def w1_zp(self) -> torch.Tensor | None: + return self.quant_config.w1_zp + + @property + def w2_zp(self) -> torch.Tensor | None: + return self.quant_config.w2_zp + + @property + def w1_bias(self) -> torch.Tensor | None: + return self.quant_config.w1_bias + + @property + def w2_bias(self) -> torch.Tensor | None: + return self.quant_config.w2_bias + + @property + def g1_alphas(self) -> torch.Tensor | None: + return self.quant_config.g1_alphas + + @property + def g2_alphas(self) -> torch.Tensor | None: + return self.quant_config.g2_alphas + + @staticmethod + def supports_lora() -> bool: + """Return True if this expert impl natively handles LoRA. + + LoRA-aware experts should mix in LoRAExpertsMixin, which flips this + to True and provides the per-forward LoRA state plumbing. + """ + return False + + def supports_packed_ue8m0_act_scales(self) -> bool: + """ + A flag indicating whether or not this class can process packed ue8m0 + activation scales. + """ + return False + + +class FusedMoEExpertsModular(FusedMoEExperts): + """ + An abstract base class for the [Permute-Experts-Unpermute] step described + above. + """ + + @staticmethod + def is_monolithic() -> bool: + return False + + def moe_problem_size( + self, + a1: torch.Tensor, + w1: torch.Tensor, + w2: torch.Tensor, + topk_ids: torch.Tensor, + ) -> tuple[int, int, int, int, int]: + """ + Extract the MoE problem size from the given tensor arguments: + - a: The hidden states, input to the MoE layer. + - w1: The first set of expert weights. + - w2: The second set of expert weights. + - topk_ids: The topk ids. + + Note: extracting the problem shape from the weight and activation + tensors is not obvious. It needs to be done this way specifically + due to subtle issues with particular kernels, e.g. the int4 kernels + divide the trailing dimension by two, so it's not "correct" to + extract N or K from the trailing dimension of w1 or w2. Similarly, + some kernels transpose the weights, so this needs to be kept in mind. + + Note: This implementation covers most cases. However, if experts + require a specialized implementation, like MarlinExperts, they are free + to override this function. + """ + assert len(w1.shape) == 3 and len(w2.shape) == 3 + E, N, _ = w1.shape + K = a1.size(-1) + + if a1.dim() == 2: + # Make sure we are using the correct a1 (pre-permute). + assert topk_ids.size(0) == a1.size(0), f"{topk_ids.size(0)} != {a1.size(0)}" + M = a1.size(0) + else: + assert a1.dim() == 3 + assert a1.size(0) == E, f"{a1.size(0)} == {E}" + M = a1.size(1) # This is max_num_tokens + + assert topk_ids.dim() == 2 + topk = topk_ids.size(1) + + return E, M, N, K, topk + + def workspace_dtype(self, act_dtype: torch.dtype) -> torch.dtype: + """ + Workspace type: The dtype to use for the workspace tensors. + """ + return act_dtype + + @abstractmethod + def workspace_shapes( + self, + M: int, + N: int, + K: int, + topk: int, + global_num_experts: int, + local_num_experts: int, + expert_tokens_meta: ExpertTokensMetadata | None, + activation: MoEActivation, + ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]: + """ + Compute the shapes for the temporary and final outputs of the two gemms + and activation in the fused expert function. Since the gemms are + independent, the workspace for the first gemm can be shared with the + workspace for the last gemm. + + Inputs: + - M: number of tokens. + - N: Row (or column) dimension of expert weights. + - K: hidden dimension + - topk: The number of top-k experts to select. + - global_num_experts: global number of experts. + - local_num_experts: local number of experts due to DP/EP. + - expert_tokens_meta: number of tokens per expert metadata for batched + format. + + Returns a tuple of: + - workspace13 shape tuple: must be large enough to hold the + result of either expert gemm. + - workspace2 shape tuple: must be large enough to hold the + result of the activation function. + - output shape tuple: must be exact size of the final gemm output. + - Note: workspace shapes can be 0 if the workspace is not needed. + But in order for activation chunking to work, the first dimension + of each tuple must be the number of tokens when the shape is + not 0. + """ + raise NotImplementedError + + @staticmethod + def adjust_N_for_activation(N: int, activation: MoEActivation) -> int: + """ + Calculate the output dimension for the activation function. + + For *_no_mul activations (e.g. relu2_no_mul), + there's no gate/up split, so output size equals input size (N). + + For regular gated activations (e.g., silu, gelu, swigluoai), + output size is N // 2 due to gate × activation(up) multiplication. + + Args: + N: The intermediate size (width of w1/w3 weights). + activation: The activation function enum. + + Returns: + The output dimension after activation. + """ + return N if not activation.is_gated else N // 2 + + def activation( + self, + activation: MoEActivation, + output: torch.Tensor, + input: torch.Tensor, + *, + clamp_limit: float | None = None, + alpha: float = 1.0, + beta: float = 0.0, + topk_ids: torch.Tensor | None = None, + expert_map: torch.Tensor | None = None, + ) -> None: + apply_moe_activation( + activation, + output, + input, + clamp_limit=clamp_limit, + alpha=alpha, + beta=beta, + topk_ids=topk_ids, + expert_map=expert_map, + ) + + @abstractmethod + def finalize_weight_and_reduce_impl(self) -> TopKWeightAndReduce: + raise NotImplementedError + + @abstractmethod + def apply( + self, + output: torch.Tensor, + hidden_states: torch.Tensor, + w1: torch.Tensor, + w2: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + activation: MoEActivation, + global_num_experts: int, + expert_map: torch.Tensor | None, + a1q_scale: torch.Tensor | None, + a2_scale: torch.Tensor | None, + workspace13: torch.Tensor, + workspace2: torch.Tensor, + expert_tokens_meta: ExpertTokensMetadata | None, + apply_router_weight_on_input: bool, + ) -> None: + """ + This function computes the intermediate result of a Mixture of Experts + (MoE) layer using two sets of weights, w1 and w2. + + Parameters: + - output: (torch.Tensor): The unweighted, unreduced output tensor. + - hidden_states: (torch.Tensor): The (quantized) input tensor to the MoE + layer. + - w1 (torch.Tensor): The first set of expert weights. + - w2 (torch.Tensor): The second set of expert weights. + - topk_weights: A map of row to expert weights. Some implementations + choose to do weight application. + - topk_ids (torch.Tensor): A map of row to expert id. + - activation (str): The activation function to apply after the first + MoE layer. + - global_num_experts (int): The total number of experts in the global + expert space. + - expert_map (Optional[torch.Tensor]): A tensor mapping expert indices + from the global expert space to the local expert space of the expert + parallel shard. + - a1q_scale (Optional[torch.Tensor]): Optional quantized scale to be + used for a1. Result of quantization from prepare/finalize and not + from the FusedMoEQuantConfig. + - workspace13 (torch.Tensor): A scratch tensor used for gemm outputs + must be large enough to hold output of either MoE gemm. + - workspace2 (torch.Tensor): A scratch tensor used for the activation + function. + - expert_tokens_meta (Optional[ExpertTokensMetadata]) - An optional + ExpertTokensMetadata object containing gpu/cpu tensors + as big as the number of local experts with the information about the + number of tokens assigned to each local expert. + - apply_router_weight_on_input: True if router weights are already + applied on the input. This is relevant if the implementation + chooses to do weight application. + """ + raise NotImplementedError + + +class FusedMoEExpertsMonolithic(FusedMoEExperts): + """ + An abstract base class for the [Permute-Experts-Unpermute] step described + above, but with the monolithic interface (accepts router logits + rather than topk ids and weights). + """ + + @staticmethod + def _supports_routing_method( + routing_method: RoutingMethodType, + weight_key: QuantKey | None, + activation_key: QuantKey | None, + ) -> bool: + """ + Whether the kernel supports a routing method (e.g. GroupedTopK). + + Monolithic kernels should explicitly opt-in to support. + """ + raise NotImplementedError + + @staticmethod + def _supports_router_logits_dtype( + router_logits_dtype: torch.dtype | None, + routing_method: RoutingMethodType, + ) -> bool: + """ + Whether the kernel supports a dtype for router logits. + + Modular kernels should opt-in to support. + """ + raise NotImplementedError + + @staticmethod + def is_monolithic() -> bool: + return True + + routing_replay_capture_fn: Callable[[torch.Tensor], None] | None = None + _routing_replay_buffer: torch.Tensor | None = None + + def supports_routing_replay_capture(self) -> bool: + """Whether this expert supports routing replay capture. + + Subclasses backed by a kernel that exposes routed expert IDs + (e.g. FlashInfer's ``routing_replay_out``) should override. + """ + return False + + def set_capture_fn( + self, + capture_fn: Callable[[torch.Tensor], None] | None, + ) -> None: + self.routing_replay_capture_fn = capture_fn + if capture_fn is None: + self._routing_replay_buffer = None + return + self._routing_replay_buffer = torch.empty( + (self.moe_config.max_num_tokens, self.moe_config.experts_per_token), + dtype=torch.int16, + device=self.moe_config.device, + ) + + def _maybe_make_routing_replay_buffer( + self, + num_tokens: int, + device: torch.device, + ) -> torch.Tensor | None: + if self.routing_replay_capture_fn is None: + return None + buf = self._routing_replay_buffer + assert buf is not None + if buf.shape[0] < num_tokens or buf.device != device: + raise ValueError( + "Routing replay buffer was initialized for " + f"{buf.shape[0]} tokens on {buf.device}, but the kernel " + f"received {num_tokens} tokens on {device}." + ) + return buf + + def _maybe_dispatch_routing_replay( + self, + routing_replay_out: torch.Tensor | None, + num_tokens: int, + ) -> None: + if routing_replay_out is None or self.routing_replay_capture_fn is None: + return + self.routing_replay_capture_fn(routing_replay_out[:num_tokens]) + + def apply( + self, + hidden_states: torch.Tensor, + w1: torch.Tensor, + w2: torch.Tensor, + router_logits: torch.Tensor, + activation: MoEActivation, + global_num_experts: int, + expert_map: torch.Tensor | None, + a1q_scale: torch.Tensor | None, + apply_router_weight_on_input: bool, + # grouped topk + fused topk bias parameters + num_expert_group: int | None = None, + e_score_correction_bias: torch.Tensor | None = None, + routed_scaling_factor: float | None = None, + topk_group: int | None = None, + ) -> torch.Tensor: + """ + Same as apply(), except uses router_logits as opposed + to the topk_ids and topk_weights. This is useful for kernels + with fused router and fused_experts (e.g. FLASHINFER_TRTLLM). + """ + raise NotImplementedError + + +################################################################################ +# Kernel +################################################################################ + + +@final +class FusedMoEKernelModularImpl: + def __init__( + self, + prepare_finalize: FusedMoEPrepareAndFinalizeModular, + fused_experts: FusedMoEExpertsModular, + ): + self.prepare_finalize = prepare_finalize + self.fused_experts = fused_experts + self.shared_experts: SharedExperts | None = None + moe_parallel_config = fused_experts.moe_config.moe_parallel_config + self.moe_parallel_config = moe_parallel_config + self.is_dp_ep = ( + moe_parallel_config is not None and moe_parallel_config.dp_size > 1 and moe_parallel_config.use_ep + ) + + def _allocate_buffers( + self, + out_dtype: torch.dtype, + device: torch.device, + M_chunk: int, + M_full: int, + N: int, + K: int, + top_k: int, + global_num_experts: int, + local_num_experts: int, + expert_tokens_meta: ExpertTokensMetadata | None, + activation: MoEActivation, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """ + Allocate temporary and output buffers for the fused experts op. + Inputs: + - out_dtype: output type of workspace and output tensors. + - device: the device of the workspace and output tensors. + See `workspace_shapes` for a description of the remainder of arguments. + Returns a tuple of (workspace13, workspace2, output) tensors. + """ + assert M_full > 0 and M_chunk > 0 + + workspace_dtype = self.fused_experts.workspace_dtype(out_dtype) + + # Get intermediate workspace shapes based off the chunked M size. + workspace13_shape, workspace2_shape, _ = self.fused_experts.workspace_shapes( + M_chunk, + N, + K, + top_k, + global_num_experts, + local_num_experts, + expert_tokens_meta, + activation, + ) + + # Get final output shape based on the full M size. + _, _, fused_out_shape = self.fused_experts.workspace_shapes( + M_full, + N, + K, + top_k, + global_num_experts, + local_num_experts, + expert_tokens_meta, + activation, + ) + + # We can reuse the memory between cache1 and cache3 because by the + # time we need cache3, we're done with cache1. + # Reuse workspace13 for the output since there is only one chunk. + max_shape_size = max(prod(workspace13_shape), prod(fused_out_shape)) + common_workspace, workspace2 = current_workspace_manager().get_simultaneous( + ((max_shape_size,), workspace_dtype), + (workspace2_shape, workspace_dtype), + ) + workspace13 = _resize_cache(common_workspace, workspace13_shape) + fused_out = _resize_cache(common_workspace, fused_out_shape) + + return workspace13, workspace2, fused_out + + def _maybe_apply_shared_experts( + self, + shared_experts: SharedExperts | None, + shared_experts_input: torch.Tensor | None, + ): + if shared_experts is not None: + assert self.prepare_finalize.supports_async() + assert shared_experts_input is not None + shared_experts( + shared_experts_input, + SharedExpertsOrder.MK_INTERNAL_OVERLAPPED, + ) + + def _prepare( + self, + hidden_states: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + global_num_experts: int, + expert_map: torch.Tensor | None, + apply_router_weight_on_input: bool, + ) -> tuple[ + torch.Tensor, + torch.Tensor | None, + ExpertTokensMetadata | None, + torch.Tensor, + torch.Tensor, + ]: + """ + The _prepare method is a wrapper around self.prepare_finalize.prepare + that handles DBO and async. + """ + + if not self.prepare_finalize.supports_async(): + # We shouldn't be running an a2a kernel that doesn't + # support async prepare/finalize + # TODO(lucas): enable in follow-up + assert not dbo_enabled() + + ( + a1q, + a1q_scale, + expert_tokens_meta, + _expert_topk_ids, + _expert_topk_weights, + ) = self.prepare_finalize.prepare( + hidden_states, + topk_weights, + topk_ids, + global_num_experts, + expert_map, + apply_router_weight_on_input, + self.fused_experts.quant_config, + defer_input_quant=self.fused_experts.expects_unquantized_inputs, + ) + else: + # Overlap shared expert compute with all2all dispatch. + dbo_maybe_run_recv_hook() + prepare_ret = self.prepare_finalize.prepare_async( + hidden_states, + topk_weights, + topk_ids, + global_num_experts, + expert_map, + apply_router_weight_on_input, + self.fused_experts.quant_config, + defer_input_quant=self.fused_experts.expects_unquantized_inputs, + ) + + # TODO(lucas): refactor this in the alternative schedules followup + # currently unpack if we have hook + receiver pair or just + # receiver (see finalize_async docstring) + hook, receiver = prepare_ret if isinstance(prepare_ret, tuple) else (None, prepare_ret) + + if hook is not None: + if dbo_enabled(): + # If DBO is being used, register the hook with the ubatch + # context and call it in dbo_maybe_run_recv_hook instead of + # passing it to the receiver. + dbo_register_recv_hook(hook) + dbo_yield() + else: + hook() + + ( + a1q, + a1q_scale, + expert_tokens_meta, + _expert_topk_ids, + _expert_topk_weights, + ) = receiver() + + # Maybe prepare gathered topk_ids and topk_weights from other EP ranks. + topk_ids = topk_ids if _expert_topk_ids is None else _expert_topk_ids + topk_weights = topk_weights if _expert_topk_weights is None else _expert_topk_weights + + return a1q, a1q_scale, expert_tokens_meta, topk_ids, topk_weights + + def _fused_experts( + self, + in_dtype: torch.dtype, + a1q: torch.Tensor, + a1q_scale: torch.Tensor | None, + w1: torch.Tensor, + w2: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + activation: MoEActivation, + global_num_experts: int, + local_num_experts: int, + expert_map: torch.Tensor | None, + apply_router_weight_on_input: bool, + expert_tokens_meta: ExpertTokensMetadata | None, + output_alias: torch.Tensor | None = None, + ) -> torch.Tensor: + _, M_full, N, K, top_k = self.fused_experts.moe_problem_size(a1q, w1, w2, topk_ids) + + # This happens when none of the tokens from the all2all reach this + # EP rank. Also, note that this is only relevant for CUDAGraph + # incompatible all2all kernels like the DeepEP high-throughput + # kernels. CUDAGraph compatible all2all kernels like the DeepEP + # low-latency kernels are always batched and can never run into + # the tensor.numel() == 0 case. + if M_full == 0: + return torch.empty_like(a1q, dtype=in_dtype) + + workspace13, workspace2, fused_out = self._allocate_buffers( + in_dtype, + a1q.device, + M_full, + M_full, + N, + K, + top_k, + global_num_experts, + local_num_experts, + expert_tokens_meta, + activation, + ) + + # 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 + # double-copy MoE write-back path). + 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() + ): + fused_out = output_alias + + self.fused_experts.apply( + output=fused_out, + hidden_states=a1q, + w1=w1, + w2=w2, + topk_weights=topk_weights, + topk_ids=topk_ids, + activation=activation, + global_num_experts=global_num_experts, + expert_map=expert_map, + a1q_scale=a1q_scale, + a2_scale=self.fused_experts.a2_scale, + workspace13=workspace13, + workspace2=workspace2, + expert_tokens_meta=expert_tokens_meta, + apply_router_weight_on_input=apply_router_weight_on_input, + ) + + return fused_out + + def _finalize( + self, + output: torch.Tensor, + fused_out: torch.Tensor, + hidden_states: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + apply_router_weight_on_input: bool, + shared_experts: SharedExperts | None, + shared_experts_input: torch.Tensor | None, + ) -> torch.Tensor: + """ + The _finalize method is a wrapper around self.prepare_finalize.finalize + that handles DBO, async and shared expert overlap. + + Args: + shared_experts: SharedExperts | None. The shared experts if any. + shared_experts_input: Optional separate input for shared experts. + When latent MoE is used, hidden_states is the latent-projected + tensor (smaller dimension) used by routed experts, while + shared_experts_input is the original hidden_states (full + dimension) needed by the shared expert MLP. + """ + if not self.prepare_finalize.supports_async(): + assert not dbo_enabled() + + self.prepare_finalize.finalize( + output, + fused_out, + topk_weights, + topk_ids, + apply_router_weight_on_input, + self.fused_experts.finalize_weight_and_reduce_impl(), + ) + else: + finalize_ret = self.prepare_finalize.finalize_async( + output, + fused_out, + topk_weights, + topk_ids, + apply_router_weight_on_input, + self.fused_experts.finalize_weight_and_reduce_impl(), + ) + self._maybe_apply_shared_experts(shared_experts, shared_experts_input) + + # TODO(lucas): refactor this in the alternative schedules followup + # currently unpack if we have hook + receiver pair or just + # receiver (see finalize_async docstring) + hook, receiver = finalize_ret if isinstance(finalize_ret, tuple) else (None, finalize_ret) + + if hook is not None: + if dbo_enabled(): + # If DBO is being used, register the hook with the ubatch + # context and call it in dbo_maybe_run_recv_hook instead of + # passing it to the receiver. + dbo_register_recv_hook(hook) + dbo_yield() + else: + hook() + + receiver() + + return output + + def apply( + self, + hidden_states: torch.Tensor, + w1: torch.Tensor, + w2: torch.Tensor, + topk_ids: torch.Tensor, + topk_weights: torch.Tensor, + activation: MoEActivation = MoEActivation.SILU, + global_num_experts: int = -1, + expert_map: torch.Tensor | None = None, + apply_router_weight_on_input: bool = False, + shared_experts: SharedExperts | None = None, + shared_experts_input: torch.Tensor | None = None, + ) -> torch.Tensor: + """ + This function computes a Mixture of Experts (MoE) layer using two sets + of weights, w1 and w2, and top-k gating mechanism. + + Parameters: + - hidden_states: (torch.Tensor): The input tensor to the MoE layer. + - w1 (torch.Tensor): The first set of expert weights. + - w2 (torch.Tensor): The second set of expert weights. + - topk_weights (torch.Tensor): The topk weights applied at the end of the layer. + - topk_ids (torch.Tensor): A map of row to expert id. + - activation (MoEActivation): The activation function to apply after the first + MoE layer. + - global_num_experts (int): The total number of experts in the global + expert space. + - expert_map (Optional[torch.Tensor]): A tensor mapping expert indices + from the global expert space to the local expert space of the expert + parallel shard. + - apply_router_weight_on_input (bool): When true, the topk weights are + applied directly on the inputs. This is only applicable when topk is + 1. + - shared_experts: SharedExperts | None. The shared experts if any. + - shared_experts_input (Optional[torch.Tensor]): Optional separate + input for shared experts. For latent MoE, this is the original + hidden_states before latent projection. + + Returns: + - torch.Tensor: The output tensor after applying the MoE layer. + """ + output = torch.empty_like(hidden_states) + + local_num_experts = w1.shape[0] + if global_num_experts == -1: + global_num_experts = local_num_experts + + a1q, a1q_scale, expert_tokens_meta, topk_ids, topk_weights = self._prepare( + hidden_states, + topk_weights, + topk_ids, + global_num_experts, + expert_map, + apply_router_weight_on_input, + ) + + # Stash the original unquantized hidden states on the LoRA context + # so apply_w13_lora sees correct-magnitude activations instead of + # the potentially quantized values produced by _prepare(). + lora_ctx = getattr(self.fused_experts, "_lora_context", None) + if lora_ctx is not None: + lora_ctx.original_hidden_states = hidden_states + + fused_out = self._fused_experts( + in_dtype=hidden_states.dtype, + a1q=a1q, + a1q_scale=a1q_scale, + w1=w1, + w2=w2, + topk_weights=topk_weights, + topk_ids=topk_ids, + activation=activation, + global_num_experts=global_num_experts, + local_num_experts=local_num_experts, + expert_map=expert_map, + apply_router_weight_on_input=apply_router_weight_on_input, + expert_tokens_meta=expert_tokens_meta, + output_alias=output, + ) + + if lora_ctx is not None: + lora_ctx.original_hidden_states = None + + return self._finalize( + output, + fused_out, + hidden_states, + topk_weights, + topk_ids, + apply_router_weight_on_input, + shared_experts=shared_experts, + shared_experts_input=shared_experts_input, + ) + + +@final +class FusedMoEKernelMonolithicImpl: + def __init__( + self, + prepare_finalize: FusedMoEPrepareAndFinalizeMonolithic, + fused_experts: FusedMoEExpertsMonolithic, + ): + self.prepare_finalize = prepare_finalize + self.fused_experts = fused_experts + + def apply( + self, + hidden_states: torch.Tensor, + w1: torch.Tensor, + w2: torch.Tensor, + router_logits: torch.Tensor, + activation: MoEActivation, + global_num_experts: int, + expert_map: torch.Tensor | None, + apply_router_weight_on_input: bool, + # grouped topk + fused topk bias parameters + num_expert_group: int | None = None, + e_score_correction_bias: torch.Tensor | None = None, + routed_scaling_factor: float | None = None, + topk_group: int | None = None, + ) -> torch.Tensor: + """ + Same as forward(), except uses router_logits as opposed + to the topk_ids and topk_weights. This is used for kernels + that have fused router + experts (e.g. FLASHINFER_TRTLLM). + """ + + a1q, a1q_scale, router_logits = self.prepare_finalize.prepare( + hidden_states, + router_logits=router_logits, + quant_config=self.fused_experts.quant_config, + defer_input_quant=self.fused_experts.expects_unquantized_inputs, + ) + + fused_out = self.fused_experts.apply( + hidden_states=a1q, + w1=w1, + w2=w2, + router_logits=router_logits, + activation=activation, + global_num_experts=global_num_experts, + expert_map=expert_map, + apply_router_weight_on_input=apply_router_weight_on_input, + a1q_scale=a1q_scale, + # grouped topk + fused topk bias parameters + num_expert_group=num_expert_group, + e_score_correction_bias=e_score_correction_bias, + routed_scaling_factor=routed_scaling_factor, + topk_group=topk_group, + ) + + output = self.prepare_finalize.finalize(fused_out) + + return output + + +@final +class FusedMoEKernel: + def __init__( + self, + prepare_finalize: FusedMoEPrepareAndFinalize, + fused_experts: FusedMoEExperts, + ): + super().__init__() + + # Initialize the implementation (monolithic or modular). + self.impl: FusedMoEKernelModularImpl | FusedMoEKernelMonolithicImpl + if isinstance(prepare_finalize, FusedMoEPrepareAndFinalizeModular) and isinstance( + fused_experts, FusedMoEExpertsModular + ): + self.impl = FusedMoEKernelModularImpl( + prepare_finalize, + fused_experts, + ) + + elif isinstance(prepare_finalize, FusedMoEPrepareAndFinalizeMonolithic) and isinstance( + fused_experts, FusedMoEExpertsMonolithic + ): + self.impl = FusedMoEKernelMonolithicImpl( + prepare_finalize, + fused_experts, + ) + + else: + raise ValueError( + "prepare_finalize and fused_experts must both be either monolithic " + f"or non-monolithic but got {prepare_finalize.__class__.__name__} " + f"and {fused_experts.__class__.__name__}" + ) + + self._post_init_setup() + + @property + def can_overlap_shared_experts(self) -> bool: + if isinstance(self.impl, FusedMoEKernelModularImpl): + return self.impl.prepare_finalize.supports_async() + else: + return False + + @property + def is_monolithic(self) -> bool: + return isinstance(self.impl, FusedMoEKernelMonolithicImpl) + + @property + def prepare_finalize(self) -> FusedMoEPrepareAndFinalize: + return self.impl.prepare_finalize + + @property + def fused_experts(self) -> FusedMoEExperts: + return self.impl.fused_experts + + @property + def moe_config(self) -> FusedMoEConfig: + return self.fused_experts.moe_config + + def supports_lora(self) -> bool: + return self.fused_experts.supports_lora() + + def _post_init_setup(self): + """ + Resolve any leftover setup dependencies between self.prepare_finalize + and self.fused_experts here. + """ + self.prepare_finalize.post_init_setup(self.impl.fused_experts) + assert self.prepare_finalize.activation_format == self.fused_experts.activation_format() + + def output_is_reduced(self) -> bool: + """ + Indicates whether or not the output of fused MoE kernel + is reduced across all ranks. + """ + return self.prepare_finalize.output_is_reduced() + + def apply_monolithic( + self, + hidden_states: torch.Tensor, + w1: torch.Tensor, + w2: torch.Tensor, + router_logits: torch.Tensor, + activation: MoEActivation, + global_num_experts: int, + expert_map: torch.Tensor | None, + apply_router_weight_on_input: bool, + # grouped topk + fused topk bias parameters + num_expert_group: int | None = None, + e_score_correction_bias: torch.Tensor | None = None, + routed_scaling_factor: float | None = None, + topk_group: int | None = None, + ) -> torch.Tensor: + assert isinstance(self.impl, FusedMoEKernelMonolithicImpl) + return self.impl.apply( + hidden_states=hidden_states, + w1=w1, + w2=w2, + router_logits=router_logits, + activation=activation, + global_num_experts=global_num_experts, + expert_map=expert_map, + apply_router_weight_on_input=apply_router_weight_on_input, + num_expert_group=num_expert_group, + e_score_correction_bias=e_score_correction_bias, + routed_scaling_factor=routed_scaling_factor, + topk_group=topk_group, + ) + + def apply( + self, + hidden_states: torch.Tensor, + w1: torch.Tensor, + w2: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + activation: MoEActivation, + global_num_experts: int, + expert_map: torch.Tensor | None, + apply_router_weight_on_input: bool, + shared_experts: SharedExperts | None = None, + shared_experts_input: torch.Tensor | None = None, + ) -> torch.Tensor: + assert isinstance(self.impl, FusedMoEKernelModularImpl) + return self.impl.apply( + hidden_states=hidden_states, + w1=w1, + w2=w2, + topk_weights=topk_weights, + topk_ids=topk_ids, + activation=activation, + global_num_experts=global_num_experts, + expert_map=expert_map, + apply_router_weight_on_input=apply_router_weight_on_input, + shared_experts=shared_experts, + shared_experts_input=shared_experts_input, + ) 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..84ff909b6f --- /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_vllm_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: + vllm_config = get_current_vllm_config() + if vllm_config.parallel_config.use_ubatching: + raise ValueError("K3 latent-MoE tail fusion does not support DBO or ubatching.") + if vllm_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.py.orig b/aphrodite/model_executor/layers/fused_moe/runner/moe_runner.py.orig new file mode 100644 index 0000000000..29a23254f8 --- /dev/null +++ b/aphrodite/model_executor/layers/fused_moe/runner/moe_runner.py.orig @@ -0,0 +1,951 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from collections.abc import Callable, Iterable +from contextlib import nullcontext +from typing import TYPE_CHECKING + +import torch +import torch.nn.functional as F + +from aphrodite.config import AphroditeConfig, get_current_aphrodite_config +from aphrodite.config.parallel import ExpertPlacementStrategy +from aphrodite.distributed import ( + get_ep_group, + get_pcp_group, + tensor_model_parallel_all_reduce, +) +from aphrodite.distributed.eplb.eplb_state import EplbLayerState +from aphrodite.forward_context import ( + ForwardContext, + get_forward_context, + is_forward_context_available, +) +from aphrodite.logger import init_logger +from aphrodite.model_executor.layers.fused_moe.activation import MoEActivation +from aphrodite.model_executor.layers.fused_moe.config import ( + FusedMoEConfig, +) +from aphrodite.model_executor.layers.fused_moe.fused_moe_method_base import ( + FusedMoEMethodBase, +) +from aphrodite.model_executor.layers.fused_moe.fused_moe_modular_method import ( + FusedMoEModularMethod, +) +from aphrodite.model_executor.layers.fused_moe.routed_experts import ( + RoutedExperts, +) +from aphrodite.model_executor.layers.fused_moe.router.fused_moe_router import ( + FusedMoERouter, +) +from aphrodite.model_executor.layers.fused_moe.router.zero_expert_router import ( + ZeroExpertRouter, +) +from aphrodite.model_executor.layers.fused_moe.runner.moe_runner_interface import ( + MoERunnerInterface, +) +from aphrodite.model_executor.layers.fused_moe.runner.shared_experts import ( + SharedExperts, + SharedExpertsOrder, +) +from aphrodite.platforms import current_platform +from aphrodite.utils.torch_utils import ( + _USE_LAYERNAME, + LayerName, + direct_register_custom_op, +) + +logger = init_logger(__name__) + + +def register_layer_for_moe_forward_op( + aphrodite_config: AphroditeConfig, + layer: "MoERunner", +): + # For smuggling this layer into the fused moe custom op + prefix = layer.layer_name + compilation_config = aphrodite_config.compilation_config + if prefix in compilation_config.static_forward_context: + raise ValueError("Duplicate layer name: {}".format(prefix)) + compilation_config.static_forward_context[prefix] = layer + compilation_config.static_all_moe_layers.append(prefix) + + +def get_layer_from_name(layer_name: str) -> MoERunnerInterface: + forward_context: ForwardContext = get_forward_context() + if not _USE_LAYERNAME and layer_name == "from_forward_context": + all_moe_layers = forward_context.all_moe_layers + assert all_moe_layers is not None + moe_layer_index = forward_context.moe_layer_index + if moe_layer_index >= len(all_moe_layers): + raise AssertionError( + "We expected the number of MOE layers in `all_moe_layers` " + "to be equal to the number of " + "{aphrodite.moe_forward, aphrodite.moe_forward_shared} calls." + ) + layer_name = all_moe_layers[moe_layer_index] + forward_context.moe_layer_index += 1 + layer = forward_context.no_compile_layers[layer_name] + assert isinstance(layer, MoERunnerInterface) + return layer + + +# On torch >= 2.11, layer_name is a hoisted LayerName opaque object; +# on older versions it remains a plain str. +if TYPE_CHECKING: + from typing import TypeAlias + + _layer_name_type: TypeAlias = str | LayerName +else: + _layer_name_type = LayerName if _USE_LAYERNAME else str + + +@torch.compiler.assume_constant_result +def _resolve_layer_name(layer_name: str | LayerName) -> str: + from torch._library.fake_class_registry import FakeScriptObject + + if isinstance(layer_name, LayerName): + return layer_name.value + elif isinstance(layer_name, FakeScriptObject): + return layer_name.real_obj.value + return layer_name + + +# Note: _moe_forward and _moe_forward_shared should not contain any +# implementation details, They should merely pass along control to +# the runner's '_forward_impl' method. +# These functions should never be called directly since they do not +# include all the functionality of the MoE layer. +def _moe_forward( + hidden_states: torch.Tensor, + router_logits: torch.Tensor, + shared_experts_input: torch.Tensor | None, + input_ids: torch.Tensor | None, + layer_name: _layer_name_type, + hidden_dim_unpadded: int, +) -> torch.Tensor: + layer = get_layer_from_name(_resolve_layer_name(layer_name)) + return layer._forward_impl( + hidden_states, + router_logits, + shared_experts_input, + input_ids, + ) + + +def _moe_forward_fake( + hidden_states: torch.Tensor, + router_logits: torch.Tensor, + shared_experts_input: torch.Tensor | None, + input_ids: torch.Tensor | None, + layer_name: _layer_name_type, + hidden_dim_unpadded: int, +) -> torch.Tensor: + # `hidden_dim_unpadded > 0` only on the TRT-LLM MXFP4 path, where the + # real kernel writes narrower than `hidden_states.shape[-1]`. Plumbed + # as an op arg (not peeked from the layer registry) to keep the fake + # a pure shape function of its inputs and preserve subgraph dedup. + if hidden_dim_unpadded > 0: + return hidden_states.new_empty((*hidden_states.shape[:-1], hidden_dim_unpadded)) + return torch.empty_like(hidden_states) + + +def _moe_forward_shared( + hidden_states: torch.Tensor, + router_logits: torch.Tensor, + shared_experts_input: torch.Tensor | None, + input_ids: torch.Tensor | None, + layer_name: _layer_name_type, + hidden_dim_unpadded: int, +) -> tuple[torch.Tensor, torch.Tensor]: + layer = get_layer_from_name(_resolve_layer_name(layer_name)) + return layer._forward_impl( + hidden_states, + router_logits, + shared_experts_input, + input_ids, + ) + + +def _moe_forward_shared_fake( + hidden_states: torch.Tensor, + router_logits: torch.Tensor, + shared_experts_input: torch.Tensor | None, + input_ids: torch.Tensor | None, + layer_name: _layer_name_type, + hidden_dim_unpadded: int, +) -> tuple[torch.Tensor, torch.Tensor]: + # `fused_out`: see `_moe_forward_fake` for hidden_dim_unpadded semantics. + # `shared_out`: matches `shared_experts_input` if provided (latent MoE), + # else `hidden_states`. + if hidden_dim_unpadded > 0: + fused_out = hidden_states.new_empty((*hidden_states.shape[:-1], hidden_dim_unpadded)) + else: + fused_out = torch.empty_like(hidden_states) + if shared_experts_input is not None: + shared_out = torch.empty_like(shared_experts_input) + else: + shared_out = torch.empty_like(hidden_states) + return shared_out, fused_out + + +# NOTE: `moe_forward` and `moe_forward_shared` being opaque custom ops is a +# load-bearing assumption for the MoE-LoRA dual-stream path. +direct_register_custom_op( + op_name="moe_forward", + op_func=_moe_forward, + mutates_args=["hidden_states"], + fake_impl=_moe_forward_fake, + tags=(torch.Tag.needs_fixed_stride_order,), +) + + +direct_register_custom_op( + op_name="moe_forward_shared", + op_func=_moe_forward_shared, + fake_impl=_moe_forward_shared_fake, + tags=(torch.Tag.needs_fixed_stride_order,), +) + + +def _unpack( + result: torch.Tensor | tuple[torch.Tensor, torch.Tensor], +) -> tuple[torch.Tensor | None, torch.Tensor]: + if isinstance(result, tuple): + return result + else: + return (None, result) + + +class MoERunner(MoERunnerInterface): + """ + Standard MoE runner implementation for executing Mixture of Experts layers. + + This is the primary concrete implementation of MoE execution logic, providing + comprehensive support for standard MoE operations. It handles: + - Expert routing and token dispatching using various routing strategies + - Shared experts computation with optional parallel execution using CUDA streams + - Tensor model parallel and expert parallel operations + - Multiple quantization methods and optimized kernel selection + - Both monolithic and decomposed expert execution paths + - Integration with various parallel execution modes (TP, EP, DP) + + The runner orchestrates the complete MoE forward pass including routing tokens + to experts, executing expert computations in parallel, and combining results. + It supports advanced features like overlapped execution of shared experts, + optimized kernels for different parallel configurations, and seamless + integration with Aphrodite's distributed execution framework. + + Eventually, this class may be split into more specialized implementations + for different configurations (e.g., with/without shared experts, gates, etc.). + """ + + def __init__( + self, + layer_name: str, + moe_config: FusedMoEConfig, + router: FusedMoERouter, + routed_experts: RoutedExperts, + enable_dbo: bool = False, + gate: torch.nn.Module | None = None, + shared_experts: torch.nn.Module | None = None, + shared_expert_gate: torch.nn.Module | None = None, + routed_input_transform: torch.nn.Module | None = None, + routed_output_transform: torch.nn.Module | None = None, + routed_scaling_factor: float = 1.0, + ): + super().__init__() + self.moe_config = moe_config + self.router = router + self.routed_input_transform = routed_input_transform + self.routed_output_transform = routed_output_transform + self.routed_scaling_factor = routed_scaling_factor + self.gate = gate + self.shared_expert_gate = shared_expert_gate + self.routed_experts = routed_experts + self.enable_dbo = enable_dbo + + # When both gates are present and FSE is enabled, fuse their + # weight matrices into [num_experts + num_shared, hidden] so one + # F.linear produces combined logits. The topk kernel can then + # apply routing softmax and shared expert activation (sigmoid) + # in a single launch. + self._fse_fuse_gate = gate is not None and shared_expert_gate is not None + self._combined_gate_weight: torch.Tensor | None = None + + self._shared_experts: SharedExperts | None = None + if shared_experts is not None: + can_overlap = lambda: self._quant_method.mk_can_overlap_shared_experts + self._shared_experts = SharedExperts( + shared_experts, + moe_config=moe_config, + enable_dbo=enable_dbo, + mk_can_overlap_shared_experts=can_overlap, + ) + + # Needed for string -> MoERunner layer lookup in custom ops. + self.layer_name = layer_name + + self._forward_entry = self._select_forward() + + # For smuggling this layer into the fused moe custom op + register_layer_for_moe_forward_op(get_current_aphrodite_config(), self) + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> Iterable[str]: + return self.routed_experts.load_weights(weights) + + def _select_forward(self) -> Callable: + if current_platform.is_tpu() or current_platform.is_cpu(): + # TODO: Once the OOM issue for the TPU backend is resolved, we + # will switch to using the moe_forward custom op. + # Note: CPU doesn't require wrapped _forward_impl. + return _moe_forward if self._shared_experts is None else _moe_forward_shared + + return ( + torch.ops.aphrodite.moe_forward if self._shared_experts is None else torch.ops.aphrodite.moe_forward_shared + ) + + @property + def shared_experts(self) -> SharedExperts | None: + return self._shared_experts + + @property + def is_internal_router(self) -> bool: + return self.gate is not None + + # TODO(bnell): Temporary hack. Get rid of this. + def _replace_quant_method(self, quant_method: FusedMoEMethodBase): + self.routed_experts._replace_quant_method(quant_method) + + # TODO(bnell): Hack for elastic_ep. Get rid of this + def _set_moe_config(self, new_moe_config: FusedMoEConfig): + self.moe_config = new_moe_config + self.routed_experts._set_moe_config(new_moe_config) + if self._shared_experts is not None: + self._shared_experts._set_moe_config(new_moe_config) + + def _maybe_fuse_gate_weights(self): + """Fuse router and shared expert gate weights on first call. + + Cannot be done at __init__ because gate weights are loaded after + module construction (via weight_loader). Called once from + _forward_impl before the first forward pass. + """ + if self._combined_gate_weight is None: + assert self.gate is not None and self.shared_expert_gate is not None + self._combined_gate_weight = torch.cat( + [self.gate.weight, self.shared_expert_gate.weight], + dim=0, + ) + + @property + 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]: + """Apply transform for routed experts (e.g., latent projection). + + This is called by FusedMoE.forward_native. The original hidden_states + is saved separately so shared experts get [S, hidden_size] while + routed experts get the transformed [S, moe_latent_size]. + + Returns (possibly transformed) hidden states and the input for shared + experts (or None if there are no shared experts). + """ + if self.routed_input_transform is not None: + result = self.routed_input_transform(hidden_states) + # ReplicatedLinear returns (output, extra_bias) tuple. + # We only need the output tensor; extra_bias is not used here. + if isinstance(result, tuple): + return result[0], hidden_states + return result, hidden_states + + return ( + hidden_states, + hidden_states if self._shared_experts is not None else None, + ) + + def apply_routed_output_transform( + self, + fused_output: torch.Tensor, + ) -> torch.Tensor: + """Apply transform to routed expert output (e.g., latent to full dim). + + Used by latent MoE models (e.g., NemotronH) where routed experts + operate in a compressed latent space and need projection back to + the full hidden dimension before combining with shared expert output. + """ + if self.routed_output_transform is not None: + r = self.routed_output_transform(fused_output) + fused_output = r[0] if isinstance(r, tuple) else r + return fused_output + + def _maybe_apply_routed_scale_to_output( + self, + shared_output: torch.Tensor | None, + fused_output: torch.Tensor, + ) -> tuple[torch.Tensor | None, torch.Tensor]: + """Apply routed_scaling_factor to the output with FP16 overflow + protection. + + Scale the fused expert output by routed_scaling_factor. For FP16, + avoid overflow by dividing shared_output by the scale instead + (the decoder layer compensates with matching divisions). + """ + if self.routed_scaling_factor != 1.0: + if fused_output.dtype != torch.float16 or shared_output is None: + fused_output *= self.routed_scaling_factor + elif shared_output is not None: + shared_output *= 1.0 / self.routed_scaling_factor + return shared_output, fused_output + + @property + def _fused_output_is_reduced(self) -> bool: + return self._quant_method.moe_kernel is not None and self._quant_method.moe_kernel.output_is_reduced() + + 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. + + * If the combine kernel does the reduction for fused_output, reduce + shared_output separately. O.w, reduce fused_output+shared_output later. + * 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: + shared_output = tensor_model_parallel_all_reduce(shared_output) + return shared_output + + def _maybe_reduce_final_output( + self, + states: torch.Tensor, + trunc_size: int | None, + ) -> torch.Tensor: + """All-reduce the combined output if needed. + + This is the "late" all-reduce path. When neither fused nor shared + output was individually reduced, the combined sum is all-reduced + here. Skipped when sequence-parallel is active (SP handles its + own reduction) or when the early path already reduced both outputs. + """ + # skip_final_all_reduce must not coexist with a pre-reduced fused + # output. This should be enforced by MoE config initialization. + if self.moe_config.skip_final_all_reduce: + assert not self._fused_output_is_reduced, "skip_final_all_reduce requires an un-reduced fused 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 output_is_reduced + ): + states = tensor_model_parallel_all_reduce(states) + + return states[..., :trunc_size] if trunc_size is not None else states + + def _encode_layer_name(self) -> str | LayerName: + if _USE_LAYERNAME: + return LayerName(self.layer_name) + # Can be unavailable or None in unittests + if is_forward_context_available() and get_forward_context().all_moe_layers is not None: + return "from_forward_context" + return self.layer_name + + def _maybe_pad_hidden_states( + self, + shared_experts_input: torch.Tensor | None, + hidden_states: torch.Tensor, + ) -> tuple[torch.Tensor, int | None, int | None]: + """Pad hidden_states to moe_config.hidden_dim and compute the + original dimension for later truncation. + + For latent MoE, the routed hidden_states may be smaller than + hidden_dim. Padding ensures uniform tensor sizes through the + fused MoE kernel. The returned trunc_size is used by + _maybe_reduce_final_output to strip the padding from the result. + """ + shared_experts_hidden_dim = shared_experts_input.shape[-1] if shared_experts_input is not None else 0 + transformed_hidden_dim: int | None = hidden_states.shape[-1] + if not self._quant_method.skip_forward_padding and self.moe_config.hidden_dim != transformed_hidden_dim: + assert transformed_hidden_dim is not None + hidden_states = F.pad( + hidden_states, + (0, self.moe_config.hidden_dim - transformed_hidden_dim), + mode="constant", + value=0.0, + ) + + # Truncation sizes for stripping kernel padding from the output. + # None means no truncation needed (no padding was applied). + # + # Two truncation points exist in forward(): + # pre_xform: applied to fused_output BEFORE routed_output_transform + # post_xform: applied to the final result AFTER all-reduce + # + # MoE with routed output transform or shared experts: + # - pre_xform applies if the transform needs unpadded routed output + # or shared+routed add needs matching hidden dims. For Nemotron-3 + # Nano, TRTLLM NVFP4 pads routed MoE hidden dim 2688->2816, while + # shared output stays 2688. + # - post_xform uses shared_experts_hidden_dim when transform and shared + # experts make the final output full hidden dim. + # + # Standard MoE / MoE without transforms (GPT-OSS, Mixtral): + # - pre_xform is None (no early truncation) + # - post_xform strips padding after all-reduce (or None if unpadded) + if transformed_hidden_dim == hidden_states.shape[-1]: + transformed_hidden_dim = None + + pre_xform_trunc_size = None + if self.routed_output_transform is not None or shared_experts_hidden_dim > 0: + pre_xform_trunc_size = transformed_hidden_dim + post_xform_trunc_size = transformed_hidden_dim + if self.routed_output_transform is not None and shared_experts_hidden_dim > 0: + post_xform_trunc_size = shared_experts_hidden_dim + + return hidden_states, pre_xform_trunc_size, post_xform_trunc_size + + def _maybe_apply_shared_experts( + self, + shared_experts_input: torch.Tensor | None, + order: SharedExpertsOrder, + ): + if self._shared_experts is not None: + assert shared_experts_input is not None + self._shared_experts(shared_experts_input, order) + + def _apply_quant_method( + self, + hidden_states: torch.Tensor, + router_logits: torch.Tensor, + shared_experts_input: torch.Tensor | None, + input_ids: torch.Tensor | None = None, + ) -> tuple[torch.Tensor | None, torch.Tensor]: + """Run expert routing and the fused MoE kernel via the quant method. + + Orchestrates shared expert execution (before/after), expert selection + via the router, and the actual fused MoE computation. Returns + (shared_expert_output, fused_expert_output). + """ + self._maybe_apply_shared_experts(shared_experts_input, SharedExpertsOrder.NO_OVERLAP) + + if self.routed_experts.quant_method.is_monolithic: + # Monolithic kernels: pass router_logits to routed_experts + fused_out = self.routed_experts.forward_monolithic( + x=hidden_states, + router_logits=router_logits, + input_ids=input_ids, + ) + else: + # Modular kernels: select experts first, then call routed_experts + topk_weights, topk_ids = self.router.select_experts( + hidden_states=hidden_states, + router_logits=router_logits, + topk_indices_dtype=self._quant_method.topk_indices_dtype, + input_ids=input_ids, + ) + + fused_out = self.routed_experts.forward_modular( + x=hidden_states, + topk_weights=topk_weights, + topk_ids=topk_ids, + shared_experts=self._shared_experts, + shared_experts_input=shared_experts_input, + ) + + self._maybe_apply_shared_experts( + shared_experts_input, + SharedExpertsOrder.MULTI_STREAM_OVERLAPPED, + ) + + return ( + self._shared_experts.output if self._shared_experts is not None else None, + fused_out, + ) + + def _sequence_parallel_context(self): + """Return a context manager for sequence-parallel token + redistribution. + + When sequence parallelism is active, returns a context that handles + local size tracking for proper token scatter/gather. Otherwise + returns a no-op context. + """ + ctx = get_forward_context() + return ctx.dp_metadata.sp_local_sizes(self.moe_config.sp_size) if ctx.dp_metadata else nullcontext() + + def _maybe_sync_shared_experts_stream( + self, + shared_experts_input: torch.Tensor | None, + ): + # If router/gate provided, then apply it here. + # (Note: This code runs only when "overlapped mode" is on to allow + # parallel execution of shared experts with the FusedMoE via + # separate cuda stream) + if self._shared_experts is not None: + assert shared_experts_input is not None + self._shared_experts.maybe_sync_shared_experts_stream(shared_experts_input) + + def _maybe_add_zero_expert_output( + self, + result: torch.Tensor, + ) -> torch.Tensor: + """Add the zero expert's contribution to the final result. + + When a ZeroExpertRouter is used, it computes a bias-like output + from the "zero expert" that is added to the combined routed+shared + expert output. + """ + if isinstance(self.router, ZeroExpertRouter): + zero_expert_output = self.router.zero_expert_output + assert zero_expert_output is not None + result = result + zero_expert_output + return result + + 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: + """Invoke the fused moe layer. + + Input: + - hidden_states + - router_logits + + Output: + - The new hidden_states. + + Calling sequence + - forward + - self._forward_entry (_moe_forward or _moe_forward_shared custom op) + - _forward_impl + + Note: The existence of _moe_forward and _moe_forward_shared custom ops are due + to the following reason: + 1. pytorch cannot handle union types in custom op signatures so + _moe_forward and _moe_forward_shared must be split. + """ + + # Apply transform for routed experts (e.g., latent projection + # for latent MoE) + 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` + # so routed output can be trimmed before + # shared+routed add / latent up proj if needed. + + 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, + ) + + # + # Note: there are two all-reduce points below. They are mutually + # exclusive, controlled by _fused_output_is_reduced + # - When True: the combine kernel already reduced fused_output, + # so we reduce shared_output here to match, then skip the + # all-reduce in _maybe_reduce_final_output. + # - When False: neither output is reduced yet, so we combine + # them first and all-reduce the sum in _maybe_reduce_final_output. + + # Extract outputs from result + shared_output, fused_output = _unpack(result) + + if og_hidden_dim_pre_xform is not None: + fused_output = fused_output[..., :og_hidden_dim_pre_xform] + + 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, fused_output_is_reduced + ) + + shared_output, fused_output = self._maybe_apply_routed_scale_to_output(shared_output, fused_output) + + # Apply output transform (e.g. latent -> full dim) + fused_output = self.apply_routed_output_transform(fused_output) + + if shared_output is not None: + result = shared_output + fused_output + else: + result = fused_output + + result = self._maybe_reduce_final_output( + result, og_hidden_dim_post_xform, fused_output_is_reduced + ) + + return self._maybe_add_zero_expert_output(result) + + @property + def do_naive_dispatch_combine(self) -> bool: + return ( + self.moe_config.dp_size > 1 or self.moe_config.is_sequence_parallel + ) and not self._quant_method.supports_internal_mk + + def _maybe_dispatch( + self, + hidden_states: torch.Tensor, + router_logits: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + # For naive dispatch/combine Dp/Ep, dispatch the hidden states and + # router logits to all experts. + # NOTE: this will be removed once all kernels are migrated into the + # MoEKernel framework. + if self.do_naive_dispatch_combine: + result = get_ep_group().dispatch_router_logits( + hidden_states, + router_logits, + self.moe_config.is_sequence_parallel, + ) + assert len(result) == 2 + hidden_states, router_logits = result + + if self.moe_config.pcp_size > 1 and not self.moe_config.moe_parallel_config.use_all2all_kernels: + hidden_states = get_pcp_group().all_gather(hidden_states, dim=0) + router_logits = get_pcp_group().all_gather(router_logits, dim=0) + + return hidden_states, router_logits + + def _maybe_combine( + self, + shared_output: torch.Tensor | None, + hidden_states: torch.Tensor, + ) -> torch.Tensor | tuple[torch.Tensor | None, torch.Tensor]: + if self.do_naive_dispatch_combine: + hidden_states = get_ep_group().combine(hidden_states, self.moe_config.is_sequence_parallel) + + if self.moe_config.pcp_size > 1 and not self.moe_config.moe_parallel_config.use_all2all_kernels: + hidden_states = get_pcp_group().reduce_scatter(hidden_states, dim=0) + + if self.shared_experts is not None: + assert shared_output is not None + return shared_output, hidden_states + else: + return hidden_states + + def _forward_impl( + self, + hidden_states: torch.Tensor, + router_logits: torch.Tensor, + shared_experts_input: torch.Tensor | None, + input_ids: torch.Tensor | None = None, + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + """Entry point called by the custom op to run the MoE computation. + + Handles pre-dispatch setup (gate application, external shared expert + triggering, quant config init) then performs the following steps + within the sequence-parallel context. + + - Performs expert routing + - fused MoE kernel execution + - shared expert computation. + + Returns a single tensor of combined fused and shared output (if present). + """ + # TODO(bnell): this can be removed after MK migration is complete. + self.routed_experts._ensure_moe_quant_config_init() + + # Sync aux and main stream for shared expert multi-stream overlap. + self._maybe_sync_shared_experts_stream(shared_experts_input) + + # If the Runner holds the gate, apply it after the stream sync, + # so it can run overlapped with the + # NOTE: in future PR, MoE runner will always hold the gate. + if self.gate is not None: + if self._fse_fuse_gate: + self._maybe_fuse_gate_weights() + router_logits = F.linear(hidden_states, self._combined_gate_weight) + else: + router_logits, _ = self.gate(hidden_states) + + with self._sequence_parallel_context(): + # TODO(bnell): parts of the dispatch/combine steps will go away once + # #32567 lands and the remaining kernels are made MKs. The PCP + # code will probably remain + hidden_states, router_logits = self._maybe_dispatch( + hidden_states, + router_logits, + ) + + shared_output, hidden_states = self._apply_quant_method( + hidden_states=hidden_states, + router_logits=router_logits, + shared_experts_input=shared_experts_input, + input_ids=input_ids, + ) + + return self._maybe_combine( + shared_output, + hidden_states, + ) + + ######################################################### + # + # Old methods from FusedMoE layer. Remove when possible. + # + ######################################################### + + # Note: maybe_init_modular_kernel should only be called by + # prepare_communication_buffer_for_model. + # This is called after all weight loading and post-processing, so it + # should be safe to swap out the quant_method. + def maybe_init_modular_kernel(self) -> None: + # NOTE(rob): WIP refactor. For quant methods that own the MK + # we create the MK during process_weights_after_loading. + if self.routed_experts.quant_method.supports_internal_mk or self.routed_experts.quant_method.is_monolithic: + return None + + self.routed_experts._ensure_moe_quant_config_init() + # routing_tables only needed for round-robin expert placement with + # DeepEP all2all backend. + routing_tables = self._expert_routing_tables() + + if isinstance(self.routed_experts.quant_method, FusedMoEModularMethod): + base_quant_method = self.routed_experts.quant_method.old_quant_method + else: + base_quant_method = self.routed_experts.quant_method + + prepare_finalize = base_quant_method.maybe_make_prepare_finalize(routing_tables=routing_tables) + if prepare_finalize is not None: + logger.debug("%s for %s(%s)", prepare_finalize.__class__.__name__, self, id(self)) + self._replace_quant_method( + FusedMoEModularMethod.make( + self.routed_experts, + base_quant_method, + prepare_finalize, + ) + ) + + # + # Properties + # + + @property + def layer_id(self): + # Delayed import to avoid circular dependency + from aphrodite.model_executor.models.utils import extract_layer_index + + return extract_layer_index(self.layer_name) + + # + # Attributes still needed by models + # + + @property + def is_monolithic(self) -> bool: + return self.routed_experts.quant_method.is_monolithic + + @property + def activation(self) -> MoEActivation: + return self.routed_experts.activation + + # + # Expert maps + # + + @property + def expert_map_manager(self): + """Forward to routed_experts.expert_map_manager for backward compatibility.""" + return self.routed_experts.expert_map_manager + + @property + def expert_placement_strategy(self) -> ExpertPlacementStrategy: + return self.expert_map_manager.placement_strategy + + @property + def expert_global_to_physical(self) -> torch.Tensor | None: + tables = self.expert_map_manager.routing_tables + return tables[0] if tables else None + + @property + def expert_physical_to_global(self) -> torch.Tensor | None: + """Routing table: physical expert ID to global expert ID.""" + tables = self.expert_map_manager.routing_tables + return tables[1] if tables else None + + @property + def expert_local_to_global(self) -> torch.Tensor | None: + """Routing table: local expert ID to global expert ID.""" + tables = self.expert_map_manager.routing_tables + return tables[2] if tables else None + + @property + def expert_map(self) -> torch.Tensor | None: + return self.routed_experts.expert_map + + def _expert_routing_tables( + self, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None: + return self.routed_experts._expert_routing_tables() + + def update_expert_map(self): + self.routed_experts.update_expert_map() + + def _map_global_expert_id_to_local_expert_id(self, expert_id: int) -> int: + """Map global expert ID to local expert ID.""" + return self.routed_experts._map_global_expert_id_to_local_expert_id(expert_id) + + def get_expert_weights(self) -> Iterable[torch.Tensor]: + return self.routed_experts.get_expert_weights() + + # + # EPLB + # + + @property + def eplb_state(self) -> EplbLayerState | None: + return self.router.eplb_state + + def set_eplb_state( + self, + moe_layer_idx: int, + expert_load_view: torch.Tensor, + logical_to_physical_map: torch.Tensor, + logical_replica_count: torch.Tensor, + ) -> None: + """ + Register the EPLB state in this layer. + + This is used later in forward pass, where we get the expert mapping + and record the load metrics in `expert_load_view`. + """ + if self.router.eplb_state is not None: + self.router.eplb_state.set_layer_state( + moe_layer_idx, + expert_load_view, + logical_to_physical_map, + logical_replica_count, + ) 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/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/causal_conv1d.py.orig b/aphrodite/model_executor/layers/mamba/ops/causal_conv1d.py.orig new file mode 100644 index 0000000000..412744f7f1 --- /dev/null +++ b/aphrodite/model_executor/layers/mamba/ops/causal_conv1d.py.orig @@ -0,0 +1,1178 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +# Copyright (c) 2024, Tri Dao. +# Adapted from https://github.com/Dao-AILab/causal-conv1d/blob/main/causal_conv1d/causal_conv1d_interface.py + + +import numpy as np +import torch + +from aphrodite.triton_utils import tl, triton +from aphrodite.v1.attention.backends.utils import NULL_BLOCK_ID, PAD_SLOT_ID + + +@triton.jit() +def _causal_conv1d_fwd_kernel( # continuous batching + # Pointers to matrices + x_ptr, # (dim, cu_seqlen) holding `batch` of actual sequences + padded sequences + w_ptr, # (dim, width) + bias_ptr, + initial_states_ptr, # conv_states_ptr + cache_indices_ptr, # (batch, n_blocks + padding) The second dimension contains + # the block indices relevant for each sequence + # plus potential 0-padding at the beginning and at the end + has_initial_states_ptr, + query_start_loc_ptr, + batch_ptr, + token_chunk_offset_ptr, + block_idx_first_scheduled_token, # (batch,) + block_idx_last_scheduled_token, # (batch,) + initial_state_idx, # (batch,) + num_computed_tokens, # (batch,) + o_ptr, # (dim, seqlen) - actually pointing to x_ptr + # Matrix dimensions + dim: tl.constexpr, + num_cache_lines: tl.constexpr, # added to support Aphrodite larger cache lines + # Strides + stride_x_dim: tl.constexpr, # stride to get to next feature-value, + stride_x_token: tl.int64, # stride to get to next token (same feature-index, same sequence-index) + stride_w_dim: tl.constexpr, # stride to get to next dim-axis value + stride_w_width: tl.constexpr, # stride to get to next width-axis value + stride_istate_seq: tl.constexpr, + stride_istate_dim: tl.constexpr, + stride_istate_token: tl.constexpr, + stride_cache_indices: tl.constexpr, + stride_o_dim: tl.constexpr, + stride_o_token: tl.int64, + stride_block_m: tl.constexpr, # Stride block to align divided by BLOCK_M + # others + pad_slot_id: tl.constexpr, + null_block_id: tl.constexpr, + # Meta-parameters + HAS_BIAS: tl.constexpr, + KERNEL_WIDTH: tl.constexpr, + SILU_ACTIVATION: tl.constexpr, + IS_APC_ENABLED: tl.constexpr, + HAS_NULL_BLOCK: tl.constexpr, + NP2_STATELEN: tl.constexpr, + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr, +): + conv_states_ptr = initial_states_ptr + conv_state_indices_ptr = cache_indices_ptr + stride_conv_state_seq = stride_istate_seq + stride_conv_state_dim = stride_istate_dim + stride_conv_state_tok = stride_istate_token + state_len = KERNEL_WIDTH - 1 # can be passed via argument if it's not the same as this value + + # one program handles one chunk in a single sequence + # rather than mixing sequences - to make updating initial_states across sequences efficiently + + # single-sequence id + idx_seq = tl.load(batch_ptr + tl.program_id(0)).to(tl.int64) + chunk_offset = tl.load(token_chunk_offset_ptr + tl.program_id(0)) + + # BLOCK_N elements along the feature-dimension (channel) + idx_feats = tl.program_id(1) * BLOCK_N + tl.arange(0, BLOCK_N) + + if idx_seq == pad_slot_id: + return + + sequence_start_index = tl.load(query_start_loc_ptr + idx_seq) + sequence_end_index = tl.load(query_start_loc_ptr + idx_seq + 1) + # find the actual sequence length + seqlen = sequence_end_index - sequence_start_index + + B_size: tl.constexpr = stride_block_m * BLOCK_M + + if IS_APC_ENABLED: + # Handle the case if prefix caching is enabled. + # In particular, if prefix caching is enabled, the program write additional cache states to "cache_indices_ptr" + + # Get the length of the completed sequence so far and compute the offset. + current_first_index = tl.load(block_idx_first_scheduled_token + idx_seq) + current_last_index = tl.load(block_idx_last_scheduled_token + idx_seq) + sequence_completed_index = tl.load(num_computed_tokens + idx_seq) + + # Compute the offset where the first stride_block_m-aligned first full block is + # Value in "token-space" + sequence_completed_offset_token = sequence_completed_index % B_size + seq_completed_offset = B_size - sequence_completed_offset_token + seq_end_offset = (seqlen - seq_completed_offset) % B_size + last_full_block_token_index = sequence_end_index - seq_end_offset + # If the sequence without the sequence_offset_index is stride_cache_chunk-aligned, then the last full chunk is the second-to-last one + if seq_end_offset == 0: + last_full_block_token_index = last_full_block_token_index - B_size + + # Get the number of blocks to be filled for the current sequence + # If n_block_to_fill = 0, then only the state at the sequence end is stored + n_block_to_fill = current_last_index - current_first_index + + # Get the index of the init block + conv_state_init_index = tl.load(initial_state_idx + idx_seq) + else: + n_block_to_fill = 0 + current_last_index = 0 + conv_state_init_index = 0 + current_first_index = 0 + last_full_block_token_index = 0 + + token_offset = BLOCK_M * chunk_offset + segment_len = min(BLOCK_M, seqlen - token_offset) + + # base of the sequence + x_base = x_ptr + sequence_start_index * stride_x_token + idx_feats * stride_x_dim # [BLOCK_N,] + + # cache_idx + conv_states_input_coord = tl.load( + conv_state_indices_ptr + idx_seq * stride_cache_indices + conv_state_init_index + ).to(tl.int64) + + if HAS_NULL_BLOCK: # noqa + if conv_states_input_coord == null_block_id: + # not processing as this is a null block (padding) + return + conv_states_base = ( + conv_states_ptr + (conv_states_input_coord * stride_conv_state_seq) + (idx_feats * stride_conv_state_dim) + ) # [BLOCK_N,] + + w_base = w_ptr + (idx_feats * stride_w_dim) # [BLOCK_N,] + + # Does 2 things: + # 1. READ prior-block init-state data - [done by every Triton programs] + # 2. update conv_state with new data [only by the Triton program handles chunk_offset=0] + if chunk_offset == 0: + # read from conv_states + load_init_state = tl.load(has_initial_states_ptr + idx_seq).to(tl.int1) + if load_init_state: + # load from conv_states + prior_tokens = conv_states_base + (state_len - 1) * stride_conv_state_tok + mask_w = idx_feats < dim + if KERNEL_WIDTH == 2: + conv_states_ptrs = prior_tokens # [BLOCK_N] + col0 = tl.load(conv_states_ptrs, mask_w, 0.0) + if KERNEL_WIDTH == 3: + conv_states_ptrs = prior_tokens # [BLOCK_N] + col1 = tl.load(conv_states_ptrs, mask_w, 0.0) + conv_states_ptrs = prior_tokens - 1 * stride_conv_state_tok # [BLOCK_N] + col0 = tl.load(conv_states_ptrs, mask_w, 0.0) + if KERNEL_WIDTH == 4: + conv_states_ptrs = prior_tokens # [BLOCK_N] + col2 = tl.load(conv_states_ptrs, mask_w, 0.0) + conv_states_ptrs = prior_tokens - 1 * stride_conv_state_tok # [BLOCK_N] + col1 = tl.load(conv_states_ptrs, mask_w, 0.0) + conv_states_ptrs = prior_tokens - 2 * stride_conv_state_tok # [BLOCK_N] + col0 = tl.load(conv_states_ptrs, mask_w, 0.0) + if KERNEL_WIDTH == 5: + conv_states_ptrs = prior_tokens # [BLOCK_N] + col3 = tl.load(conv_states_ptrs, mask_w, 0.0) + conv_states_ptrs = prior_tokens - 1 * stride_conv_state_tok # [BLOCK_N] + col2 = tl.load(conv_states_ptrs, mask_w, 0.0) + conv_states_ptrs = prior_tokens - 2 * stride_conv_state_tok # [BLOCK_N] + col1 = tl.load(conv_states_ptrs, mask_w, 0.0) + conv_states_ptrs = prior_tokens - 3 * stride_conv_state_tok # [BLOCK_N] + col0 = tl.load(conv_states_ptrs, mask_w, 0.0) + else: + # prior-tokens are zeros + if KERNEL_WIDTH >= 2: # STRATEGY1 + # first chunk and does not have prior-token, so just set to 0 + col0 = tl.zeros((BLOCK_N,), dtype=x_ptr.dtype.element_ty) + if KERNEL_WIDTH >= 3: # STRATEGY1 + col1 = tl.zeros((BLOCK_N,), dtype=x_ptr.dtype.element_ty) + if KERNEL_WIDTH >= 4: # STRATEGY1 + col2 = tl.zeros((BLOCK_N,), dtype=x_ptr.dtype.element_ty) + if KERNEL_WIDTH >= 5: # STRATEGY1 + col3 = tl.zeros((BLOCK_N,), dtype=x_ptr.dtype.element_ty) + + # STEP 2: + # here prepare data for updating conv_state + if state_len <= seqlen: # SMALL_CACHE=True (only move part of 'x' into conv_state cache) + # just read from 'x' + # copy 'x' data to conv_state + # load only 'x' data (and set 0 before 'x' if seqlen < state_len) + idx_tokens_last = (seqlen - state_len) + tl.arange(0, NP2_STATELEN) # [BLOCK_M] + x_ptrs = ( + x_ptr + + ((sequence_start_index + idx_tokens_last) * stride_x_token)[:, None] + + (idx_feats * stride_x_dim)[None, :] + ) # [BLOCK_M,BLOCK_N,] + mask_x = ( + (idx_tokens_last >= 0)[:, None] & (idx_tokens_last < seqlen)[:, None] & (idx_feats < dim)[None, :] + ) # token-index # token-index # feature-index + loaded_x = tl.load(x_ptrs, mask_x, 0.0) + idx_tokens_conv = tl.arange(0, NP2_STATELEN) # [BLOCK_M] + + # Compute the offset where the last block should be written in the conv_states + conv_states_output_coord = tl.load( + conv_state_indices_ptr + idx_seq * stride_cache_indices + current_last_index + ).to(tl.int64) + + conv_states_ptrs_target = ( + conv_states_ptr + + (conv_states_output_coord * stride_conv_state_seq) # Offset from seq + + (idx_feats * stride_conv_state_dim) + )[None, :] + ( # [BLOCK_N,] + idx_tokens_conv * stride_conv_state_tok + )[:, None] + + mask = (idx_tokens_conv < state_len)[:, None] & (idx_feats < dim)[None, :] + tl.debug_barrier() # NOTE: use this due to bug in Triton compiler + tl.store(conv_states_ptrs_target, loaded_x, mask) + + else: + if load_init_state: + # update conv_state by shifting left, i.e. take last few cols from conv_state + cols from 'x' + idx_tokens_conv = tl.arange(0, NP2_STATELEN) # [BLOCK_M] + + conv_states_ptrs_source = ( + conv_states_ptr + + (conv_states_input_coord * stride_conv_state_seq) + + (idx_feats * stride_conv_state_dim)[None, :] + + ((idx_tokens_conv + seqlen) * stride_conv_state_tok)[:, None] + ) # [BLOCK_M, BLOCK_N] + mask = ( + (conv_states_input_coord < num_cache_lines) + & ((idx_tokens_conv + seqlen) < state_len)[:, None] + & (idx_feats < dim)[None, :] + ) + conv_state = tl.load(conv_states_ptrs_source, mask, other=0.0) + + VAL = state_len - seqlen + + x_ptrs = x_base[None, :] + ((idx_tokens_conv - VAL) * stride_x_token)[:, None] # [BLOCK_M, BLOCK_N] + + mask_x = ( + (idx_tokens_conv - VAL >= 0)[:, None] + & (idx_tokens_conv - VAL < seqlen)[:, None] + & (idx_feats < dim)[None, :] + ) # token-index # token-index # feature-index + loaded_x = tl.load(x_ptrs, mask_x, 0.0) + + tl.debug_barrier() # need this due to the bug in tl.where not enforcing this when data is the result of another tl.load + new_conv_state = tl.where( + mask, conv_state, loaded_x + ) # BUG in 'tl.where' which requires a barrier before this + conv_states_ptrs_target = ( + conv_states_base + (idx_tokens_conv * stride_conv_state_tok)[:, None] + ) # [BLOCK_M, BLOCK_N] + mask = (idx_tokens_conv < state_len)[:, None] & (idx_feats < dim)[None, :] + tl.store(conv_states_ptrs_target, new_conv_state, mask) + else: # load_init_state == False + # update conv_state by shifting left, BUT + # set cols prior to 'x' as zeros + cols from 'x' + idx_tokens_conv = tl.arange(0, NP2_STATELEN) # [BLOCK_M] + + VAL = state_len - seqlen + + x_ptrs = x_base[None, :] + ((idx_tokens_conv - VAL) * stride_x_token)[:, None] # [BLOCK_M, BLOCK_N] + + mask_x = ( + (idx_tokens_conv - VAL >= 0)[:, None] + & (idx_tokens_conv - VAL < seqlen)[:, None] + & (idx_feats < dim)[None, :] + ) # token-index # token-index # feature-index + new_conv_state = tl.load(x_ptrs, mask_x, 0.0) + + conv_states_ptrs_target = ( + conv_states_base + (idx_tokens_conv * stride_conv_state_tok)[:, None] + ) # [BLOCK_M, BLOCK_N] + mask = (idx_tokens_conv < state_len)[:, None] & (idx_feats < dim)[None, :] + tl.store(conv_states_ptrs_target, new_conv_state, mask) + + else: # chunk_offset > 0 + # read prior-token data from `x` + load_init_state = True + prior_tokens = x_base + (token_offset - 1) * stride_x_token + mask_w = idx_feats < dim + if KERNEL_WIDTH == 2: + conv_states_ptrs = prior_tokens # [BLOCK_N] + col0 = tl.load(conv_states_ptrs, mask_w, 0.0, cache_modifier=".ca") + if KERNEL_WIDTH == 3: + conv_states_ptrs = prior_tokens # [BLOCK_N] + col1 = tl.load(conv_states_ptrs, mask_w, 0.0, cache_modifier=".ca") + conv_states_ptrs = prior_tokens - 1 * stride_x_token # [BLOCK_N] + col0 = tl.load(conv_states_ptrs, mask_w, 0.0, cache_modifier=".ca") + if KERNEL_WIDTH == 4: + conv_states_ptrs = prior_tokens # [BLOCK_N] + col2 = tl.load(conv_states_ptrs, mask_w, 0.0, cache_modifier=".ca") + conv_states_ptrs = prior_tokens - 1 * stride_x_token # [BLOCK_N] + col1 = tl.load(conv_states_ptrs, mask_w, 0.0, cache_modifier=".ca") + conv_states_ptrs = prior_tokens - 2 * stride_x_token # [BLOCK_N] + col0 = tl.load(conv_states_ptrs, mask_w, 0.0, cache_modifier=".ca") + if KERNEL_WIDTH == 5: + # ruff: noqa: F841 + conv_states_ptrs = prior_tokens # [BLOCK_N] + col3 = tl.load(conv_states_ptrs, mask_w, 0.0, cache_modifier=".ca") + conv_states_ptrs = prior_tokens - 1 * stride_x_token # [BLOCK_N] + col2 = tl.load(conv_states_ptrs, mask_w, 0.0, cache_modifier=".ca") + conv_states_ptrs = prior_tokens - 2 * stride_x_token # [BLOCK_N] + col1 = tl.load(conv_states_ptrs, mask_w, 0.0, cache_modifier=".ca") + conv_states_ptrs = prior_tokens - 3 * stride_x_token # [BLOCK_N] + col0 = tl.load(conv_states_ptrs, mask_w, 0.0, cache_modifier=".ca") + + # Store intermediate states aligned with stride_block_m + # The additional states are cached starting from the last stride_block_m. + # For example: + # If n_block_to_fill = 0, then only the state at the sequence end is cached and the process below is not involved. + # If n_block_to_fill > 0, then the states at the sequence end and at the n_block_to_fill-last + # stride_block_m are cached. + # For example chunk_offset = n_block_to_fill stores the state at last_full_block + if (chunk_offset - 1) < n_block_to_fill: + # Store the states at the chunk boundaries from the start of the sequence + idx_tokens_last = ( + last_full_block_token_index - (n_block_to_fill - chunk_offset) * B_size - state_len + ) + tl.arange(0, NP2_STATELEN) # [BLOCK_M] + x_ptrs = ( + x_ptr + (idx_tokens_last * stride_x_token)[:, None] + (idx_feats * stride_x_dim)[None, :] + ) # [BLOCK_M,BLOCK_N,] + + mask_x = (idx_tokens_last >= 0)[:, None] & (idx_feats < dim)[ + None, : + ] # token-index # token-index # feature-index + loaded_x = tl.load(x_ptrs, mask_x, 0.0) + idx_tokens_conv = tl.arange(0, NP2_STATELEN) # [BLOCK_M] + + # cache_idx + conv_states_output_coord = tl.load( + conv_state_indices_ptr + idx_seq * stride_cache_indices + current_first_index + (chunk_offset - 1) + ).to(tl.int64) + + conv_states_ptrs_target = ( + conv_states_ptr + + (conv_states_output_coord * stride_conv_state_seq) # Offset from seq + + (idx_feats * stride_conv_state_dim) + )[None, :] + ( # [BLOCK_N,] + idx_tokens_conv * stride_conv_state_tok + )[:, None] + + mask = (idx_tokens_conv < state_len)[:, None] & (idx_feats < dim)[None, :] + tl.debug_barrier() # NOTE: use this due to bug in Triton compiler + tl.store(conv_states_ptrs_target, loaded_x, mask) + + if HAS_BIAS: + bias = bias_ptr + idx_feats + mask_bias = idx_feats < dim + acc_preload = tl.load(bias, mask=mask_bias, other=0.0).to(tl.float32) # [BLOCK_N] + else: + acc_preload = tl.zeros((BLOCK_N,), dtype=tl.float32) + + x_base_1d = x_base + token_offset * stride_x_token # starting of chunk + + # PRE-LOAD WEIGHTS + mask_w = idx_feats < dim + if KERNEL_WIDTH >= 2: + w_ptrs = w_base + (0 * stride_w_width) # [BLOCK_N] tensor + w_col0 = tl.load(w_ptrs, mask_w, other=0.0) + w_ptrs = w_base + (1 * stride_w_width) # [BLOCK_N] tensor + w_col1 = tl.load(w_ptrs, mask_w, other=0.0) + if KERNEL_WIDTH >= 3: + w_ptrs = w_base + (2 * stride_w_width) # [BLOCK_N] tensor + w_col2 = tl.load(w_ptrs, mask_w, other=0.0) + if KERNEL_WIDTH >= 4: + w_ptrs = w_base + (3 * stride_w_width) # [BLOCK_N] tensor + w_col3 = tl.load(w_ptrs, mask_w, other=0.0) + mask_x_1d = idx_feats < dim + for idx_token in range(segment_len): + acc = acc_preload + + matrix_w = w_col0 + matrix_x = col0 + for j in tl.static_range(KERNEL_WIDTH): + if KERNEL_WIDTH == 2: + if j == 1: # KERNEL_WIDTH-1: + matrix_w = w_col1 + x_ptrs_1d = x_base_1d + idx_token * stride_x_token # [BLOCK_N] + matrix_x = tl.load(x_ptrs_1d, mask=mask_x_1d) + elif KERNEL_WIDTH == 3: + if j == 1: + matrix_w = w_col1 + matrix_x = col1 + elif j == 2: + matrix_w = w_col2 + x_ptrs_1d = x_base_1d + idx_token * stride_x_token # [BLOCK_N] + matrix_x = tl.load(x_ptrs_1d, mask=mask_x_1d) + elif KERNEL_WIDTH == 4: + if j == 1: + matrix_w = w_col1 + matrix_x = col1 + elif j == 2: + matrix_w = w_col2 + matrix_x = col2 + elif j == 3: + matrix_w = w_col3 + x_ptrs_1d = x_base_1d + idx_token * stride_x_token # [BLOCK_N] + matrix_x = tl.load(x_ptrs_1d, mask=mask_x_1d) + + acc += matrix_x * matrix_w # [BLOCK_N] + + if KERNEL_WIDTH == 2: + col0 = matrix_x + elif KERNEL_WIDTH == 3: + col0 = col1 + col1 = matrix_x + elif KERNEL_WIDTH == 4: + col0 = col1 + col1 = col2 + col2 = matrix_x + + if SILU_ACTIVATION: + acc = acc / (1 + tl.exp(-acc)) + mask_1d = (idx_token < segment_len) & (idx_feats < dim) # token-index # feature-index + o_ptrs = o_ptr + (sequence_start_index + token_offset + idx_token) * stride_o_token + (idx_feats * stride_o_dim) + + tl.store(o_ptrs, acc, mask=mask_1d) + + +def causal_conv1d_fn( + x: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor | None, + conv_states: torch.Tensor, + query_start_loc: torch.Tensor, + cache_indices: torch.Tensor | None = None, + has_initial_state: torch.Tensor | None = None, + activation: str | None = "silu", + pad_slot_id: int = PAD_SLOT_ID, + null_block_id: int = NULL_BLOCK_ID, + block_idx_first_scheduled_token: torch.Tensor | None = None, + block_idx_last_scheduled_token: torch.Tensor | None = None, + initial_state_idx: torch.Tensor | None = None, + num_computed_tokens: torch.Tensor | None = None, + block_size_to_align=0, + metadata=None, + validate_data=False, +): + """support varlen + continuous batching when x is 2D tensor + + x: (dim,cu_seq_len) + cu_seq_len = total tokens of all seqs in that batch + sequences are concatenated from left to right for varlen + weight: (dim, width) + conv_states: (...,dim,width - 1) itype + updated inplace if cache_indices are not provided + [it use `cache_indices` to get the index to the cache of conv_state for that sequence + + conv_state[cache_indices[i]] for seq-i - to be used as initial_state when has_initial_state[i] = True + and after that conv_state[cache_indices[i]] need to be shift-left and updated with values from 'x' + ] + query_start_loc: (batch + 1) int32 + The cumulative sequence lengths of the sequences in + the batch, used to index into sequence. prepended by 0. + if + x = [5, 1, 1, 1] <- continuous batching (batch=4) + then + query_start_loc = [0, 5, 6, 7, 8] <- the starting index of the next sequence; while the last value is + the ending index of the last sequence + [length(query_start_loc)-1 == batch] + for example: query_start_loc = torch.Tensor([0,10,16,17]), + x.shape=(dim,17) + cache_indices: (batch) int32 + indicates the corresponding state index, + like so: conv_state = conv_states[cache_indices[batch_id]] + has_initial_state: (batch) bool + indicates whether should the kernel take the current state as initial + state for the calculations + [single boolean for each sequence in the batch: True or False] + bias: (dim,) + activation: either None or "silu" or "swish" or True + pad_slot_id: int + if cache_indices is passed, lets the kernel identify padded + entries that will not be processed, + for example: cache_indices = [pad_slot_id, 1, 20, pad_slot_id] + in this case, the kernel will not process entries at + indices 0 and 3 + block_idx_first_scheduled_token: (batch,), dtype int32 + The pointer into cache_indices, where the first cache block to be filled is located. + block_idx_last_scheduled_token: (batch,), dtype int32 + The pointer into cache_indices, where the last cache block to be filled is located. + initial_state_idx: (batch,), dtype int32 + The pointer into cache_indices, where the cache block containing the initial state is located. + num_computed_tokens: (batch,), dtype int32 + The number of tokens already completed for each sequence + block_size_to_align: int + The block size to align the cached states to + out: same shape as `x` + """ + if isinstance(activation, bool) and activation: + activation = "silu" + + args = None + # Store original dtype to cast back at the end + original_x_dtype = x.dtype + x = x.to(conv_states.dtype) + out = torch.empty_like(x) + if metadata is not None: + nums_dict = metadata.nums_dict + args = nums_dict + batch_ptr = metadata.batch_ptr + token_chunk_offset_ptr = metadata.token_chunk_offset_ptr + else: + seqlens = query_start_loc.diff().to("cpu") + args = seqlens + MAX_NUM_PROGRAMS = 1024 + + batch_ptr = torch.full( + (MAX_NUM_PROGRAMS,), PAD_SLOT_ID, dtype=torch.int32, device=x.device + ) # tracking which seq-idx the Triton program is handling + token_chunk_offset_ptr = torch.full( + (MAX_NUM_PROGRAMS,), PAD_SLOT_ID, dtype=torch.int32, device=x.device + ) # tracking BLOCK_M-based index in the sequence the Triton program is handling + + is_channel_last = (x.stride(0) == 1) & (x.stride(1) > 1) + dim, cu_seqlen = x.shape + _, width = weight.shape + state_len = width - 1 + np2_statelen = triton.next_power_of_2(state_len) + + padded_batch = query_start_loc.size(0) - 1 + stride_x_dim = x.stride(0) + stride_x_token = x.stride(1) + stride_w_dim = weight.stride(0) + stride_w_width = weight.stride(1) + stride_istate_seq = 0 + stride_istate_dim = 0 + stride_istate_token = 0 + num_cache_lines = 0 + BLOCK_M = 8 + if conv_states is not None: + # extensions to support Aphrodite: + # 1. conv_states is used to replaced initial_states + # 2. conv_states serve as a cache with num cache lines can be larger than batch size + # 3. mapping from sequence x[idx] to a cache line at index as specified via cache_indices[idx] + # 4. computation can be skipped if cache_indices[idx] == pad_slot_id + num_cache_lines = conv_states.size(0) + assert ( + num_cache_lines == conv_states.shape[0] + and dim == conv_states.shape[1] + and width - 1 <= conv_states.shape[2] + ) + stride_istate_seq = conv_states.stride(0) + stride_istate_dim = conv_states.stride(1) + stride_istate_token = conv_states.stride(2) + if out.dim() == 2: + stride_o_dim = out.stride(0) + stride_o_token = out.stride(1) + else: + stride_o_dim = out.stride(1) + stride_o_token = out.stride(2) + stride_cache_indices = cache_indices.stride(0) if cache_indices is not None else 0 + + if validate_data: + assert x.dim() == 2 + assert query_start_loc is not None + assert query_start_loc.dim() == 1 + assert x.stride(0) == 1 or x.stride(1) == 1 + if bias is not None: + assert bias.dim() == 1 + assert dim == bias.size(0) + if cache_indices is not None: + assert cache_indices.dim() == 1 + assert padded_batch == cache_indices.size(0) + if has_initial_state is not None: + assert has_initial_state.size() == (padded_batch,) + assert conv_states is not None, "ERROR: `has_initial_state` is used, which needs also `conv_states`" + assert weight.stride(1) == 1 + assert (dim, width) == weight.shape + assert is_channel_last, "Need to run in channel-last layout" + if block_size_to_align is not None and block_size_to_align > 0: + assert (block_size_to_align % BLOCK_M) == 0, "The mamba block size needs to be divisible by the BLOCK_M" + else: + block_size_to_align = BLOCK_M + + if metadata is None: + + def num_program(META, seqlens): + tot = 0 + + mlist = [] + offsetlist = [] # type: ignore + + nums = -(-seqlens // META["BLOCK_M"]) + + tot = nums.sum().item() + mlist = np.repeat(np.arange(len(nums)), nums) + for idx, num in enumerate(nums): + offsetlist.extend(range(num)) # chunk-idx if a sequence is split into multiple chunks + + if META["batch_ptr"].nelement() < len(mlist): + newlen = len(mlist) + 1 + META["batch_ptr"].resize_(newlen).fill_(PAD_SLOT_ID) + META["token_chunk_offset_ptr"].resize_(newlen).fill_(PAD_SLOT_ID) + + if META["batch_ptr"].nelement() >= len(mlist): + META["batch_ptr"][0 : len(mlist)].copy_(torch.from_numpy(np.array(mlist))) + META["token_chunk_offset_ptr"][0 : len(mlist)].copy_(torch.from_numpy(np.array(offsetlist))) + + META["batch_ptr"] = META["batch_ptr"].to(META["x_ptr"].device) + META["token_chunk_offset_ptr"] = META["token_chunk_offset_ptr"].to(META["x_ptr"].device) + return tot + else: + + def num_program(META, nums_dict): + tot = nums_dict[META["BLOCK_M"]]["tot"] + + mlist = nums_dict[META["BLOCK_M"]]["mlist"] + mlist_len = nums_dict[META["BLOCK_M"]]["mlist_len"] + + offsetlist = nums_dict[META["BLOCK_M"]]["offsetlist"] + + if nums_dict[META["BLOCK_M"]]["batch_ptr"] is not None: + META["batch_ptr"] = nums_dict[META["BLOCK_M"]]["batch_ptr"] + META["token_chunk_offset_ptr"] = nums_dict[META["BLOCK_M"]]["token_chunk_offset_ptr"] + else: + if META["batch_ptr"].nelement() < mlist_len: + newlen = mlist_len + 1 + META["batch_ptr"].resize_(newlen).fill_(PAD_SLOT_ID) + META["token_chunk_offset_ptr"].resize_(newlen).fill_(PAD_SLOT_ID) + + if META["batch_ptr"].nelement() >= mlist_len: + META["batch_ptr"][0:mlist_len].copy_(mlist) + META["token_chunk_offset_ptr"][0:mlist_len].copy_(offsetlist) + return tot + + def grid(META): + return ( + num_program(META, args), + triton.cdiv(dim, META["BLOCK_N"]), + ) + + if batch_ptr.device != x.device: + batch_ptr = batch_ptr.to(x.device) + token_chunk_offset_ptr = token_chunk_offset_ptr.to(x.device) + + _causal_conv1d_fwd_kernel[grid]( + # Pointers to matrices + x, + weight, + bias, + conv_states, + cache_indices, + has_initial_state, + query_start_loc, + batch_ptr, + token_chunk_offset_ptr, + block_idx_first_scheduled_token, + block_idx_last_scheduled_token, + initial_state_idx, + num_computed_tokens, + out, + # Matrix dimensions + dim, + num_cache_lines, + # stride + stride_x_dim, + stride_x_token, + stride_w_dim, + stride_w_width, + stride_istate_seq, + stride_istate_dim, + stride_istate_token, + stride_cache_indices, + stride_o_dim, + stride_o_token, + block_size_to_align // BLOCK_M, + # others + pad_slot_id, + null_block_id, + # META + HAS_BIAS=bias is not None, + KERNEL_WIDTH=width, + SILU_ACTIVATION=activation in ["silu", "swish"], + IS_APC_ENABLED=block_idx_last_scheduled_token is not None, + HAS_NULL_BLOCK=null_block_id is not None, + NP2_STATELEN=np2_statelen, + # launch_cooperative_grid=True + BLOCK_M=BLOCK_M, + BLOCK_N=256, + num_stages=2, + ) + return out.to(original_x_dtype) + + +@triton.jit() +def _causal_conv1d_update_kernel( + # Pointers to matrices + x_ptr, # (batch, dim, seqlen) + w_ptr, # (dim, width) + bias_ptr, + conv_state_ptr, + conv_state_indices_ptr, + num_accepted_tokens_ptr, + query_start_loc_ptr, # (batch + 1) + block_idx_last_scheduled_token, # (batch,) + initial_state_idx, # (batch,) + o_ptr, # (batch, dim, seqlen) + # Matrix dimensions + batch: int, + dim: tl.constexpr, + seqlen: tl.constexpr, + state_len: tl.constexpr, + num_cache_lines: tl.constexpr, # added to support Aphrodite larger cache lines + # Strides + stride_x_seq: tl.constexpr, + stride_x_dim: tl.constexpr, + stride_x_token: tl.int64, + stride_w_dim: tl.constexpr, + stride_w_width: tl.constexpr, + stride_conv_state_seq: tl.constexpr, + stride_conv_state_dim: tl.constexpr, + stride_conv_state_tok: tl.constexpr, + stride_state_indices: tl.constexpr, + stride_o_seq: tl.constexpr, + stride_o_dim: tl.constexpr, + stride_o_token: tl.int64, + # others + null_block_id: tl.constexpr, + # Meta-parameters + HAS_BIAS: tl.constexpr, + KERNEL_WIDTH: tl.constexpr, + SILU_ACTIVATION: tl.constexpr, + IS_VARLEN: tl.constexpr, + IS_APC_ENABLED: tl.constexpr, + IS_SPEC_DECODING: tl.constexpr, + NP2_STATELEN: tl.constexpr, + HAS_NULL_BLOCK: tl.constexpr, + BLOCK_N: tl.constexpr, +): + # ruff: noqa: E501 + idx_seq = tl.program_id(0) + if idx_seq >= batch: + return + + # [BLOCK_N,] elements along the feature-dimension (channel) + idx_feats = tl.program_id(1) * BLOCK_N + tl.arange(0, BLOCK_N) + + if IS_APC_ENABLED: + # Get the state from the initial_state_idx + conv_state_init = tl.load(initial_state_idx + idx_seq) + current_last_index = tl.load(block_idx_last_scheduled_token + idx_seq) + else: + conv_state_init = 0 + current_last_index = 0 + + # cache_idx + conv_states_input_coord = tl.load(conv_state_indices_ptr + idx_seq * stride_state_indices + conv_state_init).to( + tl.int64 + ) + + if HAS_NULL_BLOCK: # noqa + if conv_states_input_coord == null_block_id: + # not processing as this is not the actual sequence + return + + if IS_VARLEN: + query_start_index = tl.load(query_start_loc_ptr + idx_seq).to(tl.int64) + query_end_index = tl.load(query_start_loc_ptr + (idx_seq + 1)).to(tl.int64) + # revise state_len and seqlen + state_len = state_len - (seqlen - (query_end_index - query_start_index)) + seqlen = query_end_index - query_start_index + x_offset = query_start_index * stride_x_token + o_offset = query_start_index * stride_o_token + else: + query_start_index = idx_seq * seqlen + query_end_index = query_start_index + seqlen + x_offset = idx_seq * stride_x_seq + o_offset = idx_seq * stride_o_seq + + if query_start_index == query_end_index: + return + + if IS_SPEC_DECODING: + # The rolling of conv state: + # + # Before forward, the conv_state is: + # [history1, history2, ..., historyM]. + # + # After forward, the conv_state becomes: + # [history2, ..., historyM, draft1, draft2, ..., draftN]. + # + # After acceptance, it becomes: + # + # - accept 1 tokens: [history2, ..., historyM, draft1] + # - accept 2 tokens: [history3, ..., historyM, draft1, draft2] + # - and so on. + conv_state_token_offset = tl.load(num_accepted_tokens_ptr + idx_seq).to(tl.int64) - 1 + else: + conv_state_token_offset = 0 + + # STEP 1: READ init_state data + conv_states_base = ( + conv_state_ptr + (conv_states_input_coord * stride_conv_state_seq) + (idx_feats * stride_conv_state_dim) + ) + mask_w = idx_feats < dim + + prior_tokens = conv_states_base + conv_state_token_offset * stride_conv_state_tok + if KERNEL_WIDTH >= 2: + conv_states_ptrs = prior_tokens # [BLOCK_N] + col0 = tl.load(conv_states_ptrs, mask_w, 0.0) + if KERNEL_WIDTH >= 3: + conv_states_ptrs = prior_tokens + 1 * stride_conv_state_tok # [BLOCK_N] + col1 = tl.load(conv_states_ptrs, mask_w, 0.0) + if KERNEL_WIDTH >= 4: + conv_states_ptrs = prior_tokens + 2 * stride_conv_state_tok # [BLOCK_N] + col2 = tl.load(conv_states_ptrs, mask_w, 0.0) + if KERNEL_WIDTH >= 5: + conv_states_ptrs = prior_tokens + 3 * stride_conv_state_tok # [BLOCK_N] + col3 = tl.load(conv_states_ptrs, mask_w, 0.0) + if KERNEL_WIDTH >= 6: + conv_states_ptrs = prior_tokens + 4 * stride_conv_state_tok # [BLOCK_N] + col4 = tl.load(conv_states_ptrs, mask_w, 0.0) + + # STEP 2: assume state_len > seqlen + idx_tokens = tl.arange(0, NP2_STATELEN) # [BLOCK_M] + + # With speculative decoding, the conv_state updates works in a sliding + # window manner, at each forward pass, the tokens are shift by 1, so we + # load since idx_tokens + 1. + conv_state_ptrs_source = ( + conv_state_ptr + + (conv_states_input_coord * stride_conv_state_seq) + + conv_state_token_offset * stride_conv_state_tok + + (idx_feats * stride_conv_state_dim)[None, :] + + ((idx_tokens + (1 if IS_SPEC_DECODING else seqlen)) * stride_conv_state_tok)[:, None] + ) # [BLOCK_M, BLOCK_N] + mask = ( + (conv_states_input_coord < num_cache_lines) + & ((idx_tokens + seqlen) < state_len)[:, None] + & (idx_feats < dim)[None, :] + ) + conv_state = tl.load(conv_state_ptrs_source, mask, other=0.0) + + VAL = state_len - seqlen + x_base = x_ptr + x_offset + (idx_feats * stride_x_dim) # [BLOCK_N] + + x_ptrs = x_base[None, :] + ((idx_tokens - VAL) * stride_x_token)[:, None] # [BLOCK_M, BLOCK_N] + + mask_x = ( + (idx_tokens - VAL >= 0)[:, None] & (idx_tokens - VAL < seqlen)[:, None] & (idx_feats < dim)[None, :] + ) # token-index # token-index # feature-index + loaded_x = tl.load(x_ptrs, mask_x, 0.0) + tl.debug_barrier() + + new_conv_state = tl.where(mask, conv_state, loaded_x) + + # Get the state from the initial_state_idx + # cache_idx + conv_states_offset = tl.load(conv_state_indices_ptr + idx_seq * stride_state_indices + current_last_index).to( + tl.int64 + ) + conv_state_ptrs_target = ( + conv_state_ptr + + (conv_states_offset * stride_conv_state_seq) # Offset from seq + + (idx_feats * stride_conv_state_dim) + )[None, :] + ( # [BLOCK_N,] + idx_tokens * stride_conv_state_tok + )[:, None] + mask = (idx_tokens < state_len)[:, None] & (idx_feats < dim)[None, :] + tl.store(conv_state_ptrs_target, new_conv_state, mask) + + # STEP 3: init accumulator + if HAS_BIAS: + bias = bias_ptr + idx_feats + mask_bias = idx_feats < dim + acc_preload = tl.load(bias, mask=mask_bias, other=0.0).to(tl.float32) # [BLOCK_N] + else: + acc_preload = tl.zeros((BLOCK_N,), dtype=tl.float32) + + # STEP 4: + # PRE-LOAD WEIGHTS + # first kernel column, configured for weights to handle BLOCK_N features in range + w_base = w_ptr + (idx_feats * stride_w_dim) # [BLOCK_N,] + mask_w = idx_feats < dim + if KERNEL_WIDTH >= 2: + w_ptrs = w_base + (0 * stride_w_width) # [BLOCK_N] tensor + w_col0 = tl.load(w_ptrs, mask_w, other=0.0) + w_ptrs = w_base + (1 * stride_w_width) # [BLOCK_N] tensor + w_col1 = tl.load(w_ptrs, mask_w, other=0.0) + if KERNEL_WIDTH >= 3: + w_ptrs = w_base + (2 * stride_w_width) # [BLOCK_N] tensor + w_col2 = tl.load(w_ptrs, mask_w, other=0.0) + if KERNEL_WIDTH >= 4: + w_ptrs = w_base + (3 * stride_w_width) # [BLOCK_N] tensor + w_col3 = tl.load(w_ptrs, mask_w, other=0.0) + if KERNEL_WIDTH >= 5: + w_ptrs = w_base + (4 * stride_w_width) # [BLOCK_N] tensor + w_col4 = tl.load(w_ptrs, mask_w, other=0.0) + if KERNEL_WIDTH >= 6: + w_ptrs = w_base + (5 * stride_w_width) # [BLOCK_N] tensor + w_col5 = tl.load(w_ptrs, mask_w, other=0.0) + + x_base_1d = x_base # starting of chunk [BLOCK_N] + mask_x_1d = idx_feats < dim + + # STEP 5: compute each token + for idx_token in tl.range(seqlen): + acc = acc_preload + + matrix_w = w_col0 + matrix_x = col0 + for j in tl.static_range(KERNEL_WIDTH): + if KERNEL_WIDTH == 2: + if j == 1: # KERNEL_WIDTH-1: + matrix_w = w_col1 + x_ptrs_1d = x_base_1d + idx_token * stride_x_token # [BLOCK_N] + matrix_x = tl.load(x_ptrs_1d, mask=mask_x_1d) + elif KERNEL_WIDTH == 3: + if j == 1: + matrix_w = w_col1 + matrix_x = col1 + elif j == 2: + matrix_w = w_col2 + x_ptrs_1d = x_base_1d + idx_token * stride_x_token # [BLOCK_N] + matrix_x = tl.load(x_ptrs_1d, mask=mask_x_1d) + elif KERNEL_WIDTH == 4: + if j == 1: + matrix_w = w_col1 + matrix_x = col1 + elif j == 2: + matrix_w = w_col2 + matrix_x = col2 + elif j == 3: + matrix_w = w_col3 + x_ptrs_1d = x_base_1d + idx_token * stride_x_token # [BLOCK_N] + matrix_x = tl.load(x_ptrs_1d, mask=mask_x_1d) + elif KERNEL_WIDTH == 5: + if j == 1: + matrix_w = w_col1 + matrix_x = col1 + elif j == 2: + matrix_w = w_col2 + matrix_x = col2 + elif j == 3: + matrix_w = w_col3 + matrix_x = col3 + elif j == 4: + matrix_w = w_col4 + x_ptrs_1d = x_base_1d + idx_token * stride_x_token # [BLOCK_N] + matrix_x = tl.load(x_ptrs_1d, mask=mask_x_1d) + elif KERNEL_WIDTH == 6: + if j == 1: + matrix_w = w_col1 + matrix_x = col1 + elif j == 2: + matrix_w = w_col2 + matrix_x = col2 + elif j == 3: + matrix_w = w_col3 + matrix_x = col3 + elif j == 4: + matrix_w = w_col4 + matrix_x = col4 + elif j == 5: + matrix_w = w_col5 + x_ptrs_1d = x_base_1d + idx_token * stride_x_token # [BLOCK_N] + matrix_x = tl.load(x_ptrs_1d, mask=mask_x_1d) + + acc += matrix_x * matrix_w # [BLOCK_N] + + if KERNEL_WIDTH == 2: + col0 = matrix_x + elif KERNEL_WIDTH == 3: + col0 = col1 + col1 = matrix_x + elif KERNEL_WIDTH == 4: + col0 = col1 + col1 = col2 + col2 = matrix_x + elif KERNEL_WIDTH == 5: + col0 = col1 + col1 = col2 + col2 = col3 + col3 = matrix_x + elif KERNEL_WIDTH == 6: + col0 = col1 + col1 = col2 + col2 = col3 + col3 = col4 + col4 = matrix_x + + if SILU_ACTIVATION: + acc = acc / (1 + tl.exp(-acc)) + mask_1d = (idx_token < seqlen) & (idx_feats < dim) # token-index # feature-index + o_ptrs = o_ptr + o_offset + idx_token * stride_o_token + (idx_feats * stride_o_dim) + + tl.store(o_ptrs, acc, mask=mask_1d) + + +def causal_conv1d_update( + x: torch.Tensor, + conv_state: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor | None = None, + activation: bool | str | None = None, + conv_state_indices: torch.Tensor | None = None, + num_accepted_tokens: torch.Tensor | None = None, + query_start_loc: torch.Tensor | None = None, + max_query_len: int = -1, + null_block_id: int = NULL_BLOCK_ID, + block_idx_last_scheduled_token: torch.Tensor | None = None, + initial_state_idx: torch.Tensor | None = None, + validate_data=False, +): + """ + x: Input tensor which can take the following shapes: + + - `[batch, dim]` - single token prediction + - `[batch, dim, seqlen]` - single or multiple tokens prediction + - `[num_tokens, dim]` - continuous batching, where num_tokens is + the total tokens of all sequences in that batch + + conv_state: (..., dim, state_len), where state_len >= width - 1 + weight: (dim, width) + bias: (dim,) + conv_state_indices: (batch,), dtype int32 + If not None, the conv_state is a larger tensor along the batch dim, + and we are selecting the batch coords specified by conv_state_indices. + Useful for a continuous batching scenario. + block_idx_last_scheduled_token: (batch,), dtype int32 + The pointer into conv_state_indices, where the last cache block to be filled is located. + initial_state_idx: (batch,), dtype int32 + The pointer into conv_state_indices, where the cache block containing the initial state is located. + num_accepted_tokens: (batch,), dtype int32 + If not None, it indicates the number of accepted tokens for each + sequence in the batch. + This is used in speculative decoding, where the conv_state is updated + in a sliding window manner. + query_start_loc: (batch + 1,) int32 + If not None, the inputs is given in a varlen fashion and this indicates + the starting index of each sequence in the batch. + max_query_len: int + If query_start_loc is not None, this indicates the maximum query + length in the batch. + null_block_id: int + Block ID used to identify padded entries in + conv_state_indices. Block 0 is the null block. + 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` + """ + if validate_data: + assert null_block_id is not None + assert x.stride(1) == 1 + if isinstance(activation, bool): + activation = "silu" if activation is True else None + elif activation is not None: + assert activation in ["silu", "swish"] + + original_x_dtype = x.dtype + x = x.to(conv_state.dtype) + unsqueeze = query_start_loc is None and x.dim() == 2 + if unsqueeze: + # make it (batch, dim, seqlen) with seqlen == 1 + x = x.unsqueeze(-1) + if query_start_loc is None: + batch, dim, seqlen = x.shape + else: + assert conv_state_indices is not None + batch = conv_state_indices.size(0) + dim = x.size(1) + seqlen = max_query_len + _, width = weight.shape + # conv_state: (..., dim, state_len), where state_len >= width - 1 + num_cache_lines, _, state_len = conv_state.size() + + if validate_data: + assert dim == weight.size(0) + assert state_len >= width - 1 + # when above happens, we don't shift-left to keep any records in conv_state + assert dim == conv_state.size(1) + if conv_state_indices is None: + assert conv_state.size(0) >= batch + else: + assert batch == conv_state_indices.shape[0], ( + f"ERROR: conv_state_indices should have shape ({batch},*) but got {conv_state_indices.shape}" + ) + + assert num_cache_lines >= batch + assert weight.stride(1) == 1 # Need this + + # adopt the strategy in Aphrodite that overwrite on 'x' directly, rather than creating a new tensor 'o' + out = x + stride_w_dim, stride_w_width = weight.stride() + + if query_start_loc is None: + # X (batch, dim, seqlen) + stride_x_seq, stride_x_dim, stride_x_token = x.stride() + stride_o_seq, stride_o_dim, stride_o_token = out.stride() + else: + # X (dim, cu_seqlen) + stride_x_token, stride_x_dim = x.stride() + stride_x_seq = 0 + stride_o_token, stride_o_dim = out.stride() + stride_o_seq = 0 + + stride_istate_seq, stride_istate_dim, stride_istate_token = conv_state.stride() + stride_state_indices = conv_state_indices.stride(0) if conv_state_indices is not None else 0 + if num_accepted_tokens is not None: + state_len = width - 1 + (seqlen - 1) # effective state_len needed + else: + state_len = width - 1 + np2_statelen = triton.next_power_of_2(state_len) + + def grid(META): + return ( + batch, + triton.cdiv(dim, META["BLOCK_N"]), + ) + + _causal_conv1d_update_kernel[grid]( + # Pointers to matrices + x, + weight, + bias, + conv_state, + conv_state_indices, + num_accepted_tokens, + query_start_loc, + block_idx_last_scheduled_token, + initial_state_idx, + out, + # Matrix dimensions + batch, + dim, + seqlen, + state_len, + num_cache_lines, + # stride + stride_x_seq, + stride_x_dim, + stride_x_token, + stride_w_dim, + stride_w_width, + stride_istate_seq, + stride_istate_dim, + stride_istate_token, + stride_state_indices, + stride_o_seq, + stride_o_dim, + stride_o_token, + # others + null_block_id, + # META + HAS_BIAS=bias is not None, + KERNEL_WIDTH=width, + SILU_ACTIVATION=activation in ["silu", "swish"], + IS_VARLEN=query_start_loc is not None, + IS_APC_ENABLED=block_idx_last_scheduled_token is not None, + IS_SPEC_DECODING=num_accepted_tokens is not None, + NP2_STATELEN=np2_statelen, + HAS_NULL_BLOCK=null_block_id is not None, + BLOCK_N=256, + ) + if unsqueeze: + out = out.squeeze(-1) + return out.to(original_x_dtype) + + +from aphrodite.platforms import current_platform # noqa: E402 + +if current_platform.is_cpu(): + from aphrodite.model_executor.layers.mamba.ops.cpu.causal_conv1d import ( + causal_conv1d_fn_cpu, + causal_conv1d_update_cpu, + ) + + causal_conv1d_fn = causal_conv1d_fn_cpu # type: ignore + causal_conv1d_update = causal_conv1d_update_cpu # type: ignore 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/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..23dac23c57 --- /dev/null +++ b/aphrodite/models/common/ops/fused_qk_rmsnorm.py @@ -0,0 +1,103 @@ +# 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/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/inkling/nvidia/ops/fa4_rel_attention.py.orig b/aphrodite/models/inkling/nvidia/ops/fa4_rel_attention.py.orig new file mode 100644 index 0000000000..a3fc2c0b7d --- /dev/null +++ b/aphrodite/models/inkling/nvidia/ops/fa4_rel_attention.py.orig @@ -0,0 +1,161 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from __future__ import annotations + +from collections.abc import Callable +from functools import cache +from typing import Any + +import torch + +from aphrodite.platforms import current_platform + + +def bucket_max_seqlen_q(max_seqlen_q: int) -> int: + """Round the FA4 scheduling bound up to a power of two.""" + return 1 << max(0, max_seqlen_q - 1).bit_length() + + +@cache +def _use_sheared_bias() -> bool: + capability = current_platform.get_device_capability() + return capability is not None and capability.major in (10, 11) + + +@cache +def _get_score_mod(rel_extent: int) -> Callable: + """Return the score modification that adds Inkling relative bias.""" + import cutlass.cute as cute + from cutlass.cute import Float32 + + from aphrodite.vllm_flash_attn.cute.seqlen_info import SeqlenInfoQK + + @cute.jit + def score_mod_rel_bias( + scores: cute.TensorSSA, + b_idx: cute.TensorSSA, + h_idx: cute.TensorSSA, + q_idx: cute.TensorSSA, + kv_idx: cute.TensorSSA, + seqlen_info: SeqlenInfoQK, + aux_tensors: list[cute.Tensor], + ) -> cute.TensorSSA: + rel_logits = aux_tensors[0] + + seqlen_local_offset = seqlen_info.seqlen_k - seqlen_info.seqlen_q + rel_dist = (q_idx + seqlen_local_offset) - kv_idx + global_q_idx = seqlen_info.offset_q + q_idx + + rel_dist_0 = rel_dist[0] + rel_idx = rel_dist_0 if rel_dist_0 >= 0 else 0 + rel_idx = rel_idx if rel_idx < rel_extent else (rel_extent - 1) + + rel_bias = rel_logits[global_q_idx[0], h_idx[0], rel_idx] + rel_bias = Float32(rel_bias) if rel_dist_0 == rel_idx else Float32(0.0) + return scores + rel_bias + + return score_mod_rel_bias + + +def inkling_fa4_num_splits( + *, + is_local: bool, + batch_size: int, + max_query_len: int, + num_heads: int, + num_kv_heads: int, + max_kv_len: int, +) -> int: + """Return the split-KV cap for Inkling relative attention.""" + capability = current_platform.get_device_capability() + if capability is not None and capability.major == 9: + return 1 + if is_local: + return 1 + + q_rows = max_query_len * (num_heads // num_kv_heads) + q_tiles = (q_rows + 255) // 256 + base_ctas = batch_size * num_kv_heads * q_tiles + # Shearing makes split/combine overhead more visible. Multi-tile causal + # prefill saturates around 64 CTAs. Batch-1 decode at very long context is + # memory-bound and uses a TP-specific cap measured through 1M KV tokens. + target_ctas = 256 if q_tiles == 1 and batch_size == 1 else (128 if q_tiles == 1 else 64) + max_splits = 128 + if q_tiles == 1 and batch_size == 1: + if num_kv_heads == 8: + max_splits = 16 + elif num_kv_heads == 4 or max_kv_len <= 8192: + max_splits = 32 + elif max_kv_len <= 65536: + max_splits = 64 + else: + max_splits = 128 + return max( + 1, + min(target_ctas // base_ctas, max_splits, (max_kv_len + 127) // 128), + ) + + +def inkling_fa4_rel_attention( + q: torch.Tensor, + key_cache: torch.Tensor, + value_cache: torch.Tensor, + *, + block_table: torch.Tensor, + cache_seqlens: torch.Tensor, + cu_seqlens_q: torch.Tensor, + max_seqlen_q: int, + softmax_scale: float, + causal: bool, + window_size: tuple[int, int], + rel_extent: int, + rel_logits: torch.Tensor, + num_splits: int = 32, + out: torch.Tensor | None = None, +) -> torch.Tensor: + """Paged varlen FA4 over the bound K/V cache with the Inkling relative bias. + + ``q`` is ``(num_tokens, num_heads, head_dim)``; ``key_cache`` / ``value_cache`` + are the paged caches ``(num_blocks, block_size, num_kv_heads, head_dim)``; + ``block_table`` is the per-request page table and ``cache_seqlens`` the + per-request KV lengths (``seqused_k``). ``rel_logits`` is + ``(num_tokens, num_heads, rel_extent)``. + + Hopper uses standard FA4's score-mod gather. Blackwell uses tml-fa4's + sheared relative-bias layout. + """ + # cute uses (None, None) to mean "no window". + cute_window = (None, None) if window_size == (-1, -1) else window_size + + rel_logits = rel_logits.contiguous() + if _use_sheared_bias(): + from aphrodite.third_party.tml_fa4 import 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 + + bias_kwargs = { + "score_mod": _get_score_mod(rel_extent), + "aux_tensors": [rel_logits], + } + + ret = flash_attn_varlen_func( + q=q, + k=key_cache, + v=value_cache, + cu_seqlens_q=cu_seqlens_q, + seqused_k=cache_seqlens, + max_seqlen_q=max_seqlen_q, + page_table=block_table, + softmax_scale=softmax_scale, + causal=causal, + window_size=cute_window, + num_splits=num_splits, + return_lse=False, + out=out, + **bias_kwargs, + ) + if isinstance(ret, tuple): + return ret[0] + return ret 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..21b3ea4d4f --- /dev/null +++ b/aphrodite/models/kimi_k3/amd/linear.py @@ -0,0 +1,993 @@ +# 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 CacheConfig, VllmConfig +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, + vllm_config: VllmConfig, + 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 = vllm_config.model_config + cache_config = vllm_config.cache_config + quant_config = vllm_config.quant_config + + if config.is_kda_layer(layer_idx): + self.self_attn = KimiGatedDeltaNetAttention( + config, + vllm_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, *, vllm_config: VllmConfig, prefix: str = ""): + super().__init__() + + config = vllm_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, + vllm_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, *, vllm_config: VllmConfig, prefix: str = ""): + super().__init__() + self.model_config = vllm_config.model_config + self.vllm_config = vllm_config + self.config = self.model_config.hf_config + quant_config = vllm_config.quant_config + self.quant_config = quant_config + self.model = KimiLinearModel(vllm_config=vllm_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, + vllm_config: "VllmConfig", + ) -> tuple[torch.dtype, torch.dtype]: + return MambaStateDtypeCalculator.kda_state_dtype( + vllm_config.model_config.dtype, vllm_config.cache_config.mamba_cache_dtype + ) + + @classmethod + def get_mamba_state_shape_from_config( + cls, vllm_config: "VllmConfig" + ) -> tuple[tuple[int, int], tuple[int, int, int]]: + parallel_config = vllm_config.parallel_config + hf_config = vllm_config.model_config.hf_config + tp_size = parallel_config.tensor_parallel_size + num_spec = vllm_config.speculative_config.num_speculative_tokens if vllm_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..597540b92f --- /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 VllmConfig +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, + vllm_config: VllmConfig, + prefix: str = "", + ) -> None: + super().__init__() + model_config = vllm_config.model_config + config: KimiK3Config = model_config.hf_config + self.config = config + quant_config = vllm_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(vllm_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(vllm_config): + self.language_model = init_vllm_registered_model( + vllm_config=vllm_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, vllm_config: VllmConfig): + text_config = vllm_config.model_config.hf_config.text_config + temp_vllm_config = vllm_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, vllm_config: VllmConfig): + text_config = vllm_config.model_config.hf_config.text_config + temp_vllm_config = vllm_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..eeb972dfbf --- /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 VllmConfig +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, + vllm_config: VllmConfig, + prefix: str, + ) -> None: + super().__init__() + self.config = config + quant_config = vllm_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, vllm_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, *, vllm_config: VllmConfig, prefix: str = ""): + super().__init__() + config: KimiLinearConfig = vllm_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, vllm_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, *, vllm_config: VllmConfig, prefix: str = ""): + super().__init__() + self.config = vllm_config.model_config.hf_text_config + self.quant_config = vllm_config.quant_config + self.model = KimiK3MultiTokenPredictor(vllm_config=vllm_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..3a7e33ad54 --- /dev/null +++ b/aphrodite/models/kimi_k3/common/mm_preprocess.py @@ -0,0 +1,330 @@ +# 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|>" + f"{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..65870e448f --- /dev/null +++ b/aphrodite/models/kimi_k3/common/mtp.py @@ -0,0 +1,96 @@ +# 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..fd85f2c142 --- /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 VllmConfig +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, + *, + vllm_config: VllmConfig, + config, + layer_idx: int, + start_layer_id: int, + prefix: str, + ) -> None: + super().__init__() + quant_config = get_draft_quant_config(vllm_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=vllm_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, + *, + vllm_config: VllmConfig, + start_layer_id: int, + prefix: str, + ) -> None: + super().__init__() + assert vllm_config.speculative_config is not None + self.config = vllm_config.speculative_config.draft_model_config.hf_config + self.quant_config = get_draft_quant_config(vllm_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( + vllm_config=vllm_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 = vllm_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, *, vllm_config: VllmConfig, prefix: str = "") -> None: + super().__init__() + assert vllm_config.speculative_config is not None + self.draft_model_config = vllm_config.speculative_config.draft_model_config + self.config = self.draft_model_config.hf_config + target_layer_num = vllm_config.model_config.get_num_layers(vllm_config.parallel_config) + self.model = K3DSparkModel( + vllm_config=vllm_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..3d99a9dc2c --- /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 VllmConfig +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 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, + vllm_config: VllmConfig, + prefix: str = "", + ) -> None: + super().__init__(config, vllm_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, + vllm_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 = vllm_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, + vllm_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 = vllm_config.compilation_config + if prefix in compilation_config.static_forward_context: + raise ValueError(f"Duplicate layer name: {prefix}") + compilation_config.static_forward_context[prefix] = self + + def 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..ea8afb3066 --- /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.vllm_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..d28820af71 --- /dev/null +++ b/aphrodite/models/kimi_k3/nvidia/mla.py @@ -0,0 +1,719 @@ +# 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 CacheConfig, VllmConfig, get_current_vllm_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) + + vllm_config = get_current_vllm_config() + parallel_config = vllm_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(vllm_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, + vllm_config=vllm_config, + ) + + compilation_config = vllm_config.compilation_config + if prefix in compilation_config.static_forward_context: + raise ValueError(f"Duplicate layer name: {prefix}") + compilation_config.static_forward_context[prefix] = self + self.kv_cache = torch.tensor([]) + + # ------------------------------------------------------------------ + # AttentionLayerBase interface + # ------------------------------------------------------------------ + def get_attn_backend(self) -> type[AttentionBackend]: + return self.attn_backend + + def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec: + kv_cache_dtype = kv_cache_dtype_str_to_dtype(self.kv_cache_dtype, vllm_config.model_config) + # TODO: Remove this mypy workaround once the K3 PR is fully merged. + return MLAAttentionSpec( # type: ignore[call-arg] + block_size=vllm_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..674cd87455 --- /dev/null +++ b/aphrodite/models/kimi_k3/nvidia/model.py @@ -0,0 +1,1720 @@ +# 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 VllmConfig +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, + vllm_config: VllmConfig, + 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 = vllm_config.kernel_config.moe_backend == "deep_gemm_mega_moe" + if self.use_mega_moe and not vllm_config.parallel_config.enable_expert_parallel: + raise NotImplementedError( + "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( + vllm_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, + vllm_config: VllmConfig, + 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 = vllm_config.cache_config + quant_config = vllm_config.quant_config + parallel_config = vllm_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 = vllm_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, + vllm_config, + prefix=f"{prefix}.self_attn", + ) + self._self_attn_writes_output = False + else: + self.self_attn = KimiLinearGatedDeltaNetAttention( + config, + vllm_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, + vllm_config=vllm_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, *, vllm_config: VllmConfig, prefix: str = ""): + super().__init__() + + config = vllm_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 = vllm_config.parallel_config + use_mega_moe = vllm_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, + vllm_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, *, vllm_config: VllmConfig, prefix: str = ""): + super().__init__() + self.model_config = vllm_config.model_config + self.vllm_config = vllm_config + self.config = self.model_config.hf_config + quant_config = vllm_config.quant_config + self.quant_config = quant_config + self.model = KimiLinearModel(vllm_config=vllm_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, + vllm_config: "VllmConfig", + ) -> tuple[torch.dtype, torch.dtype]: + return MambaStateDtypeCalculator.kda_state_dtype( + vllm_config.model_config.dtype, vllm_config.cache_config.mamba_cache_dtype + ) + + @classmethod + def get_mamba_state_shape_from_config( + cls, vllm_config: "VllmConfig" + ) -> tuple[tuple[int, int], tuple[int, int, int]]: + parallel_config = vllm_config.parallel_config + hf_config = vllm_config.model_config.hf_config + tp_size = parallel_config.tensor_parallel_size + num_spec = vllm_config.speculative_config.num_speculative_tokens if vllm_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, + vllm_config: VllmConfig, + prefix: str = "", + ) -> None: + super().__init__() + model_config = vllm_config.model_config + config: KimiK3Config = model_config.hf_config + self.config = config + self.model_config = model_config + quant_config = vllm_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(vllm_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=( + vllm_config.scheduler_config.max_num_seqs * mm_config.get_limit_per_prompt("image") + ), + max_seqlen=( + vllm_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(vllm_config): + self.language_model = init_vllm_registered_model( + vllm_config=vllm_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, vllm_config: VllmConfig) -> tuple[int, int]: + min_budget = 64 + max_budget = min( + vllm_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, vllm_config: VllmConfig): + text_config = vllm_config.model_config.hf_config.text_config + temp_vllm_config = vllm_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, vllm_config: VllmConfig): + text_config = vllm_config.model_config.hf_config.text_config + temp_vllm_config = vllm_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..485e051ba1 --- /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 VllmConfig +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, + vllm_config: VllmConfig, + prefix: str, + ) -> None: + super().__init__() + self.config = config + quant_config = vllm_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, vllm_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, *, vllm_config: VllmConfig, prefix: str = ""): + super().__init__() + config: KimiLinearConfig = vllm_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, vllm_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, *, vllm_config: VllmConfig, prefix: str = ""): + super().__init__() + self.config = vllm_config.model_config.hf_text_config + self.quant_config = vllm_config.quant_config + self.model = KimiK3MultiTokenPredictor(vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")) + enable_kimi_k3_low_latency_gemm(self, vllm_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..4f33f89206 --- /dev/null +++ b/aphrodite/models/kimi_k3/nvidia/ops/cute_dsl/latent_moe_tail/allreduce_rmsnorm_reduce_scatter_early_exit.py @@ -0,0 +1,1012 @@ +# 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..6a1ba9d77b --- /dev/null +++ b/aphrodite/models/kimi_k3/nvidia/ops/cute_dsl/latent_moe_tail/fused_add_multicast_gemm.py @@ -0,0 +1,1291 @@ +# 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..58d8f1ca18 --- /dev/null +++ b/aphrodite/models/kimi_k3/nvidia/ops/cute_dsl/latent_moe_tail/fused_add_multicast_skinny_gemm.py @@ -0,0 +1,438 @@ +# 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..e178e88c41 --- /dev/null +++ b/aphrodite/models/kimi_k3/nvidia/ops/cute_dsl/latent_moe_tail/lamport_copy.py @@ -0,0 +1,226 @@ +# 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..1bcb6d400c --- /dev/null +++ b/aphrodite/models/kimi_k3/nvidia/ops/cute_dsl/latent_moe_tail/primitives.py @@ -0,0 +1,437 @@ +# 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/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/third_party/flash_linear_attention/ops/kda.py.orig b/aphrodite/third_party/flash_linear_attention/ops/kda.py.orig new file mode 100644 index 0000000000..062d073b00 --- /dev/null +++ b/aphrodite/third_party/flash_linear_attention/ops/kda.py.orig @@ -0,0 +1,1591 @@ +# 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 code was licensed under the MIT license and included +# the following copyright notice: +# 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.triton_utils import tl, triton +from aphrodite.utils.math_utils import RCP_LN2, cdiv, next_power_of_2 + +from .chunk_delta_h import chunk_gated_delta_rule_fwd_h +from .cumsum import chunk_local_cumsum +from .fused_recurrent import fused_recurrent_gated_delta_rule_fwd_kernel +from .index import prepare_chunk_indices +from .l2norm import l2norm_fwd +from .op import exp2, log +from .solve_tril import solve_tril +from .utils import FLA_CHUNK_SIZE, is_amd + +BT_LIST_AUTOTUNE = [32, 64, 128] +NUM_WARPS_AUTOTUNE = [2, 4, 8, 16] if is_amd else [4, 8, 16, 32] + + +def fused_recurrent_kda_fwd( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + scale: float, + initial_state: torch.Tensor, + 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 = False, +) -> tuple[torch.Tensor, torch.Tensor]: + B, T, H, K, V = *k.shape, v.shape[-1] + HV = v.shape[2] + N = B if cu_seqlens is None else len(cu_seqlens) - 1 + BK, BV = next_power_of_2(K), min(next_power_of_2(V), 8) + NK, NV = cdiv(K, BK), cdiv(V, BV) + assert NK == 1, "NK > 1 is not supported yet" + num_stages = 3 + num_warps = 1 + + o = torch.empty_like(k) + if inplace_final_state: + final_state = initial_state + else: + final_state = q.new_empty(T, HV, V, K, dtype=initial_state.dtype) + + stride_init_state_token = initial_state.stride(0) + stride_final_state_token = final_state.stride(0) + + if ssm_state_indices is None: + stride_indices_seq, stride_indices_tok = 1, 1 + elif ssm_state_indices.ndim == 1: + stride_indices_seq, stride_indices_tok = ssm_state_indices.stride(0), 1 + else: + stride_indices_seq, stride_indices_tok = ssm_state_indices.stride() + + grid = (NK, NV, N * HV) + fused_recurrent_gated_delta_rule_fwd_kernel[grid]( + q=q, + k=k, + v=v, + g=g, + beta=beta, + o=o, + h0=initial_state, + ht=final_state, + cu_seqlens=cu_seqlens, + ssm_state_indices=ssm_state_indices, + num_accepted_tokens=num_accepted_tokens, + scale=scale, + N=N, + T=T, + B=B, + H=H, + HV=HV, + K=K, + V=V, + BK=BK, + BV=BV, + stride_init_state_token=stride_init_state_token, + stride_final_state_token=stride_final_state_token, + stride_indices_seq=stride_indices_seq, + stride_indices_tok=stride_indices_tok, + IS_BETA_HEADWISE=beta.ndim == v.ndim, + USE_QK_L2NORM_IN_KERNEL=use_qk_l2norm_in_kernel, + INPLACE_FINAL_STATE=inplace_final_state, + IS_KDA=True, + num_warps=num_warps, + num_stages=num_stages, + ) + + return o, final_state + + +def fused_recurrent_kda( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor = None, + scale: float = None, + initial_state: torch.Tensor = None, + inplace_final_state: bool = True, + use_qk_l2norm_in_kernel: bool = True, + cu_seqlens: torch.Tensor | None = None, + ssm_state_indices: torch.LongTensor | None = None, + **kwargs, +) -> tuple[torch.Tensor, torch.Tensor]: + if cu_seqlens is not None and q.shape[0] != 1: + raise ValueError( + f"The batch size is expected to be 1 rather than {q.shape[0]} when using `cu_seqlens`." + f"Please flatten variable-length inputs before processing." + ) + if scale is None: + scale = k.shape[-1] ** -0.5 + + o, final_state = fused_recurrent_kda_fwd( + q=q.contiguous(), + k=k.contiguous(), + v=v.contiguous(), + g=g.contiguous(), + beta=beta.contiguous(), + scale=scale, + initial_state=initial_state, + inplace_final_state=inplace_final_state, + cu_seqlens=cu_seqlens, + ssm_state_indices=ssm_state_indices, + num_accepted_tokens=None, + use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel, + ) + return o, final_state + + +@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 + 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, +): + i_t = tl.program_id(0) + + o_d = tl.arange(0, BD) + m_d = o_d < D + + 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 + 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) + 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, +): + 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 + + 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, +): + if residual is not None: + residual_dtype = residual.dtype + T, D = x.shape + 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.") + # heuristics for number of warps + + if D <= 512: + BT = 32 + 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, + D=D, + BD=BD, + BT=BT, + ACTIVATION=activation, + IS_RMS_NORM=is_rms_norm, + num_warps=4, + ) + 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, + ) + # 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]) + g = g.contiguous().reshape(-1, g.shape[-1]) + 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, + ) + 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/vllm/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, + ) + + +@triton.heuristics({"IS_VARLEN": lambda args: args["cu_seqlens"] is not None}) +@triton.autotune( + configs=[ + triton.Config({"BK": BK}, num_warps=num_warps, num_stages=num_stages) + for BK in [32, 64] + for num_warps in [1, 2, 4, 8] + for num_stages in [2, 3, 4] + ], + key=["BC"], +) +@triton.jit(do_not_specialize=["T"]) +def chunk_kda_scaled_dot_kkt_fwd_kernel_intra_sub_inter( + q, + k, + g, + beta, + A, + Aqk, + scale, + cu_seqlens, + chunk_indices, + T, + H: tl.constexpr, + K: tl.constexpr, + BT: tl.constexpr, + BC: tl.constexpr, + BK: tl.constexpr, + NC: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_t, i_c, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_b, i_h = i_bh // H, i_bh % H + i_i, i_j = i_c // NC, i_c % NC + if IS_VARLEN: + i_n, i_t = ( + tl.load(chunk_indices + i_t * 2).to(tl.int32), + tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32), + ) + 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 + i_i * BC >= T: + return + if i_i <= i_j: + return + + q += (bos * H + i_h) * K + k += (bos * H + i_h) * K + g += (bos * H + i_h) * K + A += (bos * H + i_h) * BT + Aqk += (bos * H + i_h) * BT + + p_b = tl.make_block_ptr(beta + bos * H + i_h, (T,), (H,), (i_t * BT + i_i * BC,), (BC,), (0,)) + b_b = tl.load(p_b, boundary_check=(0,)) + + b_A = tl.zeros([BC, BC], dtype=tl.float32) + b_Aqk = tl.zeros([BC, BC], dtype=tl.float32) + for i_k in range(tl.cdiv(K, BK)): + p_q = tl.make_block_ptr(q, (T, K), (H * K, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0)) + p_k = tl.make_block_ptr(k, (T, K), (H * K, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0)) + p_g = tl.make_block_ptr(g, (T, K), (H * K, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0)) + b_kt = tl.make_block_ptr(k, (K, T), (1, H * K), (i_k * BK, i_t * BT + i_j * BC), (BK, BC), (0, 1)) + p_gk = tl.make_block_ptr(g, (K, T), (1, H * K), (i_k * BK, i_t * BT + i_j * BC), (BK, BC), (0, 1)) + + o_k = i_k * BK + tl.arange(0, BK) + m_k = o_k < K + # [BK,] + b_gn = tl.load(g + (i_t * BT + i_i * BC) * H * K + o_k, mask=m_k, other=0) + # [BC, BK] + b_g = tl.load(p_g, boundary_check=(0, 1)) + b_k = tl.load(p_k, boundary_check=(0, 1)) * exp2(b_g - b_gn[None, :]) + # [BK, BC] + b_gk = tl.load(p_gk, boundary_check=(0, 1)) + b_kt = tl.load(b_kt, boundary_check=(0, 1)) + # [BC, BC] + b_ktg = b_kt * exp2(b_gn[:, None] - b_gk) + b_A += tl.dot(b_k, b_ktg) + + b_q = tl.load(p_q, boundary_check=(0, 1)) + b_qg = b_q * exp2(b_g - b_gn[None, :]) * scale + b_Aqk += tl.dot(b_qg, b_ktg) + + b_A *= b_b[:, None] + + p_A = tl.make_block_ptr(A, (T, BT), (H * BT, 1), (i_t * BT + i_i * BC, i_j * BC), (BC, BC), (1, 0)) + tl.store(p_A, b_A.to(A.dtype.element_ty), boundary_check=(0, 1)) + p_Aqk = tl.make_block_ptr(Aqk, (T, BT), (H * BT, 1), (i_t * BT + i_i * BC, i_j * BC), (BC, BC), (1, 0)) + tl.store(p_Aqk, b_Aqk.to(Aqk.dtype.element_ty), boundary_check=(0, 1)) + + +@triton.heuristics({"IS_VARLEN": lambda args: args["cu_seqlens"] is not None}) +@triton.autotune( + configs=[triton.Config({}, num_warps=num_warps) for num_warps in [1, 2, 4, 8]], + key=["BK", "BT"], +) +@triton.jit(do_not_specialize=["T"]) +def chunk_kda_scaled_dot_kkt_fwd_kernel_intra_sub_intra( + q, + k, + g, + beta, + A, + Aqk, + scale, + cu_seqlens, + chunk_indices, + T, + H: tl.constexpr, + K: tl.constexpr, + BT: tl.constexpr, + BC: tl.constexpr, + BK: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_t, i_i, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_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 + + if i_t * BT + i_i * BC >= T: + return + + o_i = tl.arange(0, BC) + o_k = tl.arange(0, BK) + m_k = o_k < K + m_A = (i_t * BT + i_i * BC + o_i) < T + o_A = (bos + i_t * BT + i_i * BC + o_i) * H * BT + i_h * BT + i_i * BC + + p_q = tl.make_block_ptr( + q + (bos * H + i_h) * K, + (T, K), + (H * K, 1), + (i_t * BT + i_i * BC, 0), + (BC, BK), + (1, 0), + ) + p_k = tl.make_block_ptr( + k + (bos * H + i_h) * K, + (T, K), + (H * K, 1), + (i_t * BT + i_i * BC, 0), + (BC, BK), + (1, 0), + ) + p_g = tl.make_block_ptr( + g + (bos * H + i_h) * K, + (T, K), + (H * K, 1), + (i_t * BT + i_i * BC, 0), + (BC, BK), + (1, 0), + ) + b_q = tl.load(p_q, boundary_check=(0, 1)) + b_k = tl.load(p_k, boundary_check=(0, 1)) + b_g = tl.load(p_g, boundary_check=(0, 1)) + + p_b = beta + (bos + i_t * BT + i_i * BC + o_i) * H + i_h + b_k = b_k * tl.load(p_b, mask=m_A, other=0)[:, None] + + p_kt = k + (bos + i_t * BT + i_i * BC) * H * K + i_h * K + o_k + p_gk = g + (bos + i_t * BT + i_i * BC) * H * K + i_h * K + o_k + + for j in range(0, min(BC, T - i_t * BT - i_i * BC)): + b_kt = tl.load(p_kt, mask=m_k, other=0).to(tl.float32) + b_gk = tl.load(p_gk, mask=m_k, other=0).to(tl.float32) + b_ktg = b_kt[None, :] * exp2(b_g - b_gk[None, :]) + b_A = tl.sum(b_k * b_ktg, 1) + b_A = tl.where(o_i > j, b_A, 0.0) + b_Aqk = tl.sum(b_q * b_ktg, 1) + b_Aqk = tl.where(o_i >= j, b_Aqk * scale, 0.0) + tl.store(A + o_A + j, b_A, mask=m_A) + tl.store(Aqk + o_A + j, b_Aqk, mask=m_A) + p_kt += H * K + p_gk += H * K + + +def chunk_kda_scaled_dot_kkt_fwd( + q: torch.Tensor, + k: torch.Tensor, + gk: torch.Tensor | None = None, + beta: torch.Tensor | None = None, + scale: 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 = torch.float32, +) -> tuple[torch.Tensor, torch.Tensor]: + r""" + Compute beta * K * K^T. + + Args: + k (torch.Tensor): + The key tensor of shape `[B, T, H, K]`. + beta (torch.Tensor): + The beta tensor of shape `[B, T, H]`. + gk (torch.Tensor): + The cumulative sum of the gate tensor of shape `[B, T, H, K]` applied to the key tensor. Default: `None`. + cu_seqlens (torch.Tensor): + The cumulative sequence lengths of the input tensor. + Default: None + chunk_size (int): + The chunk size. Default: 64. + output_dtype (torch.dtype): + The dtype of the output tensor. Default: `torch.float32` + + Returns: + beta * K * K^T of shape `[B, T, H, BT]` where `BT` is the chunk size. + """ + B, T, H, K = k.shape + assert K <= 256 + BT = chunk_size + 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) + + BC = min(16, BT) + NC = cdiv(BT, BC) + BK = max(next_power_of_2(K), 16) + A = torch.zeros(B, T, H, BT, device=k.device, dtype=output_dtype) + Aqk = torch.zeros(B, T, H, BT, device=k.device, dtype=output_dtype) + grid = (NT, NC * NC, B * H) + chunk_kda_scaled_dot_kkt_fwd_kernel_intra_sub_inter[grid]( + q=q, + k=k, + g=gk, + beta=beta, + A=A, + Aqk=Aqk, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + T=T, + H=H, + K=K, + BT=BT, + BC=BC, + NC=NC, + ) + + grid = (NT, NC, B * H) + chunk_kda_scaled_dot_kkt_fwd_kernel_intra_sub_intra[grid]( + q=q, + k=k, + g=gk, + beta=beta, + A=A, + Aqk=Aqk, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + T=T, + H=H, + K=K, + BT=BT, + BC=BC, + BK=BK, + ) + return A, Aqk + + +@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,)) + + 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({"BD": BD}, num_warps=num_warps) for BD in [32, 64] for num_warps in [2, 4, 8]], + key=["H", "D", "BT", "IS_VARLEN"], +) +@triton.jit(do_not_specialize=["T"]) +def kda_gate_cumsum_fwd_kernel( + g, + A, + y, + g_bias, + cu_seqlens, + chunk_indices, + cumsum_scale, + beta, + threshold, + T, + H: tl.constexpr, + D: tl.constexpr, + BT: tl.constexpr, + BD: tl.constexpr, + HAS_BIAS: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_d, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_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 + + p_g = tl.make_block_ptr( + g + (bos * H + i_h) * D, + (T, D), + (H * D, 1), + (i_t * BT, i_d * BD), + (BT, BD), + (1, 0), + ) + p_y = tl.make_block_ptr( + y + (bos * H + i_h) * D, + (T, D), + (H * D, 1), + (i_t * BT, i_d * BD), + (BT, BD), + (1, 0), + ) + + b_g = tl.load(p_g, boundary_check=(0, 1)).to(tl.float32) + if HAS_BIAS: + o_d = i_d * BD + tl.arange(0, BD) + b_bias = tl.load(g_bias + i_h * D + o_d, mask=o_d < D, other=0.0).to(tl.float32) + b_g = b_g + b_bias[None, :] + + 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 + + # Out-of-bounds rows (load returns 0, but softplus/bias can still make + # b_gate non-zero) participate in the dot product. They only contribute to + # out-of-bounds output rows, which are masked away by `boundary_check` on + # the store, so visible output matches unfused gate + chunk-local cumsum. + o_t = tl.arange(0, BT) + m_cumsum = tl.where(o_t[:, None] >= o_t[None, :], 1.0, 0.0) + b_y = tl.dot(m_cumsum, b_gate, allow_tf32=False) * cumsum_scale + tl.store(p_y, b_y.to(p_y.dtype.element_ty), boundary_check=(0, 1)) + + +def fused_kda_gate_chunk_cumsum( + raw_g: torch.Tensor, + A_log: torch.Tensor, + g_bias: torch.Tensor | None = None, + beta: float = 1.0, + threshold: float = 20.0, + 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, +) -> 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 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) + + def grid(meta): + return (cdiv(meta["D"], meta["BD"]), NT, B * H) + + kda_gate_cumsum_fwd_kernel[grid]( + g=raw_g, + A=A_log, + y=y, + g_bias=g_bias, + 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, + beta=beta, + threshold=threshold, + T=T, + H=H, + D=D, + BT=chunk_size, + ) + return y + + +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, +): + # `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. + # the intra Aqk is kept in fp32 + # the computation has very marginal effect on the entire throughput + A, Aqk = chunk_kda_scaled_dot_kkt_fwd( + q=q, + k=k, + gk=g, + beta=beta, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + chunk_size=chunk_size, + output_dtype=torch.float32, + ) + A = solve_tril(A=A, cu_seqlens=cu_seqlens, output_dtype=k.dtype) + 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, + beta: torch.Tensor, + A_log: torch.Tensor, + g_bias: torch.Tensor | None, + 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 = fused_kda_gate_chunk_cumsum( + raw_g, + A_log=A_log, + g_bias=g_bias, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + chunk_size=chunk_size, + ) + 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( + 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, + 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, + use_qk_l2norm_in_kernel: bool = False, + cu_seqlens: torch.Tensor | None = None, + **kwargs, +): + """Run chunk KDA from raw gate projection using fused gate+cumsum.""" + 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(), + beta=beta.contiguous(), + 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, + 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, + beta: tl.constexpr, + threshold: tl.constexpr, + T, + H, + D: tl.constexpr, + BT: tl.constexpr, + BD: tl.constexpr, + HAS_BIAS: 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) + + 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, :] + + # 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 + + 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, +) -> 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, + beta, + threshold, + T, + H, + head_k_dim, + BD=next_power_of_2(head_k_dim), + HAS_BIAS=g_bias is not None, + ) + + y = y.view(*orig_shape, H, head_k_dim) + return y 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/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/attention/ops/triton_merge_attn_states.py.orig b/aphrodite/v1/attention/ops/triton_merge_attn_states.py.orig new file mode 100644 index 0000000000..3fc8c66774 --- /dev/null +++ b/aphrodite/v1/attention/ops/triton_merge_attn_states.py.orig @@ -0,0 +1,267 @@ +# 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 + +float8_info = torch.finfo(current_platform.fp8_dtype()) + + +def mask_empty_context( + lse: torch.Tensor, + output: torch.Tensor, + query_start_loc: torch.Tensor, + context_start_loc: torch.Tensor, +) -> None: + """Neutralize context chunks that cover no keys before merging. + + A prefill query whose context chunk is empty attended to no keys, so its + partial attention is undefined: the backend leaves the output rows as + uninitialized scratch (which may hold NaN/Inf) even when it reports an LSE + of -inf. Sanitize both here so ``merge_attn_states`` can stay generic: + force the LSE to -inf (zero softmax weight) and zero the undefined output + rows (so a zero weight cannot combine with NaN/Inf). Emptiness is derived + from the context offsets, not from the -inf LSE, so no merge kernel has to + reason about undefined partials. + + Args: + lse: Chunk log-sum-exp, shape [num_heads, num_tokens]. + output: Chunk attention output, shape [num_tokens, num_heads, ...]. + query_start_loc: Prefill query cumulative offsets, shape [num_reqs + 1]. + context_start_loc: Chunk context cumulative offsets, + shape [num_reqs + 1]; an empty chunk has a zero-length span. + """ + num_heads, num_tokens = lse.shape + num_reqs = query_start_loc.shape[0] - 1 + block_size = 128 + num_query_blocks = num_tokens // block_size + num_reqs + is_empty = torch.zeros(num_tokens, dtype=torch.bool, device=lse.device) + mask_empty_context_kernel[(num_query_blocks,)]( + lse, + is_empty, + query_start_loc, + context_start_loc, + lse.stride(0), + lse.stride(1), + num_reqs, + NUM_HEADS=num_heads, + BLOCK_SIZE=block_size, + BLOCK_HEADS=8, + num_warps=8, + ) + output.masked_fill_(is_empty[:, None, None], 0.0) + + +@triton.jit +def mask_empty_context_kernel( + lse, + is_empty, + query_start_loc, + context_start_loc, + lse_head_stride, + lse_token_stride, + num_reqs, + NUM_HEADS: tl.constexpr, + BLOCK_SIZE: tl.constexpr, + BLOCK_HEADS: tl.constexpr, +): + query_block_idx = tl.program_id(0) + + lanes = tl.arange(0, 32) + chunk_start = 0 + req_idx = 0 + req_idx_found = False + while (chunk_start < num_reqs) & (not req_idx_found): + req_offsets = chunk_start + lanes + req_mask = req_offsets < num_reqs + query_starts = tl.load(query_start_loc + req_offsets, mask=req_mask) + req_block_starts = query_starts // BLOCK_SIZE + req_offsets + matched_idx = tl.sum((req_mask & (req_block_starts <= query_block_idx)).to(tl.int32)) + req_idx = chunk_start + matched_idx - 1 + req_idx_found = matched_idx < 32 + chunk_start += 32 + + query_start = tl.load(query_start_loc + req_idx) + query_end = tl.load(query_start_loc + req_idx + 1) + query_len = query_end - query_start + req_first_block = query_start // BLOCK_SIZE + req_idx + block_in_req = query_block_idx - req_first_block + token_offset = block_in_req * BLOCK_SIZE + if token_offset >= query_len: + return + + context_start = tl.load(context_start_loc + req_idx) + context_end = tl.load(context_start_loc + req_idx + 1) + if context_start != context_end: + return + + token_offsets = token_offset + tl.arange(0, BLOCK_SIZE) + token_indices = query_start + token_offsets + token_lse_offsets = token_indices * lse_token_stride + valid_tokens = token_offsets < query_len + tl.store(is_empty + token_indices, True, mask=valid_tokens) + head_offsets = tl.arange(0, BLOCK_HEADS) + for head_start in range(0, NUM_HEADS, BLOCK_HEADS): + head_indices = head_start + head_offsets + lse_ptrs = lse + head_indices[:, None] * lse_head_stride + token_lse_offsets[None, :] + valid_heads = head_indices < NUM_HEADS + tl.store( + lse_ptrs, + float("-inf"), + mask=valid_heads[:, None] & valid_tokens[None, :], + ) + + +# Implements section 2.2 of https://www.arxiv.org/pdf/2501.01005 +# can be used to combine partial attention results (in the split-KV case) +def merge_attn_states( + output: torch.Tensor, + prefix_output: torch.Tensor, + prefix_lse: torch.Tensor, + suffix_output: torch.Tensor, + suffix_lse: torch.Tensor, + output_lse: torch.Tensor | None = None, + prefill_tokens_with_context: int | None = None, + output_scale: torch.Tensor | None = None, +) -> None: + num_tokens = output.shape[0] + num_query_heads = output.shape[1] + head_size = output.shape[2] + padded_head_size = triton.next_power_of_2(head_size) + # We assume the output stride on num_head is not always as same as the + # `suffix_output` and `prefix_output`, as them might be padded by the + # attention backend. + prefix_head_stride = prefix_output.stride(1) + output_head_stride = output.stride(1) + + # If prefill_tokens_with_context is None, all tokens should use prefix context + if prefill_tokens_with_context is None: + prefill_tokens_with_context = num_tokens + + # TODO(woosuk): Use CUDA kernel instead of Triton to minimize CPU overhead. + merge_attn_states_kernel[(num_tokens, num_query_heads)]( + output, + output_lse, + prefix_output, + prefix_lse, + suffix_output, + suffix_lse, + prefix_head_stride, + output_head_stride, + output_scale, + prefill_tokens_with_context, + head_size, + padded_head_size, + output_lse is not None, + output_scale is not None, + ) + + +@triton.jit +def merge_attn_states_kernel( + output, # [NUM_TOKENS, NUM_HEADS, HEAD_SIZE] + output_lse, # [NUM_HEADS, NUM_TOKENS] + prefix_output, # [NUM_TOKENS, NUM_HEADS, HEAD_SIZE] + prefix_lse, # [NUM_HEADS, NUM_TOKENS] + suffix_output, # [NUM_TOKENS, NUM_HEADS, HEAD_SIZE] + suffix_lse, # [NUM_HEADS, NUM_TOKENS] + prefix_head_stride, + output_head_stride, + output_scale, # scale tensor or None + prefill_tokens_with_context, + HEAD_SIZE: tl.constexpr, + PADDED_HEAD_SIZE: tl.constexpr, + OUTPUT_LSE: tl.constexpr, + USE_FP8: tl.constexpr, + FP8_MIN: tl.constexpr = float8_info.min, + 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) + + prefix_mask = token_idx < prefill_tokens_with_context + + head_arange = tl.arange(0, PADDED_HEAD_SIZE) + head_mask = head_arange < HEAD_SIZE + + # 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) + if OUTPUT_LSE: + tl.store(output_lse + head_idx * num_tokens + token_idx, s_lse) + + s_out = tl.load( + suffix_output + token_idx * num_heads * prefix_head_stride + head_idx * prefix_head_stride + head_arange, + mask=head_mask, + ) + + if USE_FP8: + s_out = s_out * (1.0 / tl.load(output_scale)) + s_out = tl.clamp(s_out, FP8_MIN, FP8_MAX) + s_out = s_out.to(output.dtype.element_ty) + + tl.store( + output + token_idx * num_heads * output_head_stride + head_idx * output_head_stride + head_arange, + s_out, + mask=head_mask, + ) + return + + # 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) + + # 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. + # If we see an inf assume FA2 and convert inf to -inf for consistency + # and correctness. Inf generally doesn't make sense in this context outside + # of undefined-behavior/FA2-case, so I think this a safe assumption. + p_lse = float("-inf") if p_lse == float("inf") else p_lse + s_lse = float("-inf") if s_lse == float("inf") else s_lse + + max_lse = tl.maximum(p_lse, s_lse) + p_lse = p_lse - max_lse + s_lse = s_lse - max_lse + # Will reuse precomputed Exp values for scale factor computation. + p_se = tl.exp(p_lse) + s_se = tl.exp(s_lse) + out_se = p_se + s_se + + 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) + + p_out = tl.load( + prefix_output + token_idx * num_heads * prefix_head_stride + head_idx * prefix_head_stride + head_arange, + mask=head_mask, + ) + s_out = tl.load( + suffix_output + token_idx * num_heads * prefix_head_stride + head_idx * prefix_head_stride + head_arange, + mask=head_mask, + ) + + # NOTE(woosuk): Be careful with the numerical stability. + # We should compute the scale first, and then multiply it with the output. + # Do not multiply the output with tl.exp(p_lse) or tl.exp(s_lse) directly. + p_scale = p_se / out_se + s_scale = s_se / out_se + out = p_out * p_scale + s_out * s_scale + out = tl.where(max_lse == float("-inf"), 0.0, out) + + if USE_FP8: + out = out * (1.0 / tl.load(output_scale)) + out = tl.clamp(out, FP8_MIN, FP8_MAX) + out = out.to(output.dtype.element_ty) + + tl.store( + output + token_idx * num_heads * output_head_stride + head_idx * output_head_stride + head_arange, + out, + mask=head_mask, + ) 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..f970422c70 --- /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 VllmConfig, set_current_vllm_config + + with set_current_vllm_config(VllmConfig()): + 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/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_all_reduce.cuh.orig b/csrc/custom_all_reduce.cuh.orig new file mode 100644 index 0000000000..a08b25820e --- /dev/null +++ b/csrc/custom_all_reduce.cuh.orig @@ -0,0 +1,633 @@ +#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) { + 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 + // for all ranks, ensuring bitwise identical results + auto dp = *_dp; + barrier_at_start(sg, self_sg, rank); + // do the actual reduction + for (int idx = blockIdx.x * blockDim.x + threadIdx.x; idx < size; + idx += gridDim.x * blockDim.x) { + ((P*)result)[idx] = packed_reduce((const P**)&dp.ptrs[0], idx); + } + barrier_at_end(sg, self_sg, rank); +} + +template +DINLINE P* get_tmp_buf(Signal* sg) { + return (P*)(((Signal*)sg) + 1); +} + +template +__global__ void __launch_bounds__(512, 1) + cross_device_reduce_2stage(RankData* _dp, RankSignals sg, Signal* self_sg, + T* __restrict__ result, int rank, int size) { + int tid = blockIdx.x * blockDim.x + threadIdx.x; + int stride = gridDim.x * blockDim.x; + using P = typename packed_t::P; + using A = typename packed_t::A; + int part = size / ngpus; + int start = rank * part; + int end = rank == ngpus - 1 ? size : start + part; + int largest_part = part + size % ngpus; + const P* ptrs[ngpus]; + P* tmps[ngpus]; +#pragma unroll + for (int i = 0; i < ngpus; i++) { + int target = (rank + i) % ngpus; + ptrs[i] = (const P*)_dp->ptrs[target]; + tmps[i] = get_tmp_buf

(sg.signals[target]); + } + auto tmp_out = tmps[0]; + barrier_at_start(sg, self_sg, rank); + + // stage 1: reduce scatter + for (int idx = start + tid; idx < end; idx += stride) { + tmp_out[idx - start] = packed_reduce(ptrs, idx); + } + barrier_at_end(sg, self_sg, rank); + + // stage 2: allgather. Note: it's important to match the tid between + // the two stages, because visibility across devices is only guaranteed + // between threads that have the same tid. If thread i computes the sum of + // start + i in the first stage, then thread i also gathers start + i from + // all ranks. + + for (int idx = tid; idx < largest_part; idx += stride) { +#pragma unroll + for (int i = 0; i < ngpus; i++) { + int gather_from_rank = ((rank + i) % ngpus); + if (gather_from_rank == ngpus - 1 || idx < part) { + int dst_idx = gather_from_rank * part + idx; + ((P*)result)[dst_idx] = tmps[i][idx]; + } + } + } +} + +using IPC_KEY = std::array; +static_assert(sizeof(IPC_KEY) == sizeof(cudaIpcMemHandle_t)); +static_assert(alignof(IPC_KEY) == alignof(cudaIpcMemHandle_t)); + +class CustomAllreduce { + public: + int rank_; + int world_size_; + // Full NVLink or xGMI connection between GPUs. + bool fully_connected_; + + RankSignals sg_; + // Stores a map from a pointer to its peer pointers from all ranks. + std::unordered_map buffers_; + Signal* self_sg_; + + // Stores rank data from all ranks. This is mainly for cuda graph purposes. + // For cuda graph to work, all kernel arguments must be fixed during graph + // capture time. However, the peer pointers are not known during graph + // capture time. Therefore, during capture, we increment the rank data + // pointer and use that as the argument to the kernel. The kernel arguments + // are stored in graph_unreg_buffers_. The actual peer pointers will be + // filled in at the memory pointed to by the pointers in + // graph_unreg_buffers_ when the IPC handles are exchanged between ranks. + // + // The overall process looks like this: + // 1. Graph capture. + // 2. Each rank obtains the IPC handles for each addresses used during cuda + // graph capture using get_graph_buffer_ipc_meta. + // 3. (In Python) all gather the IPC handles. + // 4. Obtain the peer pointers by opening the IPC handles, and store them in + // the rank data array at corresponding positions. + RankData *d_rank_data_base_, *d_rank_data_end_; + std::vector graph_unreg_buffers_; + // a map from IPC handles to opened IPC pointers + std::map ipc_handles_; + + /** + * Signals are an array of ipc-enabled buffers from all ranks. + * For each of the buffer, the layout is as follows: + * | -- sizeof(Signal) -- | ------ a few MB ----- | + * The first section is for allreduce synchronization, and the second + * section is for storing the intermediate results required by some + * allreduce algos. + * + * Note: this class does not own any device memory. Any required buffers + * are passed in from the constructor. + */ + CustomAllreduce(Signal** signals, void* rank_data, size_t rank_data_sz, + int rank, int world_size, bool fully_connected = true) + : rank_(rank), + world_size_(world_size), + fully_connected_(fully_connected), + self_sg_(signals[rank]), + d_rank_data_base_(reinterpret_cast(rank_data)), + d_rank_data_end_(d_rank_data_base_ + rank_data_sz / sizeof(RankData)) { + for (int i = 0; i < world_size_; i++) { + sg_.signals[i] = signals[i]; + } + } + + char* open_ipc_handle(const void* ipc_handle) { + auto [it, new_handle] = + ipc_handles_.insert({*((IPC_KEY*)ipc_handle), nullptr}); + if (new_handle) { + char* ipc_ptr; + CUDACHECK(cudaIpcOpenMemHandle((void**)&ipc_ptr, + *((const cudaIpcMemHandle_t*)ipc_handle), + cudaIpcMemLazyEnablePeerAccess)); + it->second = ipc_ptr; + } + return it->second; + } + + std::pair> get_graph_buffer_ipc_meta() { + auto num_buffers = graph_unreg_buffers_.size(); + auto handle_sz = sizeof(cudaIpcMemHandle_t); + std::string handles(handle_sz * num_buffers, static_cast(0)); + std::vector offsets(num_buffers); + for (int i = 0; i < num_buffers; i++) { + auto ptr = graph_unreg_buffers_[i]; + void* base_ptr; + // note: must share the base address of each allocation, or we get wrong + // address + if (cuPointerGetAttribute(&base_ptr, rangeStartAddrAttr, + (CUdeviceptr)ptr) != CUDA_SUCCESS) + throw std::runtime_error("failed to get pointer attr"); + CUDACHECK(cudaIpcGetMemHandle( + (cudaIpcMemHandle_t*)&handles[i * handle_sz], base_ptr)); + offsets[i] = ((char*)ptr) - ((char*)base_ptr); + } + return std::make_pair(handles, offsets); + } + + void check_rank_data_capacity(size_t num = 1) { + if (d_rank_data_base_ + num > d_rank_data_end_) + throw std::runtime_error( + "Rank data buffer is overflowed by " + + std::to_string(d_rank_data_base_ + num - d_rank_data_end_)); + } + + /** + * Register already-shared IPC pointers. + */ + void register_buffer(void** ptrs) { + check_rank_data_capacity(); + RankData data; + for (int i = 0; i < world_size_; i++) { + data.ptrs[i] = ptrs[i]; + } + auto d_data = d_rank_data_base_++; + CUDACHECK( + cudaMemcpy(d_data, &data, sizeof(RankData), cudaMemcpyHostToDevice)); + buffers_[ptrs[rank_]] = d_data; + } + + // Note: when registering graph buffers, we intentionally choose to not + // deduplicate the addresses. That means if the allocator reuses some + // addresses, they will be registered again. This is to account for the + // remote possibility of different allocation patterns between ranks. For + // example, rank 1 may get the same input address for the second allreduce, + // but rank 2 got a different address. IPC handles have internal reference + // counting mechanism so overhead should be small. + void register_graph_buffers( + const std::vector& handles, + const std::vector>& offsets) { + auto num_buffers = graph_unreg_buffers_.size(); + check_rank_data_capacity(num_buffers); + std::vector rank_data(num_buffers); + for (int i = 0; i < num_buffers; i++) { + auto self_ptr = graph_unreg_buffers_[i]; + auto& rd = rank_data[i]; + for (int j = 0; j < world_size_; j++) { + if (j != rank_) { + char* handle = + open_ipc_handle(&handles[j][i * sizeof(cudaIpcMemHandle_t)]); + handle += offsets[j][i]; + rd.ptrs[j] = handle; + } else { + rd.ptrs[j] = self_ptr; + } + } + } + CUDACHECK(cudaMemcpy(d_rank_data_base_, rank_data.data(), + sizeof(RankData) * num_buffers, + cudaMemcpyHostToDevice)); + d_rank_data_base_ += num_buffers; + graph_unreg_buffers_.clear(); + } + + /** + * Performs allreduce, assuming input has already been registered. + * + * Block and grid default configs are results after careful grid search. + * Using 36 blocks give the best or close to the best runtime on the devices + * I tried: A100, A10, A30, T4, V100. You'll notice that NCCL kernels also + * only take a small amount of SMs. Not quite sure the underlying reason, + * but my guess is that too many SMs will cause contention on NVLink bus. + */ + template + void allreduce(cudaStream_t stream, T* input, T* output, int size, + int threads = 512, int block_limit = defaultBlockLimit) { + auto d = packed_t::P::size; + if (size % d != 0) + throw std::runtime_error( + "custom allreduce currently requires input length to be multiple " + "of " + + std::to_string(d)); + if (block_limit > kMaxBlocks) + throw std::runtime_error("max supported block limit is " + + std::to_string(kMaxBlocks) + ". Got " + + std::to_string(block_limit)); + + RankData* ptrs; + cudaStreamCaptureStatus status; + CUDACHECK(cudaStreamIsCapturing(stream, &status)); + if (status == cudaStreamCaptureStatusActive) { + ptrs = d_rank_data_base_ + graph_unreg_buffers_.size(); + graph_unreg_buffers_.push_back(input); + } else { + auto it = buffers_.find(input); + if (it == buffers_.end()) + throw std::runtime_error( + "buffer address " + + std::to_string(reinterpret_cast(input)) + + " is not registered!"); + ptrs = it->second; + } + + size /= d; + auto bytes = size * sizeof(typename packed_t::P); + int blocks = std::min(block_limit, (size + threads - 1) / threads); + + // Check environment variable once + const char* env_algo = std::getenv("APHRODITE_CUSTOM_ALLREDUCE_ALGO"); + bool force_1stage = false; + bool force_2stage = false; + if (env_algo != nullptr) { + if (std::strcmp(env_algo, "1stage") == 0 || + std::strcmp(env_algo, "oneshot") == 0) { + force_1stage = true; + } else if (std::strcmp(env_algo, "2stage") == 0 || + std::strcmp(env_algo, "twoshot") == 0) { + force_2stage = true; + } else { + throw std::runtime_error( + "Invalid APHRODITE_CUSTOM_ALLREDUCE_ALGO: " + + std::string(env_algo) + + ". Valid values: 1stage, oneshot, 2stage, twoshot"); + } + } + +#define KL(ngpus, name) \ + name<<>>(ptrs, sg_, self_sg_, output, \ + rank_, size); +#define REDUCE_CASE(ngpus) \ + case ngpus: { \ + if (force_1stage) { \ + KL(ngpus, cross_device_reduce_1stage); \ + } else if (force_2stage) { \ + KL(ngpus, cross_device_reduce_2stage); \ + } else { \ + if (world_size_ == 2) { \ + KL(ngpus, cross_device_reduce_1stage); \ + } else if (fully_connected_) { \ + if ((world_size_ <= 4 && bytes < 512 * 1024) || \ + (world_size_ <= 8 && bytes < 256 * 1024)) { \ + KL(ngpus, cross_device_reduce_1stage); \ + } else { \ + KL(ngpus, cross_device_reduce_2stage); \ + } \ + } \ + } \ + break; \ + } + + switch (world_size_) { + REDUCE_CASE(2) + REDUCE_CASE(4) + REDUCE_CASE(6) + REDUCE_CASE(8) + default: + throw std::runtime_error( + "custom allreduce only supports num gpus in (2,4,6,8). Actual " + "num " + "gpus = " + + std::to_string(world_size_)); + } +#undef REDUCE_CASE +#undef KL + } + + ~CustomAllreduce() { + for (auto [_, ptr] : ipc_handles_) { + CUDACHECK(cudaIpcCloseMemHandle(ptr)); + } + } +}; + +/** + * To inspect PTX/SASS, copy paste this header file to compiler explorer and + add a template instantiation: + * template void aphrodite::CustomAllreduce::allreduce(cudaStream_t, half + *, half *, int, int, int); +*/ +} // namespace aphrodite \ No newline at end of file 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/activation_kernels.cu.orig b/csrc/libtorch_stable/activation_kernels.cu.orig new file mode 100644 index 0000000000..6e3e6860e6 --- /dev/null +++ b/csrc/libtorch_stable/activation_kernels.cu.orig @@ -0,0 +1,724 @@ +#include +#include + +#include + +#include "../cuda_compat.h" +#include "cuda_vec_utils.cuh" +#include "dispatch_utils.h" +#include "torch_utils.h" + +namespace aphrodite { + +// `alpha` and `beta` are applied to opposite operands: +// - alpha lives INSIDE the activation (the activated half): the gated +// activation computes act_half * sigmoid(alpha * act_half). +// - beta is added to the OTHER (non-activated) half before the multiply. +// So the result is always ACT(act_half, alpha) * (other_half + beta). +// Which half is which depends on `act_first` (see below). Defaults +// alpha=1.0, beta=0.0 reproduce the plain SwiGLU/GeGLU behavior. +template +__device__ __forceinline__ scalar_t compute(const scalar_t& x, + const scalar_t& y, + const float limit, + const float alpha, + const float beta) { + if constexpr (act_first) { + scalar_t gate = x; + scalar_t up = y; + if constexpr (HAS_CLAMP) { + gate = (scalar_t)fminf((float)gate, limit); + up = (scalar_t)fmaxf(fminf((float)up, limit), -limit); + } + // act_first: gate is the activated half -> alpha applies to gate; + // beta is added to up (the non-activated half). + return (scalar_t)(ACT_FN(gate, alpha) * ((float)up + beta)); + } else { + scalar_t gate = x; + scalar_t up = y; + if constexpr (HAS_CLAMP) { + gate = (scalar_t)fmaxf(fminf((float)gate, limit), -limit); + up = (scalar_t)fminf((float)up, limit); + } + // !act_first: up is the activated half -> alpha applies to up; + // beta is added to gate (the non-activated half). + return (scalar_t)(((float)gate + beta) * ACT_FN(up, alpha)); + } +} + +template +__device__ __forceinline__ packed_t packed_compute(const packed_t& x, + const packed_t& y, + const float limit, + const float alpha, + const float beta) { + if constexpr (act_first) { + packed_t gate = x; + packed_t up = y; + float2 u = cast_to_float2(up); + if constexpr (HAS_CLAMP) { + float2 g = cast_to_float2(gate); + g.x = fminf(g.x, limit); + g.y = fminf(g.y, limit); + u.x = fmaxf(fminf(u.x, limit), -limit); + u.y = fmaxf(fminf(u.y, limit), -limit); + gate = cast_to_packed(g); + } + // act_first: gate is the activated half -> alpha applies to gate; + // beta is added to up (the non-activated half). + float2 activated = cast_to_float2(PACKED_ACT_FN(gate, alpha)); + activated.x *= u.x + beta; + activated.y *= u.y + beta; + return cast_to_packed(activated); + } else { + packed_t gate = x; + packed_t up = y; + float2 g = cast_to_float2(gate); + if constexpr (HAS_CLAMP) { + float2 u = cast_to_float2(up); + g.x = fmaxf(fminf(g.x, limit), -limit); + g.y = fmaxf(fminf(g.y, limit), -limit); + u.x = fminf(u.x, limit); + u.y = fminf(u.y, limit); + up = cast_to_packed(u); + } + // !act_first: up is the activated half -> alpha applies to up; + // beta is added to gate (the non-activated half). + float2 activated = cast_to_float2(PACKED_ACT_FN(up, alpha)); + activated.x *= g.x + beta; + activated.y *= g.y + beta; + return cast_to_packed(activated); + } +} + +// Activation and gating kernel template. +template +__global__ void act_and_mul_kernel( + scalar_t* __restrict__ out, // [..., d] + const scalar_t* __restrict__ input, // [..., 2, d] + const int d, const float limit, const float alpha, const float beta) { + const scalar_t* x_ptr = input + blockIdx.x * 2 * d; + const scalar_t* y_ptr = x_ptr + d; + scalar_t* out_ptr = out + blockIdx.x * d; + + if constexpr (use_vec) { + using cuda_t = typename CUDATypeConverter::Type; + using pvec_t = PackedVec; + + const pvec_t* x_vec = reinterpret_cast(x_ptr); + const pvec_t* y_vec = reinterpret_cast(y_ptr); + pvec_t* out_vec = reinterpret_cast(out_ptr); + const int num_vecs = d / 2 / pvec_t::NUM_ELTS; + + for (int i = threadIdx.x; i < num_vecs; i += blockDim.x) { + pvec_t x, y; + if constexpr (use_256b) { + ld256(x, &x_vec[i]); + ld256(y, &y_vec[i]); + } else { + ld128(x, &x_vec[i]); + ld128(y, &y_vec[i]); + } +#pragma unroll + for (int j = 0; j < pvec_t::NUM_ELTS; j++) { + x.elts[j] = + packed_compute( + x.elts[j], y.elts[j], limit, alpha, beta); + } + if constexpr (use_256b) { + st256(x, &out_vec[i]); + } else { + st128(x, &out_vec[i]); + } + } + } else { + // Scalar fallback for unaligned data or small d + for (int64_t idx = threadIdx.x; idx < d; idx += blockDim.x) { + const scalar_t x = APHRODITE_LDG(&x_ptr[idx]); + const scalar_t y = APHRODITE_LDG(&y_ptr[idx]); + out_ptr[idx] = compute( + x, y, limit, alpha, beta); + } + } +} + +// Gated activations take an `alpha` argument that scales the sigmoid input +// (`x * sigmoid(alpha * x)`). alpha defaults to 1.0 at all call sites, which +// is exactly SiLU; only the clamp path (silu_and_mul_with_clamp) passes a +// non-default alpha. Activations that do not use alpha simply ignore it. +template +__device__ __forceinline__ T silu_kernel(const T& x, const float alpha) { + // x * sigmoid(alpha * x) + return (T)(((float)x) / (1.0f + expf((float)-x * alpha))); +} + +template +__device__ __forceinline__ packed_t packed_silu_kernel(const packed_t& val, + const float alpha) { + // x * sigmoid(alpha * x) + float2 fval = cast_to_float2(val); + fval.x = fval.x / (1.0f + expf(-fval.x * alpha)); + fval.y = fval.y / (1.0f + expf(-fval.y * alpha)); + return cast_to_packed(fval); +} + +template +__device__ __forceinline__ T gelu_kernel(const T& x, const float /*alpha*/) { + // Equivalent to PyTorch GELU with 'none' approximation. + // Refer to: + // https://github.com/pytorch/pytorch/blob/8ac9b20d4b090c213799e81acf48a55ea8d437d6/aten/src/ATen/native/cuda/ActivationGeluKernel.cu#L36-L38 + const float f = (float)x; + constexpr float ALPHA = M_SQRT1_2; + return (T)(f * 0.5f * (1.0f + ::erf(f * ALPHA))); +} + +template +__device__ __forceinline__ packed_t packed_gelu_kernel(const packed_t& val, + const float /*alpha*/) { + // Equivalent to PyTorch GELU with 'none' approximation. + // Refer to: + // https://github.com/pytorch/pytorch/blob/8ac9b20d4b090c213799e81acf48a55ea8d437d6/aten/src/ATen/native/cuda/ActivationGeluKernel.cu#L36-L38 + constexpr float ALPHA = M_SQRT1_2; + float2 fval = cast_to_float2(val); + fval.x = fval.x * 0.5f * (1.0f + ::erf(fval.x * ALPHA)); + fval.y = fval.y * 0.5f * (1.0f + ::erf(fval.y * ALPHA)); + return cast_to_packed(fval); +} + +template +__device__ __forceinline__ T gelu_tanh_kernel(const T& x, + const float /*alpha*/) { + // Equivalent to PyTorch GELU with 'tanh' approximation. + // Refer to: + // https://github.com/pytorch/pytorch/blob/8ac9b20d4b090c213799e81acf48a55ea8d437d6/aten/src/ATen/native/cuda/ActivationGeluKernel.cu#L25-L30 + const float f = (float)x; + constexpr float BETA = M_SQRT2 * M_2_SQRTPI * 0.5f; + constexpr float KAPPA = 0.044715; + float x_cube = f * f * f; + float inner = BETA * (f + KAPPA * x_cube); + return (T)(0.5f * f * (1.0f + ::tanhf(inner))); +} + +template +__device__ __forceinline__ packed_t +packed_gelu_tanh_kernel(const packed_t& val, const float /*alpha*/) { + // Equivalent to PyTorch GELU with 'tanh' approximation. + // Refer to: + // https://github.com/pytorch/pytorch/blob/8ac9b20d4b090c213799e81acf48a55ea8d437d6/aten/src/ATen/native/cuda/ActivationGeluKernel.cu#L25-L30 + float2 fval = cast_to_float2(val); + constexpr float BETA = M_SQRT2 * M_2_SQRTPI * 0.5f; + constexpr float KAPPA = 0.044715; + + float x_cube = fval.x * fval.x * fval.x; + float inner = BETA * (fval.x + KAPPA * x_cube); + fval.x = 0.5f * fval.x * (1.0f + ::tanhf(inner)); + + x_cube = fval.y * fval.y * fval.y; + inner = BETA * (fval.y + KAPPA * x_cube); + fval.y = 0.5f * fval.y * (1.0f + ::tanhf(inner)); + return cast_to_packed(fval); +} + +} // namespace aphrodite + +// Launch activation and gating kernel. +// Use ACT_FIRST (bool) indicating whether to apply the activation function +// first. HAS_CLAMP (bool) enables pre-activation clamping: gate input is +// clamped (max only) and up input is clamped (both sides) before the +// activation function is applied. +#define LAUNCH_ACTIVATION_GATE_KERNEL(KERNEL, PACKED_KERNEL, ACT_FIRST, \ + HAS_CLAMP, LIMIT, ALPHA, BETA) \ + auto dtype = input.scalar_type(); \ + int d = input.size(-1) / 2; \ + int64_t num_tokens = input.numel() / input.size(-1); \ + if (num_tokens == 0) { \ + return; \ + } \ + dim3 grid(num_tokens); \ + int cc_major = get_device_prop()->major; \ + int support_vec = \ + (CUDA_VERSION >= 12090 && cc_major >= 10 && num_tokens > 128) \ + ? aphrodite::VecTraits::ARCH_MAX_VEC_SIZE \ + : aphrodite::VecTraits::ARCH_MAX_VEC_SIZE; \ + int vec_size = support_vec / input.element_size(); \ + const bool use_vec = (d % vec_size == 0); \ + const torch::stable::accelerator::DeviceGuard device_guard( \ + input.get_device_index()); \ + const cudaStream_t stream = get_current_cuda_stream(); \ + if (use_vec) { \ + dim3 block(std::min(d / vec_size, 1024)); \ + if (CUDA_VERSION >= 12090 && cc_major >= 10 && num_tokens > 128) { \ + APHRODITE_STABLE_DISPATCH_FLOATING_TYPES( \ + dtype, "act_and_mul_kernel", [&] { \ + aphrodite::act_and_mul_kernel< \ + scalar_t, \ + typename aphrodite::PackedTypeConverter::Type, \ + KERNEL, \ + PACKED_KERNEL< \ + typename aphrodite::PackedTypeConverter::Type>, \ + ACT_FIRST, true, HAS_CLAMP, true><<>>( \ + out.mutable_data_ptr(), \ + input.const_data_ptr(), d, LIMIT, ALPHA, BETA); \ + }); \ + } else { \ + APHRODITE_STABLE_DISPATCH_FLOATING_TYPES( \ + dtype, "act_and_mul_kernel", [&] { \ + aphrodite::act_and_mul_kernel< \ + scalar_t, \ + typename aphrodite::PackedTypeConverter::Type, \ + KERNEL, \ + PACKED_KERNEL< \ + typename aphrodite::PackedTypeConverter::Type>, \ + ACT_FIRST, true, HAS_CLAMP, false> \ + <<>>(out.mutable_data_ptr(), \ + input.const_data_ptr(), \ + d, LIMIT, ALPHA, BETA); \ + }); \ + } \ + } else { \ + dim3 block(std::min(d, 1024)); \ + APHRODITE_STABLE_DISPATCH_FLOATING_TYPES( \ + dtype, "act_and_mul_kernel", [&] { \ + aphrodite::act_and_mul_kernel< \ + scalar_t, \ + typename aphrodite::PackedTypeConverter::Type, \ + KERNEL, \ + PACKED_KERNEL< \ + typename aphrodite::PackedTypeConverter::Type>, \ + ACT_FIRST, false, HAS_CLAMP><<>>( \ + out.mutable_data_ptr(), \ + input.const_data_ptr(), d, LIMIT, ALPHA, BETA); \ + }); \ + } + +void silu_and_mul(torch::stable::Tensor& out, // [..., d] + torch::stable::Tensor& input) // [..., 2 * d] +{ + LAUNCH_ACTIVATION_GATE_KERNEL(aphrodite::silu_kernel, + aphrodite::packed_silu_kernel, true, false, + 0.0f, 1.0f, 0.0f); +} + +void silu_and_mul_clamp(torch::stable::Tensor& out, // [..., d] + torch::stable::Tensor& input, // [..., 2 * d] + double limit, double alpha, double beta) { + // out = (gate.clamp(max=limit) * sigmoid(alpha * gate.clamp(max=limit))) + // * (up.clamp(+-limit) + beta) + // alpha=1.0, beta=0.0 reduce this to silu(gate) * up. + LAUNCH_ACTIVATION_GATE_KERNEL(aphrodite::silu_kernel, + aphrodite::packed_silu_kernel, true, true, + (float)limit, (float)alpha, (float)beta); +} + +void mul_and_silu(torch::stable::Tensor& out, // [..., d] + torch::stable::Tensor& input) // [..., 2 * d] +{ + // The difference between mul_and_silu and silu_and_mul is that mul_and_silu + // applies the silu to the latter half of the input. + LAUNCH_ACTIVATION_GATE_KERNEL(aphrodite::silu_kernel, + aphrodite::packed_silu_kernel, false, false, + 0.0f, 1.0f, 0.0f); +} + +void gelu_and_mul(torch::stable::Tensor& out, // [..., d] + torch::stable::Tensor& input) // [..., 2 * d] +{ + LAUNCH_ACTIVATION_GATE_KERNEL(aphrodite::gelu_kernel, + aphrodite::packed_gelu_kernel, true, false, + 0.0f, 1.0f, 0.0f); +} + +void gelu_tanh_and_mul(torch::stable::Tensor& out, // [..., d] + torch::stable::Tensor& input) // [..., 2 * d] +{ + LAUNCH_ACTIVATION_GATE_KERNEL(aphrodite::gelu_tanh_kernel, + aphrodite::packed_gelu_tanh_kernel, true, false, + 0.0f, 1.0f, 0.0f); +} + +namespace aphrodite { + +template +__device__ __forceinline__ T fatrelu_kernel(const T& x, const float threshold) { + const float f = (float)x; + return (T)(f > threshold ? f : 0.0f); +} + +template +__device__ __forceinline__ packed_t +packed_fatrelu_kernel(const packed_t& val, const float threshold) { + float2 fval = cast_to_float2(val); + fval.x = fval.x > threshold ? fval.x : 0.0f; + fval.y = fval.y > threshold ? fval.y : 0.0f; + return cast_to_packed(fval); +} + +template +__global__ void act_and_mul_kernel_with_param( + scalar_t* __restrict__ out, const scalar_t* __restrict__ input, const int d, + const float param) { + const scalar_t* x_ptr = input + blockIdx.x * 2 * d; + const scalar_t* y_ptr = x_ptr + d; + scalar_t* out_ptr = out + blockIdx.x * d; + + if constexpr (use_vec) { + using cuda_t = typename CUDATypeConverter::Type; + using pvec_t = PackedVec; + + const pvec_t* x_vec = reinterpret_cast(x_ptr); + const pvec_t* y_vec = reinterpret_cast(y_ptr); + pvec_t* out_vec = reinterpret_cast(out_ptr); + const int num_vecs = d / 2 / pvec_t::NUM_ELTS; + + for (int i = threadIdx.x; i < num_vecs; i += blockDim.x) { + pvec_t x, y; + if constexpr (use_256b) { + ld256(x, &x_vec[i]); + ld256(y, &y_vec[i]); + } else { + ld128(x, &x_vec[i]); + ld128(y, &y_vec[i]); + } +#pragma unroll + for (int j = 0; j < pvec_t::NUM_ELTS; j++) { + x.elts[j] = packed_mul(PACKED_ACT_FN(x.elts[j], param), y.elts[j]); + } + if constexpr (use_256b) { + st256(x, &out_vec[i]); + } else { + st128(x, &out_vec[i]); + } + } + } else { + // Scalar fallback for unaligned data or small d + for (int64_t idx = threadIdx.x; idx < d; idx += blockDim.x) { + const scalar_t x = APHRODITE_LDG(&x_ptr[idx]); + const scalar_t y = APHRODITE_LDG(&y_ptr[idx]); + out_ptr[idx] = ACT_FN(x, param) * y; + } + } +} + +template +__device__ __forceinline__ T swigluoai_and_mul(const T& gate, const T& up, + float alpha, float limit) { + // Clamp gate to (-inf, limit] and up to [-limit, limit] + const float g = fminf((float)gate, limit); + const float u = fmaxf(fminf((float)up, limit), -limit); + // glu = gate * sigmoid(gate * alpha), then return (up + 1) * glu + return (T)((u + 1.0f) * g / (1.0f + expf(-g * alpha))); +} + +// Interleaved gate/up: input has [gate0, up0, gate1, up1, ...]. +template +__global__ void swigluoai_and_mul_kernel( + scalar_t* __restrict__ out, // [..., d] + const scalar_t* __restrict__ input, // [..., 2 * d] (interleaved) + const int d, const float alpha, const float limit) { + // For interleaved data: input has 2*d elements per token (gate/up pairs) + // output has d elements per token + constexpr int VEC_SIZE = 16 / sizeof(scalar_t); + constexpr int PAIRS = VEC_SIZE / 2; // Number of gate/up pairs per int4 load + const int64_t token_idx = blockIdx.x; + const scalar_t* in_ptr = input + token_idx * 2 * d; + scalar_t* out_ptr = out + token_idx * d; + + // Check alignment for 128-bit vectorized access on input. + // For output we use int2 (64-bit) which has 8-byte alignment requirement. + const bool in_aligned = is_16byte_aligned(in_ptr); + const bool out_aligned = + (reinterpret_cast(out_ptr) & 7) == 0; // 8-byte for int2 + + if (in_aligned && out_aligned && d >= PAIRS) { + // Fast path: vectorized loop + // Each int4 load gives VEC_SIZE elements = PAIRS gate/up pairs + // Each int2 store writes PAIRS output elements + const int4* in_vec = reinterpret_cast(in_ptr); + int2* out_vec = reinterpret_cast(out_ptr); + const int num_vecs = d / PAIRS; + const int vec_end = num_vecs * PAIRS; + + for (int i = threadIdx.x; i < num_vecs; i += blockDim.x) { + int4 v = APHRODITE_LDG(&in_vec[i]); + int2 r; + auto* vp = reinterpret_cast(&v); + auto* rp = reinterpret_cast(&r); +#pragma unroll + for (int j = 0; j < PAIRS; j++) { + rp[j] = ACT_FN(vp[2 * j], vp[2 * j + 1], alpha, limit); + } + out_vec[i] = r; + } + // Scalar cleanup for remaining elements + for (int i = vec_end + threadIdx.x; i < d; i += blockDim.x) { + out_ptr[i] = ACT_FN(APHRODITE_LDG(&in_ptr[2 * i]), + APHRODITE_LDG(&in_ptr[2 * i + 1]), alpha, limit); + } + } else { + // Scalar fallback for unaligned data or small d + for (int64_t idx = threadIdx.x; idx < d; idx += blockDim.x) { + // gate = x[..., ::2] (even indices) + const scalar_t gate = APHRODITE_LDG(&in_ptr[2 * idx]); + // up = x[..., 1::2] (odd indices) + const scalar_t up = APHRODITE_LDG(&in_ptr[2 * idx + 1]); + out_ptr[idx] = ACT_FN(gate, up, alpha, limit); + } + } +} + +} // namespace aphrodite + +#define LAUNCH_ACTIVATION_GATE_KERNEL_WITH_PARAM(KERNEL, PACKED_KERNEL, PARAM) \ + auto dtype = input.scalar_type(); \ + int d = input.size(-1) / 2; \ + int64_t num_tokens = input.numel() / input.size(-1); \ + if (num_tokens == 0) { \ + return; \ + } \ + dim3 grid(num_tokens); \ + int cc_major = get_device_prop()->major; \ + int support_vec = \ + (CUDA_VERSION >= 12090 && cc_major >= 10 && num_tokens > 128) \ + ? aphrodite::VecTraits::ARCH_MAX_VEC_SIZE \ + : aphrodite::VecTraits::ARCH_MAX_VEC_SIZE; \ + int vec_size = support_vec / input.element_size(); \ + const bool use_vec = (d % vec_size == 0); \ + const torch::stable::accelerator::DeviceGuard device_guard( \ + input.get_device_index()); \ + const cudaStream_t stream = get_current_cuda_stream(); \ + if (use_vec) { \ + dim3 block(std::min(d / vec_size, 1024)); \ + if (CUDA_VERSION >= 12090 && cc_major >= 10 && num_tokens > 128) { \ + APHRODITE_STABLE_DISPATCH_FLOATING_TYPES( \ + dtype, "act_and_mul_kernel_with_param", [&] { \ + aphrodite::act_and_mul_kernel_with_param< \ + scalar_t, \ + typename aphrodite::PackedTypeConverter::Type, \ + KERNEL, \ + PACKED_KERNEL< \ + typename aphrodite::PackedTypeConverter::Type>, \ + true, true><<>>( \ + out.mutable_data_ptr(), \ + input.const_data_ptr(), d, PARAM); \ + }); \ + } else { \ + APHRODITE_STABLE_DISPATCH_FLOATING_TYPES( \ + dtype, "act_and_mul_kernel_with_param", [&] { \ + aphrodite::act_and_mul_kernel_with_param< \ + scalar_t, \ + typename aphrodite::PackedTypeConverter::Type, \ + KERNEL, \ + PACKED_KERNEL< \ + typename aphrodite::PackedTypeConverter::Type>, \ + true, false><<>>( \ + out.mutable_data_ptr(), \ + input.const_data_ptr(), d, PARAM); \ + }); \ + } \ + } else { \ + dim3 block(std::min(d, 1024)); \ + APHRODITE_STABLE_DISPATCH_FLOATING_TYPES( \ + dtype, "act_and_mul_kernel_with_param", [&] { \ + aphrodite::act_and_mul_kernel_with_param< \ + scalar_t, \ + typename aphrodite::PackedTypeConverter::Type, \ + KERNEL, \ + PACKED_KERNEL< \ + typename aphrodite::PackedTypeConverter::Type>, \ + false><<>>( \ + out.mutable_data_ptr(), \ + input.const_data_ptr(), d, PARAM); \ + }); \ + } + +#define LAUNCH_SIGLUOAI_AND_MUL(KERNEL, ALPHA, LIMIT) \ + int d = input.size(-1) / 2; \ + int64_t num_tokens = input.numel() / input.size(-1); \ + 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(), "clamp_swiglu_kernel_with_params", [&] { \ + aphrodite::swigluoai_and_mul_kernel> \ + <<>>(out.mutable_data_ptr(), \ + input.const_data_ptr(), d, \ + ALPHA, LIMIT); \ + }); + +void fatrelu_and_mul(torch::stable::Tensor& out, // [..., d], + torch::stable::Tensor& input, // [..., 2 * d] + double threshold) { + LAUNCH_ACTIVATION_GATE_KERNEL_WITH_PARAM( + aphrodite::fatrelu_kernel, aphrodite::packed_fatrelu_kernel, threshold); +} +void swigluoai_and_mul(torch::stable::Tensor& out, // [..., d] + torch::stable::Tensor& input, // [..., 2 * d] + double alpha, double limit) { + LAUNCH_SIGLUOAI_AND_MUL(aphrodite::swigluoai_and_mul, alpha, limit); +} +namespace aphrodite { + +// Element-wise activation kernel template. +template +__global__ void activation_kernel( + scalar_t* __restrict__ out, // [..., d] + const scalar_t* __restrict__ input, // [..., d] + const int d) { + const scalar_t* in_ptr = input + blockIdx.x * d; + scalar_t* out_ptr = out + blockIdx.x * d; + + if constexpr (use_vec) { + // Fast path: 128-bit/256-bit vectorized loop + using vec_t = typename VecTraits::vec_t; + constexpr int ARCH_MAX_VEC_SIZE = VecTraits::ARCH_MAX_VEC_SIZE; + constexpr int VEC_SIZE = ARCH_MAX_VEC_SIZE / sizeof(scalar_t); + const vec_t* in_vec = reinterpret_cast(in_ptr); + vec_t* out_vec = reinterpret_cast(out_ptr); + const int num_vecs = d / VEC_SIZE; + + for (int i = threadIdx.x; i < num_vecs; i += blockDim.x) { + vec_t v; + if constexpr (use_256b) { + ld256(v, &in_vec[i]); + } else { + v = APHRODITE_LDG(&in_vec[i]); + } + auto* vp = reinterpret_cast(&v); +#pragma unroll + for (int j = 0; j < VEC_SIZE; j++) { + vp[j] = ACT_FN(vp[j]); + } + if constexpr (use_256b) { + st256(v, &out_vec[i]); + } else { + out_vec[i] = v; + } + } + } else { + // Scalar fallback for unaligned data or small d + for (int64_t idx = threadIdx.x; idx < d; idx += blockDim.x) { + const scalar_t x = APHRODITE_LDG(&in_ptr[idx]); + out_ptr[idx] = ACT_FN(x); + } + } +} + +} // namespace aphrodite + +// Launch element-wise activation kernel. +#define LAUNCH_ACTIVATION_KERNEL(KERNEL) \ + auto dtype = input.scalar_type(); \ + int d = input.size(-1); \ + int64_t num_tokens = input.numel() / input.size(-1); \ + if (num_tokens == 0) { \ + return; \ + } \ + dim3 grid(num_tokens); \ + int cc_major = get_device_prop()->major; \ + int support_vec = \ + (CUDA_VERSION >= 12090 && cc_major >= 10 && num_tokens > 128) \ + ? aphrodite::VecTraits::ARCH_MAX_VEC_SIZE \ + : aphrodite::VecTraits::ARCH_MAX_VEC_SIZE; \ + int vec_size = support_vec / input.element_size(); \ + const bool use_vec = (d % vec_size == 0); \ + const torch::stable::accelerator::DeviceGuard device_guard( \ + input.get_device_index()); \ + const cudaStream_t stream = get_current_cuda_stream(); \ + if (use_vec) { \ + dim3 block(std::min(d / vec_size, 1024)); \ + if (CUDA_VERSION >= 12090 && cc_major >= 10 && num_tokens > 128) { \ + APHRODITE_STABLE_DISPATCH_FLOATING_TYPES( \ + dtype, "activation_kernel", [&] { \ + aphrodite::activation_kernel, true, \ + true><<>>( \ + out.mutable_data_ptr(), \ + input.const_data_ptr(), d); \ + }); \ + } else { \ + APHRODITE_STABLE_DISPATCH_FLOATING_TYPES( \ + dtype, "activation_kernel", [&] { \ + aphrodite::activation_kernel, true, \ + false><<>>( \ + out.mutable_data_ptr(), \ + input.const_data_ptr(), d); \ + }); \ + } \ + } else { \ + dim3 block(std::min(d, 1024)); \ + APHRODITE_STABLE_DISPATCH_FLOATING_TYPES(dtype, "activation_kernel", [&] { \ + aphrodite::activation_kernel, false> \ + <<>>(out.mutable_data_ptr(), \ + input.const_data_ptr(), d); \ + }); \ + } + +namespace aphrodite { + +template +__device__ __forceinline__ T gelu_new_kernel(const T& x) { + const float x3 = (float)(x * x * x); + const T t = (T)tanhf((T)(0.79788456f * (float)(x + (T)(0.044715f * x3)))); + return ((T)0.5) * x * (((T)1.0) + t); +} + +template +__device__ __forceinline__ T gelu_fast_kernel(const T& x) { + const float f = (float)x; + const T t = + (T)tanhf(((T)(f * 0.79788456f)) * (((T)1.0) + (T)(0.044715f * f) * x)); + return ((T)0.5) * x * (((T)1.0) + t); +} + +template +__device__ __forceinline__ T gelu_quick_kernel(const T& x) { + // x * sigmoid(1.702 * x) + return (T)(((float)x) / (1.0f + expf(-1.702f * (float)x))); +} + +template +__device__ __forceinline__ T relu_squared_kernel(const T& x) { + // relu(x)^2 - introduced in https://arxiv.org/abs/2109.08668v2 + const float f = (float)x; + const float val = f > 0.0f ? f : 0.0f; + return (T)(val * val); +} + +} // namespace aphrodite + +void gelu_new(torch::stable::Tensor& out, // [..., d] + torch::stable::Tensor& input) // [..., d] +{ + LAUNCH_ACTIVATION_KERNEL(aphrodite::gelu_new_kernel); +} + +void gelu_fast(torch::stable::Tensor& out, // [..., d] + torch::stable::Tensor& input) // [..., d] +{ + LAUNCH_ACTIVATION_KERNEL(aphrodite::gelu_fast_kernel); +} + +void gelu_quick(torch::stable::Tensor& out, // [..., d] + torch::stable::Tensor& input) // [..., d] +{ + LAUNCH_ACTIVATION_KERNEL(aphrodite::gelu_quick_kernel); +} + +void relu_squared(torch::stable::Tensor& out, // [..., d] + torch::stable::Tensor& input) // [..., d] +{ + LAUNCH_ACTIVATION_KERNEL(aphrodite::relu_squared_kernel); +} 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/cache_kernels.cu.orig b/csrc/libtorch_stable/cache_kernels.cu.orig new file mode 100644 index 0000000000..fc8e871c9d --- /dev/null +++ b/csrc/libtorch_stable/cache_kernels.cu.orig @@ -0,0 +1,1669 @@ +#include "torch_utils.h" +#include "dispatch_utils.h" + +#include "../cuda_utils.h" +#include "../cuda_compat.h" + +#include "quantization/vectorization_utils.cuh" +#include "concat_mla_q.cuh" + +#ifdef USE_ROCM + #include "../quantization/w8a8/fp8/amd/quant_utils.cuh" +#else + #include "../quantization/w8a8/fp8/nvidia/quant_utils.cuh" +#endif + +#include +#include +#include + +#ifdef USE_ROCM + #include +typedef __hip_bfloat16 __nv_bfloat16; +#else + #include +#endif + +#if defined(__gfx942__) +constexpr float kFp8ScaleDivisor = 224.f; +#else +constexpr float kFp8ScaleDivisor = 448.f; +#endif + +void swap_blocks(torch::stable::Tensor& src, torch::stable::Tensor& dst, + int64_t block_size_in_bytes, + const torch::stable::Tensor& block_mapping) { + torch::stable::Device src_device = src.device(); + torch::stable::Device dst_device = dst.device(); + cudaMemcpyKind memcpy_type; + if (src_device.is_cuda() && dst_device.is_cuda()) { + STD_TORCH_CHECK(src_device.index() == dst_device.index(), + "src and dst must be on the same GPU"); + memcpy_type = cudaMemcpyDeviceToDevice; + } else if (src_device.is_cuda() && dst_device.is_cpu()) { + memcpy_type = cudaMemcpyDeviceToHost; + } else if (src_device.is_cpu() && dst_device.is_cuda()) { + memcpy_type = cudaMemcpyHostToDevice; + } else { + STD_TORCH_CHECK(false, "Invalid device combination"); + } + + // NOTE(youkaichao): keep in mind that `block_mapping` should be + // a cpu tensor, otherwise every `item` call will require a gpu-cpu + // synchronization. + STD_TORCH_CHECK(block_mapping.device().is_cpu(), + "block_mapping must be on CPU"); + + char* src_ptr = static_cast(src.data_ptr()); + char* dst_ptr = static_cast(dst.data_ptr()); + + auto guard_device = src_device.is_cuda() ? src_device : dst_device; + const torch::stable::accelerator::DeviceGuard device_guard( + guard_device.index()); + const cudaStream_t stream = get_current_cuda_stream(); + // NOTE(woosuk): This can be slow if the number of blocks is large. + const int64_t num_blocks = block_mapping.size(0); + const int64_t* bm_ptr = block_mapping.const_data_ptr(); + const int64_t bm_stride0 = block_mapping.stride(0); + const int64_t bm_stride1 = block_mapping.stride(1); + for (size_t i = 0; i < num_blocks; i++) { + int64_t src_block_number = bm_ptr[i * bm_stride0]; + int64_t dst_block_number = bm_ptr[i * bm_stride0 + bm_stride1]; + int64_t src_offset = src_block_number * block_size_in_bytes; + int64_t dst_offset = dst_block_number * block_size_in_bytes; + cudaMemcpyAsync(dst_ptr + dst_offset, src_ptr + src_offset, + block_size_in_bytes, memcpy_type, stream); + } +} + +void swap_blocks_batch(const torch::stable::Tensor& src_ptrs, + const torch::stable::Tensor& dst_ptrs, + const torch::stable::Tensor& sizes, + bool is_src_access_order_any) { + STD_TORCH_CHECK(src_ptrs.device().is_cpu(), "src_ptrs must be on CPU"); + STD_TORCH_CHECK(dst_ptrs.device().is_cpu(), "dst_ptrs must be on CPU"); + STD_TORCH_CHECK(sizes.device().is_cpu(), "sizes must be on CPU"); + STD_TORCH_CHECK(src_ptrs.scalar_type() == torch::headeronly::ScalarType::Long, + "src_ptrs must be int64"); + STD_TORCH_CHECK(dst_ptrs.scalar_type() == torch::headeronly::ScalarType::Long, + "dst_ptrs must be int64"); + STD_TORCH_CHECK(sizes.scalar_type() == torch::headeronly::ScalarType::Long, + "sizes must be int64"); + + const int64_t n = src_ptrs.size(0); + STD_TORCH_CHECK(dst_ptrs.size(0) == n, "dst_ptrs length must match src_ptrs"); + STD_TORCH_CHECK(sizes.size(0) == n, "sizes length must match src_ptrs"); + + if (n == 0) return; + + int64_t* src_data = src_ptrs.mutable_data_ptr(); + int64_t* dst_data = dst_ptrs.mutable_data_ptr(); + int64_t* size_data = sizes.mutable_data_ptr(); + + const cudaStream_t stream = get_current_cuda_stream(); + + // Use cuMemcpyBatchAsync / hipMemcpyBatchAsync to submit all copies in a + // single driver call, amortizing per-copy submission overhead. int64_t + // and CUdeviceptr/void*/size_t are all 8 bytes on 64-bit platforms, so we + // reinterpret_cast the tensor data directly to avoid copies. + static_assert(sizeof(size_t) == sizeof(int64_t)); +#if !defined(USE_ROCM) && defined(CUDA_VERSION) && CUDA_VERSION >= 12080 + static_assert(sizeof(CUdeviceptr) == sizeof(int64_t)); + // Resolve cuMemcpyBatchAsync at runtime via cuGetProcAddress so that + // binaries compiled with CUDA 12.8+ still work on older drivers, and + // we avoid the CUDA 13.0 header remapping (#define to _v2 signature). + // The function pointer is cached after the first call. + using BatchFn = + CUresult (*)(CUdeviceptr*, CUdeviceptr*, size_t*, size_t, + CUmemcpyAttributes*, size_t*, size_t, size_t*, CUstream); + static BatchFn batch_fn = []() -> BatchFn { + CUdriverProcAddressQueryResult sym_status; + void* fn_ptr = nullptr; + CUresult res = cuGetProcAddress("cuMemcpyBatchAsync", &fn_ptr, 12080, + CU_GET_PROC_ADDRESS_DEFAULT, &sym_status); + if (res != CUDA_SUCCESS || fn_ptr == nullptr) { + return nullptr; + } + return reinterpret_cast(fn_ptr); + }(); + + // cuMemcpyBatchAsync rejects the legacy default stream (handle 0 / + // cudaStreamLegacy) with CUDA_ERROR_INVALID_VALUE; route it to the per-copy + // fallback below, which is correct on any stream. Real and per-thread-default + // streams take the batch fast path. + const bool usable_stream = stream != nullptr && stream != cudaStreamLegacy; + if (batch_fn != nullptr && usable_stream) { + CUmemcpyAttributes attr = {}; + // ANY lets the DMA engine prefetch source bytes out of stream order, + // which is only safe when no GPU stream is concurrently writing the + // source. + attr.srcAccessOrder = is_src_access_order_any + ? CU_MEMCPY_SRC_ACCESS_ORDER_ANY + : CU_MEMCPY_SRC_ACCESS_ORDER_STREAM; + size_t attrs_idx = 0; + size_t fail_idx = 0; + CUresult result = batch_fn(reinterpret_cast(dst_data), + reinterpret_cast(src_data), + reinterpret_cast(size_data), + static_cast(n), &attr, &attrs_idx, 1, + &fail_idx, static_cast(stream)); + STD_TORCH_CHECK(result == CUDA_SUCCESS, + "cuMemcpyBatchAsync failed at index ", fail_idx, + " with error ", result); + return; + } +#elif defined(USE_ROCM) && defined(HIP_VERSION) && HIP_VERSION >= 70100000 + // ROCm 7.1+ exposes hipMemcpyBatchAsync. The 7.2.1 implementation early- + // returns hipErrorNotSupported whenever numAttrs > 0 (see ROCm/clr @ + // rocm-7.2.1 hipamd/src/hip_memory.cpp:2819-2822), so call with + // numAttrs=0. + { + hipMemcpyAttributes attr = {}; + size_t attrs_idx = 0; + size_t fail_idx = 0; + hipError_t result = hipMemcpyBatchAsync( + reinterpret_cast(dst_data), reinterpret_cast(src_data), + reinterpret_cast(size_data), static_cast(n), &attr, + &attrs_idx, 0, &fail_idx, static_cast(stream)); + STD_TORCH_CHECK(result == hipSuccess, + "hipMemcpyBatchAsync failed at index ", fail_idx, + " with error ", result); + return; + } +#endif + { + // Fallback for CUDA < 12.8, older CUDA drivers, and ROCm < 7.1: + // individual async copies. cudaMemcpyDefault lets the driver infer + // direction from pointer types. + for (int64_t i = 0; i < n; i++) { + cudaMemcpyAsync(reinterpret_cast(dst_data[i]), + reinterpret_cast(src_data[i]), + static_cast(size_data[i]), cudaMemcpyDefault, + stream); + } + } +} + +namespace aphrodite { + +// Grid: (num_layers, num_pairs) +template +__global__ void copy_blocks_kernel(int64_t* key_cache_ptrs, + int64_t* value_cache_ptrs, + const int64_t* __restrict__ block_mapping, + const int numel_per_block) { + const int layer_idx = blockIdx.x; + const int pair_idx = blockIdx.y; + + scalar_t* key_cache = reinterpret_cast(key_cache_ptrs[layer_idx]); + scalar_t* value_cache = + reinterpret_cast(value_cache_ptrs[layer_idx]); + int64_t src_block_number = block_mapping[2 * pair_idx]; + int64_t dst_block_number = block_mapping[2 * pair_idx + 1]; + + const int64_t src_block_offset = src_block_number * numel_per_block; + const int64_t dst_block_offset = dst_block_number * numel_per_block; + for (int i = threadIdx.x; i < numel_per_block; i += blockDim.x) { + int64_t src_offset = src_block_offset + i; + int64_t dst_offset = dst_block_offset + i; + key_cache[dst_offset] = key_cache[src_offset]; + } + for (int i = threadIdx.x; i < numel_per_block; i += blockDim.x) { + int64_t src_offset = src_block_offset + i; + int64_t dst_offset = dst_block_offset + i; + value_cache[dst_offset] = value_cache[src_offset]; + } +} + +// Kernel for MLA, which works on a single joint kv_cache +// Grid: (num_layers, num_pairs) +template +__global__ void copy_blocks_mla_kernel( + int64_t* cache_ptrs, const int64_t* __restrict__ block_mapping, + const int mem_footprint_per_block) { + const int layer_idx = blockIdx.x; + const int pair_idx = blockIdx.y; + scalar_t* cache = reinterpret_cast(cache_ptrs[layer_idx]); + int64_t src_block = block_mapping[2 * pair_idx]; + int64_t dst_block = block_mapping[2 * pair_idx + 1]; + int64_t src_offset = src_block * mem_footprint_per_block; + int64_t dst_offset = dst_block * mem_footprint_per_block; + for (int i = threadIdx.x; i < mem_footprint_per_block; i += blockDim.x) { + cache[dst_offset + i] = cache[src_offset + i]; + } +} + +} // namespace aphrodite + +namespace aphrodite { + +// Used to copy/convert one element +template +struct CopyWithScaleOp { + float scale; + + __device__ __forceinline__ void operator()(OutT& dst, const InT src) const { + if constexpr (kv_dt == Fp8KVCacheDataType::kAuto) { + dst = static_cast(src); + } else { + dst = fp8::scaled_convert(src, scale); + } + } +}; + +template +__global__ void reshape_and_cache_kernel( + const scalar_t* __restrict__ key, // [num_tokens, num_heads, head_size] + const scalar_t* __restrict__ value, // [num_tokens, num_heads, head_size] + cache_t* __restrict__ key_cache, // [num_blocks, num_heads, head_size/x, + // block_size, x] + cache_t* __restrict__ value_cache, // [num_blocks, num_heads, head_size, + // block_size] + const int64_t* __restrict__ slot_mapping, // [num_tokens] + const int key_stride, const int value_stride, const int num_heads, + const int head_size, const int block_size, const int x, + const float* k_scale, const float* v_scale) { + const int64_t token_idx = blockIdx.x; + const int64_t slot_idx = slot_mapping[token_idx]; + if (slot_idx < 0) { + return; + } + + const int64_t block_idx = slot_idx / block_size; + const int64_t block_offset = slot_idx % block_size; + const int h_block_count = head_size / x; // head_size//x + + const int h_block_idx = threadIdx.x; + if (h_block_idx >= num_heads * h_block_count) { + return; + } + + const int head_idx = h_block_idx / h_block_count; + const int h_block = h_block_idx % h_block_count; + + const scalar_t* __restrict__ key_src = + key + token_idx * key_stride + head_idx * head_size + h_block * x; + const int64_t src_value_start = + token_idx * value_stride + head_idx * head_size + h_block * x; + + cache_t* __restrict__ key_dst = + key_cache + block_idx * num_heads * h_block_count * block_size * x + + head_idx * h_block_count * block_size * x + h_block * block_size * x + + block_offset * x; + const int64_t tgt_value_start = + block_idx * num_heads * h_block_count * x * block_size + + head_idx * h_block_count * x * block_size + h_block * x * block_size + + block_offset; + + constexpr int VEC_SIZE = (sizeof(scalar_t) == 2) ? 8 : 4; + float k_scale_val = (kv_dt == Fp8KVCacheDataType::kAuto) ? 0.f : *k_scale; + CopyWithScaleOp k_op{k_scale_val}; + float v_scale_val = (kv_dt == Fp8KVCacheDataType::kAuto) ? 0.f : *v_scale; + CopyWithScaleOp v_op{v_scale_val}; + + vectorize_with_alignment(key_src, key_dst, x, 0, 1, k_op); + + const scalar_t* __restrict__ value_src = value + src_value_start; + cache_t* __restrict__ value_dst = value_cache + tgt_value_start; +#pragma unroll + for (int i = 0; i < x; i++) { + v_op(value_dst[i * block_size], value_src[i]); + } +} + +template +__global__ void reshape_and_cache_flash_kernel( + const scalar_t* __restrict__ key, // [num_tokens, num_heads, head_size] + const scalar_t* __restrict__ value, // [num_tokens, num_heads, head_size] + cache_t* __restrict__ key_cache, // NHD or HND, shape see comments below + cache_t* __restrict__ value_cache, // same above + const int64_t* __restrict__ slot_mapping, // [num_tokens] + const int64_t block_stride, const int64_t page_stride, + const int64_t head_stride, const int64_t key_stride, + const int64_t value_stride, const int num_heads, const int head_size, + const int block_size, const float* k_scale, const float* v_scale, + const int kv_scale_stride) { + const int64_t token_idx = blockIdx.x; + const int64_t slot_idx = slot_mapping[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; + const int n_elems = num_heads * head_size; + + // pointers to the beginning of the source row for this token. + const scalar_t* __restrict__ key_src = key + token_idx * key_stride; + const scalar_t* __restrict__ value_src = value + token_idx * value_stride; + + // find the start position inside the kv-cache for this token. + cache_t* __restrict__ key_dst = + key_cache + block_idx * block_stride + block_offset * page_stride; + cache_t* __restrict__ value_dst = + value_cache + block_idx * block_stride + block_offset * page_stride; + + // this is true for the NHD layout where `head_stride == head_size` + const bool is_contiguous_heads = (head_stride == head_size); + + constexpr int VEC_SIZE = (sizeof(scalar_t) == 2) ? 8 : 4; + + if (is_contiguous_heads && kv_scale_stride == 0) { + // NHD layout and k/v_scales are [1] (i.e. single scale for all heads) + // kv cache: [num_blocks, block_size, num_heads, head_size] + float k_scale_val = (kv_dt == Fp8KVCacheDataType::kAuto) ? 0.f : *k_scale; + float v_scale_val = (kv_dt == Fp8KVCacheDataType::kAuto) ? 0.f : *v_scale; + + CopyWithScaleOp k_op{k_scale_val}; + CopyWithScaleOp v_op{v_scale_val}; + + vectorize_with_alignment(key_src, key_dst, n_elems, threadIdx.x, + blockDim.x, k_op); + vectorize_with_alignment(value_src, value_dst, n_elems, + threadIdx.x, blockDim.x, v_op); + } else { + // HND layout OR k/v_scales are [num_heads] (i.e. per-attn-head) + // HND layout: heads are strided, but each head_size segment is contiguous + // kv cache: [num_blocks, num_heads, block_size, head_size] + const int lane = threadIdx.x & 31; // 0..31 within warp + const int warp_id = threadIdx.x >> 5; // warp index within block + const int warps_per_block = blockDim.x >> 5; + + for (int head = warp_id; head < num_heads; head += warps_per_block) { + const scalar_t* __restrict__ k_src_h = key_src + head * head_size; + const scalar_t* __restrict__ v_src_h = value_src + head * head_size; + + cache_t* __restrict__ k_dst_h = + key_dst + static_cast(head) * head_stride; + cache_t* __restrict__ v_dst_h = + value_dst + static_cast(head) * head_stride; + + float k_scale_val = (kv_dt == Fp8KVCacheDataType::kAuto) + ? 0.f + : k_scale[head * kv_scale_stride]; + float v_scale_val = (kv_dt == Fp8KVCacheDataType::kAuto) + ? 0.f + : v_scale[head * kv_scale_stride]; + + CopyWithScaleOp k_op{k_scale_val}; + CopyWithScaleOp v_op{v_scale_val}; + + // within each head, let the 32 threads of the warp perform the vector + // copy + vectorize_with_alignment(k_src_h, k_dst_h, head_size, lane, 32, + k_op); + + vectorize_with_alignment(v_src_h, v_dst_h, head_size, lane, 32, + v_op); + } + } +} + +template +__global__ void concat_and_cache_mla_kernel( + const scalar_t* __restrict__ kv_c, // [num_tokens, kv_lora_rank] + const scalar_t* __restrict__ k_pe, // [num_tokens, pe_dim] + cache_t* __restrict__ kv_cache, // [num_blocks, block_size, (kv_lora_rank + // + pe_dim)] + const int64_t* __restrict__ slot_mapping, // [num_tokens] + const int block_stride, // + const int entry_stride, // + const int kv_c_stride, // + const int k_pe_stride, // + const int kv_lora_rank, // + const int pe_dim, // + const int block_size, // + const float* scale // +) { + const int64_t token_idx = blockIdx.x; + const int64_t slot_idx = slot_mapping[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; + + auto copy = [&](const scalar_t* __restrict__ src, cache_t* __restrict__ dst, + int src_stride, int dst_stride, int size, int offset) { + for (int i = threadIdx.x; i < size; i += blockDim.x) { + const int64_t src_idx = token_idx * src_stride + i; + const int64_t dst_idx = + block_idx * block_stride + block_offset * entry_stride + i + offset; + if constexpr (kv_dt == Fp8KVCacheDataType::kAuto) { + dst[dst_idx] = src[src_idx]; + } else { + dst[dst_idx] = + fp8::scaled_convert(src[src_idx], *scale); + } + } + }; + + copy(kv_c, kv_cache, kv_c_stride, block_stride, kv_lora_rank, 0); + 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] + const scalar_t* __restrict__ k_pe, // [num_tokens, pe_dim] + cache_t* __restrict__ kv_cache, // [num_blocks, block_size, (kv_lora_rank + // + pe_dim)] + const int64_t* __restrict__ slot_mapping, // [num_tokens] + const int block_stride, // + const int entry_stride, // + const int kv_c_stride, // + const int k_pe_stride, // + const int kv_lora_rank, // + const int pe_dim, // + const int block_size, // + const float* scale // +) { + const int64_t token_idx = blockIdx.x; + const int64_t slot_idx = slot_mapping[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; + const int64_t dst_idx_start = + block_idx * block_stride + block_offset * entry_stride; + + // For the NoPE part, each tile of 128 elements is handled by half of one warp + // (16 threads). There are 4 total tiles, so 2 warps (64 threads). + // Lanes 0 and 16 of each warp write the scale values for that warp's tiles. + // The RoPE part (last 64 elements) is handled by another 1 warp (32 threads). + // So in total, we use 3 warps (96 threads) per block. + + // Cast kv_cache to 16_bit for RoPE values + scalar_t* kv_cache_16bit = + reinterpret_cast(&kv_cache[dst_idx_start]); + + // The last warp handles the RoPE part + if (threadIdx.x >= 64) { + // Each thread handles two elements of RoPE + const int8_t pe_idx_start = (threadIdx.x - 64) * 2; + const int64_t src_idx = token_idx * k_pe_stride + pe_idx_start; + // Vectorized load of two 16-bit values, performed as one 32-bit load + const int32_t vals = *reinterpret_cast(&k_pe[src_idx]); + // RoPE values start after the packed 8-bit NoPE values and the + // 32-bit scales + const int64_t dst_idx = kv_lora_rank / 2 + 8 + pe_idx_start; + // Vectorized store of two 16-bit values, performed as one 32-bit store + *reinterpret_cast(&kv_cache_16bit[dst_idx]) = vals; + return; + } + + // The first two warps handle the NoPE part + const int8_t warp_idx = threadIdx.x >> 5; + const int8_t lane_idx = threadIdx.x & 31; + const int8_t tile_idx = warp_idx * 2 + (lane_idx >> 4); + + // Each thread handles 8 elements of NoPE + // Load the NoPE elements for this thread into registers + const int64_t src_idx_start = token_idx * kv_c_stride + (threadIdx.x * 8); + // Vectorized load of eight 16-bit values, performed as an int4 load + const int4 vals_i4 = *reinterpret_cast(&kv_c[src_idx_start]); + const scalar_t* vals = reinterpret_cast(&vals_i4); + + // Max absolute value of this thread's elements + float max_abs = fmaxf(fmaxf(fmaxf(fabsf(vals[0]), fabsf(vals[1])), + fmaxf(fabsf(vals[2]), fabsf(vals[3]))), + fmaxf(fmaxf(fabsf(vals[4]), fabsf(vals[5])), + fmaxf(fabsf(vals[6]), fabsf(vals[7])))); + + // Warp-level reduction to find the max absolute value in each half-warp +#pragma unroll + for (int offset = 8; offset > 0; offset /= 2) { + max_abs = + fmaxf(max_abs, APHRODITE_SHFL_XOR_SYNC_WIDTH(max_abs, offset, 16)); + } + + // Compute the scale for the tile + float tile_scale = fmaxf(max_abs / kFp8ScaleDivisor, FLT_MIN); + + // The first lane of each half-warp writes the scale to kv_cache + if ((lane_idx == 0) || (lane_idx == 16)) { + float* kv_cache_32bit = reinterpret_cast(&kv_cache[dst_idx_start]); + const uint64_t dst_idx = kv_lora_rank / 4 + tile_idx; + kv_cache_32bit[dst_idx] = tile_scale; + } + + // Now all threads in the block scale and write their elements + // NoPE data is packed in the first kv_lora_rank/2 bytes (first 256 bytes) + const int64_t dst_idx_base = dst_idx_start + (threadIdx.x * 8); + + uint8_t result[8]; +#pragma unroll + for (int i = 0; i < 8; i++) { + result[i] = + fp8::scaled_convert( + vals[i], tile_scale); + } + + // Store as aligned 64-bit writes + *reinterpret_cast(&kv_cache[dst_idx_base]) = + *reinterpret_cast(result); +} + +template +__global__ void indexer_k_quant_and_cache_kernel( + const scalar_t* __restrict__ k, // [num_tokens, head_dim] + cache_t* __restrict__ kv_cache, // [num_blocks, block_size, cache_stride] + const int64_t* __restrict__ slot_mapping, // [num_tokens] + const int head_dim, // dimension of each head + const int quant_block_size, // quantization block size + const int cache_block_size, // cache block size + const int64_t cache_block_stride, // stride for each block in kv_cache + + const bool use_ue8m0 // use ue8m0 scale format +) { + constexpr int VEC_SIZE = 4; + const int64_t token_idx = blockIdx.x; + const int64_t head_dim_idx = (blockIdx.y * blockDim.y * blockDim.x + + threadIdx.y * blockDim.x + threadIdx.x) * + VEC_SIZE; + const int64_t slot_idx = slot_mapping[token_idx]; + const int64_t block_idx = slot_idx / cache_block_size; + const int64_t block_offset = slot_idx % cache_block_size; + + // NOTE: slot_idx can be -1 if the token is padded + if (slot_idx < 0 || (head_dim_idx >= head_dim)) { + return; + } + + float2 k_val = (reinterpret_cast( + k))[(token_idx * head_dim + head_dim_idx) / VEC_SIZE]; + scalar_t* k_val_ptr = reinterpret_cast(&k_val); + float amax = 0.0f; + for (int i = 0; i < VEC_SIZE; i++) { + amax = fmaxf(amax, fabsf(float(k_val_ptr[i]))); + } + + // Reduced amax + for (int mask = 16; mask > 0; mask /= 2) { +#ifdef USE_ROCM + amax = fmaxf(amax, __shfl_xor_sync(uint64_t(-1), amax, mask)); +#else + amax = fmaxf(amax, __shfl_xor_sync(unsigned(-1), amax, mask)); +#endif + } + + float scale = fmaxf(amax, 1e-4) / kFp8ScaleDivisor; + + if (use_ue8m0) { + scale = exp2f(ceilf(log2f(scale))); + } + + const int64_t dst_offset = + block_idx * cache_block_stride + block_offset * head_dim + head_dim_idx; + for (int i = 0; i < VEC_SIZE; i++) { + kv_cache[dst_offset + i] = + fp8::scaled_convert(k_val_ptr[i], scale); + } + if (threadIdx.x == 0) { + const int64_t dst_scale_idx = + block_idx * cache_block_stride + cache_block_size * head_dim + + (block_offset * head_dim + head_dim_idx) * 4 / quant_block_size; + reinterpret_cast(kv_cache)[dst_scale_idx / 4] = scale; + } +} + +template +__global__ void cp_gather_indexer_k_quant_cache_kernel( + const char* __restrict__ kv_cache, // [num_blocks, block_size, + // cache_stride] + char* __restrict__ dst_k, // [num_tokens, head_dim] + char* __restrict__ dst_scale, // [num_tokens, head_dim / quant_block_size * + // 4] + const int* __restrict__ block_table, // [batch_size, num_blocks] + const int* __restrict__ cu_seq_lens, // [batch_size + 1] + const int batch_size, // batch size + const int64_t token_stride, // stride for each token in dst_k + const int64_t head_dim, // dimension of each head + const int64_t block_stride, // stride for each block in kv_cache + const int64_t cache_token_stride, // stride for each token in kv_cache + const int64_t cache_block_size, // num_tokens for each block in kv_cache + const int num_blocks, // number of blocks + const int num_tokens, // number of tokens + const int quant_block_size // quantization block size +) { + constexpr int VEC_SIZE = sizeof(float4) / sizeof(char); + const int token_idx = blockIdx.x * blockDim.y + threadIdx.y; + const int head_idx = (blockIdx.y * blockDim.x + threadIdx.x) * VEC_SIZE; + // Find batch index within a block + __shared__ int batch_idx[BLOCK_Y_SIZE]; + if (threadIdx.x == 0) { + batch_idx[threadIdx.y] = -1; + } + __syncthreads(); + + for (int iter = 0; iter < cuda_utils::ceil_div(batch_size, int(blockDim.x)); + iter++) { + int tid = iter * blockDim.x + threadIdx.x; + if (tid < batch_size) { + const int seq_start = cu_seq_lens[tid]; + const int seq_end = cu_seq_lens[tid + 1]; + if (token_idx >= seq_start && token_idx < seq_end) { + batch_idx[threadIdx.y] = tid; + } + } + } + + __syncthreads(); + + // num_tokens may be an allocation upper bound when Python avoids a D2H sync. + // Only tokens covered by the exact device-side cu_seq_lens are valid to + // gather. + const int batch = batch_idx[threadIdx.y]; + if (head_idx >= head_dim || token_idx >= num_tokens || batch < 0) { + return; + } + const int inbatch_seq_idx = token_idx - cu_seq_lens[batch]; + const int block_idx = + block_table[batch * num_blocks + inbatch_seq_idx / cache_block_size]; + const int64_t src_block_offset = block_idx * block_stride; + const int64_t cache_inblock_offset = + (inbatch_seq_idx % cache_block_size) * head_dim + head_idx; + const int64_t src_inblock_offset = src_block_offset + cache_inblock_offset; + const int64_t dst_inblock_offset = token_idx * token_stride + head_idx; + + reinterpret_cast(dst_k)[dst_inblock_offset / VEC_SIZE] = + reinterpret_cast(kv_cache)[src_inblock_offset / VEC_SIZE]; + ; + if (threadIdx.x == 0) { + const int64_t src_scale_offset = + src_block_offset + cache_block_size * head_dim + + cache_inblock_offset * 4 / quant_block_size; + reinterpret_cast(dst_scale)[dst_inblock_offset / quant_block_size] = + reinterpret_cast(kv_cache)[src_scale_offset / 4]; + } +} + +} // namespace aphrodite + +// KV_T is the data type of key and value tensors. +// CACHE_T is the stored data type of kv-cache. +// KV_DTYPE is the real data type of kv-cache. +#define CALL_RESHAPE_AND_CACHE(KV_T, CACHE_T, KV_DTYPE) \ + aphrodite::reshape_and_cache_kernel \ + <<>>( \ + reinterpret_cast(key.data_ptr()), \ + reinterpret_cast(value.data_ptr()), \ + reinterpret_cast(key_cache.data_ptr()), \ + reinterpret_cast(value_cache.data_ptr()), \ + slot_mapping.const_data_ptr(), key_stride, value_stride, \ + num_heads, head_size, block_size, x, \ + reinterpret_cast(k_scale.data_ptr()), \ + reinterpret_cast(v_scale.data_ptr())); + +void reshape_and_cache( + torch::stable::Tensor& key, // [num_tokens, num_heads, head_size] + torch::stable::Tensor& value, // [num_tokens, num_heads, head_size] + torch::stable::Tensor& + key_cache, // [num_blocks, num_heads, head_size/x, block_size, x] + torch::stable::Tensor& + value_cache, // [num_blocks, num_heads, head_size, block_size] + torch::stable::Tensor& slot_mapping, // [num_tokens] + const std::string& kv_cache_dtype, torch::stable::Tensor& k_scale, + torch::stable::Tensor& v_scale) { + int num_tokens = slot_mapping.size(0); + int num_heads = key.size(1); + int head_size = key.size(2); + int block_size = key_cache.size(3); + int x = key_cache.size(4); + + int key_stride = key.stride(0); + int value_stride = value.stride(0); + int head_div_x = head_size / x; + + dim3 grid(num_tokens); + dim3 block(std::min(num_heads * head_div_x, 512)); + const torch::stable::accelerator::DeviceGuard device_guard( + key.get_device_index()); + const cudaStream_t stream = get_current_cuda_stream(); + + DISPATCH_BY_KV_CACHE_DTYPE(key.scalar_type(), kv_cache_dtype, + CALL_RESHAPE_AND_CACHE); +} + +// KV_T is the data type of key and value tensors. +// CACHE_T is the stored data type of kv-cache. +// KV_DTYPE is the real data type of kv-cache. +#define CALL_RESHAPE_AND_CACHE_FLASH(KV_T, CACHE_T, KV_DTYPE) \ + aphrodite::reshape_and_cache_flash_kernel \ + <<>>( \ + reinterpret_cast(key.data_ptr()), \ + reinterpret_cast(value.data_ptr()), \ + reinterpret_cast(key_cache.data_ptr()), \ + reinterpret_cast(value_cache.data_ptr()), \ + slot_mapping.const_data_ptr(), block_stride, page_stride, \ + head_stride, key_stride, value_stride, num_heads, head_size, \ + block_size, reinterpret_cast(k_scale.data_ptr()), \ + reinterpret_cast(v_scale.data_ptr()), \ + kv_scale_stride); + +void reshape_and_cache_flash( + torch::stable::Tensor& key, // [num_tokens, num_heads, head_size] + torch::stable::Tensor& value, // [num_tokens, num_heads, head_size] + torch::stable::Tensor& + key_cache, // [num_blocks, block_size, num_heads, head_size] + torch::stable::Tensor& + value_cache, // [num_blocks, block_size, num_heads, head_size] + torch::stable::Tensor& slot_mapping, // [num_tokens] or [num_actual_tokens] + const std::string& kv_cache_dtype, + torch::stable::Tensor& k_scale, // [1] or [num_heads] + torch::stable::Tensor& v_scale) { // [1] or [num_heads] + // NOTE(woosuk): In vLLM V1, key.size(0) can be different from + // slot_mapping.size(0) because of padding for CUDA graphs. + // In vLLM V0, key.size(0) is always equal to slot_mapping.size(0) because + // both include padding. + // In vLLM V1, however, key.size(0) can be larger than slot_mapping.size(0) + // since key includes padding for CUDA graphs, while slot_mapping does not. + // In this case, slot_mapping.size(0) represents the actual number of tokens + // before padding. + // For compatibility with both cases, we use slot_mapping.size(0) as the + // number of tokens. + int num_tokens = slot_mapping.size(0); + int num_heads = key.size(1); + int head_size = key.size(2); + + const torch::stable::accelerator::DeviceGuard device_guard( + key.get_device_index()); + const cudaStream_t stream = get_current_cuda_stream(); + + if (kv_cache_dtype == "nvfp4") { +#if defined(ENABLE_NVFP4_SM100) || defined(ENABLE_NVFP4_SM120) + // NVFP4 dispatch is compiled separately for SM100+. + extern void reshape_and_cache_nvfp4_dispatch( + torch::stable::Tensor & key, torch::stable::Tensor & value, + torch::stable::Tensor & key_cache, torch::stable::Tensor & value_cache, + torch::stable::Tensor & slot_mapping, torch::stable::Tensor & k_scale, + torch::stable::Tensor & v_scale); + reshape_and_cache_nvfp4_dispatch(key, value, key_cache, value_cache, + slot_mapping, k_scale, v_scale); + return; +#else + STD_TORCH_CHECK( + false, + "NVFP4 KV cache requires SM100+ (Blackwell). " + "Please rebuild aphrodite with a Blackwell-compatible CUDA target."); +#endif + } + + // Original FP8/auto path. + int block_size = key_cache.size(1); + + int64_t key_stride = key.stride(0); + int64_t value_stride = value.stride(0); + int64_t block_stride = key_cache.stride(0); + int64_t page_stride = key_cache.stride(1); + int64_t head_stride = key_cache.stride(2); + STD_TORCH_CHECK(key_cache.stride(0) == value_cache.stride(0)); + + STD_TORCH_CHECK(k_scale.sizes().equals(v_scale.sizes()), + "k_scale and v_scale must have the same shape"); + STD_TORCH_CHECK(k_scale.numel() == 1 || k_scale.numel() == num_heads, + "k_scale and v_scale must be of shape [1] or [num_heads]"); + int kv_scale_stride = (k_scale.numel() > 1) ? 1 : 0; + + dim3 grid(num_tokens); + dim3 block(std::min(num_heads * head_size, 512)); + + DISPATCH_BY_KV_CACHE_DTYPE(key.scalar_type(), kv_cache_dtype, + CALL_RESHAPE_AND_CACHE_FLASH); +} + +// KV_T is the data type of key and value tensors. +// CACHE_T is the stored data type of kv-cache. +// KV_DTYPE is the real data type of kv-cache. +#define CALL_CONCAT_AND_CACHE_MLA(KV_T, CACHE_T, KV_DTYPE) \ + aphrodite::concat_and_cache_mla_kernel \ + <<>>( \ + reinterpret_cast(kv_c.data_ptr()), \ + reinterpret_cast(k_pe.data_ptr()), \ + reinterpret_cast(kv_cache.data_ptr()), \ + slot_mapping.const_data_ptr(), block_stride, entry_stride, \ + kv_c_stride, k_pe_stride, kv_lora_rank, pe_dim, block_size, \ + reinterpret_cast(scale.data_ptr())); + +// KV_T is the data type of key and value tensors. +// CACHE_T is the stored data type of kv-cache. +#define CALL_CONCAT_AND_CACHE_DS_MLA(KV_T, CACHE_T, KV_DTYPE) \ + aphrodite::concat_and_cache_ds_mla_kernel \ + <<>>( \ + reinterpret_cast(kv_c.data_ptr()), \ + reinterpret_cast(k_pe.data_ptr()), \ + reinterpret_cast(kv_cache.data_ptr()), \ + slot_mapping.const_data_ptr(), block_stride, entry_stride, \ + kv_c_stride, k_pe_stride, kv_lora_rank, pe_dim, block_size, \ + reinterpret_cast(scale.data_ptr())); + +void concat_and_cache_mla( + torch::stable::Tensor& kv_c, // [num_tokens, kv_lora_rank] + torch::stable::Tensor& k_pe, // [num_tokens, pe_dim] + torch::stable::Tensor& kv_cache, // [num_blocks, block_size, (kv_lora_rank + // + pe_dim)] + torch::stable::Tensor& slot_mapping, // [num_tokens] or [num_actual_tokens] + const std::string& kv_cache_dtype, torch::stable::Tensor& scale) { + // NOTE(woosuk): In vLLM V1, key.size(0) can be different from + // slot_mapping.size(0) because of padding for CUDA graphs. + // In vLLM V0, key.size(0) is always equal to slot_mapping.size(0) because + // both include padding. + // In vLLM V1, however, key.size(0) can be larger than slot_mapping.size(0) + // since key includes padding for CUDA graphs, while slot_mapping does not. + // In this case, slot_mapping.size(0) represents the actual number of tokens + // before padding. + // For compatibility with both cases, we use slot_mapping.size(0) as the + // number of tokens. + int num_tokens = slot_mapping.size(0); + int kv_lora_rank = kv_c.size(1); + int pe_dim = k_pe.size(1); + int block_size = kv_cache.size(1); + + if (kv_cache_dtype == "fp8_ds_mla") { + STD_TORCH_CHECK(kv_lora_rank == 512, + "kv_lora_rank must be 512 for fp8_ds_mla"); + STD_TORCH_CHECK(pe_dim == 64, "pe_dim must be 64 for fp8_ds_mla"); + STD_TORCH_CHECK(kv_cache.size(2) == 656 / kv_cache.element_size(), + "kv_cache.size(2) must be 656 bytes for fp8_ds_mla"); + STD_TORCH_CHECK(kv_c.element_size() == 2, + "kv_c.element_size() must be 2 for fp8_ds_mla"); + STD_TORCH_CHECK(k_pe.element_size() == 2, + "k_pe.element_size() must be 2 for fp8_ds_mla"); + } else { + STD_TORCH_CHECK(kv_cache.size(2) == kv_lora_rank + pe_dim); + } + + int kv_c_stride = kv_c.stride(0); + int k_pe_stride = k_pe.stride(0); + int block_stride = kv_cache.stride(0); + int entry_stride = kv_cache.stride(1); + + const torch::stable::accelerator::DeviceGuard device_guard( + kv_c.get_device_index()); + const cudaStream_t stream = get_current_cuda_stream(); + + if (kv_cache_dtype == "fp8_ds_mla") { + dim3 grid(num_tokens); + // For the NoPE part, each tile of 128 elements is handled by half of one + // warp (16 threads). There are 4 total tiles, so 2 warps (64 threads). + // Lanes 0 and 16 of each warp write the scale values for that warp's tiles. + // The RoPE part (last 64 elements) is handled by another 1 warp (32 + // threads). So in total, we use 3 warps (96 threads) per block. + dim3 block(96); + DISPATCH_BY_KV_CACHE_DTYPE(kv_c.scalar_type(), kv_cache_dtype, + CALL_CONCAT_AND_CACHE_DS_MLA); + } else { + dim3 grid(num_tokens); + dim3 block(std::min(kv_lora_rank, 512)); + DISPATCH_BY_KV_CACHE_DTYPE(kv_c.scalar_type(), kv_cache_dtype, + CALL_CONCAT_AND_CACHE_MLA); + } +} + +namespace aphrodite { + +template +__global__ void convert_fp8_kernel(const Tin* __restrict__ src_cache, + Tout* __restrict__ dst_cache, + const float scale, + const int64_t block_stride) { + const int64_t block_idx = blockIdx.x; + for (int i = threadIdx.x; i < block_stride; i += blockDim.x) { + int64_t idx = block_idx * block_stride + i; + dst_cache[idx] = + fp8::scaled_convert(src_cache[idx], scale); + } +} + +} // namespace aphrodite + +#define CALL_CONVERT_FP8(Tout, Tin, KV_DTYPE) \ + aphrodite::convert_fp8_kernel \ + <<>>( \ + reinterpret_cast(src_cache.data_ptr()), \ + reinterpret_cast(dst_cache.data_ptr()), scale, block_stride); + +// Only for testing. +void convert_fp8(torch::stable::Tensor& dst_cache, + torch::stable::Tensor& src_cache, const double scale, + const std::string& kv_cache_dtype) { + torch::stable::Device src_device = src_cache.device(); + torch::stable::Device dst_device = dst_cache.device(); + STD_TORCH_CHECK(src_device.is_cuda(), "src must be on a GPU") + STD_TORCH_CHECK(dst_device.is_cuda(), "dst must be on a GPU") + STD_TORCH_CHECK(src_device.index() == dst_device.index(), + "src and dst must be on the same GPU"); + torch::stable::accelerator::DeviceGuard device_guard(src_device.index()); + + int64_t num_blocks = src_cache.size(0); + int64_t block_stride = src_cache.stride(0); + + dim3 grid(num_blocks); + dim3 block(std::min(block_stride, int64_t(512))); + const cudaStream_t stream = get_current_cuda_stream(); + + if (kv_cache_dtype == "auto") { + if (src_cache.scalar_type() == torch::headeronly::ScalarType::Float) { + CALL_CONVERT_FP8(uint8_t, float, aphrodite::Fp8KVCacheDataType::kAuto); + } else if (src_cache.scalar_type() == torch::headeronly::ScalarType::Half) { + CALL_CONVERT_FP8(uint8_t, uint16_t, aphrodite::Fp8KVCacheDataType::kAuto); + } else if (src_cache.scalar_type() == + torch::headeronly::ScalarType::BFloat16) { + CALL_CONVERT_FP8(uint8_t, __nv_bfloat16, + aphrodite::Fp8KVCacheDataType::kAuto); + } else if (dst_cache.scalar_type() == + torch::headeronly::ScalarType::Float) { + CALL_CONVERT_FP8(float, uint8_t, aphrodite::Fp8KVCacheDataType::kAuto); + } else if (dst_cache.scalar_type() == torch::headeronly::ScalarType::Half) { + CALL_CONVERT_FP8(uint16_t, uint8_t, aphrodite::Fp8KVCacheDataType::kAuto); + } else if (dst_cache.scalar_type() == + torch::headeronly::ScalarType::BFloat16) { + CALL_CONVERT_FP8(__nv_bfloat16, uint8_t, + aphrodite::Fp8KVCacheDataType::kAuto); + } + } else if (kv_cache_dtype == "fp8" || kv_cache_dtype == "fp8_e4m3") { + if (src_cache.scalar_type() == torch::headeronly::ScalarType::Float) { + CALL_CONVERT_FP8(uint8_t, float, aphrodite::Fp8KVCacheDataType::kFp8E4M3); + } else if (src_cache.scalar_type() == torch::headeronly::ScalarType::Half) { + CALL_CONVERT_FP8(uint8_t, uint16_t, + aphrodite::Fp8KVCacheDataType::kFp8E4M3); + } else if (src_cache.scalar_type() == + torch::headeronly::ScalarType::BFloat16) { + CALL_CONVERT_FP8(uint8_t, __nv_bfloat16, + aphrodite::Fp8KVCacheDataType::kFp8E4M3); + } else if (dst_cache.scalar_type() == + torch::headeronly::ScalarType::Float) { + CALL_CONVERT_FP8(float, uint8_t, aphrodite::Fp8KVCacheDataType::kFp8E4M3); + } else if (dst_cache.scalar_type() == torch::headeronly::ScalarType::Half) { + CALL_CONVERT_FP8(uint16_t, uint8_t, + aphrodite::Fp8KVCacheDataType::kFp8E4M3); + } else if (dst_cache.scalar_type() == + torch::headeronly::ScalarType::BFloat16) { + CALL_CONVERT_FP8(__nv_bfloat16, uint8_t, + aphrodite::Fp8KVCacheDataType::kFp8E4M3); + } + } else { + STD_TORCH_CHECK(false, "Unsupported data type: ", kv_cache_dtype); + } +} + +namespace aphrodite { + +// grid is launched with dimensions (batch, num_splits) +template +__global__ void gather_and_maybe_dequant_cache( + const cache_t* __restrict__ src_cache, // [NUM_BLOCKS, BLOCK_SIZE, + // ENTRIES...] + scalar_t* __restrict__ dst, // [TOT_TOKENS, ENTRIES...] + const int32_t* __restrict__ block_table, // [BATCH, BLOCK_INDICES] + const int32_t* __restrict__ cu_seq_lens, // [BATCH+1] + const int32_t* __restrict__ token_to_seq, // [MAX_TOKEN_ACROSS_CHUNK] + const int32_t num_tokens, const int32_t block_size, + const int64_t block_table_stride, const int64_t cache_block_stride, + const int64_t cache_entry_stride, const int64_t dst_entry_stride, + const float* __restrict__ scale, + const int32_t* __restrict__ seq_starts) { // Optional: starting offsets per + // batch + constexpr int vec_size = sizeof(float4) / sizeof(scalar_t); + using ltype = aphrodite::vec_n_t; + using stype = aphrodite::vec_n_t; + // We are adding this for code readability which will be optimized out when + // build in release. + assert(CTA_SIZE == blockDim.x); + +#pragma unroll + for (int token_id = blockIdx.x; token_id < num_tokens; + token_id += gridDim.x) { + int64_t batch_id = token_to_seq[token_id]; + int64_t batch_start = cu_seq_lens[batch_id]; + int64_t batch_end = cu_seq_lens[batch_id + 1]; + int32_t batch_offset = token_id - batch_start; + + if (token_id >= batch_end) return; + int32_t offset = 0; + if (seq_starts != nullptr) { + offset = seq_starts[batch_id]; + } + batch_offset += offset; + int32_t block_table_id = batch_offset / block_size; + int32_t slot_id = batch_offset % block_size; + // seq_starts may push the block index past the end of the batch's block + // table row. + if (block_table_id >= block_table_stride) continue; + int32_t block_table_offset = batch_id * block_table_stride + block_table_id; + int32_t block_id = block_table[block_table_offset]; + int64_t cache_offset = + block_id * cache_block_stride + slot_id * cache_entry_stride; + constexpr int32_t vec_iter_cnt = ENTRY_SIZE / vec_size; + scalar_t* dst_ = dst + token_id * dst_entry_stride; + cache_t* src_ = const_cast(src_cache) + cache_offset; + +#pragma unroll + for (int idx = threadIdx.x; idx < vec_iter_cnt; idx += CTA_SIZE) { + if constexpr (kv_dt == Fp8KVCacheDataType::kAuto) { + reinterpret_cast(dst_)[idx] = + static_cast(reinterpret_cast(src_)[idx]); + } else { + ltype loaded_val = reinterpret_cast(src_)[idx]; + stype store_val; +#pragma unroll + for (int j = 0; j < vec_size; ++j) { + store_val.val[j] = fp8::scaled_convert( + loaded_val.val[j], *scale); + } + reinterpret_cast(dst_)[idx] = store_val; + } + } + // process tail + constexpr int32_t tail_cnt = ENTRY_SIZE % vec_size; + dst_ = dst_ + ENTRY_SIZE - tail_cnt; + src_ = src_ + ENTRY_SIZE - tail_cnt; +#pragma unroll + for (int idx = threadIdx.x; idx < tail_cnt; idx += CTA_SIZE) { + if constexpr (kv_dt == Fp8KVCacheDataType::kAuto) { + dst_[idx] = static_cast(src_[idx]); + } else { + dst_[idx] = + fp8::scaled_convert(src_[idx], *scale); + } + } + } +} + +} // namespace aphrodite + +// Macro to dispatch the kernel based on the data type. +// SCALAR_T is the data type of the destination tensor. +// CACHE_T is the stored data type of kv-cache. +// KV_DTYPE is the real data type of kv-cache. +#define CALL_GATHER_CACHE(SCALAR_T, CACHE_T, KV_DTYPE, ENTRY_SZ) \ + aphrodite::gather_and_maybe_dequant_cache \ + <<>>( \ + reinterpret_cast(src_cache.data_ptr()), \ + reinterpret_cast(dst.data_ptr()), \ + block_table.const_data_ptr(), \ + cu_seq_lens.const_data_ptr(), \ + token_to_seq.const_data_ptr(), num_tokens, block_size, \ + block_table_stride, cache_block_stride, cache_entry_stride, \ + dst_entry_stride, reinterpret_cast(scale.data_ptr()), \ + seq_starts_ptr); + +#define CALL_GATHER_CACHE_576(SCALAR_T, CACHE_T, KV_DTYPE) \ + CALL_GATHER_CACHE(SCALAR_T, CACHE_T, KV_DTYPE, 576) + +#define CALL_GATHER_CACHE_320(SCALAR_T, CACHE_T, KV_DTYPE) \ + CALL_GATHER_CACHE(SCALAR_T, CACHE_T, KV_DTYPE, 320) + +// Gather sequences from the cache into the destination tensor. +// - cu_seq_lens contains the cumulative sequence lengths for each batch +// - block_table contains the cache block indices for each sequence +// - token_to_seq contains the back mapping from token_id to batch_id +// - Optionally, seq_starts (if provided) offsets the starting block index by +// (seq_starts[bid] / page_size) +void gather_and_maybe_dequant_cache( + torch::stable::Tensor const& + src_cache, // [NUM_BLOCKS, BLOCK_SIZE, ENTRIES...] + torch::stable::Tensor const& dst, // [TOT_TOKENS, ENTRIES...] + torch::stable::Tensor const& block_table, // [BATCH, BLOCK_INDICES] + torch::stable::Tensor const& cu_seq_lens, // [BATCH+1] + torch::stable::Tensor const& token_to_seq, // [MAX_TOKEN_ACROSS_CHUNKS] + int64_t num_tokens, const std::string& kv_cache_dtype, + torch::stable::Tensor const& scale, + std::optional seq_starts = std::nullopt) { + torch::stable::accelerator::DeviceGuard device_guard( + src_cache.get_device_index()); + const cudaStream_t stream = get_current_cuda_stream(); + + int32_t block_size = src_cache.size(1); + int32_t head_dim = dst.size(-1); + + STD_TORCH_CHECK( + block_table.scalar_type() == torch::headeronly::ScalarType::Int, + "block_table must be int32"); + STD_TORCH_CHECK( + cu_seq_lens.scalar_type() == torch::headeronly::ScalarType::Int, + "cu_seq_lens must be int32"); + if (seq_starts.has_value()) { + STD_TORCH_CHECK( + seq_starts.value().scalar_type() == torch::headeronly::ScalarType::Int, + "seq_starts must be int32"); + } + STD_TORCH_CHECK( + head_dim == 320 || head_dim == 576, + "gather_and_maybe_dequant_cache only support the head_dim to 320 or 576 " + "for better performance") + + STD_TORCH_CHECK(src_cache.device() == dst.device(), + "src_cache and dst must be on the same device"); + STD_TORCH_CHECK(src_cache.device() == block_table.device(), + "src_cache and block_table must be on the same device"); + STD_TORCH_CHECK(src_cache.device() == cu_seq_lens.device(), + "src_cache and cu_seq_lens must be on the same device"); + if (seq_starts.has_value()) { + STD_TORCH_CHECK(src_cache.device() == seq_starts.value().device(), + "src_cache and seq_starts must be on the same device"); + } + + int64_t block_table_stride = block_table.stride(0); + int64_t cache_block_stride = src_cache.stride(0); + int64_t cache_entry_stride = src_cache.stride(1); + int64_t dst_entry_stride = dst.stride(0); + + constexpr int32_t thread_block_size = 64; + dim3 grid(num_tokens); + dim3 block(thread_block_size); + + const int32_t* seq_starts_ptr = + seq_starts.has_value() ? seq_starts.value().const_data_ptr() + : nullptr; + + if (head_dim == 576) { + DISPATCH_BY_KV_CACHE_DTYPE(dst.scalar_type(), kv_cache_dtype, + CALL_GATHER_CACHE_576); + } else { + DISPATCH_BY_KV_CACHE_DTYPE(dst.scalar_type(), kv_cache_dtype, + CALL_GATHER_CACHE_320); + } +} + +namespace aphrodite { + +// Gather and upconvert FP8 KV cache tokens to BF16 workspace +// Similar to cp_gather_cache but specifically for FP8->BF16 conversion +__global__ void cp_gather_and_upconvert_fp8_kv_cache( + const uint8_t* __restrict__ src_cache, // [NUM_BLOCKS, BLOCK_SIZE, 656] + __nv_bfloat16* __restrict__ dst, // [total_tokens, 576] + const int32_t* __restrict__ block_table, // [num_reqs, BLOCK_INDICES] + const int32_t* __restrict__ workspace_starts, // [num_reqs] + const int32_t num_reqs, const int32_t block_size, + const int32_t total_tokens, const int64_t block_table_stride, + const int64_t cache_block_stride, const int64_t cache_entry_stride, + const int64_t dst_entry_stride, + const int32_t* __restrict__ seq_starts) { // Optional source offsets + const int flat_warp_id = (blockIdx.x * blockDim.x + threadIdx.x) >> 5; + if (flat_warp_id >= total_tokens) return; + const int lane_id = threadIdx.x & 31; + + // Binary search to find which request owns this output token + int lo = 0, hi = num_reqs - 1; + while (lo < hi) { + int mid = (lo + hi + 1) >> 1; + if (workspace_starts[mid] <= flat_warp_id) + lo = mid; + else + hi = mid - 1; + } + const int req_id = lo; + + // Compute physical token address via block table + const int out_token_id = flat_warp_id; + int token_offset = out_token_id - workspace_starts[req_id]; + if (seq_starts != nullptr) token_offset += seq_starts[req_id]; + const int cache_block_idx = token_offset / block_size; + const int offset_in_block = token_offset % block_size; + const int physical_block = + block_table[req_id * block_table_stride + cache_block_idx]; + + const uint8_t* token_ptr = src_cache + physical_block * cache_block_stride + + offset_in_block * cache_entry_stride; + + const int4* nope_src = reinterpret_cast(token_ptr); + const int4 fp8_data = nope_src[lane_id]; + + const float* scales_ptr = reinterpret_cast(token_ptr + 512); + const float scale = scales_ptr[lane_id >> 3]; + + const uint2 fp8_lo = make_uint2(fp8_data.x, fp8_data.y); + const uint2 fp8_hi = make_uint2(fp8_data.z, fp8_data.w); +#ifdef USE_ROCM + const bf16_8_t bf16_lo = + fp8::scaled_vec_conversion(fp8_lo, scale); + const bf16_8_t bf16_hi = + fp8::scaled_vec_conversion(fp8_hi, scale); +#else + const bf16_8_t bf16_lo = + fp8::scaled_vec_conversion(fp8_lo, scale, __NV_E4M3); + const bf16_8_t bf16_hi = + fp8::scaled_vec_conversion(fp8_hi, scale, __NV_E4M3); +#endif + + __nv_bfloat16* dst_ptr = dst + out_token_id * dst_entry_stride; + int4* nope_dst = reinterpret_cast(dst_ptr) + lane_id * 2; + nope_dst[0] = *reinterpret_cast(&bf16_lo); + nope_dst[1] = *reinterpret_cast(&bf16_hi); + + const int* rope_src = reinterpret_cast(token_ptr + 528); + int* rope_dst = reinterpret_cast(dst_ptr + 512); + rope_dst[lane_id] = rope_src[lane_id]; +} + +template +// Note(hc): The cp_gather_cache allows seq_starts to no longer be divisible by +// block_size. +__global__ void cp_gather_cache( + const scalar_t* __restrict__ src_cache, // [NUM_BLOCKS, BLOCK_SIZE, + // ENTRY_SIZE] + scalar_t* __restrict__ dst, // [TOT_TOKENS, ENTRY_SIZE] + const int32_t* __restrict__ block_table, // [BATCH, BLOCK_INDICES] + const int32_t* __restrict__ cu_seq_lens, // [BATCH+1] + const int32_t block_size, const int32_t entry_size, + const int64_t block_table_stride, const int64_t cache_block_stride, + const int64_t cache_entry_stride, const int64_t dst_entry_stride, + const int32_t* __restrict__ seq_starts // Optional: starting offsets per + // batch +) { + const int64_t bid = blockIdx.x; // Batch ID + const int32_t num_splits = gridDim.y; + const int32_t split = blockIdx.y; + const int32_t seq_start = cu_seq_lens[bid]; + const int32_t seq_end = cu_seq_lens[bid + 1]; + const int32_t seq_len = seq_end - seq_start; + const int32_t tot_slots = seq_len; + const int32_t split_slots = cuda_utils::ceil_div(tot_slots, num_splits); + + const int32_t split_start = split * split_slots; + const int32_t split_end = min((split + 1) * split_slots, tot_slots); + + const bool is_active_split = (split_start < tot_slots); + + if (!is_active_split) return; + + // Adjust the pointer for the block_table for this batch. + // If seq_starts is provided, compute an offset based on it + const int32_t batch_offset = bid * block_table_stride; + int32_t offset = split_start; + if (seq_starts != nullptr) { + offset += seq_starts[bid]; + } + int32_t offset_div = offset / block_size; + offset = offset % block_size; + const int32_t* batch_block_table = block_table + batch_offset; + + // Adjust dst pointer based on the cumulative sequence lengths. + dst += seq_start * dst_entry_stride; + + auto copy_entry = [&](const scalar_t* __restrict__ _src, + scalar_t* __restrict__ _dst) { + for (int i = threadIdx.x; i < entry_size; i += blockDim.x) + _dst[i] = _src[i]; + }; + + for (int pid = split_start; pid < split_end; ++pid) { + auto block_id = batch_block_table[offset_div]; + auto block_start_ptr = src_cache + block_id * cache_block_stride; + auto block_dst_ptr = dst + pid * dst_entry_stride; + copy_entry(block_start_ptr + offset * cache_entry_stride, block_dst_ptr); + offset += 1; + // bump to next block + if (offset == block_size) { + offset_div += 1; + offset = 0; + } + } +} +} // namespace aphrodite + +// Macro to dispatch the kernel based on the data type. +#define CALL_CP_GATHER_CACHE(CPY_DTYPE) \ + aphrodite::cp_gather_cache<<>>( \ + reinterpret_cast(src_cache.data_ptr()), \ + reinterpret_cast(dst.data_ptr()), \ + block_table.const_data_ptr(), \ + cu_seq_lens.const_data_ptr(), block_size, entry_size, \ + block_table_stride, cache_block_stride, cache_entry_stride, \ + dst_entry_stride, seq_starts_ptr); + +// Gather sequences from the cache into the destination tensor. +// - cu_seq_lens contains the cumulative sequence lengths for each batch +// - block_table contains the cache block indices for each sequence +// - Optionally, seq_starts (if provided) offsets the starting slot index by +// seq_starts[bid] +void cp_gather_cache( + torch::stable::Tensor const& + src_cache, // [NUM_BLOCKS, BLOCK_SIZE, ENTRIES...] + torch::stable::Tensor const& dst, // [TOT_TOKENS, ENTRIES...] + torch::stable::Tensor const& block_table, // [BATCH, BLOCK_INDICES] + torch::stable::Tensor const& cu_seq_lens, // [BATCH+1] + int64_t batch_size, + std::optional seq_starts = std::nullopt) { + torch::stable::accelerator::DeviceGuard device_guard( + src_cache.get_device_index()); + const cudaStream_t stream = get_current_cuda_stream(); + + int32_t block_size = src_cache.size(1); + int32_t entry_size = torch::stable::flatten(src_cache, 2, -1).size(2); + + STD_TORCH_CHECK( + block_table.scalar_type() == torch::headeronly::ScalarType::Int, + "block_table must be int32"); + STD_TORCH_CHECK( + cu_seq_lens.scalar_type() == torch::headeronly::ScalarType::Int, + "cu_seq_lens must be int32"); + if (seq_starts.has_value()) { + STD_TORCH_CHECK( + seq_starts.value().scalar_type() == torch::headeronly::ScalarType::Int, + "seq_starts must be int32"); + } + + STD_TORCH_CHECK(src_cache.device() == dst.device(), + "src_cache and dst must be on the same device"); + STD_TORCH_CHECK(src_cache.device() == block_table.device(), + "src_cache and block_table must be on the same device"); + STD_TORCH_CHECK(src_cache.device() == cu_seq_lens.device(), + "src_cache and cu_seq_lens must be on the same device"); + if (seq_starts.has_value()) { + STD_TORCH_CHECK(src_cache.device() == seq_starts.value().device(), + "src_cache and seq_starts must be on the same device"); + } + + int64_t block_table_stride = block_table.stride(0); + int64_t cache_block_stride = src_cache.stride(0); + int64_t cache_entry_stride = src_cache.stride(1); + int64_t dst_entry_stride = dst.stride(0); + + // Decide on the number of splits based on the batch size. + int num_splits = batch_size > 128 ? 2 : batch_size > 64 ? 4 : 16; + dim3 grid(batch_size, num_splits); + dim3 block(1024); + + STD_TORCH_CHECK(src_cache.scalar_type() == dst.scalar_type(), + "src_cache and dst must have the same dtype"); + + const int dtype_bits = src_cache.element_size() * 8; + const int32_t* seq_starts_ptr = + seq_starts.has_value() ? seq_starts.value().const_data_ptr() + : nullptr; + + if (dtype_bits == 32) { + CALL_CP_GATHER_CACHE(uint32_t); + } else if (dtype_bits == 16) { + CALL_CP_GATHER_CACHE(uint16_t); + } else if (dtype_bits == 8) { + CALL_CP_GATHER_CACHE(uint8_t); + } else { + STD_TORCH_CHECK(false, "Unsupported data type width: ", dtype_bits); + } +} + +void cp_gather_and_upconvert_fp8_kv_cache( + torch::stable::Tensor const& src_cache, // [NUM_BLOCKS, BLOCK_SIZE, 656] + torch::stable::Tensor const& dst, // [TOT_TOKENS, 576] + torch::stable::Tensor const& block_table, // [BATCH, BLOCK_INDICES] + torch::stable::Tensor const& workspace_starts, // [BATCH] + int64_t batch_size, + std::optional seq_starts = std::nullopt) { + torch::stable::accelerator::DeviceGuard device_guard( + src_cache.get_device_index()); + const cudaStream_t stream = get_current_cuda_stream(); + + int32_t block_size = src_cache.size(1); + int32_t head_dim = dst.size(1); + + STD_TORCH_CHECK( + block_table.scalar_type() == torch::headeronly::ScalarType::Int, + "block_table must be int32"); + STD_TORCH_CHECK( + workspace_starts.scalar_type() == torch::headeronly::ScalarType::Int, + "workspace_starts must be int32"); + if (seq_starts.has_value()) { + STD_TORCH_CHECK( + seq_starts.value().scalar_type() == torch::headeronly::ScalarType::Int, + "seq_starts must be int32"); + } + + STD_TORCH_CHECK(src_cache.device() == dst.device(), + "src_cache and dst must be on the same device"); + STD_TORCH_CHECK(src_cache.device() == block_table.device(), + "src_cache and block_table must be on the same device"); + STD_TORCH_CHECK(src_cache.device() == workspace_starts.device(), + "src_cache and workspace_starts must be on the same device"); + if (seq_starts.has_value()) { + STD_TORCH_CHECK(src_cache.device() == seq_starts.value().device(), + "src_cache and seq_starts must be on the same device"); + } + auto dtype = src_cache.scalar_type(); + STD_TORCH_CHECK( + dtype == torch::headeronly::ScalarType::Byte || // uint8 + dtype == torch::headeronly::ScalarType::Float8_e4m3fn || // fp8 e4m3 + dtype == torch::headeronly::ScalarType::Float8_e5m2, // fp8 e5m2 + "src_cache must be uint8, float8_e4m3fn, or float8_e5m2, but got ", + src_cache.scalar_type()); + STD_TORCH_CHECK(dst.scalar_type() == torch::headeronly::ScalarType::BFloat16, + "dst must be bfloat16"); + STD_TORCH_CHECK(head_dim == 576, "head_dim must be 576 for MLA"); + + int64_t block_table_stride = block_table.stride(0); + int64_t cache_block_stride = src_cache.stride(0); + int64_t cache_entry_stride = src_cache.stride(1); + int64_t dst_entry_stride = dst.stride(0); + + const uint8_t* src_ptr = nullptr; + if (dtype == torch::headeronly::ScalarType::Byte) { + src_ptr = src_cache.const_data_ptr(); + } else { + // float8_e4m3fn or float8_e5m2 + src_ptr = reinterpret_cast(src_cache.data_ptr()); + } + + const int total_tokens = dst.size(0); + constexpr int warps_per_block = 8; + const int grid_size = (total_tokens + warps_per_block - 1) / warps_per_block; + const int block_size_threads = warps_per_block * 32; // 256 threads + const int32_t* seq_starts_ptr = + seq_starts.has_value() ? seq_starts.value().const_data_ptr() + : nullptr; + + aphrodite::cp_gather_and_upconvert_fp8_kv_cache<<< + grid_size, block_size_threads, 0, stream>>>( + src_ptr, reinterpret_cast<__nv_bfloat16*>(dst.data_ptr()), + block_table.const_data_ptr(), + workspace_starts.const_data_ptr(), + static_cast(batch_size), block_size, total_tokens, + block_table_stride, cache_block_stride, cache_entry_stride, + dst_entry_stride, seq_starts_ptr); +} + +// Macro to dispatch the kernel based on the data type. +#define CALL_INDEXER_K_QUANT_AND_CACHE(KV_T, CACHE_T, KV_DTYPE) \ + aphrodite::indexer_k_quant_and_cache_kernel \ + <<>>( \ + reinterpret_cast(k.data_ptr()), \ + reinterpret_cast(kv_cache.data_ptr()), \ + slot_mapping.const_data_ptr(), head_dim, quant_block_size, \ + cache_block_size, cache_block_stride, use_ue8m0); + +void indexer_k_quant_and_cache( + torch::stable::Tensor& k, // [num_tokens, head_dim] + torch::stable::Tensor& kv_cache, // [num_blocks, block_size, cache_stride] + torch::stable::Tensor& slot_mapping, // [num_tokens] + int64_t quant_block_size, // quantization block size + const std::string& scale_fmt) { + int num_tokens = k.size(0); + int head_dim = k.size(1); + int cache_block_size = kv_cache.size(1); + int64_t cache_block_stride = kv_cache.stride(0); + bool use_ue8m0 = scale_fmt == "ue8m0"; + + STD_TORCH_CHECK(k.device() == kv_cache.device(), + "k and kv_cache must be on the same device"); + STD_TORCH_CHECK(k.device() == slot_mapping.device(), + "k and slot_mapping must be on the same device"); + STD_TORCH_CHECK(head_dim % quant_block_size == 0, + "head_dim must be divisible by quant_block_size"); + + constexpr int vec_size = 4; + dim3 grid(num_tokens, (head_dim + quant_block_size * vec_size - 1) / + (quant_block_size * vec_size)); + dim3 block(32, vec_size); + const torch::stable::accelerator::DeviceGuard device_guard( + k.get_device_index()); + const cudaStream_t stream = get_current_cuda_stream(); + + static const std::string kv_cache_dtype = "fp8_e4m3"; + DISPATCH_BY_KV_CACHE_DTYPE(k.scalar_type(), kv_cache_dtype, + CALL_INDEXER_K_QUANT_AND_CACHE); +} + +// Macro to dispatch the kernel based on the data amount. +#define CALL_CP_GATHER_INDEXER_K_QUANT_CACHE(BLOCK_Y_SIZE) \ + aphrodite::cp_gather_indexer_k_quant_cache_kernel \ + <<>>( \ + reinterpret_cast(kv_cache.data_ptr()), \ + reinterpret_cast(dst_k.data_ptr()), \ + reinterpret_cast(dst_scale.data_ptr()), \ + block_table.const_data_ptr(), \ + cu_seq_lens.const_data_ptr(), batch_size, dst_k.stride(0), \ + dst_k.size(1), kv_cache.stride(0), kv_cache.stride(1), \ + kv_cache.size(1), block_table.size(1), num_tokens, \ + quant_block_size); + +void cp_gather_indexer_k_quant_cache( + const torch::stable::Tensor& + kv_cache, // [num_blocks, block_size, cache_stride] + torch::stable::Tensor& dst_k, // [num_tokens, head_dim] + torch::stable::Tensor& + dst_scale, // [num_tokens, head_dim / quant_block_size * 4] + const torch::stable::Tensor& block_table, // [batch_size, num_blocks] + const torch::stable::Tensor& cu_seq_lens // [batch_size + 1] +) { + int batch_size = block_table.size(0); + int num_tokens = dst_k.size(0); + int head_dim = dst_k.size(1); + int quant_block_size = head_dim * 4 / dst_scale.size(1); + + STD_TORCH_CHECK(kv_cache.device() == dst_k.device(), + "kv_cache and dst_k must be on the same device"); + STD_TORCH_CHECK(kv_cache.device() == dst_scale.device(), + "kv_cache and dst_scale must be on the same device"); + STD_TORCH_CHECK(kv_cache.device() == block_table.device(), + "kv_cache and block_table must be on the same device"); + STD_TORCH_CHECK(kv_cache.device() == cu_seq_lens.device(), + "kv_cache and cu_seq_lens must be on the same device"); + STD_TORCH_CHECK(head_dim % quant_block_size == 0, + "head_dim must be divisible by quant_block_size"); + + constexpr int vec_size = 16; + const torch::stable::accelerator::DeviceGuard device_guard( + kv_cache.get_device_index()); + const cudaStream_t stream = get_current_cuda_stream(); + + if (num_tokens < 32) { + CALL_CP_GATHER_INDEXER_K_QUANT_CACHE(1); + } else if (num_tokens < 64) { + CALL_CP_GATHER_INDEXER_K_QUANT_CACHE(2); + } else if (num_tokens < 128) { + CALL_CP_GATHER_INDEXER_K_QUANT_CACHE(4); + } else if (num_tokens < 256) { + CALL_CP_GATHER_INDEXER_K_QUANT_CACHE(8); + } else if (num_tokens < 512) { + CALL_CP_GATHER_INDEXER_K_QUANT_CACHE(16); + } else { + CALL_CP_GATHER_INDEXER_K_QUANT_CACHE(32); + } +} + +// Concatenate ql_nope and q_pe into a contiguous q_out tensor for MLA/DSA. +// Replaces torch.cat((ql_nope, q_pe), dim=-1). +void concat_mla_q( + torch::stable::Tensor& ql_nope, // [num_tokens, num_heads, nope_dim] + torch::stable::Tensor& q_pe, // [num_tokens, num_heads, rope_dim] + torch::stable::Tensor& q_out // [num_tokens, num_heads, nope_dim + + // rope_dim] +) { + const int num_tokens = ql_nope.size(0); + const int num_heads = ql_nope.size(1); + const int nope_dim = ql_nope.size(2); + const int rope_dim = q_pe.size(2); + + STD_TORCH_CHECK(nope_dim % 512 == 0, + "nope_dim must be a multiple of 512, got ", nope_dim); + STD_TORCH_CHECK(rope_dim == 64, "rope_dim must be 64, got ", rope_dim); + STD_TORCH_CHECK(q_out.size(2) == nope_dim + rope_dim); + + STD_TORCH_CHECK(ql_nope.stride(2) == 1, + "ql_nope must have stride 1 in dim 2"); + STD_TORCH_CHECK(q_pe.stride(2) == 1, "q_pe must have stride 1 in dim 2"); + STD_TORCH_CHECK(q_out.stride(2) == 1, "q_out must have stride 1 in dim 2"); + STD_TORCH_CHECK( + ql_nope.scalar_type() == torch::headeronly::ScalarType::Half || + ql_nope.scalar_type() == torch::headeronly::ScalarType::BFloat16, + "ql_nope must be float16 or bfloat16 dtype"); + + if (num_tokens == 0) return; + + constexpr int warps_per_block = 8; + const int total_warps = num_tokens * num_heads; + const int grid_size = (total_warps + warps_per_block - 1) / warps_per_block; + const int block_size = warps_per_block * 32; + + const torch::stable::accelerator::DeviceGuard device_guard( + ql_nope.get_device_index()); + const cudaStream_t stream = get_current_cuda_stream(); + + APHRODITE_STABLE_DISPATCH_HALF_TYPES( + ql_nope.scalar_type(), "concat_mla_q", [&] { + aphrodite::ConcatMLAQKernel + <<>>( + q_out.mutable_data_ptr(), + ql_nope.const_data_ptr(), + q_pe.const_data_ptr(), num_tokens, num_heads, + q_out.stride(0), q_out.stride(1), ql_nope.stride(0), + ql_nope.stride(1), q_pe.stride(0), q_pe.stride(1)); + }); +} 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/custom_all_reduce.cu.orig b/csrc/libtorch_stable/custom_all_reduce.cu.orig new file mode 100644 index 0000000000..6d43964edd --- /dev/null +++ b/csrc/libtorch_stable/custom_all_reduce.cu.orig @@ -0,0 +1,202 @@ +#include "torch_utils.h" + +#include +#include +#include +#include +#include +#include + +#include "custom_all_reduce.cuh" + +// Fake pointer type, must match fptr_t type in ops.h. +// We use this type alias to indicate when pointers are passed in as int64_t. +using fptr_t = int64_t; +static_assert(sizeof(void*) == sizeof(fptr_t)); + +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 % 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]; + for (int i = 0; i < world_size; i++) { + ipc_ptrs[i] = reinterpret_cast(fake_ipc_ptrs[i]); + } + return (fptr_t) new aphrodite::CustomAllreduce( + ipc_ptrs, rank_data.mutable_data_ptr(), rank_data.numel(), rank, + world_size, fully_connected); +} + +/** + * Make sure tensor t's data lies completely within ((char)t.data_ptr()) + + * t.numel() * t.element_size(). This is slightly weaker than t.is_contiguous() + * because it allows transpose of contiguous slice (i.e. slicing the first + * dimension). Currently, we require this because stride information is not + * passed into the kernels and we treat input tensors as flat. + * + * Examples + * A = torch.zeros(3, 3, 3) + * 1. A: OK + * 2. A[1:]: OK + * 3. A.permute(2, 0, 1): OK + * 4. A[1:].permute(2, 0, 1): OK + * 5. A[None].expand(2, -1, -1, -1): Not OK + * 6. A[:, 1:, 1:]: Not OK + */ +bool _is_weak_contiguous(torch::stable::Tensor& t) { + if (t.is_contiguous()) { + return true; + } + int64_t storage_nbytes = 0; + TORCH_ERROR_CODE_CHECK(aoti_torch_get_storage_size(t.get(), &storage_nbytes)); + return storage_nbytes - t.storage_offset() * t.element_size() == + static_cast(t.numel() * t.element_size()); +} + +/** + * Performs an out-of-place allreduce and stores result in out. + * + * If _reg_buffer is null, assumes inp.data_ptr() is already IPC-registered. + * Otherwise, _reg_buffer is assumed to be IPC-registered and inp is first + * copied into _reg_buffer. + */ +void all_reduce(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()) == (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); + if (reg_buffer) { + STD_TORCH_CHECK((input_size) <= (reg_buffer_sz_bytes)); + STD_CUDA_CHECK(cudaMemcpyAsync(reg_buffer, inp.const_data_ptr(), input_size, + cudaMemcpyDeviceToDevice, stream)); + } else { + reg_buffer = inp.mutable_data_ptr(); + } + switch (out.scalar_type()) { + case torch::headeronly::ScalarType::Float: { + fa->allreduce(stream, reinterpret_cast(reg_buffer), + reinterpret_cast(out.mutable_data_ptr()), + out.numel()); + break; + } + case torch::headeronly::ScalarType::Half: { + fa->allreduce(stream, reinterpret_cast(reg_buffer), + reinterpret_cast(out.mutable_data_ptr()), + out.numel()); + break; + } +#if (__CUDA_ARCH__ >= 800 || !defined(__CUDA_ARCH__)) + case torch::headeronly::ScalarType::BFloat16: { + fa->allreduce( + stream, reinterpret_cast(reg_buffer), + reinterpret_cast(out.mutable_data_ptr()), out.numel()); + break; + } +#endif + default: + throw std::runtime_error( + "custom allreduce only supports float32, float16 and bfloat16"); + } +} + +void dispose(fptr_t _fa) { + delete reinterpret_cast(_fa); +} + +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]; + for (int i = 0; i < fake_ipc_ptrs.size(); i++) { + ipc_ptrs[i] = reinterpret_cast(fake_ipc_ptrs[i]); + } + fa->register_buffer(ipc_ptrs); +} + +// Use vector to represent byte data for python binding compatibility. +std::tuple, std::vector> +get_graph_buffer_ipc_meta(fptr_t _fa) { + auto fa = reinterpret_cast(_fa); + auto [handle, offsets] = fa->get_graph_buffer_ipc_meta(); + std::vector bytes(handle.begin(), handle.end()); + return std::make_tuple(bytes, offsets); +} + +// Use vector to represent byte data for python binding compatibility. +void register_graph_buffers(fptr_t _fa, + const std::vector>& handles, + const std::vector>& offsets) { + auto fa = reinterpret_cast(_fa); + std::vector bytes; + bytes.reserve(handles.size()); + for (int i = 0; i < handles.size(); i++) { + bytes.emplace_back(handles[i].begin(), handles[i].end()); + } + bytes.reserve(handles.size()); + fa->register_graph_buffers(bytes, offsets); +} + +std::tuple allocate_shared_buffer_and_handle( + int64_t size) { + int device_index; + STD_CUDA_CHECK(cudaGetDevice(&device_index)); + const torch::stable::accelerator::DeviceGuard device_guard(device_index); + void* buffer; + cudaStreamCaptureMode mode = cudaStreamCaptureModeRelaxed; + const cudaStream_t stream = get_current_cuda_stream(device_index); + STD_CUDA_CHECK(cudaThreadExchangeStreamCaptureMode(&mode)); + + // Allocate buffer +#if defined(USE_ROCM) + // data buffers need to be "uncached" for signal on MI200 + STD_CUDA_CHECK( + hipExtMallocWithFlags((void**)&buffer, size, hipDeviceMallocUncached)); +#else + STD_CUDA_CHECK(cudaMalloc((void**)&buffer, size)); +#endif + STD_CUDA_CHECK(cudaMemsetAsync(buffer, 0, size, stream)); + STD_CUDA_CHECK(cudaStreamSynchronize(stream)); + STD_CUDA_CHECK(cudaThreadExchangeStreamCaptureMode(&mode)); + + // Create IPC memhandle for the allocated buffer. + // Will use it in open_mem_handle. + auto handle = torch::stable::empty( + {static_cast(sizeof(cudaIpcMemHandle_t))}, + torch::headeronly::ScalarType::Byte, std::nullopt, + torch::stable::Device(torch::stable::DeviceType::CPU)); + STD_CUDA_CHECK(cudaIpcGetMemHandle( + (cudaIpcMemHandle_t*)handle.mutable_data_ptr(), buffer)); + + return std::make_tuple(reinterpret_cast(buffer), handle); +} + +fptr_t open_mem_handle(torch::stable::Tensor& mem_handle) { + void* ipc_ptr; + STD_CUDA_CHECK(cudaIpcOpenMemHandle( + (void**)&ipc_ptr, + *((const cudaIpcMemHandle_t*)mem_handle.const_data_ptr()), + cudaIpcMemLazyEnablePeerAccess)); + return reinterpret_cast(ipc_ptr); +} + +void free_shared_buffer(fptr_t buffer) { + STD_CUDA_CHECK(cudaFree(reinterpret_cast(buffer))); +} 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/grouped_topk_kernels.cu.orig b/csrc/libtorch_stable/moe/grouped_topk_kernels.cu.orig new file mode 100644 index 0000000000..dadfbc3a5c --- /dev/null +++ b/csrc/libtorch_stable/moe/grouped_topk_kernels.cu.orig @@ -0,0 +1,1124 @@ +/* + * Adapted from + * https://github.com/NVIDIA/TensorRT-LLM/blob/v1.3.0rc2/cpp/tensorrt_llm/kernels/noAuxTcKernels.cu + * Copyright (c) 2025, The vLLM team. + * SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & + * AFFILIATES. All rights reserved. SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include "moeTopKFuncs.cuh" + +#include +#include + +#include "libtorch_stable/torch_utils.h" + +#include +#include +#include +#include +#include +#include +#include +namespace cg = cooperative_groups; + +namespace aphrodite { +namespace moe { + +constexpr unsigned FULL_WARP_MASK = 0xffffffff; +static constexpr int WARP_SIZE = 32; +static constexpr int NumNemotronExperts = 512; +static constexpr int NumKimiK2Experts = 384; +static constexpr int NumDeepseekExperts = 256; +static constexpr int MaxSupportedExpertCount = + std::max({NumNemotronExperts, NumKimiK2Experts, NumDeepseekExperts}); +static constexpr int MaxNumExpertsUnit = 128; +static constexpr int NumTopGroupScores = 2; +static constexpr int DefaultMaxNumTopExperts = 8; +static constexpr int MaxSupportedTopExperts = 22; +static constexpr int MaxNumTopGroups = 4; +// The empirical value for small batch +static constexpr int PDLEnableTokens = 16; + +namespace warp_topk { + +template +__host__ __device__ constexpr T round_up_to_multiple_of(T len) { + if (len == 0) { + return 0; + } + return ((len - 1) / size + 1) * size; +} + +template +constexpr __host__ __device__ bool isPowerOf2(T v) { + return (v && !(v & (v - 1))); +} + +template +__forceinline__ __device__ bool is_better_than(T val, T baseline) { + return (val > baseline && greater) || (val < baseline && !greater); +} + +template +__forceinline__ __device__ bool is_better_than(T val, T baseline, idxT index, + idxT baseline_index) { + bool res = (val > baseline && greater) || (val < baseline && !greater); + if (val == baseline) { + res = (index < baseline_index && greater) || + (index < baseline_index && !greater); + } + return res; +} + +template +struct BitonicMerge { + // input should be a bitonic sequence, and sort it to be a monotonic sequence + __device__ static void merge(T* __restrict__ val_arr, + idxT* __restrict__ idx_arr) { + static_assert(isPowerOf2(size)); + static_assert(size >= 2 * WARP_SIZE); + constexpr int arr_len = size / WARP_SIZE; + + constexpr int stride = arr_len / 2; + for (int i = 0; i < stride; ++i) { + int const other_i = i + stride; + T& val = val_arr[i]; + T& other_val = val_arr[other_i]; + bool is_better; + if constexpr (is_stable) { + is_better = is_better_than(val, other_val, idx_arr[i], + idx_arr[other_i]); + } else { + is_better = is_better_than(val, other_val); + } + + if (is_better) { + T tmp = val; + val = other_val; + other_val = tmp; + + idxT tmp2 = idx_arr[i]; + idx_arr[i] = idx_arr[other_i]; + idx_arr[other_i] = tmp2; + } + } + + BitonicMerge::merge( + val_arr, idx_arr); + BitonicMerge::merge( + val_arr + arr_len / 2, idx_arr + arr_len / 2); + } +}; + +template +struct BitonicSort { + __device__ static void sort(T* __restrict__ val_arr, + idxT* __restrict__ idx_arr) { + static_assert(isPowerOf2(size)); + static_assert(size >= 2 * WARP_SIZE); + constexpr int arr_len = size / WARP_SIZE; + + BitonicSort::sort(val_arr, idx_arr); + BitonicSort::sort( + val_arr + arr_len / 2, idx_arr + arr_len / 2); + BitonicMerge::merge( + val_arr, idx_arr); + } +}; + +template +struct BitonicSort<32, ascending, T, idxT, is_stable> { + __device__ static void sort(T* __restrict__ val_arr, + idxT* __restrict__ idx_arr) { + int const lane = threadIdx.x % WARP_SIZE; + + // ascending doesn't matter before merging since all we need is a bitonic + // sequence + for (int stage = 0; stage < 4; ++stage) { + for (int stride = (1 << stage); stride > 0; stride /= 2) { + bool reverse = (lane >> stage) & 2; + bool is_second = lane & stride; + + T other = __shfl_xor_sync(FULL_WARP_MASK, *val_arr, stride); + idxT other_idx = __shfl_xor_sync(FULL_WARP_MASK, *idx_arr, stride); + + bool is_better; + if constexpr (is_stable) { + if constexpr (ascending) { + is_better = ((*val_arr > other) || + ((*val_arr == other) && (*idx_arr < other_idx))) != + (reverse != is_second); + } else { + is_better = ((*val_arr > other) || + ((*val_arr == other) && (*idx_arr > other_idx))) != + (reverse != is_second); + } + } else { + is_better = (*val_arr != other && + (*val_arr > other) != (reverse != is_second)); + } + if (is_better) { + *val_arr = other; + *idx_arr = other_idx; + } + } + } + + BitonicMerge<32, ascending, ascending, T, idxT, is_stable>::merge(val_arr, + idx_arr); + } +}; + +template +struct BitonicMerge<32, ascending, reverse, T, idxT, is_stable> { + __device__ static void merge(T* __restrict__ val_arr, + idxT* __restrict__ idx_arr) { + int const lane = threadIdx.x % WARP_SIZE; + for (int stride = WARP_SIZE / 2; stride > 0; stride /= 2) { + bool is_second = lane & stride; + T& val = *val_arr; + T other = __shfl_xor_sync(FULL_WARP_MASK, val, stride); + idxT& idx = *idx_arr; + idxT other_idx = __shfl_xor_sync(FULL_WARP_MASK, idx, stride); + + bool is_better; + if constexpr (is_stable) { + if constexpr (ascending) { + is_better = ((*val_arr > other) || + ((*val_arr == other) && (*idx_arr < other_idx))) == + (reverse != is_second); // for min + } else { + is_better = ((*val_arr > other) || + ((*val_arr == other) && (*idx_arr > other_idx))) == + (reverse != is_second); // for max + } + } else { + is_better = + (val != other && ((val > other) == (ascending != is_second))); + } + + if (is_better) { + val = other; + idx = other_idx; + } + } + } +}; + +template +class WarpSort { + public: + __device__ WarpSort(idxT k, T dummy) + : lane_(threadIdx.x % WARP_SIZE), k_(k), dummy_(dummy) { + static_assert(capacity >= WARP_SIZE && isPowerOf2(capacity)); + + for (int i = 0; i < max_arr_len_; ++i) { + val_arr_[i] = dummy_; + idx_arr_[i] = 0; + } + } + + // load and merge k sorted values + __device__ void load_sorted(T const* __restrict__ in, + idxT const* __restrict__ in_idx, idxT start) { + idxT idx = start + WARP_SIZE - 1 - lane_; + for (int i = max_arr_len_ - 1; i >= 0; --i, idx += WARP_SIZE) { + if (idx < start + k_) { + T t = in[idx]; + bool is_better; + if constexpr (is_stable) { + is_better = + is_better_than(t, val_arr_[i], in_idx[idx], idx_arr_[i]); + } else { + is_better = is_better_than(t, val_arr_[i]); + } + if (is_better) { + val_arr_[i] = t; + idx_arr_[i] = in_idx[idx]; + } + } + } + + BitonicMerge::merge( + val_arr_, idx_arr_); + } + + __device__ void dump(T* __restrict__ out, idxT* __restrict__ out_idx) const { + for (int i = 0; i < max_arr_len_; ++i) { + idxT out_i = i * WARP_SIZE + lane_; + if (out_i < k_) { + out[out_i] = val_arr_[i]; + out_idx[out_i] = idx_arr_[i]; + } + } + } + + __device__ void dumpIdx(idxT* __restrict__ out_idx) const { + for (int i = 0; i < max_arr_len_; ++i) { + idxT out_i = i * WARP_SIZE + lane_; + if (out_i < k_) { + out_idx[out_i] = idx_arr_[i]; + } + } + } + + // Accessors for per-lane selected value/index. + // NOTE: For the common case `capacity == WARP_SIZE`, `max_arr_len_ == 1` + // and callers should use `i == 0`. + __device__ __forceinline__ idxT get_idx(int i = 0) const { + return idx_arr_[i]; + } + + __device__ __forceinline__ T get_val(int i = 0) const { return val_arr_[i]; } + + protected: + static constexpr int max_arr_len_ = capacity / WARP_SIZE; + + T val_arr_[max_arr_len_]; + idxT idx_arr_[max_arr_len_]; + + int const lane_; + idxT const k_; + T const dummy_; + +}; // end class WarpSort + +template +class WarpSelect : public WarpSort { + public: + __device__ WarpSelect(idxT k, T dummy) + : WarpSort(k, dummy), + k_th_(dummy), + k_th_idx_(0), + k_th_lane_((k - 1) % WARP_SIZE) { + extern __shared__ char smem_buf[]; // extern __shared__ T smem_buf[]; + + int const num_of_warp = blockDim.x / WARP_SIZE; + int const warp_id = threadIdx.x / WARP_SIZE; + val_smem_ = reinterpret_cast(smem_buf); + val_smem_ += warp_id * WARP_SIZE; + idx_smem_ = reinterpret_cast( + smem_buf + + round_up_to_multiple_of<256>(num_of_warp * sizeof(T) * WARP_SIZE)); + idx_smem_ += warp_id * WARP_SIZE; + } + + __device__ void add(T const* in, idxT start, idxT end) { + idxT const end_for_fullwarp = + round_up_to_multiple_of(end - start) + start; + for (idxT i = start + lane_; i < end_for_fullwarp; i += WARP_SIZE) { + T val = (i < end) ? in[i] : dummy_; + add(val, i); + } + } + + __device__ void add(T val, idxT idx) { + bool do_add; + if constexpr (is_stable) { + do_add = is_better_than(val, k_th_, idx, k_th_idx_); + } else { + do_add = is_better_than(val, k_th_); + } + + uint32_t mask = __ballot_sync(FULL_WARP_MASK, do_add); + if (mask == 0) { + return; + } + + int pos = smem_buf_len_ + __popc(mask & ((0x1u << lane_) - 1)); + if (do_add && pos < WARP_SIZE) { + val_smem_[pos] = val; + idx_smem_[pos] = idx; + do_add = false; + } + smem_buf_len_ += __popc(mask); + if (smem_buf_len_ >= WARP_SIZE) { + __syncwarp(); + merge_buf_(val_smem_[lane_], idx_smem_[lane_]); + smem_buf_len_ -= WARP_SIZE; + } + if (do_add) { + pos -= WARP_SIZE; + val_smem_[pos] = val; + idx_smem_[pos] = idx; + } + __syncwarp(); + } + + __device__ void done() { + if (smem_buf_len_) { + T val = (lane_ < smem_buf_len_) ? val_smem_[lane_] : dummy_; + idxT idx = (lane_ < smem_buf_len_) ? idx_smem_[lane_] : 0; + merge_buf_(val, idx); + } + } + + private: + __device__ void set_k_th_() { + k_th_ = __shfl_sync(FULL_WARP_MASK, val_arr_[max_arr_len_ - 1], k_th_lane_); + if constexpr (is_stable) { + k_th_idx_ = + __shfl_sync(FULL_WARP_MASK, idx_arr_[max_arr_len_ - 1], k_th_lane_); + } + } + + __device__ void merge_buf_(T val, idxT idx) { + BitonicSort::sort(&val, &idx); + + T& old = val_arr_[max_arr_len_ - 1]; + + bool is_better; + if constexpr (is_stable) { + is_better = + is_better_than(val, old, idx, idx_arr_[max_arr_len_ - 1]); + } else { + is_better = is_better_than(val, old); + } + + if (is_better) { + old = val; + idx_arr_[max_arr_len_ - 1] = idx; + } + + BitonicMerge::merge( + val_arr_, idx_arr_); + + set_k_th_(); + } + + using WarpSort::max_arr_len_; + using WarpSort::val_arr_; + using WarpSort::idx_arr_; + using WarpSort::lane_; + using WarpSort::k_; + using WarpSort::dummy_; + + T* val_smem_; + idxT* idx_smem_; + int smem_buf_len_ = 0; + + T k_th_; + idxT k_th_idx_; + int const k_th_lane_; +}; // end class WarpSelect +} // namespace warp_topk + +template +__device__ inline T_OUT cuda_cast(T_IN val) { + return val; +} + +template <> +__device__ inline float cuda_cast(__nv_bfloat16 val) { + return __bfloat162float(val); +} + +template +__device__ inline T neg_inf() { + // cuda::std::numeric_limits::infinity() returns `0` for [T=bf16 or fp16] + // so we need to cast from fp32 + return cuda_cast(-cuda::std::numeric_limits::infinity()); +} + +template +__device__ inline bool is_finite(const T val) { +#if (__CUDACC_VER_MAJOR__ * 10000 + __CUDACC_VER_MINOR__ * 100 >= 120800) + return cuda::std::isfinite(val); +#else + return isfinite(cuda_cast(val)); +#endif +} + +// Scoring function enums +enum ScoringFunc { + SCORING_NONE = 0, // no activation function + SCORING_SIGMOID = 1 // apply sigmoid +}; + +// Efficient sigmoid approximation from TensorRT-LLM +__device__ inline float sigmoid_accurate(float x) { + return 0.5f * tanhf(0.5f * x) + 0.5f; +} + +template +__device__ inline T apply_sigmoid(T val) { + float f = cuda_cast(val); + return cuda_cast(sigmoid_accurate(f)); +} + +template +__device__ inline T apply_scoring(T val) { + if constexpr (SF == SCORING_NONE) { + return val; + } else if constexpr (SF == SCORING_SIGMOID) { + return apply_sigmoid(val); + } else { + static_assert(SF == SCORING_NONE || SF == SCORING_SIGMOID, + "Unsupported ScoringFunc in apply_scoring"); + return val; + } +} + +template +__device__ void topk_with_k2(T* output, T const* input, BiasT const* bias, + cg::thread_block_tile<32> const& tile, + int32_t const lane_id, + int const num_experts_per_group) { + // Get the top2 per thread + T largest = neg_inf(); + T second_largest = neg_inf(); + + if (num_experts_per_group > WARP_SIZE) { + for (int i = lane_id; i < num_experts_per_group; i += WARP_SIZE) { + T value = apply_scoring(input[i]); + value = value + static_cast(bias[i]); + + if (value > largest) { + second_largest = largest; + largest = value; + } else if (value > second_largest) { + second_largest = value; + } + } + } else { + for (int i = lane_id; i < num_experts_per_group; i += WARP_SIZE) { + T value = apply_scoring(input[i]); + value = value + static_cast(bias[i]); + largest = value; + } + } + // Get the top2 warpwise + T max1 = cg::reduce(tile, largest, cg::greater()); + + T max2 = max1; + bool equal_to_max1 = (max1 == largest); + + int count_max1 = __popc(__ballot_sync(FULL_WARP_MASK, equal_to_max1)); + + if (count_max1 == 1) { + largest = (largest == max1) ? second_largest : largest; + max2 = cg::reduce(tile, largest, cg::greater()); + } + + if (lane_id == 0) { + *output = max1 + max2; + } +} + +template +__global__ void grouped_topk_fused_kernel( + T* scores, float* topk_values, IdxT* topk_indices, BiasT const* bias, + int64_t const num_tokens, int64_t const num_experts, int64_t const n_group, + int64_t const topk_group, int64_t const topk, bool renormalize, + double routed_scaling_factor) { + int32_t const token_id = static_cast(blockIdx.x); + if (token_id >= num_tokens) { + return; + } + + int32_t const warp_id = threadIdx.x / WARP_SIZE; + int32_t const lane_id = threadIdx.x % WARP_SIZE; + + int32_t const n_group_i32 = static_cast(n_group); + int32_t const topk_group_i32 = static_cast(topk_group); + int32_t const topk_i32 = static_cast(topk); + int32_t const num_experts_i32 = static_cast(num_experts); + + int32_t const num_warps = blockDim.x / WARP_SIZE; + if (warp_id >= n_group_i32 || num_warps < n_group_i32) { + return; + } + + int32_t const num_experts_per_group = num_experts_i32 / n_group_i32; + + T* scores_token = scores + static_cast(token_id) * num_experts; + + cg::thread_block block = cg::this_thread_block(); + cg::thread_block_tile<32> tile = cg::tiled_partition<32>(block); + + extern __shared__ char smem_buf[]; + // warpSelect internal staging buffer layout + size_t const val_bytes = + static_cast(num_warps) * WARP_SIZE * sizeof(T); + size_t const val_bytes_aligned = + warp_topk::round_up_to_multiple_of<256>(val_bytes); + size_t const idx_bytes = + static_cast(num_warps) * WARP_SIZE * sizeof(int32_t); + size_t const internal_bytes = val_bytes_aligned + idx_bytes; + + // user-managed shared memory starts after warpSelect internal staging. + uintptr_t ptr_u = reinterpret_cast(smem_buf + internal_bytes); + ptr_u = (ptr_u + 15) & ~static_cast(15); // align to 16B + T* s_group_scores = reinterpret_cast(ptr_u); + +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) + cudaGridDependencySynchronize(); // I think all prolog can be put before + // acqbulk because it's ptr arithmetic +#endif + + // phase 1: per-group scan + int32_t const group_offset = warp_id * num_experts_per_group; + topk_with_k2(s_group_scores + warp_id, + scores_token + group_offset, bias + group_offset, + tile, lane_id, num_experts_per_group); + + __syncthreads(); + + // phase 2: warp0 selects groups + merges candidates to final topk + if (warp_id != 0) { + return; + } + + topk_values += static_cast(token_id) * topk; + topk_indices += static_cast(token_id) * topk; + + // select topk_group groups by group score + warp_topk::WarpSelect + group_sel(static_cast(topk_group_i32), neg_inf()); + + // all lanes must participate in WarpSelect::add(). + T gscore = (lane_id < n_group_i32) ? s_group_scores[lane_id] : neg_inf(); + group_sel.add(gscore, lane_id); + group_sel.done(); + + // proceed only if the k-th selected group score is not -inf + bool proceed = false; + if (topk_group_i32 > 0) { + int const kth_lane = topk_group_i32 - 1; + // broadcast the k-th selected group score to all lanes + T kth_val = __shfl_sync(FULL_WARP_MASK, group_sel.get_val(0), kth_lane); + proceed = (kth_val != neg_inf()); + } + + if (!proceed) { + for (int i = lane_id; i < topk_i32; i += WARP_SIZE) { + topk_indices[i] = static_cast(i); + topk_values[i] = 1.0f / static_cast(topk_i32); + } +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) + cudaTriggerProgrammaticLaunchCompletion(); +#endif + return; + } + + // merge per-group topk candidates for selected groups, then select topk + warp_topk::WarpSelect + expert_sel(static_cast(topk_i32), neg_inf()); + + // selected group ids reside in lanes [0, topk_group) + int32_t sel_gid_lane = (lane_id < topk_group_i32) ? group_sel.get_idx(0) : 0; + + // add candidates from selected groups to expert_sel + for (int32_t g = 0; g < topk_group_i32; ++g) { + int32_t gid = __shfl_sync(FULL_WARP_MASK, sel_gid_lane, g); + int32_t const offset = gid * num_experts_per_group; + int32_t const align_num_experts_per_group = + warp_topk::round_up_to_multiple_of(num_experts_per_group); + for (int32_t i = lane_id; i < align_num_experts_per_group; i += WARP_SIZE) { + // all lanes must call `add()` the same number of times. + T cand = neg_inf(); + int32_t idx = 0; + if (i < num_experts_per_group) { + idx = offset + i; + T input = scores_token[idx]; + if (is_finite(input)) { + T score = apply_scoring(input); + cand = score + static_cast(bias[idx]); + } + } + expert_sel.add(cand, idx); + } + } + expert_sel.done(); + + // compute unbiased routing weights + optional renorm. + float lane_unbiased = 0.0f; + IdxT lane_idx = 0; + if (lane_id < topk_i32) { + lane_idx = static_cast(expert_sel.get_idx(0)); + T in = scores_token[static_cast(lane_idx)]; + lane_unbiased = cuda_cast(apply_scoring(in)); + } + + float topk_sum = 1e-20f; + if (renormalize) { + topk_sum += cg::reduce(tile, lane_unbiased, cg::plus()); + } + + float scale = static_cast(routed_scaling_factor); + if (renormalize) { + scale /= topk_sum; + } + + if (lane_id < topk_i32) { + topk_indices[lane_id] = lane_idx; + topk_values[lane_id] = lane_unbiased * scale; + } + +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) + cudaTriggerProgrammaticLaunchCompletion(); +#endif +} + +template +__global__ void grouped_topk_fused_small_expert_count_kernel( + T* scores, float* topkValues, IdxT* topkIndices, BiasT const* routingBias, + int64_t const numTokens, int64_t const numGroup, int64_t const topkGroup, + int64_t const topk, int64_t const numExperts, + int64_t const numExpertsPerGroup, bool const renormalize, + double const routedScalingFactor) { +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) + cudaGridDependencySynchronize(); +#endif + // declare shared memory structure + // number of experts is bounded by number of threads + __shared__ float __attribute((aligned(128))) smemScoreSigmoid[MaxNumExperts]; + __shared__ float __attribute((aligned(128))) smemScoreBias[MaxNumExperts]; + // number of expert groups is bounded by number of warps + int constexpr NumWarps = MaxNumExperts / WARP_SIZE; + __shared__ float __attribute((aligned(128))) smemGroupScores[NumWarps]; + + // needed for warp reduce + auto block = cg::this_thread_block(); + auto warp = cg::tiled_partition(block); + + // for the final reduction of weight norm, only some lanes need to participate + int32_t laneIdx = threadIdx.x % WARP_SIZE; + int32_t warpIdx = __shfl_sync(0xffffffff, threadIdx.x / WARP_SIZE, 0); + + if constexpr (UseGroups) { + if (warpIdx >= numGroup) { + return; + } + } + // note that for invalid scores, we simply use a negative value: + // they work well even with the compacted format used in topK, and + // sigmoid / bias activated scores cannot be negative + const float invalidScoreFloat = float{-INFINITY}; + + // load bias already; each warp represents one expert group + auto threadExpert = threadIdx.x; + bool expertSelected = threadExpert < numExperts; + if constexpr (UseGroups) { + threadExpert = warpIdx * numExpertsPerGroup + laneIdx; + expertSelected = laneIdx < numExpertsPerGroup; + } + + auto scoreIdx = int64_t{blockIdx.x} * int64_t{numExperts} + threadExpert; + auto biasVal = expertSelected ? static_cast(routingBias[threadExpert]) + : invalidScoreFloat; + topkValues += blockIdx.x * topk; + topkIndices += blockIdx.x * topk; + + // get our assigned thread score; each warp represents one expert group + float score = + expertSelected ? static_cast(scores[scoreIdx]) : invalidScoreFloat; + auto scoreSigmoid = apply_scoring(score); + // write the sigmoid score to shared for later use + if (expertSelected) { + smemScoreSigmoid[threadExpert] = scoreSigmoid; + } + + // get the score with bias + // note that with invalid values, because sigmoid is < 1 and bias is -1, + // we must get a negative value, which is smaller than any valid value + auto scoreBias = float{scoreSigmoid + float{biasVal}}; + + if (expertSelected) { + smemScoreBias[threadExpert] = scoreBias; + } + + // registers for top group score reduction + float topExpGroupScores[NumTopGroupScores]; + [[maybe_unused]] int32_t topExpGroupIdx[NumTopGroupScores]; + float topGroups[MaxNumTopGroups]; // bound of numGroup + int32_t topGroupIdx[MaxNumTopGroups]; + float expertScoreGroup[MaxNumTopGroups]; + int32_t expertIdxGroup[MaxNumTopGroups]; + float topScores[MaxNumTopExperts]; // bound of topk + int32_t topExperts[MaxNumTopExperts]; + + if constexpr (UseGroups) { + reduce_topk::reduceTopK(warp, topExpGroupScores, topExpGroupIdx, scoreBias, + threadExpert, + /* minValue */ invalidScoreFloat); + + // get the final group score and write it to shared + if (warp.thread_rank() == 0) { + auto groupScore = topExpGroupScores[0] + topExpGroupScores[1]; + smemGroupScores[warpIdx] = groupScore; + } + } + + // make group scores available to all warps + __syncthreads(); + + if constexpr (UseGroups) { + if (warpIdx == 0) { + // a single warp performs the selection of top groups, and goes on to + // select the final experts + float groupScore = + laneIdx < numGroup ? smemGroupScores[laneIdx] : invalidScoreFloat; + + reduce_topk::reduceTopK(warp, topGroups, topGroupIdx, groupScore, laneIdx, + /* minValue */ invalidScoreFloat); + // final expert selection: get relevant indexes and scores from shared +#pragma unroll + for (int ii = 0; ii < MaxNumTopGroups; ++ii) { // bound of numGroup + auto groupIdx = topGroupIdx[ii]; + expertIdxGroup[ii] = groupIdx * numExpertsPerGroup + laneIdx; + + expertScoreGroup[ii] = (ii < topkGroup) && expertSelected + ? smemScoreBias[expertIdxGroup[ii]] + : invalidScoreFloat; + } + + reduce_topk::reduceTopK(warp, topScores, topExperts, expertScoreGroup, + expertIdxGroup, /* minValue */ invalidScoreFloat, + topk); + } + } else if constexpr (MaxNumExperts > MaxNumExpertsUnit) { + // without groups, and the expert number is larger than MaxNumExpertsUnit, + // we need to use multiple warps to calculate the intermediate topk results + + int constexpr NumExpertWarps = (MaxNumExperts - 1) / MaxNumExpertsUnit + 1; + int constexpr NumInterTopK = NumExpertWarps * MaxNumTopExperts; + __shared__ float + __attribute((aligned(128))) smemInterTopScores[NumInterTopK]; + __shared__ int32_t + __attribute((aligned(128))) smemInterTopExperts[NumInterTopK]; + if (warpIdx < NumExpertWarps) { + int offset = warpIdx * WARP_SIZE * MaxNumTopGroups; +#pragma unroll + for (int ii = 0; ii < MaxNumTopGroups; ++ii) { + auto expertIdx = ii * WARP_SIZE + laneIdx; + expertIdxGroup[ii] = offset + expertIdx; + expertScoreGroup[ii] = offset + expertIdx < numExperts + ? smemScoreBias[offset + expertIdx] + : invalidScoreFloat; + } + reduce_topk::reduceTopK(warp, topScores, topExperts, expertScoreGroup, + expertIdxGroup, + /* minValue */ invalidScoreFloat, topk); + + if (laneIdx < topk) { + smemInterTopScores[warpIdx * MaxNumTopExperts + laneIdx] = + topScores[laneIdx]; + smemInterTopExperts[warpIdx * MaxNumTopExperts + laneIdx] = + topExperts[laneIdx]; + } else if (laneIdx >= topk && laneIdx < MaxNumTopExperts) { + smemInterTopScores[warpIdx * MaxNumTopExperts + laneIdx] = + invalidScoreFloat; + smemInterTopExperts[warpIdx * MaxNumTopExperts + laneIdx] = + MaxNumExperts - 1; + } + } + __syncthreads(); + if (warpIdx == 0) { + int constexpr NumInterTopKPerThread = (NumInterTopK - 1) / WARP_SIZE + 1; + float intermediateScore[NumInterTopKPerThread]; + int32_t intermediateExpert[NumInterTopKPerThread]; + for (int i = laneIdx; i < NumInterTopKPerThread * WARP_SIZE; + i += WARP_SIZE) { + int ii = i / WARP_SIZE; + if (i < NumInterTopK) { + intermediateScore[ii] = smemInterTopScores[i]; + intermediateExpert[ii] = smemInterTopExperts[i]; + } else { + intermediateScore[ii] = invalidScoreFloat; + intermediateExpert[ii] = MaxNumExperts - 1; + } + } + reduce_topk::reduceTopK(warp, topScores, topExperts, intermediateScore, + intermediateExpert, + /* minValue */ invalidScoreFloat, topk); + } + } else { + // without groups, and the expert number is smaller than MaxNumExpertsUnit + // each thread just takes `MaxNumTopGroups` experts + if (warpIdx == 0) { +#pragma unroll + for (int ii = 0; ii < MaxNumTopGroups; ++ii) { + auto expertIdx = ii * WARP_SIZE + laneIdx; + expertIdxGroup[ii] = expertIdx; + expertScoreGroup[ii] = expertIdx < numExperts ? smemScoreBias[expertIdx] + : invalidScoreFloat; + } + reduce_topk::reduceTopK(warp, topScores, topExperts, expertScoreGroup, + expertIdxGroup, + /* minValue */ invalidScoreFloat, topk); + } + } + + if (warpIdx == 0) { + // determine our lane's expert index and write to output + int32_t expertIdx = + laneIdx < topk ? topExperts[laneIdx] : MaxNumExperts - 1; + float scoreNorm = laneIdx < topk ? smemScoreSigmoid[expertIdx] : 0.F; + float finalScore = static_cast(scoreNorm * routedScalingFactor); + // norm the value + if (renormalize) { + auto redNorm = cg::reduce(warp, scoreNorm, cg::plus{}); + finalScore /= (redNorm + 1e-20); + } + // store the topk scores and experts to output + if (laneIdx < topk) { + topkValues[laneIdx] = finalScore; + topkIndices[laneIdx] = expertIdx; + } + } + +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) + cudaTriggerProgrammaticLaunchCompletion(); +#endif +} + +template +void invokeNoAuxTc(T* scores, float* topk_values, IdxT* topk_indices, + BiasT const* bias, int64_t const num_tokens, + int64_t const num_experts, int64_t const n_group, + int64_t const topk_group, int64_t const topk, + bool const renormalize, double const routed_scaling_factor, + const bool enable_pdl = false, + cudaStream_t const stream = 0) { + cudaLaunchConfig_t config; + config.stream = stream; + cudaLaunchAttribute attrs[1]; + attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attrs[0].val.programmaticStreamSerializationAllowed = enable_pdl; + config.numAttrs = 1; + config.attrs = attrs; + + // Check if we can use the optimized + // grouped_topk_fused_small_expert_count_kernel + bool const is_single_group = + (n_group == 1) && (topk_group == 1) && + (num_experts <= MaxSupportedExpertCount) && + (topk <= DefaultMaxNumTopExperts || topk == MaxSupportedTopExperts); + + int64_t const experts_per_group = num_experts / n_group; + bool const is_multi_group = + (n_group > 1) && (num_experts <= NumDeepseekExperts) && + (experts_per_group <= WARP_SIZE) && + (experts_per_group * topk_group <= MaxNumExpertsUnit) && + (topk <= DefaultMaxNumTopExperts) && (topk_group <= MaxNumTopGroups); + + if (is_single_group || is_multi_group) { + auto* kernel_instance = + &grouped_topk_fused_small_expert_count_kernel; + int num_threads = NumDeepseekExperts; + if (is_single_group) { + // Special case for Nemotron, which selects top 22 from 512 experts, and 1 + // group only. + if (num_experts == NumNemotronExperts && n_group == 1 && + topk == MaxSupportedTopExperts) { + kernel_instance = &grouped_topk_fused_small_expert_count_kernel< + T, BiasT, IdxT, SF, NumNemotronExperts, false, + MaxSupportedTopExperts>; + num_threads = NumNemotronExperts; + } else if (num_experts > NumKimiK2Experts && + num_experts <= MaxSupportedExpertCount) { + kernel_instance = &grouped_topk_fused_small_expert_count_kernel< + T, BiasT, IdxT, SF, MaxSupportedExpertCount, false>; + num_threads = MaxSupportedExpertCount; + } else if (num_experts > MaxNumExpertsUnit && + num_experts <= NumKimiK2Experts) { + kernel_instance = &grouped_topk_fused_small_expert_count_kernel< + T, BiasT, IdxT, SF, NumKimiK2Experts, false>; + num_threads = NumKimiK2Experts; + } else { + kernel_instance = &grouped_topk_fused_small_expert_count_kernel< + T, BiasT, IdxT, SF, MaxNumExpertsUnit, false>; + num_threads = MaxNumExpertsUnit; + } + } + config.gridDim = num_tokens; + config.blockDim = num_threads; + config.dynamicSmemBytes = 0; + cudaLaunchKernelEx(&config, kernel_instance, scores, topk_values, + topk_indices, bias, num_tokens, n_group, topk_group, + topk, num_experts, num_experts / n_group, renormalize, + routed_scaling_factor); + } else { + auto* kernel_instance = &grouped_topk_fused_kernel; + // One block per token; one warp per group. + config.gridDim = static_cast(num_tokens); + config.blockDim = static_cast(n_group) * WARP_SIZE; + // Dynamic shared memory: WarpSelect staging + per-group topk buffers. + int32_t const num_warps = static_cast(n_group); + size_t const val_bytes = + static_cast(num_warps) * WARP_SIZE * sizeof(T); + size_t const val_bytes_aligned = + warp_topk::round_up_to_multiple_of<256>(val_bytes); + size_t const idx_bytes = + static_cast(num_warps) * WARP_SIZE * sizeof(int32_t); + size_t const internal_bytes = val_bytes_aligned + idx_bytes; + size_t const extra_bytes = 16 + static_cast(n_group) * sizeof(T); + config.dynamicSmemBytes = internal_bytes + extra_bytes; + cudaLaunchKernelEx(&config, kernel_instance, scores, topk_values, + topk_indices, bias, num_tokens, num_experts, n_group, + topk_group, topk, renormalize, routed_scaling_factor); + } +} + +#define INSTANTIATE_NOAUX_TC(T, BiasT, IdxT, SF) \ + template void invokeNoAuxTc( \ + T * scores, float* topk_values, IdxT* topk_indices, BiasT const* bias, \ + int64_t const num_tokens, int64_t const num_experts, \ + int64_t const n_group, int64_t const topk_group, int64_t const topk, \ + bool const renormalize, double const routed_scaling_factor, \ + const bool enable_pdl, cudaStream_t const stream); + +INSTANTIATE_NOAUX_TC(float, float, int32_t, SCORING_SIGMOID); +INSTANTIATE_NOAUX_TC(float, half, int32_t, SCORING_SIGMOID); +INSTANTIATE_NOAUX_TC(float, __nv_bfloat16, int32_t, SCORING_SIGMOID); +INSTANTIATE_NOAUX_TC(half, float, int32_t, SCORING_SIGMOID); +INSTANTIATE_NOAUX_TC(half, half, int32_t, SCORING_SIGMOID); +INSTANTIATE_NOAUX_TC(half, __nv_bfloat16, int32_t, SCORING_SIGMOID); +INSTANTIATE_NOAUX_TC(__nv_bfloat16, float, int32_t, SCORING_SIGMOID); +INSTANTIATE_NOAUX_TC(__nv_bfloat16, half, int32_t, SCORING_SIGMOID); +INSTANTIATE_NOAUX_TC(__nv_bfloat16, __nv_bfloat16, int32_t, SCORING_SIGMOID); +INSTANTIATE_NOAUX_TC(float, float, int32_t, SCORING_NONE); +INSTANTIATE_NOAUX_TC(float, half, int32_t, SCORING_NONE); +INSTANTIATE_NOAUX_TC(float, __nv_bfloat16, int32_t, SCORING_NONE); +INSTANTIATE_NOAUX_TC(half, float, int32_t, SCORING_NONE); +INSTANTIATE_NOAUX_TC(half, half, int32_t, SCORING_NONE); +INSTANTIATE_NOAUX_TC(half, __nv_bfloat16, int32_t, SCORING_NONE); +INSTANTIATE_NOAUX_TC(__nv_bfloat16, float, int32_t, SCORING_NONE); +INSTANTIATE_NOAUX_TC(__nv_bfloat16, half, int32_t, SCORING_NONE); +INSTANTIATE_NOAUX_TC(__nv_bfloat16, __nv_bfloat16, int32_t, SCORING_NONE); +} // end namespace moe +} // namespace aphrodite + +std::tuple grouped_topk( + torch::stable::Tensor const& scores, int64_t n_group, int64_t topk_group, + int64_t topk, bool renormalize, double routed_scaling_factor, + torch::stable::Tensor const& bias, int64_t scoring_func = 0) { + const auto data_type = scores.scalar_type(); + const auto bias_type = bias.scalar_type(); + STD_TORCH_CHECK(scores.dim() == 2, "scores must be a 2D Tensor"); + const int64_t num_tokens = scores.size(0); + const int64_t num_experts = scores.size(1); + STD_TORCH_CHECK(n_group > 0, "n_group must be positive"); + STD_TORCH_CHECK(topk > 0, "topk must be positive"); + STD_TORCH_CHECK(topk_group > 0, "topk_group must be positive"); + STD_TORCH_CHECK(topk_group <= n_group, "topk_group must be <= n_group"); + STD_TORCH_CHECK(num_experts % n_group == 0, + "num_experts should be divisible by n_group"); + STD_TORCH_CHECK(n_group <= 32, + "n_group should be smaller than or equal to 32 for now"); + STD_TORCH_CHECK(topk <= 32, + "topk should be smaller than or equal to 32 for now"); + STD_TORCH_CHECK(topk <= topk_group * (num_experts / n_group), + "topk must be <= topk_group * (num_experts / n_group)"); + STD_TORCH_CHECK( + scoring_func == aphrodite::moe::SCORING_NONE || + scoring_func == aphrodite::moe::SCORING_SIGMOID, + "scoring_func must be SCORING_NONE (0) or SCORING_SIGMOID (1)"); + + // Always output float32 for topk_values (eliminates Python-side conversion) + auto topk_values = torch::stable::new_empty( + scores, {num_tokens, topk}, torch::headeronly::ScalarType::Float); + auto topk_indices = torch::stable::new_empty( + scores, {num_tokens, topk}, torch::headeronly::ScalarType::Int); + + const bool pdl_flag = num_tokens <= aphrodite::moe::PDLEnableTokens; + + const torch::stable::accelerator::DeviceGuard device_guard( + scores.get_device_index()); + const cudaStream_t stream = + get_current_cuda_stream(scores.get_device_index()); + auto const sf = static_cast(scoring_func); + +#define LAUNCH_KERNEL_SF(T, BiasT, IdxT) \ + do { \ + switch (sf) { \ + case aphrodite::moe::SCORING_NONE: \ + aphrodite::moe::invokeNoAuxTc( \ + reinterpret_cast(scores.mutable_data_ptr()), \ + reinterpret_cast(topk_values.mutable_data_ptr()), \ + reinterpret_cast(topk_indices.mutable_data_ptr()), \ + reinterpret_cast(bias.data_ptr()), num_tokens, \ + num_experts, n_group, topk_group, topk, renormalize, \ + routed_scaling_factor, pdl_flag, stream); \ + break; \ + case aphrodite::moe::SCORING_SIGMOID: \ + aphrodite::moe::invokeNoAuxTc( \ + reinterpret_cast(scores.mutable_data_ptr()), \ + reinterpret_cast(topk_values.mutable_data_ptr()), \ + reinterpret_cast(topk_indices.mutable_data_ptr()), \ + reinterpret_cast(bias.data_ptr()), num_tokens, \ + num_experts, n_group, topk_group, topk, renormalize, \ + routed_scaling_factor, pdl_flag, stream); \ + break; \ + default: \ + STD_TORCH_CHECK(false, "Unsupported scoring_func"); \ + break; \ + } \ + } while (0) + +#define LAUNCH_KERNEL(T, IdxT) \ + do { \ + switch (bias_type) { \ + case torch::headeronly::ScalarType::Half: \ + LAUNCH_KERNEL_SF(T, half, IdxT); \ + break; \ + case torch::headeronly::ScalarType::Float: \ + LAUNCH_KERNEL_SF(T, float, IdxT); \ + break; \ + case torch::headeronly::ScalarType::BFloat16: \ + LAUNCH_KERNEL_SF(T, __nv_bfloat16, IdxT); \ + break; \ + default: \ + STD_TORCH_CHECK( \ + false, \ + "Invalid bias dtype, only supports float16, float32, and " \ + "bfloat16"); \ + break; \ + } \ + } while (0) + + switch (data_type) { + case torch::headeronly::ScalarType::Half: + // Handle Float16 + LAUNCH_KERNEL(half, int32_t); + break; + case torch::headeronly::ScalarType::Float: + // Handle Float32 + LAUNCH_KERNEL(float, int32_t); + break; + case torch::headeronly::ScalarType::BFloat16: + // Handle BFloat16 + LAUNCH_KERNEL(__nv_bfloat16, int32_t); + break; + default: + // Handle other data types + STD_TORCH_CHECK( + false, "Invalid dtype, only supports float16, float32, and bfloat16"); + break; + } +#undef LAUNCH_KERNEL +#undef LAUNCH_KERNEL_SF + return {topk_values, topk_indices}; +} 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/moe/moeTopKFuncs.cuh.orig b/csrc/libtorch_stable/moe/moeTopKFuncs.cuh.orig new file mode 100644 index 0000000000..81a216bd6f --- /dev/null +++ b/csrc/libtorch_stable/moe/moeTopKFuncs.cuh.orig @@ -0,0 +1,257 @@ +/* + * Adapted from + * https://github.com/NVIDIA/TensorRT-LLM/blob/v1.3.0rc2/cpp/tensorrt_llm/kernels/moeTopKFuncs.cuh + * Copyright (c) 2026, The vLLM team. + * SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION. All rights + * reserved. SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#pragma once + +#include +#include +#include + +namespace aphrodite { +namespace moe { +namespace reduce_topk { +namespace cg = cooperative_groups; +static constexpr int kWARP_SIZE = 32; + +template +struct TopKRedType { + using T = T_; + static_assert( + std::is_same_v || std::is_same_v || + std::is_same_v || std::is_same_v, + "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; + + static __host__ __device__ inline TypeCmp makeCmpVal(T val, int32_t idx = 0) { + auto valueBits = cub::Traits::TwiddleIn( + reinterpret_cast::UnsignedBits&>(val)); + TypeCmp compactTmp = valueBits; + compactTmp = (compactTmp << kMoveBits) | (0xFFFF & (kMaxIdx - idx)); + // Use 65535 minus idx to give higher priority to elements with smaller + // indices. + return compactTmp; + } + + static __host__ __device__ void unpack(T& value, int32_t& index, + TypeCmp cmp) { + // Since “65535-idx” is always smaller than 65536 and positive, we can + // directly use it as the lower 16 bits + index = kMaxIdx - static_cast((cmp & 0xFFFF)); + + auto compactTmp = cmp >> kMoveBits; + auto valueBits = cub::Traits::TwiddleOut( + reinterpret_cast::UnsignedBits&>(compactTmp)); + value = reinterpret_cast(valueBits); + } + + __host__ __device__ TopKRedType() = default; + + __host__ __device__ TopKRedType(T val, int32_t idx) + : compValIdx(makeCmpVal(val, idx)) {} + + __host__ __device__ operator TypeCmp() const noexcept { return compValIdx; } + + __device__ inline TypeCmp reduce( + cg::thread_block_tile const& warp) { + return cg::reduce(warp, compValIdx, cg::greater{}); + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template +struct TopKIdx { + // by default, empty +}; + +template +struct TopKIdx { + static constexpr int K = K_; + int32_t val[K]; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +#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 +struct Sort; + +template +struct Sort<1, RedType> { + static __device__ void run(RedType* topK) {} +}; + +template +struct Sort<2, RedType> { + static __device__ void run(RedType* topK) { TOPK_SWAP(0, 1); } +}; + +template +struct Sort<3, RedType> { + static __device__ void run(RedType* topK) { + TOPK_SWAP(0, 1); + TOPK_SWAP(1, 2); + TOPK_SWAP(0, 1); + } +}; + +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); + } +}; + +template +__forceinline__ __device__ void reduceTopK( + cg::thread_block_tile const& warp, Type (&out)[K], + int32_t (&outIdx)[K], Type value, int32_t idx, 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"); + using RedType = TopKRedType; + RedType topK{value, idx}; + 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 + 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) { + static_assert(K > 0, "Top K must have K > 0"); + 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"); + using RedType = TopKRedType; + RedType topK[N]; +#pragma unroll + for (int nn = 0; nn < N; ++nn) { + topK[nn] = RedType{value[nn], idx[nn]}; + } + + if constexpr (!IsSorted) { + Sort::run(topK); + } + typename RedType::TypeCmp packedMax{}; +#pragma unroll + for (int kk = 0; kk < actualK; ++kk) { + bool update = kk > 0 && packedMax == topK[0].compValIdx; +#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 +__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(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"); + using RedType = TopKRedType; + + 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; + + for (int ii = 0; ii < numResults; ++ii) { + topKBufferValue[ii] = minValue; + topKBufferIdx[ii] = ii * kWARP_SIZE - 1; + } + 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]; + } + } + + reduceTopKFunc(warp, out, outIdx, topKBufferValue, + topKBufferIdx, minValue, actualK); + } +}; + +#undef TOPK_SWAP + +} // namespace reduce_topk +} // namespace moe +} // namespace aphrodite 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/ops.h.orig b/csrc/libtorch_stable/ops.h.orig new file mode 100644 index 0000000000..c2bc9c4c66 --- /dev/null +++ b/csrc/libtorch_stable/ops.h.orig @@ -0,0 +1,621 @@ +#pragma once + +#include +#include +#include + +#include +#include +#include + +#include + +inline torch::stable::Tensor weak_ref_tensor(torch::stable::Tensor& tensor) { + // Ensure tensor is on CUDA + STD_TORCH_CHECK(tensor.device().is_cuda(), "Tensor must be on CUDA device"); + + // Get the raw data pointer + void* data_ptr = tensor.mutable_data_ptr(); + + /// Create a new tensor from the raw data pointer + return torch::stable::from_blob(data_ptr, tensor.sizes(), tensor.strides(), + tensor.device(), tensor.scalar_type()); +} + +void per_token_group_quant_fp8(const torch::stable::Tensor& input, + torch::stable::Tensor& output_q, + torch::stable::Tensor& output_s, + int64_t group_size, double eps, double fp8_min, + double fp8_max, bool scale_ue8m0, + bool dummy_is_scale_transposed, + bool dummy_is_tma_aligned); + +// Fused activation quantisation + DeepGEMM-compatible UE8M0-packed scales. +void per_token_group_quant_8bit_packed(const torch::stable::Tensor& input, + torch::stable::Tensor& output_q, + torch::stable::Tensor& output_s_packed, + int64_t group_size, double eps, + double min_8bit, double max_8bit); + +void per_token_group_quant_int8(const torch::stable::Tensor& input, + torch::stable::Tensor& output_q, + torch::stable::Tensor& output_s, + int64_t group_size, double eps, double int8_min, + double int8_max); + +torch::stable::Tensor permute_cols(torch::stable::Tensor const& A, + torch::stable::Tensor const& perm); + +std::tuple +dry_scan_penalties_cpu(const torch::stable::Tensor& token_history_ids, + const torch::stable::Tensor& token_history_lens, + const torch::stable::Tensor& dry_multiplier, + const torch::stable::Tensor& allowed_lengths, + const torch::stable::Tensor& sequence_breakers_ids, + const torch::stable::Tensor& ranges, + const torch::stable::Tensor& max_ngram, + const torch::stable::Tensor& max_occurrences, + const torch::stable::Tensor& early_exit_match_len, + int64_t vocab_size); + +#ifndef USE_ROCM +bool cutlass_scaled_mm_supports_fp8(int64_t cuda_device_capability); +bool cutlass_scaled_mm_supports_block_fp8(int64_t cuda_device_capability); +bool cutlass_group_gemm_supported(int64_t cuda_device_capability); + +void cutlass_scaled_mm(torch::stable::Tensor& out, + torch::stable::Tensor const& a, + torch::stable::Tensor const& b, + torch::stable::Tensor const& a_scales, + torch::stable::Tensor const& b_scales, + std::optional const& bias); + +void cutlass_moe_mm(torch::stable::Tensor& out_tensors, + torch::stable::Tensor const& a_tensors, + torch::stable::Tensor const& b_tensors, + torch::stable::Tensor const& a_scales, + torch::stable::Tensor const& b_scales, + torch::stable::Tensor const& expert_offsets, + torch::stable::Tensor const& problem_sizes, + torch::stable::Tensor const& a_strides, + torch::stable::Tensor const& b_strides, + torch::stable::Tensor const& c_strides, bool per_act_token, + bool per_out_ch); + +void cutlass_scaled_mm_azp(torch::stable::Tensor& out, + torch::stable::Tensor const& a, + torch::stable::Tensor const& b, + torch::stable::Tensor const& a_scales, + torch::stable::Tensor const& b_scales, + torch::stable::Tensor const& azp_adj, + std::optional const& azp, + std::optional const& bias); + +void get_cutlass_moe_mm_data( + const torch::stable::Tensor& topk_ids, + torch::stable::Tensor& expert_offsets, + torch::stable::Tensor& problem_sizes1, + torch::stable::Tensor& problem_sizes2, + torch::stable::Tensor& input_permutation, + torch::stable::Tensor& output_permutation, const int64_t num_experts, + const int64_t n, const int64_t k, + const std::optional& blockscale_offsets, + const bool is_gated); + +void get_cutlass_moe_mm_problem_sizes_from_expert_offsets( + const torch::stable::Tensor& expert_first_token_offset, + torch::stable::Tensor& problem_sizes1, + torch::stable::Tensor& problem_sizes2, const int64_t n, const int64_t k, + const bool swap_ab); + +void get_cutlass_batched_moe_mm_data( + torch::stable::Tensor& expert_offsets, + torch::stable::Tensor& problem_sizes1, + torch::stable::Tensor& problem_sizes2, + const torch::stable::Tensor& expert_num_tokens, + const int64_t num_local_experts, const int64_t padded_m, const int64_t n, + const int64_t k); + +// FP4/NVFP4 ops +bool cutlass_scaled_mm_supports_fp4(int64_t cuda_device_capability); + +void cutlass_scaled_fp4_mm(torch::stable::Tensor& D, + torch::stable::Tensor const& A, + torch::stable::Tensor const& B, + torch::stable::Tensor const& A_sf, + torch::stable::Tensor const& B_sf, + torch::stable::Tensor const& alpha); + +void cutlass_fp4_group_mm(torch::stable::Tensor& output, + const torch::stable::Tensor& a, + const torch::stable::Tensor& b, + const torch::stable::Tensor& a_blockscale, + const torch::stable::Tensor& b_blockscales, + const torch::stable::Tensor& alphas, + const torch::stable::Tensor& problem_sizes, + const torch::stable::Tensor& expert_offsets, + const torch::stable::Tensor& sf_offsets); + +std::tuple scaled_fp4_quant_func( + torch::stable::Tensor const& input, + torch::stable::Tensor const& input_scale, bool is_sf_swizzled_layout); + +void scaled_fp4_quant_out(torch::stable::Tensor const& input, + torch::stable::Tensor const& input_scale, + bool is_sf_swizzled_layout, + torch::stable::Tensor& output, + torch::stable::Tensor& output_scale); + +void scaled_fp4_experts_quant( + torch::stable::Tensor& output, torch::stable::Tensor& output_scale, + torch::stable::Tensor const& input, + torch::stable::Tensor const& input_global_scale, + torch::stable::Tensor const& input_offset_by_experts, + torch::stable::Tensor const& output_scale_offset_by_experts); + +void silu_and_mul_scaled_fp4_experts_quant( + torch::stable::Tensor& output, torch::stable::Tensor& output_scale, + torch::stable::Tensor const& input, + torch::stable::Tensor const& input_global_scale, + torch::stable::Tensor const& input_offset_by_experts, + torch::stable::Tensor const& output_scale_offset_by_experts); + +void silu_and_mul_nvfp4_quant(torch::stable::Tensor& out, + torch::stable::Tensor& output_block_scale, + torch::stable::Tensor& input, + torch::stable::Tensor& input_global_scale); + +void cutlass_mxfp4_group_mm(torch::stable::Tensor& output, + const torch::stable::Tensor& a, + const torch::stable::Tensor& b, + const torch::stable::Tensor& a_blockscale, + const torch::stable::Tensor& b_blockscales, + const torch::stable::Tensor& problem_sizes, + const torch::stable::Tensor& expert_offsets, + const torch::stable::Tensor& sf_offsets); + +// AWQ ops +torch::stable::Tensor awq_gemm(torch::stable::Tensor _in_feats, + torch::stable::Tensor _kernel, + torch::stable::Tensor _scaling_factors, + torch::stable::Tensor _zeros, + int64_t split_k_iters); + +torch::stable::Tensor awq_dequantize(torch::stable::Tensor _kernel, + torch::stable::Tensor _scaling_factors, + torch::stable::Tensor _zeros, + int64_t split_k_iters, int64_t thx, + int64_t thy); + +// DSV3 fused A GEMM: conditionally compiled so declaration and impl +// registration are in the source file (dsv3_fused_a_gemm.cu) + +// AllSpark ops: declarations are in the source files +// (allspark_repack.cu and allspark_qgemm_w8a16.cu) + +#endif + +// CPU tensor -> CUDA UVA view (shared CUDA/ROCm) +torch::stable::Tensor get_cuda_view_from_cpu_tensor( + torch::stable::Tensor& cpu_tensor); + +// Attention kernels (shared CUDA/ROCm) +void merge_attn_states( + torch::stable::Tensor& output, + std::optional output_lse, + const torch::stable::Tensor& prefix_output, + const torch::stable::Tensor& prefix_lse, + const torch::stable::Tensor& suffix_output, + const torch::stable::Tensor& suffix_lse, + const std::optional prefill_tokens_with_context, + const std::optional& output_scale = std::nullopt); + +torch::stable::Tensor hadacore_transform(torch::stable::Tensor& x, + bool inplace); + +// Layernorm kernels (shared CUDA/ROCm) +void rms_norm(torch::stable::Tensor& out, torch::stable::Tensor& input, + std::optional weight, double epsilon); + +void fused_add_rms_norm(torch::stable::Tensor& input, + torch::stable::Tensor& residual, + std::optional weight, + double epsilon); + +// Layernorm-quant kernels (shared CUDA/ROCm) +void rms_norm_static_fp8_quant(torch::stable::Tensor& out, + torch::stable::Tensor& input, + torch::stable::Tensor& weight, + torch::stable::Tensor& scale, double epsilon); + +void fused_add_rms_norm_static_fp8_quant(torch::stable::Tensor& out, + torch::stable::Tensor& input, + torch::stable::Tensor& residual, + torch::stable::Tensor& weight, + torch::stable::Tensor& scale, + double epsilon); + +// Fused layernorm + dynamic per-token quant kernels (shared CUDA/ROCm) +void rms_norm_dynamic_per_token_quant( + torch::stable::Tensor& out, torch::stable::Tensor const& input, + torch::stable::Tensor const& weight, torch::stable::Tensor& scales, + double const var_epsilon, std::optional scale_ub, + std::optional residual); + +void rms_norm_per_block_quant(torch::stable::Tensor& out, + torch::stable::Tensor const& input, + torch::stable::Tensor const& weight, + torch::stable::Tensor& scales, + double const var_epsilon, + std::optional scale_ub, + std::optional residual, + int64_t group_size, bool is_scale_transposed); + +void silu_and_mul_per_block_quant(torch::stable::Tensor& out, + torch::stable::Tensor const& input, + torch::stable::Tensor& scales, + int64_t group_size, + std::optional scale_ub, + bool is_scale_transposed); + +// Positional encoding kernels (shared CUDA/ROCm) +void rotary_embedding(torch::stable::Tensor& positions, + torch::stable::Tensor& query, + std::optional key, + int64_t head_size, torch::stable::Tensor& cos_sin_cache, + bool is_neox, int64_t rope_dim_offset, bool inverse); + +void fused_qk_norm_rope(torch::stable::Tensor& qkv, int64_t num_heads_q, + int64_t num_heads_k, int64_t num_heads_v, + int64_t head_dim, double eps, + torch::stable::Tensor& q_weight, + torch::stable::Tensor& k_weight, + torch::stable::Tensor& cos_sin_cache, bool is_neox, + torch::stable::Tensor& position_ids, + int64_t forced_token_heads_per_warp); + +torch::stable::Tensor fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert( + torch::stable::Tensor const& q_in, torch::stable::Tensor const& kv, + torch::stable::Tensor& k_cache, torch::stable::Tensor const& slot_mapping, + torch::stable::Tensor const& position_ids, + torch::stable::Tensor const& cos_sin_cache, int64_t q_head_padded, + double eps, int64_t cache_block_size); + +void fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_bf16_insert( + torch::stable::Tensor& q, torch::stable::Tensor const& kv, + torch::stable::Tensor& k_cache, torch::stable::Tensor const& slot_mapping, + torch::stable::Tensor const& position_ids, + torch::stable::Tensor const& cos_sin_cache, double eps, + int64_t cache_block_size); + +void fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_fp8_insert( + torch::stable::Tensor const& q, torch::stable::Tensor const& kv, + torch::stable::Tensor& q_fp8, torch::stable::Tensor& k_cache, + torch::stable::Tensor const& slot_mapping, + torch::stable::Tensor const& position_ids, + torch::stable::Tensor const& cos_sin_cache, + torch::stable::Tensor const& fp8_scale, + torch::stable::Tensor const& q_fp8_scale_inv, double eps, + int64_t cache_block_size); + +#ifndef USE_ROCM +std::tuple +minimax_allreduce_rms_qk(torch::stable::Tensor qkv, + torch::stable::Tensor const& norm_weight_q, + torch::stable::Tensor const& norm_weight_k, + torch::stable::Tensor workspace, int64_t const q_size, + int64_t const kv_size, int64_t const rank, + int64_t const nranks, double const eps); +#endif + +// Horizontally-fused MiniMax-M3 QK-norm + partial NeoX RoPE (+ optional KV / +// index-cache insert). Dense layer: norm+RoPE only; sparse layer: also packs +// the index branch and scatters k/v/index_k into their paged caches. +void fused_minimax_m3_qknorm_rope_kv_insert( + torch::stable::Tensor& qkv, torch::stable::Tensor const& q_norm_weight, + torch::stable::Tensor const& k_norm_weight, + torch::stable::Tensor const& cos_sin_cache, + torch::stable::Tensor const& positions, int64_t num_heads, + int64_t num_kv_heads, int64_t rotary_dim, double eps, + std::optional index_q_norm_weight, + std::optional index_k_norm_weight, + int64_t num_index_heads, std::optional slot_mapping, + std::optional index_slot_mapping, + std::optional kv_cache, + std::optional index_cache, int64_t block_size, + std::optional q_out, + std::optional index_q_out, + const std::string& kv_cache_dtype, bool skip_index_branch); + +#ifdef APHRODITE_ENABLE_KIMI_K3_ATTN_RES +void kimi_k3_attn_res(torch::stable::Tensor& prefix, + torch::stable::Tensor const& delta, + torch::stable::Tensor const& blocks, + torch::stable::Tensor const& norm_weight, + torch::stable::Tensor const& qk_weight, + torch::stable::Tensor const& output_norm_weight, + torch::stable::Tensor& output, int64_t num_blocks, + double eps, double output_norm_eps); +#endif + +// Sampler kernels (shared CUDA/ROCm) +void apply_repetition_penalties_( + torch::stable::Tensor& logits, const torch::stable::Tensor& prompt_mask, + const torch::stable::Tensor& output_mask, + const torch::stable::Tensor& repetition_penalties); + +void top_k_per_row_prefill(const torch::stable::Tensor& logits, + const torch::stable::Tensor& rowStarts, + const torch::stable::Tensor& rowEnds, + torch::stable::Tensor& indices, int64_t numRows, + int64_t stride0, int64_t stride1, int64_t topK); + +void top_k_per_row_decode(const torch::stable::Tensor& logits, int64_t next_n, + const torch::stable::Tensor& seqLens, + torch::stable::Tensor& indices, int64_t numRows, + int64_t stride0, int64_t stride1, int64_t topK); + +void persistent_topk(const torch::stable::Tensor& logits, + const torch::stable::Tensor& lengths, + torch::stable::Tensor& output, + torch::stable::Tensor& workspace, int64_t k, + int64_t max_seq_len); + +#ifdef APHRODITE_ENABLE_COOPERATIVE_TOPK +void cooperative_topk(const torch::stable::Tensor& logits, + const torch::stable::Tensor& lengths, + torch::stable::Tensor& output, + torch::stable::Tensor& workspace, int64_t k, + int64_t max_seq_len); +#endif + +#ifdef APHRODITE_ENABLE_SM89_DSA +// sm89 DeepSeek sparse attention kernels; impls are registered in +// csrc/libtorch_stable/attention/sm89_dsa/. +void sm89_fp8_paged_mqa_logits(const torch::stable::Tensor& q, + const torch::stable::Tensor& pool, + const torch::stable::Tensor& weights, + const torch::stable::Tensor& seq_lens, + const torch::stable::Tensor& block_table, + const torch::stable::Tensor& sched, + torch::stable::Tensor& logits, + bool clean_logits); + +void sm89_paged_mqa_logits_metadata(const torch::stable::Tensor& seq_lens, + torch::stable::Tensor& sched, + int64_t next_n); + +void sm89_fp8_mqa_logits(const torch::stable::Tensor& q, + const torch::stable::Tensor& kv, + const torch::stable::Tensor& kv_scales, + const torch::stable::Tensor& weights, + const torch::stable::Tensor& cu_seqlen_ks, + const torch::stable::Tensor& cu_seqlen_ke, + torch::stable::Tensor& logits); + +void sm89_sparse_mla_fwd(const torch::stable::Tensor& q, + const torch::stable::Tensor& pool, + const torch::stable::Tensor& indices, + torch::stable::Tensor& out, torch::stable::Tensor& lse, + double sm_scale, + std::optional topk_lens); +#endif + +void selective_scan_fwd( + const torch::stable::Tensor& u, const torch::stable::Tensor& delta, + const torch::stable::Tensor& A, const torch::stable::Tensor& B, + const torch::stable::Tensor& C, + const std::optional& D_, + const std::optional& z_, + const std::optional& delta_bias_, + bool delta_softplus, + const std::optional& query_start_loc, + const std::optional& cache_indices, + const std::optional& has_initial_state, + const torch::stable::Tensor& ssm_states, int64_t null_block_id, + int64_t block_size, + const std::optional& block_idx_first_scheduled_token, + const std::optional& block_idx_last_scheduled_token, + const std::optional& initial_state_idx, + const std::optional& cu_chunk_seqlen, + const std::optional& last_chunk_indices); + +using fptr_t = int64_t; +fptr_t init_custom_ar(const std::vector& fake_ipc_ptrs, + torch::stable::Tensor& rank_data, int64_t rank, + bool fully_connected); +void all_reduce(fptr_t _fa, torch::stable::Tensor& inp, + torch::stable::Tensor& out, fptr_t reg_buffer, + int64_t reg_buffer_sz_bytes); +void dispose(fptr_t _fa); +int64_t meta_size(); +void register_buffer(fptr_t _fa, const std::vector& fake_ipc_ptrs); +std::tuple, std::vector> +get_graph_buffer_ipc_meta(fptr_t _fa); +void register_graph_buffers(fptr_t _fa, + const std::vector>& handles, + const std::vector>& offsets); +std::tuple allocate_shared_buffer_and_handle( + int64_t size); +int64_t open_mem_handle(torch::stable::Tensor& mem_handle); +void free_shared_buffer(int64_t buffer); + +// Activation kernels (shared CUDA/ROCm) +void silu_and_mul(torch::stable::Tensor& out, torch::stable::Tensor& input); +void silu_and_mul_clamp(torch::stable::Tensor& out, + torch::stable::Tensor& input, double limit, + double alpha = 1.0, double beta = 0.0); + +void silu_and_mul_quant(torch::stable::Tensor& out, + torch::stable::Tensor& input, + torch::stable::Tensor& scale); + +void persistent_masked_m_silu_mul_quant( + const torch::stable::Tensor& input, // (E, T, 2*H) + const torch::stable::Tensor& tokens_per_expert, // (E) + torch::stable::Tensor& y_q, // (E, T, H) [OUT] + torch::stable::Tensor& y_s, // (E, T, H//group_size) [OUT] + bool use_ue8m0); + +void mul_and_silu(torch::stable::Tensor& out, torch::stable::Tensor& input); +void gelu_and_mul(torch::stable::Tensor& out, torch::stable::Tensor& input); +void gelu_tanh_and_mul(torch::stable::Tensor& out, + torch::stable::Tensor& input); +void fatrelu_and_mul(torch::stable::Tensor& out, torch::stable::Tensor& input, + double threshold); +void swigluoai_and_mul(torch::stable::Tensor& out, torch::stable::Tensor& input, + double alpha = 1.702, double limit = 7.0); +void gelu_new(torch::stable::Tensor& out, torch::stable::Tensor& input); +void gelu_fast(torch::stable::Tensor& out, torch::stable::Tensor& input); +void gelu_quick(torch::stable::Tensor& out, torch::stable::Tensor& input); +void relu_squared(torch::stable::Tensor& out, torch::stable::Tensor& input); + +// INT8 quantization kernels (shared CUDA/ROCm) +void static_scaled_int8_quant(torch::stable::Tensor& out, + torch::stable::Tensor const& input, + torch::stable::Tensor const& scale, + std::optional const& azp); + +void dynamic_scaled_int8_quant(torch::stable::Tensor& out, + torch::stable::Tensor const& input, + torch::stable::Tensor& scales, + std::optional const& azp); + +// FP8 quantization kernels (shared CUDA/ROCm) +void static_scaled_fp8_quant( + torch::stable::Tensor& out, torch::stable::Tensor const& input, + torch::stable::Tensor const& scale, + std::optional group_shape = + std::nullopt); + +void dynamic_scaled_fp8_quant(torch::stable::Tensor& out, + torch::stable::Tensor const& input, + torch::stable::Tensor& scale); + +void dynamic_per_token_scaled_fp8_quant( + torch::stable::Tensor& out, torch::stable::Tensor const& input, + torch::stable::Tensor& scale, + std::optional const& scale_ub); + +// GPTQ kernels (shared CUDA/ROCm) +torch::stable::Tensor gptq_gemm(torch::stable::Tensor a, + torch::stable::Tensor b_q_weight, + torch::stable::Tensor b_gptq_qzeros, + torch::stable::Tensor b_gptq_scales, + torch::stable::Tensor b_g_idx, bool use_exllama, + bool use_v2_format, int64_t bit); + +void gptq_shuffle(torch::stable::Tensor q_weight, torch::stable::Tensor q_perm, + int64_t bit); + +// Cache ops (shared CUDA/ROCm) +void swap_blocks(torch::stable::Tensor& src, torch::stable::Tensor& dst, + int64_t block_size_in_bytes, + const torch::stable::Tensor& block_mapping); + +// Batch swap: submit all block copies in a single driver call. +void swap_blocks_batch(const torch::stable::Tensor& src_ptrs, + const torch::stable::Tensor& dst_ptrs, + const torch::stable::Tensor& sizes, + bool is_src_access_order_any); + +void reshape_and_cache(torch::stable::Tensor& key, torch::stable::Tensor& value, + torch::stable::Tensor& key_cache, + torch::stable::Tensor& value_cache, + torch::stable::Tensor& slot_mapping, + const std::string& kv_cache_dtype, + torch::stable::Tensor& k_scale, + torch::stable::Tensor& v_scale); + +void reshape_and_cache_flash( + torch::stable::Tensor& key, torch::stable::Tensor& value, + torch::stable::Tensor& key_cache, torch::stable::Tensor& value_cache, + torch::stable::Tensor& slot_mapping, const std::string& kv_cache_dtype, + torch::stable::Tensor& k_scale, torch::stable::Tensor& v_scale); + +void concat_and_cache_mla(torch::stable::Tensor& kv_c, + torch::stable::Tensor& k_pe, + torch::stable::Tensor& kv_cache, + torch::stable::Tensor& slot_mapping, + const std::string& kv_cache_dtype, + torch::stable::Tensor& scale); + +// NOTE: k_pe and kv_c order is flipped compared to concat_and_cache_mla +void concat_and_cache_mla_rope_fused( + torch::stable::Tensor& positions, torch::stable::Tensor& q_pe, + torch::stable::Tensor& k_pe, torch::stable::Tensor& kv_c, + torch::stable::Tensor& rope_cos_sin_cache, bool rope_is_neox, + torch::stable::Tensor& slot_mapping, torch::stable::Tensor& kv_cache, + const std::string& kv_cache_dtype, + torch::stable::Tensor& kv_cache_quant_scale); + +// Just for unittest +void convert_fp8(torch::stable::Tensor& dst_cache, + torch::stable::Tensor& src_cache, const double scale, + const std::string& kv_cache_dtype); + +void gather_and_maybe_dequant_cache( + torch::stable::Tensor const& src_cache, // [NUM_BLOCKS, BLOCK_SIZE, + // ENTRIES...] + torch::stable::Tensor const& dst, // [TOT_TOKENS, ENTRIES...] + torch::stable::Tensor const& block_table, // [BATCH, BLOCK_INDICES] + torch::stable::Tensor const& cu_seq_lens, // [BATCH+1] + torch::stable::Tensor const& token_to_seq, // [MAX_TOKEN_ACROSS_CHUNKS] + int64_t num_tokens, const std::string& kv_cache_dtype, + torch::stable::Tensor const& scale, + std::optional seq_starts = std::nullopt); + +// TODO(hc): cp_gather_cache need support scaled kvcahe in the future. +void cp_gather_cache( + torch::stable::Tensor const& src_cache, // [NUM_BLOCKS, BLOCK_SIZE, + // ENTRIES...] + torch::stable::Tensor const& dst, // [TOT_TOKENS, ENTRIES...] + torch::stable::Tensor const& block_table, // [BATCH, BLOCK_INDICES] + torch::stable::Tensor const& cu_seq_lens, // [BATCH+1] + int64_t batch_size, + std::optional seq_starts = std::nullopt); + +// Gather and upconvert FP8 KV cache to BF16 workspace +void cp_gather_and_upconvert_fp8_kv_cache( + torch::stable::Tensor const& src_cache, // [NUM_BLOCKS, BLOCK_SIZE, + // 656] + torch::stable::Tensor const& dst, // [TOT_TOKENS, 576] + torch::stable::Tensor const& block_table, // [BATCH, BLOCK_INDICES] + torch::stable::Tensor const& workspace_starts, // [BATCH] + int64_t batch_size, + std::optional seq_starts = std::nullopt); + +// Indexer K quantization and cache function +void indexer_k_quant_and_cache( + torch::stable::Tensor& k, // [num_tokens, head_dim] + torch::stable::Tensor& kv_cache, // [num_blocks, block_size, + // cache_stride] + torch::stable::Tensor& slot_mapping, // [num_tokens] + int64_t quant_block_size, // quantization block size + const std::string& scale_fmt); + +// Concatenate query nope and rope for MLA/DSA attention +void concat_mla_q( + torch::stable::Tensor& ql_nope, // [num_tokens, num_heads, nope_dim] + torch::stable::Tensor& q_pe, // [num_tokens, num_heads, rope_dim] + torch::stable::Tensor& q_out); // [num_tokens, num_heads, nope_dim + + // rope_dim] + +// Extract function to gather quantized K cache +void cp_gather_indexer_k_quant_cache( + const torch::stable::Tensor& kv_cache, // [num_blocks, block_size, + // cache_stride] + torch::stable::Tensor& dst_k, // [num_tokens, head_dim] + torch::stable::Tensor& dst_scale, // [num_tokens, head_dim / + // quant_block_size * 4] + const torch::stable::Tensor& block_table, // [batch_size, num_blocks] + const torch::stable::Tensor& cu_seq_lens); // [batch_size + 1] + +// LongCat n-gram embedding index kernel (see ngram_embedding_kernels.cu). +void ngram_compute_n_gram_ids( + int64_t ne_n, int64_t ne_k, torch::stable::Tensor& ne_weights, + torch::stable::Tensor& ne_mods, + torch::stable::Tensor& exclusive_ne_embedder_size_sums, + torch::stable::Tensor& exclusive_req_len_sums, + torch::stable::Tensor& ne_token_table, torch::stable::Tensor& row_indices, + torch::stable::Tensor& column_starts, torch::stable::Tensor& n_gram_ids); diff --git a/csrc/libtorch_stable/torch_bindings.cpp b/csrc/libtorch_stable/torch_bindings.cpp index 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/libtorch_stable/torch_bindings.cpp.orig b/csrc/libtorch_stable/torch_bindings.cpp.orig new file mode 100644 index 0000000000..0a5ef9c11d --- /dev/null +++ b/csrc/libtorch_stable/torch_bindings.cpp.orig @@ -0,0 +1,1014 @@ +#include "ops.h" +#include "cuda_utils.h" +#include "core/registration.h" + +#include + +// Register ops with STABLE_TORCH_LIBRARY for libtorch stable ABI compatibility. +// Note: We register under namespace "_C" so ops are accessible as +// torch.ops._C. for compatibility with existing code. +STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) { + // Compute per-token-group FP8 quantized tensor and scaling factor. + // The dummy arguments are here so we can correctly fuse with RMSNorm. + ops.def( + "per_token_group_fp8_quant(Tensor input, Tensor! output_q, Tensor! " + "output_s, " + "int group_size, float eps, float fp8_min, float fp8_max, bool " + "scale_ue8m0, bool dummy_is_scale_transposed, bool dummy_is_tma_aligned " + ") -> ()"); + // Compute per-token-group 8-bit quantized tensor and UE8M0-packed, + // TMA-aligned scales for DeepGEMM. + ops.def( + "per_token_group_fp8_quant_packed(Tensor input, Tensor! output_q, " + "Tensor! output_s_packed, int group_size, float eps, float fp8_min, " + "float fp8_max) -> ()"); + // Compute per-token-group INT8 quantized tensor and scaling factor. + ops.def( + "per_token_group_quant_int8(Tensor input, Tensor! output_q, Tensor! " + "output_s, int group_size, float eps, float int8_min, float int8_max) -> " + "()"); + ops.def("permute_cols(Tensor A, Tensor perm) -> Tensor"); + ops.def( + "dry_scan_penalties(" + " Tensor token_history_ids," + " Tensor token_history_lens," + " Tensor dry_multiplier," + " Tensor allowed_lengths," + " Tensor sequence_breakers_ids," + " Tensor ranges," + " Tensor max_ngram," + " Tensor max_occurrences," + " Tensor early_exit_match_len," + " int vocab_size) -> (Tensor, Tensor, Tensor)"); + + ops.def("get_cuda_view_from_cpu_tensor(Tensor cpu_tensor) -> Tensor"); + +#ifndef USE_ROCM + + // Note about marlin kernel 'workspace' arguments: + // Technically these should be mutable since they are modified by the kernel. + // But since they are set back to zero once the kernel is finished we can + // hand wave and say that they have no net effect. + // + // The reason to mark 'workspace' as immutable is so that they don't interfere + // with using ScalarType arguments in the ops. If they are marked as mutable, + // pytorch throws an assert in + // 'torch._higher_order_ops._register_effectful_op' that prevents these + // kernels from being torch.compile'd. + // See the following document for more info on custom types and ops that use + // custom types: + // https://docs.google.com/document/d/18fBMPuOJ0fY5ZQ6YyrHUppw9FA332CpNtgB6SOIgyuA + + // Machete (Dense) Optimized Mixed Precision GEMM for Hopper. + ops.def( + "machete_supported_schedules(" + " ScalarType a_type," + " int b_type," + " ScalarType? maybe_group_scales_type," + " ScalarType? maybe_group_zeros_type," + " ScalarType? maybe_channel_scales_type," + " ScalarType? maybe_token_scales_type," + " ScalarType? maybe_out_type" + ") -> str[]"); + ops.def( + "machete_mm(" + " Tensor A," + " Tensor B," + " int b_type," + " ScalarType? out_type," + " Tensor? group_scales," + " Tensor? group_zeros," + " int? group_size," + " Tensor? channel_scales," + " Tensor? token_scales," + " str? schedule" + ") -> Tensor"); + ops.def( + "machete_prepack_B(" + " Tensor B," + " ScalarType a_type," + " int b_type," + " ScalarType? group_scales_type" + ") -> Tensor"); + // conditionally compiled so impl registration is in source file + + // Marlin GEMM + ops.def( + "marlin_gemm(Tensor a, Tensor? c_or_none, Tensor b_q_weight, " + "Tensor? b_bias_or_none,Tensor b_scales, " + "Tensor? a_scales, Tensor? global_scale, Tensor? b_zeros_or_none, " + "Tensor? " + "g_idx_or_none, Tensor? perm_or_none, Tensor workspace, int b_type_id, " + "SymInt size_m, SymInt size_n, SymInt size_k, bool is_k_full, " + "bool use_atomic_add, bool use_fp32_reduce, bool is_zp_float) -> Tensor"); + // conditionally compiled so impl registrations are in source file + + // gptq_marlin repack from GPTQ. + ops.def( + "gptq_marlin_repack(Tensor b_q_weight, Tensor perm, " + "SymInt size_k, SymInt size_n, int num_bits, bool is_a_8bit) -> Tensor"); + // conditionally compiled so impl registrations are in source file + + // awq_marlin repack from AWQ. + ops.def( + "awq_marlin_repack(Tensor b_q_weight, SymInt size_k, " + "SymInt size_n, int num_bits, bool is_a_8bit) -> Tensor"); + // conditionally compiled so impl registrations are in source file + + // preprocess W-int4A-fp8 weight for marlin kernel + ops.def( + "marlin_int4_fp8_preprocess(Tensor qweight, " + "Tensor? qzeros_or_none, bool inplace) -> Tensor"); + // conditionally compiled so impl registrations are in source file + + // swordfish (sm100/sm110 w4a16): pack a GPTQ int4 weight into the + // Swordfish ABI v1 block-linear layout. + ops.def( + "swordfish_prepack_B(Tensor b_q_weight, Tensor? perm, SymInt size_k, " + "SymInt size_n, int num_bits) -> Tensor"); + // conditionally compiled so impl registrations are in source file + + // swordfish w4a16 decode GEMM over the ABI v1 packed weight. group_zps + // holds prescaled (8 - zp) * scale per group when present (AWQ/HQQ). + ops.def( + "swordfish_mm(Tensor a, Tensor b_packed, Tensor group_scales, " + "Tensor? group_zps, Tensor? perm, int num_bits, int group_size, " + "SymInt size_k, SymInt size_n) -> Tensor"); + // conditionally compiled so impl registrations are in source file + + // swordfish dequant to dense fp16/bf16, optionally out-major transposed + // and expert-stacked, for the dense tiers. + ops.def( + "swordfish_dequant_dense(Tensor b_packed, Tensor group_scales, " + "Tensor? group_zps, int num_bits, int group_size, SymInt size_k, " + "SymInt size_n, bool transpose) -> Tensor"); + // conditionally compiled so impl registrations are in source file + + // swordfish fused-MoE decode GEMM: one persistent Stream-K launch over + // per-expert ABI v1 weights, token-sorted by moe_align_block_size(16). + ops.def( + "swordfish_moe_mm(Tensor a, Tensor b_packed, Tensor group_scales, " + "Tensor sorted_token_ids, Tensor expert_ids, " + "Tensor num_tokens_post_padded, Tensor? topk_weights, " + "int moe_block_size, int top_k, bool mul_topk_weights, int num_bits, " + "int group_size, SymInt size_k, SymInt size_n) -> Tensor"); + // conditionally compiled so impl registrations are in source file + + // swordfish w4a16 prefill GEMM (sm100 tcgen05 mixed-input mainloop fork) + // over the same ABI v1 packed weight. + ops.def( + "swordfish_prefill_mm(Tensor a, Tensor b_packed, Tensor group_scales, " + "Tensor? group_zps, int num_bits, int group_size, SymInt size_k, " + "SymInt size_n) -> Tensor"); + // conditionally compiled so impl registrations are in source file +#endif + +#ifndef USE_ROCM + // CUTLASS w8a8 GEMM, supporting symmetric per-tensor or per-row/column + // quantization, as well as bias + ops.def( + "cutlass_scaled_mm(Tensor! out, Tensor a," + " Tensor b, Tensor a_scales," + " Tensor b_scales, Tensor? bias) -> ()"); + + // CUTLASS w8a8 GEMM, supporting asymmetric per-tensor or per-row/column + // quantization. + ops.def( + "cutlass_scaled_mm_azp(Tensor! out, Tensor a," + " Tensor b, Tensor a_scales," + " Tensor b_scales, Tensor azp_adj," + " Tensor? azp, Tensor? bias) -> ()"); + + // Check if cutlass scaled_mm is supported for CUDA devices of the given + // capability + ops.def("cutlass_scaled_mm_supports_fp8(int cuda_device_capability) -> bool"); + + // Check if cutlass grouped gemm is supported for CUDA devices of the given + // capability + ops.def("cutlass_group_gemm_supported(int cuda_device_capability) -> bool"); + + // CUTLASS w8a8 grouped GEMM + ops.def( + "cutlass_moe_mm(Tensor! out_tensors, Tensor a_tensors, Tensor b_tensors, " + " Tensor a_scales, Tensor b_scales, Tensor expert_offsets, " + " Tensor problem_sizes, Tensor a_strides, " + " Tensor b_strides, Tensor c_strides, bool per_act_token, " + " bool per_out_ch) -> ()"); + + // A function that computes data required to run fused MoE with w8a8 grouped + // GEMM. It takes topk_ids as an input, and computes expert_offsets + // (token start indices of each expert). In addition to this, it computes + // problem sizes for each expert's multiplication used by the two mms called + // from fused MoE operation, and arrays with permutations required to shuffle + // and de-shuffle the input/output of the fused operation. + ops.def( + "get_cutlass_moe_mm_data(Tensor topk_ids, Tensor! expert_offsets, " + " Tensor! problem_sizes1, Tensor! problem_sizes2, " + " Tensor! input_permutation, " + " Tensor! output_permutation, int num_experts, " + " int n, int k, Tensor? blockscale_offsets, " + " bool is_gated) -> ()"); + + // compute per-expert problem sizes from expert_first_token_offset + // produced by vLLM's moe_permute kernel + ops.def( + "get_cutlass_moe_mm_problem_sizes_from_expert_offsets(" + " Tensor expert_first_token_offset, " + " Tensor! problem_sizes1, " + " Tensor! problem_sizes2, " + " int n, int k, bool swap_ab) -> ()"); + + // A function that computes data required to run fused MoE with w8a8 grouped + // GEMM in batched expert format. It takes expert_num_tokens + // as an input, and computes expert_offsets (token start indices of each + // expert). In addition to this, it computes problem sizes for each expert's + // multiplication used by the two mms called from fused MoE operation. + ops.def( + "get_cutlass_batched_moe_mm_data(Tensor! expert_offsets, " + " Tensor! problem_sizes1, " + " Tensor! problem_sizes2, " + " Tensor expert_num_tokens, " + " int num_local_experts, int padded_m, " + " int n, int k) -> ()"); + + // Check if cutlass scaled_mm supports block quantization (used by DeepSeekV3) + ops.def( + "cutlass_scaled_mm_supports_block_fp8(int cuda_device_capability) -> " + "bool"); + + // CUTLASS nvfp4 block scaled GEMM + ops.def( + "cutlass_scaled_fp4_mm(Tensor! out, Tensor a, Tensor b," + " Tensor block_scale_a, Tensor block_scale_b," + " Tensor alpha) -> ()"); + + // cutlass nvfp4 block scaled group GEMM + ops.def( + "cutlass_fp4_group_mm(Tensor! out, Tensor a, Tensor b," + " Tensor a_blockscale, Tensor b_blockscales, Tensor alphas," + " Tensor problem_sizes, Tensor expert_offsets, Tensor sf_offsets) -> ()"); + + // cutlass mxfp4 block scaled group GEMM (MXFP4 x MXFP4 MoE) + ops.def( + "cutlass_mxfp4_group_mm(Tensor! out, Tensor a, Tensor b," + " Tensor a_blockscale, Tensor b_blockscales," + " Tensor problem_sizes, Tensor expert_offsets, Tensor sf_offsets) -> ()"); + + // Compute NVFP4 block quantized tensor. + ops.def( + "scaled_fp4_quant(Tensor input," + " Tensor input_scale, bool " + "is_sf_swizzled_layout) -> (Tensor, Tensor)"); + + // Out variant + // TODO: Add out_variant tag once PyTorch supports it (added in 2.11) + // This registration is now migrated to stable ABI + // at::Tag::out_variant is not available in the stable ABI (enum_tag.h is not + // yet in torch/headeronly), the tag should be applied from Python + // via torch.library.Library.define(..., tags=(torch.Tag.out_variant,)) + // with the .impl remaining in C++. + // See pytorch/pytorch#176117. + ops.def( + "scaled_fp4_quant.out(Tensor input," + " Tensor input_scale, bool " + "is_sf_swizzled_layout, *, Tensor(a!) output, Tensor(b!) output_scale) " + "-> ()"); + + // Compute NVFP4 experts quantization. + ops.def( + "scaled_fp4_experts_quant(Tensor! output, Tensor! output_scale," + "Tensor input, Tensor input_global_scale, Tensor input_offset_by_experts," + "Tensor output_scale_offset_by_experts) -> ()"); + + // Fused SiLU+Mul+NVFP4 experts quantization. + ops.def( + "silu_and_mul_scaled_fp4_experts_quant(Tensor! output, Tensor! " + "output_scale," + "Tensor input, Tensor input_global_scale, Tensor input_offset_by_experts," + "Tensor output_scale_offset_by_experts) -> ()"); + + // Compute MXFP4 experts quantization (32-element blocks, E8M0 SFs). + ops.def( + "mxfp4_experts_quant(Tensor! output, Tensor! output_scale," + "Tensor input, Tensor input_offset_by_experts," + "Tensor output_scale_offset_by_experts, int n_experts) -> ()"); + + // Fused SiLU+Mul+MXFP4 experts quantization. + ops.def( + "silu_and_mul_mxfp4_experts_quant(Tensor! output, Tensor! " + "output_scale," + "Tensor input, Tensor input_offset_by_experts," + "Tensor output_scale_offset_by_experts, int n_experts) -> ()"); + + // Fused SiLU+Mul+NVFP4 quantization. + ops.def( + "silu_and_mul_nvfp4_quant(Tensor! result, Tensor! result_block_scale, " + "Tensor input, Tensor input_global_scale) -> ()"); + + // Check if cutlass_scaled_mm_fp4 is supported for CUDA devices + // of the given capability + ops.def("cutlass_scaled_mm_supports_fp4(int cuda_device_capability) -> bool"); + + // CUTLASS w4a8 GEMM + ops.def( + "cutlass_w4a8_mm(" + " Tensor A," + " Tensor B," + " Tensor group_scales," + " int group_size," + " Tensor channel_scales," + " Tensor token_scales," + " ScalarType? out_type," + " str? maybe_schedule" + ") -> Tensor"); + + // pack scales + ops.def("cutlass_pack_scale_fp8(Tensor scales) -> Tensor"); + + // encode and reorder weight matrix + ops.def("cutlass_encode_and_reorder_int4b(Tensor B) -> Tensor"); + + // CUTLASS w4a8 grouped GEMM + ops.def( + "cutlass_w4a8_moe_mm(" + " Tensor! out_tensors," + " Tensor a_tensors," + " Tensor b_tensors," + " Tensor a_scales," + " Tensor b_scales," + " Tensor b_group_scales," + " int b_group_size," + " Tensor expert_offsets," + " Tensor problem_sizes," + " Tensor a_strides," + " Tensor b_strides," + " Tensor c_strides," + " Tensor group_scale_strides," + " str? maybe_schedule" + ") -> ()"); + + ops.def( + "cutlass_encode_and_reorder_int4b_grouped(Tensor b_tensors) -> (Tensor, " + "Tensor)"); + + // SM100 CUTLASS MLA decode + // conditionally compiled so impl registrations are in source file + ops.def( + "sm100_cutlass_mla_decode(Tensor! out, Tensor! lse, Tensor q_nope," + " Tensor q_pe, Tensor kv_c_and_k_pe_cache," + " Tensor seq_lens, Tensor page_table," + " Tensor workspace, float scale," + " int num_kv_splits) -> ()"); + + ops.def( + "sm100_cutlass_mla_get_workspace_size(int max_seq_len, int num_batches," + " int sm_count, int num_kv_splits) " + "-> int"); + // Quantized GEMM for AWQ. + ops.def( + "awq_gemm(Tensor _in_feats, Tensor _kernel, Tensor _scaling_factors, " + "Tensor _zeros, SymInt split_k_iters) -> Tensor"); + + // Dequantization for AWQ. + ops.def( + "awq_dequantize(Tensor _kernel, Tensor _scaling_factors, " + "Tensor _zeros, SymInt split_k_iters, int thx, int thy) -> Tensor"); + + // DeepSeek V3 fused A GEMM (SM 9.0+, bf16 only, 1-16 tokens). + // conditionally compiled so impl registration is in source file + ops.def( + "dsv3_fused_a_gemm(Tensor! output, Tensor mat_a, Tensor mat_b) -> ()"); + + // BF16/FP32 x FP32 -> FP32 router GEMM for H=3072, E=256, M<=32 (SM90+). + // conditionally compiled so impl registration is in source file + ops.def("fp32_router_gemm(Tensor! output, Tensor mat_a, Tensor mat_b) -> ()"); + + // reorder weight for AllSpark Ampere W8A16 Fused Gemm kernel + ops.def( + "rearrange_kn_weight_as_n32k16_order(Tensor b_qweight, Tensor b_scales, " + "Tensor? b_zeros, " + "bool has_zp, Tensor! b_qweight_reorder, Tensor! b_scales_reorder, " + "Tensor!? b_zeros_reorder, " + "int K, int N, int N_32align) -> ()"); + + // AllSpark quantization ops + ops.def( + "allspark_w8a16_gemm(Tensor a, Tensor b_qweight, Tensor b_scales, " + "Tensor? b_qzeros, " + "SymInt n, SymInt group_size, SymInt sm_count, SymInt sm_version, SymInt " + "CUBLAS_M_THRESHOLD, bool has_zp, bool n32k16_reorder) -> Tensor"); +#endif + + // Merge attn states + // Implements section 2.2 of https://www.arxiv.org/pdf/2501.01005 + // can be used to combine partial attention results (in the split-KV case) + ops.def( + "merge_attn_states(" + " Tensor! output," + " Tensor!? output_lse," + " Tensor prefix_output," + " Tensor prefix_lse," + " Tensor suffix_output," + " Tensor suffix_lse," + " int!? prefill_tokens_with_context," + " Tensor? output_scale=None) -> ()"); + + // Hadamard transforms + // conditionally compiled so impl registration is in source file + ops.def("hadacore_transform(Tensor! x, bool inplace) -> Tensor"); + + // Apply Root Mean Square (RMS) Normalization to the input tensor. + ops.def( + "rms_norm(Tensor! result, Tensor input, Tensor? weight, float epsilon) " + "-> " + "()"); + + // In-place fused Add and RMS Normalization. + ops.def( + "fused_add_rms_norm(Tensor! input, Tensor! residual, Tensor? weight, " + "float epsilon) -> ()"); + + // Layernorm-quant + // Apply Root Mean Square (RMS) Normalization to the input tensor. + ops.def( + "rms_norm_static_fp8_quant(Tensor! result, Tensor input, Tensor weight, " + "Tensor scale, float epsilon) -> " + "()"); + + // In-place fused Add and RMS Normalization. + ops.def( + "fused_add_rms_norm_static_fp8_quant(Tensor! result, Tensor input, " + "Tensor! residual, Tensor weight, " + "Tensor scale, float epsilon) -> ()"); + + // Fused Layernorm + Quant kernels + ops.def( + "rms_norm_dynamic_per_token_quant(Tensor! result, Tensor input, " + "Tensor weight, Tensor! scale, float epsilon, " + "Tensor? scale_ub, Tensor!? residual) -> ()"); + + // Fused Layernorm + Block quant kernels + ops.def( + "rms_norm_per_block_quant(Tensor! result, Tensor input, " + "Tensor weight, Tensor! scale, float epsilon, " + "Tensor? scale_ub, Tensor!? residual, int group_size, " + "bool is_scale_transposed) -> ()"); + + // Fused SiLU+Mul + per-block quantization + ops.def( + "silu_and_mul_per_block_quant(" + "Tensor! out, " + "Tensor input, " + "Tensor! scales, " + "int group_size, " + "Tensor? scale_ub=None, " + "bool is_scale_transposed=False) -> ()"); + + // Rotary embedding + // Apply GPT-NeoX or GPT-J style rotary embedding to query and key. + ops.def( + "rotary_embedding(Tensor positions, Tensor! query," + " Tensor!? key, int head_size," + " Tensor cos_sin_cache, bool is_neox, int " + "rope_dim_offset=0, bool inverse=False) -> ()"); + + // Function for fused QK Norm and RoPE + ops.def( + "fused_qk_norm_rope(Tensor! qkv, int num_heads_q, " + "int num_heads_k, int num_heads_v, int head_dim, float eps, " + "Tensor q_weight, Tensor k_weight, Tensor cos_sin_cache, " + "bool is_neox, Tensor position_ids, " + "int forced_token_heads_per_warp=-1) -> ()"); + + ops.def( + "fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(" + "Tensor q_in, Tensor kv, Tensor! k_cache, " + "Tensor slot_mapping, Tensor position_ids, Tensor cos_sin_cache, " + "int q_head_padded, float eps, int cache_block_size) -> Tensor"); + + // FlashInfer V4 full-cache variants: write Q in place (bf16) or to a separate + // FP8 tensor, and KV into a contiguous 512-wide token-strided cache. + ops.def( + "fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_bf16_insert(" + "Tensor! q, Tensor kv, Tensor! k_cache, Tensor slot_mapping, " + "Tensor position_ids, Tensor cos_sin_cache, float eps, " + "int cache_block_size) -> ()"); + ops.def( + "fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_fp8_insert(" + "Tensor q, Tensor kv, Tensor! q_fp8, Tensor! k_cache, " + "Tensor slot_mapping, Tensor position_ids, Tensor cos_sin_cache, " + "Tensor fp8_scale, Tensor q_fp8_scale_inv, float eps, " + "int cache_block_size) -> ()"); + +#ifndef USE_ROCM + ops.def( + "minimax_allreduce_rms_qk(" + "Tensor qkv, Tensor norm_weight_q, Tensor norm_weight_k, " + "Tensor workspace, int q_size, int kv_size, int rank, int nranks, " + "float eps) -> (Tensor, Tensor)"); +#endif + + // Horizontally-fused MiniMax-M3 QK-norm + partial NeoX RoPE + KV-insert. + ops.def( + "fused_minimax_m3_qknorm_rope_kv_insert(" + "Tensor! qkv, Tensor q_norm_weight, Tensor k_norm_weight, " + "Tensor cos_sin_cache, Tensor positions, int num_heads, " + "int num_kv_heads, int rotary_dim, float eps, " + "Tensor? index_q_norm_weight, Tensor? index_k_norm_weight, " + "int num_index_heads, " + "Tensor? slot_mapping, Tensor? index_slot_mapping, " + "Tensor!? kv_cache, Tensor!? index_cache, " + "int block_size, Tensor!? q_out, Tensor!? index_q_out, " + "str kv_cache_dtype, bool skip_index_branch=False) -> ()"); + +#ifdef APHRODITE_ENABLE_KIMI_K3_ATTN_RES + ops.def( + "kimi_k3_attn_res(" + "Tensor! prefix, Tensor delta, Tensor blocks, Tensor norm_weight, " + "Tensor qk_weight, Tensor output_norm_weight, Tensor! output, " + "int num_blocks, float eps, float output_norm_eps) -> ()"); +#endif + + // Apply repetition penalties to logits in-place. + ops.def( + "apply_repetition_penalties_(Tensor! logits, Tensor prompt_mask, " + "Tensor output_mask, Tensor repetition_penalties) -> ()"); + + // Optimized top-k per row operations. + ops.def( + "top_k_per_row_prefill(Tensor logits, Tensor rowStarts, Tensor rowEnds, " + "Tensor! indices, int numRows, int stride0, " + "int stride1, int topK) -> ()"); + + ops.def( + "top_k_per_row_decode(Tensor logits, int next_n, " + "Tensor seq_lens, Tensor! indices, " + "int numRows, int stride0, int stride1, int topK) -> ()"); + + ops.def( + "persistent_topk(Tensor logits, Tensor lengths, Tensor! output, " + "Tensor workspace, int k, int max_seq_len) -> ()"); + +#ifdef APHRODITE_ENABLE_COOPERATIVE_TOPK + ops.def( + "cooperative_topk(Tensor logits, Tensor lengths, Tensor! output, " + "Tensor workspace, int k, int max_seq_len) -> ()"); +#endif + +#ifdef APHRODITE_ENABLE_SM89_DSA + // sm89 DeepSeek sparse attention; impls are registered in the kernel sources + ops.def( + "sm89_fp8_paged_mqa_logits(Tensor q, Tensor pool, Tensor weights, " + "Tensor seq_lens, Tensor block_table, Tensor sched, Tensor! logits, " + "bool clean_logits) -> ()"); + + ops.def( + "sm89_paged_mqa_logits_metadata(Tensor seq_lens, Tensor! sched, " + "int next_n) -> ()"); + + ops.def( + "sm89_fp8_mqa_logits(Tensor q, Tensor kv, Tensor kv_scales, " + "Tensor weights, Tensor cu_seqlen_ks, Tensor cu_seqlen_ke, " + "Tensor! logits) -> ()"); + + ops.def( + "sm89_sparse_mla_fwd(Tensor q, Tensor pool, Tensor indices, " + "Tensor! out, Tensor! lse, float sm_scale, " + "Tensor? topk_lens=None) -> ()"); +#endif + + // Activation ops + ops.def( + "persistent_masked_m_silu_mul_quant(Tensor input, Tensor counts, Tensor! " + "y_q, Tensor! y_s, bool use_ue8m0) -> ()"); + ops.def("weak_ref_tensor(Tensor input) -> Tensor"); + + // Activation function used in SwiGLU. + ops.def("silu_and_mul(Tensor! result, Tensor input) -> ()"); + + ops.def("mul_and_silu(Tensor! out, Tensor input) -> ()"); + + // SwiGLU activation with input clamping. + // alpha scales the sigmoid (gate * sigmoid(alpha * gate)); beta is added to + // the up half (up + beta). Defaults alpha=1.0, beta=0.0 give silu(gate)*up. + ops.def( + "silu_and_mul_with_clamp(Tensor! result, Tensor input, float limit, " + "float alpha=1.0, float beta=0.0) -> ()"); + + // SwiGLU activation with FP8 quantization. + ops.def( + "silu_and_mul_quant(Tensor! result, Tensor input, Tensor scale) -> ()"); + + // Activation function used in GeGLU with `none` approximation. + ops.def("gelu_and_mul(Tensor! out, Tensor input) -> ()"); + + // Activation function used in GeGLU with `tanh` approximation. + ops.def("gelu_tanh_and_mul(Tensor! out, Tensor input) -> ()"); + + // FATReLU implementation. + ops.def("fatrelu_and_mul(Tensor! out, Tensor input, float threshold) -> ()"); + + ops.def( + "swigluoai_and_mul(Tensor! out, Tensor input, float alpha=1.702, float " + "limit=7.0) " + "-> ()"); + + // GELU implementation used in GPT-2. + ops.def("gelu_new(Tensor! out, Tensor input) -> ()"); + + // Approximate GELU implementation. + ops.def("gelu_fast(Tensor! out, Tensor input) -> ()"); + + // Quick GELU implementation. + ops.def("gelu_quick(Tensor! out, Tensor input) -> ()"); + + // relu(x)^2 activation from https://arxiv.org/abs/2109.08668v2 + ops.def("relu_squared(Tensor! out, Tensor input) -> ()"); + + // Compute int8 quantized tensor for given scaling factor. + ops.def( + "static_scaled_int8_quant(Tensor! result, Tensor input, Tensor scale," + "Tensor? azp) -> ()"); + + // Compute int8 quantized tensor and scaling factor + ops.def( + "dynamic_scaled_int8_quant(Tensor! result, Tensor input, Tensor! scale, " + "Tensor!? azp) -> ()"); + + // Compute FP8 quantized tensor for given scaling factor. + // Supports per-tensor, per-channel, per-token, and arbitrary 2D group + // scaling. Optional group_m/group_n specify the group shape explicitly; + // required for 1D scales to disambiguate per-channel vs per-token. + ops.def( + "static_scaled_fp8_quant(Tensor! result, Tensor input, Tensor scale, " + "int[]? group_shape=None) -> ()"); + + // Compute dynamic-per-tensor FP8 quantized tensor and scaling factor. + ops.def( + "dynamic_scaled_fp8_quant(Tensor! result, Tensor input, Tensor! scale) " + "-> " + "()"); + + // Compute dynamic-per-token FP8 quantized tensor and scaling factor. + ops.def( + "dynamic_per_token_scaled_fp8_quant(Tensor! result, Tensor input, " + "Tensor! scale, Tensor? scale_ub) -> " + "()"); + + // Quantized GEMM for GPTQ. + // Note: even though the C++ inferred schema is correct for this op, it seems + // to prevent the meta function registry. + ops.def( + "gptq_gemm(Tensor a, Tensor b_q_weight, Tensor b_gptq_qzeros, " + "Tensor b_gptq_scales, Tensor b_g_idx, bool use_exllama, bool " + "use_v2_format, int bit) " + "-> Tensor"); + + // Post processing for GPTQ. + ops.def("gptq_shuffle(Tensor! q_weight, Tensor q_perm, int bit) -> ()"); + + // Mamba selective scan kernel + ops.def( + "selective_scan_fwd(Tensor! u, Tensor! delta," + "Tensor! A, Tensor! B, Tensor! C," + "Tensor? D_, Tensor!? z_, Tensor? delta_bias_," + "bool delta_softplus," + "Tensor? query_start_loc," + "Tensor? cache_indices," + "Tensor? has_initial_state," + "Tensor! ssm_states," + "int null_block_id," + "int block_size," + "Tensor? block_idx_first_scheduled_token," + "Tensor? block_idx_last_scheduled_token," + "Tensor? initial_state_idx," + "Tensor? cu_chunk_seqlen," + "Tensor? last_chunk_indices) -> ()"); + + // LongCat n-gram embedding index kernel. All tensor args are marked mutable + // to match the (non-const) stable-Tensor& C++ signature; only ne_token_table + // and n_gram_ids are actually written in place. + ops.def( + "ngram_compute_n_gram_ids(int ne_n, int ne_k, Tensor(a!) ne_weights, " + "Tensor(b!) ne_mods, Tensor(c!) exclusive_ne_embedder_size_sums, " + "Tensor(d!) exclusive_req_len_sums, Tensor(e!) ne_token_table, " + "Tensor(f!) row_indices, Tensor(g!) column_starts, " + "Tensor(h!) n_gram_ids) -> ()"); +} + +STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) { + // LongCat n-gram embedding index kernel. + ops.impl("ngram_compute_n_gram_ids", TORCH_BOX(&ngram_compute_n_gram_ids)); + + // Per-token group quantization + ops.impl("per_token_group_fp8_quant", TORCH_BOX(&per_token_group_quant_fp8)); + ops.impl("per_token_group_fp8_quant_packed", + TORCH_BOX(&per_token_group_quant_8bit_packed)); + ops.impl("per_token_group_quant_int8", + TORCH_BOX(&per_token_group_quant_int8)); + + ops.impl("permute_cols", TORCH_BOX(&permute_cols)); + +#ifndef USE_ROCM + // CUTLASS scaled_mm ops + ops.impl("cutlass_scaled_mm", TORCH_BOX(&cutlass_scaled_mm)); + ops.impl("cutlass_scaled_mm_azp", TORCH_BOX(&cutlass_scaled_mm_azp)); + ops.impl("cutlass_moe_mm", TORCH_BOX(&cutlass_moe_mm)); + ops.impl("get_cutlass_moe_mm_data", TORCH_BOX(&get_cutlass_moe_mm_data)); + ops.impl("get_cutlass_moe_mm_problem_sizes_from_expert_offsets", + TORCH_BOX(&get_cutlass_moe_mm_problem_sizes_from_expert_offsets)); + ops.impl("get_cutlass_batched_moe_mm_data", + TORCH_BOX(&get_cutlass_batched_moe_mm_data)); + + // FP4/NVFP4 ops + ops.impl("cutlass_scaled_fp4_mm", TORCH_BOX(&cutlass_scaled_fp4_mm)); + ops.impl("scaled_fp4_quant", TORCH_BOX(&scaled_fp4_quant_func)); + ops.impl("scaled_fp4_quant.out", TORCH_BOX(&scaled_fp4_quant_out)); + ops.impl("scaled_fp4_experts_quant", TORCH_BOX(&scaled_fp4_experts_quant)); + ops.impl("silu_and_mul_scaled_fp4_experts_quant", + TORCH_BOX(&silu_and_mul_scaled_fp4_experts_quant)); + ops.impl("silu_and_mul_nvfp4_quant", TORCH_BOX(&silu_and_mul_nvfp4_quant)); + // mxfp4_experts_quant: registered in mxfp4_experts_quant.cu (SM100 only). + // W4A8 ops: registered in w4a8_mm_entry.cu / w4a8_grouped_mm_entry.cu. + + // AWQ ops + ops.impl("awq_gemm", TORCH_BOX(&awq_gemm)); + ops.impl("awq_dequantize", TORCH_BOX(&awq_dequantize)); + + // DSV3 fused A GEMM: conditionally compiled so impl registration is in + // source file (dsv3_fused_a_gemm.cu) + + // AllSpark ops: conditionally compiled so impl registrations are in source + // files (allspark_repack.cu and allspark_qgemm_w8a16.cu) +#endif + + ops.impl("merge_attn_states", TORCH_BOX(&merge_attn_states)); + + // Layernorm kernels (shared CUDA/ROCm) + ops.impl("rms_norm", TORCH_BOX(&rms_norm)); + ops.impl("fused_add_rms_norm", TORCH_BOX(&fused_add_rms_norm)); + + // Layernorm-quant kernels (shared CUDA/ROCm) + ops.impl("rms_norm_static_fp8_quant", TORCH_BOX(&rms_norm_static_fp8_quant)); + ops.impl("fused_add_rms_norm_static_fp8_quant", + TORCH_BOX(&fused_add_rms_norm_static_fp8_quant)); + + // Fused layernorm + dynamic per-token quant kernels (shared CUDA/ROCm) + ops.impl("rms_norm_dynamic_per_token_quant", + TORCH_BOX(&rms_norm_dynamic_per_token_quant)); + ops.impl("rms_norm_per_block_quant", TORCH_BOX(&rms_norm_per_block_quant)); + ops.impl("silu_and_mul_per_block_quant", + TORCH_BOX(&silu_and_mul_per_block_quant)); + + // Positional encoding kernels (shared CUDA/ROCm) + ops.impl("rotary_embedding", TORCH_BOX(&rotary_embedding)); + ops.impl("fused_qk_norm_rope", TORCH_BOX(&fused_qk_norm_rope)); + ops.impl("fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert", + TORCH_BOX(&fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert)); + ops.impl( + "fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_bf16_insert", + TORCH_BOX(&fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_bf16_insert)); + ops.impl( + "fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_fp8_insert", + TORCH_BOX(&fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_fp8_insert)); +#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_KIMI_K3_ATTN_RES + ops.impl("kimi_k3_attn_res", TORCH_BOX(&kimi_k3_attn_res)); +#endif + + // Sampler kernels (shared CUDA/ROCm) + ops.impl("apply_repetition_penalties_", + TORCH_BOX(&apply_repetition_penalties_)); + ops.impl("top_k_per_row_prefill", TORCH_BOX(&top_k_per_row_prefill)); + ops.impl("top_k_per_row_decode", TORCH_BOX(&top_k_per_row_decode)); + ops.impl("persistent_topk", TORCH_BOX(&persistent_topk)); +#ifdef APHRODITE_ENABLE_COOPERATIVE_TOPK + ops.impl("cooperative_topk", TORCH_BOX(&cooperative_topk)); +#endif + + // Activation kernels (shared CUDA/ROCm) + ops.impl("persistent_masked_m_silu_mul_quant", + TORCH_BOX(&persistent_masked_m_silu_mul_quant)); + ops.impl("weak_ref_tensor", TORCH_BOX(&weak_ref_tensor)); + ops.impl("silu_and_mul_quant", TORCH_BOX(&silu_and_mul_quant)); + ops.impl("silu_and_mul", TORCH_BOX(&silu_and_mul)); + ops.impl("mul_and_silu", TORCH_BOX(&mul_and_silu)); + ops.impl("gelu_and_mul", TORCH_BOX(&gelu_and_mul)); + ops.impl("gelu_tanh_and_mul", TORCH_BOX(&gelu_tanh_and_mul)); + ops.impl("fatrelu_and_mul", TORCH_BOX(&fatrelu_and_mul)); + ops.impl("swigluoai_and_mul", TORCH_BOX(&swigluoai_and_mul)); + ops.impl("gelu_new", TORCH_BOX(&gelu_new)); + ops.impl("gelu_fast", TORCH_BOX(&gelu_fast)); + ops.impl("gelu_quick", TORCH_BOX(&gelu_quick)); + ops.impl("relu_squared", TORCH_BOX(&relu_squared)); + ops.impl("silu_and_mul_with_clamp", TORCH_BOX(&silu_and_mul_clamp)); + + // INT8 quantization kernels + ops.impl("static_scaled_int8_quant", TORCH_BOX(&static_scaled_int8_quant)); + ops.impl("dynamic_scaled_int8_quant", TORCH_BOX(&dynamic_scaled_int8_quant)); + + // FP8 quantization kernels + ops.impl("static_scaled_fp8_quant", TORCH_BOX(&static_scaled_fp8_quant)); + ops.impl("dynamic_scaled_fp8_quant", TORCH_BOX(&dynamic_scaled_fp8_quant)); + ops.impl("dynamic_per_token_scaled_fp8_quant", + TORCH_BOX(&dynamic_per_token_scaled_fp8_quant)); + + // GPTQ kernels + ops.impl("gptq_gemm", TORCH_BOX(&gptq_gemm)); + ops.impl("gptq_shuffle", TORCH_BOX(&gptq_shuffle)); + + // Mamba kernels + ops.impl("selective_scan_fwd", TORCH_BOX(&selective_scan_fwd)); +} + +STABLE_TORCH_LIBRARY_IMPL(_C, CPU, ops) { + ops.impl("get_cuda_view_from_cpu_tensor", + TORCH_BOX(&get_cuda_view_from_cpu_tensor)); + ops.impl("dry_scan_penalties", TORCH_BOX(&dry_scan_penalties_cpu)); +} + +STABLE_TORCH_LIBRARY_FRAGMENT(_C_cuda_utils, cuda_utils) { + cuda_utils.def("get_device_attribute(int attribute, int device_id) -> int"); + cuda_utils.def( + "get_max_shared_memory_per_block_device_attribute(int device_id) -> int"); +} + +STABLE_TORCH_LIBRARY_IMPL(_C_cuda_utils, CompositeExplicitAutograd, + cuda_utils) { + cuda_utils.impl("get_device_attribute", TORCH_BOX(&get_device_attribute)); + cuda_utils.impl("get_max_shared_memory_per_block_device_attribute", + TORCH_BOX(&get_max_shared_memory_per_block_device_attribute)); +} + +// These capability-check functions take only primitive args (no tensors), so +// there is no device to dispatch on. CompositeExplicitAutograd makes them +// available for all backends. This is the stable ABI equivalent of calling +// ops.impl("op_name", &func) without a dispatch key in the non-stable API. +STABLE_TORCH_LIBRARY_IMPL(_C, CompositeExplicitAutograd, ops) { +#ifndef USE_ROCM + ops.impl("cutlass_scaled_mm_supports_fp8", + TORCH_BOX(&cutlass_scaled_mm_supports_fp8)); + ops.impl("cutlass_group_gemm_supported", + TORCH_BOX(&cutlass_group_gemm_supported)); + ops.impl("cutlass_scaled_mm_supports_block_fp8", + TORCH_BOX(&cutlass_scaled_mm_supports_block_fp8)); + ops.impl("cutlass_scaled_mm_supports_fp4", + TORCH_BOX(&cutlass_scaled_mm_supports_fp4)); +#endif +} + +// Cache ops +STABLE_TORCH_LIBRARY_FRAGMENT(_C_cache_ops, ops) { + // Swap in (out) the cache blocks from src to dst. + ops.def( + "swap_blocks(Tensor src, Tensor! dst," + " int block_size_in_bytes, Tensor block_mapping) -> ()"); + + // Batch swap: submit all block copies in a single driver call. + ops.def( + "swap_blocks_batch(Tensor src_ptrs, Tensor dst_ptrs," + " Tensor sizes," + " bool is_src_access_order_any=False) -> ()"); + + // Reshape the key and value tensors and cache them. + ops.def( + "reshape_and_cache(Tensor key, Tensor value," + " Tensor! key_cache, Tensor! value_cache," + " Tensor slot_mapping," + " str kv_cache_dtype," + " Tensor k_scale, Tensor v_scale) -> ()"); + + // Reshape the key and value tensors and cache them. + ops.def( + "reshape_and_cache_flash(Tensor key, Tensor value," + " Tensor! key_cache," + " Tensor! value_cache," + " Tensor slot_mapping," + " str kv_cache_dtype," + " Tensor k_scale, Tensor v_scale) -> ()"); + + // Concat kv_c and k_pe and cache them. + ops.def( + "concat_and_cache_mla(Tensor kv_c, Tensor k_pe," + " Tensor! kv_cache," + " Tensor slot_mapping," + " str kv_cache_dtype," + " Tensor scale) -> ()"); + + // Rotate Q and K, then write to kv cache for MLA + ops.def( + "concat_and_cache_mla_rope_fused(" + " Tensor positions," + " Tensor! q_pe," + " Tensor! k_pe," + " Tensor kv_c," + " Tensor cos_sin_cache," + " bool is_neox," + " Tensor slot_mapping," + " Tensor! kv_cache," + " str kv_cache_dtype," + " Tensor kv_cache_scale) -> ()"); + + // Convert the key and value cache to fp8 data type. + ops.def( + "convert_fp8(Tensor! dst_cache, Tensor src_cache, float scale, " + "str kv_cache_dtype) -> ()"); + + // Gather cache blocks from src_cache to dst, dequantizing from + // src_cache's dtype to dst's dtype if necessary. + ops.def( + "gather_and_maybe_dequant_cache(Tensor src_cache, Tensor! dst, " + " Tensor block_table, Tensor cu_seq_lens, " + " Tensor token_to_seq, " + " int num_tokens, " + " str kv_cache_dtype, " + " Tensor scale, Tensor? seq_starts) -> ()"); + + ops.def( + "cp_gather_cache(Tensor src_cache, Tensor! dst, Tensor block_table, " + "Tensor cu_seq_lens, int batch_size, Tensor? seq_starts) -> ()"); + + ops.def( + "cp_gather_and_upconvert_fp8_kv_cache(Tensor src_cache, Tensor! dst, " + "Tensor block_table, Tensor workspace_starts, int batch_size, Tensor? " + "seq_starts) -> ()"); + + ops.def( + "indexer_k_quant_and_cache(Tensor k, Tensor! kv_cache, Tensor " + "slot_mapping, " + "int quant_block_size, str kv_cache_dtype) -> ()"); + + ops.def("concat_mla_q(Tensor ql_nope, Tensor q_pe, Tensor! q_out) -> ()"); + + ops.def( + "cp_gather_indexer_k_quant_cache(Tensor kv_cache, Tensor! dst_k, Tensor! " + "dst_scale, Tensor block_table, Tensor cu_seq_lens) -> ()"); +} + +STABLE_TORCH_LIBRARY_FRAGMENT(_C_custom_ar, custom_ar) { + custom_ar.def( + "init_custom_ar(int[] ipc_tensors, Tensor rank_data, " + "int rank, bool fully_connected) -> int"); + custom_ar.def( + "all_reduce(int fa, Tensor inp, Tensor! out, int reg_buffer, " + "int reg_buffer_sz_bytes) -> ()"); + custom_ar.def("dispose(int fa) -> ()"); + custom_ar.def("meta_size() -> int"); + custom_ar.def("register_buffer(int fa, int[] ipc_tensors) -> ()"); + custom_ar.def("get_graph_buffer_ipc_meta(int fa) -> (int[], int[])"); + custom_ar.def( + "register_graph_buffers(int fa, int[][] handles, int[][] offsets) -> ()"); + custom_ar.def("allocate_shared_buffer_and_handle(int size) -> (int, Tensor)"); + custom_ar.def("open_mem_handle(Tensor mem_handle) -> int"); + custom_ar.def("free_shared_buffer(int ptr) -> ()"); +} + +STABLE_TORCH_LIBRARY_IMPL(_C_custom_ar, CUDA, custom_ar) { + custom_ar.impl("init_custom_ar", TORCH_BOX(&init_custom_ar)); + custom_ar.impl("all_reduce", TORCH_BOX(&all_reduce)); +} + +STABLE_TORCH_LIBRARY_IMPL(_C_custom_ar, CPU, custom_ar) { + custom_ar.impl("open_mem_handle", TORCH_BOX(&open_mem_handle)); +} + +STABLE_TORCH_LIBRARY_IMPL(_C_custom_ar, CompositeExplicitAutograd, custom_ar) { + custom_ar.impl("dispose", TORCH_BOX(&dispose)); + custom_ar.impl("meta_size", TORCH_BOX(&meta_size)); + custom_ar.impl("register_buffer", TORCH_BOX(®ister_buffer)); + custom_ar.impl("get_graph_buffer_ipc_meta", + TORCH_BOX(&get_graph_buffer_ipc_meta)); + custom_ar.impl("register_graph_buffers", TORCH_BOX(®ister_graph_buffers)); + custom_ar.impl("allocate_shared_buffer_and_handle", + TORCH_BOX(&allocate_shared_buffer_and_handle)); + custom_ar.impl("free_shared_buffer", TORCH_BOX(&free_shared_buffer)); +} + +STABLE_TORCH_LIBRARY_IMPL(_C_cache_ops, CPU, ops) { + ops.impl("swap_blocks_batch", TORCH_BOX(&swap_blocks_batch)); +} + +STABLE_TORCH_LIBRARY_IMPL(_C_cache_ops, CUDA, ops) { + ops.impl("swap_blocks", TORCH_BOX(&swap_blocks)); + ops.impl("reshape_and_cache", TORCH_BOX(&reshape_and_cache)); + ops.impl("reshape_and_cache_flash", TORCH_BOX(&reshape_and_cache_flash)); + ops.impl("concat_and_cache_mla", TORCH_BOX(&concat_and_cache_mla)); + ops.impl("concat_and_cache_mla_rope_fused", + TORCH_BOX(&concat_and_cache_mla_rope_fused)); + ops.impl("convert_fp8", TORCH_BOX(&convert_fp8)); + ops.impl("gather_and_maybe_dequant_cache", + TORCH_BOX(&gather_and_maybe_dequant_cache)); + ops.impl("cp_gather_cache", TORCH_BOX(&cp_gather_cache)); + ops.impl("cp_gather_and_upconvert_fp8_kv_cache", + TORCH_BOX(&cp_gather_and_upconvert_fp8_kv_cache)); + ops.impl("indexer_k_quant_and_cache", TORCH_BOX(&indexer_k_quant_and_cache)); + ops.impl("concat_mla_q", TORCH_BOX(&concat_mla_q)); + ops.impl("cp_gather_indexer_k_quant_cache", + TORCH_BOX(&cp_gather_indexer_k_quant_cache)); +} + +REGISTER_EXTENSION(_C_stable_libtorch) diff --git a/csrc/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/setup.py b/setup.py index 1ca034e686..aaa45c3652 100644 --- a/setup.py +++ b/setup.py @@ -641,6 +641,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/setup.py.orig b/setup.py.orig new file mode 100644 index 0000000000..1ca034e686 --- /dev/null +++ b/setup.py.orig @@ -0,0 +1,796 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import ctypes +import importlib.util +import logging +import os +import re +import shutil +import subprocess +import sys +import sysconfig +from pathlib import Path +from shutil import which + +import torch +from packaging.version import Version, parse +from setuptools import Extension, setup +from setuptools.command.build_ext import build_ext +from setuptools_scm import get_version +from torch.utils.cpp_extension import CUDA_HOME, ROCM_HOME + + +def load_module_from_path(module_name, path): + spec = importlib.util.spec_from_file_location(module_name, path) + module = importlib.util.module_from_spec(spec) + sys.modules[module_name] = module + spec.loader.exec_module(module) + return module + + +ROOT_DIR = Path(__file__).parent +logger = logging.getLogger(__name__) + +PRECOMPILED_RUST_FRONTEND_PATH = ROOT_DIR / "aphrodite" / "aphrodite-rs" +# cannot import envs directly because it depends on aphrodite, +# which is not installed yet +envs = load_module_from_path("envs", os.path.join(ROOT_DIR, "aphrodite", "envs.py")) + +APHRODITE_TARGET_DEVICE = envs.APHRODITE_TARGET_DEVICE + +# Skips all native builds (CMake/CUDA and Rust) and extracts binaries from the +# matching per-commit wheel. +APHRODITE_USE_PRECOMPILED = envs.APHRODITE_USE_PRECOMPILED +precompiled_wheel = load_module_from_path( + "precompiled_wheel", + os.path.join(ROOT_DIR, "aphrodite", "build_utils", "precompiled.py"), +) + + +def should_require_rust_frontend() -> bool: + value = os.getenv("APHRODITE_REQUIRE_RUST_FRONTEND", "") + return value.lower() not in ("", "0", "false", "no") + + +def should_use_precompiled_rust() -> bool: + value = os.getenv("APHRODITE_USE_PRECOMPILED_RUST", "") + return value.lower() in ("1", "true", "yes") + + +def get_precompiled_rust_extension_paths() -> list[Path]: + return sorted((ROOT_DIR / "aphrodite").glob("_rust_*.so")) + + +if sys.platform.startswith("darwin") and APHRODITE_TARGET_DEVICE not in ("cpu", "metal"): + logger.warning("APHRODITE_TARGET_DEVICE automatically set to `metal` due to macOS") + APHRODITE_TARGET_DEVICE = "metal" +elif not (sys.platform.startswith("linux") or sys.platform.startswith("darwin")): + logger.warning( + "Aphrodite only supports Linux platform (including WSL) and MacOS." + "Building on %s, " + "so Aphrodite may not be able to run correctly", + sys.platform, + ) + APHRODITE_TARGET_DEVICE = "empty" +elif sys.platform.startswith("linux") and os.getenv("APHRODITE_TARGET_DEVICE") is None: + if torch.version.hip is not None: + APHRODITE_TARGET_DEVICE = "rocm" + logger.info("Auto-detected ROCm") + elif torch.version.xpu is not None: + APHRODITE_TARGET_DEVICE = "xpu" + logger.info("Auto-detected XPU") + elif torch.version.cuda is not None: + APHRODITE_TARGET_DEVICE = "cuda" + logger.info("Auto-detected CUDA") + else: + APHRODITE_TARGET_DEVICE = "cpu" + + +def is_sccache_available() -> bool: + return which("sccache") is not None and not bool(int(os.getenv("APHRODITE_DISABLE_SCCACHE", "0"))) + + +def is_ccache_available() -> bool: + return which("ccache") is not None + + +def is_ninja_available() -> bool: + return which("ninja") is not None + + +def is_freethreaded(): + return bool(sysconfig.get_config_var("Py_GIL_DISABLED")) + + +def should_bundle_tcmalloc() -> bool: + import platform + + return ( + APHRODITE_TARGET_DEVICE == "cpu" + and sys.platform.startswith("linux") + and platform.machine() in ("aarch64", "x86_64") + ) + + +def find_tcmalloc() -> Path | None: + try: + # get all shared libs the dynamic loader knows about + output = subprocess.check_output( + ["ldconfig", "-p"], + text=True, + stderr=subprocess.DEVNULL, + ) + except Exception: + return None + + # search for libtcmalloc and libtcmalloc_minimal + for library_pattern in ( + r"\blibtcmalloc_minimal\.so\.(\d+)\b", + r"\blibtcmalloc\.so\.(\d+)\b", + ): + candidates: list[tuple[int, Path]] = [] + for line in output.splitlines(): + match = re.search(library_pattern, line) + if match is None or "=>" not in line: + continue + candidate = Path(line.split("=>")[1].strip()) + if candidate.exists(): + candidates.append((int(match.group(1)), candidate)) + + if candidates: + # if multiple candidates are found, pick the one with the highest + # version number + return max(candidates, key=lambda item: item[0])[1] + + return None + + +def bundle_tcmalloc(build_lib: str) -> None: + tcmalloc_library = find_tcmalloc() + if tcmalloc_library is None: + logger.warning( + "Failed to locate tcmalloc. For best performance, " + "please install tcmalloc (e.g. `sudo apt-get " + "install -y --no-install-recommends libtcmalloc-minimal4`)" + ) + return + + bundle_dir = os.path.join(build_lib, "aphrodite", "libs") + os.makedirs(bundle_dir, exist_ok=True) + bundle_path = os.path.join(bundle_dir, tcmalloc_library.name) + shutil.copy2(tcmalloc_library, bundle_path) + logger.info("Bundled tcmalloc into wheel: %s", bundle_path) + + +class CMakeExtension(Extension): + def __init__(self, name: str, cmake_lists_dir: str = ".", **kwa) -> None: + # Default to the stable/abi3 API, but let callers opt out (e.g. the + # nanobind _paged_ops Metal kernel builds a version-specific module). + kwa.setdefault("py_limited_api", not is_freethreaded()) + super().__init__(name, sources=[], **kwa) + self.cmake_lists_dir = os.path.abspath(cmake_lists_dir) + + +class cmake_build_ext(build_ext): + # A dict of extension directories that have been configured. + did_config: dict[str, bool] = {} + + # + # Determine number of compilation jobs and optionally nvcc compile threads. + # + def compute_num_jobs(self): + # `num_jobs` is either the value of the MAX_JOBS environment variable + # (if defined) or the number of CPUs available. + num_jobs = envs.MAX_JOBS + if num_jobs is not None: + num_jobs = int(num_jobs) + logger.info("Using MAX_JOBS=%d as the number of jobs.", num_jobs) + else: + try: + # os.sched_getaffinity() isn't universally available, so fall + # back to os.cpu_count() if we get an error here. + num_jobs = len(os.sched_getaffinity(0)) + except AttributeError: + num_jobs = os.cpu_count() + + nvcc_threads = None + if _is_cuda() and CUDA_HOME is not None: + try: + nvcc_version = get_nvcc_cuda_version() + if nvcc_version >= Version("11.2"): + # `nvcc_threads` is either the value of the NVCC_THREADS + # environment variable (if defined) or 1. + # when it is set, we reduce `num_jobs` to avoid + # overloading the system. + nvcc_threads = envs.NVCC_THREADS + if nvcc_threads is not None: + nvcc_threads = int(nvcc_threads) + logger.info( + "Using NVCC_THREADS=%d as the number of nvcc threads.", + nvcc_threads, + ) + else: + nvcc_threads = 1 + num_jobs = max(1, num_jobs // nvcc_threads) + except Exception as e: + logger.warning("Failed to get NVCC version: %s", e) + + return num_jobs, nvcc_threads + + # + # Perform cmake configuration for a single extension. + # + def configure(self, ext: CMakeExtension) -> None: + # If we've already configured using the CMakeLists.txt for + # this extension, exit early. + if ext.cmake_lists_dir in cmake_build_ext.did_config: + return + + cmake_build_ext.did_config[ext.cmake_lists_dir] = True + + # Select the build type. + # Note: optimization level + debug info are set by the build type + default_cfg = "Debug" if self.debug else "RelWithDebInfo" + cfg = envs.CMAKE_BUILD_TYPE or default_cfg + + cmake_args = [ + "-DCMAKE_BUILD_TYPE={}".format(cfg), + "-DAPHRODITE_TARGET_DEVICE={}".format(APHRODITE_TARGET_DEVICE), + ] + + verbose = envs.VERBOSE + if verbose: + cmake_args += ["-DCMAKE_VERBOSE_MAKEFILE=ON"] + + if is_sccache_available(): + cmake_args += [ + "-DCMAKE_C_COMPILER_LAUNCHER=sccache", + "-DCMAKE_CXX_COMPILER_LAUNCHER=sccache", + "-DCMAKE_CUDA_COMPILER_LAUNCHER=sccache", + "-DCMAKE_HIP_COMPILER_LAUNCHER=sccache", + ] + elif is_ccache_available(): + cmake_args += [ + "-DCMAKE_C_COMPILER_LAUNCHER=ccache", + "-DCMAKE_CXX_COMPILER_LAUNCHER=ccache", + "-DCMAKE_CUDA_COMPILER_LAUNCHER=ccache", + "-DCMAKE_HIP_COMPILER_LAUNCHER=ccache", + ] + + # Pass the python executable to cmake so it can find an exact + # match. + cmake_args += ["-DAPHRODITE_PYTHON_EXECUTABLE={}".format(sys.executable)] + + # Pass the python path to cmake so it can reuse the build dependencies + # on subsequent calls to python. + cmake_args += ["-DAPHRODITE_PYTHON_PATH={}".format(":".join(sys.path))] + + # Override the base directory for FetchContent downloads to $ROOT/.deps + # This allows sharing dependencies between profiles, + # and plays more nicely with sccache. + # To override this, set the FETCHCONTENT_BASE_DIR environment variable. + fc_base_dir = os.path.join(ROOT_DIR, ".deps") + fc_base_dir = os.environ.get("FETCHCONTENT_BASE_DIR", fc_base_dir) + cmake_args += ["-DFETCHCONTENT_BASE_DIR={}".format(fc_base_dir)] + + # + # Setup parallelism and build tool + # + num_jobs, nvcc_threads = self.compute_num_jobs() + + if nvcc_threads: + cmake_args += ["-DNVCC_THREADS={}".format(nvcc_threads)] + + if is_ninja_available(): + build_tool = ["-G", "Ninja"] + cmake_args += [ + "-DCMAKE_JOB_POOL_COMPILE:STRING=compile", + "-DCMAKE_JOB_POOLS:STRING=compile={}".format(num_jobs), + ] + else: + # Default build tool to whatever cmake picks. + build_tool = [] + # Make sure we use the nvcc from CUDA_HOME + if _is_cuda() and CUDA_HOME is not None: + cmake_args += [f"-DCMAKE_CUDA_COMPILER={CUDA_HOME}/bin/nvcc"] + elif _is_hip() and ROCM_HOME is not None: + cmake_args += [f"-DROCM_PATH={ROCM_HOME}"] + + other_cmake_args = os.environ.get("CMAKE_ARGS") + if other_cmake_args: + cmake_args += other_cmake_args.split() + + subprocess.check_call( + ["cmake", ext.cmake_lists_dir, *build_tool, *cmake_args], + cwd=self.build_temp, + ) + + def build_extensions(self) -> None: + # Ensure that CMake is present and working + try: + subprocess.check_output(["cmake", "--version"]) + except OSError as e: + raise RuntimeError("Cannot find CMake executable") from e + + # Create build directory if it does not exist. + if not os.path.exists(self.build_temp): + os.makedirs(self.build_temp) + + targets = [] + + def target_name(s: str) -> str: + if s.endswith("._paged_ops"): + return "_paged_ops" + return s.removeprefix("aphrodite.").removeprefix("vllm_flash_attn.") + + # Build all the extensions + for ext in self.extensions: + self.configure(ext) + targets.append(target_name(ext.name)) + + num_jobs, _ = self.compute_num_jobs() + + build_args = [ + "--build", + ".", + f"-j={num_jobs}", + *[f"--target={name}" for name in targets], + ] + + subprocess.check_call(["cmake", *build_args], cwd=self.build_temp) + + # Install the libraries + for ext in self.extensions: + # Install the extension into the proper location + outdir = Path(self.get_ext_fullpath(ext.name)).parent.absolute() + + # Skip if the install directory is the same as the build directory + if outdir == self.build_temp: + continue + + # CMake appends the extension prefix to the install path, + # and outdir already contains that prefix, so we need to remove it. + prefix = outdir + for _ in range(ext.name.count(".")): + prefix = prefix.parent + + # prefix here should actually be the same for all components + install_args = [ + "cmake", + "--install", + ".", + "--prefix", + prefix, + "--component", + target_name(ext.name), + ] + subprocess.check_call(install_args, cwd=self.build_temp) + + def run(self): + # First, run the standard build_ext command to compile the extensions + super().run() + + # bundle tcmalloc into CPU wheels for best OOB perf + if should_bundle_tcmalloc(): + bundle_tcmalloc(self.build_lib) + + # copy aphrodite/vllm_flash_attn/**/*.py from self.build_lib to current + # directory so that they can be included in the editable build + import glob + + files = glob.glob( + os.path.join(self.build_lib, "aphrodite", "vllm_flash_attn", "**", "*.py"), + recursive=True, + ) + for file in files: + dst_file = os.path.join("aphrodite/vllm_flash_attn", file.split("aphrodite/vllm_flash_attn/")[-1]) + print(f"Copying {file} to {dst_file}") + os.makedirs(os.path.dirname(dst_file), exist_ok=True) + self.copy_file(file, dst_file) + + if _is_cuda() or _is_hip(): + # copy aphrodite/third_party/triton_kernels/**/*.py from self.build_lib + # to current directory so that they can be included in the editable + # build + print( + f"Copying {self.build_lib}/aphrodite/third_party/triton_kernels to aphrodite/third_party/triton_kernels" + ) + shutil.copytree( + f"{self.build_lib}/aphrodite/third_party/triton_kernels", + "aphrodite/third_party/triton_kernels", + dirs_exist_ok=True, + ) + + if _is_cuda(): + # copy vendored deep_gemm package from build_lib to source tree + # for editable installs + deep_gemm_build = os.path.join(self.build_lib, "aphrodite", "third_party", "deep_gemm") + if os.path.exists(deep_gemm_build): + print(f"Copying {deep_gemm_build} to aphrodite/third_party/deep_gemm") + shutil.copytree( + deep_gemm_build, + "aphrodite/third_party/deep_gemm", + dirs_exist_ok=True, + ) + + # copy vendored fmha_sm100 package from build_lib to source tree + # for editable installs + fmha_sm100_build = os.path.join(self.build_lib, "aphrodite", "third_party", "fmha_sm100") + if os.path.exists(fmha_sm100_build): + print(f"Copying {fmha_sm100_build} to aphrodite/third_party/fmha_sm100") + shutil.copytree( + fmha_sm100_build, + "aphrodite/third_party/fmha_sm100", + dirs_exist_ok=True, + ) + + # copy vendored tml-fa4 package from build_lib to source tree + # for editable installs + tml_fa4_build = os.path.join(self.build_lib, "aphrodite", "third_party", "tml_fa4") + if os.path.exists(tml_fa4_build): + print(f"Copying {tml_fa4_build} to aphrodite/third_party/tml_fa4") + shutil.copytree( + tml_fa4_build, + "aphrodite/third_party/tml_fa4", + dirs_exist_ok=True, + ) + + +def _no_device() -> bool: + return APHRODITE_TARGET_DEVICE == "empty" + + +def _is_cuda() -> bool: + has_cuda = torch.version.cuda is not None + return APHRODITE_TARGET_DEVICE == "cuda" and has_cuda and not _is_tpu() + + +def _is_hip() -> bool: + return (APHRODITE_TARGET_DEVICE == "cuda" or APHRODITE_TARGET_DEVICE == "rocm") and torch.version.hip is not None + + +def _is_tpu() -> bool: + return APHRODITE_TARGET_DEVICE == "tpu" + + +def _is_cpu() -> bool: + return APHRODITE_TARGET_DEVICE == "cpu" + + +def _is_xpu() -> bool: + return APHRODITE_TARGET_DEVICE == "xpu" + + +def _is_metal() -> bool: + return APHRODITE_TARGET_DEVICE == "metal" + + +def _build_custom_ops() -> bool: + return _is_cuda() or _is_hip() or _is_metal() + + +def get_rocm_version(): + # Get the Rocm version from the ROCM_HOME/bin/librocm-core.so + # see https://github.com/ROCm/rocm-core/blob/d11f5c20d500f729c393680a01fa902ebf92094b/rocm_version.cpp#L21 + try: + if ROCM_HOME is None: + return None + librocm_core_file = Path(ROCM_HOME) / "lib" / "librocm-core.so" + if not librocm_core_file.is_file(): + return None + librocm_core = ctypes.CDLL(librocm_core_file) + VerErrors = ctypes.c_uint32 + get_rocm_core_version = librocm_core.getROCmVersion + get_rocm_core_version.restype = VerErrors + get_rocm_core_version.argtypes = [ + ctypes.POINTER(ctypes.c_uint32), + ctypes.POINTER(ctypes.c_uint32), + ctypes.POINTER(ctypes.c_uint32), + ] + major = ctypes.c_uint32() + minor = ctypes.c_uint32() + patch = ctypes.c_uint32() + + if get_rocm_core_version(ctypes.byref(major), ctypes.byref(minor), ctypes.byref(patch)) == 0: + return f"{major.value}.{minor.value}.{patch.value}" + return None + except Exception: + return None + + +def get_nvcc_cuda_version() -> Version: + """Get the CUDA version from nvcc. + + Adapted from https://github.com/NVIDIA/apex/blob/8b7a1ff183741dd8f9b87e7bafd04cfde99cea28/setup.py + """ + assert CUDA_HOME is not None, "CUDA_HOME is not set" + nvcc_output = subprocess.check_output([CUDA_HOME + "/bin/nvcc", "-V"], universal_newlines=True) + output = nvcc_output.split() + release_idx = output.index("release") + 1 + nvcc_cuda_version = parse(output[release_idx].split(",")[0]) + return nvcc_cuda_version + + +def get_aphrodite_version() -> str: + # Allow overriding the version. This is useful to build platform-specific + # wheels (e.g. CPU, TPU) without modifying the source. + if env_version := os.getenv("APHRODITE_VERSION_OVERRIDE"): + print(f"Overriding APHRODITE version with {env_version} from APHRODITE_VERSION_OVERRIDE") + os.environ["SETUPTOOLS_SCM_PRETEND_VERSION"] = env_version + return get_version(write_to="aphrodite/_version.py") + + version = get_version(write_to="aphrodite/_version.py") + sep = "+" if "+" not in version else "." # dev versions might contain + + + if _no_device(): + if envs.APHRODITE_TARGET_DEVICE == "empty": + version += f"{sep}empty" + elif _is_cuda(): + if APHRODITE_USE_PRECOMPILED: + if not envs.APHRODITE_SKIP_PRECOMPILED_VERSION_SUFFIX: + version += f"{sep}precompiled" + else: + cuda_version = str(get_nvcc_cuda_version()) + if cuda_version != envs.APHRODITE_MAIN_CUDA_VERSION: + cuda_version_str = cuda_version.replace(".", "")[:3] + # skip this for source tarball, required for pypi + if "sdist" not in sys.argv: + version += f"{sep}cu{cuda_version_str}" + elif _is_hip(): + # Get the Rocm Version + rocm_version = get_rocm_version() or torch.version.hip + if rocm_version and rocm_version != envs.APHRODITE_MAIN_CUDA_VERSION: + version += f"{sep}rocm{rocm_version.replace('.', '')[:3]}" + elif _is_tpu(): + version += f"{sep}tpu" + elif _is_cpu(): + # Check the local APHRODITE_TARGET_DEVICE (may be set by auto-detect above), + # not envs.APHRODITE_TARGET_DEVICE, so CPU-only hosts still get `+cpu`. + if APHRODITE_TARGET_DEVICE == "cpu": + version += f"{sep}cpu" + elif _is_metal(): + version += f"{sep}metal" + elif _is_xpu(): + version += f"{sep}xpu" + else: + raise RuntimeError("Unknown runtime environment") + + return version + + +def get_requirements() -> list[str]: + """Get Python package dependencies from requirements.txt.""" + requirements_dir = ROOT_DIR / "requirements" + + def _read_requirements(filename: str) -> list[str]: + with open(requirements_dir / filename) as f: + requirements = f.read().strip().split("\n") + resolved_requirements = [] + for line in requirements: + if line.startswith("-r "): + resolved_requirements += _read_requirements(line.split()[1]) + elif not line.startswith("--") and not line.startswith("#") and line.strip() != "": + resolved_requirements.append(line) + return resolved_requirements + + if _no_device(): + requirements = _read_requirements("common.txt") + elif _is_cuda(): + requirements = _read_requirements("cuda.txt") + cuda_major, cuda_minor = torch.version.cuda.split(".") + modified_requirements = [] + for req in requirements: + if "aphrodite-flash-attn" in req and cuda_major != "12": + # aphrodite-flash-attn is built only for CUDA 12.x. + # Skip for other versions. + continue + if "flashinfer-cubin" in req: + # Not on PyPI since 0.6.14 (only https://flashinfer.ai/whl), so + # it cannot be a wheel dependency; flashinfer falls back to + # fetching cubins at runtime when the package is absent. + continue + if "nvidia-cutlass-dsl[cu13]" in req and cuda_major == "12": + # [cu13] extra is the default; strip it on CUDA 12 builds. + req = req.replace("nvidia-cutlass-dsl[cu13]", "nvidia-cutlass-dsl") + if "humming-kernels[cu13]" in req and cuda_major == "12": + req = req.replace("humming-kernels[cu13]", "humming-kernels[cu12]") + modified_requirements.append(req) + requirements = modified_requirements + elif _is_hip(): + requirements = _read_requirements("rocm.txt") + elif _is_tpu(): + requirements = _read_requirements("tpu.txt") + elif _is_cpu(): + requirements = _read_requirements("cpu.txt") + elif _is_metal(): + requirements = _read_requirements("metal.txt") + elif _is_xpu(): + requirements = _read_requirements("xpu.txt") + else: + raise ValueError("Unsupported platform, please use CUDA, ROCm, Metal, or CPU.") + return requirements + + +ext_modules = [] + +if _is_cuda() or _is_hip(): + ext_modules.append(CMakeExtension(name="aphrodite.cumem_allocator")) + # Optional since this doesn't get built (produce an .so file). This is just + # 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): + ext_modules.append(CMakeExtension(name="aphrodite.spinloop")) + ext_modules.append(CMakeExtension(name="aphrodite.fs_io_C")) + +if _is_hip(): + ext_modules.append(CMakeExtension(name="aphrodite._rocm_C")) + +if _is_cuda(): + ext_modules.append(CMakeExtension(name="aphrodite.vllm_flash_attn._vllm_fa2_C")) + if CUDA_HOME and get_nvcc_cuda_version() >= Version("12.3"): + # FA3 requires CUDA 12.3 or later + ext_modules.append(CMakeExtension(name="aphrodite.vllm_flash_attn._vllm_fa3_C")) + # FA4 CuteDSL - Python-only component for FA4's cute DSL support + # Optional since this doesn't produce a .so file, just copies Python files + ext_modules.append(CMakeExtension(name="aphrodite.vllm_flash_attn._vllm_fa4_cutedsl_C", optional=True)) + if CUDA_HOME and get_nvcc_cuda_version() >= Version("12.9"): + # FlashMLA requires CUDA 12.9 or later + # Optional since this doesn't get built (produce an .so file) when + # 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.3"): + # DeepGEMM requires CUDA 12.3+ (SM90/SM100) + # Optional since it won't build on unsupported architectures + ext_modules.append(CMakeExtension(name="aphrodite._deep_gemm_C", optional=True)) + ext_modules.append(CMakeExtension(name="aphrodite._qutlass_C", optional=True)) + # fmha_sm100 is a Python/CuTe-DSL package installed into aphrodite.third_party. + ext_modules.append(CMakeExtension(name="aphrodite.fmha_sm100", optional=True)) + # tml-fa4 is copied into an isolated aphrodite.third_party package. + ext_modules.append(CMakeExtension(name="aphrodite.tml_fa4", optional=True)) + +if _is_cpu(): + import platform + + if platform.machine() in ("x86_64", "AMD64"): + ext_modules.append(CMakeExtension(name="aphrodite._C")) + ext_modules.append(CMakeExtension(name="aphrodite._C_AVX512")) + ext_modules.append(CMakeExtension(name="aphrodite._C_AVX2")) + else: + ext_modules.append(CMakeExtension(name="aphrodite._C")) + +if _build_custom_ops(): + if _is_metal(): + # MLX/nanobind paged-attention Metal kernel (aphrodite/metal/metal/_paged_ops). + ext_modules.append(CMakeExtension(name="aphrodite.metal.metal._paged_ops", py_limited_api=False)) + if _is_hip(): + ext_modules.append(CMakeExtension(name="aphrodite._C")) + if _is_cuda() or _is_hip(): + ext_modules.append(CMakeExtension(name="aphrodite._C_stable_libtorch")) + ext_modules.append(CMakeExtension(name="aphrodite._moe_C_stable_libtorch")) + +package_data = { + "aphrodite": [ + "py.typed", + "libs/*.so*", + "model_executor/layers/fused_moe/configs/*.json", + "model_executor/layers/mamba/ops/configs/selective_state_update/*.json", + "model_executor/layers/quantization/utils/configs/*.json", + "entrypoints/serve/instrumentator/static/*.js", + "entrypoints/serve/instrumentator/static/*.css", + "distributed/kv_transfer/kv_connector/v1/hf3fs/utils/*.cpp", + "third_party/flash_linear_attention/LICENSE", + # DeepGEMM JIT include headers (vendored via cmake) + "third_party/deep_gemm/include/**/*.cuh", + "third_party/deep_gemm/include/**/*.h", + "third_party/deep_gemm/include/**/*.hpp", + # fmha_sm100 sparse CuTe-DSL helper kernels (vendored via cmake) + "third_party/fmha_sm100/csrc/**/*.cu", + "third_party/fmha_sm100/csrc/**/*.h", + "third_party/fmha_sm100/csrc/**/*.jinja", + "third_party/fmha_sm100/csrc/**/*.cu.jinja", + "third_party/fmha_sm100/cute/**/*.cu", + "third_party/fmha_sm100/cutlass/include/**/*.h", + "third_party/fmha_sm100/cutlass/include/**/*.hpp", + "third_party/fmha_sm100/cutlass/tools/util/include/**/*.h", + "third_party/fmha_sm100/cutlass/tools/util/include/**/*.hpp", + # tml-fa4 CuTe-DSL helper kernels (vendored via cmake) + "third_party/tml_fa4/**/*.py", + ] +} + + +def add_aphrodite_package_data(filename: str) -> None: + aphrodite_files = package_data.setdefault("aphrodite", []) + if filename not in aphrodite_files: + aphrodite_files.append(filename) + + +# If the rust frontend binary is already present in the source tree (e.g., +# pre-built in a separate Docker build stage), ship it as-is. +if PRECOMPILED_RUST_FRONTEND_PATH.exists(): + add_aphrodite_package_data("aphrodite-rs") +for rust_extension_path in get_precompiled_rust_extension_paths(): + add_aphrodite_package_data(rust_extension_path.name) + +if _no_device(): + ext_modules = [] + +if APHRODITE_USE_PRECOMPILED: + wheel = precompiled_wheel.determine_wheel( + ROOT_DIR, + location=envs.APHRODITE_PRECOMPILED_WHEEL_LOCATION, + commit=envs.APHRODITE_PRECOMPILED_WHEEL_COMMIT, + ) + package_data_patch = precompiled_wheel.extract_precompiled_wheel( + ROOT_DIR, + wheel, + extract_extensions=True, + extract_rust=True, + ) + for package, files in package_data_patch.items(): + package_data.setdefault(package, []).extend(files) + ext_modules = [] + +if not ext_modules: + cmdclass = {} +else: + cmdclass = {"build_ext": cmake_build_ext} + +# Rust artifacts, built via setuptools-rust and installed into the package +# directory alongside the Python modules. Imported lazily: setuptools-rust does +# not need to be installed when nothing is being built. +if APHRODITE_USE_PRECOMPILED or should_use_precompiled_rust(): + rust_extensions = [] +else: + rust_build = load_module_from_path("rust_build", os.path.join(ROOT_DIR, "tools", "build_rust.py")) + rust_extensions = rust_build.rust_extensions(optional=not should_require_rust_frontend()) + +setup( + # static metadata should rather go in pyproject.toml + version=get_aphrodite_version(), + ext_modules=ext_modules, + rust_extensions=rust_extensions, + install_requires=get_requirements(), + extras_require={ + # AMD Zen CPU optimizations via zentorch + "zen": ["zentorch==2.11.0.0"], + "bench": ["pandas", "matplotlib", "seaborn", "datasets", "scipy", "plotly"], + "tensorizer": ["tensorizer==2.10.1"], + "fastsafetensors": ["fastsafetensors >= 0.3.2"], + "instanttensor": ["instanttensor >= 0.1.9"], + "runai": ["runai-model-streamer[s3,gcs,azure] >= 0.15.7"], + "audio": [ + "av", + "scipy", + "soundfile", + "soxr", + "mistral_common[audio]", + ], # Required for audio processing + "video": [], # Kept for backwards compatibility + # NVIDIA DeepStream (NVDEC) GPU video-decode backend. Linux x86-64 + # only; also needs system GStreamer + libv4l (see docs). + "deepstream": ["nvidia-deepstream-videodecode-cu13>=9.0.2"], + "flashinfer": [], # Kept for backwards compatibility + # Optional deps for Helion kernel development + # NOTE: When updating helion version, also update CI files: + # - .buildkite/test_areas/kernels.yaml + # - .buildkite/test-amd.yaml + "helion": ["helion==1.1.0"], + # Optional deps for gRPC server (aphrodite serve --grpc) + "grpc": ["smg-grpc-servicer[aphrodite] >= 0.5.2"], + # Optional deps for OpenTelemetry tracing + "otel": [ + "opentelemetry-sdk>=1.26.0", + "opentelemetry-api>=1.26.0", + "opentelemetry-exporter-otlp>=1.26.0", + "opentelemetry-semantic-conventions-ai>=0.4.1", + ], + # extra quantization plugin + "extra-quant": ["aphrodite-gguf-plugin>=0.0.2"], + }, + cmdclass=cmdclass, + package_data=package_data, +) 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/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/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_activation.py.orig b/tests/kernels/core/test_activation.py.orig new file mode 100644 index 0000000000..eab3d942a6 --- /dev/null +++ b/tests/kernels/core/test_activation.py.orig @@ -0,0 +1,222 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import random + +import pytest +import torch + +from aphrodite.model_executor.layers.activation import ( + FastGELU, + FatreluAndMul, + GeluAndMul, + MulAndSilu, + NewGELU, + QuickGELU, + ReLUSquaredActivation, + SiluAndMul, + SiluAndMulWithClamp, + SwigluOAIAndMul, + SwigluStepAndMul, + swiglustep_and_mul_triton, +) +from aphrodite.utils.torch_utils import set_random_seed +from tests.kernels.allclose_default import get_default_atol, get_default_rtol +from tests.kernels.utils import opcheck + +DTYPES = [torch.half, torch.bfloat16, torch.float] +NUM_TOKENS = [7, 83, 2048] # Arbitrary values for testing +D = [512, 13824] # Arbitrary values for testing +SEEDS = [0] +CUDA_DEVICES = [f"cuda:{i}" for i in range(1 if torch.accelerator.device_count() == 1 else 2)] + + +@pytest.mark.parametrize( + "activation", + [ + "silu_and_mul", + "mul_and_silu", + "gelu", + "gelu_tanh", + "fatrelu", + "swigluoai_and_mul", + "swiglustep_and_mul", + ], +) +@pytest.mark.parametrize("num_tokens", NUM_TOKENS) +@pytest.mark.parametrize("d", D) +@pytest.mark.parametrize("dtype", DTYPES) +@pytest.mark.parametrize("seed", SEEDS) +@pytest.mark.parametrize("device", CUDA_DEVICES) +@torch.inference_mode() +def test_act_and_mul( + default_aphrodite_config, + activation: str, + num_tokens: int, + d: int, + dtype: torch.dtype, + seed: int, + device: str, +) -> None: + set_random_seed(seed) + torch.set_default_device(device) + x = torch.randn(num_tokens, 2 * d, dtype=dtype) + if activation == "silu_and_mul": + layer = SiluAndMul(compile_native=False) + fn = torch.ops._C.silu_and_mul + if activation == "mul_and_silu": + layer = MulAndSilu() + fn = torch.ops._C.mul_and_silu + elif activation == "gelu": + layer = GeluAndMul(approximate="none") + fn = torch.ops._C.gelu_and_mul + elif activation == "gelu_tanh": + layer = GeluAndMul(approximate="tanh") + fn = torch.ops._C.gelu_tanh_and_mul + elif activation == "fatrelu": + threshold = random.uniform(0, 1) + layer = FatreluAndMul(threshold) + fn = torch.ops._C.fatrelu_and_mul + elif activation == "swigluoai_and_mul": + layer = SwigluOAIAndMul() + fn = torch.ops._C.swigluoai_and_mul + elif activation == "swiglustep_and_mul": + layer = SwigluStepAndMul() + fn = swiglustep_and_mul_triton + out = layer(x) + ref_out = layer.forward_native(x) + if activation in ["swigluoai_and_mul", "swiglustep_and_mul"]: + rtol = { + # For fp16, change the relative tolerance from 1e-3 to 2e-3 + torch.float16: 2e-3, + torch.bfloat16: 2e-2, + torch.float: 1.3e-6, + } + + def _get_rtol(output) -> float: + return rtol[output.dtype] + + torch.testing.assert_close(out, ref_out, atol=get_default_atol(out), rtol=_get_rtol(out)) + else: + # The SiluAndMul, MulAndSilu, GELU and FatReLU implementations are + # equivalent to the native PyTorch implementations, so we can do exact + # comparison. + torch.testing.assert_close(out, ref_out, atol=0.0, rtol=0.0) + + d = x.shape[-1] // 2 + output_shape = x.shape[:-1] + (d,) + out = torch.empty(output_shape, dtype=x.dtype, device=x.device) + if activation == "fatrelu": + opcheck(fn, (out, x, threshold)) + elif activation == "swigluoai_and_mul": + opcheck(fn, (out, x, layer.alpha, layer.limit)) + elif activation != "swiglustep_and_mul": + opcheck(fn, (out, x)) + + +SWIGLU_LIMITS = [3.0, 7.0, 15.0] + + +@pytest.mark.parametrize("swiglu_limit", SWIGLU_LIMITS) +@pytest.mark.parametrize("num_tokens", NUM_TOKENS) +@pytest.mark.parametrize("d", D) +@pytest.mark.parametrize("dtype", DTYPES) +@pytest.mark.parametrize("seed", SEEDS) +@pytest.mark.parametrize("device", CUDA_DEVICES) +@torch.inference_mode() +def test_silu_and_mul_with_clamp( + default_aphrodite_config, + swiglu_limit: float, + num_tokens: int, + d: int, + dtype: torch.dtype, + seed: int, + device: str, +) -> None: + """SiluAndMulWithClamp: cuda kernel must match native reference.""" + set_random_seed(seed) + torch.set_default_device(device) + # Use large values to ensure clamping is exercised. + x = torch.randn(num_tokens, 2 * d, dtype=dtype) * swiglu_limit * 2 + + layer = SiluAndMulWithClamp(swiglu_limit, compile_native=False) + out = layer(x) + ref_out = layer.forward_native(x) + + rtol = { + torch.float16: 2e-3, + torch.bfloat16: 2e-2, + torch.float: 1.3e-6, + } + torch.testing.assert_close(out, ref_out, atol=get_default_atol(out), rtol=rtol[out.dtype]) + + # Verify clamping is actually being applied: the clamped output should + # differ from the unclamped SiluAndMul output when inputs are large. + unclamped_out = SiluAndMul.forward_native(x) + assert not torch.equal(ref_out.float(), unclamped_out.float()), ( + "Input was not large enough to exercise the clamp; increase scale" + ) + + # Verify gate clamping semantics with a controlled scalar case. + # gate=large_val is clamped to limit first, then silu(limit) * 1.0. + x_gate = torch.tensor([[swiglu_limit * 20.0, 1.0]], dtype=torch.float32, device=device) + out_gate = SiluAndMulWithClamp(swiglu_limit, compile_native=False)(x_gate) + expected_gate = torch.nn.functional.silu(torch.tensor(swiglu_limit, dtype=torch.float32)).item() + torch.testing.assert_close( + out_gate, + torch.tensor([[expected_gate]], dtype=torch.float32, device=device), + atol=1e-3, + rtol=1e-3, + ) + + # Verify up clamping semantics: up >> limit gets clamped to limit. + x_up = torch.tensor([[1.0, swiglu_limit * 20.0]], dtype=torch.float32, device=device) + out_up = SiluAndMulWithClamp(swiglu_limit, compile_native=False)(x_up) + silu_1 = torch.nn.functional.silu(torch.tensor(1.0)).item() + torch.testing.assert_close( + out_up, + torch.tensor([[silu_1 * swiglu_limit]], dtype=torch.float32, device=device), + atol=1e-3, + rtol=1e-3, + ) + + # opcheck + out_buf = torch.empty(x.shape[:-1] + (d,), dtype=dtype, device=device) + opcheck(torch.ops._C.silu_and_mul_with_clamp, (out_buf, x, swiglu_limit)) + + +@pytest.mark.parametrize( + "activation", + [ + (FastGELU, torch.ops._C.gelu_fast), + (NewGELU, torch.ops._C.gelu_new), + (QuickGELU, torch.ops._C.gelu_quick), + (ReLUSquaredActivation, torch.ops._C.relu_squared), + ], +) +@pytest.mark.parametrize("num_tokens", NUM_TOKENS) +@pytest.mark.parametrize("d", D) +@pytest.mark.parametrize("dtype", DTYPES) +@pytest.mark.parametrize("seed", SEEDS) +@pytest.mark.parametrize("device", CUDA_DEVICES) +@torch.inference_mode() +def test_activation( + default_aphrodite_config, + activation: type[torch.nn.Module], + num_tokens: int, + d: int, + dtype: torch.dtype, + seed: int, + device: str, +) -> None: + set_random_seed(seed) + torch.set_default_device(device) + x = torch.randn(num_tokens, d, dtype=dtype) + layer = activation[0]() + fn = activation[1] + out = layer(x) + ref_out = layer.forward_native(x) + torch.testing.assert_close(out, ref_out, atol=get_default_atol(out), rtol=get_default_rtol(out)) + + out = torch.empty_like(x) + opcheck(fn, (out, x)) 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_grouped_topk.py.orig b/tests/kernels/moe/test_grouped_topk.py.orig new file mode 100644 index 0000000000..72df9a933f --- /dev/null +++ b/tests/kernels/moe/test_grouped_topk.py.orig @@ -0,0 +1,316 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Tests for the MoE grouped topk kernel + +Run `pytest tests/kernels/moe/test_grouped_topk.py`. +""" + +import pytest +import torch + +import aphrodite.envs as envs +from aphrodite.config import ( + AphroditeConfig, + CompilationConfig, + get_cached_compilation_config, + set_current_aphrodite_config, +) +from aphrodite.model_executor.layers.fused_moe.router.grouped_topk_router import ( + GroupedTopk, + fused_grouped_topk, +) +from aphrodite.platforms import current_platform +from aphrodite.utils.torch_utils import set_random_seed + + +@pytest.mark.skipif(not current_platform.is_cuda(), reason="This test is skipped on non-CUDA platform.") +@pytest.mark.parametrize("n_token", [1, 33, 64]) +@pytest.mark.parametrize("n_hidden", [1024, 2048]) +@pytest.mark.parametrize( + "n_expert,topk,num_expert_group,topk_group", + [ + (16, 2, 8, 2), + (128, 2, 8, 2), + (256, 8, 8, 4), + (384, 8, 1, 1), + (512, 22, 1, 1), + ], +) +@pytest.mark.parametrize("renormalize", [True, False]) +@pytest.mark.parametrize("scoring_func", ["softmax", "sigmoid"]) +@pytest.mark.parametrize("routed_scaling_factor", [1.0, 2.5]) +@pytest.mark.parametrize("input_dtype", [torch.bfloat16, torch.float32]) +@pytest.mark.parametrize("bias_dtype", [torch.float32]) +def test_grouped_topk( + monkeypatch: pytest.MonkeyPatch, + n_token: int, + n_hidden: int, + n_expert: int, + topk: int, + num_expert_group: int, + topk_group: int, + renormalize: bool, + scoring_func: str, + routed_scaling_factor: float, + input_dtype: torch.dtype, + bias_dtype: torch.dtype, +): + aphrodite_config = AphroditeConfig(compilation_config=CompilationConfig(custom_ops=["all", "+grouped_topk"])) + get_cached_compilation_config.cache_clear() + + set_random_seed(0) + hidden_states = torch.randn((n_token, n_hidden), dtype=input_dtype, device="cuda") + gating_output = torch.randn((n_token, n_expert), dtype=input_dtype, device="cuda") + e_score_correction_bias = torch.randn((n_expert,), dtype=bias_dtype, device="cuda") + + with set_current_aphrodite_config(aphrodite_config), monkeypatch.context() as m: + m.setenv("APHRODITE_USE_FUSED_MOE_GROUPED_TOPK", "0") + m.setattr(envs, "APHRODITE_BATCH_INVARIANT", True) + grouped_topk = GroupedTopk( + topk=topk, + renormalize=renormalize, + num_expert_group=num_expert_group, + topk_group=topk_group, + scoring_func=scoring_func, + routed_scaling_factor=routed_scaling_factor, + ) + assert grouped_topk._forward_method.__name__ == "forward_cuda" + baseline_topk_weights, baseline_topk_ids = grouped_topk( + hidden_states=hidden_states, + gating_output=gating_output, + e_score_correction_bias=e_score_correction_bias, + ) + + test_topk_weights, test_topk_ids = fused_grouped_topk( + hidden_states=hidden_states, + gating_output=gating_output, + topk=topk, + renormalize=renormalize, + num_expert_group=num_expert_group, + topk_group=topk_group, + scoring_func=scoring_func, + routed_scaling_factor=routed_scaling_factor, + e_score_correction_bias=e_score_correction_bias, + ) + + 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) + + +@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_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_moe_align_block_size.py.orig b/tests/kernels/moe/test_moe_align_block_size.py.orig new file mode 100644 index 0000000000..248a799bdb --- /dev/null +++ b/tests/kernels/moe/test_moe_align_block_size.py.orig @@ -0,0 +1,366 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Tests for the MOE align block size function. + +Run `pytest tests/kernels/moe/test_moe_align_block_size.py`. +""" + +import pytest +import torch + +from aphrodite.model_executor.layers.fused_moe.moe_align_block_size import ( + batched_moe_align_block_size, + moe_align_block_size, +) +from aphrodite.utils.math_utils import cdiv, round_up +from aphrodite.utils.torch_utils import set_random_seed + +NUM_TOKENS = [1, 3, 256, 2256, 4096] +NUM_EXPERTS = [32, 160, 256, 257] +TOP_KS = [1, 2, 16, 32] +BLOCK_SIZES = [32, 128] +set_random_seed(0) + + +def _group_tokens_by_expert( + sorted_ids: torch.Tensor, + expert_ids: torch.Tensor, + block_size: int, + valid_length: int, + total_tokens: int, +) -> dict: + num_blocks = valid_length // block_size + expert_tokens: dict[int, list[int]] = {} + + for block_idx in range(num_blocks): + expert_id = expert_ids[block_idx].item() + block_start = block_idx * block_size + block_end = min(block_start + block_size, valid_length) + + block_tokens = sorted_ids[block_start:block_end] + valid_tokens = block_tokens[block_tokens < total_tokens] + + if expert_id not in expert_tokens: + expert_tokens[expert_id] = [] + expert_tokens[expert_id].extend(valid_tokens.tolist()) + return expert_tokens + + +def _verify_expert_level_sorting( + actual_sorted_ids: torch.Tensor, + golden_sorted_ids: torch.Tensor, + expert_ids: torch.Tensor, + block_size: int, + valid_length: int, + total_tokens: int, +): + """ + Verify that actual_sorted_ids follows the correct expert-level sorting. + The kerne limplementation may or may not preserve original token order + in topk_ids in the final sorted_ids however this does not impact quality. + """ + # Group tokens by expert from the golden implementation + golden_expert_tokens = _group_tokens_by_expert( + golden_sorted_ids, expert_ids, block_size, valid_length, total_tokens + ) + + actual_expert_tokens = _group_tokens_by_expert( + actual_sorted_ids, expert_ids, block_size, valid_length, total_tokens + ) + + assert set(golden_expert_tokens.keys()) == set(actual_expert_tokens.keys()), ( + f"Expert IDs mismatch: golden={set(golden_expert_tokens.keys())}, actual={set(actual_expert_tokens.keys())}" + ) + + for expert_id in golden_expert_tokens: + golden_tokens = torch.tensor(golden_expert_tokens[expert_id], device=actual_sorted_ids.device) + actual_tokens = torch.tensor(actual_expert_tokens[expert_id], device=actual_sorted_ids.device) + assert torch.equal(torch.sort(golden_tokens)[0], torch.sort(actual_tokens)[0]), ( + f"Expert {expert_id} token mismatch: " + f"golden={golden_expert_tokens[expert_id]}, " + f"actual={actual_expert_tokens[expert_id]}" + ) + + +def torch_moe_align_block_size( + topk_ids: torch.Tensor, + block_size: int, + num_experts: int, + expert_map: torch.Tensor | None = None, + pad_sorted_ids: bool = False, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """ + Golden torch implementation of moe_align_block_size. + + This function aligns the token distribution across experts to be compatible + with block size for matrix multiplication by sorting tokens by expert and + padding to block boundaries. + """ + max_num_tokens_padded = topk_ids.numel() + num_experts * (block_size - 1) + if pad_sorted_ids: + max_num_tokens_padded = round_up(max_num_tokens_padded, block_size) + if topk_ids.numel() < num_experts: + max_num_tokens_padded = topk_ids.numel() * block_size + + flattened_token_indices = torch.arange(topk_ids.numel(), device=topk_ids.device, dtype=torch.int32) + flattened_expert_ids = topk_ids.flatten() + sorted_expert_ids, sort_indices = torch.sort(flattened_expert_ids, stable=True) + sorted_token_indices = flattened_token_indices[sort_indices] + + expert_token_counts = torch.zeros(num_experts, dtype=torch.int64, device=topk_ids.device) + for expert_id in range(num_experts): + mask = sorted_expert_ids == expert_id + expert_token_counts[expert_id] = mask.sum() + + expert_padded_counts = torch.zeros(num_experts, dtype=torch.int64, device=topk_ids.device) + for expert_id in range(num_experts): + original_count = expert_token_counts[expert_id] + if expert_map is not None and expert_map[expert_id] == -1: + continue + if original_count > 0: + expert_padded_counts[expert_id] = ((original_count + block_size - 1) // block_size) * block_size + + sorted_token_ids = torch.full( + (max_num_tokens_padded,), + topk_ids.numel(), + dtype=torch.int32, + device=topk_ids.device, + ) + max_num_blocks = (max_num_tokens_padded + block_size - 1) // block_size + expert_ids = torch.full((max_num_blocks,), -1, dtype=torch.int32, device=topk_ids.device) + + current_pos = 0 + current_block = 0 + for expert_id in range(num_experts): + if expert_map is not None and expert_map[expert_id] == -1: + continue + + expert_mask = sorted_expert_ids == expert_id + expert_tokens = sorted_token_indices[expert_mask] + num_expert_tokens = expert_tokens.shape[0] + + if num_expert_tokens > 0: + sorted_token_ids[current_pos : current_pos + num_expert_tokens] = expert_tokens + + expert_blocks_needed = expert_padded_counts[expert_id] // block_size + + expert_id_new = expert_id + if expert_map is not None: + expert_id_new = expert_map[expert_id] + expert_ids[current_block : current_block + expert_blocks_needed] = expert_id_new + + current_pos += expert_padded_counts[expert_id] + current_block += expert_blocks_needed + + total_padded_tokens = expert_padded_counts.sum() + num_tokens_post_pad = torch.tensor([total_padded_tokens], dtype=torch.int32, device=topk_ids.device) + + return sorted_token_ids, expert_ids, num_tokens_post_pad + + +@pytest.mark.parametrize("m", NUM_TOKENS) +@pytest.mark.parametrize("topk", TOP_KS) +@pytest.mark.parametrize("num_experts", NUM_EXPERTS) +@pytest.mark.parametrize("block_size", BLOCK_SIZES) +@pytest.mark.parametrize("pad_sorted_ids", [False, True]) +def test_moe_align_block_size(m: int, topk: int, num_experts: int, block_size: int, pad_sorted_ids: bool): + """Test moe_align_block_size without expert mapping""" + topk_ids = torch.zeros((m, topk), device="cuda", dtype=torch.int32) + for i in range(m): + experts = torch.randperm(num_experts, device="cuda")[:topk] + topk_ids[i] = experts + + actual_sorted_ids, actual_expert_ids, actual_num_tokens = moe_align_block_size( + topk_ids=topk_ids, + block_size=block_size, + num_experts=num_experts, + pad_sorted_ids=pad_sorted_ids, + ) + golden_sorted_ids, golden_expert_ids, golden_num_tokens = torch_moe_align_block_size( + topk_ids=topk_ids, + block_size=block_size, + num_experts=num_experts, + pad_sorted_ids=pad_sorted_ids, + ) + + torch.testing.assert_close(actual_num_tokens, golden_num_tokens, atol=0, rtol=0) + torch.testing.assert_close(actual_expert_ids, golden_expert_ids, atol=0, rtol=0) + + # For sorted_token_ids, verify block-level correctness rather than exact + # order Tokens within each expert's blocks can be in any order, but expert + # regions must be correct + _verify_expert_level_sorting( + actual_sorted_ids, + golden_sorted_ids, + actual_expert_ids, + block_size, + actual_num_tokens.item(), + m * topk, + ) + + total_tokens = m * topk + assert actual_num_tokens.item() % block_size == 0, "num_tokens_post_pad should be divisible by block_size" + assert actual_num_tokens.item() >= total_tokens, "num_tokens_post_pad should be at least total_tokens" + valid_tokens = actual_sorted_ids[actual_sorted_ids < total_tokens] + assert len(valid_tokens) == total_tokens, ( + f"Should have exactly {total_tokens} valid tokens, got {len(valid_tokens)}" + ) + actual_num_blocks = cdiv(int(actual_num_tokens.item()), block_size) + assert (actual_expert_ids[:actual_num_blocks] >= 0).all() and ( + actual_expert_ids[:actual_num_blocks] < num_experts + ).all(), "expert_ids should contain valid expert indices" + + +@pytest.mark.parametrize("m", [16, 32, 2048]) +@pytest.mark.parametrize("topk", [2, 4]) +@pytest.mark.parametrize("num_experts", [8, 64]) +@pytest.mark.parametrize("block_size", [64]) +@pytest.mark.parametrize("mask_inactive_experts", [False, True]) +def test_moe_align_block_size_with_expert_map( + m: int, + topk: int, + num_experts: int, + block_size: int, + mask_inactive_experts: bool, +): + """Test moe_align_block_size with expert mapping (EP scenario)""" + expert_map = torch.full((num_experts,), -1, device="cuda", dtype=torch.int32) + local_experts = list(range(0, num_experts, 2)) + for i, expert_id in enumerate(local_experts): + expert_map[expert_id] = i + + topk_ids = torch.empty((m, topk), device="cuda", dtype=torch.int32) + for i in range(m): + 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 + + actual_sorted_ids, actual_expert_ids, actual_num_tokens = moe_align_block_size( + topk_ids=topk_ids, + block_size=block_size, + num_experts=num_experts, + expert_map=expert_map, + ignore_invalid_experts=True, + ) + golden_sorted_ids, golden_expert_ids, golden_num_tokens = torch_moe_align_block_size( + topk_ids=topk_ids, + block_size=block_size, + num_experts=num_experts, + expert_map=expert_map, + ) + + torch.testing.assert_close(actual_num_tokens, golden_num_tokens, atol=0, rtol=0) + torch.testing.assert_close(actual_expert_ids, golden_expert_ids, atol=0, rtol=0) + _verify_expert_level_sorting( + actual_sorted_ids, + golden_sorted_ids, + actual_expert_ids, + block_size, + actual_num_tokens.item(), + m * topk, + ) + + +def test_moe_align_block_size_deterministic(): + m, topk, num_experts, block_size = 128, 2, 32, 64 + + torch.manual_seed(42) + topk_ids = torch.randint(0, num_experts, (m, topk), device="cuda", dtype=torch.int32) + + # expect the results to be reproducible + results = [] + for _ in range(5): + sorted_ids, expert_ids, num_tokens = moe_align_block_size( + topk_ids=topk_ids, block_size=block_size, num_experts=num_experts + ) + results.append((sorted_ids.clone(), expert_ids.clone(), num_tokens.clone())) + + for i in range(1, len(results)): + assert torch.equal(results[0][0], results[i][0]), "sorted_ids should be deterministic" + assert torch.equal(results[0][1], results[i][1]), "expert_ids should be deterministic" + assert torch.equal(results[0][2], results[i][2]), "num_tokens should be deterministic" + + +@pytest.mark.parametrize("max_tokens_per_batch", [13, 16, 512]) +@pytest.mark.parametrize("num_experts", [8, 16, 32, 64]) +@pytest.mark.parametrize("block_size", [8, 16, 32, 64]) +@pytest.mark.parametrize("simulate_empty_batches", [False, True]) +def test_batched_moe_align_block_size( + max_tokens_per_batch: int, + num_experts: int, + block_size: int, + simulate_empty_batches: bool, +): + def ref_outputs( + expert_num_tokens: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + E = expert_num_tokens.size(0) + + # Round up so each batch can be split to blocks evenly. + Msum = round_up(max_tokens_per_batch, block_size) * E + ref_sorted_ids = torch.empty((Msum,), dtype=torch.int32) + ref_expert_ids = torch.empty((Msum // block_size,), dtype=torch.int32) + ref_num_tokens_post_pad = torch.empty((1,), dtype=torch.int32) + + # Initialize + sentinel = E * max_tokens_per_batch + ref_sorted_ids.fill_(sentinel) + ref_expert_ids.fill_(-1) + + # Fill ref_sorted_ids + i = 0 + for expert_id, expert_nt in enumerate(expert_num_tokens): + token_offset = expert_id * max_tokens_per_batch + for j in range(expert_nt): + ref_sorted_ids[i] = token_offset + j + i += 1 + # round up i to the next block_size + i = round_up(i, block_size) + + ref_num_tokens_post_pad[0] = i + + # Fill expert_ids + nt_ceil_sum = 0 + for expert_id, expert_nt in enumerate(expert_num_tokens): + expert_ids_offset = nt_ceil_sum // block_size + ceil_expert_nt = round_up(int(expert_nt.item()), block_size) + num_blocks = ceil_expert_nt // block_size + for x in range(num_blocks): + ref_expert_ids[expert_ids_offset + x] = expert_id + nt_ceil_sum += ceil_expert_nt + + return ( + ref_sorted_ids.to("cuda"), + ref_expert_ids.to("cuda"), + ref_num_tokens_post_pad.to("cuda"), + ) + + # Compute expert_num_tokens + expert_num_tokens = torch.randint( + low=0, + high=max_tokens_per_batch, + size=(num_experts,), + device="cpu", + dtype=torch.int32, + ) + if simulate_empty_batches: + # mark half the batches to have 0 tokens + zero_batches = torch.randperm(num_experts)[: num_experts // 2] + expert_num_tokens[zero_batches] = 0 + + # ref outputs + ref_sorted_ids, ref_expert_ids, ref_num_tokens_post_pad = ref_outputs(expert_num_tokens) + + # outputs + sorted_ids, expert_ids, num_tokens_post_pad = batched_moe_align_block_size( + max_tokens_per_batch, block_size, expert_num_tokens.to("cuda") + ) + + assert ref_sorted_ids.size() == sorted_ids.size(), f"{ref_sorted_ids.size()} vs {sorted_ids.size()}" + assert ref_expert_ids.size() == expert_ids.size(), f"{ref_expert_ids.size()} vs {expert_ids.size()}" + assert ref_num_tokens_post_pad.size() == num_tokens_post_pad.size(), ( + f"{ref_num_tokens_post_pad.size()} vs {num_tokens_post_pad.size()}" + ) + torch.testing.assert_close(ref_sorted_ids, sorted_ids, atol=0, rtol=0) + torch.testing.assert_close(ref_expert_ids, expert_ids, atol=0, rtol=0) + torch.testing.assert_close(ref_num_tokens_post_pad, num_tokens_post_pad, atol=0, rtol=0) 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/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_attn_res.py.orig b/tests/models/kimi_k3/test_attn_res.py.orig new file mode 100644 index 0000000000..092657e258 --- /dev/null +++ b/tests/models/kimi_k3/test_attn_res.py.orig @@ -0,0 +1,171 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import pytest +import torch +import torch.nn.functional as F + +from aphrodite.models.kimi_k3.nvidia.ops import attn_res +from aphrodite.platforms import current_platform + +HIDDEN_SIZE = 7168 +MAX_BLOCKS = 8 +EPS = 1e-5 + + +def _randn_with_row_padding(*shape: int, padding: int = 0) -> torch.Tensor: + storage = torch.randn( + *shape[:-1], + shape[-1] + padding, + device="cuda", + dtype=torch.bfloat16, + ) + return storage[..., : shape[-1]] + + +def _reference( + prefix: torch.Tensor, + delta: torch.Tensor | None, + blocks: torch.Tensor, + norm_weight: torch.Tensor, + qk_weight: torch.Tensor, + output_norm_weight: torch.Tensor | None, + num_blocks: int, +) -> tuple[torch.Tensor, torch.Tensor]: + if delta is not None: + prefix = prefix + delta + values = torch.cat((blocks[:, :num_blocks], prefix.unsqueeze(1)), dim=1) + keys = F.rms_norm(values, (HIDDEN_SIZE,), norm_weight, EPS) + probs = (keys @ qk_weight).softmax(dim=-1) + output = torch.matmul(probs.unsqueeze(1), values).squeeze(1) + if output_norm_weight is not None: + output = F.rms_norm(output, (HIDDEN_SIZE,), output_norm_weight, EPS) + return output, prefix + + +@pytest.mark.parametrize( + ( + "num_tokens", + "num_blocks", + "row_padding", + "write_block", + "has_delta", + "backend", + ), + [ + pytest.param(1, 0, 0, True, False, "triton", id="triton-empty"), + pytest.param(1, 0, 0, True, True, "triton", id="triton-empty-add"), + pytest.param(17, 5, 7, True, False, "triton", id="triton-write"), + pytest.param(17, 5, 7, True, True, "triton", id="triton-write-add"), + pytest.param(3, 8, 0, False, False, "triton", id="triton-full"), + pytest.param(3, 8, 0, False, True, "triton", id="triton-full-add"), + pytest.param(320, 1, 0, False, True, "nvidia", id="nvidia-1"), + pytest.param(320, 4, 0, False, True, "nvidia", id="nvidia-4"), + pytest.param(320, 8, 0, False, True, "nvidia", id="nvidia-8"), + ], +) +def test_attn_res( + num_tokens: int, + num_blocks: int, + row_padding: int, + write_block: bool, + has_delta: bool, + backend: str, +): + if backend == "nvidia" and not current_platform.is_device_capability_family(100): + pytest.skip("NVIDIA AttnRes requires the SM100 family") + + prefix = _randn_with_row_padding(num_tokens, HIDDEN_SIZE, padding=row_padding) + delta = _randn_with_row_padding(num_tokens, HIDDEN_SIZE, padding=row_padding) if has_delta else None + blocks = _randn_with_row_padding(num_tokens, MAX_BLOCKS, HIDDEN_SIZE, padding=row_padding) + norm_weight = 1 + 0.1 * torch.randn(HIDDEN_SIZE, device="cuda", dtype=torch.bfloat16) + qk_weight = torch.randn(HIDDEN_SIZE, device="cuda", dtype=torch.bfloat16) / HIDDEN_SIZE**0.5 + output_norm_weight = 1 + 0.1 * torch.randn(HIDDEN_SIZE, device="cuda", dtype=torch.bfloat16) + original_blocks = blocks.clone() + expected, expected_prefix = _reference( + prefix.clone(), + delta, + blocks, + norm_weight, + qk_weight, + output_norm_weight, + num_blocks, + ) + block_write_idx = num_blocks if write_block else -1 + + actual = attn_res( + prefix, + delta, + blocks, + norm_weight, + qk_weight, + output_norm_weight, + num_blocks, + block_write_idx, + EPS, + EPS, + ) + + torch.testing.assert_close(actual, expected, atol=8e-2, rtol=3e-2) + torch.testing.assert_close(prefix, expected_prefix, atol=0, rtol=0) + if write_block: + original_blocks[:, block_write_idx].copy_(expected_prefix) + torch.testing.assert_close(blocks, original_blocks, atol=0, rtol=0) + assert actual.is_contiguous() + + +@pytest.mark.parametrize("num_blocks", range(MAX_BLOCKS + 1)) +def test_attn_res_block_counts(num_blocks: int): + prefix = torch.randn(1, HIDDEN_SIZE, device="cuda", dtype=torch.bfloat16) + blocks = torch.randn(1, MAX_BLOCKS, HIDDEN_SIZE, device="cuda", dtype=torch.bfloat16) + norm_weight = torch.ones(HIDDEN_SIZE, device="cuda", dtype=torch.bfloat16) + qk_weight = torch.randn(HIDDEN_SIZE, device="cuda", dtype=torch.bfloat16) / HIDDEN_SIZE**0.5 + output_norm_weight = torch.ones_like(norm_weight) + expected, _ = _reference( + prefix.clone(), + None, + blocks, + norm_weight, + qk_weight, + output_norm_weight, + num_blocks, + ) + + actual = attn_res( + prefix, + None, + blocks, + norm_weight, + qk_weight, + output_norm_weight, + num_blocks, + -1, + EPS, + EPS, + ) + + torch.testing.assert_close(actual, expected, atol=8e-2, rtol=3e-2) + + +def test_attn_res_without_output_norm(): + prefix = torch.randn(7, HIDDEN_SIZE, device="cuda", dtype=torch.bfloat16) + delta = torch.randn_like(prefix) + blocks = torch.randn(7, MAX_BLOCKS, HIDDEN_SIZE, device="cuda", dtype=torch.bfloat16) + norm_weight = torch.randn(HIDDEN_SIZE, device="cuda", dtype=torch.bfloat16) + qk_weight = torch.randn(HIDDEN_SIZE, device="cuda", dtype=torch.bfloat16) / HIDDEN_SIZE**0.5 + expected, _ = _reference(prefix.clone(), delta, blocks, norm_weight, qk_weight, None, MAX_BLOCKS) + + actual = attn_res( + prefix, + delta, + blocks, + norm_weight, + qk_weight, + None, + MAX_BLOCKS, + -1, + EPS, + 0.0, + ) + + torch.testing.assert_close(actual, expected, atol=8e-2, rtol=3e-2) 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..cafd41a570 --- /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_common_attn_metadata, + create_vllm_config, +) + +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: + vllm_config = create_vllm_config( + model_name="Qwen/Qwen3.5-0.8B", + block_size=BLOCK_SIZE, + ) + if num_speculative_tokens: + vllm_config.speculative_config = SpeculativeConfig( + method="ngram", + num_speculative_tokens=num_speculative_tokens, + ) + vllm_config.compilation_config.cudagraph_mode = ( + CUDAGraphMode.FULL_AND_PIECEWISE if full_cuda_graph else CUDAGraphMode.NONE + ) + vllm_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"], + vllm_config=vllm_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/generation/test_common.py.orig b/tests/models/multimodal/generation/test_common.py.orig new file mode 100644 index 0000000000..859643ae1a --- /dev/null +++ b/tests/models/multimodal/generation/test_common.py.orig @@ -0,0 +1,1355 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Common tests for testing .generate() functionality for single / multiple +image, embedding, and video support for different VLMs in Aphrodite. +""" + +import math +from collections import defaultdict +from pathlib import PosixPath + +import pytest +from packaging.version import Version +from transformers import ( + AutoModel, + AutoModelForCausalLM, + AutoModelForImageTextToText, + AutoModelForTextToWaveform, +) +from transformers import __version__ as TRANSFORMERS_VERSION + +from aphrodite.platforms import current_platform +from aphrodite.utils.func_utils import identity + +from ....conftest import ( + IMAGE_ASSETS, + AphroditeRunner, + AudioTestAssets, + HfRunner, + ImageTestAssets, + VideoTestAssets, +) +from ....utils import create_new_process_for_each_test, large_gpu_mark, multi_gpu_marks +from ...utils import check_outputs_equal +from .vlm_utils import custom_inputs, model_utils, runners +from .vlm_utils.case_filtering import get_parametrized_options +from .vlm_utils.types import ( + CustomTestOptions, + ExpandableVLMTestArgs, + VLMTestInfo, + VLMTestType, +) + +COMMON_BROADCAST_SETTINGS = { + "test_type": VLMTestType.IMAGE, + "dtype": "half", + "max_tokens": 5, + "tensor_parallel_size": 2, + "hf_model_kwargs": {"device_map": "auto"}, + "image_size_factors": [(0.25, 0.5, 1.0)], + "distributed_executor_backend": ( + "ray", + "mp", + ), +} + +### Test configuration for specific models +# NOTE: The convention of the test settings below is to lead each test key +# with the name of the model arch used in the test, using underscores in place +# of hyphens; this makes it more convenient to filter tests for a specific kind +# of model. For example.... +# +# To run all test types for a specific key: +# use the k flag to substring match with a leading square bracket; if the +# model arch happens to be a substring of another one, you can add a +# trailing hyphen. E.g., +# - pytest $TEST_FILE -k "[llava-" +# prevents matching on "[llava_next-" & will match just the enabled cases +# for llava, i.e., single image, image embedding, and custom input tests. +# +# To run a test for a Test Info for just one of multiple models: +# use the k flag to substring match the model name, e.g., +# - pytest $TEST_FILE -k OpenGVLab/InternVL2-1B +# prevents matching on nGVLab/InternVL2-2B. +# +# You can also combine substrings to match more granularly. +# ex 1: +# pytest $TEST_FILE -k "test_single_image and OpenGVLab/InternVL2-1B" +# will run only test_single_image* for OpenGVLab/InternVL2-1B; this would +# match both wrappers for single image tests, since it also matches +# test_single_image_heavy (which forks if we have a distributed backend) +# ex 2: +# pytest $TEST_FILE -k "[llava- or [intern_vl-" +# will run all of the tests for only llava & internvl. +# +# NOTE you can add --collect-only to any of the above commands to see +# which cases would be selected and deselected by pytest. In general, +# this is a good idea for checking your command first, since tests are slow. + + +def _granite4_vision_aphrodite_to_hf_output(aphrodite_output, model): + """Post-processor for granite4_vision Aphrodite output. + + Self-contained to avoid calling AutoConfig/AutoTokenizer without + trust_remote_code (needed while the model is not in upstream HF). + """ + output_ids, output_str, out_logprobs = aphrodite_output + mm_token_id = 100352 + hf_output_ids = [ + token_id + for idx, token_id in enumerate(output_ids) + if token_id != mm_token_id or idx == 0 or output_ids[idx - 1] != mm_token_id + ] + hf_output_str = output_str[1:] if output_str and output_str[0] == " " else output_str + eos_token_id = 100257 + if hf_output_ids and hf_output_ids[-1] == eos_token_id: + hf_output_str = hf_output_str + "<|end_of_text|>" + return hf_output_ids, hf_output_str, out_logprobs + + +VLM_TEST_SETTINGS = { + #### Core tests to always run in the CI + "llava": VLMTestInfo( + models=["llava-hf/llava-1.5-7b-hf"], + test_type=(VLMTestType.EMBEDDING, VLMTestType.IMAGE, VLMTestType.CUSTOM_INPUTS), + prompt_formatter=lambda img_prompt: f"USER: {img_prompt}\nASSISTANT:", + convert_assets_to_embeddings=model_utils.get_llava_embeddings, + max_model_len=4096, + auto_cls=AutoModelForImageTextToText, + aphrodite_output_post_proc=model_utils.llava_image_aphrodite_to_hf_output, + custom_test_opts=[ + CustomTestOptions( + inputs=custom_inputs.multi_image_multi_aspect_ratio_inputs( + formatter=lambda img_prompt: f"USER: {img_prompt}\nASSISTANT:" + ), + limit_mm_per_prompt={"image": 4}, + ) + ], + aphrodite_runner_kwargs={"enable_mm_embeds": True}, + marks=[pytest.mark.core_model, pytest.mark.cpu_model], + ), + "paligemma": VLMTestInfo( + models=["google/paligemma-3b-mix-224"], + test_type=VLMTestType.IMAGE, + prompt_formatter=identity, + img_idx_to_prompt=lambda idx: "", + # Paligemma uses its own sample prompts because the default one fails + single_image_prompts=IMAGE_ASSETS.prompts( + { + "stop_sign": "caption es", + "cherry_blossom": "What is in the picture?", + } + ), + auto_cls=AutoModelForImageTextToText, + aphrodite_output_post_proc=model_utils.paligemma_aphrodite_to_hf_output, + ), + "qwen2_5_vl": VLMTestInfo( + models=["Qwen/Qwen2.5-VL-3B-Instruct"], + test_type=(VLMTestType.IMAGE, VLMTestType.MULTI_IMAGE, VLMTestType.VIDEO), + prompt_formatter=lambda img_prompt: f"<|im_start|>User\n{img_prompt}<|im_end|>\n<|im_start|>assistant\n", # noqa: E501 + img_idx_to_prompt=lambda idx: "<|vision_start|><|image_pad|><|vision_end|>", + video_idx_to_prompt=lambda idx: "<|vision_start|><|video_pad|><|vision_end|>", + enforce_eager=False, + max_model_len=4096, + max_num_seqs=2, + auto_cls=AutoModelForImageTextToText, + aphrodite_output_post_proc=model_utils.qwen2_aphrodite_to_hf_output, + image_size_factors=[(0.25,), (0.25, 0.25, 0.25), (0.25, 0.2, 0.15)], + marks=[pytest.mark.core_model, pytest.mark.cpu_model], + ), + "qwen2_5_omni": VLMTestInfo( + models=["Qwen/Qwen2.5-Omni-3B"], + test_type=(VLMTestType.IMAGE, VLMTestType.MULTI_IMAGE, VLMTestType.VIDEO), + prompt_formatter=lambda img_prompt: f"<|im_start|>User\n{img_prompt}<|im_end|>\n<|im_start|>assistant\n", # noqa: E501 + img_idx_to_prompt=lambda idx: "<|vision_bos|><|IMAGE|><|vision_eos|>", + video_idx_to_prompt=lambda idx: "<|vision_bos|><|VIDEO|><|vision_eos|>", + max_model_len=4096, + max_num_seqs=2, + num_logprobs=6 if current_platform.is_cpu() else 5, + auto_cls=AutoModelForTextToWaveform, + aphrodite_output_post_proc=model_utils.qwen2_aphrodite_to_hf_output, + patch_hf_runner=model_utils.qwen2_5_omni_patch_hf_runner, + image_size_factors=[(0.25,), (0.25, 0.25, 0.25), (0.25, 0.2, 0.15)], + marks=[pytest.mark.core_model, pytest.mark.cpu_model], + ), + "qwen3_vl": VLMTestInfo( + models=["Qwen/Qwen3-VL-4B-Instruct"], + test_type=( + VLMTestType.IMAGE, + VLMTestType.MULTI_IMAGE, + VLMTestType.VIDEO, + ), + enforce_eager=False, + needs_video_metadata=True, + prompt_formatter=lambda img_prompt: f"<|im_start|>User\n{img_prompt}<|im_end|>\n<|im_start|>assistant\n", # noqa: E501 + img_idx_to_prompt=lambda idx: "<|vision_start|><|image_pad|><|vision_end|>", # noqa: E501 + video_idx_to_prompt=lambda idx: "<|vision_start|><|video_pad|><|vision_end|>", # noqa: E501 + max_model_len=4096, + max_num_seqs=2, + num_logprobs=20, + auto_cls=AutoModelForImageTextToText, + aphrodite_output_post_proc=model_utils.qwen2_aphrodite_to_hf_output, + patch_hf_runner=model_utils.qwen3_vl_patch_hf_runner, + image_size_factors=[(0.25,), (0.25, 0.25, 0.25), (0.25, 0.2, 0.15)], + marks=[ + pytest.mark.core_model, + ], + aphrodite_runner_kwargs={"attention_backend": "TRITON_ATTN"} if current_platform.is_rocm() else {}, + ), + "ultravox": VLMTestInfo( + models=["fixie-ai/ultravox-v0_5-llama-3_2-1b"], + test_type=VLMTestType.AUDIO, + prompt_formatter=lambda audio_prompt: f"<|begin_of_text|><|start_header_id|>user<|end_header_id|>\n\n{audio_prompt}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n", # noqa: E501 + audio_idx_to_prompt=lambda idx: "<|audio|>", + max_model_len=4096, + max_num_seqs=2, + auto_cls=AutoModel, + hf_output_post_proc=model_utils.ultravox_trunc_hf_output, + marks=[ + pytest.mark.core_model, + pytest.mark.cpu_model, + # TODO: Remove skip once model has been upstreamed to Transformers + pytest.mark.skip(reason="Custom model code is not compatible with Transformers v5"), + ], + ), + #### Transformers fallback to test + ## To reduce test burden, we only test batching arbitrary image size + # Dynamic image length and number of patches + "llava-onevision-transformers": VLMTestInfo( + models=["llava-hf/llava-onevision-qwen2-0.5b-ov-hf"], + test_type=VLMTestType.IMAGE, + prompt_formatter=lambda vid_prompt: f"<|im_start|>user\n{vid_prompt}<|im_end|>\n<|im_start|>assistant\n", # noqa: E501 + max_model_len=16384, + hf_model_kwargs=model_utils.llava_onevision_hf_model_kwargs("llava-hf/llava-onevision-qwen2-0.5b-ov-hf"), + auto_cls=AutoModelForImageTextToText, + aphrodite_output_post_proc=model_utils.llava_onevision_aphrodite_to_hf_output, + image_size_factors=[(0.25, 0.5, 1.0)], + aphrodite_runner_kwargs={ + "model_impl": "transformers", + "default_torch_num_threads": 1, + }, + marks=[pytest.mark.core_model], + ), + # Gemma3 has bidirectional mask on images + "gemma3-transformers": VLMTestInfo( + models=["google/gemma-3-4b-it"], + test_type=VLMTestType.IMAGE, + prompt_formatter=lambda vid_prompt: f"<'user\n{vid_prompt}\nmodel\n", # noqa: E501 + max_model_len=4096, + auto_cls=AutoModelForImageTextToText, + aphrodite_output_post_proc=model_utils.gemma3_aphrodite_to_hf_output, + image_size_factors=[(0.25, 0.5, 1.0)], + aphrodite_runner_kwargs={ + "model_impl": "transformers", + }, + marks=[ + pytest.mark.core_model, + *([large_gpu_mark(min_gb=80)] if current_platform.is_rocm() else []), + ], + ), + "idefics3-transformers": VLMTestInfo( + models=["HuggingFaceTB/SmolVLM-256M-Instruct"], + test_type=(VLMTestType.IMAGE, VLMTestType.MULTI_IMAGE), + prompt_formatter=lambda img_prompt: f"<|begin_of_text|>User:{img_prompt}\nAssistant:", # noqa: E501 + img_idx_to_prompt=lambda idx: "", + max_model_len=8192, + max_num_seqs=2, + auto_cls=AutoModelForImageTextToText, + hf_output_post_proc=model_utils.idefics3_trunc_hf_output, + image_size_factors=[(0.25, 0.5, 1.0)], + aphrodite_runner_kwargs={ + "model_impl": "transformers", + }, + marks=[pytest.mark.core_model], + ), + # Pixel values from processor are not 4D or 5D arrays + "qwen2_5_vl-transformers": VLMTestInfo( + models=["Qwen/Qwen2.5-VL-3B-Instruct"], + test_type=VLMTestType.IMAGE, + prompt_formatter=lambda img_prompt: f"<|im_start|>User\n{img_prompt}<|im_end|>\n<|im_start|>assistant\n", # noqa: E501 + img_idx_to_prompt=lambda idx: "<|vision_start|><|image_pad|><|vision_end|>", + max_model_len=4096, + max_num_seqs=2, + auto_cls=AutoModelForImageTextToText, + aphrodite_output_post_proc=model_utils.qwen2_aphrodite_to_hf_output, + image_size_factors=[(0.25, 0.2, 0.15)], + aphrodite_runner_kwargs={ + "model_impl": "transformers", + # TODO: [ROCm] Revert this once issue #30167 is resolved + **( + { + "mm_processor_kwargs": { + "min_pixels": 256 * 28 * 28, + "max_pixels": 1280 * 28 * 28, + }, + } + if current_platform.is_rocm() + else {} + ), + }, + marks=[large_gpu_mark(min_gb=80 if current_platform.is_rocm() else 32)], + ), + #### Extended model tests + "aria": VLMTestInfo( + models=["rhymes-ai/Aria"], + test_type=(VLMTestType.IMAGE, VLMTestType.MULTI_IMAGE), + prompt_formatter=lambda img_prompt: f"<|im_start|>user\n{img_prompt}<|im_end|>\n<|im_start|>assistant\n ", # noqa: E501 + img_idx_to_prompt=lambda idx: "<|img|>\n", + max_model_len=4096, + max_num_seqs=2, + auto_cls=AutoModelForImageTextToText, + single_image_prompts=IMAGE_ASSETS.prompts( + { + "stop_sign": "Please describe the image shortly.", + "cherry_blossom": "Please infer the season with reason.", + } + ), + multi_image_prompt="Describe the two images shortly.", + stop_str=["<|im_end|>"], + image_size_factors=[(0.10, 0.15)], + max_tokens=64, + marks=[ + pytest.mark.skip( + reason="Aria needs to update for latest transformers, " + "must have a vision_processor.py." + "An issue has been filed:" + "https://huggingface.co/rhymes-ai/Aria/discussions/23" + ), + large_gpu_mark(min_gb=64), + ], + ), + "blip2": VLMTestInfo( + models=["Salesforce/blip2-opt-2.7b"], + test_type=VLMTestType.IMAGE, + prompt_formatter=lambda img_prompt: f"Question: {img_prompt} Answer:", + img_idx_to_prompt=lambda idx: "", + auto_cls=AutoModelForImageTextToText, + aphrodite_output_post_proc=model_utils.blip2_aphrodite_to_hf_output, + # FIXME: https://github.com/huggingface/transformers/pull/38510 + marks=[pytest.mark.skip("Model is broken")], + ), + "chameleon": VLMTestInfo( + models=["facebook/chameleon-7b"], + test_type=VLMTestType.IMAGE, + prompt_formatter=lambda img_prompt: f"USER: {img_prompt}\nASSISTANT:", + max_model_len=4096, + max_num_seqs=2, + auto_cls=AutoModelForImageTextToText, + # For chameleon, we only compare the sequences + aphrodite_output_post_proc=lambda aphrodite_output, model: aphrodite_output[:2], + hf_output_post_proc=lambda hf_output, model: hf_output[:2], + comparator=check_outputs_equal, + max_tokens=8, + dtype="bfloat16", + ), + "cosmos3": VLMTestInfo( + models=["nvidia/Cosmos3-Nano"], + test_type=( + VLMTestType.IMAGE, + VLMTestType.MULTI_IMAGE, + VLMTestType.VIDEO, + ), + enforce_eager=False, + needs_video_metadata=True, + prompt_formatter=lambda img_prompt: f"<|im_start|>User\n{img_prompt}<|im_end|>\n<|im_start|>assistant\n", # noqa: E501 + img_idx_to_prompt=lambda idx: "<|vision_start|><|image_pad|><|vision_end|>", # noqa: E501 + video_idx_to_prompt=lambda idx: "<|vision_start|><|video_pad|><|vision_end|>", # noqa: E501 + max_model_len=4096, + max_num_seqs=2, + num_logprobs=20, + auto_cls=AutoModelForImageTextToText, + aphrodite_output_post_proc=model_utils.qwen2_aphrodite_to_hf_output, + patch_hf_runner=model_utils.qwen3_vl_patch_hf_runner, + image_size_factors=[(0.25,), (0.25, 0.25, 0.25), (0.25, 0.2, 0.15)], + marks=[ + pytest.mark.skip( + reason="Transformers has no cosmos3_omni mapping, so the HF " + "reference runner cannot load the checkpoint." + ) + ], + ), + "deepseek_vl_v2": VLMTestInfo( + models=["Isotr0py/deepseek-vl2-tiny"], # model repo using dynamic module + test_type=(VLMTestType.IMAGE, VLMTestType.MULTI_IMAGE), + prompt_formatter=lambda img_prompt: f"<|User|>: {img_prompt}\n\n<|Assistant|>: ", # noqa: E501 + max_model_len=4096, + max_num_seqs=2, + single_image_prompts=IMAGE_ASSETS.prompts( + { + "stop_sign": "\nWhat's the content in the center of the image?", + "cherry_blossom": "\nPlease infer the season with reason in details.", # noqa: E501 + } + ), + multi_image_prompt="image_1:\nimage_2:\nWhich image can we see the car and the tower?", # noqa: E501 + patch_hf_runner=model_utils.deepseekvl2_patch_hf_runner, + hf_output_post_proc=model_utils.deepseekvl2_trunc_hf_output, + stop_str=["<｜end▁of▁sentence｜>", "<｜begin▁of▁sentence｜>"], + image_size_factors=[(1.0,), (1.0, 1.0, 1.0), (0.1, 0.5, 1.0)], + ), + "gemma3": VLMTestInfo( + models=["google/gemma-3-4b-it"], + test_type=(VLMTestType.IMAGE, VLMTestType.MULTI_IMAGE), + prompt_formatter=lambda img_prompt: f"user\n{img_prompt}\nmodel\n", # noqa: E501 + single_image_prompts=IMAGE_ASSETS.prompts( + { + "stop_sign": "What's the content in the center of the image?", # noqa: E501 + "cherry_blossom": "What is the season?", + } + ), + multi_image_prompt="Describe the two images in detail.", # noqa: E501 + max_model_len=4096, + max_num_seqs=2, + auto_cls=AutoModelForImageTextToText, + aphrodite_runner_kwargs={"mm_processor_kwargs": {"do_pan_and_scan": True}}, + patch_hf_runner=model_utils.gemma3_patch_hf_runner, + ), + "gemma4": VLMTestInfo( + models=["google/gemma-4-E2B-it"], + test_type=(VLMTestType.IMAGE, VLMTestType.MULTI_IMAGE), + prompt_formatter=lambda img_prompt: f"<|turn>user\n{img_prompt}\n<|turn>model\n", # noqa: E501 + single_image_prompts=IMAGE_ASSETS.prompts( + { + "stop_sign": "<|image|>What's the content in the center of the image?", # noqa: E501 + "cherry_blossom": "<|image|>What is the season?", + } + ), + multi_image_prompt="<|image|><|image|>Describe the two images in detail.", # noqa: E501 + max_model_len=4096, + max_num_seqs=2, + auto_cls=AutoModelForImageTextToText, + aphrodite_runner_kwargs={"limit_mm_per_prompt": {"image": 4}}, + ), + "granite_vision": VLMTestInfo( + models=["ibm-granite/granite-vision-3.3-2b"], + test_type=(VLMTestType.IMAGE), + prompt_formatter=lambda img_prompt: f"<|user|>\n{img_prompt}\n<|assistant|>\n", + max_model_len=8192, + auto_cls=AutoModelForImageTextToText, + aphrodite_output_post_proc=model_utils.llava_image_aphrodite_to_hf_output, + ), + "glm4v": VLMTestInfo( + models=["zai-org/glm-4v-9b"], + test_type=VLMTestType.IMAGE, + prompt_formatter=lambda img_prompt: f"<|user|>\n{img_prompt}<|assistant|>", + single_image_prompts=IMAGE_ASSETS.prompts( + { + "stop_sign": "<|begin_of_image|><|endoftext|><|end_of_image|>What's the content in the center of the image?", # noqa: E501 + "cherry_blossom": "<|begin_of_image|><|endoftext|><|end_of_image|>What is the season?", # noqa: E501 + } + ), + max_model_len=2048, + max_num_seqs=2, + get_stop_token_ids=lambda tok: [151329, 151336, 151338], + patch_hf_runner=model_utils.glm4v_patch_hf_runner, + # The image embeddings match with HF but the outputs of the language + # decoder are only consistent up to 2 decimal places. + # So, we need to reduce the number of tokens for the test to pass. + max_tokens=8, + num_logprobs=10, + auto_cls=AutoModelForCausalLM, + marks=[ + pytest.mark.skip( + reason="The code for this model has a bug." + "Please see the issue here:" + "https://huggingface.co/zai-org/glm-4v-9b/discussions/46." + ), + large_gpu_mark(min_gb=32), + ], + ), + "glm4_1v": VLMTestInfo( + models=["zai-org/GLM-4.1V-9B-Thinking"], + test_type=(VLMTestType.IMAGE, VLMTestType.MULTI_IMAGE), + prompt_formatter=lambda img_prompt: f"[gMASK]<|user|>\n{img_prompt}<|assistant|>\n", # noqa: E501 + img_idx_to_prompt=lambda idx: "<|begin_of_image|><|image|><|end_of_image|>", + video_idx_to_prompt=lambda idx: "<|begin_of_video|><|video|><|end_of_video|>", + max_model_len=2048, + max_num_seqs=2, + get_stop_token_ids=lambda tok: [151329, 151336, 151338], + num_logprobs=10, + image_size_factors=[(0.25,), (0.25, 0.25, 0.25), (0.25, 0.2, 0.15)], + auto_cls=AutoModelForImageTextToText, + marks=[large_gpu_mark(min_gb=32)], + ), + "glm4_1v-video": VLMTestInfo( + models=["zai-org/GLM-4.1V-9B-Thinking"], + # GLM4.1V require include video metadata for input + test_type=VLMTestType.CUSTOM_INPUTS, + prompt_formatter=lambda vid_prompt: f"[gMASK]<|user|>\n{vid_prompt}<|assistant|>\n", # noqa: E501 + max_model_len=4096, + max_num_seqs=2, + auto_cls=AutoModelForImageTextToText, + patch_hf_runner=model_utils.glm4_1v_patch_hf_runner, + custom_test_opts=[ + CustomTestOptions( + inputs=custom_inputs.video_with_metadata_glm4_1v(), + limit_mm_per_prompt={"video": 1}, + ) + ], + marks=[large_gpu_mark(min_gb=32)], + ), + "glm_ocr": VLMTestInfo( + models=["zai-org/GLM-OCR"], + test_type=(VLMTestType.IMAGE, VLMTestType.MULTI_IMAGE), + prompt_formatter=lambda img_prompt: f"[gMASK]<|user|>\n{img_prompt}<|assistant|>\n", # noqa: E501 + img_idx_to_prompt=lambda idx: "<|begin_of_image|><|image|><|end_of_image|>", + video_idx_to_prompt=lambda idx: "<|begin_of_video|><|video|><|end_of_video|>", + max_model_len=2048, + max_num_seqs=2, + get_stop_token_ids=lambda tok: [151329, 151336, 151338], + num_logprobs=10, + image_size_factors=[(0.25,), (0.25, 0.25, 0.25), (0.25, 0.2, 0.15)], + auto_cls=AutoModelForImageTextToText, + marks=[ + pytest.mark.skip( + reason="This test fails on both AMD and NV" + "hardware. please see the issue:" + "https://github.com/vllm-project/vllm/issues/42016" + ), + large_gpu_mark(min_gb=32), + ], + ), + "granite4_vision": VLMTestInfo( + models=["ibm-granite/granite-vision-4.1-4b"], + test_type=(VLMTestType.IMAGE), + prompt_formatter=lambda img_prompt: f"<|user|>\n{img_prompt}\n<|assistant|>\n", + max_model_len=8192, + auto_cls=AutoModelForImageTextToText, + aphrodite_output_post_proc=_granite4_vision_aphrodite_to_hf_output, + image_size_factors=[(1.0,)], + aphrodite_runner_kwargs={ + "enable_lora": True, + "max_lora_rank": 256, + "default_mm_loras": {"image": "ibm-granite/granite-vision-4.1-4b"}, + }, + ), + "h2ovl": VLMTestInfo( + models=[ + "h2oai/h2ovl-mississippi-800m", + "h2oai/h2ovl-mississippi-2b", + ], + test_type=(VLMTestType.IMAGE, VLMTestType.MULTI_IMAGE), + prompt_formatter=lambda img_prompt: f"<|prompt|>{img_prompt}<|end|><|answer|>", + single_image_prompts=IMAGE_ASSETS.prompts( + { + "stop_sign": "\nWhat's the content in the center of the image?", + "cherry_blossom": "\nWhat is the season?", + } + ), + multi_image_prompt="Image-1: \nImage-2: \nDescribe the two images in short.", # noqa: E501 + max_model_len=8192, + use_tokenizer_eos=True, + num_logprobs=10, + patch_hf_runner=model_utils.h2ovl_patch_hf_runner, + ), + "idefics3": VLMTestInfo( + models=["HuggingFaceTB/SmolVLM-256M-Instruct"], + test_type=(VLMTestType.IMAGE, VLMTestType.MULTI_IMAGE), + prompt_formatter=lambda img_prompt: f"<|begin_of_text|>User:{img_prompt}\nAssistant:", # noqa: E501 + img_idx_to_prompt=lambda idx: "", + max_model_len=8192, + max_num_seqs=2, + auto_cls=AutoModelForImageTextToText, + hf_output_post_proc=model_utils.idefics3_trunc_hf_output, + ), + "intern_vl": VLMTestInfo( + models=[ + "OpenGVLab/InternVL2-1B", + "OpenGVLab/InternVL2-2B", + ], + test_type=(VLMTestType.IMAGE, VLMTestType.MULTI_IMAGE), + prompt_formatter=lambda img_prompt: f"<|im_start|>User\n{img_prompt}<|im_end|>\n<|im_start|>Assistant\n", # noqa: E501 + single_image_prompts=IMAGE_ASSETS.prompts( + { + "stop_sign": "\nWhat's the content in the center of the image?", + "cherry_blossom": "\nWhat is the season?", + } + ), + multi_image_prompt="Image-1: \nImage-2: \nDescribe the two images in short.", # noqa: E501 + max_model_len=4096, + use_tokenizer_eos=True, + patch_hf_runner=model_utils.internvl_patch_hf_runner, + # TODO: Remove skip once model has been upstreamed to Transformers + marks=[pytest.mark.skip(reason="Custom model code tries to access data from meta-tensor")], + ), + "intern_vl-video": VLMTestInfo( + models=[ + "OpenGVLab/InternVL3-1B", + ], + test_type=VLMTestType.VIDEO, + prompt_formatter=lambda img_prompt: f"<|im_start|>User\n{img_prompt}<|im_end|>\n<|im_start|>Assistant\n", # noqa: E501 + video_idx_to_prompt=lambda idx: "

(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) { - 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 - // for all ranks, ensuring bitwise identical results - auto dp = *_dp; - barrier_at_start(sg, self_sg, rank); - // do the actual reduction - for (int idx = blockIdx.x * blockDim.x + threadIdx.x; idx < size; - idx += gridDim.x * blockDim.x) { - ((P*)result)[idx] = packed_reduce((const P**)&dp.ptrs[0], idx); - } - barrier_at_end(sg, self_sg, rank); -} - -template -DINLINE P* get_tmp_buf(Signal* sg) { - return (P*)(((Signal*)sg) + 1); -} - -template -__global__ void __launch_bounds__(512, 1) - cross_device_reduce_2stage(RankData* _dp, RankSignals sg, Signal* self_sg, - T* __restrict__ result, int rank, int size) { - int tid = blockIdx.x * blockDim.x + threadIdx.x; - int stride = gridDim.x * blockDim.x; - using P = typename packed_t::P; - using A = typename packed_t::A; - int part = size / ngpus; - int start = rank * part; - int end = rank == ngpus - 1 ? size : start + part; - int largest_part = part + size % ngpus; - const P* ptrs[ngpus]; - P* tmps[ngpus]; -#pragma unroll - for (int i = 0; i < ngpus; i++) { - int target = (rank + i) % ngpus; - ptrs[i] = (const P*)_dp->ptrs[target]; - tmps[i] = get_tmp_buf

(sg.signals[target]); - } - auto tmp_out = tmps[0]; - barrier_at_start(sg, self_sg, rank); - - // stage 1: reduce scatter - for (int idx = start + tid; idx < end; idx += stride) { - tmp_out[idx - start] = packed_reduce(ptrs, idx); - } - barrier_at_end(sg, self_sg, rank); - - // stage 2: allgather. Note: it's important to match the tid between - // the two stages, because visibility across devices is only guaranteed - // between threads that have the same tid. If thread i computes the sum of - // start + i in the first stage, then thread i also gathers start + i from - // all ranks. - - for (int idx = tid; idx < largest_part; idx += stride) { -#pragma unroll - for (int i = 0; i < ngpus; i++) { - int gather_from_rank = ((rank + i) % ngpus); - if (gather_from_rank == ngpus - 1 || idx < part) { - int dst_idx = gather_from_rank * part + idx; - ((P*)result)[dst_idx] = tmps[i][idx]; - } - } - } -} - -using IPC_KEY = std::array; -static_assert(sizeof(IPC_KEY) == sizeof(cudaIpcMemHandle_t)); -static_assert(alignof(IPC_KEY) == alignof(cudaIpcMemHandle_t)); - -class CustomAllreduce { - public: - int rank_; - int world_size_; - // Full NVLink or xGMI connection between GPUs. - bool fully_connected_; - - RankSignals sg_; - // Stores a map from a pointer to its peer pointers from all ranks. - std::unordered_map buffers_; - Signal* self_sg_; - - // Stores rank data from all ranks. This is mainly for cuda graph purposes. - // For cuda graph to work, all kernel arguments must be fixed during graph - // capture time. However, the peer pointers are not known during graph - // capture time. Therefore, during capture, we increment the rank data - // pointer and use that as the argument to the kernel. The kernel arguments - // are stored in graph_unreg_buffers_. The actual peer pointers will be - // filled in at the memory pointed to by the pointers in - // graph_unreg_buffers_ when the IPC handles are exchanged between ranks. - // - // The overall process looks like this: - // 1. Graph capture. - // 2. Each rank obtains the IPC handles for each addresses used during cuda - // graph capture using get_graph_buffer_ipc_meta. - // 3. (In Python) all gather the IPC handles. - // 4. Obtain the peer pointers by opening the IPC handles, and store them in - // the rank data array at corresponding positions. - RankData *d_rank_data_base_, *d_rank_data_end_; - std::vector graph_unreg_buffers_; - // a map from IPC handles to opened IPC pointers - std::map ipc_handles_; - - /** - * Signals are an array of ipc-enabled buffers from all ranks. - * For each of the buffer, the layout is as follows: - * | -- sizeof(Signal) -- | ------ a few MB ----- | - * The first section is for allreduce synchronization, and the second - * section is for storing the intermediate results required by some - * allreduce algos. - * - * Note: this class does not own any device memory. Any required buffers - * are passed in from the constructor. - */ - CustomAllreduce(Signal** signals, void* rank_data, size_t rank_data_sz, - int rank, int world_size, bool fully_connected = true) - : rank_(rank), - world_size_(world_size), - fully_connected_(fully_connected), - self_sg_(signals[rank]), - d_rank_data_base_(reinterpret_cast(rank_data)), - d_rank_data_end_(d_rank_data_base_ + rank_data_sz / sizeof(RankData)) { - for (int i = 0; i < world_size_; i++) { - sg_.signals[i] = signals[i]; - } - } - - char* open_ipc_handle(const void* ipc_handle) { - auto [it, new_handle] = - ipc_handles_.insert({*((IPC_KEY*)ipc_handle), nullptr}); - if (new_handle) { - char* ipc_ptr; - CUDACHECK(cudaIpcOpenMemHandle((void**)&ipc_ptr, - *((const cudaIpcMemHandle_t*)ipc_handle), - cudaIpcMemLazyEnablePeerAccess)); - it->second = ipc_ptr; - } - return it->second; - } - - std::pair> get_graph_buffer_ipc_meta() { - auto num_buffers = graph_unreg_buffers_.size(); - auto handle_sz = sizeof(cudaIpcMemHandle_t); - std::string handles(handle_sz * num_buffers, static_cast(0)); - std::vector offsets(num_buffers); - for (int i = 0; i < num_buffers; i++) { - auto ptr = graph_unreg_buffers_[i]; - void* base_ptr; - // note: must share the base address of each allocation, or we get wrong - // address - if (cuPointerGetAttribute(&base_ptr, rangeStartAddrAttr, - (CUdeviceptr)ptr) != CUDA_SUCCESS) - throw std::runtime_error("failed to get pointer attr"); - CUDACHECK(cudaIpcGetMemHandle( - (cudaIpcMemHandle_t*)&handles[i * handle_sz], base_ptr)); - offsets[i] = ((char*)ptr) - ((char*)base_ptr); - } - return std::make_pair(handles, offsets); - } - - void check_rank_data_capacity(size_t num = 1) { - if (d_rank_data_base_ + num > d_rank_data_end_) - throw std::runtime_error( - "Rank data buffer is overflowed by " + - std::to_string(d_rank_data_base_ + num - d_rank_data_end_)); - } - - /** - * Register already-shared IPC pointers. - */ - void register_buffer(void** ptrs) { - check_rank_data_capacity(); - RankData data; - for (int i = 0; i < world_size_; i++) { - data.ptrs[i] = ptrs[i]; - } - auto d_data = d_rank_data_base_++; - CUDACHECK( - cudaMemcpy(d_data, &data, sizeof(RankData), cudaMemcpyHostToDevice)); - buffers_[ptrs[rank_]] = d_data; - } - - // Note: when registering graph buffers, we intentionally choose to not - // deduplicate the addresses. That means if the allocator reuses some - // addresses, they will be registered again. This is to account for the - // remote possibility of different allocation patterns between ranks. For - // example, rank 1 may get the same input address for the second allreduce, - // but rank 2 got a different address. IPC handles have internal reference - // counting mechanism so overhead should be small. - void register_graph_buffers( - const std::vector& handles, - const std::vector>& offsets) { - auto num_buffers = graph_unreg_buffers_.size(); - check_rank_data_capacity(num_buffers); - std::vector rank_data(num_buffers); - for (int i = 0; i < num_buffers; i++) { - auto self_ptr = graph_unreg_buffers_[i]; - auto& rd = rank_data[i]; - for (int j = 0; j < world_size_; j++) { - if (j != rank_) { - char* handle = - open_ipc_handle(&handles[j][i * sizeof(cudaIpcMemHandle_t)]); - handle += offsets[j][i]; - rd.ptrs[j] = handle; - } else { - rd.ptrs[j] = self_ptr; - } - } - } - CUDACHECK(cudaMemcpy(d_rank_data_base_, rank_data.data(), - sizeof(RankData) * num_buffers, - cudaMemcpyHostToDevice)); - d_rank_data_base_ += num_buffers; - graph_unreg_buffers_.clear(); - } - - /** - * Performs allreduce, assuming input has already been registered. - * - * Block and grid default configs are results after careful grid search. - * Using 36 blocks give the best or close to the best runtime on the devices - * I tried: A100, A10, A30, T4, V100. You'll notice that NCCL kernels also - * only take a small amount of SMs. Not quite sure the underlying reason, - * but my guess is that too many SMs will cause contention on NVLink bus. - */ - template - void allreduce(cudaStream_t stream, T* input, T* output, int size, - int threads = 512, int block_limit = defaultBlockLimit) { - auto d = packed_t::P::size; - if (size % d != 0) - throw std::runtime_error( - "custom allreduce currently requires input length to be multiple " - "of " + - std::to_string(d)); - if (block_limit > kMaxBlocks) - throw std::runtime_error("max supported block limit is " + - std::to_string(kMaxBlocks) + ". Got " + - std::to_string(block_limit)); - - RankData* ptrs; - cudaStreamCaptureStatus status; - CUDACHECK(cudaStreamIsCapturing(stream, &status)); - if (status == cudaStreamCaptureStatusActive) { - ptrs = d_rank_data_base_ + graph_unreg_buffers_.size(); - graph_unreg_buffers_.push_back(input); - } else { - auto it = buffers_.find(input); - if (it == buffers_.end()) - throw std::runtime_error( - "buffer address " + - std::to_string(reinterpret_cast(input)) + - " is not registered!"); - ptrs = it->second; - } - - size /= d; - auto bytes = size * sizeof(typename packed_t::P); - int blocks = std::min(block_limit, (size + threads - 1) / threads); - - // Check environment variable once - const char* env_algo = std::getenv("APHRODITE_CUSTOM_ALLREDUCE_ALGO"); - bool force_1stage = false; - bool force_2stage = false; - if (env_algo != nullptr) { - if (std::strcmp(env_algo, "1stage") == 0 || - std::strcmp(env_algo, "oneshot") == 0) { - force_1stage = true; - } else if (std::strcmp(env_algo, "2stage") == 0 || - std::strcmp(env_algo, "twoshot") == 0) { - force_2stage = true; - } else { - throw std::runtime_error( - "Invalid APHRODITE_CUSTOM_ALLREDUCE_ALGO: " + - std::string(env_algo) + - ". Valid values: 1stage, oneshot, 2stage, twoshot"); - } - } - -#define KL(ngpus, name) \ - name<<>>(ptrs, sg_, self_sg_, output, \ - rank_, size); -#define REDUCE_CASE(ngpus) \ - case ngpus: { \ - if (force_1stage) { \ - KL(ngpus, cross_device_reduce_1stage); \ - } else if (force_2stage) { \ - KL(ngpus, cross_device_reduce_2stage); \ - } else { \ - if (world_size_ == 2) { \ - KL(ngpus, cross_device_reduce_1stage); \ - } else if (fully_connected_) { \ - if ((world_size_ <= 4 && bytes < 512 * 1024) || \ - (world_size_ <= 8 && bytes < 256 * 1024)) { \ - KL(ngpus, cross_device_reduce_1stage); \ - } else { \ - KL(ngpus, cross_device_reduce_2stage); \ - } \ - } \ - } \ - break; \ - } - - switch (world_size_) { - REDUCE_CASE(2) - REDUCE_CASE(4) - REDUCE_CASE(6) - REDUCE_CASE(8) - default: - throw std::runtime_error( - "custom allreduce only supports num gpus in (2,4,6,8). Actual " - "num " - "gpus = " + - std::to_string(world_size_)); - } -#undef REDUCE_CASE -#undef KL - } - - ~CustomAllreduce() { - for (auto [_, ptr] : ipc_handles_) { - CUDACHECK(cudaIpcCloseMemHandle(ptr)); - } - } -}; - -/** - * To inspect PTX/SASS, copy paste this header file to compiler explorer and - add a template instantiation: - * template void aphrodite::CustomAllreduce::allreduce(cudaStream_t, half - *, half *, int, int, int); -*/ -} // namespace aphrodite \ No newline at end of file diff --git a/csrc/libtorch_stable/activation_kernels.cu.orig b/csrc/libtorch_stable/activation_kernels.cu.orig deleted file mode 100644 index 6e3e6860e6..0000000000 --- a/csrc/libtorch_stable/activation_kernels.cu.orig +++ /dev/null @@ -1,724 +0,0 @@ -#include -#include - -#include - -#include "../cuda_compat.h" -#include "cuda_vec_utils.cuh" -#include "dispatch_utils.h" -#include "torch_utils.h" - -namespace aphrodite { - -// `alpha` and `beta` are applied to opposite operands: -// - alpha lives INSIDE the activation (the activated half): the gated -// activation computes act_half * sigmoid(alpha * act_half). -// - beta is added to the OTHER (non-activated) half before the multiply. -// So the result is always ACT(act_half, alpha) * (other_half + beta). -// Which half is which depends on `act_first` (see below). Defaults -// alpha=1.0, beta=0.0 reproduce the plain SwiGLU/GeGLU behavior. -template -__device__ __forceinline__ scalar_t compute(const scalar_t& x, - const scalar_t& y, - const float limit, - const float alpha, - const float beta) { - if constexpr (act_first) { - scalar_t gate = x; - scalar_t up = y; - if constexpr (HAS_CLAMP) { - gate = (scalar_t)fminf((float)gate, limit); - up = (scalar_t)fmaxf(fminf((float)up, limit), -limit); - } - // act_first: gate is the activated half -> alpha applies to gate; - // beta is added to up (the non-activated half). - return (scalar_t)(ACT_FN(gate, alpha) * ((float)up + beta)); - } else { - scalar_t gate = x; - scalar_t up = y; - if constexpr (HAS_CLAMP) { - gate = (scalar_t)fmaxf(fminf((float)gate, limit), -limit); - up = (scalar_t)fminf((float)up, limit); - } - // !act_first: up is the activated half -> alpha applies to up; - // beta is added to gate (the non-activated half). - return (scalar_t)(((float)gate + beta) * ACT_FN(up, alpha)); - } -} - -template -__device__ __forceinline__ packed_t packed_compute(const packed_t& x, - const packed_t& y, - const float limit, - const float alpha, - const float beta) { - if constexpr (act_first) { - packed_t gate = x; - packed_t up = y; - float2 u = cast_to_float2(up); - if constexpr (HAS_CLAMP) { - float2 g = cast_to_float2(gate); - g.x = fminf(g.x, limit); - g.y = fminf(g.y, limit); - u.x = fmaxf(fminf(u.x, limit), -limit); - u.y = fmaxf(fminf(u.y, limit), -limit); - gate = cast_to_packed(g); - } - // act_first: gate is the activated half -> alpha applies to gate; - // beta is added to up (the non-activated half). - float2 activated = cast_to_float2(PACKED_ACT_FN(gate, alpha)); - activated.x *= u.x + beta; - activated.y *= u.y + beta; - return cast_to_packed(activated); - } else { - packed_t gate = x; - packed_t up = y; - float2 g = cast_to_float2(gate); - if constexpr (HAS_CLAMP) { - float2 u = cast_to_float2(up); - g.x = fmaxf(fminf(g.x, limit), -limit); - g.y = fmaxf(fminf(g.y, limit), -limit); - u.x = fminf(u.x, limit); - u.y = fminf(u.y, limit); - up = cast_to_packed(u); - } - // !act_first: up is the activated half -> alpha applies to up; - // beta is added to gate (the non-activated half). - float2 activated = cast_to_float2(PACKED_ACT_FN(up, alpha)); - activated.x *= g.x + beta; - activated.y *= g.y + beta; - return cast_to_packed(activated); - } -} - -// Activation and gating kernel template. -template -__global__ void act_and_mul_kernel( - scalar_t* __restrict__ out, // [..., d] - const scalar_t* __restrict__ input, // [..., 2, d] - const int d, const float limit, const float alpha, const float beta) { - const scalar_t* x_ptr = input + blockIdx.x * 2 * d; - const scalar_t* y_ptr = x_ptr + d; - scalar_t* out_ptr = out + blockIdx.x * d; - - if constexpr (use_vec) { - using cuda_t = typename CUDATypeConverter::Type; - using pvec_t = PackedVec; - - const pvec_t* x_vec = reinterpret_cast(x_ptr); - const pvec_t* y_vec = reinterpret_cast(y_ptr); - pvec_t* out_vec = reinterpret_cast(out_ptr); - const int num_vecs = d / 2 / pvec_t::NUM_ELTS; - - for (int i = threadIdx.x; i < num_vecs; i += blockDim.x) { - pvec_t x, y; - if constexpr (use_256b) { - ld256(x, &x_vec[i]); - ld256(y, &y_vec[i]); - } else { - ld128(x, &x_vec[i]); - ld128(y, &y_vec[i]); - } -#pragma unroll - for (int j = 0; j < pvec_t::NUM_ELTS; j++) { - x.elts[j] = - packed_compute( - x.elts[j], y.elts[j], limit, alpha, beta); - } - if constexpr (use_256b) { - st256(x, &out_vec[i]); - } else { - st128(x, &out_vec[i]); - } - } - } else { - // Scalar fallback for unaligned data or small d - for (int64_t idx = threadIdx.x; idx < d; idx += blockDim.x) { - const scalar_t x = APHRODITE_LDG(&x_ptr[idx]); - const scalar_t y = APHRODITE_LDG(&y_ptr[idx]); - out_ptr[idx] = compute( - x, y, limit, alpha, beta); - } - } -} - -// Gated activations take an `alpha` argument that scales the sigmoid input -// (`x * sigmoid(alpha * x)`). alpha defaults to 1.0 at all call sites, which -// is exactly SiLU; only the clamp path (silu_and_mul_with_clamp) passes a -// non-default alpha. Activations that do not use alpha simply ignore it. -template -__device__ __forceinline__ T silu_kernel(const T& x, const float alpha) { - // x * sigmoid(alpha * x) - return (T)(((float)x) / (1.0f + expf((float)-x * alpha))); -} - -template -__device__ __forceinline__ packed_t packed_silu_kernel(const packed_t& val, - const float alpha) { - // x * sigmoid(alpha * x) - float2 fval = cast_to_float2(val); - fval.x = fval.x / (1.0f + expf(-fval.x * alpha)); - fval.y = fval.y / (1.0f + expf(-fval.y * alpha)); - return cast_to_packed(fval); -} - -template -__device__ __forceinline__ T gelu_kernel(const T& x, const float /*alpha*/) { - // Equivalent to PyTorch GELU with 'none' approximation. - // Refer to: - // https://github.com/pytorch/pytorch/blob/8ac9b20d4b090c213799e81acf48a55ea8d437d6/aten/src/ATen/native/cuda/ActivationGeluKernel.cu#L36-L38 - const float f = (float)x; - constexpr float ALPHA = M_SQRT1_2; - return (T)(f * 0.5f * (1.0f + ::erf(f * ALPHA))); -} - -template -__device__ __forceinline__ packed_t packed_gelu_kernel(const packed_t& val, - const float /*alpha*/) { - // Equivalent to PyTorch GELU with 'none' approximation. - // Refer to: - // https://github.com/pytorch/pytorch/blob/8ac9b20d4b090c213799e81acf48a55ea8d437d6/aten/src/ATen/native/cuda/ActivationGeluKernel.cu#L36-L38 - constexpr float ALPHA = M_SQRT1_2; - float2 fval = cast_to_float2(val); - fval.x = fval.x * 0.5f * (1.0f + ::erf(fval.x * ALPHA)); - fval.y = fval.y * 0.5f * (1.0f + ::erf(fval.y * ALPHA)); - return cast_to_packed(fval); -} - -template -__device__ __forceinline__ T gelu_tanh_kernel(const T& x, - const float /*alpha*/) { - // Equivalent to PyTorch GELU with 'tanh' approximation. - // Refer to: - // https://github.com/pytorch/pytorch/blob/8ac9b20d4b090c213799e81acf48a55ea8d437d6/aten/src/ATen/native/cuda/ActivationGeluKernel.cu#L25-L30 - const float f = (float)x; - constexpr float BETA = M_SQRT2 * M_2_SQRTPI * 0.5f; - constexpr float KAPPA = 0.044715; - float x_cube = f * f * f; - float inner = BETA * (f + KAPPA * x_cube); - return (T)(0.5f * f * (1.0f + ::tanhf(inner))); -} - -template -__device__ __forceinline__ packed_t -packed_gelu_tanh_kernel(const packed_t& val, const float /*alpha*/) { - // Equivalent to PyTorch GELU with 'tanh' approximation. - // Refer to: - // https://github.com/pytorch/pytorch/blob/8ac9b20d4b090c213799e81acf48a55ea8d437d6/aten/src/ATen/native/cuda/ActivationGeluKernel.cu#L25-L30 - float2 fval = cast_to_float2(val); - constexpr float BETA = M_SQRT2 * M_2_SQRTPI * 0.5f; - constexpr float KAPPA = 0.044715; - - float x_cube = fval.x * fval.x * fval.x; - float inner = BETA * (fval.x + KAPPA * x_cube); - fval.x = 0.5f * fval.x * (1.0f + ::tanhf(inner)); - - x_cube = fval.y * fval.y * fval.y; - inner = BETA * (fval.y + KAPPA * x_cube); - fval.y = 0.5f * fval.y * (1.0f + ::tanhf(inner)); - return cast_to_packed(fval); -} - -} // namespace aphrodite - -// Launch activation and gating kernel. -// Use ACT_FIRST (bool) indicating whether to apply the activation function -// first. HAS_CLAMP (bool) enables pre-activation clamping: gate input is -// clamped (max only) and up input is clamped (both sides) before the -// activation function is applied. -#define LAUNCH_ACTIVATION_GATE_KERNEL(KERNEL, PACKED_KERNEL, ACT_FIRST, \ - HAS_CLAMP, LIMIT, ALPHA, BETA) \ - auto dtype = input.scalar_type(); \ - int d = input.size(-1) / 2; \ - int64_t num_tokens = input.numel() / input.size(-1); \ - if (num_tokens == 0) { \ - return; \ - } \ - dim3 grid(num_tokens); \ - int cc_major = get_device_prop()->major; \ - int support_vec = \ - (CUDA_VERSION >= 12090 && cc_major >= 10 && num_tokens > 128) \ - ? aphrodite::VecTraits::ARCH_MAX_VEC_SIZE \ - : aphrodite::VecTraits::ARCH_MAX_VEC_SIZE; \ - int vec_size = support_vec / input.element_size(); \ - const bool use_vec = (d % vec_size == 0); \ - const torch::stable::accelerator::DeviceGuard device_guard( \ - input.get_device_index()); \ - const cudaStream_t stream = get_current_cuda_stream(); \ - if (use_vec) { \ - dim3 block(std::min(d / vec_size, 1024)); \ - if (CUDA_VERSION >= 12090 && cc_major >= 10 && num_tokens > 128) { \ - APHRODITE_STABLE_DISPATCH_FLOATING_TYPES( \ - dtype, "act_and_mul_kernel", [&] { \ - aphrodite::act_and_mul_kernel< \ - scalar_t, \ - typename aphrodite::PackedTypeConverter::Type, \ - KERNEL, \ - PACKED_KERNEL< \ - typename aphrodite::PackedTypeConverter::Type>, \ - ACT_FIRST, true, HAS_CLAMP, true><<>>( \ - out.mutable_data_ptr(), \ - input.const_data_ptr(), d, LIMIT, ALPHA, BETA); \ - }); \ - } else { \ - APHRODITE_STABLE_DISPATCH_FLOATING_TYPES( \ - dtype, "act_and_mul_kernel", [&] { \ - aphrodite::act_and_mul_kernel< \ - scalar_t, \ - typename aphrodite::PackedTypeConverter::Type, \ - KERNEL, \ - PACKED_KERNEL< \ - typename aphrodite::PackedTypeConverter::Type>, \ - ACT_FIRST, true, HAS_CLAMP, false> \ - <<>>(out.mutable_data_ptr(), \ - input.const_data_ptr(), \ - d, LIMIT, ALPHA, BETA); \ - }); \ - } \ - } else { \ - dim3 block(std::min(d, 1024)); \ - APHRODITE_STABLE_DISPATCH_FLOATING_TYPES( \ - dtype, "act_and_mul_kernel", [&] { \ - aphrodite::act_and_mul_kernel< \ - scalar_t, \ - typename aphrodite::PackedTypeConverter::Type, \ - KERNEL, \ - PACKED_KERNEL< \ - typename aphrodite::PackedTypeConverter::Type>, \ - ACT_FIRST, false, HAS_CLAMP><<>>( \ - out.mutable_data_ptr(), \ - input.const_data_ptr(), d, LIMIT, ALPHA, BETA); \ - }); \ - } - -void silu_and_mul(torch::stable::Tensor& out, // [..., d] - torch::stable::Tensor& input) // [..., 2 * d] -{ - LAUNCH_ACTIVATION_GATE_KERNEL(aphrodite::silu_kernel, - aphrodite::packed_silu_kernel, true, false, - 0.0f, 1.0f, 0.0f); -} - -void silu_and_mul_clamp(torch::stable::Tensor& out, // [..., d] - torch::stable::Tensor& input, // [..., 2 * d] - double limit, double alpha, double beta) { - // out = (gate.clamp(max=limit) * sigmoid(alpha * gate.clamp(max=limit))) - // * (up.clamp(+-limit) + beta) - // alpha=1.0, beta=0.0 reduce this to silu(gate) * up. - LAUNCH_ACTIVATION_GATE_KERNEL(aphrodite::silu_kernel, - aphrodite::packed_silu_kernel, true, true, - (float)limit, (float)alpha, (float)beta); -} - -void mul_and_silu(torch::stable::Tensor& out, // [..., d] - torch::stable::Tensor& input) // [..., 2 * d] -{ - // The difference between mul_and_silu and silu_and_mul is that mul_and_silu - // applies the silu to the latter half of the input. - LAUNCH_ACTIVATION_GATE_KERNEL(aphrodite::silu_kernel, - aphrodite::packed_silu_kernel, false, false, - 0.0f, 1.0f, 0.0f); -} - -void gelu_and_mul(torch::stable::Tensor& out, // [..., d] - torch::stable::Tensor& input) // [..., 2 * d] -{ - LAUNCH_ACTIVATION_GATE_KERNEL(aphrodite::gelu_kernel, - aphrodite::packed_gelu_kernel, true, false, - 0.0f, 1.0f, 0.0f); -} - -void gelu_tanh_and_mul(torch::stable::Tensor& out, // [..., d] - torch::stable::Tensor& input) // [..., 2 * d] -{ - LAUNCH_ACTIVATION_GATE_KERNEL(aphrodite::gelu_tanh_kernel, - aphrodite::packed_gelu_tanh_kernel, true, false, - 0.0f, 1.0f, 0.0f); -} - -namespace aphrodite { - -template -__device__ __forceinline__ T fatrelu_kernel(const T& x, const float threshold) { - const float f = (float)x; - return (T)(f > threshold ? f : 0.0f); -} - -template -__device__ __forceinline__ packed_t -packed_fatrelu_kernel(const packed_t& val, const float threshold) { - float2 fval = cast_to_float2(val); - fval.x = fval.x > threshold ? fval.x : 0.0f; - fval.y = fval.y > threshold ? fval.y : 0.0f; - return cast_to_packed(fval); -} - -template -__global__ void act_and_mul_kernel_with_param( - scalar_t* __restrict__ out, const scalar_t* __restrict__ input, const int d, - const float param) { - const scalar_t* x_ptr = input + blockIdx.x * 2 * d; - const scalar_t* y_ptr = x_ptr + d; - scalar_t* out_ptr = out + blockIdx.x * d; - - if constexpr (use_vec) { - using cuda_t = typename CUDATypeConverter::Type; - using pvec_t = PackedVec; - - const pvec_t* x_vec = reinterpret_cast(x_ptr); - const pvec_t* y_vec = reinterpret_cast(y_ptr); - pvec_t* out_vec = reinterpret_cast(out_ptr); - const int num_vecs = d / 2 / pvec_t::NUM_ELTS; - - for (int i = threadIdx.x; i < num_vecs; i += blockDim.x) { - pvec_t x, y; - if constexpr (use_256b) { - ld256(x, &x_vec[i]); - ld256(y, &y_vec[i]); - } else { - ld128(x, &x_vec[i]); - ld128(y, &y_vec[i]); - } -#pragma unroll - for (int j = 0; j < pvec_t::NUM_ELTS; j++) { - x.elts[j] = packed_mul(PACKED_ACT_FN(x.elts[j], param), y.elts[j]); - } - if constexpr (use_256b) { - st256(x, &out_vec[i]); - } else { - st128(x, &out_vec[i]); - } - } - } else { - // Scalar fallback for unaligned data or small d - for (int64_t idx = threadIdx.x; idx < d; idx += blockDim.x) { - const scalar_t x = APHRODITE_LDG(&x_ptr[idx]); - const scalar_t y = APHRODITE_LDG(&y_ptr[idx]); - out_ptr[idx] = ACT_FN(x, param) * y; - } - } -} - -template -__device__ __forceinline__ T swigluoai_and_mul(const T& gate, const T& up, - float alpha, float limit) { - // Clamp gate to (-inf, limit] and up to [-limit, limit] - const float g = fminf((float)gate, limit); - const float u = fmaxf(fminf((float)up, limit), -limit); - // glu = gate * sigmoid(gate * alpha), then return (up + 1) * glu - return (T)((u + 1.0f) * g / (1.0f + expf(-g * alpha))); -} - -// Interleaved gate/up: input has [gate0, up0, gate1, up1, ...]. -template -__global__ void swigluoai_and_mul_kernel( - scalar_t* __restrict__ out, // [..., d] - const scalar_t* __restrict__ input, // [..., 2 * d] (interleaved) - const int d, const float alpha, const float limit) { - // For interleaved data: input has 2*d elements per token (gate/up pairs) - // output has d elements per token - constexpr int VEC_SIZE = 16 / sizeof(scalar_t); - constexpr int PAIRS = VEC_SIZE / 2; // Number of gate/up pairs per int4 load - const int64_t token_idx = blockIdx.x; - const scalar_t* in_ptr = input + token_idx * 2 * d; - scalar_t* out_ptr = out + token_idx * d; - - // Check alignment for 128-bit vectorized access on input. - // For output we use int2 (64-bit) which has 8-byte alignment requirement. - const bool in_aligned = is_16byte_aligned(in_ptr); - const bool out_aligned = - (reinterpret_cast(out_ptr) & 7) == 0; // 8-byte for int2 - - if (in_aligned && out_aligned && d >= PAIRS) { - // Fast path: vectorized loop - // Each int4 load gives VEC_SIZE elements = PAIRS gate/up pairs - // Each int2 store writes PAIRS output elements - const int4* in_vec = reinterpret_cast(in_ptr); - int2* out_vec = reinterpret_cast(out_ptr); - const int num_vecs = d / PAIRS; - const int vec_end = num_vecs * PAIRS; - - for (int i = threadIdx.x; i < num_vecs; i += blockDim.x) { - int4 v = APHRODITE_LDG(&in_vec[i]); - int2 r; - auto* vp = reinterpret_cast(&v); - auto* rp = reinterpret_cast(&r); -#pragma unroll - for (int j = 0; j < PAIRS; j++) { - rp[j] = ACT_FN(vp[2 * j], vp[2 * j + 1], alpha, limit); - } - out_vec[i] = r; - } - // Scalar cleanup for remaining elements - for (int i = vec_end + threadIdx.x; i < d; i += blockDim.x) { - out_ptr[i] = ACT_FN(APHRODITE_LDG(&in_ptr[2 * i]), - APHRODITE_LDG(&in_ptr[2 * i + 1]), alpha, limit); - } - } else { - // Scalar fallback for unaligned data or small d - for (int64_t idx = threadIdx.x; idx < d; idx += blockDim.x) { - // gate = x[..., ::2] (even indices) - const scalar_t gate = APHRODITE_LDG(&in_ptr[2 * idx]); - // up = x[..., 1::2] (odd indices) - const scalar_t up = APHRODITE_LDG(&in_ptr[2 * idx + 1]); - out_ptr[idx] = ACT_FN(gate, up, alpha, limit); - } - } -} - -} // namespace aphrodite - -#define LAUNCH_ACTIVATION_GATE_KERNEL_WITH_PARAM(KERNEL, PACKED_KERNEL, PARAM) \ - auto dtype = input.scalar_type(); \ - int d = input.size(-1) / 2; \ - int64_t num_tokens = input.numel() / input.size(-1); \ - if (num_tokens == 0) { \ - return; \ - } \ - dim3 grid(num_tokens); \ - int cc_major = get_device_prop()->major; \ - int support_vec = \ - (CUDA_VERSION >= 12090 && cc_major >= 10 && num_tokens > 128) \ - ? aphrodite::VecTraits::ARCH_MAX_VEC_SIZE \ - : aphrodite::VecTraits::ARCH_MAX_VEC_SIZE; \ - int vec_size = support_vec / input.element_size(); \ - const bool use_vec = (d % vec_size == 0); \ - const torch::stable::accelerator::DeviceGuard device_guard( \ - input.get_device_index()); \ - const cudaStream_t stream = get_current_cuda_stream(); \ - if (use_vec) { \ - dim3 block(std::min(d / vec_size, 1024)); \ - if (CUDA_VERSION >= 12090 && cc_major >= 10 && num_tokens > 128) { \ - APHRODITE_STABLE_DISPATCH_FLOATING_TYPES( \ - dtype, "act_and_mul_kernel_with_param", [&] { \ - aphrodite::act_and_mul_kernel_with_param< \ - scalar_t, \ - typename aphrodite::PackedTypeConverter::Type, \ - KERNEL, \ - PACKED_KERNEL< \ - typename aphrodite::PackedTypeConverter::Type>, \ - true, true><<>>( \ - out.mutable_data_ptr(), \ - input.const_data_ptr(), d, PARAM); \ - }); \ - } else { \ - APHRODITE_STABLE_DISPATCH_FLOATING_TYPES( \ - dtype, "act_and_mul_kernel_with_param", [&] { \ - aphrodite::act_and_mul_kernel_with_param< \ - scalar_t, \ - typename aphrodite::PackedTypeConverter::Type, \ - KERNEL, \ - PACKED_KERNEL< \ - typename aphrodite::PackedTypeConverter::Type>, \ - true, false><<>>( \ - out.mutable_data_ptr(), \ - input.const_data_ptr(), d, PARAM); \ - }); \ - } \ - } else { \ - dim3 block(std::min(d, 1024)); \ - APHRODITE_STABLE_DISPATCH_FLOATING_TYPES( \ - dtype, "act_and_mul_kernel_with_param", [&] { \ - aphrodite::act_and_mul_kernel_with_param< \ - scalar_t, \ - typename aphrodite::PackedTypeConverter::Type, \ - KERNEL, \ - PACKED_KERNEL< \ - typename aphrodite::PackedTypeConverter::Type>, \ - false><<>>( \ - out.mutable_data_ptr(), \ - input.const_data_ptr(), d, PARAM); \ - }); \ - } - -#define LAUNCH_SIGLUOAI_AND_MUL(KERNEL, ALPHA, LIMIT) \ - int d = input.size(-1) / 2; \ - int64_t num_tokens = input.numel() / input.size(-1); \ - 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(), "clamp_swiglu_kernel_with_params", [&] { \ - aphrodite::swigluoai_and_mul_kernel> \ - <<>>(out.mutable_data_ptr(), \ - input.const_data_ptr(), d, \ - ALPHA, LIMIT); \ - }); - -void fatrelu_and_mul(torch::stable::Tensor& out, // [..., d], - torch::stable::Tensor& input, // [..., 2 * d] - double threshold) { - LAUNCH_ACTIVATION_GATE_KERNEL_WITH_PARAM( - aphrodite::fatrelu_kernel, aphrodite::packed_fatrelu_kernel, threshold); -} -void swigluoai_and_mul(torch::stable::Tensor& out, // [..., d] - torch::stable::Tensor& input, // [..., 2 * d] - double alpha, double limit) { - LAUNCH_SIGLUOAI_AND_MUL(aphrodite::swigluoai_and_mul, alpha, limit); -} -namespace aphrodite { - -// Element-wise activation kernel template. -template -__global__ void activation_kernel( - scalar_t* __restrict__ out, // [..., d] - const scalar_t* __restrict__ input, // [..., d] - const int d) { - const scalar_t* in_ptr = input + blockIdx.x * d; - scalar_t* out_ptr = out + blockIdx.x * d; - - if constexpr (use_vec) { - // Fast path: 128-bit/256-bit vectorized loop - using vec_t = typename VecTraits::vec_t; - constexpr int ARCH_MAX_VEC_SIZE = VecTraits::ARCH_MAX_VEC_SIZE; - constexpr int VEC_SIZE = ARCH_MAX_VEC_SIZE / sizeof(scalar_t); - const vec_t* in_vec = reinterpret_cast(in_ptr); - vec_t* out_vec = reinterpret_cast(out_ptr); - const int num_vecs = d / VEC_SIZE; - - for (int i = threadIdx.x; i < num_vecs; i += blockDim.x) { - vec_t v; - if constexpr (use_256b) { - ld256(v, &in_vec[i]); - } else { - v = APHRODITE_LDG(&in_vec[i]); - } - auto* vp = reinterpret_cast(&v); -#pragma unroll - for (int j = 0; j < VEC_SIZE; j++) { - vp[j] = ACT_FN(vp[j]); - } - if constexpr (use_256b) { - st256(v, &out_vec[i]); - } else { - out_vec[i] = v; - } - } - } else { - // Scalar fallback for unaligned data or small d - for (int64_t idx = threadIdx.x; idx < d; idx += blockDim.x) { - const scalar_t x = APHRODITE_LDG(&in_ptr[idx]); - out_ptr[idx] = ACT_FN(x); - } - } -} - -} // namespace aphrodite - -// Launch element-wise activation kernel. -#define LAUNCH_ACTIVATION_KERNEL(KERNEL) \ - auto dtype = input.scalar_type(); \ - int d = input.size(-1); \ - int64_t num_tokens = input.numel() / input.size(-1); \ - if (num_tokens == 0) { \ - return; \ - } \ - dim3 grid(num_tokens); \ - int cc_major = get_device_prop()->major; \ - int support_vec = \ - (CUDA_VERSION >= 12090 && cc_major >= 10 && num_tokens > 128) \ - ? aphrodite::VecTraits::ARCH_MAX_VEC_SIZE \ - : aphrodite::VecTraits::ARCH_MAX_VEC_SIZE; \ - int vec_size = support_vec / input.element_size(); \ - const bool use_vec = (d % vec_size == 0); \ - const torch::stable::accelerator::DeviceGuard device_guard( \ - input.get_device_index()); \ - const cudaStream_t stream = get_current_cuda_stream(); \ - if (use_vec) { \ - dim3 block(std::min(d / vec_size, 1024)); \ - if (CUDA_VERSION >= 12090 && cc_major >= 10 && num_tokens > 128) { \ - APHRODITE_STABLE_DISPATCH_FLOATING_TYPES( \ - dtype, "activation_kernel", [&] { \ - aphrodite::activation_kernel, true, \ - true><<>>( \ - out.mutable_data_ptr(), \ - input.const_data_ptr(), d); \ - }); \ - } else { \ - APHRODITE_STABLE_DISPATCH_FLOATING_TYPES( \ - dtype, "activation_kernel", [&] { \ - aphrodite::activation_kernel, true, \ - false><<>>( \ - out.mutable_data_ptr(), \ - input.const_data_ptr(), d); \ - }); \ - } \ - } else { \ - dim3 block(std::min(d, 1024)); \ - APHRODITE_STABLE_DISPATCH_FLOATING_TYPES(dtype, "activation_kernel", [&] { \ - aphrodite::activation_kernel, false> \ - <<>>(out.mutable_data_ptr(), \ - input.const_data_ptr(), d); \ - }); \ - } - -namespace aphrodite { - -template -__device__ __forceinline__ T gelu_new_kernel(const T& x) { - const float x3 = (float)(x * x * x); - const T t = (T)tanhf((T)(0.79788456f * (float)(x + (T)(0.044715f * x3)))); - return ((T)0.5) * x * (((T)1.0) + t); -} - -template -__device__ __forceinline__ T gelu_fast_kernel(const T& x) { - const float f = (float)x; - const T t = - (T)tanhf(((T)(f * 0.79788456f)) * (((T)1.0) + (T)(0.044715f * f) * x)); - return ((T)0.5) * x * (((T)1.0) + t); -} - -template -__device__ __forceinline__ T gelu_quick_kernel(const T& x) { - // x * sigmoid(1.702 * x) - return (T)(((float)x) / (1.0f + expf(-1.702f * (float)x))); -} - -template -__device__ __forceinline__ T relu_squared_kernel(const T& x) { - // relu(x)^2 - introduced in https://arxiv.org/abs/2109.08668v2 - const float f = (float)x; - const float val = f > 0.0f ? f : 0.0f; - return (T)(val * val); -} - -} // namespace aphrodite - -void gelu_new(torch::stable::Tensor& out, // [..., d] - torch::stable::Tensor& input) // [..., d] -{ - LAUNCH_ACTIVATION_KERNEL(aphrodite::gelu_new_kernel); -} - -void gelu_fast(torch::stable::Tensor& out, // [..., d] - torch::stable::Tensor& input) // [..., d] -{ - LAUNCH_ACTIVATION_KERNEL(aphrodite::gelu_fast_kernel); -} - -void gelu_quick(torch::stable::Tensor& out, // [..., d] - torch::stable::Tensor& input) // [..., d] -{ - LAUNCH_ACTIVATION_KERNEL(aphrodite::gelu_quick_kernel); -} - -void relu_squared(torch::stable::Tensor& out, // [..., d] - torch::stable::Tensor& input) // [..., d] -{ - LAUNCH_ACTIVATION_KERNEL(aphrodite::relu_squared_kernel); -} diff --git a/csrc/libtorch_stable/cache_kernels.cu.orig b/csrc/libtorch_stable/cache_kernels.cu.orig deleted file mode 100644 index fc8e871c9d..0000000000 --- a/csrc/libtorch_stable/cache_kernels.cu.orig +++ /dev/null @@ -1,1669 +0,0 @@ -#include "torch_utils.h" -#include "dispatch_utils.h" - -#include "../cuda_utils.h" -#include "../cuda_compat.h" - -#include "quantization/vectorization_utils.cuh" -#include "concat_mla_q.cuh" - -#ifdef USE_ROCM - #include "../quantization/w8a8/fp8/amd/quant_utils.cuh" -#else - #include "../quantization/w8a8/fp8/nvidia/quant_utils.cuh" -#endif - -#include -#include -#include - -#ifdef USE_ROCM - #include -typedef __hip_bfloat16 __nv_bfloat16; -#else - #include -#endif - -#if defined(__gfx942__) -constexpr float kFp8ScaleDivisor = 224.f; -#else -constexpr float kFp8ScaleDivisor = 448.f; -#endif - -void swap_blocks(torch::stable::Tensor& src, torch::stable::Tensor& dst, - int64_t block_size_in_bytes, - const torch::stable::Tensor& block_mapping) { - torch::stable::Device src_device = src.device(); - torch::stable::Device dst_device = dst.device(); - cudaMemcpyKind memcpy_type; - if (src_device.is_cuda() && dst_device.is_cuda()) { - STD_TORCH_CHECK(src_device.index() == dst_device.index(), - "src and dst must be on the same GPU"); - memcpy_type = cudaMemcpyDeviceToDevice; - } else if (src_device.is_cuda() && dst_device.is_cpu()) { - memcpy_type = cudaMemcpyDeviceToHost; - } else if (src_device.is_cpu() && dst_device.is_cuda()) { - memcpy_type = cudaMemcpyHostToDevice; - } else { - STD_TORCH_CHECK(false, "Invalid device combination"); - } - - // NOTE(youkaichao): keep in mind that `block_mapping` should be - // a cpu tensor, otherwise every `item` call will require a gpu-cpu - // synchronization. - STD_TORCH_CHECK(block_mapping.device().is_cpu(), - "block_mapping must be on CPU"); - - char* src_ptr = static_cast(src.data_ptr()); - char* dst_ptr = static_cast(dst.data_ptr()); - - auto guard_device = src_device.is_cuda() ? src_device : dst_device; - const torch::stable::accelerator::DeviceGuard device_guard( - guard_device.index()); - const cudaStream_t stream = get_current_cuda_stream(); - // NOTE(woosuk): This can be slow if the number of blocks is large. - const int64_t num_blocks = block_mapping.size(0); - const int64_t* bm_ptr = block_mapping.const_data_ptr(); - const int64_t bm_stride0 = block_mapping.stride(0); - const int64_t bm_stride1 = block_mapping.stride(1); - for (size_t i = 0; i < num_blocks; i++) { - int64_t src_block_number = bm_ptr[i * bm_stride0]; - int64_t dst_block_number = bm_ptr[i * bm_stride0 + bm_stride1]; - int64_t src_offset = src_block_number * block_size_in_bytes; - int64_t dst_offset = dst_block_number * block_size_in_bytes; - cudaMemcpyAsync(dst_ptr + dst_offset, src_ptr + src_offset, - block_size_in_bytes, memcpy_type, stream); - } -} - -void swap_blocks_batch(const torch::stable::Tensor& src_ptrs, - const torch::stable::Tensor& dst_ptrs, - const torch::stable::Tensor& sizes, - bool is_src_access_order_any) { - STD_TORCH_CHECK(src_ptrs.device().is_cpu(), "src_ptrs must be on CPU"); - STD_TORCH_CHECK(dst_ptrs.device().is_cpu(), "dst_ptrs must be on CPU"); - STD_TORCH_CHECK(sizes.device().is_cpu(), "sizes must be on CPU"); - STD_TORCH_CHECK(src_ptrs.scalar_type() == torch::headeronly::ScalarType::Long, - "src_ptrs must be int64"); - STD_TORCH_CHECK(dst_ptrs.scalar_type() == torch::headeronly::ScalarType::Long, - "dst_ptrs must be int64"); - STD_TORCH_CHECK(sizes.scalar_type() == torch::headeronly::ScalarType::Long, - "sizes must be int64"); - - const int64_t n = src_ptrs.size(0); - STD_TORCH_CHECK(dst_ptrs.size(0) == n, "dst_ptrs length must match src_ptrs"); - STD_TORCH_CHECK(sizes.size(0) == n, "sizes length must match src_ptrs"); - - if (n == 0) return; - - int64_t* src_data = src_ptrs.mutable_data_ptr(); - int64_t* dst_data = dst_ptrs.mutable_data_ptr(); - int64_t* size_data = sizes.mutable_data_ptr(); - - const cudaStream_t stream = get_current_cuda_stream(); - - // Use cuMemcpyBatchAsync / hipMemcpyBatchAsync to submit all copies in a - // single driver call, amortizing per-copy submission overhead. int64_t - // and CUdeviceptr/void*/size_t are all 8 bytes on 64-bit platforms, so we - // reinterpret_cast the tensor data directly to avoid copies. - static_assert(sizeof(size_t) == sizeof(int64_t)); -#if !defined(USE_ROCM) && defined(CUDA_VERSION) && CUDA_VERSION >= 12080 - static_assert(sizeof(CUdeviceptr) == sizeof(int64_t)); - // Resolve cuMemcpyBatchAsync at runtime via cuGetProcAddress so that - // binaries compiled with CUDA 12.8+ still work on older drivers, and - // we avoid the CUDA 13.0 header remapping (#define to _v2 signature). - // The function pointer is cached after the first call. - using BatchFn = - CUresult (*)(CUdeviceptr*, CUdeviceptr*, size_t*, size_t, - CUmemcpyAttributes*, size_t*, size_t, size_t*, CUstream); - static BatchFn batch_fn = []() -> BatchFn { - CUdriverProcAddressQueryResult sym_status; - void* fn_ptr = nullptr; - CUresult res = cuGetProcAddress("cuMemcpyBatchAsync", &fn_ptr, 12080, - CU_GET_PROC_ADDRESS_DEFAULT, &sym_status); - if (res != CUDA_SUCCESS || fn_ptr == nullptr) { - return nullptr; - } - return reinterpret_cast(fn_ptr); - }(); - - // cuMemcpyBatchAsync rejects the legacy default stream (handle 0 / - // cudaStreamLegacy) with CUDA_ERROR_INVALID_VALUE; route it to the per-copy - // fallback below, which is correct on any stream. Real and per-thread-default - // streams take the batch fast path. - const bool usable_stream = stream != nullptr && stream != cudaStreamLegacy; - if (batch_fn != nullptr && usable_stream) { - CUmemcpyAttributes attr = {}; - // ANY lets the DMA engine prefetch source bytes out of stream order, - // which is only safe when no GPU stream is concurrently writing the - // source. - attr.srcAccessOrder = is_src_access_order_any - ? CU_MEMCPY_SRC_ACCESS_ORDER_ANY - : CU_MEMCPY_SRC_ACCESS_ORDER_STREAM; - size_t attrs_idx = 0; - size_t fail_idx = 0; - CUresult result = batch_fn(reinterpret_cast(dst_data), - reinterpret_cast(src_data), - reinterpret_cast(size_data), - static_cast(n), &attr, &attrs_idx, 1, - &fail_idx, static_cast(stream)); - STD_TORCH_CHECK(result == CUDA_SUCCESS, - "cuMemcpyBatchAsync failed at index ", fail_idx, - " with error ", result); - return; - } -#elif defined(USE_ROCM) && defined(HIP_VERSION) && HIP_VERSION >= 70100000 - // ROCm 7.1+ exposes hipMemcpyBatchAsync. The 7.2.1 implementation early- - // returns hipErrorNotSupported whenever numAttrs > 0 (see ROCm/clr @ - // rocm-7.2.1 hipamd/src/hip_memory.cpp:2819-2822), so call with - // numAttrs=0. - { - hipMemcpyAttributes attr = {}; - size_t attrs_idx = 0; - size_t fail_idx = 0; - hipError_t result = hipMemcpyBatchAsync( - reinterpret_cast(dst_data), reinterpret_cast(src_data), - reinterpret_cast(size_data), static_cast(n), &attr, - &attrs_idx, 0, &fail_idx, static_cast(stream)); - STD_TORCH_CHECK(result == hipSuccess, - "hipMemcpyBatchAsync failed at index ", fail_idx, - " with error ", result); - return; - } -#endif - { - // Fallback for CUDA < 12.8, older CUDA drivers, and ROCm < 7.1: - // individual async copies. cudaMemcpyDefault lets the driver infer - // direction from pointer types. - for (int64_t i = 0; i < n; i++) { - cudaMemcpyAsync(reinterpret_cast(dst_data[i]), - reinterpret_cast(src_data[i]), - static_cast(size_data[i]), cudaMemcpyDefault, - stream); - } - } -} - -namespace aphrodite { - -// Grid: (num_layers, num_pairs) -template -__global__ void copy_blocks_kernel(int64_t* key_cache_ptrs, - int64_t* value_cache_ptrs, - const int64_t* __restrict__ block_mapping, - const int numel_per_block) { - const int layer_idx = blockIdx.x; - const int pair_idx = blockIdx.y; - - scalar_t* key_cache = reinterpret_cast(key_cache_ptrs[layer_idx]); - scalar_t* value_cache = - reinterpret_cast(value_cache_ptrs[layer_idx]); - int64_t src_block_number = block_mapping[2 * pair_idx]; - int64_t dst_block_number = block_mapping[2 * pair_idx + 1]; - - const int64_t src_block_offset = src_block_number * numel_per_block; - const int64_t dst_block_offset = dst_block_number * numel_per_block; - for (int i = threadIdx.x; i < numel_per_block; i += blockDim.x) { - int64_t src_offset = src_block_offset + i; - int64_t dst_offset = dst_block_offset + i; - key_cache[dst_offset] = key_cache[src_offset]; - } - for (int i = threadIdx.x; i < numel_per_block; i += blockDim.x) { - int64_t src_offset = src_block_offset + i; - int64_t dst_offset = dst_block_offset + i; - value_cache[dst_offset] = value_cache[src_offset]; - } -} - -// Kernel for MLA, which works on a single joint kv_cache -// Grid: (num_layers, num_pairs) -template -__global__ void copy_blocks_mla_kernel( - int64_t* cache_ptrs, const int64_t* __restrict__ block_mapping, - const int mem_footprint_per_block) { - const int layer_idx = blockIdx.x; - const int pair_idx = blockIdx.y; - scalar_t* cache = reinterpret_cast(cache_ptrs[layer_idx]); - int64_t src_block = block_mapping[2 * pair_idx]; - int64_t dst_block = block_mapping[2 * pair_idx + 1]; - int64_t src_offset = src_block * mem_footprint_per_block; - int64_t dst_offset = dst_block * mem_footprint_per_block; - for (int i = threadIdx.x; i < mem_footprint_per_block; i += blockDim.x) { - cache[dst_offset + i] = cache[src_offset + i]; - } -} - -} // namespace aphrodite - -namespace aphrodite { - -// Used to copy/convert one element -template -struct CopyWithScaleOp { - float scale; - - __device__ __forceinline__ void operator()(OutT& dst, const InT src) const { - if constexpr (kv_dt == Fp8KVCacheDataType::kAuto) { - dst = static_cast(src); - } else { - dst = fp8::scaled_convert(src, scale); - } - } -}; - -template -__global__ void reshape_and_cache_kernel( - const scalar_t* __restrict__ key, // [num_tokens, num_heads, head_size] - const scalar_t* __restrict__ value, // [num_tokens, num_heads, head_size] - cache_t* __restrict__ key_cache, // [num_blocks, num_heads, head_size/x, - // block_size, x] - cache_t* __restrict__ value_cache, // [num_blocks, num_heads, head_size, - // block_size] - const int64_t* __restrict__ slot_mapping, // [num_tokens] - const int key_stride, const int value_stride, const int num_heads, - const int head_size, const int block_size, const int x, - const float* k_scale, const float* v_scale) { - const int64_t token_idx = blockIdx.x; - const int64_t slot_idx = slot_mapping[token_idx]; - if (slot_idx < 0) { - return; - } - - const int64_t block_idx = slot_idx / block_size; - const int64_t block_offset = slot_idx % block_size; - const int h_block_count = head_size / x; // head_size//x - - const int h_block_idx = threadIdx.x; - if (h_block_idx >= num_heads * h_block_count) { - return; - } - - const int head_idx = h_block_idx / h_block_count; - const int h_block = h_block_idx % h_block_count; - - const scalar_t* __restrict__ key_src = - key + token_idx * key_stride + head_idx * head_size + h_block * x; - const int64_t src_value_start = - token_idx * value_stride + head_idx * head_size + h_block * x; - - cache_t* __restrict__ key_dst = - key_cache + block_idx * num_heads * h_block_count * block_size * x + - head_idx * h_block_count * block_size * x + h_block * block_size * x + - block_offset * x; - const int64_t tgt_value_start = - block_idx * num_heads * h_block_count * x * block_size + - head_idx * h_block_count * x * block_size + h_block * x * block_size + - block_offset; - - constexpr int VEC_SIZE = (sizeof(scalar_t) == 2) ? 8 : 4; - float k_scale_val = (kv_dt == Fp8KVCacheDataType::kAuto) ? 0.f : *k_scale; - CopyWithScaleOp k_op{k_scale_val}; - float v_scale_val = (kv_dt == Fp8KVCacheDataType::kAuto) ? 0.f : *v_scale; - CopyWithScaleOp v_op{v_scale_val}; - - vectorize_with_alignment(key_src, key_dst, x, 0, 1, k_op); - - const scalar_t* __restrict__ value_src = value + src_value_start; - cache_t* __restrict__ value_dst = value_cache + tgt_value_start; -#pragma unroll - for (int i = 0; i < x; i++) { - v_op(value_dst[i * block_size], value_src[i]); - } -} - -template -__global__ void reshape_and_cache_flash_kernel( - const scalar_t* __restrict__ key, // [num_tokens, num_heads, head_size] - const scalar_t* __restrict__ value, // [num_tokens, num_heads, head_size] - cache_t* __restrict__ key_cache, // NHD or HND, shape see comments below - cache_t* __restrict__ value_cache, // same above - const int64_t* __restrict__ slot_mapping, // [num_tokens] - const int64_t block_stride, const int64_t page_stride, - const int64_t head_stride, const int64_t key_stride, - const int64_t value_stride, const int num_heads, const int head_size, - const int block_size, const float* k_scale, const float* v_scale, - const int kv_scale_stride) { - const int64_t token_idx = blockIdx.x; - const int64_t slot_idx = slot_mapping[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; - const int n_elems = num_heads * head_size; - - // pointers to the beginning of the source row for this token. - const scalar_t* __restrict__ key_src = key + token_idx * key_stride; - const scalar_t* __restrict__ value_src = value + token_idx * value_stride; - - // find the start position inside the kv-cache for this token. - cache_t* __restrict__ key_dst = - key_cache + block_idx * block_stride + block_offset * page_stride; - cache_t* __restrict__ value_dst = - value_cache + block_idx * block_stride + block_offset * page_stride; - - // this is true for the NHD layout where `head_stride == head_size` - const bool is_contiguous_heads = (head_stride == head_size); - - constexpr int VEC_SIZE = (sizeof(scalar_t) == 2) ? 8 : 4; - - if (is_contiguous_heads && kv_scale_stride == 0) { - // NHD layout and k/v_scales are [1] (i.e. single scale for all heads) - // kv cache: [num_blocks, block_size, num_heads, head_size] - float k_scale_val = (kv_dt == Fp8KVCacheDataType::kAuto) ? 0.f : *k_scale; - float v_scale_val = (kv_dt == Fp8KVCacheDataType::kAuto) ? 0.f : *v_scale; - - CopyWithScaleOp k_op{k_scale_val}; - CopyWithScaleOp v_op{v_scale_val}; - - vectorize_with_alignment(key_src, key_dst, n_elems, threadIdx.x, - blockDim.x, k_op); - vectorize_with_alignment(value_src, value_dst, n_elems, - threadIdx.x, blockDim.x, v_op); - } else { - // HND layout OR k/v_scales are [num_heads] (i.e. per-attn-head) - // HND layout: heads are strided, but each head_size segment is contiguous - // kv cache: [num_blocks, num_heads, block_size, head_size] - const int lane = threadIdx.x & 31; // 0..31 within warp - const int warp_id = threadIdx.x >> 5; // warp index within block - const int warps_per_block = blockDim.x >> 5; - - for (int head = warp_id; head < num_heads; head += warps_per_block) { - const scalar_t* __restrict__ k_src_h = key_src + head * head_size; - const scalar_t* __restrict__ v_src_h = value_src + head * head_size; - - cache_t* __restrict__ k_dst_h = - key_dst + static_cast(head) * head_stride; - cache_t* __restrict__ v_dst_h = - value_dst + static_cast(head) * head_stride; - - float k_scale_val = (kv_dt == Fp8KVCacheDataType::kAuto) - ? 0.f - : k_scale[head * kv_scale_stride]; - float v_scale_val = (kv_dt == Fp8KVCacheDataType::kAuto) - ? 0.f - : v_scale[head * kv_scale_stride]; - - CopyWithScaleOp k_op{k_scale_val}; - CopyWithScaleOp v_op{v_scale_val}; - - // within each head, let the 32 threads of the warp perform the vector - // copy - vectorize_with_alignment(k_src_h, k_dst_h, head_size, lane, 32, - k_op); - - vectorize_with_alignment(v_src_h, v_dst_h, head_size, lane, 32, - v_op); - } - } -} - -template -__global__ void concat_and_cache_mla_kernel( - const scalar_t* __restrict__ kv_c, // [num_tokens, kv_lora_rank] - const scalar_t* __restrict__ k_pe, // [num_tokens, pe_dim] - cache_t* __restrict__ kv_cache, // [num_blocks, block_size, (kv_lora_rank - // + pe_dim)] - const int64_t* __restrict__ slot_mapping, // [num_tokens] - const int block_stride, // - const int entry_stride, // - const int kv_c_stride, // - const int k_pe_stride, // - const int kv_lora_rank, // - const int pe_dim, // - const int block_size, // - const float* scale // -) { - const int64_t token_idx = blockIdx.x; - const int64_t slot_idx = slot_mapping[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; - - auto copy = [&](const scalar_t* __restrict__ src, cache_t* __restrict__ dst, - int src_stride, int dst_stride, int size, int offset) { - for (int i = threadIdx.x; i < size; i += blockDim.x) { - const int64_t src_idx = token_idx * src_stride + i; - const int64_t dst_idx = - block_idx * block_stride + block_offset * entry_stride + i + offset; - if constexpr (kv_dt == Fp8KVCacheDataType::kAuto) { - dst[dst_idx] = src[src_idx]; - } else { - dst[dst_idx] = - fp8::scaled_convert(src[src_idx], *scale); - } - } - }; - - copy(kv_c, kv_cache, kv_c_stride, block_stride, kv_lora_rank, 0); - 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] - const scalar_t* __restrict__ k_pe, // [num_tokens, pe_dim] - cache_t* __restrict__ kv_cache, // [num_blocks, block_size, (kv_lora_rank - // + pe_dim)] - const int64_t* __restrict__ slot_mapping, // [num_tokens] - const int block_stride, // - const int entry_stride, // - const int kv_c_stride, // - const int k_pe_stride, // - const int kv_lora_rank, // - const int pe_dim, // - const int block_size, // - const float* scale // -) { - const int64_t token_idx = blockIdx.x; - const int64_t slot_idx = slot_mapping[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; - const int64_t dst_idx_start = - block_idx * block_stride + block_offset * entry_stride; - - // For the NoPE part, each tile of 128 elements is handled by half of one warp - // (16 threads). There are 4 total tiles, so 2 warps (64 threads). - // Lanes 0 and 16 of each warp write the scale values for that warp's tiles. - // The RoPE part (last 64 elements) is handled by another 1 warp (32 threads). - // So in total, we use 3 warps (96 threads) per block. - - // Cast kv_cache to 16_bit for RoPE values - scalar_t* kv_cache_16bit = - reinterpret_cast(&kv_cache[dst_idx_start]); - - // The last warp handles the RoPE part - if (threadIdx.x >= 64) { - // Each thread handles two elements of RoPE - const int8_t pe_idx_start = (threadIdx.x - 64) * 2; - const int64_t src_idx = token_idx * k_pe_stride + pe_idx_start; - // Vectorized load of two 16-bit values, performed as one 32-bit load - const int32_t vals = *reinterpret_cast(&k_pe[src_idx]); - // RoPE values start after the packed 8-bit NoPE values and the - // 32-bit scales - const int64_t dst_idx = kv_lora_rank / 2 + 8 + pe_idx_start; - // Vectorized store of two 16-bit values, performed as one 32-bit store - *reinterpret_cast(&kv_cache_16bit[dst_idx]) = vals; - return; - } - - // The first two warps handle the NoPE part - const int8_t warp_idx = threadIdx.x >> 5; - const int8_t lane_idx = threadIdx.x & 31; - const int8_t tile_idx = warp_idx * 2 + (lane_idx >> 4); - - // Each thread handles 8 elements of NoPE - // Load the NoPE elements for this thread into registers - const int64_t src_idx_start = token_idx * kv_c_stride + (threadIdx.x * 8); - // Vectorized load of eight 16-bit values, performed as an int4 load - const int4 vals_i4 = *reinterpret_cast(&kv_c[src_idx_start]); - const scalar_t* vals = reinterpret_cast(&vals_i4); - - // Max absolute value of this thread's elements - float max_abs = fmaxf(fmaxf(fmaxf(fabsf(vals[0]), fabsf(vals[1])), - fmaxf(fabsf(vals[2]), fabsf(vals[3]))), - fmaxf(fmaxf(fabsf(vals[4]), fabsf(vals[5])), - fmaxf(fabsf(vals[6]), fabsf(vals[7])))); - - // Warp-level reduction to find the max absolute value in each half-warp -#pragma unroll - for (int offset = 8; offset > 0; offset /= 2) { - max_abs = - fmaxf(max_abs, APHRODITE_SHFL_XOR_SYNC_WIDTH(max_abs, offset, 16)); - } - - // Compute the scale for the tile - float tile_scale = fmaxf(max_abs / kFp8ScaleDivisor, FLT_MIN); - - // The first lane of each half-warp writes the scale to kv_cache - if ((lane_idx == 0) || (lane_idx == 16)) { - float* kv_cache_32bit = reinterpret_cast(&kv_cache[dst_idx_start]); - const uint64_t dst_idx = kv_lora_rank / 4 + tile_idx; - kv_cache_32bit[dst_idx] = tile_scale; - } - - // Now all threads in the block scale and write their elements - // NoPE data is packed in the first kv_lora_rank/2 bytes (first 256 bytes) - const int64_t dst_idx_base = dst_idx_start + (threadIdx.x * 8); - - uint8_t result[8]; -#pragma unroll - for (int i = 0; i < 8; i++) { - result[i] = - fp8::scaled_convert( - vals[i], tile_scale); - } - - // Store as aligned 64-bit writes - *reinterpret_cast(&kv_cache[dst_idx_base]) = - *reinterpret_cast(result); -} - -template -__global__ void indexer_k_quant_and_cache_kernel( - const scalar_t* __restrict__ k, // [num_tokens, head_dim] - cache_t* __restrict__ kv_cache, // [num_blocks, block_size, cache_stride] - const int64_t* __restrict__ slot_mapping, // [num_tokens] - const int head_dim, // dimension of each head - const int quant_block_size, // quantization block size - const int cache_block_size, // cache block size - const int64_t cache_block_stride, // stride for each block in kv_cache - - const bool use_ue8m0 // use ue8m0 scale format -) { - constexpr int VEC_SIZE = 4; - const int64_t token_idx = blockIdx.x; - const int64_t head_dim_idx = (blockIdx.y * blockDim.y * blockDim.x + - threadIdx.y * blockDim.x + threadIdx.x) * - VEC_SIZE; - const int64_t slot_idx = slot_mapping[token_idx]; - const int64_t block_idx = slot_idx / cache_block_size; - const int64_t block_offset = slot_idx % cache_block_size; - - // NOTE: slot_idx can be -1 if the token is padded - if (slot_idx < 0 || (head_dim_idx >= head_dim)) { - return; - } - - float2 k_val = (reinterpret_cast( - k))[(token_idx * head_dim + head_dim_idx) / VEC_SIZE]; - scalar_t* k_val_ptr = reinterpret_cast(&k_val); - float amax = 0.0f; - for (int i = 0; i < VEC_SIZE; i++) { - amax = fmaxf(amax, fabsf(float(k_val_ptr[i]))); - } - - // Reduced amax - for (int mask = 16; mask > 0; mask /= 2) { -#ifdef USE_ROCM - amax = fmaxf(amax, __shfl_xor_sync(uint64_t(-1), amax, mask)); -#else - amax = fmaxf(amax, __shfl_xor_sync(unsigned(-1), amax, mask)); -#endif - } - - float scale = fmaxf(amax, 1e-4) / kFp8ScaleDivisor; - - if (use_ue8m0) { - scale = exp2f(ceilf(log2f(scale))); - } - - const int64_t dst_offset = - block_idx * cache_block_stride + block_offset * head_dim + head_dim_idx; - for (int i = 0; i < VEC_SIZE; i++) { - kv_cache[dst_offset + i] = - fp8::scaled_convert(k_val_ptr[i], scale); - } - if (threadIdx.x == 0) { - const int64_t dst_scale_idx = - block_idx * cache_block_stride + cache_block_size * head_dim + - (block_offset * head_dim + head_dim_idx) * 4 / quant_block_size; - reinterpret_cast(kv_cache)[dst_scale_idx / 4] = scale; - } -} - -template -__global__ void cp_gather_indexer_k_quant_cache_kernel( - const char* __restrict__ kv_cache, // [num_blocks, block_size, - // cache_stride] - char* __restrict__ dst_k, // [num_tokens, head_dim] - char* __restrict__ dst_scale, // [num_tokens, head_dim / quant_block_size * - // 4] - const int* __restrict__ block_table, // [batch_size, num_blocks] - const int* __restrict__ cu_seq_lens, // [batch_size + 1] - const int batch_size, // batch size - const int64_t token_stride, // stride for each token in dst_k - const int64_t head_dim, // dimension of each head - const int64_t block_stride, // stride for each block in kv_cache - const int64_t cache_token_stride, // stride for each token in kv_cache - const int64_t cache_block_size, // num_tokens for each block in kv_cache - const int num_blocks, // number of blocks - const int num_tokens, // number of tokens - const int quant_block_size // quantization block size -) { - constexpr int VEC_SIZE = sizeof(float4) / sizeof(char); - const int token_idx = blockIdx.x * blockDim.y + threadIdx.y; - const int head_idx = (blockIdx.y * blockDim.x + threadIdx.x) * VEC_SIZE; - // Find batch index within a block - __shared__ int batch_idx[BLOCK_Y_SIZE]; - if (threadIdx.x == 0) { - batch_idx[threadIdx.y] = -1; - } - __syncthreads(); - - for (int iter = 0; iter < cuda_utils::ceil_div(batch_size, int(blockDim.x)); - iter++) { - int tid = iter * blockDim.x + threadIdx.x; - if (tid < batch_size) { - const int seq_start = cu_seq_lens[tid]; - const int seq_end = cu_seq_lens[tid + 1]; - if (token_idx >= seq_start && token_idx < seq_end) { - batch_idx[threadIdx.y] = tid; - } - } - } - - __syncthreads(); - - // num_tokens may be an allocation upper bound when Python avoids a D2H sync. - // Only tokens covered by the exact device-side cu_seq_lens are valid to - // gather. - const int batch = batch_idx[threadIdx.y]; - if (head_idx >= head_dim || token_idx >= num_tokens || batch < 0) { - return; - } - const int inbatch_seq_idx = token_idx - cu_seq_lens[batch]; - const int block_idx = - block_table[batch * num_blocks + inbatch_seq_idx / cache_block_size]; - const int64_t src_block_offset = block_idx * block_stride; - const int64_t cache_inblock_offset = - (inbatch_seq_idx % cache_block_size) * head_dim + head_idx; - const int64_t src_inblock_offset = src_block_offset + cache_inblock_offset; - const int64_t dst_inblock_offset = token_idx * token_stride + head_idx; - - reinterpret_cast(dst_k)[dst_inblock_offset / VEC_SIZE] = - reinterpret_cast(kv_cache)[src_inblock_offset / VEC_SIZE]; - ; - if (threadIdx.x == 0) { - const int64_t src_scale_offset = - src_block_offset + cache_block_size * head_dim + - cache_inblock_offset * 4 / quant_block_size; - reinterpret_cast(dst_scale)[dst_inblock_offset / quant_block_size] = - reinterpret_cast(kv_cache)[src_scale_offset / 4]; - } -} - -} // namespace aphrodite - -// KV_T is the data type of key and value tensors. -// CACHE_T is the stored data type of kv-cache. -// KV_DTYPE is the real data type of kv-cache. -#define CALL_RESHAPE_AND_CACHE(KV_T, CACHE_T, KV_DTYPE) \ - aphrodite::reshape_and_cache_kernel \ - <<>>( \ - reinterpret_cast(key.data_ptr()), \ - reinterpret_cast(value.data_ptr()), \ - reinterpret_cast(key_cache.data_ptr()), \ - reinterpret_cast(value_cache.data_ptr()), \ - slot_mapping.const_data_ptr(), key_stride, value_stride, \ - num_heads, head_size, block_size, x, \ - reinterpret_cast(k_scale.data_ptr()), \ - reinterpret_cast(v_scale.data_ptr())); - -void reshape_and_cache( - torch::stable::Tensor& key, // [num_tokens, num_heads, head_size] - torch::stable::Tensor& value, // [num_tokens, num_heads, head_size] - torch::stable::Tensor& - key_cache, // [num_blocks, num_heads, head_size/x, block_size, x] - torch::stable::Tensor& - value_cache, // [num_blocks, num_heads, head_size, block_size] - torch::stable::Tensor& slot_mapping, // [num_tokens] - const std::string& kv_cache_dtype, torch::stable::Tensor& k_scale, - torch::stable::Tensor& v_scale) { - int num_tokens = slot_mapping.size(0); - int num_heads = key.size(1); - int head_size = key.size(2); - int block_size = key_cache.size(3); - int x = key_cache.size(4); - - int key_stride = key.stride(0); - int value_stride = value.stride(0); - int head_div_x = head_size / x; - - dim3 grid(num_tokens); - dim3 block(std::min(num_heads * head_div_x, 512)); - const torch::stable::accelerator::DeviceGuard device_guard( - key.get_device_index()); - const cudaStream_t stream = get_current_cuda_stream(); - - DISPATCH_BY_KV_CACHE_DTYPE(key.scalar_type(), kv_cache_dtype, - CALL_RESHAPE_AND_CACHE); -} - -// KV_T is the data type of key and value tensors. -// CACHE_T is the stored data type of kv-cache. -// KV_DTYPE is the real data type of kv-cache. -#define CALL_RESHAPE_AND_CACHE_FLASH(KV_T, CACHE_T, KV_DTYPE) \ - aphrodite::reshape_and_cache_flash_kernel \ - <<>>( \ - reinterpret_cast(key.data_ptr()), \ - reinterpret_cast(value.data_ptr()), \ - reinterpret_cast(key_cache.data_ptr()), \ - reinterpret_cast(value_cache.data_ptr()), \ - slot_mapping.const_data_ptr(), block_stride, page_stride, \ - head_stride, key_stride, value_stride, num_heads, head_size, \ - block_size, reinterpret_cast(k_scale.data_ptr()), \ - reinterpret_cast(v_scale.data_ptr()), \ - kv_scale_stride); - -void reshape_and_cache_flash( - torch::stable::Tensor& key, // [num_tokens, num_heads, head_size] - torch::stable::Tensor& value, // [num_tokens, num_heads, head_size] - torch::stable::Tensor& - key_cache, // [num_blocks, block_size, num_heads, head_size] - torch::stable::Tensor& - value_cache, // [num_blocks, block_size, num_heads, head_size] - torch::stable::Tensor& slot_mapping, // [num_tokens] or [num_actual_tokens] - const std::string& kv_cache_dtype, - torch::stable::Tensor& k_scale, // [1] or [num_heads] - torch::stable::Tensor& v_scale) { // [1] or [num_heads] - // NOTE(woosuk): In vLLM V1, key.size(0) can be different from - // slot_mapping.size(0) because of padding for CUDA graphs. - // In vLLM V0, key.size(0) is always equal to slot_mapping.size(0) because - // both include padding. - // In vLLM V1, however, key.size(0) can be larger than slot_mapping.size(0) - // since key includes padding for CUDA graphs, while slot_mapping does not. - // In this case, slot_mapping.size(0) represents the actual number of tokens - // before padding. - // For compatibility with both cases, we use slot_mapping.size(0) as the - // number of tokens. - int num_tokens = slot_mapping.size(0); - int num_heads = key.size(1); - int head_size = key.size(2); - - const torch::stable::accelerator::DeviceGuard device_guard( - key.get_device_index()); - const cudaStream_t stream = get_current_cuda_stream(); - - if (kv_cache_dtype == "nvfp4") { -#if defined(ENABLE_NVFP4_SM100) || defined(ENABLE_NVFP4_SM120) - // NVFP4 dispatch is compiled separately for SM100+. - extern void reshape_and_cache_nvfp4_dispatch( - torch::stable::Tensor & key, torch::stable::Tensor & value, - torch::stable::Tensor & key_cache, torch::stable::Tensor & value_cache, - torch::stable::Tensor & slot_mapping, torch::stable::Tensor & k_scale, - torch::stable::Tensor & v_scale); - reshape_and_cache_nvfp4_dispatch(key, value, key_cache, value_cache, - slot_mapping, k_scale, v_scale); - return; -#else - STD_TORCH_CHECK( - false, - "NVFP4 KV cache requires SM100+ (Blackwell). " - "Please rebuild aphrodite with a Blackwell-compatible CUDA target."); -#endif - } - - // Original FP8/auto path. - int block_size = key_cache.size(1); - - int64_t key_stride = key.stride(0); - int64_t value_stride = value.stride(0); - int64_t block_stride = key_cache.stride(0); - int64_t page_stride = key_cache.stride(1); - int64_t head_stride = key_cache.stride(2); - STD_TORCH_CHECK(key_cache.stride(0) == value_cache.stride(0)); - - STD_TORCH_CHECK(k_scale.sizes().equals(v_scale.sizes()), - "k_scale and v_scale must have the same shape"); - STD_TORCH_CHECK(k_scale.numel() == 1 || k_scale.numel() == num_heads, - "k_scale and v_scale must be of shape [1] or [num_heads]"); - int kv_scale_stride = (k_scale.numel() > 1) ? 1 : 0; - - dim3 grid(num_tokens); - dim3 block(std::min(num_heads * head_size, 512)); - - DISPATCH_BY_KV_CACHE_DTYPE(key.scalar_type(), kv_cache_dtype, - CALL_RESHAPE_AND_CACHE_FLASH); -} - -// KV_T is the data type of key and value tensors. -// CACHE_T is the stored data type of kv-cache. -// KV_DTYPE is the real data type of kv-cache. -#define CALL_CONCAT_AND_CACHE_MLA(KV_T, CACHE_T, KV_DTYPE) \ - aphrodite::concat_and_cache_mla_kernel \ - <<>>( \ - reinterpret_cast(kv_c.data_ptr()), \ - reinterpret_cast(k_pe.data_ptr()), \ - reinterpret_cast(kv_cache.data_ptr()), \ - slot_mapping.const_data_ptr(), block_stride, entry_stride, \ - kv_c_stride, k_pe_stride, kv_lora_rank, pe_dim, block_size, \ - reinterpret_cast(scale.data_ptr())); - -// KV_T is the data type of key and value tensors. -// CACHE_T is the stored data type of kv-cache. -#define CALL_CONCAT_AND_CACHE_DS_MLA(KV_T, CACHE_T, KV_DTYPE) \ - aphrodite::concat_and_cache_ds_mla_kernel \ - <<>>( \ - reinterpret_cast(kv_c.data_ptr()), \ - reinterpret_cast(k_pe.data_ptr()), \ - reinterpret_cast(kv_cache.data_ptr()), \ - slot_mapping.const_data_ptr(), block_stride, entry_stride, \ - kv_c_stride, k_pe_stride, kv_lora_rank, pe_dim, block_size, \ - reinterpret_cast(scale.data_ptr())); - -void concat_and_cache_mla( - torch::stable::Tensor& kv_c, // [num_tokens, kv_lora_rank] - torch::stable::Tensor& k_pe, // [num_tokens, pe_dim] - torch::stable::Tensor& kv_cache, // [num_blocks, block_size, (kv_lora_rank - // + pe_dim)] - torch::stable::Tensor& slot_mapping, // [num_tokens] or [num_actual_tokens] - const std::string& kv_cache_dtype, torch::stable::Tensor& scale) { - // NOTE(woosuk): In vLLM V1, key.size(0) can be different from - // slot_mapping.size(0) because of padding for CUDA graphs. - // In vLLM V0, key.size(0) is always equal to slot_mapping.size(0) because - // both include padding. - // In vLLM V1, however, key.size(0) can be larger than slot_mapping.size(0) - // since key includes padding for CUDA graphs, while slot_mapping does not. - // In this case, slot_mapping.size(0) represents the actual number of tokens - // before padding. - // For compatibility with both cases, we use slot_mapping.size(0) as the - // number of tokens. - int num_tokens = slot_mapping.size(0); - int kv_lora_rank = kv_c.size(1); - int pe_dim = k_pe.size(1); - int block_size = kv_cache.size(1); - - if (kv_cache_dtype == "fp8_ds_mla") { - STD_TORCH_CHECK(kv_lora_rank == 512, - "kv_lora_rank must be 512 for fp8_ds_mla"); - STD_TORCH_CHECK(pe_dim == 64, "pe_dim must be 64 for fp8_ds_mla"); - STD_TORCH_CHECK(kv_cache.size(2) == 656 / kv_cache.element_size(), - "kv_cache.size(2) must be 656 bytes for fp8_ds_mla"); - STD_TORCH_CHECK(kv_c.element_size() == 2, - "kv_c.element_size() must be 2 for fp8_ds_mla"); - STD_TORCH_CHECK(k_pe.element_size() == 2, - "k_pe.element_size() must be 2 for fp8_ds_mla"); - } else { - STD_TORCH_CHECK(kv_cache.size(2) == kv_lora_rank + pe_dim); - } - - int kv_c_stride = kv_c.stride(0); - int k_pe_stride = k_pe.stride(0); - int block_stride = kv_cache.stride(0); - int entry_stride = kv_cache.stride(1); - - const torch::stable::accelerator::DeviceGuard device_guard( - kv_c.get_device_index()); - const cudaStream_t stream = get_current_cuda_stream(); - - if (kv_cache_dtype == "fp8_ds_mla") { - dim3 grid(num_tokens); - // For the NoPE part, each tile of 128 elements is handled by half of one - // warp (16 threads). There are 4 total tiles, so 2 warps (64 threads). - // Lanes 0 and 16 of each warp write the scale values for that warp's tiles. - // The RoPE part (last 64 elements) is handled by another 1 warp (32 - // threads). So in total, we use 3 warps (96 threads) per block. - dim3 block(96); - DISPATCH_BY_KV_CACHE_DTYPE(kv_c.scalar_type(), kv_cache_dtype, - CALL_CONCAT_AND_CACHE_DS_MLA); - } else { - dim3 grid(num_tokens); - dim3 block(std::min(kv_lora_rank, 512)); - DISPATCH_BY_KV_CACHE_DTYPE(kv_c.scalar_type(), kv_cache_dtype, - CALL_CONCAT_AND_CACHE_MLA); - } -} - -namespace aphrodite { - -template -__global__ void convert_fp8_kernel(const Tin* __restrict__ src_cache, - Tout* __restrict__ dst_cache, - const float scale, - const int64_t block_stride) { - const int64_t block_idx = blockIdx.x; - for (int i = threadIdx.x; i < block_stride; i += blockDim.x) { - int64_t idx = block_idx * block_stride + i; - dst_cache[idx] = - fp8::scaled_convert(src_cache[idx], scale); - } -} - -} // namespace aphrodite - -#define CALL_CONVERT_FP8(Tout, Tin, KV_DTYPE) \ - aphrodite::convert_fp8_kernel \ - <<>>( \ - reinterpret_cast(src_cache.data_ptr()), \ - reinterpret_cast(dst_cache.data_ptr()), scale, block_stride); - -// Only for testing. -void convert_fp8(torch::stable::Tensor& dst_cache, - torch::stable::Tensor& src_cache, const double scale, - const std::string& kv_cache_dtype) { - torch::stable::Device src_device = src_cache.device(); - torch::stable::Device dst_device = dst_cache.device(); - STD_TORCH_CHECK(src_device.is_cuda(), "src must be on a GPU") - STD_TORCH_CHECK(dst_device.is_cuda(), "dst must be on a GPU") - STD_TORCH_CHECK(src_device.index() == dst_device.index(), - "src and dst must be on the same GPU"); - torch::stable::accelerator::DeviceGuard device_guard(src_device.index()); - - int64_t num_blocks = src_cache.size(0); - int64_t block_stride = src_cache.stride(0); - - dim3 grid(num_blocks); - dim3 block(std::min(block_stride, int64_t(512))); - const cudaStream_t stream = get_current_cuda_stream(); - - if (kv_cache_dtype == "auto") { - if (src_cache.scalar_type() == torch::headeronly::ScalarType::Float) { - CALL_CONVERT_FP8(uint8_t, float, aphrodite::Fp8KVCacheDataType::kAuto); - } else if (src_cache.scalar_type() == torch::headeronly::ScalarType::Half) { - CALL_CONVERT_FP8(uint8_t, uint16_t, aphrodite::Fp8KVCacheDataType::kAuto); - } else if (src_cache.scalar_type() == - torch::headeronly::ScalarType::BFloat16) { - CALL_CONVERT_FP8(uint8_t, __nv_bfloat16, - aphrodite::Fp8KVCacheDataType::kAuto); - } else if (dst_cache.scalar_type() == - torch::headeronly::ScalarType::Float) { - CALL_CONVERT_FP8(float, uint8_t, aphrodite::Fp8KVCacheDataType::kAuto); - } else if (dst_cache.scalar_type() == torch::headeronly::ScalarType::Half) { - CALL_CONVERT_FP8(uint16_t, uint8_t, aphrodite::Fp8KVCacheDataType::kAuto); - } else if (dst_cache.scalar_type() == - torch::headeronly::ScalarType::BFloat16) { - CALL_CONVERT_FP8(__nv_bfloat16, uint8_t, - aphrodite::Fp8KVCacheDataType::kAuto); - } - } else if (kv_cache_dtype == "fp8" || kv_cache_dtype == "fp8_e4m3") { - if (src_cache.scalar_type() == torch::headeronly::ScalarType::Float) { - CALL_CONVERT_FP8(uint8_t, float, aphrodite::Fp8KVCacheDataType::kFp8E4M3); - } else if (src_cache.scalar_type() == torch::headeronly::ScalarType::Half) { - CALL_CONVERT_FP8(uint8_t, uint16_t, - aphrodite::Fp8KVCacheDataType::kFp8E4M3); - } else if (src_cache.scalar_type() == - torch::headeronly::ScalarType::BFloat16) { - CALL_CONVERT_FP8(uint8_t, __nv_bfloat16, - aphrodite::Fp8KVCacheDataType::kFp8E4M3); - } else if (dst_cache.scalar_type() == - torch::headeronly::ScalarType::Float) { - CALL_CONVERT_FP8(float, uint8_t, aphrodite::Fp8KVCacheDataType::kFp8E4M3); - } else if (dst_cache.scalar_type() == torch::headeronly::ScalarType::Half) { - CALL_CONVERT_FP8(uint16_t, uint8_t, - aphrodite::Fp8KVCacheDataType::kFp8E4M3); - } else if (dst_cache.scalar_type() == - torch::headeronly::ScalarType::BFloat16) { - CALL_CONVERT_FP8(__nv_bfloat16, uint8_t, - aphrodite::Fp8KVCacheDataType::kFp8E4M3); - } - } else { - STD_TORCH_CHECK(false, "Unsupported data type: ", kv_cache_dtype); - } -} - -namespace aphrodite { - -// grid is launched with dimensions (batch, num_splits) -template -__global__ void gather_and_maybe_dequant_cache( - const cache_t* __restrict__ src_cache, // [NUM_BLOCKS, BLOCK_SIZE, - // ENTRIES...] - scalar_t* __restrict__ dst, // [TOT_TOKENS, ENTRIES...] - const int32_t* __restrict__ block_table, // [BATCH, BLOCK_INDICES] - const int32_t* __restrict__ cu_seq_lens, // [BATCH+1] - const int32_t* __restrict__ token_to_seq, // [MAX_TOKEN_ACROSS_CHUNK] - const int32_t num_tokens, const int32_t block_size, - const int64_t block_table_stride, const int64_t cache_block_stride, - const int64_t cache_entry_stride, const int64_t dst_entry_stride, - const float* __restrict__ scale, - const int32_t* __restrict__ seq_starts) { // Optional: starting offsets per - // batch - constexpr int vec_size = sizeof(float4) / sizeof(scalar_t); - using ltype = aphrodite::vec_n_t; - using stype = aphrodite::vec_n_t; - // We are adding this for code readability which will be optimized out when - // build in release. - assert(CTA_SIZE == blockDim.x); - -#pragma unroll - for (int token_id = blockIdx.x; token_id < num_tokens; - token_id += gridDim.x) { - int64_t batch_id = token_to_seq[token_id]; - int64_t batch_start = cu_seq_lens[batch_id]; - int64_t batch_end = cu_seq_lens[batch_id + 1]; - int32_t batch_offset = token_id - batch_start; - - if (token_id >= batch_end) return; - int32_t offset = 0; - if (seq_starts != nullptr) { - offset = seq_starts[batch_id]; - } - batch_offset += offset; - int32_t block_table_id = batch_offset / block_size; - int32_t slot_id = batch_offset % block_size; - // seq_starts may push the block index past the end of the batch's block - // table row. - if (block_table_id >= block_table_stride) continue; - int32_t block_table_offset = batch_id * block_table_stride + block_table_id; - int32_t block_id = block_table[block_table_offset]; - int64_t cache_offset = - block_id * cache_block_stride + slot_id * cache_entry_stride; - constexpr int32_t vec_iter_cnt = ENTRY_SIZE / vec_size; - scalar_t* dst_ = dst + token_id * dst_entry_stride; - cache_t* src_ = const_cast(src_cache) + cache_offset; - -#pragma unroll - for (int idx = threadIdx.x; idx < vec_iter_cnt; idx += CTA_SIZE) { - if constexpr (kv_dt == Fp8KVCacheDataType::kAuto) { - reinterpret_cast(dst_)[idx] = - static_cast(reinterpret_cast(src_)[idx]); - } else { - ltype loaded_val = reinterpret_cast(src_)[idx]; - stype store_val; -#pragma unroll - for (int j = 0; j < vec_size; ++j) { - store_val.val[j] = fp8::scaled_convert( - loaded_val.val[j], *scale); - } - reinterpret_cast(dst_)[idx] = store_val; - } - } - // process tail - constexpr int32_t tail_cnt = ENTRY_SIZE % vec_size; - dst_ = dst_ + ENTRY_SIZE - tail_cnt; - src_ = src_ + ENTRY_SIZE - tail_cnt; -#pragma unroll - for (int idx = threadIdx.x; idx < tail_cnt; idx += CTA_SIZE) { - if constexpr (kv_dt == Fp8KVCacheDataType::kAuto) { - dst_[idx] = static_cast(src_[idx]); - } else { - dst_[idx] = - fp8::scaled_convert(src_[idx], *scale); - } - } - } -} - -} // namespace aphrodite - -// Macro to dispatch the kernel based on the data type. -// SCALAR_T is the data type of the destination tensor. -// CACHE_T is the stored data type of kv-cache. -// KV_DTYPE is the real data type of kv-cache. -#define CALL_GATHER_CACHE(SCALAR_T, CACHE_T, KV_DTYPE, ENTRY_SZ) \ - aphrodite::gather_and_maybe_dequant_cache \ - <<>>( \ - reinterpret_cast(src_cache.data_ptr()), \ - reinterpret_cast(dst.data_ptr()), \ - block_table.const_data_ptr(), \ - cu_seq_lens.const_data_ptr(), \ - token_to_seq.const_data_ptr(), num_tokens, block_size, \ - block_table_stride, cache_block_stride, cache_entry_stride, \ - dst_entry_stride, reinterpret_cast(scale.data_ptr()), \ - seq_starts_ptr); - -#define CALL_GATHER_CACHE_576(SCALAR_T, CACHE_T, KV_DTYPE) \ - CALL_GATHER_CACHE(SCALAR_T, CACHE_T, KV_DTYPE, 576) - -#define CALL_GATHER_CACHE_320(SCALAR_T, CACHE_T, KV_DTYPE) \ - CALL_GATHER_CACHE(SCALAR_T, CACHE_T, KV_DTYPE, 320) - -// Gather sequences from the cache into the destination tensor. -// - cu_seq_lens contains the cumulative sequence lengths for each batch -// - block_table contains the cache block indices for each sequence -// - token_to_seq contains the back mapping from token_id to batch_id -// - Optionally, seq_starts (if provided) offsets the starting block index by -// (seq_starts[bid] / page_size) -void gather_and_maybe_dequant_cache( - torch::stable::Tensor const& - src_cache, // [NUM_BLOCKS, BLOCK_SIZE, ENTRIES...] - torch::stable::Tensor const& dst, // [TOT_TOKENS, ENTRIES...] - torch::stable::Tensor const& block_table, // [BATCH, BLOCK_INDICES] - torch::stable::Tensor const& cu_seq_lens, // [BATCH+1] - torch::stable::Tensor const& token_to_seq, // [MAX_TOKEN_ACROSS_CHUNKS] - int64_t num_tokens, const std::string& kv_cache_dtype, - torch::stable::Tensor const& scale, - std::optional seq_starts = std::nullopt) { - torch::stable::accelerator::DeviceGuard device_guard( - src_cache.get_device_index()); - const cudaStream_t stream = get_current_cuda_stream(); - - int32_t block_size = src_cache.size(1); - int32_t head_dim = dst.size(-1); - - STD_TORCH_CHECK( - block_table.scalar_type() == torch::headeronly::ScalarType::Int, - "block_table must be int32"); - STD_TORCH_CHECK( - cu_seq_lens.scalar_type() == torch::headeronly::ScalarType::Int, - "cu_seq_lens must be int32"); - if (seq_starts.has_value()) { - STD_TORCH_CHECK( - seq_starts.value().scalar_type() == torch::headeronly::ScalarType::Int, - "seq_starts must be int32"); - } - STD_TORCH_CHECK( - head_dim == 320 || head_dim == 576, - "gather_and_maybe_dequant_cache only support the head_dim to 320 or 576 " - "for better performance") - - STD_TORCH_CHECK(src_cache.device() == dst.device(), - "src_cache and dst must be on the same device"); - STD_TORCH_CHECK(src_cache.device() == block_table.device(), - "src_cache and block_table must be on the same device"); - STD_TORCH_CHECK(src_cache.device() == cu_seq_lens.device(), - "src_cache and cu_seq_lens must be on the same device"); - if (seq_starts.has_value()) { - STD_TORCH_CHECK(src_cache.device() == seq_starts.value().device(), - "src_cache and seq_starts must be on the same device"); - } - - int64_t block_table_stride = block_table.stride(0); - int64_t cache_block_stride = src_cache.stride(0); - int64_t cache_entry_stride = src_cache.stride(1); - int64_t dst_entry_stride = dst.stride(0); - - constexpr int32_t thread_block_size = 64; - dim3 grid(num_tokens); - dim3 block(thread_block_size); - - const int32_t* seq_starts_ptr = - seq_starts.has_value() ? seq_starts.value().const_data_ptr() - : nullptr; - - if (head_dim == 576) { - DISPATCH_BY_KV_CACHE_DTYPE(dst.scalar_type(), kv_cache_dtype, - CALL_GATHER_CACHE_576); - } else { - DISPATCH_BY_KV_CACHE_DTYPE(dst.scalar_type(), kv_cache_dtype, - CALL_GATHER_CACHE_320); - } -} - -namespace aphrodite { - -// Gather and upconvert FP8 KV cache tokens to BF16 workspace -// Similar to cp_gather_cache but specifically for FP8->BF16 conversion -__global__ void cp_gather_and_upconvert_fp8_kv_cache( - const uint8_t* __restrict__ src_cache, // [NUM_BLOCKS, BLOCK_SIZE, 656] - __nv_bfloat16* __restrict__ dst, // [total_tokens, 576] - const int32_t* __restrict__ block_table, // [num_reqs, BLOCK_INDICES] - const int32_t* __restrict__ workspace_starts, // [num_reqs] - const int32_t num_reqs, const int32_t block_size, - const int32_t total_tokens, const int64_t block_table_stride, - const int64_t cache_block_stride, const int64_t cache_entry_stride, - const int64_t dst_entry_stride, - const int32_t* __restrict__ seq_starts) { // Optional source offsets - const int flat_warp_id = (blockIdx.x * blockDim.x + threadIdx.x) >> 5; - if (flat_warp_id >= total_tokens) return; - const int lane_id = threadIdx.x & 31; - - // Binary search to find which request owns this output token - int lo = 0, hi = num_reqs - 1; - while (lo < hi) { - int mid = (lo + hi + 1) >> 1; - if (workspace_starts[mid] <= flat_warp_id) - lo = mid; - else - hi = mid - 1; - } - const int req_id = lo; - - // Compute physical token address via block table - const int out_token_id = flat_warp_id; - int token_offset = out_token_id - workspace_starts[req_id]; - if (seq_starts != nullptr) token_offset += seq_starts[req_id]; - const int cache_block_idx = token_offset / block_size; - const int offset_in_block = token_offset % block_size; - const int physical_block = - block_table[req_id * block_table_stride + cache_block_idx]; - - const uint8_t* token_ptr = src_cache + physical_block * cache_block_stride + - offset_in_block * cache_entry_stride; - - const int4* nope_src = reinterpret_cast(token_ptr); - const int4 fp8_data = nope_src[lane_id]; - - const float* scales_ptr = reinterpret_cast(token_ptr + 512); - const float scale = scales_ptr[lane_id >> 3]; - - const uint2 fp8_lo = make_uint2(fp8_data.x, fp8_data.y); - const uint2 fp8_hi = make_uint2(fp8_data.z, fp8_data.w); -#ifdef USE_ROCM - const bf16_8_t bf16_lo = - fp8::scaled_vec_conversion(fp8_lo, scale); - const bf16_8_t bf16_hi = - fp8::scaled_vec_conversion(fp8_hi, scale); -#else - const bf16_8_t bf16_lo = - fp8::scaled_vec_conversion(fp8_lo, scale, __NV_E4M3); - const bf16_8_t bf16_hi = - fp8::scaled_vec_conversion(fp8_hi, scale, __NV_E4M3); -#endif - - __nv_bfloat16* dst_ptr = dst + out_token_id * dst_entry_stride; - int4* nope_dst = reinterpret_cast(dst_ptr) + lane_id * 2; - nope_dst[0] = *reinterpret_cast(&bf16_lo); - nope_dst[1] = *reinterpret_cast(&bf16_hi); - - const int* rope_src = reinterpret_cast(token_ptr + 528); - int* rope_dst = reinterpret_cast(dst_ptr + 512); - rope_dst[lane_id] = rope_src[lane_id]; -} - -template -// Note(hc): The cp_gather_cache allows seq_starts to no longer be divisible by -// block_size. -__global__ void cp_gather_cache( - const scalar_t* __restrict__ src_cache, // [NUM_BLOCKS, BLOCK_SIZE, - // ENTRY_SIZE] - scalar_t* __restrict__ dst, // [TOT_TOKENS, ENTRY_SIZE] - const int32_t* __restrict__ block_table, // [BATCH, BLOCK_INDICES] - const int32_t* __restrict__ cu_seq_lens, // [BATCH+1] - const int32_t block_size, const int32_t entry_size, - const int64_t block_table_stride, const int64_t cache_block_stride, - const int64_t cache_entry_stride, const int64_t dst_entry_stride, - const int32_t* __restrict__ seq_starts // Optional: starting offsets per - // batch -) { - const int64_t bid = blockIdx.x; // Batch ID - const int32_t num_splits = gridDim.y; - const int32_t split = blockIdx.y; - const int32_t seq_start = cu_seq_lens[bid]; - const int32_t seq_end = cu_seq_lens[bid + 1]; - const int32_t seq_len = seq_end - seq_start; - const int32_t tot_slots = seq_len; - const int32_t split_slots = cuda_utils::ceil_div(tot_slots, num_splits); - - const int32_t split_start = split * split_slots; - const int32_t split_end = min((split + 1) * split_slots, tot_slots); - - const bool is_active_split = (split_start < tot_slots); - - if (!is_active_split) return; - - // Adjust the pointer for the block_table for this batch. - // If seq_starts is provided, compute an offset based on it - const int32_t batch_offset = bid * block_table_stride; - int32_t offset = split_start; - if (seq_starts != nullptr) { - offset += seq_starts[bid]; - } - int32_t offset_div = offset / block_size; - offset = offset % block_size; - const int32_t* batch_block_table = block_table + batch_offset; - - // Adjust dst pointer based on the cumulative sequence lengths. - dst += seq_start * dst_entry_stride; - - auto copy_entry = [&](const scalar_t* __restrict__ _src, - scalar_t* __restrict__ _dst) { - for (int i = threadIdx.x; i < entry_size; i += blockDim.x) - _dst[i] = _src[i]; - }; - - for (int pid = split_start; pid < split_end; ++pid) { - auto block_id = batch_block_table[offset_div]; - auto block_start_ptr = src_cache + block_id * cache_block_stride; - auto block_dst_ptr = dst + pid * dst_entry_stride; - copy_entry(block_start_ptr + offset * cache_entry_stride, block_dst_ptr); - offset += 1; - // bump to next block - if (offset == block_size) { - offset_div += 1; - offset = 0; - } - } -} -} // namespace aphrodite - -// Macro to dispatch the kernel based on the data type. -#define CALL_CP_GATHER_CACHE(CPY_DTYPE) \ - aphrodite::cp_gather_cache<<>>( \ - reinterpret_cast(src_cache.data_ptr()), \ - reinterpret_cast(dst.data_ptr()), \ - block_table.const_data_ptr(), \ - cu_seq_lens.const_data_ptr(), block_size, entry_size, \ - block_table_stride, cache_block_stride, cache_entry_stride, \ - dst_entry_stride, seq_starts_ptr); - -// Gather sequences from the cache into the destination tensor. -// - cu_seq_lens contains the cumulative sequence lengths for each batch -// - block_table contains the cache block indices for each sequence -// - Optionally, seq_starts (if provided) offsets the starting slot index by -// seq_starts[bid] -void cp_gather_cache( - torch::stable::Tensor const& - src_cache, // [NUM_BLOCKS, BLOCK_SIZE, ENTRIES...] - torch::stable::Tensor const& dst, // [TOT_TOKENS, ENTRIES...] - torch::stable::Tensor const& block_table, // [BATCH, BLOCK_INDICES] - torch::stable::Tensor const& cu_seq_lens, // [BATCH+1] - int64_t batch_size, - std::optional seq_starts = std::nullopt) { - torch::stable::accelerator::DeviceGuard device_guard( - src_cache.get_device_index()); - const cudaStream_t stream = get_current_cuda_stream(); - - int32_t block_size = src_cache.size(1); - int32_t entry_size = torch::stable::flatten(src_cache, 2, -1).size(2); - - STD_TORCH_CHECK( - block_table.scalar_type() == torch::headeronly::ScalarType::Int, - "block_table must be int32"); - STD_TORCH_CHECK( - cu_seq_lens.scalar_type() == torch::headeronly::ScalarType::Int, - "cu_seq_lens must be int32"); - if (seq_starts.has_value()) { - STD_TORCH_CHECK( - seq_starts.value().scalar_type() == torch::headeronly::ScalarType::Int, - "seq_starts must be int32"); - } - - STD_TORCH_CHECK(src_cache.device() == dst.device(), - "src_cache and dst must be on the same device"); - STD_TORCH_CHECK(src_cache.device() == block_table.device(), - "src_cache and block_table must be on the same device"); - STD_TORCH_CHECK(src_cache.device() == cu_seq_lens.device(), - "src_cache and cu_seq_lens must be on the same device"); - if (seq_starts.has_value()) { - STD_TORCH_CHECK(src_cache.device() == seq_starts.value().device(), - "src_cache and seq_starts must be on the same device"); - } - - int64_t block_table_stride = block_table.stride(0); - int64_t cache_block_stride = src_cache.stride(0); - int64_t cache_entry_stride = src_cache.stride(1); - int64_t dst_entry_stride = dst.stride(0); - - // Decide on the number of splits based on the batch size. - int num_splits = batch_size > 128 ? 2 : batch_size > 64 ? 4 : 16; - dim3 grid(batch_size, num_splits); - dim3 block(1024); - - STD_TORCH_CHECK(src_cache.scalar_type() == dst.scalar_type(), - "src_cache and dst must have the same dtype"); - - const int dtype_bits = src_cache.element_size() * 8; - const int32_t* seq_starts_ptr = - seq_starts.has_value() ? seq_starts.value().const_data_ptr() - : nullptr; - - if (dtype_bits == 32) { - CALL_CP_GATHER_CACHE(uint32_t); - } else if (dtype_bits == 16) { - CALL_CP_GATHER_CACHE(uint16_t); - } else if (dtype_bits == 8) { - CALL_CP_GATHER_CACHE(uint8_t); - } else { - STD_TORCH_CHECK(false, "Unsupported data type width: ", dtype_bits); - } -} - -void cp_gather_and_upconvert_fp8_kv_cache( - torch::stable::Tensor const& src_cache, // [NUM_BLOCKS, BLOCK_SIZE, 656] - torch::stable::Tensor const& dst, // [TOT_TOKENS, 576] - torch::stable::Tensor const& block_table, // [BATCH, BLOCK_INDICES] - torch::stable::Tensor const& workspace_starts, // [BATCH] - int64_t batch_size, - std::optional seq_starts = std::nullopt) { - torch::stable::accelerator::DeviceGuard device_guard( - src_cache.get_device_index()); - const cudaStream_t stream = get_current_cuda_stream(); - - int32_t block_size = src_cache.size(1); - int32_t head_dim = dst.size(1); - - STD_TORCH_CHECK( - block_table.scalar_type() == torch::headeronly::ScalarType::Int, - "block_table must be int32"); - STD_TORCH_CHECK( - workspace_starts.scalar_type() == torch::headeronly::ScalarType::Int, - "workspace_starts must be int32"); - if (seq_starts.has_value()) { - STD_TORCH_CHECK( - seq_starts.value().scalar_type() == torch::headeronly::ScalarType::Int, - "seq_starts must be int32"); - } - - STD_TORCH_CHECK(src_cache.device() == dst.device(), - "src_cache and dst must be on the same device"); - STD_TORCH_CHECK(src_cache.device() == block_table.device(), - "src_cache and block_table must be on the same device"); - STD_TORCH_CHECK(src_cache.device() == workspace_starts.device(), - "src_cache and workspace_starts must be on the same device"); - if (seq_starts.has_value()) { - STD_TORCH_CHECK(src_cache.device() == seq_starts.value().device(), - "src_cache and seq_starts must be on the same device"); - } - auto dtype = src_cache.scalar_type(); - STD_TORCH_CHECK( - dtype == torch::headeronly::ScalarType::Byte || // uint8 - dtype == torch::headeronly::ScalarType::Float8_e4m3fn || // fp8 e4m3 - dtype == torch::headeronly::ScalarType::Float8_e5m2, // fp8 e5m2 - "src_cache must be uint8, float8_e4m3fn, or float8_e5m2, but got ", - src_cache.scalar_type()); - STD_TORCH_CHECK(dst.scalar_type() == torch::headeronly::ScalarType::BFloat16, - "dst must be bfloat16"); - STD_TORCH_CHECK(head_dim == 576, "head_dim must be 576 for MLA"); - - int64_t block_table_stride = block_table.stride(0); - int64_t cache_block_stride = src_cache.stride(0); - int64_t cache_entry_stride = src_cache.stride(1); - int64_t dst_entry_stride = dst.stride(0); - - const uint8_t* src_ptr = nullptr; - if (dtype == torch::headeronly::ScalarType::Byte) { - src_ptr = src_cache.const_data_ptr(); - } else { - // float8_e4m3fn or float8_e5m2 - src_ptr = reinterpret_cast(src_cache.data_ptr()); - } - - const int total_tokens = dst.size(0); - constexpr int warps_per_block = 8; - const int grid_size = (total_tokens + warps_per_block - 1) / warps_per_block; - const int block_size_threads = warps_per_block * 32; // 256 threads - const int32_t* seq_starts_ptr = - seq_starts.has_value() ? seq_starts.value().const_data_ptr() - : nullptr; - - aphrodite::cp_gather_and_upconvert_fp8_kv_cache<<< - grid_size, block_size_threads, 0, stream>>>( - src_ptr, reinterpret_cast<__nv_bfloat16*>(dst.data_ptr()), - block_table.const_data_ptr(), - workspace_starts.const_data_ptr(), - static_cast(batch_size), block_size, total_tokens, - block_table_stride, cache_block_stride, cache_entry_stride, - dst_entry_stride, seq_starts_ptr); -} - -// Macro to dispatch the kernel based on the data type. -#define CALL_INDEXER_K_QUANT_AND_CACHE(KV_T, CACHE_T, KV_DTYPE) \ - aphrodite::indexer_k_quant_and_cache_kernel \ - <<>>( \ - reinterpret_cast(k.data_ptr()), \ - reinterpret_cast(kv_cache.data_ptr()), \ - slot_mapping.const_data_ptr(), head_dim, quant_block_size, \ - cache_block_size, cache_block_stride, use_ue8m0); - -void indexer_k_quant_and_cache( - torch::stable::Tensor& k, // [num_tokens, head_dim] - torch::stable::Tensor& kv_cache, // [num_blocks, block_size, cache_stride] - torch::stable::Tensor& slot_mapping, // [num_tokens] - int64_t quant_block_size, // quantization block size - const std::string& scale_fmt) { - int num_tokens = k.size(0); - int head_dim = k.size(1); - int cache_block_size = kv_cache.size(1); - int64_t cache_block_stride = kv_cache.stride(0); - bool use_ue8m0 = scale_fmt == "ue8m0"; - - STD_TORCH_CHECK(k.device() == kv_cache.device(), - "k and kv_cache must be on the same device"); - STD_TORCH_CHECK(k.device() == slot_mapping.device(), - "k and slot_mapping must be on the same device"); - STD_TORCH_CHECK(head_dim % quant_block_size == 0, - "head_dim must be divisible by quant_block_size"); - - constexpr int vec_size = 4; - dim3 grid(num_tokens, (head_dim + quant_block_size * vec_size - 1) / - (quant_block_size * vec_size)); - dim3 block(32, vec_size); - const torch::stable::accelerator::DeviceGuard device_guard( - k.get_device_index()); - const cudaStream_t stream = get_current_cuda_stream(); - - static const std::string kv_cache_dtype = "fp8_e4m3"; - DISPATCH_BY_KV_CACHE_DTYPE(k.scalar_type(), kv_cache_dtype, - CALL_INDEXER_K_QUANT_AND_CACHE); -} - -// Macro to dispatch the kernel based on the data amount. -#define CALL_CP_GATHER_INDEXER_K_QUANT_CACHE(BLOCK_Y_SIZE) \ - aphrodite::cp_gather_indexer_k_quant_cache_kernel \ - <<>>( \ - reinterpret_cast(kv_cache.data_ptr()), \ - reinterpret_cast(dst_k.data_ptr()), \ - reinterpret_cast(dst_scale.data_ptr()), \ - block_table.const_data_ptr(), \ - cu_seq_lens.const_data_ptr(), batch_size, dst_k.stride(0), \ - dst_k.size(1), kv_cache.stride(0), kv_cache.stride(1), \ - kv_cache.size(1), block_table.size(1), num_tokens, \ - quant_block_size); - -void cp_gather_indexer_k_quant_cache( - const torch::stable::Tensor& - kv_cache, // [num_blocks, block_size, cache_stride] - torch::stable::Tensor& dst_k, // [num_tokens, head_dim] - torch::stable::Tensor& - dst_scale, // [num_tokens, head_dim / quant_block_size * 4] - const torch::stable::Tensor& block_table, // [batch_size, num_blocks] - const torch::stable::Tensor& cu_seq_lens // [batch_size + 1] -) { - int batch_size = block_table.size(0); - int num_tokens = dst_k.size(0); - int head_dim = dst_k.size(1); - int quant_block_size = head_dim * 4 / dst_scale.size(1); - - STD_TORCH_CHECK(kv_cache.device() == dst_k.device(), - "kv_cache and dst_k must be on the same device"); - STD_TORCH_CHECK(kv_cache.device() == dst_scale.device(), - "kv_cache and dst_scale must be on the same device"); - STD_TORCH_CHECK(kv_cache.device() == block_table.device(), - "kv_cache and block_table must be on the same device"); - STD_TORCH_CHECK(kv_cache.device() == cu_seq_lens.device(), - "kv_cache and cu_seq_lens must be on the same device"); - STD_TORCH_CHECK(head_dim % quant_block_size == 0, - "head_dim must be divisible by quant_block_size"); - - constexpr int vec_size = 16; - const torch::stable::accelerator::DeviceGuard device_guard( - kv_cache.get_device_index()); - const cudaStream_t stream = get_current_cuda_stream(); - - if (num_tokens < 32) { - CALL_CP_GATHER_INDEXER_K_QUANT_CACHE(1); - } else if (num_tokens < 64) { - CALL_CP_GATHER_INDEXER_K_QUANT_CACHE(2); - } else if (num_tokens < 128) { - CALL_CP_GATHER_INDEXER_K_QUANT_CACHE(4); - } else if (num_tokens < 256) { - CALL_CP_GATHER_INDEXER_K_QUANT_CACHE(8); - } else if (num_tokens < 512) { - CALL_CP_GATHER_INDEXER_K_QUANT_CACHE(16); - } else { - CALL_CP_GATHER_INDEXER_K_QUANT_CACHE(32); - } -} - -// Concatenate ql_nope and q_pe into a contiguous q_out tensor for MLA/DSA. -// Replaces torch.cat((ql_nope, q_pe), dim=-1). -void concat_mla_q( - torch::stable::Tensor& ql_nope, // [num_tokens, num_heads, nope_dim] - torch::stable::Tensor& q_pe, // [num_tokens, num_heads, rope_dim] - torch::stable::Tensor& q_out // [num_tokens, num_heads, nope_dim + - // rope_dim] -) { - const int num_tokens = ql_nope.size(0); - const int num_heads = ql_nope.size(1); - const int nope_dim = ql_nope.size(2); - const int rope_dim = q_pe.size(2); - - STD_TORCH_CHECK(nope_dim % 512 == 0, - "nope_dim must be a multiple of 512, got ", nope_dim); - STD_TORCH_CHECK(rope_dim == 64, "rope_dim must be 64, got ", rope_dim); - STD_TORCH_CHECK(q_out.size(2) == nope_dim + rope_dim); - - STD_TORCH_CHECK(ql_nope.stride(2) == 1, - "ql_nope must have stride 1 in dim 2"); - STD_TORCH_CHECK(q_pe.stride(2) == 1, "q_pe must have stride 1 in dim 2"); - STD_TORCH_CHECK(q_out.stride(2) == 1, "q_out must have stride 1 in dim 2"); - STD_TORCH_CHECK( - ql_nope.scalar_type() == torch::headeronly::ScalarType::Half || - ql_nope.scalar_type() == torch::headeronly::ScalarType::BFloat16, - "ql_nope must be float16 or bfloat16 dtype"); - - if (num_tokens == 0) return; - - constexpr int warps_per_block = 8; - const int total_warps = num_tokens * num_heads; - const int grid_size = (total_warps + warps_per_block - 1) / warps_per_block; - const int block_size = warps_per_block * 32; - - const torch::stable::accelerator::DeviceGuard device_guard( - ql_nope.get_device_index()); - const cudaStream_t stream = get_current_cuda_stream(); - - APHRODITE_STABLE_DISPATCH_HALF_TYPES( - ql_nope.scalar_type(), "concat_mla_q", [&] { - aphrodite::ConcatMLAQKernel - <<>>( - q_out.mutable_data_ptr(), - ql_nope.const_data_ptr(), - q_pe.const_data_ptr(), num_tokens, num_heads, - q_out.stride(0), q_out.stride(1), ql_nope.stride(0), - ql_nope.stride(1), q_pe.stride(0), q_pe.stride(1)); - }); -} diff --git a/csrc/libtorch_stable/custom_all_reduce.cu.orig b/csrc/libtorch_stable/custom_all_reduce.cu.orig deleted file mode 100644 index 6d43964edd..0000000000 --- a/csrc/libtorch_stable/custom_all_reduce.cu.orig +++ /dev/null @@ -1,202 +0,0 @@ -#include "torch_utils.h" - -#include -#include -#include -#include -#include -#include - -#include "custom_all_reduce.cuh" - -// Fake pointer type, must match fptr_t type in ops.h. -// We use this type alias to indicate when pointers are passed in as int64_t. -using fptr_t = int64_t; -static_assert(sizeof(void*) == sizeof(fptr_t)); - -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 % 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]; - for (int i = 0; i < world_size; i++) { - ipc_ptrs[i] = reinterpret_cast(fake_ipc_ptrs[i]); - } - return (fptr_t) new aphrodite::CustomAllreduce( - ipc_ptrs, rank_data.mutable_data_ptr(), rank_data.numel(), rank, - world_size, fully_connected); -} - -/** - * Make sure tensor t's data lies completely within ((char)t.data_ptr()) + - * t.numel() * t.element_size(). This is slightly weaker than t.is_contiguous() - * because it allows transpose of contiguous slice (i.e. slicing the first - * dimension). Currently, we require this because stride information is not - * passed into the kernels and we treat input tensors as flat. - * - * Examples - * A = torch.zeros(3, 3, 3) - * 1. A: OK - * 2. A[1:]: OK - * 3. A.permute(2, 0, 1): OK - * 4. A[1:].permute(2, 0, 1): OK - * 5. A[None].expand(2, -1, -1, -1): Not OK - * 6. A[:, 1:, 1:]: Not OK - */ -bool _is_weak_contiguous(torch::stable::Tensor& t) { - if (t.is_contiguous()) { - return true; - } - int64_t storage_nbytes = 0; - TORCH_ERROR_CODE_CHECK(aoti_torch_get_storage_size(t.get(), &storage_nbytes)); - return storage_nbytes - t.storage_offset() * t.element_size() == - static_cast(t.numel() * t.element_size()); -} - -/** - * Performs an out-of-place allreduce and stores result in out. - * - * If _reg_buffer is null, assumes inp.data_ptr() is already IPC-registered. - * Otherwise, _reg_buffer is assumed to be IPC-registered and inp is first - * copied into _reg_buffer. - */ -void all_reduce(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()) == (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); - if (reg_buffer) { - STD_TORCH_CHECK((input_size) <= (reg_buffer_sz_bytes)); - STD_CUDA_CHECK(cudaMemcpyAsync(reg_buffer, inp.const_data_ptr(), input_size, - cudaMemcpyDeviceToDevice, stream)); - } else { - reg_buffer = inp.mutable_data_ptr(); - } - switch (out.scalar_type()) { - case torch::headeronly::ScalarType::Float: { - fa->allreduce(stream, reinterpret_cast(reg_buffer), - reinterpret_cast(out.mutable_data_ptr()), - out.numel()); - break; - } - case torch::headeronly::ScalarType::Half: { - fa->allreduce(stream, reinterpret_cast(reg_buffer), - reinterpret_cast(out.mutable_data_ptr()), - out.numel()); - break; - } -#if (__CUDA_ARCH__ >= 800 || !defined(__CUDA_ARCH__)) - case torch::headeronly::ScalarType::BFloat16: { - fa->allreduce( - stream, reinterpret_cast(reg_buffer), - reinterpret_cast(out.mutable_data_ptr()), out.numel()); - break; - } -#endif - default: - throw std::runtime_error( - "custom allreduce only supports float32, float16 and bfloat16"); - } -} - -void dispose(fptr_t _fa) { - delete reinterpret_cast(_fa); -} - -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]; - for (int i = 0; i < fake_ipc_ptrs.size(); i++) { - ipc_ptrs[i] = reinterpret_cast(fake_ipc_ptrs[i]); - } - fa->register_buffer(ipc_ptrs); -} - -// Use vector to represent byte data for python binding compatibility. -std::tuple, std::vector> -get_graph_buffer_ipc_meta(fptr_t _fa) { - auto fa = reinterpret_cast(_fa); - auto [handle, offsets] = fa->get_graph_buffer_ipc_meta(); - std::vector bytes(handle.begin(), handle.end()); - return std::make_tuple(bytes, offsets); -} - -// Use vector to represent byte data for python binding compatibility. -void register_graph_buffers(fptr_t _fa, - const std::vector>& handles, - const std::vector>& offsets) { - auto fa = reinterpret_cast(_fa); - std::vector bytes; - bytes.reserve(handles.size()); - for (int i = 0; i < handles.size(); i++) { - bytes.emplace_back(handles[i].begin(), handles[i].end()); - } - bytes.reserve(handles.size()); - fa->register_graph_buffers(bytes, offsets); -} - -std::tuple allocate_shared_buffer_and_handle( - int64_t size) { - int device_index; - STD_CUDA_CHECK(cudaGetDevice(&device_index)); - const torch::stable::accelerator::DeviceGuard device_guard(device_index); - void* buffer; - cudaStreamCaptureMode mode = cudaStreamCaptureModeRelaxed; - const cudaStream_t stream = get_current_cuda_stream(device_index); - STD_CUDA_CHECK(cudaThreadExchangeStreamCaptureMode(&mode)); - - // Allocate buffer -#if defined(USE_ROCM) - // data buffers need to be "uncached" for signal on MI200 - STD_CUDA_CHECK( - hipExtMallocWithFlags((void**)&buffer, size, hipDeviceMallocUncached)); -#else - STD_CUDA_CHECK(cudaMalloc((void**)&buffer, size)); -#endif - STD_CUDA_CHECK(cudaMemsetAsync(buffer, 0, size, stream)); - STD_CUDA_CHECK(cudaStreamSynchronize(stream)); - STD_CUDA_CHECK(cudaThreadExchangeStreamCaptureMode(&mode)); - - // Create IPC memhandle for the allocated buffer. - // Will use it in open_mem_handle. - auto handle = torch::stable::empty( - {static_cast(sizeof(cudaIpcMemHandle_t))}, - torch::headeronly::ScalarType::Byte, std::nullopt, - torch::stable::Device(torch::stable::DeviceType::CPU)); - STD_CUDA_CHECK(cudaIpcGetMemHandle( - (cudaIpcMemHandle_t*)handle.mutable_data_ptr(), buffer)); - - return std::make_tuple(reinterpret_cast(buffer), handle); -} - -fptr_t open_mem_handle(torch::stable::Tensor& mem_handle) { - void* ipc_ptr; - STD_CUDA_CHECK(cudaIpcOpenMemHandle( - (void**)&ipc_ptr, - *((const cudaIpcMemHandle_t*)mem_handle.const_data_ptr()), - cudaIpcMemLazyEnablePeerAccess)); - return reinterpret_cast(ipc_ptr); -} - -void free_shared_buffer(fptr_t buffer) { - STD_CUDA_CHECK(cudaFree(reinterpret_cast(buffer))); -} diff --git a/csrc/libtorch_stable/moe/grouped_topk_kernels.cu.orig b/csrc/libtorch_stable/moe/grouped_topk_kernels.cu.orig deleted file mode 100644 index dadfbc3a5c..0000000000 --- a/csrc/libtorch_stable/moe/grouped_topk_kernels.cu.orig +++ /dev/null @@ -1,1124 +0,0 @@ -/* - * Adapted from - * https://github.com/NVIDIA/TensorRT-LLM/blob/v1.3.0rc2/cpp/tensorrt_llm/kernels/noAuxTcKernels.cu - * Copyright (c) 2025, The vLLM team. - * SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & - * AFFILIATES. All rights reserved. SPDX-License-Identifier: Apache-2.0 - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ -#include "moeTopKFuncs.cuh" - -#include -#include - -#include "libtorch_stable/torch_utils.h" - -#include -#include -#include -#include -#include -#include -#include -namespace cg = cooperative_groups; - -namespace aphrodite { -namespace moe { - -constexpr unsigned FULL_WARP_MASK = 0xffffffff; -static constexpr int WARP_SIZE = 32; -static constexpr int NumNemotronExperts = 512; -static constexpr int NumKimiK2Experts = 384; -static constexpr int NumDeepseekExperts = 256; -static constexpr int MaxSupportedExpertCount = - std::max({NumNemotronExperts, NumKimiK2Experts, NumDeepseekExperts}); -static constexpr int MaxNumExpertsUnit = 128; -static constexpr int NumTopGroupScores = 2; -static constexpr int DefaultMaxNumTopExperts = 8; -static constexpr int MaxSupportedTopExperts = 22; -static constexpr int MaxNumTopGroups = 4; -// The empirical value for small batch -static constexpr int PDLEnableTokens = 16; - -namespace warp_topk { - -template -__host__ __device__ constexpr T round_up_to_multiple_of(T len) { - if (len == 0) { - return 0; - } - return ((len - 1) / size + 1) * size; -} - -template -constexpr __host__ __device__ bool isPowerOf2(T v) { - return (v && !(v & (v - 1))); -} - -template -__forceinline__ __device__ bool is_better_than(T val, T baseline) { - return (val > baseline && greater) || (val < baseline && !greater); -} - -template -__forceinline__ __device__ bool is_better_than(T val, T baseline, idxT index, - idxT baseline_index) { - bool res = (val > baseline && greater) || (val < baseline && !greater); - if (val == baseline) { - res = (index < baseline_index && greater) || - (index < baseline_index && !greater); - } - return res; -} - -template -struct BitonicMerge { - // input should be a bitonic sequence, and sort it to be a monotonic sequence - __device__ static void merge(T* __restrict__ val_arr, - idxT* __restrict__ idx_arr) { - static_assert(isPowerOf2(size)); - static_assert(size >= 2 * WARP_SIZE); - constexpr int arr_len = size / WARP_SIZE; - - constexpr int stride = arr_len / 2; - for (int i = 0; i < stride; ++i) { - int const other_i = i + stride; - T& val = val_arr[i]; - T& other_val = val_arr[other_i]; - bool is_better; - if constexpr (is_stable) { - is_better = is_better_than(val, other_val, idx_arr[i], - idx_arr[other_i]); - } else { - is_better = is_better_than(val, other_val); - } - - if (is_better) { - T tmp = val; - val = other_val; - other_val = tmp; - - idxT tmp2 = idx_arr[i]; - idx_arr[i] = idx_arr[other_i]; - idx_arr[other_i] = tmp2; - } - } - - BitonicMerge::merge( - val_arr, idx_arr); - BitonicMerge::merge( - val_arr + arr_len / 2, idx_arr + arr_len / 2); - } -}; - -template -struct BitonicSort { - __device__ static void sort(T* __restrict__ val_arr, - idxT* __restrict__ idx_arr) { - static_assert(isPowerOf2(size)); - static_assert(size >= 2 * WARP_SIZE); - constexpr int arr_len = size / WARP_SIZE; - - BitonicSort::sort(val_arr, idx_arr); - BitonicSort::sort( - val_arr + arr_len / 2, idx_arr + arr_len / 2); - BitonicMerge::merge( - val_arr, idx_arr); - } -}; - -template -struct BitonicSort<32, ascending, T, idxT, is_stable> { - __device__ static void sort(T* __restrict__ val_arr, - idxT* __restrict__ idx_arr) { - int const lane = threadIdx.x % WARP_SIZE; - - // ascending doesn't matter before merging since all we need is a bitonic - // sequence - for (int stage = 0; stage < 4; ++stage) { - for (int stride = (1 << stage); stride > 0; stride /= 2) { - bool reverse = (lane >> stage) & 2; - bool is_second = lane & stride; - - T other = __shfl_xor_sync(FULL_WARP_MASK, *val_arr, stride); - idxT other_idx = __shfl_xor_sync(FULL_WARP_MASK, *idx_arr, stride); - - bool is_better; - if constexpr (is_stable) { - if constexpr (ascending) { - is_better = ((*val_arr > other) || - ((*val_arr == other) && (*idx_arr < other_idx))) != - (reverse != is_second); - } else { - is_better = ((*val_arr > other) || - ((*val_arr == other) && (*idx_arr > other_idx))) != - (reverse != is_second); - } - } else { - is_better = (*val_arr != other && - (*val_arr > other) != (reverse != is_second)); - } - if (is_better) { - *val_arr = other; - *idx_arr = other_idx; - } - } - } - - BitonicMerge<32, ascending, ascending, T, idxT, is_stable>::merge(val_arr, - idx_arr); - } -}; - -template -struct BitonicMerge<32, ascending, reverse, T, idxT, is_stable> { - __device__ static void merge(T* __restrict__ val_arr, - idxT* __restrict__ idx_arr) { - int const lane = threadIdx.x % WARP_SIZE; - for (int stride = WARP_SIZE / 2; stride > 0; stride /= 2) { - bool is_second = lane & stride; - T& val = *val_arr; - T other = __shfl_xor_sync(FULL_WARP_MASK, val, stride); - idxT& idx = *idx_arr; - idxT other_idx = __shfl_xor_sync(FULL_WARP_MASK, idx, stride); - - bool is_better; - if constexpr (is_stable) { - if constexpr (ascending) { - is_better = ((*val_arr > other) || - ((*val_arr == other) && (*idx_arr < other_idx))) == - (reverse != is_second); // for min - } else { - is_better = ((*val_arr > other) || - ((*val_arr == other) && (*idx_arr > other_idx))) == - (reverse != is_second); // for max - } - } else { - is_better = - (val != other && ((val > other) == (ascending != is_second))); - } - - if (is_better) { - val = other; - idx = other_idx; - } - } - } -}; - -template -class WarpSort { - public: - __device__ WarpSort(idxT k, T dummy) - : lane_(threadIdx.x % WARP_SIZE), k_(k), dummy_(dummy) { - static_assert(capacity >= WARP_SIZE && isPowerOf2(capacity)); - - for (int i = 0; i < max_arr_len_; ++i) { - val_arr_[i] = dummy_; - idx_arr_[i] = 0; - } - } - - // load and merge k sorted values - __device__ void load_sorted(T const* __restrict__ in, - idxT const* __restrict__ in_idx, idxT start) { - idxT idx = start + WARP_SIZE - 1 - lane_; - for (int i = max_arr_len_ - 1; i >= 0; --i, idx += WARP_SIZE) { - if (idx < start + k_) { - T t = in[idx]; - bool is_better; - if constexpr (is_stable) { - is_better = - is_better_than(t, val_arr_[i], in_idx[idx], idx_arr_[i]); - } else { - is_better = is_better_than(t, val_arr_[i]); - } - if (is_better) { - val_arr_[i] = t; - idx_arr_[i] = in_idx[idx]; - } - } - } - - BitonicMerge::merge( - val_arr_, idx_arr_); - } - - __device__ void dump(T* __restrict__ out, idxT* __restrict__ out_idx) const { - for (int i = 0; i < max_arr_len_; ++i) { - idxT out_i = i * WARP_SIZE + lane_; - if (out_i < k_) { - out[out_i] = val_arr_[i]; - out_idx[out_i] = idx_arr_[i]; - } - } - } - - __device__ void dumpIdx(idxT* __restrict__ out_idx) const { - for (int i = 0; i < max_arr_len_; ++i) { - idxT out_i = i * WARP_SIZE + lane_; - if (out_i < k_) { - out_idx[out_i] = idx_arr_[i]; - } - } - } - - // Accessors for per-lane selected value/index. - // NOTE: For the common case `capacity == WARP_SIZE`, `max_arr_len_ == 1` - // and callers should use `i == 0`. - __device__ __forceinline__ idxT get_idx(int i = 0) const { - return idx_arr_[i]; - } - - __device__ __forceinline__ T get_val(int i = 0) const { return val_arr_[i]; } - - protected: - static constexpr int max_arr_len_ = capacity / WARP_SIZE; - - T val_arr_[max_arr_len_]; - idxT idx_arr_[max_arr_len_]; - - int const lane_; - idxT const k_; - T const dummy_; - -}; // end class WarpSort - -template -class WarpSelect : public WarpSort { - public: - __device__ WarpSelect(idxT k, T dummy) - : WarpSort(k, dummy), - k_th_(dummy), - k_th_idx_(0), - k_th_lane_((k - 1) % WARP_SIZE) { - extern __shared__ char smem_buf[]; // extern __shared__ T smem_buf[]; - - int const num_of_warp = blockDim.x / WARP_SIZE; - int const warp_id = threadIdx.x / WARP_SIZE; - val_smem_ = reinterpret_cast(smem_buf); - val_smem_ += warp_id * WARP_SIZE; - idx_smem_ = reinterpret_cast( - smem_buf + - round_up_to_multiple_of<256>(num_of_warp * sizeof(T) * WARP_SIZE)); - idx_smem_ += warp_id * WARP_SIZE; - } - - __device__ void add(T const* in, idxT start, idxT end) { - idxT const end_for_fullwarp = - round_up_to_multiple_of(end - start) + start; - for (idxT i = start + lane_; i < end_for_fullwarp; i += WARP_SIZE) { - T val = (i < end) ? in[i] : dummy_; - add(val, i); - } - } - - __device__ void add(T val, idxT idx) { - bool do_add; - if constexpr (is_stable) { - do_add = is_better_than(val, k_th_, idx, k_th_idx_); - } else { - do_add = is_better_than(val, k_th_); - } - - uint32_t mask = __ballot_sync(FULL_WARP_MASK, do_add); - if (mask == 0) { - return; - } - - int pos = smem_buf_len_ + __popc(mask & ((0x1u << lane_) - 1)); - if (do_add && pos < WARP_SIZE) { - val_smem_[pos] = val; - idx_smem_[pos] = idx; - do_add = false; - } - smem_buf_len_ += __popc(mask); - if (smem_buf_len_ >= WARP_SIZE) { - __syncwarp(); - merge_buf_(val_smem_[lane_], idx_smem_[lane_]); - smem_buf_len_ -= WARP_SIZE; - } - if (do_add) { - pos -= WARP_SIZE; - val_smem_[pos] = val; - idx_smem_[pos] = idx; - } - __syncwarp(); - } - - __device__ void done() { - if (smem_buf_len_) { - T val = (lane_ < smem_buf_len_) ? val_smem_[lane_] : dummy_; - idxT idx = (lane_ < smem_buf_len_) ? idx_smem_[lane_] : 0; - merge_buf_(val, idx); - } - } - - private: - __device__ void set_k_th_() { - k_th_ = __shfl_sync(FULL_WARP_MASK, val_arr_[max_arr_len_ - 1], k_th_lane_); - if constexpr (is_stable) { - k_th_idx_ = - __shfl_sync(FULL_WARP_MASK, idx_arr_[max_arr_len_ - 1], k_th_lane_); - } - } - - __device__ void merge_buf_(T val, idxT idx) { - BitonicSort::sort(&val, &idx); - - T& old = val_arr_[max_arr_len_ - 1]; - - bool is_better; - if constexpr (is_stable) { - is_better = - is_better_than(val, old, idx, idx_arr_[max_arr_len_ - 1]); - } else { - is_better = is_better_than(val, old); - } - - if (is_better) { - old = val; - idx_arr_[max_arr_len_ - 1] = idx; - } - - BitonicMerge::merge( - val_arr_, idx_arr_); - - set_k_th_(); - } - - using WarpSort::max_arr_len_; - using WarpSort::val_arr_; - using WarpSort::idx_arr_; - using WarpSort::lane_; - using WarpSort::k_; - using WarpSort::dummy_; - - T* val_smem_; - idxT* idx_smem_; - int smem_buf_len_ = 0; - - T k_th_; - idxT k_th_idx_; - int const k_th_lane_; -}; // end class WarpSelect -} // namespace warp_topk - -template -__device__ inline T_OUT cuda_cast(T_IN val) { - return val; -} - -template <> -__device__ inline float cuda_cast(__nv_bfloat16 val) { - return __bfloat162float(val); -} - -template -__device__ inline T neg_inf() { - // cuda::std::numeric_limits::infinity() returns `0` for [T=bf16 or fp16] - // so we need to cast from fp32 - return cuda_cast(-cuda::std::numeric_limits::infinity()); -} - -template -__device__ inline bool is_finite(const T val) { -#if (__CUDACC_VER_MAJOR__ * 10000 + __CUDACC_VER_MINOR__ * 100 >= 120800) - return cuda::std::isfinite(val); -#else - return isfinite(cuda_cast(val)); -#endif -} - -// Scoring function enums -enum ScoringFunc { - SCORING_NONE = 0, // no activation function - SCORING_SIGMOID = 1 // apply sigmoid -}; - -// Efficient sigmoid approximation from TensorRT-LLM -__device__ inline float sigmoid_accurate(float x) { - return 0.5f * tanhf(0.5f * x) + 0.5f; -} - -template -__device__ inline T apply_sigmoid(T val) { - float f = cuda_cast(val); - return cuda_cast(sigmoid_accurate(f)); -} - -template -__device__ inline T apply_scoring(T val) { - if constexpr (SF == SCORING_NONE) { - return val; - } else if constexpr (SF == SCORING_SIGMOID) { - return apply_sigmoid(val); - } else { - static_assert(SF == SCORING_NONE || SF == SCORING_SIGMOID, - "Unsupported ScoringFunc in apply_scoring"); - return val; - } -} - -template -__device__ void topk_with_k2(T* output, T const* input, BiasT const* bias, - cg::thread_block_tile<32> const& tile, - int32_t const lane_id, - int const num_experts_per_group) { - // Get the top2 per thread - T largest = neg_inf(); - T second_largest = neg_inf(); - - if (num_experts_per_group > WARP_SIZE) { - for (int i = lane_id; i < num_experts_per_group; i += WARP_SIZE) { - T value = apply_scoring(input[i]); - value = value + static_cast(bias[i]); - - if (value > largest) { - second_largest = largest; - largest = value; - } else if (value > second_largest) { - second_largest = value; - } - } - } else { - for (int i = lane_id; i < num_experts_per_group; i += WARP_SIZE) { - T value = apply_scoring(input[i]); - value = value + static_cast(bias[i]); - largest = value; - } - } - // Get the top2 warpwise - T max1 = cg::reduce(tile, largest, cg::greater()); - - T max2 = max1; - bool equal_to_max1 = (max1 == largest); - - int count_max1 = __popc(__ballot_sync(FULL_WARP_MASK, equal_to_max1)); - - if (count_max1 == 1) { - largest = (largest == max1) ? second_largest : largest; - max2 = cg::reduce(tile, largest, cg::greater()); - } - - if (lane_id == 0) { - *output = max1 + max2; - } -} - -template -__global__ void grouped_topk_fused_kernel( - T* scores, float* topk_values, IdxT* topk_indices, BiasT const* bias, - int64_t const num_tokens, int64_t const num_experts, int64_t const n_group, - int64_t const topk_group, int64_t const topk, bool renormalize, - double routed_scaling_factor) { - int32_t const token_id = static_cast(blockIdx.x); - if (token_id >= num_tokens) { - return; - } - - int32_t const warp_id = threadIdx.x / WARP_SIZE; - int32_t const lane_id = threadIdx.x % WARP_SIZE; - - int32_t const n_group_i32 = static_cast(n_group); - int32_t const topk_group_i32 = static_cast(topk_group); - int32_t const topk_i32 = static_cast(topk); - int32_t const num_experts_i32 = static_cast(num_experts); - - int32_t const num_warps = blockDim.x / WARP_SIZE; - if (warp_id >= n_group_i32 || num_warps < n_group_i32) { - return; - } - - int32_t const num_experts_per_group = num_experts_i32 / n_group_i32; - - T* scores_token = scores + static_cast(token_id) * num_experts; - - cg::thread_block block = cg::this_thread_block(); - cg::thread_block_tile<32> tile = cg::tiled_partition<32>(block); - - extern __shared__ char smem_buf[]; - // warpSelect internal staging buffer layout - size_t const val_bytes = - static_cast(num_warps) * WARP_SIZE * sizeof(T); - size_t const val_bytes_aligned = - warp_topk::round_up_to_multiple_of<256>(val_bytes); - size_t const idx_bytes = - static_cast(num_warps) * WARP_SIZE * sizeof(int32_t); - size_t const internal_bytes = val_bytes_aligned + idx_bytes; - - // user-managed shared memory starts after warpSelect internal staging. - uintptr_t ptr_u = reinterpret_cast(smem_buf + internal_bytes); - ptr_u = (ptr_u + 15) & ~static_cast(15); // align to 16B - T* s_group_scores = reinterpret_cast(ptr_u); - -#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) - cudaGridDependencySynchronize(); // I think all prolog can be put before - // acqbulk because it's ptr arithmetic -#endif - - // phase 1: per-group scan - int32_t const group_offset = warp_id * num_experts_per_group; - topk_with_k2(s_group_scores + warp_id, - scores_token + group_offset, bias + group_offset, - tile, lane_id, num_experts_per_group); - - __syncthreads(); - - // phase 2: warp0 selects groups + merges candidates to final topk - if (warp_id != 0) { - return; - } - - topk_values += static_cast(token_id) * topk; - topk_indices += static_cast(token_id) * topk; - - // select topk_group groups by group score - warp_topk::WarpSelect - group_sel(static_cast(topk_group_i32), neg_inf()); - - // all lanes must participate in WarpSelect::add(). - T gscore = (lane_id < n_group_i32) ? s_group_scores[lane_id] : neg_inf(); - group_sel.add(gscore, lane_id); - group_sel.done(); - - // proceed only if the k-th selected group score is not -inf - bool proceed = false; - if (topk_group_i32 > 0) { - int const kth_lane = topk_group_i32 - 1; - // broadcast the k-th selected group score to all lanes - T kth_val = __shfl_sync(FULL_WARP_MASK, group_sel.get_val(0), kth_lane); - proceed = (kth_val != neg_inf()); - } - - if (!proceed) { - for (int i = lane_id; i < topk_i32; i += WARP_SIZE) { - topk_indices[i] = static_cast(i); - topk_values[i] = 1.0f / static_cast(topk_i32); - } -#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) - cudaTriggerProgrammaticLaunchCompletion(); -#endif - return; - } - - // merge per-group topk candidates for selected groups, then select topk - warp_topk::WarpSelect - expert_sel(static_cast(topk_i32), neg_inf()); - - // selected group ids reside in lanes [0, topk_group) - int32_t sel_gid_lane = (lane_id < topk_group_i32) ? group_sel.get_idx(0) : 0; - - // add candidates from selected groups to expert_sel - for (int32_t g = 0; g < topk_group_i32; ++g) { - int32_t gid = __shfl_sync(FULL_WARP_MASK, sel_gid_lane, g); - int32_t const offset = gid * num_experts_per_group; - int32_t const align_num_experts_per_group = - warp_topk::round_up_to_multiple_of(num_experts_per_group); - for (int32_t i = lane_id; i < align_num_experts_per_group; i += WARP_SIZE) { - // all lanes must call `add()` the same number of times. - T cand = neg_inf(); - int32_t idx = 0; - if (i < num_experts_per_group) { - idx = offset + i; - T input = scores_token[idx]; - if (is_finite(input)) { - T score = apply_scoring(input); - cand = score + static_cast(bias[idx]); - } - } - expert_sel.add(cand, idx); - } - } - expert_sel.done(); - - // compute unbiased routing weights + optional renorm. - float lane_unbiased = 0.0f; - IdxT lane_idx = 0; - if (lane_id < topk_i32) { - lane_idx = static_cast(expert_sel.get_idx(0)); - T in = scores_token[static_cast(lane_idx)]; - lane_unbiased = cuda_cast(apply_scoring(in)); - } - - float topk_sum = 1e-20f; - if (renormalize) { - topk_sum += cg::reduce(tile, lane_unbiased, cg::plus()); - } - - float scale = static_cast(routed_scaling_factor); - if (renormalize) { - scale /= topk_sum; - } - - if (lane_id < topk_i32) { - topk_indices[lane_id] = lane_idx; - topk_values[lane_id] = lane_unbiased * scale; - } - -#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) - cudaTriggerProgrammaticLaunchCompletion(); -#endif -} - -template -__global__ void grouped_topk_fused_small_expert_count_kernel( - T* scores, float* topkValues, IdxT* topkIndices, BiasT const* routingBias, - int64_t const numTokens, int64_t const numGroup, int64_t const topkGroup, - int64_t const topk, int64_t const numExperts, - int64_t const numExpertsPerGroup, bool const renormalize, - double const routedScalingFactor) { -#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) - cudaGridDependencySynchronize(); -#endif - // declare shared memory structure - // number of experts is bounded by number of threads - __shared__ float __attribute((aligned(128))) smemScoreSigmoid[MaxNumExperts]; - __shared__ float __attribute((aligned(128))) smemScoreBias[MaxNumExperts]; - // number of expert groups is bounded by number of warps - int constexpr NumWarps = MaxNumExperts / WARP_SIZE; - __shared__ float __attribute((aligned(128))) smemGroupScores[NumWarps]; - - // needed for warp reduce - auto block = cg::this_thread_block(); - auto warp = cg::tiled_partition(block); - - // for the final reduction of weight norm, only some lanes need to participate - int32_t laneIdx = threadIdx.x % WARP_SIZE; - int32_t warpIdx = __shfl_sync(0xffffffff, threadIdx.x / WARP_SIZE, 0); - - if constexpr (UseGroups) { - if (warpIdx >= numGroup) { - return; - } - } - // note that for invalid scores, we simply use a negative value: - // they work well even with the compacted format used in topK, and - // sigmoid / bias activated scores cannot be negative - const float invalidScoreFloat = float{-INFINITY}; - - // load bias already; each warp represents one expert group - auto threadExpert = threadIdx.x; - bool expertSelected = threadExpert < numExperts; - if constexpr (UseGroups) { - threadExpert = warpIdx * numExpertsPerGroup + laneIdx; - expertSelected = laneIdx < numExpertsPerGroup; - } - - auto scoreIdx = int64_t{blockIdx.x} * int64_t{numExperts} + threadExpert; - auto biasVal = expertSelected ? static_cast(routingBias[threadExpert]) - : invalidScoreFloat; - topkValues += blockIdx.x * topk; - topkIndices += blockIdx.x * topk; - - // get our assigned thread score; each warp represents one expert group - float score = - expertSelected ? static_cast(scores[scoreIdx]) : invalidScoreFloat; - auto scoreSigmoid = apply_scoring(score); - // write the sigmoid score to shared for later use - if (expertSelected) { - smemScoreSigmoid[threadExpert] = scoreSigmoid; - } - - // get the score with bias - // note that with invalid values, because sigmoid is < 1 and bias is -1, - // we must get a negative value, which is smaller than any valid value - auto scoreBias = float{scoreSigmoid + float{biasVal}}; - - if (expertSelected) { - smemScoreBias[threadExpert] = scoreBias; - } - - // registers for top group score reduction - float topExpGroupScores[NumTopGroupScores]; - [[maybe_unused]] int32_t topExpGroupIdx[NumTopGroupScores]; - float topGroups[MaxNumTopGroups]; // bound of numGroup - int32_t topGroupIdx[MaxNumTopGroups]; - float expertScoreGroup[MaxNumTopGroups]; - int32_t expertIdxGroup[MaxNumTopGroups]; - float topScores[MaxNumTopExperts]; // bound of topk - int32_t topExperts[MaxNumTopExperts]; - - if constexpr (UseGroups) { - reduce_topk::reduceTopK(warp, topExpGroupScores, topExpGroupIdx, scoreBias, - threadExpert, - /* minValue */ invalidScoreFloat); - - // get the final group score and write it to shared - if (warp.thread_rank() == 0) { - auto groupScore = topExpGroupScores[0] + topExpGroupScores[1]; - smemGroupScores[warpIdx] = groupScore; - } - } - - // make group scores available to all warps - __syncthreads(); - - if constexpr (UseGroups) { - if (warpIdx == 0) { - // a single warp performs the selection of top groups, and goes on to - // select the final experts - float groupScore = - laneIdx < numGroup ? smemGroupScores[laneIdx] : invalidScoreFloat; - - reduce_topk::reduceTopK(warp, topGroups, topGroupIdx, groupScore, laneIdx, - /* minValue */ invalidScoreFloat); - // final expert selection: get relevant indexes and scores from shared -#pragma unroll - for (int ii = 0; ii < MaxNumTopGroups; ++ii) { // bound of numGroup - auto groupIdx = topGroupIdx[ii]; - expertIdxGroup[ii] = groupIdx * numExpertsPerGroup + laneIdx; - - expertScoreGroup[ii] = (ii < topkGroup) && expertSelected - ? smemScoreBias[expertIdxGroup[ii]] - : invalidScoreFloat; - } - - reduce_topk::reduceTopK(warp, topScores, topExperts, expertScoreGroup, - expertIdxGroup, /* minValue */ invalidScoreFloat, - topk); - } - } else if constexpr (MaxNumExperts > MaxNumExpertsUnit) { - // without groups, and the expert number is larger than MaxNumExpertsUnit, - // we need to use multiple warps to calculate the intermediate topk results - - int constexpr NumExpertWarps = (MaxNumExperts - 1) / MaxNumExpertsUnit + 1; - int constexpr NumInterTopK = NumExpertWarps * MaxNumTopExperts; - __shared__ float - __attribute((aligned(128))) smemInterTopScores[NumInterTopK]; - __shared__ int32_t - __attribute((aligned(128))) smemInterTopExperts[NumInterTopK]; - if (warpIdx < NumExpertWarps) { - int offset = warpIdx * WARP_SIZE * MaxNumTopGroups; -#pragma unroll - for (int ii = 0; ii < MaxNumTopGroups; ++ii) { - auto expertIdx = ii * WARP_SIZE + laneIdx; - expertIdxGroup[ii] = offset + expertIdx; - expertScoreGroup[ii] = offset + expertIdx < numExperts - ? smemScoreBias[offset + expertIdx] - : invalidScoreFloat; - } - reduce_topk::reduceTopK(warp, topScores, topExperts, expertScoreGroup, - expertIdxGroup, - /* minValue */ invalidScoreFloat, topk); - - if (laneIdx < topk) { - smemInterTopScores[warpIdx * MaxNumTopExperts + laneIdx] = - topScores[laneIdx]; - smemInterTopExperts[warpIdx * MaxNumTopExperts + laneIdx] = - topExperts[laneIdx]; - } else if (laneIdx >= topk && laneIdx < MaxNumTopExperts) { - smemInterTopScores[warpIdx * MaxNumTopExperts + laneIdx] = - invalidScoreFloat; - smemInterTopExperts[warpIdx * MaxNumTopExperts + laneIdx] = - MaxNumExperts - 1; - } - } - __syncthreads(); - if (warpIdx == 0) { - int constexpr NumInterTopKPerThread = (NumInterTopK - 1) / WARP_SIZE + 1; - float intermediateScore[NumInterTopKPerThread]; - int32_t intermediateExpert[NumInterTopKPerThread]; - for (int i = laneIdx; i < NumInterTopKPerThread * WARP_SIZE; - i += WARP_SIZE) { - int ii = i / WARP_SIZE; - if (i < NumInterTopK) { - intermediateScore[ii] = smemInterTopScores[i]; - intermediateExpert[ii] = smemInterTopExperts[i]; - } else { - intermediateScore[ii] = invalidScoreFloat; - intermediateExpert[ii] = MaxNumExperts - 1; - } - } - reduce_topk::reduceTopK(warp, topScores, topExperts, intermediateScore, - intermediateExpert, - /* minValue */ invalidScoreFloat, topk); - } - } else { - // without groups, and the expert number is smaller than MaxNumExpertsUnit - // each thread just takes `MaxNumTopGroups` experts - if (warpIdx == 0) { -#pragma unroll - for (int ii = 0; ii < MaxNumTopGroups; ++ii) { - auto expertIdx = ii * WARP_SIZE + laneIdx; - expertIdxGroup[ii] = expertIdx; - expertScoreGroup[ii] = expertIdx < numExperts ? smemScoreBias[expertIdx] - : invalidScoreFloat; - } - reduce_topk::reduceTopK(warp, topScores, topExperts, expertScoreGroup, - expertIdxGroup, - /* minValue */ invalidScoreFloat, topk); - } - } - - if (warpIdx == 0) { - // determine our lane's expert index and write to output - int32_t expertIdx = - laneIdx < topk ? topExperts[laneIdx] : MaxNumExperts - 1; - float scoreNorm = laneIdx < topk ? smemScoreSigmoid[expertIdx] : 0.F; - float finalScore = static_cast(scoreNorm * routedScalingFactor); - // norm the value - if (renormalize) { - auto redNorm = cg::reduce(warp, scoreNorm, cg::plus{}); - finalScore /= (redNorm + 1e-20); - } - // store the topk scores and experts to output - if (laneIdx < topk) { - topkValues[laneIdx] = finalScore; - topkIndices[laneIdx] = expertIdx; - } - } - -#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) - cudaTriggerProgrammaticLaunchCompletion(); -#endif -} - -template -void invokeNoAuxTc(T* scores, float* topk_values, IdxT* topk_indices, - BiasT const* bias, int64_t const num_tokens, - int64_t const num_experts, int64_t const n_group, - int64_t const topk_group, int64_t const topk, - bool const renormalize, double const routed_scaling_factor, - const bool enable_pdl = false, - cudaStream_t const stream = 0) { - cudaLaunchConfig_t config; - config.stream = stream; - cudaLaunchAttribute attrs[1]; - attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; - attrs[0].val.programmaticStreamSerializationAllowed = enable_pdl; - config.numAttrs = 1; - config.attrs = attrs; - - // Check if we can use the optimized - // grouped_topk_fused_small_expert_count_kernel - bool const is_single_group = - (n_group == 1) && (topk_group == 1) && - (num_experts <= MaxSupportedExpertCount) && - (topk <= DefaultMaxNumTopExperts || topk == MaxSupportedTopExperts); - - int64_t const experts_per_group = num_experts / n_group; - bool const is_multi_group = - (n_group > 1) && (num_experts <= NumDeepseekExperts) && - (experts_per_group <= WARP_SIZE) && - (experts_per_group * topk_group <= MaxNumExpertsUnit) && - (topk <= DefaultMaxNumTopExperts) && (topk_group <= MaxNumTopGroups); - - if (is_single_group || is_multi_group) { - auto* kernel_instance = - &grouped_topk_fused_small_expert_count_kernel; - int num_threads = NumDeepseekExperts; - if (is_single_group) { - // Special case for Nemotron, which selects top 22 from 512 experts, and 1 - // group only. - if (num_experts == NumNemotronExperts && n_group == 1 && - topk == MaxSupportedTopExperts) { - kernel_instance = &grouped_topk_fused_small_expert_count_kernel< - T, BiasT, IdxT, SF, NumNemotronExperts, false, - MaxSupportedTopExperts>; - num_threads = NumNemotronExperts; - } else if (num_experts > NumKimiK2Experts && - num_experts <= MaxSupportedExpertCount) { - kernel_instance = &grouped_topk_fused_small_expert_count_kernel< - T, BiasT, IdxT, SF, MaxSupportedExpertCount, false>; - num_threads = MaxSupportedExpertCount; - } else if (num_experts > MaxNumExpertsUnit && - num_experts <= NumKimiK2Experts) { - kernel_instance = &grouped_topk_fused_small_expert_count_kernel< - T, BiasT, IdxT, SF, NumKimiK2Experts, false>; - num_threads = NumKimiK2Experts; - } else { - kernel_instance = &grouped_topk_fused_small_expert_count_kernel< - T, BiasT, IdxT, SF, MaxNumExpertsUnit, false>; - num_threads = MaxNumExpertsUnit; - } - } - config.gridDim = num_tokens; - config.blockDim = num_threads; - config.dynamicSmemBytes = 0; - cudaLaunchKernelEx(&config, kernel_instance, scores, topk_values, - topk_indices, bias, num_tokens, n_group, topk_group, - topk, num_experts, num_experts / n_group, renormalize, - routed_scaling_factor); - } else { - auto* kernel_instance = &grouped_topk_fused_kernel; - // One block per token; one warp per group. - config.gridDim = static_cast(num_tokens); - config.blockDim = static_cast(n_group) * WARP_SIZE; - // Dynamic shared memory: WarpSelect staging + per-group topk buffers. - int32_t const num_warps = static_cast(n_group); - size_t const val_bytes = - static_cast(num_warps) * WARP_SIZE * sizeof(T); - size_t const val_bytes_aligned = - warp_topk::round_up_to_multiple_of<256>(val_bytes); - size_t const idx_bytes = - static_cast(num_warps) * WARP_SIZE * sizeof(int32_t); - size_t const internal_bytes = val_bytes_aligned + idx_bytes; - size_t const extra_bytes = 16 + static_cast(n_group) * sizeof(T); - config.dynamicSmemBytes = internal_bytes + extra_bytes; - cudaLaunchKernelEx(&config, kernel_instance, scores, topk_values, - topk_indices, bias, num_tokens, num_experts, n_group, - topk_group, topk, renormalize, routed_scaling_factor); - } -} - -#define INSTANTIATE_NOAUX_TC(T, BiasT, IdxT, SF) \ - template void invokeNoAuxTc( \ - T * scores, float* topk_values, IdxT* topk_indices, BiasT const* bias, \ - int64_t const num_tokens, int64_t const num_experts, \ - int64_t const n_group, int64_t const topk_group, int64_t const topk, \ - bool const renormalize, double const routed_scaling_factor, \ - const bool enable_pdl, cudaStream_t const stream); - -INSTANTIATE_NOAUX_TC(float, float, int32_t, SCORING_SIGMOID); -INSTANTIATE_NOAUX_TC(float, half, int32_t, SCORING_SIGMOID); -INSTANTIATE_NOAUX_TC(float, __nv_bfloat16, int32_t, SCORING_SIGMOID); -INSTANTIATE_NOAUX_TC(half, float, int32_t, SCORING_SIGMOID); -INSTANTIATE_NOAUX_TC(half, half, int32_t, SCORING_SIGMOID); -INSTANTIATE_NOAUX_TC(half, __nv_bfloat16, int32_t, SCORING_SIGMOID); -INSTANTIATE_NOAUX_TC(__nv_bfloat16, float, int32_t, SCORING_SIGMOID); -INSTANTIATE_NOAUX_TC(__nv_bfloat16, half, int32_t, SCORING_SIGMOID); -INSTANTIATE_NOAUX_TC(__nv_bfloat16, __nv_bfloat16, int32_t, SCORING_SIGMOID); -INSTANTIATE_NOAUX_TC(float, float, int32_t, SCORING_NONE); -INSTANTIATE_NOAUX_TC(float, half, int32_t, SCORING_NONE); -INSTANTIATE_NOAUX_TC(float, __nv_bfloat16, int32_t, SCORING_NONE); -INSTANTIATE_NOAUX_TC(half, float, int32_t, SCORING_NONE); -INSTANTIATE_NOAUX_TC(half, half, int32_t, SCORING_NONE); -INSTANTIATE_NOAUX_TC(half, __nv_bfloat16, int32_t, SCORING_NONE); -INSTANTIATE_NOAUX_TC(__nv_bfloat16, float, int32_t, SCORING_NONE); -INSTANTIATE_NOAUX_TC(__nv_bfloat16, half, int32_t, SCORING_NONE); -INSTANTIATE_NOAUX_TC(__nv_bfloat16, __nv_bfloat16, int32_t, SCORING_NONE); -} // end namespace moe -} // namespace aphrodite - -std::tuple grouped_topk( - torch::stable::Tensor const& scores, int64_t n_group, int64_t topk_group, - int64_t topk, bool renormalize, double routed_scaling_factor, - torch::stable::Tensor const& bias, int64_t scoring_func = 0) { - const auto data_type = scores.scalar_type(); - const auto bias_type = bias.scalar_type(); - STD_TORCH_CHECK(scores.dim() == 2, "scores must be a 2D Tensor"); - const int64_t num_tokens = scores.size(0); - const int64_t num_experts = scores.size(1); - STD_TORCH_CHECK(n_group > 0, "n_group must be positive"); - STD_TORCH_CHECK(topk > 0, "topk must be positive"); - STD_TORCH_CHECK(topk_group > 0, "topk_group must be positive"); - STD_TORCH_CHECK(topk_group <= n_group, "topk_group must be <= n_group"); - STD_TORCH_CHECK(num_experts % n_group == 0, - "num_experts should be divisible by n_group"); - STD_TORCH_CHECK(n_group <= 32, - "n_group should be smaller than or equal to 32 for now"); - STD_TORCH_CHECK(topk <= 32, - "topk should be smaller than or equal to 32 for now"); - STD_TORCH_CHECK(topk <= topk_group * (num_experts / n_group), - "topk must be <= topk_group * (num_experts / n_group)"); - STD_TORCH_CHECK( - scoring_func == aphrodite::moe::SCORING_NONE || - scoring_func == aphrodite::moe::SCORING_SIGMOID, - "scoring_func must be SCORING_NONE (0) or SCORING_SIGMOID (1)"); - - // Always output float32 for topk_values (eliminates Python-side conversion) - auto topk_values = torch::stable::new_empty( - scores, {num_tokens, topk}, torch::headeronly::ScalarType::Float); - auto topk_indices = torch::stable::new_empty( - scores, {num_tokens, topk}, torch::headeronly::ScalarType::Int); - - const bool pdl_flag = num_tokens <= aphrodite::moe::PDLEnableTokens; - - const torch::stable::accelerator::DeviceGuard device_guard( - scores.get_device_index()); - const cudaStream_t stream = - get_current_cuda_stream(scores.get_device_index()); - auto const sf = static_cast(scoring_func); - -#define LAUNCH_KERNEL_SF(T, BiasT, IdxT) \ - do { \ - switch (sf) { \ - case aphrodite::moe::SCORING_NONE: \ - aphrodite::moe::invokeNoAuxTc( \ - reinterpret_cast(scores.mutable_data_ptr()), \ - reinterpret_cast(topk_values.mutable_data_ptr()), \ - reinterpret_cast(topk_indices.mutable_data_ptr()), \ - reinterpret_cast(bias.data_ptr()), num_tokens, \ - num_experts, n_group, topk_group, topk, renormalize, \ - routed_scaling_factor, pdl_flag, stream); \ - break; \ - case aphrodite::moe::SCORING_SIGMOID: \ - aphrodite::moe::invokeNoAuxTc( \ - reinterpret_cast(scores.mutable_data_ptr()), \ - reinterpret_cast(topk_values.mutable_data_ptr()), \ - reinterpret_cast(topk_indices.mutable_data_ptr()), \ - reinterpret_cast(bias.data_ptr()), num_tokens, \ - num_experts, n_group, topk_group, topk, renormalize, \ - routed_scaling_factor, pdl_flag, stream); \ - break; \ - default: \ - STD_TORCH_CHECK(false, "Unsupported scoring_func"); \ - break; \ - } \ - } while (0) - -#define LAUNCH_KERNEL(T, IdxT) \ - do { \ - switch (bias_type) { \ - case torch::headeronly::ScalarType::Half: \ - LAUNCH_KERNEL_SF(T, half, IdxT); \ - break; \ - case torch::headeronly::ScalarType::Float: \ - LAUNCH_KERNEL_SF(T, float, IdxT); \ - break; \ - case torch::headeronly::ScalarType::BFloat16: \ - LAUNCH_KERNEL_SF(T, __nv_bfloat16, IdxT); \ - break; \ - default: \ - STD_TORCH_CHECK( \ - false, \ - "Invalid bias dtype, only supports float16, float32, and " \ - "bfloat16"); \ - break; \ - } \ - } while (0) - - switch (data_type) { - case torch::headeronly::ScalarType::Half: - // Handle Float16 - LAUNCH_KERNEL(half, int32_t); - break; - case torch::headeronly::ScalarType::Float: - // Handle Float32 - LAUNCH_KERNEL(float, int32_t); - break; - case torch::headeronly::ScalarType::BFloat16: - // Handle BFloat16 - LAUNCH_KERNEL(__nv_bfloat16, int32_t); - break; - default: - // Handle other data types - STD_TORCH_CHECK( - false, "Invalid dtype, only supports float16, float32, and bfloat16"); - break; - } -#undef LAUNCH_KERNEL -#undef LAUNCH_KERNEL_SF - return {topk_values, topk_indices}; -} diff --git a/csrc/libtorch_stable/moe/moeTopKFuncs.cuh.orig b/csrc/libtorch_stable/moe/moeTopKFuncs.cuh.orig deleted file mode 100644 index 81a216bd6f..0000000000 --- a/csrc/libtorch_stable/moe/moeTopKFuncs.cuh.orig +++ /dev/null @@ -1,257 +0,0 @@ -/* - * Adapted from - * https://github.com/NVIDIA/TensorRT-LLM/blob/v1.3.0rc2/cpp/tensorrt_llm/kernels/moeTopKFuncs.cuh - * Copyright (c) 2026, The vLLM team. - * SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION. All rights - * reserved. SPDX-License-Identifier: Apache-2.0 - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ -#pragma once - -#include -#include -#include - -namespace aphrodite { -namespace moe { -namespace reduce_topk { -namespace cg = cooperative_groups; -static constexpr int kWARP_SIZE = 32; - -template -struct TopKRedType { - using T = T_; - static_assert( - std::is_same_v || std::is_same_v || - std::is_same_v || std::is_same_v, - "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; - - static __host__ __device__ inline TypeCmp makeCmpVal(T val, int32_t idx = 0) { - auto valueBits = cub::Traits::TwiddleIn( - reinterpret_cast::UnsignedBits&>(val)); - TypeCmp compactTmp = valueBits; - compactTmp = (compactTmp << kMoveBits) | (0xFFFF & (kMaxIdx - idx)); - // Use 65535 minus idx to give higher priority to elements with smaller - // indices. - return compactTmp; - } - - static __host__ __device__ void unpack(T& value, int32_t& index, - TypeCmp cmp) { - // Since “65535-idx” is always smaller than 65536 and positive, we can - // directly use it as the lower 16 bits - index = kMaxIdx - static_cast((cmp & 0xFFFF)); - - auto compactTmp = cmp >> kMoveBits; - auto valueBits = cub::Traits::TwiddleOut( - reinterpret_cast::UnsignedBits&>(compactTmp)); - value = reinterpret_cast(valueBits); - } - - __host__ __device__ TopKRedType() = default; - - __host__ __device__ TopKRedType(T val, int32_t idx) - : compValIdx(makeCmpVal(val, idx)) {} - - __host__ __device__ operator TypeCmp() const noexcept { return compValIdx; } - - __device__ inline TypeCmp reduce( - cg::thread_block_tile const& warp) { - return cg::reduce(warp, compValIdx, cg::greater{}); - } -}; - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -template -struct TopKIdx { - // by default, empty -}; - -template -struct TopKIdx { - static constexpr int K = K_; - int32_t val[K]; -}; - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -#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 -struct Sort; - -template -struct Sort<1, RedType> { - static __device__ void run(RedType* topK) {} -}; - -template -struct Sort<2, RedType> { - static __device__ void run(RedType* topK) { TOPK_SWAP(0, 1); } -}; - -template -struct Sort<3, RedType> { - static __device__ void run(RedType* topK) { - TOPK_SWAP(0, 1); - TOPK_SWAP(1, 2); - TOPK_SWAP(0, 1); - } -}; - -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); - } -}; - -template -__forceinline__ __device__ void reduceTopK( - cg::thread_block_tile const& warp, Type (&out)[K], - int32_t (&outIdx)[K], Type value, int32_t idx, 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"); - using RedType = TopKRedType; - RedType topK{value, idx}; - 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 - 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) { - static_assert(K > 0, "Top K must have K > 0"); - 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"); - using RedType = TopKRedType; - RedType topK[N]; -#pragma unroll - for (int nn = 0; nn < N; ++nn) { - topK[nn] = RedType{value[nn], idx[nn]}; - } - - if constexpr (!IsSorted) { - Sort::run(topK); - } - typename RedType::TypeCmp packedMax{}; -#pragma unroll - for (int kk = 0; kk < actualK; ++kk) { - bool update = kk > 0 && packedMax == topK[0].compValIdx; -#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 -__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(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"); - using RedType = TopKRedType; - - 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; - - for (int ii = 0; ii < numResults; ++ii) { - topKBufferValue[ii] = minValue; - topKBufferIdx[ii] = ii * kWARP_SIZE - 1; - } - 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]; - } - } - - reduceTopKFunc(warp, out, outIdx, topKBufferValue, - topKBufferIdx, minValue, actualK); - } -}; - -#undef TOPK_SWAP - -} // namespace reduce_topk -} // namespace moe -} // namespace aphrodite diff --git a/csrc/libtorch_stable/ops.h.orig b/csrc/libtorch_stable/ops.h.orig deleted file mode 100644 index c2bc9c4c66..0000000000 --- a/csrc/libtorch_stable/ops.h.orig +++ /dev/null @@ -1,621 +0,0 @@ -#pragma once - -#include -#include -#include - -#include -#include -#include - -#include - -inline torch::stable::Tensor weak_ref_tensor(torch::stable::Tensor& tensor) { - // Ensure tensor is on CUDA - STD_TORCH_CHECK(tensor.device().is_cuda(), "Tensor must be on CUDA device"); - - // Get the raw data pointer - void* data_ptr = tensor.mutable_data_ptr(); - - /// Create a new tensor from the raw data pointer - return torch::stable::from_blob(data_ptr, tensor.sizes(), tensor.strides(), - tensor.device(), tensor.scalar_type()); -} - -void per_token_group_quant_fp8(const torch::stable::Tensor& input, - torch::stable::Tensor& output_q, - torch::stable::Tensor& output_s, - int64_t group_size, double eps, double fp8_min, - double fp8_max, bool scale_ue8m0, - bool dummy_is_scale_transposed, - bool dummy_is_tma_aligned); - -// Fused activation quantisation + DeepGEMM-compatible UE8M0-packed scales. -void per_token_group_quant_8bit_packed(const torch::stable::Tensor& input, - torch::stable::Tensor& output_q, - torch::stable::Tensor& output_s_packed, - int64_t group_size, double eps, - double min_8bit, double max_8bit); - -void per_token_group_quant_int8(const torch::stable::Tensor& input, - torch::stable::Tensor& output_q, - torch::stable::Tensor& output_s, - int64_t group_size, double eps, double int8_min, - double int8_max); - -torch::stable::Tensor permute_cols(torch::stable::Tensor const& A, - torch::stable::Tensor const& perm); - -std::tuple -dry_scan_penalties_cpu(const torch::stable::Tensor& token_history_ids, - const torch::stable::Tensor& token_history_lens, - const torch::stable::Tensor& dry_multiplier, - const torch::stable::Tensor& allowed_lengths, - const torch::stable::Tensor& sequence_breakers_ids, - const torch::stable::Tensor& ranges, - const torch::stable::Tensor& max_ngram, - const torch::stable::Tensor& max_occurrences, - const torch::stable::Tensor& early_exit_match_len, - int64_t vocab_size); - -#ifndef USE_ROCM -bool cutlass_scaled_mm_supports_fp8(int64_t cuda_device_capability); -bool cutlass_scaled_mm_supports_block_fp8(int64_t cuda_device_capability); -bool cutlass_group_gemm_supported(int64_t cuda_device_capability); - -void cutlass_scaled_mm(torch::stable::Tensor& out, - torch::stable::Tensor const& a, - torch::stable::Tensor const& b, - torch::stable::Tensor const& a_scales, - torch::stable::Tensor const& b_scales, - std::optional const& bias); - -void cutlass_moe_mm(torch::stable::Tensor& out_tensors, - torch::stable::Tensor const& a_tensors, - torch::stable::Tensor const& b_tensors, - torch::stable::Tensor const& a_scales, - torch::stable::Tensor const& b_scales, - torch::stable::Tensor const& expert_offsets, - torch::stable::Tensor const& problem_sizes, - torch::stable::Tensor const& a_strides, - torch::stable::Tensor const& b_strides, - torch::stable::Tensor const& c_strides, bool per_act_token, - bool per_out_ch); - -void cutlass_scaled_mm_azp(torch::stable::Tensor& out, - torch::stable::Tensor const& a, - torch::stable::Tensor const& b, - torch::stable::Tensor const& a_scales, - torch::stable::Tensor const& b_scales, - torch::stable::Tensor const& azp_adj, - std::optional const& azp, - std::optional const& bias); - -void get_cutlass_moe_mm_data( - const torch::stable::Tensor& topk_ids, - torch::stable::Tensor& expert_offsets, - torch::stable::Tensor& problem_sizes1, - torch::stable::Tensor& problem_sizes2, - torch::stable::Tensor& input_permutation, - torch::stable::Tensor& output_permutation, const int64_t num_experts, - const int64_t n, const int64_t k, - const std::optional& blockscale_offsets, - const bool is_gated); - -void get_cutlass_moe_mm_problem_sizes_from_expert_offsets( - const torch::stable::Tensor& expert_first_token_offset, - torch::stable::Tensor& problem_sizes1, - torch::stable::Tensor& problem_sizes2, const int64_t n, const int64_t k, - const bool swap_ab); - -void get_cutlass_batched_moe_mm_data( - torch::stable::Tensor& expert_offsets, - torch::stable::Tensor& problem_sizes1, - torch::stable::Tensor& problem_sizes2, - const torch::stable::Tensor& expert_num_tokens, - const int64_t num_local_experts, const int64_t padded_m, const int64_t n, - const int64_t k); - -// FP4/NVFP4 ops -bool cutlass_scaled_mm_supports_fp4(int64_t cuda_device_capability); - -void cutlass_scaled_fp4_mm(torch::stable::Tensor& D, - torch::stable::Tensor const& A, - torch::stable::Tensor const& B, - torch::stable::Tensor const& A_sf, - torch::stable::Tensor const& B_sf, - torch::stable::Tensor const& alpha); - -void cutlass_fp4_group_mm(torch::stable::Tensor& output, - const torch::stable::Tensor& a, - const torch::stable::Tensor& b, - const torch::stable::Tensor& a_blockscale, - const torch::stable::Tensor& b_blockscales, - const torch::stable::Tensor& alphas, - const torch::stable::Tensor& problem_sizes, - const torch::stable::Tensor& expert_offsets, - const torch::stable::Tensor& sf_offsets); - -std::tuple scaled_fp4_quant_func( - torch::stable::Tensor const& input, - torch::stable::Tensor const& input_scale, bool is_sf_swizzled_layout); - -void scaled_fp4_quant_out(torch::stable::Tensor const& input, - torch::stable::Tensor const& input_scale, - bool is_sf_swizzled_layout, - torch::stable::Tensor& output, - torch::stable::Tensor& output_scale); - -void scaled_fp4_experts_quant( - torch::stable::Tensor& output, torch::stable::Tensor& output_scale, - torch::stable::Tensor const& input, - torch::stable::Tensor const& input_global_scale, - torch::stable::Tensor const& input_offset_by_experts, - torch::stable::Tensor const& output_scale_offset_by_experts); - -void silu_and_mul_scaled_fp4_experts_quant( - torch::stable::Tensor& output, torch::stable::Tensor& output_scale, - torch::stable::Tensor const& input, - torch::stable::Tensor const& input_global_scale, - torch::stable::Tensor const& input_offset_by_experts, - torch::stable::Tensor const& output_scale_offset_by_experts); - -void silu_and_mul_nvfp4_quant(torch::stable::Tensor& out, - torch::stable::Tensor& output_block_scale, - torch::stable::Tensor& input, - torch::stable::Tensor& input_global_scale); - -void cutlass_mxfp4_group_mm(torch::stable::Tensor& output, - const torch::stable::Tensor& a, - const torch::stable::Tensor& b, - const torch::stable::Tensor& a_blockscale, - const torch::stable::Tensor& b_blockscales, - const torch::stable::Tensor& problem_sizes, - const torch::stable::Tensor& expert_offsets, - const torch::stable::Tensor& sf_offsets); - -// AWQ ops -torch::stable::Tensor awq_gemm(torch::stable::Tensor _in_feats, - torch::stable::Tensor _kernel, - torch::stable::Tensor _scaling_factors, - torch::stable::Tensor _zeros, - int64_t split_k_iters); - -torch::stable::Tensor awq_dequantize(torch::stable::Tensor _kernel, - torch::stable::Tensor _scaling_factors, - torch::stable::Tensor _zeros, - int64_t split_k_iters, int64_t thx, - int64_t thy); - -// DSV3 fused A GEMM: conditionally compiled so declaration and impl -// registration are in the source file (dsv3_fused_a_gemm.cu) - -// AllSpark ops: declarations are in the source files -// (allspark_repack.cu and allspark_qgemm_w8a16.cu) - -#endif - -// CPU tensor -> CUDA UVA view (shared CUDA/ROCm) -torch::stable::Tensor get_cuda_view_from_cpu_tensor( - torch::stable::Tensor& cpu_tensor); - -// Attention kernels (shared CUDA/ROCm) -void merge_attn_states( - torch::stable::Tensor& output, - std::optional output_lse, - const torch::stable::Tensor& prefix_output, - const torch::stable::Tensor& prefix_lse, - const torch::stable::Tensor& suffix_output, - const torch::stable::Tensor& suffix_lse, - const std::optional prefill_tokens_with_context, - const std::optional& output_scale = std::nullopt); - -torch::stable::Tensor hadacore_transform(torch::stable::Tensor& x, - bool inplace); - -// Layernorm kernels (shared CUDA/ROCm) -void rms_norm(torch::stable::Tensor& out, torch::stable::Tensor& input, - std::optional weight, double epsilon); - -void fused_add_rms_norm(torch::stable::Tensor& input, - torch::stable::Tensor& residual, - std::optional weight, - double epsilon); - -// Layernorm-quant kernels (shared CUDA/ROCm) -void rms_norm_static_fp8_quant(torch::stable::Tensor& out, - torch::stable::Tensor& input, - torch::stable::Tensor& weight, - torch::stable::Tensor& scale, double epsilon); - -void fused_add_rms_norm_static_fp8_quant(torch::stable::Tensor& out, - torch::stable::Tensor& input, - torch::stable::Tensor& residual, - torch::stable::Tensor& weight, - torch::stable::Tensor& scale, - double epsilon); - -// Fused layernorm + dynamic per-token quant kernels (shared CUDA/ROCm) -void rms_norm_dynamic_per_token_quant( - torch::stable::Tensor& out, torch::stable::Tensor const& input, - torch::stable::Tensor const& weight, torch::stable::Tensor& scales, - double const var_epsilon, std::optional scale_ub, - std::optional residual); - -void rms_norm_per_block_quant(torch::stable::Tensor& out, - torch::stable::Tensor const& input, - torch::stable::Tensor const& weight, - torch::stable::Tensor& scales, - double const var_epsilon, - std::optional scale_ub, - std::optional residual, - int64_t group_size, bool is_scale_transposed); - -void silu_and_mul_per_block_quant(torch::stable::Tensor& out, - torch::stable::Tensor const& input, - torch::stable::Tensor& scales, - int64_t group_size, - std::optional scale_ub, - bool is_scale_transposed); - -// Positional encoding kernels (shared CUDA/ROCm) -void rotary_embedding(torch::stable::Tensor& positions, - torch::stable::Tensor& query, - std::optional key, - int64_t head_size, torch::stable::Tensor& cos_sin_cache, - bool is_neox, int64_t rope_dim_offset, bool inverse); - -void fused_qk_norm_rope(torch::stable::Tensor& qkv, int64_t num_heads_q, - int64_t num_heads_k, int64_t num_heads_v, - int64_t head_dim, double eps, - torch::stable::Tensor& q_weight, - torch::stable::Tensor& k_weight, - torch::stable::Tensor& cos_sin_cache, bool is_neox, - torch::stable::Tensor& position_ids, - int64_t forced_token_heads_per_warp); - -torch::stable::Tensor fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert( - torch::stable::Tensor const& q_in, torch::stable::Tensor const& kv, - torch::stable::Tensor& k_cache, torch::stable::Tensor const& slot_mapping, - torch::stable::Tensor const& position_ids, - torch::stable::Tensor const& cos_sin_cache, int64_t q_head_padded, - double eps, int64_t cache_block_size); - -void fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_bf16_insert( - torch::stable::Tensor& q, torch::stable::Tensor const& kv, - torch::stable::Tensor& k_cache, torch::stable::Tensor const& slot_mapping, - torch::stable::Tensor const& position_ids, - torch::stable::Tensor const& cos_sin_cache, double eps, - int64_t cache_block_size); - -void fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_fp8_insert( - torch::stable::Tensor const& q, torch::stable::Tensor const& kv, - torch::stable::Tensor& q_fp8, torch::stable::Tensor& k_cache, - torch::stable::Tensor const& slot_mapping, - torch::stable::Tensor const& position_ids, - torch::stable::Tensor const& cos_sin_cache, - torch::stable::Tensor const& fp8_scale, - torch::stable::Tensor const& q_fp8_scale_inv, double eps, - int64_t cache_block_size); - -#ifndef USE_ROCM -std::tuple -minimax_allreduce_rms_qk(torch::stable::Tensor qkv, - torch::stable::Tensor const& norm_weight_q, - torch::stable::Tensor const& norm_weight_k, - torch::stable::Tensor workspace, int64_t const q_size, - int64_t const kv_size, int64_t const rank, - int64_t const nranks, double const eps); -#endif - -// Horizontally-fused MiniMax-M3 QK-norm + partial NeoX RoPE (+ optional KV / -// index-cache insert). Dense layer: norm+RoPE only; sparse layer: also packs -// the index branch and scatters k/v/index_k into their paged caches. -void fused_minimax_m3_qknorm_rope_kv_insert( - torch::stable::Tensor& qkv, torch::stable::Tensor const& q_norm_weight, - torch::stable::Tensor const& k_norm_weight, - torch::stable::Tensor const& cos_sin_cache, - torch::stable::Tensor const& positions, int64_t num_heads, - int64_t num_kv_heads, int64_t rotary_dim, double eps, - std::optional index_q_norm_weight, - std::optional index_k_norm_weight, - int64_t num_index_heads, std::optional slot_mapping, - std::optional index_slot_mapping, - std::optional kv_cache, - std::optional index_cache, int64_t block_size, - std::optional q_out, - std::optional index_q_out, - const std::string& kv_cache_dtype, bool skip_index_branch); - -#ifdef APHRODITE_ENABLE_KIMI_K3_ATTN_RES -void kimi_k3_attn_res(torch::stable::Tensor& prefix, - torch::stable::Tensor const& delta, - torch::stable::Tensor const& blocks, - torch::stable::Tensor const& norm_weight, - torch::stable::Tensor const& qk_weight, - torch::stable::Tensor const& output_norm_weight, - torch::stable::Tensor& output, int64_t num_blocks, - double eps, double output_norm_eps); -#endif - -// Sampler kernels (shared CUDA/ROCm) -void apply_repetition_penalties_( - torch::stable::Tensor& logits, const torch::stable::Tensor& prompt_mask, - const torch::stable::Tensor& output_mask, - const torch::stable::Tensor& repetition_penalties); - -void top_k_per_row_prefill(const torch::stable::Tensor& logits, - const torch::stable::Tensor& rowStarts, - const torch::stable::Tensor& rowEnds, - torch::stable::Tensor& indices, int64_t numRows, - int64_t stride0, int64_t stride1, int64_t topK); - -void top_k_per_row_decode(const torch::stable::Tensor& logits, int64_t next_n, - const torch::stable::Tensor& seqLens, - torch::stable::Tensor& indices, int64_t numRows, - int64_t stride0, int64_t stride1, int64_t topK); - -void persistent_topk(const torch::stable::Tensor& logits, - const torch::stable::Tensor& lengths, - torch::stable::Tensor& output, - torch::stable::Tensor& workspace, int64_t k, - int64_t max_seq_len); - -#ifdef APHRODITE_ENABLE_COOPERATIVE_TOPK -void cooperative_topk(const torch::stable::Tensor& logits, - const torch::stable::Tensor& lengths, - torch::stable::Tensor& output, - torch::stable::Tensor& workspace, int64_t k, - int64_t max_seq_len); -#endif - -#ifdef APHRODITE_ENABLE_SM89_DSA -// sm89 DeepSeek sparse attention kernels; impls are registered in -// csrc/libtorch_stable/attention/sm89_dsa/. -void sm89_fp8_paged_mqa_logits(const torch::stable::Tensor& q, - const torch::stable::Tensor& pool, - const torch::stable::Tensor& weights, - const torch::stable::Tensor& seq_lens, - const torch::stable::Tensor& block_table, - const torch::stable::Tensor& sched, - torch::stable::Tensor& logits, - bool clean_logits); - -void sm89_paged_mqa_logits_metadata(const torch::stable::Tensor& seq_lens, - torch::stable::Tensor& sched, - int64_t next_n); - -void sm89_fp8_mqa_logits(const torch::stable::Tensor& q, - const torch::stable::Tensor& kv, - const torch::stable::Tensor& kv_scales, - const torch::stable::Tensor& weights, - const torch::stable::Tensor& cu_seqlen_ks, - const torch::stable::Tensor& cu_seqlen_ke, - torch::stable::Tensor& logits); - -void sm89_sparse_mla_fwd(const torch::stable::Tensor& q, - const torch::stable::Tensor& pool, - const torch::stable::Tensor& indices, - torch::stable::Tensor& out, torch::stable::Tensor& lse, - double sm_scale, - std::optional topk_lens); -#endif - -void selective_scan_fwd( - const torch::stable::Tensor& u, const torch::stable::Tensor& delta, - const torch::stable::Tensor& A, const torch::stable::Tensor& B, - const torch::stable::Tensor& C, - const std::optional& D_, - const std::optional& z_, - const std::optional& delta_bias_, - bool delta_softplus, - const std::optional& query_start_loc, - const std::optional& cache_indices, - const std::optional& has_initial_state, - const torch::stable::Tensor& ssm_states, int64_t null_block_id, - int64_t block_size, - const std::optional& block_idx_first_scheduled_token, - const std::optional& block_idx_last_scheduled_token, - const std::optional& initial_state_idx, - const std::optional& cu_chunk_seqlen, - const std::optional& last_chunk_indices); - -using fptr_t = int64_t; -fptr_t init_custom_ar(const std::vector& fake_ipc_ptrs, - torch::stable::Tensor& rank_data, int64_t rank, - bool fully_connected); -void all_reduce(fptr_t _fa, torch::stable::Tensor& inp, - torch::stable::Tensor& out, fptr_t reg_buffer, - int64_t reg_buffer_sz_bytes); -void dispose(fptr_t _fa); -int64_t meta_size(); -void register_buffer(fptr_t _fa, const std::vector& fake_ipc_ptrs); -std::tuple, std::vector> -get_graph_buffer_ipc_meta(fptr_t _fa); -void register_graph_buffers(fptr_t _fa, - const std::vector>& handles, - const std::vector>& offsets); -std::tuple allocate_shared_buffer_and_handle( - int64_t size); -int64_t open_mem_handle(torch::stable::Tensor& mem_handle); -void free_shared_buffer(int64_t buffer); - -// Activation kernels (shared CUDA/ROCm) -void silu_and_mul(torch::stable::Tensor& out, torch::stable::Tensor& input); -void silu_and_mul_clamp(torch::stable::Tensor& out, - torch::stable::Tensor& input, double limit, - double alpha = 1.0, double beta = 0.0); - -void silu_and_mul_quant(torch::stable::Tensor& out, - torch::stable::Tensor& input, - torch::stable::Tensor& scale); - -void persistent_masked_m_silu_mul_quant( - const torch::stable::Tensor& input, // (E, T, 2*H) - const torch::stable::Tensor& tokens_per_expert, // (E) - torch::stable::Tensor& y_q, // (E, T, H) [OUT] - torch::stable::Tensor& y_s, // (E, T, H//group_size) [OUT] - bool use_ue8m0); - -void mul_and_silu(torch::stable::Tensor& out, torch::stable::Tensor& input); -void gelu_and_mul(torch::stable::Tensor& out, torch::stable::Tensor& input); -void gelu_tanh_and_mul(torch::stable::Tensor& out, - torch::stable::Tensor& input); -void fatrelu_and_mul(torch::stable::Tensor& out, torch::stable::Tensor& input, - double threshold); -void swigluoai_and_mul(torch::stable::Tensor& out, torch::stable::Tensor& input, - double alpha = 1.702, double limit = 7.0); -void gelu_new(torch::stable::Tensor& out, torch::stable::Tensor& input); -void gelu_fast(torch::stable::Tensor& out, torch::stable::Tensor& input); -void gelu_quick(torch::stable::Tensor& out, torch::stable::Tensor& input); -void relu_squared(torch::stable::Tensor& out, torch::stable::Tensor& input); - -// INT8 quantization kernels (shared CUDA/ROCm) -void static_scaled_int8_quant(torch::stable::Tensor& out, - torch::stable::Tensor const& input, - torch::stable::Tensor const& scale, - std::optional const& azp); - -void dynamic_scaled_int8_quant(torch::stable::Tensor& out, - torch::stable::Tensor const& input, - torch::stable::Tensor& scales, - std::optional const& azp); - -// FP8 quantization kernels (shared CUDA/ROCm) -void static_scaled_fp8_quant( - torch::stable::Tensor& out, torch::stable::Tensor const& input, - torch::stable::Tensor const& scale, - std::optional group_shape = - std::nullopt); - -void dynamic_scaled_fp8_quant(torch::stable::Tensor& out, - torch::stable::Tensor const& input, - torch::stable::Tensor& scale); - -void dynamic_per_token_scaled_fp8_quant( - torch::stable::Tensor& out, torch::stable::Tensor const& input, - torch::stable::Tensor& scale, - std::optional const& scale_ub); - -// GPTQ kernels (shared CUDA/ROCm) -torch::stable::Tensor gptq_gemm(torch::stable::Tensor a, - torch::stable::Tensor b_q_weight, - torch::stable::Tensor b_gptq_qzeros, - torch::stable::Tensor b_gptq_scales, - torch::stable::Tensor b_g_idx, bool use_exllama, - bool use_v2_format, int64_t bit); - -void gptq_shuffle(torch::stable::Tensor q_weight, torch::stable::Tensor q_perm, - int64_t bit); - -// Cache ops (shared CUDA/ROCm) -void swap_blocks(torch::stable::Tensor& src, torch::stable::Tensor& dst, - int64_t block_size_in_bytes, - const torch::stable::Tensor& block_mapping); - -// Batch swap: submit all block copies in a single driver call. -void swap_blocks_batch(const torch::stable::Tensor& src_ptrs, - const torch::stable::Tensor& dst_ptrs, - const torch::stable::Tensor& sizes, - bool is_src_access_order_any); - -void reshape_and_cache(torch::stable::Tensor& key, torch::stable::Tensor& value, - torch::stable::Tensor& key_cache, - torch::stable::Tensor& value_cache, - torch::stable::Tensor& slot_mapping, - const std::string& kv_cache_dtype, - torch::stable::Tensor& k_scale, - torch::stable::Tensor& v_scale); - -void reshape_and_cache_flash( - torch::stable::Tensor& key, torch::stable::Tensor& value, - torch::stable::Tensor& key_cache, torch::stable::Tensor& value_cache, - torch::stable::Tensor& slot_mapping, const std::string& kv_cache_dtype, - torch::stable::Tensor& k_scale, torch::stable::Tensor& v_scale); - -void concat_and_cache_mla(torch::stable::Tensor& kv_c, - torch::stable::Tensor& k_pe, - torch::stable::Tensor& kv_cache, - torch::stable::Tensor& slot_mapping, - const std::string& kv_cache_dtype, - torch::stable::Tensor& scale); - -// NOTE: k_pe and kv_c order is flipped compared to concat_and_cache_mla -void concat_and_cache_mla_rope_fused( - torch::stable::Tensor& positions, torch::stable::Tensor& q_pe, - torch::stable::Tensor& k_pe, torch::stable::Tensor& kv_c, - torch::stable::Tensor& rope_cos_sin_cache, bool rope_is_neox, - torch::stable::Tensor& slot_mapping, torch::stable::Tensor& kv_cache, - const std::string& kv_cache_dtype, - torch::stable::Tensor& kv_cache_quant_scale); - -// Just for unittest -void convert_fp8(torch::stable::Tensor& dst_cache, - torch::stable::Tensor& src_cache, const double scale, - const std::string& kv_cache_dtype); - -void gather_and_maybe_dequant_cache( - torch::stable::Tensor const& src_cache, // [NUM_BLOCKS, BLOCK_SIZE, - // ENTRIES...] - torch::stable::Tensor const& dst, // [TOT_TOKENS, ENTRIES...] - torch::stable::Tensor const& block_table, // [BATCH, BLOCK_INDICES] - torch::stable::Tensor const& cu_seq_lens, // [BATCH+1] - torch::stable::Tensor const& token_to_seq, // [MAX_TOKEN_ACROSS_CHUNKS] - int64_t num_tokens, const std::string& kv_cache_dtype, - torch::stable::Tensor const& scale, - std::optional seq_starts = std::nullopt); - -// TODO(hc): cp_gather_cache need support scaled kvcahe in the future. -void cp_gather_cache( - torch::stable::Tensor const& src_cache, // [NUM_BLOCKS, BLOCK_SIZE, - // ENTRIES...] - torch::stable::Tensor const& dst, // [TOT_TOKENS, ENTRIES...] - torch::stable::Tensor const& block_table, // [BATCH, BLOCK_INDICES] - torch::stable::Tensor const& cu_seq_lens, // [BATCH+1] - int64_t batch_size, - std::optional seq_starts = std::nullopt); - -// Gather and upconvert FP8 KV cache to BF16 workspace -void cp_gather_and_upconvert_fp8_kv_cache( - torch::stable::Tensor const& src_cache, // [NUM_BLOCKS, BLOCK_SIZE, - // 656] - torch::stable::Tensor const& dst, // [TOT_TOKENS, 576] - torch::stable::Tensor const& block_table, // [BATCH, BLOCK_INDICES] - torch::stable::Tensor const& workspace_starts, // [BATCH] - int64_t batch_size, - std::optional seq_starts = std::nullopt); - -// Indexer K quantization and cache function -void indexer_k_quant_and_cache( - torch::stable::Tensor& k, // [num_tokens, head_dim] - torch::stable::Tensor& kv_cache, // [num_blocks, block_size, - // cache_stride] - torch::stable::Tensor& slot_mapping, // [num_tokens] - int64_t quant_block_size, // quantization block size - const std::string& scale_fmt); - -// Concatenate query nope and rope for MLA/DSA attention -void concat_mla_q( - torch::stable::Tensor& ql_nope, // [num_tokens, num_heads, nope_dim] - torch::stable::Tensor& q_pe, // [num_tokens, num_heads, rope_dim] - torch::stable::Tensor& q_out); // [num_tokens, num_heads, nope_dim + - // rope_dim] - -// Extract function to gather quantized K cache -void cp_gather_indexer_k_quant_cache( - const torch::stable::Tensor& kv_cache, // [num_blocks, block_size, - // cache_stride] - torch::stable::Tensor& dst_k, // [num_tokens, head_dim] - torch::stable::Tensor& dst_scale, // [num_tokens, head_dim / - // quant_block_size * 4] - const torch::stable::Tensor& block_table, // [batch_size, num_blocks] - const torch::stable::Tensor& cu_seq_lens); // [batch_size + 1] - -// LongCat n-gram embedding index kernel (see ngram_embedding_kernels.cu). -void ngram_compute_n_gram_ids( - int64_t ne_n, int64_t ne_k, torch::stable::Tensor& ne_weights, - torch::stable::Tensor& ne_mods, - torch::stable::Tensor& exclusive_ne_embedder_size_sums, - torch::stable::Tensor& exclusive_req_len_sums, - torch::stable::Tensor& ne_token_table, torch::stable::Tensor& row_indices, - torch::stable::Tensor& column_starts, torch::stable::Tensor& n_gram_ids); diff --git a/csrc/libtorch_stable/torch_bindings.cpp.orig b/csrc/libtorch_stable/torch_bindings.cpp.orig deleted file mode 100644 index 0a5ef9c11d..0000000000 --- a/csrc/libtorch_stable/torch_bindings.cpp.orig +++ /dev/null @@ -1,1014 +0,0 @@ -#include "ops.h" -#include "cuda_utils.h" -#include "core/registration.h" - -#include - -// Register ops with STABLE_TORCH_LIBRARY for libtorch stable ABI compatibility. -// Note: We register under namespace "_C" so ops are accessible as -// torch.ops._C. for compatibility with existing code. -STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) { - // Compute per-token-group FP8 quantized tensor and scaling factor. - // The dummy arguments are here so we can correctly fuse with RMSNorm. - ops.def( - "per_token_group_fp8_quant(Tensor input, Tensor! output_q, Tensor! " - "output_s, " - "int group_size, float eps, float fp8_min, float fp8_max, bool " - "scale_ue8m0, bool dummy_is_scale_transposed, bool dummy_is_tma_aligned " - ") -> ()"); - // Compute per-token-group 8-bit quantized tensor and UE8M0-packed, - // TMA-aligned scales for DeepGEMM. - ops.def( - "per_token_group_fp8_quant_packed(Tensor input, Tensor! output_q, " - "Tensor! output_s_packed, int group_size, float eps, float fp8_min, " - "float fp8_max) -> ()"); - // Compute per-token-group INT8 quantized tensor and scaling factor. - ops.def( - "per_token_group_quant_int8(Tensor input, Tensor! output_q, Tensor! " - "output_s, int group_size, float eps, float int8_min, float int8_max) -> " - "()"); - ops.def("permute_cols(Tensor A, Tensor perm) -> Tensor"); - ops.def( - "dry_scan_penalties(" - " Tensor token_history_ids," - " Tensor token_history_lens," - " Tensor dry_multiplier," - " Tensor allowed_lengths," - " Tensor sequence_breakers_ids," - " Tensor ranges," - " Tensor max_ngram," - " Tensor max_occurrences," - " Tensor early_exit_match_len," - " int vocab_size) -> (Tensor, Tensor, Tensor)"); - - ops.def("get_cuda_view_from_cpu_tensor(Tensor cpu_tensor) -> Tensor"); - -#ifndef USE_ROCM - - // Note about marlin kernel 'workspace' arguments: - // Technically these should be mutable since they are modified by the kernel. - // But since they are set back to zero once the kernel is finished we can - // hand wave and say that they have no net effect. - // - // The reason to mark 'workspace' as immutable is so that they don't interfere - // with using ScalarType arguments in the ops. If they are marked as mutable, - // pytorch throws an assert in - // 'torch._higher_order_ops._register_effectful_op' that prevents these - // kernels from being torch.compile'd. - // See the following document for more info on custom types and ops that use - // custom types: - // https://docs.google.com/document/d/18fBMPuOJ0fY5ZQ6YyrHUppw9FA332CpNtgB6SOIgyuA - - // Machete (Dense) Optimized Mixed Precision GEMM for Hopper. - ops.def( - "machete_supported_schedules(" - " ScalarType a_type," - " int b_type," - " ScalarType? maybe_group_scales_type," - " ScalarType? maybe_group_zeros_type," - " ScalarType? maybe_channel_scales_type," - " ScalarType? maybe_token_scales_type," - " ScalarType? maybe_out_type" - ") -> str[]"); - ops.def( - "machete_mm(" - " Tensor A," - " Tensor B," - " int b_type," - " ScalarType? out_type," - " Tensor? group_scales," - " Tensor? group_zeros," - " int? group_size," - " Tensor? channel_scales," - " Tensor? token_scales," - " str? schedule" - ") -> Tensor"); - ops.def( - "machete_prepack_B(" - " Tensor B," - " ScalarType a_type," - " int b_type," - " ScalarType? group_scales_type" - ") -> Tensor"); - // conditionally compiled so impl registration is in source file - - // Marlin GEMM - ops.def( - "marlin_gemm(Tensor a, Tensor? c_or_none, Tensor b_q_weight, " - "Tensor? b_bias_or_none,Tensor b_scales, " - "Tensor? a_scales, Tensor? global_scale, Tensor? b_zeros_or_none, " - "Tensor? " - "g_idx_or_none, Tensor? perm_or_none, Tensor workspace, int b_type_id, " - "SymInt size_m, SymInt size_n, SymInt size_k, bool is_k_full, " - "bool use_atomic_add, bool use_fp32_reduce, bool is_zp_float) -> Tensor"); - // conditionally compiled so impl registrations are in source file - - // gptq_marlin repack from GPTQ. - ops.def( - "gptq_marlin_repack(Tensor b_q_weight, Tensor perm, " - "SymInt size_k, SymInt size_n, int num_bits, bool is_a_8bit) -> Tensor"); - // conditionally compiled so impl registrations are in source file - - // awq_marlin repack from AWQ. - ops.def( - "awq_marlin_repack(Tensor b_q_weight, SymInt size_k, " - "SymInt size_n, int num_bits, bool is_a_8bit) -> Tensor"); - // conditionally compiled so impl registrations are in source file - - // preprocess W-int4A-fp8 weight for marlin kernel - ops.def( - "marlin_int4_fp8_preprocess(Tensor qweight, " - "Tensor? qzeros_or_none, bool inplace) -> Tensor"); - // conditionally compiled so impl registrations are in source file - - // swordfish (sm100/sm110 w4a16): pack a GPTQ int4 weight into the - // Swordfish ABI v1 block-linear layout. - ops.def( - "swordfish_prepack_B(Tensor b_q_weight, Tensor? perm, SymInt size_k, " - "SymInt size_n, int num_bits) -> Tensor"); - // conditionally compiled so impl registrations are in source file - - // swordfish w4a16 decode GEMM over the ABI v1 packed weight. group_zps - // holds prescaled (8 - zp) * scale per group when present (AWQ/HQQ). - ops.def( - "swordfish_mm(Tensor a, Tensor b_packed, Tensor group_scales, " - "Tensor? group_zps, Tensor? perm, int num_bits, int group_size, " - "SymInt size_k, SymInt size_n) -> Tensor"); - // conditionally compiled so impl registrations are in source file - - // swordfish dequant to dense fp16/bf16, optionally out-major transposed - // and expert-stacked, for the dense tiers. - ops.def( - "swordfish_dequant_dense(Tensor b_packed, Tensor group_scales, " - "Tensor? group_zps, int num_bits, int group_size, SymInt size_k, " - "SymInt size_n, bool transpose) -> Tensor"); - // conditionally compiled so impl registrations are in source file - - // swordfish fused-MoE decode GEMM: one persistent Stream-K launch over - // per-expert ABI v1 weights, token-sorted by moe_align_block_size(16). - ops.def( - "swordfish_moe_mm(Tensor a, Tensor b_packed, Tensor group_scales, " - "Tensor sorted_token_ids, Tensor expert_ids, " - "Tensor num_tokens_post_padded, Tensor? topk_weights, " - "int moe_block_size, int top_k, bool mul_topk_weights, int num_bits, " - "int group_size, SymInt size_k, SymInt size_n) -> Tensor"); - // conditionally compiled so impl registrations are in source file - - // swordfish w4a16 prefill GEMM (sm100 tcgen05 mixed-input mainloop fork) - // over the same ABI v1 packed weight. - ops.def( - "swordfish_prefill_mm(Tensor a, Tensor b_packed, Tensor group_scales, " - "Tensor? group_zps, int num_bits, int group_size, SymInt size_k, " - "SymInt size_n) -> Tensor"); - // conditionally compiled so impl registrations are in source file -#endif - -#ifndef USE_ROCM - // CUTLASS w8a8 GEMM, supporting symmetric per-tensor or per-row/column - // quantization, as well as bias - ops.def( - "cutlass_scaled_mm(Tensor! out, Tensor a," - " Tensor b, Tensor a_scales," - " Tensor b_scales, Tensor? bias) -> ()"); - - // CUTLASS w8a8 GEMM, supporting asymmetric per-tensor or per-row/column - // quantization. - ops.def( - "cutlass_scaled_mm_azp(Tensor! out, Tensor a," - " Tensor b, Tensor a_scales," - " Tensor b_scales, Tensor azp_adj," - " Tensor? azp, Tensor? bias) -> ()"); - - // Check if cutlass scaled_mm is supported for CUDA devices of the given - // capability - ops.def("cutlass_scaled_mm_supports_fp8(int cuda_device_capability) -> bool"); - - // Check if cutlass grouped gemm is supported for CUDA devices of the given - // capability - ops.def("cutlass_group_gemm_supported(int cuda_device_capability) -> bool"); - - // CUTLASS w8a8 grouped GEMM - ops.def( - "cutlass_moe_mm(Tensor! out_tensors, Tensor a_tensors, Tensor b_tensors, " - " Tensor a_scales, Tensor b_scales, Tensor expert_offsets, " - " Tensor problem_sizes, Tensor a_strides, " - " Tensor b_strides, Tensor c_strides, bool per_act_token, " - " bool per_out_ch) -> ()"); - - // A function that computes data required to run fused MoE with w8a8 grouped - // GEMM. It takes topk_ids as an input, and computes expert_offsets - // (token start indices of each expert). In addition to this, it computes - // problem sizes for each expert's multiplication used by the two mms called - // from fused MoE operation, and arrays with permutations required to shuffle - // and de-shuffle the input/output of the fused operation. - ops.def( - "get_cutlass_moe_mm_data(Tensor topk_ids, Tensor! expert_offsets, " - " Tensor! problem_sizes1, Tensor! problem_sizes2, " - " Tensor! input_permutation, " - " Tensor! output_permutation, int num_experts, " - " int n, int k, Tensor? blockscale_offsets, " - " bool is_gated) -> ()"); - - // compute per-expert problem sizes from expert_first_token_offset - // produced by vLLM's moe_permute kernel - ops.def( - "get_cutlass_moe_mm_problem_sizes_from_expert_offsets(" - " Tensor expert_first_token_offset, " - " Tensor! problem_sizes1, " - " Tensor! problem_sizes2, " - " int n, int k, bool swap_ab) -> ()"); - - // A function that computes data required to run fused MoE with w8a8 grouped - // GEMM in batched expert format. It takes expert_num_tokens - // as an input, and computes expert_offsets (token start indices of each - // expert). In addition to this, it computes problem sizes for each expert's - // multiplication used by the two mms called from fused MoE operation. - ops.def( - "get_cutlass_batched_moe_mm_data(Tensor! expert_offsets, " - " Tensor! problem_sizes1, " - " Tensor! problem_sizes2, " - " Tensor expert_num_tokens, " - " int num_local_experts, int padded_m, " - " int n, int k) -> ()"); - - // Check if cutlass scaled_mm supports block quantization (used by DeepSeekV3) - ops.def( - "cutlass_scaled_mm_supports_block_fp8(int cuda_device_capability) -> " - "bool"); - - // CUTLASS nvfp4 block scaled GEMM - ops.def( - "cutlass_scaled_fp4_mm(Tensor! out, Tensor a, Tensor b," - " Tensor block_scale_a, Tensor block_scale_b," - " Tensor alpha) -> ()"); - - // cutlass nvfp4 block scaled group GEMM - ops.def( - "cutlass_fp4_group_mm(Tensor! out, Tensor a, Tensor b," - " Tensor a_blockscale, Tensor b_blockscales, Tensor alphas," - " Tensor problem_sizes, Tensor expert_offsets, Tensor sf_offsets) -> ()"); - - // cutlass mxfp4 block scaled group GEMM (MXFP4 x MXFP4 MoE) - ops.def( - "cutlass_mxfp4_group_mm(Tensor! out, Tensor a, Tensor b," - " Tensor a_blockscale, Tensor b_blockscales," - " Tensor problem_sizes, Tensor expert_offsets, Tensor sf_offsets) -> ()"); - - // Compute NVFP4 block quantized tensor. - ops.def( - "scaled_fp4_quant(Tensor input," - " Tensor input_scale, bool " - "is_sf_swizzled_layout) -> (Tensor, Tensor)"); - - // Out variant - // TODO: Add out_variant tag once PyTorch supports it (added in 2.11) - // This registration is now migrated to stable ABI - // at::Tag::out_variant is not available in the stable ABI (enum_tag.h is not - // yet in torch/headeronly), the tag should be applied from Python - // via torch.library.Library.define(..., tags=(torch.Tag.out_variant,)) - // with the .impl remaining in C++. - // See pytorch/pytorch#176117. - ops.def( - "scaled_fp4_quant.out(Tensor input," - " Tensor input_scale, bool " - "is_sf_swizzled_layout, *, Tensor(a!) output, Tensor(b!) output_scale) " - "-> ()"); - - // Compute NVFP4 experts quantization. - ops.def( - "scaled_fp4_experts_quant(Tensor! output, Tensor! output_scale," - "Tensor input, Tensor input_global_scale, Tensor input_offset_by_experts," - "Tensor output_scale_offset_by_experts) -> ()"); - - // Fused SiLU+Mul+NVFP4 experts quantization. - ops.def( - "silu_and_mul_scaled_fp4_experts_quant(Tensor! output, Tensor! " - "output_scale," - "Tensor input, Tensor input_global_scale, Tensor input_offset_by_experts," - "Tensor output_scale_offset_by_experts) -> ()"); - - // Compute MXFP4 experts quantization (32-element blocks, E8M0 SFs). - ops.def( - "mxfp4_experts_quant(Tensor! output, Tensor! output_scale," - "Tensor input, Tensor input_offset_by_experts," - "Tensor output_scale_offset_by_experts, int n_experts) -> ()"); - - // Fused SiLU+Mul+MXFP4 experts quantization. - ops.def( - "silu_and_mul_mxfp4_experts_quant(Tensor! output, Tensor! " - "output_scale," - "Tensor input, Tensor input_offset_by_experts," - "Tensor output_scale_offset_by_experts, int n_experts) -> ()"); - - // Fused SiLU+Mul+NVFP4 quantization. - ops.def( - "silu_and_mul_nvfp4_quant(Tensor! result, Tensor! result_block_scale, " - "Tensor input, Tensor input_global_scale) -> ()"); - - // Check if cutlass_scaled_mm_fp4 is supported for CUDA devices - // of the given capability - ops.def("cutlass_scaled_mm_supports_fp4(int cuda_device_capability) -> bool"); - - // CUTLASS w4a8 GEMM - ops.def( - "cutlass_w4a8_mm(" - " Tensor A," - " Tensor B," - " Tensor group_scales," - " int group_size," - " Tensor channel_scales," - " Tensor token_scales," - " ScalarType? out_type," - " str? maybe_schedule" - ") -> Tensor"); - - // pack scales - ops.def("cutlass_pack_scale_fp8(Tensor scales) -> Tensor"); - - // encode and reorder weight matrix - ops.def("cutlass_encode_and_reorder_int4b(Tensor B) -> Tensor"); - - // CUTLASS w4a8 grouped GEMM - ops.def( - "cutlass_w4a8_moe_mm(" - " Tensor! out_tensors," - " Tensor a_tensors," - " Tensor b_tensors," - " Tensor a_scales," - " Tensor b_scales," - " Tensor b_group_scales," - " int b_group_size," - " Tensor expert_offsets," - " Tensor problem_sizes," - " Tensor a_strides," - " Tensor b_strides," - " Tensor c_strides," - " Tensor group_scale_strides," - " str? maybe_schedule" - ") -> ()"); - - ops.def( - "cutlass_encode_and_reorder_int4b_grouped(Tensor b_tensors) -> (Tensor, " - "Tensor)"); - - // SM100 CUTLASS MLA decode - // conditionally compiled so impl registrations are in source file - ops.def( - "sm100_cutlass_mla_decode(Tensor! out, Tensor! lse, Tensor q_nope," - " Tensor q_pe, Tensor kv_c_and_k_pe_cache," - " Tensor seq_lens, Tensor page_table," - " Tensor workspace, float scale," - " int num_kv_splits) -> ()"); - - ops.def( - "sm100_cutlass_mla_get_workspace_size(int max_seq_len, int num_batches," - " int sm_count, int num_kv_splits) " - "-> int"); - // Quantized GEMM for AWQ. - ops.def( - "awq_gemm(Tensor _in_feats, Tensor _kernel, Tensor _scaling_factors, " - "Tensor _zeros, SymInt split_k_iters) -> Tensor"); - - // Dequantization for AWQ. - ops.def( - "awq_dequantize(Tensor _kernel, Tensor _scaling_factors, " - "Tensor _zeros, SymInt split_k_iters, int thx, int thy) -> Tensor"); - - // DeepSeek V3 fused A GEMM (SM 9.0+, bf16 only, 1-16 tokens). - // conditionally compiled so impl registration is in source file - ops.def( - "dsv3_fused_a_gemm(Tensor! output, Tensor mat_a, Tensor mat_b) -> ()"); - - // BF16/FP32 x FP32 -> FP32 router GEMM for H=3072, E=256, M<=32 (SM90+). - // conditionally compiled so impl registration is in source file - ops.def("fp32_router_gemm(Tensor! output, Tensor mat_a, Tensor mat_b) -> ()"); - - // reorder weight for AllSpark Ampere W8A16 Fused Gemm kernel - ops.def( - "rearrange_kn_weight_as_n32k16_order(Tensor b_qweight, Tensor b_scales, " - "Tensor? b_zeros, " - "bool has_zp, Tensor! b_qweight_reorder, Tensor! b_scales_reorder, " - "Tensor!? b_zeros_reorder, " - "int K, int N, int N_32align) -> ()"); - - // AllSpark quantization ops - ops.def( - "allspark_w8a16_gemm(Tensor a, Tensor b_qweight, Tensor b_scales, " - "Tensor? b_qzeros, " - "SymInt n, SymInt group_size, SymInt sm_count, SymInt sm_version, SymInt " - "CUBLAS_M_THRESHOLD, bool has_zp, bool n32k16_reorder) -> Tensor"); -#endif - - // Merge attn states - // Implements section 2.2 of https://www.arxiv.org/pdf/2501.01005 - // can be used to combine partial attention results (in the split-KV case) - ops.def( - "merge_attn_states(" - " Tensor! output," - " Tensor!? output_lse," - " Tensor prefix_output," - " Tensor prefix_lse," - " Tensor suffix_output," - " Tensor suffix_lse," - " int!? prefill_tokens_with_context," - " Tensor? output_scale=None) -> ()"); - - // Hadamard transforms - // conditionally compiled so impl registration is in source file - ops.def("hadacore_transform(Tensor! x, bool inplace) -> Tensor"); - - // Apply Root Mean Square (RMS) Normalization to the input tensor. - ops.def( - "rms_norm(Tensor! result, Tensor input, Tensor? weight, float epsilon) " - "-> " - "()"); - - // In-place fused Add and RMS Normalization. - ops.def( - "fused_add_rms_norm(Tensor! input, Tensor! residual, Tensor? weight, " - "float epsilon) -> ()"); - - // Layernorm-quant - // Apply Root Mean Square (RMS) Normalization to the input tensor. - ops.def( - "rms_norm_static_fp8_quant(Tensor! result, Tensor input, Tensor weight, " - "Tensor scale, float epsilon) -> " - "()"); - - // In-place fused Add and RMS Normalization. - ops.def( - "fused_add_rms_norm_static_fp8_quant(Tensor! result, Tensor input, " - "Tensor! residual, Tensor weight, " - "Tensor scale, float epsilon) -> ()"); - - // Fused Layernorm + Quant kernels - ops.def( - "rms_norm_dynamic_per_token_quant(Tensor! result, Tensor input, " - "Tensor weight, Tensor! scale, float epsilon, " - "Tensor? scale_ub, Tensor!? residual) -> ()"); - - // Fused Layernorm + Block quant kernels - ops.def( - "rms_norm_per_block_quant(Tensor! result, Tensor input, " - "Tensor weight, Tensor! scale, float epsilon, " - "Tensor? scale_ub, Tensor!? residual, int group_size, " - "bool is_scale_transposed) -> ()"); - - // Fused SiLU+Mul + per-block quantization - ops.def( - "silu_and_mul_per_block_quant(" - "Tensor! out, " - "Tensor input, " - "Tensor! scales, " - "int group_size, " - "Tensor? scale_ub=None, " - "bool is_scale_transposed=False) -> ()"); - - // Rotary embedding - // Apply GPT-NeoX or GPT-J style rotary embedding to query and key. - ops.def( - "rotary_embedding(Tensor positions, Tensor! query," - " Tensor!? key, int head_size," - " Tensor cos_sin_cache, bool is_neox, int " - "rope_dim_offset=0, bool inverse=False) -> ()"); - - // Function for fused QK Norm and RoPE - ops.def( - "fused_qk_norm_rope(Tensor! qkv, int num_heads_q, " - "int num_heads_k, int num_heads_v, int head_dim, float eps, " - "Tensor q_weight, Tensor k_weight, Tensor cos_sin_cache, " - "bool is_neox, Tensor position_ids, " - "int forced_token_heads_per_warp=-1) -> ()"); - - ops.def( - "fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(" - "Tensor q_in, Tensor kv, Tensor! k_cache, " - "Tensor slot_mapping, Tensor position_ids, Tensor cos_sin_cache, " - "int q_head_padded, float eps, int cache_block_size) -> Tensor"); - - // FlashInfer V4 full-cache variants: write Q in place (bf16) or to a separate - // FP8 tensor, and KV into a contiguous 512-wide token-strided cache. - ops.def( - "fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_bf16_insert(" - "Tensor! q, Tensor kv, Tensor! k_cache, Tensor slot_mapping, " - "Tensor position_ids, Tensor cos_sin_cache, float eps, " - "int cache_block_size) -> ()"); - ops.def( - "fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_fp8_insert(" - "Tensor q, Tensor kv, Tensor! q_fp8, Tensor! k_cache, " - "Tensor slot_mapping, Tensor position_ids, Tensor cos_sin_cache, " - "Tensor fp8_scale, Tensor q_fp8_scale_inv, float eps, " - "int cache_block_size) -> ()"); - -#ifndef USE_ROCM - ops.def( - "minimax_allreduce_rms_qk(" - "Tensor qkv, Tensor norm_weight_q, Tensor norm_weight_k, " - "Tensor workspace, int q_size, int kv_size, int rank, int nranks, " - "float eps) -> (Tensor, Tensor)"); -#endif - - // Horizontally-fused MiniMax-M3 QK-norm + partial NeoX RoPE + KV-insert. - ops.def( - "fused_minimax_m3_qknorm_rope_kv_insert(" - "Tensor! qkv, Tensor q_norm_weight, Tensor k_norm_weight, " - "Tensor cos_sin_cache, Tensor positions, int num_heads, " - "int num_kv_heads, int rotary_dim, float eps, " - "Tensor? index_q_norm_weight, Tensor? index_k_norm_weight, " - "int num_index_heads, " - "Tensor? slot_mapping, Tensor? index_slot_mapping, " - "Tensor!? kv_cache, Tensor!? index_cache, " - "int block_size, Tensor!? q_out, Tensor!? index_q_out, " - "str kv_cache_dtype, bool skip_index_branch=False) -> ()"); - -#ifdef APHRODITE_ENABLE_KIMI_K3_ATTN_RES - ops.def( - "kimi_k3_attn_res(" - "Tensor! prefix, Tensor delta, Tensor blocks, Tensor norm_weight, " - "Tensor qk_weight, Tensor output_norm_weight, Tensor! output, " - "int num_blocks, float eps, float output_norm_eps) -> ()"); -#endif - - // Apply repetition penalties to logits in-place. - ops.def( - "apply_repetition_penalties_(Tensor! logits, Tensor prompt_mask, " - "Tensor output_mask, Tensor repetition_penalties) -> ()"); - - // Optimized top-k per row operations. - ops.def( - "top_k_per_row_prefill(Tensor logits, Tensor rowStarts, Tensor rowEnds, " - "Tensor! indices, int numRows, int stride0, " - "int stride1, int topK) -> ()"); - - ops.def( - "top_k_per_row_decode(Tensor logits, int next_n, " - "Tensor seq_lens, Tensor! indices, " - "int numRows, int stride0, int stride1, int topK) -> ()"); - - ops.def( - "persistent_topk(Tensor logits, Tensor lengths, Tensor! output, " - "Tensor workspace, int k, int max_seq_len) -> ()"); - -#ifdef APHRODITE_ENABLE_COOPERATIVE_TOPK - ops.def( - "cooperative_topk(Tensor logits, Tensor lengths, Tensor! output, " - "Tensor workspace, int k, int max_seq_len) -> ()"); -#endif - -#ifdef APHRODITE_ENABLE_SM89_DSA - // sm89 DeepSeek sparse attention; impls are registered in the kernel sources - ops.def( - "sm89_fp8_paged_mqa_logits(Tensor q, Tensor pool, Tensor weights, " - "Tensor seq_lens, Tensor block_table, Tensor sched, Tensor! logits, " - "bool clean_logits) -> ()"); - - ops.def( - "sm89_paged_mqa_logits_metadata(Tensor seq_lens, Tensor! sched, " - "int next_n) -> ()"); - - ops.def( - "sm89_fp8_mqa_logits(Tensor q, Tensor kv, Tensor kv_scales, " - "Tensor weights, Tensor cu_seqlen_ks, Tensor cu_seqlen_ke, " - "Tensor! logits) -> ()"); - - ops.def( - "sm89_sparse_mla_fwd(Tensor q, Tensor pool, Tensor indices, " - "Tensor! out, Tensor! lse, float sm_scale, " - "Tensor? topk_lens=None) -> ()"); -#endif - - // Activation ops - ops.def( - "persistent_masked_m_silu_mul_quant(Tensor input, Tensor counts, Tensor! " - "y_q, Tensor! y_s, bool use_ue8m0) -> ()"); - ops.def("weak_ref_tensor(Tensor input) -> Tensor"); - - // Activation function used in SwiGLU. - ops.def("silu_and_mul(Tensor! result, Tensor input) -> ()"); - - ops.def("mul_and_silu(Tensor! out, Tensor input) -> ()"); - - // SwiGLU activation with input clamping. - // alpha scales the sigmoid (gate * sigmoid(alpha * gate)); beta is added to - // the up half (up + beta). Defaults alpha=1.0, beta=0.0 give silu(gate)*up. - ops.def( - "silu_and_mul_with_clamp(Tensor! result, Tensor input, float limit, " - "float alpha=1.0, float beta=0.0) -> ()"); - - // SwiGLU activation with FP8 quantization. - ops.def( - "silu_and_mul_quant(Tensor! result, Tensor input, Tensor scale) -> ()"); - - // Activation function used in GeGLU with `none` approximation. - ops.def("gelu_and_mul(Tensor! out, Tensor input) -> ()"); - - // Activation function used in GeGLU with `tanh` approximation. - ops.def("gelu_tanh_and_mul(Tensor! out, Tensor input) -> ()"); - - // FATReLU implementation. - ops.def("fatrelu_and_mul(Tensor! out, Tensor input, float threshold) -> ()"); - - ops.def( - "swigluoai_and_mul(Tensor! out, Tensor input, float alpha=1.702, float " - "limit=7.0) " - "-> ()"); - - // GELU implementation used in GPT-2. - ops.def("gelu_new(Tensor! out, Tensor input) -> ()"); - - // Approximate GELU implementation. - ops.def("gelu_fast(Tensor! out, Tensor input) -> ()"); - - // Quick GELU implementation. - ops.def("gelu_quick(Tensor! out, Tensor input) -> ()"); - - // relu(x)^2 activation from https://arxiv.org/abs/2109.08668v2 - ops.def("relu_squared(Tensor! out, Tensor input) -> ()"); - - // Compute int8 quantized tensor for given scaling factor. - ops.def( - "static_scaled_int8_quant(Tensor! result, Tensor input, Tensor scale," - "Tensor? azp) -> ()"); - - // Compute int8 quantized tensor and scaling factor - ops.def( - "dynamic_scaled_int8_quant(Tensor! result, Tensor input, Tensor! scale, " - "Tensor!? azp) -> ()"); - - // Compute FP8 quantized tensor for given scaling factor. - // Supports per-tensor, per-channel, per-token, and arbitrary 2D group - // scaling. Optional group_m/group_n specify the group shape explicitly; - // required for 1D scales to disambiguate per-channel vs per-token. - ops.def( - "static_scaled_fp8_quant(Tensor! result, Tensor input, Tensor scale, " - "int[]? group_shape=None) -> ()"); - - // Compute dynamic-per-tensor FP8 quantized tensor and scaling factor. - ops.def( - "dynamic_scaled_fp8_quant(Tensor! result, Tensor input, Tensor! scale) " - "-> " - "()"); - - // Compute dynamic-per-token FP8 quantized tensor and scaling factor. - ops.def( - "dynamic_per_token_scaled_fp8_quant(Tensor! result, Tensor input, " - "Tensor! scale, Tensor? scale_ub) -> " - "()"); - - // Quantized GEMM for GPTQ. - // Note: even though the C++ inferred schema is correct for this op, it seems - // to prevent the meta function registry. - ops.def( - "gptq_gemm(Tensor a, Tensor b_q_weight, Tensor b_gptq_qzeros, " - "Tensor b_gptq_scales, Tensor b_g_idx, bool use_exllama, bool " - "use_v2_format, int bit) " - "-> Tensor"); - - // Post processing for GPTQ. - ops.def("gptq_shuffle(Tensor! q_weight, Tensor q_perm, int bit) -> ()"); - - // Mamba selective scan kernel - ops.def( - "selective_scan_fwd(Tensor! u, Tensor! delta," - "Tensor! A, Tensor! B, Tensor! C," - "Tensor? D_, Tensor!? z_, Tensor? delta_bias_," - "bool delta_softplus," - "Tensor? query_start_loc," - "Tensor? cache_indices," - "Tensor? has_initial_state," - "Tensor! ssm_states," - "int null_block_id," - "int block_size," - "Tensor? block_idx_first_scheduled_token," - "Tensor? block_idx_last_scheduled_token," - "Tensor? initial_state_idx," - "Tensor? cu_chunk_seqlen," - "Tensor? last_chunk_indices) -> ()"); - - // LongCat n-gram embedding index kernel. All tensor args are marked mutable - // to match the (non-const) stable-Tensor& C++ signature; only ne_token_table - // and n_gram_ids are actually written in place. - ops.def( - "ngram_compute_n_gram_ids(int ne_n, int ne_k, Tensor(a!) ne_weights, " - "Tensor(b!) ne_mods, Tensor(c!) exclusive_ne_embedder_size_sums, " - "Tensor(d!) exclusive_req_len_sums, Tensor(e!) ne_token_table, " - "Tensor(f!) row_indices, Tensor(g!) column_starts, " - "Tensor(h!) n_gram_ids) -> ()"); -} - -STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) { - // LongCat n-gram embedding index kernel. - ops.impl("ngram_compute_n_gram_ids", TORCH_BOX(&ngram_compute_n_gram_ids)); - - // Per-token group quantization - ops.impl("per_token_group_fp8_quant", TORCH_BOX(&per_token_group_quant_fp8)); - ops.impl("per_token_group_fp8_quant_packed", - TORCH_BOX(&per_token_group_quant_8bit_packed)); - ops.impl("per_token_group_quant_int8", - TORCH_BOX(&per_token_group_quant_int8)); - - ops.impl("permute_cols", TORCH_BOX(&permute_cols)); - -#ifndef USE_ROCM - // CUTLASS scaled_mm ops - ops.impl("cutlass_scaled_mm", TORCH_BOX(&cutlass_scaled_mm)); - ops.impl("cutlass_scaled_mm_azp", TORCH_BOX(&cutlass_scaled_mm_azp)); - ops.impl("cutlass_moe_mm", TORCH_BOX(&cutlass_moe_mm)); - ops.impl("get_cutlass_moe_mm_data", TORCH_BOX(&get_cutlass_moe_mm_data)); - ops.impl("get_cutlass_moe_mm_problem_sizes_from_expert_offsets", - TORCH_BOX(&get_cutlass_moe_mm_problem_sizes_from_expert_offsets)); - ops.impl("get_cutlass_batched_moe_mm_data", - TORCH_BOX(&get_cutlass_batched_moe_mm_data)); - - // FP4/NVFP4 ops - ops.impl("cutlass_scaled_fp4_mm", TORCH_BOX(&cutlass_scaled_fp4_mm)); - ops.impl("scaled_fp4_quant", TORCH_BOX(&scaled_fp4_quant_func)); - ops.impl("scaled_fp4_quant.out", TORCH_BOX(&scaled_fp4_quant_out)); - ops.impl("scaled_fp4_experts_quant", TORCH_BOX(&scaled_fp4_experts_quant)); - ops.impl("silu_and_mul_scaled_fp4_experts_quant", - TORCH_BOX(&silu_and_mul_scaled_fp4_experts_quant)); - ops.impl("silu_and_mul_nvfp4_quant", TORCH_BOX(&silu_and_mul_nvfp4_quant)); - // mxfp4_experts_quant: registered in mxfp4_experts_quant.cu (SM100 only). - // W4A8 ops: registered in w4a8_mm_entry.cu / w4a8_grouped_mm_entry.cu. - - // AWQ ops - ops.impl("awq_gemm", TORCH_BOX(&awq_gemm)); - ops.impl("awq_dequantize", TORCH_BOX(&awq_dequantize)); - - // DSV3 fused A GEMM: conditionally compiled so impl registration is in - // source file (dsv3_fused_a_gemm.cu) - - // AllSpark ops: conditionally compiled so impl registrations are in source - // files (allspark_repack.cu and allspark_qgemm_w8a16.cu) -#endif - - ops.impl("merge_attn_states", TORCH_BOX(&merge_attn_states)); - - // Layernorm kernels (shared CUDA/ROCm) - ops.impl("rms_norm", TORCH_BOX(&rms_norm)); - ops.impl("fused_add_rms_norm", TORCH_BOX(&fused_add_rms_norm)); - - // Layernorm-quant kernels (shared CUDA/ROCm) - ops.impl("rms_norm_static_fp8_quant", TORCH_BOX(&rms_norm_static_fp8_quant)); - ops.impl("fused_add_rms_norm_static_fp8_quant", - TORCH_BOX(&fused_add_rms_norm_static_fp8_quant)); - - // Fused layernorm + dynamic per-token quant kernels (shared CUDA/ROCm) - ops.impl("rms_norm_dynamic_per_token_quant", - TORCH_BOX(&rms_norm_dynamic_per_token_quant)); - ops.impl("rms_norm_per_block_quant", TORCH_BOX(&rms_norm_per_block_quant)); - ops.impl("silu_and_mul_per_block_quant", - TORCH_BOX(&silu_and_mul_per_block_quant)); - - // Positional encoding kernels (shared CUDA/ROCm) - ops.impl("rotary_embedding", TORCH_BOX(&rotary_embedding)); - ops.impl("fused_qk_norm_rope", TORCH_BOX(&fused_qk_norm_rope)); - ops.impl("fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert", - TORCH_BOX(&fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert)); - ops.impl( - "fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_bf16_insert", - TORCH_BOX(&fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_bf16_insert)); - ops.impl( - "fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_fp8_insert", - TORCH_BOX(&fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_fp8_insert)); -#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_KIMI_K3_ATTN_RES - ops.impl("kimi_k3_attn_res", TORCH_BOX(&kimi_k3_attn_res)); -#endif - - // Sampler kernels (shared CUDA/ROCm) - ops.impl("apply_repetition_penalties_", - TORCH_BOX(&apply_repetition_penalties_)); - ops.impl("top_k_per_row_prefill", TORCH_BOX(&top_k_per_row_prefill)); - ops.impl("top_k_per_row_decode", TORCH_BOX(&top_k_per_row_decode)); - ops.impl("persistent_topk", TORCH_BOX(&persistent_topk)); -#ifdef APHRODITE_ENABLE_COOPERATIVE_TOPK - ops.impl("cooperative_topk", TORCH_BOX(&cooperative_topk)); -#endif - - // Activation kernels (shared CUDA/ROCm) - ops.impl("persistent_masked_m_silu_mul_quant", - TORCH_BOX(&persistent_masked_m_silu_mul_quant)); - ops.impl("weak_ref_tensor", TORCH_BOX(&weak_ref_tensor)); - ops.impl("silu_and_mul_quant", TORCH_BOX(&silu_and_mul_quant)); - ops.impl("silu_and_mul", TORCH_BOX(&silu_and_mul)); - ops.impl("mul_and_silu", TORCH_BOX(&mul_and_silu)); - ops.impl("gelu_and_mul", TORCH_BOX(&gelu_and_mul)); - ops.impl("gelu_tanh_and_mul", TORCH_BOX(&gelu_tanh_and_mul)); - ops.impl("fatrelu_and_mul", TORCH_BOX(&fatrelu_and_mul)); - ops.impl("swigluoai_and_mul", TORCH_BOX(&swigluoai_and_mul)); - ops.impl("gelu_new", TORCH_BOX(&gelu_new)); - ops.impl("gelu_fast", TORCH_BOX(&gelu_fast)); - ops.impl("gelu_quick", TORCH_BOX(&gelu_quick)); - ops.impl("relu_squared", TORCH_BOX(&relu_squared)); - ops.impl("silu_and_mul_with_clamp", TORCH_BOX(&silu_and_mul_clamp)); - - // INT8 quantization kernels - ops.impl("static_scaled_int8_quant", TORCH_BOX(&static_scaled_int8_quant)); - ops.impl("dynamic_scaled_int8_quant", TORCH_BOX(&dynamic_scaled_int8_quant)); - - // FP8 quantization kernels - ops.impl("static_scaled_fp8_quant", TORCH_BOX(&static_scaled_fp8_quant)); - ops.impl("dynamic_scaled_fp8_quant", TORCH_BOX(&dynamic_scaled_fp8_quant)); - ops.impl("dynamic_per_token_scaled_fp8_quant", - TORCH_BOX(&dynamic_per_token_scaled_fp8_quant)); - - // GPTQ kernels - ops.impl("gptq_gemm", TORCH_BOX(&gptq_gemm)); - ops.impl("gptq_shuffle", TORCH_BOX(&gptq_shuffle)); - - // Mamba kernels - ops.impl("selective_scan_fwd", TORCH_BOX(&selective_scan_fwd)); -} - -STABLE_TORCH_LIBRARY_IMPL(_C, CPU, ops) { - ops.impl("get_cuda_view_from_cpu_tensor", - TORCH_BOX(&get_cuda_view_from_cpu_tensor)); - ops.impl("dry_scan_penalties", TORCH_BOX(&dry_scan_penalties_cpu)); -} - -STABLE_TORCH_LIBRARY_FRAGMENT(_C_cuda_utils, cuda_utils) { - cuda_utils.def("get_device_attribute(int attribute, int device_id) -> int"); - cuda_utils.def( - "get_max_shared_memory_per_block_device_attribute(int device_id) -> int"); -} - -STABLE_TORCH_LIBRARY_IMPL(_C_cuda_utils, CompositeExplicitAutograd, - cuda_utils) { - cuda_utils.impl("get_device_attribute", TORCH_BOX(&get_device_attribute)); - cuda_utils.impl("get_max_shared_memory_per_block_device_attribute", - TORCH_BOX(&get_max_shared_memory_per_block_device_attribute)); -} - -// These capability-check functions take only primitive args (no tensors), so -// there is no device to dispatch on. CompositeExplicitAutograd makes them -// available for all backends. This is the stable ABI equivalent of calling -// ops.impl("op_name", &func) without a dispatch key in the non-stable API. -STABLE_TORCH_LIBRARY_IMPL(_C, CompositeExplicitAutograd, ops) { -#ifndef USE_ROCM - ops.impl("cutlass_scaled_mm_supports_fp8", - TORCH_BOX(&cutlass_scaled_mm_supports_fp8)); - ops.impl("cutlass_group_gemm_supported", - TORCH_BOX(&cutlass_group_gemm_supported)); - ops.impl("cutlass_scaled_mm_supports_block_fp8", - TORCH_BOX(&cutlass_scaled_mm_supports_block_fp8)); - ops.impl("cutlass_scaled_mm_supports_fp4", - TORCH_BOX(&cutlass_scaled_mm_supports_fp4)); -#endif -} - -// Cache ops -STABLE_TORCH_LIBRARY_FRAGMENT(_C_cache_ops, ops) { - // Swap in (out) the cache blocks from src to dst. - ops.def( - "swap_blocks(Tensor src, Tensor! dst," - " int block_size_in_bytes, Tensor block_mapping) -> ()"); - - // Batch swap: submit all block copies in a single driver call. - ops.def( - "swap_blocks_batch(Tensor src_ptrs, Tensor dst_ptrs," - " Tensor sizes," - " bool is_src_access_order_any=False) -> ()"); - - // Reshape the key and value tensors and cache them. - ops.def( - "reshape_and_cache(Tensor key, Tensor value," - " Tensor! key_cache, Tensor! value_cache," - " Tensor slot_mapping," - " str kv_cache_dtype," - " Tensor k_scale, Tensor v_scale) -> ()"); - - // Reshape the key and value tensors and cache them. - ops.def( - "reshape_and_cache_flash(Tensor key, Tensor value," - " Tensor! key_cache," - " Tensor! value_cache," - " Tensor slot_mapping," - " str kv_cache_dtype," - " Tensor k_scale, Tensor v_scale) -> ()"); - - // Concat kv_c and k_pe and cache them. - ops.def( - "concat_and_cache_mla(Tensor kv_c, Tensor k_pe," - " Tensor! kv_cache," - " Tensor slot_mapping," - " str kv_cache_dtype," - " Tensor scale) -> ()"); - - // Rotate Q and K, then write to kv cache for MLA - ops.def( - "concat_and_cache_mla_rope_fused(" - " Tensor positions," - " Tensor! q_pe," - " Tensor! k_pe," - " Tensor kv_c," - " Tensor cos_sin_cache," - " bool is_neox," - " Tensor slot_mapping," - " Tensor! kv_cache," - " str kv_cache_dtype," - " Tensor kv_cache_scale) -> ()"); - - // Convert the key and value cache to fp8 data type. - ops.def( - "convert_fp8(Tensor! dst_cache, Tensor src_cache, float scale, " - "str kv_cache_dtype) -> ()"); - - // Gather cache blocks from src_cache to dst, dequantizing from - // src_cache's dtype to dst's dtype if necessary. - ops.def( - "gather_and_maybe_dequant_cache(Tensor src_cache, Tensor! dst, " - " Tensor block_table, Tensor cu_seq_lens, " - " Tensor token_to_seq, " - " int num_tokens, " - " str kv_cache_dtype, " - " Tensor scale, Tensor? seq_starts) -> ()"); - - ops.def( - "cp_gather_cache(Tensor src_cache, Tensor! dst, Tensor block_table, " - "Tensor cu_seq_lens, int batch_size, Tensor? seq_starts) -> ()"); - - ops.def( - "cp_gather_and_upconvert_fp8_kv_cache(Tensor src_cache, Tensor! dst, " - "Tensor block_table, Tensor workspace_starts, int batch_size, Tensor? " - "seq_starts) -> ()"); - - ops.def( - "indexer_k_quant_and_cache(Tensor k, Tensor! kv_cache, Tensor " - "slot_mapping, " - "int quant_block_size, str kv_cache_dtype) -> ()"); - - ops.def("concat_mla_q(Tensor ql_nope, Tensor q_pe, Tensor! q_out) -> ()"); - - ops.def( - "cp_gather_indexer_k_quant_cache(Tensor kv_cache, Tensor! dst_k, Tensor! " - "dst_scale, Tensor block_table, Tensor cu_seq_lens) -> ()"); -} - -STABLE_TORCH_LIBRARY_FRAGMENT(_C_custom_ar, custom_ar) { - custom_ar.def( - "init_custom_ar(int[] ipc_tensors, Tensor rank_data, " - "int rank, bool fully_connected) -> int"); - custom_ar.def( - "all_reduce(int fa, Tensor inp, Tensor! out, int reg_buffer, " - "int reg_buffer_sz_bytes) -> ()"); - custom_ar.def("dispose(int fa) -> ()"); - custom_ar.def("meta_size() -> int"); - custom_ar.def("register_buffer(int fa, int[] ipc_tensors) -> ()"); - custom_ar.def("get_graph_buffer_ipc_meta(int fa) -> (int[], int[])"); - custom_ar.def( - "register_graph_buffers(int fa, int[][] handles, int[][] offsets) -> ()"); - custom_ar.def("allocate_shared_buffer_and_handle(int size) -> (int, Tensor)"); - custom_ar.def("open_mem_handle(Tensor mem_handle) -> int"); - custom_ar.def("free_shared_buffer(int ptr) -> ()"); -} - -STABLE_TORCH_LIBRARY_IMPL(_C_custom_ar, CUDA, custom_ar) { - custom_ar.impl("init_custom_ar", TORCH_BOX(&init_custom_ar)); - custom_ar.impl("all_reduce", TORCH_BOX(&all_reduce)); -} - -STABLE_TORCH_LIBRARY_IMPL(_C_custom_ar, CPU, custom_ar) { - custom_ar.impl("open_mem_handle", TORCH_BOX(&open_mem_handle)); -} - -STABLE_TORCH_LIBRARY_IMPL(_C_custom_ar, CompositeExplicitAutograd, custom_ar) { - custom_ar.impl("dispose", TORCH_BOX(&dispose)); - custom_ar.impl("meta_size", TORCH_BOX(&meta_size)); - custom_ar.impl("register_buffer", TORCH_BOX(®ister_buffer)); - custom_ar.impl("get_graph_buffer_ipc_meta", - TORCH_BOX(&get_graph_buffer_ipc_meta)); - custom_ar.impl("register_graph_buffers", TORCH_BOX(®ister_graph_buffers)); - custom_ar.impl("allocate_shared_buffer_and_handle", - TORCH_BOX(&allocate_shared_buffer_and_handle)); - custom_ar.impl("free_shared_buffer", TORCH_BOX(&free_shared_buffer)); -} - -STABLE_TORCH_LIBRARY_IMPL(_C_cache_ops, CPU, ops) { - ops.impl("swap_blocks_batch", TORCH_BOX(&swap_blocks_batch)); -} - -STABLE_TORCH_LIBRARY_IMPL(_C_cache_ops, CUDA, ops) { - ops.impl("swap_blocks", TORCH_BOX(&swap_blocks)); - ops.impl("reshape_and_cache", TORCH_BOX(&reshape_and_cache)); - ops.impl("reshape_and_cache_flash", TORCH_BOX(&reshape_and_cache_flash)); - ops.impl("concat_and_cache_mla", TORCH_BOX(&concat_and_cache_mla)); - ops.impl("concat_and_cache_mla_rope_fused", - TORCH_BOX(&concat_and_cache_mla_rope_fused)); - ops.impl("convert_fp8", TORCH_BOX(&convert_fp8)); - ops.impl("gather_and_maybe_dequant_cache", - TORCH_BOX(&gather_and_maybe_dequant_cache)); - ops.impl("cp_gather_cache", TORCH_BOX(&cp_gather_cache)); - ops.impl("cp_gather_and_upconvert_fp8_kv_cache", - TORCH_BOX(&cp_gather_and_upconvert_fp8_kv_cache)); - ops.impl("indexer_k_quant_and_cache", TORCH_BOX(&indexer_k_quant_and_cache)); - ops.impl("concat_mla_q", TORCH_BOX(&concat_mla_q)); - ops.impl("cp_gather_indexer_k_quant_cache", - TORCH_BOX(&cp_gather_indexer_k_quant_cache)); -} - -REGISTER_EXTENSION(_C_stable_libtorch) diff --git a/setup.py.orig b/setup.py.orig deleted file mode 100644 index 1ca034e686..0000000000 --- a/setup.py.orig +++ /dev/null @@ -1,796 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -import ctypes -import importlib.util -import logging -import os -import re -import shutil -import subprocess -import sys -import sysconfig -from pathlib import Path -from shutil import which - -import torch -from packaging.version import Version, parse -from setuptools import Extension, setup -from setuptools.command.build_ext import build_ext -from setuptools_scm import get_version -from torch.utils.cpp_extension import CUDA_HOME, ROCM_HOME - - -def load_module_from_path(module_name, path): - spec = importlib.util.spec_from_file_location(module_name, path) - module = importlib.util.module_from_spec(spec) - sys.modules[module_name] = module - spec.loader.exec_module(module) - return module - - -ROOT_DIR = Path(__file__).parent -logger = logging.getLogger(__name__) - -PRECOMPILED_RUST_FRONTEND_PATH = ROOT_DIR / "aphrodite" / "aphrodite-rs" -# cannot import envs directly because it depends on aphrodite, -# which is not installed yet -envs = load_module_from_path("envs", os.path.join(ROOT_DIR, "aphrodite", "envs.py")) - -APHRODITE_TARGET_DEVICE = envs.APHRODITE_TARGET_DEVICE - -# Skips all native builds (CMake/CUDA and Rust) and extracts binaries from the -# matching per-commit wheel. -APHRODITE_USE_PRECOMPILED = envs.APHRODITE_USE_PRECOMPILED -precompiled_wheel = load_module_from_path( - "precompiled_wheel", - os.path.join(ROOT_DIR, "aphrodite", "build_utils", "precompiled.py"), -) - - -def should_require_rust_frontend() -> bool: - value = os.getenv("APHRODITE_REQUIRE_RUST_FRONTEND", "") - return value.lower() not in ("", "0", "false", "no") - - -def should_use_precompiled_rust() -> bool: - value = os.getenv("APHRODITE_USE_PRECOMPILED_RUST", "") - return value.lower() in ("1", "true", "yes") - - -def get_precompiled_rust_extension_paths() -> list[Path]: - return sorted((ROOT_DIR / "aphrodite").glob("_rust_*.so")) - - -if sys.platform.startswith("darwin") and APHRODITE_TARGET_DEVICE not in ("cpu", "metal"): - logger.warning("APHRODITE_TARGET_DEVICE automatically set to `metal` due to macOS") - APHRODITE_TARGET_DEVICE = "metal" -elif not (sys.platform.startswith("linux") or sys.platform.startswith("darwin")): - logger.warning( - "Aphrodite only supports Linux platform (including WSL) and MacOS." - "Building on %s, " - "so Aphrodite may not be able to run correctly", - sys.platform, - ) - APHRODITE_TARGET_DEVICE = "empty" -elif sys.platform.startswith("linux") and os.getenv("APHRODITE_TARGET_DEVICE") is None: - if torch.version.hip is not None: - APHRODITE_TARGET_DEVICE = "rocm" - logger.info("Auto-detected ROCm") - elif torch.version.xpu is not None: - APHRODITE_TARGET_DEVICE = "xpu" - logger.info("Auto-detected XPU") - elif torch.version.cuda is not None: - APHRODITE_TARGET_DEVICE = "cuda" - logger.info("Auto-detected CUDA") - else: - APHRODITE_TARGET_DEVICE = "cpu" - - -def is_sccache_available() -> bool: - return which("sccache") is not None and not bool(int(os.getenv("APHRODITE_DISABLE_SCCACHE", "0"))) - - -def is_ccache_available() -> bool: - return which("ccache") is not None - - -def is_ninja_available() -> bool: - return which("ninja") is not None - - -def is_freethreaded(): - return bool(sysconfig.get_config_var("Py_GIL_DISABLED")) - - -def should_bundle_tcmalloc() -> bool: - import platform - - return ( - APHRODITE_TARGET_DEVICE == "cpu" - and sys.platform.startswith("linux") - and platform.machine() in ("aarch64", "x86_64") - ) - - -def find_tcmalloc() -> Path | None: - try: - # get all shared libs the dynamic loader knows about - output = subprocess.check_output( - ["ldconfig", "-p"], - text=True, - stderr=subprocess.DEVNULL, - ) - except Exception: - return None - - # search for libtcmalloc and libtcmalloc_minimal - for library_pattern in ( - r"\blibtcmalloc_minimal\.so\.(\d+)\b", - r"\blibtcmalloc\.so\.(\d+)\b", - ): - candidates: list[tuple[int, Path]] = [] - for line in output.splitlines(): - match = re.search(library_pattern, line) - if match is None or "=>" not in line: - continue - candidate = Path(line.split("=>")[1].strip()) - if candidate.exists(): - candidates.append((int(match.group(1)), candidate)) - - if candidates: - # if multiple candidates are found, pick the one with the highest - # version number - return max(candidates, key=lambda item: item[0])[1] - - return None - - -def bundle_tcmalloc(build_lib: str) -> None: - tcmalloc_library = find_tcmalloc() - if tcmalloc_library is None: - logger.warning( - "Failed to locate tcmalloc. For best performance, " - "please install tcmalloc (e.g. `sudo apt-get " - "install -y --no-install-recommends libtcmalloc-minimal4`)" - ) - return - - bundle_dir = os.path.join(build_lib, "aphrodite", "libs") - os.makedirs(bundle_dir, exist_ok=True) - bundle_path = os.path.join(bundle_dir, tcmalloc_library.name) - shutil.copy2(tcmalloc_library, bundle_path) - logger.info("Bundled tcmalloc into wheel: %s", bundle_path) - - -class CMakeExtension(Extension): - def __init__(self, name: str, cmake_lists_dir: str = ".", **kwa) -> None: - # Default to the stable/abi3 API, but let callers opt out (e.g. the - # nanobind _paged_ops Metal kernel builds a version-specific module). - kwa.setdefault("py_limited_api", not is_freethreaded()) - super().__init__(name, sources=[], **kwa) - self.cmake_lists_dir = os.path.abspath(cmake_lists_dir) - - -class cmake_build_ext(build_ext): - # A dict of extension directories that have been configured. - did_config: dict[str, bool] = {} - - # - # Determine number of compilation jobs and optionally nvcc compile threads. - # - def compute_num_jobs(self): - # `num_jobs` is either the value of the MAX_JOBS environment variable - # (if defined) or the number of CPUs available. - num_jobs = envs.MAX_JOBS - if num_jobs is not None: - num_jobs = int(num_jobs) - logger.info("Using MAX_JOBS=%d as the number of jobs.", num_jobs) - else: - try: - # os.sched_getaffinity() isn't universally available, so fall - # back to os.cpu_count() if we get an error here. - num_jobs = len(os.sched_getaffinity(0)) - except AttributeError: - num_jobs = os.cpu_count() - - nvcc_threads = None - if _is_cuda() and CUDA_HOME is not None: - try: - nvcc_version = get_nvcc_cuda_version() - if nvcc_version >= Version("11.2"): - # `nvcc_threads` is either the value of the NVCC_THREADS - # environment variable (if defined) or 1. - # when it is set, we reduce `num_jobs` to avoid - # overloading the system. - nvcc_threads = envs.NVCC_THREADS - if nvcc_threads is not None: - nvcc_threads = int(nvcc_threads) - logger.info( - "Using NVCC_THREADS=%d as the number of nvcc threads.", - nvcc_threads, - ) - else: - nvcc_threads = 1 - num_jobs = max(1, num_jobs // nvcc_threads) - except Exception as e: - logger.warning("Failed to get NVCC version: %s", e) - - return num_jobs, nvcc_threads - - # - # Perform cmake configuration for a single extension. - # - def configure(self, ext: CMakeExtension) -> None: - # If we've already configured using the CMakeLists.txt for - # this extension, exit early. - if ext.cmake_lists_dir in cmake_build_ext.did_config: - return - - cmake_build_ext.did_config[ext.cmake_lists_dir] = True - - # Select the build type. - # Note: optimization level + debug info are set by the build type - default_cfg = "Debug" if self.debug else "RelWithDebInfo" - cfg = envs.CMAKE_BUILD_TYPE or default_cfg - - cmake_args = [ - "-DCMAKE_BUILD_TYPE={}".format(cfg), - "-DAPHRODITE_TARGET_DEVICE={}".format(APHRODITE_TARGET_DEVICE), - ] - - verbose = envs.VERBOSE - if verbose: - cmake_args += ["-DCMAKE_VERBOSE_MAKEFILE=ON"] - - if is_sccache_available(): - cmake_args += [ - "-DCMAKE_C_COMPILER_LAUNCHER=sccache", - "-DCMAKE_CXX_COMPILER_LAUNCHER=sccache", - "-DCMAKE_CUDA_COMPILER_LAUNCHER=sccache", - "-DCMAKE_HIP_COMPILER_LAUNCHER=sccache", - ] - elif is_ccache_available(): - cmake_args += [ - "-DCMAKE_C_COMPILER_LAUNCHER=ccache", - "-DCMAKE_CXX_COMPILER_LAUNCHER=ccache", - "-DCMAKE_CUDA_COMPILER_LAUNCHER=ccache", - "-DCMAKE_HIP_COMPILER_LAUNCHER=ccache", - ] - - # Pass the python executable to cmake so it can find an exact - # match. - cmake_args += ["-DAPHRODITE_PYTHON_EXECUTABLE={}".format(sys.executable)] - - # Pass the python path to cmake so it can reuse the build dependencies - # on subsequent calls to python. - cmake_args += ["-DAPHRODITE_PYTHON_PATH={}".format(":".join(sys.path))] - - # Override the base directory for FetchContent downloads to $ROOT/.deps - # This allows sharing dependencies between profiles, - # and plays more nicely with sccache. - # To override this, set the FETCHCONTENT_BASE_DIR environment variable. - fc_base_dir = os.path.join(ROOT_DIR, ".deps") - fc_base_dir = os.environ.get("FETCHCONTENT_BASE_DIR", fc_base_dir) - cmake_args += ["-DFETCHCONTENT_BASE_DIR={}".format(fc_base_dir)] - - # - # Setup parallelism and build tool - # - num_jobs, nvcc_threads = self.compute_num_jobs() - - if nvcc_threads: - cmake_args += ["-DNVCC_THREADS={}".format(nvcc_threads)] - - if is_ninja_available(): - build_tool = ["-G", "Ninja"] - cmake_args += [ - "-DCMAKE_JOB_POOL_COMPILE:STRING=compile", - "-DCMAKE_JOB_POOLS:STRING=compile={}".format(num_jobs), - ] - else: - # Default build tool to whatever cmake picks. - build_tool = [] - # Make sure we use the nvcc from CUDA_HOME - if _is_cuda() and CUDA_HOME is not None: - cmake_args += [f"-DCMAKE_CUDA_COMPILER={CUDA_HOME}/bin/nvcc"] - elif _is_hip() and ROCM_HOME is not None: - cmake_args += [f"-DROCM_PATH={ROCM_HOME}"] - - other_cmake_args = os.environ.get("CMAKE_ARGS") - if other_cmake_args: - cmake_args += other_cmake_args.split() - - subprocess.check_call( - ["cmake", ext.cmake_lists_dir, *build_tool, *cmake_args], - cwd=self.build_temp, - ) - - def build_extensions(self) -> None: - # Ensure that CMake is present and working - try: - subprocess.check_output(["cmake", "--version"]) - except OSError as e: - raise RuntimeError("Cannot find CMake executable") from e - - # Create build directory if it does not exist. - if not os.path.exists(self.build_temp): - os.makedirs(self.build_temp) - - targets = [] - - def target_name(s: str) -> str: - if s.endswith("._paged_ops"): - return "_paged_ops" - return s.removeprefix("aphrodite.").removeprefix("vllm_flash_attn.") - - # Build all the extensions - for ext in self.extensions: - self.configure(ext) - targets.append(target_name(ext.name)) - - num_jobs, _ = self.compute_num_jobs() - - build_args = [ - "--build", - ".", - f"-j={num_jobs}", - *[f"--target={name}" for name in targets], - ] - - subprocess.check_call(["cmake", *build_args], cwd=self.build_temp) - - # Install the libraries - for ext in self.extensions: - # Install the extension into the proper location - outdir = Path(self.get_ext_fullpath(ext.name)).parent.absolute() - - # Skip if the install directory is the same as the build directory - if outdir == self.build_temp: - continue - - # CMake appends the extension prefix to the install path, - # and outdir already contains that prefix, so we need to remove it. - prefix = outdir - for _ in range(ext.name.count(".")): - prefix = prefix.parent - - # prefix here should actually be the same for all components - install_args = [ - "cmake", - "--install", - ".", - "--prefix", - prefix, - "--component", - target_name(ext.name), - ] - subprocess.check_call(install_args, cwd=self.build_temp) - - def run(self): - # First, run the standard build_ext command to compile the extensions - super().run() - - # bundle tcmalloc into CPU wheels for best OOB perf - if should_bundle_tcmalloc(): - bundle_tcmalloc(self.build_lib) - - # copy aphrodite/vllm_flash_attn/**/*.py from self.build_lib to current - # directory so that they can be included in the editable build - import glob - - files = glob.glob( - os.path.join(self.build_lib, "aphrodite", "vllm_flash_attn", "**", "*.py"), - recursive=True, - ) - for file in files: - dst_file = os.path.join("aphrodite/vllm_flash_attn", file.split("aphrodite/vllm_flash_attn/")[-1]) - print(f"Copying {file} to {dst_file}") - os.makedirs(os.path.dirname(dst_file), exist_ok=True) - self.copy_file(file, dst_file) - - if _is_cuda() or _is_hip(): - # copy aphrodite/third_party/triton_kernels/**/*.py from self.build_lib - # to current directory so that they can be included in the editable - # build - print( - f"Copying {self.build_lib}/aphrodite/third_party/triton_kernels to aphrodite/third_party/triton_kernels" - ) - shutil.copytree( - f"{self.build_lib}/aphrodite/third_party/triton_kernels", - "aphrodite/third_party/triton_kernels", - dirs_exist_ok=True, - ) - - if _is_cuda(): - # copy vendored deep_gemm package from build_lib to source tree - # for editable installs - deep_gemm_build = os.path.join(self.build_lib, "aphrodite", "third_party", "deep_gemm") - if os.path.exists(deep_gemm_build): - print(f"Copying {deep_gemm_build} to aphrodite/third_party/deep_gemm") - shutil.copytree( - deep_gemm_build, - "aphrodite/third_party/deep_gemm", - dirs_exist_ok=True, - ) - - # copy vendored fmha_sm100 package from build_lib to source tree - # for editable installs - fmha_sm100_build = os.path.join(self.build_lib, "aphrodite", "third_party", "fmha_sm100") - if os.path.exists(fmha_sm100_build): - print(f"Copying {fmha_sm100_build} to aphrodite/third_party/fmha_sm100") - shutil.copytree( - fmha_sm100_build, - "aphrodite/third_party/fmha_sm100", - dirs_exist_ok=True, - ) - - # copy vendored tml-fa4 package from build_lib to source tree - # for editable installs - tml_fa4_build = os.path.join(self.build_lib, "aphrodite", "third_party", "tml_fa4") - if os.path.exists(tml_fa4_build): - print(f"Copying {tml_fa4_build} to aphrodite/third_party/tml_fa4") - shutil.copytree( - tml_fa4_build, - "aphrodite/third_party/tml_fa4", - dirs_exist_ok=True, - ) - - -def _no_device() -> bool: - return APHRODITE_TARGET_DEVICE == "empty" - - -def _is_cuda() -> bool: - has_cuda = torch.version.cuda is not None - return APHRODITE_TARGET_DEVICE == "cuda" and has_cuda and not _is_tpu() - - -def _is_hip() -> bool: - return (APHRODITE_TARGET_DEVICE == "cuda" or APHRODITE_TARGET_DEVICE == "rocm") and torch.version.hip is not None - - -def _is_tpu() -> bool: - return APHRODITE_TARGET_DEVICE == "tpu" - - -def _is_cpu() -> bool: - return APHRODITE_TARGET_DEVICE == "cpu" - - -def _is_xpu() -> bool: - return APHRODITE_TARGET_DEVICE == "xpu" - - -def _is_metal() -> bool: - return APHRODITE_TARGET_DEVICE == "metal" - - -def _build_custom_ops() -> bool: - return _is_cuda() or _is_hip() or _is_metal() - - -def get_rocm_version(): - # Get the Rocm version from the ROCM_HOME/bin/librocm-core.so - # see https://github.com/ROCm/rocm-core/blob/d11f5c20d500f729c393680a01fa902ebf92094b/rocm_version.cpp#L21 - try: - if ROCM_HOME is None: - return None - librocm_core_file = Path(ROCM_HOME) / "lib" / "librocm-core.so" - if not librocm_core_file.is_file(): - return None - librocm_core = ctypes.CDLL(librocm_core_file) - VerErrors = ctypes.c_uint32 - get_rocm_core_version = librocm_core.getROCmVersion - get_rocm_core_version.restype = VerErrors - get_rocm_core_version.argtypes = [ - ctypes.POINTER(ctypes.c_uint32), - ctypes.POINTER(ctypes.c_uint32), - ctypes.POINTER(ctypes.c_uint32), - ] - major = ctypes.c_uint32() - minor = ctypes.c_uint32() - patch = ctypes.c_uint32() - - if get_rocm_core_version(ctypes.byref(major), ctypes.byref(minor), ctypes.byref(patch)) == 0: - return f"{major.value}.{minor.value}.{patch.value}" - return None - except Exception: - return None - - -def get_nvcc_cuda_version() -> Version: - """Get the CUDA version from nvcc. - - Adapted from https://github.com/NVIDIA/apex/blob/8b7a1ff183741dd8f9b87e7bafd04cfde99cea28/setup.py - """ - assert CUDA_HOME is not None, "CUDA_HOME is not set" - nvcc_output = subprocess.check_output([CUDA_HOME + "/bin/nvcc", "-V"], universal_newlines=True) - output = nvcc_output.split() - release_idx = output.index("release") + 1 - nvcc_cuda_version = parse(output[release_idx].split(",")[0]) - return nvcc_cuda_version - - -def get_aphrodite_version() -> str: - # Allow overriding the version. This is useful to build platform-specific - # wheels (e.g. CPU, TPU) without modifying the source. - if env_version := os.getenv("APHRODITE_VERSION_OVERRIDE"): - print(f"Overriding APHRODITE version with {env_version} from APHRODITE_VERSION_OVERRIDE") - os.environ["SETUPTOOLS_SCM_PRETEND_VERSION"] = env_version - return get_version(write_to="aphrodite/_version.py") - - version = get_version(write_to="aphrodite/_version.py") - sep = "+" if "+" not in version else "." # dev versions might contain + - - if _no_device(): - if envs.APHRODITE_TARGET_DEVICE == "empty": - version += f"{sep}empty" - elif _is_cuda(): - if APHRODITE_USE_PRECOMPILED: - if not envs.APHRODITE_SKIP_PRECOMPILED_VERSION_SUFFIX: - version += f"{sep}precompiled" - else: - cuda_version = str(get_nvcc_cuda_version()) - if cuda_version != envs.APHRODITE_MAIN_CUDA_VERSION: - cuda_version_str = cuda_version.replace(".", "")[:3] - # skip this for source tarball, required for pypi - if "sdist" not in sys.argv: - version += f"{sep}cu{cuda_version_str}" - elif _is_hip(): - # Get the Rocm Version - rocm_version = get_rocm_version() or torch.version.hip - if rocm_version and rocm_version != envs.APHRODITE_MAIN_CUDA_VERSION: - version += f"{sep}rocm{rocm_version.replace('.', '')[:3]}" - elif _is_tpu(): - version += f"{sep}tpu" - elif _is_cpu(): - # Check the local APHRODITE_TARGET_DEVICE (may be set by auto-detect above), - # not envs.APHRODITE_TARGET_DEVICE, so CPU-only hosts still get `+cpu`. - if APHRODITE_TARGET_DEVICE == "cpu": - version += f"{sep}cpu" - elif _is_metal(): - version += f"{sep}metal" - elif _is_xpu(): - version += f"{sep}xpu" - else: - raise RuntimeError("Unknown runtime environment") - - return version - - -def get_requirements() -> list[str]: - """Get Python package dependencies from requirements.txt.""" - requirements_dir = ROOT_DIR / "requirements" - - def _read_requirements(filename: str) -> list[str]: - with open(requirements_dir / filename) as f: - requirements = f.read().strip().split("\n") - resolved_requirements = [] - for line in requirements: - if line.startswith("-r "): - resolved_requirements += _read_requirements(line.split()[1]) - elif not line.startswith("--") and not line.startswith("#") and line.strip() != "": - resolved_requirements.append(line) - return resolved_requirements - - if _no_device(): - requirements = _read_requirements("common.txt") - elif _is_cuda(): - requirements = _read_requirements("cuda.txt") - cuda_major, cuda_minor = torch.version.cuda.split(".") - modified_requirements = [] - for req in requirements: - if "aphrodite-flash-attn" in req and cuda_major != "12": - # aphrodite-flash-attn is built only for CUDA 12.x. - # Skip for other versions. - continue - if "flashinfer-cubin" in req: - # Not on PyPI since 0.6.14 (only https://flashinfer.ai/whl), so - # it cannot be a wheel dependency; flashinfer falls back to - # fetching cubins at runtime when the package is absent. - continue - if "nvidia-cutlass-dsl[cu13]" in req and cuda_major == "12": - # [cu13] extra is the default; strip it on CUDA 12 builds. - req = req.replace("nvidia-cutlass-dsl[cu13]", "nvidia-cutlass-dsl") - if "humming-kernels[cu13]" in req and cuda_major == "12": - req = req.replace("humming-kernels[cu13]", "humming-kernels[cu12]") - modified_requirements.append(req) - requirements = modified_requirements - elif _is_hip(): - requirements = _read_requirements("rocm.txt") - elif _is_tpu(): - requirements = _read_requirements("tpu.txt") - elif _is_cpu(): - requirements = _read_requirements("cpu.txt") - elif _is_metal(): - requirements = _read_requirements("metal.txt") - elif _is_xpu(): - requirements = _read_requirements("xpu.txt") - else: - raise ValueError("Unsupported platform, please use CUDA, ROCm, Metal, or CPU.") - return requirements - - -ext_modules = [] - -if _is_cuda() or _is_hip(): - ext_modules.append(CMakeExtension(name="aphrodite.cumem_allocator")) - # Optional since this doesn't get built (produce an .so file). This is just - # 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): - ext_modules.append(CMakeExtension(name="aphrodite.spinloop")) - ext_modules.append(CMakeExtension(name="aphrodite.fs_io_C")) - -if _is_hip(): - ext_modules.append(CMakeExtension(name="aphrodite._rocm_C")) - -if _is_cuda(): - ext_modules.append(CMakeExtension(name="aphrodite.vllm_flash_attn._vllm_fa2_C")) - if CUDA_HOME and get_nvcc_cuda_version() >= Version("12.3"): - # FA3 requires CUDA 12.3 or later - ext_modules.append(CMakeExtension(name="aphrodite.vllm_flash_attn._vllm_fa3_C")) - # FA4 CuteDSL - Python-only component for FA4's cute DSL support - # Optional since this doesn't produce a .so file, just copies Python files - ext_modules.append(CMakeExtension(name="aphrodite.vllm_flash_attn._vllm_fa4_cutedsl_C", optional=True)) - if CUDA_HOME and get_nvcc_cuda_version() >= Version("12.9"): - # FlashMLA requires CUDA 12.9 or later - # Optional since this doesn't get built (produce an .so file) when - # 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.3"): - # DeepGEMM requires CUDA 12.3+ (SM90/SM100) - # Optional since it won't build on unsupported architectures - ext_modules.append(CMakeExtension(name="aphrodite._deep_gemm_C", optional=True)) - ext_modules.append(CMakeExtension(name="aphrodite._qutlass_C", optional=True)) - # fmha_sm100 is a Python/CuTe-DSL package installed into aphrodite.third_party. - ext_modules.append(CMakeExtension(name="aphrodite.fmha_sm100", optional=True)) - # tml-fa4 is copied into an isolated aphrodite.third_party package. - ext_modules.append(CMakeExtension(name="aphrodite.tml_fa4", optional=True)) - -if _is_cpu(): - import platform - - if platform.machine() in ("x86_64", "AMD64"): - ext_modules.append(CMakeExtension(name="aphrodite._C")) - ext_modules.append(CMakeExtension(name="aphrodite._C_AVX512")) - ext_modules.append(CMakeExtension(name="aphrodite._C_AVX2")) - else: - ext_modules.append(CMakeExtension(name="aphrodite._C")) - -if _build_custom_ops(): - if _is_metal(): - # MLX/nanobind paged-attention Metal kernel (aphrodite/metal/metal/_paged_ops). - ext_modules.append(CMakeExtension(name="aphrodite.metal.metal._paged_ops", py_limited_api=False)) - if _is_hip(): - ext_modules.append(CMakeExtension(name="aphrodite._C")) - if _is_cuda() or _is_hip(): - ext_modules.append(CMakeExtension(name="aphrodite._C_stable_libtorch")) - ext_modules.append(CMakeExtension(name="aphrodite._moe_C_stable_libtorch")) - -package_data = { - "aphrodite": [ - "py.typed", - "libs/*.so*", - "model_executor/layers/fused_moe/configs/*.json", - "model_executor/layers/mamba/ops/configs/selective_state_update/*.json", - "model_executor/layers/quantization/utils/configs/*.json", - "entrypoints/serve/instrumentator/static/*.js", - "entrypoints/serve/instrumentator/static/*.css", - "distributed/kv_transfer/kv_connector/v1/hf3fs/utils/*.cpp", - "third_party/flash_linear_attention/LICENSE", - # DeepGEMM JIT include headers (vendored via cmake) - "third_party/deep_gemm/include/**/*.cuh", - "third_party/deep_gemm/include/**/*.h", - "third_party/deep_gemm/include/**/*.hpp", - # fmha_sm100 sparse CuTe-DSL helper kernels (vendored via cmake) - "third_party/fmha_sm100/csrc/**/*.cu", - "third_party/fmha_sm100/csrc/**/*.h", - "third_party/fmha_sm100/csrc/**/*.jinja", - "third_party/fmha_sm100/csrc/**/*.cu.jinja", - "third_party/fmha_sm100/cute/**/*.cu", - "third_party/fmha_sm100/cutlass/include/**/*.h", - "third_party/fmha_sm100/cutlass/include/**/*.hpp", - "third_party/fmha_sm100/cutlass/tools/util/include/**/*.h", - "third_party/fmha_sm100/cutlass/tools/util/include/**/*.hpp", - # tml-fa4 CuTe-DSL helper kernels (vendored via cmake) - "third_party/tml_fa4/**/*.py", - ] -} - - -def add_aphrodite_package_data(filename: str) -> None: - aphrodite_files = package_data.setdefault("aphrodite", []) - if filename not in aphrodite_files: - aphrodite_files.append(filename) - - -# If the rust frontend binary is already present in the source tree (e.g., -# pre-built in a separate Docker build stage), ship it as-is. -if PRECOMPILED_RUST_FRONTEND_PATH.exists(): - add_aphrodite_package_data("aphrodite-rs") -for rust_extension_path in get_precompiled_rust_extension_paths(): - add_aphrodite_package_data(rust_extension_path.name) - -if _no_device(): - ext_modules = [] - -if APHRODITE_USE_PRECOMPILED: - wheel = precompiled_wheel.determine_wheel( - ROOT_DIR, - location=envs.APHRODITE_PRECOMPILED_WHEEL_LOCATION, - commit=envs.APHRODITE_PRECOMPILED_WHEEL_COMMIT, - ) - package_data_patch = precompiled_wheel.extract_precompiled_wheel( - ROOT_DIR, - wheel, - extract_extensions=True, - extract_rust=True, - ) - for package, files in package_data_patch.items(): - package_data.setdefault(package, []).extend(files) - ext_modules = [] - -if not ext_modules: - cmdclass = {} -else: - cmdclass = {"build_ext": cmake_build_ext} - -# Rust artifacts, built via setuptools-rust and installed into the package -# directory alongside the Python modules. Imported lazily: setuptools-rust does -# not need to be installed when nothing is being built. -if APHRODITE_USE_PRECOMPILED or should_use_precompiled_rust(): - rust_extensions = [] -else: - rust_build = load_module_from_path("rust_build", os.path.join(ROOT_DIR, "tools", "build_rust.py")) - rust_extensions = rust_build.rust_extensions(optional=not should_require_rust_frontend()) - -setup( - # static metadata should rather go in pyproject.toml - version=get_aphrodite_version(), - ext_modules=ext_modules, - rust_extensions=rust_extensions, - install_requires=get_requirements(), - extras_require={ - # AMD Zen CPU optimizations via zentorch - "zen": ["zentorch==2.11.0.0"], - "bench": ["pandas", "matplotlib", "seaborn", "datasets", "scipy", "plotly"], - "tensorizer": ["tensorizer==2.10.1"], - "fastsafetensors": ["fastsafetensors >= 0.3.2"], - "instanttensor": ["instanttensor >= 0.1.9"], - "runai": ["runai-model-streamer[s3,gcs,azure] >= 0.15.7"], - "audio": [ - "av", - "scipy", - "soundfile", - "soxr", - "mistral_common[audio]", - ], # Required for audio processing - "video": [], # Kept for backwards compatibility - # NVIDIA DeepStream (NVDEC) GPU video-decode backend. Linux x86-64 - # only; also needs system GStreamer + libv4l (see docs). - "deepstream": ["nvidia-deepstream-videodecode-cu13>=9.0.2"], - "flashinfer": [], # Kept for backwards compatibility - # Optional deps for Helion kernel development - # NOTE: When updating helion version, also update CI files: - # - .buildkite/test_areas/kernels.yaml - # - .buildkite/test-amd.yaml - "helion": ["helion==1.1.0"], - # Optional deps for gRPC server (aphrodite serve --grpc) - "grpc": ["smg-grpc-servicer[aphrodite] >= 0.5.2"], - # Optional deps for OpenTelemetry tracing - "otel": [ - "opentelemetry-sdk>=1.26.0", - "opentelemetry-api>=1.26.0", - "opentelemetry-exporter-otlp>=1.26.0", - "opentelemetry-semantic-conventions-ai>=0.4.1", - ], - # extra quantization plugin - "extra-quant": ["aphrodite-gguf-plugin>=0.0.2"], - }, - cmdclass=cmdclass, - package_data=package_data, -) diff --git a/tests/kernels/core/test_activation.py.orig b/tests/kernels/core/test_activation.py.orig deleted file mode 100644 index eab3d942a6..0000000000 --- a/tests/kernels/core/test_activation.py.orig +++ /dev/null @@ -1,222 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -import random - -import pytest -import torch - -from aphrodite.model_executor.layers.activation import ( - FastGELU, - FatreluAndMul, - GeluAndMul, - MulAndSilu, - NewGELU, - QuickGELU, - ReLUSquaredActivation, - SiluAndMul, - SiluAndMulWithClamp, - SwigluOAIAndMul, - SwigluStepAndMul, - swiglustep_and_mul_triton, -) -from aphrodite.utils.torch_utils import set_random_seed -from tests.kernels.allclose_default import get_default_atol, get_default_rtol -from tests.kernels.utils import opcheck - -DTYPES = [torch.half, torch.bfloat16, torch.float] -NUM_TOKENS = [7, 83, 2048] # Arbitrary values for testing -D = [512, 13824] # Arbitrary values for testing -SEEDS = [0] -CUDA_DEVICES = [f"cuda:{i}" for i in range(1 if torch.accelerator.device_count() == 1 else 2)] - - -@pytest.mark.parametrize( - "activation", - [ - "silu_and_mul", - "mul_and_silu", - "gelu", - "gelu_tanh", - "fatrelu", - "swigluoai_and_mul", - "swiglustep_and_mul", - ], -) -@pytest.mark.parametrize("num_tokens", NUM_TOKENS) -@pytest.mark.parametrize("d", D) -@pytest.mark.parametrize("dtype", DTYPES) -@pytest.mark.parametrize("seed", SEEDS) -@pytest.mark.parametrize("device", CUDA_DEVICES) -@torch.inference_mode() -def test_act_and_mul( - default_aphrodite_config, - activation: str, - num_tokens: int, - d: int, - dtype: torch.dtype, - seed: int, - device: str, -) -> None: - set_random_seed(seed) - torch.set_default_device(device) - x = torch.randn(num_tokens, 2 * d, dtype=dtype) - if activation == "silu_and_mul": - layer = SiluAndMul(compile_native=False) - fn = torch.ops._C.silu_and_mul - if activation == "mul_and_silu": - layer = MulAndSilu() - fn = torch.ops._C.mul_and_silu - elif activation == "gelu": - layer = GeluAndMul(approximate="none") - fn = torch.ops._C.gelu_and_mul - elif activation == "gelu_tanh": - layer = GeluAndMul(approximate="tanh") - fn = torch.ops._C.gelu_tanh_and_mul - elif activation == "fatrelu": - threshold = random.uniform(0, 1) - layer = FatreluAndMul(threshold) - fn = torch.ops._C.fatrelu_and_mul - elif activation == "swigluoai_and_mul": - layer = SwigluOAIAndMul() - fn = torch.ops._C.swigluoai_and_mul - elif activation == "swiglustep_and_mul": - layer = SwigluStepAndMul() - fn = swiglustep_and_mul_triton - out = layer(x) - ref_out = layer.forward_native(x) - if activation in ["swigluoai_and_mul", "swiglustep_and_mul"]: - rtol = { - # For fp16, change the relative tolerance from 1e-3 to 2e-3 - torch.float16: 2e-3, - torch.bfloat16: 2e-2, - torch.float: 1.3e-6, - } - - def _get_rtol(output) -> float: - return rtol[output.dtype] - - torch.testing.assert_close(out, ref_out, atol=get_default_atol(out), rtol=_get_rtol(out)) - else: - # The SiluAndMul, MulAndSilu, GELU and FatReLU implementations are - # equivalent to the native PyTorch implementations, so we can do exact - # comparison. - torch.testing.assert_close(out, ref_out, atol=0.0, rtol=0.0) - - d = x.shape[-1] // 2 - output_shape = x.shape[:-1] + (d,) - out = torch.empty(output_shape, dtype=x.dtype, device=x.device) - if activation == "fatrelu": - opcheck(fn, (out, x, threshold)) - elif activation == "swigluoai_and_mul": - opcheck(fn, (out, x, layer.alpha, layer.limit)) - elif activation != "swiglustep_and_mul": - opcheck(fn, (out, x)) - - -SWIGLU_LIMITS = [3.0, 7.0, 15.0] - - -@pytest.mark.parametrize("swiglu_limit", SWIGLU_LIMITS) -@pytest.mark.parametrize("num_tokens", NUM_TOKENS) -@pytest.mark.parametrize("d", D) -@pytest.mark.parametrize("dtype", DTYPES) -@pytest.mark.parametrize("seed", SEEDS) -@pytest.mark.parametrize("device", CUDA_DEVICES) -@torch.inference_mode() -def test_silu_and_mul_with_clamp( - default_aphrodite_config, - swiglu_limit: float, - num_tokens: int, - d: int, - dtype: torch.dtype, - seed: int, - device: str, -) -> None: - """SiluAndMulWithClamp: cuda kernel must match native reference.""" - set_random_seed(seed) - torch.set_default_device(device) - # Use large values to ensure clamping is exercised. - x = torch.randn(num_tokens, 2 * d, dtype=dtype) * swiglu_limit * 2 - - layer = SiluAndMulWithClamp(swiglu_limit, compile_native=False) - out = layer(x) - ref_out = layer.forward_native(x) - - rtol = { - torch.float16: 2e-3, - torch.bfloat16: 2e-2, - torch.float: 1.3e-6, - } - torch.testing.assert_close(out, ref_out, atol=get_default_atol(out), rtol=rtol[out.dtype]) - - # Verify clamping is actually being applied: the clamped output should - # differ from the unclamped SiluAndMul output when inputs are large. - unclamped_out = SiluAndMul.forward_native(x) - assert not torch.equal(ref_out.float(), unclamped_out.float()), ( - "Input was not large enough to exercise the clamp; increase scale" - ) - - # Verify gate clamping semantics with a controlled scalar case. - # gate=large_val is clamped to limit first, then silu(limit) * 1.0. - x_gate = torch.tensor([[swiglu_limit * 20.0, 1.0]], dtype=torch.float32, device=device) - out_gate = SiluAndMulWithClamp(swiglu_limit, compile_native=False)(x_gate) - expected_gate = torch.nn.functional.silu(torch.tensor(swiglu_limit, dtype=torch.float32)).item() - torch.testing.assert_close( - out_gate, - torch.tensor([[expected_gate]], dtype=torch.float32, device=device), - atol=1e-3, - rtol=1e-3, - ) - - # Verify up clamping semantics: up >> limit gets clamped to limit. - x_up = torch.tensor([[1.0, swiglu_limit * 20.0]], dtype=torch.float32, device=device) - out_up = SiluAndMulWithClamp(swiglu_limit, compile_native=False)(x_up) - silu_1 = torch.nn.functional.silu(torch.tensor(1.0)).item() - torch.testing.assert_close( - out_up, - torch.tensor([[silu_1 * swiglu_limit]], dtype=torch.float32, device=device), - atol=1e-3, - rtol=1e-3, - ) - - # opcheck - out_buf = torch.empty(x.shape[:-1] + (d,), dtype=dtype, device=device) - opcheck(torch.ops._C.silu_and_mul_with_clamp, (out_buf, x, swiglu_limit)) - - -@pytest.mark.parametrize( - "activation", - [ - (FastGELU, torch.ops._C.gelu_fast), - (NewGELU, torch.ops._C.gelu_new), - (QuickGELU, torch.ops._C.gelu_quick), - (ReLUSquaredActivation, torch.ops._C.relu_squared), - ], -) -@pytest.mark.parametrize("num_tokens", NUM_TOKENS) -@pytest.mark.parametrize("d", D) -@pytest.mark.parametrize("dtype", DTYPES) -@pytest.mark.parametrize("seed", SEEDS) -@pytest.mark.parametrize("device", CUDA_DEVICES) -@torch.inference_mode() -def test_activation( - default_aphrodite_config, - activation: type[torch.nn.Module], - num_tokens: int, - d: int, - dtype: torch.dtype, - seed: int, - device: str, -) -> None: - set_random_seed(seed) - torch.set_default_device(device) - x = torch.randn(num_tokens, d, dtype=dtype) - layer = activation[0]() - fn = activation[1] - out = layer(x) - ref_out = layer.forward_native(x) - torch.testing.assert_close(out, ref_out, atol=get_default_atol(out), rtol=get_default_rtol(out)) - - out = torch.empty_like(x) - opcheck(fn, (out, x)) diff --git a/tests/kernels/moe/test_grouped_topk.py.orig b/tests/kernels/moe/test_grouped_topk.py.orig deleted file mode 100644 index 72df9a933f..0000000000 --- a/tests/kernels/moe/test_grouped_topk.py.orig +++ /dev/null @@ -1,316 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Tests for the MoE grouped topk kernel - -Run `pytest tests/kernels/moe/test_grouped_topk.py`. -""" - -import pytest -import torch - -import aphrodite.envs as envs -from aphrodite.config import ( - AphroditeConfig, - CompilationConfig, - get_cached_compilation_config, - set_current_aphrodite_config, -) -from aphrodite.model_executor.layers.fused_moe.router.grouped_topk_router import ( - GroupedTopk, - fused_grouped_topk, -) -from aphrodite.platforms import current_platform -from aphrodite.utils.torch_utils import set_random_seed - - -@pytest.mark.skipif(not current_platform.is_cuda(), reason="This test is skipped on non-CUDA platform.") -@pytest.mark.parametrize("n_token", [1, 33, 64]) -@pytest.mark.parametrize("n_hidden", [1024, 2048]) -@pytest.mark.parametrize( - "n_expert,topk,num_expert_group,topk_group", - [ - (16, 2, 8, 2), - (128, 2, 8, 2), - (256, 8, 8, 4), - (384, 8, 1, 1), - (512, 22, 1, 1), - ], -) -@pytest.mark.parametrize("renormalize", [True, False]) -@pytest.mark.parametrize("scoring_func", ["softmax", "sigmoid"]) -@pytest.mark.parametrize("routed_scaling_factor", [1.0, 2.5]) -@pytest.mark.parametrize("input_dtype", [torch.bfloat16, torch.float32]) -@pytest.mark.parametrize("bias_dtype", [torch.float32]) -def test_grouped_topk( - monkeypatch: pytest.MonkeyPatch, - n_token: int, - n_hidden: int, - n_expert: int, - topk: int, - num_expert_group: int, - topk_group: int, - renormalize: bool, - scoring_func: str, - routed_scaling_factor: float, - input_dtype: torch.dtype, - bias_dtype: torch.dtype, -): - aphrodite_config = AphroditeConfig(compilation_config=CompilationConfig(custom_ops=["all", "+grouped_topk"])) - get_cached_compilation_config.cache_clear() - - set_random_seed(0) - hidden_states = torch.randn((n_token, n_hidden), dtype=input_dtype, device="cuda") - gating_output = torch.randn((n_token, n_expert), dtype=input_dtype, device="cuda") - e_score_correction_bias = torch.randn((n_expert,), dtype=bias_dtype, device="cuda") - - with set_current_aphrodite_config(aphrodite_config), monkeypatch.context() as m: - m.setenv("APHRODITE_USE_FUSED_MOE_GROUPED_TOPK", "0") - m.setattr(envs, "APHRODITE_BATCH_INVARIANT", True) - grouped_topk = GroupedTopk( - topk=topk, - renormalize=renormalize, - num_expert_group=num_expert_group, - topk_group=topk_group, - scoring_func=scoring_func, - routed_scaling_factor=routed_scaling_factor, - ) - assert grouped_topk._forward_method.__name__ == "forward_cuda" - baseline_topk_weights, baseline_topk_ids = grouped_topk( - hidden_states=hidden_states, - gating_output=gating_output, - e_score_correction_bias=e_score_correction_bias, - ) - - test_topk_weights, test_topk_ids = fused_grouped_topk( - hidden_states=hidden_states, - gating_output=gating_output, - topk=topk, - renormalize=renormalize, - num_expert_group=num_expert_group, - topk_group=topk_group, - scoring_func=scoring_func, - routed_scaling_factor=routed_scaling_factor, - e_score_correction_bias=e_score_correction_bias, - ) - - 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) - - -@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_align_block_size.py.orig b/tests/kernels/moe/test_moe_align_block_size.py.orig deleted file mode 100644 index 248a799bdb..0000000000 --- a/tests/kernels/moe/test_moe_align_block_size.py.orig +++ /dev/null @@ -1,366 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Tests for the MOE align block size function. - -Run `pytest tests/kernels/moe/test_moe_align_block_size.py`. -""" - -import pytest -import torch - -from aphrodite.model_executor.layers.fused_moe.moe_align_block_size import ( - batched_moe_align_block_size, - moe_align_block_size, -) -from aphrodite.utils.math_utils import cdiv, round_up -from aphrodite.utils.torch_utils import set_random_seed - -NUM_TOKENS = [1, 3, 256, 2256, 4096] -NUM_EXPERTS = [32, 160, 256, 257] -TOP_KS = [1, 2, 16, 32] -BLOCK_SIZES = [32, 128] -set_random_seed(0) - - -def _group_tokens_by_expert( - sorted_ids: torch.Tensor, - expert_ids: torch.Tensor, - block_size: int, - valid_length: int, - total_tokens: int, -) -> dict: - num_blocks = valid_length // block_size - expert_tokens: dict[int, list[int]] = {} - - for block_idx in range(num_blocks): - expert_id = expert_ids[block_idx].item() - block_start = block_idx * block_size - block_end = min(block_start + block_size, valid_length) - - block_tokens = sorted_ids[block_start:block_end] - valid_tokens = block_tokens[block_tokens < total_tokens] - - if expert_id not in expert_tokens: - expert_tokens[expert_id] = [] - expert_tokens[expert_id].extend(valid_tokens.tolist()) - return expert_tokens - - -def _verify_expert_level_sorting( - actual_sorted_ids: torch.Tensor, - golden_sorted_ids: torch.Tensor, - expert_ids: torch.Tensor, - block_size: int, - valid_length: int, - total_tokens: int, -): - """ - Verify that actual_sorted_ids follows the correct expert-level sorting. - The kerne limplementation may or may not preserve original token order - in topk_ids in the final sorted_ids however this does not impact quality. - """ - # Group tokens by expert from the golden implementation - golden_expert_tokens = _group_tokens_by_expert( - golden_sorted_ids, expert_ids, block_size, valid_length, total_tokens - ) - - actual_expert_tokens = _group_tokens_by_expert( - actual_sorted_ids, expert_ids, block_size, valid_length, total_tokens - ) - - assert set(golden_expert_tokens.keys()) == set(actual_expert_tokens.keys()), ( - f"Expert IDs mismatch: golden={set(golden_expert_tokens.keys())}, actual={set(actual_expert_tokens.keys())}" - ) - - for expert_id in golden_expert_tokens: - golden_tokens = torch.tensor(golden_expert_tokens[expert_id], device=actual_sorted_ids.device) - actual_tokens = torch.tensor(actual_expert_tokens[expert_id], device=actual_sorted_ids.device) - assert torch.equal(torch.sort(golden_tokens)[0], torch.sort(actual_tokens)[0]), ( - f"Expert {expert_id} token mismatch: " - f"golden={golden_expert_tokens[expert_id]}, " - f"actual={actual_expert_tokens[expert_id]}" - ) - - -def torch_moe_align_block_size( - topk_ids: torch.Tensor, - block_size: int, - num_experts: int, - expert_map: torch.Tensor | None = None, - pad_sorted_ids: bool = False, -) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - """ - Golden torch implementation of moe_align_block_size. - - This function aligns the token distribution across experts to be compatible - with block size for matrix multiplication by sorting tokens by expert and - padding to block boundaries. - """ - max_num_tokens_padded = topk_ids.numel() + num_experts * (block_size - 1) - if pad_sorted_ids: - max_num_tokens_padded = round_up(max_num_tokens_padded, block_size) - if topk_ids.numel() < num_experts: - max_num_tokens_padded = topk_ids.numel() * block_size - - flattened_token_indices = torch.arange(topk_ids.numel(), device=topk_ids.device, dtype=torch.int32) - flattened_expert_ids = topk_ids.flatten() - sorted_expert_ids, sort_indices = torch.sort(flattened_expert_ids, stable=True) - sorted_token_indices = flattened_token_indices[sort_indices] - - expert_token_counts = torch.zeros(num_experts, dtype=torch.int64, device=topk_ids.device) - for expert_id in range(num_experts): - mask = sorted_expert_ids == expert_id - expert_token_counts[expert_id] = mask.sum() - - expert_padded_counts = torch.zeros(num_experts, dtype=torch.int64, device=topk_ids.device) - for expert_id in range(num_experts): - original_count = expert_token_counts[expert_id] - if expert_map is not None and expert_map[expert_id] == -1: - continue - if original_count > 0: - expert_padded_counts[expert_id] = ((original_count + block_size - 1) // block_size) * block_size - - sorted_token_ids = torch.full( - (max_num_tokens_padded,), - topk_ids.numel(), - dtype=torch.int32, - device=topk_ids.device, - ) - max_num_blocks = (max_num_tokens_padded + block_size - 1) // block_size - expert_ids = torch.full((max_num_blocks,), -1, dtype=torch.int32, device=topk_ids.device) - - current_pos = 0 - current_block = 0 - for expert_id in range(num_experts): - if expert_map is not None and expert_map[expert_id] == -1: - continue - - expert_mask = sorted_expert_ids == expert_id - expert_tokens = sorted_token_indices[expert_mask] - num_expert_tokens = expert_tokens.shape[0] - - if num_expert_tokens > 0: - sorted_token_ids[current_pos : current_pos + num_expert_tokens] = expert_tokens - - expert_blocks_needed = expert_padded_counts[expert_id] // block_size - - expert_id_new = expert_id - if expert_map is not None: - expert_id_new = expert_map[expert_id] - expert_ids[current_block : current_block + expert_blocks_needed] = expert_id_new - - current_pos += expert_padded_counts[expert_id] - current_block += expert_blocks_needed - - total_padded_tokens = expert_padded_counts.sum() - num_tokens_post_pad = torch.tensor([total_padded_tokens], dtype=torch.int32, device=topk_ids.device) - - return sorted_token_ids, expert_ids, num_tokens_post_pad - - -@pytest.mark.parametrize("m", NUM_TOKENS) -@pytest.mark.parametrize("topk", TOP_KS) -@pytest.mark.parametrize("num_experts", NUM_EXPERTS) -@pytest.mark.parametrize("block_size", BLOCK_SIZES) -@pytest.mark.parametrize("pad_sorted_ids", [False, True]) -def test_moe_align_block_size(m: int, topk: int, num_experts: int, block_size: int, pad_sorted_ids: bool): - """Test moe_align_block_size without expert mapping""" - topk_ids = torch.zeros((m, topk), device="cuda", dtype=torch.int32) - for i in range(m): - experts = torch.randperm(num_experts, device="cuda")[:topk] - topk_ids[i] = experts - - actual_sorted_ids, actual_expert_ids, actual_num_tokens = moe_align_block_size( - topk_ids=topk_ids, - block_size=block_size, - num_experts=num_experts, - pad_sorted_ids=pad_sorted_ids, - ) - golden_sorted_ids, golden_expert_ids, golden_num_tokens = torch_moe_align_block_size( - topk_ids=topk_ids, - block_size=block_size, - num_experts=num_experts, - pad_sorted_ids=pad_sorted_ids, - ) - - torch.testing.assert_close(actual_num_tokens, golden_num_tokens, atol=0, rtol=0) - torch.testing.assert_close(actual_expert_ids, golden_expert_ids, atol=0, rtol=0) - - # For sorted_token_ids, verify block-level correctness rather than exact - # order Tokens within each expert's blocks can be in any order, but expert - # regions must be correct - _verify_expert_level_sorting( - actual_sorted_ids, - golden_sorted_ids, - actual_expert_ids, - block_size, - actual_num_tokens.item(), - m * topk, - ) - - total_tokens = m * topk - assert actual_num_tokens.item() % block_size == 0, "num_tokens_post_pad should be divisible by block_size" - assert actual_num_tokens.item() >= total_tokens, "num_tokens_post_pad should be at least total_tokens" - valid_tokens = actual_sorted_ids[actual_sorted_ids < total_tokens] - assert len(valid_tokens) == total_tokens, ( - f"Should have exactly {total_tokens} valid tokens, got {len(valid_tokens)}" - ) - actual_num_blocks = cdiv(int(actual_num_tokens.item()), block_size) - assert (actual_expert_ids[:actual_num_blocks] >= 0).all() and ( - actual_expert_ids[:actual_num_blocks] < num_experts - ).all(), "expert_ids should contain valid expert indices" - - -@pytest.mark.parametrize("m", [16, 32, 2048]) -@pytest.mark.parametrize("topk", [2, 4]) -@pytest.mark.parametrize("num_experts", [8, 64]) -@pytest.mark.parametrize("block_size", [64]) -@pytest.mark.parametrize("mask_inactive_experts", [False, True]) -def test_moe_align_block_size_with_expert_map( - m: int, - topk: int, - num_experts: int, - block_size: int, - mask_inactive_experts: bool, -): - """Test moe_align_block_size with expert mapping (EP scenario)""" - expert_map = torch.full((num_experts,), -1, device="cuda", dtype=torch.int32) - local_experts = list(range(0, num_experts, 2)) - for i, expert_id in enumerate(local_experts): - expert_map[expert_id] = i - - topk_ids = torch.empty((m, topk), device="cuda", dtype=torch.int32) - for i in range(m): - 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 - - actual_sorted_ids, actual_expert_ids, actual_num_tokens = moe_align_block_size( - topk_ids=topk_ids, - block_size=block_size, - num_experts=num_experts, - expert_map=expert_map, - ignore_invalid_experts=True, - ) - golden_sorted_ids, golden_expert_ids, golden_num_tokens = torch_moe_align_block_size( - topk_ids=topk_ids, - block_size=block_size, - num_experts=num_experts, - expert_map=expert_map, - ) - - torch.testing.assert_close(actual_num_tokens, golden_num_tokens, atol=0, rtol=0) - torch.testing.assert_close(actual_expert_ids, golden_expert_ids, atol=0, rtol=0) - _verify_expert_level_sorting( - actual_sorted_ids, - golden_sorted_ids, - actual_expert_ids, - block_size, - actual_num_tokens.item(), - m * topk, - ) - - -def test_moe_align_block_size_deterministic(): - m, topk, num_experts, block_size = 128, 2, 32, 64 - - torch.manual_seed(42) - topk_ids = torch.randint(0, num_experts, (m, topk), device="cuda", dtype=torch.int32) - - # expect the results to be reproducible - results = [] - for _ in range(5): - sorted_ids, expert_ids, num_tokens = moe_align_block_size( - topk_ids=topk_ids, block_size=block_size, num_experts=num_experts - ) - results.append((sorted_ids.clone(), expert_ids.clone(), num_tokens.clone())) - - for i in range(1, len(results)): - assert torch.equal(results[0][0], results[i][0]), "sorted_ids should be deterministic" - assert torch.equal(results[0][1], results[i][1]), "expert_ids should be deterministic" - assert torch.equal(results[0][2], results[i][2]), "num_tokens should be deterministic" - - -@pytest.mark.parametrize("max_tokens_per_batch", [13, 16, 512]) -@pytest.mark.parametrize("num_experts", [8, 16, 32, 64]) -@pytest.mark.parametrize("block_size", [8, 16, 32, 64]) -@pytest.mark.parametrize("simulate_empty_batches", [False, True]) -def test_batched_moe_align_block_size( - max_tokens_per_batch: int, - num_experts: int, - block_size: int, - simulate_empty_batches: bool, -): - def ref_outputs( - expert_num_tokens: torch.Tensor, - ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - E = expert_num_tokens.size(0) - - # Round up so each batch can be split to blocks evenly. - Msum = round_up(max_tokens_per_batch, block_size) * E - ref_sorted_ids = torch.empty((Msum,), dtype=torch.int32) - ref_expert_ids = torch.empty((Msum // block_size,), dtype=torch.int32) - ref_num_tokens_post_pad = torch.empty((1,), dtype=torch.int32) - - # Initialize - sentinel = E * max_tokens_per_batch - ref_sorted_ids.fill_(sentinel) - ref_expert_ids.fill_(-1) - - # Fill ref_sorted_ids - i = 0 - for expert_id, expert_nt in enumerate(expert_num_tokens): - token_offset = expert_id * max_tokens_per_batch - for j in range(expert_nt): - ref_sorted_ids[i] = token_offset + j - i += 1 - # round up i to the next block_size - i = round_up(i, block_size) - - ref_num_tokens_post_pad[0] = i - - # Fill expert_ids - nt_ceil_sum = 0 - for expert_id, expert_nt in enumerate(expert_num_tokens): - expert_ids_offset = nt_ceil_sum // block_size - ceil_expert_nt = round_up(int(expert_nt.item()), block_size) - num_blocks = ceil_expert_nt // block_size - for x in range(num_blocks): - ref_expert_ids[expert_ids_offset + x] = expert_id - nt_ceil_sum += ceil_expert_nt - - return ( - ref_sorted_ids.to("cuda"), - ref_expert_ids.to("cuda"), - ref_num_tokens_post_pad.to("cuda"), - ) - - # Compute expert_num_tokens - expert_num_tokens = torch.randint( - low=0, - high=max_tokens_per_batch, - size=(num_experts,), - device="cpu", - dtype=torch.int32, - ) - if simulate_empty_batches: - # mark half the batches to have 0 tokens - zero_batches = torch.randperm(num_experts)[: num_experts // 2] - expert_num_tokens[zero_batches] = 0 - - # ref outputs - ref_sorted_ids, ref_expert_ids, ref_num_tokens_post_pad = ref_outputs(expert_num_tokens) - - # outputs - sorted_ids, expert_ids, num_tokens_post_pad = batched_moe_align_block_size( - max_tokens_per_batch, block_size, expert_num_tokens.to("cuda") - ) - - assert ref_sorted_ids.size() == sorted_ids.size(), f"{ref_sorted_ids.size()} vs {sorted_ids.size()}" - assert ref_expert_ids.size() == expert_ids.size(), f"{ref_expert_ids.size()} vs {expert_ids.size()}" - assert ref_num_tokens_post_pad.size() == num_tokens_post_pad.size(), ( - f"{ref_num_tokens_post_pad.size()} vs {num_tokens_post_pad.size()}" - ) - torch.testing.assert_close(ref_sorted_ids, sorted_ids, atol=0, rtol=0) - torch.testing.assert_close(ref_expert_ids, expert_ids, atol=0, rtol=0) - torch.testing.assert_close(ref_num_tokens_post_pad, num_tokens_post_pad, atol=0, rtol=0) diff --git a/tests/models/kimi_k3/test_attn_res.py.orig b/tests/models/kimi_k3/test_attn_res.py.orig deleted file mode 100644 index 092657e258..0000000000 --- a/tests/models/kimi_k3/test_attn_res.py.orig +++ /dev/null @@ -1,171 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -import pytest -import torch -import torch.nn.functional as F - -from aphrodite.models.kimi_k3.nvidia.ops import attn_res -from aphrodite.platforms import current_platform - -HIDDEN_SIZE = 7168 -MAX_BLOCKS = 8 -EPS = 1e-5 - - -def _randn_with_row_padding(*shape: int, padding: int = 0) -> torch.Tensor: - storage = torch.randn( - *shape[:-1], - shape[-1] + padding, - device="cuda", - dtype=torch.bfloat16, - ) - return storage[..., : shape[-1]] - - -def _reference( - prefix: torch.Tensor, - delta: torch.Tensor | None, - blocks: torch.Tensor, - norm_weight: torch.Tensor, - qk_weight: torch.Tensor, - output_norm_weight: torch.Tensor | None, - num_blocks: int, -) -> tuple[torch.Tensor, torch.Tensor]: - if delta is not None: - prefix = prefix + delta - values = torch.cat((blocks[:, :num_blocks], prefix.unsqueeze(1)), dim=1) - keys = F.rms_norm(values, (HIDDEN_SIZE,), norm_weight, EPS) - probs = (keys @ qk_weight).softmax(dim=-1) - output = torch.matmul(probs.unsqueeze(1), values).squeeze(1) - if output_norm_weight is not None: - output = F.rms_norm(output, (HIDDEN_SIZE,), output_norm_weight, EPS) - return output, prefix - - -@pytest.mark.parametrize( - ( - "num_tokens", - "num_blocks", - "row_padding", - "write_block", - "has_delta", - "backend", - ), - [ - pytest.param(1, 0, 0, True, False, "triton", id="triton-empty"), - pytest.param(1, 0, 0, True, True, "triton", id="triton-empty-add"), - pytest.param(17, 5, 7, True, False, "triton", id="triton-write"), - pytest.param(17, 5, 7, True, True, "triton", id="triton-write-add"), - pytest.param(3, 8, 0, False, False, "triton", id="triton-full"), - pytest.param(3, 8, 0, False, True, "triton", id="triton-full-add"), - pytest.param(320, 1, 0, False, True, "nvidia", id="nvidia-1"), - pytest.param(320, 4, 0, False, True, "nvidia", id="nvidia-4"), - pytest.param(320, 8, 0, False, True, "nvidia", id="nvidia-8"), - ], -) -def test_attn_res( - num_tokens: int, - num_blocks: int, - row_padding: int, - write_block: bool, - has_delta: bool, - backend: str, -): - if backend == "nvidia" and not current_platform.is_device_capability_family(100): - pytest.skip("NVIDIA AttnRes requires the SM100 family") - - prefix = _randn_with_row_padding(num_tokens, HIDDEN_SIZE, padding=row_padding) - delta = _randn_with_row_padding(num_tokens, HIDDEN_SIZE, padding=row_padding) if has_delta else None - blocks = _randn_with_row_padding(num_tokens, MAX_BLOCKS, HIDDEN_SIZE, padding=row_padding) - norm_weight = 1 + 0.1 * torch.randn(HIDDEN_SIZE, device="cuda", dtype=torch.bfloat16) - qk_weight = torch.randn(HIDDEN_SIZE, device="cuda", dtype=torch.bfloat16) / HIDDEN_SIZE**0.5 - output_norm_weight = 1 + 0.1 * torch.randn(HIDDEN_SIZE, device="cuda", dtype=torch.bfloat16) - original_blocks = blocks.clone() - expected, expected_prefix = _reference( - prefix.clone(), - delta, - blocks, - norm_weight, - qk_weight, - output_norm_weight, - num_blocks, - ) - block_write_idx = num_blocks if write_block else -1 - - actual = attn_res( - prefix, - delta, - blocks, - norm_weight, - qk_weight, - output_norm_weight, - num_blocks, - block_write_idx, - EPS, - EPS, - ) - - torch.testing.assert_close(actual, expected, atol=8e-2, rtol=3e-2) - torch.testing.assert_close(prefix, expected_prefix, atol=0, rtol=0) - if write_block: - original_blocks[:, block_write_idx].copy_(expected_prefix) - torch.testing.assert_close(blocks, original_blocks, atol=0, rtol=0) - assert actual.is_contiguous() - - -@pytest.mark.parametrize("num_blocks", range(MAX_BLOCKS + 1)) -def test_attn_res_block_counts(num_blocks: int): - prefix = torch.randn(1, HIDDEN_SIZE, device="cuda", dtype=torch.bfloat16) - blocks = torch.randn(1, MAX_BLOCKS, HIDDEN_SIZE, device="cuda", dtype=torch.bfloat16) - norm_weight = torch.ones(HIDDEN_SIZE, device="cuda", dtype=torch.bfloat16) - qk_weight = torch.randn(HIDDEN_SIZE, device="cuda", dtype=torch.bfloat16) / HIDDEN_SIZE**0.5 - output_norm_weight = torch.ones_like(norm_weight) - expected, _ = _reference( - prefix.clone(), - None, - blocks, - norm_weight, - qk_weight, - output_norm_weight, - num_blocks, - ) - - actual = attn_res( - prefix, - None, - blocks, - norm_weight, - qk_weight, - output_norm_weight, - num_blocks, - -1, - EPS, - EPS, - ) - - torch.testing.assert_close(actual, expected, atol=8e-2, rtol=3e-2) - - -def test_attn_res_without_output_norm(): - prefix = torch.randn(7, HIDDEN_SIZE, device="cuda", dtype=torch.bfloat16) - delta = torch.randn_like(prefix) - blocks = torch.randn(7, MAX_BLOCKS, HIDDEN_SIZE, device="cuda", dtype=torch.bfloat16) - norm_weight = torch.randn(HIDDEN_SIZE, device="cuda", dtype=torch.bfloat16) - qk_weight = torch.randn(HIDDEN_SIZE, device="cuda", dtype=torch.bfloat16) / HIDDEN_SIZE**0.5 - expected, _ = _reference(prefix.clone(), delta, blocks, norm_weight, qk_weight, None, MAX_BLOCKS) - - actual = attn_res( - prefix, - delta, - blocks, - norm_weight, - qk_weight, - None, - MAX_BLOCKS, - -1, - EPS, - 0.0, - ) - - torch.testing.assert_close(actual, expected, atol=8e-2, rtol=3e-2) diff --git a/tests/models/multimodal/generation/test_common.py.orig b/tests/models/multimodal/generation/test_common.py.orig deleted file mode 100644 index 859643ae1a..0000000000 --- a/tests/models/multimodal/generation/test_common.py.orig +++ /dev/null @@ -1,1355 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Common tests for testing .generate() functionality for single / multiple -image, embedding, and video support for different VLMs in Aphrodite. -""" - -import math -from collections import defaultdict -from pathlib import PosixPath - -import pytest -from packaging.version import Version -from transformers import ( - AutoModel, - AutoModelForCausalLM, - AutoModelForImageTextToText, - AutoModelForTextToWaveform, -) -from transformers import __version__ as TRANSFORMERS_VERSION - -from aphrodite.platforms import current_platform -from aphrodite.utils.func_utils import identity - -from ....conftest import ( - IMAGE_ASSETS, - AphroditeRunner, - AudioTestAssets, - HfRunner, - ImageTestAssets, - VideoTestAssets, -) -from ....utils import create_new_process_for_each_test, large_gpu_mark, multi_gpu_marks -from ...utils import check_outputs_equal -from .vlm_utils import custom_inputs, model_utils, runners -from .vlm_utils.case_filtering import get_parametrized_options -from .vlm_utils.types import ( - CustomTestOptions, - ExpandableVLMTestArgs, - VLMTestInfo, - VLMTestType, -) - -COMMON_BROADCAST_SETTINGS = { - "test_type": VLMTestType.IMAGE, - "dtype": "half", - "max_tokens": 5, - "tensor_parallel_size": 2, - "hf_model_kwargs": {"device_map": "auto"}, - "image_size_factors": [(0.25, 0.5, 1.0)], - "distributed_executor_backend": ( - "ray", - "mp", - ), -} - -### Test configuration for specific models -# NOTE: The convention of the test settings below is to lead each test key -# with the name of the model arch used in the test, using underscores in place -# of hyphens; this makes it more convenient to filter tests for a specific kind -# of model. For example.... -# -# To run all test types for a specific key: -# use the k flag to substring match with a leading square bracket; if the -# model arch happens to be a substring of another one, you can add a -# trailing hyphen. E.g., -# - pytest $TEST_FILE -k "[llava-" -# prevents matching on "[llava_next-" & will match just the enabled cases -# for llava, i.e., single image, image embedding, and custom input tests. -# -# To run a test for a Test Info for just one of multiple models: -# use the k flag to substring match the model name, e.g., -# - pytest $TEST_FILE -k OpenGVLab/InternVL2-1B -# prevents matching on nGVLab/InternVL2-2B. -# -# You can also combine substrings to match more granularly. -# ex 1: -# pytest $TEST_FILE -k "test_single_image and OpenGVLab/InternVL2-1B" -# will run only test_single_image* for OpenGVLab/InternVL2-1B; this would -# match both wrappers for single image tests, since it also matches -# test_single_image_heavy (which forks if we have a distributed backend) -# ex 2: -# pytest $TEST_FILE -k "[llava- or [intern_vl-" -# will run all of the tests for only llava & internvl. -# -# NOTE you can add --collect-only to any of the above commands to see -# which cases would be selected and deselected by pytest. In general, -# this is a good idea for checking your command first, since tests are slow. - - -def _granite4_vision_aphrodite_to_hf_output(aphrodite_output, model): - """Post-processor for granite4_vision Aphrodite output. - - Self-contained to avoid calling AutoConfig/AutoTokenizer without - trust_remote_code (needed while the model is not in upstream HF). - """ - output_ids, output_str, out_logprobs = aphrodite_output - mm_token_id = 100352 - hf_output_ids = [ - token_id - for idx, token_id in enumerate(output_ids) - if token_id != mm_token_id or idx == 0 or output_ids[idx - 1] != mm_token_id - ] - hf_output_str = output_str[1:] if output_str and output_str[0] == " " else output_str - eos_token_id = 100257 - if hf_output_ids and hf_output_ids[-1] == eos_token_id: - hf_output_str = hf_output_str + "<|end_of_text|>" - return hf_output_ids, hf_output_str, out_logprobs - - -VLM_TEST_SETTINGS = { - #### Core tests to always run in the CI - "llava": VLMTestInfo( - models=["llava-hf/llava-1.5-7b-hf"], - test_type=(VLMTestType.EMBEDDING, VLMTestType.IMAGE, VLMTestType.CUSTOM_INPUTS), - prompt_formatter=lambda img_prompt: f"USER: {img_prompt}\nASSISTANT:", - convert_assets_to_embeddings=model_utils.get_llava_embeddings, - max_model_len=4096, - auto_cls=AutoModelForImageTextToText, - aphrodite_output_post_proc=model_utils.llava_image_aphrodite_to_hf_output, - custom_test_opts=[ - CustomTestOptions( - inputs=custom_inputs.multi_image_multi_aspect_ratio_inputs( - formatter=lambda img_prompt: f"USER: {img_prompt}\nASSISTANT:" - ), - limit_mm_per_prompt={"image": 4}, - ) - ], - aphrodite_runner_kwargs={"enable_mm_embeds": True}, - marks=[pytest.mark.core_model, pytest.mark.cpu_model], - ), - "paligemma": VLMTestInfo( - models=["google/paligemma-3b-mix-224"], - test_type=VLMTestType.IMAGE, - prompt_formatter=identity, - img_idx_to_prompt=lambda idx: "", - # Paligemma uses its own sample prompts because the default one fails - single_image_prompts=IMAGE_ASSETS.prompts( - { - "stop_sign": "caption es", - "cherry_blossom": "What is in the picture?", - } - ), - auto_cls=AutoModelForImageTextToText, - aphrodite_output_post_proc=model_utils.paligemma_aphrodite_to_hf_output, - ), - "qwen2_5_vl": VLMTestInfo( - models=["Qwen/Qwen2.5-VL-3B-Instruct"], - test_type=(VLMTestType.IMAGE, VLMTestType.MULTI_IMAGE, VLMTestType.VIDEO), - prompt_formatter=lambda img_prompt: f"<|im_start|>User\n{img_prompt}<|im_end|>\n<|im_start|>assistant\n", # noqa: E501 - img_idx_to_prompt=lambda idx: "<|vision_start|><|image_pad|><|vision_end|>", - video_idx_to_prompt=lambda idx: "<|vision_start|><|video_pad|><|vision_end|>", - enforce_eager=False, - max_model_len=4096, - max_num_seqs=2, - auto_cls=AutoModelForImageTextToText, - aphrodite_output_post_proc=model_utils.qwen2_aphrodite_to_hf_output, - image_size_factors=[(0.25,), (0.25, 0.25, 0.25), (0.25, 0.2, 0.15)], - marks=[pytest.mark.core_model, pytest.mark.cpu_model], - ), - "qwen2_5_omni": VLMTestInfo( - models=["Qwen/Qwen2.5-Omni-3B"], - test_type=(VLMTestType.IMAGE, VLMTestType.MULTI_IMAGE, VLMTestType.VIDEO), - prompt_formatter=lambda img_prompt: f"<|im_start|>User\n{img_prompt}<|im_end|>\n<|im_start|>assistant\n", # noqa: E501 - img_idx_to_prompt=lambda idx: "<|vision_bos|><|IMAGE|><|vision_eos|>", - video_idx_to_prompt=lambda idx: "<|vision_bos|><|VIDEO|><|vision_eos|>", - max_model_len=4096, - max_num_seqs=2, - num_logprobs=6 if current_platform.is_cpu() else 5, - auto_cls=AutoModelForTextToWaveform, - aphrodite_output_post_proc=model_utils.qwen2_aphrodite_to_hf_output, - patch_hf_runner=model_utils.qwen2_5_omni_patch_hf_runner, - image_size_factors=[(0.25,), (0.25, 0.25, 0.25), (0.25, 0.2, 0.15)], - marks=[pytest.mark.core_model, pytest.mark.cpu_model], - ), - "qwen3_vl": VLMTestInfo( - models=["Qwen/Qwen3-VL-4B-Instruct"], - test_type=( - VLMTestType.IMAGE, - VLMTestType.MULTI_IMAGE, - VLMTestType.VIDEO, - ), - enforce_eager=False, - needs_video_metadata=True, - prompt_formatter=lambda img_prompt: f"<|im_start|>User\n{img_prompt}<|im_end|>\n<|im_start|>assistant\n", # noqa: E501 - img_idx_to_prompt=lambda idx: "<|vision_start|><|image_pad|><|vision_end|>", # noqa: E501 - video_idx_to_prompt=lambda idx: "<|vision_start|><|video_pad|><|vision_end|>", # noqa: E501 - max_model_len=4096, - max_num_seqs=2, - num_logprobs=20, - auto_cls=AutoModelForImageTextToText, - aphrodite_output_post_proc=model_utils.qwen2_aphrodite_to_hf_output, - patch_hf_runner=model_utils.qwen3_vl_patch_hf_runner, - image_size_factors=[(0.25,), (0.25, 0.25, 0.25), (0.25, 0.2, 0.15)], - marks=[ - pytest.mark.core_model, - ], - aphrodite_runner_kwargs={"attention_backend": "TRITON_ATTN"} if current_platform.is_rocm() else {}, - ), - "ultravox": VLMTestInfo( - models=["fixie-ai/ultravox-v0_5-llama-3_2-1b"], - test_type=VLMTestType.AUDIO, - prompt_formatter=lambda audio_prompt: f"<|begin_of_text|><|start_header_id|>user<|end_header_id|>\n\n{audio_prompt}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n", # noqa: E501 - audio_idx_to_prompt=lambda idx: "<|audio|>", - max_model_len=4096, - max_num_seqs=2, - auto_cls=AutoModel, - hf_output_post_proc=model_utils.ultravox_trunc_hf_output, - marks=[ - pytest.mark.core_model, - pytest.mark.cpu_model, - # TODO: Remove skip once model has been upstreamed to Transformers - pytest.mark.skip(reason="Custom model code is not compatible with Transformers v5"), - ], - ), - #### Transformers fallback to test - ## To reduce test burden, we only test batching arbitrary image size - # Dynamic image length and number of patches - "llava-onevision-transformers": VLMTestInfo( - models=["llava-hf/llava-onevision-qwen2-0.5b-ov-hf"], - test_type=VLMTestType.IMAGE, - prompt_formatter=lambda vid_prompt: f"<|im_start|>user\n{vid_prompt}<|im_end|>\n<|im_start|>assistant\n", # noqa: E501 - max_model_len=16384, - hf_model_kwargs=model_utils.llava_onevision_hf_model_kwargs("llava-hf/llava-onevision-qwen2-0.5b-ov-hf"), - auto_cls=AutoModelForImageTextToText, - aphrodite_output_post_proc=model_utils.llava_onevision_aphrodite_to_hf_output, - image_size_factors=[(0.25, 0.5, 1.0)], - aphrodite_runner_kwargs={ - "model_impl": "transformers", - "default_torch_num_threads": 1, - }, - marks=[pytest.mark.core_model], - ), - # Gemma3 has bidirectional mask on images - "gemma3-transformers": VLMTestInfo( - models=["google/gemma-3-4b-it"], - test_type=VLMTestType.IMAGE, - prompt_formatter=lambda vid_prompt: f"<'user\n{vid_prompt}\nmodel\n", # noqa: E501 - max_model_len=4096, - auto_cls=AutoModelForImageTextToText, - aphrodite_output_post_proc=model_utils.gemma3_aphrodite_to_hf_output, - image_size_factors=[(0.25, 0.5, 1.0)], - aphrodite_runner_kwargs={ - "model_impl": "transformers", - }, - marks=[ - pytest.mark.core_model, - *([large_gpu_mark(min_gb=80)] if current_platform.is_rocm() else []), - ], - ), - "idefics3-transformers": VLMTestInfo( - models=["HuggingFaceTB/SmolVLM-256M-Instruct"], - test_type=(VLMTestType.IMAGE, VLMTestType.MULTI_IMAGE), - prompt_formatter=lambda img_prompt: f"<|begin_of_text|>User:{img_prompt}\nAssistant:", # noqa: E501 - img_idx_to_prompt=lambda idx: "", - max_model_len=8192, - max_num_seqs=2, - auto_cls=AutoModelForImageTextToText, - hf_output_post_proc=model_utils.idefics3_trunc_hf_output, - image_size_factors=[(0.25, 0.5, 1.0)], - aphrodite_runner_kwargs={ - "model_impl": "transformers", - }, - marks=[pytest.mark.core_model], - ), - # Pixel values from processor are not 4D or 5D arrays - "qwen2_5_vl-transformers": VLMTestInfo( - models=["Qwen/Qwen2.5-VL-3B-Instruct"], - test_type=VLMTestType.IMAGE, - prompt_formatter=lambda img_prompt: f"<|im_start|>User\n{img_prompt}<|im_end|>\n<|im_start|>assistant\n", # noqa: E501 - img_idx_to_prompt=lambda idx: "<|vision_start|><|image_pad|><|vision_end|>", - max_model_len=4096, - max_num_seqs=2, - auto_cls=AutoModelForImageTextToText, - aphrodite_output_post_proc=model_utils.qwen2_aphrodite_to_hf_output, - image_size_factors=[(0.25, 0.2, 0.15)], - aphrodite_runner_kwargs={ - "model_impl": "transformers", - # TODO: [ROCm] Revert this once issue #30167 is resolved - **( - { - "mm_processor_kwargs": { - "min_pixels": 256 * 28 * 28, - "max_pixels": 1280 * 28 * 28, - }, - } - if current_platform.is_rocm() - else {} - ), - }, - marks=[large_gpu_mark(min_gb=80 if current_platform.is_rocm() else 32)], - ), - #### Extended model tests - "aria": VLMTestInfo( - models=["rhymes-ai/Aria"], - test_type=(VLMTestType.IMAGE, VLMTestType.MULTI_IMAGE), - prompt_formatter=lambda img_prompt: f"<|im_start|>user\n{img_prompt}<|im_end|>\n<|im_start|>assistant\n ", # noqa: E501 - img_idx_to_prompt=lambda idx: "<|img|>\n", - max_model_len=4096, - max_num_seqs=2, - auto_cls=AutoModelForImageTextToText, - single_image_prompts=IMAGE_ASSETS.prompts( - { - "stop_sign": "Please describe the image shortly.", - "cherry_blossom": "Please infer the season with reason.", - } - ), - multi_image_prompt="Describe the two images shortly.", - stop_str=["<|im_end|>"], - image_size_factors=[(0.10, 0.15)], - max_tokens=64, - marks=[ - pytest.mark.skip( - reason="Aria needs to update for latest transformers, " - "must have a vision_processor.py." - "An issue has been filed:" - "https://huggingface.co/rhymes-ai/Aria/discussions/23" - ), - large_gpu_mark(min_gb=64), - ], - ), - "blip2": VLMTestInfo( - models=["Salesforce/blip2-opt-2.7b"], - test_type=VLMTestType.IMAGE, - prompt_formatter=lambda img_prompt: f"Question: {img_prompt} Answer:", - img_idx_to_prompt=lambda idx: "", - auto_cls=AutoModelForImageTextToText, - aphrodite_output_post_proc=model_utils.blip2_aphrodite_to_hf_output, - # FIXME: https://github.com/huggingface/transformers/pull/38510 - marks=[pytest.mark.skip("Model is broken")], - ), - "chameleon": VLMTestInfo( - models=["facebook/chameleon-7b"], - test_type=VLMTestType.IMAGE, - prompt_formatter=lambda img_prompt: f"USER: {img_prompt}\nASSISTANT:", - max_model_len=4096, - max_num_seqs=2, - auto_cls=AutoModelForImageTextToText, - # For chameleon, we only compare the sequences - aphrodite_output_post_proc=lambda aphrodite_output, model: aphrodite_output[:2], - hf_output_post_proc=lambda hf_output, model: hf_output[:2], - comparator=check_outputs_equal, - max_tokens=8, - dtype="bfloat16", - ), - "cosmos3": VLMTestInfo( - models=["nvidia/Cosmos3-Nano"], - test_type=( - VLMTestType.IMAGE, - VLMTestType.MULTI_IMAGE, - VLMTestType.VIDEO, - ), - enforce_eager=False, - needs_video_metadata=True, - prompt_formatter=lambda img_prompt: f"<|im_start|>User\n{img_prompt}<|im_end|>\n<|im_start|>assistant\n", # noqa: E501 - img_idx_to_prompt=lambda idx: "<|vision_start|><|image_pad|><|vision_end|>", # noqa: E501 - video_idx_to_prompt=lambda idx: "<|vision_start|><|video_pad|><|vision_end|>", # noqa: E501 - max_model_len=4096, - max_num_seqs=2, - num_logprobs=20, - auto_cls=AutoModelForImageTextToText, - aphrodite_output_post_proc=model_utils.qwen2_aphrodite_to_hf_output, - patch_hf_runner=model_utils.qwen3_vl_patch_hf_runner, - image_size_factors=[(0.25,), (0.25, 0.25, 0.25), (0.25, 0.2, 0.15)], - marks=[ - pytest.mark.skip( - reason="Transformers has no cosmos3_omni mapping, so the HF " - "reference runner cannot load the checkpoint." - ) - ], - ), - "deepseek_vl_v2": VLMTestInfo( - models=["Isotr0py/deepseek-vl2-tiny"], # model repo using dynamic module - test_type=(VLMTestType.IMAGE, VLMTestType.MULTI_IMAGE), - prompt_formatter=lambda img_prompt: f"<|User|>: {img_prompt}\n\n<|Assistant|>: ", # noqa: E501 - max_model_len=4096, - max_num_seqs=2, - single_image_prompts=IMAGE_ASSETS.prompts( - { - "stop_sign": "\nWhat's the content in the center of the image?", - "cherry_blossom": "\nPlease infer the season with reason in details.", # noqa: E501 - } - ), - multi_image_prompt="image_1:\nimage_2:\nWhich image can we see the car and the tower?", # noqa: E501 - patch_hf_runner=model_utils.deepseekvl2_patch_hf_runner, - hf_output_post_proc=model_utils.deepseekvl2_trunc_hf_output, - stop_str=["<｜end▁of▁sentence｜>", "<｜begin▁of▁sentence｜>"], - image_size_factors=[(1.0,), (1.0, 1.0, 1.0), (0.1, 0.5, 1.0)], - ), - "gemma3": VLMTestInfo( - models=["google/gemma-3-4b-it"], - test_type=(VLMTestType.IMAGE, VLMTestType.MULTI_IMAGE), - prompt_formatter=lambda img_prompt: f"user\n{img_prompt}\nmodel\n", # noqa: E501 - single_image_prompts=IMAGE_ASSETS.prompts( - { - "stop_sign": "What's the content in the center of the image?", # noqa: E501 - "cherry_blossom": "What is the season?", - } - ), - multi_image_prompt="Describe the two images in detail.", # noqa: E501 - max_model_len=4096, - max_num_seqs=2, - auto_cls=AutoModelForImageTextToText, - aphrodite_runner_kwargs={"mm_processor_kwargs": {"do_pan_and_scan": True}}, - patch_hf_runner=model_utils.gemma3_patch_hf_runner, - ), - "gemma4": VLMTestInfo( - models=["google/gemma-4-E2B-it"], - test_type=(VLMTestType.IMAGE, VLMTestType.MULTI_IMAGE), - prompt_formatter=lambda img_prompt: f"<|turn>user\n{img_prompt}\n<|turn>model\n", # noqa: E501 - single_image_prompts=IMAGE_ASSETS.prompts( - { - "stop_sign": "<|image|>What's the content in the center of the image?", # noqa: E501 - "cherry_blossom": "<|image|>What is the season?", - } - ), - multi_image_prompt="<|image|><|image|>Describe the two images in detail.", # noqa: E501 - max_model_len=4096, - max_num_seqs=2, - auto_cls=AutoModelForImageTextToText, - aphrodite_runner_kwargs={"limit_mm_per_prompt": {"image": 4}}, - ), - "granite_vision": VLMTestInfo( - models=["ibm-granite/granite-vision-3.3-2b"], - test_type=(VLMTestType.IMAGE), - prompt_formatter=lambda img_prompt: f"<|user|>\n{img_prompt}\n<|assistant|>\n", - max_model_len=8192, - auto_cls=AutoModelForImageTextToText, - aphrodite_output_post_proc=model_utils.llava_image_aphrodite_to_hf_output, - ), - "glm4v": VLMTestInfo( - models=["zai-org/glm-4v-9b"], - test_type=VLMTestType.IMAGE, - prompt_formatter=lambda img_prompt: f"<|user|>\n{img_prompt}<|assistant|>", - single_image_prompts=IMAGE_ASSETS.prompts( - { - "stop_sign": "<|begin_of_image|><|endoftext|><|end_of_image|>What's the content in the center of the image?", # noqa: E501 - "cherry_blossom": "<|begin_of_image|><|endoftext|><|end_of_image|>What is the season?", # noqa: E501 - } - ), - max_model_len=2048, - max_num_seqs=2, - get_stop_token_ids=lambda tok: [151329, 151336, 151338], - patch_hf_runner=model_utils.glm4v_patch_hf_runner, - # The image embeddings match with HF but the outputs of the language - # decoder are only consistent up to 2 decimal places. - # So, we need to reduce the number of tokens for the test to pass. - max_tokens=8, - num_logprobs=10, - auto_cls=AutoModelForCausalLM, - marks=[ - pytest.mark.skip( - reason="The code for this model has a bug." - "Please see the issue here:" - "https://huggingface.co/zai-org/glm-4v-9b/discussions/46." - ), - large_gpu_mark(min_gb=32), - ], - ), - "glm4_1v": VLMTestInfo( - models=["zai-org/GLM-4.1V-9B-Thinking"], - test_type=(VLMTestType.IMAGE, VLMTestType.MULTI_IMAGE), - prompt_formatter=lambda img_prompt: f"[gMASK]<|user|>\n{img_prompt}<|assistant|>\n", # noqa: E501 - img_idx_to_prompt=lambda idx: "<|begin_of_image|><|image|><|end_of_image|>", - video_idx_to_prompt=lambda idx: "<|begin_of_video|><|video|><|end_of_video|>", - max_model_len=2048, - max_num_seqs=2, - get_stop_token_ids=lambda tok: [151329, 151336, 151338], - num_logprobs=10, - image_size_factors=[(0.25,), (0.25, 0.25, 0.25), (0.25, 0.2, 0.15)], - auto_cls=AutoModelForImageTextToText, - marks=[large_gpu_mark(min_gb=32)], - ), - "glm4_1v-video": VLMTestInfo( - models=["zai-org/GLM-4.1V-9B-Thinking"], - # GLM4.1V require include video metadata for input - test_type=VLMTestType.CUSTOM_INPUTS, - prompt_formatter=lambda vid_prompt: f"[gMASK]<|user|>\n{vid_prompt}<|assistant|>\n", # noqa: E501 - max_model_len=4096, - max_num_seqs=2, - auto_cls=AutoModelForImageTextToText, - patch_hf_runner=model_utils.glm4_1v_patch_hf_runner, - custom_test_opts=[ - CustomTestOptions( - inputs=custom_inputs.video_with_metadata_glm4_1v(), - limit_mm_per_prompt={"video": 1}, - ) - ], - marks=[large_gpu_mark(min_gb=32)], - ), - "glm_ocr": VLMTestInfo( - models=["zai-org/GLM-OCR"], - test_type=(VLMTestType.IMAGE, VLMTestType.MULTI_IMAGE), - prompt_formatter=lambda img_prompt: f"[gMASK]<|user|>\n{img_prompt}<|assistant|>\n", # noqa: E501 - img_idx_to_prompt=lambda idx: "<|begin_of_image|><|image|><|end_of_image|>", - video_idx_to_prompt=lambda idx: "<|begin_of_video|><|video|><|end_of_video|>", - max_model_len=2048, - max_num_seqs=2, - get_stop_token_ids=lambda tok: [151329, 151336, 151338], - num_logprobs=10, - image_size_factors=[(0.25,), (0.25, 0.25, 0.25), (0.25, 0.2, 0.15)], - auto_cls=AutoModelForImageTextToText, - marks=[ - pytest.mark.skip( - reason="This test fails on both AMD and NV" - "hardware. please see the issue:" - "https://github.com/vllm-project/vllm/issues/42016" - ), - large_gpu_mark(min_gb=32), - ], - ), - "granite4_vision": VLMTestInfo( - models=["ibm-granite/granite-vision-4.1-4b"], - test_type=(VLMTestType.IMAGE), - prompt_formatter=lambda img_prompt: f"<|user|>\n{img_prompt}\n<|assistant|>\n", - max_model_len=8192, - auto_cls=AutoModelForImageTextToText, - aphrodite_output_post_proc=_granite4_vision_aphrodite_to_hf_output, - image_size_factors=[(1.0,)], - aphrodite_runner_kwargs={ - "enable_lora": True, - "max_lora_rank": 256, - "default_mm_loras": {"image": "ibm-granite/granite-vision-4.1-4b"}, - }, - ), - "h2ovl": VLMTestInfo( - models=[ - "h2oai/h2ovl-mississippi-800m", - "h2oai/h2ovl-mississippi-2b", - ], - test_type=(VLMTestType.IMAGE, VLMTestType.MULTI_IMAGE), - prompt_formatter=lambda img_prompt: f"<|prompt|>{img_prompt}<|end|><|answer|>", - single_image_prompts=IMAGE_ASSETS.prompts( - { - "stop_sign": "\nWhat's the content in the center of the image?", - "cherry_blossom": "\nWhat is the season?", - } - ), - multi_image_prompt="Image-1: \nImage-2: \nDescribe the two images in short.", # noqa: E501 - max_model_len=8192, - use_tokenizer_eos=True, - num_logprobs=10, - patch_hf_runner=model_utils.h2ovl_patch_hf_runner, - ), - "idefics3": VLMTestInfo( - models=["HuggingFaceTB/SmolVLM-256M-Instruct"], - test_type=(VLMTestType.IMAGE, VLMTestType.MULTI_IMAGE), - prompt_formatter=lambda img_prompt: f"<|begin_of_text|>User:{img_prompt}\nAssistant:", # noqa: E501 - img_idx_to_prompt=lambda idx: "", - max_model_len=8192, - max_num_seqs=2, - auto_cls=AutoModelForImageTextToText, - hf_output_post_proc=model_utils.idefics3_trunc_hf_output, - ), - "intern_vl": VLMTestInfo( - models=[ - "OpenGVLab/InternVL2-1B", - "OpenGVLab/InternVL2-2B", - ], - test_type=(VLMTestType.IMAGE, VLMTestType.MULTI_IMAGE), - prompt_formatter=lambda img_prompt: f"<|im_start|>User\n{img_prompt}<|im_end|>\n<|im_start|>Assistant\n", # noqa: E501 - single_image_prompts=IMAGE_ASSETS.prompts( - { - "stop_sign": "\nWhat's the content in the center of the image?", - "cherry_blossom": "\nWhat is the season?", - } - ), - multi_image_prompt="Image-1: \nImage-2: \nDescribe the two images in short.", # noqa: E501 - max_model_len=4096, - use_tokenizer_eos=True, - patch_hf_runner=model_utils.internvl_patch_hf_runner, - # TODO: Remove skip once model has been upstreamed to Transformers - marks=[pytest.mark.skip(reason="Custom model code tries to access data from meta-tensor")], - ), - "intern_vl-video": VLMTestInfo( - models=[ - "OpenGVLab/InternVL3-1B", - ], - test_type=VLMTestType.VIDEO, - prompt_formatter=lambda img_prompt: f"<|im_start|>User\n{img_prompt}<|im_end|>\n<|im_start|>Assistant\n", # noqa: E501 - video_idx_to_prompt=lambda idx: "