[ROCm][CI] Fix Kimi K3 KDA on ROCm (#50262)

Signed-off-by: Stefan Koncarevic <stefan.koncarevic@amd.com>

.buildkite/check-torch-abi.py

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

.buildkite/intel_jobs/test-intel.yaml

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

.buildkite/performance-benchmarks/tests/serving-tests.json

@@ -28,11 +28,6 @@
         "dataset_path": "./ShareGPT_V3_unfiltered_cleaned_split.json"
       }
     },
-    {
-        "dataset_name": "sharegpt",
-        "dataset_path": "./ShareGPT_V3_unfiltered_cleaned_split.json"
-      }
-    },
     {
       "test_name": "serving_llama8B_tp1_random_128_128",
       "server_parameters": {

.buildkite/scripts/hardware_ci/run-intel-test.sh

@@ -369,7 +369,7 @@ export HF_TOKEN ZE_AFFINITY_MASK
     -e CMDS \
     --name "${container_name}" \
     "${IMAGE}" \
-    bash -c 'set -e; source /opt/intel/oneapi/setvars.sh --force; source /opt/intel/oneapi/ccl/2021.15/env/vars.sh --force; echo "ZE_AFFINITY_MASK is ${ZE_AFFINITY_MASK:-}"; eval "$CMDS"' \
+    bash -c 'set -e; echo "ZE_AFFINITY_MASK is ${ZE_AFFINITY_MASK:-}"; eval "$CMDS"' \
     >/dev/null
 } 9>/tmp/docker-pull.lock
 

.buildkite/test-amd.yaml

@@ -1576,10 +1576,14 @@ steps:
   - 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
+  - vllm/models/kimi_k3/nvidia/kda.py
+  - vllm/models/kimi_k3/nvidia/kda_metadata.py
+  - vllm/models/kimi_k3/nvidia/ops/third_party/kda/
+  - tests/models/kimi_k3/test_kda.py
+  - tests/models/kimi_k3/test_kda_metadata.py
   - vllm/platforms/rocm.py
   commands:
-  - pytest -v -s kernels/test_kda.py
+  - pytest -v -s models/kimi_k3/test_kda.py models/kimi_k3/test_kda_metadata.py
 
 - label: Kernels Mamba Test # TBD
   timeout_in_minutes: 180
@@ -2289,7 +2293,7 @@ steps:
   - export VLLM_USE_RUST_FRONTEND=1
   - export VLLM_WORKER_MULTIPROC_METHOD=spawn
   - pytest -v -s benchmarks/test_serve_cli.py -k "not insecure and not (test_bench_serve and not test_bench_serve_chat)"
-  - pytest -v -s entrypoints/openai/chat_completion/test_chat_completion.py -k "not test_invalid_json_schema and not test_invalid_regex"
+  - pytest -v -s entrypoints/openai/chat_completion/test_chat_completion.py -k "not test_invalid_json_schema and not test_invalid_regex and not test_kv_transfer_prompt_token_ids_round_trip and not test_kv_transfer_prompt_token_ids_streaming"
   - pytest -v -s entrypoints/openai/chat_completion/test_chat_logit_bias_validation.py -k "not multiple"
   - pytest -v -s entrypoints/openai/completion/test_shutdown.py -k "not engine_failure and not test_abort_timeout_exits_quickly"
   - pytest -v -s entrypoints/openai/test_return_token_ids.py -k "not test_comparison"

.buildkite/test_areas/kernels.yaml

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

.buildkite/test_areas/lora.yaml

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

.buildkite/test_areas/models_basic.yaml

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

.buildkite/test_areas/rust_frontend.yaml

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

.github/workflows/new_pr_bot.yml

@@ -80,9 +80,9 @@ jobs:
                   '',
                   '\u{1f4ac} Join our developer Slack at https://slack.vllm.ai to discuss your PR in `#pr-reviews`, coordinate on features in `#feat-` channels, or join special interest groups in `#sig-` channels.',
                   '',
-                  'PRs do not trigger a full CI run by default. Once the PR is approved and ready to go, your PR reviewer(s) can run CI to test the changes comprehensively before merging.',
+                  'PRs do not trigger a full CI run by default. Reviewers with write access and configured trusted contributors can comment `/ci run` whenever CI signals are needed.',
                   '',
-                  'To run CI, PR reviewers can either: Add `ready` label to the PR or enable auto-merge.',
+                  'Once the PR is approved or has the `ready` label, the PR author can also use `/ci run` or `/ci retry`. New commits do not start CI automatically.',
                   '',
                   'If you have any questions, please reach out to us on Slack at https://slack.vllm.ai.',
                   '',

.github/workflows/pre-commit.yml

@@ -41,7 +41,7 @@ jobs:
           if (hasReadyLabel || hasVerifiedLabel || mergedCount >= 4) {
             core.info(`Check passed: verified label=${hasVerifiedLabel}, ready label=${hasReadyLabel}, 4+ merged PRs=${mergedCount >= 4}`);
           } else {
-            core.setFailed(`PR must have the 'verified', 'ready', or 'ready-run-all-tests' label (the ready labels also trigger tests) or the author must have at least 4 merged PRs (found ${mergedCount}).`);
+            core.setFailed(`PR must have the 'verified', 'ready', or 'ready-run-all-tests' label to run pre-commit, or the author must have at least 4 merged PRs (found ${mergedCount}).`);
           }
 
   pre-commit:

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

@@ -0,0 +1,40 @@
+name: Run CI from PR comment
+
+on:
+  issue_comment:
+    types: [created]
+
+concurrency:
+  group: run-ci-comment-${{ github.event.issue.number }}
+  cancel-in-progress: false
+
+permissions:
+  contents: read
+  issues: write
+  pull-requests: write
+
+jobs:
+  run-ci-command:
+    if: >-
+      github.event.issue.pull_request &&
+      (github.event.comment.body == '/ci run' ||
+      github.event.comment.body == '/ci retry')
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0
+        with:
+          ref: ${{ github.event.repository.default_branch }}
+          persist-credentials: false
+      - uses: astral-sh/setup-uv@37802adc94f370d6bfd71619e3f0bf239e1f3b78 # v7.6.0
+        with:
+          python-version: "3.12"
+      - name: Authorize and run CI command
+        run: >-
+          uv run --no-project --python 3.12
+          .github/workflows/scripts/run_ci_command.py
+        env:
+          BUILDKITE_API_TOKEN: ${{ secrets.BUILDKITE_API_TOKEN }}
+          BUILDKITE_ORGANIZATION: vllm
+          BUILDKITE_PIPELINE: ci
+          CI_TRUSTED_USERS: ${{ vars.CI_TRUSTED_USERS }}
+          GH_TOKEN: ${{ github.token }}

.github/workflows/scripts/run_ci_command.py

@@ -0,0 +1,581 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import json
+import os
+import sys
+import urllib.error
+import urllib.parse
+import urllib.request
+from collections.abc import Mapping, Sequence
+from typing import Any
+
+COMMAND_RUN_CI = "/ci run"
+COMMAND_RETRY_FAILED = "/ci retry"
+READY_LABELS = {"ready", "ready-run-all-tests"}
+TRUSTED_PERMISSIONS = {"admin", "maintain", "write"}
+ACTIVE_BUILD_STATES = {
+    "blocked",
+    "creating",
+    "scheduled",
+    "running",
+    "failing",
+    "canceling",
+    "waiting",
+    "waiting_failed",
+}
+RETRY_STATES = "failed,timed_out,expired"
+
+
+class ApiError(RuntimeError):
+    def __init__(self, status: int | None, message: str) -> None:
+        super().__init__(message)
+        self.status = status
+
+
+class HttpTransport:
+    def request(
+        self,
+        url: str,
+        *,
+        body: Mapping[str, Any] | None = None,
+        headers: Mapping[str, str] | None = None,
+        method: str = "GET",
+    ) -> Any:
+        data = None if body is None else json.dumps(body).encode()
+        request = urllib.request.Request(
+            url,
+            data=data,
+            headers=dict(headers or {}),
+            method=method,
+        )
+        try:
+            with urllib.request.urlopen(request, timeout=30) as response:
+                response_body = response.read().decode()
+        except urllib.error.HTTPError as error:
+            response_body = error.read().decode()
+            message = self._error_message(response_body, error.reason)
+            raise ApiError(
+                error.code,
+                f"API returned {error.code}: {message}",
+            ) from error
+        except urllib.error.URLError as error:
+            raise ApiError(None, f"API request failed: {error.reason}") from error
+
+        if not response_body:
+            return None
+        try:
+            return json.loads(response_body)
+        except json.JSONDecodeError as error:
+            raise ApiError(None, "API returned a non-JSON response.") from error
+
+    @staticmethod
+    def _error_message(response_body: str, fallback: str) -> str:
+        try:
+            parsed = json.loads(response_body)
+        except json.JSONDecodeError:
+            return fallback
+        return str(parsed.get("message", fallback))
+
+
+class GitHubClient:
+    def __init__(
+        self,
+        token: str,
+        repository: str,
+        transport: HttpTransport | None = None,
+    ) -> None:
+        if not token:
+            raise RuntimeError("GH_TOKEN is not set.")
+        self.owner, self.repo = repository.split("/", maxsplit=1)
+        self.transport = transport or HttpTransport()
+        self.headers = {
+            "Accept": "application/vnd.github+json",
+            "Authorization": f"Bearer {token}",
+            "Content-Type": "application/json",
+            "User-Agent": "vllm-ci-command",
+            "X-GitHub-Api-Version": "2022-11-28",
+        }
+
+    def _request(
+        self,
+        path: str,
+        *,
+        body: Mapping[str, Any] | None = None,
+        method: str = "GET",
+    ) -> Any:
+        return self.transport.request(
+            f"https://api.github.com{path}",
+            body=body,
+            headers=self.headers,
+            method=method,
+        )
+
+    def _repo_path(self, suffix: str) -> str:
+        owner = urllib.parse.quote(self.owner, safe="")
+        repo = urllib.parse.quote(self.repo, safe="")
+        return f"/repos/{owner}/{repo}{suffix}"
+
+    def _paginate(self, path: str) -> list[dict[str, Any]]:
+        results: list[dict[str, Any]] = []
+        separator = "&" if "?" in path else "?"
+        for page in range(1, 101):
+            response = self._request(f"{path}{separator}per_page=100&page={page}")
+            if not isinstance(response, list):
+                raise ApiError(None, "GitHub API returned an invalid list response.")
+            results.extend(response)
+            if len(response) < 100:
+                return results
+        raise ApiError(None, "GitHub API pagination exceeded 10,000 results.")
+
+    def get_pr(self, number: int) -> dict[str, Any]:
+        return self._request(self._repo_path(f"/pulls/{number}"))
+
+    def get_permission(self, actor: str) -> str:
+        username = urllib.parse.quote(actor, safe="")
+        try:
+            response = self._request(
+                self._repo_path(f"/collaborators/{username}/permission")
+            )
+        except ApiError as error:
+            if error.status == 404:
+                return "none"
+            raise
+        return str(response["permission"])
+
+    def get_review_decision(self, number: int) -> str | None:
+        query = """
+          query($owner: String!, $repo: String!, $number: Int!) {
+            repository(owner: $owner, name: $repo) {
+              pullRequest(number: $number) {
+                reviewDecision
+              }
+            }
+          }
+        """
+        response = self._request(
+            "/graphql",
+            body={
+                "query": query,
+                "variables": {
+                    "number": number,
+                    "owner": self.owner,
+                    "repo": self.repo,
+                },
+            },
+            method="POST",
+        )
+        return response["data"]["repository"]["pullRequest"]["reviewDecision"]
+
+    def list_reviews(self, number: int) -> list[dict[str, Any]]:
+        return self._paginate(self._repo_path(f"/pulls/{number}/reviews"))
+
+    def list_reactions(self, comment_id: int) -> list[dict[str, Any]]:
+        return self._paginate(
+            self._repo_path(f"/issues/comments/{comment_id}/reactions")
+        )
+
+    def add_reaction(self, comment_id: int, content: str) -> None:
+        self._request(
+            self._repo_path(f"/issues/comments/{comment_id}/reactions"),
+            body={"content": content},
+            method="POST",
+        )
+
+    def add_comment(self, issue_number: int, body: str) -> None:
+        self._request(
+            self._repo_path(f"/issues/{issue_number}/comments"),
+            body={"body": body},
+            method="POST",
+        )
+
+
+class BuildkiteClient:
+    def __init__(
+        self,
+        token: str,
+        organization: str,
+        pipeline: str,
+        transport: HttpTransport | None = None,
+    ) -> None:
+        self.token = token
+        self.transport = transport or HttpTransport()
+        organization = urllib.parse.quote(organization, safe="")
+        pipeline = urllib.parse.quote(pipeline, safe="")
+        self.base_url = (
+            "https://api.buildkite.com/v2/organizations/"
+            f"{organization}/pipelines/{pipeline}/builds"
+        )
+
+    def _request(
+        self,
+        *,
+        body: Mapping[str, Any] | None = None,
+        method: str = "GET",
+        path: str = "",
+        query: Sequence[tuple[str, str]] = (),
+    ) -> Any:
+        if not self.token:
+            raise RuntimeError("The BUILDKITE_API_TOKEN repository secret is not set.")
+        url = f"{self.base_url}{path}"
+        if query:
+            url = f"{url}?{urllib.parse.urlencode(query)}"
+        return self.transport.request(
+            url,
+            body=body,
+            headers={
+                "Authorization": f"Bearer {self.token}",
+                "Content-Type": "application/json",
+                "User-Agent": "vllm-ci-command",
+            },
+            method=method,
+        )
+
+    def list_builds(
+        self,
+        commit: str,
+        *,
+        metadata: tuple[str, str] | None = None,
+    ) -> list[dict[str, Any]]:
+        query = [
+            ("commit", commit),
+            ("exclude_jobs", "true"),
+            ("exclude_pipeline", "true"),
+            ("per_page", "100"),
+        ]
+        if metadata:
+            key, value = metadata
+            query.append((f"meta_data[{key}]", value))
+        response = self._request(query=query)
+        if not isinstance(response, list):
+            raise ApiError(None, "Buildkite API returned an invalid build list.")
+        return response
+
+    def create_build(self, body: Mapping[str, Any]) -> dict[str, Any]:
+        return self._request(body=body, method="POST")
+
+    def retry_failed_jobs(
+        self,
+        build_number: int,
+        states: str,
+    ) -> dict[str, Any]:
+        number = urllib.parse.quote(str(build_number), safe="")
+        return self._request(
+            body={"states": states},
+            method="PUT",
+            path=f"/{number}/retry_failed_jobs",
+        )
+
+
+def parse_command(body: str) -> str | None:
+    if body in {COMMAND_RUN_CI, COMMAND_RETRY_FAILED}:
+        return body
+    return None
+
+
+def parse_trusted_users(value: str = "") -> set[str]:
+    return {
+        user.casefold() for item in value.split(",") for user in item.split() if user
+    }
+
+
+def has_ready_label(pr: Mapping[str, Any]) -> bool:
+    return any(label["name"] in READY_LABELS for label in pr["labels"])
+
+
+def is_trusted_permission(permission: str) -> bool:
+    return permission in TRUSTED_PERMISSIONS
+
+
+def authorize(
+    *,
+    actor: str,
+    permission: str,
+    pr: Mapping[str, Any],
+    trusted_approval: bool = False,
+    trusted_users: set[str] | None = None,
+) -> tuple[bool, str]:
+    trusted_users = trusted_users or set()
+    if is_trusted_permission(permission):
+        return True, f"repository {permission} permission"
+    if actor.casefold() in trusted_users:
+        return True, "configured trusted contributor"
+    if actor.casefold() != pr["user"]["login"].casefold():
+        return (
+            False,
+            "Only reviewers with write access can run CI before it is "
+            "delegated to the PR author.",
+        )
+    if pr["draft"]:
+        return False, "PR authors cannot run CI while the PR is a draft."
+    if has_ready_label(pr):
+        return True, "ready label"
+    if trusted_approval:
+        return True, "approval from a trusted reviewer"
+    return (
+        False,
+        "A reviewer with write access must run `/ci run`, approve the PR, "
+        "or add the `ready` label first.",
+    )
+
+
+def has_trusted_approval(
+    github: GitHubClient,
+    number: int,
+    trusted_users: set[str],
+) -> bool:
+    if github.get_review_decision(number) != "APPROVED":
+        return False
+
+    latest_review_states: dict[str, tuple[str, str]] = {}
+    for review in github.list_reviews(number):
+        user = review.get("user") or {}
+        login = user.get("login")
+        state = review.get("state")
+        if login and state in {"APPROVED", "CHANGES_REQUESTED", "DISMISSED"}:
+            latest_review_states[login.casefold()] = (login, state)
+
+    for login, state in latest_review_states.values():
+        if state != "APPROVED":
+            continue
+        if login.casefold() in trusted_users:
+            return True
+        if is_trusted_permission(github.get_permission(login)):
+            return True
+    return False
+
+
+def is_build_for_pr(build: Mapping[str, Any], pr_number: int) -> bool:
+    pull_request = build.get("pull_request")
+    if isinstance(pull_request, Mapping):
+        build_pr_number = pull_request.get("id", pull_request.get("number"))
+        if build_pr_number is not None:
+            return str(build_pr_number) == str(pr_number)
+    metadata = build.get("meta_data") or {}
+    return str(metadata.get("github-pr-number")) == str(pr_number)
+
+
+def is_active_build(build: Mapping[str, Any]) -> bool:
+    return bool(build.get("blocked")) or build.get("state") in ACTIVE_BUILD_STATES
+
+
+def select_latest_build(
+    builds: Sequence[dict[str, Any]],
+    pr_number: int,
+) -> dict[str, Any] | None:
+    matching = [build for build in builds if is_build_for_pr(build, pr_number)]
+    return max(matching, key=lambda build: build.get("created_at", ""), default=None)
+
+
+def create_build_payload(
+    *,
+    actor: str,
+    comment_id: int,
+    pr: Mapping[str, Any],
+) -> dict[str, Any]:
+    return {
+        "commit": pr["head"]["sha"],
+        "branch": pr["head"]["ref"],
+        "message": f"PR #{pr['number']} {COMMAND_RUN_CI} by @{actor}",
+        "pull_request_id": pr["number"],
+        "pull_request_base_branch": pr["base"]["ref"],
+        "pull_request_repository": pr["head"]["repo"]["clone_url"],
+        "pull_request_labels": [label["name"] for label in pr["labels"]],
+        "ignore_pipeline_branch_filters": True,
+        "env": {
+            "VLLM_CI_GITHUB_COMMENT_ID": str(comment_id),
+            "VLLM_CI_TRIGGERED_BY": actor,
+        },
+        "meta_data": {
+            "github-comment-id": str(comment_id),
+            "github-pr-number": str(pr["number"]),
+            "github-triggered-by": actor,
+        },
+    }
+
+
+def add_reaction_safely(
+    github: GitHubClient,
+    comment_id: int,
+    content: str,
+) -> None:
+    try:
+        github.add_reaction(comment_id, content)
+    except Exception as error:
+        print(f"Could not add {content} reaction: {error}", file=sys.stderr)
+
+
+def is_already_handled(github: GitHubClient, comment_id: int) -> bool:
+    return any(
+        reaction.get("content") in {"rocket", "-1"}
+        and (reaction.get("user") or {}).get("login") == "github-actions[bot]"
+        for reaction in github.list_reactions(comment_id)
+    )
+
+
+def handle_run_ci(
+    *,
+    actor: str,
+    buildkite: BuildkiteClient,
+    comment_id: int,
+    github: GitHubClient,
+    pr: Mapping[str, Any],
+) -> str:
+    duplicate_builds = buildkite.list_builds(
+        pr["head"]["sha"],
+        metadata=("github-comment-id", str(comment_id)),
+    )
+    duplicate = select_latest_build(duplicate_builds, pr["number"])
+    if duplicate:
+        return f"CI was already requested by this comment: {duplicate['web_url']}"
+
+    current_builds = buildkite.list_builds(pr["head"]["sha"])
+    active_build = next(
+        (
+            build
+            for build in current_builds
+            if is_build_for_pr(build, pr["number"]) and is_active_build(build)
+        ),
+        None,
+    )
+    if active_build:
+        return f"CI is already running for this commit: {active_build['web_url']}"
+
+    current_pr = github.get_pr(pr["number"])
+    if current_pr["state"] != "open" or current_pr["head"]["sha"] != pr["head"]["sha"]:
+        return (
+            "The PR head changed while processing the command. Comment `/ci run` again."
+        )
+
+    build = buildkite.create_build(
+        create_build_payload(
+            actor=actor,
+            comment_id=comment_id,
+            pr=current_pr,
+        )
+    )
+    return (
+        f"Triggered [Buildkite CI #{build['number']}]({build['web_url']}) "
+        f"for commit `{current_pr['head']['sha'][:12]}`."
+    )
+
+
+def handle_retry_failed(
+    *,
+    buildkite: BuildkiteClient,
+    pr: Mapping[str, Any],
+) -> str:
+    builds = buildkite.list_builds(pr["head"]["sha"])
+    build = select_latest_build(builds, pr["number"])
+    if not build:
+        return "No CI build exists for the current PR commit. Use `/ci run` first."
+    if not build.get("finished_at") or is_active_build(build):
+        return f"CI is still running for this commit: {build['web_url']}"
+
+    retried = buildkite.retry_failed_jobs(build["number"], RETRY_STATES)
+    if retried["retried_jobs_count"] == 0:
+        return (
+            f"No failed, timed-out, or expired jobs need retrying: {build['web_url']}"
+        )
+    return (
+        f"Queued {retried['retried_jobs_count']} failed job(s) for retry in "
+        f"[Buildkite CI #{build['number']}]({build['web_url']})."
+    )
+
+
+def run(
+    event: Mapping[str, Any],
+    github: GitHubClient,
+    buildkite: BuildkiteClient,
+    trusted_users_value: str = "",
+) -> None:
+    command = parse_command(event["comment"]["body"])
+    if not command or "pull_request" not in event["issue"]:
+        return
+
+    issue_number = event["issue"]["number"]
+    comment_id = event["comment"]["id"]
+    actor = event["comment"]["user"]["login"]
+
+    if is_already_handled(github, comment_id):
+        print(f"Comment {comment_id} was already handled.")
+        return
+    add_reaction_safely(github, comment_id, "eyes")
+
+    try:
+        pr = github.get_pr(issue_number)
+        permission = github.get_permission(actor)
+        if pr["state"] != "open":
+            github.add_comment(issue_number, "CI commands require an open PR.")
+            return
+
+        trusted_users = parse_trusted_users(trusted_users_value)
+        should_check_approval = (
+            not is_trusted_permission(permission)
+            and actor.casefold() not in trusted_users
+            and actor.casefold() == pr["user"]["login"].casefold()
+            and not pr["draft"]
+            and not has_ready_label(pr)
+        )
+        trusted_approval = should_check_approval and has_trusted_approval(
+            github,
+            issue_number,
+            trusted_users,
+        )
+        allowed, reason = authorize(
+            actor=actor,
+            permission=permission,
+            pr=pr,
+            trusted_approval=trusted_approval,
+            trusted_users=trusted_users,
+        )
+        if not allowed:
+            add_reaction_safely(github, comment_id, "-1")
+            github.add_comment(issue_number, f"@{actor}, {reason}")
+            return
+
+        print(f"Authorized @{actor}: {reason}")
+        if command == COMMAND_RUN_CI:
+            message = handle_run_ci(
+                actor=actor,
+                buildkite=buildkite,
+                comment_id=comment_id,
+                github=github,
+                pr=pr,
+            )
+        else:
+            message = handle_retry_failed(buildkite=buildkite, pr=pr)
+        add_reaction_safely(github, comment_id, "rocket")
+        github.add_comment(issue_number, message)
+    except Exception:
+        add_reaction_safely(github, comment_id, "confused")
+        raise
+
+
+def main() -> None:
+    event_path = os.environ["GITHUB_EVENT_PATH"]
+    with open(event_path, encoding="utf-8") as event_file:
+        event = json.load(event_file)
+
+    if not parse_command(event["comment"]["body"]):
+        return
+
+    github = GitHubClient(
+        os.environ.get("GH_TOKEN", ""),
+        os.environ["GITHUB_REPOSITORY"],
+    )
+    buildkite = BuildkiteClient(
+        os.environ.get("BUILDKITE_API_TOKEN", ""),
+        os.environ.get("BUILDKITE_ORGANIZATION", "vllm"),
+        os.environ.get("BUILDKITE_PIPELINE", "ci"),
+    )
+    run(
+        event,
+        github,
+        buildkite,
+        os.environ.get("CI_TRUSTED_USERS", ""),
+    )
+
+
+if __name__ == "__main__":
+    main()

.github/workflows/scripts/test_run_ci_command.py

@@ -0,0 +1,363 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import unittest
+from typing import Any
+
+from run_ci_command import (
+    COMMAND_RETRY_FAILED,
+    COMMAND_RUN_CI,
+    RETRY_STATES,
+    BuildkiteClient,
+    authorize,
+    create_build_payload,
+    has_trusted_approval,
+    is_active_build,
+    is_build_for_pr,
+    parse_command,
+    parse_trusted_users,
+    run,
+    select_latest_build,
+)
+
+
+def make_pr(**overrides: Any) -> dict[str, Any]:
+    pr = {
+        "base": {"ref": "main"},
+        "draft": False,
+        "head": {
+            "ref": "feature",
+            "repo": {"clone_url": "https://github.com/contributor/vllm.git"},
+            "sha": "0123456789abcdef",
+        },
+        "labels": [],
+        "number": 42,
+        "state": "open",
+        "user": {"login": "author"},
+    }
+    pr.update(overrides)
+    return pr
+
+
+def make_event(command: str, actor: str = "reviewer") -> dict[str, Any]:
+    return {
+        "comment": {
+            "body": command,
+            "id": 99,
+            "user": {"login": actor},
+        },
+        "issue": {
+            "number": 42,
+            "pull_request": {},
+        },
+    }
+
+
+class FakeGitHub:
+    def __init__(
+        self,
+        *,
+        permission: str = "write",
+        permissions: dict[str, str] | None = None,
+        pr: dict[str, Any] | None = None,
+        review_decision: str = "REVIEW_REQUIRED",
+        reviews: list[dict[str, Any]] | None = None,
+    ) -> None:
+        self.comments: list[str] = []
+        self.permission = permission
+        self.permissions = permissions or {}
+        self.pr = pr or make_pr()
+        self.reactions: list[str] = []
+        self.review_decision = review_decision
+        self.reviews = reviews or []
+
+    def get_pr(self, number: int) -> dict[str, Any]:
+        return self.pr
+
+    def get_permission(self, actor: str) -> str:
+        return self.permissions.get(actor, self.permission)
+
+    def get_review_decision(self, number: int) -> str:
+        return self.review_decision
+
+    def list_reviews(self, number: int) -> list[dict[str, Any]]:
+        return self.reviews
+
+    def list_reactions(self, comment_id: int) -> list[dict[str, Any]]:
+        return []
+
+    def add_reaction(self, comment_id: int, content: str) -> None:
+        self.reactions.append(content)
+
+    def add_comment(self, issue_number: int, body: str) -> None:
+        self.comments.append(body)
+
+
+class FakeBuildkite:
+    def __init__(
+        self,
+        build_lists: list[list[dict[str, Any]]] | None = None,
+    ) -> None:
+        self.build_lists = build_lists or []
+        self.created_builds: list[dict[str, Any]] = []
+        self.list_calls: list[tuple[str, tuple[str, str] | None]] = []
+        self.retry_calls: list[tuple[int, str]] = []
+
+    def list_builds(
+        self,
+        commit: str,
+        *,
+        metadata: tuple[str, str] | None = None,
+    ) -> list[dict[str, Any]]:
+        self.list_calls.append((commit, metadata))
+        return self.build_lists.pop(0)
+
+    def create_build(self, body: dict[str, Any]) -> dict[str, Any]:
+        self.created_builds.append(body)
+        return {
+            "number": 123,
+            "web_url": "https://buildkite.example/builds/123",
+        }
+
+    def retry_failed_jobs(
+        self,
+        build_number: int,
+        states: str,
+    ) -> dict[str, Any]:
+        self.retry_calls.append((build_number, states))
+        return {"retried_jobs_count": 3}
+
+
+class FakeTransport:
+    def __init__(self, response: Any) -> None:
+        self.calls: list[dict[str, Any]] = []
+        self.response = response
+
+    def request(self, url: str, **kwargs: Any) -> Any:
+        self.calls.append({"url": url, **kwargs})
+        return self.response
+
+
+class RunCiCommandTest(unittest.TestCase):
+    def test_only_exact_ci_commands_are_accepted(self) -> None:
+        self.assertEqual(parse_command(COMMAND_RUN_CI), COMMAND_RUN_CI)
+        self.assertEqual(
+            parse_command(COMMAND_RETRY_FAILED),
+            COMMAND_RETRY_FAILED,
+        )
+        self.assertIsNone(parse_command("/ci run please"))
+        self.assertIsNone(parse_command(" /ci run"))
+
+    def test_write_access_authorizes_reviewers_and_authors(self) -> None:
+        allowed, _ = authorize(
+            actor="reviewer",
+            permission="write",
+            pr=make_pr(),
+        )
+        self.assertTrue(allowed)
+
+    def test_configured_trusted_contributors_can_run_ci(self) -> None:
+        trusted_users = parse_trusted_users("trusted-one, TRUSTED-TWO")
+        allowed, _ = authorize(
+            actor="trusted-two",
+            permission="read",
+            pr=make_pr(),
+            trusted_users=trusted_users,
+        )
+        self.assertTrue(allowed)
+
+    def test_authors_need_an_approval_or_ready_label(self) -> None:
+        pending, _ = authorize(
+            actor="author",
+            permission="read",
+            pr=make_pr(),
+        )
+        approved, _ = authorize(
+            actor="author",
+            permission="read",
+            pr=make_pr(),
+            trusted_approval=True,
+        )
+        ready, _ = authorize(
+            actor="author",
+            permission="read",
+            pr=make_pr(labels=[{"name": "ready"}]),
+        )
+        self.assertFalse(pending)
+        self.assertTrue(approved)
+        self.assertTrue(ready)
+
+    def test_non_author_contributors_without_write_are_denied(self) -> None:
+        allowed, _ = authorize(
+            actor="contributor",
+            permission="read",
+            pr=make_pr(),
+            trusted_approval=True,
+        )
+        self.assertFalse(allowed)
+
+    def test_authors_cannot_use_ready_state_on_draft_prs(self) -> None:
+        allowed, _ = authorize(
+            actor="author",
+            permission="read",
+            pr=make_pr(draft=True, labels=[{"name": "ready"}]),
+            trusted_approval=True,
+        )
+        self.assertFalse(allowed)
+
+    def test_only_trusted_reviewers_can_delegate_through_approval(self) -> None:
+        approved_review = {
+            "state": "APPROVED",
+            "user": {"login": "reviewer"},
+        }
+        trusted = FakeGitHub(
+            permission="read",
+            permissions={"reviewer": "write"},
+            review_decision="APPROVED",
+            reviews=[approved_review],
+        )
+        untrusted = FakeGitHub(
+            permission="read",
+            review_decision="APPROVED",
+            reviews=[approved_review],
+        )
+        self.assertTrue(has_trusted_approval(trusted, 42, set()))
+        self.assertFalse(has_trusted_approval(untrusted, 42, set()))
+
+    def test_build_matching_is_scoped_to_the_pr(self) -> None:
+        self.assertTrue(is_build_for_pr({"pull_request": {"id": 42}}, 42))
+        self.assertFalse(is_build_for_pr({"pull_request": {"id": 43}}, 42))
+        self.assertTrue(
+            is_build_for_pr(
+                {"meta_data": {"github-pr-number": "42"}},
+                42,
+            )
+        )
+
+    def test_latest_build_selection_ignores_other_prs(self) -> None:
+        latest = select_latest_build(
+            [
+                {
+                    "created_at": "2026-07-28T02:00:00Z",
+                    "number": 3,
+                    "pull_request": {"id": 43},
+                },
+                {
+                    "created_at": "2026-07-28T01:00:00Z",
+                    "number": 2,
+                    "pull_request": {"id": 42},
+                },
+                {
+                    "created_at": "2026-07-28T00:00:00Z",
+                    "number": 1,
+                    "pull_request": {"id": 42},
+                },
+            ],
+            42,
+        )
+        self.assertEqual(latest["number"], 2)
+
+    def test_active_build_states_prevent_duplicate_runs(self) -> None:
+        self.assertTrue(is_active_build({"state": "scheduled"}))
+        self.assertTrue(is_active_build({"state": "running"}))
+        self.assertTrue(is_active_build({"state": "waiting"}))
+        self.assertTrue(is_active_build({"blocked": True, "state": "passed"}))
+        self.assertFalse(is_active_build({"state": "failed"}))
+
+    def test_build_payload_preserves_pr_context(self) -> None:
+        payload = create_build_payload(
+            actor="reviewer",
+            comment_id=99,
+            pr=make_pr(labels=[{"name": "ready"}, {"name": "v1"}]),
+        )
+        self.assertEqual(
+            payload,
+            {
+                "commit": "0123456789abcdef",
+                "branch": "feature",
+                "message": "PR #42 /ci run by @reviewer",
+                "pull_request_id": 42,
+                "pull_request_base_branch": "main",
+                "pull_request_repository": ("https://github.com/contributor/vllm.git"),
+                "pull_request_labels": ["ready", "v1"],
+                "ignore_pipeline_branch_filters": True,
+                "env": {
+                    "VLLM_CI_GITHUB_COMMENT_ID": "99",
+                    "VLLM_CI_TRIGGERED_BY": "reviewer",
+                },
+                "meta_data": {
+                    "github-comment-id": "99",
+                    "github-pr-number": "42",
+                    "github-triggered-by": "reviewer",
+                },
+            },
+        )
+
+    def test_ci_run_dispatches_build_with_current_pr_metadata(self) -> None:
+        github = FakeGitHub()
+        buildkite = FakeBuildkite([[], []])
+        run(make_event(COMMAND_RUN_CI), github, buildkite)
+
+        self.assertEqual(len(buildkite.created_builds), 1)
+        self.assertEqual(
+            buildkite.created_builds[0]["message"],
+            "PR #42 /ci run by @reviewer",
+        )
+        self.assertEqual(github.reactions, ["eyes", "rocket"])
+        self.assertIn("Buildkite CI #123", github.comments[0])
+
+    def test_unapproved_authors_are_denied_without_buildkite(self) -> None:
+        github = FakeGitHub(
+            permission="read",
+            pr=make_pr(),
+            review_decision="REVIEW_REQUIRED",
+        )
+        buildkite = FakeBuildkite()
+        run(make_event(COMMAND_RUN_CI, "author"), github, buildkite)
+
+        self.assertEqual(buildkite.list_calls, [])
+        self.assertEqual(github.reactions, ["eyes", "-1"])
+        self.assertIn("approve the PR", github.comments[0])
+
+    def test_ci_retry_uses_latest_current_sha_build(self) -> None:
+        github = FakeGitHub(
+            permission="read",
+            pr=make_pr(labels=[{"name": "ready"}]),
+        )
+        buildkite = FakeBuildkite(
+            [
+                [
+                    {
+                        "created_at": "2026-07-28T01:00:00Z",
+                        "finished_at": "2026-07-28T02:00:00Z",
+                        "number": 123,
+                        "pull_request": {"id": 42},
+                        "state": "failed",
+                        "web_url": "https://buildkite.example/builds/123",
+                    }
+                ]
+            ]
+        )
+        run(make_event(COMMAND_RETRY_FAILED, "author"), github, buildkite)
+
+        self.assertEqual(buildkite.retry_calls, [(123, RETRY_STATES)])
+        self.assertIn("Queued 3 failed job", github.comments[0])
+
+    def test_buildkite_retry_uses_retry_failed_jobs_endpoint(self) -> None:
+        transport = FakeTransport({"retried_jobs_count": 2})
+        client = BuildkiteClient(
+            "secret",
+            "vllm",
+            "ci",
+            transport=transport,
+        )
+        client.retry_failed_jobs(123, RETRY_STATES)
+
+        call = transport.calls[0]
+        self.assertEqual(call["method"], "PUT")
+        self.assertTrue(call["url"].endswith("/123/retry_failed_jobs"))
+        self.assertEqual(call["body"], {"states": RETRY_STATES})
+
+
+if __name__ == "__main__":
+    unittest.main()

.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/cpu/cpu_attn.cpp

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

csrc/cpu/cpu_types_vxe.hpp

@@ -269,7 +269,7 @@ struct FP32Vec4 : public Vec<FP32Vec4> {
 
   explicit FP32Vec4(__vector float data) : reg(data) {}
 
-  explicit FP32Vec4(const FP32Vec4& data) : reg(data.reg) {}
+  FP32Vec4(const FP32Vec4& data) : reg(data.reg) {}
 };
 
 struct FP32Vec8 : public Vec<FP32Vec8> {
@@ -298,7 +298,7 @@ struct FP32Vec8 : public Vec<FP32Vec8> {
 
   explicit FP32Vec8(f32x4x2_t data) : reg(data) {}
 
-  explicit FP32Vec8(const FP32Vec8& data) {
+  FP32Vec8(const FP32Vec8& data) {
     reg.val[0] = data.reg.val[0];
     reg.val[1] = data.reg.val[1];
   }
@@ -643,7 +643,7 @@ struct FP32Vec16 : public Vec<FP32Vec16> {
 
   explicit FP32Vec16(f32x4x4_t data) : reg(data) {}
 
-  explicit FP32Vec16(const FP32Vec16& data) {
+  FP32Vec16(const FP32Vec16& data) {
     reg.val[0] = data.reg.val[0];
     reg.val[1] = data.reg.val[1];
     reg.val[2] = data.reg.val[2];

csrc/cpu/torch_bindings.cpp

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

csrc/custom_all_gather_reduce_scatter.cuh

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

csrc/custom_all_reduce.cuh

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

csrc/custom_collective_common.cuh

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

csrc/flashkda_registration.cpp

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

csrc/libtorch_stable/activation_kernels.cu

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

csrc/libtorch_stable/attention/merge_attn_states.cu

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

csrc/libtorch_stable/cache_kernels.cu

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

csrc/libtorch_stable/custom_all_gather_reduce_scatter.cu

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

csrc/libtorch_stable/custom_all_gather_reduce_scatter_ops.cpp

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

csrc/libtorch_stable/custom_all_reduce.cu

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

csrc/libtorch_stable/dsv3_fused_a_gemm.cu

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

csrc/libtorch_stable/moe/grouped_topk_kernels.cu

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

csrc/libtorch_stable/moe/moeTopKFuncs.cuh

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

csrc/libtorch_stable/moe/moe_align_sum_kernels.cu

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

csrc/libtorch_stable/ops.h

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

csrc/libtorch_stable/torch_bindings.cpp

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

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

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

csrc/spinloop.cpp

@@ -7,7 +7,7 @@ extern "C" {
 
 #if defined(__i386__) || defined(__x86_64__)
   #include <cpuid.h>
-  #include <mwaitxintrin.h>
+  #include <x86intrin.h>
 #endif
 
 #if defined(CLOCK_MONOTONIC_RAW)

docker/Dockerfile.s390x

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

docker/Dockerfile.xpu

@@ -63,21 +63,6 @@ RUN apt-get update -y && \
     python3-pip && \
     rm -rf /var/lib/apt/lists/*
 
-# Add oneAPI repo, pin oneAPI to 2025.3, then install pinned packages in one layer.
-RUN wget -O- https://apt.repos.intel.com/intel-gpg-keys/GPG-PUB-KEY-INTEL-SW-PRODUCTS.PUB | gpg --dearmor | tee /usr/share/keyrings/oneapi-archive-keyring.gpg > /dev/null && \
-    echo "deb [signed-by=/usr/share/keyrings/oneapi-archive-keyring.gpg] https://apt.repos.intel.com/oneapi all main" | tee /etc/apt/sources.list.d/oneAPI.list && \
-    printf '%s\n' \
-    'Package: intel-oneapi-* intel-deep-learning-essentials* intel-pti*' \
-    'Pin: version 2025.3*' \
-    'Pin-Priority: 1001' \
-    > /etc/apt/preferences.d/oneapi-2025.3.pref && \
-    apt-get update -y && \
-    apt-get install -y --no-install-recommends \
-    intel-oneapi-compiler-dpcpp-cpp-2025.3 \
-    intel-oneapi-mkl-devel-2025.3 \
-    intel-oneapi-dnnl-devel-2025.3 && \
-    rm -rf /var/lib/apt/lists/*
-
 # Install UMD
 RUN mkdir neo && \
     cd neo && \
@@ -100,22 +85,6 @@ RUN curl -LsSf https://astral.sh/uv/install.sh | sh \
     && uv venv --python ${PYTHON_VERSION} --seed ${VIRTUAL_ENV}
 ENV PATH="$VIRTUAL_ENV/bin:$PATH"
 
-# This oneccl contains the BMG support which is not the case for default version of oneapi 2025.3.
-ARG ONECCL_INSTALLER="intel-oneccl-2021.15.9.14_offline.sh"
-RUN wget "https://github.com/uxlfoundation/oneCCL/releases/download/2021.15.9/${ONECCL_INSTALLER}" && \
-    bash "${ONECCL_INSTALLER}" -a --silent --eula accept && \
-    rm "${ONECCL_INSTALLER}" && \
-    echo "source /opt/intel/oneapi/setvars.sh --force" >> /root/.bashrc && \
-    echo "source /opt/intel/oneapi/ccl/2021.15/env/vars.sh --force" >> /root/.bashrc && \
-    rm -f /opt/intel/oneapi/ccl/latest && \
-    ln -s /opt/intel/oneapi/ccl/2021.15 /opt/intel/oneapi/ccl/latest && \
-    printf '%s\n' \
-    '/opt/intel/oneapi/ccl/2021.15/lib' \
-    '/opt/intel/oneapi/mpi/2021.15/lib' \
-    '/opt/intel/oneapi/compiler/2025.3/lib' \
-    > /etc/ld.so.conf.d/oneapi-ccl.conf && \
-    ldconfig
-
 SHELL ["bash", "-c"]
 CMD ["bash", "-c", "source /root/.bashrc && exec bash"]
 
@@ -135,8 +104,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \
     uv pip install --upgrade pip
 
 
-
-ENV LD_LIBRARY_PATH=/opt/intel/oneapi/ccl/2021.15/lib:/opt/intel/oneapi/mpi/2021.15/lib:/opt/intel/oneapi/compiler/2025.3/lib:/usr/local/lib
+ENV LD_LIBRARY_PATH=/opt/venv/lib:/usr/local/lib
 CMD ["/bin/bash"]
 
 ######################### UCX + NIXL BUILD STAGE #########################
@@ -216,8 +184,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \
     uv pip install -r /workspace/vllm/requirements/xpu.txt && \
     uv pip install --no-build-isolation -r /workspace/vllm/requirements/test/xpu.txt && \
     uv pip uninstall triton triton-xpu && \
-    uv pip install triton-xpu==3.7.1 && \
-    uv pip uninstall oneccl oneccl-devel
+    uv pip install triton-xpu==3.7.2
 
 # Keep source-dependent layers near the end so frequent code-only changes
 # don't invalidate heavy dependency and UCX/NIXL layers.

docs/contributing/README.md

@@ -301,8 +301,10 @@ review process:
   isn't clear or you disagree with a suggestion, feel free to ask for
   clarification or discuss the suggestion.
 - Note that not all CI checks will be executed due to limited computational
-  resources. The reviewer will add `ready` label to the PR when the PR is
-  ready to merge or a full CI run is needed.
+  resources. Reviewers with write access and configured trusted contributors
+  can comment `/ci run` when CI signals are needed before a PR is ready. After
+  the PR is approved or has the `ready` label, the PR author can use `/ci run`
+  or `/ci retry`. New commits do not start CI automatically.
 
 ### Pull Request Limits and Escalation
 

docs/features/disagg_prefill.md

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

docs/features/kv_offloading_usage.md

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

docs/getting_started/installation/gpu.xpu.inc.md

@@ -42,12 +42,12 @@ pip install -v -r requirements/xpu.txt
 
     ```bash
     pip uninstall -y triton triton-xpu
-    pip install triton-xpu==3.7.1 --extra-index-url https://download.pytorch.org/whl/xpu
+    pip install triton-xpu==3.7.2 --extra-index-url https://download.pytorch.org/whl/xpu
     ```
 
     !!! note
         - `triton` (without suffix) is for NVIDIA GPUs only. On XPU, using it instead of `triton-xpu` can cause correctness or runtime issues.
-        - For torch 2.12 (the version used in `requirements/xpu.txt`), the matching package is `triton-xpu==3.7.1`. If you use a different version of torch, check the corresponding `triton-xpu` version in [docker/Dockerfile.xpu](https://github.com/vllm-project/vllm/blob/main/docker/Dockerfile.xpu).
+        - For torch 2.13 (the version used in `requirements/xpu.txt`), the matching package is `triton-xpu==3.7.2`. If you use a different version of torch, check the corresponding `triton-xpu` version in [docker/Dockerfile.xpu](https://github.com/vllm-project/vllm/blob/main/docker/Dockerfile.xpu).
 
 - Finally, build and install vLLM XPU backend:
 

pyproject.toml

@@ -129,10 +129,9 @@ extend-exclude = ["tests/models/fixtures/*", "tests/prompts/*", "tests/tokenizer
     "tests/entrypoints/speech_to_text/transcription/test_transcription_validation.py",
     "docs/governance/process.md", "docs/assets/contributing/vllm_bench_serve_timeline.html", 
     "tests/v1/engine/test_fast_incdec_prefix_err.py", ".git/*", "csrc/cpu/sgl-kernels/*",
-    "rust/src/chat/src/renderer/deepseek_v32/fixtures/*",
-    "rust/src/parser/src/tool/gemma4.rs", "rust/src/parser/src/unified/gemma4.rs",
+    "rust/src/chat/src/renderer/deepseek_v32/fixtures/*", "rust/src/parser/**",
     "rust/src/text/src/output/decoded.rs",
-    "rust/src/tokenizer/src/incremental.rs", "rust/src/parser/src/reasoning/tests.rs"]
+    "rust/src/tokenizer/src/incremental.rs"]
 ignore-hidden = false
 
 [tool.typos.default]

requirements/test/xpu.txt

@@ -140,7 +140,7 @@ docopt==0.6.2
     # via num2words
 docstring-parser==0.18.0
     # via anthropic
-dpcpp-cpp-rt==2025.3.2
+dpcpp-cpp-rt==2026.0.0
     # via
     #   onemkl-sycl-blas
     #   onemkl-sycl-dft
@@ -253,27 +253,27 @@ ijson==3.5.0
     # via -r requirements/test/../common.txt
 imageio==2.37.3
     # via scikit-image
-impi-rt==2021.17.2
+impi-rt==2021.18.0
     # via
     #   oneccl
     #   torch
 iniconfig==2.3.0
     # via pytest
-intel-cmplr-lib-rt==2025.3.2
+intel-cmplr-lib-rt==2026.0.0
     # via
     #   intel-sycl-rt
     #   torch
-intel-cmplr-lib-ur==2025.3.2
+intel-cmplr-lib-ur==2026.0.0
     # via
     #   intel-openmp
     #   intel-sycl-rt
     #   torch
-intel-cmplr-lic-rt==2025.3.2
+intel-cmplr-lic-rt==2026.0.0
     # via
     #   intel-opencl-rt
     #   intel-sycl-rt
     #   torch
-intel-opencl-rt==2025.3.2
+intel-opencl-rt==2026.0.0
     # via
     #   dpcpp-cpp-rt
     #   onemkl-sycl-blas
@@ -282,14 +282,14 @@ intel-opencl-rt==2025.3.2
     #   onemkl-sycl-rng
     #   onemkl-sycl-sparse
     #   torch
-intel-openmp==2025.3.2
+intel-openmp==2026.0.0
     # via
     #   dpcpp-cpp-rt
     #   mkl
     #   torch
-intel-pti==0.16.0
+intel-pti==0.17.0
     # via torch
-intel-sycl-rt==2025.3.2
+intel-sycl-rt==2026.0.0
     # via
     #   dpcpp-cpp-rt
     #   oneccl
@@ -378,7 +378,7 @@ mistral-common==1.11.5
     #   -c requirements/common.txt
     #   -r requirements/test/../common.txt
     #   -r requirements/test/xpu.in
-mkl==2025.3.1
+mkl==2026.0.0
     # via
     #   onemkl-sycl-blas
     #   onemkl-sycl-dft
@@ -453,28 +453,28 @@ numpy==2.2.6
     #   torchvision
     #   transformers
     #   xgrammar
-oneccl==2021.17.2
+oneccl==2022.0.0
     # via
     #   oneccl-devel
     #   torch
-oneccl-devel==2021.17.2
+oneccl-devel==2022.0.0
     # via torch
-onemkl-license==2025.3.1
+onemkl-license==2026.0.0
     # via
     #   mkl
     #   torch
-onemkl-sycl-blas==2025.3.1
+onemkl-sycl-blas==2026.0.0
     # via
     #   onemkl-sycl-lapack
     #   onemkl-sycl-sparse
     #   torch
-onemkl-sycl-dft==2025.3.1
+onemkl-sycl-dft==2026.0.0
     # via torch
-onemkl-sycl-lapack==2025.3.1
+onemkl-sycl-lapack==2026.0.0
     # via torch
-onemkl-sycl-rng==2025.3.1
+onemkl-sycl-rng==2026.0.0
     # via torch
-onemkl-sycl-sparse==2025.3.1
+onemkl-sycl-sparse==2026.0.0
     # via torch
 openai==2.44.0
     # via
@@ -719,6 +719,8 @@ pyyaml==6.0.3
     #   timm
     #   transformers
     #   uvicorn
+pyzes==0.1.1
+    # via torch
 pyzmq==27.1.0
     # via
     #   -c requirements/common.txt
@@ -871,14 +873,14 @@ tabledata==1.3.4
     # via pytablewriter
 tabulate==0.10.0
     # via sacrebleu
-tbb==2022.3.1
+tbb==2023.0.0
     # via
     #   intel-opencl-rt
     #   mkl
     #   torch
 tblib==3.1.0
     # via -r requirements/test/xpu.in
-tcmlib==1.4.1
+tcmlib==1.5.0
     # via
     #   tbb
     #   torch
@@ -910,7 +912,7 @@ tokenizers==0.22.2
     #   -c requirements/common.txt
     #   -r requirements/test/../common.txt
     #   transformers
-torch==2.12.0+xpu
+torch==2.13.0+xpu
     # via
     #   -c requirements/xpu.txt
     #   accelerate
@@ -920,7 +922,7 @@ torch==2.12.0+xpu
     #   timm
     #   torchvision
     #   xgrammar
-torchvision==0.27.0+xpu
+torchvision==0.28.0+xpu
     # via timm
 tqdm==4.67.3
     # via
@@ -946,7 +948,7 @@ transformers==5.14.1
     #   xgrammar
 triton==3.7.1
     # via xgrammar
-triton-xpu==3.7.1
+triton-xpu==3.7.2
     # via torch
 typepy==1.3.4
     # via
@@ -1001,7 +1003,7 @@ typing-inspection==0.4.2
     #   mcp
     #   pydantic
     #   pydantic-settings
-umf==1.0.3
+umf==1.1.0
     # via
     #   intel-cmplr-lib-ur
     #   torch

requirements/xpu.txt

@@ -12,10 +12,10 @@ jinja2>=3.1.6
 datasets # for benchmark scripts
 numba == 0.65.0 # Required for N-gram speculative decoding
 --extra-index-url=https://download.pytorch.org/whl/xpu
-torch==2.12.0
+torch==2.13.0
 torchaudio
 torchvision
 torchcodec >= 0.14 # Required for the torchcodec video decoding backend
 
-auto_round_lib==0.14.1
-vllm_xpu_kernels @ https://github.com/vllm-project/vllm-xpu-kernels/releases/download/v0.1.11.1/vllm_xpu_kernels-0.1.11.1-cp38-abi3-manylinux_2_28_x86_64.whl
+auto_round_lib==0.14.2
+vllm_xpu_kernels @ https://github.com/vllm-project/vllm-xpu-kernels/releases/download/v0.1.12/vllm_xpu_kernels-0.1.12-cp38-abi3-manylinux_2_28_x86_64.whl

rust/Cargo.lock

@@ -2220,7 +2220,7 @@ checksum = "11d3d7f243d5c5a8b9bb5d6dd2b1602c0cb0b9db1621bafc7ed66e35ff9fe092"
 [[package]]
 name = "llm-multimodal"
 version = "1.7.1"
-source = "git+https://github.com/smg-project/llm-multimodal?rev=5390032d6dc8a3e6fdc83acd320260367eb4b9b5#5390032d6dc8a3e6fdc83acd320260367eb4b9b5"
+source = "git+https://github.com/smg-project/llm-multimodal?rev=15adba5e025d8636ba4a334fb379b1371f6196a1#15adba5e025d8636ba4a334fb379b1371f6196a1"
 dependencies = [
  "anyhow",
  "base64 0.22.1",
@@ -5503,9 +5503,9 @@ dependencies = [
  "tokio",
  "tokio-stream",
  "tracing",
- "tracing-subscriber",
  "url",
  "uuid",
+ "vllm-tracing",
 ]
 
 [[package]]
@@ -5570,17 +5570,16 @@ dependencies = [
  "serde_json",
  "serde_with",
  "thiserror-ext",
- "time",
  "tokio",
  "tokio-util",
  "tracing",
- "tracing-subscriber",
  "uuid",
  "vllm-bench",
  "vllm-chat",
  "vllm-engine-core-client",
  "vllm-managed-engine",
  "vllm-server",
+ "vllm-tracing",
 ]
 
 [[package]]
@@ -5827,6 +5826,15 @@ dependencies = [
  "vllm-parser",
 ]
 
+[[package]]
+name = "vllm-tracing"
+version = "0.1.0"
+dependencies = [
+ "time",
+ "tracing",
+ "tracing-subscriber",
+]
+
 [[package]]
 name = "walkdir"
 version = "2.5.0"

rust/Cargo.toml

@@ -13,6 +13,7 @@ members = [
     "src/server",
     "src/text",
     "src/tokenizer",
+    "src/tracing",
 ]
 resolver = "3"
 
@@ -58,7 +59,7 @@ indexmap = "2.13.0"
 indicatif = "0.18.4"
 itertools = "0.14.0"
 libc = "0.2.177"
-llm-multimodal = { git = "https://github.com/smg-project/llm-multimodal", rev = "5390032d6dc8a3e6fdc83acd320260367eb4b9b5", default-features = false, features = ["native-tls"] }
+llm-multimodal = { git = "https://github.com/smg-project/llm-multimodal", rev = "15adba5e025d8636ba4a334fb379b1371f6196a1", default-features = false, features = ["native-tls"] }
 mimalloc = "0.1.52"
 minijinja = { version = "2.0", features = ["unstable_machinery", "json", "builtins", "loader", "loop_controls", "preserve_order"] }
 minijinja-contrib = { version = "2.0", features = ["pycompat"] }
@@ -143,6 +144,7 @@ vllm-parser = { path = "src/parser" }
 vllm-server = { path = "src/server" }
 vllm-text = { path = "src/text" }
 vllm-tokenizer = { path = "src/tokenizer" }
+vllm-tracing = { path = "src/tracing" }
 winnow = { version = "1.0.2", features = ["simd"] }
 xgrammar-structural-tag = "0.2.0"
 zeromq = { version = "0.6.0", default-features = false, features = [

rust/src/bench/Cargo.toml

@@ -32,9 +32,9 @@ tokenizers.workspace = true
 tokio.workspace = true
 tokio-stream.workspace = true
 tracing.workspace = true
-tracing-subscriber.workspace = true
 url.workspace = true
 uuid.workspace = true
+vllm-tracing.workspace = true
 
 [lints]
 workspace = true

rust/src/bench/src/main.rs

@@ -19,18 +19,8 @@ struct Cli {
     args: vllm_bench::BenchServeArgs,
 }
 
-// TODO: unify the tracing subscriber used by different binaries.
-fn init_tracing() {
-    let filter = tracing_subscriber::EnvFilter::try_from_default_env()
-        .unwrap_or_else(|_| tracing_subscriber::EnvFilter::new("info"));
-    let _ = tracing_subscriber::fmt()
-        .with_env_filter(filter)
-        .with_writer(std::io::stderr)
-        .try_init();
-}
-
 fn main() -> anyhow::Result<()> {
-    init_tracing();
+    vllm_tracing::init_tracing("Bench");
 
     let cli = Cli::parse();
     vllm_bench::prepare_process();

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

@@ -20,7 +20,7 @@ use crate::output::{
 use crate::renderer::hf::{HfChatRenderer, MultimodalRenderInfo};
 use crate::renderer::{
     DeepSeekV4ChatRenderer, DeepSeekV32ChatRenderer, DynChatRenderer, HarmonyChatRenderer,
-    InklingChatRenderer,
+    InklingChatRenderer, KimiK3ChatRenderer,
 };
 use crate::request::ChatRequest;
 use crate::{DynChatOutputProcessor, RendererSelection};
@@ -57,6 +57,7 @@ impl HfChatBackend {
                     processor_config: files.processor_config_path.as_deref(),
                 },
                 tokenizer.clone(),
+                options.limit_mm_per_prompt.clone(),
             )?
         };
         let multimodal_render_info = resolve_multimodal_render_info(multimodal_model_info.as_ref());
@@ -73,6 +74,7 @@ impl HfChatBackend {
             RendererSelection::DeepSeekV4 => Arc::new(DeepSeekV4ChatRenderer::new()),
             RendererSelection::Harmony => Arc::new(HarmonyChatRenderer::new()?),
             RendererSelection::Inkling => Arc::new(InklingChatRenderer::new(tokenizer.clone())?),
+            RendererSelection::KimiK3 => Arc::new(KimiK3ChatRenderer::new(tokenizer.clone())),
         };
 
         info!(
@@ -230,6 +232,7 @@ mod tests {
                 chat_template_content_format: Default::default(),
                 chat_template: None,
                 default_chat_template_kwargs: HashMap::new(),
+                limit_mm_per_prompt: HashMap::new(),
             },
             test_tokenizer(),
         )

rust/src/chat/src/backend/mod.rs

@@ -8,7 +8,7 @@ use serde_json::Value;
 use vllm_text::{DynTextBackend, TextBackend};
 
 use crate::error::Result;
-use crate::multimodal::MultimodalModelInfo;
+use crate::multimodal::{MmLimitPerPrompt, MultimodalModelInfo};
 use crate::output::DynChatOutputProcessor;
 use crate::renderer::DynChatRenderer;
 use crate::request::ChatRequest;
@@ -74,6 +74,9 @@ pub struct LoadModelBackendsOptions {
     /// Optional server-default keyword arguments merged into every
     /// chat-template render before request-level `chat_template_kwargs`.
     pub default_chat_template_kwargs: HashMap<String, Value>,
+    /// Maximum number of input items allowed per prompt for each modality.
+    /// Unspecified modalities are unlimited.
+    pub limit_mm_per_prompt: MmLimitPerPrompt,
 }
 
 /// Shared backends loaded from a model id.

rust/src/chat/src/error.rs

@@ -23,6 +23,8 @@ pub enum Error {
     UnsupportedMultimodalContent(&'static str),
     #[error("`{modality}` input is not supported by this model")]
     UnsupportedModality { modality: String },
+    #[error("At most {limit} {modality}(s) may be provided in one prompt.")]
+    MmLimitExceeded { modality: String, limit: usize },
     #[error("multimodal preprocessing error: {0}")]
     Multimodal(#[message] String),
     #[error("{kind} parsing is not available for model `{model_id}`")]
@@ -87,7 +89,8 @@ impl Error {
             Self::Text(error) => error.is_request_validation_error(),
             Self::UnsupportedMultimodalRenderer
             | Self::UnsupportedMultimodalContent(_)
-            | Self::UnsupportedModality { .. } => true,
+            | Self::UnsupportedModality { .. }
+            | Self::MmLimitExceeded { .. } => true,
 
             _ => false,
         }

rust/src/chat/src/lib.rs

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

rust/src/chat/src/multimodal.rs

@@ -11,7 +11,7 @@
 //! Raw media stays above `vllm-text`; this module lowers it into token IDs and
 //! opaque tensor payloads before the request is handed to text generation.
 
-use std::collections::HashSet;
+use std::collections::{BTreeMap, HashMap, HashSet};
 use std::fs;
 use std::path::Path;
 use std::sync::{Arc, LazyLock};
@@ -24,6 +24,7 @@ use llm_multimodal::{
     PromptReplacement, Tokenizer as TokenResolver, TrackedMedia, VideoClip, VisionPreProcessor,
     VisionProcessorRegistry,
 };
+use serde::{Deserialize, Serialize};
 use thiserror_ext::AsReport as _;
 use tracing::warn;
 use vllm_engine_core_client::protocol::dtype::ModelDtype;
@@ -52,6 +53,71 @@ pub struct MultimodalModelInfo {
     video: Option<ModalitySupport>,
     audio: Option<AudioModalitySupport>,
     media_connector: Arc<MediaConnector>,
+    /// Maximum number of input items allowed per prompt for each modality.
+    limit_mm_per_prompt: MmLimitPerPrompt,
+}
+
+/// Per-modality item-count limits configured by `--limit-mm-per-prompt`.
+///
+/// Modalities absent from the map are unlimited.
+pub type MmLimitPerPrompt = HashMap<MmLimitModality, MmLimitSpec>;
+
+/// Modalities that `--limit-mm-per-prompt` can be keyed by.
+///
+/// Closed on purpose: these are exactly the keys Python accepts, per
+/// `MultiModalDummyOptionsBuiltins` in `vllm/config/multimodal.py`.
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
+#[serde(rename_all = "snake_case")]
+pub enum MmLimitModality {
+    Image,
+    Audio,
+    Video,
+}
+
+impl MmLimitModality {
+    /// The wire name, matching Python's modality strings.
+    pub fn as_str(self) -> &'static str {
+        match self {
+            Self::Image => "image",
+            Self::Audio => "audio",
+            Self::Video => "video",
+        }
+    }
+}
+
+/// One modality's limit, in either of the two shapes Python accepts.
+#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
+#[serde(
+    untagged,
+    expecting = "an item count, or an object with an optional `count` field"
+)]
+pub enum MmLimitSpec {
+    /// Legacy form: `"image": 16`
+    Count(usize),
+
+    /// Configurable form:
+    /// `"video": {"count": 1, "num_frames": 32}`
+    Options {
+        /// Absent means unlimited, matching an absent modality.
+        #[serde(default, skip_serializing_if = "Option::is_none")]
+        count: Option<usize>,
+
+        /// Preserve Python-owned options for forwarding. Never interpreted
+        /// here: they size the engine's dummy-profiling encoder cache, which
+        /// has no Rust counterpart.
+        #[serde(flatten)]
+        extra: BTreeMap<String, serde_json::Value>,
+    },
+}
+
+impl MmLimitSpec {
+    /// The configured item count, or `None` when this modality is unlimited.
+    pub fn count(&self) -> Option<usize> {
+        match self {
+            Self::Count(count) => Some(*count),
+            Self::Options { count, .. } => *count,
+        }
+    }
 }
 
 /// Model metadata and tokenizer access shared by all multimodal specs.
@@ -287,6 +353,7 @@ impl MultimodalModelInfo {
         model_type: Option<String>,
         files: MultimodalConfigFiles<'_>,
         tokenizer: DynTokenizer,
+        limit_mm_per_prompt: MmLimitPerPrompt,
     ) -> Result<Option<Self>> {
         let config = match files.config {
             Some(path) => {
@@ -319,7 +386,12 @@ impl MultimodalModelInfo {
             tokenizer: TokenizerResolver(tokenizer),
         };
 
-        Self::from_loaded(context, preprocessor_config, video_preprocessor_config)
+        Self::from_loaded(
+            context,
+            preprocessor_config,
+            video_preprocessor_config,
+            limit_mm_per_prompt,
+        )
     }
 
     /// Resolve multimodal support from an assembled context and parsed
@@ -328,6 +400,7 @@ impl MultimodalModelInfo {
         context: MultimodalModelContext,
         preprocessor_config: PreProcessorConfig,
         video_preprocessor_config: PreProcessorConfig,
+        limit_mm_per_prompt: MmLimitPerPrompt,
     ) -> Result<Option<Self>> {
         let (image, video) = Self::resolve_vision_lanes(
             &context,
@@ -356,6 +429,7 @@ impl MultimodalModelInfo {
             video,
             audio,
             media_connector,
+            limit_mm_per_prompt,
         }))
     }
 
@@ -567,7 +641,55 @@ fn input_audio_data_url(data: &str, format: Option<&str>) -> Result<String> {
     Ok(format!("data:{mime_type};base64,{data}"))
 }
 
+/// The modality a content part counts against, or `None` for plain text.
+///
+/// Embedding inputs share their base modality's budget rather than getting one
+/// of their own, matching Python's `modality.replace("_embeds", "")` in
+/// `vllm/entrypoints/chat_utils.py`.
+fn media_part_limit_modality(part: &MediaContentPart) -> Option<MmLimitModality> {
+    match part {
+        MediaContentPart::Text { .. } => None,
+        MediaContentPart::ImageUrl { .. }
+        | MediaContentPart::ImageData { .. }
+        | MediaContentPart::ImageEmbeds { .. } => Some(MmLimitModality::Image),
+        MediaContentPart::AudioUrl { .. } | MediaContentPart::AudioData { .. } => {
+            Some(MmLimitModality::Audio)
+        }
+        MediaContentPart::VideoUrl { .. } | MediaContentPart::VideoData { .. } => {
+            Some(MmLimitModality::Video)
+        }
+    }
+}
+
 impl MultimodalModelInfo {
+    /// Reject requests exceeding `--limit-mm-per-prompt`'s configured
+    /// per-modality item count, before any fetch/decode work is spent on them.
+    ///
+    /// Modalities without a configured count are unlimited.
+    fn validate_mm_limits(&self, media_parts: &[MediaContentPart]) -> Result<()> {
+        let mut counts: HashMap<MmLimitModality, usize> = HashMap::new();
+        for part in media_parts {
+            if let Some(modality) = media_part_limit_modality(part) {
+                *counts.entry(modality).or_default() += 1;
+            }
+        }
+
+        for (modality, count) in counts {
+            let Some(limit) = self.limit_mm_per_prompt.get(&modality).and_then(MmLimitSpec::count)
+            else {
+                continue;
+            };
+            if count > limit {
+                return Err(Error::MmLimitExceeded {
+                    modality: modality.as_str().to_string(),
+                    limit,
+                });
+            }
+        }
+
+        Ok(())
+    }
+
     /// Run media fetch, per-modality preprocessing, prompt expansion, and
     /// feature build.
     ///
@@ -584,8 +706,7 @@ impl MultimodalModelInfo {
         if media_parts_len == 0 {
             return Ok(Vec::new());
         }
-        // TODO: enforce per-modality item-count limits, aligned with the
-        // engine's `--limit-mm-per-prompt` semantics.
+        self.validate_mm_limits(&media_parts)?;
         let fetched = self.fetch_media(media_parts).await?;
 
         let mut prepared = Vec::new();
@@ -753,10 +874,11 @@ mod tests {
             .with_regular_token("<|video_pad|>", QWEN3_VIDEO_PAD_ID)
     }
 
-    fn test_info(
+    fn test_info_with_limits(
         model_type: &str,
         config: serde_json::Value,
         tokenizer: TestTokenizer,
+        limit_mm_per_prompt: MmLimitPerPrompt,
     ) -> MultimodalModelInfo {
         let context = MultimodalModelContext {
             model_id: format!("{model_type}-test"),
@@ -769,11 +891,20 @@ mod tests {
             context,
             PreProcessorConfig::default(),
             PreProcessorConfig::default(),
+            limit_mm_per_prompt,
         )
         .unwrap()
         .unwrap_or_else(|| panic!("{model_type} multimodal support should resolve"))
     }
 
+    fn test_info(
+        model_type: &str,
+        config: serde_json::Value,
+        tokenizer: TestTokenizer,
+    ) -> MultimodalModelInfo {
+        test_info_with_limits(model_type, config, tokenizer, HashMap::new())
+    }
+
     fn llama4_info() -> MultimodalModelInfo {
         let config = serde_json::json!({
             "model_type": "llama4",
@@ -783,16 +914,19 @@ mod tests {
         test_info("llama4", config, llama4_tokenizer())
     }
 
-    pub(super) fn qwen3_vl_info() -> MultimodalModelInfo {
-        let config = serde_json::json!({
+    fn qwen3_vl_config() -> serde_json::Value {
+        serde_json::json!({
             "model_type": "qwen3_vl",
             "image_token_id": QWEN3_IMAGE_PAD_ID,
             "video_token_id": QWEN3_VIDEO_PAD_ID,
             "vision_start_token_id": 151652,
             "vision_end_token_id": 151653,
             "vision_config": {"patch_size": 16}
-        });
-        test_info("qwen3_vl", config, qwen3_vl_tokenizer())
+        })
+    }
+
+    pub(super) fn qwen3_vl_info() -> MultimodalModelInfo {
+        test_info("qwen3_vl", qwen3_vl_config(), qwen3_vl_tokenizer())
     }
 
     #[test]
@@ -853,4 +987,123 @@ mod tests {
         );
         assert!(input_audio_data_url("AAAA", Some("flac")).is_err());
     }
+
+    fn image_url_part() -> MediaContentPart {
+        MediaContentPart::ImageUrl {
+            url: "https://example.com/image.png".to_string(),
+            detail: None,
+            uuid: None,
+        }
+    }
+
+    fn qwen3_vl_info_with_limits(limit_mm_per_prompt: MmLimitPerPrompt) -> MultimodalModelInfo {
+        test_info_with_limits(
+            "qwen3_vl",
+            qwen3_vl_config(),
+            qwen3_vl_tokenizer(),
+            limit_mm_per_prompt,
+        )
+    }
+
+    #[test]
+    fn validate_mm_limits_ignores_text_parts() {
+        let info = qwen3_vl_info();
+        let parts = vec![
+            MediaContentPart::Text {
+                text: "hello".to_string(),
+            },
+            MediaContentPart::Text {
+                text: "world".to_string(),
+            },
+        ];
+        assert!(info.validate_mm_limits(&parts).is_ok());
+    }
+
+    #[test]
+    fn validate_mm_limits_leaves_unconfigured_modalities_unlimited() {
+        let info = qwen3_vl_info();
+        let parts: Vec<_> = std::iter::repeat_with(image_url_part).take(1_000).collect();
+
+        assert!(info.validate_mm_limits(&parts).is_ok());
+    }
+
+    #[test]
+    fn validate_mm_limits_enforces_configured_limit_at_the_boundary() {
+        let info = qwen3_vl_info_with_limits(HashMap::from([(
+            MmLimitModality::Image,
+            MmLimitSpec::Count(1),
+        )]));
+        assert!(info.validate_mm_limits(&[image_url_part()]).is_ok());
+
+        let error = info.validate_mm_limits(&[image_url_part(), image_url_part()]).unwrap_err();
+        assert_eq!(
+            error.to_report_string(),
+            "At most 1 image(s) may be provided in one prompt."
+        );
+        // Confirms the HTTP-mapping bug found during implementation stays fixed:
+        // this must map to 400, not 500.
+        assert!(error.is_request_validation_error());
+    }
+
+    #[test]
+    fn validate_mm_limits_counts_image_embeds_against_the_image_limit() {
+        let info = qwen3_vl_info_with_limits(HashMap::from([(
+            MmLimitModality::Image,
+            MmLimitSpec::Count(1),
+        )]));
+        let image_embeds_part = MediaContentPart::ImageEmbeds {
+            payload: serde_json::Value::String("AAAA".to_string()),
+            uuid: None,
+        };
+
+        let error = info.validate_mm_limits(&[image_url_part(), image_embeds_part]).unwrap_err();
+        assert_eq!(
+            error.to_report_string(),
+            "At most 1 image(s) may be provided in one prompt."
+        );
+    }
+
+    /// An options object without a `count` carries only profiling keys, which
+    /// say nothing about how many items are allowed.
+    #[test]
+    fn validate_mm_limits_treats_a_count_less_options_object_as_unlimited() {
+        let info = qwen3_vl_info_with_limits(HashMap::from([(
+            MmLimitModality::Image,
+            MmLimitSpec::Options {
+                count: None,
+                extra: BTreeMap::from([("width".to_string(), serde_json::json!(512))]),
+            },
+        )]));
+
+        assert!(info.validate_mm_limits(&[image_url_part(), image_url_part()]).is_ok());
+    }
+
+    fn parse_limits(json: &str) -> MmLimitPerPrompt {
+        serde_json::from_str(json).expect("limit map should parse")
+    }
+
+    #[test]
+    fn limit_map_parses_both_shapes_python_accepts() {
+        let limits = parse_limits(r#"{"image": 16, "video": {"count": 1, "num_frames": 32}}"#);
+
+        assert_eq!(limits[&MmLimitModality::Image].count(), Some(16));
+        assert_eq!(limits[&MmLimitModality::Video].count(), Some(1));
+        assert_eq!(limits.get(&MmLimitModality::Audio), None);
+    }
+
+    #[test]
+    fn limit_map_rejects_keys_python_does_not_accept() {
+        assert!(serde_json::from_str::<MmLimitPerPrompt>(r#"{"image_embeds": 1}"#).is_err());
+    }
+
+    /// Managed mode forwards this map back to Python as JSON, where
+    /// `BaseDummyOptions.count` is a non-optional `int` under `extra="forbid"`.
+    /// Emitting `"count": null` would make the engine subprocess fail to start.
+    #[test]
+    fn limit_map_round_trips_without_emitting_a_null_count() {
+        let source = r#"{"video":{"num_frames":32}}"#;
+        let encoded = serde_json::to_string(&parse_limits(source)).expect("map should serialize");
+
+        assert_eq!(encoded, source);
+    }
 }

rust/src/chat/src/multimodal/audio.rs

@@ -107,6 +107,7 @@ mod tests {
             context,
             PreProcessorConfig::default(),
             PreProcessorConfig::default(),
+            HashMap::new(),
         )
         .unwrap()
         .expect("Inkling multimodal support")
@@ -126,6 +127,7 @@ mod tests {
             context,
             PreProcessorConfig::default(),
             PreProcessorConfig::default(),
+            HashMap::new(),
         )
         .unwrap()
         .expect("Qwen3-ASR multimodal support")

rust/src/chat/src/multimodal/expand.rs

@@ -222,7 +222,9 @@ mod tests {
         assert_eq!(tensor.shape, vec![expected.len()]);
         assert_eq!(
             tensor.data,
-            WireArrayData::RawView(expected.iter().map(|value| u8::from(*value)).collect())
+            WireArrayData::RawView(
+                expected.iter().map(|value| u8::from(*value)).collect::<Vec<_>>().into(),
+            )
         );
     }
 

rust/src/chat/src/multimodal/tensor.rs

@@ -2,28 +2,58 @@
 // SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 use std::collections::HashMap;
+use std::mem::size_of;
 
 use half::{bf16, f16};
 use llm_multimodal::{ModelSpecificValue, PreprocessedEncoderInputs};
 use vllm_engine_core_client::protocol::dtype::ModelDtype;
 use vllm_engine_core_client::protocol::multimodal::MmKwargValue as ProtocolKwargValue;
-use vllm_engine_core_client::protocol::tensor::{ShapeExt as _, WireTensor};
+use vllm_engine_core_client::protocol::tensor::{ShapeExt as _, WireArrayData, WireTensor};
 
 use crate::error::{Error, Result, bail_multimodal, multimodal};
 
+/// Element type retained alongside an encoded tensor during multimodal lowering.
+#[derive(Debug, Clone, Copy)]
+pub(super) enum TensorKind {
+    /// 32-bit floating point.
+    F32,
+    /// IEEE 16-bit floating point.
+    F16,
+    /// Brain floating point.
+    Bf16,
+    /// Signed 64-bit integer.
+    I64,
+    /// Unsigned 32-bit integer.
+    U32,
+}
+
+impl TensorKind {
+    const fn element_size(self) -> usize {
+        match self {
+            Self::F32 => size_of::<f32>(),
+            Self::F16 => size_of::<f16>(),
+            Self::Bf16 => size_of::<bf16>(),
+            Self::I64 => size_of::<i64>(),
+            Self::U32 => size_of::<u32>(),
+        }
+    }
+
+    const fn wire_dtype(self) -> &'static str {
+        match self {
+            Self::F32 => "float32",
+            Self::F16 => "float16",
+            Self::Bf16 => "bfloat16",
+            Self::I64 => "int64",
+            Self::U32 => "uint32",
+        }
+    }
+}
+
 /// Representation for multimodal kwarg values for transformation.
 #[derive(Debug)]
 pub(super) enum KwargValue {
-    /// Float tensor with row-major flat data and shape.
-    F32Tensor { data: Vec<f32>, shape: Vec<usize> },
-    /// Float16 tensor with row-major flat data and shape.
-    F16Tensor { data: Vec<f16>, shape: Vec<usize> },
-    /// BFloat16 tensor with row-major flat data and shape.
-    Bf16Tensor { data: Vec<bf16>, shape: Vec<usize> },
-    /// Signed integer tensor with row-major flat data and shape.
-    I64Tensor { data: Vec<i64>, shape: Vec<usize> },
-    /// Unsigned integer tensor with row-major flat data and shape.
-    U32Tensor { data: Vec<u32>, shape: Vec<usize> },
+    /// Tensor with row-major flat data and shape.
+    Tensor { kind: TensorKind, wire: WireTensor },
     /// Non-tensor kwarg value that is shared or copied as-is.
     Passthrough(ProtocolKwargValue),
 }
@@ -67,8 +97,14 @@ impl KwargValue {
             ModelSpecificValue::Tensor { data, shape } => {
                 Self::from_f32_tensor(data, shape, float_dtype)?
             }
-            ModelSpecificValue::IntTensor { data, shape } => Self::I64Tensor { data, shape },
-            ModelSpecificValue::UintTensor { data, shape } => Self::U32Tensor { data, shape },
+            ModelSpecificValue::IntTensor { data, shape } => {
+                let wire = WireTensor::from_i64(shape, data).map_err(Error::Multimodal)?;
+                Self::tensor(TensorKind::I64, wire)
+            }
+            ModelSpecificValue::UintTensor { data, shape } => {
+                let wire = WireTensor::from_u32(shape, data).map_err(Error::Multimodal)?;
+                Self::tensor(TensorKind::U32, wire)
+            }
             ModelSpecificValue::Int(value) => Self::Passthrough(Int(value)),
             ModelSpecificValue::Float(value) => Self::Passthrough(Float(value)),
             ModelSpecificValue::IntVec(values) => {
@@ -93,17 +129,23 @@ impl KwargValue {
     /// Convert a float tensor to the target float dtype if needed, keeping the
     /// same shape.
     fn from_f32_tensor(data: Vec<f32>, shape: Vec<usize>, float_dtype: ModelDtype) -> Result<Self> {
-        match float_dtype {
-            ModelDtype::Float16 => Ok(Self::F16Tensor {
-                data: data.into_iter().map(f16::from_f32).collect(),
-                shape,
-            }),
-            ModelDtype::BFloat16 => Ok(Self::Bf16Tensor {
-                data: data.into_iter().map(bf16::from_f32).collect(),
-                shape,
-            }),
-            ModelDtype::Float32 => Ok(Self::F32Tensor { data, shape }),
-        }
+        let (kind, wire) = match float_dtype {
+            ModelDtype::Float16 => (
+                TensorKind::F16,
+                WireTensor::from_f16(shape, data.into_iter().map(f16::from_f32).collect()),
+            ),
+            ModelDtype::BFloat16 => (
+                TensorKind::Bf16,
+                WireTensor::from_bf16(shape, data.into_iter().map(bf16::from_f32).collect()),
+            ),
+            ModelDtype::Float32 => (TensorKind::F32, WireTensor::from_f32(shape, data)),
+        };
+        wire.map(|wire| Self::tensor(kind, wire)).map_err(Error::Multimodal)
+    }
+
+    fn tensor(kind: TensorKind, wire: WireTensor) -> Self {
+        debug_assert_eq!(wire.dtype, kind.wire_dtype());
+        Self::Tensor { kind, wire }
     }
 }
 
@@ -111,15 +153,11 @@ impl TryFrom<&KwargValue> for ProtocolKwargValue {
     type Error = Error;
 
     fn try_from(value: &KwargValue) -> Result<Self> {
-        let tensor = match value {
-            KwargValue::F32Tensor { data, shape } => WireTensor::from_f32(shape.clone(), data),
-            KwargValue::F16Tensor { data, shape } => WireTensor::from_f16(shape.clone(), data),
-            KwargValue::Bf16Tensor { data, shape } => WireTensor::from_bf16(shape.clone(), data),
-            KwargValue::I64Tensor { data, shape } => WireTensor::from_i64(shape.clone(), data),
-            KwargValue::U32Tensor { data, shape } => WireTensor::from_u32(shape.clone(), data),
+        let wire = match value {
+            KwargValue::Tensor { wire, .. } => wire.clone(),
             KwargValue::Passthrough(value) => return Ok(value.clone()),
         };
-        tensor.map(ProtocolKwargValue::Tensor).map_err(Error::Multimodal)
+        Ok(ProtocolKwargValue::Tensor(wire))
     }
 }
 
@@ -127,11 +165,7 @@ impl KwargValue {
     /// First-axis length for tensor values; `None` for passthrough kwargs.
     pub(super) fn first_dim(&self) -> Option<usize> {
         match self {
-            Self::F32Tensor { shape, .. }
-            | Self::F16Tensor { shape, .. }
-            | Self::Bf16Tensor { shape, .. }
-            | Self::I64Tensor { shape, .. }
-            | Self::U32Tensor { shape, .. } => shape.first().copied(),
+            Self::Tensor { wire, .. } => wire.shape.first().copied(),
             Self::Passthrough(_) => None,
         }
     }
@@ -161,30 +195,13 @@ impl KwargValue {
         end: usize,
         drop_axis: bool,
     ) -> Result<ProtocolKwargValue> {
-        let tensor = match self {
-            Self::F32Tensor { data, shape } => {
-                let (shape, data) = slice_first_axis_range(shape, data, start, end, drop_axis)?;
-                WireTensor::from_f32(shape, data)
-            }
-            Self::F16Tensor { data, shape } => {
-                let (shape, data) = slice_first_axis_range(shape, data, start, end, drop_axis)?;
-                WireTensor::from_f16(shape, data)
-            }
-            Self::Bf16Tensor { data, shape } => {
-                let (shape, data) = slice_first_axis_range(shape, data, start, end, drop_axis)?;
-                WireTensor::from_bf16(shape, data)
-            }
-            Self::I64Tensor { data, shape } => {
-                let (shape, data) = slice_first_axis_range(shape, data, start, end, drop_axis)?;
-                WireTensor::from_i64(shape, data)
-            }
-            Self::U32Tensor { data, shape } => {
-                let (shape, data) = slice_first_axis_range(shape, data, start, end, drop_axis)?;
-                WireTensor::from_u32(shape, data)
+        let wire = match self {
+            Self::Tensor { kind, wire } => {
+                slice_first_axis_range(wire, kind.element_size(), start, end, drop_axis)
             }
             Self::Passthrough(value) => return Ok(value.clone()),
         };
-        tensor.map(ProtocolKwargValue::Tensor).map_err(Error::Multimodal)
+        wire.map(ProtocolKwargValue::Tensor)
     }
 }
 
@@ -208,44 +225,53 @@ pub(super) fn flat_range_for_index(
 /// Read a tensor value as per-image sizes for flat slicing.
 fn tensor_as_usize_vec(tensor: &KwargValue) -> Result<Vec<usize>> {
     match tensor {
-        KwargValue::I64Tensor { data, .. } => data
-            .iter()
+        KwargValue::Tensor {
+            kind: TensorKind::I64,
+            wire,
+        } => raw_tensor_bytes(wire, size_of::<i64>())?
+            .chunks_exact(size_of::<i64>())
+            .map(|bytes| i64::from_ne_bytes(bytes.try_into().expect("exact int64 chunk")))
             .map(|value| {
-                usize::try_from(*value)
+                usize::try_from(value)
                     .map_err(|_| multimodal!("negative flat tensor size `{value}`"))
             })
             .collect(),
-        KwargValue::U32Tensor { data, .. } => {
-            Ok(data.iter().map(|value| *value as usize).collect())
-        }
+        KwargValue::Tensor {
+            kind: TensorKind::U32,
+            wire,
+        } => Ok(raw_tensor_bytes(wire, size_of::<u32>())?
+            .chunks_exact(size_of::<u32>())
+            .map(|bytes| u32::from_ne_bytes(bytes.try_into().expect("exact uint32 chunk")) as usize)
+            .collect()),
         _ => Err(multimodal!("flat tensor sizes must be int64 or uint32")),
     }
 }
 
 /// Slice a flat row-major tensor along its first axis.
-fn slice_first_axis_range<'a, T>(
-    shape: &[usize],
-    data: &'a [T],
+fn slice_first_axis_range(
+    tensor: &WireTensor,
+    element_size: usize,
     start: usize,
     end: usize,
     drop_axis: bool,
-) -> Result<(Vec<usize>, &'a [T])> {
+) -> Result<WireTensor> {
+    let shape = tensor.shape.as_slice();
+    raw_tensor_bytes(tensor, element_size)?;
     let first_dim = *shape.first().ok_or_else(|| multimodal!("tensor has no first dimension"))?;
     if start > end || end > first_dim {
         bail_multimodal!("invalid tensor slice {start}..{end} for first dimension {first_dim}");
     }
-    let expected_len = shape
-        .checked_numel()
-        .ok_or_else(|| multimodal!("tensor shape {shape:?} has too many elements"))?;
-    if expected_len != data.len() {
-        bail_multimodal!(
-            "tensor shape {shape:?} expects {expected_len} elements, got {}",
-            data.len()
-        );
-    }
-    let stride = shape[1..].iter().product::<usize>();
-    let data_start = start * stride;
-    let data_end = end * stride;
+    let stride = shape[1..]
+        .iter()
+        .try_fold(1usize, |acc, dim| acc.checked_mul(*dim))
+        .and_then(|stride| stride.checked_mul(element_size))
+        .ok_or_else(|| multimodal!("tensor shape {shape:?} byte stride overflowed usize"))?;
+    let data_start = start
+        .checked_mul(stride)
+        .ok_or_else(|| multimodal!("tensor slice start byte offset overflowed usize"))?;
+    let data_end = end
+        .checked_mul(stride)
+        .ok_or_else(|| multimodal!("tensor slice end byte offset overflowed usize"))?;
     let out_shape = if drop_axis {
         shape[1..].to_vec()
     } else {
@@ -253,7 +279,38 @@ fn slice_first_axis_range<'a, T>(
         shape[0] = end - start;
         shape
     };
-    Ok((out_shape, &data[data_start..data_end]))
+    let WireArrayData::RawView(data) = &tensor.data else {
+        return Err(multimodal!("cannot slice an aux tensor buffer"));
+    };
+    Ok(WireTensor::from_raw_bytes(
+        tensor.dtype.clone(),
+        out_shape,
+        data.slice(data_start..data_end),
+    ))
+}
+
+fn raw_tensor_bytes(tensor: &WireTensor, element_size: usize) -> Result<&[u8]> {
+    let WireArrayData::RawView(data) = &tensor.data else {
+        return Err(multimodal!("expected an inline tensor buffer"));
+    };
+    let expected_bytes = tensor
+        .shape
+        .checked_numel()
+        .and_then(|numel| numel.checked_mul(element_size))
+        .ok_or_else(|| {
+            multimodal!(
+                "tensor shape {:?} byte length overflowed usize",
+                tensor.shape
+            )
+        })?;
+    if expected_bytes != data.len() {
+        bail_multimodal!(
+            "tensor shape {:?} expects {expected_bytes} bytes, got {}",
+            tensor.shape,
+            data.len()
+        );
+    }
+    Ok(data)
 }
 
 #[cfg(test)]
@@ -262,28 +319,32 @@ mod tests {
 
     #[test]
     fn batched_wire_value_at_drops_first_axis() {
-        let value = KwargValue::F32Tensor {
-            data: vec![1.0, 2.0, 3.0, 4.0],
-            shape: vec![2, 2],
-        };
+        let data = vec![1.0_f32, 2.0, 3.0, 4.0];
+        let expected_ptr = data.as_ptr().cast::<u8>().wrapping_add(2 * size_of::<f32>());
+        let value = KwargValue::tensor(
+            TensorKind::F32,
+            WireTensor::from_f32(vec![2, 2], data).unwrap(),
+        );
 
         let ProtocolKwargValue::Tensor(tensor) = value.batched_wire_value_at(1).unwrap() else {
             panic!("expected tensor");
         };
 
         assert_eq!(tensor.shape, vec![2]);
+        let raw_view = tensor.data.into_raw_view().unwrap();
+        assert_eq!(raw_view.as_ptr(), expected_ptr);
         assert_eq!(
-            tensor.data.into_raw_view().unwrap(),
+            raw_view,
             [3.0_f32, 4.0].into_iter().flat_map(f32::to_ne_bytes).collect::<Vec<_>>()
         );
     }
 
     #[test]
     fn flat_wire_value_range_keeps_first_axis() {
-        let value = KwargValue::U32Tensor {
-            data: (0..10).collect(),
-            shape: vec![5, 2],
-        };
+        let value = KwargValue::tensor(
+            TensorKind::U32,
+            WireTensor::from_u32(vec![5, 2], (0..10_u32).collect()).unwrap(),
+        );
 
         let ProtocolKwargValue::Tensor(tensor) = value.flat_wire_value_range(1, 3).unwrap() else {
             panic!("expected tensor");
@@ -298,10 +359,10 @@ mod tests {
 
     #[test]
     fn flat_range_for_index_uses_size_tensor() {
-        let sizes = KwargValue::I64Tensor {
-            data: vec![2, 3, 4],
-            shape: vec![3],
-        };
+        let sizes = KwargValue::tensor(
+            TensorKind::I64,
+            WireTensor::from_i64(vec![3], vec![2_i64, 3, 4]).unwrap(),
+        );
 
         assert_eq!(
             flat_range_for_index(&sizes, "image_grid_thw", 1).unwrap(),
@@ -311,10 +372,11 @@ mod tests {
 
     #[test]
     fn slice_first_axis_range_errors_on_shape_data_mismatch() {
-        let error = slice_first_axis_range(&[2, 2], &[1.0_f32, 2.0, 3.0], 0, 1, true).unwrap_err();
+        let tensor = WireTensor::from_raw("float32", vec![2, 2], vec![0; 3 * size_of::<f32>()]);
+        let error = slice_first_axis_range(&tensor, size_of::<f32>(), 0, 1, true).unwrap_err();
 
         assert!(
-            matches!(error, Error::Multimodal(message) if message.contains("expects 4 elements"))
+            matches!(error, Error::Multimodal(message) if message.contains("expects 16 bytes"))
         );
     }
 

rust/src/chat/src/multimodal/video.rs

@@ -207,6 +207,7 @@ mod tests {
                 Some("qwen3_vl".to_string()),
                 files,
                 Arc::new(qwen3_vl_tokenizer()),
+                std::collections::HashMap::new(),
             )
         };
 

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

@@ -0,0 +1,267 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+//! Golden fixtures generated from HF remote-code `encoding_k3.py`.
+
+use std::path::PathBuf;
+use std::sync::Arc;
+
+use expect_test::{expect, expect_file};
+use serde_json::json;
+use vllm_text::Prompt;
+use vllm_text::tokenizer::DynTokenizer;
+use vllm_tokenizer::Tokenizer;
+use vllm_tokenizer::test_utils::TestTokenizer;
+
+use super::KimiK3ChatRenderer;
+use crate::AssistantContentBlock;
+use crate::ChatRenderer;
+use crate::renderer::kimi_k3::encoding::{CLOSE, END_OF_MSG, IMAGE_PLACEHOLDER, OPEN, SEP};
+use crate::renderer::test_utils::{FixtureRequestOptions, fixture_chat_request};
+use crate::request::{ChatContentPart, ChatMessage, GenerationPromptMode, ReasoningEffort};
+
+const OPEN_ID: u32 = 256;
+const CLOSE_ID: u32 = 257;
+const SEP_ID: u32 = 258;
+const END_OF_MSG_ID: u32 = 259;
+const MEDIA_ID: u32 = 260;
+
+fn test_tokenizer() -> TestTokenizer {
+    TestTokenizer::new()
+        .with_special_token(OPEN, OPEN_ID)
+        .with_special_token(CLOSE, CLOSE_ID)
+        .with_special_token(SEP, SEP_ID)
+        .with_special_token(END_OF_MSG, END_OF_MSG_ID)
+        .with_special_token(IMAGE_PLACEHOLDER, MEDIA_ID)
+}
+
+fn render_token_ids(request: &crate::request::ChatRequest, tokenizer: DynTokenizer) -> Vec<u32> {
+    let prompt = KimiK3ChatRenderer::new(tokenizer).render(request).unwrap().prompt;
+    let Prompt::TokenIds(token_ids) = prompt else {
+        panic!("kimi k3 renderer should return token IDs")
+    };
+    token_ids
+}
+
+fn render_request(request: &crate::request::ChatRequest) -> String {
+    let tokenizer: DynTokenizer = Arc::new(test_tokenizer());
+    let token_ids = render_token_ids(request, tokenizer.clone());
+    tokenizer.decode(&token_ids, false).unwrap()
+}
+
+fn fixture_path(name: &str) -> PathBuf {
+    PathBuf::from(env!("CARGO_MANIFEST_DIR"))
+        .join("src/renderer/kimi_k3/fixtures")
+        .join(name)
+}
+
+fn kimi_k3_fixture_options() -> FixtureRequestOptions {
+    FixtureRequestOptions {
+        // Fixture JSON owns thinking via `template_kwargs`.
+        enable_thinking: None,
+        no_generation_prompt_when_last_assistant: false,
+    }
+}
+
+fn assert_golden(name: &str) {
+    let input_name = format!("{name}_input.json");
+    let request = fixture_chat_request(&fixture_path(&input_name), kimi_k3_fixture_options());
+    let rendered = render_request(&request);
+    expect_file![format!("fixtures/{name}_output.txt")].assert_eq(&rendered);
+}
+
+#[test]
+fn golden_history_preserve_and_image() {
+    assert_golden("history_preserve_and_image");
+}
+
+#[test]
+fn golden_tools_history_and_required() {
+    assert_golden("tools_history_and_required");
+}
+
+#[test]
+fn golden_controls_thinking_off() {
+    assert_golden("controls_thinking_off");
+}
+
+#[test]
+fn golden_dynamic_system_tool_declare() {
+    assert_golden("dynamic_system_tool_declare");
+}
+
+#[test]
+fn token_writer_protects_literal_control_and_media_markers() {
+    let tokenizer = Arc::new(test_tokenizer());
+    let user_text = format!("literal {OPEN} and {}", super::encoding::IMAGE_PLACEHOLDER);
+    let mut request = crate::request::ChatRequest::for_test();
+    request.messages = vec![ChatMessage::user(vec![
+        ChatContentPart::text(user_text),
+        ChatContentPart::image_url("data:image/png;base64,test"),
+    ])];
+    request
+        .chat_options
+        .template_kwargs
+        .insert("thinking".to_string(), json!(false));
+    request.chat_options.generation_prompt_mode = GenerationPromptMode::NoGenerationPrompt;
+
+    let token_ids = render_token_ids(&request, tokenizer.clone());
+
+    assert_eq!(
+        token_ids.iter().filter(|&&token_id| token_id == OPEN_ID).count(),
+        1
+    );
+    assert_eq!(
+        token_ids.iter().filter(|&&token_id| token_id == MEDIA_ID).count(),
+        1
+    );
+
+    let flattened = tokenizer.decode(&token_ids, false).unwrap();
+    let flattened_ids = tokenizer.encode(&flattened, false).unwrap();
+    assert_eq!(
+        flattened_ids.iter().filter(|&&token_id| token_id == OPEN_ID).count(),
+        2
+    );
+    assert_eq!(
+        flattened_ids.iter().filter(|&&token_id| token_id == MEDIA_ID).count(),
+        2
+    );
+}
+
+#[test]
+fn thinking_history_renders_empty_think_channel() {
+    let mut request = crate::request::ChatRequest::for_test();
+    request.messages = vec![
+        ChatMessage::user("question"),
+        ChatMessage::assistant_text("answer"),
+        ChatMessage::user("follow-up"),
+    ];
+    request.chat_options.generation_prompt_mode = GenerationPromptMode::NoGenerationPrompt;
+
+    let rendered = render_request(&request);
+
+    assert!(rendered.contains(
+        "<|open|>message role=\"assistant\"<|sep|>\
+         <|open|>think<|sep|><|close|>think<|sep|>\
+         <|open|>response<|sep|>answer<|close|>response<|sep|>"
+    ));
+}
+
+#[test]
+fn non_thinking_history_omits_reasoning_channel() {
+    let mut request = crate::request::ChatRequest::for_test();
+    request.messages = vec![
+        ChatMessage::user("question"),
+        ChatMessage::assistant_blocks(vec![
+            AssistantContentBlock::Reasoning {
+                text: "hidden reasoning".to_string(),
+            },
+            AssistantContentBlock::Text {
+                text: "answer".to_string(),
+            },
+        ]),
+        ChatMessage::user("follow-up"),
+    ];
+    request
+        .chat_options
+        .template_kwargs
+        .insert("thinking".to_string(), json!(false));
+    request.chat_options.generation_prompt_mode = GenerationPromptMode::NoGenerationPrompt;
+
+    let rendered = render_request(&request);
+
+    assert!(!rendered.contains("hidden reasoning"));
+    assert!(!rendered.contains("<|open|>think<|sep|>"));
+    assert!(rendered.contains(
+        "<|open|>message role=\"assistant\"<|sep|>\
+         <|open|>response<|sep|>answer<|close|>response<|sep|>"
+    ));
+}
+
+#[test]
+fn defaults_thinking_effort_to_max() {
+    let rendered = render_request(&crate::request::ChatRequest::for_test());
+
+    expect![[r#"<|open|>message role="system" type="thinking-effort"<|sep|>`thinking_effort` guides on how much to think in your thinking channel (not including the response channel), supported values include `low`, `medium`, `high`, and `max`.
+Now the system is invoked with `thinking_effort=max`.<|close|>message<|sep|><|end_of_msg|><|open|>message role="user"<|sep|>test<|close|>message<|sep|><|end_of_msg|><|open|>message role="assistant"<|sep|><|open|>think<|sep|>"#]]
+        .assert_eq(&rendered);
+}
+
+#[test]
+fn translates_standard_thinking_kwargs() {
+    let mut request = crate::request::ChatRequest::for_test();
+    request
+        .chat_options
+        .template_kwargs
+        .insert("enable_thinking".to_string(), json!(true));
+    request
+        .chat_options
+        .template_kwargs
+        .insert("reasoning_effort".to_string(), json!("high"));
+
+    let rendered = render_request(&request);
+
+    assert!(rendered.contains("thinking_effort=high"));
+    assert!(rendered.ends_with("<|open|>think<|sep|>"));
+}
+
+#[test]
+fn native_k3_kwargs_take_precedence() {
+    let mut request = crate::request::ChatRequest::for_test();
+    request.chat_options.template_kwargs.extend([
+        ("thinking".to_string(), json!(true)),
+        ("enable_thinking".to_string(), json!(false)),
+        ("thinking_effort".to_string(), json!("low")),
+        ("reasoning_effort".to_string(), json!("high")),
+    ]);
+
+    let rendered = render_request(&request);
+
+    assert!(rendered.contains("thinking_effort=low"));
+    assert!(!rendered.contains("thinking_effort=high"));
+}
+
+#[test]
+fn standard_none_disables_thinking() {
+    let mut request = crate::request::ChatRequest::for_test();
+    request.chat_options.template_kwargs.extend([
+        ("enable_thinking".to_string(), json!(false)),
+        ("reasoning_effort".to_string(), json!("none")),
+    ]);
+
+    let rendered = render_request(&request);
+
+    assert!(!rendered.contains("type=\"thinking-effort\""));
+    assert!(rendered.ends_with("<|open|>response<|sep|>"));
+}
+
+#[test]
+fn typed_none_disables_thinking() {
+    let mut request = crate::request::ChatRequest::for_test();
+    request.chat_options.reasoning_effort = Some(ReasoningEffort::None);
+
+    let rendered = render_request(&request);
+
+    assert!(!rendered.contains("type=\"thinking-effort\""));
+    assert!(rendered.ends_with("<|open|>response<|sep|>"));
+}
+
+#[test]
+fn rejects_removed_medium_thinking_effort() {
+    let mut request = crate::request::ChatRequest::for_test();
+    request
+        .chat_options
+        .template_kwargs
+        .insert("thinking_effort".to_string(), json!("medium"));
+
+    let error = KimiK3ChatRenderer::new(Arc::new(test_tokenizer()))
+        .render(&request)
+        .unwrap_err();
+
+    expect![[r#"
+        ChatTemplate(
+            "unsupported thinking_effort=\"medium\"; supported values are `low`, `high`, and `max`",
+        )
+    "#]]
+    .assert_debug_eq(&error);
+}

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

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

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

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

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

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

rust/src/chat/src/request.rs

@@ -397,6 +397,10 @@ pub struct ChatOptions {
     /// Effort level exposed to chat templates for reasoning models.
     pub reasoning_effort: Option<ReasoningEffort>,
 
+    /// Standard response format available to model-specific renderers.
+    #[serde(default)]
+    pub response_format: Option<Value>,
+
     /// Additional keyword arguments exposed to the chat template.
     pub template_kwargs: HashMap<String, Value>,
 }
@@ -407,6 +411,7 @@ impl Default for ChatOptions {
             generation_prompt_mode: GenerationPromptMode::StartNewAssistant,
             chat_template: None,
             reasoning_effort: None,
+            response_format: None,
             template_kwargs: HashMap::new(),
         }
     }

rust/src/chat/tests/roundtrip.rs

@@ -212,6 +212,23 @@ impl RoundtripCase {
         }
     }
 
+    /// Kimi K3 XTML tool/reasoning channels (native renderer + unified parser).
+    ///
+    /// Needs HF tokenizer files under `HF_HOME` (`tiktoken.model` +
+    /// `tokenizer_config.json`). Weights are not required for this text-level
+    /// roundtrip.
+    fn kimi_k3() -> Self {
+        Self {
+            model_id: "moonshotai/Kimi-K3",
+            assistant_stop_suffix: "<|end_of_msg|>",
+            tool_call_parser: ParserSelection::Auto,
+            reasoning_parser: ParserSelection::Auto,
+            thinking_behavior: ThinkingBehavior::Toggleable { default: true },
+            json_fmt: compact_json_fmt(),
+            sort_json_keys: false,
+        }
+    }
+
     /// SeedOSS with `<seed:think>` / `</seed:think>` reasoning tags.
     fn seed_oss() -> Self {
         Self {
@@ -255,7 +272,7 @@ impl RoundtripCase {
     fn gpt_oss() -> Self {
         Self {
             model_id: "openai/gpt-oss-20b",
-            assistant_stop_suffix: "", // not applicable for token-id cases
+            assistant_stop_suffix: "",
             tool_call_parser: ParserSelection::Auto,
             reasoning_parser: ParserSelection::Auto,
             thinking_behavior: ThinkingBehavior::Always { value: true },
@@ -268,7 +285,7 @@ impl RoundtripCase {
     fn inkling() -> Self {
         Self {
             model_id: "thinkingmachines/Inkling",
-            assistant_stop_suffix: "",
+            assistant_stop_suffix: "<|content_model_end_sampling|>",
             tool_call_parser: ParserSelection::Auto,
             reasoning_parser: ParserSelection::Auto,
             thinking_behavior: ThinkingBehavior::Always { value: true },
@@ -312,6 +329,8 @@ roundtrip_tests! {
     nemotron_v3 => [reasoning_and_content],
     gemma4 => [tool_call_mix], // Gemma4 strips reasoning in history if there's no tool call
     kimi_k25 => [tool_call_mix], // Kimi K2.5 strips reasoning in history
+    // K3 drops plain-assistant reasoning in history; tool-call turns keep it.
+    kimi_k3 => [tool_call_mix],
     gpt_oss => [tool_call_mix], // Harmony strips reasoning in history if there's no tool call
     inkling => [reasoning_and_content, tool_call_mix],
 }
@@ -670,14 +689,23 @@ fn decoded_completion_stream(
                 .collect()
         }
         Prompt::TokenIds(token_ids) => {
-            ensure!(
-                assistant_stop_suffix.is_empty(),
-                "token-id roundtrip cases do not support text stop suffixes"
-            );
+            let body = if assistant_stop_suffix.is_empty() {
+                token_ids.as_slice()
+            } else {
+                let stop_token_ids = tokenizer
+                    .encode(assistant_stop_suffix, false)
+                    .context("failed to encode token-id completion stop suffix")?;
+                token_ids.strip_suffix(stop_token_ids.as_slice()).with_context(|| {
+                    format!(
+                        "token-id completion did not end with {:?}: {:?}",
+                        assistant_stop_suffix, token_ids
+                    )
+                })?
+            };
             incremental_decode_chunks(
                 tokenizer,
                 &prompt_token_ids,
-                token_ids,
+                body,
                 TOKEN_COMPLETION_CHUNK_TOKENS,
             )?
         }

rust/src/cmd/Cargo.toml

@@ -23,17 +23,16 @@ serde.workspace = true
 serde_json.workspace = true
 serde_with.workspace = true
 thiserror-ext.workspace = true
-time.workspace = true
 tokio = { workspace = true, features = ["signal"] }
 tokio-util.workspace = true
 tracing.workspace = true
-tracing-subscriber.workspace = true
 uuid.workspace = true
 vllm-bench.workspace = true
 vllm-chat.workspace = true
 vllm-engine-core-client.workspace = true
 vllm-managed-engine.workspace = true
 vllm-server.workspace = true
+vllm-tracing.workspace = true
 
 [dev-dependencies]
 expect-test.workspace = true

rust/src/cmd/src/cli.rs

@@ -23,6 +23,7 @@ use serde_with::{DefaultOnNull, OneOrMany, serde_as};
 use thiserror_ext::AsReport as _;
 use uuid::Uuid;
 use vllm_chat::ReasoningParserFactory;
+use vllm_chat::multimodal::MmLimitPerPrompt;
 use vllm_engine_core_client::TransportMode;
 use vllm_managed_engine::ManagedEngineConfig;
 use vllm_managed_engine::cli::{ManagedEngineArgs, repartition_managed_engine_args};
@@ -186,6 +187,17 @@ pub struct SharedRuntimeArgs {
     #[serde(default)]
     pub default_chat_template_kwargs: Option<HashMap<String, Value>>,
 
+    /// The maximum number of input items allowed per prompt for each
+    /// modality, as a JSON object (e.g. `{"image": 16, "video": 2}`).
+    ///
+    /// Also accepts the engine's configurable form
+    /// (e.g. `{"video": {"count": 1, "num_frames": 32}}`); the extra
+    /// profiling options are forwarded to the engine untouched.
+    /// Unspecified modalities are unlimited.
+    #[arg(long, value_parser = parse_json::<MmLimitPerPrompt>, value_name = "JSON", default_value = "{}")]
+    #[serde(default)]
+    pub limit_mm_per_prompt: MmLimitPerPrompt,
+
     /// The format to render message content within a chat template.
     ///
     /// * "auto" detects the format from the template
@@ -348,6 +360,18 @@ impl SharedRuntimeArgs {
             .expect("profiler config serialization should not fail")
     }
 
+    /// Return the per-modality limits as JSON for managed Python engine
+    /// forwarding, or `None` when nothing is configured.
+    ///
+    /// Round-tripping the parsed map rather than the raw argument keeps the
+    /// engine's own profiling options (`num_frames`, `width`, ...) intact.
+    pub fn limit_mm_per_prompt_json(&self) -> Option<String> {
+        (!self.limit_mm_per_prompt.is_empty()).then(|| {
+            serde_json::to_string(&self.limit_mm_per_prompt)
+                .expect("limit-mm-per-prompt serialization should not fail")
+        })
+    }
+
     /// Apply fallback logic for API key configuration from env variables.
     fn apply_env_api_key_fallback(&mut self) {
         if self.api_key.is_empty()
@@ -399,6 +423,7 @@ impl SharedRuntimeArgs {
             language_model_only: self.language_model_only,
             chat_template: self.chat_template,
             default_chat_template_kwargs: self.default_chat_template_kwargs,
+            limit_mm_per_prompt: self.limit_mm_per_prompt,
             chat_template_content_format: self.chat_template_content_format,
             max_logprobs: self.max_logprobs,
             api_server_options,
@@ -451,6 +476,7 @@ impl SharedRuntimeArgs {
             language_model_only: self.language_model_only,
             chat_template: self.chat_template,
             default_chat_template_kwargs: self.default_chat_template_kwargs,
+            limit_mm_per_prompt: self.limit_mm_per_prompt,
             chat_template_content_format: self.chat_template_content_format,
             max_logprobs: self.max_logprobs,
             api_server_options,
@@ -666,6 +692,7 @@ impl ServeArgs {
             self.runtime.disable_log_stats,
             self.runtime.shutdown_timeout,
             handshake_port,
+            self.runtime.limit_mm_per_prompt_json(),
         )
     }
 }

rust/src/cmd/src/cli/tests.rs

@@ -64,6 +64,7 @@ fn serve_args_forward_python_flags_with_separator() {
                         http_timeout_keep_alive: None,
                         chat_template: None,
                         default_chat_template_kwargs: None,
+                        limit_mm_per_prompt: {},
                         chat_template_content_format: Auto,
                         enable_log_requests: false,
                         enable_prompt_tokens_details: false,
@@ -666,7 +667,7 @@ fn serve_args_reject_unknown_renderer_value() {
     .unwrap_err();
 
     expect![[r#"
-        error: invalid value 'definitely_missing' for '--tokenizer-mode <RENDERER>': unknown renderer `definitely_missing` (expected one of: auto, hf, deepseek_v32, deepseek_v4, harmony, inkling)
+        error: invalid value 'definitely_missing' for '--tokenizer-mode <RENDERER>': unknown renderer `definitely_missing` (expected one of: auto, hf, deepseek_v32, deepseek_v4, harmony, inkling, kimi_k3)
 
         For more information, try '--help'.
     "#]]
@@ -762,6 +763,7 @@ fn frontend_args_accept_json() {
                         http_timeout_keep_alive: None,
                         chat_template: None,
                         default_chat_template_kwargs: None,
+                        limit_mm_per_prompt: {},
                         chat_template_content_format: Auto,
                         enable_log_requests: false,
                         enable_prompt_tokens_details: false,
@@ -1117,6 +1119,24 @@ fn frontend_args_json_rejects_malformed_json() {
     "#]].assert_eq(&error.to_string());
 }
 
+#[test]
+fn serve_args_reject_unsupported_modality_in_limit_mm_per_prompt() {
+    let error = Cli::try_parse_from([
+        "vllm-rs",
+        "serve",
+        "Qwen/Qwen3-0.6B",
+        "--limit-mm-per-prompt",
+        r#"{"unsupported_modality": 1}"#,
+    ])
+    .unwrap_err();
+
+    expect![[r#"
+        error: invalid value '{"unsupported_modality": 1}' for '--limit-mm-per-prompt <JSON>': invalid JSON object: unknown variant `unsupported_modality`, expected one of `image`, `audio`, `video` at line 1 column 23
+
+        For more information, try '--help'.
+    "#]].assert_eq(&error.to_string());
+}
+
 #[test]
 fn serve_args_reject_flags_before_model() {
     let error = Cli::try_parse_from(["vllm-rs", "serve", "--python", "python3", "Qwen/Qwen3-0.6B"])
@@ -1331,6 +1351,7 @@ fn serve_args_accept_handshake_aliases() {
                         http_timeout_keep_alive: None,
                         chat_template: None,
                         default_chat_template_kwargs: None,
+                        limit_mm_per_prompt: {},
                         chat_template_content_format: Auto,
                         enable_log_requests: false,
                         enable_prompt_tokens_details: false,
@@ -1474,6 +1495,7 @@ fn serve_frontend_config_uses_dp_address_as_advertised_host() {
             language_model_only: false,
             chat_template: None,
             default_chat_template_kwargs: None,
+            limit_mm_per_prompt: {},
             chat_template_content_format: Auto,
             max_logprobs: None,
             api_server_options: ApiServerOptions {
@@ -1558,6 +1580,7 @@ fn serve_frontend_config_keeps_tcp_transport_for_non_local_only_topology() {
             language_model_only: false,
             chat_template: None,
             default_chat_template_kwargs: None,
+            limit_mm_per_prompt: {},
             chat_template_content_format: Auto,
             max_logprobs: None,
             api_server_options: ApiServerOptions {
@@ -1660,6 +1683,7 @@ fn frontend_config_uses_external_coordinator_when_coordinator_address_is_present
             language_model_only: false,
             chat_template: None,
             default_chat_template_kwargs: None,
+            limit_mm_per_prompt: {},
             chat_template_content_format: Auto,
             max_logprobs: None,
             api_server_options: ApiServerOptions {

rust/src/cmd/src/cli/unsupported.rs

@@ -299,11 +299,6 @@ pub struct EngineUnsupportedArgs {
     )]
     pub kv_sharing_fast_prefill: Option<Unsupported>,
 
-    /// The maximum number of input items and options allowed per
-    /// prompt for each modality.
-    #[arg(long)]
-    pub limit_mm_per_prompt: Option<Unsupported>,
-
     /// Additional args passed to process media inputs, keyed by modalities.
     #[arg(long)]
     pub media_io_kwargs: Option<Unsupported>,

rust/src/cmd/src/main.rs

@@ -2,7 +2,6 @@
 // SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 mod cli;
-mod logging;
 
 use std::env;
 use std::ffi::OsStr;
@@ -89,7 +88,7 @@ fn main() -> Result<()> {
             "serve" | "frontend" => "RustFrontend",
             _ => "Rust",
         };
-    logging::init_tracing(process_label);
+    vllm_tracing::init_tracing(process_label);
 
     let cli = Cli::parse();
 

rust/src/engine-core-client/src/client.rs

@@ -519,7 +519,7 @@ impl EngineCoreClient {
                 "registered request to engine"
             );
 
-            self.inner.send_to_engine(&engine_id, EngineCoreRequestType::Add, &req).await?;
+            self.inner.send_request_to_engine(&engine_id, req).await?;
             Ok(())
         }
         .await;

rust/src/engine-core-client/src/client/imp.rs

@@ -6,6 +6,7 @@ use std::sync::Arc;
 use std::time::{SystemTime, UNIX_EPOCH};
 
 use arc_swap::ArcSwapOption;
+use bytes::Bytes;
 use parking_lot::Mutex;
 use thiserror_ext::AsReport as _;
 use tokio::runtime::Handle;
@@ -21,17 +22,30 @@ use crate::error::{client_closed, dispatcher_closed, unexpected_dispatcher_outpu
 use crate::metrics::{LoraInfoExporter, SchedulerStatsRecorder};
 use crate::protocol::encode_msgpack;
 use crate::protocol::output::{EngineCoreOutput, EngineCoreOutputs};
-use crate::protocol::request::EngineCoreRequestType;
+use crate::protocol::request::{EngineCoreRequest, EngineCoreRequestType};
 use crate::protocol::stats::SchedulerStats;
 use crate::protocol::utility::UtilityOutput;
 use crate::transport::{ConnectedEngine, EngineId};
 use crate::{Error, Result, transport};
 
+const MSGPACK_ZERO_COPY_THRESHOLD_ENV: &str = "VLLM_MSGPACK_ZERO_COPY_THRESHOLD";
+const DEFAULT_MSGPACK_ZERO_COPY_THRESHOLD: usize = 256;
+
+fn msgpack_zero_copy_threshold() -> usize {
+    std::env::var(MSGPACK_ZERO_COPY_THRESHOLD_ENV)
+        .ok()
+        .and_then(|value| value.parse().ok())
+        .unwrap_or(DEFAULT_MSGPACK_ZERO_COPY_THRESHOLD)
+}
+
 pub(crate) struct ClientInner {
     input_send: RouterSendHalf,
     /// The runtime handle used for sending messages to the engine.
     handle: Handle,
     model_name: String,
+    /// Per-tensor byte threshold loaded from env variable
+    /// `VLLM_MSGPACK_ZERO_COPY_THRESHOLD` when this inner client is created.
+    msgpack_zero_copy_threshold: usize,
     scheduler_stats_recorder: SchedulerStatsRecorder,
     request_reg: Mutex<RequestRegistry>,
     utility_reg: Mutex<UtilityRegistry>,
@@ -54,6 +68,7 @@ impl ClientInner {
             input_send,
             handle,
             model_name,
+            msgpack_zero_copy_threshold: msgpack_zero_copy_threshold(),
             scheduler_stats_recorder,
             request_reg: Mutex::new(RequestRegistry::new(engines)),
             utility_reg: Mutex::new(UtilityRegistry::default()),
@@ -229,9 +244,29 @@ impl ClientInner {
     where
         T: serde::Serialize + std::fmt::Debug,
     {
-        // TODO: for `EngineCoreRequest`, split outbound tensor raw views into aux
-        // frames instead of always producing a single msgpack frame.
-        let payload = encode_msgpack(payload)?;
+        let payload = Bytes::from(encode_msgpack(payload)?);
+        self.send_encoded_to_engine(engine_id, request_type, payload, Vec::new()).await
+    }
+
+    /// Send an add request, moving large tensor buffers into auxiliary frames.
+    pub async fn send_request_to_engine(
+        &self,
+        engine_id: &EngineId,
+        mut payload: EngineCoreRequest,
+    ) -> Result<()> {
+        let aux_frames = payload.extract_aux_frames(self.msgpack_zero_copy_threshold);
+        let payload = Bytes::from(encode_msgpack(&payload)?);
+        self.send_encoded_to_engine(engine_id, EngineCoreRequestType::Add, payload, aux_frames)
+            .await
+    }
+
+    async fn send_encoded_to_engine(
+        &self,
+        engine_id: &EngineId,
+        request_type: EngineCoreRequestType,
+        payload: Bytes,
+        aux_frames: Vec<Bytes>,
+    ) -> Result<()> {
         let mut input_send = self.input_send.clone();
         let engine_id = engine_id.clone();
 
@@ -242,6 +277,7 @@ impl ClientInner {
                     &engine_id,
                     request_type.to_frame(),
                     payload,
+                    aux_frames,
                 )
                 .await
             })

rust/src/engine-core-client/src/protocol/logprobs.rs

@@ -209,17 +209,17 @@ impl WireLogprobs {
             logprob_token_ids: WireNdArray {
                 dtype: "<i8".to_string(),
                 shape: vec![rows, cols],
-                data: WireArrayData::RawView(token_ids),
+                data: WireArrayData::RawView(token_ids.into()),
             },
             logprobs: WireNdArray {
                 dtype: "<f4".to_string(),
                 shape: vec![rows, cols],
-                data: WireArrayData::RawView(logprobs),
+                data: WireArrayData::RawView(logprobs.into()),
             },
             token_ranks: WireNdArray {
                 dtype: "<i8".to_string(),
                 shape: vec![rows],
-                data: WireArrayData::RawView(token_ranks),
+                data: WireArrayData::RawView(token_ranks.into()),
             },
             cu_num_generated_tokens: None,
         })

rust/src/engine-core-client/src/protocol/logprobs/array.rs

@@ -4,6 +4,7 @@
 use std::io::Cursor;
 
 use byteorder::{BigEndian, LittleEndian, NativeEndian, ReadBytesExt};
+use bytes::Bytes;
 use itertools::Itertools as _;
 
 use crate::error::{Error, Result, ext_value_decode};
@@ -126,7 +127,7 @@ pub(super) fn decode_array_metadata<Frame>(
     field: &str,
     frames: &[Frame],
     expected_scalars: &[ScalarType],
-) -> Result<(Vec<usize>, Vec<u8>, ScalarType, Endianness)>
+) -> Result<(Vec<usize>, Bytes, ScalarType, Endianness)>
 where
     Frame: AsRef<[u8]>,
 {
@@ -171,7 +172,7 @@ pub(super) fn resolve_array_bytes<Frame>(
     value: WireArrayData,
     field: &str,
     frames: &[Frame],
-) -> Result<Vec<u8>>
+) -> Result<Bytes>
 where
     Frame: AsRef<[u8]>,
 {
@@ -187,7 +188,7 @@ where
                     ),
                 )
             })?;
-            Ok(frame.as_ref().to_vec())
+            Ok(Bytes::copy_from_slice(frame.as_ref()))
         }
     }
 }

rust/src/engine-core-client/src/protocol/multimodal.rs

@@ -3,6 +3,7 @@
 
 use std::collections::BTreeMap;
 
+use bytes::Bytes;
 use serde::{Deserialize, Serialize};
 use serde_tuple::{Deserialize_tuple, Serialize_tuple};
 
@@ -105,6 +106,44 @@ pub enum MmKwargValue {
     List(Vec<MmKwargValue>),
 }
 
+impl MmFeatureSpec {
+    /// Extract large tensor buffers from this feature in serialized field order.
+    pub(crate) fn extract_aux_frames(&mut self, aux_frames: &mut Vec<Bytes>, threshold: usize) {
+        if let Some(data) = &mut self.data {
+            for elem in data.values_mut() {
+                elem.extract_aux_frames(aux_frames, threshold);
+            }
+        }
+        if let Some(is_embed) = &mut self.mm_position.is_embed {
+            is_embed.extract_aux_frame(aux_frames, threshold);
+        }
+    }
+}
+
+impl MmFieldElem {
+    /// Extract large tensor buffers from this field element.
+    fn extract_aux_frames(&mut self, aux_frames: &mut Vec<Bytes>, threshold: usize) {
+        if let Some(data) = &mut self.data {
+            data.extract_aux_frames(aux_frames, threshold);
+        }
+    }
+}
+
+impl MmKwargValue {
+    /// Recursively extract large tensor buffers from this nested value.
+    fn extract_aux_frames(&mut self, aux_frames: &mut Vec<Bytes>, threshold: usize) {
+        match self {
+            Self::Tensor(tensor) => tensor.extract_aux_frame(aux_frames, threshold),
+            Self::List(values) => {
+                for value in values {
+                    value.extract_aux_frames(aux_frames, threshold);
+                }
+            }
+            Self::Int(_) | Self::Float(_) => {}
+        }
+    }
+}
+
 /// Defines how to interpret tensor data belonging to a keyword argument for
 /// `MultiModalKwargsItems`, and vice versa.
 ///

rust/src/engine-core-client/src/protocol/request.rs

@@ -137,6 +137,17 @@ impl EngineCoreRequest {
         }
         Ok(())
     }
+
+    /// Extract large request tensors into ordered auxiliary frames.
+    pub(crate) fn extract_aux_frames(&mut self, threshold: usize) -> Vec<Bytes> {
+        let mut aux_frames = Vec::new();
+        if let Some(features) = &mut self.mm_features {
+            for feature in features {
+                feature.extract_aux_frames(&mut aux_frames, threshold);
+            }
+        }
+        aux_frames
+    }
 }
 
 #[cfg(test)]
@@ -144,9 +155,15 @@ mod tests {
     use rmpv::Value;
 
     use super::*;
+    use crate::protocol::multimodal::{
+        MmBatchedField, MmFeatureSpec, MmField, MmFieldElem, MmKwargValue, PlaceholderRange,
+    };
     use crate::protocol::sampling::EngineCoreSamplingParams;
+    use crate::protocol::tensor::{WireArrayData, WireTensor};
     use crate::protocol::{decode_value, encode_msgpack};
 
+    const AUX_FRAME_THRESHOLD: usize = 256;
+
     #[test]
     fn engine_core_request_serializes_as_full_array() {
         let request = EngineCoreRequest {
@@ -175,4 +192,84 @@ mod tests {
         assert_eq!(array[10], Value::Nil);
         assert_eq!(array[11], Value::from(7));
     }
+
+    #[test]
+    fn engine_core_request_extracts_large_nested_tensors_in_wire_order() {
+        let inline = vec![1_u8; AUX_FRAME_THRESHOLD - 1];
+        let first_aux = vec![2_u8; AUX_FRAME_THRESHOLD];
+        let second_aux = vec![3_u8; AUX_FRAME_THRESHOLD + 1];
+        let first_aux_ptr = first_aux.as_ptr();
+        let second_aux_ptr = second_aux.as_ptr();
+        let mut request = EngineCoreRequest {
+            mm_features: Some(vec![MmFeatureSpec {
+                data: Some(BTreeMap::from([
+                    (
+                        "inline".to_string(),
+                        MmFieldElem {
+                            data: Some(MmKwargValue::Tensor(WireTensor::from_raw(
+                                "uint8",
+                                vec![inline.len()],
+                                inline,
+                            ))),
+                            field: MmField::Batched(MmBatchedField { keep_on_cpu: false }),
+                        },
+                    ),
+                    (
+                        "nested".to_string(),
+                        MmFieldElem {
+                            data: Some(MmKwargValue::List(vec![
+                                MmKwargValue::Int(7),
+                                MmKwargValue::Tensor(WireTensor::from_raw(
+                                    "uint8",
+                                    vec![first_aux.len()],
+                                    first_aux,
+                                )),
+                            ])),
+                            field: MmField::Batched(MmBatchedField { keep_on_cpu: false }),
+                        },
+                    ),
+                ])),
+                modality: "image".to_string(),
+                identifier: "id".to_string(),
+                mm_position: PlaceholderRange {
+                    offset: 0,
+                    length: second_aux.len(),
+                    is_embed: Some(WireTensor::from_raw(
+                        "bool",
+                        vec![second_aux.len()],
+                        second_aux,
+                    )),
+                },
+                mm_hash: None,
+            }]),
+            ..EngineCoreRequest::default()
+        };
+
+        let aux_frames = request.extract_aux_frames(AUX_FRAME_THRESHOLD);
+
+        assert_eq!(aux_frames.len(), 2);
+        assert_eq!(aux_frames[0].as_ptr(), first_aux_ptr);
+        assert_eq!(aux_frames[1].as_ptr(), second_aux_ptr);
+        let feature = &request.mm_features.as_ref().unwrap()[0];
+        let MmKwargValue::Tensor(inline) =
+            feature.data.as_ref().unwrap()["inline"].data.as_ref().unwrap()
+        else {
+            panic!("expected inline tensor");
+        };
+        assert!(matches!(inline.data, WireArrayData::RawView(_)));
+        let MmKwargValue::List(nested) =
+            feature.data.as_ref().unwrap()["nested"].data.as_ref().unwrap()
+        else {
+            panic!("expected nested tensor list");
+        };
+        let MmKwargValue::Tensor(nested_tensor) = &nested[1] else {
+            panic!("expected nested tensor");
+        };
+        assert_eq!(nested_tensor.data, WireArrayData::AuxIndex(1));
+        assert_eq!(
+            feature.mm_position.is_embed.as_ref().unwrap().data,
+            WireArrayData::AuxIndex(2)
+        );
+        assert!(request.extract_aux_frames(AUX_FRAME_THRESHOLD).is_empty());
+    }
 }

rust/src/engine-core-client/src/protocol/tensor.rs

@@ -1,7 +1,8 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
-use bytemuck::allocation::pod_collect_to_vec;
+use bytemuck::{Pod, cast_slice};
+use bytes::Bytes;
 use enum_as_inner::EnumAsInner;
 use half::{bf16, f16};
 use rmpv::Value;
@@ -20,6 +21,21 @@ struct MsgpackExtRef<'a>((i8, ByteSlice<'a>));
 
 struct ByteSlice<'a>(&'a [u8]);
 
+struct PodVec<T: Pod>(Vec<T>);
+
+impl<T: Pod> AsRef<[u8]> for PodVec<T> {
+    fn as_ref(&self) -> &[u8] {
+        cast_slice(&self.0)
+    }
+}
+
+fn bytes_from_pod_vec<T>(data: Vec<T>) -> Bytes
+where
+    T: Pod + Send + 'static,
+{
+    Bytes::from_owner(PodVec(data))
+}
+
 impl Serialize for ByteSlice<'_> {
     fn serialize<S>(&self, serializer: S) -> std::result::Result<S::Ok, S::Error>
     where
@@ -55,58 +71,63 @@ pub struct WireNdArray {
 
 impl WireNdArray {
     /// Build a float32 tensor/ndarray backed by native-endian raw-view bytes.
-    pub fn from_f32(shape: Vec<usize>, data: impl AsRef<[f32]>) -> Result<Self, String> {
-        let data = data.as_ref();
+    ///
+    /// Takes ownership of the backing buffer without copying its data.
+    pub fn from_f32(shape: Vec<usize>, data: Vec<f32>) -> Result<Self, String> {
         validate_element_count(&shape, data.len())?;
-        Ok(Self {
-            dtype: "float32".to_string(),
+        Ok(Self::from_raw_bytes(
+            "float32",
             shape,
-            data: WireArrayData::RawView(pod_collect_to_vec::<f32, u8>(data)),
-        })
+            bytes_from_pod_vec(data),
+        ))
     }
 
     /// Build a float16 tensor/ndarray backed by native-endian raw-view bytes.
-    pub fn from_f16(shape: Vec<usize>, data: impl AsRef<[f16]>) -> Result<Self, String> {
-        let data = data.as_ref();
+    ///
+    /// Takes ownership of the backing buffer without copying its data.
+    pub fn from_f16(shape: Vec<usize>, data: Vec<f16>) -> Result<Self, String> {
         validate_element_count(&shape, data.len())?;
-        Ok(Self {
-            dtype: "float16".to_string(),
+        Ok(Self::from_raw_bytes(
+            "float16",
             shape,
-            data: WireArrayData::RawView(pod_collect_to_vec::<f16, u8>(data)),
-        })
+            bytes_from_pod_vec(data),
+        ))
     }
 
     /// Build a bfloat16 tensor/ndarray backed by native-endian raw-view bytes.
-    pub fn from_bf16(shape: Vec<usize>, data: impl AsRef<[bf16]>) -> Result<Self, String> {
-        let data = data.as_ref();
+    ///
+    /// Takes ownership of the backing buffer without copying its data.
+    pub fn from_bf16(shape: Vec<usize>, data: Vec<bf16>) -> Result<Self, String> {
         validate_element_count(&shape, data.len())?;
-        Ok(Self {
-            dtype: "bfloat16".to_string(),
+        Ok(Self::from_raw_bytes(
+            "bfloat16",
             shape,
-            data: WireArrayData::RawView(pod_collect_to_vec::<bf16, u8>(data)),
-        })
+            bytes_from_pod_vec(data),
+        ))
     }
 
     /// Build an int64 tensor/ndarray backed by native-endian raw-view bytes.
-    pub fn from_i64(shape: Vec<usize>, data: impl AsRef<[i64]>) -> Result<Self, String> {
-        let data = data.as_ref();
+    ///
+    /// Takes ownership of the backing buffer without copying its data.
+    pub fn from_i64(shape: Vec<usize>, data: Vec<i64>) -> Result<Self, String> {
         validate_element_count(&shape, data.len())?;
-        Ok(Self {
-            dtype: "int64".to_string(),
+        Ok(Self::from_raw_bytes(
+            "int64",
             shape,
-            data: WireArrayData::RawView(pod_collect_to_vec::<i64, u8>(data)),
-        })
+            bytes_from_pod_vec(data),
+        ))
     }
 
     /// Build a uint32 tensor/ndarray backed by native-endian raw-view bytes.
-    pub fn from_u32(shape: Vec<usize>, data: impl AsRef<[u32]>) -> Result<Self, String> {
-        let data = data.as_ref();
+    ///
+    /// Takes ownership of the backing buffer without copying its data.
+    pub fn from_u32(shape: Vec<usize>, data: Vec<u32>) -> Result<Self, String> {
         validate_element_count(&shape, data.len())?;
-        Ok(Self {
-            dtype: "uint32".to_string(),
+        Ok(Self::from_raw_bytes(
+            "uint32",
             shape,
-            data: WireArrayData::RawView(pod_collect_to_vec::<u32, u8>(data)),
-        })
+            bytes_from_pod_vec(data),
+        ))
     }
 
     /// Build a bool tensor/ndarray backed by raw-view bytes.
@@ -118,7 +139,9 @@ impl WireNdArray {
         Ok(Self {
             dtype: "bool".to_string(),
             shape,
-            data: WireArrayData::RawView(data.iter().map(|value| u8::from(*value)).collect()),
+            data: WireArrayData::RawView(Bytes::from(
+                data.into_iter().map(u8::from).collect::<Vec<_>>(),
+            )),
         })
     }
 
@@ -127,12 +150,22 @@ impl WireNdArray {
     /// Use this as an escape hatch when the caller already owns bytes that
     /// match the requested `dtype` and `shape`.
     pub fn from_raw(dtype: impl Into<String>, shape: Vec<usize>, data: Vec<u8>) -> Self {
+        Self::from_raw_bytes(dtype, shape, Bytes::from(data))
+    }
+
+    /// Build a tensor/ndarray from an owned immutable raw-view buffer.
+    pub fn from_raw_bytes(dtype: impl Into<String>, shape: Vec<usize>, data: Bytes) -> Self {
         Self {
             dtype: dtype.into(),
             shape,
             data: WireArrayData::RawView(data),
         }
     }
+
+    /// Move a sufficiently large inline buffer into the ordered auxiliary-frame list.
+    pub(crate) fn extract_aux_frame(&mut self, aux_frames: &mut Vec<Bytes>, threshold: usize) {
+        self.data.extract_aux_frame(aux_frames, threshold);
+    }
 }
 
 /// Validate that the number of elements implied by the shape matches the length
@@ -170,7 +203,25 @@ pub enum WireArrayData {
     /// stored.
     AuxIndex(usize),
     /// The raw bytes of this array/tensor.
-    RawView(Vec<u8>),
+    RawView(Bytes),
+}
+
+impl WireArrayData {
+    /// Replace a sufficiently large raw view with its one-based auxiliary-frame index.
+    fn extract_aux_frame(&mut self, aux_frames: &mut Vec<Bytes>, threshold: usize) {
+        let Self::RawView(bytes) = self else {
+            return;
+        };
+        if bytes.len() < threshold {
+            return;
+        }
+
+        let index = aux_frames.len() + 1;
+        let bytes = std::mem::replace(self, Self::AuxIndex(index))
+            .into_raw_view()
+            .expect("raw view was matched above");
+        aux_frames.push(bytes);
+    }
 }
 
 impl<'de> Deserialize<'de> for WireArrayData {
@@ -180,7 +231,9 @@ impl<'de> Deserialize<'de> for WireArrayData {
     {
         let value = Value::deserialize(deserializer)?;
         match value {
-            Value::Ext(tag, bytes) if tag == CUSTOM_TYPE_RAW_VIEW => Ok(Self::RawView(bytes)),
+            Value::Ext(tag, bytes) if tag == CUSTOM_TYPE_RAW_VIEW => {
+                Ok(Self::RawView(Bytes::from(bytes)))
+            }
             Value::Ext(tag, _) => Err(serde::de::Error::custom(format!(
                 "unsupported extension type code {tag}"
             ))),
@@ -201,9 +254,6 @@ impl Serialize for WireArrayData {
     where
         S: Serializer,
     {
-        // TODO: outbound request serialization currently only supports inline
-        // raw-view bytes. Emitting aux frames needs transport-level plumbing;
-        // serializing `AuxIndex` here only preserves an already-built reference.
         match self {
             Self::AuxIndex(index) => serializer.serialize_u64(*index as u64),
             Self::RawView(bytes) => {
@@ -221,7 +271,7 @@ mod tests {
     fn raw_view_serializes_as_msgpack_ext() {
         let bytes = vec![1, 2, 3, 4];
         let encoded =
-            rmp_serde::to_vec_named(&WireArrayData::RawView(bytes.clone())).expect("encode");
+            rmp_serde::to_vec_named(&WireArrayData::RawView(bytes.clone().into())).expect("encode");
         let expected = rmp_serde::to_vec_named(&Value::Ext(CUSTOM_TYPE_RAW_VIEW, bytes.clone()))
             .expect("encode expected");
 
@@ -234,11 +284,15 @@ mod tests {
 
     #[test]
     fn constructors_build_raw_view_tensors() {
-        let f32_tensor = WireNdArray::from_f32(vec![2], vec![1.0, 2.5]).unwrap();
+        let f32_data = vec![1.0, 2.5];
+        let f32_data_ptr = f32_data.as_ptr().cast::<u8>();
+        let f32_tensor = WireNdArray::from_f32(vec![2], f32_data).unwrap();
         assert_eq!(f32_tensor.dtype, "float32");
         assert_eq!(f32_tensor.shape, vec![2]);
+        let f32_raw_view = f32_tensor.data.into_raw_view().expect("raw view");
+        assert_eq!(f32_raw_view.as_ptr(), f32_data_ptr);
         assert_eq!(
-            f32_tensor.data.into_raw_view().expect("raw view"),
+            f32_raw_view,
             [1.0_f32, 2.5].into_iter().flat_map(f32::to_ne_bytes).collect::<Vec<_>>()
         );
 
@@ -258,15 +312,15 @@ mod tests {
         let i64_tensor = WireNdArray::from_i64(vec![1], vec![-7]).unwrap();
         assert_eq!(i64_tensor.dtype, "int64");
         assert_eq!(
-            i64_tensor.data.into_raw_view().expect("raw view"),
-            (-7_i64).to_ne_bytes()
+            i64_tensor.data.into_raw_view().expect("raw view").as_ref(),
+            (-7_i64).to_ne_bytes().as_ref()
         );
 
         let u32_tensor = WireNdArray::from_u32(vec![1], vec![42]).unwrap();
         assert_eq!(u32_tensor.dtype, "uint32");
         assert_eq!(
-            u32_tensor.data.into_raw_view().expect("raw view"),
-            42_u32.to_ne_bytes()
+            u32_tensor.data.into_raw_view().expect("raw view").as_ref(),
+            42_u32.to_ne_bytes().as_ref()
         );
 
         let bool_tensor = WireNdArray::from_bool(vec![2], vec![false, true]).unwrap();