kimi-k3

Co-authored-by: Isotr0py <Isotr0py@outlook.com> Co-authored-by: Thien Tran <gau.nernst@yahoo.com.sg> Co-authored-by: Bugen Zhao <i@bugenzhao.com> Co-authored-by: Yongye Zhu <zyy1102000@gmail.com> Co-authored-by: Ziming Huang <zelda.huanghuang@gmail.com> Co-authored-by: Roger Wang <hey@rogerw.io> Co-authored-by: Isotr0py <mozf@inferact.ai> Co-authored-by: aoshen02 <aoshen02@users.noreply.github.com> Co-authored-by: Woosuk Kwon <woosuk@inferact.ai> Co-authored-by: Jee Jee Li <jeejeelee@inferact.ai> Co-authored-by: aoshen02 <aoshen@inferact.ai> Co-authored-by: Jiangyun Zhu <riverclouds.zhu@qq.com> Co-authored-by: Summer Yang <girasoleyang@gmail.com> Co-authored-by: Kevin H. Luu <khluu000@gmail.com> Co-authored-by: khluu <khluu000@gmail.com> Co-authored-by: Bowen Wang <abmfy@icloud.com> Co-authored-by: gnovack <novackgm@gmail.com> Co-authored-by: Nick Hill <nickhill123@gmail.com> Co-authored-by: Yifan Qiao <yifanqiao@inferact.ai> Co-authored-by: xiaozhoupy <peiyuanzhou1994@gmail.com> Co-authored-by: Roy Wang <yasong.wang@inferact.ai> Co-authored-by: Jeff (Junze) Ma <93145857+majunze2001@users.noreply.github.com> Signed-off-by: zjy0516 <riverclouds.zhu@qq.com>
Stabilize GPU memory teardown between ROCm CI tests (#49242 )
2026-07-27 07:01:52 +00:00 · 2026-07-25 04:21:36 +00:00 · 2026-07-24 23:15:19 -05:00 · 2026-07-24 23:14:37 -05:00 · 2026-07-24 21:02:22 -07:00 · 2026-07-25 02:26:15 +00:00
392 changed files with 45312 additions and 4365 deletions
@@ -35,7 +35,7 @@ set -o pipefail
 : "${PY_COLORS:=1}"
 : "${ROCM_DOCKER_TTY:=1}"
 : "${PYTHONFAULTHANDLER:=1}"
-: "${PYTEST_TIMEOUT:=2100}"
+: "${PYTEST_TIMEOUT:=2400}"
 if [[ " ${PYTEST_ADDOPTS:-} " != *" --color"* ]]; then
  PYTEST_ADDOPTS="${PYTEST_ADDOPTS:+${PYTEST_ADDOPTS} }--color=yes"
 fi
@@ -45,9 +45,9 @@ fi
 if [[ " ${PYTEST_ADDOPTS:-} " != *" --durations-min="* ]]; then
  PYTEST_ADDOPTS="${PYTEST_ADDOPTS:+${PYTEST_ADDOPTS} }--durations-min=1.0"
 fi
-# Dump stacks after 15 minutes, then stop an individual test after 35 minutes.
+# Dump stacks after 25 minutes, then stop an individual test after 40 minutes.
 if [[ " ${PYTEST_ADDOPTS:-} " != *" faulthandler_timeout="* ]]; then
-  PYTEST_ADDOPTS="${PYTEST_ADDOPTS:+${PYTEST_ADDOPTS} }-o faulthandler_timeout=900"
+  PYTEST_ADDOPTS="${PYTEST_ADDOPTS:+${PYTEST_ADDOPTS} }-o faulthandler_timeout=1500"
 fi
 if [[ " ${PYTEST_ADDOPTS:-} " != *" --timeout-method="* &&
  " ${PYTEST_ADDOPTS:-} " != *" --timeout-method "* ]]; then
@@ -400,10 +400,10 @@ initialize_native_environment() {
  native_root="/tmp/vllm-native-${job_id}"
  TMPDIR="/tmp/vllm-${job_id_suffix}/tmp"
  VLLM_RPC_BASE_PATH="/tmp"
-  : "${TORCHINDUCTOR_CACHE_DIR:=${native_root}/cache/torchinductor}"
-  : "${TRITON_CACHE_DIR:=${native_root}/cache/triton}"
-  : "${VLLM_CACHE_ROOT:=${native_root}/cache/vllm}"
-  : "${XDG_CACHE_HOME:=${native_root}/cache/xdg}"
+  TORCHINDUCTOR_CACHE_DIR="${native_root}/cache/torchinductor"
+  TRITON_CACHE_DIR="${native_root}/cache/triton"
+  VLLM_CACHE_ROOT="${native_root}/cache/vllm"
+  XDG_CACHE_HOME="${native_root}/cache/xdg"
  : "${HF_HOME:=/home/buildkite-agent/huggingface}"
  : "${HF_HUB_DOWNLOAD_TIMEOUT:=300}"
  : "${HF_HUB_ETAG_TIMEOUT:=60}"
@@ -419,6 +419,8 @@ initialize_native_environment() {
    "${XDG_CACHE_HOME}" \
    "${HF_HOME}" || return 1

+  echo "Native compile caches: VLLM_CACHE_ROOT=${VLLM_CACHE_ROOT} TORCHINDUCTOR_CACHE_DIR=${TORCHINDUCTOR_CACHE_DIR}"
+
  if [[ "${VLLM_CI_REQUIRE_PERSISTENT_HF_CACHE:-0}" == "1" ]]; then
    if ! command -v findmnt >/dev/null 2>&1; then
      echo "findmnt is required to verify the native Hugging Face cache mount" >&2
@@ -40,7 +40,7 @@
 #####################################################################################################################################
 #                                                                                                                                   #
 # IMPORTANT:                                                                                                                        #
-#   * Currently AMD CI has MI250 agents, MI300 agents, MI325 agents, and MI355 agents. All upcoming feature improvements are        #
+#   * Currently AMD CI has MI250 agents, MI300 agents, and MI355 agents. All upcoming feature improvements are                      #
 #         tracked in: https://github.com/vllm-project/vllm/issues/34994                                                             #
 #                                                                                                                                   #
 #-----------------------------------------------------------------------------------------------------------------------------------#
@@ -81,10 +81,8 @@
 #                                                     the above test.) Also run if model initialization test file is modified.      #
 #   * [Language Models Tests (Extra Standard) %N]: Shard slow subset of standard language models tests. Only run when model         #
 #                                                  source is modified, or when specified test files are modified.                   #
-#   * [Language Models Tests (Hybrid) %N]: Install fast path packages for testing against transformers (mamba, conv1d) and to       #
-#                                          run plamo2 model in vLLM.                                                                #
-#   * [Language Models Test (Extended Generation)]: Install fast path packages for testing against transformers (mamba, conv1d)     #
-#                                                   and to run plamo2 model in vLLM.                                                #
+#   * [Language Models Tests (Hybrid) %N]: Install fast path packages for testing against transformers (mamba, conv1d).             #
+#   * [Language Models Test (Extended Generation)]: Install fast path packages for testing against transformers (mamba, conv1d).    #
 #   * [Multi-Modal Models (Standard) 1-4]:                                                                                          #
 #     - Do NOT remove `VLLM_WORKER_MULTIPROC_METHOD=spawn` setting as ROCm requires this for certain models to function.            #
 #   * [Transformers Nightly Models]: Whisper needs `VLLM_WORKER_MULTIPROC_METHOD=spawn` to avoid deadlock.                          #
@@ -207,6 +205,7 @@ steps:
  timeout_in_minutes: 180
  mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250]
  agent_pool: mi250_1
+  optional: true
  working_dir: "/vllm-workspace/tests"
  source_file_dependencies:
  - vllm/
@@ -240,6 +239,20 @@ steps:
  - pytest -v -s models/multimodal/generation/test_common.py -m core_model -k "not qwen2 and not qwen3 and not gemma"
  - pytest -v -s models/multimodal/generation/test_qwen2_vl.py -m core_model

+- label: Multi-Modal Processor (CPU) %N # TBD
+  timeout_in_minutes: 180
+  mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250]
+  agent_pool: mi250_1
+  parallelism: 6
+  optional: true
+  working_dir: "/vllm-workspace/tests"
+  source_file_dependencies:
+  - vllm/
+  - tests/models/multimodal
+  - tests/models/registry.py
+  commands:
+  - pytest -v -s models/multimodal/processing --ignore models/multimodal/processing/test_tensor_schema.py --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB
+
 #------------------------------------------------------------  mi250 · v1  -------------------------------------------------------------#

 - label: Batch Invariance (H100-MI250) # TBD
@@ -367,6 +380,22 @@ steps:
  commands:
  - pytest -v -s v1/attention

+- label: V1 others (CPU) # TBD
+  timeout_in_minutes: 180
+  mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250]
+  agent_pool: mi250_1
+  optional: true
+  working_dir: "/vllm-workspace/tests"
+  source_file_dependencies:
+  - vllm/
+  - tests/v1
+  commands:
+  - pytest -v -s -m 'cpu_test' v1/core
+  - pytest -v -s v1/structured_output
+  - pytest -v -s v1/test_serial_utils.py
+  - pytest -v -s -m 'cpu_test' v1/kv_connector/unit
+  - pytest -v -s -m 'cpu_test' v1/metrics
+
 #-------------------------------------------------------------  mi250 · misc  ------------------------------------------------------------#

 - label: Async Engine, Inputs, Utils, Worker, Config (CPU) # TBD
@@ -408,6 +437,19 @@ steps:
  - pytest -v -s transformers_utils
  - pytest -v -s config

+- label: Python-only Installation # TBD
+  timeout_in_minutes: 180
+  mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250]
+  agent_pool: mi250_1
+  optional: true
+  working_dir: "/vllm-workspace/tests"
+  source_file_dependencies:
+  - tests/standalone_tests/python_only_compile.sh
+  - setup.py
+  - vllm/platforms/rocm.py
+  commands:
+  - bash standalone_tests/python_only_compile.sh
+
 #------------------------------------------------------------  mi250 · rust  -----------------------------------------------------------#

 - label: Rust Frontend Cargo Style + Clippy # TBD
@@ -445,6 +487,7 @@ steps:
  mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250]
  agent_pool: mi250_1
  no_gpu: true
+  optional: true
  working_dir: "/vllm-workspace/tests"
  source_file_dependencies:
  - .buildkite/scripts/docker-build-metadata-args.sh
@@ -507,7 +550,7 @@ steps:
  - tests/models/
  commands:
  - TARGET_TEST_SUITE=MI300 pytest basic_correctness/ -v -s -m 'distributed(num_gpus=2)'
-  - CUDA_VISIBLE_DEVICES=0,1 pytest -v -s model_executor/model_loader/test_sharded_state_loader.py -m '(not slow_test)'
+  - HIP_VISIBLE_DEVICES=0,1 pytest -v -s model_executor/model_loader/test_sharded_state_loader.py -m '(not slow_test)'
  - pytest models/transformers/test_backend.py -v -s -m 'distributed(num_gpus=2)'
  - pytest models/language -v -s -m 'distributed(num_gpus=2)'
  - pytest models/multimodal -v -s -m 'distributed(num_gpus=2)' --ignore models/multimodal/generation/test_whisper.py --ignore models/multimodal/generation/test_phi4siglip.py
@@ -657,6 +700,30 @@ steps:
  - VLLM_TEST_CLEAN_GPU_MEMORY=1 pytest -v -s tests/compile/passes/distributed/test_async_tp.py
  - pytest -v -s tests/compile/fusions_e2e/test_tp2_ar_rms.py::test_tp2_ar_rms_fusions

+- label: Distributed Compile + RPC Tests (2 GPUs) # TBD
+  timeout_in_minutes: 180
+  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
+  dind: false
+  agent_pool: mi300_2
+  num_gpus: 2
+  working_dir: "/vllm-workspace/tests"
+  source_file_dependencies:
+  - vllm/compilation/
+  - vllm/distributed/
+  - vllm/engine/
+  - vllm/executor/
+  - vllm/worker/worker_base.py
+  - vllm/v1/engine/
+  - vllm/v1/worker/
+  - tests/compile/fullgraph/test_basic_correctness.py
+  - tests/compile/test_wrapper.py
+  - tests/entrypoints/llm/test_collective_rpc.py
+  - vllm/platforms/rocm.py
+  commands:
+  - pytest -v -s entrypoints/llm/test_collective_rpc.py
+  - pytest -v -s ./compile/fullgraph/test_basic_correctness.py
+  - pytest -v -s ./compile/test_wrapper.py
+
 #-----------------------------------------------------------  mi300 · cuda  ------------------------------------------------------------#

 - label: Platform Tests # TBD
@@ -679,6 +746,7 @@ steps:
  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
  dind: false
  agent_pool: mi300_1
+  optional: true
  working_dir: "/vllm-workspace/tests"
  source_file_dependencies:
  - vllm/
@@ -870,6 +938,71 @@ steps:
  commands:
  - torchrun --nproc-per-node=8 ../examples/features/torchrun/torchrun_dp_example_offline.py --tp-size=2 --pp-size=1 --dp-size=4 --enable-ep

+- label: Distributed Torchrun + Shutdown Tests (2 GPUs) # TBD
+  timeout_in_minutes: 180
+  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
+  dind: false
+  agent_pool: mi300_2
+  num_gpus: 2
+  working_dir: "/vllm-workspace/tests"
+  source_file_dependencies:
+  - vllm/distributed/
+  - vllm/engine/
+  - vllm/executor/
+  - vllm/worker/worker_base.py
+  - vllm/v1/engine/
+  - vllm/v1/worker/
+  - tests/distributed/
+  - tests/v1/shutdown
+  - tests/v1/worker/test_worker_memory_snapshot.py
+  - vllm/platforms/rocm.py
+  commands:
+  - VLLM_TEST_SAME_HOST=1 torchrun --nproc-per-node=4 distributed/test_same_node.py | grep 'Same node test passed'
+  - VLLM_TEST_SAME_HOST=1 VLLM_TEST_WITH_DEFAULT_DEVICE_SET=1 torchrun --nproc-per-node=4 distributed/test_same_node.py | grep 'Same node test passed'
+  - HIP_VISIBLE_DEVICES=0,1 pytest -v -s v1/shutdown
+  - pytest -v -s v1/worker/test_worker_memory_snapshot.py
+
+- label: Distributed Compile + Comm (4 GPUs) # TBD
+  timeout_in_minutes: 180
+  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
+  dind: false
+  agent_pool: mi300_4
+  num_gpus: 4
+  working_dir: "/vllm-workspace/tests"
+  source_file_dependencies:
+  - vllm/distributed/
+  - tests/distributed/test_pynccl
+  - tests/distributed/test_events
+  - tests/compile/fullgraph/test_basic_correctness.py
+  - tests/distributed/test_symm_mem_allreduce.py
+  - tests/distributed/test_multiproc_executor.py
+  - vllm/platforms/rocm.py
+  commands:
+  - pytest -v -s compile/fullgraph/test_basic_correctness.py
+  - pytest -v -s distributed/test_pynccl.py
+  - pytest -v -s distributed/test_events.py
+  - pytest -v -s distributed/test_symm_mem_allreduce.py
+  - pytest -v -s distributed/test_multiproc_executor.py::test_multiproc_executor_multi_node
+
+#----------------------------------------------------------  mi300 · engine  -----------------------------------------------------------#
+
+- label: Engine # TBD
+  timeout_in_minutes: 180
+  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
+  dind: false
+  agent_pool: mi300_1
+  optional: true
+  working_dir: "/vllm-workspace/tests"
+  source_file_dependencies:
+  - vllm/
+  - tests/engine
+  - tests/test_sequence
+  - tests/test_config
+  - tests/test_logger
+  - tests/test_vllm_port
+  commands:
+  - pytest -v -s engine test_sequence.py test_config.py test_logger.py test_vllm_port.py test_jit_monitor.py
+
 #--------------------------------------------------------  mi300 · entrypoints  --------------------------------------------------------#

 - label: Entrypoints Unit Tests # TBD
@@ -878,6 +1011,7 @@ steps:
  dind: false
  agent_pool: mi300_1
  fast_check: true
+  optional: true
  working_dir: "/vllm-workspace/tests"
  source_file_dependencies:
  - vllm/entrypoints
@@ -982,6 +1116,7 @@ steps:
  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
  dind: false
  agent_pool: mi300_1
+  optional: true
  fast_check: true
  working_dir: "/vllm-workspace/tests"
  source_file_dependencies:
@@ -1010,6 +1145,7 @@ steps:
  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
  dind: false
  agent_pool: mi300_1
+  optional: true
  fast_check: true
  working_dir: "/vllm-workspace/tests"
  source_file_dependencies:
@@ -1024,6 +1160,7 @@ steps:
  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
  dind: false
  agent_pool: mi300_1
+  optional: true
  fast_check: true
  working_dir: "/vllm-workspace/tests"
  source_file_dependencies:
@@ -1345,6 +1482,27 @@ steps:
  - export VLLM_WORKER_MULTIPROC_METHOD=spawn
  - pytest -s -v test_lm_eval_correctness.py --config-list-file=configs/models-large-rocm.txt --tp-size=8

+- label: LM Eval Large Models (4xH100-4xMI300) # TBD
+  timeout_in_minutes: 180
+  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
+  dind: false
+  agent_pool: mi300_4
+  num_gpus: 4
+  optional: true
+  working_dir: "/vllm-workspace/.buildkite/lm-eval-harness"
+  source_file_dependencies:
+  - csrc/
+  - vllm/model_executor/layers/quantization
+  - vllm/model_executor/models/
+  - vllm/model_executor/model_loader/
+  - vllm/v1/attention/backends/
+  - vllm/v1/attention/selector.py
+  - vllm/_aiter_ops.py
+  - vllm/platforms/rocm.py
+  commands:
+  - export VLLM_USE_DEEP_GEMM=0
+  - pytest -s -v test_lm_eval_correctness.py --config-list-file=configs/models-large-rocm-fp8.txt --tp-size=4
+
 #---------------------------------------------------------  mi300 · examples  ----------------------------------------------------------#

 - label: Examples # TBD
@@ -1593,7 +1751,7 @@ steps:
  - set -x
  - export VLLM_USE_V2_MODEL_RUNNER=1
  - pytest -v -s v1/engine/test_llm_engine.py -k "not test_engine_metrics"
-  - ENFORCE_EAGER=1 pytest -v -s v1/e2e/general/test_async_scheduling.py -k "not ngram"
+  - pytest -v -s v1/e2e/general/test_async_scheduling.py -k "not ngram"
  - pytest -v -s v1/e2e/general/test_context_length.py
  - pytest -v -s v1/e2e/general/test_min_tokens.py
  - pytest -v -s entrypoints/llm/test_struct_output_generate.py -k "xgrammar and not speculative_config6 and not speculative_config7 and not speculative_config8 and not speculative_config0"
@@ -1797,6 +1955,37 @@ steps:
  - pip freeze | grep -E 'torch'
  - pytest -v -s models/language -m 'core_model and slow_test' --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB

+- label: Language Models Test (Extended Generation) # TBD
+  timeout_in_minutes: 180
+  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
+  dind: false
+  agent_pool: mi300_1
+  optional: true
+  working_dir: "/vllm-workspace/tests"
+  source_file_dependencies:
+  - vllm/
+  - tests/models/language/generation
+  commands:
+  - uv pip install --system --no-build-isolation 'git+https://github.com/AndreasKaratzas/mamba@fix-rocm-7.0-warp-size-constexpr'
+  - uv pip install --system --no-build-isolation 'git+https://github.com/Dao-AILab/causal-conv1d@v1.6.0'
+  - pytest -v -s models/language/generation -m '(not core_model) and (not hybrid_model)'
+
+- label: Language Models Tests (Hybrid) %N # TBD
+  timeout_in_minutes: 180
+  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
+  dind: false
+  agent_pool: mi300_1
+  parallelism: 2
+  optional: true
+  working_dir: "/vllm-workspace/tests"
+  source_file_dependencies:
+  - vllm/
+  - tests/models/language/generation
+  commands:
+  - uv pip install --system --no-build-isolation 'git+https://github.com/AndreasKaratzas/mamba@fix-rocm-7.0-warp-size-constexpr'
+  - uv pip install --system --no-build-isolation 'git+https://github.com/Dao-AILab/causal-conv1d@v1.6.0'
+  - pytest -v -s models/language/generation -m hybrid_model --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB
+
 #----------------------------------------------------  mi300 · models / multimodal  ----------------------------------------------------#

 - label: Multi-Modal Models (Extended Generation 1) # TBD
@@ -1899,20 +2088,32 @@ steps:
  commands:
  - pytest -v -s models/multimodal/processing/test_tensor_schema.py

- label: Multi-Modal Processor (CPU) %N # TBD
+- label: Multi-Modal Models (Extended Pooling) # TBD
+  timeout_in_minutes: 180
+  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
+  dind: false
+  agent_pool: mi300_1
+  optional: true
+  working_dir: "/vllm-workspace/tests"
+  source_file_dependencies:
+  - vllm/
+  - tests/models/multimodal/pooling
+  commands:
+  - pytest -v -s models/multimodal/pooling -m 'not core_model'
+
+- label: "Multi-Modal Models (Standard) 2: qwen3 + gemma" # TBD
  timeout_in_minutes: 180
  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
  dind: false
  agent_pool: mi300_1
-  parallelism: 4
  optional: true
  working_dir: "/vllm-workspace/tests"
  source_file_dependencies:
  - vllm/
  - tests/models/multimodal
-  - tests/models/registry.py
  commands:
-  - pytest -v -s models/multimodal/processing --ignore models/multimodal/processing/test_tensor_schema.py --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB
+  - pytest -v -s models/multimodal/generation/test_common.py -m core_model -k "qwen3 or gemma"
+  - pytest -v -s models/multimodal/generation/test_qwen2_5_vl.py -m core_model

 #-----------------------------------------------------  mi300 · models / quantized  -----------------------------------------------------#

@@ -1933,7 +2134,29 @@ steps:

 #--------------------------------------------------  mi300 · models / transformers  ---------------------------------------------------#

- label: Transformers Nightly Models (Shardable) %N # TBD
+- label: Transformers Nightly Models (Initialization) %N # TBD
+  timeout_in_minutes: 180
+  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
+  dind: false
+  agent_pool: mi300_1
+  parallelism: 6
+  optional: true
+  working_dir: "/vllm-workspace/"
+  source_file_dependencies:
+  - vllm/model_executor/models/
+  - vllm/model_executor/model_loader/
+  - vllm/multimodal/
+  - vllm/model_executor/layers/
+  - vllm/v1/attention/backends/
+  - vllm/v1/attention/selector.py
+  - vllm/_aiter_ops.py
+  - vllm/platforms/rocm.py
+  - tests/models/
+  commands:
+  - pip install --upgrade git+https://github.com/huggingface/transformers
+  - pytest -v -s tests/models/test_initialization.py --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB
+
+- label: Transformers Nightly Models (Processing) %N # TBD
  timeout_in_minutes: 180
  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
  dind: false
@@ -1953,7 +2176,6 @@ steps:
  - tests/models/
  commands:
  - pip install --upgrade git+https://github.com/huggingface/transformers
-  - pytest -v -s tests/models/test_initialization.py --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB
  - pytest -v -s tests/models/multimodal/processing/ --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB

 - label: Transformers Nightly Models (Single) # TBD
@@ -2517,23 +2739,6 @@ steps:
  # - export HSA_NO_SCRATCH_RECLAIM=1
  - pytest -v -s entrypoints/openai/correctness/test_lmeval.py::test_lm_eval_accuracy_v1_engine

- label: V1 others (CPU) # TBD
-  timeout_in_minutes: 180
-  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
-  dind: false
-  agent_pool: mi300_1
-  optional: true
-  working_dir: "/vllm-workspace/tests"
-  source_file_dependencies:
-  - vllm/
-  - tests/v1
-  commands:
-  - pytest -v -s -m 'cpu_test' v1/core
-  - pytest -v -s v1/structured_output
-  - pytest -v -s v1/test_serial_utils.py
-  - pytest -v -s -m 'cpu_test' v1/kv_connector/unit
-  - pytest -v -s -m 'cpu_test' v1/metrics
-
 - label: V1 Sample + Logits # TBD
  timeout_in_minutes: 180
  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
@@ -2830,195 +3035,6 @@ steps:
  commands:
  - bash weight_loading/run_model_weight_loading_test.sh -c weight_loading/models-large-amd.txt

