kimi-k3

Co-authored-by: Isotr0py <Isotr0py@outlook.com>
Co-authored-by: Thien Tran <gau.nernst@yahoo.com.sg>
Co-authored-by: Bugen Zhao <i@bugenzhao.com>
Co-authored-by: Yongye Zhu <zyy1102000@gmail.com>
Co-authored-by: Ziming Huang <zelda.huanghuang@gmail.com>
Co-authored-by: Roger Wang <hey@rogerw.io>
Co-authored-by: Isotr0py <mozf@inferact.ai>
Co-authored-by: aoshen02 <aoshen02@users.noreply.github.com>
Co-authored-by: Woosuk Kwon <woosuk@inferact.ai>
Co-authored-by: Jee Jee Li <jeejeelee@inferact.ai>
Co-authored-by: aoshen02 <aoshen@inferact.ai>
Co-authored-by: Jiangyun Zhu <riverclouds.zhu@qq.com>
Co-authored-by: Summer Yang <girasoleyang@gmail.com>
Co-authored-by: Kevin H. Luu <khluu000@gmail.com>
Co-authored-by: khluu <khluu000@gmail.com>
Co-authored-by: Bowen Wang <abmfy@icloud.com>
Co-authored-by: gnovack <novackgm@gmail.com>
Co-authored-by: Nick Hill <nickhill123@gmail.com>
Co-authored-by: Yifan Qiao <yifanqiao@inferact.ai>
Co-authored-by: xiaozhoupy <peiyuanzhou1994@gmail.com>
Co-authored-by: Roy Wang <yasong.wang@inferact.ai>
Co-authored-by: Jeff (Junze) Ma <93145857+majunze2001@users.noreply.github.com>
Signed-off-by: zjy0516 <riverclouds.zhu@qq.com>

.dockerignore

@@ -3,6 +3,7 @@
 dist
 vllm/*.so
 vllm/vllm-rs
+.git
 
 # Byte-compiled / optimized / DLL files
 __pycache__/

.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

@@ -416,8 +416,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
@@ -1074,6 +1077,41 @@ 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}")
+  endif()
+  if(KIMI_K3_ATTN_RES_ARCHS)
+    set(KIMI_K3_ATTN_RES_SRC
+      "csrc/libtorch_stable/kimi_k3/attn_res_kernel.cu")
+    set_gencode_flags_for_srcs(
+      SRCS "${KIMI_K3_ATTN_RES_SRC}"
+      CUDA_ARCHS "${KIMI_K3_ATTN_RES_ARCHS}")
+    set_property(SOURCE ${KIMI_K3_ATTN_RES_SRC} APPEND PROPERTY
+      COMPILE_OPTIONS
+      "$<$<COMPILE_LANGUAGE:CUDA>:--expt-relaxed-constexpr;--expt-extended-lambda;--use_fast_math>")
+    list(APPEND VLLM_STABLE_EXT_SRC "${KIMI_K3_ATTN_RES_SRC}")
+    message(STATUS
+      "Building Kimi K3 AttnRes for archs: ${KIMI_K3_ATTN_RES_ARCHS}")
+  endif()
+
   # Hadacore kernels
   cuda_archs_loose_intersection(HADACORE_ARCHS "8.0+PTX;9.0+PTX" "${CUDA_ARCHS}")
   if(HADACORE_ARCHS)
@@ -1115,6 +1153,14 @@ 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)
+    endif()
     # Needed by CUTLASS kernels
     target_compile_definitions(_C_stable_libtorch PRIVATE
       CUTLASS_ENABLE_DIRECT_CUDA_DRIVER_CALL=1)
@@ -1412,6 +1458,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.cuda.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.cuda.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.cuda.set_device(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.cuda.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 "git@github.com:Inferact/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 git@github.com:Inferact/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/custom_all_gather_reduce_scatter.cuh

@@ -0,0 +1,326 @@
+#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,48 @@ 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 +751,86 @@ 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/kimi_k3/attn_res_kernel.cu

@@ -0,0 +1,954 @@
+/*
+ * Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved.
+ */
+
+// Production AttnRes forward for Blackwell (SM100).
+//
+// Warp-specialized online softmax + residual + RMSNorm:
+//   - 1 producer warp issues cp.async.bulk row loads into shared memory.
+//   - 8 consumer warps compute reductions and output.
+//   - Q=res_weight*rms_weight remains in registers across persistent tokens.
+//   - V rows are converted once and cached as FP32 in TMEM between passes.
+//
+// Integration contract: Kimi K3 H=7168, 1<=num_blocks<=8, and token-major
+// block residual storage.
+
+#include "../torch_utils.h"
+
+#include <cfloat>
+#include <cstdint>
+#include <cstdio>
+#include <cuda_runtime.h>
+#include <type_traits>
+
+using bf16_t = __nv_bfloat16;
+
+namespace sm100 {
+namespace fwd_prod_v2 {
+
+constexpr int K_TILE = 1024;
+constexpr int N_CHUNK_DEFAULT = 4;
+constexpr int CHUNK_DEPTH = 2;
+constexpr int BLK = 288;  // 1 producer warp + 8 consumer warps
+constexpr int CONSUMER_THREADS = BLK - 32;  // 256
+constexpr int CONSUMER_WARPS = CONSUMER_THREADS / 32;
+constexpr int CONSUMER_GROUPS = 2;  // two 128-thread consumer groups
+constexpr int CONSUMER_THREADS_PER_GROUP = CONSUMER_THREADS / CONSUMER_GROUPS;
+constexpr int FIRST_USER_NAMED_BARRIER = 8;
+
+__device__ __forceinline__ const bf16_t* residual_addr(
+    const bf16_t* block_res, const bf16_t* layer_res, int source, int N,
+    int token, int block_stride_m, int block_stride_r, int H) {
+  if (source < N - 1) {
+    return block_res + static_cast<long long>(token) * block_stride_m +
+           source * block_stride_r;
+  }
+  return layer_res + static_cast<long long>(token) * H;
+}
+
+__device__ __forceinline__ uint32_t elect_one_sync() {
+  uint32_t pred = 0;
+  uint32_t laneid = 0;
+  asm volatile(
+      "{\n"
+      ".reg .b32 %%rx;\n"
+      ".reg .pred %%px;\n"
+      "     elect.sync %%rx|%%px, %2;\n"
+      "@%%px mov.s32 %1, 1;\n"
+      "     mov.s32 %0, %%rx;\n"
+      "}\n"
+      : "+r"(laneid), "+r"(pred)
+      : "r"(0xffffffff));
+  return pred;
+}
+
+__device__ __forceinline__ void mbarrier_init(uint64_t& barrier,
+                                              int thread_count) {
+  uint32_t const barrier_addr =
+      static_cast<uint32_t>(__cvta_generic_to_shared(&barrier));
+  asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;\n" ::"r"(barrier_addr),
+               "r"(thread_count));
+}
+
+__device__ __forceinline__ void mbarrier_expect_tx(uint64_t& barrier,
+                                                   uint32_t bytes) {
+  uint32_t const barrier_addr =
+      static_cast<uint32_t>(__cvta_generic_to_shared(&barrier));
+  asm volatile("mbarrier.arrive.expect_tx.shared::cta.b64 _, [%0], %1;\n" ::"r"(
+                   barrier_addr),
+               "r"(bytes));
+}
+
+__device__ __forceinline__ void mbarrier_wait(uint64_t& barrier, int phase) {
+  uint32_t const barrier_addr =
+      static_cast<uint32_t>(__cvta_generic_to_shared(&barrier));
+  asm volatile(
+      "{\n"
+      ".reg .pred p;\n"
+      "WAIT:\n"
+      "mbarrier.try_wait.parity.shared::cta.b64 p, [%0], %1;\n"
+      "@p bra DONE;\n"
+      "bra WAIT;\n"
+      "DONE:\n"
+      "}\n" ::"r"(barrier_addr),
+      "r"(phase));
+}
+
+__device__ __forceinline__ void mbarrier_arrive(uint64_t& barrier) {
+  uint32_t const barrier_addr =
+      static_cast<uint32_t>(__cvta_generic_to_shared(&barrier));
+  asm volatile(
+      "{\n"
+      ".reg .b64 state;\n"
+      "mbarrier.arrive.shared::cta.b64 state, [%0];\n"
+      "}\n" ::"r"(barrier_addr));
+}
+
+__device__ __forceinline__ void fence_mbarrier_init() {
+  asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory");
+}
+
+__device__ __forceinline__ void named_barrier_sync(uint32_t num_threads,
+                                                   uint32_t user_barrier_id) {
+  asm volatile(
+      "bar.sync %0, %1;" ::"r"(user_barrier_id + FIRST_USER_NAMED_BARRIER),
+      "r"(num_threads)
+      : "memory");
+}
+
+__device__ __forceinline__ void tmem_allocate(int num_columns, uint32_t* dst) {
+  uint32_t const dst_addr =
+      static_cast<uint32_t>(__cvta_generic_to_shared(dst));
+  asm volatile(
+      "tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" ::"r"(
+          dst_addr),
+      "r"(num_columns));
+}
+
+__device__ __forceinline__ void tmem_free(uint32_t tmem_ptr, int num_columns) {
+  asm volatile(
+      "tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" ::"r"(tmem_ptr),
+      "r"(num_columns));
+}
+
+__device__ __forceinline__ void tmem_release_allocation_lock() {
+  asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;");
+}
+
+__device__ __forceinline__ void tmem_store_wait() {
+  asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory");
+}
+
+template <int N, typename T>
+__device__ __forceinline__ void tmem_load(uint32_t src_addr, T* dst) {
+  uint32_t* values = reinterpret_cast<uint32_t*>(dst);
+  if constexpr (N == 8) {
+    asm volatile(
+        "tcgen05.ld.sync.aligned.32x32b.x8.b32"
+        "{%0, %1, %2, %3, %4, %5, %6, %7}, [%8];\n"
+        : "=r"(values[0]), "=r"(values[1]), "=r"(values[2]), "=r"(values[3]),
+          "=r"(values[4]), "=r"(values[5]), "=r"(values[6]), "=r"(values[7])
+        : "r"(src_addr));
+  } else {
+    static_assert(N == 4, "AttnRes TMEM helpers support x4 and x8");
+    asm volatile(
+        "tcgen05.ld.sync.aligned.32x32b.x4.b32"
+        "{%0, %1, %2, %3}, [%4];\n"
+        : "=r"(values[0]), "=r"(values[1]), "=r"(values[2]), "=r"(values[3])
+        : "r"(src_addr));
+  }
+}
+
+template <int N, typename T>
+__device__ __forceinline__ void tmem_store(uint32_t dst_addr, T* src) {
+  uint32_t* values = reinterpret_cast<uint32_t*>(src);
+  if constexpr (N == 8) {
+    asm volatile(
+        "tcgen05.st.sync.aligned.32x32b.x8.b32"
+        "[%8], {%0, %1, %2, %3, %4, %5, %6, %7};\n" ::"r"(values[0]),
+        "r"(values[1]), "r"(values[2]), "r"(values[3]), "r"(values[4]),
+        "r"(values[5]), "r"(values[6]), "r"(values[7]), "r"(dst_addr));
+  } else {
+    static_assert(N == 4, "AttnRes TMEM helpers support x4 and x8");
+    asm volatile(
+        "tcgen05.st.sync.aligned.32x32b.x4.b32"
+        "[%4], {%0, %1, %2, %3};\n" ::"r"(values[0]),
+        "r"(values[1]), "r"(values[2]), "r"(values[3]), "r"(dst_addr));
+  }
+}
+
+__device__ __forceinline__ float2 float2_add(const float2& a, const float2& b) {
+  float2 result;
+  asm volatile("add.rn.f32x2 %0, %1, %2;\n"
+               : "=l"(reinterpret_cast<uint64_t&>(result))
+               : "l"(reinterpret_cast<uint64_t const&>(a)),
+                 "l"(reinterpret_cast<uint64_t const&>(b)));
+  return result;
+}
+
+__device__ __forceinline__ float2 float2_mul(const float2& a, const float2& b) {
+  float2 result;
+  asm volatile("mul.f32x2 %0, %1, %2;\n"
+               : "=l"(reinterpret_cast<uint64_t&>(result))
+               : "l"(reinterpret_cast<uint64_t const&>(a)),
+                 "l"(reinterpret_cast<uint64_t const&>(b)));
+  return result;
+}
+
+__device__ __forceinline__ float2 float2_fma(const float2& a, const float2& b,
+                                             const float2& c) {
+  float2 result;
+  asm volatile("fma.rn.f32x2 %0, %1, %2, %3;\n"
+               : "=l"(reinterpret_cast<uint64_t&>(result))
+               : "l"(reinterpret_cast<uint64_t const&>(a)),
+                 "l"(reinterpret_cast<uint64_t const&>(b)),
+                 "l"(reinterpret_cast<uint64_t const&>(c)));
+  return result;
+}
+
+template <int NC>
+struct FwdSmemPlan {
+  alignas(16) uint64_t bar_ready[CHUNK_DEPTH];
+  alignas(16) uint64_t bar_consumed[CHUNK_DEPTH];
+  alignas(16) uint64_t bar_output_norm_ready;
+  alignas(16) float2 ws_stats[CONSUMER_WARPS][NC];
+  uint32_t tmem_base;
+};
+
+__device__ __forceinline__ void cp_async_bulk(void* smem_dst,
+                                              const void* gmem_src, int bytes,
+                                              uint64_t& mbar) {
+  uint32_t const s = static_cast<uint32_t>(__cvta_generic_to_shared(smem_dst));
+  uint32_t const m = static_cast<uint32_t>(__cvta_generic_to_shared(&mbar));
+  asm volatile(
+      "cp.async.bulk.shared::cta.global.mbarrier::complete_tx::bytes [%0], "
+      "[%1], %2, [%3];\n" ::"r"(s),
+      "l"(gmem_src), "r"(bytes), "r"(m)
+      : "memory");
+}
+
+template <int H, int NC = N_CHUNK_DEFAULT, bool RELEASE_TMEM = false,
+          bool HAS_DELTA = false, bool HAS_OUTPUT_NORM = false,
+          bool OUTPUT_NORM_IN_SMEM = false>
+__global__ void __launch_bounds__(BLK, 1) attn_res_fwd_online_v2_kernel(
+    const bf16_t* __restrict__ block_res, bf16_t* __restrict__ layer_res,
+    const bf16_t* __restrict__ delta, const bf16_t* __restrict__ res_w,
+    const bf16_t* __restrict__ rms_w, bf16_t* __restrict__ output, int N, int T,
+    int B, int block_stride_m, int block_stride_r, float rms_eps,
+    const bf16_t* __restrict__ output_norm_weight, float output_norm_eps) {
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 1000 && __CUDA_ARCH__ < 1100
+  constexpr float LOG2_E = 1.4426950408889634f;
+  constexpr int N_CHUNK = NC;
+  // The two-source specialization only consumes half of the TMEM columns.
+  constexpr int TMEM_COLS_ALLOC = NC == 2 ? 128 : 256;
+  constexpr int NUM_BUFS = CHUNK_DEPTH * NC;
+  constexpr int NHT = H / K_TILE;
+  constexpr int SLICES_PER_GROUP =
+      (NHT + CONSUMER_GROUPS - 1) / CONSUMER_GROUPS;
+  constexpr int VEC = 8;
+  constexpr int ACC_PER_THREAD = H == 7168 ? 28 : SLICES_PER_GROUP * VEC;
+  constexpr int TMEM_V_COLS_PER_GROUP = SLICES_PER_GROUP * N_CHUNK * VEC;
+  constexpr int TMEM_V_COLS_TOTAL = CONSUMER_GROUPS * TMEM_V_COLS_PER_GROUP;
+  static_assert(TMEM_V_COLS_TOTAL <= TMEM_COLS_ALLOC);
+  static_assert(H >= 4096 && H <= 8192);
+  static_assert(H % K_TILE == 0);
+
+  const int tid = threadIdx.x;
+  const int wid = tid >> 5;
+  const int lane = tid & 31;
+  const int TB = T * B;
+  const int num_ctas = gridDim.x;
+  const int num_chunks = (N + N_CHUNK - 1) / N_CHUNK;
+
+  const int comp_wid = wid - 1;
+  const int comp_tid = tid - 32;
+  const int group = (comp_wid >= 4) ? 1 : 0;
+  const int ct_in_group =
+      (comp_tid >= 0) ? (comp_tid & (CONSUMER_THREADS_PER_GROUP - 1)) : -1;
+  const int k_local = ct_in_group * VEC;
+
+  constexpr size_t V_BYTES = (size_t)NUM_BUFS * H * sizeof(bf16_t);
+  constexpr size_t DELTA_BYTES =
+      HAS_DELTA ? (size_t)CHUNK_DEPTH * H * sizeof(bf16_t) : 0;
+  constexpr size_t OUTPUT_NORM_BYTES =
+      OUTPUT_NORM_IN_SMEM ? (size_t)H * sizeof(bf16_t) : 0;
+  extern __shared__ __align__(16) char smem_raw[];
+  bf16_t* v_bufs = reinterpret_cast<bf16_t*>(smem_raw);  // [NUM_BUFS][H]
+  bf16_t* delta_bufs = reinterpret_cast<bf16_t*>(smem_raw + V_BYTES);
+  bf16_t* output_norm_buf =
+      reinterpret_cast<bf16_t*>(smem_raw + V_BYTES + DELTA_BYTES);
+  FwdSmemPlan<NC>& plan = *reinterpret_cast<FwdSmemPlan<NC>*>(
+      smem_raw + V_BYTES + DELTA_BYTES + OUTPUT_NORM_BYTES);
+
+  auto slot_of = [](long long gci, int n) {
+    return (int)(gci % CHUNK_DEPTH) * N_CHUNK + n;
+  };
+  auto phase_of = [](long long gci) { return (int)((gci / CHUNK_DEPTH) & 1); };
+  auto buf_ptr = [&](int slot) -> bf16_t* { return v_bufs + slot * H; };
+  auto delta_buf_ptr = [&](int chunk_slot) -> bf16_t* {
+    return delta_bufs + chunk_slot * H;
+  };
+
+  if (wid == 0 && elect_one_sync()) {
+  #pragma unroll
+    for (int i = 0; i < CHUNK_DEPTH; i++) {
+      mbarrier_init(plan.bar_ready[i], 1);
+      mbarrier_init(plan.bar_consumed[i], CONSUMER_WARPS);
+    }
+    if constexpr (OUTPUT_NORM_IN_SMEM) {
+      mbarrier_init(plan.bar_output_norm_ready, 1);
+    }
+    fence_mbarrier_init();
+  }
+
+  // gdc wait BEFORE tmem alloc
+  cudaGridDependencySynchronize();
+
+  if (wid == 1) {
+    tmem_allocate(TMEM_COLS_ALLOC, &plan.tmem_base);
+    if constexpr (RELEASE_TMEM) {
+      tmem_release_allocation_lock();
+    }
+  }
+  __syncthreads();
+
+  if constexpr (OUTPUT_NORM_IN_SMEM) {
+    if (wid == 0 && elect_one_sync()) {
+      mbarrier_expect_tx(plan.bar_output_norm_ready, H * (int)sizeof(bf16_t));
+      cp_async_bulk(output_norm_buf, output_norm_weight, H * sizeof(bf16_t),
+                    plan.bar_output_norm_ready);
+    }
+  }
+
+  const uint32_t my_v_tmem =
+      comp_tid >= 0 ? plan.tmem_base + group * TMEM_V_COLS_PER_GROUP : 0;
+  float q_cache[ACC_PER_THREAD];
+  if (comp_tid >= 0) {
+  #pragma unroll
+    for (int si = 0; si < SLICES_PER_GROUP; si++) {
+      if constexpr (H == 7168) {
+        if (si == SLICES_PER_GROUP - 1) {
+          int h_base = 6 * K_TILE + group * (K_TILE / 2) + ct_in_group * 4;
+  #pragma unroll
+          for (int j = 0; j < 4; j++) {
+            int h = h_base + j;
+            q_cache[si * VEC + j] =
+                __bfloat162float(rms_w[h]) * __bfloat162float(res_w[h]);
+          }
+          continue;
+        }
+      }
+      int dt = si * CONSUMER_GROUPS + group;
+      if (dt >= NHT) continue;
+      int h_base = dt * K_TILE + k_local;
+  #pragma unroll
+      for (int j = 0; j < VEC; j++) {
+        int h = h_base + j;
+        q_cache[si * VEC + j] =
+            __bfloat162float(rms_w[h]) * __bfloat162float(res_w[h]);
+      }
+    }
+  }
+
+  if (wid == 0) {
+    if (elect_one_sync()) {
+      long long gci = 0;
+      for (int tb = blockIdx.x; tb < TB; tb += num_ctas) {
+        const int t = tb / B;
+        for (int ci = 0; ci < num_chunks; ci++, gci++) {
+          int ns = ci * N_CHUNK;
+          int an = min(N_CHUNK, N - ns);
+          int chunk_slot = (int)(gci % CHUNK_DEPTH);
+          int pc = phase_of(gci);
+          mbarrier_wait(plan.bar_consumed[chunk_slot], pc ^ 1);
+          int transaction_bytes = an * H * (int)sizeof(bf16_t);
+          if constexpr (HAS_DELTA) {
+            int prefix_n = N - 1 - ns;
+            if (prefix_n >= 0 && prefix_n < an) {
+              transaction_bytes += H * (int)sizeof(bf16_t);
+            }
+          }
+          mbarrier_expect_tx(plan.bar_ready[chunk_slot], transaction_bytes);
+  #pragma unroll
+          for (int n = 0; n < N_CHUNK; n++) {
+            if (n >= an) continue;
+            int slot = slot_of(gci, n);
+            const bf16_t* src =
+                residual_addr(block_res, layer_res, ns + n, N, t,
+                              block_stride_m, block_stride_r, H);
+            cp_async_bulk(buf_ptr(slot), src, H * sizeof(bf16_t),
+                          plan.bar_ready[chunk_slot]);
+          }
+          if constexpr (HAS_DELTA) {
+            int prefix_n = N - 1 - ns;
+            if (prefix_n >= 0 && prefix_n < an) {
+              cp_async_bulk(delta_buf_ptr(chunk_slot),
+                            delta + (long long)tb * H, H * sizeof(bf16_t),
+                            plan.bar_ready[chunk_slot]);
+            }
+          }
+        }
+      }
+    }
+  } else {
+    float acc32[ACC_PER_THREAD] = {};
+    float eps_cache;
+    asm volatile("mov.b32 %0, %1;" : "=f"(eps_cache) : "f"(rms_eps));
+
+    long long gci = 0;
+    for (int tb = blockIdx.x; tb < TB; tb += num_ctas) {
+      float m_running = -FLT_MAX;
+      float s_running = 0.f;
+  #pragma unroll
+      for (int i = 0; i < ACC_PER_THREAD; i++) {
+        acc32[i] = 0.f;
+      }
+
+      for (int ci = 0; ci < num_chunks; ci++, gci++) {
+        int ns = ci * N_CHUNK;
+        int an = min(N_CHUNK, N - ns);
+        int chunk_slot = (int)(gci % CHUNK_DEPTH);
+        int pr = phase_of(gci);
+        mbarrier_wait(plan.bar_ready[chunk_slot], pr);
+
+        float2 sq_local[N_CHUNK] = {};
+        float2 dot_local[N_CHUNK] = {};
+
+        auto pass_A_body = [&](auto AN_TOK) {
+          constexpr int AN = decltype(AN_TOK)::value;
+  #pragma unroll
+          for (int si = 0; si < SLICES_PER_GROUP; si++) {
+            if constexpr (H == 7168) {
+              if (si == SLICES_PER_GROUP - 1) {
+                int h_base =
+                    6 * K_TILE + group * (K_TILE / 2) + ct_in_group * 4;
+                const float* qv = &q_cache[si * VEC];
+  #pragma unroll
+                for (int n = 0; n < AN; n++) {
+                  int slot = slot_of(gci, n);
+                  int2 vp =
+                      *reinterpret_cast<const int2*>(buf_ptr(slot) + h_base);
+                  auto* v2 = reinterpret_cast<__nv_bfloat162*>(&vp);
+                  if constexpr (HAS_DELTA) {
+                    int prefix_n = N - 1 - ns;
+                    if (n == prefix_n) {
+                      const bf16_t* delta_ptr =
+                          delta_buf_ptr(chunk_slot) + h_base;
+  #pragma unroll
+                      for (int j = 0; j < 2; j++) {
+                        auto delta2 = *reinterpret_cast<const __nv_bfloat162*>(
+                            delta_ptr + 2 * j);
+                        v2[j] = __hadd2(v2[j], delta2);
+                      }
+                      *reinterpret_cast<int2*>(layer_res + (long long)tb * H +
+                                               h_base) = vp;
+                    }
+                  }
+                  float2 f[2] = {__bfloat1622float2(v2[0]),
+                                 __bfloat1622float2(v2[1])};
+                  tmem_store<4>(my_v_tmem + (si * N_CHUNK + n) * VEC, f);
+                  sq_local[n] = float2_fma(f[0], f[0], sq_local[n]);
+                  sq_local[n] = float2_fma(f[1], f[1], sq_local[n]);
+                  dot_local[n] =
+                      float2_fma(f[0], make_float2(qv[0], qv[1]), dot_local[n]);
+                  dot_local[n] =
+                      float2_fma(f[1], make_float2(qv[2], qv[3]), dot_local[n]);
+                }
+                continue;
+              }
+            }
+            int dt = si * CONSUMER_GROUPS + group;
+            if (dt >= NHT) continue;
+            int h_base = dt * K_TILE + k_local;
+            const float* qv = &q_cache[si * VEC];
+
+  #pragma unroll
+            for (int n = 0; n < AN; n++) {
+              int slot = slot_of(gci, n);
+              int4 vp = *reinterpret_cast<const int4*>(buf_ptr(slot) + h_base);
+              auto* v2 = reinterpret_cast<__nv_bfloat162*>(&vp);
+              if constexpr (HAS_DELTA) {
+                int prefix_n = N - 1 - ns;
+                if (n == prefix_n) {
+                  const bf16_t* delta_ptr = delta_buf_ptr(chunk_slot) + h_base;
+  #pragma unroll
+                  for (int j = 0; j < VEC / 2; j++) {
+                    auto delta2 = *reinterpret_cast<const __nv_bfloat162*>(
+                        delta_ptr + 2 * j);
+                    v2[j] = __hadd2(v2[j], delta2);
+                  }
+                  *reinterpret_cast<int4*>(layer_res + (long long)tb * H +
+                                           h_base) = vp;
+                }
+              }
+              float2 f[4] = {
+                  __bfloat1622float2(v2[0]), __bfloat1622float2(v2[1]),
+                  __bfloat1622float2(v2[2]), __bfloat1622float2(v2[3])};
+              tmem_store<VEC>(my_v_tmem + (si * N_CHUNK + n) * VEC, f);
+  #pragma unroll
+              for (int j = 0; j < VEC / 2; j++) {
+                sq_local[n] = float2_fma(f[j], f[j], sq_local[n]);
+                dot_local[n] = float2_fma(
+                    f[j], make_float2(qv[2 * j], qv[2 * j + 1]), dot_local[n]);
+              }
+            }
+          }
+        };
+        if constexpr (NC == 4) {
+          switch (an) {
+            case 4:
+              pass_A_body(std::integral_constant<int, 4>{});
+              break;
+            case 3:
+              pass_A_body(std::integral_constant<int, 3>{});
+              break;
+            case 2:
+              pass_A_body(std::integral_constant<int, 2>{});
+              break;
+            case 1:
+              pass_A_body(std::integral_constant<int, 1>{});
+              break;
+            default:
+              __builtin_unreachable();
+          }
+        } else if constexpr (NC == 3) {
+          switch (an) {
+            case 3:
+              pass_A_body(std::integral_constant<int, 3>{});
+              break;
+            case 2:
+              pass_A_body(std::integral_constant<int, 2>{});
+              break;
+            case 1:
+              pass_A_body(std::integral_constant<int, 1>{});
+              break;
+            default:
+              __builtin_unreachable();
+          }
+        } else {
+          static_assert(NC == 2);
+          switch (an) {
+            case 2:
+              pass_A_body(std::integral_constant<int, 2>{});
+              break;
+            case 1:
+              pass_A_body(std::integral_constant<int, 1>{});
+              break;
+            default:
+              __builtin_unreachable();
+          }
+        }
+        if (lane == 0) {
+          mbarrier_arrive(plan.bar_consumed[chunk_slot]);
+        }
+        tmem_store_wait();
+
+        float2 reduce_pair[N_CHUNK];
+  #pragma unroll
+        for (int n = 0; n < N_CHUNK; n++) {
+          reduce_pair[n] = make_float2(sq_local[n].x + sq_local[n].y,
+                                       dot_local[n].x + dot_local[n].y);
+        }
+  #pragma unroll
+        for (int offset = 16; offset > 0; offset >>= 1) {
+  #pragma unroll
+          for (int n = 0; n < N_CHUNK; n++) {
+            uint64_t packed = reinterpret_cast<uint64_t&>(reduce_pair[n]);
+            packed = __shfl_xor_sync(0xffffffff, packed, offset);
+            float2 other = reinterpret_cast<float2&>(packed);
+            reduce_pair[n] = float2_add(reduce_pair[n], other);
+          }
+        }
+        if (lane == 0) {
+  #pragma unroll
+          for (int n = 0; n < N_CHUNK; n++) {
+            plan.ws_stats[comp_wid][n] = reduce_pair[n];
+          }
+        }
+        named_barrier_sync(CONSUMER_THREADS, 0);
+
+        float local_rsig = 0.f;
+        float local_logit = 0.f;
+        int stat_n = lane / CONSUMER_WARPS;
+        int stat_w = lane % CONSUMER_WARPS;
+        float2 totals = {};
+        if (stat_n < N_CHUNK) {
+          totals = plan.ws_stats[stat_w][stat_n];
+        }
+  #pragma unroll
+        for (int offset = CONSUMER_WARPS / 2; offset > 0; offset >>= 1) {
+          totals.x +=
+              __shfl_down_sync(0xffffffff, totals.x, offset, CONSUMER_WARPS);
+          totals.y +=
+              __shfl_down_sync(0xffffffff, totals.y, offset, CONSUMER_WARPS);
+        }
+        if (stat_n < N_CHUNK && stat_w == 0) {
+          local_rsig = rsqrtf(totals.x / H + eps_cache);
+          local_logit = totals.y * local_rsig;
+        }
+        float logit_n[N_CHUNK];
+  #pragma unroll
+        for (int n = 0; n < N_CHUNK; n++) {
+          logit_n[n] = __shfl_sync(0xffffffff, local_logit, n * CONSUMER_WARPS);
+        }
+
+        float m_chunk = -FLT_MAX;
+  #pragma unroll
+        for (int n = 0; n < N_CHUNK; n++) {
+          if (n < an) m_chunk = fmaxf(m_chunk, logit_n[n]);
+        }
+        float m_new = fmaxf(m_running, m_chunk);
+        float corr = exp2f((m_running - m_new) * LOG2_E);
+        float w_n[N_CHUNK] = {};
+        float w_sum = 0.f;
+  #pragma unroll
+        for (int n = 0; n < N_CHUNK; n++) {
+          if (n < an) {
+            w_n[n] = exp2f((logit_n[n] - m_new) * LOG2_E);
+            w_sum += w_n[n];
+          }
+        }
+
+        auto pass_B_body = [&](auto AN_TOK) {
+          constexpr int AN = decltype(AN_TOK)::value;
+  #pragma unroll
+          for (int si = 0; si < SLICES_PER_GROUP; si++) {
+            if constexpr (H == 7168) {
+              if (si == SLICES_PER_GROUP - 1) {
+                float2 corr2 = make_float2(corr, corr);
+                float2 a[2];
+  #pragma unroll
+                for (int j = 0; j < 2; j++) {
+                  float2 old = make_float2(acc32[si * VEC + 2 * j],
+                                           acc32[si * VEC + 2 * j + 1]);
+                  a[j] = float2_mul(old, corr2);
+                }
+                float2 f_cache[AN][2];
+  #pragma unroll
+                for (int n = 0; n < AN; n++) {
+                  tmem_load<4>(my_v_tmem + (si * N_CHUNK + n) * VEC,
+                               f_cache[n]);
+                }
+  #pragma unroll
+                for (int n = 0; n < AN; n++) {
+                  float2 wn = make_float2(w_n[n], w_n[n]);
+  #pragma unroll
+                  for (int j = 0; j < 2; j++) {
+                    a[j] = float2_fma(wn, f_cache[n][j], a[j]);
+                  }
+                }
+  #pragma unroll
+                for (int j = 0; j < 2; j++) {
+                  acc32[si * VEC + 2 * j] = a[j].x;
+                  acc32[si * VEC + 2 * j + 1] = a[j].y;
+                }
+                continue;
+              }
+            }
+            int dt = si * CONSUMER_GROUPS + group;
+            if (dt >= NHT) continue;
+            float2 corr2 = make_float2(corr, corr);
+            float2 a[VEC / 2];
+  #pragma unroll
+            for (int j = 0; j < VEC / 2; j++) {
+              float2 old = make_float2(acc32[si * VEC + 2 * j],
+                                       acc32[si * VEC + 2 * j + 1]);
+              a[j] = float2_mul(old, corr2);
+            }
+            float2 f_cache[AN][VEC / 2];
+  #pragma unroll
+            for (int n = 0; n < AN; n++) {
+              tmem_load<VEC>(my_v_tmem + (si * N_CHUNK + n) * VEC, f_cache[n]);
+            }
+  #pragma unroll
+            for (int n = 0; n < AN; n++) {
+              float2 wn = make_float2(w_n[n], w_n[n]);
+  #pragma unroll
+              for (int j = 0; j < VEC / 2; j++) {
+                a[j] = float2_fma(wn, f_cache[n][j], a[j]);
+              }
+            }
+  #pragma unroll
+            for (int j = 0; j < VEC / 2; j++) {
+              acc32[si * VEC + 2 * j] = a[j].x;
+              acc32[si * VEC + 2 * j + 1] = a[j].y;
+            }
+          }
+        };
+        if constexpr (NC == 4) {
+          switch (an) {
+            case 4:
+              pass_B_body(std::integral_constant<int, 4>{});
+              break;
+            case 3:
+              pass_B_body(std::integral_constant<int, 3>{});
+              break;
+            case 2:
+              pass_B_body(std::integral_constant<int, 2>{});
+              break;
+            case 1:
+              pass_B_body(std::integral_constant<int, 1>{});
+              break;
+            default:
+              __builtin_unreachable();
+          }
+        } else if constexpr (NC == 3) {
+          switch (an) {
+            case 3:
+              pass_B_body(std::integral_constant<int, 3>{});
+              break;
+            case 2:
+              pass_B_body(std::integral_constant<int, 2>{});
+              break;
+            case 1:
+              pass_B_body(std::integral_constant<int, 1>{});
+              break;
+            default:
+              __builtin_unreachable();
+          }
+        } else {
+          static_assert(NC == 2);
+          switch (an) {
+            case 2:
+              pass_B_body(std::integral_constant<int, 2>{});
+              break;
+            case 1:
+              pass_B_body(std::integral_constant<int, 1>{});
+              break;
+            default:
+              __builtin_unreachable();
+          }
+        }
+
+        s_running = s_running * corr + w_sum;
+        m_running = m_new;
+      }
+
+      float inv_s = 1.f / s_running;
+      bf16_t* out_ptr = output + (long long)tb * H;
+      float2 output_sq_pair = {};
+      // When output RMSNorm is fused, the softmax denominator cancels:
+      // (acc / s) * rsqrt(mean((acc / s)^2) + eps)
+      //   = acc * rsqrt(mean(acc^2) + eps * s^2).
+  #pragma unroll
+      for (int si = 0; si < SLICES_PER_GROUP; si++) {
+        if constexpr (H == 7168) {
+          if (si == SLICES_PER_GROUP - 1) {
+            int h_base = 6 * K_TILE + group * (K_TILE / 2) + ct_in_group * 4;
+            uint2 packed;
+            auto* ov2 = reinterpret_cast<__nv_bfloat162*>(&packed);
+            float2 inv2 = make_float2(inv_s, inv_s);
+  #pragma unroll
+            for (int j = 0; j < 2; j++) {
+              float2 old = make_float2(acc32[si * VEC + 2 * j],
+                                       acc32[si * VEC + 2 * j + 1]);
+              if constexpr (HAS_OUTPUT_NORM) {
+                output_sq_pair = float2_fma(old, old, output_sq_pair);
+              } else {
+                float2 mixed = float2_mul(old, inv2);
+                ov2[j] = __float22bfloat162_rn(mixed);
+              }
+            }
+            if constexpr (!HAS_OUTPUT_NORM) {
+              *reinterpret_cast<uint2*>(out_ptr + h_base) = packed;
+            }
+            continue;
+          }
+        }
+        int dt = si * CONSUMER_GROUPS + group;
+        if (dt >= NHT) continue;
+        int h_base = dt * K_TILE + k_local;
+        uint4 packed;
+        auto* ov2 = reinterpret_cast<__nv_bfloat162*>(&packed);
+        float2 inv2 = make_float2(inv_s, inv_s);
+  #pragma unroll
+        for (int j = 0; j < VEC / 2; j++) {
+          float2 old =
+              make_float2(acc32[si * VEC + 2 * j], acc32[si * VEC + 2 * j + 1]);
+          if constexpr (HAS_OUTPUT_NORM) {
+            output_sq_pair = float2_fma(old, old, output_sq_pair);
+          } else {
+            float2 mixed = float2_mul(old, inv2);
+            ov2[j] = __float22bfloat162_rn(mixed);
+          }
+        }
+        if constexpr (!HAS_OUTPUT_NORM) {
+          *reinterpret_cast<uint4*>(out_ptr + h_base) = packed;
+        }
+      }
+
+      if constexpr (HAS_OUTPUT_NORM) {
+        if constexpr (OUTPUT_NORM_IN_SMEM) {
+          // The immutable weight copy is acquired once, at its first use.
+          if (tb == blockIdx.x) {
+            mbarrier_wait(plan.bar_output_norm_ready, 0);
+          }
+        }
+        float output_sq = output_sq_pair.x + output_sq_pair.y;
+  #pragma unroll
+        for (int offset = 16; offset > 0; offset >>= 1) {
+          output_sq += __shfl_xor_sync(0xffffffff, output_sq, offset);
+        }
+        if (lane == 0) {
+          plan.ws_stats[comp_wid][0] = make_float2(output_sq, 0.f);
+        }
+        named_barrier_sync(CONSUMER_THREADS, 0);
+        float total_sq = lane < CONSUMER_WARPS ? plan.ws_stats[lane][0].x : 0.f;
+  #pragma unroll
+        for (int offset = CONSUMER_WARPS / 2; offset > 0; offset >>= 1) {
+          total_sq +=
+              __shfl_down_sync(0xffffffff, total_sq, offset, CONSUMER_WARPS);
+        }
+        if (lane == 0) {
+          total_sq =
+              rsqrtf(total_sq / H + output_norm_eps * s_running * s_running);
+        }
+        float output_rsigma = __shfl_sync(0xffffffff, total_sq, 0);
+  #pragma unroll
+        for (int si = 0; si < SLICES_PER_GROUP; si++) {
+          if constexpr (H == 7168) {
+            if (si == SLICES_PER_GROUP - 1) {
+              int h_base = 6 * K_TILE + group * (K_TILE / 2) + ct_in_group * 4;
+              uint2 packed;
+              auto* values = reinterpret_cast<bf16_t*>(&packed);
+  #pragma unroll
+              for (int j = 0; j < 4; j++) {
+                const bf16_t* weight_ptr =
+                    OUTPUT_NORM_IN_SMEM ? output_norm_buf : output_norm_weight;
+                float weight = __bfloat162float(weight_ptr[h_base + j]);
+                values[j] = __float2bfloat16(acc32[si * VEC + j] *
+                                             output_rsigma * weight);
+              }
+              *reinterpret_cast<uint2*>(out_ptr + h_base) = packed;
+              continue;
+            }
+          }
+          int dt = si * CONSUMER_GROUPS + group;
+          if (dt >= NHT) continue;
+          int h_base = dt * K_TILE + k_local;
+          uint4 packed;
+          auto* values = reinterpret_cast<bf16_t*>(&packed);
+  #pragma unroll
+          for (int j = 0; j < VEC; j++) {
+            const bf16_t* weight_ptr =
+                OUTPUT_NORM_IN_SMEM ? output_norm_buf : output_norm_weight;
+            float weight = __bfloat162float(weight_ptr[h_base + j]);
+            values[j] =
+                __float2bfloat16(acc32[si * VEC + j] * output_rsigma * weight);
+          }
+          *reinterpret_cast<uint4*>(out_ptr + h_base) = packed;
+        }
+      }
+    }
+  }
+
+  cudaTriggerProgrammaticLaunchCompletion();
+  __syncthreads();
+  if (wid == 1) {
+    tmem_free(plan.tmem_base, TMEM_COLS_ALLOC);
+  }
+#else
+  if (threadIdx.x == 0) {
+    printf("attn_res_fwd_online_v2_kernel requires sm_10x\n");
+  }
+#endif
+}
+
+template <int H, int NC = N_CHUNK_DEFAULT, bool RELEASE_TMEM = false,
+          bool HAS_DELTA = false, bool HAS_OUTPUT_NORM = false,
+          bool OUTPUT_NORM_IN_SMEM = false>
+static void launch_fwd(const bf16_t* block_residual, bf16_t* layer_residual,
+                       const bf16_t* delta, const bf16_t* res_weight,
+                       const bf16_t* rms_weight, bf16_t* output, int N, int T,
+                       int B, float rms_eps, int num_sm, cudaStream_t stream,
+                       const bf16_t* output_norm_weight = nullptr,
+                       float output_norm_eps = 0.f, int block_stride_m = 0,
+                       int block_stride_r = 0) {
+  constexpr size_t smem_size =
+      ((size_t)CHUNK_DEPTH * (NC + (HAS_DELTA ? 1 : 0)) * H * sizeof(bf16_t) +
+       (OUTPUT_NORM_IN_SMEM ? (size_t)H * sizeof(bf16_t) : 0) +
+       sizeof(FwdSmemPlan<NC>) + 15) &
+      ~size_t(15);
+  auto kernel =
+      &attn_res_fwd_online_v2_kernel<H, NC, RELEASE_TMEM, HAS_DELTA,
+                                     HAS_OUTPUT_NORM, OUTPUT_NORM_IN_SMEM>;
+  static bool attrs_set = false;
+  if (!attrs_set) {
+    if (smem_size > 48 * 1024) {
+      cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize,
+                           smem_size);
+    }
+    attrs_set = true;
+  }
+  int grid = RELEASE_TMEM ? num_sm * 2 : num_sm;
+  cudaLaunchConfig_t config{};
+  config.gridDim = grid;
+  config.blockDim = BLK;
+  config.dynamicSmemBytes = smem_size;
+  config.stream = stream;
+  cudaLaunchAttribute attrs[1];
+  attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization;
+  attrs[0].val.programmaticStreamSerializationAllowed = 1;
+  config.attrs = attrs;
+  config.numAttrs = 1;
+  cudaLaunchKernelEx(&config, kernel, block_residual, layer_residual, delta,
+                     res_weight, rms_weight, output, N, T, B, block_stride_m,
+                     block_stride_r, rms_eps, output_norm_weight,
+                     output_norm_eps);
+}
+
+}  // namespace fwd_prod_v2
+}  // namespace sm100
+
+void kimi_k3_attn_res(torch::stable::Tensor& prefix,
+                      torch::stable::Tensor const& delta,
+                      torch::stable::Tensor const& blocks,
+                      torch::stable::Tensor const& norm_weight,
+                      torch::stable::Tensor const& qk_weight,
+                      torch::stable::Tensor const& output_norm_weight,
+                      torch::stable::Tensor& output, int64_t num_blocks,
+                      double eps, double output_norm_eps) {
+  int const num_tokens = static_cast<int>(prefix.size(0));
+  int const device = prefix.get_device_index();
+  torch::stable::accelerator::DeviceGuard const device_guard(device);
+  cudaDeviceProp const* properties = get_device_prop();
+  STD_TORCH_CHECK(properties->major == 10,
+                  "Kimi K3 AttnRes requires the SM100 family");
+
+  using namespace sm100::fwd_prod_v2;
+  // Two-source chunks and two resident CTAs are beneficial once setup is
+  // amortized by the long, full eight-block prefill workload.
+  if (num_blocks == 8 && num_tokens >= 4096) {
+    launch_fwd<7168, 2, true, true, true, true>(
+        static_cast<bf16_t const*>(blocks.data_ptr()),
+        static_cast<bf16_t*>(prefix.data_ptr()),
+        static_cast<bf16_t const*>(delta.data_ptr()),
+        static_cast<bf16_t const*>(qk_weight.data_ptr()),
+        static_cast<bf16_t const*>(norm_weight.data_ptr()),
+        static_cast<bf16_t*>(output.data_ptr()),
+        static_cast<int>(num_blocks) + 1, num_tokens, 1,
+        static_cast<float>(eps), properties->multiProcessorCount,
+        get_current_cuda_stream(device),
+        static_cast<bf16_t const*>(output_norm_weight.data_ptr()),
+        static_cast<float>(output_norm_eps), static_cast<int>(blocks.stride(0)),
+        static_cast<int>(blocks.stride(1)));
+  } else {
+    launch_fwd<7168, 4, false, true, true, true>(
+        static_cast<bf16_t const*>(blocks.data_ptr()),
+        static_cast<bf16_t*>(prefix.data_ptr()),
+        static_cast<bf16_t const*>(delta.data_ptr()),
+        static_cast<bf16_t const*>(qk_weight.data_ptr()),
+        static_cast<bf16_t const*>(norm_weight.data_ptr()),
+        static_cast<bf16_t*>(output.data_ptr()),
+        static_cast<int>(num_blocks) + 1, num_tokens, 1,
+        static_cast<float>(eps), properties->multiProcessorCount,
+        get_current_cuda_stream(device),
+        static_cast<bf16_t const*>(output_norm_weight.data_ptr()),
+        static_cast<float>(output_norm_eps), static_cast<int>(blocks.stride(0)),
+        static_cast<int>(blocks.stride(1)));
+  }
+  cudaError_t const error = cudaGetLastError();
+  STD_TORCH_CHECK(
+      error == cudaSuccess,
+      "Kimi K3 AttnRes kernel launch failed: ", cudaGetErrorString(error));
+}

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,30 @@ 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,
+                      torch::stable::Tensor const& blocks,
+                      torch::stable::Tensor const& norm_weight,
+                      torch::stable::Tensor const& qk_weight,
+                      torch::stable::Tensor const& output_norm_weight,
+                      torch::stable::Tensor& output, int64_t num_blocks,
+                      double eps, double output_norm_eps);
+#endif
+
 // Sampler kernels (shared CUDA/ROCm)
 void apply_repetition_penalties_(
     torch::stable::Tensor& logits, const torch::stable::Tensor& prompt_mask,
@@ -372,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);
@@ -410,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);
@@ -486,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,24 @@ 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("
+      "Tensor! prefix, Tensor delta, Tensor blocks, Tensor norm_weight, "
+      "Tensor qk_weight, Tensor output_norm_weight, Tensor! output, "
+      "int num_blocks, float eps, float output_norm_eps) -> ()");
+#endif
+
   // Apply repetition penalties to logits in-place.
   ops.def(
       "apply_repetition_penalties_(Tensor! logits, Tensor prompt_mask, "
@@ -530,6 +591,14 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
       "limit=7.0) "
       "-> ()");
 
