disable k3dspark test

Signed-off-by: Isotr0py <Isotr0py@outlook.com>

.buildkite/check-torch-abi.py

@@ -13,6 +13,7 @@ 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, ...] = (
+    "_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",

.buildkite/intel_jobs/test-intel.yaml

@@ -147,6 +147,30 @@ steps:
         'cd tests &&
         pytest -v -s quantization/test_auto_round.py &&
         pytest -v -s quantization/test_online.py'
+  - label: "XPU GPQA Eval (GPT-OSS)"
+    depends_on:
+      - image-build-xpu
+    timeout_in_minutes: 60
+    device: intel_gpu
+    agent_tags:
+      label: production
+      gpu: 1+
+      mem: 24+
+    no_plugin: true
+    env:
+      REGISTRY: "public.ecr.aws/q9t5s3a7"
+      REPO: "vllm-ci-test-repo"
+      VLLM_TEST_DEVICE: "xpu"
+    source_file_dependencies:
+      - vllm/
+      - tests/evals/gpt_oss/
+      - .buildkite/intel_jobs/test-intel.yaml
+    commands:
+      - >-
+        bash .buildkite/scripts/hardware_ci/run-intel-test.sh
+        'pip install "gpt-oss[eval]==0.0.5" &&
+        cd tests &&
+        pytest -s -v evals/gpt_oss/test_gpqa_correctness.py --config-list-file=configs/models-xpu.txt'
   - label: "XPU compressed tensors FP8 test"
     depends_on:
       - image-build-xpu

.buildkite/test-amd.yaml

@@ -2289,7 +2289,7 @@ steps:
   - export VLLM_USE_RUST_FRONTEND=1
   - export VLLM_WORKER_MULTIPROC_METHOD=spawn
   - pytest -v -s benchmarks/test_serve_cli.py -k "not insecure and not (test_bench_serve and not test_bench_serve_chat)"
-  - pytest -v -s entrypoints/openai/chat_completion/test_chat_completion.py -k "not test_invalid_json_schema and not test_invalid_regex"
+  - pytest -v -s entrypoints/openai/chat_completion/test_chat_completion.py -k "not test_invalid_json_schema and not test_invalid_regex and not test_kv_transfer_prompt_token_ids_round_trip and not test_kv_transfer_prompt_token_ids_streaming"
   - pytest -v -s entrypoints/openai/chat_completion/test_chat_logit_bias_validation.py -k "not multiple"
   - pytest -v -s entrypoints/openai/completion/test_shutdown.py -k "not engine_failure and not test_abort_timeout_exits_quickly"
   - pytest -v -s entrypoints/openai/test_return_token_ids.py -k "not test_comparison"

.buildkite/test_areas/kernels.yaml

@@ -214,17 +214,6 @@ steps:
   commands:
     - pytest -v -s kernels/mamba
 
-- label: Kernels KDA Test
-  timeout_in_minutes: 25
-  device: h200_18gb
-  source_file_dependencies:
-  - vllm/third_party/flash_linear_attention/ops/kda.py
-  - vllm/third_party/flash_linear_attention/ops/chunk_delta_h.py
-  - vllm/third_party/flash_linear_attention/ops/l2norm.py
-  - tests/kernels/test_kda.py
-  commands:
-    - pytest -v -s kernels/test_kda.py
-
 - label: Kernels DeepGEMM Test (H100)
   key: kernels-deepgemm-test-h100
   timeout_in_minutes: 35

.buildkite/test_areas/lora.yaml

@@ -16,7 +16,6 @@ steps:
     amd:
       dind: false
       device: mi300_1
-      soft_fail: true
       working_dir: "/vllm-workspace/tests"
       timeout_in_minutes: 85
       source_file_dependencies:

.buildkite/test_areas/models_basic.yaml

@@ -69,9 +69,11 @@ steps:
   - vllm/models/kimi_k3/
   - csrc/libtorch_stable/kimi_k3/
   - tests/models/kimi_k3/
+  - tests/kernels/attention/test_kimi_k3_mla_fused_epilogue.py
+  - tests/kernels/test_bf16_skinny_gemm.py
   commands:
-    # The native NVIDIA AttnRes kernel requires the SM100 family.
-    - pytest -v -s models/kimi_k3
+    # The native NVIDIA Kimi K3 kernels require the SM100 family.
+    - pytest -v -s models/kimi_k3 kernels/attention/test_kimi_k3_mla_fused_epilogue.py kernels/test_bf16_skinny_gemm.py
 
 - label: Basic Models Test (Other CPU) # 5min
   key: basic-models-test-other-cpu

.buildkite/test_areas/rust_frontend.yaml

@@ -26,7 +26,7 @@ steps:
   - export VLLM_USE_RUST_FRONTEND=1
   - export VLLM_WORKER_MULTIPROC_METHOD=spawn
   - pytest -v -s benchmarks/test_serve_cli.py -k "not insecure and not (test_bench_serve and not test_bench_serve_chat)"
-  - pytest -v -s entrypoints/openai/chat_completion/test_chat_completion.py -k "not test_invalid_json_schema and not test_invalid_regex"
+  - pytest -v -s entrypoints/openai/chat_completion/test_chat_completion.py -k "not test_invalid_json_schema and not test_invalid_regex and not test_kv_transfer_prompt_token_ids_round_trip and not test_kv_transfer_prompt_token_ids_streaming"
   - pytest -v -s entrypoints/openai/chat_completion/test_chat_logit_bias_validation.py -k "not multiple"
 
   # - pytest -v -s entrypoints/openai/completion/test_prompt_validation.py -k "not prompt_embeds"

.github/workflows/run-ci-command.yml

@@ -11,7 +11,7 @@ concurrency:
 permissions:
   contents: read
   issues: write
-  pull-requests: read
+  pull-requests: write
 
 jobs:
   run-ci-command:

.github/workflows/scripts/run_ci_command.py

@@ -381,6 +381,7 @@ def create_build_payload(
         "pull_request_base_branch": pr["base"]["ref"],
         "pull_request_repository": pr["head"]["repo"]["clone_url"],
         "pull_request_labels": [label["name"] for label in pr["labels"]],
+        "ignore_pipeline_branch_filters": True,
         "env": {
             "VLLM_CI_GITHUB_COMMENT_ID": str(comment_id),
             "VLLM_CI_TRIGGERED_BY": actor,

.github/workflows/scripts/test_run_ci_command.py

@@ -280,6 +280,7 @@ class RunCiCommandTest(unittest.TestCase):
                 "pull_request_base_branch": "main",
                 "pull_request_repository": ("https://github.com/contributor/vllm.git"),
                 "pull_request_labels": ["ready", "v1"],
+                "ignore_pipeline_branch_filters": True,
                 "env": {
                     "VLLM_CI_GITHUB_COMMENT_ID": "99",
                     "VLLM_CI_TRIGGERED_BY": "reviewer",

.pre-commit-config.yaml

@@ -4,7 +4,7 @@ default_install_hook_types:
 default_stages:
   - pre-commit # Run locally
   - manual # Run in CI
-exclude: 'vllm/third_party/.*'
+exclude: 'vllm/third_party/.*|vllm/models/kimi_k3/nvidia/ops/third_party/.*|vllm/models/kimi_k3/amd/ops/third_party/.*'
 repos:
 - repo: https://github.com/astral-sh/ruff-pre-commit
   rev: v0.14.0

CMakeLists.txt

@@ -421,8 +421,11 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_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(VLLM_GPU_LANG STREQUAL "CUDA" AND
      DEFINED CMAKE_CUDA_COMPILER_VERSION AND
@@ -1079,6 +1082,23 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_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 "$<$<COMPILE_LANGUAGE:CUDA>:--use_fast_math>")
+    list(APPEND VLLM_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}")
@@ -1138,6 +1158,10 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
       target_compile_definitions(_C_stable_libtorch PRIVATE
         VLLM_ENABLE_COOPERATIVE_TOPK=1)
     endif()
+    if(FUSED_KDA_DECODE_ARCHS)
+      target_compile_definitions(_C_stable_libtorch PRIVATE
+        VLLM_ENABLE_FUSED_KDA_DECODE=1)
+    endif()
     if(KIMI_K3_ATTN_RES_ARCHS)
       target_compile_definitions(_C_stable_libtorch PRIVATE
         VLLM_ENABLE_KIMI_K3_ATTN_RES=1)
@@ -1439,6 +1463,7 @@ if (VLLM_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)
 

benchmarks/kernels/benchmark_k3_cutedsl_residual.py

@@ -0,0 +1,367 @@
+# 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 vllm.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=(
+            "--enable-tvm-ffi --keep-cubin "
+            f"--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()

benchmarks/kernels/benchmark_kimi_k3_latent_moe_tail.py

@@ -0,0 +1,806 @@
+# 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 vllm.distributed import get_tp_group
+from vllm.distributed.parallel_state import (
+    init_distributed_environment,
+    initialize_model_parallel,
+    set_custom_all_reduce,
+)
+from vllm.model_executor.warmup.cutedsl_warmup import cutedsl_warmup
+from vllm.models.kimi_k3.nvidia.ops import latent_moe_tail
+from vllm.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 vllm.config import VllmConfig, set_current_vllm_config
+
+    with set_current_vllm_config(VllmConfig()):
+        benchmark_whole_tail(args)
+
+
+if __name__ == "__main__":
+    main()

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 vllm._custom_ops as ops
+from vllm.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()

cmake/external_projects/deepgemm.cmake

@@ -28,9 +28,9 @@ if(DEEPGEMM_SRC_DIR)
   message(STATUS "DeepGEMM using local DEEPGEMM_SRC_DIR: ${deepgemm_SOURCE_DIR}")
 else()
   # Keep in sync with tools/install_deepgemm.sh
-  set(_DEEPGEMM_UPSTREAM_REPO "https://github.com/deepseek-ai/DeepGEMM.git")
+  set(_DEEPGEMM_UPSTREAM_REPO "https://github.com/vllm-project/DeepGEMM.git")
   # NOTE: This is currently targeting nv-dev branch due to sm120 support
-  set(_DEEPGEMM_UPSTREAM_TAG "a6b593d2826719dcf4892609af7b84ee23aaf32a")
+  set(_DEEPGEMM_UPSTREAM_TAG "f5a76426fa084087169693fd0cd815223576d6e9")
 
   set(_deepgemm_fc_root "${FETCHCONTENT_BASE_DIR}")
   if(NOT _deepgemm_fc_root)

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 ${VLLM_GPU_LANG}
+    SOURCES ${FLASH_KDA_SOURCES}
+    COMPILE_FLAGS ${VLLM_GPU_FLAGS}
+    ARCHITECTURES ${VLLM_GPU_ARCHES}
+    INCLUDE_DIRECTORIES ${FLASH_KDA_INCLUDES}
+    USE_SABI 3
+    WITH_SOABI)
+
+  target_compile_options(_flashkda_C PRIVATE
+    $<$<COMPILE_LANGUAGE:CUDA>:-UPy_LIMITED_API --expt-relaxed-constexpr --expt-extended-lambda --use_fast_math -O3>
+    $<$<COMPILE_LANGUAGE:CXX>:-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()

csrc/cpu/cpu_attn.cpp

@@ -30,7 +30,8 @@ torch::Tensor get_scheduler_metadata(
     const torch::Tensor& query_start_loc, const bool causal,
     const int64_t window_size, const std::string& isa_hint,
     const bool enable_kv_split,
-    const std::optional<torch::Tensor>& dynamic_causal) {
+    const std::optional<torch::Tensor>& dynamic_causal,
+    const std::string& kv_cache_dtype) {
   cpu_attention::ISA isa;
   if (isa_hint == "amx") {
     isa = cpu_attention::ISA::AMX;
@@ -65,9 +66,11 @@ torch::Tensor get_scheduler_metadata(
   input.dynamic_causal =
       dynamic_causal.has_value() ? dynamic_causal->data_ptr<bool>() : nullptr;
 
+  const int64_t kv_cache_idx =
+      static_cast<int64_t>(parse_fp8_kv_dtype(kv_cache_dtype));
   VLLM_DISPATCH_FLOATING_TYPES(dtype, "get_scheduler_metadata", [&]() {
-    CPU_ATTN_DISPATCH(head_dim, isa, 0, [&]() {
-      input.elem_size = sizeof(scalar_t);
+    CPU_ATTN_DISPATCH(head_dim, isa, kv_cache_idx, [&]() {
+      input.elem_size = sizeof(attn_impl::kv_cache_t);
       input.q_buffer_elem_size = sizeof(attn_impl::q_buffer_t);
       input.logits_buffer_elem_size = sizeof(attn_impl::logits_buffer_t);
       input.output_buffer_elem_size =

csrc/cpu/cpu_types_vxe.hpp

@@ -269,7 +269,7 @@ struct FP32Vec4 : public Vec<FP32Vec4> {
 
   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<FP32Vec8> {
@@ -298,7 +298,7 @@ struct FP32Vec8 : public Vec<FP32Vec8> {
 
   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];
   }
@@ -643,7 +643,7 @@ struct FP32Vec16 : public Vec<FP32Vec16> {
 
   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];

csrc/cpu/torch_bindings.cpp

@@ -163,7 +163,8 @@ torch::Tensor get_scheduler_metadata(
     const torch::Tensor& query_start_loc, const bool casual,
     const int64_t window_size, const std::string& isa_hint,
     const bool enable_kv_split,
-    const std::optional<torch::Tensor>& dynamic_causal);
+    const std::optional<torch::Tensor>& dynamic_causal,
+    const std::string& kv_cache_dtype);
 
 void cpu_attn_reshape_and_cache(const torch::Tensor& key,
                                 const torch::Tensor& value,
@@ -577,7 +578,8 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
       "get_scheduler_metadata(int num_req, int num_heads_q, int num_heads_kv, "
       "int head_dim, Tensor seq_lens, ScalarType dtype, Tensor "
       "query_start_loc, bool casual, int window_size, str isa_hint, bool "
-      "enable_kv_split, Tensor? dynamic_causal) -> Tensor",
+      "enable_kv_split, Tensor? dynamic_causal, "
+      "str kv_cache_dtype=\"auto\") -> Tensor",
       &get_scheduler_metadata);
   ops.def(
       "cpu_attn_reshape_and_cache(Tensor key, Tensor value, Tensor(a2!) "

csrc/custom_all_gather_reduce_scatter.cuh

@@ -0,0 +1,327 @@
+#pragma once
+
+#include "custom_collective_common.cuh"
+
+namespace vllm {
+
+constexpr int kMnnvlLamportAgThreads = 128;
+constexpr int kMnnvlLamportRsThreads = 256;
+constexpr int kMnnvlLamportConcurrentPollMaxPacks = 8192;
+
+using CopyPack = array_t<uint64_t, 2>;
+
+template <int ngpus>
+__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<ngpus>(sg, self_sg, rank);
+#pragma unroll
+  for (int src_rank = 0; src_rank < ngpus; ++src_rank) {
+    auto src = reinterpret_cast<const CopyPack*>(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<ngpus, true>(sg, self_sg, rank);
+}
+
+template <typename T, int ngpus>
+__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<T>::P;
+  using A = typename packed_t<T>::A;
+  auto dp = *_dp;
+  auto offset = rank * size_per_rank;
+  barrier_at_start<ngpus>(sg, self_sg, rank);
+  for (int idx = blockIdx.x * blockDim.x + threadIdx.x; idx < size_per_rank;
+       idx += gridDim.x * blockDim.x) {
+    reinterpret_cast<P*>(result)[idx] =
+        packed_reduce<P, ngpus, A>((const P**)&dp.ptrs[0], offset + idx);
+  }
+  barrier_at_end<ngpus, true>(sg, self_sg, rank);
+}
+
+template <typename P>
+union LamportPack {
+  P packed;
+  uint32_t words[sizeof(P) / sizeof(uint32_t)];
+};
+
+template <typename P>
+DINLINE LamportPack<P> load_lamport_pack(const P* ptr) {
+  static_assert(sizeof(P) == 16);
+  LamportPack<P> 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<const volatile uint32_t*>(ptr);
+  #pragma unroll
+  for (int i = 0; i < sizeof(P) / sizeof(uint32_t); ++i) {
+    value.words[i] = src[i];
+  }
+#endif
+  return value;
+}
+
+template <typename P>
+DINLINE bool is_lamport_dirty(const LamportPack<P>& value) {
+#pragma unroll
+  for (int i = 0; i < sizeof(P) / sizeof(uint32_t); ++i) {
+    if (value.words[i] == 0x80000000U) return true;
+  }
+  return false;
+}
+
+template <typename P>
+DINLINE P lamport_sentinel() {
+  LamportPack<P> value;
+#pragma unroll
+  for (int i = 0; i < sizeof(P) / sizeof(uint32_t); ++i) {
+    value.words[i] = 0x80000000U;
+  }
+  return value.packed;
+}
+
+template <typename P>
+DINLINE P sanitize_lamport_payload(P packed) {
+  LamportPack<P> 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 <typename P>
+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 <typename P, int ngpus>
+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 <typename P, typename A, int ngpus>
+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<P>(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 <typename T, int ngpus>
+__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<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<P*>(const_cast<void*>(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<P*>(multicast_buffer) + current_stage * stage_size;
+  auto packed_input = reinterpret_cast<const P*>(input);
+  auto packed_result = reinterpret_cast<P*>(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<P>();
+  }
+
+  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<volatile uint32_t*>(&epochs[1]) < gridDim.x);
+    epochs[2 + current_stage] = total_size;
+    epochs[0] = epoch + 1;
+    epochs[1] = 0;
+  }
+}
+
+template <typename T, int ngpus>
+__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<T>::P;
+  using A = typename packed_t<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<P*>(const_cast<void*>(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<const P*>(input);
+
+  if (idx < size_per_rank && dst_rank != rank) {
+    auto dst = reinterpret_cast<P*>(const_cast<void*>(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<P>();
+  }
+
+  if (idx < size_per_rank && dst_rank == rank) {
+    if constexpr (ngpus == 4) {
+      if (size_per_rank > kMnnvlLamportConcurrentPollMaxPacks) {
+        reinterpret_cast<P*>(result)[idx] =
+            reduce_lamport_payloads<P, A, ngpus>(current_local, packed_input,
+                                                 rank, size_per_rank, idx);
+      } else {
+        P values[ngpus];
+        wait_lamport_payloads<P, ngpus>(
+            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<P*>(result)[idx] =
+            sanitize_lamport_payload(downcast<P>(tmp));
+      }
+    } else {
+      reinterpret_cast<P*>(result)[idx] = reduce_lamport_payloads<P, A, ngpus>(
+          current_local, packed_input, rank, size_per_rank, idx);
+    }
+  }
+
+  if (tid == 0) {
+    while (*reinterpret_cast<volatile uint32_t*>(&epochs[1]) < gridDim.x);
+    epochs[2 + current_stage] = size_per_rank * ngpus;
+    epochs[0] = epoch + 1;
+    epochs[1] = 0;
+  }
+}
+
+}  // namespace vllm

csrc/custom_all_reduce.cuh

@@ -1,299 +1,8 @@
 #pragma once
 
-#include <cuda.h>
-#include <cuda_bf16.h>
-#include <cuda_fp16.h>
-#include <cuda_runtime.h>
-
-#if defined(USE_ROCM)
-typedef __hip_bfloat16 nv_bfloat16;
-#endif
-
-#include <iostream>
-#include <array>
-#include <limits>
-#include <map>
-#include <unordered_map>
-#include <vector>
-#include <cstdlib>
-#include <cstring>
+#include "custom_collective_common.cuh"
 
 namespace vllm {
-#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 <typename T, int sz>
-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 <typename T>
-struct packed_t {
-  // the (P)acked type for load/store
-  using P = array_t<T, 16 / sizeof(T)>;
-  // the (A)ccumulator type for reduction
-  using A = array_t<float, 16 / sizeof(T)>;
-};
-
-#define DINLINE __device__ __forceinline__
-
-// scalar cast functions
-DINLINE float upcast_s(half val) { return __half2float(val); }
-
-template <typename T>
-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 <typename T, int N>
-DINLINE array_t<T, N>& packed_assign_add(array_t<T, N>& a, array_t<T, N> b) {
-#pragma unroll
-  for (int i = 0; i < N; i++) {
-    assign_add(a.data[i], b.data[i]);
-  }
-  return a;
-}
-
-template <typename T, int N>
-DINLINE array_t<float, N> upcast(array_t<T, N> val) {
-  if constexpr (std::is_same<T, float>::value) {
-    return val;
-  } else {
-    array_t<float, N> out;
-#pragma unroll
-    for (int i = 0; i < N; i++) {
-      out.data[i] = upcast_s(val.data[i]);
-    }
-    return out;
-  }
-}
-
-template <typename O>
-DINLINE O downcast(array_t<float, O::size> val) {
-  if constexpr (std::is_same<typename O::type, float>::value) {
-    return val;
-  } else {
-    O out;
-#pragma unroll
-    for (int i = 0; i < O::size; i++) {
-      out.data[i] = downcast_s<typename O::type>(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 <int ngpus>
-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 <int ngpus, bool final_sync = false>
-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 <int ngpus>
-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 <int ngpus, bool final_sync = false>
-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 <typename P, int ngpus, typename A>
-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<P>(tmp);
-}
 
 template <typename T, int ngpus>
 __global__ void __launch_bounds__(512, 1)
@@ -616,6 +325,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 <typename T>
+  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 <typename T>
+  void reduce_scatter(cudaStream_t stream, T* input, T* output, int size,
+                      int threads = 512, int block_limit = defaultBlockLimit);
+  template <typename T>
+  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));
@@ -625,8 +349,8 @@ class CustomAllreduce {
 
 /**
  * To inspect PTX/SASS, copy paste this header file to compiler explorer and
- add a template instantiation:
+ * add a template instantiation:
  * template void vllm::CustomAllreduce::allreduce<half>(cudaStream_t, half *,
- half *, int, int, int);
-*/
-}  // namespace vllm
+ *                                                       half *, int, int, int);
+ */
+}  // namespace vllm

csrc/custom_collective_common.cuh

@@ -0,0 +1,332 @@
+#pragma once
+
+#include <cuda.h>
+#include <cuda_bf16.h>
+#include <cuda_fp16.h>
+#include <cuda_runtime.h>
+
+#if defined(USE_ROCM)
+typedef __hip_bfloat16 nv_bfloat16;
+#endif
+
+#include <iostream>
+#include <array>
+#include <limits>
+#include <map>
+#include <unordered_map>
+#include <vector>
+#include <cstdlib>
+#include <cstring>
+
+namespace vllm {
+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 <typename T, int sz>
+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 <typename T>
+struct packed_t {
+  // the (P)acked type for load/store
+  using P = array_t<T, 16 / sizeof(T)>;
+  // the (A)ccumulator type for reduction
+  using A = array_t<float, 16 / sizeof(T)>;
+};
+
+#define DINLINE __device__ __forceinline__
+
+// scalar cast functions
+DINLINE float upcast_s(half val) { return __half2float(val); }
+
+template <typename T>
+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 <typename T, int N>
+DINLINE array_t<T, N>& packed_assign_add(array_t<T, N>& a, array_t<T, N> b) {
+#pragma unroll
+  for (int i = 0; i < N; i++) {
+    assign_add(a.data[i], b.data[i]);
+  }
+  return a;
+}
+
+template <typename T, int N>
+DINLINE array_t<float, N> upcast(array_t<T, N> val) {
+  if constexpr (std::is_same<T, float>::value) {
+    return val;
+  } else {
+    array_t<float, N> out;
+#pragma unroll
+    for (int i = 0; i < N; i++) {
+      out.data[i] = upcast_s(val.data[i]);
+    }
+    return out;
+  }
+}
+
+template <typename O>
+DINLINE O downcast(array_t<float, O::size> val) {
+  if constexpr (std::is_same<typename O::type, float>::value) {
+    return val;
+  } else {
+    O out;
+#pragma unroll
+    for (int i = 0; i < O::size; i++) {
+      out.data[i] = downcast_s<typename O::type>(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 <int ngpus>
+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 <int ngpus>
+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 <int ngpus, bool final_sync = false>
+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 <int ngpus>
+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 <int ngpus>
+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 <int ngpus, bool final_sync = false>
+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 <typename P, int ngpus, typename A>
+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<P>(tmp);
+}
+
+}  // namespace vllm

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)

csrc/libtorch_stable/activation_kernels.cu

@@ -464,6 +464,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 <typename scalar_t>
+__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)VLLM_LDG(&gate_ptr[idx]);
+    const float u = (float)VLLM_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 <typename scalar_t>
+__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<int64_t>(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)VLLM_LDG(&gate_ptr[idx]);
+    const float u = (float)VLLM_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 vllm
 
 #define LAUNCH_ACTIVATION_GATE_KERNEL_WITH_PARAM(KERNEL, PACKED_KERNEL, PARAM) \
@@ -553,6 +613,54 @@ void swigluoai_and_mul(torch::stable::Tensor& out,    // [..., d]
                        double alpha, double limit) {
   LAUNCH_SIGLUOAI_AND_MUL(vllm::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();
+  VLLM_STABLE_DISPATCH_FLOATING_TYPES(
+      input.scalar_type(), "situ_and_mul_kernel", [&] {
+        vllm::situ_and_mul_kernel<scalar_t><<<grid, block, 0, stream>>>(
+            out.mutable_data_ptr<scalar_t>(), input.const_data_ptr<scalar_t>(),
+            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();
+  VLLM_STABLE_DISPATCH_FLOATING_TYPES(
+      input.scalar_type(), "masked_situ_and_mul_kernel", [&] {
+        vllm::masked_situ_and_mul_kernel<scalar_t><<<grid, block, 0, stream>>>(
+            out.mutable_data_ptr<scalar_t>(), input.const_data_ptr<scalar_t>(),
+            expert_num_tokens.const_data_ptr<int>(), max_num_tokens, d,
+            (float)beta, (float)linear_beta);
+      });
+}
 namespace vllm {
 
 // Element-wise activation kernel template.

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<float>::infinity() : p_lse;
   s_lse = std::isinf(s_lse) ? -std::numeric_limits<float>::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<scalar_t*>(suffix_output.data_ptr()),          \
             reinterpret_cast<float*>(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,

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 <typename scalar_t>
+__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<scalar_t*>(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 <typename scalar_t, typename cache_t, Fp8KVCacheDataType kv_dt>
 __global__ void concat_and_cache_ds_mla_kernel(
     const scalar_t* __restrict__ kv_c,  // [num_tokens, kv_lora_rank]
@@ -902,6 +951,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));
+  vllm::concat_and_cache_mla_grouped_kernel<uint16_t>
+      <<<grid, block, 0, stream>>>(
+          reinterpret_cast<const uint16_t*>(kv_c.data_ptr()),
+          reinterpret_cast<const uint16_t*>(k_pe.data_ptr()),
+          kv_cache_ptrs.const_data_ptr<int64_t>(),
+          slot_mapping.const_data_ptr<int64_t>(), 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 vllm {
 
 template <typename Tout, typename Tin, Fp8KVCacheDataType kv_dt>

csrc/libtorch_stable/custom_all_gather_reduce_scatter.cu

@@ -0,0 +1,362 @@
+#include "torch_utils.h"
+
+#include <torch/csrc/stable/macros.h>
+#include <torch/csrc/stable/accelerator.h>
+#include <torch/csrc/stable/tensor.h>
+#include <torch/headeronly/core/ScalarType.h>
+
+#include "custom_all_reduce.cuh"
+#include "custom_all_gather_reduce_scatter.cuh"
+
+namespace vllm {
+
+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<ngpus><<<blocks, threads, 0, stream>>>(      \
+        ptrs, sg_, self_sg_, reinterpret_cast<CopyPack*>(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 <typename T>
+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<T>::P) != 0 ||
+      stage_size_bytes % sizeof(typename packed_t<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<T>::P);
+  int stage_size = stage_size_bytes / sizeof(typename packed_t<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<T, ngpus>, \
+                                 ptrs, input, output,                          \
+                                 reinterpret_cast<T*>(multicast_buffer),       \
+                                 epochs, rank_, size_per_rank, stage_size))
+#else
+  #define MNNVL_LAMPORT_AG_LAUNCH(ngpus)                                 \
+    mnnvl_lamport_all_gather<T, ngpus>                                   \
+        <<<blocks, kMnnvlLamportAgThreads, 0, stream>>>(                 \
+            ptrs, input, output, reinterpret_cast<T*>(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 <typename T>
+void CustomAllreduce::reduce_scatter(cudaStream_t stream, T* input, T* output,
+                                     int size, int threads, int block_limit) {
+  auto packed_size = packed_t<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<T, ngpus><<<blocks, threads, 0, stream>>>( \
+        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 <typename T>
+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<T>::P::size;
+  if (size % (packed_size * world_size_) != 0 ||
+      stage_size_bytes % sizeof(typename packed_t<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<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<T, ngpus>, ptrs, input, \
+        output, epochs, rank_, size_per_rank, stage_size))
+#else
+  #define MNNVL_LAMPORT_RS_LAUNCH(ngpus)                 \
+    mnnvl_lamport_reduce_scatter_kernel<T, ngpus>        \
+        <<<blocks, kMnnvlLamportRsThreads, 0, stream>>>( \
+            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 vllm
+
+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<vllm::CustomAllreduce*>(_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<void*>(_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<vllm::CustomAllreduce*>(_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<void*>(_local_buffer);
+  auto multicast_buffer = reinterpret_cast<void*>(_multicast_buffer);
+  auto epochs = reinterpret_cast<uint32_t*>(_epoch_buffer);
+  switch (out.scalar_type()) {
+    case torch::headeronly::ScalarType::Float: {
+      fa->mnnvl_lamport_allgather<float>(
+          stream, reinterpret_cast<float*>(inp.mutable_data_ptr()),
+          reinterpret_cast<float*>(out.mutable_data_ptr()), local_buffer,
+          multicast_buffer, epochs, input_size, stage_sz_bytes);
+      break;
+    }
+    case torch::headeronly::ScalarType::Half: {
+      fa->mnnvl_lamport_allgather<half>(
+          stream, reinterpret_cast<half*>(inp.mutable_data_ptr()),
+          reinterpret_cast<half*>(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<nv_bfloat16>(
+          stream, reinterpret_cast<nv_bfloat16*>(inp.mutable_data_ptr()),
+          reinterpret_cast<nv_bfloat16*>(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<vllm::CustomAllreduce*>(_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<void*>(_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<float>(
+          stream, reinterpret_cast<float*>(reg_buffer),
+          reinterpret_cast<float*>(out.mutable_data_ptr()), inp.numel());
+      break;
+    }
+    case torch::headeronly::ScalarType::Half: {
+      fa->reduce_scatter<half>(stream, reinterpret_cast<half*>(reg_buffer),
+                               reinterpret_cast<half*>(out.mutable_data_ptr()),
+                               inp.numel());
+      break;
+    }
+#if (__CUDA_ARCH__ >= 800 || !defined(__CUDA_ARCH__))
+    case torch::headeronly::ScalarType::BFloat16: {
+      fa->reduce_scatter<nv_bfloat16>(
+          stream, reinterpret_cast<nv_bfloat16*>(reg_buffer),
+          reinterpret_cast<nv_bfloat16*>(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<vllm::CustomAllreduce*>(_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<void*>(_local_buffer);
+  auto epochs = reinterpret_cast<uint32_t*>(_epoch_buffer);
+  switch (out.scalar_type()) {
+    case torch::headeronly::ScalarType::Float: {
+      fa->mnnvl_lamport_reduce_scatter<float>(
+          stream, reinterpret_cast<float*>(inp.mutable_data_ptr()),
+          reinterpret_cast<float*>(out.mutable_data_ptr()), local_buffer,
+          epochs, inp.numel(), stage_sz_bytes);
+      break;
+    }
+    case torch::headeronly::ScalarType::Half: {
+      fa->mnnvl_lamport_reduce_scatter<half>(
+          stream, reinterpret_cast<half*>(inp.mutable_data_ptr()),
+          reinterpret_cast<half*>(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<nv_bfloat16>(
+          stream, reinterpret_cast<nv_bfloat16*>(inp.mutable_data_ptr()),
+          reinterpret_cast<nv_bfloat16*>(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");
+  }
+}

csrc/libtorch_stable/custom_all_gather_reduce_scatter_ops.cpp

@@ -0,0 +1,29 @@
+#include "ops.h"
+#include "core/registration.h"
+
+#include <torch/csrc/stable/library.h>
+
+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));
+}

csrc/libtorch_stable/custom_all_reduce.cu

@@ -18,14 +18,14 @@ fptr_t init_custom_ar(const std::vector<fptr_t>& 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 > vllm::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");
 
-  vllm::Signal* ipc_ptrs[8];
+  vllm::Signal* ipc_ptrs[vllm::kMaxCustomCollectiveRanks];
   for (int i = 0; i < world_size; i++) {
     ipc_ptrs[i] = reinterpret_cast<vllm::Signal*>(fake_ipc_ptrs[i]);
   }
@@ -124,7 +124,7 @@ int64_t meta_size() { return sizeof(vllm::Signal); }
 void register_buffer(fptr_t _fa, const std::vector<fptr_t>& fake_ipc_ptrs) {
   auto fa = reinterpret_cast<vllm::CustomAllreduce*>(_fa);
   STD_TORCH_CHECK(fake_ipc_ptrs.size() == fa->world_size_);
-  void* ipc_ptrs[8];
+  void* ipc_ptrs[vllm::kMaxCustomCollectiveRanks];
   for (int i = 0; i < fake_ipc_ptrs.size(); i++) {
     ipc_ptrs[i] = reinterpret_cast<void*>(fake_ipc_ptrs[i]);
   }

csrc/libtorch_stable/dsv3_fused_a_gemm.cu

@@ -647,17 +647,17 @@ __global__ __launch_bounds__(256, 1) void fused_a_gemm_kernel(
 #endif
 }
 
-template <typename T, int kHdIn, int kHdOut, int kTileN>
+template <typename T, int kHdIn, int kHdOut, int kTileN, int kTileK = 256>
 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 <typename T, int kHdIn, int kHdOut, int kTileK = 256>
+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<T, kHdIn, kHdOut, 8, kTileK>(
+        output, mat_a, mat_b, num_tokens, stream, enable_pdl);
+  } else {
+    invokeFusedAGemm<T, kHdIn, kHdOut, 16, kTileK>(
+        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 VLLM_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 VLLM_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) {

csrc/libtorch_stable/moe/grouped_topk_kernels.cu

@@ -25,6 +25,7 @@
 #include "libtorch_stable/torch_utils.h"
 
 #include <cmath>
+#include <cstdint>
 #include <tuple>
 #include <cuda_fp16.h>
 #include <cuda_bf16.h>
@@ -448,7 +449,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;
 }
@@ -890,6 +892,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 <int MaxNumExperts, int MaxNumTopExperts>
+static constexpr bool UseTunedBlockPath =
+    MaxNumTopExperts == 16 && (MaxNumExperts == 896 || MaxNumExperts == 1024);
+
+template <typename T, typename BiasT, ScoringFunc SF>
+__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<float, T>(input);
+  float const bias = cuda_cast<float, BiasT>(correction_bias);
+  if (!is_finite(input_float) || !is_finite(bias)) {
+    return;
+  }
+
+  float const unbiased = apply_scoring<SF>(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 <typename IdxT>
+__device__ __forceinline__ void write_outputs(
+    cg::thread_block_tile<WARP_SIZE> 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<float>{});
+
+  if (lane < topk) {
+    float output = 0.0F;
+    if (finite_mask == 0) {
+      if (renormalize) {
+        output = 1.0F / static_cast<float>(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<IdxT>(lane_expert);
+  }
+}
+
+template <typename T, typename BiasT, typename IdxT, ScoringFunc SF,
+          int MaxNumExperts, int MaxNumTopExperts>
+__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<MaxNumExperts, MaxNumTopExperts> ? 8 : 4;
+  static constexpr int ExpertsPerWorkerWarp = WorkerValuesPerLane * WARP_SIZE;
+  using LaneOwnedRange =
+      reduce_topk::HighExpertLaneOwnedTopKRange<MaxNumExperts,
+                                                MaxNumTopExperts>;
+  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<int32_t>(blockIdx.x);
+  int32_t const lane = static_cast<int32_t>(threadIdx.x) % WARP_SIZE;
+  int32_t const warp_id = static_cast<int32_t>(threadIdx.x) / WARP_SIZE;
+  int32_t const num_experts_i32 = static_cast<int32_t>(num_experts);
+  int32_t const topk_i32 = static_cast<int32_t>(topk);
+  T const* token_scores = scores + int64_t{token} * num_experts;
+
+  for (int32_t expert = static_cast<int32_t>(threadIdx.x);
+       expert < num_experts_i32; expert += BlockDim) {
+    preprocess_score<T, BiasT, SF>(token_scores[expert], bias[expert],
+                                   unbiased_scores[expert],
+                                   biased_scores[expert]);
+  }
+  __syncthreads();
+
+  auto warp = cg::tiled_partition<WARP_SIZE>(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<MaxNumTopExperts>(
+          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<MaxNumTopExperts>(
+        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 <int MaxNumExperts>
+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<uint32_t>(selected > 0 ? selected : 1);
+  }
+};
+
+template <typename T, typename BiasT, typename IdxT, ScoringFunc SF,
+          int MaxNumExperts, int MaxNumTopExperts>
+__global__ void __launch_bounds__(WarpTopKLaunchConfig<MaxNumExperts>::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<MaxNumExperts>::BlockDim;
+  using LaneOwnedRange =
+      reduce_topk::HighExpertLaneOwnedTopKRange<MaxNumExperts,
+                                                MaxNumTopExperts>;
+
+  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<int32_t>(threadIdx.x) % WARP_SIZE;
+  int32_t const warp_id = static_cast<int32_t>(threadIdx.x) / WARP_SIZE;
+  int32_t const global_warp =
+      static_cast<int32_t>(blockIdx.x) * WarpBlockDim / WARP_SIZE + warp_id;
+  int32_t const global_warp_stride =
+      static_cast<int32_t>(gridDim.x) * WarpBlockDim / WARP_SIZE;
+  int32_t const num_experts_i32 = static_cast<int32_t>(num_experts);
+  int32_t const topk_i32 = static_cast<int32_t>(topk);
+  auto warp = cg::tiled_partition<WARP_SIZE>(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<T, BiasT, SF>(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<MaxNumTopExperts>(
+          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<float, BiasT>(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 <int Experts, int TopK>
+struct Tier {
+  static constexpr int kExperts = Experts;
+  static constexpr int kTopK = TopK;
+};
+
+template <typename... Tiers>
+struct TierList {};
+
+using SigmoidBiasTiers =
+    TierList<Tier<128, 8>, 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, 4>, 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 <typename T, typename BiasT, typename IdxT, ScoringFunc SF,
+          int MaxNumExperts, int MaxNumTopExperts>
+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, MaxNumTopExperts> ||
+      MaxNumExperts > 1024 || num_experts >= 1024 ||
+      (num_experts >= 256 && num_tokens <= 1024);
+  if (use_block_kernel) {
+    config.gridDim = static_cast<uint32_t>(num_tokens);
+    config.blockDim = BlockDim;
+    cudaLaunchKernelEx(
+        &config,
+        &single_group_topk_block_kernel<T, BiasT, IdxT, SF, MaxNumExperts,
+                                        MaxNumTopExperts>,
+        scores, topk_values, topk_indices, bias, num_experts, topk, renormalize,
+        static_cast<float>(routed_scaling_factor), enable_pdl);
+  } else {
+    using WarpConfig = WarpTopKLaunchConfig<MaxNumExperts>;
+    config.gridDim = WarpConfig::grid_dim(num_tokens);
+    config.blockDim = WarpConfig::BlockDim;
+    cudaLaunchKernelEx(
+        &config,
+        &single_group_topk_warp_kernel<T, BiasT, IdxT, SF, MaxNumExperts,
+                                       MaxNumTopExperts>,
+        scores, topk_values, topk_indices, bias, num_tokens, num_experts, topk,
+        renormalize, static_cast<float>(routed_scaling_factor), enable_pdl);
+  }
+}
+
+template <typename T, typename BiasT, typename IdxT, ScoringFunc SF>
+bool dispatch(TierList<>*, T*, float*, IdxT*, BiasT const*, int64_t, int64_t,
+              int64_t, bool, double, bool, cudaLaunchConfig_t&) {
+  return false;
+}
+
+template <typename T, typename BiasT, typename IdxT, ScoringFunc SF,
+          typename First, typename... Rest>
+bool dispatch(TierList<First, Rest...>*, 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<T, BiasT, IdxT, SF, First::kExperts, First::kTopK>(
+        scores, topk_values, topk_indices, bias, num_tokens, num_experts, topk,
+        renormalize, routed_scaling_factor, enable_pdl, config);
+    return true;
+  }
+  return dispatch<T, BiasT, IdxT, SF>(
+      static_cast<TierList<Rest...>*>(nullptr), scores, topk_values,
+      topk_indices, bias, num_tokens, num_experts, topk, renormalize,
+      routed_scaling_factor, enable_pdl, config);
+}
+
+}  // namespace detail
+
+template <typename T, typename BiasT, typename IdxT, ScoringFunc SF>
+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<T, BiasT, IdxT, SF>(
+        static_cast<detail::SigmoidBiasTiers*>(nullptr), scores, topk_values,
+        topk_indices, bias, num_tokens, num_experts, topk, renormalize,
+        routed_scaling_factor, enable_pdl, config);
+  } else {
+    return detail::dispatch<T, BiasT, IdxT, SF>(
+        static_cast<detail::PrecomputedSoftmaxBiasTiers*>(nullptr), scores,
+        topk_values, topk_indices, bias, num_tokens, num_experts, topk,
+        renormalize, routed_scaling_factor, enable_pdl, config);
+  }
+}
+
+}  // namespace single_group_topk
+
 template <typename T, typename BiasT, typename IdxT, ScoringFunc SF>
 void invokeNoAuxTc(T* scores, float* topk_values, IdxT* topk_indices,
                    BiasT const* bias, int64_t const num_tokens,
@@ -905,6 +1335,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<T, BiasT, IdxT, SF>(
+          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

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 <cooperative_groups/reduce.h>
 #include <cub/cub.cuh>
 
+#include <cstdint>
+#include <type_traits>
+
 namespace vllm {
 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<sizeof(T) == 4, uint64_t, uint32_t>;
-  using IdxT = std::conditional_t<sizeof(T) == 4, int32_t, int16_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<T>::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<kWARP_SIZE> const& warp) {
-    return cg::reduce(warp, compValIdx, cg::greater<TypeCmp>{});
+#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<TypeCmp>{});
+    } else if constexpr (sizeof(TypeCmp) == 8) {
+      uint32_t hi = static_cast<uint32_t>(compVal >> 32);
+      uint32_t lo = static_cast<uint32_t>(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<TypeCmp>(maxHi) << 32) | static_cast<TypeCmp>(maxLo);
+    } else {
+      TypeCmp result;
+      asm volatile("redux.sync.max.u32 %0, %1, 0xffffffff;\n"
+                   : "=r"(result)
+                   : "r"(compVal));
+      return result;
+    }
   }
 };
 
-////////////////////////////////////////////////////////////////////////////////////////////////////
-
-template <int K_, bool Enable_>
-struct TopKIdx {
-  // by default, empty
+template <int N>
+struct IsPowerOf2 {
+  static constexpr bool value = N > 0 && (N & (N - 1)) == 0;
 };
 
-template <int K_>
-struct TopKIdx<K_, true> {
-  static constexpr int K = K_;
-  int32_t val[K];
+template <int N>
+struct NextPow2 {
+ private:
+  static constexpr unsigned u = static_cast<unsigned>(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<int>(s5 + 1);
 };
 
-////////////////////////////////////////////////////////////////////////////////////////////////////
-
-#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 <int A, int B, int Size, typename T>
+__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;
   }
+}
+
+template <int I, int End, int Step, int PairStride, int Size, typename T>
+__device__ __forceinline__ void topkMergePairs(T* a) {
+  if constexpr (I + Step < End) {
+    topkCompareSwap<I, I + Step, Size, T>(a);
+    topkMergePairs<I + PairStride, End, Step, PairStride, Size, T>(a);
+  } else {
+    (void)a;
+  }
+}
+
+template <int Lo, int N, int R, int Size, typename T>
+__device__ __forceinline__ void topkOEM(T* a) {
+  constexpr int M = R * 2;
+  if constexpr (M < N) {
+    topkOEM<Lo, N, M, Size, T>(a);
+    topkOEM<Lo + R, N - R, M, Size, T>(a);
+    topkMergePairs<Lo + R, Lo + N, R, M, Size, T>(a);
+  } else if constexpr (R < N) {
+    topkCompareSwap<Lo, Lo + R, Size, T>(a);
+  } else {
+    (void)a;
+  }
+}
+
+template <int Lo, int N, int Size, typename T>
+__device__ __forceinline__ void topkSortBatcher(T* a) {
+  if constexpr (N > 1) {
+    constexpr int Half = N / 2;
+    topkSortBatcher<Lo, Half, Size, T>(a);
+    topkSortBatcher<Lo + Half, N - Half, Size, T>(a);
+    topkOEM<Lo, N, 1, Size, T>(a);
+  } else {
+    (void)a;
+  }
+}
 
 template <int N, typename RedType>
-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<N>::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<N>::value;
+      topkSortBatcher<0, P, N, RedType>(topK);
+    }
+  }
+};
 
 template <typename RedType>
 struct Sort<1, RedType> {
-  static __device__ void run(RedType* topK) {}
+  static __device__ void run(RedType*) {}
 };
 
 template <typename RedType>
 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 <typename RedType>
 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 <typename RedType>
 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,110 +256,112 @@ __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 <int K, typename Type, int N, bool IsSorted = false>
-__device__ void reduceTopKFunc(cg::thread_block_tile<kWARP_SIZE> 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<Type>;
-  RedType topK[N];
-#pragma unroll
-  for (int nn = 0; nn < N; ++nn) {
-    topK[nn] = RedType{value[nn], idx[nn]};
-  }
-
-  if constexpr (!IsSorted) {
-    Sort<N, RedType>::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 <int K, typename Type, int N>
 __forceinline__ __device__ void reduceTopK(
     cg::thread_block_tile<kWARP_SIZE> 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 <= 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<Type>;
+  RedType topK[N];
+#pragma unroll
+  for (int nn = 0; nn < N; ++nn) {
+    topK[nn] = RedType{value[nn], idx[nn]};
+  }
 
-  if constexpr (N <= 4) {
-    reduceTopKFunc<K, Type, N>(warp, out, outIdx, value, idx, minValue,
-                               actualK);
-  } else {
-    constexpr int numLoops = N / 4;
-    constexpr int numResults = (numLoops * K - 1) / kWARP_SIZE + 1;
+  Sort<N, RedType>::run(topK);
 
-    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;
+  typename RedType::TypeCmp packedMax{};
+  for (int kk = 0; kk < actualK; ++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<K, Type, 4>(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<K, Type, numResults>(warp, out, outIdx, topKBufferValue,
-                                        topKBufferIdx, minValue, actualK);
+    packedMax = topK[0].reduce(warp);
+    RedType::unpack(out[kk], outIdx[kk], packedMax);
   }
 };
 
-#undef TOPK_SWAP
+template <int NumExperts, int NumTopExperts, int MinExperts, int MaxExperts,
+          int MinTopExperts, int MaxTopExperts>
+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 <int NumExperts, int NumTopExperts>
+using HighExpertLaneOwnedTopKRange =
+    LaneOwnedTopKRange<NumExperts, NumTopExperts,
+                       kHIGH_EXPERT_LANE_OWNED_TOPK_MIN_EXPERTS,
+                       kHIGH_EXPERT_LANE_OWNED_TOPK_MAX_EXPERTS,
+                       kHIGH_EXPERT_LANE_OWNED_TOPK_MIN_TOP_EXPERTS,
+                       kHIGH_EXPERT_LANE_OWNED_TOPK_MAX_TOP_EXPERTS>;
+
+template <int K, typename Type, int N>
+__forceinline__ __device__ void reduceTopKForLane(
+    cg::thread_block_tile<kWARP_SIZE> 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(N > 0, "Top K must have N > 0");
+  static_assert(N <= 64,
+                "Only support candidates number less than or equal to "
+                "64*32=2048");
+  using RedType = TopKRedType<Type>;
+  RedType topK[N];
+#pragma unroll
+  for (int nn = 0; nn < N; ++nn) {
+    topK[nn] = RedType{value[nn], idx[nn]};
+  }
+
+  Sort<N, RedType>::run(topK);
+
+  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];
+    }
+    packedMax = topK[0].reduce(warp);
+    if (laneIdx == kk) {
+      lanePacked = packedMax;
+    }
+  }
+
+  if (laneIdx < K) {
+    RedType::unpack(out, outIdx, lanePacked);
+  } else {
+    out = minValue;
+    outIdx = -1;
+  }
+}
 
 }  // namespace reduce_topk
 }  // namespace moe

csrc/libtorch_stable/moe/moe_align_sum_kernels.cu

@@ -1086,4 +1086,4 @@ void moe_lora_align_block_size(
               has_expert_map);
         }
       });
-}
+}

csrc/libtorch_stable/ops.h

@@ -276,6 +276,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<torch::stable::Tensor> position_ids,
+    std::optional<torch::stable::Tensor> 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<torch::stable::Tensor> position_ids,
+    std::optional<torch::stable::Tensor> 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<torch::stable::Tensor> position_ids,
+    std::optional<torch::stable::Tensor> 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<torch::stable::Tensor> position_ids,
+    std::optional<torch::stable::Tensor> 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<torch::stable::Tensor> position_ids,
+    std::optional<torch::stable::Tensor> 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<torch::stable::Tensor> position_ids,
+    std::optional<torch::stable::Tensor> 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,
@@ -315,6 +370,19 @@ void fused_minimax_m3_qknorm_rope_kv_insert(
     std::optional<torch::stable::Tensor> index_q_out,
     const std::string& kv_cache_dtype, bool skip_index_branch);
 
+#ifdef VLLM_ENABLE_FUSED_KDA_DECODE
+void fused_kda_decode(
+    torch::stable::Tensor const& x, torch::stable::Tensor const& weight,
+    std::optional<torch::stable::Tensor> 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<double> lower_bound,
+    std::optional<torch::stable::Tensor> output_gate,
+    std::optional<torch::stable::Tensor> norm_weight, double norm_eps);
+#endif
+
 #ifdef VLLM_ENABLE_KIMI_K3_ATTN_RES
 void kimi_k3_attn_res(torch::stable::Tensor& prefix,
                       torch::stable::Tensor const& delta,
@@ -383,6 +451,20 @@ fptr_t init_custom_ar(const std::vector<int64_t>& 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<int64_t>& fake_ipc_ptrs);
@@ -421,6 +503,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);
@@ -497,6 +585,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,

csrc/libtorch_stable/torch_bindings.cpp

@@ -324,7 +324,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
@@ -447,6 +448,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("
@@ -468,6 +511,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 VLLM_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 VLLM_ENABLE_KIMI_K3_ATTN_RES
   ops.def(
       "kimi_k3_attn_res("
@@ -538,6 +591,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) -> ()");
 
@@ -696,11 +757,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 VLLM_ENABLE_FUSED_KDA_DECODE
+  ops.impl("fused_kda_decode", TORCH_BOX(&fused_kda_decode));
+#endif
+
 #ifdef VLLM_ENABLE_KIMI_K3_ATTN_RES
   ops.impl("kimi_k3_attn_res", TORCH_BOX(&kimi_k3_attn_res));
 #endif
@@ -726,6 +803,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));
@@ -823,6 +902,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("
@@ -921,6 +1009,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));

csrc/quantization/w8a8/fp8/nvidia/quant_utils.cuh

@@ -13,10 +13,18 @@ namespace vllm {
 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 <typename>
+inline constexpr bool _no_conversion_specialization = false;
+
 template <typename Tout, typename Tin>
 __inline__ __device__ Tout vec_conversion(
     const Tin& x, const __nv_fp8_interpretation_t fp8_type = __NV_E4M3) {
-  return x;
+  static_assert(_no_conversion_specialization<Tin>,
+                "no vec_conversion specialization for this (Tout, Tin) pair");
 }
 
 // float -> c10::Float8_e4m3fn
@@ -301,7 +309,9 @@ __inline__ __device__ bf16_8_t vec_conversion<bf16_8_t, Float8_>(
 template <typename Tout, typename Tin>
 __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<Tin>,
+      "no scaled_vec_conversion specialization for this (Tout, Tin) pair");
 }
 
 // fp8 -> half
@@ -492,6 +502,25 @@ __inline__ __device__ uint8_t scaled_vec_conversion<uint8_t, __nv_bfloat16>(
   __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<uint8_t, c10::BFloat16>(
+    const c10::BFloat16& a, const float scale,
+    const __nv_fp8_interpretation_t fp8_type) {
+  return scaled_vec_conversion<uint8_t, __nv_bfloat16>(
+      reinterpret_cast<const __nv_bfloat16&>(a), scale, fp8_type);
+}
+
+template <>
+__inline__ __device__ uint8_t scaled_vec_conversion<uint8_t, c10::Half>(
+    const c10::Half& a, const float scale,
+    const __nv_fp8_interpretation_t fp8_type) {
+  return scaled_vec_conversion<uint8_t, uint16_t>(
+      reinterpret_cast<const uint16_t&>(a), scale, fp8_type);
+}
+
 // float -> fp8
 template <>
 __inline__ __device__ uint8_t scaled_vec_conversion<uint8_t, float>(

docker/Dockerfile.s390x

@@ -61,13 +61,13 @@ ENV C_INCLUDE_PATH="/usr/local/include:$C_INCLUDE_PATH"
 
 FROM python-install AS torch-vision
 # Install torchvision
-ARG TORCH_VISION_VERSION=v0.26.0
+ARG TORCH_VISION_VERSION=v0.28.0
 WORKDIR /tmp
 RUN --mount=type=cache,target=/root/.cache/uv \
     git clone https://github.com/pytorch/vision.git && \
     cd vision && \
     git checkout $TORCH_VISION_VERSION && \
-    uv pip install torch==2.11.0 --index-url https://download.pytorch.org/whl/cpu && \
+    uv pip install torch==2.13.0 --index-url https://download.pytorch.org/whl/cpu && \
     python setup.py bdist_wheel
 
 FROM python-install AS hf-xet-builder

docs/features/disagg_prefill.md

@@ -49,6 +49,34 @@ Now supports 9 types of connectors:
   --kv-transfer-config '{"kv_connector":"FlexKVConnectorV1","kv_role":"kv_both"}'
   ```
 
+## Reusing prefill token ids on decode
+
+!!! note
+    This applies to disaggregated prefill and decode serving on the `/v1/chat/completions` endpoint, using a KV connector configured as in the Usage example above. It is experimental and subject to change.
+
+In disaggregated serving, the prefill and decode stages both render the chat prompt from `messages` and tokenize it. Because the prefill stage has already produced the token ids, the decode stage can reuse them and skip its own templating and tokenization. The output is otherwise identical to a normal chat completion: it is detokenized to text, and tool and reasoning parsing, streaming, and structured output constraints all still apply.
+
+The token ids are passed to the decode stage through `kv_transfer_params`, the dict already attached to the decode request to coordinate the transfer:
+
+1. Send the prefill request with `return_token_ids` enabled, and read `prompt_token_ids` from the response.
+2. Set `kv_transfer_params["prompt_token_ids"]` to those ids on the decode request. `messages` is still required, but its content is not tokenized when the ids are present.
+
+```python
+prefill = client.chat.completions.create(
+    model=model,
+    messages=messages,
+    extra_body={"return_token_ids": True, "kv_transfer_params": {"do_remote_decode": True}},
+)
+ids = prefill.prompt_token_ids
+
+decode = client.chat.completions.create(
+    model=model,
+    messages=messages,
+    stream=True,
+    extra_body={"kv_transfer_params": {"do_remote_prefill": True, "prompt_token_ids": ids}},
+)
+```
+
 ## Development
 
 We implement disaggregated prefilling by running 2 vLLM instances. One for prefill (we call it prefill instance) and one for decode (we call it decode instance), and then use a connector to transfer the prefill KV caches and results from prefill instance to decode instance.

docs/features/kv_offloading_usage.md

@@ -70,7 +70,8 @@ vllm serve <model> \
 | `cpu_bytes_to_use` | yes | — | both | Total bytes of host memory reserved for the CPU tier across all workers (not per-worker). |
 | `block_size` | no | GPU block size | both | Offloaded block size in tokens; must be a multiple of the GPU block size. Mutually exclusive with `blocks_per_chunk`. |
 | `blocks_per_chunk` | no | `1` | both | Offloaded chunk size in GPU blocks; must be > 0. Alternative to `block_size` for models whose KV cache groups have different block sizes. |
-| `eviction_policy` | no | `lru` | both | Primary tier policy: `lru` or `arc`. |
+| `eviction_policy` | no | `lru` | both | Primary tier policy: built-in `lru`/`arc`, or a custom `CachePolicy` name (see [Custom Eviction Policies](#custom-eviction-policies)). |
+| `cache_policy_module_path` | no | — | both | Python import path for a custom `CachePolicy` not in the built-in registry. Required only when `eviction_policy` is not built-in and wasn't pre-registered via `CachePolicyFactory` (advanced). |
 | `store_threshold` | no | `0` | single-tier | Min lookups before a block is offloaded. Values ≥ 2 are rejected by `TieringOffloadingSpec`. |
 | `max_tracker_size` | no | `64000` | single-tier | Max entries in the lookup tracker. |
 | `secondary_tiers` | no | `[]` | multi-tier | List of secondary tier configs (see below). |
@@ -78,6 +79,36 @@ vllm serve <model> \
 | `self_describing_kv_events` | no | `false` | both | Opt-in. When `true` *and* KV cache events are enabled (`--kv-events-config` with `enable_kv_cache_events`), the connector emits self-describing block-granular `BlockStored`/`BlockRemoved` payloads (constituent block hashes, whole-chunk `token_ids`, per-block `block_size`, parent hash, LoRA + group/cache-spec metadata) instead of the placeholder fallback, so external KV-event consumers can index offloaded blocks. Inert unless events are enabled. With `TieringOffloadingSpec`, a CPU promotion is self-describing when a local request observes its primary-tier `HIT` before event translation; otherwise its stored event may retain the placeholder, while a later `HIT` can backfill metadata for removal. Pending-removal/re-promotion races and externally initiated promotions may also produce placeholders, and consumers must ignore removals for unknown hashes. Full-attention groups only; sliding-window/SSM groups keep the placeholder fallback. In chunk mode (`block_size` > GPU block size, or `blocks_per_chunk` > 1), overlapping chunks re-announce shared per-block hashes, so consumers must reference-count (deduplicate) repeated store/remove announcements. |
 | `spec_module_path` | no | — | both | Python import path for a custom `OffloadingSpec` not in the built-in registry. Required only when `spec_name` is not built-in (advanced). |
 
+## Custom Eviction Policies
+
+`eviction_policy` resolves through `CachePolicyFactory` (`vllm/v1/kv_offload/cpu/policies/factory.py`), which pre-registers the built-in `lru` and `arc` policies.
+
+### Out-of-tree (recommended)
+
+Implement `CachePolicy` (`vllm/v1/kv_offload/cpu/policies/base.py`) in your own package — no vLLM fork or patch required — and point `kv_connector_extra_config` at it directly:
+
+```json
+{
+  "cpu_bytes_to_use": 10737418240,
+  "eviction_policy": "MyCachePolicy",
+  "cache_policy_module_path": "my_package.my_module"
+}
+```
+
+`eviction_policy` is checked against the built-in registry first; if it isn't a registered name, vLLM imports `cache_policy_module_path` and looks up `eviction_policy` as a class name in that module — the same fallback `spec_module_path` provides for a custom `OffloadingSpec`. No import or registration call needs to run before the server starts.
+
+### Registering a friendly short name (in-process only)
+
+If you control the process that constructs the vLLM engine (e.g. an embedding application), you can register a short name once at startup instead of repeating the module path in every config:
+
+```python
+from vllm.v1.kv_offload.cpu.policies.factory import CachePolicyFactory
+
+CachePolicyFactory.register_cache_policy("my_policy", "my_package.my_module", "MyCachePolicy")
+```
+
+Then set `"eviction_policy": "my_policy"` in `kv_connector_extra_config`, the same as `"lru"`/`"arc"`. This only takes effect within the process that ran the `register_cache_policy` call — it does not help when the server is launched as a separate process (e.g. via the `vllm serve` CLI), where the out-of-tree `cache_policy_module_path` config above is the only option.
+
 ## Secondary Tiers
 
 Each entry in `secondary_tiers` is a dict with a required `type` field plus tier-specific fields.

docs/models/supported_models.md

@@ -547,6 +547,7 @@ These models primarily accept the [`LLM.generate`](./generative_models.md#llmgen
 | `KeyeVL1_5ForConditionalGeneration` | Keye-VL-1_5-8B | T + I<sup>E+</sup> + V<sup>E+</sup> | `Kwai-Keye/Keye-VL-1_5-8B` | ✅︎ | ✅︎ |
 | `KimiAudioForConditionalGeneration` | Kimi-Audio | T + A<sup>+</sup> | `moonshotai/Kimi-Audio-7B-Instruct` | | ✅︎ |
 | `KimiK25ForConditionalGeneration` | Kimi-K2.5 | T + I<sup>+</sup> | `moonshotai/Kimi-K2.5` | | ✅︎ |
+| `KimiK3ForConditionalGeneration` | Kimi-K3 | T + I<sup>+</sup> | `moonshotai/Kimi-K3` | | ✅︎ |
 | `KimiVLForConditionalGeneration` | Kimi-VL-A3B-Instruct, Kimi-VL-A3B-Thinking | T + I<sup>+</sup> | `moonshotai/Kimi-VL-A3B-Instruct`, `moonshotai/Kimi-VL-A3B-Thinking` | | ✅︎ |
 | `LightOnOCRForConditionalGeneration` | LightOnOCR-1B | T + I<sup>+</sup> | `lightonai/LightOnOCR-1B`, etc | ✅︎ | ✅︎ |
 | `Lfm2VlForConditionalGeneration` | LFM2-VL | T + I<sup>+</sup> | `LiquidAI/LFM2-VL-450M`, `LiquidAI/LFM2-VL-3B`, `LiquidAI/LFM2-VL-8B-A1B`, etc. | ✅︎ | ✅︎ |

pyproject.toml

@@ -155,6 +155,8 @@ view_seperator = "view_seperator"
 inverse_std_variences = "inverse_std_variences"
 
 [tool.typos.default.extend-words]
+Hel = "Hel"
+wether = "wether"
 iy = "iy"
 indx = "indx"
 # intel cpu features

requirements/cuda.txt

@@ -16,8 +16,8 @@ PyNvVideoCodec==2.0.4
 --extra-index-url https://flashinfer.ai/whl/
 flashinfer-python==0.6.15.post1
 flashinfer-cubin==0.6.15.post1
-apache-tvm-ffi==0.1.10
-tilelang==0.1.9
+apache-tvm-ffi==0.1.11
+tilelang==0.1.12
 nvidia-cudnn-frontend>=1.19.1
 # Required for LLM_NVTX_SCOPES_FOR_PROFILING=1
 nvtx==0.2.15

requirements/test/cuda.txt

@@ -45,7 +45,7 @@ anyio==4.14.1
     #   sse-starlette
     #   starlette
     #   watchfiles
-apache-tvm-ffi==0.1.10
+apache-tvm-ffi==0.1.11
     # via
     #   -c requirements/cuda.txt
     #   xgrammar

setup.py

@@ -783,6 +783,7 @@ class precompiled_wheel_utils:
                             "vllm/_qutlass_C.abi3.so",
                             "vllm/_flashmla_C.abi3.so",
                             "vllm/_flashmla_extension_C.abi3.so",
+                            "vllm/_flashkda_C.abi3.so",
                             "vllm/_sparse_flashmla_C.abi3.so",
                             "vllm/vllm_flash_attn/_vllm_fa2_C.abi3.so",
                             "vllm/vllm_flash_attn/_vllm_fa3_C.abi3.so",
@@ -1150,6 +1151,10 @@ if _is_cuda():
         ext_modules.append(
             CMakeExtension(name="vllm._flashmla_extension_C", optional=True)
         )
+    if USE_PRECOMPILED_EXTENSIONS or (
+        CUDA_HOME and get_nvcc_cuda_version() >= Version("12.0")
+    ):
+        ext_modules.append(CMakeExtension(name="vllm._flashkda_C", optional=True))
     if envs.VLLM_USE_PRECOMPILED or (
         CUDA_HOME and get_nvcc_cuda_version() >= Version("12.3")
     ):

tests/compile/test_dynamic_shapes_compilation.py

@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
-import gc
 import tempfile
 from contextlib import contextmanager
 
@@ -9,8 +8,7 @@ import pytest
 import torch
 
 from tests.models.utils import check_logprobs_close
-from tests.utils import wait_for_rocm_memory_to_settle
-from vllm import LLM, SamplingParams
+from vllm import SamplingParams
 from vllm.compilation.decorators import support_torch_compile
 from vllm.config import CompilationConfig, VllmConfig, set_current_vllm_config
 from vllm.config.compilation import (
@@ -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,
+    vllm_runner,
     model_name,
     shapes_type,
     use_aot_compile,
@@ -79,9 +78,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"]
+
+    # VllmRunner shuts down the engine core on exit, so the eager model
+    # below never races a lingering compiled engine for GPU memory.
+    with vllm_runner(
+        model_name,
         compilation_config={
             "mode": CompilationMode.VLLM_COMPILE,
             "dynamic_shapes_config": {
@@ -90,33 +93,25 @@ def test_dynamic_shapes_compilation(
             },
         },
         max_model_len=1024,
-    )
+        enable_chunked_prefill=None,
+    ) as vllm_model:
+        compiled_outputs = []
+        for p in test_prompts:
+            output = vllm_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))
 
-    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()
+    with vllm_runner(
+        model_name,
+        enforce_eager=True,
+        max_model_len=1024,
+        enable_chunked_prefill=None,
+    ) as vllm_model:
+        eager_outputs = []
+        for p in test_prompts:
+            output = vllm_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,
@@ -241,44 +236,39 @@ def test_model_specialization_with_evaluate_guards(
 
 
 @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(vllm_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 vllm_runner(
+        "Qwen/Qwen3-0.6B",
         max_model_len=512,
         max_num_batched_tokens=1,
+        enable_chunked_prefill=None,
         compilation_config={
             "mode": CompilationMode.VLLM_COMPILE,
             "dynamic_shapes_config": {
                 "type": DynamicShapesType.BACKED.value,
             },
         },
-    )
+    ) as vllm_model:
+        sampling_params = SamplingParams(temperature=0, top_p=0.95, max_tokens=10)
 
-    sampling_params = SamplingParams(temperature=0, top_p=0.95, max_tokens=10)
+        # Generate with static shape inputs
+        output = vllm_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 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()
+        # Generate again to verify compilation works with empty sym_shape_indices
+        output = vllm_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"

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
 import torch.distributed as dist
 
 from vllm.distributed.communication_op import tensor_model_parallel_all_reduce  # noqa
+from vllm.distributed.device_communicators import custom_all_reduce
+from vllm.distributed.device_communicators.custom_all_reduce import CustomAllreduce
 from vllm.distributed.parallel_state import get_tp_group, graph_capture
 
 from ..utils import (
@@ -23,6 +27,59 @@ for i, v in enumerate(test_sizes):
     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 +137,32 @@ 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 +193,29 @@ 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 +236,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)

tests/engine/test_short_mm_context.py

@@ -3,6 +3,8 @@
 
 import pytest
 
+from vllm.exceptions import VLLMValidationError
+
 from ..conftest import IMAGE_ASSETS
 
 HF_IMAGE_PROMPTS = IMAGE_ASSETS.prompts(
@@ -19,7 +21,9 @@ models = ["llava-hf/llava-1.5-7b-hf"]
 def test_context_length_too_short(vllm_runner, image_assets, model):
     images = [asset.pil_image for asset in image_assets]
 
-    with pytest.raises(ValueError, match="longer than the maximum model length"):
+    with pytest.raises(
+        VLLMValidationError, match="longer than the maximum model length"
+    ):
         vllm_model = vllm_runner(
             model,
             # LLaVA has a feature size of 576

tests/entrypoints/llm/test_chat.py

@@ -6,6 +6,7 @@ import pytest
 
 from vllm import LLM
 from vllm.distributed import cleanup_dist_env_and_memory
+from vllm.exceptions import VLLMValidationError
 from vllm.sampling_params import SamplingParams
 
 
@@ -157,7 +158,7 @@ def test_chat_batch_failure_cleanup(llm_for_failure_test):
     batch_2 = [valid_msg, valid_msg]
     sampling_params = SamplingParams(temperature=0, max_tokens=10)
 
-    with pytest.raises(ValueError, match="maximum context length is"):
+    with pytest.raises(VLLMValidationError, match="maximum context length is"):
         llm.chat(batch_1, sampling_params=sampling_params)
     assert llm.llm_engine.get_num_unfinished_requests() == 0
 

tests/entrypoints/llm/test_prompt_validation.py

@@ -5,17 +5,18 @@ import pytest
 import torch
 
 from vllm import LLM
+from vllm.exceptions import VLLMValidationError
 
 
 def test_empty_prompt():
     llm = LLM(model="openai-community/gpt2", enforce_eager=True)
-    with pytest.raises(ValueError, match="decoder prompt cannot be empty"):
+    with pytest.raises(VLLMValidationError, match="decoder prompt cannot be empty"):
         llm.generate([""])
 
 
 def test_out_of_vocab_token():
     llm = LLM(model="openai-community/gpt2", enforce_eager=True)
-    with pytest.raises(ValueError, match="out of vocabulary"):
+    with pytest.raises(VLLMValidationError, match="out of vocabulary"):
         llm.generate({"prompt_token_ids": [999999]})
 
 

tests/entrypoints/multimodal/llm/test_mm_embeds_only.py

@@ -6,6 +6,7 @@ import pytest
 from tests.entrypoints.multimodal.conftest import managed_llm
 from vllm import LLM, SamplingParams
 from vllm.assets.image import ImageAsset
+from vllm.exceptions import VLLMValidationError
 
 MODEL = "llava-hf/llava-1.5-7b-hf"
 PROMPT = "USER: <image>\nDescribe this image briefly.\nASSISTANT:"
@@ -42,7 +43,7 @@ def test_generate_with_embedding(llm: LLM):
 def test_raw_image_rejected(llm: LLM):
     """Raw image input is still rejected when limit=0."""
     raw_image = ImageAsset("stop_sign").pil_image
-    with pytest.raises(ValueError, match=r"At most 0 image\(s\)"):
+    with pytest.raises(VLLMValidationError, match=r"At most 0 image\(s\)"):
         llm.generate(
             {"prompt": PROMPT, "multi_modal_data": {"image": raw_image}},
             sampling_params=SamplingParams(max_tokens=16),

tests/entrypoints/openai/chat_completion/test_chat.py

@@ -18,6 +18,7 @@ from tests.utils import RemoteOpenAIServer
 from vllm.entrypoints.openai.chat_completion.protocol import (
     ChatCompletionRequest,
 )
+from vllm.exceptions import VLLMValidationError
 from vllm.sampling_params import SamplingParams
 
 # any model with a chat template should work here
@@ -1074,7 +1075,7 @@ def test_chat_completion_request_n_parameter_exceeds_default_limit(
         max_tokens=10,
     )
 
-    with pytest.raises(ValueError, match="n must be at most"):
+    with pytest.raises(VLLMValidationError, match="n must be at most"):
         request.to_sampling_params(
             max_tokens=10,
             default_sampling_params={},
@@ -1136,7 +1137,7 @@ def test_chat_completion_request_n_parameter_custom_limit(
         max_tokens=10,
     )
 
-    with pytest.raises(ValueError, match="n must be at most 128"):
+    with pytest.raises(VLLMValidationError, match="n must be at most 128"):
         request_over.to_sampling_params(
             max_tokens=10,
             default_sampling_params={},
@@ -1160,7 +1161,7 @@ def test_chat_completion_request_n_parameter_massive_value(
         max_tokens=1,
     )
 
-    with pytest.raises(ValueError, match="n must be at most"):
+    with pytest.raises(VLLMValidationError, match="n must be at most"):
         request.to_sampling_params(
             max_tokens=1,
             default_sampling_params={},

tests/entrypoints/openai/chat_completion/test_chat_completion.py

@@ -158,3 +158,92 @@ async def test_empty_grammar(client: openai.AsyncOpenAI, model_name: str) -> Non
             ],
             extra_body={"structured_outputs": {"grammar": ""}},
         )
+
+
+# Decode-side token reuse for disaggregated serving. The router forwards the
+# prefill stage's prompt token ids in kv_transfer_params so the decode stage
+# skips re-tokenizing.
+
+TOKEN_IN_MESSAGES = [{"role": "user", "content": "Hello, how are you today?"}]
+DECODE_MESSAGES = [{"role": "user", "content": "unrelated decode-side text"}]
+
+
+@pytest.mark.asyncio
+async def test_kv_transfer_prompt_token_ids_round_trip(client: openai.AsyncOpenAI):
+    """Ids forwarded in kv_transfer_params are used verbatim, skipping tokenize.
+
+    The decode request carries different messages, so a response whose
+    prompt_token_ids match the forwarded ids proves the ids were used rather
+    than the request's own messages. Generated text is not compared across
+    requests because vLLM greedy decoding is not bitwise-reproducible.
+    """
+    baseline = await client.chat.completions.create(
+        model=MODEL_NAME,
+        messages=TOKEN_IN_MESSAGES,
+        max_completion_tokens=16,
+        temperature=0,
+        extra_body={"return_token_ids": True},
+    )
+    reused_ids = baseline.prompt_token_ids
+    assert reused_ids
+
+    decode = await client.chat.completions.create(
+        model=MODEL_NAME,
+        messages=DECODE_MESSAGES,
+        max_completion_tokens=16,
+        temperature=0,
+        extra_body={
+            "kv_transfer_params": {"prompt_token_ids": reused_ids},
+            "return_token_ids": True,
+        },
+    )
+
+    # The engine saw the forwarded ids, not the decode request's own messages.
+    assert decode.prompt_token_ids == reused_ids
+    # text-out: reuse still yields a detokenized message.
+    assert decode.choices[0].message.content
+
+
+@pytest.mark.asyncio
+async def test_kv_transfer_prompt_token_ids_streaming(client: openai.AsyncOpenAI):
+    """Decode-side token reuse streams chat-formatted text-out."""
+    baseline = await client.chat.completions.create(
+        model=MODEL_NAME,
+        messages=TOKEN_IN_MESSAGES,
+        max_completion_tokens=16,
+        temperature=0,
+        extra_body={"return_token_ids": True},
+    )
+    reused_ids = baseline.prompt_token_ids
+    assert reused_ids
+
+    stream = await client.chat.completions.create(
+        model=MODEL_NAME,
+        messages=DECODE_MESSAGES,
+        max_completion_tokens=16,
+        temperature=0,
+        stream=True,
+        extra_body={
+            "kv_transfer_params": {"prompt_token_ids": reused_ids},
+            "return_token_ids": True,
+        },
+    )
+
+    content = ""
+    delta_token_ids: list[int] = []
+    first_chunk = True
+    async for chunk in stream:
+        if first_chunk:
+            # prompt_token_ids arrives once, on the first chunk.
+            assert chunk.prompt_token_ids == reused_ids
+            first_chunk = False
+        if not chunk.choices:
+            continue
+        if chunk.choices[0].delta.content:
+            content += chunk.choices[0].delta.content
+        if tids := getattr(chunk.choices[0], "token_ids", None):
+            delta_token_ids.extend(tids)
+
+    # streamed text-out, reconstructed from deltas, with generated token ids.
+    assert content
+    assert delta_token_ids

tests/entrypoints/openai/chat_completion/test_chat_error.py

@@ -6,7 +6,6 @@ from typing import Any
 from unittest.mock import AsyncMock, MagicMock, patch
 
 import pytest
-from pydantic import ValidationError
 
 from vllm.config.multimodal import MultiModalConfig
 from vllm.entrypoints.openai.chat_completion.protocol import (
@@ -18,6 +17,7 @@ from vllm.entrypoints.openai.engine.protocol import GenerationError
 from vllm.entrypoints.openai.models.protocol import BaseModelPath
 from vllm.entrypoints.openai.models.serving import OpenAIServingModels
 from vllm.entrypoints.scale_out.render.serving import ServingRender
+from vllm.exceptions import VLLMValidationError
 from vllm.outputs import CompletionOutput, RequestOutput
 from vllm.renderers.hf import HfRenderer
 from vllm.renderers.online_renderer import OnlineRenderer
@@ -479,7 +479,7 @@ def test_json_schema_response_format_missing_schema():
 def test_structural_tag_response_format_invalid(format_value):
     """Malformed structural tags should be rejected during request validation."""
     with pytest.raises(
-        ValidationError,
+        VLLMValidationError,
         match="Invalid response_format structural_tag",
     ):
         ChatCompletionRequest(
@@ -493,7 +493,7 @@ def test_structural_tag_response_format_invalid(format_value):
 def test_batch_structural_tag_response_format_invalid(format_value):
     """Batch chat should reject malformed structural tags at request parsing."""
     with pytest.raises(
-        ValidationError,
+        VLLMValidationError,
         match="Invalid response_format structural_tag",
     ):
         BatchChatCompletionRequest(
@@ -507,7 +507,7 @@ def test_batch_structural_tag_response_format_invalid(format_value):
 def test_structured_outputs_structural_tag_invalid(structural_tag):
     """Malformed direct structured_outputs structural tags should be rejected."""
     with pytest.raises(
-        ValidationError,
+        VLLMValidationError,
         match="Invalid structured_outputs structural_tag",
     ):
         ChatCompletionRequest(
@@ -521,7 +521,7 @@ def test_structured_outputs_structural_tag_invalid(structural_tag):
 def test_non_numeric_logprobs_rejected(field_name):
     """A non-numeric logprobs value must be a clean 400 validation error, not a
     TypeError from the mode='before' comparison (which surfaces as HTTP 500)."""
-    with pytest.raises(ValidationError, match=f"`{field_name}` must be an integer"):
+    with pytest.raises(VLLMValidationError, match=f"`{field_name}` must be an integer"):
         ChatCompletionRequest(
             model=MODEL_NAME,
             messages=[{"role": "user", "content": "hello"}],

tests/entrypoints/openai/chat_completion/test_logprob_token_ids.py

@@ -14,11 +14,11 @@ digit-token vocab id).
 import math
 
 import pytest
-from pydantic import ValidationError
 
 from tests.utils import RemoteOpenAIServer
 from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest
 from vllm.entrypoints.openai.completion.protocol import CompletionRequest
+from vllm.exceptions import VLLMValidationError
 
 MODEL_NAME = "Qwen/Qwen2.5-1.5B-Instruct"
 
@@ -87,7 +87,7 @@ def test_completion_request_decouples_top_k_from_explicit_token_ids():
 
 
 def test_completion_rejects_explicit_token_ids_without_generated_tokens():
-    with pytest.raises(ValidationError, match="no output tokens are generated"):
+    with pytest.raises(VLLMValidationError, match="no output tokens are generated"):
         CompletionRequest(
             model=MODEL_NAME,
             prompt="Hello",
@@ -99,7 +99,7 @@ def test_completion_rejects_explicit_token_ids_without_generated_tokens():
 
 
 def test_requests_reject_explicit_token_ids_with_beam_search():
-    with pytest.raises(ValidationError, match="not supported with beam search"):
+    with pytest.raises(VLLMValidationError, match="not supported with beam search"):
         ChatCompletionRequest(
             model=MODEL_NAME,
             messages=[{"role": "user", "content": "Hello"}],
@@ -108,7 +108,7 @@ def test_requests_reject_explicit_token_ids_with_beam_search():
             use_beam_search=True,
         )
 
-    with pytest.raises(ValidationError, match="not supported with beam search"):
+    with pytest.raises(VLLMValidationError, match="not supported with beam search"):
         CompletionRequest(
             model=MODEL_NAME,
             prompt="Hello",

tests/entrypoints/openai/chat_completion/test_serving_chat.py

@@ -2314,3 +2314,34 @@ async def test_streaming_n_gt1_independent_tool_parsers():
             f"Choice {choice_idx}: expected finish_reason='tool_calls', "
             f"got '{reasons[0]}'"
         )
+
+
+def test_make_request_with_harmony_reuses_kv_transfer_prompt_token_ids():
+    """The Harmony reuse branch honors ids forwarded in kv_transfer_params.
+
+    A GPT-OSS server is impractical to stand up here, so this exercises the
+    branch directly on a harmony-configured renderer.
+    """
+    engine = MockEngine()
+    engine.model_config.hf_config = MockHFConfig(model_type="gpt_oss")
+    models = OpenAIServingModels(engine, BASE_MODEL_PATHS)
+    online_renderer = _build_online_renderer(engine, models.registry)
+    assert online_renderer.use_harmony
+
+    request = ChatCompletionRequest(
+        model=MODEL_NAME,
+        messages=[{"role": "user", "content": "hi"}],
+        kv_transfer_params={
+            "prompt_token_ids": [10, 20, 30],
+            "do_remote_prefill": True,
+        },
+    )
+    conversation, engine_inputs = online_renderer._make_request_with_harmony(request)
+
+    assert conversation == []
+    assert len(engine_inputs) == 1
+    engine_input = engine_inputs[0]
+    assert engine_input["type"] == "token"
+    assert engine_input["prompt_token_ids"] == [10, 20, 30]
+    # The reuse key is consumed and other kv_transfer_params are preserved.
+    assert request.kv_transfer_params == {"do_remote_prefill": True}

tests/entrypoints/openai/chat_completion/test_thinking_token_budget_validation.py

@@ -2,15 +2,15 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 import pytest
-from pydantic import ValidationError
 
 from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest
 from vllm.entrypoints.openai.completion.protocol import CompletionRequest
+from vllm.exceptions import VLLMValidationError
 
 
 @pytest.mark.parametrize("raw_value", [-2, 0.6, 10.5])
 def test_chat_completion_request_rejects_invalid_thinking_token_budget(raw_value):
-    with pytest.raises(ValidationError, match="thinking_token_budget"):
+    with pytest.raises(VLLMValidationError, match="thinking_token_budget"):
         ChatCompletionRequest.model_validate(
             {
                 "model": "qwen",
@@ -44,7 +44,7 @@ def test_chat_completion_request_accepts_minus_one_as_unlimited():
 
 @pytest.mark.parametrize("raw_value", [0.6, 3.14, -2])
 def test_completion_request_rejects_invalid_thinking_token_budget(raw_value):
-    with pytest.raises(ValidationError, match="thinking_token_budget"):
+    with pytest.raises(VLLMValidationError, match="thinking_token_budget"):
         CompletionRequest.model_validate(
             {
                 "model": "qwen",

tests/entrypoints/openai/completion/test_completion_error.py

@@ -6,7 +6,6 @@ from typing import Any
 from unittest.mock import AsyncMock, MagicMock
 
 import pytest
-from pydantic import ValidationError
 
 from vllm.config.multimodal import MultiModalConfig
 from vllm.entrypoints.openai.completion.protocol import CompletionRequest
@@ -18,6 +17,7 @@ from vllm.entrypoints.openai.engine.protocol import (
 from vllm.entrypoints.openai.models.protocol import BaseModelPath
 from vllm.entrypoints.openai.models.serving import OpenAIServingModels
 from vllm.entrypoints.scale_out.render.serving import ServingRender
+from vllm.exceptions import VLLMValidationError
 from vllm.outputs import CompletionOutput, RequestOutput
 from vllm.renderers.hf import HfRenderer
 from vllm.renderers.online_renderer import OnlineRenderer
@@ -430,7 +430,7 @@ def test_json_schema_response_format_missing_schema():
 def test_structural_tag_response_format_invalid(format_value):
     """Malformed structural tags should be rejected during request validation."""
     with pytest.raises(
-        ValidationError,
+        VLLMValidationError,
         match="Invalid response_format structural_tag",
     ):
         CompletionRequest(
@@ -445,7 +445,7 @@ def test_structural_tag_response_format_invalid(format_value):
 def test_structured_outputs_structural_tag_invalid(structural_tag):
     """Malformed direct structured_outputs structural tags should be rejected."""
     with pytest.raises(
-        ValidationError,
+        VLLMValidationError,
         match="Invalid structured_outputs structural_tag",
     ):
         CompletionRequest(
@@ -616,7 +616,7 @@ class TestCompletionPromptListLimit:
 def test_non_numeric_logprobs_rejected(field_name):
     """A non-numeric logprobs value must be a clean 400 validation error, not a
     TypeError from the mode='before' comparison (which surfaces as HTTP 500)."""
-    with pytest.raises(ValidationError, match=f"`{field_name}` must be an integer"):
+    with pytest.raises(VLLMValidationError, match=f"`{field_name}` must be an integer"):
         CompletionRequest(
             model=MODEL_NAME,
             prompt="Test prompt",

tests/entrypoints/openai/completion/test_prompt_validation.py

@@ -13,6 +13,7 @@ import torch
 
 from tests.utils import RemoteOpenAIServer
 from vllm.config import ModelConfig
+from vllm.exceptions import VLLMValidationError
 from vllm.renderers.embed_utils import safe_load_prompt_embeds
 
 
@@ -111,5 +112,5 @@ def test_disable_prompt_embeds(dtype: torch.dtype, seq_len: int, hidden_size: in
     buffer.seek(0)
     encoded_tensor = pybase64.b64encode(buffer.getvalue())
 
-    with pytest.raises(ValueError, match="--enable-prompt-embeds"):
+    with pytest.raises(VLLMValidationError, match="--enable-prompt-embeds"):
         safe_load_prompt_embeds(model_config, encoded_tensor)

tests/entrypoints/openai/responses/test_sampling_params.py

@@ -14,6 +14,7 @@ from vllm.entrypoints.openai.responses.protocol import (
     ResponsesRequest,
     ResponseTextConfig,
 )
+from vllm.exceptions import VLLMValidationError
 from vllm.sampling_params import StructuredOutputsParams
 
 
@@ -163,7 +164,7 @@ class TestResponsesRequestSamplingParams:
             text=text_config,
         )
 
-        with pytest.raises(ValueError) as exc_info:
+        with pytest.raises(VLLMValidationError) as exc_info:
             request.to_sampling_params(default_max_tokens=1000)
 
         assert "Cannot specify both structured_outputs and text.format" in str(

tests/entrypoints/serve/instrumentator/test_http_status_metrics.py

@@ -8,54 +8,95 @@ PrometheusInstrumentatorMiddleware before being caught by ServerErrorMiddleware.
 """
 
 from argparse import Namespace
-from http import HTTPStatus
 
 import httpx
 import pytest
-from fastapi import FastAPI, HTTPException, Request
-from fastapi.exceptions import RequestValidationError
-from fastapi.responses import JSONResponse
+from fastapi import HTTPException
 from prometheus_client import CollectorRegistry
-from prometheus_fastapi_instrumentator import Instrumentator
 
-from vllm.entrypoints.serve.utils.server_utils import exception_handler
-from vllm.exceptions import VLLMNotFoundError, VLLMValidationError
+from vllm.entrypoints.openai.api_server import build_app
+from vllm.exceptions import (
+    VLLMNotFoundError,
+    VLLMServerError,
+    VLLMValidationError,
+)
 
 
-@pytest.fixture
+@pytest.fixture(scope="module")
+def should_do_global_cleanup_after_test() -> bool:
+    # This suite never initializes distributed/accelerator state.
+    return False
+
+
+def _build_args() -> Namespace:
+    """Minimal args for ``build_app``; avoids ``make_arg_parser`` device probing."""
+    return Namespace(
+        disable_fastapi_docs=True,
+        enable_offline_docs=False,
+        root_path=None,
+        allowed_origins=["*"],
+        allow_credentials=False,
+        allowed_methods=["*"],
+        allowed_headers=["*"],
+        api_key=None,
+        enable_request_id_headers=False,
+        enable_fault_tolerance=False,
+        middleware=[],
+        log_error_stack=False,
+    )
+
+
+@pytest.fixture(scope="module")
 def registry():
-    """Create a fresh Prometheus registry for each test."""
+    """Shared Prometheus registry for the module-scoped app."""
     return CollectorRegistry()
 
 
-@pytest.fixture
+@pytest.fixture(scope="module")
 def app(registry):
-    """Create a minimal FastAPI app that mirrors vLLM's exception handler
-    and Prometheus middleware setup."""
+    """Build the real vLLM FastAPI app once and attach probe routes that raise.
 
-    app = FastAPI()
+    Patch the name used by ``attach_router`` (imported into the instrumentator
+    metrics module), not ``vllm.v1.metrics.prometheus`` alone — that binding is
+    captured at import time.
+    """
+    import vllm.entrypoints.serve.instrumentator.metrics as metrics_mod
 
-    # Mock app state that exception_handler needs
-    app.state.args = Namespace(log_error_stack=False)
+    original = metrics_mod.get_prometheus_registry
+    metrics_mod.get_prometheus_registry = lambda: registry
+    try:
+        app = build_app(_build_args(), supported_tasks=())
+    finally:
+        metrics_mod.get_prometheus_registry = original
 
-    # Register exception handlers exactly as vLLM does in build_app()
-    app.exception_handler(HTTPException)(_http_exception_handler)
-    app.exception_handler(RequestValidationError)(_validation_exception_handler)
-    app.exception_handler(ValueError)(exception_handler)
-    app.exception_handler(TypeError)(exception_handler)
-    app.exception_handler(OverflowError)(exception_handler)
-    app.exception_handler(NotImplementedError)(exception_handler)
-    app.exception_handler(VLLMValidationError)(exception_handler)
-    app.exception_handler(VLLMNotFoundError)(exception_handler)
-    app.exception_handler(Exception)(exception_handler)
+    @app.get("/raise_http_exception_400")
+    async def raise_http_exception_400():
+        raise HTTPException(status_code=400, detail="bad request")
 
-    # Instrument with Prometheus (same as vLLM's attach_router)
-    Instrumentator(
-        excluded_handlers=["/metrics"],
-        registry=registry,
-    ).add().instrument(app)
+    @app.get("/raise_http_exception_404")
+    async def raise_http_exception_404():
+        raise HTTPException(status_code=404, detail="not found")
+
+    @app.get("/raise_request_validation_error")
+    async def raise_request_validation_error(n: int):
+        # Invalid ``n`` triggers FastAPI's RequestValidationError.
+        return {"n": n}
+
+    @app.get("/raise_vllm_validation_error")
+    async def raise_vllm_validation_error():
+        raise VLLMValidationError("bad parameter", parameter="temperature")
+
+    @app.get("/raise_vllm_not_found_error")
+    async def raise_vllm_not_found_error():
+        raise VLLMNotFoundError("model not found")
+
+    @app.get("/raise_vllm_server_error")
+    async def raise_vllm_server_error():
+        # Bare VLLMServerError goes through vllm_error_handler → 500.
+        # EngineGenerateError / EngineDeadError are not used here: they call
+        # terminate_if_errored and need engine/server state.
+        raise VLLMServerError("internal server failure")
 
-    # Test routes that raise different exception types
     @app.get("/raise_value_error")
     async def raise_value_error():
         raise ValueError("invalid input value")
@@ -72,22 +113,6 @@ def app(registry):
     async def raise_not_implemented_error():
         raise NotImplementedError("feature not supported")
 
-    @app.get("/raise_vllm_validation_error")
-    async def raise_vllm_validation_error():
-        raise VLLMValidationError("bad parameter", parameter="temperature")
-
-    @app.get("/raise_vllm_not_found_error")
-    async def raise_vllm_not_found_error():
-        raise VLLMNotFoundError("model not found")
-
-    @app.get("/raise_http_exception_400")
-    async def raise_http_exception_400():
-        raise HTTPException(status_code=400, detail="bad request")
-
-    @app.get("/raise_http_exception_404")
-    async def raise_http_exception_404():
-        raise HTTPException(status_code=404, detail="not found")
-
     @app.get("/raise_runtime_error")
     async def raise_runtime_error():
         raise RuntimeError("unexpected server error")
@@ -99,14 +124,6 @@ def app(registry):
     return app
 
 
-async def _http_exception_handler(req: Request, exc: HTTPException):
-    return JSONResponse({"error": exc.detail}, status_code=exc.status_code)
-
-
-async def _validation_exception_handler(req: Request, exc: RequestValidationError):
-    return JSONResponse({"error": str(exc)}, status_code=HTTPStatus.BAD_REQUEST)
-
-
 def _get_http_requests_total(registry, method: str, handler: str):
     """Extract the http_requests_total metric values grouped by status.
 
@@ -128,31 +145,31 @@ def _get_http_requests_total(registry, method: str, handler: str):
 
 @pytest.mark.asyncio
 @pytest.mark.parametrize(
-    "endpoint,expected_status_group,expected_http_code",
+    "endpoint,expected_status_group,expected_http_code,request_kwargs",
     [
-        # These should record as 4xx in Prometheus
-        ("/raise_value_error", "4xx", 400),
-        ("/raise_type_error", "4xx", 400),
-        ("/raise_overflow_error", "4xx", 400),
-        ("/raise_vllm_validation_error", "4xx", 400),
-        ("/raise_vllm_not_found_error", "4xx", 404),
-        ("/raise_http_exception_400", "4xx", 400),
-        ("/raise_http_exception_404", "4xx", 404),
-        # NotImplementedError returns 501 which is still 5xx group
-        ("/raise_not_implemented_error", "5xx", 501),
-        # These should record as 5xx in Prometheus (genuine server errors)
-        ("/raise_runtime_error", "5xx", 500),
-        # Successful requests should record as 2xx
-        ("/success", "2xx", 200),
+        ("/raise_http_exception_400", "4xx", 400, {}),
+        ("/raise_http_exception_404", "4xx", 404, {}),
+        ("/raise_request_validation_error", "4xx", 400, {"params": {"n": "x"}}),
+        ("/raise_vllm_validation_error", "4xx", 400, {}),
+        ("/raise_vllm_not_found_error", "4xx", 404, {}),
+        ("/raise_vllm_server_error", "5xx", 500, {}),
+        ("/raise_value_error", "4xx", 400, {}),
+        ("/raise_type_error", "4xx", 400, {}),
+        ("/raise_overflow_error", "4xx", 400, {}),
+        ("/raise_not_implemented_error", "5xx", 501, {}),
+        ("/raise_runtime_error", "5xx", 500, {}),
+        ("/success", "2xx", 200, {}),
     ],
     ids=[
+        "HTTPException(400)->4xx",
+        "HTTPException(404)->4xx",
+        "RequestValidationError->4xx",
+        "VLLMValidationError->4xx",
+        "VLLMNotFoundError->4xx",
+        "VLLMServerError->5xx",
         "ValueError->4xx",
         "TypeError->4xx",
         "OverflowError->4xx",
-        "VLLMValidationError->4xx",
-        "VLLMNotFoundError->4xx",
-        "HTTPException(400)->4xx",
-        "HTTPException(404)->4xx",
         "NotImplementedError->5xx",
         "RuntimeError->5xx",
         "success->2xx",
@@ -164,6 +181,7 @@ async def test_http_requests_total_records_correct_status(
     endpoint,
     expected_status_group,
     expected_http_code,
+    request_kwargs,
 ):
     """Verify that http_requests_total records the correct status group.
 
@@ -177,7 +195,7 @@ async def test_http_requests_total_records_correct_status(
     async with httpx.AsyncClient(
         transport=transport, base_url="http://testserver"
     ) as client:
-        response = await client.get(endpoint)
+        response = await client.get(endpoint, **request_kwargs)
 
     # Verify the HTTP response code returned to the client is correct
     assert response.status_code == expected_http_code, (

tests/entrypoints/unit_tests/test_chat_utils.py

@@ -21,6 +21,7 @@ from vllm.entrypoints.chat_utils import (
     parse_chat_messages,
     parse_chat_messages_async,
 )
+from vllm.exceptions import VLLMValidationError
 from vllm.inputs import MultiModalDataDict, MultiModalUUIDDict
 from vllm.multimodal.utils import (
     encode_audio_url,
@@ -1504,7 +1505,7 @@ def test_parse_chat_messages_rejects_too_many_images_in_one_message(
             "ignore",
             message="coroutine 'async_get_and_parse_image' was never awaited",
         )
-        with pytest.raises(ValueError, match="At most"):
+        with pytest.raises(VLLMValidationError, match="At most"):
             parse_chat_messages(
                 [
                     {
@@ -1540,7 +1541,7 @@ def test_parse_chat_messages_rejects_too_many_images_across_messages(
             "ignore",
             message="coroutine 'async_get_and_parse_image' was never awaited",
         )
-        with pytest.raises(ValueError, match="At most"):
+        with pytest.raises(VLLMValidationError, match="At most"):
             parse_chat_messages(
                 [
                     {

tests/evals/gpt_oss/configs/gpt-oss-20b-xpu-baseline.yaml

@@ -0,0 +1,5 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+model_name: openai/gpt-oss-20b
+metric_threshold: 0.568
+reasoning_effort: low

tests/evals/gpt_oss/configs/gpt-oss-20b-xpu-triton-attn.yaml

@@ -0,0 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+model_name: openai/gpt-oss-20b
+metric_threshold: 0.568
+reasoning_effort: low
+server_args: "--attention-backend TRITON_ATTN"

tests/evals/gpt_oss/configs/models-xpu.txt

@@ -0,0 +1,3 @@
+# Intel XPU model configurations for GPQA evaluation
+gpt-oss-20b-xpu-baseline.yaml
+gpt-oss-20b-xpu-triton-attn.yaml

tests/kernels/attention/test_cpu_attn.py

@@ -534,6 +534,7 @@ def varlen_with_paged_kv(
         isa=isa,
         enable_kv_split=False,
         dynamic_causal=dynamic_causal_tensor,
+        kv_cache_dtype=kv_cache_dtype,
     )
 
     out_without_split = torch.empty_like(query)
@@ -569,6 +570,7 @@ def varlen_with_paged_kv(
         isa=isa,
         enable_kv_split=True,
         dynamic_causal=dynamic_causal_tensor,
+        kv_cache_dtype=kv_cache_dtype,
     )
 
     out_with_split = torch.empty_like(query)
@@ -803,6 +805,24 @@ def test_varlen_with_paged_kv_normal_amx(
     )
 
 
+@pytest.mark.skipif(not torch.cpu._is_amx_tile_supported(), reason="no AMX support.")
+def test_varlen_with_paged_kv_fp8_large_prefill_amx() -> None:
+    varlen_with_paged_kv(
+        seq_lens=[(1024, 1024)] * 4,
+        num_heads=(16, 2),
+        head_size=256,
+        sliding_window=None,
+        dtype=torch.bfloat16,
+        block_size=2176,
+        soft_cap=None,
+        num_blocks=4,
+        use_alibi=False,
+        use_sink=False,
+        isa="amx",
+        kv_cache_dtype="fp8_e4m3",
+    )
+
+
 @pytest.mark.parametrize("seq_lens", SEQ_LENS)
 @pytest.mark.parametrize("num_heads", NUM_HEADS)
 @pytest.mark.parametrize("head_size", HEAD_SIZES_VEC16)

tests/kernels/attention/test_kimi_k3_mla_fused_epilogue.py

@@ -0,0 +1,221 @@
+# 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 vllm.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 vllm.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)
+    )

tests/kernels/core/test_activation.py

@@ -197,6 +197,46 @@ 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",
     [

tests/kernels/core/test_fused_q_kv_rmsnorm.py

@@ -13,7 +13,7 @@ from __future__ import annotations
 import pytest
 import torch
 
-from vllm.models.deepseek_v4.common.ops import fused_q_kv_rmsnorm
+from vllm.models.common.ops import fused_q_kv_rmsnorm
 from vllm.platforms import current_platform
 
 pytestmark = pytest.mark.skipif(

tests/kernels/core/test_fused_rms_norm_gated.py

@@ -7,7 +7,9 @@ matching the eager triton kernel output."""
 import pytest
 import torch
 
-from vllm.third_party.flash_linear_attention.ops.kda import FusedRMSNormGated
+from vllm.third_party.flash_linear_attention.ops.fused_norm_gate import (
+    FusedRMSNormGated,
+)
 from vllm.utils.torch_utils import set_random_seed
 
 DTYPES = [torch.bfloat16]

tests/kernels/moe/test_flashinfer_cutedsl_nvfp4_moe.py

@@ -0,0 +1,230 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Tests for FlashInfer CuTeDSL NVFP4 MoE."""
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+import vllm.model_executor.layers.fused_moe.modular_kernel as mk
+from tests.kernels.quantization.nvfp4_utils import (
+    FLOAT4_E2M1_MAX,
+    FLOAT8_E4M3_MAX,
+    break_fp4_bytes,
+)
+from tests.kernels.utils import torch_moe
+from vllm import _custom_ops as ops
+from vllm.config import ParallelConfig, VllmConfig, set_current_vllm_config
+from vllm.model_executor.layers.fused_moe import fused_topk
+from vllm.model_executor.layers.fused_moe.activation import MoEActivation
+from vllm.model_executor.layers.fused_moe.all2all_utils import (
+    maybe_make_prepare_finalize,
+)
+from vllm.model_executor.layers.fused_moe.config import (
+    FusedMoEConfig,
+    FusedMoEParallelConfig,
+    RoutingMethodType,
+    nvfp4_moe_quant_config,
+)
+from vllm.model_executor.layers.fused_moe.experts.flashinfer_cutedsl_moe import (
+    FlashInferCuteDSLExperts,
+)
+from vllm.model_executor.layers.quantization.utils.flashinfer_fp4_moe import (
+    prepare_nvfp4_moe_layer_for_flashinfer_cutedsl,
+)
+from vllm.platforms import current_platform
+from vllm.utils.flashinfer import has_flashinfer_cutedsl_moe_nvfp4
+from vllm.utils.math_utils import next_power_of_2
+from vllm.utils.torch_utils import set_random_seed
+
+if not has_flashinfer_cutedsl_moe_nvfp4() or not (
+    current_platform.is_device_capability_family(100)
+):
+    pytest.skip(
+        "Requires FlashInfer CuTeDSL NVFP4 MoE on SM100",
+        allow_module_level=True,
+    )
+
+
+def _quantize_nvfp4_linear(
+    weight: torch.Tensor,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    weights_q = []
+    scales = []
+    global_scales = []
+    for expert_weight in weight:
+        global_scale = (
+            FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX / expert_weight.abs().max()
+        ).to(torch.float32)
+        weight_q, scale = ops.scaled_fp4_quant(
+            expert_weight,
+            global_scale,
+            is_sf_swizzled_layout=False,
+        )
+        weights_q.append(weight_q)
+        scales.append(scale)
+        global_scales.append(global_scale)
+    return torch.stack(weights_q), torch.stack(scales), torch.stack(global_scales)
+
+
+def _dequantize_nvfp4_linear(
+    tensor_fp4: torch.Tensor,
+    tensor_sf: torch.Tensor,
+    global_scale: torch.Tensor,
+    dtype: torch.dtype,
+) -> torch.Tensor:
+    assert tensor_fp4.dtype == torch.uint8
+    m, packed_k = tensor_fp4.shape
+    k = packed_k * 2
+    tensor_f32 = break_fp4_bytes(tensor_fp4, torch.float32)
+    tensor_f32 = tensor_f32.reshape(m, k // 16, 16)
+    tensor_sf = tensor_sf.view(torch.float8_e4m3fn).to(torch.float32)
+    tensor_sf = tensor_sf[:, : k // 16] / global_scale
+    return (tensor_f32 * tensor_sf.unsqueeze(-1)).reshape(m, k).to(dtype)
+
+
+@pytest.mark.parametrize("m,n,k,e,topk", [(16, 128, 512, 4, 2)])
+@pytest.mark.parametrize("dtype", [torch.bfloat16])
+@torch.inference_mode()
+def test_flashinfer_cutedsl_fp4_moe_relu2_no_mul(
+    m: int,
+    n: int,
+    k: int,
+    e: int,
+    topk: int,
+    dtype: torch.dtype,
+    workspace_init,
+):
+    set_random_seed(7)
+    with set_current_vllm_config(
+        VllmConfig(parallel_config=ParallelConfig(pipeline_parallel_size=1))
+    ):
+        hidden_states = torch.randn((m, k), device="cuda", dtype=dtype) / 10
+
+        w1 = torch.randn((e, n, k), device="cuda", dtype=dtype) / 15
+        w2 = torch.randn((e, k, n), device="cuda", dtype=dtype) / 15
+        w1_q, w1_scale, w1_global_scale = _quantize_nvfp4_linear(w1)
+        w2_q, w2_scale, w2_global_scale = _quantize_nvfp4_linear(w2)
+
+        score = torch.randn((m, e), device="cuda", dtype=dtype)
+        topk_weights, topk_ids, _ = fused_topk(
+            hidden_states, score, topk, renormalize=False
+        )
+
+        activation = MoEActivation.RELU2_NO_MUL
+        fake_layer = SimpleNamespace(activation=activation)
+        a1_scale = torch.ones(1, device="cuda", dtype=torch.float32)
+        a2_scale = torch.ones(1, device="cuda", dtype=torch.float32)
+        (
+            w1_cutedsl,
+            w1_scale_cutedsl,
+            w1_alpha,
+            a1_scale,
+            w2_cutedsl,
+            w2_scale_cutedsl,
+            w2_alpha,
+            a2_scale,
+        ) = prepare_nvfp4_moe_layer_for_flashinfer_cutedsl(
+            layer=fake_layer,
+            w13=w1_q,
+            w13_scale=w1_scale,
+            w13_scale_2=(1.0 / w1_global_scale),
+            a13_scale=a1_scale,
+            w2=w2_q,
+            w2_scale=w2_scale,
+            w2_scale_2=(1.0 / w2_global_scale),
+            a2_scale=a2_scale,
+        )
+        quant_config = nvfp4_moe_quant_config(
+            g1_alphas=w1_alpha,
+            g2_alphas=w2_alpha,
+            a1_gscale=(1.0 / a1_scale),
+            a2_gscale=(1.0 / a2_scale),
+            w1_scale=w1_scale_cutedsl,
+            w2_scale=w2_scale_cutedsl,
+            is_scale_swizzled=False,
+        )
+        moe_config = FusedMoEConfig(
+            num_experts=e,
+            experts_per_token=topk,
+            hidden_dim=k,
+            intermediate_size=n,
+            num_local_experts=e,
+            num_logical_experts=e,
+            activation=activation,
+            device="cuda",
+            moe_parallel_config=FusedMoEParallelConfig.make_no_parallel(),
+            in_dtype=dtype,
+            routing_method=RoutingMethodType.TopK,
+            max_num_tokens=next_power_of_2(m),
+        )
+
+        cutedsl_experts = mk.FusedMoEKernel(
+            maybe_make_prepare_finalize(
+                moe=moe_config,
+                quant_config=quant_config,
+                allow_new_interface=True,
+                use_monolithic=False,
+            ),
+            FlashInferCuteDSLExperts(
+                moe_config=moe_config,
+                quant_config=quant_config,
+            ),
+        )
+
+        cutedsl_output = cutedsl_experts.apply(
+            hidden_states=hidden_states,
+            w1=w1_cutedsl,
+            w2=w2_cutedsl,
+            topk_weights=topk_weights,
+            topk_ids=topk_ids,
+            activation=activation,
+            global_num_experts=e,
+            expert_map=None,
+            apply_router_weight_on_input=False,
+        )
+
+        a_global_scale = torch.ones(1, device="cuda", dtype=torch.float32)
+        a_q, a_scale = ops.scaled_fp4_quant(
+            hidden_states,
+            a_global_scale,
+            is_sf_swizzled_layout=False,
+        )
+        a_in_dtype = _dequantize_nvfp4_linear(
+            a_q,
+            a_scale,
+            a_global_scale,
+            dtype=dtype,
+        )
+
+        w1_d = torch.empty((e, n, k), device="cuda", dtype=dtype)
+        w2_d = torch.empty((e, k, n), device="cuda", dtype=dtype)
+        for idx in range(e):
+            w1_d[idx] = _dequantize_nvfp4_linear(
+                w1_q[idx],
+                w1_scale[idx],
+                w1_global_scale[idx],
+                dtype=dtype,
+            )
+            w2_d[idx] = _dequantize_nvfp4_linear(
+                w2_q[idx],
+                w2_scale[idx],
+                w2_global_scale[idx],
+                dtype=dtype,
+            )
+
+        torch_output = torch_moe(
+            a_in_dtype,
+            w1_d,
+            w2_d,
+            score,
+            topk,
+            activation=activation,
+        )
+        torch.testing.assert_close(
+            torch_output,
+            cutedsl_output,
+            atol=2e-1,
+            rtol=2e-1,
+        )

tests/kernels/moe/test_grouped_topk.py

@@ -23,6 +23,53 @@ from vllm.platforms import current_platform
 from vllm.utils.torch_utils import set_random_seed
 
 
+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."
 )
@@ -101,3 +148,222 @@ def test_grouped_topk(
             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)

tests/kernels/moe/test_moe.py

@@ -36,6 +36,9 @@ from vllm.model_executor.layers.fused_moe.experts.marlin_moe import (
     batched_fused_marlin_moe,
     fused_marlin_moe,
 )
+from vllm.model_executor.layers.fused_moe.utils import (
+    moe_use_td_hw_supported,
+)
 from vllm.model_executor.layers.quantization.utils.marlin_utils import (
     marlin_permute_bias,
 )
@@ -53,9 +56,12 @@ from vllm.model_executor.layers.quantization.utils.marlin_utils_test import (
 from vllm.model_executor.layers.quantization.utils.quant_utils import quantize_weights
 from vllm.platforms import current_platform
 from vllm.scalar_type import ScalarType, scalar_types
+from vllm.triton_utils import tl
 from vllm.utils.math_utils import next_power_of_2
 from vllm.utils.torch_utils import set_random_seed
 
+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 (tcgen05/TMEM), which ptxas "
+            "rejects on Hopper (sm90) and earlier"
+        )
+    monkeypatch.setenv("VLLM_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:

tests/kernels/moe/test_moe_align_block_size.py

@@ -269,6 +269,7 @@ def test_moe_align_block_size_with_expert_map(
                 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,

tests/kernels/test_bf16_skinny_gemm.py

@@ -0,0 +1,649 @@
+# 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 vllm.model_executor.kernels.linear.cute_dsl.skinny_gemm import SkinnyGemmConfig
+from vllm.models.kimi_k3.nvidia import low_latency_gemm as k3_gemm
+from vllm.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 VLLM_K3_BENCH_SHAPES and are not built
+    # by default, so they must not count as available.
+    production_macros = source.split("#ifdef VLLM_K3_BENCH_SHAPES")[0]
+    explicit = source.split("#undef DISPATCH_DSV3_SHAPE")[1].split(
+        "#ifdef VLLM_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 vllm.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

tests/kernels/test_kda.py

@@ -1,226 +0,0 @@
-# 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 vllm.third_party.flash_linear_attention.ops.kda import (
-    chunk_kda,
-    chunk_kda_with_fused_gate,
-    fused_kda_gate,
-)
-from vllm.third_party.flash_linear_attention.ops.l2norm import l2norm_fwd
-
-DEVICE = "cuda"
-
-
-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"),
-    [
-        ([0, 64], torch.float16),
-        ([0, 15, 100, 300], torch.bfloat16),
-    ],
-)
-@torch.inference_mode()
-def test_chunk_kda_fused_gate_cumsum_matches_unfused(
-    cu_seqlens: list[int],
-    dtype: torch.dtype,
-):
-    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 = torch.rand(1, T, H, dtype=dtype, device=DEVICE).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,
-    ).unsqueeze(0)
-    old_o, old_ht = chunk_kda(
-        q=q.clone(),
-        k=k.clone(),
-        v=v.clone(),
-        g=gate,
-        beta=beta.clone(),
-        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,
-        beta=beta.clone(),
-        A_log=A_log,
-        g_bias=dt_bias,
-        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)

tests/lora/test_qwenvl.py

@@ -186,6 +186,18 @@ QWEN25VL_MODEL_PATH = "Qwen/Qwen2.5-VL-3B-Instruct"
 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("VLLM_BATCH_INVARIANT", "1")
+    # The kernel configuration reads VLLM_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("VLLM_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)
@@ -250,7 +262,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,
@@ -273,6 +290,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,
@@ -283,6 +301,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

tests/models/inkling/test_moe_weight_layout.py

@@ -187,6 +187,32 @@ def test_moe_loads_compressed_tensors_global_scale(
     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)

tests/models/kimi_k3/test_attn_res.py

@@ -5,6 +5,7 @@ import pytest
 import torch
 import torch.nn.functional as F
 
+from vllm.models.kimi_k3.common.mtp import fused_mtp_input
 from vllm.models.kimi_k3.nvidia.ops import attn_res
 from vllm.platforms import current_platform
 
@@ -191,3 +192,35 @@ 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()

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 vllm.model_executor.models.interfaces import supports_eagle3
+from vllm.models.kimi_k3.nvidia import model as kimi_model
+from vllm.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

tests/models/kimi_k3/test_kda.py

@@ -0,0 +1,757 @@
+# 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 vllm import _custom_ops as ops
+from vllm.model_executor.layers.mamba.ops.causal_conv1d import causal_conv1d_update
+from vllm.model_executor.layers.mamba.ops.gather_initial_states import (
+    gather_initial_states,
+)
+from vllm.models.kimi_k3.nvidia.kda import (
+    is_flashkda_supported,
+    is_fused_kda_decode_supported,
+)
+from vllm.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 vllm.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 vllm._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)

tests/models/kimi_k3/test_kda_metadata.py

@@ -0,0 +1,411 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from dataclasses import fields
+
+import pytest
+import torch
+
+from tests.v1.attention.utils import (
+    BatchSpec,
+    create_common_attn_metadata,
+    create_vllm_config,
+)
+from vllm.config import SpeculativeConfig
+from vllm.config.compilation import CUDAGraphMode
+from vllm.models.kimi_k3.nvidia.kda_metadata import (
+    KimiK3KDAAttentionBackend,
+    KimiK3KDAMetadata,
+    KimiK3KDAMetadataBuilder,
+    _mamba_get_block_table_tensor,
+    stage_spec_decode_metadata,
+)
+from vllm.v1.attention.backend import AttentionMetadataBuilder
+from vllm.v1.attention.backends.gdn_attn import (
+    GDNAttentionBackend,
+    GDNAttentionMetadata,
+    GDNAttentionMetadataBuilder,
+)
+from vllm.v1.attention.backends.utils import (
+    NULL_BLOCK_ID,
+    mamba_get_block_table_tensor,
+)
+from vllm.v1.kv_cache_interface import MambaSpec
+
+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)

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 tests.utils import (
+    init_test_distributed_environment,
+    multi_gpu_test,
+    multi_process_parallel,
+)
+from vllm.distributed import get_tp_group
+from vllm.model_executor.warmup.cutedsl_warmup import cutedsl_warmup
+from vllm.models.kimi_k3.nvidia.ops.latent_moe_tail import KimiK3LatentMoETailOp
+from vllm.platforms import current_platform
+
+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)

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 vllm.config import ParallelConfig
+from vllm.models.kimi_k3.nvidia import model as kimi_model
+from vllm.models.kimi_k3.nvidia import mtp as kimi_mtp
+from vllm.models.kimi_k3.nvidia.ops import sequence_parallel as sp_ops
+from vllm.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)

tests/models/multimodal/generation/test_common.py

@@ -1242,7 +1242,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,

tests/models/multimodal/processing/test_gemma4.py

@@ -7,6 +7,7 @@ import pytest
 import torch
 from PIL import Image as PILImage
 
+from vllm.exceptions import VLLMValidationError
 from vllm.model_executor.models.gemma4_mm import (
     Gemma4ForConditionalGeneration,
     Gemma4ImagePixelInputs,
@@ -222,7 +223,7 @@ def test_limit_mm_per_prompt(
     mm_data = {"image": images}
 
     # Expect ValueError when exceeding limit
-    with pytest.raises(ValueError, match="At most 1 image"):
+    with pytest.raises(VLLMValidationError, match="At most 1 image"):
         processor(
             prompt,
             mm_items=processor.info.parse_mm_data(mm_data),

tests/models/multimodal/processing/test_gemma4_unified.py

@@ -7,6 +7,7 @@ import pytest
 import torch
 from PIL import Image as PILImage
 
+from vllm.exceptions import VLLMValidationError
 from vllm.model_executor.models.gemma4_mm import Gemma4ImagePixelInputs
 from vllm.multimodal import MULTIMODAL_REGISTRY
 from vllm.multimodal.inputs import MultiModalFieldConfig
@@ -197,7 +198,7 @@ def test_limit_mm_per_prompt(
 
     mm_data = {"image": images}
 
-    with pytest.raises(ValueError, match="At most 1 image"):
+    with pytest.raises(VLLMValidationError, match="At most 1 image"):
         processor(
             prompt,
             mm_items=processor.info.parse_mm_data(mm_data),

tests/models/registry.py

@@ -1024,6 +1024,10 @@ _MULTIMODAL_EXAMPLE_MODELS = {
         "moonshotai/Kimi-K2.5",
         trust_remote_code=True,
     ),
+    "KimiK3ForConditionalGeneration": _HfExamplesInfo(
+        "moonshotai/Kimi-K3",
+        trust_remote_code=True,
+    ),
     "KimiVLForConditionalGeneration": _HfExamplesInfo(
         "moonshotai/Kimi-VL-A3B-Instruct",
         extras={"thinking": "moonshotai/Kimi-VL-A3B-Thinking"},
@@ -1464,6 +1468,14 @@ _SPECULATIVE_DECODING_EXAMPLE_MODELS = {
         is_available_online=False,
         use_original_num_layers=True,  # DSpark has >1 draft block
     ),
+    "K3DSparkModel": _HfExamplesInfo(
+        "moonshotai/Kimi-K3",
+        speculative_model="Inferact/Kimi-K3-DSpark",
+        use_original_num_layers=True,  # DSpark has >1 draft block
+        trust_remote_code=True,
+        # FIXME: Investigate the NVML failure in CI.
+        is_available_online=False,
+    ),
     "Qwen3DSparkModel": _HfExamplesInfo(
         "Qwen/Qwen3-8B",
         speculative_model="deepseek-ai/dspark_qwen3_8b_block7",
@@ -1662,6 +1674,12 @@ _SPECULATIVE_DECODING_EXAMPLE_MODELS = {
         trust_remote_code=True,
         max_model_len=4096,
     ),
+    "KimiK3MTPModel": _HfExamplesInfo(
+        "moonshotai/Kimi-K3",
+        speculative_model="moonshotai/Kimi-K3",
+        trust_remote_code=True,
+        is_available_online=False,
+    ),
     "LongCatFlashMTPModel": _HfExamplesInfo(
         "meituan-longcat/LongCat-Flash-Chat",
         trust_remote_code=True,

tests/models/test_dspark_mla.py

@@ -0,0 +1,153 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+import torch.nn as nn
+
+from vllm.compilation.wrapper import TorchCompileWithNoGuardsWrapper
+from vllm.model_executor.layers.logits_processor import LogitsProcessor
+from vllm.model_executor.models.qwen3_dspark import DSparkMarkovHead
+from vllm.model_executor.models.registry import ModelRegistry
+from vllm.models.kimi_k3.nvidia import dspark_mla
+from vllm.models.kimi_k3.nvidia.dspark_mla import (
+    K3DSparkForCausalLM,
+    K3DSparkModel,
+    ReplicatedDSparkMarkovHead,
+)
+
+
+def test_dspark_mla_uses_compile_free_model_entrypoint():
+    assert ModelRegistry._try_load_model_cls("K3DSparkModel") is K3DSparkForCausalLM
+    assert not issubclass(K3DSparkModel, TorchCompileWithNoGuardsWrapper)
+
+
+@pytest.mark.parametrize(
+    ("checkpoint_name", "runtime_name", "shard_id"),
+    [
+        (
+            "layers.0.self_attn.q_a_proj.weight",
+            "model.layers.0.self_attn.fused_qkv_a_proj.weight",
+            0,
+        ),
+        (
+            "layers.0.self_attn.kv_a_proj_with_mqa.weight",
+            "model.layers.0.self_attn.fused_qkv_a_proj.weight",
+            1,
+        ),
+        (
+            "layers.0.mlp.gate_proj.weight",
+            "model.layers.0.mlp.gate_up_proj.weight",
+            0,
+        ),
+        (
+            "layers.0.mlp.up_proj.weight",
+            "model.layers.0.mlp.gate_up_proj.weight",
+            1,
+        ),
+        ("context_proj.weight", "model.context_proj.weight", None),
+    ],
+)
+def test_dspark_mla_checkpoint_weight_mapping(checkpoint_name, runtime_name, shard_id):
+    assert K3DSparkForCausalLM.hf_to_vllm_mapper._map_name_with_shard(
+        checkpoint_name
+    ) == (runtime_name, shard_id)
+
+
+def test_dspark_mla_shares_frozen_target_weights_and_skips_training_head():
+    assert not K3DSparkForCausalLM.has_own_embed_tokens
+    assert not K3DSparkForCausalLM.has_own_lm_head
+    assert set(K3DSparkForCausalLM.checkpoint_skip_substrs) == {
+        "confidence_head",
+        "embed_tokens",
+        "lm_head",
+    }
+
+
+@pytest.mark.cpu_test
+def test_dspark_markov_head_replication_is_opt_in(
+    monkeypatch: pytest.MonkeyPatch,
+):
+    from vllm.model_executor.layers import logits_processor, vocab_parallel_embedding
+
+    monkeypatch.setattr(
+        vocab_parallel_embedding, "get_tensor_model_parallel_rank", lambda: 3
+    )
+    monkeypatch.setattr(
+        vocab_parallel_embedding,
+        "get_tensor_model_parallel_world_size",
+        lambda: 8,
+    )
+    monkeypatch.setattr(
+        logits_processor,
+        "get_current_vllm_config",
+        lambda: SimpleNamespace(model_config=None),
+    )
+
+    sharded = DSparkMarkovHead(128, 128, 8, prefix="markov_head")
+    assert sharded.markov_w1.tp_size == 8
+    assert sharded.markov_w2.tp_size == 8
+
+    replicated = ReplicatedDSparkMarkovHead(128, 128, 8, prefix="markov_head")
+    assert isinstance(replicated, DSparkMarkovHead)
+    assert replicated.markov_w1.tp_size == 1
+    assert replicated.markov_w2.tp_size == 1
+    assert replicated.markov_w1.weight.shape == (128, 8)
+    assert replicated.markov_w2.weight.shape == (128, 8)
+
+    def fail_collective(*args, **kwargs):
+        raise AssertionError("replicated Markov head must not invoke TP collectives")
+
+    monkeypatch.setattr(
+        vocab_parallel_embedding,
+        "tensor_model_parallel_all_reduce",
+        fail_collective,
+    )
+    logits_processor = LogitsProcessor(128)
+    monkeypatch.setattr(logits_processor, "_gather_logits", fail_collective)
+
+    markov_embed = replicated.embed(torch.tensor([1, 2]))
+    bias = replicated.bias(markov_embed, logits_processor)
+    assert markov_embed.shape == (2, 8)
+    assert bias.shape == (2, 128)
+
+
+@pytest.mark.cpu_test
+def test_k3_dspark_uses_replicated_markov_head(monkeypatch: pytest.MonkeyPatch):
+    markov_head_calls = []
+
+    class DummyModule(nn.Module):
+        def __init__(self, *args, **kwargs):
+            super().__init__()
+
+    def make_markov_head(*args, **kwargs):
+        markov_head_calls.append((args, kwargs))
+        return DummyModule()
+
+    monkeypatch.setattr(dspark_mla, "get_draft_quant_config", lambda _: None)
+    monkeypatch.setattr(dspark_mla, "ReplicatedLinear", DummyModule)
+    monkeypatch.setattr(dspark_mla, "RMSNorm", DummyModule)
+    monkeypatch.setattr(dspark_mla, "K3DSparkDecoderLayer", DummyModule)
+    monkeypatch.setattr(dspark_mla, "ReplicatedDSparkMarkovHead", make_markov_head)
+
+    config = SimpleNamespace(
+        target_hidden_size=16,
+        num_target_layers=2,
+        hidden_size=8,
+        rms_norm_eps=1e-6,
+        num_hidden_layers=1,
+        vocab_size=128,
+        draft_vocab_size=128,
+        markov_rank=4,
+    )
+    vllm_config = SimpleNamespace(
+        speculative_config=SimpleNamespace(
+            draft_model_config=SimpleNamespace(hf_config=config)
+        )
+    )
+
+    K3DSparkModel(vllm_config=vllm_config, start_layer_id=0, prefix="model")
+
+    assert len(markov_head_calls) == 1

tests/multimodal/media/test_unprocessable_entity_error.py

@@ -14,7 +14,7 @@ import aiohttp
 import pytest
 
 from vllm.entrypoints.serve.utils.error_response import create_error_response
-from vllm.exceptions import VLLMUnprocessableEntityError
+from vllm.exceptions import VLLMClientError, VLLMUnprocessableEntityError
 from vllm.multimodal.media import MediaConnector
 
 
@@ -35,9 +35,9 @@ class TestVLLMUnprocessableEntityError:
         assert "parameter=image_url" in str(exc)
         assert "value=https://example.com/image.jpg" in str(exc)
 
-    def test_is_value_error_subclass(self):
+    def test_is_client_error_subclass(self):
         exc = VLLMUnprocessableEntityError("Test")
-        assert isinstance(exc, ValueError)
+        assert isinstance(exc, VLLMClientError)
 
 
 class TestMediaConnectorErrorHandling:

tests/multimodal/test_processing.py

@@ -8,6 +8,7 @@ import numpy as np
 import pytest
 
 from vllm.config import ModelConfig
+from vllm.exceptions import VLLMValidationError
 from vllm.multimodal import MULTIMODAL_REGISTRY
 from vllm.multimodal.processing.context import InputProcessingContext
 from vllm.multimodal.processing.processor import (
@@ -931,7 +932,11 @@ def test_limit_mm_per_prompt_apply(model_id, num_images, limit, is_valid):
     else:
         mm_data = {"image": [image] * num_images}
 
-    exc_ctx = nullcontext() if is_valid else pytest.raises(ValueError, match="At most")
+    exc_ctx = (
+        nullcontext()
+        if is_valid
+        else pytest.raises(VLLMValidationError, match="At most")
+    )
 
     with exc_ctx:
         processor(

tests/quantization/test_compressed_tensors.py

@@ -473,6 +473,11 @@ def test_compressed_tensors_w4a8_fp8(vllm_runner, args):
             "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(

tests/reasoning/test_kimi_k3_reasoning_parser.py

@@ -0,0 +1,250 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import pytest
+
+from vllm.entrypoints.openai.chat_completion.protocol import (
+    ChatCompletionRequest,
+)
+from vllm.entrypoints.openai.engine.protocol import DeltaMessage
+from vllm.parser.kimi_k3 import KimiK3Parser
+from vllm.parser.parser_manager import ParserManager
+from vllm.reasoning.kimi_k3_reasoning_parser import KimiK3ReasoningParser
+
+pytestmark = pytest.mark.skip_global_cleanup
+
+OPEN = "<|open|>"
+CLOSE = "<|close|>"
+SEP = "<|sep|>"
+THINK_OPEN = f"{OPEN}think{SEP}"
+THINK_CLOSE = f"{CLOSE}think{SEP}"
+RESPONSE_OPEN = f"{OPEN}response{SEP}"
+
+
+class DummyTokenizer:
+    def get_vocab(self) -> dict[str, int]:
+        return {}
+
+    def encode(self, text: str, add_special_tokens: bool = False) -> list[int]:
+        if text == THINK_OPEN:
+            return [1, 2, 3]
+        if text == THINK_CLOSE:
+            return [4, 2, 3]
+        return [ord(ch) for ch in text]
+
+
+class ReasoningOnlyParser(KimiK3Parser):
+    reasoning_parser_cls = KimiK3ReasoningParser
+
+
+def test_parser_manager_selects_kimi_k3_parser_for_reasoning_only():
+    parser_cls = ParserManager.get_parser(reasoning_parser_name="kimi_k3")
+
+    assert parser_cls is not None
+    assert issubclass(parser_cls, KimiK3Parser)
+    assert parser_cls.reasoning_parser_cls is KimiK3ReasoningParser
+    assert parser_cls.tool_parser_cls is None
+
+
+def test_parser_selection_thinking_disabled():
+    parser = KimiK3ReasoningParser(
+        DummyTokenizer(), chat_template_kwargs={"thinking": False}
+    )
+
+    assert parser._thinking_enabled is False
+
+
+def test_extract_reasoning_with_xtml_tags():
+    parser = KimiK3ReasoningParser(DummyTokenizer())
+    request = ChatCompletionRequest(model="test-model", messages=[])
+
+    reasoning, content = parser.extract_reasoning_content(
+        f"{THINK_OPEN}step{THINK_CLOSE}{RESPONSE_OPEN}answer",
+        request,
+    )
+
+    assert reasoning == "step"
+    assert content == "answer"
+
+
+def test_extract_reasoning_with_generation_prefix_consumed():
+    parser = KimiK3ReasoningParser(DummyTokenizer())
+    request = ChatCompletionRequest(model="test-model", messages=[])
+
+    reasoning, content = parser.extract_reasoning_content(
+        f"step{THINK_CLOSE}{RESPONSE_OPEN}answer",
+        request,
+    )
+
+    assert reasoning == "step"
+    assert content == "answer"
+
+
+def test_delegating_parser_strips_response_wrapper_without_tool_parser():
+    parser = ReasoningOnlyParser(DummyTokenizer())
+    request = ChatCompletionRequest(model="test-model", messages=[])
+
+    reasoning, content, tool_calls = parser.parse(
+        f"{THINK_OPEN}step{THINK_CLOSE}{RESPONSE_OPEN}answer",
+        request,
+    )
+
+    assert reasoning == "step"
+    assert content == "answer"
+    assert tool_calls == []
+
+
+def test_is_reasoning_end_uses_full_input_ids():
+    parser = KimiK3ReasoningParser(DummyTokenizer())
+
+    assert not parser.is_reasoning_end([4, 2])
+    assert parser.is_reasoning_end([4, 2, 3])
+
+
+def test_is_reasoning_end_ignores_stale_close_from_prior_turn():
+    # DummyTokenizer: THINK_OPEN -> [1, 2, 3], THINK_CLOSE -> [4, 2, 3].
+    # Multi-turn / agent continuation: a prior turn's think channel (its close
+    # marker) is kept in the prompt, then the current turn opens a new think
+    # block that has not closed yet. Reasoning must read as NOT ended, otherwise
+    # the structured-output gate constrains the current turn's reasoning.
+    parser = KimiK3ReasoningParser(DummyTokenizer())
+
+    stale_close = [4, 2, 3]
+    new_open = [1, 2, 3]
+    # prior close, then current-turn open still unclosed -> not ended
+    assert not parser.is_reasoning_end([*stale_close, *new_open])
+    # ...then the current turn emits its own close -> ended
+    assert parser.is_reasoning_end([*stale_close, *new_open, *stale_close])
+    # open with no close yet -> not ended
+    assert not parser.is_reasoning_end([*new_open])
+
+
+def test_streaming_split_open_marker_is_held_back():
+    parser = KimiK3ReasoningParser(DummyTokenizer())
+
+    first = parser.extract_reasoning_content_streaming(
+        previous_text="",
+        current_text=OPEN,
+        delta_text=OPEN,
+        previous_token_ids=[],
+        current_token_ids=[1],
+        delta_token_ids=[1],
+    )
+    second = parser.extract_reasoning_content_streaming(
+        previous_text=OPEN,
+        current_text=f"{OPEN}think",
+        delta_text="think",
+        previous_token_ids=[1],
+        current_token_ids=[1, 2],
+        delta_token_ids=[2],
+    )
+    third = parser.extract_reasoning_content_streaming(
+        previous_text=f"{OPEN}think",
+        current_text=THINK_OPEN + "step",
+        delta_text=f"{SEP}step",
+        previous_token_ids=[1, 2],
+        current_token_ids=[1, 2, 3, 9],
+        delta_token_ids=[3, 9],
+    )
+
+    assert first is None
+    assert second is None
+    assert isinstance(third, DeltaMessage)
+    assert third.reasoning == "step"
+
+
+def test_streaming_split_close_marker_hands_content_downstream():
+    parser = KimiK3ReasoningParser(DummyTokenizer())
+
+    previous_text = f"{THINK_OPEN}step"
+    partial_close = parser.extract_reasoning_content_streaming(
+        previous_text=previous_text,
+        current_text=previous_text + CLOSE,
+        delta_text=CLOSE,
+        previous_token_ids=[1, 2, 3, 9],
+        current_token_ids=[1, 2, 3, 9, 4],
+        delta_token_ids=[4],
+    )
+    closed = parser.extract_reasoning_content_streaming(
+        previous_text=previous_text + CLOSE,
+        current_text=previous_text + f"{THINK_CLOSE}{RESPONSE_OPEN}answer",
+        delta_text=f"think{SEP}{RESPONSE_OPEN}answer",
+        previous_token_ids=[1, 2, 3, 9, 4],
+        current_token_ids=[1, 2, 3, 9, 4, 2, 3, 10],
+        delta_token_ids=[2, 3, 10],
+    )
+
+    assert partial_close is None
+    assert isinstance(closed, DeltaMessage)
+    assert closed.reasoning is None
+    assert closed.content == f"{RESPONSE_OPEN}answer"
+    assert parser.extract_content_ids([2, 3, 10]) == [10]
+
+
+def test_thinking_disabled_streams_content():
+    parser = KimiK3ReasoningParser(
+        DummyTokenizer(), chat_template_kwargs={"enable_thinking": False}
+    )
+
+    delta = parser.extract_reasoning_content_streaming(
+        previous_text="",
+        current_text=f"{RESPONSE_OPEN}answer",
+        delta_text=f"{RESPONSE_OPEN}answer",
+        previous_token_ids=[],
+        current_token_ids=[1],
+        delta_token_ids=[1],
+    )
+
+    assert isinstance(delta, DeltaMessage)
+    assert delta.content == f"{RESPONSE_OPEN}answer"
+    assert delta.reasoning is None
+
+
+def test_delegating_parser_thinking_false_streams_response_content():
+    parser = ReasoningOnlyParser(
+        DummyTokenizer(), chat_template_kwargs={"thinking": False}
+    )
+    request = ChatCompletionRequest(
+        model="test-model",
+        messages=[],
+        chat_template_kwargs={"thinking": False},
+    )
+
+    first = parser.parse_delta(
+        delta_text="OK",
+        delta_token_ids=[10],
+        request=request,
+        prompt_token_ids=[1],
+        finished=False,
+    )
+    partial_close = parser.parse_delta(
+        delta_text=CLOSE,
+        delta_token_ids=[2],
+        request=request,
+        prompt_token_ids=[1],
+        finished=False,
+    )
+    closed = parser.parse_delta(
+        delta_text=f"response{SEP}",
+        delta_token_ids=[3, 4],
+        request=request,
+        prompt_token_ids=[1],
+        finished=False,
+    )
+
+    assert first is not None
+    assert first.content == "OK"
+    assert first.reasoning is None
+    assert partial_close is None
+    assert closed is None
+
+
+def test_adjust_request_keeps_xtml_markers_contiguous():
+    parser = KimiK3ReasoningParser(DummyTokenizer())
+    request = ChatCompletionRequest(model="test-model", messages=[])
+
+    adjusted = parser.adjust_request(request)
+
+    assert adjusted.skip_special_tokens is False
+    if hasattr(adjusted, "spaces_between_special_tokens"):
+        assert adjusted.spaces_between_special_tokens is False

tests/renderers/test_chat_utils_prompt_embeds.py

@@ -26,6 +26,7 @@ from vllm.entrypoints.chat_utils import (
     parse_chat_messages,
     parse_chat_messages_async,
 )
+from vllm.exceptions import VLLMValidationError
 from vllm.renderers.hf import (
     _PROMPT_EMBEDS_PLACEHOLDER_SPAN_MISMATCH_ERROR,
     _build_mixed_prompt_embeds,
@@ -264,7 +265,7 @@ def test_parse_chat_messages_requires_flag():
             "content": [{"type": "prompt_embeds", "data": b64}],
         }
     ]
-    with pytest.raises(ValueError, match=_ENABLE_PROMPT_EMBEDS_ERROR):
+    with pytest.raises(VLLMValidationError, match=_ENABLE_PROMPT_EMBEDS_ERROR):
         parse_chat_messages(
             messages,
             mc,
@@ -283,7 +284,7 @@ def test_parse_chat_messages_rejects_missing_data():
             "content": [{"type": "prompt_embeds"}],  # no `data`
         }
     ]
-    with pytest.raises(ValueError, match=_PROMPT_EMBEDS_MISSING_DATA_ERROR):
+    with pytest.raises(VLLMValidationError, match=_PROMPT_EMBEDS_MISSING_DATA_ERROR):
         parse_chat_messages(
             messages,
             mc,

tests/renderers/test_completions.py

@@ -11,6 +11,7 @@ import pytest
 import torch
 
 from vllm.config import ModelConfig
+from vllm.exceptions import VLLMValidationError
 from vllm.inputs import SingletonPrompt
 from vllm.renderers import TokenizeParams
 from vllm.renderers.hf import HfRenderer
@@ -286,7 +287,7 @@ class TestRenderPrompt:
         )
 
         with pytest.raises(
-            ValueError,
+            VLLMValidationError,
             match="maximum context length is",
         ):
             renderer.tokenize_prompts(
@@ -307,7 +308,7 @@ class TestRenderPrompt:
         )
 
         with pytest.raises(
-            ValueError,
+            VLLMValidationError,
             match="maximum context length is",
         ):
             renderer.tokenize_prompts(
@@ -328,7 +329,7 @@ class TestRenderPrompt:
         )
 
         with pytest.raises(
-            ValueError,
+            VLLMValidationError,
             match="maximum context length is",
         ):
             renderer.tokenize_prompts(

tests/renderers/test_kimi_k3.py

@@ -0,0 +1,353 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+from dataclasses import dataclass, field
+from typing import Any
+
+import pytest
+
+from vllm.exceptions import VLLMValidationError
+from vllm.renderers import ChatParams
+from vllm.renderers.kimi_k3 import KimiK3Renderer, _merge_k3_media_io_kwargs
+from vllm.renderers.registry import RENDERER_REGISTRY
+from vllm.tokenizers.registry import TokenizerRegistry
+
+
+class StubTokenizer:
+    """Stands in for the model's TikTokenTokenizer.
+
+    Records the kwargs it is called with and returns fixed token ids, so tests
+    can assert how the renderer drives ``apply_chat_template`` without the real
+    (git-LFS) tiktoken vocabulary.
+    """
+
+    def __init__(self, token_ids: list[int]) -> None:
+        self.token_ids = token_ids
+        self.calls: list[dict[str, Any]] = []
+        self.conversations: list[list[dict[str, Any]]] = []
+
+    def apply_chat_template(self, conversation, **kwargs) -> list[int]:
+        self.conversations.append(conversation)
+        self.calls.append(kwargs)
+        return list(self.token_ids)
+
+
+@dataclass
+class MockHFConfig:
+    model_type: str = "kimi_k3"
+
+
+@dataclass
+class MockModelConfig:
+    runner_type: str = "generate"
+    is_multimodal_model: bool = False
+    multimodal_config: Any = None
+    hf_config: MockHFConfig = field(default_factory=MockHFConfig)
+    allowed_local_media_path: str = ""
+    allowed_media_domains: Any = None
+    enable_prompt_embeds: bool = False
+    renderer_num_workers: int = 1
+
+
+@dataclass
+class MockParallelConfig:
+    _api_process_rank: int = 0
+
+
+@dataclass
+class MockVllmConfig:
+    model_config: MockModelConfig
+    parallel_config: MockParallelConfig
+
+
+def _make_renderer(tokenizer: StubTokenizer) -> KimiK3Renderer:
+    config = MockVllmConfig(MockModelConfig(), MockParallelConfig())
+    return KimiK3Renderer(config, tokenizer)
+
+
+def test_kimi_k3_registered():
+    assert RENDERER_REGISTRY.load_renderer_cls("kimi_k3").__name__ == "KimiK3Renderer"
+    assert (
+        TokenizerRegistry.load_tokenizer_cls("kimi_k3").__name__ == "CachedHfTokenizer"
+    )
+
+
+def test_k3_media_io_defaults_preserve_original_mode():
+    # Default: K3 keeps the original image mode (no background flattening).
+    assert _merge_k3_media_io_kwargs(None) == {"image": {"image_mode": None}}
+
+    # Server-/request-level values take precedence over the K3 default.
+    assert _merge_k3_media_io_kwargs({"image": {"image_mode": "RGB"}}) == {
+        "image": {"image_mode": "RGB"}
+    }
+
+    # Unrelated image kwargs are merged with the default.
+    assert _merge_k3_media_io_kwargs(
+        {"image": {"rgba_background_color": (0, 0, 0)}}
+    ) == {"image": {"image_mode": None, "rgba_background_color": (0, 0, 0)}}
+
+
+def test_apply_chat_template_forces_tokenize_and_pins_return_dict():
+    tokenizer = StubTokenizer([7, 8, 9])
+    renderer = _make_renderer(tokenizer)
+    tools = [{"type": "function", "function": {"name": "search"}}]
+    params = ChatParams(
+        chat_template_kwargs={"tools": tools, "tokenize": False, "thinking": True}
+    )
+
+    token_ids = renderer._apply_chat_template(
+        [{"role": "user", "content": "hi"}], params
+    )
+
+    assert token_ids == [7, 8, 9]
+    kwargs = tokenizer.calls[-1]
+    # tokenize is forced on even though the request asked for False, so K3 keeps
+    # the special-vs-ordinary token distinction instead of re-tokenizing a string.
+    assert kwargs["tokenize"] is True
+    # return_dict is pinned False so we always get a flat list of ids.
+    assert kwargs["return_dict"] is False
+    assert kwargs["tools"] == tools
+    assert kwargs["thinking"] is True
+
+
+def test_apply_chat_template_translates_standard_thinking_kwargs():
+    # Standard enable_thinking/reasoning_effort kwargs must be translated
+    # to K3's native thinking/thinking_effort.
+    tokenizer = StubTokenizer([7, 8, 9])
+    renderer = _make_renderer(tokenizer)
+    params = ChatParams(
+        chat_template_kwargs={"enable_thinking": False, "reasoning_effort": "none"}
+    )
+
+    renderer._apply_chat_template([{"role": "user", "content": "hi"}], params)
+
+    kwargs = tokenizer.calls[-1]
+    assert kwargs["thinking"] is False
+    assert "thinking_effort" not in kwargs
+    assert "enable_thinking" not in kwargs
+    assert "reasoning_effort" not in kwargs
+
+
+@pytest.mark.parametrize("reasoning_effort", ["low", "high", "max"])
+def test_apply_chat_template_translates_supported_reasoning_effort(
+    reasoning_effort: str,
+):
+    tokenizer = StubTokenizer([7, 8, 9])
+    renderer = _make_renderer(tokenizer)
+    params = ChatParams(chat_template_kwargs={"reasoning_effort": reasoning_effort})
+
+    renderer._apply_chat_template([{"role": "user", "content": "hi"}], params)
+
+    kwargs = tokenizer.calls[-1]
+    assert kwargs["thinking_effort"] == reasoning_effort
+    assert "reasoning_effort" not in kwargs
+
+
+@pytest.mark.parametrize("reasoning_effort", ["minimal", "medium", "xhigh"])
+def test_apply_chat_template_rejects_unsupported_reasoning_effort(
+    reasoning_effort: str,
+):
+    tokenizer = StubTokenizer([7, 8, 9])
+    renderer = _make_renderer(tokenizer)
+    params = ChatParams(chat_template_kwargs={"reasoning_effort": reasoning_effort})
+
+    with pytest.raises(VLLMValidationError, match="thinking_effort") as exc_info:
+        renderer._apply_chat_template([{"role": "user", "content": "hi"}], params)
+
+    assert exc_info.value.parameter == "thinking_effort"
+    assert exc_info.value.value == reasoning_effort
+    assert tokenizer.calls == []
+
+
+def test_apply_chat_template_validates_canonical_native_thinking_effort():
+    tokenizer = StubTokenizer([7, 8, 9])
+    renderer = _make_renderer(tokenizer)
+    params = ChatParams(
+        chat_template_kwargs={
+            "thinking_effort": "low",
+            "reasoning_effort": "medium",
+        }
+    )
+
+    renderer._apply_chat_template([{"role": "user", "content": "hi"}], params)
+
+    assert tokenizer.calls[-1]["thinking_effort"] == "low"
+
+
+@pytest.mark.parametrize("thinking_effort", ["none", "minimal", "medium", "xhigh"])
+def test_apply_chat_template_rejects_unsupported_native_thinking_effort(
+    thinking_effort: str,
+):
+    tokenizer = StubTokenizer([7, 8, 9])
+    renderer = _make_renderer(tokenizer)
+    params = ChatParams(chat_template_kwargs={"thinking_effort": thinking_effort})
+
+    with pytest.raises(VLLMValidationError, match="thinking_effort") as exc_info:
+        renderer._apply_chat_template([{"role": "user", "content": "hi"}], params)
+
+    assert exc_info.value.parameter == "thinking_effort"
+    assert exc_info.value.value == thinking_effort
+    assert tokenizer.calls == []
+
+
+def test_apply_chat_template_native_k3_kwargs_take_precedence():
+    tokenizer = StubTokenizer([7, 8, 9])
+    renderer = _make_renderer(tokenizer)
+    params = ChatParams(
+        chat_template_kwargs={
+            "thinking": True,
+            "enable_thinking": False,
+            "thinking_effort": "low",
+            "reasoning_effort": "high",
+        }
+    )
+
+    renderer._apply_chat_template([{"role": "user", "content": "hi"}], params)
+
+    kwargs = tokenizer.calls[-1]
+    assert kwargs["thinking"] is True
+    assert kwargs["thinking_effort"] == "low"
+
+
+def test_apply_chat_template_adds_k3_api_metadata():
+    tokenizer = StubTokenizer([7, 8, 9])
+    renderer = _make_renderer(tokenizer)
+    response_format = {"type": "json_object"}
+    params = ChatParams(
+        tool_choice="required",
+        response_format=response_format,
+    )
+
+    renderer._apply_chat_template([{"role": "user", "content": "hi"}], params)
+
+    kwargs = tokenizer.calls[-1]
+    assert kwargs["tool_choice"] == "required"
+    assert kwargs["response_format"] == response_format
+
+
+def test_apply_chat_template_auto_tool_choice_keeps_template_kwarg():
+    tokenizer = StubTokenizer([7, 8, 9])
+    renderer = _make_renderer(tokenizer)
+    params = ChatParams(
+        chat_template_kwargs={"tool_choice": "required"},
+        tool_choice="auto",
+    )
+
+    renderer._apply_chat_template([{"role": "user", "content": "hi"}], params)
+
+    assert tokenizer.calls[-1]["tool_choice"] == "required"
+
+
+def test_apply_chat_template_omits_tool_choice_without_tools():
+    tokenizer = StubTokenizer([7, 8, 9])
+    renderer = _make_renderer(tokenizer)
+
+    renderer._apply_chat_template(
+        [{"role": "user", "content": "hi"}], ChatParams(tool_choice=None)
+    )
+
+    assert "tool_choice" not in tokenizer.calls[-1]
+
+
+def test_render_messages_returns_token_prompt():
+    renderer = _make_renderer(StubTokenizer([1, 2, 3]))
+
+    conversation, prompt = renderer.render_messages(
+        [{"role": "user", "content": "hi"}], ChatParams()
+    )
+
+    assert prompt == {"prompt_token_ids": [1, 2, 3]}
+    assert "multi_modal_data" not in prompt
+    assert conversation[0]["role"] == "user"
+
+
+def test_render_messages_derives_private_xtml_tool_attrs():
+    tokenizer = StubTokenizer([1, 2, 3])
+    renderer = _make_renderer(tokenizer)
+
+    conversation, _ = renderer.render_messages(
+        [
+            {
+                "role": "assistant",
+                "content": "",
+                "tool_calls": [
+                    {
+                        "id": "lookup:0",
+                        "type": "function",
+                        "function": {"name": "lookup", "arguments": "{}"},
+                    },
+                    {
+                        "id": "lookup:1",
+                        "type": "function",
+                        "function": {"name": "lookup", "arguments": "{}"},
+                    },
+                ],
+            },
+            {
+                "role": "tool",
+                "tool_call_id": "lookup:1",
+                "tool": "client-supplied-name",
+                "index": 99,
+                "content": "second",
+            },
+            {
+                "role": "tool",
+                "tool_call_id": "lookup:0",
+                "content": "first",
+            },
+        ],
+        ChatParams(),
+    )
+
+    assert [message["content"] for message in conversation[1:]] == [
+        "first",
+        "second",
+    ]
+    assert conversation[1]["tool"] == "lookup"
+    assert conversation[1]["index"] == 1
+    assert conversation[2]["tool"] == "lookup"
+    assert conversation[2]["index"] == 2
+    assert tokenizer.conversations[-1] == conversation
+
+
+def test_render_messages_ignores_client_supplied_xtml_tool_attrs():
+    tokenizer = StubTokenizer([1, 2, 3])
+    renderer = _make_renderer(tokenizer)
+
+    conversation, _ = renderer.render_messages(
+        [
+            {
+                "role": "assistant",
+                "content": "",
+                "tool_calls": [
+                    {
+                        "id": "lookup:0",
+                        "type": "function",
+                        "function": {"name": "lookup", "arguments": "{}"},
+                    }
+                ],
+            },
+            {
+                "role": "tool",
+                "tool_call_id": "unknown",
+                "tool": "lookup",
+                "index": 1,
+                "content": "result",
+            },
+        ],
+        ChatParams(),
+    )
+
+    assert "tool" not in conversation[1]
+    assert "index" not in conversation[1]
+
+
+@pytest.mark.asyncio
+async def test_render_messages_async_returns_token_prompt():
+    renderer = _make_renderer(StubTokenizer([4, 5]))
+
+    conversation, prompt = await renderer.render_messages_async(
+        [{"role": "user", "content": "hi"}], ChatParams()
+    )
+
+    assert prompt == {"prompt_token_ids": [4, 5]}
+    assert conversation[0]["role"] == "user"

tests/samplers/test_non_finite_params.py

@@ -42,7 +42,7 @@ class TestNonFiniteRepetitionPenalty:
         ids=["nan", "inf", "-inf", "math.nan", "math.inf"],
     )
     def test_non_finite_repetition_penalty_rejected(self, value: float):
-        with pytest.raises(ValueError, match="repetition_penalty"):
+        with pytest.raises(VLLMValidationError, match="repetition_penalty"):
             SamplingParams(repetition_penalty=value)
 
     def test_finite_repetition_penalty_accepted(self):

tests/test_config.py

@@ -1235,6 +1235,26 @@ def test_vllm_config_defaults_are_none():
                 assert getattr(config.compilation_config, k) is None
 
 
+def test_validate_mamba_align_subblock_prefill():
+    """Align mode permits configured prefill chunks smaller than a block."""
+    config = SimpleNamespace(
+        cache_config=SimpleNamespace(
+            block_size=11392,
+            mamba_cache_mode="align",
+        ),
+        parallel_config=SimpleNamespace(
+            decode_context_parallel_size=1,
+        ),
+        scheduler_config=SimpleNamespace(
+            max_num_batched_tokens=8192,
+            long_prefill_token_threshold=4096,
+            disable_chunked_mm_input=False,
+        ),
+    )
+
+    VllmConfig.validate_block_size(config)
+
+
 @pytest.mark.parametrize(
     ("model_id", "compilation_config", "optimization_level"),
     [