-#########################################################################################################################################
-#                                                                                                                                       #
-#                                                         MI325 (gfx942) tests                                                          #
-#                                                                                                                                       #
-#########################################################################################################################################
-
-#----------------------------------------------------------  mi325 · compile  ----------------------------------------------------------#
-
- label: Distributed Compile + RPC Tests (2 GPUs) # TBD
-  timeout_in_minutes: 180
-  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325]
-  agent_pool: mi325_2
-  num_gpus: 2
-  working_dir: "/vllm-workspace/tests"
-  source_file_dependencies:
-  - vllm/compilation/
-  - vllm/distributed/
-  - vllm/engine/
-  - vllm/executor/
-  - vllm/worker/worker_base.py
-  - vllm/v1/engine/
-  - vllm/v1/worker/
-  - tests/compile/fullgraph/test_basic_correctness.py
-  - tests/compile/test_wrapper.py
-  - tests/entrypoints/llm/test_collective_rpc.py
-  - vllm/platforms/rocm.py
-  commands:
-  - pytest -v -s entrypoints/llm/test_collective_rpc.py
-  - pytest -v -s ./compile/fullgraph/test_basic_correctness.py
-  - pytest -v -s ./compile/test_wrapper.py
-
-#--------------------------------------------------------  mi325 · distributed  --------------------------------------------------------#
-
- label: Distributed Torchrun + Shutdown Tests (2 GPUs) # TBD
-  timeout_in_minutes: 180
-  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325]
-  agent_pool: mi325_2
-  num_gpus: 2
-  working_dir: "/vllm-workspace/tests"
-  source_file_dependencies:
-  - vllm/distributed/
-  - vllm/engine/
-  - vllm/executor/
-  - vllm/worker/worker_base.py
-  - vllm/v1/engine/
-  - vllm/v1/worker/
-  - tests/distributed/
-  - tests/v1/shutdown
-  - tests/v1/worker/test_worker_memory_snapshot.py
-  - vllm/platforms/rocm.py
-  commands:
-  - VLLM_TEST_SAME_HOST=1 torchrun --nproc-per-node=4 distributed/test_same_node.py | grep 'Same node test passed'
-  - VLLM_TEST_SAME_HOST=1 VLLM_TEST_WITH_DEFAULT_DEVICE_SET=1 torchrun --nproc-per-node=4 distributed/test_same_node.py | grep 'Same node test passed'
-  - CUDA_VISIBLE_DEVICES=0,1 pytest -v -s v1/shutdown
-  - pytest -v -s v1/worker/test_worker_memory_snapshot.py
-
- label: Distributed Compile + Comm (4 GPUs) # TBD
-  timeout_in_minutes: 180
-  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325]
-  agent_pool: mi325_4
-  num_gpus: 4
-  working_dir: "/vllm-workspace/tests"
-  source_file_dependencies:
-  - vllm/distributed/
-  - tests/distributed/test_pynccl
-  - tests/distributed/test_events
-  - tests/compile/fullgraph/test_basic_correctness.py
-  - tests/distributed/test_symm_mem_allreduce.py
-  - tests/distributed/test_multiproc_executor.py
-  - vllm/platforms/rocm.py
-  commands:
-  - pytest -v -s compile/fullgraph/test_basic_correctness.py
-  - pytest -v -s distributed/test_pynccl.py
-  - pytest -v -s distributed/test_events.py
-  - pytest -v -s distributed/test_symm_mem_allreduce.py
-  - pytest -v -s distributed/test_multiproc_executor.py::test_multiproc_executor_multi_node
-
-#----------------------------------------------------------  mi325 · engine  -----------------------------------------------------------#
-
- label: Engine # TBD
-  timeout_in_minutes: 180
-  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325]
-  agent_pool: mi325_1
-  optional: true
-  working_dir: "/vllm-workspace/tests"
-  source_file_dependencies:
-  - vllm/
-  - tests/engine
-  - tests/test_sequence
-  - tests/test_config
-  - tests/test_logger
-  - tests/test_vllm_port
-  commands:
-  - pytest -v -s engine test_sequence.py test_config.py test_logger.py test_vllm_port.py test_jit_monitor.py
-
-#-----------------------------------------------------------  mi325 · evals  -----------------------------------------------------------#
-
- label: LM Eval Large Models (4xH100-4xMI325) # TBD
-  timeout_in_minutes: 180
-  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325]
-  agent_pool: mi325_4
-  num_gpus: 4
-  optional: true
-  working_dir: "/vllm-workspace/.buildkite/lm-eval-harness"
-  source_file_dependencies:
-  - csrc/
-  - vllm/model_executor/layers/quantization
-  - vllm/model_executor/models/
-  - vllm/model_executor/model_loader/
-  - vllm/v1/attention/backends/
-  - vllm/v1/attention/selector.py
-  - vllm/_aiter_ops.py
-  - vllm/platforms/rocm.py
-  commands:
-  - export VLLM_USE_DEEP_GEMM=0
-  - pytest -s -v test_lm_eval_correctness.py --config-list-file=configs/models-large-rocm-fp8.txt --tp-size=4
-
-#-----------------------------------------------------  mi325 · models / language  -----------------------------------------------------#
-
- label: Language Models Test (Extended Generation) # TBD
-  timeout_in_minutes: 180
-  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325]
-  agent_pool: mi325_1
-  working_dir: "/vllm-workspace/tests"
-  source_file_dependencies:
-  - vllm/
-  - tests/models/language/generation
-  commands:
-  - uv pip install --system --no-build-isolation 'git+https://github.com/AndreasKaratzas/mamba@fix-rocm-7.0-warp-size-constexpr'
-  - uv pip install --system --no-build-isolation 'git+https://github.com/Dao-AILab/causal-conv1d@v1.6.0'
-  - pytest -v -s models/language/generation -m '(not core_model) and (not hybrid_model)'
-
- label: Language Models Tests (Hybrid) %N # TBD
-  timeout_in_minutes: 180
-  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325]
-  agent_pool: mi325_1
-  parallelism: 2
-  optional: true
-  working_dir: "/vllm-workspace/tests"
-  source_file_dependencies:
-  - vllm/
-  - tests/models/language/generation
-  commands:
-  - uv pip install --system --no-build-isolation 'git+https://github.com/AndreasKaratzas/mamba@fix-rocm-7.0-warp-size-constexpr'
-  - uv pip install --system --no-build-isolation 'git+https://github.com/Dao-AILab/causal-conv1d@v1.6.0'
-  - pytest -v -s models/language/generation -m hybrid_model --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB
-
-#----------------------------------------------------  mi325 · models / multimodal  ----------------------------------------------------#
-
- label: Multi-Modal Models (Extended Pooling) # TBD
-  timeout_in_minutes: 180
-  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325]
-  agent_pool: mi325_1
-  optional: true
-  working_dir: "/vllm-workspace/tests"
-  source_file_dependencies:
-  - vllm/
-  - tests/models/multimodal/pooling
-  commands:
-  - pytest -v -s models/multimodal/pooling -m 'not core_model'
-
- label: "Multi-Modal Models (Standard) 2: qwen3 + gemma" # TBD
-  timeout_in_minutes: 180
-  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325]
-  agent_pool: mi325_1
-  optional: true
-  working_dir: "/vllm-workspace/tests"
-  source_file_dependencies:
-  - vllm/
-  - tests/models/multimodal
-  commands:
-  - pytest -v -s models/multimodal/generation/test_common.py -m core_model -k "qwen3 or gemma"
-  - pytest -v -s models/multimodal/generation/test_qwen2_5_vl.py -m core_model
-
-#-----------------------------------------------------------  mi325 · misc  ------------------------------------------------------------#
-
- label: Python-only Installation # TBD
-  timeout_in_minutes: 180
-  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325]
-  agent_pool: mi325_1
-  optional: true
-  working_dir: "/vllm-workspace/tests"
-  source_file_dependencies:
-  - tests/standalone_tests/python_only_compile.sh
-  - setup.py
-  - vllm/platforms/rocm.py
-  commands:
-  - bash standalone_tests/python_only_compile.sh
-
 #########################################################################################################################################
 #                                                                                                                                       #
 #                                                         MI355 (gfx950) tests                                                          #
@@ -3030,6 +3046,7 @@ steps:
 - label: Attention Benchmarks Smoke Test (B200-MI355) # TBD
  timeout_in_minutes: 180
  mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
+  dind: false
  agent_pool: mi355_2
  num_gpus: 2
  working_dir: "/vllm-workspace/"
@@ -3046,6 +3063,7 @@ steps:
 - label: Distributed Tests (2xH100-2xMI355) # TBD
  timeout_in_minutes: 180
  mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
+  dind: false
  agent_pool: mi355_2
  num_gpus: 2
  optional: true
@@ -3090,6 +3108,7 @@ steps:
 - label: Entrypoints Integration (API Server) # TBD
  timeout_in_minutes: 180
  mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
+  dind: false
  agent_pool: mi355_1
  optional: true
  fast_check: true
@@ -3107,6 +3126,7 @@ steps:
 - label: Entrypoints Integration (API Server OpenAI - Part 1) # TBD
  timeout_in_minutes: 180
  mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
+  dind: false
  agent_pool: mi355_1
  fast_check: true
  optional: true
@@ -3122,6 +3142,7 @@ steps:
 - label: Entrypoints Integration (API Server OpenAI - Part 2) # TBD
  timeout_in_minutes: 180
  mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
+  dind: false
  agent_pool: mi355_1
  fast_check: true
  optional: true
@@ -3138,6 +3159,7 @@ steps:
 - label: Entrypoints Integration (API Server Generate) # TBD
  timeout_in_minutes: 180
  mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
+  dind: false
  agent_pool: mi355_1
  fast_check: true
  optional: true
@@ -3158,6 +3180,7 @@ steps:
 - label: Entrypoints Integration (Speech to Text) # TBD
  timeout_in_minutes: 180
  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi355]
+  dind: false
  agent_pool: mi355_1
  fast_check: true
  working_dir: "/vllm-workspace/tests"
@@ -3171,6 +3194,7 @@ steps:
 - label: Entrypoints Integration (Multimodal)
  timeout_in_minutes: 180
  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi355]
+  dind: false
  agent_pool: mi355_1
  fast_check: true
  working_dir: "/vllm-workspace/tests"
@@ -3184,6 +3208,7 @@ steps:
 - label: Entrypoints Integration (Pooling) # TBD
  timeout_in_minutes: 180
  mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
+  dind: false
  agent_pool: mi355_1
  fast_check: true
  working_dir: "/vllm-workspace/tests"
@@ -3199,6 +3224,7 @@ steps:
 - label: GPQA Eval (GPT-OSS) (2xB200-2xMI355) # TBD
  timeout_in_minutes: 180
  mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
+  dind: false
  agent_pool: mi355_2
  num_gpus: 2
  optional: true
@@ -3221,6 +3247,7 @@ steps:
 - label: LM Eval Qwen3-5 Models (B200-MI355) # TBD
  timeout_in_minutes: 180
  mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
+  dind: false
  agent_pool: mi355_2
  num_gpus: 2
  optional: true
@@ -3243,6 +3270,7 @@ steps:
 - label: LM Eval Small Models (2xB200-2xMI355) # TBD
  timeout_in_minutes: 180
  mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
+  dind: false
  agent_pool: mi355_2
  num_gpus: 2
  optional: true
@@ -3262,6 +3290,7 @@ steps:
 - label: Qwen3-30B-A3B-FP8-block Sync EPLB Accuracy (B200-MI355) # TBD
  timeout_in_minutes: 180
  mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
+  dind: false
  agent_pool: mi355_2
  num_gpus: 2
  working_dir: "/vllm-workspace"
@@ -3282,6 +3311,7 @@ steps:
 - label: LM Eval Large Models (4xH100-4xMI355) # TBD
  timeout_in_minutes: 180
  mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
+  dind: false
  agent_pool: mi355_4
  num_gpus: 4
  optional: true
@@ -3304,6 +3334,7 @@ steps:
 - label: Examples # TBD
  timeout_in_minutes: 180
  mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
+  dind: false
  agent_pool: mi355_1
  working_dir: "/vllm-workspace/examples"
  source_file_dependencies:
@@ -3339,6 +3370,7 @@ steps:
 - label: Kernels (B200-MI355) # TBD
  timeout_in_minutes: 180
  mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
+  dind: false
  agent_pool: mi355_1
  working_dir: "/vllm-workspace/"
  source_file_dependencies:
@@ -3364,6 +3396,7 @@ steps:
 - label: Kernels Attention Test %N # TBD
  timeout_in_minutes: 180
  mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
+  dind: false
  agent_pool: mi355_1
  parallelism: 2
  working_dir: "/vllm-workspace/tests"
@@ -3381,6 +3414,7 @@ steps:
 - label: Kernels MoE Test %N # TBD
  timeout_in_minutes: 180
  mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
+  dind: false
  agent_pool: mi355_1
  parallelism: 5
  working_dir: "/vllm-workspace/tests"
@@ -3401,6 +3435,7 @@ steps:
 - label: Kernels Quantization Test %N # TBD
  timeout_in_minutes: 180
  mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
+  dind: false
  agent_pool: mi355_1
  parallelism: 2
  working_dir: "/vllm-workspace/tests"
@@ -3418,6 +3453,7 @@ steps:
 - label: Kernels FP8 MoE Test (2xH100-2xMI355) # TBD
  timeout_in_minutes: 180
  mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
+  dind: false
  agent_pool: mi355_2
  num_gpus: 2
  working_dir: "/vllm-workspace/tests"
@@ -3437,6 +3473,7 @@ steps:
 - label: Language Models Test (Extended Generation) # TBD
  timeout_in_minutes: 180
  mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
+  dind: false
  agent_pool: mi355_1
  working_dir: "/vllm-workspace/tests"
  source_file_dependencies:
@@ -3450,6 +3487,7 @@ steps:
 - label: Language Models Test (Extended Pooling) # TBD
  timeout_in_minutes: 180
  mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
+  dind: false
  agent_pool: mi355_1
  optional: true
  working_dir: "/vllm-workspace/tests"
@@ -3462,7 +3500,9 @@ steps:
 - label: Language Models Test (PPL) # TBD
  timeout_in_minutes: 180
  mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
+  dind: false
  agent_pool: mi355_1
+  optional: true
  working_dir: "/vllm-workspace/tests"
  source_file_dependencies:
  - vllm/model_executor/models/qwen3_5.py
@@ -3489,6 +3529,7 @@ steps:
 - label: Language Models Tests (Standard) # TBD
  timeout_in_minutes: 180
  mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
+  dind: false
  agent_pool: mi355_1
  working_dir: "/vllm-workspace/tests"
  source_file_dependencies:
@@ -3503,6 +3544,7 @@ steps:
 - label: Multi-Modal Models (Extended Generation 1) # TBD
  timeout_in_minutes: 180
  mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
+  dind: false
  agent_pool: mi355_1
  optional: true
  working_dir: "/vllm-workspace/tests"
@@ -3517,6 +3559,7 @@ steps:
 - label: Multi-Modal Models (Extended Generation 3) # TBD
  timeout_in_minutes: 180
  mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
+  dind: false
  agent_pool: mi355_1
  optional: true
  working_dir: "/vllm-workspace/tests"
@@ -3529,6 +3572,7 @@ steps:
 - label: Multi-Modal Models (Extended Pooling) # TBD
  timeout_in_minutes: 180
  mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
+  dind: false
  agent_pool: mi355_1
  optional: true
  working_dir: "/vllm-workspace/tests"
@@ -3541,6 +3585,7 @@ steps:
 - label: "Multi-Modal Models (Standard) 1: qwen2" # TBD
  timeout_in_minutes: 180
  mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
+  dind: false
  agent_pool: mi355_1
  optional: true
  working_dir: "/vllm-workspace/tests"
@@ -3554,6 +3599,7 @@ steps:
 - label: "Multi-Modal Models (Standard) 4: other + whisper" # TBD
  timeout_in_minutes: 180
  mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
+  dind: false
  agent_pool: mi355_1
  optional: true
  working_dir: "/vllm-workspace/tests"
@@ -3570,6 +3616,7 @@ steps:
 - label: Quantized Models Test # TBD
  timeout_in_minutes: 180
  mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
+  dind: false
  agent_pool: mi355_1
  working_dir: "/vllm-workspace/tests"
  source_file_dependencies:
@@ -3586,6 +3633,7 @@ steps:
 - label: Quantization # TBD
  timeout_in_minutes: 180
  mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
+  dind: false
  agent_pool: mi355_1
  working_dir: "/vllm-workspace/tests"
  source_file_dependencies:
@@ -3602,6 +3650,7 @@ steps:
 # - label: Quantized MoE Test (B200-MI355) # TBD
 #   timeout_in_minutes: 180
 #   mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
+#   dind: false
 #   agent_pool: mi355_1
 #   working_dir: "/vllm-workspace/"
 #   source_file_dependencies:
@@ -3630,6 +3679,7 @@ steps:
 - label: V1 attention (B200-MI355) # TBD
  timeout_in_minutes: 180
  mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
+  dind: false
  agent_pool: mi355_1
  working_dir: "/vllm-workspace/tests"
  source_file_dependencies:
@@ -3646,6 +3696,7 @@ steps:
 - label: V1 Core + KV + Metrics # TBD
  timeout_in_minutes: 180
  mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
+  dind: false
  agent_pool: mi355_1
  optional: true
  working_dir: "/vllm-workspace/tests"
@@ -3672,6 +3723,7 @@ steps:
 - label: V1 Sample + Logits # TBD
  timeout_in_minutes: 180
  mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
+  dind: false
  agent_pool: mi355_1
  optional: true
  working_dir: "/vllm-workspace/tests"
@@ -3692,6 +3744,7 @@ steps:
 - label: V1 Spec Decode # TBD
  timeout_in_minutes: 180
  mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
+  dind: false
  agent_pool: mi355_1
  working_dir: "/vllm-workspace/tests"
  source_file_dependencies:
@@ -3705,6 +3758,7 @@ steps:
 - label: Weight Loading Multiple GPU # TBD
  timeout_in_minutes: 180
  mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
+  dind: false
  agent_pool: mi355_2
  num_gpus: 2
  working_dir: "/vllm-workspace/tests"
@@ -3717,6 +3771,7 @@ steps:
 - label: Weight Loading Multiple GPU - Large Models # TBD
  timeout_in_minutes: 180
  mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
+  dind: false
  agent_pool: mi355_2
  working_dir: "/vllm-workspace/tests"
  num_gpus: 2
@@ -3732,6 +3787,7 @@ steps:
 - label: Regression # TBD
  timeout_in_minutes: 180
  mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
+  dind: false
  agent_pool: mi355_1
  optional: true
  working_dir: "/vllm-workspace/tests"
@@ -16,8 +16,9 @@ steps:
  parallelism: 2
  mirror:
    amd:
-      device: mi325_1
-      timeout_in_minutes: 95
+      dind: false
+      device: mi300_1
+      timeout_in_minutes: 125
      depends_on:
      - image-build-amd
      source_file_dependencies:
@@ -18,7 +18,8 @@ steps:
  - pytest -v -s basic_correctness/test_cpu_offload.py
  mirror:
    amd:
-      device: mi325_1
-      timeout_in_minutes: 70
+      dind: false
+      device: mi300_1
+      timeout_in_minutes: 60
      depends_on:
      - image-build-amd
@@ -15,6 +15,7 @@ steps:
    amd:
      dind: false
      device: mi300_1
+      timeout_in_minutes: 40
      depends_on:
      - image-build-amd

@@ -17,7 +17,7 @@ steps:
    amd:
      dind: false
      device: mi300_4
-      timeout_in_minutes: 85
+      timeout_in_minutes: 60
      depends_on:
        - image-build-amd
      source_file_dependencies:
@@ -68,7 +68,7 @@ steps:
    amd:
      dind: false
      device: mi300_4
-      timeout_in_minutes: 60
+      timeout_in_minutes: 40
      depends_on:
        - image-build-amd
      source_file_dependencies:
@@ -94,7 +94,7 @@ steps:
    amd:
      dind: false
      device: mi300_4
-      timeout_in_minutes: 85
+      timeout_in_minutes: 60
      depends_on:
        - image-build-amd
      source_file_dependencies:
@@ -120,7 +120,7 @@ steps:
    amd:
      dind: false
      device: mi300_4
-      timeout_in_minutes: 80
+      timeout_in_minutes: 55
      depends_on:
        - image-build-amd
      source_file_dependencies:
@@ -177,7 +177,7 @@ steps:
    amd:
      dind: false
      device: mi300_2
-      timeout_in_minutes: 70
+      timeout_in_minutes: 45
      depends_on:
        - image-build-amd
      source_file_dependencies:
@@ -41,6 +41,7 @@ steps:
    amd:
      dind: false
      device: mi300_2
+      timeout_in_minutes: 45
      depends_on:
      - image-build-amd
      source_file_dependencies:
@@ -28,8 +28,9 @@ steps:
  - pytest -v -s engine test_sequence.py test_config.py test_logger.py test_vllm_port.py test_jit_monitor.py
  mirror:
    amd:
-      device: mi325_1
-      timeout_in_minutes: 50
+      dind: false
+      device: mi300_1
+      timeout_in_minutes: 40
      depends_on:
      - image-build-amd

@@ -44,8 +45,8 @@ steps:
    - pytest -v -s v1/engine --ignore v1/engine/test_preprocess_error_handling.py
  mirror:
    amd:
-      device: mi325_1
-      timeout_in_minutes: 55
+      device: mi250_1
+      timeout_in_minutes: 45
      depends_on:
      - image-build-amd

@@ -60,8 +61,8 @@ steps:
    - pytest -v -s v1/e2e/general/test_async_scheduling.py
  mirror:
    amd:
-      device: mi325_1
-      timeout_in_minutes: 70
+      device: mi250_1
+      timeout_in_minutes: 55
      depends_on:
      - image-build-amd

@@ -76,8 +77,8 @@ steps:
    - pytest -v -s v1/e2e/general --ignore v1/e2e/general/test_async_scheduling.py
  mirror:
    amd:
-      device: mi325_1
-      timeout_in_minutes: 60
+      device: mi250_1
+      timeout_in_minutes: 50
      depends_on:
      - image-build-amd
      source_file_dependencies:
@@ -116,6 +117,7 @@ steps:
    amd:
      dind: false
      device: mi300_2
+      timeout_in_minutes: 30
      depends_on:
      - image-build-amd

@@ -30,9 +30,9 @@ steps:
  - pytest -v -s entrypoints/llm/offline_mode # Needs to avoid interference with other tests
  mirror:
    amd:
-      device: mi325_1
-      # TODO(akaratza): Test after Torch >= 2.12 bump
-      soft_fail: true
+      dind: false
+      device: mi300_1
+      timeout_in_minutes: 55
      depends_on:
      - image-build-amd

@@ -52,7 +52,9 @@ steps:
  - pytest -v -s entrypoints/scale_out
  mirror:
    amd:
-      device: mi325_1
+      dind: false
+      device: mi300_1
+      timeout_in_minutes: 65
      depends_on:
      - image-build-amd

@@ -70,7 +72,8 @@ steps:
  - pytest -v -s entrypoints/openai --ignore=entrypoints/openai/completion --ignore=entrypoints/openai/chat_completion --ignore=entrypoints/openai/responses --ignore=entrypoints/openai/correctness
  mirror:
    amd:
-      device: mi325_1
+      dind: false
+      device: mi300_1
      timeout_in_minutes: 65
      depends_on:
      - image-build-amd
@@ -90,8 +93,9 @@ steps:
  - pytest -v -s entrypoints/openai/completion --ignore=entrypoints/openai/completion/test_tensorizer_entrypoint.py
  mirror:
    amd:
-      device: mi325_1
-      timeout_in_minutes: 80
+      dind: false
+      device: mi300_1
+      timeout_in_minutes: 70
      depends_on:
      - image-build-amd

@@ -113,7 +117,8 @@ steps:
  - pytest -v -s entrypoints/anthropic
  mirror:
    amd:
-      device: mi325_1
+      dind: false
+      device: mi300_1
      timeout_in_minutes: 65
      depends_on:
      - image-build-amd
@@ -176,7 +181,9 @@ steps:
  - pytest -s entrypoints/openai/correctness/
  mirror:
    amd:
-      device: mi325_1
+      dind: false
+      device: mi300_1
+      timeout_in_minutes: 30
      depends_on:
      - image-build-amd
      source_file_dependencies:
@@ -18,6 +18,7 @@ steps:
    amd:
      dind: false
      device: mi300_1
+      timeout_in_minutes: 30
      depends_on:
      - image-build-amd
      source_file_dependencies:
@@ -80,7 +80,8 @@ steps:
  parallelism: 2
  mirror:
    amd:
-      device: mi325_1
+      dind: false
+      device: mi300_1
      timeout_in_minutes: 90
      depends_on:
      - image-build-amd
@@ -118,7 +119,9 @@ steps:
  parallelism: 2
  mirror:
    amd:
-      device: mi325_1
+      dind: false
+      device: mi300_1
+      timeout_in_minutes: 120
      source_file_dependencies:
      - csrc/quantization/
      - vllm/model_executor/layers/quantization
@@ -148,8 +151,9 @@ steps:
  parallelism: 5
  mirror:
    amd:
-      device: mi325_1
-      timeout_in_minutes: 65
+      dind: false
+      device: mi300_1
+      timeout_in_minutes: 55
      source_file_dependencies:
      - csrc/quantization/cutlass_w8a8/moe/
      - csrc/moe/
@@ -14,8 +14,9 @@ steps:
  - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-small.txt
  mirror:
    amd:
-      device: mi325_1
-      timeout_in_minutes: 55
+      dind: false
+      device: mi300_1
+      timeout_in_minutes: 45
      depends_on:
      - image-build-amd
      source_file_dependencies:
@@ -141,7 +142,7 @@ steps:
    amd:
      dind: false
      device: mi300_8
-      timeout_in_minutes: 60
+      timeout_in_minutes: 40
      depends_on:
      - image-build-amd
      commands:
@@ -14,9 +14,10 @@ steps:
  parallelism: 4
  mirror:
    amd:
-      device: mi325_1
+      dind: false
+      device: mi300_1
      working_dir: "/vllm-workspace/tests"
-      timeout_in_minutes: 65
+      timeout_in_minutes: 85
      source_file_dependencies:
      - vllm/lora
      - tests/lora
@@ -46,4 +47,4 @@ steps:
    - pytest -v -s -x lora/test_qwen3_with_multi_loras.py
    - pytest -v -s -x lora/test_olmoe_tp.py
    - pytest -v -s -x lora/test_gptoss_tp.py
-    - pytest -v -s -x lora/test_qwen35_densemodel_lora.py
+    - pytest -v -s -x lora/test_qwen35_densemodel_lora.py
@@ -25,7 +25,7 @@ steps:
    amd:
      dind: false
      device: mi300_1
-      timeout_in_minutes: 75
+      timeout_in_minutes: 50
      depends_on:
      - image-build-amd

@@ -59,7 +59,9 @@ steps:
    - pytest -v -s v1/test_outputs.py
  mirror:
    amd:
-      device: mi325_1
+      dind: false
+      device: mi300_1
+      timeout_in_minutes: 70
      depends_on:
      - image-build-amd

