[Model] Add Kimi K3 support: model files and kernels [1/N] (#50089)

Signed-off-by: zjy0516 <riverclouds.zhu@qq.com>
Signed-off-by: Isotr0py <Isotr0py@outlook.com>
Signed-off-by: Bugen Zhao <i@bugenzhao.com>
Signed-off-by: Yongye Zhu <zyy1102000@gmail.com>
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: Summer Yang <girasoleyang@gmail.com>
Co-authored-by: Kevin H. Luu <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>
Co-authored-by: Codex <codex@openai.com>

.buildkite/check-torch-abi.py

@@ -13,6 +13,7 @@ from torch_abi_audit.report import ExtensionReport, PackageReport
 # Temporary allowlist of extensions not yet on the stable ABI.
 # Shrink and remove over time.
 ALLOWED_UNSTABLE_LIBRARIES: tuple[str, ...] = (
+    "_flashkda_C.abi3.so",
     "vllm_flash_attn/_vllm_fa2_C.abi3.so",
     "vllm_flash_attn/_vllm_fa3_C.abi3.so",
     "third_party/deep_gemm/_C*.so",

.buildkite/test_areas/kernels.yaml

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

.buildkite/test_areas/models_basic.yaml

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

.pre-commit-config.yaml

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

CMakeLists.txt

@@ -421,8 +421,11 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
     "csrc/libtorch_stable/mamba/selective_scan_fwd.cu"
     "csrc/libtorch_stable/cache_kernels.cu"
     "csrc/libtorch_stable/cache_kernels_fused.cu"
+    "csrc/libtorch_stable/custom_all_gather_reduce_scatter.cu"
+    "csrc/libtorch_stable/custom_all_gather_reduce_scatter_ops.cpp"
     "csrc/libtorch_stable/custom_all_reduce.cu"
-    "csrc/libtorch_stable/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu")
+    "csrc/libtorch_stable/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu"
+    "csrc/libtorch_stable/fused_kimi_k3_mla_key_concat_kv_cache_kernel.cu")
 
   if(VLLM_GPU_LANG STREQUAL "CUDA" AND
      DEFINED CMAKE_CUDA_COMPILER_VERSION AND
@@ -1079,6 +1082,23 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
     set(MLA_ARCHS)
   endif()
 
+  if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
+    cuda_archs_loose_intersection(FUSED_KDA_DECODE_ARCHS
+      "9.0a;10.0f;12.0f" "${CUDA_ARCHS}")
+  endif()
+  if(FUSED_KDA_DECODE_ARCHS)
+    set(FUSED_KDA_DECODE_SRC
+      "csrc/libtorch_stable/kimi_k3/fused_kda_decode_kernel.cu")
+    set_gencode_flags_for_srcs(
+      SRCS "${FUSED_KDA_DECODE_SRC}"
+      CUDA_ARCHS "${FUSED_KDA_DECODE_ARCHS}")
+    set_property(SOURCE ${FUSED_KDA_DECODE_SRC} APPEND PROPERTY
+      COMPILE_OPTIONS "$<$<COMPILE_LANGUAGE:CUDA>:--use_fast_math>")
+    list(APPEND VLLM_STABLE_EXT_SRC "${FUSED_KDA_DECODE_SRC}")
+    message(STATUS
+      "Building fused KDA decode for archs: ${FUSED_KDA_DECODE_ARCHS}")
+  endif()
+
   if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
     cuda_archs_loose_intersection(KIMI_K3_ATTN_RES_ARCHS
       "10.0f" "${CUDA_ARCHS}")
@@ -1138,6 +1158,10 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
       target_compile_definitions(_C_stable_libtorch PRIVATE
         VLLM_ENABLE_COOPERATIVE_TOPK=1)
     endif()
+    if(FUSED_KDA_DECODE_ARCHS)
+      target_compile_definitions(_C_stable_libtorch PRIVATE
+        VLLM_ENABLE_FUSED_KDA_DECODE=1)
+    endif()
     if(KIMI_K3_ATTN_RES_ARCHS)
       target_compile_definitions(_C_stable_libtorch PRIVATE
         VLLM_ENABLE_KIMI_K3_ATTN_RES=1)
@@ -1439,6 +1463,7 @@ if (VLLM_GPU_LANG STREQUAL "CUDA")
     include(cmake/external_projects/deepgemm.cmake)
     include(cmake/external_projects/fmha_sm100.cmake)
     include(cmake/external_projects/flashmla.cmake)
+    include(cmake/external_projects/flashkda.cmake)
     include(cmake/external_projects/qutlass.cmake)
     include(cmake/external_projects/tml_fa4.cmake)
 

benchmarks/kernels/benchmark_k3_cutedsl_residual.py

@@ -0,0 +1,367 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Benchmark the Kimi-K3 latent MoE addmm against CuTe residual GEMM.
+
+The benchmark covers ``BF16[M, 3584] @ BF16[7168, 3584].T + BF16[M, 7168]``
+with FP32 accumulation and BF16 output. Both backends execute through CUDA
+Graph replay. Weights and residuals rotate across buffers exceeding L2 so the
+comparison models the full latent MoE projection-and-add path.
+"""
+
+from __future__ import annotations
+
+import argparse
+import dataclasses
+import importlib.util
+import json
+import math
+import statistics
+from collections.abc import Callable, Sequence
+from pathlib import Path
+from typing import Any
+
+import cutlass
+import cutlass.cute as cute
+import torch
+from cuda.bindings import driver as cuda
+from cuda.bindings.driver import CUstream
+from quack.compile_utils import make_fake_tensor
+
+N = 7168
+K = 3584
+
+
+@dataclasses.dataclass(frozen=True, slots=True)
+class Config:
+    block_size: int
+    outputs_per_block: int
+    k_unroll: int
+    vector_width: int = 8
+
+
+def parse_config(value: str) -> Config:
+    try:
+        parts = [int(part) for part in value.split(",")]
+    except ValueError as error:
+        raise argparse.ArgumentTypeError(
+            "config must be BLOCK,OUTPUTS,K_UNROLL[,VECTOR_WIDTH]"
+        ) from error
+    if len(parts) == 3:
+        return Config(*parts)
+    if len(parts) == 4:
+        return Config(*parts)
+    raise argparse.ArgumentTypeError(
+        "config must be BLOCK,OUTPUTS,K_UNROLL[,VECTOR_WIDTH]"
+    )
+
+
+def production_residual_config(m: int) -> Config | None:
+    """The measured Latent-MoE residual config for M, from the K3 table."""
+    from vllm.models.kimi_k3.nvidia.low_latency_gemm import KIMI_K3_PROJECTIONS
+
+    spec = KIMI_K3_PROJECTIONS.get((N, K))
+    config = spec.residual_config(m) if spec is not None else None
+    if config is None:
+        return None
+    return Config(
+        config.block_size,
+        config.outputs_per_block,
+        config.k_unroll,
+        config.vector_width,
+    )
+
+
+def candidate_configs(mode: str, selected: Config | None, m: int) -> list[Config]:
+    if mode == "selected":
+        if selected is not None:
+            return [selected]
+        # No explicit --config: fall back to the production table for this M.
+        config = production_residual_config(m)
+        return [config] if config is not None else []
+    if mode == "baseline":
+        return [Config(224, 4, 2)]
+    return [
+        Config(block_size, outputs_per_block, k_unroll, vector_width)
+        for vector_width in (4, 8)
+        for block_size in (32, 64, 128, 224, 448)
+        if block_size % 32 == 0 and K % (block_size * vector_width) == 0
+        for outputs_per_block in (1, 2, 4, 7, 8)
+        if N % outputs_per_block == 0
+        for k_unroll in (1, 2, 4)
+    ]
+
+
+def load_kernel_class(path: Path):
+    spec = importlib.util.spec_from_file_location("cute_skinny_device", path)
+    if spec is None or spec.loader is None:
+        raise RuntimeError(f"cannot load CuTe kernel from {path}")
+    module = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(module)
+    return module.CuteSkinnyGemm
+
+
+def stream() -> CUstream:
+    return CUstream(torch.cuda.current_stream().cuda_stream)
+
+
+def compile_kernel(kernel_class, m: int, config: Config, max_registers: int):
+    element_type = cutlass.BFloat16
+    n = cute.sym_int(divisibility=config.outputs_per_block)
+    k = cute.sym_int(divisibility=config.block_size * config.vector_width)
+    a = make_fake_tensor(element_type, (m, k), divisibility=config.vector_width)
+    b = make_fake_tensor(element_type, (n, k), divisibility=config.vector_width)
+    residual = make_fake_tensor(element_type, (m, n), divisibility=1)
+    c = make_fake_tensor(element_type, (m, n), divisibility=1)
+    kernel = kernel_class(
+        element_type=element_type,
+        num_rows=m,
+        block_size=config.block_size,
+        outputs_per_block=config.outputs_per_block,
+        vector_width=config.vector_width,
+        k_unroll=config.k_unroll,
+        has_residual=True,
+        use_pdl=True,
+    )
+    return cute.compile(
+        kernel,
+        a,
+        b,
+        residual,
+        c,
+        stream(),
+        options=(
+            "--enable-tvm-ffi --keep-cubin "
+            f"--ptxas-options -maxrregcount={max_registers} "
+            "--ptxas-options -lineinfo"
+        ),
+    )
+
+
+def resource_usage(compiled) -> dict[str, Any]:
+    executor = getattr(compiled, "_default_executor", None)
+    context = getattr(executor, "exec_context", None)
+    functions = getattr(context, "kernel_functions", None)
+    if not functions:
+        return {"resource_metrics_available": False}
+
+    def attribute(name, function) -> int:
+        error, value = cuda.cuFuncGetAttribute(name, function)
+        if error != cuda.CUresult.CUDA_SUCCESS:
+            raise RuntimeError(f"cuFuncGetAttribute failed with {error}")
+        return int(value)
+
+    registers = [
+        attribute(cuda.CUfunction_attribute.CU_FUNC_ATTRIBUTE_NUM_REGS, function)
+        for function in functions
+    ]
+    local_bytes = [
+        attribute(
+            cuda.CUfunction_attribute.CU_FUNC_ATTRIBUTE_LOCAL_SIZE_BYTES,
+            function,
+        )
+        for function in functions
+    ]
+    return {
+        "resource_metrics_available": True,
+        "registers_per_thread": max(registers, default=0),
+        "spill_bytes": max(local_bytes, default=0),
+    }
+
+
+def rotating_buffer_count(m: int, multiplier: float, limit: int) -> int:
+    properties = torch.cuda.get_device_properties(0)
+    bytes_per_pair = (N * K + m * N) * 2
+    target = math.ceil(multiplier * properties.L2_cache_size)
+    return max(2, min(limit, math.ceil(target / bytes_per_pair)))
+
+
+def graph_samples(
+    launch: Callable[[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor], None],
+    activation: torch.Tensor,
+    weights: Sequence[torch.Tensor],
+    residuals: Sequence[torch.Tensor],
+    repeats: int,
+    replays: int,
+) -> tuple[list[float], list[torch.Tensor]]:
+    outputs = [torch.empty_like(residual) for residual in residuals]
+    for weight, residual, output in zip(weights, residuals, outputs):
+        launch(activation, weight, residual, output)
+    torch.accelerator.synchronize()
+
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph):
+        for weight, residual, output in zip(weights, residuals, outputs):
+            launch(activation, weight, residual, output)
+    for _ in range(20):
+        graph.replay()
+    torch.accelerator.synchronize()
+
+    samples = []
+    for _ in range(repeats):
+        start = torch.cuda.Event(enable_timing=True)
+        end = torch.cuda.Event(enable_timing=True)
+        start.record()
+        for _ in range(replays):
+            graph.replay()
+        end.record()
+        end.synchronize()
+        samples.append(start.elapsed_time(end) * 1000.0 / (replays * len(weights)))
+    return samples, outputs
+
+
+def summarize(samples: Sequence[float]) -> dict[str, Any]:
+    ordered = sorted(samples)
+
+    def percentile(fraction: float) -> float:
+        position = fraction * (len(ordered) - 1)
+        lower = math.floor(position)
+        upper = math.ceil(position)
+        if lower == upper:
+            return ordered[lower]
+        weight = position - lower
+        return ordered[lower] * (1.0 - weight) + ordered[upper] * weight
+
+    mean = statistics.mean(samples)
+    return {
+        "median_us": statistics.median(samples),
+        "p10_us": percentile(0.1),
+        "p90_us": percentile(0.9),
+        "mean_us": mean,
+        "cv_pct": statistics.pstdev(samples) / mean * 100.0,
+        "samples_us": list(samples),
+    }
+
+
+def correctness(
+    output: torch.Tensor,
+    activation: torch.Tensor,
+    weight: torch.Tensor,
+    residual: torch.Tensor,
+) -> dict[str, Any]:
+    actual = output.float()
+    reference = activation.float() @ weight.float().t() + residual.float()
+    error = (actual - reference).abs()
+    scaled_error = error / (reference.abs() + 1.0)
+    cosine = torch.nn.functional.cosine_similarity(
+        actual.flatten(), reference.flatten(), dim=0
+    ).item()
+    return {
+        "valid": cosine > 0.999,
+        "cosine": cosine,
+        "max_abs_error": error.max().item(),
+        "max_scaled_error": scaled_error.max().item(),
+    }
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--kernel", type=Path, required=True)
+    parser.add_argument("--output", type=Path, required=True)
+    parser.add_argument(
+        "--mode", choices=("baseline", "sweep", "selected"), default="baseline"
+    )
+    parser.add_argument("--config", type=parse_config)
+    parser.add_argument("--m", type=int, action="append")
+    parser.add_argument("--config-shard", type=int, default=0)
+    parser.add_argument("--num-config-shards", type=int, default=1)
+    parser.add_argument("--repeats", type=int, default=21)
+    parser.add_argument("--replays", type=int, default=200)
+    parser.add_argument("--cache-multiplier", type=float, default=3.0)
+    parser.add_argument("--max-buffers", type=int, default=32)
+    parser.add_argument("--max-registers", type=int, default=64)
+    args = parser.parse_args()
+
+    token_counts = args.m or list(range(1, 17))
+    if any(not 1 <= m <= 16 for m in token_counts):
+        raise ValueError("expected 1 <= M <= 16")
+    if not 0 <= args.config_shard < args.num_config_shards:
+        raise ValueError("config shard must be in [0, num_config_shards)")
+    torch.accelerator.set_device_index(0)
+    if torch.cuda.get_device_capability() != (10, 3):
+        raise RuntimeError("this benchmark requires SM103")
+
+    kernel_class = load_kernel_class(args.kernel)
+    properties = torch.cuda.get_device_properties(0)
+    metadata = {
+        "device": properties.name,
+        "compute_capability": list(torch.cuda.get_device_capability()),
+        "torch_version": torch.__version__,
+        "cuda_version": torch.version.cuda,
+    }
+    args.output.parent.mkdir(parents=True, exist_ok=True)
+    with args.output.open("w", encoding="utf-8") as output_file:
+        for m in token_counts:
+            configs = candidate_configs(args.mode, args.config, m)
+            torch.manual_seed(20260722 + m)
+            count = rotating_buffer_count(m, args.cache_multiplier, args.max_buffers)
+            activation = torch.randn((m, K), device="cuda", dtype=torch.bfloat16)
+            weights = [
+                torch.randn((N, K), device="cuda", dtype=torch.bfloat16)
+                for _ in range(count)
+            ]
+            residuals = [
+                torch.randn((m, N), device="cuda", dtype=torch.bfloat16)
+                for _ in range(count)
+            ]
+            candidates: list[tuple[str, Config | None]] = [("cublas_addmm", None)]
+            candidates.extend(
+                ("cute_residual", config)
+                for index, config in enumerate(configs)
+                if index % args.num_config_shards == args.config_shard
+            )
+            for backend, config in candidates:
+                row: dict[str, Any] = {
+                    "m": m,
+                    "n": N,
+                    "k": K,
+                    "backend": backend,
+                    "mode": args.mode,
+                    "config": dataclasses.asdict(config) if config else {},
+                    "num_buffers": count,
+                    "cache_multiplier": args.cache_multiplier,
+                    **metadata,
+                }
+                try:
+                    if backend == "cublas_addmm":
+                        launch = lambda a, b, residual, c: torch.addmm(
+                            residual, a, b.t(), out=c
+                        )
+                    else:
+                        if config is None:
+                            raise AssertionError("missing CuTe config")
+                        compiled = compile_kernel(
+                            kernel_class, m, config, args.max_registers
+                        )
+                        launch = lambda a, b, residual, c, fn=compiled: fn(
+                            a, b, residual, c, stream()
+                        )
+                        row.update(resource_usage(compiled))
+                    samples, outputs = graph_samples(
+                        launch,
+                        activation,
+                        weights,
+                        residuals,
+                        args.repeats,
+                        args.replays,
+                    )
+                    row.update(
+                        correctness(outputs[0], activation, weights[0], residuals[0])
+                    )
+                    row.update(summarize(samples))
+                except Exception as error:  # noqa: BLE001
+                    row.update(
+                        {
+                            "valid": False,
+                            "error": f"{type(error).__name__}: {error}",
+                        }
+                    )
+                output_file.write(json.dumps(row, sort_keys=True) + "\n")
+                output_file.flush()
+                print(json.dumps(row, sort_keys=True), flush=True)
+
+            del activation, weights, residuals
+            torch.accelerator.empty_cache()
+
+
+if __name__ == "__main__":
+    main()

benchmarks/kernels/benchmark_kimi_k3_latent_moe_tail.py

@@ -0,0 +1,806 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Benchmark the Kimi K3 latent-MoE tail and its up-projection kernels.
+
+The ``up-projection`` subcommand isolates the TP-local dynamic and static-M
+skinny GEMMs. It rotates weights through a working set larger than L2 to model
+successive model layers.
+
+The ``whole-tail`` subcommand measures the distributed operator. Its reference
+path includes two AllReduces, RMSNorm, the replicated up-projection, and the
+final add. CUDA-event samples report the slowest rank so cross-rank skew is
+included.
+
+Examples:
+
+.. code-block:: console
+
+    .venv/bin/python \
+      benchmarks/kernels/benchmark_kimi_k3_latent_moe_tail.py up-projection
+
+    torchrun --nproc-per-node=8 \
+      benchmarks/kernels/benchmark_kimi_k3_latent_moe_tail.py whole-tail
+
+For multi-node runs, launch one ``torchrun`` agent per node and use a shared
+rendezvous endpoint.
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import math
+import os
+import statistics
+from collections.abc import Callable, Sequence
+from dataclasses import asdict
+from pathlib import Path
+from typing import Any
+
+import cutlass
+import cutlass.utils as utils
+import torch
+import torch.distributed as dist
+import torch.nn.functional as F
+from cuda.bindings import driver as cuda
+
+from vllm.distributed import get_tp_group
+from vllm.distributed.parallel_state import (
+    init_distributed_environment,
+    initialize_model_parallel,
+    set_custom_all_reduce,
+)
+from vllm.model_executor.warmup.cutedsl_warmup import cutedsl_warmup
+from vllm.models.kimi_k3.nvidia.ops import latent_moe_tail
+from vllm.models.kimi_k3.nvidia.ops.cute_dsl.latent_moe_tail import (
+    fused_add_multicast_gemm,
+    fused_add_multicast_skinny_gemm,
+)
+
+HIDDEN_SIZE = 7168
+LATENT_SIZE = 3584
+RMS_EPS = 0.1
+MAX_NUM_TOKENS = 16
+MMA_TILER_MN = (64, 32)
+CLUSTER_SHAPE_MN = (1, 8)
+B_PRIME_STAGES = 2
+
+
+def parse_up_projection_config(
+    value: str,
+) -> fused_add_multicast_skinny_gemm.SkinnyConfig:
+    try:
+        values = [int(part) for part in value.split(",")]
+    except ValueError as error:
+        raise argparse.ArgumentTypeError(
+            "config must be BLOCK,OUTPUTS,K_UNROLL[,VECTOR_WIDTH[,PREFETCH_B]]"
+        ) from error
+    if len(values) in (3, 4):
+        return fused_add_multicast_skinny_gemm.SkinnyConfig(*values)
+    if len(values) == 5 and values[4] in (0, 1):
+        return fused_add_multicast_skinny_gemm.SkinnyConfig(
+            *values[:4],
+            prefetch_b_before_pdl=bool(values[4]),
+        )
+    raise argparse.ArgumentTypeError(
+        "config must be BLOCK,OUTPUTS,K_UNROLL"
+        "[,VECTOR_WIDTH[,PREFETCH_B]], where PREFETCH_B is 0 or 1"
+    )
+
+
+def parse_tail_skinny_config(
+    value: str,
+) -> tuple[int, fused_add_multicast_skinny_gemm.SkinnyConfig]:
+    try:
+        values = [int(part) for part in value.split(",")]
+    except ValueError as error:
+        raise argparse.ArgumentTypeError(
+            "config must be M,BLOCK,OUTPUTS,K_UNROLL[,VECTOR_WIDTH[,PREFETCH_B]]"
+        ) from error
+    if len(values) == 4:
+        num_tokens, *config = values
+        return num_tokens, fused_add_multicast_skinny_gemm.SkinnyConfig(*config)
+    if len(values) == 5:
+        num_tokens, *config = values
+        return num_tokens, fused_add_multicast_skinny_gemm.SkinnyConfig(*config)
+    if len(values) == 6 and values[5] in (0, 1):
+        num_tokens, block, outputs, unroll, vector_width, prefetch = values
+        return num_tokens, fused_add_multicast_skinny_gemm.SkinnyConfig(
+            block,
+            outputs,
+            unroll,
+            vector_width,
+            bool(prefetch),
+        )
+    raise argparse.ArgumentTypeError(
+        "config must be M,BLOCK,OUTPUTS,K_UNROLL"
+        "[,VECTOR_WIDTH[,PREFETCH_B]], where PREFETCH_B is 0 or 1"
+    )
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description=__doc__)
+    subparsers = parser.add_subparsers(dest="scope", required=True)
+
+    up_projection = subparsers.add_parser(
+        "up-projection",
+        help="Benchmark the isolated TP-local up-projection kernels.",
+    )
+    up_projection.add_argument(
+        "--backend",
+        choices=("dynamic", "skinny", "both"),
+        default="both",
+    )
+    up_projection.add_argument("--tp-size", type=int, default=16)
+    up_projection.add_argument(
+        "--num-tokens",
+        type=int,
+        nargs="+",
+        default=[*range(1, 9), 16],
+    )
+    up_projection.add_argument(
+        "--skinny-config",
+        type=parse_up_projection_config,
+        action="append",
+        help="Benchmark a static-M config for every selected token count.",
+    )
+    up_projection.add_argument("--cache-multiplier", type=float, default=2.0)
+    up_projection.add_argument("--max-weights", type=int, default=64)
+    up_projection.add_argument("--warmup-replays", type=int, default=10)
+    up_projection.add_argument("--samples", type=int, default=31)
+    up_projection.add_argument("--output", type=Path)
+
+    whole_tail = subparsers.add_parser(
+        "whole-tail",
+        help="Benchmark the distributed latent-MoE tail operator.",
+    )
+    whole_tail.add_argument(
+        "--backend",
+        choices=("reference", "fused", "both"),
+        default="both",
+    )
+    whole_tail.add_argument(
+        "--num-tokens",
+        type=int,
+        nargs="+",
+        default=[1, 5, 8, 16],
+    )
+    whole_tail.add_argument("--warmup-replays", type=int, default=20)
+    whole_tail.add_argument("--samples", type=int, default=51)
+    whole_tail.add_argument(
+        "--skinny-max-num-tokens",
+        type=int,
+        nargs="+",
+        help="Override the fused operator's static-M cutoff; use 0 for dynamic-only.",
+    )
+    whole_tail.add_argument(
+        "--skinny-config",
+        type=parse_tail_skinny_config,
+        action="append",
+        help="Override one static-M config for tuning.",
+    )
+    whole_tail.add_argument("--output", type=Path)
+    return parser.parse_args()
+
+
+def percentile(samples: Sequence[float], fraction: float) -> float:
+    ordered = sorted(samples)
+    position = fraction * (len(ordered) - 1)
+    lower = math.floor(position)
+    upper = math.ceil(position)
+    if lower == upper:
+        return ordered[lower]
+    upper_weight = position - lower
+    return ordered[lower] * (1.0 - upper_weight) + ordered[upper] * upper_weight
+
+
+def summarize(samples_us: Sequence[float]) -> dict[str, Any]:
+    mean_us = statistics.mean(samples_us)
+    return {
+        "median_us": statistics.median(samples_us),
+        "p10_us": percentile(samples_us, 0.1),
+        "p90_us": percentile(samples_us, 0.9),
+        "mean_us": mean_us,
+        "cv_pct": statistics.pstdev(samples_us) / mean_us * 100.0,
+        "samples_us": list(samples_us),
+    }
+
+
+def rotating_weight_count(
+    shard_size: int,
+    cache_multiplier: float,
+    limit: int,
+) -> int:
+    properties = torch.cuda.get_device_properties(
+        torch.accelerator.current_device_index()
+    )
+    weight_bytes = shard_size * LATENT_SIZE * 2
+    target_bytes = math.ceil(properties.L2_cache_size * cache_multiplier)
+    return max(2, min(limit, math.ceil(target_bytes / weight_bytes)))
+
+
+def capture_up_projection_graph(
+    launches: Sequence[Callable[[], None]],
+) -> torch.cuda.CUDAGraph:
+    for launch in launches:
+        launch()
+    torch.accelerator.synchronize()
+
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph):
+        for launch in launches:
+            launch()
+    torch.accelerator.synchronize()
+    return graph
+
+
+def benchmark_up_projection_graph(
+    graph: torch.cuda.CUDAGraph,
+    *,
+    operations_per_replay: int,
+    warmup_replays: int,
+    samples: int,
+) -> dict[str, Any]:
+    for _ in range(warmup_replays):
+        graph.replay()
+    torch.accelerator.synchronize()
+
+    samples_us = []
+    start = torch.cuda.Event(enable_timing=True)
+    end = torch.cuda.Event(enable_timing=True)
+    for _ in range(samples):
+        start.record()
+        graph.replay()
+        end.record()
+        end.synchronize()
+        samples_us.append(start.elapsed_time(end) * 1000.0 / operations_per_replay)
+    return summarize(samples_us)
+
+
+class DynamicKernel:
+    def __init__(
+        self,
+        shard_size: int,
+        mailbox: torch.Tensor,
+        shared_shard: torch.Tensor,
+    ) -> None:
+        self.shard_size = shard_size
+        self.mailbox = mailbox
+        self.mailbox_c = fused_add_multicast_gemm._as_cute(mailbox)
+        compile_latent = torch.empty(
+            (1, MAX_NUM_TOKENS, LATENT_SIZE),
+            dtype=torch.bfloat16,
+            device=mailbox.device,
+        )
+        compile_weight = torch.empty(
+            (1, shard_size, LATENT_SIZE),
+            dtype=torch.bfloat16,
+            device=mailbox.device,
+        )
+        cluster_size = math.prod(CLUSTER_SHAPE_MN)
+        max_active_clusters = utils.HardwareInfo().get_max_active_clusters(cluster_size)
+        self.compiled = fused_add_multicast_gemm.compile_kernel(
+            (MAX_NUM_TOKENS, shard_size, LATENT_SIZE, 1),
+            fused_add_multicast_gemm._as_cute(
+                compile_latent,
+                dynamic_m=True,
+            ),
+            fused_add_multicast_gemm._as_cute(compile_weight),
+            self.mailbox_c,
+            fused_add_multicast_gemm._as_cute(shared_shard),
+            HIDDEN_SIZE,
+            shard_size,
+            MMA_TILER_MN,
+            CLUSTER_SHAPE_MN,
+            max_active_clusters,
+            B_PRIME_STAGES,
+        )
+
+    def launch(
+        self,
+        latent: torch.Tensor,
+        weight: torch.Tensor,
+        shared_shard: torch.Tensor,
+    ) -> None:
+        stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream)
+        self.compiled(
+            fused_add_multicast_gemm._as_cute(
+                latent.unsqueeze(0),
+                dynamic_m=True,
+            ),
+            fused_add_multicast_gemm._as_cute(weight.unsqueeze(0)),
+            self.mailbox_c,
+            fused_add_multicast_gemm._as_cute(shared_shard),
+            cutlass.Int64(latent.shape[0]),
+            cutlass.Int64(self.mailbox.data_ptr()),
+            stream,
+        )
+
+
+class SkinnyKernel:
+    def __init__(
+        self,
+        num_tokens: int,
+        shard_size: int,
+        config: fused_add_multicast_skinny_gemm.SkinnyConfig,
+    ) -> None:
+        self.compiled = fused_add_multicast_skinny_gemm.compile_kernel(
+            num_rows=num_tokens,
+            latent_dim=LATENT_SIZE,
+            hidden_dim=HIDDEN_SIZE,
+            shard_dim=shard_size,
+            config=config,
+        )
+
+    def launch(
+        self,
+        latent: torch.Tensor,
+        weight: torch.Tensor,
+        shared_shard: torch.Tensor,
+        mailbox: torch.Tensor,
+    ) -> None:
+        self.compiled(
+            fused_add_multicast_skinny_gemm._as_cute(latent),
+            fused_add_multicast_skinny_gemm._as_cute(weight),
+            fused_add_multicast_skinny_gemm._as_cute(shared_shard),
+            cutlass.Int64(mailbox.data_ptr()),
+            cuda.CUstream(torch.cuda.current_stream().cuda_stream),
+        )
+
+
+def check_up_projection_output(
+    actual: torch.Tensor,
+    latent: torch.Tensor,
+    weight: torch.Tensor,
+    shared_shard: torch.Tensor,
+) -> None:
+    gemm = F.linear(latent.float(), weight.float()).to(torch.bfloat16)
+    expected = (gemm.float() + shared_shard.float()).to(torch.bfloat16)
+    torch.testing.assert_close(actual, expected, atol=8e-2, rtol=3e-2)
+
+
+def make_up_projection_launches(
+    launch: Callable[[torch.Tensor, torch.Tensor, torch.Tensor], None],
+    latent: torch.Tensor,
+    weights: Sequence[torch.Tensor],
+    shared_shard: torch.Tensor,
+) -> list[Callable[[], None]]:
+    return [
+        lambda weight=weight: launch(latent, weight, shared_shard) for weight in weights
+    ]
+
+
+def benchmark_up_projection(args: argparse.Namespace) -> None:
+    if args.tp_size <= 0 or HIDDEN_SIZE % args.tp_size:
+        raise ValueError("TP size must be positive and divide the hidden size")
+    if any(not 1 <= num_tokens <= MAX_NUM_TOKENS for num_tokens in args.num_tokens):
+        raise ValueError("--num-tokens values must be in [1, 16]")
+    if args.cache_multiplier <= 0 or args.max_weights <= 0:
+        raise ValueError("cache multiplier and max weights must be positive")
+    if args.warmup_replays < 0 or args.samples <= 0:
+        raise ValueError("warmup replays must be nonnegative and samples positive")
+
+    torch.accelerator.set_device_index(0)
+    device = torch.device("cuda", 0)
+    if torch.cuda.get_device_capability(device)[0] != 10:
+        raise RuntimeError("Kimi K3 latent-MoE tail requires SM100")
+
+    shard_size = HIDDEN_SIZE // args.tp_size
+    weight_count = rotating_weight_count(
+        shard_size,
+        args.cache_multiplier,
+        args.max_weights,
+    )
+    torch.manual_seed(20260726)
+    weights = [
+        torch.randn(
+            (shard_size, LATENT_SIZE),
+            dtype=torch.bfloat16,
+            device=device,
+        )
+        / LATENT_SIZE**0.5
+        for _ in range(weight_count)
+    ]
+    mailbox = torch.empty(
+        (1, MAX_NUM_TOKENS, HIDDEN_SIZE),
+        dtype=torch.bfloat16,
+        device=device,
+    )
+    shared = torch.randn(
+        (MAX_NUM_TOKENS, HIDDEN_SIZE),
+        dtype=torch.bfloat16,
+        device=device,
+    )
+    shared_shard = shared[:, :shard_size]
+    use_dynamic = args.backend in ("dynamic", "both")
+    use_skinny = args.backend in ("skinny", "both")
+    dynamic_kernel = (
+        DynamicKernel(shard_size, mailbox, shared_shard) if use_dynamic else None
+    )
+
+    results = []
+    for num_tokens in args.num_tokens:
+        latent = torch.randn(
+            (num_tokens, LATENT_SIZE),
+            dtype=torch.bfloat16,
+            device=device,
+        )
+        result: dict[str, Any] = {"num_tokens": num_tokens}
+        if dynamic_kernel is not None:
+            launches = make_up_projection_launches(
+                dynamic_kernel.launch,
+                latent,
+                weights,
+                shared_shard,
+            )
+            graph = capture_up_projection_graph(launches)
+            result["dynamic"] = benchmark_up_projection_graph(
+                graph,
+                operations_per_replay=len(launches),
+                warmup_replays=args.warmup_replays,
+                samples=args.samples,
+            )
+            check_up_projection_output(
+                mailbox[0, :num_tokens, :shard_size],
+                latent,
+                weights[-1],
+                shared_shard[:num_tokens],
+            )
+        if use_skinny:
+            configs = args.skinny_config or [
+                fused_add_multicast_skinny_gemm.config_for_m(
+                    num_tokens,
+                    shard_size,
+                )
+            ]
+            skinny_results = []
+            for config in configs:
+                skinny_kernel = SkinnyKernel(num_tokens, shard_size, config)
+
+                def launch_skinny(
+                    latent: torch.Tensor,
+                    weight: torch.Tensor,
+                    shared_shard: torch.Tensor,
+                    *,
+                    skinny_kernel: SkinnyKernel = skinny_kernel,
+                    num_tokens: int = num_tokens,
+                ) -> None:
+                    skinny_kernel.launch(
+                        latent,
+                        weight,
+                        shared_shard[:num_tokens],
+                        mailbox,
+                    )
+
+                launches = make_up_projection_launches(
+                    launch_skinny,
+                    latent,
+                    weights,
+                    shared_shard,
+                )
+                graph = capture_up_projection_graph(launches)
+                timing = benchmark_up_projection_graph(
+                    graph,
+                    operations_per_replay=len(launches),
+                    warmup_replays=args.warmup_replays,
+                    samples=args.samples,
+                )
+                check_up_projection_output(
+                    mailbox[0, :num_tokens, :shard_size],
+                    latent,
+                    weights[-1],
+                    shared_shard[:num_tokens],
+                )
+                skinny_results.append(
+                    {
+                        "config": asdict(config),
+                        **timing,
+                    }
+                )
+            result["skinny"] = skinny_results
+        results.append(result)
+
+    properties = torch.cuda.get_device_properties(device)
+    report = {
+        "scope": "up-projection",
+        "device": properties.name,
+        "compute_capability": list(torch.cuda.get_device_capability(device)),
+        "tp_size": args.tp_size,
+        "shard_size": shard_size,
+        "weight_count": weight_count,
+        "cache_multiplier": args.cache_multiplier,
+        "warmup_replays": args.warmup_replays,
+        "samples": args.samples,
+        "results": results,
+    }
+    rendered = json.dumps(report, indent=2)
+    print(rendered, flush=True)
+    if args.output is not None:
+        args.output.parent.mkdir(parents=True, exist_ok=True)
+        args.output.write_text(rendered + "\n", encoding="utf-8")
+
+
+def capture_tail_graph(
+    operation: Callable[[], torch.Tensor],
+    cpu_group: dist.ProcessGroup,
+) -> tuple[torch.cuda.CUDAGraph, torch.Tensor]:
+    for _ in range(3):
+        dist.barrier(group=cpu_group)
+        output = operation()
+    torch.accelerator.synchronize()
+
+    dist.barrier(group=cpu_group)
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph):
+        output = operation()
+    torch.accelerator.synchronize()
+    return graph, output
+
+
+def benchmark_tail_graph(
+    graph: torch.cuda.CUDAGraph,
+    *,
+    warmup_replays: int,
+    samples: int,
+    device_group: dist.ProcessGroup,
+    cpu_group: dist.ProcessGroup,
+) -> dict[str, Any]:
+    for _ in range(warmup_replays):
+        graph.replay()
+    torch.accelerator.synchronize()
+
+    dist.barrier(group=cpu_group)
+    starts = [torch.cuda.Event(enable_timing=True) for _ in range(samples + 1)]
+    ends = [torch.cuda.Event(enable_timing=True) for _ in range(samples + 1)]
+    for start, end in zip(starts, ends):
+        start.record()
+        graph.replay()
+        end.record()
+    torch.accelerator.synchronize()
+
+    samples_us = torch.tensor(
+        [start.elapsed_time(end) * 1000.0 for start, end in zip(starts, ends)],
+        dtype=torch.float64,
+        device=torch.accelerator.current_device_index(),
+    )
+    dist.all_reduce(samples_us, op=dist.ReduceOp.MAX, group=device_group)
+    return summarize(samples_us[1:].tolist())
+
+
+def make_inputs(
+    num_tokens: int,
+    rank: int,
+    device: torch.device,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    torch.manual_seed(20260726 + 100 * num_tokens + rank)
+    routed = torch.randn(
+        (num_tokens, LATENT_SIZE),
+        dtype=torch.bfloat16,
+        device=device,
+    ).mul_(0.01)
+    shared = torch.randn(
+        (num_tokens, HIDDEN_SIZE),
+        dtype=torch.bfloat16,
+        device=device,
+    )
+    return routed, shared
+
+
+def make_reference(
+    routed: torch.Tensor,
+    shared: torch.Tensor,
+    rms_weight: torch.Tensor,
+    up_weight: torch.Tensor,
+    device_group: dist.ProcessGroup,
+) -> Callable[[], torch.Tensor]:
+    routed_workspace = torch.empty_like(routed)
+    shared_workspace = torch.empty_like(shared)
+
+    def reference() -> torch.Tensor:
+        routed_workspace.copy_(routed)
+        dist.all_reduce(routed_workspace, group=device_group)
+        normalized = F.rms_norm(
+            routed_workspace,
+            (LATENT_SIZE,),
+            rms_weight,
+            RMS_EPS,
+        )
+        projected = F.linear(normalized, up_weight)
+        shared_workspace.copy_(shared)
+        dist.all_reduce(shared_workspace, group=device_group)
+        return projected.add(shared_workspace)
+
+    return reference
+
+
+def check_fused_output(
+    fused_output: torch.Tensor,
+    reference: Callable[[], torch.Tensor],
+    cpu_group: dist.ProcessGroup,
+) -> None:
+    dist.barrier(group=cpu_group)
+    expected = reference()
+    torch.testing.assert_close(fused_output, expected, atol=8e-2, rtol=3e-2)
+
+
+def benchmark_whole_tail(args: argparse.Namespace) -> None:
+    if any(not 1 <= num_tokens <= 16 for num_tokens in args.num_tokens):
+        raise ValueError("--num-tokens values must be in [1, 16]")
+    if args.warmup_replays < 0 or args.samples <= 0:
+        raise ValueError("warmup replays must be nonnegative and samples positive")
+    if args.skinny_max_num_tokens is not None and any(
+        not 0 <= cutoff <= 8 for cutoff in args.skinny_max_num_tokens
+    ):
+        raise ValueError("--skinny-max-num-tokens must be in [0, 8]")
+    skinny_configs = dict(args.skinny_config or ())
+    if len(skinny_configs) != len(args.skinny_config or ()):
+        raise ValueError("--skinny-config must not repeat an M value")
+    if any(not 1 <= num_tokens <= 8 for num_tokens in skinny_configs):
+        raise ValueError("--skinny-config M values must be in [1, 8]")
+    if not {"RANK", "WORLD_SIZE", "LOCAL_RANK"} <= os.environ.keys():
+        raise RuntimeError("launch this benchmark with torchrun")
+
+    rank = int(os.environ["RANK"])
+    world_size = int(os.environ["WORLD_SIZE"])
+    local_rank = int(os.environ["LOCAL_RANK"])
+    device = torch.device("cuda", local_rank)
+    torch.accelerator.set_device_index(device)
+    init_distributed_environment()
+    if world_size > 8:
+        set_custom_all_reduce(False)
+    initialize_model_parallel(tensor_model_parallel_size=world_size)
+    device_group = get_tp_group().device_group
+    cpu_group = dist.new_group(backend="gloo")
+
+    if torch.cuda.get_device_capability(device)[0] != 10:
+        raise RuntimeError("Kimi K3 latent-MoE tail requires SM100")
+
+    torch.manual_seed(20260726)
+    rms_weight = 1 + 0.1 * torch.randn(
+        LATENT_SIZE,
+        dtype=torch.bfloat16,
+        device=device,
+    )
+    up_weight = (
+        torch.randn(
+            (HIDDEN_SIZE, LATENT_SIZE),
+            dtype=torch.bfloat16,
+            device=device,
+        )
+        / LATENT_SIZE**0.5
+    )
+
+    use_reference = args.backend in ("reference", "both")
+    use_fused = args.backend in ("fused", "both")
+    fused_ops = []
+    if use_fused:
+        production_config_for_m = fused_add_multicast_skinny_gemm.config_for_m
+
+        def config_for_m(
+            num_rows: int,
+            shard_dim: int = 896,
+        ) -> fused_add_multicast_skinny_gemm.SkinnyConfig:
+            config = skinny_configs.get(num_rows)
+            if config is not None:
+                return config
+            return production_config_for_m(num_rows, shard_dim)
+
+        fused_add_multicast_skinny_gemm.config_for_m = config_for_m
+        cutoffs = args.skinny_max_num_tokens or [latent_moe_tail._SKINNY_MAX_NUM_TOKENS]
+        for cutoff in cutoffs:
+            latent_moe_tail._SKINNY_MAX_NUM_TOKENS = cutoff
+            latent_moe_tail.KimiK3LatentMoETailOp._instances.clear()
+            fused_ops.append(
+                (
+                    cutoff,
+                    latent_moe_tail.KimiK3LatentMoETailOp.initialize(
+                        hidden_size=HIDDEN_SIZE,
+                        latent_size=LATENT_SIZE,
+                        dtype=torch.bfloat16,
+                        device=device,
+                        rms_eps=RMS_EPS,
+                    ),
+                )
+            )
+        cutedsl_warmup()
+
+    results = []
+    for num_tokens in args.num_tokens:
+        routed, shared = make_inputs(num_tokens, rank, device)
+        reference = make_reference(
+            routed,
+            shared,
+            rms_weight,
+            up_weight,
+            device_group,
+        )
+        result: dict[str, Any] = {"num_tokens": num_tokens}
+        if use_reference:
+            reference_graph, _ = capture_tail_graph(reference, cpu_group)
+            result["reference"] = benchmark_tail_graph(
+                reference_graph,
+                warmup_replays=args.warmup_replays,
+                samples=args.samples,
+                device_group=device_group,
+                cpu_group=cpu_group,
+            )
+        for cutoff, fused_op in fused_ops:
+
+            def fused(
+                routed: torch.Tensor = routed,
+                shared: torch.Tensor = shared,
+                fused_op: latent_moe_tail.KimiK3LatentMoETailOp = fused_op,
+            ) -> torch.Tensor:
+                return fused_op(routed, shared, rms_weight, up_weight)
+
+            fused_graph, fused_output = capture_tail_graph(fused, cpu_group)
+            fused_key = "fused" if len(fused_ops) == 1 else f"fused_skinny_max_{cutoff}"
+            result[fused_key] = benchmark_tail_graph(
+                fused_graph,
+                warmup_replays=args.warmup_replays,
+                samples=args.samples,
+                device_group=device_group,
+                cpu_group=cpu_group,
+            )
+            check_fused_output(fused_output, reference, cpu_group)
+            if "reference" in result:
+                speedup = (
+                    result["reference"]["median_us"] / result[fused_key]["median_us"]
+                )
+                if len(fused_ops) == 1:
+                    result["speedup"] = speedup
+                else:
+                    result[f"{fused_key}_speedup"] = speedup
+        results.append(result)
+
+    properties = torch.cuda.get_device_properties(device)
+    report = {
+        "scope": "whole-tail",
+        "device": properties.name,
+        "compute_capability": list(torch.cuda.get_device_capability(device)),
+        "world_size": world_size,
+        "torch_version": torch.__version__,
+        "cuda_version": torch.version.cuda,
+        "warmup_replays": args.warmup_replays,
+        "samples": args.samples,
+        "skinny_max_num_tokens": [cutoff for cutoff, _ in fused_ops],
+        "skinny_configs": {
+            str(num_tokens): asdict(config)
+            for num_tokens, config in skinny_configs.items()
+        },
+        "timing_scope": {
+            "reference": (
+                "two input copies, two AllReduces, RMSNorm, full replicated "
+                "up-projection GEMM, and final add"
+            ),
+            "fused": (
+                "routed AllReduce/RMSNorm plus shared ReduceScatter, sharded "
+                "up-projection/multicast, and Lamport copy"
+            ),
+        },
+        "results": results,
+    }
+    if rank == 0:
+        rendered = json.dumps(report, indent=2)
+        print(rendered, flush=True)
+        if args.output is not None:
+            args.output.parent.mkdir(parents=True, exist_ok=True)
+            args.output.write_text(rendered + "\n", encoding="utf-8")
+
+    dist.barrier(group=cpu_group)
+
+
+def main() -> None:
+    args = parse_args()
+    if args.scope == "up-projection":
+        benchmark_up_projection(args)
+        return
+
+    from vllm.config import VllmConfig, set_current_vllm_config
+
+    with set_current_vllm_config(VllmConfig()):
+        benchmark_whole_tail(args)
+
+
+if __name__ == "__main__":
+    main()

benchmarks/kernels/benchmark_kimi_k3_sp_collectives.py

@@ -0,0 +1,239 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import argparse
+import json
+import os
+import statistics
+from collections.abc import Callable
+
+import torch
+import torch.distributed as dist
+
+import vllm._custom_ops as ops
+from vllm.distributed.device_communicators.custom_all_reduce import CustomAllreduce
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--tokens", type=int, nargs="+", default=[8, 32, 128, 1024])
+    parser.add_argument("--hidden-size", type=int, default=7168)
+    parser.add_argument("--graph-repeats", type=int, default=20)
+    parser.add_argument("--warmup-replays", type=int, default=5)
+    parser.add_argument("--samples", type=int, default=15)
+    return parser.parse_args()
+
+
+def capture_graph(op: Callable[[], None], repeats: int) -> torch.cuda.CUDAGraph:
+    stream = torch.cuda.Stream()
+    stream.wait_stream(torch.cuda.current_stream())
+    with torch.cuda.stream(stream):
+        for _ in range(3):
+            op()
+    stream.synchronize()
+
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph, stream=stream):
+        for _ in range(repeats):
+            op()
+    torch.cuda.current_stream().wait_stream(stream)
+    return graph
+
+
+def max_rank_graph_time(
+    graph: torch.cuda.CUDAGraph,
+    repeats: int,
+    warmup_replays: int,
+    samples: int,
+    device_group: dist.ProcessGroup,
+    cpu_group: dist.ProcessGroup,
+) -> float:
+    for _ in range(warmup_replays):
+        graph.replay()
+    torch.accelerator.synchronize()
+
+    timings = []
+    start = torch.cuda.Event(enable_timing=True)
+    end = torch.cuda.Event(enable_timing=True)
+    for _ in range(samples):
+        dist.barrier(group=cpu_group)
+        start.record()
+        graph.replay()
+        end.record()
+        end.synchronize()
+        elapsed = torch.tensor(
+            start.elapsed_time(end) / repeats,
+            dtype=torch.float64,
+            device=torch.accelerator.current_device_index(),
+        )
+        dist.all_reduce(elapsed, op=dist.ReduceOp.MAX, group=device_group)
+        timings.append(elapsed.item())
+    return statistics.median(timings)
+
+
+def check_outputs(
+    comm: CustomAllreduce,
+    local: torch.Tensor,
+    reduce_input: torch.Tensor,
+    device_group: dist.ProcessGroup,
+) -> None:
+    expected_gather = torch.empty(
+        (local.shape[0] * dist.get_world_size(), local.shape[1]),
+        dtype=local.dtype,
+        device=local.device,
+    )
+    dist.all_gather_into_tensor(expected_gather, local, group=device_group)
+    gathered = comm.custom_all_gather(local)
+    assert gathered is not None
+    torch.testing.assert_close(gathered, expected_gather)
+
+    expected_scatter = torch.empty_like(local)
+    dist.reduce_scatter_tensor(
+        expected_scatter,
+        reduce_input.clone(),
+        group=device_group,
+    )
+    scattered = comm.custom_reduce_scatter(reduce_input)
+    assert scattered is not None
+    torch.testing.assert_close(scattered, expected_scatter)
+
+
+def benchmark_shape(
+    comm: CustomAllreduce,
+    global_tokens: int,
+    hidden_size: int,
+    graph_repeats: int,
+    warmup_replays: int,
+    samples: int,
+    device_group: dist.ProcessGroup,
+    cpu_group: dist.ProcessGroup,
+) -> dict[str, float | int]:
+    world_size = dist.get_world_size()
+    rank = dist.get_rank()
+    padded_tokens = (global_tokens + world_size - 1) // world_size * world_size
+    local_tokens = padded_tokens // world_size
+    local = torch.full(
+        (local_tokens, hidden_size),
+        rank + 1,
+        dtype=torch.bfloat16,
+        device=torch.accelerator.current_device_index(),
+    )
+    reduce_input = torch.full(
+        (padded_tokens, hidden_size),
+        rank + 1,
+        dtype=torch.bfloat16,
+        device=local.device,
+    )
+    check_outputs(comm, local, reduce_input, device_group)
+
+    custom_gather_out = torch.empty(
+        (padded_tokens, hidden_size),
+        dtype=local.dtype,
+        device=local.device,
+    )
+    custom_scatter_out = torch.empty_like(local)
+    nccl_gather_out = torch.empty_like(custom_gather_out)
+    nccl_scatter_out = torch.empty_like(local)
+
+    def custom_ag() -> None:
+        ops.mnnvl_lamport_all_gather(
+            comm._ptr,
+            local,
+            custom_gather_out,
+            comm.mnnvl_lamport_ag_local_ptr,
+            comm.mnnvl_lamport_ag_multicast_ptr,
+            comm.mnnvl_lamport_ag_epoch_ptr,
+            comm.mnnvl_buffer_size,
+        )
+
+    def custom_rs() -> None:
+        ops.mnnvl_lamport_reduce_scatter(
+            comm._ptr,
+            reduce_input,
+            custom_scatter_out,
+            comm.mnnvl_lamport_rs_local_ptr,
+            comm.mnnvl_lamport_rs_epoch_ptr,
+            comm.mnnvl_buffer_size,
+        )
+
+    def nccl_ag() -> None:
+        dist.all_gather_into_tensor(nccl_gather_out, local, group=device_group)
+
+    def nccl_rs() -> None:
+        dist.reduce_scatter_tensor(
+            nccl_scatter_out,
+            reduce_input,
+            group=device_group,
+        )
+
+    graphs = {
+        "custom_ag_us": capture_graph(custom_ag, graph_repeats),
+        "nccl_ag_us": capture_graph(nccl_ag, graph_repeats),
+        "custom_rs_us": capture_graph(custom_rs, graph_repeats),
+        "nccl_rs_us": capture_graph(nccl_rs, graph_repeats),
+    }
+    times = {
+        name: max_rank_graph_time(
+            graph,
+            graph_repeats,
+            warmup_replays,
+            samples,
+            device_group,
+            cpu_group,
+        )
+        * 1000
+        for name, graph in graphs.items()
+    }
+    torch.testing.assert_close(custom_gather_out, nccl_gather_out)
+    torch.testing.assert_close(custom_scatter_out, nccl_scatter_out)
+    return {
+        "global_tokens": global_tokens,
+        "padded_tokens": padded_tokens,
+        "local_bytes": local.nbytes,
+        "full_bytes": reduce_input.nbytes,
+        **times,
+        "ag_speedup": times["nccl_ag_us"] / times["custom_ag_us"],
+        "rs_speedup": times["nccl_rs_us"] / times["custom_rs_us"],
+    }
+
+
+def main() -> None:
+    args = parse_args()
+    local_rank = int(os.environ["LOCAL_RANK"])
+    torch.accelerator.set_device_index(local_rank)
+    dist.init_process_group("nccl")
+    device_group = dist.group.WORLD
+    cpu_group = dist.new_group(backend="gloo")
+
+    comm = CustomAllreduce(
+        group=cpu_group,
+        device=torch.device("cuda", local_rank),
+    )
+    assert not comm.disabled
+    assert comm.world_size == 16
+    assert comm.mnnvl_only
+    assert comm.mnnvl_multicast_ptr
+
+    results = [
+        benchmark_shape(
+            comm,
+            tokens,
+            args.hidden_size,
+            args.graph_repeats,
+            args.warmup_replays,
+            args.samples,
+            device_group,
+            cpu_group,
+        )
+        for tokens in args.tokens
+    ]
+    if dist.get_rank() == 0:
+        print(json.dumps(results, indent=2), flush=True)
+
+    comm.close()
+    dist.destroy_process_group(cpu_group)
+    dist.destroy_process_group()
+
+
+if __name__ == "__main__":
+    main()

cmake/external_projects/flashkda.cmake

@@ -0,0 +1,74 @@
+include(FetchContent)
+
+if(DEFINED ENV{FLASH_KDA_SRC_DIR})
+  set(FLASH_KDA_SRC_DIR $ENV{FLASH_KDA_SRC_DIR})
+endif()
+
+if(FLASH_KDA_SRC_DIR)
+  FetchContent_Declare(
+    flashkda
+    SOURCE_DIR ${FLASH_KDA_SRC_DIR}
+  )
+else()
+  FetchContent_Declare(
+    flashkda
+    GIT_REPOSITORY https://github.com/vllm-project/FlashKDA.git
+    GIT_TAG a3e42bbbece3bb38f7c426b880315294a336e82f
+    GIT_PROGRESS TRUE
+    GIT_SUBMODULES cutlass
+  )
+endif()
+
+FetchContent_MakeAvailable(flashkda)
+message(STATUS "FlashKDA is available at ${flashkda_SOURCE_DIR}")
+
+set(FLASH_KDA_SUPPORT_ARCHS)
+if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.0)
+  list(APPEND FLASH_KDA_SUPPORT_ARCHS "9.0a")
+endif()
+if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
+  list(APPEND FLASH_KDA_SUPPORT_ARCHS "10.0f" "12.0f")
+elseif(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.9)
+  list(APPEND FLASH_KDA_SUPPORT_ARCHS "10.0a" "10.3a" "12.0a")
+endif()
+
+cuda_archs_loose_intersection(
+  FLASH_KDA_ARCHS "${FLASH_KDA_SUPPORT_ARCHS}" "${CUDA_ARCHS}")
+
+if(FLASH_KDA_ARCHS)
+  message(STATUS "FlashKDA CUDA architectures: ${FLASH_KDA_ARCHS}")
+
+  set(FLASH_KDA_SOURCES
+    csrc/flashkda_registration.cpp
+    ${flashkda_SOURCE_DIR}/csrc/flash_kda.cpp
+    ${flashkda_SOURCE_DIR}/csrc/smxx/fwd_launch.cu)
+  set(FLASH_KDA_INCLUDES
+    ${flashkda_SOURCE_DIR}/csrc
+    ${flashkda_SOURCE_DIR}/cutlass/include
+    ${flashkda_SOURCE_DIR}/cutlass/examples/common
+    ${flashkda_SOURCE_DIR}/cutlass/tools/util/include)
+
+  set_gencode_flags_for_srcs(
+    SRCS "${FLASH_KDA_SOURCES}"
+    CUDA_ARCHS "${FLASH_KDA_ARCHS}")
+
+  define_extension_target(
+    _flashkda_C
+    DESTINATION vllm
+    LANGUAGE ${VLLM_GPU_LANG}
+    SOURCES ${FLASH_KDA_SOURCES}
+    COMPILE_FLAGS ${VLLM_GPU_FLAGS}
+    ARCHITECTURES ${VLLM_GPU_ARCHES}
+    INCLUDE_DIRECTORIES ${FLASH_KDA_INCLUDES}
+    USE_SABI 3
+    WITH_SOABI)
+
+  target_compile_options(_flashkda_C PRIVATE
+    $<$<COMPILE_LANGUAGE:CUDA>:-UPy_LIMITED_API --expt-relaxed-constexpr --expt-extended-lambda --use_fast_math -O3>
+    $<$<COMPILE_LANGUAGE:CXX>:-UPy_LIMITED_API>)
+else()
+  message(STATUS
+    "FlashKDA will not compile: CUDA >=12.0 and a supported architecture "
+    "(SM90, SM10x, or SM12x) are required")
+  add_custom_target(_flashkda_C)
+endif()

csrc/custom_all_gather_reduce_scatter.cuh

@@ -0,0 +1,327 @@
+#pragma once
+
+#include "custom_collective_common.cuh"
+
+namespace vllm {
+
+constexpr int kMnnvlLamportAgThreads = 128;
+constexpr int kMnnvlLamportRsThreads = 256;
+constexpr int kMnnvlLamportConcurrentPollMaxPacks = 8192;
+
+using CopyPack = array_t<uint64_t, 2>;
+
+template <int ngpus>
+__global__ void __launch_bounds__(512, 1)
+    cross_device_all_gather(RankData* _dp, RankSignals sg, Signal* self_sg,
+                            CopyPack* __restrict__ result, int rank,
+                            int size_per_rank) {
+  auto dp = *_dp;
+  int tid = blockIdx.x * blockDim.x + threadIdx.x;
+  int stride = gridDim.x * blockDim.x;
+  barrier_at_start<ngpus>(sg, self_sg, rank);
+#pragma unroll
+  for (int src_rank = 0; src_rank < ngpus; ++src_rank) {
+    auto src = reinterpret_cast<const CopyPack*>(dp.ptrs[src_rank]);
+    auto dst = result + src_rank * size_per_rank;
+    for (int idx = tid; idx < size_per_rank; idx += stride) {
+      dst[idx] = src[idx];
+    }
+  }
+  barrier_at_end<ngpus, true>(sg, self_sg, rank);
+}
+
+template <typename T, int ngpus>
+__global__ void __launch_bounds__(512, 1)
+    cross_device_reduce_scatter(RankData* _dp, RankSignals sg, Signal* self_sg,
+                                T* __restrict__ result, int rank,
+                                int size_per_rank) {
+  using P = typename packed_t<T>::P;
+  using A = typename packed_t<T>::A;
+  auto dp = *_dp;
+  auto offset = rank * size_per_rank;
+  barrier_at_start<ngpus>(sg, self_sg, rank);
+  for (int idx = blockIdx.x * blockDim.x + threadIdx.x; idx < size_per_rank;
+       idx += gridDim.x * blockDim.x) {
+    reinterpret_cast<P*>(result)[idx] =
+        packed_reduce<P, ngpus, A>((const P**)&dp.ptrs[0], offset + idx);
+  }
+  barrier_at_end<ngpus, true>(sg, self_sg, rank);
+}
+
+template <typename P>
+union LamportPack {
+  P packed;
+  uint32_t words[sizeof(P) / sizeof(uint32_t)];
+};
+
+template <typename P>
+DINLINE LamportPack<P> load_lamport_pack(const P* ptr) {
+  static_assert(sizeof(P) == 16);
+  LamportPack<P> value;
+#if !defined(USE_ROCM)
+  asm volatile("ld.volatile.global.v4.u32 {%0, %1, %2, %3}, [%4];"
+               : "=r"(value.words[0]), "=r"(value.words[1]),
+                 "=r"(value.words[2]), "=r"(value.words[3])
+               : "l"(ptr)
+               : "memory");
+#else
+  const volatile uint32_t* src =
+      reinterpret_cast<const volatile uint32_t*>(ptr);
+  #pragma unroll
+  for (int i = 0; i < sizeof(P) / sizeof(uint32_t); ++i) {
+    value.words[i] = src[i];
+  }
+#endif
+  return value;
+}
+
+template <typename P>
+DINLINE bool is_lamport_dirty(const LamportPack<P>& value) {
+#pragma unroll
+  for (int i = 0; i < sizeof(P) / sizeof(uint32_t); ++i) {
+    if (value.words[i] == 0x80000000U) return true;
+  }
+  return false;
+}
+
+template <typename P>
+DINLINE P lamport_sentinel() {
+  LamportPack<P> value;
+#pragma unroll
+  for (int i = 0; i < sizeof(P) / sizeof(uint32_t); ++i) {
+    value.words[i] = 0x80000000U;
+  }
+  return value.packed;
+}
+
+template <typename P>
+DINLINE P sanitize_lamport_payload(P packed) {
+  LamportPack<P> value{.packed = packed};
+#pragma unroll
+  for (int i = 0; i < sizeof(P) / sizeof(uint32_t); ++i) {
+    if (value.words[i] == 0x80000000U) value.words[i] = 0;
+  }
+  return value.packed;
+}
+
+template <typename P>
+DINLINE P wait_lamport_payload(const P* ptr) {
+  auto value = load_lamport_pack(ptr);
+  while (is_lamport_dirty(value)) value = load_lamport_pack(ptr);
+  return value.packed;
+}
+
+template <typename P, int ngpus>
+DINLINE void wait_lamport_payloads(const P* base, int rank, int rank_stride,
+                                   P local_value, P (&values)[ngpus]) {
+  bool ready[ngpus];
+#pragma unroll
+  for (int src_rank = 0; src_rank < ngpus; ++src_rank) {
+    ready[src_rank] = src_rank == rank;
+    if (src_rank == rank) values[src_rank] = local_value;
+  }
+
+  int remaining = ngpus - 1;
+  while (remaining != 0) {
+#pragma unroll
+    for (int src_rank = 0; src_rank < ngpus; ++src_rank) {
+      if (!ready[src_rank]) {
+        auto value = load_lamport_pack(base + src_rank * rank_stride);
+        if (!is_lamport_dirty(value)) {
+          values[src_rank] = value.packed;
+          ready[src_rank] = true;
+          --remaining;
+        }
+      }
+    }
+  }
+}
+
+template <typename P, typename A, int ngpus>
+DINLINE P reduce_lamport_payloads(const P* current_local, const P* packed_input,
+                                  int rank, int size_per_rank, int idx) {
+  P source_zero =
+      rank == 0 ? packed_input[idx] : wait_lamport_payload(current_local + idx);
+  A tmp = upcast(source_zero);
+#pragma unroll
+  for (int src_rank = 1; src_rank < ngpus; ++src_rank) {
+    P value = src_rank == rank
+                  ? packed_input[rank * size_per_rank + idx]
+                  : wait_lamport_payload(current_local +
+                                         src_rank * size_per_rank + idx);
+    packed_assign_add(tmp, upcast(value));
+  }
+  return sanitize_lamport_payload(downcast<P>(tmp));
+}
+
+DINLINE void lamport_cta_arrive(uint32_t* counter) {
+#if !defined(USE_ROCM)
+  if (threadIdx.x < 32) {
+    asm volatile("barrier.cta.sync 1, %0;" : : "r"(blockDim.x) : "memory");
+    if (threadIdx.x == 0) {
+  #if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 1000
+      asm volatile("red.async.release.global.gpu.add.u32 [%0], 1;"
+                   :
+                   : "l"(counter)
+                   : "memory");
+  #elif defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 700
+      asm volatile("red.release.global.gpu.add.u32 [%0], 1;"
+                   :
+                   : "l"(counter)
+                   : "memory");
+  #else
+      atomicAdd(counter, 1);
+  #endif
+    }
+  } else {
+    asm volatile("barrier.cta.arrive 1, %0;" : : "r"(blockDim.x) : "memory");
+  }
+#else
+  __syncthreads();
+  if (threadIdx.x == 0) atomicAdd(counter, 1);
+#endif
+}
+
+template <typename T, int ngpus>
+__global__ void __launch_bounds__(kMnnvlLamportAgThreads, 1)
+    mnnvl_lamport_all_gather(RankData* _dp, const T* __restrict__ input,
+                             T* __restrict__ result,
+                             T* __restrict__ multicast_buffer,
+                             uint32_t* __restrict__ epochs, int rank,
+                             int size_per_rank, int stage_size) {
+  using P = typename packed_t<T>::P;
+#if !defined(USE_ROCM) && CUDA_VERSION >= 12000 && defined(__CUDA_ARCH__) && \
+    (__CUDA_ARCH__ >= 900)
+  cudaGridDependencySynchronize();
+#endif
+  auto dp = *_dp;
+  int tid = blockIdx.x * blockDim.x + threadIdx.x;
+  int stride = gridDim.x * blockDim.x;
+  uint32_t epoch = epochs[0];
+  int current_stage = epoch % 3;
+  int dirty_stage = (epoch + 1) % 3;
+  int dirty_size = epochs[2 + dirty_stage];
+  auto local_buffer = reinterpret_cast<P*>(const_cast<void*>(dp.ptrs[rank]));
+  auto current_local = local_buffer + current_stage * stage_size;
+  auto dirty_local = local_buffer + dirty_stage * stage_size;
+  auto current_multicast =
+      reinterpret_cast<P*>(multicast_buffer) + current_stage * stage_size;
+  auto packed_input = reinterpret_cast<const P*>(input);
+  auto packed_result = reinterpret_cast<P*>(result);
+
+  int total_size = size_per_rank * ngpus;
+  P local_value;
+  if (tid < size_per_rank) {
+    local_value = packed_input[tid];
+    current_multicast[rank * size_per_rank + tid] =
+        sanitize_lamport_payload(local_value);
+  }
+#if !defined(USE_ROCM) && CUDA_VERSION >= 12000 && defined(__CUDA_ARCH__) && \
+    (__CUDA_ARCH__ >= 900)
+  cudaTriggerProgrammaticLaunchCompletion();
+#endif
+
+  lamport_cta_arrive(&epochs[1]);
+
+  for (int idx = tid; idx < dirty_size; idx += stride) {
+    dirty_local[idx] = lamport_sentinel<P>();
+  }
+
+  if (tid < size_per_rank) {
+#pragma unroll
+    for (int src_rank = 0; src_rank < ngpus; ++src_rank) {
+      int output_idx = src_rank * size_per_rank + tid;
+      P value = src_rank == rank
+                    ? local_value
+                    : wait_lamport_payload(current_local + output_idx);
+      packed_result[output_idx] = value;
+    }
+  }
+
+  if (tid == 0) {
+    while (*reinterpret_cast<volatile uint32_t*>(&epochs[1]) < gridDim.x);
+    epochs[2 + current_stage] = total_size;
+    epochs[0] = epoch + 1;
+    epochs[1] = 0;
+  }
+}
+
+template <typename T, int ngpus>
+__global__ void __launch_bounds__(kMnnvlLamportRsThreads, 1)
+    mnnvl_lamport_reduce_scatter_kernel(RankData* _dp,
+                                        const T* __restrict__ input,
+                                        T* __restrict__ result,
+                                        uint32_t* __restrict__ epochs, int rank,
+                                        int size_per_rank, int stage_size) {
+  using P = typename packed_t<T>::P;
+  using A = typename packed_t<T>::A;
+#if !defined(USE_ROCM) && CUDA_VERSION >= 12000 && defined(__CUDA_ARCH__) && \
+    (__CUDA_ARCH__ >= 900)
+  cudaGridDependencySynchronize();
+#endif
+  auto dp = *_dp;
+  int dst_rank = blockIdx.x % ngpus;
+  int tile = blockIdx.x / ngpus;
+  int idx = tile * blockDim.x + threadIdx.x;
+  int tid = blockIdx.x * blockDim.x + threadIdx.x;
+  int stride = gridDim.x * blockDim.x;
+  uint32_t epoch = epochs[0];
+  int current_stage = epoch % 3;
+  int dirty_stage = (epoch + 1) % 3;
+  int dirty_size = epochs[2 + dirty_stage];
+  auto local_buffer = reinterpret_cast<P*>(const_cast<void*>(dp.ptrs[rank]));
+  auto current_local = local_buffer + current_stage * stage_size;
+  auto dirty_local = local_buffer + dirty_stage * stage_size;
+  auto packed_input = reinterpret_cast<const P*>(input);
+
+  if (idx < size_per_rank && dst_rank != rank) {
+    auto dst = reinterpret_cast<P*>(const_cast<void*>(dp.ptrs[dst_rank])) +
+               current_stage * stage_size + rank * size_per_rank;
+    auto src = packed_input + dst_rank * size_per_rank;
+    dst[idx] = sanitize_lamport_payload(src[idx]);
+  }
+#if !defined(USE_ROCM) && CUDA_VERSION >= 12000 && defined(__CUDA_ARCH__) && \
+    (__CUDA_ARCH__ >= 900)
+  cudaTriggerProgrammaticLaunchCompletion();
+#endif
+
+  lamport_cta_arrive(&epochs[1]);
+
+  for (int idx = tid; idx < dirty_size; idx += stride) {
+    dirty_local[idx] = lamport_sentinel<P>();
+  }
+
+  if (idx < size_per_rank && dst_rank == rank) {
+    if constexpr (ngpus == 4) {
+      if (size_per_rank > kMnnvlLamportConcurrentPollMaxPacks) {
+        reinterpret_cast<P*>(result)[idx] =
+            reduce_lamport_payloads<P, A, ngpus>(current_local, packed_input,
+                                                 rank, size_per_rank, idx);
+      } else {
+        P values[ngpus];
+        wait_lamport_payloads<P, ngpus>(
+            current_local + idx, rank, size_per_rank,
+            packed_input[rank * size_per_rank + idx], values);
+        A tmp = upcast(values[0]);
+#pragma unroll
+        for (int src_rank = 1; src_rank < ngpus; ++src_rank) {
+          packed_assign_add(tmp, upcast(values[src_rank]));
+        }
+        reinterpret_cast<P*>(result)[idx] =
+            sanitize_lamport_payload(downcast<P>(tmp));
+      }
+    } else {
+      reinterpret_cast<P*>(result)[idx] = reduce_lamport_payloads<P, A, ngpus>(
+          current_local, packed_input, rank, size_per_rank, idx);
+    }
+  }
+
+  if (tid == 0) {
+    while (*reinterpret_cast<volatile uint32_t*>(&epochs[1]) < gridDim.x);
+    epochs[2 + current_stage] = size_per_rank * ngpus;
+    epochs[0] = epoch + 1;
+    epochs[1] = 0;
+  }
+}
+
+}  // namespace vllm

csrc/custom_all_reduce.cuh

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

csrc/custom_collective_common.cuh

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

csrc/flashkda_registration.cpp

@@ -0,0 +1,17 @@
+#include "core/registration.h"
+#include "flash_kda.h"
+
+TORCH_LIBRARY(_flashkda_C, m) {
+  m.def("get_workspace_size(int T_total, int H, int N=1) -> int",
+        &get_workspace_size);
+  m.def(
+      "fwd(Tensor q, Tensor k, Tensor v, Tensor g, Tensor beta, float scale, "
+      "Tensor(a!) out, Tensor workspace, Tensor A_log, Tensor dt_bias, "
+      "float lower_bound, "
+      "Tensor? initial_state=None, Tensor(b!)? final_state=None, "
+      "Tensor? cu_seqlens=None) -> ()");
+}
+
+TORCH_LIBRARY_IMPL(_flashkda_C, CUDA, m) { m.impl("fwd", &fwd); }
+
+REGISTER_EXTENSION(_flashkda_C)

csrc/libtorch_stable/activation_kernels.cu

@@ -464,6 +464,66 @@ __global__ void swigluoai_and_mul_kernel(
   }
 }
 
+// SITU (Kimi SituGLU) gated activation. Non-interleaved layout:
+// input = [gate(d), up(d)] per token.
+//   gate_out = beta * tanh(gate / beta) * sigmoid(gate)
+//   up_out   = (linear_beta > 0) ? linear_beta * tanh(up / linear_beta) : up
+//   out      = gate_out * up_out
+// Compute is done in fp32 and written straight to `out` -- no intermediate
+// tensors and no full-tensor fp32 upcast (the pure-torch forward_native
+// allocated ~8 fp32 temporaries per call, which blows up MoE profiling).
+template <typename scalar_t>
+__global__ void situ_and_mul_kernel(
+    scalar_t* __restrict__ out,          // [..., d]
+    const scalar_t* __restrict__ input,  // [..., 2, d]
+    const int d, const float beta, const float linear_beta) {
+  const int64_t row = blockIdx.x;
+  const scalar_t* gate_ptr = input + row * 2 * d;
+  const scalar_t* up_ptr = gate_ptr + d;
+  scalar_t* out_ptr = out + row * d;
+  const bool clamp_up = linear_beta > 0.0f;
+  const float inv_beta = 1.0f / beta;
+  const float inv_linear_beta = clamp_up ? 1.0f / linear_beta : 0.0f;
+  for (int64_t idx = threadIdx.x; idx < d; idx += blockDim.x) {
+    const float g = (float)VLLM_LDG(&gate_ptr[idx]);
+    const float u = (float)VLLM_LDG(&up_ptr[idx]);
+    const float gate_out = beta * tanhf(g * inv_beta) / (1.0f + expf(-g));
+    const float up_out =
+        clamp_up ? linear_beta * tanhf(u * inv_linear_beta) : u;
+    out_ptr[idx] = (scalar_t)(gate_out * up_out);
+  }
+}
+
+template <typename scalar_t>
+__global__ void masked_situ_and_mul_kernel(
+    scalar_t* __restrict__ out, const scalar_t* __restrict__ input,
+    const int* __restrict__ expert_num_tokens, const int max_num_tokens,
+    const int d, const float beta, const float linear_beta) {
+  const int expert = blockIdx.y;
+  const int num_tokens = expert_num_tokens[expert];
+  const int idx = blockIdx.x * blockDim.x + threadIdx.x;
+  if (idx >= d || num_tokens == 0) {
+    return;
+  }
+
+  const bool clamp_up = linear_beta > 0.0f;
+  const float inv_beta = 1.0f / beta;
+  const float inv_linear_beta = clamp_up ? 1.0f / linear_beta : 0.0f;
+  const int64_t expert_row = static_cast<int64_t>(expert) * max_num_tokens;
+  for (int token = 0; token < num_tokens; ++token) {
+    const int64_t row = expert_row + token;
+    const scalar_t* gate_ptr = input + row * 2 * d;
+    const scalar_t* up_ptr = gate_ptr + d;
+    scalar_t* out_ptr = out + row * d;
+    const float g = (float)VLLM_LDG(&gate_ptr[idx]);
+    const float u = (float)VLLM_LDG(&up_ptr[idx]);
+    const float gate_out = beta * tanhf(g * inv_beta) / (1.0f + expf(-g));
+    const float up_out =
+        clamp_up ? linear_beta * tanhf(u * inv_linear_beta) : u;
+    out_ptr[idx] = (scalar_t)(gate_out * up_out);
+  }
+}
+
 }  // namespace vllm
 
 #define LAUNCH_ACTIVATION_GATE_KERNEL_WITH_PARAM(KERNEL, PACKED_KERNEL, PARAM) \
@@ -553,6 +613,54 @@ void swigluoai_and_mul(torch::stable::Tensor& out,    // [..., d]
                        double alpha, double limit) {
   LAUNCH_SIGLUOAI_AND_MUL(vllm::swigluoai_and_mul, alpha, limit);
 }
+
+// Kimi SITU gated activation. `linear_beta <= 0` means "unset" (up passed
+// through), matching SituAndMul(linear_beta=None) on the Python side.
+void situ_and_mul(torch::stable::Tensor& out,    // [..., d]
+                  torch::stable::Tensor& input,  // [..., 2 * d]
+                  double beta, double linear_beta) {
+  int d = input.size(-1) / 2;
+  int64_t num_tokens = input.numel() / input.size(-1);
+  if (num_tokens == 0) {
+    return;
+  }
+  dim3 grid(num_tokens);
+  dim3 block(std::min(d, 1024));
+  const torch::stable::accelerator::DeviceGuard device_guard(
+      input.get_device_index());
+  const cudaStream_t stream = get_current_cuda_stream();
+  VLLM_STABLE_DISPATCH_FLOATING_TYPES(
+      input.scalar_type(), "situ_and_mul_kernel", [&] {
+        vllm::situ_and_mul_kernel<scalar_t><<<grid, block, 0, stream>>>(
+            out.mutable_data_ptr<scalar_t>(), input.const_data_ptr<scalar_t>(),
+            d, (float)beta, (float)linear_beta);
+      });
+}
+
+void masked_situ_and_mul(torch::stable::Tensor& out,    // [E, T, d]
+                         torch::stable::Tensor& input,  // [E, T, 2 * d]
+                         const torch::stable::Tensor& expert_num_tokens,
+                         double beta, double linear_beta) {
+  int num_experts = input.size(0);
+  int max_num_tokens = input.size(1);
+  int d = input.size(2) / 2;
+  if (num_experts == 0 || max_num_tokens == 0) {
+    return;
+  }
+  constexpr int block_size = 256;
+  dim3 grid((d + block_size - 1) / block_size, num_experts);
+  dim3 block(block_size);
+  const torch::stable::accelerator::DeviceGuard device_guard(
+      input.get_device_index());
+  const cudaStream_t stream = get_current_cuda_stream();
+  VLLM_STABLE_DISPATCH_FLOATING_TYPES(
+      input.scalar_type(), "masked_situ_and_mul_kernel", [&] {
+        vllm::masked_situ_and_mul_kernel<scalar_t><<<grid, block, 0, stream>>>(
+            out.mutable_data_ptr<scalar_t>(), input.const_data_ptr<scalar_t>(),
+            expert_num_tokens.const_data_ptr<int>(), max_num_tokens, d,
+            (float)beta, (float)linear_beta);
+      });
+}
 namespace vllm {
 
 // Element-wise activation kernel template.

csrc/libtorch_stable/attention/merge_attn_states.cu

@@ -21,7 +21,10 @@ __global__ void merge_attn_states_kernel(
     const float* prefix_lse, const scalar_t* suffix_output,
     const float* suffix_lse, const uint num_tokens, const uint num_heads,
     const uint head_size, const uint prefix_head_stride,
-    const uint output_head_stride, const uint prefix_num_tokens,
+    const uint output_head_stride, const uint prefix_lse_head_stride,
+    const uint prefix_lse_token_stride, const uint suffix_lse_head_stride,
+    const uint suffix_lse_token_stride, const uint output_lse_head_stride,
+    const uint output_lse_token_stride, const uint prefix_num_tokens,
     const float* output_scale) {
   // Inputs always load 128-bit packs (pack_size elements of scalar_t).
   // Outputs store pack_size elements of output_t, which is smaller for FP8.
@@ -84,15 +87,19 @@ __global__ void merge_attn_states_kernel(
       }
     }
     if (output_lse != nullptr && pack_idx == 0) {
-      float s_lse = suffix_lse[head_idx * num_tokens + token_idx];
-      output_lse[head_idx * num_tokens + token_idx] = s_lse;
+      float s_lse = suffix_lse[head_idx * suffix_lse_head_stride +
+                               token_idx * suffix_lse_token_stride];
+      output_lse[head_idx * output_lse_head_stride +
+                 token_idx * output_lse_token_stride] = s_lse;
     }
     return;
   }
 
   // For tokens within prefix range, merge prefix and suffix
-  float p_lse = prefix_lse[head_idx * num_tokens + token_idx];
-  float s_lse = suffix_lse[head_idx * num_tokens + token_idx];
+  float p_lse = prefix_lse[head_idx * prefix_lse_head_stride +
+                           token_idx * prefix_lse_token_stride];
+  float s_lse = suffix_lse[head_idx * suffix_lse_head_stride +
+                           token_idx * suffix_lse_token_stride];
   p_lse = std::isinf(p_lse) ? -std::numeric_limits<float>::infinity() : p_lse;
   s_lse = std::isinf(s_lse) ? -std::numeric_limits<float>::infinity() : s_lse;
 
@@ -132,7 +139,8 @@ __global__ void merge_attn_states_kernel(
     }
     // We only need to write to output_lse once per head.
     if (output_lse != nullptr && pack_idx == 0) {
-      output_lse[head_idx * num_tokens + token_idx] = max_lse;
+      output_lse[head_idx * output_lse_head_stride +
+                 token_idx * output_lse_token_stride] = max_lse;
     }
     return;
   }
@@ -187,7 +195,8 @@ __global__ void merge_attn_states_kernel(
   // We only need to write to output_lse once per head.
   if (output_lse != nullptr && pack_idx == 0) {
     float out_lse = logf(out_se) + max_lse;
-    output_lse[head_idx * num_tokens + token_idx] = out_lse;
+    output_lse[head_idx * output_lse_head_stride +
+               token_idx * output_lse_token_stride] = out_lse;
   }
 }
 
@@ -221,6 +230,9 @@ __global__ void merge_attn_states_kernel(
             reinterpret_cast<scalar_t*>(suffix_output.data_ptr()),          \
             reinterpret_cast<float*>(suffix_lse.data_ptr()), num_tokens,    \
             num_heads, head_size, prefix_head_stride, output_head_stride,   \
+            prefix_lse_head_stride, prefix_lse_token_stride,                \
+            suffix_lse_head_stride, suffix_lse_token_stride,                \
+            output_lse_head_stride, output_lse_token_stride,                \
             prefix_num_tokens, output_scale_ptr);                           \
   }
 
@@ -259,6 +271,19 @@ void merge_attn_states_launcher(
   const uint head_size = output.size(2);
   const uint prefix_head_stride = prefix_output.stride(1);
   const uint output_head_stride = output.stride(1);
+  // lse tensors are [NUM_HEADS, NUM_TOKENS] but may be non-contiguous views
+  // (e.g. a transpose of a backend's [NUM_TOKENS, NUM_HEADS] output), so index
+  // them by their actual strides rather than assuming a contiguous layout.
+  const uint prefix_lse_head_stride = prefix_lse.stride(0);
+  const uint prefix_lse_token_stride = prefix_lse.stride(1);
+  const uint suffix_lse_head_stride = suffix_lse.stride(0);
+  const uint suffix_lse_token_stride = suffix_lse.stride(1);
+  uint output_lse_head_stride = 0;
+  uint output_lse_token_stride = 0;
+  if (output_lse.has_value()) {
+    output_lse_head_stride = output_lse.value().stride(0);
+    output_lse_token_stride = output_lse.value().stride(1);
+  }
   // Thread mapping is based on input BF16 pack_size
   const uint pack_size = 16 / sizeof(scalar_t);
   STD_TORCH_CHECK(head_size % pack_size == 0,

csrc/libtorch_stable/cache_kernels.cu

@@ -443,6 +443,55 @@ __global__ void concat_and_cache_mla_kernel(
   copy(k_pe, kv_cache, k_pe_stride, block_stride, pe_dim, kv_lora_rank);
 }
 
+// Grouped variant of concat_and_cache_mla: inserts the context K/V for every
+// draft layer in a single launch. Grid is (num_tokens, num_layers); each layer
+// reads its own cache base pointer from kv_cache_ptrs (same pointer-array
+// pattern as copy_blocks_kernel). bf16 only, so it is a raw 16-bit copy with no
+// scaling or quantization; scalar_t is uint16_t for portability.
+template <typename scalar_t>
+__global__ void concat_and_cache_mla_grouped_kernel(
+    const scalar_t* __restrict__ kv_c,  // [num_layers, num_tokens,
+                                        // kv_lora_rank]
+    const scalar_t* __restrict__ k_pe,  // [num_layers, num_tokens, pe_dim]
+    const int64_t* __restrict__ kv_cache_ptrs,  // [num_layers]
+    const int64_t* __restrict__ slot_mapping,   // [num_layers, num_tokens]
+    const int64_t kv_c_layer_stride, const int64_t kv_c_token_stride,
+    const int64_t k_pe_layer_stride, const int64_t k_pe_token_stride,
+    const int64_t slot_layer_stride, const int64_t block_stride,
+    const int64_t entry_stride, const int kv_lora_rank, const int pe_dim,
+    const int block_size) {
+  const int64_t token_idx = blockIdx.x;
+  const int64_t layer_idx = blockIdx.y;
+  const int64_t slot_idx =
+      slot_mapping[layer_idx * slot_layer_stride + token_idx];
+  // NOTE: slot_idx can be -1 if the token is padded
+  if (slot_idx < 0) {
+    return;
+  }
+  const int64_t block_idx = slot_idx / block_size;
+  const int64_t block_offset = slot_idx % block_size;
+
+  scalar_t* __restrict__ kv_cache =
+      reinterpret_cast<scalar_t*>(kv_cache_ptrs[layer_idx]);
+  const scalar_t* __restrict__ kv_c_layer =
+      kv_c + layer_idx * kv_c_layer_stride;
+  const scalar_t* __restrict__ k_pe_layer =
+      k_pe + layer_idx * k_pe_layer_stride;
+
+  auto copy = [&](const scalar_t* __restrict__ src, int64_t src_token_stride,
+                  int size, int offset) {
+    for (int i = threadIdx.x; i < size; i += blockDim.x) {
+      const int64_t src_idx = token_idx * src_token_stride + i;
+      const int64_t dst_idx =
+          block_idx * block_stride + block_offset * entry_stride + i + offset;
+      kv_cache[dst_idx] = src[src_idx];
+    }
+  };
+
+  copy(kv_c_layer, kv_c_token_stride, kv_lora_rank, 0);
+  copy(k_pe_layer, k_pe_token_stride, pe_dim, kv_lora_rank);
+}
+
 template <typename scalar_t, typename cache_t, Fp8KVCacheDataType kv_dt>
 __global__ void concat_and_cache_ds_mla_kernel(
     const scalar_t* __restrict__ kv_c,  // [num_tokens, kv_lora_rank]
@@ -902,6 +951,53 @@ void concat_and_cache_mla(
   }
 }
 
+void concat_and_cache_mla_grouped(
+    torch::stable::Tensor& kv_c,  // [num_layers, num_tokens, kv_lora_rank]
+    torch::stable::Tensor& k_pe,  // [num_layers, num_tokens, pe_dim]
+    torch::stable::Tensor& kv_cache_ptrs,  // [num_layers] int64, on device
+    torch::stable::Tensor& slot_mapping,   // [num_layers, num_tokens] int64
+    int64_t block_size, int64_t block_stride, int64_t entry_stride) {
+  int num_layers = kv_c.size(0);
+  int num_tokens = kv_c.size(1);
+  int kv_lora_rank = kv_c.size(2);
+  int pe_dim = k_pe.size(2);
+
+  STD_TORCH_CHECK(
+      kv_c.scalar_type() == torch::headeronly::ScalarType::BFloat16 &&
+          k_pe.scalar_type() == torch::headeronly::ScalarType::BFloat16,
+      "concat_and_cache_mla_grouped only supports a bf16 KV cache; got kv_c=",
+      kv_c.scalar_type(), ", k_pe=", k_pe.scalar_type());
+  STD_TORCH_CHECK(
+      kv_cache_ptrs.scalar_type() == torch::headeronly::ScalarType::Long,
+      "kv_cache_ptrs must be int64");
+
+  if (num_tokens == 0 || num_layers == 0) {
+    return;
+  }
+
+  const int64_t kv_c_layer_stride = kv_c.stride(0);
+  const int64_t kv_c_token_stride = kv_c.stride(1);
+  const int64_t k_pe_layer_stride = k_pe.stride(0);
+  const int64_t k_pe_token_stride = k_pe.stride(1);
+  const int64_t slot_layer_stride = slot_mapping.stride(0);
+
+  const torch::stable::accelerator::DeviceGuard device_guard(
+      kv_c.get_device_index());
+  const cudaStream_t stream = get_current_cuda_stream();
+
+  dim3 grid(num_tokens, num_layers);
+  dim3 block(std::min(kv_lora_rank, 512));
+  vllm::concat_and_cache_mla_grouped_kernel<uint16_t>
+      <<<grid, block, 0, stream>>>(
+          reinterpret_cast<const uint16_t*>(kv_c.data_ptr()),
+          reinterpret_cast<const uint16_t*>(k_pe.data_ptr()),
+          kv_cache_ptrs.const_data_ptr<int64_t>(),
+          slot_mapping.const_data_ptr<int64_t>(), kv_c_layer_stride,
+          kv_c_token_stride, k_pe_layer_stride, k_pe_token_stride,
+          slot_layer_stride, block_stride, entry_stride, kv_lora_rank, pe_dim,
+          block_size);
+}
+
 namespace vllm {
 
 template <typename Tout, typename Tin, Fp8KVCacheDataType kv_dt>

csrc/libtorch_stable/custom_all_gather_reduce_scatter.cu

@@ -0,0 +1,362 @@
+#include "torch_utils.h"
+
+#include <torch/csrc/stable/macros.h>
+#include <torch/csrc/stable/accelerator.h>
+#include <torch/csrc/stable/tensor.h>
+#include <torch/headeronly/core/ScalarType.h>
+
+#include "custom_all_reduce.cuh"
+#include "custom_all_gather_reduce_scatter.cuh"
+
+namespace vllm {
+
+void CustomAllreduce::allgather(cudaStream_t stream, void* input, void* output,
+                                int size_bytes, int threads, int block_limit) {
+  if (size_bytes % sizeof(CopyPack) != 0)
+    throw std::runtime_error(
+        "custom allgather requires input byte size to be a multiple of " +
+        std::to_string(sizeof(CopyPack)));
+
+  auto ptrs = buffers_.at(input);
+  int size_per_rank = size_bytes / sizeof(CopyPack);
+  int total_size = size_per_rank * world_size_;
+  int blocks = std::min(block_limit, (total_size + threads - 1) / threads);
+
+#define AG_CASE(ngpus)                                                   \
+  case ngpus:                                                            \
+    cross_device_all_gather<ngpus><<<blocks, threads, 0, stream>>>(      \
+        ptrs, sg_, self_sg_, reinterpret_cast<CopyPack*>(output), rank_, \
+        size_per_rank);                                                  \
+    break;
+
+  switch (world_size_) {
+    AG_CASE(2)
+    AG_CASE(4)
+    AG_CASE(6)
+    AG_CASE(8)
+    default:
+      throw std::runtime_error(
+          "custom allgather only supports num gpus in (2,4,6,8)");
+  }
+#undef AG_CASE
+}
+
+template <typename T>
+void CustomAllreduce::mnnvl_lamport_allgather(cudaStream_t stream, T* input,
+                                              T* output, void* local_buffer,
+                                              void* multicast_buffer,
+                                              uint32_t* epochs, int size_bytes,
+                                              int stage_size_bytes) {
+  if (size_bytes % sizeof(typename packed_t<T>::P) != 0 ||
+      stage_size_bytes % sizeof(typename packed_t<T>::P) != 0)
+    throw std::runtime_error(
+        "MNNVL Lamport allgather requires 16-byte aligned sizes");
+
+  auto ptrs = buffers_.at(local_buffer);
+  int size_per_rank = size_bytes / sizeof(typename packed_t<T>::P);
+  int stage_size = stage_size_bytes / sizeof(typename packed_t<T>::P);
+  int blocks =
+      (size_per_rank + kMnnvlLamportAgThreads - 1) / kMnnvlLamportAgThreads;
+
+#if !defined(USE_ROCM) && CUDA_VERSION >= 12000
+  cudaLaunchAttribute attributes[1]{};
+  attributes[0].id = cudaLaunchAttributeProgrammaticStreamSerialization;
+  attributes[0].val.programmaticStreamSerializationAllowed = 1;
+  cudaLaunchConfig_t config{.gridDim = dim3(blocks),
+                            .blockDim = dim3(kMnnvlLamportAgThreads),
+                            .dynamicSmemBytes = 0,
+                            .stream = stream,
+                            .attrs = attributes,
+                            .numAttrs = 1};
+  #define MNNVL_LAMPORT_AG_LAUNCH(ngpus)                                       \
+    CUDACHECK(cudaLaunchKernelEx(&config, &mnnvl_lamport_all_gather<T, ngpus>, \
+                                 ptrs, input, output,                          \
+                                 reinterpret_cast<T*>(multicast_buffer),       \
+                                 epochs, rank_, size_per_rank, stage_size))
+#else
+  #define MNNVL_LAMPORT_AG_LAUNCH(ngpus)                                 \
+    mnnvl_lamport_all_gather<T, ngpus>                                   \
+        <<<blocks, kMnnvlLamportAgThreads, 0, stream>>>(                 \
+            ptrs, input, output, reinterpret_cast<T*>(multicast_buffer), \
+            epochs, rank_, size_per_rank, stage_size)
+#endif
+
+#define MNNVL_LAMPORT_AG_CASE(ngpus) \
+  case ngpus:                        \
+    MNNVL_LAMPORT_AG_LAUNCH(ngpus);  \
+    break;
+
+  switch (world_size_) {
+    MNNVL_LAMPORT_AG_CASE(2)
+    MNNVL_LAMPORT_AG_CASE(4)
+    MNNVL_LAMPORT_AG_CASE(6)
+    MNNVL_LAMPORT_AG_CASE(8)
+    MNNVL_LAMPORT_AG_CASE(16)
+    default:
+      throw std::runtime_error(
+          "MNNVL Lamport allgather only supports num gpus in (2,4,6,8,16)");
+  }
+#undef MNNVL_LAMPORT_AG_CASE
+#undef MNNVL_LAMPORT_AG_LAUNCH
+}
+
+template <typename T>
+void CustomAllreduce::reduce_scatter(cudaStream_t stream, T* input, T* output,
+                                     int size, int threads, int block_limit) {
+  auto packed_size = packed_t<T>::P::size;
+  if (size % (packed_size * world_size_) != 0)
+    throw std::runtime_error(
+        "custom reduce-scatter requires each output shard byte size to be "
+        "a multiple of 16");
+
+  auto ptrs = buffers_.at(input);
+  int size_per_rank = size / packed_size / world_size_;
+  int blocks = std::min(block_limit, (size_per_rank + threads - 1) / threads);
+
+#define RS_CASE(ngpus)                                                     \
+  case ngpus:                                                              \
+    cross_device_reduce_scatter<T, ngpus><<<blocks, threads, 0, stream>>>( \
+        ptrs, sg_, self_sg_, output, rank_, size_per_rank);                \
+    break;
+
+  switch (world_size_) {
+    RS_CASE(2)
+    RS_CASE(4)
+    RS_CASE(6)
+    RS_CASE(8)
+    default:
+      throw std::runtime_error(
+          "custom reduce-scatter only supports num gpus in (2,4,6,8)");
+  }
+#undef RS_CASE
+}
+
+template <typename T>
+void CustomAllreduce::mnnvl_lamport_reduce_scatter(cudaStream_t stream,
+                                                   T* input, T* output,
+                                                   void* local_buffer,
+                                                   uint32_t* epochs, int size,
+                                                   int stage_size_bytes) {
+  auto packed_size = packed_t<T>::P::size;
+  if (size % (packed_size * world_size_) != 0 ||
+      stage_size_bytes % sizeof(typename packed_t<T>::P) != 0)
+    throw std::runtime_error(
+        "MNNVL Lamport reduce-scatter requires 16-byte aligned sizes");
+
+  auto ptrs = buffers_.at(local_buffer);
+  int size_per_rank = size / packed_size / world_size_;
+  int stage_size = stage_size_bytes / sizeof(typename packed_t<T>::P);
+  int blocks_per_rank =
+      (size_per_rank + kMnnvlLamportRsThreads - 1) / kMnnvlLamportRsThreads;
+  int blocks = blocks_per_rank * world_size_;
+
+#if !defined(USE_ROCM) && CUDA_VERSION >= 12000
+  cudaLaunchAttribute attributes[1]{};
+  attributes[0].id = cudaLaunchAttributeProgrammaticStreamSerialization;
+  attributes[0].val.programmaticStreamSerializationAllowed = 1;
+  cudaLaunchConfig_t config{.gridDim = dim3(blocks),
+                            .blockDim = dim3(kMnnvlLamportRsThreads),
+                            .dynamicSmemBytes = 0,
+                            .stream = stream,
+                            .attrs = attributes,
+                            .numAttrs = 1};
+  #define MNNVL_LAMPORT_RS_LAUNCH(ngpus)                                      \
+    CUDACHECK(cudaLaunchKernelEx(                                             \
+        &config, &mnnvl_lamport_reduce_scatter_kernel<T, ngpus>, ptrs, input, \
+        output, epochs, rank_, size_per_rank, stage_size))
+#else
+  #define MNNVL_LAMPORT_RS_LAUNCH(ngpus)                 \
+    mnnvl_lamport_reduce_scatter_kernel<T, ngpus>        \
+        <<<blocks, kMnnvlLamportRsThreads, 0, stream>>>( \
+            ptrs, input, output, epochs, rank_, size_per_rank, stage_size)
+#endif
+
+#define MNNVL_LAMPORT_RS_CASE(ngpus) \
+  case ngpus:                        \
+    MNNVL_LAMPORT_RS_LAUNCH(ngpus);  \
+    break;
+
+  switch (world_size_) {
+    MNNVL_LAMPORT_RS_CASE(2)
+    MNNVL_LAMPORT_RS_CASE(4)
+    MNNVL_LAMPORT_RS_CASE(6)
+    MNNVL_LAMPORT_RS_CASE(8)
+    MNNVL_LAMPORT_RS_CASE(16)
+    default:
+      throw std::runtime_error(
+          "MNNVL Lamport reduce-scatter only supports num gpus in "
+          "(2,4,6,8,16)");
+  }
+#undef MNNVL_LAMPORT_RS_CASE
+#undef MNNVL_LAMPORT_RS_LAUNCH
+}
+
+}  // namespace vllm
+
+using fptr_t = int64_t;
+static_assert(sizeof(void*) == sizeof(fptr_t));
+
+bool _is_weak_contiguous(torch::stable::Tensor& t);
+
+void custom_all_gather(fptr_t _fa, torch::stable::Tensor& inp,
+                       torch::stable::Tensor& out, fptr_t _reg_buffer,
+                       int64_t reg_buffer_sz_bytes) {
+  auto fa = reinterpret_cast<vllm::CustomAllreduce*>(_fa);
+  const torch::stable::accelerator::DeviceGuard device_guard(
+      inp.get_device_index());
+  const cudaStream_t stream = get_current_cuda_stream(inp.get_device_index());
+
+  STD_TORCH_CHECK((inp.scalar_type()) == (out.scalar_type()));
+  STD_TORCH_CHECK((inp.numel() * fa->world_size_) == (out.numel()));
+  STD_TORCH_CHECK(_is_weak_contiguous(out));
+  STD_TORCH_CHECK(_is_weak_contiguous(inp));
+  auto input_size = inp.numel() * inp.element_size();
+  auto reg_buffer = reinterpret_cast<void*>(_reg_buffer);
+  STD_TORCH_CHECK(reg_buffer != nullptr);
+  STD_TORCH_CHECK((input_size) <= (reg_buffer_sz_bytes));
+  STD_CUDA_CHECK(cudaMemcpyAsync(reg_buffer, inp.const_data_ptr(), input_size,
+                                 cudaMemcpyDeviceToDevice, stream));
+  fa->allgather(stream, reg_buffer, out.mutable_data_ptr(), input_size);
+}
+
+void mnnvl_lamport_all_gather(fptr_t _fa, torch::stable::Tensor& inp,
+                              torch::stable::Tensor& out, fptr_t _local_buffer,
+                              fptr_t _multicast_buffer, fptr_t _epoch_buffer,
+                              int64_t stage_sz_bytes) {
+  auto fa = reinterpret_cast<vllm::CustomAllreduce*>(_fa);
+  const torch::stable::accelerator::DeviceGuard device_guard(
+      inp.get_device_index());
+  const cudaStream_t stream = get_current_cuda_stream(inp.get_device_index());
+
+  STD_TORCH_CHECK((inp.scalar_type()) == (out.scalar_type()));
+  STD_TORCH_CHECK((inp.numel() * fa->world_size_) == (out.numel()));
+  STD_TORCH_CHECK(_is_weak_contiguous(out));
+  STD_TORCH_CHECK(_is_weak_contiguous(inp));
+  auto input_size = inp.numel() * inp.element_size();
+  STD_TORCH_CHECK((input_size * fa->world_size_) <= stage_sz_bytes);
+  auto local_buffer = reinterpret_cast<void*>(_local_buffer);
+  auto multicast_buffer = reinterpret_cast<void*>(_multicast_buffer);
+  auto epochs = reinterpret_cast<uint32_t*>(_epoch_buffer);
+  switch (out.scalar_type()) {
+    case torch::headeronly::ScalarType::Float: {
+      fa->mnnvl_lamport_allgather<float>(
+          stream, reinterpret_cast<float*>(inp.mutable_data_ptr()),
+          reinterpret_cast<float*>(out.mutable_data_ptr()), local_buffer,
+          multicast_buffer, epochs, input_size, stage_sz_bytes);
+      break;
+    }
+    case torch::headeronly::ScalarType::Half: {
+      fa->mnnvl_lamport_allgather<half>(
+          stream, reinterpret_cast<half*>(inp.mutable_data_ptr()),
+          reinterpret_cast<half*>(out.mutable_data_ptr()), local_buffer,
+          multicast_buffer, epochs, input_size, stage_sz_bytes);
+      break;
+    }
+#if (__CUDA_ARCH__ >= 800 || !defined(__CUDA_ARCH__))
+    case torch::headeronly::ScalarType::BFloat16: {
+      fa->mnnvl_lamport_allgather<nv_bfloat16>(
+          stream, reinterpret_cast<nv_bfloat16*>(inp.mutable_data_ptr()),
+          reinterpret_cast<nv_bfloat16*>(out.mutable_data_ptr()), local_buffer,
+          multicast_buffer, epochs, input_size, stage_sz_bytes);
+      break;
+    }
+#endif
+    default:
+      throw std::runtime_error(
+          "MNNVL Lamport allgather only supports float32, float16 and "
+          "bfloat16");
+  }
+}
+
+void custom_reduce_scatter(fptr_t _fa, torch::stable::Tensor& inp,
+                           torch::stable::Tensor& out, fptr_t _reg_buffer,
+                           int64_t reg_buffer_sz_bytes) {
+  auto fa = reinterpret_cast<vllm::CustomAllreduce*>(_fa);
+  const torch::stable::accelerator::DeviceGuard device_guard(
+      inp.get_device_index());
+  const cudaStream_t stream = get_current_cuda_stream(inp.get_device_index());
+
+  STD_TORCH_CHECK((inp.scalar_type()) == (out.scalar_type()));
+  STD_TORCH_CHECK((out.numel() * fa->world_size_) == (inp.numel()));
+  STD_TORCH_CHECK(_is_weak_contiguous(out));
+  STD_TORCH_CHECK(_is_weak_contiguous(inp));
+  auto input_size = inp.numel() * inp.element_size();
+  auto reg_buffer = reinterpret_cast<void*>(_reg_buffer);
+  STD_TORCH_CHECK(reg_buffer != nullptr);
+  STD_TORCH_CHECK((input_size) <= (reg_buffer_sz_bytes));
+  STD_CUDA_CHECK(cudaMemcpyAsync(reg_buffer, inp.const_data_ptr(), input_size,
+                                 cudaMemcpyDeviceToDevice, stream));
+  switch (out.scalar_type()) {
+    case torch::headeronly::ScalarType::Float: {
+      fa->reduce_scatter<float>(
+          stream, reinterpret_cast<float*>(reg_buffer),
+          reinterpret_cast<float*>(out.mutable_data_ptr()), inp.numel());
+      break;
+    }
+    case torch::headeronly::ScalarType::Half: {
+      fa->reduce_scatter<half>(stream, reinterpret_cast<half*>(reg_buffer),
+                               reinterpret_cast<half*>(out.mutable_data_ptr()),
+                               inp.numel());
+      break;
+    }
+#if (__CUDA_ARCH__ >= 800 || !defined(__CUDA_ARCH__))
+    case torch::headeronly::ScalarType::BFloat16: {
+      fa->reduce_scatter<nv_bfloat16>(
+          stream, reinterpret_cast<nv_bfloat16*>(reg_buffer),
+          reinterpret_cast<nv_bfloat16*>(out.mutable_data_ptr()), inp.numel());
+      break;
+    }
+#endif
+    default:
+      throw std::runtime_error(
+          "custom reduce-scatter only supports float32, float16 and bfloat16");
+  }
+}
+
+void mnnvl_lamport_reduce_scatter(fptr_t _fa, torch::stable::Tensor& inp,
+                                  torch::stable::Tensor& out,
+                                  fptr_t _local_buffer, fptr_t _epoch_buffer,
+                                  int64_t stage_sz_bytes) {
+  auto fa = reinterpret_cast<vllm::CustomAllreduce*>(_fa);
+  const torch::stable::accelerator::DeviceGuard device_guard(
+      inp.get_device_index());
+  const cudaStream_t stream = get_current_cuda_stream(inp.get_device_index());
+
+  STD_TORCH_CHECK((inp.scalar_type()) == (out.scalar_type()));
+  STD_TORCH_CHECK((out.numel() * fa->world_size_) == (inp.numel()));
+  STD_TORCH_CHECK(_is_weak_contiguous(out));
+  STD_TORCH_CHECK(_is_weak_contiguous(inp));
+  auto input_size = inp.numel() * inp.element_size();
+  STD_TORCH_CHECK(input_size <= stage_sz_bytes);
+  auto local_buffer = reinterpret_cast<void*>(_local_buffer);
+  auto epochs = reinterpret_cast<uint32_t*>(_epoch_buffer);
+  switch (out.scalar_type()) {
+    case torch::headeronly::ScalarType::Float: {
+      fa->mnnvl_lamport_reduce_scatter<float>(
+          stream, reinterpret_cast<float*>(inp.mutable_data_ptr()),
+          reinterpret_cast<float*>(out.mutable_data_ptr()), local_buffer,
+          epochs, inp.numel(), stage_sz_bytes);
+      break;
+    }
+    case torch::headeronly::ScalarType::Half: {
+      fa->mnnvl_lamport_reduce_scatter<half>(
+          stream, reinterpret_cast<half*>(inp.mutable_data_ptr()),
+          reinterpret_cast<half*>(out.mutable_data_ptr()), local_buffer, epochs,
+          inp.numel(), stage_sz_bytes);
+      break;
+    }
+#if (__CUDA_ARCH__ >= 800 || !defined(__CUDA_ARCH__))
+    case torch::headeronly::ScalarType::BFloat16: {
+      fa->mnnvl_lamport_reduce_scatter<nv_bfloat16>(
+          stream, reinterpret_cast<nv_bfloat16*>(inp.mutable_data_ptr()),
+          reinterpret_cast<nv_bfloat16*>(out.mutable_data_ptr()), local_buffer,
+          epochs, inp.numel(), stage_sz_bytes);
+      break;
+    }
+#endif
+    default:
+      throw std::runtime_error(
+          "MNNVL Lamport reduce-scatter only supports float32, float16 and "
+          "bfloat16");
+  }
+}

csrc/libtorch_stable/custom_all_gather_reduce_scatter_ops.cpp

@@ -0,0 +1,29 @@
+#include "ops.h"
+#include "core/registration.h"
+
+#include <torch/csrc/stable/library.h>
+
+STABLE_TORCH_LIBRARY_FRAGMENT(_C_custom_ar, custom_ag_rs) {
+  custom_ag_rs.def(
+      "custom_all_gather(int fa, Tensor inp, Tensor! out, int reg_buffer, "
+      "int reg_buffer_sz_bytes) -> ()");
+  custom_ag_rs.def(
+      "mnnvl_lamport_all_gather(int fa, Tensor inp, Tensor! out, int "
+      "local_buffer, int multicast_buffer, int epoch_buffer, int "
+      "stage_sz_bytes) -> ()");
+  custom_ag_rs.def(
+      "custom_reduce_scatter(int fa, Tensor inp, Tensor! out, int reg_buffer, "
+      "int reg_buffer_sz_bytes) -> ()");
+  custom_ag_rs.def(
+      "mnnvl_lamport_reduce_scatter(int fa, Tensor inp, Tensor! out, int "
+      "local_buffer, int epoch_buffer, int stage_sz_bytes) -> ()");
+}
+
+STABLE_TORCH_LIBRARY_IMPL(_C_custom_ar, CUDA, custom_ag_rs) {
+  custom_ag_rs.impl("custom_all_gather", TORCH_BOX(&custom_all_gather));
+  custom_ag_rs.impl("mnnvl_lamport_all_gather",
+                    TORCH_BOX(&mnnvl_lamport_all_gather));
+  custom_ag_rs.impl("custom_reduce_scatter", TORCH_BOX(&custom_reduce_scatter));
+  custom_ag_rs.impl("mnnvl_lamport_reduce_scatter",
+                    TORCH_BOX(&mnnvl_lamport_reduce_scatter));
+}

csrc/libtorch_stable/custom_all_reduce.cu

@@ -18,14 +18,14 @@ fptr_t init_custom_ar(const std::vector<fptr_t>& fake_ipc_ptrs,
                       torch::stable::Tensor& rank_data, int64_t rank,
                       bool fully_connected) {
   int world_size = fake_ipc_ptrs.size();
-  if (world_size > 8)
-    throw std::invalid_argument("world size > 8 is not supported");
+  if (world_size > vllm::kMaxCustomCollectiveRanks)
+    throw std::invalid_argument("world size > 16 is not supported");
   if (world_size % 2 != 0)
     throw std::invalid_argument("Odd num gpus is not supported for now");
   if (rank < 0 || rank >= world_size)
     throw std::invalid_argument("invalid rank passed in");
 
-  vllm::Signal* ipc_ptrs[8];
+  vllm::Signal* ipc_ptrs[vllm::kMaxCustomCollectiveRanks];
   for (int i = 0; i < world_size; i++) {
     ipc_ptrs[i] = reinterpret_cast<vllm::Signal*>(fake_ipc_ptrs[i]);
   }
@@ -124,7 +124,7 @@ int64_t meta_size() { return sizeof(vllm::Signal); }
 void register_buffer(fptr_t _fa, const std::vector<fptr_t>& fake_ipc_ptrs) {
   auto fa = reinterpret_cast<vllm::CustomAllreduce*>(_fa);
   STD_TORCH_CHECK(fake_ipc_ptrs.size() == fa->world_size_);
-  void* ipc_ptrs[8];
+  void* ipc_ptrs[vllm::kMaxCustomCollectiveRanks];
   for (int i = 0; i < fake_ipc_ptrs.size(); i++) {
     ipc_ptrs[i] = reinterpret_cast<void*>(fake_ipc_ptrs[i]);
   }

csrc/libtorch_stable/dsv3_fused_a_gemm.cu

@@ -647,17 +647,17 @@ __global__ __launch_bounds__(256, 1) void fused_a_gemm_kernel(
 #endif
 }
 
-template <typename T, int kHdIn, int kHdOut, int kTileN>
+template <typename T, int kHdIn, int kHdOut, int kTileN, int kTileK = 256>
 void invokeFusedAGemm(T* output, T const* mat_a, T const* mat_b, int num_tokens,
-                      cudaStream_t const stream) {
-  constexpr int gemm_m = kHdOut;  // 2112
-  int const gemm_n = num_tokens;  // 1-16
-  constexpr int gemm_k = kHdIn;   // 7168
+                      cudaStream_t const stream, bool enable_pdl) {
+  constexpr int gemm_m = kHdOut;
+  int const gemm_n = num_tokens;
+  constexpr int gemm_k = kHdIn;
   constexpr int batch_size = 1;
   std::swap(mat_a, mat_b);
   constexpr int tile_m = 16;
-  constexpr int tile_n = kTileN;                        // 8 or 16
-  constexpr int tile_k = std::max(256, 1024 / tile_n);  // 256
+  constexpr int tile_n = kTileN;
+  constexpr int tile_k = kTileK;
   constexpr int max_stage_cnt =
       1024 * 192 / ((tile_m + tile_n) * tile_k * sizeof(bf16_t));
   constexpr int k_iter_cnt = gemm_k / tile_k;
@@ -679,7 +679,8 @@ void invokeFusedAGemm(T* output, T const* mat_a, T const* mat_b, int num_tokens,
   config.stream = stream;
   cudaLaunchAttribute attrs[1];
   attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization;
-  attrs[0].val.programmaticStreamSerializationAllowed = getEnvEnablePDL();
+  attrs[0].val.programmaticStreamSerializationAllowed =
+      enable_pdl || getEnvEnablePDL();
   config.numAttrs = 1;
   config.attrs = attrs;
   if (smem_bytes >= (48 * 1024)) {
@@ -694,36 +695,50 @@ void invokeFusedAGemm(T* output, T const* mat_a, T const* mat_b, int num_tokens,
                      output, mat_a, mat_b, gemm_n);
 }
 
-template void invokeFusedAGemm<__nv_bfloat16, 7168, 2112, 8>(
-    __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, int num_tokens,
-    cudaStream_t);
-
-template void invokeFusedAGemm<__nv_bfloat16, 7168, 2112, 16>(
-    __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, int num_tokens,
-    cudaStream_t);
+template <typename T, int kHdIn, int kHdOut, int kTileK = 256>
+void invokeFusedAGemmForTokens(T* output, T const* mat_a, T const* mat_b,
+                               int num_tokens, cudaStream_t const stream,
+                               bool enable_pdl) {
+  if (num_tokens <= 8) {
+    invokeFusedAGemm<T, kHdIn, kHdOut, 8, kTileK>(
+        output, mat_a, mat_b, num_tokens, stream, enable_pdl);
+  } else {
+    invokeFusedAGemm<T, kHdIn, kHdOut, 16, kTileK>(
+        output, mat_a, mat_b, num_tokens, stream, enable_pdl);
+  }
+}
 
 void dsv3_fused_a_gemm(torch::stable::Tensor& output,
                        torch::stable::Tensor const& mat_a,
-                       torch::stable::Tensor const& mat_b) {
+                       torch::stable::Tensor const& mat_b, bool enable_pdl) {
   STD_TORCH_CHECK(mat_a.dim() == 2 && mat_b.dim() == 2 && output.dim() == 2);
   int const num_tokens = mat_a.size(0);
   int const hd_in = mat_a.size(1);
   int const hd_out = mat_b.size(1);
 
-  constexpr int kHdIn = 7168;
-  constexpr int kHdOut = 2112;
   STD_TORCH_CHECK(num_tokens >= 1 && num_tokens <= 16,
                   "required 1 <= mat_a.shape[0] <= 16");
-  STD_TORCH_CHECK(hd_in == kHdIn, "required mat_a.shape[1] == 7168");
-  STD_TORCH_CHECK(hd_out == kHdOut, "required mat_b.shape[1] == 2112");
   STD_TORCH_CHECK(output.size(0) == num_tokens,
                   "required output.shape[0] == mat_a.shape[0]");
   STD_TORCH_CHECK(output.size(1) == hd_out,
                   "required output.shape[1] == mat_b.shape[1]");
+  STD_TORCH_CHECK(mat_b.size(0) == hd_in,
+                  "required mat_b.shape[0] == mat_a.shape[1]");
 
-  STD_TORCH_CHECK(mat_a.stride(1) == 1, "mat_a must be a row major tensor");
-  STD_TORCH_CHECK(output.stride(1) == 1, "output must be a row major tensor");
-  STD_TORCH_CHECK(mat_b.stride(0) == 1, "mat_b must be a column major tensor");
+  STD_TORCH_CHECK(mat_a.get_device_index() == mat_b.get_device_index() &&
+                      mat_a.get_device_index() == output.get_device_index(),
+                  "mat_a, mat_b, and output must be on the same device");
+
+  // The kernels index global memory with raw pointers and packed strides, so
+  // reject any padded or transposed view rather than reading out of bounds.
+  STD_TORCH_CHECK(
+      mat_a.stride(0) == hd_in && mat_a.stride(1) == 1,
+      "mat_a must be a packed row-major [num_tokens, hd_in] tensor");
+  STD_TORCH_CHECK(
+      output.stride(0) == hd_out && output.stride(1) == 1,
+      "output must be a packed row-major [num_tokens, hd_out] tensor");
+  STD_TORCH_CHECK(mat_b.stride(0) == 1 && mat_b.stride(1) == hd_in,
+                  "mat_b must be a packed column-major [hd_in, hd_out] tensor");
 
   STD_TORCH_CHECK(
       mat_a.scalar_type() == torch::headeronly::ScalarType::BFloat16 &&
@@ -738,19 +753,85 @@ void dsv3_fused_a_gemm(torch::stable::Tensor& output,
   STD_TORCH_CHECK(getSMVersion() >= 90, "required CUDA ARCH >= SM_90");
 
   auto stream = get_current_cuda_stream(mat_a.get_device_index());
-  if (num_tokens <= 8) {
-    invokeFusedAGemm<__nv_bfloat16, kHdIn, kHdOut, 8>(
-        reinterpret_cast<__nv_bfloat16*>(output.mutable_data_ptr()),
-        reinterpret_cast<__nv_bfloat16 const*>(mat_a.data_ptr()),
-        reinterpret_cast<__nv_bfloat16 const*>(mat_b.data_ptr()), num_tokens,
-        stream);
-  } else {
-    invokeFusedAGemm<__nv_bfloat16, kHdIn, kHdOut, 16>(
-        reinterpret_cast<__nv_bfloat16*>(output.mutable_data_ptr()),
-        reinterpret_cast<__nv_bfloat16 const*>(mat_a.data_ptr()),
-        reinterpret_cast<__nv_bfloat16 const*>(mat_b.data_ptr()), num_tokens,
-        stream);
+  auto* output_ptr =
+      reinterpret_cast<__nv_bfloat16*>(output.mutable_data_ptr());
+  auto const* mat_a_ptr =
+      reinterpret_cast<__nv_bfloat16 const*>(mat_a.data_ptr());
+  auto const* mat_b_ptr =
+      reinterpret_cast<__nv_bfloat16 const*>(mat_b.data_ptr());
+
+#define DISPATCH_DSV3_SHAPE(HD_IN, HD_OUT)                                 \
+  if (hd_in == HD_IN && hd_out == HD_OUT) {                                \
+    invokeFusedAGemmForTokens<__nv_bfloat16, HD_IN, HD_OUT>(               \
+        output_ptr, mat_a_ptr, mat_b_ptr, num_tokens, stream, enable_pdl); \
+    return;                                                                \
   }
+
+  // Shapes the Kimi-K3 selector routes to dsv3_fused_a (see the dsv3 winners
+  // in KIMI_K3_PROJECTIONS) plus the DeepSeek V2/V3 QKV A-projection.
+  DISPATCH_DSV3_SHAPE(7168, 1536)
+  DISPATCH_DSV3_SHAPE(7168, 2112)
+  DISPATCH_DSV3_SHAPE(1536, 2304)
+  DISPATCH_DSV3_SHAPE(1536, 4608)
+  DISPATCH_DSV3_SHAPE(7168, 3584)
+  DISPATCH_DSV3_SHAPE(768, 7168)
+  // TP16 dsv3 winners, as (hd_in=K, hd_out=N). TP16 dense down_proj is absent
+  // because hd_in=2112 is not a multiple of any supported tile_k.
+  DISPATCH_DSV3_SHAPE(1536, 1152)
+  DISPATCH_DSV3_SHAPE(7168, 768)
+  DISPATCH_DSV3_SHAPE(7168, 3216)
+  DISPATCH_DSV3_SHAPE(7168, 4224)
+
+#ifdef VLLM_K3_BENCH_SHAPES
+  // The selector routes these shapes to CuTe or the default GEMM, so they are
+  // never reached in production. They are compiled only for offline
+  // DSV3-vs-CuTe benchmarking.
+  DISPATCH_DSV3_SHAPE(7168, 6288)
+  DISPATCH_DSV3_SHAPE(1536, 7168)
+  DISPATCH_DSV3_SHAPE(3584, 7168)
+  DISPATCH_DSV3_SHAPE(7168, 8448)
+  DISPATCH_DSV3_SHAPE(7168, 20480)
+  DISPATCH_DSV3_SHAPE(7168, 3072)
+  DISPATCH_DSV3_SHAPE(7168, 12448)
+  DISPATCH_DSV3_SHAPE(3072, 7168)
+  DISPATCH_DSV3_SHAPE(8448, 7168)
+  DISPATCH_DSV3_SHAPE(7168, 16896)
+  DISPATCH_DSV3_SHAPE(7168, 40960)
+#endif
+
+#undef DISPATCH_DSV3_SHAPE
+
+  if (hd_in == 128 && hd_out == 1536) {
+    invokeFusedAGemmForTokens<__nv_bfloat16, 128, 1536, 128>(
+        output_ptr, mat_a_ptr, mat_b_ptr, num_tokens, stream, enable_pdl);
+    return;
+  }
+  if (hd_in == 128 && hd_out == 3072) {
+    invokeFusedAGemmForTokens<__nv_bfloat16, 128, 3072, 128>(
+        output_ptr, mat_a_ptr, mat_b_ptr, num_tokens, stream, enable_pdl);
+    return;
+  }
+  // TP16 KDA f_b_proj and shared_expert down_proj. Neither hd_in is a multiple
+  // of 256, so both need the 128 tile_k.
+  if (hd_in == 128 && hd_out == 768) {
+    invokeFusedAGemmForTokens<__nv_bfloat16, 128, 768, 128>(
+        output_ptr, mat_a_ptr, mat_b_ptr, num_tokens, stream, enable_pdl);
+    return;
+  }
+  if (hd_in == 384 && hd_out == 7168) {
+    invokeFusedAGemmForTokens<__nv_bfloat16, 384, 7168, 128>(
+        output_ptr, mat_a_ptr, mat_b_ptr, num_tokens, stream, enable_pdl);
+    return;
+  }
+#ifdef VLLM_K3_BENCH_SHAPES
+  if (hd_in == 4224 && hd_out == 7168) {
+    invokeFusedAGemmForTokens<__nv_bfloat16, 4224, 7168, 128>(
+        output_ptr, mat_a_ptr, mat_b_ptr, num_tokens, stream, enable_pdl);
+    return;
+  }
+#endif
+
+  STD_TORCH_CHECK(false, "unsupported DSV3 fused-A GEMM shape");
 }
 
 STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, m) {

csrc/libtorch_stable/moe/grouped_topk_kernels.cu

@@ -25,6 +25,7 @@
 #include "libtorch_stable/torch_utils.h"
 
 #include <cmath>
+#include <cstdint>
 #include <tuple>
 #include <cuda_fp16.h>
 #include <cuda_bf16.h>
@@ -448,7 +449,8 @@ enum ScoringFunc {
   SCORING_SIGMOID = 1  // apply sigmoid
 };
 
-// Efficient sigmoid approximation from TensorRT-LLM
+// Adapted from
+// https://github.com/NVIDIA/TensorRT-LLM/blob/v1.3.0rc2/cpp/tensorrt_llm/kernels/noAuxTcKernels.cu
 __device__ inline float sigmoid_accurate(float x) {
   return 0.5f * tanhf(0.5f * x) + 0.5f;
 }
@@ -890,6 +892,434 @@ __global__ void grouped_topk_fused_small_expert_count_kernel(
 #endif
 }
 
+// Adapted from
+// https://github.com/flashinfer-ai/flashinfer/blob/06400d062a2d51564bbe781f6f811d0b75ca593e/include/flashinfer/trtllm/fused_moe/RoutingKernelTopK.cuh
+namespace single_group_topk {
+namespace detail {
+
+static constexpr int BlockDim = 256;
+static constexpr uint32_t FullWarpMask = 0xffffffffU;
+static constexpr float InvalidScore = -INFINITY;
+
+// TopK-only tuning: use wider workers and keep these tiers on the block path.
+template <int MaxNumExperts, int MaxNumTopExperts>
+static constexpr bool UseTunedBlockPath =
+    MaxNumTopExperts == 16 && (MaxNumExperts == 896 || MaxNumExperts == 1024);
+
+template <typename T, typename BiasT, ScoringFunc SF>
+__device__ __forceinline__ void preprocess_score(T input, BiasT correction_bias,
+                                                 float& unbiased_score,
+                                                 float& selection_score) {
+  unbiased_score = 0.0F;
+  selection_score = InvalidScore;
+  float const input_float = cuda_cast<float, T>(input);
+  float const bias = cuda_cast<float, BiasT>(correction_bias);
+  if (!is_finite(input_float) || !is_finite(bias)) {
+    return;
+  }
+
+  float const unbiased = apply_scoring<SF>(input_float);
+  float const biased = unbiased + bias;
+  if constexpr (SF == SCORING_NONE) {
+    if (!is_finite(biased)) {
+      return;
+    }
+  }
+  unbiased_score = unbiased;
+  selection_score = biased == 0.0F ? 0.0F : biased;
+}
+
+template <typename IdxT>
+__device__ __forceinline__ void write_outputs(
+    cg::thread_block_tile<WARP_SIZE> const& warp, float lane_selection_score,
+    float lane_unbiased, int32_t lane_expert, int32_t lane, int32_t token,
+    int32_t topk, float* topk_values, IdxT* topk_indices, bool renormalize,
+    float routed_scaling_factor) {
+  bool const finite_selection =
+      lane < topk && lane_selection_score != InvalidScore;
+  lane_unbiased = finite_selection ? lane_unbiased : 0.0F;
+  unsigned const finite_mask = __ballot_sync(FullWarpMask, finite_selection);
+  float const sum = cg::reduce(warp, lane_unbiased, cg::plus<float>{});
+
+  if (lane < topk) {
+    float output = 0.0F;
+    if (finite_mask == 0) {
+      if (renormalize) {
+        output = 1.0F / static_cast<float>(topk);
+      }
+    } else if (finite_selection) {
+      float scale = routed_scaling_factor;
+      if (renormalize) {
+        scale /= sum + 1e-20F;
+      }
+      output = lane_unbiased * scale;
+    }
+
+    int64_t const output_index = int64_t{token} * topk + lane;
+    topk_values[output_index] = output;
+    topk_indices[output_index] = static_cast<IdxT>(lane_expert);
+  }
+}
+
+template <typename T, typename BiasT, typename IdxT, ScoringFunc SF,
+          int MaxNumExperts, int MaxNumTopExperts>
+__global__ void __launch_bounds__(BlockDim)
+    single_group_topk_block_kernel(T const* scores, float* topk_values,
+                                   IdxT* topk_indices, BiasT const* bias,
+                                   int64_t num_experts, int64_t topk,
+                                   bool renormalize,
+                                   float routed_scaling_factor,
+                                   bool enable_pdl) {
+  static constexpr int NumChunks = (MaxNumExperts + WARP_SIZE - 1) / WARP_SIZE;
+  static constexpr int WorkerValuesPerLane =
+      UseTunedBlockPath<MaxNumExperts, MaxNumTopExperts> ? 8 : 4;
+  static constexpr int ExpertsPerWorkerWarp = WorkerValuesPerLane * WARP_SIZE;
+  using LaneOwnedRange =
+      reduce_topk::HighExpertLaneOwnedTopKRange<MaxNumExperts,
+                                                MaxNumTopExperts>;
+  static constexpr int NumWorkerWarps =
+      (MaxNumExperts + ExpertsPerWorkerWarp - 1) / ExpertsPerWorkerWarp;
+  static constexpr int NumIntermediate = NumWorkerWarps * MaxNumTopExperts;
+  static constexpr int MergeValuesPerLane =
+      (NumIntermediate + WARP_SIZE - 1) / WARP_SIZE;
+  static constexpr bool LaneOwnedResourcesFit =
+      NumWorkerWarps <= BlockDim / WARP_SIZE && MergeValuesPerLane <= 64;
+  static constexpr bool UseHierarchicalLaneTopK =
+      LaneOwnedRange::kEnabled && LaneOwnedResourcesFit;
+
+  static_assert(NumChunks <= 64);
+  static_assert(MaxNumTopExperts <= WARP_SIZE);
+
+#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
+  if (enable_pdl) {
+    cudaGridDependencySynchronize();
+  }
+#endif
+
+  __shared__ float __attribute((aligned(128))) biased_scores[MaxNumExperts];
+  __shared__ float __attribute((aligned(128))) unbiased_scores[MaxNumExperts];
+
+  int32_t const token = static_cast<int32_t>(blockIdx.x);
+  int32_t const lane = static_cast<int32_t>(threadIdx.x) % WARP_SIZE;
+  int32_t const warp_id = static_cast<int32_t>(threadIdx.x) / WARP_SIZE;
+  int32_t const num_experts_i32 = static_cast<int32_t>(num_experts);
+  int32_t const topk_i32 = static_cast<int32_t>(topk);
+  T const* token_scores = scores + int64_t{token} * num_experts;
+
+  for (int32_t expert = static_cast<int32_t>(threadIdx.x);
+       expert < num_experts_i32; expert += BlockDim) {
+    preprocess_score<T, BiasT, SF>(token_scores[expert], bias[expert],
+                                   unbiased_scores[expert],
+                                   biased_scores[expert]);
+  }
+  __syncthreads();
+
+  auto warp = cg::tiled_partition<WARP_SIZE>(cg::this_thread_block());
+
+  if constexpr (UseHierarchicalLaneTopK) {
+    __shared__ float
+        __attribute((aligned(128))) intermediate_scores[NumIntermediate];
+    __shared__ int32_t
+        __attribute((aligned(128))) intermediate_indices[NumIntermediate];
+
+    if (warp_id < NumWorkerWarps) {
+      float local_scores[WorkerValuesPerLane];
+      int32_t local_indices[WorkerValuesPerLane];
+#pragma unroll
+      for (int index = 0; index < WorkerValuesPerLane; ++index) {
+        int32_t const expert =
+            warp_id * ExpertsPerWorkerWarp + index * WARP_SIZE + lane;
+        local_scores[index] =
+            expert < num_experts_i32 ? biased_scores[expert] : InvalidScore;
+        local_indices[index] = expert;
+      }
+
+      float lane_score;
+      int32_t lane_expert;
+      reduce_topk::reduceTopKForLane<MaxNumTopExperts>(
+          warp, lane_score, lane_expert, local_scores, local_indices,
+          InvalidScore, lane);
+      if (lane < MaxNumTopExperts) {
+        int32_t const intermediate = warp_id * MaxNumTopExperts + lane;
+        bool const active = lane < topk_i32;
+        intermediate_scores[intermediate] = active ? lane_score : InvalidScore;
+        intermediate_indices[intermediate] =
+            active ? lane_expert : MaxNumExperts;
+      }
+    }
+    __syncthreads();
+
+    if (warp_id != 0) {
+      return;
+    }
+
+    float merge_scores[MergeValuesPerLane];
+    int32_t merge_indices[MergeValuesPerLane];
+#pragma unroll
+    for (int index = 0; index < MergeValuesPerLane; ++index) {
+      int32_t const intermediate = index * WARP_SIZE + lane;
+      bool const active = intermediate < NumIntermediate;
+      merge_scores[index] =
+          active ? intermediate_scores[intermediate] : InvalidScore;
+      merge_indices[index] =
+          active ? intermediate_indices[intermediate] : MaxNumExperts;
+    }
+
+    float lane_score;
+    int32_t lane_expert;
+    reduce_topk::reduceTopKForLane<MaxNumTopExperts>(
+        warp, lane_score, lane_expert, merge_scores, merge_indices,
+        InvalidScore, lane);
+    float const lane_unbiased =
+        lane < topk_i32 && lane_expert >= 0 && lane_expert < num_experts_i32
+            ? unbiased_scores[lane_expert]
+            : 0.0F;
+    write_outputs(warp, lane_score, lane_unbiased, lane_expert, lane, token,
+                  topk_i32, topk_values, topk_indices, renormalize,
+                  routed_scaling_factor);
+  } else {
+    if (warp_id != 0) {
+      return;
+    }
+
+    float local_scores[NumChunks];
+    int32_t local_indices[NumChunks];
+#pragma unroll
+    for (int index = 0; index < NumChunks; ++index) {
+      int32_t const expert = index * WARP_SIZE + lane;
+      local_scores[index] =
+          expert < num_experts_i32 ? biased_scores[expert] : InvalidScore;
+      local_indices[index] = expert;
+    }
+
+    float top_scores[MaxNumTopExperts];
+    int32_t top_experts[MaxNumTopExperts];
+    reduce_topk::reduceTopK(warp, top_scores, top_experts, local_scores,
+                            local_indices, InvalidScore, topk_i32);
+    float const lane_score = lane < topk_i32 ? top_scores[lane] : InvalidScore;
+    int32_t const lane_expert = lane < topk_i32 ? top_experts[lane] : -1;
+    float const lane_unbiased =
+        lane < topk_i32 && lane_expert >= 0 && lane_expert < num_experts_i32
+            ? unbiased_scores[lane_expert]
+            : 0.0F;
+    write_outputs(warp, lane_score, lane_unbiased, lane_expert, lane, token,
+                  topk_i32, topk_values, topk_indices, renormalize,
+                  routed_scaling_factor);
+  }
+
+#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
+  if (enable_pdl) {
+    cudaTriggerProgrammaticLaunchCompletion();
+  }
+#endif
+}
+
+template <int MaxNumExperts>
+struct WarpTopKLaunchConfig {
+  static constexpr int DefaultBlockDim =
+      MaxNumExperts <= 1024 ? MaxNumExperts : 1024;
+  static constexpr int BlockDim = DefaultBlockDim > 256 ? 256 : DefaultBlockDim;
+  static constexpr int NumWarps = BlockDim / WARP_SIZE;
+  static constexpr int MaxBlockScale =
+      (DefaultBlockDim + BlockDim - 1) / BlockDim;
+  static constexpr int MaxBlocks = 1024 * MaxBlockScale;
+
+  static_assert(BlockDim % WARP_SIZE == 0);
+
+  static uint32_t grid_dim(int64_t num_tokens) {
+    int64_t const token_blocks = (num_tokens + NumWarps - 1) / NumWarps;
+    int64_t const selected =
+        token_blocks < MaxBlocks ? token_blocks : MaxBlocks;
+    return static_cast<uint32_t>(selected > 0 ? selected : 1);
+  }
+};
+
+template <typename T, typename BiasT, typename IdxT, ScoringFunc SF,
+          int MaxNumExperts, int MaxNumTopExperts>
+__global__ void __launch_bounds__(WarpTopKLaunchConfig<MaxNumExperts>::BlockDim)
+    single_group_topk_warp_kernel(T const* scores, float* topk_values,
+                                  IdxT* topk_indices, BiasT const* bias,
+                                  int64_t num_tokens, int64_t num_experts,
+                                  int64_t topk, bool renormalize,
+                                  float routed_scaling_factor,
+                                  bool enable_pdl) {
+  static constexpr int NumChunks = (MaxNumExperts + WARP_SIZE - 1) / WARP_SIZE;
+  static constexpr int WarpBlockDim =
+      WarpTopKLaunchConfig<MaxNumExperts>::BlockDim;
+  using LaneOwnedRange =
+      reduce_topk::HighExpertLaneOwnedTopKRange<MaxNumExperts,
+                                                MaxNumTopExperts>;
+
+  static_assert(NumChunks <= 64);
+  static_assert(MaxNumTopExperts <= WARP_SIZE);
+
+#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
+  if (enable_pdl) {
+    cudaGridDependencySynchronize();
+  }
+#endif
+
+  int32_t const lane = static_cast<int32_t>(threadIdx.x) % WARP_SIZE;
+  int32_t const warp_id = static_cast<int32_t>(threadIdx.x) / WARP_SIZE;
+  int32_t const global_warp =
+      static_cast<int32_t>(blockIdx.x) * WarpBlockDim / WARP_SIZE + warp_id;
+  int32_t const global_warp_stride =
+      static_cast<int32_t>(gridDim.x) * WarpBlockDim / WARP_SIZE;
+  int32_t const num_experts_i32 = static_cast<int32_t>(num_experts);
+  int32_t const topk_i32 = static_cast<int32_t>(topk);
+  auto warp = cg::tiled_partition<WARP_SIZE>(cg::this_thread_block());
+
+  for (int32_t token = global_warp; token < num_tokens;
+       token += global_warp_stride) {
+    T const* token_scores = scores + int64_t{token} * num_experts;
+    float local_scores[NumChunks];
+    int32_t local_indices[NumChunks];
+#pragma unroll
+    for (int index = 0; index < NumChunks; ++index) {
+      int32_t const expert = index * WARP_SIZE + lane;
+      float unbiased;
+      float selection;
+      if (expert < num_experts_i32) {
+        preprocess_score<T, BiasT, SF>(token_scores[expert], bias[expert],
+                                       unbiased, selection);
+      } else {
+        selection = InvalidScore;
+      }
+      local_scores[index] = selection;
+      local_indices[index] = expert;
+    }
+
+    float lane_score;
+    int32_t lane_expert;
+    if constexpr (LaneOwnedRange::kEnabled) {
+      reduce_topk::reduceTopKForLane<MaxNumTopExperts>(
+          warp, lane_score, lane_expert, local_scores, local_indices,
+          InvalidScore, lane);
+    } else {
+      float top_scores[MaxNumTopExperts];
+      int32_t top_experts[MaxNumTopExperts];
+      reduce_topk::reduceTopK(warp, top_scores, top_experts, local_scores,
+                              local_indices, InvalidScore, topk_i32);
+      lane_score = lane < topk_i32 ? top_scores[lane] : InvalidScore;
+      lane_expert = lane < topk_i32 ? top_experts[lane] : -1;
+    }
+
+    float lane_unbiased = 0.0F;
+    if (lane < topk_i32 && lane_expert >= 0 && lane_expert < num_experts_i32) {
+      lane_unbiased = lane_score - cuda_cast<float, BiasT>(bias[lane_expert]);
+    }
+    write_outputs(warp, lane_score, lane_unbiased, lane_expert, lane, token,
+                  topk_i32, topk_values, topk_indices, renormalize,
+                  routed_scaling_factor);
+  }
+
+#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
+  if (enable_pdl) {
+    cudaTriggerProgrammaticLaunchCompletion();
+  }
+#endif
+}
+
+template <int Experts, int TopK>
+struct Tier {
+  static constexpr int kExperts = Experts;
+  static constexpr int kTopK = TopK;
+};
+
+template <typename... Tiers>
+struct TierList {};
+
+using SigmoidBiasTiers =
+    TierList<Tier<128, 8>, Tier<256, 8>, Tier<384, 8>, Tier<512, 8>,
+             Tier<512, 22>, Tier<768, 16>, Tier<896, 16>, Tier<1024, 16>>;
+
+using PrecomputedSoftmaxBiasTiers =
+    TierList<Tier<128, 4>, Tier<128, 8>, Tier<160, 8>, Tier<256, 8>,
+             Tier<256, 16>, Tier<512, 8>, Tier<512, 16>, Tier<512, 22>,
+             Tier<512, 32>, Tier<576, 8>, Tier<768, 16>, Tier<896, 16>,
+             Tier<1024, 16>>;
+
+template <typename T, typename BiasT, typename IdxT, ScoringFunc SF,
+          int MaxNumExperts, int MaxNumTopExperts>
+void launch(T* scores, float* topk_values, IdxT* topk_indices,
+            BiasT const* bias, int64_t num_tokens, int64_t num_experts,
+            int64_t topk, bool renormalize, double routed_scaling_factor,
+            bool enable_pdl, cudaLaunchConfig_t& config) {
+  config.dynamicSmemBytes = 0;
+  bool const use_block_kernel =
+      UseTunedBlockPath<MaxNumExperts, MaxNumTopExperts> ||
+      MaxNumExperts > 1024 || num_experts >= 1024 ||
+      (num_experts >= 256 && num_tokens <= 1024);
+  if (use_block_kernel) {
+    config.gridDim = static_cast<uint32_t>(num_tokens);
+    config.blockDim = BlockDim;
+    cudaLaunchKernelEx(
+        &config,
+        &single_group_topk_block_kernel<T, BiasT, IdxT, SF, MaxNumExperts,
+                                        MaxNumTopExperts>,
+        scores, topk_values, topk_indices, bias, num_experts, topk, renormalize,
+        static_cast<float>(routed_scaling_factor), enable_pdl);
+  } else {
+    using WarpConfig = WarpTopKLaunchConfig<MaxNumExperts>;
+    config.gridDim = WarpConfig::grid_dim(num_tokens);
+    config.blockDim = WarpConfig::BlockDim;
+    cudaLaunchKernelEx(
+        &config,
+        &single_group_topk_warp_kernel<T, BiasT, IdxT, SF, MaxNumExperts,
+                                       MaxNumTopExperts>,
+        scores, topk_values, topk_indices, bias, num_tokens, num_experts, topk,
+        renormalize, static_cast<float>(routed_scaling_factor), enable_pdl);
+  }
+}
+
+template <typename T, typename BiasT, typename IdxT, ScoringFunc SF>
+bool dispatch(TierList<>*, T*, float*, IdxT*, BiasT const*, int64_t, int64_t,
+              int64_t, bool, double, bool, cudaLaunchConfig_t&) {
+  return false;
+}
+
+template <typename T, typename BiasT, typename IdxT, ScoringFunc SF,
+          typename First, typename... Rest>
+bool dispatch(TierList<First, Rest...>*, T* scores, float* topk_values,
+              IdxT* topk_indices, BiasT const* bias, int64_t num_tokens,
+              int64_t num_experts, int64_t topk, bool renormalize,
+              double routed_scaling_factor, bool enable_pdl,
+              cudaLaunchConfig_t& config) {
+  if (num_experts <= First::kExperts && topk <= First::kTopK) {
+    launch<T, BiasT, IdxT, SF, First::kExperts, First::kTopK>(
+        scores, topk_values, topk_indices, bias, num_tokens, num_experts, topk,
+        renormalize, routed_scaling_factor, enable_pdl, config);
+    return true;
+  }
+  return dispatch<T, BiasT, IdxT, SF>(
+      static_cast<TierList<Rest...>*>(nullptr), scores, topk_values,
+      topk_indices, bias, num_tokens, num_experts, topk, renormalize,
+      routed_scaling_factor, enable_pdl, config);
+}
+
+}  // namespace detail
+
+template <typename T, typename BiasT, typename IdxT, ScoringFunc SF>
+bool invoke(T* scores, float* topk_values, IdxT* topk_indices,
+            BiasT const* bias, int64_t num_tokens, int64_t num_experts,
+            int64_t topk, bool renormalize, double routed_scaling_factor,
+            bool enable_pdl, cudaLaunchConfig_t& config) {
+  static_assert(SF == SCORING_NONE || SF == SCORING_SIGMOID);
+  if constexpr (SF == SCORING_SIGMOID) {
+    return detail::dispatch<T, BiasT, IdxT, SF>(
+        static_cast<detail::SigmoidBiasTiers*>(nullptr), scores, topk_values,
+        topk_indices, bias, num_tokens, num_experts, topk, renormalize,
+        routed_scaling_factor, enable_pdl, config);
+  } else {
+    return detail::dispatch<T, BiasT, IdxT, SF>(
+        static_cast<detail::PrecomputedSoftmaxBiasTiers*>(nullptr), scores,
+        topk_values, topk_indices, bias, num_tokens, num_experts, topk,
+        renormalize, routed_scaling_factor, enable_pdl, config);
+  }
+}
+
+}  // namespace single_group_topk
+
 template <typename T, typename BiasT, typename IdxT, ScoringFunc SF>
 void invokeNoAuxTc(T* scores, float* topk_values, IdxT* topk_indices,
                    BiasT const* bias, int64_t const num_tokens,
@@ -905,6 +1335,12 @@ void invokeNoAuxTc(T* scores, float* topk_values, IdxT* topk_indices,
   attrs[0].val.programmaticStreamSerializationAllowed = enable_pdl;
   config.numAttrs = 1;
   config.attrs = attrs;
+  if (n_group == 1 && topk_group == 1 &&
+      single_group_topk::invoke<T, BiasT, IdxT, SF>(
+          scores, topk_values, topk_indices, bias, num_tokens, num_experts,
+          topk, renormalize, routed_scaling_factor, enable_pdl, config)) {
+    return;
+  }
 
   // Check if we can use the optimized
   // grouped_topk_fused_small_expert_count_kernel

csrc/libtorch_stable/moe/moeTopKFuncs.cuh

@@ -1,6 +1,7 @@
 /*
  * Adapted from
  * https://github.com/NVIDIA/TensorRT-LLM/blob/v1.3.0rc2/cpp/tensorrt_llm/kernels/moeTopKFuncs.cuh
+ * https://github.com/flashinfer-ai/flashinfer/blob/06400d062a2d51564bbe781f6f811d0b75ca593e/include/flashinfer/trtllm/fused_moe/RoutingKernelTopK.cuh
  * Copyright (c) 2026, The vLLM team.
  * SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION. All rights
  * reserved. SPDX-License-Identifier: Apache-2.0
@@ -23,6 +24,9 @@
 #include <cooperative_groups/reduce.h>
 #include <cub/cub.cuh>
 
+#include <cstdint>
+#include <type_traits>
+
 namespace vllm {
 namespace moe {
 namespace reduce_topk {
@@ -38,11 +42,10 @@ struct TopKRedType {
       "Top K reduction only implemented for int, float, float16 and bfloat16");
 
   using TypeCmp = std::conditional_t<sizeof(T) == 4, uint64_t, uint32_t>;
-  using IdxT = std::conditional_t<sizeof(T) == 4, int32_t, int16_t>;
 
   static constexpr int kMoveBits = (sizeof(T) == 4) ? 32 : 16;
   static constexpr int kMaxIdx = 65535;
-  TypeCmp compValIdx;
+  TypeCmp compVal;
 
   static __host__ __device__ inline TypeCmp makeCmpVal(T val, int32_t idx = 0) {
     auto valueBits = cub::Traits<T>::TwiddleIn(
@@ -69,69 +72,175 @@ struct TopKRedType {
   __host__ __device__ TopKRedType() = default;
 
   __host__ __device__ TopKRedType(T val, int32_t idx)
-      : compValIdx(makeCmpVal(val, idx)) {}
+      : compVal(makeCmpVal(val, idx)) {}
 
-  __host__ __device__ operator TypeCmp() const noexcept { return compValIdx; }
+  __host__ __device__ operator TypeCmp() const noexcept { return compVal; }
 
   __device__ inline TypeCmp reduce(
       cg::thread_block_tile<kWARP_SIZE> const& warp) {
-    return cg::reduce(warp, compValIdx, cg::greater<TypeCmp>{});
+#ifdef __CUDA_ARCH__
+    static constexpr bool kHAS_FAST_REDUX = (__CUDA_ARCH__ / 100) >= 10;
+#else
+    static constexpr bool kHAS_FAST_REDUX = false;
+#endif
+    if constexpr (!kHAS_FAST_REDUX) {
+      return cg::reduce(warp, compVal, cg::greater<TypeCmp>{});
+    } else if constexpr (sizeof(TypeCmp) == 8) {
+      uint32_t hi = static_cast<uint32_t>(compVal >> 32);
+      uint32_t lo = static_cast<uint32_t>(compVal & 0xffffffffu);
+      uint32_t maxHi;
+      asm volatile("redux.sync.max.u32 %0, %1, 0xffffffff;\n"
+                   : "=r"(maxHi)
+                   : "r"(hi));
+      uint32_t loContrib = hi == maxHi ? lo : 0u;
+      uint32_t maxLo;
+      asm volatile("redux.sync.max.u32 %0, %1, 0xffffffff;\n"
+                   : "=r"(maxLo)
+                   : "r"(loContrib));
+      return (static_cast<TypeCmp>(maxHi) << 32) | static_cast<TypeCmp>(maxLo);
+    } else {
+      TypeCmp result;
+      asm volatile("redux.sync.max.u32 %0, %1, 0xffffffff;\n"
+                   : "=r"(result)
+                   : "r"(compVal));
+      return result;
+    }
   }
 };
 
-////////////////////////////////////////////////////////////////////////////////////////////////////
-
-template <int K_, bool Enable_>
-struct TopKIdx {
-  // by default, empty
+template <int N>
+struct IsPowerOf2 {
+  static constexpr bool value = N > 0 && (N & (N - 1)) == 0;
 };
 
-template <int K_>
-struct TopKIdx<K_, true> {
-  static constexpr int K = K_;
-  int32_t val[K];
+template <int N>
+struct NextPow2 {
+ private:
+  static constexpr unsigned u = static_cast<unsigned>(N - 1);
+  static constexpr unsigned s1 = u | (u >> 1);
+  static constexpr unsigned s2 = s1 | (s1 >> 2);
+  static constexpr unsigned s3 = s2 | (s2 >> 4);
+  static constexpr unsigned s4 = s3 | (s3 >> 8);
+  static constexpr unsigned s5 = s4 | (s4 >> 16);
+
+ public:
+  static constexpr int value = N <= 1 ? 1 : static_cast<int>(s5 + 1);
 };
 
-////////////////////////////////////////////////////////////////////////////////////////////////////
-
-#define TOPK_SWAP(I, J)                                         \
-  {                                                             \
-    auto pairMin = min(topK[I].compValIdx, topK[J].compValIdx); \
-    auto pairMax = max(topK[I].compValIdx, topK[J].compValIdx); \
-    topK[I].compValIdx = pairMax;                               \
-    topK[J].compValIdx = pairMin;                               \
+template <int A, int B, int Size, typename T>
+__device__ __forceinline__ void topkCompareSwap(T* a) {
+  if constexpr (A < Size && B < Size) {
+    if (a[A] < a[B]) {
+      T tmp = a[A];
+      a[A] = a[B];
+      a[B] = tmp;
+    }
+  } else {
+    (void)a;
   }
+}
+
+template <int I, int End, int Step, int PairStride, int Size, typename T>
+__device__ __forceinline__ void topkMergePairs(T* a) {
+  if constexpr (I + Step < End) {
+    topkCompareSwap<I, I + Step, Size, T>(a);
+    topkMergePairs<I + PairStride, End, Step, PairStride, Size, T>(a);
+  } else {
+    (void)a;
+  }
+}
+
+template <int Lo, int N, int R, int Size, typename T>
+__device__ __forceinline__ void topkOEM(T* a) {
+  constexpr int M = R * 2;
+  if constexpr (M < N) {
+    topkOEM<Lo, N, M, Size, T>(a);
+    topkOEM<Lo + R, N - R, M, Size, T>(a);
+    topkMergePairs<Lo + R, Lo + N, R, M, Size, T>(a);
+  } else if constexpr (R < N) {
+    topkCompareSwap<Lo, Lo + R, Size, T>(a);
+  } else {
+    (void)a;
+  }
+}
+
+template <int Lo, int N, int Size, typename T>
+__device__ __forceinline__ void topkSortBatcher(T* a) {
+  if constexpr (N > 1) {
+    constexpr int Half = N / 2;
+    topkSortBatcher<Lo, Half, Size, T>(a);
+    topkSortBatcher<Lo + Half, N - Half, Size, T>(a);
+    topkOEM<Lo, N, 1, Size, T>(a);
+  } else {
+    (void)a;
+  }
+}
 
 template <int N, typename RedType>
-struct Sort;
+struct Sort {
+  static_assert(N > 0 && N <= 64, "Sort only supports N in range [1, 64]");
+
+  static __device__ void run(RedType* topK) {
+    if constexpr (IsPowerOf2<N>::value) {
+#pragma unroll
+      for (int k = 2; k <= N; k *= 2) {
+#pragma unroll
+        for (int j = k / 2; j > 0; j /= 2) {
+#pragma unroll
+          for (int i = 0; i < N; ++i) {
+            int ixj = i ^ j;
+            if (ixj > i) {
+              if ((i & k) == 0) {
+                if (topK[i].compVal < topK[ixj].compVal) {
+                  auto tmp = topK[i].compVal;
+                  topK[i].compVal = topK[ixj].compVal;
+                  topK[ixj].compVal = tmp;
+                }
+              } else {
+                if (topK[i].compVal > topK[ixj].compVal) {
+                  auto tmp = topK[i].compVal;
+                  topK[i].compVal = topK[ixj].compVal;
+                  topK[ixj].compVal = tmp;
+                }
+              }
+            }
+          }
+        }
+      }
+    } else {
+      constexpr int P = NextPow2<N>::value;
+      topkSortBatcher<0, P, N, RedType>(topK);
+    }
+  }
+};
 
 template <typename RedType>
 struct Sort<1, RedType> {
-  static __device__ void run(RedType* topK) {}
+  static __device__ void run(RedType*) {}
 };
 
 template <typename RedType>
 struct Sort<2, RedType> {
-  static __device__ void run(RedType* topK) { TOPK_SWAP(0, 1); }
+  static __device__ void run(RedType* topK) { topkCompareSwap<0, 1, 2>(topK); }
 };
 
 template <typename RedType>
 struct Sort<3, RedType> {
   static __device__ void run(RedType* topK) {
-    TOPK_SWAP(0, 1);
-    TOPK_SWAP(1, 2);
-    TOPK_SWAP(0, 1);
+    topkCompareSwap<0, 1, 3>(topK);
+    topkCompareSwap<1, 2, 3>(topK);
+    topkCompareSwap<0, 1, 3>(topK);
   }
 };
 
 template <typename RedType>
 struct Sort<4, RedType> {
   static __device__ void run(RedType* topK) {
-    TOPK_SWAP(0, 2);
-    TOPK_SWAP(1, 3);
-    TOPK_SWAP(0, 1);
-    TOPK_SWAP(2, 3);
-    TOPK_SWAP(1, 2);
+    topkCompareSwap<0, 2, 4>(topK);
+    topkCompareSwap<1, 3, 4>(topK);
+    topkCompareSwap<0, 1, 4>(topK);
+    topkCompareSwap<2, 3, 4>(topK);
+    topkCompareSwap<1, 2, 4>(topK);
   }
 };
 
@@ -147,110 +256,112 @@ __forceinline__ __device__ void reduceTopK(
   typename RedType::TypeCmp packedMax{};
 #pragma unroll
   for (int kk = 0; kk < actualK; ++kk) {
-    topK =
-        kk > 0 && packedMax == topK.compValIdx ? RedType{minValue, idx} : topK;
-    // get the next largest value
+    topK = kk > 0 && packedMax == topK.compVal ? RedType{minValue, idx} : topK;
     packedMax = topK.reduce(warp);
     RedType::unpack(out[kk], outIdx[kk], packedMax);
   }
 };
 
-template <int K, typename Type, int N, bool IsSorted = false>
-__device__ void reduceTopKFunc(cg::thread_block_tile<kWARP_SIZE> const& warp,
-                               Type (&out)[K], int32_t (&outIdx)[K],
-                               Type (&value)[N], int32_t (&idx)[N],
-                               Type minValue, int actualK = K) {
-  static_assert(K > 0, "Top K must have K > 0");
-  static_assert(K < kWARP_SIZE, "Top K must have K < kWARP_SIZE");
-  static_assert(N > 0, "Top K must have N > 0");
-  static_assert(N < 5,
-                "Only support candidates number less than or equal to 128");
-  using RedType = TopKRedType<Type>;
-  RedType topK[N];
-#pragma unroll
-  for (int nn = 0; nn < N; ++nn) {
-    topK[nn] = RedType{value[nn], idx[nn]};
-  }
-
-  if constexpr (!IsSorted) {
-    Sort<N, RedType>::run(topK);
-  }
-  typename RedType::TypeCmp packedMax{};
-#pragma unroll
-  for (int kk = 0; kk < actualK; ++kk) {
-    bool update = kk > 0 && packedMax == topK[0].compValIdx;
-#pragma unroll
-    for (int nn = 0; nn < N; ++nn) {
-      topK[nn] = update && nn == N - 1 ? RedType{minValue, idx[nn]}
-                 : update              ? topK[nn + 1]
-                                       : topK[nn];
-    }
-    // get the next largest value
-    packedMax = topK[0].reduce(warp);
-    RedType::unpack(out[kk], outIdx[kk], packedMax);
-  }
-};
-
 template <int K, typename Type, int N>
 __forceinline__ __device__ void reduceTopK(
     cg::thread_block_tile<kWARP_SIZE> const& warp, Type (&out)[K],
     int32_t (&outIdx)[K], Type (&value)[N], int32_t (&idx)[N],
     Type const minValue, int actualK = K) {
   static_assert(K > 0, "Top K must have K > 0");
-  static_assert(K < kWARP_SIZE, "Top K must have K < kWARP_SIZE");
+  static_assert(K <= kWARP_SIZE, "Top K must have K <= kWARP_SIZE");
   static_assert(N > 0, "Top K must have N > 0");
-  static_assert(
-      N <= 16,
-      "Only support candidates number less than or equal to 16*32=512");
-  static_assert(N <= 4 || N % 4 == 0,
-                "Only support candidates number is a multiple of 4*32=128 or "
-                "less than or equal to 4");
+  static_assert(N <= 64,
+                "Only support candidates number less than or equal to "
+                "64*32=2048");
   using RedType = TopKRedType<Type>;
+  RedType topK[N];
+#pragma unroll
+  for (int nn = 0; nn < N; ++nn) {
+    topK[nn] = RedType{value[nn], idx[nn]};
+  }
 
-  if constexpr (N <= 4) {
-    reduceTopKFunc<K, Type, N>(warp, out, outIdx, value, idx, minValue,
-                               actualK);
-  } else {
-    constexpr int numLoops = N / 4;
-    constexpr int numResults = (numLoops * K - 1) / kWARP_SIZE + 1;
+  Sort<N, RedType>::run(topK);
 
-    Type topKBufferValue[numResults];
-    int32_t topKBufferIdx[numResults];
-    int32_t laneIdx = threadIdx.x % kWARP_SIZE;
-
-    for (int ii = 0; ii < numResults; ++ii) {
-      topKBufferValue[ii] = minValue;
-      topKBufferIdx[ii] = ii * kWARP_SIZE - 1;
+  typename RedType::TypeCmp packedMax{};
+  for (int kk = 0; kk < actualK; ++kk) {
+    bool update = kk > 0 && packedMax == topK[0].compVal;
+#pragma unroll
+    for (int nn = 0; nn < N; ++nn) {
+      topK[nn] = update && nn == N - 1 ? RedType{minValue, idx[nn]}
+                 : update              ? topK[nn + 1]
+                                       : topK[nn];
     }
-    for (int loop = 0; loop < numLoops; ++loop) {
-      int start = loop * 4;
-      Type topKValue[K];
-      int32_t topKIdx[K];
-      Type inValue[4];
-      int32_t inIdx[4];
-      for (int i = 0; i < 4; ++i) {
-        inValue[i] = value[start + i];
-        inIdx[i] = idx[start + i];
-      }
-      reduceTopKFunc<K, Type, 4>(warp, topKValue, topKIdx, inValue, inIdx,
-                                 minValue, actualK);
-      int inOffset = laneIdx % K;
-      if (laneIdx >= loop * K && laneIdx < (loop + 1) * K) {
-        topKBufferValue[0] = topKValue[inOffset];
-        topKBufferIdx[0] = topKIdx[inOffset];
-      }
-      if (loop == numLoops - 1 && (laneIdx < (numLoops * K - kWARP_SIZE))) {
-        topKBufferValue[1] = topKValue[inOffset];
-        topKBufferIdx[1] = topKIdx[inOffset];
-      }
-    }
-
-    reduceTopKFunc<K, Type, numResults>(warp, out, outIdx, topKBufferValue,
-                                        topKBufferIdx, minValue, actualK);
+    packedMax = topK[0].reduce(warp);
+    RedType::unpack(out[kk], outIdx[kk], packedMax);
   }
 };
 
-#undef TOPK_SWAP
+template <int NumExperts, int NumTopExperts, int MinExperts, int MaxExperts,
+          int MinTopExperts, int MaxTopExperts>
+struct LaneOwnedTopKRange {
+  static_assert(MinExperts > 0 && MinExperts <= MaxExperts);
+  static_assert(MinTopExperts > 0 && MinTopExperts <= MaxTopExperts);
+  static constexpr bool kEnabled =
+      NumExperts >= MinExperts && NumExperts <= MaxExperts &&
+      NumTopExperts >= MinTopExperts && NumTopExperts <= MaxTopExperts;
+};
+
+static constexpr int kHIGH_EXPERT_LANE_OWNED_TOPK_MIN_EXPERTS = 512;
+static constexpr int kHIGH_EXPERT_LANE_OWNED_TOPK_MAX_EXPERTS = 1024;
+static constexpr int kHIGH_EXPERT_LANE_OWNED_TOPK_MIN_TOP_EXPERTS = 9;
+static constexpr int kHIGH_EXPERT_LANE_OWNED_TOPK_MAX_TOP_EXPERTS = 16;
+
+template <int NumExperts, int NumTopExperts>
+using HighExpertLaneOwnedTopKRange =
+    LaneOwnedTopKRange<NumExperts, NumTopExperts,
+                       kHIGH_EXPERT_LANE_OWNED_TOPK_MIN_EXPERTS,
+                       kHIGH_EXPERT_LANE_OWNED_TOPK_MAX_EXPERTS,
+                       kHIGH_EXPERT_LANE_OWNED_TOPK_MIN_TOP_EXPERTS,
+                       kHIGH_EXPERT_LANE_OWNED_TOPK_MAX_TOP_EXPERTS>;
+
+template <int K, typename Type, int N>
+__forceinline__ __device__ void reduceTopKForLane(
+    cg::thread_block_tile<kWARP_SIZE> const& warp, Type& out, int32_t& outIdx,
+    Type (&value)[N], int32_t (&idx)[N], Type const minValue, int32_t laneIdx) {
+  static_assert(K > 0, "Top K must have K > 0");
+  static_assert(K <= kWARP_SIZE, "Top K must have K <= kWARP_SIZE");
+  static_assert(N > 0, "Top K must have N > 0");
+  static_assert(N <= 64,
+                "Only support candidates number less than or equal to "
+                "64*32=2048");
+  using RedType = TopKRedType<Type>;
+  RedType topK[N];
+#pragma unroll
+  for (int nn = 0; nn < N; ++nn) {
+    topK[nn] = RedType{value[nn], idx[nn]};
+  }
+
+  Sort<N, RedType>::run(topK);
+
+  typename RedType::TypeCmp packedMax{};
+  typename RedType::TypeCmp lanePacked{};
+#pragma unroll
+  for (int kk = 0; kk < K; ++kk) {
+    bool update = kk > 0 && packedMax == topK[0].compVal;
+#pragma unroll
+    for (int nn = 0; nn < N; ++nn) {
+      topK[nn] = update && nn == N - 1 ? RedType{minValue, idx[nn]}
+                 : update              ? topK[nn + 1]
+                                       : topK[nn];
+    }
+    packedMax = topK[0].reduce(warp);
+    if (laneIdx == kk) {
+      lanePacked = packedMax;
+    }
+  }
+
+  if (laneIdx < K) {
+    RedType::unpack(out, outIdx, lanePacked);
+  } else {
+    out = minValue;
+    outIdx = -1;
+  }
+}
 
 }  // namespace reduce_topk
 }  // namespace moe

csrc/libtorch_stable/moe/moe_align_sum_kernels.cu

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

csrc/libtorch_stable/ops.h

@@ -276,6 +276,61 @@ void fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_bf16_insert(
     torch::stable::Tensor const& cos_sin_cache, double eps,
     int64_t cache_block_size);
 
+void fused_kimi_k3_mla_key_concat_kv_cache_insert(
+    torch::stable::Tensor& q, torch::stable::Tensor const& k_nope,
+    torch::stable::Tensor const& k_pe, torch::stable::Tensor const& kv_c_normed,
+    torch::stable::Tensor& k_out, torch::stable::Tensor& k_cache,
+    torch::stable::Tensor const& slot_mapping, int64_t cache_block_size,
+    std::optional<torch::stable::Tensor> position_ids,
+    std::optional<torch::stable::Tensor> cos_sin_cache);
+
+void fused_kimi_k3_mla_key_concat_ds_mla_insert(
+    torch::stable::Tensor& q, torch::stable::Tensor const& k_nope,
+    torch::stable::Tensor const& k_pe, torch::stable::Tensor const& kv_c_normed,
+    torch::stable::Tensor& k_out, torch::stable::Tensor& k_cache,
+    torch::stable::Tensor const& slot_mapping, int64_t cache_block_size,
+    std::optional<torch::stable::Tensor> position_ids,
+    std::optional<torch::stable::Tensor> cos_sin_cache);
+
+void fused_kimi_k3_mla_qkv_quant_kv_cache_fp8_insert(
+    torch::stable::Tensor const& q, torch::stable::Tensor const& k_nope,
+    torch::stable::Tensor const& k_pe, torch::stable::Tensor const& kv_c_normed,
+    torch::stable::Tensor const& v, torch::stable::Tensor& q_fp8,
+    torch::stable::Tensor& k_fp8, torch::stable::Tensor& v_fp8,
+    torch::stable::Tensor& k_cache, torch::stable::Tensor const& slot_mapping,
+    torch::stable::Tensor const& q_scale_inv,
+    torch::stable::Tensor const& k_scale_inv,
+    torch::stable::Tensor const& v_scale_inv,
+    torch::stable::Tensor const& cache_scale_inv, int64_t cache_block_size,
+    std::optional<torch::stable::Tensor> position_ids,
+    std::optional<torch::stable::Tensor> cos_sin_cache);
+
+void fused_kimi_k3_mla_decode_q_concat_kv_cache_insert(
+    torch::stable::Tensor const& ql_nope, torch::stable::Tensor const& q_pe,
+    torch::stable::Tensor const& kv_c_normed, torch::stable::Tensor const& k_pe,
+    torch::stable::Tensor& mqa_q, torch::stable::Tensor& k_cache,
+    torch::stable::Tensor const& slot_mapping, int64_t cache_block_size,
+    std::optional<torch::stable::Tensor> position_ids,
+    std::optional<torch::stable::Tensor> cos_sin_cache);
+
+void fused_kimi_k3_mla_decode_q_concat_kv_cache_fp8_insert(
+    torch::stable::Tensor const& ql_nope, torch::stable::Tensor const& q_pe,
+    torch::stable::Tensor const& kv_c_normed, torch::stable::Tensor const& k_pe,
+    torch::stable::Tensor& mqa_q, torch::stable::Tensor& k_cache,
+    torch::stable::Tensor const& slot_mapping,
+    torch::stable::Tensor const& q_scale_inv,
+    torch::stable::Tensor const& cache_scale_inv, int64_t cache_block_size,
+    std::optional<torch::stable::Tensor> position_ids,
+    std::optional<torch::stable::Tensor> cos_sin_cache);
+
+void fused_kimi_k3_mla_decode_q_concat_ds_mla_insert(
+    torch::stable::Tensor const& ql_nope, torch::stable::Tensor const& q_pe,
+    torch::stable::Tensor const& kv_c_normed, torch::stable::Tensor const& k_pe,
+    torch::stable::Tensor& mqa_q, torch::stable::Tensor& k_cache,
+    torch::stable::Tensor const& slot_mapping, int64_t cache_block_size,
+    std::optional<torch::stable::Tensor> position_ids,
+    std::optional<torch::stable::Tensor> cos_sin_cache);
+
 void fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_fp8_insert(
     torch::stable::Tensor const& q, torch::stable::Tensor const& kv,
     torch::stable::Tensor& q_fp8, torch::stable::Tensor& k_cache,
@@ -315,6 +370,19 @@ void fused_minimax_m3_qknorm_rope_kv_insert(
     std::optional<torch::stable::Tensor> index_q_out,
     const std::string& kv_cache_dtype, bool skip_index_branch);
 
+#ifdef VLLM_ENABLE_FUSED_KDA_DECODE
+void fused_kda_decode(
+    torch::stable::Tensor const& x, torch::stable::Tensor const& weight,
+    std::optional<torch::stable::Tensor> bias,
+    torch::stable::Tensor& conv_state, torch::stable::Tensor const& raw_g,
+    torch::stable::Tensor const& raw_beta, torch::stable::Tensor const& a_log,
+    torch::stable::Tensor const& dt_bias,
+    torch::stable::Tensor const& state_indices, torch::stable::Tensor& state,
+    torch::stable::Tensor& out, std::optional<double> lower_bound,
+    std::optional<torch::stable::Tensor> output_gate,
+    std::optional<torch::stable::Tensor> norm_weight, double norm_eps);
+#endif
+
 #ifdef VLLM_ENABLE_KIMI_K3_ATTN_RES
 void kimi_k3_attn_res(torch::stable::Tensor& prefix,
                       torch::stable::Tensor const& delta,
@@ -383,6 +451,20 @@ fptr_t init_custom_ar(const std::vector<int64_t>& fake_ipc_ptrs,
 void all_reduce(fptr_t _fa, torch::stable::Tensor& inp,
                 torch::stable::Tensor& out, fptr_t reg_buffer,
                 int64_t reg_buffer_sz_bytes);
+void custom_all_gather(fptr_t _fa, torch::stable::Tensor& inp,
+                       torch::stable::Tensor& out, fptr_t reg_buffer,
+                       int64_t reg_buffer_sz_bytes);
+void mnnvl_lamport_all_gather(fptr_t _fa, torch::stable::Tensor& inp,
+                              torch::stable::Tensor& out, fptr_t local_buffer,
+                              fptr_t multicast_buffer, fptr_t epoch_buffer,
+                              int64_t stage_sz_bytes);
+void custom_reduce_scatter(fptr_t _fa, torch::stable::Tensor& inp,
+                           torch::stable::Tensor& out, fptr_t reg_buffer,
+                           int64_t reg_buffer_sz_bytes);
+void mnnvl_lamport_reduce_scatter(fptr_t _fa, torch::stable::Tensor& inp,
+                                  torch::stable::Tensor& out,
+                                  fptr_t local_buffer, fptr_t epoch_buffer,
+                                  int64_t stage_sz_bytes);
 void dispose(fptr_t _fa);
 int64_t meta_size();
 void register_buffer(fptr_t _fa, const std::vector<int64_t>& fake_ipc_ptrs);
@@ -421,6 +503,12 @@ void fatrelu_and_mul(torch::stable::Tensor& out, torch::stable::Tensor& input,
                      double threshold);
 void swigluoai_and_mul(torch::stable::Tensor& out, torch::stable::Tensor& input,
                        double alpha = 1.702, double limit = 7.0);
+void situ_and_mul(torch::stable::Tensor& out, torch::stable::Tensor& input,
+                  double beta = 1.0, double linear_beta = -1.0);
+void masked_situ_and_mul(torch::stable::Tensor& out,
+                         torch::stable::Tensor& input,
+                         const torch::stable::Tensor& expert_num_tokens,
+                         double beta = 1.0, double linear_beta = -1.0);
 void gelu_new(torch::stable::Tensor& out, torch::stable::Tensor& input);
 void gelu_fast(torch::stable::Tensor& out, torch::stable::Tensor& input);
 void gelu_quick(torch::stable::Tensor& out, torch::stable::Tensor& input);
@@ -497,6 +585,13 @@ void concat_and_cache_mla(torch::stable::Tensor& kv_c,
                           const std::string& kv_cache_dtype,
                           torch::stable::Tensor& scale);
 
+void concat_and_cache_mla_grouped(torch::stable::Tensor& kv_c,
+                                  torch::stable::Tensor& k_pe,
+                                  torch::stable::Tensor& kv_cache_ptrs,
+                                  torch::stable::Tensor& slot_mapping,
+                                  int64_t block_size, int64_t block_stride,
+                                  int64_t entry_stride);
+
 // NOTE: k_pe and kv_c order is flipped compared to concat_and_cache_mla
 void concat_and_cache_mla_rope_fused(
     torch::stable::Tensor& positions, torch::stable::Tensor& q_pe,

csrc/libtorch_stable/torch_bindings.cpp

@@ -324,7 +324,8 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
   // DeepSeek V3 fused A GEMM (SM 9.0+, bf16 only, 1-16 tokens).
   // conditionally compiled so impl registration is in source file
   ops.def(
-      "dsv3_fused_a_gemm(Tensor! output, Tensor mat_a, Tensor mat_b) -> ()");
+      "dsv3_fused_a_gemm(Tensor! output, Tensor mat_a, Tensor mat_b, "
+      "bool enable_pdl=False) -> ()");
 
   // BF16/FP32 x FP32 -> FP32 router GEMM for H=3072, E=256, M<=32 (SM90+).
   // conditionally compiled so impl registration is in source file
@@ -447,6 +448,48 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
       "Tensor fp8_scale, Tensor q_fp8_scale_inv, float eps, "
       "int cache_block_size) -> ()");
 
+  // Kimi-K3 MLA epilogues: optional RoPE followed by concat/cache insertion.
+  ops.def(
+      "fused_kimi_k3_mla_key_concat_kv_cache_insert("
+      "Tensor! q, Tensor k_nope, Tensor k_pe, Tensor kv_c_normed, "
+      "Tensor! k_out, Tensor! k_cache, Tensor slot_mapping, "
+      "int cache_block_size, Tensor? position_ids=None, "
+      "Tensor? cos_sin_cache=None) -> ()");
+  ops.def(
+      "fused_kimi_k3_mla_key_concat_ds_mla_insert("
+      "Tensor! q, Tensor k_nope, Tensor k_pe, Tensor kv_c_normed, "
+      "Tensor! k_out, Tensor! k_cache, Tensor slot_mapping, "
+      "int cache_block_size, Tensor? position_ids=None, "
+      "Tensor? cos_sin_cache=None) -> ()");
+  ops.def(
+      "fused_kimi_k3_mla_qkv_quant_kv_cache_fp8_insert("
+      "Tensor q, Tensor k_nope, Tensor k_pe, Tensor kv_c_normed, Tensor v, "
+      "Tensor! q_fp8, Tensor! k_fp8, Tensor! v_fp8, Tensor! k_cache, "
+      "Tensor slot_mapping, Tensor q_scale_inv, Tensor k_scale_inv, "
+      "Tensor v_scale_inv, Tensor cache_scale_inv, int cache_block_size, "
+      "Tensor? position_ids=None, Tensor? cos_sin_cache=None) -> ()");
+
+  // Kimi-K3 MLA decode epilogue: concat mqa_q = [ql_nope | q_pe] and insert the
+  // latent [kv_c_normed | k_pe] into the paged cache (bf16 / fp8 / fp8_ds_mla).
+  ops.def(
+      "fused_kimi_k3_mla_decode_q_concat_kv_cache_insert("
+      "Tensor ql_nope, Tensor q_pe, Tensor kv_c_normed, Tensor k_pe, "
+      "Tensor! mqa_q, Tensor! k_cache, Tensor slot_mapping, "
+      "int cache_block_size, Tensor? position_ids=None, "
+      "Tensor? cos_sin_cache=None) -> ()");
+  ops.def(
+      "fused_kimi_k3_mla_decode_q_concat_kv_cache_fp8_insert("
+      "Tensor ql_nope, Tensor q_pe, Tensor kv_c_normed, Tensor k_pe, "
+      "Tensor! mqa_q, Tensor! k_cache, Tensor slot_mapping, "
+      "Tensor q_scale_inv, Tensor cache_scale_inv, int cache_block_size, "
+      "Tensor? position_ids=None, Tensor? cos_sin_cache=None) -> ()");
+  ops.def(
+      "fused_kimi_k3_mla_decode_q_concat_ds_mla_insert("
+      "Tensor ql_nope, Tensor q_pe, Tensor kv_c_normed, Tensor k_pe, "
+      "Tensor! mqa_q, Tensor! k_cache, Tensor slot_mapping, "
+      "int cache_block_size, Tensor? position_ids=None, "
+      "Tensor? cos_sin_cache=None) -> ()");
+
 #ifndef USE_ROCM
   ops.def(
       "minimax_allreduce_rms_qk("
@@ -468,6 +511,16 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
       "int block_size, Tensor!? q_out, Tensor!? index_q_out, "
       "str kv_cache_dtype, bool skip_index_branch=False) -> ()");
 
+#ifdef VLLM_ENABLE_FUSED_KDA_DECODE
+  ops.def(
+      "fused_kda_decode("
+      "Tensor x, Tensor weight, Tensor? bias, Tensor! conv_state, "
+      "Tensor raw_g, Tensor raw_beta, Tensor A_log, Tensor dt_bias, "
+      "Tensor state_indices, Tensor! state, Tensor! out, "
+      "float? lower_bound=None, Tensor? output_gate=None, "
+      "Tensor? norm_weight=None, float norm_eps=1e-5) -> ()");
+#endif
+
 #ifdef VLLM_ENABLE_KIMI_K3_ATTN_RES
   ops.def(
       "kimi_k3_attn_res("
@@ -538,6 +591,14 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
       "limit=7.0) "
       "-> ()");
 
+  // SituGLU implementation used in Kimi models.
+  ops.def(
+      "situ_and_mul(Tensor! out, Tensor input, float beta=1.0, float "
+      "linear_beta=-1.0) -> ()");
+  ops.def(
+      "masked_situ_and_mul(Tensor! out, Tensor input, Tensor "
+      "expert_num_tokens, float beta=1.0, float linear_beta=-1.0) -> ()");
+
   // GELU implementation used in GPT-2.
   ops.def("gelu_new(Tensor! out, Tensor input) -> ()");
 
@@ -696,11 +757,27 @@ STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) {
   ops.impl(
       "fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_fp8_insert",
       TORCH_BOX(&fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_fp8_insert));
+  ops.impl("fused_kimi_k3_mla_key_concat_kv_cache_insert",
+           TORCH_BOX(&fused_kimi_k3_mla_key_concat_kv_cache_insert));
+  ops.impl("fused_kimi_k3_mla_key_concat_ds_mla_insert",
+           TORCH_BOX(&fused_kimi_k3_mla_key_concat_ds_mla_insert));
+  ops.impl("fused_kimi_k3_mla_qkv_quant_kv_cache_fp8_insert",
+           TORCH_BOX(&fused_kimi_k3_mla_qkv_quant_kv_cache_fp8_insert));
+  ops.impl("fused_kimi_k3_mla_decode_q_concat_kv_cache_insert",
+           TORCH_BOX(&fused_kimi_k3_mla_decode_q_concat_kv_cache_insert));
+  ops.impl("fused_kimi_k3_mla_decode_q_concat_kv_cache_fp8_insert",
+           TORCH_BOX(&fused_kimi_k3_mla_decode_q_concat_kv_cache_fp8_insert));
+  ops.impl("fused_kimi_k3_mla_decode_q_concat_ds_mla_insert",
+           TORCH_BOX(&fused_kimi_k3_mla_decode_q_concat_ds_mla_insert));
 #ifndef USE_ROCM
   ops.impl("minimax_allreduce_rms_qk", TORCH_BOX(&minimax_allreduce_rms_qk));
 #endif
   ops.impl("fused_minimax_m3_qknorm_rope_kv_insert",
            TORCH_BOX(&fused_minimax_m3_qknorm_rope_kv_insert));
+#ifdef VLLM_ENABLE_FUSED_KDA_DECODE
+  ops.impl("fused_kda_decode", TORCH_BOX(&fused_kda_decode));
+#endif
+
 #ifdef VLLM_ENABLE_KIMI_K3_ATTN_RES
   ops.impl("kimi_k3_attn_res", TORCH_BOX(&kimi_k3_attn_res));
 #endif
@@ -726,6 +803,8 @@ STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) {
   ops.impl("gelu_tanh_and_mul", TORCH_BOX(&gelu_tanh_and_mul));
   ops.impl("fatrelu_and_mul", TORCH_BOX(&fatrelu_and_mul));
   ops.impl("swigluoai_and_mul", TORCH_BOX(&swigluoai_and_mul));
+  ops.impl("situ_and_mul", TORCH_BOX(&situ_and_mul));
+  ops.impl("masked_situ_and_mul", TORCH_BOX(&masked_situ_and_mul));
   ops.impl("gelu_new", TORCH_BOX(&gelu_new));
   ops.impl("gelu_fast", TORCH_BOX(&gelu_fast));
   ops.impl("gelu_quick", TORCH_BOX(&gelu_quick));
@@ -823,6 +902,15 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C_cache_ops, ops) {
       "                     str kv_cache_dtype,"
       "                     Tensor scale) -> ()");
 
+  // Grouped concat_and_cache_mla across all layers (bf16 only). Each
+  // layer's cache base pointer is read from kv_cache_ptrs.
+  ops.def(
+      "concat_and_cache_mla_grouped(Tensor kv_c, Tensor k_pe,"
+      "                             Tensor kv_cache_ptrs,"
+      "                             Tensor slot_mapping,"
+      "                             int block_size, int block_stride,"
+      "                             int entry_stride) -> ()");
+
   // Rotate Q and K, then write to kv cache for MLA
   ops.def(
       "concat_and_cache_mla_rope_fused("
@@ -921,6 +1009,8 @@ STABLE_TORCH_LIBRARY_IMPL(_C_cache_ops, CUDA, ops) {
   ops.impl("reshape_and_cache", TORCH_BOX(&reshape_and_cache));
   ops.impl("reshape_and_cache_flash", TORCH_BOX(&reshape_and_cache_flash));
   ops.impl("concat_and_cache_mla", TORCH_BOX(&concat_and_cache_mla));
+  ops.impl("concat_and_cache_mla_grouped",
+           TORCH_BOX(&concat_and_cache_mla_grouped));
   ops.impl("concat_and_cache_mla_rope_fused",
            TORCH_BOX(&concat_and_cache_mla_rope_fused));
   ops.impl("convert_fp8", TORCH_BOX(&convert_fp8));

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

@@ -13,10 +13,18 @@ namespace vllm {
 namespace fp8 {
   #ifdef ENABLE_FP8
 
+// Unspecialized conversions are a compile error: the old passthrough
+// (`return x;`) silently skipped fp8 encoding for any (Tout, Tin) pair
+// without a specialization below (e.g. the torch stable-ABI scalar types),
+// corrupting quantized data with no runtime signal.
+template <typename>
+inline constexpr bool _no_conversion_specialization = false;
+
 template <typename Tout, typename Tin>
 __inline__ __device__ Tout vec_conversion(
     const Tin& x, const __nv_fp8_interpretation_t fp8_type = __NV_E4M3) {
-  return x;
+  static_assert(_no_conversion_specialization<Tin>,
+                "no vec_conversion specialization for this (Tout, Tin) pair");
 }
 
 // float -> c10::Float8_e4m3fn
@@ -301,7 +309,9 @@ __inline__ __device__ bf16_8_t vec_conversion<bf16_8_t, Float8_>(
 template <typename Tout, typename Tin>
 __inline__ __device__ Tout scaled_vec_conversion(
     const Tin& x, const float scale, const __nv_fp8_interpretation_t fp8_type) {
-  return x;
+  static_assert(
+      _no_conversion_specialization<Tin>,
+      "no scaled_vec_conversion specialization for this (Tout, Tin) pair");
 }
 
 // fp8 -> half
@@ -492,6 +502,25 @@ __inline__ __device__ uint8_t scaled_vec_conversion<uint8_t, __nv_bfloat16>(
   __builtin_unreachable();  // Suppress missing return statement warning
 }
 
+// torch stable-ABI (headeronly) scalar types delegate to the CUDA-native
+// conversions, so libtorch_stable kernels dispatched on c10::BFloat16 /
+// c10::Half quantize correctly without manual casts.
+template <>
+__inline__ __device__ uint8_t scaled_vec_conversion<uint8_t, c10::BFloat16>(
+    const c10::BFloat16& a, const float scale,
+    const __nv_fp8_interpretation_t fp8_type) {
+  return scaled_vec_conversion<uint8_t, __nv_bfloat16>(
+      reinterpret_cast<const __nv_bfloat16&>(a), scale, fp8_type);
+}
+
+template <>
+__inline__ __device__ uint8_t scaled_vec_conversion<uint8_t, c10::Half>(
+    const c10::Half& a, const float scale,
+    const __nv_fp8_interpretation_t fp8_type) {
+  return scaled_vec_conversion<uint8_t, uint16_t>(
+      reinterpret_cast<const uint16_t&>(a), scale, fp8_type);
+}
+
 // float -> fp8
 template <>
 __inline__ __device__ uint8_t scaled_vec_conversion<uint8_t, float>(

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/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,649 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Tests for the Kimi-K3 SM103 decode GEMM selector (shape-only dispatch)."""
+
+from pathlib import Path
+from types import SimpleNamespace
+from typing import Any
+
+import pytest
+import regex as re
+import torch
+from torch import nn
+
+from vllm.model_executor.kernels.linear.cute_dsl.skinny_gemm import SkinnyGemmConfig
+from vllm.models.kimi_k3.nvidia import low_latency_gemm as k3_gemm
+from vllm.models.kimi_k3.nvidia.low_latency_gemm import KIMI_K3_PROJECTIONS
+
+# Keyed by local (N, K): (cute token counts, dsv3 token counts). 1536x7168 is
+# the unified shared_gate_up_proj/mla_g_proj entry (dsv3 M1..16).
+EXPECTED_SELECTIONS = {
+    (1536, 128): (set(), set(range(1, 17))),
+    (3072, 128): (set(), set(range(1, 17))),
+    (1536, 7168): (set(), set(range(1, 17))),
+    (3072, 7168): (set(range(1, 6)), set()),
+    (2112, 7168): (set(), set(range(1, 17))),
+    (2304, 1536): (set(), set(range(1, 17))),
+    (4608, 1536): (set(), set(range(1, 17))),
+    (3584, 7168): ({1}, set(range(2, 9))),
+    (6288, 7168): (set(range(1, 5)), set()),
+    (12448, 7168): (set(range(1, 4)), set()),
+    (7168, 768): (set(), set(range(1, 17))),
+    (7168, 1536): ({1}, set()),
+    (7168, 3072): ({1, 2}, set()),
+    (7168, 3584): ({1, 2}, set()),
+    (7168, 4224): ({1}, set()),
+    (7168, 8448): (set(range(1, 4)), set()),
+    (8448, 7168): ({1, 2}, set()),
+    (16896, 7168): ({1, 2}, set()),
+    (20480, 7168): (set(range(1, 5)), set()),
+    (40960, 7168): (set(range(1, 5)), set()),
+    # TP16.
+    (3216, 7168): (set(range(1, 6)), set(range(9, 16))),
+    (768, 7168): (set(range(1, 5)), set(range(5, 17))),
+    (1152, 1536): ({1}, set(range(2, 17))),
+    (768, 128): (set(), set(range(1, 17))),
+    (7168, 384): (set(), set(range(1, 9))),
+    (4224, 7168): (set(range(1, 4)), set(range(4, 9))),
+    (10240, 7168): (set(range(1, 5)), set()),
+}
+
+CUTE_CASES = [
+    (spec.n, spec.k, num_tokens)
+    for spec in k3_gemm.KIMI_K3_PROJECTIONS.values()
+    for num_tokens, _ in spec.cute_configs
+]
+
+RESIDUAL_CUTE_CASES = [
+    (spec.n, spec.k, num_tokens)
+    for spec in k3_gemm.KIMI_K3_PROJECTIONS.values()
+    for num_tokens, _ in spec.residual_configs
+]
+
+EXPECTED_CUTE_CONFIGS = {
+    (3072, 7168, 1): (224, 3, 4, 8),
+    (3072, 7168, 2): (128, 3, 2, 8),
+    (3072, 7168, 3): (128, 2, 1, 8),
+    (3072, 7168, 4): (64, 2, 2, 8),
+    (3072, 7168, 5): (128, 3, 1, 8),
+    (3584, 7168, 1): (224, 2, 4, 8),
+    (6288, 7168, 1): (224, 3, 4, 8),
+    (6288, 7168, 2): (64, 3, 2, 8),
+    (6288, 7168, 3): (32, 3, 4, 8),
+    (6288, 7168, 4): (128, 6, 1, 8),
+    (12448, 7168, 1): (224, 4, 2, 8),
+    (12448, 7168, 2): (64, 4, 2, 8),
+    (12448, 7168, 3): (64, 2, 2, 8),
+    (7168, 1536, 1): (96, 4, 2, 8),
+    (7168, 3072, 1): (96, 2, 4, 8),
+    (7168, 3072, 2): (32, 4, 4, 8),
+    (7168, 3584, 1): (224, 4, 2, 8),
+    (7168, 3584, 2): (64, 4, 2, 8),
+    (7168, 4224, 1): (96, 4, 2, 4),
+    (7168, 8448, 1): (32, 4, 4, 8),
+    (7168, 8448, 2): (96, 4, 1, 8),
+    (7168, 8448, 3): (96, 4, 1, 8),
+    (8448, 7168, 1): (224, 3, 4, 8),
+    (8448, 7168, 2): (32, 4, 4, 8),
+    (16896, 7168, 1): (224, 6, 4, 8),
+    (16896, 7168, 2): (32, 4, 4, 8),
+    (20480, 7168, 1): (224, 4, 2, 8),
+    (20480, 7168, 2): (64, 4, 2, 8),
+    (20480, 7168, 3): (64, 2, 2, 8),
+    (20480, 7168, 4): (64, 4, 1, 8),
+    (40960, 7168, 1): (128, 4, 2, 8),
+    (40960, 7168, 2): (64, 4, 2, 8),
+    (40960, 7168, 3): (64, 2, 2, 8),
+    (40960, 7168, 4): (64, 4, 1, 8),
+    # TP16.
+    (3216, 7168, 1): (224, 3, 4, 8),
+    (3216, 7168, 2): (128, 4, 2, 8),
+    (3216, 7168, 3): (128, 2, 1, 8),
+    (3216, 7168, 4): (64, 2, 2, 8),
+    (3216, 7168, 5): (128, 3, 1, 8),
+    (768, 7168, 1): (224, 2, 4, 8),
+    (768, 7168, 2): (224, 2, 2, 8),
+    (768, 7168, 3): (224, 2, 2, 8),
+    (768, 7168, 4): (224, 2, 2, 8),
+    (1152, 1536, 1): (192, 3, 4, 8),
+    (4224, 7168, 1): (224, 3, 4, 8),
+    (4224, 7168, 2): (128, 2, 1, 8),
+    (4224, 7168, 3): (64, 2, 2, 8),
+    (10240, 7168, 1): (224, 4, 2, 8),
+    (10240, 7168, 2): (32, 2, 4, 8),
+    (10240, 7168, 3): (64, 4, 1, 8),
+    (10240, 7168, 4): (64, 4, 1, 8),
+}
+
+EXPECTED_RESIDUAL_CUTE_CONFIGS = {
+    (7168, 3584, 1): (64, 4, 2, 8),
+    (7168, 3584, 2): (64, 7, 2, 8),
+    (7168, 3584, 3): (64, 2, 1, 8),
+    (7168, 3584, 4): (64, 2, 1, 8),
+}
+
+
+def _config_tuple(config) -> tuple[int, int, int, int]:
+    return (
+        config.block_size,
+        config.outputs_per_block,
+        config.k_unroll,
+        config.vector_width,
+    )
+
+
+def test_table_is_keyed_by_shape() -> None:
+    for (n, k), spec in k3_gemm.KIMI_K3_PROJECTIONS.items():
+        assert (spec.n, spec.k) == (n, k)
+
+
+def test_every_dsv3_routed_shape_is_instantiated() -> None:
+    """dsv3_fused_a_gemm specializes on (K, N); an unlisted shape raises.
+
+    The table routes by shape while the kernel is built per shape, so a missing
+    instantiation only shows up at the token counts that route to dsv3. Checking
+    it here needs no GPU, which is the point -- a GPU-only check is exactly what
+    let (3216, 7168) ship without its DISPATCH_DSV3_SHAPE(7168, 3216).
+    """
+    source = (
+        Path(__file__).resolve().parents[2]
+        / "csrc"
+        / "libtorch_stable"
+        / "dsv3_fused_a_gemm.cu"
+    ).read_text(encoding="utf-8")
+    # Benchmark-only shapes live behind VLLM_K3_BENCH_SHAPES and are not built
+    # by default, so they must not count as available.
+    production_macros = source.split("#ifdef VLLM_K3_BENCH_SHAPES")[0]
+    explicit = source.split("#undef DISPATCH_DSV3_SHAPE")[1].split(
+        "#ifdef VLLM_K3_BENCH_SHAPES"
+    )[0]
+    compiled = {
+        (int(hd_in), int(hd_out))
+        for hd_in, hd_out in re.findall(
+            r"DISPATCH_DSV3_SHAPE\((\d+),\s*(\d+)\)", production_macros
+        )
+    } | {
+        (int(hd_in), int(hd_out))
+        for hd_in, hd_out in re.findall(r"hd_in == (\d+) && hd_out == (\d+)", explicit)
+    }
+    assert compiled, "failed to parse the dispatch list"
+
+    missing = sorted(
+        (spec.n, spec.k)
+        for spec in KIMI_K3_PROJECTIONS.values()
+        if spec.dsv3_tokens and (spec.k, spec.n) not in compiled
+    )
+    assert not missing, (
+        f"routed to dsv3 with no instantiation: {missing}; add "
+        "DISPATCH_DSV3_SHAPE(K, N) for each"
+    )
+
+
+def test_packed_row_major_rejects_single_row_slice() -> None:
+    packed = torch.empty(1, 128)
+    sliced = torch.empty(1, 144)[:, :128]
+
+    assert packed.is_contiguous()
+    assert sliced.is_contiguous()
+    assert k3_gemm._is_packed_row_major(packed)
+    assert not k3_gemm._is_packed_row_major(sliced)
+
+
+def test_cute_configs_match_measured_table() -> None:
+    actual = {
+        (spec.n, spec.k, num_tokens): _config_tuple(config)
+        for spec in k3_gemm.KIMI_K3_PROJECTIONS.values()
+        for num_tokens, config in spec.cute_configs
+    }
+    assert actual == EXPECTED_CUTE_CONFIGS
+
+
+def test_residual_cute_configs_match_measured_table() -> None:
+    actual = {
+        (spec.n, spec.k, num_tokens): _config_tuple(config)
+        for spec in k3_gemm.KIMI_K3_PROJECTIONS.values()
+        for num_tokens, config in spec.residual_configs
+    }
+    assert actual == EXPECTED_RESIDUAL_CUTE_CONFIGS
+
+
+@pytest.mark.parametrize("key", EXPECTED_SELECTIONS)
+def test_sm103_selector_table(key: tuple[int, int]) -> None:
+    n, k = key
+    cute_tokens, dsv3_tokens = EXPECTED_SELECTIONS[key]
+    for num_tokens in range(1, 17):
+        backend = k3_gemm.select_kimi_k3_backend(num_tokens, n, k)
+        if num_tokens in cute_tokens:
+            assert backend == "cute"
+        elif num_tokens in dsv3_tokens:
+            assert backend == "dsv3_fused_a"
+        else:
+            assert backend is None
+
+
+@pytest.mark.parametrize("key", EXPECTED_SELECTIONS)
+def test_selector_requires_supported_shape_and_tokens(key: tuple[int, int]) -> None:
+    n, k = key
+    assert k3_gemm.select_kimi_k3_backend(0, n, k) is None
+    assert k3_gemm.select_kimi_k3_backend(17, n, k) is None
+    assert k3_gemm.select_kimi_k3_backend(1, n + 1, k) is None
+    assert k3_gemm.select_kimi_k3_backend(1, n, k + 1) is None
+
+
+def test_unlisted_shape_and_unselected_tokens_fall_back() -> None:
+    # Shape absent from the table.
+    assert k3_gemm.select_kimi_k3_backend(1, 1000, 1000) is None
+    # o_proj (7168,1536) is CuTe M1 only; M2+ falls back.
+    assert k3_gemm.select_kimi_k3_backend(2, 7168, 1536) is None
+
+
+@pytest.mark.parametrize("num_tokens", range(1, 17))
+def test_sm103_residual_selector_table(num_tokens: int) -> None:
+    backend = k3_gemm.select_kimi_k3_backend(num_tokens, 7168, 3584, has_residual=True)
+    assert backend == ("cute" if num_tokens <= 4 else None)
+
+
+def test_build_plan_matches_selector() -> None:
+    for spec in k3_gemm.KIMI_K3_PROJECTIONS.values():
+        plan = k3_gemm._build_plan(spec)
+        for num_tokens in range(1, 17):
+            backend = k3_gemm.select_kimi_k3_backend(num_tokens, spec.n, spec.k)
+            if backend is None:
+                assert num_tokens not in plan
+            else:
+                assert plan[num_tokens][0] == backend
+
+
+def test_installation_is_shape_specific_and_unquantized(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    class FakeLinear(nn.Module):
+        def __init__(self, quant_method: object, n: int, k: int) -> None:
+            super().__init__()
+            self.quant_method = quant_method
+            self.weight = torch.empty(n, k)
+
+    class FakeHead(nn.Module):
+        def __init__(self, n: int, k: int) -> None:
+            super().__init__()
+            self.quant_method = k3_gemm.UnquantizedEmbeddingMethod()
+            self.weight = torch.empty(n, k)
+
+    root = nn.Module()
+    # dsv3-only shape (no cute warmup contribution).
+    root.dsv3_only = FakeLinear(k3_gemm.UnquantizedLinearMethod(), 2304, 1536)
+    # quantized: must be left untouched.
+    quantized_method = object()
+    root.quantized = FakeLinear(quantized_method, 6288, 7168)
+    # cute shape.
+    root.cute = FakeLinear(k3_gemm.UnquantizedLinearMethod(), 6288, 7168)
+    # cute + residual shape.
+    root.residual = FakeLinear(k3_gemm.UnquantizedLinearMethod(), 7168, 3584)
+    # shape absent from the table: must be left untouched.
+    root.unlisted = FakeLinear(k3_gemm.UnquantizedLinearMethod(), 1234, 5678)
+    root.lm_head = FakeHead(20480, 7168)
+
+    monkeypatch.setattr(k3_gemm, "LinearBase", FakeLinear)
+    monkeypatch.setattr(k3_gemm, "ParallelLMHead", FakeHead)
+    monkeypatch.setattr(k3_gemm, "_is_sm103", lambda: True)
+    warmup_configs: set[SkinnyGemmConfig] = set()
+    residual_warmup_configs: set[SkinnyGemmConfig] = set()
+    monkeypatch.setattr(k3_gemm.shape_dynamic_skinny_gemm, "is_available", lambda: True)
+
+    def request_warmup_configs(dtype, configs, *, has_residual=False):
+        target = residual_warmup_configs if has_residual else warmup_configs
+        target.update(configs)
+
+    monkeypatch.setattr(
+        k3_gemm.shape_dynamic_skinny_gemm,
+        "request_warmup_configs",
+        request_warmup_configs,
+    )
+
+    k3_gemm.enable_kimi_k3_low_latency_gemm(root, torch.bfloat16)
+
+    assert isinstance(root.dsv3_only.quant_method, k3_gemm.KimiK3LowLatencyLinearMethod)
+    assert isinstance(root.cute.quant_method, k3_gemm.KimiK3LowLatencyLinearMethod)
+    assert isinstance(root.residual.quant_method, k3_gemm.KimiK3LowLatencyLinearMethod)
+    assert root.quantized.quant_method is quantized_method
+    assert type(root.unlisted.quant_method) is k3_gemm.UnquantizedLinearMethod
+    assert isinstance(
+        root.lm_head.quant_method, k3_gemm.KimiK3LowLatencyEmbeddingMethod
+    )
+    # Warmup covers only the installed modules' local (N, K).
+    assert warmup_configs == {
+        config
+        for key in ((6288, 7168), (7168, 3584), (20480, 7168))
+        for _, config in k3_gemm.KIMI_K3_PROJECTIONS[key].cute_configs
+    }
+    assert residual_warmup_configs == {
+        config
+        for _, config in k3_gemm.KIMI_K3_PROJECTIONS[(7168, 3584)].residual_configs
+    }
+
+
+@pytest.mark.parametrize(
+    "dtype,platform_enabled",
+    [(torch.float16, True), (torch.bfloat16, False)],
+)
+def test_installation_requires_bf16_sm103(
+    monkeypatch: pytest.MonkeyPatch,
+    dtype: torch.dtype,
+    platform_enabled: bool,
+) -> None:
+    class FakeLinear(nn.Module):
+        def __init__(self) -> None:
+            super().__init__()
+            self.quant_method = k3_gemm.UnquantizedLinearMethod()
+            self.weight = torch.empty(2304, 1536)
+
+    root = nn.Module()
+    root.projection = FakeLinear()
+    monkeypatch.setattr(k3_gemm, "LinearBase", FakeLinear)
+    monkeypatch.setattr(k3_gemm, "_is_sm103", lambda: platform_enabled)
+
+    k3_gemm.enable_kimi_k3_low_latency_gemm(root, dtype)
+
+    assert type(root.projection.quant_method) is k3_gemm.UnquantizedLinearMethod
+
+
+def _require_sm103_and_dsv3() -> None:
+    if not torch.cuda.is_available() or torch.cuda.get_device_capability() != (10, 3):
+        pytest.skip("Kimi-K3 production selection requires SM103")
+    if not hasattr(torch.ops._C, "dsv3_fused_a_gemm"):
+        pytest.skip("dsv3_fused_a_gemm was not built")
+
+
+def _require_sm103_and_cute() -> None:
+    if not torch.cuda.is_available() or torch.cuda.get_device_capability() != (10, 3):
+        pytest.skip("Kimi-K3 production selection requires SM103")
+    if not k3_gemm.shape_dynamic_skinny_gemm.is_available():
+        pytest.skip("CuTe DSL is not available")
+
+
+@pytest.mark.parametrize("n,k,num_tokens", CUTE_CASES)
+def test_cute_selected_shapes(n: int, k: int, num_tokens: int) -> None:
+    _require_sm103_and_cute()
+    torch.manual_seed(42)
+    x = torch.randn(num_tokens, k, dtype=torch.bfloat16, device="cuda")
+    weight = torch.randn(n, k, dtype=torch.bfloat16, device="cuda")
+
+    output = k3_gemm.try_low_latency_gemm(x, weight)
+
+    assert output is not None
+    reference = torch.nn.functional.linear(x, weight)
+    cosine = torch.nn.functional.cosine_similarity(
+        output.float().flatten(), reference.float().flatten(), dim=0
+    ).item()
+    assert cosine > 0.999
+
+
+def _dsv3_probe_tokens(tokens: frozenset[int]) -> set[int]:
+    """Extremes, plus both sides of the kernel's num_tokens<=8 tile_n branch."""
+    if not tokens:
+        return set()
+    return {min(tokens), max(tokens)} | ({8, 9} & set(tokens))
+
+
+# Derived from the table rather than hand-listed, so a shape routed to dsv3
+# cannot be added without being exercised here.
+DSV3_CASES = sorted(
+    (num_tokens, spec.n, spec.k)
+    for spec in KIMI_K3_PROJECTIONS.values()
+    for num_tokens in _dsv3_probe_tokens(spec.dsv3_tokens)
+)
+
+
+@pytest.mark.parametrize("num_tokens,n,k", DSV3_CASES)
+def test_dsv3_selected_shapes(num_tokens: int, n: int, k: int) -> None:
+    _require_sm103_and_dsv3()
+    spec = k3_gemm.KIMI_K3_PROJECTIONS[(n, k)]
+    assert num_tokens in spec.dsv3_tokens
+    torch.manual_seed(42)
+    x = torch.randn(num_tokens, k, dtype=torch.bfloat16, device="cuda")
+    weight = torch.randn(n, k, dtype=torch.bfloat16, device="cuda")
+
+    output = k3_gemm.try_low_latency_gemm(x, weight)
+
+    assert output is not None
+    reference = torch.nn.functional.linear(x, weight)
+    cosine = torch.nn.functional.cosine_similarity(
+        output.float().flatten(), reference.float().flatten(), dim=0
+    ).item()
+    assert cosine > 0.999
+
+
+def test_nonpacked_single_token_dsv3_falls_back() -> None:
+    _require_sm103_and_dsv3()
+    n, k = 1536, 128
+    storage = torch.randn(1, k + 16, dtype=torch.bfloat16, device="cuda")
+    x = storage[:, :k]
+    weight = torch.randn(n, k, dtype=torch.bfloat16, device="cuda")
+    spec = k3_gemm.KIMI_K3_PROJECTIONS[(n, k)]
+    method = k3_gemm.KimiK3LowLatencyLinearMethod(
+        k3_gemm._build_plan(spec), k3_gemm._build_residual_plan(spec)
+    )
+
+    assert x.is_contiguous()
+    assert x.stride() == (k + 16, 1)
+    assert not k3_gemm._runtime_ok(x, weight)  # strict guard rejects the view
+    output = method.apply(SimpleNamespace(weight=weight), x)
+
+    reference = torch.nn.functional.linear(x, weight)
+    torch.testing.assert_close(output, reference)
+
+
+def test_selected_kernels_cuda_graph_capture() -> None:
+    _require_sm103_and_cute()
+    _require_sm103_and_dsv3()
+    cute_spec = k3_gemm.KIMI_K3_PROJECTIONS[(6288, 7168)]
+    dsv3_spec = k3_gemm.KIMI_K3_PROJECTIONS[(1536, 128)]
+    cute_x = torch.randn(1, cute_spec.k, dtype=torch.bfloat16, device="cuda")
+    cute_weight = torch.randn(
+        cute_spec.n, cute_spec.k, dtype=torch.bfloat16, device="cuda"
+    )
+    dsv3_x = torch.randn(1, dsv3_spec.k, dtype=torch.bfloat16, device="cuda")
+    dsv3_weight = torch.randn(
+        dsv3_spec.n, dsv3_spec.k, dtype=torch.bfloat16, device="cuda"
+    )
+    k3_gemm.try_low_latency_gemm(cute_x, cute_weight)
+    k3_gemm.try_low_latency_gemm(dsv3_x, dsv3_weight)
+    torch.accelerator.synchronize()
+
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph):
+        cute_output = k3_gemm.try_low_latency_gemm(cute_x, cute_weight)
+        dsv3_output = k3_gemm.try_low_latency_gemm(dsv3_x, dsv3_weight)
+    graph.replay()
+    torch.accelerator.synchronize()
+
+    assert cute_output is not None
+    assert dsv3_output is not None
+    for output, activation, weight in (
+        (cute_output, cute_x, cute_weight),
+        (dsv3_output, dsv3_x, dsv3_weight),
+    ):
+        reference = torch.nn.functional.linear(activation, weight)
+        cosine = torch.nn.functional.cosine_similarity(
+            output.float().flatten(), reference.float().flatten(), dim=0
+        ).item()
+        assert cosine > 0.999
+
+
+@pytest.mark.parametrize("num_tokens", [1, 8, 9, 16])
+def test_dsv3_cuda_graph_capture_tile_branches(num_tokens: int) -> None:
+    """Capture DSV3 across the num_tokens<=8 vs >8 tile_n branch."""
+    _require_sm103_and_dsv3()
+    spec = k3_gemm.KIMI_K3_PROJECTIONS[(1536, 128)]
+    x = torch.randn(num_tokens, spec.k, dtype=torch.bfloat16, device="cuda")
+    weight = torch.randn(spec.n, spec.k, dtype=torch.bfloat16, device="cuda")
+    k3_gemm.try_low_latency_gemm(x, weight)
+    torch.accelerator.synchronize()
+
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph):
+        output = k3_gemm.try_low_latency_gemm(x, weight)
+    graph.replay()
+    torch.accelerator.synchronize()
+
+    assert output is not None
+    reference = torch.nn.functional.linear(x, weight)
+    cosine = torch.nn.functional.cosine_similarity(
+        output.float().flatten(), reference.float().flatten(), dim=0
+    ).item()
+    assert cosine > 0.999
+
+
+@pytest.mark.parametrize("n,k,num_tokens", RESIDUAL_CUTE_CASES)
+def test_cute_residual_epilogue(n: int, k: int, num_tokens: int) -> None:
+    _require_sm103_and_cute()
+    torch.manual_seed(42 + num_tokens)
+    x = torch.randn(num_tokens, k, dtype=torch.bfloat16, device="cuda")
+    weight = torch.randn(n, k, dtype=torch.bfloat16, device="cuda")
+    residual = torch.randn(num_tokens, n, dtype=torch.bfloat16, device="cuda")
+    spec = k3_gemm.KIMI_K3_PROJECTIONS[(n, k)]
+    config = spec.residual_config(num_tokens)
+    assert config is not None
+
+    output = k3_gemm.shape_dynamic_skinny_gemm(x, weight, config, residual)
+
+    reference = x.float() @ weight.float().t() + residual.float()
+    cosine = torch.nn.functional.cosine_similarity(
+        output.float().flatten(), reference.flatten(), dim=0
+    ).item()
+    assert cosine > 0.999
+
+
+@pytest.mark.parametrize("num_tokens", range(1, 17))
+def test_cute_residual_epilogue_all_supported_token_counts(num_tokens: int) -> None:
+    _require_sm103_and_cute()
+    from vllm.model_executor.kernels.linear.cute_dsl.skinny_gemm import (
+        ShapeDynamicSkinnyGemm,
+    )
+
+    n, k = 64, 512
+    x = torch.randn(num_tokens, k, dtype=torch.bfloat16, device="cuda")
+    weight = torch.randn(n, k, dtype=torch.bfloat16, device="cuda")
+    residual = torch.randn(num_tokens, n, dtype=torch.bfloat16, device="cuda")
+    config = ShapeDynamicSkinnyGemm._config(num_tokens, n, k)
+
+    output = k3_gemm.shape_dynamic_skinny_gemm(x, weight, config, residual)
+
+    reference = x.float() @ weight.float().t() + residual.float()
+    torch.testing.assert_close(output.float(), reference, rtol=2e-2, atol=2e-1)
+
+
+@pytest.mark.parametrize("num_tokens", range(1, 5))
+def test_cute_residual_epilogue_cuda_graph_capture(num_tokens: int) -> None:
+    _require_sm103_and_cute()
+    spec = k3_gemm.KIMI_K3_PROJECTIONS[(7168, 3584)]
+    config = spec.residual_config(num_tokens)
+    assert config is not None
+    x = torch.randn(num_tokens, spec.k, dtype=torch.bfloat16, device="cuda")
+    weight = torch.randn(spec.n, spec.k, dtype=torch.bfloat16, device="cuda")
+    residual = torch.randn(num_tokens, spec.n, dtype=torch.bfloat16, device="cuda")
+    k3_gemm.shape_dynamic_skinny_gemm(x, weight, config, residual)
+    torch.accelerator.synchronize()
+
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph):
+        output = k3_gemm.shape_dynamic_skinny_gemm(x, weight, config, residual)
+    graph.replay()
+    torch.accelerator.synchronize()
+
+    reference = x.float() @ weight.float().t() + residual.float()
+    cosine = torch.nn.functional.cosine_similarity(
+        output.float().flatten(), reference.flatten(), dim=0
+    ).item()
+    assert cosine > 0.999
+
+
+class _SkinnyGemmSpy:
+    """Wraps the skinny-GEMM singleton to record whether CuTe was invoked."""
+
+    def __init__(self, real: Any) -> None:
+        self._real = real
+        self.calls: list[int] = []
+
+    def __call__(self, a, b, config=None, residual=None):
+        self.calls.append(a.shape[0])
+        return self._real(a, b, config, residual)
+
+    def is_available(self) -> bool:
+        return self._real.is_available()
+
+
+@pytest.mark.parametrize("num_tokens", [1, 2, 3, 4])
+def test_latent_moe_production_layout_residual(
+    monkeypatch: pytest.MonkeyPatch,
+    num_tokens: int,
+) -> None:
+    """The real Latent-MoE residual is a non-packed slice of a cat buffer.
+
+    The strict packed-row-major guard rejects such a slice at every token count
+    (a size-1 leading dim reads as contiguous but its stride is not packed), so
+    the CuTe residual epilogue never fires for this production layout and the
+    method falls back to addmm. Output is correct regardless of the path.
+    """
+    _require_sm103_and_cute()
+    latent_dim, shared_dim = 3584, 7168  # routed_expert_up_proj K, N
+    torch.manual_seed(7 + num_tokens)
+    buf = torch.randn(
+        num_tokens, latent_dim + shared_dim, dtype=torch.bfloat16, device="cuda"
+    )
+    latent = buf[:, :latent_dim]  # non-contiguous view (row stride = full width)
+    residual = buf[:, latent_dim:]  # non-contiguous view
+    weight = torch.randn(shared_dim, latent_dim, dtype=torch.bfloat16, device="cuda")
+
+    spec = k3_gemm.KIMI_K3_PROJECTIONS[(shared_dim, latent_dim)]
+    method = k3_gemm.KimiK3LowLatencyLinearMethod(
+        k3_gemm._build_plan(spec), k3_gemm._build_residual_plan(spec)
+    )
+    spy = _SkinnyGemmSpy(k3_gemm.shape_dynamic_skinny_gemm)
+    monkeypatch.setattr(k3_gemm, "shape_dynamic_skinny_gemm", spy)
+
+    layer = SimpleNamespace(weight=weight)
+    output = method.apply_with_residual(layer, latent, residual)
+
+    reference = latent.float() @ weight.float().t() + residual.float()
+    cosine = torch.nn.functional.cosine_similarity(
+        output.float().flatten(), reference.flatten(), dim=0
+    ).item()
+    assert cosine > 0.999  # correct regardless of the path taken
+    assert not spy.calls, (
+        "non-packed buf-slice residual must fall back to addmm at every M"
+    )
+
+
+def test_residual_dispatch_falls_back_to_addmm(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    fallback = torch.randn(2, 3)
+    residual = torch.randn(2, 3)
+    x = torch.randn(2, 4)
+    weight = torch.randn(3, 4)
+    monkeypatch.setattr(torch, "addmm", lambda *args: fallback)
+    # CPU tensors fail the runtime check, forcing the addmm fallback.
+    method = k3_gemm.KimiK3LowLatencyLinearMethod({}, {})
+
+    output = method.apply_with_residual(SimpleNamespace(weight=weight), x, residual)
+
+    assert output is fallback
+
+
+def test_fallback_preserves_default_method(monkeypatch: pytest.MonkeyPatch) -> None:
+    fallback = torch.empty(2, 8)
+    monkeypatch.setattr(
+        k3_gemm.UnquantizedLinearMethod,
+        "apply",
+        lambda *args: fallback,
+    )
+    # 1-D input fails the runtime check, forcing the base-method fallback.
+    method = k3_gemm.KimiK3LowLatencyLinearMethod({}, {})
+
+    output = method.apply(
+        SimpleNamespace(weight=torch.empty(0)),
+        torch.empty(0),
+    )
+
+    assert output is fallback

tests/kernels/test_kda.py

@@ -1,226 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Precision tests for vllm's chunk_kda Triton operator.
-
-Compares chunk_kda against a naive recurrent reference (float32).
-Uses torch.rand for q/k/v to match FLA's test pattern.
-"""
-
-import pytest
-import torch
-import torch.nn.functional as F
-
-from vllm.third_party.flash_linear_attention.ops.kda import (
-    chunk_kda,
-    chunk_kda_with_fused_gate,
-    fused_kda_gate,
-)
-from vllm.third_party.flash_linear_attention.ops.l2norm import l2norm_fwd
-
-DEVICE = "cuda"
-
-
-def naive_recurrent_kda(
-    q: torch.Tensor,
-    k: torch.Tensor,
-    v: torch.Tensor,
-    g: torch.Tensor,
-    beta: torch.Tensor,
-    scale: float | None = None,
-    initial_state: torch.Tensor | None = None,
-    output_final_state: bool = False,
-) -> tuple[torch.Tensor, torch.Tensor | None]:
-    """Naive recurrent KDA reference, ported from FLA's naive.py."""
-    dtype = v.dtype
-    B, T, H, K = q.shape
-    V = v.shape[-1]
-    if scale is None:
-        scale = K**-0.5
-
-    q, k, v, g, beta = (x.to(torch.float) for x in [q, k, v, g, beta])
-    q = q * scale
-
-    S = k.new_zeros(B, H, K, V).to(q)
-    if initial_state is not None:
-        S += initial_state
-    o = torch.zeros_like(v)
-    for i in range(T):
-        q_i, k_i, v_i, g_i, b_i = q[:, i], k[:, i], v[:, i], g[:, i], beta[:, i]
-        S = S * g_i[..., None].exp()
-        S = S + torch.einsum(
-            "bhk,bhv->bhkv",
-            b_i[..., None] * k_i,
-            v_i - (k_i[..., None] * S).sum(-2),
-        )
-        o[:, i] = torch.einsum("bhk,bhkv->bhv", q_i, S)
-    if not output_final_state:
-        S = None
-    return o.to(dtype), S
-
-
-def assert_close(
-    name: str,
-    ref: torch.Tensor,
-    tri: torch.Tensor,
-    ratio: float,
-    err_atol: float = 1e-6,
-):
-    """RMSE-based relative error comparison."""
-    abs_err = (ref.detach() - tri.detach()).flatten().abs().max().item()
-    rmse_diff = (ref.detach() - tri.detach()).flatten().square().mean().sqrt().item()
-    rmse_base = ref.detach().flatten().square().mean().sqrt().item()
-    rel_err = rmse_diff / (rmse_base + 1e-8)
-    print(f"{name:>4} | abs={abs_err:.6f} | rmse={rel_err:.6f} | thr={ratio}")
-    if abs_err <= err_atol:
-        return
-    assert not torch.isnan(ref).any(), f"{name}: NaN detected in ref"
-    assert not torch.isnan(tri).any(), f"{name}: NaN detected in tri"
-    assert rel_err < ratio, (
-        f"{name}: max abs err {abs_err:.6f}, rmse ratio {rel_err:.6f} >= {ratio}"
-    )
-
-
-@pytest.mark.parametrize(
-    ("H", "D", "cu_seqlens", "dtype"),
-    [
-        pytest.param(
-            *test,
-            id="H{}-D{}-cu{}-{}".format(*test),
-        )
-        for test in [
-            (32, 128, [0, 64], torch.float16),
-            (32, 128, [0, 1024], torch.float16),
-            (32, 128, [0, 15], torch.float16),
-            (32, 128, [0, 256, 512, 768, 1024], torch.float16),
-            (32, 128, [0, 15, 100, 300, 1200], torch.float16),
-            (64, 128, [0, 256, 500, 1000], torch.float16),
-            (32, 128, [0, 8192], torch.float16),
-            (32, 128, [0, 256, 500, 1000], torch.bfloat16),
-        ]
-    ],
-)
-@torch.inference_mode()
-def test_chunk_kda(
-    H: int,
-    D: int,
-    cu_seqlens: list[int],
-    dtype: torch.dtype,
-):
-    T = cu_seqlens[-1]
-    torch.manual_seed(42)
-    B = 1
-    cu_seqlens_t = torch.LongTensor(cu_seqlens).to(DEVICE)
-    N = len(cu_seqlens) - 1
-
-    q = torch.rand(B, T, H, D, dtype=dtype, device=DEVICE)
-    k = torch.rand(B, T, H, D, dtype=dtype, device=DEVICE)
-    v = torch.rand(B, T, H, D, dtype=dtype, device=DEVICE)
-    g = F.logsigmoid(torch.randn(B, T, H, D, dtype=torch.float32, device=DEVICE)).to(
-        dtype
-    )
-    beta = torch.rand(B, T, H, dtype=dtype, device=DEVICE).sigmoid()
-    h0 = torch.randn(N, H, D, D, dtype=torch.float32, device=DEVICE)
-
-    # Naive reference with l2norm_fwd (same kernel as chunk_kda)
-    ref_outputs = []
-    ref_states = []
-    for i in range(N):
-        s, e = cu_seqlens[i], cu_seqlens[i + 1]
-        q_i = l2norm_fwd(q[:, s:e].contiguous())
-        k_i = l2norm_fwd(k[:, s:e].contiguous())
-        o_i, ht_i = naive_recurrent_kda(
-            q_i,
-            k_i,
-            v[:, s:e],
-            g[:, s:e],
-            beta[:, s:e],
-            initial_state=h0[i],
-            output_final_state=True,
-        )
-        ref_outputs.append(o_i)
-        ref_states.append(ht_i)
-    ref_o = torch.cat(ref_outputs, dim=1)
-    ref_ht = torch.cat(ref_states, dim=0)
-
-    # h0 transposed to (V, K) layout for the kernel; naive uses (K, V)
-    tri_o, tri_ht = chunk_kda(
-        q=q.clone(),
-        k=k.clone(),
-        v=v.clone(),
-        g=g.clone(),
-        beta=beta.clone(),
-        initial_state=h0.transpose(-1, -2).contiguous().clone(),
-        output_final_state=True,
-        cu_seqlens=cu_seqlens_t,
-        use_qk_l2norm_in_kernel=True,
-    )
-
-    assert not torch.isnan(tri_o).any(), "Triton output o contains NaN"
-    assert not torch.isnan(tri_ht).any(), "Triton output ht contains NaN"
-    assert_close("o", ref_o, tri_o, 0.005)
-    assert_close("ht", ref_ht, tri_ht.transpose(-1, -2).contiguous(), 0.005)
-
-
-@pytest.mark.parametrize(
-    ("cu_seqlens", "dtype"),
-    [
-        ([0, 64], torch.float16),
-        ([0, 15, 100, 300], torch.bfloat16),
-    ],
-)
-@torch.inference_mode()
-def test_chunk_kda_fused_gate_cumsum_matches_unfused(
-    cu_seqlens: list[int],
-    dtype: torch.dtype,
-):
-    H, D = 8, 64
-    T = cu_seqlens[-1]
-    N = len(cu_seqlens) - 1
-    torch.manual_seed(123)
-
-    cu_seqlens_t = torch.tensor(cu_seqlens, dtype=torch.int32, device=DEVICE)
-    q = torch.randn(1, T, H, D, dtype=dtype, device=DEVICE)
-    k = torch.randn(1, T, H, D, dtype=dtype, device=DEVICE)
-    v = torch.randn(1, T, H, D, dtype=dtype, device=DEVICE)
-    raw_g = torch.randn(1, T, H, D, dtype=dtype, device=DEVICE)
-    beta = torch.rand(1, T, H, dtype=dtype, device=DEVICE).sigmoid()
-    A_log = (torch.randn(H, dtype=torch.float32, device=DEVICE) * 0.5).contiguous()
-    dt_bias = (
-        torch.randn(H * D, dtype=torch.float32, device=DEVICE) * 0.1
-    ).contiguous()
-    h0 = torch.randn(N, H, D, D, dtype=torch.float32, device=DEVICE)
-    initial_state = h0.transpose(-1, -2).contiguous()
-
-    gate = fused_kda_gate(
-        raw_g.reshape(T, H * D),
-        A_log,
-        D,
-        g_bias=dt_bias,
-    ).unsqueeze(0)
-    old_o, old_ht = chunk_kda(
-        q=q.clone(),
-        k=k.clone(),
-        v=v.clone(),
-        g=gate,
-        beta=beta.clone(),
-        initial_state=initial_state.clone(),
-        output_final_state=True,
-        cu_seqlens=cu_seqlens_t,
-        use_qk_l2norm_in_kernel=True,
-    )
-    new_o, new_ht = chunk_kda_with_fused_gate(
-        q=q.clone(),
-        k=k.clone(),
-        v=v.clone(),
-        raw_g=raw_g,
-        beta=beta.clone(),
-        A_log=A_log,
-        g_bias=dt_bias,
-        initial_state=initial_state.clone(),
-        output_final_state=True,
-        cu_seqlens=cu_seqlens_t,
-        use_qk_l2norm_in_kernel=True,
-    )
-
-    assert_close("o", old_o, new_o, 1e-3, err_atol=1e-3)
-    assert_close("ht", old_ht, new_ht, 1e-3, err_atol=1e-3)

tests/models/kimi_k3/test_attn_res.py

@@ -5,6 +5,7 @@ import pytest
 import torch
 import torch.nn.functional as F
 
+from vllm.models.kimi_k3.common.mtp import fused_mtp_input
 from vllm.models.kimi_k3.nvidia.ops import attn_res
 from vllm.platforms import current_platform
 
@@ -191,3 +192,35 @@ def test_attn_res_without_output_norm():
     )
 
     torch.testing.assert_close(actual, expected, atol=8e-2, rtol=3e-2)
+
+
+@pytest.mark.parametrize("num_tokens", [0, 1, 17])
+def test_fused_mtp_input(num_tokens: int):
+    positions = torch.arange(num_tokens, device="cuda")
+    inputs_embeds = _randn_with_row_padding(num_tokens, HIDDEN_SIZE, padding=7)
+    previous_hidden_states = _randn_with_row_padding(
+        num_tokens, HIDDEN_SIZE, padding=11
+    )
+    enorm_weight = torch.randn(HIDDEN_SIZE, device="cuda", dtype=torch.bfloat16)
+    hnorm_weight = torch.randn(HIDDEN_SIZE, device="cuda", dtype=torch.bfloat16)
+
+    masked_inputs_embeds = torch.where(positions.unsqueeze(-1) == 0, 0, inputs_embeds)
+    expected = torch.cat(
+        (
+            F.rms_norm(masked_inputs_embeds, (HIDDEN_SIZE,), enorm_weight, EPS),
+            F.rms_norm(previous_hidden_states, (HIDDEN_SIZE,), hnorm_weight, EPS),
+        ),
+        dim=-1,
+    )
+    actual = fused_mtp_input(
+        positions,
+        inputs_embeds,
+        previous_hidden_states,
+        enorm_weight,
+        hnorm_weight,
+        EPS,
+    )
+
+    torch.testing.assert_close(actual, expected, atol=2e-2, rtol=2e-2)
+    assert actual.shape == (num_tokens, 2 * HIDDEN_SIZE)
+    assert actual.is_contiguous()

tests/models/kimi_k3/test_eagle3.py

@@ -0,0 +1,130 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from types import SimpleNamespace
+from unittest.mock import Mock
+
+import torch
+
+from vllm.model_executor.models.interfaces import supports_eagle3
+from vllm.models.kimi_k3.nvidia import model as kimi_model
+from vllm.models.kimi_k3.nvidia.model import (
+    KimiK3ForConditionalGeneration,
+    KimiLinearModel,
+)
+
+
+def _make_kimi_linear_model() -> KimiLinearModel:
+    model = object.__new__(KimiLinearModel)
+    object.__setattr__(model, "aux_hidden_state_layers", (2,))
+    object.__setattr__(model, "use_sequence_parallel", False)
+    return model
+
+
+def test_kimi_k3_advertises_eagle3_support():
+    assert supports_eagle3(KimiK3ForConditionalGeneration)
+
+
+def test_kimi_k3_uses_shared_eagle3_layer_configuration():
+    target = object.__new__(KimiK3ForConditionalGeneration)
+    torch.nn.Module.__init__(target)
+    model = _make_kimi_linear_model()
+    object.__setattr__(model, "layers", [None] * 93)
+    language_model = SimpleNamespace(
+        embed_input_ids=lambda _: None,
+        model=model,
+    )
+    object.__setattr__(target, "language_model", language_model)
+    object.__setattr__(target, "_language_model_names", ["language_model"])
+
+    target.set_aux_hidden_state_layers((2, 46, 90))
+
+    assert model.aux_hidden_state_layers == (2, 46, 90)
+    assert target.get_eagle3_default_aux_hidden_state_layers() == (
+        2,
+        46,
+        90,
+    )
+
+
+def test_kimi_linear_forward_extracts_standard_aux_hidden_states(monkeypatch):
+    model = _make_kimi_linear_model()
+    initial_hidden_states = torch.tensor([[1.0, 2.0]])
+    layer_hidden_states = torch.tensor([[3.0, 4.0]])
+    layer_residual = torch.tensor([[5.0, 6.0]])
+
+    object.__setattr__(model, "start_layer", 0)
+    object.__setattr__(model, "end_layer", 1)
+    object.__setattr__(
+        model,
+        "layers",
+        [Mock(return_value=(layer_hidden_states, None, layer_residual))],
+    )
+    object.__setattr__(model, "aux_hidden_state_layers", (0, 1))
+    object.__setattr__(model, "use_attn_res", False)
+    monkeypatch.setattr(
+        kimi_model,
+        "get_pp_group",
+        lambda: SimpleNamespace(is_first_rank=True, is_last_rank=True),
+    )
+
+    output, aux_hidden_states = model.forward(
+        input_ids=None,
+        positions=torch.tensor([0]),
+        intermediate_tensors=None,
+        inputs_embeds=initial_hidden_states,
+    )
+
+    expected_layer_output = layer_hidden_states + layer_residual
+    torch.testing.assert_close(output, expected_layer_output)
+    torch.testing.assert_close(aux_hidden_states[0], initial_hidden_states)
+    torch.testing.assert_close(aux_hidden_states[1], expected_layer_output)
+
+
+def test_kimi_linear_forward_extracts_attn_res_aux_hidden_states(monkeypatch):
+    model = _make_kimi_linear_model()
+    initial_hidden_states = torch.tensor([[1.0, 2.0]])
+    layer_hidden_states = torch.tensor([[3.0, 4.0]])
+    prefix_sum = torch.tensor([[5.0, 6.0]])
+    block_residual = torch.tensor([[[7.0, 8.0]]])
+    final_hidden_states = torch.tensor([[9.0, 10.0]])
+
+    object.__setattr__(model, "start_layer", 0)
+    object.__setattr__(model, "end_layer", 1)
+    object.__setattr__(
+        model,
+        "layers",
+        [Mock(return_value=(layer_hidden_states, prefix_sum, block_residual))],
+    )
+    object.__setattr__(model, "aux_hidden_state_layers", (0, 1))
+    object.__setattr__(model, "use_attn_res", True)
+    object.__setattr__(model, "num_attn_res_blocks", 1)
+    object.__setattr__(
+        model,
+        "output_attn_res_norm",
+        SimpleNamespace(weight=torch.ones(2), variance_epsilon=1e-5),
+    )
+    object.__setattr__(
+        model,
+        "output_attn_res_proj",
+        SimpleNamespace(weight=torch.ones(1, 2)),
+    )
+    monkeypatch.setattr(
+        kimi_model,
+        "get_pp_group",
+        lambda: SimpleNamespace(is_first_rank=True, is_last_rank=True),
+    )
+    final_attn_res = Mock(return_value=final_hidden_states)
+    monkeypatch.setattr(kimi_model, "attn_res", final_attn_res)
+
+    output, aux_hidden_states = model.forward(
+        input_ids=None,
+        positions=torch.tensor([0]),
+        intermediate_tensors=None,
+        inputs_embeds=initial_hidden_states,
+    )
+
+    torch.testing.assert_close(output, final_hidden_states)
+    torch.testing.assert_close(aux_hidden_states[0], initial_hidden_states)
+    torch.testing.assert_close(aux_hidden_states[1], prefix_sum + layer_hidden_states)
+    assert final_attn_res.call_args.args[2] is block_residual

tests/models/kimi_k3/test_kda.py

@@ -0,0 +1,757 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Precision tests for vllm's chunk_kda Triton operator.
+
+Compares chunk_kda against a naive recurrent reference (float32).
+Uses torch.rand for q/k/v to match FLA's test pattern.
+"""
+
+import pytest
+import torch
+import torch.nn.functional as F
+
+from vllm import _custom_ops as ops
+from vllm.model_executor.layers.mamba.ops.causal_conv1d import causal_conv1d_update
+from vllm.model_executor.layers.mamba.ops.gather_initial_states import (
+    gather_initial_states,
+)
+from vllm.models.kimi_k3.nvidia.kda import (
+    is_flashkda_supported,
+    is_fused_kda_decode_supported,
+)
+from vllm.models.kimi_k3.nvidia.ops.third_party.kda import (
+    chunk_kda,
+    chunk_kda_with_fused_gate,
+    fused_kda_gate,
+    fused_recurrent_kda,
+    fused_recurrent_kda_fwd,
+    fused_recurrent_kda_packed_decode,
+)
+from vllm.third_party.flash_linear_attention.ops.l2norm import l2norm_fwd
+
+DEVICE = "cuda"
+
+
+@torch.inference_mode()
+def test_gather_initial_states_correctness():
+    row_size = 8 * 128 * 128
+    storage = torch.randn(5, row_size + 256, dtype=torch.float32, device=DEVICE)
+    state = storage[:, :row_size].view(5, 8, 128, 128)
+    assert not state.is_contiguous()
+    assert state[0].is_contiguous()
+    indices = torch.tensor([4, 1, 3], dtype=torch.int32, device=DEVICE)
+    has_initial_state = torch.tensor([True, False, True], device=DEVICE)
+
+    expected = state[indices].clone()
+    expected[~has_initial_state] = 0
+
+    torch.testing.assert_close(
+        gather_initial_states(state, indices, has_initial_state),
+        expected,
+    )
+
+
+def naive_recurrent_kda(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    g: torch.Tensor,
+    beta: torch.Tensor,
+    scale: float | None = None,
+    initial_state: torch.Tensor | None = None,
+    output_final_state: bool = False,
+) -> tuple[torch.Tensor, torch.Tensor | None]:
+    """Naive recurrent KDA reference, ported from FLA's naive.py."""
+    dtype = v.dtype
+    B, T, H, K = q.shape
+    V = v.shape[-1]
+    if scale is None:
+        scale = K**-0.5
+
+    q, k, v, g, beta = (x.to(torch.float) for x in [q, k, v, g, beta])
+    q = q * scale
+
+    S = k.new_zeros(B, H, K, V).to(q)
+    if initial_state is not None:
+        S += initial_state
+    o = torch.zeros_like(v)
+    for i in range(T):
+        q_i, k_i, v_i, g_i, b_i = q[:, i], k[:, i], v[:, i], g[:, i], beta[:, i]
+        S = S * g_i[..., None].exp()
+        S = S + torch.einsum(
+            "bhk,bhv->bhkv",
+            b_i[..., None] * k_i,
+            v_i - (k_i[..., None] * S).sum(-2),
+        )
+        o[:, i] = torch.einsum("bhk,bhkv->bhv", q_i, S)
+    if not output_final_state:
+        S = None
+    return o.to(dtype), S
+
+
+def assert_close(
+    name: str,
+    ref: torch.Tensor,
+    tri: torch.Tensor,
+    ratio: float,
+    err_atol: float = 1e-6,
+):
+    """RMSE-based relative error comparison."""
+    abs_err = (ref.detach() - tri.detach()).flatten().abs().max().item()
+    rmse_diff = (ref.detach() - tri.detach()).flatten().square().mean().sqrt().item()
+    rmse_base = ref.detach().flatten().square().mean().sqrt().item()
+    rel_err = rmse_diff / (rmse_base + 1e-8)
+    print(f"{name:>4} | abs={abs_err:.6f} | rmse={rel_err:.6f} | thr={ratio}")
+    if abs_err <= err_atol:
+        return
+    assert not torch.isnan(ref).any(), f"{name}: NaN detected in ref"
+    assert not torch.isnan(tri).any(), f"{name}: NaN detected in tri"
+    assert rel_err < ratio, (
+        f"{name}: max abs err {abs_err:.6f}, rmse ratio {rel_err:.6f} >= {ratio}"
+    )
+
+
+@pytest.mark.parametrize(
+    ("H", "D", "cu_seqlens", "dtype"),
+    [
+        pytest.param(
+            *test,
+            id="H{}-D{}-cu{}-{}".format(*test),
+        )
+        for test in [
+            (32, 128, [0, 64], torch.float16),
+            (32, 128, [0, 1024], torch.float16),
+            (32, 128, [0, 15], torch.float16),
+            (32, 128, [0, 256, 512, 768, 1024], torch.float16),
+            (32, 128, [0, 15, 100, 300, 1200], torch.float16),
+            (64, 128, [0, 256, 500, 1000], torch.float16),
+            (32, 128, [0, 8192], torch.float16),
+            (32, 128, [0, 256, 500, 1000], torch.bfloat16),
+        ]
+    ],
+)
+@torch.inference_mode()
+def test_chunk_kda(
+    H: int,
+    D: int,
+    cu_seqlens: list[int],
+    dtype: torch.dtype,
+):
+    T = cu_seqlens[-1]
+    torch.manual_seed(42)
+    B = 1
+    cu_seqlens_t = torch.LongTensor(cu_seqlens).to(DEVICE)
+    N = len(cu_seqlens) - 1
+
+    q = torch.rand(B, T, H, D, dtype=dtype, device=DEVICE)
+    k = torch.rand(B, T, H, D, dtype=dtype, device=DEVICE)
+    v = torch.rand(B, T, H, D, dtype=dtype, device=DEVICE)
+    g = F.logsigmoid(torch.randn(B, T, H, D, dtype=torch.float32, device=DEVICE)).to(
+        dtype
+    )
+    beta = torch.rand(B, T, H, dtype=dtype, device=DEVICE).sigmoid()
+    h0 = torch.randn(N, H, D, D, dtype=torch.float32, device=DEVICE)
+
+    # Naive reference with l2norm_fwd (same kernel as chunk_kda)
+    ref_outputs = []
+    ref_states = []
+    for i in range(N):
+        s, e = cu_seqlens[i], cu_seqlens[i + 1]
+        q_i = l2norm_fwd(q[:, s:e].contiguous())
+        k_i = l2norm_fwd(k[:, s:e].contiguous())
+        o_i, ht_i = naive_recurrent_kda(
+            q_i,
+            k_i,
+            v[:, s:e],
+            g[:, s:e],
+            beta[:, s:e],
+            initial_state=h0[i],
+            output_final_state=True,
+        )
+        ref_outputs.append(o_i)
+        ref_states.append(ht_i)
+    ref_o = torch.cat(ref_outputs, dim=1)
+    ref_ht = torch.cat(ref_states, dim=0)
+
+    # h0 transposed to (V, K) layout for the kernel; naive uses (K, V)
+    tri_o, tri_ht = chunk_kda(
+        q=q.clone(),
+        k=k.clone(),
+        v=v.clone(),
+        g=g.clone(),
+        beta=beta.clone(),
+        initial_state=h0.transpose(-1, -2).contiguous().clone(),
+        output_final_state=True,
+        cu_seqlens=cu_seqlens_t,
+        use_qk_l2norm_in_kernel=True,
+    )
+
+    assert not torch.isnan(tri_o).any(), "Triton output o contains NaN"
+    assert not torch.isnan(tri_ht).any(), "Triton output ht contains NaN"
+    assert_close("o", ref_o, tri_o, 0.005)
+    assert_close("ht", ref_ht, tri_ht.transpose(-1, -2).contiguous(), 0.005)
+
+
+@pytest.mark.parametrize(
+    ("cu_seqlens", "dtype", "lower_bound"),
+    [
+        ([0, 64], torch.float16, None),
+        ([0, 15, 100, 300], torch.bfloat16, None),
+        ([0, 15, 100, 300], torch.bfloat16, -3.0),
+    ],
+)
+@torch.inference_mode()
+def test_chunk_kda_fused_gate_cumsum_matches_unfused(
+    cu_seqlens: list[int],
+    dtype: torch.dtype,
+    lower_bound: float | None,
+):
+    H, D = 8, 64
+    T = cu_seqlens[-1]
+    N = len(cu_seqlens) - 1
+    torch.manual_seed(123)
+
+    cu_seqlens_t = torch.tensor(cu_seqlens, dtype=torch.int32, device=DEVICE)
+    q = torch.randn(1, T, H, D, dtype=dtype, device=DEVICE)
+    k = torch.randn(1, T, H, D, dtype=dtype, device=DEVICE)
+    v = torch.randn(1, T, H, D, dtype=dtype, device=DEVICE)
+    raw_g = torch.randn(1, T, H, D, dtype=dtype, device=DEVICE)
+    beta_storage = torch.randn(1, T, 2 * H + 3, dtype=dtype, device=DEVICE)
+    raw_beta = beta_storage[..., 1 : 2 * H + 1 : 2]
+    beta = raw_beta.float().sigmoid()
+    A_log = (torch.randn(H, dtype=torch.float32, device=DEVICE) * 0.5).contiguous()
+    dt_bias = (
+        torch.randn(H * D, dtype=torch.float32, device=DEVICE) * 0.1
+    ).contiguous()
+    h0 = torch.randn(N, H, D, D, dtype=torch.float32, device=DEVICE)
+    initial_state = h0.transpose(-1, -2).contiguous()
+
+    gate = fused_kda_gate(
+        raw_g.reshape(T, H * D),
+        A_log,
+        D,
+        g_bias=dt_bias,
+        lower_bound=lower_bound,
+    )
+    if lower_bound is not None:
+        expected_gate = lower_bound * torch.sigmoid(
+            A_log.exp()[None, :, None]
+            * (raw_g.float().view(T, H, D) + dt_bias.view(H, D))
+        )
+        torch.testing.assert_close(gate, expected_gate)
+    gate = gate.unsqueeze(0)
+    old_o, old_ht = chunk_kda(
+        q=q.clone(),
+        k=k.clone(),
+        v=v.clone(),
+        g=gate,
+        beta=beta,
+        initial_state=initial_state.clone(),
+        output_final_state=True,
+        cu_seqlens=cu_seqlens_t,
+        use_qk_l2norm_in_kernel=True,
+    )
+    new_o, new_ht = chunk_kda_with_fused_gate(
+        q=q.clone(),
+        k=k.clone(),
+        v=v.clone(),
+        raw_g=raw_g,
+        raw_beta=raw_beta,
+        A_log=A_log,
+        g_bias=dt_bias,
+        lower_bound=lower_bound,
+        initial_state=initial_state.clone(),
+        output_final_state=True,
+        cu_seqlens=cu_seqlens_t,
+        use_qk_l2norm_in_kernel=True,
+    )
+
+    assert_close("o", old_o, new_o, 1e-3, err_atol=1e-3)
+    assert_close("ht", old_ht, new_ht, 1e-3, err_atol=1e-3)
+
+
+@pytest.mark.parametrize("num_seqs", [1, 8, 32])
+@pytest.mark.parametrize("lower_bound", [-5.0, None])
+@pytest.mark.parametrize("state_indices_stride", [1, 8])
+@torch.inference_mode()
+def test_packed_kda_decode_correctness(
+    num_seqs: int,
+    lower_bound: float | None,
+    state_indices_stride: int,
+):
+    H, D = 8, 128
+    torch.manual_seed(321)
+
+    packed_storage = torch.randn(
+        num_seqs,
+        3 * H * D + 1,
+        dtype=torch.bfloat16,
+        device=DEVICE,
+    )
+    mixed_qkv = packed_storage[:, : 3 * H * D]
+    assert mixed_qkv.stride(0) == 3 * H * D + 1
+    q, k, v = (
+        x.contiguous().view(1, num_seqs, H, D) for x in mixed_qkv.split(H * D, dim=-1)
+    )
+    raw_g = torch.randn(
+        1,
+        num_seqs,
+        H,
+        D,
+        dtype=torch.bfloat16,
+        device=DEVICE,
+    )
+    raw_beta = torch.randn(
+        1,
+        num_seqs,
+        H,
+        dtype=torch.bfloat16,
+        device=DEVICE,
+    )
+    beta = raw_beta.float().sigmoid()
+    A_log = torch.randn(H, dtype=torch.float32, device=DEVICE) * 0.5
+    dt_bias = torch.randn(H, D, dtype=torch.float32, device=DEVICE) * 0.1
+    state_storage = torch.randn(
+        num_seqs + 1,
+        H * D * D + 17,
+        dtype=torch.float32,
+        device=DEVICE,
+    )
+    state = state_storage[:, : H * D * D].view(num_seqs + 1, H, D, D)
+    assert not state.is_contiguous()
+    assert state.stride()[1:] == (D * D, D, 1)
+    state_indices_storage = torch.zeros(
+        num_seqs,
+        state_indices_stride,
+        dtype=torch.int32,
+        device=DEVICE,
+    )
+    state_indices = state_indices_storage[:, 0]
+    state_indices.copy_(
+        torch.arange(
+            1,
+            num_seqs + 1,
+            dtype=torch.int32,
+            device=DEVICE,
+        )
+    )
+    gate = fused_kda_gate(
+        raw_g.reshape(num_seqs, H * D),
+        A_log,
+        D,
+        g_bias=dt_bias,
+        lower_bound=lower_bound,
+    ).unsqueeze(0)
+    dense_state = state.clone()
+    dense_out, _ = fused_recurrent_kda_fwd(
+        q=q,
+        k=k,
+        v=v,
+        g=gate,
+        beta=beta,
+        scale=D**-0.5,
+        initial_state=dense_state,
+        inplace_final_state=True,
+        cu_seqlens=torch.arange(
+            num_seqs + 1,
+            dtype=torch.int32,
+            device=DEVICE,
+        ),
+        ssm_state_indices=state_indices,
+        use_qk_l2norm_in_kernel=True,
+    )
+    packed_state = state
+    packed_out, _ = fused_recurrent_kda_packed_decode(
+        mixed_qkv=mixed_qkv,
+        raw_g=raw_g,
+        raw_beta=raw_beta,
+        A_log=A_log,
+        dt_bias=dt_bias,
+        lower_bound=lower_bound,
+        initial_state=packed_state,
+        state_indices=state_indices,
+    )
+
+    assert_close("o", dense_out, packed_out, 1e-3, err_atol=1e-3)
+    assert_close("ht", dense_state, packed_state, 1e-3, err_atol=1e-3)
+
+
+@pytest.mark.parametrize(
+    ("H", "fuse_gate"),
+    [(12, True), (12, False), (12, None), (96, None)],
+)
+@pytest.mark.parametrize("lower_bound", [-5.0, None])
+@torch.inference_mode()
+def test_kda_spec_decode_correctness(
+    H: int,
+    fuse_gate: bool | None,
+    lower_bound: float | None,
+):
+    num_seqs, query_len, D = 3, 3, 128
+    T = num_seqs * query_len
+    torch.manual_seed(1234)
+
+    qkv_storage = torch.randn(
+        1,
+        T,
+        3 * H * D + 7,
+        dtype=torch.bfloat16,
+        device=DEVICE,
+    )
+    packed_qkv = qkv_storage[..., : 3 * H * D]
+    q, k, v = (x.view(1, T, H, D) for x in packed_qkv.split(H * D, dim=-1))
+    gate_storage = torch.randn(
+        1,
+        T,
+        H * D + 5,
+        dtype=torch.bfloat16,
+        device=DEVICE,
+    )
+    raw_g = gate_storage[..., : H * D].view(1, T, H, D)
+    beta_storage = torch.randn(
+        1,
+        T,
+        H + 1,
+        dtype=torch.bfloat16,
+        device=DEVICE,
+    )
+    raw_beta = beta_storage[..., :H]
+    A_log = 0.5 * torch.randn(H, dtype=torch.float32, device=DEVICE)
+    dt_bias = 0.1 * torch.randn(H, D, dtype=torch.float32, device=DEVICE)
+    cu_seqlens = torch.arange(
+        0,
+        T + 1,
+        query_len,
+        dtype=torch.int32,
+        device=DEVICE,
+    )
+    state_indices = torch.arange(
+        1,
+        T + 1,
+        dtype=torch.int32,
+        device=DEVICE,
+    ).view(num_seqs, query_len)
+    num_accepted_tokens = torch.tensor(
+        [1, 2, 3],
+        dtype=torch.int32,
+        device=DEVICE,
+    )
+    state_storage = 0.01 * torch.randn(
+        T + 1,
+        H * D * D + 17,
+        dtype=torch.float32,
+        device=DEVICE,
+    )
+    state = state_storage[:, : H * D * D].view(T + 1, H, D, D)
+    output_storage = torch.full(
+        (1, T, H * D + 11),
+        torch.nan,
+        dtype=torch.bfloat16,
+        device=DEVICE,
+    )
+    output = output_storage[..., : H * D].view(1, T, H, D)
+
+    gate = fused_kda_gate(
+        raw_g.contiguous().view(T, H * D),
+        A_log,
+        D,
+        g_bias=dt_bias,
+        lower_bound=lower_bound,
+    ).unsqueeze(0)
+    beta = raw_beta.float().sigmoid()
+    q_norm = l2norm_fwd(q.contiguous())
+    k_norm = l2norm_fwd(k.contiguous())
+    expected_state = state.clone()
+    expected_outputs = []
+    for seq, accepted in enumerate(num_accepted_tokens.tolist()):
+        recurrent_state = expected_state[state_indices[seq, accepted - 1]].transpose(
+            -1, -2
+        )
+        start = seq * query_len
+        for token in range(query_len):
+            token_slice = slice(start + token, start + token + 1)
+            token_output, recurrent_state = naive_recurrent_kda(
+                q_norm[:, token_slice],
+                k_norm[:, token_slice],
+                v[:, token_slice],
+                gate[:, token_slice],
+                beta[:, token_slice],
+                initial_state=recurrent_state,
+                output_final_state=True,
+            )
+            assert recurrent_state is not None
+            expected_outputs.append(token_output)
+            expected_state[state_indices[seq, token]] = recurrent_state.transpose(
+                -1, -2
+            )
+    expected = torch.cat(expected_outputs, dim=1)
+
+    actual_state = state.clone()
+    actual, _ = fused_recurrent_kda(
+        q=q,
+        k=k,
+        v=v,
+        raw_g=raw_g,
+        raw_beta=raw_beta,
+        A_log=A_log,
+        dt_bias=dt_bias,
+        lower_bound=lower_bound,
+        initial_state=actual_state,
+        cu_seqlens=cu_seqlens,
+        ssm_state_indices=state_indices,
+        num_accepted_tokens=num_accepted_tokens,
+        out=output,
+        fuse_gate=fuse_gate,
+    )
+
+    assert actual.data_ptr() == output.data_ptr()
+    assert_close("o", expected, actual, 1e-3, err_atol=1e-3)
+    used_states = state_indices.flatten().long()
+    assert_close(
+        "ht",
+        expected_state[used_states],
+        actual_state[used_states],
+        3e-3,
+        err_atol=3e-3,
+    )
+    assert torch.isnan(output_storage[..., H * D :]).all()
+
+
+@pytest.mark.parametrize(
+    ("num_heads", "num_seqs", "lower_bound", "fuse_output_norm"),
+    [
+        (12, 1, -5.0, True),
+        (12, 4, None, False),
+        (24, 4, None, False),
+        (48, 1, -5.0, True),
+        (96, 1, -5.0, True),
+    ],
+)
+@torch.inference_mode()
+def test_fused_kda_decode_correctness(
+    num_heads: int,
+    num_seqs: int,
+    lower_bound: float | None,
+    fuse_output_norm: bool,
+):
+    D, W = 128, 4
+    if not is_fused_kda_decode_supported(
+        num_heads,
+        D,
+        W,
+        num_spec=0,
+        input_dtype=torch.bfloat16,
+        conv_state_dtype=torch.bfloat16,
+    ):
+        pytest.skip("Fused KDA decode is not supported on this platform")
+    torch.manual_seed(967 + num_heads + num_seqs)
+    dim = num_heads * D
+    slots = num_seqs + 2
+    packed_x_storage = torch.randn(
+        num_seqs, 3 * dim + 17, dtype=torch.bfloat16, device=DEVICE
+    )
+    packed_x = packed_x_storage[:, : 3 * dim]
+    weight = 0.1 * torch.randn(3 * dim, W, dtype=torch.float32, device=DEVICE)
+    conv_seed = 0.1 * torch.randn(
+        slots,
+        W - 1,
+        3 * dim,
+        dtype=torch.bfloat16,
+        device=DEVICE,
+    ).transpose(1, 2)
+    raw_g = torch.randn(
+        1,
+        num_seqs,
+        num_heads,
+        D,
+        dtype=torch.bfloat16,
+        device=DEVICE,
+    )
+    raw_beta_storage = torch.randn(
+        1,
+        num_seqs,
+        num_heads + 1,
+        dtype=torch.bfloat16,
+        device=DEVICE,
+    )
+    raw_beta = raw_beta_storage[:, :, :num_heads]
+    output_gate_storage = torch.randn(
+        num_seqs,
+        dim + 7,
+        dtype=torch.bfloat16,
+        device=DEVICE,
+    )
+    output_gate = output_gate_storage[:, :dim].view(num_seqs, num_heads, D)
+    norm_weight = torch.randn(D, dtype=torch.float32, device=DEVICE)
+    norm_eps = 1e-5
+    A_log = 0.5 * torch.randn(num_heads, dtype=torch.float32, device=DEVICE)
+    dt_bias = 0.1 * torch.randn(dim, dtype=torch.float32, device=DEVICE)
+    state_indices = torch.arange(
+        num_seqs,
+        0,
+        -1,
+        dtype=torch.int32,
+        device=DEVICE,
+    )
+    state_seed = 0.01 * torch.randn(
+        slots,
+        num_heads,
+        D,
+        D,
+        dtype=torch.float32,
+        device=DEVICE,
+    )
+
+    conv_ref = conv_seed.clone()
+    state_ref = state_seed.clone()
+    mixed_qkv = causal_conv1d_update(
+        packed_x,
+        conv_ref,
+        weight,
+        activation="silu",
+        conv_state_indices=state_indices,
+        validate_data=True,
+        out=torch.empty_like(packed_x),
+    )
+    expected, _ = fused_recurrent_kda_packed_decode(
+        mixed_qkv=mixed_qkv,
+        raw_g=raw_g,
+        raw_beta=raw_beta,
+        A_log=A_log,
+        dt_bias=dt_bias,
+        lower_bound=lower_bound,
+        initial_state=state_ref,
+        state_indices=state_indices,
+    )
+    if fuse_output_norm:
+        expected_float = expected.float()
+        expected = (
+            expected_float
+            * torch.rsqrt(expected_float.square().mean(dim=-1, keepdim=True) + norm_eps)
+            * norm_weight
+            * output_gate.float().sigmoid().unsqueeze(0)
+        ).to(expected.dtype)
+
+    conv_slot_elements = 3 * dim * (W - 1)
+    state_slot_elements = num_heads * D * D
+    conv_slot_bytes = conv_slot_elements * torch.bfloat16.itemsize
+    page_bytes = conv_slot_bytes + state_slot_elements * torch.float32.itemsize
+    cache_storage = torch.empty(slots * page_bytes, dtype=torch.uint8, device=DEVICE)
+    conv_actual = torch.as_strided(
+        cache_storage.view(torch.bfloat16),
+        size=(slots, 3 * dim, W - 1),
+        stride=(page_bytes // torch.bfloat16.itemsize, 1, 3 * dim),
+    )
+    state_actual = torch.as_strided(
+        cache_storage.view(torch.float32),
+        size=(slots, num_heads, D, D),
+        stride=(page_bytes // torch.float32.itemsize, D * D, D, 1),
+        storage_offset=conv_slot_bytes // torch.float32.itemsize,
+    )
+    conv_actual.copy_(conv_seed)
+    state_actual.copy_(state_seed)
+    fused_weight = weight.reshape(3, dim, W).transpose(1, 2).contiguous()
+    actual = ops.fused_kda_decode(
+        x=packed_x,
+        weight=fused_weight,
+        bias=None,
+        conv_state=conv_actual,
+        raw_g=raw_g,
+        raw_beta=raw_beta,
+        A_log=A_log,
+        dt_bias=dt_bias,
+        state_indices=state_indices,
+        state=state_actual,
+        lower_bound=lower_bound,
+        output_gate=output_gate if fuse_output_norm else None,
+        norm_weight=norm_weight if fuse_output_norm else None,
+        norm_eps=norm_eps,
+    )
+
+    torch.testing.assert_close(actual, expected, atol=3e-2, rtol=3e-2)
+    torch.testing.assert_close(conv_actual, conv_ref, atol=0, rtol=0)
+    torch.testing.assert_close(state_actual, state_ref, atol=3e-2, rtol=3e-2)
+
+
+def test_fused_kda_decode_rejects_speculative_conv_state():
+    assert not is_fused_kda_decode_supported(
+        num_heads=12,
+        head_dim=128,
+        conv_width=4,
+        num_spec=2,
+        input_dtype=torch.bfloat16,
+        conv_state_dtype=torch.bfloat16,
+    )
+
+
+@torch.inference_mode()
+def test_flashkda_correctness():
+    if not is_flashkda_supported(128, torch.bfloat16, -3.0):
+        pytest.skip("FlashKDA is not supported on this platform")
+
+    import vllm._flashkda_C  # noqa: F401
+
+    B, T, H, D = 1, 48, 2, 128
+    torch.manual_seed(11)
+    q, k, v, raw_g = [
+        torch.randn(B, T, H, D, dtype=torch.bfloat16, device=DEVICE) for _ in range(4)
+    ]
+    beta_logits = torch.randn(B, T, H, dtype=torch.bfloat16, device=DEVICE)
+    A_log = torch.randn(H, dtype=torch.float32, device=DEVICE) * 0.5
+    dt_bias = torch.randn(H, D, dtype=torch.float32, device=DEVICE) * 0.1
+    initial_state = torch.randn(2, H, D, D, dtype=torch.float32, device=DEVICE)
+    cu_seqlens = torch.tensor([0, 17, T], dtype=torch.int32, device=DEVICE)
+    lower_bound = -3.0
+
+    gate = lower_bound * torch.sigmoid(
+        A_log.exp()[None, None, :, None] * (raw_g.float() + dt_bias[None, None, :, :])
+    )
+    beta = beta_logits.float().sigmoid()
+    q_norm = l2norm_fwd(q.contiguous())
+    k_norm = l2norm_fwd(k.contiguous())
+
+    expected_outputs = []
+    expected_states = []
+    for i, (start, end) in enumerate(
+        zip(cu_seqlens[:-1].tolist(), cu_seqlens[1:].tolist())
+    ):
+        output, final_state = naive_recurrent_kda(
+            q_norm[:, start:end],
+            k_norm[:, start:end],
+            v[:, start:end],
+            gate[:, start:end],
+            beta[:, start:end],
+            initial_state=initial_state[i].transpose(-1, -2),
+            output_final_state=True,
+        )
+        expected_outputs.append(output)
+        expected_states.append(final_state)
+    expected_out = torch.cat(expected_outputs, dim=1)
+    expected_state = torch.cat(expected_states).transpose(-1, -2).contiguous()
+
+    actual_out = torch.empty_like(v)
+    actual_state = torch.empty_like(initial_state)
+    workspace = torch.empty(
+        torch.ops._flashkda_C.get_workspace_size(T, H, cu_seqlens.numel() - 1),
+        dtype=torch.uint8,
+        device=DEVICE,
+    )
+    torch.ops._flashkda_C.fwd(
+        q,
+        k,
+        v,
+        raw_g,
+        beta_logits,
+        D**-0.5,
+        actual_out,
+        workspace,
+        A_log,
+        dt_bias,
+        lower_bound,
+        initial_state,
+        actual_state,
+        cu_seqlens,
+    )
+
+    assert_close("o", expected_out, actual_out, 0.01)
+    assert_close("ht", expected_state, actual_state, 0.01)

tests/models/kimi_k3/test_kda_metadata.py

@@ -0,0 +1,411 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from dataclasses import fields
+
+import pytest
+import torch
+
+from tests.v1.attention.utils import (
+    BatchSpec,
+    create_common_attn_metadata,
+    create_vllm_config,
+)
+from vllm.config import SpeculativeConfig
+from vllm.config.compilation import CUDAGraphMode
+from vllm.models.kimi_k3.nvidia.kda_metadata import (
+    KimiK3KDAAttentionBackend,
+    KimiK3KDAMetadata,
+    KimiK3KDAMetadataBuilder,
+    _mamba_get_block_table_tensor,
+    stage_spec_decode_metadata,
+)
+from vllm.v1.attention.backend import AttentionMetadataBuilder
+from vllm.v1.attention.backends.gdn_attn import (
+    GDNAttentionBackend,
+    GDNAttentionMetadata,
+    GDNAttentionMetadataBuilder,
+)
+from vllm.v1.attention.backends.utils import (
+    NULL_BLOCK_ID,
+    mamba_get_block_table_tensor,
+)
+from vllm.v1.kv_cache_interface import MambaSpec
+
+BLOCK_SIZE = 16
+DEVICE = torch.device("cpu")
+PRUNED_METADATA_FIELDS = {
+    "chunk_indices",
+    "chunk_offsets",
+    "prefill_query_start_loc",
+    "prefill_state_indices",
+    "prefill_has_initial_state",
+    "spec_sequence_masks",
+}
+
+
+def _assert_matches_shared_gdn(reference, actual: KimiK3KDAMetadata):
+    for field in fields(KimiK3KDAMetadata):
+        actual_value = getattr(actual, field.name)
+        expected_value = getattr(reference, field.name)
+        if field.name in PRUNED_METADATA_FIELDS:
+            assert actual_value is None
+            continue
+        if (
+            field.name in {"spec_token_indx", "non_spec_token_indx"}
+            and actual.num_spec_decodes > 0
+            and actual.num_prefills == 0
+            and actual.num_decodes == 0
+        ):
+            assert actual_value is None
+            continue
+        if isinstance(actual_value, torch.Tensor):
+            torch.testing.assert_close(actual_value, expected_value)
+        elif field.name == "nums_dict":
+            assert (actual_value is None) == (expected_value is None)
+            if actual_value is not None:
+                assert actual_value[8]["tot"] == expected_value[8]["tot"]
+                torch.testing.assert_close(
+                    actual_value[8]["nums"], expected_value[8]["nums"]
+                )
+        else:
+            assert actual_value == expected_value
+
+
+def _make_builder(
+    builder_cls: type[AttentionMetadataBuilder],
+    num_speculative_tokens: int,
+    full_cuda_graph: bool,
+    device: torch.device = DEVICE,
+    mamba_cache_mode: str = "none",
+) -> AttentionMetadataBuilder:
+    vllm_config = create_vllm_config(
+        model_name="Qwen/Qwen3.5-0.8B",
+        block_size=BLOCK_SIZE,
+    )
+    if num_speculative_tokens:
+        vllm_config.speculative_config = SpeculativeConfig(
+            method="ngram",
+            num_speculative_tokens=num_speculative_tokens,
+        )
+    vllm_config.compilation_config.cudagraph_mode = (
+        CUDAGraphMode.FULL_AND_PIECEWISE if full_cuda_graph else CUDAGraphMode.NONE
+    )
+    vllm_config.cache_config.mamba_cache_mode = mamba_cache_mode
+    return builder_cls(
+        kv_cache_spec=MambaSpec(
+            block_size=BLOCK_SIZE,
+            shapes=((16, 64),),
+            dtypes=(torch.float16,),
+            num_speculative_blocks=num_speculative_tokens,
+        ),
+        layer_names=["layer.0"],
+        vllm_config=vllm_config,
+        device=device,
+    )
+
+
+@pytest.mark.parametrize(
+    (
+        "batch",
+        "num_decode_draft_tokens",
+        "num_speculative_tokens",
+        "full_cuda_graph",
+        "is_prefilling",
+    ),
+    [
+        pytest.param(
+            BatchSpec(seq_lens=[50, 30], query_lens=[3, 3]),
+            [2, 2],
+            2,
+            False,
+            [False, False],
+            id="pure-spec-decode",
+        ),
+        pytest.param(
+            BatchSpec(seq_lens=[100, 65, 20], query_lens=[50, 1, 3]),
+            [-1, -1, 2],
+            2,
+            False,
+            [True, False, False],
+            id="mixed-prefill-and-spec-decode",
+        ),
+        pytest.param(
+            BatchSpec(seq_lens=[40, 30], query_lens=[1, 1]),
+            None,
+            0,
+            False,
+            [False, False],
+            id="regular-decode",
+        ),
+        pytest.param(
+            BatchSpec(seq_lens=[40, 30], query_lens=[1, 1]),
+            [0, 0],
+            2,
+            False,
+            [False, False],
+            id="no-scheduled-draft-tokens",
+        ),
+    ],
+)
+def test_kimi_k3_kda_metadata_matches_shared_gdn(
+    batch: BatchSpec,
+    num_decode_draft_tokens: list[int] | None,
+    num_speculative_tokens: int,
+    full_cuda_graph: bool,
+    is_prefilling: list[bool],
+):
+    kwargs: dict[str, torch.Tensor] = {}
+    if num_decode_draft_tokens is not None:
+        kwargs = {
+            "num_decode_draft_tokens_cpu": torch.tensor(
+                num_decode_draft_tokens, dtype=torch.int32
+            ),
+            "num_accepted_tokens": torch.ones(
+                batch.batch_size, dtype=torch.int32, device=DEVICE
+            ),
+        }
+
+    common_attn_metadata = create_common_attn_metadata(
+        batch, BLOCK_SIZE, DEVICE
+    ).replace(is_prefilling=torch.tensor(is_prefilling, dtype=torch.bool))
+    reference = _make_builder(
+        GDNAttentionMetadataBuilder,
+        num_speculative_tokens,
+        full_cuda_graph,
+    ).build(
+        0,
+        common_attn_metadata,
+        **kwargs,
+    )
+    actual = _make_builder(
+        KimiK3KDAMetadataBuilder,
+        num_speculative_tokens,
+        full_cuda_graph,
+    ).build(0, common_attn_metadata, **kwargs)
+
+    assert isinstance(actual, KimiK3KDAMetadata)
+    _assert_matches_shared_gdn(reference, actual)
+
+
+def test_mixed_regular_and_spec_decode_uses_packed_decode_metadata():
+    batch = BatchSpec(seq_lens=[100, 65, 20], query_lens=[1, 1, 3])
+    common_attn_metadata = create_common_attn_metadata(
+        batch, BLOCK_SIZE, DEVICE
+    ).replace(is_prefilling=torch.tensor([False, False, False]))
+    actual = _make_builder(
+        KimiK3KDAMetadataBuilder,
+        num_speculative_tokens=2,
+        full_cuda_graph=False,
+    ).build(
+        0,
+        common_attn_metadata,
+        num_decode_draft_tokens_cpu=torch.tensor([-1, -1, 2], dtype=torch.int32),
+        num_accepted_tokens=torch.ones(3, dtype=torch.int32, device=DEVICE),
+    )
+
+    # The K3 layer dispatches the non-spec subgroup to packed decode whenever
+    # it contains no prefill request.
+    assert actual.num_decodes == 2
+    assert actual.num_decode_tokens == 2
+    assert actual.num_prefills == 0
+    assert actual.num_prefill_tokens == 0
+    assert actual.has_initial_state is None
+    assert actual.nums_dict is None
+    assert actual.non_spec_query_start_loc is None
+    torch.testing.assert_close(actual.non_spec_token_indx, torch.tensor([0, 1]))
+    torch.testing.assert_close(actual.spec_token_indx, torch.tensor([2, 3, 4]))
+    torch.testing.assert_close(
+        actual.spec_query_start_loc,
+        torch.tensor([0, 3], dtype=torch.int32),
+    )
+
+
+def test_mixed_regular_and_spec_decode_excludes_request_padding():
+    batch = BatchSpec(seq_lens=[16, 65, 20], query_lens=[0, 1, 3])
+    common_attn_metadata = create_common_attn_metadata(
+        batch, BLOCK_SIZE, DEVICE
+    ).replace(is_prefilling=torch.tensor([False, False, False]))
+    actual = _make_builder(
+        KimiK3KDAMetadataBuilder,
+        num_speculative_tokens=2,
+        full_cuda_graph=False,
+    ).build(
+        0,
+        common_attn_metadata,
+        num_decode_draft_tokens_cpu=torch.tensor([-1, -1, 2], dtype=torch.int32),
+        num_accepted_tokens=torch.ones(3, dtype=torch.int32, device=DEVICE),
+    )
+
+    assert actual.num_decodes == 1
+    assert actual.non_spec_state_indices_tensor is not None
+    assert actual.non_spec_state_indices_tensor.shape == (1,)
+    torch.testing.assert_close(actual.non_spec_token_indx, torch.tensor([0]))
+    torch.testing.assert_close(actual.spec_token_indx, torch.tensor([1, 2, 3]))
+
+
+@pytest.mark.parametrize(
+    ("seq_len", "expected_has_initial_state"),
+    [
+        pytest.param(1, False, id="first-token-prefill"),
+        pytest.param(65, True, id="final-one-token-prefill-chunk"),
+    ],
+)
+def test_mixed_one_token_prefill_and_spec_decode_uses_prefill_metadata(
+    seq_len: int,
+    expected_has_initial_state: bool,
+):
+    batch = BatchSpec(seq_lens=[seq_len, 20], query_lens=[1, 3])
+    common_attn_metadata = create_common_attn_metadata(
+        batch, BLOCK_SIZE, DEVICE
+    ).replace(is_prefilling=torch.tensor([True, False]))
+    actual = _make_builder(
+        KimiK3KDAMetadataBuilder,
+        num_speculative_tokens=2,
+        full_cuda_graph=False,
+    ).build(
+        0,
+        common_attn_metadata,
+        num_decode_draft_tokens_cpu=torch.tensor([-1, 2], dtype=torch.int32),
+        num_accepted_tokens=torch.ones(2, dtype=torch.int32, device=DEVICE),
+    )
+
+    assert actual.num_prefills == 1
+    assert actual.num_prefill_tokens == 1
+    assert actual.num_decodes == 0
+    assert actual.num_decode_tokens == 0
+    assert actual.has_initial_state is not None
+    assert actual.has_initial_state.tolist() == [expected_has_initial_state]
+    assert actual.non_spec_query_start_loc is not None
+    torch.testing.assert_close(
+        actual.non_spec_query_start_loc,
+        torch.tensor([0, 1], dtype=torch.int32),
+    )
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA")
+def test_kimi_k3_kda_cudagraph_capture_matches_shared_gdn():
+    device = torch.device("cuda")
+    batch = BatchSpec(seq_lens=[50, 30], query_lens=[3, 3])
+    common_attn_metadata = create_common_attn_metadata(
+        batch, BLOCK_SIZE, device
+    ).replace(is_prefilling=torch.tensor([False, False]))
+    reference = _make_builder(
+        GDNAttentionMetadataBuilder,
+        num_speculative_tokens=2,
+        full_cuda_graph=True,
+        device=device,
+    ).build_for_cudagraph_capture(common_attn_metadata)
+    actual = _make_builder(
+        KimiK3KDAMetadataBuilder,
+        num_speculative_tokens=2,
+        full_cuda_graph=True,
+        device=device,
+    ).build_for_cudagraph_capture(common_attn_metadata)
+
+    assert isinstance(actual, KimiK3KDAMetadata)
+    _assert_matches_shared_gdn(reference, actual)
+
+
+def test_kimi_k3_kda_backend_uses_private_metadata_builder():
+    assert KimiK3KDAAttentionBackend.get_builder_cls() is KimiK3KDAMetadataBuilder
+    assert KimiK3KDAAttentionBackend.is_ssm()
+    assert issubclass(KimiK3KDAAttentionBackend, GDNAttentionBackend)
+    assert issubclass(KimiK3KDAMetadata, GDNAttentionMetadata)
+    assert issubclass(KimiK3KDAMetadataBuilder, GDNAttentionMetadataBuilder)
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA")
+def test_stage_spec_decode_metadata_matches_pytorch():
+    device = torch.device("cuda")
+    num_spec_decodes = 33
+    batch_size = 65
+    num_state_slots = 3
+    state_indices = torch.arange(
+        num_spec_decodes * 32,
+        dtype=torch.int32,
+        device=device,
+    ).reshape(num_spec_decodes, 32)[:, :num_state_slots]
+    query_start_loc = (
+        torch.arange(num_spec_decodes + 1, dtype=torch.int32, device=device)
+        * num_state_slots
+    )
+    num_accepted_tokens = (
+        torch.arange(num_spec_decodes, dtype=torch.int32, device=device)
+        % num_state_slots
+        + 1
+    )
+
+    staged_state_indices = torch.empty(
+        (batch_size, num_state_slots), dtype=torch.int32, device=device
+    )
+    staged_query_start_loc = torch.empty(
+        batch_size + 1, dtype=torch.int32, device=device
+    )
+    staged_num_accepted_tokens = torch.empty(
+        batch_size, dtype=torch.int32, device=device
+    )
+    stage_spec_decode_metadata(
+        state_indices,
+        query_start_loc,
+        num_accepted_tokens,
+        staged_state_indices,
+        staged_query_start_loc,
+        staged_num_accepted_tokens,
+        num_spec_decodes=num_spec_decodes,
+    )
+
+    expected_state_indices = torch.full_like(staged_state_indices, NULL_BLOCK_ID)
+    expected_state_indices[:num_spec_decodes] = state_indices
+    expected_query_start_loc = torch.full(
+        (batch_size + 1,),
+        query_start_loc[-1],
+        dtype=torch.int32,
+        device=device,
+    )
+    expected_query_start_loc[: num_spec_decodes + 1] = query_start_loc
+    expected_num_accepted_tokens = torch.ones(
+        batch_size, dtype=torch.int32, device=device
+    )
+    expected_num_accepted_tokens[:num_spec_decodes] = num_accepted_tokens
+
+    torch.testing.assert_close(staged_state_indices, expected_state_indices)
+    torch.testing.assert_close(staged_query_start_loc, expected_query_start_loc)
+    torch.testing.assert_close(staged_num_accepted_tokens, expected_num_accepted_tokens)
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA")
+def test_aligned_block_table_matches_shared_gdn():
+    device = torch.device("cuda")
+    seq_lens = torch.tensor(
+        [0, 1, 15, 16, 17, 31, 32, 33, 511, 512, 513],
+        dtype=torch.int32,
+        device=device,
+    ).repeat(6)[:65]
+    block_table_storage = torch.arange(
+        seq_lens.numel() * 128,
+        dtype=torch.int32,
+        device=device,
+    ).reshape(seq_lens.numel(), 128)
+    block_table = block_table_storage[:, ::2]
+    kv_cache_spec = MambaSpec(
+        block_size=BLOCK_SIZE,
+        shapes=((16, 64),),
+        dtypes=(torch.float16,),
+        num_speculative_blocks=2,
+    )
+
+    expected = mamba_get_block_table_tensor(
+        block_table,
+        seq_lens,
+        kv_cache_spec,
+        "align",
+    )
+    actual = _mamba_get_block_table_tensor(
+        block_table,
+        seq_lens,
+        kv_cache_spec,
+        "align",
+    )
+
+    torch.testing.assert_close(actual, expected)

tests/models/kimi_k3/test_latent_moe_tail.py

@@ -0,0 +1,145 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import pytest
+import ray
+import torch
+import torch.distributed as dist
+import torch.nn.functional as F
+
+from tests.utils import (
+    init_test_distributed_environment,
+    multi_gpu_test,
+    multi_process_parallel,
+)
+from vllm.distributed import get_tp_group
+from vllm.model_executor.warmup.cutedsl_warmup import cutedsl_warmup
+from vllm.models.kimi_k3.nvidia.ops.latent_moe_tail import KimiK3LatentMoETailOp
+from vllm.platforms import current_platform
+
+HIDDEN_SIZE = 7168
+LATENT_SIZE = 3584
+EPS = 0.1
+
+
+@ray.remote(num_gpus=1, max_calls=1)
+def _test_latent_moe_tail_worker(
+    monkeypatch: pytest.MonkeyPatch,
+    tp_size: int,
+    pp_size: int,
+    rank: int,
+    distributed_init_port: str,
+) -> None:
+    monkeypatch.delenv("CUDA_VISIBLE_DEVICES", raising=False)
+    device = torch.device(f"cuda:{rank}")
+    torch.accelerator.set_device_index(device)
+    init_test_distributed_environment(
+        tp_size,
+        pp_size,
+        rank,
+        distributed_init_port,
+    )
+
+    torch.manual_seed(0)
+    rms_weight = 1 + 0.1 * torch.randn(
+        LATENT_SIZE,
+        device=device,
+        dtype=torch.bfloat16,
+    )
+    up_weight = (
+        torch.randn(
+            HIDDEN_SIZE,
+            LATENT_SIZE,
+            device=device,
+            dtype=torch.bfloat16,
+        )
+        / LATENT_SIZE**0.5
+    )
+
+    group = get_tp_group().device_group
+    op = KimiK3LatentMoETailOp.initialize(
+        hidden_size=HIDDEN_SIZE,
+        latent_size=LATENT_SIZE,
+        dtype=torch.bfloat16,
+        device=device,
+        rms_eps=EPS,
+    )
+    cutedsl_warmup()
+
+    for iteration, num_tokens in enumerate((1, 5, 8, 16, 5)):
+        torch.manual_seed(100 * iteration + rank + 1)
+        routed_output = torch.randn(
+            num_tokens,
+            LATENT_SIZE,
+            device=device,
+            dtype=torch.bfloat16,
+        ).mul_(0.01)
+        shared_output = torch.randn(
+            num_tokens,
+            HIDDEN_SIZE,
+            device=device,
+            dtype=torch.bfloat16,
+        )
+
+        routed_reference = routed_output.clone()
+        shared_reference = shared_output.clone()
+        dist.all_reduce(routed_reference, group=group)
+        dist.all_reduce(shared_reference, group=group)
+        expected = F.linear(
+            F.rms_norm(
+                routed_reference,
+                (LATENT_SIZE,),
+                rms_weight,
+                EPS,
+            ),
+            up_weight,
+        )
+        expected.add_(shared_reference)
+
+        actual = op(
+            routed_output,
+            shared_output,
+            rms_weight,
+            up_weight,
+        )
+        torch.testing.assert_close(actual, expected, atol=8e-2, rtol=3e-2)
+        assert actual.is_contiguous()
+
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph):
+        graph_output = op(
+            routed_output,
+            shared_output,
+            rms_weight,
+            up_weight,
+        )
+    graph.replay()
+    torch.testing.assert_close(graph_output, expected, atol=8e-2, rtol=3e-2)
+
+
+def _run_latent_moe_tail_test(
+    monkeypatch: pytest.MonkeyPatch,
+    tp_size: int,
+) -> None:
+    if not current_platform.is_device_capability_family(100):
+        pytest.skip("K3 latent-MoE tail fusion requires SM100")
+    multi_process_parallel(
+        monkeypatch,
+        tp_size,
+        1,
+        _test_latent_moe_tail_worker,
+    )
+
+
+@multi_gpu_test(num_gpus=8)
+def test_latent_moe_tail_tp8_matches_native_path(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    _run_latent_moe_tail_test(monkeypatch, 8)
+
+
+@multi_gpu_test(num_gpus=16)
+def test_latent_moe_tail_tp16_matches_native_path(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    _run_latent_moe_tail_test(monkeypatch, 16)

tests/models/kimi_k3/test_sequence_parallel.py

@@ -0,0 +1,343 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from types import MethodType, SimpleNamespace
+from unittest.mock import Mock
+
+import pytest
+import torch
+from torch import nn
+
+from vllm.config import ParallelConfig
+from vllm.models.kimi_k3.nvidia import model as kimi_model
+from vllm.models.kimi_k3.nvidia import mtp as kimi_mtp
+from vllm.models.kimi_k3.nvidia.ops import sequence_parallel as sp_ops
+from vllm.platforms import current_platform
+
+
+class _IdentityNorm(nn.Module):
+    def __init__(self, hidden_size: int = 2) -> None:
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(hidden_size), requires_grad=False)
+        self.variance_epsilon = 1e-5
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        residual: torch.Tensor | None = None,
+    ):
+        if residual is None:
+            return hidden_states
+        return hidden_states, residual
+
+
+class _RecordingMoE(nn.Module):
+    def __init__(self) -> None:
+        super().__init__()
+        self.num_tokens = 0
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        self.num_tokens = hidden_states.shape[0]
+        return hidden_states
+
+
+class _Projection(nn.Module):
+    def __init__(self, hidden_size: int = 2) -> None:
+        super().__init__()
+        self.weight = nn.Parameter(
+            torch.ones(1, hidden_size),
+            requires_grad=False,
+        )
+
+
+class _SequenceParallelMTPBlock:
+    use_sequence_parallel = True
+
+    def __call__(
+        self,
+        *,
+        positions: torch.Tensor,
+        hidden_states: torch.Tensor,
+        residual: torch.Tensor | None,
+    ):
+        assert residual is None
+        return hidden_states * 2, None, hidden_states * 3
+
+
+def _mock_sequence_parallel_collectives(monkeypatch):
+    monkeypatch.setattr(
+        kimi_model,
+        "sp_reduce_scatter",
+        lambda tensor: tensor.chunk(2, dim=0)[0],
+    )
+    monkeypatch.setattr(
+        kimi_model,
+        "sp_shard",
+        lambda tensor: torch.nn.functional.pad(tensor, (0, 0, 0, 1))[:2],
+    )
+    monkeypatch.setattr(
+        kimi_model,
+        "sp_all_gather",
+        lambda tensor: torch.cat([tensor, tensor], dim=0),
+    )
+
+
+@pytest.mark.parametrize(
+    ("num_tokens", "is_padding", "tp_rank", "expected"),
+    [
+        (1, None, 0, [False]),
+        (1, None, 1, [True]),
+        (5, None, 2, [False, True]),
+        (5, None, 3, [True, True]),
+        (5, [False, True, False, False, False], 0, [False, True]),
+    ],
+)
+def test_sp_padding_mask_marks_added_rows(
+    monkeypatch,
+    num_tokens: int,
+    is_padding: list[bool] | None,
+    tp_rank: int,
+    expected: list[bool],
+):
+    monkeypatch.setattr(sp_ops, "get_tensor_model_parallel_world_size", lambda: 4)
+    monkeypatch.setattr(sp_ops, "get_tensor_model_parallel_rank", lambda: tp_rank)
+
+    hidden_states = torch.empty(num_tokens, 2)
+    padding = torch.tensor(is_padding) if is_padding is not None else None
+    actual = sp_ops.sp_padding_mask(padding, hidden_states)
+
+    torch.testing.assert_close(actual, torch.tensor(expected))
+
+
+@pytest.mark.parametrize(
+    ("data_parallel_size", "expected"),
+    [
+        (1, False),
+        (2, True),
+    ],
+)
+def test_moe_sequence_parallel_requires_data_parallel(
+    monkeypatch,
+    data_parallel_size: int,
+    expected: bool,
+):
+    monkeypatch.setattr(current_platform, "device_count", lambda: 2)
+    parallel_config = ParallelConfig(
+        tensor_parallel_size=2,
+        data_parallel_size=data_parallel_size,
+        enable_expert_parallel=True,
+        all2all_backend="allgather_reducescatter",
+    )
+
+    assert parallel_config.use_sequence_parallel_moe is expected
+
+
+def test_kimi_decoder_layer_keeps_moe_states_sequence_sharded(monkeypatch):
+    layer = object.__new__(kimi_model.KimiDecoderLayer)
+    nn.Module.__init__(layer)
+    layer.use_attn_res = False
+    layer.use_sequence_parallel = True
+    layer.input_layernorm = _IdentityNorm()
+    layer.post_attention_layernorm = _IdentityNorm()
+    layer.mlp = _RecordingMoE()
+    layer._run_self_attn = MethodType(
+        lambda self, positions, hidden_states: hidden_states,
+        layer,
+    )
+
+    _mock_sequence_parallel_collectives(monkeypatch)
+
+    positions = torch.arange(3)
+    full_hidden_states = torch.arange(6, dtype=torch.float32).view(3, 2)
+    hidden_states = kimi_model.sp_shard(full_hidden_states)
+    hidden_states, prefix_sum, residual = layer(
+        positions=positions,
+        hidden_states=hidden_states,
+        residual=None,
+    )
+
+    assert prefix_sum is None
+    assert hidden_states.shape == residual.shape == (2, 2)
+    assert layer.mlp.num_tokens == 2
+
+    hidden_states, prefix_sum, residual = layer(
+        positions=positions,
+        hidden_states=hidden_states,
+        residual=residual,
+    )
+
+    assert prefix_sum is None
+    assert hidden_states.shape == residual.shape == (2, 2)
+    assert layer.mlp.num_tokens == 2
+
+
+def test_kimi_attn_residual_states_stay_sequence_sharded(monkeypatch):
+    layer = object.__new__(kimi_model.KimiDecoderLayer)
+    nn.Module.__init__(layer)
+    layer.use_attn_res = True
+    layer.use_sequence_parallel = True
+    layer.prev_valid_blocks = 0
+    layer.block_write_idx = 0
+    layer.is_block_write_layer = False
+    layer.input_layernorm = _IdentityNorm()
+    layer.post_attention_layernorm = _IdentityNorm()
+    layer.self_attention_res_norm = _IdentityNorm()
+    layer.mlp_res_norm = _IdentityNorm()
+    layer.self_attention_res_proj = _Projection()
+    layer.mlp_res_proj = _Projection()
+    layer.mlp = _RecordingMoE()
+    layer._run_self_attn = MethodType(
+        lambda self, positions, hidden_states: hidden_states,
+        layer,
+    )
+
+    _mock_sequence_parallel_collectives(monkeypatch)
+    monkeypatch.setattr(
+        kimi_model,
+        "attn_res",
+        lambda prefix_sum, hidden_states, *args, **kwargs: (
+            prefix_sum if hidden_states is None else prefix_sum + hidden_states
+        ),
+    )
+
+    prefix_sum = kimi_model.sp_shard(torch.arange(6, dtype=torch.float32).view(3, 2))
+    block_residual = torch.zeros(2, 1, 2)
+    hidden_states, prefix_sum, block_residual = layer(
+        positions=torch.arange(3),
+        hidden_states=None,
+        prefix_sum=prefix_sum,
+        residual=block_residual,
+    )
+
+    assert hidden_states.shape == prefix_sum.shape == (2, 2)
+    assert block_residual.shape == (2, 1, 2)
+    assert layer.mlp.num_tokens == 2
+
+
+def test_kimi_mtp_restores_sequence_parallel_output(monkeypatch):
+    layer = object.__new__(kimi_mtp.KimiK3MultiTokenPredictorLayer)
+    nn.Module.__init__(layer)
+    layer.enorm = _IdentityNorm()
+    layer.hnorm = _IdentityNorm()
+    layer.eh_proj = nn.Identity()
+    object.__setattr__(layer, "mtp_block", _SequenceParallelMTPBlock())
+
+    final_norm = Mock(side_effect=lambda hidden_states: hidden_states + 1)
+    object.__setattr__(
+        layer,
+        "shared_head",
+        SimpleNamespace(norm=final_norm),
+    )
+
+    monkeypatch.setattr(
+        kimi_mtp,
+        "fused_mtp_input",
+        lambda positions, inputs_embeds, *args: inputs_embeds,
+    )
+    monkeypatch.setattr(
+        kimi_mtp,
+        "sp_shard",
+        lambda tensor: torch.nn.functional.pad(tensor, (0, 0, 0, 1))[:2],
+    )
+    monkeypatch.setattr(
+        kimi_mtp,
+        "sp_all_gather",
+        lambda tensor: torch.cat([tensor, tensor], dim=0),
+    )
+
+    inputs_embeds = torch.arange(6, dtype=torch.float32).view(3, 2)
+    logits_hidden_states, hidden_states = layer(
+        input_ids=torch.zeros(3, dtype=torch.long),
+        positions=torch.arange(3),
+        previous_hidden_states=torch.zeros_like(inputs_embeds),
+        inputs_embeds=inputs_embeds,
+    )
+
+    sharded_states = torch.nn.functional.pad(inputs_embeds, (0, 0, 0, 1))[:2]
+    expected_hidden_states = torch.cat(
+        [sharded_states * 5, sharded_states * 5],
+        dim=0,
+    )[:3]
+    torch.testing.assert_close(hidden_states, expected_hidden_states)
+    torch.testing.assert_close(logits_hidden_states, expected_hidden_states + 1)
+    final_norm.assert_called_once()
+    torch.testing.assert_close(final_norm.call_args.args[0], expected_hidden_states)
+
+
+def test_sp_all_gather_uses_custom_kernel(monkeypatch):
+    hidden_states = torch.arange(4, dtype=torch.float32).view(2, 2)
+    expected = torch.cat([hidden_states, hidden_states])
+    custom_all_gather = Mock(return_value=expected)
+    device_communicator = SimpleNamespace(
+        custom_all_gather=custom_all_gather,
+    )
+    monkeypatch.setattr(
+        sp_ops,
+        "get_tp_group",
+        lambda: SimpleNamespace(device_communicator=device_communicator),
+    )
+    fallback = Mock(side_effect=AssertionError("unexpected fallback"))
+    monkeypatch.setattr(sp_ops, "tensor_model_parallel_all_gather", fallback)
+
+    output = sp_ops.sp_all_gather(hidden_states)
+
+    torch.testing.assert_close(output, expected)
+    custom_all_gather.assert_called_once_with(hidden_states)
+    fallback.assert_not_called()
+
+
+def test_sp_reduce_scatter_uses_custom_kernel_after_padding(monkeypatch):
+    hidden_states = torch.arange(6, dtype=torch.float32).view(3, 2)
+    expected = torch.arange(4, dtype=torch.float32).view(2, 2)
+    custom_reduce_scatter = Mock(return_value=expected)
+    device_communicator = SimpleNamespace(
+        custom_reduce_scatter=custom_reduce_scatter,
+    )
+    monkeypatch.setattr(
+        sp_ops,
+        "get_tp_group",
+        lambda: SimpleNamespace(device_communicator=device_communicator),
+    )
+    monkeypatch.setattr(
+        sp_ops,
+        "get_tensor_model_parallel_world_size",
+        lambda: 2,
+    )
+    fallback = Mock(side_effect=AssertionError("unexpected fallback"))
+    monkeypatch.setattr(sp_ops, "tensor_model_parallel_reduce_scatter", fallback)
+
+    output = sp_ops.sp_reduce_scatter(hidden_states)
+
+    torch.testing.assert_close(output, expected)
+    padded = custom_reduce_scatter.call_args.args[0]
+    assert padded.shape == (4, 2)
+    torch.testing.assert_close(padded[:3], hidden_states)
+    torch.testing.assert_close(padded[3], torch.zeros(2))
+    fallback.assert_not_called()
+
+
+def test_sp_collectives_fall_back_without_custom_kernel(monkeypatch):
+    hidden_states = torch.arange(4, dtype=torch.float32).view(2, 2)
+    monkeypatch.setattr(
+        sp_ops,
+        "get_tp_group",
+        lambda: SimpleNamespace(device_communicator=None),
+    )
+    monkeypatch.setattr(
+        sp_ops,
+        "get_tensor_model_parallel_world_size",
+        lambda: 2,
+    )
+    all_gather = Mock(return_value=hidden_states)
+    reduce_scatter = Mock(return_value=hidden_states)
+    monkeypatch.setattr(sp_ops, "tensor_model_parallel_all_gather", all_gather)
+    monkeypatch.setattr(
+        sp_ops,
+        "tensor_model_parallel_reduce_scatter",
+        reduce_scatter,
+    )
+
+    torch.testing.assert_close(sp_ops.sp_all_gather(hidden_states), hidden_states)
+    torch.testing.assert_close(sp_ops.sp_reduce_scatter(hidden_states), hidden_states)
+    all_gather.assert_called_once_with(hidden_states, 0)
+    reduce_scatter.assert_called_once_with(hidden_states, 0)

tests/models/multimodal/generation/test_common.py

@@ -1242,7 +1242,7 @@ def test_custom_inputs_models(
         create_new_process_for_each_test=True,
     ),
 )
-@create_new_process_for_each_test()
+@create_new_process_for_each_test("spawn")
 def test_single_image_models_heavy(
     tmp_path: PosixPath,
     model_type: str,

vllm/_custom_ops.py

@@ -2733,6 +2733,52 @@ def fused_minimax_m3_qknorm_rope_kv_insert(
     )
 
 
+def fused_kda_decode(
+    x: torch.Tensor,
+    weight: torch.Tensor,
+    bias: torch.Tensor | None,
+    conv_state: torch.Tensor,
+    raw_g: torch.Tensor,
+    raw_beta: torch.Tensor,
+    A_log: torch.Tensor,
+    dt_bias: torch.Tensor,
+    state_indices: torch.Tensor,
+    state: torch.Tensor,
+    out: torch.Tensor | None = None,
+    lower_bound: float | None = None,
+    output_gate: torch.Tensor | None = None,
+    norm_weight: torch.Tensor | None = None,
+    norm_eps: float = 1e-5,
+) -> torch.Tensor:
+    if out is None:
+        out = torch.empty(
+            1,
+            x.shape[0],
+            raw_g.shape[2],
+            raw_g.shape[3],
+            dtype=x.dtype,
+            device=x.device,
+        )
+    torch.ops._C.fused_kda_decode(
+        x,
+        weight,
+        bias,
+        conv_state,
+        raw_g,
+        raw_beta,
+        A_log,
+        dt_bias,
+        state_indices,
+        state,
+        out,
+        lower_bound,
+        output_gate,
+        norm_weight,
+        norm_eps,
+    )
+    return out
+
+
 def concat_and_cache_mla(
     kv_c: torch.Tensor,
     k_pe: torch.Tensor,
@@ -2746,6 +2792,26 @@ def concat_and_cache_mla(
     )
 
 
+def concat_and_cache_mla_grouped(
+    kv_c: torch.Tensor,
+    k_pe: torch.Tensor,
+    kv_cache_ptrs: torch.Tensor,
+    slot_mapping: torch.Tensor,
+    block_size: int,
+    block_stride: int,
+    entry_stride: int,
+) -> None:
+    torch.ops._C_cache_ops.concat_and_cache_mla_grouped(
+        kv_c,
+        k_pe,
+        kv_cache_ptrs,
+        slot_mapping,
+        block_size,
+        block_stride,
+        entry_stride,
+    )
+
+
 def kimi_k3_attn_res(
     prefix: torch.Tensor,
     delta: torch.Tensor,
@@ -3040,6 +3106,63 @@ def all_reduce(
     torch.ops._C_custom_ar.all_reduce(fa, inp, out, reg_buffer, reg_buffer_sz_bytes)
 
 
+def custom_all_gather(
+    fa: int,
+    inp: torch.Tensor,
+    out: torch.Tensor,
+    reg_buffer: int,
+    reg_buffer_sz_bytes: int,
+) -> None:
+    torch.ops._C_custom_ar.custom_all_gather(
+        fa, inp, out, reg_buffer, reg_buffer_sz_bytes
+    )
+
+
+def mnnvl_lamport_all_gather(
+    fa: int,
+    inp: torch.Tensor,
+    out: torch.Tensor,
+    local_buffer: int,
+    multicast_buffer: int,
+    epoch_buffer: int,
+    stage_sz_bytes: int,
+) -> None:
+    torch.ops._C_custom_ar.mnnvl_lamport_all_gather(
+        fa,
+        inp,
+        out,
+        local_buffer,
+        multicast_buffer,
+        epoch_buffer,
+        stage_sz_bytes,
+    )
+
+
+def custom_reduce_scatter(
+    fa: int,
+    inp: torch.Tensor,
+    out: torch.Tensor,
+    reg_buffer: int,
+    reg_buffer_sz_bytes: int,
+) -> None:
+    torch.ops._C_custom_ar.custom_reduce_scatter(
+        fa, inp, out, reg_buffer, reg_buffer_sz_bytes
+    )
+
+
+def mnnvl_lamport_reduce_scatter(
+    fa: int,
+    inp: torch.Tensor,
+    out: torch.Tensor,
+    local_buffer: int,
+    epoch_buffer: int,
+    stage_sz_bytes: int,
+) -> None:
+    torch.ops._C_custom_ar.mnnvl_lamport_reduce_scatter(
+        fa, inp, out, local_buffer, epoch_buffer, stage_sz_bytes
+    )
+
+
 def dispose(fa: int) -> None:
     torch.ops._C_custom_ar.dispose(fa)
 
@@ -3144,18 +3267,18 @@ def dsv3_fused_a_gemm(
     output: torch.Tensor,
     mat_a: torch.Tensor,
     mat_b: torch.Tensor,
+    enable_pdl: bool = False,
 ) -> None:
-    """DeepSeek V3 fused A GEMM (SM 9.0+, bf16 only, 1-16 tokens).
+    """Low-latency fused-A-style GEMM (SM 9.0+, BF16, 1-16 tokens).
 
-    Computes output = mat_a @ mat_b.T where:
-      mat_a: [num_tokens, 7168] row-major bf16 (hidden states)
-      mat_b: [7168, 2112] column-major bf16 (weight transposed)
-      output: [num_tokens, 2112] row-major bf16
+    Computes ``output = mat_a @ mat_b`` for the compiled Kimi K3 and
+    DeepSeek V3 projection shapes. ``mat_a`` and ``output`` are row-major;
+    ``mat_b`` is the column-major transposed weight. ``enable_pdl`` permits
+    programmatic dependent launch for callers that have validated it.
 
-    Optimized for the DeepSeek V2/V3 QKV A-projection at small batch sizes.
-    Requires SM 9.0+ (Hopper).
+    Requires SM 9.0+.
     """
-    torch.ops._C.dsv3_fused_a_gemm(output, mat_a, mat_b)
+    torch.ops._C.dsv3_fused_a_gemm(output, mat_a, mat_b, enable_pdl)
 
 
 if hasattr(torch.ops._C, "weight_packed_linear"):

vllm/envs.py

@@ -269,6 +269,7 @@ if TYPE_CHECKING:
     VLLM_NCCL_INCLUDE_PATH: str | None = None
     VLLM_GC_DEBUG: str = ""
     VLLM_DEBUG_WORKSPACE: bool = False
+    VLLM_ENABLE_K3_LATENT_MOE_TAIL_FUSION: bool = False
     VLLM_DISABLE_SHARED_EXPERTS_STREAM: bool = False
     VLLM_SHARED_EXPERTS_STREAM_TOKEN_THRESHOLD: int = 256
     VLLM_MULTI_STREAM_GEMM_TOKEN_THRESHOLD: int = 1024
@@ -1894,6 +1895,11 @@ environment_variables: dict[str, Callable[[], Any]] = {
     # Debug workspace allocations.
     # logging of workspace resize operations.
     "VLLM_DEBUG_WORKSPACE": lambda: bool(int(os.getenv("VLLM_DEBUG_WORKSPACE", "0"))),
+    # Enable the experimental Kimi K3 latent-MoE tail fusion.
+    # Currently supported only on SM100 with TP=8/16 and BF16.
+    "VLLM_ENABLE_K3_LATENT_MOE_TAIL_FUSION": lambda: bool(
+        int(os.getenv("VLLM_ENABLE_K3_LATENT_MOE_TAIL_FUSION", "0"))
+    ),
     # Disables parallel execution of shared_experts via separate cuda stream
     "VLLM_DISABLE_SHARED_EXPERTS_STREAM": lambda: bool(
         int(os.getenv("VLLM_DISABLE_SHARED_EXPERTS_STREAM", "0"))

vllm/model_executor/kernels/linear/cute_dsl/_skinny_gemm.py

@@ -0,0 +1,180 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from __future__ import annotations
+
+import cutlass
+import cutlass.cute as cute
+from cuda.bindings.driver import CUstream
+from cutlass import const_expr
+
+
+class CuteSkinnyGemm:
+    """Shape-dynamic low-latency GEMM for small token counts.
+
+    Computes ``C[M, N] = A[M, K] @ B[N, K].T + residual`` with BF16 or FP16
+    inputs, FP32 accumulators, and an output matching the input dtype. The
+    residual term is optional and is added before the output conversion. N and
+    K are runtime values. The tiny M dimension is fully unrolled alongside a
+    small set of tuning parameters.
+    """
+
+    def __init__(
+        self,
+        *,
+        element_type,
+        num_rows: int,
+        block_size: int,
+        outputs_per_block: int,
+        vector_width: int = 8,
+        k_unroll: int = 1,
+        has_residual: bool = False,
+        use_pdl: bool = False,
+    ) -> None:
+        if block_size % cute.arch.WARP_SIZE != 0:
+            raise ValueError("block_size must be a multiple of the warp size")
+        self.element_type = element_type
+        self.num_rows = num_rows
+        self.block_size = block_size
+        self.outputs_per_block = outputs_per_block
+        self.vector_width = vector_width
+        self.k_unroll = k_unroll
+        self.has_residual = has_residual
+        self.use_pdl = use_pdl
+        self.num_warps = block_size // cute.arch.WARP_SIZE
+
+    @cute.jit
+    def __call__(
+        self,
+        gA: cute.Tensor,
+        gB: cute.Tensor,
+        gResidual: cute.Tensor,
+        gC: cute.Tensor,
+        stream: CUstream,
+    ) -> None:
+        n = cute.size(gB, mode=[0])
+        k = cute.size(gA, mode=[1])
+        copy_a = cute.make_copy_atom(
+            cute.nvgpu.CopyG2ROp(),
+            self.element_type,
+            num_bits_per_copy=self.vector_width * self.element_type.width,
+            load_cache_mode=cute.nvgpu.LoadCacheMode.ALWAYS,
+        )
+        copy_b = cute.make_copy_atom(
+            cute.nvgpu.CopyG2ROp(),
+            self.element_type,
+            num_bits_per_copy=self.vector_width * self.element_type.width,
+            load_cache_mode=cute.nvgpu.LoadCacheMode.STREAMING,
+        )
+        self.kernel(gA, gB, gResidual, gC, k, copy_a, copy_b).launch(
+            grid=[cute.ceil_div(n, self.outputs_per_block), 1, 1],
+            block=[self.block_size, 1, 1],
+            smem=self.num_rows * self.outputs_per_block * self.num_warps * 4,
+            stream=stream,
+            use_pdl=self.use_pdl,
+            min_blocks_per_mp=1,
+        )
+
+    @cute.kernel
+    def kernel(
+        self,
+        gA: cute.Tensor,
+        gB: cute.Tensor,
+        gResidual: cute.Tensor,
+        gC: cute.Tensor,
+        k_extent: cutlass.Int32,
+        copy_a: cute.CopyAtom,
+        copy_b: cute.CopyAtom,
+    ) -> None:
+        tidx, _, _ = cute.arch.thread_idx()
+        block_idx, _, _ = cute.arch.block_idx()
+        warp_idx = cute.arch.warp_idx()
+
+        num_rows: cutlass.Constexpr = self.num_rows
+        outputs_per_block: cutlass.Constexpr = self.outputs_per_block
+        vector_width: cutlass.Constexpr = self.vector_width
+        block_size: cutlass.Constexpr = self.block_size
+        num_warps: cutlass.Constexpr = self.num_warps
+
+        acc_layout = cute.make_layout(
+            (num_rows, outputs_per_block), stride=(outputs_per_block, 1)
+        )
+        acc = cute.make_rmem_tensor(acc_layout, cutlass.Float32)
+        acc.fill(0.0)
+
+        if const_expr(self.use_pdl):
+            cute.arch.griddepcontrol_wait()
+
+        n_base = block_idx * outputs_per_block
+        k_tile_size: cutlass.Constexpr = block_size * vector_width
+        num_k_tiles = k_extent // k_tile_size
+
+        gA_vec = cute.logical_divide(gA, (None, vector_width))
+        gB_vec = cute.logical_divide(gB, (None, vector_width))
+        # Layout after both divides is (M/N, K_TILE, K_LANE, K_VEC).
+        tA_all = cute.logical_divide(gA_vec, (None, (None, block_size)))
+        tB_all = cute.logical_divide(gB_vec, (None, (None, block_size)))
+        tA = tA_all[None, (None, (tidx, None))]
+
+        a_regs = cute.make_rmem_tensor(
+            cute.make_layout((num_rows, vector_width), stride=(vector_width, 1)),
+            self.element_type,
+        )
+        b_regs = cute.make_rmem_tensor(
+            cute.make_layout(
+                (outputs_per_block, vector_width), stride=(vector_width, 1)
+            ),
+            self.element_type,
+        )
+
+        for k_tile in cutlass.range(num_k_tiles, unroll=self.k_unroll):
+            for mi in cutlass.range_constexpr(num_rows):
+                cute.copy(copy_a, tA[mi, None, k_tile], a_regs[mi, None])
+
+            for ni in cutlass.range_constexpr(outputs_per_block):
+                n_idx = n_base + ni
+                tB = tB_all[n_idx, (None, (tidx, None))]
+                cute.copy(copy_b, tB[None, k_tile], b_regs[ni, None])
+
+            for vi in cutlass.range_constexpr(vector_width):
+                for mi in cutlass.range_constexpr(num_rows):
+                    for ni in cutlass.range_constexpr(outputs_per_block):
+                        acc[mi, ni] = acc[mi, ni] + a_regs[mi, vi].to(
+                            cutlass.Float32
+                        ) * b_regs[ni, vi].to(cutlass.Float32)
+
+        for mi in cutlass.range_constexpr(num_rows):
+            for ni in cutlass.range_constexpr(outputs_per_block):
+                acc[mi, ni] = cute.arch.warp_reduction_sum(acc[mi, ni])
+
+        smem_layout = cute.make_layout(
+            (num_rows, outputs_per_block, num_warps),
+            stride=(outputs_per_block * num_warps, num_warps, 1),
+        )
+        smem = cutlass.utils.SmemAllocator()
+        partials = smem.allocate_tensor(cutlass.Float32, smem_layout, byte_alignment=16)
+        with cute.arch.elect_one():
+            for mi in cutlass.range_constexpr(num_rows):
+                for ni in cutlass.range_constexpr(outputs_per_block):
+                    partials[mi, ni, warp_idx] = acc[mi, ni]
+
+        cute.arch.sync_threads()
+        if tidx == 0:
+            for mi in cutlass.range_constexpr(num_rows):
+                for ni in cutlass.range_constexpr(outputs_per_block):
+                    n_idx = n_base + ni
+                    total = (
+                        partials[mi, ni, None]
+                        .load()
+                        .reduce(
+                            cute.ReductionOp.ADD,
+                            init_val=cutlass.Float32(0.0),
+                            reduction_profile=0,
+                        )
+                    )
+                    if const_expr(self.has_residual):
+                        total += gResidual[mi, n_idx].to(cutlass.Float32)
+                    gC[mi, n_idx] = cutlass.Float32(total).to(self.element_type)
+
+        if const_expr(self.use_pdl):
+            cute.arch.griddepcontrol_launch_dependents()

vllm/model_executor/kernels/linear/cute_dsl/skinny_gemm.py

@@ -0,0 +1,252 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from __future__ import annotations
+
+import logging
+from collections.abc import Iterable
+from dataclasses import dataclass
+from functools import partial
+from typing import Any
+
+import torch
+
+logger = logging.getLogger(__name__)
+
+_cutedsl_available: bool | None = None
+
+
+@dataclass(frozen=True, slots=True)
+class SkinnyGemmConfig:
+    num_rows: int
+    block_size: int
+    outputs_per_block: int
+    k_unroll: int = 1
+    vector_width: int = 8
+
+
+class ShapeDynamicSkinnyGemm:
+    def __init__(self) -> None:
+        self._compiled: dict[tuple[torch.dtype, SkinnyGemmConfig, bool], Any] = {}
+        self._warmup_configs: set[tuple[torch.dtype, SkinnyGemmConfig, bool]] = set()
+        self._warmup_registered = False
+
+    @staticmethod
+    def is_available() -> bool:
+        global _cutedsl_available
+        if _cutedsl_available is not None:
+            return _cutedsl_available
+        try:
+            import cutlass  # noqa: F401
+            import cutlass.cute  # noqa: F401
+
+            _cutedsl_available = True
+        except ImportError:
+            _cutedsl_available = False
+            logger.info("cuteDSL is not available; skinny GEMM is disabled")
+        return _cutedsl_available
+
+    @staticmethod
+    def _config(m: int, n: int, k: int) -> SkinnyGemmConfig:
+        num_rows = m
+        wide_block = 224
+        if m == 1 and k >= 7168 and k % (wide_block * 8) == 0:
+            if n % 3 == 0:
+                k_unroll = 2 if n <= 2304 else 4
+                return SkinnyGemmConfig(num_rows, wide_block, 3, k_unroll)
+            if 2304 < n < 4096 and n % 2 == 0:
+                return SkinnyGemmConfig(num_rows, wide_block, 2, k_unroll=4)
+
+        if k <= 2048 or k % (128 * 8) != 0:
+            outputs_per_block = 2 if k <= 2048 else 4
+            if n % outputs_per_block:
+                outputs_per_block = 1
+            if k % (64 * 8) == 0:
+                return SkinnyGemmConfig(num_rows, 64, outputs_per_block, 2)
+            if k % (32 * 8) == 0:
+                return SkinnyGemmConfig(num_rows, 32, outputs_per_block, 2)
+            return SkinnyGemmConfig(num_rows, 32, outputs_per_block, 2, vector_width=4)
+
+        block_size = 64 if 4096 <= n < 8192 else 128
+        outputs_per_block = 1 if m == 1 and n <= 2304 else 2
+        if n % outputs_per_block:
+            outputs_per_block = 1
+        k_unroll = 2 if n <= 2304 or n >= 16384 else 1
+        return SkinnyGemmConfig(
+            num_rows,
+            block_size,
+            outputs_per_block,
+            k_unroll=k_unroll,
+        )
+
+    @staticmethod
+    def _cutlass_dtype(dtype: torch.dtype):
+        from cutlass import BFloat16, Float16
+
+        return BFloat16 if dtype == torch.bfloat16 else Float16
+
+    @staticmethod
+    def _stream():
+        from cuda.bindings.driver import CUstream
+
+        from vllm.utils.torch_utils import current_stream
+
+        return CUstream(current_stream().cuda_stream)
+
+    @staticmethod
+    def _use_pdl() -> bool:
+        from vllm.platforms import current_platform
+
+        return current_platform.is_arch_support_pdl()
+
+    def _compile(
+        self,
+        dtype: torch.dtype,
+        config: SkinnyGemmConfig,
+        has_residual: bool,
+    ) -> None:
+        import cutlass.cute as cute
+        from quack.compile_utils import make_fake_tensor
+
+        from ._skinny_gemm import CuteSkinnyGemm
+
+        element_type = self._cutlass_dtype(dtype)
+        n = cute.sym_int(divisibility=config.outputs_per_block)
+        k = cute.sym_int(divisibility=config.block_size * config.vector_width)
+        a = make_fake_tensor(
+            element_type,
+            (config.num_rows, k),
+            divisibility=config.vector_width,
+        )
+        b = make_fake_tensor(
+            element_type,
+            (n, k),
+            divisibility=config.vector_width,
+        )
+        c = make_fake_tensor(element_type, (config.num_rows, n), divisibility=1)
+        residual = make_fake_tensor(element_type, (config.num_rows, n), divisibility=1)
+        kernel = CuteSkinnyGemm(
+            element_type=element_type,
+            num_rows=config.num_rows,
+            block_size=config.block_size,
+            outputs_per_block=config.outputs_per_block,
+            vector_width=config.vector_width,
+            k_unroll=config.k_unroll,
+            has_residual=has_residual,
+            use_pdl=self._use_pdl(),
+        )
+        self._compiled[(dtype, config, has_residual)] = cute.compile(
+            kernel,
+            a,
+            b,
+            residual,
+            c,
+            self._stream(),
+            options="--enable-tvm-ffi --ptxas-options -maxrregcount=64",
+        )
+
+    def request_warmup(
+        self,
+        dtype: torch.dtype,
+        shapes: Iterable[tuple[int, int, int]],
+    ) -> None:
+        """Request compilation of ``(M, N, K)`` shapes before graph capture."""
+        self.request_warmup_configs(
+            dtype,
+            (self._config(m, n, k) for m, n, k in shapes),
+        )
+
+    def request_warmup_configs(
+        self,
+        dtype: torch.dtype,
+        configs: Iterable[SkinnyGemmConfig],
+        *,
+        has_residual: bool = False,
+    ) -> None:
+        """Request compilation of explicit measured configs before capture."""
+        self._warmup_configs.update((dtype, config, has_residual) for config in configs)
+        if self._warmup_registered:
+            return
+        from vllm.model_executor.warmup.cutedsl_warmup import (
+            register_cutedsl_warmup_provider,
+        )
+
+        register_cutedsl_warmup_provider(self)
+        self._warmup_registered = True
+
+    def get_cutedsl_warmup_compile_units(self):
+        from vllm.model_executor.warmup.cutedsl_warmup import CuTeDSLCompileUnit
+
+        return tuple(
+            CuTeDSLCompileUnit(
+                name=(
+                    "shape-dynamic skinny GEMM with residual"
+                    if has_residual
+                    else "shape-dynamic skinny GEMM"
+                ),
+                key=("shape-dynamic-skinny-gemm", dtype, config, has_residual),
+                compile=partial(self._compile, dtype, config, has_residual),
+            )
+            for dtype, config, has_residual in sorted(
+                self._warmup_configs,
+                key=lambda item: (
+                    str(item[0]),
+                    item[1].num_rows,
+                    item[1].block_size,
+                    item[1].outputs_per_block,
+                    item[1].k_unroll,
+                    item[1].vector_width,
+                    item[2],
+                ),
+            )
+        )
+
+    def __call__(
+        self,
+        a: torch.Tensor,
+        b: torch.Tensor,
+        config: SkinnyGemmConfig | None = None,
+        residual: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        if a.dim() != 2 or b.dim() != 2:
+            raise ValueError("a and b must be 2D tensors")
+        if a.dtype not in (torch.bfloat16, torch.float16) or b.dtype != a.dtype:
+            raise ValueError("a and b must have the same BF16 or FP16 dtype")
+        if not a.is_cuda or not b.is_cuda or a.device != b.device:
+            raise ValueError("a and b must be CUDA tensors on the same device")
+        if not a.is_contiguous() or not b.is_contiguous():
+            raise ValueError("a and b must be contiguous")
+        if a.shape[1] != b.shape[1]:
+            raise ValueError("a and b must have matching K dimensions")
+        if not 1 <= a.shape[0] <= 16:
+            raise ValueError("shape-dynamic skinny GEMM requires 1 <= M <= 16")
+        if residual is not None:
+            if residual.dim() != 2 or residual.shape != (a.shape[0], b.shape[0]):
+                raise ValueError("residual must have shape (M, N)")
+            if residual.dtype != a.dtype:
+                raise ValueError("residual must have the same dtype as a and b")
+            if residual.device != a.device or not residual.is_cuda:
+                raise ValueError("residual must be on the same CUDA device")
+            if not residual.is_contiguous():
+                raise ValueError("residual must be contiguous")
+
+        config = config or self._config(a.shape[0], b.shape[0], a.shape[1])
+        if config.num_rows != a.shape[0]:
+            raise ValueError("config num_rows must match M")
+        if b.shape[0] % config.outputs_per_block != 0:
+            raise ValueError("N must be divisible by outputs_per_block")
+        if a.shape[1] % (config.block_size * config.vector_width) != 0:
+            raise ValueError(
+                "K must be divisible by block_size * vector_width for this config"
+            )
+        has_residual = residual is not None
+        cache_key = (a.dtype, config, has_residual)
+        if cache_key not in self._compiled:
+            self._compile(a.dtype, config, has_residual)
+        output = torch.empty((a.shape[0], b.shape[0]), dtype=a.dtype, device=a.device)
+        residual_arg = output if residual is None else residual
+        self._compiled[cache_key](a, b, residual_arg, output, self._stream())
+        return output
+
+
+shape_dynamic_skinny_gemm = ShapeDynamicSkinnyGemm()

vllm/model_executor/layers/activation.py

@@ -158,6 +158,51 @@ class SiluAndMul(CustomOp):
         return self.forward_native(x)
 
 
+@CustomOp.register("situ_and_mul")
+class SituAndMul(CustomOp):
+    """SituGLU activation used by Kimi models.
+
+    Computes beta * tanh(gate / beta) * sigmoid(gate) * up.  When
+    ``linear_beta`` is set, the up projection is also softly clipped with
+    linear_beta * tanh(up / linear_beta).
+    """
+
+    def __init__(
+        self,
+        beta: float = 1.0,
+        linear_beta: float | None = None,
+        *,
+        compile_native: bool = True,
+    ):
+        super().__init__(compile_native=compile_native)
+        self.beta = float(beta)
+        self.linear_beta = None if linear_beta is None else float(linear_beta)
+        if current_platform.is_cuda_alike():
+            self.op = torch.ops._C.situ_and_mul
+
+    def forward_native(self, x: torch.Tensor) -> torch.Tensor:
+        d = x.shape[-1] // 2
+        gate = x[..., :d].float()
+        up = x[..., d:].float()
+        gate = self.beta * torch.tanh(gate / self.beta) * torch.sigmoid(gate)
+        if self.linear_beta is not None:
+            up = self.linear_beta * torch.tanh(up / self.linear_beta)
+        return (gate * up).to(x.dtype)
+
+    def forward_cuda(self, x: torch.Tensor) -> torch.Tensor:
+        # Fused CUDA kernel: writes straight to `out`, no fp32 temporaries.
+        # linear_beta<=0 signals "unset" to the kernel (up passed through).
+        d = x.shape[-1] // 2
+        out = torch.empty(x.shape[:-1] + (d,), dtype=x.dtype, device=x.device)
+        self.op(
+            out, x, self.beta, -1.0 if self.linear_beta is None else self.linear_beta
+        )
+        return out
+
+    def forward_xpu(self, x: torch.Tensor) -> torch.Tensor:
+        return self.forward_native(x)
+
+
 @CustomOp.register("silu_and_mul_with_clamp")
 class SiluAndMulWithClamp(CustomOp):
     """SwiGLU activation with input clamping (used by some MoE shared experts).

vllm/model_executor/layers/fused_moe/activation.py

@@ -22,6 +22,7 @@ class MoEActivation(Enum):
     # expects the *packed* layout ([all gates; all ups]), as produced by a
     # MergedColumnParallelLinear gate_up_proj (e.g. MiniMax-M3).
     SWIGLUOAI = "swigluoai"
+    SITU = "situ"
     SWIGLUOAI_UNINTERLEAVE = "swigluoai_uninterleave"
     SWIGLUSTEP = "swiglustep"
 
@@ -77,6 +78,7 @@ _CUSTOM_OP_NAMES: dict[MoEActivation, str] = {
     MoEActivation.SILU: "silu_and_mul",
     MoEActivation.GELU: "gelu_and_mul",
     MoEActivation.GELU_TANH: "gelu_tanh_and_mul",
+    MoEActivation.SITU: "situ_and_mul",
     MoEActivation.SWIGLUOAI: "swigluoai_and_mul",
     MoEActivation.SWIGLUOAI_UNINTERLEAVE: "silu_and_mul_with_clamp",
     MoEActivation.SWIGLUSTEP: "swiglustep_and_mul",
@@ -133,6 +135,8 @@ def apply_moe_activation(
     beta: float = 0.0,
     topk_ids: torch.Tensor | None = None,
     expert_map: torch.Tensor | None = None,
+    activation_situ_beta: float | None = None,
+    activation_situ_linear_beta: float | None = None,
 ) -> torch.Tensor:
     """Apply MoE activation function.
 
@@ -163,6 +167,25 @@ def apply_moe_activation(
         torch.ops._C.gelu_and_mul(output, input)
     elif activation == MoEActivation.GELU_TANH:
         torch.ops._C.gelu_tanh_and_mul(output, input)
+    elif activation == MoEActivation.SITU:
+        # Fused CUDA kernel: writes straight to `output`, no fp32 temporaries.
+        # (The pure-torch fallback below upcast both halves to fp32 and
+        # allocated ~8 temporaries per call, blowing up MoE memory.)
+        # Both betas come from FusedMoEConfig; a missing beta means the caller
+        # bypassed the config plumbing, so fail rather than silently use 1.0.
+        # linear_beta is genuinely optional: <= 0 signals "unset" to the kernel
+        # (up passed through), matching SituAndMul(linear_beta=None).
+        assert activation_situ_beta is not None, (
+            "SITU requires activation_situ_beta from FusedMoEConfig"
+        )
+        torch.ops._C.situ_and_mul(
+            output,
+            input,
+            activation_situ_beta,
+            -1.0
+            if activation_situ_linear_beta is None
+            else activation_situ_linear_beta,
+        )
     elif activation == MoEActivation.SWIGLUOAI:
         torch.ops._C.swigluoai_and_mul(output, input)
     elif activation == MoEActivation.SWIGLUOAI_UNINTERLEAVE:

vllm/model_executor/layers/fused_moe/config.py

@@ -1314,6 +1314,10 @@ class FusedMoEConfig:
     swiglu_alpha: float | None = None
     swiglu_beta: float | None = None
 
+    # SituGLU parameters used by Kimi sit(u/v2) activations.
+    activation_situ_beta: float | None = None
+    activation_situ_linear_beta: float | None = None
+
     max_capture_size: int = 0
 
     # Set by __post_init__

vllm/model_executor/layers/fused_moe/experts/deep_gemm_moe.py

@@ -444,7 +444,14 @@ class DeepGemmFP4Experts(mk.FusedMoEExpertsModular):
 
     @staticmethod
     def _supports_activation(activation: MoEActivation) -> bool:
-        return activation in [MoEActivation.SILU, MoEActivation.SWIGLUSTEP]
+        # SILU has fused gate+mul+quant kernels; SWIGLUSTEP/SITU take the
+        # general path (activation applied via self.activation, which forwards
+        # the situ betas, then FP8 requant).
+        return activation in [
+            MoEActivation.SILU,
+            MoEActivation.SWIGLUSTEP,
+            MoEActivation.SITU,
+        ]
 
     @staticmethod
     def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bool:
@@ -488,16 +495,15 @@ class DeepGemmFP4Experts(mk.FusedMoEExpertsModular):
         M_sum, N = input.size()
         activation_out_dim = self.adjust_N_for_activation(N, activation)
 
-        if scale_fmt == DeepGemmQuantScaleFMT.UE8M0:
-            assert activation == MoEActivation.SILU
-            return fused_silu_mul_fp8_quant_packed(
-                input=input,
-                output_q=output,
-                group_size=block_k,
-                clamp_limit=self.gemm1_clamp_limit,
-            )
-
         if activation == MoEActivation.SILU:
+            # Fused gate+mul+quant kernels for the common SILU case.
+            if scale_fmt == DeepGemmQuantScaleFMT.UE8M0:
+                return fused_silu_mul_fp8_quant_packed(
+                    input=input,
+                    output_q=output,
+                    group_size=block_k,
+                    clamp_limit=self.gemm1_clamp_limit,
+                )
             use_ue8m0 = scale_fmt == DeepGemmQuantScaleFMT.FLOAT32_CEIL_UE8M0
             return silu_mul_per_token_group_quant_fp8_colmajor(
                 input=input,
@@ -506,12 +512,24 @@ class DeepGemmFP4Experts(mk.FusedMoEExpertsModular):
                 clamp_limit=self.gemm1_clamp_limit,
             )
 
+        # General gated activations (SWIGLUSTEP, SITU): apply the activation
+        # (self.activation forwards the situ betas from moe_config) then
+        # FP8-requant into the layout DeepGEMM expects for this scale format.
         act_out = torch.empty(
             (M_sum, activation_out_dim), dtype=input.dtype, device=input.device
         )
         self.activation(activation, act_out, input)
+        if scale_fmt == DeepGemmQuantScaleFMT.UE8M0:
+            return per_token_group_quant_fp8_packed_for_deepgemm(
+                act_out, block_k, use_ue8m0=True, out_q=output
+            )
+        use_ue8m0 = scale_fmt == DeepGemmQuantScaleFMT.FLOAT32_CEIL_UE8M0
         return per_token_group_quant_fp8(
-            act_out, block_k, column_major_scales=True, out_q=output
+            act_out,
+            block_k,
+            column_major_scales=True,
+            out_q=output,
+            use_ue8m0=use_ue8m0,
         )
 
     def apply(

vllm/model_executor/layers/fused_moe/experts/marlin_moe.py

@@ -93,6 +93,8 @@ def _fused_marlin_moe(
     clamp_limit: float | None = None,
     gemm1_alpha: float = 1.0,
     gemm1_beta: float = 0.0,
+    activation_situ_beta: float | None = None,
+    activation_situ_linear_beta: float | None = None,
 ) -> torch.Tensor:
     assert hidden_states.ndim == 2
     M, K = hidden_states.size()
@@ -171,6 +173,8 @@ def _fused_marlin_moe(
         beta=gemm1_beta,
         topk_ids=topk_ids,
         expert_map=expert_map,
+        activation_situ_beta=activation_situ_beta,
+        activation_situ_linear_beta=activation_situ_linear_beta,
     )
 
     if output is None:
@@ -256,6 +260,8 @@ def fused_marlin_moe(
     clamp_limit: float | None = None,
     gemm1_alpha: float = 1.0,
     gemm1_beta: float = 0.0,
+    activation_situ_beta: float | None = None,
+    activation_situ_linear_beta: float | None = None,
 ) -> torch.Tensor:
     """
     This function computes a Mixture of Experts (MoE) layer using two sets of
@@ -357,6 +363,8 @@ def fused_marlin_moe(
         num_tokens_post_padded=num_tokens_post_padded,
         activation=activation,
         activation_func=activation_func,
+        activation_situ_beta=activation_situ_beta,
+        activation_situ_linear_beta=activation_situ_linear_beta,
         input_global_scale1=input_global_scale1,
         input_global_scale2=input_global_scale2,
         global_scale1=global_scale1,
@@ -423,6 +431,9 @@ def batched_fused_marlin_moe(
     clamp_limit: float | None = None,
     gemm1_alpha: float = 1.0,
     gemm1_beta: float = 0.0,
+    activation_func: Callable[..., None] = apply_moe_activation,
+    activation_situ_beta: float | None = None,
+    activation_situ_linear_beta: float | None = None,
 ) -> torch.Tensor:
     """
     This function massages the inputs so the batched hidden_states can be
@@ -530,6 +541,9 @@ def batched_fused_marlin_moe(
         quant_type=quant_type,
         apply_router_weight_on_input=apply_router_weight_on_input,
         activation=activation,
+        activation_func=activation_func,
+        activation_situ_beta=activation_situ_beta,
+        activation_situ_linear_beta=activation_situ_linear_beta,
         expert_map=expert_map,
         block_size_m=block_size_m,
         sorted_token_ids=sorted_token_ids,
@@ -645,6 +659,7 @@ class MarlinExpertsBase(mk.FusedMoEExpertsModular):
             MoEActivation.SILU,
             MoEActivation.GELU,
             MoEActivation.GELU_TANH,
+            MoEActivation.SITU,
             MoEActivation.SWIGLUOAI,
             MoEActivation.SWIGLUOAI_UNINTERLEAVE,
             MoEActivation.SWIGLUSTEP,
@@ -793,6 +808,10 @@ class MarlinExperts(LoRAExpertsMixin, MarlinExpertsBase):
                 global_num_experts=global_num_experts,
                 activation=activation,
                 activation_func=self.activation,
+                activation_situ_beta=self.moe_config.activation_situ_beta,
+                activation_situ_linear_beta=(
+                    self.moe_config.activation_situ_linear_beta
+                ),
                 moe_sum=self.moe_sum,
                 expert_map=expert_map,
                 output=output,
@@ -833,6 +852,8 @@ class MarlinExperts(LoRAExpertsMixin, MarlinExpertsBase):
             beta: float = 0.0,
             topk_ids: torch.Tensor | None = None,
             expert_map: torch.Tensor | None = None,
+            activation_situ_beta: float | None = None,
+            activation_situ_linear_beta: float | None = None,
         ) -> None:
             # act_input  = intermediate_cache1 (M*topk, 2N for gated)
             # act_output = intermediate_cache2 (M*topk, N)
@@ -871,6 +892,8 @@ class MarlinExperts(LoRAExpertsMixin, MarlinExpertsBase):
                 beta=beta,
                 topk_ids=topk_ids,
                 expert_map=expert_map,
+                activation_situ_beta=activation_situ_beta,
+                activation_situ_linear_beta=activation_situ_linear_beta,
             )
             lora_state["cache2"] = act_output
 
@@ -918,6 +941,8 @@ class MarlinExperts(LoRAExpertsMixin, MarlinExpertsBase):
             global_num_experts=global_num_experts,
             activation=activation,
             activation_func=activation_with_lora,
+            activation_situ_beta=self.moe_config.activation_situ_beta,
+            activation_situ_linear_beta=self.moe_config.activation_situ_linear_beta,
             moe_sum=moe_sum_with_lora,
             expert_map=expert_map,
             output=output,
@@ -1021,6 +1046,41 @@ class BatchedMarlinExperts(MarlinExpertsBase):
         apply_router_weight_on_input: bool,
     ):
         assert expert_tokens_meta is not None, "Num valid tokens per batch is required"
+
+        def activation_func(
+            act: MoEActivation,
+            act_output: torch.Tensor,
+            act_input: torch.Tensor,
+            *,
+            activation_situ_beta: float | None = None,
+            activation_situ_linear_beta: float | None = None,
+            **kwargs,
+        ) -> None:
+            if act != MoEActivation.SITU:
+                self.activation(
+                    act,
+                    act_output,
+                    act_input,
+                    activation_situ_beta=activation_situ_beta,
+                    activation_situ_linear_beta=activation_situ_linear_beta,
+                    **kwargs,
+                )
+                return
+
+            num_experts, max_num_tokens = hidden_states.shape[:2]
+            beta = activation_situ_beta
+            linear_beta = activation_situ_linear_beta
+            assert beta is not None, (
+                "SITU requires activation_situ_beta from FusedMoEConfig"
+            )
+            torch.ops._C.masked_situ_and_mul(
+                act_output.view(num_experts, max_num_tokens, -1),
+                act_input.view(num_experts, max_num_tokens, -1),
+                expert_tokens_meta.expert_num_tokens,
+                beta,
+                -1.0 if linear_beta is None else linear_beta,
+            )
+
         return batched_fused_marlin_moe(
             hidden_states=hidden_states,
             expert_num_tokens=expert_tokens_meta.expert_num_tokens,
@@ -1051,4 +1111,7 @@ class BatchedMarlinExperts(MarlinExpertsBase):
             clamp_limit=self.gemm1_clamp_limit,
             gemm1_alpha=self.gemm1_alpha,
             gemm1_beta=self.gemm1_beta,
+            activation_func=activation_func,
+            activation_situ_beta=self.moe_config.activation_situ_beta,
+            activation_situ_linear_beta=self.moe_config.activation_situ_linear_beta,
         )

vllm/model_executor/layers/fused_moe/experts/trtllm_mxfp4_moe.py

@@ -79,6 +79,36 @@ class TrtLlmMxfp4ExpertsBase:
         else:
             self.gemm1_clamp_limit = None
 
+        # SITU (SituGLU) TRTLLM-Gen kernel computes
+        #   left  = alpha * tanh(x0 / alpha) * sigmoid(x0)   # gate (x0)
+        #   right = beta  * tanh(x1 / beta)                  # up   (x1)
+        # which matches vLLM's situ_and_mul with (beta, linear_beta), so map
+        # situ beta -> gatedActAlpha (gemm1_alpha) and situ linear_beta ->
+        # gatedActBeta (gemm1_beta). Both must be > 0.
+        if moe_config.activation == MoEActivation.SITU:
+            situ_beta = moe_config.activation_situ_beta
+            situ_linear_beta = moe_config.activation_situ_linear_beta
+            assert situ_beta is not None and situ_beta > 0, (
+                "SITU requires activation_situ_beta > 0"
+            )
+            assert situ_linear_beta is not None and situ_linear_beta > 0, (
+                "TRTLLM SiTuGlu requires activation_situ_linear_beta > 0 "
+                "(the private cubin has no up-passthrough path)"
+            )
+            self.gemm1_alpha = torch.full(
+                (self.local_num_experts,),
+                float(situ_beta),
+                dtype=torch.float32,
+                device=device,
+            )
+            self.gemm1_beta = torch.full(
+                (self.local_num_experts,),
+                float(situ_linear_beta),
+                dtype=torch.float32,
+                device=device,
+            )
+            self.gemm1_clamp_limit = None
+
         self.max_capture_size = moe_config.max_capture_size
 
     @staticmethod
@@ -103,7 +133,19 @@ class TrtLlmMxfp4ExpertsBase:
 
     @staticmethod
     def _supports_activation(activation: MoEActivation) -> bool:
-        return activation in (MoEActivation.SWIGLUOAI, MoEActivation.SILU)
+        return activation in (
+            MoEActivation.SWIGLUOAI,
+            MoEActivation.SILU,
+            MoEActivation.SITU,
+        )
+
+    @staticmethod
+    def _flashinfer_activation_type(activation: MoEActivation) -> int:
+        from flashinfer.fused_moe.core import ActivationType
+
+        if activation == MoEActivation.SITU:
+            return ActivationType.Situ.value
+        return ActivationType.Swiglu.value
 
     @staticmethod
     def activation_format() -> mk.FusedMoEActivationFormat:
@@ -142,6 +184,7 @@ class TrtLlmMxfp4ExpertsMonolithic(
         activation_key: QuantKey | None,
     ) -> bool:
         return routing_method in [
+            RoutingMethodType.DeepSeekV3,
             RoutingMethodType.Renormalize,
             RoutingMethodType.RenormalizeNaive,
         ]
@@ -192,8 +235,8 @@ class TrtLlmMxfp4ExpertsMonolithic(
             device=hidden_states.device,
         )
         trtllm_fp4_block_scale_moe(
-            routing_logits=router_logits.to(torch.bfloat16),
-            routing_bias=None,
+            routing_logits=router_logits,
+            routing_bias=e_score_correction_bias,
             hidden_states=x_quant,
             hidden_states_scale=x_scale,
             gemm1_weights=w1,
@@ -210,14 +253,15 @@ class TrtLlmMxfp4ExpertsMonolithic(
             output2_scale_scalar=None,
             num_experts=global_num_experts,
             top_k=self.topk,
-            n_group=None,
-            topk_group=None,
+            n_group=(num_expert_group or 0),
+            topk_group=(topk_group or 0),
             intermediate_size=self.intermediate_size_per_partition,
             local_expert_offset=self.ep_rank * self.local_num_experts,
             local_num_experts=self.local_num_experts,
-            routed_scaling_factor=None,
+            routed_scaling_factor=routed_scaling_factor,
             routing_method_type=self.routing_method_type,
             do_finalize=True,
+            activation_type=self._flashinfer_activation_type(activation),
             tune_max_num_tokens=max(self.max_capture_size, 1),
             output=output,
             routing_replay_out=routing_replay_out,
@@ -338,6 +382,7 @@ class TrtLlmMxfp4ExpertsModular(TrtLlmMxfp4ExpertsBase, mk.FusedMoEExpertsModula
             "routing_method_type": RoutingMethodType.Renormalize,
             "do_finalize": True,
             "enable_pdl": True,
+            "activation_type": self._flashinfer_activation_type(activation),
             "output": output,
             "tune_max_num_tokens": max(self.max_capture_size, 1),
         }

vllm/model_executor/layers/fused_moe/experts/trtllm_nvfp4_moe.py

@@ -110,6 +110,27 @@ class TrtLlmNvFp4ExpertsBase:
             self.gemm1_alpha = None
             self.gemm1_beta = None
 
+        # SITU (Kimi SituGLU) TRTLLM-Gen kernel computes
+        # left=alpha*tanh(x0/alpha)*sigmoid(x0), right=beta*tanh(x1/beta),
+        # matching vLLM's situ_and_mul, so map situ beta -> gatedActAlpha
+        # (gemm1_alpha) and situ linear_beta -> gatedActBeta (gemm1_beta).
+        # These operate on the dequantized gate/up, so they are NOT folded by
+        # g1_alphas in process_weights_after_loading.
+        self.is_situ = moe_config.activation == MoEActivation.SITU
+        if self.is_situ:
+            situ_beta = moe_config.activation_situ_beta
+            situ_linear_beta = moe_config.activation_situ_linear_beta
+            assert situ_beta is not None and situ_beta > 0, (
+                "SITU requires activation_situ_beta > 0"
+            )
+            assert situ_linear_beta is not None and situ_linear_beta > 0, (
+                "TRTLLM SiTuGlu requires activation_situ_linear_beta > 0 "
+                "(the private cubin has no up-passthrough path)"
+            )
+            self.gemm1_alpha = _per_expert(situ_beta)
+            self.gemm1_beta = _per_expert(situ_linear_beta)
+            self.gemm1_clamp_limit = None
+
         logger.debug_once(
             "activation=%s, gemm1_alpha=%s, gemm1_beta=%s, gemm1_clamp_limit=%s",
             moe_config.activation,
@@ -142,7 +163,10 @@ class TrtLlmNvFp4ExpertsBase:
         # (via the in-place mul above) and never changes again, so this is a
         # static, per-expert constant. Register on the layer so EPLB
         # rearranges it alongside the other expert tensors.
-        if self.gemm1_clamp_limit is not None:
+        # SITU alpha/beta act on the dequantized gate/up (tanh clamps), not the
+        # raw GEMM1 accumulator, so they are registered as-is without the
+        # g1_alphas fold used by the SwiGLU-OAI clamp/beta below.
+        if self.gemm1_clamp_limit is not None and not self.is_situ:
             gemm1_clamp_limit = self.gemm1_clamp_limit / self.quant_config.g1_alphas
             layer.register_parameter(
                 "gemm1_clamp_limit",
@@ -155,7 +179,11 @@ class TrtLlmNvFp4ExpertsBase:
         # raw. Register both on the layer so EPLB rearranges them with the
         # other per-expert tensors.
         if self.gemm1_beta is not None:
-            gemm1_beta = self.gemm1_beta / self.quant_config.g1_alphas
+            gemm1_beta = (
+                self.gemm1_beta
+                if self.is_situ
+                else self.gemm1_beta / self.quant_config.g1_alphas
+            )
             layer.register_parameter(
                 "gemm1_beta",
                 torch.nn.Parameter(gemm1_beta, requires_grad=False),
@@ -197,13 +225,15 @@ class TrtLlmNvFp4ExpertsBase:
 
     @staticmethod
     def _supports_activation(activation: MoEActivation) -> bool:
-        """Supports SiLU, RELU^2 non-gated, GELU, and clamped SwiGLU-OAI."""
+        """Supports SiLU, RELU^2 non-gated, GELU, clamped SwiGLU-OAI, and the
+        private SituGLU (SITU) cubin."""
         return activation in [
             MoEActivation.SILU,
             MoEActivation.RELU2_NO_MUL,
             MoEActivation.GELU,
             MoEActivation.GELU_TANH,
             MoEActivation.SWIGLUOAI_UNINTERLEAVE,
+            MoEActivation.SITU,
         ]
 
     @staticmethod

vllm/model_executor/layers/fused_moe/layer.py

@@ -120,6 +120,8 @@ def FusedMoE(
     swiglu_limit: float | None = None,
     swiglu_alpha: float | None = None,
     swiglu_beta: float | None = None,
+    activation_situ_beta: float | None = None,
+    activation_situ_linear_beta: float | None = None,
     e_score_correction_bias: torch.Tensor | None = None,
     apply_router_weight_on_input: bool = False,
     activation: str = "silu",
@@ -178,6 +180,8 @@ def FusedMoE(
         scoring_func: Scoring function for routing ("softmax" or others)
         routed_scaling_factor: Scaling factor applied to topk_weights or output
         swiglu_limit: SwiGLU activation limit
+        activation_situ_beta: SituGLU activation beta
+        activation_situ_linear_beta: SituGLU linear beta
         e_score_correction_bias: Expert score correction bias tensor
         apply_router_weight_on_input: Whether to apply router weights on input
         activation: Activation function name ("silu", "gelu", etc.)
@@ -352,6 +356,8 @@ def FusedMoE(
         swiglu_limit=swiglu_limit,
         swiglu_alpha=swiglu_alpha,
         swiglu_beta=swiglu_beta,
+        activation_situ_beta=activation_situ_beta,
+        activation_situ_linear_beta=activation_situ_linear_beta,
         max_capture_size=vllm_config.compilation_config.max_cudagraph_capture_size,
         skip_final_all_reduce=skip_final_all_reduce,
     )

vllm/model_executor/layers/fused_moe/modular_kernel.py

@@ -890,6 +890,8 @@ class FusedMoEExpertsModular(FusedMoEExperts):
         beta: float = 0.0,
         topk_ids: torch.Tensor | None = None,
         expert_map: torch.Tensor | None = None,
+        activation_situ_beta: float | None = None,
+        activation_situ_linear_beta: float | None = None,
     ) -> None:
         apply_moe_activation(
             activation,
@@ -900,6 +902,16 @@ class FusedMoEExpertsModular(FusedMoEExperts):
             beta=beta,
             topk_ids=topk_ids,
             expert_map=expert_map,
+            activation_situ_beta=(
+                self.moe_config.activation_situ_beta
+                if activation_situ_beta is None
+                else activation_situ_beta
+            ),
+            activation_situ_linear_beta=(
+                self.moe_config.activation_situ_linear_beta
+                if activation_situ_linear_beta is None
+                else activation_situ_linear_beta
+            ),
         )
 
     @abstractmethod
@@ -1307,6 +1319,14 @@ class FusedMoEKernelModularImpl:
             activation,
         )
 
+        use_output_alias = (
+            output_alias is not None
+            and output_alias.shape == fused_out.shape
+            and output_alias.dtype == fused_out.dtype
+            and output_alias.device == fused_out.device
+            and output_alias.is_contiguous()
+        )
+
         # If caller's output buffer already matches fused_out shape/dtype, alias
         # to skip the redundant copy in TopKWeightAndReduceNoOP.apply downstream.
         # This eliminates ~94% of __amd_rocclr_copyBuffer events (Copy 2 of the
@@ -1314,15 +1334,10 @@ class FusedMoEKernelModularImpl:
         if current_platform.is_rocm():
             from vllm._aiter_ops import rocm_aiter_ops
 
-            if (
-                rocm_aiter_ops.is_fused_moe_enabled()
-                and output_alias is not None
-                and output_alias.shape == fused_out.shape
-                and output_alias.dtype == fused_out.dtype
-                and output_alias.device == fused_out.device
-                and output_alias.is_contiguous()
-            ):
+            if use_output_alias and rocm_aiter_ops.is_fused_moe_enabled():
                 fused_out = output_alias
+        elif use_output_alias:
+            fused_out = output_alias
 
         self.fused_experts.apply(
             output=fused_out,

vllm/model_executor/layers/fused_moe/runner/latent_moe_runner.py

@@ -0,0 +1,255 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import torch
+
+import vllm.envs as envs
+from vllm.config import get_current_vllm_config
+from vllm.distributed import tensor_model_parallel_all_reduce
+from vllm.logger import init_logger
+from vllm.model_executor.layers.fused_allreduce_gemma_rms_norm import (
+    _AR_RESIDUAL_RMS_NORM,
+    _can_use_flashinfer,
+    flashinfer_trtllm_fused_allreduce_norm,
+)
+from vllm.model_executor.layers.layernorm import RMSNorm
+from vllm.utils.torch_utils import aux_stream, current_stream
+
+from .moe_runner import MoERunner, _unpack
+
+logger = init_logger(__name__)
+
+
+class LatentMoERunner(MoERunner):
+    """MoE runner for latent MoE with a replicated routed up-projection.
+
+    Fused path (tp>1, un-reduced combine output, shared expert, no SP):
+    concatenates the un-reduced latent partial (dim d) and the un-reduced
+    shared partial (dim D) into one contiguous buffer, all-reduces once, then
+    splits. The latent half is normed and up-projected locally (replicated
+    up-proj -> full hidden), and the shared add folds into the GEMM epilogue
+    (``torch.addmm``). One collective total, no post-reduction communication.
+
+    Native path: the replicated up-proj produces the full hidden dim on every
+    rank, so the base runner combines routed + shared correctly at any TP size
+    (using two collectives instead of the fused path's one).
+    """
+
+    def __init__(
+        self,
+        *args,
+        enable_k3_latent_moe_tail_fusion: bool = False,
+        **kwargs,
+    ) -> None:
+        super().__init__(*args, **kwargs)
+        self.enable_k3_latent_moe_tail_fusion = enable_k3_latent_moe_tail_fusion
+        use_fused_path = self._use_fused_path()
+        if (
+            self.enable_k3_latent_moe_tail_fusion
+            and use_fused_path
+            and self.moe_config.tp_size not in (8, 16)
+        ):
+            logger.warning_once(
+                "K3 latent-MoE tail fusion currently supports TP=8 and TP=16, "
+                "but TP=%d is configured. Falling back to the default path.",
+                self.moe_config.tp_size,
+            )
+            self.enable_k3_latent_moe_tail_fusion = False
+
+        if self.enable_k3_latent_moe_tail_fusion and use_fused_path:
+            vllm_config = get_current_vllm_config()
+            if vllm_config.parallel_config.use_ubatching:
+                raise ValueError(
+                    "K3 latent-MoE tail fusion does not support DBO or ubatching."
+                )
+            if vllm_config.model_config.enable_sleep_mode:
+                raise ValueError(
+                    "K3 latent-MoE tail fusion does not support sleep mode."
+                )
+            transform = self.routed_output_transform
+            assert transform is not None
+            norm = transform.norm
+            assert norm is not None
+            from vllm.models.kimi_k3.nvidia.ops.latent_moe_tail import (
+                KimiK3LatentMoETailOp,
+            )
+
+            op = KimiK3LatentMoETailOp.initialize(
+                hidden_size=transform.up_proj.weight.shape[0],
+                latent_size=norm.weight.shape[0],
+                dtype=norm.weight.dtype,
+                device=norm.weight.device,
+                rms_eps=norm.variance_epsilon,
+            )
+            self._k3_latent_moe_tail_op = op
+
+    def _get_zero_residual(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        """Read-only zero ``residual_in`` for the fused AR+RMSNorm kernel.
+
+        flashinfer requires a residual buffer even when there is no residual to
+        add.
+        """
+        numel = hidden_states.numel()
+        buf = getattr(self, "_zero_residual", None)
+        if (
+            buf is None
+            or buf.dtype != hidden_states.dtype
+            or buf.device != hidden_states.device
+            or buf.numel() < numel
+        ):
+            buf = torch.zeros(
+                numel, dtype=hidden_states.dtype, device=hidden_states.device
+            )
+            self._zero_residual = buf
+        return buf[:numel].view_as(hidden_states)
+
+    def _use_fused_path(self) -> bool:
+        # The fused path merges the latent and shared reductions into one
+        # all-reduce, so it needs actual TP parallelism, a shared expert (to
+        # concat), an un-reduced combine output, and no sequence parallelism.
+        return (
+            self.moe_config.tp_size > 1
+            and self._shared_experts is not None
+            and not self._fused_output_is_reduced
+            and not self.moe_config.is_sequence_parallel
+        )
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        router_logits: torch.Tensor,
+        input_ids: torch.Tensor | None = None,
+        shared_experts_input: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        if self._use_fused_path():
+            return self._fused_forward(
+                hidden_states, router_logits, input_ids, shared_experts_input
+            )
+        return super().forward(
+            hidden_states, router_logits, input_ids, shared_experts_input
+        )
+
+    def _fused_forward(
+        self,
+        hidden_states: torch.Tensor,
+        router_logits: torch.Tensor,
+        input_ids: torch.Tensor | None = None,
+        shared_experts_input: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        # When the caller pre-applies the routed input transform outside the
+        # runner (e.g. to overlap it on a separate stream), it passes the
+        # already-transformed routed input as ``hidden_states`` and the original
+        # hidden states as ``shared_experts_input``; skip the transform then.
+        if shared_experts_input is None:
+            hidden_states, shared_experts_input = self.apply_routed_input_transform(
+                hidden_states
+            )
+        hidden_states, og_hidden_dim_pre_xform, og_hidden_dim_post_xform = (
+            self._maybe_pad_hidden_states(
+                shared_experts_input,
+                hidden_states,
+            )
+        )
+
+        result = self._forward_entry(
+            hidden_states,
+            router_logits,
+            shared_experts_input,
+            input_ids,
+            self._encode_layer_name(),
+            self.moe_config.hidden_dim_unpadded
+            if self._quant_method.has_unpadded_output
+            else 0,
+        )
+
+        shared_output, fused_output = _unpack(result)
+        assert shared_output is not None
+
+        if og_hidden_dim_pre_xform is not None:
+            fused_output = fused_output[..., :og_hidden_dim_pre_xform]
+
+        transform = self.routed_output_transform
+        assert transform is not None
+
+        if self.enable_k3_latent_moe_tail_fusion:
+            op = self._k3_latent_moe_tail_op
+            if 0 < fused_output.shape[0] <= op.contract.max_num_tokens:
+                norm = transform.norm
+                assert norm is not None
+                result = op(
+                    fused_output,
+                    shared_output,
+                    norm.weight,
+                    transform.up_proj.weight,
+                )
+                result = self._maybe_reduce_final_output(
+                    result, og_hidden_dim_post_xform, output_is_reduced=True
+                )
+                return self._maybe_add_zero_expert_output(result)
+
+        fused_latent = None
+        if transform.norm is not None:
+            fused_latent = self.allreduce_norm_latent_out(fused_output, transform.norm)
+        else:
+            fused_latent = tensor_model_parallel_all_reduce(fused_output)
+
+        shared_expert_stream = (
+            aux_stream()
+            if shared_output.size(0) <= envs.VLLM_SHARED_EXPERTS_STREAM_TOKEN_THRESHOLD
+            else None
+        )
+        if shared_expert_stream is not None:
+            # overlap shared expert allreduce with latent up_proj
+            main = current_stream()
+            shared_output.record_stream(shared_expert_stream)
+            shared_expert_stream.wait_stream(main)
+            with torch.cuda.stream(shared_expert_stream):
+                shared_output = tensor_model_parallel_all_reduce(shared_output)
+            result = torch.mm(fused_latent, transform.up_proj.weight.t())
+            main.wait_stream(shared_expert_stream)
+        else:
+            shared_output = tensor_model_parallel_all_reduce(shared_output)
+            result = torch.mm(fused_latent, transform.up_proj.weight.t())
+        result.add_(shared_output)
+
+        # Output is already fully reduced; this only strips padding.
+        result = self._maybe_reduce_final_output(
+            result, og_hidden_dim_post_xform, output_is_reduced=True
+        )
+
+        return self._maybe_add_zero_expert_output(result)
+
+    def allreduce_norm_latent_out(
+        self,
+        hidden_states: torch.Tensor,
+        norm: RMSNorm,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """All-reduce + add residual + (standard) RMSNorm, fused via flashinfer."""
+        if self.moe_config.tp_size == 1:
+            return norm(hidden_states)
+
+        if flashinfer_trtllm_fused_allreduce_norm is not None:
+            ok, max_token_num = _can_use_flashinfer(
+                hidden_states, self.moe_config.tp_size
+            )
+            if ok:
+                norm_out = torch.empty_like(hidden_states)
+                # With norm_out provided, the kernel writes the new residual
+                # (all_reduce(hidden_states) + residual) into the hidden_states
+                # buffer and the normalized result into norm_out.
+                flashinfer_trtllm_fused_allreduce_norm(
+                    allreduce_in=hidden_states,
+                    residual=self._get_zero_residual(hidden_states),
+                    rms_gamma=norm.weight,
+                    rms_eps=norm.variance_epsilon,
+                    world_size=self.moe_config.tp_size,
+                    weight_bias=0.0,
+                    launch_with_pdl=True,
+                    fp32_acc=True,
+                    max_token_num=max_token_num,
+                    pattern_code=_AR_RESIDUAL_RMS_NORM,
+                    norm_out=norm_out,
+                )
+                return norm_out
+
+        reduced = tensor_model_parallel_all_reduce(hidden_states)
+        return norm(reduced)

vllm/model_executor/layers/fused_moe/runner/moe_runner.py

@@ -348,7 +348,8 @@ class MoERunner(MoERunnerInterface):
         return self.routed_experts.quant_method
 
     def apply_routed_input_transform(
-        self, hidden_states: torch.Tensor
+        self,
+        hidden_states: torch.Tensor,
     ) -> tuple[torch.Tensor, torch.Tensor | None]:
         """Apply transform for routed experts (e.g., latent projection).
 
@@ -416,6 +417,7 @@ class MoERunner(MoERunnerInterface):
     def _maybe_reduce_shared_expert_output(
         self,
         shared_output: torch.Tensor | None,
+        fused_output_is_reduced: bool | None = None,
     ) -> torch.Tensor | None:
         """All-reduce shared expert output when the combine kernel already
         reduced fused output.
@@ -425,18 +427,43 @@ class MoERunner(MoERunnerInterface):
         * If we have SP (TP=N, DP=M, EP), there is a separate AG step handled
           in the model.
         """
+        if fused_output_is_reduced is None:
+            fused_output_is_reduced = self._fused_output_is_reduced
+
         if (
             shared_output is not None
             and not self.moe_config.is_sequence_parallel
-            and self._fused_output_is_reduced
+            and fused_output_is_reduced
         ):
             shared_output = tensor_model_parallel_all_reduce(shared_output)
         return shared_output
 
+    def _maybe_reduce_routed_output_before_transform(
+        self,
+        fused_output: torch.Tensor,
+        fused_output_is_reduced: bool,
+    ) -> tuple[torch.Tensor, bool]:
+        """All-reduce latent routed output before its output transform.
+
+        Latent MoE output transforms may contain non-linear ops, e.g. RMSNorm.
+        TP partial routed outputs must be summed in latent space before such
+        transforms are applied.
+        """
+        if (
+            self.routed_output_transform is not None
+            and not self.moe_config.is_sequence_parallel
+            and (self.moe_config.tp_size > 1 or self.moe_config.ep_size > 1)
+            and not fused_output_is_reduced
+        ):
+            fused_output = tensor_model_parallel_all_reduce(fused_output)
+            fused_output_is_reduced = True
+        return fused_output, fused_output_is_reduced
+
     def _maybe_reduce_final_output(
         self,
         states: torch.Tensor,
         trunc_size: int | None,
+        output_is_reduced: bool | None = None,
     ) -> torch.Tensor:
         """All-reduce the combined output if needed.
 
@@ -455,11 +482,14 @@ class MoERunner(MoERunnerInterface):
         # We don't need to reduce the final output if:
         # - We are not running with TP or DP
         # - The MK already reduced the fused output itself.
+        if output_is_reduced is None:
+            output_is_reduced = self._fused_output_is_reduced
+
         if (
             not self.moe_config.is_sequence_parallel
             and not self.moe_config.skip_final_all_reduce
             and (self.moe_config.tp_size > 1 or self.moe_config.ep_size > 1)
-            and not self._fused_output_is_reduced
+            and not output_is_reduced
         ):
             states = tensor_model_parallel_all_reduce(states)
 
@@ -643,6 +673,7 @@ class MoERunner(MoERunnerInterface):
         hidden_states: torch.Tensor,
         router_logits: torch.Tensor,
         input_ids: torch.Tensor | None = None,
+        shared_experts_input: torch.Tensor | None = None,
     ) -> torch.Tensor:
         """Invoke the fused moe layer.
 
@@ -664,11 +695,16 @@ class MoERunner(MoERunnerInterface):
            _moe_forward and _moe_forward_shared must be split.
         """
 
-        # Apply transform for routed experts (e.g., latent projection
-        # for latent MoE)
-        hidden_states, shared_experts_input = self.apply_routed_input_transform(
-            hidden_states
-        )
+        # Apply transform for routed experts (e.g., latent projection for
+        # latent MoE). When the caller pre-applies the routed input transform
+        # outside the runner (e.g. to overlap it on a separate stream), it
+        # passes the already-transformed routed input as ``hidden_states`` and
+        # the original hidden states as ``shared_experts_input``; skip the
+        # transform in that case so shared experts still see the original input.
+        if shared_experts_input is None:
+            hidden_states, shared_experts_input = self.apply_routed_input_transform(
+                hidden_states
+            )
 
         # Record before `_maybe_pad_hidden_states` pads activations to match
         # `moe_config.hidden_dim`, e.g. after `align_trtllm_fp4_moe_hidden_dim_for_fi`
@@ -708,9 +744,22 @@ class MoERunner(MoERunnerInterface):
         if og_hidden_dim_pre_xform is not None:
             fused_output = fused_output[..., :og_hidden_dim_pre_xform]
 
-        # If combine kernel already reduced fused, reduce shared to match.
+        fused_output_is_reduced = self._fused_output_is_reduced
+
+        # Latent routed output has to be reduced before output transform,
+        # because the transform may include non-linear normalization.
+        fused_output, fused_output_is_reduced = (
+            self._maybe_reduce_routed_output_before_transform(
+                fused_output,
+                fused_output_is_reduced,
+            )
+        )
+
+        # If routed output is already reduced, reduce shared to match.
         # See note above re: the two all-reduce points.
-        shared_output = self._maybe_reduce_shared_expert_output(shared_output)
+        shared_output = self._maybe_reduce_shared_expert_output(
+            shared_output, fused_output_is_reduced
+        )
 
         shared_output, fused_output = self._maybe_apply_routed_scale_to_output(
             shared_output, fused_output
@@ -724,7 +773,9 @@ class MoERunner(MoERunnerInterface):
         else:
             result = fused_output
 
-        result = self._maybe_reduce_final_output(result, og_hidden_dim_post_xform)
+        result = self._maybe_reduce_final_output(
+            result, og_hidden_dim_post_xform, fused_output_is_reduced
+        )
 
         return self._maybe_add_zero_expert_output(result)
 

vllm/model_executor/layers/fused_moe/runner/moe_runner_interface.py

@@ -36,6 +36,7 @@ class MoERunnerInterface(PluggableLayer, ABC):
         hidden_states: torch.Tensor,
         router_logits: torch.Tensor,
         input_ids: torch.Tensor | None = None,
+        shared_experts_input: torch.Tensor | None = None,
     ) -> torch.Tensor:
         raise NotImplementedError
 

vllm/model_executor/layers/mamba/ops/causal_conv1d.py

@@ -1080,6 +1080,7 @@ def causal_conv1d_update(
     block_idx_last_scheduled_token: torch.Tensor | None = None,
     initial_state_idx: torch.Tensor | None = None,
     validate_data=False,
+    out: torch.Tensor | None = None,
 ):
     """
     x: Input tensor which can take the following shapes:
@@ -1117,7 +1118,8 @@ def causal_conv1d_update(
             for example: conv_state_indices = [null_block_id, 1, 20, null_block_id]
             in this case, the kernel will not process entries at
             indices 0 and 3
-    out: (batch, dim) or (batch, dim, seqlen) or (num_tokens, dim), same shape as `x`
+    out: optional output tensor with the same shape as `x`. When omitted,
+        the input is overwritten.
     """
     if validate_data:
         assert null_block_id is not None
@@ -1129,10 +1131,22 @@ def causal_conv1d_update(
 
     original_x_dtype = x.dtype
     x = x.to(conv_state.dtype)
+    if out is None:
+        out = x
+    else:
+        if out.shape != x.shape:
+            raise ValueError(
+                f"`out` shape {tuple(out.shape)} must match `x` shape {tuple(x.shape)}."
+            )
+        if out.dtype != original_x_dtype or out.device != x.device:
+            raise ValueError(
+                "`out` must have the same dtype and device as the input `x`."
+            )
     unsqueeze = query_start_loc is None and x.dim() == 2
     if unsqueeze:
         # make it (batch, dim, seqlen) with seqlen == 1
         x = x.unsqueeze(-1)
+        out = out.unsqueeze(-1)
     if query_start_loc is None:
         batch, dim, seqlen = x.shape
     else:
@@ -1159,8 +1173,6 @@ def causal_conv1d_update(
         assert num_cache_lines >= batch
         assert weight.stride(1) == 1  # Need this
 
-    # adopt the strategy in vLLM that overwrite on 'x' directly, rather than creating a new tensor 'o'
-    out = x
     stride_w_dim, stride_w_width = weight.stride()
 
     if query_start_loc is None:

vllm/model_executor/layers/mamba/ops/gather_initial_states.py

@@ -0,0 +1,83 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import torch
+
+from vllm.platforms import current_platform
+from vllm.triton_utils import tl, triton
+
+
+@triton.jit
+def _gather_initial_states_kernel(
+    state_ptr,
+    indices_ptr,
+    has_initial_state_ptr,
+    output_ptr,
+    stride_state_batch,
+    stride_indices,
+    stride_has_initial_state,
+    row_size: tl.constexpr,
+    BLOCK_SIZE: tl.constexpr,
+    launch_pdl: tl.constexpr,
+):
+    block_idx = tl.program_id(0)
+    batch_idx = tl.program_id(1)
+    offsets = block_idx * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
+    mask = offsets < row_size
+
+    if launch_pdl:
+        tl.extra.cuda.gdc_wait()
+        tl.extra.cuda.gdc_launch_dependents()
+
+    has_initial_state = tl.load(
+        has_initial_state_ptr + batch_idx * stride_has_initial_state
+    ).to(tl.int1)
+    state_idx = tl.load(indices_ptr + batch_idx * stride_indices).to(tl.int64)
+    state_idx = tl.where(has_initial_state, state_idx, 0)
+    values = tl.load(
+        state_ptr + state_idx * stride_state_batch + offsets,
+        mask=mask & has_initial_state,
+        other=0.0,
+    )
+    tl.store(output_ptr + batch_idx * row_size + offsets, values, mask=mask)
+
+
+def gather_initial_states(
+    state: torch.Tensor,
+    indices: torch.Tensor,
+    has_initial_state: torch.Tensor,
+) -> torch.Tensor:
+    """Gather dense state rows, replacing uninitialized rows with zeros."""
+    assert state.ndim >= 2
+    assert state.is_cuda
+    assert indices.ndim == 1 and has_initial_state.ndim == 1
+    assert indices.shape == has_initial_state.shape
+    assert indices.device == state.device
+    assert has_initial_state.device == state.device
+    assert indices.dtype in (torch.int32, torch.int64)
+    assert has_initial_state.dtype == torch.bool
+
+    row_size = state[0].numel()
+    # Mamba pages may pad stride(0), but each state row remains dense.
+    assert state[0].is_contiguous()
+    output = torch.empty(
+        (indices.numel(), *state.shape[1:]),
+        dtype=state.dtype,
+        device=state.device,
+    )
+    block_size = min(triton.next_power_of_2(row_size), 1024)
+    grid = (triton.cdiv(row_size, block_size), indices.numel())
+    _gather_initial_states_kernel[grid](
+        state,
+        indices,
+        has_initial_state,
+        output,
+        state.stride(0),
+        indices.stride(0),
+        has_initial_state.stride(0),
+        row_size=row_size,
+        BLOCK_SIZE=block_size,
+        num_warps=8,
+        launch_pdl=current_platform.is_arch_support_pdl(),
+    )
+    return output

vllm/models/common/__init__.py

@@ -0,0 +1,2 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

vllm/models/common/ops/__init__.py

@@ -0,0 +1,9 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Ops shared across model implementations."""
+
+from .fused_qk_rmsnorm import fused_q_kv_rmsnorm
+
+__all__ = [
+    "fused_q_kv_rmsnorm",
+]

vllm/models/common/ops/fused_qk_rmsnorm.py

@@ -0,0 +1,103 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import torch
+
+from vllm.platforms import current_platform
+from vllm.triton_utils import tl, triton
+
+
+@triton.jit
+def _fused_q_kv_rmsnorm_kernel(
+    q_ptr,
+    q_out_ptr,
+    q_weight_ptr,
+    q_in_stride,
+    q_out_stride,
+    kv_ptr,
+    kv_out_ptr,
+    kv_weight_ptr,
+    kv_in_stride,
+    kv_out_stride,
+    eps,
+    Q_SIZE: tl.constexpr,
+    KV_SIZE: tl.constexpr,
+    BLOCK_SIZE: tl.constexpr,
+    launch_pdl: tl.constexpr,
+):
+    # num_tokens goes on grid-x (max 2**31 - 1); task goes on grid-y.
+    # CUDA's grid-y/z are capped at 65535, so putting num_tokens there crashes
+    # the launch at max-num-batched-tokens >= 65536 with "invalid argument".
+    # int64: q_in_stride can be ~24K (128 heads × 192) and overflows int32
+    # past num_tokens ~87K under large chunked prefill.
+    token_idx = tl.program_id(0).to(tl.int64)
+    pid_task = tl.program_id(1)
+
+    if pid_task == 0:
+        SIZE = Q_SIZE
+        row_in = q_ptr + token_idx * q_in_stride
+        weight_ptr = q_weight_ptr
+        row_out = q_out_ptr + token_idx * q_out_stride
+    else:
+        SIZE = KV_SIZE
+        row_in = kv_ptr + token_idx * kv_in_stride
+        weight_ptr = kv_weight_ptr
+        row_out = kv_out_ptr + token_idx * kv_out_stride
+
+    if launch_pdl:
+        tl.extra.cuda.gdc_wait()
+        tl.extra.cuda.gdc_launch_dependents()
+
+    # RMSNorm in fp32 throughout — matches csrc/layernorm_kernels.cu's
+    # `(scalar_t)(x * s_variance * w)` and DeepseekV4's compressor kernel, which
+    # keep x, rrms, and w all in fp32 and perform a single cast at store.
+    block = tl.arange(0, BLOCK_SIZE)
+    mask = block < SIZE
+    x = tl.load(row_in + block, mask=mask, other=0.0).to(tl.float32)
+    variance = tl.sum(x * x, axis=0) / SIZE
+    rrms = tl.rsqrt(variance + eps)
+    w = tl.load(weight_ptr + block, mask=mask, other=0.0).to(tl.float32)
+    y = x * rrms * w
+    tl.store(row_out + block, y.to(row_out.dtype.element_ty), mask=mask)
+
+
+def fused_q_kv_rmsnorm(
+    qr: torch.Tensor,
+    kv: torch.Tensor,
+    q_weight: torch.Tensor,
+    kv_weight: torch.Tensor,
+    eps: float,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    assert qr.ndim == 2 and kv.ndim == 2
+    assert qr.shape[0] == kv.shape[0], (
+        f"token dim mismatch: qr={qr.shape}, kv={kv.shape}"
+    )
+    assert qr.stride(-1) == 1 and kv.stride(-1) == 1
+    assert q_weight.is_contiguous() and kv_weight.is_contiguous()
+
+    q_size = qr.shape[1]
+    kv_size = kv.shape[1]
+    num_tokens = qr.shape[0]
+    qr_out = torch.empty_like(qr)
+    kv_out = torch.empty_like(kv)
+    if num_tokens == 0:
+        return qr_out, kv_out
+
+    block_size = triton.next_power_of_2(max(q_size, kv_size))
+    _fused_q_kv_rmsnorm_kernel[(num_tokens, 2)](
+        qr,
+        qr_out,
+        q_weight,
+        qr.stride(0),
+        qr_out.stride(0),
+        kv,
+        kv_out,
+        kv_weight,
+        kv.stride(0),
+        kv_out.stride(0),
+        eps,
+        Q_SIZE=q_size,
+        KV_SIZE=kv_size,
+        BLOCK_SIZE=block_size,
+        launch_pdl=current_platform.is_arch_support_pdl(),
+    )
+    return qr_out, kv_out

vllm/models/inkling/nvidia/ops/fa4_rel_attention.py

@@ -130,13 +130,20 @@ def inkling_fa4_rel_attention(
     cute_window = (None, None) if window_size == (-1, -1) else window_size
 
     rel_logits = rel_logits.contiguous()
+    flash_attn_varlen_func: Callable[..., Any]
     if _use_sheared_bias():
-        from vllm.third_party.tml_fa4 import flash_attn_varlen_func
+        from vllm.third_party.tml_fa4 import (
+            flash_attn_varlen_func as tml_flash_attn_varlen_func,
+        )
 
+        flash_attn_varlen_func = tml_flash_attn_varlen_func
         bias_kwargs: dict[str, Any] = {"rel_bias": rel_logits}
     else:
-        from vllm.vllm_flash_attn.cute import flash_attn_varlen_func
+        from vllm.vllm_flash_attn.cute import (
+            flash_attn_varlen_func as cute_flash_attn_varlen_func,
+        )
 
+        flash_attn_varlen_func = cute_flash_attn_varlen_func
         bias_kwargs = {
             "score_mod": _get_score_mod(rel_extent),
             "aux_tensors": [rel_logits],

vllm/models/kimi_k3/__init__.py

@@ -1,2 +1,28 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Kimi K3 model — hardware-isolated entry point.
+
+The implementation lives under ``nvidia/`` and ``amd/``; this module picks the
+right one for the current platform and re-exports the public classes used by
+the model registry. (Mirrors ``vllm.models.minimax_m3``.)
+"""
+
+from typing import TYPE_CHECKING
+
+from vllm.platforms import current_platform
+
+# The NVIDIA branch is the static default that type-checkers see; the ROCm
+# branch overrides it at runtime (kept type-compatible via type: ignore).
+if TYPE_CHECKING or not current_platform.is_rocm():
+    from .nvidia.model import KimiK3ForConditionalGeneration, KimiLinearForCausalLM
+    from .nvidia.mtp import KimiK3MTP
+else:
+    from .amd.linear import KimiLinearForCausalLM  # type: ignore[assignment]
+    from .amd.model import KimiK3ForConditionalGeneration  # type: ignore[assignment]
+    from .amd.mtp import KimiK3MTP  # type: ignore[assignment]
+
+__all__ = [
+    "KimiK3ForConditionalGeneration",
+    "KimiK3MTP",
+    "KimiLinearForCausalLM",
+]

vllm/models/kimi_k3/amd/model.py

@@ -0,0 +1,249 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Kimi-K3 multimodal model implementation for vLLM."""
+
+from collections.abc import Iterable
+from typing import cast
+
+import torch
+from torch import nn
+
+from vllm.config import VllmConfig
+from vllm.model_executor.layers.quantization import QuantizationConfig
+from vllm.model_executor.layers.quantization.compressed_tensors import (
+    compressed_tensors,
+)
+from vllm.model_executor.models.interfaces import (
+    HasInnerState,
+    IsHybrid,
+    SupportsEagle3,
+    SupportsMultiModal,
+    SupportsPP,
+    SupportsQuant,
+)
+from vllm.model_executor.models.kimi_k25 import KimiK25MediaPixelInputs
+from vllm.model_executor.models.kimi_k25_vit import (
+    KimiK25MultiModalProjector,
+    MoonViT3dPretrainedModel,
+    vision_tower_forward,
+)
+from vllm.model_executor.models.utils import (
+    AutoWeightsLoader,
+    WeightsMapper,
+    init_vllm_registered_model,
+    maybe_prefix,
+)
+from vllm.model_executor.models.vision import is_vit_use_data_parallel
+from vllm.multimodal import MULTIMODAL_REGISTRY
+from vllm.multimodal.inputs import NestedTensors
+from vllm.platforms import current_platform
+from vllm.sequence import IntermediateTensors
+from vllm.transformers_utils.configs.kimi_k3 import KimiK3Config
+
+from ..common.mm_preprocess import (
+    KimiK3DummyInputsBuilder,
+    KimiK3MultiModalProcessor,
+    KimiK3ProcessingInfo,
+)
+from .linear import KimiLinearForCausalLM
+
+
+@MULTIMODAL_REGISTRY.register_processor(
+    KimiK3MultiModalProcessor,
+    info=KimiK3ProcessingInfo,
+    dummy_inputs=KimiK3DummyInputsBuilder,
+)
+class KimiK3ForConditionalGeneration(
+    nn.Module,
+    SupportsMultiModal,
+    SupportsPP,
+    SupportsQuant,
+    SupportsEagle3,
+    HasInnerState,
+    IsHybrid,
+):
+    """Kimi-K3 model with Kimi-K2.5 vision and KimiLinear text."""
+
+    supports_encoder_tp_data = True
+
+    hf_to_vllm_mapper = WeightsMapper(
+        orig_to_new_prefix={
+            "language_model.layers.": "language_model.model.layers.",
+            "mm_projector.proj.0": "mm_projector.linear_1",
+            "mm_projector.proj.2": "mm_projector.linear_2",
+        }
+    )
+
+    @classmethod
+    def get_placeholder_str(cls, modality: str, i: int) -> str | None:
+        if modality == "image":
+            return "<|kimi_image_placeholder|>"
+        raise ValueError(f"Unsupported modality: {modality}")
+
+    def __init__(
+        self,
+        vllm_config: VllmConfig,
+        prefix: str = "",
+    ) -> None:
+        super().__init__()
+        model_config = vllm_config.model_config
+        config: KimiK3Config = model_config.hf_config
+        self.config = config
+        quant_config = vllm_config.quant_config
+
+        multimodal_config = model_config.multimodal_config
+        assert multimodal_config is not None
+        self.use_data_parallel = is_vit_use_data_parallel(
+            config.vision_config.num_attention_heads
+        )
+        self.hidden_size = config.text_config.hidden_size
+        self.device = current_platform.current_device()
+
+        with self._mark_tower_model(vllm_config, "image"):
+            self.vision_tower = MoonViT3dPretrainedModel(
+                config.vision_config,
+                quant_config=self._maybe_ignore_quant_config(quant_config),
+                prefix=maybe_prefix(prefix, "vision_tower"),
+            )
+            if self._maybe_ignore_quant_config(quant_config) is not None:
+                self.vision_tower = self.vision_tower.to(device=self.device)
+            else:
+                self.vision_tower = self.vision_tower.to(
+                    device=self.device, dtype=model_config.dtype
+                )
+
+            self.mm_projector = KimiK25MultiModalProjector(
+                config=config.vision_config,
+                use_data_parallel=self.use_data_parallel,
+                quant_config=self._maybe_ignore_quant_config(quant_config),
+                prefix=maybe_prefix(prefix, "mm_projector"),
+            )
+            self.mm_projector = self.mm_projector.to(
+                device=self.device, dtype=model_config.dtype
+            )
+
+        self.quant_config = quant_config
+        with self._mark_language_model(vllm_config):
+            self.language_model = init_vllm_registered_model(
+                vllm_config=vllm_config,
+                hf_config=config.text_config,
+                prefix=maybe_prefix(prefix, "language_model"),
+                architectures=["KimiLinearForCausalLM"],
+            )
+        self.make_empty_intermediate_tensors = (  # type: ignore[method-assign]
+            self.language_model.make_empty_intermediate_tensors
+        )
+        self.media_placeholder: int = self.config.media_placeholder_token_id
+
+    def _maybe_ignore_quant_config(
+        self, quant_config: QuantizationConfig | None
+    ) -> QuantizationConfig | None:
+        if isinstance(quant_config, compressed_tensors.CompressedTensorsConfig):
+            return None
+        return quant_config
+
+    def _parse_and_validate_media_input(
+        self, **kwargs: object
+    ) -> KimiK25MediaPixelInputs | None:
+        pixel_values = kwargs.pop("pixel_values", None)
+        grid_thws = kwargs.pop("grid_thws", None)
+        if pixel_values is None:
+            return None
+
+        if isinstance(pixel_values, list):
+            pixel_values = torch.cat(cast(list[torch.Tensor], pixel_values), dim=0)
+        if not isinstance(pixel_values, torch.Tensor):
+            raise TypeError(
+                "pixel_values must be a tensor or a list of tensors, "
+                f"got {type(pixel_values)}"
+            )
+
+        if len(pixel_values.shape) == 5 or len(pixel_values.shape) == 3:
+            pixel_values = pixel_values.reshape(
+                pixel_values.shape[0] * pixel_values.shape[1], *pixel_values.shape[2:]
+            )
+
+        target_dtype = next(self.vision_tower.parameters()).dtype
+        pixel_values = pixel_values.to(target_dtype)
+        assert isinstance(grid_thws, torch.Tensor), (
+            f"expect grid_thws to be a tensor, got {type(grid_thws)}"
+        )
+        grid_thws = grid_thws.reshape(-1, grid_thws.shape[-1])
+        assert grid_thws.ndim == 2 and grid_thws.size(1) == 3, (
+            f"unexpected shape for grid_thws: {grid_thws.shape}"
+        )
+
+        return KimiK25MediaPixelInputs(
+            type="pixel_values",
+            pixel_values=pixel_values,
+            grid_thws=grid_thws,
+        )
+
+    def _process_media_input(
+        self, media_input: KimiK25MediaPixelInputs
+    ) -> list[torch.Tensor]:
+        media_features = vision_tower_forward(
+            self.vision_tower,
+            media_input["pixel_values"],
+            media_input["grid_thws"],
+            mm_projector=self.mm_projector,
+            use_data_parallel=self.use_data_parallel,
+        )
+        return media_features
+
+    def embed_multimodal(self, **kwargs: object) -> NestedTensors | None:
+        media_input = self._parse_and_validate_media_input(**kwargs)
+        if media_input is None:
+            return None
+        return self._process_media_input(media_input)
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        positions: torch.Tensor,
+        intermediate_tensors: IntermediateTensors | None = None,
+        inputs_embeds: torch.Tensor | None = None,
+        **kwargs: object,
+    ) -> IntermediateTensors:
+        if intermediate_tensors is not None:
+            inputs_embeds = None
+        hidden_states = self.language_model(
+            input_ids=input_ids,
+            positions=positions,
+            intermediate_tensors=intermediate_tensors,
+            inputs_embeds=inputs_embeds,
+        )
+        return hidden_states
+
+    def compute_logits(self, hidden_states: torch.Tensor, **kwargs) -> torch.Tensor:
+        return self.language_model.compute_logits(hidden_states)
+
+    def copy_inputs_before_cuda_graphs(self, input_buffers, **kwargs):
+        return self.language_model.mamba_cache.copy_inputs_before_cuda_graphs(
+            input_buffers, **kwargs
+        )
+
+    def get_seqlen_agnostic_capture_inputs(self, batch_size: int):
+        return self.language_model.mamba_cache.get_seqlen_agnostic_capture_inputs(
+            batch_size
+        )
+
+    @classmethod
+    def get_mamba_state_dtype_from_config(cls, vllm_config: VllmConfig):
+        text_config = vllm_config.model_config.hf_config.text_config
+        temp_vllm_config = vllm_config.with_hf_config(text_config)
+        return KimiLinearForCausalLM.get_mamba_state_dtype_from_config(temp_vllm_config)
+
+    @classmethod
+    def get_mamba_state_shape_from_config(cls, vllm_config: VllmConfig):
+        text_config = vllm_config.model_config.hf_config.text_config
+        temp_vllm_config = vllm_config.with_hf_config(text_config)
+        return KimiLinearForCausalLM.get_mamba_state_shape_from_config(temp_vllm_config)
+
+    @classmethod
+    def get_mamba_state_copy_func(cls):
+        return KimiLinearForCausalLM.get_mamba_state_copy_func()
+
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):
+        loader = AutoWeightsLoader(self)
+        return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)

vllm/models/kimi_k3/amd/mtp.py

@@ -0,0 +1,403 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Inference-only Kimi-K3 Multi-Token-Prediction (MTP) draft model."""
+
+import copy
+from collections.abc import Iterable
+
+import torch
+import torch.nn as nn
+
+from vllm.config import VllmConfig
+from vllm.logger import init_logger
+from vllm.model_executor.layers.fused_moe import (
+    fused_moe_make_expert_params_mapping,
+)
+from vllm.model_executor.layers.layernorm import RMSNorm
+from vllm.model_executor.layers.logits_processor import LogitsProcessor
+from vllm.model_executor.layers.quantization import QuantizationConfig
+from vllm.model_executor.layers.vocab_parallel_embedding import (
+    ParallelLMHead,
+    VocabParallelEmbedding,
+)
+from vllm.model_executor.model_loader.weight_utils import (
+    default_weight_loader,
+    maybe_remap_kv_scale_name,
+)
+from vllm.model_executor.models.utils import get_pp_missing_layer_names, maybe_prefix
+from vllm.sequence import IntermediateTensors
+from vllm.transformers_utils.configs.kimi_linear import KimiLinearConfig
+
+from ..common.mtp import fused_mtp_input
+from .linear import KimiDecoderLayer, get_spec_layer_idx_from_weight_name
+
+logger = init_logger(__name__)
+
+
+class SharedHead(nn.Module):
+    def __init__(
+        self,
+        config: KimiLinearConfig,
+        prefix: str,
+        quant_config: QuantizationConfig | None = None,
+    ) -> None:
+        super().__init__()
+        self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.head = ParallelLMHead(
+            config.vocab_size,
+            config.hidden_size,
+            quant_config=quant_config,
+            prefix=maybe_prefix(prefix, "head"),
+        )
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        return self.norm(hidden_states)
+
+
+class KimiK3MultiTokenPredictorLayer(nn.Module):
+    def __init__(
+        self,
+        config: KimiLinearConfig,
+        vllm_config: VllmConfig,
+        prefix: str,
+    ) -> None:
+        super().__init__()
+        self.config = config
+        quant_config = vllm_config.quant_config
+
+        self.enorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.hnorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.eh_proj = nn.Linear(config.hidden_size * 2, config.hidden_size, bias=False)
+
+        self.shared_head = SharedHead(
+            config=config, prefix=prefix, quant_config=quant_config
+        )
+        # The MTP block is a standard KimiDecoderLayer, but it must NOT use the
+        # attn-residual (block-residual) scheme even when the base model does:
+        # the draft starts from the target's hidden state, without the main
+        # model's accumulated per-block residual tensor. We disable it by
+        # shallow-copying the config with ``attn_res_block_size=None``.
+        block_config = copy.copy(config)
+        block_config.attn_res_block_size = None
+        # NOTE: the prefix must end in the numeric spec-layer index so that
+        # KimiDecoderLayer can parse ``layer_idx`` and pick MLA (full attn).
+        self.mtp_block = KimiDecoderLayer(block_config, vllm_config, prefix=prefix)
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        positions: torch.Tensor,
+        previous_hidden_states: torch.Tensor,
+        inputs_embeds: torch.Tensor | None = None,
+        spec_step_index: int = 0,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        assert inputs_embeds is not None
+        hidden_states = self.eh_proj(
+            fused_mtp_input(
+                positions,
+                inputs_embeds,
+                previous_hidden_states,
+                self.enorm.weight,
+                self.hnorm.weight,
+                self.enorm.variance_epsilon,
+            )
+        )
+
+        hidden_states, residual = self.mtp_block(
+            positions=positions,
+            hidden_states=hidden_states,
+            residual=None,
+        )
+        # Produce the normalized logits input and the pre-norm recurrent state
+        # in one fused add-RMSNorm launch.
+        logits_hidden_states, hidden_states = self.shared_head.norm(
+            hidden_states, residual
+        )
+        return logits_hidden_states, hidden_states
+
+
+class KimiK3MultiTokenPredictor(nn.Module):
+    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
+        super().__init__()
+        config: KimiLinearConfig = vllm_config.model_config.hf_text_config
+        self.config = config
+        self.mtp_start_layer_idx = config.num_hidden_layers
+        self.num_mtp_layers = config.num_nextn_predict_layers
+
+        self.layers = torch.nn.ModuleDict(
+            {
+                str(idx): KimiK3MultiTokenPredictorLayer(
+                    config, vllm_config, f"{prefix}.layers.{idx}"
+                )
+                for idx in range(
+                    self.mtp_start_layer_idx,
+                    self.mtp_start_layer_idx + self.num_mtp_layers,
+                )
+            }
+        )
+        self.embed_tokens = VocabParallelEmbedding(
+            config.vocab_size,
+            config.hidden_size,
+            prefix=maybe_prefix(prefix, "embed_tokens"),
+        )
+        self.logits_processor = LogitsProcessor(config.vocab_size)
+
+    def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
+        return self.embed_tokens(input_ids)
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        positions: torch.Tensor,
+        previous_hidden_states: torch.Tensor,
+        inputs_embeds: torch.Tensor | None = None,
+        spec_step_idx: int = 0,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        if inputs_embeds is None:
+            inputs_embeds = self.embed_tokens(input_ids)
+        current_step_idx = spec_step_idx % self.num_mtp_layers
+        return self.layers[str(self.mtp_start_layer_idx + current_step_idx)](
+            input_ids,
+            positions,
+            previous_hidden_states,
+            inputs_embeds,
+            current_step_idx,
+        )
+
+    def compute_logits(
+        self,
+        hidden_states: torch.Tensor,
+        spec_step_idx: int = 0,
+    ) -> torch.Tensor:
+        current_step_idx = spec_step_idx % self.num_mtp_layers
+        mtp_layer = self.layers[str(self.mtp_start_layer_idx + current_step_idx)]
+        logits = self.logits_processor(mtp_layer.shared_head.head, hidden_states)
+        return logits
+
+
+class KimiK3MTP(nn.Module):
+    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
+        super().__init__()
+        self.config = vllm_config.model_config.hf_text_config
+        self.quant_config = vllm_config.quant_config
+        self.model = KimiK3MultiTokenPredictor(
+            vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
+        )
+
+    def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
+        return self.model.embed_input_ids(input_ids)
+
+    def forward(
+        self,
+        input_ids: torch.Tensor | None,
+        positions: torch.Tensor,
+        hidden_states: torch.Tensor,
+        intermediate_tensors: IntermediateTensors | None = None,
+        inputs_embeds: torch.Tensor | None = None,
+        spec_step_idx: int = 0,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        return self.model(
+            input_ids,
+            positions,
+            hidden_states,
+            inputs_embeds,
+            spec_step_idx,
+        )
+
+    def compute_logits(
+        self,
+        hidden_states: torch.Tensor,
+        spec_step_idx: int = 0,
+    ) -> torch.Tensor | None:
+        return self.model.compute_logits(hidden_states, spec_step_idx)
+
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+        # Mirror KimiLinearForCausalLM.load_weights naming: leading-dot shard
+        # names, q_lora-conditional fused QKV, and w1/w2/w3 expert weights.
+        kda_config = self.config.linear_attn_config
+        use_full_rank_gate = bool(
+            kda_config and kda_config.get("use_full_rank_gate", False)
+        )
+        beta_shard_id = 5 if use_full_rank_gate else 3
+        stacked_params_mapping = [
+            # (param_name, shard_name, shard_id)
+            (".in_proj_qkvgfab", ".q_proj", 0),
+            (".in_proj_qkvgfab", ".k_proj", 1),
+            (".in_proj_qkvgfab", ".v_proj", 2),
+            (".in_proj_qkvgfab", ".b_proj", beta_shard_id),
+            (".in_proj_qkvgfab", ".f_a_proj", 4),
+            (".conv1d", ".q_conv1d", 0),
+            (".conv1d", ".k_conv1d", 1),
+            (".conv1d", ".v_conv1d", 2),
+            (".gate_up_proj", ".gate_proj", 0),
+            (".gate_up_proj", ".up_proj", 1),
+        ]
+        if use_full_rank_gate:
+            stacked_params_mapping.append((".in_proj_qkvgfab", ".g_proj", 3))
+        if getattr(self.config, "q_lora_rank", None) is not None:
+            stacked_params_mapping += [
+                (".fused_qkv_a_proj", ".q_a_proj", 0),
+                (".fused_qkv_a_proj", ".kv_a_proj_with_mqa", 1),
+            ]
+
+        expert_params_mapping = (
+            fused_moe_make_expert_params_mapping(
+                self,
+                ckpt_gate_proj_name="w1",
+                ckpt_down_proj_name="w2",
+                ckpt_up_proj_name="w3",
+                num_experts=self.config.num_experts,
+            )
+            if self.config.is_moe
+            else []
+        )
+
+        pp_missing_layer_names = get_pp_missing_layer_names(self)
+        params_dict = dict(self.named_parameters())
+        # Under the MXFP4 quant interface the routed experts register unpacked
+        # params (``w13_weight``), while the compressed-tensors checkpoint names
+        # them ``.weight_packed``. Rebind so the expert mapping resolves; scales
+        # already share the ``.weight_scale`` suffix.
+        experts_unpacked = not any(n.endswith("w13_weight_packed") for n in params_dict)
+        loaded_params: set[str] = set()
+        for name, loaded_weight in weights:
+            if "rotary_emb.inv_freq" in name:
+                continue
+            # The multimodal checkpoint prefixes text weights with
+            # ``language_model.``; strip it so names match this draft model's
+            # parameter paths (``model.layers.{i}.``). Non-text weights
+            # (vision_tower, mm_projector, ...) never match a spec layer below.
+            if name.startswith("language_model."):
+                name = name[len("language_model.") :]
+            if experts_unpacked and name.endswith(".weight_packed"):
+                name = name.replace(".weight_packed", ".weight")
+            spec_layer = get_spec_layer_idx_from_weight_name(self.config, name)
+            if spec_layer is None:
+                continue
+            name = self._rewrite_spec_layer_name(spec_layer, name)
+
+            for param_name, weight_name, shard_id in stacked_params_mapping:
+                if weight_name not in name:
+                    continue
+                # Routed experts (``.experts.{i}.w1/w2/w3``) are handled by the
+                # expert mapping below; skip them here. Shared experts
+                # (``.shared_experts.``) use gate/up_proj and fall through.
+                if ".experts." in name:
+                    continue
+                name_mapped = name.replace(weight_name, param_name)
+                # Only take this mapping if the fused destination actually
+                # exists (e.g. QKV fusion is only present when q_lora is used).
+                if name_mapped not in params_dict:
+                    continue
+                if name_mapped in pp_missing_layer_names:
+                    continue
+                name = name_mapped
+                param = params_dict[name]
+                weight_loader = param.weight_loader
+                weight_loader(param, loaded_weight, shard_id)
+                break
+            else:
+                for (
+                    expert_param_name,
+                    expert_weight_name,
+                    expert_id,
+                    expert_shard_id,
+                ) in expert_params_mapping:
+                    if expert_weight_name not in name:
+                        continue
+                    name_mapped = name.replace(expert_weight_name, expert_param_name)
+                    if name_mapped in pp_missing_layer_names:
+                        continue
+                    param = params_dict[name_mapped]
+                    weight_loader = param.weight_loader
+                    weight_loader(
+                        param,
+                        loaded_weight,
+                        name_mapped,
+                        shard_id=expert_shard_id,
+                        expert_id=expert_id,
+                    )
+                    name = name_mapped
+                    break
+                else:
+                    if name.endswith(".bias") and name not in params_dict:
+                        continue
+                    remapped_name = maybe_remap_kv_scale_name(name, params_dict)
+                    if remapped_name is None:
+                        continue
+                    name = remapped_name
+
+                    # The embedding is shared across MTP layers; only the first
+                    # spec layer carries the hoisted (non-".layers") copy.
+                    if spec_layer != self.model.mtp_start_layer_idx and (
+                        ".layers" not in name
+                    ):
+                        continue
+                    if name in pp_missing_layer_names:
+                        continue
+                    # The base model uses an attn-residual scheme whose per-layer
+                    # weights (self_attention_res_*, mlp_res_*) are not used by
+                    # the draft block; such names have no matching parameter and
+                    # are safely skipped.
+                    if name not in params_dict:
+                        continue
+
+                    param = params_dict[name]
+                    weight_loader = getattr(
+                        param, "weight_loader", default_weight_loader
+                    )
+                    weight_loader(param, loaded_weight)
+            loaded_params.add(name)
+
+        # Validate that weights were loaded for each expected MTP layer.
+        loaded_layers: set[int] = set()
+        for param_name in loaded_params:
+            spec_layer = get_spec_layer_idx_from_weight_name(self.config, param_name)
+            if spec_layer is not None:
+                loaded_layers.add(spec_layer)
+        for layer_idx in range(
+            self.model.mtp_start_layer_idx,
+            self.model.mtp_start_layer_idx + self.model.num_mtp_layers,
+        ):
+            if layer_idx not in loaded_layers:
+                raise ValueError(
+                    f"MTP speculative decoding layer {layer_idx} weights "
+                    f"missing from checkpoint. The checkpoint may not include "
+                    f"the MTP layer weights. Use a checkpoint that includes "
+                    f"MTP layer weights, or disable speculative decoding."
+                )
+
+        return loaded_params
+
+    def _rewrite_spec_layer_name(self, spec_layer: int, name: str) -> str:
+        """Rewrite a checkpoint weight name to this module's parameter path.
+
+        Top-level MTP submodules (enorm/hnorm/eh_proj/shared_head) stay under
+        ``model.layers.{spec_layer}.*``; the shared ``embed_tokens`` is hoisted
+        to ``model.*``; everything else is a transformer-block weight and gets
+        ``.mtp_block`` inserted.
+        """
+        spec_layer_weight_names = [
+            "embed_tokens",
+            "enorm",
+            "hnorm",
+            "eh_proj",
+            "shared_head",
+        ]
+        shared_weight_names = ["embed_tokens"]
+        spec_layer_weight = False
+        shared_weight = False
+        for weight_name in spec_layer_weight_names:
+            if weight_name in name:
+                spec_layer_weight = True
+                if weight_name in shared_weight_names:
+                    shared_weight = True
+                break
+        if not spec_layer_weight:
+            name = name.replace(
+                f"model.layers.{spec_layer}.",
+                f"model.layers.{spec_layer}.mtp_block.",
+            )
+        elif shared_weight:
+            name = name.replace(f"model.layers.{spec_layer}.", "model.")
+        return name

vllm/models/kimi_k3/amd/ops/third_party/__init__.py

@@ -0,0 +1,2 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

vllm/models/kimi_k3/amd/ops/third_party/kda/__init__.py

@@ -0,0 +1,47 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+#
+# AMD/ROCm vendored copy of the Kimi-K3 KDA triton kernels.
+#
+# Provenance: mirror of vllm/models/kimi_k3/nvidia/ops/third_party/kda (tracker
+# mke-tracker @ 7adebfcf9; FLA vendored per PRs #39/#86). Split per-vendor so
+# AMD can carry gfx950-specific kernel changes without touching the NVIDIA copy.
+#
+# fla-org/flash-linear-attention#869 (the unmerged ROCm fixes our earlier
+# amd_fla shim carried against FLA 0.5.0) is covered here by the *newer* vendored
+# FLA rather than the literal patch:
+#   - transpose-state-layout workaround: N/A (kernels rewritten; no
+#     transpose_state_layout path remains),
+#   - AMD autotune configs: present (is_amd num_warps/num_stages branches),
+#   - OOB-mask correctness fix: present (all tl.load use mask=..., other=0).
+# Validated on gfx950: no core-dump, gsm8k 94.1%.
+#
+# AMD-specific deltas vs the NVIDIA copy: NONE yet (byte-identical). Keep in sync
+# with the NVIDIA copy on FLA updates; any divergence should be an intentional,
+# documented gfx950-specific change (a #869-style AMD-only fix).
+
+from .chunk import (
+    chunk_kda,
+    chunk_kda_fwd,
+    chunk_kda_with_fused_gate,
+    chunk_kda_with_fused_gate_fwd,
+    fused_kda_gate,
+    fused_kda_gate_chunk_cumsum,
+)
+from .fused_recurrent import (
+    fused_recurrent_kda,
+    fused_recurrent_kda_fwd,
+    fused_recurrent_kda_packed_decode,
+)
+
+__all__ = [
+    "chunk_kda",
+    "chunk_kda_fwd",
+    "chunk_kda_with_fused_gate",
+    "chunk_kda_with_fused_gate_fwd",
+    "fused_kda_gate",
+    "fused_kda_gate_chunk_cumsum",
+    "fused_recurrent_kda",
+    "fused_recurrent_kda_fwd",
+    "fused_recurrent_kda_packed_decode",
+]

vllm/models/kimi_k3/amd/ops/third_party/kda/chunk.py

@@ -0,0 +1,935 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Songlin Yang, Yu Zhang, Zhiyuan Li
+#
+# This file contains code copied from the flash-linear-attention project.
+# The original source code was licensed under the MIT license and included
+# the following copyright notice:
+# Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li
+# ruff: noqa: E501
+
+
+import torch
+
+from vllm.third_party.flash_linear_attention.ops.chunk_delta_h import (
+    chunk_gated_delta_rule_fwd_h,
+)
+from vllm.third_party.flash_linear_attention.ops.cumsum import chunk_local_cumsum
+from vllm.third_party.flash_linear_attention.ops.index import prepare_chunk_indices
+from vllm.third_party.flash_linear_attention.ops.l2norm import l2norm_fwd
+from vllm.third_party.flash_linear_attention.ops.op import exp2, log
+from vllm.third_party.flash_linear_attention.ops.utils import FLA_CHUNK_SIZE, is_amd
+from vllm.triton_utils import tl, triton
+from vllm.utils.math_utils import RCP_LN2, cdiv, next_power_of_2
+
+from .chunk_intra import chunk_kda_fwd_intra
+
+BT_LIST_AUTOTUNE = [32, 64, 128]
+NUM_WARPS_AUTOTUNE = [2, 4, 8, 16] if is_amd else [4, 8, 16, 32]
+
+
+@triton.heuristics(
+    {
+        "STORE_QG": lambda args: args["qg"] is not None,
+        "STORE_KG": lambda args: args["kg"] is not None,
+        "IS_VARLEN": lambda args: args["cu_seqlens"] is not None,
+    }
+)
+@triton.autotune(
+    configs=[
+        triton.Config({}, num_warps=num_warps, num_stages=num_stages)
+        for num_warps in [2, 4, 8]
+        for num_stages in [2, 3, 4]
+    ],
+    key=["H", "K", "V", "BT", "BK", "BV", "IS_VARLEN"],
+)
+@triton.jit(do_not_specialize=["T"])
+def recompute_w_u_fwd_kernel(
+    q,
+    k,
+    qg,
+    kg,
+    v,
+    beta,
+    w,
+    u,
+    A,
+    gk,
+    cu_seqlens,
+    chunk_indices,
+    T,
+    H: tl.constexpr,
+    K: tl.constexpr,
+    V: tl.constexpr,
+    BT: tl.constexpr,
+    BK: tl.constexpr,
+    BV: tl.constexpr,
+    STORE_QG: tl.constexpr,
+    STORE_KG: tl.constexpr,
+    IS_VARLEN: tl.constexpr,
+    DOT_PRECISION: tl.constexpr,
+):
+    i_t, i_bh = tl.program_id(0), tl.program_id(1)
+    i_b, i_h = i_bh // H, i_bh % H
+    if IS_VARLEN:
+        i_n, i_t = (
+            tl.load(chunk_indices + i_t * 2).to(tl.int32),
+            tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32),
+        )
+        bos, eos = (
+            tl.load(cu_seqlens + i_n).to(tl.int32),
+            tl.load(cu_seqlens + i_n + 1).to(tl.int32),
+        )
+        T = eos - bos
+    else:
+        bos, eos = i_b * T, i_b * T + T
+    p_b = tl.make_block_ptr(beta + bos * H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,))
+    b_b = tl.load(p_b, boundary_check=(0,)).to(tl.float32)
+
+    p_A = tl.make_block_ptr(
+        A + (bos * H + i_h) * BT, (T, BT), (H * BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)
+    )
+    b_A = tl.load(p_A, boundary_check=(0, 1))
+
+    for i_v in range(tl.cdiv(V, BV)):
+        p_v = tl.make_block_ptr(
+            v + (bos * H + i_h) * V,
+            (T, V),
+            (H * V, 1),
+            (i_t * BT, i_v * BV),
+            (BT, BV),
+            (1, 0),
+        )
+        p_u = tl.make_block_ptr(
+            u + (bos * H + i_h) * V,
+            (T, V),
+            (H * V, 1),
+            (i_t * BT, i_v * BV),
+            (BT, BV),
+            (1, 0),
+        )
+        b_v = tl.load(p_v, boundary_check=(0, 1))
+        b_vb = (b_v * b_b[:, None]).to(b_v.dtype)
+        b_u = tl.dot(b_A, b_vb, input_precision=DOT_PRECISION)
+        tl.store(p_u, b_u.to(p_u.dtype.element_ty), boundary_check=(0, 1))
+
+    for i_k in range(tl.cdiv(K, BK)):
+        p_w = tl.make_block_ptr(
+            w + (bos * H + i_h) * K,
+            (T, K),
+            (H * K, 1),
+            (i_t * BT, i_k * BK),
+            (BT, BK),
+            (1, 0),
+        )
+        p_k = tl.make_block_ptr(
+            k + (bos * H + i_h) * K,
+            (T, K),
+            (H * K, 1),
+            (i_t * BT, i_k * BK),
+            (BT, BK),
+            (1, 0),
+        )
+        b_k = tl.load(p_k, boundary_check=(0, 1))
+        b_kb = b_k * b_b[:, None]
+
+        p_gk = tl.make_block_ptr(
+            gk + (bos * H + i_h) * K,
+            (T, K),
+            (H * K, 1),
+            (i_t * BT, i_k * BK),
+            (BT, BK),
+            (1, 0),
+        )
+        b_gk = tl.load(p_gk, boundary_check=(0, 1))
+        b_kb *= exp2(b_gk)
+        if STORE_QG:
+            p_q = tl.make_block_ptr(
+                q + (bos * H + i_h) * K,
+                (T, K),
+                (H * K, 1),
+                (i_t * BT, i_k * BK),
+                (BT, BK),
+                (1, 0),
+            )
+            p_qg = tl.make_block_ptr(
+                qg + (bos * H + i_h) * K,
+                (T, K),
+                (H * K, 1),
+                (i_t * BT, i_k * BK),
+                (BT, BK),
+                (1, 0),
+            )
+            b_q = tl.load(p_q, boundary_check=(0, 1))
+            b_qg = b_q * exp2(b_gk)
+            tl.store(p_qg, b_qg.to(p_qg.dtype.element_ty), boundary_check=(0, 1))
+        if STORE_KG:
+            last_idx = min(i_t * BT + BT, T) - 1
+
+            o_k = i_k * BK + tl.arange(0, BK)
+            m_k = o_k < K
+            b_gn = tl.load(
+                gk + ((bos + last_idx) * H + i_h) * K + o_k, mask=m_k, other=0.0
+            )
+            b_kg = b_k * exp2(b_gn - b_gk)
+
+            p_kg = tl.make_block_ptr(
+                kg + (bos * H + i_h) * K,
+                (T, K),
+                (H * K, 1),
+                (i_t * BT, i_k * BK),
+                (BT, BK),
+                (1, 0),
+            )
+            tl.store(p_kg, b_kg.to(p_kg.dtype.element_ty), boundary_check=(0, 1))
+
+        b_w = tl.dot(b_A, b_kb.to(b_k.dtype))
+        tl.store(p_w, b_w.to(p_w.dtype.element_ty), boundary_check=(0, 1))
+
+
+def recompute_w_u_fwd(
+    k: torch.Tensor,
+    v: torch.Tensor,
+    beta: torch.Tensor,
+    A: torch.Tensor,
+    q: torch.Tensor | None = None,
+    gk: torch.Tensor | None = None,
+    cu_seqlens: torch.Tensor | None = None,
+    chunk_indices: torch.Tensor | None = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    B, T, H, K, V = *k.shape, v.shape[-1]
+    BT = A.shape[-1]
+    BK = 64
+    BV = 64
+
+    if chunk_indices is None and cu_seqlens is not None:
+        chunk_indices = prepare_chunk_indices(cu_seqlens, BT)
+    NT = cdiv(T, BT) if cu_seqlens is None else len(chunk_indices)
+
+    w = torch.empty_like(k)
+    u = torch.empty_like(v)
+    kg = torch.empty_like(k) if gk is not None else None
+    recompute_w_u_fwd_kernel[(NT, B * H)](
+        q=q,
+        k=k,
+        qg=None,
+        kg=kg,
+        v=v,
+        beta=beta,
+        w=w,
+        u=u,
+        A=A,
+        gk=gk,
+        cu_seqlens=cu_seqlens,
+        chunk_indices=chunk_indices,
+        T=T,
+        H=H,
+        K=K,
+        V=V,
+        BT=BT,
+        BK=BK,
+        BV=BV,
+        DOT_PRECISION="ieee",
+    )
+    return w, u, None, kg
+
+
+@triton.heuristics({"IS_VARLEN": lambda args: args["cu_seqlens"] is not None})
+@triton.autotune(
+    configs=[
+        triton.Config({"BK": BK, "BV": BV}, num_warps=num_warps, num_stages=num_stages)
+        for BK in [32, 64]
+        for BV in [64, 128]
+        for num_warps in [2, 4, 8]
+        for num_stages in [2, 3, 4]
+    ],
+    key=["BT"],
+)
+@triton.jit(do_not_specialize=["T"])
+def chunk_gla_fwd_kernel_o(
+    q,
+    v,
+    g,
+    h,
+    o,
+    A,
+    cu_seqlens,
+    chunk_indices,
+    scale,
+    T,
+    H: tl.constexpr,
+    K: tl.constexpr,
+    V: tl.constexpr,
+    BT: tl.constexpr,
+    BK: tl.constexpr,
+    BV: tl.constexpr,
+    IS_VARLEN: tl.constexpr,
+):
+    i_v, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2)
+    i_b, i_h = i_bh // H, i_bh % H
+    if IS_VARLEN:
+        i_tg = i_t
+        i_n, i_t = (
+            tl.load(chunk_indices + i_t * 2).to(tl.int32),
+            tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32),
+        )
+        bos, eos = (
+            tl.load(cu_seqlens + i_n).to(tl.int32),
+            tl.load(cu_seqlens + i_n + 1).to(tl.int32),
+        )
+        T = eos - bos
+        NT = tl.cdiv(T, BT)
+    else:
+        NT = tl.cdiv(T, BT)
+        i_tg = i_b * NT + i_t
+        bos, eos = i_b * T, i_b * T + T
+
+    m_s = tl.arange(0, BT)[:, None] >= tl.arange(0, BT)[None, :]
+
+    b_o = tl.zeros([BT, BV], dtype=tl.float32)
+    for i_k in range(tl.cdiv(K, BK)):
+        p_q = tl.make_block_ptr(
+            q + (bos * H + i_h) * K,
+            (T, K),
+            (H * K, 1),
+            (i_t * BT, i_k * BK),
+            (BT, BK),
+            (1, 0),
+        )
+        p_g = tl.make_block_ptr(
+            g + (bos * H + i_h) * K,
+            (T, K),
+            (H * K, 1),
+            (i_t * BT, i_k * BK),
+            (BT, BK),
+            (1, 0),
+        )
+        p_h = tl.make_block_ptr(
+            h + (i_tg * H + i_h) * K * V,
+            (V, K),
+            (K, 1),
+            (i_v * BV, i_k * BK),
+            (BV, BK),
+            (1, 0),
+        )
+
+        # [BT, BK]
+        b_q = tl.load(p_q, boundary_check=(0, 1))
+        b_q = (b_q * scale).to(b_q.dtype)
+        # [BT, BK]
+        b_g = tl.load(p_g, boundary_check=(0, 1))
+        # [BT, BK]
+        b_qg = (b_q * exp2(b_g)).to(b_q.dtype)
+        # [BV, BK]
+        b_h = tl.load(p_h, boundary_check=(0, 1))
+        # [BT, BV]
+        if i_k >= 0:
+            b_o += tl.dot(b_qg, tl.trans(b_h).to(b_qg.dtype))
+    p_v = tl.make_block_ptr(
+        v + (bos * H + i_h) * V,
+        (T, V),
+        (H * V, 1),
+        (i_t * BT, i_v * BV),
+        (BT, BV),
+        (1, 0),
+    )
+    p_o = tl.make_block_ptr(
+        o + (bos * H + i_h) * V,
+        (T, V),
+        (H * V, 1),
+        (i_t * BT, i_v * BV),
+        (BT, BV),
+        (1, 0),
+    )
+    p_A = tl.make_block_ptr(
+        A + (bos * H + i_h) * BT, (T, BT), (H * BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)
+    )
+    # [BT, BV]
+    b_v = tl.load(p_v, boundary_check=(0, 1))
+    # [BT, BT]
+    b_A = tl.load(p_A, boundary_check=(0, 1))
+    b_A = tl.where(m_s, b_A, 0.0).to(b_v.dtype)
+    b_o += tl.dot(b_A, b_v, allow_tf32=False)
+    tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1))
+
+
+def chunk_gla_fwd_o_gk(
+    q: torch.Tensor,
+    v: torch.Tensor,
+    g: torch.Tensor,
+    A: torch.Tensor,
+    h: torch.Tensor,
+    o: torch.Tensor,
+    scale: float,
+    cu_seqlens: torch.Tensor | None = None,
+    chunk_indices: torch.Tensor | None = None,
+    chunk_size: int = FLA_CHUNK_SIZE,
+):
+    B, T, H, K, V = *q.shape, v.shape[-1]
+    BT = chunk_size
+
+    if chunk_indices is None and cu_seqlens is not None:
+        chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size)
+    NT = cdiv(T, BT) if cu_seqlens is None else len(chunk_indices)
+
+    def grid(meta):
+        return (cdiv(V, meta["BV"]), NT, B * H)
+
+    chunk_gla_fwd_kernel_o[grid](
+        q=q,
+        v=v,
+        g=g,
+        h=h,
+        o=o,
+        A=A,
+        cu_seqlens=cu_seqlens,
+        chunk_indices=chunk_indices,
+        scale=scale,
+        T=T,
+        H=H,
+        K=K,
+        V=V,
+        BT=BT,
+    )
+    return o
+
+
+@triton.heuristics(
+    {
+        "HAS_BIAS": lambda args: args["g_bias"] is not None,
+        "IS_VARLEN": lambda args: args["cu_seqlens"] is not None,
+    }
+)
+@triton.autotune(
+    configs=[
+        triton.Config({"BS": BS}, num_warps=num_warps)
+        for BS in [32, 64]
+        for num_warps in [2, 4, 8]
+    ],
+    key=["H", "S", "BT", "IS_VARLEN"],
+)
+@triton.jit(do_not_specialize=["T"])
+def kda_gate_chunk_cumsum_vector_kernel(
+    s,
+    raw_beta,
+    A_log,
+    g_bias,
+    o,
+    beta_out,
+    cu_seqlens,
+    chunk_indices,
+    cumsum_scale,
+    lower_bound,
+    beta,
+    threshold,
+    T,
+    stride_beta_batch,
+    stride_beta_token,
+    stride_beta_head,
+    H: tl.constexpr,
+    S: tl.constexpr,
+    BT: tl.constexpr,
+    BS: tl.constexpr,
+    HAS_BIAS: tl.constexpr,
+    IS_VARLEN: tl.constexpr,
+    USE_LOWER_BOUND: tl.constexpr,
+):
+    i_s, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2)
+    i_b, i_h = i_bh // H, i_bh % H
+    if IS_VARLEN:
+        i_n, i_t = (
+            tl.load(chunk_indices + i_t * 2).to(tl.int32),
+            tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32),
+        )
+        bos, eos = (
+            tl.load(cu_seqlens + i_n).to(tl.int32),
+            tl.load(cu_seqlens + i_n + 1).to(tl.int32),
+        )
+        T = eos - bos
+    else:
+        bos = i_b * T
+
+    if i_s == 0:
+        o_beta_t = tl.arange(0, BT)
+        m_beta = i_t * BT + o_beta_t < T
+        if IS_VARLEN:
+            p_beta = (
+                raw_beta
+                + (bos + i_t * BT + o_beta_t) * stride_beta_token
+                + i_h * stride_beta_head
+            )
+        else:
+            p_beta = (
+                raw_beta
+                + i_b * stride_beta_batch
+                + (i_t * BT + o_beta_t) * stride_beta_token
+                + i_h * stride_beta_head
+            )
+        b_beta = tl.load(p_beta, mask=m_beta, other=0.0).to(tl.float32)
+        p_beta_out = beta_out + (bos + i_t * BT + o_beta_t) * H + i_h
+        tl.store(p_beta_out, tl.sigmoid(b_beta), mask=m_beta)
+        return
+
+    i_s -= 1
+
+    p_s = tl.make_block_ptr(
+        s + (bos * H + i_h) * S,
+        (T, S),
+        (H * S, 1),
+        (i_t * BT, i_s * BS),
+        (BT, BS),
+        (1, 0),
+    )
+    p_o = tl.make_block_ptr(
+        o + (bos * H + i_h) * S,
+        (T, S),
+        (H * S, 1),
+        (i_t * BT, i_s * BS),
+        (BT, BS),
+        (1, 0),
+    )
+
+    b_s = tl.load(p_s, boundary_check=(0, 1)).to(tl.float32)
+    if HAS_BIAS:
+        p_bias = tl.make_block_ptr(
+            g_bias + i_h * S,
+            (S,),
+            (1,),
+            (i_s * BS,),
+            (BS,),
+            (0,),
+        )
+        b_bias = tl.load(p_bias, boundary_check=(0,)).to(tl.float32)
+        b_s += b_bias[None, :]
+
+    b_a = tl.exp(tl.load(A_log + i_h).to(tl.float32))
+    if USE_LOWER_BOUND:
+        b_gate = lower_bound * tl.sigmoid(b_a * b_s)
+    else:
+        b_g_scaled = b_s * beta
+        b_softplus = tl.where(
+            b_g_scaled > threshold,
+            b_s,
+            (1.0 / beta) * log(1.0 + tl.exp(b_g_scaled)),
+        )
+        b_gate = -b_a * b_softplus
+
+    # Boundary loads return zero, but bias and gate activation can make padded
+    # rows nonzero. Padding trails valid rows, so it only affects masked stores.
+    b_o = tl.cumsum(b_gate, axis=0) * cumsum_scale
+    tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1))
+
+
+def fused_kda_gate_chunk_cumsum(
+    raw_g: torch.Tensor,
+    raw_beta: torch.Tensor,
+    A_log: torch.Tensor,
+    g_bias: torch.Tensor | None = None,
+    beta: float = 1.0,
+    threshold: float = 20.0,
+    lower_bound: float | None = None,
+    cu_seqlens: torch.Tensor | None = None,
+    chunk_indices: torch.Tensor | None = None,
+    chunk_size: int = FLA_CHUNK_SIZE,
+    output_dtype: torch.dtype | None = torch.float,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    if cu_seqlens is not None:
+        assert raw_g.shape[0] == 1, (
+            "Only batch size 1 is supported when cu_seqlens are provided"
+        )
+    B, T, H, D = raw_g.shape
+    if raw_beta.shape != (B, T, H):
+        raise ValueError(f"Expected raw_beta shape {(B, T, H)}, got {raw_beta.shape}")
+    if chunk_indices is None and cu_seqlens is not None:
+        chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size)
+    NT = cdiv(T, chunk_size) if cu_seqlens is None else len(chunk_indices)
+
+    A_log = A_log.reshape(-1)
+    if g_bias is not None:
+        g_bias = g_bias.reshape(-1)
+    y = torch.empty_like(raw_g, dtype=output_dtype or raw_g.dtype)
+    beta_out = torch.empty(raw_beta.shape, device=raw_beta.device, dtype=torch.float32)
+
+    def grid(meta):
+        # For each (chunk, head), program 0 computes beta without extending a
+        # gate tile's critical path. The remaining programs cover the gate dim.
+        return (cdiv(meta["S"], meta["BS"]) + 1, NT, B * H)
+
+    kda_gate_chunk_cumsum_vector_kernel[grid](
+        s=raw_g,
+        raw_beta=raw_beta,
+        A_log=A_log,
+        g_bias=g_bias,
+        o=y,
+        beta_out=beta_out,
+        cu_seqlens=cu_seqlens,
+        chunk_indices=chunk_indices,
+        # RCP_LN2 folds in the natural-log -> log2 conversion so downstream
+        # exp2-based kernels reproduce exp(g). Keep this in sync with the
+        # `use_exp2=True` path in `_chunk_kda_fwd_with_cumulative_g`.
+        cumsum_scale=RCP_LN2,
+        lower_bound=lower_bound or 0.0,
+        beta=beta,
+        threshold=threshold,
+        T=T,
+        stride_beta_batch=raw_beta.stride(0),
+        stride_beta_token=raw_beta.stride(1),
+        stride_beta_head=raw_beta.stride(2),
+        H=H,
+        S=D,
+        BT=chunk_size,
+        USE_LOWER_BOUND=lower_bound is not None,
+    )
+    return y, beta_out
+
+
+def _chunk_kda_fwd_with_cumulative_g(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    g: torch.Tensor,
+    beta: torch.Tensor,
+    scale: float,
+    initial_state: torch.Tensor,
+    output_final_state: bool,
+    cu_seqlens: torch.Tensor | None = None,
+    chunk_indices: torch.Tensor | None = None,
+    chunk_size: int = FLA_CHUNK_SIZE,
+    safe_gate: bool = False,
+):
+    # `g` must already be chunk-local cumulatively-summed AND scaled by
+    # RCP_LN2 (so the downstream exp2-based kernels reproduce exp(g)).
+    # Use `chunk_kda_fwd` or `chunk_kda_with_fused_gate_fwd` instead of
+    # calling this helper directly unless that invariant is upheld.
+    Aqk, A = chunk_kda_fwd_intra(
+        q=q,
+        k=k,
+        gk=g,
+        beta=beta,
+        scale=scale,
+        cu_seqlens=cu_seqlens,
+        chunk_indices=chunk_indices,
+        chunk_size=chunk_size,
+        safe_gate=safe_gate,
+    )
+    w, u, _, kg = recompute_w_u_fwd(
+        k=k,
+        v=v,
+        beta=beta,
+        A=A,
+        gk=g,
+        cu_seqlens=cu_seqlens,
+        chunk_indices=chunk_indices,
+    )
+    del A
+    h, v_new, final_state = chunk_gated_delta_rule_fwd_h(
+        k=kg,
+        w=w,
+        u=u,
+        gk=g,
+        initial_state=initial_state,
+        output_final_state=output_final_state,
+        cu_seqlens=cu_seqlens,
+        chunk_indices=chunk_indices,
+        use_exp2=True,
+    )
+    del w, u, kg
+    o = chunk_gla_fwd_o_gk(
+        q=q,
+        v=v_new,
+        g=g,
+        A=Aqk,
+        h=h,
+        o=v,
+        scale=scale,
+        cu_seqlens=cu_seqlens,
+        chunk_indices=chunk_indices,
+        chunk_size=chunk_size,
+    )
+    del Aqk, v_new, h
+    return o, final_state
+
+
+def chunk_kda_fwd(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    g: torch.Tensor,
+    beta: torch.Tensor,
+    scale: float,
+    initial_state: torch.Tensor,
+    output_final_state: bool,
+    cu_seqlens: torch.Tensor | None = None,
+):
+    chunk_size = FLA_CHUNK_SIZE
+    chunk_indices = (
+        prepare_chunk_indices(cu_seqlens, chunk_size)
+        if cu_seqlens is not None
+        else None
+    )
+    g = chunk_local_cumsum(
+        g,
+        chunk_size=chunk_size,
+        cu_seqlens=cu_seqlens,
+        chunk_indices=chunk_indices,
+    )
+    # KDA evaluates cumulative gate decays with exp2. Convert from natural-log
+    # space so exp(x) is preserved as exp2(x / ln(2)).
+    g = g * RCP_LN2
+    return _chunk_kda_fwd_with_cumulative_g(
+        q=q,
+        k=k,
+        v=v,
+        g=g,
+        beta=beta,
+        scale=scale,
+        initial_state=initial_state,
+        output_final_state=output_final_state,
+        cu_seqlens=cu_seqlens,
+        chunk_indices=chunk_indices,
+        chunk_size=chunk_size,
+    )
+
+
+def chunk_kda_with_fused_gate_fwd(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    raw_g: torch.Tensor,
+    raw_beta: torch.Tensor,
+    A_log: torch.Tensor,
+    g_bias: torch.Tensor | None,
+    scale: float,
+    initial_state: torch.Tensor,
+    output_final_state: bool,
+    lower_bound: float | None = None,
+    cu_seqlens: torch.Tensor | None = None,
+):
+    chunk_size = FLA_CHUNK_SIZE
+    chunk_indices = (
+        prepare_chunk_indices(cu_seqlens, chunk_size)
+        if cu_seqlens is not None
+        else None
+    )
+    g, beta = fused_kda_gate_chunk_cumsum(
+        raw_g,
+        raw_beta=raw_beta,
+        A_log=A_log,
+        g_bias=g_bias,
+        cu_seqlens=cu_seqlens,
+        chunk_indices=chunk_indices,
+        chunk_size=chunk_size,
+        lower_bound=lower_bound,
+    )
+    return _chunk_kda_fwd_with_cumulative_g(
+        q=q,
+        k=k,
+        v=v,
+        g=g,
+        beta=beta,
+        scale=scale,
+        initial_state=initial_state,
+        output_final_state=output_final_state,
+        cu_seqlens=cu_seqlens,
+        chunk_indices=chunk_indices,
+        chunk_size=chunk_size,
+        safe_gate=lower_bound is not None,
+    )
+
+
+def chunk_kda(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    g: torch.Tensor,
+    beta: torch.Tensor,
+    scale: float = None,
+    initial_state: torch.Tensor = None,
+    output_final_state: bool = False,
+    use_qk_l2norm_in_kernel: bool = False,
+    cu_seqlens: torch.Tensor | None = None,
+    **kwargs,
+):
+    if scale is None:
+        scale = k.shape[-1] ** -0.5
+
+    if use_qk_l2norm_in_kernel:
+        q = l2norm_fwd(q.contiguous())
+        k = l2norm_fwd(k.contiguous())
+
+    o, final_state = chunk_kda_fwd(
+        q=q,
+        k=k,
+        v=v.contiguous(),
+        g=g.contiguous(),
+        beta=beta.contiguous(),
+        scale=scale,
+        initial_state=initial_state.contiguous(),
+        output_final_state=output_final_state,
+        cu_seqlens=cu_seqlens,
+    )
+    return o, final_state
+
+
+def chunk_kda_with_fused_gate(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    raw_g: torch.Tensor,
+    raw_beta: torch.Tensor,
+    A_log: torch.Tensor,
+    g_bias: torch.Tensor | None,
+    scale: float | None = None,
+    initial_state: torch.Tensor | None = None,
+    output_final_state: bool = False,
+    lower_bound: float | None = None,
+    use_qk_l2norm_in_kernel: bool = False,
+    cu_seqlens: torch.Tensor | None = None,
+    **kwargs,
+):
+    """Run chunk KDA from raw gate and beta projections."""
+    if scale is None:
+        scale = k.shape[-1] ** -0.5
+
+    if use_qk_l2norm_in_kernel:
+        q = l2norm_fwd(q.contiguous())
+        k = l2norm_fwd(k.contiguous())
+
+    o, final_state = chunk_kda_with_fused_gate_fwd(
+        q=q,
+        k=k,
+        v=v.contiguous(),
+        raw_g=raw_g.contiguous(),
+        raw_beta=raw_beta,
+        A_log=A_log,
+        g_bias=g_bias,
+        scale=scale,
+        initial_state=initial_state.contiguous() if initial_state is not None else None,
+        output_final_state=output_final_state,
+        lower_bound=lower_bound,
+        cu_seqlens=cu_seqlens,
+    )
+    return o, final_state
+
+
+@triton.autotune(
+    configs=[
+        triton.Config({"BT": bt}, num_warps=nw, num_stages=ns)
+        for bt in BT_LIST_AUTOTUNE
+        for nw in NUM_WARPS_AUTOTUNE
+        for ns in [2, 3]
+    ],
+    key=["H", "D"],
+)
+@triton.jit
+def kda_gate_fwd_kernel(
+    g,
+    A,
+    y,
+    g_bias,
+    lower_bound,
+    beta: tl.constexpr,
+    threshold: tl.constexpr,
+    T,
+    H,
+    D: tl.constexpr,
+    BT: tl.constexpr,
+    BD: tl.constexpr,
+    HAS_BIAS: tl.constexpr,
+    USE_LOWER_BOUND: tl.constexpr,
+):
+    i_t, i_h = tl.program_id(0), tl.program_id(1)
+    n_t = i_t * BT
+
+    b_a = tl.exp(tl.load(A + i_h).to(tl.float32))
+
+    stride_row = H * D
+    stride_col = 1
+
+    g_ptr = tl.make_block_ptr(
+        base=g + i_h * D,
+        shape=(T, D),
+        strides=(stride_row, stride_col),
+        offsets=(n_t, 0),
+        block_shape=(BT, BD),
+        order=(1, 0),
+    )
+
+    y_ptr = tl.make_block_ptr(
+        base=y + i_h * D,
+        shape=(T, D),
+        strides=(stride_row, stride_col),
+        offsets=(n_t, 0),
+        block_shape=(BT, BD),
+        order=(1, 0),
+    )
+
+    b_g = tl.load(g_ptr, boundary_check=(0, 1)).to(tl.float32)
+
+    if HAS_BIAS:
+        n_d = tl.arange(0, BD)
+        bias_mask = n_d < D
+        b_bias = tl.load(g_bias + i_h * D + n_d, mask=bias_mask, other=0.0).to(
+            tl.float32
+        )
+        b_g = b_g + b_bias[None, :]
+
+    if USE_LOWER_BOUND:
+        b_y = lower_bound * tl.sigmoid(b_a * b_g)
+    else:
+        g_scaled = b_g * beta
+        use_linear = g_scaled > threshold
+        sp = tl.where(use_linear, b_g, (1.0 / beta) * log(1.0 + tl.exp(g_scaled)))
+        b_y = -b_a * sp
+
+    tl.store(y_ptr, b_y.to(y.dtype.element_ty), boundary_check=(0, 1))
+
+
+def fused_kda_gate(
+    g: torch.Tensor,
+    A: torch.Tensor,
+    head_k_dim: int,
+    g_bias: torch.Tensor | None = None,
+    beta: float = 1.0,
+    threshold: float = 20.0,
+    lower_bound: float | None = None,
+) -> torch.Tensor:
+    """
+    Forward pass for KDA gate:
+      input g: [..., H*D]
+      param A: [H] or [1, 1, H, 1]
+      beta: softplus beta parameter
+      threshold: softplus threshold parameter
+      return  : [..., H, D]
+    """
+    orig_shape = g.shape[:-1]
+
+    g = g.view(-1, g.shape[-1])
+    T = g.shape[0]
+    HD = g.shape[1]
+    H = A.numel()
+    assert H * head_k_dim == HD
+
+    y = torch.empty_like(g, dtype=torch.float32)
+
+    def grid(meta):
+        return (cdiv(T, meta["BT"]), H)
+
+    kda_gate_fwd_kernel[grid](
+        g,
+        A,
+        y,
+        g_bias,
+        lower_bound or 0.0,
+        beta,
+        threshold,
+        T,
+        H,
+        head_k_dim,
+        BD=next_power_of_2(head_k_dim),
+        HAS_BIAS=g_bias is not None,
+        USE_LOWER_BOUND=lower_bound is not None,
+    )
+
+    y = y.view(*orig_shape, H, head_k_dim)
+    return y

vllm/models/kimi_k3/amd/ops/third_party/kda/chunk_intra.py

@@ -0,0 +1,662 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Songlin Yang, Yu Zhang, Zhiyuan Li
+#
+# This file contains code copied from the flash-linear-attention project.
+# The original source was licensed under the MIT license.
+# Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li
+# Forward-only adaptation of flash-linear-attention 0.5.0.
+# ruff: noqa: E501
+
+import torch
+
+from vllm.platforms import current_platform
+from vllm.third_party.flash_linear_attention.ops.index import prepare_chunk_indices
+from vllm.third_party.flash_linear_attention.ops.op import exp2, gather
+from vllm.third_party.flash_linear_attention.ops.utils import is_gather_supported
+from vllm.triton_utils import tl, triton
+
+from .chunk_intra_token_parallel import chunk_kda_fwd_intra_token_parallel
+
+################################################################################
+# Fused inter + solve_tril kernel: compute off-diagonal Akk and solve in one pass
+################################################################################
+
+
+@triton.heuristics(
+    {
+        "IS_VARLEN": lambda args: args["cu_seqlens"] is not None,
+    }
+)
+@triton.autotune(
+    configs=[
+        triton.Config({"BK": BK}, num_warps=num_warps)
+        for BK in [32, 64]
+        for num_warps in [1, 2, 4]
+    ],
+    key=["H", "HV", "K", "BC"],
+)
+@triton.jit(do_not_specialize=["T"])
+def chunk_kda_fwd_kernel_inter_solve_fused(
+    q,
+    k,
+    g,
+    beta,
+    Aqk,
+    Akkd,
+    Akk,
+    scale,
+    cu_seqlens,
+    chunk_indices,
+    T,
+    H: tl.constexpr,
+    HV: tl.constexpr,
+    K: tl.constexpr,
+    BT: tl.constexpr,
+    BC: tl.constexpr,
+    BK: tl.constexpr,
+    IS_VARLEN: tl.constexpr,
+    USE_SAFE_GATE: tl.constexpr,
+    SOLVE_TRIL_DOT_PRECISION: tl.constexpr,
+):
+    """
+    Fused kernel: compute inter-subchunk Akk + solve_tril in one pass.
+    Prerequisite: token_parallel has already computed diagonal Akk blocks in Akkd.
+
+    This kernel:
+    1. Computes off-diagonal Aqk blocks -> writes to global
+    2. Computes off-diagonal Akk blocks -> keeps in registers
+    3. Loads diagonal Akk blocks from Akkd (fp32)
+    4. Does forward substitution on diagonals
+    5. Computes merged Akk_inv
+    6. Writes Akk_inv to Akk
+    """
+    i_t, i_bh = tl.program_id(0), tl.program_id(1)
+    i_b, i_hv = i_bh // HV, i_bh % HV
+    i_h = i_hv // (HV // H)
+
+    if IS_VARLEN:
+        i_n, i_t = (
+            tl.load(chunk_indices + i_t * 2).to(tl.int32),
+            tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32),
+        )
+        bos, eos = (
+            tl.load(cu_seqlens + i_n).to(tl.int32),
+            tl.load(cu_seqlens + i_n + 1).to(tl.int32),
+        )
+        T = eos - bos
+    else:
+        bos, eos = i_b * T, i_b * T + T
+
+    if i_t * BT >= T:
+        return
+
+    i_tc0 = i_t * BT
+    i_tc1 = i_t * BT + BC
+    i_tc2 = i_t * BT + 2 * BC
+    i_tc3 = i_t * BT + 3 * BC
+
+    q += (bos * H + i_h) * K
+    k += (bos * H + i_h) * K
+    g += (bos * HV + i_hv) * K
+    Aqk += (bos * HV + i_hv) * BT
+    Akk += (bos * HV + i_hv) * BT
+    Akkd += (bos * HV + i_hv) * BC
+
+    o_i = tl.arange(0, BC)
+    m_tc1 = (i_tc1 + o_i) < T
+    m_tc2 = (i_tc2 + o_i) < T
+    m_tc3 = (i_tc3 + o_i) < T
+
+    b_Aqk10 = tl.zeros([BC, BC], dtype=tl.float32)
+    b_Akk10 = tl.zeros([BC, BC], dtype=tl.float32)
+
+    b_Aqk20 = tl.zeros([BC, BC], dtype=tl.float32)
+    b_Akk20 = tl.zeros([BC, BC], dtype=tl.float32)
+    b_Aqk21 = tl.zeros([BC, BC], dtype=tl.float32)
+    b_Akk21 = tl.zeros([BC, BC], dtype=tl.float32)
+
+    b_Aqk30 = tl.zeros([BC, BC], dtype=tl.float32)
+    b_Akk30 = tl.zeros([BC, BC], dtype=tl.float32)
+    b_Aqk31 = tl.zeros([BC, BC], dtype=tl.float32)
+    b_Akk31 = tl.zeros([BC, BC], dtype=tl.float32)
+    b_Aqk32 = tl.zeros([BC, BC], dtype=tl.float32)
+    b_Akk32 = tl.zeros([BC, BC], dtype=tl.float32)
+
+    ################################################################################
+    # off-diagonal blocks
+    ################################################################################
+    for i_k in range(tl.cdiv(K, BK)):
+        o_k = i_k * BK + tl.arange(0, BK)
+        m_k = o_k < K
+
+        p_k0 = tl.make_block_ptr(
+            k, (T, K), (H * K, 1), (i_tc0, i_k * BK), (BC, BK), (1, 0)
+        )
+        p_g0 = tl.make_block_ptr(
+            g, (T, K), (HV * K, 1), (i_tc0, i_k * BK), (BC, BK), (1, 0)
+        )
+        b_k0 = tl.load(p_k0, boundary_check=(0, 1)).to(tl.float32)
+        b_g0 = tl.load(p_g0, boundary_check=(0, 1)).to(tl.float32)
+
+        if i_tc1 < T:
+            p_q1 = tl.make_block_ptr(
+                q, (T, K), (H * K, 1), (i_tc1, i_k * BK), (BC, BK), (1, 0)
+            )
+            p_k1 = tl.make_block_ptr(
+                k, (T, K), (H * K, 1), (i_tc1, i_k * BK), (BC, BK), (1, 0)
+            )
+            p_g1 = tl.make_block_ptr(
+                g, (T, K), (HV * K, 1), (i_tc1, i_k * BK), (BC, BK), (1, 0)
+            )
+            # [BC, BK]
+            b_q1 = tl.load(p_q1, boundary_check=(0, 1)).to(tl.float32)
+            b_k1 = tl.load(p_k1, boundary_check=(0, 1)).to(tl.float32)
+            b_g1 = tl.load(p_g1, boundary_check=(0, 1)).to(tl.float32)
+            # [BK]
+            b_gn1 = tl.load(g + i_tc1 * HV * K + o_k, mask=m_k, other=0).to(tl.float32)
+            # [BC, BK]
+            b_gqn = tl.where(m_tc1[:, None], exp2(b_g1 - b_gn1[None, :]), 0)
+            # [BK, BC]
+            b_kgt = tl.trans(b_k0 * exp2(b_gn1[None, :] - b_g0))
+            # [BC, BC]
+            b_Aqk10 += tl.dot(b_q1 * b_gqn, b_kgt)
+            b_Akk10 += tl.dot(b_k1 * b_gqn, b_kgt)
+
+            if i_tc2 < T:
+                p_q2 = tl.make_block_ptr(
+                    q, (T, K), (H * K, 1), (i_tc2, i_k * BK), (BC, BK), (1, 0)
+                )
+                p_k2 = tl.make_block_ptr(
+                    k, (T, K), (H * K, 1), (i_tc2, i_k * BK), (BC, BK), (1, 0)
+                )
+                p_g2 = tl.make_block_ptr(
+                    g, (T, K), (HV * K, 1), (i_tc2, i_k * BK), (BC, BK), (1, 0)
+                )
+                # [BC, BK]
+                b_q2 = tl.load(p_q2, boundary_check=(0, 1)).to(tl.float32)
+                b_k2 = tl.load(p_k2, boundary_check=(0, 1)).to(tl.float32)
+                b_g2 = tl.load(p_g2, boundary_check=(0, 1)).to(tl.float32)
+                # [BK]
+                b_gn2 = tl.load(g + i_tc2 * HV * K + o_k, mask=m_k, other=0).to(
+                    tl.float32
+                )
+                # [BC, BK]
+                b_gqn2 = tl.where(m_tc2[:, None], exp2(b_g2 - b_gn2[None, :]), 0)
+                b_qg2 = b_q2 * b_gqn2
+                b_kg2 = b_k2 * b_gqn2
+                # [BK, BC]
+                b_kgt = tl.trans(b_k0 * exp2(b_gn2[None, :] - b_g0))
+                b_Aqk20 += tl.dot(b_qg2, b_kgt)
+                b_Akk20 += tl.dot(b_kg2, b_kgt)
+                # [BC, BC]
+                b_kgt = tl.trans(b_k1 * exp2(b_gn2[None, :] - b_g1))
+                # [BC, BC]
+                b_Aqk21 += tl.dot(b_qg2, b_kgt)
+                b_Akk21 += tl.dot(b_kg2, b_kgt)
+
+                if i_tc3 < T:
+                    p_q3 = tl.make_block_ptr(
+                        q, (T, K), (H * K, 1), (i_tc3, i_k * BK), (BC, BK), (1, 0)
+                    )
+                    p_k3 = tl.make_block_ptr(
+                        k, (T, K), (H * K, 1), (i_tc3, i_k * BK), (BC, BK), (1, 0)
+                    )
+                    p_g3 = tl.make_block_ptr(
+                        g, (T, K), (HV * K, 1), (i_tc3, i_k * BK), (BC, BK), (1, 0)
+                    )
+                    # [BC, BK]
+                    b_q3 = tl.load(p_q3, boundary_check=(0, 1)).to(tl.float32)
+                    b_k3 = tl.load(p_k3, boundary_check=(0, 1)).to(tl.float32)
+                    b_g3 = tl.load(p_g3, boundary_check=(0, 1)).to(tl.float32)
+                    # [BK]
+                    b_gn3 = tl.load(g + i_tc3 * HV * K + o_k, mask=m_k, other=0).to(
+                        tl.float32
+                    )
+                    # [BC, BK]
+                    b_gqn3 = tl.where(m_tc3[:, None], exp2(b_g3 - b_gn3[None, :]), 0)
+                    b_qg3 = b_q3 * b_gqn3
+                    b_kg3 = b_k3 * b_gqn3
+                    # [BK, BC]
+                    b_kgt = tl.trans(b_k0 * exp2(b_gn3[None, :] - b_g0))
+                    # [BC, BC]
+                    b_Aqk30 += tl.dot(b_qg3, b_kgt)
+                    b_Akk30 += tl.dot(b_kg3, b_kgt)
+                    # [BK, BC]
+                    b_kgt = tl.trans(b_k1 * exp2(b_gn3[None, :] - b_g1))
+                    # [BC, BC]
+                    b_Aqk31 += tl.dot(b_qg3, b_kgt)
+                    b_Akk31 += tl.dot(b_kg3, b_kgt)
+                    # [BK, BC]
+                    b_kgt = tl.trans(b_k2 * exp2(b_gn3[None, :] - b_g2))
+                    # [BC, BC]
+                    b_Aqk32 += tl.dot(b_qg3, b_kgt)
+                    b_Akk32 += tl.dot(b_kg3, b_kgt)
+
+    ################################################################################
+    # save off-diagonal Aqk blocks and prepare Akk
+    ################################################################################
+    if i_tc1 < T:
+        p_Aqk10 = tl.make_block_ptr(
+            Aqk, (T, BT), (HV * BT, 1), (i_tc1, 0), (BC, BC), (1, 0)
+        )
+        tl.store(
+            p_Aqk10, (b_Aqk10 * scale).to(Aqk.dtype.element_ty), boundary_check=(0, 1)
+        )
+
+        p_b1 = tl.make_block_ptr(
+            beta + bos * HV + i_hv, (T,), (HV,), (i_tc1,), (BC,), (0,)
+        )
+        b_b1 = tl.load(p_b1, boundary_check=(0,)).to(tl.float32)
+        b_Akk10 = b_Akk10 * b_b1[:, None]
+    if i_tc2 < T:
+        p_Aqk20 = tl.make_block_ptr(
+            Aqk, (T, BT), (HV * BT, 1), (i_tc2, 0), (BC, BC), (1, 0)
+        )
+        p_Aqk21 = tl.make_block_ptr(
+            Aqk, (T, BT), (HV * BT, 1), (i_tc2, BC), (BC, BC), (1, 0)
+        )
+        tl.store(
+            p_Aqk20, (b_Aqk20 * scale).to(Aqk.dtype.element_ty), boundary_check=(0, 1)
+        )
+        tl.store(
+            p_Aqk21, (b_Aqk21 * scale).to(Aqk.dtype.element_ty), boundary_check=(0, 1)
+        )
+
+        p_b2 = tl.make_block_ptr(
+            beta + bos * HV + i_hv, (T,), (HV,), (i_tc2,), (BC,), (0,)
+        )
+        b_b2 = tl.load(p_b2, boundary_check=(0,)).to(tl.float32)
+        b_Akk20 = b_Akk20 * b_b2[:, None]
+        b_Akk21 = b_Akk21 * b_b2[:, None]
+    if i_tc3 < T:
+        p_Aqk30 = tl.make_block_ptr(
+            Aqk, (T, BT), (HV * BT, 1), (i_tc3, 0), (BC, BC), (1, 0)
+        )
+        p_Aqk31 = tl.make_block_ptr(
+            Aqk, (T, BT), (HV * BT, 1), (i_tc3, BC), (BC, BC), (1, 0)
+        )
+        p_Aqk32 = tl.make_block_ptr(
+            Aqk, (T, BT), (HV * BT, 1), (i_tc3, 2 * BC), (BC, BC), (1, 0)
+        )
+        tl.store(
+            p_Aqk30, (b_Aqk30 * scale).to(Aqk.dtype.element_ty), boundary_check=(0, 1)
+        )
+        tl.store(
+            p_Aqk31, (b_Aqk31 * scale).to(Aqk.dtype.element_ty), boundary_check=(0, 1)
+        )
+        tl.store(
+            p_Aqk32, (b_Aqk32 * scale).to(Aqk.dtype.element_ty), boundary_check=(0, 1)
+        )
+
+        p_b3 = tl.make_block_ptr(
+            beta + bos * HV + i_hv, (T,), (HV,), (i_tc3,), (BC,), (0,)
+        )
+        b_b3 = tl.load(p_b3, boundary_check=(0,)).to(tl.float32)
+        b_Akk30 = b_Akk30 * b_b3[:, None]
+        b_Akk31 = b_Akk31 * b_b3[:, None]
+        b_Akk32 = b_Akk32 * b_b3[:, None]
+
+    p_Akk00 = tl.make_block_ptr(
+        Akkd, (T, BC), (HV * BC, 1), (i_tc0, 0), (BC, BC), (1, 0)
+    )
+    p_Akk11 = tl.make_block_ptr(
+        Akkd, (T, BC), (HV * BC, 1), (i_tc1, 0), (BC, BC), (1, 0)
+    )
+    p_Akk22 = tl.make_block_ptr(
+        Akkd, (T, BC), (HV * BC, 1), (i_tc2, 0), (BC, BC), (1, 0)
+    )
+    p_Akk33 = tl.make_block_ptr(
+        Akkd, (T, BC), (HV * BC, 1), (i_tc3, 0), (BC, BC), (1, 0)
+    )
+    b_Ai00 = tl.load(p_Akk00, boundary_check=(0, 1)).to(tl.float32)
+    b_Ai11 = tl.load(p_Akk11, boundary_check=(0, 1)).to(tl.float32)
+    b_Ai22 = tl.load(p_Akk22, boundary_check=(0, 1)).to(tl.float32)
+    b_Ai33 = tl.load(p_Akk33, boundary_check=(0, 1)).to(tl.float32)
+
+    ################################################################################
+    # forward substitution on diagonals
+    ################################################################################
+
+    if not USE_SAFE_GATE:
+        m_A = o_i[:, None] > o_i[None, :]
+        m_I = o_i[:, None] == o_i[None, :]
+
+        b_Ai00 = -tl.where(m_A, b_Ai00, 0)
+        b_Ai11 = -tl.where(m_A, b_Ai11, 0)
+        b_Ai22 = -tl.where(m_A, b_Ai22, 0)
+        b_Ai33 = -tl.where(m_A, b_Ai33, 0)
+
+        for i in range(2, min(BC, T - i_tc0)):
+            b_a00 = -tl.load(Akkd + (i_tc0 + i) * HV * BC + o_i)
+            b_a00 = tl.where(o_i < i, b_a00, 0.0)
+            b_a00 += tl.sum(b_a00[:, None] * b_Ai00, 0)
+            b_Ai00 = tl.where((o_i == i)[:, None], b_a00, b_Ai00)
+        for i in range(BC + 2, min(2 * BC, T - i_tc0)):
+            b_a11 = -tl.load(Akkd + (i_tc0 + i) * HV * BC + o_i)
+            b_a11 = tl.where(o_i < i - BC, b_a11, 0.0)
+            b_a11 += tl.sum(b_a11[:, None] * b_Ai11, 0)
+            b_Ai11 = tl.where((o_i == i - BC)[:, None], b_a11, b_Ai11)
+        for i in range(2 * BC + 2, min(3 * BC, T - i_tc0)):
+            b_a22 = -tl.load(Akkd + (i_tc0 + i) * HV * BC + o_i)
+            b_a22 = tl.where(o_i < i - 2 * BC, b_a22, 0.0)
+            b_a22 += tl.sum(b_a22[:, None] * b_Ai22, 0)
+            b_Ai22 = tl.where((o_i == i - 2 * BC)[:, None], b_a22, b_Ai22)
+        for i in range(3 * BC + 2, min(4 * BC, T - i_tc0)):
+            b_a33 = -tl.load(Akkd + (i_tc0 + i) * HV * BC + o_i)
+            b_a33 = tl.where(o_i < i - 3 * BC, b_a33, 0.0)
+            b_a33 += tl.sum(b_a33[:, None] * b_Ai33, 0)
+            b_Ai33 = tl.where((o_i == i - 3 * BC)[:, None], b_a33, b_Ai33)
+
+        b_Ai00 += m_I
+        b_Ai11 += m_I
+        b_Ai22 += m_I
+        b_Ai33 += m_I
+
+    ################################################################################
+    # compute merged inverse using off-diagonals
+    ################################################################################
+
+    # we used tf32 to maintain matrix inverse's precision whenever possible.
+    b_Ai10 = -tl.dot(
+        tl.dot(b_Ai11, b_Akk10, input_precision=SOLVE_TRIL_DOT_PRECISION),
+        b_Ai00,
+        input_precision=SOLVE_TRIL_DOT_PRECISION,
+    )
+    b_Ai21 = -tl.dot(
+        tl.dot(b_Ai22, b_Akk21, input_precision=SOLVE_TRIL_DOT_PRECISION),
+        b_Ai11,
+        input_precision=SOLVE_TRIL_DOT_PRECISION,
+    )
+    b_Ai32 = -tl.dot(
+        tl.dot(b_Ai33, b_Akk32, input_precision=SOLVE_TRIL_DOT_PRECISION),
+        b_Ai22,
+        input_precision=SOLVE_TRIL_DOT_PRECISION,
+    )
+
+    b_Ai20 = -tl.dot(
+        b_Ai22,
+        tl.dot(b_Akk20, b_Ai00, input_precision=SOLVE_TRIL_DOT_PRECISION)
+        + tl.dot(b_Akk21, b_Ai10, input_precision=SOLVE_TRIL_DOT_PRECISION),
+        input_precision=SOLVE_TRIL_DOT_PRECISION,
+    )
+    b_Ai31 = -tl.dot(
+        b_Ai33,
+        tl.dot(b_Akk31, b_Ai11, input_precision=SOLVE_TRIL_DOT_PRECISION)
+        + tl.dot(b_Akk32, b_Ai21, input_precision=SOLVE_TRIL_DOT_PRECISION),
+        input_precision=SOLVE_TRIL_DOT_PRECISION,
+    )
+    b_Ai30 = -tl.dot(
+        b_Ai33,
+        tl.dot(b_Akk30, b_Ai00, input_precision=SOLVE_TRIL_DOT_PRECISION)
+        + tl.dot(b_Akk31, b_Ai10, input_precision=SOLVE_TRIL_DOT_PRECISION)
+        + tl.dot(b_Akk32, b_Ai20, input_precision=SOLVE_TRIL_DOT_PRECISION),
+        input_precision=SOLVE_TRIL_DOT_PRECISION,
+    )
+
+    ################################################################################
+    # store full Akk_inv to Akk
+    ################################################################################
+
+    p_Akk00 = tl.make_block_ptr(
+        Akk, (T, BT), (HV * BT, 1), (i_tc0, 0), (BC, BC), (1, 0)
+    )
+    p_Akk10 = tl.make_block_ptr(
+        Akk, (T, BT), (HV * BT, 1), (i_tc1, 0), (BC, BC), (1, 0)
+    )
+    p_Akk11 = tl.make_block_ptr(
+        Akk, (T, BT), (HV * BT, 1), (i_tc1, BC), (BC, BC), (1, 0)
+    )
+    p_Akk20 = tl.make_block_ptr(
+        Akk, (T, BT), (HV * BT, 1), (i_tc2, 0), (BC, BC), (1, 0)
+    )
+    p_Akk21 = tl.make_block_ptr(
+        Akk, (T, BT), (HV * BT, 1), (i_tc2, BC), (BC, BC), (1, 0)
+    )
+    p_Akk22 = tl.make_block_ptr(
+        Akk, (T, BT), (HV * BT, 1), (i_tc2, 2 * BC), (BC, BC), (1, 0)
+    )
+    p_Akk30 = tl.make_block_ptr(
+        Akk, (T, BT), (HV * BT, 1), (i_tc3, 0), (BC, BC), (1, 0)
+    )
+    p_Akk31 = tl.make_block_ptr(
+        Akk, (T, BT), (HV * BT, 1), (i_tc3, BC), (BC, BC), (1, 0)
+    )
+    p_Akk32 = tl.make_block_ptr(
+        Akk, (T, BT), (HV * BT, 1), (i_tc3, 2 * BC), (BC, BC), (1, 0)
+    )
+    p_Akk33 = tl.make_block_ptr(
+        Akk, (T, BT), (HV * BT, 1), (i_tc3, 3 * BC), (BC, BC), (1, 0)
+    )
+
+    tl.store(p_Akk00, b_Ai00.to(Akk.dtype.element_ty), boundary_check=(0, 1))
+    tl.store(p_Akk10, b_Ai10.to(Akk.dtype.element_ty), boundary_check=(0, 1))
+    tl.store(p_Akk11, b_Ai11.to(Akk.dtype.element_ty), boundary_check=(0, 1))
+    tl.store(p_Akk20, b_Ai20.to(Akk.dtype.element_ty), boundary_check=(0, 1))
+    tl.store(p_Akk21, b_Ai21.to(Akk.dtype.element_ty), boundary_check=(0, 1))
+    tl.store(p_Akk22, b_Ai22.to(Akk.dtype.element_ty), boundary_check=(0, 1))
+    tl.store(p_Akk30, b_Ai30.to(Akk.dtype.element_ty), boundary_check=(0, 1))
+    tl.store(p_Akk31, b_Ai31.to(Akk.dtype.element_ty), boundary_check=(0, 1))
+    tl.store(p_Akk32, b_Ai32.to(Akk.dtype.element_ty), boundary_check=(0, 1))
+    tl.store(p_Akk33, b_Ai33.to(Akk.dtype.element_ty), boundary_check=(0, 1))
+
+
+@triton.heuristics(
+    {
+        "IS_VARLEN": lambda args: args["cu_seqlens"] is not None,
+    }
+)
+@triton.autotune(
+    configs=[
+        triton.Config({}, num_warps=num_warps, num_stages=num_stages)
+        for num_warps in [1, 2, 4, 8]
+        for num_stages in [2, 3, 4]
+    ],
+    key=["BK", "NC", "BT", "HV"],
+)
+@triton.jit(do_not_specialize=["B", "T"])
+def chunk_kda_fwd_kernel_intra_sub_chunk(
+    q,
+    k,
+    g,
+    beta,
+    Aqk,
+    Akk,
+    scale,
+    cu_seqlens,
+    chunk_indices,
+    T,
+    H: tl.constexpr,
+    HV: tl.constexpr,
+    K: tl.constexpr,
+    BT: tl.constexpr,
+    BC: tl.constexpr,
+    BK: tl.constexpr,
+    IS_VARLEN: tl.constexpr,
+    USE_GATHER: tl.constexpr,
+):
+    i_t, i_i, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2)
+    i_b, i_hv = i_bh // HV, i_bh % HV
+    i_h = i_hv // (HV // H)
+
+    if IS_VARLEN:
+        i_n, i_t = (
+            tl.load(chunk_indices + i_t * 2).to(tl.int32),
+            tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32),
+        )
+        bos, eos = (
+            tl.load(cu_seqlens + i_n).to(tl.int32),
+            tl.load(cu_seqlens + i_n + 1).to(tl.int32),
+        )
+        T = eos - bos
+    else:
+        bos, eos = i_b * T, i_b * T + T
+
+    i_ti = i_t * BT + i_i * BC
+    if i_ti >= T:
+        return
+
+    o_c = i_ti + tl.arange(0, BC)
+    m_c = o_c < T
+
+    q = q + (bos * H + i_h) * K
+    k = k + (bos * H + i_h) * K
+    g = g + (bos * HV + i_hv) * K
+    beta = beta + bos * HV + i_hv
+    Aqk = Aqk + (bos * HV + i_hv) * BT
+    Akk = Akk + (bos * HV + i_hv) * BC
+
+    p_q = tl.make_block_ptr(q, (T, K), (H * K, 1), (i_ti, 0), (BC, BK), (1, 0))
+    p_k = tl.make_block_ptr(k, (T, K), (H * K, 1), (i_ti, 0), (BC, BK), (1, 0))
+    p_g = tl.make_block_ptr(g, (T, K), (HV * K, 1), (i_ti, 0), (BC, BK), (1, 0))
+
+    p_beta = tl.make_block_ptr(beta, (T,), (HV,), (i_ti,), (BC,), (0,))
+
+    b_q = tl.load(p_q, boundary_check=(0, 1))
+    b_k = tl.load(p_k, boundary_check=(0, 1))
+    b_g = tl.load(p_g, boundary_check=(0, 1))
+    b_beta = tl.load(p_beta, boundary_check=(0,)).to(tl.float32)
+
+    if USE_GATHER:
+        b_gn = gather(
+            b_g, tl.full([1, BK], min(BC // 2, T - i_ti - 1), dtype=tl.int16), axis=0
+        )
+    else:
+        # caculate offset
+        p_gn = g + (i_ti + min(BC // 2, T - i_ti - 1)) * HV * K + tl.arange(0, BK)
+        b_gn = tl.load(p_gn, mask=tl.arange(0, BK) < K, other=0.0)
+        b_gn = b_gn[None, :]
+
+    # current block, keep numerical stability by subtracting the left boundary
+    # less than 85 to avoid overflow in exp2
+    b_gm = (b_g - b_gn).to(tl.float32)
+
+    b_gq = tl.where(m_c[:, None], exp2(b_gm), 0.0)
+    b_gk = tl.where(m_c[:, None], exp2(-b_gm), 0.0)
+
+    b_kgt = tl.trans(b_k * b_gk)
+
+    b_Aqk = tl.dot(b_q * b_gq, b_kgt) * scale
+    b_Akk = tl.dot(b_k * b_gq, b_kgt) * b_beta[:, None]
+
+    o_i = tl.arange(0, BC)
+    m_Aqk = o_i[:, None] >= o_i[None, :]
+    m_Akk = o_i[:, None] > o_i[None, :]
+    m_I = o_i[:, None] == o_i[None, :]
+
+    b_Aqk = tl.where(m_Aqk, b_Aqk, 0.0)
+    b_Akk = tl.where(m_Akk, b_Akk, 0.0)
+
+    p_Aqk = tl.make_block_ptr(
+        Aqk, (T, BT), (HV * BT, 1), (i_ti, i_i * BC), (BC, BC), (1, 0)
+    )
+    p_Akk = tl.make_block_ptr(Akk, (T, BC), (HV * BC, 1), (i_ti, 0), (BC, BC), (1, 0))
+    tl.store(p_Aqk, b_Aqk.to(Aqk.dtype.element_ty), boundary_check=(0, 1))
+    tl.store(p_Akk, b_Akk.to(Akk.dtype.element_ty), boundary_check=(0, 1))
+
+    tl.debug_barrier()
+
+    ################################################################################
+    # forward substitution
+    ################################################################################
+
+    b_Ai = -b_Akk
+    for i in range(2, min(BC, T - i_ti)):
+        b_a = -tl.load(Akk + (i_ti + i) * HV * BC + o_i)
+        b_a = tl.where(o_i < i, b_a, 0.0)
+        b_a += tl.sum(b_a[:, None] * b_Ai, 0)
+        b_Ai = tl.where((o_i == i)[:, None], b_a, b_Ai)
+    b_Ai += m_I
+    tl.store(p_Akk, b_Ai.to(Akk.dtype.element_ty), boundary_check=(0, 1))
+
+
+def chunk_kda_fwd_intra(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    gk: torch.Tensor | None = None,
+    beta: torch.Tensor | None = None,
+    scale: float | None = None,
+    cu_seqlens: torch.LongTensor | None = None,
+    chunk_size: int = 64,
+    chunk_indices: torch.LongTensor | None = None,
+    safe_gate: bool = False,
+):
+    B, T, H, K, HV = *k.shape, gk.shape[2]
+    BT = chunk_size
+    BC = 16
+    if chunk_indices is None and cu_seqlens is not None:
+        chunk_indices = prepare_chunk_indices(cu_seqlens, BT)
+    NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices)
+    NC = triton.cdiv(BT, BC)
+
+    Aqk = torch.empty(B, T, HV, BT, device=k.device, dtype=k.dtype)
+    # Akk must be zero-initialized - kernel only writes lower triangular
+    Akk = torch.zeros(B, T, HV, BT, device=k.device, dtype=k.dtype)
+    # Separate fp32 buffer for diagonal 16x16 blocks (for precision in solve_tril)
+    Akkd = torch.empty(B, T, HV, BC, device=k.device, dtype=torch.float32)
+
+    # Compute diagonal blocks into Akkd in fp32.
+    if safe_gate:
+        grid = (NT, NC, B * HV)
+        BK = triton.next_power_of_2(K)
+        chunk_kda_fwd_kernel_intra_sub_chunk[grid](
+            q=q,
+            k=k,
+            g=gk,
+            beta=beta,
+            Aqk=Aqk,
+            Akk=Akkd,
+            scale=scale,
+            cu_seqlens=cu_seqlens,
+            chunk_indices=chunk_indices,
+            T=T,
+            H=H,
+            HV=HV,
+            K=K,
+            BT=BT,
+            BC=BC,
+            BK=BK,
+            USE_GATHER=is_gather_supported,
+        )
+    else:
+        Aqk, Akkd = chunk_kda_fwd_intra_token_parallel(
+            q=q,
+            k=k,
+            gk=gk,
+            beta=beta,
+            Aqk=Aqk,
+            Akk=Akkd,
+            scale=scale,
+            cu_seqlens=cu_seqlens,
+            chunk_size=BT,
+            sub_chunk_size=BC,
+        )
+
+    # Step 2: Fused inter + solve_tril (works for both fixed-len and varlen)
+    solve_tril_dot_precision = (
+        "tf32"
+        if current_platform.is_cuda() and current_platform.has_device_capability(80)
+        else "ieee"
+    )
+    grid = (NT, B * HV)
+    chunk_kda_fwd_kernel_inter_solve_fused[grid](
+        q=q,
+        k=k,
+        g=gk,
+        beta=beta,
+        Aqk=Aqk,
+        Akkd=Akkd,
+        Akk=Akk,
+        scale=scale,
+        cu_seqlens=cu_seqlens,
+        chunk_indices=chunk_indices,
+        T=T,
+        H=H,
+        HV=HV,
+        K=K,
+        BT=BT,
+        BC=BC,
+        USE_SAFE_GATE=safe_gate,
+        SOLVE_TRIL_DOT_PRECISION=solve_tril_dot_precision,
+    )
+    return Aqk, Akk

vllm/models/kimi_k3/amd/ops/third_party/kda/chunk_intra_token_parallel.py

@@ -0,0 +1,197 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Songlin Yang, Yu Zhang, Zhiyuan Li
+#
+# This file contains code copied from the flash-linear-attention project.
+# The original source was licensed under the MIT license.
+# Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li
+# Forward-only adaptation of flash-linear-attention 0.5.0.
+# ruff: noqa: E501
+
+# Token-parallel implementation of KDA intra chunk kernel
+
+import torch
+
+from vllm.third_party.flash_linear_attention.ops.op import exp2
+from vllm.triton_utils import tl, triton
+
+
+@triton.heuristics(
+    {
+        "IS_VARLEN": lambda args: args["cu_seqlens"] is not None,
+    }
+)
+@triton.autotune(
+    configs=[
+        triton.Config({"BH": BH}, num_warps=num_warps)
+        for BH in [1, 2, 4, 8]
+        for num_warps in [1, 2, 4, 8]
+    ],
+    key=["K", "H", "HV"],
+)
+@triton.jit(do_not_specialize=["T", "N"])
+def chunk_kda_fwd_kernel_intra_token_parallel(
+    q,
+    k,
+    g,
+    beta,
+    Aqk,
+    Akk,
+    scale,
+    cu_seqlens,
+    N,
+    T,
+    H: tl.constexpr,
+    HV: tl.constexpr,
+    K: tl.constexpr,
+    BT: tl.constexpr,
+    BC: tl.constexpr,
+    BH: tl.constexpr,
+    IS_VARLEN: tl.constexpr,
+):
+    i_tg, i_hg = tl.program_id(0), tl.program_id(1)
+
+    if IS_VARLEN:
+        i_n = 0
+        left, right = 0, N
+
+        # Unrolled binary search (max B=2^32)
+        # We can limit iterations based on expected max batch size if needed
+        # 20 iterations covers B=1M, usually enough
+        for _ in range(20):
+            if left < right:
+                mid = (left + right) // 2
+                if i_tg < tl.load(cu_seqlens + mid + 1).to(tl.int32):
+                    right = mid
+                else:
+                    left = mid + 1
+        i_n = left
+
+        bos, eos = (
+            tl.load(cu_seqlens + i_n).to(tl.int32),
+            tl.load(cu_seqlens + i_n + 1).to(tl.int32),
+        )
+        T = eos - bos
+        i_t = i_tg - bos
+    else:
+        bos = (i_tg // T) * T
+        i_t = i_tg % T
+
+    if i_t >= T:
+        return
+
+    i_c = i_t // BT
+    i_s = (i_t % BT) // BC
+    i_tc = i_c * BT
+    i_ts = i_tc + i_s * BC
+
+    G: tl.constexpr = HV // H
+
+    q += bos * H * K
+    k += bos * H * K
+    g += bos * HV * K
+    Aqk += bos * HV * BT
+    Akk += bos * HV * BC
+    beta += bos * HV
+
+    BK: tl.constexpr = triton.next_power_of_2(K)
+    o_hv = i_hg * BH + tl.arange(0, BH)
+    o_h = o_hv // G
+    o_k = tl.arange(0, BK)
+    m_hv = o_hv < HV
+    m_k = o_k < K
+    m_hk = m_hv[:, None] & m_k[None, :]
+
+    # q/k: [B, T, H, K], manual load via mapped qk head index
+    p_qk = o_h[:, None] * K + o_k[None, :]
+    b_q = tl.load(q + i_t * H * K + p_qk, mask=m_hk, other=0).to(tl.float32)
+    b_k = tl.load(k + i_t * H * K + p_qk, mask=m_hk, other=0).to(tl.float32)
+
+    # g: [B, T, HV, K], beta: [B, T, HV]
+    p_g = tl.make_block_ptr(
+        g + i_t * HV * K, (HV, K), (K, 1), (i_hg * BH, 0), (BH, BK), (1, 0)
+    )
+    p_beta = tl.make_block_ptr(beta + i_t * HV, (HV,), (1,), (i_hg * BH,), (BH,), (0,))
+    b_g = tl.load(p_g, boundary_check=(0, 1)).to(tl.float32)
+    b_beta = tl.load(p_beta, boundary_check=(0,)).to(tl.float32)
+    b_k *= b_beta[:, None]
+
+    for j in range(i_ts, min(i_t + 1, min(T, i_ts + BC))):
+        b_kj = tl.load(k + j * H * K + p_qk, mask=m_hk, other=0).to(tl.float32)
+        p_gj = tl.make_block_ptr(
+            g + j * HV * K, (HV, K), (K, 1), (i_hg * BH, 0), (BH, BK), (1, 0)
+        )
+        b_gj = tl.load(p_gj, boundary_check=(0, 1)).to(tl.float32)
+
+        b_kgj = tl.where(m_k[None, :], b_kj * exp2(b_g - b_gj), 0.0)
+        b_Aqk = tl.sum(b_q * b_kgj, axis=1) * scale
+        b_Akk = tl.sum(b_k * b_kgj, axis=1) * tl.where(j < i_t, 1.0, 0.0)
+
+        tl.store(
+            Aqk + i_t * HV * BT + o_hv * BT + j % BT,
+            b_Aqk.to(Aqk.dtype.element_ty),
+            mask=m_hv,
+        )
+        tl.store(
+            Akk + i_t * HV * BC + o_hv * BC + j - i_ts,
+            b_Akk.to(Akk.dtype.element_ty),
+            mask=m_hv,
+        )
+
+
+def chunk_kda_fwd_intra_token_parallel(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    gk: torch.Tensor,
+    beta: torch.Tensor,
+    Aqk: torch.Tensor,
+    Akk: torch.Tensor,
+    scale: float,
+    cu_seqlens: torch.LongTensor | None = None,
+    chunk_size: int = 64,
+    sub_chunk_size: int = 16,
+) -> None:
+    """
+    Token-parallel implementation: each token gets its own thread block.
+    Supports both fixed-length and variable-length sequences.
+    Reduces wasted computation on padding.
+
+    Writes directly to Aqk and Akk tensors (in-place).
+
+    Args:
+        q: [B, T, H, K]
+        k: [B, T, H, K]
+        gk: [B, T, HV, K] cumsum of gates (HV >= H for GVA)
+        beta: [B, T, HV]
+        Aqk: [B, T, HV, BT] output tensor to write to
+        Akk: [B, T, HV, BC] output tensor for diagonal blocks (fp32)
+        scale: attention scale
+        chunk_size: BT (default 64)
+        sub_chunk_size: BC (default 16)
+    """
+    B, T, H, K, HV = *q.shape, gk.shape[2]
+    N = len(cu_seqlens) - 1 if cu_seqlens is not None else B
+    BT = chunk_size
+    BC = sub_chunk_size
+
+    def grid(meta):
+        return (B * T, triton.cdiv(HV, meta["BH"]))
+
+    chunk_kda_fwd_kernel_intra_token_parallel[grid](
+        q=q,
+        k=k,
+        g=gk,
+        beta=beta,
+        Aqk=Aqk,
+        Akk=Akk,
+        scale=scale,
+        cu_seqlens=cu_seqlens,
+        N=N,
+        T=T,
+        H=H,
+        HV=HV,
+        K=K,
+        BT=BT,
+        BC=BC,
+    )
+    return Aqk, Akk

vllm/models/kimi_k3/amd/ops/third_party/kda/fused_recurrent.py

@@ -0,0 +1,621 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Songlin Yang, Yu Zhang, Zhiyuan Li
+#
+# This file contains code adapted from the flash-linear-attention project.
+# The original source was licensed under the MIT license.
+# Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li
+# ruff: noqa: E501
+
+import torch
+
+from vllm.third_party.flash_linear_attention.ops.op import exp, log
+from vllm.triton_utils import tl, triton
+from vllm.utils.math_utils import cdiv, next_power_of_2
+
+
+@triton.heuristics(
+    {
+        "HAS_DT_BIAS": lambda args: args["dt_bias"] is not None,
+        "USE_LOWER_BOUND": lambda args: args["lower_bound"] is not None,
+    }
+)
+@triton.jit
+def _kda_gate_beta_fwd_kernel(
+    raw_g,
+    raw_beta,
+    A_log,
+    dt_bias,
+    gate,
+    beta_out,
+    lower_bound,
+    softplus_beta: tl.constexpr,
+    softplus_threshold: tl.constexpr,
+    T,
+    stride_g_token: tl.constexpr,
+    stride_beta_token: tl.constexpr,
+    H: tl.constexpr,
+    D: tl.constexpr,
+    BT: tl.constexpr,
+    BD: tl.constexpr,
+    HAS_DT_BIAS: tl.constexpr,
+    USE_LOWER_BOUND: tl.constexpr,
+):
+    i_t, i_h = tl.program_id(0), tl.program_id(1)
+    o_t = i_t * BT + tl.arange(0, BT)
+    o_d = tl.arange(0, BD)
+    m_t = o_t < T
+    m_d = o_d < D
+
+    p_g = raw_g + o_t[:, None] * stride_g_token + i_h * D + o_d[None, :]
+    b_g = tl.load(p_g, mask=m_t[:, None] & m_d[None, :], other=0.0).to(tl.float32)
+    if HAS_DT_BIAS:
+        b_bias = tl.load(
+            dt_bias + i_h * D + o_d,
+            mask=m_d,
+            other=0.0,
+        ).to(tl.float32)
+        b_g += b_bias[None, :]
+
+    b_a = exp(tl.load(A_log + i_h).to(tl.float32))
+    if USE_LOWER_BOUND:
+        b_gate = lower_bound * tl.sigmoid(b_a * b_g)
+    else:
+        b_scaled = b_g * softplus_beta
+        b_softplus = tl.where(
+            b_scaled > softplus_threshold,
+            b_g,
+            log(1.0 + tl.exp(b_scaled)) / softplus_beta,
+        )
+        b_gate = -b_a * b_softplus
+
+    p_gate = gate + (o_t[:, None] * H + i_h) * D + o_d[None, :]
+    tl.store(
+        p_gate,
+        b_gate,
+        mask=m_t[:, None] & m_d[None, :],
+    )
+
+    b_beta = tl.load(
+        raw_beta + o_t * stride_beta_token + i_h,
+        mask=m_t,
+        other=0.0,
+    ).to(tl.float32)
+    tl.store(beta_out + o_t * H + i_h, tl.sigmoid(b_beta), mask=m_t)
+
+
+def _fused_kda_gate_beta(
+    raw_g: torch.Tensor,
+    raw_beta: torch.Tensor,
+    A_log: torch.Tensor,
+    dt_bias: torch.Tensor | None,
+    lower_bound: float | None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    B, T, H, D = raw_g.shape
+    assert B == 1
+    assert raw_beta.shape == (B, T, H)
+    assert raw_g.stride()[2:] == (D, 1)
+    assert raw_beta.stride(2) == 1
+    gate = torch.empty((B, T, H, D), dtype=torch.float32, device=raw_g.device)
+    beta = torch.empty((B, T, H), dtype=torch.float32, device=raw_beta.device)
+
+    BT = 16
+    _kda_gate_beta_fwd_kernel[(cdiv(T, BT), H)](
+        raw_g=raw_g,
+        raw_beta=raw_beta,
+        A_log=A_log,
+        dt_bias=dt_bias,
+        gate=gate,
+        beta_out=beta,
+        lower_bound=lower_bound,
+        softplus_beta=1.0,
+        softplus_threshold=20.0,
+        T=T,
+        stride_g_token=raw_g.stride(1),
+        stride_beta_token=raw_beta.stride(1),
+        H=H,
+        D=D,
+        BT=BT,
+        BD=next_power_of_2(D),
+        num_warps=4,
+    )
+    return gate, beta
+
+
+@triton.heuristics(
+    {
+        "IS_SPEC_DECODING": lambda args: args["num_accepted_tokens"] is not None,
+        "HAS_DT_BIAS": lambda args: args["dt_bias"] is not None,
+        "USE_LOWER_BOUND": lambda args: args["lower_bound"] is not None,
+    }
+)
+@triton.jit(do_not_specialize=["N", "T"])
+def fused_recurrent_kda_fwd_kernel(
+    q,
+    k,
+    v,
+    g,
+    beta,
+    A_log,
+    dt_bias,
+    out,
+    state,
+    cu_seqlens,
+    state_indices,
+    num_accepted_tokens,
+    lower_bound,
+    scale: tl.constexpr,
+    N: tl.int64,
+    T: tl.int64,
+    H: tl.constexpr,
+    K: tl.constexpr,
+    V: tl.constexpr,
+    BK: tl.constexpr,
+    BV: tl.constexpr,
+    stride_qkv_token: tl.constexpr,
+    stride_g_token: tl.constexpr,
+    stride_beta_token: tl.constexpr,
+    stride_out_token: tl.constexpr,
+    stride_state_token: tl.constexpr,
+    stride_indices_seq: tl.constexpr,
+    IS_SPEC_DECODING: tl.constexpr,
+    USE_QK_L2NORM_IN_KERNEL: tl.constexpr,
+    USE_GATE_IN_KERNEL: tl.constexpr,
+    APPLY_BETA_SIGMOID: tl.constexpr,
+    HAS_DT_BIAS: tl.constexpr,
+    USE_LOWER_BOUND: tl.constexpr,
+    num_stages: tl.constexpr,
+):
+    pid = tl.program_id(0)
+    i_v = pid % tl.cdiv(V, BV)
+    i_nh = pid // tl.cdiv(V, BV)
+    i_n, i_h = i_nh // H, i_nh % H
+    bos = tl.load(cu_seqlens + i_n).to(tl.int64)
+    eos = tl.load(cu_seqlens + i_n + 1).to(tl.int64)
+    sequence_length = eos - bos
+    if sequence_length == 0:
+        return
+
+    o_k = tl.arange(0, BK)
+    o_v = i_v * BV + tl.arange(0, BV)
+    m_k = o_k < K
+    m_v = o_v < V
+    m_state = m_v[:, None] & m_k[None, :]
+
+    if IS_SPEC_DECODING:
+        initial_token = tl.load(num_accepted_tokens + i_n).to(tl.int64) - 1
+    else:
+        initial_token = 0
+    state_index = tl.load(state_indices + i_n * stride_indices_seq + initial_token).to(
+        tl.int64
+    )
+    p_out = out + bos * stride_out_token + i_h * V + o_v
+    if state_index <= 0:
+        tl.store(p_out, tl.zeros([BV], dtype=tl.float32), mask=m_v)
+        return
+
+    p_state = (
+        state
+        + state_index * stride_state_token
+        + i_h * V * K
+        + o_v[:, None] * K
+        + o_k[None, :]
+    )
+    b_state = tl.load(p_state, mask=m_state, other=0.0).to(tl.float32)
+
+    p_q = q + bos * stride_qkv_token + i_h * K + o_k
+    p_k = k + bos * stride_qkv_token + i_h * K + o_k
+    p_v = v + bos * stride_qkv_token + i_h * V + o_v
+    p_g = g + bos * stride_g_token + i_h * K + o_k
+    p_beta = beta + bos * stride_beta_token + i_h
+    for i_t in tl.range(0, sequence_length, num_stages=num_stages):
+        b_q = tl.load(p_q, mask=m_k, other=0.0, eviction_policy="evict_last").to(
+            tl.float32
+        )
+        b_k = tl.load(p_k, mask=m_k, other=0.0, eviction_policy="evict_last").to(
+            tl.float32
+        )
+        b_v = tl.load(p_v, mask=m_v, other=0.0, eviction_policy="evict_first").to(
+            tl.float32
+        )
+        if USE_QK_L2NORM_IN_KERNEL:
+            b_q = b_q / tl.sqrt(tl.sum(b_q * b_q) + 1e-6)
+            b_k = b_k / tl.sqrt(tl.sum(b_k * b_k) + 1e-6)
+        b_q *= scale
+
+        b_gate = tl.load(
+            p_g,
+            mask=m_k,
+            other=0.0,
+            eviction_policy="evict_last",
+        ).to(tl.float32)
+        if USE_GATE_IN_KERNEL:
+            if HAS_DT_BIAS:
+                b_bias = tl.load(
+                    dt_bias + i_h * K + o_k,
+                    mask=m_k,
+                    other=0.0,
+                ).to(tl.float32)
+                b_gate += b_bias
+            b_a = exp(tl.load(A_log + i_h).to(tl.float32))
+            if USE_LOWER_BOUND:
+                b_gate = lower_bound * tl.sigmoid(b_a * b_gate)
+            else:
+                b_softplus = tl.where(
+                    b_gate > 20.0,
+                    b_gate,
+                    log(1.0 + tl.exp(b_gate)),
+                )
+                b_gate = -b_a * b_softplus
+
+        b_state *= exp(b_gate[None, :])
+        b_v -= tl.sum(b_state * b_k[None, :], axis=1)
+        b_beta = tl.load(p_beta, eviction_policy="evict_last").to(tl.float32)
+        if APPLY_BETA_SIGMOID:
+            b_beta = tl.sigmoid(b_beta)
+        b_v *= b_beta
+        b_state += b_v[:, None] * b_k[None, :]
+        b_out = tl.sum(b_state * b_q[None, :], axis=1)
+        tl.store(
+            p_out,
+            b_out.to(p_out.dtype.element_ty),
+            mask=m_v,
+            eviction_policy="evict_first",
+        )
+
+        final_state_index = tl.load(state_indices + i_n * stride_indices_seq + i_t).to(
+            tl.int64
+        )
+        if final_state_index > 0:
+            p_final_state = (
+                state
+                + final_state_index * stride_state_token
+                + i_h * V * K
+                + o_v[:, None] * K
+                + o_k[None, :]
+            )
+            tl.store(
+                p_final_state,
+                b_state.to(p_final_state.dtype.element_ty),
+                mask=m_state,
+            )
+
+        p_q += stride_qkv_token
+        p_k += stride_qkv_token
+        p_v += stride_qkv_token
+        p_g += stride_g_token
+        p_beta += stride_beta_token
+        p_out += stride_out_token
+
+
+def fused_recurrent_kda_fwd(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    g: torch.Tensor,
+    beta: torch.Tensor,
+    scale: float | None = None,
+    initial_state: torch.Tensor | None = None,
+    inplace_final_state: bool = True,
+    cu_seqlens: torch.Tensor | None = None,
+    ssm_state_indices: torch.Tensor | None = None,
+    num_accepted_tokens: torch.Tensor | None = None,
+    use_qk_l2norm_in_kernel: bool = True,
+    A_log: torch.Tensor | None = None,
+    dt_bias: torch.Tensor | None = None,
+    lower_bound: float | None = None,
+    use_gate_in_kernel: bool = False,
+    use_beta_sigmoid_in_kernel: bool = False,
+    out: torch.Tensor | None = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Launch recurrent KDA with dense inner dimensions and row strides."""
+    B, T, H, K = q.shape
+    V = v.shape[-1]
+    assert B == 1 and k.shape == q.shape
+    assert v.shape == (B, T, H, V) and g.shape == (B, T, H, K)
+    assert beta.shape == (B, T, H)
+    assert initial_state is not None
+    assert cu_seqlens is not None
+    assert ssm_state_indices is not None
+    assert inplace_final_state
+    if out is None:
+        out = torch.empty_like(v)
+    assert out.shape == v.shape
+    assert initial_state.shape[1:] == (H, V, K)
+    assert ssm_state_indices.ndim in (1, 2)
+
+    assert q.stride()[2:] == k.stride()[2:] == (K, 1)
+    assert v.stride()[2:] == out.stride()[2:] == (V, 1)
+    assert g.stride()[2:] == (K, 1)
+    assert beta.stride(2) == 1
+    assert q.stride(1) == k.stride(1) == v.stride(1)
+    assert initial_state.stride()[1:] == (V * K, K, 1)
+    N = cu_seqlens.numel() - 1
+    if ssm_state_indices.ndim == 1:
+        assert T == N
+        assert num_accepted_tokens is None
+    else:
+        assert ssm_state_indices.stride(1) == 1
+    assert cu_seqlens.is_contiguous()
+    if use_gate_in_kernel:
+        assert A_log is not None and A_log.is_contiguous()
+        assert dt_bias is None or dt_bias.is_contiguous()
+
+    if scale is None:
+        scale = K**-0.5
+
+    BV = 32 if use_gate_in_kernel else 8
+    num_warps = 4 if use_gate_in_kernel else 1
+    grid = (cdiv(V, BV) * N * H,)
+    fused_recurrent_kda_fwd_kernel[grid](
+        q=q,
+        k=k,
+        v=v,
+        g=g,
+        beta=beta,
+        A_log=A_log,
+        dt_bias=dt_bias,
+        out=out,
+        state=initial_state,
+        cu_seqlens=cu_seqlens,
+        state_indices=ssm_state_indices,
+        num_accepted_tokens=num_accepted_tokens,
+        lower_bound=lower_bound,
+        scale=scale,
+        N=N,
+        T=T,
+        H=H,
+        K=K,
+        V=V,
+        BK=next_power_of_2(K),
+        BV=BV,
+        stride_qkv_token=q.stride(1),
+        stride_g_token=g.stride(1),
+        stride_beta_token=beta.stride(1),
+        stride_out_token=out.stride(1),
+        stride_state_token=initial_state.stride(0),
+        stride_indices_seq=ssm_state_indices.stride(0),
+        IS_SPEC_DECODING=num_accepted_tokens is not None,
+        USE_QK_L2NORM_IN_KERNEL=use_qk_l2norm_in_kernel,
+        USE_GATE_IN_KERNEL=use_gate_in_kernel,
+        APPLY_BETA_SIGMOID=use_beta_sigmoid_in_kernel,
+        num_warps=num_warps,
+        num_stages=2,
+    )
+    return out, initial_state
+
+
+def fused_recurrent_kda(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    raw_g: torch.Tensor,
+    raw_beta: torch.Tensor,
+    A_log: torch.Tensor,
+    dt_bias: torch.Tensor | None,
+    lower_bound: float | None,
+    initial_state: torch.Tensor,
+    cu_seqlens: torch.Tensor,
+    ssm_state_indices: torch.Tensor,
+    num_accepted_tokens: torch.Tensor | None = None,
+    out: torch.Tensor | None = None,
+    fuse_gate: bool | None = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Run recurrent KDA from raw gate and beta inputs.
+
+    This vLLM wrapper applies the gate activation and beta sigmoid, selecting
+    whether to materialize them before launching the recurrent kernel.
+    """
+    if fuse_gate is None:
+        # gfx950: always fuse the gate.
+        fuse_gate = True
+
+    if fuse_gate:
+        gate = raw_g
+        beta = raw_beta
+    else:
+        gate, beta = _fused_kda_gate_beta(
+            raw_g,
+            raw_beta,
+            A_log,
+            dt_bias,
+            lower_bound,
+        )
+    return fused_recurrent_kda_fwd(
+        q=q,
+        k=k,
+        v=v,
+        g=gate,
+        beta=beta,
+        scale=q.shape[-1] ** -0.5,
+        initial_state=initial_state,
+        inplace_final_state=True,
+        cu_seqlens=cu_seqlens,
+        ssm_state_indices=ssm_state_indices,
+        num_accepted_tokens=num_accepted_tokens,
+        use_qk_l2norm_in_kernel=True,
+        A_log=A_log if fuse_gate else None,
+        dt_bias=dt_bias if fuse_gate else None,
+        lower_bound=lower_bound if fuse_gate else None,
+        use_gate_in_kernel=fuse_gate,
+        use_beta_sigmoid_in_kernel=fuse_gate,
+        out=out,
+    )
+
+
+@triton.jit
+def fused_recurrent_kda_packed_decode_kernel(
+    mixed_qkv,
+    raw_g,
+    raw_beta,
+    A_log,
+    dt_bias,
+    out,
+    state,
+    state_indices,
+    lower_bound,
+    scale: tl.constexpr,
+    stride_mixed_token: tl.constexpr,
+    stride_g_token: tl.constexpr,
+    stride_beta_token: tl.constexpr,
+    stride_state_token: tl.constexpr,
+    H: tl.constexpr,
+    K: tl.constexpr,
+    V: tl.constexpr,
+    BK: tl.constexpr,
+    BV: tl.constexpr,
+    SOFTPLUS_THRESHOLD: tl.constexpr,
+    USE_LOWER_BOUND: tl.constexpr,
+):
+    i_v, i_nh = tl.program_id(0), tl.program_id(1)
+    i_n, i_h = i_nh // H, i_nh % H
+
+    o_k = tl.arange(0, BK)
+    o_v = i_v * BV + tl.arange(0, BV)
+    mask_k = o_k < K
+    mask_v = o_v < V
+    mask_state = mask_v[:, None] & mask_k[None, :]
+
+    state_idx = tl.load(state_indices + i_n).to(tl.int64)
+    p_out = out + (i_n * H + i_h) * V + o_v
+    if state_idx <= 0:
+        tl.store(p_out, tl.zeros([BV], dtype=tl.float32), mask=mask_v)
+        return
+
+    p_state = state + state_idx * stride_state_token
+    p_state += i_h * V * K + o_v[:, None] * K + o_k[None, :]
+    b_state = tl.load(p_state, mask=mask_state, other=0).to(tl.float32)
+
+    # Q, K, and V occupy consecutive channel ranges, while the token stride
+    # may also include the output-gate projection that follows packed QKV.
+    p_mixed = mixed_qkv + i_n * stride_mixed_token
+    b_q = tl.load(p_mixed + i_h * K + o_k, mask=mask_k, other=0).to(tl.float32)
+    b_k = tl.load(
+        p_mixed + H * K + i_h * K + o_k,
+        mask=mask_k,
+        other=0,
+    ).to(tl.float32)
+    b_v = tl.load(
+        p_mixed + 2 * H * K + i_h * V + o_v,
+        mask=mask_v,
+        other=0,
+    ).to(tl.float32)
+
+    b_q /= tl.sqrt(tl.sum(b_q * b_q) + 1e-6)
+    b_k /= tl.sqrt(tl.sum(b_k * b_k) + 1e-6)
+    b_q *= scale
+
+    p_g = raw_g + i_n * stride_g_token + i_h * K + o_k
+    b_g = tl.load(p_g, mask=mask_k, other=0).to(tl.float32)
+    b_bias = tl.load(dt_bias + i_h * K + o_k, mask=mask_k, other=0).to(tl.float32)
+    b_a = exp(tl.load(A_log + i_h).to(tl.float32))
+    b_g += b_bias
+    if USE_LOWER_BOUND:
+        b_gate = lower_bound * tl.sigmoid(b_a * b_g)
+    else:
+        b_softplus = tl.where(
+            b_g > SOFTPLUS_THRESHOLD,
+            b_g,
+            log(1.0 + tl.exp(b_g)),
+        )
+        b_gate = -b_a * b_softplus
+
+    b_state *= exp(b_gate[None, :])
+    b_v -= tl.sum(b_state * b_k[None, :], axis=1)
+    b_beta = tl.sigmoid(
+        tl.load(raw_beta + i_n * stride_beta_token + i_h).to(tl.float32)
+    )
+    b_v *= b_beta
+    b_state += b_v[:, None] * b_k[None, :]
+    b_out = tl.sum(b_state * b_q[None, :], axis=1)
+
+    tl.store(p_out, b_out.to(p_out.dtype.element_ty), mask=mask_v)
+    tl.store(p_state, b_state.to(p_state.dtype.element_ty), mask=mask_state)
+
+
+def fused_recurrent_kda_packed_decode(
+    mixed_qkv: torch.Tensor,
+    raw_g: torch.Tensor,
+    raw_beta: torch.Tensor,
+    A_log: torch.Tensor,
+    dt_bias: torch.Tensor,
+    lower_bound: float | None,
+    initial_state: torch.Tensor,
+    state_indices: torch.Tensor,
+    scale: float | None = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Run one-token KDA decode directly from packed post-conv QKV."""
+    if mixed_qkv.ndim != 2 or mixed_qkv.stride(-1) != 1:
+        raise ValueError("`mixed_qkv` must be 2D and contiguous in its last dim.")
+    if raw_g.ndim != 4 or raw_g.shape[0] != 1:
+        raise ValueError("`raw_g` must have shape [1, B, H, K].")
+    if raw_beta.ndim != 3 or raw_beta.shape[0] != 1:
+        raise ValueError("`raw_beta` must have shape [1, B, H].")
+    if initial_state.ndim != 4:
+        raise ValueError("`initial_state` must have shape [cache, H, V, K].")
+    _, H, V, K = initial_state.shape
+    if raw_g.stride()[2:] != (K, 1):
+        raise ValueError("`raw_g` must be contiguous within each token.")
+    if raw_beta.stride(2) != 1:
+        raise ValueError("`raw_beta` heads must be contiguous.")
+    if initial_state.stride()[1:] != (V * K, K, 1):
+        raise ValueError("`initial_state` must be contiguous within each cache slot.")
+    if state_indices.ndim != 1 or state_indices.stride(0) != 1:
+        raise ValueError("`state_indices` must be contiguous and one-dimensional.")
+    if A_log.ndim != 1 or not A_log.is_contiguous():
+        raise ValueError("`A_log` must be contiguous and one-dimensional.")
+    if not dt_bias.is_contiguous():
+        raise ValueError("`dt_bias` must be contiguous.")
+
+    device = mixed_qkv.device
+    if any(
+        x.device != device
+        for x in (raw_g, raw_beta, A_log, dt_bias, initial_state, state_indices)
+    ):
+        raise ValueError("All packed KDA inputs must be on the same device.")
+
+    B = mixed_qkv.shape[0]
+    if raw_g.shape != (1, B, H, K):
+        raise ValueError(f"Unexpected raw gate shape {tuple(raw_g.shape)}.")
+    if raw_beta.shape != (1, B, H):
+        raise ValueError(f"Unexpected raw beta shape {tuple(raw_beta.shape)}.")
+    if mixed_qkv.shape[1] != 2 * H * K + H * V:
+        raise ValueError(f"Unexpected packed QKV shape {tuple(mixed_qkv.shape)}.")
+    if A_log.numel() != H or dt_bias.numel() != H * K:
+        raise ValueError("`A_log` or `dt_bias` has an incompatible shape.")
+    if state_indices.shape[0] != B:
+        raise ValueError("`state_indices` must contain one entry per token.")
+
+    BK = next_power_of_2(K)
+    BV = min(next_power_of_2(V), 32)
+    if scale is None:
+        scale = K**-0.5
+
+    out = torch.empty((1, B, H, V), dtype=mixed_qkv.dtype, device=device)
+    grid = (cdiv(V, BV), B * H)
+    fused_recurrent_kda_packed_decode_kernel[grid](
+        mixed_qkv=mixed_qkv,
+        raw_g=raw_g,
+        raw_beta=raw_beta,
+        A_log=A_log,
+        dt_bias=dt_bias,
+        out=out,
+        state=initial_state,
+        state_indices=state_indices,
+        lower_bound=lower_bound or 0.0,
+        scale=scale,
+        stride_mixed_token=mixed_qkv.stride(0),
+        stride_g_token=raw_g.stride(1),
+        stride_beta_token=raw_beta.stride(1),
+        stride_state_token=initial_state.stride(0),
+        H=H,
+        K=K,
+        V=V,
+        BK=BK,
+        BV=BV,
+        SOFTPLUS_THRESHOLD=20.0,
+        USE_LOWER_BOUND=lower_bound is not None,
+        num_warps=4,
+        num_stages=2,
+    )
+    return out, initial_state

vllm/models/kimi_k3/common/__init__.py

@@ -0,0 +1,2 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

vllm/models/kimi_k3/common/mm_preprocess.py

@@ -0,0 +1,330 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Shared Kimi-K3 multimodal preprocessing."""
+
+import math
+from collections.abc import Mapping, Sequence
+from typing import Any, cast
+
+import torch
+from transformers import BatchFeature
+
+from vllm.config.multimodal import BaseDummyOptions, ImageDummyOptions
+from vllm.inputs import MultiModalDataDict
+from vllm.logger import init_logger
+from vllm.multimodal.inputs import (
+    MultiModalFieldConfig,
+    MultiModalKwargsItems,
+)
+from vllm.multimodal.parse import ImageProcessorItems, ImageSize, MultiModalDataItems
+from vllm.multimodal.processing import (
+    BaseDummyInputsBuilder,
+    BaseMultiModalProcessor,
+    BaseProcessingInfo,
+    InputProcessingContext,
+    PromptReplacement,
+    PromptUpdate,
+    PromptUpdateDetails,
+)
+from vllm.transformers_utils.configs.kimi_k3 import KimiK3Config
+from vllm.transformers_utils.processor import cached_get_image_processor
+from vllm.transformers_utils.processors.kimi_k3 import KimiK3Processor
+
+logger = init_logger(__name__)
+
+
+def navit_resize_image(
+    width: int,
+    height: int,
+    patch_size: int,
+    merge_kernel_size: int,
+    in_patch_limit: int,
+    patch_limit_on_one_side: int,
+    fixed_output_tokens: int | None,
+):
+    # Apply the patch limits.
+    s1 = math.sqrt(
+        in_patch_limit
+        / (max(1.0, width // patch_size) * max(1.0, height // patch_size))
+    )
+    s2 = patch_limit_on_one_side * patch_size / width
+    s3 = patch_limit_on_one_side * patch_size / height
+    scale = min(1.0, s1, s2, s3)
+    new_w, new_h = max(1, int(width * scale)), max(1, int(height * scale))
+    new_w = min(new_w, patch_limit_on_one_side * patch_size)
+    new_h = min(new_h, patch_limit_on_one_side * patch_size)
+
+    factor = merge_kernel_size * patch_size
+
+    pad_height = (factor - new_h % factor) % factor
+    pad_width = (factor - new_w % factor) % factor
+
+    if fixed_output_tokens is not None:
+        num_tokens = fixed_output_tokens
+    else:
+        # Calculate new dimensions after padding and patching
+        token_height = (new_h + pad_height) // factor
+        token_width = (new_w + pad_width) // factor
+
+        assert token_height * merge_kernel_size <= patch_limit_on_one_side, (
+            f"token_height {token_height} * merge_kernel_size {merge_kernel_size} > "
+            f"patch_limit_on_one_side {patch_limit_on_one_side}"
+        )
+        assert token_width * merge_kernel_size <= patch_limit_on_one_side, (
+            f"token_width {token_width} * merge_kernel_size {merge_kernel_size} > "
+            f"patch_limit_on_one_side {patch_limit_on_one_side}"
+        )
+
+        num_tokens = token_height * token_width
+    return {
+        "num_tokens": num_tokens,
+        "new_width": new_w,
+        "new_height": new_h,
+        "pad_width": pad_width,
+        "pad_height": pad_height,
+        "sampled_nframes": 1,
+    }
+
+
+class KimiK3ProcessingInfo(BaseProcessingInfo):
+    """Processing information for the image-only Kimi-K3 model.
+
+    K3 uses the standard ``image`` modality (unlike K2.5's unified
+    ``vision_chunk``), so it builds its own ``KimiK3Processor`` wrapper around
+    the checkpoint's image processor and resolves the ``<|media_pad|>`` token
+    id the same way K2.5 does.
+    """
+
+    def __init__(self, ctx: InputProcessingContext) -> None:
+        super().__init__(ctx)
+
+        self.hf_config = hf_config = self.get_hf_config()
+
+        tokenizer = self.get_tokenizer()
+        image_processor = cached_get_image_processor(
+            self.ctx.model_config.model,
+            revision=self.ctx.model_config.revision,
+            trust_remote_code=self.ctx.model_config.trust_remote_code,
+        )
+
+        # Resolve token ID from the tokenizer because transformers v5
+        # may remap token IDs vs config.json.
+        config_token_id = hf_config.media_placeholder_token_id
+        resolved_token_id = tokenizer.convert_tokens_to_ids("<|media_pad|>")
+        unk_token_id = getattr(tokenizer, "unk_token_id", None)
+        is_valid_resolved = isinstance(resolved_token_id, int) and (
+            unk_token_id is None or resolved_token_id != unk_token_id
+        )
+        if is_valid_resolved and resolved_token_id != config_token_id:
+            logger.warning_once(
+                "Kimi-K3 config.media_placeholder_token_id (%d) disagrees "
+                "with tokenizer mapping for <|media_pad|> (%d). "
+                "Using tokenizer value.",
+                config_token_id,
+                resolved_token_id,
+            )
+            media_token_id = resolved_token_id
+            # Patch config so downstream code also sees the correct ID.
+            hf_config.media_placeholder_token_id = resolved_token_id
+        else:
+            media_token_id = config_token_id
+
+        self.media_token_id = media_token_id
+        self.media_token = tokenizer.decode(media_token_id)
+
+        self.image_processor = image_processor
+        self.hf_processor = KimiK3Processor(
+            tokenizer=tokenizer,
+            image_processor=image_processor,
+        )
+        self.media_tokens_calculator = image_processor.media_tokens_calculator
+
+    def get_hf_processor(self, **kwargs: object) -> KimiK3Processor:
+        return self.hf_processor
+
+    def get_hf_config(self) -> KimiK3Config:
+        return self.ctx.get_hf_config(KimiK3Config)
+
+    def get_supported_mm_limits(self) -> Mapping[str, int | None]:
+        # None means unlimited
+        return {"image": None}
+
+    @classmethod
+    def get_max_image_size(
+        cls,
+        patch_size: int,
+        merge_kernel_size: int,
+        in_patch_limit: int,
+        patch_limit_on_one_side: int,
+        fixed_output_tokens: int | None,
+    ) -> ImageSize:
+        max_side = patch_limit_on_one_side * patch_size
+        best_score = (-1, -1)
+        best_size = (max_side, max_side)
+
+        for width_patches in range(patch_limit_on_one_side + 1):
+            width = min((width_patches + 1) * patch_size - 1, max_side)
+            for height_patches in range(width_patches, patch_limit_on_one_side + 1):
+                height = min((height_patches + 1) * patch_size - 1, max_side)
+                resize_config = navit_resize_image(
+                    width,
+                    height,
+                    patch_size,
+                    merge_kernel_size,
+                    in_patch_limit,
+                    patch_limit_on_one_side,
+                    fixed_output_tokens,
+                )
+                padded_width = resize_config["new_width"] + resize_config["pad_width"]
+                padded_height = (
+                    resize_config["new_height"] + resize_config["pad_height"]
+                )
+                num_patches = padded_width // patch_size * (padded_height // patch_size)
+                score = (resize_config["num_tokens"], num_patches)
+                if score > best_score:
+                    best_score = score
+                    best_size = (width, height)
+        return ImageSize(width=best_size[0], height=best_size[1])
+
+
+class KimiK3DummyInputsBuilder(BaseDummyInputsBuilder[KimiK3ProcessingInfo]):
+    """Builds image-based dummy inputs for K3 profiling.
+
+    The dummy text is made of ``<|kimi_image_placeholder|>`` tokens — exactly
+    the placeholder that K3's ``_get_prompt_updates`` expands — and the dummy
+    mm data is a plain list of PIL images under the ``image`` key.
+    """
+
+    def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str:
+        num_images = mm_counts.get("image", 0)
+        return self.info.get_hf_config().image_placeholder * num_images
+
+    def get_dummy_mm_data(
+        self,
+        seq_len: int,
+        mm_counts: Mapping[str, int],
+        mm_options: Mapping[str, BaseDummyOptions] | None = None,
+    ) -> MultiModalDataDict:
+        media_proc_cfg = self.info.image_processor.media_proc_cfg
+        max_size = self.info.get_max_image_size(
+            media_proc_cfg["patch_size"],
+            media_proc_cfg["merge_kernel_size"],
+            media_proc_cfg["in_patch_limit"],
+            media_proc_cfg["patch_limit_on_one_side"],
+            media_proc_cfg["fixed_output_tokens"],
+        )
+        num_images = mm_counts.get("image", 0)
+        image_overrides = cast(
+            ImageDummyOptions | None,
+            mm_options.get("image") if mm_options else None,
+        )
+        return {
+            "image": self._get_dummy_images(
+                width=max_size.width,
+                height=max_size.height,
+                num_images=num_images,
+                overrides=image_overrides,
+            )
+        }
+
+
+class KimiK3MultiModalProcessor(BaseMultiModalProcessor[KimiK3ProcessingInfo]):
+    """Image-only multi-modal processor for Kimi-K3."""
+
+    def _call_hf_processor(
+        self,
+        prompt: str,
+        mm_data: Mapping[str, object],
+        mm_kwargs: Mapping[str, object],
+        tok_kwargs: Mapping[str, object],
+    ) -> BatchFeature:
+        # Override so the base always routes through the text+mm path
+        # (`KimiK3Processor.__call__`). Otherwise the mm-only fast path calls
+        # the checkpoint image processor directly with bare PIL images, but it
+        # requires `{"type": "image", "image": PIL}` media dicts that only our
+        # wrapper builds.
+        return super()._call_hf_processor(prompt, mm_data, mm_kwargs, tok_kwargs)
+
+    def _hf_processor_applies_updates(
+        self,
+        prompt_text: str,
+        mm_items: MultiModalDataItems,
+        hf_processor_mm_kwargs: Mapping[str, object],
+        tokenization_kwargs: Mapping[str, object],
+    ) -> bool:
+        return False
+
+    def _get_mm_fields_config(
+        self,
+        hf_inputs: BatchFeature,
+        hf_processor_mm_kwargs: Mapping[str, object],
+    ) -> Mapping[str, MultiModalFieldConfig]:
+        """Slice the flattened patch tensor back into per-image items.
+
+        ``pixel_values`` holds all patches from every image concatenated; each
+        image's patch count is ``prod(grid_thws[i])``. ``grid_thws`` is one
+        ``[N_t, N_h, N_w]`` row per image.
+        """
+        grid_thws = hf_inputs.get("grid_thws", torch.empty((0, 3)))
+        grid_sizes = grid_thws.prod(-1)
+
+        return dict(
+            pixel_values=MultiModalFieldConfig.flat_from_sizes("image", grid_sizes),
+            grid_thws=MultiModalFieldConfig.batched("image", keep_on_cpu=True),
+        )
+
+    def _get_prompt_updates(
+        self,
+        mm_items: MultiModalDataItems,
+        hf_processor_mm_kwargs: Mapping[str, Any],
+        out_mm_kwargs: MultiModalKwargsItems,
+    ) -> Sequence[PromptUpdate]:
+        """Expand each K3 image placeholder into a resolution-aware update.
+
+        K3's prompt carries a single ``<|kimi_image_placeholder|>`` token per
+        image. This replaces that token with
+        ``<|media_begin|>image {w}x{h}<|media_content|>{pads}<|media_end|>``,
+        embedding the per-image resolution in the prompt and marking only the
+        ``<|media_pad|>`` positions as embedding slots (the number of pads is
+        the feature size returned by ``media_tokens_calculator``).
+        """
+        media_token_id = self.info.media_token_id
+        media_token = self.info.media_token
+        image_placeholder = self.info.get_hf_config().image_placeholder
+
+        def get_replacement(item_idx: int) -> PromptUpdateDetails[str]:
+            images = mm_items.get_items("image", ImageProcessorItems)
+            image = images.get(item_idx)
+            if image is None:
+                raise ValueError(f"Missing image data at index {item_idx}")
+
+            # The checkpoint image processor works on media dicts, so wrap the
+            # PIL image before asking it for the token count.
+            num_media_token = self.info.media_tokens_calculator(
+                {"type": "image", "image": image}
+            )
+            pads = media_token * num_media_token
+
+            # NOTE: `width`/`height` are the ORIGINAL upload dimensions, not the
+            # post-preprocess (smart-resized) ones. `image` comes from the
+            # untouched parsed `mm_items`; the checkpoint image processor
+            # (`KimiK3VisionProcessor.preprocess`) only produces new tensors via
+            # `image.resize(...)` and never mutates the stored PIL. This matches
+            # the reference HF processor (`KimiK3Processor.preprocess_medias`),
+            # which also builds the prompt from the original `img.size`. The
+            # resize is reflected only in the pad count above.
+            width, height = images.get_image_size(item_idx)
+            full = (
+                f"<|media_begin|>image {width}x{height}<|media_content|>"
+                f"{pads}<|media_end|>"
+            )
+
+            return PromptUpdateDetails.select_token_id(full, media_token_id)
+
+        return [
+            PromptReplacement(
+                modality="image",
+                target=image_placeholder,
+                replacement=get_replacement,
+            ),
+        ]

vllm/models/kimi_k3/common/mtp.py

@@ -0,0 +1,96 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Fused Kimi-K3 MTP input preparation."""
+
+import torch
+
+from vllm.triton_utils import tl, triton
+
+
+@triton.jit
+def _rms_norm(x, weight, eps, hidden_size: tl.constexpr):
+    x = x.to(tl.float32)
+    variance = tl.sum(x * x, axis=0) / hidden_size
+    return x * tl.rsqrt(variance + eps) * weight.to(tl.float32)
+
+
+@triton.jit
+def _fused_mtp_input_kernel(
+    positions_ptr,
+    inputs_embeds_ptr,
+    previous_hidden_states_ptr,
+    enorm_weight_ptr,
+    hnorm_weight_ptr,
+    output_ptr,
+    eps,
+    inputs_embeds_stride,
+    previous_hidden_states_stride,
+    output_stride,
+    hidden_size: tl.constexpr,
+    block_size: tl.constexpr,
+):
+    token_idx = tl.program_id(0).to(tl.int64)
+    input_idx = tl.program_id(1)
+    offsets = tl.arange(0, block_size)
+    mask = offsets < hidden_size
+
+    if input_idx == 0:
+        position = tl.load(positions_ptr + token_idx)
+        values = tl.load(
+            inputs_embeds_ptr + token_idx * inputs_embeds_stride + offsets,
+            mask=mask & (position != 0),
+            other=0.0,
+        )
+        weight = tl.load(enorm_weight_ptr + offsets, mask=mask, other=0.0)
+    else:
+        values = tl.load(
+            previous_hidden_states_ptr
+            + token_idx * previous_hidden_states_stride
+            + offsets,
+            mask=mask,
+            other=0.0,
+        )
+        weight = tl.load(hnorm_weight_ptr + offsets, mask=mask, other=0.0)
+
+    output = _rms_norm(values, weight, eps, hidden_size)
+    tl.store(
+        output_ptr + token_idx * output_stride + input_idx * hidden_size + offsets,
+        output,
+        mask=mask,
+    )
+
+
+def fused_mtp_input(
+    positions: torch.Tensor,
+    inputs_embeds: torch.Tensor,
+    previous_hidden_states: torch.Tensor,
+    enorm_weight: torch.Tensor,
+    hnorm_weight: torch.Tensor,
+    eps: float,
+) -> torch.Tensor:
+    """Mask and normalize both MTP inputs into the projection layout."""
+    num_tokens, hidden_size = inputs_embeds.shape
+    output = torch.empty(
+        num_tokens,
+        2 * hidden_size,
+        dtype=inputs_embeds.dtype,
+        device=inputs_embeds.device,
+    )
+    if num_tokens == 0:
+        return output
+
+    _fused_mtp_input_kernel[(num_tokens, 2)](
+        positions,
+        inputs_embeds,
+        previous_hidden_states,
+        enorm_weight,
+        hnorm_weight,
+        output,
+        eps,
+        inputs_embeds.stride(0),
+        previous_hidden_states.stride(0),
+        output.stride(0),
+        hidden_size,
+        triton.next_power_of_2(hidden_size),
+    )
+    return output

vllm/models/kimi_k3/nvidia/dspark_mla.py

@@ -0,0 +1,527 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""K3 dense MLA draft model for DSpark speculative decoding."""
+
+from collections.abc import Iterable
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+import vllm._custom_ops as ops
+from vllm.config import VllmConfig
+from vllm.model_executor.layers.layernorm import RMSNorm
+from vllm.model_executor.layers.linear import ReplicatedLinear
+from vllm.model_executor.layers.logits_processor import LogitsProcessor
+from vllm.model_executor.models.qwen3_dspark import DSparkMarkovHead
+from vllm.model_executor.models.utils import (
+    AutoWeightsLoader,
+    WeightsMapper,
+    get_draft_quant_config,
+    maybe_prefix,
+)
+from vllm.models.kimi_k3.nvidia.mla import MultiHeadLatentAttention
+from vllm.models.kimi_k3.nvidia.model import KimiMLP
+from vllm.utils.torch_utils import is_quantized_kv_cache
+
+
+class ReplicatedDSparkMarkovHead(DSparkMarkovHead):
+    """DSpark Markov head with full weights on every TP rank."""
+
+    def __init__(
+        self,
+        vocab_size: int,
+        draft_vocab_size: int,
+        markov_rank: int,
+        prefix: str,
+    ) -> None:
+        # TODO: Remove this mypy workaround once the K3 PR is fully merged.
+        super().__init__(  # type: ignore[call-arg]
+            vocab_size,
+            draft_vocab_size,
+            markov_rank,
+            prefix,
+            replicated=True,
+        )
+
+
+class K3DSparkDecoderLayer(nn.Module):
+    def __init__(
+        self,
+        *,
+        vllm_config: VllmConfig,
+        config,
+        layer_idx: int,
+        start_layer_id: int,
+        prefix: str,
+    ) -> None:
+        super().__init__()
+        quant_config = get_draft_quant_config(vllm_config)
+        self.self_attn = MultiHeadLatentAttention(
+            config=config,
+            hidden_size=config.hidden_size,
+            num_heads=config.num_attention_heads,
+            qk_nope_head_dim=config.qk_nope_head_dim,
+            qk_rope_head_dim=config.qk_rope_head_dim,
+            v_head_dim=config.v_head_dim,
+            q_lora_rank=config.q_lora_rank,
+            kv_lora_rank=config.kv_lora_rank,
+            cache_config=vllm_config.cache_config,
+            quant_config=quant_config,
+            prefix=maybe_prefix(
+                prefix, f"layers.{start_layer_id + layer_idx}.self_attn"
+            ),
+            use_rope=True,
+            non_causal_multi_token_decode=True,
+        )
+        self.mlp = KimiMLP(
+            hidden_size=config.hidden_size,
+            intermediate_size=config.intermediate_size,
+            hidden_act=config.hidden_act,
+            quant_config=quant_config,
+            prefix=maybe_prefix(prefix, f"layers.{start_layer_id + layer_idx}.mlp"),
+        )
+        self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.post_attention_layernorm = RMSNorm(
+            config.hidden_size, eps=config.rms_norm_eps
+        )
+
+    def forward(
+        self,
+        positions: torch.Tensor,
+        hidden_states: torch.Tensor,
+        residual: torch.Tensor | None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        if residual is None:
+            residual = hidden_states
+            hidden_states = self.input_layernorm(hidden_states)
+        else:
+            hidden_states, residual = self.input_layernorm(hidden_states, residual)
+
+        hidden_states = self.self_attn(
+            positions=positions,
+            hidden_states=hidden_states,
+        )
+        hidden_states, residual = self.post_attention_layernorm(hidden_states, residual)
+        hidden_states = self.mlp(hidden_states)
+        return hidden_states, residual
+
+
+class K3DSparkModel(nn.Module):
+    def __init__(
+        self,
+        *,
+        vllm_config: VllmConfig,
+        start_layer_id: int,
+        prefix: str,
+    ) -> None:
+        super().__init__()
+        assert vllm_config.speculative_config is not None
+        self.config = vllm_config.speculative_config.draft_model_config.hf_config
+        self.quant_config = get_draft_quant_config(vllm_config)
+
+        # The frozen target embedding is aliased after the draft checkpoint loads.
+        self.embed_tokens: nn.Module | None = None
+
+        self.context_proj = ReplicatedLinear(
+            self.config.target_hidden_size * self.config.num_target_layers,
+            self.config.hidden_size,
+            bias=False,
+            return_bias=False,
+            quant_config=self.quant_config,
+            prefix=maybe_prefix(prefix, "context_proj"),
+        )
+        self.context_norm = RMSNorm(
+            self.config.hidden_size, eps=self.config.rms_norm_eps
+        )
+
+        self.layers = nn.ModuleList(
+            [
+                K3DSparkDecoderLayer(
+                    vllm_config=vllm_config,
+                    config=self.config,
+                    layer_idx=layer_idx,
+                    start_layer_id=start_layer_id,
+                    prefix=prefix,
+                )
+                for layer_idx in range(self.config.num_hidden_layers)
+            ]
+        )
+        self.final_norm = RMSNorm(self.config.hidden_size, eps=self.config.rms_norm_eps)
+        self.markov_head = ReplicatedDSparkMarkovHead(
+            self.config.vocab_size,
+            self.config.draft_vocab_size,
+            self.config.markov_rank,
+            prefix=maybe_prefix(prefix, "markov_head"),
+        )
+        self._context_kv_fusion_available: bool | None = None
+        self._max_num_context_tokens = (
+            vllm_config.scheduler_config.max_num_batched_tokens
+        )
+
+    def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
+        assert self.embed_tokens is not None
+        return self.embed_tokens(input_ids)
+
+    def combine_hidden_states(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        return self.context_norm(self.context_proj(hidden_states))
+
+    @torch.inference_mode()
+    def precompute_and_store_context_kv(
+        self,
+        context_states: torch.Tensor,
+        context_positions: torch.Tensor,
+        context_slot_mapping: torch.Tensor | list[torch.Tensor | None] | None = None,
+    ) -> None:
+        """Project target-derived context into each draft layer's latent cache."""
+        if self._context_kv_fusion_available is None:
+            self._build_fused_context_kv_buffers()
+        if self._context_kv_fusion_available:
+            self._precompute_fused_context_kv(
+                context_states, context_positions, context_slot_mapping
+            )
+            return
+
+        # Quantized fallback. Directly invoking the projection modules preserves
+        # their quantization methods, at the cost of also computing unused Q rows.
+        for layer_idx, layer in enumerate(self.layers):
+            attn = layer.self_attn
+            assert attn.fused_qkv_a_proj is not None
+            assert attn.q_lora_rank is not None
+            assert attn.rotary_emb is not None
+            qkv_lora = attn.fused_qkv_a_proj(context_states)[0]
+            kv_lora = qkv_lora[..., attn.q_lora_rank :]
+            kv_c, k_pe = kv_lora.split(
+                [attn.kv_lora_rank, attn.qk_rope_head_dim], dim=-1
+            )
+            kv_c = attn.kv_a_layernorm(kv_c)
+            k_pe = k_pe.unsqueeze(1)
+            # DeepSeek YaRN's FlashInfer path requires paired Q/K tensors.
+            # The vLLM CUDA op supports rotating one tensor in place and
+            # consumes the same (possibly scaled fp32) cos/sin cache.
+            rotary_emb = attn.rotary_emb
+            ops.rotary_embedding(
+                context_positions,
+                k_pe,
+                None,
+                rotary_emb.head_size,
+                rotary_emb.cos_sin_cache,
+                rotary_emb.is_neox_style,
+            )
+
+            slot_mapping = (
+                context_slot_mapping[layer_idx]
+                if isinstance(context_slot_mapping, (list, tuple))
+                else context_slot_mapping
+            )
+            if slot_mapping is None:
+                continue
+            attn.impl.do_kv_cache_update(
+                kv_c,
+                k_pe,
+                attn.kv_cache,
+                slot_mapping,
+                attn.kv_cache_dtype,
+                attn._k_scale,
+            )
+
+    def _build_fused_context_kv_buffers(self) -> None:
+        """Build a cross-layer KV-only A projection after checkpoint loading."""
+        if self.quant_config is not None:
+            self._context_kv_fusion_available = False
+            return
+
+        attentions = [layer.self_attn for layer in self.layers]
+        if not attentions or any(
+            attn.fused_qkv_a_proj is None
+            or not hasattr(attn.fused_qkv_a_proj, "weight")
+            for attn in attentions
+        ):
+            self._context_kv_fusion_available = False
+            return
+
+        attn0 = attentions[0]
+        assert attn0.q_lora_rank is not None
+        kv_width = attn0.kv_lora_rank + attn0.qk_rope_head_dim
+        kv_weights = []
+        for attn in attentions:
+            assert attn.q_lora_rank is not None
+            assert (
+                attn.q_lora_rank == attn0.q_lora_rank
+                and attn.kv_lora_rank == attn0.kv_lora_rank
+                and attn.qk_rope_head_dim == attn0.qk_rope_head_dim
+                and attn.kv_a_layernorm.variance_epsilon
+                == attn0.kv_a_layernorm.variance_epsilon
+            ), "All MLA DSpark layers must share their latent KV geometry."
+            kv_weights.append(
+                attn.fused_qkv_a_proj.weight.detach().narrow(
+                    0, attn.q_lora_rank, kv_width
+                )
+            )
+
+        # [L * (kv_lora_rank + rope_dim), hidden_size]. The underlying fused
+        # A weights are replicated (`disable_tp=True`), so this is valid on
+        # every TP rank without communication.
+        self._fused_context_kv_weight = torch.cat(kv_weights, dim=0)
+        self._context_kv_norm_weights = torch.stack(
+            [attn.kv_a_layernorm.weight.detach() for attn in attentions], dim=0
+        ).contiguous()
+        self._num_context_layers = len(attentions)
+        self._context_kv_width = kv_width
+        self._context_kv_lora_rank = attn0.kv_lora_rank
+        self._context_rope_dim = attn0.qk_rope_head_dim
+        self._context_rms_norm_eps = attn0.kv_a_layernorm.variance_epsilon
+        self._context_positions_repeated = torch.empty(
+            self._num_context_layers * self._max_num_context_tokens,
+            dtype=torch.int64,
+            device=self._fused_context_kv_weight.device,
+        )
+        self._context_kv_fusion_available = True
+
+    def _precompute_fused_context_kv(
+        self,
+        context_states: torch.Tensor,
+        context_positions: torch.Tensor,
+        context_slot_mapping: torch.Tensor | list[torch.Tensor | None] | None,
+    ) -> None:
+        num_ctx = context_states.shape[0]
+        num_layers = self._num_context_layers
+
+        # One KV-only GEMM replaces five full Q+KV GEMMs. For K3 this projects
+        # 5*576 rows rather than 5*2112 rows (72.7% fewer A-projection FLOPs).
+        all_kv = F.linear(context_states, self._fused_context_kv_weight)
+        all_kv = all_kv.view(num_ctx, num_layers, self._context_kv_width)
+        all_kv_c = all_kv[..., : self._context_kv_lora_rank]
+        all_k_pe = all_kv[..., self._context_kv_lora_rank :]
+
+        # Layer-major layout lets the 2-D RMSNorm weights select a distinct row
+        # for each draft layer in one grouped kernel.
+        all_kv_c = all_kv_c.permute(1, 0, 2).contiguous()
+        all_kv_c_normed = torch.empty_like(all_kv_c)
+        ops.rms_norm(
+            all_kv_c_normed,
+            all_kv_c,
+            self._context_kv_norm_weights,
+            self._context_rms_norm_eps,
+        )
+
+        all_k_pe = all_k_pe.permute(1, 0, 2).contiguous()
+        all_k_pe_flat = all_k_pe.view(num_layers * num_ctx, 1, self._context_rope_dim)
+        repeated_positions = self._context_positions_repeated[: num_layers * num_ctx]
+        repeated_positions.view(num_layers, num_ctx).copy_(context_positions)
+        # Keep the single-tensor context RoPE on vLLM's optimized CUDA op;
+        # DeepSeek YaRN's FlashInfer wrapper assumes a non-null key tensor.
+        rotary_emb = self.layers[0].self_attn.rotary_emb
+        assert rotary_emb is not None
+        ops.rotary_embedding(
+            repeated_positions,
+            all_k_pe_flat,
+            None,
+            rotary_emb.head_size,
+            rotary_emb.cos_sin_cache,
+            rotary_emb.is_neox_style,
+        )
+        all_k_pe = all_k_pe_flat.view(num_layers, num_ctx, 1, self._context_rope_dim)
+
+        if context_slot_mapping is None:
+            return
+
+        cache_layers = [layer.self_attn for layer in self.layers]
+        if (
+            not is_quantized_kv_cache(cache_layers[0].kv_cache_dtype)
+            and self._has_uniform_block_layout(cache_layers)
+            and (
+                isinstance(context_slot_mapping, torch.Tensor)
+                or all(s is not None for s in context_slot_mapping)
+            )
+        ):
+            # Grouped context KV insert only supports unquantized (bf16) KV cache
+            # and assumes that all layers share the same block layout.
+
+            if isinstance(context_slot_mapping, (list, tuple)):
+                per_layer_slot_mappings = [
+                    s for s in context_slot_mapping if s is not None
+                ]
+                if len({s.data_ptr() for s in per_layer_slot_mappings}) == 1:
+                    # All rows alias to the same slot mapping.
+                    slot_mapping = (
+                        per_layer_slot_mappings[0].unsqueeze(0).expand(num_layers, -1)
+                    )
+                else:
+                    slot_mapping = torch.stack(per_layer_slot_mappings, dim=0)
+            else:
+                # Broadcast the single shared context_slot_mapping tensor.
+                slot_mapping = context_slot_mapping.unsqueeze(0).expand(num_layers, -1)
+
+            ref_cache = cache_layers[0].kv_cache
+            ops.concat_and_cache_mla_grouped(
+                all_kv_c_normed,
+                all_k_pe.squeeze(2),
+                self._get_context_kv_cache_ptrs(cache_layers),
+                slot_mapping,
+                ref_cache.size(1),
+                ref_cache.stride(0),
+                ref_cache.stride(1),
+            )
+            return
+
+        for layer_idx, layer in enumerate(self.layers):
+            slot_mapping = (
+                context_slot_mapping[layer_idx]
+                if isinstance(context_slot_mapping, (list, tuple))
+                else context_slot_mapping
+            )
+            if slot_mapping is None:
+                continue
+            attn = layer.self_attn
+            attn.impl.do_kv_cache_update(
+                all_kv_c_normed[layer_idx],
+                all_k_pe[layer_idx],
+                attn.kv_cache,
+                slot_mapping,
+                attn.kv_cache_dtype,
+                attn._k_scale,
+            )
+
+    def _has_uniform_block_layout(
+        self,
+        cache_layers: list[MultiHeadLatentAttention],
+    ) -> bool:
+        if not hasattr(self, "_layers_share_kv_block_layout"):
+            ref_cache = cache_layers[0].kv_cache
+            self._layers_share_kv_block_layout = all(
+                cl.kv_cache.size(1) == ref_cache.size(1)
+                and cl.kv_cache.stride(0) == ref_cache.stride(0)
+                and cl.kv_cache.stride(1) == ref_cache.stride(1)
+                for cl in cache_layers
+            )
+        return self._layers_share_kv_block_layout
+
+    def _get_context_kv_cache_ptrs(
+        self,
+        cache_layers: list[MultiHeadLatentAttention],
+    ) -> torch.Tensor:
+        # The per-layer KV cache base pointers are stable after allocation, so
+        # build the pointer array once and return it on every call.
+        if not hasattr(self, "_context_cache_ptrs"):
+            ref_cache = cache_layers[0].kv_cache
+            cache_ptrs = torch.tensor(
+                [cl.kv_cache.data_ptr() for cl in cache_layers],
+                dtype=torch.int64,
+                device=ref_cache.device,
+            )
+            self._context_cache_ptrs = cache_ptrs
+        return self._context_cache_ptrs
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        positions: torch.Tensor,
+        inputs_embeds: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        if inputs_embeds is None:
+            inputs_embeds = self.embed_input_ids(input_ids)
+
+        hidden_states = inputs_embeds
+        residual = None
+        for layer in self.layers:
+            hidden_states, residual = layer(
+                positions=positions,
+                hidden_states=hidden_states,
+                residual=residual,
+            )
+        hidden_states, _ = self.final_norm(hidden_states, residual)
+        return hidden_states
+
+
+class K3DSparkForCausalLM(nn.Module):
+    has_own_embed_tokens = False
+    has_own_lm_head = False
+    draft_id_to_target_id = None
+    checkpoint_skip_substrs = ("confidence_head", "embed_tokens", "lm_head")
+
+    hf_to_vllm_mapper = WeightsMapper(
+        orig_to_new_prefix={"": "model."},
+        orig_to_new_stacked={
+            ".gate_proj": (".gate_up_proj", 0),
+            ".up_proj": (".gate_up_proj", 1),
+            ".q_a_proj": (".fused_qkv_a_proj", 0),
+            ".kv_a_proj_with_mqa": (".fused_qkv_a_proj", 1),
+        },
+    )
+
+    def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
+        super().__init__()
+        assert vllm_config.speculative_config is not None
+        self.draft_model_config = vllm_config.speculative_config.draft_model_config
+        self.config = self.draft_model_config.hf_config
+        target_layer_num = vllm_config.model_config.get_num_layers(
+            vllm_config.parallel_config
+        )
+        self.model = K3DSparkModel(
+            vllm_config=vllm_config,
+            start_layer_id=target_layer_num,
+            prefix=maybe_prefix(prefix, "model"),
+        )
+
+        # Assigned by load_dspark_model from the target. Keeping no placeholder
+        # avoids a transient full-vocabulary allocation for this 163k-vocab model.
+        self.lm_head: nn.Module | None = None
+        logit_scale = getattr(self.config, "logit_scale", 1.0)
+        self.logits_processor = LogitsProcessor(
+            self.config.draft_vocab_size, scale=logit_scale
+        )
+
+    def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
+        return self.model.embed_input_ids(input_ids)
+
+    def combine_hidden_states(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        return self.model.combine_hidden_states(hidden_states)
+
+    def get_draft_kv_cache_layer_names(self) -> list[str]:
+        return [layer.self_attn.layer_name for layer in self.model.layers]
+
+    def precompute_and_store_context_kv(
+        self,
+        context_states: torch.Tensor,
+        context_positions: torch.Tensor,
+        context_slot_mapping: torch.Tensor | list[torch.Tensor | None] | None = None,
+    ) -> None:
+        self.model.precompute_and_store_context_kv(
+            context_states, context_positions, context_slot_mapping
+        )
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        positions: torch.Tensor,
+        inputs_embeds: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        return self.model(input_ids, positions, inputs_embeds)
+
+    def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        assert self.lm_head is not None
+        return self.logits_processor(self.lm_head, hidden_states)
+
+    def compute_draft_logits(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        return self.compute_logits(hidden_states)
+
+    def map_draft_to_target(self, draft_ids: torch.Tensor) -> torch.Tensor:
+        return draft_ids
+
+    def markov_embed(self, token_ids: torch.Tensor) -> torch.Tensor:
+        return self.model.markov_head.embed(token_ids)
+
+    def markov_bias(self, markov_embed: torch.Tensor) -> torch.Tensor:
+        return self.model.markov_head.bias(markov_embed, self.logits_processor)
+
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+        # confidence_head is training-only. The frozen target embedding and LM
+        # head are shared after this draft-specific checkpoint is loaded.
+        loader = AutoWeightsLoader(
+            self,
+            skip_substrs=list(self.checkpoint_skip_substrs),
+        )
+        loaded_weights = loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
+        self.model._build_fused_context_kv_buffers()
+        return loaded_weights

vllm/models/kimi_k3/nvidia/kda.py

@@ -0,0 +1,775 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from collections.abc import Callable
+
+import torch
+from einops import rearrange
+from torch import nn
+from torch.nn.parameter import Parameter
+
+from vllm import _custom_ops as ops
+from vllm.compilation.breakable_cudagraph import eager_break_during_capture
+from vllm.config import VllmConfig
+from vllm.distributed import divide, get_tensor_model_parallel_rank
+from vllm.forward_context import get_forward_context
+from vllm.logger import init_logger
+from vllm.model_executor.layers.linear import (
+    ColumnParallelLinear,
+    MergedColumnParallelLinear,
+    RowParallelLinear,
+)
+from vllm.model_executor.layers.mamba.gdn.base import GatedDeltaNetAttention
+from vllm.model_executor.layers.mamba.mamba_utils import (
+    MambaStateDtypeCalculator,
+    MambaStateShapeCalculator,
+    is_conv_state_dim_first,
+)
+from vllm.model_executor.layers.mamba.ops.causal_conv1d import (
+    causal_conv1d_fn,
+    causal_conv1d_update,
+)
+from vllm.model_executor.layers.mamba.ops.gather_initial_states import (
+    gather_initial_states,
+)
+from vllm.model_executor.model_loader.weight_utils import (
+    default_weight_loader,
+    sharded_weight_loader,
+)
+from vllm.model_executor.parameter import BasevLLMParameter
+from vllm.model_executor.utils import set_weight_attrs
+from vllm.models.kimi_k3.nvidia.kda_metadata import (
+    KimiK3KDAAttentionBackend,
+    KimiK3KDAMetadata,
+)
+from vllm.platforms import current_platform
+from vllm.third_party.flash_linear_attention.ops.kda import FusedRMSNormGated
+from vllm.transformers_utils.configs.kimi_linear import KimiLinearConfig
+from vllm.v1.attention.backend import AttentionBackend
+
+logger = init_logger(__name__)
+
+_KDA_GATE_LOGBOUND_MIN = -5.0
+
+
+def a_log_weight_loader(
+    shard_axis: int,
+) -> Callable[[torch.Tensor, torch.Tensor], None]:
+    """Load KDA A_log stored as either old 4D or current 1D weights."""
+
+    def loader(param: torch.Tensor, loaded_weight: torch.Tensor) -> None:
+        tp_rank = get_tensor_model_parallel_rank()
+        shard_size = param.data.shape[shard_axis]
+        start_idx = tp_rank * shard_size
+
+        if loaded_weight.dim() == 4:
+            assert loaded_weight.shape[:2] == (1, 1), (
+                f"Expected old A_log shape (1, 1, H, 1), got {loaded_weight.shape}"
+            )
+            assert loaded_weight.shape[-1] == 1, (
+                f"Expected old A_log last dim to be 1, got {loaded_weight.shape}"
+            )
+            loaded_weight = loaded_weight.view(loaded_weight.shape[2])
+
+        loaded_weight = loaded_weight.narrow(shard_axis, start_idx, shard_size)
+        return default_weight_loader(param, loaded_weight)
+
+    return loader
+
+
+class _KimiGDNMergedColumnParallelLinear(MergedColumnParallelLinear):
+    """Merged projection with one output replicated across TP ranks."""
+
+    def __init__(
+        self,
+        input_size: int,
+        output_sizes: list[int],
+        replicated_shard_id: int,
+        tp_size: int,
+        **kwargs,
+    ) -> None:
+        self.replicated_shard_id = replicated_shard_id
+        output_sizes = output_sizes.copy()
+        output_sizes[replicated_shard_id] *= tp_size
+        super().__init__(input_size, output_sizes, **kwargs)
+
+    def weight_loader(
+        self,
+        param: Parameter,
+        loaded_weight: torch.Tensor,
+        loaded_shard_id: tuple[int, ...] | int | None = None,
+    ) -> None:
+        tp_rank = self.tp_rank
+        param_tp_rank = getattr(param, "tp_rank", None)
+        if loaded_shard_id == self.replicated_shard_id:
+            self.tp_rank = 0
+            if param_tp_rank is not None:
+                param.tp_rank = 0
+        try:
+            super().weight_loader(param, loaded_weight, loaded_shard_id)
+        finally:
+            self.tp_rank = tp_rank
+            if param_tp_rank is not None:
+                param.tp_rank = param_tp_rank
+
+    def weight_loader_v2(
+        self,
+        param: BasevLLMParameter,
+        loaded_weight: torch.Tensor,
+        loaded_shard_id: tuple[int, ...] | int | None = None,
+    ) -> None:
+        tp_rank = self.tp_rank
+        param_tp_rank = getattr(param, "tp_rank", None)
+        if loaded_shard_id == self.replicated_shard_id:
+            self.tp_rank = 0
+            if param_tp_rank is not None:
+                param.tp_rank = 0
+        try:
+            super().weight_loader_v2(param, loaded_weight, loaded_shard_id)
+        finally:
+            self.tp_rank = tp_rank
+            if param_tp_rank is not None:
+                param.tp_rank = param_tp_rank
+
+
+def is_fused_kda_decode_supported(
+    num_heads: int,
+    head_dim: int,
+    conv_width: int,
+    num_spec: int,
+    input_dtype: torch.dtype,
+    conv_state_dtype: torch.dtype,
+) -> bool:
+    if (
+        num_heads not in (12, 24, 48, 96)
+        or head_dim != 128
+        or conv_width != 4
+        or num_spec != 0
+        or input_dtype != torch.bfloat16
+        or conv_state_dtype != torch.bfloat16
+        or is_conv_state_dim_first()
+        or not hasattr(torch.ops._C, "fused_kda_decode")
+    ):
+        return False
+    # SM90 is architecture-specific; SM10x and SM12x use family binaries.
+    return (
+        current_platform.is_device_capability(90)
+        or current_platform.is_device_capability_family(100)
+        or current_platform.is_device_capability_family(120)
+    )
+
+
+def is_flashkda_supported(
+    head_dim: int,
+    dtype: torch.dtype,
+    lower_bound: float | None,
+) -> bool:
+    capability = current_platform.get_device_capability()
+    return (
+        capability is not None
+        and capability.major in (9, 10, 12)
+        and head_dim == 128
+        and dtype == torch.bfloat16
+        and lower_bound is not None
+    )
+
+
+def _flashkda_prefill(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    g: torch.Tensor,
+    beta: torch.Tensor,
+    A_log: torch.Tensor,
+    dt_bias: torch.Tensor,
+    lower_bound: float,
+    initial_state: torch.Tensor,
+    cu_seqlens: torch.Tensor,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    import vllm._flashkda_C  # noqa: F401
+
+    out = torch.empty(v.shape, dtype=v.dtype, device=v.device)
+    final_state = torch.empty_like(initial_state)
+    workspace = torch.empty(
+        torch.ops._flashkda_C.get_workspace_size(
+            q.shape[0] * q.shape[1],
+            q.shape[2],
+            cu_seqlens.numel() - 1,
+        ),
+        dtype=torch.uint8,
+        device=q.device,
+    )
+    # FlashKDA hardcodes dense Q/K/V/G strides. Beta may be row-strided because
+    # FlashKDA materializes its transposed [H, T] layout internally.
+    # TODO: Teach FlashKDA to consume beta in [T, H] layout directly instead
+    # of transposing it to contiguous [H, T] storage internally.
+    torch.ops._flashkda_C.fwd(
+        q.contiguous(),
+        k.contiguous(),
+        v.contiguous(),
+        g.contiguous(),
+        beta,
+        q.shape[-1] ** -0.5,
+        out,
+        workspace,
+        A_log.contiguous(),
+        dt_bias.view(-1, q.shape[-1]).contiguous(),
+        lower_bound,
+        initial_state.contiguous(),
+        final_state,
+        cu_seqlens.contiguous(),
+    )
+    return out, final_state
+
+
+def resolve_kda_prefill_backend(
+    backend: str,
+    head_dim: int,
+    dtype: torch.dtype,
+    lower_bound: float | None,
+) -> str:
+    if backend not in ("auto", "triton", "flashkda"):
+        raise ValueError(f"Unsupported KDA prefill backend: {backend}")
+    supported = is_flashkda_supported(head_dim, dtype, lower_bound)
+    if backend == "flashkda" and not supported:
+        raise RuntimeError(
+            "FlashKDA requires CUDA SM90/SM10x/SM12x, bfloat16, "
+            "head_dim=128, and a bounded KDA gate."
+        )
+    if supported and backend != "triton":
+        logger.info_once("Using FlashKDA KDA prefill backend.")
+        return "flashkda"
+    return "triton"
+
+
+def _make_decode_conv1d_weight_loader(
+    dims: list[int],
+    tp_size: int,
+    tp_rank: int,
+    decode_conv1d_weight: torch.Tensor | None,
+) -> Callable[..., None]:
+    sharded_dims = [dim // tp_size for dim in dims]
+
+    def weight_loader(
+        param: torch.Tensor,
+        loaded_weight: torch.Tensor,
+        loaded_shard_id: int,
+    ) -> None:
+        if loaded_weight.dim() == 2:
+            loaded_weight = loaded_weight.unsqueeze(1)
+        shard_size = sharded_dims[loaded_shard_id]
+        source_start = tp_rank * shard_size
+        target_start = sum(sharded_dims[:loaded_shard_id])
+        loaded_shard = loaded_weight[source_start : source_start + shard_size]
+        param.data[target_start : target_start + shard_size].copy_(loaded_shard)
+        if decode_conv1d_weight is not None and not param.is_meta:
+            decode_conv1d_weight[loaded_shard_id].copy_(
+                loaded_shard.squeeze(1).transpose(0, 1)
+            )
+
+    return weight_loader
+
+
+def _make_decode_norm_weight_loader(
+    decode_norm_weight: torch.Tensor,
+) -> Callable[..., None]:
+    def weight_loader(param: torch.Tensor, loaded_weight: torch.Tensor) -> None:
+        default_weight_loader(param, loaded_weight)
+        if not param.is_meta:
+            decode_norm_weight.copy_(param.data)
+
+    return weight_loader
+
+
+class KimiK3DeltaAttention(GatedDeltaNetAttention):
+    def get_attn_backend(self) -> type[AttentionBackend]:
+        return KimiK3KDAAttentionBackend
+
+    def get_state_dtype(
+        self,
+    ) -> tuple[torch.dtype, torch.dtype]:
+        if self.model_config is None or self.cache_config is None:
+            raise ValueError("model_config and cache_config must be set")
+        return MambaStateDtypeCalculator.kda_state_dtype(
+            self.model_config.dtype, self.cache_config.mamba_cache_dtype
+        )
+
+    def get_state_shape(
+        self,
+    ) -> tuple[tuple[int, ...], tuple[int, ...]]:
+        return MambaStateShapeCalculator.kda_state_shape(
+            self.tp_size,
+            self.num_heads,
+            self.head_dim,
+            conv_kernel_size=self.conv_size,
+            num_spec=self.num_spec,
+        )
+
+    def __init__(
+        self,
+        config: KimiLinearConfig,
+        vllm_config: VllmConfig,
+        prefix: str = "",
+    ) -> None:
+        super().__init__(config, vllm_config, prefix)
+
+        kda_config = config.linear_attn_config  # type: ignore[attr-defined]
+        assert kda_config is not None, "linear_attn_config must be set"
+        self.head_dim = kda_config["head_dim"]
+        self.num_heads = kda_config["num_heads"]
+        assert self.num_heads % self.tp_size == 0
+        self.local_num_heads = divide(self.num_heads, self.tp_size)
+        self.projection_size = self.head_dim * self.num_heads
+        self.local_projection_size = divide(self.projection_size, self.tp_size)
+        self.conv_size = kda_config["short_conv_kernel_size"]
+        assert kda_config.get("use_full_rank_gate", False), (
+            "KimiK3DeltaAttention requires a full-rank gate"
+        )
+
+        # Keep f_a before the narrow beta shard, then pad each TP-local row
+        # to select the aligned BF16 GEMM path.
+        qkvg_output_sizes = [self.projection_size] * 4
+        in_proj_output_sizes = qkvg_output_sizes + [
+            self.head_dim,
+            self.num_heads,
+        ]
+        local_output_size = (
+            4 * self.local_projection_size + self.head_dim + self.local_num_heads
+        )
+        self.in_proj_padding = -local_output_size % 16
+        if self.in_proj_padding:
+            in_proj_output_sizes.append(self.in_proj_padding * self.tp_size)
+        self.in_proj_qkvgfab = _KimiGDNMergedColumnParallelLinear(
+            self.hidden_size,
+            in_proj_output_sizes,
+            replicated_shard_id=4,
+            tp_size=self.tp_size,
+            bias=False,
+            quant_config=self.quant_config,
+            prefix=f"{prefix}.in_proj_qkvgfab",
+        )
+        if self.in_proj_padding:
+            self.in_proj_qkvgfab.weight.data[-self.in_proj_padding :].zero_()
+
+        self.f_b_proj = ColumnParallelLinear(
+            self.head_dim,
+            self.projection_size,
+            bias=False,
+            quant_config=self.quant_config,
+            prefix=f"{prefix}.f_b_proj",
+        )
+        self.dt_bias = nn.Parameter(
+            torch.empty(self.local_projection_size, dtype=torch.float32)
+        )
+        set_weight_attrs(self.dt_bias, {"weight_loader": sharded_weight_loader(0)})
+
+        # One packed parameter and cache let decode run a single conv update.
+        # Prefill slices them back into Q/K/V to obtain dense outputs cheaply.
+        self.conv1d = ColumnParallelLinear(
+            input_size=self.conv_size,
+            output_size=3 * self.projection_size,
+            bias=False,
+            params_dtype=torch.float32,
+            prefix=f"{prefix}.conv1d",
+        )
+        self.conv1d.weight.data = self.conv1d.weight.data.unsqueeze(1)
+        # Keep a width-major copy for fused decode without changing the layout
+        # consumed by the prefill and fallback decode kernels.
+        conv_state_dtype, _ = self.get_state_dtype()
+        decode_conv1d_weight = None
+        if is_fused_kda_decode_supported(
+            self.local_num_heads,
+            self.head_dim,
+            self.conv_size,
+            self.num_spec,
+            vllm_config.model_config.dtype,
+            conv_state_dtype,
+        ):
+            logger.info_once("Fused KDA decode kernel (conv+KDA+norm) is enabled.")
+            decode_conv1d_weight = torch.empty(
+                3,
+                self.conv_size,
+                self.local_projection_size,
+                dtype=self.conv1d.weight.dtype,
+                device=self.conv1d.weight.device,
+            )
+        self.register_buffer(
+            "decode_conv1d_weight", decode_conv1d_weight, persistent=False
+        )
+        delattr(self.conv1d.weight, "weight_loader")
+        set_weight_attrs(
+            self.conv1d.weight,
+            {
+                "weight_loader": _make_decode_conv1d_weight_loader(
+                    [self.projection_size] * 3,
+                    self.tp_size,
+                    self.tp_rank,
+                    decode_conv1d_weight,
+                )
+            },
+        )
+
+        self.A_log = nn.Parameter(
+            torch.empty(self.local_num_heads, dtype=torch.float32)
+        )
+        set_weight_attrs(self.A_log, {"weight_loader": a_log_weight_loader(0)})
+
+        self.gate_lower_bound: float | None = kda_config.get("gate_lower_bound", None)
+        if self.gate_lower_bound is not None:
+            assert _KDA_GATE_LOGBOUND_MIN <= self.gate_lower_bound < 0, (
+                "KDA gate lower bound must be in "
+                f"[{_KDA_GATE_LOGBOUND_MIN}, 0). "
+                f"Got {self.gate_lower_bound}."
+            )
+
+        additional_config = vllm_config.additional_config
+        backend = (
+            additional_config.get("kda_prefill_backend", "auto")
+            if isinstance(additional_config, dict)
+            else "auto"
+        )
+        self.kda_prefill_backend = resolve_kda_prefill_backend(
+            backend,
+            self.head_dim,
+            vllm_config.model_config.dtype,
+            self.gate_lower_bound,
+        )
+
+        self.o_norm = FusedRMSNormGated(self.head_dim, activation="sigmoid")
+        decode_norm_weight = None
+        if decode_conv1d_weight is not None:
+            decode_norm_weight = torch.empty(
+                self.head_dim,
+                dtype=torch.float32,
+                device=self.o_norm.weight.device,
+            )
+        self.register_buffer("decode_norm_weight", decode_norm_weight, persistent=False)
+        if decode_norm_weight is not None:
+            # Upcast once while loading; direct BF16 norm weights slow the
+            # fully fused decode kernel.
+            if hasattr(self.o_norm.weight, "weight_loader"):
+                delattr(self.o_norm.weight, "weight_loader")
+            set_weight_attrs(
+                self.o_norm.weight,
+                {"weight_loader": _make_decode_norm_weight_loader(decode_norm_weight)},
+            )
+        self.o_proj = RowParallelLinear(
+            self.projection_size,
+            self.hidden_size,
+            bias=False,
+            quant_config=self.quant_config,
+            prefix=f"{prefix}.o_proj",
+        )
+
+        compilation_config = vllm_config.compilation_config
+        if prefix in compilation_config.static_forward_context:
+            raise ValueError(f"Duplicate layer name: {prefix}")
+        compilation_config.static_forward_context[prefix] = self
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        positions: torch.Tensor,
+    ) -> torch.Tensor:
+        num_tokens = hidden_states.size(0)
+        projected_qkvgfab = self.in_proj_qkvgfab(hidden_states)[0]
+        split_sizes = [
+            3 * self.local_projection_size,
+            self.local_projection_size,
+            self.head_dim,
+            self.local_num_heads,
+        ]
+        if self.in_proj_padding:
+            split_sizes.append(self.in_proj_padding)
+        projected = projected_qkvgfab.split(split_sizes, dim=-1)
+        mixed_qkv, g_proj_states, f_a, beta = projected[:4]
+
+        g1 = self.f_b_proj(f_a)[0]
+        beta = beta.unsqueeze(0)
+        g1 = rearrange(g1, "n (h d) -> 1 n h d", d=self.head_dim)
+        g2 = rearrange(g_proj_states, "... (h d) -> ... h d", d=self.head_dim)
+        core_attn_out = torch.empty(
+            (1, num_tokens, self.local_num_heads, self.head_dim),
+            dtype=hidden_states.dtype,
+            device=hidden_states.device,
+        )
+        self._forward(
+            mixed_qkv=mixed_qkv,
+            g1=g1,
+            g2=g2,
+            beta=beta,
+            core_attn_out=core_attn_out,
+        )
+        core_attn_out = rearrange(core_attn_out, "1 n h d -> n (h d)")
+        return self.o_proj(core_attn_out)[0]
+
+    @eager_break_during_capture
+    def _forward(
+        self,
+        mixed_qkv: torch.Tensor,
+        g1: torch.Tensor,
+        g2: torch.Tensor,
+        beta: torch.Tensor,
+        core_attn_out: torch.Tensor,
+    ) -> None:
+        forward_context = get_forward_context()
+        attn_metadata_raw = forward_context.attn_metadata
+        if attn_metadata_raw is None:
+            return
+
+        from vllm.models.kimi_k3.nvidia.ops.third_party.kda import (
+            chunk_kda_with_fused_gate,
+            fused_recurrent_kda,
+            fused_recurrent_kda_packed_decode,
+        )
+
+        assert isinstance(attn_metadata_raw, dict)
+        attn_metadata_narrowed = attn_metadata_raw[self.prefix]
+        assert isinstance(attn_metadata_narrowed, KimiK3KDAMetadata)
+        m = attn_metadata_narrowed
+        has_initial_state = m.has_initial_state
+        non_spec_query_start_loc = m.non_spec_query_start_loc
+        non_spec_state_indices_tensor = m.non_spec_state_indices_tensor
+        spec_token_indx = m.spec_token_indx
+        non_spec_token_indx = m.non_spec_token_indx
+        spec_state_indices_tensor = m.spec_state_indices_tensor
+        spec_query_start_loc = m.spec_query_start_loc
+        num_accepted_tokens = m.num_accepted_tokens
+        num_actual_tokens = m.num_actual_tokens
+        has_spec_decode = m.num_spec_decodes > 0
+        mixed_qkv = mixed_qkv[:num_actual_tokens]
+        g1 = g1[:, :num_actual_tokens]
+        beta = beta[:, :num_actual_tokens]
+
+        conv_state, recurrent_state = self.kv_cache
+        # The convolution kernels consume (..., dim, width - 1).
+        if not is_conv_state_dim_first():
+            conv_state = conv_state.transpose(-1, -2)
+
+        if (
+            self.decode_conv1d_weight is not None
+            and self.decode_norm_weight is not None
+            and not has_spec_decode
+            and m.num_prefills == 0
+            and m.num_decodes > 0
+        ):
+            assert non_spec_state_indices_tensor is not None
+            ops.fused_kda_decode(
+                x=mixed_qkv,
+                weight=self.decode_conv1d_weight,
+                bias=self.conv1d.bias,
+                conv_state=conv_state,
+                raw_g=g1,
+                raw_beta=beta,
+                A_log=self.A_log,
+                dt_bias=self.dt_bias,
+                state_indices=non_spec_state_indices_tensor[:num_actual_tokens],
+                state=recurrent_state,
+                out=core_attn_out[:, :num_actual_tokens],
+                lower_bound=self.gate_lower_bound,
+                output_gate=g2[:num_actual_tokens],
+                norm_weight=self.decode_norm_weight,
+                norm_eps=self.o_norm.eps,
+            )
+            return
+
+        conv_weights = self.conv1d.weight.view(
+            self.conv1d.weight.size(0), self.conv1d.weight.size(2)
+        )
+        q_conv_weight, k_conv_weight, v_conv_weight = conv_weights.split(
+            self.local_projection_size, dim=0
+        )
+        q_conv_state, k_conv_state, v_conv_state = conv_state.split(
+            self.local_projection_size, dim=-2
+        )
+
+        # Separate multi-query speculative tokens from prefill/plain decode.
+        if has_spec_decode:
+            if m.num_prefills == 0 and m.num_decodes == 0:
+                mixed_qkv_spec = mixed_qkv
+                g1_spec, beta_spec = g1, beta
+                mixed_qkv_ns = g1_ns = beta_ns = None
+            else:
+                assert spec_token_indx is not None
+                assert non_spec_token_indx is not None
+                mixed_qkv_spec = mixed_qkv.index_select(0, spec_token_indx)
+                g1_spec = g1.index_select(1, spec_token_indx)
+                beta_spec = beta.index_select(1, spec_token_indx)
+                mixed_qkv_ns = mixed_qkv.index_select(0, non_spec_token_indx)
+                g1_ns = g1.index_select(1, non_spec_token_indx)
+                beta_ns = beta.index_select(1, non_spec_token_indx)
+        else:
+            mixed_qkv_spec = g1_spec = beta_spec = None
+            mixed_qkv_ns, g1_ns, beta_ns = mixed_qkv, g1, beta
+
+        # Spec-decode multi-query path.
+        core_attn_out_spec = None
+        if has_spec_decode:
+            assert spec_state_indices_tensor is not None
+            assert spec_query_start_loc is not None
+            spec_conv_indices = spec_state_indices_tensor[:, 0][: m.num_spec_decodes]
+            spec_max_query_len = spec_state_indices_tensor.size(-1)
+            spec_conv_out = torch.empty_like(mixed_qkv_spec)
+            mixed_qkv_spec = causal_conv1d_update(
+                mixed_qkv_spec,
+                conv_state,
+                conv_weights,
+                self.conv1d.bias,
+                activation="silu",
+                conv_state_indices=spec_conv_indices,
+                num_accepted_tokens=num_accepted_tokens,
+                query_start_loc=spec_query_start_loc,
+                max_query_len=spec_max_query_len,
+                validate_data=False,
+                out=spec_conv_out,
+            )
+            q_spec, k_spec, v_spec = (
+                rearrange(x, "n (h d) -> 1 n h d", d=self.head_dim)
+                for x in mixed_qkv_spec.split(self.local_projection_size, dim=-1)
+            )
+            spec_cu_seqlens = spec_query_start_loc[: m.num_spec_decodes + 1]
+            spec_out = (
+                core_attn_out[:, : q_spec.shape[1]]
+                if m.num_prefills == 0 and m.num_decodes == 0
+                else None
+            )
+            core_attn_out_spec, _ = fused_recurrent_kda(
+                q=q_spec,
+                k=k_spec,
+                v=v_spec,
+                raw_g=g1_spec,
+                raw_beta=beta_spec,
+                A_log=self.A_log,
+                dt_bias=self.dt_bias,
+                lower_bound=self.gate_lower_bound,
+                initial_state=recurrent_state,
+                cu_seqlens=spec_cu_seqlens,
+                ssm_state_indices=spec_state_indices_tensor,
+                num_accepted_tokens=num_accepted_tokens,
+                out=spec_out,
+            )
+
+        # Prefill or plain-decode path.
+        core_attn_out_non_spec = None
+        if mixed_qkv_ns is not None:
+            assert g1_ns is not None and beta_ns is not None
+            if m.num_prefills > 0:
+                q_ns, k_ns, v_ns = mixed_qkv_ns.split(
+                    self.local_projection_size, dim=-1
+                )
+
+                # Separate convolution calls accept row-strided packed inputs
+                # and produce dense Q/K/V without an additional V copy.
+                def _prefill_conv(
+                    x: torch.Tensor,
+                    state: torch.Tensor,
+                    weight: torch.Tensor,
+                ) -> torch.Tensor:
+                    return causal_conv1d_fn(
+                        x.transpose(0, 1),
+                        weight,
+                        None,
+                        activation="silu",
+                        conv_states=state,
+                        has_initial_state=has_initial_state,
+                        cache_indices=non_spec_state_indices_tensor,
+                        query_start_loc=non_spec_query_start_loc,
+                        metadata=m,
+                    ).transpose(0, 1)
+
+                q_ns = _prefill_conv(q_ns, q_conv_state, q_conv_weight)
+                k_ns = _prefill_conv(k_ns, k_conv_state, k_conv_weight)
+                v_ns = _prefill_conv(v_ns, v_conv_state, v_conv_weight)
+                q_ns, k_ns, v_ns = (
+                    rearrange(x, "n (h d) -> 1 n h d", d=self.head_dim)
+                    for x in (q_ns, k_ns, v_ns)
+                )
+
+                assert non_spec_state_indices_tensor is not None
+                assert has_initial_state is not None
+                initial_state = gather_initial_states(
+                    recurrent_state,
+                    non_spec_state_indices_tensor,
+                    has_initial_state,
+                )
+                if self.kda_prefill_backend == "flashkda":
+                    assert self.gate_lower_bound is not None
+                    (
+                        core_attn_out_non_spec,
+                        last_recurrent_state,
+                    ) = _flashkda_prefill(
+                        q=q_ns,
+                        k=k_ns,
+                        v=v_ns,
+                        g=g1_ns,
+                        beta=beta_ns,
+                        A_log=self.A_log,
+                        dt_bias=self.dt_bias,
+                        lower_bound=self.gate_lower_bound,
+                        initial_state=initial_state,
+                        cu_seqlens=non_spec_query_start_loc,
+                    )
+                else:
+                    (
+                        core_attn_out_non_spec,
+                        last_recurrent_state,
+                    ) = chunk_kda_with_fused_gate(
+                        q=q_ns,
+                        k=k_ns,
+                        v=v_ns,
+                        raw_g=g1_ns,
+                        raw_beta=beta_ns,
+                        A_log=self.A_log,
+                        g_bias=self.dt_bias,
+                        lower_bound=self.gate_lower_bound,
+                        initial_state=initial_state,
+                        output_final_state=True,
+                        use_qk_l2norm_in_kernel=True,
+                        cu_seqlens=non_spec_query_start_loc,
+                    )
+                recurrent_state[non_spec_state_indices_tensor] = last_recurrent_state
+            else:
+                # Pure non-speculative decode.
+                assert non_spec_state_indices_tensor is not None
+                decode_conv_indices = non_spec_state_indices_tensor[
+                    : mixed_qkv_ns.size(0)
+                ]
+                packed_conv_out = torch.empty_like(mixed_qkv_ns)
+                mixed_qkv_ns = causal_conv1d_update(
+                    mixed_qkv_ns,
+                    conv_state,
+                    conv_weights,
+                    self.conv1d.bias,
+                    activation="silu",
+                    conv_state_indices=decode_conv_indices,
+                    validate_data=True,
+                    out=packed_conv_out,
+                )
+                (
+                    core_attn_out_non_spec,
+                    _,
+                ) = fused_recurrent_kda_packed_decode(
+                    mixed_qkv=mixed_qkv_ns,
+                    raw_g=g1_ns,
+                    raw_beta=beta_ns,
+                    A_log=self.A_log,
+                    dt_bias=self.dt_bias,
+                    lower_bound=self.gate_lower_bound,
+                    initial_state=recurrent_state,
+                    state_indices=decode_conv_indices,
+                )
+
+        # Restore the scheduler's original token order for mixed batches.
+        if core_attn_out_spec is not None and core_attn_out_non_spec is not None:
+            core_attn_out.index_copy_(1, spec_token_indx, core_attn_out_spec)
+            core_attn_out.index_copy_(1, non_spec_token_indx, core_attn_out_non_spec)
+        elif core_attn_out_non_spec is not None:
+            # TODO: prefill and decode kernels write directly to core_attn_out
+            core_attn_out[0, :num_actual_tokens] = core_attn_out_non_spec[
+                0, :num_actual_tokens
+            ]
+        else:
+            assert core_attn_out_spec is not None
+        # Triton normalizes in place, so this is a self-copy with no device
+        # work. Keep it for the out-of-place native implementation.
+        core_attn_out.copy_(self.o_norm(core_attn_out, g2))

vllm/models/kimi_k3/nvidia/kda_metadata.py

@@ -0,0 +1,496 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Kimi-K3 specialization of GDN attention metadata.
+
+The request classification and cudagraph staging intentionally mirror
+``GDNAttentionMetadataBuilder``. Kimi-K3 builds the metadata required by its
+prefill KDA kernel internally, so this builder omits the shared FLA chunk
+metadata construction.
+"""
+
+from dataclasses import dataclass
+from functools import cache
+
+import torch
+
+from vllm.platforms import current_platform
+from vllm.triton_utils import tl, triton
+from vllm.utils.torch_utils import async_tensor_h2d
+from vllm.v1.attention.backend import CommonAttentionMetadata
+from vllm.v1.attention.backends.gdn_attn import (
+    GDNAttentionBackend,
+    GDNAttentionMetadata,
+    GDNAttentionMetadataBuilder,
+)
+from vllm.v1.attention.backends.utils import (
+    NULL_BLOCK_ID,
+    compute_causal_conv1d_metadata,
+    split_decodes_and_prefills,
+)
+from vllm.v1.kv_cache_interface import MambaSpec
+
+
+@cache
+def _metadata_launch_pdl() -> bool:
+    return current_platform.is_arch_support_pdl()
+
+
+@triton.jit(do_not_specialize=["num_requests"])
+def _get_aligned_state_indices_kernel(
+    block_table_ptr,
+    seq_lens_ptr,
+    state_indices_ptr,
+    block_table_stride_0: tl.constexpr,
+    block_table_stride_1: tl.constexpr,
+    seq_lens_stride: tl.constexpr,
+    state_indices_stride_0: tl.constexpr,
+    state_indices_stride_1: tl.constexpr,
+    num_requests,
+    CACHE_BLOCK_SIZE: tl.constexpr,
+    NUM_STATE_SLOTS: tl.constexpr,
+    BLOCK_STATE_SLOTS: tl.constexpr,
+    BLOCK_ROWS: tl.constexpr,
+    launch_pdl: tl.constexpr,
+):
+    if launch_pdl:
+        tl.extra.cuda.gdc_wait()
+        tl.extra.cuda.gdc_launch_dependents()
+
+    rows = tl.program_id(0) * BLOCK_ROWS + tl.arange(0, BLOCK_ROWS)
+    valid_row = rows < num_requests
+    seq_lens = tl.load(
+        seq_lens_ptr + rows * seq_lens_stride,
+        mask=valid_row,
+        other=1,
+    )
+    # Triton truncates signed division toward zero, unlike PyTorch floor
+    # division. Clamping makes both semantics equivalent for seq_lens <= 0.
+    first_state_slot = tl.maximum((seq_lens - 1) // CACHE_BLOCK_SIZE, 0)
+
+    state_slots = tl.arange(0, BLOCK_STATE_SLOTS)
+    valid_state_slot = state_slots < NUM_STATE_SLOTS
+    state_indices = tl.load(
+        block_table_ptr
+        + rows[:, None] * block_table_stride_0
+        + (first_state_slot[:, None] + state_slots[None, :]) * block_table_stride_1,
+        mask=valid_row[:, None] & valid_state_slot[None, :],
+    )
+    tl.store(
+        state_indices_ptr
+        + rows[:, None] * state_indices_stride_0
+        + state_slots[None, :] * state_indices_stride_1,
+        state_indices,
+        mask=valid_row[:, None] & valid_state_slot[None, :],
+    )
+
+
+def _mamba_get_block_table_tensor(
+    block_table: torch.Tensor,
+    seq_lens: torch.Tensor,
+    kv_cache_spec: MambaSpec,
+    mamba_cache_mode: str,
+) -> torch.Tensor:
+    if mamba_cache_mode in ("all", "none"):
+        return block_table
+
+    assert block_table.is_cuda and seq_lens.is_cuda
+    num_requests = block_table.shape[0]
+    num_state_slots = 1 + kv_cache_spec.num_speculative_blocks
+    state_indices = torch.empty(
+        (num_requests, num_state_slots),
+        dtype=block_table.dtype,
+        device=block_table.device,
+    )
+    BLOCK_ROWS = 32
+    grid = (triton.cdiv(num_requests, BLOCK_ROWS),)
+    _get_aligned_state_indices_kernel[grid](
+        block_table,
+        seq_lens,
+        state_indices,
+        block_table.stride(0),
+        block_table.stride(1),
+        seq_lens.stride(0),
+        state_indices.stride(0),
+        state_indices.stride(1),
+        num_requests,
+        CACHE_BLOCK_SIZE=kv_cache_spec.block_size,
+        NUM_STATE_SLOTS=num_state_slots,
+        BLOCK_STATE_SLOTS=triton.next_power_of_2(num_state_slots),
+        BLOCK_ROWS=BLOCK_ROWS,
+        num_warps=1,
+        launch_pdl=_metadata_launch_pdl(),
+    )
+    return state_indices
+
+
+@triton.jit(do_not_specialize=["num_spec_decodes", "batch_size"])
+def _stage_spec_decode_metadata_kernel(
+    state_indices_ptr,
+    query_start_loc_ptr,
+    num_accepted_tokens_ptr,
+    staged_state_indices_ptr,
+    staged_query_start_loc_ptr,
+    staged_num_accepted_tokens_ptr,
+    state_indices_stride_0: tl.constexpr,
+    state_indices_stride_1: tl.constexpr,
+    staged_state_indices_stride_0: tl.constexpr,
+    staged_state_indices_stride_1: tl.constexpr,
+    num_spec_decodes,
+    batch_size,
+    NUM_STATE_SLOTS: tl.constexpr,
+    BLOCK_STATE_SLOTS: tl.constexpr,
+    NULL_STATE_ID: tl.constexpr,
+    BLOCK_ROWS: tl.constexpr,
+    launch_pdl: tl.constexpr,
+):
+    if launch_pdl:
+        tl.extra.cuda.gdc_wait()
+        tl.extra.cuda.gdc_launch_dependents()
+
+    rows = tl.program_id(0) * BLOCK_ROWS + tl.arange(0, BLOCK_ROWS)
+    real_request = rows < num_spec_decodes
+
+    state_slots = tl.arange(0, BLOCK_STATE_SLOTS)
+    valid_state_slot = state_slots < NUM_STATE_SLOTS
+    state_indices = tl.load(
+        state_indices_ptr
+        + rows[:, None] * state_indices_stride_0
+        + state_slots[None, :] * state_indices_stride_1,
+        mask=real_request[:, None] & valid_state_slot[None, :],
+        other=NULL_STATE_ID,
+    )
+    tl.store(
+        staged_state_indices_ptr
+        + rows[:, None] * staged_state_indices_stride_0
+        + state_slots[None, :] * staged_state_indices_stride_1,
+        state_indices,
+        mask=(rows < batch_size)[:, None] & valid_state_slot[None, :],
+    )
+
+    query_row = tl.minimum(rows, num_spec_decodes)
+    query_start_loc = tl.load(
+        query_start_loc_ptr + query_row,
+        mask=rows <= batch_size,
+    )
+    tl.store(
+        staged_query_start_loc_ptr + rows,
+        query_start_loc,
+        mask=rows <= batch_size,
+    )
+
+    num_accepted_tokens = tl.load(
+        num_accepted_tokens_ptr + rows,
+        mask=real_request,
+        other=1,
+    )
+    tl.store(
+        staged_num_accepted_tokens_ptr + rows,
+        num_accepted_tokens,
+        mask=rows < batch_size,
+    )
+
+
+def stage_spec_decode_metadata(
+    state_indices: torch.Tensor,
+    query_start_loc: torch.Tensor,
+    num_accepted_tokens: torch.Tensor,
+    staged_state_indices: torch.Tensor,
+    staged_query_start_loc: torch.Tensor,
+    staged_num_accepted_tokens: torch.Tensor,
+    *,
+    num_spec_decodes: int,
+) -> None:
+    """Stage speculative-decode metadata into CUDA-graph buffers."""
+    assert state_indices.is_cuda
+    assert state_indices.ndim == 2
+    batch_size, num_state_slots = staged_state_indices.shape
+    BLOCK_ROWS = 32
+    grid = (triton.cdiv(batch_size + 1, BLOCK_ROWS),)
+    _stage_spec_decode_metadata_kernel[grid](
+        state_indices,
+        query_start_loc,
+        num_accepted_tokens,
+        staged_state_indices,
+        staged_query_start_loc,
+        staged_num_accepted_tokens,
+        state_indices.stride(0),
+        state_indices.stride(1),
+        staged_state_indices.stride(0),
+        staged_state_indices.stride(1),
+        num_spec_decodes,
+        batch_size,
+        NUM_STATE_SLOTS=num_state_slots,
+        BLOCK_STATE_SLOTS=triton.next_power_of_2(num_state_slots),
+        NULL_STATE_ID=NULL_BLOCK_ID,
+        BLOCK_ROWS=BLOCK_ROWS,
+        num_warps=1,
+        launch_pdl=_metadata_launch_pdl(),
+    )
+
+
+@dataclass
+class KimiK3KDAMetadata(GDNAttentionMetadata):
+    pass
+
+
+class KimiK3KDAMetadataBuilder(GDNAttentionMetadataBuilder):
+    def build(  # type: ignore[override]
+        self,
+        common_prefix_len: int,
+        common_attn_metadata: CommonAttentionMetadata,
+        num_accepted_tokens: torch.Tensor | None = None,
+        num_decode_draft_tokens_cpu: torch.Tensor | None = None,
+        fast_build: bool = False,
+    ) -> KimiK3KDAMetadata:
+        m = common_attn_metadata
+        query_start_loc = m.query_start_loc
+        query_start_loc_cpu = m.query_start_loc_cpu
+        assert isinstance(self.kv_cache_spec, MambaSpec)
+        # Equivalent PyTorch "align" path:
+        #   start = ((seq_lens - 1) // block_size).clamp_(min=0)
+        #   offsets = torch.arange(1 + num_speculative_blocks, dtype=torch.int32)
+        #   indices = (start[:, None] + offsets).to(torch.int64)
+        #   block_table_tensor = torch.gather(block_table, 1, indices)
+        block_table_tensor = _mamba_get_block_table_tensor(
+            m.block_table_tensor,
+            m.seq_lens,
+            self.kv_cache_spec,
+            self.vllm_config.cache_config.mamba_cache_mode,
+        )
+
+        if not self.use_spec_decode or num_decode_draft_tokens_cpu is None:
+            spec_sequence_masks_cpu = None
+            num_spec_decodes = 0
+        else:
+            spec_sequence_masks_cpu = num_decode_draft_tokens_cpu >= 0
+            # A nonnegative entry identifies a spec request. If no draft token
+            # was scheduled, process the whole batch as non-spec instead.
+            if num_decode_draft_tokens_cpu[spec_sequence_masks_cpu].sum().item() == 0:
+                spec_sequence_masks_cpu = None
+                num_spec_decodes = 0
+            else:
+                num_spec_decodes = spec_sequence_masks_cpu.sum().item()
+
+        if num_spec_decodes == 0:
+            # The runner orders ordinary decodes before prefills.
+            num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens = (
+                split_decodes_and_prefills(m, decode_threshold=1)
+            )
+            num_spec_decode_tokens = 0
+            spec_token_indx = None
+            non_spec_token_indx = None
+            spec_state_indices_tensor = None
+            non_spec_state_indices_tensor = block_table_tensor[:, 0]
+            spec_query_start_loc = None
+            non_spec_query_start_loc = query_start_loc
+            non_spec_query_start_loc_cpu = query_start_loc_cpu
+            num_accepted_tokens = None
+        else:
+            assert spec_sequence_masks_cpu is not None
+            assert num_accepted_tokens is not None
+            query_lens_cpu = query_start_loc_cpu.diff()
+            num_query_tokens = query_start_loc_cpu[-1].item()
+
+            # Exclude zero-length cudagraph padding from request-indexed
+            # non-spec metadata.
+            active_non_spec_mask_cpu = (~spec_sequence_masks_cpu) & (query_lens_cpu > 0)
+            non_spec_query_lens_cpu = query_lens_cpu[active_non_spec_mask_cpu]
+            num_non_spec_requests = non_spec_query_lens_cpu.numel()
+            num_non_spec_tokens = non_spec_query_lens_cpu.sum().item()
+
+            # Query length alone cannot distinguish a true decode from a
+            # one-token prefill chunk. Packed decode is safe only when every
+            # active non-spec request is a true, one-token decode.
+            assert m.is_prefilling is not None
+            assert m.is_prefilling.device.type == "cpu"
+            non_spec_is_prefilling = m.is_prefilling[active_non_spec_mask_cpu]
+            use_prefill = torch.any(
+                non_spec_is_prefilling | (non_spec_query_lens_cpu != 1)
+            ).item()
+            if use_prefill:
+                num_prefills = num_non_spec_requests
+                num_prefill_tokens = num_non_spec_tokens
+                num_decodes = 0
+                num_decode_tokens = 0
+            else:
+                num_prefills = 0
+                num_prefill_tokens = 0
+                num_decodes = num_non_spec_requests
+                num_decode_tokens = num_non_spec_tokens
+
+            num_spec_decode_tokens = num_query_tokens - num_non_spec_tokens
+
+            if num_prefills == 0 and num_decodes == 0:
+                spec_token_indx = None
+                non_spec_token_indx = None
+                # Real requests precede trailing cudagraph padding.
+                spec_state_indices_tensor = block_table_tensor[
+                    :num_spec_decodes, : self.num_spec + 1
+                ]
+                non_spec_state_indices_tensor = None
+                # Padding trails real requests, so this prefix already contains
+                # the correct cumulative token counts.
+                spec_query_start_loc = query_start_loc[: num_spec_decodes + 1]
+                non_spec_query_start_loc = None
+                non_spec_query_start_loc_cpu = None
+                num_accepted_tokens = num_accepted_tokens[:num_spec_decodes]
+            else:
+                query_lens = query_start_loc.diff()
+                spec_sequence_masks_gpu = async_tensor_h2d(
+                    spec_sequence_masks_cpu, device=query_start_loc.device
+                )
+                spec_token_masks = torch.repeat_interleave(
+                    spec_sequence_masks_gpu,
+                    query_lens,
+                    output_size=num_query_tokens,
+                )
+                # Stable partitioning preserves request-local token order in
+                # both subgroup tensors.
+                index = torch.argsort(spec_token_masks, stable=True)
+                num_non_spec_tokens = num_prefill_tokens + num_decode_tokens
+                non_spec_token_indx = index[:num_non_spec_tokens]
+                spec_token_indx = index[num_non_spec_tokens:]
+
+                # Spec requests carry one state slot per speculative step;
+                # non-spec requests use only their current state slot.
+                spec_state_indices_tensor = block_table_tensor[
+                    spec_sequence_masks_cpu, : self.num_spec + 1
+                ]
+                non_spec_state_indices_tensor = block_table_tensor[
+                    active_non_spec_mask_cpu, 0
+                ]
+
+                spec_query_lens = query_lens[spec_sequence_masks_cpu]
+                spec_query_start_loc = torch.zeros(
+                    num_spec_decodes + 1,
+                    dtype=torch.int32,
+                    device=query_start_loc.device,
+                )
+                torch.cumsum(
+                    spec_query_lens,
+                    dim=0,
+                    out=spec_query_start_loc[1:],
+                )
+                if num_prefills > 0:
+                    non_spec_query_lens = query_lens[active_non_spec_mask_cpu]
+                    non_spec_query_start_loc = torch.zeros(
+                        non_spec_query_lens.size(0) + 1,
+                        dtype=torch.int32,
+                        device=query_start_loc.device,
+                    )
+                    torch.cumsum(
+                        non_spec_query_lens,
+                        dim=0,
+                        out=non_spec_query_start_loc[1:],
+                    )
+                    non_spec_query_start_loc_cpu = torch.zeros(
+                        non_spec_query_lens_cpu.size(0) + 1,
+                        dtype=torch.int32,
+                    )
+                    torch.cumsum(
+                        non_spec_query_lens_cpu,
+                        dim=0,
+                        out=non_spec_query_start_loc_cpu[1:],
+                    )
+                else:
+                    # Packed decode consumes one row per request and does not
+                    # use cumulative sequence lengths.
+                    non_spec_query_start_loc = None
+                    non_spec_query_start_loc_cpu = None
+
+                num_accepted_tokens = num_accepted_tokens[spec_sequence_masks_cpu]
+
+        # Unlike the shared GDN layer, Kimi-K3's prefill KDA wrapper prepares
+        # its own chunk indices. Only causal-convolution metadata is needed here.
+        nums_dict, batch_ptr, token_chunk_offset_ptr = None, None, None
+        if num_prefills > 0:
+            has_initial_state = m.compute_num_computed_tokens() > 0
+            if spec_sequence_masks_cpu is not None:
+                has_initial_state = has_initial_state[active_non_spec_mask_cpu]
+                assert non_spec_query_start_loc_cpu is not None
+            nums_dict, batch_ptr, token_chunk_offset_ptr = (
+                compute_causal_conv1d_metadata(
+                    non_spec_query_start_loc_cpu,
+                    device=query_start_loc.device,
+                )
+            )
+        else:
+            has_initial_state = None
+
+        # Prepare per-request tensors for cudagraph replay. num_actual_tokens
+        # may be token-padded, while state/query/acceptance metadata is indexed
+        # by request.
+        batch_size = m.num_reqs
+        if (
+            self.use_full_cuda_graph
+            and num_spec_decodes > 0
+            and num_prefills == 0
+            and num_decodes == 0
+            and num_spec_decodes <= self.decode_cudagraph_max_bs
+            and num_spec_decode_tokens <= self.decode_cudagraph_max_bs
+        ):
+            # Equivalent PyTorch staging:
+            #   state[:N].copy_(state_src); state[N:].fill_(NULL_BLOCK_ID)
+            #   qsl[:N + 1].copy_(qsl_src); qsl[N + 1:].fill_(qsl_src[-1])
+            #   accepted[:N].copy_(accepted_src); accepted[N:].fill_(1)
+            stage_spec_decode_metadata(
+                state_indices=spec_state_indices_tensor,
+                query_start_loc=spec_query_start_loc,
+                num_accepted_tokens=num_accepted_tokens,
+                staged_state_indices=self.spec_state_indices_tensor[:batch_size],
+                staged_query_start_loc=self.spec_query_start_loc[: batch_size + 1],
+                staged_num_accepted_tokens=self.num_accepted_tokens[:batch_size],
+                num_spec_decodes=num_spec_decodes,
+            )
+
+            spec_state_indices_tensor = self.spec_state_indices_tensor[:batch_size]
+            spec_query_start_loc = self.spec_query_start_loc[: batch_size + 1]
+            num_accepted_tokens = self.num_accepted_tokens[:batch_size]
+
+        if (
+            self.use_full_cuda_graph
+            and num_prefills == 0
+            and num_spec_decodes == 0
+            and num_decodes <= self.decode_cudagraph_max_bs
+        ):
+            self.non_spec_state_indices_tensor[:num_decodes].copy_(
+                non_spec_state_indices_tensor, non_blocking=True
+            )
+            self.non_spec_state_indices_tensor[num_decodes:batch_size].fill_(
+                NULL_BLOCK_ID
+            )
+            non_spec_state_indices_tensor = self.non_spec_state_indices_tensor[
+                :batch_size
+            ]
+
+        return KimiK3KDAMetadata(
+            num_prefills=num_prefills,
+            num_prefill_tokens=num_prefill_tokens,
+            num_decodes=num_decodes,
+            num_decode_tokens=num_decode_tokens,
+            num_spec_decodes=num_spec_decodes,
+            num_spec_decode_tokens=num_spec_decode_tokens,
+            num_actual_tokens=m.num_actual_tokens,
+            has_initial_state=has_initial_state,
+            spec_query_start_loc=spec_query_start_loc,
+            non_spec_query_start_loc=non_spec_query_start_loc,
+            spec_state_indices_tensor=spec_state_indices_tensor,
+            non_spec_state_indices_tensor=non_spec_state_indices_tensor,
+            spec_sequence_masks=None,
+            spec_token_indx=spec_token_indx,
+            non_spec_token_indx=non_spec_token_indx,
+            num_accepted_tokens=num_accepted_tokens,
+            nums_dict=nums_dict,
+            batch_ptr=batch_ptr,
+            token_chunk_offset_ptr=token_chunk_offset_ptr,
+        )
+
+
+class KimiK3KDAAttentionBackend(GDNAttentionBackend):
+    @staticmethod
+    def get_name() -> str:
+        return "KIMI_K3_KDA"
+
+    @staticmethod
+    def get_builder_cls() -> type[KimiK3KDAMetadataBuilder]:
+        return KimiK3KDAMetadataBuilder

vllm/models/kimi_k3/nvidia/low_latency_gemm.py

@@ -0,0 +1,513 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Kimi-K3 decode GEMM selection for unquantized BF16 on SM103.
+
+Dispatch is purely by local ``(N, K)`` shape and token count ``M`` — the module
+name plays no role. Each measured shape maps to a :class:`ProjectionSpec`
+holding the winning backend per token count. The static part of the decision is
+resolved once per module at install time into a small ``{M: call}`` plan, so the
+per-forward path is a single dict lookup.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Literal
+
+import torch
+from torch import nn
+
+import vllm.envs as envs
+from vllm import _custom_ops as ops
+from vllm.model_executor.kernels.linear.cute_dsl.skinny_gemm import (
+    SkinnyGemmConfig,
+    shape_dynamic_skinny_gemm,
+)
+from vllm.model_executor.layers.linear import LinearBase, UnquantizedLinearMethod
+from vllm.model_executor.layers.vocab_parallel_embedding import (
+    ParallelLMHead,
+    UnquantizedEmbeddingMethod,
+)
+from vllm.platforms import current_platform
+
+Backend = Literal["cute", "dsv3_fused_a"]
+# A resolved per-token-count call: the backend plus its CuTe config (None for
+# dsv3, which needs no config).
+ResolvedCall = tuple[Backend, SkinnyGemmConfig | None]
+
+
+@dataclass(frozen=True, slots=True)
+class ProjectionSpec:
+    n: int
+    k: int
+    dsv3_tokens: frozenset[int] = frozenset()
+    cute_configs: tuple[tuple[int, SkinnyGemmConfig], ...] = ()
+    residual_configs: tuple[tuple[int, SkinnyGemmConfig], ...] = ()
+    name: str = ""  # optional debug label; never used for dispatch
+
+    def cute_config(self, num_tokens: int) -> SkinnyGemmConfig | None:
+        return dict(self.cute_configs).get(num_tokens)
+
+    def residual_config(self, num_tokens: int) -> SkinnyGemmConfig | None:
+        return dict(self.residual_configs).get(num_tokens)
+
+
+def _cute(
+    num_tokens: int,
+    block_size: int,
+    outputs_per_block: int,
+    k_unroll: int,
+    vector_width: int = 8,
+) -> SkinnyGemmConfig:
+    return SkinnyGemmConfig(
+        num_tokens,
+        block_size,
+        outputs_per_block,
+        k_unroll,
+        vector_width,
+    )
+
+
+_M1_TO_16 = frozenset(range(1, 17))
+_M1 = frozenset({1})
+
+# Keyed by local (N, K). Where two projections share a shape (only 1536x7168:
+# shared_gate_up_proj and mla_g_proj) the entry is unified.
+KIMI_K3_PROJECTIONS: dict[tuple[int, int], ProjectionSpec] = {
+    (1536, 128): ProjectionSpec(1536, 128, _M1_TO_16, name="f_b_proj"),
+    (3072, 128): ProjectionSpec(3072, 128, _M1_TO_16, name="f_b_proj"),
+    # 1536x7168 is shared by shared_gate_up_proj and mla_g_proj. dsv3 M1..16 is
+    # only crash-safe once the mla_g aux-stream/PDL capture fix lands (subtask
+    # task_7388aba1); the fallback if it cannot be fixed is dsv3_tokens=_M1.
+    (1536, 7168): ProjectionSpec(
+        1536, 7168, _M1_TO_16, name="shared_gate_up_proj/mla_g_proj"
+    ),
+    (3072, 7168): ProjectionSpec(
+        3072,
+        7168,
+        cute_configs=(
+            (1, _cute(1, 224, 3, 4)),
+            (2, _cute(2, 128, 3, 2)),
+            (3, _cute(3, 128, 2, 1)),
+            (4, _cute(4, 64, 2, 2)),
+            (5, _cute(5, 128, 3, 1)),
+        ),
+        name="shared_gate_up_proj",
+    ),
+    (2112, 7168): ProjectionSpec(2112, 7168, _M1_TO_16, name="fused_qkv_a_proj"),
+    (2304, 1536): ProjectionSpec(2304, 1536, _M1_TO_16, name="q_b_proj"),
+    (4608, 1536): ProjectionSpec(4608, 1536, _M1_TO_16, name="q_b_proj"),
+    (3584, 7168): ProjectionSpec(
+        3584,
+        7168,
+        frozenset(range(2, 9)),
+        ((1, _cute(1, 224, 2, 4)),),
+        name="routed_expert_down_proj",
+    ),
+    (6288, 7168): ProjectionSpec(
+        6288,
+        7168,
+        cute_configs=(
+            (1, _cute(1, 224, 3, 4)),
+            (2, _cute(2, 64, 3, 2)),
+            (3, _cute(3, 32, 3, 4)),
+            (4, _cute(4, 128, 6, 1)),
+        ),
+        name="in_proj_qkvgfab",
+    ),
+    (12448, 7168): ProjectionSpec(
+        12448,
+        7168,
+        cute_configs=(
+            (1, _cute(1, 224, 4, 2)),
+            (2, _cute(2, 64, 4, 2)),
+            (3, _cute(3, 64, 2, 2)),
+        ),
+        name="in_proj_qkvgfab",
+    ),
+    (7168, 768): ProjectionSpec(7168, 768, _M1_TO_16, name="shared_down_proj"),
+    (7168, 1536): ProjectionSpec(
+        7168, 1536, cute_configs=((1, _cute(1, 96, 4, 2)),), name="o_proj"
+    ),
+    (7168, 3072): ProjectionSpec(
+        7168,
+        3072,
+        cute_configs=(
+            (1, _cute(1, 96, 2, 4)),
+            (2, _cute(2, 32, 4, 4)),
+        ),
+        name="o_proj",
+    ),
+    (7168, 3584): ProjectionSpec(
+        7168,
+        3584,
+        cute_configs=(
+            (1, _cute(1, 224, 4, 2)),
+            (2, _cute(2, 64, 4, 2)),
+        ),
+        residual_configs=(
+            (1, _cute(1, 64, 4, 2)),
+            (2, _cute(2, 64, 7, 2)),
+            (3, _cute(3, 64, 2, 1)),
+            (4, _cute(4, 64, 2, 1)),
+        ),
+        name="routed_expert_up_proj",
+    ),
+    (7168, 4224): ProjectionSpec(
+        7168,
+        4224,
+        cute_configs=((1, _cute(1, 96, 4, 2, 4)),),
+        name="dense_down_proj",
+    ),
+    (7168, 8448): ProjectionSpec(
+        7168,
+        8448,
+        cute_configs=(
+            (1, _cute(1, 32, 4, 4)),
+            (2, _cute(2, 96, 4, 1)),
+            (3, _cute(3, 96, 4, 1)),
+        ),
+        name="dense_down_proj",
+    ),
+    (8448, 7168): ProjectionSpec(
+        8448,
+        7168,
+        cute_configs=(
+            (1, _cute(1, 224, 3, 4)),
+            (2, _cute(2, 32, 4, 4)),
+        ),
+        name="dense_gate_up_proj",
+    ),
+    (16896, 7168): ProjectionSpec(
+        16896,
+        7168,
+        cute_configs=(
+            (1, _cute(1, 224, 6, 4)),
+            (2, _cute(2, 32, 4, 4)),
+        ),
+        name="dense_gate_up_proj",
+    ),
+    (20480, 7168): ProjectionSpec(
+        20480,
+        7168,
+        cute_configs=(
+            (1, _cute(1, 224, 4, 2)),
+            (2, _cute(2, 64, 4, 2)),
+            (3, _cute(3, 64, 2, 2)),
+            (4, _cute(4, 64, 4, 1)),
+        ),
+        name="lm_head",
+    ),
+    (40960, 7168): ProjectionSpec(
+        40960,
+        7168,
+        cute_configs=(
+            (1, _cute(1, 128, 4, 2)),
+            (2, _cute(2, 64, 4, 2)),
+            (3, _cute(3, 64, 2, 2)),
+            (4, _cute(4, 64, 4, 1)),
+        ),
+        name="lm_head",
+    ),
+    # TP16. Measured on B300 over M=1..16 with the same >=5% threshold as the
+    # entries above. The replicated projections (2112x7168, 3584x7168,
+    # 7168x3584) keep their shapes at TP16 and reuse the entries above, and
+    # o_proj lands on 7168x768, which shared_down_proj already covers.
+    (3216, 7168): ProjectionSpec(
+        3216,
+        7168,
+        # Both gaps in this range are measured, not oversights: dsv3 is only
+        # 4% ahead at M6..M8, and at M16 cuBLAS switches to a faster kernel
+        # (11.42us vs dsv3's 11.83us) after trailing it by 6-8% at M9..M15.
+        frozenset(range(9, 16)),
+        cute_configs=(
+            (1, _cute(1, 224, 3, 4)),
+            (2, _cute(2, 128, 4, 2)),
+            (3, _cute(3, 128, 2, 1)),
+            (4, _cute(4, 64, 2, 2)),
+            (5, _cute(5, 128, 3, 1)),
+        ),
+        name="in_proj_qkvgfab",
+    ),
+    (768, 7168): ProjectionSpec(
+        768,
+        7168,
+        frozenset(range(5, 17)),
+        cute_configs=(
+            (1, _cute(1, 224, 2, 4)),
+            (2, _cute(2, 224, 2, 2)),
+            (3, _cute(3, 224, 2, 2)),
+            (4, _cute(4, 224, 2, 2)),
+        ),
+        name="mla_g_proj/shared_gate_up_proj",
+    ),
+    (1152, 1536): ProjectionSpec(
+        1152,
+        1536,
+        frozenset(range(2, 17)),
+        ((1, _cute(1, 192, 3, 4)),),
+        name="q_b_proj",
+    ),
+    (768, 128): ProjectionSpec(768, 128, _M1_TO_16, name="f_b_proj"),
+    # dsv3 drops under 5% from M9 on for this shape.
+    (7168, 384): ProjectionSpec(
+        7168, 384, frozenset(range(1, 9)), name="shared_down_proj"
+    ),
+    (4224, 7168): ProjectionSpec(
+        4224,
+        7168,
+        frozenset(range(4, 9)),
+        cute_configs=(
+            (1, _cute(1, 224, 3, 4)),
+            (2, _cute(2, 128, 2, 1)),
+            (3, _cute(3, 64, 2, 2)),
+        ),
+        name="dense_gate_up_proj",
+    ),
+    (10240, 7168): ProjectionSpec(
+        10240,
+        7168,
+        cute_configs=(
+            (1, _cute(1, 224, 4, 2)),
+            (2, _cute(2, 32, 2, 4)),
+            (3, _cute(3, 64, 4, 1)),
+            (4, _cute(4, 64, 4, 1)),
+        ),
+        name="lm_head",
+    ),
+    # 7168x2112 (TP16 dense down_proj) has no entry on purpose: K=2112 divides
+    # none of the fused-A tile_k values, and the CuTe kernel is left with
+    # vector_width=2, which measured slower than cuBLAS.
+}
+
+
+def _backend_for(
+    spec: ProjectionSpec, num_tokens: int, has_residual: bool
+) -> Backend | None:
+    if has_residual:
+        return "cute" if spec.residual_config(num_tokens) is not None else None
+    if spec.cute_config(num_tokens) is not None:
+        return "cute"
+    if num_tokens in spec.dsv3_tokens:
+        return "dsv3_fused_a"
+    return None
+
+
+def select_kimi_k3_backend(
+    num_tokens: int,
+    n: int,
+    k: int,
+    *,
+    has_residual: bool = False,
+) -> Backend | None:
+    """Backend for a local ``(N, K)`` at ``num_tokens``, or None to fall back."""
+    spec = KIMI_K3_PROJECTIONS.get((n, k))
+    return _backend_for(spec, num_tokens, has_residual) if spec is not None else None
+
+
+def _build_plan(spec: ProjectionSpec) -> dict[int, ResolvedCall]:
+    plan: dict[int, ResolvedCall] = {}
+    for num_tokens in range(1, 17):
+        backend = _backend_for(spec, num_tokens, has_residual=False)
+        if backend == "cute":
+            plan[num_tokens] = ("cute", spec.cute_config(num_tokens))
+        elif backend == "dsv3_fused_a":
+            plan[num_tokens] = ("dsv3_fused_a", None)
+    return plan
+
+
+def _build_residual_plan(spec: ProjectionSpec) -> dict[int, SkinnyGemmConfig]:
+    return {num_tokens: config for num_tokens, config in spec.residual_configs}
+
+
+def _is_sm103() -> bool:
+    return current_platform.is_device_capability((10, 3))
+
+
+def _is_packed_row_major(tensor: torch.Tensor) -> bool:
+    return tensor.dim() == 2 and tensor.stride() == (tensor.shape[1], 1)
+
+
+def _runtime_ok(x: torch.Tensor, weight: torch.Tensor) -> bool:
+    return (
+        _is_packed_row_major(x)
+        and _is_packed_row_major(weight)
+        and x.dtype == torch.bfloat16
+        and weight.dtype == torch.bfloat16
+        and x.is_cuda
+        and weight.is_cuda
+        and x.device == weight.device
+        and x.shape[1] == weight.shape[1]
+    )
+
+
+def _residual_ok(x: torch.Tensor, weight: torch.Tensor, residual: torch.Tensor) -> bool:
+    return (
+        residual.dim() == 2
+        and residual.dtype == torch.bfloat16
+        and residual.is_cuda
+        and residual.device == x.device
+        and residual.is_contiguous()
+        and residual.shape == (x.shape[0], weight.shape[0])
+    )
+
+
+def _run_plan(
+    plan: dict[int, ResolvedCall], x: torch.Tensor, weight: torch.Tensor
+) -> torch.Tensor | None:
+    entry = plan.get(x.shape[0])
+    if entry is None:
+        return None
+    backend, config = entry
+    if backend == "cute":
+        if not shape_dynamic_skinny_gemm.is_available():
+            return None
+        return shape_dynamic_skinny_gemm(x, weight, config, None)
+    if not hasattr(torch.ops._C, "dsv3_fused_a_gemm"):
+        return None
+    output = torch.empty((x.shape[0], weight.shape[0]), dtype=x.dtype, device=x.device)
+    ops.dsv3_fused_a_gemm(output, x, weight.t(), enable_pdl=True)
+    return output
+
+
+def _run_residual_plan(
+    residual_plan: dict[int, SkinnyGemmConfig],
+    x: torch.Tensor,
+    weight: torch.Tensor,
+    residual: torch.Tensor,
+) -> torch.Tensor | None:
+    config = residual_plan.get(x.shape[0])
+    if config is None or not shape_dynamic_skinny_gemm.is_available():
+        return None
+    return shape_dynamic_skinny_gemm(x, weight, config, residual)
+
+
+def try_low_latency_gemm(
+    x: torch.Tensor,
+    weight: torch.Tensor,
+    residual: torch.Tensor | None = None,
+) -> torch.Tensor | None:
+    """Run the shape-selected low-latency kernel, or None to fall back.
+
+    Resolves the plan from the shape table on each call; production installs a
+    precomputed plan (see :func:`enable_kimi_k3_low_latency_gemm`) and does not
+    use this path.
+    """
+    if envs.VLLM_BATCH_INVARIANT or not _is_sm103() or not _runtime_ok(x, weight):
+        return None
+    spec = KIMI_K3_PROJECTIONS.get((weight.shape[0], weight.shape[1]))
+    if spec is None:
+        return None
+    if residual is None:
+        return _run_plan(_build_plan(spec), x, weight)
+    if not _residual_ok(x, weight, residual):
+        return None
+    return _run_residual_plan(_build_residual_plan(spec), x, weight, residual)
+
+
+class _KimiK3LowLatencyApply:
+    """Mixin: try the precomputed plan, else defer to the base method."""
+
+    def __init__(self, plan: dict[int, ResolvedCall]) -> None:
+        self._plan = plan
+
+    def apply(
+        self,
+        layer: nn.Module,
+        x: torch.Tensor,
+        bias: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        if (
+            bias is None
+            and not envs.VLLM_BATCH_INVARIANT
+            and _runtime_ok(x, layer.weight)
+        ):
+            output = _run_plan(self._plan, x, layer.weight)
+            if output is not None:
+                return output
+        return super().apply(layer, x, bias)  # type: ignore[misc]
+
+
+class KimiK3LowLatencyLinearMethod(_KimiK3LowLatencyApply, UnquantizedLinearMethod):
+    def __init__(
+        self,
+        plan: dict[int, ResolvedCall],
+        residual_plan: dict[int, SkinnyGemmConfig],
+    ) -> None:
+        super().__init__(plan)
+        self._residual_plan = residual_plan
+
+    def apply_with_residual(
+        self,
+        layer: nn.Module,
+        x: torch.Tensor,
+        residual: torch.Tensor,
+    ) -> torch.Tensor:
+        if (
+            not envs.VLLM_BATCH_INVARIANT
+            and _runtime_ok(x, layer.weight)
+            and _residual_ok(x, layer.weight, residual)
+        ):
+            output = _run_residual_plan(self._residual_plan, x, layer.weight, residual)
+            if output is not None:
+                return output
+        return torch.addmm(residual, x, layer.weight.t())
+
+
+class KimiK3LowLatencyEmbeddingMethod(
+    _KimiK3LowLatencyApply, UnquantizedEmbeddingMethod
+):
+    pass
+
+
+def enable_kimi_k3_low_latency_gemm(
+    module: nn.Module,
+    dtype: torch.dtype,
+) -> None:
+    """Install shape-selected low-latency GEMMs and register CuTe warmups.
+
+    Modules are matched purely by type, an exactly-unquantized method, and a
+    local ``(N, K)`` present in :data:`KIMI_K3_PROJECTIONS`.
+    """
+    if dtype != torch.bfloat16 or not _is_sm103():
+        return
+
+    warmup_configs: set[SkinnyGemmConfig] = set()
+    residual_warmup_configs: set[SkinnyGemmConfig] = set()
+    for child in module.modules():
+        is_linear = (
+            isinstance(child, LinearBase)
+            and type(child.quant_method) is UnquantizedLinearMethod
+        )
+        # ParallelLMHead is a VocabParallelEmbedding subclass; embed_tokens is
+        # the parent type, so isinstance already excludes it.
+        is_head = (
+            isinstance(child, ParallelLMHead)
+            and type(child.quant_method) is UnquantizedEmbeddingMethod
+        )
+        if not (is_linear or is_head):
+            continue
+        weight = getattr(child, "weight", None)
+        if weight is None or weight.dim() != 2:
+            continue
+        spec = KIMI_K3_PROJECTIONS.get((weight.shape[0], weight.shape[1]))
+        if spec is None:
+            continue
+        if is_linear:
+            child.quant_method = KimiK3LowLatencyLinearMethod(
+                _build_plan(spec), _build_residual_plan(spec)
+            )
+        else:
+            child.quant_method = KimiK3LowLatencyEmbeddingMethod(_build_plan(spec))
+        # Warm up only the configs measured for this module's local (N, K) so a
+        # TP8 deployment does not compile TP4 configs and vice versa.
+        warmup_configs.update(config for _, config in spec.cute_configs)
+        residual_warmup_configs.update(config for _, config in spec.residual_configs)
+
+    if shape_dynamic_skinny_gemm.is_available():
+        if warmup_configs:
+            shape_dynamic_skinny_gemm.request_warmup_configs(dtype, warmup_configs)
+        if residual_warmup_configs:
+            shape_dynamic_skinny_gemm.request_warmup_configs(
+                dtype, residual_warmup_configs, has_residual=True
+            )

vllm/models/kimi_k3/nvidia/mla.py

@@ -0,0 +1,764 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Clean Multi-head Latent Attention for Kimi-K3 (NVIDIA).
+
+This is a self-contained MLA layer that owns the full attention path:
+
+    hidden_states
+      -> fused pre-attention ops (fused_qkv_a_proj / norms / q_b_proj)
+      -> explicit prefill / decode split
+           prefill: fused key-concat + cache-insert kernel -> run_prefill_new_tokens
+                    (+ chunked-context merge); dispatched by cache dtype
+                    (bf16 / plain fp8 / fp8_ds_mla)
+           decode : W_UK absorb (BMM1) -> fused q-concat + cache-insert kernel
+                    -> impl.forward_mqa -> W_UV up-proj (MQA)
+      -> optional output gate
+      -> o_proj
+
+Unlike ``MultiHeadLatentAttentionWrapper`` (which delegates orchestration to
+``MLAAttention.forward``), this class *is* the ``AttentionLayerBase``: it selects
+the backend, builds the impl, registers itself in the forward context, owns the
+KV cache, and absorbs ``kv_b_proj`` into ``W_UK_T`` / ``W_UV`` -- mirroring the
+``DeepseekV4Attention`` structure.
+
+K3 specifics: optional rotary embedding (disabled for the target model's NoPE
+layers, enabled for DSpark) and an optional sigmoid output gate (``g_proj``).
+
+Out of scope (extension points, not wired here): context parallelism (DCP/PCP),
+sparse/indexer MLA, and the ROCm/aiter fp8/fp4 BMM fast paths.
+"""
+
+import math
+from typing import TYPE_CHECKING, cast
+
+import torch
+from torch import nn
+
+from vllm.compilation.breakable_cudagraph import eager_break_during_capture
+from vllm.config import CacheConfig, VllmConfig, get_current_vllm_config
+from vllm.distributed import get_tensor_model_parallel_world_size
+from vllm.forward_context import get_forward_context
+from vllm.logger import init_logger
+from vllm.model_executor.layers.attention.attention import (
+    _init_kv_cache_quant,
+    set_default_quant_scales,
+    should_load_quant_weights,
+)
+from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
+from vllm.model_executor.layers.layernorm import RMSNorm
+from vllm.model_executor.layers.linear import (
+    ColumnParallelLinear,
+    MergedColumnParallelLinear,
+    ReplicatedLinear,
+    RowParallelLinear,
+)
+from vllm.model_executor.layers.quantization import QuantizationConfig
+from vllm.model_executor.layers.quantization.utils.quant_utils import (
+    get_and_maybe_dequant_weights,
+)
+from vllm.model_executor.layers.rotary_embedding import RotaryEmbedding, get_rope
+from vllm.model_executor.utils import replace_parameter
+from vllm.models.common.ops import fused_q_kv_rmsnorm
+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
+from vllm.transformers_utils.configs.kimi_linear import KimiLinearConfig
+from vllm.utils.multi_stream_utils import maybe_execute_in_parallel
+from vllm.utils.torch_utils import (
+    is_quantized_kv_cache,
+    kv_cache_dtype_str_to_dtype,
+)
+from vllm.v1.attention.backend import (
+    AttentionBackend,
+    AttentionType,
+    MLAAttentionImpl,
+)
+from vllm.v1.attention.backends.mla.prefill import get_mla_prefill_backend
+from vllm.v1.attention.ops.merge_attn_states import merge_attn_states
+from vllm.v1.attention.selector import get_attn_backend
+from vllm.v1.kv_cache_interface import KVCacheSpec, MLAAttentionSpec, get_kv_quant_mode
+
+if TYPE_CHECKING:
+    from vllm.model_executor.layers.attention.mla_attention import MLACommonMetadata
+
+logger = init_logger(__name__)
+
+# Below this many tokens, overlap the g_proj GEMM on the aux stream with the
+# attention front-end (the GEMM is small and launch-bound, so the overlap
+# hides it); at or above it, run the gate on the main stream.
+_GATE_MULTI_STREAM_TOKEN_THRESHOLD = 512
+
+
+@torch.compile(backend=current_platform.simple_compile_backend)
+def _gate_sigmoid_mul(attn_out: torch.Tensor, gate: torch.Tensor) -> torch.Tensor:
+    """Apply the sigmoid output gate to a precomputed ``g_proj`` projection."""
+    return attn_out * gate.sigmoid()
+
+
+class MultiHeadLatentAttention(nn.Module, AttentionLayerBase):
+    """Kimi-K3 Multi-head Latent Attention with optional RoPE and output gate."""
+
+    def __init__(
+        self,
+        config: KimiLinearConfig,
+        hidden_size: int,
+        num_heads: int,
+        qk_nope_head_dim: int,
+        qk_rope_head_dim: int,
+        v_head_dim: int,
+        q_lora_rank: int | None,
+        kv_lora_rank: int,
+        use_output_gate: bool = False,
+        cache_config: CacheConfig | None = None,
+        quant_config: QuantizationConfig | None = None,
+        prefix: str = "",
+        aux_stream: torch.cuda.Stream | None = None,
+        use_rope: bool = False,
+        non_causal_multi_token_decode: bool = False,
+    ) -> None:
+        super().__init__()
+        self.hidden_size = hidden_size
+        self.qk_nope_head_dim = qk_nope_head_dim
+        self.qk_rope_head_dim = qk_rope_head_dim
+        self.qk_head_dim = qk_nope_head_dim + qk_rope_head_dim
+        self.v_head_dim = v_head_dim
+        self.q_lora_rank = q_lora_rank
+        self.kv_lora_rank = kv_lora_rank
+        self.non_causal_multi_token_decode = non_causal_multi_token_decode
+        # Latent "head" seen by the attention kernel / KV cache.
+        self.head_size = kv_lora_rank + qk_rope_head_dim
+        self.scale = self.qk_head_dim**-0.5
+        self.rms_norm_eps = config.rms_norm_eps
+        self.layer_name = prefix
+
+        self.rotary_emb: RotaryEmbedding | None = None
+        if use_rope:
+            rope_parameters = dict(config.rope_parameters)
+            if rope_parameters["rope_type"] != "default":
+                rope_parameters["rope_type"] = (
+                    "deepseek_yarn"
+                    if rope_parameters.get("apply_yarn_scaling", True)
+                    else "deepseek_llama_scaling"
+                )
+            self.rotary_emb = get_rope(
+                qk_rope_head_dim,
+                max_position=config.max_position_embeddings,
+                rope_parameters=rope_parameters,
+                is_neox_style=False,
+                dtype=torch.float32,
+            )
+            if rope_parameters["rope_type"] == "deepseek_yarn":
+                mscale_all_dim = rope_parameters.get("mscale_all_dim", False)
+                scaling_factor = rope_parameters["factor"]
+                mscale = (
+                    1.0
+                    if scaling_factor <= 1
+                    else 0.1 * float(mscale_all_dim) * math.log(scaling_factor) + 1.0
+                )
+                self.scale *= mscale * mscale
+            # The fused epilogues read the cos/sin table directly in fp32 and run
+            # the RoPE math in fp32, so there is no per-forward dtype cast (and no
+            # precision loss). deepseek_yarn builds cos_sin_cache in fp32 already;
+            # dtype=torch.float32 above forces it for the default rope too (the
+            # DSpark draft, which has no yarn scaling).
+            assert self.rotary_emb.cos_sin_cache.dtype == torch.float32, (
+                "K3 fused MLA RoPE requires an fp32 cos/sin cache; got "
+                f"{self.rotary_emb.cos_sin_cache.dtype}."
+            )
+
+        tp_size = get_tensor_model_parallel_world_size()
+        assert num_heads % tp_size == 0
+        self.num_heads = num_heads
+        self.num_local_heads = num_heads // tp_size
+
+        # ---- Pre-attention projections (fusable front-end) ----
+        # Two query variants: a low-rank q-LoRA path (Kimi-K3) fused with the
+        # kv-down proj, or an uncompressed q path (Kimi-Linear, ``q_lora_rank``
+        # None) with a standalone ``q_proj`` and separate ``kv_a_proj_with_mqa``.
+        if self.q_lora_rank is not None:
+            # Fused q-down + kv-down projection. Replicated (disable_tp) because
+            # the low-rank latents are shared across TP ranks; TP splitting
+            # happens at q_b_proj / kv_b_proj. Checkpoint weights ``q_a_proj``
+            # and ``kv_a_proj_with_mqa`` map onto shards 0 and 1 respectively.
+            self.fused_qkv_a_proj = MergedColumnParallelLinear(
+                self.hidden_size,
+                [self.q_lora_rank, self.kv_lora_rank + self.qk_rope_head_dim],
+                bias=False,
+                quant_config=quant_config,
+                prefix=f"{prefix}.fused_qkv_a_proj",
+                disable_tp=True,
+            )
+            self.q_a_layernorm = RMSNorm(self.q_lora_rank, eps=config.rms_norm_eps)
+            self.q_b_proj = ColumnParallelLinear(
+                self.q_lora_rank,
+                self.num_heads * self.qk_head_dim,
+                bias=False,
+                quant_config=quant_config,
+                prefix=f"{prefix}.q_b_proj",
+            )
+        else:
+            # Uncompressed query: full-rank q_proj (TP-split over heads) plus a
+            # replicated kv-down projection (shared latent across TP ranks).
+            self.q_proj = ColumnParallelLinear(
+                self.hidden_size,
+                self.num_heads * self.qk_head_dim,
+                bias=False,
+                quant_config=quant_config,
+                prefix=f"{prefix}.q_proj",
+            )
+            self.kv_a_proj_with_mqa = ReplicatedLinear(
+                self.hidden_size,
+                self.kv_lora_rank + self.qk_rope_head_dim,
+                bias=False,
+                quant_config=quant_config,
+                prefix=f"{prefix}.kv_a_proj_with_mqa",
+            )
+        self.kv_a_layernorm = RMSNorm(self.kv_lora_rank, eps=config.rms_norm_eps)
+        self.kv_b_proj = ColumnParallelLinear(
+            self.kv_lora_rank,
+            self.num_heads * (self.qk_nope_head_dim + self.v_head_dim),
+            bias=False,
+            quant_config=quant_config,
+            prefix=f"{prefix}.kv_b_proj",
+        )
+
+        # ---- Post-attention projections ----
+        self.use_output_gate = use_output_gate
+        self.g_proj = (
+            ColumnParallelLinear(
+                self.hidden_size,
+                self.num_heads * self.v_head_dim,
+                bias=False,
+                quant_config=quant_config,
+                prefix=f"{prefix}.g_proj",
+            )
+            if use_output_gate
+            else None
+        )
+        # Aux stream (created at the model level, DeepseekV4 convention) for
+        # overlapping the g_proj GEMM with the attention front-end. None on
+        # ROCm/non-cuda -> maybe_execute_in_parallel falls back to sequential.
+        self.aux_stream = aux_stream
+        self._gate_events = (
+            [torch.cuda.Event(), torch.cuda.Event()]
+            if self.g_proj is not None and current_platform.is_cuda_alike()
+            else None
+        )
+        self.o_proj = RowParallelLinear(
+            self.num_heads * self.v_head_dim,
+            self.hidden_size,
+            bias=False,
+            quant_config=quant_config,
+            prefix=f"{prefix}.o_proj",
+        )
+
+        # ---- Attention backend / impl / KV cache ----
+        self.quant_config = quant_config
+        if cache_config is not None:
+            self.kv_cache_dtype = cache_config.cache_dtype
+        else:
+            self.kv_cache_dtype = "auto"
+
+        dtype = torch.get_default_dtype()
+        self.attn_backend = get_attn_backend(
+            self.head_size,
+            dtype,
+            self.kv_cache_dtype,
+            use_mla=True,
+            use_sparse=False,
+            num_heads=self.num_local_heads,
+        )
+        _init_kv_cache_quant(self, quant_config, prefix)
+        # Unit (1.0) scale for the fused fp8 prefill path: q/k/v are cast
+        # unscaled to match forward_mha (the prefill flash path does not
+        # dequantize); only the cache uses _k_scale.
+        self.register_buffer(
+            "_one_scale", torch.ones(1, dtype=torch.float32), persistent=False
+        )
+
+        impl_cls = cast(type[MLAAttentionImpl], self.attn_backend.get_impl_cls())
+        self.impl = impl_cls(  # type: ignore[assignment]
+            num_heads=self.num_local_heads,
+            head_size=self.head_size,
+            scale=self.scale,
+            num_kv_heads=1,
+            alibi_slopes=None,
+            sliding_window=None,
+            kv_cache_dtype=self.kv_cache_dtype,
+            logits_soft_cap=None,
+            attn_type=AttentionType.DECODER,
+            kv_sharing_target_layer_name=None,
+            q_lora_rank=self.q_lora_rank,
+            kv_lora_rank=self.kv_lora_rank,
+            qk_nope_head_dim=self.qk_nope_head_dim,
+            qk_rope_head_dim=self.qk_rope_head_dim,
+            qk_head_dim=self.qk_head_dim,
+            v_head_dim=self.v_head_dim,
+            kv_b_proj=self.kv_b_proj,
+            indexer=None,
+        )
+        self.q_pad_num_heads = getattr(self.impl, "q_pad_num_heads", None)
+
+        vllm_config = get_current_vllm_config()
+        parallel_config = vllm_config.parallel_config
+        assert (
+            parallel_config.decode_context_parallel_size <= 1
+            and parallel_config.prefill_context_parallel_size <= 1
+        ), "Kimi-K3 MultiHeadLatentAttention does not support context parallelism."
+        self.prefill_backend = get_mla_prefill_backend(vllm_config)(
+            num_heads=self.num_local_heads,
+            scale=self.scale,
+            kv_lora_rank=self.kv_lora_rank,
+            qk_nope_head_dim=self.qk_nope_head_dim,
+            qk_rope_head_dim=self.qk_rope_head_dim,
+            v_head_dim=self.v_head_dim,
+            vllm_config=vllm_config,
+        )
+
+        compilation_config = vllm_config.compilation_config
+        if prefix in compilation_config.static_forward_context:
+            raise ValueError(f"Duplicate layer name: {prefix}")
+        compilation_config.static_forward_context[prefix] = self
+        self.kv_cache = torch.tensor([])
+
+    # ------------------------------------------------------------------
+    # AttentionLayerBase interface
+    # ------------------------------------------------------------------
+    def get_attn_backend(self) -> type[AttentionBackend]:
+        return self.attn_backend
+
+    def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec:
+        kv_cache_dtype = kv_cache_dtype_str_to_dtype(
+            self.kv_cache_dtype, vllm_config.model_config
+        )
+        # TODO: Remove this mypy workaround once the K3 PR is fully merged.
+        return MLAAttentionSpec(  # type: ignore[call-arg]
+            block_size=vllm_config.cache_config.block_size,
+            num_kv_heads=1,
+            head_size=self.head_size,
+            dtype=kv_cache_dtype,
+            cache_dtype_str=self.kv_cache_dtype,
+            kv_quant_mode=get_kv_quant_mode(self.kv_cache_dtype),
+            non_causal_multi_token_decode=self.non_causal_multi_token_decode,
+        )
+
+    def process_weights_after_loading(self, act_dtype: torch.dtype) -> None:
+        """Absorb ``kv_b_proj`` into decode-time ``W_UK_T`` / ``W_UV`` bmm weights.
+
+        ``kv_b_proj`` produces ``[k_nope; v]`` per head from the ``kv_lora_rank``
+        latent. For the MQA decode path we pre-split it so that queries are
+        projected into latent space by ``W_UK_T`` and the attention output is
+        projected back to ``v`` by ``W_UV`` -- avoiding materializing full K/V.
+        """
+        kv_b_proj_weight = get_and_maybe_dequant_weights(
+            self.kv_b_proj, out_dtype=act_dtype
+        ).T
+        assert kv_b_proj_weight.shape == (
+            self.kv_lora_rank,
+            self.num_local_heads * (self.qk_nope_head_dim + self.v_head_dim),
+        ), f"{kv_b_proj_weight.shape=}"
+        kv_b_proj_weight = kv_b_proj_weight.view(
+            self.kv_lora_rank,
+            self.num_local_heads,
+            self.qk_nope_head_dim + self.v_head_dim,
+        )
+        W_UK, W_UV = kv_b_proj_weight.split(
+            [self.qk_nope_head_dim, self.v_head_dim], dim=-1
+        )
+        # (L, N, V) -> (N, L, V)
+        replace_parameter(self, "W_UV", W_UV.transpose(0, 1), prefer_copy=True)
+        # (L, N, P) -> (N, P, L)
+        replace_parameter(self, "W_UK_T", W_UK.permute(1, 2, 0), prefer_copy=True)
+
+        quant_method = (
+            self.quant_config.get_quant_method(self, prefix=self.layer_name)
+            if self.quant_config
+            else None
+        )
+        if not should_load_quant_weights(quant_method):
+            set_default_quant_scales(self, register_buffer=False)
+
+        # Precompute reciprocal scales once here (scales are final after load;
+        # K3 has no runtime calculate_kv_scales path) so the fp8 fused kernels
+        # in the decode/prefill hot path take a ready inverse instead of
+        # launching a per-step reciprocal kernel.
+        self.register_buffer(
+            "_q_scale_inv", self._q_scale.reciprocal().reshape(1), persistent=False
+        )
+        self.register_buffer(
+            "_k_scale_inv", self._k_scale.reciprocal().reshape(1), persistent=False
+        )
+
+    def _v_up_proj(self, x: torch.Tensor, out: torch.Tensor) -> None:
+        """Project latent attention output back to ``v`` via ``W_UV`` (bmm)."""
+        # (B, N, L) -> (N, B, L)
+        x = x.view(-1, self.num_local_heads, self.kv_lora_rank).transpose(0, 1)
+        out = out.view(-1, self.num_local_heads, self.v_head_dim)
+        # (N, B, L) x (N, L, V) -> (N, B, V) written transposed into (B, N, V)
+        torch.bmm(x, self.W_UV, out=out.transpose(0, 1))
+
+    def _attn_read_kv_cache(self) -> torch.Tensor:
+        """Latent cache as seen by the attention read kernels (decode / context).
+
+        A plain per-tensor fp8 cache is stored as ``uint8``; view it as fp8 so
+        the backend reads it as E4M3 rather than fp4/E2M1 -- the latter doubles
+        the perceived head dim (``head_size * 2``) and fails the kernel's
+        ``head_dim_k == head_dim_q`` check. Mirrors ``MLAAttention.forward``;
+        the fp8_ds_mla layout keeps its native uint8 view.
+        """
+        cache = self.kv_cache
+        if (
+            is_quantized_kv_cache(self.kv_cache_dtype)
+            and self.kv_cache_dtype != "fp8_ds_mla"
+        ):
+            return cache.view(current_platform.fp8_dtype())
+        return cache
+
+    # ------------------------------------------------------------------
+    # Forward
+    # ------------------------------------------------------------------
+    def _forward_attn(
+        self,
+        positions: torch.Tensor,
+        hidden_states: torch.Tensor,
+    ) -> torch.Tensor:
+        """Attention front-end: fused qkv-a proj -> norms -> q_b -> attention.
+
+        Returns the pre-gate attention output ``[num_tokens,
+        num_local_heads * v_head_dim]``. On a profile/dummy run
+        it returns a zeroed buffer.
+        """
+        if self.q_lora_rank is not None:
+            qkv_lora = self.fused_qkv_a_proj(hidden_states)[0]
+            q_c, kv_c, k_pe = qkv_lora.split(
+                [self.q_lora_rank, self.kv_lora_rank, self.qk_rope_head_dim], dim=-1
+            )
+            q_c, kv_c_normed = fused_q_kv_rmsnorm(
+                q_c,
+                kv_c,
+                self.q_a_layernorm.weight.data,
+                self.kv_a_layernorm.weight.data,
+                self.rms_norm_eps,
+            )
+            q = self.q_b_proj(q_c)[0].view(-1, self.num_local_heads, self.qk_head_dim)
+        else:
+            # Uncompressed query: project directly (no q-LoRA, no q norm) and
+            # normalize only the kv latent.
+            q = self.q_proj(hidden_states)[0].view(
+                -1, self.num_local_heads, self.qk_head_dim
+            )
+            kv_lora = self.kv_a_proj_with_mqa(hidden_states)[0]
+            kv_c, k_pe = kv_lora.split(
+                [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1
+            )
+            kv_c_normed = self.kv_a_layernorm(kv_c)
+        k_pe = k_pe.unsqueeze(1)
+
+        attn_out = torch.empty(
+            (hidden_states.shape[0], self.num_local_heads * self.v_head_dim),
+            dtype=hidden_states.dtype,
+            device=hidden_states.device,
+        )
+        self._attention(positions, q, kv_c_normed, k_pe, attn_out)
+        return attn_out
+
+    def forward(
+        self,
+        positions: torch.Tensor,
+        hidden_states: torch.Tensor,
+    ) -> torch.Tensor:
+        # Both branches produce (attn_out, gate); they differ only in whether
+        # the g_proj GEMM is overlapped on the aux stream.
+        g_proj = self.g_proj
+        events = self._gate_events
+        if (
+            g_proj is not None
+            and events is not None
+            and self.aux_stream is not None
+            and hidden_states.shape[0] < _GATE_MULTI_STREAM_TOKEN_THRESHOLD
+        ):
+            attn_out, gate = maybe_execute_in_parallel(
+                lambda: self._forward_attn(positions, hidden_states),
+                lambda: g_proj(hidden_states)[0],
+                events[0],
+                events[1],
+                self.aux_stream,
+            )
+        else:
+            attn_out = self._forward_attn(positions, hidden_states)
+            gate = g_proj(hidden_states)[0] if g_proj is not None else None
+
+        if gate is not None:
+            attn_out = _gate_sigmoid_mul(attn_out, gate)
+
+        # ``o_proj`` (RowParallelLinear + out-of-place all-reduce) returns a
+        # fresh private tensor, so return it directly rather than copying into a
+        # caller buffer -- the previous ``output[:] = ...`` convention forced an
+        # extra [num_tokens, hidden] copy per layer.
+        return self.o_proj(attn_out)[0]
+
+    @eager_break_during_capture
+    def _attention(
+        self,
+        positions: torch.Tensor,
+        q: torch.Tensor,
+        kv_c_normed: torch.Tensor,
+        k_pe: torch.Tensor,
+        attn_out: torch.Tensor,
+    ) -> None:
+        forward_context = get_forward_context()
+        attn_metadata_by_layer = forward_context.attn_metadata
+        if attn_metadata_by_layer is None:
+            attn_out.zero_()
+            return
+        assert isinstance(attn_metadata_by_layer, dict)
+        attn_metadata = cast(
+            "MLACommonMetadata", attn_metadata_by_layer[self.layer_name]
+        )
+
+        num_actual_toks = attn_metadata.num_actual_tokens
+        slot_mapping_by_layer = forward_context.slot_mapping
+        assert isinstance(slot_mapping_by_layer, dict)
+        slot_mapping = slot_mapping_by_layer[self.layer_name]
+
+        q = q[:num_actual_toks]
+        kv_c_normed = kv_c_normed[:num_actual_toks]
+        k_pe = k_pe[:num_actual_toks]
+        positions = positions[:num_actual_toks]
+        attn_out = attn_out[:num_actual_toks]
+
+        cos_sin_cache = None
+        rope_positions = None
+        if self.rotary_emb is not None:
+            # Pass the fp32 cos/sin table straight to the fused epilogue (it reads
+            # fp32 and does the RoPE math in fp32) -- no per-forward dtype cast.
+            cos_sin_cache = self.rotary_emb.cos_sin_cache
+            rope_positions = positions
+
+        # Decode tokens are laid out first, prefill tokens after. The fused
+        # prefill covers every supported config (bf16 / plain-fp8 /
+        # fp8_ds_mla), so there is no dense-MHA (forward_mha) fallback.
+        num_mqa_tokens = attn_metadata.num_decode_tokens
+        num_mha_tokens = q.size(0) - num_mqa_tokens
+
+        # Both the prefill and decode fused epilogues write their own cache
+        # slice, so there is no separate do_kv_cache_update.
+
+        # ---- Prefill: fused key-concat + cache-insert + attention ----
+        if num_mha_tokens > 0:
+            self._forward_prefill_fused(
+                q[num_mqa_tokens:],
+                kv_c_normed[num_mqa_tokens:],
+                k_pe[num_mqa_tokens:],
+                rope_positions[num_mqa_tokens:] if rope_positions is not None else None,
+                cos_sin_cache,
+                slot_mapping[num_mqa_tokens:num_actual_toks],
+                attn_metadata,
+                attn_out[num_mqa_tokens:],
+            )
+
+        # ---- Decode: latent multi-query attention ----
+        if num_mqa_tokens > 0:
+            mqa_q_nope, mqa_q_pe = q[:num_mqa_tokens].split(
+                [self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1
+            )
+            # BMM1: absorb q_nope into latent space. (N,B,P) x (N,P,L) -> (B,N,L)
+            ql_nope = torch.bmm(mqa_q_nope.transpose(0, 1), self.W_UK_T).transpose(0, 1)
+            # Fused: concat mqa_q = [ql_nope | q_pe] and insert the decode-token
+            # latent into the paged cache (one launch, right before forward_mqa).
+            mqa_q = self._decode_concat_cache(
+                ql_nope,
+                mqa_q_pe,
+                kv_c_normed[:num_mqa_tokens],
+                k_pe[:num_mqa_tokens],
+                rope_positions[:num_mqa_tokens] if rope_positions is not None else None,
+                cos_sin_cache,
+                slot_mapping[:num_mqa_tokens],
+            )
+            latent_out, _lse = self.impl.forward_mqa(  # type: ignore[attr-defined]
+                mqa_q, self._attn_read_kv_cache(), attn_metadata, self
+            )
+            self._v_up_proj(latent_out, out=attn_out[:num_mqa_tokens])
+
+    def _decode_concat_cache(
+        self,
+        ql_nope: torch.Tensor,
+        q_pe: torch.Tensor,
+        kv_c_normed: torch.Tensor,
+        k_pe: torch.Tensor,
+        positions: torch.Tensor | None,
+        cos_sin_cache: torch.Tensor | None,
+        slot_mapping: torch.Tensor,
+    ) -> torch.Tensor:
+        """Fused decode query-concat + latent cache insert, dispatched by cache
+        dtype (same policy as prefill: fp8 cache -> fp8 query)."""
+        if self.kv_cache_dtype == "fp8_ds_mla":
+            cache = self.kv_cache
+            if cache.dtype != torch.uint8:
+                cache = cache.view(torch.uint8)
+            return fused_mla_decode_q_concat_kv_cache_insert(
+                ql_nope,
+                q_pe,
+                kv_c_normed,
+                k_pe,
+                cache,
+                slot_mapping,
+                ds_mla=True,
+                positions=positions,
+                cos_sin_cache=cos_sin_cache,
+            )
+        if is_quantized_kv_cache(self.kv_cache_dtype):
+            assert self.impl.supports_quant_query_input, (  # type: ignore[attr-defined]
+                "Kimi-K3 fp8 KV cache decode requires a backend that accepts an "
+                "fp8 (quantized) query input."
+            )
+            cache = self.kv_cache
+            if cache.dtype != torch.float8_e4m3fn:
+                cache = cache.view(torch.float8_e4m3fn)
+            return fused_mla_decode_q_concat_kv_cache_insert(
+                ql_nope,
+                q_pe,
+                kv_c_normed,
+                k_pe,
+                cache,
+                slot_mapping,
+                q_scale_inv=self._q_scale_inv,
+                cache_scale_inv=self._k_scale_inv,
+                positions=positions,
+                cos_sin_cache=cos_sin_cache,
+            )
+        return fused_mla_decode_q_concat_kv_cache_insert(
+            ql_nope,
+            q_pe,
+            kv_c_normed,
+            k_pe,
+            self.kv_cache,
+            slot_mapping,
+            positions=positions,
+            cos_sin_cache=cos_sin_cache,
+        )
+
+    def _forward_prefill_fused(
+        self,
+        q: torch.Tensor,
+        kv_c_normed: torch.Tensor,
+        k_pe: torch.Tensor,
+        positions: torch.Tensor | None,
+        cos_sin_cache: torch.Tensor | None,
+        slot_mapping: torch.Tensor,
+        attn_metadata,
+        out: torch.Tensor,
+    ) -> None:
+        """Prefill using the fused key-concat + cache-insert kernel.
+
+        Replaces ``_concat_k_nope_k_pe`` and the prefill cache write with one
+        fused kernel launch, dispatched by cache dtype. The chunked context
+        gather + online-softmax merge are delegated to the impl.
+
+        Supported configs (K3 fp8 policy):
+          - bf16 cache        -> bf16 prefill query
+          - plain fp8 cache   -> fp8 prefill query (unscaled q/k/v; cache _k_scale)
+          - fp8_ds_mla cache  -> bf16 prefill query (656B per-tile self-scaled)
+        """
+        prefill = attn_metadata.prefill
+        has_context = prefill.chunked_context is not None
+        fp8_prefill = prefill.q_data_type == current_platform.fp8_dtype()
+
+        kv_nope = self.kv_b_proj(kv_c_normed)[0].view(
+            -1, self.num_local_heads, self.qk_nope_head_dim + self.v_head_dim
+        )
+        k_nope, v = kv_nope.split([self.qk_nope_head_dim, self.v_head_dim], dim=-1)
+
+        if self.kv_cache_dtype == "fp8_ds_mla":
+            # fp8_ds_mla cache (656B, per-tile self-scaled); bf16 attention.
+            assert not fp8_prefill, (
+                "Kimi-K3 fp8_ds_mla uses a bf16 prefill query; fp8 prefill "
+                "query is not supported with fp8_ds_mla."
+            )
+            kv_cache = self.kv_cache
+            if kv_cache.dtype != torch.uint8:
+                kv_cache = kv_cache.view(torch.uint8)
+            k = fused_mla_key_concat_ds_mla_insert(
+                q,
+                k_nope,
+                k_pe,
+                kv_c_normed,
+                kv_cache,
+                slot_mapping,
+                positions,
+                cos_sin_cache,
+            )
+        elif is_quantized_kv_cache(self.kv_cache_dtype):
+            assert fp8_prefill, (
+                "Kimi-K3 fp8 KV cache requires an fp8 prefill query; enable "
+                "--attention-config '{\"use_prefill_query_quantization\": true}'."
+            )
+            # Plain per-tensor fp8: quant q/k/v (unscaled, matching forward_mha's
+            # unscaled `.to(fp8)`) and insert the fp8 latent (scaled by _k_scale).
+            kv_cache = self.kv_cache
+            if kv_cache.dtype != torch.float8_e4m3fn:
+                kv_cache = kv_cache.view(torch.float8_e4m3fn)
+            q, k, v = fused_mla_qkv_quant_kv_cache_fp8_insert(
+                q,
+                k_nope,
+                k_pe,
+                kv_c_normed,
+                v,
+                kv_cache,
+                slot_mapping,
+                self._one_scale,
+                self._one_scale,
+                self._one_scale,
+                self._k_scale_inv,
+                positions,
+                cos_sin_cache,
+            )
+        else:
+            # Concat full K = [k_nope | k_pe] and insert [kv_c_normed | k_pe]
+            # into the paged cache for these prefill tokens, in one launch.
+            k = fused_mla_key_concat_kv_cache_insert(
+                q,
+                k_nope,
+                k_pe,
+                kv_c_normed,
+                self.kv_cache,
+                slot_mapping,
+                positions,
+                cos_sin_cache,
+            )
+
+        # When there is no chunked context, backends that honor `out` write the
+        # attention result straight into it, avoiding a slice+flatten+copy.
+        writes_out = not has_context and prefill.prefill_backend.supports_out()
+        output_prefill = prefill.prefill_backend.run_prefill_new_tokens(
+            q=q,
+            k=k,
+            v=v,
+            return_softmax_lse=has_context,
+            out=(
+                out.view(-1, self.num_local_heads, self.v_head_dim)
+                if writes_out
+                else None
+            ),
+        )
+
+        if has_context:
+            context_output, context_lse = self.impl._compute_prefill_context(  # type: ignore[attr-defined]
+                q, self._attn_read_kv_cache(), attn_metadata, self._k_scale
+            )
+            suffix_output, suffix_lse = output_prefill
+            out = out.view(-1, self.num_local_heads, self.v_head_dim)
+            merge_attn_states(
+                output=out,
+                prefix_output=context_output[..., : self.v_head_dim],
+                prefix_lse=context_lse,
+                suffix_output=suffix_output[..., : self.v_head_dim],
+                suffix_lse=suffix_lse,
+                prefill_tokens_with_context=prefill.chunked_context.prefill_tokens_with_context,
+            )
+        elif not writes_out:
+            out.copy_(output_prefill[..., : self.v_head_dim].flatten(start_dim=-2))

vllm/models/kimi_k3/nvidia/mtp.py

@@ -0,0 +1,443 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Inference-only Kimi-K3 Multi-Token-Prediction (MTP) draft model."""
+
+import copy
+from collections.abc import Iterable
+
+import torch
+import torch.nn as nn
+
+import vllm.envs as envs
+from vllm.config import VllmConfig
+from vllm.forward_context import get_forward_context, is_forward_context_available
+from vllm.logger import init_logger
+from vllm.model_executor.layers.fused_moe import (
+    fused_moe_make_expert_params_mapping,
+)
+from vllm.model_executor.layers.layernorm import RMSNorm
+from vllm.model_executor.layers.logits_processor import LogitsProcessor
+from vllm.model_executor.layers.quantization import QuantizationConfig
+from vllm.model_executor.layers.vocab_parallel_embedding import (
+    ParallelLMHead,
+    VocabParallelEmbedding,
+)
+from vllm.model_executor.model_loader.weight_utils import (
+    default_weight_loader,
+    maybe_remap_kv_scale_name,
+)
+from vllm.model_executor.models.utils import get_pp_missing_layer_names, maybe_prefix
+from vllm.sequence import IntermediateTensors
+from vllm.transformers_utils.configs.kimi_linear import KimiLinearConfig
+
+from ..common.mtp import fused_mtp_input
+from .low_latency_gemm import enable_kimi_k3_low_latency_gemm
+from .model import (
+    KimiDecoderLayer,
+    KimiMoE,
+    get_spec_layer_idx_from_weight_name,
+    make_kimi_k3_mega_moe_expert_params_mapping,
+)
+from .ops.sequence_parallel import sp_all_gather, sp_padding_mask, sp_shard
+
+logger = init_logger(__name__)
+
+
+class SharedHead(nn.Module):
+    def __init__(
+        self,
+        config: KimiLinearConfig,
+        prefix: str,
+        quant_config: QuantizationConfig | None = None,
+    ) -> None:
+        super().__init__()
+        self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.head = ParallelLMHead(
+            config.vocab_size,
+            config.hidden_size,
+            quant_config=quant_config,
+            prefix=maybe_prefix(prefix, "head"),
+        )
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        return self.norm(hidden_states)
+
+
+class KimiK3MultiTokenPredictorLayer(nn.Module):
+    def __init__(
+        self,
+        config: KimiLinearConfig,
+        vllm_config: VllmConfig,
+        prefix: str,
+    ) -> None:
+        super().__init__()
+        self.config = config
+        quant_config = vllm_config.quant_config
+
+        self.enorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.hnorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.eh_proj = nn.Linear(config.hidden_size * 2, config.hidden_size, bias=False)
+
+        self.shared_head = SharedHead(
+            config=config, prefix=prefix, quant_config=quant_config
+        )
+        # The MTP block starts without the base model's AttnRes state.
+        block_config = copy.copy(config)
+        block_config.attn_res_block_size = None
+        # NOTE: the prefix must end in the numeric spec-layer index so that
+        # KimiDecoderLayer can parse ``layer_idx`` and pick MLA (full attn).
+        # Own aux stream for the MLA g_proj output-gate overlap (DeepseekV4
+        # convention: the MTP block creates its own stream).
+        aux_stream = torch.cuda.Stream()
+        self.mtp_block = KimiDecoderLayer(
+            block_config, vllm_config, prefix=prefix, aux_stream=aux_stream
+        )
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        positions: torch.Tensor,
+        previous_hidden_states: torch.Tensor,
+        inputs_embeds: torch.Tensor | None = None,
+        spec_step_index: int = 0,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        assert inputs_embeds is not None
+        hidden_states = self.eh_proj(
+            fused_mtp_input(
+                positions,
+                inputs_embeds,
+                previous_hidden_states,
+                self.enorm.weight,
+                self.hnorm.weight,
+                self.enorm.variance_epsilon,
+            )
+        )
+        if self.mtp_block.use_sequence_parallel:
+            if envs.VLLM_MOE_SKIP_PADDING and is_forward_context_available():
+                forward_context = get_forward_context()
+                forward_context.is_padding = sp_padding_mask(
+                    forward_context.is_padding, hidden_states
+                )
+            hidden_states = sp_shard(hidden_states)
+
+        hidden_states, _, residual = self.mtp_block(
+            positions=positions,
+            hidden_states=hidden_states,
+            residual=None,
+        )
+        if self.mtp_block.use_sequence_parallel:
+            assert residual is not None
+            hidden_states = hidden_states + residual
+            hidden_states = sp_all_gather(hidden_states)[: positions.shape[0]]
+            logits_hidden_states = self.shared_head.norm(hidden_states)
+            return logits_hidden_states, hidden_states
+
+        # Produce the normalized logits input and the pre-norm recurrent state
+        # in one fused add-RMSNorm launch.
+        logits_hidden_states, hidden_states = self.shared_head.norm(
+            hidden_states, residual
+        )
+        return logits_hidden_states, hidden_states
+
+
+class KimiK3MultiTokenPredictor(nn.Module):
+    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
+        super().__init__()
+        config: KimiLinearConfig = vllm_config.model_config.hf_text_config
+        self.config = config
+        self.mtp_start_layer_idx = config.num_hidden_layers
+        self.num_mtp_layers = config.num_nextn_predict_layers
+
+        self.layers = torch.nn.ModuleDict(
+            {
+                str(idx): KimiK3MultiTokenPredictorLayer(
+                    config, vllm_config, f"{prefix}.layers.{idx}"
+                )
+                for idx in range(
+                    self.mtp_start_layer_idx,
+                    self.mtp_start_layer_idx + self.num_mtp_layers,
+                )
+            }
+        )
+        self.embed_tokens = VocabParallelEmbedding(
+            config.vocab_size,
+            config.hidden_size,
+            prefix=maybe_prefix(prefix, "embed_tokens"),
+        )
+        self.logits_processor = LogitsProcessor(config.vocab_size)
+
+    def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
+        return self.embed_tokens(input_ids)
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        positions: torch.Tensor,
+        previous_hidden_states: torch.Tensor,
+        inputs_embeds: torch.Tensor | None = None,
+        spec_step_idx: int = 0,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        if inputs_embeds is None:
+            inputs_embeds = self.embed_tokens(input_ids)
+        current_step_idx = spec_step_idx % self.num_mtp_layers
+        return self.layers[str(self.mtp_start_layer_idx + current_step_idx)](
+            input_ids,
+            positions,
+            previous_hidden_states,
+            inputs_embeds,
+            current_step_idx,
+        )
+
+    def compute_logits(
+        self,
+        hidden_states: torch.Tensor,
+        spec_step_idx: int = 0,
+    ) -> torch.Tensor:
+        current_step_idx = spec_step_idx % self.num_mtp_layers
+        mtp_layer = self.layers[str(self.mtp_start_layer_idx + current_step_idx)]
+        logits = self.logits_processor(mtp_layer.shared_head.head, hidden_states)
+        return logits
+
+
+class KimiK3MTP(nn.Module):
+    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
+        super().__init__()
+        self.config = vllm_config.model_config.hf_text_config
+        self.quant_config = vllm_config.quant_config
+        self.model = KimiK3MultiTokenPredictor(
+            vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
+        )
+        enable_kimi_k3_low_latency_gemm(self, vllm_config.model_config.dtype)
+
+    def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
+        return self.model.embed_input_ids(input_ids)
+
+    def forward(
+        self,
+        input_ids: torch.Tensor | None,
+        positions: torch.Tensor,
+        hidden_states: torch.Tensor,
+        intermediate_tensors: IntermediateTensors | None = None,
+        inputs_embeds: torch.Tensor | None = None,
+        spec_step_idx: int = 0,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        return self.model(
+            input_ids,
+            positions,
+            hidden_states,
+            inputs_embeds,
+            spec_step_idx,
+        )
+
+    def compute_logits(
+        self,
+        hidden_states: torch.Tensor,
+        spec_step_idx: int = 0,
+    ) -> torch.Tensor | None:
+        return self.model.compute_logits(hidden_states, spec_step_idx)
+
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+        # Mirror KimiLinearForCausalLM.load_weights naming: leading-dot shard
+        # names, q_lora-conditional fused QKV, and w1/w2/w3 expert weights.
+        kda_config = self.config.linear_attn_config
+        use_full_rank_gate = bool(
+            kda_config and kda_config.get("use_full_rank_gate", False)
+        )
+        beta_shard_id = 5 if use_full_rank_gate else 3
+        stacked_params_mapping = [
+            # (param_name, shard_name, shard_id)
+            (".in_proj_qkvgfab", ".q_proj", 0),
+            (".in_proj_qkvgfab", ".k_proj", 1),
+            (".in_proj_qkvgfab", ".v_proj", 2),
+            (".in_proj_qkvgfab", ".b_proj", beta_shard_id),
+            (".in_proj_qkvgfab", ".f_a_proj", 4),
+            (".conv1d", ".q_conv1d", 0),
+            (".conv1d", ".k_conv1d", 1),
+            (".conv1d", ".v_conv1d", 2),
+            (".gate_up_proj", ".gate_proj", 0),
+            (".gate_up_proj", ".up_proj", 1),
+        ]
+        if use_full_rank_gate:
+            stacked_params_mapping.append((".in_proj_qkvgfab", ".g_proj", 3))
+        if getattr(self.config, "q_lora_rank", None) is not None:
+            stacked_params_mapping += [
+                (".fused_qkv_a_proj", ".q_a_proj", 0),
+                (".fused_qkv_a_proj", ".kv_a_proj_with_mqa", 1),
+            ]
+
+        use_mega_moe = any(
+            module.use_mega_moe
+            for module in self.modules()
+            if isinstance(module, KimiMoE)
+        )
+        if self.config.is_moe and use_mega_moe:
+            expert_params_mapping = make_kimi_k3_mega_moe_expert_params_mapping(
+                self.config.num_experts
+            )
+        elif self.config.is_moe:
+            expert_params_mapping = fused_moe_make_expert_params_mapping(
+                self,
+                ckpt_gate_proj_name="w1",
+                ckpt_down_proj_name="w2",
+                ckpt_up_proj_name="w3",
+                num_experts=self.config.num_experts,
+            )
+        else:
+            expert_params_mapping = []
+
+        pp_missing_layer_names = get_pp_missing_layer_names(self)
+        params_dict = dict(self.named_parameters())
+        # Under the MXFP4 quant interface the routed experts register unpacked
+        # params (``w13_weight``), while the compressed-tensors checkpoint names
+        # them ``.weight_packed``. Rebind so the expert mapping resolves; scales
+        # already share the ``.weight_scale`` suffix.
+        experts_unpacked = not use_mega_moe and not any(
+            n.endswith("w13_weight_packed") for n in params_dict
+        )
+        loaded_params: set[str] = set()
+        for name, loaded_weight in weights:
+            if "rotary_emb.inv_freq" in name:
+                continue
+            # The multimodal checkpoint prefixes text weights with
+            # ``language_model.``; strip it so names match this draft model's
+            # parameter paths (``model.layers.{i}.``). Non-text weights
+            # (vision_tower, mm_projector, ...) never match a spec layer below.
+            if name.startswith("language_model."):
+                name = name[len("language_model.") :]
+            if experts_unpacked and name.endswith(".weight_packed"):
+                name = name.replace(".weight_packed", ".weight")
+            spec_layer = get_spec_layer_idx_from_weight_name(self.config, name)
+            if spec_layer is None:
+                continue
+            name = self._rewrite_spec_layer_name(spec_layer, name)
+
+            for param_name, weight_name, shard_id in stacked_params_mapping:
+                if weight_name not in name:
+                    continue
+                # Routed experts (``.experts.{i}.w1/w2/w3``) are handled by the
+                # expert mapping below; skip them here. Shared experts
+                # (``.shared_experts.``) use gate/up_proj and fall through.
+                if ".experts." in name:
+                    continue
+                name_mapped = name.replace(weight_name, param_name)
+                # Only take this mapping if the fused destination actually
+                # exists (e.g. QKV fusion is only present when q_lora is used).
+                if name_mapped not in params_dict:
+                    continue
+                if name_mapped in pp_missing_layer_names:
+                    continue
+                name = name_mapped
+                param = params_dict[name]
+                weight_loader = param.weight_loader
+                weight_loader(param, loaded_weight, shard_id)
+                break
+            else:
+                for (
+                    expert_param_name,
+                    expert_weight_name,
+                    expert_id,
+                    expert_shard_id,
+                ) in expert_params_mapping:
+                    if expert_weight_name not in name:
+                        continue
+                    name_mapped = name.replace(expert_weight_name, expert_param_name)
+                    if name_mapped in pp_missing_layer_names:
+                        continue
+                    param = params_dict[name_mapped]
+                    weight_loader = param.weight_loader
+                    weight_loader(
+                        param,
+                        loaded_weight,
+                        name_mapped,
+                        shard_id=expert_shard_id,
+                        expert_id=expert_id,
+                    )
+                    name = name_mapped
+                    break
+                else:
+                    if name.endswith(".bias") and name not in params_dict:
+                        continue
+                    remapped_name = maybe_remap_kv_scale_name(name, params_dict)
+                    if remapped_name is None:
+                        continue
+                    name = remapped_name
+
+                    # The embedding is shared across MTP layers; only the first
+                    # spec layer carries the hoisted (non-".layers") copy.
+                    if spec_layer != self.model.mtp_start_layer_idx and (
+                        ".layers" not in name
+                    ):
+                        continue
+                    if name in pp_missing_layer_names:
+                        continue
+                    # The base model uses an attn-residual scheme whose per-layer
+                    # weights (self_attention_res_*, mlp_res_*) are not used by
+                    # the draft block; such names have no matching parameter and
+                    # are safely skipped.
+                    if name not in params_dict:
+                        continue
+
+                    param = params_dict[name]
+                    weight_loader = getattr(
+                        param, "weight_loader", default_weight_loader
+                    )
+                    weight_loader(param, loaded_weight)
+            loaded_params.add(name)
+
+        # Validate that weights were loaded for each expected MTP layer.
+        loaded_layers: set[int] = set()
+        for param_name in loaded_params:
+            spec_layer = get_spec_layer_idx_from_weight_name(self.config, param_name)
+            if spec_layer is not None:
+                loaded_layers.add(spec_layer)
+        for layer_idx in range(
+            self.model.mtp_start_layer_idx,
+            self.model.mtp_start_layer_idx + self.model.num_mtp_layers,
+        ):
+            if layer_idx not in loaded_layers:
+                raise ValueError(
+                    f"MTP speculative decoding layer {layer_idx} weights "
+                    f"missing from checkpoint. The checkpoint may not include "
+                    f"the MTP layer weights. Use a checkpoint that includes "
+                    f"MTP layer weights, or disable speculative decoding."
+                )
+
+        if use_mega_moe:
+            for module in self.modules():
+                if isinstance(module, KimiMoE) and module.use_mega_moe:
+                    module.experts.finalize_weights()
+
+        return loaded_params
+
+    def _rewrite_spec_layer_name(self, spec_layer: int, name: str) -> str:
+        """Rewrite a checkpoint weight name to this module's parameter path.
+
+        Top-level MTP submodules (enorm/hnorm/eh_proj/shared_head) stay under
+        ``model.layers.{spec_layer}.*``; the shared ``embed_tokens`` is hoisted
+        to ``model.*``; everything else is a transformer-block weight and gets
+        ``.mtp_block`` inserted.
+        """
+        spec_layer_weight_names = [
+            "embed_tokens",
+            "enorm",
+            "hnorm",
+            "eh_proj",
+            "shared_head",
+        ]
+        shared_weight_names = ["embed_tokens"]
+        spec_layer_weight = False
+        shared_weight = False
+        for weight_name in spec_layer_weight_names:
+            if weight_name in name:
+                spec_layer_weight = True
+                if weight_name in shared_weight_names:
+                    shared_weight = True
+                break
+        if not spec_layer_weight:
+            name = name.replace(
+                f"model.layers.{spec_layer}.",
+                f"model.layers.{spec_layer}.mtp_block.",
+            )
+        elif shared_weight:
+            name = name.replace(f"model.layers.{spec_layer}.", "model.")
+        return name

vllm/models/kimi_k3/nvidia/ops/cute_dsl/__init__.py

@@ -0,0 +1,2 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

vllm/models/kimi_k3/nvidia/ops/cute_dsl/latent_moe_tail/__init__.py

@@ -0,0 +1,14 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""CuTe DSL kernels for KimiK3LatentMoETailOp."""
+
+from .allreduce_rmsnorm_reduce_scatter_early_exit import CollectiveKernel
+from .fused_add_multicast_gemm import AdaptiveUpProjectionKernel
+from .lamport_copy import LamportCopyKernel
+
+__all__ = [
+    "AdaptiveUpProjectionKernel",
+    "CollectiveKernel",
+    "LamportCopyKernel",
+]

vllm/models/kimi_k3/nvidia/ops/cute_dsl/latent_moe_tail/fused_add_multicast_skinny_gemm.py

@@ -0,0 +1,438 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""Static-M SIMT skinny GEMM with a shared-add multicast epilogue."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import cuda.bindings.driver as cuda
+import cutlass
+import cutlass.cute as cute
+import torch
+from cutlass import BFloat16, Float32, Int64, const_expr
+from cutlass.cute.runtime import from_dlpack
+
+from .primitives import (
+    CUDAGraphCompatibleWrapper,
+    bf16x2_to_u32,
+    bf16x4_to_packed_u32x2,
+    bf16x8_to_packed_u32x4,
+    sanitize_negative_zero,
+    sanitize_negative_zero_u32,
+    sanitize_negative_zero_u32x2,
+    store_global_u32,
+    store_global_u32x2,
+    store_global_u32x4,
+)
+
+
+@dataclass(frozen=True)
+class SkinnyConfig:
+    block_size: int = 224
+    outputs_per_block: int = 2
+    k_unroll: int = 2
+    vector_width: int = 8
+    prefetch_b_before_pdl: bool = True
+
+
+def config_for_m(num_rows: int, shard_dim: int = 896) -> SkinnyConfig:
+    if shard_dim == 448:
+        if num_rows >= 6:
+            return SkinnyConfig(
+                block_size=224,
+                outputs_per_block=2,
+                k_unroll=1,
+            )
+        outputs_per_block = 2 if num_rows <= 3 else 4
+        return SkinnyConfig(
+            block_size=448,
+            outputs_per_block=outputs_per_block,
+            k_unroll=1,
+        )
+    if num_rows == 1:
+        return SkinnyConfig(outputs_per_block=8, k_unroll=1)
+    return SkinnyConfig(outputs_per_block=4)
+
+
+def _as_cute(tensor: torch.Tensor):
+    return from_dlpack(
+        CUDAGraphCompatibleWrapper(tensor.detach()),
+        assumed_align=16,
+    )
+
+
+class FusedAddMulticastSkinnyGemm:
+    """SIMT GEMM adapted from the existing Skinny GEMM."""
+
+    def __init__(
+        self,
+        *,
+        num_rows: int,
+        hidden_dim: int,
+        config: SkinnyConfig,
+    ) -> None:
+        if config.block_size % 32:
+            raise ValueError("skinny block_size must be a multiple of 32")
+        self.num_rows = num_rows
+        self.hidden_dim = hidden_dim
+        self.block_size = config.block_size
+        self.outputs_per_block = config.outputs_per_block
+        self.k_unroll = config.k_unroll
+        self.vector_width = config.vector_width
+        self.prefetch_b_before_pdl = config.prefetch_b_before_pdl
+        self.num_warps = config.block_size // 32
+
+    @cute.jit
+    def __call__(
+        self,
+        gA: cute.Tensor,
+        gB: cute.Tensor,
+        gShared: cute.Tensor,
+        output_multicast_ptr: Int64,
+        stream: cuda.CUstream,
+    ) -> None:
+        n = cute.size(gB, mode=[0])
+        k = cute.size(gA, mode=[1])
+        copy_a = cute.make_copy_atom(
+            cute.nvgpu.CopyG2ROp(),
+            BFloat16,
+            num_bits_per_copy=self.vector_width * BFloat16.width,
+            load_cache_mode=cute.nvgpu.LoadCacheMode.ALWAYS,
+        )
+        copy_b = cute.make_copy_atom(
+            cute.nvgpu.CopyG2ROp(),
+            BFloat16,
+            num_bits_per_copy=self.vector_width * BFloat16.width,
+            load_cache_mode=cute.nvgpu.LoadCacheMode.STREAMING,
+        )
+        self.kernel(
+            gA,
+            gB,
+            gShared,
+            output_multicast_ptr,
+            k,
+            copy_a,
+            copy_b,
+        ).launch(
+            grid=[cute.ceil_div(n, self.outputs_per_block), 1, 1],
+            block=[self.block_size, 1, 1],
+            smem=(self.num_rows * self.outputs_per_block * self.num_warps * 4),
+            stream=stream,
+            use_pdl=True,
+            min_blocks_per_mp=1,
+        )
+
+    @cute.kernel
+    def kernel(
+        self,
+        gA: cute.Tensor,
+        gB: cute.Tensor,
+        gShared: cute.Tensor,
+        output_multicast_ptr: Int64,
+        k_extent: cutlass.Int32,
+        copy_a: cute.CopyAtom,
+        copy_b: cute.CopyAtom,
+    ) -> None:
+        tidx, _, _ = cute.arch.thread_idx()
+        block_idx, _, _ = cute.arch.block_idx()
+        warp_idx = cute.arch.warp_idx()
+
+        outputs_per_block: cutlass.Constexpr = self.outputs_per_block
+        vector_width: cutlass.Constexpr = self.vector_width
+        block_size: cutlass.Constexpr = self.block_size
+        num_warps: cutlass.Constexpr = self.num_warps
+        num_rows: cutlass.Constexpr = self.num_rows
+
+        acc = cute.make_rmem_tensor(
+            cute.make_layout(
+                (num_rows, outputs_per_block),
+                stride=(outputs_per_block, 1),
+            ),
+            Float32,
+        )
+        acc.fill(0.0)
+
+        n_base = block_idx * outputs_per_block
+        k_tile_size: cutlass.Constexpr = block_size * vector_width
+        num_k_tiles = k_extent // k_tile_size
+        gA_vec = cute.logical_divide(gA, (None, vector_width))
+        gB_vec = cute.logical_divide(gB, (None, vector_width))
+        tA_all = cute.logical_divide(gA_vec, (None, (None, block_size)))
+        tB_all = cute.logical_divide(gB_vec, (None, (None, block_size)))
+        tA = tA_all[None, (None, (tidx, None))]
+
+        a_regs = cute.make_rmem_tensor(
+            cute.make_layout(
+                (num_rows, vector_width),
+                stride=(vector_width, 1),
+            ),
+            BFloat16,
+        )
+        b_regs = cute.make_rmem_tensor(
+            cute.make_layout(
+                (outputs_per_block, vector_width),
+                stride=(vector_width, 1),
+            ),
+            BFloat16,
+        )
+
+        if const_expr(self.prefetch_b_before_pdl):
+            for ni in cutlass.range_constexpr(outputs_per_block):
+                tB = tB_all[n_base + ni, (None, (tidx, None))]
+                cute.copy(copy_b, tB[None, 0], b_regs[ni, None])
+
+        cute.arch.griddepcontrol_wait()
+
+        for mi in cutlass.range_constexpr(num_rows):
+            cute.copy(copy_a, tA[mi, None, 0], a_regs[mi, None])
+        if const_expr(not self.prefetch_b_before_pdl):
+            for ni in cutlass.range_constexpr(outputs_per_block):
+                tB = tB_all[n_base + ni, (None, (tidx, None))]
+                cute.copy(copy_b, tB[None, 0], b_regs[ni, None])
+        for vi in cutlass.range_constexpr(vector_width):
+            for mi in cutlass.range_constexpr(num_rows):
+                for ni in cutlass.range_constexpr(outputs_per_block):
+                    acc[mi, ni] = acc[mi, ni] + a_regs[mi, vi].to(Float32) * b_regs[
+                        ni, vi
+                    ].to(Float32)
+
+        for k_tile in cutlass.range(1, num_k_tiles, unroll=self.k_unroll):
+            for mi in cutlass.range_constexpr(num_rows):
+                cute.copy(
+                    copy_a,
+                    tA[mi, None, k_tile],
+                    a_regs[mi, None],
+                )
+            for ni in cutlass.range_constexpr(outputs_per_block):
+                tB = tB_all[n_base + ni, (None, (tidx, None))]
+                cute.copy(copy_b, tB[None, k_tile], b_regs[ni, None])
+            for vi in cutlass.range_constexpr(vector_width):
+                for mi in cutlass.range_constexpr(num_rows):
+                    for ni in cutlass.range_constexpr(outputs_per_block):
+                        acc[mi, ni] = acc[mi, ni] + a_regs[mi, vi].to(Float32) * b_regs[
+                            ni, vi
+                        ].to(Float32)
+
+        for mi in cutlass.range_constexpr(num_rows):
+            for ni in cutlass.range_constexpr(outputs_per_block):
+                acc[mi, ni] = cute.arch.warp_reduction_sum(acc[mi, ni])
+
+        smem_layout = cute.make_layout(
+            (num_rows, outputs_per_block, num_warps),
+            stride=(outputs_per_block * num_warps, num_warps, 1),
+        )
+        smem = cutlass.utils.SmemAllocator()
+        partials = smem.allocate_tensor(
+            Float32,
+            smem_layout,
+            byte_alignment=16,
+        )
+        with cute.arch.elect_one():
+            for mi in cutlass.range_constexpr(num_rows):
+                for ni in cutlass.range_constexpr(outputs_per_block):
+                    partials[mi, ni, warp_idx] = acc[mi, ni]
+
+        cute.arch.sync_threads()
+        if tidx == 0:
+            fused = cute.make_rmem_tensor(
+                cute.make_layout((outputs_per_block,)),
+                BFloat16,
+            )
+            for mi in cutlass.range_constexpr(num_rows):
+                for ni in cutlass.range_constexpr(outputs_per_block):
+                    total = (
+                        partials[mi, ni, None]
+                        .load()
+                        .reduce(
+                            cute.ReductionOp.ADD,
+                            init_val=Float32(0.0),
+                            reduction_profile=0,
+                        )
+                    )
+                    gemm_value = Float32(total).to(BFloat16)
+                    fused[ni] = (
+                        gemm_value.to(Float32) + gShared[mi, n_base + ni].to(Float32)
+                    ).to(BFloat16)
+                output_offset = Int64((mi * self.hidden_dim + n_base) * 2)
+                if const_expr(outputs_per_block == 2):
+                    packed = sanitize_negative_zero_u32(bf16x2_to_u32(fused.load()))
+                    store_global_u32(
+                        output_multicast_ptr + output_offset,
+                        packed,
+                    )
+                elif const_expr(outputs_per_block == 4):
+                    packed = sanitize_negative_zero_u32x2(
+                        bf16x4_to_packed_u32x2(fused.load())
+                    )
+                    store_global_u32x2(
+                        output_multicast_ptr + output_offset,
+                        packed,
+                    )
+                else:
+                    packed = sanitize_negative_zero(
+                        bf16x8_to_packed_u32x4(fused.load())
+                    )
+                    store_global_u32x4(
+                        output_multicast_ptr + output_offset,
+                        packed,
+                    )
+
+        cute.arch.griddepcontrol_launch_dependents()
+
+
+_COMPILED: dict[tuple[object, ...], object] = {}
+
+
+def compile_kernel(
+    *,
+    num_rows: int,
+    latent_dim: int,
+    hidden_dim: int,
+    shard_dim: int,
+    config: SkinnyConfig,
+):
+    key = (
+        torch.accelerator.current_device_index(),
+        num_rows,
+        latent_dim,
+        hidden_dim,
+        shard_dim,
+        config,
+    )
+    if key in _COMPILED:
+        return _COMPILED[key]
+
+    from cutlass.cute.runtime import make_fake_stream, make_fake_tensor
+
+    if shard_dim % config.outputs_per_block:
+        raise ValueError("shard_dim must be divisible by outputs_per_block")
+    if latent_dim % (config.block_size * config.vector_width):
+        raise ValueError("latent_dim must be divisible by block_size * vector_width")
+
+    a = make_fake_tensor(
+        BFloat16,
+        (num_rows, latent_dim),
+        stride=(latent_dim, 1),
+        assumed_align=16,
+    )
+    b = make_fake_tensor(
+        BFloat16,
+        (shard_dim, latent_dim),
+        stride=(latent_dim, 1),
+        assumed_align=16,
+    )
+    shared = make_fake_tensor(
+        BFloat16,
+        (num_rows, shard_dim),
+        stride=(hidden_dim, 1),
+        assumed_align=16,
+    )
+    compiled = cute.compile(
+        FusedAddMulticastSkinnyGemm(
+            num_rows=num_rows,
+            hidden_dim=hidden_dim,
+            config=config,
+        ),
+        a,
+        b,
+        shared,
+        Int64(0),
+        make_fake_stream(),
+        options="--ptxas-options -maxrregcount=128",
+    )
+    _COMPILED[key] = compiled
+    return compiled
+
+
+class FusedAddMulticastSkinnyGemmKernel:
+    """Buffer-free compiled launcher for one static M."""
+
+    def __init__(
+        self,
+        *,
+        rank: int,
+        tp_size: int,
+        latent_dim: int,
+        hidden_dim: int,
+        num_rows: int,
+    ) -> None:
+        if not 1 <= num_rows <= 8:
+            raise ValueError("skinny backend requires static M in [1, 8]")
+        if hidden_dim % tp_size:
+            raise ValueError("hidden_dim must be divisible by TP size")
+        self.rank = rank
+        self.latent_dim = latent_dim
+        self.hidden_dim = hidden_dim
+        self.shard_dim = hidden_dim // tp_size
+        self.num_rows = num_rows
+        self._skinny = compile_kernel(
+            num_rows=num_rows,
+            latent_dim=latent_dim,
+            hidden_dim=hidden_dim,
+            shard_dim=self.shard_dim,
+            config=config_for_m(num_rows, self.shard_dim),
+        )
+
+    def __call__(
+        self,
+        latent: torch.Tensor,
+        weight: torch.Tensor,
+        shared_shard: torch.Tensor,
+        mailbox: torch.Tensor,
+        mailbox_multicast_ptr: int,
+    ) -> torch.Tensor:
+        device = mailbox.device
+        expected = (
+            (
+                latent,
+                (self.num_rows, self.latent_dim),
+                torch.bfloat16,
+                "latent",
+            ),
+            (
+                weight,
+                (self.shard_dim, self.latent_dim),
+                torch.bfloat16,
+                "weight",
+            ),
+            (
+                shared_shard,
+                (mailbox.shape[1], self.shard_dim),
+                torch.bfloat16,
+                "shared_shard",
+            ),
+            (
+                mailbox,
+                (1, mailbox.shape[1], self.hidden_dim),
+                torch.bfloat16,
+                "mailbox",
+            ),
+        )
+        for tensor, shape, dtype, name in expected:
+            if (
+                tensor.shape != shape
+                or tensor.dtype != dtype
+                or tensor.device != device
+            ):
+                raise ValueError(f"{name} must be CUDA {dtype} {list(shape)}")
+        if (
+            not latent.is_contiguous()
+            or not weight.is_contiguous()
+            or shared_shard.stride() != (self.hidden_dim, 1)
+            or not mailbox.is_contiguous()
+        ):
+            raise ValueError("skinny up-projection inputs have unsupported strides")
+        if mailbox.shape[1] < self.num_rows:
+            raise ValueError("mailbox capacity is smaller than runtime M")
+
+        with torch.accelerator.device_index(device.index):
+            self._skinny(
+                _as_cute(latent),
+                _as_cute(weight),
+                _as_cute(shared_shard[: self.num_rows]),
+                Int64(mailbox_multicast_ptr + self.rank * self.shard_dim * 2),
+                cuda.CUstream(torch.cuda.current_stream(device).cuda_stream),
+            )
+        return mailbox

vllm/models/kimi_k3/nvidia/ops/cute_dsl/latent_moe_tail/lamport_copy.py

@@ -0,0 +1,226 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""Lamport mailbox-to-local copy and reset."""
+
+from __future__ import annotations
+
+import functools
+
+import cuda.bindings.driver as cuda
+import cutlass
+import cutlass.cute as cute
+import torch
+from cutlass.cute.runtime import make_fake_compact_tensor, make_fake_stream
+
+from .primitives import (
+    VEC_BF16,
+    fragment_is_dirty,
+    load_global_u32x4,
+    store_global_u32x4,
+    store_lamport_sentinel_128,
+    to_cute,
+    to_cute_dynamic_m,
+)
+
+
+class LamportCopy:
+    """Consume the local physical copy of an NVLS-multicast mailbox."""
+
+    def __init__(self, hidden_dim: int, ctas: int, threads: int):
+        self.hidden_dim = hidden_dim
+        self.ctas = ctas
+        self.threads = threads
+
+    @cute.jit
+    def __call__(
+        self,
+        symmetric_mailbox: cute.Tensor,
+        local_output: cute.Tensor,
+        m: cutlass.Int32,
+        stream: cuda.CUstream,
+    ):
+        self.kernel(symmetric_mailbox, local_output, m).launch(
+            grid=(self.ctas, 1, 1),
+            block=(self.threads, 1, 1),
+            stream=stream,
+            use_pdl=True,
+        )
+
+    @cute.kernel
+    def kernel(
+        self,
+        symmetric_mailbox: cute.Tensor,
+        local_output: cute.Tensor,
+        m: cutlass.Int32,
+    ):
+        # The CTA may be scheduled early, but mailbox inspection must not pass
+        # the producer GEMM's programmatic completion point.
+        cute.arch.griddepcontrol_wait()
+
+        tidx, _, _ = cute.arch.thread_idx()
+        block, _, _ = cute.arch.block_idx()
+        thread = cutlass.Int64(block * self.threads + tidx)
+        stride = cutlass.Int64(self.ctas * self.threads)
+        fragments = cutlass.Int64(m) * cutlass.Int64(self.hidden_dim // VEC_BF16)
+
+        fragment = thread
+        while fragment < fragments:
+            element = fragment * VEC_BF16
+            source = cute.make_ptr(
+                cutlass.BFloat16,
+                (symmetric_mailbox.iterator + element).llvm_ptr,
+                cute.AddressSpace.gmem,
+                assumed_align=16,
+            )
+            packed = load_global_u32x4(source, volatile=True)
+            while fragment_is_dirty(packed):
+                packed = load_global_u32x4(source, volatile=True)
+
+            destination = cutlass.Int64((local_output.iterator + element).toint())
+            store_global_u32x4(destination, packed, volatile=False)
+            fragment = fragment + stride
+
+        # The returned ordinary tensor is complete. A same-stream successor
+        # may overlap the mailbox cleanup below.
+        cute.arch.griddepcontrol_launch_dependents()
+
+        fragment = thread
+        while fragment < fragments:
+            element = fragment * VEC_BF16
+            source = cute.make_ptr(
+                cutlass.BFloat16,
+                (symmetric_mailbox.iterator + element).llvm_ptr,
+                cute.AddressSpace.gmem,
+                assumed_align=16,
+            )
+            store_lamport_sentinel_128(source)
+            fragment = fragment + stride
+
+
+@functools.cache
+def compile_kernel(
+    hidden_dim: int,
+    max_m: int,
+    ctas: int,
+    threads: int,
+    device_index: int,
+):
+    if hidden_dim <= 0 or hidden_dim % VEC_BF16:
+        raise ValueError("hidden_dim must be a positive multiple of 8")
+    if max_m <= 0:
+        raise ValueError("max_m must be positive")
+    if ctas <= 0:
+        raise ValueError("Lamport copy ctas must be positive")
+    if not 1 <= threads <= 1024:
+        raise ValueError("Lamport copy threads must be in [1, 1024]")
+    with torch.accelerator.device_index(device_index):
+        mailbox = make_fake_compact_tensor(
+            cutlass.BFloat16, (max_m * hidden_dim,), assumed_align=16
+        )
+        output = make_fake_compact_tensor(
+            cutlass.BFloat16,
+            (cute.sym_int32(divisibility=VEC_BF16),),
+            assumed_align=16,
+        )
+        return cute.compile(
+            LamportCopy(hidden_dim, ctas, threads),
+            mailbox,
+            output,
+            cutlass.Int32(max_m),
+            make_fake_stream(),
+        )
+
+
+def launch(
+    symmetric_mailbox: torch.Tensor,
+    local_output: torch.Tensor,
+    *,
+    m: int,
+    hidden_dim: int,
+    max_m: int,
+    ctas: int,
+    threads: int,
+) -> None:
+    if not 1 <= m <= max_m:
+        raise ValueError(f"runtime M={m} must be in [1, {max_m}]")
+    if (
+        symmetric_mailbox.ndim != 3
+        or symmetric_mailbox.dtype != torch.bfloat16
+        or not symmetric_mailbox.is_cuda
+        or not symmetric_mailbox.is_contiguous()
+        or symmetric_mailbox.shape[0] != 1
+        or symmetric_mailbox.shape[2] != hidden_dim
+    ):
+        raise ValueError(
+            f"symmetric_mailbox must be contiguous CUDA BF16 [1,M,{hidden_dim}]"
+        )
+    if symmetric_mailbox.shape[1] != max_m:
+        raise ValueError("symmetric mailbox must retain its full max_m capacity")
+    if (
+        local_output.shape != (1, m, hidden_dim)
+        or local_output.dtype != torch.bfloat16
+        or local_output.device != symmetric_mailbox.device
+        or not local_output.is_contiguous()
+    ):
+        raise ValueError("local_output must be contiguous CUDA BF16 [1,M,H]")
+
+    device_index = symmetric_mailbox.device.index
+    stream = cuda.CUstream(
+        torch.cuda.current_stream(symmetric_mailbox.device).cuda_stream
+    )
+    compile_kernel(hidden_dim, max_m, ctas, threads, device_index)(
+        to_cute(symmetric_mailbox.flatten(), 16),
+        to_cute_dynamic_m(
+            local_output.flatten(),
+            mode=0,
+            assumed_align=16,
+        ),
+        cutlass.Int32(m),
+        stream,
+    )
+
+
+class LamportCopyKernel:
+    """Copy a borrowed symmetric mailbox into a fresh local tensor."""
+
+    def __init__(
+        self,
+        *,
+        hidden_dim: int,
+        max_m: int,
+        ctas: int,
+        threads: int,
+    ) -> None:
+        self.hidden_dim = hidden_dim
+        self.max_m = max_m
+        self.ctas = ctas
+        self.threads = threads
+        compile_kernel(
+            hidden_dim,
+            max_m,
+            ctas,
+            threads,
+            torch.accelerator.current_device_index(),
+        )
+
+    def __call__(self, symmetric_mailbox: torch.Tensor, *, m: int) -> torch.Tensor:
+        if not symmetric_mailbox.is_cuda:
+            raise ValueError("symmetric_mailbox must be a CUDA tensor")
+        device = symmetric_mailbox.device
+        with torch.accelerator.device_index(device.index):
+            output = torch.empty(
+                (1, m, self.hidden_dim),
+                dtype=torch.bfloat16,
+                device=device,
+            )
+            launch(
+                symmetric_mailbox,
+                output,
+                m=m,
+                hidden_dim=self.hidden_dim,
+                max_m=self.max_m,
+                ctas=self.ctas,
+                threads=self.threads,
+            )
+        return output

vllm/models/kimi_k3/nvidia/ops/cute_dsl/latent_moe_tail/primitives.py

@@ -0,0 +1,437 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""Shared CuTe DSL primitives; this module does not define a CUDA kernel."""
+
+from __future__ import annotations
+
+import cutlass
+import cutlass.cute as cute
+import torch
+from cutlass import BFloat16, Float32, Int32, Int64, Uint16, Uint32
+from cutlass._mlir import ir
+from cutlass._mlir.dialects import llvm, vector
+from cutlass.cute.runtime import from_dlpack
+from cutlass.cutlass_dsl import T, dsl_user_op
+
+VEC_BF16 = 8
+PACKED_BYTES = 16
+NUM_LAMPORT_BUFFERS = 3
+NEG_ZERO_F32_BITS = 0x80000000
+NEG_ZERO_BF16_BITS = 0x8000
+
+
+class CUDAGraphCompatibleWrapper:
+    """DLPack view that does not synchronize with the producer stream."""
+
+    def __init__(self, tensor: torch.Tensor):
+        self.tensor = tensor
+
+    def __dlpack__(self, stream=None):
+        return self.tensor.__dlpack__(stream=-1)
+
+    def __dlpack_device__(self):
+        return self.tensor.__dlpack_device__()
+
+
+def to_cute(tensor: torch.Tensor, assumed_align: int = 16) -> cute.Tensor:
+    return from_dlpack(
+        CUDAGraphCompatibleWrapper(tensor.detach()), assumed_align=assumed_align
+    )
+
+
+def to_cute_dynamic_m(
+    tensor: torch.Tensor,
+    *,
+    mode: int,
+    assumed_align: int = 16,
+) -> cute.Tensor:
+    """Expose exactly one compact tensor mode as a runtime shape.
+
+    Model dimensions remain part of the compiled tensor type. Only the token
+    mode is symbolic, so changing M within an Op's capacity reuses the same
+    compiled kernel.
+    """
+
+    return to_cute(tensor, assumed_align).mark_compact_shape_dynamic(
+        mode=mode,
+        stride_order=tensor.dim_order(),
+    )
+
+
+@dsl_user_op
+def load_global_u32x4(
+    pointer: cute.Pointer,
+    *,
+    volatile: cutlass.Constexpr[bool] = False,
+    loc=None,
+    ip=None,
+):
+    """Load one 128-bit fragment as four u32 registers.
+
+    The volatile form is the Lamport polling load.  Marking the asm
+    side-effecting prevents loop-invariant motion and common-subexpression
+    elimination across polling iterations.
+    """
+
+    address = pointer.toint(loc=loc, ip=ip)
+    opcode = "ld.volatile.global.v4.u32" if volatile else "ld.global.v4.u32"
+    out = llvm.inline_asm(
+        llvm.StructType.get_literal([T.i32()] * 4),
+        [address.ir_value(loc=loc, ip=ip)],
+        f"{opcode} {{$0, $1, $2, $3}}, [$4];",
+        "=r,=r,=r,=r,l",
+        has_side_effects=volatile,
+        is_align_stack=False,
+        asm_dialect=llvm.AsmDialect.AD_ATT,
+        loc=loc,
+        ip=ip,
+    )
+    packed = vector.from_elements(
+        ir.VectorType.get([4], T.i32(), loc=loc),
+        [llvm.extractvalue(T.i32(), out, [i], loc=loc, ip=ip) for i in range(4)],
+        loc=loc,
+        ip=ip,
+    )
+    return cute.TensorSSA(packed, 4, Uint32)
+
+
+@dsl_user_op
+def store_global_u32x4(
+    address: Int64,
+    packed,
+    *,
+    volatile: cutlass.Constexpr[bool] = False,
+    loc=None,
+    ip=None,
+) -> None:
+    """Store four packed words to an ordinary or NVLS multicast global VA."""
+
+    words = [packed[i].ir_value(loc=loc, ip=ip) for i in range(4)]
+    opcode = "st.volatile.global.v4.u32" if volatile else "st.global.v4.u32"
+    llvm.inline_asm(
+        None,
+        [address.ir_value(loc=loc, ip=ip), *words],
+        f"{opcode} [$0], {{$1, $2, $3, $4}};",
+        "l,r,r,r,r",
+        has_side_effects=True,
+        is_align_stack=False,
+        asm_dialect=llvm.AsmDialect.AD_ATT,
+        loc=loc,
+        ip=ip,
+    )
+
+
+@dsl_user_op
+def store_global_u32x2(
+    address: Int64,
+    packed,
+    *,
+    loc=None,
+    ip=None,
+) -> None:
+    words = [packed[i].ir_value(loc=loc, ip=ip) for i in range(2)]
+    llvm.inline_asm(
+        None,
+        [address.ir_value(loc=loc, ip=ip), *words],
+        "st.global.v2.u32 [$0], {$1, $2};",
+        "l,r,r",
+        has_side_effects=True,
+        is_align_stack=False,
+        asm_dialect=llvm.AsmDialect.AD_ATT,
+        loc=loc,
+        ip=ip,
+    )
+
+
+@dsl_user_op
+def store_global_u32(
+    address: Int64,
+    word: Uint32,
+    *,
+    loc=None,
+    ip=None,
+) -> None:
+    llvm.inline_asm(
+        None,
+        [
+            address.ir_value(loc=loc, ip=ip),
+            word.ir_value(loc=loc, ip=ip),
+        ],
+        "st.global.u32 [$0], $1;",
+        "l,r",
+        has_side_effects=True,
+        is_align_stack=False,
+        asm_dialect=llvm.AsmDialect.AD_ATT,
+        loc=loc,
+        ip=ip,
+    )
+
+
+@dsl_user_op
+def store_lamport_sentinel_128(pointer: cute.Pointer, *, loc=None, ip=None) -> None:
+    """Reset one Lamport fragment to four FP32 negative-zero bit patterns."""
+
+    address = pointer.toint(loc=loc, ip=ip)
+    value = Uint32(NEG_ZERO_F32_BITS).ir_value(loc=loc, ip=ip)
+    llvm.inline_asm(
+        None,
+        [address.ir_value(loc=loc, ip=ip), value, value, value, value],
+        "st.global.v4.u32 [$0], {$1, $2, $3, $4};",
+        "l,r,r,r,r",
+        has_side_effects=True,
+        is_align_stack=False,
+        asm_dialect=llvm.AsmDialect.AD_ATT,
+        loc=loc,
+        ip=ip,
+    )
+
+
+@dsl_user_op
+def red_async_release_gpu_add_u32(
+    pointer: cute.Pointer, value: Uint32, *, loc=None, ip=None
+) -> None:
+    """The exact SM100 arrival primitive used by upstream LamportFlags."""
+
+    address = pointer.toint(loc=loc, ip=ip)
+    llvm.inline_asm(
+        None,
+        [
+            address.ir_value(loc=loc, ip=ip),
+            value.ir_value(loc=loc, ip=ip),
+        ],
+        "red.async.release.global.gpu.add.u32 [$0], $1;",
+        "l,r",
+        has_side_effects=True,
+        is_align_stack=False,
+        asm_dialect=llvm.AsmDialect.AD_ATT,
+        loc=loc,
+        ip=ip,
+    )
+
+
+@dsl_user_op
+def load_volatile_u32(pointer: cute.Pointer, *, loc=None, ip=None) -> Uint32:
+    address = pointer.toint(loc=loc, ip=ip)
+    return Uint32(
+        llvm.inline_asm(
+            T.i32(),
+            [address.ir_value(loc=loc, ip=ip)],
+            "ld.volatile.global.u32 $0, [$1];",
+            "=r,l",
+            has_side_effects=True,
+            is_align_stack=False,
+            asm_dialect=llvm.AsmDialect.AD_ATT,
+            loc=loc,
+            ip=ip,
+        )
+    )
+
+
+@dsl_user_op
+def map_shared_to_peer(
+    smem_ptr: cute.Pointer,
+    peer_rank: Int32,
+    *,
+    loc=None,
+    ip=None,
+) -> Int32:
+    """Map a local shared-memory slot to the same slot in a peer CTA."""
+
+    smem_address = smem_ptr.toint(loc=loc, ip=ip).ir_value()
+    return Int32(
+        llvm.inline_asm(
+            T.i32(),
+            [smem_address, peer_rank.ir_value(loc=loc, ip=ip)],
+            "mapa.shared::cluster.u32 $0, $1, $2;",
+            "=r,r,r",
+            has_side_effects=False,
+            is_align_stack=False,
+            asm_dialect=llvm.AsmDialect.AD_ATT,
+            loc=loc,
+            ip=ip,
+        )
+    )
+
+
+@dsl_user_op
+def store_shared_cluster_f32(
+    remote_address: Int32,
+    value: Float32,
+    *,
+    loc=None,
+    ip=None,
+) -> None:
+    llvm.inline_asm(
+        None,
+        [
+            remote_address.ir_value(loc=loc, ip=ip),
+            value.ir_value(loc=loc, ip=ip),
+        ],
+        "st.shared::cluster.f32 [$0], $1;",
+        "r,f",
+        has_side_effects=True,
+        is_align_stack=False,
+        asm_dialect=llvm.AsmDialect.AD_ATT,
+        loc=loc,
+        ip=ip,
+    )
+
+
+@dsl_user_op
+def packed_u32x4_to_bf16x8(packed, *, loc=None, ip=None):
+    target_type = ir.VectorType.get([VEC_BF16], BFloat16.mlir_type, loc=loc)
+    values = llvm.bitcast(
+        target_type,
+        packed.ir_value(loc=loc, ip=ip),
+        loc=loc,
+        ip=ip,
+    )
+    return cute.TensorSSA(values, VEC_BF16, BFloat16)
+
+
+@dsl_user_op
+def bf16x8_to_packed_u32x4(values, *, loc=None, ip=None):
+    target_type = ir.VectorType.get([4], T.i32(), loc=loc)
+    packed = llvm.bitcast(
+        target_type,
+        values.ir_value(loc=loc, ip=ip),
+        loc=loc,
+        ip=ip,
+    )
+    return cute.TensorSSA(packed, 4, Uint32)
+
+
+@dsl_user_op
+def bf16x4_to_packed_u32x2(values, *, loc=None, ip=None):
+    target_type = ir.VectorType.get([2], T.i32(), loc=loc)
+    packed = llvm.bitcast(
+        target_type,
+        values.ir_value(loc=loc, ip=ip),
+        loc=loc,
+        ip=ip,
+    )
+    return cute.TensorSSA(packed, 2, Uint32)
+
+
+@dsl_user_op
+def bf16x2_to_u32(values, *, loc=None, ip=None):
+    packed = llvm.bitcast(
+        T.i32(),
+        values.ir_value(loc=loc, ip=ip),
+        loc=loc,
+        ip=ip,
+    )
+    return Uint32(packed)
+
+
+@cute.jit
+def sanitize_negative_zero_u32(word):
+    low = Uint16(word & Uint32(0xFFFF))
+    high = Uint16(word >> Uint32(16))
+    if low == Uint16(NEG_ZERO_BF16_BITS):
+        word = word & Uint32(0xFFFF0000)
+    if high == Uint16(NEG_ZERO_BF16_BITS):
+        word = word & Uint32(0x0000FFFF)
+    return word
+
+
+@cute.jit
+def sanitize_negative_zero_u32x2(packed):
+    result = cute.make_rmem_tensor(cute.make_layout((2,)), Uint32)
+    for i in cutlass.range_constexpr(2):
+        result[i] = sanitize_negative_zero_u32(packed[i])
+    return result.load()
+
+
+@cute.jit
+def sanitize_negative_zero(packed):
+    """Turn real BF16 -0 into +0 so it cannot equal the empty sentinel."""
+
+    result = cute.make_rmem_tensor(cute.make_layout((4,)), Uint32)
+    for i in cutlass.range_constexpr(4):
+        word = packed[i]
+        low = Uint16(word & Uint32(0xFFFF))
+        high = Uint16(word >> Uint32(16))
+        if low == Uint16(NEG_ZERO_BF16_BITS):
+            word = word & Uint32(0xFFFF0000)
+        if high == Uint16(NEG_ZERO_BF16_BITS):
+            word = word & Uint32(0x0000FFFF)
+        result[i] = word
+    return result.load()
+
+
+@cute.jit
+def fragment_is_dirty(packed):
+    """Bit-exact upstream sentinel check: one comparison per 32-bit word."""
+
+    dirty = packed[0] == Uint32(NEG_ZERO_F32_BITS)
+    for i in cutlass.range_constexpr(1, 4):
+        dirty = dirty | (packed[i] == Uint32(NEG_ZERO_F32_BITS))
+    return dirty
+
+
+@cute.jit
+def warp_sum_specialized(
+    value: Float32,
+    warp_idx: Int32,
+    lane: Int32,
+    warps: cutlass.Constexpr[int],
+    last_warp_lanes: cutlass.Constexpr[int],
+    last_warp_mask: cutlass.Constexpr[int],
+) -> Float32:
+    """Warp sum supporting a compile-time partial final warp."""
+
+    if warp_idx == Int32(warps - 1) and cutlass.const_expr(last_warp_lanes < 32):
+        for offset in cutlass.range_constexpr(16, 0, -1):
+            # range_constexpr does not provide powers-of-two stepping.
+            if cutlass.const_expr(offset in (16, 8, 4, 2, 1)):
+                other = cute.arch.shuffle_sync_bfly(
+                    value,
+                    offset=offset,
+                    mask=last_warp_mask,
+                    mask_and_clamp=31,
+                )
+                if (lane ^ Int32(offset)) < Int32(last_warp_lanes):
+                    value = value + other
+    else:
+        for offset in cutlass.range_constexpr(16, 0, -1):
+            if cutlass.const_expr(offset in (16, 8, 4, 2, 1)):
+                value = value + cute.arch.shuffle_sync_bfly(
+                    value,
+                    offset=offset,
+                    mask=-1,
+                    mask_and_clamp=31,
+                )
+    return value
+
+
+@cute.jit
+def block_sum_specialized(
+    value: Float32,
+    warp_sums: cute.Tensor,
+    tidx: Int32,
+    warps: cutlass.Constexpr[int],
+    last_warp_lanes: cutlass.Constexpr[int],
+    last_warp_mask: cutlass.Constexpr[int],
+) -> Float32:
+    """Upstream-equivalent FP32 block reduction."""
+
+    lane = cute.arch.lane_idx()
+    warp_idx = cute.arch.warp_idx()
+    value = warp_sum_specialized(
+        value, warp_idx, lane, warps, last_warp_lanes, last_warp_mask
+    )
+    if lane == 0:
+        warp_sums[warp_idx] = value
+    cute.arch.barrier()
+
+    block_sum = Float32(0.0)
+    if warp_idx == 0:
+        if lane < Int32(warps):
+            block_sum = warp_sums[lane]
+        block_sum = cute.arch.warp_reduction_sum(block_sum)
+        if lane == 0:
+            warp_sums[0] = block_sum
+    cute.arch.barrier()
+    return warp_sums[0]

vllm/models/kimi_k3/nvidia/ops/fused_mla_key_concat_kv_cache.py

@@ -0,0 +1,244 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Fused MLA prefill and decode epilogues for Kimi-K3.
+
+Thin wrappers over the CUDA ops in
+``csrc/libtorch_stable/fused_kimi_k3_mla_key_concat_kv_cache_kernel.cu``, which
+mirror ``fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_{bf16,fp8}_insert``.
+
+- ``fused_mla_key_concat_kv_cache_insert`` (bf16): optionally apply RoPE,
+  concat the full per-head key ``[k_nope | k_pe]`` into ``k_out``, and insert
+  the latent ``[kv_c_normed | k_pe]`` into the paged cache.
+- ``fused_mla_qkv_quant_kv_cache_fp8_insert`` (fp8): additionally quantize
+  ``q``/``k``/``v`` to E4M3 with ``q_scale`` / ``k_scale`` / ``v_scale`` (the
+  cache shares ``k_scale``, as in ``concat_and_cache_mla``).
+
+The optional ``positions`` / ``cos_sin_cache`` pair enables GPT-J-style RoPE
+inside the epilogue. Omitting both keeps the K3 NoPE fast path. The kernels use
+Programmatic Dependent Launch to overlap the tail of the producing GEMMs on
+sm_90+.
+"""
+
+import torch
+
+
+def fused_mla_key_concat_kv_cache_insert(
+    q: torch.Tensor,  # [Tp, H, qk_head_dim], RoPE is applied in place
+    k_nope: torch.Tensor,  # [Tp, H, qk_nope_head_dim]
+    k_pe: torch.Tensor,  # [Tp, qk_rope_head_dim] or [Tp, 1, qk_rope_head_dim]
+    kv_c_normed: torch.Tensor,  # [Tp, kv_lora_rank]
+    kv_cache: torch.Tensor,  # [num_blocks, block_size, kv_lora_rank + rope]
+    slot_mapping: torch.Tensor,  # [Tp] int64
+    positions: torch.Tensor | None = None,  # [Tp] int64
+    cos_sin_cache: torch.Tensor | None = None,  # [max_position, rope]
+) -> torch.Tensor:
+    """Apply optional RoPE, concat K, and insert the paged latent (bf16).
+
+    Returns the full key ``[Tp, H, qk_nope_head_dim + qk_rope_head_dim]``;
+    optionally rotates ``q`` and writes ``kv_cache`` in place.
+    """
+    k_pe = k_pe.reshape(k_pe.shape[0], -1)
+    tp, num_heads, qk_nope_head_dim = k_nope.shape
+    qk_head_dim = qk_nope_head_dim + k_pe.shape[1]
+    k_out = torch.empty(
+        (tp, num_heads, qk_head_dim), dtype=k_nope.dtype, device=k_nope.device
+    )
+    if tp == 0:
+        return k_out
+    torch.ops._C.fused_kimi_k3_mla_key_concat_kv_cache_insert(
+        q,
+        k_nope,
+        k_pe,
+        kv_c_normed,
+        k_out,
+        kv_cache,
+        slot_mapping,
+        kv_cache.shape[1],
+        positions,
+        cos_sin_cache,
+    )
+    return k_out
+
+
+def fused_mla_key_concat_ds_mla_insert(
+    q: torch.Tensor,  # [Tp, H, qk_head_dim], RoPE is applied in place
+    k_nope: torch.Tensor,  # [Tp, H, qk_nope_head_dim]
+    k_pe: torch.Tensor,  # [Tp, qk_rope_head_dim] or [Tp, 1, qk_rope_head_dim]
+    kv_c_normed: torch.Tensor,  # [Tp, kv_lora_rank]
+    kv_cache: torch.Tensor,  # [num_blocks, block_size, 656] uint8 (fp8_ds_mla)
+    slot_mapping: torch.Tensor,  # [Tp] int64
+    positions: torch.Tensor | None = None,  # [Tp] int64
+    cos_sin_cache: torch.Tensor | None = None,  # [max_position, rope]
+) -> torch.Tensor:
+    """Concat full K (bf16) and insert the latent in the fp8_ds_mla layout.
+
+    The cache uses DeepSeek's 656-byte block-scaled layout (NoPE in 4 tiles of
+    128 with per-tile dynamic fp8 scales, RoPE as bf16) -- self-scaling, so no
+    scale argument. Returns the bf16 full key; optionally rotates ``q`` and
+    writes ``kv_cache`` in place.
+    """
+    k_pe = k_pe.reshape(k_pe.shape[0], -1)
+    tp, num_heads, qk_nope_head_dim = k_nope.shape
+    qk_head_dim = qk_nope_head_dim + k_pe.shape[1]
+    k_out = torch.empty(
+        (tp, num_heads, qk_head_dim), dtype=k_nope.dtype, device=k_nope.device
+    )
+    if tp == 0:
+        return k_out
+    torch.ops._C.fused_kimi_k3_mla_key_concat_ds_mla_insert(
+        q,
+        k_nope,
+        k_pe,
+        kv_c_normed,
+        k_out,
+        kv_cache,
+        slot_mapping,
+        kv_cache.shape[1],
+        positions,
+        cos_sin_cache,
+    )
+    return k_out
+
+
+def fused_mla_qkv_quant_kv_cache_fp8_insert(
+    q: torch.Tensor,  # [Tp, H, qk_head_dim]
+    k_nope: torch.Tensor,  # [Tp, H, qk_nope_head_dim]
+    k_pe: torch.Tensor,  # [Tp, qk_rope_head_dim] or [Tp, 1, qk_rope_head_dim]
+    kv_c_normed: torch.Tensor,  # [Tp, kv_lora_rank]
+    v: torch.Tensor,  # [Tp, H, v_head_dim]
+    kv_cache: torch.Tensor,  # [num_blocks, block_size, kv_lora_rank + rope] fp8
+    slot_mapping: torch.Tensor,  # [Tp] int64
+    q_scale_inv: torch.Tensor,  # scalar fp32, 1 / q scale (attention query)
+    k_scale_inv: torch.Tensor,  # scalar fp32, 1 / k scale (attention key)
+    v_scale_inv: torch.Tensor,  # scalar fp32, 1 / v scale (attention value)
+    cache_scale_inv: torch.Tensor,  # scalar fp32, 1 / kv scale (cache latent)
+    positions: torch.Tensor | None = None,  # [Tp] int64
+    cos_sin_cache: torch.Tensor | None = None,  # [max_position, rope]
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    """Quantize q/k/v to fp8 and insert the fp8 latent into the paged cache.
+
+    The attention key ``k_fp8`` and the cache latent use *separate* scales
+    (``k_scale_inv`` vs ``cache_scale_inv``): the cache must be quantized with
+    ``_k_scale`` (read back by decode / context), while the prefill attention
+    q/k/v currently stay unscaled (the prefill flash path does not dequantize).
+
+    Returns ``(q_fp8, k_fp8, v_fp8)``; writes the fp8 ``kv_cache`` in place.
+    """
+    k_pe = k_pe.reshape(k_pe.shape[0], -1)
+    tp, num_heads, _ = q.shape
+    qk_head_dim = q.shape[2]
+    v_head_dim = v.shape[2]
+    fp8 = torch.float8_e4m3fn
+    q_fp8 = torch.empty((tp, num_heads, qk_head_dim), dtype=fp8, device=q.device)
+    k_fp8 = torch.empty((tp, num_heads, qk_head_dim), dtype=fp8, device=q.device)
+    v_fp8 = torch.empty((tp, num_heads, v_head_dim), dtype=fp8, device=q.device)
+    if tp == 0:
+        return q_fp8, k_fp8, v_fp8
+    torch.ops._C.fused_kimi_k3_mla_qkv_quant_kv_cache_fp8_insert(
+        q,
+        k_nope,
+        k_pe,
+        kv_c_normed,
+        v,
+        q_fp8,
+        k_fp8,
+        v_fp8,
+        kv_cache,
+        slot_mapping,
+        q_scale_inv,
+        k_scale_inv,
+        v_scale_inv,
+        cache_scale_inv,
+        kv_cache.shape[1],
+        positions,
+        cos_sin_cache,
+    )
+    return q_fp8, k_fp8, v_fp8
+
+
+def fused_mla_decode_q_concat_kv_cache_insert(
+    ql_nope: torch.Tensor,  # [B, H, kv_lora_rank]  (BMM1 output, absorbed q)
+    q_pe: torch.Tensor,  # [B, H, qk_rope_head_dim]
+    kv_c_normed: torch.Tensor,  # [B, kv_lora_rank]
+    k_pe: torch.Tensor,  # [B, qk_rope_head_dim] or [B, 1, qk_rope_head_dim]
+    kv_cache: torch.Tensor,  # [num_blocks, block_size, entry]
+    slot_mapping: torch.Tensor,  # [B] int64
+    *,
+    ds_mla: bool = False,
+    q_scale_inv: torch.Tensor | None = None,  # scalar fp32, 1 / q scale
+    cache_scale_inv: torch.Tensor | None = None,  # scalar fp32, 1 / kv scale
+    positions: torch.Tensor | None = None,  # [B] int64
+    cos_sin_cache: torch.Tensor | None = None,  # [max_position, rope]
+) -> torch.Tensor:
+    """Concat the latent decode query ``mqa_q = [ql_nope | q_pe]`` and insert the
+    latent ``[kv_c_normed | k_pe]`` into the paged cache, in one launch (runs
+    right before ``forward_mqa``).
+
+    Dispatched by cache format:
+      - bf16          -> bf16 mqa_q, bf16 cache
+      - plain fp8     -> fp8 mqa_q (q_scale_inv), fp8 cache (cache_scale_inv)
+      - fp8_ds_mla    -> bf16 mqa_q, 656B block-scaled cache
+
+    Returns ``mqa_q`` of shape ``[B, H, kv_lora_rank + qk_rope_head_dim]``;
+    writes ``kv_cache`` in place.
+    """
+    k_pe = k_pe.reshape(k_pe.shape[0], -1)
+    b, num_heads, kv_lora_rank = ql_nope.shape
+    entry = kv_lora_rank + q_pe.shape[-1]
+    fp8_q = q_scale_inv is not None
+    out_dtype = torch.float8_e4m3fn if fp8_q else ql_nope.dtype
+    mqa_q = torch.empty((b, num_heads, entry), dtype=out_dtype, device=ql_nope.device)
+    if b == 0:
+        return mqa_q
+
+    if ds_mla:
+        cache = (
+            kv_cache if kv_cache.dtype == torch.uint8 else kv_cache.view(torch.uint8)
+        )
+        torch.ops._C.fused_kimi_k3_mla_decode_q_concat_ds_mla_insert(
+            ql_nope,
+            q_pe,
+            kv_c_normed,
+            k_pe,
+            mqa_q,
+            cache,
+            slot_mapping,
+            cache.shape[1],
+            positions,
+            cos_sin_cache,
+        )
+    elif fp8_q:
+        assert cache_scale_inv is not None, "fp8 decode requires cache_scale_inv"
+        cache = (
+            kv_cache
+            if kv_cache.dtype == torch.float8_e4m3fn
+            else kv_cache.view(torch.float8_e4m3fn)
+        )
+        torch.ops._C.fused_kimi_k3_mla_decode_q_concat_kv_cache_fp8_insert(
+            ql_nope,
+            q_pe,
+            kv_c_normed,
+            k_pe,
+            mqa_q,
+            cache,
+            slot_mapping,
+            q_scale_inv,
+            cache_scale_inv,
+            cache.shape[1],
+            positions,
+            cos_sin_cache,
+        )
+    else:
+        torch.ops._C.fused_kimi_k3_mla_decode_q_concat_kv_cache_insert(
+            ql_nope,
+            q_pe,
+            kv_c_normed,
+            k_pe,
+            mqa_q,
+            kv_cache,
+            slot_mapping,
+            kv_cache.shape[1],
+            positions,
+            cos_sin_cache,
+        )
+    return mqa_q

vllm/models/kimi_k3/nvidia/ops/latent_moe_tail.py

@@ -0,0 +1,265 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from dataclasses import dataclass
+from functools import partial
+from typing import ClassVar
+
+import torch
+import torch.distributed as dist
+
+from vllm.distributed import get_tp_group
+from vllm.model_executor.warmup.cutedsl_warmup import (
+    CuTeDSLCompileUnit,
+    register_cutedsl_warmup_provider,
+)
+
+_MAX_NUM_TOKENS = 16
+_SKINNY_MAX_NUM_TOKENS = 5
+_MMA_TILER_MN = (64, 32)
+_GEMM_CLUSTER_MN = (1, 8)
+_B_PRIME_STAGES = 2
+_COLLECTIVE_TOKEN_CTAS = 8
+_LAMPORT_COPY_CTAS = 32
+_LAMPORT_COPY_THREADS = 224
+_SUPPORTED_TP_SIZES = (8, 16)
+
+
+@dataclass(frozen=True)
+class KimiK3LatentMoETailContract:
+    tp_group_id: int
+    tp_size: int
+    device: torch.device
+    dtype: torch.dtype
+    hidden_size: int
+    latent_size: int
+    max_num_tokens: int
+    rms_eps: float
+
+
+class KimiK3LatentMoETailOp:
+    """Process-wide cached K3 latent-MoE tail implementation."""
+
+    _instances: ClassVar[
+        dict[KimiK3LatentMoETailContract, "KimiK3LatentMoETailOp"]
+    ] = {}
+
+    @classmethod
+    def _contract_and_group(
+        cls,
+        *,
+        hidden_size: int,
+        latent_size: int,
+        dtype: torch.dtype,
+        device: torch.device,
+        rms_eps: float,
+    ) -> tuple[KimiK3LatentMoETailContract, dist.ProcessGroup]:
+        tp = get_tp_group()
+        group = tp.device_group
+        device = torch.device(device)
+        tp_device = torch.device(tp.device)
+        if device != tp_device:
+            raise ValueError(
+                f"Input device {device} does not match TP device {tp_device}."
+            )
+        return (
+            KimiK3LatentMoETailContract(
+                tp_group_id=id(group),
+                tp_size=dist.get_world_size(group),
+                device=device,
+                dtype=dtype,
+                hidden_size=hidden_size,
+                latent_size=latent_size,
+                max_num_tokens=_MAX_NUM_TOKENS,
+                rms_eps=float(rms_eps),
+            ),
+            group,
+        )
+
+    @classmethod
+    def initialize(
+        cls,
+        *,
+        hidden_size: int,
+        latent_size: int,
+        dtype: torch.dtype,
+        device: torch.device,
+        rms_eps: float,
+    ) -> "KimiK3LatentMoETailOp":
+        contract, group = cls._contract_and_group(
+            hidden_size=hidden_size,
+            latent_size=latent_size,
+            dtype=dtype,
+            device=device,
+            rms_eps=rms_eps,
+        )
+        op = cls._instances.get(contract)
+        if op is None:
+            op = cls(contract, group)
+            cls._instances[contract] = op
+        return op
+
+    def __init__(
+        self,
+        contract: KimiK3LatentMoETailContract,
+        group: dist.ProcessGroup,
+    ) -> None:
+        if contract.tp_size not in _SUPPORTED_TP_SIZES:
+            raise ValueError(
+                "K3 latent-MoE tail fusion requires TP in "
+                f"{_SUPPORTED_TP_SIZES}, got {contract.tp_size}."
+            )
+        if contract.device.type != "cuda":
+            raise ValueError("K3 latent-MoE tail fusion requires a CUDA device.")
+        if torch.cuda.get_device_capability(contract.device)[0] != 10:
+            raise ValueError("K3 latent-MoE tail fusion requires SM100.")
+        if contract.dtype != torch.bfloat16:
+            raise ValueError("K3 latent-MoE tail fusion requires bfloat16.")
+        if (contract.hidden_size, contract.latent_size) != (7168, 3584):
+            raise ValueError(
+                "K3 latent-MoE tail fusion requires hidden_size=7168 and "
+                "latent_size=3584."
+            )
+
+        self.contract = contract
+        self.rank = dist.get_rank(group)
+        from .cute_dsl.latent_moe_tail import (
+            AdaptiveUpProjectionKernel,
+            CollectiveKernel,
+            LamportCopyKernel,
+        )
+
+        with torch.accelerator.device_index(contract.device.index):
+            self._collective = CollectiveKernel(
+                group=group,
+                rank=self.rank,
+                tp_size=contract.tp_size,
+                latent_dim=contract.latent_size,
+                hidden_dim=contract.hidden_size,
+                max_m=contract.max_num_tokens,
+                max_token_ctas=_COLLECTIVE_TOKEN_CTAS,
+                rms_eps=contract.rms_eps,
+                fp32_internal=False,
+            )
+            self._up_projection = AdaptiveUpProjectionKernel(
+                group=group,
+                rank=self.rank,
+                tp_size=contract.tp_size,
+                latent_dim=contract.latent_size,
+                hidden_dim=contract.hidden_size,
+                max_m=contract.max_num_tokens,
+                skinny_max_m=_SKINNY_MAX_NUM_TOKENS,
+                mma_tiler_mn=_MMA_TILER_MN,
+                cluster_shape_mn=_GEMM_CLUSTER_MN,
+                b_prime_stages=_B_PRIME_STAGES,
+            )
+            self._lamport_copy = LamportCopyKernel(
+                hidden_dim=contract.hidden_size,
+                max_m=contract.max_num_tokens,
+                ctas=_LAMPORT_COPY_CTAS,
+                threads=_LAMPORT_COPY_THREADS,
+            )
+        register_cutedsl_warmup_provider(self)
+
+    def get_cutedsl_warmup_compile_units(self) -> tuple[CuTeDSLCompileUnit, ...]:
+        contract = self.contract
+        skinny_units = tuple(
+            CuTeDSLCompileUnit(
+                name="K3 latent MoE tail Skinny up-projection",
+                key=(
+                    "k3-latent-moe-tail-skinny-up-projection",
+                    contract,
+                    m,
+                ),
+                compile=partial(self._up_projection.compile_skinny, m),
+            )
+            for m in range(1, self._up_projection.skinny_max_m + 1)
+        )
+        return skinny_units + (
+            CuTeDSLCompileUnit(
+                name="K3 latent MoE tail dynamic up-projection",
+                key=(
+                    "k3-latent-moe-tail-dynamic-up-projection",
+                    contract,
+                    _MMA_TILER_MN,
+                    _GEMM_CLUSTER_MN,
+                    _B_PRIME_STAGES,
+                ),
+                compile=self._up_projection.compile_dynamic,
+            ),
+        )
+
+    def __call__(
+        self,
+        routed_output: torch.Tensor,
+        shared_output: torch.Tensor,
+        rms_weight: torch.Tensor,
+        up_weight: torch.Tensor,
+    ) -> torch.Tensor:
+        self._validate_inputs(
+            routed_output,
+            shared_output,
+            rms_weight,
+            up_weight,
+        )
+        self._up_projection.ensure_compiled(routed_output.shape[0])
+        latent, shared_shard = self._collective(
+            routed_output,
+            shared_output,
+            rms_weight,
+        )
+        local_hidden_size = self.contract.hidden_size // self.contract.tp_size
+        local_up_weight = up_weight.narrow(
+            0,
+            self.rank * local_hidden_size,
+            local_hidden_size,
+        )
+        mailbox = self._up_projection(
+            latent,
+            local_up_weight,
+            shared_shard,
+        )
+        return self._lamport_copy(
+            mailbox,
+            m=routed_output.shape[0],
+        ).squeeze(0)
+
+    def _validate_inputs(
+        self,
+        routed_output: torch.Tensor,
+        shared_output: torch.Tensor,
+        rms_weight: torch.Tensor,
+        up_weight: torch.Tensor,
+    ) -> None:
+        contract = self.contract
+        if routed_output.ndim != 2:
+            raise ValueError("routed_output must be a 2D tensor.")
+        num_tokens = routed_output.shape[0]
+        if routed_output.shape != (num_tokens, contract.latent_size):
+            raise ValueError(
+                f"routed_output must have shape [M, {contract.latent_size}]."
+            )
+        if shared_output.shape != (num_tokens, contract.hidden_size):
+            raise ValueError(
+                f"shared_output must have shape [M, {contract.hidden_size}]."
+            )
+        if rms_weight.shape != (contract.latent_size,):
+            raise ValueError(f"rms_weight must have shape [{contract.latent_size}].")
+        if up_weight.shape != (contract.hidden_size, contract.latent_size):
+            raise ValueError(
+                "up_weight must have shape "
+                f"[{contract.hidden_size}, {contract.latent_size}]."
+            )
+        if not 1 <= num_tokens <= contract.max_num_tokens:
+            raise ValueError(
+                "K3 latent-MoE tail fusion requires between 1 and "
+                f"{contract.max_num_tokens} tokens."
+            )
+
+        tensors = (routed_output, shared_output, rms_weight, up_weight)
+        if any(tensor.device != contract.device for tensor in tensors):
+            raise ValueError("All inputs must be on the contract device.")
+        if any(tensor.dtype != contract.dtype for tensor in tensors):
+            raise ValueError("All inputs must use the contract dtype.")
+        if any(not tensor.is_contiguous() for tensor in tensors):
+            raise ValueError("All inputs must be contiguous.")

vllm/models/kimi_k3/nvidia/ops/sequence_parallel.py

@@ -0,0 +1,73 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import torch
+
+from vllm.distributed import (
+    get_tensor_model_parallel_rank,
+    get_tensor_model_parallel_world_size,
+    get_tp_group,
+    tensor_model_parallel_all_gather,
+    tensor_model_parallel_reduce_scatter,
+)
+
+
+def _custom_collective(
+    name: str,
+    x: torch.Tensor,
+) -> torch.Tensor | None:
+    device_communicator = get_tp_group().device_communicator
+    if device_communicator is None:
+        return None
+    collective = getattr(device_communicator, name, None)
+    if collective is None:
+        return None
+    return collective(x)
+
+
+def sp_all_gather(x: torch.Tensor) -> torch.Tensor:
+    output = _custom_collective("custom_all_gather", x)
+    if output is not None:
+        return output
+    return tensor_model_parallel_all_gather(x, 0)
+
+
+def sp_reduce_scatter(x: torch.Tensor) -> torch.Tensor:
+    assert x.ndim == 2
+    tp_size = get_tensor_model_parallel_world_size()
+    sp_pad = (-x.shape[0]) % tp_size
+    if sp_pad > 0:
+        x = torch.nn.functional.pad(x, (0, 0, 0, sp_pad))
+    output = _custom_collective("custom_reduce_scatter", x)
+    if output is not None:
+        return output
+    return tensor_model_parallel_reduce_scatter(x, 0)
+
+
+def sp_shard(x: torch.Tensor) -> torch.Tensor:
+    assert x.ndim == 2
+    tp_size = get_tensor_model_parallel_world_size()
+    tp_rank = get_tensor_model_parallel_rank()
+    sp_pad = (-x.shape[0]) % tp_size
+    if sp_pad > 0:
+        x = torch.nn.functional.pad(x, (0, 0, 0, sp_pad))
+    chunk = x.shape[0] // tp_size
+    return x[tp_rank * chunk : (tp_rank + 1) * chunk]
+
+
+def sp_padding_mask(
+    is_padding: torch.Tensor | None,
+    hidden_states: torch.Tensor,
+) -> torch.Tensor:
+    num_tokens = hidden_states.shape[0]
+    if is_padding is None:
+        is_padding = hidden_states.new_zeros(num_tokens, dtype=torch.bool)
+    assert is_padding.shape[0] == num_tokens
+
+    tp_size = get_tensor_model_parallel_world_size()
+    sp_pad = (-num_tokens) % tp_size
+    if sp_pad > 0:
+        is_padding = torch.nn.functional.pad(is_padding, (0, sp_pad), value=True)
+    chunk = is_padding.shape[0] // tp_size
+    tp_rank = get_tensor_model_parallel_rank()
+    return is_padding[tp_rank * chunk : (tp_rank + 1) * chunk]

vllm/models/kimi_k3/nvidia/ops/third_party/__init__.py

@@ -0,0 +1,2 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

vllm/models/kimi_k3/nvidia/ops/third_party/kda/__init__.py

@@ -0,0 +1,28 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from .chunk import (
+    chunk_kda,
+    chunk_kda_fwd,
+    chunk_kda_with_fused_gate,
+    chunk_kda_with_fused_gate_fwd,
+    fused_kda_gate,
+    fused_kda_gate_chunk_cumsum,
+)
+from .fused_recurrent import (
+    fused_recurrent_kda,
+    fused_recurrent_kda_fwd,
+    fused_recurrent_kda_packed_decode,
+)
+
+__all__ = [
+    "chunk_kda",
+    "chunk_kda_fwd",
+    "chunk_kda_with_fused_gate",
+    "chunk_kda_with_fused_gate_fwd",
+    "fused_kda_gate",
+    "fused_kda_gate_chunk_cumsum",
+    "fused_recurrent_kda",
+    "fused_recurrent_kda_fwd",
+    "fused_recurrent_kda_packed_decode",
+]

vllm/models/kimi_k3/nvidia/ops/third_party/kda/chunk.py

@@ -0,0 +1,938 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Songlin Yang, Yu Zhang, Zhiyuan Li
+#
+# This file contains code copied from the flash-linear-attention project.
+# The original source code was licensed under the MIT license and included
+# the following copyright notice:
+# Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li
+# ruff: noqa: E501
+
+
+import torch
+
+from vllm.third_party.flash_linear_attention.ops.chunk_delta_h import (
+    chunk_gated_delta_rule_fwd_h,
+)
+from vllm.third_party.flash_linear_attention.ops.cumsum import chunk_local_cumsum
+from vllm.third_party.flash_linear_attention.ops.index import prepare_chunk_indices
+from vllm.third_party.flash_linear_attention.ops.l2norm import l2norm_fwd
+from vllm.third_party.flash_linear_attention.ops.op import exp2, log
+from vllm.third_party.flash_linear_attention.ops.utils import FLA_CHUNK_SIZE, is_amd
+from vllm.triton_utils import tl, triton
+from vllm.utils.math_utils import RCP_LN2, cdiv, next_power_of_2
+
+from .chunk_intra import chunk_kda_fwd_intra
+
+BT_LIST_AUTOTUNE = [32, 64, 128]
+NUM_WARPS_AUTOTUNE = [2, 4, 8, 16] if is_amd else [4, 8, 16, 32]
+
+
+@triton.heuristics(
+    {
+        "STORE_QG": lambda args: args["qg"] is not None,
+        "STORE_KG": lambda args: args["kg"] is not None,
+        "IS_VARLEN": lambda args: args["cu_seqlens"] is not None,
+    }
+)
+@triton.autotune(
+    configs=[
+        triton.Config({}, num_warps=num_warps, num_stages=num_stages)
+        for num_warps in [2, 4, 8]
+        for num_stages in [2, 3, 4]
+    ],
+    key=["H", "K", "V", "BT", "BK", "BV", "IS_VARLEN"],
+)
+@triton.jit(do_not_specialize=["T"])
+def recompute_w_u_fwd_kernel(
+    q,
+    k,
+    qg,
+    kg,
+    v,
+    beta,
+    w,
+    u,
+    A,
+    gk,
+    cu_seqlens,
+    chunk_indices,
+    T,
+    H: tl.constexpr,
+    K: tl.constexpr,
+    V: tl.constexpr,
+    BT: tl.constexpr,
+    BK: tl.constexpr,
+    BV: tl.constexpr,
+    STORE_QG: tl.constexpr,
+    STORE_KG: tl.constexpr,
+    IS_VARLEN: tl.constexpr,
+    DOT_PRECISION: tl.constexpr,
+):
+    i_t, i_bh = tl.program_id(0), tl.program_id(1)
+    i_b, i_h = i_bh // H, i_bh % H
+    if IS_VARLEN:
+        i_n, i_t = (
+            tl.load(chunk_indices + i_t * 2).to(tl.int32),
+            tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32),
+        )
+        bos, eos = (
+            tl.load(cu_seqlens + i_n).to(tl.int32),
+            tl.load(cu_seqlens + i_n + 1).to(tl.int32),
+        )
+        T = eos - bos
+    else:
+        bos, eos = i_b * T, i_b * T + T
+    p_b = tl.make_block_ptr(beta + bos * H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,))
+    b_b = tl.load(p_b, boundary_check=(0,)).to(tl.float32)
+
+    p_A = tl.make_block_ptr(
+        A + (bos * H + i_h) * BT, (T, BT), (H * BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)
+    )
+    b_A = tl.load(p_A, boundary_check=(0, 1))
+
+    for i_v in range(tl.cdiv(V, BV)):
+        p_v = tl.make_block_ptr(
+            v + (bos * H + i_h) * V,
+            (T, V),
+            (H * V, 1),
+            (i_t * BT, i_v * BV),
+            (BT, BV),
+            (1, 0),
+        )
+        p_u = tl.make_block_ptr(
+            u + (bos * H + i_h) * V,
+            (T, V),
+            (H * V, 1),
+            (i_t * BT, i_v * BV),
+            (BT, BV),
+            (1, 0),
+        )
+        b_v = tl.load(p_v, boundary_check=(0, 1))
+        b_vb = (b_v * b_b[:, None]).to(b_v.dtype)
+        b_u = tl.dot(b_A, b_vb, input_precision=DOT_PRECISION)
+        tl.store(p_u, b_u.to(p_u.dtype.element_ty), boundary_check=(0, 1))
+
+    for i_k in range(tl.cdiv(K, BK)):
+        p_w = tl.make_block_ptr(
+            w + (bos * H + i_h) * K,
+            (T, K),
+            (H * K, 1),
+            (i_t * BT, i_k * BK),
+            (BT, BK),
+            (1, 0),
+        )
+        p_k = tl.make_block_ptr(
+            k + (bos * H + i_h) * K,
+            (T, K),
+            (H * K, 1),
+            (i_t * BT, i_k * BK),
+            (BT, BK),
+            (1, 0),
+        )
+        b_k = tl.load(p_k, boundary_check=(0, 1))
+        b_kb = b_k * b_b[:, None]
+
+        p_gk = tl.make_block_ptr(
+            gk + (bos * H + i_h) * K,
+            (T, K),
+            (H * K, 1),
+            (i_t * BT, i_k * BK),
+            (BT, BK),
+            (1, 0),
+        )
+        b_gk = tl.load(p_gk, boundary_check=(0, 1))
+        b_kb *= exp2(b_gk)
+        if STORE_QG:
+            p_q = tl.make_block_ptr(
+                q + (bos * H + i_h) * K,
+                (T, K),
+                (H * K, 1),
+                (i_t * BT, i_k * BK),
+                (BT, BK),
+                (1, 0),
+            )
+            p_qg = tl.make_block_ptr(
+                qg + (bos * H + i_h) * K,
+                (T, K),
+                (H * K, 1),
+                (i_t * BT, i_k * BK),
+                (BT, BK),
+                (1, 0),
+            )
+            b_q = tl.load(p_q, boundary_check=(0, 1))
+            b_qg = b_q * exp2(b_gk)
+            tl.store(p_qg, b_qg.to(p_qg.dtype.element_ty), boundary_check=(0, 1))
+        if STORE_KG:
+            last_idx = min(i_t * BT + BT, T) - 1
+
+            o_k = i_k * BK + tl.arange(0, BK)
+            m_k = o_k < K
+            b_gn = tl.load(
+                gk + ((bos + last_idx) * H + i_h) * K + o_k, mask=m_k, other=0.0
+            )
+            b_kg = b_k * exp2(b_gn - b_gk)
+
+            p_kg = tl.make_block_ptr(
+                kg + (bos * H + i_h) * K,
+                (T, K),
+                (H * K, 1),
+                (i_t * BT, i_k * BK),
+                (BT, BK),
+                (1, 0),
+            )
+            tl.store(p_kg, b_kg.to(p_kg.dtype.element_ty), boundary_check=(0, 1))
+
+        b_w = tl.dot(b_A, b_kb.to(b_k.dtype))
+        tl.store(p_w, b_w.to(p_w.dtype.element_ty), boundary_check=(0, 1))
+
+
+def recompute_w_u_fwd(
+    k: torch.Tensor,
+    v: torch.Tensor,
+    beta: torch.Tensor,
+    A: torch.Tensor,
+    q: torch.Tensor | None = None,
+    gk: torch.Tensor | None = None,
+    cu_seqlens: torch.Tensor | None = None,
+    chunk_indices: torch.Tensor | None = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    B, T, H, K, V = *k.shape, v.shape[-1]
+    BT = A.shape[-1]
+    BK = 64
+    BV = 64
+
+    if chunk_indices is None and cu_seqlens is not None:
+        chunk_indices = prepare_chunk_indices(cu_seqlens, BT)
+    NT = cdiv(T, BT) if cu_seqlens is None else len(chunk_indices)
+
+    w = torch.empty_like(k)
+    u = torch.empty_like(v)
+    kg = torch.empty_like(k) if gk is not None else None
+    recompute_w_u_fwd_kernel[(NT, B * H)](
+        q=q,
+        k=k,
+        qg=None,
+        kg=kg,
+        v=v,
+        beta=beta,
+        w=w,
+        u=u,
+        A=A,
+        gk=gk,
+        cu_seqlens=cu_seqlens,
+        chunk_indices=chunk_indices,
+        T=T,
+        H=H,
+        K=K,
+        V=V,
+        BT=BT,
+        BK=BK,
+        BV=BV,
+        DOT_PRECISION="ieee",
+    )
+    return w, u, None, kg
+
+
+@triton.heuristics({"IS_VARLEN": lambda args: args["cu_seqlens"] is not None})
+@triton.autotune(
+    configs=[
+        triton.Config({"BK": BK, "BV": BV}, num_warps=num_warps, num_stages=num_stages)
+        for BK in [32, 64]
+        for BV in [64, 128]
+        for num_warps in [2, 4, 8]
+        for num_stages in [2, 3, 4]
+    ],
+    key=["BT"],
+)
+@triton.jit(do_not_specialize=["T"])
+def chunk_gla_fwd_kernel_o(
+    q,
+    v,
+    g,
+    h,
+    o,
+    A,
+    cu_seqlens,
+    chunk_indices,
+    scale,
+    T,
+    H: tl.constexpr,
+    K: tl.constexpr,
+    V: tl.constexpr,
+    BT: tl.constexpr,
+    BK: tl.constexpr,
+    BV: tl.constexpr,
+    IS_VARLEN: tl.constexpr,
+):
+    i_v, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2)
+    i_b, i_h = i_bh // H, i_bh % H
+    if IS_VARLEN:
+        i_tg = i_t
+        i_n, i_t = (
+            tl.load(chunk_indices + i_t * 2).to(tl.int32),
+            tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32),
+        )
+        bos, eos = (
+            tl.load(cu_seqlens + i_n).to(tl.int32),
+            tl.load(cu_seqlens + i_n + 1).to(tl.int32),
+        )
+        T = eos - bos
+        NT = tl.cdiv(T, BT)
+    else:
+        NT = tl.cdiv(T, BT)
+        i_tg = i_b * NT + i_t
+        bos, eos = i_b * T, i_b * T + T
+
+    m_s = tl.arange(0, BT)[:, None] >= tl.arange(0, BT)[None, :]
+
+    b_o = tl.zeros([BT, BV], dtype=tl.float32)
+    for i_k in range(tl.cdiv(K, BK)):
+        p_q = tl.make_block_ptr(
+            q + (bos * H + i_h) * K,
+            (T, K),
+            (H * K, 1),
+            (i_t * BT, i_k * BK),
+            (BT, BK),
+            (1, 0),
+        )
+        p_g = tl.make_block_ptr(
+            g + (bos * H + i_h) * K,
+            (T, K),
+            (H * K, 1),
+            (i_t * BT, i_k * BK),
+            (BT, BK),
+            (1, 0),
+        )
+        p_h = tl.make_block_ptr(
+            h + (i_tg * H + i_h) * K * V,
+            (V, K),
+            (K, 1),
+            (i_v * BV, i_k * BK),
+            (BV, BK),
+            (1, 0),
+        )
+
+        # [BT, BK]
+        b_q = tl.load(p_q, boundary_check=(0, 1))
+        b_q = (b_q * scale).to(b_q.dtype)
+        # [BT, BK]
+        b_g = tl.load(p_g, boundary_check=(0, 1))
+        # [BT, BK]
+        b_qg = (b_q * exp2(b_g)).to(b_q.dtype)
+        # [BV, BK]
+        b_h = tl.load(p_h, boundary_check=(0, 1))
+        # [BT, BV]
+        if i_k >= 0:
+            b_o += tl.dot(b_qg, tl.trans(b_h).to(b_qg.dtype))
+    p_v = tl.make_block_ptr(
+        v + (bos * H + i_h) * V,
+        (T, V),
+        (H * V, 1),
+        (i_t * BT, i_v * BV),
+        (BT, BV),
+        (1, 0),
+    )
+    p_o = tl.make_block_ptr(
+        o + (bos * H + i_h) * V,
+        (T, V),
+        (H * V, 1),
+        (i_t * BT, i_v * BV),
+        (BT, BV),
+        (1, 0),
+    )
+    p_A = tl.make_block_ptr(
+        A + (bos * H + i_h) * BT, (T, BT), (H * BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)
+    )
+    # [BT, BV]
+    b_v = tl.load(p_v, boundary_check=(0, 1))
+    # [BT, BT]
+    b_A = tl.load(p_A, boundary_check=(0, 1))
+    b_A = tl.where(m_s, b_A, 0.0).to(b_v.dtype)
+    b_o += tl.dot(b_A, b_v, allow_tf32=False)
+    tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1))
+
+
+def chunk_gla_fwd_o_gk(
+    q: torch.Tensor,
+    v: torch.Tensor,
+    g: torch.Tensor,
+    A: torch.Tensor,
+    h: torch.Tensor,
+    o: torch.Tensor,
+    scale: float,
+    cu_seqlens: torch.Tensor | None = None,
+    chunk_indices: torch.Tensor | None = None,
+    chunk_size: int = FLA_CHUNK_SIZE,
+):
+    B, T, H, K, V = *q.shape, v.shape[-1]
+    BT = chunk_size
+
+    if chunk_indices is None and cu_seqlens is not None:
+        chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size)
+    NT = cdiv(T, BT) if cu_seqlens is None else len(chunk_indices)
+
+    def grid(meta):
+        return (cdiv(V, meta["BV"]), NT, B * H)
+
+    chunk_gla_fwd_kernel_o[grid](
+        q=q,
+        v=v,
+        g=g,
+        h=h,
+        o=o,
+        A=A,
+        cu_seqlens=cu_seqlens,
+        chunk_indices=chunk_indices,
+        scale=scale,
+        T=T,
+        H=H,
+        K=K,
+        V=V,
+        BT=BT,
+    )
+    return o
+
+
+@triton.heuristics(
+    {
+        "HAS_BIAS": lambda args: args["g_bias"] is not None,
+        "IS_VARLEN": lambda args: args["cu_seqlens"] is not None,
+    }
+)
+@triton.autotune(
+    configs=[
+        triton.Config({"BS": BS}, num_warps=num_warps)
+        for BS in [32, 64]
+        for num_warps in [2, 4, 8]
+    ],
+    key=["H", "S", "BT", "IS_VARLEN"],
+)
+@triton.jit(do_not_specialize=["T"])
+def kda_gate_chunk_cumsum_vector_kernel(
+    s,
+    raw_beta,
+    A_log,
+    g_bias,
+    o,
+    beta_out,
+    cu_seqlens,
+    chunk_indices,
+    cumsum_scale,
+    lower_bound,
+    beta,
+    threshold,
+    T,
+    stride_beta_batch,
+    stride_beta_token,
+    stride_beta_head,
+    H: tl.constexpr,
+    S: tl.constexpr,
+    BT: tl.constexpr,
+    BS: tl.constexpr,
+    HAS_BIAS: tl.constexpr,
+    IS_VARLEN: tl.constexpr,
+    USE_LOWER_BOUND: tl.constexpr,
+):
+    i_s, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2)
+    i_b, i_h = i_bh // H, i_bh % H
+    if IS_VARLEN:
+        i_n, i_t = (
+            tl.load(chunk_indices + i_t * 2).to(tl.int32),
+            tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32),
+        )
+        bos, eos = (
+            tl.load(cu_seqlens + i_n).to(tl.int32),
+            tl.load(cu_seqlens + i_n + 1).to(tl.int32),
+        )
+        T = eos - bos
+    else:
+        bos = i_b * T
+
+    if i_s == 0:
+        o_beta_t = tl.arange(0, BT)
+        m_beta = i_t * BT + o_beta_t < T
+        if IS_VARLEN:
+            p_beta = (
+                raw_beta
+                + (bos + i_t * BT + o_beta_t) * stride_beta_token
+                + i_h * stride_beta_head
+            )
+        else:
+            p_beta = (
+                raw_beta
+                + i_b * stride_beta_batch
+                + (i_t * BT + o_beta_t) * stride_beta_token
+                + i_h * stride_beta_head
+            )
+        b_beta = tl.load(p_beta, mask=m_beta, other=0.0).to(tl.float32)
+        p_beta_out = beta_out + (bos + i_t * BT + o_beta_t) * H + i_h
+        tl.store(p_beta_out, tl.sigmoid(b_beta), mask=m_beta)
+        return
+
+    i_s -= 1
+
+    p_s = tl.make_block_ptr(
+        s + (bos * H + i_h) * S,
+        (T, S),
+        (H * S, 1),
+        (i_t * BT, i_s * BS),
+        (BT, BS),
+        (1, 0),
+    )
+    p_o = tl.make_block_ptr(
+        o + (bos * H + i_h) * S,
+        (T, S),
+        (H * S, 1),
+        (i_t * BT, i_s * BS),
+        (BT, BS),
+        (1, 0),
+    )
+
+    b_s = tl.load(p_s, boundary_check=(0, 1)).to(tl.float32)
+    if HAS_BIAS:
+        p_bias = tl.make_block_ptr(
+            g_bias + i_h * S,
+            (S,),
+            (1,),
+            (i_s * BS,),
+            (BS,),
+            (0,),
+        )
+        b_bias = tl.load(p_bias, boundary_check=(0,)).to(tl.float32)
+        b_s += b_bias[None, :]
+
+    b_a = tl.exp(tl.load(A_log + i_h).to(tl.float32))
+    if USE_LOWER_BOUND:
+        b_gate = lower_bound * tl.sigmoid(b_a * b_s)
+    else:
+        b_g_scaled = b_s * beta
+        b_softplus = tl.where(
+            b_g_scaled > threshold,
+            b_s,
+            (1.0 / beta) * log(1.0 + tl.exp(b_g_scaled)),
+        )
+        b_gate = -b_a * b_softplus
+
+    # Boundary loads return zero, but bias and gate activation can make padded
+    # rows nonzero. Padding trails valid rows, so it only affects masked stores.
+    b_o = tl.cumsum(b_gate, axis=0) * cumsum_scale
+    tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1))
+
+
+def fused_kda_gate_chunk_cumsum(
+    raw_g: torch.Tensor,
+    raw_beta: torch.Tensor,
+    A_log: torch.Tensor,
+    g_bias: torch.Tensor | None = None,
+    beta: float = 1.0,
+    threshold: float = 20.0,
+    lower_bound: float | None = None,
+    cu_seqlens: torch.Tensor | None = None,
+    chunk_indices: torch.Tensor | None = None,
+    chunk_size: int = FLA_CHUNK_SIZE,
+    output_dtype: torch.dtype | None = torch.float,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    if cu_seqlens is not None:
+        assert raw_g.shape[0] == 1, (
+            "Only batch size 1 is supported when cu_seqlens are provided"
+        )
+    B, T, H, D = raw_g.shape
+    if raw_beta.shape != (B, T, H):
+        raise ValueError(
+            f"Expected raw_beta shape {(B, T, H)}, got {raw_beta.shape}"
+        )
+    if chunk_indices is None and cu_seqlens is not None:
+        chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size)
+    NT = cdiv(T, chunk_size) if cu_seqlens is None else len(chunk_indices)
+
+    A_log = A_log.reshape(-1)
+    if g_bias is not None:
+        g_bias = g_bias.reshape(-1)
+    y = torch.empty_like(raw_g, dtype=output_dtype or raw_g.dtype)
+    beta_out = torch.empty(raw_beta.shape, device=raw_beta.device, dtype=torch.float32)
+
+    def grid(meta):
+        # For each (chunk, head), program 0 computes beta without extending a
+        # gate tile's critical path. The remaining programs cover the gate dim.
+        return (cdiv(meta["S"], meta["BS"]) + 1, NT, B * H)
+
+    kda_gate_chunk_cumsum_vector_kernel[grid](
+        s=raw_g,
+        raw_beta=raw_beta,
+        A_log=A_log,
+        g_bias=g_bias,
+        o=y,
+        beta_out=beta_out,
+        cu_seqlens=cu_seqlens,
+        chunk_indices=chunk_indices,
+        # RCP_LN2 folds in the natural-log -> log2 conversion so downstream
+        # exp2-based kernels reproduce exp(g). Keep this in sync with the
+        # `use_exp2=True` path in `_chunk_kda_fwd_with_cumulative_g`.
+        cumsum_scale=RCP_LN2,
+        lower_bound=lower_bound or 0.0,
+        beta=beta,
+        threshold=threshold,
+        T=T,
+        stride_beta_batch=raw_beta.stride(0),
+        stride_beta_token=raw_beta.stride(1),
+        stride_beta_head=raw_beta.stride(2),
+        H=H,
+        S=D,
+        BT=chunk_size,
+        USE_LOWER_BOUND=lower_bound is not None,
+    )
+    return y, beta_out
+
+
+def _chunk_kda_fwd_with_cumulative_g(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    g: torch.Tensor,
+    beta: torch.Tensor,
+    scale: float,
+    initial_state: torch.Tensor,
+    output_final_state: bool,
+    cu_seqlens: torch.Tensor | None = None,
+    chunk_indices: torch.Tensor | None = None,
+    chunk_size: int = FLA_CHUNK_SIZE,
+    safe_gate: bool = False,
+):
+    # `g` must already be chunk-local cumulatively-summed AND scaled by
+    # RCP_LN2 (so the downstream exp2-based kernels reproduce exp(g)).
+    # Use `chunk_kda_fwd` or `chunk_kda_with_fused_gate_fwd` instead of
+    # calling this helper directly unless that invariant is upheld.
+    Aqk, A = chunk_kda_fwd_intra(
+        q=q,
+        k=k,
+        gk=g,
+        beta=beta,
+        scale=scale,
+        cu_seqlens=cu_seqlens,
+        chunk_indices=chunk_indices,
+        chunk_size=chunk_size,
+        safe_gate=safe_gate,
+    )
+    w, u, _, kg = recompute_w_u_fwd(
+        k=k,
+        v=v,
+        beta=beta,
+        A=A,
+        gk=g,
+        cu_seqlens=cu_seqlens,
+        chunk_indices=chunk_indices,
+    )
+    del A
+    h, v_new, final_state = chunk_gated_delta_rule_fwd_h(
+        k=kg,
+        w=w,
+        u=u,
+        gk=g,
+        initial_state=initial_state,
+        output_final_state=output_final_state,
+        cu_seqlens=cu_seqlens,
+        chunk_indices=chunk_indices,
+        use_exp2=True,
+    )
+    del w, u, kg
+    o = chunk_gla_fwd_o_gk(
+        q=q,
+        v=v_new,
+        g=g,
+        A=Aqk,
+        h=h,
+        o=v,
+        scale=scale,
+        cu_seqlens=cu_seqlens,
+        chunk_indices=chunk_indices,
+        chunk_size=chunk_size,
+    )
+    del Aqk, v_new, h
+    return o, final_state
+
+
+def chunk_kda_fwd(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    g: torch.Tensor,
+    beta: torch.Tensor,
+    scale: float,
+    initial_state: torch.Tensor,
+    output_final_state: bool,
+    cu_seqlens: torch.Tensor | None = None,
+):
+    chunk_size = FLA_CHUNK_SIZE
+    chunk_indices = (
+        prepare_chunk_indices(cu_seqlens, chunk_size)
+        if cu_seqlens is not None
+        else None
+    )
+    g = chunk_local_cumsum(
+        g,
+        chunk_size=chunk_size,
+        cu_seqlens=cu_seqlens,
+        chunk_indices=chunk_indices,
+    )
+    # KDA evaluates cumulative gate decays with exp2. Convert from natural-log
+    # space so exp(x) is preserved as exp2(x / ln(2)).
+    g = g * RCP_LN2
+    return _chunk_kda_fwd_with_cumulative_g(
+        q=q,
+        k=k,
+        v=v,
+        g=g,
+        beta=beta,
+        scale=scale,
+        initial_state=initial_state,
+        output_final_state=output_final_state,
+        cu_seqlens=cu_seqlens,
+        chunk_indices=chunk_indices,
+        chunk_size=chunk_size,
+    )
+
+
+def chunk_kda_with_fused_gate_fwd(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    raw_g: torch.Tensor,
+    raw_beta: torch.Tensor,
+    A_log: torch.Tensor,
+    g_bias: torch.Tensor | None,
+    scale: float,
+    initial_state: torch.Tensor,
+    output_final_state: bool,
+    lower_bound: float | None = None,
+    cu_seqlens: torch.Tensor | None = None,
+):
+    chunk_size = FLA_CHUNK_SIZE
+    chunk_indices = (
+        prepare_chunk_indices(cu_seqlens, chunk_size)
+        if cu_seqlens is not None
+        else None
+    )
+    g, beta = fused_kda_gate_chunk_cumsum(
+        raw_g,
+        raw_beta=raw_beta,
+        A_log=A_log,
+        g_bias=g_bias,
+        cu_seqlens=cu_seqlens,
+        chunk_indices=chunk_indices,
+        chunk_size=chunk_size,
+        lower_bound=lower_bound,
+    )
+    return _chunk_kda_fwd_with_cumulative_g(
+        q=q,
+        k=k,
+        v=v,
+        g=g,
+        beta=beta,
+        scale=scale,
+        initial_state=initial_state,
+        output_final_state=output_final_state,
+        cu_seqlens=cu_seqlens,
+        chunk_indices=chunk_indices,
+        chunk_size=chunk_size,
+        safe_gate=lower_bound is not None,
+    )
+
+
+def chunk_kda(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    g: torch.Tensor,
+    beta: torch.Tensor,
+    scale: float = None,
+    initial_state: torch.Tensor = None,
+    output_final_state: bool = False,
+    use_qk_l2norm_in_kernel: bool = False,
+    cu_seqlens: torch.Tensor | None = None,
+    **kwargs,
+):
+    if scale is None:
+        scale = k.shape[-1] ** -0.5
+
+    if use_qk_l2norm_in_kernel:
+        q = l2norm_fwd(q.contiguous())
+        k = l2norm_fwd(k.contiguous())
+
+    o, final_state = chunk_kda_fwd(
+        q=q,
+        k=k,
+        v=v.contiguous(),
+        g=g.contiguous(),
+        beta=beta.contiguous(),
+        scale=scale,
+        initial_state=initial_state.contiguous(),
+        output_final_state=output_final_state,
+        cu_seqlens=cu_seqlens,
+    )
+    return o, final_state
+
+
+def chunk_kda_with_fused_gate(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    raw_g: torch.Tensor,
+    raw_beta: torch.Tensor,
+    A_log: torch.Tensor,
+    g_bias: torch.Tensor | None,
+    scale: float | None = None,
+    initial_state: torch.Tensor | None = None,
+    output_final_state: bool = False,
+    lower_bound: float | None = None,
+    use_qk_l2norm_in_kernel: bool = False,
+    cu_seqlens: torch.Tensor | None = None,
+    **kwargs,
+):
+    """Run chunk KDA from raw gate and beta projections."""
+    if scale is None:
+        scale = k.shape[-1] ** -0.5
+
+    if use_qk_l2norm_in_kernel:
+        q = l2norm_fwd(q.contiguous())
+        k = l2norm_fwd(k.contiguous())
+
+    o, final_state = chunk_kda_with_fused_gate_fwd(
+        q=q,
+        k=k,
+        v=v.contiguous(),
+        raw_g=raw_g.contiguous(),
+        raw_beta=raw_beta,
+        A_log=A_log,
+        g_bias=g_bias,
+        scale=scale,
+        initial_state=initial_state.contiguous() if initial_state is not None else None,
+        output_final_state=output_final_state,
+        lower_bound=lower_bound,
+        cu_seqlens=cu_seqlens,
+    )
+    return o, final_state
+
+
+@triton.autotune(
+    configs=[
+        triton.Config({"BT": bt}, num_warps=nw, num_stages=ns)
+        for bt in BT_LIST_AUTOTUNE
+        for nw in NUM_WARPS_AUTOTUNE
+        for ns in [2, 3]
+    ],
+    key=["H", "D"],
+)
+@triton.jit
+def kda_gate_fwd_kernel(
+    g,
+    A,
+    y,
+    g_bias,
+    lower_bound,
+    beta: tl.constexpr,
+    threshold: tl.constexpr,
+    T,
+    H,
+    D: tl.constexpr,
+    BT: tl.constexpr,
+    BD: tl.constexpr,
+    HAS_BIAS: tl.constexpr,
+    USE_LOWER_BOUND: tl.constexpr,
+):
+    i_t, i_h = tl.program_id(0), tl.program_id(1)
+    n_t = i_t * BT
+
+    b_a = tl.exp(tl.load(A + i_h).to(tl.float32))
+
+    stride_row = H * D
+    stride_col = 1
+
+    g_ptr = tl.make_block_ptr(
+        base=g + i_h * D,
+        shape=(T, D),
+        strides=(stride_row, stride_col),
+        offsets=(n_t, 0),
+        block_shape=(BT, BD),
+        order=(1, 0),
+    )
+
+    y_ptr = tl.make_block_ptr(
+        base=y + i_h * D,
+        shape=(T, D),
+        strides=(stride_row, stride_col),
+        offsets=(n_t, 0),
+        block_shape=(BT, BD),
+        order=(1, 0),
+    )
+
+    b_g = tl.load(g_ptr, boundary_check=(0, 1)).to(tl.float32)
+
+    if HAS_BIAS:
+        n_d = tl.arange(0, BD)
+        bias_mask = n_d < D
+        b_bias = tl.load(g_bias + i_h * D + n_d, mask=bias_mask, other=0.0).to(
+            tl.float32
+        )
+        b_g = b_g + b_bias[None, :]
+
+    if USE_LOWER_BOUND:
+        b_y = lower_bound * tl.sigmoid(b_a * b_g)
+    else:
+        g_scaled = b_g * beta
+        use_linear = g_scaled > threshold
+        sp = tl.where(use_linear, b_g, (1.0 / beta) * log(1.0 + tl.exp(g_scaled)))
+        b_y = -b_a * sp
+
+    tl.store(y_ptr, b_y.to(y.dtype.element_ty), boundary_check=(0, 1))
+
+
+def fused_kda_gate(
+    g: torch.Tensor,
+    A: torch.Tensor,
+    head_k_dim: int,
+    g_bias: torch.Tensor | None = None,
+    beta: float = 1.0,
+    threshold: float = 20.0,
+    lower_bound: float | None = None,
+) -> torch.Tensor:
+    """
+    Forward pass for KDA gate:
+      input g: [..., H*D]
+      param A: [H] or [1, 1, H, 1]
+      beta: softplus beta parameter
+      threshold: softplus threshold parameter
+      return  : [..., H, D]
+    """
+    orig_shape = g.shape[:-1]
+
+    g = g.view(-1, g.shape[-1])
+    T = g.shape[0]
+    HD = g.shape[1]
+    H = A.numel()
+    assert H * head_k_dim == HD
+    assert g.stride() == (HD, 1)
+
+    y = torch.empty_like(g, dtype=torch.float32)
+
+    def grid(meta):
+        return (cdiv(T, meta["BT"]), H)
+
+    kda_gate_fwd_kernel[grid](
+        g,
+        A,
+        y,
+        g_bias,
+        lower_bound or 0.0,
+        beta,
+        threshold,
+        T,
+        H,
+        head_k_dim,
+        BD=next_power_of_2(head_k_dim),
+        HAS_BIAS=g_bias is not None,
+        USE_LOWER_BOUND=lower_bound is not None,
+    )
+
+    y = y.view(*orig_shape, H, head_k_dim)
+    return y