+  // Kimi SITU (SituGLU) gated activation. linear_beta<=0 means unset.
+  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) -> ()");
 
@@ -688,11 +757,30 @@ 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
 
   // Sampler kernels (shared CUDA/ROCm)
   ops.impl("apply_repetition_penalties_",
@@ -715,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));
@@ -812,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("
@@ -910,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>(

docs/design/attention_backends.md

@@ -222,7 +222,7 @@ MLA decode backends are selected using the standard
 | Backend | Dtypes | KV Dtypes | Block Sizes | Head Sizes | Sink | Non-Causal | Sparse | MM Prefix | DCP | Attention Types | Compute Cap. |
 | ------- | ------ | --------- | ----------- | ---------- | ---- | ---------- | ------ | --------- | --- | --------------- | ------------ |
 | `CUTLASS_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 128 | Any | ❌ | ❌ | ❌ | ❌ | ✅ | Decoder | 10.x |
-| `FLASHINFER_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 32, 64 | Any | ❌ | ❌ | ❌ | ❌ | ✅ | Decoder | 10.x |
+| `FLASHINFER_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 32, 64 | Any | ❌ | ✅ | ❌ | ❌ | ✅ | Decoder | 10.x |
 | `FLASHINFER_MLA_SPARSE` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 32, 64 | Any | ❌ | ❌ | ❌ | ❌ | ✅ | Decoder | 10.x |
 | `FLASHINFER_MLA_SPARSE_SM120` | bf16 | `auto`, `fp8`, `fp8_e4m3`, `fp8_ds_mla` | 64, 256 | Any | ❌ | ❌ | ❌ | ❌ | ❌ | Decoder | 12.x |
 | `FLASHMLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 64 | Any | ❌ | ❌ | ❌ | ❌ | ✅ | Decoder | 9.x-10.x |

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. | ✅︎ | ✅︎ |

requirements/cuda.txt

@@ -17,7 +17,7 @@ PyNvVideoCodec==2.0.4
 flashinfer-python==0.6.15.post1
 flashinfer-cubin==0.6.15.post1
 apache-tvm-ffi==0.1.10
-tilelang==0.1.9
+tilelang==0.1.12
 nvidia-cudnn-frontend>=1.19.1
 # Required for LLM_NVTX_SCOPES_FOR_PROFILING=1
 nvtx==0.2.15
@@ -33,3 +33,6 @@ tokenspeed-mla==0.1.8; platform_system == "Linux"
 
 # Humming kernels for quantization gemm
 humming-kernels[cu13]==0.1.10
+
+# KDA
+flash-linear-attention==0.5.0

rust/Cargo.lock

@@ -2220,7 +2220,7 @@ checksum = "11d3d7f243d5c5a8b9bb5d6dd2b1602c0cb0b9db1621bafc7ed66e35ff9fe092"
 [[package]]
 name = "llm-multimodal"
 version = "1.7.1"
-source = "git+https://github.com/smg-project/llm-multimodal?rev=5390032d6dc8a3e6fdc83acd320260367eb4b9b5#5390032d6dc8a3e6fdc83acd320260367eb4b9b5"
+source = "git+ssh://git@github.com/Inferact/llm-multimodal-internal.git?branch=k3-image#ceec43ec6beea5812d5ae59af712d9ff616496ef"
 dependencies = [
  "anyhow",
  "base64 0.22.1",
@@ -2235,7 +2235,6 @@ dependencies = [
  "once_cell",
  "pkg-config",
  "rayon",
- "realfft",
  "reqwest 0.13.4",
  "rustfft",
  "serde",

rust/Cargo.toml

@@ -58,7 +58,7 @@ indexmap = "2.13.0"
 indicatif = "0.18.4"
 itertools = "0.14.0"
 libc = "0.2.177"
-llm-multimodal = { git = "https://github.com/smg-project/llm-multimodal", rev = "5390032d6dc8a3e6fdc83acd320260367eb4b9b5", default-features = false, features = ["native-tls"] }
+llm-multimodal = { git = "ssh://git@github.com/Inferact/llm-multimodal-internal.git", branch = "k3-image", default-features = false, features = ["native-tls"] }
 mimalloc = "0.1.52"
 minijinja = { version = "2.0", features = ["unstable_machinery", "json", "builtins", "loader", "loop_controls", "preserve_order"] }
 minijinja-contrib = { version = "2.0", features = ["pycompat"] }

rust/src/chat/src/backend/hf.rs

@@ -20,7 +20,7 @@ use crate::output::{
 use crate::renderer::hf::{HfChatRenderer, MultimodalRenderInfo};
 use crate::renderer::{
     DeepSeekV4ChatRenderer, DeepSeekV32ChatRenderer, DynChatRenderer, HarmonyChatRenderer,
-    InklingChatRenderer,
+    InklingChatRenderer, KimiK3ChatRenderer,
 };
 use crate::request::ChatRequest;
 use crate::{DynChatOutputProcessor, RendererSelection};
@@ -73,6 +73,7 @@ impl HfChatBackend {
             RendererSelection::DeepSeekV4 => Arc::new(DeepSeekV4ChatRenderer::new()),
             RendererSelection::Harmony => Arc::new(HarmonyChatRenderer::new()?),
             RendererSelection::Inkling => Arc::new(InklingChatRenderer::new(tokenizer.clone())?),
+            RendererSelection::KimiK3 => Arc::new(KimiK3ChatRenderer::new()),
         };
 
         info!(

rust/src/chat/src/lib.rs

@@ -33,7 +33,8 @@ pub use parser::tool::{ToolParser, ToolParserError, ToolParserFactory};
 pub use renderer::hf::ChatTemplateContentFormatOption;
 pub use renderer::{
     ChatRenderer, DeepSeekV4ChatRenderer, DeepSeekV32ChatRenderer, DynChatRenderer,
-    HarmonyChatRenderer, InklingChatRenderer, RenderedPrompt, RendererSelection,
+    HarmonyChatRenderer, InklingChatRenderer, KimiK3ChatRenderer, RenderedPrompt,
+    RendererSelection,
 };
 pub use request::{
     ChatContent, ChatContentPart, ChatMessage, ChatOptions, ChatRequest, ChatRole, ChatTool,
@@ -326,6 +327,12 @@ mod tests {
         .unwrap();
     }
 
+    #[test]
+    fn validate_parser_overrides_accepts_explicit_kimi_k3() {
+        let selection = ParserSelection::Explicit("kimi_k3".to_string());
+        validate_parser_overrides(&selection, &selection).unwrap();
+    }
+
     #[test]
     fn validate_parser_overrides_accepts_auto_and_none() {
         validate_parser_overrides(&ParserSelection::Auto, &ParserSelection::None).unwrap();
@@ -339,7 +346,7 @@ mod tests {
         )
         .unwrap_err();
 
-        expect_test::expect!["tool parser `definitely_missing_tool_parser` is not registered (choose from: deepseek_v3, deepseek_v31, deepseek_v32, deepseek_v4, gemma4, glm45, glm47, granite4, hermes, hy_v3, inkling, internlm, kimi_k2, llama3_json, llama4_json, minimax_m2, minimax_m3, mistral, phi4_mini_json, qwen3_coder, qwen3_xml, seed_oss)"].assert_eq(&error.to_report_string());
+        expect_test::expect!["tool parser `definitely_missing_tool_parser` is not registered (choose from: deepseek_v3, deepseek_v31, deepseek_v32, deepseek_v4, gemma4, glm45, glm47, granite4, hermes, hy_v3, inkling, internlm, kimi_k2, kimi_k3, llama3_json, llama4_json, minimax_m2, minimax_m3, mistral, phi4_mini_json, qwen3_coder, qwen3_xml, seed_oss)"].assert_eq(&error.to_report_string());
     }
 
     #[test]
@@ -350,6 +357,6 @@ mod tests {
         )
         .unwrap_err();
 
-        expect_test::expect!["reasoning parser `definitely_missing_reasoning_parser` is not registered (choose from: cohere_cmd, deepseek_r1, deepseek_v3, deepseek_v4, gemma4, glm45, inkling, kimi, kimi_k2, minimax_m2, minimax_m3, nemotron_v3, qwen3, seed_oss, step3, step3p5)"].assert_eq(&error.to_report_string());
+        expect_test::expect!["reasoning parser `definitely_missing_reasoning_parser` is not registered (choose from: cohere_cmd, deepseek_r1, deepseek_v3, deepseek_v4, gemma4, glm45, inkling, kimi, kimi_k2, kimi_k3, minimax_m2, minimax_m3, nemotron_v3, qwen3, seed_oss, step3, step3p5)"].assert_eq(&error.to_report_string());
     }
 }

rust/src/chat/src/parser/reasoning/mod.rs

@@ -27,6 +27,7 @@ pub mod names {
     pub const GLM45: &str = "glm45";
     pub const KIMI: &str = "kimi";
     pub const KIMI_K2: &str = "kimi_k2";
+    pub const KIMI_K3: &str = "kimi_k3";
     pub const MINIMAX_M2: &str = "minimax_m2";
     pub const MINIMAX_M3: &str = "minimax_m3";
     pub const NEMOTRON_V3: &str = "nemotron_v3";
@@ -68,6 +69,7 @@ impl ReasoningParserFactory {
             .register_parser::<Glm45ReasoningParser>(names::GLM45)
             .register_parser::<KimiReasoningParser>(names::KIMI)
             .register_parser::<KimiK2ReasoningParser>(names::KIMI_K2)
+            .register_unified_dummy(names::KIMI_K3)
             .register_parser::<MiniMaxM2ReasoningParser>(names::MINIMAX_M2)
             .register_parser::<MiniMaxM3ReasoningParser>(names::MINIMAX_M3)
             .register_parser::<NemotronV3ReasoningParser>(names::NEMOTRON_V3)

rust/src/chat/src/parser/tool/mod.rs

@@ -33,6 +33,7 @@ pub mod names {
     // also routes to `Internlm2ToolParser` despite the version-agnostic name.
     pub const INTERNLM: &str = "internlm";
     pub const KIMI_K2: &str = "kimi_k2";
+    pub const KIMI_K3: &str = "kimi_k3";
     pub const LLAMA3_JSON: &str = "llama3_json";
     pub const LLAMA4_JSON: &str = "llama4_json";
     pub const MINIMAX_M2: &str = "minimax_m2";
@@ -78,6 +79,7 @@ impl ToolParserFactory {
             .register_parser::<HyV3ToolParser>(names::HY_V3)
             .register_parser::<Internlm2ToolParser>(names::INTERNLM)
             .register_parser::<KimiK2ToolParser>(names::KIMI_K2)
+            .register_unified_dummy(names::KIMI_K3)
             .register_parser::<Llama3JsonToolParser>(names::LLAMA3_JSON)
             .register_parser::<Llama3JsonToolParser>(names::LLAMA4_JSON)
             .register_parser::<MinimaxM2ToolParser>(names::MINIMAX_M2)

rust/src/chat/src/parser/unified.rs

@@ -5,7 +5,9 @@
 
 use std::sync::LazyLock;
 
-pub use vllm_parser::unified::{Gemma4UnifiedParser, InklingUnifiedParser, UnifiedParser};
+pub use vllm_parser::unified::{
+    Gemma4UnifiedParser, InklingUnifiedParser, KimiK3UnifiedParser, UnifiedParser,
+};
 use vllm_tokenizer::DynTokenizer;
 
 use crate::parser::ParserFactory;
@@ -15,6 +17,7 @@ use crate::request::ChatTool;
 pub mod names {
     pub const GEMMA4: &str = "gemma4";
     pub const INKLING: &str = "inkling";
+    pub const KIMI_K3: &str = "kimi_k3";
 }
 
 /// Constructor signature for one registered unified parser implementation.
@@ -39,11 +42,14 @@ impl UnifiedParserFactory {
 
         factory.register_parser::<Gemma4UnifiedParser>(names::GEMMA4);
         factory.register_parser::<InklingUnifiedParser>(names::INKLING);
+        factory.register_parser::<KimiK3UnifiedParser>(names::KIMI_K3);
 
         factory
             .register_pattern("gemma-4", names::GEMMA4)
             .register_pattern("gemma4", names::GEMMA4)
-            .register_pattern("inkling", names::INKLING);
+            .register_pattern("inkling", names::INKLING)
+            .register_pattern("kimi-k3", names::KIMI_K3)
+            .register_pattern("kimi_k3", names::KIMI_K3);
 
         factory
     }
@@ -121,4 +127,20 @@ mod tests {
         );
         factory.create(names::INKLING, &[], Arc::new(inkling_tokenizer())).unwrap();
     }
+
+    #[test]
+    fn factory_registers_kimi_k3() {
+        let factory = UnifiedParserFactory::new();
+        let tokenizer = TestTokenizer::new()
+            .with_regular_token("<|open|>", 1001)
+            .with_regular_token("<|close|>", 1002)
+            .with_regular_token("<|sep|>", 1003);
+
+        assert!(factory.contains(names::KIMI_K3));
+        assert_eq!(
+            factory.resolve_name_for_model("moonshotai/Kimi-K3"),
+            Some(names::KIMI_K3)
+        );
+        factory.create(names::KIMI_K3, &[], Arc::new(tokenizer)).unwrap();
+    }
 }

rust/src/chat/src/renderer/deepseek_v32/tests.rs

@@ -59,7 +59,7 @@ fn fixture_request(input_name: &str) -> ChatRequest {
 
 fn deepseek_fixture_options() -> FixtureRequestOptions {
     FixtureRequestOptions {
-        enable_thinking: true,
+        enable_thinking: Some(true),
         no_generation_prompt_when_last_assistant: true,
     }
 }

rust/src/chat/src/renderer/deepseek_v4/tests.rs

@@ -27,7 +27,7 @@ fn fixture_request(input_name: &str) -> ChatRequest {
 
 fn deepseek_fixture_options() -> FixtureRequestOptions {
     FixtureRequestOptions {
-        enable_thinking: true,
+        enable_thinking: Some(true),
         no_generation_prompt_when_last_assistant: true,
     }
 }

rust/src/chat/src/renderer/harmony/tests.rs

@@ -22,7 +22,7 @@ fn fixture_request(input_name: &str) -> ChatRequest {
     fixture_chat_request(
         &fixture_path(input_name),
         FixtureRequestOptions {
-            enable_thinking: false,
+            enable_thinking: Some(false),
             no_generation_prompt_when_last_assistant: false,
         },
     )

rust/src/chat/src/renderer/inkling/tests.rs

@@ -101,7 +101,7 @@ fn fixture_request(name: &str) -> ChatRequest {
 
 fn inkling_fixture_options() -> FixtureRequestOptions {
     FixtureRequestOptions {
-        enable_thinking: false,
+        enable_thinking: Some(false),
         no_generation_prompt_when_last_assistant: false,
     }
 }

rust/src/chat/src/renderer/kimi_k3/encoding.rs

@@ -0,0 +1,568 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+//! Kimi K3 XTML prompt renderer.
+//!
+//! Port of Moonshot remote-code `encoding_k3.py::build_chat_segments()`.
+
+use std::collections::HashMap;
+use std::fmt::Write as _;
+
+use serde_json::{Map, Value, json};
+
+use crate::error::{Error, Result};
+use crate::request::{
+    ChatContent, ChatContentPart, ChatMessage, ChatRequest, ChatTool, ChatToolChoice,
+};
+use crate::{AssistantContentBlock, AssistantToolCall};
+
+pub(super) const OPEN: &str = "<|open|>";
+pub(super) const CLOSE: &str = "<|close|>";
+pub(super) const SEP: &str = "<|sep|>";
+pub(super) const END_OF_MSG: &str = "<|end_of_msg|>";
+pub(super) const IMAGE_PLACEHOLDER: &str = "<|media_pad|>";
+
+const DEFAULT_THINKING_EFFORT: &str = "max";
+const VALID_THINKING_EFFORTS: &[&str] = &["low", "high", "max"];
+
+/// Render one chat request into the K3 XTML prompt string.
+pub(super) fn render_request(request: &ChatRequest) -> Result<String> {
+    let thinking = thinking_enabled(request)?;
+    let thinking_effort = thinking.then(|| thinking_effort(request)).transpose()?;
+    let tools = request_tools(request);
+    let mut out = String::new();
+
+    if !tools.is_empty() {
+        write_tool_declare(&mut out, tools, false)?;
+    }
+
+    if let Some(effort) = thinking_effort {
+        // Preserve the checkpoint's literal guidance text: it still names
+        // `medium`, although the validator above no longer accepts it.
+        write_internal_system(
+            &mut out,
+            "thinking-effort",
+            &format!(
+                "`thinking_effort` guides on how much to think in your \
+                 thinking channel (not including the response channel), \
+                 supported values include `low`, `medium`, `high`, and `max`.\n\
+                 Now the system is invoked with `thinking_effort={effort}`."
+            ),
+        );
+    }
+
+    // Track prior assistant tool-call ids for tool-result reordering / naming.
+    let mut tool_call_id_index: HashMap<String, (usize, String)> = HashMap::new();
+    let mut pending_tool_run: Vec<(usize, ChatMessage)> = Vec::new();
+
+    let flush_tool_run = |out: &mut String,
+                          run: &mut Vec<(usize, ChatMessage)>,
+                          id_index: &HashMap<String, (usize, String)>|
+     -> Result<()> {
+        if run.is_empty() {
+            return Ok(());
+        }
+
+        let mut resolved = Vec::with_capacity(run.len());
+        let mut unresolved = false;
+        for (offset, message) in run.drain(..) {
+            let ChatMessage::ToolResponse {
+                content,
+                tool_call_id,
+            } = message
+            else {
+                unreachable!("pending tool run only holds tool responses");
+            };
+            match id_index.get(&tool_call_id) {
+                Some(&(position, ref name)) => {
+                    resolved.push((position, offset, content, Some(name.clone())));
+                }
+                None => {
+                    unresolved = true;
+                    resolved.push((usize::MAX, offset, content, None));
+                }
+            }
+        }
+
+        if unresolved {
+            // Preserve original order when the run cannot be fully matched.
+            resolved.sort_by_key(|item| item.1);
+        } else {
+            resolved.sort_by_key(|item| (item.0, item.1));
+        }
+
+        for (xtml_index, (_, _, content, name)) in resolved.into_iter().enumerate() {
+            let tool_name = name.as_deref().ok_or_else(|| {
+                Error::ChatTemplate(
+                    "Kimi K3 tool messages need a resolvable tool name: \
+                         carry a matching tool_call_id against a preceding \
+                         assistant tool_call"
+                        .to_string(),
+                )
+            })?;
+            write_tool_message(out, tool_name, xtml_index + 1, &content)?;
+        }
+        Ok(())
+    };
+
+    for (message_index, message) in request.messages.iter().enumerate() {
+        match message {
+            ChatMessage::ToolResponse { .. } => {
+                pending_tool_run.push((message_index, message.clone()));
+                continue;
+            }
+            _ => {
+                flush_tool_run(&mut out, &mut pending_tool_run, &tool_call_id_index)?;
+            }
+        }
+
+        match message {
+            // Python: role=system with a `tools` field renders a dynamic
+            // tool-declare (`## New Tools Available`). Map that to Developer
+            // messages that carry tools (OpenAI "developer" / system-tools).
+            ChatMessage::Developer {
+                content,
+                tools: Some(local_tools),
+            } if !local_tools.is_empty() => {
+                write_tool_declare(&mut out, local_tools, true)?;
+                if !content_is_empty(content) {
+                    write_role_message(&mut out, "system", None, content)?;
+                }
+            }
+            ChatMessage::System { content } | ChatMessage::Developer { content, .. } => {
+                write_role_message(&mut out, "system", None, content)?;
+            }
+            ChatMessage::User { content } => {
+                write_role_message(&mut out, "user", None, content)?;
+            }
+            ChatMessage::Assistant { content } => {
+                tool_call_id_index.clear();
+                let mut call_position = 0usize;
+                for block in content {
+                    if let AssistantContentBlock::ToolCall(call) = block {
+                        call_position += 1;
+                        if !call.id.is_empty() {
+                            tool_call_id_index
+                                .entry(call.id.clone())
+                                .or_insert((call_position, call.name.clone()));
+                        }
+                    }
+                }
+
+                write_assistant_message(&mut out, content, thinking)?;
+            }
+            ChatMessage::ToolResponse { .. } => unreachable!("handled above"),
+        }
+    }
+    flush_tool_run(&mut out, &mut pending_tool_run, &tool_call_id_index)?;
+
+    match &request.tool_choice {
+        ChatToolChoice::Required => {
+            write_internal_system(
+                &mut out,
+                "tool-choice",
+                "The system is invoked with `tool_choice=required`.\n\
+                 You MUST call tools in the next message.",
+            );
+        }
+        // Emit only when tools are present: Rust defaults tool_choice to None
+        // for tool-free requests, which must not inject a tool-choice message.
+        ChatToolChoice::None if !request.tools.is_empty() => {
+            write_internal_system(
+                &mut out,
+                "tool-choice",
+                "The system is invoked with `tool_choice=none`.\n\
+                 You MUST NOT call any tools in the next message.",
+            );
+        }
+        ChatToolChoice::None | ChatToolChoice::Auto | ChatToolChoice::Function { .. } => {}
+    }
+
+    write_response_format(&mut out, request)?;
+
+    if request.chat_options.add_generation_prompt() {
+        write_open_tag(&mut out, "message", &[("role", "assistant")]);
+        write_open_tag(&mut out, if thinking { "think" } else { "response" }, &[]);
+    }
+
+    Ok(out)
+}
+
+fn request_tools(request: &ChatRequest) -> &[ChatTool] {
+    // Declare tools whenever the request carries them. K3 tool-declare is
+    // independent of tool_choice; tool_choice only injects control messages.
+    request.tools.as_slice()
+}
+
+fn thinking_enabled(request: &ChatRequest) -> Result<bool> {
+    if let Some(thinking) = request.parse_template_bool("thinking")? {
+        return Ok(thinking);
+    }
+    if let Some(enable_thinking) = request.parse_template_bool("enable_thinking")? {
+        return Ok(enable_thinking);
+    }
+    Ok(request
+        .chat_options
+        .reasoning_effort
+        .map(|effort| effort != crate::request::ReasoningEffort::None)
+        .unwrap_or(true))
+}
+
+fn thinking_effort(request: &ChatRequest) -> Result<String> {
+    let effort = if let Some(value) = request.chat_options.template_kwargs.get("thinking_effort") {
+        value.as_str().ok_or_else(|| {
+            Error::ChatTemplate(format!(
+                "template kwarg `thinking_effort` must be a string, got {value}"
+            ))
+        })?
+    } else if let Some(effort) = request.chat_options.reasoning_effort {
+        effort.as_str()
+    } else if let Some(value) = request.chat_options.template_kwargs.get("reasoning_effort") {
+        value.as_str().ok_or_else(|| {
+            Error::ChatTemplate(format!(
+                "template kwarg `reasoning_effort` must be a string, got {value}"
+            ))
+        })?
+    } else {
+        DEFAULT_THINKING_EFFORT
+    };
+
+    if !VALID_THINKING_EFFORTS.contains(&effort) {
+        return Err(Error::ChatTemplate(format!(
+            "unsupported thinking_effort={effort:?}; supported values are `low`, `high`, and `max`"
+        )));
+    }
+    Ok(effort.to_string())
+}
+
+fn content_is_empty(content: &ChatContent) -> bool {
+    match content {
+        ChatContent::Text(text) => text.is_empty(),
+        ChatContent::Parts(parts) => parts.iter().all(|part| match part {
+            ChatContentPart::Text { text } => text.is_empty(),
+            ChatContentPart::ImageUrl { .. }
+            | ChatContentPart::VideoUrl { .. }
+            | ChatContentPart::InputAudio { .. }
+            | ChatContentPart::AudioUrl { .. } => false,
+        }),
+    }
+}
+
+fn write_tool_declare(out: &mut String, tools: &[ChatTool], dynamic: bool) -> Result<()> {
+    let mut specs = Vec::with_capacity(tools.len());
+    for tool in tools {
+        let mut function = Map::new();
+        function.insert(
+            "description".to_string(),
+            Value::String(tool.description.clone().unwrap_or_default()),
+        );
+        function.insert("name".to_string(), Value::String(tool.name.clone()));
+        function.insert("parameters".to_string(), sort_json(&tool.parameters));
+        specs.push(json!({
+            "function": Value::Object(function),
+            "type": "function",
+        }));
+    }
+    let payload = compact_json(&sort_json(&Value::Array(specs)))?;
+
+    let body = if dynamic {
+        format!(
+            "## New Tools Available\n\
+             The system dynamically extends the toolset via lazy-loading.\n\
+             You have access to all existing and extended tools.\n\
+             Here are the specs for the extended tools.\n\n\
+             ```json\n\
+             {payload}\n\
+             ```"
+        )
+    } else {
+        format!(
+            "# Tools\n\
+             Here are the available tools, described in JSONSchema.\n\n\
+             ```json\n\
+             {payload}\n\
+             ```"
+        )
+    };
+
+    write_internal_system(out, "tool-declare", &body);
+    Ok(())
+}
+
+fn write_internal_system(out: &mut String, message_type: &str, body: &str) {
+    write_open_tag(
+        out,
+        "message",
+        &[("role", "system"), ("type", message_type)],
+    );
+    out.push_str(body.trim());
+    write_close_tag(out, "message");
+    out.push_str(END_OF_MSG);
+}
+
+fn write_role_message(
+    out: &mut String,
+    role: &str,
+    name: Option<&str>,
+    content: &ChatContent,
+) -> Result<()> {
+    let mut attrs = vec![("role", role.to_string())];
+    if let Some(name) = name.filter(|name| !name.is_empty()) {
+        attrs.push(("name", name.to_string()));
+    }
+    let attr_refs: Vec<(&str, &str)> = attrs.iter().map(|(k, v)| (*k, v.as_str())).collect();
+    write_open_tag(out, "message", &attr_refs);
+    write_content(out, content)?;
+    write_close_tag(out, "message");
+    out.push_str(END_OF_MSG);
+    Ok(())
+}
+
+fn write_tool_message(
+    out: &mut String,
+    tool_name: &str,
+    index: usize,
+    content: &ChatContent,
+) -> Result<()> {
+    let index_str = index.to_string();
+    write_open_tag(
+        out,
+        "message",
+        &[("role", "tool"), ("tool", tool_name), ("index", &index_str)],
+    );
+    write_content(out, content)?;
+    write_close_tag(out, "message");
+    out.push_str(END_OF_MSG);
+    Ok(())
+}
+
+fn write_assistant_message(
+    out: &mut String,
+    content: &[AssistantContentBlock],
+    thinking: bool,
+) -> Result<()> {
+    write_open_tag(out, "message", &[("role", "assistant")]);
+
+    let mut reasoning = String::new();
+    let mut response = String::new();
+    let mut tool_calls = Vec::new();
+    for block in content {
+        match block {
+            AssistantContentBlock::Reasoning { text } => reasoning.push_str(text),
+            AssistantContentBlock::Text { text } => response.push_str(text),
+            AssistantContentBlock::ToolCall(call) => tool_calls.push(call),
+        }
+    }
+
+    // The think channel is structural: in thinking mode every assistant
+    // message carries open/close tags even when there is no reasoning content.
+    // In non-thinking mode the channel is dropped entirely.
+    if thinking {
+        write_open_tag(out, "think", &[]);
+        if !reasoning.trim().is_empty() {
+            out.push_str(&reasoning);
+        }
+        write_close_tag(out, "think");
+    }
+
+    write_open_tag(out, "response", &[]);
+    out.push_str(&response);
+    write_close_tag(out, "response");
+
+    if !tool_calls.is_empty() {
+        write_open_tag(out, "tools", &[]);
+        for (index, tool_call) in tool_calls.into_iter().enumerate() {
+            write_assistant_tool_call(out, tool_call, index + 1)?;
+        }
+        write_close_tag(out, "tools");
+    }
+
+    write_close_tag(out, "message");
+    out.push_str(END_OF_MSG);
+    Ok(())
+}
+
+fn write_assistant_tool_call(
+    out: &mut String,
+    tool_call: &AssistantToolCall,
+    index: usize,
+) -> Result<()> {
+    let index_str = index.to_string();
+    write_open_tag(
+        out,
+        "call",
+        &[
+            ("tool", tool_call.name.as_str()),
+            ("index", index_str.as_str()),
+        ],
+    );
+
+    let (args, json_block) = normalize_tool_arguments(&tool_call.arguments)?;
+    if let Some(raw) = json_block {
+        write_open_tag(out, "json", &[("type", "object")]);
+        out.push_str(&raw);
+        write_close_tag(out, "json");
+    } else {
+        for (key, value) in args {
+            let typ = xtml_type(&value);
+            write_open_tag(out, "argument", &[("key", key.as_str()), ("type", typ)]);
+            out.push_str(&xtml_value(&value));
+            write_close_tag(out, "argument");
+        }
+    }
+
+    write_close_tag(out, "call");
+    Ok(())
+}
+
+fn write_content(out: &mut String, content: &ChatContent) -> Result<()> {
+    match content {
+        ChatContent::Text(text) => write_text_with_images(out, text),
+        ChatContent::Parts(parts) => {
+            for part in parts {
+                match part {
+                    ChatContentPart::Text { text } => write_text_with_images(out, text)?,
+                    ChatContentPart::ImageUrl { .. } => out.push_str(IMAGE_PLACEHOLDER),
+                    ChatContentPart::VideoUrl { .. } => {
+                        return Err(Error::UnsupportedMultimodalContent("video_url"));
+                    }
+                    ChatContentPart::InputAudio { .. } => {
+                        return Err(Error::UnsupportedMultimodalContent("input_audio"));
+                    }
+                    ChatContentPart::AudioUrl { .. } => {
+                        return Err(Error::UnsupportedMultimodalContent("audio_url"));
+                    }
+                }
+            }
+            Ok(())
+        }
+    }
+}
+
+fn write_text_with_images(out: &mut String, text: &str) -> Result<()> {
+    // Placeholder expansion is left as the literal K3 image token; multimodal
+    // preprocessing can replace it once image prompts are known.
+    out.push_str(text);
+    Ok(())
+}
+
+fn write_response_format(out: &mut String, request: &ChatRequest) -> Result<()> {
+    let Some(rf) = request.chat_options.template_kwargs.get("response_format") else {
+        return Ok(());
+    };
+
+    let rf_type = rf.get("type").and_then(Value::as_str).or_else(|| rf.as_str()).unwrap_or("");
+
+    match rf_type {
+        "json_object" => {
+            write_internal_system(
+                out,
+                "response-format",
+                "The system is invoked with `response_format=json_object`.\n\
+                 Your response must be raw JSON data without markdown code \
+                 blocks (```json) or any additional formatting.",
+            );
+        }
+        "json_schema" => {
+            let schema = extract_response_schema(rf);
+            let schema_json = compact_json(&sort_json(&schema.unwrap_or(Value::Null)))?;
+            write_internal_system(
+                out,
+                "response-format",
+                &format!(
+                    "The system is invoked with `response_format=json_schema`.\n\
+                     Your response must be raw JSON data without markdown code \
+                     blocks (```json) or any additional formatting.\n\
+                     The JSON data must match the following schema:\n\
+                     ```json\n\
+                     {schema_json}\n\
+                     ```"
+                ),
+            );
+        }
+        _ => {}
+    }
+    Ok(())
+}
+
+fn extract_response_schema(response_format: &Value) -> Option<Value> {
+    let json_schema = response_format.get("json_schema")?;
+    if let Some(schema) = json_schema.get("schema") {
+        return Some(schema.clone());
+    }
+    if let Some(schema) = json_schema.get("json_schema") {
+        return Some(schema.clone());
+    }
+    Some(json_schema.clone())
+}
+
+fn normalize_tool_arguments(arguments: &str) -> Result<(Map<String, Value>, Option<String>)> {
+    let trimmed = arguments.trim();
+    if trimmed.is_empty() {
+        return Ok((Map::new(), None));
+    }
+    match serde_json::from_str::<Value>(trimmed) {
+        Ok(Value::Object(map)) => Ok((map, None)),
+        Ok(_) => Err(Error::ChatTemplate(
+            "Kimi K3 tool call arguments must be a JSON object".to_string(),
+        )),
+        Err(_) => Ok((Map::new(), Some(arguments.to_string()))),
+    }
+}
+
+fn xtml_type(value: &Value) -> &'static str {
+    match value {
+        Value::Bool(_) => "boolean",
+        Value::Null => "null",
+        Value::Number(_) => "number",
+        Value::String(_) => "string",
+        Value::Object(_) => "object",
+        Value::Array(_) => "array",
+    }
+}
+
+fn xtml_value(value: &Value) -> String {
+    match value {
+        Value::String(text) => text.clone(),
+        other => compact_json(other).unwrap_or_else(|_| other.to_string()),
+    }
+}
+
+fn write_open_tag(out: &mut String, tag: &str, attrs: &[(&str, &str)]) {
+    out.push_str(OPEN);
+    out.push_str(tag);
+    for (key, value) in attrs {
+        let _ = write!(out, " {key}=\"{}\"", escape_attr_value(value));
+    }
+    out.push_str(SEP);
+}
+
+fn write_close_tag(out: &mut String, tag: &str) {
+    out.push_str(CLOSE);
+    out.push_str(tag);
+    out.push_str(SEP);
+}
+
+fn escape_attr_value(value: &str) -> String {
+    value.replace('&', "&amp;").replace('"', "&quot;")
+}
+
+fn compact_json(value: &Value) -> Result<String> {
+    serde_json::to_string(value).map_err(|error| Error::ChatTemplate(error.to_string()))
+}
+
+fn sort_json(value: &Value) -> Value {
+    match value {
+        Value::Array(items) => Value::Array(items.iter().map(sort_json).collect()),
+        Value::Object(map) => {
+            let mut sorted = Map::new();
+            let mut keys = map.keys().collect::<Vec<_>>();
+            keys.sort();
+            for key in keys {
+                sorted.insert(key.clone(), sort_json(&map[key]));
+            }
+            Value::Object(sorted)
+        }
+        _ => value.clone(),
+    }
+}

rust/src/chat/src/renderer/kimi_k3/fixtures/controls_thinking_off_input.json

@@ -0,0 +1,15 @@
+{
+  "messages": [
+    {
+      "role": "user",
+      "content": "json pls"
+    }
+  ],
+  "add_generation_prompt": true,
+  "template_kwargs": {
+    "thinking": false,
+    "response_format": {
+      "type": "json_object"
+    }
+  }
+}

rust/src/chat/src/renderer/kimi_k3/fixtures/controls_thinking_off_output.txt

@@ -0,0 +1,2 @@
+<|open|>message role="user"<|sep|>json pls<|close|>message<|sep|><|end_of_msg|><|open|>message role="system" type="response-format"<|sep|>The system is invoked with `response_format=json_object`.
+Your response must be raw JSON data without markdown code blocks (```json) or any additional formatting.<|close|>message<|sep|><|end_of_msg|><|open|>message role="assistant"<|sep|><|open|>response<|sep|>

rust/src/chat/src/renderer/kimi_k3/fixtures/dynamic_system_tool_declare_input.json

@@ -0,0 +1,62 @@
+{
+  "tools": [
+    {
+      "type": "function",
+      "function": {
+        "name": "get_weather",
+        "description": "weather",
+        "parameters": {
+          "type": "object",
+          "properties": {
+            "city": {
+              "type": "string"
+            },
+            "days": {
+              "type": "number"
+            }
+          },
+          "required": [
+            "city"
+          ]
+        }
+      }
+    }
+  ],
+  "messages": [
+    {
+      "role": "user",
+      "content": "hi"
+    },
+    {
+      "role": "developer",
+      "tools": [
+        {
+          "type": "function",
+          "function": {
+            "name": "calc",
+            "description": "calculator",
+            "parameters": {
+              "type": "object",
+              "properties": {
+                "x": {
+                  "type": "number"
+                }
+              },
+              "required": [
+                "x"
+              ]
+            }
+          }
+        }
+      ]
+    },
+    {
+      "role": "user",
+      "content": "use calc"
+    }
+  ],
+  "add_generation_prompt": true,
+  "template_kwargs": {
+    "thinking": true
+  }
+}

rust/src/chat/src/renderer/kimi_k3/fixtures/dynamic_system_tool_declare_output.txt

@@ -0,0 +1,14 @@
+<|open|>message role="system" type="tool-declare"<|sep|># Tools
+Here are the available tools, described in JSONSchema.
+
+```json
+[{"function":{"description":"weather","name":"get_weather","parameters":{"properties":{"city":{"type":"string"},"days":{"type":"number"}},"required":["city"],"type":"object"}},"type":"function"}]
+```<|close|>message<|sep|><|end_of_msg|><|open|>message role="system" type="thinking-effort"<|sep|>`thinking_effort` guides on how much to think in your thinking channel (not including the response channel), supported values include `low`, `medium`, `high`, and `max`.
+Now the system is invoked with `thinking_effort=max`.<|close|>message<|sep|><|end_of_msg|><|open|>message role="user"<|sep|>hi<|close|>message<|sep|><|end_of_msg|><|open|>message role="system" type="tool-declare"<|sep|>## New Tools Available
+The system dynamically extends the toolset via lazy-loading.
+You have access to all existing and extended tools.
+Here are the specs for the extended tools.
+
+```json
+[{"function":{"description":"calculator","name":"calc","parameters":{"properties":{"x":{"type":"number"}},"required":["x"],"type":"object"}},"type":"function"}]
+```<|close|>message<|sep|><|end_of_msg|><|open|>message role="user"<|sep|>use calc<|close|>message<|sep|><|end_of_msg|><|open|>message role="assistant"<|sep|><|open|>think<|sep|>

rust/src/chat/src/renderer/kimi_k3/fixtures/history_preserve_and_image_input.json

@@ -0,0 +1,31 @@
+{
+  "messages": [
+    {
+      "role": "user",
+      "content": [
+        {
+          "type": "text",
+          "text": "see this"
+        },
+        {
+          "type": "image_url",
+          "image_url": "https://example.com/a.png"
+        }
+      ]
+    },
+    {
+      "role": "assistant",
+      "content": "old answer",
+      "reasoning_content": "old reasoning"
+    },
+    {
+      "role": "user",
+      "content": "continue"
+    }
+  ],
+  "add_generation_prompt": true,
+  "template_kwargs": {
+    "thinking": true,
+    "thinking_effort": "high"
+  }
+}

rust/src/chat/src/renderer/kimi_k3/fixtures/history_preserve_and_image_output.txt

@@ -0,0 +1,2 @@
+<|open|>message role="system" type="thinking-effort"<|sep|>`thinking_effort` guides on how much to think in your thinking channel (not including the response channel), supported values include `low`, `medium`, `high`, and `max`.
+Now the system is invoked with `thinking_effort=high`.<|close|>message<|sep|><|end_of_msg|><|open|>message role="user"<|sep|>see this<|media_pad|><|close|>message<|sep|><|end_of_msg|><|open|>message role="assistant"<|sep|><|open|>think<|sep|>old reasoning<|close|>think<|sep|><|open|>response<|sep|>old answer<|close|>response<|sep|><|close|>message<|sep|><|end_of_msg|><|open|>message role="user"<|sep|>continue<|close|>message<|sep|><|end_of_msg|><|open|>message role="assistant"<|sep|><|open|>think<|sep|>

rust/src/chat/src/renderer/kimi_k3/fixtures/tools_history_and_required_input.json

@@ -0,0 +1,87 @@
+{
+  "tools": [
+    {
+      "type": "function",
+      "function": {
+        "name": "get_weather",
+        "description": "weather",
+        "parameters": {
+          "type": "object",
+          "properties": {
+            "city": {
+              "type": "string"
+            },
+            "days": {
+              "type": "number"
+            }
+          },
+          "required": [
+            "city"
+          ]
+        }
+      }
+    },
+    {
+      "type": "function",
+      "function": {
+        "name": "calc",
+        "description": "calculator",
+        "parameters": {
+          "type": "object",
+          "properties": {
+            "x": {
+              "type": "number"
+            }
+          },
+          "required": [
+            "x"
+          ]
+        }
+      }
+    }
+  ],
+  "messages": [
+    {
+      "role": "user",
+      "content": "Hangzhou weather and calc"
+    },
+    {
+      "role": "assistant",
+      "content": "I'll check.",
+      "reasoning_content": "Need tools.",
+      "tool_calls": [
+        {
+          "id": "get_weather:0",
+          "type": "function",
+          "function": {
+            "name": "get_weather",
+            "arguments": "{\"city\":\"Hangzhou\",\"days\":1}"
+          }
+        },
+        {
+          "id": "calc:1",
+          "type": "function",
+          "function": {
+            "name": "calc",
+            "arguments": "{\"x\":2}"
+          }
+        }
+      ]
+    },
+    {
+      "role": "tool",
+      "tool_call_id": "calc:1",
+      "content": "2"
+    },
+    {
+      "role": "tool",
+      "tool_call_id": "get_weather:0",
+      "content": "{\"city\":\"Hangzhou\",\"condition\":\"rain\"}"
+    }
+  ],
+  "add_generation_prompt": true,
+  "tool_choice": "required",
+  "template_kwargs": {
+    "thinking": true
+  }
+}

rust/src/chat/src/renderer/kimi_k3/fixtures/tools_history_and_required_output.txt

@@ -0,0 +1,8 @@
+<|open|>message role="system" type="tool-declare"<|sep|># Tools
+Here are the available tools, described in JSONSchema.
+
+```json
+[{"function":{"description":"weather","name":"get_weather","parameters":{"properties":{"city":{"type":"string"},"days":{"type":"number"}},"required":["city"],"type":"object"}},"type":"function"},{"function":{"description":"calculator","name":"calc","parameters":{"properties":{"x":{"type":"number"}},"required":["x"],"type":"object"}},"type":"function"}]
+```<|close|>message<|sep|><|end_of_msg|><|open|>message role="system" type="thinking-effort"<|sep|>`thinking_effort` guides on how much to think in your thinking channel (not including the response channel), supported values include `low`, `medium`, `high`, and `max`.
+Now the system is invoked with `thinking_effort=max`.<|close|>message<|sep|><|end_of_msg|><|open|>message role="user"<|sep|>Hangzhou weather and calc<|close|>message<|sep|><|end_of_msg|><|open|>message role="assistant"<|sep|><|open|>think<|sep|>Need tools.<|close|>think<|sep|><|open|>response<|sep|>I'll check.<|close|>response<|sep|><|open|>tools<|sep|><|open|>call tool="get_weather" index="1"<|sep|><|open|>argument key="city" type="string"<|sep|>Hangzhou<|close|>argument<|sep|><|open|>argument key="days" type="number"<|sep|>1<|close|>argument<|sep|><|close|>call<|sep|><|open|>call tool="calc" index="2"<|sep|><|open|>argument key="x" type="number"<|sep|>2<|close|>argument<|sep|><|close|>call<|sep|><|close|>tools<|sep|><|close|>message<|sep|><|end_of_msg|><|open|>message role="tool" tool="get_weather" index="1"<|sep|>{"city":"Hangzhou","condition":"rain"}<|close|>message<|sep|><|end_of_msg|><|open|>message role="tool" tool="calc" index="2"<|sep|>2<|close|>message<|sep|><|end_of_msg|><|open|>message role="system" type="tool-choice"<|sep|>The system is invoked with `tool_choice=required`.
+You MUST call tools in the next message.<|close|>message<|sep|><|end_of_msg|><|open|>message role="assistant"<|sep|><|open|>think<|sep|>

rust/src/chat/src/renderer/kimi_k3/mod.rs

@@ -0,0 +1,36 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+//! Native Kimi K3 XTML chat renderer.
+
+mod encoding;
+#[cfg(test)]
+mod tests;
+
+use vllm_text::Prompt;
+
+use super::{ChatRenderer, RenderedPrompt, request_template_kwargs};
+use crate::Result;
+use crate::request::ChatRequest;
+
+/// Dedicated Kimi K3 XTML renderer.
+#[derive(Debug, Clone, Copy, Default)]
+pub struct KimiK3ChatRenderer;
+
+impl KimiK3ChatRenderer {
+    /// Create a Kimi K3 renderer.
+    pub fn new() -> Self {
+        Self
+    }
+}
+
+impl ChatRenderer for KimiK3ChatRenderer {
+    fn render(&self, request: &ChatRequest) -> Result<RenderedPrompt> {
+        request.validate()?;
+
+        Ok(RenderedPrompt {
+            prompt: Prompt::Text(encoding::render_request(request)?),
+            effective_template_kwargs: request_template_kwargs(request),
+        })
+    }
+}

rust/src/chat/src/renderer/kimi_k3/tests.rs

@@ -0,0 +1,201 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+//! Golden fixtures generated from HF remote-code `encoding_k3.py`.
+
+use std::path::PathBuf;
+
+use expect_test::{expect, expect_file};
+use serde_json::json;
+
+use super::KimiK3ChatRenderer;
+use crate::AssistantContentBlock;
+use crate::ChatRenderer;
+use crate::renderer::test_utils::{FixtureRequestOptions, fixture_chat_request};
+use crate::request::{ChatMessage, GenerationPromptMode, ReasoningEffort};
+
+fn render_request(request: &crate::request::ChatRequest) -> String {
+    KimiK3ChatRenderer::new()
+        .render(request)
+        .unwrap()
+        .prompt
+        .into_text()
+        .expect("kimi k3 renderer should return text prompt")
+}
+
+fn fixture_path(name: &str) -> PathBuf {
+    PathBuf::from(env!("CARGO_MANIFEST_DIR"))
+        .join("src/renderer/kimi_k3/fixtures")
+        .join(name)
+}
+
+fn kimi_k3_fixture_options() -> FixtureRequestOptions {
+    FixtureRequestOptions {
+        // Fixture JSON owns thinking via `template_kwargs`.
+        enable_thinking: None,
+        no_generation_prompt_when_last_assistant: false,
+    }
+}
+
+fn assert_golden(name: &str) {
+    let input_name = format!("{name}_input.json");
+    let request = fixture_chat_request(&fixture_path(&input_name), kimi_k3_fixture_options());
+    let rendered = render_request(&request);
+    expect_file![format!("fixtures/{name}_output.txt")].assert_eq(&rendered);
+}
+
+#[test]
+fn golden_history_preserve_and_image() {
+    assert_golden("history_preserve_and_image");
+}
+
+#[test]
+fn golden_tools_history_and_required() {
+    assert_golden("tools_history_and_required");
+}
+
+#[test]
+fn golden_controls_thinking_off() {
+    assert_golden("controls_thinking_off");
+}
+
+#[test]
+fn golden_dynamic_system_tool_declare() {
+    assert_golden("dynamic_system_tool_declare");
+}
+
+#[test]
+fn thinking_history_renders_empty_think_channel() {
+    let mut request = crate::request::ChatRequest::for_test();
+    request.messages = vec![
+        ChatMessage::user("question"),
+        ChatMessage::assistant_text("answer"),
+        ChatMessage::user("follow-up"),
+    ];
+    request.chat_options.generation_prompt_mode = GenerationPromptMode::NoGenerationPrompt;
+
+    let rendered = render_request(&request);
+
+    assert!(rendered.contains(
+        "<|open|>message role=\"assistant\"<|sep|>\
+         <|open|>think<|sep|><|close|>think<|sep|>\
+         <|open|>response<|sep|>answer<|close|>response<|sep|>"
+    ));
+}
+
+#[test]
+fn non_thinking_history_omits_reasoning_channel() {
+    let mut request = crate::request::ChatRequest::for_test();
+    request.messages = vec![
+        ChatMessage::user("question"),
+        ChatMessage::assistant_blocks(vec![
+            AssistantContentBlock::Reasoning {
+                text: "hidden reasoning".to_string(),
+            },
+            AssistantContentBlock::Text {
+                text: "answer".to_string(),
+            },
+        ]),
+        ChatMessage::user("follow-up"),
+    ];
+    request
+        .chat_options
+        .template_kwargs
+        .insert("thinking".to_string(), json!(false));
+    request.chat_options.generation_prompt_mode = GenerationPromptMode::NoGenerationPrompt;
+
+    let rendered = render_request(&request);
+
+    assert!(!rendered.contains("hidden reasoning"));
+    assert!(!rendered.contains("<|open|>think<|sep|>"));
+    assert!(rendered.contains(
+        "<|open|>message role=\"assistant\"<|sep|>\
+         <|open|>response<|sep|>answer<|close|>response<|sep|>"
+    ));
+}
+
+#[test]
+fn defaults_thinking_effort_to_max() {
+    let rendered = render_request(&crate::request::ChatRequest::for_test());
+
+    expect![[r#"<|open|>message role="system" type="thinking-effort"<|sep|>`thinking_effort` guides on how much to think in your thinking channel (not including the response channel), supported values include `low`, `medium`, `high`, and `max`.
+Now the system is invoked with `thinking_effort=max`.<|close|>message<|sep|><|end_of_msg|><|open|>message role="user"<|sep|>test<|close|>message<|sep|><|end_of_msg|><|open|>message role="assistant"<|sep|><|open|>think<|sep|>"#]]
+        .assert_eq(&rendered);
+}
+
+#[test]
+fn translates_standard_thinking_kwargs() {
+    let mut request = crate::request::ChatRequest::for_test();
+    request
+        .chat_options
+        .template_kwargs
+        .insert("enable_thinking".to_string(), json!(true));
+    request
+        .chat_options
+        .template_kwargs
+        .insert("reasoning_effort".to_string(), json!("high"));
+
+    let rendered = render_request(&request);
+
+    assert!(rendered.contains("thinking_effort=high"));
+    assert!(rendered.ends_with("<|open|>think<|sep|>"));
+}
+
+#[test]
+fn native_k3_kwargs_take_precedence() {
+    let mut request = crate::request::ChatRequest::for_test();
+    request.chat_options.template_kwargs.extend([
+        ("thinking".to_string(), json!(true)),
+        ("enable_thinking".to_string(), json!(false)),
+        ("thinking_effort".to_string(), json!("low")),
+        ("reasoning_effort".to_string(), json!("high")),
+    ]);
+
+    let rendered = render_request(&request);
+
+    assert!(rendered.contains("thinking_effort=low"));
+    assert!(!rendered.contains("thinking_effort=high"));
+}
+
+#[test]
+fn standard_none_disables_thinking() {
+    let mut request = crate::request::ChatRequest::for_test();
+    request.chat_options.template_kwargs.extend([
+        ("enable_thinking".to_string(), json!(false)),
+        ("reasoning_effort".to_string(), json!("none")),
+    ]);
+
+    let rendered = render_request(&request);
+
+    assert!(!rendered.contains("type=\"thinking-effort\""));
+    assert!(rendered.ends_with("<|open|>response<|sep|>"));
+}
+
+#[test]
+fn typed_none_disables_thinking() {
+    let mut request = crate::request::ChatRequest::for_test();
+    request.chat_options.reasoning_effort = Some(ReasoningEffort::None);
+
+    let rendered = render_request(&request);
+
+    assert!(!rendered.contains("type=\"thinking-effort\""));
+    assert!(rendered.ends_with("<|open|>response<|sep|>"));
+}
+
+#[test]
+fn rejects_removed_medium_thinking_effort() {
+    let mut request = crate::request::ChatRequest::for_test();
+    request
+        .chat_options
+        .template_kwargs
+        .insert("thinking_effort".to_string(), json!("medium"));
+
+    let error = KimiK3ChatRenderer::new().render(&request).unwrap_err();
+
+    expect![[r#"
+        ChatTemplate(
+            "unsupported thinking_effort=\"medium\"; supported values are `low`, `high`, and `max`",
+        )
+    "#]]
+    .assert_debug_eq(&error);
+}

rust/src/chat/src/renderer/mod.rs

@@ -15,6 +15,7 @@ pub mod deepseek_v4;
 pub mod harmony;
 pub mod hf;
 mod inkling;
+mod kimi_k3;
 mod selection;
 #[cfg(test)]
 mod test_utils;
@@ -23,6 +24,7 @@ pub use deepseek_v4::DeepSeekV4ChatRenderer;
 pub use deepseek_v32::DeepSeekV32ChatRenderer;
 pub use harmony::HarmonyChatRenderer;
 pub use inkling::InklingChatRenderer;
+pub use kimi_k3::KimiK3ChatRenderer;
 pub use selection::RendererSelection;
 
 /// Rendered chat prompt submitted to the text backend.

rust/src/chat/src/renderer/selection.rs

@@ -26,6 +26,8 @@ pub enum RendererSelection {
     Harmony,
     /// Force the Inkling native token renderer.
     Inkling,
+    /// Force the Kimi K3 XTML renderer.
+    KimiK3,
 }
 
 impl RendererSelection {
@@ -37,6 +39,7 @@ impl RendererSelection {
     pub const HF_LITERAL: &str = "hf";
     pub const INKLING_LITERAL: &str = "inkling";
     pub const INKLING_MODEL_TYPE: &str = "inkling_mm_model";
+    pub const KIMI_K3_LITERAL: &str = "kimi_k3";
 
     /// Resolve the renderer selection using the given model type string, if
     /// it's `Auto`.
@@ -47,6 +50,7 @@ impl RendererSelection {
                 Self::DEEPSEEK_V4_LITERAL => Self::DeepSeekV4,
                 Self::GPT_OSS_MODEL_TYPE => Self::Harmony,
                 Self::INKLING_MODEL_TYPE => Self::Inkling,
+                Self::KIMI_K3_LITERAL => Self::KimiK3,
                 _ => Self::Hf,
             },
             selection => selection,
@@ -70,6 +74,8 @@ impl FromStr for RendererSelection {
             Ok(Self::Harmony)
         } else if value.eq_ignore_ascii_case(Self::INKLING_LITERAL) {
             Ok(Self::Inkling)
+        } else if value.eq_ignore_ascii_case(Self::KIMI_K3_LITERAL) {
+            Ok(Self::KimiK3)
         } else {
             Err(format!(
                 "unknown renderer `{value}` (expected one of: {})",
@@ -88,6 +94,7 @@ impl fmt::Display for RendererSelection {
             Self::DeepSeekV4 => f.write_str(Self::DEEPSEEK_V4_LITERAL),
             Self::Harmony => f.write_str(Self::HARMONY_LITERAL),
             Self::Inkling => f.write_str(Self::INKLING_LITERAL),
+            Self::KimiK3 => f.write_str(Self::KIMI_K3_LITERAL),
         }
     }
 }
@@ -114,7 +121,7 @@ mod tests {
     fn renderer_selection_expected_error_message() {
         let err = RendererSelection::from_str("unknown").unwrap_err();
         expect_test::expect![
-            "unknown renderer `unknown` (expected one of: auto, hf, deepseek_v32, deepseek_v4, harmony, inkling)"
+            "unknown renderer `unknown` (expected one of: auto, hf, deepseek_v32, deepseek_v4, harmony, inkling, kimi_k3)"
         ]
         .assert_eq(&err);
     }

rust/src/chat/src/renderer/test_utils.rs

@@ -1,6 +1,7 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
+use std::collections::HashMap;
 use std::fs;
 use std::path::Path;
 
@@ -16,8 +17,9 @@ use crate::request::{
 /// Options for constructing a [`ChatRequest`] from a fixture file.
 #[derive(Debug, Clone, Copy)]
 pub(crate) struct FixtureRequestOptions {
-    /// Whether to set the template kwarg `[enable_]thinking=true`.
-    pub enable_thinking: bool,
+    /// Optional thinking toggle applied only when the fixture does not already
+    /// set `thinking` / `enable_thinking` in [`FixtureRequest::template_kwargs`].
+    pub enable_thinking: Option<bool>,
     /// Whether fixtures ending in an assistant message should omit the
     /// trailing generation prompt.
     pub no_generation_prompt_when_last_assistant: bool,
@@ -46,6 +48,12 @@ pub(crate) struct FixtureRequest {
     messages: Vec<FixtureMessage>,
     add_generation_prompt: Option<bool>,
     reasoning_effort: Option<ReasoningEffort>,
+    /// Extra chat-template kwargs (thinking, preserve_thinking, response_format, …).
+    #[serde(default)]
+    template_kwargs: HashMap<String, Value>,
+    /// When omitted, defaults to `auto` if tools are present, otherwise `none`.
+    #[serde(default)]
+    tool_choice: Option<ChatToolChoice>,
 }
 
 impl FixtureFile {
@@ -57,6 +65,8 @@ impl FixtureFile {
                 messages,
                 add_generation_prompt: None,
                 reasoning_effort: None,
+                template_kwargs: HashMap::new(),
+                tool_choice: None,
             },
         }
     }
@@ -69,6 +79,7 @@ pub(crate) enum FixtureMessage {
         content: FixtureContent,
     },
     Developer {
+        #[serde(default)]
         content: FixtureContent,
         #[serde(default)]
         tools: Vec<FixtureTool>,
@@ -98,6 +109,12 @@ pub(crate) enum FixtureContent {
     Parts(Vec<FixtureContentPart>),
 }
 
+impl Default for FixtureContent {
+    fn default() -> Self {
+        Self::Text(String::new())
+    }
+}
+
 #[derive(Debug, Deserialize)]
 #[serde(tag = "type", rename_all = "snake_case")]
 pub(crate) enum FixtureContentPart {
@@ -136,6 +153,13 @@ struct FixtureToolCallFunction {
 
 impl FixtureRequest {
     fn into_chat_request(self, options: FixtureRequestOptions) -> ChatRequest {
+        let tools = to_chat_tools(&self.tools);
+        let tool_choice = self.tool_choice.unwrap_or(if tools.is_empty() {
+            ChatToolChoice::None
+        } else {
+            ChatToolChoice::Auto
+        });
+
         let mut request = ChatRequest {
             request_id: "renderer-fixture".to_string(),
             messages: self
@@ -144,12 +168,8 @@ impl FixtureRequest {
                 .enumerate()
                 .map(|(index, message)| fixture_message_to_chat_message(index, message))
                 .collect(),
-            tools: to_chat_tools(&self.tools),
-            tool_choice: if self.tools.is_empty() {
-                ChatToolChoice::None
-            } else {
-                ChatToolChoice::Auto
-            },
+            tools,
+            tool_choice,
             ..ChatRequest::for_test()
         };
 
@@ -162,9 +182,14 @@ impl FixtureRequest {
             request.chat_options.generation_prompt_mode = GenerationPromptMode::NoGenerationPrompt;
         }
         request.chat_options.reasoning_effort = self.reasoning_effort;
-        if options.enable_thinking {
-            for key in ["thinking", "enable_thinking"] {
-                request.chat_options.template_kwargs.insert(key.to_string(), Value::Bool(true));
+        request.chat_options.template_kwargs.extend(self.template_kwargs);
+
+        // Options supply a default thinking toggle only when the fixture did not.
+        if let Some(thinking) = options.enable_thinking {
+            let kwargs = &mut request.chat_options.template_kwargs;
+            if !kwargs.contains_key("thinking") && !kwargs.contains_key("enable_thinking") {
+                kwargs.insert("thinking".to_string(), Value::Bool(thinking));
+                kwargs.insert("enable_thinking".to_string(), Value::Bool(thinking));
             }
         }
 

rust/src/chat/tests/roundtrip.rs

@@ -212,6 +212,23 @@ impl RoundtripCase {
         }
     }
 
+    /// Kimi K3 XTML tool/reasoning channels (native renderer + unified parser).
+    ///
+    /// Needs HF tokenizer files under `HF_HOME` (`tiktoken.model` +
+    /// `tokenizer_config.json`). Weights are not required for this text-level
+    /// roundtrip.
+    fn kimi_k3() -> Self {
+        Self {
+            model_id: "moonshotai/Kimi-K3",
+            assistant_stop_suffix: "<|end_of_msg|>",
+            tool_call_parser: ParserSelection::Auto,
+            reasoning_parser: ParserSelection::Auto,
+            thinking_behavior: ThinkingBehavior::Toggleable { default: true },
+            json_fmt: compact_json_fmt(),
+            sort_json_keys: false,
+        }
+    }
+
     /// SeedOSS with `<seed:think>` / `</seed:think>` reasoning tags.
     fn seed_oss() -> Self {
         Self {
@@ -312,6 +329,8 @@ roundtrip_tests! {
     nemotron_v3 => [reasoning_and_content],
     gemma4 => [tool_call_mix], // Gemma4 strips reasoning in history if there's no tool call
     kimi_k25 => [tool_call_mix], // Kimi K2.5 strips reasoning in history
+    // K3 drops plain-assistant reasoning in history; tool-call turns keep it.
+    kimi_k3 => [tool_call_mix],
     gpt_oss => [tool_call_mix], // Harmony strips reasoning in history if there's no tool call
     inkling => [reasoning_and_content, tool_call_mix],
 }

rust/src/parser/src/unified/kimi_k3/structural_tag.rs

@@ -0,0 +1,500 @@
+//! Structural-tag grammar for Kimi K3 XTML tool calls.
+
+use serde_json::{Map, Value};
+use xgrammar_structural_tag::Result;
+use xgrammar_structural_tag::builders::{
+    StructuralTagBuilder, StructuralTagContext, StructuralTagOptions,
+};
+use xgrammar_structural_tag::format::{Format, JsonSchemaFormat, StructuralTag, TagFormat};
+use xgrammar_structural_tag::tool::{BuilderToolChoice, FunctionToolParam, function_parameters};
+
+use super::{
+    ARG_CLOSE, CALL_CLOSE, END_OF_MSG, JSON_CLOSE, JSON_OPEN, MESSAGE_CLOSE, OPEN, RESPONSE_CLOSE,
+    RESPONSE_OPEN, SEP, THINK_CLOSE, TOOLS_CLOSE, TOOLS_OPEN,
+};
+
+pub(super) static KIMI_K3_STRUCTURAL_TAG_BUILDER: KimiK3StructuralTagBuilder =
+    KimiK3StructuralTagBuilder;
+
+const XTML_TYPES: &[&str] = &["string", "number", "boolean", "null", "object", "array"];
+
+/// Kimi K3 XTML structural-tag builder.
+#[derive(Debug, Clone, Copy, Default)]
+pub struct KimiK3StructuralTagBuilder;
+
+impl StructuralTagBuilder for KimiK3StructuralTagBuilder {
+    fn build(&self, ctx: StructuralTagContext<'_>) -> Result<StructuralTag> {
+        let mut elements = response_prefix(ctx.options.reasoning);
+
+        let tools = match ctx.tool_choice {
+            // Serving lowering filters empty tools before calling the builder,
+            // while direct `build_structural_tag(..., auto, ...)` calls do not.
+            BuilderToolChoice::Auto if ctx.function_tools.is_empty() => None,
+            BuilderToolChoice::Auto => Some(Format::optional(tools_channel(
+                ctx.function_tools,
+                ctx.options,
+            ))),
+            BuilderToolChoice::Forced | BuilderToolChoice::Required => {
+                Some(tools_channel(ctx.function_tools, ctx.options))
+            }
+        };
+        if let Some(tools) = tools {
+            elements.push(tools);
+        }
+        elements.push(Format::optional(Format::const_string(MESSAGE_CLOSE)));
+
+        Ok(StructuralTag::new(Format::sequence(elements)))
+    }
+}
+
+fn response_prefix(reasoning: bool) -> Vec<Format> {
+    let mut elements = Vec::new();
+    if reasoning {
+        elements.push(Format::tag(
+            "",
+            Format::any_text_excluding(&[THINK_CLOSE, END_OF_MSG]),
+            THINK_CLOSE,
+        ));
+        elements.push(Format::const_string(RESPONSE_OPEN));
+    } else {
+        elements.push(Format::optional(Format::const_string(RESPONSE_OPEN)));
+    }
+    elements.push(Format::tag(
+        "",
+        Format::any_text_excluding(&[RESPONSE_CLOSE, TOOLS_OPEN, MESSAGE_CLOSE, END_OF_MSG]),
+        RESPONSE_CLOSE,
+    ));
+    elements
+}
+
+fn tools_channel(tools: &[FunctionToolParam], options: StructuralTagOptions) -> Format {
+    let calls = tools.iter().map(|tool| call_tag(tool, options)).collect();
+    Format::tag(
+        TOOLS_OPEN,
+        Format::tags_with_separator(calls, "", true, false),
+        TOOLS_CLOSE,
+    )
+}
+
+fn call_tag(tool: &FunctionToolParam, options: StructuralTagOptions) -> TagFormat {
+    let parameters = function_parameters(&tool.function);
+    let call_body = Format::or(vec![
+        typed_arguments(&parameters, options),
+        raw_json_arguments(&parameters, options),
+    ]);
+
+    TagFormat::new(
+        format!(
+            "{OPEN}call tool=\"{}\" index=\"",
+            escape_attr_value(&tool.function.name)
+        ),
+        Format::sequence(vec![
+            Format::regex("[1-9][0-9]*"),
+            Format::const_string(format!("\"{SEP}")),
+            call_body,
+        ]),
+        CALL_CLOSE,
+    )
+}
+
+fn typed_arguments(parameters: &Value, options: StructuralTagOptions) -> Format {
+    let Some(schema) = parameters.as_object() else {
+        return if parameters == &Value::Bool(false) {
+            Format::const_string("")
+        } else {
+            Format::star(permissive_argument())
+        };
+    };
+    let Some(properties) = schema.get("properties").and_then(Value::as_object) else {
+        return Format::star(permissive_argument());
+    };
+    if properties.is_empty() {
+        return Format::star(permissive_argument());
+    }
+
+    let root_defs = root_definitions(schema);
+    let arguments = properties
+        .iter()
+        .flat_map(|(key, schema)| argument_tags(key, schema, &root_defs, options))
+        .map(Format::Tag)
+        .collect::<Vec<_>>();
+    let arguments = match arguments.as_slice() {
+        [argument] => argument.clone(),
+        _ => Format::or(arguments),
+    };
+    // Keep typed arguments order-agnostic and non-unique, but do not allow an
+    // empty call when the root schema declares required properties.
+    if schema
+        .get("required")
+        .and_then(Value::as_array)
+        .is_some_and(|required| !required.is_empty())
+    {
+        Format::plus(arguments)
+    } else {
+        Format::star(arguments)
+    }
+}
+
+fn argument_tags(
+    key: &str,
+    schema: &Value,
+    root_defs: &Map<String, Value>,
+    options: StructuralTagOptions,
+) -> Vec<TagFormat> {
+    let types = schema_types(schema);
+    types
+        .into_iter()
+        .map(|xtml_type| {
+            let content = if xtml_type == "string" {
+                string_argument_content(schema)
+            } else {
+                json_schema(
+                    attach_root_definitions(&narrow_schema_type(schema, xtml_type), root_defs),
+                    options,
+                )
+            };
+            TagFormat::new(
+                format!(
+                    "{OPEN}argument key=\"{}\" type=\"{xtml_type}\"{SEP}",
+                    escape_attr_value(key)
+                ),
+                content,
+                ARG_CLOSE,
+            )
+        })
+        .collect()
+}
+
+fn schema_types(schema: &Value) -> Vec<&'static str> {
+    let Some(schema) = schema.as_object() else {
+        return XTML_TYPES.to_vec();
+    };
+    let mut types = Vec::new();
+    match schema.get("type") {
+        Some(Value::String(value)) => push_schema_type(&mut types, value),
+        Some(Value::Array(values)) => {
+            for value in values.iter().filter_map(Value::as_str) {
+                push_schema_type(&mut types, value);
+            }
+        }
+        _ => {}
+    }
+    if types.is_empty()
+        && let Some(value) = schema.get("const")
+    {
+        push_value_type(&mut types, value);
+    }
+    if types.is_empty()
+        && let Some(values) = schema.get("enum").and_then(Value::as_array)
+    {
+        for value in values {
+            push_value_type(&mut types, value);
+        }
+    }
+    if types.is_empty() {
+        XTML_TYPES.to_vec()
+    } else {
+        types
+    }
+}
+
+fn push_value_type(types: &mut Vec<&'static str>, value: &Value) {
+    let xtml_type = match value {
+        Value::String(_) => "string",
+        Value::Number(_) => "number",
+        Value::Bool(_) => "boolean",
+        Value::Null => "null",
+        Value::Object(_) => "object",
+        Value::Array(_) => "array",
+    };
+    if !types.contains(&xtml_type) {
+        types.push(xtml_type);
+    }
+}
+
+fn narrow_schema_type(schema: &Value, xtml_type: &str) -> Value {
+    let Some(mut schema) = schema.as_object().cloned() else {
+        return schema.clone();
+    };
+    let json_type = if xtml_type == "number" && explicitly_integer_only(&schema) {
+        "integer"
+    } else {
+        xtml_type
+    };
+    schema.insert("type".to_string(), Value::String(json_type.to_string()));
+    Value::Object(schema)
+}
+
+fn explicitly_integer_only(schema: &Map<String, Value>) -> bool {
+    match schema.get("type") {
+        Some(Value::String(value)) => value == "integer",
+        Some(Value::Array(values)) => {
+            let values = values.iter().filter_map(Value::as_str).collect::<Vec<_>>();
+            values.contains(&"integer") && !values.contains(&"number")
+        }
+        _ => false,
+    }
+}
+
+fn push_schema_type(types: &mut Vec<&'static str>, json_type: &str) {
+    let xtml_type = match json_type {
+        "string" => Some("string"),
+        "integer" | "number" => Some("number"),
+        "boolean" => Some("boolean"),
+        "null" => Some("null"),
+        "object" => Some("object"),
+        "array" => Some("array"),
+        _ => None,
+    };
+    if let Some(xtml_type) = xtml_type
+        && !types.contains(&xtml_type)
+    {
+        types.push(xtml_type);
+    }
+}
+
+fn string_argument_content(schema: &Value) -> Format {
+    let Some(schema) = schema.as_object() else {
+        return Format::any_text_excluding(&[ARG_CLOSE, CALL_CLOSE]);
+    };
+    let values = schema
+        .get("enum")
+        .and_then(Value::as_array)
+        .cloned()
+        .or_else(|| schema.get("const").cloned().map(|value| vec![value]));
+    let Some(values) = values else {
+        return Format::any_text_excluding(&[ARG_CLOSE, CALL_CLOSE]);
+    };
+    if values.is_empty()
+        || values.len() > 256
+        || values
+            .iter()
+            .any(|value| value.as_str().is_none_or(|value| value.contains("<|")))
+    {
+        return Format::any_text_excluding(&[ARG_CLOSE, CALL_CLOSE]);
+    }
+    let values = values.iter().filter_map(Value::as_str).collect::<Vec<_>>();
+    match values.as_slice() {
+        [value] => Format::const_string(*value),
+        _ => Format::or(values.into_iter().map(Format::const_string).collect()),
+    }
+}
+
+fn raw_json_arguments(parameters: &Value, options: StructuralTagOptions) -> Format {
+    Format::tag(
+        format!("{JSON_OPEN} type=\"object\"{SEP}"),
+        json_schema(parameters.clone(), options),
+        JSON_CLOSE,
+    )
+}
+
+fn permissive_argument() -> Format {
+    let key = Format::regex(r#"(?:[^<\"&]|&(?:amp|quot);|<[^|])*"#);
+    let alternatives = XTML_TYPES
+        .iter()
+        .map(|xtml_type| {
+            Format::sequence(vec![
+                key.clone(),
+                Format::const_string(format!("\" type=\"{xtml_type}\"{SEP}")),
+                if *xtml_type == "string" {
+                    Format::any_text_excluding(&[ARG_CLOSE, CALL_CLOSE])
+                } else {
+                    Format::json_schema(Value::Bool(true))
+                },
+            ])
+        })
+        .collect();
+    Format::tag(
+        format!("{OPEN}argument key=\""),
+        Format::or(alternatives),
+        ARG_CLOSE,
+    )
+}
+
+fn json_schema(schema: Value, options: StructuralTagOptions) -> Format {
+    Format::JsonSchema(
+        JsonSchemaFormat::new(schema)
+            .with_any_order(options.any_order)
+            .with_max_whitespace_cnt(options.max_whitespace_cnt),
+    )
+}
+
+fn root_definitions(schema: &Map<String, Value>) -> Map<String, Value> {
+    ["$defs", "definitions"]
+        .into_iter()
+        .filter_map(|key| schema.get(key).map(|value| (key.to_string(), value.clone())))
+        .collect()
+}
+
+fn attach_root_definitions(schema: &Value, root_defs: &Map<String, Value>) -> Value {
+    let Some(mut schema) = schema.as_object().cloned() else {
+        return schema.clone();
+    };
+    for (key, value) in root_defs {
+        schema.entry(key.clone()).or_insert_with(|| value.clone());
+    }
+    Value::Object(schema)
+}
+
+fn escape_attr_value(value: &str) -> String {
+    value.replace('&', "&amp;").replace('"', "&quot;")
+}
+
+#[cfg(test)]
+mod tests {
+    use expect_test::expect;
+    use serde_json::json;
+    use xgrammar_structural_tag::builders::StructuralTagOptions;
+    use xgrammar_structural_tag::{
+        FunctionDefinition, FunctionToolParam, ToolChoice, ToolParam, build_structural_tag,
+    };
+
+    use super::KimiK3StructuralTagBuilder;
+
+    fn tool(name: &str, parameters: serde_json::Value) -> ToolParam {
+        ToolParam::Function(FunctionToolParam::new(
+            FunctionDefinition::new(name).with_parameters(parameters),
+        ))
+    }
+
+    #[test]
+    fn required_structural_tag_matches_xtml_channels() {
+        let tools = vec![tool(
+            "get_weather",
+            json!({
+                "$defs": {
+                    "place": { "type": "object", "properties": { "city": { "type": "string" } } }
+                },
+                "type": "object",
+                "properties": {
+                    "unit": { "type": "string", "enum": ["celsius", "fahrenheit"] },
+                    "place": { "$ref": "#/$defs/place", "type": "object" }
+                },
+                "required": ["place"]
+            }),
+        )];
+        let tag = build_structural_tag(
+            KimiK3StructuralTagBuilder,
+            &tools,
+            ToolChoice::required(),
+            StructuralTagOptions::default().with_reasoning(false),
+        )
+        .unwrap();
+
+        expect![[r##"{"type":"structural_tag","format":{"type":"sequence","elements":[{"type":"optional","content":{"type":"const_string","value":"<|open|>response<|sep|>"}},{"type":"tag","begin":"","content":{"type":"any_text","excludes":["<|close|>response<|sep|>","<|open|>tools<|sep|>","<|close|>message<|sep|>","<|end_of_msg|>"]},"end":"<|close|>response<|sep|>"},{"type":"tag","begin":"<|open|>tools<|sep|>","content":{"type":"tags_with_separator","tags":[{"begin":"<|open|>call tool=\"get_weather\" index=\"","content":{"type":"sequence","elements":[{"type":"regex","pattern":"[1-9][0-9]*"},{"type":"const_string","value":"\"<|sep|>"},{"type":"or","elements":[{"type":"plus","content":{"type":"or","elements":[{"type":"tag","begin":"<|open|>argument key=\"unit\" type=\"string\"<|sep|>","content":{"type":"or","elements":[{"type":"const_string","value":"celsius"},{"type":"const_string","value":"fahrenheit"}]},"end":"<|close|>argument<|sep|>"},{"type":"tag","begin":"<|open|>argument key=\"place\" type=\"object\"<|sep|>","content":{"type":"json_schema","json_schema":{"$ref":"#/$defs/place","type":"object","$defs":{"place":{"type":"object","properties":{"city":{"type":"string"}}}}},"style":"json","any_order":false,"max_whitespace_cnt":null},"end":"<|close|>argument<|sep|>"}]}},{"type":"tag","begin":"<|open|>json type=\"object\"<|sep|>","content":{"type":"json_schema","json_schema":{"$defs":{"place":{"type":"object","properties":{"city":{"type":"string"}}}},"type":"object","properties":{"unit":{"type":"string","enum":["celsius","fahrenheit"]},"place":{"$ref":"#/$defs/place","type":"object"}},"required":["place"]},"style":"json","any_order":false,"max_whitespace_cnt":null},"end":"<|close|>json<|sep|>"}]}]},"end":"<|close|>call<|sep|>"}],"separator":"","at_least_one":true,"stop_after_first":false},"end":"<|close|>tools<|sep|>"},{"type":"optional","content":{"type":"const_string","value":"<|close|>message<|sep|>"}}]}}"##]].assert_eq(&tag.to_json_string().unwrap());
+    }
+
+    #[test]
+    fn typed_arguments_require_one_tag_only_for_nonempty_required() {
+        let required = super::typed_arguments(
+            &json!({
+                "type": "object",
+                "properties": { "query": { "type": "string" } },
+                "required": ["query"]
+            }),
+            StructuralTagOptions::default(),
+        );
+        let optional = super::typed_arguments(
+            &json!({
+                "type": "object",
+                "properties": { "query": { "type": "string" } }
+            }),
+            StructuralTagOptions::default(),
+        );
+        let empty_required = super::typed_arguments(
+            &json!({
+                "type": "object",
+                "properties": { "query": { "type": "string" } },
+                "required": []
+            }),
+            StructuralTagOptions::default(),
+        );
+
+        assert_eq!(serde_json::to_value(required).unwrap()["type"], "plus");
+        assert_eq!(serde_json::to_value(optional).unwrap()["type"], "star");
+        assert_eq!(
+            serde_json::to_value(empty_required).unwrap()["type"],
+            "star"
+        );
+    }
+
+    #[test]
+    fn reasoning_grammar_starts_inside_prefilled_think_channel() {
+        let tools = vec![tool("ping", json!({ "type": "object", "properties": {} }))];
+        let tag = build_structural_tag(
+            KimiK3StructuralTagBuilder,
+            &tools,
+            ToolChoice::auto(),
+            StructuralTagOptions::default().with_reasoning(true),
+        )
+        .unwrap();
+        let value = serde_json::to_value(tag).unwrap();
+
+        assert_eq!(
+            value["format"]["elements"][0]["end"],
+            "<|close|>think<|sep|>"
+        );
+        assert_eq!(
+            value["format"]["elements"][1]["value"],
+            "<|open|>response<|sep|>"
+        );
+        assert_eq!(value["format"]["elements"][3]["type"], "optional");
+    }
+
+    #[test]
+    fn forced_choice_keeps_only_the_named_tool() {
+        let tools = vec![
+            tool("search", json!({ "type": "object" })),
+            tool("lookup", json!({ "type": "object" })),
+        ];
+        let tag = build_structural_tag(
+            KimiK3StructuralTagBuilder,
+            &tools,
+            ToolChoice::function("lookup"),
+            StructuralTagOptions::default().with_reasoning(false),
+        )
+        .unwrap()
+        .to_json_string()
+        .unwrap();
+
+        assert!(tag.contains("lookup"));
+        assert!(!tag.contains("search"));
+    }
+
+    #[test]
+    fn union_argument_content_matches_its_xtml_type() {
+        let tools = vec![tool(
+            "set_count",
+            json!({
+                "type": "object",
+                "properties": {
+                    "count": { "type": ["integer", "null"] }
+                }
+            }),
+        )];
+        let tag = build_structural_tag(
+            KimiK3StructuralTagBuilder,
+            &tools,
+            ToolChoice::required(),
+            StructuralTagOptions::default(),
+        )
+        .unwrap()
+        .to_json_string()
+        .unwrap();
+
+        assert!(tag.contains(r#"type=\"number\""#));
+        assert!(tag.contains(r#""json_schema":{"type":"integer"}"#), "{tag}");
+        assert!(tag.contains(r#"type=\"null\""#));
+        assert!(tag.contains(r#""json_schema":{"type":"null"}"#), "{tag}");
+    }
+
+    #[test]
+    fn unsafe_string_enum_falls_back_as_a_whole() {
+        let format = super::string_argument_content(&json!({
+            "type": "string",
+            "enum": ["safe", "<|unsafe"]
+        }));
+
+        assert_eq!(serde_json::to_value(format).unwrap()["type"], "any_text");
+    }
+}

rust/src/parser/src/unified/mod.rs

@@ -6,10 +6,12 @@
 mod combined;
 mod gemma4;
 mod inkling;
+mod kimi_k3;
 
 pub use combined::CombinedParser;
 pub use gemma4::Gemma4UnifiedParser;
 pub use inkling::InklingUnifiedParser;
+pub use kimi_k3::{KimiK3StructuralTagBuilder, KimiK3UnifiedParser};
 use thiserror::Error;
 use thiserror_ext::Macro;
 use vllm_tokenizer::DynTokenizer;

rust/src/tokenizer/src/tiktoken.rs

@@ -525,7 +525,7 @@ fn detect_bpe_pattern(config: &TiktokenModelConfig) -> &'static str {
     let model_type = config.effective_model_type();
 
     match model_type {
-        Some("kimi" | "kimi_k2" | "kimi_k25" | "deepseek_v3") => KIMI_PATTERN,
+        Some("kimi" | "kimi_k2" | "kimi_k25" | "kimi_k3" | "deepseek_v3") => KIMI_PATTERN,
         _ => CL100K_BASE_PATTERN,
     }
 }
@@ -777,6 +777,7 @@ mod tests {
     #[test]
     fn tiktoken_detects_kimi_pattern_from_model_type() {
         let kimi = config_json!({ "model_type": "kimi_k25" });
+        let kimi_k3 = config_json!({ "model_type": "kimi_k3" });
         let baseten_kimi = config_json!({ "model_type": "deepseek_v3" });
         let nested_kimi = config_json!({
             "model_type": "composite_wrapper",
@@ -790,6 +791,7 @@ mod tests {
         let missing = config_json!({ "text_config": {} });
 
         assert_eq!(detect_bpe_pattern(&kimi), KIMI_PATTERN);
+        assert_eq!(detect_bpe_pattern(&kimi_k3), KIMI_PATTERN);
         assert_eq!(detect_bpe_pattern(&baseten_kimi), KIMI_PATTERN);
         assert_eq!(detect_bpe_pattern(&nested_kimi), CL100K_BASE_PATTERN);
         assert_eq!(detect_bpe_pattern(&generic), CL100K_BASE_PATTERN);

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/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/entrypoints/test_kimi_k3_chat_utils.py

@@ -0,0 +1,95 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from vllm.models.kimi_k3.nvidia.model import KimiK3MultiModalProcessor
+
+
+def _tool_call(name: str, call_id: str | None = None) -> dict:
+    tool_call = {
+        "type": "function",
+        "function": {"name": name, "arguments": "{}"},
+    }
+    if call_id is not None:
+        tool_call["id"] = call_id
+    return tool_call
+
+
+_preprocess = KimiK3MultiModalProcessor.preprocess_messages
+
+
+def test_preprocess_messages_orders_and_fills_attrs():
+    messages = [
+        {
+            "role": "assistant",
+            "content": "",
+            "tool_calls": [
+                _tool_call("lookup", "lookup:0"),
+                _tool_call("lookup", "lookup:1"),
+            ],
+        },
+        {"role": "tool", "tool_call_id": "lookup:1", "content": "second"},
+        {"role": "tool", "tool_call_id": "lookup:0", "content": "first"},
+    ]
+
+    normalized = _preprocess(messages)
+
+    assert [message["content"] for message in normalized[1:]] == [
+        "first",
+        "second",
+    ]
+    assert normalized[1]["tool"] == "lookup"
+    assert normalized[1]["index"] == 1
+    assert normalized[2]["tool"] == "lookup"
+    assert normalized[2]["index"] == 2
+
+
+def test_preprocess_messages_accepts_synthetic_aliases():
+    messages = [
+        {
+            "role": "assistant",
+            "content": "",
+            "tool_calls": [
+                _tool_call("search"),
+                _tool_call("fetch"),
+            ],
+        },
+        {"role": "tool", "tool_call_id": "fetch:1", "content": "b"},
+        {"role": "tool", "tool_call_id": "search:0", "content": "a"},
+    ]
+
+    normalized = _preprocess(messages)
+
+    assert [message["content"] for message in normalized[1:]] == ["a", "b"]
+    assert normalized[1]["tool"] == "search"
+    assert normalized[1]["index"] == 1
+    assert normalized[2]["tool"] == "fetch"
+    assert normalized[2]["index"] == 2
+
+
+def test_preprocess_messages_fallback_on_unknown_id():
+    messages = [
+        {
+            "role": "assistant",
+            "content": "",
+            "tool_calls": [
+                _tool_call("search"),
+                _tool_call("fetch"),
+            ],
+        },
+        {"role": "tool", "tool_call_id": "functions.fetch:1", "content": "b"},
+    ]
+
+    assert _preprocess(messages) == messages
+
+
+def test_preprocess_messages_leaves_unknown_block_unchanged():
+    messages = [
+        {
+            "role": "assistant",
+            "content": "",
+            "tool_calls": [_tool_call("lookup", "lookup:0")],
+        },
+        {"role": "tool", "tool_call_id": "unknown", "content": "result"},
+    ]
+
+    assert _preprocess(messages) == messages

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_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_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,657 @@
+# 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)."""
+
+import re
+from pathlib import Path
+from types import SimpleNamespace
+
+import pytest
+import torch
+from torch import nn
+
+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()
+    residual_warmup_configs = 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.cuda.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.cuda.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.cuda.synchronize()
+
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph):
+        output = k3_gemm.try_low_latency_gemm(x, weight)
+    graph.replay()
+    torch.cuda.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.cuda.synchronize()
+
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph):
+        output = k3_gemm.shape_dynamic_skinny_gemm(x, weight, config, residual)
+    graph.replay()
+    torch.cuda.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: object) -> 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/models/kimi_k3/test_amd_attn_res.py

@@ -0,0 +1,102 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import pytest
+import torch
+import torch.nn.functional as F
+
+from vllm.models.kimi_k3.amd.ops.attn_res import attn_res
+from vllm.platforms import current_platform
+
+pytestmark = pytest.mark.skipif(
+    not current_platform.is_rocm(),
+    reason="AMD AttnRes requires ROCm",
+)
+
+
+def _randn_with_row_padding(*shape: int, padding: int = 0) -> torch.Tensor:
+    storage = torch.randn(
+        *shape[:-1],
+        shape[-1] + padding,
+        device="cuda",
+        dtype=torch.bfloat16,
+    )
+    return storage[..., : shape[-1]]
+
+
+def _reference(
+    prefix: torch.Tensor,
+    blocks: torch.Tensor,
+    norm_weight: torch.Tensor,
+    qk_weight: torch.Tensor,
+    num_blocks: int,
+    eps: float,
+) -> torch.Tensor:
+    hidden_size = prefix.shape[-1]
+    values = torch.cat((blocks[:, :num_blocks], prefix.unsqueeze(1)), dim=1)
+    keys = F.rms_norm(values, (hidden_size,), norm_weight, eps)
+    probs = (keys @ qk_weight).softmax(dim=-1)
+    return torch.matmul(probs.unsqueeze(1), values).squeeze(1)
+
+
+@pytest.mark.parametrize(
+    (
+        "num_tokens",
+        "num_blocks",
+        "block_capacity",
+        "hidden_size",
+        "row_padding",
+    ),
+    [
+        pytest.param(0, 3, 5, 128, 0, id="empty"),
+        pytest.param(1, 1, 2, 128, 0, id="decode-single"),
+        pytest.param(17, 4, 6, 1024, 7, id="decode-padded"),
+        pytest.param(320, 8, 10, 7168, 0, id="prefill-full"),
+    ],
+)
+def test_amd_attn_res_matches_reference(
+    num_tokens: int,
+    num_blocks: int,
+    block_capacity: int,
+    hidden_size: int,
+    row_padding: int,
+) -> None:
+    eps = 1e-5
+    prefix = _randn_with_row_padding(num_tokens, hidden_size, padding=row_padding)
+    blocks = _randn_with_row_padding(
+        num_tokens,
+        block_capacity,
+        hidden_size,
+        padding=row_padding,
+    )
+    norm_weight = 1 + 0.1 * torch.randn(
+        hidden_size, device="cuda", dtype=torch.bfloat16
+    )
+    qk_weight = (
+        torch.randn(hidden_size, device="cuda", dtype=torch.bfloat16) / hidden_size**0.5
+    )
+    expected = _reference(
+        prefix,
+        blocks,
+        norm_weight,
+        qk_weight,
+        num_blocks,
+        eps,
+    )
+    original_prefix = prefix.clone()
+    original_blocks = blocks.clone()
+
+    actual = attn_res(
+        prefix,
+        blocks,
+        norm_weight,
+        qk_weight,
+        num_blocks,
+        eps,
+    )
+
+    torch.testing.assert_close(actual, expected, atol=8e-2, rtol=3e-2)
+    torch.testing.assert_close(prefix, original_prefix, atol=0, rtol=0)
+    torch.testing.assert_close(blocks, original_blocks, atol=0, rtol=0)
+    assert actual.shape == prefix.shape
+    assert actual.is_contiguous()

tests/models/kimi_k3/test_attn_res.py

@@ -0,0 +1,226 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import pytest
+import torch
+import torch.nn.functional as F
+
+from 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
+
+HIDDEN_SIZE = 7168
+MAX_BLOCKS = 8
+EPS = 1e-5
+
+
+def _randn_with_row_padding(*shape: int, padding: int = 0) -> torch.Tensor:
+    storage = torch.randn(
+        *shape[:-1],
+        shape[-1] + padding,
+        device="cuda",
+        dtype=torch.bfloat16,
+    )
+    return storage[..., : shape[-1]]
+
+
+def _reference(
+    prefix: torch.Tensor,
+    delta: torch.Tensor | None,
+    blocks: torch.Tensor,
+    norm_weight: torch.Tensor,
+    qk_weight: torch.Tensor,
+    output_norm_weight: torch.Tensor | None,
+    num_blocks: int,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    if delta is not None:
+        prefix = prefix + delta
+    values = torch.cat((blocks[:, :num_blocks], prefix.unsqueeze(1)), dim=1)
+    keys = F.rms_norm(values, (HIDDEN_SIZE,), norm_weight, EPS)
+    probs = (keys @ qk_weight).softmax(dim=-1)
+    output = torch.matmul(probs.unsqueeze(1), values).squeeze(1)
+    if output_norm_weight is not None:
+        output = F.rms_norm(output, (HIDDEN_SIZE,), output_norm_weight, EPS)
+    return output, prefix
+
+
+@pytest.mark.parametrize(
+    (
+        "num_tokens",
+        "num_blocks",
+        "row_padding",
+        "write_block",
+        "has_delta",
+        "backend",
+    ),
+    [
+        pytest.param(1, 0, 0, True, False, "triton", id="triton-empty"),
+        pytest.param(1, 0, 0, True, True, "triton", id="triton-empty-add"),
+        pytest.param(17, 5, 7, True, False, "triton", id="triton-write"),
+        pytest.param(17, 5, 7, True, True, "triton", id="triton-write-add"),
+        pytest.param(3, 8, 0, False, False, "triton", id="triton-full"),
+        pytest.param(3, 8, 0, False, True, "triton", id="triton-full-add"),
+        pytest.param(320, 1, 0, False, True, "nvidia", id="nvidia-1"),
+        pytest.param(320, 4, 0, False, True, "nvidia", id="nvidia-4"),
+        pytest.param(320, 8, 0, False, True, "nvidia", id="nvidia-8"),
+    ],
+)
+def test_attn_res(
+    num_tokens: int,
+    num_blocks: int,
+    row_padding: int,
+    write_block: bool,
+    has_delta: bool,
+    backend: str,
+):
+    if backend == "nvidia" and not current_platform.is_device_capability_family(100):
+        pytest.skip("NVIDIA AttnRes requires the SM100 family")
+
+    prefix = _randn_with_row_padding(num_tokens, HIDDEN_SIZE, padding=row_padding)
+    delta = (
+        _randn_with_row_padding(num_tokens, HIDDEN_SIZE, padding=row_padding)
+        if has_delta
+        else None
+    )
+    blocks = _randn_with_row_padding(
+        num_tokens, MAX_BLOCKS, HIDDEN_SIZE, padding=row_padding
+    )
+    norm_weight = 1 + 0.1 * torch.randn(
+        HIDDEN_SIZE, device="cuda", dtype=torch.bfloat16
+    )
+    qk_weight = (
+        torch.randn(HIDDEN_SIZE, device="cuda", dtype=torch.bfloat16) / HIDDEN_SIZE**0.5
+    )
+    output_norm_weight = 1 + 0.1 * torch.randn(
+        HIDDEN_SIZE, device="cuda", dtype=torch.bfloat16
+    )
+    original_blocks = blocks.clone()
+    expected, expected_prefix = _reference(
+        prefix.clone(),
+        delta,
+        blocks,
+        norm_weight,
+        qk_weight,
+        output_norm_weight,
+        num_blocks,
+    )
+    block_write_idx = num_blocks if write_block else -1
+
+    actual = attn_res(
+        prefix,
+        delta,
+        blocks,
+        norm_weight,
+        qk_weight,
+        output_norm_weight,
+        num_blocks,
+        block_write_idx,
+        EPS,
+        EPS,
+    )
+
+    torch.testing.assert_close(actual, expected, atol=8e-2, rtol=3e-2)
+    torch.testing.assert_close(prefix, expected_prefix, atol=0, rtol=0)
+    if write_block:
+        original_blocks[:, block_write_idx].copy_(expected_prefix)
+    torch.testing.assert_close(blocks, original_blocks, atol=0, rtol=0)
+    assert actual.is_contiguous()
+
+
+@pytest.mark.parametrize("num_blocks", range(MAX_BLOCKS + 1))
+def test_attn_res_block_counts(num_blocks: int):
+    prefix = torch.randn(1, HIDDEN_SIZE, device="cuda", dtype=torch.bfloat16)
+    blocks = torch.randn(
+        1, MAX_BLOCKS, HIDDEN_SIZE, device="cuda", dtype=torch.bfloat16
+    )
+    norm_weight = torch.ones(HIDDEN_SIZE, device="cuda", dtype=torch.bfloat16)
+    qk_weight = (
+        torch.randn(HIDDEN_SIZE, device="cuda", dtype=torch.bfloat16) / HIDDEN_SIZE**0.5
+    )
+    output_norm_weight = torch.ones_like(norm_weight)
+    expected, _ = _reference(
+        prefix.clone(),
+        None,
+        blocks,
+        norm_weight,
+        qk_weight,
+        output_norm_weight,
+        num_blocks,
+    )
+
+    actual = attn_res(
+        prefix,
+        None,
+        blocks,
+        norm_weight,
+        qk_weight,
+        output_norm_weight,
+        num_blocks,
+        -1,
+        EPS,
+        EPS,
+    )
+
+    torch.testing.assert_close(actual, expected, atol=8e-2, rtol=3e-2)
+
+
+def test_attn_res_without_output_norm():
+    prefix = torch.randn(7, HIDDEN_SIZE, device="cuda", dtype=torch.bfloat16)
+    delta = torch.randn_like(prefix)
+    blocks = torch.randn(
+        7, MAX_BLOCKS, HIDDEN_SIZE, device="cuda", dtype=torch.bfloat16
+    )
+    norm_weight = torch.randn(HIDDEN_SIZE, device="cuda", dtype=torch.bfloat16)
+    qk_weight = (
+        torch.randn(HIDDEN_SIZE, device="cuda", dtype=torch.bfloat16) / HIDDEN_SIZE**0.5
+    )
+    expected, _ = _reference(
+        prefix.clone(), delta, blocks, norm_weight, qk_weight, None, MAX_BLOCKS
+    )
+
+    actual = attn_res(
+        prefix,
+        delta,
+        blocks,
+        norm_weight,
+        qk_weight,
+        None,
+        MAX_BLOCKS,
+        -1,
+        EPS,
+        0.0,
+    )
+
+    torch.testing.assert_close(actual, expected, atol=8e-2, rtol=3e-2)
+
+
+@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,330 @@
+# 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
+
+
+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))
+
+
+def test_moe_sequence_parallel_is_available_without_data_parallel():
+    parallel_config = ParallelConfig(
+        tensor_parallel_size=2,
+        data_parallel_size=1,
+        enable_expert_parallel=True,
+        all2all_backend="allgather_reducescatter",
+    )
+
+    assert parallel_config.use_sequence_parallel_moe
+
+
+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/registry.py

@@ -1024,6 +1024,11 @@ _MULTIMODAL_EXAMPLE_MODELS = {
         "moonshotai/Kimi-K2.5",
         trust_remote_code=True,
     ),
+    "KimiK3ForConditionalGeneration": _HfExamplesInfo(
+        "moonshotai/Kimi-K3",
+        trust_remote_code=True,
+        is_available_online=False,
+    ),
     "KimiVLForConditionalGeneration": _HfExamplesInfo(
         "moonshotai/Kimi-VL-A3B-Instruct",
         extras={"thinking": "moonshotai/Kimi-VL-A3B-Thinking"},
@@ -1658,6 +1663,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/parser/test_streaming.py

@@ -13,6 +13,8 @@ from vllm.parser.engine.registered_adapters import Qwen3ParserReasoningAdapter
 from vllm.reasoning.basic_parsers import BaseThinkingReasoningParser
 from vllm.tool_parsers.hermes_tool_parser import Hermes2ProToolParser
 
+pytestmark = pytest.mark.skip_global_cleanup
+
 
 class ThinkReasoningParser(BaseThinkingReasoningParser):
     @property
@@ -251,6 +253,95 @@ def test_parse_delta_reasoning_boundary_no_tool_parser(tokenizer, request_obj):
     assert "get_weather" in content
 
 
+class _SplitCloseReasoningParser:
+    close_ids = [2, 3]
+
+    @staticmethod
+    def _close_index(input_ids: list[int]) -> int:
+        for i in range(len(input_ids) - 1):
+            if input_ids[i : i + 2] == _SplitCloseReasoningParser.close_ids:
+                return i
+        return -1
+
+    def is_reasoning_end(self, input_ids: list[int]) -> bool:
+        return self._close_index(input_ids) != -1
+
+    def is_reasoning_end_streaming(
+        self,
+        input_ids: list[int],
+        delta_ids: list[int],
+    ) -> bool:
+        return self.is_reasoning_end(input_ids)
+
+    def extract_content_ids(self, input_ids: list[int]) -> list[int]:
+        close_index = self._close_index(input_ids)
+        if close_index == -1:
+            return []
+        return input_ids[close_index + len(self.close_ids) :]
+
+    def extract_reasoning_streaming(self, *args, **kwargs):
+        current_token_ids = kwargs["current_token_ids"] if kwargs else args[4]
+        if self.is_reasoning_end(current_token_ids):
+            return DeltaMessage(content="tail")
+        return DeltaMessage(reasoning=kwargs["delta_text"] if kwargs else args[2])
+
+    def extract_reasoning(self, model_output, request):
+        return None, model_output
+
+
+class _RecordingToolParser:
+    supports_required_and_named = False
+
+    def __init__(self):
+        self.current_token_ids: list[list[int]] = []
+
+    def extract_tool_calls(self, model_output, request):
+        raise AssertionError("non-streaming path is not used")
+
+    def extract_tool_calls_streaming(
+        self,
+        previous_text,
+        current_text,
+        delta_text,
+        previous_token_ids,
+        current_token_ids,
+        delta_token_ids,
+        request,
+    ):
+        self.current_token_ids.append(list(current_token_ids))
+        return DeltaMessage(content=delta_text)
+
+
+class _SplitCloseParser(DelegatingParser):
+    def __init__(self):
+        super().__init__(tokenizer=None)
+        self.reasoning_parser = _SplitCloseReasoningParser()
+        self.recording_tool_parser = _RecordingToolParser()
+        self.tool_parser = self.recording_tool_parser
+
+
+def test_parse_delta_extracts_content_ids_from_full_current_ids(request_obj):
+    parser = _SplitCloseParser()
+
+    parser.parse_delta(
+        "reasoning[close]",
+        [9, 2],
+        request_obj,
+        prompt_token_ids=[],
+        finished=False,
+    )
+    result = parser.parse_delta(
+        "think[sep]tail",
+        [3, 4],
+        request_obj,
+        finished=False,
+    )
+
+    assert result is not None
+    assert result.content == "tail"
+    assert parser.recording_tool_parser.current_token_ids == [[4]]
+
+
 def test_parse_delta_reasoning_only_no_think_leak(tokenizer, request_obj):
     """Regression: </think> must not leak into content when streaming
     token-by-token with reasoning=True, tool=False."""

tests/reasoning/test_kimi_k3_reasoning_parser.py

@@ -0,0 +1,239 @@
+# 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.abstract_parser import DelegatingParser
+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(DelegatingParser):
+    reasoning_parser_cls = KimiK3ReasoningParser
+
+
+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"
+
+
+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_kimi_k3.py

@@ -0,0 +1,167 @@
+# 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.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]] = []
+
+    def apply_chat_template(self, conversation, **kwargs) -> list[int]:
+        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 kwargs["thinking_effort"] == "none"
+    assert "enable_thinking" not in kwargs
+    assert "reasoning_effort" not in kwargs
+
+
+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_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"
+
+
+@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/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"),
     [

tests/tool_parsers/test_kimi_k3_named_tool_choice.py

@@ -0,0 +1,62 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Named tool choice for Kimi K3: allowed when the XTML structural tag is
+attached (strict tool calling), rejected otherwise."""
+
+import pytest
+
+from vllm.entrypoints.openai.chat_completion.protocol import (
+    ChatCompletionRequest,
+)
+from vllm.exceptions import VLLMValidationError
+from vllm.sampling_params import StructuredOutputsParams
+
+
+class _DummyTokenizer:
+    def get_vocab(self):
+        return {}
+
+    def encode(self, text, add_special_tokens=False):
+        return [ord(c) for c in text]
+
+
+def _request(with_tag: bool) -> ChatCompletionRequest:
+    req = ChatCompletionRequest(
+        model="k3",
+        messages=[{"role": "user", "content": "hi"}],
+        tools=[
+            {
+                "type": "function",
+                "function": {
+                    "name": "get_weather",
+                    "parameters": {
+                        "type": "object",
+                        "properties": {"city": {"type": "string"}},
+                        "required": ["city"],
+                    },
+                },
+            }
+        ],
+        tool_choice={"type": "function", "function": {"name": "get_weather"}},
+    )
+    if with_tag:
+        req.structured_outputs = StructuredOutputsParams(
+            structural_tag='{"type": "structural_tag", "format": {}}'
+        )
+    return req
+
+
+def _parser():
+    from vllm.tool_parsers.kimi_k3_tool_parser import KimiK3ToolParser
+
+    return KimiK3ToolParser(_DummyTokenizer())
+
+
+def test_named_choice_allowed_with_structural_tag():
+    req = _parser().adjust_request(_request(with_tag=True))
+    assert req.skip_special_tokens is False
+
+
+def test_named_choice_rejected_without_structural_tag():
+    with pytest.raises(VLLMValidationError):
+        _parser().adjust_request(_request(with_tag=False))

tests/tool_parsers/test_structural_tag_registry.py

@@ -5,7 +5,8 @@ from types import SimpleNamespace
 from unittest.mock import MagicMock
 
 import pytest
-from xgrammar import StructuralTag
+from xgrammar import Grammar, StructuralTag
+from xgrammar.testing import _is_grammar_accept_string
 
 from vllm.entrypoints.openai.chat_completion.protocol import (
     ChatCompletionNamedFunction,
@@ -24,6 +25,7 @@ from vllm.tool_parsers.deepseekv32_engine_tool_parser import (
 from vllm.tool_parsers.glm47_moe_tool_parser import Glm47MoeModelToolParser
 from vllm.tool_parsers.hermes_tool_parser import Hermes2ProToolParser
 from vllm.tool_parsers.kimi_k2_tool_parser import KimiK2ToolParser
+from vllm.tool_parsers.kimi_k3_tool_parser import KimiK3ToolParser
 from vllm.tool_parsers.llama_tool_parser import Llama3JsonToolParser
 from vllm.tool_parsers.minimax_m2_tool_parser import MinimaxM2ToolParser
 from vllm.tool_parsers.qwen3_engine_tool_parser import Qwen3EngineToolParser
@@ -165,6 +167,196 @@ def test_hermes_required_tool_calls_use_empty_separator():
     assert tag.format.separator == ""
 
 
+# ---------------------------------------------------------------------------
+# Kimi K3 (XTML channel format) structural tag
+# ---------------------------------------------------------------------------
+_K3_RESPONSE_OPEN = "<|open|>response<|sep|>"
+_K3_RESPONSE_CLOSE = "<|close|>response<|sep|>"
+_K3_TOOLS_OPEN = "<|open|>tools<|sep|>"
+_K3_TOOLS_CLOSE = "<|close|>tools<|sep|>"
+_K3_CALL_CLOSE = "<|close|>call<|sep|>"
+_K3_ARG_CLOSE = "<|close|>argument<|sep|>"
+_K3_MESSAGE_CLOSE = "<|close|>message<|sep|>"
+
+
+def _k3_tools_by_name() -> list[ChatCompletionToolsParam]:
+    return [
+        ChatCompletionToolsParam(
+            type="function",
+            function={
+                "name": "get_weather",
+                "parameters": {
+                    "type": "object",
+                    "properties": {
+                        "city": {"type": "string"},
+                        "days": {"type": "integer"},
+                    },
+                    "required": ["city"],
+                },
+            },
+        ),
+        ChatCompletionToolsParam(
+            type="function",
+            function={
+                "name": "run_command",
+                "parameters": {
+                    "type": "object",
+                    "properties": {"command": {"type": "string"}},
+                    "required": ["command"],
+                },
+            },
+        ),
+    ]
+
+
+def _k3_arg(key: str, typ: str, val: str) -> str:
+    return f'<|open|>argument key="{key}" type="{typ}"<|sep|>{val}{_K3_ARG_CLOSE}'
+
+
+def _k3_call(name: str, args: str, idx: int = 1) -> str:
+    return f'<|open|>call tool="{name}" index="{idx}"<|sep|>{args}{_K3_CALL_CLOSE}'
+
+
+def _k3_response(content: str = "") -> str:
+    return f"{_K3_RESPONSE_OPEN}{content}{_K3_RESPONSE_CLOSE}"
+
+
+def _k3_tools(*calls: str) -> str:
+    return f"{_K3_TOOLS_OPEN}{''.join(calls)}{_K3_TOOLS_CLOSE}"
+
+
+def _k3_grammar(tool_choice, tools=None):
+    tag = get_model_structural_tag(
+        model="kimi_k3",
+        tools=tools if tools is not None else _k3_tools_by_name(),
+        tool_choice=tool_choice,
+        reasoning=False,
+    )
+    assert isinstance(tag, StructuralTag)
+    return Grammar.from_structural_tag(tag)
+
+
+def test_kimi_k3_registered_as_vllm_builtin():
+    assert "kimi_k3" in VLLM_BUILTIN_STRUCTURAL_TAG_MODELS
+    assert KimiK3ToolParser.structural_tag_model == "kimi_k3"
+
+
+def test_kimi_k3_auto_without_strict_is_unconstrained():
+    # auto + no strict tool => no structural tag (matches the strict gate).
+    tag = get_model_structural_tag(
+        model="kimi_k3",
+        tools=_k3_tools_by_name(),
+        tool_choice="auto",
+        reasoning=False,
+    )
+    assert tag is None
+
+
+@pytest.mark.parametrize(
+    "body",
+    [
+        # single required arg
+        _k3_response()
+        + _k3_tools(_k3_call("get_weather", _k3_arg("city", "string", "Paris"))),
+        # response content + two args (string + number)
+        _k3_response("Checking.")
+        + _k3_tools(
+            _k3_call(
+                "get_weather",
+                _k3_arg("city", "string", "Paris") + _k3_arg("days", "number", "3"),
+            )
+        ),
+        # args in reverse order (parser is order-agnostic)
+        _k3_response()
+        + _k3_tools(
+            _k3_call(
+                "get_weather",
+                _k3_arg("days", "number", "3") + _k3_arg("city", "string", "Paris"),
+            )
+        ),
+        # two calls, second tool
+        _k3_response()
+        + _k3_tools(
+            _k3_call("get_weather", _k3_arg("city", "string", "Paris"), 1),
+            _k3_call("run_command", _k3_arg("command", "string", "ls -la"), 2),
+        ),
+        # string value with regex metacharacters / spaces
+        _k3_response()
+        + _k3_tools(
+            _k3_call(
+                "run_command", _k3_arg("command", "string", "grep -E 'a|b{2,}' x.py")
+            )
+        ),
+        # trailing message-close marker (model's natural turn terminator)
+        _k3_response()
+        + _k3_tools(_k3_call("get_weather", _k3_arg("city", "string", "Paris")))
+        + _K3_MESSAGE_CLOSE,
+        # non-thinking mode: response-open is the prompt prefix, so it is absent
+        _K3_RESPONSE_CLOSE
+        + _k3_tools(_k3_call("get_weather", _k3_arg("city", "string", "Paris"))),
+    ],
+)
+def test_kimi_k3_required_accepts_valid_tool_calls(body: str):
+    assert _is_grammar_accept_string(_k3_grammar("required"), body)
+
+
+@pytest.mark.parametrize(
+    "body",
+    [
+        # unknown tool name
+        _k3_response()
+        + _k3_tools(_k3_call("get_wether", _k3_arg("city", "string", "x"))),
+        # number arg given a non-numeric JSON value
+        _k3_response()
+        + _k3_tools(_k3_call("get_weather", _k3_arg("days", "number", "abc"))),
+        # undeclared argument key
+        _k3_response()
+        + _k3_tools(
+            _k3_call(
+                "get_weather",
+                _k3_arg("city", "string", "Paris") + _k3_arg("zzz", "string", "x"),
+            )
+        ),
+        # required schema but no argument tags
+        _k3_response() + _k3_tools(_k3_call("get_weather", "")),
+        # missing tools close marker
+        _k3_response()
+        + _K3_TOOLS_OPEN
+        + _k3_call("get_weather", _k3_arg("city", "string", "Paris")),
+        # required but no tool call
+        _k3_response("hello"),
+    ],
+)
+def test_kimi_k3_required_rejects_invalid(body: str):
+    assert not _is_grammar_accept_string(_k3_grammar("required"), body)
+
+
+def test_kimi_k3_schema_without_required_accepts_empty_call():
+    tools = [
+        ChatCompletionToolsParam(
+            type="function",
+            function={
+                "name": "get_weather",
+                "parameters": {
+                    "type": "object",
+                    "properties": {"city": {"type": "string"}},
+                },
+            },
+        )
+    ]
+    grammar = _k3_grammar("required", tools=tools)
+    body = _k3_response() + _k3_tools(_k3_call("get_weather", ""))
+
+    assert _is_grammar_accept_string(grammar, body)
+
+
+def test_kimi_k3_auto_strict_allows_response_only(sample_tools_strict):
+    # With a strict tool the tag is built; the tools channel is optional so a
+    # plain response (no tool call) is still valid.
+    grammar = _k3_grammar("auto", tools=sample_tools_strict)
+    assert _is_grammar_accept_string(grammar, _k3_response("Just answering."))
+
+
 @pytest.mark.parametrize("model", sorted(XGRAMMAR_BUILTIN_STRUCTURAL_TAG_MODELS))
 def test_get_model_structural_tag_supports_named_tool_choice(
     model: str,
@@ -339,3 +531,151 @@ def test_get_function_parameters_relaxes_function_strict_false():
     )
 
     assert _get_function_parameters(function) is True
+
+
+def _k3_tools_with_root_defs() -> list[ChatCompletionToolsParam]:
+    return [
+        ChatCompletionToolsParam(
+            type="function",
+            function={
+                "name": "make_config",
+                "parameters": {
+                    "type": "object",
+                    "properties": {
+                        "config": {
+                            "type": "object",
+                            "properties": {
+                                "build": {"$ref": "#/$defs/build"},
+                                "index": {"type": "string"},
+                            },
+                            "required": ["index"],
+                            "additionalProperties": False,
+                        },
+                    },
+                    "required": ["config"],
+                    "$defs": {
+                        "build": {
+                            "type": "object",
+                            "properties": {"outDir": {"type": "string"}},
+                            "additionalProperties": False,
+                        }
+                    },
+                },
+            },
+        )
+    ]
+
+
+def test_kimi_k3_property_ref_to_root_defs_compiles_and_accepts():
+    # Root-level $defs referenced from inside a property schema (the walle
+    # TestReferences shape). Slicing the property out of the parameters
+    # document orphans "#/$defs/..." unless the builder re-attaches $defs;
+    # before the fix Grammar.from_structural_tag raised on the dangling ref.
+    grammar = _k3_grammar("required", tools=_k3_tools_with_root_defs())
+
+    body = _k3_response() + _k3_tools(
+        _k3_call(
+            "make_config",
+            _k3_arg(
+                "config",
+                "object",
+                '{"build": {"outDir": "dist"}, "index": "a.html"}',
+            ),
+        )
+    )
+    assert _is_grammar_accept_string(grammar, body)
+
+
+def _k3_tools_with_string_enum() -> list[ChatCompletionToolsParam]:
+    return [
+        ChatCompletionToolsParam(
+            type="function",
+            function={
+                "name": "set_unit",
+                "parameters": {
+                    "type": "object",
+                    "properties": {
+                        "unit": {
+                            "type": "string",
+                            "enum": ["celsius", " fahrenheit", "\tkelvin"],
+                        },
+                    },
+                    "required": ["unit"],
+                },
+            },
+        )
+    ]
+
+
+@pytest.mark.parametrize("value", ["celsius", " fahrenheit", "\tkelvin"])
+def test_kimi_k3_string_enum_accepts_exact_values(value: str):
+    # Raw string channel with enum: constrained to the exact enum values,
+    # including leading-whitespace variants the model otherwise flubs.
+    grammar = _k3_grammar("required", tools=_k3_tools_with_string_enum())
+    body = _k3_response() + _k3_tools(
+        _k3_call("set_unit", _k3_arg("unit", "string", value))
+    )
+    assert _is_grammar_accept_string(grammar, body)
+
+
+@pytest.mark.parametrize("value", ["kelvin", "Celsius", "celsius ", ""])
+def test_kimi_k3_string_enum_rejects_non_members(value: str):
+    grammar = _k3_grammar("required", tools=_k3_tools_with_string_enum())
+    body = _k3_response() + _k3_tools(
+        _k3_call("set_unit", _k3_arg("unit", "string", value))
+    )
+    assert not _is_grammar_accept_string(grammar, body)
+
+
+def _k3_tools_with_maxlen() -> list[ChatCompletionToolsParam]:
+    return [
+        ChatCompletionToolsParam(
+            type="function",
+            function={
+                "name": "set_note",
+                "parameters": {
+                    "type": "object",
+                    "properties": {
+                        "note": {"type": "string", "maxLength": 8, "minLength": 2},
+                    },
+                    "required": ["note"],
+                },
+            },
+        )
+    ]
+
+
+def test_kimi_k3_string_maxlength_bounds_raw_channel():
+    # Raw string channel with maxLength/minLength: enforced via a bounded
+    # regex that keeps the "<|" marker prefix unambiguous but still allows
+    # a bare '<' inside values.
+    grammar = _k3_grammar("required", tools=_k3_tools_with_maxlen())
+
+    def body(val: str) -> str:
+        return _k3_response() + _k3_tools(
+            _k3_call("set_note", _k3_arg("note", "string", val))
+        )
+
+    assert _is_grammar_accept_string(grammar, body("ab"))
+    assert _is_grammar_accept_string(grammar, body("a<b then"))
+    assert not _is_grammar_accept_string(grammar, body("way too long note"))
+    assert not _is_grammar_accept_string(grammar, body("a"))  # under minLength
+
+
+def test_kimi_k3_forced_tool_choice_builds_single_mandatory_call():
+    # Named tool choice normalizes to "forced": the tag must require exactly
+    # the named tool's call (no response-only escape).
+    grammar = _k3_grammar(
+        ChatCompletionNamedToolChoiceParam(
+            type="function",
+            function=ChatCompletionNamedFunction(name="get_weather"),
+        ),
+        tools=_k3_tools_by_name(),
+    )
+
+    ok = _k3_response() + _k3_tools(
+        _k3_call("get_weather", _k3_arg("city", "string", "Paris"))
+    )
+    response_only = _k3_response("no call here")
+    assert _is_grammar_accept_string(grammar, ok)
+    assert not _is_grammar_accept_string(grammar, response_only)

tests/tool_use/test_kimi_k3_tool_parser.py

@@ -0,0 +1,530 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import json
+
+import pytest
+
+from vllm.entrypoints.openai.chat_completion.protocol import (
+    ChatCompletionRequest,
+)
+from vllm.entrypoints.openai.engine.protocol import DeltaMessage
+from vllm.entrypoints.openai.responses.protocol import ResponsesRequest
+from vllm.entrypoints.openai.responses.utils import build_response_output_items
+from vllm.exceptions import VLLMValidationError
+from vllm.parser.abstract_parser import DelegatingParser
+from vllm.reasoning.kimi_k3_reasoning_parser import KimiK3ReasoningParser
+from vllm.tool_parsers.kimi_k3_tool_parser import KimiK3ToolParser
+
+OPEN = "<|open|>"
+CLOSE = "<|close|>"
+SEP = "<|sep|>"
+THINK_OPEN = f"{OPEN}think{SEP}"
+THINK_CLOSE = f"{CLOSE}think{SEP}"
+RESPONSE_CLOSE = f"{CLOSE}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 KimiK3DelegatingParser(DelegatingParser):
+    reasoning_parser_cls = KimiK3ReasoningParser
+    tool_parser_cls = KimiK3ToolParser
+
+
+def _request() -> ChatCompletionRequest:
+    return ChatCompletionRequest(
+        model="test-model",
+        messages=[],
+        tools=[
+            {
+                "type": "function",
+                "function": {
+                    "name": "calc",
+                    "parameters": {"type": "object", "properties": {}},
+                },
+            }
+        ],
+        tool_choice="auto",
+    )
+
+
+def _named_request() -> ChatCompletionRequest:
+    return ChatCompletionRequest(
+        model="test-model",
+        messages=[],
+        tools=[
+            {
+                "type": "function",
+                "function": {
+                    "name": "calc",
+                    "parameters": {"type": "object", "properties": {}},
+                },
+            }
+        ],
+        tool_choice={"type": "function", "function": {"name": "calc"}},
+    )
+
+
+def _responses_request(*, tool_choice="auto") -> ResponsesRequest:
+    return ResponsesRequest.model_validate(
+        {
+            "model": "test-model",
+            "input": "Call the calc tool.",
+            "tools": [
+                {
+                    "type": "function",
+                    "name": "calc",
+                    "parameters": {"type": "object", "properties": {}},
+                }
+            ],
+            "tool_choice": tool_choice,
+        }
+    )
+
+
+def _arg(key: str, typ: str, value: str) -> str:
+    return f'{OPEN}argument key="{key}" type="{typ}"{SEP}{value}{CLOSE}argument{SEP}'
+
+
+def _call(tool: str, index: int, *args: str) -> str:
+    body = "".join(args)
+    return f'{OPEN}call tool="{tool}" index="{index}"{SEP}{body}{CLOSE}call{SEP}'
+
+
+def _response(content: str) -> str:
+    return f"{OPEN}response{SEP}{content}{RESPONSE_CLOSE}"
+
+
+def _tools(*calls: str) -> str:
+    return f"{OPEN}tools{SEP}{''.join(calls)}{CLOSE}tools{SEP}"
+
+
+def test_extract_tool_calls_with_response_and_typed_arguments():
+    parser = KimiK3ToolParser(DummyTokenizer())
+
+    output = _response("answer") + _tools(
+        _call(
+            "calc",
+            1,
+            _arg("x", "number", "1"),
+            _arg("flag", "boolean", "true"),
+            _arg("text", "string", "raw"),
+        )
+    )
+    extracted = parser.extract_tool_calls(output, _request())
+
+    assert extracted.tools_called is True
+    assert extracted.content == "answer"
+    assert len(extracted.tool_calls) == 1
+    tool_call = extracted.tool_calls[0]
+    assert tool_call.id == "calc:0"
+    assert tool_call.function.name == "calc"
+    assert json.loads(tool_call.function.arguments) == {
+        "x": 1,
+        "flag": True,
+        "text": "raw",
+    }
+
+
+def test_delegating_parser_preserves_tool_calls_after_reasoning():
+    parser = KimiK3DelegatingParser(DummyTokenizer())
+    output = (
+        f"{THINK_OPEN}step{THINK_CLOSE}"
+        + _response("answer")
+        + _tools(_call("calc", 1, _arg("x", "number", "1")))
+    )
+
+    reasoning, content, tool_calls = parser.parse(
+        output,
+        _request(),
+        enable_auto_tools=True,
+    )
+
+    assert reasoning == "step"
+    assert content == "answer"
+    assert tool_calls is not None
+    assert len(tool_calls) == 1
+    assert tool_calls[0].id == "calc:0"
+    assert tool_calls[0].name == "calc"
+    assert json.loads(tool_calls[0].arguments) == {"x": 1}
+
+
+def test_delegating_parser_required_tool_choice_uses_xtml_parser():
+    parser = KimiK3DelegatingParser(DummyTokenizer())
+    request = _request().model_copy(update={"tool_choice": "required"})
+    output = (
+        f"{THINK_OPEN}step{THINK_CLOSE}"
+        + _response("")
+        + _tools(_call("calc", 1, _arg("x", "number", "1")))
+    )
+
+    reasoning, content, tool_calls = parser.parse(
+        output,
+        request,
+        enable_auto_tools=False,
+    )
+
+    assert reasoning == "step"
+    assert content is None
+    assert tool_calls is not None
+    assert len(tool_calls) == 1
+    assert tool_calls[0].name == "calc"
+    assert json.loads(tool_calls[0].arguments) == {"x": 1}
+
+
+def test_delegating_parser_named_tool_choice_uses_xtml_parser():
+    parser = KimiK3DelegatingParser(DummyTokenizer())
+    output = (
+        f"{THINK_OPEN}step{THINK_CLOSE}"
+        + _response("")
+        + _tools(_call("calc", 1, _arg("x", "number", "1")))
+    )
+
+    reasoning, content, tool_calls = parser.parse(
+        output,
+        _named_request(),
+        enable_auto_tools=False,
+    )
+
+    assert reasoning == "step"
+    assert content is None
+    assert tool_calls is not None
+    assert len(tool_calls) == 1
+    assert tool_calls[0].name == "calc"
+    assert json.loads(tool_calls[0].arguments) == {"x": 1}
+
+
+def test_delegating_parser_auto_no_call_strips_consumed_response_prefix():
+    parser = KimiK3DelegatingParser(
+        DummyTokenizer(), chat_template_kwargs={"thinking": False}
+    )
+    request = _request().model_copy(
+        update={"chat_template_kwargs": {"thinking": False}}
+    )
+
+    reasoning, content, tool_calls = parser.parse(
+        f"answer{RESPONSE_CLOSE}",
+        request,
+        enable_auto_tools=True,
+    )
+
+    assert reasoning is None
+    assert content == "answer"
+    assert tool_calls is None
+
+
+def test_delegating_parser_required_call_strips_consumed_response_prefix():
+    parser = KimiK3DelegatingParser(
+        DummyTokenizer(), chat_template_kwargs={"thinking": False}
+    )
+    request = _request().model_copy(
+        update={
+            "tool_choice": "required",
+            "chat_template_kwargs": {"thinking": False},
+        }
+    )
+    output = RESPONSE_CLOSE + _tools(_call("calc", 1, _arg("x", "number", "1")))
+
+    reasoning, content, tool_calls = parser.parse(
+        output,
+        request,
+        enable_auto_tools=False,
+    )
+
+    assert reasoning is None
+    assert content is None
+    assert tool_calls is not None
+    assert len(tool_calls) == 1
+    assert tool_calls[0].name == "calc"
+    assert json.loads(tool_calls[0].arguments) == {"x": 1}
+
+
+def test_delegating_parser_truncated_tools_do_not_leak_xtml():
+    parser = KimiK3DelegatingParser(
+        DummyTokenizer(), chat_template_kwargs={"thinking": False}
+    )
+    request = _request().model_copy(
+        update={
+            "tool_choice": "required",
+            "chat_template_kwargs": {"thinking": False},
+        }
+    )
+
+    reasoning, content, tool_calls = parser.parse(
+        (f'{RESPONSE_CLOSE}{OPEN}tools{SEP}{OPEN}call tool="calc" index="1"'),
+        request,
+        enable_auto_tools=False,
+    )
+
+    assert reasoning is None
+    assert content is None
+    assert tool_calls is None
+
+
+def test_extract_tool_calls_unescapes_attributes():
+    parser = KimiK3ToolParser(DummyTokenizer())
+
+    output = _tools(_call("a&amp;b&quot;c", 1, _arg("k&amp;q", "string", "v")))
+    extracted = parser.extract_tool_calls(output, _request())
+
+    assert extracted.tools_called is True
+    assert extracted.tool_calls[0].function.name == 'a&b"c'
+    assert json.loads(extracted.tool_calls[0].function.arguments) == {"k&q": "v"}
+
+
+def test_extract_tool_calls_allows_less_than_in_attributes():
+    parser = KimiK3ToolParser(DummyTokenizer())
+
+    output = _tools(_call("calc<beta", 1, _arg("foo<bar", "string", "raw")))
+    extracted = parser.extract_tool_calls(output, _request())
+
+    assert extracted.tools_called is True
+    assert extracted.tool_calls[0].function.name == "calc<beta"
+    assert json.loads(extracted.tool_calls[0].function.arguments) == {"foo<bar": "raw"}
+
+
+def test_extract_content_from_whitespace_degraded_markers():
+    parser = KimiK3ToolParser(DummyTokenizer())
+
+    extracted = parser.extract_tool_calls(
+        f"{OPEN} response {SEP}answer{CLOSE} response {SEP}",
+        _request(),
+    )
+
+    assert extracted.tools_called is False
+    assert extracted.content == "answer"
+
+
+def test_streaming_split_markers_do_not_leak():
+    parser = KimiK3ToolParser(DummyTokenizer())
+    request = _request()
+    previous_text = ""
+    previous_ids: list[int] = []
+    messages: list[DeltaMessage] = []
+    chunks = [
+        OPEN,
+        "response",
+        f"{SEP}Hi",
+        OPEN,
+        "tools",
+        SEP,
+        f'{OPEN}call tool="calc" index="1"{SEP}',
+        _arg("x", "number", "1"),
+        f"{CLOSE}call",
+        SEP,
+    ]
+
+    for i, chunk in enumerate(chunks, start=1):
+        current_text = previous_text + chunk
+        current_ids = previous_ids + [i]
+        delta = parser.extract_tool_calls_streaming(
+            previous_text=previous_text,
+            current_text=current_text,
+            delta_text=chunk,
+            previous_token_ids=previous_ids,
+            current_token_ids=current_ids,
+            delta_token_ids=[i],
+            request=request,
+        )
+        if delta is not None:
+            messages.append(delta)
+        previous_text = current_text
+        previous_ids = current_ids
+
+    content = "".join(message.content or "" for message in messages)
+    tool_deltas = [
+        tool_call for message in messages for tool_call in (message.tool_calls or [])
+    ]
+
+    assert content == "Hi"
+    assert OPEN not in content
+    assert SEP not in content
+    assert len(tool_deltas) == 1
+    assert tool_deltas[0].id == "calc:0"
+    assert tool_deltas[0].function.name == "calc"
+    assert json.loads(tool_deltas[0].function.arguments) == {"x": 1}
+
+
+def test_streaming_consumed_response_prefix_no_call_keeps_content():
+    parser = KimiK3ToolParser(DummyTokenizer())
+    request = _request()
+    previous_text = ""
+    previous_ids: list[int] = []
+    messages: list[DeltaMessage] = []
+    chunks = ["O", "K", CLOSE, f"response{SEP}"]
+
+    for i, chunk in enumerate(chunks, start=1):
+        current_text = previous_text + chunk
+        current_ids = previous_ids + [i]
+        delta = parser.extract_tool_calls_streaming(
+            previous_text=previous_text,
+            current_text=current_text,
+            delta_text=chunk,
+            previous_token_ids=previous_ids,
+            current_token_ids=current_ids,
+            delta_token_ids=[i],
+            request=request,
+        )
+        if delta is not None:
+            messages.append(delta)
+        previous_text = current_text
+        previous_ids = current_ids
+
+    assert "".join(message.content or "" for message in messages) == "OK"
+    assert all(CLOSE not in (message.content or "") for message in messages)
+
+
+def test_adjust_request_keeps_xtml_markers_contiguous():
+    parser = KimiK3ToolParser(DummyTokenizer())
+    request = _request()
+
+    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
+    assert KimiK3ToolParser.supports_required_and_named is False
+
+
+def test_adjust_request_required_uses_xtml_parser_not_json_guidance():
+    parser = KimiK3ToolParser(DummyTokenizer())
+    request = _request().model_copy(update={"tool_choice": "required"})
+
+    adjusted = parser.adjust_request(request)
+
+    assert adjusted.structured_outputs is None
+    assert adjusted.skip_special_tokens is False
+    if hasattr(adjusted, "spaces_between_special_tokens"):
+        assert adjusted.spaces_between_special_tokens is False
+
+
+@pytest.mark.parametrize(
+    "tool_request",
+    [
+        _named_request(),
+        _responses_request(
+            tool_choice={"type": "function", "name": "calc"},
+        ),
+    ],
+)
+def test_adjust_request_rejects_named_tool_choice(tool_request):
+    parser = KimiK3ToolParser(DummyTokenizer())
+
+    with pytest.raises(VLLMValidationError) as exc_info:
+        parser.adjust_request(tool_request)
+
+    assert exc_info.value.parameter == "tool_choice"
+    assert "requires strict tool calling" in str(exc_info.value)
+
+
+def test_responses_chat_params_passes_tool_choice_to_template():
+    request = _responses_request(tool_choice="required")
+
+    chat_params = request.build_chat_params(
+        default_template=None,
+        default_template_content_format="auto",
+    )
+
+    assert chat_params.chat_template_kwargs["tool_choice"] == "required"
+
+
+def test_responses_chat_params_keeps_request_template_tool_choice_when_api_auto():
+    request = _responses_request().model_copy(
+        update={"chat_template_kwargs": {"tool_choice": "required"}}
+    )
+
+    chat_params = request.build_chat_params(
+        default_template=None,
+        default_template_content_format="auto",
+    )
+
+    assert chat_params.chat_template_kwargs["tool_choice"] == "required"
+
+
+def test_responses_required_tool_choice_uses_xtml_parser():
+    parser = KimiK3DelegatingParser(
+        DummyTokenizer(), chat_template_kwargs={"thinking": False}
+    )
+    request = _responses_request(tool_choice="required").model_copy(
+        update={"chat_template_kwargs": {"thinking": False}}
+    )
+    output = RESPONSE_CLOSE + _tools(_call("calc", 1, _arg("x", "number", "1")))
+
+    reasoning, content, tool_calls = parser.parse(
+        output, request, enable_auto_tools=False, model_output_token_ids=[]
+    )
+    response_outputs = build_response_output_items(
+        reasoning=reasoning,
+        content=content,
+        tool_calls=tool_calls,
+        tools=request.tools,
+    )
+
+    assert len(response_outputs) == 1
+    tool_call = response_outputs[0]
+    assert tool_call.type == "function_call"
+    assert tool_call.name == "calc"
+    assert json.loads(tool_call.arguments) == {"x": 1}
+
+
+def test_responses_named_tool_choice_uses_xtml_parser():
+    parser = KimiK3DelegatingParser(
+        DummyTokenizer(), chat_template_kwargs={"thinking": False}
+    )
+    request = _responses_request(
+        tool_choice={"type": "function", "name": "calc"}
+    ).model_copy(update={"chat_template_kwargs": {"thinking": False}})
+    output = RESPONSE_CLOSE + _tools(_call("calc", 1, _arg("x", "number", "1")))
+
+    reasoning, content, tool_calls = parser.parse(
+        output, request, enable_auto_tools=False, model_output_token_ids=[]
+    )
+    response_outputs = build_response_output_items(
+        reasoning=reasoning,
+        content=content,
+        tool_calls=tool_calls,
+        tools=request.tools,
+    )
+
+    assert len(response_outputs) == 1
+    tool_call = response_outputs[0]
+    assert tool_call.type == "function_call"
+    assert tool_call.name == "calc"
+    assert json.loads(tool_call.arguments) == {"x": 1}
+
+
+def test_chat_params_passes_tool_choice_to_template():
+    request = _request().model_copy(update={"tool_choice": "required"})
+
+    chat_params = request.build_chat_params(
+        default_template=None,
+        default_template_content_format="auto",
+    )
+
+    assert chat_params.chat_template_kwargs["tool_choice"] == "required"
+
+
+def test_chat_params_keeps_request_template_tool_choice_when_api_auto():
+    request = _request().model_copy(
+        update={
+            "tool_choice": "auto",
+            "chat_template_kwargs": {"tool_choice": "required"},
+        }
+    )
+
+    chat_params = request.build_chat_params(
+        default_template=None,
+        default_template_content_format="auto",
+    )
+
+    assert chat_params.chat_template_kwargs["tool_choice"] == "required"

tests/transformers_utils/test_dspark_mla_config.py

@@ -0,0 +1,165 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import json
+
+import pytest
+
+from vllm.config import ModelConfig, ParallelConfig, SpeculativeConfig
+from vllm.transformers_utils.config import get_config
+from vllm.transformers_utils.configs.k3_dspark import K3DSparkConfig
+
+
+def _write_dspark_config(path, **overrides):
+    path.mkdir()
+    config = {
+        "architectures": ["K3DSparkModel"],
+        "model_type": "k3_dspark",
+        "hidden_size": 7168,
+        "intermediate_size": 14336,
+        "num_hidden_layers": 5,
+        "num_attention_heads": 64,
+        "num_key_value_heads": 64,
+        "q_lora_rank": 1536,
+        "kv_lora_rank": 512,
+        "qk_nope_head_dim": 128,
+        "qk_rope_head_dim": 64,
+        "v_head_dim": 128,
+        "vocab_size": 163840,
+        "rms_norm_eps": 1e-5,
+        "max_position_embeddings": 32768,
+        "rope_theta": 50000.0,
+        "num_target_layers": 5,
+        "target_hidden_size": 7168,
+        "target_num_hidden_layers": 93,
+        "target_layer_ids": [2, 23, 47, 71, 89],
+        "markov_rank": 256,
+        "draft_vocab_size": 163840,
+        "torch_dtype": "bfloat16",
+    }
+    config.update(overrides)
+    (path / "config.json").write_text(json.dumps(config))
+
+
+def _write_target_config(path):
+    path.mkdir()
+    config = {
+        "architectures": ["LlamaForCausalLM"],
+        "model_type": "llama",
+        "hidden_size": 7168,
+        "intermediate_size": 14336,
+        "num_hidden_layers": 93,
+        "num_attention_heads": 56,
+        "num_key_value_heads": 8,
+        "vocab_size": 163840,
+        "max_position_embeddings": 32768,
+        "torch_dtype": "bfloat16",
+    }
+    (path / "config.json").write_text(json.dumps(config))
+
+
+def test_dspark_mla_config_loads_from_local_json(tmp_path):
+    draft_path = tmp_path / "draft"
+    _write_dspark_config(draft_path)
+
+    config = get_config(draft_path, trust_remote_code=False)
+
+    assert isinstance(config, K3DSparkConfig)
+    assert config.model_type == "k3_dspark"
+    assert config.architectures == ["K3DSparkModel"]
+    assert config.hidden_act == "silu"
+    assert config.rope_parameters == {
+        "rope_type": "default",
+        "rope_theta": 50000.0,
+    }
+    assert config.n_routed_experts == 0
+    assert config.draft_vocab_size == config.vocab_size
+
+
+@pytest.mark.parametrize(
+    "overrides",
+    [
+        {"mla_use_nope": True},
+        {"mla_use_output_gate": True},
+        {"mla_use_qk_norm": True},
+        {"dspark_bonus_anchor": True},
+        {"q_lora_rank": None},
+        {"draft_vocab_size": 8192},
+        {"target_layer_ids": []},
+        {"num_target_layers": 4},
+    ],
+)
+def test_dspark_mla_rejects_unsupported_checkpoint_options(tmp_path, overrides):
+    draft_path = tmp_path / "draft"
+    _write_dspark_config(draft_path, **overrides)
+
+    with pytest.raises(ValueError, match="MLA DSpark"):
+        get_config(draft_path, trust_remote_code=False)
+
+
+def test_dspark_mla_uses_latent_kv_geometry(tmp_path):
+    draft_path = tmp_path / "draft"
+    _write_dspark_config(draft_path)
+    model_config = ModelConfig(
+        model=str(draft_path),
+        tokenizer_mode="skip",
+        runner="draft",
+        max_model_len=32768,
+    )
+
+    assert model_config.is_deepseek_mla
+    assert model_config.use_mla
+    assert model_config.get_head_size() == 576
+    # external_launcher skips ParallelConfig's local-GPU-count check so the
+    # config logic can be exercised at TP8 on a single-GPU test node.
+    parallel_config = ParallelConfig(
+        tensor_parallel_size=8, distributed_executor_backend="external_launcher"
+    )
+    assert model_config.get_num_kv_heads(parallel_config) == 1
+    assert model_config.get_num_attention_heads(parallel_config) == 8
+    assert model_config.get_num_experts() == 0
+
+
+def test_dspark_mla_speculative_config_preserves_architecture(tmp_path):
+    target_path = tmp_path / "target"
+    draft_path = tmp_path / "draft"
+    _write_target_config(target_path)
+    _write_dspark_config(draft_path)
+    target_config = ModelConfig(
+        model=str(target_path), tokenizer_mode="skip", max_model_len=32768
+    )
+    speculative_config = SpeculativeConfig(
+        model=str(draft_path),
+        method="dspark",
+        num_speculative_tokens=8,
+        target_model_config=target_config,
+        target_parallel_config=ParallelConfig(),
+    )
+
+    assert speculative_config.parallel_drafting
+    assert speculative_config.draft_model_config.architectures == ["K3DSparkModel"]
+    assert speculative_config.draft_model_config.hf_config.model_type == "k3_dspark"
+    assert speculative_config.draft_model_config.use_mla
+
+
+def test_dspark_mla_rejects_decode_context_parallelism(tmp_path):
+    target_path = tmp_path / "target"
+    draft_path = tmp_path / "draft"
+    _write_target_config(target_path)
+    _write_dspark_config(draft_path)
+    target_config = ModelConfig(
+        model=str(target_path), tokenizer_mode="skip", max_model_len=32768
+    )
+
+    with pytest.raises(ValueError, match="does not currently support decode context"):
+        SpeculativeConfig(
+            model=str(draft_path),
+            method="dspark",
+            num_speculative_tokens=8,
+            target_model_config=target_config,
+            target_parallel_config=ParallelConfig(
+                tensor_parallel_size=2,
+                decode_context_parallel_size=2,
+                distributed_executor_backend="external_launcher",
+            ),
+        )

tests/v1/attention/test_mla_noncausal.py

@@ -0,0 +1,124 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from vllm.model_executor.layers.attention.mla_attention import (
+    MLACommonMetadata,
+    MLACommonMetadataBuilder,
+    QueryLenSupport,
+)
+from vllm.v1.attention.backend import CommonAttentionMetadata
+from vllm.v1.kv_cache_interface import MLAAttentionSpec
+
+
+class _NonCausalMLAMetadataBuilder(MLACommonMetadataBuilder[MLACommonMetadata]):
+    supports_non_causal_multi_token_decode = True
+
+
+def _metadata(
+    query_start_loc: list[int], num_tokens: int | None = None
+) -> CommonAttentionMetadata:
+    num_reqs = len(query_start_loc) - 1
+    num_tokens = query_start_loc[-1] if num_tokens is None else num_tokens
+    return CommonAttentionMetadata(
+        query_start_loc=torch.tensor(query_start_loc, dtype=torch.int32),
+        query_start_loc_cpu=torch.tensor(query_start_loc, dtype=torch.int32),
+        seq_lens=torch.arange(1, num_reqs + 1, dtype=torch.int32) * 100 + 8,
+        num_reqs=num_reqs,
+        num_actual_tokens=num_tokens,
+        max_query_len=max(
+            end - start for start, end in zip(query_start_loc, query_start_loc[1:])
+        ),
+        max_seq_len=num_reqs * 100 + 8,
+        block_table_tensor=torch.arange(num_reqs * 3, dtype=torch.int32).view(
+            num_reqs, 3
+        ),
+        slot_mapping=torch.arange(num_tokens),
+        causal=False,
+        seq_lens_cpu_upper_bound=None,
+    )
+
+
+def _builder(marked: bool = True) -> _NonCausalMLAMetadataBuilder:
+    builder = object.__new__(_NonCausalMLAMetadataBuilder)
+    builder.device = torch.device("cpu")
+    builder.reorder_batch_threshold = 1
+    builder.query_len_support = QueryLenSupport.SINGLE_ONLY
+    builder.non_causal_multi_token_decode = marked
+    builder.dcp_world_size = 1
+    builder.metadata_cls = MLACommonMetadata
+    builder.model_config = SimpleNamespace(
+        dtype=torch.bfloat16, get_head_size=lambda: 576
+    )
+    return builder
+
+
+def test_noncausal_block_uses_decode_without_cpu_lengths():
+    common_metadata = _metadata([0, 8, 16])
+    metadata = _builder().build(0, common_metadata)
+
+    assert metadata.num_decodes == 2
+    assert metadata.num_decode_tokens == 16
+    assert metadata.num_prefills == 0
+    assert metadata.prefill is None
+    assert metadata.decode is not None
+    assert metadata.decode.block_table.shape == (2, 3)
+    assert metadata.decode.seq_lens.shape == (2,)
+    assert torch.equal(metadata.decode.block_table, common_metadata.block_table_tensor)
+    assert torch.equal(metadata.decode.seq_lens, common_metadata.seq_lens)
+    assert not metadata.causal
+
+
+def test_noncausal_support_is_explicit_and_uniform():
+    with pytest.raises(ValueError, match="explicitly supported"):
+        _builder(marked=False).build(0, _metadata([0, 8, 16]))
+    with pytest.raises(ValueError, match="uniform query block"):
+        _builder().build(0, _metadata([0, 3, 8]))
+
+
+def test_noncausal_block_allows_trailing_cudagraph_padding():
+    common_metadata = _metadata([0, 8, 16, 16], num_tokens=24)
+    common_metadata.seq_lens[-1] = 0
+
+    metadata = _builder().build(0, common_metadata)
+
+    assert metadata.num_decodes == 3
+    assert metadata.num_decode_tokens == 24
+    assert metadata.decode is not None
+    assert metadata.decode.seq_lens.tolist() == [108, 208, 0]
+
+
+def test_noncausal_block_rejects_non_trailing_padding():
+    with pytest.raises(ValueError, match="uniform query block"):
+        _builder().build(0, _metadata([0, 8, 8, 16], num_tokens=24))
+
+
+def test_noncausal_decode_metadata_keeps_live_request_buffers():
+    common_metadata = _metadata([0, 8, 16])
+    metadata = _builder().build(0, common_metadata)
+
+    assert metadata.decode is not None
+    assert metadata.decode.seq_lens.data_ptr() == common_metadata.seq_lens.data_ptr()
+    assert (
+        metadata.decode.block_table.data_ptr()
+        == common_metadata.block_table_tensor.data_ptr()
+    )
+
+
+def test_mla_cache_marker_is_preserved_and_cannot_be_mixed():
+    kwargs = {
+        "block_size": 64,
+        "num_kv_heads": 1,
+        "head_size": 576,
+        "dtype": torch.bfloat16,
+    }
+    marked = MLAAttentionSpec(**kwargs, non_causal_multi_token_decode=True)
+    unmarked = MLAAttentionSpec(**kwargs)
+
+    assert MLAAttentionSpec.merge([marked, marked]).non_causal_multi_token_decode
+    with pytest.raises(AssertionError, match="non-causal decode mode"):
+        MLAAttentionSpec.merge([marked, unmarked])

tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py

@@ -40,6 +40,8 @@ def test_mamba_align_split_partial_tail_schedule():
     hash_block_size = 32
     mock = SimpleNamespace(
         cache_config=SimpleNamespace(block_size=block_size),
+        max_num_scheduled_tokens=8192,
+        scheduler_config=SimpleNamespace(long_prefill_token_threshold=0),
         use_eagle=False,
         hash_block_size=hash_block_size,
         mamba_partial_cache_hit=True,
@@ -75,6 +77,78 @@ def test_mamba_align_split_partial_tail_schedule():
     assert split(self=mock, request=req2, num_new_tokens=1000) == 512
 
 
+def test_mamba_align_split_when_block_exceeds_scheduling_budget():
+    """Sub-block chunks make progress only when no step can fit a full block."""
+    block_size = 11392
+    token_budget = 8192
+    prompt_length = 30000
+    mock = SimpleNamespace(
+        cache_config=SimpleNamespace(block_size=block_size),
+        max_num_scheduled_tokens=token_budget,
+        scheduler_config=SimpleNamespace(long_prefill_token_threshold=0),
+        use_eagle=False,
+        hash_block_size=32,
+        mamba_partial_cache_hit=False,
+    )
+    req = make_request("0", [0] * prompt_length, 32, sha256)
+    split = Scheduler._mamba_block_aligned_split
+
+    mock.max_num_scheduled_tokens = block_size
+    assert split(self=mock, request=req, num_new_tokens=token_budget) == 0
+    mock.max_num_scheduled_tokens = token_budget
+
+    scheduled_chunks = []
+    while req.num_computed_tokens < prompt_length:
+        num_new_tokens = min(token_budget, prompt_length - req.num_computed_tokens)
+        num_scheduled_tokens = split(
+            self=mock,
+            request=req,
+            num_new_tokens=num_new_tokens,
+        )
+        assert 0 < num_scheduled_tokens <= token_budget
+        scheduled_chunks.append(num_scheduled_tokens)
+        req.num_computed_tokens += num_scheduled_tokens
+
+    assert scheduled_chunks == [8192, 3200, 8192, 3200, 7216]
+
+
+def test_mamba_align_split_when_block_exceeds_long_prefill_threshold():
+    """A long-prefill cap below the block size permits sub-block progress."""
+    block_size = 512
+    token_budget = 8192
+    long_prefill_threshold = 384
+    prompt_length = 1300
+    mock = SimpleNamespace(
+        cache_config=SimpleNamespace(block_size=block_size),
+        max_num_scheduled_tokens=token_budget,
+        scheduler_config=SimpleNamespace(
+            long_prefill_token_threshold=long_prefill_threshold
+        ),
+        use_eagle=False,
+        hash_block_size=32,
+        mamba_partial_cache_hit=False,
+    )
+    req = make_request("0", [0] * prompt_length, 32, sha256)
+    split = Scheduler._mamba_block_aligned_split
+
+    scheduled_chunks = []
+    while req.num_computed_tokens < prompt_length:
+        num_new_tokens = min(
+            long_prefill_threshold,
+            prompt_length - req.num_computed_tokens,
+        )
+        num_scheduled_tokens = split(
+            self=mock,
+            request=req,
+            num_new_tokens=num_new_tokens,
+        )
+        assert 0 < num_scheduled_tokens <= long_prefill_threshold
+        scheduled_chunks.append(num_scheduled_tokens)
+        req.num_computed_tokens += num_scheduled_tokens
+
+    assert scheduled_chunks == [384, 128, 384, 128, 276]
+
+
 def test_hybrid_mamba_align_partial_hash_hit():
     hash_block_size = 2
     mamba_block_size = 2 * hash_block_size

tests/v1/core/test_prefix_caching.py

@@ -1078,6 +1078,8 @@ def test_hybrid_cache_mamba_align_shared_prefix_detection():
     # Create minimal mock with just the needed attributes
     mock = SimpleNamespace(
         cache_config=SimpleNamespace(block_size=block_size),
+        max_num_scheduled_tokens=3 * block_size,
+        scheduler_config=SimpleNamespace(long_prefill_token_threshold=0),
         use_eagle=False,
         hash_block_size=block_size,
         mamba_partial_cache_hit=False,

tests/v1/cudagraph/test_breakable_cudagraph.py

@@ -333,6 +333,54 @@ def test_decorator_breaks_when_invoked_inside_capture(cuda_capture_stream):
     assert torch.equal(x, torch.full((4,), 15.0, device="cuda"))
 
 
+def test_eager_attention_inside_multistream_overlap(cuda_capture_stream):
+    """Handle an eager attention break inside a multi-stream overlap region."""
+    from vllm.compilation.breakable_cudagraph import (
+        BreakableCUDAGraphCapture,
+        eager_break_during_capture,
+    )
+    from vllm.utils.multi_stream_utils import maybe_execute_in_parallel
+
+    x = torch.zeros(1024, device="cuda")
+    output = torch.empty_like(x)
+    aux_stream = torch.cuda.Stream()
+    main_event = torch.cuda.Event()
+    aux_event = torch.cuda.Event()
+
+    @eager_break_during_capture
+    def attention(query: torch.Tensor, out: torch.Tensor) -> None:
+        torch.add(query, 3.0, out=out)
+
+    def attention_frontend() -> torch.Tensor:
+        query = x * 2.0
+        attention_output = torch.empty_like(x)
+        attention(query, attention_output)
+        return attention_output
+
+    cap = BreakableCUDAGraphCapture()
+    with cap:
+        attention_output, gate = maybe_execute_in_parallel(
+            attention_frontend,
+            lambda: x * 5.0,
+            main_event,
+            aux_event,
+            aux_stream,
+        )
+        torch.add(attention_output, gate, out=output)
+
+    assert cap.num_graphs == 2
+    assert cap.num_eager_breaks == 1
+
+    for value in (1.0, 7.0, 19.0):
+        x.fill_(value)
+        cap.replay()
+        cuda_capture_stream.synchronize()
+        torch.testing.assert_close(
+            output,
+            torch.full_like(output, value * 7.0 + 3.0),
+        )
+
+
 # ---------------------------------------------------------------------------
 # Replay ordering
 # ---------------------------------------------------------------------------

tests/v1/kv_connector/unit/test_nixl_connector.py

@@ -479,8 +479,9 @@ class FakeNixlConnectorWorker(NixlConnectorWorker):
         super().__init__(*args, kv_cache_config=kv_cache_config, **kwargs)
         self._hand_shake_latency = hand_shake_latency
         self.kv_cache_layout = kv_cache_layout
-        # Mock register_kv_caches attribute needed for tests that do not call it.
+        # Mock register_kv_caches attributes needed for tests that do not call it.
         self.src_xfer_handles_by_block_size = {self.block_size: 1}
+        self.src_blocks_data = np.empty((0, 3), dtype=np.uint64)
         test_shape = self.attn_backends[0].get_kv_cache_shape(
             num_blocks=1, block_size=16, num_kv_heads=1, head_size=1
         )
@@ -765,8 +766,9 @@ class TestNixlHandshake:
             assert remote_info.remote_tp_size == remote_tp_size
             assert -tp_ratio == worker.transfer_topo.tp_ratio(remote_tp_size)
             # ensure src_xfer_handles_by_tp_ratio is populated with tpratio chunks
-            assert -tp_ratio in worker.src_xfer_handles_by_tp_ratio
-            assert len(worker.src_xfer_handles_by_tp_ratio[-tp_ratio]) == tp_ratio
+            split_key = (-tp_ratio, worker.block_size)
+            assert split_key in worker.src_xfer_handles_by_tp_ratio
+            assert len(worker.src_xfer_handles_by_tp_ratio[split_key]) == tp_ratio
             assert remote_engine_id in worker.dst_xfer_side_handles
             assert set(worker.dst_xfer_side_handles[remote_engine_id].keys()) == set(
                 range(tp_ratio)
@@ -1408,6 +1410,53 @@ def test_kv_connector_stats(default_vllm_config, dist_init):
     assert stats_after_reset is None
 
 
+def test_reqs_to_send_deadline_rebased_to_worker_clock(default_vllm_config, dist_init):
+    """reqs_to_send deadlines are stamped with the scheduler process's
+    perf_counter, whose epoch differs across processes and (by boot-time
+    deltas) across nodes. Without rebasing, a P worker on a node whose
+    monotonic clock is ahead of the scheduler's by more than the TTL
+    expires the lease on arrival and reports done_sending before D has
+    read the blocks — the freed blocks can then be reallocated and the
+    remote read pulls another request's data (silent accuracy corruption).
+    The worker must anchor the remaining TTL to its own clock.
+    """
+    vllm_config = create_vllm_config()
+    connector = NixlConnector(
+        vllm_config, KVConnectorRole.WORKER, make_kv_cache_config(block_size=16)
+    )
+    connector.connector_worker = FakeNixlConnectorWorker(
+        vllm_config, connector.engine_id, hand_shake_latency=0
+    )
+    worker = connector.connector_worker
+
+    req_id = "req-lease-clock"
+    ttl = 480.0
+    # Simulate a scheduler whose monotonic clock is 10,000 s behind this
+    # worker's (e.g. its node booted much later): the raw deadline is
+    # then already far in the past in this worker's clock domain.
+    scheduler_clock = time.perf_counter() - 10_000.0
+
+    metadata = NixlConnectorMetadata()
+    metadata.reqs_in_batch = {req_id}
+    metadata.reqs_to_send = {req_id: scheduler_clock + ttl}
+    metadata.scheduler_clock = scheduler_clock
+    connector.bind_connector_metadata(metadata)
+    dummy_ctx = ForwardContext(
+        no_compile_layers={},
+        attn_metadata={},
+        slot_mapping={},
+    )
+    connector.start_load_kv(dummy_ctx)
+
+    remaining = worker._reqs_to_send[req_id] - time.perf_counter()
+    assert ttl - 5.0 < remaining <= ttl + 5.0
+
+    # The expiry sweep must not release the request.
+    done_sending, _ = worker.get_finished()
+    assert req_id not in done_sending
+    assert req_id in worker._reqs_to_process
+
+
 def test_kv_connector_stats_aggregation():
     """
     Test KV transfer stats aggregation across TP ranks using
@@ -2103,7 +2152,7 @@ def test_shutdown_cleans_up_resources(default_vllm_config, dist_init):
         # Mock register_kv_cache which registers local handle
         worker.src_xfer_handles_by_block_size = {worker.block_size: 455}
         # P TP = 2 * D TP case, we should register 2 local handles
-        worker.src_xfer_handles_by_tp_ratio = {-2: [456, 457]}
+        worker.src_xfer_handles_by_tp_ratio = {(-2, 16): [456, 457]}
         worker.dst_xfer_side_handles = {"engine1": {0: 789}}
         worker._remote_agents = {"engine1": {(0, 0): "agent1"}}
         # _cleanup_remote_engine (called by shutdown) also clears these:

tests/v1/kv_connector/unit/test_nixl_connector_hma.py

@@ -708,6 +708,120 @@ def test_get_block_descs_ids_kernel_block_mismatch():
     assert list(result) == expected, f"Expected {expected}, got {list(result)}"
 
 
+@pytest.mark.cpu_test
+def test_get_block_descs_ids_hetero_block_size_hybrid():
+    """With a block-size ratio, FA desc ids are ratio-expanded while SSM
+    desc ids keep the unexpanded logical stride (state blocks are never
+    sub-split)."""
+    from vllm.v1.kv_cache_interface import FullAttentionSpec, MambaSpec
+
+    worker = _make_mock_worker_for_desc_ids(
+        num_regions=2,
+        has_mamba=True,
+        group_spec_types=(FullAttentionSpec, MambaSpec),
+        block_len_per_layer=[100],
+    )
+
+    ratio = 4
+    # FA ids are already remote-granularity (expanded) sub-block ids.
+    fa_sub_blocks = [3, 5]
+    ssm_blocks = [1]
+    result = worker._compute_desc_ids(
+        block_ids=(fa_sub_blocks, ssm_blocks),
+        dst_num_blocks=100,
+        block_size_ratio=ratio,
+        physical_blocks_per_logical=1,
+    )
+
+    # FA regions have 100*4 entries each; SSM regions (4 per layer) start at
+    # 2*400 and stride by the unexpanded 100 logical blocks.
+    expected = [3, 5, 403, 405, 801, 901, 1001, 1101]
+    assert list(result) == expected, f"Expected {expected}, got {list(result)}"
+
+
+def _bind_worker_method(worker, name):
+    from vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker import (
+        NixlConnectorWorker,
+    )
+
+    method = getattr(NixlConnectorWorker, name)
+    setattr(worker, name, method.__get__(worker, NixlConnectorWorker))
+
+
+@pytest.mark.cpu_test
+def test_map_block_ids_for_block_size_ratio_hybrid():
+    """Attention groups expand to remote granularity and clip to the remote
+    coverage; mamba state blocks pass through 1:1."""
+    from unittest.mock import MagicMock
+
+    from vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker import (
+        NixlConnectorWorker,
+    )
+    from vllm.v1.kv_cache_interface import FullAttentionSpec, MambaSpec
+
+    worker = MagicMock(spec=NixlConnectorWorker)
+    worker._group_spec_types = (FullAttentionSpec, MambaSpec)
+    _bind_worker_method(worker, "get_mapped_blocks")
+    _bind_worker_method(worker, "_map_block_ids_for_block_size_ratio")
+
+    local, remote = worker._map_block_ids_for_block_size_ratio(
+        [[1, 2, 3], [7]],
+        [list(range(30, 40)), [42]],
+        4,
+    )
+    # [1, 2, 3] expand to sub-blocks [4..15], clipped to the 10 remote blocks.
+    assert local == [list(range(4, 14)), [7]]
+    assert remote == [list(range(30, 40)), [42]]
+
+    # Attention-only full prefix hit: empty local list is preserved.
+    worker._group_spec_types = (FullAttentionSpec,)
+    local, remote = worker._map_block_ids_for_block_size_ratio([[]], [[30, 31]], 4)
+    assert local == []
+
+
+@pytest.mark.cpu_test
+def test_post_process_zeroes_untransferred_tail():
+    """The untransferred sub-blocks of the last local block are zeroed on
+    receive; mamba state caches are untouched by the attention permute."""
+    from unittest.mock import MagicMock
+
+    from vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker import (
+        NixlConnectorWorker,
+    )
+    from vllm.v1.kv_cache_interface import FullAttentionSpec, MambaSpec
+
+    ratio = 4
+    block_tokens = 8  # 2 tokens per remote sub-block
+
+    worker = MagicMock(spec=NixlConnectorWorker)
+    worker._group_spec_types = (FullAttentionSpec, MambaSpec)
+    worker.transfer_topo = MagicMock()
+    worker.device_type = "cpu"
+    worker.enable_permute_local_kv = False
+    attn_cache = torch.ones(6, block_tokens, 2, 4)
+    mamba_cache = torch.ones(6, 16)
+    worker.device_kv_caches = {"attn.0": attn_cache, "mamba.0": mamba_cache}
+    fa_group = MagicMock(layer_names=["attn.0"])
+    ssm_group = MagicMock(layer_names=["mamba.0"])
+    worker.kv_cache_config = MagicMock(kv_cache_groups=[fa_group, ssm_group])
+    # The cached property filters mamba layers out of the permuted caches.
+    attn_caches = NixlConnectorWorker._attention_kv_caches.func(worker)
+    assert len(attn_caches) == 1 and attn_caches[0] is attn_cache
+    worker._attention_kv_caches = attn_caches
+    _bind_worker_method(worker, "post_process_device_kv_on_receive")
+
+    # Request occupies blocks [2, 3]; only 6 of 8 sub-blocks were received.
+    worker.post_process_device_kv_on_receive(ratio, [([2, 3], 6)])
+
+    # Block 2 fully covered; block 3 covered for 2 sub-blocks (4 tokens).
+    assert torch.all(attn_cache[2] == 1)
+    assert torch.all(attn_cache[3, :4] == 1)
+    assert torch.all(attn_cache[3, 4:] == 0)
+    # Untouched blocks and the mamba cache keep their content.
+    assert torch.all(attn_cache[4] == 1)
+    assert torch.all(mamba_cache == 1)
+
+
 @pytest.mark.cpu_test
 def test_nixl_metadata_hybrid_ssm_block_ids():
     """Test NixlConnectorMetadata correctly stores block IDs for FA + SSM
@@ -1266,3 +1380,194 @@ def test_logical_to_kernel_block_ids_with_remote_ratio(
     assert list(result) == expected_kernel_block_ids, (
         f"Expected {expected_kernel_block_ids}, got {result}"
     )
+
+
+# ── Hybrid MLA+SSM (KimiLinear-shaped KDA+MLA) tests ─────────────────────
+
+
+def _make_hybrid_mla_kv_cache_config(num_blocks: int = 4):
+    """KimiLinear-shaped config: one MLA group and two KDA (GDN-typed
+    MambaSpec) groups whose layers share the same HMA tensors, with a
+    mamba-aligned unified page and an MLA kernel block smaller than the
+    logical block."""
+    from vllm.v1.attention.backends.registry import MambaAttentionBackendEnum
+    from vllm.v1.kv_cache_interface import (
+        KVCacheConfig,
+        KVCacheGroupSpec,
+        KVCacheTensor,
+        MambaSpec,
+        MLAAttentionSpec,
+    )
+
+    # 12-token logical blocks over a 4-token MLA kernel block.
+    mla_spec = MLAAttentionSpec(
+        block_size=12, num_kv_heads=1, head_size=6, dtype=torch.float16
+    )
+    unified_page = mla_spec.page_size_bytes
+    kda_spec = MambaSpec(
+        block_size=12,
+        # GDN-decomposable conv (Q|K|V = 2|2|4 cols x 3 rows) + fp32 temporal.
+        shapes=((8, 3), (1, 4, 4)),
+        dtypes=(torch.float16, torch.float32),
+        page_size_padded=unified_page,
+        mamba_type=MambaAttentionBackendEnum.GDN_ATTN,
+    )
+    assert kda_spec.page_size_bytes == unified_page
+    return KVCacheConfig(
+        num_blocks=num_blocks,
+        kv_cache_tensors=[
+            KVCacheTensor(
+                size=num_blocks * unified_page,
+                shared_by=[f"mla.{i}", f"kda_a.{i}", f"kda_b.{i}"],
+            )
+            for i in range(2)
+        ],
+        kv_cache_groups=[
+            KVCacheGroupSpec(["mla.0", "mla.1"], mla_spec),
+            KVCacheGroupSpec(["kda_a.0", "kda_a.1"], kda_spec),
+            KVCacheGroupSpec(["kda_b.0", "kda_b.1"], kda_spec),
+        ],
+    )
+
+
+@pytest.mark.cpu_test
+def test_register_kv_caches_hybrid_mla_dual_purpose_regions():
+    """Hybrid MLA+KDA registration: HMA tensors shared by both layer types
+    must be flagged as MLA regions even when a KDA layer registers them
+    first, expose TP-independent kernel-granularity block lens, and build
+    FA + mamba descriptors for every region."""
+    from unittest.mock import MagicMock
+
+    from vllm.config import set_current_vllm_config
+    from vllm.distributed.kv_transfer.kv_connector.v1.nixl import base_worker as bw
+    from vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker import (
+        NixlConnectorWorker,
+    )
+
+    kv_cache_config = _make_hybrid_mla_kv_cache_config()
+    unified_page = kv_cache_config.kv_cache_groups[0].kv_cache_spec.page_size_bytes
+    vllm_config = create_vllm_config(block_size=12)
+
+    fake_backend = MagicMock()
+    fake_backend.get_supported_kernel_block_sizes.return_value = [4]
+    fake_backend.get_name.return_value = "FLASHMLA"
+    fake_backend.full_cls_name.return_value = "fake.FLASHMLA"
+    fake_platform = MagicMock()
+    fake_platform.device_type = "cuda"
+    fake_platform.get_nixl_memory_type.return_value = "VRAM"
+
+    with (
+        patch.object(bw, "NixlWrapper"),
+        patch.object(bw, "get_tensor_model_parallel_rank", return_value=0),
+        patch.object(bw, "get_tensor_model_parallel_world_size", return_value=1),
+        patch.object(bw, "get_current_attn_backends", return_value=[fake_backend]),
+        patch.object(bw, "current_platform", fake_platform),
+        patch(
+            "vllm.model_executor.layers.mamba.mamba_utils.get_conv_state_layout",
+            return_value="DS",
+        ),
+        set_current_vllm_config(vllm_config),
+    ):
+        worker = NixlConnectorWorker(vllm_config, "test-engine", kv_cache_config)
+        worker.use_mla = True  # opt-125m test config is not MLA; force the flag
+        worker.nixl_wrapper.get_agent_metadata.return_value = b"fake-agent-metadata"
+
+        tensors = [torch.zeros(4 * unified_page, dtype=torch.uint8) for _ in range(2)]
+        # KDA layer first per tensor: exercises the dual-purpose flag merge.
+        worker.register_kv_caches(
+            {
+                "kda_a.0": tensors[0],
+                "mla.0": tensors[0],
+                "kda_b.0": tensors[0],
+                "kda_a.1": tensors[1],
+                "mla.1": tensors[1],
+                "kda_b.1": tensors[1],
+            }
+        )
+
+    # 12-token logical blocks over the 4-token MLA kernel block.
+    assert worker._physical_blocks_per_logical_kv_block == 3
+    assert worker.block_size == 4 and worker.num_blocks == 12
+    # Both shared tensors are dual-purpose: their FA view is MLA even though
+    # a KDA layer registered them first.
+    assert worker._region_is_mla == [True, True]
+    assert worker.num_regions == 2 and worker.num_descs == 24
+    # Kernel-granularity block lens; TP-independent for MLA hybrids.
+    assert worker.block_len_per_layer == [unified_page // 3] * 2
+    # Split handles must replicate every FA descriptor (MLA isn't head-sharded).
+    assert worker._fa_desc_replicated(worker.num_descs) == [True] * 24
+    # FA descs: 2 regions x 12 kernel blocks, page stride = kernel page.
+    # Mamba descs: 2 regions x (3 conv sub-projections + 1 ssm) x 4 blocks.
+    assert worker.src_blocks_data.shape == (24 + 32, 3)
+    fa_descs = worker.src_blocks_data[:24]
+    assert fa_descs[1][0] - fa_descs[0][0] == unified_page // 3
+    assert all(size == unified_page // 3 for size in fa_descs[:, 1])
+
+
+@pytest.mark.cpu_test
+def test_push_write_hybrid_mla_replicates_attention():
+    """Hybrid MLA+SSM push with P_TP < D_TP: attention blocks must be
+    written to every covered D rank (replicated MLA latent) while SSM state
+    is written per-rank through the split handles."""
+    import threading
+    from collections import defaultdict
+    from unittest.mock import MagicMock
+
+    from vllm.distributed.kv_transfer.kv_connector.v1.nixl.push_worker import (
+        NixlPushConnectorWorker,
+    )
+    from vllm.distributed.kv_transfer.kv_connector.v1.nixl.tp_mapping import (
+        TPMapping,
+    )
+    from vllm.v1.kv_cache_interface import MambaSpec, MLAAttentionSpec
+
+    worker = object.__new__(NixlPushConnectorWorker)
+    worker.shutdown = lambda: None  # skeleton worker: silence __del__
+    worker.use_mla = True
+    worker._has_mamba = True
+    worker._group_spec_types = (MLAAttentionSpec, MambaSpec)
+    worker.transfer_topo = MagicMock()
+    worker.transfer_topo.tp_ratio.return_value = -2
+    remote_info = MagicMock()
+    remote_info.remote_physical_blocks_per_logical = 1
+    remote_info.remote_block_size = 4
+    worker.transfer_topo.get_engine_info.return_value = remote_info
+
+    engine_id = "remote-engine"
+    # Read-oriented mapping collapses the replicated attention group to one
+    # source rank; the SSM state is sharded across both covered D ranks.
+    worker.tp_mappings = {
+        engine_id: TPMapping(
+            source_ranks_per_group=((0,), (0, 1)),
+            all_source_ranks=(0, 1),
+            rank_to_attention_slot={0: 0, 1: 0},
+            rank_offset_factor=0,
+        )
+    }
+    worker.dst_xfer_side_handles = {engine_id: {0: 100, 1: 101}}
+    worker.src_xfer_handles_by_tp_ratio = {(-2, 4): [200, 201]}
+    worker.src_xfer_handles_by_block_size = {4: 300}
+    worker._sending_transfers = defaultdict(list)
+    worker._sending_transfers_lock = threading.Lock()
+    worker.kv_cache_config = _make_hybrid_mla_kv_cache_config()
+    worker._xfer_blocks = MagicMock(return_value=1)
+
+    meta = MagicMock()
+    meta.remote.engine_id = engine_id
+    meta.remote.block_ids = [[7, 8], [3]]
+    meta.local_physical_block_ids = [[1, 2], [5]]
+
+    worker._xfer_blocks_for_req("req-1", meta)
+
+    calls = worker._xfer_blocks.call_args_list
+    assert len(calls) == 2
+    for call, rank, local_handle, remote_handle in zip(
+        calls, (0, 1), (200, 201), (100, 101)
+    ):
+        spec = call.kwargs["read_spec"]
+        assert spec.remote_rank == rank
+        # Attention group replicated to every rank, SSM by membership.
+        assert spec.local_block_ids == [[1, 2], [5]]
+        assert spec.remote_block_ids == [[7, 8], [3]]
+        assert call.kwargs["local_xfer_side_handle"] == local_handle
+        assert call.kwargs["remote_xfer_side_handle"] == remote_handle

tests/v1/kv_connector/unit/test_nixl_desc_geometry.py

@@ -0,0 +1,573 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""End-to-end NIXL descriptor geometry invariants for hybrid MLA+SSM models
+under heterogeneous P/D block geometry (TP-sharded KDA-style state, so
+the mamba-aligned logical block size differs between P and D while the
+kernel-granularity pages stay equal).
+
+The invariant under test: every LOCAL byte range a request's READ transfers
+into must lie within that request's own blocks. A violation means an
+incoming transfer can overwrite a co-resident request's KV or mamba state
+mid-decode (silent corruption of an unrelated request).
+"""
+
+from unittest.mock import patch
+
+import numpy as np
+import pytest
+import torch
+
+from .utils import create_vllm_config
+
+
+class _RecordingNixl:
+    """Minimal NIXL wrapper stand-in that records descriptor lists and
+    prepared transfers so tests can resolve desc ids to byte ranges."""
+
+    def __init__(self, *args, **kwargs):
+        self.dlists: dict[int, np.ndarray] = {}
+        self.xfers: list[tuple] = []
+        self._next_handle = 1
+
+    def get_reg_descs(self, caches_data, mem_type):
+        return caches_data
+
+    def register_memory(self, descs, backends=None):
+        pass
+
+    def deregister_memory(self, descs):
+        pass
+
+    def get_agent_metadata(self):
+        return b"agent-meta"
+
+    def get_xfer_descs(self, blocks_data, mem_type):
+        return blocks_data
+
+    def prep_xfer_dlist(self, agent, descs):
+        handle = self._next_handle
+        self._next_handle += 1
+        self.dlists[handle] = np.asarray(descs, dtype=np.uint64).reshape(-1, 3)
+        return handle
+
+    def add_remote_agent(self, metadata):
+        return "remote-agent"
+
+    def make_prepped_xfer(
+        self, op, local_handle, local_ids, remote_handle, remote_ids, notif_msg=None
+    ):
+        handle = self._next_handle
+        self._next_handle += 1
+        self.xfers.append(
+            (
+                op,
+                local_handle,
+                np.asarray(local_ids),
+                remote_handle,
+                np.asarray(remote_ids),
+            )
+        )
+        return handle
+
+    def transfer(self, handle):
+        pass
+
+    def check_xfer_state(self, handle):
+        return "DONE"
+
+    def get_xfer_telemetry(self, handle):
+        from types import SimpleNamespace
+
+        return SimpleNamespace(
+            xferDuration=1.0, postDuration=1.0, totalBytes=1, descCount=1
+        )
+
+    def release_xfer_handle(self, handle):
+        pass
+
+    def release_dlist_handle(self, handle):
+        pass
+
+    def send_notif(self, agent, notif_msg=None):
+        pass
+
+    def get_new_notifs(self):
+        return {}
+
+    def remove_remote_agent(self, agent):
+        pass
+
+
+def _make_mla_hybrid_worker(local_block_size, kernel_block_size, num_logical_blocks):
+    """Build a real pull worker with a hybrid MLA + 2xKDA HMA layout."""
+    from vllm.distributed.kv_transfer.kv_connector.v1.nixl import (
+        base_worker as bw,
+    )
+    from vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker import (
+        NixlConnectorWorker,
+    )
+    from vllm.v1.attention.backends.registry import MambaAttentionBackendEnum
+    from vllm.v1.kv_cache_interface import (
+        KVCacheConfig,
+        KVCacheGroupSpec,
+        KVCacheTensor,
+        MambaSpec,
+        MLAAttentionSpec,
+    )
+
+    mla_spec = MLAAttentionSpec(
+        block_size=local_block_size,
+        num_kv_heads=1,
+        head_size=6,
+        dtype=torch.float16,
+    )
+    unified_page = mla_spec.page_size_bytes
+    kda_spec = MambaSpec(
+        block_size=local_block_size,
+        shapes=((8, 3), (1, 4, 4)),
+        dtypes=(torch.float16, torch.float32),
+        page_size_padded=unified_page,
+        mamba_type=MambaAttentionBackendEnum.GDN_ATTN,
+    )
+    kv_cache_config = KVCacheConfig(
+        num_blocks=num_logical_blocks,
+        kv_cache_tensors=[
+            KVCacheTensor(
+                size=num_logical_blocks * unified_page,
+                shared_by=[f"mla.{i}", f"kda_a.{i}", f"kda_b.{i}"],
+            )
+            for i in range(2)
+        ],
+        kv_cache_groups=[
+            KVCacheGroupSpec(["mla.0", "mla.1"], mla_spec),
+            KVCacheGroupSpec(["kda_a.0", "kda_a.1"], kda_spec),
+            KVCacheGroupSpec(["kda_b.0", "kda_b.1"], kda_spec),
+        ],
+    )
+
+    vllm_config = create_vllm_config(block_size=local_block_size)
+    vllm_config.cache_config.enable_prefix_caching = False
+
+    from unittest.mock import MagicMock
+
+    fake_backend = MagicMock()
+    fake_backend.get_supported_kernel_block_sizes.return_value = [kernel_block_size]
+    fake_backend.get_name.return_value = "FLASHMLA"
+    fake_backend.full_cls_name.return_value = "fake.FLASHMLA"
+    fake_platform = MagicMock()
+    fake_platform.device_type = "cuda"
+    fake_platform.get_nixl_memory_type.return_value = "VRAM"
+
+    from vllm.config import set_current_vllm_config
+
+    with (
+        patch.object(bw, "NixlWrapper", _RecordingNixl),
+        patch.object(bw, "get_tensor_model_parallel_rank", return_value=0),
+        patch.object(bw, "get_tensor_model_parallel_world_size", return_value=1),
+        patch.object(bw, "get_current_attn_backends", return_value=[fake_backend]),
+        patch.object(bw, "current_platform", fake_platform),
+        patch(
+            "vllm.model_executor.layers.mamba.mamba_utils.get_conv_state_layout",
+            return_value="DS",
+        ),
+        set_current_vllm_config(vllm_config),
+    ):
+        worker = NixlConnectorWorker(vllm_config, "local-engine", kv_cache_config)
+        worker.use_mla = True
+
+        # Attention caches are kernel-block granular on dim 0, as the
+        # receive post-process assumes.
+        ppl = local_block_size // kernel_block_size
+        tensors = [
+            torch.zeros(
+                num_logical_blocks * ppl, unified_page // ppl, dtype=torch.uint8
+            )
+            for _ in range(2)
+        ]
+        worker.register_kv_caches(
+            {
+                "kda_a.0": tensors[0],
+                "mla.0": tensors[0],
+                "kda_b.0": tensors[0],
+                "kda_a.1": tensors[1],
+                "mla.1": tensors[1],
+                "kda_b.1": tensors[1],
+            }
+        )
+    # Keep tensors alive alongside the worker; flat views for byte checks.
+    worker._test_tensors = [t.view(-1) for t in tensors]
+    worker._test_tensors_2d = tensors
+    worker._test_unified_page = unified_page
+    return worker
+
+
+def _make_remote_meta(
+    worker,
+    remote_block_size,
+    remote_kernel_block_size,
+    remote_num_logical,
+    remote_ssm_sizes,
+):
+    from vllm.distributed.kv_transfer.kv_connector.v1.nixl.metadata import (
+        NixlAgentMetadata,
+    )
+
+    remote_ppl = remote_block_size // remote_kernel_block_size
+    # Kernel-granularity pages are TP-independent for MLA hybrids and must
+    # match the local ones for the handshake to pass.
+    kernel_page = worker.block_len_per_layer[0]
+    return NixlAgentMetadata(
+        engine_id="remote-engine",
+        agent_metadata=b"remote-agent-meta",
+        device_id=0,
+        kv_caches_base_addr=[0x10_000_000, 0x20_000_000],
+        num_blocks=remote_num_logical * remote_ppl,
+        block_lens=[kernel_page, kernel_page],
+        kv_cache_layout=worker.kv_cache_layout,
+        block_size=remote_kernel_block_size,
+        ssm_sizes=remote_ssm_sizes,
+        attn_backend_name=worker.backend_name,
+        physical_blocks_per_logical_kv_block=remote_ppl,
+    )
+
+
+def _owned_byte_ranges(worker, group_logical_ids):
+    """Byte ranges owned by a request: for each HMA region tensor, every
+    logical block id of every group maps to one unified page."""
+    unified_page = worker._test_unified_page
+    bases = [t.data_ptr() for t in worker._test_tensors]
+    owned = []
+    for base in bases:
+        for ids in group_logical_ids:
+            for b in ids:
+                owned.append((base + b * unified_page, base + (b + 1) * unified_page))
+    return owned
+
+
+def _assert_local_writes_within(worker, owned_ranges):
+    nixl = worker.nixl_wrapper
+    assert nixl.xfers, "no transfers were posted"
+    violations = []
+    total_descs = 0
+    for op, local_handle, local_ids, _, remote_ids in nixl.xfers:
+        assert len(local_ids) == len(remote_ids)
+        desc_arr = nixl.dlists[local_handle]
+        for i in local_ids:
+            addr, length, _dev = desc_arr[int(i)]
+            addr, length = int(addr), int(length)
+            total_descs += 1
+            if not any(lo <= addr and addr + length <= hi for lo, hi in owned_ranges):
+                violations.append((int(i), hex(addr), length))
+    assert not violations, (
+        f"{len(violations)}/{total_descs} local descriptors write outside "
+        f"the request's own blocks: {violations[:10]}"
+    )
+    return total_descs
+
+
+@pytest.mark.cpu_test
+def test_hetero_ppl_multi_read_writes_stay_within_request_blocks():
+    """MLA-hybrid hetero geometry: local (D, TP1) logical blocks of 12 tokens
+    (kernel 4, ppl=3) vs remote (P, TP2) logical blocks of 8 tokens (ppl=2),
+    equal kernel pages, tp_ratio=-2 multi-read with replicated MLA and
+    TP-sharded KDA state. Every local descriptor of the request's reads must
+    stay within its own blocks."""
+    from vllm.distributed.kv_transfer.kv_connector.v1.nixl.metadata import (
+        NixlConnectorMetadata,
+    )
+
+    worker = _make_mla_hybrid_worker(
+        local_block_size=12, kernel_block_size=4, num_logical_blocks=8
+    )
+    assert worker._physical_blocks_per_logical_kv_block == 3
+
+    meta_r = _make_remote_meta(
+        worker,
+        remote_block_size=8,
+        remote_kernel_block_size=4,
+        remote_num_logical=12,
+        remote_ssm_sizes=(24, 32),
+    )
+    for rank in (0, 1):
+        worker.add_remote_agent(meta_r, remote_tp_rank=rank, remote_tp_size=2)
+
+    # Request B: 17 matched tokens. Local: 2 logical blocks (24 tok
+    # capacity); remote: 16 prefilled tokens -> 2 remote logical blocks.
+    # Sparse, non-contiguous ids so neighbor blocks exist on all sides.
+    local_ids = ([2, 5], [1], [7])
+    remote_ids = [[1, 4], [5], [2]]
+
+    metadata = NixlConnectorMetadata()
+    metadata.add_new_req_to_recv(
+        request_id="req-b",
+        local_block_ids=local_ids,
+        kv_transfer_params={
+            "remote_block_ids": remote_ids,
+            "remote_engine_id": "remote-engine",
+            "remote_request_id": "prefill-req-b",
+            "remote_host": "localhost",
+            "remote_port": 1234,
+            "tp_size": 2,
+        },
+    )
+    meta = metadata.reqs_to_recv["req-b"]
+    meta.local_physical_block_ids = worker._logical_to_kernel_block_ids(
+        meta.local_block_ids, worker._physical_blocks_per_logical_kv_block
+    )
+    worker._recving_metadata["req-b"] = meta
+
+    worker._read_blocks_for_req("req-b", meta)
+
+    owned = _owned_byte_ranges(worker, local_ids)
+    total = _assert_local_writes_within(worker, owned)
+    # Multi-read: rank 0 carries the replicated MLA + its SSM shard,
+    # rank 1 carries only its SSM shard.
+    assert len(worker.nixl_wrapper.xfers) == 2
+    assert total > 0
+
+
+def _resolve(
+    desc_arr, idx, bases, region_size, unified_page, kernel_page, logical_ids_attn
+):
+    """Resolve a desc id to (region, kind, token_start) where kind is 'attn'
+    (kernel-page sized, block-aligned, in the request's attention blocks) or
+    'mamba'. token_start is the request-relative kernel-block index."""
+    addr, length, _ = (int(x) for x in desc_arr[int(idx)])
+    for region, base in enumerate(bases):
+        off = addr - base
+        if 0 <= off < region_size:
+            b = off // unified_page
+            rem = off % unified_page
+            if (
+                length == kernel_page
+                and rem % kernel_page == 0
+                and (b in logical_ids_attn)
+            ):
+                pos = logical_ids_attn.index(b)
+                tokens_per_block = unified_page // kernel_page
+                sub = rem // kernel_page
+                return (region, "attn", (pos * tokens_per_block + sub))
+            return (region, "mamba", None)
+    raise AssertionError(f"desc {idx} addr {addr:#x} not in any region")
+
+
+def _run_hetero_case(local_block, kernel, remote_block, num_tokens, tp_size=2):
+    """Full pull-path run for one geometry; returns pairing records."""
+    from vllm.distributed.kv_transfer.kv_connector.v1.nixl.metadata import (
+        NixlConnectorMetadata,
+    )
+
+    remote_ppl = remote_block // kernel
+    matched = num_tokens - 1  # mamba N-1 rule
+    n_local = -(-num_tokens // local_block)
+    n_remote = -(-matched // remote_block)
+
+    worker = _make_mla_hybrid_worker(
+        local_block_size=local_block,
+        kernel_block_size=kernel,
+        num_logical_blocks=max(2 * n_local + 4, 8),
+    )
+    # Local KDA state pages are (48, 64) bytes; the remote holds 1/tp_size
+    # shards of each.
+    meta_r = _make_remote_meta(
+        worker,
+        remote_block_size=remote_block,
+        remote_kernel_block_size=kernel,
+        remote_num_logical=max(2 * n_remote + 4, 8),
+        remote_ssm_sizes=(48 // tp_size, 64 // tp_size),
+    )
+    for rank in range(tp_size):
+        worker.add_remote_agent(meta_r, remote_tp_rank=rank, remote_tp_size=tp_size)
+
+    # Sparse ids so neighbors exist between the request's blocks.
+    local_attn = [2 * i + 1 for i in range(n_local)]
+    remote_attn = [2 * i + 2 for i in range(n_remote)]
+    local_ids = (local_attn, [0], [2 * n_local + 2])
+    remote_ids = [remote_attn, [1], [0]]
+
+    metadata = NixlConnectorMetadata()
+    metadata.add_new_req_to_recv(
+        request_id="req-b",
+        local_block_ids=local_ids,
+        kv_transfer_params={
+            "remote_block_ids": remote_ids,
+            "remote_engine_id": "remote-engine",
+            "remote_request_id": "prefill-req-b",
+            "remote_host": "localhost",
+            "remote_port": 1234,
+            "tp_size": tp_size,
+        },
+    )
+    meta = metadata.reqs_to_recv["req-b"]
+    meta.local_physical_block_ids = worker._logical_to_kernel_block_ids(
+        meta.local_block_ids, worker._physical_blocks_per_logical_kv_block
+    )
+    worker._recving_metadata["req-b"] = meta
+
+    # Sentinel-fill the local KV so untouched bytes are detectable.
+    for t in worker._test_tensors:
+        t.fill_(0xAA)
+
+    worker._read_blocks_for_req("req-b", meta)
+
+    # Invariant 1: all local writes within the request's own blocks.
+    owned = _owned_byte_ranges(worker, local_ids)
+    _assert_local_writes_within(worker, owned)
+
+    # Invariant 2: local<->remote attention pairs are token-aligned.
+    nixl = worker.nixl_wrapper
+    local_bases = [t.data_ptr() for t in worker._test_tensors]
+    remote_bases = [0x10_000_000, 0x20_000_000]
+    local_unified = worker._test_unified_page
+    remote_unified = (local_unified // local_block) * remote_block
+    kernel_page = worker.block_len_per_layer[0]
+    meta_r_num_blocks_bytes = (meta_r.num_blocks // remote_ppl) * remote_unified
+    covered_tokens = set()
+    for op, lh, lids, rh, rids in nixl.xfers:
+        larr, rarr = nixl.dlists[lh], nixl.dlists[rh]
+        for li, ri in zip(lids, rids):
+            lreg, lkind, ltok = _resolve(
+                larr,
+                li,
+                local_bases,
+                len(worker._test_tensors[0]),
+                local_unified,
+                kernel_page,
+                local_attn,
+            )
+            rreg, rkind, rtok = _resolve(
+                rarr,
+                ri,
+                remote_bases,
+                meta_r_num_blocks_bytes,
+                remote_unified,
+                kernel_page,
+                remote_attn,
+            )
+            assert lkind == rkind, (
+                f"pair kind mismatch: local {lkind} vs remote {rkind} "
+                f"(local desc {li}, remote desc {ri})"
+            )
+            assert lreg == rreg, (
+                f"region mismatch: local {lreg} vs remote {rreg} for "
+                f"tokens {ltok} vs {rtok}"
+            )
+            if lkind == "attn":
+                assert ltok == rtok, (
+                    f"TOKEN MISALIGNMENT: local kernel block holds tokens "
+                    f"[{ltok * kernel}..) but receives remote tokens "
+                    f"[{rtok * kernel}..) "
+                    f"(geometry local_block={local_block}, "
+                    f"remote_block={remote_block}, N={num_tokens})"
+                )
+                covered_tokens.add(ltok * kernel)
+
+    # Invariant 3: full coverage of the matched tokens.
+    needed = {t for t in range(0, matched - matched % kernel, kernel)}
+    missing = needed - covered_tokens
+    assert not missing, (
+        f"tokens never transferred: {sorted(missing)[:8]} "
+        f"(geometry local_block={local_block}, remote_block={remote_block}, "
+        f"N={num_tokens}, matched={matched})"
+    )
+
+    # Invariant 4: no stale bytes after receive completion. The scheduler
+    # excludes the blocks covering the matched tokens from alloc-time KV
+    # zeroing (the zeroing would race the RDMA write), so every byte of
+    # those blocks must be either written by the transfer or zeroed by the
+    # receive post-process. Stale bytes surface as mid-response garbage
+    # once decode grows into the untransferred tail.
+    for op, lh, lids, rh, rids in nixl.xfers:
+        larr = nixl.dlists[lh]
+        for li in lids:
+            addr, length, _ = (int(x) for x in larr[int(li)])
+            for t in worker._test_tensors:
+                off = addr - t.data_ptr()
+                if 0 <= off < t.numel():
+                    t[off : off + length] = 0  # simulate the RDMA write
+                    break
+    done_sending, done_recving = worker.get_finished()
+    assert "req-b" in done_recving
+    n_excluded = -(-matched // local_block)
+    stale = []
+    for b in local_attn[:n_excluded]:
+        for region, t in enumerate(worker._test_tensors):
+            page = t[b * local_unified : (b + 1) * local_unified]
+            n_stale = int((page == 0xAA).sum())
+            if n_stale:
+                stale.append((region, b, n_stale))
+    assert not stale, (
+        f"stale (unzeroed, untransferred) bytes in matched-range attention "
+        f"blocks (region, block, bytes): {stale} "
+        f"(geometry local_block={local_block}, remote_block={remote_block}, "
+        f"N={num_tokens}, matched={matched})"
+    )
+
+
+@pytest.mark.cpu_test
+@pytest.mark.parametrize(
+    "local_block,remote_block",
+    [
+        (12, 8),  # ppl 3 vs 2
+        (36, 8),  # ppl 9 vs 2 (large ppl asymmetry, scaled)
+        (24, 4),  # ppl 6 vs 1
+        (16, 24),  # remote larger than local (D_TP > P_TP direction)
+    ],
+)
+@pytest.mark.parametrize("num_tokens", list(range(2, 40)))
+def test_hetero_ppl_token_alignment_sweep(local_block, remote_block, num_tokens):
+    """Sweep prompt lengths across block-boundary residues for several
+    hetero-ppl geometries; assert neighbor-safety, token alignment, and
+    coverage of every transferred kernel block."""
+    _run_hetero_case(
+        local_block, kernel=4, remote_block=remote_block, num_tokens=num_tokens
+    )
+
+
+@pytest.mark.cpu_test
+@pytest.mark.parametrize(
+    "num_tokens",
+    # Residues around every geometric boundary: kernel block (64), remote
+    # logical block (768), local logical block (5760), plus odd offsets.
+    [
+        2,
+        63,
+        64,
+        65,
+        127,
+        128,
+        300,
+        640,
+        767,
+        768,
+        769,
+        831,
+        832,
+        1000,
+        1535,
+        1536,
+        1537,
+        2303,
+        2304,
+        2305,
+        3001,
+        5759,
+        5760,
+        5761,
+        5824,
+        6528,
+        6529,
+    ],
+)
+def test_mla_hybrid_large_ppl_geometry(num_tokens):
+    """KimiLinear-scale MLA-hybrid geometry (TP8 prefill -> TP1 decode):
+    decode (local) logical block 5760 / kernel 64 (ppl=90), prefill
+    (remote) logical block 768 (ppl=12), tp_ratio=-8 multi-read with
+    replicated MLA and 8-way TP-sharded KDA state."""
+    _run_hetero_case(
+        local_block=5760,
+        kernel=64,
+        remote_block=768,
+        num_tokens=num_tokens,
+        tp_size=8,
+    )

tests/v1/kv_connector/unit/test_tp_mapping.py

@@ -161,3 +161,57 @@ class TestMambaPlanSplitHandles:
         # FA: chunk=200//1=200, slot=0 (skip_fa) → (1000, 200, 0), (2000, 200, 0)
         # SSM: chunk=400//2=200, idx=1 → (3200, 200, 0)
         assert splits[1] == [(1000, 200, 0), (2000, 200, 0), (3200, 200, 0)]
+
+    def test_hetero_block_size_splits(self):
+        """With a block-size ratio, single-source FA sub-block descs pass
+        through whole; SSM descs are unexpanded and split per source."""
+        plan = TPMapping(
+            source_ranks_per_group=((0,), (0, 1)),
+            all_source_ranks=(0, 1),
+            rank_to_attention_slot={0: 0, 1: 0},
+            rank_offset_factor=0,
+        )
+
+        worker = _make_mock_worker_for_splits((FullAttentionSpec, MambaSpec))
+        # 2 FA blocks x ratio 2 sub-blocks + 1 SSM desc (never expanded).
+        src_blocks_data = np.array(
+            [
+                (1000, 100, 0),
+                (1100, 100, 0),
+                (2000, 100, 0),
+                (2100, 100, 0),
+                (3000, 400, 0),
+            ],
+            dtype=np.uint64,
+        )
+
+        splits = list(worker._build_local_splits_from_plan(plan, src_blocks_data, 4, 2))
+
+        assert len(splits) == 2
+        fa_passthrough = [
+            (1000, 100, 0),
+            (1100, 100, 0),
+            (2000, 100, 0),
+            (2100, 100, 0),
+        ]
+        assert splits[0] == fa_passthrough + [(3000, 200, 0)]
+        assert splits[1] == fa_passthrough + [(3200, 200, 0)]
+
+    def test_hetero_block_size_head_sharded_asserts(self):
+        """Head-sharded FA reads (multiple FA sources) are incompatible with
+        a block-size mismatch and must fail loudly."""
+        plan = TPMapping(
+            source_ranks_per_group=((0, 1), (0, 1)),
+            all_source_ranks=(0, 1),
+            rank_to_attention_slot={0: 0, 1: 1},
+            rank_offset_factor=0,
+        )
+
+        worker = _make_mock_worker_for_splits((FullAttentionSpec, MambaSpec))
+        src_blocks_data = np.array(
+            [(1000, 100, 0), (1100, 100, 0), (3000, 400, 0)],
+            dtype=np.uint64,
+        )
+
+        with pytest.raises(AssertionError, match="Head-sharded"):
+            list(worker._build_local_splits_from_plan(plan, src_blocks_data, 2, 2))

tests/v1/spec_decode/test_llm_base_proposer_sampling.py

@@ -1,6 +1,9 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
+from types import SimpleNamespace
+
+import pytest
 import torch
 
 from vllm.platforms import current_platform
@@ -8,6 +11,7 @@ from vllm.utils.torch_utils import set_random_seed
 from vllm.v1.sample.logits_processor import LogitsProcessors
 from vllm.v1.sample.metadata import SamplingMetadata
 from vllm.v1.spec_decode.llm_base_proposer import (
+    SpecDecodeBaseProposer,
     compute_probs_and_sample_next_token,
 )
 
@@ -68,3 +72,21 @@ def test_compute_probs_and_sample_next_token_uses_fp64_exponential_race():
 
     assert torch.equal(actual_ids, expected_ids)
     assert torch.allclose(actual_probs, probs)
+
+
+@pytest.mark.parametrize(
+    ("architecture", "expected"),
+    [
+        ("DeepSeekMTPModel", True),
+        ("KimiK3MTPModel", True),
+        ("MiniMaxM3ForCausalLM", False),
+    ],
+)
+def test_mtp_model_returns_tuple(architecture: str, expected: bool):
+    proposer = object.__new__(SpecDecodeBaseProposer)
+    proposer.method = "mtp"
+    proposer.draft_model_config = SimpleNamespace(
+        hf_config=SimpleNamespace(architectures=[architecture])
+    )
+
+    assert proposer.model_returns_tuple() is expected

tests/v1/spec_decode/test_mtp_structured_output.py

@@ -341,3 +341,84 @@ def test_trim_reasoning_for_advance():
     next_step = [post, post]
     request.append_output_token_ids(next_step)
     assert manager.trim_reasoning_for_advance(request, next_step) == next_step
+
+
+def test_should_advance_boundary_applies_to_json_grammars():
+    """With speculative decoding, the boundary step carries tokens sampled
+    after the reasoning-end marker for *every* grammar kind, not just
+    structural tags. Deferring the advance desyncs the FSM from the emitted
+    tokens (the next step re-derives from the initial state and duplicates
+    the grammar prefix, e.g. '{"' + '{"city": ...}').
+    """
+    from vllm.v1.structured_output.backend_types import StructuredOutputOptions
+
+    tokenizer, manager, request, prompt, marker = _setup_boundary_request("xgrammar")
+    structured_req = request.structured_output_request
+    structured_req.__dict__["structured_output_key"] = (
+        StructuredOutputOptions.JSON,
+        "{}",
+    )
+
+    pre = tokenizer.encode(" ")[0]
+    post = tokenizer.encode("{")[0]
+    step_tokens = [pre, marker, post]
+    request.append_output_token_ids(step_tokens)
+
+    assert manager.should_advance(request)
+    assert structured_req.reasoning_ended
+    assert structured_req.reasoning_end_token_index == len(prompt) + 1
+    # The scheduler advances with the trimmed suffix: only the post-marker
+    # token reaches the grammar.
+    assert manager.trim_reasoning_for_advance(request, step_tokens) == [post]
+
+
+class _K3StyleReasoner:
+    """Multi-token markers with last-close-after-last-open semantics,
+    mirroring KimiK3ReasoningParser (close=[8,2,9], open=[7,2,9])."""
+
+    _CLOSE = [8, 2, 9]
+    _OPEN = [7, 2, 9]
+
+    @staticmethod
+    def _last(ids: list[int], sub: list[int]) -> int:
+        for i in range(len(ids) - len(sub), -1, -1):
+            if ids[i : i + len(sub)] == sub:
+                return i
+        return -1
+
+    def is_reasoning_end(self, input_ids) -> bool:
+        ids = list(input_ids)
+        last_close = self._last(ids, self._CLOSE)
+        last_open = self._last(ids, self._OPEN)
+        if last_open == -1:
+            return last_close != -1
+        return last_close > last_open
+
+    def is_reasoning_end_streaming(self, input_ids, delta_ids) -> bool:
+        return self.is_reasoning_end(input_ids)
+
+
+def test_find_reasoning_end_index_is_frame_independent():
+    """Regression: the boundary index must anchor at the prompt end, not the
+    step frame. With spec decode + async scheduling the frame can start past
+    the marker; a frame-anchored walk fires immediately for full-sequence
+    parsers and reports the frame start, so the trim drops grammar tokens
+    (or keeps marker tokens) and the FSM desyncs from the emitted stream.
+    """
+    # Prompt: stale close from a prior turn, then the current think-open
+    # (the agent-continuation shape).
+    prompt = [1, 8, 2, 9, 5, 7, 2, 9]
+    # Output: reasoning tokens, close marker, then grammar tokens.
+    out = [11, 12, 8, 2, 9, 30, 31]
+    all_ids = prompt + out
+    marker_end = len(prompt) + 4
+
+    reasoner = _K3StyleReasoner()
+    assert not reasoner.is_reasoning_end(prompt)
+    assert reasoner.is_reasoning_end(all_ids)
+
+    # Independent of any frame: always the marker's final-token index.
+    got = StructuredOutputManager._find_reasoning_end_index(
+        reasoner, all_ids, len(prompt)
+    )
+    assert got == marker_end

tests/v1/structured_output/test_reasoning_structured_output.py

@@ -330,6 +330,9 @@ class TestReasoningStructuredOutput:
             def __init__(self, *_, **__):
                 pass
 
+            def is_reasoning_end(self, input_ids):
+                return marker in list(input_ids)
+
             def is_reasoning_end_streaming(self, input_ids, delta_ids):
                 return marker in list(delta_ids)