@@ -113,8 +115,9 @@ steps:
    - pytest -v -s entrypoints/openai/correctness/test_lmeval.py::test_lm_eval_accuracy_v1_engine
  mirror:
    amd:
-      device: mi325_1
-      timeout_in_minutes: 75
+      dind: false
+      device: mi300_1
+      timeout_in_minutes: 65
      depends_on:
      - image-build-amd

@@ -209,7 +212,7 @@ steps:
  - vllm/multimodal
  - examples/
  commands:
-    - pip install tensorizer # for tensorizer test
+    - pip install --no-deps tensorizer # for tensorizer test
     # for basic
    - python3 basic/offline_inference/chat.py
    - python3 basic/offline_inference/generate.py --model facebook/opt-125m
@@ -233,7 +236,9 @@ steps:
    - python3 features/speculative_decoding/spec_decode_offline.py --test --method eagle3 --num_spec_tokens 3 --dataset-name hf --dataset-path philschmid/mt-bench --num-prompts 80 --temp 0 --top-p 1.0 --top-k -1 --tp 1 --enable-chunked-prefill --max-model-len 1536
  mirror:
    amd:
-      device: mi325_1
+      dind: false
+      device: mi300_1
+      timeout_in_minutes: 75
      source_file_dependencies:
      - vllm/entrypoints
      - vllm/multimodal
@@ -269,10 +274,11 @@ steps:
  - pytest -v -s v1/tracing
  mirror:
    amd:
-      device: mi325_2
+      dind: false
+      device: mi300_2
+      timeout_in_minutes: 30
      depends_on:
      - image-build-amd
-      optional: true

 - label: Python-only Installation
  key: python-only-installation
@@ -287,8 +293,8 @@ steps:
  - bash standalone_tests/python_only_compile.sh
  mirror:
    amd:
-      device: mi325_1
-      timeout_in_minutes: 45
+      device: mi250_1
+      timeout_in_minutes: 55
      depends_on:
      - image-build-amd
      source_file_dependencies:
@@ -30,6 +30,7 @@ steps:
    amd:
      dind: false
      device: mi300_1
+      timeout_in_minutes: 60
      depends_on:
      - image-build-amd
      source_file_dependencies:
@@ -41,7 +41,7 @@ steps:
  commands:
    - set -x
    - export VLLM_USE_V2_MODEL_RUNNER=1
-    - pip install tensorizer # for tensorizer test
+    - pip install --no-deps tensorizer # for tensorizer test
    - python3 basic/offline_inference/chat.py # for basic
    - python3 basic/offline_inference/generate.py --model facebook/opt-125m
    #- python3 basic/offline_inference/generate.py --model meta-llama/Llama-2-13b-chat-hf --cpu-offload-gb 10  # TODO
@@ -42,7 +42,9 @@ steps:
    - pytest -v -s models/test_terratorch.py models/transformers/test_backend.py models/test_registry.py
  mirror:
    amd:
-      device: mi325_1
+      dind: false
+      device: mi300_1
+      timeout_in_minutes: 50
      depends_on:
      - image-build-amd

@@ -17,6 +17,7 @@ steps:
    amd:
      dind: false
      device: mi300_1
+      timeout_in_minutes: 45
      depends_on:
      - image-build-amd

@@ -39,6 +40,7 @@ steps:
    amd:
      dind: false
      device: mi300_1
+      timeout_in_minutes: 40
      depends_on:
      - image-build-amd
      source_file_dependencies:
@@ -61,7 +63,6 @@ steps:
  - tests/models/language/generation
  commands:
    # Install fast path packages for testing against transformers
-    # Note: also needed to run plamo2 model in vLLM
    - uv pip install --system --no-build-isolation 'git+https://github.com/state-spaces/mamba@v2.3.0'
    - uv pip install --system --no-build-isolation 'git+https://github.com/Dao-AILab/causal-conv1d@v1.6.0'
    # Shard the hybrid language model tests that are numerically stable on Hopper.
@@ -69,8 +70,9 @@ steps:
  parallelism: 2
  mirror:
    amd:
-      device: mi325_1
-      timeout_in_minutes: 70
+      dind: false
+      device: mi300_1
+      timeout_in_minutes: 60
      depends_on:
      - image-build-amd
      commands:
@@ -102,7 +104,6 @@ steps:
  - tests/models/language/generation
  commands:
    # Install fast path packages for testing against transformers
-    # Note: also needed to run plamo2 model in vLLM
    - uv pip install --system --no-build-isolation 'git+https://github.com/state-spaces/mamba@v2.3.0'
    - uv pip install --system --no-build-isolation 'git+https://github.com/Dao-AILab/causal-conv1d@v1.6.0'
    - pytest -v -s models/language/generation -m '(not core_model) and (not hybrid_model)'
@@ -130,8 +131,9 @@ steps:
    - pytest -v -s models/language/pooling -m 'not core_model'
  mirror:
    amd:
-      device: mi325_1
-      timeout_in_minutes: 120
+      dind: false
+      device: mi300_1
+      timeout_in_minutes: 95
      depends_on:
      - image-build-amd

@@ -14,7 +14,9 @@ steps:
    - pytest -v -s models/multimodal/generation/test_ultravox.py -m core_model
  mirror:
    amd:
-      device: mi325_1
+      dind: false
+      device: mi300_1
+      timeout_in_minutes: 65
      depends_on:
      - image-build-amd

@@ -31,7 +33,9 @@ steps:
    - pytest -v -s models/multimodal/generation/test_qwen2_5_vl.py -m core_model
  mirror:
    amd:
-      device: mi325_1
+      dind: false
+      device: mi300_1
+      timeout_in_minutes: 55
      depends_on:
      - image-build-amd

@@ -47,7 +51,8 @@ steps:
    - pytest -v -s models/multimodal/generation/test_qwen2_vl.py -m core_model
  mirror:
    amd:
-      device: mi325_1
+      device: mi250_1
+      timeout_in_minutes: 55
      depends_on:
      - image-build-amd

@@ -65,7 +70,9 @@ steps:
    - cd .. && VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s tests/models/multimodal/generation/test_whisper.py -m core_model  # Otherwise, mp_method="spawn" doesn't work
  mirror:
    amd:
-      device: mi325_1
+      dind: false
+      device: mi300_1
+      timeout_in_minutes: 50
      depends_on:
      - image-build-amd

@@ -109,6 +116,7 @@ steps:
    amd:
      dind: false
      device: mi300_1
+      timeout_in_minutes: 35
      depends_on:
      - image-build-amd
      source_file_dependencies:
@@ -131,7 +139,9 @@ steps:
    - pytest -v -s models/multimodal/test_mapping.py
  mirror:
    amd:
-      device: mi325_1
+      dind: false
+      device: mi300_1
+      timeout_in_minutes: 90
      depends_on:
      - image-build-amd

@@ -166,8 +176,9 @@ steps:
    - pytest -v -s models/multimodal/pooling -m 'not core_model'
  mirror:
    amd:
-      device: mi325_1
-      timeout_in_minutes: 75
+      dind: false
+      device: mi300_1
+      timeout_in_minutes: 60
      depends_on:
      - image-build-amd
      source_file_dependencies:
@@ -107,13 +107,6 @@ steps:
    - tests/compile/passes
  commands:
  - pytest -s -v compile/passes --ignore compile/passes/distributed
-  mirror:
-    amd:
-      dind: false
-      device: mi300_1
-      timeout_in_minutes: 65
-      depends_on:
-      - image-build-amd

 - label: PyTorch Fullgraph Smoke Test
  device: h200_35gb
@@ -236,6 +229,7 @@ steps:
    amd:
      dind: false
      device: mi300_1
+      timeout_in_minutes: 30
      depends_on:
      - image-build-amd
      source_file_dependencies:
@@ -24,8 +24,7 @@ steps:
  - uv pip install --system conch-triton-kernels
  # The SM90-only checkpoint currently contains a removed weight_chan_scale
  # parameter. It was not exercised by the previous L4 job.
-  - VLLM_TEST_FORCE_LOAD_FORMAT=auto pytest -v -s quantization/ --ignore quantization/test_blackwell_moe.py -k 'not test_compressed_tensors_w4a8_fp8' --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT
-  parallelism: 8
+  - VLLM_TEST_FORCE_LOAD_FORMAT=auto pytest -v -s quantization/ --ignore quantization/test_blackwell_moe.py -k 'not test_compressed_tensors_w4a8_fp8'

 - label: Quantized Fusions
  device: h200_35gb
@@ -68,5 +67,4 @@ steps:
  - vllm/model_executor/layers/quantization
  - tests/models/quantization
  commands:
-    - pytest -v -s models/quantization --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT
-  parallelism: 3
+    - pytest -v -s models/quantization
@@ -19,8 +19,18 @@ steps:
    - VLLM_USE_FLASHINFER_SAMPLER=1 pytest -v -s samplers
  mirror:
    amd:
-      device: mi325_1
+      device: mi250_1
+      timeout_in_minutes: 40
      depends_on:
      - image-build-amd
+      source_file_dependencies:
+      - vllm/model_executor/layers
+      - vllm/sampling_metadata.py
+      - vllm/v1/sample/
+      - vllm/entrypoints/generate/beam_search/
+      - tests/samplers
+      - tests/conftest.py
+      - vllm/_aiter_ops.py
+      - vllm/platforms/rocm.py
      commands:
      - pytest -v -s samplers
@@ -14,8 +14,9 @@ steps:
    - pytest -v -s v1/e2e/spec_decode -k "eagle_correctness"
  mirror:
    amd:
-      device: mi325_1
-      timeout_in_minutes: 60
+      dind: false
+      device: mi300_1
+      timeout_in_minutes: 55
      depends_on:
      - image-build-amd
      source_file_dependencies:
@@ -53,8 +54,9 @@ steps:
    - pytest -v -s v1/e2e/spec_decode -k "speculators or mtp_correctness"
  mirror:
    amd:
-      device: mi325_1
-      timeout_in_minutes: 65
+      dind: false
+      device: mi300_1
+      timeout_in_minutes: 75
      depends_on:
      - image-build-amd
      source_file_dependencies:
@@ -92,10 +94,9 @@ steps:
    - pytest -v -s v1/e2e/spec_decode -k "ngram or suffix"
  mirror:
    amd:
-      device: mi325_1
-      timeout_in_minutes: 55
-      # TODO(akaratza): Test after Torch >= 2.12 bump
-      soft_fail: true
+      dind: false
+      device: mi300_1
+      timeout_in_minutes: 35
      depends_on:
      - image-build-amd
      source_file_dependencies:
@@ -119,7 +120,8 @@ steps:
    - pytest -v -s v1/e2e/spec_decode -k "draft_model or no_sync or batch_inference"
  mirror:
    amd:
-      device: mi325_1
+      dind: false
+      device: mi300_1
      timeout_in_minutes: 55
      depends_on:
      - image-build-amd
@@ -17,6 +17,7 @@ steps:
    amd:
      dind: false
      device: mi300_2
+      timeout_in_minutes: 35
      depends_on:
      - image-build-amd
      commands:
@@ -3,6 +3,7 @@
 dist
 vllm/*.so
 vllm/vllm-rs
+.git

 # Byte-compiled / optimized / DLL files
 __pycache__/
@@ -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
@@ -114,6 +114,11 @@ find_package(Torch REQUIRED)
 # Supported NVIDIA architectures.
 # This check must happen after find_package(Torch) because that's when CMAKE_CUDA_COMPILER_VERSION gets defined
 if(DEFINED CMAKE_CUDA_COMPILER_VERSION AND
+   CMAKE_CUDA_COMPILER_VERSION VERSION_GREATER_EQUAL 13.4)
+  # Rubin (10.7) can run SM100 family code, but CUDA 13.4 also supports
+  # targeting it directly.
+  set(CUDA_SUPPORTED_ARCHS "7.5;8.0;8.6;8.7;8.9;9.0;10.0;10.7;11.0;12.0")
+elseif(DEFINED CMAKE_CUDA_COMPILER_VERSION AND
   CMAKE_CUDA_COMPILER_VERSION VERSION_GREATER_EQUAL 13.0)
   # starting from CUDA 12.9 and Blackwell (10.0), we use family-specific targets (10.0f, 12.0f, etc)
   # to support the whole generation without specifying all sub-architectures
@@ -411,8 +416,11 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
    "csrc/libtorch_stable/mamba/selective_scan_fwd.cu"
    "csrc/libtorch_stable/cache_kernels.cu"
    "csrc/libtorch_stable/cache_kernels_fused.cu"
+    "csrc/libtorch_stable/custom_all_gather_reduce_scatter.cu"
+    "csrc/libtorch_stable/custom_all_gather_reduce_scatter_ops.cpp"
    "csrc/libtorch_stable/custom_all_reduce.cu"
-    "csrc/libtorch_stable/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu")
+    "csrc/libtorch_stable/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu"
+    "csrc/libtorch_stable/fused_kimi_k3_mla_key_concat_kv_cache_kernel.cu")

  if(VLLM_GPU_LANG STREQUAL "CUDA" AND
     DEFINED CMAKE_CUDA_COMPILER_VERSION AND
@@ -420,7 +428,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")

    if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
      cuda_archs_loose_intersection(COOPERATIVE_TOPK_ARCHS
-        "9.0a;10.0f;10.1f;10.3f;11.0f;12.0f;12.1f" "${CUDA_ARCHS}")
+        "9.0a;10.0f;10.1f;10.3f;10.7f;11.0f;12.0f;12.1f" "${CUDA_ARCHS}")
    else()
      cuda_archs_loose_intersection(COOPERATIVE_TOPK_ARCHS
        "9.0a;10.0a;10.1a;10.3a;12.0a;12.1a" "${CUDA_ARCHS}")
@@ -695,7 +703,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")

  # DeepSeek V3 fused A GEMM kernel (requires SM 9.0+, Hopper and later)
  if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
-    cuda_archs_loose_intersection(DSV3_FUSED_A_GEMM_ARCHS "9.0a;10.0f;11.0f;12.0f" "${CUDA_ARCHS}")
+    cuda_archs_loose_intersection(DSV3_FUSED_A_GEMM_ARCHS "9.0a;10.0f;10.7f;11.0f;12.0f" "${CUDA_ARCHS}")
  else()
    cuda_archs_loose_intersection(DSV3_FUSED_A_GEMM_ARCHS "9.0a;10.0a;10.1a;10.3a;12.0a;12.1a" "${CUDA_ARCHS}")
  endif()
@@ -815,7 +823,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
  # The cutlass_scaled_mm kernels for Blackwell SM100 (c3x, i.e. CUTLASS 3.x)
  # require CUDA 12.8 or later
  if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
-    cuda_archs_loose_intersection(SCALED_MM_ARCHS "10.0f;11.0f" "${CUDA_ARCHS}")
+    cuda_archs_loose_intersection(SCALED_MM_ARCHS "10.0f;10.7f;11.0f" "${CUDA_ARCHS}")
  else()
    cuda_archs_loose_intersection(SCALED_MM_ARCHS "10.0a;10.1a;10.3a" "${CUDA_ARCHS}")
  endif()
@@ -899,7 +907,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
  endif()

  if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
-    cuda_archs_loose_intersection(SCALED_MM_ARCHS "10.0f;11.0f" "${CUDA_ARCHS}")
+    cuda_archs_loose_intersection(SCALED_MM_ARCHS "10.0f;10.7f;11.0f" "${CUDA_ARCHS}")
  else()
    cuda_archs_loose_intersection(SCALED_MM_ARCHS "10.0a;10.1a;10.3a" "${CUDA_ARCHS}")
  endif()
@@ -924,7 +932,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")

  # moe_data.cu is used by all CUTLASS MoE kernels.
  if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
-    cuda_archs_loose_intersection(CUTLASS_MOE_DATA_ARCHS "9.0a;10.0f;11.0f;12.0f" "${CUDA_ARCHS}")
+    cuda_archs_loose_intersection(CUTLASS_MOE_DATA_ARCHS "9.0a;10.0f;10.7f;11.0f;12.0f" "${CUDA_ARCHS}")
  else()
    cuda_archs_loose_intersection(CUTLASS_MOE_DATA_ARCHS "9.0a;10.0a;10.1a;10.3a;12.0a;12.1a" "${CUDA_ARCHS}")
  endif()
@@ -981,7 +989,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
  # SM10x/11x FP4 kernels. MXFP4 experts quantization is currently compiled
  # only in this block; SM12x has separate NVFP4 matmul/MoE kernels above.
  if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
-    cuda_archs_loose_intersection(FP4_SM100_ARCHS "10.0f;11.0f" "${CUDA_ARCHS}")
+    cuda_archs_loose_intersection(FP4_SM100_ARCHS "10.0f;10.7f;11.0f" "${CUDA_ARCHS}")
  else()
    cuda_archs_loose_intersection(FP4_SM100_ARCHS "10.0a;10.1a;10.3a" "${CUDA_ARCHS}")
  endif()
@@ -1047,7 +1055,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
  # Runtime dispatch is gated in
  # vllm/v1/attention/backends/mla/cutlass_mla.py.
  if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
-    cuda_archs_loose_intersection(MLA_ARCHS "10.0f;11.0f" "${CUDA_ARCHS}")
+    cuda_archs_loose_intersection(MLA_ARCHS "10.0f;10.7f;11.0f" "${CUDA_ARCHS}")
  else()
    cuda_archs_loose_intersection(MLA_ARCHS "10.0a;10.1a;10.3a" "${CUDA_ARCHS}")
  endif()
@@ -1069,6 +1077,41 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
    set(MLA_ARCHS)
  endif()

+  if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
+    cuda_archs_loose_intersection(FUSED_KDA_DECODE_ARCHS
+      "9.0a;10.0f;12.0f" "${CUDA_ARCHS}")
+  endif()
+  if(FUSED_KDA_DECODE_ARCHS)
+    set(FUSED_KDA_DECODE_SRC
+      "csrc/libtorch_stable/kimi_k3/fused_kda_decode_kernel.cu")
+    set_gencode_flags_for_srcs(
+      SRCS "${FUSED_KDA_DECODE_SRC}"
+      CUDA_ARCHS "${FUSED_KDA_DECODE_ARCHS}")
+    set_property(SOURCE ${FUSED_KDA_DECODE_SRC} APPEND PROPERTY
+      COMPILE_OPTIONS "$<$<COMPILE_LANGUAGE:CUDA>:--use_fast_math>")
+    list(APPEND VLLM_STABLE_EXT_SRC "${FUSED_KDA_DECODE_SRC}")
+    message(STATUS
+      "Building fused KDA decode for archs: ${FUSED_KDA_DECODE_ARCHS}")
+  endif()
+
+  if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
+    cuda_archs_loose_intersection(KIMI_K3_ATTN_RES_ARCHS
+      "10.0f" "${CUDA_ARCHS}")
+  endif()
+  if(KIMI_K3_ATTN_RES_ARCHS)
+    set(KIMI_K3_ATTN_RES_SRC
+      "csrc/libtorch_stable/kimi_k3/attn_res_kernel.cu")
+    set_gencode_flags_for_srcs(
+      SRCS "${KIMI_K3_ATTN_RES_SRC}"
+      CUDA_ARCHS "${KIMI_K3_ATTN_RES_ARCHS}")
+    set_property(SOURCE ${KIMI_K3_ATTN_RES_SRC} APPEND PROPERTY
+      COMPILE_OPTIONS
+      "$<$<COMPILE_LANGUAGE:CUDA>:--expt-relaxed-constexpr;--expt-extended-lambda;--use_fast_math>")
+    list(APPEND VLLM_STABLE_EXT_SRC "${KIMI_K3_ATTN_RES_SRC}")
+    message(STATUS
+      "Building Kimi K3 AttnRes for archs: ${KIMI_K3_ATTN_RES_ARCHS}")
+  endif()
+
  # Hadacore kernels
  cuda_archs_loose_intersection(HADACORE_ARCHS "8.0+PTX;9.0+PTX" "${CUDA_ARCHS}")
  if(HADACORE_ARCHS)
@@ -1110,6 +1153,14 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
      target_compile_definitions(_C_stable_libtorch PRIVATE
        VLLM_ENABLE_COOPERATIVE_TOPK=1)
    endif()
+    if(FUSED_KDA_DECODE_ARCHS)
+      target_compile_definitions(_C_stable_libtorch PRIVATE
+        VLLM_ENABLE_FUSED_KDA_DECODE=1)
+    endif()
+    if(KIMI_K3_ATTN_RES_ARCHS)
+      target_compile_definitions(_C_stable_libtorch PRIVATE
+        VLLM_ENABLE_KIMI_K3_ATTN_RES=1)
+    endif()
    # Needed by CUTLASS kernels
    target_compile_definitions(_C_stable_libtorch PRIVATE
      CUTLASS_ENABLE_DIRECT_CUDA_DRIVER_CALL=1)
@@ -1407,6 +1458,7 @@ if (VLLM_GPU_LANG STREQUAL "CUDA")
    include(cmake/external_projects/deepgemm.cmake)
    include(cmake/external_projects/fmha_sm100.cmake)
    include(cmake/external_projects/flashmla.cmake)
+    include(cmake/external_projects/flashkda.cmake)
    include(cmake/external_projects/qutlass.cmake)
    include(cmake/external_projects/tml_fa4.cmake)

@@ -0,0 +1,367 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Benchmark the Kimi-K3 latent MoE addmm against CuTe residual GEMM.
+
+The benchmark covers ``BF16[M, 3584] @ BF16[7168, 3584].T + BF16[M, 7168]``
+with FP32 accumulation and BF16 output. Both backends execute through CUDA
+Graph replay. Weights and residuals rotate across buffers exceeding L2 so the
+comparison models the full latent MoE projection-and-add path.
+"""
+
+from __future__ import annotations
+
+import argparse
+import dataclasses
+import importlib.util
+import json
+import math
+import statistics
+from collections.abc import Callable, Sequence
+from pathlib import Path
+from typing import Any
+
+import cutlass
+import cutlass.cute as cute
+import torch
+from cuda.bindings import driver as cuda
+from cuda.bindings.driver import CUstream
+from quack.compile_utils import make_fake_tensor
+
+N = 7168
+K = 3584
+
+
+@dataclasses.dataclass(frozen=True, slots=True)
+class Config:
+    block_size: int
+    outputs_per_block: int
+    k_unroll: int
+    vector_width: int = 8
+
+
+def parse_config(value: str) -> Config:
+    try:
+        parts = [int(part) for part in value.split(",")]
+    except ValueError as error:
+        raise argparse.ArgumentTypeError(
+            "config must be BLOCK,OUTPUTS,K_UNROLL[,VECTOR_WIDTH]"
+        ) from error
+    if len(parts) == 3:
+        return Config(*parts)
+    if len(parts) == 4:
+        return Config(*parts)
+    raise argparse.ArgumentTypeError(
+        "config must be BLOCK,OUTPUTS,K_UNROLL[,VECTOR_WIDTH]"
+    )
+
+
+def production_residual_config(m: int) -> Config | None:
+    """The measured Latent-MoE residual config for M, from the K3 table."""
+    from vllm.models.kimi_k3.nvidia.low_latency_gemm import KIMI_K3_PROJECTIONS
+
+    spec = KIMI_K3_PROJECTIONS.get((N, K))
+    config = spec.residual_config(m) if spec is not None else None
+    if config is None:
+        return None
+    return Config(
+        config.block_size,
+        config.outputs_per_block,
+        config.k_unroll,
+        config.vector_width,
+    )
+
+
+def candidate_configs(mode: str, selected: Config | None, m: int) -> list[Config]:
+    if mode == "selected":
+        if selected is not None:
+            return [selected]
+        # No explicit --config: fall back to the production table for this M.
+        config = production_residual_config(m)
+        return [config] if config is not None else []
+    if mode == "baseline":
+        return [Config(224, 4, 2)]
+    return [
+        Config(block_size, outputs_per_block, k_unroll, vector_width)
+        for vector_width in (4, 8)
+        for block_size in (32, 64, 128, 224, 448)
+        if block_size % 32 == 0 and K % (block_size * vector_width) == 0
+        for outputs_per_block in (1, 2, 4, 7, 8)
+        if N % outputs_per_block == 0
+        for k_unroll in (1, 2, 4)
+    ]
+
+
+def load_kernel_class(path: Path):
+    spec = importlib.util.spec_from_file_location("cute_skinny_device", path)
+    if spec is None or spec.loader is None:
+        raise RuntimeError(f"cannot load CuTe kernel from {path}")
+    module = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(module)
+    return module.CuteSkinnyGemm
+
+
+def stream() -> CUstream:
+    return CUstream(torch.cuda.current_stream().cuda_stream)
+
+
+def compile_kernel(kernel_class, m: int, config: Config, max_registers: int):
+    element_type = cutlass.BFloat16
+    n = cute.sym_int(divisibility=config.outputs_per_block)
+    k = cute.sym_int(divisibility=config.block_size * config.vector_width)
+    a = make_fake_tensor(element_type, (m, k), divisibility=config.vector_width)
+    b = make_fake_tensor(element_type, (n, k), divisibility=config.vector_width)
+    residual = make_fake_tensor(element_type, (m, n), divisibility=1)
+    c = make_fake_tensor(element_type, (m, n), divisibility=1)
+    kernel = kernel_class(
+        element_type=element_type,
+        num_rows=m,
+        block_size=config.block_size,
+        outputs_per_block=config.outputs_per_block,
+        vector_width=config.vector_width,
+        k_unroll=config.k_unroll,
+        has_residual=True,
+        use_pdl=True,
+    )
+    return cute.compile(
+        kernel,
+        a,
+        b,
+        residual,
+        c,
+        stream(),
+        options=(
+            "--enable-tvm-ffi --keep-cubin "
+            f"--ptxas-options -maxrregcount={max_registers} "
+            "--ptxas-options -lineinfo"
+        ),
+    )
+
+
+def resource_usage(compiled) -> dict[str, Any]:
+    executor = getattr(compiled, "_default_executor", None)
+    context = getattr(executor, "exec_context", None)
+    functions = getattr(context, "kernel_functions", None)
+    if not functions:
+        return {"resource_metrics_available": False}
+
+    def attribute(name, function) -> int:
+        error, value = cuda.cuFuncGetAttribute(name, function)
+        if error != cuda.CUresult.CUDA_SUCCESS:
+            raise RuntimeError(f"cuFuncGetAttribute failed with {error}")
+        return int(value)
+
+    registers = [
+        attribute(cuda.CUfunction_attribute.CU_FUNC_ATTRIBUTE_NUM_REGS, function)
+        for function in functions
+    ]
+    local_bytes = [
+        attribute(
+            cuda.CUfunction_attribute.CU_FUNC_ATTRIBUTE_LOCAL_SIZE_BYTES,
+            function,
+        )
+        for function in functions
+    ]
+    return {
+        "resource_metrics_available": True,
+        "registers_per_thread": max(registers, default=0),
+        "spill_bytes": max(local_bytes, default=0),
+    }
+
+
+def rotating_buffer_count(m: int, multiplier: float, limit: int) -> int:
+    properties = torch.cuda.get_device_properties(0)
+    bytes_per_pair = (N * K + m * N) * 2
+    target = math.ceil(multiplier * properties.L2_cache_size)
+    return max(2, min(limit, math.ceil(target / bytes_per_pair)))
+
+
+def graph_samples(
+    launch: Callable[[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor], None],
+    activation: torch.Tensor,
+    weights: Sequence[torch.Tensor],
+    residuals: Sequence[torch.Tensor],
+    repeats: int,
+    replays: int,
+) -> tuple[list[float], list[torch.Tensor]]:
+    outputs = [torch.empty_like(residual) for residual in residuals]
+    for weight, residual, output in zip(weights, residuals, outputs):
+        launch(activation, weight, residual, output)
+    torch.cuda.synchronize()
+
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph):
+        for weight, residual, output in zip(weights, residuals, outputs):
+            launch(activation, weight, residual, output)
+    for _ in range(20):
+        graph.replay()
+    torch.cuda.synchronize()
+
+    samples = []
+    for _ in range(repeats):
+        start = torch.cuda.Event(enable_timing=True)
+        end = torch.cuda.Event(enable_timing=True)
+        start.record()
+        for _ in range(replays):
+            graph.replay()
+        end.record()
+        end.synchronize()
+        samples.append(start.elapsed_time(end) * 1000.0 / (replays * len(weights)))
+    return samples, outputs
+
+
+def summarize(samples: Sequence[float]) -> dict[str, Any]:
+    ordered = sorted(samples)
+
+    def percentile(fraction: float) -> float:
+        position = fraction * (len(ordered) - 1)
+        lower = math.floor(position)
+        upper = math.ceil(position)
+        if lower == upper:
+            return ordered[lower]
+        weight = position - lower
+        return ordered[lower] * (1.0 - weight) + ordered[upper] * weight
+
+    mean = statistics.mean(samples)
+    return {
+        "median_us": statistics.median(samples),
+        "p10_us": percentile(0.1),
+        "p90_us": percentile(0.9),
+        "mean_us": mean,
+        "cv_pct": statistics.pstdev(samples) / mean * 100.0,
+        "samples_us": list(samples),
+    }
+
+
+def correctness(
+    output: torch.Tensor,
+    activation: torch.Tensor,
+    weight: torch.Tensor,
+    residual: torch.Tensor,
+) -> dict[str, Any]:
+    actual = output.float()
+    reference = activation.float() @ weight.float().t() + residual.float()
+    error = (actual - reference).abs()
+    scaled_error = error / (reference.abs() + 1.0)
+    cosine = torch.nn.functional.cosine_similarity(
+        actual.flatten(), reference.flatten(), dim=0
+    ).item()
+    return {
+        "valid": cosine > 0.999,
+        "cosine": cosine,
+        "max_abs_error": error.max().item(),
+        "max_scaled_error": scaled_error.max().item(),
+    }
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--kernel", type=Path, required=True)
+    parser.add_argument("--output", type=Path, required=True)
+    parser.add_argument(
+        "--mode", choices=("baseline", "sweep", "selected"), default="baseline"
+    )
+    parser.add_argument("--config", type=parse_config)
+    parser.add_argument("--m", type=int, action="append")
+    parser.add_argument("--config-shard", type=int, default=0)
+    parser.add_argument("--num-config-shards", type=int, default=1)
+    parser.add_argument("--repeats", type=int, default=21)
+    parser.add_argument("--replays", type=int, default=200)
+    parser.add_argument("--cache-multiplier", type=float, default=3.0)
+    parser.add_argument("--max-buffers", type=int, default=32)
+    parser.add_argument("--max-registers", type=int, default=64)
+    args = parser.parse_args()
+
+    token_counts = args.m or list(range(1, 17))
+    if any(not 1 <= m <= 16 for m in token_counts):
+        raise ValueError("expected 1 <= M <= 16")
+    if not 0 <= args.config_shard < args.num_config_shards:
+        raise ValueError("config shard must be in [0, num_config_shards)")
+    torch.cuda.set_device(0)
+    if torch.cuda.get_device_capability() != (10, 3):
+        raise RuntimeError("this benchmark requires SM103")
+
+    kernel_class = load_kernel_class(args.kernel)
+    properties = torch.cuda.get_device_properties(0)
+    metadata = {
+        "device": properties.name,
+        "compute_capability": list(torch.cuda.get_device_capability()),
+        "torch_version": torch.__version__,
+        "cuda_version": torch.version.cuda,
+    }
+    args.output.parent.mkdir(parents=True, exist_ok=True)
+    with args.output.open("w", encoding="utf-8") as output_file:
+        for m in token_counts:
+            configs = candidate_configs(args.mode, args.config, m)
+            torch.manual_seed(20260722 + m)
+            count = rotating_buffer_count(m, args.cache_multiplier, args.max_buffers)
+            activation = torch.randn((m, K), device="cuda", dtype=torch.bfloat16)
+            weights = [
+                torch.randn((N, K), device="cuda", dtype=torch.bfloat16)
+                for _ in range(count)
+            ]
+            residuals = [
+                torch.randn((m, N), device="cuda", dtype=torch.bfloat16)
+                for _ in range(count)
+            ]
+            candidates: list[tuple[str, Config | None]] = [("cublas_addmm", None)]
+            candidates.extend(
+                ("cute_residual", config)
+                for index, config in enumerate(configs)
+                if index % args.num_config_shards == args.config_shard
+            )
+            for backend, config in candidates:
+                row: dict[str, Any] = {
+                    "m": m,
+                    "n": N,
+                    "k": K,
+                    "backend": backend,
+                    "mode": args.mode,
+                    "config": dataclasses.asdict(config) if config else {},
+                    "num_buffers": count,
+                    "cache_multiplier": args.cache_multiplier,
+                    **metadata,
+                }
+                try:
+                    if backend == "cublas_addmm":
+                        launch = lambda a, b, residual, c: torch.addmm(
+                            residual, a, b.t(), out=c
+                        )
+                    else:
+                        if config is None:
+                            raise AssertionError("missing CuTe config")
+                        compiled = compile_kernel(
+                            kernel_class, m, config, args.max_registers
+                        )
+                        launch = lambda a, b, residual, c, fn=compiled: fn(
+                            a, b, residual, c, stream()
+                        )
+                        row.update(resource_usage(compiled))
+                    samples, outputs = graph_samples(
+                        launch,
+                        activation,
+                        weights,
+                        residuals,
+                        args.repeats,
+                        args.replays,
+                    )
+                    row.update(
+                        correctness(outputs[0], activation, weights[0], residuals[0])
+                    )
+                    row.update(summarize(samples))
+                except Exception as error:  # noqa: BLE001
+                    row.update(
+                        {
+                            "valid": False,
+                            "error": f"{type(error).__name__}: {error}",
+                        }
+                    )
+                output_file.write(json.dumps(row, sort_keys=True) + "\n")
+                output_file.flush()
+                print(json.dumps(row, sort_keys=True), flush=True)
+
+            del activation, weights, residuals
+            torch.cuda.empty_cache()
+
+
+if __name__ == "__main__":
+    main()
@@ -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()
@@ -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()
@@ -0,0 +1,267 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""End-to-end autoregressive decode benchmark: ReplaySSM vs the standard SSM kernel.
+
+Loads a hybrid Mamba2 model, replicates one prompt across the batch, and times a
+long greedy decode (CUDA graphs on) once with the standard kernel and once with
+ReplaySSM, then reports the per-step / throughput speedup. The two modes run in
+separate subprocesses so each gets a clean CUDA context.
+
+The FlashInfer FP4-MoE autotuner is disabled by default (it is unstable under
+CUDA-graph capture on the pre-release Blackwell FP4 path); pass
+--no-disable-flashinfer-autotune for non-FP4 models.
+
+Examples:
+    python e2e_decode_speedup.py --model-id nvidia/NVIDIA-Nemotron-3-Nano-4B-BF16
+    python e2e_decode_speedup.py --dtype auto --buffer-len 16 \
+        --model-id nvidia/NVIDIA-Nemotron-3-Super-120B-A12B-NVFP4   # B300 NVFP4
+"""
+
+import argparse
+import json
+import os
+import subprocess
+import sys
+import time
+
+DEFAULT_PROMPT = "My cat wrote all this CUDA code for a new language model and"
+
+MODE_LABEL = {"standard": "standard", "replayssm": "ReplaySSM"}
+
+
+def parse_args():
+    p = argparse.ArgumentParser(
+        description="E2E decode speedup: ReplaySSM vs the standard SSM kernel."
+    )
+    p.add_argument("--model-id", default="nvidia/NVIDIA-Nemotron-3-Nano-4B-BF16")
+    p.add_argument("--prompt", default=DEFAULT_PROMPT)
+    p.add_argument("--batch-size", type=int, default=256)
+    p.add_argument("--num-steps", type=int, default=1000)
+    p.add_argument("--warmup-steps", type=int, default=128)
+    p.add_argument("--repeats", type=int, default=1)
+    p.add_argument(
+        "--buffer-len", type=int, default=16, help="ReplaySSM input-buffer length."
+    )
+    p.add_argument(
+        "--dtype",
+        default="bfloat16",
+        choices=["bfloat16", "float16", "float32", "auto"],
+    )
+    p.add_argument("--gpu-memory-utilization", type=float, default=0.9)
+    p.add_argument("--max-model-len", type=int, default=None)
+    p.add_argument(
+        "--disable-flashinfer-autotune",
+        action=argparse.BooleanOptionalAction,
+        default=True,
+        help="Disable the FlashInfer FP4-MoE autotuner (default: on). "
+        "It is unstable under CUDA-graph capture on the "
+        "pre-release Blackwell FP4 path; pass "
+        "--no-disable-flashinfer-autotune for non-FP4 models.",
+    )
+    p.add_argument(
+        "--mamba-ssm-cache-dtype",
+        default="auto",
+        choices=["auto", "float32", "float16", "bfloat16"],
+        help="SSM state dtype (both modes). 'auto' = config-driven; "
+        "'float32' = fp32 state, 'bfloat16' = s16 state.",
+    )
+    p.add_argument(
+        "--baseline-ssm-config",
+        default="",
+        help="Pin the STANDARD baseline's SSM launch config as "
+        "'bsm,nw' via override_ssm_config (forces the in-process "
+        "engine so the override reaches the kernel). Empty = off.",
+    )
+    p.add_argument(
+        "--worker",
+        choices=["standard", "replayssm"],
+        default=None,
+        help=argparse.SUPPRESS,
+    )
+    return p.parse_args()
+
+
+def resolve_max_model_len(args) -> int:
+    if args.max_model_len is not None:
+        return args.max_model_len
+    return args.num_steps + 256
+
+
+def run_worker(args):
+    # override_ssm_config is a module global; it only reaches the model if the
+    # engine runs in-process (default V1 spawns a separate EngineCore). Force it.
+    if args.worker == "standard" and args.baseline_ssm_config:
+        os.environ["VLLM_ENABLE_V1_MULTIPROCESSING"] = "0"
+
+    import torch
+
+    from vllm import LLM, SamplingParams
+
+    mode = args.worker
+    max_model_len = resolve_max_model_len(args)
+
+    llm_kwargs = dict(
+        model=args.model_id,
+        tensor_parallel_size=1,
+        dtype=args.dtype,
+        max_model_len=max_model_len,
+        trust_remote_code=True,
+        enable_prefix_caching=False,
+        enable_chunked_prefill=False,
+        max_num_seqs=args.batch_size,
+        max_num_batched_tokens=max(max_model_len, args.batch_size * 64),
+        enforce_eager=False,
+        disable_log_stats=True,
+        gpu_memory_utilization=args.gpu_memory_utilization,
+        # SSM state dtype (applies to both standard and ReplaySSM).
+        mamba_ssm_cache_dtype=args.mamba_ssm_cache_dtype,
+    )
+    if args.disable_flashinfer_autotune:
+        # FP4-MoE autotuner is unstable under CUDA-graph capture on Blackwell;
+        # re-enable (--no-disable-flashinfer-autotune) only for non-FP4 models.
+        llm_kwargs["kernel_config"] = {"enable_flashinfer_autotune": False}
+    if mode == "replayssm":
+        llm_kwargs.update(use_replayssm=True, replayssm_buffer_len=args.buffer_len)
+
+    _ssm_cm = None
+    if mode == "standard" and args.baseline_ssm_config:
+        from vllm.model_executor.layers.mamba.ops.mamba_ssm import override_ssm_config
+
+        _bsm, _nw = (int(x) for x in args.baseline_ssm_config.split(","))
+        _ssm_cm = override_ssm_config((_bsm, _nw))
+        _ssm_cm.__enter__()  # active through LLM() graph capture + decode
+        print(
+            f"[{mode}] override_ssm_config -> (BLOCK_SIZE_M={_bsm}, num_warps={_nw})",
+            flush=True,
+        )
+
+    llm = LLM(**llm_kwargs)
+    prompts = [args.prompt] * args.batch_size
+
+    def timed_generate(n_tokens):
+        sp = SamplingParams(
+            n=1,
+            temperature=0.0,
+            ignore_eos=True,
+            min_tokens=n_tokens,
+            max_tokens=n_tokens,
+        )
+        if torch.accelerator.is_available():
+            torch.accelerator.synchronize()
+        t0 = time.perf_counter()
+        outs = llm.generate(prompts, sp, use_tqdm=False)
+        if torch.accelerator.is_available():
+            torch.accelerator.synchronize()
+        elapsed = time.perf_counter() - t0
+        produced = min(len(o.outputs[0].token_ids) for o in outs)
+        assert produced == n_tokens, f"expected {n_tokens} tokens, got {produced}"
+        return elapsed
+
+    timed_generate(args.warmup_steps)
+
+    best = None
+    for _ in range(args.repeats):
+        elapsed = timed_generate(args.num_steps)
+        tok_s = args.batch_size * args.num_steps / elapsed
+        per_step_ms = elapsed / args.num_steps * 1e3
+        print(
+            f"[{mode}] {elapsed:.3f}s  {tok_s:,.0f} tok/s  {per_step_ms:.3f} ms/step",
+            flush=True,
+        )
+        if best is None or elapsed < best["elapsed_s"]:
+            best = {
+                "mode": mode,
+                "elapsed_s": elapsed,
+                "tok_s": tok_s,
+                "per_step_ms": per_step_ms,
+            }
+
+    print("RESULT_JSON " + json.dumps(best), flush=True)
+    if _ssm_cm is not None:
+        _ssm_cm.__exit__(None, None, None)
+
+
+def run_one_mode(args, mode) -> dict:
+    cmd = [
+        sys.executable,
+        __file__,
+        "--worker",
+        mode,
+        "--model-id",
+        args.model_id,
+        "--prompt",
+        args.prompt,
+        "--batch-size",
+        str(args.batch_size),
+        "--num-steps",
+        str(args.num_steps),
+        "--warmup-steps",
+        str(args.warmup_steps),
+        "--repeats",
+        str(args.repeats),
+        "--buffer-len",
+        str(args.buffer_len),
+        "--dtype",
+        args.dtype,
+        "--gpu-memory-utilization",
+        str(args.gpu_memory_utilization),
+        "--mamba-ssm-cache-dtype",
+        args.mamba_ssm_cache_dtype,
+        "--baseline-ssm-config",
+        args.baseline_ssm_config,
+    ]
+    cmd.append(
+        "--disable-flashinfer-autotune"
+        if args.disable_flashinfer_autotune
+        else "--no-disable-flashinfer-autotune"
+    )
+    if args.max_model_len is not None:
+        cmd += ["--max-model-len", str(args.max_model_len)]
+
+    result = None
+    proc = subprocess.Popen(
+        cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True, bufsize=1
+    )
+    for line in proc.stdout:
+        sys.stdout.write(line)
+        sys.stdout.flush()
+        if line.startswith("RESULT_JSON "):
+            result = json.loads(line[len("RESULT_JSON ") :])
+    proc.wait()
+    if proc.returncode != 0:
+        raise RuntimeError(f"mode '{mode}' worker exited with {proc.returncode}")
+    if result is None:
+        raise RuntimeError(f"mode '{mode}' produced no RESULT_JSON line")
+    return result
+
+
+def main():
+    args = parse_args()
+    if args.worker is not None:
+        run_worker(args)
+        return
+
+    print(
+        f"model={args.model_id}  batch_size={args.batch_size}  "
+        f"steps={args.num_steps}  buffer_len={args.buffer_len}  dtype={args.dtype}"
+    )
+
+    std = run_one_mode(args, "standard")
+    fla = run_one_mode(args, "replayssm")
+    speedup = std["per_step_ms"] / fla["per_step_ms"]
+
+    print()
+    header = f"{'mode':<10}{'ms/step':>12}{'tok/s':>16}{'wall (s)':>12}"
+    print(header)
+    print("-" * len(header))
+    for r in (std, fla):
+        print(
+            f"{MODE_LABEL[r['mode']]:<10}{r['per_step_ms']:>12.3f}"
+            f"{r['tok_s']:>16,.0f}{r['elapsed_s']:>12.3f}"
+        )
+    print("-" * len(header))
+    print(f"speedup (standard / ReplaySSM, per step): {speedup:.3f}x")
+
+
+if __name__ == "__main__":
+    main()
@@ -28,9 +28,9 @@ if(DEEPGEMM_SRC_DIR)
  message(STATUS "DeepGEMM using local DEEPGEMM_SRC_DIR: ${deepgemm_SOURCE_DIR}")
 else()
  # Keep in sync with tools/install_deepgemm.sh
-  set(_DEEPGEMM_UPSTREAM_REPO "https://github.com/deepseek-ai/DeepGEMM.git")
+  set(_DEEPGEMM_UPSTREAM_REPO "git@github.com:Inferact/DeepGEMM.git")
  # NOTE: This is currently targeting nv-dev branch due to sm120 support
-  set(_DEEPGEMM_UPSTREAM_TAG "a6b593d2826719dcf4892609af7b84ee23aaf32a")
+  set(_DEEPGEMM_UPSTREAM_TAG "f5a76426fa084087169693fd0cd815223576d6e9")

  set(_deepgemm_fc_root "${FETCHCONTENT_BASE_DIR}")
  if(NOT _deepgemm_fc_root)
@@ -68,6 +68,9 @@ endif()
 if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8)
  if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.9)
    list(APPEND DEEPGEMM_SUPPORT_ARCHS "10.0f")
+    if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.4)
+      list(APPEND DEEPGEMM_SUPPORT_ARCHS "10.7f")
+    endif()
  else()
    list(APPEND DEEPGEMM_SUPPORT_ARCHS "10.0a")
  endif()
@@ -0,0 +1,74 @@
+include(FetchContent)
+
+if(DEFINED ENV{FLASH_KDA_SRC_DIR})
+  set(FLASH_KDA_SRC_DIR $ENV{FLASH_KDA_SRC_DIR})
+endif()
+
+if(FLASH_KDA_SRC_DIR)
+  FetchContent_Declare(
+    flashkda
+    SOURCE_DIR ${FLASH_KDA_SRC_DIR}
+  )
+else()
+  FetchContent_Declare(
+    flashkda
+    GIT_REPOSITORY git@github.com:Inferact/FlashKDA.git
+    GIT_TAG a3e42bbbece3bb38f7c426b880315294a336e82f
+    GIT_PROGRESS TRUE
+    GIT_SUBMODULES cutlass
+  )
+endif()
+
+FetchContent_MakeAvailable(flashkda)
+message(STATUS "FlashKDA is available at ${flashkda_SOURCE_DIR}")
+
+set(FLASH_KDA_SUPPORT_ARCHS)
+if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.0)
+  list(APPEND FLASH_KDA_SUPPORT_ARCHS "9.0a")
+endif()
+if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
+  list(APPEND FLASH_KDA_SUPPORT_ARCHS "10.0f" "12.0f")
+elseif(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.9)
+  list(APPEND FLASH_KDA_SUPPORT_ARCHS "10.0a" "10.3a" "12.0a")
+endif()
+
+cuda_archs_loose_intersection(
+  FLASH_KDA_ARCHS "${FLASH_KDA_SUPPORT_ARCHS}" "${CUDA_ARCHS}")
+
+if(FLASH_KDA_ARCHS)
+  message(STATUS "FlashKDA CUDA architectures: ${FLASH_KDA_ARCHS}")
+
+  set(FLASH_KDA_SOURCES
+    csrc/flashkda_registration.cpp
+    ${flashkda_SOURCE_DIR}/csrc/flash_kda.cpp
+    ${flashkda_SOURCE_DIR}/csrc/smxx/fwd_launch.cu)
+  set(FLASH_KDA_INCLUDES
+    ${flashkda_SOURCE_DIR}/csrc
+    ${flashkda_SOURCE_DIR}/cutlass/include
+    ${flashkda_SOURCE_DIR}/cutlass/examples/common
+    ${flashkda_SOURCE_DIR}/cutlass/tools/util/include)
+
+  set_gencode_flags_for_srcs(
+    SRCS "${FLASH_KDA_SOURCES}"
+    CUDA_ARCHS "${FLASH_KDA_ARCHS}")
+
+  define_extension_target(
+    _flashkda_C
+    DESTINATION vllm
+    LANGUAGE ${VLLM_GPU_LANG}
+    SOURCES ${FLASH_KDA_SOURCES}
+    COMPILE_FLAGS ${VLLM_GPU_FLAGS}
+    ARCHITECTURES ${VLLM_GPU_ARCHES}
+    INCLUDE_DIRECTORIES ${FLASH_KDA_INCLUDES}
+    USE_SABI 3
+    WITH_SOABI)
+
+  target_compile_options(_flashkda_C PRIVATE
+    $<$<COMPILE_LANGUAGE:CUDA>:-UPy_LIMITED_API --expt-relaxed-constexpr --expt-extended-lambda --use_fast_math -O3>
+    $<$<COMPILE_LANGUAGE:CXX>:-UPy_LIMITED_API>)
+else()
+  message(STATUS
+    "FlashKDA will not compile: CUDA >=12.0 and a supported architecture "
+    "(SM90, SM10x, or SM12x) are required")
+  add_custom_target(_flashkda_C)
+endif()
@@ -60,6 +60,9 @@ if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.9)
    # CUDA 12.9 has introduced "Family-Specific Architecture Features"
    # this supports all compute_10x family
    list(APPEND SUPPORT_ARCHS "10.0f")
+    if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.4)
+        list(APPEND SUPPORT_ARCHS "10.7f")
+    endif()
 elseif(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8)
    list(APPEND SUPPORT_ARCHS "10.0a")
 endif()
@@ -188,4 +191,3 @@ else()
    add_custom_target(_flashmla_C)
    add_custom_target(_flashmla_extension_C)
 endif()
-
@@ -55,7 +55,11 @@ message(STATUS "[QUTLASS] QuTLASS is available at ${qutlass_SOURCE_DIR}")

 if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
  cuda_archs_loose_intersection(QUTLASS_SM120_ARCHS "12.0f" "${CUDA_ARCHS}")
-  cuda_archs_loose_intersection(QUTLASS_SM100_ARCHS "10.0f" "${CUDA_ARCHS}")
+  if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.4)
+    cuda_archs_loose_intersection(QUTLASS_SM100_ARCHS "10.0f;10.7f" "${CUDA_ARCHS}")
+  else()
+    cuda_archs_loose_intersection(QUTLASS_SM100_ARCHS "10.0f" "${CUDA_ARCHS}")
+  endif()
 else()
  cuda_archs_loose_intersection(QUTLASS_SM120_ARCHS "12.0a;12.1a" "${CUDA_ARCHS}")
  cuda_archs_loose_intersection(QUTLASS_SM100_ARCHS "10.0a;10.3a" "${CUDA_ARCHS}")
@@ -396,14 +396,24 @@ function(cuda_archs_loose_intersection OUT_CUDA_ARCHS SRC_CUDA_ARCHS TGT_CUDA_AR
  # match — e.g. SRC="12.0f" matches TGT="12.1a" since SM121 is in the SM12x
  # family. The output uses TGT's value to preserve the user's compilation flags.
  set(_CUDA_ARCHS)
+  # Resolve exact base matches before family fallbacks so a generic entry such
+  # as 10.0f cannot consume a 10.7 target that has a 10.7f source entry.
+  foreach(_arch ${_SRC_CUDA_ARCHS})
+    if(_arch MATCHES "[af]$")
+      string(REGEX REPLACE "[af]$" "" _base "${_arch}")
+      if("${_base}" IN_LIST _TGT_CUDA_ARCHS)
+        list(REMOVE_ITEM _SRC_CUDA_ARCHS "${_arch}")
+        list(REMOVE_ITEM _TGT_CUDA_ARCHS "${_base}")
+        list(APPEND _CUDA_ARCHS "${_arch}")
+      endif()
+    endif()
+  endforeach()
+
  foreach(_arch ${_SRC_CUDA_ARCHS})
    if(_arch MATCHES "[af]$")
      list(REMOVE_ITEM _SRC_CUDA_ARCHS "${_arch}")
      string(REGEX REPLACE "[af]$" "" _base "${_arch}")
-      if ("${_base}" IN_LIST TGT_CUDA_ARCHS)
-        list(REMOVE_ITEM _TGT_CUDA_ARCHS "${_base}")
-        list(APPEND _CUDA_ARCHS "${_arch}")
-      elseif("${_base}a" IN_LIST _TGT_CUDA_ARCHS)
+      if("${_base}a" IN_LIST _TGT_CUDA_ARCHS)
        list(REMOVE_ITEM _TGT_CUDA_ARCHS "${_base}a")
        list(APPEND _CUDA_ARCHS "${_base}a")
      elseif("${_base}f" IN_LIST _TGT_CUDA_ARCHS)
@@ -487,7 +497,7 @@ endfunction()

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

-template <typename T, int kHdIn, int kHdOut, int kTileN>
+template <typename T, int kHdIn, int kHdOut, int kTileN, int kTileK = 256>
 void invokeFusedAGemm(T* output, T const* mat_a, T const* mat_b, int num_tokens,
-                      cudaStream_t const stream) {
-  constexpr int gemm_m = kHdOut;  // 2112
-  int const gemm_n = num_tokens;  // 1-16
-  constexpr int gemm_k = kHdIn;   // 7168
+                      cudaStream_t const stream, bool enable_pdl) {
+  constexpr int gemm_m = kHdOut;
+  int const gemm_n = num_tokens;
+  constexpr int gemm_k = kHdIn;
  constexpr int batch_size = 1;
  std::swap(mat_a, mat_b);
  constexpr int tile_m = 16;
-  constexpr int tile_n = kTileN;                        // 8 or 16
-  constexpr int tile_k = std::max(256, 1024 / tile_n);  // 256
+  constexpr int tile_n = kTileN;
+  constexpr int tile_k = kTileK;
  constexpr int max_stage_cnt =
      1024 * 192 / ((tile_m + tile_n) * tile_k * sizeof(bf16_t));
  constexpr int k_iter_cnt = gemm_k / tile_k;
@@ -679,7 +679,8 @@ void invokeFusedAGemm(T* output, T const* mat_a, T const* mat_b, int num_tokens,
  config.stream = stream;
  cudaLaunchAttribute attrs[1];
  attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization;
-  attrs[0].val.programmaticStreamSerializationAllowed = getEnvEnablePDL();
+  attrs[0].val.programmaticStreamSerializationAllowed =
+      enable_pdl || getEnvEnablePDL();
  config.numAttrs = 1;
  config.attrs = attrs;
  if (smem_bytes >= (48 * 1024)) {
@@ -694,36 +695,48 @@ void invokeFusedAGemm(T* output, T const* mat_a, T const* mat_b, int num_tokens,
                     output, mat_a, mat_b, gemm_n);
 }

-template void invokeFusedAGemm<__nv_bfloat16, 7168, 2112, 8>(
-    __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, int num_tokens,
-    cudaStream_t);
-
-template void invokeFusedAGemm<__nv_bfloat16, 7168, 2112, 16>(
-    __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, int num_tokens,
-    cudaStream_t);
+template <typename T, int kHdIn, int kHdOut, int kTileK = 256>
+void invokeFusedAGemmForTokens(T* output, T const* mat_a, T const* mat_b,
+                              int num_tokens, cudaStream_t const stream,
+                              bool enable_pdl) {
+  if (num_tokens <= 8) {
+    invokeFusedAGemm<T, kHdIn, kHdOut, 8, kTileK>(
+        output, mat_a, mat_b, num_tokens, stream, enable_pdl);
+  } else {
+    invokeFusedAGemm<T, kHdIn, kHdOut, 16, kTileK>(
+        output, mat_a, mat_b, num_tokens, stream, enable_pdl);
+  }
+}

 void dsv3_fused_a_gemm(torch::stable::Tensor& output,
                       torch::stable::Tensor const& mat_a,
-                       torch::stable::Tensor const& mat_b) {
+                       torch::stable::Tensor const& mat_b, bool enable_pdl) {
  STD_TORCH_CHECK(mat_a.dim() == 2 && mat_b.dim() == 2 && output.dim() == 2);
  int const num_tokens = mat_a.size(0);
  int const hd_in = mat_a.size(1);
  int const hd_out = mat_b.size(1);

-  constexpr int kHdIn = 7168;
-  constexpr int kHdOut = 2112;
  STD_TORCH_CHECK(num_tokens >= 1 && num_tokens <= 16,
                  "required 1 <= mat_a.shape[0] <= 16");
-  STD_TORCH_CHECK(hd_in == kHdIn, "required mat_a.shape[1] == 7168");
-  STD_TORCH_CHECK(hd_out == kHdOut, "required mat_b.shape[1] == 2112");
  STD_TORCH_CHECK(output.size(0) == num_tokens,
                  "required output.shape[0] == mat_a.shape[0]");
  STD_TORCH_CHECK(output.size(1) == hd_out,
                  "required output.shape[1] == mat_b.shape[1]");
+  STD_TORCH_CHECK(mat_b.size(0) == hd_in,
+                  "required mat_b.shape[0] == mat_a.shape[1]");

-  STD_TORCH_CHECK(mat_a.stride(1) == 1, "mat_a must be a row major tensor");
-  STD_TORCH_CHECK(output.stride(1) == 1, "output must be a row major tensor");
-  STD_TORCH_CHECK(mat_b.stride(0) == 1, "mat_b must be a column major tensor");
+  STD_TORCH_CHECK(mat_a.get_device_index() == mat_b.get_device_index() &&
+                      mat_a.get_device_index() == output.get_device_index(),
+                  "mat_a, mat_b, and output must be on the same device");
+
+  // The kernels index global memory with raw pointers and packed strides, so
+  // reject any padded or transposed view rather than reading out of bounds.
+  STD_TORCH_CHECK(mat_a.stride(0) == hd_in && mat_a.stride(1) == 1,
+                  "mat_a must be a packed row-major [num_tokens, hd_in] tensor");
+  STD_TORCH_CHECK(output.stride(0) == hd_out && output.stride(1) == 1,
+                  "output must be a packed row-major [num_tokens, hd_out] tensor");
+  STD_TORCH_CHECK(mat_b.stride(0) == 1 && mat_b.stride(1) == hd_in,
+                  "mat_b must be a packed column-major [hd_in, hd_out] tensor");

  STD_TORCH_CHECK(
      mat_a.scalar_type() == torch::headeronly::ScalarType::BFloat16 &&
@@ -738,19 +751,86 @@ void dsv3_fused_a_gemm(torch::stable::Tensor& output,
  STD_TORCH_CHECK(getSMVersion() >= 90, "required CUDA ARCH >= SM_90");

  auto stream = get_current_cuda_stream(mat_a.get_device_index());
-  if (num_tokens <= 8) {
-    invokeFusedAGemm<__nv_bfloat16, kHdIn, kHdOut, 8>(
-        reinterpret_cast<__nv_bfloat16*>(output.mutable_data_ptr()),
-        reinterpret_cast<__nv_bfloat16 const*>(mat_a.data_ptr()),
-        reinterpret_cast<__nv_bfloat16 const*>(mat_b.data_ptr()), num_tokens,
-        stream);
-  } else {
-    invokeFusedAGemm<__nv_bfloat16, kHdIn, kHdOut, 16>(
-        reinterpret_cast<__nv_bfloat16*>(output.mutable_data_ptr()),
-        reinterpret_cast<__nv_bfloat16 const*>(mat_a.data_ptr()),
-        reinterpret_cast<__nv_bfloat16 const*>(mat_b.data_ptr()), num_tokens,
-        stream);
+  auto* output_ptr =
+      reinterpret_cast<__nv_bfloat16*>(output.mutable_data_ptr());
+  auto const* mat_a_ptr =
+      reinterpret_cast<__nv_bfloat16 const*>(mat_a.data_ptr());
+  auto const* mat_b_ptr =
+      reinterpret_cast<__nv_bfloat16 const*>(mat_b.data_ptr());
+
+#define DISPATCH_DSV3_SHAPE(HD_IN, HD_OUT)                       \
+  if (hd_in == HD_IN && hd_out == HD_OUT) {                      \
+    invokeFusedAGemmForTokens<__nv_bfloat16, HD_IN, HD_OUT>(     \
+        output_ptr, mat_a_ptr, mat_b_ptr, num_tokens, stream,    \
+        enable_pdl);                                             \
+    return;                                                      \
  }
+
+  // Shapes the Kimi-K3 selector routes to dsv3_fused_a (see the dsv3 winners
+  // in KIMI_K3_PROJECTIONS) plus the DeepSeek V2/V3 QKV A-projection.
+  DISPATCH_DSV3_SHAPE(7168, 1536)
+  DISPATCH_DSV3_SHAPE(7168, 2112)
+  DISPATCH_DSV3_SHAPE(1536, 2304)
+  DISPATCH_DSV3_SHAPE(1536, 4608)
+  DISPATCH_DSV3_SHAPE(7168, 3584)
+  DISPATCH_DSV3_SHAPE(768, 7168)
+  // TP16 dsv3 winners, as (hd_in=K, hd_out=N). TP16 dense down_proj is absent
+  // because hd_in=2112 is not a multiple of any supported tile_k.
+  DISPATCH_DSV3_SHAPE(1536, 1152)
+  DISPATCH_DSV3_SHAPE(7168, 768)
+  DISPATCH_DSV3_SHAPE(7168, 3216)
+  DISPATCH_DSV3_SHAPE(7168, 4224)
+
+#ifdef VLLM_K3_BENCH_SHAPES
+  // The selector routes these shapes to CuTe or the default GEMM, so they are
+  // never reached in production. They are compiled only for offline
+  // DSV3-vs-CuTe benchmarking.
+  DISPATCH_DSV3_SHAPE(7168, 6288)
+  DISPATCH_DSV3_SHAPE(1536, 7168)
+  DISPATCH_DSV3_SHAPE(3584, 7168)
+  DISPATCH_DSV3_SHAPE(7168, 8448)
+  DISPATCH_DSV3_SHAPE(7168, 20480)
+  DISPATCH_DSV3_SHAPE(7168, 3072)
+  DISPATCH_DSV3_SHAPE(7168, 12448)
+  DISPATCH_DSV3_SHAPE(3072, 7168)
+  DISPATCH_DSV3_SHAPE(8448, 7168)
+  DISPATCH_DSV3_SHAPE(7168, 16896)
+  DISPATCH_DSV3_SHAPE(7168, 40960)
+#endif
+
+#undef DISPATCH_DSV3_SHAPE
+
+  if (hd_in == 128 && hd_out == 1536) {
+    invokeFusedAGemmForTokens<__nv_bfloat16, 128, 1536, 128>(
+        output_ptr, mat_a_ptr, mat_b_ptr, num_tokens, stream, enable_pdl);
+    return;
+  }
+  if (hd_in == 128 && hd_out == 3072) {
+    invokeFusedAGemmForTokens<__nv_bfloat16, 128, 3072, 128>(
+        output_ptr, mat_a_ptr, mat_b_ptr, num_tokens, stream, enable_pdl);
+    return;
+  }
+  // TP16 KDA f_b_proj and shared_expert down_proj. Neither hd_in is a multiple
+  // of 256, so both need the 128 tile_k.
+  if (hd_in == 128 && hd_out == 768) {
+    invokeFusedAGemmForTokens<__nv_bfloat16, 128, 768, 128>(
+        output_ptr, mat_a_ptr, mat_b_ptr, num_tokens, stream, enable_pdl);
+    return;
+  }
+  if (hd_in == 384 && hd_out == 7168) {
+    invokeFusedAGemmForTokens<__nv_bfloat16, 384, 7168, 128>(
+        output_ptr, mat_a_ptr, mat_b_ptr, num_tokens, stream, enable_pdl);
+    return;
+  }
+#ifdef VLLM_K3_BENCH_SHAPES
+  if (hd_in == 4224 && hd_out == 7168) {
+    invokeFusedAGemmForTokens<__nv_bfloat16, 4224, 7168, 128>(
+        output_ptr, mat_a_ptr, mat_b_ptr, num_tokens, stream, enable_pdl);
+    return;
+  }
+#endif
+
+  STD_TORCH_CHECK(false, "unsupported DSV3 fused-A GEMM shape");
 }

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

+void fused_kimi_k3_mla_key_concat_kv_cache_insert(
+    torch::stable::Tensor& q, torch::stable::Tensor const& k_nope,
+    torch::stable::Tensor const& k_pe, torch::stable::Tensor const& kv_c_normed,
+    torch::stable::Tensor& k_out, torch::stable::Tensor& k_cache,
+    torch::stable::Tensor const& slot_mapping, int64_t cache_block_size,
+    std::optional<torch::stable::Tensor> position_ids,
+    std::optional<torch::stable::Tensor> cos_sin_cache);
+
+void fused_kimi_k3_mla_key_concat_ds_mla_insert(
+    torch::stable::Tensor& q, torch::stable::Tensor const& k_nope,
+    torch::stable::Tensor const& k_pe, torch::stable::Tensor const& kv_c_normed,
+    torch::stable::Tensor& k_out, torch::stable::Tensor& k_cache,
+    torch::stable::Tensor const& slot_mapping, int64_t cache_block_size,
+    std::optional<torch::stable::Tensor> position_ids,
+    std::optional<torch::stable::Tensor> cos_sin_cache);
+
+void fused_kimi_k3_mla_qkv_quant_kv_cache_fp8_insert(
+    torch::stable::Tensor const& q, torch::stable::Tensor const& k_nope,
+    torch::stable::Tensor const& k_pe, torch::stable::Tensor const& kv_c_normed,
+    torch::stable::Tensor const& v, torch::stable::Tensor& q_fp8,
+    torch::stable::Tensor& k_fp8, torch::stable::Tensor& v_fp8,
+    torch::stable::Tensor& k_cache, torch::stable::Tensor const& slot_mapping,
+    torch::stable::Tensor const& q_scale_inv,
+    torch::stable::Tensor const& k_scale_inv,
+    torch::stable::Tensor const& v_scale_inv,
+    torch::stable::Tensor const& cache_scale_inv, int64_t cache_block_size,
+    std::optional<torch::stable::Tensor> position_ids,
+    std::optional<torch::stable::Tensor> cos_sin_cache);
+
+void fused_kimi_k3_mla_decode_q_concat_kv_cache_insert(
+    torch::stable::Tensor const& ql_nope, torch::stable::Tensor const& q_pe,
+    torch::stable::Tensor const& kv_c_normed, torch::stable::Tensor const& k_pe,
+    torch::stable::Tensor& mqa_q, torch::stable::Tensor& k_cache,
+    torch::stable::Tensor const& slot_mapping, int64_t cache_block_size,
+    std::optional<torch::stable::Tensor> position_ids,
+    std::optional<torch::stable::Tensor> cos_sin_cache);
+
+void fused_kimi_k3_mla_decode_q_concat_kv_cache_fp8_insert(
+    torch::stable::Tensor const& ql_nope, torch::stable::Tensor const& q_pe,
+    torch::stable::Tensor const& kv_c_normed, torch::stable::Tensor const& k_pe,
+    torch::stable::Tensor& mqa_q, torch::stable::Tensor& k_cache,
+    torch::stable::Tensor const& slot_mapping,
+    torch::stable::Tensor const& q_scale_inv,
+    torch::stable::Tensor const& cache_scale_inv, int64_t cache_block_size,
+    std::optional<torch::stable::Tensor> position_ids,
+    std::optional<torch::stable::Tensor> cos_sin_cache);
+
+void fused_kimi_k3_mla_decode_q_concat_ds_mla_insert(
+    torch::stable::Tensor const& ql_nope, torch::stable::Tensor const& q_pe,
+    torch::stable::Tensor const& kv_c_normed, torch::stable::Tensor const& k_pe,
+    torch::stable::Tensor& mqa_q, torch::stable::Tensor& k_cache,
+    torch::stable::Tensor const& slot_mapping, int64_t cache_block_size,
+    std::optional<torch::stable::Tensor> position_ids,
+    std::optional<torch::stable::Tensor> cos_sin_cache);
+
 void fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_fp8_insert(
    torch::stable::Tensor const& q, torch::stable::Tensor const& kv,
    torch::stable::Tensor& q_fp8, torch::stable::Tensor& k_cache,
@@ -315,6 +370,30 @@ void fused_minimax_m3_qknorm_rope_kv_insert(
    std::optional<torch::stable::Tensor> index_q_out,
    const std::string& kv_cache_dtype, bool skip_index_branch);

+#ifdef VLLM_ENABLE_FUSED_KDA_DECODE
+void fused_kda_decode(
+    torch::stable::Tensor const& x, torch::stable::Tensor const& weight,
+    std::optional<torch::stable::Tensor> bias,
+    torch::stable::Tensor& conv_state, torch::stable::Tensor const& raw_g,
+    torch::stable::Tensor const& raw_beta, torch::stable::Tensor const& a_log,
+    torch::stable::Tensor const& dt_bias,
+    torch::stable::Tensor const& state_indices, torch::stable::Tensor& state,
+    torch::stable::Tensor& out, std::optional<double> lower_bound,
+    std::optional<torch::stable::Tensor> output_gate,
+    std::optional<torch::stable::Tensor> norm_weight, double norm_eps);
+#endif
+
+#ifdef VLLM_ENABLE_KIMI_K3_ATTN_RES
+void kimi_k3_attn_res(torch::stable::Tensor& prefix,
+                      torch::stable::Tensor const& delta,
+                      torch::stable::Tensor const& blocks,
+                      torch::stable::Tensor const& norm_weight,
+                      torch::stable::Tensor const& qk_weight,
+                      torch::stable::Tensor const& output_norm_weight,
+                      torch::stable::Tensor& output, int64_t num_blocks,
+                      double eps, double output_norm_eps);
+#endif
+
 // Sampler kernels (shared CUDA/ROCm)
 void apply_repetition_penalties_(
    torch::stable::Tensor& logits, const torch::stable::Tensor& prompt_mask,
@@ -372,6 +451,20 @@ fptr_t init_custom_ar(const std::vector<int64_t>& fake_ipc_ptrs,
 void all_reduce(fptr_t _fa, torch::stable::Tensor& inp,
                torch::stable::Tensor& out, fptr_t reg_buffer,
                int64_t reg_buffer_sz_bytes);
+void custom_all_gather(fptr_t _fa, torch::stable::Tensor& inp,
+                       torch::stable::Tensor& out, fptr_t reg_buffer,
+                       int64_t reg_buffer_sz_bytes);
+void mnnvl_lamport_all_gather(fptr_t _fa, torch::stable::Tensor& inp,
+                              torch::stable::Tensor& out, fptr_t local_buffer,
+                              fptr_t multicast_buffer, fptr_t epoch_buffer,
+                              int64_t stage_sz_bytes);
+void custom_reduce_scatter(fptr_t _fa, torch::stable::Tensor& inp,
+                           torch::stable::Tensor& out, fptr_t reg_buffer,
+                           int64_t reg_buffer_sz_bytes);
+void mnnvl_lamport_reduce_scatter(fptr_t _fa, torch::stable::Tensor& inp,
+                                  torch::stable::Tensor& out,
+                                  fptr_t local_buffer, fptr_t epoch_buffer,
+                                  int64_t stage_sz_bytes);
 void dispose(fptr_t _fa);
 int64_t meta_size();
 void register_buffer(fptr_t _fa, const std::vector<int64_t>& fake_ipc_ptrs);
@@ -410,6 +503,12 @@ void fatrelu_and_mul(torch::stable::Tensor& out, torch::stable::Tensor& input,
                     double threshold);
 void swigluoai_and_mul(torch::stable::Tensor& out, torch::stable::Tensor& input,
                       double alpha = 1.702, double limit = 7.0);
+void situ_and_mul(torch::stable::Tensor& out, torch::stable::Tensor& input,
+                  double beta = 1.0, double linear_beta = -1.0);
+void masked_situ_and_mul(torch::stable::Tensor& out,
+                         torch::stable::Tensor& input,
+                         const torch::stable::Tensor& expert_num_tokens,
+                         double beta = 1.0, double linear_beta = -1.0);
 void gelu_new(torch::stable::Tensor& out, torch::stable::Tensor& input);
 void gelu_fast(torch::stable::Tensor& out, torch::stable::Tensor& input);
 void gelu_quick(torch::stable::Tensor& out, torch::stable::Tensor& input);
@@ -486,6 +585,13 @@ void concat_and_cache_mla(torch::stable::Tensor& kv_c,
                          const std::string& kv_cache_dtype,
                          torch::stable::Tensor& scale);

+void concat_and_cache_mla_grouped(torch::stable::Tensor& kv_c,
+                                  torch::stable::Tensor& k_pe,
+                                  torch::stable::Tensor& kv_cache_ptrs,
+                                  torch::stable::Tensor& slot_mapping,
+                                  int64_t block_size, int64_t block_stride,
+                                  int64_t entry_stride);
+
 // NOTE: k_pe and kv_c order is flipped compared to concat_and_cache_mla
 void concat_and_cache_mla_rope_fused(
    torch::stable::Tensor& positions, torch::stable::Tensor& q_pe,
@@ -324,7 +324,8 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
  // DeepSeek V3 fused A GEMM (SM 9.0+, bf16 only, 1-16 tokens).
  // conditionally compiled so impl registration is in source file
  ops.def(
-      "dsv3_fused_a_gemm(Tensor! output, Tensor mat_a, Tensor mat_b) -> ()");
+      "dsv3_fused_a_gemm(Tensor! output, Tensor mat_a, Tensor mat_b, "
+      "bool enable_pdl=False) -> ()");

  // BF16/FP32 x FP32 -> FP32 router GEMM for H=3072, E=256, M<=32 (SM90+).
  // conditionally compiled so impl registration is in source file
@@ -447,6 +448,48 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
      "Tensor fp8_scale, Tensor q_fp8_scale_inv, float eps, "
      "int cache_block_size) -> ()");

+  // Kimi-K3 MLA epilogues: optional RoPE followed by concat/cache insertion.
+  ops.def(
+      "fused_kimi_k3_mla_key_concat_kv_cache_insert("
+      "Tensor! q, Tensor k_nope, Tensor k_pe, Tensor kv_c_normed, "
+      "Tensor! k_out, Tensor! k_cache, Tensor slot_mapping, "
+      "int cache_block_size, Tensor? position_ids=None, "
+      "Tensor? cos_sin_cache=None) -> ()");
+  ops.def(
+      "fused_kimi_k3_mla_key_concat_ds_mla_insert("
+      "Tensor! q, Tensor k_nope, Tensor k_pe, Tensor kv_c_normed, "
+      "Tensor! k_out, Tensor! k_cache, Tensor slot_mapping, "
+      "int cache_block_size, Tensor? position_ids=None, "
+      "Tensor? cos_sin_cache=None) -> ()");
+  ops.def(
+      "fused_kimi_k3_mla_qkv_quant_kv_cache_fp8_insert("
+      "Tensor q, Tensor k_nope, Tensor k_pe, Tensor kv_c_normed, Tensor v, "
+      "Tensor! q_fp8, Tensor! k_fp8, Tensor! v_fp8, Tensor! k_cache, "
+      "Tensor slot_mapping, Tensor q_scale_inv, Tensor k_scale_inv, "
+      "Tensor v_scale_inv, Tensor cache_scale_inv, int cache_block_size, "
+      "Tensor? position_ids=None, Tensor? cos_sin_cache=None) -> ()");
+
+  // Kimi-K3 MLA decode epilogue: concat mqa_q = [ql_nope | q_pe] and insert the
+  // latent [kv_c_normed | k_pe] into the paged cache (bf16 / fp8 / fp8_ds_mla).
+  ops.def(
+      "fused_kimi_k3_mla_decode_q_concat_kv_cache_insert("
+      "Tensor ql_nope, Tensor q_pe, Tensor kv_c_normed, Tensor k_pe, "
+      "Tensor! mqa_q, Tensor! k_cache, Tensor slot_mapping, "
+      "int cache_block_size, Tensor? position_ids=None, "
+      "Tensor? cos_sin_cache=None) -> ()");
+  ops.def(
+      "fused_kimi_k3_mla_decode_q_concat_kv_cache_fp8_insert("
+      "Tensor ql_nope, Tensor q_pe, Tensor kv_c_normed, Tensor k_pe, "
+      "Tensor! mqa_q, Tensor! k_cache, Tensor slot_mapping, "
+      "Tensor q_scale_inv, Tensor cache_scale_inv, int cache_block_size, "
+      "Tensor? position_ids=None, Tensor? cos_sin_cache=None) -> ()");
+  ops.def(
+      "fused_kimi_k3_mla_decode_q_concat_ds_mla_insert("
+      "Tensor ql_nope, Tensor q_pe, Tensor kv_c_normed, Tensor k_pe, "
+      "Tensor! mqa_q, Tensor! k_cache, Tensor slot_mapping, "
+      "int cache_block_size, Tensor? position_ids=None, "
+      "Tensor? cos_sin_cache=None) -> ()");
+
 #ifndef USE_ROCM
  ops.def(
      "minimax_allreduce_rms_qk("
@@ -468,6 +511,24 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
      "int block_size, Tensor!? q_out, Tensor!? index_q_out, "
      "str kv_cache_dtype, bool skip_index_branch=False) -> ()");

+#ifdef VLLM_ENABLE_FUSED_KDA_DECODE
+  ops.def(
+      "fused_kda_decode("
+      "Tensor x, Tensor weight, Tensor? bias, Tensor! conv_state, "
+      "Tensor raw_g, Tensor raw_beta, Tensor A_log, Tensor dt_bias, "
+      "Tensor state_indices, Tensor! state, Tensor! out, "
+      "float? lower_bound=None, Tensor? output_gate=None, "
+      "Tensor? norm_weight=None, float norm_eps=1e-5) -> ()");
+#endif
+
+#ifdef VLLM_ENABLE_KIMI_K3_ATTN_RES
+  ops.def(
+      "kimi_k3_attn_res("
+      "Tensor! prefix, Tensor delta, Tensor blocks, Tensor norm_weight, "
+      "Tensor qk_weight, Tensor output_norm_weight, Tensor! output, "
+      "int num_blocks, float eps, float output_norm_eps) -> ()");
+#endif
+
  // Apply repetition penalties to logits in-place.
  ops.def(
      "apply_repetition_penalties_(Tensor! logits, Tensor prompt_mask, "
@@ -530,6 +591,14 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
      "limit=7.0) "
      "-> ()");

+  // Kimi SITU (SituGLU) gated activation. linear_beta<=0 means unset.
+  ops.def(
+      "situ_and_mul(Tensor! out, Tensor input, float beta=1.0, float "
+      "linear_beta=-1.0) -> ()");
+  ops.def(
+      "masked_situ_and_mul(Tensor! out, Tensor input, Tensor "
+      "expert_num_tokens, float beta=1.0, float linear_beta=-1.0) -> ()");
+
  // GELU implementation used in GPT-2.
  ops.def("gelu_new(Tensor! out, Tensor input) -> ()");

@@ -688,11 +757,30 @@ STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) {
  ops.impl(
      "fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_fp8_insert",
      TORCH_BOX(&fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_fp8_insert));
+  ops.impl("fused_kimi_k3_mla_key_concat_kv_cache_insert",
+           TORCH_BOX(&fused_kimi_k3_mla_key_concat_kv_cache_insert));
+  ops.impl("fused_kimi_k3_mla_key_concat_ds_mla_insert",
+           TORCH_BOX(&fused_kimi_k3_mla_key_concat_ds_mla_insert));
+  ops.impl("fused_kimi_k3_mla_qkv_quant_kv_cache_fp8_insert",
+           TORCH_BOX(&fused_kimi_k3_mla_qkv_quant_kv_cache_fp8_insert));
+  ops.impl("fused_kimi_k3_mla_decode_q_concat_kv_cache_insert",
+           TORCH_BOX(&fused_kimi_k3_mla_decode_q_concat_kv_cache_insert));
+  ops.impl("fused_kimi_k3_mla_decode_q_concat_kv_cache_fp8_insert",
+           TORCH_BOX(&fused_kimi_k3_mla_decode_q_concat_kv_cache_fp8_insert));
+  ops.impl("fused_kimi_k3_mla_decode_q_concat_ds_mla_insert",
+           TORCH_BOX(&fused_kimi_k3_mla_decode_q_concat_ds_mla_insert));
 #ifndef USE_ROCM
  ops.impl("minimax_allreduce_rms_qk", TORCH_BOX(&minimax_allreduce_rms_qk));
 #endif
  ops.impl("fused_minimax_m3_qknorm_rope_kv_insert",
           TORCH_BOX(&fused_minimax_m3_qknorm_rope_kv_insert));
+#ifdef VLLM_ENABLE_FUSED_KDA_DECODE
+  ops.impl("fused_kda_decode", TORCH_BOX(&fused_kda_decode));
+#endif
+
+#ifdef VLLM_ENABLE_KIMI_K3_ATTN_RES
+  ops.impl("kimi_k3_attn_res", TORCH_BOX(&kimi_k3_attn_res));
+#endif

  // Sampler kernels (shared CUDA/ROCm)
  ops.impl("apply_repetition_penalties_",
@@ -715,6 +803,8 @@ STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) {
  ops.impl("gelu_tanh_and_mul", TORCH_BOX(&gelu_tanh_and_mul));
  ops.impl("fatrelu_and_mul", TORCH_BOX(&fatrelu_and_mul));
  ops.impl("swigluoai_and_mul", TORCH_BOX(&swigluoai_and_mul));
+  ops.impl("situ_and_mul", TORCH_BOX(&situ_and_mul));
+  ops.impl("masked_situ_and_mul", TORCH_BOX(&masked_situ_and_mul));
  ops.impl("gelu_new", TORCH_BOX(&gelu_new));
  ops.impl("gelu_fast", TORCH_BOX(&gelu_fast));
  ops.impl("gelu_quick", TORCH_BOX(&gelu_quick));
@@ -812,6 +902,15 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C_cache_ops, ops) {
      "                     str kv_cache_dtype,"
      "                     Tensor scale) -> ()");

+  // Grouped concat_and_cache_mla across all layers (bf16 only). Each
+  // layer's cache base pointer is read from kv_cache_ptrs.
+  ops.def(
+      "concat_and_cache_mla_grouped(Tensor kv_c, Tensor k_pe,"
+      "                             Tensor kv_cache_ptrs,"
+      "                             Tensor slot_mapping,"
+      "                             int block_size, int block_stride,"
+      "                             int entry_stride) -> ()");
+
  // Rotate Q and K, then write to kv cache for MLA
  ops.def(
      "concat_and_cache_mla_rope_fused("
@@ -910,6 +1009,8 @@ STABLE_TORCH_LIBRARY_IMPL(_C_cache_ops, CUDA, ops) {
  ops.impl("reshape_and_cache", TORCH_BOX(&reshape_and_cache));
  ops.impl("reshape_and_cache_flash", TORCH_BOX(&reshape_and_cache_flash));
  ops.impl("concat_and_cache_mla", TORCH_BOX(&concat_and_cache_mla));
+  ops.impl("concat_and_cache_mla_grouped",
+           TORCH_BOX(&concat_and_cache_mla_grouped));
  ops.impl("concat_and_cache_mla_rope_fused",
           TORCH_BOX(&concat_and_cache_mla_rope_fused));
  ops.impl("convert_fp8", TORCH_BOX(&convert_fp8));
@@ -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>(
@@ -25,6 +25,10 @@
 ARG CUDA_VERSION=13.0.3
 ARG PYTHON_VERSION=3.12
 ARG UBUNTU_VERSION=22.04
+# DeepEPv2 requires NCCL >= 2.30.4 (GIN backend).
+# This version is only used for CUDA 13+ builds; CUDA 12 falls back to
+# the default NCCL version shipped with the base image.
+ARG NCCL_VERSION=2.30.7

 # By parameterizing the base images, we allow third-party to use their own
 # base images. One use case is hermetic builds with base images stored in
@@ -477,10 +481,17 @@ WORKDIR /workspace
 # Build DeepEP wheels
 COPY tools/ep_kernels/install_python_libraries.sh /tmp/install_python_libraries.sh
 # Defaults moved here from tools/ep_kernels/install_python_libraries.sh for centralized version management
-ARG DEEPEP_COMMIT_HASH=73b6ea4
+ARG DEEPEP_COMMIT_HASH=d4f41e4e93
 ARG NVSHMEM_VER
+ARG NCCL_VERSION
 RUN --mount=type=cache,target=/opt/uv/cache \
    mkdir -p /tmp/ep_kernels_workspace/dist && \
+    CUDA_MAJOR=$(echo $CUDA_VERSION | cut -d. -f1) && \
+    if [ "$CUDA_MAJOR" -ge 13 ] && [ -n "$NCCL_VERSION" ]; then \
+        echo "nvidia-nccl-cu${CUDA_MAJOR}==${NCCL_VERSION}" \
+            > /tmp/nccl-override.txt && \
+        export UV_OVERRIDE=/tmp/nccl-override.txt; \
+    fi && \
    export TORCH_CUDA_ARCH_LIST='9.0a 10.0a' && \
    /tmp/install_python_libraries.sh \
        --workspace /tmp/ep_kernels_workspace \
@@ -644,6 +655,7 @@ FROM ${FINAL_BASE_IMAGE} AS vllm-base

 ARG CUDA_VERSION
 ARG PYTHON_VERSION
+ARG NCCL_VERSION
 ARG DEADSNAKES_MIRROR_URL
 ARG DEADSNAKES_GPGKEY_URL
 ARG GET_PIP_URL
@@ -696,7 +708,6 @@ RUN apt-get update -y \
 # Install CUDA development tools for runtime JIT compilation
 # (FlashInfer, DeepGEMM, EP kernels all require compilation at runtime)
 RUN CUDA_VERSION_DASH=$(echo $CUDA_VERSION | cut -d. -f1,2 | tr '.' '-') && \
-    CUDA_VERSION_SHORT=$(echo $CUDA_VERSION | cut -d. -f1,2) && \
    apt-get update -y && \
    apt-get install -y --no-install-recommends --allow-change-held-packages \
        cuda-nvcc-${CUDA_VERSION_DASH} \
@@ -709,12 +720,6 @@ RUN CUDA_VERSION_DASH=$(echo $CUDA_VERSION | cut -d. -f1,2 | tr '.' '-') && \
        libnuma-dev \
        # numactl CLI for NUMA binding at runtime
        numactl && \
-    # Fixes nccl_allocator requiring nccl.h at runtime
-    # https://github.com/vllm-project/vllm/blob/1336a1ea244fa8bfd7e72751cabbdb5b68a0c11a/vllm/distributed/device_communicators/pynccl_allocator.py#L22
-    # NCCL packages don't use the cuda-MAJOR-MINOR naming convention,
-    # so we pin the version to match our CUDA version
-    NCCL_VER=$(apt-cache madison libnccl-dev | grep "+cuda${CUDA_VERSION_SHORT}" | head -1 | awk -F'|' '{gsub(/^ +| +$/, "", $2); print $2}') && \
-    apt-get install -y --no-install-recommends --allow-change-held-packages libnccl-dev=${NCCL_VER} libnccl2=${NCCL_VER} && \
    rm -rf /var/lib/apt/lists/*

 # Install uv for faster pip installs
@@ -734,6 +739,18 @@ RUN mkdir -p "${UV_PYTHON_INSTALL_DIR}" "${UV_CACHE_DIR}" \
    && chgrp -R 0 /opt/uv \
    && chmod -R g+rwX,a+rX /opt/uv

+# DeepEPv2 GIN requires NCCL >= 2.30.4 at both compile and runtime. torch pins
+# an older version as a transitive dep; this override forces uv to use our
+# pinned version whenever nvidia-nccl-cu* is resolved. Empty on CUDA 12 (no-op).
+RUN CUDA_MAJOR=$(echo $CUDA_VERSION | cut -d. -f1) && \
+    if [ "$CUDA_MAJOR" -ge 13 ]; then \
+        echo "nvidia-nccl-cu${CUDA_MAJOR}==${NCCL_VERSION}" \
+            > /etc/uv-overrides.txt; \
+    else \
+        touch /etc/uv-overrides.txt; \
+    fi
+ENV UV_OVERRIDE=/etc/uv-overrides.txt
+
 # ----------------------------------------------------------------------
 # Non-root support (opt-in)
 # ----------------------------------------------------------------------
@@ -745,6 +745,7 @@ ENV MIOPEN_DEBUG_CONV_GEMM=0
 # See: https://github.com/ROCm/rocm-libraries/issues/6266
 ENV HSA_ENABLE_IPC_MODE_LEGACY=1
 ENV UCX_RMA_PPLN_ENABLE=y
+ENV UCX_ROCM_COPY_SIGPOOL_MAX_ELEMS=inf

 # ROCm profiler limits workaround.
 RUN echo "ROCTRACER_MAX_EVENTS=10000000" > ${COMMON_WORKDIR}/libkineto.conf
@@ -835,6 +836,7 @@ COPY --from=export_vllm /docker ${COMMON_WORKDIR}/vllm/docker
 # See: https://github.com/ROCm/rocm-libraries/issues/6266
 ENV HSA_ENABLE_IPC_MODE_LEGACY=1
 ENV UCX_RMA_PPLN_ENABLE=y
+ENV UCX_ROCM_COPY_SIGPOOL_MAX_ELEMS=inf

 ENV TOKENIZERS_PARALLELISM=false

@@ -30,7 +30,7 @@ ENV LD_LIBRARY_PATH=/opt/rocm/lib:/usr/local/lib:
 ARG PYTORCH_ROCM_ARCH=gfx90a;gfx942;gfx950;gfx1100;gfx1101;gfx1200;gfx1201;gfx1150;gfx1151
 ENV PYTORCH_ROCM_ARCH=${PYTORCH_ROCM_ARCH}
 ENV AITER_ROCM_ARCH=gfx942;gfx950
-ENV MORI_GPU_ARCHS=gfx942;gfx950
+# Note: Do not set MORI_GPU_ARCHS here, it is automatically inferred at runtime

 # Required for RCCL in ROCm7.1
 ENV HSA_NO_SCRATCH_RECLAIM=1
@@ -10,6 +10,9 @@
    "UBUNTU_VERSION": {
      "default": "22.04"
    },
+    "NCCL_VERSION": {
+      "default": "2.30.7"
+    },
    "BUILD_BASE_IMAGE": {
      "default": "nvidia/cuda:13.0.3-devel-ubuntu22.04"
    },
@@ -56,7 +59,7 @@
      "default": "cuda"
    },
    "DEEPEP_COMMIT_HASH": {
-      "default": "73b6ea4"
+      "default": "d4f41e4e93"
    },
    "GIT_REPO_CHECK": {
      "default": "0"
@@ -222,7 +222,7 @@ MLA decode backends are selected using the standard
 | Backend | Dtypes | KV Dtypes | Block Sizes | Head Sizes | Sink | Non-Causal | Sparse | MM Prefix | DCP | Attention Types | Compute Cap. |
 | ------- | ------ | --------- | ----------- | ---------- | ---- | ---------- | ------ | --------- | --- | --------------- | ------------ |
 | `CUTLASS_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 128 | Any | ❌ | ❌ | ❌ | ❌ | ✅ | Decoder | 10.x |
-| `FLASHINFER_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 32, 64 | Any | ❌ | ❌ | ❌ | ❌ | ✅ | Decoder | 10.x |
+| `FLASHINFER_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 32, 64 | Any | ❌ | ✅ | ❌ | ❌ | ✅ | Decoder | 10.x |
 | `FLASHINFER_MLA_SPARSE` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 32, 64 | Any | ❌ | ❌ | ❌ | ❌ | ✅ | Decoder | 10.x |
 | `FLASHINFER_MLA_SPARSE_SM120` | bf16 | `auto`, `fp8`, `fp8_e4m3`, `fp8_ds_mla` | 64, 256 | Any | ❌ | ❌ | ❌ | ❌ | ❌ | Decoder | 12.x |
 | `FLASHMLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 64 | Any | ❌ | ❌ | ❌ | ❌ | ✅ | Decoder | 9.x-10.x |
@@ -122,8 +122,6 @@ For example:

 --8<-- "vllm/model_executor/layers/mamba/mamba_mixer2.py:mixer2_gated_rms_norm"

--8<-- "vllm/model_executor/models/plamo2.py:plamo2_mamba_mixer"
-
 --8<-- "vllm/model_executor/layers/mamba/short_conv.py:short_conv"
 ```

@@ -242,4 +242,4 @@ See [Fused MoE Kernel features](./moe_kernel_features.md#fused-moe-modular-all2a

 ## FusedMoEExpertsModular

-See [Fused MoE Kernel features](./moe_kernel_features.md#fused-moe-experts-kernels) for a list of all the available modular experts.
+See [Fused MoE Kernel features](./moe_kernel_features.md#fused-experts-kernels) for a list of all the available modular experts.
@@ -47,7 +47,8 @@ sequenceDiagram
    else only one side present
        PWriter->>PWriter: stash and wait, self-poll only when blocks unmatched
    end
-    PWriter->>PWriter: ensure D handshake (one-time)
+    PWriter->>PWriter: _ensure_handshake to D (async; defer WRITE)
+    PWriter->>PWriter: handshake callback re-queues on _deferred_push_inbox, wake
    PWriter->>DWriter: NIXL WRITE direct to D GPU + completion notif

    note over DWorker,DWriter: D side - completion accounting
@@ -100,15 +101,20 @@ event:
   D, completion notifs after a WRITE, late-arriving ``PUSH_REG``)
   even when there is no new metadata to act on.
 3. **Handshake-completion callback** (background handshake executor
-   thread) — when a deferred D→P handshake finishes successfully, the
-   future's done-callback re-enqueues the registration onto
-   ``_reg_send_inbox`` and sets the wake so the corresponding
-   ``send_notif`` runs on the writer (we never call ``send_notif`` from
-   the executor thread). On this second pass ``_ensure_handshake``
-   returns ``None`` (the agent is now connected), so the writer sends
-   the ``PUSH_REG`` directly. If the handshake *failed*, the callback
-   fails the request instead of re-enqueuing, so there is no retry
-   loop.
+   thread) — both handshakes run on the executor and never block the
+   writer; their done-callbacks re-enqueue the deferred op and set the
+   wake, since neither ``send_notif`` nor the NIXL WRITE may run off the
+   writer thread:
+    * the **D→P** handshake (before sending ``PUSH_REG``) re-enqueues the
+      registration onto ``_reg_send_inbox``;
+    * the **P→D** handshake (before a WRITE) re-enqueues the matched
+      ``(req_id, blocks, reg_data)`` onto ``_deferred_push_inbox``.
+
+   On this second pass ``_ensure_handshake`` returns ``None`` (the agent
+   is now connected), so the writer sends the ``PUSH_REG`` / issues the
+   WRITE directly. If a handshake *failed*, the callback fails or drops
+   the request instead of re-enqueuing, so there is no retry loop (see
+   Failure handling).

 In addition to event-driven wakes, the writer self-polls at
 ``_PUSH_WRITER_POLL_INTERVAL_MS = 1.0`` ms while there are P-side
@@ -227,6 +233,13 @@ Two per-request timers are armed on the scheduler:
 * **D-side ``send_notif`` failure when shipping the PUSH_REG to P** —
  identical handling: ``_handle_failed_transfer`` marks the recv as
  failed.
+* **P-side handshake failure (P→D handshake before a WRITE)** — the
+  future's done-callback logs ``push_handshake_failed`` and drops the
+  request without re-queuing. It deliberately does *not* call
+  ``_handle_failed_transfer`` (there is no ``_recving_metadata`` entry to
+  invalidate on the producer side, same reasoning as the WRITE-submission
+  failure below). P's blocks are reclaimed by the ``_kv_lease_duration``
+  lease and D's stale registration by its watchdog.
 * **P-side WRITE submission failure** — the WRITE handle (if any) is
  released and ``xfer_stats.record_failed_transfer()`` bumps the
  failure counter. We deliberately do not call
@@ -246,10 +259,11 @@ existing NIXL connector:
  class — all subclasses of the existing base classes;
 * one dedicated background thread per worker;
 * a few cross-thread queues, each with a single consumer (the writer);
-  most have one producer, except ``_reg_send_inbox``, which is fed both
-  by the engine main thread (new registrations) and by the
-  handshake-completion callback (registrations replayed after their
-  D→P handshake finishes);
+  most have one producer, except the two replay queues fed by both the
+  engine main thread and a handshake-completion callback:
+  ``_reg_send_inbox`` (registrations replayed after their D→P handshake)
+  and ``_deferred_push_inbox`` (matched pushes replayed after their P→D
+  handshake);
 * one new notification type (`PUSH_REG:<msgpack>`).

 Behavior on the engine main thread is otherwise unchanged. The writer
@@ -184,7 +184,7 @@ Our online Server provides endpoints that correspond to the offline APIs:
    - [Classification API](classify.md#online-serving)(`/classify`)
 - Corresponding to `LLM.score`:
    - [Score API](scoring.md#score-api) (`/score`, `/v1/score`)
-    - [Cohere Rerank API](scoring.md#rerank-api) (`/rerank`, `/v1/rerank`, `/v2/rerank`)
+    - [Cohere Rerank API](scoring.md#cohere-rerank-api) (`/rerank`, `/v1/rerank`, `/v2/rerank`)
 - Pooling API (`/pooling`) is similar to `LLM.encode`, being applicable to all types of pooling models.

 The following introduces the Pooling API. For other APIs, please refer to the link above.
@@ -20,7 +20,7 @@ The score models is designed to compute similarity scores between two input prom
    - `LLM.score`
 - Online APIs:
    - [Score API](scoring.md#score-api) (`/score`, `/v1/score`)
-    - [Cohere Rerank API](scoring.md#rerank-api) (`/rerank`, `/v1/rerank`, `/v2/rerank`)
+    - [Cohere Rerank API](scoring.md#cohere-rerank-api) (`/rerank`, `/v1/rerank`, `/v2/rerank`)

 !!! note
    Only when a classification model outputs num_labels equal to 1 can it be used as a scoring model and have its scoring API enabled.
@@ -435,7 +435,6 @@ th {
 | `PhiForCausalLM` | Phi | `microsoft/phi-1_5`, `microsoft/phi-2`, etc. | ✅︎ | ✅︎ |
 | `Phi3ForCausalLM` | Phi-4, Phi-3 | `microsoft/Phi-4-mini-instruct`, `microsoft/Phi-4`, `microsoft/Phi-3-mini-4k-instruct`, `microsoft/Phi-3-mini-128k-instruct`, `microsoft/Phi-3-medium-128k-instruct`, etc. | ✅︎ | ✅︎ |
 | `PhiMoEForCausalLM` | Phi-3.5-MoE | `microsoft/Phi-3.5-MoE-instruct`, etc. | ✅︎ | ✅︎ |
-| `Plamo2ForCausalLM` | PLaMo2 | `pfnet/plamo-2-1b`, `pfnet/plamo-2-8b`, etc. | ✅ | ✅︎ |
 | `Plamo3ForCausalLM` | PLaMo3 | `pfnet/plamo-3-nict-2b-base`, `pfnet/plamo-3-nict-8b-base`, etc. | ✅ | ✅︎ |
 | `Qwen2ForCausalLM` | QwQ, Qwen2 | `Qwen/QwQ-32B-Preview`, `Qwen/Qwen2-7B-Instruct`, `Qwen/Qwen2-7B`, etc. | ✅︎ | ✅︎ |
 | `Qwen2MoeForCausalLM` | Qwen2MoE | `Qwen/Qwen1.5-MoE-A2.7B`, `Qwen/Qwen1.5-MoE-A2.7B-Chat`, etc. | ✅︎ | ✅︎ |
@@ -548,6 +547,7 @@ These models primarily accept the [`LLM.generate`](./generative_models.md#llmgen
 | `KeyeVL1_5ForConditionalGeneration` | Keye-VL-1_5-8B | T + I<sup>E+</sup> + V<sup>E+</sup> | `Kwai-Keye/Keye-VL-1_5-8B` | ✅︎ | ✅︎ |
 | `KimiAudioForConditionalGeneration` | Kimi-Audio | T + A<sup>+</sup> | `moonshotai/Kimi-Audio-7B-Instruct` | | ✅︎ |
 | `KimiK25ForConditionalGeneration` | Kimi-K2.5 | T + I<sup>+</sup> | `moonshotai/Kimi-K2.5` | | ✅︎ |
+| `KimiK3ForConditionalGeneration` | Kimi-K3 | T + I<sup>+</sup> | `moonshotai/Kimi-K3` | | ✅︎ |
 | `KimiVLForConditionalGeneration` | Kimi-VL-A3B-Instruct, Kimi-VL-A3B-Thinking | T + I<sup>+</sup> | `moonshotai/Kimi-VL-A3B-Instruct`, `moonshotai/Kimi-VL-A3B-Thinking` | | ✅︎ |
 | `LightOnOCRForConditionalGeneration` | LightOnOCR-1B | T + I<sup>+</sup> | `lightonai/LightOnOCR-1B`, etc | ✅︎ | ✅︎ |
 | `Lfm2VlForConditionalGeneration` | LFM2-VL | T + I<sup>+</sup> | `LiquidAI/LFM2-VL-450M`, `LiquidAI/LFM2-VL-3B`, `LiquidAI/LFM2-VL-8B-A1B`, etc. | ✅︎ | ✅︎ |
@@ -12,6 +12,11 @@ Before using EP, you need to install the necessary dependencies. We are actively
 2. **Install DeepGEMM library**: Follow the [official instructions](https://github.com/deepseek-ai/DeepGEMM#installation).
 3. **For disaggregated serving**: Install `gdrcopy` by running the [`install_gdrcopy.sh`](../../tools/install_gdrcopy.sh) script (e.g., `install_gdrcopy.sh "${GDRCOPY_OS_VERSION}" "12.8" "x64"`). You can find available OS versions [here](https://developer.download.nvidia.com/compute/redist/gdrcopy/CUDA%2012.8/).

+!!! note "NCCL version (CUDA 13+)"
+    The `deepep_v2` backend requires NCCL >= 2.30.4. PyTorch ships an older
+    NCCL, so you must upgrade it before building or running DeepEP. See the
+    [EP kernels guide](../../tools/ep_kernels) for instructions.
+
 ### Backend Selection Guide

 vLLM provides multiple communication backends for EP. Use `--all2all-backend` to select one:
@@ -10,7 +10,7 @@ We currently support the following OpenAI APIs:
    - Only applicable to [text generation models](../../models/generative_models.md).
    - *Note: `suffix` parameter is not supported.*
 - [Chat Completions API](./openai_compatible_server.md#chat-api) (`/v1/chat/completions`)
-    - Only applicable to [text generation models](../../models/generative_models.md) with a [chat template](./openai_compatible_server.md#chat-template).
+    - Only applicable to [text generation models](../../models/generative_models.md) with a [chat template](#chat-template).
    - *Note: `user` parameter is ignored.*
    - *Note:* Setting the `parallel_tool_calls` parameter to `false` ensures vLLM only returns zero or one tool call per request. Setting it to `true` (the default) allows returning more than one tool call per request. There is no guarantee more than one tool call will be returned if this is set to `true`, as that behavior is model dependent and not all models are designed to support parallel tool calls.
 - [Chat Completions batch API](./openai_compatible_server.md#chat-api) (`/v1/chat/completions/batch`)
@@ -32,7 +32,7 @@ We currently support the following OpenAI APIs:
 - [Cohere Embed API](../../models/pooling_models/embed.md#cohere-embed-api) (`/v2/embed`)
    - Compatible with [Cohere's Embed API](https://docs.cohere.com/reference/embed)
    - Works with any [embedding model](../../models/pooling_models/embed.md#supported-models), including multimodal models.
- [Cohere Rerank API](../../models/pooling_models/scoring.md#rerank-api) (`/rerank`, `/v1/rerank`, `/v2/rerank`)
+- [Cohere Rerank API](../../models/pooling_models/scoring.md#cohere-rerank-api) (`/rerank`, `/v1/rerank`, `/v2/rerank`)
    - Implements [Jina AI's v1 rerank API](https://jina.ai/reranker/)
    - compatible with [Cohere's v1 & v2 rerank APIs](https://docs.cohere.com/v2/reference/rerank)

@@ -49,7 +49,7 @@ For further details on pooling models, please refer to [this page](../../models/
    - Only applicable to [embedding models](../../models/pooling_models/embed.md).
 - [Scoring Usages](../../models/pooling_models/scoring.md)
    - [Score API](../../models/pooling_models/scoring.md#score-api) (`/score`, `/v1/score`)
-    - [Cohere Rerank API](../../models/pooling_models/scoring.md#rerank-api) (`/rerank`, `/v1/rerank`, `/v2/rerank`)
+    - [Cohere Rerank API](../../models/pooling_models/scoring.md#cohere-rerank-api) (`/rerank`, `/v1/rerank`, `/v2/rerank`)
    - Applicable to [score models](../../models/pooling_models/scoring.md) (cross-encoder, bi-encoder, late-interaction).
 - [Pooling API](../../models/pooling_models/README.md#pooling-api) (`/pooling`)
    - Applicable to all [pooling models](../../models/pooling_models/README.md).
@@ -73,7 +73,7 @@ For further details on speech to text, please refer to [this page](speech_to_tex
    - Applicable to [score models](../../models/pooling_models/scoring.md) (cross-encoder, bi-encoder, late-interaction).
 - [Pooling API](../../models/pooling_models/README.md#pooling-api) (`/pooling`)
    - Applicable to all [pooling models](../../models/pooling_models/README.md).
- [Generative Scoring API](generative_scoring.md#generative-scoring-api) (`/generative_scoring`)
+- [Generative Scoring API](generative_scoring.md) (`/generative_scoring`)
    - Applicable to [CausalLM models](../../models/generative_models.md) (task `"generate"`).
    - Computes next-token probabilities for specified `label_token_ids`.

@@ -126,7 +126,7 @@ Models using selective state-space mechanisms instead of standard transformer at
 Models that use Mamba-2 and Mamba-1 layers (e.g., `Mamba2ForCausalLM`, `MambaForCausalLM`, `FalconMambaForCausalLM`) are supported.

 Hybrid models that combine Mamba-2 and Mamba-1 layers with standard attention layers are also supported (e.g.,
-`Zamba2ForCausalLM`, `NemotronHForCausalLM`, `FalconH1ForCausalLM` and `GraniteMoeHybridForCausalLM`, `JambaForCausalLM`, `Plamo2ForCausalLM`).
+`Zamba2ForCausalLM`, `NemotronHForCausalLM`, `FalconH1ForCausalLM` and `GraniteMoeHybridForCausalLM`, `JambaForCausalLM`).

 Hybrid models with mechanisms different to Mamba are also supported (e.g, `Lfm2ForCausalLM`).

@@ -17,7 +17,7 @@ PyNvVideoCodec==2.0.4
 flashinfer-python==0.6.15.post1
 flashinfer-cubin==0.6.15.post1
 apache-tvm-ffi==0.1.10
-tilelang==0.1.9
+tilelang==0.1.12
 nvidia-cudnn-frontend>=1.19.1
 # Required for LLM_NVTX_SCOPES_FOR_PROFILING=1
 nvtx==0.2.15
@@ -33,3 +33,6 @@ tokenspeed-mla==0.1.8; platform_system == "Linux"

 # Humming kernels for quantization gemm
 humming-kernels[cu13]==0.1.10
+
+# KDA
+flash-linear-attention==0.5.0
@@ -1157,7 +1157,7 @@ tqdm==4.67.3
    #   segmentation-models-pytorch
    #   sentence-transformers
    #   transformers
-transformers==5.13.1
+transformers==5.14.1
    # via
    #   -r requirements/test/../common.txt
    #   -r requirements/test/cuda.in
@@ -39,7 +39,7 @@ open_clip_torch==2.32.0 # Required for nemotron_vl test, Nemotron Parse in test_
 datamodel_code_generator # required for minicpm3 test
 lm-eval[api]>=0.4.12 # required for model evaluation test
 mteb[bm25s]>=2, <3 # required for mteb test
-transformers==5.13.1
+transformers==5.14.1
 tokenizers==0.22.2
 schemathesis>=4.0.0 # Required for openai schema test.
 # quantization
@@ -1259,7 +1259,7 @@ tqdm==4.67.3
    #   segmentation-models-pytorch
    #   sentence-transformers
    #   transformers
-transformers==5.13.1
+transformers==5.14.1
    # via
    #   -c requirements/common.txt
    #   -r requirements/test/../common.txt
@@ -29,7 +29,7 @@ opencv-python-headless >= 4.13.0 # required for video test
 datamodel_code_generator # required for minicpm3 test
 lm-eval[api]>=0.4.12 # required for model evaluation test
 mteb[bm25s]>=2, <3 # required for mteb test
-transformers==5.13.1
+transformers==5.14.1
 tokenizers==0.22.2
 schemathesis>=4.0.0 # Required for openai schema test.
 # quantization
@@ -35,7 +35,7 @@ open_clip_torch==2.32.0 # Required for nemotron_vl test, Nemotron Parse in test_
 datamodel_code_generator # required for minicpm3 test
 lm-eval[api]>=0.4.12 # required for model evaluation test
 mteb[bm25s]>=2, <3 # required for mteb test
-transformers==5.13.1
+transformers==5.14.1
 tokenizers==0.22.2
 schemathesis>=4.0.0 # Required for openai schema test
 # quantization
@@ -1214,7 +1214,7 @@ tqdm==4.67.3
    #   sentence-transformers
    #   tilelang
    #   transformers
-transformers==5.13.1
+transformers==5.14.1
    # via
    #   -c requirements/common.txt
    #   -r requirements/test/../common.txt
@@ -17,7 +17,7 @@ accelerate
 arctic-inference
 lm_eval[api]>=0.4.12
 modelscope<1.38
-transformers==5.13.1
+transformers==5.14.1

 # --- Audio Processing ---
 librosa
@@ -936,7 +936,7 @@ tqdm==4.67.3
    #   pqdm
    #   sentence-transformers
    #   transformers
-transformers==5.13.1
+transformers==5.14.1
    # via
    #   -c requirements/common.txt
    #   -r requirements/test/../common.txt
@@ -2220,7 +2220,7 @@ checksum = "11d3d7f243d5c5a8b9bb5d6dd2b1602c0cb0b9db1621bafc7ed66e35ff9fe092"
 [[package]]
 name = "llm-multimodal"
 version = "1.7.1"
-source = "git+https://github.com/smg-project/llm-multimodal?rev=5390032d6dc8a3e6fdc83acd320260367eb4b9b5#5390032d6dc8a3e6fdc83acd320260367eb4b9b5"
+source = "git+ssh://git@github.com/Inferact/llm-multimodal-internal.git?branch=k3-image#ceec43ec6beea5812d5ae59af712d9ff616496ef"
 dependencies = [
 "anyhow",
 "base64 0.22.1",
@@ -2235,7 +2235,6 @@ dependencies = [
 "once_cell",
 "pkg-config",
 "rayon",
- "realfft",
 "reqwest 0.13.4",
 "rustfft",
 "serde",
@@ -58,7 +58,7 @@ indexmap = "2.13.0"
 indicatif = "0.18.4"
 itertools = "0.14.0"
 libc = "0.2.177"
-llm-multimodal = { git = "https://github.com/smg-project/llm-multimodal", rev = "5390032d6dc8a3e6fdc83acd320260367eb4b9b5", default-features = false, features = ["native-tls"] }
+llm-multimodal = { git = "ssh://git@github.com/Inferact/llm-multimodal-internal.git", branch = "k3-image", default-features = false, features = ["native-tls"] }
 mimalloc = "0.1.52"
 minijinja = { version = "2.0", features = ["unstable_machinery", "json", "builtins", "loader", "loop_controls", "preserve_order"] }
 minijinja-contrib = { version = "2.0", features = ["pycompat"] }
@@ -20,7 +20,7 @@ use crate::output::{
 use crate::renderer::hf::{HfChatRenderer, MultimodalRenderInfo};
 use crate::renderer::{
    DeepSeekV4ChatRenderer, DeepSeekV32ChatRenderer, DynChatRenderer, HarmonyChatRenderer,
-    InklingChatRenderer,
+    InklingChatRenderer, KimiK3ChatRenderer,
 };
 use crate::request::ChatRequest;
 use crate::{DynChatOutputProcessor, RendererSelection};
@@ -73,6 +73,7 @@ impl HfChatBackend {
            RendererSelection::DeepSeekV4 => Arc::new(DeepSeekV4ChatRenderer::new()),
            RendererSelection::Harmony => Arc::new(HarmonyChatRenderer::new()?),
            RendererSelection::Inkling => Arc::new(InklingChatRenderer::new(tokenizer.clone())?),
+            RendererSelection::KimiK3 => Arc::new(KimiK3ChatRenderer::new()),
        };

        info!(
@@ -33,7 +33,8 @@ pub use parser::tool::{ToolParser, ToolParserError, ToolParserFactory};
 pub use renderer::hf::ChatTemplateContentFormatOption;
 pub use renderer::{
    ChatRenderer, DeepSeekV4ChatRenderer, DeepSeekV32ChatRenderer, DynChatRenderer,
-    HarmonyChatRenderer, InklingChatRenderer, RenderedPrompt, RendererSelection,
+    HarmonyChatRenderer, InklingChatRenderer, KimiK3ChatRenderer, RenderedPrompt,
+    RendererSelection,
 };
 pub use request::{
    ChatContent, ChatContentPart, ChatMessage, ChatOptions, ChatRequest, ChatRole, ChatTool,
@@ -326,6 +327,12 @@ mod tests {
        .unwrap();
    }

+    #[test]
+    fn validate_parser_overrides_accepts_explicit_kimi_k3() {
+        let selection = ParserSelection::Explicit("kimi_k3".to_string());
+        validate_parser_overrides(&selection, &selection).unwrap();
+    }
+
    #[test]
    fn validate_parser_overrides_accepts_auto_and_none() {
        validate_parser_overrides(&ParserSelection::Auto, &ParserSelection::None).unwrap();
@@ -339,7 +346,7 @@ mod tests {
        )
        .unwrap_err();

-        expect_test::expect!["tool parser `definitely_missing_tool_parser` is not registered (choose from: deepseek_v3, deepseek_v31, deepseek_v32, deepseek_v4, gemma4, glm45, glm47, granite4, hermes, hy_v3, inkling, internlm, kimi_k2, llama3_json, llama4_json, minimax_m2, minimax_m3, mistral, phi4_mini_json, qwen3_coder, qwen3_xml, seed_oss)"].assert_eq(&error.to_report_string());
+        expect_test::expect!["tool parser `definitely_missing_tool_parser` is not registered (choose from: deepseek_v3, deepseek_v31, deepseek_v32, deepseek_v4, gemma4, glm45, glm47, granite4, hermes, hy_v3, inkling, internlm, kimi_k2, kimi_k3, llama3_json, llama4_json, minimax_m2, minimax_m3, mistral, phi4_mini_json, qwen3_coder, qwen3_xml, seed_oss)"].assert_eq(&error.to_report_string());
    }

    #[test]
@@ -350,6 +357,6 @@ mod tests {
        )
        .unwrap_err();

-        expect_test::expect!["reasoning parser `definitely_missing_reasoning_parser` is not registered (choose from: cohere_cmd, deepseek_r1, deepseek_v3, deepseek_v4, gemma4, glm45, inkling, kimi, kimi_k2, minimax_m2, minimax_m3, nemotron_v3, qwen3, seed_oss, step3, step3p5)"].assert_eq(&error.to_report_string());
+        expect_test::expect!["reasoning parser `definitely_missing_reasoning_parser` is not registered (choose from: cohere_cmd, deepseek_r1, deepseek_v3, deepseek_v4, gemma4, glm45, inkling, kimi, kimi_k2, kimi_k3, minimax_m2, minimax_m3, nemotron_v3, qwen3, seed_oss, step3, step3p5)"].assert_eq(&error.to_report_string());
    }
 }
@@ -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)
@@ -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)
@@ -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();
+    }
 }
@@ -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,
    }
 }
@@ -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,
    }
 }
@@ -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,
        },
    )
@@ -101,7 +101,7 @@ fn fixture_request(name: &str) -> ChatRequest {

 fn inkling_fixture_options() -> FixtureRequestOptions {
    FixtureRequestOptions {
-        enable_thinking: false,
+        enable_thinking: Some(false),
        no_generation_prompt_when_last_assistant: false,
    }
 }
@@ -0,0 +1,568 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+//! Kimi K3 XTML prompt renderer.
+//!
+//! Port of Moonshot remote-code `encoding_k3.py::build_chat_segments()`.
+
+use std::collections::HashMap;
+use std::fmt::Write as _;
+
+use serde_json::{Map, Value, json};
+
+use crate::error::{Error, Result};
+use crate::request::{
+    ChatContent, ChatContentPart, ChatMessage, ChatRequest, ChatTool, ChatToolChoice,
+};
+use crate::{AssistantContentBlock, AssistantToolCall};
+
+pub(super) const OPEN: &str = "<|open|>";
+pub(super) const CLOSE: &str = "<|close|>";
+pub(super) const SEP: &str = "<|sep|>";
+pub(super) const END_OF_MSG: &str = "<|end_of_msg|>";
+pub(super) const IMAGE_PLACEHOLDER: &str = "<|media_pad|>";
+
+const DEFAULT_THINKING_EFFORT: &str = "max";
+const VALID_THINKING_EFFORTS: &[&str] = &["low", "high", "max"];
+
+/// Render one chat request into the K3 XTML prompt string.
+pub(super) fn render_request(request: &ChatRequest) -> Result<String> {
+    let thinking = thinking_enabled(request)?;
+    let thinking_effort = thinking.then(|| thinking_effort(request)).transpose()?;
+    let tools = request_tools(request);
+    let mut out = String::new();
+
+    if !tools.is_empty() {
+        write_tool_declare(&mut out, tools, false)?;
+    }
+
+    if let Some(effort) = thinking_effort {
+        // Preserve the checkpoint's literal guidance text: it still names
+        // `medium`, although the validator above no longer accepts it.
+        write_internal_system(
+            &mut out,
+            "thinking-effort",
+            &format!(
+                "`thinking_effort` guides on how much to think in your \
+                 thinking channel (not including the response channel), \
+                 supported values include `low`, `medium`, `high`, and `max`.\n\
+                 Now the system is invoked with `thinking_effort={effort}`."
+            ),
+        );
+    }
+
+    // Track prior assistant tool-call ids for tool-result reordering / naming.
+    let mut tool_call_id_index: HashMap<String, (usize, String)> = HashMap::new();
+    let mut pending_tool_run: Vec<(usize, ChatMessage)> = Vec::new();
+
+    let flush_tool_run = |out: &mut String,
+                          run: &mut Vec<(usize, ChatMessage)>,
+                          id_index: &HashMap<String, (usize, String)>|
+     -> Result<()> {
+        if run.is_empty() {
+            return Ok(());
+        }
+
+        let mut resolved = Vec::with_capacity(run.len());
+        let mut unresolved = false;
+        for (offset, message) in run.drain(..) {
+            let ChatMessage::ToolResponse {
+                content,
+                tool_call_id,
+            } = message
+            else {
+                unreachable!("pending tool run only holds tool responses");
+            };
+            match id_index.get(&tool_call_id) {
+                Some(&(position, ref name)) => {
+                    resolved.push((position, offset, content, Some(name.clone())));
+                }
+                None => {
+                    unresolved = true;
+                    resolved.push((usize::MAX, offset, content, None));
+                }
+            }
+        }
+
+        if unresolved {
+            // Preserve original order when the run cannot be fully matched.
+            resolved.sort_by_key(|item| item.1);
+        } else {
+            resolved.sort_by_key(|item| (item.0, item.1));
+        }
+
+        for (xtml_index, (_, _, content, name)) in resolved.into_iter().enumerate() {
+            let tool_name = name.as_deref().ok_or_else(|| {
+                Error::ChatTemplate(
+                    "Kimi K3 tool messages need a resolvable tool name: \
+                         carry a matching tool_call_id against a preceding \
+                         assistant tool_call"
+                        .to_string(),
+                )
+            })?;
+            write_tool_message(out, tool_name, xtml_index + 1, &content)?;
+        }
+        Ok(())
+    };
+
+    for (message_index, message) in request.messages.iter().enumerate() {
+        match message {
+            ChatMessage::ToolResponse { .. } => {
+                pending_tool_run.push((message_index, message.clone()));
+                continue;
+            }
+            _ => {
+                flush_tool_run(&mut out, &mut pending_tool_run, &tool_call_id_index)?;
+            }
+        }
+
+        match message {
+            // Python: role=system with a `tools` field renders a dynamic
+            // tool-declare (`## New Tools Available`). Map that to Developer
+            // messages that carry tools (OpenAI "developer" / system-tools).
+            ChatMessage::Developer {
+                content,
+                tools: Some(local_tools),
+            } if !local_tools.is_empty() => {
+                write_tool_declare(&mut out, local_tools, true)?;
+                if !content_is_empty(content) {
+                    write_role_message(&mut out, "system", None, content)?;
+                }
+            }
+            ChatMessage::System { content } | ChatMessage::Developer { content, .. } => {
+                write_role_message(&mut out, "system", None, content)?;
+            }
+            ChatMessage::User { content } => {
+                write_role_message(&mut out, "user", None, content)?;
+            }
+            ChatMessage::Assistant { content } => {
+                tool_call_id_index.clear();
+                let mut call_position = 0usize;
+                for block in content {
+                    if let AssistantContentBlock::ToolCall(call) = block {
+                        call_position += 1;
+                        if !call.id.is_empty() {
+                            tool_call_id_index
+                                .entry(call.id.clone())
+                                .or_insert((call_position, call.name.clone()));
+                        }
+                    }
+                }
+
+                write_assistant_message(&mut out, content, thinking)?;
+            }
+            ChatMessage::ToolResponse { .. } => unreachable!("handled above"),
+        }
+    }
+    flush_tool_run(&mut out, &mut pending_tool_run, &tool_call_id_index)?;
+
+    match &request.tool_choice {
+        ChatToolChoice::Required => {
+            write_internal_system(
+                &mut out,
+                "tool-choice",
+                "The system is invoked with `tool_choice=required`.\n\
+                 You MUST call tools in the next message.",
+            );
+        }
+        // Emit only when tools are present: Rust defaults tool_choice to None
+        // for tool-free requests, which must not inject a tool-choice message.
+        ChatToolChoice::None if !request.tools.is_empty() => {
+            write_internal_system(
+                &mut out,
+                "tool-choice",
+                "The system is invoked with `tool_choice=none`.\n\
+                 You MUST NOT call any tools in the next message.",
+            );
+        }
+        ChatToolChoice::None | ChatToolChoice::Auto | ChatToolChoice::Function { .. } => {}
+    }
+
+    write_response_format(&mut out, request)?;
+
+    if request.chat_options.add_generation_prompt() {
+        write_open_tag(&mut out, "message", &[("role", "assistant")]);
+        write_open_tag(&mut out, if thinking { "think" } else { "response" }, &[]);
+    }
+
+    Ok(out)
+}
+
+fn request_tools(request: &ChatRequest) -> &[ChatTool] {
+    // Declare tools whenever the request carries them. K3 tool-declare is
+    // independent of tool_choice; tool_choice only injects control messages.
+    request.tools.as_slice()
+}
+
+fn thinking_enabled(request: &ChatRequest) -> Result<bool> {
+    if let Some(thinking) = request.parse_template_bool("thinking")? {
+        return Ok(thinking);
+    }
+    if let Some(enable_thinking) = request.parse_template_bool("enable_thinking")? {
+        return Ok(enable_thinking);
+    }
+    Ok(request
+        .chat_options
+        .reasoning_effort
+        .map(|effort| effort != crate::request::ReasoningEffort::None)
+        .unwrap_or(true))
+}
+
+fn thinking_effort(request: &ChatRequest) -> Result<String> {
+    let effort = if let Some(value) = request.chat_options.template_kwargs.get("thinking_effort") {
+        value.as_str().ok_or_else(|| {
+            Error::ChatTemplate(format!(
+                "template kwarg `thinking_effort` must be a string, got {value}"
+            ))
+        })?
+    } else if let Some(effort) = request.chat_options.reasoning_effort {
+        effort.as_str()
+    } else if let Some(value) = request.chat_options.template_kwargs.get("reasoning_effort") {
+        value.as_str().ok_or_else(|| {
+            Error::ChatTemplate(format!(
+                "template kwarg `reasoning_effort` must be a string, got {value}"
+            ))
+        })?
+    } else {
+        DEFAULT_THINKING_EFFORT
+    };
+
+    if !VALID_THINKING_EFFORTS.contains(&effort) {
+        return Err(Error::ChatTemplate(format!(
+            "unsupported thinking_effort={effort:?}; supported values are `low`, `high`, and `max`"
+        )));
+    }
+    Ok(effort.to_string())
+}
+
+fn content_is_empty(content: &ChatContent) -> bool {
+    match content {
+        ChatContent::Text(text) => text.is_empty(),
+        ChatContent::Parts(parts) => parts.iter().all(|part| match part {
+            ChatContentPart::Text { text } => text.is_empty(),
+            ChatContentPart::ImageUrl { .. }
+            | ChatContentPart::VideoUrl { .. }
+            | ChatContentPart::InputAudio { .. }
+            | ChatContentPart::AudioUrl { .. } => false,
+        }),
+    }
+}
+
+fn write_tool_declare(out: &mut String, tools: &[ChatTool], dynamic: bool) -> Result<()> {
+    let mut specs = Vec::with_capacity(tools.len());
+    for tool in tools {
+        let mut function = Map::new();
+        function.insert(
+            "description".to_string(),
+            Value::String(tool.description.clone().unwrap_or_default()),
+        );
+        function.insert("name".to_string(), Value::String(tool.name.clone()));
+        function.insert("parameters".to_string(), sort_json(&tool.parameters));
+        specs.push(json!({
+            "function": Value::Object(function),
+            "type": "function",
+        }));
+    }
+    let payload = compact_json(&sort_json(&Value::Array(specs)))?;
+
+    let body = if dynamic {
+        format!(
+            "## New Tools Available\n\
+             The system dynamically extends the toolset via lazy-loading.\n\
+             You have access to all existing and extended tools.\n\
+             Here are the specs for the extended tools.\n\n\
+             ```json\n\
+             {payload}\n\
+             ```"
+        )
+    } else {
+        format!(
+            "# Tools\n\
+             Here are the available tools, described in JSONSchema.\n\n\
+             ```json\n\
+             {payload}\n\
+             ```"
+        )
+    };
+
+    write_internal_system(out, "tool-declare", &body);
+    Ok(())
+}
+
+fn write_internal_system(out: &mut String, message_type: &str, body: &str) {
+    write_open_tag(
+        out,
+        "message",
+        &[("role", "system"), ("type", message_type)],
+    );
+    out.push_str(body.trim());
+    write_close_tag(out, "message");
+    out.push_str(END_OF_MSG);
+}
+
+fn write_role_message(
+    out: &mut String,
+    role: &str,
+    name: Option<&str>,
+    content: &ChatContent,
+) -> Result<()> {
+    let mut attrs = vec![("role", role.to_string())];
+    if let Some(name) = name.filter(|name| !name.is_empty()) {
+        attrs.push(("name", name.to_string()));
+    }
+    let attr_refs: Vec<(&str, &str)> = attrs.iter().map(|(k, v)| (*k, v.as_str())).collect();
+    write_open_tag(out, "message", &attr_refs);
+    write_content(out, content)?;
+    write_close_tag(out, "message");
+    out.push_str(END_OF_MSG);
+    Ok(())
+}
+
+fn write_tool_message(
+    out: &mut String,
+    tool_name: &str,
+    index: usize,
+    content: &ChatContent,
+) -> Result<()> {
+    let index_str = index.to_string();
+    write_open_tag(
+        out,
+        "message",
+        &[("role", "tool"), ("tool", tool_name), ("index", &index_str)],
+    );
+    write_content(out, content)?;
+    write_close_tag(out, "message");
+    out.push_str(END_OF_MSG);
+    Ok(())
+}
+
+fn write_assistant_message(
+    out: &mut String,
+    content: &[AssistantContentBlock],
+    thinking: bool,
+) -> Result<()> {
+    write_open_tag(out, "message", &[("role", "assistant")]);
+
+    let mut reasoning = String::new();
+    let mut response = String::new();
+    let mut tool_calls = Vec::new();
+    for block in content {
+        match block {
+            AssistantContentBlock::Reasoning { text } => reasoning.push_str(text),
+            AssistantContentBlock::Text { text } => response.push_str(text),
+            AssistantContentBlock::ToolCall(call) => tool_calls.push(call),
+        }
+    }
+
+    // The think channel is structural: in thinking mode every assistant
+    // message carries open/close tags even when there is no reasoning content.
+    // In non-thinking mode the channel is dropped entirely.
+    if thinking {
+        write_open_tag(out, "think", &[]);
+        if !reasoning.trim().is_empty() {
+            out.push_str(&reasoning);
+        }
+        write_close_tag(out, "think");
+    }
+
+    write_open_tag(out, "response", &[]);
+    out.push_str(&response);
+    write_close_tag(out, "response");
+
+    if !tool_calls.is_empty() {
+        write_open_tag(out, "tools", &[]);
+        for (index, tool_call) in tool_calls.into_iter().enumerate() {
+            write_assistant_tool_call(out, tool_call, index + 1)?;
+        }
+        write_close_tag(out, "tools");
+    }
+
+    write_close_tag(out, "message");
+    out.push_str(END_OF_MSG);
+    Ok(())
+}
+
+fn write_assistant_tool_call(
+    out: &mut String,
+    tool_call: &AssistantToolCall,
+    index: usize,
+) -> Result<()> {
+    let index_str = index.to_string();
+    write_open_tag(
+        out,
+        "call",
+        &[
+            ("tool", tool_call.name.as_str()),
+            ("index", index_str.as_str()),
+        ],
+    );
+
+    let (args, json_block) = normalize_tool_arguments(&tool_call.arguments)?;
+    if let Some(raw) = json_block {
+        write_open_tag(out, "json", &[("type", "object")]);
+        out.push_str(&raw);
+        write_close_tag(out, "json");
+    } else {
+        for (key, value) in args {
+            let typ = xtml_type(&value);
+            write_open_tag(out, "argument", &[("key", key.as_str()), ("type", typ)]);
+            out.push_str(&xtml_value(&value));
+            write_close_tag(out, "argument");
+        }
+    }
+
+    write_close_tag(out, "call");
+    Ok(())
+}
+
+fn write_content(out: &mut String, content: &ChatContent) -> Result<()> {
+    match content {
+        ChatContent::Text(text) => write_text_with_images(out, text),
+        ChatContent::Parts(parts) => {
+            for part in parts {
+                match part {
+                    ChatContentPart::Text { text } => write_text_with_images(out, text)?,
+                    ChatContentPart::ImageUrl { .. } => out.push_str(IMAGE_PLACEHOLDER),
+                    ChatContentPart::VideoUrl { .. } => {
+                        return Err(Error::UnsupportedMultimodalContent("video_url"));
+                    }
+                    ChatContentPart::InputAudio { .. } => {
+                        return Err(Error::UnsupportedMultimodalContent("input_audio"));
+                    }
+                    ChatContentPart::AudioUrl { .. } => {
+                        return Err(Error::UnsupportedMultimodalContent("audio_url"));
+                    }
+                }
+            }
+            Ok(())
+        }
+    }
+}
+
+fn write_text_with_images(out: &mut String, text: &str) -> Result<()> {
+    // Placeholder expansion is left as the literal K3 image token; multimodal
+    // preprocessing can replace it once image prompts are known.
+    out.push_str(text);
+    Ok(())
+}
+
+fn write_response_format(out: &mut String, request: &ChatRequest) -> Result<()> {
+    let Some(rf) = request.chat_options.template_kwargs.get("response_format") else {
+        return Ok(());
+    };
+
+    let rf_type = rf.get("type").and_then(Value::as_str).or_else(|| rf.as_str()).unwrap_or("");
+
+    match rf_type {
+        "json_object" => {
+            write_internal_system(
+                out,
+                "response-format",
+                "The system is invoked with `response_format=json_object`.\n\
+                 Your response must be raw JSON data without markdown code \
+                 blocks (```json) or any additional formatting.",
+            );
+        }
+        "json_schema" => {
+            let schema = extract_response_schema(rf);
+            let schema_json = compact_json(&sort_json(&schema.unwrap_or(Value::Null)))?;
+            write_internal_system(
+                out,
+                "response-format",
+                &format!(
+                    "The system is invoked with `response_format=json_schema`.\n\
+                     Your response must be raw JSON data without markdown code \
+                     blocks (```json) or any additional formatting.\n\
+                     The JSON data must match the following schema:\n\
+                     ```json\n\
+                     {schema_json}\n\
+                     ```"
+                ),
+            );
+        }
+        _ => {}
+    }
+    Ok(())
+}
+
+fn extract_response_schema(response_format: &Value) -> Option<Value> {
+    let json_schema = response_format.get("json_schema")?;
+    if let Some(schema) = json_schema.get("schema") {
+        return Some(schema.clone());
+    }
+    if let Some(schema) = json_schema.get("json_schema") {
+        return Some(schema.clone());
+    }
+    Some(json_schema.clone())
+}
+
+fn normalize_tool_arguments(arguments: &str) -> Result<(Map<String, Value>, Option<String>)> {
+    let trimmed = arguments.trim();
+    if trimmed.is_empty() {
+        return Ok((Map::new(), None));
+    }
+    match serde_json::from_str::<Value>(trimmed) {
+        Ok(Value::Object(map)) => Ok((map, None)),
+        Ok(_) => Err(Error::ChatTemplate(
+            "Kimi K3 tool call arguments must be a JSON object".to_string(),
+        )),
+        Err(_) => Ok((Map::new(), Some(arguments.to_string()))),
+    }
+}
+
+fn xtml_type(value: &Value) -> &'static str {
+    match value {
+        Value::Bool(_) => "boolean",
+        Value::Null => "null",
+        Value::Number(_) => "number",
+        Value::String(_) => "string",
+        Value::Object(_) => "object",
+        Value::Array(_) => "array",
+    }
+}
+
+fn xtml_value(value: &Value) -> String {
+    match value {
+        Value::String(text) => text.clone(),
+        other => compact_json(other).unwrap_or_else(|_| other.to_string()),
+    }
+}
+
+fn write_open_tag(out: &mut String, tag: &str, attrs: &[(&str, &str)]) {
+    out.push_str(OPEN);
+    out.push_str(tag);
+    for (key, value) in attrs {
+        let _ = write!(out, " {key}=\"{}\"", escape_attr_value(value));
+    }
+    out.push_str(SEP);
+}
+
+fn write_close_tag(out: &mut String, tag: &str) {
+    out.push_str(CLOSE);
+    out.push_str(tag);
+    out.push_str(SEP);
+}
+
+fn escape_attr_value(value: &str) -> String {
+    value.replace('&', "&amp;").replace('"', "&quot;")
+}
+
+fn compact_json(value: &Value) -> Result<String> {
+    serde_json::to_string(value).map_err(|error| Error::ChatTemplate(error.to_string()))
+}
+
+fn sort_json(value: &Value) -> Value {
+    match value {
+        Value::Array(items) => Value::Array(items.iter().map(sort_json).collect()),
+        Value::Object(map) => {
+            let mut sorted = Map::new();
+            let mut keys = map.keys().collect::<Vec<_>>();
+            keys.sort();
+            for key in keys {
+                sorted.insert(key.clone(), sort_json(&map[key]));
+            }
+            Value::Object(sorted)
+        }
+        _ => value.clone(),
+    }
+}
@@ -0,0 +1,15 @@
+{
+  "messages": [
+    {
+      "role": "user",
+      "content": "json pls"
+    }
+  ],
+  "add_generation_prompt": true,
+  "template_kwargs": {
+    "thinking": false,
+    "response_format": {
+      "type": "json_object"
+    }
+  }
+}
@@ -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|>
@@ -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
+  }
+}
@@ -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|>
@@ -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"
+  }
+}
@@ -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|>
@@ -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
+  }
+}
@@ -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|>
@@ -0,0 +1,36 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+//! Native Kimi K3 XTML chat renderer.
+
+mod encoding;
+#[cfg(test)]
+mod tests;
+
+use vllm_text::Prompt;
+
+use super::{ChatRenderer, RenderedPrompt, request_template_kwargs};
+use crate::Result;
+use crate::request::ChatRequest;
+
+/// Dedicated Kimi K3 XTML renderer.
+#[derive(Debug, Clone, Copy, Default)]
+pub struct KimiK3ChatRenderer;
+
+impl KimiK3ChatRenderer {
+    /// Create a Kimi K3 renderer.
+    pub fn new() -> Self {
+        Self
+    }
+}
+
+impl ChatRenderer for KimiK3ChatRenderer {
+    fn render(&self, request: &ChatRequest) -> Result<RenderedPrompt> {
+        request.validate()?;
+
+        Ok(RenderedPrompt {
+            prompt: Prompt::Text(encoding::render_request(request)?),
+            effective_template_kwargs: request_template_kwargs(request),
+        })
+    }
+}
--- a/Show More
+++ b/Show More
Author	SHA1	Message	Date
+11 zjy0516 Isotr0py Thien Tran Bugen Zhao Yongye Zhu Ziming Huang Roger Wang Isotr0py aoshen02 Woosuk Kwon Jee Jee Li aoshen02 Summer Yang Kevin H. Luu Bowen Wang gnovack Nick Hill Yifan Qiao xiaozhoupy Roy Wang Jeff Ma	f68f4fddea	kimi-k3 Co-authored-by: Isotr0py <Isotr0py@outlook.com> Co-authored-by: Thien Tran <gau.nernst@yahoo.com.sg> Co-authored-by: Bugen Zhao <i@bugenzhao.com> Co-authored-by: Yongye Zhu <zyy1102000@gmail.com> Co-authored-by: Ziming Huang <zelda.huanghuang@gmail.com> Co-authored-by: Roger Wang <hey@rogerw.io> Co-authored-by: Isotr0py <mozf@inferact.ai> Co-authored-by: aoshen02 <aoshen02@users.noreply.github.com> Co-authored-by: Woosuk Kwon <woosuk@inferact.ai> Co-authored-by: Jee Jee Li <jeejeelee@inferact.ai> Co-authored-by: aoshen02 <aoshen@inferact.ai> Co-authored-by: Jiangyun Zhu <riverclouds.zhu@qq.com> Co-authored-by: Summer Yang <girasoleyang@gmail.com> Co-authored-by: Kevin H. Luu <khluu000@gmail.com> Co-authored-by: khluu <khluu000@gmail.com> Co-authored-by: Bowen Wang <abmfy@icloud.com> Co-authored-by: gnovack <novackgm@gmail.com> Co-authored-by: Nick Hill <nickhill123@gmail.com> Co-authored-by: Yifan Qiao <yifanqiao@inferact.ai> Co-authored-by: xiaozhoupy <peiyuanzhou1994@gmail.com> Co-authored-by: Roy Wang <yasong.wang@inferact.ai> Co-authored-by: Jeff (Junze) Ma <93145857+majunze2001@users.noreply.github.com> Signed-off-by: zjy0516 <riverclouds.zhu@qq.com>	2026-07-27 07:01:52 +00:00
Aarushi Jain GitHub Andreas Karatzas Andreas Karatzas OpenAI Codex	0ba2aa35a8	Stabilize GPU memory teardown between ROCm CI tests (#49242 ) Signed-off-by: aarushjain29 <Aarushi.Jain2@amd.com> Signed-off-by: Andreas Karatzas <Andreas.Karatzas@amd.com> Co-authored-by: Andreas Karatzas <akaratza@amd.com> Co-authored-by: Andreas Karatzas <Andreas.Karatzas@amd.com> Co-authored-by: OpenAI Codex <codex@openai.com>	2026-07-25 04:21:36 +00:00
Divakar VermaandGitHub	aaaeda98dc	[CI] fix compile test \| refactor VLLM_DISABLE_COMPILE_CACHE for tests (#49770 ) Signed-off-by: Divakar Verma <divakar.verma@amd.com>	2026-07-24 23:15:19 -05:00
Aarushi Jain GitHub Andreas Karatzas	d9cd774198	[ROCm][CI] Force native compile caches onto local disk (#49763 ) Signed-off-by: aarushjain29 <Aarushi.Jain2@amd.com> Co-authored-by: Andreas Karatzas <akaratza@amd.com>	2026-07-24 23:14:37 -05:00
Yifan Qiao GitHub Claude	70052fb924	[Bugfix][KV Connector][Mooncake] Keep TP-sharded Mamba state out of the KV-head dedup (#49499 ) Signed-off-by: Yifan Qiao <yifanqiao@inferact.ai> Co-authored-by: Claude <noreply@anthropic.com>	2026-07-24 21:02:22 -07:00
liuzhenweiandGitHub	318b527cc2	[XPU] add warning for xpu graph limitations (#49419 ) Signed-off-by: zhenwei-intel <zhenwei.liu@intel.com>	2026-07-25 02:26:15 +00:00
liuzhenwei GitHub Kunshang Ji	6a1acac3fe	[BUGFIX] Fix log capture in KV test (#49655 ) Signed-off-by: zhenwei-intel <zhenwei.liu@intel.com> Co-authored-by: Kunshang Ji <kunshang.ji@intel.com>	2026-07-25 00:36:53 +00:00
Woosuk KwonandGitHub	213f681f81	Revert "[Perf][GLM-5.2] Blackwell decode optimizations" (#49768 )	2026-07-24 16:04:59 -07:00
Aarushi JainandGitHub	33c4f3551c	[ROCm][CI] Wait for ROCm VRAM to settle between compiled and eager LL… (#49739 ) Signed-off-by: aarushjain29 <Aarushi.Jain2@amd.com>	2026-07-24 17:19:46 -05:00
Rohan Potdar GitHub Claude Opus 4.8	caa9cad31e	[ROCm][Docker] Drop MORI_GPU_ARCHS so MoRI autodetects the device arch (#49737 ) Signed-off-by: Rohan Potdar <rohan.potdar@amd.com> Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>	2026-07-24 21:14:03 +00:00
djramic GitHub Andreas Karatzas OpenAI Codex	7513d071bd	[ROCm][CI] Fix XPASS(strict) on mixed audio embeds test (#49733 ) Signed-off-by: Djordje Ramic <djoramic@amd.com> Signed-off-by: Andreas Karatzas <akaratza@amd.com> Co-authored-by: Andreas Karatzas <akaratza@amd.com> Co-authored-by: OpenAI Codex <codex@openai.com>	2026-07-24 16:09:18 -05:00
Harry MellorandGitHub	89f6aa3a9e	[KV Offload][CI] Fall back to buffered I/O without O_DIRECT; fix flaky api-server test (#49734 ) Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>	2026-07-24 13:56:57 -07:00
Harry MellorandGitHub	84d26b9ee3	[Model] Remove Plamo2 (#49729 ) Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>	2026-07-24 13:54:24 -07:00
Jiangyun Zhu GitHub OpenAI Codex	9e6746b3c7	[CI] Stabilize memory-sensitive compile and structured output tests (#49749 ) Signed-off-by: zjy0516 <riverclouds.zhu@qq.com> Co-authored-by: OpenAI Codex <noreply@openai.com>	2026-07-24 16:45:09 -04:00
Elvir Crnčević GitHub Claude Opus 4.6	972848f276	[Bugfix] Support non-uniform page sizes in KVBlockZeroer (#49704 ) Signed-off-by: Elvir Crncevic <elvircrn@gmail.com> Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>	2026-07-24 13:38:59 -07:00
Oxana Korzh GitHub Claude Opus 4.8 Andreas Karatzas	2279575cd9	[AMD][Bugfix][EPLB] Fix elastic EP scaling accuracy on ROCm (#47206 ) Signed-off-by: okorzh <okorzh@amd.com> Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Co-authored-by: Andreas Karatzas <akaratza@amd.com>	2026-07-24 13:35:27 -07:00
Tyler Michael Smith GitHub Claude Ilya Markov OpenAI Codex Codex	5d8e90a966	[WideEP] Update NCCL to 2.30.7 to enable DeepEPv2 in the vllm/vllm-openai image (#45321 ) Signed-off-by: Tyler Michael Smith <tlrmchlsmth@gmail.com> Signed-off-by: Tyler Michael Smith <tyler@vllm.ai> Signed-off-by: Tyler Michael Smith <tyler@tylermsmith.com> Co-authored-by: Claude <noreply@anthropic.com> Co-authored-by: Ilya Markov <ilmarkov@users.noreply.github.com> Co-authored-by: OpenAI Codex <codex@openai.com> Co-authored-by: Codex <noreply@openai.com>	2026-07-24 13:00:02 -07:00
djramicandGitHub	e222c33f2f	[Bugfix] Register axk1 config to fix A.X-K1 init (#49727 ) Signed-off-by: Djordje Ramic <djoramic@amd.com>	2026-07-24 12:57:49 -07:00
Aarushi Jain GitHub Andreas Karatzas	c064fa52b6	Fix GLM-4.1V video placeholder token ID handling. (#49484 ) Signed-off-by: aarushjain29 <Aarushi.Jain2@amd.com> Co-authored-by: Andreas Karatzas <akaratza@amd.com>	2026-07-24 12:35:16 -07:00
Andreas KaratzasandGitHub	7e51939e25	[CI] Avoid unnecessary Hugging Face metadata requests (#49508 ) Signed-off-by: Andreas Karatzas <Andreas.Karatzas@amd.com>	2026-07-24 12:33:53 -07:00
Andreas KaratzasandGitHub	9863102ed9	[CI] Reuse loaded config for cached tokenizer (#49509 ) Signed-off-by: Andreas Karatzas <Andreas.Karatzas@amd.com>	2026-07-24 12:32:40 -07:00
Tyler Michael Smith GitHub OpenAI Codex	8c13ee5735	Add `sm_107` for Rubin (#49387 ) Signed-off-by: Tyler Michael Smith <tlrmchlsmth@gmail.com> Co-authored-by: OpenAI Codex <noreply@openai.com>	2026-07-24 11:59:49 -07:00
Andreas KaratzasandGitHub	41798069f3	[CI][AMD] Deprecate DinD for MI355 tests (#49257 ) Signed-off-by: Andreas Karatzas <Andreas.Karatzas@amd.com>	2026-07-24 13:41:08 -05:00
Johnny-Liou GitHub Claude Opus 4.8mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>tomeras91 Cyrus Leung	866fea2b99	[Kernel] ReplaySSM: cache SSM inputs for faster Mamba2 standard decode (#48018 ) Signed-off-by: Johnny-Liou <a897111@gmail.com> Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> Co-authored-by: tomeras91 <57313761+tomeras91@users.noreply.github.com> Co-authored-by: Cyrus Leung <tlleungac@connect.ust.hk>	2026-07-24 09:39:49 -07:00
Thomas Fahrner GitHub Claude Fable 5	d02df748bf	[Bugfix] Accept RFC 2397 parameters in base64 data URLs (#48973 ) Signed-off-by: Thomas Fahrner <thomas.fahrner@parasail.io> Co-authored-by: Claude Fable 5 <noreply@anthropic.com>	2026-07-24 08:23:08 -07:00
Taneem Ibrahim GitHub Wentao Ye	453f01783d	[UX] Improve data-parallel launch validation (#49124 ) Signed-off-by: Taneem Ibrahim <taneem.ibrahim@gmail.com> Co-authored-by: Wentao Ye <44945378+yewentao256@users.noreply.github.com>	2026-07-24 07:16:39 -07:00
Taneem Ibrahim GitHubmergify[bot] <37929162+mergify[bot]@users.noreply.github.com>	7b40fb9645	[UX] Reject incompatible nested runtime overrides (#49247 ) Signed-off-by: Taneem Ibrahim <taneem.ibrahim@gmail.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>	2026-07-24 07:16:22 -07:00
BadrBasowidandGitHub	8eac21a602	[ROCM] Fix AITER Fused AllReduce RMSNorm for Transformers Backend (#49673 ) Signed-off-by: BadrBasowid <badr.basowid@gmail.com>	2026-07-24 07:16:17 -07:00
mgazz GitHub Kevin H. Luu	a454a1dd25	[Bugfix][Benchmarks] Restore --skip-tokenizer-init with custom dataset (#49180 ) Signed-off-by: Michele Gazzetti <michele.gazzetti1@ibm.com> Co-authored-by: Kevin H. Luu <khluu000@gmail.com>	2026-07-24 04:30:19 -07:00
hotTea GitHub hanxi-java 韩熙	833483f357	Encoder cache extension hooks (#48218 ) Signed-off-by: hotTea <958436561@qq.com> Signed-off-by: hanxi-java <634498162@qq.com> Co-authored-by: hanxi-java <634498162@qq.com> Co-authored-by: 韩熙 <63780107+hanxi-java@users.noreply.github.com>	2026-07-24 02:52:00 -07:00
li-jinpeng GitHub xymli	163ecba377	[Bugfix] Skip linear bias in layerwise reload to avoid corruption (#49586 ) Signed-off-by: li-jinpeng <3332126450@qq.com> Signed-off-by: xymli <xymli@tencent.com> Co-authored-by: xymli <xymli@tencent.com>	2026-07-24 16:53:48 +08:00
Nicolò Lucchesi GitHubmergify[bot] <37929162+mergify[bot]@users.noreply.github.com>	589a5b884b	[PD][NixlPush][Bugfix] Fix blocking handshake call on writer thread (#49221 ) Signed-off-by: NickLucche <nicolo.lucchesi@mistral.ai> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>	2026-07-24 01:41:37 -07:00
Kevin H. Luu GitHub OpenAI Codex	5c5434e2d8	Remove Quantization test parallelism (#49693 ) Signed-off-by: khluu <khluu000@gmail.com> Co-authored-by: OpenAI Codex <noreply@openai.com>	2026-07-24 01:14:48 -07:00
xiaozhoupy GitHub Claude Giancarlo Delfin Thien Tran OpenAI Codex zhou stefankoncarevic	dd72658e7d	[Perf][GLM-5.2] Blackwell decode optimizations (#48597 ) Signed-off-by: Peiyuan Zhou <peiyuanzhou1994@gmail.com> Signed-off-by: zhou <zhou@zhoudeMacBook-Pro.local> Signed-off-by: Stefan Koncarevic <stefan.koncarevic@amd.com> Co-authored-by: Claude <noreply@anthropic.com> Co-authored-by: Giancarlo Delfin <32987265+TheEpicDolphin@users.noreply.github.com> Co-authored-by: Thien Tran <gau.nernst@yahoo.com.sg> Co-authored-by: OpenAI Codex <codex@openai.com> Co-authored-by: zhou <zhou@zhoudeMacBook-Pro.local> Co-authored-by: stefankoncarevic <Stefan.Koncarevic@amd.com>	2026-07-23 21:36:27 -07:00
Harry MellorandGitHub	0d77325b10	Bump Transformers version to 5.14.1 (#49223 ) Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>	2026-07-23 21:16:49 -07:00
Andreas KaratzasandGitHub	2ac125123a	[CI] Use explicit devices in IR tests (#49513 ) Signed-off-by: Andreas Karatzas <Andreas.Karatzas@amd.com>	2026-07-23 20:08:14 -07:00
Kyle Sayers GitHub Claude Sonnet 4.6	7bdf8cc37c	[Bugfix] Fix humming kernel crash when layer.has_bias is None (#48769 ) Signed-off-by: Kyle Sayers <kylesayrs@gmail.com> Co-authored-by: Claude Sonnet 4.6 <noreply@anthropic.com>	2026-07-23 20:07:45 -07:00
Kyle SayersandGitHub	bf27e34ebb	[CompressedTensors] DeepSeek4 CT Quantization Support (#41276 ) Signed-off-by: Kyle Sayers <kylesayrs@gmail.com>	2026-07-23 20:07:31 -07:00
Michael Goin GitHub OpenAI Codex	275556c35c	[Bugfix] Detect mixed precision in packed KV cache specs (#49623 ) Signed-off-by: mgoin <mgoin64@gmail.com> Signed-off-by: Michael Goin <mgoin64@gmail.com> Co-authored-by: OpenAI Codex <noreply@openai.com>	2026-07-23 20:07:06 -07:00
Michael GoinandGitHub	d65acd83d8	[Model] Support llm-compressor Inkling NVFP4 weights (#49258 ) Signed-off-by: mgoin <mgoin64@gmail.com>	2026-07-23 20:06:01 -07:00
amd-sourjya GitHub amd-sourjya Cursor Andreas Karatzas	80c9d5d5e0	[ROCm][Quantization] Add Quark W4A8 (INT4-FP8) MoE CI coverage (#48050 ) Signed-off-by: amd-sourjya <amd-sourjya@users.noreply.github.com> Co-authored-by: amd-sourjya <amd-sourjya@users.noreply.github.com> Co-authored-by: Cursor <cursoragent@cursor.com> Co-authored-by: Andreas Karatzas <akaratza@amd.com>	2026-07-23 21:56:40 -05:00
Euisuh JeongandGitHub	da54a5bf05	[Docs] Fix broken anchor links in serving/pooling/MoE docs (#49654 ) Signed-off-by: euisuh <euisuh.jeong@gmail.com>	2026-07-24 02:42:56 +00:00
Andreas Karatzas GitHub OpenAI Codex	1479bd9e9d	[ROCm][CI] Prepare AMD mirrors for regating (#49270 ) Signed-off-by: Andreas Karatzas <Andreas.Karatzas@amd.com> Signed-off-by: Andreas Karatzas <akaratza@amd.com> Co-authored-by: OpenAI Codex <codex@openai.com>	2026-07-23 18:42:25 -07:00
Chaojun Zhang GitHubmergify[bot] <37929162+mergify[bot]@users.noreply.github.com>	0231dd5467	[BugFix][LoRA] Skip marlin-backend gpt-oss LoRA tests on XPU (#49385 ) Signed-off-by: Chaojun Zhang <chaojun.zhang@intel.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>	2026-07-24 09:16:00 +08:00
Nicolò Lucchesi GitHubmergify[bot] <37929162+mergify[bot]@users.noreply.github.com>	2659467497	[CI][PD] Add hybrid SSM P_TP>D_TP accuracy sweep entry (#49593 ) Signed-off-by: NickLucche <nicolo.lucchesi@mistral.ai> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>	2026-07-23 22:49:53 +01:00