From 90460f10f8ec7c38877fd38116a0b4c75bb7cc11 Mon Sep 17 00:00:00 2001 From: Yichao Zhu Date: Thu, 22 Oct 2026 11:39:47 +0800 Subject: [PATCH 1/6] feat(kimik3): refresh PD sweep on ROCm nightly MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 更新五个 Kimi-K3 PD 点位：吞吐点启用 BF16 SSM，移除 1P1D c56 和 1P3D c12，增加 FP32 SSM 的 1P1D c1 K7 monokernel 点。临时 kernel 补丁单独提交。 --- .../mi355x-fp4/agentx/disagg-variants.yaml | 45 ++++++++----------- inferencex-e2e/configs/amd-master.yaml | 37 ++++----------- inferencex-e2e/perf-changelog.yaml | 7 +++ 3 files changed, 33 insertions(+), 56 deletions(-) diff --git a/inferencex-e2e/benchmarks/multi_node/srt-slurm-recipes/kimik3/vllm/mi355x-fp4/agentx/disagg-variants.yaml b/inferencex-e2e/benchmarks/multi_node/srt-slurm-recipes/kimik3/vllm/mi355x-fp4/agentx/disagg-variants.yaml index 5517332e9f..c024b6879e 100644 --- a/inferencex-e2e/benchmarks/multi_node/srt-slurm-recipes/kimik3/vllm/mi355x-fp4/agentx/disagg-variants.yaml +++ b/inferencex-e2e/benchmarks/multi_node/srt-slurm-recipes/kimik3/vllm/mi355x-fp4/agentx/disagg-variants.yaml @@ -65,7 +65,7 @@ base: enable-prefix-caching: true prefix-match-unit: 128 kv-cache-dtype: fp8 - mamba-ssm-cache-dtype: float32 + mamba-ssm-cache-dtype: bfloat16 block-size: 128 attention-backend: ROCM_AITER_MLA attention-config: '{"mla_prefill_backend":"ROCM_AITER_FA","use_prefill_query_quantization":true}' @@ -118,22 +118,29 @@ base: placement: node: head -override_1p3d_c12: +override_1p1d_c1: roles: prefill: args: - max-num-seqs: 24 + decode-context-parallel-size: 8 + mamba-ssm-cache-dtype: float32 + max-num-seqs: 1 max-num-batched-tokens: 8192 compilation-config: '{"mode":3,"cudagraph_mode":"NONE","max_cudagraph_capture_size":0,"custom_ops":["+fused_rms_norm_gated","+quant_fp8","+grouped_topk","+sparse_attn_indexer","none"]}' - env: - HSA_NO_SCRATCH_RECLAIM: '0' + speculative-config: &spec_k7 '{"model":"/models/Inferact-Kimi-K3-DSpark","method":"dspark","num_speculative_tokens":7,"attention_backend":"ROCM_AITER_MLA","kv_cache_dtype":"fp8","draft_sample_method":"probabilistic","rejection_sample_method":"block"}' + kv-transfer-config: '{"kv_connector":"MoRIIOConnector","kv_role":"kv_producer","kv_load_failure_policy":"fail","kv_connector_extra_config":{"backend":"rdma"}}' decode: - nodes: 3 - workers: 3 args: - max-num-seqs: 8 + decode-context-parallel-size: 8 + mamba-ssm-cache-dtype: float32 + max-num-seqs: 1 max-num-batched-tokens: 512 - compilation-config: '{"mode":3,"cudagraph_mode":"FULL_DECODE_ONLY","max_cudagraph_capture_size":40,"cudagraph_capture_sizes":[1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19,20,21,22,23,24,25,26,27,28,29,30,31,32,33,34,35,36,37,38,39,40],"custom_ops":["+fused_rms_norm_gated","+quant_fp8","+grouped_topk","+sparse_attn_indexer","none"]}' + additional-config: '{"decode_backend":"latency_specialized"}' + speculative-config: *spec_k7 + # K7 verifies eight tokens for the single active request. + compilation-config: '{"mode":3,"cudagraph_mode":"FULL_DECODE_ONLY","max_cudagraph_capture_size":8,"cudagraph_capture_sizes":[1,8],"custom_ops":["+fused_rms_norm_gated","+quant_fp8","+grouped_topk","+sparse_attn_indexer","none"]}' + env: + VLLM_USE_V2_MODEL_RUNNER: '1' override_1p2d_c24: &tp8_c24 roles: &tp8_c24_roles @@ -167,8 +174,8 @@ override_1p3d_c24: override_1p1d_c48: &tp8_c48 roles: &tp8_c48_roles - prefill: &tp8_c48_prefill - args: &tp8_c48_prefill_args + prefill: + args: decode-context-parallel-size: 8 max-num-seqs: 96 max-num-batched-tokens: 8192 @@ -183,22 +190,6 @@ override_1p1d_c48: &tp8_c48 max-num-batched-tokens: 512 compilation-config: '{"mode":3,"cudagraph_mode":"FULL_DECODE_ONLY","max_cudagraph_capture_size":240,"cudagraph_capture_sizes":[1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19,20,21,22,23,24,25,26,27,28,29,30,31,32,33,34,35,36,37,38,39,40,41,42,43,44,45,46,47,48,49,50,51,52,53,54,55,56,57,58,59,60,61,62,63,64,65,66,67,68,69,70,71,72,73,74,75,76,77,78,79,80,81,82,83,84,85,86,87,88,89,90,91,92,93,94,95,96,97,98,99,100,101,102,103,104,105,106,107,108,109,110,111,112,113,114,115,116,117,118,119,120,121,122,123,124,125,126,127,128,130,140,150,160,170,180,190,200,210,220,230,240],"custom_ops":["+fused_rms_norm_gated","+quant_fp8","+grouped_topk","+sparse_attn_indexer","none"]}' -override_1p1d_c56: - <<: *tp8_c48 - roles: - <<: *tp8_c48_roles - prefill: - <<: *tp8_c48_prefill - args: - <<: *tp8_c48_prefill_args - max-num-seqs: 112 - decode: - <<: *tp8_c48_decode - args: - <<: *tp8_c48_decode_args - max-num-seqs: 112 - compilation-config: '{"mode":3,"cudagraph_mode":"FULL_DECODE_ONLY","max_cudagraph_capture_size":280,"cudagraph_capture_sizes":[1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19,20,21,22,23,24,25,26,27,28,29,30,31,32,33,34,35,36,37,38,39,40,41,42,43,44,45,46,47,48,49,50,51,52,53,54,55,56,57,58,59,60,61,62,63,64,65,66,67,68,69,70,71,72,73,74,75,76,77,78,79,80,81,82,83,84,85,86,87,88,89,90,91,92,93,94,95,96,97,98,99,100,101,102,103,104,105,106,107,108,109,110,111,112,113,114,115,116,117,118,119,120,121,122,123,124,125,126,127,128,130,140,150,160,170,180,190,200,210,220,230,240,250,260,270,280],"custom_ops":["+fused_rms_norm_gated","+quant_fp8","+grouped_topk","+sparse_attn_indexer","none"]}' - override_1p2d_c48: <<: *tp8_c48 roles: diff --git a/inferencex-e2e/configs/amd-master.yaml b/inferencex-e2e/configs/amd-master.yaml index ef46e58983..ae3eddfa3d 100644 --- a/inferencex-e2e/configs/amd-master.yaml +++ b/inferencex-e2e/configs/amd-master.yaml @@ -1682,7 +1682,7 @@ musespark-fp4-mi355x-sglang-agentic-dflash: # Native Kimi-K3 PD with TP8 workers and MoRIIO READ. kimik3-fp4-mi355x-vllm-disagg-agentic: - image: vllm/vllm-openai-rocm:nightly-81198e97ba7eee2a22540caaa756b7fdddcb4d93@sha256:a401e4f46872dcde77e07ae3d7fa2d3a71b7f2773714316116532e2b39deb179 + image: vllm/vllm-openai-rocm:nightly model: moonshotai/Kimi-K3 model-prefix: kimik3 runner: cluster:mi355x-amds @@ -1697,20 +1697,6 @@ kimik3-fp4-mi355x-vllm-disagg-agentic: agentic-coding: - dram-utilization: 0.6 search-space: - - spec-decoding: mtp - conc-list: [12] - kv-offloading: dram - kv-offload-backend: { name: vllm-simple } - srt-recipe: disagg-variants.yaml:override_1p3d_c12 - prefill: - num-worker: 1 - tp: 8 - dcp-size: 1 - ep: 1 - dp-attn: false - additional-settings: - - TOTAL_CPU_DRAM_GB=1799 - decode: { num-worker: 3, tp: 8, dcp-size: 1, ep: 1, dp-attn: false } - spec-decoding: mtp conc-list: [24] kv-offloading: dram @@ -1754,20 +1740,6 @@ kimik3-fp4-mi355x-vllm-disagg-agentic: # Recipe enables P-only scratch reclaim; D and GMU0.90 are unchanged. - TOTAL_CPU_DRAM_GB=1799 decode: { num-worker: 1, tp: 8, dcp-size: 8, ep: 1, dp-attn: false } - - spec-decoding: mtp - conc-list: [56] - kv-offloading: dram - kv-offload-backend: { name: vllm-simple } - srt-recipe: disagg-variants.yaml:override_1p1d_c56 - prefill: - num-worker: 1 - tp: 8 - dcp-size: 8 - ep: 1 - dp-attn: false - additional-settings: - - TOTAL_CPU_DRAM_GB=1799 - decode: { num-worker: 1, tp: 8, dcp-size: 8, ep: 1, dp-attn: false } - spec-decoding: mtp conc-list: [48] kv-offloading: dram @@ -1783,3 +1755,10 @@ kimik3-fp4-mi355x-vllm-disagg-agentic: # Inherits the same P-only scratch-reclaim setting as 1P1D c48. - TOTAL_CPU_DRAM_GB=1799 decode: { num-worker: 2, tp: 8, dcp-size: 8, ep: 1, dp-attn: false } + - search-space: + - spec-decoding: mtp + conc-list: [1] + kv-offloading: none + srt-recipe: disagg-variants.yaml:override_1p1d_c1 + prefill: { num-worker: 1, tp: 8, dcp-size: 8, ep: 1, dp-attn: false } + decode: { num-worker: 1, tp: 8, dcp-size: 8, ep: 1, dp-attn: false } diff --git a/inferencex-e2e/perf-changelog.yaml b/inferencex-e2e/perf-changelog.yaml index a50995ce2b..fa3b27dcd1 100644 --- a/inferencex-e2e/perf-changelog.yaml +++ b/inferencex-e2e/perf-changelog.yaml @@ -11100,3 +11100,10 @@ - "Search space: Balanced TP4 at conc 1-32; TP8 + EP8 with attention DP2 at conc 64-128." pr-link: https://github.com/SemiAnalysisAI/InferenceX/pull/4466 scenario-type: [agentic-coding] + +- config-keys: + - kimik3-fp4-mi355x-vllm-disagg-agentic + description: + - "Use the ROCm nightly image and BF16 SSM states for the throughput points." + - "Replace 1P1D c56 and 1P3D c12 with 1P1D c1, K7 specialized decode, and FP32 SSM states." + pr-link: https://github.com/SemiAnalysisAI/InferenceX/pull/XXX From 8b4392e51542d99ef52003d40c4efe98a3680953 Mon Sep 17 00:00:00 2001 From: Yichao Zhu Date: Thu, 22 Oct 2026 11:42:27 +0800 Subject: [PATCH 2/6] build(kimik3): cherry-pick upstream kernel support MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 导入 vLLM #60683、#58743、#60808、#63031 的公开提交，保留原始 SHA 与 cherry-pick 来源。安装器使用运行镜像对应的源码和 ROCm SDK 重建 HIP 扩展，便于 nightly 收录后整体移除。 --- .../configs/kimik3-nightly-kernels.sh | 73 + .../kimik3-nightly/01-vllm-60683.patch | 6667 +++++++++++++++++ .../kimik3-nightly/02-vllm-58743-state.patch | 155 + .../03-vllm-58743-checkpoints.patch | 251 + .../04-vllm-58743-native-prefill.patch | 512 ++ .../05-vllm-58743-backends.patch | 181 + .../kimik3-nightly/06-vllm-58743-tests.patch | 115 + .../kimik3-nightly/07-vllm-60808.patch | 569 ++ .../kimik3-nightly/08-vllm-63031.patch | 172 + .../mi355x-fp4/agentx/disagg-variants.yaml | 3 + inferencex-e2e/perf-changelog.yaml | 3 +- 11 files changed, 8700 insertions(+), 1 deletion(-) create mode 100644 inferencex-e2e/benchmarks/multi_node/srt-slurm-recipes/configs/kimik3-nightly-kernels.sh create mode 100644 inferencex-e2e/benchmarks/multi_node/srt-slurm-recipes/configs/patches/kimik3-nightly/01-vllm-60683.patch create mode 100644 inferencex-e2e/benchmarks/multi_node/srt-slurm-recipes/configs/patches/kimik3-nightly/02-vllm-58743-state.patch create mode 100644 inferencex-e2e/benchmarks/multi_node/srt-slurm-recipes/configs/patches/kimik3-nightly/03-vllm-58743-checkpoints.patch create mode 100644 inferencex-e2e/benchmarks/multi_node/srt-slurm-recipes/configs/patches/kimik3-nightly/04-vllm-58743-native-prefill.patch create mode 100644 inferencex-e2e/benchmarks/multi_node/srt-slurm-recipes/configs/patches/kimik3-nightly/05-vllm-58743-backends.patch create mode 100644 inferencex-e2e/benchmarks/multi_node/srt-slurm-recipes/configs/patches/kimik3-nightly/06-vllm-58743-tests.patch create mode 100644 inferencex-e2e/benchmarks/multi_node/srt-slurm-recipes/configs/patches/kimik3-nightly/07-vllm-60808.patch create mode 100644 inferencex-e2e/benchmarks/multi_node/srt-slurm-recipes/configs/patches/kimik3-nightly/08-vllm-63031.patch diff --git a/inferencex-e2e/benchmarks/multi_node/srt-slurm-recipes/configs/kimik3-nightly-kernels.sh b/inferencex-e2e/benchmarks/multi_node/srt-slurm-recipes/configs/kimik3-nightly-kernels.sh new file mode 100644 index 0000000000..9049ace8b4 --- /dev/null +++ b/inferencex-e2e/benchmarks/multi_node/srt-slurm-recipes/configs/kimik3-nightly-kernels.sh @@ -0,0 +1,73 @@ +#!/usr/bin/env bash +set -eo pipefail + +# The router uses a separate image without vLLM. +if ! python3 -c 'import importlib.metadata; importlib.metadata.version("vllm")' >/dev/null 2>&1; then + exit 0 +fi +source /infmax-workspace/benchmarks/check_env.sh +check_env_vars KIMIK3_KERNEL_BUILD_JOBS ROCM_PATH +[[ "$KIMIK3_KERNEL_BUILD_JOBS" =~ ^[1-9][0-9]*$ ]] + +package=$(python3 -c 'import importlib.util; print(next(iter(importlib.util.find_spec("vllm").submodule_search_locations)))') +version=$(python3 -c 'import importlib.metadata; print(importlib.metadata.version("vllm"))') +revision=$(python3 -c 'import re, sys; match = re.search(r"\+g([0-9a-f]{7,40})(?:\.|$)", sys.argv[1]); assert match, f"Cannot resolve nightly source: {sys.argv[1]}"; print(match[1])' "$version") +patch_dir=$(cd "$(dirname "$0")/patches/kimik3-nightly" && pwd) +patches=("$patch_dir"/*.patch) +[[ -f "${patches[0]}" ]] +signature=$({ printf '%s\n' "$version"; sha256sum "${patches[@]}"; } | sha256sum | cut -d' ' -f1) +marker="$package/.infx-kimik3-kernels" +exec 9> /tmp/infx-kimik3-kernels.lock +flock 9 +if [[ -f "$marker" && $(< "$marker") == "$signature" ]]; then + echo "Kimi-K3 kernel patches already installed: $signature" + exit 0 +fi + +build_root=$(mktemp -d /tmp/infx-kimik3-kernels.XXXXXX) +curl --fail --location --retry 3 --max-time 300 \ + "https://codeload.github.com/vllm-project/vllm/tar.gz/$revision" \ + -o "$build_root/source.tgz" +mkdir "$build_root/source" +tar -xzf "$build_root/source.tgz" --strip-components=1 -C "$build_root/source" +git -C "$build_root/source" init -q +applied=() +for patch in "${patches[@]}"; do + if git -C "$build_root/source" apply --reverse --check "$patch" 2>/dev/null; then + echo "Already included in $version: $(basename "$patch")" + else + git -C "$build_root/source" apply --check "$patch" + git -C "$build_root/source" apply "$patch" + applied+=("$patch") + fi +done + +if (( ${#applied[@]} )); then + cp "$package"/*.so "$build_root/source/vllm/" + cp "$package/_version.py" "$build_root/source/vllm/_version.py" + uv venv --system-site-packages --python python3 "$build_root/venv" + uv pip install --python "$build_root/venv/bin/python" cmake ninja + export PATH="$build_root/venv/bin:$PATH" + cmake -S "$build_root/source" -B "$build_root/build" -G Ninja \ + -DVLLM_TARGET_DEVICE=rocm \ + -DVLLM_PYTHON_EXECUTABLE="$build_root/venv/bin/python" \ + -DCMAKE_BUILD_TYPE=RelWithDebInfo \ + -DCMAKE_MODULE_LINKER_FLAGS=-fuse-ld=lld \ + -DCMAKE_SHARED_LINKER_FLAGS=-fuse-ld=lld \ + -DCMAKE_INSTALL_PREFIX="$build_root/stage" \ + -DVLLM_GPU_ARCHES=gfx950 -DCMAKE_HIP_ARCHITECTURES=gfx950 \ + -DROCM_PATH="$ROCM_PATH" + cmake --build "$build_root/build" --target _C_stable_libtorch -j "$KIMIK3_KERNEL_BUILD_JOBS" + cmake --install "$build_root/build" --component _C_stable_libtorch + install -m 755 "$build_root/stage/vllm/_C_stable_libtorch.abi3.so" "$package/" + for patch in "${applied[@]}"; do + while IFS=$'\t' read -r additions deletions path; do + if [[ "$path" == vllm/*.py ]]; then + install -D -m 644 "$build_root/source/$path" "$package/${path#vllm/}" + fi + done < <(git apply --numstat "$patch") + done + python3 -c 'import torch; import vllm._C_stable_libtorch' +fi +printf '%s\n' "$signature" > "$marker" +echo "Kimi-K3 kernel patches installed on $version; build evidence: $build_root" diff --git a/inferencex-e2e/benchmarks/multi_node/srt-slurm-recipes/configs/patches/kimik3-nightly/01-vllm-60683.patch b/inferencex-e2e/benchmarks/multi_node/srt-slurm-recipes/configs/patches/kimik3-nightly/01-vllm-60683.patch new file mode 100644 index 0000000000..9f72da7fb6 --- /dev/null +++ b/inferencex-e2e/benchmarks/multi_node/srt-slurm-recipes/configs/patches/kimik3-nightly/01-vllm-60683.patch @@ -0,0 +1,6667 @@ +From 59781f807cc71208ceac80d8eb43e264dbe35f3f Mon Sep 17 00:00:00 2001 +From: Hyukjoon Lee +Date: Wed, 21 Oct 2026 15:53:29 +0900 +Subject: [PATCH] [ROCm][Perf][Kimi-K3] Specialized decode for Kimi-K3 on + MI355X (#60683) + +Signed-off-by: hyukjlee +Co-authored-by: Claude Opus 5.5 +Co-authored-by: vllm-agent +(cherry picked from commit a9bf45ed8a856a0b1839197152d0d26fabb6599c) +--- + .buildkite/test_areas/models_basic.yaml | 15 + + .../amd_specialized_decode/__init__.py | 2 + + .../test_engine_fault.py | 153 +++ + .../amd_specialized_decode/test_numerics.py | 575 +++++++++ + .../amd_specialized_decode/test_runtime.py | 318 +++++ + vllm/models/kimi_k3/amd/linear.py | 27 +- + vllm/models/kimi_k3/amd/mla.py | 3 + + vllm/models/kimi_k3/amd/model.py | 3 + + .../kimi_k3/amd/specialized_decode/README.md | 108 ++ + .../amd/specialized_decode/__init__.py | 12 + + .../kimi_k3/amd/specialized_decode/hooks.py | 63 + + .../amd/specialized_decode/model_state.py | 27 + + .../amd/specialized_decode/mono/__init__.py | 2 + + .../amd/specialized_decode/mono/abi.py | 26 + + .../mono/attention/__init__.py | 2 + + .../specialized_decode/mono/attention/back.py | 178 +++ + .../specialized_decode/mono/attention/kda.py | 1021 ++++++++++++++++ + .../mono/common/__init__.py | 2 + + .../specialized_decode/mono/common/debug.py | 38 + + .../amd/specialized_decode/mono/common/mx.py | 84 ++ + .../amd/specialized_decode/mono/common/ops.py | 435 +++++++ + .../specialized_decode/mono/common/plan.py | 237 ++++ + .../specialized_decode/mono/common/ranks.py | 33 + + .../specialized_decode/mono/common/sync.py | 255 ++++ + .../amd/specialized_decode/mono/layer.py | 573 +++++++++ + .../mono/stages/__init__.py | 2 + + .../specialized_decode/mono/stages/gemv.py | 126 ++ + .../amd/specialized_decode/mono/stages/moe.py | 1025 +++++++++++++++++ + .../specialized_decode/runtime/__init__.py | 219 ++++ + .../amd/specialized_decode/runtime/decode.py | 198 ++++ + .../amd/specialized_decode/runtime/guards.py | 154 +++ + .../amd/specialized_decode/runtime/layers.py | 221 ++++ + .../amd/specialized_decode/runtime/peer.py | 112 ++ + .../amd/specialized_decode/runtime/runner.py | 76 ++ + 34 files changed, 6318 insertions(+), 7 deletions(-) + create mode 100644 tests/models/kimi_k3/amd_specialized_decode/__init__.py + create mode 100644 tests/models/kimi_k3/amd_specialized_decode/test_engine_fault.py + create mode 100644 tests/models/kimi_k3/amd_specialized_decode/test_numerics.py + create mode 100644 tests/models/kimi_k3/amd_specialized_decode/test_runtime.py + create mode 100644 vllm/models/kimi_k3/amd/specialized_decode/README.md + create mode 100644 vllm/models/kimi_k3/amd/specialized_decode/__init__.py + create mode 100644 vllm/models/kimi_k3/amd/specialized_decode/hooks.py + create mode 100644 vllm/models/kimi_k3/amd/specialized_decode/model_state.py + create mode 100644 vllm/models/kimi_k3/amd/specialized_decode/mono/__init__.py + create mode 100644 vllm/models/kimi_k3/amd/specialized_decode/mono/abi.py + create mode 100644 vllm/models/kimi_k3/amd/specialized_decode/mono/attention/__init__.py + create mode 100644 vllm/models/kimi_k3/amd/specialized_decode/mono/attention/back.py + create mode 100644 vllm/models/kimi_k3/amd/specialized_decode/mono/attention/kda.py + create mode 100644 vllm/models/kimi_k3/amd/specialized_decode/mono/common/__init__.py + create mode 100644 vllm/models/kimi_k3/amd/specialized_decode/mono/common/debug.py + create mode 100644 vllm/models/kimi_k3/amd/specialized_decode/mono/common/mx.py + create mode 100644 vllm/models/kimi_k3/amd/specialized_decode/mono/common/ops.py + create mode 100644 vllm/models/kimi_k3/amd/specialized_decode/mono/common/plan.py + create mode 100644 vllm/models/kimi_k3/amd/specialized_decode/mono/common/ranks.py + create mode 100644 vllm/models/kimi_k3/amd/specialized_decode/mono/common/sync.py + create mode 100644 vllm/models/kimi_k3/amd/specialized_decode/mono/layer.py + create mode 100644 vllm/models/kimi_k3/amd/specialized_decode/mono/stages/__init__.py + create mode 100644 vllm/models/kimi_k3/amd/specialized_decode/mono/stages/gemv.py + create mode 100644 vllm/models/kimi_k3/amd/specialized_decode/mono/stages/moe.py + create mode 100644 vllm/models/kimi_k3/amd/specialized_decode/runtime/__init__.py + create mode 100644 vllm/models/kimi_k3/amd/specialized_decode/runtime/decode.py + create mode 100644 vllm/models/kimi_k3/amd/specialized_decode/runtime/guards.py + create mode 100644 vllm/models/kimi_k3/amd/specialized_decode/runtime/layers.py + create mode 100644 vllm/models/kimi_k3/amd/specialized_decode/runtime/peer.py + create mode 100644 vllm/models/kimi_k3/amd/specialized_decode/runtime/runner.py + +diff --git a/.buildkite/test_areas/models_basic.yaml b/.buildkite/test_areas/models_basic.yaml +index cfcc351a5..80c95e95b 100644 +--- a/.buildkite/test_areas/models_basic.yaml ++++ b/.buildkite/test_areas/models_basic.yaml +@@ -324,3 +324,18 @@ steps: + commands: + - pytest -v -s tests/models/test_dsv41_mono_numerics.py + - pytest -v -s tests/models/test_dsv41_mono_layers.py ++ ++- label: ':amd: (MI355) Kimi-K3 Specialized Decode' ++ key: kimi-k3-specialized-decode-amd ++ dind: false ++ device: mi355_8 ++ num_devices: 8 ++ timeout_in_minutes: 60 ++ working_dir: /vllm-workspace/ ++ depends_on: ++ - image-build-amd ++ source_file_dependencies: ++ - vllm/models/kimi_k3/amd/ ++ - tests/models/kimi_k3/amd_specialized_decode/ ++ commands: ++ - pytest -v -s tests/models/kimi_k3/amd_specialized_decode +diff --git a/tests/models/kimi_k3/amd_specialized_decode/__init__.py b/tests/models/kimi_k3/amd_specialized_decode/__init__.py +new file mode 100644 +index 000000000..208f01a7c +--- /dev/null ++++ b/tests/models/kimi_k3/amd_specialized_decode/__init__.py +@@ -0,0 +1,2 @@ ++# SPDX-License-Identifier: Apache-2.0 ++# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +diff --git a/tests/models/kimi_k3/amd_specialized_decode/test_engine_fault.py b/tests/models/kimi_k3/amd_specialized_decode/test_engine_fault.py +new file mode 100644 +index 000000000..485dba75f +--- /dev/null ++++ b/tests/models/kimi_k3/amd_specialized_decode/test_engine_fault.py +@@ -0,0 +1,153 @@ ++# SPDX-License-Identifier: Apache-2.0 ++# SPDX-FileCopyrightText: Copyright contributors to the vLLM project ++"""Fail-stop end to end: a TP rank that stalls inside a served Kimi-K3 ++specialized decode step fails a request, and the engine stops instead of ++hanging: every TP rank must stop at the same step (``guards.FailStop``), or the ++ranks that went on would wait in the next collective for ever.""" ++ ++import json ++import os ++import time ++ ++import pytest ++import requests ++ ++from tests.models.kimi_k3.amd_specialized_decode.test_numerics import needs_gfx950 ++from tests.utils import RemoteOpenAIServer, multi_gpu_test ++ ++MODEL = "riverclouds/Kimi-K3-8L-dummy" ++DRAFT = "riverclouds/Kimi-K3-DSpark-for-8L-dummy" ++STALLED = 3 ++FAULT_EPOCH = 256 # a mono decode step well past warm-up and capture ++STARTUP_TIMEOUT = 1800 ++FAULT_DEADLINE = 900 ++EXIT_TIMEOUT = 300 ++ ++# Builds every rank's K2 with the fault hook (``K2Build.fault_*``): CTA 0 of ++# rank STALLED stalls past every in-kernel wait's bound in the step at ++# FAULT_EPOCH. Patched into ``MonoRunner.fault`` as the runner module loads, in ++# every vLLM process. ++_SITECUSTOMIZE = """\ ++import importlib.abc ++import importlib.machinery ++import sys ++ ++_MODULE = "vllm.models.kimi_k3.amd.specialized_decode.runtime.runner" ++ ++ ++def _patch(module): ++ def fault(runner): ++ return {{ ++ "fault_rank": {rank}, ++ "fault_epoch": {epoch}, ++ "fault_layer": 128 + runner.first_layer, # the step's first K2 ++ "fault_ticks": {ticks}, ++ }} ++ ++ module.MonoRunner.fault = property(fault, lambda runner, value: None) ++ ++ ++class _Finder(importlib.abc.MetaPathFinder): ++ def find_spec(self, name, path, target=None): ++ if name != _MODULE: ++ return None ++ spec = importlib.machinery.PathFinder.find_spec(name, path) ++ exec_module = spec.loader.exec_module ++ ++ def patched_exec(module): ++ exec_module(module) ++ _patch(module) ++ ++ spec.loader.exec_module = patched_exec ++ return spec ++ ++ ++sys.meta_path.insert(0, _Finder()) ++""" ++ ++SERVER_ARGS = [ ++ "--trust-remote-code", ++ "--language-model-only", ++ "--load-format", ++ "dummy", ++ "--tensor-parallel-size", ++ "8", ++ "--decode-context-parallel-size", ++ "8", ++ "--dcp-comm-backend", ++ "a2a", ++ "--kv-cache-dtype", ++ "fp8", ++ "--max-model-len", ++ "8192", ++ "--max-num-seqs", ++ "2", ++ "--attention-config", ++ json.dumps({"mla_prefill_backend": "ROCM_AITER_FA"}), ++ "--speculative-config", ++ json.dumps( ++ { ++ "method": "dspark", ++ "model": DRAFT, ++ "num_speculative_tokens": 7, ++ "attention_backend": "ROCM_AITER_MLA", ++ "draft_sample_method": "probabilistic", ++ "rejection_sample_method": "block", ++ } ++ ), ++ "--compilation-config", ++ json.dumps({"cudagraph_mode": "FULL_AND_PIECEWISE"}), ++ "--additional-config", ++ json.dumps({"decode_backend": "latency_specialized"}), ++] ++ ++ ++@needs_gfx950 ++@multi_gpu_test(num_gpus=8) ++def test_a_stalled_rank_stops_the_engine(tmp_path): ++ from vllm.models.kimi_k3.amd.specialized_decode.mono.common.sync import WAIT_TICKS ++ ++ site = tmp_path / "mono_fault" ++ site.mkdir() ++ (site / "sitecustomize.py").write_text( ++ _SITECUSTOMIZE.format( ++ rank=STALLED, epoch=FAULT_EPOCH, ticks=WAIT_TICKS * 3 // 2 ++ ) ++ ) ++ env = { ++ "PYTHONPATH": os.pathsep.join( ++ filter(None, [str(site), os.getenv("PYTHONPATH")]) ++ ), ++ "VLLM_USE_V2_MODEL_RUNNER": "1", ++ "VLLM_ROCM_USE_AITER": "1", ++ "VLLM_ROCM_AITER_MLA_ASM_PADDING": "asm", ++ } ++ body = { ++ "model": MODEL, ++ "prompt": "hello " * 64, ++ "max_tokens": 128, ++ "temperature": 0, ++ "ignore_eos": True, ++ } ++ with RemoteOpenAIServer( ++ MODEL, SERVER_ARGS, env_dict=env, max_wait_seconds=STARTUP_TIMEOUT ++ ) as server: ++ url = server.url_root ++ served = 0 ++ deadline = time.monotonic() + FAULT_DEADLINE ++ while time.monotonic() < deadline: ++ try: ++ r = requests.post(f"{url}/v1/completions", json=body, timeout=300) ++ except requests.ConnectionError: ++ break ++ if r.status_code != 200: ++ break ++ served += 1 ++ else: ++ # the stall only exists inside the kernels: also catches a decode ++ # that never ran them ++ pytest.fail(f"no request failed in {FAULT_DEADLINE}s: the fault is lost") ++ ++ assert served > 0, "the engine failed before the injected step" ++ # the failed step stops the engine and the server exits ++ assert server.proc.wait(timeout=EXIT_TIMEOUT) is not None +diff --git a/tests/models/kimi_k3/amd_specialized_decode/test_numerics.py b/tests/models/kimi_k3/amd_specialized_decode/test_numerics.py +new file mode 100644 +index 000000000..dc396b099 +--- /dev/null ++++ b/tests/models/kimi_k3/amd_specialized_decode/test_numerics.py +@@ -0,0 +1,575 @@ ++# SPDX-License-Identifier: Apache-2.0 ++# SPDX-FileCopyrightText: Copyright contributors to the vLLM project ++"""Kimi-K3's MonoKernels on MI355X against the ops vLLM runs for the same ++spec-verify step, random weights in the checkpoint's formats at TP8 shapes: ++ ++- K1 (one GPU): AttnRes -> in_proj -> f_b -> conv -> KDA recurrence -> gated ++ norm, the conv / SSM state updates and a block-write layer's block store. ++- K2 (eight GPUs): o_proj + all-reduce + the MLP AttnRes + the latent MoE. The ++ MoE is compared on the kernel's own MLP input, where routing cannot differ. ++ Then a CUDA graph of K2 replayed 300 times, bit-identical, and a stalled rank ++ (the fault hook): every rank's launch finishes and the output rank's flags ++ fail the step. ++""" ++ ++import pytest ++import ray ++import torch ++import torch.distributed as dist ++import torch.nn.functional as F ++from einops import rearrange ++ ++from tests.utils import ( ++ init_test_distributed_environment, ++ multi_gpu_test, ++ multi_process_parallel, ++) ++from vllm.platforms import current_platform ++ ++bf = torch.bfloat16 ++COS_MIN = 0.9999 ++ ++ ++def _on_gfx950() -> bool: ++ if not current_platform.is_rocm(): ++ return False ++ from vllm.platforms.rocm import on_gfx950 ++ ++ return on_gfx950() ++ ++ ++needs_gfx950 = pytest.mark.skipif( ++ not _on_gfx950(), reason="the K3 mono kernels need gfx950 (MI355X)" ++) ++ ++ ++def _cos(a: torch.Tensor, b: torch.Tensor) -> float: ++ a, b = a.float().reshape(-1), b.float().reshape(-1) ++ return F.cosine_similarity(a, b, dim=0).item() ++ ++ ++def _rel_l2(a: torch.Tensor, b: torch.Tensor) -> float: ++ a, b = a.float(), b.float() ++ return ((a - b).norm() / b.norm()).item() ++ ++ ++def _randn(g, *shape, scale=1.0, dtype=bf): ++ return (torch.randn(*shape, generator=g, device=g.device) * scale).to(dtype) ++ ++ ++# ---------------------------------------------------------------- K1 ++ ++ ++def _kda_inputs(nb: int, nacc: int, dim_first: bool, S: int, slots: int = 16): ++ from vllm.models.kimi_k3.amd.specialized_decode.mono.attention.kda import ( ++ HD, ++ HIDDEN, ++ NH, ++ NPROJ, ++ PROJ, ++ QKV, ++ ) ++ ++ g = torch.Generator(device="cuda").manual_seed(0) ++ sl = 3 + (S - 1) ++ w = { ++ "prefix": _randn(g, S, HIDDEN, scale=2.0), ++ "delta": _randn(g, S, HIDDEN, scale=0.5), ++ "blocks": _randn(g, S, 8, HIDDEN, scale=2.0), ++ "ares_nw": (1 + _randn(g, HIDDEN, scale=0.1, dtype=torch.float32)).to(bf), ++ "ares_qk": _randn(g, HIDDEN, scale=0.05), ++ "in_nw": (1 + _randn(g, HIDDEN, scale=0.1, dtype=torch.float32)).to(bf), ++ "w_in": _randn(g, NPROJ, HIDDEN, scale=0.02), ++ "w_fb": _randn(g, PROJ, HD, scale=0.08), ++ "conv_w": _randn(g, QKV, 4, scale=0.4, dtype=torch.float32), ++ "a_log": _randn(g, NH, scale=0.5, dtype=torch.float32), ++ "dt_bias": _randn(g, PROJ, scale=0.5, dtype=torch.float32), ++ "on_w": (1 + _randn(g, HD, scale=0.1, dtype=torch.float32)).to(bf), ++ } ++ w["w_in"][NPROJ - 4 :].zero_() ++ # vLLM packs a layer's conv and SSM state in one padded page: strided slots ++ pad = _randn(g, slots, NH * HD * HD + 24576, scale=0.05, dtype=torch.float32) ++ w["rstate"] = pad[:, : NH * HD * HD].view(slots, NH, HD, HD) ++ conv = _randn(g, slots, QKV * sl + 4096)[:, : QKV * sl] ++ if dim_first: ++ w["conv_state"] = conv.view(slots, QKV, sl) ++ else: ++ w["conv_state"] = conv.view(slots, sl, QKV).transpose(-1, -2) ++ w["st_idx"] = torch.arange(3, 3 + S, dtype=torch.int32, device="cuda").view(1, S) ++ w["num_acc"] = torch.full((1,), nacc, dtype=torch.int32, device="cuda") ++ return w ++ ++ ++def _kda_reference(w, nb: int, S: int, write_idx: int): ++ from vllm.config import VllmConfig, set_current_vllm_config ++ from vllm.model_executor.layers.mamba.ops.causal_conv1d import ( ++ causal_conv1d_update, ++ ) ++ from vllm.models.kimi_k3.amd.ops.attn_res import attn_res ++ from vllm.models.kimi_k3.amd.ops.third_party.kda import fused_recurrent_kda ++ from vllm.models.kimi_k3.amd.specialized_decode.mono.attention.kda import ( ++ HD, ++ NH, ++ PROJ, ++ QKV, ++ ) ++ from vllm.third_party.flash_linear_attention.ops.kda import FusedRMSNormGated ++ ++ prefix, blocks = w["prefix"].clone(), w["blocks"].clone() ++ x = attn_res( ++ prefix, ++ w["delta"], ++ blocks, ++ w["ares_nw"], ++ w["ares_qk"], ++ w["in_nw"], ++ nb, ++ write_idx, ++ 1e-5, ++ 1e-5, ++ ) ++ proj = F.linear(x, w["w_in"]) ++ mixed_qkv, g2, f_a, beta = proj.split([QKV, PROJ, HD, NH + 4], dim=-1)[:4] ++ g1 = F.linear(f_a, w["w_fb"]) ++ conv_state, rstate = w["conv_state"].clone(), w["rstate"].clone() ++ qsl = torch.tensor([0, S], dtype=torch.int32, device="cuda") ++ conv = causal_conv1d_update( ++ mixed_qkv.contiguous(), ++ conv_state, ++ w["conv_w"], ++ None, ++ activation="silu", ++ conv_state_indices=w["st_idx"][:, 0].contiguous(), ++ num_accepted_tokens=w["num_acc"], ++ query_start_loc=qsl, ++ max_query_len=S, ++ validate_data=False, ++ out=torch.empty_like(mixed_qkv), ++ ) ++ q, k, v = ( ++ rearrange(t, "n (h d) -> 1 n h d", d=HD) for t in conv.split(PROJ, dim=-1) ++ ) ++ core = torch.empty(1, S, NH, HD, dtype=bf, device="cuda") ++ fused_recurrent_kda( ++ q=q, ++ k=k, ++ v=v, ++ raw_g=rearrange(g1, "n (h d) -> 1 n h d", d=HD), ++ raw_beta=beta[:, :NH].unsqueeze(0), ++ A_log=w["a_log"], ++ dt_bias=w["dt_bias"], ++ lower_bound=-5.0, ++ initial_state=rstate, ++ cu_seqlens=qsl, ++ ssm_state_indices=w["st_idx"], ++ num_accepted_tokens=w["num_acc"], ++ uniform_sequence_length=S, ++ out=core, ++ ) ++ with set_current_vllm_config(VllmConfig()): ++ norm = FusedRMSNormGated(HD, activation="sigmoid").cuda() ++ norm.weight.data = w["on_w"].clone() ++ out = norm.forward_cuda(core, rearrange(g2, "n (h d) -> 1 n h d", d=HD)) ++ return { ++ "prefix": prefix, ++ "blocks": blocks, ++ "conv_state": conv_state, ++ "rstate": rstate, ++ "core": out.reshape(S, PROJ), ++ } ++ ++ ++@needs_gfx950 ++@pytest.mark.parametrize("dim_first", [True, False]) ++@pytest.mark.parametrize("nacc", [1, 3]) ++@pytest.mark.parametrize("write", [False, True]) ++def test_k1_matches_vllm_ops(dim_first: bool, nacc: int, write: bool) -> None: ++ """K1 against attn_res -> in_proj -> f_b -> causal_conv1d_update -> ++ fused_recurrent_kda -> FusedRMSNormGated on one spec-verify step (one ++ request, 8 tokens), including the state each step leaves behind.""" ++ from vllm.models.kimi_k3.amd.specialized_decode.mono.attention.kda import ( ++ PROJ, ++ KdaPreBuild, ++ kda_pre, ++ scratch_bytes, ++ ) ++ from vllm.models.kimi_k3.amd.specialized_decode.mono.common.plan import HEALTH_WORDS ++ ++ S, nb = 8, 4 ++ # a block-write layer appends block nb (block_write_idx == prev_valid_blocks) ++ write_idx = nb if write else -1 ++ w = _kda_inputs(nb, nacc, dim_first, S) ++ ref = _kda_reference(w, nb, S, write_idx) ++ ++ key = KdaPreBuild( ++ tokens=S, ++ qlen=S, ++ nblocks=nb, ++ delta=True, ++ state_len=w["conv_state"].size(-1), ++ write_idx=write_idx, ++ ) ++ prefix, blocks = w["prefix"].clone(), w["blocks"].clone() ++ conv_state, rstate = w["conv_state"].clone(), w["rstate"].clone() ++ core = torch.zeros(S, PROJ, dtype=bf, device="cuda") ++ health = torch.zeros(HEALTH_WORDS, dtype=torch.int32, device="cuda") ++ flags = torch.zeros(1, dtype=torch.int32, device="cuda") ++ kda_pre( ++ key, ++ prefix=prefix, ++ delta=w["delta"], ++ blocks=blocks, ++ ares_nw=w["ares_nw"], ++ ares_qk=w["ares_qk"], ++ in_nw=w["in_nw"], ++ w_in=w["w_in"], ++ w_fb=w["w_fb"], ++ conv_w=w["conv_w"], ++ conv_state=conv_state, ++ a_log=w["a_log"], ++ dt_bias=w["dt_bias"], ++ on_w=w["on_w"], ++ rstate=rstate, ++ st_idx=w["st_idx"], ++ num_acc=w["num_acc"], ++ core_out=core, ++ scratch=torch.zeros(scratch_bytes(key), dtype=torch.uint8, device="cuda"), ++ layer=0, ++ epoch=torch.ones(1, dtype=torch.int32, device="cuda"), ++ health=health, ++ alarm=flags.data_ptr(), ++ ) ++ torch.accelerator.synchronize() ++ assert not health.any() and not flags.any() ++ ++ torch.testing.assert_close(prefix, ref["prefix"], atol=0, rtol=0) ++ torch.testing.assert_close(blocks, ref["blocks"], atol=0, rtol=0) ++ # the conv output feeds the recurrence in fp32 here, bf16 in the reference ++ errs = { ++ "conv_state": _rel_l2(conv_state, ref["conv_state"]), ++ "rstate": _rel_l2(rstate, ref["rstate"]), ++ "core": _rel_l2(core, ref["core"]), ++ } ++ assert errs["conv_state"] < 1e-3 and errs["rstate"] < 1e-3, errs ++ assert _cos(core, ref["core"]) > COS_MIN, errs ++ ++ ++# ---------------------------------------------------------------- K2 / MoE ++ ++ ++def _situ(g, u): ++ g, u = g.float(), u.float() ++ a = (2.0 * torch.tanh(g * 0.25)) * (1.0 + torch.tanh(g * 0.5)) ++ return (a * (25.0 * torch.tanh(u * 0.04))).to(bf) ++ ++ ++def _all_reduce_fp32(t: torch.Tensor) -> torch.Tensor: ++ """The ranks' bf16 partials summed in fp32, as vLLM's custom all-reduce.""" ++ t = t.float() ++ dist.all_reduce(t) ++ return t.to(bf) ++ ++ ++def _moe_reference(x, mw, rank: int): ++ """KimiMoE's decode path on ``x``: sigmoid router + bias top-16, the latent ++ down, aiter's MXFP4 experts, the latent all-reduce + RMSNorm, this rank's ++ up_proj rows into the shared experts' output, the final all-reduce.""" ++ from aiter import ActivationType, QuantType ++ from aiter.fused_moe import fused_moe ++ ++ from vllm.models.kimi_k3.amd.specialized_decode.mono.stages.moe import ( ++ SI, ++ TOPK, ++ UP_N, ++ ) ++ ++ sig = torch.sigmoid(torch.mm(x.float(), mw["w_gate"].float().t())) ++ ids = torch.topk(sig + mw["bias"], TOPK, dim=-1, sorted=True).indices ++ wts = sig.gather(-1, ids) ++ wts = wts / wts.sum(-1, keepdim=True) ++ gu = F.linear(x, mw["w_sgu"]) ++ shared = F.linear(_situ(gu[:, :SI], gu[:, SI:]), mw["w_sd"]) ++ y = fused_moe( ++ F.linear(x, mw["w_ld"]), ++ mw["w13"], ++ mw["w2"], ++ wts.float(), ++ ids.to(torch.int32), ++ None, ++ ActivationType.Situv2, ++ QuantType.per_1x32, ++ False, ++ mw["w13s"], ++ mw["w2s"], ++ None, ++ None, ++ dtype=bf, ++ gate_mode="separated", ++ beta=4.0, ++ linear_beta=25.0, ++ ) ++ lat = _all_reduce_fp32(y).float() ++ ln = lat * torch.rsqrt(lat.pow(2).mean(-1, keepdim=True) + 1e-5) * mw["ln_w"] ++ lo = rank * UP_N ++ up = ln.to(bf).float() @ mw["w_up_shard"].float().t() ++ shared[:, lo : lo + UP_N] = (shared[:, lo : lo + UP_N].float() + up).to(bf) ++ return _all_reduce_fp32(shared) ++ ++ ++def _moe_weights(rank: int, device): ++ from aiter.utility.fp4_utils import e8m0_shuffle ++ ++ from vllm._aiter_ops import rocm_aiter_ops ++ from vllm.models.kimi_k3.amd.specialized_decode.mono.stages.moe import ( ++ HIDDEN, ++ LAT, ++ RI, ++ SI, ++ UP_N, ++ E, ++ ) ++ ++ gc = torch.Generator(device=device).manual_seed(7) # replicated weights ++ gr = torch.Generator(device=device).manual_seed(100 + rank) # TP shards ++ ++ def u8(g, *shape, lo=0, hi=256): ++ return torch.randint(lo, hi, shape, generator=g, device=device).to(torch.uint8) ++ ++ fp4, e8 = torch.float4_e2m1fn_x2, torch.float8_e8m0fnu ++ w_up = _randn(gc, HIDDEN, LAT, scale=0.02) ++ return { ++ "w_gate": _randn(gc, E, HIDDEN, scale=0.02), ++ "bias": _randn(gc, E, scale=0.02, dtype=torch.float32), ++ "w_ld": _randn(gc, LAT, HIDDEN, scale=0.02), ++ "ln_w": (1 + _randn(gc, LAT, scale=0.1, dtype=torch.float32)).to(bf), ++ "w_up_shard": w_up[rank * UP_N : (rank + 1) * UP_N].contiguous(), ++ "w_sgu": _randn(gr, 2 * SI, HIDDEN, scale=0.02), ++ "w_sd": _randn(gr, HIDDEN, SI, scale=0.03), ++ "w13": rocm_aiter_ops.shuffle_weight_a16w4( ++ u8(gr, E, 2 * RI, LAT // 2).view(fp4), 16, False ++ ), ++ "w2": rocm_aiter_ops.shuffle_weight_a16w4( ++ u8(gr, E, LAT, RI // 2).view(fp4), 16, False ++ ), ++ "w13s": rocm_aiter_ops.shuffle_scale_a16w4( ++ u8(gr, E, 2 * RI, LAT // 32, lo=118, hi=123).view(e8).view(-1, LAT // 32), ++ E, ++ False, ++ ), ++ "w2s": e8m0_shuffle( ++ u8(gr, E, LAT, RI // 32, lo=118, hi=123).view(e8).view(-1, RI // 32) ++ ), ++ } ++ ++ ++def _moe_args(mw): ++ return { ++ "w_gate": mw["w_gate"], ++ "bias": mw["bias"], ++ "w_ld": mw["w_ld"], ++ "w_sgu": mw["w_sgu"], ++ "w_sd": mw["w_sd"], ++ "w13": mw["w13"], ++ "w13s": mw["w13s"], ++ "w2": mw["w2"], ++ "w2s": mw["w2s"], ++ "ln_w": mw["ln_w"], ++ "w_up": mw["w_up_shard"], ++ } ++ ++ ++def _k2_inputs(rank: int, device): ++ from vllm.models.kimi_k3.amd.specialized_decode.mono.attention.kda import ( ++ HIDDEN, ++ PROJ, ++ ) ++ ++ gc = torch.Generator(device=device).manual_seed(11) ++ gr = torch.Generator(device=device).manual_seed(200 + rank) ++ return { ++ "core": _randn(gr, 8, PROJ), ++ "w_o": _randn(gr, HIDDEN, PROJ, scale=0.02), ++ "prefix": _randn(gc, 8, HIDDEN, scale=2.0), ++ "blocks": _randn(gc, 8, 8, HIDDEN, scale=2.0), ++ "ares_nw": (1 + _randn(gc, HIDDEN, scale=0.1, dtype=torch.float32)).to(bf), ++ "ares_qk": _randn(gc, HIDDEN, scale=0.05), ++ "in_nw": (1 + _randn(gc, HIDDEN, scale=0.1, dtype=torch.float32)).to(bf), ++ } ++ ++ ++def _k2_launch(runner, w, mw, prefix, out, reset: bool, nb: int = 4): ++ from vllm.models.kimi_k3.amd.specialized_decode.mono import layer as k2 ++ ++ key = k2.K2Build(tokens=8, nblocks=nb + int(reset), reset=reset, **runner.fault) ++ k2.k2_launch( ++ key, ++ core=w["core"], ++ w_o=w["w_o"], ++ prefix=prefix, ++ blocks=w["blocks"], ++ ares_nw=w["ares_nw"], ++ ares_qk=w["ares_qk"], ++ in_nw=w["in_nw"], ++ out=out, ++ layer=3, ++ **_moe_args(mw), ++ **runner.launch_args(), ++ **runner.peer_args(), ++ ) ++ return key ++ ++ ++def _init_worker(monkeypatch, tp_size, pp_size, rank, distributed_init_port): ++ monkeypatch.delenv("CUDA_VISIBLE_DEVICES", raising=False) ++ device = torch.device(f"cuda:{rank}") ++ torch.accelerator.set_device_index(device) ++ init_test_distributed_environment(tp_size, pp_size, rank, distributed_init_port) ++ return device ++ ++ ++@ray.remote(num_gpus=1, max_calls=1) ++def _k2_worker( ++ monkeypatch: pytest.MonkeyPatch, ++ tp_size: int, ++ pp_size: int, ++ rank: int, ++ distributed_init_port: str, ++) -> None: ++ device = _init_worker(monkeypatch, tp_size, pp_size, rank, distributed_init_port) ++ ++ from vllm.models.kimi_k3.amd.ops.attn_res import attn_res ++ from vllm.models.kimi_k3.amd.specialized_decode.mono import layer as k2 ++ from vllm.models.kimi_k3.amd.specialized_decode.mono.attention.kda import HIDDEN ++ from vllm.models.kimi_k3.amd.specialized_decode.runtime.runner import MonoRunner ++ ++ S, nb = 8, 4 ++ runner = MonoRunner(device) ++ mw = _moe_weights(rank, device) ++ w = _k2_inputs(rank, device) ++ ++ for reset in (False, True): ++ runner.step_begin() ++ prefix = w["prefix"].clone() ++ out = torch.empty(S, HIDDEN, dtype=bf, device=device) ++ key = _k2_launch(runner, w, mw, prefix, out, reset) ++ torch.accelerator.synchronize() ++ ++ attn = _all_reduce_fp32(F.linear(w["core"], w["w_o"])) ++ # a block-write layer's prefix restarts at the attention output ++ ref_prefix = attn if reset else w["prefix"].clone() ++ delta = None if reset else attn ++ x_ref = attn_res( ++ ref_prefix, ++ delta, ++ w["blocks"], ++ w["ares_nw"], ++ w["ares_qk"], ++ w["in_nw"], ++ nb + int(reset), ++ -1, ++ 1e-5, ++ 1e-5, ++ ) ++ # the kernel sums the ranks in rank order, NCCL in its own order ++ assert _rel_l2(prefix, ref_prefix) < 1e-2 ++ lay = k2.scratch_layout(key)["xrow"][0] ++ x = runner.scratch[lay : lay + S * HIDDEN * 2].view(bf).view(S, HIDDEN) ++ assert _cos(x, x_ref) > COS_MIN ++ assert _cos(out, _moe_reference(x, mw, rank)) > COS_MIN ++ outs = [torch.empty_like(out) for _ in range(tp_size)] ++ dist.all_gather(outs, out) ++ assert all(torch.equal(outs[0], o) for o in outs), "ranks disagree" ++ ++ # A graph of one step (the epoch bump captured with it), replayed: a ++ # block-write launch writes its prefix, so every replay must match the first. ++ prefix = torch.empty_like(w["prefix"]) ++ out = torch.empty(S, HIDDEN, dtype=bf, device=device) ++ graph = torch.cuda.CUDAGraph() ++ with torch.cuda.graph(graph): ++ runner.step_begin() ++ _k2_launch(runner, w, mw, prefix, out, reset=True) ++ graph.replay() ++ torch.accelerator.synchronize() ++ want_prefix, want_out = prefix.clone(), out.clone() ++ for _ in range(300): ++ graph.replay() ++ torch.accelerator.synchronize() ++ assert torch.equal(prefix, want_prefix) and torch.equal(out, want_out) ++ assert not runner.health.any() and not runner.flags.any() ++ ++ ++@needs_gfx950 ++@multi_gpu_test(num_gpus=8) ++def test_k2_tp8_matches_vllm_ops_and_replays(monkeypatch: pytest.MonkeyPatch) -> None: ++ multi_process_parallel(monkeypatch, 8, 1, _k2_worker) ++ ++ ++STALLED = 3 ++ ++ ++@ray.remote(num_gpus=1, max_calls=1) ++def _fault_worker( ++ monkeypatch: pytest.MonkeyPatch, ++ tp_size: int, ++ pp_size: int, ++ rank: int, ++ distributed_init_port: str, ++) -> None: ++ device = _init_worker(monkeypatch, tp_size, pp_size, rank, distributed_init_port) ++ ++ from vllm.models.kimi_k3.amd.specialized_decode.mono import layer as k2 ++ from vllm.models.kimi_k3.amd.specialized_decode.mono.attention.kda import HIDDEN ++ from vllm.models.kimi_k3.amd.specialized_decode.mono.common.debug import ( ++ given_up_waits, ++ ) ++ from vllm.models.kimi_k3.amd.specialized_decode.mono.common.sync import WAIT_TICKS ++ from vllm.models.kimi_k3.amd.specialized_decode.runtime.runner import MonoRunner ++ ++ runner = MonoRunner(device) ++ # rank STALLED's CTA 0 stalls past every wait's bound in the epoch 2 launch ++ runner.fault = { ++ "fault_rank": STALLED, ++ "fault_epoch": 2, ++ "fault_layer": 3, ++ "fault_ticks": WAIT_TICKS * 3 // 2, ++ } ++ mw = _moe_weights(rank, device) ++ w = _k2_inputs(rank, device) ++ out = torch.empty(8, HIDDEN, dtype=bf, device=device) ++ runner.step_begin() # epoch 1: no stall ++ _k2_launch(runner, w, mw, w["prefix"].clone(), out, reset=True) ++ torch.accelerator.synchronize() ++ assert not runner.flags.any() ++ ++ runner.step_begin() # epoch 2: rank STALLED stalls; every launch ends ++ _k2_launch(runner, w, mw, w["prefix"].clone(), out, reset=True) ++ torch.accelerator.synchronize() ++ waits = given_up_waits(runner.health.cpu(), k2.REGIONS) ++ gave_up = torch.tensor([int(bool(waits))], device=device) ++ dist.all_reduce(gave_up) ++ assert gave_up.item() > 0, "no rank's wait gave up" ++ if rank == 0: ++ assert runner.flags.any(), "the output rank's flags missed the failure" ++ ++ ++@needs_gfx950 ++@multi_gpu_test(num_gpus=8) ++def test_stalled_rank_fails_the_step_without_a_hang( ++ monkeypatch: pytest.MonkeyPatch, ++) -> None: ++ multi_process_parallel(monkeypatch, 8, 1, _fault_worker) ++ ++ ++@needs_gfx950 ++def test_warmup_compiles_without_launching() -> None: ++ """The JIT warm-up compiles K1 and K2 before capture without running them: ++ a launch on its null arguments would fault.""" ++ from vllm.models.kimi_k3.amd.specialized_decode.mono.attention.kda import ( ++ KdaPreBuild, ++ ) ++ from vllm.models.kimi_k3.amd.specialized_decode.mono.layer import K2Build ++ from vllm.models.kimi_k3.amd.specialized_decode.runtime.layers import compile_build ++ ++ compile_build(KdaPreBuild(tokens=8, qlen=8, nblocks=4, delta=True, state_len=10)) ++ compile_build(K2Build(tokens=8, nblocks=4)) ++ torch.accelerator.synchronize() +diff --git a/tests/models/kimi_k3/amd_specialized_decode/test_runtime.py b/tests/models/kimi_k3/amd_specialized_decode/test_runtime.py +new file mode 100644 +index 000000000..6cc5c793c +--- /dev/null ++++ b/tests/models/kimi_k3/amd_specialized_decode/test_runtime.py +@@ -0,0 +1,318 @@ ++# SPDX-License-Identifier: Apache-2.0 ++# SPDX-FileCopyrightText: Copyright contributors to the vLLM project ++"""Kimi-K3's specialized decode on CPU: the switch, its widths and layers, its ++refusals (config, loaded weights, caches) and plan() routing.""" ++ ++from types import SimpleNamespace ++ ++import pytest ++import torch ++ ++from vllm.config import CompilationMode, CUDAGraphMode ++from vllm.models.kimi_k3.amd.specialized_decode import runtime ++from vllm.models.kimi_k3.amd.specialized_decode.runtime import ( ++ FFN, ++ WHOLE, ++ DecodeConfig, ++ layer_kinds, ++ specialized_decode_config, ++) ++from vllm.models.kimi_k3.amd.specialized_decode.runtime import decode as decode_mod ++from vllm.models.kimi_k3.amd.specialized_decode.runtime.decode import ( ++ SpecializedDecode, ++ step_reason, ++) ++from vllm.models.kimi_k3.amd.specialized_decode.runtime.guards import ( ++ cache_reasons, ++ weight_reasons, ++) ++ ++ ++def _vllm_config(**over) -> SimpleNamespace: ++ """A supported config: TP8, DCP8, 7 speculative tokens, FULL graphs, MRV2. ++ ``over`` sets dotted paths, e.g. ``{"parallel_config.data_parallel_size": 2}``.""" ++ spec = SimpleNamespace( ++ num_speculative_tokens=7, ++ enable_adaptive_verification=False, ++ dynamic=False, ++ ) ++ spec.uses_dynamic_speculative_decoding = lambda: spec.dynamic ++ c = SimpleNamespace( ++ additional_config={"decode_backend": "latency_specialized"}, ++ parallel_config=SimpleNamespace( ++ tensor_parallel_size=8, ++ pipeline_parallel_size=1, ++ data_parallel_size=1, ++ enable_expert_parallel=False, ++ enable_eplb=False, ++ decode_context_parallel_size=8, ++ prefill_context_parallel_size=1, ++ use_ubatching=False, ++ ), ++ model_config=SimpleNamespace( ++ dtype=torch.bfloat16, ++ enable_sleep_mode=False, ++ hf_text_config=SimpleNamespace( ++ hidden_size=7168, ++ num_experts=896, ++ attn_res_block_size=12, ++ linear_attn_config={ ++ "num_heads": 96, ++ "head_dim": 128, ++ "short_conv_kernel_size": 4, ++ }, ++ ), ++ ), ++ cache_config=SimpleNamespace( ++ mamba_cache_dtype="auto", use_replayssm=False, use_kda_recoverssm=False ++ ), ++ lora_config=None, ++ kv_transfer_config=None, ++ aux_output_config=SimpleNamespace(enable_return_routed_experts=False), ++ offload_config=SimpleNamespace( ++ uva=SimpleNamespace(cpu_offload_gb=0), ++ prefetch=SimpleNamespace(offload_group_size=0), ++ ), ++ weight_transfer_config=None, ++ speculative_config=spec, ++ use_v2_model_runner=True, ++ compilation_config=SimpleNamespace( ++ mode=CompilationMode.VLLM_COMPILE, ++ cudagraph_mode=CUDAGraphMode.FULL_AND_PIECEWISE, ++ ), ++ ) ++ for path, v in over.items(): ++ obj = c ++ *head, last = path.split(".") ++ for h in head: ++ obj = getattr(obj, h) ++ setattr(obj, last, v) ++ return c ++ ++ ++@pytest.fixture ++def on_mi355x(monkeypatch): ++ monkeypatch.setattr(runtime, "platform_reasons", lambda: []) ++ ++ ++@pytest.mark.parametrize("k, want", [(7, (8,)), (3, (4, 8)), (1, (2, 4, 6, 8))]) ++def test_widths_are_whole_requests_up_to_eight_rows(k, want): ++ assert DecodeConfig(q=k + 1).widths == want ++ ++ ++def _k3_text_config(): ++ """Kimi-K3's layer pattern: 93 layers, MLA every 4th, AttnRes blocks of 12, ++ layer 0 dense.""" ++ mla = {i for i in range(93) if i % 4 == 3} | {92} ++ return SimpleNamespace( ++ num_hidden_layers=93, ++ attn_res_block_size=12, ++ num_experts=896, ++ first_k_dense_replace=1, ++ moe_layer_freq=1, ++ is_kda_layer=lambda i: i not in mla, ++ ) ++ ++ ++def test_kernel_layers_follow_the_k3_pattern(): ++ kinds = layer_kinds(_k3_text_config()) ++ assert list(kinds) == list(range(1, 93)) ++ assert kinds[1] == WHOLE and kinds[12] == WHOLE ++ assert kinds[3] == FFN and kinds[92] == FFN ++ ++ ++def test_auto_runs_the_stock_decode(): ++ """Without the switch nothing is checked, even on an unsupported deployment.""" ++ cfg = _vllm_config( ++ additional_config={}, **{"parallel_config.tensor_parallel_size": 4} ++ ) ++ assert specialized_decode_config(cfg) is None ++ cfg.additional_config = {"decode_backend": "auto"} ++ assert specialized_decode_config(cfg) is None ++ ++ ++def test_an_unknown_backend_is_refused(): ++ cfg = _vllm_config(additional_config={"decode_backend": "fast"}) ++ with pytest.raises(ValueError, match="decode_backend must be"): ++ specialized_decode_config(cfg) ++ ++ ++def test_supported_config_is_accepted(on_mi355x): ++ assert specialized_decode_config(_vllm_config()) == DecodeConfig(q=8) ++ ++ ++@pytest.mark.parametrize( ++ "over, reason", ++ [ ++ ( ++ { ++ "parallel_config.tensor_parallel_size": 4, ++ "parallel_config.decode_context_parallel_size": 4, ++ }, ++ "TP 4", ++ ), ++ ({"parallel_config.pipeline_parallel_size": 2}, "pipeline parallelism"), ++ ({"parallel_config.decode_context_parallel_size": 2}, "decode context"), ++ ({"kv_transfer_config": object()}, "KV transfer connectors"), ++ ({"speculative_config": None}, "without speculative decoding"), ++ ({"speculative_config.num_speculative_tokens": 8}, "more than 7"), ++ ({"speculative_config.enable_adaptive_verification": True}, "adaptive"), ++ ({"speculative_config.dynamic": True}, "dynamic speculative"), ++ ({"cache_config.mamba_cache_dtype": "float32"}, "KDA cache dtypes"), ++ ({"cache_config.use_replayssm": True}, "ReplaySSM"), ++ ({"use_v2_model_runner": False}, "V1 model runner"), ++ ( ++ {"compilation_config.mode": CompilationMode.STOCK_TORCH_COMPILE}, ++ "torch.compile", ++ ), ++ ({"model_config.hf_text_config.hidden_size": 4096}, "model shapes"), ++ ], ++) ++def test_refusals(on_mi355x, over, reason): ++ reasons = runtime.unsupported(_vllm_config(**over)) ++ assert len(reasons) == 1 and reason in reasons[0], reasons ++ ++ ++def test_every_refusal_is_reported_at_once(monkeypatch): ++ monkeypatch.setattr( ++ runtime, "platform_reasons", lambda: ["this GPU (gfx950 / MI355X only)"] ++ ) ++ cfg = _vllm_config( ++ **{"parallel_config.tensor_parallel_size": 4, "use_v2_model_runner": False} ++ ) ++ with pytest.raises(ValueError) as e: ++ specialized_decode_config(cfg) ++ for reason in ("TP 4", "V1 model runner", "this GPU", "decode_backend=auto"): ++ assert reason in str(e.value) ++ ++ ++def test_piecewise_graphs_only_warn(on_mi355x): ++ """They are not refused: plan() runs every PIECEWISE step stock.""" ++ cfg = _vllm_config(**{"compilation_config.cudagraph_mode": CUDAGraphMode.PIECEWISE}) ++ assert runtime.piecewise_graphs_only(cfg) ++ assert specialized_decode_config(cfg) is not None ++ ++ ++def test_cache_refusal_names_the_cache(): ++ pytest.importorskip("flydsl") ++ conv = torch.zeros(2, 4, 3, dtype=torch.float32) ++ rstate = torch.zeros(2, 12, 128, 128, dtype=torch.bfloat16) ++ reasons = cache_reasons({"layers.5.attn": (conv, rstate)}) ++ assert len(reasons) == 2 and all(r.startswith("layers.5.attn:") for r in reasons) ++ ++ ++def _kda_metadata(**md): ++ """One request of 4 tokens unless overridden.""" ++ m = SimpleNamespace( ++ spec_sequence_masks=torch.ones(1, dtype=torch.bool), ++ num_prefills=0, ++ num_decodes=0, ++ num_spec_decodes=1, ++ spec_state_indices_tensor=torch.zeros(1, 4, dtype=torch.int32), ++ uniform_spec_sequence_length=4, ++ ) ++ for k, v in md.items(): ++ setattr(m, k, v) ++ return {"kda": m} ++ ++ ++@pytest.mark.parametrize( ++ "width, md, why", ++ [ ++ (4, {}, None), ++ (8, {"num_spec_decodes": 2}, None), ++ (4, {"num_prefills": 1}, "not pure decode"), ++ (4, {"num_decodes": 1}, "not pure decode"), ++ (4, {"spec_sequence_masks": None}, "not spec-verify"), ++ (8, {}, "not uniform"), ++ (4, {"uniform_spec_sequence_length": None}, "not uniform"), ++ ], ++) ++def test_steps_must_be_pure_uniform_spec_verify(width, md, why): ++ assert step_reason(_kda_metadata(**md), "kda", width) == why ++ ++ ++def test_steps_need_metadata(): ++ assert step_reason(None, "kda", 4) == "no metadata" ++ ++ ++class _Kernel: ++ def launch(self, step, positions, hidden_states, residual, prefix_delta): ++ return "kernel", step.width ++ ++ ++def _decode(monkeypatch, mode=CUDAGraphMode.FULL, md=None): ++ """An executor with kernels on layers 1 and 2 at widths 4 and 8, bound to a ++ KV cache, under a forward context of ``mode``.""" ++ decode = object.__new__(SpecializedDecode) ++ decode.config = DecodeConfig(q=4) ++ decode.runner = SimpleNamespace(kda_prefix="kda") ++ decode.kernels = {w: {1: _Kernel(), 2: _Kernel()} for w in (4, 8)} ++ decode._kv = (1, 2) ++ monkeypatch.setattr(SpecializedDecode, "_kv_key", lambda self: (1, 2)) ++ context = SimpleNamespace( ++ cudagraph_runtime_mode=mode, ++ attn_metadata=_kda_metadata() if md is None else md, ++ ) ++ monkeypatch.setattr(decode_mod, "is_forward_context_available", lambda: True) ++ monkeypatch.setattr(decode_mod, "get_forward_context", lambda: context) ++ return decode ++ ++ ++def test_plan_runs_the_kernels_of_its_width(monkeypatch): ++ step = _decode(monkeypatch).plan(torch.zeros(4), None) ++ assert step is not None and step.width == 4 ++ assert step.layer(1)(None, None, None, None) == ("kernel", 4) ++ assert step.layer(0) is None and step.layer(3) is None ++ ++ ++@pytest.mark.parametrize( ++ "mode, tokens, md, intermediate", ++ [ ++ (CUDAGraphMode.PIECEWISE, 4, None, None), ++ (CUDAGraphMode.FULL, 16, None, None), # no kernels that wide ++ (CUDAGraphMode.FULL, 4, _kda_metadata(num_prefills=1), None), ++ (CUDAGraphMode.FULL, 4, None, object()), # a pipeline stage's input ++ ], ++) ++def test_plan_runs_other_steps_stock(monkeypatch, mode, tokens, md, intermediate): ++ decode = _decode(monkeypatch, mode, md) ++ assert decode.plan(torch.zeros(tokens), intermediate) is None ++ ++ ++def test_plan_runs_stock_before_a_kv_cache_is_bound(monkeypatch): ++ decode = _decode(monkeypatch) ++ decode._kv = None ++ assert decode.plan(torch.zeros(4), None) is None ++ ++ ++def test_plan_refuses_a_kv_cache_bound_behind_its_back(monkeypatch): ++ decode = _decode(monkeypatch) ++ monkeypatch.setattr(SpecializedDecode, "_kv_key", lambda self: (3, 4)) ++ with pytest.raises(RuntimeError, match="without prepare_kv"): ++ decode.plan(torch.zeros(4), None) ++ ++ ++def test_weight_views_reject_every_mismatch(): ++ """weight_reasons names each layer whose weights the kernels cannot read as ++ they are, and a MoE backend that lays them out otherwise.""" ++ pytest.importorskip("flydsl") ++ from vllm.models.kimi_k3.amd.specialized_decode.mono.stages import moe ++ ++ bad_moe = SimpleNamespace( ++ use_latent_moe=True, ++ shared_experts=object(), ++ routed_scaling_factor=2.0, ++ experts=SimpleNamespace(quant_method=SimpleNamespace()), ++ ) ++ no_latent = SimpleNamespace(use_latent_moe=False, shared_experts=None) ++ layers = [None] + [SimpleNamespace(mlp=no_latent)] * 92 ++ layers[5] = SimpleNamespace(mlp=bad_moe) ++ model = SimpleNamespace( ++ config=_k3_text_config(), layers=layers, start_layer=0, end_layer=93 ++ ) ++ reasons = weight_reasons(model) ++ assert "layer 1: needs a latent MoE with shared experts" in reasons ++ assert "layer 5: routed_scaling_factor 2.0" in reasons ++ assert any(r.startswith("layer 5: MoE backend must be") for r in reasons) ++ assert moe.E == 896 +diff --git a/vllm/models/kimi_k3/amd/linear.py b/vllm/models/kimi_k3/amd/linear.py +index 682773189..029636146 100644 +--- a/vllm/models/kimi_k3/amd/linear.py ++++ b/vllm/models/kimi_k3/amd/linear.py +@@ -74,6 +74,10 @@ from vllm.models.kimi_k3.amd.kda import KimiK3DeltaAttention + from vllm.models.kimi_k3.amd.latent_moe_runner import ROCmLatentMoERunner + from vllm.models.kimi_k3.amd.mla import KimiK3MultiHeadLatentAttentionWrapper + from vllm.models.kimi_k3.amd.ops.attn_res import attn_res ++from vllm.models.kimi_k3.amd.specialized_decode.hooks import SpecializedDecodeMixin ++from vllm.models.kimi_k3.amd.specialized_decode.runtime import ( ++ specialized_decode_config, ++) + from vllm.sequence import IntermediateTensors + from vllm.transformers_utils.configs.kimi_linear import KimiLinearConfig + from vllm.utils.math_utils import cdiv +@@ -828,6 +832,7 @@ class KimiLinearModel(nn.Module, EagleModelMixin): + positions: torch.Tensor, + intermediate_tensors: IntermediateTensors | None, + inputs_embeds: torch.Tensor | None = None, ++ step=None, # specialized_decode's plan + **kwargs, + ) -> torch.Tensor | IntermediateTensors | tuple[torch.Tensor, list[torch.Tensor]]: + if get_pp_group().is_first_rank: +@@ -885,11 +890,9 @@ class KimiLinearModel(nn.Module, EagleModelMixin): + self.layers[self.start_layer : self.end_layer], + start=self.start_layer, + ): +- hidden_states, residual, prefix_delta = layer( +- positions=positions, +- hidden_states=hidden_states, +- residual=residual, +- prefix_delta=prefix_delta, ++ kernel = step.layer(layer_idx) if step is not None else None ++ hidden_states, residual, prefix_delta = (kernel or layer)( ++ positions, hidden_states, residual, prefix_delta + ) + if (layer_idx + 1) in self.aux_hidden_state_layers: + self._maybe_add_hidden_state( +@@ -1057,10 +1060,18 @@ class KimiLinearModel(nn.Module, EagleModelMixin): + + + class KimiLinearForCausalLM( +- nn.Module, HasInnerState, SupportsPP, MixtureOfExperts, IsHybrid ++ nn.Module, ++ HasInnerState, ++ SupportsPP, ++ MixtureOfExperts, ++ IsHybrid, ++ SpecializedDecodeMixin, + ): + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): ++ # raises for an unsupported deployment before any weight is allocated ++ decode_config = specialized_decode_config(vllm_config) + super().__init__() ++ self.init_specialized_decode(decode_config) + self.model_config = vllm_config.model_config + self.vllm_config = vllm_config + self.config = self.model_config.hf_config +@@ -1102,8 +1113,9 @@ class KimiLinearForCausalLM( + inputs_embeds: torch.Tensor | None = None, + **kwargs, + ) -> torch.Tensor | IntermediateTensors: ++ step = self.plan_specialized_decode(positions, intermediate_tensors) + hidden_states = self.model( +- input_ids, positions, intermediate_tensors, inputs_embeds, **kwargs ++ input_ids, positions, intermediate_tensors, inputs_embeds, step, **kwargs + ) + return hidden_states + +@@ -1148,6 +1160,7 @@ class KimiLinearForCausalLM( + ) -> torch.Tensor | None: + # The model's final norm is applied here (not at the end of forward) so + # that the pre-norm hidden states can be fed to the MTP draft model. ++ self.after_specialized_decode_step() + hidden_states = self.model.norm(hidden_states, None) + return self.logits_processor(self.lm_head, hidden_states) + +diff --git a/vllm/models/kimi_k3/amd/mla.py b/vllm/models/kimi_k3/amd/mla.py +index 48ea1792b..2db1575a9 100644 +--- a/vllm/models/kimi_k3/amd/mla.py ++++ b/vllm/models/kimi_k3/amd/mla.py +@@ -196,6 +196,7 @@ class KimiK3MultiHeadLatentAttentionWrapper(MultiHeadLatentAttentionWrapper): + positions: torch.Tensor, + hidden_states: torch.Tensor | QuantizedActivation, + llama_4_scaling: torch.Tensor | None = None, ++ skip_o_proj: bool = False, + ) -> torch.Tensor: + q_c = None + num_tokens = ( +@@ -311,4 +312,6 @@ class KimiK3MultiHeadLatentAttentionWrapper(MultiHeadLatentAttentionWrapper): + else: + attn_out = attn_out * gate.sigmoid() + ++ if skip_o_proj: # the specialized decode runs o_proj in its kernel ++ return attn_out + return self.o_proj(attn_out)[0] +diff --git a/vllm/models/kimi_k3/amd/model.py b/vllm/models/kimi_k3/amd/model.py +index 124fca3e3..aacbd086b 100644 +--- a/vllm/models/kimi_k3/amd/model.py ++++ b/vllm/models/kimi_k3/amd/model.py +@@ -220,6 +220,9 @@ class KimiK3ForConditionalGeneration( + def compute_logits(self, hidden_states: torch.Tensor, **kwargs) -> torch.Tensor: + return self.language_model.compute_logits(hidden_states) + ++ def get_model_state_cls(self) -> type: ++ return self.language_model.get_model_state_cls() ++ + def copy_inputs_before_cuda_graphs(self, input_buffers, **kwargs): + return self.language_model.mamba_cache.copy_inputs_before_cuda_graphs( + input_buffers, **kwargs +diff --git a/vllm/models/kimi_k3/amd/specialized_decode/README.md b/vllm/models/kimi_k3/amd/specialized_decode/README.md +new file mode 100644 +index 000000000..3e0e048b8 +--- /dev/null ++++ b/vllm/models/kimi_k3/amd/specialized_decode/README.md +@@ -0,0 +1,108 @@ ++# Kimi-K3 specialized decode (MI355X) ++ ++An opt-in decode path for Kimi-K3 at TP8 on MI355X (gfx950), following RFC #60904. ++A pure spec-verify step of at most 8 rows runs each MoE decoder layer as persistent ++FlyDSL MonoKernels (256 CTAs, all resident), in "layer mode", inside the stock ++layer loop. The TP all-reduces run in-kernel over peer memory. ++ ++- **Whole layers** (KDA): K1 (AttnRes, in_proj, f_b, spec conv, KDA recurrence, ++ gated RMSNorm), then K2. ++- **FFN layers** (MLA): vLLM's AttnRes and MLA up to o_proj's input, then K2. ++- **K2**: o_proj + all-reduce + MLP AttnRes, then the latent MoE (router, ++ top-16, MXFP4 experts, shared experts, latent all-reduce + RMSNorm, up_proj, ++ final all-reduce). ++ ++The kernels read vLLM's loaded weights and KDA states in place. Nothing is copied ++or re-quantized. ++ ++## Enabling ++ ++```bash ++VLLM_USE_V2_MODEL_RUNNER=1 vllm serve -tp 8 \ ++ --speculative-config '{"method": "dspark", "num_speculative_tokens": 7, ...}' \ ++ --additional-config '{"decode_backend": "latency_specialized"}' ++``` ++ ++There are no other options. `decode_backend=auto` (the default) runs vLLM's decode. ++ ++## Supported deployments ++ ++TP8 on MI355X with all 256 CUs, decode context parallel 1 or 8, speculative ++decoding with at most 7 speculative tokens, Model Runner V2, BF16 conv / FP32 ++recurrent KDA states. ++ ++If you request the specialized decode for anything else, startup fails with one ++`ValueError` that lists every unsupported setting: ++ ++- **Config, in the model's `__init__`:** pipeline, data, expert or prefill ++ context parallelism, EPLB, DBO, LoRA, KV connectors, routed-expert capture, ++ sleep mode, CPU weight offload, weight updates, no speculative decoding, ++ adaptive verification, dynamic speculative decoding, ReplaySSM / RecoverSSM, ++ the V1 model runner, `STOCK_TORCH_COMPILE`, another model's shapes, another GPU. ++- **Loaded weights, when the model state builds the executor:** each kernel ++ layer's weight shapes and dtypes, and the MoE backend (`AITER_MXFP4_BF16`, K3 ++ SiTU). ++- **KV caches, at every bind:** the KDA state dtypes and layout. ++ ++Every TP rank agrees on each verdict before anything collective runs. ++ ++## Steps ++ ++`plan()` runs once per step. It picks the kernels for a pure spec-verify step of ++whole uniform requests whose width is one of the built widths: every multiple of ++`num_speculative_tokens + 1` up to 8 rows. Every other step runs vLLM's layers: ++prefill, mixed, wider and PIECEWISE steps. ++ ++Under FULL graphs the plan is baked in at capture. vLLM's capture sizes are not ++changed. A replay with fewer requests than its graph pads whole requests. Those ++requests get the reserved null state slot and one accepted token, and the kernels ++serve them like real ones. For example, at `num_speculative_tokens=1` three ++requests replay the 4-request graph. To capture every width exactly, pass them in ++`--cudagraph-capture-sizes`. ++ ++## Fail-stop ++ ++Every in-kernel wait gives up after 10 s. The rank records the wait in its health ++words (sticky: its later waits exit at once) and stops waiting, so a stalled peer ++fails every rank instead of hanging them. ++ ++Each step, `compute_logits` queues a non-blocking copy of the rank's `FAILED` ++word. Every 16 steps the TP ranks agree on what has landed over the TP CPU group, ++and then every rank raises at the same step, so the engine stops. They must stop ++together: a rank that raised alone would leave the others waiting in the next ++collective (the drafter's) for ever. ++ ++Detection takes up to about 16 steps, so tokens of the failed steps can reach the ++client before the engine stops. A zero-delay check needs a core hook and is left to ++a follow-up. ++ ++`tests/models/kimi_k3/amd_specialized_decode/test_engine_fault.py` stalls a rank in ++a served model and checks that the engine stops. ++ ++## Layout ++ ++Everything lives in this folder. The model files only call in: `linear.py` ++(the config check in `__init__`, `plan_specialized_decode` in `forward`, ++`step.layer(i)` in the layer loop, `after_specialized_decode_step` in ++`compute_logits`), `model.py` (`get_model_state_cls`) and `mla.py` ++(`skip_o_proj`, o_proj's input for the MLA layers' K2). ++ ++```text ++hooks.py SpecializedDecodeMixin: the model-side hooks ++model_state.py KimiK3ModelState: builds the executor after load, binds KV caches ++runtime/ ++ __init__.py DecodeConfig, specialized_decode_config(), unsupported() ++ decode.py SpecializedDecode: plan() -> Step, Step.layer(i), prepare_kv(), ++ after_step() ++ guards.py weight and KV cache refusals, the fail-stop watch ++ layers.py the layer kernels: vLLM tensors -> launches ++ runner.py per-rank launch state: scratch, step epoch, health, peers ++ peer.py peer memory and TP agreement (adapted from ROCm/ATOM, MIT) ++mono/ kernel source ++ layer.py K2 ++ attention/ kda.py (K1), back.py ++ stages/ moe.py, gemv.py ++ common/ device ops, mailbox sync, MX helpers, plan / build keys ++ (adapted from ROCm/ATOM atom/mono, MIT) ++ abi.py the error-word layout ++``` +diff --git a/vllm/models/kimi_k3/amd/specialized_decode/__init__.py b/vllm/models/kimi_k3/amd/specialized_decode/__init__.py +new file mode 100644 +index 000000000..4a88231b5 +--- /dev/null ++++ b/vllm/models/kimi_k3/amd/specialized_decode/__init__.py +@@ -0,0 +1,12 @@ ++# SPDX-License-Identifier: Apache-2.0 ++# SPDX-FileCopyrightText: Copyright contributors to the vLLM project ++"""Kimi-K3 specialized decode for ROCm (MI355X, TP8), layer mode: FlyDSL ++MonoKernels per MoE decoder layer on pure spec-verify steps, inside the stock layer ++loop. Enabled by ``additional_config={"decode_backend": "latency_specialized"}``. ++ ++``hooks.py`` holds the model-side hooks (``SpecializedDecodeMixin``) and ++``model_state.py`` the model runner state that builds the executor. ``runtime/`` is ++the host side: the config check (``runtime/__init__.py``), the executor ++(``runtime/decode.py``), the layer kernels, the per-rank launch state and the later ++refusals. ``mono/`` is the kernel source. Tests live in ++``tests/models/kimi_k3/amd_specialized_decode/``.""" +diff --git a/vllm/models/kimi_k3/amd/specialized_decode/hooks.py b/vllm/models/kimi_k3/amd/specialized_decode/hooks.py +new file mode 100644 +index 000000000..32b448050 +--- /dev/null ++++ b/vllm/models/kimi_k3/amd/specialized_decode/hooks.py +@@ -0,0 +1,63 @@ ++# SPDX-License-Identifier: Apache-2.0 ++# SPDX-FileCopyrightText: Copyright contributors to the vLLM project ++"""The specialized decode's hooks on ``KimiLinearForCausalLM``, kept here so the ++model only calls them: ``init_specialized_decode`` in ``__init__`` (after the config ++check, which runs before any module allocates weights), ``plan_specialized_decode`` ++in ``forward`` and ``after_specialized_decode_step`` in ``compute_logits``. The ++model state (``model_state.KimiK3ModelState``) builds the executor after load.""" ++ ++from typing import TYPE_CHECKING ++ ++import torch ++ ++from vllm.models.kimi_k3.amd.specialized_decode.runtime import DecodeConfig ++ ++if TYPE_CHECKING: ++ from vllm.models.kimi_k3.amd.specialized_decode.runtime.decode import ( ++ SpecializedDecode, ++ Step, ++ ) ++ from vllm.v1.worker.gpu.model_states.interface import ModelState ++ ++ ++class SpecializedDecodeMixin: ++ """``self.model`` is the ``KimiLinearModel`` the executor runs over.""" ++ ++ model: torch.nn.Module ++ specialized_decode_config: DecodeConfig | None ++ specialized_decode: "SpecializedDecode | None" ++ ++ def init_specialized_decode(self, config: DecodeConfig | None) -> None: ++ """``config``: ``specialized_decode_config(vllm_config)``, None for the ++ stock decode.""" ++ self.specialized_decode_config = config ++ self.specialized_decode = None ++ ++ def build_specialized_decode(self) -> None: ++ from vllm.models.kimi_k3.amd.specialized_decode.runtime.decode import ( ++ SpecializedDecode, ++ ) ++ ++ assert self.specialized_decode_config is not None ++ self.specialized_decode = SpecializedDecode( ++ self.model, self.specialized_decode_config ++ ) ++ ++ def get_model_state_cls(self) -> "type[ModelState]": ++ from vllm.models.kimi_k3.amd.specialized_decode.model_state import ( ++ KimiK3ModelState, ++ ) ++ ++ return KimiK3ModelState ++ ++ def plan_specialized_decode( ++ self, positions: torch.Tensor, intermediate_tensors ++ ) -> "Step | None": ++ """This step's plan (``Step.layer(i)`` per decoder layer), None for the ++ stock forward.""" ++ decode = self.specialized_decode ++ return None if decode is None else decode.plan(positions, intermediate_tensors) ++ ++ def after_specialized_decode_step(self) -> None: ++ if self.specialized_decode is not None: ++ self.specialized_decode.after_step() +diff --git a/vllm/models/kimi_k3/amd/specialized_decode/model_state.py b/vllm/models/kimi_k3/amd/specialized_decode/model_state.py +new file mode 100644 +index 000000000..b1e3898a4 +--- /dev/null ++++ b/vllm/models/kimi_k3/amd/specialized_decode/model_state.py +@@ -0,0 +1,27 @@ ++# SPDX-License-Identifier: Apache-2.0 ++# SPDX-FileCopyrightText: Copyright contributors to the vLLM project ++"""Kimi-K3's model runner state on ROCm: the hybrid (KDA + MLA) state, plus the ++lifecycle of the specialized decode when ``decode_backend=latency_specialized`` ++selects it.""" ++ ++from vllm.v1.worker.gpu.model_states.mamba_hybrid import MambaHybridModelState ++ ++ ++class KimiK3ModelState(MambaHybridModelState): ++ """Builds the specialized decode once the weights are loaded, and prepares it ++ for every bound KV cache (the CUDA-graph memory-profiling cache, then the real ++ one).""" ++ ++ def __init__(self, vllm_config, model, encoder_cache, device): ++ super().__init__(vllm_config, model, encoder_cache, device) ++ if hasattr(model, "get_language_model"): ++ model = model.get_language_model() ++ self.language_model = model ++ if model.specialized_decode_config is not None: ++ model.build_specialized_decode() ++ ++ def initialize_kv_cache(self, kv_cache_config, block_tables): ++ super().initialize_kv_cache(kv_cache_config, block_tables) ++ decode = self.language_model.specialized_decode ++ if decode is not None: ++ decode.prepare_kv() +diff --git a/vllm/models/kimi_k3/amd/specialized_decode/mono/__init__.py b/vllm/models/kimi_k3/amd/specialized_decode/mono/__init__.py +new file mode 100644 +index 000000000..208f01a7c +--- /dev/null ++++ b/vllm/models/kimi_k3/amd/specialized_decode/mono/__init__.py +@@ -0,0 +1,2 @@ ++# SPDX-License-Identifier: Apache-2.0 ++# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +diff --git a/vllm/models/kimi_k3/amd/specialized_decode/mono/abi.py b/vllm/models/kimi_k3/amd/specialized_decode/mono/abi.py +new file mode 100644 +index 000000000..7cf0c121f +--- /dev/null ++++ b/vllm/models/kimi_k3/amd/specialized_decode/mono/abi.py +@@ -0,0 +1,26 @@ ++# SPDX-License-Identifier: Apache-2.0 ++# SPDX-FileCopyrightText: Copyright contributors to the vLLM project ++"""The layout every MonoKernel's error words follow. ++ ++A rank's health words: a header whose word FAILED is set once any of its ++in-kernel waits gave up (sticky: a failed rank stays failed, and its later waits ++exit at once), then one record a CTA of the first wait it gave up on (flag, ++region id, pair index, tag, seen tag, CTA, clock, pad). ++ ++The output rank's flags: one int32 word a TP rank in the first PEER_HEAD bytes ++of the output rank's peer buffer. A rank whose wait gives up writes its word ++there, drains the store, then sets FAILED, so nothing it stores later can ++overtake the flag. The runner reads these words after every step. ++""" ++ ++MONO_ABI_VERSION = 1 ++ ++RECORD_WORDS = 8 ++FAILED = 0 ++PEER_HEAD = 256 # bytes: the output rank's flags, a word a rank ++WAIT_SECONDS = 10 # an in-kernel wait gives up after this long ++ ++ ++def health_words(ctas: int) -> int: ++ """A rank's health words for a grid of ``ctas``: the header, a record a CTA.""" ++ return (ctas + 1) * RECORD_WORDS +diff --git a/vllm/models/kimi_k3/amd/specialized_decode/mono/attention/__init__.py b/vllm/models/kimi_k3/amd/specialized_decode/mono/attention/__init__.py +new file mode 100644 +index 000000000..208f01a7c +--- /dev/null ++++ b/vllm/models/kimi_k3/amd/specialized_decode/mono/attention/__init__.py +@@ -0,0 +1,2 @@ ++# SPDX-License-Identifier: Apache-2.0 ++# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +diff --git a/vllm/models/kimi_k3/amd/specialized_decode/mono/attention/back.py b/vllm/models/kimi_k3/amd/specialized_decode/mono/attention/back.py +new file mode 100644 +index 000000000..9efe5d7f4 +--- /dev/null ++++ b/vllm/models/kimi_k3/amd/specialized_decode/mono/attention/back.py +@@ -0,0 +1,178 @@ ++# SPDX-License-Identifier: Apache-2.0 ++# SPDX-FileCopyrightText: Copyright contributors to the vLLM project ++"""K2a ``attn_post``: one Kimi-K3 KDA layer from the gated core output to the ++MoE's input row -- o_proj, its TP all-reduce inside the kernel, and the MLP ++AttnRes seam. ++ ++One launch a layer and rank, ``BLOCKS`` x ``THREADS``: ++ ++ gemv CTAs 56 .. o_proj (448 x 16 rows, K = 1536 local): bf16 GEMV of the ++ core rows (plain loads: the previous launch wrote them) -> ++ bf16 partial -> pushed to every rank's AR region (system ++ scope, this rank's slot) ++ attn CTAs 0 .. 55 per 128 columns: the 8 ranks' partials summed in rank ++ order (fp32 -> bf16: the custom all-reduce's result) -> ++ LDS delta; then K1's AttnRes stages with the MLP's weights: ++ prefix + delta -> prefix out, the sources' softmax mix, ++ post_attention_layernorm -> the MoE input x ++ ++The AR region is double-buffered by the step epoch's parity: a rank a step ++ahead writes the other half, so no peer's unread partials are overwritten. ++""" ++ ++import flydsl.expr as fx ++from aiter.ops.flydsl.kernels import buffer_ops as bo ++from flydsl.expr import gpu, range_constexpr ++from flydsl.expr.typing import T ++ ++from vllm.models.kimi_k3.amd.specialized_decode.mono.attention.kda import ( ++ COLS, ++ HIDDEN, ++ LDS_PAD, ++ NA, ++ PROJ, ++ ROWS, ++) ++from vllm.models.kimi_k3.amd.specialized_decode.mono.common.ops import ( ++ CM_NT, ++ bf16_round, ++ mfma_bf16, ++ row_sum, ++ rsrc, ++ traced, ++) ++from vllm.models.kimi_k3.amd.specialized_decode.mono.common.plan import ( ++ BLOCKS, ++ THREADS, ++ WAVES, ++ pair_layout, ++) ++from vllm.models.kimi_k3.amd.specialized_decode.mono.common.ranks import sum_partials ++ ++TP = 8 ++OROWS = HIDDEN # o_proj rows (full hidden, a rank's K slice) ++OTASKS = OROWS // ROWS # 448 ++OK = PROJ # K = 1536 a rank ++OKCH = OK // 128 # 12 chunks ++GEMV0 = NA # the first gemv CTA ++GEMV_N = BLOCKS - NA # 200 ++ ++_STREAM = fx.Stream(None) ++ ++ ++def ar_half_pairs(s): ++ """One parity's AR region: ``[src rank][token][column pair]``.""" ++ return TP * s * HIDDEN // 2 ++ ++ ++def scratch_layout(s: int, nblocks: int) -> dict: ++ ns = nblocks + 1 ++ pairs = { ++ "part": NA * s * ns * 2, ++ "wgt": s * (ns + 1), ++ "msq": NA * s, ++ "xrdy": NA, ++ } ++ return pair_layout(pairs.items()) ++ ++ ++def o_loads(c, rg): ++ """Rows 16 rg .. of o_proj, this wave's K chunks (wave, wave + 8 < 12): a lane ++ 64 B of its row a chunk.""" ++ lane, wave = c["lane"], c["wave"] ++ r_w = rsrc(c["args"]["w_o"]) ++ row = rg * ROWS + lane % ROWS ++ g = lane // ROWS ++ out = [] ++ for i in range(2): ++ ch = fx.min(wave + WAVES * i, OKCH - 1) ++ base = (row * OK + ch * 128 + g * 8) // 2 ++ out.append( ++ [ ++ fx.Vector( ++ bo.buffer_load( ++ r_w, ++ base + 16 * q, ++ vec_width=4, ++ dtype=T.i32, ++ cache_modifier=CM_NT, ++ ) ++ ) ++ for q in range(4) ++ ] ++ ) ++ return out ++ ++ ++@traced ++def load_core(c): ++ """The core rows (S x 1536 bf16) -> LDS, rows padded.""" ++ s, tid = c["S"], c["tid"] ++ words = s * OK // 8 ++ row = OK + LDS_PAD ++ for i in range_constexpr((words + THREADS - 1) // THREADS): ++ e = fx.min(tid + THREADS * i, words - 1) ++ v = fx.Vector( ++ bo.buffer_load(rsrc(c["args"]["core"]), e * 4, vec_width=4, dtype=T.i32) ++ ) ++ t = e // (OK // 8) ++ k = e % (OK // 8) * 8 ++ if tid + THREADS * i < words: ++ fx.ptr_store(v, c["xl"] + (t * row + k) // 2) ++ gpu.barrier() ++ ++ ++@traced ++def stage_o(c, rg, ops): ++ """Rows 16 rg ..: GEMV -> bf16 partials -> every rank's AR region, this ++ rank's slot (thread: token, row pair; a put a peer).""" ++ s, tid, lane, wave, red = c["S"], c["tid"], c["lane"], c["wave"], c["red"] ++ g = lane // ROWS ++ t = fx.min(lane % ROWS, s - 1) ++ row = OK + LDS_PAD ++ acc = fx.Vector.filled(4, 0.0, fx.Float32) ++ for i in range_constexpr(2): ++ ch = wave + WAVES * i ++ live = ch < OKCH ++ k = fx.min(ch, OKCH - 1) * 128 + g * 8 ++ for q in range_constexpr(4): ++ xb = fx.Vector( ++ fx.ptr_load( ++ c["xl"] + (t * row + k + 32 * q) // 2, ++ result_type=fx.Vector.make_type(4, fx.Int32), ++ ) ++ ) ++ xb = fx.Vector.from_elements( ++ [live.select(xb[d], fx.Int32(0)) for d in range(4)], fx.Int32 ++ ) ++ acc = mfma_bf16( ++ ops[i][q].bitcast(fx.BFloat16), xb.bitcast(fx.BFloat16), acc ++ ) ++ fx.ptr_store(acc, red + (wave * 64 + lane) * 4) ++ gpu.barrier() ++ if tid < s * (ROWS // 2): ++ tt = tid // (ROWS // 2) ++ rp = tid % (ROWS // 2) ++ v0 = bf16_round(row_sum(red, 2 * rp, tt)) ++ v1 = bf16_round(row_sum(red, 2 * rp + 1, tt)) ++ at = (c["rank"] * s + tt) * HIDDEN + rg * ROWS + 2 * rp ++ for p in range_constexpr(TP): ++ c["put_bf"](c["ar_peer"][p], at, [v0, v1]) ++ gpu.barrier() ++ ++ ++@traced ++def stage_reduce(c, task): ++ """Columns 128 task ..: the 8 ranks' partials, rank order, -> bf16 delta in ++ LDS (thread: token, column pair).""" ++ s, tid = c["S"], c["tid"] ++ if tid < s * (COLS // 2): ++ t = tid // (COLS // 2) ++ cp = tid % (COLS // 2) ++ col = task * COLS + 2 * cp ++ lo, hi = sum_partials( ++ c["poll"], c["ar_own"], lambda src: ((src * s + t) * HIDDEN + col) // 2, TP ++ ) ++ fx.ptr_store(bf16_round(lo), c["delta_lds"] + (t * COLS + 2 * cp)) ++ fx.ptr_store(bf16_round(hi), c["delta_lds"] + (t * COLS + 2 * cp + 1)) ++ gpu.barrier() +diff --git a/vllm/models/kimi_k3/amd/specialized_decode/mono/attention/kda.py b/vllm/models/kimi_k3/amd/specialized_decode/mono/attention/kda.py +new file mode 100644 +index 000000000..43a624f7c +--- /dev/null ++++ b/vllm/models/kimi_k3/amd/specialized_decode/mono/attention/kda.py +@@ -0,0 +1,1021 @@ ++# SPDX-License-Identifier: Apache-2.0 ++# SPDX-FileCopyrightText: Copyright contributors to the vLLM project ++"""K1 ``kda_pre``: one Kimi-K3 KDA layer from the AttnRes seam to the gated core ++output, for a spec-verify step (R requests x L tokens, S = R L rows). ++ ++One launch a layer, ``BLOCKS`` CTAs x ``THREADS`` threads, in three roles so that ++no CTA that polls a mailbox has weight loads in flight (a poll's load retires ++behind every earlier load): ++ ++ attn CTAs 0 .. 55 slice (128 hidden columns each): prefix (+ delta) -> ++ bf16 prefix out, per-source partial sum of squares and ++ qk dot (PART); gate (token t on CTA t): softmax weights ++ (WGT); mix: the softmax mix, its partial sum of squares ++ (MSQ), every slice's -> output RMSNorm -> bf16 x rows ++ (plain device-scope stores, an XRDY flag a slice) ++ gemv CTAs 56 .. in_proj, two 16-row groups each: the first group's ++ weights issued at kernel start, the second's as x lands ++ -> q|k|v|g|f_a|beta (PROJ) ++ kda tasks (R x 12 heads x 8 v-chunks) on CTAs 0 ..: inputs that ++ need no hand-off loaded first (metadata, conv window, ++ state rows, f_b rows), then conv update of the owned ++ channels (q / k out through QK), f_b MFMA + lower-bound ++ gate, l2norm, the recurrence on 16 state rows (a state a ++ token written), partial norms (ONRM), gated RMSNorm -> ++ core_out ++ ++Rounding follows the original path where it is elementwise (the AttnRes prefix ++add, every bf16 hand-off between the original kernels); split-K reductions ++differ only in their order. ++""" ++ ++import math ++from dataclasses import dataclass ++ ++import flydsl.compiler as flyc ++import flydsl.expr as fx ++import torch ++from aiter.ops.flydsl.kernels import buffer_ops as bo ++from flydsl.expr import const_expr, gpu, range_constexpr ++from flydsl.expr.typing import Int32, Int64, T ++ ++from vllm.models.kimi_k3.amd.specialized_decode.mono.common.debug import ( ++ region_ids, ++) ++from vllm.models.kimi_k3.amd.specialized_decode.mono.common.ops import ( ++ CM_DEV, ++ CM_NT, ++ bf_hi, ++ bf_lo, ++ butterfly, ++ hw_exp2, ++ hw_rsq, ++ ld_i32, ++ mfma_bf16, ++ row_sum, ++ rsrc, ++ traced, ++ uniform, ++ wave_max, ++) ++from vllm.models.kimi_k3.amd.specialized_decode.mono.common.plan import ( ++ BLOCKS, ++ THREADS, ++ WAVES, ++ KernelAbi, ++ first_task, ++ key_tuple, ++ pair_layout, ++ source_digest, ++) ++from vllm.models.kimi_k3.amd.specialized_decode.mono.common.sync import ( ++ Mailbox, ++ publish, ++ sreg, ++) ++ ++HIDDEN = 7168 ++NA = 56 # attn CTAs ++COLS = HIDDEN // NA # 128 columns a slice ++QC = COLS // 4 # threads a (token, source) row: 4 columns each ++GSL = 7 # slices a gate / mix-norm thread sums (NA = 8 x 7) ++HD = 128 ++NH = 12 # local heads at TP8 ++PROJ = NH * HD # 1536 ++QKV = 3 * PROJ ++G2_0 = QKV ++FA_0 = G2_0 + PROJ # 6144 ++BETA_0 = FA_0 + HD # 6272 ++NPROJ = 6288 # in_proj rows a rank holds (6284 + 4 pad) ++ROWS = 16 ++GEMV_TASKS = NPROJ // ROWS # 393 ++GEMV_CTAS = (GEMV_TASKS + 1) // 2 # 197 ++KCH = 128 # K a chunk: 4 MFMAs of 32 ++NKCH = HIDDEN // KCH # 56 ++PER_WAVE = NKCH // WAVES # 7 ++VCH = 16 # state rows a KDA task ++NVCH = HD // VCH # 8 ++CONV_W = 4 ++LOG2E = 1.0 / math.log(2.0) ++LDS_PAD = 8 # bf16 a padded x row ++ ++assert NA + GEMV_CTAS <= BLOCKS and NA == 8 * GSL ++ ++_STREAM = fx.Stream(None) ++REGIONS = ("part", "wgt", "msq", "xrdy", "proj", "qk", "onrm") ++# every mono source file: in every build's JIT cache key ++_SOURCES = source_digest(".") ++ ++ ++@dataclass(frozen=True) ++class KdaPreBuild: ++ tokens: int # S ++ qlen: int # L, tokens a request ++ nblocks: int # AttnRes block sources (the prefix is one more) ++ delta: bool ++ state_len: int # conv state entries: CONV_W - 1 + num_spec ++ eps: float = 1e-5 # AttnRes norm ++ out_eps: float = 1e-5 # input_layernorm ++ onorm_eps: float = 1e-5 ++ lower_bound: float = -5.0 ++ write_idx: int = -1 # a block-write layer's block: the updated prefix stored there ++ kda0: int = 0 # the first KDA task's CTA ++ ++ ++ABI = KernelAbi( ++ ( ++ "prefix", ++ "delta", ++ "blocks", ++ "blk_sm", ++ "blk_sr", ++ "ares_nw", ++ "ares_qk", ++ "in_nw", ++ "w_in", ++ "w_fb", ++ "conv_w", ++ "conv_st", ++ "cs_seq", ++ "cs_dim", ++ "cs_tok", ++ "a_log", ++ "dt_bias", ++ "on_w", ++ "rstate", ++ "rs_seq", ++ "st_idx", ++ "st_stride", ++ "num_acc", ++ "core_out", ++ "scratch", ++ "layer", ++ "epoch", ++ "health", ++ "alarm", ++ ) ++) ++ ++ ++def scratch_layout(key: KdaPreBuild) -> dict: ++ s, ns = key.tokens, key.nblocks + 1 ++ assert s % key.qlen == 0 ++ pairs = { ++ "part": NA * s * ns * 2, ++ "wgt": s * (ns + 1), ++ "msq": NA * s, ++ "xrdy": NA, ++ "proj": s * NPROJ, ++ "qk": s * NH * 2 * HD, ++ "onrm": s * NH * NVCH, ++ } ++ lay = pair_layout(pairs.items()) ++ end = max(o + n for o, n in lay.values()) ++ # the plain bf16 x rows (not pairs), 16 B aligned ++ lay["x"] = ((end + 15) // 16 * 16, s * HIDDEN * 2) ++ return lay ++ ++ ++def scratch_bytes(key: KdaPreBuild) -> int: ++ return max(o + n for o, n in scratch_layout(key).values()) ++ ++ ++TAG_SLOTS = 256 # a step's mailbox tags: epoch x 256 + a launch's slot ++ ++ ++def step_tag(epoch, slot): ++ """This launch's mailbox tag: the step's epoch (a device counter the runner ++ bumps once a step, in the graph) and the launch's slot, so no step's or ++ other launch's pair ever passes for this one's -- nothing is zeroed.""" ++ e = uniform(ld_i32(epoch, 0)) ++ return (e % (1 << 22)) * TAG_SLOTS + slot + 1 ++ ++ ++def ld_bf(ptr, i, cm=0): ++ return fx.Float32( ++ fx.BFloat16( ++ bo.buffer_load(rsrc(ptr), i, vec_width=1, dtype=T.bf16, cache_modifier=cm) ++ ) ++ ) ++ ++ ++def ld_bf4(ptr, i): ++ """bf16 elements i .. i + 4 (i a multiple of 4) as four f32.""" ++ w = fx.Vector(bo.buffer_load(rsrc(ptr), i // 2, vec_width=2, dtype=T.i32)) ++ return [bf_lo(w[0]), bf_hi(w[0]), bf_lo(w[1]), bf_hi(w[1])] ++ ++ ++def bf4_words(v): ++ """Four f32 -> two words of packed bf16.""" ++ return fx.Vector.from_elements(v, fx.Float32).to(fx.BFloat16).bitcast(fx.Int32) ++ ++ ++def ld_f32(ptr, i): ++ return fx.Float32(bo.buffer_load(rsrc(ptr), i, vec_width=1, dtype=T.f32)) ++ ++ ++def bf16r(v): ++ return fx.Float32(fx.Float32(v).to(fx.BFloat16)) ++ ++ ++def sigmoid(z): ++ return fx.Float32(1.0) / (fx.Float32(1.0) + hw_exp2(z * fx.Float32(-LOG2E))) ++ ++ ++def f32_of(w): ++ return w.bitcast(fx.Float32) ++ ++ ++def sum_polled(vals): ++ acc = f32_of(vals[0][0]) ++ for k in range(1, len(vals)): ++ acc = acc + f32_of(vals[k][0]) ++ return acc ++ ++ ++# ---------------------------------------------------------------- attn CTAs ++@traced ++def stage_slice(c, task): ++ """Columns 128 task ..: every token's sources into LDS ``vals`` (4 columns a ++ thread), the prefix (+ delta) rounded to bf16 and stored back, per-(token, ++ source) partial sum of squares and qk dot -> PART.""" ++ s, ns, nb, tid, lane = c["S"], c["NS"], c["NB"], c["tid"], c["lane"] ++ a = c["args"] ++ c0 = task * COLS ++ n = s * ns * QC # threads' work items ++ for i in range_constexpr((n + THREADS - 1) // THREADS): ++ e = fx.min(tid + THREADS * i, n - 1) ++ live = (tid + THREADS * i) < n ++ row = e // QC # (token, source) ++ t = row // ns ++ src = row % ns ++ col = c0 + (e % QC) * 4 ++ is_pre = src == nb ++ blk = ld_bf4( ++ a["blocks"], t * a["blk_sm"] + fx.min(src, nb - 1) * a["blk_sr"] + col ++ ) ++ pre = ld_bf4(a["prefix"], t * HIDDEN + col) ++ if const_expr(c["delta"]): ++ if const_expr(c.get("delta_lds") is not None): ++ # the delta reduced into LDS first (K2a: the all-reduced o_proj) ++ dl = [ ++ fx.ptr_load(c["delta_lds"] + (t * COLS + (e % QC) * 4 + j)) ++ for j in range(4) ++ ] ++ else: ++ dl = ld_bf4(a["delta"], t * HIDDEN + col) ++ if const_expr(c.get("reset", False)): ++ # a block-write layer's MLP seam: the prefix restarts at the ++ # attention's output (the input prefix is not read) ++ pre = dl ++ else: ++ pre = [bf16r(pre[j] + dl[j]) for j in range(4)] ++ if live & is_pre: ++ bo.buffer_store( ++ bf4_words(pre), rsrc(a["prefix"]), (t * HIDDEN + col) // 2 ++ ) ++ if const_expr(c.get("write_idx", -1) >= 0): # noqa: SIM102 (traced: `and` does not combine traced values) ++ # the block-write layer's attention seam: the updated prefix is ++ # block write_idx ++ if live & is_pre: ++ bo.buffer_store( ++ bf4_words(pre), ++ rsrc(a["blocks"]), ++ (t * a["blk_sm"] + c["write_idx"] * a["blk_sr"] + col) // 2, ++ ) ++ v = [ ++ live.select(is_pre.select(pre[j], blk[j]), fx.Float32(0.0)) ++ for j in range(4) ++ ] ++ w = ld_bf4(a["ares_nw"], col) ++ q = ld_bf4(a["ares_qk"], col) ++ # a thread past the end holds a clamped copy of the last element: it ++ # must not store (it raced the owner with zeros) ++ if live: ++ for j in range_constexpr(4): ++ fx.ptr_store(v[j], c["vals"] + (row * COLS + (e % QC) * 4 + j)) ++ sq = v[0] * v[0] + v[1] * v[1] + v[2] * v[2] + v[3] * v[3] ++ dt = ( ++ v[0] * (w[0] * q[0]) ++ + v[1] * (w[1] * q[1]) ++ + v[2] * (w[2] * q[2]) ++ + v[3] * (w[3] * q[3]) ++ ) ++ sq = butterfly(sq, (16, 8, 4, 2, 1)) ++ dt = butterfly(dt, (16, 8, 4, 2, 1)) ++ if live & (lane % 32 == 0): ++ at = ((task * s + t) * ns + src) * 2 ++ c["put"](c["part"], at, sq) ++ c["put"](c["part"], at + 1, dt) ++ gpu.barrier() ++ ++ ++@traced ++def stage_gate(c, t): ++ """Token t: the 56 slices' partials -> source weights e_s and their sum.""" ++ s, ns, tid, lane, wave, red = ( ++ c["S"], ++ c["NS"], ++ c["tid"], ++ c["lane"], ++ c["wave"], ++ c["red"], ++ ) ++ nq = ns * 2 ++ if tid < 8 * nq: ++ g = tid // nq ++ q = tid % nq ++ vals = c["poll"]( ++ [(c["part"], (((g * GSL + k) * s + t) * ns) * 2 + q, 1) for k in range(GSL)] ++ ) ++ fx.ptr_store(sum_polled(vals), red + (g * nq + q)) ++ gpu.barrier() ++ if wave == 0: ++ src = fx.min(lane, ns - 1) ++ sq = fx.ptr_load(red + src * 2) ++ dt = fx.ptr_load(red + src * 2 + 1) ++ for g in range_constexpr(1, 8): ++ sq = sq + fx.ptr_load(red + (g * nq + src * 2)) ++ dt = dt + fx.ptr_load(red + (g * nq + src * 2 + 1)) ++ rstd = hw_rsq(sq * fx.Float32(1.0 / HIDDEN) + fx.Float32(c["eps"])) ++ logit = dt * rstd * fx.Float32(LOG2E) ++ logit = (lane < ns).select(logit, fx.Float32(-3.0e38)) ++ m = wave_max(logit) ++ ex = (lane < ns).select(hw_exp2(logit - m), fx.Float32(0.0)) ++ den = butterfly(ex, (32, 16, 8, 4, 2, 1)) ++ if lane < ns: ++ c["put"](c["wgt"], t * (ns + 1) + lane, ex) ++ if lane == 0: ++ c["put"](c["wgt"], t * (ns + 1) + ns, den) ++ gpu.barrier() ++ ++ ++@traced ++def stage_mix(c, task): ++ """Columns 128 task ..: the softmax mix of ``vals`` (4 columns a thread), its ++ sum of squares (MSQ), every slice's -> rstd, then bf16 x = mixed rstd w -> ++ the plain x rows, and the slice's XRDY flag.""" ++ s, ns, tid, lane, red = c["S"], c["NS"], c["tid"], c["lane"], c["red"] ++ a = c["args"] ++ c0 = task * COLS ++ nw = s * (ns + 1) ++ if tid < nw: ++ v = c["poll"]([(c["wgt"], tid, 1)])[0][0] ++ fx.ptr_store(f32_of(v), red + tid) ++ gpu.barrier() ++ mine = tid < s * QC ++ t = fx.min(tid // QC, s - 1) ++ cc = (tid % QC) * 4 ++ mixed = [] ++ for j in range_constexpr(4): ++ acc = fx.Float32(0.0) ++ for src in range_constexpr(ns): ++ acc = acc + fx.ptr_load(red + (t * (ns + 1) + src)) * fx.ptr_load( ++ c["vals"] + ((t * ns + src) * COLS + cc + j) ++ ) ++ mixed.append(acc / fx.ptr_load(red + (t * (ns + 1) + ns))) ++ sq = ( ++ mixed[0] * mixed[0] ++ + mixed[1] * mixed[1] ++ + mixed[2] * mixed[2] ++ + mixed[3] * mixed[3] ++ ) ++ sq = butterfly(sq, (16, 8, 4, 2, 1)) ++ if mine & (lane % 32 == 0): ++ c["put"](c["msq"], task * s + t, sq) ++ gpu.barrier() ++ if tid < 8 * s: ++ g = tid // s ++ tt = tid % s ++ vals = c["poll"]([(c["msq"], (g * GSL + k) * s + tt, 1) for k in range(GSL)]) ++ fx.ptr_store(sum_polled(vals), red + (128 + g * s + tt)) ++ gpu.barrier() ++ tot = fx.ptr_load(red + (128 + t)) ++ for g in range_constexpr(1, 8): ++ tot = tot + fx.ptr_load(red + (128 + g * s + t)) ++ r = hw_rsq(tot * fx.Float32(1.0 / HIDDEN) + fx.Float32(c["out_eps"])) ++ wn = ld_bf4(a["in_nw"], c0 + cc) ++ xw = bf4_words([mixed[j] * r * wn[j] for j in range(4)]) ++ if mine: ++ at = (t * HIDDEN + c0 + cc) // 2 ++ bo.buffer_store(xw, rsrc(c["xbuf"]), at, cache_modifier=CM_DEV) ++ publish(c["put"], c["xrdy"], task, fx.Int32(1), tid == 0) ++ gpu.barrier() ++ ++ ++# ---------------------------------------------------------------- gemv CTAs ++def gemv_loads(c, rg, chunks): ++ """This wave's weight operands of rows 16 rg .. for its K chunks ``chunks`` ++ (indices into wave, wave + 8, ..): a lane 64 contiguous bytes of its row a ++ chunk (4 dwordx4). ``rg`` past the last group loads the last group.""" ++ lane, wave = c["lane"], c["wave"] ++ r_w = rsrc(c["args"]["w_in"]) ++ row = fx.min(rg, GEMV_TASKS - 1) * ROWS + lane % ROWS ++ g = lane // ROWS ++ out = [] ++ for i in chunks: ++ k = (wave + WAVES * i) * KCH + g * 8 ++ base = (row * HIDDEN + k) // 2 # dwords ++ out.append( ++ [ ++ fx.Vector( ++ bo.buffer_load( ++ r_w, ++ base + 16 * q, ++ vec_width=4, ++ dtype=T.i32, ++ cache_modifier=CM_NT, ++ ) ++ ) ++ for q in range(4) ++ ] ++ ) ++ return out ++ ++ ++@traced ++def load_x(c): ++ """Every slice's x ready -> the bf16 x rows into LDS (rows padded).""" ++ s, tid = c["S"], c["tid"] ++ if tid < NA: ++ c["poll"]([(c["xrdy"], tid, 1)]) ++ gpu.barrier() ++ words = s * HIDDEN // 8 # 16 B chunks ++ row = HIDDEN + LDS_PAD ++ for i in range_constexpr((words + THREADS - 1) // THREADS): ++ e = fx.min(tid + THREADS * i, words - 1) ++ v = fx.Vector( ++ bo.buffer_load( ++ rsrc(c["xbuf"]), e * 4, vec_width=4, dtype=T.i32, cache_modifier=CM_DEV ++ ) ++ ) ++ t = e // (HIDDEN // 8) ++ k = e % (HIDDEN // 8) * 8 ++ if tid + THREADS * i < words: ++ fx.ptr_store(v, c["xl"] + (t * row + k) // 2) ++ gpu.barrier() ++ ++ ++@traced ++def stage_gemv(c, rg, ops): ++ """Rows 16 rg .. of in_proj against the LDS x -> bf16 PROJ.""" ++ s, tid, lane, wave, red = c["S"], c["tid"], c["lane"], c["wave"], c["red"] ++ g = lane // ROWS ++ t = fx.min(lane % ROWS, s - 1) ++ row = HIDDEN + LDS_PAD ++ acc = fx.Vector.filled(4, 0.0, fx.Float32) ++ for i in range_constexpr(PER_WAVE): ++ k = (wave + WAVES * i) * KCH + g * 8 ++ for q in range_constexpr(4): ++ xb = fx.Vector( ++ fx.ptr_load( ++ c["xl"] + (t * row + k + 32 * q) // 2, ++ result_type=fx.Vector.make_type(4, fx.Int32), ++ ) ++ ) ++ acc = mfma_bf16( ++ ops[i][q].bitcast(fx.BFloat16), xb.bitcast(fx.BFloat16), acc ++ ) ++ fx.ptr_store(acc, red + (wave * 64 + lane) * 4) ++ gpu.barrier() ++ if (tid < ROWS * s) & (rg < GEMV_TASKS): ++ r = tid % ROWS ++ tt = tid // ROWS ++ v = bf16r(row_sum(red, r, tt)) ++ c["put"](c["proj"], tt * NPROJ + rg * ROWS + r, v) ++ gpu.barrier() ++ ++ ++# ---------------------------------------------------------------- kda ++def kda_stage(c, req, h, ch): ++ """The task's inputs that need no hand-off, issued before its first poll: ++ metadata, the owned channels' conv window / weights, the initial state rows, ++ the f_b rows of this wave and the per-head constants.""" ++ L, sl, tid, lane, wave = c["L"], c["SL"], c["tid"], c["lane"], c["wave"] ++ a = c["args"] ++ k = {} ++ # every value a buffer resource is built on wave-uniform (else a waterfall) ++ k["acc_idx"] = uniform(ld_i32(a["num_acc"], req)) - 1 ++ k["slot_c"] = uniform(ld_i32(a["st_idx"], req * a["st_stride"])) ++ k["slot0"] = uniform(ld_i32(a["st_idx"], req * a["st_stride"] + k["acc_idx"])) ++ k["slots"] = [ ++ uniform(ld_i32(a["st_idx"], req * a["st_stride"] + j)) for j in range(L) ++ ] ++ # conv: thread < 48 owns channel (part, d) ++ ct = fx.min(tid, 3 * VCH - 1) ++ part = ct // VCH ++ chan = part * PROJ + h * HD + ch * VCH + ct % VCH ++ cs = a["conv_st"] + fx.Int64(fx.max(k["slot_c"], 0)) * fx.Int64(a["cs_seq"]) * 2 ++ co = chan * a["cs_dim"] # the channel's first entry (bf16 elements) ++ k["cs"], k["co"], k["chan"], k["part"] = cs, co, chan, part ++ k["cw"] = fx.Vector( ++ bo.buffer_load(rsrc(a["conv_w"]), chan * CONV_W, vec_width=4, dtype=T.f32) ++ ) ++ k["win"] = [ ++ ld_bf(cs, co + (k["acc_idx"] + j) * a["cs_tok"]) for j in range(CONV_W - 1) ++ ] ++ k["old"] = [ ++ ld_bf(cs, co + (k["acc_idx"] + 1 + j) * a["cs_tok"]) for j in range(sl - L) ++ ] ++ # recurrence: thread -> state row v = tid / 32, k columns (tid % 32) 4 .. ++ vr = ch * VCH + tid // 32 ++ k0 = (tid % 32) * 4 ++ k["head_off"] = (h * HD + vr) * HD + k0 ++ base0 = a["rstate"] + fx.Int64(fx.max(k["slot0"], 0)) * fx.Int64(a["rs_seq"]) * 4 ++ k["st"] = fx.Vector( ++ bo.buffer_load(rsrc(base0), k["head_off"], vec_width=4, dtype=T.f32) ++ ) ++ # f_b: wave w's rows 16 w .. of head h, a lane 64 B of its row ++ fr = h * HD + wave * 16 + lane % 16 ++ fb = (fr * HD + (lane // 16) * 32) // 2 ++ k["fb"] = [ ++ fx.Vector(bo.buffer_load(rsrc(a["w_fb"]), fb + 4 * q, vec_width=4, dtype=T.i32)) ++ for q in range(4) ++ ] ++ k["aa"] = hw_exp2(ld_f32(a["a_log"], h) * fx.Float32(LOG2E)) ++ k["dtb"] = [ ++ ld_f32(a["dt_bias"], h * HD + wave * 16 + 4 * (lane // 16) + i) ++ for i in range(4) ++ ] ++ return k ++ ++ ++@traced ++def kda_conv(c, k, req, h, ch): ++ """Owned channels: the spec conv update of request ``req``'s L tokens -> new ++ conv state; q / k out through QK, v to LDS.""" ++ L, sl, tid = c["L"], c["SL"], c["tid"] ++ a = c["args"] ++ if tid < 3 * VCH: ++ xs = c["poll"]( ++ [(c["proj"], (req * L + j) * NPROJ + k["chan"], 1) for j in range(L)] ++ ) ++ xv = [f32_of(xs[j][0]) for j in range(L)] ++ rs = rsrc(k["cs"]) ++ if k["slot_c"] > 0: ++ for j in range_constexpr(sl): ++ nv = k["old"][j] if j < sl - L else xv[j - (sl - L)] ++ bo.buffer_store( ++ fx.Float32(nv).to(fx.BFloat16), rs, k["co"] + j * a["cs_tok"] ++ ) ++ seq = k["win"] + xv ++ for j in range_constexpr(L): ++ acc = fx.Float32(0.0) ++ for q in range_constexpr(CONV_W): ++ acc = acc + k["cw"][q] * seq[j + q] ++ o = bf16r(acc * sigmoid(acc)) ++ if k["part"] < 2: ++ c["put"]( ++ c["qk"], ++ ((req * NH + h) * L + j) * 2 * HD ++ + k["part"] * HD ++ + ch * VCH ++ + tid % VCH, ++ o, ++ ) ++ if k["part"] == 2: ++ fx.ptr_store(o, c["vloc"] + (j * VCH + tid % VCH)) ++ gpu.barrier() ++ ++ ++@traced ++def kda_gates(c, k, req, h): ++ """f_a, beta and the head's q / k (one poll batch a thread); g1 = f_b f_a by ++ MFMA (a wave 16 rows) -> lower-bound gate -> decay in LDS; q / k l2-normed ++ (q scaled).""" ++ L, tid, lane, wave = c["L"], c["tid"], c["lane"], c["wave"] ++ n_fa = L * HD ++ n = n_fa + L * 2 * HD ++ per = (n + THREADS - 1) // THREADS ++ specs = [] ++ for i in range_constexpr(per): ++ e = fx.min(tid + THREADS * i, n - 1) ++ fa_at = (req * L + e // HD) * NPROJ + FA_0 + e % HD ++ qk_at = (req * NH + h) * L * 2 * HD + (e - n_fa) ++ specs.append((e, fa_at, qk_at)) ++ # two regions, two batches: the f_a pairs, then the q / k pairs ++ got_fa = c["poll"]( ++ [(c["proj"], fx.min(fa, (req * L + L) * NPROJ - 1), 1) for _, fa, _ in specs] ++ ) ++ got_qk = c["poll"]( ++ [(c["qk"], fx.max(qk, (req * NH + h) * L * 2 * HD), 1) for _, _, qk in specs] ++ ) ++ for i in range_constexpr(per): ++ e = specs[i][0] ++ if tid + THREADS * i < n_fa: ++ fx.ptr_store(f32_of(got_fa[i][0]), c["fa"] + e) ++ if (tid + THREADS * i >= n_fa) & (tid + THREADS * i < n): ++ fx.ptr_store(f32_of(got_qk[i][0]), c["qkl"] + (e - n_fa)) ++ if tid < L: ++ b = c["poll"]([(c["proj"], (req * L + tid) * NPROJ + BETA_0 + h, 1)])[0][0] ++ fx.ptr_store(sigmoid(f32_of(b)), c["beta"] + tid) ++ gpu.barrier() ++ # g1 rows 16 wave .. x tokens: A = f_b rows (staged), B = f_a [token][K] ++ t = fx.min(lane % 16, L - 1) ++ acc = fx.Vector.filled(4, 0.0, fx.Float32) ++ for q in range_constexpr(4): ++ kb = (lane // 16) * 32 + 8 * q ++ xs = [fx.ptr_load(c["fa"] + (t * HD + kb + e)) for e in range(8)] ++ xb = fx.Vector.from_elements(xs, fx.Float32).to(fx.BFloat16) ++ acc = mfma_bf16(k["fb"][q].bitcast(fx.BFloat16), xb, acc) ++ for i in range_constexpr(4): ++ kk = wave * 16 + 4 * (lane // 16) + i ++ g = bf16r(acc[i]) + k["dtb"][i] ++ gate = fx.Float32(c["lb"]) * sigmoid(k["aa"] * g) ++ if lane % 16 < L: ++ fx.ptr_store(hw_exp2(gate * fx.Float32(LOG2E)), c["decay"] + (t * HD + kk)) ++ # l2norm: q / k row r (2 L rows of 128) on wave r % 8 ++ for rr in range_constexpr((2 * L + WAVES - 1) // WAVES): ++ row = fx.min(wave + WAVES * rr, 2 * L - 1) ++ x0 = fx.ptr_load(c["qkl"] + (row * HD + lane)) ++ x1 = fx.ptr_load(c["qkl"] + (row * HD + 64 + lane)) ++ ss = butterfly(x0 * x0 + x1 * x1, (32, 16, 8, 4, 2, 1)) ++ sc = hw_rsq(ss + fx.Float32(1e-6)) ++ sc = (row % 2 == 0).select(sc * fx.Float32(HD**-0.5), sc) ++ if wave + WAVES * rr < 2 * L: ++ fx.ptr_store(x0 * sc, c["qkl"] + (row * HD + lane)) ++ fx.ptr_store(x1 * sc, c["qkl"] + (row * HD + 64 + lane)) ++ gpu.barrier() ++ ++ ++@traced ++def kda_recur(c, k, req, h, ch): ++ """State rows 16 ch .. of head h: the recurrence over the L tokens, a state ++ written a token, bf16 outputs to LDS ``ol`` and their partial norms (ONRM).""" ++ L, tid, lane = c["L"], c["tid"], c["lane"] ++ a = c["args"] ++ v = tid // 32 ++ k0 = (tid % 32) * 4 ++ live0 = k["slot0"] > 0 ++ stv = [live0.select(k["st"][e], fx.Float32(0.0)) for e in range(4)] ++ for j in range_constexpr(L): ++ qv = [fx.ptr_load(c["qkl"] + ((2 * j) * HD + k0 + e)) for e in range(4)] ++ kv = [fx.ptr_load(c["qkl"] + ((2 * j + 1) * HD + k0 + e)) for e in range(4)] ++ dk = [fx.ptr_load(c["decay"] + (j * HD + k0 + e)) for e in range(4)] ++ stv = [stv[e] * dk[e] for e in range(4)] ++ p = stv[0] * kv[0] + stv[1] * kv[1] + stv[2] * kv[2] + stv[3] * kv[3] ++ p = butterfly(p, (16, 8, 4, 2, 1)) ++ vn = (fx.ptr_load(c["vloc"] + (j * VCH + v)) - p) * fx.ptr_load(c["beta"] + j) ++ stv = [stv[e] + vn * kv[e] for e in range(4)] ++ o = stv[0] * qv[0] + stv[1] * qv[1] + stv[2] * qv[2] + stv[3] * qv[3] ++ o = bf16r(butterfly(o, (16, 8, 4, 2, 1))) ++ o = live0.select(o, fx.Float32(0.0)) ++ slot = k["slots"][j] ++ if (slot > 0) & live0: ++ dst = rsrc(a["rstate"] + fx.Int64(slot) * fx.Int64(a["rs_seq"]) * 4) ++ bo.buffer_store( ++ fx.Vector.from_elements(stv, fx.Float32), dst, k["head_off"] ++ ) ++ if lane % 32 == 0: ++ fx.ptr_store(o, c["ol"] + (j * VCH + v)) ++ gpu.barrier() ++ if tid < L: ++ ss = fx.Float32(0.0) ++ for e in range_constexpr(VCH): ++ x = fx.ptr_load(c["ol"] + (tid * VCH + e)) ++ ss = ss + x * x ++ c["put"](c["onrm"], ((req * NH + h) * L + tid) * NVCH + ch, ss) ++ gpu.barrier() ++ ++ ++@traced ++def kda_norm(c, req, h, ch): ++ """Rows 16 ch ..: gated RMSNorm over the head's 128 (every chunk's partial) ++ -> core_out.""" ++ L, tid = c["L"], c["tid"] ++ a = c["args"] ++ if tid < L * VCH: ++ j = tid // VCH ++ e = tid % VCH ++ vr = ch * VCH + e ++ got = c["poll"]( ++ [(c["onrm"], ((req * NH + h) * L + j) * NVCH + q, 1) for q in range(NVCH)] ++ + [(c["proj"], (req * L + j) * NPROJ + G2_0 + h * HD + vr, 1)] ++ ) ++ ss = sum_polled(got[:NVCH]) ++ r = hw_rsq(ss * fx.Float32(1.0 / HD) + fx.Float32(c["onorm_eps"])) ++ y = ( ++ fx.ptr_load(c["ol"] + (j * VCH + e)) ++ * r ++ * ld_bf(a["on_w"], vr) ++ * sigmoid(f32_of(got[NVCH][0])) ++ ) ++ bo.buffer_store( ++ y.to(fx.BFloat16), rsrc(a["core_out"]), (req * L + j) * PROJ + h * HD + vr ++ ) ++ gpu.barrier() ++ ++ ++@traced ++def kda_tasks(c): ++ s, L, bid = c["S"], c["L"], c["bid"] ++ for task in range(first_task(bid, c["kda0"]), (s // L) * NH * NVCH, BLOCKS): ++ req = task // (NH * NVCH) ++ h = (task // NVCH) % NH ++ ch = task % NVCH ++ k = kda_stage(c, req, h, ch) ++ kda_conv(c, k, req, h, ch) ++ kda_gates(c, k, req, h) ++ kda_recur(c, k, req, h, ch) ++ kda_norm(c, req, h, ch) ++ ++ ++# ---------------------------------------------------------------- kernel ++_BUILDS: dict = {} ++ ++ ++def build(key: KdaPreBuild): ++ if key in _BUILDS: ++ return _BUILDS[key] ++ s, L = key.tokens, key.qlen ++ ns = key.nblocks + 1 ++ lay = scratch_layout(key) ++ xrow = HIDDEN + LDS_PAD ++ assert s * xrow * 2 <= 136 * 1024, "x rows exceed LDS" ++ assert L <= 16 and key.nblocks >= 1 ++ ++ @fx.struct ++ class AttnLds: ++ vals: fx.Array[fx.Float32, s * ns * COLS, 16] ++ ++ @fx.struct ++ class GemvLds: ++ xl: fx.Array[fx.Int32, s * xrow // 2, 16] ++ ++ @fx.struct ++ class KdaLds: ++ fa: fx.Array[fx.Float32, L * HD, 16] ++ qkl: fx.Array[fx.Float32, L * 2 * HD, 16] ++ decay: fx.Array[fx.Float32, L * HD, 16] ++ vloc: fx.Array[fx.Float32, L * VCH, 16] ++ ol: fx.Array[fx.Float32, L * VCH, 16] ++ beta: fx.Array[fx.Float32, 16, 16] ++ ++ @fx.union ++ class RoleLds: ++ attn: AttnLds ++ gemv: GemvLds ++ kda: KdaLds ++ ++ @fx.struct ++ class Smem: ++ red: fx.Array[fx.Float32, WAVES * 64 * 4, 16] ++ ++ keyed = key_tuple(key, _SOURCES) ++ name = f"k3_mono_kda_pre_s{s}_l{L}_nb{key.nblocks}_d{int(key.delta)}" ++ ++ @flyc.kernel(name=name, known_block_size=[THREADS, 1, 1]) ++ def kda_pre( ++ prefix: Int64, ++ delta: Int64, ++ blocks: Int64, ++ blk_sm: Int32, ++ blk_sr: Int32, ++ ares_nw: Int64, ++ ares_qk: Int64, ++ in_nw: Int64, ++ w_in: Int64, ++ w_fb: Int64, ++ conv_w: Int64, ++ conv_st: Int64, ++ cs_seq: Int32, ++ cs_dim: Int32, ++ cs_tok: Int32, ++ a_log: Int64, ++ dt_bias: Int64, ++ on_w: Int64, ++ rstate: Int64, ++ rs_seq: Int32, ++ st_idx: Int64, ++ st_stride: Int32, ++ num_acc: Int64, ++ core_out: Int64, ++ scratch: Int64, ++ layer: Int32, ++ epoch: Int64, ++ health: Int64, ++ alarm: Int64, ++ ): ++ _ = keyed ++ tid = fx.thread_idx.x ++ bid = fx.block_idx.x ++ alloc = fx.SharedAllocator() ++ lds = alloc.allocate(Smem).peek() ++ role = alloc.allocate(RoleLds) ++ al, gl, kl = role.attn.peek(), role.gemv.peek(), role.kda.peek() ++ mb = Mailbox(step_tag(epoch, layer) - 1, health, alarm, region_ids(REGIONS)) ++ c = { ++ "S": s, ++ "L": L, ++ "NS": ns, ++ "NB": key.nblocks, ++ "SL": key.state_len, ++ "tid": tid, ++ "bid": bid, ++ "lane": tid % 64, ++ "wave": tid // 64, ++ "delta": key.delta, ++ "eps": key.eps, ++ "out_eps": key.out_eps, ++ "onorm_eps": key.onorm_eps, ++ "lb": key.lower_bound, ++ "kda0": key.kda0, ++ "write_idx": key.write_idx, ++ "put": mb.put, ++ "poll": mb.poll, ++ "red": lds.red.ptr, ++ "vals": al.vals.ptr, ++ "xl": gl.xl.ptr, ++ "fa": kl.fa.ptr, ++ "qkl": kl.qkl.ptr, ++ "decay": kl.decay.ptr, ++ "vloc": kl.vloc.ptr, ++ "ol": kl.ol.ptr, ++ "beta": kl.beta.ptr, ++ "xbuf": scratch + fx.Int64(lay["x"][0]), ++ "args": { ++ "prefix": prefix, ++ "delta": delta, ++ "blocks": blocks, ++ "blk_sm": blk_sm, ++ "blk_sr": blk_sr, ++ "ares_nw": ares_nw, ++ "ares_qk": ares_qk, ++ "in_nw": in_nw, ++ "w_in": w_in, ++ "w_fb": w_fb, ++ "conv_w": conv_w, ++ "conv_st": conv_st, ++ "cs_seq": cs_seq, ++ "cs_dim": cs_dim, ++ "cs_tok": cs_tok, ++ "a_log": a_log, ++ "dt_bias": dt_bias, ++ "on_w": on_w, ++ "rstate": rstate, ++ "rs_seq": rs_seq, ++ "st_idx": st_idx, ++ "st_stride": st_stride, ++ "num_acc": num_acc, ++ "core_out": core_out, ++ }, ++ } ++ for region in REGIONS: ++ c[region] = sreg(scratch, lay[region][0], region) ++ if bid < NA: ++ stage_slice(c, bid) ++ if bid < s: ++ stage_gate(c, bid) ++ stage_mix(c, bid) ++ if bid >= NA: ++ gi = bid - NA ++ rg0 = 2 * gi ++ rg1 = 2 * gi + 1 ++ ops0 = gemv_loads(c, rg0, range(PER_WAVE)) ++ load_x(c) ++ # the second group's first half in flight under the first's MFMAs ++ ops1 = gemv_loads(c, rg1, range(4)) ++ stage_gemv(c, rg0, ops0) ++ ops1 = ops1 + gemv_loads(c, rg1, range(4, PER_WAVE)) ++ stage_gemv(c, rg1, ops1) ++ kda_tasks(c) ++ ++ @flyc.jit ++ def launch( ++ prefix: Int64, ++ delta: Int64, ++ blocks: Int64, ++ blk_sm: Int32, ++ blk_sr: Int32, ++ ares_nw: Int64, ++ ares_qk: Int64, ++ in_nw: Int64, ++ w_in: Int64, ++ w_fb: Int64, ++ conv_w: Int64, ++ conv_st: Int64, ++ cs_seq: Int32, ++ cs_dim: Int32, ++ cs_tok: Int32, ++ a_log: Int64, ++ dt_bias: Int64, ++ on_w: Int64, ++ rstate: Int64, ++ rs_seq: Int32, ++ st_idx: Int64, ++ st_stride: Int32, ++ num_acc: Int64, ++ core_out: Int64, ++ scratch: Int64, ++ layer: Int32, ++ epoch: Int64, ++ health: Int64, ++ alarm: Int64, ++ stream: fx.Stream = _STREAM, ++ ): ++ _ = keyed ++ kda_pre( ++ prefix, ++ delta, ++ blocks, ++ blk_sm, ++ blk_sr, ++ ares_nw, ++ ares_qk, ++ in_nw, ++ w_in, ++ w_fb, ++ conv_w, ++ conv_st, ++ cs_seq, ++ cs_dim, ++ cs_tok, ++ a_log, ++ dt_bias, ++ on_w, ++ rstate, ++ rs_seq, ++ st_idx, ++ st_stride, ++ num_acc, ++ core_out, ++ scratch, ++ layer, ++ epoch, ++ health, ++ alarm, ++ ).launch(grid=(BLOCKS,), block=(THREADS,), stream=stream) ++ ++ ABI.check(kda_pre, launch) ++ _BUILDS[key] = launch ++ return launch ++ ++ ++def ptr(t): ++ return 0 if t is None else t.data_ptr() ++ ++ ++def kda_pre( ++ key: KdaPreBuild, ++ *, ++ prefix, ++ delta, ++ blocks, ++ ares_nw, ++ ares_qk, ++ in_nw, ++ w_in, ++ w_fb, ++ conv_w, ++ conv_state, ++ a_log, ++ dt_bias, ++ on_w, ++ rstate, ++ st_idx, ++ num_acc, ++ core_out, ++ scratch, ++ layer, ++ epoch, ++ health, ++ alarm, ++): ++ """Host launcher. ``conv_state`` [slots, dim, state_len] or a transposed view ++ of [slots, state_len, dim] (strides read); ``rstate`` [slots, 12, 128, 128] ++ fp32; ``st_idx`` [R, >= L] int32; ``scratch`` zeroed before the step's first ++ layer (tags are a layer's); ``epoch`` a device int32 bumped once a step and ++ ``layer`` this launch's slot below ``TAG_SLOTS``.""" ++ f = build(key) ++ f( ++ prefix.data_ptr(), ++ ptr(delta), ++ blocks.data_ptr(), ++ blocks.stride(0), ++ blocks.stride(1), ++ ares_nw.data_ptr(), ++ ares_qk.data_ptr(), ++ in_nw.data_ptr(), ++ w_in.data_ptr(), ++ w_fb.data_ptr(), ++ conv_w.data_ptr(), ++ conv_state.data_ptr(), ++ conv_state.stride(0), ++ conv_state.stride(1), ++ conv_state.stride(2), ++ a_log.data_ptr(), ++ dt_bias.data_ptr(), ++ on_w.data_ptr(), ++ rstate.data_ptr(), ++ rstate.stride(0), ++ st_idx.data_ptr(), ++ st_idx.stride(0), ++ num_acc.data_ptr(), ++ core_out.data_ptr(), ++ scratch.data_ptr(), ++ layer, ++ epoch.data_ptr(), ++ health.data_ptr(), ++ alarm, ++ stream=torch.cuda.current_stream(), ++ ) +diff --git a/vllm/models/kimi_k3/amd/specialized_decode/mono/common/__init__.py b/vllm/models/kimi_k3/amd/specialized_decode/mono/common/__init__.py +new file mode 100644 +index 000000000..208f01a7c +--- /dev/null ++++ b/vllm/models/kimi_k3/amd/specialized_decode/mono/common/__init__.py +@@ -0,0 +1,2 @@ ++# SPDX-License-Identifier: Apache-2.0 ++# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +diff --git a/vllm/models/kimi_k3/amd/specialized_decode/mono/common/debug.py b/vllm/models/kimi_k3/amd/specialized_decode/mono/common/debug.py +new file mode 100644 +index 000000000..18cdcff38 +--- /dev/null ++++ b/vllm/models/kimi_k3/amd/specialized_decode/mono/common/debug.py +@@ -0,0 +1,38 @@ ++# SPDX-License-Identifier: Apache-2.0 ++# SPDX-FileCopyrightText: Copyright contributors to the vLLM project ++# Adapted from ROCm/ATOM (https://github.com/ROCm/ATOM) at a526f0d, MIT License, ++# Copyright (C) 2026, Advanced Micro Devices, Inc.: ++# atom/mono/runtime/debug.py ++"""Wait records: which mailbox region a CTA's given-up wait was on.""" ++ ++import zlib ++ ++import torch ++ ++from vllm.models.kimi_k3.amd.specialized_decode.mono.abi import RECORD_WORDS ++from vllm.models.kimi_k3.amd.specialized_decode.mono.common.plan import BLOCKS ++ ++ ++def region_id(name: str) -> int: ++ """A mailbox region's id in a wait record: its name's CRC-32 (31 bits), the ++ same in every kernel and table that names it.""" ++ return zlib.crc32(name.encode()) & 0x7FFFFFFF ++ ++ ++def region_ids(names) -> dict[str, int]: ++ """Name -> ``region_id``: a ``Mailbox``'s table.""" ++ return {n: region_id(n) for n in names} ++ ++ ++def given_up_waits(health: torch.Tensor, names) -> list[str]: ++ """Each CTA's first wait that gave up, the earliest first, from a rank's ++ ``health`` words; ``names`` the regions a record may name.""" ++ by_id = {region_id(n): n for n in names} ++ records = health.view(BLOCKS + 1, RECORD_WORDS)[1:] ++ records = records[records[:, 0] != 0].tolist() ++ t0 = min((r[6] for r in records), default=0) ++ return [ ++ f"+{(at - t0) / 100:.0f} us CTA {cta}: {by_id.get(rid, rid)} pair {idx}" ++ f" waited for tag {tag}, saw {seen}" ++ for _, rid, idx, tag, seen, cta, at, _ in sorted(records, key=lambda r: r[6]) ++ ] +diff --git a/vllm/models/kimi_k3/amd/specialized_decode/mono/common/mx.py b/vllm/models/kimi_k3/amd/specialized_decode/mono/common/mx.py +new file mode 100644 +index 000000000..908afeaa7 +--- /dev/null ++++ b/vllm/models/kimi_k3/amd/specialized_decode/mono/common/mx.py +@@ -0,0 +1,84 @@ ++# SPDX-License-Identifier: Apache-2.0 ++# SPDX-FileCopyrightText: Copyright contributors to the vLLM project ++# Adapted from ROCm/ATOM (https://github.com/ROCm/ATOM) at a526f0d, MIT License, ++# Copyright (C) 2026, Advanced Micro Devices, Inc.: ++# atom/mono/device/mx.py ++"""FP8 / MX arithmetic the mono kernels share: the E4M3 range, E8M0 scale ++codes, the byte of a scale in aiter's ``shuffle_scale`` layout, an FP4 weight ++tile's operand loads and the scaled MFMA that takes them. ++ ++An E8M0 code is a biased exponent; ``pow2(code)`` is the scale it stands for. ++""" ++ ++import flydsl.expr as fx ++from aiter.ops.flydsl.kernels import buffer_ops as bo ++from flydsl.expr import rocdl ++from flydsl.expr.typing import T ++ ++from vllm.models.kimi_k3.amd.specialized_decode.mono.common.ops import CM_NT ++ ++FP8_MAX = 448.0 ++# the f8f6f4 MFMA's operand formats (cbsz / blgp) ++FP8, FP4 = 0, 4 ++UNIT_SCALE = 127 # the E8M0 code of 1.0 ++ ++ ++def clamp_fp8(v): ++ """``v`` clamped to the E4M3 range.""" ++ return fx.max(fx.min(v, fx.Float32(FP8_MAX)), fx.Float32(-FP8_MAX)) ++ ++ ++def code_ceil(v): ++ """Biased exponent of the smallest power of two >= ``v`` (v >= 0, finite), ++ clamped to [1, 254]: the codes a scale may take.""" ++ bits = v.bitcast(fx.Int32) ++ e = ((bits >> 23) & 0xFF) + ((bits & 0x7FFFFF) != 0).select( ++ fx.Int32(1), fx.Int32(0) ++ ) ++ return fx.max(fx.min(e, fx.Int32(254)), fx.Int32(1)) ++ ++ ++def pow2(code): ++ """The fp32 power of two a biased exponent stands for.""" ++ return (code << 23).bitcast(fx.Float32) ++ ++ ++def rcp_pow2(code): ++ """1 / ``pow2(code)``, exact (code in [1, 253]): x times it is x / scale ++ bit for bit, a multiply where the division is a dozen instructions.""" ++ return ((254 - code) << 23).bitcast(fx.Float32) ++ ++ ++def scale_index(row, col, cols): ++ """Byte of e8m0 scale (row, col) in aiter's ``shuffle_scale`` layout: ++ ``view(rows/32, 2, 16, cols/8, 2, 4).permute(0, 3, 5, 2, 4, 1)``.""" ++ r32, a, b = row // 32, (row // 16) % 2, row % 16 ++ c8, d, e = col // 8, (col // 4) % 2, col % 4 ++ return ((((r32 * (cols // 8) + c8) * 4 + e) * 16 + b) * 2 + d) * 2 + a ++ ++ ++def fp4_tile_load(r_w, base, rg, k, st, at): ++ """A lane's 16 B of an FP4 weight's (16, 16)-shuffled 16-row group ``rg`` ++ (K = ``k``) at 128-k step ``st``: dword ``base`` + the group's 2 k + 256 a ++ step + ``at`` (``lane * 4``: 32 fp4, K block lane / 16 of row lane % 16). ++ Nontemporal: a weight a step reads once.""" ++ return fx.Vector( ++ bo.buffer_load( ++ r_w, ++ base + rg * (2 * k) + st * 256 + at, ++ vec_width=4, ++ dtype=T.i32, ++ cache_modifier=CM_NT, ++ ) ++ ) ++ ++ ++def mfma_scaled(a, b, c, sa, sb, a_fmt=FP8, b_fmt=FP8, sa_byte=0, sb_byte=0): ++ """One 16x16x128 scaled MFMA (gfx950 f8f6f4): A ``a`` (``a_fmt``) with E8M0 ++ ``sa`` (its byte ``sa_byte``), B ``b`` (``b_fmt``) with E8M0 ``sb`` (its ++ byte ``sb_byte``).""" ++ return fx.Vector( ++ rocdl.mfma_scale_f32_16x16x128_f8f6f4( ++ T.vec(4, T.f32), [a, b, c, a_fmt, b_fmt, sa_byte, sa, sb_byte, sb] ++ ) ++ ) +diff --git a/vllm/models/kimi_k3/amd/specialized_decode/mono/common/ops.py b/vllm/models/kimi_k3/amd/specialized_decode/mono/common/ops.py +new file mode 100644 +index 000000000..4b7b62dbf +--- /dev/null ++++ b/vllm/models/kimi_k3/amd/specialized_decode/mono/common/ops.py +@@ -0,0 +1,435 @@ ++# SPDX-License-Identifier: Apache-2.0 ++# SPDX-FileCopyrightText: Copyright contributors to the vLLM project ++# Adapted from ROCm/ATOM (https://github.com/ROCm/ATOM) at a526f0d, MIT License, ++# Copyright (C) 2026, Advanced Micro Devices, Inc.: ++# atom/mono/device/ops.py ++"""Device-side primitives of the mono kernels: wave reductions and shuffles, the ++approximate hardware math, bf16 / fp8 packing, buffer resources, the ++execution-model constants and the kernel symbol naming. ++ ++Ported from the FlyDSL GLM-5 layer kernel (ROCm/FlyDSL branch ++``glm5-mla-moe-layer-monokernel`` @ 1e28fd8a, ++``kernels/mla_moe_layer/shared_reuse_moe_kernel.py``). Only published flydsl APIs ++and aiter's vendored ``buffer_ops`` / ``dpp_utils`` are used. ++ ++Every function but ``kernel_symbol`` is traced inside a ``@flyc.kernel`` body. ++ ++Raw boundaries kept, for want of an equivalent on the pinned FlyDSL: ++``buffer_ops`` loads / stores (the kernels take raw Int64 addresses, no tensors ++for ``make_buffer_tensor``); raw MFMA intrinsics (hand-laid operands: permlane ++transposes, fp8 packing, MXFP4 scales); the LLVM intrinsics here (``expect``, ++``s_memrealtime``, the approximate rsq / rcp / exp2) and the struct unpacking of ++``permlane*_swap``. ++""" ++ ++from __future__ import annotations ++ ++import flydsl.expr as fx ++from aiter.ops.flydsl.kernels import buffer_ops as bo ++from aiter.ops.flydsl.kernels.dpp_utils import update_dpp_i32 ++from aiter.ops.flydsl.kernels.kernels_common import kernel_signature ++from flydsl._mlir import ir ++from flydsl._mlir.dialects import llvm ++from flydsl.compiler.ast_rewriter import ASTRewriter ++from flydsl.expr import gpu, range_constexpr, rocdl ++from flydsl.expr import math as fmath ++from flydsl.expr.typing import T, as_ir_value ++ ++from vllm.models.kimi_k3.amd.specialized_decode.mono.common.plan import THREADS, WAVES ++ ++# gfx95x cache policy bits (LLVM CPol): SC0 = 1, SC1 = 16. SC1 alone is device ++# scope (past the per-XCD caches); SC0 | SC1 is system scope (peers over XGMI). ++CM_DEV = 16 ++ ++ ++CM_SYS = 17 ++ ++ ++# NT = 2: nontemporal. A weight a layer reads once streams past the caches it ++# would only evict the step's reused data from (K2b at 48 rows: 143 -> 133 us) ++CM_NT = 2 ++ ++ ++def kernel_symbol(stem, **params): ++ """Kernel symbol that profilers demangle to ``atom::_``. The ++ amdhsa assembler rejects ``::`` in a symbol, so it is Itanium-mangled, as ++ aiter's ``_ZN5aiter...E`` kernels are.""" ++ name = f"{stem}_{kernel_signature(**params)}" ++ return f"_ZN4atom{len(name)}{name}E" ++ ++ ++def rsrc(addr, nbytes=None): ++ return bo.create_buffer_resource_from_addr(addr, num_records_bytes=nbytes) ++ ++ ++def permlane_swap(off, x, y): ++ """v_permlane{32,16}_swap on registers (x, y): off 32 trades x's upper-half ++ lanes with y's lower-half lanes; off 16 trades x's odd 16-lane rows with y's ++ even rows (each within its half). -> (x', y').""" ++ swap = rocdl.permlane32_swap if off == 32 else rocdl.permlane16_swap ++ pr = swap( ++ llvm.StructType.get_literal([T.i32, T.i32]), ++ as_ir_value(fx.Int32(x)), ++ as_ir_value(fx.Int32(y)), ++ False, ++ False, ++ ) ++ return tuple(fx.Int32(llvm.extractvalue(T.i32, pr, [j])) for j in range(2)) ++ ++ ++def lane_gather(v, src_lane): ++ """Lane ``src_lane``'s value of v, a source lane a lane (ds_bpermute).""" ++ return fx.Int32( ++ rocdl.ds_bpermute(T.i32, fx.Int32(src_lane) * 4, fx.Int32(v).ir_value()) ++ ) ++ ++ ++def readlane(v, src_lane): ++ """Lane ``src_lane``'s value of v (wave-uniform).""" ++ bits = v.bitcast(fx.Int32) if isinstance(v, fx.Float32) else fx.Int32(v) ++ r = fx.Int32(rocdl.readlane(T.i32, bits, fx.Int32(src_lane))) ++ return r.bitcast(fx.Float32) if isinstance(v, fx.Float32) else r ++ ++ ++def memrealtime(): ++ """s_memrealtime: the 100 MHz device clock (no FlyDSL wrapper).""" ++ return fx.Int64(llvm.call_intrinsic(T.i64, "llvm.amdgcn.s.memrealtime", [], [], [])) ++ ++ ++def unlikely(cond): ++ """``cond`` marked cold (llvm.expect): LLVM lays the branch out of the hot ++ path. A rare path's code inline among the hot stages costs them even when it ++ never runs (the long-context selection: S = 16, 3k context +1.4 us).""" ++ return fx.Boolean( ++ llvm.call_intrinsic( ++ ir.IntegerType.get_signless(1), ++ "llvm.expect.i1", ++ [as_ir_value(cond), as_ir_value(fx.Boolean(False))], ++ [], ++ [], ++ ) ++ ) ++ ++ ++def fresh(v): ++ """``v`` (a per-lane Int32) through an opaque move, here: the lane math built ++ on it is computed at its use, where LLVM would otherwise hoist it out of a task ++ loop or CSE it across stages into registers held for the whole kernel (spills, ++ and a kernel with scratch is refused: ``runtime.widths``).""" ++ return fx.Int32( ++ llvm.inline_asm( ++ T.i32, [fx.Int32(v).ir_value()], "v_mov_b32 $0, $1", "=v,v", ++ has_side_effects=True, ++ ) ++ ) # fmt: skip ++ ++ ++def uniform(v): ++ return fx.Int32(rocdl.readfirstlane(T.i32, fx.Int32(v).ir_value())) ++ ++ ++def _hw_f32(name, x): ++ return fx.Float32( ++ llvm.call_intrinsic(T.f32, name, [fx.Float32(x).ir_value()], [], []) ++ ) ++ ++ ++def hw_rsq(x): ++ """v_rsq_f32 (~1 ulp, no libm range fixups).""" ++ return _hw_f32("llvm.amdgcn.rsq.f32", x) ++ ++ ++def hw_rcp(x): ++ return _hw_f32("llvm.amdgcn.rcp.f32", x) ++ ++ ++def hw_exp2(x): ++ return _hw_f32("llvm.amdgcn.exp2.f32", x) ++ ++ ++def div_rn(x, d, r): ++ """``x / d`` rounded to nearest, bit-identical to the IEEE division, given ++ ``r = 1.0 / d`` from one true division per divisor (Markstein: a quotient ++ within 1 ulp plus its exact fma residual times RN(1/d) rounds correctly).""" ++ q = x * r ++ return fx.Float32(fmath.fma(fx.Float32(fmath.fma(-q, d, x)), r, q)) ++ ++ ++def xshfl(v, off): ++ """Value of lane ``lane ^ off``: offsets 32 / 16 via v_permlane*_swap, in-row ++ offsets via DPP (VALU latency, not LDS). Powers of two only: any other ++ offset would take the lane ^ 1 path.""" ++ assert off in (1, 2, 4, 8, 16, 32), off ++ if off >= 16: ++ return v.shuffle_xor(off, 64) ++ is_f = isinstance(v, fx.Float32) ++ x = v.bitcast(fx.Int32) if is_f else fx.Int32(v) ++ if off == 8: ++ y = fx.Int32(update_dpp_i32(x, x, 0x118, 0xF, 0xC, False)) ++ y = fx.Int32(update_dpp_i32(y, x, 0x108, 0xF, 0x3, False)) ++ elif off == 4: ++ y = fx.Int32(update_dpp_i32(x, x, 0x114, 0xF, 0xA, False)) ++ y = fx.Int32(update_dpp_i32(y, x, 0x104, 0xF, 0x5, False)) ++ elif off == 2: ++ y = fx.Int32(update_dpp_i32(x, x, 0x4E, 0xF, 0xF, False)) ++ else: ++ y = fx.Int32(update_dpp_i32(x, x, 0xB1, 0xF, 0xF, False)) ++ return y.bitcast(fx.Float32) if is_f else y ++ ++ ++def xred(v, off, op): ++ """op(v, lane ^ off) for a symmetric op; 32 / 16 take both halves of one swap.""" ++ if off < 16: ++ return op(v, xshfl(v, off)) ++ is_f = isinstance(v, fx.Float32) ++ x = v.bitcast(fx.Int32) if is_f else fx.Int32(v) ++ a, b = permlane_swap(off, x, x) ++ if is_f: ++ return op(a.bitcast(fx.Float32), b.bitcast(fx.Float32)) ++ return op(type(v)(a), type(v)(b)) ++ ++ ++def _add(a, b): ++ return a + b ++ ++ ++def butterfly(v, offsets, op=_add): ++ """Butterfly reduction over the given xor offsets, in the given order. ++ ++ The order is the numerical contract: the aiter per-head norm reduces 32 ++ lanes with offsets 16, 8, 4, 2, 1, and matching it keeps rstd bit-exact.""" ++ for off in offsets: ++ v = xred(v, off, op) ++ return v ++ ++ ++def wave_sum(v): ++ return butterfly(v, (32, 16, 8, 4, 2, 1)) ++ ++ ++def traced(fn): ++ """Apply flydsl's AST rewriting to a helper, so its ``while`` / ``if`` over ++ traced values lower to scf ops. Kernels rewrite only their own body.""" ++ return ASTRewriter.transform(fn) ++ ++ ++def wave_max(v): ++ return butterfly(v, (32, 16, 8, 4, 2, 1), fx.max) ++ ++ ++@traced ++def block_max(v, lane, wave, red): ++ """Max of v over the CTA's WAVES waves; ``red`` needs WAVES words.""" ++ wv = wave_max(v) ++ gpu.barrier() ++ if lane == 0: ++ fx.ptr_store(wv, red + wave) ++ gpu.barrier() ++ t = fx.ptr_load(red + 0) ++ for w in range_constexpr(1, WAVES): ++ t = fx.max(t, fx.ptr_load(red + w)) ++ gpu.barrier() ++ return t ++ ++ ++def row_shr(v, lane, k): ++ """DPP row_shr:k of an f32: lane - k of the same 16-lane row, else 0.""" ++ moved = lane_gather(v.bitcast(fx.Int32), fx.max(lane - k, 0)).bitcast(fx.Float32) ++ return fx.Float32((lane % 16 >= k).select(moved, fx.Float32(0.0))) ++ ++ ++def _xpartner(v, off, lane): ++ """Lane ``lane ^ off``'s v: DPP within a row, v_permlane*_swap across rows ++ (``xshfl`` takes ``shuffle_xor`` there).""" ++ if off < 16: ++ return xshfl(v, off) ++ a, b = permlane_swap(off, v, v) ++ return ((lane & off) != 0).select(a, b) ++ ++ ++def wave_incl_scan(v, lane): ++ """Inclusive prefix sum of an Int32 over the wave, in lane order.""" ++ blk = v ++ for off in (1, 2, 4, 8, 16, 32): ++ o = _xpartner(blk, off, lane) ++ v = v + ((lane & off) != 0).select(o, 0) ++ blk = blk + o ++ return v ++ ++ ++@traced ++def block_excl_scan(v, lane, wave, buf): ++ """(exclusive prefix of v over the CTA in thread order, CTA total); ``buf`` ++ needs WAVES words.""" ++ incl = wave_incl_scan(v, lane) ++ gpu.barrier() ++ if lane == 63: ++ fx.ptr_store(incl, buf + wave) ++ gpu.barrier() ++ off = fx.Int32(0) ++ tot = fx.Int32(0) ++ for w in range_constexpr(WAVES): ++ cnt = fx.ptr_load(buf + w) ++ off = off + (w < wave).select(cnt, 0) ++ tot = tot + cnt ++ return incl - v + off, tot ++ ++ ++def global_load(addr, words=1, nt=False): ++ """``words`` i32 (1, 2 or 4) at the 64-bit global address ``addr``. A pool ++ past 4 GiB is reached only this way: a buffer load's offset (index x stride ++ included) is 32 bits. ``nt``: nontemporal.""" ++ i32 = ir.IntegerType.get_signless(32) ++ ty = i32 if words == 1 else ir.VectorType.get([words], i32) ++ ptr = llvm.IntToPtrOp(ir.Type.parse("!llvm.ptr<1>"), as_ir_value(fx.Int64(addr))) ++ return llvm.LoadOp( ++ ty, ptr.result, alignment=4 * words, nontemporal=nt or None ++ ).result ++ ++ ++def ld_i32(ptr, i): ++ """Word ``i`` of a buffer, a plain Int32 load.""" ++ return fx.Int32(bo.buffer_load(rsrc(ptr), i, vec_width=1, dtype=T.i32)) ++ ++ ++def row_live(batch_ids, k): ++ """Whether row k of the step is a request's: its ``batch_id_per_q_token`` ++ entry, -1 on a CUDA-graph pad row (wave-uniform for a uniform k). A pad row ++ runs like any other but joins no decision another row's work depends on.""" ++ return uniform(ld_i32(batch_ids, k)) >= 0 ++ ++ ++def ld_dev(ptr, i, words=1): ++ """``words`` i32 (1, 2 or 4) of a buffer from word ``i``, at device scope ++ (past the per-XCD L2).""" ++ return bo.buffer_load( ++ rsrc(ptr), i, vec_width=words, dtype=T.i32, cache_modifier=CM_DEV ++ ) ++ ++ ++LOAD_BATCH = 4 # rounds of a strided loop whose loads go out together ++ ++ ++@traced ++def batched_rounds(tid, n, load, use, batch=LOAD_BATCH): ++ """A CTA's strided loop over elements v = tid + THREADS i < n (n a Python ++ int): ``load(v)`` of ``batch`` rounds issued before any ``use(v, loaded)``. ++ A round whose ``use`` stores to global memory keeps the next round's loads ++ behind it (they may alias), so the rounds' loads would go out in series. A ++ round past n loads element n - 1 (in range) and uses nothing.""" ++ rounds = -(-n // THREADS) ++ for b0 in range_constexpr(0, rounds, batch): ++ idx = range(b0, min(b0 + batch, rounds)) ++ got = [load(fx.min(tid + THREADS * i, n - 1)) for i in idx] ++ for i, g in zip(idx, got): ++ v = tid + THREADS * i ++ if v < n: ++ use(v, g) ++ ++ ++def rows_to_lds(tid, src, rows, row_words, dst, dst_row): ++ """``rows`` rows of ``row_words`` words (a multiple of 4) read at device scope ++ from address ``src`` -> LDS ``dst``, a row every ``dst_row`` words: 16 B a ++ lane, the loads batched (``batched_rounds``).""" ++ assert row_words % 4 == 0, row_words ++ batched_rounds( ++ tid, rows * row_words // 4, ++ lambda v: fx.Vector(ld_dev(src, v * 4, 4)), ++ lambda v, w: fx.ptr_store( ++ w, dst + (v * 4 // row_words * dst_row + v * 4 % row_words) ++ ), ++ ) # fmt: skip ++ ++ ++def mfma_bf16(a, b, c): ++ return fx.Vector(rocdl.mfma_f32_16x16x32_bf16(T.vec(4, T.f32), [a, b, c])) ++ ++ ++def row_sum(red, r, t, waves=None): ++ """Row r, token t of the waves' 16x16 MFMA accumulators (wave w's at ``red`` ++ + (w * 64 + lane) * 4), summed in wave order from 0 (``waves``: which, in ++ order; all of them by default).""" ++ tot = fx.Float32(0.0) ++ for w in waves if waves is not None else range(WAVES): ++ tot = tot + fx.ptr_load(red + ((w * 64 + 16 * (r // 4) + t) * 4 + r % 4)) ++ return tot ++ ++ ++def mfma_fp8(a, b, c): ++ """16x16x32 fp8 x fp8 MFMA, unscaled: ``a`` / ``b`` an i64 of 8 E4M3 each.""" ++ return fx.Vector( ++ rocdl.mfma_f32_16x16x32_fp8_fp8(T.vec(4, T.f32), [a, b, c, 0, 0, 0]) ++ ) ++ ++ ++def ballot(pred): ++ """The wave's lanes where ``pred`` holds, a bit a lane.""" ++ return fx.Int64(rocdl.ballot(T.i64, as_ir_value(pred))) ++ ++ ++def popcount(x): ++ return fx.Int32(fx.ctpop(x)) ++ ++ ++def lanes_below(bal, lane): ++ """How many of the lanes in ballot ``bal`` are below ``lane``.""" ++ return popcount(bal & ((fx.Int64(1) << fx.Int64(lane)) - 1)) ++ ++ ++def load_ptr64(table, i): ++ """Entry ``i`` of an i64 table in memory, wave-uniform (each half through ++ ``uniform``): a peer's buffer base, a kernel argument block's pointer.""" ++ pw = fx.Vector(bo.buffer_load(rsrc(table), i * 2, vec_width=2, dtype=T.i32)) ++ return (fx.Int64(uniform(pw[1])) << 32) | fx.Int64(fx.Uint32(uniform(pw[0]))) ++ ++ ++def fp8x8_bf16_pk(words): ++ """Eight fp8 of two dwords -> bf16 MFMA operand (exact) by ++ v_cvt_scalef32_pk_bf16_fp8 (scale 1), two values an instruction.""" ++ out = [] ++ for w in words: ++ for hi in (False, True): ++ p = rocdl.cvt_scalef32_pk_bf16_fp8( ++ T.bf16x2, as_ir_value(fx.Int32(w)), as_ir_value(fx.Float32(1.0)), hi ++ ) ++ out.append(fx.Vector(p).bitcast(fx.Int32)[0]) ++ return fx.Vector.from_elements(out, fx.Int32).bitcast(fx.BFloat16) ++ ++ ++def bf16_round(a): ++ return fx.Float32(fx.Float32(a).to(fx.BFloat16)) ++ ++ ++def bf16_pair(a, b): ++ """Two f32 -> one f32-typed word holding (bf16(a), bf16(b)).""" ++ return ( ++ fx.Vector.from_elements([a, b], fx.Float32) ++ .to(fx.BFloat16) ++ .bitcast(fx.Float32)[0] ++ ) ++ ++ ++def lds_bytes(struct) -> int: ++ """The LDS an ``fx.struct`` class takes, alignment included.""" ++ return struct.__dsl_size_of__() ++ ++ ++def bf_lo(w): ++ """The f32 of a packed bf16 pair word's low half.""" ++ return (w << 16).bitcast(fx.Float32) ++ ++ ++def bf_hi(w): ++ """The f32 of a packed bf16 pair word's high half.""" ++ return (w & fx.Int32(-65536)).bitcast(fx.Float32) ++ ++ ++def bf2_f32(w): ++ """Packed bf16 pair word -> (f32 low, f32 high).""" ++ return bf_lo(w), bf_hi(w) ++ ++ ++def fp8_pack4(a, b, c, d): ++ """Four f32 (already in range) -> one dword of four E4M3 bytes, RNE.""" ++ lo = fx.Int32(rocdl.cvt_pk_fp8_f32(T.i32, a, b, fx.Int32(0), False)) ++ return fx.Int32(rocdl.cvt_pk_fp8_f32(T.i32, c, d, lo, True)) +diff --git a/vllm/models/kimi_k3/amd/specialized_decode/mono/common/plan.py b/vllm/models/kimi_k3/amd/specialized_decode/mono/common/plan.py +new file mode 100644 +index 000000000..0e769475a +--- /dev/null ++++ b/vllm/models/kimi_k3/amd/specialized_decode/mono/common/plan.py +@@ -0,0 +1,237 @@ ++# SPDX-License-Identifier: Apache-2.0 ++# SPDX-FileCopyrightText: Copyright contributors to the vLLM project ++# Adapted from ROCm/ATOM (https://github.com/ROCm/ATOM) at a526f0d, MIT License, ++# Copyright (C) 2026, Advanced Micro Devices, Inc.: ++# atom/mono/plan/build_key.py ++# atom/mono/plan/execution.py ++# atom/mono/plan/layout.py ++# atom/mono/runtime/abi.py ++# atom/mono/plan/trace.py ++"""A kernel build's parameters, all of them in FlyDSL's JIT cache key. ++ ++FlyDSL keys a compiled kernel by the scalar values (int, float, bool, str, enum, ++tuple) its launcher's closures capture; an object in a closure is not keyed, and ++a parameter used only by build-time Python (a placement table, a layout) is in no ++closure at all. Two builds differing only there would share one binary -- every ++TP rank ran rank 3's kernel once that way. So a build's parameters are one frozen ++dataclass of scalars, ``key_tuple`` turns it into one (name, value) tuple, and the ++launcher references that tuple: every field is keyed, whatever the body uses. ++ ++The tuple also carries the model's kernel sources (``source_digest``): FlyDSL ++keys a helper's source only when its walk reaches it, and it does not reach a ++helper called inside a traced function's dynamic loop or branch (the rewriter ++moves those bodies into nested functions it does not scan) nor one outside the ++kernel's directory -- an edit there kept the old binary. ++""" ++ ++import enum ++import functools ++import hashlib ++import inspect ++from contextlib import contextmanager ++from dataclasses import dataclass, field, fields ++from enum import Enum ++from pathlib import Path ++ ++from vllm.models.kimi_k3.amd.specialized_decode.mono.abi import health_words ++ ++_SCALARS = (int, float, bool, str, enum.Enum) ++_ROOT = Path(__file__).resolve().parents[1] ++ ++ ++@functools.cache ++def source_digest(*paths: str) -> str: ++ """sha256 of every ``.py`` file under ``paths`` (files or directories, ++ relative to the ``mono`` package), each file's path and contents: one model's ++ kernel sources, so an edit to another model's rebuilds none of these.""" ++ h = hashlib.sha256() ++ for rel in paths: ++ root = _ROOT / rel ++ files = [root] if root.is_file() else sorted(root.rglob("*.py")) ++ if not files: ++ raise FileNotFoundError(f"no kernel sources at {root}") ++ for f in files: ++ h.update(str(f.relative_to(_ROOT)).encode()) ++ h.update(f.read_bytes()) ++ return h.hexdigest()[:16] ++ ++ ++def key_tuple(key, sources: str) -> tuple: ++ """``key``'s (field name, value) pairs, every value a scalar, then ++ ("sources", ``sources``): the model's ``source_digest``.""" ++ out = [] ++ for f in fields(key): ++ value = getattr(key, f.name) ++ if not isinstance(value, _SCALARS): ++ raise TypeError( ++ f"{type(key).__name__}.{f.name} = {value!r}: a build parameter must be" ++ " a scalar to reach the JIT cache key" ++ ) ++ out.append((f.name, value)) ++ out.append(("sources", sources)) ++ return tuple(out) ++ ++ ++def symbol_params(key) -> dict: ++ """The fields that name the kernel symbol (``metadata={"sym": short name}``), ++ short name -> value, in field order.""" ++ return { ++ f.metadata["sym"]: getattr(key, f.name) ++ for f in fields(key) ++ if "sym" in f.metadata ++ } ++ ++ ++# The execution model every mono kernel is written for. ++# ++# One CTA per CU of an MI355X, all of them co-resident -- a spin wait on another ++# CTA can only make progress if that CTA is running, so the runner refuses a GPU ++# with another CU count -- and THREADS threads each. ++# ++ ++BLOCKS = 256 ++THREADS = 512 ++WAVES = THREADS // 64 ++LDS_BYTES = 160 * 1024 # a CU's LDS: one CTA's whole ++ ++ ++def first_task(bid, base): ++ """CTA ``bid``'s first task of a stage placed from CTA ``base`` (any count of ++ CTAs past 0: placements that add task counts run past the grid), then every ++ BLOCKS-th.""" ++ return (bid - base % BLOCKS + BLOCKS) % BLOCKS ++ ++ ++# A kernel's mailbox regions laid out in its scratch or peer buffer. ++ ++PAIR_BYTES = 8 # a (value, tag) pair ++# A rank's health words (``mono.abi``), for this grid. ++HEALTH_WORDS = health_words(BLOCKS) ++ ++ ++def pair_layout(items, align: int = PAIR_BYTES, start: int = 0) -> dict: ++ """``items`` (name, pairs) laid out in order: name -> (byte offset, bytes), ++ each region from an ``align``-byte boundary at or past ``start``.""" ++ out, off = {}, start ++ for name, pairs in items: ++ off = (off + align - 1) // align * align ++ out[name] = (off, pairs * PAIR_BYTES) ++ off += pairs * PAIR_BYTES ++ return out ++ ++ ++# A kernel's argument list, named once. ++# ++# A FlyDSL kernel and its launcher each spell out their positional parameters, and ++# the host passes them by position: an argument out of order is a pointer read as ++# another, silently. ``KernelAbi`` is the one list of names; the build checks the ++# kernel's and the launcher's parameters against it, and the host packs by name. ++# ++ ++ ++def _python_params(fn) -> tuple[str, ...]: ++ """Parameter names of a ``@flyc.kernel`` / ``@flyc.jit`` object's Python ++ function (FlyDSL keeps it as ``_func`` / ``func``).""" ++ py = getattr(fn, "func", None) or getattr(fn, "_func", None) ++ if py is None: ++ raise TypeError(f"no Python function behind {fn!r}") ++ return tuple(inspect.signature(py).parameters) ++ ++ ++@dataclass(frozen=True) ++class KernelAbi: ++ names: tuple[str, ...] ++ ++ def check(self, kernel, launcher) -> None: ++ """The kernel takes exactly ``names``; the launcher too, then ``stream``.""" ++ got = _python_params(kernel) ++ if got != self.names: ++ raise TypeError(f"kernel parameters {got} != ABI {self.names}") ++ got = _python_params(launcher) ++ if got != (*self.names, "stream"): ++ raise TypeError(f"launcher parameters {got} != ABI {self.names} + stream") ++ ++ def pack(self, values: dict) -> list: ++ """``values`` (every name, nothing else) in ABI order.""" ++ try: ++ if len(values) == len(self.names): ++ return [values[n] for n in self.names] ++ except KeyError: ++ pass ++ extra = sorted(set(values) - set(self.names)) ++ missing = sorted(set(self.names) - set(values)) ++ raise TypeError(f"ABI mismatch: missing {missing}, unexpected {extra}") ++ ++ def zeros(self) -> list: ++ """Every argument 0: what a compile-only call passes.""" ++ return [0] * len(self.names) ++ ++ ++# A kernel's mailbox accesses, recorded while FlyDSL traces it. ++# ++# Tracing runs the kernel body's Python once, so every mailbox put / poll executes ++# once as Python: noting it here costs the device nothing (no IR is emitted). The ++# stage currently being traced is set by ``enter_stage``, called where the program ++# starts each stage, so the record is ``(stage, region, put | poll)`` in program ++# order. ++# ++# Everything here is plain functions on a module-level recorder: FlyDSL's AST ++# rewriter carries every local whose method is called inside a dynamic if / for as ++# loop state, so a kernel body must not call methods on these objects. ++# ++ ++ ++class Space(Enum): ++ SCRATCH = "scratch" # this GPU's memory: device-scope hand-off between CTAs ++ PEER = "peer" # a peer-symmetric region: system-scope hand-off between GPUs ++ ++ ++@dataclass(frozen=True) ++class Access: ++ stage: str ++ region: str ++ space: Space ++ kind: str # "put" | "poll" ++ ++ ++@dataclass ++class Record: ++ stages: list[str] = field(default_factory=list) # in the order entered ++ accesses: list[Access] = field(default_factory=list) ++ unnamed: list[tuple[str, str]] = field(default_factory=list) # (stage, kind) ++ ++ ++_record: Record | None = None ++ ++ ++@contextmanager ++def recording(): ++ """Record the accesses of kernels traced inside this block.""" ++ global _record ++ outer, _record = _record, Record() ++ try: ++ yield _record ++ finally: ++ _record = outer ++ ++ ++def enter_stage(name: str) -> None: ++ if _record is not None: ++ _record.stages.append(name) ++ ++ ++def _stage() -> str: ++ assert _record is not None ++ return _record.stages[-1] if _record.stages else "" ++ ++ ++def note(region: str, space: Space, kind: str) -> None: ++ if _record is not None: ++ _record.accesses.append(Access(_stage(), region, space, kind)) ++ ++ ++def note_unnamed(kind: str) -> None: ++ """A put / poll on a raw address: the check reports it (every mailbox access of ++ a checked kernel names its region).""" ++ if _record is not None: ++ _record.unnamed.append((_stage(), kind)) +diff --git a/vllm/models/kimi_k3/amd/specialized_decode/mono/common/ranks.py b/vllm/models/kimi_k3/amd/specialized_decode/mono/common/ranks.py +new file mode 100644 +index 000000000..ba75dba7c +--- /dev/null ++++ b/vllm/models/kimi_k3/amd/specialized_decode/mono/common/ranks.py +@@ -0,0 +1,33 @@ ++# SPDX-License-Identifier: Apache-2.0 ++# SPDX-FileCopyrightText: Copyright contributors to the vLLM project ++# Adapted from ROCm/ATOM (https://github.com/ROCm/ATOM) at a526f0d, MIT License, ++# Copyright (C) 2026, Advanced Micro Devices, Inc.: ++# atom/mono/device/ranks.py ++"""A TP group seen from inside a kernel: every rank's peer buffer, and an ++all-reduce's sum of the ranks' partials. ``tp`` is a build-time constant.""" ++ ++from flydsl.expr import range_constexpr ++ ++from vllm.models.kimi_k3.amd.specialized_decode.mono.common.ops import ( ++ bf2_f32, ++ load_ptr64, ++) ++ ++ ++def peer_bases(peers, tp): ++ """Every rank's peer buffer base from the ``peers`` table (an i64 a rank), ++ loaded once: a put's address must not wait on a load of its own.""" ++ return [load_ptr64(peers, p) for p in range(tp)] ++ ++ ++def sum_partials(poll, own, pair_of, tp): ++ """The ``tp`` ranks' bf16 pair partials at ``pair_of(src)`` of the mailbox ++ region ``own``, summed in fp32 in rank order 0 .. tp - 1: one poll batch -> ++ (lo sum, hi sum). Every rank sums in the same order, so they agree.""" ++ ws = poll([(own, pair_of(src), 1) for src in range(tp)]) ++ acc0, acc1 = bf2_f32(ws[0][0]) ++ for src in range_constexpr(1, tp): ++ lo, hi = bf2_f32(ws[src][0]) ++ acc0 = acc0 + lo ++ acc1 = acc1 + hi ++ return acc0, acc1 +diff --git a/vllm/models/kimi_k3/amd/specialized_decode/mono/common/sync.py b/vllm/models/kimi_k3/amd/specialized_decode/mono/common/sync.py +new file mode 100644 +index 000000000..4dfa087e3 +--- /dev/null ++++ b/vllm/models/kimi_k3/amd/specialized_decode/mono/common/sync.py +@@ -0,0 +1,255 @@ ++# SPDX-License-Identifier: Apache-2.0 ++# SPDX-FileCopyrightText: Copyright contributors to the vLLM project ++# Adapted from ROCm/ATOM (https://github.com/ROCm/ATOM) at a526f0d, MIT License, ++# Copyright (C) 2026, Advanced Micro Devices, Inc.: ++# atom/mono/device/sync.py ++"""The tagged mailbox: hand-offs between the CTAs of one GPU (device scope) and ++between GPUs (system scope), every address a named region (``Addr``) recorded ++when the kernel is traced (``trace``).""" ++ ++from __future__ import annotations ++ ++from dataclasses import dataclass ++ ++import flydsl.expr as fx ++from aiter.ops.flydsl.kernels import buffer_ops as bo ++from flydsl.expr import const_expr, gpu, range_constexpr, rocdl ++from flydsl.expr.typing import T ++ ++from vllm.models.kimi_k3.amd.specialized_decode.mono.abi import ( ++ FAILED, ++ RECORD_WORDS, ++ WAIT_SECONDS, ++) ++from vllm.models.kimi_k3.amd.specialized_decode.mono.common.ops import ( ++ CM_DEV, ++ CM_SYS, ++ bf16_pair, ++ ld_i32, ++ memrealtime, ++ rsrc, ++ traced, ++ uniform, ++) ++from vllm.models.kimi_k3.amd.specialized_decode.mono.common.plan import Space, note ++ ++POLL_MAX = 12 # mailbox specs polled per batch (bounds live registers) ++ ++WAIT_TICKS = WAIT_SECONDS * 100_000_000 # s_memrealtime ticks at 100 MHz ++ ++ ++@traced ++def _spin(load_all, pending, health, alarm, region_id, index, tag): ++ """Re-issue ``load_all`` until ``pending`` clears, for at most WAIT_TICKS, and ++ not at all once this rank has failed (``health`` word FAILED), so a failed ++ step runs to its end instead of hanging. A side-effecting op in the loop keeps ++ the loads from being hoisted. ++ ++ A CTA's first give-up is recorded in its ``health`` record; it then raises the ++ output rank's flag for this rank (``alarm``) and drains that store before it ++ sets FAILED, so no later store of this rank's can overtake the flag.""" ++ t0 = memrealtime() ++ v = load_all() ++ why = fx.Int32(0) # 1: timed out, 2: the rank had failed ++ while pending(v) & (why == 0): ++ rocdl.s_sleep(1) ++ v = load_all() ++ failed = ld_health(health, FAILED) != 0 ++ late = memrealtime() - t0 > WAIT_TICKS ++ why = failed.select(fx.Int32(2), late.select(fx.Int32(1), fx.Int32(0))) ++ if why == 1: ++ bid = fx.block_idx.x ++ rec = (bid + 1) * RECORD_WORDS ++ if ld_health(health, rec) == 0: ++ at = fx.Int32(memrealtime() & 0x7FFFFFFF) ++ record = [1, region_id, index, tag, v[1], bid, at, 0] ++ for h in range_constexpr(2): # 16 B a store ++ bo.buffer_store( ++ fx.Vector.from_elements( ++ [fx.Int32(w) for w in record[4 * h : 4 * h + 4]], fx.Int32 ++ ), ++ rsrc(health), ++ rec + 4 * h, ++ cache_modifier=CM_SYS, ++ ) ++ bo.buffer_store(fx.Int32(1), rsrc(alarm), 0, cache_modifier=CM_SYS) ++ rocdl.s_waitcnt(vmcnt=0) ++ bo.buffer_store(fx.Int32(1), rsrc(health), FAILED, cache_modifier=CM_SYS) ++ return v ++ ++ ++@traced ++def stall(fault_rank, fault_epoch, fault_layer, ticks, rank, epoch, layer): ++ """The fault hook (tests only): CTA 0 of rank ``fault_rank`` stalls ``ticks`` ++ once, in the launch for slot ``fault_layer`` of the step at epoch ++ ``fault_epoch``.""" ++ if ( ++ (rank == fault_rank) ++ & (uniform(ld_i32(epoch, 0)) == fault_epoch) ++ & (layer == fault_layer) ++ & (fx.block_idx.x == 0) ++ ): ++ t0 = memrealtime() ++ now = memrealtime() ++ while now - t0 < ticks: ++ rocdl.s_sleep(127) ++ now = memrealtime() ++ ++ ++def ld_health(health, i): ++ return fx.Int32( ++ bo.buffer_load(rsrc(health), i, vec_width=1, dtype=T.i32, cache_modifier=CM_SYS) ++ ) ++ ++ ++@dataclass(frozen=True) ++class Addr: ++ """A mailbox region's address: the region's name and space, recorded when the ++ kernel is traced (``trace``), and the traced address. The space also ++ sets the hand-off's scope: device for SCRATCH, system for PEER.""" ++ ++ region: str ++ space: Space ++ value: object ++ ++ ++def sreg(base, offset, region): ++ """Scratch region ``region`` at byte ``offset`` of ``base``.""" ++ return Addr(region, Space.SCRATCH, base + fx.Int64(offset)) ++ ++ ++def preg(base, offset, region): ++ """Peer-symmetric region ``region`` at byte ``offset`` of a rank's buffer ++ ``base``.""" ++ return Addr(region, Space.PEER, base + fx.Int64(offset)) ++ ++ ++def shift(a, delta): ++ """``a`` moved by ``delta`` bytes (a traced or Python integer), same region.""" ++ return Addr(a.region, a.space, a.value + delta) ++ ++ ++# the store / load cache policy of a hand-off in each space ++_SCOPE_CM = {Space.SCRATCH: CM_DEV, Space.PEER: CM_SYS} ++ ++ ++def _noted(a, kind): ++ note(a.region, a.space, kind) ++ return a.value ++ ++ ++class Mailbox: ++ """Tagged-pair hand-off between CTAs (device scope) and GPUs (system scope). ++ ++ Every 32-bit value is stored next to this launch's tag; a consumer ++ polls the payload until all tags match, so a hand-off costs one round trip ++ (no store drain, no separate flag). The tag is ``layer + 1``: it tells the ++ layers of one step apart, and the runner zeroes every mailbox at a step's ++ start (``StepMailboxes.begin_step``), so a pair from an earlier step never ++ carries a live tag. Every address is an ``Addr``: its space picks the scope. ++ """ ++ ++ def __init__(self, layer, health, alarm, region_ids): ++ """``health``: this rank's health words (FAILED, then a wait record a ++ CTA); ``alarm``: this rank's flag word on the output rank; ++ ``region_ids``: region name -> the id a record names it by.""" ++ self.tag = layer + 1 ++ self.health = health ++ self.alarm = alarm ++ self.region_ids = region_ids ++ ++ def put(self, a, i, v): ++ bits = v.bitcast(fx.Int32) if isinstance(v, fx.Float32) else fx.Int32(v) ++ bo.buffer_store( ++ fx.Vector.from_elements([bits, self.tag], fx.Int32), ++ rsrc(_noted(a, "put")), ++ i * 2, ++ cache_modifier=_SCOPE_CM[a.space], ++ ) ++ ++ def put_bf(self, a, i, vs): ++ """Elements i .. i+len(vs) (2 or 4, i aligned) as packed bf16 pairs, one ++ store.""" ++ words = [] ++ for j in range_constexpr(len(vs) // 2): ++ words += [bf16_pair(vs[2 * j], vs[2 * j + 1]).bitcast(fx.Int32), self.tag] ++ bo.buffer_store( ++ fx.Vector.from_elements(words, fx.Int32), ++ rsrc(_noted(a, "put")), ++ i, ++ cache_modifier=_SCOPE_CM[a.space], ++ ) ++ ++ def put_words(self, a, i, words): ++ """Raw 32-bit words (1 or 2) at pairs i.. as (word, tag). Integer payloads ++ (fp8 words) must come this way: through an f32 value a NaN bit pattern is ++ not kept.""" ++ vals = [] ++ for w in words: ++ vals += [fx.Int32(w), self.tag] ++ bo.buffer_store( ++ fx.Vector.from_elements(vals, fx.Int32), ++ rsrc(_noted(a, "put")), ++ i * 2, ++ cache_modifier=_SCOPE_CM[a.space], ++ ) ++ ++ def poll(self, specs, batch=POLL_MAX): ++ """Batched poll of pairs ``specs`` = [(Addr, pair index, npairs in {1, 2})], ++ all of one space. ++ ++ All pairs are re-loaded together while any tag is stale, so a batch costs ++ one round trip after its last producer lands. Returns the value words.""" ++ if const_expr(len(specs) == 0): ++ return [] ++ if const_expr(len(specs) > batch): ++ return self.poll(specs[:batch], batch) + self.poll(specs[batch:], batch) ++ spaces = {a.space for a, _, _ in specs} ++ assert len(spaces) == 1, f"one poll batch spans {spaces}" ++ cm = _SCOPE_CM[spaces.pop()] ++ bases = [_noted(a, "poll") for a, _, _ in specs] ++ tag = self.tag ++ ++ def load_all(): ++ words = [] ++ for b, (_, i, n) in zip(bases, specs): ++ w = fx.Vector( ++ bo.buffer_load( ++ rsrc(b), ++ fx.Int32(i) * 2, ++ vec_width=2 * n, ++ dtype=T.i32, ++ cache_modifier=cm, ++ ) ++ ) ++ words += [w[e] for e in range(2 * n)] ++ return fx.Vector.from_elements(words, fx.Int32) ++ ++ nw = sum(2 * n for _, _, n in specs) ++ ++ def pending(v): ++ bad = v[1] != tag ++ for e in range_constexpr(3, nw, 2): ++ bad = bad | (v[e] != tag) ++ return bad ++ ++ rid = self.region_ids.get(specs[0][0].region, -1) ++ v = _spin(load_all, pending, self.health, self.alarm, rid, specs[0][1], tag) ++ outs, e = [], 0 ++ for _, _, n in specs: ++ outs.append([v[e + 2 * q] for q in range(n)]) ++ e += 2 * n ++ return outs ++ ++ ++@traced ++def publish(put, a, i, v, who): ++ """This CTA's plain stores made readable by every CTA, then (thread ``who``) ++ the mailbox pair (``a``, i) := v. The stores are drained first: a flag may ++ not overtake the data it announces, and a workgroup-scope release fence does ++ not wait for them (LLVM's AMDGPU memory model omits vmcnt outside tgsplit ++ mode: the workgroup's waves share one CU, the readers here do not).""" ++ rocdl.s_waitcnt(vmcnt=0) ++ gpu.barrier() ++ if who: ++ put(a, i, v) +diff --git a/vllm/models/kimi_k3/amd/specialized_decode/mono/layer.py b/vllm/models/kimi_k3/amd/specialized_decode/mono/layer.py +new file mode 100644 +index 000000000..c89c2b309 +--- /dev/null ++++ b/vllm/models/kimi_k3/amd/specialized_decode/mono/layer.py +@@ -0,0 +1,573 @@ ++# SPDX-License-Identifier: Apache-2.0 ++# SPDX-FileCopyrightText: Copyright contributors to the vLLM project ++"""K2 ``layer_post``: one Kimi-K3 layer from the attention's core output (the ++o_proj input) to the layer's MoE output, in one launch -- K2a's o_proj, its ++all-reduce and the MLP AttnRes seam, then K2b's latent MoE (``moe``). ++ ++As ATOM's ``layer_post``, the two halves share the launch so that the MoE's ++static weights stream while the AttnRes chain (partials -> reduce -> slice -> ++gate -> mix) is still settling: the CTAs that will run the router, latent and ++shared gate/up GEMVs issue their weights before they first wait. ++ ++ CTAs 0 .. 55 K2a's reduce + AttnRes (slice, gate, mix -> x rows + XRDY); ++ then route (0 .. S-1) and the shared SiTU (8 .. 8+S-1) ++ CTAs 56 .. 159 o_proj (448 x 16 rows) -> partials to every rank; then ++ the shared gate/up (96 tasks, weights right after o_proj) ++ CTAs 160 .. 243 router parts (28) and latent parts (56): weights first ++ every CTA ug; down (0 .. 223) or the shared down (224 .. 255); ++ lnorm (56 .. 83), up (0 .. 55), reduce (200 .. 255) ++""" ++ ++from dataclasses import dataclass ++ ++import flydsl.compiler as flyc ++import flydsl.expr as fx ++import torch ++from flydsl.expr import const_expr, range_constexpr ++from flydsl.expr.typing import Int32, Int64 ++ ++from vllm.models.kimi_k3.amd.specialized_decode.mono.attention import back as k2a ++from vllm.models.kimi_k3.amd.specialized_decode.mono.attention.kda import ( ++ _SOURCES, ++ COLS, ++ HIDDEN, ++ NA, ++ stage_gate, ++ stage_mix, ++ stage_slice, ++ step_tag, ++) ++from vllm.models.kimi_k3.amd.specialized_decode.mono.common.debug import ( ++ region_ids, ++) ++from vllm.models.kimi_k3.amd.specialized_decode.mono.common.ops import ( ++ ld_i32, ++ load_ptr64, ++ traced, ++ uniform, ++) ++from vllm.models.kimi_k3.amd.specialized_decode.mono.common.plan import ( ++ BLOCKS, ++ THREADS, ++ WAVES, ++ KernelAbi, ++ key_tuple, ++ pair_layout, ++) ++from vllm.models.kimi_k3.amd.specialized_decode.mono.common.ranks import peer_bases ++from vllm.models.kimi_k3.amd.specialized_decode.mono.common.sync import ( ++ Mailbox, ++ preg, ++ shift, ++ sreg, ++ stall, ++) ++from vllm.models.kimi_k3.amd.specialized_decode.mono.stages import gemv as gemv ++from vllm.models.kimi_k3.amd.specialized_decode.mono.stages import moe as k2b ++ ++O0, O1 = NA, BLOCKS # o_proj CTAs: every CTA past AttnRes's ++O_N = O1 - O0 # 200 ++OTASKS = k2a.OTASKS # 448 ++O_ROUNDS = (OTASKS + O_N - 1) // O_N # 5 ++E0 = 160 # router / latent CTAs (their weights in flight under their o_proj) ++H0 = 8 # the shared SiTU's CTAs (after their AttnRes work) ++SD0 = k2b.DOWN_TASKS # 224: the shared down's CTAs ++ ++assert E0 + k2b.R_TASKS + k2b.L_TASKS <= BLOCKS and k2b.G_TASKS <= E0 - O0 ++ ++REGIONS = ("part", "wgt", "msq", "xrdy", "ar") + k2b.REGIONS ++ ++_STREAM = fx.Stream(None) ++ ++ ++@dataclass(frozen=True) ++class K2Build: ++ tokens: int ++ nblocks: int # the MLP AttnRes block sources ++ eps: float = 1e-5 # mlp_res_norm ++ out_eps: float = 1e-5 # post_attention_layernorm ++ ln_eps: float = 1e-5 # the latent RMSNorm ++ reduce: bool = True # the final reduce here (else the next layer's K1) ++ reset: bool = False # a block-write layer: the prefix restarts at the o_proj output ++ # the fault hook (sync.stall): rank fault_rank stalls fault_ticks once, in ++ # slot fault_layer's launch of the step at epoch fault_epoch ++ fault_rank: int = -1 ++ fault_epoch: int = 0 ++ fault_layer: int = 0 ++ fault_ticks: int = 0 ++ ++ ++ABI = KernelAbi( ++ ( ++ "core", ++ "w_o", ++ "prefix", ++ "blocks", ++ "blk_sm", ++ "blk_sr", ++ "ares_nw", ++ "ares_qk", ++ "in_nw", ++ "w_gate", ++ "bias", ++ "w_ld", ++ "w_sgu", ++ "w_sd", ++ "w13", ++ "w13s", ++ "w2", ++ "w2s", ++ "ln_w", ++ "w_up", ++ "out", ++ "scratch", ++ "peers", ++ "rank", ++ "epoch", ++ "layer", ++ "health", ++ "alarm", ++ ) ++) ++ ++ ++def peer_layout(s): ++ """K2a's AR region, then K2b's: one parity's (name -> (byte offset, bytes)).""" ++ pairs = {"ar": k2a.ar_half_pairs(s), **k2b.peer_pairs(s)} ++ return pair_layout(pairs.items()) ++ ++ ++def peer_bytes(s): ++ lay = peer_layout(s) ++ return 2 * max(o + n for o, n in lay.values()) ++ ++ ++def scratch_layout(key: K2Build) -> dict: ++ s = key.tokens ++ a = k2a.scratch_layout(s, key.nblocks) ++ end = max(o + n for o, n in a.values()) ++ b = k2b.scratch_layout(s) ++ lay = dict(a) ++ for name, (o, n) in b.items(): ++ lay[name] = ((end + 15) // 16 * 16 + o, n) ++ end = max(o + n for o, n in lay.values()) ++ lay["xrow"] = ((end + 15) // 16 * 16, s * HIDDEN * 2) ++ return lay ++ ++ ++def scratch_bytes(key: K2Build) -> int: ++ return max(o + n for o, n in scratch_layout(key).values()) ++ ++ ++@traced ++def stage_oproj(c, gi): ++ """o_proj tasks gi, gi + 200, gi + 400 (< 448): every task's weights in ++ flight at once (one memory latency, not three), then the MFMAs.""" ++ ops = [k2a.o_loads(c, fx.min(gi + O_N * j, OTASKS - 1)) for j in range(O_ROUNDS)] ++ k2a.load_core(c) ++ for j in range_constexpr(O_ROUNDS): ++ t = gi + O_N * j ++ if t < OTASKS: ++ k2a.stage_o(c, t, ops[j]) ++ ++ ++_BUILDS: dict = {} ++ ++ ++def build(key: K2Build): ++ if key in _BUILDS: ++ return _BUILDS[key] ++ s, ns = key.tokens, key.nblocks + 1 ++ assert 1 <= s <= 8 ++ lay = scratch_layout(key) ++ play = peer_layout(s) ++ half = peer_bytes(s) // 2 ++ ++ @fx.struct ++ class RouteLds: ++ route: fx.Array[fx.Int32, s * 2 * k2b.TOPK, 16] ++ flag: fx.Array[fx.Int32, k2b.E, 16] ++ expert: fx.Array[fx.Int32, k2b.MAX_U, 16] ++ kof: fx.Array[fx.Int32, k2b.MAX_U * 8, 16] ++ scan: fx.Array[fx.Int32, WAVES, 16] ++ ++ @fx.struct ++ class EarlyLds: ++ xl: fx.Array[fx.Int32, s * (HIDDEN + gemv.LDS_PAD) // 2, 16] ++ ++ @fx.struct ++ class AttnLds: ++ vals: fx.Array[fx.Float32, s * ns * COLS, 16] ++ dl: fx.Array[fx.Float32, s * COLS, 16] ++ ++ @fx.struct ++ class UgLds: ++ latq: fx.Array[fx.Int32, s * k2b.LAT // 8, 16] ++ latsc: fx.Array[fx.Int32, s * k2b.LAT // 32, 16] ++ av: fx.Array[fx.Float32, THREADS, 16] ++ ++ @fx.struct ++ class DownLds: ++ interq: fx.Array[fx.Int32, s * k2b.TOPK * 48, 16] ++ intsc: fx.Array[fx.Int32, s * k2b.TOPK * k2b.UG_GROUPS, 16] ++ contrib: fx.Array[fx.Float32, s * k2b.TOPK * k2b.ROWS, 16] ++ ++ @fx.union ++ class StageLds: ++ early: EarlyLds ++ attn: AttnLds ++ ug: UgLds ++ down: DownLds ++ ++ @fx.struct ++ class Smem: ++ red: fx.Array[fx.Float32, WAVES * 64 * 4, 16] ++ ++ keyed = key_tuple(key, _SOURCES) ++ name = f"k3_mono_k2_s{s}_nb{key.nblocks}_x{int(key.reset)}_r{int(key.reduce)}" ++ if key.fault_rank >= 0: ++ name += ( ++ f"_f{key.fault_rank}_{key.fault_epoch}_{key.fault_layer}_{key.fault_ticks}" ++ ) ++ ++ @flyc.kernel(name=name, known_block_size=[THREADS, 1, 1]) ++ def k2( ++ core: Int64, ++ w_o: Int64, ++ prefix: Int64, ++ blocks: Int64, ++ blk_sm: Int32, ++ blk_sr: Int32, ++ ares_nw: Int64, ++ ares_qk: Int64, ++ in_nw: Int64, ++ w_gate: Int64, ++ bias: Int64, ++ w_ld: Int64, ++ w_sgu: Int64, ++ w_sd: Int64, ++ w13: Int64, ++ w13s: Int64, ++ w2: Int64, ++ w2s: Int64, ++ ln_w: Int64, ++ w_up: Int64, ++ out: Int64, ++ scratch: Int64, ++ peers: Int64, ++ rank: Int32, ++ epoch: Int64, ++ layer: Int32, ++ health: Int64, ++ alarm: Int64, ++ ): ++ _ = keyed ++ tid = fx.thread_idx.x ++ bid = fx.block_idx.x ++ alloc = fx.SharedAllocator() ++ lds = alloc.allocate(Smem).peek() ++ rl = alloc.allocate(RouteLds).peek() ++ st = alloc.allocate(StageLds) ++ el, al, ul, dl = st.early.peek(), st.attn.peek(), st.ug.peek(), st.down.peek() ++ if const_expr(key.fault_rank >= 0): ++ stall( ++ key.fault_rank, ++ key.fault_epoch, ++ key.fault_layer, ++ key.fault_ticks, ++ rank, ++ epoch, ++ layer, ++ ) ++ mb = Mailbox(step_tag(epoch, layer) - 1, health, alarm, region_ids(REGIONS)) ++ par = uniform(ld_i32(epoch, 0)) % 2 ++ bases = peer_bases(peers, TP_) ++ own = load_ptr64(peers, rank) ++ peer = {} ++ for region in ("rlog", "lat", "latr", "final", "ar"): ++ off = play[region][0] ++ peer[region] = [ ++ shift(preg(b, off, region), fx.Int64(par) * half) for b in bases ++ ] ++ peer[region + "_own"] = shift(preg(own, off, region), fx.Int64(par) * half) ++ c = { ++ "S": s, ++ "NS": ns, ++ "NB": key.nblocks, ++ "tid": tid, ++ "bid": bid, ++ "lane": tid % 64, ++ "wave": tid // 64, ++ "rank": rank, ++ "delta": True, ++ "reset": key.reset, ++ "eps": key.eps, ++ "out_eps": key.out_eps, ++ "ln_eps": key.ln_eps, ++ "put": mb.put, ++ "put_bf": mb.put_bf, ++ "put_words": mb.put_words, ++ "poll": mb.poll, ++ "red": lds.red.ptr, ++ "peer": peer, ++ "ar_peer": peer["ar"], ++ "ar_own": peer["ar_own"], ++ "xl": el.xl.ptr, ++ "vals": al.vals.ptr, ++ "delta_lds": al.dl.ptr, ++ "latq": ul.latq.ptr, ++ "latsc": ul.latsc.ptr, ++ "av": ul.av.ptr, ++ "interq": dl.interq.ptr, ++ "intsc": dl.intsc.ptr, ++ "contrib": dl.contrib.ptr, ++ "scan": rl.scan.ptr, ++ "rt": { ++ "route": rl.route.ptr, ++ "flag": rl.flag.ptr, ++ "expert": rl.expert.ptr, ++ "kof": rl.kof.ptr, ++ }, ++ "xbuf": scratch + fx.Int64(lay["xrow"][0]), ++ "xflags": NA, ++ "ln": scratch + fx.Int64(lay["ln"][0]), ++ "shp": scratch + fx.Int64(lay["shp"][0]), ++ "args": { ++ "core": core, ++ "w_o": w_o, ++ "prefix": prefix, ++ "blocks": blocks, ++ "blk_sm": blk_sm, ++ "blk_sr": blk_sr, ++ "ares_nw": ares_nw, ++ "ares_qk": ares_qk, ++ "in_nw": in_nw, ++ "x_dbg": 0, ++ "w_gate": w_gate, ++ "bias": bias, ++ "w_ld": w_ld, ++ "w_sgu": w_sgu, ++ "w_sd": w_sd, ++ "w13": w13, ++ "w13s": w13s, ++ "w2": w2, ++ "w2s": w2s, ++ "ln_w": ln_w, ++ "w_up": w_up, ++ "out": out, ++ }, ++ } ++ for region in ( ++ "part", ++ "wgt", ++ "msq", ++ "xrdy", ++ "sgu", ++ "h", ++ "rpart", ++ "lpart", ++ "route", ++ "inter", ++ "isc", ++ "lmsq", ++ "lnrdy", ++ "shrdy", ++ ): ++ c[region] = sreg(scratch, lay[region][0], region) ++ # ---- o_proj on every CTA past AttnRes's ++ eb = bid - E0 ++ lb = eb - k2b.R_TASKS ++ is_r = (bid >= E0) & (eb < k2b.R_TASKS) ++ is_l = (bid >= E0 + k2b.R_TASKS) & (lb < k2b.L_TASKS) ++ # o_proj first (every AttnRes CTA waits on the slowest partial; loads ++ # retire in order), the early weights after: they land before x does ++ if is_r: ++ stage_oproj(c, bid - O0) ++ rops = k2b.early_loads(c, 0, eb // k2b.R_PARTS, eb % k2b.R_PARTS) ++ k2b.early_compute(c, 0, eb // k2b.R_PARTS, eb % k2b.R_PARTS, rops) ++ if is_l: ++ stage_oproj(c, bid - O0) ++ lops = k2b.early_loads(c, 1, lb // k2b.L_PARTS, lb % k2b.L_PARTS) ++ k2b.early_compute(c, 1, lb // k2b.L_PARTS, lb % k2b.L_PARTS, lops) ++ if (bid >= O0) & ~is_r & ~is_l: ++ stage_oproj(c, bid - O0) ++ gi = bid - O0 ++ if gi < k2b.G_TASKS: ++ k2b.early_compute(c, 2, gi, 0, k2b.early_loads(c, 2, gi, 0)) ++ # ---- AttnRes: the reduced o_proj -> x; then route / SiTU ++ if bid < NA: ++ k2a.stage_reduce(c, bid) ++ stage_slice(c, bid) ++ if bid < s: ++ stage_gate(c, bid) ++ stage_mix(c, bid) ++ if bid < s: ++ k2b.stage_route(c, bid) ++ if (bid >= H0) & (bid < H0 + s): ++ k2b.stage_h(c, bid - H0) ++ # ---- the routed experts ++ # the latent (ready ~10 us before the route) quantized first ++ k2b.load_latq(c) ++ nu = k2b.load_route(c) ++ total = nu * k2b.UG_GROUPS ++ last = total - 1 ++ # reversed placement: the round's leftover tasks go to the high CTAs, ++ # whose early work ends first (the AttnRes / route CTAs start ug last) ++ task = BLOCKS - 1 - bid ++ ops = k2b.ug_loads(c, fx.min(task, last)) ++ while task < total: ++ nxt = task + BLOCKS ++ ops_n = k2b.ug_loads(c, fx.min(nxt, last)) ++ k2b.stage_ug(c, task, ops) ++ task = nxt ++ ops = ops_n ++ if bid < k2b.DOWN_TASKS: ++ dops = k2b.down_loads(c, bid, nu, fx.min(c["wave"], nu - 1)) ++ k2b.load_inter(c) ++ k2b.stage_down(c, bid, nu, dops) ++ k2b.run_sdown_rest(c, bid, w_sd) ++ k2b.run_sdown(c, bid, w_sd) ++ if (bid >= k2b.LN0) & (bid < k2b.LN0 + k2b.LN_TASKS): ++ k2b.stage_lnorm(c, bid - k2b.LN0) ++ if bid < k2b.UP_TASKS: ++ k2b.stage_up(c, bid) ++ if const_expr(key.reduce): # noqa: SIM102 (traced: `and` does not combine traced values) ++ if (bid >= k2b.RED0) & (bid < k2b.RED0 + k2b.RED_TASKS): ++ k2b.stage_reduce(c, bid - k2b.RED0) ++ ++ @flyc.jit ++ def launch( ++ core: Int64, ++ w_o: Int64, ++ prefix: Int64, ++ blocks: Int64, ++ blk_sm: Int32, ++ blk_sr: Int32, ++ ares_nw: Int64, ++ ares_qk: Int64, ++ in_nw: Int64, ++ w_gate: Int64, ++ bias: Int64, ++ w_ld: Int64, ++ w_sgu: Int64, ++ w_sd: Int64, ++ w13: Int64, ++ w13s: Int64, ++ w2: Int64, ++ w2s: Int64, ++ ln_w: Int64, ++ w_up: Int64, ++ out: Int64, ++ scratch: Int64, ++ peers: Int64, ++ rank: Int32, ++ epoch: Int64, ++ layer: Int32, ++ health: Int64, ++ alarm: Int64, ++ stream: fx.Stream = _STREAM, ++ ): ++ _ = keyed ++ k2( ++ core, ++ w_o, ++ prefix, ++ blocks, ++ blk_sm, ++ blk_sr, ++ ares_nw, ++ ares_qk, ++ in_nw, ++ w_gate, ++ bias, ++ w_ld, ++ w_sgu, ++ w_sd, ++ w13, ++ w13s, ++ w2, ++ w2s, ++ ln_w, ++ w_up, ++ out, ++ scratch, ++ peers, ++ rank, ++ epoch, ++ layer, ++ health, ++ alarm, ++ ).launch(grid=(BLOCKS,), block=(THREADS,), stream=stream) ++ ++ ABI.check(k2, launch) ++ _BUILDS[key] = launch ++ return launch ++ ++ ++TP_ = k2b.TP ++ ++ ++def k2_launch( ++ key: K2Build, ++ *, ++ core, ++ w_o, ++ prefix, ++ blocks, ++ ares_nw, ++ ares_qk, ++ in_nw, ++ w_gate, ++ bias, ++ w_ld, ++ w_sgu, ++ w_sd, ++ w13, ++ w13s, ++ w2, ++ w2s, ++ ln_w, ++ w_up, ++ out, ++ scratch, ++ peers, ++ rank, ++ epoch, ++ layer, ++ health, ++ alarm, ++): ++ """Host launcher: ``core`` [S, 1536] (o_proj's input), ``prefix`` updated in ++ place, ``out`` the MoE output [S, 7168]; MoE weights as ``moe``.""" ++ f = build(key) ++ f( ++ core.data_ptr(), ++ w_o.data_ptr(), ++ prefix.data_ptr(), ++ blocks.data_ptr(), ++ blocks.stride(0), ++ blocks.stride(1), ++ ares_nw.data_ptr(), ++ ares_qk.data_ptr(), ++ in_nw.data_ptr(), ++ w_gate.data_ptr(), ++ bias.data_ptr(), ++ w_ld.data_ptr(), ++ w_sgu.data_ptr(), ++ w_sd.data_ptr(), ++ w13.data_ptr(), ++ w13s.data_ptr(), ++ w2.data_ptr(), ++ w2s.data_ptr(), ++ ln_w.data_ptr(), ++ w_up.data_ptr(), ++ out.data_ptr(), ++ scratch.data_ptr(), ++ peers.data_ptr(), ++ rank, ++ epoch.data_ptr(), ++ layer, ++ health.data_ptr(), ++ alarm, ++ stream=torch.cuda.current_stream(), ++ ) +diff --git a/vllm/models/kimi_k3/amd/specialized_decode/mono/stages/__init__.py b/vllm/models/kimi_k3/amd/specialized_decode/mono/stages/__init__.py +new file mode 100644 +index 000000000..208f01a7c +--- /dev/null ++++ b/vllm/models/kimi_k3/amd/specialized_decode/mono/stages/__init__.py +@@ -0,0 +1,2 @@ ++# SPDX-License-Identifier: Apache-2.0 ++# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +diff --git a/vllm/models/kimi_k3/amd/specialized_decode/mono/stages/gemv.py b/vllm/models/kimi_k3/amd/specialized_decode/mono/stages/gemv.py +new file mode 100644 +index 000000000..26acddb5a +--- /dev/null ++++ b/vllm/models/kimi_k3/amd/specialized_decode/mono/stages/gemv.py +@@ -0,0 +1,126 @@ ++# SPDX-License-Identifier: Apache-2.0 ++# SPDX-FileCopyrightText: Copyright contributors to the vLLM project ++"""bf16 GEMV pieces the K3 mono kernels share: 16 weight rows a task against up ++to 16 token rows in LDS, one ``mfma_f32_16x16x32_bf16`` per 32 K. ++ ++K is walked in 128-wide chunks, wave w taking chunks w, w + 8, ...; a lane reads ++64 contiguous bytes of its row a chunk (4 dwordx4 at q 32 + g 8 elements: lane ++group g = lane / 16), and the token rows are read at the same K. The K order ++inside a chunk is a permutation both operands share, so the dot is unchanged. ++""" ++ ++import flydsl.expr as fx ++from aiter.ops.flydsl.kernels import buffer_ops as bo ++from flydsl.expr import gpu, range_constexpr ++from flydsl.expr.typing import T ++ ++from vllm.models.kimi_k3.amd.specialized_decode.mono.common.ops import ( ++ CM_DEV, ++ CM_NT, ++ mfma_bf16, ++ rsrc, ++ traced, ++) ++from vllm.models.kimi_k3.amd.specialized_decode.mono.common.plan import THREADS, WAVES ++ ++ROWS = 16 ++KCH = 128 ++LDS_PAD = 8 # bf16 a padded token row ++ ++ ++def chunks(k): ++ """A wave's K chunks for K = ``k``: (count a wave, total).""" ++ n = k // KCH ++ return (n + WAVES - 1) // WAVES, n ++ ++ ++def loads(lane, wave, w, row0, k, nt=True, ldk=None, k0=0): ++ """This wave's operands of weight rows ``row0`` .. + 16 (row-major, ``ldk`` ++ bf16 a row at ``w``, K columns ``k0`` .. ``k0 + k``): a list a chunk ++ (clamped into K: ``mfmas`` zeroes the extra).""" ++ per, n = chunks(k) ++ ldk = k if ldk is None else ldk ++ r_w = rsrc(w) ++ row = row0 + lane % ROWS ++ g = lane // ROWS ++ out = [] ++ for i in range(per): ++ ch = fx.min(wave + WAVES * i, n - 1) ++ base = (row * ldk + k0 + ch * KCH + g * 8) // 2 ++ out.append( ++ [ ++ fx.Vector( ++ bo.buffer_load( ++ r_w, ++ base + 16 * q, ++ vec_width=4, ++ dtype=T.i32, ++ cache_modifier=CM_NT if nt else 0, ++ ) ++ ) ++ for q in range(4) ++ ] ++ ) ++ return out ++ ++ ++def mfmas(lane, wave, xl, xrow, k, s, ops): ++ """``loads``' operands against token rows 0 .. s - 1 of the LDS rows ``xl`` ++ (``xrow`` bf16 a row) -> this wave's 16x16 accumulator.""" ++ per, n = chunks(k) ++ g = lane // ROWS ++ t = fx.min(lane % ROWS, s - 1) ++ acc = fx.Vector.filled(4, 0.0, fx.Float32) ++ for i in range_constexpr(per): ++ ch = wave + WAVES * i ++ live = ch < n ++ kk = fx.min(ch, n - 1) * KCH + g * 8 ++ for q in range_constexpr(4): ++ xb = fx.Vector( ++ fx.ptr_load( ++ xl + (t * xrow + kk + 32 * q) // 2, ++ result_type=fx.Vector.make_type(4, fx.Int32), ++ ) ++ ) ++ if const_live(per, n): ++ xb = fx.Vector.from_elements( ++ [live.select(xb[d], fx.Int32(0)) for d in range(4)], fx.Int32 ++ ) ++ acc = mfma_bf16( ++ ops[i][q].bitcast(fx.BFloat16), xb.bitcast(fx.BFloat16), acc ++ ) ++ return acc ++ ++ ++def const_live(per, n): ++ """Whether some wave's last chunk is past K (its MFMA must add zero).""" ++ return per * WAVES != n ++ ++ ++@traced ++def rows_to_lds(tid, src, k, s, xl, xrow, cm=0, ldk=None, k0=0): ++ """Columns k0 .. k0 + k of token rows 0 .. s - 1 of the bf16 [s, ldk] at ++ ``src`` -> LDS (``xrow`` a row), 16 B a lane.""" ++ ldk = k if ldk is None else ldk ++ words = s * k // 8 ++ for i in range_constexpr((words + THREADS - 1) // THREADS): ++ e = fx.min(tid + THREADS * i, words - 1) ++ t = e // (k // 8) ++ kk = e % (k // 8) * 8 ++ v = fx.Vector( ++ bo.buffer_load( ++ rsrc(src), ++ (t * ldk + k0 + kk) // 2, ++ vec_width=4, ++ dtype=T.i32, ++ cache_modifier=cm, ++ ) ++ ) ++ if tid + THREADS * i < words: ++ fx.ptr_store(v, xl + (t * xrow + kk) // 2) ++ gpu.barrier() ++ ++ ++def dev_rows_to_lds(tid, src, k, s, xl, xrow): ++ """``rows_to_lds`` at device scope: rows another CTA of this launch wrote.""" ++ rows_to_lds(tid, src, k, s, xl, xrow, CM_DEV) +diff --git a/vllm/models/kimi_k3/amd/specialized_decode/mono/stages/moe.py b/vllm/models/kimi_k3/amd/specialized_decode/mono/stages/moe.py +new file mode 100644 +index 000000000..01a0f861e +--- /dev/null ++++ b/vllm/models/kimi_k3/amd/specialized_decode/mono/stages/moe.py +@@ -0,0 +1,1025 @@ ++# SPDX-License-Identifier: Apache-2.0 ++# SPDX-FileCopyrightText: Copyright contributors to the vLLM project ++"""K2b ``moe``: one Kimi-K3 latent-MoE layer, from the MoE input row x to the ++layer's output, every TP reduction inside the kernel. ++ ++One launch a layer and rank, ``BLOCKS`` x ``THREADS``, S <= 8 rows (a ++spec-verify step). The original path (vLLM ROCm, see MOE_SPEC.md) replicates the ++router and the latent down projection on every rank; here each rank computes ++its eighth and the rows are all-gathered over the peer buffer. ++ ++ early (CTAs 0 .. 130, one task each, x from the input row; weights first): ++ router (7 x 16 rows of this rank's 112 experts): fp32 logits -> every ++ rank's RLOG ++ latent (28 x 16 rows of this rank's 448): bf16 -> every rank's LAT ++ sgu (96 x 16 rows of the shared gate | up): bf16 -> SGU ++ route (S, CTAs 131 ..): sigmoid (v_rcp), + bias, top-16, the pre-bias ++ sigmoids renormalized (IEEE 1 / sum) -> ROUTE (every rank alike) ++ sdown (448 x 16 rows, K 768, CTAs 131 ..): SiTU of SGU (vLLM's kernel) -> ++ bf16 -> GEMV -> bf16 partial; this rank's up columns -> SHP (plain, ++ SHRDY flags), the rest -> every rank's FINAL ++ ug (|U| x 12, every CTA): the union U of the step's picks; an expert's ++ 32 gate + 32 up rows against the MXFP4 latent (aiter's inline quant), ++ SiTUv2 (aiter's fp32 epilogue), MXFP4 -> INTER / ISC a pick ++ down (224 x 16 latent rows, every CTA): w2 of every expert of U against ++ its picks' INTER, bf16(w acc) folded in each token's top-k order with ++ a bf16 rounding an add -> every rank's LATR ++ lnorm (28 x 128 columns): the ranks' LATR in rank order, bf16, RMSNorm ++ (the fused AR + RMSNorm form) -> LN (plain, LNRDY flags) ++ up (56 x 16 rows of this rank's 896): LN GEMV + SHP -> bf16 -> every ++ rank's FINAL ++ reduce (56 x 128 columns, ``reduce``): FINAL in rank order -> bf16 out ++ ++Peer regions are double-buffered by the epoch's parity. ++""" ++ ++import math ++ ++import flydsl.expr as fx ++from aiter.ops.flydsl.kernels import buffer_ops as bo ++from flydsl.expr import const_expr, gpu, range_constexpr, rocdl ++from flydsl.expr.typing import T, as_ir_value ++ ++from vllm.models.kimi_k3.amd.specialized_decode.mono.common.mx import ( ++ FP4, ++ UNIT_SCALE, ++ fp4_tile_load, ++ mfma_scaled, ++ scale_index, ++) ++from vllm.models.kimi_k3.amd.specialized_decode.mono.common.ops import ( ++ CM_DEV, ++ ballot, ++ bf16_round, ++ bf_hi, ++ bf_lo, ++ block_excl_scan, ++ butterfly, ++ hw_exp2, ++ hw_rcp, ++ hw_rsq, ++ lanes_below, ++ popcount, ++ row_sum, ++ rsrc, ++ traced, ++ uniform, ++ wave_sum, ++) ++from vllm.models.kimi_k3.amd.specialized_decode.mono.common.plan import ( ++ BLOCKS, ++ THREADS, ++ WAVES, ++ pair_layout, ++) ++from vllm.models.kimi_k3.amd.specialized_decode.mono.common.ranks import sum_partials ++from vllm.models.kimi_k3.amd.specialized_decode.mono.common.sync import publish ++from vllm.models.kimi_k3.amd.specialized_decode.mono.stages import gemv as gemv ++ ++REGIONS = ( ++ "sgu", ++ "h", ++ "rpart", ++ "lpart", ++ "route", ++ "inter", ++ "isc", ++ "lmsq", ++ "lnrdy", ++ "shrdy", ++ "rlog", ++ "lat", ++ "latr", ++ "final", ++) ++ ++ ++TP = 8 ++HIDDEN = 7168 ++LAT = 3584 ++E = 896 ++TOPK = 16 ++RI = 384 # routed intermediate a rank ++SI = 768 # shared intermediate a rank ++ROWS = 16 ++LOG2E = 1.0 / math.log(2.0) ++BETA, LBETA = 4.0, 25.0 ++ ++R_GROUPS = E // TP // ROWS # 7 router row groups a rank ++R_PARTS = 4 # K quarters: the router is on the path to the route ++L_GROUPS = LAT // TP // ROWS # 28 latent row groups a rank ++L_PARTS = 2 # K halves: the latent is on the path to ug ++R_TASKS = R_GROUPS * R_PARTS # 28 ++L_TASKS = L_GROUPS * L_PARTS # 56 ++G_TASKS = 2 * SI // ROWS # 96 ++EARLY = R_TASKS + L_TASKS + G_TASKS # 180 ++LATE0 = EARLY # route from here ++H_CTA = LATE0 + 8 # the shared experts' SiTU ++SPARE0 = H_CTA + 8 # 196: CTAs with no early task run the shared down first ++SD_TASKS = HIDDEN // ROWS # 448 ++UP_N = HIDDEN // TP # 896 columns a rank up-projects ++UP_TASKS = UP_N // ROWS # 56 ++LN_TASKS = LAT // 128 # 28 ++LN0 = 56 # lnorm's CTAs ++RED_TASKS = HIDDEN // 128 # 56 ++RED0 = 200 ++UG_GROUPS = RI // 32 # 12 ++W13_K = LAT # 3584 ++W13_STEPS = W13_K // 128 # 28 ++W2_STEPS = RI // 128 # 3 ++DOWN_TASKS = LAT // ROWS # 224 ++DOWN_BATCH = 5 # a wave's experts whose w2 is in flight together ++MAX_U = 8 * TOPK ++ ++_STREAM = fx.Stream(None) ++ ++ ++def peer_pairs(s): ++ """One parity's peer regions, pairs: name -> count.""" ++ return { ++ "rlog": TP * s * (E // TP), ++ "lat": s * LAT // 2, ++ "latr": TP * s * LAT // 2, ++ "final": TP * s * HIDDEN // 2, ++ } ++ ++ ++def scratch_layout(s: int) -> dict: ++ pairs = { ++ "sgu": s * SI, ++ "h": s * SI // 2, ++ "rpart": R_GROUPS * R_PARTS * s * ROWS, ++ "lpart": L_GROUPS * L_PARTS * s * ROWS, ++ "route": s * 2 * TOPK, ++ "inter": s * TOPK * 48, ++ "isc": s * TOPK * UG_GROUPS, ++ "lmsq": LN_TASKS * s, ++ "lnrdy": LN_TASKS, ++ "shrdy": UP_TASKS, ++ } ++ lay = pair_layout(pairs.items()) ++ end = max(o + n for o, n in lay.values()) ++ lay["ln"] = ((end + 15) // 16 * 16, s * LAT * 2) ++ end = lay["ln"][0] + lay["ln"][1] ++ lay["shp"] = ((end + 15) // 16 * 16, s * UP_N * 2) ++ return lay ++ ++ ++def f32_of(w): ++ return w.bitcast(fx.Float32) ++ ++ ++def tanh_libm(x): ++ """The ``tanhf`` of x to a few ulp: (1 - e) / (1 + e), e = exp(-2|x|), x's sign.""" ++ e = hw_exp2(abs(x) * fx.Float32(-2.0 * LOG2E)) ++ t = (fx.Float32(1.0) - e) / (fx.Float32(1.0) + e) ++ return (x < fx.Float32(0.0)).select(-t, t) ++ ++ ++def situ_vllm(g, u): ++ """The ``situ_and_mul`` of vLLM on bf16 g, u: ++ ((2 tanh(g/4)) (1 + tanh(g/2))) (25 tanh(u/25)).""" ++ a = fx.Float32(0.5 * BETA) * tanh_libm(g * fx.Float32(1.0 / BETA)) ++ a = a * (fx.Float32(1.0) + tanh_libm(g * fx.Float32(0.5))) ++ return a * (fx.Float32(LBETA) * tanh_libm(u * fx.Float32(1.0 / LBETA))) ++ ++ ++def tanh_aiter(x): ++ """The ``act.tanh_f32`` of aiter: (1 - e) rcp(1 + e), e = exp2(|x| (-2 log2 e)).""" ++ e = hw_exp2(abs(x) * fx.Float32(-2.0 * LOG2E)) ++ t = (fx.Float32(1.0) - e) * hw_rcp(fx.Float32(1.0) + e) ++ return (x > fx.Float32(0.0)).select(t, -t) ++ ++ ++def situ_aiter(g, u): ++ """The situv2 epilogue of aiter on fp32 accumulators.""" ++ sg = hw_rcp(fx.Float32(1.0) + hw_exp2(g * fx.Float32(-LOG2E))) ++ a = (fx.Float32(BETA) * tanh_aiter(g * fx.Float32(1.0 / BETA))) * sg ++ return (a * fx.Float32(LBETA)) * tanh_aiter(u * fx.Float32(1.0 / LBETA)) ++ ++ ++def e8m0_roundup(amax): ++ """The ``_e8m0_roundup(amax, 6)`` of aiter: ceil code of amax * f32(1/6), ++ at most 254.""" ++ wi = (amax * fx.Float32(1.0 / 6.0)).bitcast(fx.Int32) ++ b = fx.Int32(fx.Uint32(wi + fx.Int32(0x7FFFFF)) >> fx.Uint32(23)) & 0xFF ++ return fx.min(b, fx.Int32(254)) ++ ++ ++def pk_fp4_f32(old, a, b, scale, j): ++ """Two f32 / scale -> E2M1 nibbles into byte j of ``old`` (RNE, saturating).""" ++ return fx.Int32( ++ rocdl.cvt_scalef32_pk_fp4_f32( ++ T.i32, ++ as_ir_value(fx.Int32(old)), ++ as_ir_value(fx.Float32(a)), ++ as_ir_value(fx.Float32(b)), ++ as_ir_value(fx.Float32(scale)), ++ j, ++ ) ++ ) ++ ++ ++def pk_fp4_bf16(old, word, scale, j): ++ """A packed bf16 pair word / scale -> E2M1 nibbles into byte j of ``old``.""" ++ v = fx.Vector.from_elements([fx.Int32(word)], fx.Int32).bitcast(fx.BFloat16) ++ return fx.Int32( ++ rocdl.cvt_scalef32_pk_fp4_bf16( ++ T.i32, ++ as_ir_value(fx.Int32(old)), ++ as_ir_value(v), ++ as_ir_value(fx.Float32(scale)), ++ j, ++ ) ++ ) ++ ++ ++# ---------------------------------------------------------------- early ++@traced ++def early_out(c, kind, i, q, red): ++ """Task (i, q) of ``kind`` (0 router, 1 latent, 2 sgu): its 16 rows. A K-split ++ kind's parts q > 0 publish fp32 partials; part 0 adds them in part order and ++ publishes the rows (router: fp32 to every rank's RLOG, latent: bf16 to every ++ rank's LAT); sgu bf16 to SGU.""" ++ s, tid = c["S"], c["tid"] ++ if const_expr(kind == 2): ++ if tid < (ROWS // 2) * s: ++ rp = tid % (ROWS // 2) ++ t = tid // (ROWS // 2) ++ v0 = bf16_round(row_sum(red, 2 * rp, t)) ++ v1 = bf16_round(row_sum(red, 2 * rp + 1, t)) ++ c["put_bf"](c["sgu"], t * 2 * SI + i * ROWS + 2 * rp, [v0, v1]) ++ else: ++ parts = R_PARTS if kind == 0 else L_PARTS ++ region = c["rpart"] if kind == 0 else c["lpart"] ++ if tid < ROWS * s: ++ r = tid % ROWS ++ t = tid // ROWS ++ v = row_sum(red, r, t) ++ at = lambda qq: ((i * parts + qq) * s + t) * ROWS + r # noqa: E731 ++ if q > 0: ++ c["put"](region, at(q), v) ++ if q == 0: ++ got = c["poll"]([(region, at(qq), 1) for qq in range(1, parts)]) ++ for qq in range_constexpr(parts - 1): ++ v = v + f32_of(got[qq][0]) ++ fx.ptr_store(v, c["av"] + tid) ++ gpu.barrier() ++ if const_expr(kind == 0): ++ if (q == 0) & (tid < ROWS * s): ++ r = tid % ROWS ++ t = tid // ROWS ++ v = fx.ptr_load(c["av"] + tid) ++ for p in range_constexpr(TP): ++ c["put"]( ++ c["peer"]["rlog"][p], ++ (c["rank"] * s + t) * (E // TP) + i * ROWS + r, ++ v, ++ ) ++ else: ++ if (q == 0) & (tid < (ROWS // 2) * s): ++ rp = tid % (ROWS // 2) ++ t = tid // (ROWS // 2) ++ v0 = bf16_round(fx.ptr_load(c["av"] + (t * ROWS + 2 * rp))) ++ v1 = bf16_round(fx.ptr_load(c["av"] + (t * ROWS + 2 * rp + 1))) ++ col = c["rank"] * (LAT // TP) + i * ROWS + 2 * rp ++ for p in range_constexpr(TP): ++ c["put_bf"](c["peer"]["lat"][p], t * LAT + col, [v0, v1]) ++ gpu.barrier() ++ ++ ++def early_shape(kind): ++ """(K a part, parts) of an early kind: router a quarter, latent a half, sgu all.""" ++ return { ++ 0: (HIDDEN // R_PARTS, R_PARTS), ++ 1: (HIDDEN // L_PARTS, L_PARTS), ++ 2: (HIDDEN, 1), ++ }[kind] ++ ++ ++def early_loads(c, kind, i, q): ++ """Early task (i, q) of ``kind``'s weights for this wave: issued as early as ++ the CTA can, then ``early_compute`` takes them.""" ++ a = c["args"] ++ klen, _ = early_shape(kind) ++ if kind == 0: ++ w, row0 = a["w_gate"], c["rank"] * (E // TP) + i * ROWS ++ elif kind == 1: ++ w, row0 = a["w_ld"], c["rank"] * (LAT // TP) + i * ROWS ++ else: ++ w, row0 = a["w_sgu"], i * ROWS ++ return gemv.loads(c["lane"], c["wave"], w, row0, klen, ldk=HIDDEN, k0=q * klen) ++ ++ ++@traced ++def early_compute(c, kind, i, q, ops): ++ """Early task (i, q): x's K slice (the input row, or -- ``xflags`` -- the ++ row this launch's AttnRes stage publishes) against ``ops`` -> ``early_out``.""" ++ s, tid, lane, wave, red = c["S"], c["tid"], c["lane"], c["wave"], c["red"] ++ klen, _ = early_shape(kind) ++ xrow = klen + gemv.LDS_PAD ++ if const_expr(c.get("xflags") is not None): ++ if tid < c["xflags"]: ++ c["poll"]([(c["xrdy"], tid, 1)]) ++ gpu.barrier() ++ gemv.rows_to_lds( ++ tid, c["xbuf"], klen, s, c["xl"], xrow, CM_DEV, ldk=HIDDEN, k0=q * klen ++ ) ++ else: ++ gemv.rows_to_lds( ++ tid, c["args"]["x"], klen, s, c["xl"], xrow, ldk=HIDDEN, k0=q * klen ++ ) ++ acc = gemv.mfmas(lane, wave, c["xl"], xrow, klen, s, ops) ++ fx.ptr_store(acc, red + (wave * 64 + lane) * 4) ++ gpu.barrier() ++ early_out(c, kind, i, q, red) ++ ++ ++@traced ++def stage_early(c, bid): ++ """K2b alone: CTA bid's early task (router parts, latent parts, sgu).""" ++ lb = bid - R_TASKS ++ gb = bid - R_TASKS - L_TASKS ++ if bid < R_TASKS: ++ early_compute( ++ c, ++ 0, ++ bid // R_PARTS, ++ bid % R_PARTS, ++ early_loads(c, 0, bid // R_PARTS, bid % R_PARTS), ++ ) ++ if (bid >= R_TASKS) & (bid < R_TASKS + L_TASKS): ++ early_compute( ++ c, ++ 1, ++ lb // L_PARTS, ++ lb % L_PARTS, ++ early_loads(c, 1, lb // L_PARTS, lb % L_PARTS), ++ ) ++ if bid >= R_TASKS + L_TASKS: ++ early_compute(c, 2, gb, 0, early_loads(c, 2, gb, 0)) ++ ++ ++# ---------------------------------------------------------------- route ++def _okey(v): ++ """An order-preserving uint32 key of an f32 (larger float, larger key).""" ++ b = v.bitcast(fx.Int32) ++ neg = b < fx.Int32(0) ++ return fx.Uint32(neg.select(~b, b ^ fx.Int32(-2147483648))) ++ ++ ++@traced ++def stage_route(c, t): ++ """Token t (wave 0): 896 logits, 14 a lane (expert lane + 64 i) -> sigmoid ++ (v_rcp) + bias -> the top-16 by a radix select (the 16th largest key by a ++ 32-step bit search of ballot counts; equal keys: the lower experts) -> the ++ 16 ranked by score, descending (aiter's slot order) -> the pre-bias ++ sigmoids renormalized in that order (IEEE 1 / sum).""" ++ lane, wave, red = c["lane"], c["wave"], c["red"] ++ a = c["args"] ++ if wave == 0: ++ s = c["S"] ++ per = E // 64 ++ specs = [] ++ for i in range_constexpr(per): ++ e = lane + 64 * i ++ src = e // (E // TP) ++ specs.append( ++ (c["peer"]["rlog_own"], (src * s + t) * (E // TP) + e % (E // TP), 1) ++ ) ++ got = c["poll"](specs) ++ keys, sigs = [], [] ++ for i in range_constexpr(per): ++ e = lane + 64 * i ++ lg = f32_of(got[i][0]) ++ sg = hw_rcp(fx.Float32(1.0) + hw_exp2(lg * fx.Float32(-LOG2E))) ++ sc = sg + bo_ld_f32(a["bias"], e) ++ # a NaN logit (a warmup's garbage row) ranks last ++ keys.append(_okey((sc == sc).select(sc, fx.Float32(-3.0e38)))) ++ sigs.append((sg == sg).select(sg, fx.Float32(0.0))) ++ ++ def count_ge(th): ++ n = fx.Int32(0) ++ for i in range_constexpr(per): ++ n = n + popcount(ballot(keys[i] >= th)) ++ return n ++ ++ th = fx.Uint32(0) ++ for bit in range_constexpr(31, -1, -1): ++ cand = th | fx.Uint32(1 << bit) ++ th = (count_ge(cand) >= TOPK).select(cand, th) ++ # selected: key > th, then the lowest experts with key == th ++ n_gt = fx.Int32(0) ++ for i in range_constexpr(per): ++ n_gt = n_gt + popcount(ballot(keys[i] > th)) ++ need = fx.Int32(TOPK) - n_gt ++ eq_before = fx.Int32(0) ++ sel_before = fx.Int32(0) ++ for i in range_constexpr(per): ++ eqb = ballot(keys[i] == th) ++ eq_rank = eq_before + lanes_below(eqb, lane) ++ sel = (keys[i] > th) | ((keys[i] == th) & (eq_rank < need)) ++ selb = ballot(sel) ++ pos = sel_before + lanes_below(selb, lane) ++ if sel: ++ fx.ptr_store(keys[i].bitcast(fx.Float32), red + pos * 3) ++ fx.ptr_store( ++ fx.Int32(lane + 64 * i).bitcast(fx.Float32), red + (pos * 3 + 1) ++ ) ++ fx.ptr_store(sigs[i], red + (pos * 3 + 2)) ++ eq_before = eq_before + popcount(eqb) ++ sel_before = sel_before + popcount(selb) ++ # rank the 16 by key descending, the lower expert first on a tie ++ q = fx.min(lane, TOPK - 1) ++ mk = fx.Uint32(fx.ptr_load(red + q * 3).bitcast(fx.Int32)) ++ me = fx.ptr_load(red + (q * 3 + 1)).bitcast(fx.Int32) ++ ms = fx.ptr_load(red + (q * 3 + 2)) ++ rank = fx.Int32(0) ++ for j in range_constexpr(TOPK): ++ ok = fx.Uint32(fx.ptr_load(red + j * 3).bitcast(fx.Int32)) ++ oe = fx.ptr_load(red + (j * 3 + 1)).bitcast(fx.Int32) ++ before = (ok > mk) | ((ok == mk) & (oe < me)) ++ rank = rank + before.select(fx.Int32(1), fx.Int32(0)) ++ # the sigmoids in rank order, summed serially (aiter's renorm) ++ if lane < TOPK: ++ fx.ptr_store(ms, red + (64 + rank)) ++ total = fx.Float32(0.0) ++ for j in range_constexpr(TOPK): ++ total = total + fx.ptr_load(red + (64 + j)) ++ inv = fx.Float32(1.0) / total ++ if lane < TOPK: ++ c["put"]( ++ c["route"], t * 2 * TOPK + rank, fx.min(fx.max(me, 0), fx.Int32(E - 1)) ++ ) ++ c["put"](c["route"], t * 2 * TOPK + TOPK + rank, ms * inv) ++ gpu.barrier() ++ ++ ++def bo_ld_f32(ptr, i): ++ return fx.Float32(bo.buffer_load(rsrc(ptr), i, vec_width=1, dtype=T.f32)) ++ ++ ++@traced ++def load_route(c): ++ """ROUTE -> LDS (ids, weights); the union U: ``slot_of[e]``, ++ ``expert_of[slot]``, ``kof[slot][t]`` (the token's top-k index, -1 if not ++ picked); returns |U|.""" ++ s, tid, lane, wave = c["S"], c["tid"], c["lane"], c["wave"] ++ rt = c["rt"] ++ n = s * 2 * TOPK ++ if tid < n: ++ v = c["poll"]([(c["route"], tid, 1)])[0][0] ++ fx.ptr_store(v, rt["route"] + tid) ++ for i in range_constexpr((E + THREADS - 1) // THREADS): ++ e = tid + THREADS * i ++ if e < E: ++ fx.ptr_store(fx.Int32(0), rt["flag"] + e) ++ for i in range_constexpr((MAX_U * 8 + THREADS - 1) // THREADS): ++ j = tid + THREADS * i ++ if j < MAX_U * 8: ++ fx.ptr_store(fx.Int32(-1), rt["kof"] + j) ++ gpu.barrier() ++ if tid < s * TOPK: ++ t = tid // TOPK ++ k = tid % TOPK ++ e = fx.max(fx.min(fx.ptr_load(rt["route"] + (t * 2 * TOPK + k)), E - 1), 0) ++ fx.ptr_store(fx.Int32(1), rt["flag"] + e) ++ gpu.barrier() ++ e0 = fx.min(2 * tid, E - 1) ++ f0 = (2 * tid < E).select(fx.ptr_load(rt["flag"] + e0), fx.Int32(0)) ++ f1 = (2 * tid + 1 < E).select( ++ fx.ptr_load(rt["flag"] + fx.min(2 * tid + 1, E - 1)), fx.Int32(0) ++ ) ++ off, tot = block_excl_scan(f0 + f1, lane, wave, c["scan"]) ++ gpu.barrier() ++ if f0 == 1: ++ fx.ptr_store(off, rt["flag"] + 2 * tid) ++ fx.ptr_store(fx.Int32(2 * tid), rt["expert"] + off) ++ if f1 == 1: ++ fx.ptr_store(off + f0, rt["flag"] + 2 * tid + 1) ++ fx.ptr_store(fx.Int32(2 * tid + 1), rt["expert"] + off + f0) ++ gpu.barrier() ++ if tid < s * TOPK: ++ t = tid // TOPK ++ k = tid % TOPK ++ e = fx.max(fx.min(fx.ptr_load(rt["route"] + (t * 2 * TOPK + k)), E - 1), 0) ++ slot = fx.ptr_load(rt["flag"] + e) ++ fx.ptr_store(fx.Int32(k), rt["kof"] + (slot * 8 + t)) ++ gpu.barrier() ++ return uniform(tot) ++ ++ ++@traced ++def load_latq(c): ++ """LAT (every rank's 448 columns) -> MXFP4 as aiter's inline quant: a ++ thread a (token, 32-block): amax of |bf16| bits, the ceil code of amax / 6, ++ v_cvt_scalef32_pk_fp4_bf16 -> latq words, latsc codes.""" ++ s, tid = c["S"], c["tid"] ++ nblk = s * (LAT // 32) ++ for i in range_constexpr((nblk + THREADS - 1) // THREADS): ++ b = fx.min(tid + THREADS * i, nblk - 1) ++ t = b // (LAT // 32) ++ kb = b % (LAT // 32) ++ base = (t * LAT + kb * 32) // 2 ++ got = c["poll"]([(c["peer"]["lat_own"], base + j, 1) for j in range(16)]) ++ ws = [got[j][0] for j in range(16)] ++ am = fx.Int32(0) ++ for w in ws: ++ am = fx.max(am, w & 0x7FFF) ++ am = fx.max(am, (w >> 16) & 0x7FFF) ++ amax = (am << 16).bitcast(fx.Float32) ++ code = e8m0_roundup(amax) ++ scale = (code << 23).bitcast(fx.Float32) ++ if tid + THREADS * i < nblk: ++ for q in range_constexpr(4): ++ pk = fx.Int32(0) ++ for j in range_constexpr(4): ++ pk = pk_fp4_bf16(pk, ws[4 * q + j], scale, j) ++ fx.ptr_store(pk, c["latq"] + (t * (LAT // 8) + kb * 4 + q)) ++ fx.ptr_store(code, c["latsc"] + (t * (LAT // 32) + kb)) ++ gpu.barrier() ++ ++ ++# ---------------------------------------------------------------- ug ++def w13_scale_word(e, rg, st, lane): ++ """The dword of aiter's gate_up=False ``shuffle_scale`` holding (row rg 16 + ++ lane % 16, K block 4 st + lane / 16) for st's pair and rg's pair: byte ++ (st % 2) 2 + rg % 2.""" ++ n1 = rg // 2 ++ k1 = st // 2 ++ return ( ++ e * (2 * RI * (W13_K // 32) // 4) ++ + ((n1 * (W13_STEPS // 2) + k1) * 4 + lane // 16) * 16 ++ + lane % 16 ++ ) ++ ++ ++def ug_loads(c, task): ++ """Task ``task``'s w13 operands for this wave (its row group, its K half) ++ and their scale words: issued a task ahead (``stage_ug`` takes them).""" ++ lane, wave = c["lane"], c["wave"] ++ a, rt = c["args"], c["rt"] ++ slot = task // UG_GROUPS ++ g = task % UG_GROUPS ++ e = uniform(fx.ptr_load(rt["expert"] + slot)) ++ rgl = wave % 4 # 0, 1 gate; 2, 3 up ++ half = wave // 4 ++ rg = (rgl // 2) * (RI // ROWS) + 2 * g + rgl % 2 ++ r_w = rsrc(a["w13"]) ++ base = e * (2 * RI * W13_K // 8) ++ steps = [half * (W13_STEPS // 2) + i for i in range(W13_STEPS // 2)] ++ av = [fp4_tile_load(r_w, base, rg, W13_K, st, lane * 4) for st in steps] ++ r_s = rsrc(a["w13s"]) ++ sw = [ ++ fx.Int32( ++ bo.buffer_load( ++ r_s, w13_scale_word(e, rg, st, lane), vec_width=1, dtype=T.i32 ++ ) ++ ) ++ for st in steps[::2] ++ ] ++ return av + sw ++ ++ ++@traced ++def stage_ug(c, task, ops): ++ """Slot task / 12, 32-column group task % 12: gate rows 32 g .. and up rows ++ 384 + 32 g .. of the slot's expert (``ops``: ``ug_loads``) against the MXFP4 ++ latent; SiTUv2; MXFP4 a 32-column group -> INTER / ISC of each picking ++ token.""" ++ s, tid, lane, wave, red = c["S"], c["tid"], c["lane"], c["wave"], c["red"] ++ rt = c["rt"] ++ slot = task // UG_GROUPS ++ g = task % UG_GROUPS ++ rgl = wave % 4 ++ half = wave // 4 ++ rg = (rgl // 2) * (RI // ROWS) + 2 * g + rgl % 2 ++ steps = [half * (W13_STEPS // 2) + i for i in range(W13_STEPS // 2)] ++ av = ops[: len(steps)] ++ sw = ops[len(steps) :] ++ t = fx.min(lane % 16, s - 1) ++ kb = lane // 16 ++ acc = fx.Vector.filled(4, 0.0, fx.Float32) ++ for i in range_constexpr(len(steps)): ++ st = steps[i] ++ bw = fx.Vector( ++ fx.ptr_load( ++ c["latq"] + (t * (LAT // 8) + st * 16 + kb * 4), ++ result_type=fx.Vector.make_type(4, fx.Int32), ++ ) ++ ) ++ sb = fx.ptr_load(c["latsc"] + (t * (LAT // 32) + st * 4 + kb)) ++ # the code's byte (st % 2) 2 + rg % 2 of the word; rg % 2 is the wave's ++ sa = (sw[i // 2] >> ((st % 2) * 16 + (rg % 2) * 8)) & 0xFF ++ acc = mfma_scaled(av[i], bw, acc, sa, sb, FP4, FP4, 0, 0) ++ fx.ptr_store(acc, red + (wave * 64 + lane) * 4) ++ gpu.barrier() ++ # thread (token, column): g / u summed over the K halves, SiTUv2 ++ tt = tid // 32 ++ col = tid % 32 ++ live = tt < s ++ tt = fx.min(tt, s - 1) ++ first = col < 16 # gate / up row group 0 (waves 0, 2) or 1 (1, 3); +4: K's 2nd half ++ r = col % 16 ++ gv = first.select( ++ row_sum(red, r, tt, waves=(0, 4)), row_sum(red, r, tt, waves=(1, 5)) ++ ) ++ uv = first.select( ++ row_sum(red, r, tt, waves=(2, 6)), row_sum(red, r, tt, waves=(3, 7)) ++ ) ++ v = situ_aiter(gv, uv) ++ amax = butterfly(abs(v), (16, 8, 4, 2, 1), fx.max) ++ code = e8m0_roundup(amax) ++ gpu.barrier() ++ fx.ptr_store(v, c["av"] + tid) ++ gpu.barrier() ++ k = fx.ptr_load(rt["kof"] + (slot * 8 + tt)) ++ if live & (col < 4) & (k >= 0): ++ scale = (code << 23).bitcast(fx.Float32) ++ pk = fx.Int32(0) ++ for j in range_constexpr(4): ++ x0 = fx.ptr_load(c["av"] + (tt * 32 + col * 8 + 2 * j)) ++ x1 = fx.ptr_load(c["av"] + (tt * 32 + col * 8 + 2 * j + 1)) ++ pk = pk_fp4_f32(pk, x0, x1, scale, j) ++ pick = tt * TOPK + k ++ c["put_words"](c["inter"], pick * 48 + g * 4 + col, [pk]) ++ if col == 0: ++ c["put"](c["isc"], pick * UG_GROUPS + g, code) ++ gpu.barrier() ++ ++ ++# ---------------------------------------------------------------- queues ++@traced ++def load_inter(c): ++ """Every pick's INTER words and ISC codes -> LDS.""" ++ s, tid = c["S"], c["tid"] ++ nw = s * TOPK * 48 ++ per = (nw + THREADS - 1) // THREADS ++ got = c["poll"]( ++ [(c["inter"], fx.min(tid + THREADS * i, nw - 1), 1) for i in range(per)] ++ ) ++ for i in range_constexpr(per): ++ if tid + THREADS * i < nw: ++ fx.ptr_store(got[i][0], c["interq"] + (tid + THREADS * i)) ++ nc = s * TOPK * UG_GROUPS ++ perc = (nc + THREADS - 1) // THREADS ++ gc = c["poll"]( ++ [(c["isc"], fx.min(tid + THREADS * i, nc - 1), 1) for i in range(perc)] ++ ) ++ for i in range_constexpr(perc): ++ if tid + THREADS * i < nc: ++ fx.ptr_store(gc[i][0], c["intsc"] + (tid + THREADS * i)) ++ gpu.barrier() ++ ++ ++def w2_scale(c, e, rgd, st, lane): ++ """The e8m0 code of w2 (row e 3584 + rgd 16 + lane % 16, K block 4 st + ++ lane / 16) in ``e8m0_shuffle``'s layout (16 codes a row).""" ++ row = e * LAT + rgd * ROWS + lane % 16 ++ idx = scale_index(row, st * 4 + lane // 16, 16) ++ w = fx.Int32( ++ bo.buffer_load(rsrc(c["args"]["w2s"]), idx // 4, vec_width=1, dtype=T.i32) ++ ) ++ return (w >> ((idx % 4) * 8)) & 0xFF ++ ++ ++def poll_inter(c, pick, kb): ++ """A lane's INTER words (3 steps x 4) and ISC codes of ``pick``, polled at ++ its use: a slot whose ug tasks are done is multiplied while later slots' ++ still run (no wait for all of U).""" ++ return c["poll"]( ++ [ ++ (c["inter"], pick * 48 + st * 16 + kb * 4 + d, 1) ++ for st in range(W2_STEPS) ++ for d in range(4) ++ ] ++ + [(c["isc"], pick * UG_GROUPS + st * 4 + kb, 1) for st in range(W2_STEPS)] ++ ) ++ ++ ++def down_loads(c, rgd, nu, b0): ++ """The w2 tiles and scale codes of a wave's batch from slot b0 (slots b0, ++ b0 + 8, ...; past |U| the last slot's): flat, ``stage_down`` unpacks.""" ++ lane = c["lane"] ++ r_w = rsrc(c["args"]["w2"]) ++ rt = c["rt"] ++ slots = [fx.min(b0 + WAVES * j, nu - 1) for j in range(DOWN_BATCH)] ++ es = [uniform(fx.ptr_load(rt["expert"] + sl)) for sl in slots] ++ av = [ ++ fp4_tile_load(r_w, es[j] * (LAT * RI // 8), rgd, RI, st, lane * 4) ++ for j in range(DOWN_BATCH) ++ for st in range(W2_STEPS) ++ ] ++ sa = [ ++ w2_scale(c, es[j], rgd, st, lane) ++ for j in range(DOWN_BATCH) ++ for st in range(W2_STEPS) ++ ] ++ return av + sa ++ ++ ++@traced ++def stage_down(c, rgd, nu, ops0): ++ """Latent rows 16 rgd ..: every slot's w2 against its picks (wave w the slots ++ w, w + 8, ...; a batch of DOWN_BATCH in flight, the next issued under this ++ one's MFMAs; ``ops0`` the first, issued before INTER was waited for) -> ++ bf16(w acc) a (token, k) in LDS -> folded in top-k order with a bf16 ++ rounding an add -> every rank's LATR.""" ++ s, tid, lane, wave = c["S"], c["tid"], c["lane"], c["wave"] ++ rt = c["rt"] ++ t = fx.min(lane % 16, s - 1) ++ kb = lane // 16 ++ nbs = DOWN_BATCH * W2_STEPS ++ b0 = wave ++ ops = ops0 ++ while b0 < nu: ++ nb = b0 + WAVES * DOWN_BATCH ++ ops_n = down_loads(c, rgd, nu, fx.min(nb, nu - 1)) ++ slots = [fx.min(b0 + WAVES * j, nu - 1) for j in range(DOWN_BATCH)] ++ for j in range_constexpr(DOWN_BATCH): ++ live = (b0 + WAVES * j) < nu ++ k = fx.ptr_load(rt["kof"] + (slots[j] * 8 + t)) ++ picked = (k >= 0) & (lane % 16 < s) ++ pick = t * TOPK + fx.max(k, 0) ++ acc = fx.Vector.filled(4, 0.0, fx.Float32) ++ # this slot's picks' INTER, polled here (not all of U first): a slot ++ # whose ug tasks are done is multiplied while later ones still run ++ for st in range_constexpr(W2_STEPS): ++ bw = fx.Vector( ++ fx.ptr_load( ++ c["interq"] + (pick * 48 + st * 16 + kb * 4), ++ result_type=fx.Vector.make_type(4, fx.Int32), ++ ) ++ ) ++ scode = fx.ptr_load(c["intsc"] + (pick * UG_GROUPS + st * 4 + kb)) ++ bw = fx.Vector.from_elements( ++ [picked.select(bw[d], fx.Int32(0)) for d in range(4)], fx.Int32 ++ ) ++ sb = picked.select(scode, fx.Int32(UNIT_SCALE)) ++ acc = mfma_scaled( ++ ops[j * W2_STEPS + st], ++ bw, ++ acc, ++ ops[nbs + j * W2_STEPS + st], ++ sb, ++ FP4, ++ FP4, ++ 0, ++ 0, ++ ) ++ # lane: rows 4 (lane / 16) + i of token lane % 16 ++ if live & picked: ++ wt = fx.ptr_load(rt["route"] + (t * 2 * TOPK + TOPK + k)).bitcast( ++ fx.Float32 ++ ) ++ for i in range_constexpr(4): ++ row = 4 * (lane // 16) + i ++ fx.ptr_store( ++ bf16_round(acc[i] * wt), ++ c["contrib"] + ((t * TOPK + k) * ROWS + row), ++ ) ++ b0 = nb ++ ops = ops_n ++ gpu.barrier() ++ if tid < (ROWS // 2) * s: ++ rp = tid % (ROWS // 2) ++ tt = tid // (ROWS // 2) ++ y0 = fx.Float32(0.0) ++ y1 = fx.Float32(0.0) ++ for k in range_constexpr(TOPK): ++ y0 = bf16_round( ++ y0 + fx.ptr_load(c["contrib"] + ((tt * TOPK + k) * ROWS + 2 * rp)) ++ ) ++ y1 = bf16_round( ++ y1 + fx.ptr_load(c["contrib"] + ((tt * TOPK + k) * ROWS + 2 * rp + 1)) ++ ) ++ for p in range_constexpr(TP): ++ c["put_bf"]( ++ c["peer"]["latr"][p], ++ (c["rank"] * s + tt) * LAT + rgd * ROWS + 2 * rp, ++ [y0, y1], ++ ) ++ gpu.barrier() ++ ++ ++# ---------------------------------------------------------------- shared down ++@traced ++def stage_h(c, t): ++ """Token t's shared h = SiTU(g, u) (vLLM's situ_and_mul, bf16) of SGU, once ++ -> H; a thread a column pair, its polls in one batch.""" ++ tid = c["tid"] ++ n = SI // 2 ++ if tid < n: ++ got = c["poll"]( ++ [ ++ (c["sgu"], (t * 2 * SI) // 2 + tid, 1), ++ (c["sgu"], (t * 2 * SI + SI) // 2 + tid, 1), ++ ] ++ ) ++ g0, g1 = bf_lo(got[0][0]), bf_hi(got[0][0]) ++ u0, u1 = bf_lo(got[1][0]), bf_hi(got[1][0]) ++ c["put_bf"](c["h"], t * SI + 2 * tid, [situ_vllm(g0, u0), situ_vllm(g1, u1)]) ++ gpu.barrier() ++ ++ ++@traced ++def load_h(c): ++ """H -> LDS rows (for the shared down GEMV), once a CTA.""" ++ s, tid = c["S"], c["tid"] ++ hrow = SI + gemv.LDS_PAD ++ n = s * SI // 2 ++ per = (n + THREADS - 1) // THREADS ++ got = c["poll"]([(c["h"], fx.min(tid + THREADS * i, n - 1), 1) for i in range(per)]) ++ for i in range_constexpr(per): ++ pidx = tid + THREADS * i ++ t = fx.min(pidx, n - 1) // (SI // 2) ++ cp = fx.min(pidx, n - 1) % (SI // 2) ++ if pidx < n: ++ fx.ptr_store(got[i][0], c["xl"] + (t * hrow + 2 * cp) // 2) ++ gpu.barrier() ++ ++ ++@traced ++def stage_sdown(c, j, ops): ++ """Shared down rows 16 j .. (K 768) against h in LDS (``load_h``): bf16 ++ partial -> SHP (this rank's up columns) or every rank's FINAL.""" ++ s, tid, lane, wave, red = c["S"], c["tid"], c["lane"], c["wave"], c["red"] ++ hrow = SI + gemv.LDS_PAD ++ acc = gemv.mfmas(lane, wave, c["xl"], hrow, SI, s, ops) ++ fx.ptr_store(acc, red + (wave * 64 + lane) * 4) ++ gpu.barrier() ++ lo = c["rank"] * UP_TASKS ++ mine = (j >= lo) & (j < lo + UP_TASKS) ++ if tid < (ROWS // 2) * s: ++ rp = tid % (ROWS // 2) ++ tt = tid // (ROWS // 2) ++ v0 = bf16_round(row_sum(red, 2 * rp, tt)) ++ v1 = bf16_round(row_sum(red, 2 * rp + 1, tt)) ++ if mine: ++ w = ( ++ fx.Vector.from_elements([v0, v1], fx.Float32) ++ .to(fx.BFloat16) ++ .bitcast(fx.Int32)[0] ++ ) ++ bo.buffer_store( ++ w, ++ rsrc(c["shp"]), ++ (tt * UP_N + (j - lo) * ROWS + 2 * rp) // 2, ++ cache_modifier=CM_DEV, ++ ) ++ else: ++ for p in range_constexpr(TP): ++ c["put_bf"]( ++ c["peer"]["final"][p], ++ (c["rank"] * s + tt) * HIDDEN + j * ROWS + 2 * rp, ++ [v0, v1], ++ ) ++ publish( ++ c["put"], ++ c["shrdy"], ++ fx.max(fx.min(j - lo, UP_TASKS - 1), 0), ++ fx.Int32(1), ++ (tid == 0) & mine, ++ ) ++ gpu.barrier() ++ ++ ++@traced ++def run_sdown(c, bid, w_sd): ++ """This rank's 56 shared-down tasks whose columns the up stage adds to (k -> ++ row group (k + 56 rank) % 448), on the CTAs past the down tasks: up waits ++ on them.""" ++ load_h(c) ++ lo = c["rank"] * UP_TASKS ++ if bid >= DOWN_TASKS: ++ for k in range(bid - DOWN_TASKS, UP_TASKS, BLOCKS - DOWN_TASKS): ++ j = (k + lo) % SD_TASKS ++ stage_sdown(c, j, gemv.loads(c["lane"], c["wave"], w_sd, j * ROWS, SI)) ++ ++ ++@traced ++def run_sdown_rest(c, bid, w_sd): ++ """The other 392 shared-down tasks (only the final reduce waits on them), on ++ the CTAs past lnorm / up's (84 ..), after their down: off the latent-norm -> ++ up critical path.""" ++ r0 = LN0 + LN_TASKS # 84: past the lnorm / up CTAs ++ if bid >= r0: ++ load_h(c) ++ lo = c["rank"] * UP_TASKS ++ for k in range(UP_TASKS + bid - r0, SD_TASKS, BLOCKS - r0): ++ j = (k + lo) % SD_TASKS ++ stage_sdown(c, j, gemv.loads(c["lane"], c["wave"], w_sd, j * ROWS, SI)) ++ ++ ++# ---------------------------------------------------------------- lnorm / up / reduce ++@traced ++def stage_lnorm(c, j): ++ """Columns 128 j ..: the ranks' LATR (rank order, fp32 -> bf16), every ++ task's sum of squares -> rstd -> bf16(acc rstd w) -> LN, LNRDY.""" ++ s, tid, lane, red = c["S"], c["tid"], c["lane"], c["red"] ++ a = c["args"] ++ mine = tid < s * 64 ++ t = fx.min(tid // 64, s - 1) ++ cp = tid % 64 ++ col = j * 128 + 2 * cp ++ lo, hi = sum_partials( ++ c["poll"], ++ c["peer"]["latr_own"], ++ lambda src: ((src * s + t) * LAT + col) // 2, ++ TP, ++ ) ++ lo, hi = bf16_round(lo), bf16_round(hi) ++ sq = wave_sum(lo * lo + hi * hi) ++ if mine & (lane == 0): ++ c["put"](c["lmsq"], j * s + t, sq) ++ gpu.barrier() ++ if mine & (lane < LN_TASKS): ++ v = c["poll"]([(c["lmsq"], lane * s + t, 1)])[0][0] ++ fx.ptr_store(f32_of(v), red + (t * 32 + lane)) ++ gpu.barrier() ++ tot = fx.ptr_load(red + t * 32) ++ for q in range_constexpr(1, LN_TASKS): ++ tot = tot + fx.ptr_load(red + (t * 32 + q)) ++ r = hw_rsq(tot * fx.Float32(1.0 / LAT) + fx.Float32(c["ln_eps"])) ++ w0 = fx.Float32( ++ fx.BFloat16(bo.buffer_load(rsrc(a["ln_w"]), col, vec_width=1, dtype=T.bf16)) ++ ) ++ w1 = fx.Float32( ++ fx.BFloat16(bo.buffer_load(rsrc(a["ln_w"]), col + 1, vec_width=1, dtype=T.bf16)) ++ ) ++ y = ( ++ fx.Vector.from_elements([lo * r * w0, hi * r * w1], fx.Float32) ++ .to(fx.BFloat16) ++ .bitcast(fx.Int32)[0] ++ ) ++ if mine: ++ bo.buffer_store(y, rsrc(c["ln"]), (t * LAT + col) // 2, cache_modifier=CM_DEV) ++ publish(c["put"], c["lnrdy"], j, fx.Int32(1), tid == 0) ++ gpu.barrier() ++ ++ ++@traced ++def stage_up(c, j): ++ """This rank's up rows 16 j .. of its shard (hidden rows 896 rank ..): LN GEMV ++ + SHP -> bf16 -> FINAL.""" ++ s, tid, lane, wave, red = c["S"], c["tid"], c["lane"], c["wave"], c["red"] ++ a = c["args"] ++ row0 = c["rank"] * UP_N + j * ROWS # the hidden column; w_up is the shard ++ ops = gemv.loads(lane, wave, a["w_up"], j * ROWS, LAT) ++ if tid < LN_TASKS: ++ c["poll"]([(c["lnrdy"], tid, 1)]) ++ if tid == 0: ++ c["poll"]([(c["shrdy"], j, 1)]) ++ gpu.barrier() ++ lrow = LAT + gemv.LDS_PAD ++ gemv.dev_rows_to_lds(tid, c["ln"], LAT, s, c["xl"], lrow) ++ acc = gemv.mfmas(lane, wave, c["xl"], lrow, LAT, s, ops) ++ fx.ptr_store(acc, red + (wave * 64 + lane) * 4) ++ gpu.barrier() ++ if tid < (ROWS // 2) * s: ++ rp = tid % (ROWS // 2) ++ tt = tid // (ROWS // 2) ++ w = fx.Int32( ++ bo.buffer_load( ++ rsrc(c["shp"]), ++ (tt * UP_N + j * ROWS + 2 * rp) // 2, ++ vec_width=1, ++ dtype=T.i32, ++ cache_modifier=CM_DEV, ++ ) ++ ) ++ v0 = bf16_round(bf_lo(w) + row_sum(red, 2 * rp, tt)) ++ v1 = bf16_round(bf_hi(w) + row_sum(red, 2 * rp + 1, tt)) ++ for p in range_constexpr(TP): ++ c["put_bf"]( ++ c["peer"]["final"][p], ++ (c["rank"] * s + tt) * HIDDEN + row0 + 2 * rp, ++ [v0, v1], ++ ) ++ gpu.barrier() ++ ++ ++@traced ++def stage_reduce(c, j): ++ """Columns 128 j ..: FINAL in rank order -> bf16 out.""" ++ s, tid = c["S"], c["tid"] ++ if tid < s * 64: ++ t = tid // 64 ++ cp = tid % 64 ++ col = j * 128 + 2 * cp ++ lo, hi = sum_partials( ++ c["poll"], ++ c["peer"]["final_own"], ++ lambda src: ((src * s + t) * HIDDEN + col) // 2, ++ TP, ++ ) ++ w = ( ++ fx.Vector.from_elements([lo, hi], fx.Float32) ++ .to(fx.BFloat16) ++ .bitcast(fx.Int32)[0] ++ ) ++ bo.buffer_store(w, rsrc(c["args"]["out"]), (t * HIDDEN + col) // 2) ++ gpu.barrier() ++ ++ ++# ---------------------------------------------------------------- kernel +diff --git a/vllm/models/kimi_k3/amd/specialized_decode/runtime/__init__.py b/vllm/models/kimi_k3/amd/specialized_decode/runtime/__init__.py +new file mode 100644 +index 000000000..e17a0e3ff +--- /dev/null ++++ b/vllm/models/kimi_k3/amd/specialized_decode/runtime/__init__.py +@@ -0,0 +1,219 @@ ++# SPDX-License-Identifier: Apache-2.0 ++# SPDX-FileCopyrightText: Copyright contributors to the vLLM project ++"""Kimi-K3 specialized decode on MI355X (layer mode), enabled by ++``additional_config={"decode_backend": "latency_specialized"}``. ++ ++``specialized_decode_config`` runs in the model's ``__init__``, before the weights ++load. It returns None for the stock decode, and raises with every reason from ++``unsupported`` when the specialized decode is requested for a deployment it cannot ++serve. Checks that need the loaded weights run when the executor is built ++(``decode.SpecializedDecode``), those that need the KV cache in ``prepare_kv``.""" ++ ++from dataclasses import dataclass ++ ++import torch ++ ++from vllm.logger import init_logger ++ ++logger = init_logger(__name__) ++ ++AUTO, SPECIALIZED = "auto", "latency_specialized" ++TP = 8 ++# a launch's x rows fit one CTA's LDS ++MAX_TOKENS = 8 ++WHOLE, FFN = "whole", "ffn" ++# Kimi-K3's text shapes the kernels are written for (hidden size, routed experts, ++# KDA heads, KDA head dim, KDA short-conv width) ++K3_SHAPES = (7168, 896, 96, 128, 4) ++ ++ ++@dataclass(frozen=True) ++class DecodeConfig: ++ # rows a request verifies: num_speculative_tokens + 1 ++ q: int ++ ++ @property ++ def widths(self) -> tuple[int, ...]: ++ """The rows a launch serves: whole requests, at most MAX_TOKENS rows. Every ++ uniform verify step up to MAX_TOKENS rows is one of them, so an eager step ++ never pads; a FULL replay pads whole requests the way vLLM's stock decode ++ does.""" ++ return tuple(range(self.q, MAX_TOKENS + 1, self.q)) ++ ++ ++def decode_backend(vllm_config) -> str: ++ """The requested decode backend: the one place that reads the switch.""" ++ extra = vllm_config.additional_config ++ return extra.get("decode_backend", AUTO) if isinstance(extra, dict) else AUTO ++ ++ ++def specialized_decode_config(vllm_config) -> DecodeConfig | None: ++ """The deployment's specialized decode settings, None for the stock decode. ++ Raises when the specialized decode is requested for an unsupported deployment, ++ listing every reason.""" ++ backend = decode_backend(vllm_config) ++ if backend == AUTO: ++ return None ++ if backend != SPECIALIZED: ++ raise ValueError( ++ f"decode_backend must be '{AUTO}' or '{SPECIALIZED}', got {backend!r}" ++ ) ++ reasons = unsupported(vllm_config) ++ if reasons: ++ raise ValueError( ++ "Kimi-K3 specialized decode (decode_backend=latency_specialized) does " ++ "not support: " + "; ".join(reasons) + f" (decode_backend={AUTO} runs " ++ "vLLM's decode)" ++ ) ++ if piecewise_graphs_only(vllm_config): ++ logger.warning( ++ "Kimi-K3 specialized decode: PIECEWISE-only CUDA graphs capture every " ++ "decode step as stock; use FULL decode graphs" ++ ) ++ spec = vllm_config.speculative_config ++ return DecodeConfig(q=spec.num_speculative_tokens + 1) ++ ++ ++def unsupported(vllm_config) -> list[str]: ++ """Every reason this deployment cannot run the specialized decode (empty: it ++ can). Config only: the loaded weights are checked when the executor is built. ++ Default-deny: every feature below is refused, the kernels serve none of them.""" ++ from vllm.config import CompilationMode ++ from vllm.model_executor.layers.mamba.mamba_utils import ( ++ MambaStateDtypeCalculator, ++ ) ++ ++ mc, pc, cc = ( ++ vllm_config.model_config, ++ vllm_config.parallel_config, ++ vllm_config.cache_config, ++ ) ++ spec, oc = vllm_config.speculative_config, vllm_config.offload_config ++ checks = ( ++ (pc.tensor_parallel_size != TP, f"TP {pc.tensor_parallel_size} (TP{TP} only)"), ++ (pc.pipeline_parallel_size != 1, "pipeline parallelism"), ++ (pc.data_parallel_size != 1, "data parallelism"), ++ (pc.enable_expert_parallel, "expert parallelism"), ++ (pc.enable_eplb, "EPLB"), ++ ( ++ pc.decode_context_parallel_size not in (1, pc.tensor_parallel_size), ++ "decode context parallelism other than 1 or TP", ++ ), ++ (pc.prefill_context_parallel_size != 1, "prefill context parallelism"), ++ # the kernels read one batch's attention metadata per step ++ (pc.use_ubatching, "DBO (dual-batch overlap)"), ++ (vllm_config.lora_config is not None, "LoRA"), ++ (vllm_config.kv_transfer_config is not None, "KV transfer connectors"), ++ ( ++ vllm_config.aux_output_config.enable_return_routed_experts, ++ "routed-expert capture", ++ ), ++ # the kernels read the loaded weights in place ++ (mc.enable_sleep_mode, "sleep mode"), ++ ( ++ oc.uva.cpu_offload_gb > 0 or oc.prefetch.offload_group_size > 0, ++ "CPU weight offload", ++ ), ++ (vllm_config.weight_transfer_config is not None, "weight updates"), ++ # the kernels serve spec-verify steps of whole uniform requests only ++ (spec is None, "decoding without speculative decoding"), ++ ( ++ spec is not None and spec.num_speculative_tokens + 1 > MAX_TOKENS, ++ f"more than {MAX_TOKENS - 1} speculative tokens", ++ ), ++ ( ++ spec is not None and spec.enable_adaptive_verification, ++ "adaptive verification", ++ ), ++ ( ++ spec is not None and spec.uses_dynamic_speculative_decoding(), ++ "dynamic speculative decoding", ++ ), ++ # K1 keeps a BF16 conv state and an FP32 recurrent state in place; it ++ # maintains no ReplaySSM / RecoverSSM records ++ ( ++ MambaStateDtypeCalculator.kda_state_dtype(mc.dtype, cc.mamba_cache_dtype) ++ != (torch.bfloat16, torch.float32), ++ "KDA cache dtypes other than a BF16 conv state and an FP32 recurrent " ++ "state (mamba cache dtype)", ++ ), ++ (cc.use_replayssm or cc.use_kda_recoverssm, "ReplaySSM / RecoverSSM"), ++ # the model state hooks build the executor and bind it to each KV cache ++ ( ++ not vllm_config.use_v2_model_runner, ++ "the V1 model runner (set VLLM_USE_V2_MODEL_RUNNER=1)", ++ ), ++ # plan() and the layer kernels run in forward; a compiled forward would ++ # trace them ++ ( ++ vllm_config.compilation_config.mode == CompilationMode.STOCK_TORCH_COMPILE, ++ "torch.compile (compilation mode STOCK_TORCH_COMPILE)", ++ ), ++ ) ++ reasons = [why for bad, why in checks if bad] ++ shapes = text_shapes(mc.hf_text_config) ++ if shapes != K3_SHAPES: ++ reasons.append( ++ f"model shapes {shapes} (hidden, experts, KDA heads, KDA head dim, conv " ++ f"width; Kimi-K3 is {K3_SHAPES})" ++ ) ++ if getattr(mc.hf_text_config, "attn_res_block_size", None) is None: ++ reasons.append("models without AttnRes") ++ return reasons + platform_reasons() ++ ++ ++def text_shapes(text) -> tuple: ++ kda = getattr(text, "linear_attn_config", None) or {} ++ return ( ++ getattr(text, "hidden_size", None), ++ getattr(text, "num_experts", None), ++ kda.get("num_heads"), ++ kda.get("head_dim"), ++ kda.get("short_conv_kernel_size"), ++ ) ++ ++ ++def platform_reasons() -> list[str]: ++ from vllm.platforms import current_platform ++ ++ if not current_platform.is_rocm(): ++ return ["this platform (ROCm only)"] ++ from vllm.platforms.rocm import on_gfx950 ++ ++ if not on_gfx950(): ++ return ["this GPU (gfx950 / MI355X only)"] ++ from vllm.models.kimi_k3.amd.specialized_decode.mono.common.plan import BLOCKS ++ ++ cus = torch.cuda.get_device_properties().multi_processor_count ++ if cus != BLOCKS: ++ # every CTA of a launch must be resident: a partitioned GPU deadlocks ++ return [f"{cus} CUs (a partitioned GPU; the kernels need {BLOCKS})"] ++ return [] ++ ++ ++def piecewise_graphs_only(vllm_config) -> bool: ++ """Decode steps captured as PIECEWISE graphs only: plan() runs them stock.""" ++ from vllm.config import CUDAGraphMode ++ ++ mode = vllm_config.compilation_config.cudagraph_mode ++ if mode is None or mode == CUDAGraphMode.NONE: ++ return False ++ return mode.decode_mode() != CUDAGraphMode.FULL ++ ++ ++def layer_kinds(config) -> dict[int, str]: ++ """Decoder layer -> kind, for every layer a kernel serves: the MoE layers with ++ AttnRes and at least one block before them, minus the MLA block-write layers ++ (none in Kimi-K3).""" ++ block = config.attn_res_block_size ++ if block is None or config.num_experts is None: ++ return {} ++ out = {} ++ for i in range(1, config.num_hidden_layers): ++ if i < config.first_k_dense_replace or i % config.moe_layer_freq: ++ continue ++ if config.is_kda_layer(i): ++ out[i] = WHOLE ++ elif i % block: ++ out[i] = FFN ++ return out +diff --git a/vllm/models/kimi_k3/amd/specialized_decode/runtime/decode.py b/vllm/models/kimi_k3/amd/specialized_decode/runtime/decode.py +new file mode 100644 +index 000000000..fa62155fd +--- /dev/null ++++ b/vllm/models/kimi_k3/amd/specialized_decode/runtime/decode.py +@@ -0,0 +1,198 @@ ++# SPDX-License-Identifier: Apache-2.0 ++# SPDX-FileCopyrightText: Copyright contributors to the vLLM project ++"""The Kimi-K3 specialized decode executor (layer mode): one routing decision per ++step, then the layer kernels inside the stock layer loop. ++ ++Built by the model state once the weights are loaded (``KimiK3ModelState``): checks ++the loaded weights, builds the per-rank launch state (``runner``) and the layer ++kernels of every width, and compiles every build, before memory profiling and graph ++capture. ``prepare_kv`` binds it to every KV cache vLLM allocates. Each step ++``plan`` returns a ``Step`` or None; ``Step.layer(i)`` is the kernel replacing ++decoder layer i (None for vLLM's layers). Weight reloads are not supported.""" ++ ++from collections.abc import Callable, Mapping ++from functools import partial ++from typing import Any, NamedTuple ++ ++import torch ++ ++from vllm.config import CUDAGraphMode ++from vllm.distributed import get_tp_group ++from vllm.forward_context import get_forward_context, is_forward_context_available ++from vllm.logger import init_logger ++from vllm.models.kimi_k3.amd.specialized_decode.mono.abi import FAILED ++from vllm.models.kimi_k3.amd.specialized_decode.runtime import ( ++ WHOLE, ++ DecodeConfig, ++ layer_kinds, ++) ++from vllm.models.kimi_k3.amd.specialized_decode.runtime.guards import ( ++ FailStop, ++ cache_reasons, ++ weight_reasons, ++) ++from vllm.models.kimi_k3.amd.specialized_decode.runtime.layers import ( ++ FfnLayer, ++ WholeLayer, ++ compile_build, ++) ++from vllm.models.kimi_k3.amd.specialized_decode.runtime.peer import tp_agree ++from vllm.models.kimi_k3.amd.specialized_decode.runtime.runner import MonoRunner ++ ++logger = init_logger(__name__) ++ ++Kernel = WholeLayer | FfnLayer ++ ++ ++class Step(NamedTuple): ++ """A planned step: ``width`` rows of whole uniform spec-verify requests.""" ++ ++ width: int ++ attn_metadata: Mapping[str, Any] # by layer name ++ kernels: Mapping[int, Kernel] # by decoder layer, at this width ++ ++ def layer(self, idx: int) -> Callable | None: ++ """The kernel replacing decoder layer idx in this step, or None (vLLM's).""" ++ kernel = self.kernels.get(idx) ++ return None if kernel is None else partial(kernel.launch, self) ++ ++ ++def step_reason(attn_metadata, kda_prefix: str, width: int) -> str | None: ++ """None if a step's KDA metadata (identical on every TP rank) is a pure ++ spec-verify step of whole uniform requests filling ``width`` rows, else why ++ not.""" ++ if not isinstance(attn_metadata, dict): ++ return "no metadata" ++ m: Any = attn_metadata.get(kda_prefix) ++ if m is None or m.spec_sequence_masks is None: ++ return "not spec-verify" ++ if m.num_prefills or m.num_decodes or m.num_spec_decodes < 1: ++ return "not pure decode" ++ q = m.spec_state_indices_tensor.size(-1) ++ if m.uniform_spec_sequence_length != q or m.num_spec_decodes * q != width: ++ return "not uniform" ++ return None ++ ++ ++def agree(reasons: list[str], group, what: str) -> None: ++ """Raises on every TP rank if any rank has ``reasons``: a rank that refused ++ alone would leave the others spinning in-kernel.""" ++ if not tp_agree(not reasons, group): ++ why = "; ".join(reasons) if reasons else "another TP rank refused" ++ raise ValueError( ++ f"Kimi-K3 specialized decode does not support {what}: {why} " ++ "(decode_backend=auto runs vLLM's decode)" ++ ) ++ ++ ++class SpecializedDecode: ++ """The executor: builds the layer kernels, binds them to each KV cache and ++ plans every step.""" ++ ++ def __init__(self, backbone, config: DecodeConfig): ++ tp = get_tp_group() ++ agree(weight_reasons(backbone), tp.cpu_group, "the loaded weights") ++ self.backbone, self.config = backbone, config ++ kinds = { ++ i: kind ++ for i, kind in layer_kinds(backbone.config).items() ++ if backbone.start_layer <= i < backbone.end_layer ++ } ++ runner = self.runner = MonoRunner(next(backbone.parameters()).device) ++ runner.num_spec = config.q - 1 ++ runner.first_layer = min(kinds) ++ whole = [i for i, kind in kinds.items() if kind == WHOLE] ++ runner.kda_prefix = backbone.layers[whole[0]].self_attn.prefix ++ self.kernels: dict[int, dict[int, Kernel]] = { ++ w: { ++ i: (WholeLayer if kind == WHOLE else FfnLayer)( ++ backbone.layers[i], runner, w ++ ) ++ for i, kind in kinds.items() ++ } ++ for w in config.widths ++ } ++ self.kda: list[WholeLayer] = [ ++ k for k in self.kernels[config.widths[0]].values() ++ if isinstance(k, WholeLayer) ++ ] # fmt: skip ++ builds = { ++ b for by in self.kernels.values() for k in by.values() for b in k.builds() ++ } ++ for build in builds: ++ compile_build(build) ++ self.failstop = FailStop( ++ runner.health[FAILED : FAILED + 1], runner.rank, tp.world_size, tp.cpu_group ++ ) ++ # (data_ptr, numel) of the first KDA layer's state prepare_kv bound ++ self._kv: tuple[int, int] | None = None ++ logger.info( ++ "Kimi-K3 specialized decode: %d layers, widths %s, %d builds; other " ++ "steps run vLLM's decode", ++ len(kinds), ++ list(config.widths), ++ len(builds), ++ ) ++ ++ def _kv_key(self) -> tuple[int, int] | None: ++ state = self.kda[0].layer.self_attn.kv_cache ++ if not isinstance(state, (list, tuple)) or len(state) != 2: ++ return None ++ if state[1].numel() == 0: ++ return None ++ return state[1].data_ptr(), state[1].numel() ++ ++ def prepare_kv(self) -> None: ++ """Checks and binds the KDA states vLLM just bound (the CUDA-graph memory ++ profiling cache, then the real one).""" ++ self._kv = None ++ if self._kv_key() is None: ++ return ++ caches = {k.layer.self_attn.prefix: k.caches() for k in self.kda} ++ agree(cache_reasons(caches), get_tp_group().cpu_group, "the KV caches") ++ self._kv = self._kv_key() ++ ++ def plan(self, positions: torch.Tensor, intermediate_tensors) -> Step | None: ++ """This step's plan, or None to run vLLM's layers. ++ ++ Steps that run vLLM's layers: ++ - pipeline-parallel stage inputs, and steps before a KV cache is bound ++ (profiling, dummy runs); ++ - PIECEWISE steps (captured without attention metadata); ++ - widths without kernels (more than MAX_TOKENS rows); ++ - steps that are not pure spec-verify steps of whole uniform requests ++ (``step_reason``). ++ ++ Eager steps run the kernels too. Every predicate is the same on all TP ++ ranks: the kernels' mailboxes need every rank on one path. Under FULL ++ graphs this runs at capture only and the plan is baked in; a replay with ++ fewer requests than its graph pads whole requests (NULL_BLOCK_ID state, ++ one accepted token), which the kernels serve like real ones. ++ """ ++ if intermediate_tensors is not None or self._kv is None: ++ return None ++ if not is_forward_context_available(): ++ return None ++ context = get_forward_context() ++ if context.cudagraph_runtime_mode == CUDAGraphMode.PIECEWISE: ++ return None ++ width = positions.size(-1) ++ kernels = self.kernels.get(width) ++ if kernels is None: ++ return None ++ md = context.attn_metadata ++ if step_reason(md, self.runner.kda_prefix, width) is not None: ++ return None ++ assert isinstance(md, dict) ++ if self._kv_key() != self._kv: ++ raise RuntimeError( ++ "Kimi-K3 specialized decode: the KV cache was re-bound without " ++ "prepare_kv (ModelState.initialize_kv_cache)" ++ ) ++ return Step(width, md, kernels) ++ ++ def after_step(self) -> None: ++ """Once per executed step, from compute_logits (eager on every rank, also ++ under graph replay). Raises on every rank once a rank gave up an in-kernel ++ wait.""" ++ self.failstop.after_step() +diff --git a/vllm/models/kimi_k3/amd/specialized_decode/runtime/guards.py b/vllm/models/kimi_k3/amd/specialized_decode/runtime/guards.py +new file mode 100644 +index 000000000..9d0808d72 +--- /dev/null ++++ b/vllm/models/kimi_k3/amd/specialized_decode/runtime/guards.py +@@ -0,0 +1,154 @@ ++# SPDX-License-Identifier: Apache-2.0 ++# SPDX-FileCopyrightText: Copyright contributors to the vLLM project ++"""The refusals that need the loaded model (``weight_reasons``) or the bound KV ++caches (``cache_reasons``), and the fail-stop watch on the kernels' error words. ++Each refusal reports every reason, not only the first.""" ++ ++from collections.abc import Mapping ++from typing import Any ++ ++import torch ++ ++from vllm.models.kimi_k3.amd.specialized_decode.runtime import WHOLE, layer_kinds ++ ++ ++def weight_reasons(model) -> list[str]: ++ """``model``: the loaded ``KimiLinearModel``; every kernel layer's weights as ++ the kernels read them, and the MoE backend that lays them out.""" ++ from vllm.model_executor.layers.fused_moe.oracle.mxfp4 import Mxfp4MoeBackend ++ from vllm.models.kimi_k3.amd.specialized_decode.mono.attention import back, kda ++ from vllm.models.kimi_k3.amd.specialized_decode.mono.stages import moe ++ ++ reasons: list[str] = [] ++ bf, f32 = torch.bfloat16, torch.float32 ++ ++ def need(name, t, shape, dtype): ++ if t is None: ++ reasons.append(f"{name}: missing") ++ elif tuple(t.shape) != shape or t.dtype != dtype or not t.is_contiguous(): ++ reasons.append( ++ f"{name}: {tuple(t.shape)} {t.dtype}" ++ f"{'' if t.is_contiguous() else ' strided'}," ++ f" kernel reads contiguous {shape} {dtype}" ++ ) ++ ++ kinds = layer_kinds(model.config) ++ if not kinds: ++ return ["no layer has an AttnRes latent MoE"] ++ for i, kind in kinds.items(): ++ if not model.start_layer <= i < model.end_layer: ++ continue ++ layer = model.layers[i] ++ p = f"layer {i}" ++ m = layer.mlp ++ if not getattr(m, "use_latent_moe", False) or m.shared_experts is None: ++ reasons.append(f"{p}: needs a latent MoE with shared experts") ++ continue ++ if m.routed_scaling_factor != 1.0: ++ reasons.append(f"{p}: routed_scaling_factor {m.routed_scaling_factor}") ++ ex = getattr(m.experts, "routed_experts", m.experts) ++ qm = getattr(ex, "quant_method", None) ++ if getattr( ++ qm, "mxfp4_backend", None ++ ) != Mxfp4MoeBackend.AITER_MXFP4_BF16 or not getattr( ++ qm, "is_k3_situ_aiter", False ++ ): ++ reasons.append(f"{p}: MoE backend must be AITER_MXFP4_BF16 (K3 SiTU)") ++ continue ++ H, LAT, E = moe.HIDDEN, moe.LAT, moe.E ++ need(f"{p} gate", m.gate.weight, (E, H), bf) ++ need(f"{p} gate bias", m.gate.e_score_correction_bias, (E,), f32) ++ need(f"{p} latent down", m.routed_expert_down_proj.weight, (LAT, H), bf) ++ need(f"{p} latent norm", m.routed_expert_norm.weight, (LAT,), bf) ++ need(f"{p} shared gate_up", m.shared_experts.gate_up_proj.weight, ++ (2 * moe.SI, H), bf) # fmt: skip ++ need(f"{p} shared down", m.shared_experts.down_proj.weight, (H, moe.SI), bf) ++ sharded = getattr(m.routed_output_transform, "row_sharded", False) ++ need(f"{p} up_proj", m.routed_expert_up_proj.weight, ++ (moe.UP_N if sharded else H, LAT), bf) # fmt: skip ++ for name in ("w13_weight", "w13_weight_scale", "w2_weight", "w2_weight_scale"): ++ if getattr(ex, name, None) is None: ++ reasons.append(f"{p}: experts have no {name}") ++ need(f"{p} mlp_res_norm", layer.mlp_res_norm.weight, (H,), bf) ++ need(f"{p} mlp_res_proj", layer.mlp_res_proj.weight, (1, H), bf) ++ need(f"{p} post_attention_layernorm", layer.post_attention_layernorm.weight, ++ (H,), bf) # fmt: skip ++ need(f"{p} o_proj", layer.self_attn.o_proj.weight, (H, back.OK), bf) ++ if kind != WHOLE: ++ continue ++ a = layer.self_attn ++ if a.local_num_heads != kda.NH or a.head_dim != kda.HD or a.conv_size != 4: ++ reasons.append(f"{p}: KDA heads/head_dim/conv_size") ++ continue ++ need(f"{p} in_proj", a.in_proj_qkvgfab.weight, (kda.NPROJ, H), bf) ++ need(f"{p} f_b_proj", a.f_b_proj.weight, (kda.PROJ, kda.HD), bf) ++ need(f"{p} o_norm", a.o_norm.weight, (kda.HD,), bf) ++ need(f"{p} conv1d", a.conv1d.weight.view(-1, 4), (kda.QKV, 4), f32) ++ need(f"{p} A_log", a.A_log.view(-1), (kda.NH,), f32) ++ need(f"{p} dt_bias", a.dt_bias.view(-1), (kda.PROJ,), f32) ++ need(f"{p} input_layernorm", layer.input_layernorm.weight, (H,), bf) ++ need(f"{p} self_attention_res_norm", layer.self_attention_res_norm.weight, ++ (H,), bf) # fmt: skip ++ need(f"{p} self_attention_res_proj", layer.self_attention_res_proj.weight, ++ (1, H), bf) # fmt: skip ++ return reasons ++ ++ ++def cache_reasons(caches: Mapping[str, Any]) -> list[str]: ++ """Every KDA layer's (conv, recurrent) state as K1 reads them.""" ++ from vllm.models.kimi_k3.amd.specialized_decode.mono.attention import kda ++ ++ reasons = [] ++ for name, (conv_state, rstate) in caches.items(): ++ if conv_state.dtype != torch.bfloat16: ++ reasons.append(f"{name}: conv state {conv_state.dtype}, needs bf16") ++ if rstate.dtype != torch.float32: ++ reasons.append(f"{name}: recurrent state {rstate.dtype}, needs fp32") ++ elif ( ++ rstate[0].numel() != kda.NH * kda.HD * kda.HD ++ or not rstate[0].is_contiguous() ++ ): ++ reasons.append(f"{name}: a slot's recurrent state must be dense") ++ return reasons ++ ++ ++class FailStop: ++ """Stops every TP rank at the same step once any of them gave up an in-kernel ++ wait (its health word FAILED, ``mono.abi``). Each step queues a D2H copy of ++ this rank's word and keeps the newest copy that has landed, without blocking; ++ every CHECK_EVERY steps the ranks agree on those over the TP CPU group and all ++ raise together. A rank that raised alone would leave the others in the next ++ collective (the drafter's) for ever. A failure is seen within about ++ CHECK_EVERY steps.""" ++ ++ CHECK_EVERY = 16 ++ ++ def __init__(self, failed: torch.Tensor, rank: int, world: int, group): ++ self.failed, self.rank, self.world, self.group = failed, rank, world, group ++ self.host = torch.zeros(1, dtype=torch.int32, pin_memory=True) ++ self.copied: torch.cuda.Event | None = None ++ self.seen = False ++ self.steps = 0 ++ ++ def after_step(self) -> None: ++ """Once per step on every TP rank (compute_logits, itself collective).""" ++ if torch.cuda.is_current_stream_capturing(): ++ return ++ if self.copied is None or self.copied.query(): ++ if self.copied is not None: ++ self.seen |= bool(self.host[0]) ++ self.host.copy_(self.failed, non_blocking=True) ++ self.copied = torch.cuda.Event() ++ self.copied.record() ++ self.steps += 1 ++ if self.steps % self.CHECK_EVERY: ++ return ++ votes = torch.zeros(self.world, dtype=torch.int32) ++ votes[self.rank] = int(self.seen) ++ torch.distributed.all_reduce(votes, group=self.group) ++ if votes.any(): ++ ranks = [r for r, v in enumerate(votes.tolist()) if v] ++ raise RuntimeError( ++ f"Kimi-K3 specialized decode: TP ranks {ranks} gave up an in-kernel " ++ "wait; stopping the engine (decode_backend=auto runs vLLM's decode)" ++ ) +diff --git a/vllm/models/kimi_k3/amd/specialized_decode/runtime/layers.py b/vllm/models/kimi_k3/amd/specialized_decode/runtime/layers.py +new file mode 100644 +index 000000000..843e66dc5 +--- /dev/null ++++ b/vllm/models/kimi_k3/amd/specialized_decode/runtime/layers.py +@@ -0,0 +1,221 @@ ++# SPDX-License-Identifier: Apache-2.0 ++# SPDX-FileCopyrightText: Copyright contributors to the vLLM project ++"""The layer kernels: each takes and returns what ``KimiDecoderLayer.forward`` ++does with AttnRes, ``(prefix_sum, block_residual, prefix_delta)``. ++ ++- ``WholeLayer`` (KDA layers): K1 (AttnRes, in_proj, f_b, spec conv, the KDA ++ recurrence, the gated norm) writes o_proj's input; K2 runs o_proj, its ++ all-reduce, the MLP AttnRes and the latent MoE. ++- ``FfnLayer`` (MLA layers): vLLM's AttnRes and MLA up to o_proj's input, then ++ K2. ++ ++Weights are views of the loaded parameters, taken at every launch. ++""" ++ ++from typing import TYPE_CHECKING, Any ++ ++import torch ++ ++from vllm.model_executor.layers.mamba.mamba_utils import is_conv_state_dim_first ++from vllm.models.kimi_k3.amd.specialized_decode.runtime.runner import MonoRunner ++ ++if TYPE_CHECKING: ++ from vllm.models.kimi_k3.amd.specialized_decode.runtime.decode import Step ++ ++ ++def _k2_key(layer, runner: MonoRunner, tokens: int, reset: bool): ++ from vllm.models.kimi_k3.amd.specialized_decode.mono import layer as k2 ++ ++ m = layer.mlp ++ return k2.K2Build( ++ tokens=tokens, ++ nblocks=layer.prev_valid_blocks + (1 if reset else 0), ++ reset=reset, ++ eps=layer.mlp_res_norm.variance_epsilon, ++ out_eps=layer.post_attention_layernorm.variance_epsilon, ++ ln_eps=m.routed_expert_norm.variance_epsilon, ++ **runner.fault, ++ ) ++ ++ ++def _k2(layer, runner: MonoRunner, core, prefix, blocks, reset=False): ++ """K2 from ``core`` (o_proj's input): ``prefix`` updated in place (written, ++ when ``reset``), the MoE output returned.""" ++ from vllm.models.kimi_k3.amd.specialized_decode.mono import layer as k2 ++ from vllm.models.kimi_k3.amd.specialized_decode.mono.attention.kda import HIDDEN ++ from vllm.models.kimi_k3.amd.specialized_decode.mono.stages import moe ++ ++ m = layer.mlp ++ ex = getattr(m.experts, "routed_experts", m.experts) ++ w_up = m.routed_expert_up_proj.weight ++ if not getattr(m.routed_output_transform, "row_sharded", False): ++ w_up = w_up.narrow(0, runner.rank * moe.UP_N, moe.UP_N) ++ out = torch.empty_like(prefix) ++ k2.k2_launch( ++ _k2_key(layer, runner, core.size(0), reset), ++ core=core, ++ w_o=layer.self_attn.o_proj.weight, ++ prefix=prefix, ++ blocks=blocks, ++ ares_nw=layer.mlp_res_norm.weight, ++ ares_qk=layer.mlp_res_proj.weight.view(HIDDEN), ++ in_nw=layer.post_attention_layernorm.weight, ++ w_gate=m.gate.weight, ++ bias=m.gate.e_score_correction_bias, ++ w_ld=m.routed_expert_down_proj.weight, ++ w_sgu=m.shared_experts.gate_up_proj.weight, ++ w_sd=m.shared_experts.down_proj.weight, ++ w13=ex.w13_weight, ++ w13s=ex.w13_weight_scale, ++ w2=ex.w2_weight, ++ w2s=ex.w2_weight_scale, ++ ln_w=m.routed_expert_norm.weight, ++ w_up=w_up, ++ out=out, ++ layer=128 + layer.layer_idx, ++ **runner.launch_args(), ++ **runner.peer_args(), ++ ) ++ return out ++ ++ ++def compile_build(key) -> None: ++ """Compiles one launch's kernel (a ``K2Build`` or ``KdaPreBuild``) without ++ running it: FlyDSL's compile-only mode returns before the launch.""" ++ import os ++ ++ from vllm.models.kimi_k3.amd.specialized_decode.mono import layer as k2 ++ from vllm.models.kimi_k3.amd.specialized_decode.mono.attention import kda ++ ++ mod = kda if isinstance(key, kda.KdaPreBuild) else k2 ++ launch = mod.build(key) ++ # FlyDSL's compile-only switch has no FLYDSL_ prefix ++ old = os.environ.get("COMPILE_ONLY") ++ os.environ["COMPILE_ONLY"] = "1" ++ try: ++ launch(*[0] * len(mod.ABI.names), stream=torch.cuda.current_stream()) ++ finally: ++ if old is None: ++ del os.environ["COMPILE_ONLY"] ++ else: ++ os.environ["COMPILE_ONLY"] = old ++ ++ ++class _Layer: ++ def __init__(self, layer, runner: MonoRunner, width: int): ++ self.layer = layer ++ self.runner = runner ++ self.width = width ++ # the step's first launch bumps the epoch (a graph captures it) ++ self.first = layer.layer_idx == runner.first_layer ++ ++ def builds(self) -> list: ++ """The builds this layer launches at its width, compiled before capture.""" ++ reset = getattr(self.layer, "is_block_write_layer", False) ++ return [_k2_key(self.layer, self.runner, self.width, reset)] ++ ++ ++class WholeLayer(_Layer): ++ def _k1_key(self, qlen: int, state_len: int, delta: bool): ++ from vllm.models.kimi_k3.amd.specialized_decode.mono.attention import kda ++ ++ layer = self.layer ++ attn = layer.self_attn ++ return kda.KdaPreBuild( ++ tokens=self.width, ++ qlen=qlen, ++ nblocks=layer.prev_valid_blocks, ++ delta=delta, ++ state_len=state_len, ++ eps=layer.self_attention_res_norm.variance_epsilon, ++ out_eps=layer.input_layernorm.variance_epsilon, ++ onorm_eps=attn.o_norm.eps, ++ lower_bound=attn.gate_lower_bound, ++ write_idx=layer.block_write_idx if layer.is_block_write_layer else -1, ++ ) ++ ++ def builds(self) -> list: ++ from vllm.models.kimi_k3.amd.specialized_decode.mono.attention import kda ++ ++ q = self.runner.num_spec + 1 ++ # the conv state keeps CONV_W - 1 + num_spec entries; a layer after ++ # another always has a pending delta ++ k1 = self._k1_key(q, kda.CONV_W - 1 + self.runner.num_spec, True) ++ return [k1, *super().builds()] ++ ++ def caches(self) -> tuple[torch.Tensor, torch.Tensor]: ++ """The conv state as [slots, dim, state_len] (a view) and the recurrent ++ state.""" ++ conv_state, rstate = self.layer.self_attn.kv_cache ++ if not is_conv_state_dim_first(): ++ conv_state = conv_state.transpose(-1, -2) ++ return conv_state, rstate ++ ++ def launch(self, step: "Step", positions, prefix_sum, block_residual, prefix_delta): ++ from vllm.models.kimi_k3.amd.specialized_decode.mono.attention import kda ++ ++ layer, runner = self.layer, self.runner ++ attn = layer.self_attn ++ m: Any = step.attn_metadata[attn.prefix] ++ if self.first: ++ runner.step_begin() ++ conv_state, rstate = self.caches() ++ s = prefix_sum.size(0) ++ core = prefix_sum.new_empty(s, kda.PROJ) ++ kda.kda_pre( ++ self._k1_key( ++ m.spec_state_indices_tensor.size(-1), ++ conv_state.size(-1), ++ prefix_delta is not None, ++ ), ++ prefix=prefix_sum, ++ delta=prefix_delta, ++ blocks=block_residual, ++ ares_nw=layer.self_attention_res_norm.weight, ++ ares_qk=layer.self_attention_res_proj.weight.view(kda.HIDDEN), ++ in_nw=layer.input_layernorm.weight, ++ w_in=attn.in_proj_qkvgfab.weight, ++ w_fb=attn.f_b_proj.weight, ++ conv_w=attn.conv1d.weight.view(kda.QKV, 4), ++ conv_state=conv_state, ++ a_log=attn.A_log, ++ dt_bias=attn.dt_bias, ++ on_w=attn.o_norm.weight, ++ rstate=rstate, ++ st_idx=m.spec_state_indices_tensor, ++ num_acc=m.num_accepted_tokens, ++ core_out=core, ++ layer=layer.layer_idx, ++ **runner.launch_args(), ++ ) ++ if layer.is_block_write_layer: ++ # the prefix restarts at the attention's output ++ prefix_sum = torch.empty_like(prefix_sum) ++ out = _k2(layer, runner, core, prefix_sum, block_residual, reset=True) ++ else: ++ out = _k2(layer, runner, core, prefix_sum, block_residual) ++ return prefix_sum, block_residual, out ++ ++ ++class FfnLayer(_Layer): ++ def launch(self, step: "Step", positions, prefix_sum, block_residual, prefix_delta): ++ from vllm.models.kimi_k3.amd.linear import _apply_attn_res ++ ++ layer = self.layer ++ if self.first: ++ self.runner.step_begin() ++ hidden_states = _apply_attn_res( ++ prefix_sum, ++ block_residual, ++ layer.self_attention_res_proj, ++ layer.self_attention_res_norm, ++ layer.prev_valid_blocks, ++ delta=prefix_delta, ++ output_norm=layer.input_layernorm, ++ quant_key=layer.self_attn.get_input_quant_key(), ++ ) ++ core = layer.self_attn.mla_attn.forward( ++ positions, hidden_states, skip_o_proj=True ++ ) ++ out = _k2(layer, self.runner, core, prefix_sum, block_residual) ++ return prefix_sum, block_residual, out +diff --git a/vllm/models/kimi_k3/amd/specialized_decode/runtime/peer.py b/vllm/models/kimi_k3/amd/specialized_decode/runtime/peer.py +new file mode 100644 +index 000000000..8e07955eb +--- /dev/null ++++ b/vllm/models/kimi_k3/amd/specialized_decode/runtime/peer.py +@@ -0,0 +1,112 @@ ++# SPDX-License-Identifier: Apache-2.0 ++# SPDX-FileCopyrightText: Copyright contributors to the vLLM project ++# Adapted from ROCm/ATOM (https://github.com/ROCm/ATOM) at a526f0d, MIT License, ++# Copyright (C) 2026, Advanced Micro Devices, Inc.: ++# atom/mono/runtime/peer_memory.py ++# atom/mono/runtime/consensus.py ++"""Peer memory for in-kernel TP collectives: one uncached buffer a rank, every ++rank holding every peer's address, and the TP agreements that must precede ++anything collective.""" ++ ++import torch ++import torch.distributed as dist ++ ++ ++class _DeviceBytes: ++ """``__cuda_array_interface__`` of ``nbytes`` bytes at a raw device address, for a ++ non-owning torch view.""" ++ ++ def __init__(self, ptr: int, nbytes: int): ++ self.__cuda_array_interface__ = { ++ "shape": (nbytes,), ++ "typestr": "|u1", ++ "data": (ptr, False), ++ "version": 3, ++ } ++ ++ ++class PeerBuffer: ++ """One uncached buffer per rank, every rank holding every peer's address ++ (``addresses``, int64, indexed by TP rank). What lives where, and why no rank ++ overwrites a peer's unread data, is the model's layout's to say. ++ ++ ``bytes`` is the whole buffer: the kernels tag every pair with the step's ++ epoch, so it is never zeroed between steps. ++ ++ The constructor is collective over ``group`` (the handle exchange, then an ++ agreement that every rank mapped every peer): every rank must reach it, so a ++ runner reaches it only after ``tp_agree``. A failed mapping on any rank raises ++ ``RuntimeError`` on all of them. ++ """ ++ ++ def __init__( ++ self, ++ nbytes: int, ++ group, ++ rank: int, ++ npes: int, ++ device: torch.device, ++ ): ++ from aiter.ops.flydsl.quick_allreduce_int4_ipc import UncachedIpcHeap ++ ++ self._heap = UncachedIpcHeap ++ self.local = UncachedIpcHeap.alloc_uncached(nbytes) ++ # a view, not an owner: close() frees the memory under it ++ self.bytes = torch.as_tensor(_DeviceBytes(self.local, nbytes), device=device) ++ self.rank = rank ++ self._opened: list[int] = [] ++ addresses = [self.local] ++ if npes > 1: ++ handles = UncachedIpcHeap.gather_object_list_via_broadcast( ++ group, UncachedIpcHeap.get_mem_handle_bytes(self.local) ++ ) ++ addresses, failure = [], None ++ try: ++ for peer, handle in enumerate(handles): ++ if peer == rank: ++ addresses.append(self.local) ++ else: ++ base = UncachedIpcHeap.open_mem_handle(handle) ++ self._opened.append(base) ++ addresses.append(base) ++ except RuntimeError as err: # every rank must still reach the agreement ++ failure = err ++ if not tp_agree(failure is None, group): ++ self.close() ++ raise RuntimeError( ++ "peer memory handshake failed on " ++ f"{'this' if failure else 'another'}" ++ " TP rank" ++ ) from failure ++ self.addresses = torch.tensor(addresses, dtype=torch.int64, device=device) ++ ++ def close(self) -> None: ++ """Close the peer mappings and free the local buffer; idempotent. Only once no ++ kernel of any rank can still touch it (the runner is being dropped).""" ++ opened, self._opened = self._opened, [] ++ for base in opened: ++ self._heap.close_mem_handle(base) ++ if self.local: ++ self.bytes = None ++ self._heap.free_device_mem(self.local) ++ self.local = 0 ++ ++ ++# Decisions every TP rank must take alike. ++# ++# A mono kernel waits in-kernel on every peer rank, so a rank that turned mono off ++# alone -- or failed to build what the others launch -- leaves the others spinning ++# on the GPU for ever. Each such decision is a local verdict agreed on over the TP ++# group's CPU group before anything collective (peer handshake, launch) follows. ++# ++ ++ ++def tp_agree(ok: bool, group) -> bool: ++ """True when every rank of ``group`` passed ``ok`` true. Every rank must call ++ it at the same point; ``group`` is a CPU (gloo) process group, or None for a ++ single rank.""" ++ if group is None: ++ return ok ++ flag = torch.tensor([int(ok)], dtype=torch.int32) ++ dist.all_reduce(flag, op=dist.ReduceOp.MIN, group=group) ++ return bool(flag.item()) +diff --git a/vllm/models/kimi_k3/amd/specialized_decode/runtime/runner.py b/vllm/models/kimi_k3/amd/specialized_decode/runtime/runner.py +new file mode 100644 +index 000000000..85119c7bd +--- /dev/null ++++ b/vllm/models/kimi_k3/amd/specialized_decode/runtime/runner.py +@@ -0,0 +1,76 @@ ++# SPDX-License-Identifier: Apache-2.0 ++# SPDX-FileCopyrightText: Copyright contributors to the vLLM project ++"""The specialized decode's per-rank state: one scratch buffer every launch shares ++(their mailbox tags differ), the device step epoch, the health words and the TP ++group's peer memory. ++ ++Mailbox tags carry the epoch the model bumps once a forward (``step_begin``, ++which a CUDA graph captures), so no buffer is zeroed between steps. ++ ++The peer buffer's first ``PEER_HEAD`` bytes hold the output rank's flags, a ++word a rank (``alarm``): a rank whose wait gives up sets its word there before it ++stores anything else. The kernels see the buffer past the head. ++""" ++ ++import torch ++ ++from vllm.distributed import get_tp_group ++from vllm.models.kimi_k3.amd.specialized_decode.mono.abi import PEER_HEAD ++from vllm.models.kimi_k3.amd.specialized_decode.mono.common.plan import HEALTH_WORDS ++from vllm.models.kimi_k3.amd.specialized_decode.runtime import MAX_TOKENS ++from vllm.models.kimi_k3.amd.specialized_decode.runtime.peer import PeerBuffer ++ ++ ++class MonoRunner: ++ """Collective: every TP rank builds it at the same point.""" ++ ++ def __init__(self, device: torch.device): ++ from vllm.models.kimi_k3.amd.specialized_decode.mono import layer ++ from vllm.models.kimi_k3.amd.specialized_decode.mono.attention import kda ++ ++ self.scratch = torch.zeros( ++ max( ++ kda.scratch_bytes( ++ kda.KdaPreBuild( ++ tokens=MAX_TOKENS, ++ qlen=MAX_TOKENS, ++ nblocks=8, ++ delta=True, ++ state_len=3 + MAX_TOKENS, ++ ) ++ ), ++ layer.scratch_bytes(layer.K2Build(tokens=MAX_TOKENS, nblocks=8)), ++ ), ++ dtype=torch.uint8, ++ device=device, ++ ) ++ self.epoch = torch.zeros(1, dtype=torch.int32, device=device) ++ self.health = torch.zeros(HEALTH_WORDS, dtype=torch.int32, device=device) ++ ++ tp = get_tp_group() ++ self.rank = tp.rank_in_group ++ nbytes = PEER_HEAD + layer.peer_bytes(MAX_TOKENS) ++ self.peers = PeerBuffer(nbytes, tp.cpu_group, self.rank, tp.world_size, device) ++ self.peers.bytes.zero_() ++ self.peer_table = self.peers.addresses + PEER_HEAD ++ self.alarm = int(self.peers.addresses[0].item()) + 4 * self.rank ++ self.flags = self.peers.bytes[: 4 * tp.world_size].view(torch.int32) ++ self.fault: dict = {} # K2Build's fault_* fields (tests only) ++ self.first_layer = -1 # the step's first kernel layer: bumps the epoch ++ self.kda_prefix = "" # a KDA layer's name: the step's metadata ++ self.num_spec = 0 # speculative tokens a request ++ ++ def step_begin(self) -> None: ++ """Once a forward, before its first mono launch (a graph captures it).""" ++ self.epoch.add_(1) ++ ++ def launch_args(self) -> dict: ++ return { ++ "scratch": self.scratch, ++ "epoch": self.epoch, ++ "health": self.health, ++ "alarm": self.alarm, ++ } ++ ++ def peer_args(self) -> dict: ++ return {"peers": self.peer_table, "rank": self.rank} diff --git a/inferencex-e2e/benchmarks/multi_node/srt-slurm-recipes/configs/patches/kimik3-nightly/02-vllm-58743-state.patch b/inferencex-e2e/benchmarks/multi_node/srt-slurm-recipes/configs/patches/kimik3-nightly/02-vllm-58743-state.patch new file mode 100644 index 0000000000..1e94c67316 --- /dev/null +++ b/inferencex-e2e/benchmarks/multi_node/srt-slurm-recipes/configs/patches/kimik3-nightly/02-vllm-58743-state.patch @@ -0,0 +1,155 @@ +From a884f79a30c723c7fbbe5a83f099e8c7ab97d1f9 Mon Sep 17 00:00:00 2001 +From: Yichao Zhu +Date: Sat, 26 Sep 2026 07:35:45 +0800 +Subject: [PATCH] [ROCm][Kimi-K3] Support BF16 KDA recurrent state + +Signed-off-by: Yichao Zhu +(cherry picked from commit e09e870c38649ac9db758133054483611349ab53) +--- + tests/models/kimi_k3/test_amd_kda_chunk.py | 23 ++++++++++++++++ + vllm/models/kimi_k3/amd/kda.py | 4 ++- + vllm/models/kimi_k3/amd/linear.py | 4 ++- + vllm/models/kimi_k3/amd/ops/kda_prefill.py | 31 +++++++++++++++------- + 4 files changed, 51 insertions(+), 11 deletions(-) + +diff --git a/tests/models/kimi_k3/test_amd_kda_chunk.py b/tests/models/kimi_k3/test_amd_kda_chunk.py +index 2dcf886c9..0d0944032 100644 +--- a/tests/models/kimi_k3/test_amd_kda_chunk.py ++++ b/tests/models/kimi_k3/test_amd_kda_chunk.py +@@ -701,6 +701,29 @@ def test_paged_state_matches_the_dense_pair( + ) + + ++@pytest.mark.parametrize("use_fused", [False, True]) ++def test_paged_bf16_state_matches_dense_pair(use_fused: bool) -> None: ++ """The FP32 walk may load and store a BF16 persistent cache.""" ++ if use_fused: ++ _requires_kernel() ++ seqlens = [513, 64, 1] ++ inp = _inputs(seqlens, seed=17) ++ rows = [2, 7, 11] ++ cache, idx, warm = _paged_cache(inp, [True] * len(seqlens), rows, slots=16) ++ cache = cache.to(torch.bfloat16) ++ ++ dense = dict(inp) ++ dense["h0"] = cache[idx.long()].float() ++ o_ref, ht_ref = _run(dense, use_fused=use_fused) ++ o_got, ht_got = _run_paged(inp, cache, idx, warm, use_fused=use_fused) ++ ++ assert ht_got is None ++ torch.testing.assert_close(o_got, o_ref, rtol=5e-3, atol=5e-3) ++ torch.testing.assert_close( ++ cache[idx.long()], ht_ref.to(torch.bfloat16), rtol=5e-3, atol=5e-3 ++ ) ++ ++ + @pytest.mark.parametrize("use_fused", [False, True]) + def test_cold_rows_start_from_zero_not_from_cache_junk(use_fused: bool) -> None: + """`has_initial_state=False` must ignore whatever the row already holds.""" +diff --git a/vllm/models/kimi_k3/amd/kda.py b/vllm/models/kimi_k3/amd/kda.py +index 73e357e44..4216c6f94 100644 +--- a/vllm/models/kimi_k3/amd/kda.py ++++ b/vllm/models/kimi_k3/amd/kda.py +@@ -78,7 +78,9 @@ class KimiK3DeltaAttention(GatedDeltaNetAttention): + if self.model_config is None or self.cache_config is None: + raise ValueError("model_config and cache_config must be set") + return MambaStateDtypeCalculator.kda_state_dtype( +- self.model_config.dtype, self.cache_config.mamba_cache_dtype ++ self.model_config.dtype, ++ self.cache_config.mamba_cache_dtype, ++ self.cache_config.mamba_ssm_cache_dtype, + ) + + def get_state_shape(self) -> tuple[tuple[int, ...], tuple[int, ...]]: +diff --git a/vllm/models/kimi_k3/amd/linear.py b/vllm/models/kimi_k3/amd/linear.py +index 029636146..d080c34e7 100644 +--- a/vllm/models/kimi_k3/amd/linear.py ++++ b/vllm/models/kimi_k3/amd/linear.py +@@ -1125,7 +1125,9 @@ class KimiLinearForCausalLM( + vllm_config: "VllmConfig", + ) -> tuple[torch.dtype, torch.dtype]: + return MambaStateDtypeCalculator.kda_state_dtype( +- vllm_config.model_config.dtype, vllm_config.cache_config.mamba_cache_dtype ++ vllm_config.model_config.dtype, ++ vllm_config.cache_config.mamba_cache_dtype, ++ vllm_config.cache_config.mamba_ssm_cache_dtype, + ) + + @classmethod +diff --git a/vllm/models/kimi_k3/amd/ops/kda_prefill.py b/vllm/models/kimi_k3/amd/ops/kda_prefill.py +index 040f3afe6..18d90b811 100644 +--- a/vllm/models/kimi_k3/amd/ops/kda_prefill.py ++++ b/vllm/models/kimi_k3/amd/ops/kda_prefill.py +@@ -86,9 +86,10 @@ def chunk_kda_prefill( + checkpoint_offsets: per-sequence token offset to snapshot at, ``0`` + for none. + checkpoint_state_indices: optional per-sequence destination row. +- state_cache: the paged recurrent state. When given, the fused backend +- reads and writes it in place and neither a gather nor a scatter is +- needed around this call; the returned final state is ``None``. ++ state_cache: the paged recurrent state. The fused backend addresses an ++ FP32 cache directly. Other cache dtypes are gathered into the FP32 ++ accumulator format and cast back after the walk. The returned final ++ state is ``None`` in either case. + state_indices: per-sequence cache row. + has_initial_state: per-sequence flag; false starts from a zero state. + +@@ -150,11 +151,11 @@ def chunk_kda_prefill( + has_initial_state=has_initial_state, + ) + +- # Only the fused walk addresses the paged rows directly. The Triton path +- # gathers the rows it needs and scatters the results back here, so callers +- # pass the cache the same way for either backend. ++ # Only the fused walk addresses FP32 paged rows directly. The Triton path ++ # and lower-precision persistent caches gather the active rows and scatter ++ # the final state back here, so callers use one cache API for every backend. + scatter_to: torch.Tensor | None = None +- if state_cache is not None and not fused: ++ if state_cache is not None and (not fused or state_cache.dtype != torch.float32): + initial_state = gather_initial_states( + state_cache, state_indices, has_initial_state + ) +@@ -165,6 +166,11 @@ def chunk_kda_prefill( + # Restated for the type checker; `fused` already implies all three. + assert cu_seqlens is not None and lower_bound is not None + assert g_bias is not None ++ # The fused HIP prefill kernel accumulates recurrent state in FP32. ++ # Convert only active rows; the persistent cache keeps its configured ++ # dtype and receives a cast-back below. ++ if initial_state is not None and initial_state.dtype != torch.float32: ++ initial_state = initial_state.to(torch.float32) + logger.info_once( + "Kimi-K3 KDA prefill: dispatching the fused ROCm chunk kernel." + ) +@@ -181,7 +187,7 @@ def chunk_kda_prefill( + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + ) +- return fused_kda_chunk( ++ o, final_state = fused_kda_chunk( + qg=ws["qg"], + w=ws["w"], + u=ws["u"], +@@ -208,6 +214,12 @@ def chunk_kda_prefill( + state_indices=state_indices, + has_initial_state=has_initial_state, + ) ++ if scatter_to is not None: ++ assert state_indices is not None ++ assert final_state is not None ++ scatter_to[state_indices.long()] = final_state.to(scatter_to.dtype) ++ return o, None ++ return o, final_state + + o, final_state = chunk_kda_with_fused_gate( + q=q, +@@ -231,6 +243,7 @@ def chunk_kda_prefill( + o = out + if scatter_to is not None: + assert state_indices is not None +- scatter_to[state_indices.long()] = final_state ++ assert final_state is not None ++ scatter_to[state_indices.long()] = final_state.to(scatter_to.dtype) + return o, None + return o, final_state diff --git a/inferencex-e2e/benchmarks/multi_node/srt-slurm-recipes/configs/patches/kimik3-nightly/03-vllm-58743-checkpoints.patch b/inferencex-e2e/benchmarks/multi_node/srt-slurm-recipes/configs/patches/kimik3-nightly/03-vllm-58743-checkpoints.patch new file mode 100644 index 0000000000..1e1f5c075e --- /dev/null +++ b/inferencex-e2e/benchmarks/multi_node/srt-slurm-recipes/configs/patches/kimik3-nightly/03-vllm-58743-checkpoints.patch @@ -0,0 +1,251 @@ +From c8a4c72073306a2d2b7370e45bd91be11c0df796 Mon Sep 17 00:00:00 2001 +From: Yichao Zhu +Date: Sat, 17 Oct 2026 12:35:33 +0800 +Subject: [PATCH] fix(kimik3): support BF16 prefill checkpoint stores on ROCm + +Preserve FP32 accumulation and write valid checkpoint rows in their persistent cache dtype. Extend the existing checkpoint regressions for BF16 destinations and padded rows. + +Signed-off-by: Yichao Zhu +(cherry picked from commit 53bfc127fdc8e87d81e1a50dbb773f197e9df875) +--- + csrc/kimi_k3/fused_kda_chunk_kernel_rocm.cu | 31 +++++++++---- + tests/models/kimi_k3/test_amd_kda_chunk.py | 48 ++++++++++++++------- + 2 files changed, 55 insertions(+), 24 deletions(-) + +diff --git a/csrc/kimi_k3/fused_kda_chunk_kernel_rocm.cu b/csrc/kimi_k3/fused_kda_chunk_kernel_rocm.cu +index bb9485dd7..3ed350c85 100644 +--- a/csrc/kimi_k3/fused_kda_chunk_kernel_rocm.cu ++++ b/csrc/kimi_k3/fused_kda_chunk_kernel_rocm.cu +@@ -48,6 +48,7 @@ + + using bf16_t = __bf16; + using bf16x2 = __bf16 __attribute__((ext_vector_type(2))); ++using bf16x4 = __bf16 __attribute__((ext_vector_type(4))); + using bf16x8 = __bf16 __attribute__((ext_vector_type(8))); + using f32x4 = float __attribute__((ext_vector_type(4))); + using u32x2 = uint32_t __attribute__((ext_vector_type(2))); +@@ -151,7 +152,8 @@ struct Params { + // replaying the tail as a second model forward. `ckidx`, when present, + // redirects sequence n's snapshot to row `ckidx[n]`, which lets the store go + // straight into the paged state cache; a negative row disables it. +- float* ckpt; ++ void* ckpt; ++ bool ckpt_bf16; + const int* ckoff; + const int* ckidx; + int G; +@@ -842,11 +844,18 @@ __global__ __launch_bounds__(64 * WAVES * KS, + // Same store as the final state below, at an interior boundary instead of + // the last one. + auto store_checkpoint = [&]() { +- float* dst = p.ckpt + ck_dst + (v0 + lrow) * kK + 4 * lgrp; ++ const int64_t base = ck_dst + (v0 + lrow) * kK + 4 * lgrp; + for (int vt = 0; vt < kNV; ++vt) +- for (int kt = 0; kt < kNKTL; ++kt) +- *reinterpret_cast(dst + vt * 16 * kK + (kt0 + kt) * 16) = +- S[kt][vt]; ++ for (int kt = 0; kt < kNKTL; ++kt) { ++ const int64_t offset = base + vt * 16 * kK + (kt0 + kt) * 16; ++ if (p.ckpt_bf16) { ++ *reinterpret_cast(static_cast(p.ckpt) + offset) = ++ __builtin_convertvector(S[kt][vt], bf16x4); ++ } else { ++ *reinterpret_cast(static_cast(p.ckpt) + offset) = ++ S[kt][vt]; ++ } ++ } + }; + + const int tail = last_group ? T - (nt_all - 1) * kBT : kBT; +@@ -2112,7 +2121,7 @@ void fused_kda_chunk( + // Prefill checkpoint export. `checkpoint_state` is either a [N, H, V, K] + // staging buffer, as FlashKDA takes, or the paged state cache itself when + // `checkpoint_state_indices` says which row each sequence owns. +- float* ckpt_ptr = nullptr; ++ void* ckpt_ptr = nullptr; + int const* ckoff_ptr = nullptr; + int const* ckidx_ptr = nullptr; + if (checkpoint_state.has_value() || checkpoint_offsets.has_value()) { +@@ -2121,7 +2130,8 @@ void fused_kda_chunk( + "checkpoint_state and checkpoint_offsets must be given together"); + STD_TORCH_CHECK( + checkpoint_state->is_cuda() && +- checkpoint_state->scalar_type() == ScalarType::Float && ++ (checkpoint_state->scalar_type() == ScalarType::Float || ++ checkpoint_state->scalar_type() == ScalarType::BFloat16) && + checkpoint_state->dim() == 4 && + checkpoint_state->size(1) == num_heads && + checkpoint_state->size(2) == kV && +@@ -2129,7 +2139,8 @@ void fused_kda_chunk( + checkpoint_state->stride(3) == 1 && + checkpoint_state->stride(2) == kK && + checkpoint_state->stride(1) == kV * kK, +- "checkpoint_state must be an fp32 [rows, H, 128, 128] tensor whose " ++ "checkpoint_state must be an fp32 or bf16 [rows, H, 128, 128] tensor " ++ "whose " + "rows are dense; stride(0) may be padded, as the paged state cache " + "it usually is"); + STD_TORCH_CHECK(checkpoint_offsets->is_cuda() && +@@ -2154,7 +2165,7 @@ void fused_kda_chunk( + checkpoint_state->size(0), " rows for ", num_seqs, + " sequences"); + } +- ckpt_ptr = static_cast(checkpoint_state->data_ptr()); ++ ckpt_ptr = checkpoint_state->data_ptr(); + ckoff_ptr = static_cast(checkpoint_offsets->data_ptr()); + } + +@@ -2238,6 +2249,8 @@ void fused_kda_chunk( + params.ckrow = checkpoint_state.has_value() ? checkpoint_state->stride(0) + : num_heads * kV * kK; + params.ckpt = ckpt_ptr; ++ params.ckpt_bf16 = checkpoint_state.has_value() && ++ checkpoint_state->scalar_type() == ScalarType::BFloat16; + params.ckoff = ckoff_ptr; + params.ckidx = ckidx_ptr; + +diff --git a/tests/models/kimi_k3/test_amd_kda_chunk.py b/tests/models/kimi_k3/test_amd_kda_chunk.py +index 0d0944032..e3e000137 100644 +--- a/tests/models/kimi_k3/test_amd_kda_chunk.py ++++ b/tests/models/kimi_k3/test_amd_kda_chunk.py +@@ -447,13 +447,13 @@ def test_accepts_chunk_metadata_from_the_attention_builder( + torch.testing.assert_close(state_got, state_ref, rtol=1e-3, atol=1e-3) + + +-def _checkpoint_buffer(rows: int) -> torch.Tensor: ++def _checkpoint_buffer(rows: int, dtype: torch.dtype = torch.float32) -> torch.Tensor: + """A destination filled with NaN, so an unwritten row is unmistakable.""" + return torch.full( + (rows, NUM_HEADS, HEAD_DIM, HEAD_DIM), + float("nan"), + device="cuda", +- dtype=torch.float32, ++ dtype=dtype, + ) + + +@@ -470,9 +470,13 @@ def _prefix_inputs(inp: dict, seq: int, num_tokens: int) -> dict: + + + def _run_with_checkpoints( +- inp: dict, offsets: list[int], rows: int | None = None, **kw ++ inp: dict, ++ offsets: list[int], ++ rows: int | None = None, ++ dtype: torch.dtype = torch.float32, ++ **kw, + ) -> tuple[torch.Tensor, tuple]: +- ckpt = _checkpoint_buffer(len(offsets) if rows is None else rows) ++ ckpt = _checkpoint_buffer(len(offsets) if rows is None else rows, dtype) + result = _run( + inp, + use_fused=True, +@@ -534,7 +538,10 @@ def test_checkpoint_leaves_the_output_and_final_state_alone() -> None: + + + @pytest.mark.parametrize("offset", [100, 63, 65, 1024]) +-def test_unusable_checkpoint_offsets_write_nothing(offset: int) -> None: ++@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) ++def test_unusable_checkpoint_offsets_write_nothing( ++ offset: int, dtype: torch.dtype ++) -> None: + """Offsets the walk never reaches must leave the destination untouched. + + The recurrent state only exists at chunk boundaries, so an offset that is +@@ -550,12 +557,13 @@ def test_unusable_checkpoint_offsets_write_nothing(offset: int) -> None: + + assert offset % KDA_CHECKPOINT_ALIGNMENT != 0 or offset > 512 + inp = _inputs([512], seed=31) +- ckpt, _ = _run_with_checkpoints(inp, [offset]) ++ ckpt, _ = _run_with_checkpoints(inp, [offset], dtype=dtype) + + assert ckpt.isnan().all() + + +-def test_checkpoint_state_indices_pick_the_destination_row() -> None: ++@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) ++def test_checkpoint_state_indices_pick_the_destination_row(dtype: torch.dtype) -> None: + """With indices the export lands straight in the paged state cache. + + FlashKDA writes one row per sequence into a staging buffer that the caller +@@ -574,6 +582,7 @@ def test_checkpoint_state_indices_pick_the_destination_row() -> None: + inp, + [64, 64, 128], + rows=rows, ++ dtype=dtype, + checkpoint_state_indices=torch.tensor( + indices, device="cuda", dtype=torch.int32 + ), +@@ -586,7 +595,7 @@ def test_checkpoint_state_indices_pick_the_destination_row() -> None: + continue + seq, offset = written[row] + _, truncated = _run(_prefix_inputs(inp, seq, offset), use_fused=True) +- torch.testing.assert_close(ckpt[row], truncated[0], rtol=0, atol=0) ++ torch.testing.assert_close(ckpt[row], truncated[0].to(dtype), rtol=0, atol=0) + + + @pytest.mark.parametrize("groups", [2, 4]) +@@ -811,11 +820,14 @@ def test_paged_state_survives_the_group_split(groups: int, monkeypatch) -> None: + torch.testing.assert_close(cache[idx.long()], ht_ref, rtol=5e-3, atol=5e-3) + + +-def test_paged_state_and_checkpoint_export_together() -> None: ++@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) ++def test_paged_state_and_checkpoint_export_together(dtype: torch.dtype) -> None: + """The checkpoint lands in its own row while the walk owns the state row.""" + _requires_kernel() + inp = _inputs([1024], seed=14) ++ inp["h0"] = inp["h0"].to(dtype).float() + cache, idx, warm = _paged_cache(inp, [True], [3], slots=12) ++ cache = cache.to(dtype) + ckidx = torch.tensor([8], device="cuda", dtype=torch.int32) + + o_ref, ht_ref = _run(inp, use_fused=True) +@@ -833,8 +845,10 @@ def test_paged_state_and_checkpoint_export_together() -> None: + checkpoint_state_indices=ckidx, + ) + torch.testing.assert_close(o_got, o_ref, rtol=5e-3, atol=5e-3) +- torch.testing.assert_close(cache[idx.long()], ht_ref, rtol=5e-3, atol=5e-3) +- torch.testing.assert_close(cache[ckidx.long()], ck_ref, rtol=5e-3, atol=5e-3) ++ torch.testing.assert_close(cache[idx.long()].float(), ht_ref, rtol=5e-3, atol=5e-3) ++ torch.testing.assert_close( ++ cache[ckidx.long()].float(), ck_ref, rtol=5e-3, atol=5e-3 ++ ) + + + def test_paged_state_rejects_a_dense_initial_state() -> None: +@@ -898,7 +912,8 @@ def test_padded_cache_stride_is_honoured(use_fused: bool) -> None: + ) + + +-def test_padded_checkpoint_stride_is_honoured() -> None: ++@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) ++def test_padded_checkpoint_stride_is_honoured(dtype: torch.dtype) -> None: + """The checkpoint destination is the paged cache, which pads stride(0). + + Every other checkpoint case here builds a contiguous buffer, so none of +@@ -908,9 +923,10 @@ def test_padded_checkpoint_stride_is_honoured() -> None: + _requires_kernel() + seqlens = [1024] + inp = _inputs(seqlens, seed=31) ++ inp["h0"] = inp["h0"].to(dtype).float() + row = NUM_HEADS * HEAD_DIM * HEAD_DIM + slots, pad = 16, 512 +- raw = torch.full((slots, row + pad), -7.0, device="cuda", dtype=torch.float32) ++ raw = torch.full((slots, row + pad), -7.0, device="cuda", dtype=dtype) + cache = raw.as_strided( + (slots, NUM_HEADS, HEAD_DIM, HEAD_DIM), + (row + pad, HEAD_DIM * HEAD_DIM, HEAD_DIM, 1), +@@ -942,8 +958,10 @@ def test_padded_checkpoint_stride_is_honoured() -> None: + ) + + torch.testing.assert_close(o_got, o_ref, rtol=5e-3, atol=5e-3) +- torch.testing.assert_close(cache[idx.long()], ht_ref, rtol=5e-3, atol=5e-3) +- torch.testing.assert_close(cache[ck_rows[0]], ck_ref[0], rtol=5e-3, atol=5e-3) ++ torch.testing.assert_close(cache[idx.long()].float(), ht_ref, rtol=5e-3, atol=5e-3) ++ torch.testing.assert_close( ++ cache[ck_rows[0]].float(), ck_ref[0], rtol=5e-3, atol=5e-3 ++ ) + + touched = torch.zeros(slots, dtype=torch.bool, device="cuda") + touched[state_rows + ck_rows] = True diff --git a/inferencex-e2e/benchmarks/multi_node/srt-slurm-recipes/configs/patches/kimik3-nightly/04-vllm-58743-native-prefill.patch b/inferencex-e2e/benchmarks/multi_node/srt-slurm-recipes/configs/patches/kimik3-nightly/04-vllm-58743-native-prefill.patch new file mode 100644 index 0000000000..c01df8737b --- /dev/null +++ b/inferencex-e2e/benchmarks/multi_node/srt-slurm-recipes/configs/patches/kimik3-nightly/04-vllm-58743-native-prefill.patch @@ -0,0 +1,512 @@ +From 7542922112f99a0af8a3fb636a2d1bb5476c61e6 Mon Sep 17 00:00:00 2001 +From: Yichao Zhu +Date: Mon, 19 Oct 2026 17:23:49 +0800 +Subject: [PATCH] Support native BF16 paged states in ROCm KDA prefill + +Read and write BF16 persistent states in the HIP chunk kernel while retaining FP32 accumulation and scan workspaces. Avoid gathering and scattering paged state, preserving inactive cache rows. Extend the existing tests for BF16 paging and checkpoint resume under graph replay. + +Signed-off-by: Yichao Zhu +(cherry picked from commit 85c7ecc6eae032f6af5ab44485ca42fb76ca31a3) +--- + csrc/kimi_k3/fused_kda_chunk_kernel_rocm.cu | 98 +++++++++------- + tests/models/kimi_k3/test_amd_kda_chunk.py | 123 +++++++++++++++++--- + vllm/models/kimi_k3/amd/ops/kda_chunk.py | 4 +- + vllm/models/kimi_k3/amd/ops/kda_prefill.py | 26 ++--- + 4 files changed, 172 insertions(+), 79 deletions(-) + +diff --git a/csrc/kimi_k3/fused_kda_chunk_kernel_rocm.cu b/csrc/kimi_k3/fused_kda_chunk_kernel_rocm.cu +index 3ed350c85..f6f9845dd 100644 +--- a/csrc/kimi_k3/fused_kda_chunk_kernel_rocm.cu ++++ b/csrc/kimi_k3/fused_kda_chunk_kernel_rocm.cu +@@ -121,11 +121,12 @@ struct Params { + // separate pointers so that a paged cache (both equal) and a dense + // [N, H, V, K] pair (two tensors, identity rows) share one code path. + // Either pointer may be null. +- const float* cache_in; +- float* cache_out; ++ const void* cache_in; ++ void* cache_out; ++ bool cache_bf16; + const int* sidx; + const bool* hh0; +- // Floats between consecutive rows: the tensor's stride(0), which a paged ++ // Elements between consecutive rows: the tensor's stride(0), which a paged + // cache may pad beyond H * kV * kK. Only each row itself is dense. + int64_t crow; + int64_t ckrow; // the same, for the checkpoint destination +@@ -165,6 +166,17 @@ struct Params { + // with a leading dimension of 0, so nothing inside the chunk loop changes. + [[maybe_unused]] __device__ bf16_t g_uzero[kK]; + ++__device__ __forceinline__ void store_state4(void* dst, int64_t offset, ++ f32x4 value, bool bf16) { ++ if (bf16) { ++ const bf16x4 packed = __builtin_convertvector(value, bf16x4); ++ __builtin_memcpy(static_cast(dst) + offset, &packed, ++ sizeof(packed)); ++ } else { ++ *reinterpret_cast(static_cast(dst) + offset) = value; ++ } ++} ++ + union Frag { + bf16x8 v; + uint32_t d[4]; +@@ -760,13 +772,22 @@ __global__ __launch_bounds__(64 * WAVES * KS, + // state the scan composed for it; in pass one they start at zero, which is + // the B_g the scan expects. Transfer blocks start at the identity. + if (HAS_H0 && g == 0 && warm && !from_sin) { +- const float* s0 = +- p.cache_in + h0_row + +- ((static_cast(h) * kV + v0 + lrow) * kK + 4 * lgrp); ++ const int64_t base = ++ h0_row + ((static_cast(h) * kV + v0 + lrow) * kK + 4 * lgrp); + for (int vt = 0; vt < kNV; ++vt) +- for (int kt = 0; kt < kNKTL; ++kt) +- S[kt][vt] = *reinterpret_cast(s0 + vt * 16 * kK + +- (kt0 + kt) * 16); ++ for (int kt = 0; kt < kNKTL; ++kt) { ++ const int64_t offset = base + vt * 16 * kK + (kt0 + kt) * 16; ++ if (p.cache_bf16) { ++ bf16x4 packed; ++ __builtin_memcpy(&packed, ++ static_cast(p.cache_in) + offset, ++ sizeof(packed)); ++ S[kt][vt] = __builtin_convertvector(packed, f32x4); ++ } else { ++ S[kt][vt] = *reinterpret_cast( ++ static_cast(p.cache_in) + offset); ++ } ++ } + } else if (PASS2 && (g > 0 || (from_sin && warm))) { + const float* s0 = p.sin_ + + (static_cast(g) * p.NH() + nh) * +@@ -848,13 +869,7 @@ __global__ __launch_bounds__(64 * WAVES * KS, + for (int vt = 0; vt < kNV; ++vt) + for (int kt = 0; kt < kNKTL; ++kt) { + const int64_t offset = base + vt * 16 * kK + (kt0 + kt) * 16; +- if (p.ckpt_bf16) { +- *reinterpret_cast(static_cast(p.ckpt) + offset) = +- __builtin_convertvector(S[kt][vt], bf16x4); +- } else { +- *reinterpret_cast(static_cast(p.ckpt) + offset) = +- S[kt][vt]; +- } ++ store_state4(p.ckpt, offset, S[kt][vt], p.ckpt_bf16); + } + }; + +@@ -1058,15 +1073,18 @@ __global__ __launch_bounds__(64 * WAVES * KS, + + // With G == 1 this is the final state. With G > 1 it is what the group + // reaches from a zero start, which the scan then composes. +- float* s1 = nullptr; +- int64_t st_vt = 16 * kK, st_kt = 16; // the cache row is [H][V][K] fp32 + if constexpr (PASS2) { +- // Only the last group holds the sequence's final state. The row's inner +- // layout is [V][K], so this is one f32x4 per lane per tile, strictly +- // inside the row this block owns. +- if (last_group && p.cache_out != nullptr) +- s1 = p.cache_out + h0_row + +- ((static_cast(h) * kV + v0 + lrow) * kK + 4 * lgrp); ++ // Only the last group writes the persistent state; group workspaces ++ // retain the accumulator precision regardless of the cache dtype. ++ if (last_group && p.cache_out != nullptr) { ++ const int64_t base = ++ h0_row + ((static_cast(h) * kV + v0 + lrow) * kK + 4 * lgrp); ++ for (int vt = 0; vt < kNV; ++vt) ++ for (int kt = 0; kt < kNKTL; ++kt) ++ store_state4(p.cache_out, base + vt * 16 * kK + (kt0 + kt) * 16, ++ S[kt][vt], p.cache_bf16); ++ } ++ return; + } else if (xfer) { + // Publish C_g = M_g - diag(A_g). A_g is broadcast through the + // double-buffer half the last chunk did not use, so the diagonal comes off +@@ -1095,18 +1113,15 @@ __global__ __launch_bounds__(64 * WAVES * KS, + } else { + if (slice == 0 && tid < kK) + p.ag[(static_cast(g) * p.NH() + nh) * kK + tid] = aacc; +- s1 = p.bg + +- (static_cast(g) * p.NH() + nh) * +- (static_cast(kV) * kK) + +- static_cast(v0 / 16) * kStTile + lane * 4; +- st_vt = kStTile; +- st_kt = kStKt; +- } +- if (s1 != nullptr) ++ float* const s1 = p.bg + ++ (static_cast(g) * p.NH() + nh) * ++ (static_cast(kV) * kK) + ++ static_cast(v0 / 16) * kStTile + lane * 4; + for (int vt = 0; vt < kNV; ++vt) + for (int kt = 0; kt < kNKTL; ++kt) +- *reinterpret_cast(s1 + vt * st_vt + (kt0 + kt) * st_kt) = ++ *reinterpret_cast(s1 + vt * kStTile + (kt0 + kt) * kStKt) = + S[kt][vt]; ++ } + #else + kda_unsupported_arch("kda_chunk_fused"); + #endif // ifdef ROCM_USE_FUSED_PREFILL_KDA +@@ -2173,8 +2188,8 @@ void fused_kda_chunk( + // reads and writes row `state_indices[n]` of a paged cache. Without it, the + // dense `initial_state` / `final_state` pair is used, which is the same + // layout with one row per sequence in order. +- float const* cache_in = h0_ptr; +- float* cache_out = ht_ptr; ++ void const* cache_in = h0_ptr; ++ void* cache_out = ht_ptr; + int const* sidx_ptr = nullptr; + bool const* hh0_ptr = nullptr; + if (state_cache.has_value()) { +@@ -2186,13 +2201,14 @@ void fused_kda_chunk( + "state_cache needs state_indices and has_initial_state"); + STD_TORCH_CHECK( + state_cache->is_cuda() && +- state_cache->scalar_type() == ScalarType::Float && ++ (state_cache->scalar_type() == ScalarType::Float || ++ state_cache->scalar_type() == ScalarType::BFloat16) && + state_cache->dim() == 4 && state_cache->size(1) == num_heads && + state_cache->size(2) == kV && state_cache->size(3) == kK && + state_cache->stride(3) == 1 && state_cache->stride(2) == kK && + state_cache->stride(1) == kV * kK, +- "state_cache must be an fp32 [slots, H, 128, 128] tensor whose rows " +- "are dense; stride(0) may be padded"); ++ "state_cache must be an fp32 or bf16 [slots, H, 128, 128] tensor whose " ++ "rows are dense; stride(0) may be padded"); + STD_TORCH_CHECK(state_indices->is_cuda() && + state_indices->is_contiguous() && + state_indices->scalar_type() == ScalarType::Int && +@@ -2208,8 +2224,8 @@ void fused_kda_chunk( + // A block reads row `state_indices[n]` at entry and writes it at exit, so + // `checkpoint_state_indices` must name no row in `state_indices`; a + // checkpoint aimed at one would be a write racing that read. +- cache_in = static_cast(state_cache->data_ptr()); +- cache_out = static_cast(state_cache->data_ptr()); ++ cache_in = state_cache->data_ptr(); ++ cache_out = state_cache->data_ptr(); + sidx_ptr = static_cast(state_indices->data_ptr()); + hh0_ptr = static_cast(has_initial_state->data_ptr()); + } +@@ -2242,6 +2258,8 @@ void fused_kda_chunk( + params.NHv = static_cast(nh_total); + params.cache_in = cache_in; + params.cache_out = cache_out; ++ params.cache_bf16 = state_cache.has_value() && ++ state_cache->scalar_type() == ScalarType::BFloat16; + params.sidx = sidx_ptr; + params.hh0 = hh0_ptr; + params.crow = +diff --git a/tests/models/kimi_k3/test_amd_kda_chunk.py b/tests/models/kimi_k3/test_amd_kda_chunk.py +index e3e000137..c7a2c511f 100644 +--- a/tests/models/kimi_k3/test_amd_kda_chunk.py ++++ b/tests/models/kimi_k3/test_amd_kda_chunk.py +@@ -645,13 +645,15 @@ def test_triton_fallback_refuses_a_checkpoint_export() -> None: + # --------------------------------------------------------------------------- + + +-def _paged_cache(inp: dict, warm: list[bool], rows: list[int], slots: int): ++def _paged_cache( ++ inp: dict, warm: list[bool], rows: list[int], slots: int, dtype=torch.float32 ++): + """A cache holding each sequence's h0 at its row, a fill pattern elsewhere.""" + cache = torch.full( + (slots, NUM_HEADS, HEAD_DIM, HEAD_DIM), + -7.0, + device="cuda", +- dtype=torch.float32, ++ dtype=dtype, + ) + idx = torch.tensor(rows, device="cuda", dtype=torch.int32) + for n, row in enumerate(rows): +@@ -734,14 +736,18 @@ def test_paged_bf16_state_matches_dense_pair(use_fused: bool) -> None: + + + @pytest.mark.parametrize("use_fused", [False, True]) +-def test_cold_rows_start_from_zero_not_from_cache_junk(use_fused: bool) -> None: ++@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) ++def test_cold_rows_start_from_zero_not_from_cache_junk( ++ use_fused: bool, dtype: torch.dtype ++) -> None: + """`has_initial_state=False` must ignore whatever the row already holds.""" + _requires_kernel() + seqlens = [256, 256, 256] + inp = _inputs(seqlens, seed=12) ++ inp["h0"] = inp["h0"].to(dtype).float() + warm = [True, False, True] + rows = [4, 9, 1] +- cache, idx, warm_t = _paged_cache(inp, warm, rows, slots=16) ++ cache, idx, warm_t = _paged_cache(inp, warm, rows, slots=16, dtype=dtype) + + # The dense reference zeroes exactly the cold sequences' initial states. + ref = dict(inp) +@@ -751,11 +757,14 @@ def test_cold_rows_start_from_zero_not_from_cache_junk(use_fused: bool) -> None: + o_got, _ = _run_paged(inp, cache, idx, warm_t, use_fused=use_fused) + + torch.testing.assert_close(o_got, o_ref, rtol=5e-3, atol=5e-3) +- torch.testing.assert_close(cache[idx.long()], ht_ref, rtol=5e-3, atol=5e-3) ++ torch.testing.assert_close(cache[idx.long()].float(), ht_ref, rtol=5e-3, atol=5e-3) + + + @pytest.mark.parametrize("use_fused", [False, True]) +-def test_state_indices_accept_a_block_table_column(use_fused: bool) -> None: ++@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) ++def test_state_indices_accept_a_block_table_column( ++ use_fused: bool, dtype: torch.dtype ++) -> None: + # The layer passes block_table[:, 0], which is contiguous only while the + # table has one column. Speculative decoding reserves a state page per + # draft token, so the table widens and that column becomes strided. +@@ -764,7 +773,9 @@ def test_state_indices_accept_a_block_table_column(use_fused: bool) -> None: + inp = _inputs(seqlens, seed=43) + n = len(seqlens) + rows = [3 * i + 5 for i in range(n)] +- cache, idx, warm = _paged_cache(inp, [True] * n, rows, slots=3 * n + 12) ++ cache, idx, warm = _paged_cache( ++ inp, [True] * n, rows, slots=3 * n + 12, dtype=dtype ++ ) + ref_cache = cache.clone() + + table = torch.zeros(n, 9, dtype=torch.int32, device="cuda") +@@ -780,7 +791,10 @@ def test_state_indices_accept_a_block_table_column(use_fused: bool) -> None: + + + @pytest.mark.parametrize("use_fused", [False, True]) +-def test_padded_zero_length_sequence_is_inert(use_fused: bool) -> None: ++@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) ++def test_padded_zero_length_sequence_is_inert( ++ use_fused: bool, dtype: torch.dtype, monkeypatch ++) -> None: + # A batch carrying speculative decodes pads its non-spec group with + # zero-length sequences whose cache row still belongs to a request, so + # consuming no tokens has to mean writing nothing and perturbing nobody. +@@ -791,11 +805,21 @@ def test_padded_zero_length_sequence_is_inert(use_fused: bool) -> None: + padded, unpadded = _inputs([*live, 0], seed=41), _inputs(live, seed=41) + rows = [3 * i + 5 for i in range(4)] + +- cache, idx, warm = _paged_cache(padded, [True] * 4, rows, slots=24) ++ cache, idx, warm = _paged_cache(padded, [True] * 4, rows, slots=24, dtype=dtype) + before = cache.clone() +- _run_paged(padded, cache, idx, warm, use_fused=use_fused) +- +- ref_cache, ref_idx, ref_warm = _paged_cache(unpadded, [True] * 3, rows[:3], 24) ++ # Poison unwritten outputs rather than relying on allocator contents. ++ deterministic = torch.are_deterministic_algorithms_enabled() ++ warn_only = torch.is_deterministic_algorithms_warn_only_enabled() ++ monkeypatch.setattr(torch.utils.deterministic, "fill_uninitialized_memory", True) ++ try: ++ torch.use_deterministic_algorithms(True) ++ _run_paged(padded, cache, idx, warm, use_fused=use_fused) ++ finally: ++ torch.use_deterministic_algorithms(deterministic, warn_only=warn_only) ++ ++ ref_cache, ref_idx, ref_warm = _paged_cache( ++ unpadded, [True] * 3, rows[:3], 24, dtype=dtype ++ ) + _run_paged(unpadded, ref_cache, ref_idx, ref_warm, use_fused=use_fused) + + assert torch.equal(cache[rows[3]], before[rows[3]]), ( +@@ -805,19 +829,23 @@ def test_padded_zero_length_sequence_is_inert(use_fused: bool) -> None: + + + @pytest.mark.parametrize("groups", [2, 4]) +-def test_paged_state_survives_the_group_split(groups: int, monkeypatch) -> None: ++@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) ++def test_paged_state_survives_the_group_split( ++ groups: int, dtype: torch.dtype, monkeypatch ++) -> None: + """At G>1 the read and the write are different blocks of one launch.""" + _requires_kernel() + _force_groups(monkeypatch, groups) + inp = _inputs([1024, 512], seed=13) ++ inp["h0"] = inp["h0"].to(dtype).float() + rows = [6, 2] +- cache, idx, warm = _paged_cache(inp, [True, True], rows, slots=12) ++ cache, idx, warm = _paged_cache(inp, [True, True], rows, slots=12, dtype=dtype) + + o_ref, ht_ref = _run(inp, use_fused=True) + o_got, _ = _run_paged(inp, cache, idx, warm, use_fused=True) + + torch.testing.assert_close(o_got, o_ref, rtol=5e-3, atol=5e-3) +- torch.testing.assert_close(cache[idx.long()], ht_ref, rtol=5e-3, atol=5e-3) ++ torch.testing.assert_close(cache[idx.long()].float(), ht_ref, rtol=5e-3, atol=5e-3) + + + @pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +@@ -867,8 +895,64 @@ def test_paged_state_rejects_a_dense_initial_state() -> None: + ) + + ++@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) ++@pytest.mark.parametrize("groups", [1, 2, 4]) ++def test_paged_checkpoint_can_resume_with_graph_replay( ++ dtype: torch.dtype, groups: int, monkeypatch ++) -> None: ++ """Resume from a stored checkpoint, including cold zero-token padding.""" ++ _requires_kernel() ++ _force_groups(monkeypatch, groups) ++ inp = _inputs([1024], seed=47) ++ cache, idx, warm = _paged_cache(inp, [True], [3], slots=12, dtype=dtype) ++ ckidx = torch.tensor([8], device="cuda", dtype=torch.int32) ++ _run_paged( ++ inp, ++ cache, ++ idx, ++ warm, ++ use_fused=True, ++ checkpoint_state=cache, ++ checkpoint_offsets=torch.tensor([512], device="cuda", dtype=torch.int32), ++ checkpoint_state_indices=ckidx, ++ ) ++ ++ # Compare against the same rounded checkpoint, not an uninterrupted FP32 walk. ++ tail = dict(inp) ++ for key in ("q", "k", "v", "raw_g", "raw_beta"): ++ tail[key] = inp[key][:, 512:].contiguous() ++ tail["cu"] = torch.tensor([0, 512], device="cuda", dtype=torch.int32) ++ tail["h0"] = cache[ckidx.long()].float() ++ o_ref, ht_ref = _run(tail, use_fused=True) ++ ++ tail["cu"] = torch.tensor([0, 512, 512], device="cuda", dtype=torch.int32) ++ resume_idx = torch.tensor([8, 10], device="cuda", dtype=torch.int32) ++ resume_warm = torch.tensor([True, False], device="cuda", dtype=torch.bool) ++ before = cache.clone() ++ stream = torch.cuda.Stream() ++ stream.wait_stream(torch.cuda.current_stream()) ++ with torch.cuda.stream(stream): ++ _run_paged(tail, cache, resume_idx, resume_warm, use_fused=True) ++ torch.cuda.current_stream().wait_stream(stream) ++ cache.copy_(before) ++ graph = torch.cuda.CUDAGraph() ++ with torch.cuda.graph(graph): ++ o_got, ht_got = _run_paged(tail, cache, resume_idx, resume_warm, use_fused=True) ++ for _ in range(2): ++ cache.copy_(before) ++ graph.replay() ++ assert ht_got is None ++ torch.testing.assert_close(o_got, o_ref, rtol=5e-3, atol=5e-3) ++ torch.testing.assert_close( ++ cache[ckidx.long()].float(), ht_ref, rtol=5e-3, atol=5e-3 ++ ) ++ untouched = torch.arange(cache.shape[0], device="cuda") != 8 ++ assert torch.equal(cache[untouched], before[untouched]) ++ ++ + @pytest.mark.parametrize("use_fused", [False, True]) +-def test_padded_cache_stride_is_honoured(use_fused: bool) -> None: ++@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) ++def test_padded_cache_stride_is_honoured(use_fused: bool, dtype: torch.dtype) -> None: + """A mamba page may pad stride(0); only each state row is dense. + + The real `recurrent_state` is not a contiguous tensor, so a kernel that +@@ -879,10 +963,11 @@ def test_padded_cache_stride_is_honoured(use_fused: bool) -> None: + _requires_kernel() + seqlens = [1024, 512] + inp = _inputs(seqlens, seed=21) ++ inp["h0"] = inp["h0"].to(dtype).float() + n = len(seqlens) + row = NUM_HEADS * HEAD_DIM * HEAD_DIM +- slots, pad = 16, 512 +- raw = torch.full((slots, row + pad), -7.0, device="cuda", dtype=torch.float32) ++ slots, pad = 16, 513 if dtype == torch.bfloat16 else 512 ++ raw = torch.full((slots, row + pad), -7.0, device="cuda", dtype=dtype) + cache = raw.as_strided( + (slots, NUM_HEADS, HEAD_DIM, HEAD_DIM), + (row + pad, HEAD_DIM * HEAD_DIM, HEAD_DIM, 1), +@@ -901,7 +986,7 @@ def test_padded_cache_stride_is_honoured(use_fused: bool) -> None: + + assert ht_got is None + torch.testing.assert_close(o_got, o_ref, rtol=5e-3, atol=5e-3) +- torch.testing.assert_close(cache[idx.long()], ht_ref, rtol=5e-3, atol=5e-3) ++ torch.testing.assert_close(cache[idx.long()].float(), ht_ref, rtol=5e-3, atol=5e-3) + + # The padding between rows, and every unowned row, must be untouched. + untouched = torch.ones(slots, dtype=torch.bool, device="cuda") +diff --git a/vllm/models/kimi_k3/amd/ops/kda_chunk.py b/vllm/models/kimi_k3/amd/ops/kda_chunk.py +index a3fbada90..1562b69fc 100644 +--- a/vllm/models/kimi_k3/amd/ops/kda_chunk.py ++++ b/vllm/models/kimi_k3/amd/ops/kda_chunk.py +@@ -235,7 +235,7 @@ def fused_kda_chunk( + output_final_state: whether to return the final recurrent state. + chunk_offsets: int32 per-sequence first chunk index. + checkpoint_state: destination for the mid-prefill state snapshots, +- fp32 ``[rows, H, 128, 128]``. Without ++ fp32 or bf16 ``[rows, H, 128, 128]``. Without + ``checkpoint_state_indices`` it is a staging buffer indexed by + sequence; with them it can be the paged state cache itself. + checkpoint_offsets: int32 ``[N]`` token offsets, each relative to its +@@ -246,7 +246,7 @@ def fused_kda_chunk( + rather than rely on the store happening. + checkpoint_state_indices: optional int32 ``[N]`` destination row per + sequence. A negative row disables the export for that sequence. +- state_cache: the paged recurrent state, fp32 ++ state_cache: the paged recurrent state, fp32 or bf16 + ``[slots, H, 128, 128]``. When given, the walk reads sequence + ``n``'s initial state from row ``state_indices[n]`` and writes its + final state back to the same row, which removes the gather and the +diff --git a/vllm/models/kimi_k3/amd/ops/kda_prefill.py b/vllm/models/kimi_k3/amd/ops/kda_prefill.py +index 18d90b811..e1900b5e5 100644 +--- a/vllm/models/kimi_k3/amd/ops/kda_prefill.py ++++ b/vllm/models/kimi_k3/amd/ops/kda_prefill.py +@@ -86,10 +86,9 @@ def chunk_kda_prefill( + checkpoint_offsets: per-sequence token offset to snapshot at, ``0`` + for none. + checkpoint_state_indices: optional per-sequence destination row. +- state_cache: the paged recurrent state. The fused backend addresses an +- FP32 cache directly. Other cache dtypes are gathered into the FP32 +- accumulator format and cast back after the walk. The returned final +- state is ``None`` in either case. ++ state_cache: the paged recurrent state. The fused backend reads and ++ writes FP32 or BF16 cache rows directly, accumulating in FP32. ++ The returned final state is ``None``. + state_indices: per-sequence cache row. + has_initial_state: per-sequence flag; false starts from a zero state. + +@@ -151,11 +150,10 @@ def chunk_kda_prefill( + has_initial_state=has_initial_state, + ) + +- # Only the fused walk addresses FP32 paged rows directly. The Triton path +- # and lower-precision persistent caches gather the active rows and scatter +- # the final state back here, so callers use one cache API for every backend. ++ # Only the fused walk addresses paged rows directly. The Triton path ++ # gathers the active rows and scatters the final state back here. + scatter_to: torch.Tensor | None = None +- if state_cache is not None and (not fused or state_cache.dtype != torch.float32): ++ if state_cache is not None and not fused: + initial_state = gather_initial_states( + state_cache, state_indices, has_initial_state + ) +@@ -166,9 +164,7 @@ def chunk_kda_prefill( + # Restated for the type checker; `fused` already implies all three. + assert cu_seqlens is not None and lower_bound is not None + assert g_bias is not None +- # The fused HIP prefill kernel accumulates recurrent state in FP32. +- # Convert only active rows; the persistent cache keeps its configured +- # dtype and receives a cast-back below. ++ # The dense state API remains FP32; paged caches bypass this conversion. + if initial_state is not None and initial_state.dtype != torch.float32: + initial_state = initial_state.to(torch.float32) + logger.info_once( +@@ -187,7 +183,7 @@ def chunk_kda_prefill( + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + ) +- o, final_state = fused_kda_chunk( ++ return fused_kda_chunk( + qg=ws["qg"], + w=ws["w"], + u=ws["u"], +@@ -214,12 +210,6 @@ def chunk_kda_prefill( + state_indices=state_indices, + has_initial_state=has_initial_state, + ) +- if scatter_to is not None: +- assert state_indices is not None +- assert final_state is not None +- scatter_to[state_indices.long()] = final_state.to(scatter_to.dtype) +- return o, None +- return o, final_state + + o, final_state = chunk_kda_with_fused_gate( + q=q, diff --git a/inferencex-e2e/benchmarks/multi_node/srt-slurm-recipes/configs/patches/kimik3-nightly/05-vllm-58743-backends.patch b/inferencex-e2e/benchmarks/multi_node/srt-slurm-recipes/configs/patches/kimik3-nightly/05-vllm-58743-backends.patch new file mode 100644 index 0000000000..bb862b66ad --- /dev/null +++ b/inferencex-e2e/benchmarks/multi_node/srt-slurm-recipes/configs/patches/kimik3-nightly/05-vllm-58743-backends.patch @@ -0,0 +1,181 @@ +From 983548e8aa73a29b1d8b5543ac83d231350440f8 Mon Sep 17 00:00:00 2001 +From: Yichao Zhu +Date: Mon, 19 Oct 2026 18:34:00 +0800 +Subject: [PATCH] Select compatible ROCm KDA backends for BF16 states + +Keep FP32-only Gluon and HIP decode paths out of BF16 recurrent-state execution. Preserve FP32 backend priority and validate BF16 decode state writes and backend selection. + +Signed-off-by: Yichao Zhu +(cherry picked from commit 149262f3ca7c3a07837797942342d45f0abe9f64) +--- + tests/models/kimi_k3/test_amd_kda_decode.py | 59 ++++++++++++++++++++- + vllm/models/kimi_k3/amd/kda.py | 14 +++-- + vllm/models/kimi_k3/amd/ops/kda_decode.py | 2 + + 3 files changed, 71 insertions(+), 4 deletions(-) + +diff --git a/tests/models/kimi_k3/test_amd_kda_decode.py b/tests/models/kimi_k3/test_amd_kda_decode.py +index fc6039ee5..8b12d31c8 100644 +--- a/tests/models/kimi_k3/test_amd_kda_decode.py ++++ b/tests/models/kimi_k3/test_amd_kda_decode.py +@@ -14,6 +14,8 @@ Every output of each is compared against the unfused path the AMD layer falls + back to, with the argument layout the layer itself uses. + """ + ++from functools import partial ++ + import pytest + import torch + +@@ -335,7 +337,11 @@ def test_fused_kda_decode_skips_null_block_padding() -> None: + + + @requires_hip_kernel_arch +-def test_decode_weights_staged_after_loading(dist_init) -> None: ++@pytest.mark.parametrize("state_dtype", ["float32", "bfloat16"]) ++@pytest.mark.parametrize("aiter_available", [False, True]) ++def test_decode_weights_staged_after_loading( ++ dist_init, monkeypatch, state_dtype: str, aiter_available: bool ++) -> None: + """The fused-decode weight copies must not depend on the weight loaders. + + Loaders that bypass ``weight_loader`` (e.g. ``--load-format dummy``) left +@@ -344,11 +350,20 @@ def test_decode_weights_staged_after_loading(dist_init) -> None: + """ + _requires_kernel() + from tests.v1.attention.utils import create_vllm_config ++ from vllm._aiter_ops import rocm_aiter_ops + from vllm.config import set_current_vllm_config + from vllm.model_executor.model_loader.utils import process_weights_after_loading + from vllm.models.kimi_k3.amd.kda import KimiK3DeltaAttention ++ from vllm.models.kimi_k3.amd.ops.kda_chunk import is_fused_kda_chunk_supported + from vllm.transformers_utils.configs.kimi_linear import KimiLinearConfig + ++ monkeypatch.setenv("VLLM_SSM_CONV_STATE_LAYOUT", "SD") ++ monkeypatch.setattr( ++ rocm_aiter_ops, "is_kda_gluon_decode_available", lambda: aiter_available ++ ) ++ monkeypatch.setattr( ++ rocm_aiter_ops, "is_kda_gluon_prefill_available", lambda: aiter_available ++ ) + config = KimiLinearConfig( + hidden_size=256, + linear_attn_config={ +@@ -362,10 +377,26 @@ def test_decode_weights_staged_after_loading(dist_init) -> None: + }, + ) + vllm_config = create_vllm_config(model_name="Qwen/Qwen3.5-0.8B", dtype=DTYPE) ++ vllm_config.cache_config.mamba_ssm_cache_dtype = state_dtype + with set_current_vllm_config(vllm_config), torch.device("cuda"): + layer = KimiK3DeltaAttention( + config, vllm_config, prefix="model.layers.0.self_attn" + ) ++ use_aiter = aiter_available and state_dtype == "float32" ++ assert layer.use_aiter_decode == use_aiter ++ expected_prefill = ( ++ "gluon" ++ if use_aiter ++ else "fused" ++ if is_fused_kda_chunk_supported() ++ else "triton" ++ ) ++ assert layer.kda_prefill_backend == expected_prefill ++ if state_dtype == "bfloat16": ++ # Neither FP32-only decode backend may stage weights or take its fast path. ++ assert layer.decode_conv1d_weight is None ++ assert layer.decode_norm_weight is None ++ return + assert layer.decode_conv1d_weight is not None + # The AITER Gluon decode path takes o_norm.weight directly, so only the HIP + # kernel gets an fp32 norm copy. +@@ -625,6 +656,32 @@ def _run_aiter_spec( + return core_attn_out, conv_state, recurrent_state + + ++@pytest.mark.parametrize("num_spec", [0, 3]) ++@torch.inference_mode() ++def test_triton_decode_with_bf16_recurrent_state(num_spec: int) -> None: ++ """BF16 fallback reads rounded states and preserves paged writeback semantics.""" ++ if num_spec: ++ batch = SpecBatch([4, 2, 1], [1, 2, 4], 12, num_spec, seed=5) ++ inp = batch.inp ++ run = partial(_run_spec_triton_chain, batch) ++ else: ++ inp = KdaDecodeInputs(3, 12, num_slots=8, seed=5) ++ run = partial(_run_triton_chain, inp) ++ ++ initial_state = inp.recurrent_state.bfloat16() ++ inp.recurrent_state = initial_state.float() ++ expected_out, expected_conv, expected_state = run() ++ inp.recurrent_state = initial_state ++ actual_out, actual_conv, actual_state = run() ++ ++ assert actual_state.dtype == torch.bfloat16 ++ torch.testing.assert_close(actual_out, expected_out, atol=3e-2, rtol=3e-2) ++ torch.testing.assert_close(actual_conv, expected_conv, atol=0, rtol=0) ++ torch.testing.assert_close( ++ actual_state, expected_state.bfloat16(), atol=2e-3, rtol=3e-2 ++ ) ++ ++ + @requires_aiter_gluon_kda + @pytest.mark.parametrize("num_heads", [12, 24]) + @pytest.mark.parametrize( +diff --git a/vllm/models/kimi_k3/amd/kda.py b/vllm/models/kimi_k3/amd/kda.py +index 4216c6f94..8460c09f7 100644 +--- a/vllm/models/kimi_k3/amd/kda.py ++++ b/vllm/models/kimi_k3/amd/kda.py +@@ -194,8 +194,11 @@ class KimiK3DeltaAttention(GatedDeltaNetAttention): + # ROCm can fuse the whole decode step (conv + recurrence + gated norm) + # into one kernel, which wants a width-major fp32 conv weight staged at + # load time. Everything else keeps the [channel, width] layout. +- conv_state_dtype, _ = self.get_state_dtype() +- self.use_aiter_decode = bool(rocm_aiter_ops.is_kda_gluon_decode_available()) ++ conv_state_dtype, recurrent_state_dtype = self.get_state_dtype() ++ # AITER's Gluon kernels currently require FP32 recurrent states. ++ self.use_aiter_decode = recurrent_state_dtype == torch.float32 and bool( ++ rocm_aiter_ops.is_kda_gluon_decode_available() ++ ) + use_hip_decode = not self.use_aiter_decode and is_fused_kda_decode_supported( + self.local_num_heads, + self.head_dim, +@@ -203,6 +206,7 @@ class KimiK3DeltaAttention(GatedDeltaNetAttention): + self.num_spec, + vllm_config.model_config.dtype, + conv_state_dtype, ++ recurrent_state_dtype, + ) + if use_hip_decode: + logger.info_once("Fused KDA decode kernel (conv+KDA+norm) is enabled.") +@@ -256,7 +260,11 @@ class KimiK3DeltaAttention(GatedDeltaNetAttention): + ) + self.kda_prefill_backend = backend + if backend == "auto": +- if self.use_safe_gate and rocm_aiter_ops.is_kda_gluon_prefill_available(): ++ if ( ++ recurrent_state_dtype == torch.float32 ++ and self.use_safe_gate ++ and rocm_aiter_ops.is_kda_gluon_prefill_available() ++ ): + self.kda_prefill_backend = "gluon" + elif is_fused_kda_chunk_supported(): + self.kda_prefill_backend = "fused" +diff --git a/vllm/models/kimi_k3/amd/ops/kda_decode.py b/vllm/models/kimi_k3/amd/ops/kda_decode.py +index d944e1eb0..236a95d72 100644 +--- a/vllm/models/kimi_k3/amd/ops/kda_decode.py ++++ b/vllm/models/kimi_k3/amd/ops/kda_decode.py +@@ -30,6 +30,7 @@ def is_fused_kda_decode_supported( + num_spec: int, + input_dtype: torch.dtype, + conv_state_dtype: torch.dtype, ++ recurrent_state_dtype: torch.dtype, + ) -> bool: + """Whether the fused decode kernel can serve this layer on this device.""" + from vllm.platforms.rocm import on_gfx942, on_gfx950 +@@ -41,6 +42,7 @@ def is_fused_kda_decode_supported( + or num_spec != 0 + or input_dtype != torch.bfloat16 + or conv_state_dtype != torch.bfloat16 ++ or recurrent_state_dtype != torch.float32 + or is_conv_state_dim_first() + or not hasattr(torch.ops._C, "fused_kda_decode") + ): diff --git a/inferencex-e2e/benchmarks/multi_node/srt-slurm-recipes/configs/patches/kimik3-nightly/06-vllm-58743-tests.patch b/inferencex-e2e/benchmarks/multi_node/srt-slurm-recipes/configs/patches/kimik3-nightly/06-vllm-58743-tests.patch new file mode 100644 index 0000000000..330ad5bf00 --- /dev/null +++ b/inferencex-e2e/benchmarks/multi_node/srt-slurm-recipes/configs/patches/kimik3-nightly/06-vllm-58743-tests.patch @@ -0,0 +1,115 @@ +From a02a204017e4e0ef88e2ae0d1eac76e8ac1b9332 Mon Sep 17 00:00:00 2001 +From: Yichao Zhu +Date: Mon, 19 Oct 2026 19:30:48 +0800 +Subject: [PATCH] Keep the KDA dense state contract and consolidate BF16 tests + +Remove the implicit dense initial-state conversion now that HIP prefill handles BF16 paged states natively. Extend the existing paged-state comparison instead of maintaining a separate BF16 case. + +Signed-off-by: Yichao Zhu +(cherry picked from commit ae6bb332600412fdf95fec94548e091a7638be99) +--- + tests/models/kimi_k3/test_amd_kda_chunk.py | 36 ++++++---------------- + vllm/models/kimi_k3/amd/ops/kda_prefill.py | 8 ++--- + 2 files changed, 13 insertions(+), 31 deletions(-) + +diff --git a/tests/models/kimi_k3/test_amd_kda_chunk.py b/tests/models/kimi_k3/test_amd_kda_chunk.py +index c7a2c511f..6865dea65 100644 +--- a/tests/models/kimi_k3/test_amd_kda_chunk.py ++++ b/tests/models/kimi_k3/test_amd_kda_chunk.py +@@ -687,15 +687,19 @@ def _run_paged(inp: dict, cache, idx, warm, use_fused: bool, **kw): + + @pytest.mark.parametrize("use_fused", [False, True]) + @pytest.mark.parametrize("seqlens", [[512], [2048], [320, 64], [513, 64, 1]]) ++@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) + def test_paged_state_matches_the_dense_pair( +- use_fused: bool, seqlens: list[int] ++ use_fused: bool, seqlens: list[int], dtype: torch.dtype + ) -> None: + """A paged cache must give the same result as the dense pair.""" + _requires_kernel() + inp = _inputs(seqlens, seed=11) ++ inp["h0"] = inp["h0"].to(dtype).float() + n = len(seqlens) + rows = [3 * i + 5 for i in range(n)] +- cache, idx, warm = _paged_cache(inp, [True] * n, rows, slots=3 * n + 12) ++ cache, idx, warm = _paged_cache( ++ inp, [True] * n, rows, slots=3 * n + 12, dtype=dtype ++ ) + before = cache.clone() + + o_ref, ht_ref = _run(inp, use_fused=use_fused) +@@ -703,7 +707,9 @@ def test_paged_state_matches_the_dense_pair( + + assert ht_got is None, "the final state is in the cache, not returned" + torch.testing.assert_close(o_got, o_ref, rtol=5e-3, atol=5e-3) +- torch.testing.assert_close(cache[idx.long()], ht_ref, rtol=5e-3, atol=5e-3) ++ torch.testing.assert_close( ++ cache[idx.long()], ht_ref.to(dtype), rtol=5e-3, atol=5e-3 ++ ) + + untouched = torch.ones(cache.shape[0], dtype=torch.bool, device="cuda") + untouched[idx.long()] = False +@@ -712,29 +718,6 @@ def test_paged_state_matches_the_dense_pair( + ) + + +-@pytest.mark.parametrize("use_fused", [False, True]) +-def test_paged_bf16_state_matches_dense_pair(use_fused: bool) -> None: +- """The FP32 walk may load and store a BF16 persistent cache.""" +- if use_fused: +- _requires_kernel() +- seqlens = [513, 64, 1] +- inp = _inputs(seqlens, seed=17) +- rows = [2, 7, 11] +- cache, idx, warm = _paged_cache(inp, [True] * len(seqlens), rows, slots=16) +- cache = cache.to(torch.bfloat16) +- +- dense = dict(inp) +- dense["h0"] = cache[idx.long()].float() +- o_ref, ht_ref = _run(dense, use_fused=use_fused) +- o_got, ht_got = _run_paged(inp, cache, idx, warm, use_fused=use_fused) +- +- assert ht_got is None +- torch.testing.assert_close(o_got, o_ref, rtol=5e-3, atol=5e-3) +- torch.testing.assert_close( +- cache[idx.long()], ht_ref.to(torch.bfloat16), rtol=5e-3, atol=5e-3 +- ) +- +- + @pytest.mark.parametrize("use_fused", [False, True]) + @pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) + def test_cold_rows_start_from_zero_not_from_cache_junk( +@@ -966,6 +949,7 @@ def test_padded_cache_stride_is_honoured(use_fused: bool, dtype: torch.dtype) -> + inp["h0"] = inp["h0"].to(dtype).float() + n = len(seqlens) + row = NUM_HEADS * HEAD_DIM * HEAD_DIM ++ # An odd BF16 row stride exercises unaligned vector loads and stores. + slots, pad = 16, 513 if dtype == torch.bfloat16 else 512 + raw = torch.full((slots, row + pad), -7.0, device="cuda", dtype=dtype) + cache = raw.as_strided( +diff --git a/vllm/models/kimi_k3/amd/ops/kda_prefill.py b/vllm/models/kimi_k3/amd/ops/kda_prefill.py +index e1900b5e5..ecdb36f4b 100644 +--- a/vllm/models/kimi_k3/amd/ops/kda_prefill.py ++++ b/vllm/models/kimi_k3/amd/ops/kda_prefill.py +@@ -150,8 +150,9 @@ def chunk_kda_prefill( + has_initial_state=has_initial_state, + ) + +- # Only the fused walk addresses paged rows directly. The Triton path +- # gathers the active rows and scatters the final state back here. ++ # Only the fused walk addresses the paged rows directly. The Triton path ++ # gathers the rows it needs and scatters the results back here, so callers ++ # pass the cache the same way for either backend. + scatter_to: torch.Tensor | None = None + if state_cache is not None and not fused: + initial_state = gather_initial_states( +@@ -164,9 +165,6 @@ def chunk_kda_prefill( + # Restated for the type checker; `fused` already implies all three. + assert cu_seqlens is not None and lower_bound is not None + assert g_bias is not None +- # The dense state API remains FP32; paged caches bypass this conversion. +- if initial_state is not None and initial_state.dtype != torch.float32: +- initial_state = initial_state.to(torch.float32) + logger.info_once( + "Kimi-K3 KDA prefill: dispatching the fused ROCm chunk kernel." + ) diff --git a/inferencex-e2e/benchmarks/multi_node/srt-slurm-recipes/configs/patches/kimik3-nightly/07-vllm-60808.patch b/inferencex-e2e/benchmarks/multi_node/srt-slurm-recipes/configs/patches/kimik3-nightly/07-vllm-60808.patch new file mode 100644 index 0000000000..c3d485cbee --- /dev/null +++ b/inferencex-e2e/benchmarks/multi_node/srt-slurm-recipes/configs/patches/kimik3-nightly/07-vllm-60808.patch @@ -0,0 +1,569 @@ +From a0ef2bdc9879a40cd1d6e6566e505622194fdf32 Mon Sep 17 00:00:00 2001 +From: Mir Mustafa Ali +Date: Wed, 21 Oct 2026 07:11:16 -0500 +Subject: [PATCH] [Bugfix][KV Connector][ROCm] MoRIIO READ: await every read + before the forward (#60808) + +Signed-off-by: Mir Mustafa Ali +Co-authored-by: Claude +(cherry picked from commit db3b0f0fae4bc98492b3863685e541e62bc59576) +--- + .../kv_connector/unit/test_moriio_tp_ack.py | 278 +++++++++++++++++- + .../v1/moriio/moriio_connector.py | 116 ++------ + .../kv_connector/v1/moriio/moriio_engine.py | 21 +- + 3 files changed, 316 insertions(+), 99 deletions(-) + +diff --git a/tests/v1/kv_connector/unit/test_moriio_tp_ack.py b/tests/v1/kv_connector/unit/test_moriio_tp_ack.py +index 34d193e27..72453f743 100644 +--- a/tests/v1/kv_connector/unit/test_moriio_tp_ack.py ++++ b/tests/v1/kv_connector/unit/test_moriio_tp_ack.py +@@ -3,9 +3,11 @@ + + import threading + from types import SimpleNamespace ++from unittest.mock import Mock + + import pytest + ++from vllm.distributed.kv_transfer.kv_connector.v1.moriio import moriio_engine + from vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_common import ( + MoRIIOMode, + MoRIIOTransferAck, +@@ -460,9 +462,9 @@ def test_failed_read_reports_blocks_only_without_hma(has_mamba, expected_invalid + assert worker.get_block_ids_with_load_errors() == expected_invalid + + +-def test_hybrid_step_barrier_fails_closed(monkeypatch): ++def test_hybrid_step_barrier_fails_closed(): + class FailingWrapper: +- def waiting_for_transfer_complete(self, _statuses): ++ def waiting_for_transfer_complete(self, _statuses, fail_fast=True): + raise TransferError("failed") + + def shutdown(self): +@@ -470,18 +472,278 @@ def test_hybrid_step_barrier_fails_closed(monkeypatch): + + worker = MoRIIOConnectorWorker.__new__(MoRIIOConnectorWorker) + worker._has_mamba = True +- worker._reads_issued_this_step = [object()] +- worker._mamba_reads_this_step = [object()] ++ worker._reads_issued_this_step = [_ReadStatus("failed")] + worker.moriio_wrapper = FailingWrapper() +- monkeypatch.setattr( +- "vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_connector.get_forward_context", +- lambda: SimpleNamespace(cudagraph_runtime_mode=None), +- ) + + with pytest.raises(TransferError, match="failed"): + worker._await_reads_issued_this_step() + + ++def test_step_barrier_awaits_every_read_before_forward(): ++ class RecordingWrapper: ++ def __init__(self): ++ self.awaited = [] ++ ++ def waiting_for_transfer_complete(self, statuses, fail_fast=True): ++ self.awaited.append(list(statuses)) ++ ++ def shutdown(self): ++ pass ++ ++ attn, indexer, mtp_attn = object(), object(), object() ++ worker = MoRIIOConnectorWorker.__new__(MoRIIOConnectorWorker) ++ worker._has_mamba = False ++ worker._reads_issued_this_step = [attn, indexer, mtp_attn] ++ worker.moriio_wrapper = RecordingWrapper() ++ ++ worker._await_reads_issued_this_step() ++ worker._await_reads_issued_this_step() ++ ++ assert worker.moriio_wrapper.awaited == [[attn, indexer, mtp_attn]] ++ assert worker._reads_issued_this_step == [] ++ ++ ++def test_wait_for_layer_load_is_noop(): ++ class _Boom: ++ def __getattr__(self, name): ++ raise AssertionError("wait_for_layer_load must not touch the worker") ++ ++ connector = MoRIIOConnector.__new__(MoRIIOConnector) ++ connector.connector_worker = _Boom() ++ ++ assert connector.wait_for_layer_load("layer") is None ++ ++ ++def test_step_barrier_without_mamba_does_not_raise_on_failure(): ++ # Without HMA, failures go to get_block_ids_with_load_errors. ++ class FailingWrapper: ++ def waiting_for_transfer_complete(self, _statuses, fail_fast=True): ++ raise TransferError("failed") ++ ++ def shutdown(self): ++ pass ++ ++ worker = MoRIIOConnectorWorker.__new__(MoRIIOConnectorWorker) ++ worker._has_mamba = False ++ worker._reads_issued_this_step = [_ReadStatus("failed")] ++ worker.moriio_wrapper = FailingWrapper() ++ ++ worker._await_reads_issued_this_step() ++ assert worker._reads_issued_this_step == [] ++ ++ ++class _ReadStatus: ++ def __init__(self, state: str): ++ self.state = state ++ ++ def Succeeded(self): ++ return self.state == "success" ++ ++ def Failed(self): ++ return self.state == "failed" ++ ++ def Message(self): ++ return "SQ full" ++ ++ def Code(self): ++ return -1 ++ ++ ++class _PolledStatus(_ReadStatus): ++ """Lands on its succeed_after-th Succeeded() poll; never if None.""" ++ ++ def __init__(self, succeed_after: int | None = None): ++ super().__init__("in_progress") ++ self.succeed_after = succeed_after ++ self.polls = 0 ++ ++ def Succeeded(self): ++ self.polls += 1 ++ if self.polls == self.succeed_after: ++ self.state = "success" ++ return super().Succeeded() ++ ++ ++class _FakeClock: ++ def __init__(self): ++ self.now = 0.0 ++ self.sleeps = 0 ++ ++ def monotonic(self): ++ return self.now ++ ++ def sleep(self, seconds): ++ self.sleeps += 1 ++ self.now += seconds ++ ++ ++@pytest.fixture ++def fake_clock(monkeypatch): ++ clock = _FakeClock() ++ monkeypatch.setattr(moriio_engine, "time", clock) ++ return clock ++ ++ ++def _worker_with_failed_and_pending_read(has_mamba, pending_b, monkeypatch): ++ failed_a = _ReadStatus("failed") ++ worker = MoRIIOConnectorWorker.__new__(MoRIIOConnectorWorker) ++ worker._has_mamba = has_mamba ++ worker.world_size = 1 ++ worker.moriio_config = SimpleNamespace(recv_abort_timeout=600.0) ++ worker.moriio_wrapper = MoRIIOWrapper(moriio_engine=object(), transfer_timeout=5.0) ++ monkeypatch.setattr(worker.moriio_wrapper, "send_notify", lambda *a, **kw: None) ++ worker._reads_issued_this_step = [failed_a, pending_b] ++ worker._recving_transfers = { ++ "a": {"layer0": [failed_a]}, ++ "b": {"layer0": [pending_b]}, ++ } ++ worker._recving_transfers_callback_addr = { ++ "a": ("127.0.0.1", "7000", "tx-a"), ++ "b": ("127.0.0.1", "7000", "tx-b"), ++ } ++ worker._recving_transfers_start = {} ++ worker._recving_local_blocks = {"a": [1, 2], "b": [3, 4]} ++ worker._invalid_block_ids = set() ++ return worker ++ ++ ++def test_step_barrier_failure_still_awaits_other_requests(monkeypatch, fake_clock): ++ late_b = _PolledStatus(succeed_after=5) ++ worker = _worker_with_failed_and_pending_read(False, late_b, monkeypatch) ++ ++ worker._await_reads_issued_this_step() ++ ++ assert late_b.state == "success" ++ assert late_b.polls >= 5 ++ assert fake_clock.sleeps >= 4 ++ assert worker._pop_done_transfers() == {"tx-b"} ++ assert worker.get_block_ids_with_load_errors() == {1, 2} ++ ++ ++def test_hybrid_step_barrier_failure_does_not_wait(monkeypatch, fake_clock): ++ pending_b = _PolledStatus() ++ worker = _worker_with_failed_and_pending_read(True, pending_b, monkeypatch) ++ ++ with pytest.raises(TransferError, match="SQ full"): ++ worker._await_reads_issued_this_step() ++ ++ assert pending_b.state == "in_progress" ++ assert fake_clock.sleeps == 0 ++ ++ ++@pytest.mark.parametrize("has_mamba", [False, True]) ++def test_step_barrier_timeout_fails_closed_while_reads_in_flight(has_mamba, fake_clock): ++ worker = MoRIIOConnectorWorker.__new__(MoRIIOConnectorWorker) ++ worker._has_mamba = has_mamba ++ worker.moriio_wrapper = MoRIIOWrapper(moriio_engine=object(), transfer_timeout=0.05) ++ worker._reads_issued_this_step = [_PolledStatus(), _PolledStatus()] ++ ++ with pytest.raises(TransferError, match="timed out"): ++ worker._await_reads_issued_this_step() ++ ++ assert fake_clock.now > 0.05 ++ ++ ++@pytest.mark.skipif( ++ moriio_engine.StatusCode is None, reason="mori batched wait unavailable" ++) ++@pytest.mark.parametrize("has_mamba", [False, True]) ++def test_step_barrier_native_timeout_fails_closed_while_reads_in_flight(has_mamba): ++ engine = SimpleNamespace( ++ wait_all=Mock(return_value=moriio_engine.StatusCode.IN_PROGRESS) ++ ) ++ worker = MoRIIOConnectorWorker.__new__(MoRIIOConnectorWorker) ++ worker._has_mamba = has_mamba ++ worker.moriio_wrapper = MoRIIOWrapper(moriio_engine=engine, transfer_timeout=0.05) ++ worker._reads_issued_this_step = [_PolledStatus(), _PolledStatus()] ++ ++ with pytest.raises(TransferError, match="timed out"): ++ worker._await_reads_issued_this_step() ++ ++ engine.wait_all.assert_called_once() ++ ++ ++def test_step_barrier_noop_when_no_reads_posted(): ++ class RecordingWrapper: ++ def __init__(self): ++ self.called = False ++ ++ def waiting_for_transfer_complete(self, _statuses): ++ self.called = True ++ ++ def shutdown(self): ++ pass ++ ++ worker = MoRIIOConnectorWorker.__new__(MoRIIOConnectorWorker) ++ worker._has_mamba = False ++ worker._reads_issued_this_step = [] ++ worker.moriio_wrapper = RecordingWrapper() ++ ++ worker._await_reads_issued_this_step() ++ ++ assert worker.moriio_wrapper.called is False ++ ++ ++def test_write_mode_read_blocks_posts_no_read(): ++ class RecordingWrapper: ++ def __init__(self): ++ self.sent = [] ++ ++ def send_notify(self, *args, **kwargs): ++ self.sent.append((args, kwargs)) ++ ++ def shutdown(self): ++ pass ++ ++ worker = MoRIIOConnectorWorker.__new__(MoRIIOConnectorWorker) ++ worker.mode = MoRIIOMode.WRITE ++ worker._reads_issued_this_step = [] ++ worker.moriio_wrapper = RecordingWrapper() ++ ++ worker._read_blocks( ++ local_block_ids=[1, 2], ++ remote_block_ids=[[10, 11], [90]], ++ dst_engine_id="prefill", ++ request_id="req-write", ++ transfer_id="tx-write", ++ remote_host="127.0.0.1", ++ remote_notify_port=7000, ++ remote_tp_size=8, ++ ) ++ ++ assert worker._reads_issued_this_step == [] ++ assert worker.moriio_wrapper.sent == [] ++ ++ ++def test_producer_start_load_kv_posts_no_read(): ++ class RecordingWrapper: ++ def __init__(self): ++ self.waited = False ++ self.async_waited = False ++ ++ def async_wait_reqid(self): ++ self.async_waited = True ++ ++ def waiting_for_transfer_complete(self, _statuses): ++ self.waited = True ++ ++ def shutdown(self): ++ pass ++ ++ worker = MoRIIOConnectorWorker.__new__(MoRIIOConnectorWorker) ++ worker.is_producer = True ++ worker._reads_issued_this_step = [] ++ worker._consumer_notification_counts = {} ++ worker._completed_consumer_notifications = set() ++ worker.moriio_wrapper = RecordingWrapper() ++ ++ worker.start_load_kv(SimpleNamespace(transfer_id_to_request_id={})) ++ ++ assert worker._reads_issued_this_step == [] ++ assert worker.moriio_wrapper.async_waited is True ++ assert worker.moriio_wrapper.waited is False ++ ++ + def test_requested_cudagraph_mode_is_never_overridden(): + # The configured cudagraph mode is always honored: the barrier fires when + # the operator sets cudagraph_mode=PIECEWISE, and READ mode with full +diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py +index 4f19d73d9..d6039e09e 100644 +--- a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py ++++ b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py +@@ -17,7 +17,7 @@ import numpy as np + import torch + import zmq + +-from vllm.config import CUDAGraphMode, KVTransferConfig, VllmConfig ++from vllm.config import KVTransferConfig, VllmConfig + from vllm.distributed.kv_transfer.kv_connector.utils import ( + BlockIds, + clip_ssm_state_blocks, +@@ -97,7 +97,7 @@ from vllm.distributed.parallel_state import ( + get_tp_group, + get_world_group, + ) +-from vllm.forward_context import ForwardContext, get_forward_context ++from vllm.forward_context import ForwardContext + from vllm.logger import init_logger + from vllm.utils.math_utils import cdiv + from vllm.utils.network_utils import ( +@@ -416,8 +416,8 @@ class MoRIIOConnector(KVConnectorBase_V1, SupportsHMA): + self.connector_worker.start_load_kv(self._connector_metadata) + + def wait_for_layer_load(self, layer_name: str) -> None: +- assert self.connector_worker is not None +- self.connector_worker.wait_for_layer_load(layer_name) ++ """READs are awaited in start_load_kv, before the forward.""" ++ pass + + def save_kv_layer( + self, +@@ -833,8 +833,8 @@ class MoRIIOConnectorScheduler: + # + # The second element declares who waits for the KV, not whether it has + # arrived. False keeps the wait on this side: the forward runs in the +- # same step, and the connector blocks in wait_for_layer_load (or in +- # start_load_kv under full-graph capture) until the KV is really there. ++ # same step, and the connector blocks in start_load_kv, before the ++ # forward, until the KV is really there. + return max(num_external_tokens - 1, 0), False + + def on_new_request(self, request: "Request") -> None: +@@ -1907,13 +1907,8 @@ class MoRIIOConnectorWorker: + self._recving_local_blocks: dict[ReqId, list[int]] = {} + # Drained by get_block_ids_with_load_errors. + self._invalid_block_ids: set[int] = set() +- # This step's posted reads, awaited before the forward when the +- # per-layer barrier cannot run. ++ # This step's posted reads, awaited in start_load_kv before the forward. + self._reads_issued_this_step: list = [] +- # This step's posted recurrent-state reads. Mamba layers are not +- # wrapped by the attention KV-transfer decorator, so wait_for_layer_load +- # never fires for them and these must be awaited in start_load_kv. +- self._mamba_reads_this_step: list = [] + + # Track the expiration time of requests that are waiting to be sent. + self._reqs_to_send: dict[ReqId, float] = {} +@@ -2900,44 +2895,6 @@ class MoRIIOConnectorWorker: + failed_recving=failed_recving, + ) + +- def wait_for_layer_load(self, layer_name: str) -> None: +- """Block until all in-flight READs of this layer have landed. +- +- A host-side blocking wait must not run during full-graph capture. +- """ +- if self.is_producer or self.mode != MoRIIOMode.READ: +- return +- +- if get_forward_context().cudagraph_runtime_mode == CUDAGraphMode.FULL: +- return +- +- with self.moriio_wrapper.lock: +- pending = [ +- status +- for status_by_layer in self._recving_transfers.values() +- for status in status_by_layer.get(layer_name, ()) +- ] +- +- if not pending: +- return +- +- try: +- self.moriio_wrapper.waiting_for_transfer_complete(pending) +- except TransferError: +- if self._has_mamba: +- logger.exception( +- "MoRIIO hybrid READ failed before layer %s; aborting the " +- "forward because HMA load-error recovery is unsupported.", +- layer_name, +- ) +- raise +- logger.warning( +- "MoRIIO READ barrier did not complete for layer %s; proceeding " +- "to transfer cleanup.", +- layer_name, +- exc_info=True, +- ) +- + def _pop_done_transfers(self) -> set[str]: + done_req_ids: set[str] = set() + _xfer_timeout = self.moriio_config.recv_abort_timeout +@@ -3020,40 +2977,37 @@ class MoRIIOConnectorWorker: + return done_req_ids + + def _await_reads_issued_this_step(self) -> None: +- """Step-level READ barrier, for the reads no per-layer barrier covers. +- +- load_kv_async=False promises the KV is in place before the forward runs. +- wait_for_layer_load keeps that promise per layer, letting early layers +- overlap later transfers, but it only fires for layers that go through +- the attention KV-transfer decorator. Two cases fall through it: +- +- * mamba/KDA layers have no such hook at all, so their conv+ssm reads +- would otherwise be consumed while still in flight; +- * under CUDAGraphMode.FULL a host-side blocking wait is illegal +- anywhere in the graph, so every read is awaited here instead. +- +- start_load_kv runs outside the graph, so blocking here is safe. +- Attention reads keep their per-layer overlap in the non-FULL case. +- +- Attention-only failures are reported via +- get_block_ids_with_load_errors. Hybrid failures are re-raised because +- the scheduler's recovery path does not support multiple cache groups; +- this fails the step closed instead of running forward with incomplete +- recurrent state. ++ """Block until every READ posted this step has landed. ++ ++ load_kv_async=False promises the KV is in place before the forward. A ++ per-layer barrier can't keep it: illegal under FULL cudagraph, and absent ++ for mamba/KDA state and indexer caches. ++ ++ A READ still in flight when the wait gives up fails the step in both ++ modes: RDMA may still write into its blocks, so the forward must not ++ consume them. If every READ is terminal, attention-only failures are ++ reported via get_block_ids_with_load_errors. Hybrid failures are ++ re-raised because the scheduler's recovery path does not support ++ multiple cache groups. + """ +- all_statuses = self._reads_issued_this_step ++ statuses = self._reads_issued_this_step + self._reads_issued_this_step = [] +- mamba_statuses = self._mamba_reads_this_step +- self._mamba_reads_this_step = [] +- if get_forward_context().cudagraph_runtime_mode == CUDAGraphMode.FULL: +- statuses = all_statuses +- else: +- statuses = mamba_statuses + if not statuses: + return + try: +- self.moriio_wrapper.waiting_for_transfer_complete(statuses) ++ self.moriio_wrapper.waiting_for_transfer_complete( ++ statuses, fail_fast=self._has_mamba ++ ) + except TransferError: ++ in_flight = sum(not (s.Succeeded() or s.Failed()) for s in statuses) ++ if in_flight: ++ logger.exception( ++ "MoRIIO READ barrier failed with %d/%d reads still in " ++ "flight; failing the step.", ++ in_flight, ++ len(statuses), ++ ) ++ raise + if self._has_mamba: + logger.exception( + "MoRIIO hybrid READ failed before the forward; aborting " +@@ -4065,13 +4019,7 @@ class MoRIIOConnectorWorker: + self._recving_local_blocks.setdefault( + request_id, list(failed_recv_blocks) + ) +- # Awaited before the forward when the per-layer barrier cannot +- # run (see _await_reads_issued_this_step). + self._reads_issued_this_step.extend(statuses) +- if self._is_mamba_layer(layer_name): +- # No per-layer barrier exists for mamba layers, so these are +- # always awaited in start_load_kv. +- self._mamba_reads_this_step.extend(statuses) + self._recving_transfers_callback_addr[request_id] = ( + remote_host, + str( +diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_engine.py b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_engine.py +index 9929341b3..61e2c8d38 100644 +--- a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_engine.py ++++ b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_engine.py +@@ -893,7 +893,9 @@ class MoRIIOWrapper: + state = TransferBatchState.PENDING + return state + +- def waiting_for_transfer_complete(self, transfer_statuses: list[Any] | None = None): ++ def waiting_for_transfer_complete( ++ self, transfer_statuses: list[Any] | None = None, fail_fast: bool = True ++ ): + if transfer_statuses is None: + with self.lock: + transfers_to_wait = self.transfer_status[:] +@@ -906,9 +908,11 @@ class MoRIIOWrapper: + + wait_all = getattr(self.moriio_engine, "wait_all", None) + if wait_all is None or StatusCode is None: +- self._poll_transfers_until_done(transfers_to_wait) ++ self._poll_transfers_until_done(transfers_to_wait, fail_fast) + return + ++ # wait_all waits on every status until the deadline; a failure does not ++ # end it early, so fail_fast only applies to the polling fallback. + # Zero means a nonblocking progress call in Mori, not an immediate + # deadline. Round up so even sub-millisecond waits remain blocking. + timeout_ms = ceil(self._transfer_timeout * 1000) +@@ -927,12 +931,15 @@ class MoRIIOWrapper: + f"MoRIIO batch wait failed ({result}): " + "; ".join(details) + ) + +- def _poll_transfers_until_done(self, transfers_to_wait: list[Any]) -> None: ++ def _poll_transfers_until_done( ++ self, transfers_to_wait: list[Any], fail_fast: bool = True ++ ) -> None: + """Fallback for mori builds without the batched wait. + +- Gives failure precedence over pending: once any status has failed the +- request is lost, so waiting out the remaining ones would only delay the +- barrier by up to transfer_timeout. ++ With fail_fast, failure takes precedence over pending: once any status ++ has failed the request is lost, so waiting out the remaining ones would ++ only delay the barrier by up to transfer_timeout. Pass False when the ++ statuses span several requests. + """ + timeout = self._transfer_timeout + deadline = time.monotonic() + timeout +@@ -959,7 +966,7 @@ class MoRIIOWrapper: + ) + else: + still_waiting.append(status) +- if errors: ++ if errors and fail_fast: + break + remaining = still_waiting + if remaining: diff --git a/inferencex-e2e/benchmarks/multi_node/srt-slurm-recipes/configs/patches/kimik3-nightly/08-vllm-63031.patch b/inferencex-e2e/benchmarks/multi_node/srt-slurm-recipes/configs/patches/kimik3-nightly/08-vllm-63031.patch new file mode 100644 index 0000000000..3a69f6047c --- /dev/null +++ b/inferencex-e2e/benchmarks/multi_node/srt-slurm-recipes/configs/patches/kimik3-nightly/08-vllm-63031.patch @@ -0,0 +1,172 @@ +From fcd71bb9a36adb63b9ce9f43c7e4170ffadb595c Mon Sep 17 00:00:00 2001 +From: Yichao Zhu +Date: Thu, 22 Oct 2026 12:08:07 +0800 +Subject: [PATCH] [ROCm][Kimi-K3] Support MoRIIO consumers with specialized + decode + +Allow the built-in MoRIIO consumer in READ and WRITE modes. Reuse the existing transfer completion guarantees, and keep other connectors and producer-side specialized execution unsupported. Add configuration coverage and document the supported deployment. + +Signed-off-by: Yichao Zhu +(cherry picked from commit e34bf5f4784569ca3c2381312f2e7e4d512434de) +--- + docs/features/moriio_connector_usage.md | 8 ++++ + .../amd_specialized_decode/test_runtime.py | 46 ++++++++++++++++++- + .../kimi_k3/amd/specialized_decode/README.md | 9 +++- + .../specialized_decode/runtime/__init__.py | 16 ++++++- + 4 files changed, 75 insertions(+), 4 deletions(-) + +diff --git a/docs/features/moriio_connector_usage.md b/docs/features/moriio_connector_usage.md +index 2fe1ffa63..7da33d6a4 100644 +--- a/docs/features/moriio_connector_usage.md ++++ b/docs/features/moriio_connector_usage.md +@@ -144,6 +144,14 @@ For DSpark, configure the same draft model on both instances through `--speculat + + Hybrid READ keeps its existing single-transferable-attention-group requirement. Plain-attention WRITE retains its existing notification format. + ++### Kimi-K3 specialized decode ++ ++Kimi-K3 on MI355X supports `--additional-config '{"decode_backend":"latency_specialized"}'` on a MoRIIO READ or WRITE consumer. Use `VLLM_USE_V2_MODEL_RUNNER=1`, TP8, DCP1 or DCP8, DSpark with up to seven speculative tokens, and BF16 convolution / FP32 recurrent state on both instances. Keep the producer on the default decode backend and use the same transfer mode and cache layouts on both instances. ++ ++The specialized kernels serve pure speculative verification steps of at most eight token rows. For example, with seven speculative tokens and `--max-num-seqs 1`, a verification step has eight rows; use `--compilation-config '{"cudagraph_mode":"FULL_DECODE_ONLY","cudagraph_capture_sizes":[1,8]}'` on the consumer. Prefill, mixed, and wider steps use the regular model layers. All transferred attention and recurrent states are ready before the forward, so FULL graph replay does not need per-layer transfer hooks. ++ ++Producers using the specialized backend and other connectors remain unsupported. ++ + ### Different prefill and decode TP sizes + + In WRITE mode, the prefill and decode instances can use different tensor-parallel sizes, as long as one is a multiple of the other. Each prefill rank writes its shard into the decode ranks that hold it: +diff --git a/tests/models/kimi_k3/amd_specialized_decode/test_runtime.py b/tests/models/kimi_k3/amd_specialized_decode/test_runtime.py +index 6cc5c793c..5fde9359c 100644 +--- a/tests/models/kimi_k3/amd_specialized_decode/test_runtime.py ++++ b/tests/models/kimi_k3/amd_specialized_decode/test_runtime.py +@@ -8,7 +8,7 @@ from types import SimpleNamespace + import pytest + import torch + +-from vllm.config import CompilationMode, CUDAGraphMode ++from vllm.config import CompilationMode, CUDAGraphMode, KVTransferConfig + from vllm.models.kimi_k3.amd.specialized_decode import runtime + from vllm.models.kimi_k3.amd.specialized_decode.runtime import ( + FFN, +@@ -141,6 +141,41 @@ def test_supported_config_is_accepted(on_mi355x): + assert specialized_decode_config(_vllm_config()) == DecodeConfig(q=8) + + ++@pytest.mark.parametrize("read_mode", [False, True]) ++@pytest.mark.parametrize("backend", ["rdma", "xgmi"]) ++def test_moriio_consumer_is_supported(read_mode, backend): ++ config = KVTransferConfig( ++ kv_connector="MoRIIOConnector", ++ kv_role="kv_consumer", ++ kv_connector_extra_config={"read_mode": read_mode, "backend": backend}, ++ ) ++ assert runtime.supports_kv_transfer(config) ++ ++ ++@pytest.mark.parametrize( ++ "overrides", ++ [ ++ {"kv_connector": "NixlConnector"}, ++ {"kv_connector": "MultiConnector"}, ++ {"kv_connector_module_path": "custom.connector"}, ++ {"kv_role": "kv_producer"}, ++ {"kv_role": "kv_both"}, ++ ], ++) ++def test_unsupported_kv_transfer_is_refused(overrides): ++ config = KVTransferConfig( ++ **( ++ { ++ "kv_connector": "MoRIIOConnector", ++ "kv_role": "kv_consumer", ++ "kv_connector_extra_config": {"read_mode": True}, ++ } ++ | overrides ++ ) ++ ) ++ assert not runtime.supports_kv_transfer(config) ++ ++ + @pytest.mark.parametrize( + "over, reason", + [ +@@ -153,7 +188,14 @@ def test_supported_config_is_accepted(on_mi355x): + ), + ({"parallel_config.pipeline_parallel_size": 2}, "pipeline parallelism"), + ({"parallel_config.decode_context_parallel_size": 2}, "decode context"), +- ({"kv_transfer_config": object()}, "KV transfer connectors"), ++ ( ++ { ++ "kv_transfer_config": KVTransferConfig( ++ kv_connector="NixlConnector", kv_role="kv_consumer" ++ ) ++ }, ++ "KV transfer connectors", ++ ), + ({"speculative_config": None}, "without speculative decoding"), + ({"speculative_config.num_speculative_tokens": 8}, "more than 7"), + ({"speculative_config.enable_adaptive_verification": True}, "adaptive"), +diff --git a/vllm/models/kimi_k3/amd/specialized_decode/README.md b/vllm/models/kimi_k3/amd/specialized_decode/README.md +index 3e0e048b8..602fb2613 100644 +--- a/vllm/models/kimi_k3/amd/specialized_decode/README.md ++++ b/vllm/models/kimi_k3/amd/specialized_decode/README.md +@@ -31,11 +31,18 @@ TP8 on MI355X with all 256 CUs, decode context parallel 1 or 8, speculative + decoding with at most 7 speculative tokens, Model Runner V2, BF16 conv / FP32 + recurrent KDA states. + ++MoRIIO READ and WRITE consumers are supported. READ finishes all attention and ++recurrent-state loads before the forward; WRITE requests wait for transfer ++completion before admission. Neither requires per-layer transfer hooks during ++FULL graph replay. Keep the producer on the stock decode backend and use ++matching cache layouts. Other KV connector configurations remain unsupported. ++ + If you request the specialized decode for anything else, startup fails with one + `ValueError` that lists every unsupported setting: + + - **Config, in the model's `__init__`:** pipeline, data, expert or prefill +- context parallelism, EPLB, DBO, LoRA, KV connectors, routed-expert capture, ++ context parallelism, EPLB, DBO, LoRA, unsupported KV connectors, ++ routed-expert capture, + sleep mode, CPU weight offload, weight updates, no speculative decoding, + adaptive verification, dynamic speculative decoding, ReplaySSM / RecoverSSM, + the V1 model runner, `STOCK_TORCH_COMPILE`, another model's shapes, another GPU. +diff --git a/vllm/models/kimi_k3/amd/specialized_decode/runtime/__init__.py b/vllm/models/kimi_k3/amd/specialized_decode/runtime/__init__.py +index e17a0e3ff..4dd549911 100644 +--- a/vllm/models/kimi_k3/amd/specialized_decode/runtime/__init__.py ++++ b/vllm/models/kimi_k3/amd/specialized_decode/runtime/__init__.py +@@ -13,6 +13,7 @@ from dataclasses import dataclass + + import torch + ++from vllm.config import KVTransferConfig + from vllm.logger import init_logger + + logger = init_logger(__name__) +@@ -103,7 +104,10 @@ def unsupported(vllm_config) -> list[str]: + # the kernels read one batch's attention metadata per step + (pc.use_ubatching, "DBO (dual-batch overlap)"), + (vllm_config.lora_config is not None, "LoRA"), +- (vllm_config.kv_transfer_config is not None, "KV transfer connectors"), ++ ( ++ not supports_kv_transfer(vllm_config.kv_transfer_config), ++ "KV transfer connectors other than a MoRIIO consumer", ++ ), + ( + vllm_config.aux_output_config.enable_return_routed_experts, + "routed-expert capture", +@@ -162,6 +166,16 @@ def unsupported(vllm_config) -> list[str]: + return reasons + platform_reasons() + + ++def supports_kv_transfer(config: KVTransferConfig | None) -> bool: ++ # MoRIIO completes READs before forward and WRITE receives before admission; ++ # neither consumer needs the per-layer hooks specialized layers bypass. ++ return config is None or ( ++ config.kv_connector == "MoRIIOConnector" ++ and config.kv_connector_module_path is None ++ and config.kv_role == "kv_consumer" ++ ) ++ ++ + def text_shapes(text) -> tuple: + kda = getattr(text, "linear_attn_config", None) or {} + return ( diff --git a/inferencex-e2e/benchmarks/multi_node/srt-slurm-recipes/kimik3/vllm/mi355x-fp4/agentx/disagg-variants.yaml b/inferencex-e2e/benchmarks/multi_node/srt-slurm-recipes/kimik3/vllm/mi355x-fp4/agentx/disagg-variants.yaml index c024b6879e..3683082660 100644 --- a/inferencex-e2e/benchmarks/multi_node/srt-slurm-recipes/kimik3/vllm/mi355x-fp4/agentx/disagg-variants.yaml +++ b/inferencex-e2e/benchmarks/multi_node/srt-slurm-recipes/kimik3/vllm/mi355x-fp4/agentx/disagg-variants.yaml @@ -1,6 +1,9 @@ base: schema: 2 name: kimik3-fp4-mi355x-vllm-disagg-agentic + setup_script: kimik3-nightly-kernels.sh + environment: + KIMIK3_KERNEL_BUILD_JOBS: '12' slurm: time_limit: '08:00:00' sbatch_directives: diff --git a/inferencex-e2e/perf-changelog.yaml b/inferencex-e2e/perf-changelog.yaml index fa3b27dcd1..f1b68aa05d 100644 --- a/inferencex-e2e/perf-changelog.yaml +++ b/inferencex-e2e/perf-changelog.yaml @@ -11106,4 +11106,5 @@ description: - "Use the ROCm nightly image and BF16 SSM states for the throughput points." - "Replace 1P1D c56 and 1P3D c12 with 1P1D c1, K7 specialized decode, and FP32 SSM states." - pr-link: https://github.com/SemiAnalysisAI/InferenceX/pull/XXX + - "Apply vLLM #60683, #58743, #60808, and #63031 from their source commits." + pr-link: https://github.com/SemiAnalysisAI/InferenceX/pull/4391 From 0902c816684160ae1a8ebac98fd5f8cda249b4ee Mon Sep 17 00:00:00 2001 From: Yichao Zhu Date: Sat, 24 Oct 2026 22:30:48 +0800 Subject: [PATCH 3/6] build(kimik3): use upstream nightly without kernel overlays MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 官方 nightly 已包含四个依赖 PR，移除临时补丁、安装器及构建入口，仅保留 recipe 配置更新。 --- .../configs/kimik3-nightly-kernels.sh | 73 - .../kimik3-nightly/01-vllm-60683.patch | 6667 ----------------- .../kimik3-nightly/02-vllm-58743-state.patch | 155 - .../03-vllm-58743-checkpoints.patch | 251 - .../04-vllm-58743-native-prefill.patch | 512 -- .../05-vllm-58743-backends.patch | 181 - .../kimik3-nightly/06-vllm-58743-tests.patch | 115 - .../kimik3-nightly/07-vllm-60808.patch | 569 -- .../kimik3-nightly/08-vllm-63031.patch | 172 - .../mi355x-fp4/agentx/disagg-variants.yaml | 3 - inferencex-e2e/perf-changelog.yaml | 1 - 11 files changed, 8699 deletions(-) delete mode 100644 inferencex-e2e/benchmarks/multi_node/srt-slurm-recipes/configs/kimik3-nightly-kernels.sh delete mode 100644 inferencex-e2e/benchmarks/multi_node/srt-slurm-recipes/configs/patches/kimik3-nightly/01-vllm-60683.patch delete mode 100644 inferencex-e2e/benchmarks/multi_node/srt-slurm-recipes/configs/patches/kimik3-nightly/02-vllm-58743-state.patch delete mode 100644 inferencex-e2e/benchmarks/multi_node/srt-slurm-recipes/configs/patches/kimik3-nightly/03-vllm-58743-checkpoints.patch delete mode 100644 inferencex-e2e/benchmarks/multi_node/srt-slurm-recipes/configs/patches/kimik3-nightly/04-vllm-58743-native-prefill.patch delete mode 100644 inferencex-e2e/benchmarks/multi_node/srt-slurm-recipes/configs/patches/kimik3-nightly/05-vllm-58743-backends.patch delete mode 100644 inferencex-e2e/benchmarks/multi_node/srt-slurm-recipes/configs/patches/kimik3-nightly/06-vllm-58743-tests.patch delete mode 100644 inferencex-e2e/benchmarks/multi_node/srt-slurm-recipes/configs/patches/kimik3-nightly/07-vllm-60808.patch delete mode 100644 inferencex-e2e/benchmarks/multi_node/srt-slurm-recipes/configs/patches/kimik3-nightly/08-vllm-63031.patch diff --git a/inferencex-e2e/benchmarks/multi_node/srt-slurm-recipes/configs/kimik3-nightly-kernels.sh b/inferencex-e2e/benchmarks/multi_node/srt-slurm-recipes/configs/kimik3-nightly-kernels.sh deleted file mode 100644 index 9049ace8b4..0000000000 --- a/inferencex-e2e/benchmarks/multi_node/srt-slurm-recipes/configs/kimik3-nightly-kernels.sh +++ /dev/null @@ -1,73 +0,0 @@ -#!/usr/bin/env bash -set -eo pipefail - -# The router uses a separate image without vLLM. -if ! python3 -c 'import importlib.metadata; importlib.metadata.version("vllm")' >/dev/null 2>&1; then - exit 0 -fi -source /infmax-workspace/benchmarks/check_env.sh -check_env_vars KIMIK3_KERNEL_BUILD_JOBS ROCM_PATH -[[ "$KIMIK3_KERNEL_BUILD_JOBS" =~ ^[1-9][0-9]*$ ]] - -package=$(python3 -c 'import importlib.util; print(next(iter(importlib.util.find_spec("vllm").submodule_search_locations)))') -version=$(python3 -c 'import importlib.metadata; print(importlib.metadata.version("vllm"))') -revision=$(python3 -c 'import re, sys; match = re.search(r"\+g([0-9a-f]{7,40})(?:\.|$)", sys.argv[1]); assert match, f"Cannot resolve nightly source: {sys.argv[1]}"; print(match[1])' "$version") -patch_dir=$(cd "$(dirname "$0")/patches/kimik3-nightly" && pwd) -patches=("$patch_dir"/*.patch) -[[ -f "${patches[0]}" ]] -signature=$({ printf '%s\n' "$version"; sha256sum "${patches[@]}"; } | sha256sum | cut -d' ' -f1) -marker="$package/.infx-kimik3-kernels" -exec 9> /tmp/infx-kimik3-kernels.lock -flock 9 -if [[ -f "$marker" && $(< "$marker") == "$signature" ]]; then - echo "Kimi-K3 kernel patches already installed: $signature" - exit 0 -fi - -build_root=$(mktemp -d /tmp/infx-kimik3-kernels.XXXXXX) -curl --fail --location --retry 3 --max-time 300 \ - "https://codeload.github.com/vllm-project/vllm/tar.gz/$revision" \ - -o "$build_root/source.tgz" -mkdir "$build_root/source" -tar -xzf "$build_root/source.tgz" --strip-components=1 -C "$build_root/source" -git -C "$build_root/source" init -q -applied=() -for patch in "${patches[@]}"; do - if git -C "$build_root/source" apply --reverse --check "$patch" 2>/dev/null; then - echo "Already included in $version: $(basename "$patch")" - else - git -C "$build_root/source" apply --check "$patch" - git -C "$build_root/source" apply "$patch" - applied+=("$patch") - fi -done - -if (( ${#applied[@]} )); then - cp "$package"/*.so "$build_root/source/vllm/" - cp "$package/_version.py" "$build_root/source/vllm/_version.py" - uv venv --system-site-packages --python python3 "$build_root/venv" - uv pip install --python "$build_root/venv/bin/python" cmake ninja - export PATH="$build_root/venv/bin:$PATH" - cmake -S "$build_root/source" -B "$build_root/build" -G Ninja \ - -DVLLM_TARGET_DEVICE=rocm \ - -DVLLM_PYTHON_EXECUTABLE="$build_root/venv/bin/python" \ - -DCMAKE_BUILD_TYPE=RelWithDebInfo \ - -DCMAKE_MODULE_LINKER_FLAGS=-fuse-ld=lld \ - -DCMAKE_SHARED_LINKER_FLAGS=-fuse-ld=lld \ - -DCMAKE_INSTALL_PREFIX="$build_root/stage" \ - -DVLLM_GPU_ARCHES=gfx950 -DCMAKE_HIP_ARCHITECTURES=gfx950 \ - -DROCM_PATH="$ROCM_PATH" - cmake --build "$build_root/build" --target _C_stable_libtorch -j "$KIMIK3_KERNEL_BUILD_JOBS" - cmake --install "$build_root/build" --component _C_stable_libtorch - install -m 755 "$build_root/stage/vllm/_C_stable_libtorch.abi3.so" "$package/" - for patch in "${applied[@]}"; do - while IFS=$'\t' read -r additions deletions path; do - if [[ "$path" == vllm/*.py ]]; then - install -D -m 644 "$build_root/source/$path" "$package/${path#vllm/}" - fi - done < <(git apply --numstat "$patch") - done - python3 -c 'import torch; import vllm._C_stable_libtorch' -fi -printf '%s\n' "$signature" > "$marker" -echo "Kimi-K3 kernel patches installed on $version; build evidence: $build_root" diff --git a/inferencex-e2e/benchmarks/multi_node/srt-slurm-recipes/configs/patches/kimik3-nightly/01-vllm-60683.patch b/inferencex-e2e/benchmarks/multi_node/srt-slurm-recipes/configs/patches/kimik3-nightly/01-vllm-60683.patch deleted file mode 100644 index 9f72da7fb6..0000000000 --- a/inferencex-e2e/benchmarks/multi_node/srt-slurm-recipes/configs/patches/kimik3-nightly/01-vllm-60683.patch +++ /dev/null @@ -1,6667 +0,0 @@ -From 59781f807cc71208ceac80d8eb43e264dbe35f3f Mon Sep 17 00:00:00 2001 -From: Hyukjoon Lee -Date: Wed, 21 Oct 2026 15:53:29 +0900 -Subject: [PATCH] [ROCm][Perf][Kimi-K3] Specialized decode for Kimi-K3 on - MI355X (#60683) - -Signed-off-by: hyukjlee -Co-authored-by: Claude Opus 5.5 -Co-authored-by: vllm-agent -(cherry picked from commit a9bf45ed8a856a0b1839197152d0d26fabb6599c) ---- - .buildkite/test_areas/models_basic.yaml | 15 + - .../amd_specialized_decode/__init__.py | 2 + - .../test_engine_fault.py | 153 +++ - .../amd_specialized_decode/test_numerics.py | 575 +++++++++ - .../amd_specialized_decode/test_runtime.py | 318 +++++ - vllm/models/kimi_k3/amd/linear.py | 27 +- - vllm/models/kimi_k3/amd/mla.py | 3 + - vllm/models/kimi_k3/amd/model.py | 3 + - .../kimi_k3/amd/specialized_decode/README.md | 108 ++ - .../amd/specialized_decode/__init__.py | 12 + - .../kimi_k3/amd/specialized_decode/hooks.py | 63 + - .../amd/specialized_decode/model_state.py | 27 + - .../amd/specialized_decode/mono/__init__.py | 2 + - .../amd/specialized_decode/mono/abi.py | 26 + - .../mono/attention/__init__.py | 2 + - .../specialized_decode/mono/attention/back.py | 178 +++ - .../specialized_decode/mono/attention/kda.py | 1021 ++++++++++++++++ - .../mono/common/__init__.py | 2 + - .../specialized_decode/mono/common/debug.py | 38 + - .../amd/specialized_decode/mono/common/mx.py | 84 ++ - .../amd/specialized_decode/mono/common/ops.py | 435 +++++++ - .../specialized_decode/mono/common/plan.py | 237 ++++ - .../specialized_decode/mono/common/ranks.py | 33 + - .../specialized_decode/mono/common/sync.py | 255 ++++ - .../amd/specialized_decode/mono/layer.py | 573 +++++++++ - .../mono/stages/__init__.py | 2 + - .../specialized_decode/mono/stages/gemv.py | 126 ++ - .../amd/specialized_decode/mono/stages/moe.py | 1025 +++++++++++++++++ - .../specialized_decode/runtime/__init__.py | 219 ++++ - .../amd/specialized_decode/runtime/decode.py | 198 ++++ - .../amd/specialized_decode/runtime/guards.py | 154 +++ - .../amd/specialized_decode/runtime/layers.py | 221 ++++ - .../amd/specialized_decode/runtime/peer.py | 112 ++ - .../amd/specialized_decode/runtime/runner.py | 76 ++ - 34 files changed, 6318 insertions(+), 7 deletions(-) - create mode 100644 tests/models/kimi_k3/amd_specialized_decode/__init__.py - create mode 100644 tests/models/kimi_k3/amd_specialized_decode/test_engine_fault.py - create mode 100644 tests/models/kimi_k3/amd_specialized_decode/test_numerics.py - create mode 100644 tests/models/kimi_k3/amd_specialized_decode/test_runtime.py - create mode 100644 vllm/models/kimi_k3/amd/specialized_decode/README.md - create mode 100644 vllm/models/kimi_k3/amd/specialized_decode/__init__.py - create mode 100644 vllm/models/kimi_k3/amd/specialized_decode/hooks.py - create mode 100644 vllm/models/kimi_k3/amd/specialized_decode/model_state.py - create mode 100644 vllm/models/kimi_k3/amd/specialized_decode/mono/__init__.py - create mode 100644 vllm/models/kimi_k3/amd/specialized_decode/mono/abi.py - create mode 100644 vllm/models/kimi_k3/amd/specialized_decode/mono/attention/__init__.py - create mode 100644 vllm/models/kimi_k3/amd/specialized_decode/mono/attention/back.py - create mode 100644 vllm/models/kimi_k3/amd/specialized_decode/mono/attention/kda.py - create mode 100644 vllm/models/kimi_k3/amd/specialized_decode/mono/common/__init__.py - create mode 100644 vllm/models/kimi_k3/amd/specialized_decode/mono/common/debug.py - create mode 100644 vllm/models/kimi_k3/amd/specialized_decode/mono/common/mx.py - create mode 100644 vllm/models/kimi_k3/amd/specialized_decode/mono/common/ops.py - create mode 100644 vllm/models/kimi_k3/amd/specialized_decode/mono/common/plan.py - create mode 100644 vllm/models/kimi_k3/amd/specialized_decode/mono/common/ranks.py - create mode 100644 vllm/models/kimi_k3/amd/specialized_decode/mono/common/sync.py - create mode 100644 vllm/models/kimi_k3/amd/specialized_decode/mono/layer.py - create mode 100644 vllm/models/kimi_k3/amd/specialized_decode/mono/stages/__init__.py - create mode 100644 vllm/models/kimi_k3/amd/specialized_decode/mono/stages/gemv.py - create mode 100644 vllm/models/kimi_k3/amd/specialized_decode/mono/stages/moe.py - create mode 100644 vllm/models/kimi_k3/amd/specialized_decode/runtime/__init__.py - create mode 100644 vllm/models/kimi_k3/amd/specialized_decode/runtime/decode.py - create mode 100644 vllm/models/kimi_k3/amd/specialized_decode/runtime/guards.py - create mode 100644 vllm/models/kimi_k3/amd/specialized_decode/runtime/layers.py - create mode 100644 vllm/models/kimi_k3/amd/specialized_decode/runtime/peer.py - create mode 100644 vllm/models/kimi_k3/amd/specialized_decode/runtime/runner.py - -diff --git a/.buildkite/test_areas/models_basic.yaml b/.buildkite/test_areas/models_basic.yaml -index cfcc351a5..80c95e95b 100644 ---- a/.buildkite/test_areas/models_basic.yaml -+++ b/.buildkite/test_areas/models_basic.yaml -@@ -324,3 +324,18 @@ steps: - commands: - - pytest -v -s tests/models/test_dsv41_mono_numerics.py - - pytest -v -s tests/models/test_dsv41_mono_layers.py -+ -+- label: ':amd: (MI355) Kimi-K3 Specialized Decode' -+ key: kimi-k3-specialized-decode-amd -+ dind: false -+ device: mi355_8 -+ num_devices: 8 -+ timeout_in_minutes: 60 -+ working_dir: /vllm-workspace/ -+ depends_on: -+ - image-build-amd -+ source_file_dependencies: -+ - vllm/models/kimi_k3/amd/ -+ - tests/models/kimi_k3/amd_specialized_decode/ -+ commands: -+ - pytest -v -s tests/models/kimi_k3/amd_specialized_decode -diff --git a/tests/models/kimi_k3/amd_specialized_decode/__init__.py b/tests/models/kimi_k3/amd_specialized_decode/__init__.py -new file mode 100644 -index 000000000..208f01a7c ---- /dev/null -+++ b/tests/models/kimi_k3/amd_specialized_decode/__init__.py -@@ -0,0 +1,2 @@ -+# SPDX-License-Identifier: Apache-2.0 -+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -diff --git a/tests/models/kimi_k3/amd_specialized_decode/test_engine_fault.py b/tests/models/kimi_k3/amd_specialized_decode/test_engine_fault.py -new file mode 100644 -index 000000000..485dba75f ---- /dev/null -+++ b/tests/models/kimi_k3/amd_specialized_decode/test_engine_fault.py -@@ -0,0 +1,153 @@ -+# SPDX-License-Identifier: Apache-2.0 -+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -+"""Fail-stop end to end: a TP rank that stalls inside a served Kimi-K3 -+specialized decode step fails a request, and the engine stops instead of -+hanging: every TP rank must stop at the same step (``guards.FailStop``), or the -+ranks that went on would wait in the next collective for ever.""" -+ -+import json -+import os -+import time -+ -+import pytest -+import requests -+ -+from tests.models.kimi_k3.amd_specialized_decode.test_numerics import needs_gfx950 -+from tests.utils import RemoteOpenAIServer, multi_gpu_test -+ -+MODEL = "riverclouds/Kimi-K3-8L-dummy" -+DRAFT = "riverclouds/Kimi-K3-DSpark-for-8L-dummy" -+STALLED = 3 -+FAULT_EPOCH = 256 # a mono decode step well past warm-up and capture -+STARTUP_TIMEOUT = 1800 -+FAULT_DEADLINE = 900 -+EXIT_TIMEOUT = 300 -+ -+# Builds every rank's K2 with the fault hook (``K2Build.fault_*``): CTA 0 of -+# rank STALLED stalls past every in-kernel wait's bound in the step at -+# FAULT_EPOCH. Patched into ``MonoRunner.fault`` as the runner module loads, in -+# every vLLM process. -+_SITECUSTOMIZE = """\ -+import importlib.abc -+import importlib.machinery -+import sys -+ -+_MODULE = "vllm.models.kimi_k3.amd.specialized_decode.runtime.runner" -+ -+ -+def _patch(module): -+ def fault(runner): -+ return {{ -+ "fault_rank": {rank}, -+ "fault_epoch": {epoch}, -+ "fault_layer": 128 + runner.first_layer, # the step's first K2 -+ "fault_ticks": {ticks}, -+ }} -+ -+ module.MonoRunner.fault = property(fault, lambda runner, value: None) -+ -+ -+class _Finder(importlib.abc.MetaPathFinder): -+ def find_spec(self, name, path, target=None): -+ if name != _MODULE: -+ return None -+ spec = importlib.machinery.PathFinder.find_spec(name, path) -+ exec_module = spec.loader.exec_module -+ -+ def patched_exec(module): -+ exec_module(module) -+ _patch(module) -+ -+ spec.loader.exec_module = patched_exec -+ return spec -+ -+ -+sys.meta_path.insert(0, _Finder()) -+""" -+ -+SERVER_ARGS = [ -+ "--trust-remote-code", -+ "--language-model-only", -+ "--load-format", -+ "dummy", -+ "--tensor-parallel-size", -+ "8", -+ "--decode-context-parallel-size", -+ "8", -+ "--dcp-comm-backend", -+ "a2a", -+ "--kv-cache-dtype", -+ "fp8", -+ "--max-model-len", -+ "8192", -+ "--max-num-seqs", -+ "2", -+ "--attention-config", -+ json.dumps({"mla_prefill_backend": "ROCM_AITER_FA"}), -+ "--speculative-config", -+ json.dumps( -+ { -+ "method": "dspark", -+ "model": DRAFT, -+ "num_speculative_tokens": 7, -+ "attention_backend": "ROCM_AITER_MLA", -+ "draft_sample_method": "probabilistic", -+ "rejection_sample_method": "block", -+ } -+ ), -+ "--compilation-config", -+ json.dumps({"cudagraph_mode": "FULL_AND_PIECEWISE"}), -+ "--additional-config", -+ json.dumps({"decode_backend": "latency_specialized"}), -+] -+ -+ -+@needs_gfx950 -+@multi_gpu_test(num_gpus=8) -+def test_a_stalled_rank_stops_the_engine(tmp_path): -+ from vllm.models.kimi_k3.amd.specialized_decode.mono.common.sync import WAIT_TICKS -+ -+ site = tmp_path / "mono_fault" -+ site.mkdir() -+ (site / "sitecustomize.py").write_text( -+ _SITECUSTOMIZE.format( -+ rank=STALLED, epoch=FAULT_EPOCH, ticks=WAIT_TICKS * 3 // 2 -+ ) -+ ) -+ env = { -+ "PYTHONPATH": os.pathsep.join( -+ filter(None, [str(site), os.getenv("PYTHONPATH")]) -+ ), -+ "VLLM_USE_V2_MODEL_RUNNER": "1", -+ "VLLM_ROCM_USE_AITER": "1", -+ "VLLM_ROCM_AITER_MLA_ASM_PADDING": "asm", -+ } -+ body = { -+ "model": MODEL, -+ "prompt": "hello " * 64, -+ "max_tokens": 128, -+ "temperature": 0, -+ "ignore_eos": True, -+ } -+ with RemoteOpenAIServer( -+ MODEL, SERVER_ARGS, env_dict=env, max_wait_seconds=STARTUP_TIMEOUT -+ ) as server: -+ url = server.url_root -+ served = 0 -+ deadline = time.monotonic() + FAULT_DEADLINE -+ while time.monotonic() < deadline: -+ try: -+ r = requests.post(f"{url}/v1/completions", json=body, timeout=300) -+ except requests.ConnectionError: -+ break -+ if r.status_code != 200: -+ break -+ served += 1 -+ else: -+ # the stall only exists inside the kernels: also catches a decode -+ # that never ran them -+ pytest.fail(f"no request failed in {FAULT_DEADLINE}s: the fault is lost") -+ -+ assert served > 0, "the engine failed before the injected step" -+ # the failed step stops the engine and the server exits -+ assert server.proc.wait(timeout=EXIT_TIMEOUT) is not None -diff --git a/tests/models/kimi_k3/amd_specialized_decode/test_numerics.py b/tests/models/kimi_k3/amd_specialized_decode/test_numerics.py -new file mode 100644 -index 000000000..dc396b099 ---- /dev/null -+++ b/tests/models/kimi_k3/amd_specialized_decode/test_numerics.py -@@ -0,0 +1,575 @@ -+# SPDX-License-Identifier: Apache-2.0 -+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -+"""Kimi-K3's MonoKernels on MI355X against the ops vLLM runs for the same -+spec-verify step, random weights in the checkpoint's formats at TP8 shapes: -+ -+- K1 (one GPU): AttnRes -> in_proj -> f_b -> conv -> KDA recurrence -> gated -+ norm, the conv / SSM state updates and a block-write layer's block store. -+- K2 (eight GPUs): o_proj + all-reduce + the MLP AttnRes + the latent MoE. The -+ MoE is compared on the kernel's own MLP input, where routing cannot differ. -+ Then a CUDA graph of K2 replayed 300 times, bit-identical, and a stalled rank -+ (the fault hook): every rank's launch finishes and the output rank's flags -+ fail the step. -+""" -+ -+import pytest -+import ray -+import torch -+import torch.distributed as dist -+import torch.nn.functional as F -+from einops import rearrange -+ -+from tests.utils import ( -+ init_test_distributed_environment, -+ multi_gpu_test, -+ multi_process_parallel, -+) -+from vllm.platforms import current_platform -+ -+bf = torch.bfloat16 -+COS_MIN = 0.9999 -+ -+ -+def _on_gfx950() -> bool: -+ if not current_platform.is_rocm(): -+ return False -+ from vllm.platforms.rocm import on_gfx950 -+ -+ return on_gfx950() -+ -+ -+needs_gfx950 = pytest.mark.skipif( -+ not _on_gfx950(), reason="the K3 mono kernels need gfx950 (MI355X)" -+) -+ -+ -+def _cos(a: torch.Tensor, b: torch.Tensor) -> float: -+ a, b = a.float().reshape(-1), b.float().reshape(-1) -+ return F.cosine_similarity(a, b, dim=0).item() -+ -+ -+def _rel_l2(a: torch.Tensor, b: torch.Tensor) -> float: -+ a, b = a.float(), b.float() -+ return ((a - b).norm() / b.norm()).item() -+ -+ -+def _randn(g, *shape, scale=1.0, dtype=bf): -+ return (torch.randn(*shape, generator=g, device=g.device) * scale).to(dtype) -+ -+ -+# ---------------------------------------------------------------- K1 -+ -+ -+def _kda_inputs(nb: int, nacc: int, dim_first: bool, S: int, slots: int = 16): -+ from vllm.models.kimi_k3.amd.specialized_decode.mono.attention.kda import ( -+ HD, -+ HIDDEN, -+ NH, -+ NPROJ, -+ PROJ, -+ QKV, -+ ) -+ -+ g = torch.Generator(device="cuda").manual_seed(0) -+ sl = 3 + (S - 1) -+ w = { -+ "prefix": _randn(g, S, HIDDEN, scale=2.0), -+ "delta": _randn(g, S, HIDDEN, scale=0.5), -+ "blocks": _randn(g, S, 8, HIDDEN, scale=2.0), -+ "ares_nw": (1 + _randn(g, HIDDEN, scale=0.1, dtype=torch.float32)).to(bf), -+ "ares_qk": _randn(g, HIDDEN, scale=0.05), -+ "in_nw": (1 + _randn(g, HIDDEN, scale=0.1, dtype=torch.float32)).to(bf), -+ "w_in": _randn(g, NPROJ, HIDDEN, scale=0.02), -+ "w_fb": _randn(g, PROJ, HD, scale=0.08), -+ "conv_w": _randn(g, QKV, 4, scale=0.4, dtype=torch.float32), -+ "a_log": _randn(g, NH, scale=0.5, dtype=torch.float32), -+ "dt_bias": _randn(g, PROJ, scale=0.5, dtype=torch.float32), -+ "on_w": (1 + _randn(g, HD, scale=0.1, dtype=torch.float32)).to(bf), -+ } -+ w["w_in"][NPROJ - 4 :].zero_() -+ # vLLM packs a layer's conv and SSM state in one padded page: strided slots -+ pad = _randn(g, slots, NH * HD * HD + 24576, scale=0.05, dtype=torch.float32) -+ w["rstate"] = pad[:, : NH * HD * HD].view(slots, NH, HD, HD) -+ conv = _randn(g, slots, QKV * sl + 4096)[:, : QKV * sl] -+ if dim_first: -+ w["conv_state"] = conv.view(slots, QKV, sl) -+ else: -+ w["conv_state"] = conv.view(slots, sl, QKV).transpose(-1, -2) -+ w["st_idx"] = torch.arange(3, 3 + S, dtype=torch.int32, device="cuda").view(1, S) -+ w["num_acc"] = torch.full((1,), nacc, dtype=torch.int32, device="cuda") -+ return w -+ -+ -+def _kda_reference(w, nb: int, S: int, write_idx: int): -+ from vllm.config import VllmConfig, set_current_vllm_config -+ from vllm.model_executor.layers.mamba.ops.causal_conv1d import ( -+ causal_conv1d_update, -+ ) -+ from vllm.models.kimi_k3.amd.ops.attn_res import attn_res -+ from vllm.models.kimi_k3.amd.ops.third_party.kda import fused_recurrent_kda -+ from vllm.models.kimi_k3.amd.specialized_decode.mono.attention.kda import ( -+ HD, -+ NH, -+ PROJ, -+ QKV, -+ ) -+ from vllm.third_party.flash_linear_attention.ops.kda import FusedRMSNormGated -+ -+ prefix, blocks = w["prefix"].clone(), w["blocks"].clone() -+ x = attn_res( -+ prefix, -+ w["delta"], -+ blocks, -+ w["ares_nw"], -+ w["ares_qk"], -+ w["in_nw"], -+ nb, -+ write_idx, -+ 1e-5, -+ 1e-5, -+ ) -+ proj = F.linear(x, w["w_in"]) -+ mixed_qkv, g2, f_a, beta = proj.split([QKV, PROJ, HD, NH + 4], dim=-1)[:4] -+ g1 = F.linear(f_a, w["w_fb"]) -+ conv_state, rstate = w["conv_state"].clone(), w["rstate"].clone() -+ qsl = torch.tensor([0, S], dtype=torch.int32, device="cuda") -+ conv = causal_conv1d_update( -+ mixed_qkv.contiguous(), -+ conv_state, -+ w["conv_w"], -+ None, -+ activation="silu", -+ conv_state_indices=w["st_idx"][:, 0].contiguous(), -+ num_accepted_tokens=w["num_acc"], -+ query_start_loc=qsl, -+ max_query_len=S, -+ validate_data=False, -+ out=torch.empty_like(mixed_qkv), -+ ) -+ q, k, v = ( -+ rearrange(t, "n (h d) -> 1 n h d", d=HD) for t in conv.split(PROJ, dim=-1) -+ ) -+ core = torch.empty(1, S, NH, HD, dtype=bf, device="cuda") -+ fused_recurrent_kda( -+ q=q, -+ k=k, -+ v=v, -+ raw_g=rearrange(g1, "n (h d) -> 1 n h d", d=HD), -+ raw_beta=beta[:, :NH].unsqueeze(0), -+ A_log=w["a_log"], -+ dt_bias=w["dt_bias"], -+ lower_bound=-5.0, -+ initial_state=rstate, -+ cu_seqlens=qsl, -+ ssm_state_indices=w["st_idx"], -+ num_accepted_tokens=w["num_acc"], -+ uniform_sequence_length=S, -+ out=core, -+ ) -+ with set_current_vllm_config(VllmConfig()): -+ norm = FusedRMSNormGated(HD, activation="sigmoid").cuda() -+ norm.weight.data = w["on_w"].clone() -+ out = norm.forward_cuda(core, rearrange(g2, "n (h d) -> 1 n h d", d=HD)) -+ return { -+ "prefix": prefix, -+ "blocks": blocks, -+ "conv_state": conv_state, -+ "rstate": rstate, -+ "core": out.reshape(S, PROJ), -+ } -+ -+ -+@needs_gfx950 -+@pytest.mark.parametrize("dim_first", [True, False]) -+@pytest.mark.parametrize("nacc", [1, 3]) -+@pytest.mark.parametrize("write", [False, True]) -+def test_k1_matches_vllm_ops(dim_first: bool, nacc: int, write: bool) -> None: -+ """K1 against attn_res -> in_proj -> f_b -> causal_conv1d_update -> -+ fused_recurrent_kda -> FusedRMSNormGated on one spec-verify step (one -+ request, 8 tokens), including the state each step leaves behind.""" -+ from vllm.models.kimi_k3.amd.specialized_decode.mono.attention.kda import ( -+ PROJ, -+ KdaPreBuild, -+ kda_pre, -+ scratch_bytes, -+ ) -+ from vllm.models.kimi_k3.amd.specialized_decode.mono.common.plan import HEALTH_WORDS -+ -+ S, nb = 8, 4 -+ # a block-write layer appends block nb (block_write_idx == prev_valid_blocks) -+ write_idx = nb if write else -1 -+ w = _kda_inputs(nb, nacc, dim_first, S) -+ ref = _kda_reference(w, nb, S, write_idx) -+ -+ key = KdaPreBuild( -+ tokens=S, -+ qlen=S, -+ nblocks=nb, -+ delta=True, -+ state_len=w["conv_state"].size(-1), -+ write_idx=write_idx, -+ ) -+ prefix, blocks = w["prefix"].clone(), w["blocks"].clone() -+ conv_state, rstate = w["conv_state"].clone(), w["rstate"].clone() -+ core = torch.zeros(S, PROJ, dtype=bf, device="cuda") -+ health = torch.zeros(HEALTH_WORDS, dtype=torch.int32, device="cuda") -+ flags = torch.zeros(1, dtype=torch.int32, device="cuda") -+ kda_pre( -+ key, -+ prefix=prefix, -+ delta=w["delta"], -+ blocks=blocks, -+ ares_nw=w["ares_nw"], -+ ares_qk=w["ares_qk"], -+ in_nw=w["in_nw"], -+ w_in=w["w_in"], -+ w_fb=w["w_fb"], -+ conv_w=w["conv_w"], -+ conv_state=conv_state, -+ a_log=w["a_log"], -+ dt_bias=w["dt_bias"], -+ on_w=w["on_w"], -+ rstate=rstate, -+ st_idx=w["st_idx"], -+ num_acc=w["num_acc"], -+ core_out=core, -+ scratch=torch.zeros(scratch_bytes(key), dtype=torch.uint8, device="cuda"), -+ layer=0, -+ epoch=torch.ones(1, dtype=torch.int32, device="cuda"), -+ health=health, -+ alarm=flags.data_ptr(), -+ ) -+ torch.accelerator.synchronize() -+ assert not health.any() and not flags.any() -+ -+ torch.testing.assert_close(prefix, ref["prefix"], atol=0, rtol=0) -+ torch.testing.assert_close(blocks, ref["blocks"], atol=0, rtol=0) -+ # the conv output feeds the recurrence in fp32 here, bf16 in the reference -+ errs = { -+ "conv_state": _rel_l2(conv_state, ref["conv_state"]), -+ "rstate": _rel_l2(rstate, ref["rstate"]), -+ "core": _rel_l2(core, ref["core"]), -+ } -+ assert errs["conv_state"] < 1e-3 and errs["rstate"] < 1e-3, errs -+ assert _cos(core, ref["core"]) > COS_MIN, errs -+ -+ -+# ---------------------------------------------------------------- K2 / MoE -+ -+ -+def _situ(g, u): -+ g, u = g.float(), u.float() -+ a = (2.0 * torch.tanh(g * 0.25)) * (1.0 + torch.tanh(g * 0.5)) -+ return (a * (25.0 * torch.tanh(u * 0.04))).to(bf) -+ -+ -+def _all_reduce_fp32(t: torch.Tensor) -> torch.Tensor: -+ """The ranks' bf16 partials summed in fp32, as vLLM's custom all-reduce.""" -+ t = t.float() -+ dist.all_reduce(t) -+ return t.to(bf) -+ -+ -+def _moe_reference(x, mw, rank: int): -+ """KimiMoE's decode path on ``x``: sigmoid router + bias top-16, the latent -+ down, aiter's MXFP4 experts, the latent all-reduce + RMSNorm, this rank's -+ up_proj rows into the shared experts' output, the final all-reduce.""" -+ from aiter import ActivationType, QuantType -+ from aiter.fused_moe import fused_moe -+ -+ from vllm.models.kimi_k3.amd.specialized_decode.mono.stages.moe import ( -+ SI, -+ TOPK, -+ UP_N, -+ ) -+ -+ sig = torch.sigmoid(torch.mm(x.float(), mw["w_gate"].float().t())) -+ ids = torch.topk(sig + mw["bias"], TOPK, dim=-1, sorted=True).indices -+ wts = sig.gather(-1, ids) -+ wts = wts / wts.sum(-1, keepdim=True) -+ gu = F.linear(x, mw["w_sgu"]) -+ shared = F.linear(_situ(gu[:, :SI], gu[:, SI:]), mw["w_sd"]) -+ y = fused_moe( -+ F.linear(x, mw["w_ld"]), -+ mw["w13"], -+ mw["w2"], -+ wts.float(), -+ ids.to(torch.int32), -+ None, -+ ActivationType.Situv2, -+ QuantType.per_1x32, -+ False, -+ mw["w13s"], -+ mw["w2s"], -+ None, -+ None, -+ dtype=bf, -+ gate_mode="separated", -+ beta=4.0, -+ linear_beta=25.0, -+ ) -+ lat = _all_reduce_fp32(y).float() -+ ln = lat * torch.rsqrt(lat.pow(2).mean(-1, keepdim=True) + 1e-5) * mw["ln_w"] -+ lo = rank * UP_N -+ up = ln.to(bf).float() @ mw["w_up_shard"].float().t() -+ shared[:, lo : lo + UP_N] = (shared[:, lo : lo + UP_N].float() + up).to(bf) -+ return _all_reduce_fp32(shared) -+ -+ -+def _moe_weights(rank: int, device): -+ from aiter.utility.fp4_utils import e8m0_shuffle -+ -+ from vllm._aiter_ops import rocm_aiter_ops -+ from vllm.models.kimi_k3.amd.specialized_decode.mono.stages.moe import ( -+ HIDDEN, -+ LAT, -+ RI, -+ SI, -+ UP_N, -+ E, -+ ) -+ -+ gc = torch.Generator(device=device).manual_seed(7) # replicated weights -+ gr = torch.Generator(device=device).manual_seed(100 + rank) # TP shards -+ -+ def u8(g, *shape, lo=0, hi=256): -+ return torch.randint(lo, hi, shape, generator=g, device=device).to(torch.uint8) -+ -+ fp4, e8 = torch.float4_e2m1fn_x2, torch.float8_e8m0fnu -+ w_up = _randn(gc, HIDDEN, LAT, scale=0.02) -+ return { -+ "w_gate": _randn(gc, E, HIDDEN, scale=0.02), -+ "bias": _randn(gc, E, scale=0.02, dtype=torch.float32), -+ "w_ld": _randn(gc, LAT, HIDDEN, scale=0.02), -+ "ln_w": (1 + _randn(gc, LAT, scale=0.1, dtype=torch.float32)).to(bf), -+ "w_up_shard": w_up[rank * UP_N : (rank + 1) * UP_N].contiguous(), -+ "w_sgu": _randn(gr, 2 * SI, HIDDEN, scale=0.02), -+ "w_sd": _randn(gr, HIDDEN, SI, scale=0.03), -+ "w13": rocm_aiter_ops.shuffle_weight_a16w4( -+ u8(gr, E, 2 * RI, LAT // 2).view(fp4), 16, False -+ ), -+ "w2": rocm_aiter_ops.shuffle_weight_a16w4( -+ u8(gr, E, LAT, RI // 2).view(fp4), 16, False -+ ), -+ "w13s": rocm_aiter_ops.shuffle_scale_a16w4( -+ u8(gr, E, 2 * RI, LAT // 32, lo=118, hi=123).view(e8).view(-1, LAT // 32), -+ E, -+ False, -+ ), -+ "w2s": e8m0_shuffle( -+ u8(gr, E, LAT, RI // 32, lo=118, hi=123).view(e8).view(-1, RI // 32) -+ ), -+ } -+ -+ -+def _moe_args(mw): -+ return { -+ "w_gate": mw["w_gate"], -+ "bias": mw["bias"], -+ "w_ld": mw["w_ld"], -+ "w_sgu": mw["w_sgu"], -+ "w_sd": mw["w_sd"], -+ "w13": mw["w13"], -+ "w13s": mw["w13s"], -+ "w2": mw["w2"], -+ "w2s": mw["w2s"], -+ "ln_w": mw["ln_w"], -+ "w_up": mw["w_up_shard"], -+ } -+ -+ -+def _k2_inputs(rank: int, device): -+ from vllm.models.kimi_k3.amd.specialized_decode.mono.attention.kda import ( -+ HIDDEN, -+ PROJ, -+ ) -+ -+ gc = torch.Generator(device=device).manual_seed(11) -+ gr = torch.Generator(device=device).manual_seed(200 + rank) -+ return { -+ "core": _randn(gr, 8, PROJ), -+ "w_o": _randn(gr, HIDDEN, PROJ, scale=0.02), -+ "prefix": _randn(gc, 8, HIDDEN, scale=2.0), -+ "blocks": _randn(gc, 8, 8, HIDDEN, scale=2.0), -+ "ares_nw": (1 + _randn(gc, HIDDEN, scale=0.1, dtype=torch.float32)).to(bf), -+ "ares_qk": _randn(gc, HIDDEN, scale=0.05), -+ "in_nw": (1 + _randn(gc, HIDDEN, scale=0.1, dtype=torch.float32)).to(bf), -+ } -+ -+ -+def _k2_launch(runner, w, mw, prefix, out, reset: bool, nb: int = 4): -+ from vllm.models.kimi_k3.amd.specialized_decode.mono import layer as k2 -+ -+ key = k2.K2Build(tokens=8, nblocks=nb + int(reset), reset=reset, **runner.fault) -+ k2.k2_launch( -+ key, -+ core=w["core"], -+ w_o=w["w_o"], -+ prefix=prefix, -+ blocks=w["blocks"], -+ ares_nw=w["ares_nw"], -+ ares_qk=w["ares_qk"], -+ in_nw=w["in_nw"], -+ out=out, -+ layer=3, -+ **_moe_args(mw), -+ **runner.launch_args(), -+ **runner.peer_args(), -+ ) -+ return key -+ -+ -+def _init_worker(monkeypatch, tp_size, pp_size, rank, distributed_init_port): -+ monkeypatch.delenv("CUDA_VISIBLE_DEVICES", raising=False) -+ device = torch.device(f"cuda:{rank}") -+ torch.accelerator.set_device_index(device) -+ init_test_distributed_environment(tp_size, pp_size, rank, distributed_init_port) -+ return device -+ -+ -+@ray.remote(num_gpus=1, max_calls=1) -+def _k2_worker( -+ monkeypatch: pytest.MonkeyPatch, -+ tp_size: int, -+ pp_size: int, -+ rank: int, -+ distributed_init_port: str, -+) -> None: -+ device = _init_worker(monkeypatch, tp_size, pp_size, rank, distributed_init_port) -+ -+ from vllm.models.kimi_k3.amd.ops.attn_res import attn_res -+ from vllm.models.kimi_k3.amd.specialized_decode.mono import layer as k2 -+ from vllm.models.kimi_k3.amd.specialized_decode.mono.attention.kda import HIDDEN -+ from vllm.models.kimi_k3.amd.specialized_decode.runtime.runner import MonoRunner -+ -+ S, nb = 8, 4 -+ runner = MonoRunner(device) -+ mw = _moe_weights(rank, device) -+ w = _k2_inputs(rank, device) -+ -+ for reset in (False, True): -+ runner.step_begin() -+ prefix = w["prefix"].clone() -+ out = torch.empty(S, HIDDEN, dtype=bf, device=device) -+ key = _k2_launch(runner, w, mw, prefix, out, reset) -+ torch.accelerator.synchronize() -+ -+ attn = _all_reduce_fp32(F.linear(w["core"], w["w_o"])) -+ # a block-write layer's prefix restarts at the attention output -+ ref_prefix = attn if reset else w["prefix"].clone() -+ delta = None if reset else attn -+ x_ref = attn_res( -+ ref_prefix, -+ delta, -+ w["blocks"], -+ w["ares_nw"], -+ w["ares_qk"], -+ w["in_nw"], -+ nb + int(reset), -+ -1, -+ 1e-5, -+ 1e-5, -+ ) -+ # the kernel sums the ranks in rank order, NCCL in its own order -+ assert _rel_l2(prefix, ref_prefix) < 1e-2 -+ lay = k2.scratch_layout(key)["xrow"][0] -+ x = runner.scratch[lay : lay + S * HIDDEN * 2].view(bf).view(S, HIDDEN) -+ assert _cos(x, x_ref) > COS_MIN -+ assert _cos(out, _moe_reference(x, mw, rank)) > COS_MIN -+ outs = [torch.empty_like(out) for _ in range(tp_size)] -+ dist.all_gather(outs, out) -+ assert all(torch.equal(outs[0], o) for o in outs), "ranks disagree" -+ -+ # A graph of one step (the epoch bump captured with it), replayed: a -+ # block-write launch writes its prefix, so every replay must match the first. -+ prefix = torch.empty_like(w["prefix"]) -+ out = torch.empty(S, HIDDEN, dtype=bf, device=device) -+ graph = torch.cuda.CUDAGraph() -+ with torch.cuda.graph(graph): -+ runner.step_begin() -+ _k2_launch(runner, w, mw, prefix, out, reset=True) -+ graph.replay() -+ torch.accelerator.synchronize() -+ want_prefix, want_out = prefix.clone(), out.clone() -+ for _ in range(300): -+ graph.replay() -+ torch.accelerator.synchronize() -+ assert torch.equal(prefix, want_prefix) and torch.equal(out, want_out) -+ assert not runner.health.any() and not runner.flags.any() -+ -+ -+@needs_gfx950 -+@multi_gpu_test(num_gpus=8) -+def test_k2_tp8_matches_vllm_ops_and_replays(monkeypatch: pytest.MonkeyPatch) -> None: -+ multi_process_parallel(monkeypatch, 8, 1, _k2_worker) -+ -+ -+STALLED = 3 -+ -+ -+@ray.remote(num_gpus=1, max_calls=1) -+def _fault_worker( -+ monkeypatch: pytest.MonkeyPatch, -+ tp_size: int, -+ pp_size: int, -+ rank: int, -+ distributed_init_port: str, -+) -> None: -+ device = _init_worker(monkeypatch, tp_size, pp_size, rank, distributed_init_port) -+ -+ from vllm.models.kimi_k3.amd.specialized_decode.mono import layer as k2 -+ from vllm.models.kimi_k3.amd.specialized_decode.mono.attention.kda import HIDDEN -+ from vllm.models.kimi_k3.amd.specialized_decode.mono.common.debug import ( -+ given_up_waits, -+ ) -+ from vllm.models.kimi_k3.amd.specialized_decode.mono.common.sync import WAIT_TICKS -+ from vllm.models.kimi_k3.amd.specialized_decode.runtime.runner import MonoRunner -+ -+ runner = MonoRunner(device) -+ # rank STALLED's CTA 0 stalls past every wait's bound in the epoch 2 launch -+ runner.fault = { -+ "fault_rank": STALLED, -+ "fault_epoch": 2, -+ "fault_layer": 3, -+ "fault_ticks": WAIT_TICKS * 3 // 2, -+ } -+ mw = _moe_weights(rank, device) -+ w = _k2_inputs(rank, device) -+ out = torch.empty(8, HIDDEN, dtype=bf, device=device) -+ runner.step_begin() # epoch 1: no stall -+ _k2_launch(runner, w, mw, w["prefix"].clone(), out, reset=True) -+ torch.accelerator.synchronize() -+ assert not runner.flags.any() -+ -+ runner.step_begin() # epoch 2: rank STALLED stalls; every launch ends -+ _k2_launch(runner, w, mw, w["prefix"].clone(), out, reset=True) -+ torch.accelerator.synchronize() -+ waits = given_up_waits(runner.health.cpu(), k2.REGIONS) -+ gave_up = torch.tensor([int(bool(waits))], device=device) -+ dist.all_reduce(gave_up) -+ assert gave_up.item() > 0, "no rank's wait gave up" -+ if rank == 0: -+ assert runner.flags.any(), "the output rank's flags missed the failure" -+ -+ -+@needs_gfx950 -+@multi_gpu_test(num_gpus=8) -+def test_stalled_rank_fails_the_step_without_a_hang( -+ monkeypatch: pytest.MonkeyPatch, -+) -> None: -+ multi_process_parallel(monkeypatch, 8, 1, _fault_worker) -+ -+ -+@needs_gfx950 -+def test_warmup_compiles_without_launching() -> None: -+ """The JIT warm-up compiles K1 and K2 before capture without running them: -+ a launch on its null arguments would fault.""" -+ from vllm.models.kimi_k3.amd.specialized_decode.mono.attention.kda import ( -+ KdaPreBuild, -+ ) -+ from vllm.models.kimi_k3.amd.specialized_decode.mono.layer import K2Build -+ from vllm.models.kimi_k3.amd.specialized_decode.runtime.layers import compile_build -+ -+ compile_build(KdaPreBuild(tokens=8, qlen=8, nblocks=4, delta=True, state_len=10)) -+ compile_build(K2Build(tokens=8, nblocks=4)) -+ torch.accelerator.synchronize() -diff --git a/tests/models/kimi_k3/amd_specialized_decode/test_runtime.py b/tests/models/kimi_k3/amd_specialized_decode/test_runtime.py -new file mode 100644 -index 000000000..6cc5c793c ---- /dev/null -+++ b/tests/models/kimi_k3/amd_specialized_decode/test_runtime.py -@@ -0,0 +1,318 @@ -+# SPDX-License-Identifier: Apache-2.0 -+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -+"""Kimi-K3's specialized decode on CPU: the switch, its widths and layers, its -+refusals (config, loaded weights, caches) and plan() routing.""" -+ -+from types import SimpleNamespace -+ -+import pytest -+import torch -+ -+from vllm.config import CompilationMode, CUDAGraphMode -+from vllm.models.kimi_k3.amd.specialized_decode import runtime -+from vllm.models.kimi_k3.amd.specialized_decode.runtime import ( -+ FFN, -+ WHOLE, -+ DecodeConfig, -+ layer_kinds, -+ specialized_decode_config, -+) -+from vllm.models.kimi_k3.amd.specialized_decode.runtime import decode as decode_mod -+from vllm.models.kimi_k3.amd.specialized_decode.runtime.decode import ( -+ SpecializedDecode, -+ step_reason, -+) -+from vllm.models.kimi_k3.amd.specialized_decode.runtime.guards import ( -+ cache_reasons, -+ weight_reasons, -+) -+ -+ -+def _vllm_config(**over) -> SimpleNamespace: -+ """A supported config: TP8, DCP8, 7 speculative tokens, FULL graphs, MRV2. -+ ``over`` sets dotted paths, e.g. ``{"parallel_config.data_parallel_size": 2}``.""" -+ spec = SimpleNamespace( -+ num_speculative_tokens=7, -+ enable_adaptive_verification=False, -+ dynamic=False, -+ ) -+ spec.uses_dynamic_speculative_decoding = lambda: spec.dynamic -+ c = SimpleNamespace( -+ additional_config={"decode_backend": "latency_specialized"}, -+ parallel_config=SimpleNamespace( -+ tensor_parallel_size=8, -+ pipeline_parallel_size=1, -+ data_parallel_size=1, -+ enable_expert_parallel=False, -+ enable_eplb=False, -+ decode_context_parallel_size=8, -+ prefill_context_parallel_size=1, -+ use_ubatching=False, -+ ), -+ model_config=SimpleNamespace( -+ dtype=torch.bfloat16, -+ enable_sleep_mode=False, -+ hf_text_config=SimpleNamespace( -+ hidden_size=7168, -+ num_experts=896, -+ attn_res_block_size=12, -+ linear_attn_config={ -+ "num_heads": 96, -+ "head_dim": 128, -+ "short_conv_kernel_size": 4, -+ }, -+ ), -+ ), -+ cache_config=SimpleNamespace( -+ mamba_cache_dtype="auto", use_replayssm=False, use_kda_recoverssm=False -+ ), -+ lora_config=None, -+ kv_transfer_config=None, -+ aux_output_config=SimpleNamespace(enable_return_routed_experts=False), -+ offload_config=SimpleNamespace( -+ uva=SimpleNamespace(cpu_offload_gb=0), -+ prefetch=SimpleNamespace(offload_group_size=0), -+ ), -+ weight_transfer_config=None, -+ speculative_config=spec, -+ use_v2_model_runner=True, -+ compilation_config=SimpleNamespace( -+ mode=CompilationMode.VLLM_COMPILE, -+ cudagraph_mode=CUDAGraphMode.FULL_AND_PIECEWISE, -+ ), -+ ) -+ for path, v in over.items(): -+ obj = c -+ *head, last = path.split(".") -+ for h in head: -+ obj = getattr(obj, h) -+ setattr(obj, last, v) -+ return c -+ -+ -+@pytest.fixture -+def on_mi355x(monkeypatch): -+ monkeypatch.setattr(runtime, "platform_reasons", lambda: []) -+ -+ -+@pytest.mark.parametrize("k, want", [(7, (8,)), (3, (4, 8)), (1, (2, 4, 6, 8))]) -+def test_widths_are_whole_requests_up_to_eight_rows(k, want): -+ assert DecodeConfig(q=k + 1).widths == want -+ -+ -+def _k3_text_config(): -+ """Kimi-K3's layer pattern: 93 layers, MLA every 4th, AttnRes blocks of 12, -+ layer 0 dense.""" -+ mla = {i for i in range(93) if i % 4 == 3} | {92} -+ return SimpleNamespace( -+ num_hidden_layers=93, -+ attn_res_block_size=12, -+ num_experts=896, -+ first_k_dense_replace=1, -+ moe_layer_freq=1, -+ is_kda_layer=lambda i: i not in mla, -+ ) -+ -+ -+def test_kernel_layers_follow_the_k3_pattern(): -+ kinds = layer_kinds(_k3_text_config()) -+ assert list(kinds) == list(range(1, 93)) -+ assert kinds[1] == WHOLE and kinds[12] == WHOLE -+ assert kinds[3] == FFN and kinds[92] == FFN -+ -+ -+def test_auto_runs_the_stock_decode(): -+ """Without the switch nothing is checked, even on an unsupported deployment.""" -+ cfg = _vllm_config( -+ additional_config={}, **{"parallel_config.tensor_parallel_size": 4} -+ ) -+ assert specialized_decode_config(cfg) is None -+ cfg.additional_config = {"decode_backend": "auto"} -+ assert specialized_decode_config(cfg) is None -+ -+ -+def test_an_unknown_backend_is_refused(): -+ cfg = _vllm_config(additional_config={"decode_backend": "fast"}) -+ with pytest.raises(ValueError, match="decode_backend must be"): -+ specialized_decode_config(cfg) -+ -+ -+def test_supported_config_is_accepted(on_mi355x): -+ assert specialized_decode_config(_vllm_config()) == DecodeConfig(q=8) -+ -+ -+@pytest.mark.parametrize( -+ "over, reason", -+ [ -+ ( -+ { -+ "parallel_config.tensor_parallel_size": 4, -+ "parallel_config.decode_context_parallel_size": 4, -+ }, -+ "TP 4", -+ ), -+ ({"parallel_config.pipeline_parallel_size": 2}, "pipeline parallelism"), -+ ({"parallel_config.decode_context_parallel_size": 2}, "decode context"), -+ ({"kv_transfer_config": object()}, "KV transfer connectors"), -+ ({"speculative_config": None}, "without speculative decoding"), -+ ({"speculative_config.num_speculative_tokens": 8}, "more than 7"), -+ ({"speculative_config.enable_adaptive_verification": True}, "adaptive"), -+ ({"speculative_config.dynamic": True}, "dynamic speculative"), -+ ({"cache_config.mamba_cache_dtype": "float32"}, "KDA cache dtypes"), -+ ({"cache_config.use_replayssm": True}, "ReplaySSM"), -+ ({"use_v2_model_runner": False}, "V1 model runner"), -+ ( -+ {"compilation_config.mode": CompilationMode.STOCK_TORCH_COMPILE}, -+ "torch.compile", -+ ), -+ ({"model_config.hf_text_config.hidden_size": 4096}, "model shapes"), -+ ], -+) -+def test_refusals(on_mi355x, over, reason): -+ reasons = runtime.unsupported(_vllm_config(**over)) -+ assert len(reasons) == 1 and reason in reasons[0], reasons -+ -+ -+def test_every_refusal_is_reported_at_once(monkeypatch): -+ monkeypatch.setattr( -+ runtime, "platform_reasons", lambda: ["this GPU (gfx950 / MI355X only)"] -+ ) -+ cfg = _vllm_config( -+ **{"parallel_config.tensor_parallel_size": 4, "use_v2_model_runner": False} -+ ) -+ with pytest.raises(ValueError) as e: -+ specialized_decode_config(cfg) -+ for reason in ("TP 4", "V1 model runner", "this GPU", "decode_backend=auto"): -+ assert reason in str(e.value) -+ -+ -+def test_piecewise_graphs_only_warn(on_mi355x): -+ """They are not refused: plan() runs every PIECEWISE step stock.""" -+ cfg = _vllm_config(**{"compilation_config.cudagraph_mode": CUDAGraphMode.PIECEWISE}) -+ assert runtime.piecewise_graphs_only(cfg) -+ assert specialized_decode_config(cfg) is not None -+ -+ -+def test_cache_refusal_names_the_cache(): -+ pytest.importorskip("flydsl") -+ conv = torch.zeros(2, 4, 3, dtype=torch.float32) -+ rstate = torch.zeros(2, 12, 128, 128, dtype=torch.bfloat16) -+ reasons = cache_reasons({"layers.5.attn": (conv, rstate)}) -+ assert len(reasons) == 2 and all(r.startswith("layers.5.attn:") for r in reasons) -+ -+ -+def _kda_metadata(**md): -+ """One request of 4 tokens unless overridden.""" -+ m = SimpleNamespace( -+ spec_sequence_masks=torch.ones(1, dtype=torch.bool), -+ num_prefills=0, -+ num_decodes=0, -+ num_spec_decodes=1, -+ spec_state_indices_tensor=torch.zeros(1, 4, dtype=torch.int32), -+ uniform_spec_sequence_length=4, -+ ) -+ for k, v in md.items(): -+ setattr(m, k, v) -+ return {"kda": m} -+ -+ -+@pytest.mark.parametrize( -+ "width, md, why", -+ [ -+ (4, {}, None), -+ (8, {"num_spec_decodes": 2}, None), -+ (4, {"num_prefills": 1}, "not pure decode"), -+ (4, {"num_decodes": 1}, "not pure decode"), -+ (4, {"spec_sequence_masks": None}, "not spec-verify"), -+ (8, {}, "not uniform"), -+ (4, {"uniform_spec_sequence_length": None}, "not uniform"), -+ ], -+) -+def test_steps_must_be_pure_uniform_spec_verify(width, md, why): -+ assert step_reason(_kda_metadata(**md), "kda", width) == why -+ -+ -+def test_steps_need_metadata(): -+ assert step_reason(None, "kda", 4) == "no metadata" -+ -+ -+class _Kernel: -+ def launch(self, step, positions, hidden_states, residual, prefix_delta): -+ return "kernel", step.width -+ -+ -+def _decode(monkeypatch, mode=CUDAGraphMode.FULL, md=None): -+ """An executor with kernels on layers 1 and 2 at widths 4 and 8, bound to a -+ KV cache, under a forward context of ``mode``.""" -+ decode = object.__new__(SpecializedDecode) -+ decode.config = DecodeConfig(q=4) -+ decode.runner = SimpleNamespace(kda_prefix="kda") -+ decode.kernels = {w: {1: _Kernel(), 2: _Kernel()} for w in (4, 8)} -+ decode._kv = (1, 2) -+ monkeypatch.setattr(SpecializedDecode, "_kv_key", lambda self: (1, 2)) -+ context = SimpleNamespace( -+ cudagraph_runtime_mode=mode, -+ attn_metadata=_kda_metadata() if md is None else md, -+ ) -+ monkeypatch.setattr(decode_mod, "is_forward_context_available", lambda: True) -+ monkeypatch.setattr(decode_mod, "get_forward_context", lambda: context) -+ return decode -+ -+ -+def test_plan_runs_the_kernels_of_its_width(monkeypatch): -+ step = _decode(monkeypatch).plan(torch.zeros(4), None) -+ assert step is not None and step.width == 4 -+ assert step.layer(1)(None, None, None, None) == ("kernel", 4) -+ assert step.layer(0) is None and step.layer(3) is None -+ -+ -+@pytest.mark.parametrize( -+ "mode, tokens, md, intermediate", -+ [ -+ (CUDAGraphMode.PIECEWISE, 4, None, None), -+ (CUDAGraphMode.FULL, 16, None, None), # no kernels that wide -+ (CUDAGraphMode.FULL, 4, _kda_metadata(num_prefills=1), None), -+ (CUDAGraphMode.FULL, 4, None, object()), # a pipeline stage's input -+ ], -+) -+def test_plan_runs_other_steps_stock(monkeypatch, mode, tokens, md, intermediate): -+ decode = _decode(monkeypatch, mode, md) -+ assert decode.plan(torch.zeros(tokens), intermediate) is None -+ -+ -+def test_plan_runs_stock_before_a_kv_cache_is_bound(monkeypatch): -+ decode = _decode(monkeypatch) -+ decode._kv = None -+ assert decode.plan(torch.zeros(4), None) is None -+ -+ -+def test_plan_refuses_a_kv_cache_bound_behind_its_back(monkeypatch): -+ decode = _decode(monkeypatch) -+ monkeypatch.setattr(SpecializedDecode, "_kv_key", lambda self: (3, 4)) -+ with pytest.raises(RuntimeError, match="without prepare_kv"): -+ decode.plan(torch.zeros(4), None) -+ -+ -+def test_weight_views_reject_every_mismatch(): -+ """weight_reasons names each layer whose weights the kernels cannot read as -+ they are, and a MoE backend that lays them out otherwise.""" -+ pytest.importorskip("flydsl") -+ from vllm.models.kimi_k3.amd.specialized_decode.mono.stages import moe -+ -+ bad_moe = SimpleNamespace( -+ use_latent_moe=True, -+ shared_experts=object(), -+ routed_scaling_factor=2.0, -+ experts=SimpleNamespace(quant_method=SimpleNamespace()), -+ ) -+ no_latent = SimpleNamespace(use_latent_moe=False, shared_experts=None) -+ layers = [None] + [SimpleNamespace(mlp=no_latent)] * 92 -+ layers[5] = SimpleNamespace(mlp=bad_moe) -+ model = SimpleNamespace( -+ config=_k3_text_config(), layers=layers, start_layer=0, end_layer=93 -+ ) -+ reasons = weight_reasons(model) -+ assert "layer 1: needs a latent MoE with shared experts" in reasons -+ assert "layer 5: routed_scaling_factor 2.0" in reasons -+ assert any(r.startswith("layer 5: MoE backend must be") for r in reasons) -+ assert moe.E == 896 -diff --git a/vllm/models/kimi_k3/amd/linear.py b/vllm/models/kimi_k3/amd/linear.py -index 682773189..029636146 100644 ---- a/vllm/models/kimi_k3/amd/linear.py -+++ b/vllm/models/kimi_k3/amd/linear.py -@@ -74,6 +74,10 @@ from vllm.models.kimi_k3.amd.kda import KimiK3DeltaAttention - from vllm.models.kimi_k3.amd.latent_moe_runner import ROCmLatentMoERunner - from vllm.models.kimi_k3.amd.mla import KimiK3MultiHeadLatentAttentionWrapper - from vllm.models.kimi_k3.amd.ops.attn_res import attn_res -+from vllm.models.kimi_k3.amd.specialized_decode.hooks import SpecializedDecodeMixin -+from vllm.models.kimi_k3.amd.specialized_decode.runtime import ( -+ specialized_decode_config, -+) - from vllm.sequence import IntermediateTensors - from vllm.transformers_utils.configs.kimi_linear import KimiLinearConfig - from vllm.utils.math_utils import cdiv -@@ -828,6 +832,7 @@ class KimiLinearModel(nn.Module, EagleModelMixin): - positions: torch.Tensor, - intermediate_tensors: IntermediateTensors | None, - inputs_embeds: torch.Tensor | None = None, -+ step=None, # specialized_decode's plan - **kwargs, - ) -> torch.Tensor | IntermediateTensors | tuple[torch.Tensor, list[torch.Tensor]]: - if get_pp_group().is_first_rank: -@@ -885,11 +890,9 @@ class KimiLinearModel(nn.Module, EagleModelMixin): - self.layers[self.start_layer : self.end_layer], - start=self.start_layer, - ): -- hidden_states, residual, prefix_delta = layer( -- positions=positions, -- hidden_states=hidden_states, -- residual=residual, -- prefix_delta=prefix_delta, -+ kernel = step.layer(layer_idx) if step is not None else None -+ hidden_states, residual, prefix_delta = (kernel or layer)( -+ positions, hidden_states, residual, prefix_delta - ) - if (layer_idx + 1) in self.aux_hidden_state_layers: - self._maybe_add_hidden_state( -@@ -1057,10 +1060,18 @@ class KimiLinearModel(nn.Module, EagleModelMixin): - - - class KimiLinearForCausalLM( -- nn.Module, HasInnerState, SupportsPP, MixtureOfExperts, IsHybrid -+ nn.Module, -+ HasInnerState, -+ SupportsPP, -+ MixtureOfExperts, -+ IsHybrid, -+ SpecializedDecodeMixin, - ): - def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): -+ # raises for an unsupported deployment before any weight is allocated -+ decode_config = specialized_decode_config(vllm_config) - super().__init__() -+ self.init_specialized_decode(decode_config) - self.model_config = vllm_config.model_config - self.vllm_config = vllm_config - self.config = self.model_config.hf_config -@@ -1102,8 +1113,9 @@ class KimiLinearForCausalLM( - inputs_embeds: torch.Tensor | None = None, - **kwargs, - ) -> torch.Tensor | IntermediateTensors: -+ step = self.plan_specialized_decode(positions, intermediate_tensors) - hidden_states = self.model( -- input_ids, positions, intermediate_tensors, inputs_embeds, **kwargs -+ input_ids, positions, intermediate_tensors, inputs_embeds, step, **kwargs - ) - return hidden_states - -@@ -1148,6 +1160,7 @@ class KimiLinearForCausalLM( - ) -> torch.Tensor | None: - # The model's final norm is applied here (not at the end of forward) so - # that the pre-norm hidden states can be fed to the MTP draft model. -+ self.after_specialized_decode_step() - hidden_states = self.model.norm(hidden_states, None) - return self.logits_processor(self.lm_head, hidden_states) - -diff --git a/vllm/models/kimi_k3/amd/mla.py b/vllm/models/kimi_k3/amd/mla.py -index 48ea1792b..2db1575a9 100644 ---- a/vllm/models/kimi_k3/amd/mla.py -+++ b/vllm/models/kimi_k3/amd/mla.py -@@ -196,6 +196,7 @@ class KimiK3MultiHeadLatentAttentionWrapper(MultiHeadLatentAttentionWrapper): - positions: torch.Tensor, - hidden_states: torch.Tensor | QuantizedActivation, - llama_4_scaling: torch.Tensor | None = None, -+ skip_o_proj: bool = False, - ) -> torch.Tensor: - q_c = None - num_tokens = ( -@@ -311,4 +312,6 @@ class KimiK3MultiHeadLatentAttentionWrapper(MultiHeadLatentAttentionWrapper): - else: - attn_out = attn_out * gate.sigmoid() - -+ if skip_o_proj: # the specialized decode runs o_proj in its kernel -+ return attn_out - return self.o_proj(attn_out)[0] -diff --git a/vllm/models/kimi_k3/amd/model.py b/vllm/models/kimi_k3/amd/model.py -index 124fca3e3..aacbd086b 100644 ---- a/vllm/models/kimi_k3/amd/model.py -+++ b/vllm/models/kimi_k3/amd/model.py -@@ -220,6 +220,9 @@ class KimiK3ForConditionalGeneration( - def compute_logits(self, hidden_states: torch.Tensor, **kwargs) -> torch.Tensor: - return self.language_model.compute_logits(hidden_states) - -+ def get_model_state_cls(self) -> type: -+ return self.language_model.get_model_state_cls() -+ - def copy_inputs_before_cuda_graphs(self, input_buffers, **kwargs): - return self.language_model.mamba_cache.copy_inputs_before_cuda_graphs( - input_buffers, **kwargs -diff --git a/vllm/models/kimi_k3/amd/specialized_decode/README.md b/vllm/models/kimi_k3/amd/specialized_decode/README.md -new file mode 100644 -index 000000000..3e0e048b8 ---- /dev/null -+++ b/vllm/models/kimi_k3/amd/specialized_decode/README.md -@@ -0,0 +1,108 @@ -+# Kimi-K3 specialized decode (MI355X) -+ -+An opt-in decode path for Kimi-K3 at TP8 on MI355X (gfx950), following RFC #60904. -+A pure spec-verify step of at most 8 rows runs each MoE decoder layer as persistent -+FlyDSL MonoKernels (256 CTAs, all resident), in "layer mode", inside the stock -+layer loop. The TP all-reduces run in-kernel over peer memory. -+ -+- **Whole layers** (KDA): K1 (AttnRes, in_proj, f_b, spec conv, KDA recurrence, -+ gated RMSNorm), then K2. -+- **FFN layers** (MLA): vLLM's AttnRes and MLA up to o_proj's input, then K2. -+- **K2**: o_proj + all-reduce + MLP AttnRes, then the latent MoE (router, -+ top-16, MXFP4 experts, shared experts, latent all-reduce + RMSNorm, up_proj, -+ final all-reduce). -+ -+The kernels read vLLM's loaded weights and KDA states in place. Nothing is copied -+or re-quantized. -+ -+## Enabling -+ -+```bash -+VLLM_USE_V2_MODEL_RUNNER=1 vllm serve -tp 8 \ -+ --speculative-config '{"method": "dspark", "num_speculative_tokens": 7, ...}' \ -+ --additional-config '{"decode_backend": "latency_specialized"}' -+``` -+ -+There are no other options. `decode_backend=auto` (the default) runs vLLM's decode. -+ -+## Supported deployments -+ -+TP8 on MI355X with all 256 CUs, decode context parallel 1 or 8, speculative -+decoding with at most 7 speculative tokens, Model Runner V2, BF16 conv / FP32 -+recurrent KDA states. -+ -+If you request the specialized decode for anything else, startup fails with one -+`ValueError` that lists every unsupported setting: -+ -+- **Config, in the model's `__init__`:** pipeline, data, expert or prefill -+ context parallelism, EPLB, DBO, LoRA, KV connectors, routed-expert capture, -+ sleep mode, CPU weight offload, weight updates, no speculative decoding, -+ adaptive verification, dynamic speculative decoding, ReplaySSM / RecoverSSM, -+ the V1 model runner, `STOCK_TORCH_COMPILE`, another model's shapes, another GPU. -+- **Loaded weights, when the model state builds the executor:** each kernel -+ layer's weight shapes and dtypes, and the MoE backend (`AITER_MXFP4_BF16`, K3 -+ SiTU). -+- **KV caches, at every bind:** the KDA state dtypes and layout. -+ -+Every TP rank agrees on each verdict before anything collective runs. -+ -+## Steps -+ -+`plan()` runs once per step. It picks the kernels for a pure spec-verify step of -+whole uniform requests whose width is one of the built widths: every multiple of -+`num_speculative_tokens + 1` up to 8 rows. Every other step runs vLLM's layers: -+prefill, mixed, wider and PIECEWISE steps. -+ -+Under FULL graphs the plan is baked in at capture. vLLM's capture sizes are not -+changed. A replay with fewer requests than its graph pads whole requests. Those -+requests get the reserved null state slot and one accepted token, and the kernels -+serve them like real ones. For example, at `num_speculative_tokens=1` three -+requests replay the 4-request graph. To capture every width exactly, pass them in -+`--cudagraph-capture-sizes`. -+ -+## Fail-stop -+ -+Every in-kernel wait gives up after 10 s. The rank records the wait in its health -+words (sticky: its later waits exit at once) and stops waiting, so a stalled peer -+fails every rank instead of hanging them. -+ -+Each step, `compute_logits` queues a non-blocking copy of the rank's `FAILED` -+word. Every 16 steps the TP ranks agree on what has landed over the TP CPU group, -+and then every rank raises at the same step, so the engine stops. They must stop -+together: a rank that raised alone would leave the others waiting in the next -+collective (the drafter's) for ever. -+ -+Detection takes up to about 16 steps, so tokens of the failed steps can reach the -+client before the engine stops. A zero-delay check needs a core hook and is left to -+a follow-up. -+ -+`tests/models/kimi_k3/amd_specialized_decode/test_engine_fault.py` stalls a rank in -+a served model and checks that the engine stops. -+ -+## Layout -+ -+Everything lives in this folder. The model files only call in: `linear.py` -+(the config check in `__init__`, `plan_specialized_decode` in `forward`, -+`step.layer(i)` in the layer loop, `after_specialized_decode_step` in -+`compute_logits`), `model.py` (`get_model_state_cls`) and `mla.py` -+(`skip_o_proj`, o_proj's input for the MLA layers' K2). -+ -+```text -+hooks.py SpecializedDecodeMixin: the model-side hooks -+model_state.py KimiK3ModelState: builds the executor after load, binds KV caches -+runtime/ -+ __init__.py DecodeConfig, specialized_decode_config(), unsupported() -+ decode.py SpecializedDecode: plan() -> Step, Step.layer(i), prepare_kv(), -+ after_step() -+ guards.py weight and KV cache refusals, the fail-stop watch -+ layers.py the layer kernels: vLLM tensors -> launches -+ runner.py per-rank launch state: scratch, step epoch, health, peers -+ peer.py peer memory and TP agreement (adapted from ROCm/ATOM, MIT) -+mono/ kernel source -+ layer.py K2 -+ attention/ kda.py (K1), back.py -+ stages/ moe.py, gemv.py -+ common/ device ops, mailbox sync, MX helpers, plan / build keys -+ (adapted from ROCm/ATOM atom/mono, MIT) -+ abi.py the error-word layout -+``` -diff --git a/vllm/models/kimi_k3/amd/specialized_decode/__init__.py b/vllm/models/kimi_k3/amd/specialized_decode/__init__.py -new file mode 100644 -index 000000000..4a88231b5 ---- /dev/null -+++ b/vllm/models/kimi_k3/amd/specialized_decode/__init__.py -@@ -0,0 +1,12 @@ -+# SPDX-License-Identifier: Apache-2.0 -+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -+"""Kimi-K3 specialized decode for ROCm (MI355X, TP8), layer mode: FlyDSL -+MonoKernels per MoE decoder layer on pure spec-verify steps, inside the stock layer -+loop. Enabled by ``additional_config={"decode_backend": "latency_specialized"}``. -+ -+``hooks.py`` holds the model-side hooks (``SpecializedDecodeMixin``) and -+``model_state.py`` the model runner state that builds the executor. ``runtime/`` is -+the host side: the config check (``runtime/__init__.py``), the executor -+(``runtime/decode.py``), the layer kernels, the per-rank launch state and the later -+refusals. ``mono/`` is the kernel source. Tests live in -+``tests/models/kimi_k3/amd_specialized_decode/``.""" -diff --git a/vllm/models/kimi_k3/amd/specialized_decode/hooks.py b/vllm/models/kimi_k3/amd/specialized_decode/hooks.py -new file mode 100644 -index 000000000..32b448050 ---- /dev/null -+++ b/vllm/models/kimi_k3/amd/specialized_decode/hooks.py -@@ -0,0 +1,63 @@ -+# SPDX-License-Identifier: Apache-2.0 -+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -+"""The specialized decode's hooks on ``KimiLinearForCausalLM``, kept here so the -+model only calls them: ``init_specialized_decode`` in ``__init__`` (after the config -+check, which runs before any module allocates weights), ``plan_specialized_decode`` -+in ``forward`` and ``after_specialized_decode_step`` in ``compute_logits``. The -+model state (``model_state.KimiK3ModelState``) builds the executor after load.""" -+ -+from typing import TYPE_CHECKING -+ -+import torch -+ -+from vllm.models.kimi_k3.amd.specialized_decode.runtime import DecodeConfig -+ -+if TYPE_CHECKING: -+ from vllm.models.kimi_k3.amd.specialized_decode.runtime.decode import ( -+ SpecializedDecode, -+ Step, -+ ) -+ from vllm.v1.worker.gpu.model_states.interface import ModelState -+ -+ -+class SpecializedDecodeMixin: -+ """``self.model`` is the ``KimiLinearModel`` the executor runs over.""" -+ -+ model: torch.nn.Module -+ specialized_decode_config: DecodeConfig | None -+ specialized_decode: "SpecializedDecode | None" -+ -+ def init_specialized_decode(self, config: DecodeConfig | None) -> None: -+ """``config``: ``specialized_decode_config(vllm_config)``, None for the -+ stock decode.""" -+ self.specialized_decode_config = config -+ self.specialized_decode = None -+ -+ def build_specialized_decode(self) -> None: -+ from vllm.models.kimi_k3.amd.specialized_decode.runtime.decode import ( -+ SpecializedDecode, -+ ) -+ -+ assert self.specialized_decode_config is not None -+ self.specialized_decode = SpecializedDecode( -+ self.model, self.specialized_decode_config -+ ) -+ -+ def get_model_state_cls(self) -> "type[ModelState]": -+ from vllm.models.kimi_k3.amd.specialized_decode.model_state import ( -+ KimiK3ModelState, -+ ) -+ -+ return KimiK3ModelState -+ -+ def plan_specialized_decode( -+ self, positions: torch.Tensor, intermediate_tensors -+ ) -> "Step | None": -+ """This step's plan (``Step.layer(i)`` per decoder layer), None for the -+ stock forward.""" -+ decode = self.specialized_decode -+ return None if decode is None else decode.plan(positions, intermediate_tensors) -+ -+ def after_specialized_decode_step(self) -> None: -+ if self.specialized_decode is not None: -+ self.specialized_decode.after_step() -diff --git a/vllm/models/kimi_k3/amd/specialized_decode/model_state.py b/vllm/models/kimi_k3/amd/specialized_decode/model_state.py -new file mode 100644 -index 000000000..b1e3898a4 ---- /dev/null -+++ b/vllm/models/kimi_k3/amd/specialized_decode/model_state.py -@@ -0,0 +1,27 @@ -+# SPDX-License-Identifier: Apache-2.0 -+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -+"""Kimi-K3's model runner state on ROCm: the hybrid (KDA + MLA) state, plus the -+lifecycle of the specialized decode when ``decode_backend=latency_specialized`` -+selects it.""" -+ -+from vllm.v1.worker.gpu.model_states.mamba_hybrid import MambaHybridModelState -+ -+ -+class KimiK3ModelState(MambaHybridModelState): -+ """Builds the specialized decode once the weights are loaded, and prepares it -+ for every bound KV cache (the CUDA-graph memory-profiling cache, then the real -+ one).""" -+ -+ def __init__(self, vllm_config, model, encoder_cache, device): -+ super().__init__(vllm_config, model, encoder_cache, device) -+ if hasattr(model, "get_language_model"): -+ model = model.get_language_model() -+ self.language_model = model -+ if model.specialized_decode_config is not None: -+ model.build_specialized_decode() -+ -+ def initialize_kv_cache(self, kv_cache_config, block_tables): -+ super().initialize_kv_cache(kv_cache_config, block_tables) -+ decode = self.language_model.specialized_decode -+ if decode is not None: -+ decode.prepare_kv() -diff --git a/vllm/models/kimi_k3/amd/specialized_decode/mono/__init__.py b/vllm/models/kimi_k3/amd/specialized_decode/mono/__init__.py -new file mode 100644 -index 000000000..208f01a7c ---- /dev/null -+++ b/vllm/models/kimi_k3/amd/specialized_decode/mono/__init__.py -@@ -0,0 +1,2 @@ -+# SPDX-License-Identifier: Apache-2.0 -+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -diff --git a/vllm/models/kimi_k3/amd/specialized_decode/mono/abi.py b/vllm/models/kimi_k3/amd/specialized_decode/mono/abi.py -new file mode 100644 -index 000000000..7cf0c121f ---- /dev/null -+++ b/vllm/models/kimi_k3/amd/specialized_decode/mono/abi.py -@@ -0,0 +1,26 @@ -+# SPDX-License-Identifier: Apache-2.0 -+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -+"""The layout every MonoKernel's error words follow. -+ -+A rank's health words: a header whose word FAILED is set once any of its -+in-kernel waits gave up (sticky: a failed rank stays failed, and its later waits -+exit at once), then one record a CTA of the first wait it gave up on (flag, -+region id, pair index, tag, seen tag, CTA, clock, pad). -+ -+The output rank's flags: one int32 word a TP rank in the first PEER_HEAD bytes -+of the output rank's peer buffer. A rank whose wait gives up writes its word -+there, drains the store, then sets FAILED, so nothing it stores later can -+overtake the flag. The runner reads these words after every step. -+""" -+ -+MONO_ABI_VERSION = 1 -+ -+RECORD_WORDS = 8 -+FAILED = 0 -+PEER_HEAD = 256 # bytes: the output rank's flags, a word a rank -+WAIT_SECONDS = 10 # an in-kernel wait gives up after this long -+ -+ -+def health_words(ctas: int) -> int: -+ """A rank's health words for a grid of ``ctas``: the header, a record a CTA.""" -+ return (ctas + 1) * RECORD_WORDS -diff --git a/vllm/models/kimi_k3/amd/specialized_decode/mono/attention/__init__.py b/vllm/models/kimi_k3/amd/specialized_decode/mono/attention/__init__.py -new file mode 100644 -index 000000000..208f01a7c ---- /dev/null -+++ b/vllm/models/kimi_k3/amd/specialized_decode/mono/attention/__init__.py -@@ -0,0 +1,2 @@ -+# SPDX-License-Identifier: Apache-2.0 -+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -diff --git a/vllm/models/kimi_k3/amd/specialized_decode/mono/attention/back.py b/vllm/models/kimi_k3/amd/specialized_decode/mono/attention/back.py -new file mode 100644 -index 000000000..9efe5d7f4 ---- /dev/null -+++ b/vllm/models/kimi_k3/amd/specialized_decode/mono/attention/back.py -@@ -0,0 +1,178 @@ -+# SPDX-License-Identifier: Apache-2.0 -+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -+"""K2a ``attn_post``: one Kimi-K3 KDA layer from the gated core output to the -+MoE's input row -- o_proj, its TP all-reduce inside the kernel, and the MLP -+AttnRes seam. -+ -+One launch a layer and rank, ``BLOCKS`` x ``THREADS``: -+ -+ gemv CTAs 56 .. o_proj (448 x 16 rows, K = 1536 local): bf16 GEMV of the -+ core rows (plain loads: the previous launch wrote them) -> -+ bf16 partial -> pushed to every rank's AR region (system -+ scope, this rank's slot) -+ attn CTAs 0 .. 55 per 128 columns: the 8 ranks' partials summed in rank -+ order (fp32 -> bf16: the custom all-reduce's result) -> -+ LDS delta; then K1's AttnRes stages with the MLP's weights: -+ prefix + delta -> prefix out, the sources' softmax mix, -+ post_attention_layernorm -> the MoE input x -+ -+The AR region is double-buffered by the step epoch's parity: a rank a step -+ahead writes the other half, so no peer's unread partials are overwritten. -+""" -+ -+import flydsl.expr as fx -+from aiter.ops.flydsl.kernels import buffer_ops as bo -+from flydsl.expr import gpu, range_constexpr -+from flydsl.expr.typing import T -+ -+from vllm.models.kimi_k3.amd.specialized_decode.mono.attention.kda import ( -+ COLS, -+ HIDDEN, -+ LDS_PAD, -+ NA, -+ PROJ, -+ ROWS, -+) -+from vllm.models.kimi_k3.amd.specialized_decode.mono.common.ops import ( -+ CM_NT, -+ bf16_round, -+ mfma_bf16, -+ row_sum, -+ rsrc, -+ traced, -+) -+from vllm.models.kimi_k3.amd.specialized_decode.mono.common.plan import ( -+ BLOCKS, -+ THREADS, -+ WAVES, -+ pair_layout, -+) -+from vllm.models.kimi_k3.amd.specialized_decode.mono.common.ranks import sum_partials -+ -+TP = 8 -+OROWS = HIDDEN # o_proj rows (full hidden, a rank's K slice) -+OTASKS = OROWS // ROWS # 448 -+OK = PROJ # K = 1536 a rank -+OKCH = OK // 128 # 12 chunks -+GEMV0 = NA # the first gemv CTA -+GEMV_N = BLOCKS - NA # 200 -+ -+_STREAM = fx.Stream(None) -+ -+ -+def ar_half_pairs(s): -+ """One parity's AR region: ``[src rank][token][column pair]``.""" -+ return TP * s * HIDDEN // 2 -+ -+ -+def scratch_layout(s: int, nblocks: int) -> dict: -+ ns = nblocks + 1 -+ pairs = { -+ "part": NA * s * ns * 2, -+ "wgt": s * (ns + 1), -+ "msq": NA * s, -+ "xrdy": NA, -+ } -+ return pair_layout(pairs.items()) -+ -+ -+def o_loads(c, rg): -+ """Rows 16 rg .. of o_proj, this wave's K chunks (wave, wave + 8 < 12): a lane -+ 64 B of its row a chunk.""" -+ lane, wave = c["lane"], c["wave"] -+ r_w = rsrc(c["args"]["w_o"]) -+ row = rg * ROWS + lane % ROWS -+ g = lane // ROWS -+ out = [] -+ for i in range(2): -+ ch = fx.min(wave + WAVES * i, OKCH - 1) -+ base = (row * OK + ch * 128 + g * 8) // 2 -+ out.append( -+ [ -+ fx.Vector( -+ bo.buffer_load( -+ r_w, -+ base + 16 * q, -+ vec_width=4, -+ dtype=T.i32, -+ cache_modifier=CM_NT, -+ ) -+ ) -+ for q in range(4) -+ ] -+ ) -+ return out -+ -+ -+@traced -+def load_core(c): -+ """The core rows (S x 1536 bf16) -> LDS, rows padded.""" -+ s, tid = c["S"], c["tid"] -+ words = s * OK // 8 -+ row = OK + LDS_PAD -+ for i in range_constexpr((words + THREADS - 1) // THREADS): -+ e = fx.min(tid + THREADS * i, words - 1) -+ v = fx.Vector( -+ bo.buffer_load(rsrc(c["args"]["core"]), e * 4, vec_width=4, dtype=T.i32) -+ ) -+ t = e // (OK // 8) -+ k = e % (OK // 8) * 8 -+ if tid + THREADS * i < words: -+ fx.ptr_store(v, c["xl"] + (t * row + k) // 2) -+ gpu.barrier() -+ -+ -+@traced -+def stage_o(c, rg, ops): -+ """Rows 16 rg ..: GEMV -> bf16 partials -> every rank's AR region, this -+ rank's slot (thread: token, row pair; a put a peer).""" -+ s, tid, lane, wave, red = c["S"], c["tid"], c["lane"], c["wave"], c["red"] -+ g = lane // ROWS -+ t = fx.min(lane % ROWS, s - 1) -+ row = OK + LDS_PAD -+ acc = fx.Vector.filled(4, 0.0, fx.Float32) -+ for i in range_constexpr(2): -+ ch = wave + WAVES * i -+ live = ch < OKCH -+ k = fx.min(ch, OKCH - 1) * 128 + g * 8 -+ for q in range_constexpr(4): -+ xb = fx.Vector( -+ fx.ptr_load( -+ c["xl"] + (t * row + k + 32 * q) // 2, -+ result_type=fx.Vector.make_type(4, fx.Int32), -+ ) -+ ) -+ xb = fx.Vector.from_elements( -+ [live.select(xb[d], fx.Int32(0)) for d in range(4)], fx.Int32 -+ ) -+ acc = mfma_bf16( -+ ops[i][q].bitcast(fx.BFloat16), xb.bitcast(fx.BFloat16), acc -+ ) -+ fx.ptr_store(acc, red + (wave * 64 + lane) * 4) -+ gpu.barrier() -+ if tid < s * (ROWS // 2): -+ tt = tid // (ROWS // 2) -+ rp = tid % (ROWS // 2) -+ v0 = bf16_round(row_sum(red, 2 * rp, tt)) -+ v1 = bf16_round(row_sum(red, 2 * rp + 1, tt)) -+ at = (c["rank"] * s + tt) * HIDDEN + rg * ROWS + 2 * rp -+ for p in range_constexpr(TP): -+ c["put_bf"](c["ar_peer"][p], at, [v0, v1]) -+ gpu.barrier() -+ -+ -+@traced -+def stage_reduce(c, task): -+ """Columns 128 task ..: the 8 ranks' partials, rank order, -> bf16 delta in -+ LDS (thread: token, column pair).""" -+ s, tid = c["S"], c["tid"] -+ if tid < s * (COLS // 2): -+ t = tid // (COLS // 2) -+ cp = tid % (COLS // 2) -+ col = task * COLS + 2 * cp -+ lo, hi = sum_partials( -+ c["poll"], c["ar_own"], lambda src: ((src * s + t) * HIDDEN + col) // 2, TP -+ ) -+ fx.ptr_store(bf16_round(lo), c["delta_lds"] + (t * COLS + 2 * cp)) -+ fx.ptr_store(bf16_round(hi), c["delta_lds"] + (t * COLS + 2 * cp + 1)) -+ gpu.barrier() -diff --git a/vllm/models/kimi_k3/amd/specialized_decode/mono/attention/kda.py b/vllm/models/kimi_k3/amd/specialized_decode/mono/attention/kda.py -new file mode 100644 -index 000000000..43a624f7c ---- /dev/null -+++ b/vllm/models/kimi_k3/amd/specialized_decode/mono/attention/kda.py -@@ -0,0 +1,1021 @@ -+# SPDX-License-Identifier: Apache-2.0 -+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -+"""K1 ``kda_pre``: one Kimi-K3 KDA layer from the AttnRes seam to the gated core -+output, for a spec-verify step (R requests x L tokens, S = R L rows). -+ -+One launch a layer, ``BLOCKS`` CTAs x ``THREADS`` threads, in three roles so that -+no CTA that polls a mailbox has weight loads in flight (a poll's load retires -+behind every earlier load): -+ -+ attn CTAs 0 .. 55 slice (128 hidden columns each): prefix (+ delta) -> -+ bf16 prefix out, per-source partial sum of squares and -+ qk dot (PART); gate (token t on CTA t): softmax weights -+ (WGT); mix: the softmax mix, its partial sum of squares -+ (MSQ), every slice's -> output RMSNorm -> bf16 x rows -+ (plain device-scope stores, an XRDY flag a slice) -+ gemv CTAs 56 .. in_proj, two 16-row groups each: the first group's -+ weights issued at kernel start, the second's as x lands -+ -> q|k|v|g|f_a|beta (PROJ) -+ kda tasks (R x 12 heads x 8 v-chunks) on CTAs 0 ..: inputs that -+ need no hand-off loaded first (metadata, conv window, -+ state rows, f_b rows), then conv update of the owned -+ channels (q / k out through QK), f_b MFMA + lower-bound -+ gate, l2norm, the recurrence on 16 state rows (a state a -+ token written), partial norms (ONRM), gated RMSNorm -> -+ core_out -+ -+Rounding follows the original path where it is elementwise (the AttnRes prefix -+add, every bf16 hand-off between the original kernels); split-K reductions -+differ only in their order. -+""" -+ -+import math -+from dataclasses import dataclass -+ -+import flydsl.compiler as flyc -+import flydsl.expr as fx -+import torch -+from aiter.ops.flydsl.kernels import buffer_ops as bo -+from flydsl.expr import const_expr, gpu, range_constexpr -+from flydsl.expr.typing import Int32, Int64, T -+ -+from vllm.models.kimi_k3.amd.specialized_decode.mono.common.debug import ( -+ region_ids, -+) -+from vllm.models.kimi_k3.amd.specialized_decode.mono.common.ops import ( -+ CM_DEV, -+ CM_NT, -+ bf_hi, -+ bf_lo, -+ butterfly, -+ hw_exp2, -+ hw_rsq, -+ ld_i32, -+ mfma_bf16, -+ row_sum, -+ rsrc, -+ traced, -+ uniform, -+ wave_max, -+) -+from vllm.models.kimi_k3.amd.specialized_decode.mono.common.plan import ( -+ BLOCKS, -+ THREADS, -+ WAVES, -+ KernelAbi, -+ first_task, -+ key_tuple, -+ pair_layout, -+ source_digest, -+) -+from vllm.models.kimi_k3.amd.specialized_decode.mono.common.sync import ( -+ Mailbox, -+ publish, -+ sreg, -+) -+ -+HIDDEN = 7168 -+NA = 56 # attn CTAs -+COLS = HIDDEN // NA # 128 columns a slice -+QC = COLS // 4 # threads a (token, source) row: 4 columns each -+GSL = 7 # slices a gate / mix-norm thread sums (NA = 8 x 7) -+HD = 128 -+NH = 12 # local heads at TP8 -+PROJ = NH * HD # 1536 -+QKV = 3 * PROJ -+G2_0 = QKV -+FA_0 = G2_0 + PROJ # 6144 -+BETA_0 = FA_0 + HD # 6272 -+NPROJ = 6288 # in_proj rows a rank holds (6284 + 4 pad) -+ROWS = 16 -+GEMV_TASKS = NPROJ // ROWS # 393 -+GEMV_CTAS = (GEMV_TASKS + 1) // 2 # 197 -+KCH = 128 # K a chunk: 4 MFMAs of 32 -+NKCH = HIDDEN // KCH # 56 -+PER_WAVE = NKCH // WAVES # 7 -+VCH = 16 # state rows a KDA task -+NVCH = HD // VCH # 8 -+CONV_W = 4 -+LOG2E = 1.0 / math.log(2.0) -+LDS_PAD = 8 # bf16 a padded x row -+ -+assert NA + GEMV_CTAS <= BLOCKS and NA == 8 * GSL -+ -+_STREAM = fx.Stream(None) -+REGIONS = ("part", "wgt", "msq", "xrdy", "proj", "qk", "onrm") -+# every mono source file: in every build's JIT cache key -+_SOURCES = source_digest(".") -+ -+ -+@dataclass(frozen=True) -+class KdaPreBuild: -+ tokens: int # S -+ qlen: int # L, tokens a request -+ nblocks: int # AttnRes block sources (the prefix is one more) -+ delta: bool -+ state_len: int # conv state entries: CONV_W - 1 + num_spec -+ eps: float = 1e-5 # AttnRes norm -+ out_eps: float = 1e-5 # input_layernorm -+ onorm_eps: float = 1e-5 -+ lower_bound: float = -5.0 -+ write_idx: int = -1 # a block-write layer's block: the updated prefix stored there -+ kda0: int = 0 # the first KDA task's CTA -+ -+ -+ABI = KernelAbi( -+ ( -+ "prefix", -+ "delta", -+ "blocks", -+ "blk_sm", -+ "blk_sr", -+ "ares_nw", -+ "ares_qk", -+ "in_nw", -+ "w_in", -+ "w_fb", -+ "conv_w", -+ "conv_st", -+ "cs_seq", -+ "cs_dim", -+ "cs_tok", -+ "a_log", -+ "dt_bias", -+ "on_w", -+ "rstate", -+ "rs_seq", -+ "st_idx", -+ "st_stride", -+ "num_acc", -+ "core_out", -+ "scratch", -+ "layer", -+ "epoch", -+ "health", -+ "alarm", -+ ) -+) -+ -+ -+def scratch_layout(key: KdaPreBuild) -> dict: -+ s, ns = key.tokens, key.nblocks + 1 -+ assert s % key.qlen == 0 -+ pairs = { -+ "part": NA * s * ns * 2, -+ "wgt": s * (ns + 1), -+ "msq": NA * s, -+ "xrdy": NA, -+ "proj": s * NPROJ, -+ "qk": s * NH * 2 * HD, -+ "onrm": s * NH * NVCH, -+ } -+ lay = pair_layout(pairs.items()) -+ end = max(o + n for o, n in lay.values()) -+ # the plain bf16 x rows (not pairs), 16 B aligned -+ lay["x"] = ((end + 15) // 16 * 16, s * HIDDEN * 2) -+ return lay -+ -+ -+def scratch_bytes(key: KdaPreBuild) -> int: -+ return max(o + n for o, n in scratch_layout(key).values()) -+ -+ -+TAG_SLOTS = 256 # a step's mailbox tags: epoch x 256 + a launch's slot -+ -+ -+def step_tag(epoch, slot): -+ """This launch's mailbox tag: the step's epoch (a device counter the runner -+ bumps once a step, in the graph) and the launch's slot, so no step's or -+ other launch's pair ever passes for this one's -- nothing is zeroed.""" -+ e = uniform(ld_i32(epoch, 0)) -+ return (e % (1 << 22)) * TAG_SLOTS + slot + 1 -+ -+ -+def ld_bf(ptr, i, cm=0): -+ return fx.Float32( -+ fx.BFloat16( -+ bo.buffer_load(rsrc(ptr), i, vec_width=1, dtype=T.bf16, cache_modifier=cm) -+ ) -+ ) -+ -+ -+def ld_bf4(ptr, i): -+ """bf16 elements i .. i + 4 (i a multiple of 4) as four f32.""" -+ w = fx.Vector(bo.buffer_load(rsrc(ptr), i // 2, vec_width=2, dtype=T.i32)) -+ return [bf_lo(w[0]), bf_hi(w[0]), bf_lo(w[1]), bf_hi(w[1])] -+ -+ -+def bf4_words(v): -+ """Four f32 -> two words of packed bf16.""" -+ return fx.Vector.from_elements(v, fx.Float32).to(fx.BFloat16).bitcast(fx.Int32) -+ -+ -+def ld_f32(ptr, i): -+ return fx.Float32(bo.buffer_load(rsrc(ptr), i, vec_width=1, dtype=T.f32)) -+ -+ -+def bf16r(v): -+ return fx.Float32(fx.Float32(v).to(fx.BFloat16)) -+ -+ -+def sigmoid(z): -+ return fx.Float32(1.0) / (fx.Float32(1.0) + hw_exp2(z * fx.Float32(-LOG2E))) -+ -+ -+def f32_of(w): -+ return w.bitcast(fx.Float32) -+ -+ -+def sum_polled(vals): -+ acc = f32_of(vals[0][0]) -+ for k in range(1, len(vals)): -+ acc = acc + f32_of(vals[k][0]) -+ return acc -+ -+ -+# ---------------------------------------------------------------- attn CTAs -+@traced -+def stage_slice(c, task): -+ """Columns 128 task ..: every token's sources into LDS ``vals`` (4 columns a -+ thread), the prefix (+ delta) rounded to bf16 and stored back, per-(token, -+ source) partial sum of squares and qk dot -> PART.""" -+ s, ns, nb, tid, lane = c["S"], c["NS"], c["NB"], c["tid"], c["lane"] -+ a = c["args"] -+ c0 = task * COLS -+ n = s * ns * QC # threads' work items -+ for i in range_constexpr((n + THREADS - 1) // THREADS): -+ e = fx.min(tid + THREADS * i, n - 1) -+ live = (tid + THREADS * i) < n -+ row = e // QC # (token, source) -+ t = row // ns -+ src = row % ns -+ col = c0 + (e % QC) * 4 -+ is_pre = src == nb -+ blk = ld_bf4( -+ a["blocks"], t * a["blk_sm"] + fx.min(src, nb - 1) * a["blk_sr"] + col -+ ) -+ pre = ld_bf4(a["prefix"], t * HIDDEN + col) -+ if const_expr(c["delta"]): -+ if const_expr(c.get("delta_lds") is not None): -+ # the delta reduced into LDS first (K2a: the all-reduced o_proj) -+ dl = [ -+ fx.ptr_load(c["delta_lds"] + (t * COLS + (e % QC) * 4 + j)) -+ for j in range(4) -+ ] -+ else: -+ dl = ld_bf4(a["delta"], t * HIDDEN + col) -+ if const_expr(c.get("reset", False)): -+ # a block-write layer's MLP seam: the prefix restarts at the -+ # attention's output (the input prefix is not read) -+ pre = dl -+ else: -+ pre = [bf16r(pre[j] + dl[j]) for j in range(4)] -+ if live & is_pre: -+ bo.buffer_store( -+ bf4_words(pre), rsrc(a["prefix"]), (t * HIDDEN + col) // 2 -+ ) -+ if const_expr(c.get("write_idx", -1) >= 0): # noqa: SIM102 (traced: `and` does not combine traced values) -+ # the block-write layer's attention seam: the updated prefix is -+ # block write_idx -+ if live & is_pre: -+ bo.buffer_store( -+ bf4_words(pre), -+ rsrc(a["blocks"]), -+ (t * a["blk_sm"] + c["write_idx"] * a["blk_sr"] + col) // 2, -+ ) -+ v = [ -+ live.select(is_pre.select(pre[j], blk[j]), fx.Float32(0.0)) -+ for j in range(4) -+ ] -+ w = ld_bf4(a["ares_nw"], col) -+ q = ld_bf4(a["ares_qk"], col) -+ # a thread past the end holds a clamped copy of the last element: it -+ # must not store (it raced the owner with zeros) -+ if live: -+ for j in range_constexpr(4): -+ fx.ptr_store(v[j], c["vals"] + (row * COLS + (e % QC) * 4 + j)) -+ sq = v[0] * v[0] + v[1] * v[1] + v[2] * v[2] + v[3] * v[3] -+ dt = ( -+ v[0] * (w[0] * q[0]) -+ + v[1] * (w[1] * q[1]) -+ + v[2] * (w[2] * q[2]) -+ + v[3] * (w[3] * q[3]) -+ ) -+ sq = butterfly(sq, (16, 8, 4, 2, 1)) -+ dt = butterfly(dt, (16, 8, 4, 2, 1)) -+ if live & (lane % 32 == 0): -+ at = ((task * s + t) * ns + src) * 2 -+ c["put"](c["part"], at, sq) -+ c["put"](c["part"], at + 1, dt) -+ gpu.barrier() -+ -+ -+@traced -+def stage_gate(c, t): -+ """Token t: the 56 slices' partials -> source weights e_s and their sum.""" -+ s, ns, tid, lane, wave, red = ( -+ c["S"], -+ c["NS"], -+ c["tid"], -+ c["lane"], -+ c["wave"], -+ c["red"], -+ ) -+ nq = ns * 2 -+ if tid < 8 * nq: -+ g = tid // nq -+ q = tid % nq -+ vals = c["poll"]( -+ [(c["part"], (((g * GSL + k) * s + t) * ns) * 2 + q, 1) for k in range(GSL)] -+ ) -+ fx.ptr_store(sum_polled(vals), red + (g * nq + q)) -+ gpu.barrier() -+ if wave == 0: -+ src = fx.min(lane, ns - 1) -+ sq = fx.ptr_load(red + src * 2) -+ dt = fx.ptr_load(red + src * 2 + 1) -+ for g in range_constexpr(1, 8): -+ sq = sq + fx.ptr_load(red + (g * nq + src * 2)) -+ dt = dt + fx.ptr_load(red + (g * nq + src * 2 + 1)) -+ rstd = hw_rsq(sq * fx.Float32(1.0 / HIDDEN) + fx.Float32(c["eps"])) -+ logit = dt * rstd * fx.Float32(LOG2E) -+ logit = (lane < ns).select(logit, fx.Float32(-3.0e38)) -+ m = wave_max(logit) -+ ex = (lane < ns).select(hw_exp2(logit - m), fx.Float32(0.0)) -+ den = butterfly(ex, (32, 16, 8, 4, 2, 1)) -+ if lane < ns: -+ c["put"](c["wgt"], t * (ns + 1) + lane, ex) -+ if lane == 0: -+ c["put"](c["wgt"], t * (ns + 1) + ns, den) -+ gpu.barrier() -+ -+ -+@traced -+def stage_mix(c, task): -+ """Columns 128 task ..: the softmax mix of ``vals`` (4 columns a thread), its -+ sum of squares (MSQ), every slice's -> rstd, then bf16 x = mixed rstd w -> -+ the plain x rows, and the slice's XRDY flag.""" -+ s, ns, tid, lane, red = c["S"], c["NS"], c["tid"], c["lane"], c["red"] -+ a = c["args"] -+ c0 = task * COLS -+ nw = s * (ns + 1) -+ if tid < nw: -+ v = c["poll"]([(c["wgt"], tid, 1)])[0][0] -+ fx.ptr_store(f32_of(v), red + tid) -+ gpu.barrier() -+ mine = tid < s * QC -+ t = fx.min(tid // QC, s - 1) -+ cc = (tid % QC) * 4 -+ mixed = [] -+ for j in range_constexpr(4): -+ acc = fx.Float32(0.0) -+ for src in range_constexpr(ns): -+ acc = acc + fx.ptr_load(red + (t * (ns + 1) + src)) * fx.ptr_load( -+ c["vals"] + ((t * ns + src) * COLS + cc + j) -+ ) -+ mixed.append(acc / fx.ptr_load(red + (t * (ns + 1) + ns))) -+ sq = ( -+ mixed[0] * mixed[0] -+ + mixed[1] * mixed[1] -+ + mixed[2] * mixed[2] -+ + mixed[3] * mixed[3] -+ ) -+ sq = butterfly(sq, (16, 8, 4, 2, 1)) -+ if mine & (lane % 32 == 0): -+ c["put"](c["msq"], task * s + t, sq) -+ gpu.barrier() -+ if tid < 8 * s: -+ g = tid // s -+ tt = tid % s -+ vals = c["poll"]([(c["msq"], (g * GSL + k) * s + tt, 1) for k in range(GSL)]) -+ fx.ptr_store(sum_polled(vals), red + (128 + g * s + tt)) -+ gpu.barrier() -+ tot = fx.ptr_load(red + (128 + t)) -+ for g in range_constexpr(1, 8): -+ tot = tot + fx.ptr_load(red + (128 + g * s + t)) -+ r = hw_rsq(tot * fx.Float32(1.0 / HIDDEN) + fx.Float32(c["out_eps"])) -+ wn = ld_bf4(a["in_nw"], c0 + cc) -+ xw = bf4_words([mixed[j] * r * wn[j] for j in range(4)]) -+ if mine: -+ at = (t * HIDDEN + c0 + cc) // 2 -+ bo.buffer_store(xw, rsrc(c["xbuf"]), at, cache_modifier=CM_DEV) -+ publish(c["put"], c["xrdy"], task, fx.Int32(1), tid == 0) -+ gpu.barrier() -+ -+ -+# ---------------------------------------------------------------- gemv CTAs -+def gemv_loads(c, rg, chunks): -+ """This wave's weight operands of rows 16 rg .. for its K chunks ``chunks`` -+ (indices into wave, wave + 8, ..): a lane 64 contiguous bytes of its row a -+ chunk (4 dwordx4). ``rg`` past the last group loads the last group.""" -+ lane, wave = c["lane"], c["wave"] -+ r_w = rsrc(c["args"]["w_in"]) -+ row = fx.min(rg, GEMV_TASKS - 1) * ROWS + lane % ROWS -+ g = lane // ROWS -+ out = [] -+ for i in chunks: -+ k = (wave + WAVES * i) * KCH + g * 8 -+ base = (row * HIDDEN + k) // 2 # dwords -+ out.append( -+ [ -+ fx.Vector( -+ bo.buffer_load( -+ r_w, -+ base + 16 * q, -+ vec_width=4, -+ dtype=T.i32, -+ cache_modifier=CM_NT, -+ ) -+ ) -+ for q in range(4) -+ ] -+ ) -+ return out -+ -+ -+@traced -+def load_x(c): -+ """Every slice's x ready -> the bf16 x rows into LDS (rows padded).""" -+ s, tid = c["S"], c["tid"] -+ if tid < NA: -+ c["poll"]([(c["xrdy"], tid, 1)]) -+ gpu.barrier() -+ words = s * HIDDEN // 8 # 16 B chunks -+ row = HIDDEN + LDS_PAD -+ for i in range_constexpr((words + THREADS - 1) // THREADS): -+ e = fx.min(tid + THREADS * i, words - 1) -+ v = fx.Vector( -+ bo.buffer_load( -+ rsrc(c["xbuf"]), e * 4, vec_width=4, dtype=T.i32, cache_modifier=CM_DEV -+ ) -+ ) -+ t = e // (HIDDEN // 8) -+ k = e % (HIDDEN // 8) * 8 -+ if tid + THREADS * i < words: -+ fx.ptr_store(v, c["xl"] + (t * row + k) // 2) -+ gpu.barrier() -+ -+ -+@traced -+def stage_gemv(c, rg, ops): -+ """Rows 16 rg .. of in_proj against the LDS x -> bf16 PROJ.""" -+ s, tid, lane, wave, red = c["S"], c["tid"], c["lane"], c["wave"], c["red"] -+ g = lane // ROWS -+ t = fx.min(lane % ROWS, s - 1) -+ row = HIDDEN + LDS_PAD -+ acc = fx.Vector.filled(4, 0.0, fx.Float32) -+ for i in range_constexpr(PER_WAVE): -+ k = (wave + WAVES * i) * KCH + g * 8 -+ for q in range_constexpr(4): -+ xb = fx.Vector( -+ fx.ptr_load( -+ c["xl"] + (t * row + k + 32 * q) // 2, -+ result_type=fx.Vector.make_type(4, fx.Int32), -+ ) -+ ) -+ acc = mfma_bf16( -+ ops[i][q].bitcast(fx.BFloat16), xb.bitcast(fx.BFloat16), acc -+ ) -+ fx.ptr_store(acc, red + (wave * 64 + lane) * 4) -+ gpu.barrier() -+ if (tid < ROWS * s) & (rg < GEMV_TASKS): -+ r = tid % ROWS -+ tt = tid // ROWS -+ v = bf16r(row_sum(red, r, tt)) -+ c["put"](c["proj"], tt * NPROJ + rg * ROWS + r, v) -+ gpu.barrier() -+ -+ -+# ---------------------------------------------------------------- kda -+def kda_stage(c, req, h, ch): -+ """The task's inputs that need no hand-off, issued before its first poll: -+ metadata, the owned channels' conv window / weights, the initial state rows, -+ the f_b rows of this wave and the per-head constants.""" -+ L, sl, tid, lane, wave = c["L"], c["SL"], c["tid"], c["lane"], c["wave"] -+ a = c["args"] -+ k = {} -+ # every value a buffer resource is built on wave-uniform (else a waterfall) -+ k["acc_idx"] = uniform(ld_i32(a["num_acc"], req)) - 1 -+ k["slot_c"] = uniform(ld_i32(a["st_idx"], req * a["st_stride"])) -+ k["slot0"] = uniform(ld_i32(a["st_idx"], req * a["st_stride"] + k["acc_idx"])) -+ k["slots"] = [ -+ uniform(ld_i32(a["st_idx"], req * a["st_stride"] + j)) for j in range(L) -+ ] -+ # conv: thread < 48 owns channel (part, d) -+ ct = fx.min(tid, 3 * VCH - 1) -+ part = ct // VCH -+ chan = part * PROJ + h * HD + ch * VCH + ct % VCH -+ cs = a["conv_st"] + fx.Int64(fx.max(k["slot_c"], 0)) * fx.Int64(a["cs_seq"]) * 2 -+ co = chan * a["cs_dim"] # the channel's first entry (bf16 elements) -+ k["cs"], k["co"], k["chan"], k["part"] = cs, co, chan, part -+ k["cw"] = fx.Vector( -+ bo.buffer_load(rsrc(a["conv_w"]), chan * CONV_W, vec_width=4, dtype=T.f32) -+ ) -+ k["win"] = [ -+ ld_bf(cs, co + (k["acc_idx"] + j) * a["cs_tok"]) for j in range(CONV_W - 1) -+ ] -+ k["old"] = [ -+ ld_bf(cs, co + (k["acc_idx"] + 1 + j) * a["cs_tok"]) for j in range(sl - L) -+ ] -+ # recurrence: thread -> state row v = tid / 32, k columns (tid % 32) 4 .. -+ vr = ch * VCH + tid // 32 -+ k0 = (tid % 32) * 4 -+ k["head_off"] = (h * HD + vr) * HD + k0 -+ base0 = a["rstate"] + fx.Int64(fx.max(k["slot0"], 0)) * fx.Int64(a["rs_seq"]) * 4 -+ k["st"] = fx.Vector( -+ bo.buffer_load(rsrc(base0), k["head_off"], vec_width=4, dtype=T.f32) -+ ) -+ # f_b: wave w's rows 16 w .. of head h, a lane 64 B of its row -+ fr = h * HD + wave * 16 + lane % 16 -+ fb = (fr * HD + (lane // 16) * 32) // 2 -+ k["fb"] = [ -+ fx.Vector(bo.buffer_load(rsrc(a["w_fb"]), fb + 4 * q, vec_width=4, dtype=T.i32)) -+ for q in range(4) -+ ] -+ k["aa"] = hw_exp2(ld_f32(a["a_log"], h) * fx.Float32(LOG2E)) -+ k["dtb"] = [ -+ ld_f32(a["dt_bias"], h * HD + wave * 16 + 4 * (lane // 16) + i) -+ for i in range(4) -+ ] -+ return k -+ -+ -+@traced -+def kda_conv(c, k, req, h, ch): -+ """Owned channels: the spec conv update of request ``req``'s L tokens -> new -+ conv state; q / k out through QK, v to LDS.""" -+ L, sl, tid = c["L"], c["SL"], c["tid"] -+ a = c["args"] -+ if tid < 3 * VCH: -+ xs = c["poll"]( -+ [(c["proj"], (req * L + j) * NPROJ + k["chan"], 1) for j in range(L)] -+ ) -+ xv = [f32_of(xs[j][0]) for j in range(L)] -+ rs = rsrc(k["cs"]) -+ if k["slot_c"] > 0: -+ for j in range_constexpr(sl): -+ nv = k["old"][j] if j < sl - L else xv[j - (sl - L)] -+ bo.buffer_store( -+ fx.Float32(nv).to(fx.BFloat16), rs, k["co"] + j * a["cs_tok"] -+ ) -+ seq = k["win"] + xv -+ for j in range_constexpr(L): -+ acc = fx.Float32(0.0) -+ for q in range_constexpr(CONV_W): -+ acc = acc + k["cw"][q] * seq[j + q] -+ o = bf16r(acc * sigmoid(acc)) -+ if k["part"] < 2: -+ c["put"]( -+ c["qk"], -+ ((req * NH + h) * L + j) * 2 * HD -+ + k["part"] * HD -+ + ch * VCH -+ + tid % VCH, -+ o, -+ ) -+ if k["part"] == 2: -+ fx.ptr_store(o, c["vloc"] + (j * VCH + tid % VCH)) -+ gpu.barrier() -+ -+ -+@traced -+def kda_gates(c, k, req, h): -+ """f_a, beta and the head's q / k (one poll batch a thread); g1 = f_b f_a by -+ MFMA (a wave 16 rows) -> lower-bound gate -> decay in LDS; q / k l2-normed -+ (q scaled).""" -+ L, tid, lane, wave = c["L"], c["tid"], c["lane"], c["wave"] -+ n_fa = L * HD -+ n = n_fa + L * 2 * HD -+ per = (n + THREADS - 1) // THREADS -+ specs = [] -+ for i in range_constexpr(per): -+ e = fx.min(tid + THREADS * i, n - 1) -+ fa_at = (req * L + e // HD) * NPROJ + FA_0 + e % HD -+ qk_at = (req * NH + h) * L * 2 * HD + (e - n_fa) -+ specs.append((e, fa_at, qk_at)) -+ # two regions, two batches: the f_a pairs, then the q / k pairs -+ got_fa = c["poll"]( -+ [(c["proj"], fx.min(fa, (req * L + L) * NPROJ - 1), 1) for _, fa, _ in specs] -+ ) -+ got_qk = c["poll"]( -+ [(c["qk"], fx.max(qk, (req * NH + h) * L * 2 * HD), 1) for _, _, qk in specs] -+ ) -+ for i in range_constexpr(per): -+ e = specs[i][0] -+ if tid + THREADS * i < n_fa: -+ fx.ptr_store(f32_of(got_fa[i][0]), c["fa"] + e) -+ if (tid + THREADS * i >= n_fa) & (tid + THREADS * i < n): -+ fx.ptr_store(f32_of(got_qk[i][0]), c["qkl"] + (e - n_fa)) -+ if tid < L: -+ b = c["poll"]([(c["proj"], (req * L + tid) * NPROJ + BETA_0 + h, 1)])[0][0] -+ fx.ptr_store(sigmoid(f32_of(b)), c["beta"] + tid) -+ gpu.barrier() -+ # g1 rows 16 wave .. x tokens: A = f_b rows (staged), B = f_a [token][K] -+ t = fx.min(lane % 16, L - 1) -+ acc = fx.Vector.filled(4, 0.0, fx.Float32) -+ for q in range_constexpr(4): -+ kb = (lane // 16) * 32 + 8 * q -+ xs = [fx.ptr_load(c["fa"] + (t * HD + kb + e)) for e in range(8)] -+ xb = fx.Vector.from_elements(xs, fx.Float32).to(fx.BFloat16) -+ acc = mfma_bf16(k["fb"][q].bitcast(fx.BFloat16), xb, acc) -+ for i in range_constexpr(4): -+ kk = wave * 16 + 4 * (lane // 16) + i -+ g = bf16r(acc[i]) + k["dtb"][i] -+ gate = fx.Float32(c["lb"]) * sigmoid(k["aa"] * g) -+ if lane % 16 < L: -+ fx.ptr_store(hw_exp2(gate * fx.Float32(LOG2E)), c["decay"] + (t * HD + kk)) -+ # l2norm: q / k row r (2 L rows of 128) on wave r % 8 -+ for rr in range_constexpr((2 * L + WAVES - 1) // WAVES): -+ row = fx.min(wave + WAVES * rr, 2 * L - 1) -+ x0 = fx.ptr_load(c["qkl"] + (row * HD + lane)) -+ x1 = fx.ptr_load(c["qkl"] + (row * HD + 64 + lane)) -+ ss = butterfly(x0 * x0 + x1 * x1, (32, 16, 8, 4, 2, 1)) -+ sc = hw_rsq(ss + fx.Float32(1e-6)) -+ sc = (row % 2 == 0).select(sc * fx.Float32(HD**-0.5), sc) -+ if wave + WAVES * rr < 2 * L: -+ fx.ptr_store(x0 * sc, c["qkl"] + (row * HD + lane)) -+ fx.ptr_store(x1 * sc, c["qkl"] + (row * HD + 64 + lane)) -+ gpu.barrier() -+ -+ -+@traced -+def kda_recur(c, k, req, h, ch): -+ """State rows 16 ch .. of head h: the recurrence over the L tokens, a state -+ written a token, bf16 outputs to LDS ``ol`` and their partial norms (ONRM).""" -+ L, tid, lane = c["L"], c["tid"], c["lane"] -+ a = c["args"] -+ v = tid // 32 -+ k0 = (tid % 32) * 4 -+ live0 = k["slot0"] > 0 -+ stv = [live0.select(k["st"][e], fx.Float32(0.0)) for e in range(4)] -+ for j in range_constexpr(L): -+ qv = [fx.ptr_load(c["qkl"] + ((2 * j) * HD + k0 + e)) for e in range(4)] -+ kv = [fx.ptr_load(c["qkl"] + ((2 * j + 1) * HD + k0 + e)) for e in range(4)] -+ dk = [fx.ptr_load(c["decay"] + (j * HD + k0 + e)) for e in range(4)] -+ stv = [stv[e] * dk[e] for e in range(4)] -+ p = stv[0] * kv[0] + stv[1] * kv[1] + stv[2] * kv[2] + stv[3] * kv[3] -+ p = butterfly(p, (16, 8, 4, 2, 1)) -+ vn = (fx.ptr_load(c["vloc"] + (j * VCH + v)) - p) * fx.ptr_load(c["beta"] + j) -+ stv = [stv[e] + vn * kv[e] for e in range(4)] -+ o = stv[0] * qv[0] + stv[1] * qv[1] + stv[2] * qv[2] + stv[3] * qv[3] -+ o = bf16r(butterfly(o, (16, 8, 4, 2, 1))) -+ o = live0.select(o, fx.Float32(0.0)) -+ slot = k["slots"][j] -+ if (slot > 0) & live0: -+ dst = rsrc(a["rstate"] + fx.Int64(slot) * fx.Int64(a["rs_seq"]) * 4) -+ bo.buffer_store( -+ fx.Vector.from_elements(stv, fx.Float32), dst, k["head_off"] -+ ) -+ if lane % 32 == 0: -+ fx.ptr_store(o, c["ol"] + (j * VCH + v)) -+ gpu.barrier() -+ if tid < L: -+ ss = fx.Float32(0.0) -+ for e in range_constexpr(VCH): -+ x = fx.ptr_load(c["ol"] + (tid * VCH + e)) -+ ss = ss + x * x -+ c["put"](c["onrm"], ((req * NH + h) * L + tid) * NVCH + ch, ss) -+ gpu.barrier() -+ -+ -+@traced -+def kda_norm(c, req, h, ch): -+ """Rows 16 ch ..: gated RMSNorm over the head's 128 (every chunk's partial) -+ -> core_out.""" -+ L, tid = c["L"], c["tid"] -+ a = c["args"] -+ if tid < L * VCH: -+ j = tid // VCH -+ e = tid % VCH -+ vr = ch * VCH + e -+ got = c["poll"]( -+ [(c["onrm"], ((req * NH + h) * L + j) * NVCH + q, 1) for q in range(NVCH)] -+ + [(c["proj"], (req * L + j) * NPROJ + G2_0 + h * HD + vr, 1)] -+ ) -+ ss = sum_polled(got[:NVCH]) -+ r = hw_rsq(ss * fx.Float32(1.0 / HD) + fx.Float32(c["onorm_eps"])) -+ y = ( -+ fx.ptr_load(c["ol"] + (j * VCH + e)) -+ * r -+ * ld_bf(a["on_w"], vr) -+ * sigmoid(f32_of(got[NVCH][0])) -+ ) -+ bo.buffer_store( -+ y.to(fx.BFloat16), rsrc(a["core_out"]), (req * L + j) * PROJ + h * HD + vr -+ ) -+ gpu.barrier() -+ -+ -+@traced -+def kda_tasks(c): -+ s, L, bid = c["S"], c["L"], c["bid"] -+ for task in range(first_task(bid, c["kda0"]), (s // L) * NH * NVCH, BLOCKS): -+ req = task // (NH * NVCH) -+ h = (task // NVCH) % NH -+ ch = task % NVCH -+ k = kda_stage(c, req, h, ch) -+ kda_conv(c, k, req, h, ch) -+ kda_gates(c, k, req, h) -+ kda_recur(c, k, req, h, ch) -+ kda_norm(c, req, h, ch) -+ -+ -+# ---------------------------------------------------------------- kernel -+_BUILDS: dict = {} -+ -+ -+def build(key: KdaPreBuild): -+ if key in _BUILDS: -+ return _BUILDS[key] -+ s, L = key.tokens, key.qlen -+ ns = key.nblocks + 1 -+ lay = scratch_layout(key) -+ xrow = HIDDEN + LDS_PAD -+ assert s * xrow * 2 <= 136 * 1024, "x rows exceed LDS" -+ assert L <= 16 and key.nblocks >= 1 -+ -+ @fx.struct -+ class AttnLds: -+ vals: fx.Array[fx.Float32, s * ns * COLS, 16] -+ -+ @fx.struct -+ class GemvLds: -+ xl: fx.Array[fx.Int32, s * xrow // 2, 16] -+ -+ @fx.struct -+ class KdaLds: -+ fa: fx.Array[fx.Float32, L * HD, 16] -+ qkl: fx.Array[fx.Float32, L * 2 * HD, 16] -+ decay: fx.Array[fx.Float32, L * HD, 16] -+ vloc: fx.Array[fx.Float32, L * VCH, 16] -+ ol: fx.Array[fx.Float32, L * VCH, 16] -+ beta: fx.Array[fx.Float32, 16, 16] -+ -+ @fx.union -+ class RoleLds: -+ attn: AttnLds -+ gemv: GemvLds -+ kda: KdaLds -+ -+ @fx.struct -+ class Smem: -+ red: fx.Array[fx.Float32, WAVES * 64 * 4, 16] -+ -+ keyed = key_tuple(key, _SOURCES) -+ name = f"k3_mono_kda_pre_s{s}_l{L}_nb{key.nblocks}_d{int(key.delta)}" -+ -+ @flyc.kernel(name=name, known_block_size=[THREADS, 1, 1]) -+ def kda_pre( -+ prefix: Int64, -+ delta: Int64, -+ blocks: Int64, -+ blk_sm: Int32, -+ blk_sr: Int32, -+ ares_nw: Int64, -+ ares_qk: Int64, -+ in_nw: Int64, -+ w_in: Int64, -+ w_fb: Int64, -+ conv_w: Int64, -+ conv_st: Int64, -+ cs_seq: Int32, -+ cs_dim: Int32, -+ cs_tok: Int32, -+ a_log: Int64, -+ dt_bias: Int64, -+ on_w: Int64, -+ rstate: Int64, -+ rs_seq: Int32, -+ st_idx: Int64, -+ st_stride: Int32, -+ num_acc: Int64, -+ core_out: Int64, -+ scratch: Int64, -+ layer: Int32, -+ epoch: Int64, -+ health: Int64, -+ alarm: Int64, -+ ): -+ _ = keyed -+ tid = fx.thread_idx.x -+ bid = fx.block_idx.x -+ alloc = fx.SharedAllocator() -+ lds = alloc.allocate(Smem).peek() -+ role = alloc.allocate(RoleLds) -+ al, gl, kl = role.attn.peek(), role.gemv.peek(), role.kda.peek() -+ mb = Mailbox(step_tag(epoch, layer) - 1, health, alarm, region_ids(REGIONS)) -+ c = { -+ "S": s, -+ "L": L, -+ "NS": ns, -+ "NB": key.nblocks, -+ "SL": key.state_len, -+ "tid": tid, -+ "bid": bid, -+ "lane": tid % 64, -+ "wave": tid // 64, -+ "delta": key.delta, -+ "eps": key.eps, -+ "out_eps": key.out_eps, -+ "onorm_eps": key.onorm_eps, -+ "lb": key.lower_bound, -+ "kda0": key.kda0, -+ "write_idx": key.write_idx, -+ "put": mb.put, -+ "poll": mb.poll, -+ "red": lds.red.ptr, -+ "vals": al.vals.ptr, -+ "xl": gl.xl.ptr, -+ "fa": kl.fa.ptr, -+ "qkl": kl.qkl.ptr, -+ "decay": kl.decay.ptr, -+ "vloc": kl.vloc.ptr, -+ "ol": kl.ol.ptr, -+ "beta": kl.beta.ptr, -+ "xbuf": scratch + fx.Int64(lay["x"][0]), -+ "args": { -+ "prefix": prefix, -+ "delta": delta, -+ "blocks": blocks, -+ "blk_sm": blk_sm, -+ "blk_sr": blk_sr, -+ "ares_nw": ares_nw, -+ "ares_qk": ares_qk, -+ "in_nw": in_nw, -+ "w_in": w_in, -+ "w_fb": w_fb, -+ "conv_w": conv_w, -+ "conv_st": conv_st, -+ "cs_seq": cs_seq, -+ "cs_dim": cs_dim, -+ "cs_tok": cs_tok, -+ "a_log": a_log, -+ "dt_bias": dt_bias, -+ "on_w": on_w, -+ "rstate": rstate, -+ "rs_seq": rs_seq, -+ "st_idx": st_idx, -+ "st_stride": st_stride, -+ "num_acc": num_acc, -+ "core_out": core_out, -+ }, -+ } -+ for region in REGIONS: -+ c[region] = sreg(scratch, lay[region][0], region) -+ if bid < NA: -+ stage_slice(c, bid) -+ if bid < s: -+ stage_gate(c, bid) -+ stage_mix(c, bid) -+ if bid >= NA: -+ gi = bid - NA -+ rg0 = 2 * gi -+ rg1 = 2 * gi + 1 -+ ops0 = gemv_loads(c, rg0, range(PER_WAVE)) -+ load_x(c) -+ # the second group's first half in flight under the first's MFMAs -+ ops1 = gemv_loads(c, rg1, range(4)) -+ stage_gemv(c, rg0, ops0) -+ ops1 = ops1 + gemv_loads(c, rg1, range(4, PER_WAVE)) -+ stage_gemv(c, rg1, ops1) -+ kda_tasks(c) -+ -+ @flyc.jit -+ def launch( -+ prefix: Int64, -+ delta: Int64, -+ blocks: Int64, -+ blk_sm: Int32, -+ blk_sr: Int32, -+ ares_nw: Int64, -+ ares_qk: Int64, -+ in_nw: Int64, -+ w_in: Int64, -+ w_fb: Int64, -+ conv_w: Int64, -+ conv_st: Int64, -+ cs_seq: Int32, -+ cs_dim: Int32, -+ cs_tok: Int32, -+ a_log: Int64, -+ dt_bias: Int64, -+ on_w: Int64, -+ rstate: Int64, -+ rs_seq: Int32, -+ st_idx: Int64, -+ st_stride: Int32, -+ num_acc: Int64, -+ core_out: Int64, -+ scratch: Int64, -+ layer: Int32, -+ epoch: Int64, -+ health: Int64, -+ alarm: Int64, -+ stream: fx.Stream = _STREAM, -+ ): -+ _ = keyed -+ kda_pre( -+ prefix, -+ delta, -+ blocks, -+ blk_sm, -+ blk_sr, -+ ares_nw, -+ ares_qk, -+ in_nw, -+ w_in, -+ w_fb, -+ conv_w, -+ conv_st, -+ cs_seq, -+ cs_dim, -+ cs_tok, -+ a_log, -+ dt_bias, -+ on_w, -+ rstate, -+ rs_seq, -+ st_idx, -+ st_stride, -+ num_acc, -+ core_out, -+ scratch, -+ layer, -+ epoch, -+ health, -+ alarm, -+ ).launch(grid=(BLOCKS,), block=(THREADS,), stream=stream) -+ -+ ABI.check(kda_pre, launch) -+ _BUILDS[key] = launch -+ return launch -+ -+ -+def ptr(t): -+ return 0 if t is None else t.data_ptr() -+ -+ -+def kda_pre( -+ key: KdaPreBuild, -+ *, -+ prefix, -+ delta, -+ blocks, -+ ares_nw, -+ ares_qk, -+ in_nw, -+ w_in, -+ w_fb, -+ conv_w, -+ conv_state, -+ a_log, -+ dt_bias, -+ on_w, -+ rstate, -+ st_idx, -+ num_acc, -+ core_out, -+ scratch, -+ layer, -+ epoch, -+ health, -+ alarm, -+): -+ """Host launcher. ``conv_state`` [slots, dim, state_len] or a transposed view -+ of [slots, state_len, dim] (strides read); ``rstate`` [slots, 12, 128, 128] -+ fp32; ``st_idx`` [R, >= L] int32; ``scratch`` zeroed before the step's first -+ layer (tags are a layer's); ``epoch`` a device int32 bumped once a step and -+ ``layer`` this launch's slot below ``TAG_SLOTS``.""" -+ f = build(key) -+ f( -+ prefix.data_ptr(), -+ ptr(delta), -+ blocks.data_ptr(), -+ blocks.stride(0), -+ blocks.stride(1), -+ ares_nw.data_ptr(), -+ ares_qk.data_ptr(), -+ in_nw.data_ptr(), -+ w_in.data_ptr(), -+ w_fb.data_ptr(), -+ conv_w.data_ptr(), -+ conv_state.data_ptr(), -+ conv_state.stride(0), -+ conv_state.stride(1), -+ conv_state.stride(2), -+ a_log.data_ptr(), -+ dt_bias.data_ptr(), -+ on_w.data_ptr(), -+ rstate.data_ptr(), -+ rstate.stride(0), -+ st_idx.data_ptr(), -+ st_idx.stride(0), -+ num_acc.data_ptr(), -+ core_out.data_ptr(), -+ scratch.data_ptr(), -+ layer, -+ epoch.data_ptr(), -+ health.data_ptr(), -+ alarm, -+ stream=torch.cuda.current_stream(), -+ ) -diff --git a/vllm/models/kimi_k3/amd/specialized_decode/mono/common/__init__.py b/vllm/models/kimi_k3/amd/specialized_decode/mono/common/__init__.py -new file mode 100644 -index 000000000..208f01a7c ---- /dev/null -+++ b/vllm/models/kimi_k3/amd/specialized_decode/mono/common/__init__.py -@@ -0,0 +1,2 @@ -+# SPDX-License-Identifier: Apache-2.0 -+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -diff --git a/vllm/models/kimi_k3/amd/specialized_decode/mono/common/debug.py b/vllm/models/kimi_k3/amd/specialized_decode/mono/common/debug.py -new file mode 100644 -index 000000000..18cdcff38 ---- /dev/null -+++ b/vllm/models/kimi_k3/amd/specialized_decode/mono/common/debug.py -@@ -0,0 +1,38 @@ -+# SPDX-License-Identifier: Apache-2.0 -+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -+# Adapted from ROCm/ATOM (https://github.com/ROCm/ATOM) at a526f0d, MIT License, -+# Copyright (C) 2026, Advanced Micro Devices, Inc.: -+# atom/mono/runtime/debug.py -+"""Wait records: which mailbox region a CTA's given-up wait was on.""" -+ -+import zlib -+ -+import torch -+ -+from vllm.models.kimi_k3.amd.specialized_decode.mono.abi import RECORD_WORDS -+from vllm.models.kimi_k3.amd.specialized_decode.mono.common.plan import BLOCKS -+ -+ -+def region_id(name: str) -> int: -+ """A mailbox region's id in a wait record: its name's CRC-32 (31 bits), the -+ same in every kernel and table that names it.""" -+ return zlib.crc32(name.encode()) & 0x7FFFFFFF -+ -+ -+def region_ids(names) -> dict[str, int]: -+ """Name -> ``region_id``: a ``Mailbox``'s table.""" -+ return {n: region_id(n) for n in names} -+ -+ -+def given_up_waits(health: torch.Tensor, names) -> list[str]: -+ """Each CTA's first wait that gave up, the earliest first, from a rank's -+ ``health`` words; ``names`` the regions a record may name.""" -+ by_id = {region_id(n): n for n in names} -+ records = health.view(BLOCKS + 1, RECORD_WORDS)[1:] -+ records = records[records[:, 0] != 0].tolist() -+ t0 = min((r[6] for r in records), default=0) -+ return [ -+ f"+{(at - t0) / 100:.0f} us CTA {cta}: {by_id.get(rid, rid)} pair {idx}" -+ f" waited for tag {tag}, saw {seen}" -+ for _, rid, idx, tag, seen, cta, at, _ in sorted(records, key=lambda r: r[6]) -+ ] -diff --git a/vllm/models/kimi_k3/amd/specialized_decode/mono/common/mx.py b/vllm/models/kimi_k3/amd/specialized_decode/mono/common/mx.py -new file mode 100644 -index 000000000..908afeaa7 ---- /dev/null -+++ b/vllm/models/kimi_k3/amd/specialized_decode/mono/common/mx.py -@@ -0,0 +1,84 @@ -+# SPDX-License-Identifier: Apache-2.0 -+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -+# Adapted from ROCm/ATOM (https://github.com/ROCm/ATOM) at a526f0d, MIT License, -+# Copyright (C) 2026, Advanced Micro Devices, Inc.: -+# atom/mono/device/mx.py -+"""FP8 / MX arithmetic the mono kernels share: the E4M3 range, E8M0 scale -+codes, the byte of a scale in aiter's ``shuffle_scale`` layout, an FP4 weight -+tile's operand loads and the scaled MFMA that takes them. -+ -+An E8M0 code is a biased exponent; ``pow2(code)`` is the scale it stands for. -+""" -+ -+import flydsl.expr as fx -+from aiter.ops.flydsl.kernels import buffer_ops as bo -+from flydsl.expr import rocdl -+from flydsl.expr.typing import T -+ -+from vllm.models.kimi_k3.amd.specialized_decode.mono.common.ops import CM_NT -+ -+FP8_MAX = 448.0 -+# the f8f6f4 MFMA's operand formats (cbsz / blgp) -+FP8, FP4 = 0, 4 -+UNIT_SCALE = 127 # the E8M0 code of 1.0 -+ -+ -+def clamp_fp8(v): -+ """``v`` clamped to the E4M3 range.""" -+ return fx.max(fx.min(v, fx.Float32(FP8_MAX)), fx.Float32(-FP8_MAX)) -+ -+ -+def code_ceil(v): -+ """Biased exponent of the smallest power of two >= ``v`` (v >= 0, finite), -+ clamped to [1, 254]: the codes a scale may take.""" -+ bits = v.bitcast(fx.Int32) -+ e = ((bits >> 23) & 0xFF) + ((bits & 0x7FFFFF) != 0).select( -+ fx.Int32(1), fx.Int32(0) -+ ) -+ return fx.max(fx.min(e, fx.Int32(254)), fx.Int32(1)) -+ -+ -+def pow2(code): -+ """The fp32 power of two a biased exponent stands for.""" -+ return (code << 23).bitcast(fx.Float32) -+ -+ -+def rcp_pow2(code): -+ """1 / ``pow2(code)``, exact (code in [1, 253]): x times it is x / scale -+ bit for bit, a multiply where the division is a dozen instructions.""" -+ return ((254 - code) << 23).bitcast(fx.Float32) -+ -+ -+def scale_index(row, col, cols): -+ """Byte of e8m0 scale (row, col) in aiter's ``shuffle_scale`` layout: -+ ``view(rows/32, 2, 16, cols/8, 2, 4).permute(0, 3, 5, 2, 4, 1)``.""" -+ r32, a, b = row // 32, (row // 16) % 2, row % 16 -+ c8, d, e = col // 8, (col // 4) % 2, col % 4 -+ return ((((r32 * (cols // 8) + c8) * 4 + e) * 16 + b) * 2 + d) * 2 + a -+ -+ -+def fp4_tile_load(r_w, base, rg, k, st, at): -+ """A lane's 16 B of an FP4 weight's (16, 16)-shuffled 16-row group ``rg`` -+ (K = ``k``) at 128-k step ``st``: dword ``base`` + the group's 2 k + 256 a -+ step + ``at`` (``lane * 4``: 32 fp4, K block lane / 16 of row lane % 16). -+ Nontemporal: a weight a step reads once.""" -+ return fx.Vector( -+ bo.buffer_load( -+ r_w, -+ base + rg * (2 * k) + st * 256 + at, -+ vec_width=4, -+ dtype=T.i32, -+ cache_modifier=CM_NT, -+ ) -+ ) -+ -+ -+def mfma_scaled(a, b, c, sa, sb, a_fmt=FP8, b_fmt=FP8, sa_byte=0, sb_byte=0): -+ """One 16x16x128 scaled MFMA (gfx950 f8f6f4): A ``a`` (``a_fmt``) with E8M0 -+ ``sa`` (its byte ``sa_byte``), B ``b`` (``b_fmt``) with E8M0 ``sb`` (its -+ byte ``sb_byte``).""" -+ return fx.Vector( -+ rocdl.mfma_scale_f32_16x16x128_f8f6f4( -+ T.vec(4, T.f32), [a, b, c, a_fmt, b_fmt, sa_byte, sa, sb_byte, sb] -+ ) -+ ) -diff --git a/vllm/models/kimi_k3/amd/specialized_decode/mono/common/ops.py b/vllm/models/kimi_k3/amd/specialized_decode/mono/common/ops.py -new file mode 100644 -index 000000000..4b7b62dbf ---- /dev/null -+++ b/vllm/models/kimi_k3/amd/specialized_decode/mono/common/ops.py -@@ -0,0 +1,435 @@ -+# SPDX-License-Identifier: Apache-2.0 -+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -+# Adapted from ROCm/ATOM (https://github.com/ROCm/ATOM) at a526f0d, MIT License, -+# Copyright (C) 2026, Advanced Micro Devices, Inc.: -+# atom/mono/device/ops.py -+"""Device-side primitives of the mono kernels: wave reductions and shuffles, the -+approximate hardware math, bf16 / fp8 packing, buffer resources, the -+execution-model constants and the kernel symbol naming. -+ -+Ported from the FlyDSL GLM-5 layer kernel (ROCm/FlyDSL branch -+``glm5-mla-moe-layer-monokernel`` @ 1e28fd8a, -+``kernels/mla_moe_layer/shared_reuse_moe_kernel.py``). Only published flydsl APIs -+and aiter's vendored ``buffer_ops`` / ``dpp_utils`` are used. -+ -+Every function but ``kernel_symbol`` is traced inside a ``@flyc.kernel`` body. -+ -+Raw boundaries kept, for want of an equivalent on the pinned FlyDSL: -+``buffer_ops`` loads / stores (the kernels take raw Int64 addresses, no tensors -+for ``make_buffer_tensor``); raw MFMA intrinsics (hand-laid operands: permlane -+transposes, fp8 packing, MXFP4 scales); the LLVM intrinsics here (``expect``, -+``s_memrealtime``, the approximate rsq / rcp / exp2) and the struct unpacking of -+``permlane*_swap``. -+""" -+ -+from __future__ import annotations -+ -+import flydsl.expr as fx -+from aiter.ops.flydsl.kernels import buffer_ops as bo -+from aiter.ops.flydsl.kernels.dpp_utils import update_dpp_i32 -+from aiter.ops.flydsl.kernels.kernels_common import kernel_signature -+from flydsl._mlir import ir -+from flydsl._mlir.dialects import llvm -+from flydsl.compiler.ast_rewriter import ASTRewriter -+from flydsl.expr import gpu, range_constexpr, rocdl -+from flydsl.expr import math as fmath -+from flydsl.expr.typing import T, as_ir_value -+ -+from vllm.models.kimi_k3.amd.specialized_decode.mono.common.plan import THREADS, WAVES -+ -+# gfx95x cache policy bits (LLVM CPol): SC0 = 1, SC1 = 16. SC1 alone is device -+# scope (past the per-XCD caches); SC0 | SC1 is system scope (peers over XGMI). -+CM_DEV = 16 -+ -+ -+CM_SYS = 17 -+ -+ -+# NT = 2: nontemporal. A weight a layer reads once streams past the caches it -+# would only evict the step's reused data from (K2b at 48 rows: 143 -> 133 us) -+CM_NT = 2 -+ -+ -+def kernel_symbol(stem, **params): -+ """Kernel symbol that profilers demangle to ``atom::_``. The -+ amdhsa assembler rejects ``::`` in a symbol, so it is Itanium-mangled, as -+ aiter's ``_ZN5aiter...E`` kernels are.""" -+ name = f"{stem}_{kernel_signature(**params)}" -+ return f"_ZN4atom{len(name)}{name}E" -+ -+ -+def rsrc(addr, nbytes=None): -+ return bo.create_buffer_resource_from_addr(addr, num_records_bytes=nbytes) -+ -+ -+def permlane_swap(off, x, y): -+ """v_permlane{32,16}_swap on registers (x, y): off 32 trades x's upper-half -+ lanes with y's lower-half lanes; off 16 trades x's odd 16-lane rows with y's -+ even rows (each within its half). -> (x', y').""" -+ swap = rocdl.permlane32_swap if off == 32 else rocdl.permlane16_swap -+ pr = swap( -+ llvm.StructType.get_literal([T.i32, T.i32]), -+ as_ir_value(fx.Int32(x)), -+ as_ir_value(fx.Int32(y)), -+ False, -+ False, -+ ) -+ return tuple(fx.Int32(llvm.extractvalue(T.i32, pr, [j])) for j in range(2)) -+ -+ -+def lane_gather(v, src_lane): -+ """Lane ``src_lane``'s value of v, a source lane a lane (ds_bpermute).""" -+ return fx.Int32( -+ rocdl.ds_bpermute(T.i32, fx.Int32(src_lane) * 4, fx.Int32(v).ir_value()) -+ ) -+ -+ -+def readlane(v, src_lane): -+ """Lane ``src_lane``'s value of v (wave-uniform).""" -+ bits = v.bitcast(fx.Int32) if isinstance(v, fx.Float32) else fx.Int32(v) -+ r = fx.Int32(rocdl.readlane(T.i32, bits, fx.Int32(src_lane))) -+ return r.bitcast(fx.Float32) if isinstance(v, fx.Float32) else r -+ -+ -+def memrealtime(): -+ """s_memrealtime: the 100 MHz device clock (no FlyDSL wrapper).""" -+ return fx.Int64(llvm.call_intrinsic(T.i64, "llvm.amdgcn.s.memrealtime", [], [], [])) -+ -+ -+def unlikely(cond): -+ """``cond`` marked cold (llvm.expect): LLVM lays the branch out of the hot -+ path. A rare path's code inline among the hot stages costs them even when it -+ never runs (the long-context selection: S = 16, 3k context +1.4 us).""" -+ return fx.Boolean( -+ llvm.call_intrinsic( -+ ir.IntegerType.get_signless(1), -+ "llvm.expect.i1", -+ [as_ir_value(cond), as_ir_value(fx.Boolean(False))], -+ [], -+ [], -+ ) -+ ) -+ -+ -+def fresh(v): -+ """``v`` (a per-lane Int32) through an opaque move, here: the lane math built -+ on it is computed at its use, where LLVM would otherwise hoist it out of a task -+ loop or CSE it across stages into registers held for the whole kernel (spills, -+ and a kernel with scratch is refused: ``runtime.widths``).""" -+ return fx.Int32( -+ llvm.inline_asm( -+ T.i32, [fx.Int32(v).ir_value()], "v_mov_b32 $0, $1", "=v,v", -+ has_side_effects=True, -+ ) -+ ) # fmt: skip -+ -+ -+def uniform(v): -+ return fx.Int32(rocdl.readfirstlane(T.i32, fx.Int32(v).ir_value())) -+ -+ -+def _hw_f32(name, x): -+ return fx.Float32( -+ llvm.call_intrinsic(T.f32, name, [fx.Float32(x).ir_value()], [], []) -+ ) -+ -+ -+def hw_rsq(x): -+ """v_rsq_f32 (~1 ulp, no libm range fixups).""" -+ return _hw_f32("llvm.amdgcn.rsq.f32", x) -+ -+ -+def hw_rcp(x): -+ return _hw_f32("llvm.amdgcn.rcp.f32", x) -+ -+ -+def hw_exp2(x): -+ return _hw_f32("llvm.amdgcn.exp2.f32", x) -+ -+ -+def div_rn(x, d, r): -+ """``x / d`` rounded to nearest, bit-identical to the IEEE division, given -+ ``r = 1.0 / d`` from one true division per divisor (Markstein: a quotient -+ within 1 ulp plus its exact fma residual times RN(1/d) rounds correctly).""" -+ q = x * r -+ return fx.Float32(fmath.fma(fx.Float32(fmath.fma(-q, d, x)), r, q)) -+ -+ -+def xshfl(v, off): -+ """Value of lane ``lane ^ off``: offsets 32 / 16 via v_permlane*_swap, in-row -+ offsets via DPP (VALU latency, not LDS). Powers of two only: any other -+ offset would take the lane ^ 1 path.""" -+ assert off in (1, 2, 4, 8, 16, 32), off -+ if off >= 16: -+ return v.shuffle_xor(off, 64) -+ is_f = isinstance(v, fx.Float32) -+ x = v.bitcast(fx.Int32) if is_f else fx.Int32(v) -+ if off == 8: -+ y = fx.Int32(update_dpp_i32(x, x, 0x118, 0xF, 0xC, False)) -+ y = fx.Int32(update_dpp_i32(y, x, 0x108, 0xF, 0x3, False)) -+ elif off == 4: -+ y = fx.Int32(update_dpp_i32(x, x, 0x114, 0xF, 0xA, False)) -+ y = fx.Int32(update_dpp_i32(y, x, 0x104, 0xF, 0x5, False)) -+ elif off == 2: -+ y = fx.Int32(update_dpp_i32(x, x, 0x4E, 0xF, 0xF, False)) -+ else: -+ y = fx.Int32(update_dpp_i32(x, x, 0xB1, 0xF, 0xF, False)) -+ return y.bitcast(fx.Float32) if is_f else y -+ -+ -+def xred(v, off, op): -+ """op(v, lane ^ off) for a symmetric op; 32 / 16 take both halves of one swap.""" -+ if off < 16: -+ return op(v, xshfl(v, off)) -+ is_f = isinstance(v, fx.Float32) -+ x = v.bitcast(fx.Int32) if is_f else fx.Int32(v) -+ a, b = permlane_swap(off, x, x) -+ if is_f: -+ return op(a.bitcast(fx.Float32), b.bitcast(fx.Float32)) -+ return op(type(v)(a), type(v)(b)) -+ -+ -+def _add(a, b): -+ return a + b -+ -+ -+def butterfly(v, offsets, op=_add): -+ """Butterfly reduction over the given xor offsets, in the given order. -+ -+ The order is the numerical contract: the aiter per-head norm reduces 32 -+ lanes with offsets 16, 8, 4, 2, 1, and matching it keeps rstd bit-exact.""" -+ for off in offsets: -+ v = xred(v, off, op) -+ return v -+ -+ -+def wave_sum(v): -+ return butterfly(v, (32, 16, 8, 4, 2, 1)) -+ -+ -+def traced(fn): -+ """Apply flydsl's AST rewriting to a helper, so its ``while`` / ``if`` over -+ traced values lower to scf ops. Kernels rewrite only their own body.""" -+ return ASTRewriter.transform(fn) -+ -+ -+def wave_max(v): -+ return butterfly(v, (32, 16, 8, 4, 2, 1), fx.max) -+ -+ -+@traced -+def block_max(v, lane, wave, red): -+ """Max of v over the CTA's WAVES waves; ``red`` needs WAVES words.""" -+ wv = wave_max(v) -+ gpu.barrier() -+ if lane == 0: -+ fx.ptr_store(wv, red + wave) -+ gpu.barrier() -+ t = fx.ptr_load(red + 0) -+ for w in range_constexpr(1, WAVES): -+ t = fx.max(t, fx.ptr_load(red + w)) -+ gpu.barrier() -+ return t -+ -+ -+def row_shr(v, lane, k): -+ """DPP row_shr:k of an f32: lane - k of the same 16-lane row, else 0.""" -+ moved = lane_gather(v.bitcast(fx.Int32), fx.max(lane - k, 0)).bitcast(fx.Float32) -+ return fx.Float32((lane % 16 >= k).select(moved, fx.Float32(0.0))) -+ -+ -+def _xpartner(v, off, lane): -+ """Lane ``lane ^ off``'s v: DPP within a row, v_permlane*_swap across rows -+ (``xshfl`` takes ``shuffle_xor`` there).""" -+ if off < 16: -+ return xshfl(v, off) -+ a, b = permlane_swap(off, v, v) -+ return ((lane & off) != 0).select(a, b) -+ -+ -+def wave_incl_scan(v, lane): -+ """Inclusive prefix sum of an Int32 over the wave, in lane order.""" -+ blk = v -+ for off in (1, 2, 4, 8, 16, 32): -+ o = _xpartner(blk, off, lane) -+ v = v + ((lane & off) != 0).select(o, 0) -+ blk = blk + o -+ return v -+ -+ -+@traced -+def block_excl_scan(v, lane, wave, buf): -+ """(exclusive prefix of v over the CTA in thread order, CTA total); ``buf`` -+ needs WAVES words.""" -+ incl = wave_incl_scan(v, lane) -+ gpu.barrier() -+ if lane == 63: -+ fx.ptr_store(incl, buf + wave) -+ gpu.barrier() -+ off = fx.Int32(0) -+ tot = fx.Int32(0) -+ for w in range_constexpr(WAVES): -+ cnt = fx.ptr_load(buf + w) -+ off = off + (w < wave).select(cnt, 0) -+ tot = tot + cnt -+ return incl - v + off, tot -+ -+ -+def global_load(addr, words=1, nt=False): -+ """``words`` i32 (1, 2 or 4) at the 64-bit global address ``addr``. A pool -+ past 4 GiB is reached only this way: a buffer load's offset (index x stride -+ included) is 32 bits. ``nt``: nontemporal.""" -+ i32 = ir.IntegerType.get_signless(32) -+ ty = i32 if words == 1 else ir.VectorType.get([words], i32) -+ ptr = llvm.IntToPtrOp(ir.Type.parse("!llvm.ptr<1>"), as_ir_value(fx.Int64(addr))) -+ return llvm.LoadOp( -+ ty, ptr.result, alignment=4 * words, nontemporal=nt or None -+ ).result -+ -+ -+def ld_i32(ptr, i): -+ """Word ``i`` of a buffer, a plain Int32 load.""" -+ return fx.Int32(bo.buffer_load(rsrc(ptr), i, vec_width=1, dtype=T.i32)) -+ -+ -+def row_live(batch_ids, k): -+ """Whether row k of the step is a request's: its ``batch_id_per_q_token`` -+ entry, -1 on a CUDA-graph pad row (wave-uniform for a uniform k). A pad row -+ runs like any other but joins no decision another row's work depends on.""" -+ return uniform(ld_i32(batch_ids, k)) >= 0 -+ -+ -+def ld_dev(ptr, i, words=1): -+ """``words`` i32 (1, 2 or 4) of a buffer from word ``i``, at device scope -+ (past the per-XCD L2).""" -+ return bo.buffer_load( -+ rsrc(ptr), i, vec_width=words, dtype=T.i32, cache_modifier=CM_DEV -+ ) -+ -+ -+LOAD_BATCH = 4 # rounds of a strided loop whose loads go out together -+ -+ -+@traced -+def batched_rounds(tid, n, load, use, batch=LOAD_BATCH): -+ """A CTA's strided loop over elements v = tid + THREADS i < n (n a Python -+ int): ``load(v)`` of ``batch`` rounds issued before any ``use(v, loaded)``. -+ A round whose ``use`` stores to global memory keeps the next round's loads -+ behind it (they may alias), so the rounds' loads would go out in series. A -+ round past n loads element n - 1 (in range) and uses nothing.""" -+ rounds = -(-n // THREADS) -+ for b0 in range_constexpr(0, rounds, batch): -+ idx = range(b0, min(b0 + batch, rounds)) -+ got = [load(fx.min(tid + THREADS * i, n - 1)) for i in idx] -+ for i, g in zip(idx, got): -+ v = tid + THREADS * i -+ if v < n: -+ use(v, g) -+ -+ -+def rows_to_lds(tid, src, rows, row_words, dst, dst_row): -+ """``rows`` rows of ``row_words`` words (a multiple of 4) read at device scope -+ from address ``src`` -> LDS ``dst``, a row every ``dst_row`` words: 16 B a -+ lane, the loads batched (``batched_rounds``).""" -+ assert row_words % 4 == 0, row_words -+ batched_rounds( -+ tid, rows * row_words // 4, -+ lambda v: fx.Vector(ld_dev(src, v * 4, 4)), -+ lambda v, w: fx.ptr_store( -+ w, dst + (v * 4 // row_words * dst_row + v * 4 % row_words) -+ ), -+ ) # fmt: skip -+ -+ -+def mfma_bf16(a, b, c): -+ return fx.Vector(rocdl.mfma_f32_16x16x32_bf16(T.vec(4, T.f32), [a, b, c])) -+ -+ -+def row_sum(red, r, t, waves=None): -+ """Row r, token t of the waves' 16x16 MFMA accumulators (wave w's at ``red`` -+ + (w * 64 + lane) * 4), summed in wave order from 0 (``waves``: which, in -+ order; all of them by default).""" -+ tot = fx.Float32(0.0) -+ for w in waves if waves is not None else range(WAVES): -+ tot = tot + fx.ptr_load(red + ((w * 64 + 16 * (r // 4) + t) * 4 + r % 4)) -+ return tot -+ -+ -+def mfma_fp8(a, b, c): -+ """16x16x32 fp8 x fp8 MFMA, unscaled: ``a`` / ``b`` an i64 of 8 E4M3 each.""" -+ return fx.Vector( -+ rocdl.mfma_f32_16x16x32_fp8_fp8(T.vec(4, T.f32), [a, b, c, 0, 0, 0]) -+ ) -+ -+ -+def ballot(pred): -+ """The wave's lanes where ``pred`` holds, a bit a lane.""" -+ return fx.Int64(rocdl.ballot(T.i64, as_ir_value(pred))) -+ -+ -+def popcount(x): -+ return fx.Int32(fx.ctpop(x)) -+ -+ -+def lanes_below(bal, lane): -+ """How many of the lanes in ballot ``bal`` are below ``lane``.""" -+ return popcount(bal & ((fx.Int64(1) << fx.Int64(lane)) - 1)) -+ -+ -+def load_ptr64(table, i): -+ """Entry ``i`` of an i64 table in memory, wave-uniform (each half through -+ ``uniform``): a peer's buffer base, a kernel argument block's pointer.""" -+ pw = fx.Vector(bo.buffer_load(rsrc(table), i * 2, vec_width=2, dtype=T.i32)) -+ return (fx.Int64(uniform(pw[1])) << 32) | fx.Int64(fx.Uint32(uniform(pw[0]))) -+ -+ -+def fp8x8_bf16_pk(words): -+ """Eight fp8 of two dwords -> bf16 MFMA operand (exact) by -+ v_cvt_scalef32_pk_bf16_fp8 (scale 1), two values an instruction.""" -+ out = [] -+ for w in words: -+ for hi in (False, True): -+ p = rocdl.cvt_scalef32_pk_bf16_fp8( -+ T.bf16x2, as_ir_value(fx.Int32(w)), as_ir_value(fx.Float32(1.0)), hi -+ ) -+ out.append(fx.Vector(p).bitcast(fx.Int32)[0]) -+ return fx.Vector.from_elements(out, fx.Int32).bitcast(fx.BFloat16) -+ -+ -+def bf16_round(a): -+ return fx.Float32(fx.Float32(a).to(fx.BFloat16)) -+ -+ -+def bf16_pair(a, b): -+ """Two f32 -> one f32-typed word holding (bf16(a), bf16(b)).""" -+ return ( -+ fx.Vector.from_elements([a, b], fx.Float32) -+ .to(fx.BFloat16) -+ .bitcast(fx.Float32)[0] -+ ) -+ -+ -+def lds_bytes(struct) -> int: -+ """The LDS an ``fx.struct`` class takes, alignment included.""" -+ return struct.__dsl_size_of__() -+ -+ -+def bf_lo(w): -+ """The f32 of a packed bf16 pair word's low half.""" -+ return (w << 16).bitcast(fx.Float32) -+ -+ -+def bf_hi(w): -+ """The f32 of a packed bf16 pair word's high half.""" -+ return (w & fx.Int32(-65536)).bitcast(fx.Float32) -+ -+ -+def bf2_f32(w): -+ """Packed bf16 pair word -> (f32 low, f32 high).""" -+ return bf_lo(w), bf_hi(w) -+ -+ -+def fp8_pack4(a, b, c, d): -+ """Four f32 (already in range) -> one dword of four E4M3 bytes, RNE.""" -+ lo = fx.Int32(rocdl.cvt_pk_fp8_f32(T.i32, a, b, fx.Int32(0), False)) -+ return fx.Int32(rocdl.cvt_pk_fp8_f32(T.i32, c, d, lo, True)) -diff --git a/vllm/models/kimi_k3/amd/specialized_decode/mono/common/plan.py b/vllm/models/kimi_k3/amd/specialized_decode/mono/common/plan.py -new file mode 100644 -index 000000000..0e769475a ---- /dev/null -+++ b/vllm/models/kimi_k3/amd/specialized_decode/mono/common/plan.py -@@ -0,0 +1,237 @@ -+# SPDX-License-Identifier: Apache-2.0 -+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -+# Adapted from ROCm/ATOM (https://github.com/ROCm/ATOM) at a526f0d, MIT License, -+# Copyright (C) 2026, Advanced Micro Devices, Inc.: -+# atom/mono/plan/build_key.py -+# atom/mono/plan/execution.py -+# atom/mono/plan/layout.py -+# atom/mono/runtime/abi.py -+# atom/mono/plan/trace.py -+"""A kernel build's parameters, all of them in FlyDSL's JIT cache key. -+ -+FlyDSL keys a compiled kernel by the scalar values (int, float, bool, str, enum, -+tuple) its launcher's closures capture; an object in a closure is not keyed, and -+a parameter used only by build-time Python (a placement table, a layout) is in no -+closure at all. Two builds differing only there would share one binary -- every -+TP rank ran rank 3's kernel once that way. So a build's parameters are one frozen -+dataclass of scalars, ``key_tuple`` turns it into one (name, value) tuple, and the -+launcher references that tuple: every field is keyed, whatever the body uses. -+ -+The tuple also carries the model's kernel sources (``source_digest``): FlyDSL -+keys a helper's source only when its walk reaches it, and it does not reach a -+helper called inside a traced function's dynamic loop or branch (the rewriter -+moves those bodies into nested functions it does not scan) nor one outside the -+kernel's directory -- an edit there kept the old binary. -+""" -+ -+import enum -+import functools -+import hashlib -+import inspect -+from contextlib import contextmanager -+from dataclasses import dataclass, field, fields -+from enum import Enum -+from pathlib import Path -+ -+from vllm.models.kimi_k3.amd.specialized_decode.mono.abi import health_words -+ -+_SCALARS = (int, float, bool, str, enum.Enum) -+_ROOT = Path(__file__).resolve().parents[1] -+ -+ -+@functools.cache -+def source_digest(*paths: str) -> str: -+ """sha256 of every ``.py`` file under ``paths`` (files or directories, -+ relative to the ``mono`` package), each file's path and contents: one model's -+ kernel sources, so an edit to another model's rebuilds none of these.""" -+ h = hashlib.sha256() -+ for rel in paths: -+ root = _ROOT / rel -+ files = [root] if root.is_file() else sorted(root.rglob("*.py")) -+ if not files: -+ raise FileNotFoundError(f"no kernel sources at {root}") -+ for f in files: -+ h.update(str(f.relative_to(_ROOT)).encode()) -+ h.update(f.read_bytes()) -+ return h.hexdigest()[:16] -+ -+ -+def key_tuple(key, sources: str) -> tuple: -+ """``key``'s (field name, value) pairs, every value a scalar, then -+ ("sources", ``sources``): the model's ``source_digest``.""" -+ out = [] -+ for f in fields(key): -+ value = getattr(key, f.name) -+ if not isinstance(value, _SCALARS): -+ raise TypeError( -+ f"{type(key).__name__}.{f.name} = {value!r}: a build parameter must be" -+ " a scalar to reach the JIT cache key" -+ ) -+ out.append((f.name, value)) -+ out.append(("sources", sources)) -+ return tuple(out) -+ -+ -+def symbol_params(key) -> dict: -+ """The fields that name the kernel symbol (``metadata={"sym": short name}``), -+ short name -> value, in field order.""" -+ return { -+ f.metadata["sym"]: getattr(key, f.name) -+ for f in fields(key) -+ if "sym" in f.metadata -+ } -+ -+ -+# The execution model every mono kernel is written for. -+# -+# One CTA per CU of an MI355X, all of them co-resident -- a spin wait on another -+# CTA can only make progress if that CTA is running, so the runner refuses a GPU -+# with another CU count -- and THREADS threads each. -+# -+ -+BLOCKS = 256 -+THREADS = 512 -+WAVES = THREADS // 64 -+LDS_BYTES = 160 * 1024 # a CU's LDS: one CTA's whole -+ -+ -+def first_task(bid, base): -+ """CTA ``bid``'s first task of a stage placed from CTA ``base`` (any count of -+ CTAs past 0: placements that add task counts run past the grid), then every -+ BLOCKS-th.""" -+ return (bid - base % BLOCKS + BLOCKS) % BLOCKS -+ -+ -+# A kernel's mailbox regions laid out in its scratch or peer buffer. -+ -+PAIR_BYTES = 8 # a (value, tag) pair -+# A rank's health words (``mono.abi``), for this grid. -+HEALTH_WORDS = health_words(BLOCKS) -+ -+ -+def pair_layout(items, align: int = PAIR_BYTES, start: int = 0) -> dict: -+ """``items`` (name, pairs) laid out in order: name -> (byte offset, bytes), -+ each region from an ``align``-byte boundary at or past ``start``.""" -+ out, off = {}, start -+ for name, pairs in items: -+ off = (off + align - 1) // align * align -+ out[name] = (off, pairs * PAIR_BYTES) -+ off += pairs * PAIR_BYTES -+ return out -+ -+ -+# A kernel's argument list, named once. -+# -+# A FlyDSL kernel and its launcher each spell out their positional parameters, and -+# the host passes them by position: an argument out of order is a pointer read as -+# another, silently. ``KernelAbi`` is the one list of names; the build checks the -+# kernel's and the launcher's parameters against it, and the host packs by name. -+# -+ -+ -+def _python_params(fn) -> tuple[str, ...]: -+ """Parameter names of a ``@flyc.kernel`` / ``@flyc.jit`` object's Python -+ function (FlyDSL keeps it as ``_func`` / ``func``).""" -+ py = getattr(fn, "func", None) or getattr(fn, "_func", None) -+ if py is None: -+ raise TypeError(f"no Python function behind {fn!r}") -+ return tuple(inspect.signature(py).parameters) -+ -+ -+@dataclass(frozen=True) -+class KernelAbi: -+ names: tuple[str, ...] -+ -+ def check(self, kernel, launcher) -> None: -+ """The kernel takes exactly ``names``; the launcher too, then ``stream``.""" -+ got = _python_params(kernel) -+ if got != self.names: -+ raise TypeError(f"kernel parameters {got} != ABI {self.names}") -+ got = _python_params(launcher) -+ if got != (*self.names, "stream"): -+ raise TypeError(f"launcher parameters {got} != ABI {self.names} + stream") -+ -+ def pack(self, values: dict) -> list: -+ """``values`` (every name, nothing else) in ABI order.""" -+ try: -+ if len(values) == len(self.names): -+ return [values[n] for n in self.names] -+ except KeyError: -+ pass -+ extra = sorted(set(values) - set(self.names)) -+ missing = sorted(set(self.names) - set(values)) -+ raise TypeError(f"ABI mismatch: missing {missing}, unexpected {extra}") -+ -+ def zeros(self) -> list: -+ """Every argument 0: what a compile-only call passes.""" -+ return [0] * len(self.names) -+ -+ -+# A kernel's mailbox accesses, recorded while FlyDSL traces it. -+# -+# Tracing runs the kernel body's Python once, so every mailbox put / poll executes -+# once as Python: noting it here costs the device nothing (no IR is emitted). The -+# stage currently being traced is set by ``enter_stage``, called where the program -+# starts each stage, so the record is ``(stage, region, put | poll)`` in program -+# order. -+# -+# Everything here is plain functions on a module-level recorder: FlyDSL's AST -+# rewriter carries every local whose method is called inside a dynamic if / for as -+# loop state, so a kernel body must not call methods on these objects. -+# -+ -+ -+class Space(Enum): -+ SCRATCH = "scratch" # this GPU's memory: device-scope hand-off between CTAs -+ PEER = "peer" # a peer-symmetric region: system-scope hand-off between GPUs -+ -+ -+@dataclass(frozen=True) -+class Access: -+ stage: str -+ region: str -+ space: Space -+ kind: str # "put" | "poll" -+ -+ -+@dataclass -+class Record: -+ stages: list[str] = field(default_factory=list) # in the order entered -+ accesses: list[Access] = field(default_factory=list) -+ unnamed: list[tuple[str, str]] = field(default_factory=list) # (stage, kind) -+ -+ -+_record: Record | None = None -+ -+ -+@contextmanager -+def recording(): -+ """Record the accesses of kernels traced inside this block.""" -+ global _record -+ outer, _record = _record, Record() -+ try: -+ yield _record -+ finally: -+ _record = outer -+ -+ -+def enter_stage(name: str) -> None: -+ if _record is not None: -+ _record.stages.append(name) -+ -+ -+def _stage() -> str: -+ assert _record is not None -+ return _record.stages[-1] if _record.stages else "" -+ -+ -+def note(region: str, space: Space, kind: str) -> None: -+ if _record is not None: -+ _record.accesses.append(Access(_stage(), region, space, kind)) -+ -+ -+def note_unnamed(kind: str) -> None: -+ """A put / poll on a raw address: the check reports it (every mailbox access of -+ a checked kernel names its region).""" -+ if _record is not None: -+ _record.unnamed.append((_stage(), kind)) -diff --git a/vllm/models/kimi_k3/amd/specialized_decode/mono/common/ranks.py b/vllm/models/kimi_k3/amd/specialized_decode/mono/common/ranks.py -new file mode 100644 -index 000000000..ba75dba7c ---- /dev/null -+++ b/vllm/models/kimi_k3/amd/specialized_decode/mono/common/ranks.py -@@ -0,0 +1,33 @@ -+# SPDX-License-Identifier: Apache-2.0 -+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -+# Adapted from ROCm/ATOM (https://github.com/ROCm/ATOM) at a526f0d, MIT License, -+# Copyright (C) 2026, Advanced Micro Devices, Inc.: -+# atom/mono/device/ranks.py -+"""A TP group seen from inside a kernel: every rank's peer buffer, and an -+all-reduce's sum of the ranks' partials. ``tp`` is a build-time constant.""" -+ -+from flydsl.expr import range_constexpr -+ -+from vllm.models.kimi_k3.amd.specialized_decode.mono.common.ops import ( -+ bf2_f32, -+ load_ptr64, -+) -+ -+ -+def peer_bases(peers, tp): -+ """Every rank's peer buffer base from the ``peers`` table (an i64 a rank), -+ loaded once: a put's address must not wait on a load of its own.""" -+ return [load_ptr64(peers, p) for p in range(tp)] -+ -+ -+def sum_partials(poll, own, pair_of, tp): -+ """The ``tp`` ranks' bf16 pair partials at ``pair_of(src)`` of the mailbox -+ region ``own``, summed in fp32 in rank order 0 .. tp - 1: one poll batch -> -+ (lo sum, hi sum). Every rank sums in the same order, so they agree.""" -+ ws = poll([(own, pair_of(src), 1) for src in range(tp)]) -+ acc0, acc1 = bf2_f32(ws[0][0]) -+ for src in range_constexpr(1, tp): -+ lo, hi = bf2_f32(ws[src][0]) -+ acc0 = acc0 + lo -+ acc1 = acc1 + hi -+ return acc0, acc1 -diff --git a/vllm/models/kimi_k3/amd/specialized_decode/mono/common/sync.py b/vllm/models/kimi_k3/amd/specialized_decode/mono/common/sync.py -new file mode 100644 -index 000000000..4dfa087e3 ---- /dev/null -+++ b/vllm/models/kimi_k3/amd/specialized_decode/mono/common/sync.py -@@ -0,0 +1,255 @@ -+# SPDX-License-Identifier: Apache-2.0 -+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -+# Adapted from ROCm/ATOM (https://github.com/ROCm/ATOM) at a526f0d, MIT License, -+# Copyright (C) 2026, Advanced Micro Devices, Inc.: -+# atom/mono/device/sync.py -+"""The tagged mailbox: hand-offs between the CTAs of one GPU (device scope) and -+between GPUs (system scope), every address a named region (``Addr``) recorded -+when the kernel is traced (``trace``).""" -+ -+from __future__ import annotations -+ -+from dataclasses import dataclass -+ -+import flydsl.expr as fx -+from aiter.ops.flydsl.kernels import buffer_ops as bo -+from flydsl.expr import const_expr, gpu, range_constexpr, rocdl -+from flydsl.expr.typing import T -+ -+from vllm.models.kimi_k3.amd.specialized_decode.mono.abi import ( -+ FAILED, -+ RECORD_WORDS, -+ WAIT_SECONDS, -+) -+from vllm.models.kimi_k3.amd.specialized_decode.mono.common.ops import ( -+ CM_DEV, -+ CM_SYS, -+ bf16_pair, -+ ld_i32, -+ memrealtime, -+ rsrc, -+ traced, -+ uniform, -+) -+from vllm.models.kimi_k3.amd.specialized_decode.mono.common.plan import Space, note -+ -+POLL_MAX = 12 # mailbox specs polled per batch (bounds live registers) -+ -+WAIT_TICKS = WAIT_SECONDS * 100_000_000 # s_memrealtime ticks at 100 MHz -+ -+ -+@traced -+def _spin(load_all, pending, health, alarm, region_id, index, tag): -+ """Re-issue ``load_all`` until ``pending`` clears, for at most WAIT_TICKS, and -+ not at all once this rank has failed (``health`` word FAILED), so a failed -+ step runs to its end instead of hanging. A side-effecting op in the loop keeps -+ the loads from being hoisted. -+ -+ A CTA's first give-up is recorded in its ``health`` record; it then raises the -+ output rank's flag for this rank (``alarm``) and drains that store before it -+ sets FAILED, so no later store of this rank's can overtake the flag.""" -+ t0 = memrealtime() -+ v = load_all() -+ why = fx.Int32(0) # 1: timed out, 2: the rank had failed -+ while pending(v) & (why == 0): -+ rocdl.s_sleep(1) -+ v = load_all() -+ failed = ld_health(health, FAILED) != 0 -+ late = memrealtime() - t0 > WAIT_TICKS -+ why = failed.select(fx.Int32(2), late.select(fx.Int32(1), fx.Int32(0))) -+ if why == 1: -+ bid = fx.block_idx.x -+ rec = (bid + 1) * RECORD_WORDS -+ if ld_health(health, rec) == 0: -+ at = fx.Int32(memrealtime() & 0x7FFFFFFF) -+ record = [1, region_id, index, tag, v[1], bid, at, 0] -+ for h in range_constexpr(2): # 16 B a store -+ bo.buffer_store( -+ fx.Vector.from_elements( -+ [fx.Int32(w) for w in record[4 * h : 4 * h + 4]], fx.Int32 -+ ), -+ rsrc(health), -+ rec + 4 * h, -+ cache_modifier=CM_SYS, -+ ) -+ bo.buffer_store(fx.Int32(1), rsrc(alarm), 0, cache_modifier=CM_SYS) -+ rocdl.s_waitcnt(vmcnt=0) -+ bo.buffer_store(fx.Int32(1), rsrc(health), FAILED, cache_modifier=CM_SYS) -+ return v -+ -+ -+@traced -+def stall(fault_rank, fault_epoch, fault_layer, ticks, rank, epoch, layer): -+ """The fault hook (tests only): CTA 0 of rank ``fault_rank`` stalls ``ticks`` -+ once, in the launch for slot ``fault_layer`` of the step at epoch -+ ``fault_epoch``.""" -+ if ( -+ (rank == fault_rank) -+ & (uniform(ld_i32(epoch, 0)) == fault_epoch) -+ & (layer == fault_layer) -+ & (fx.block_idx.x == 0) -+ ): -+ t0 = memrealtime() -+ now = memrealtime() -+ while now - t0 < ticks: -+ rocdl.s_sleep(127) -+ now = memrealtime() -+ -+ -+def ld_health(health, i): -+ return fx.Int32( -+ bo.buffer_load(rsrc(health), i, vec_width=1, dtype=T.i32, cache_modifier=CM_SYS) -+ ) -+ -+ -+@dataclass(frozen=True) -+class Addr: -+ """A mailbox region's address: the region's name and space, recorded when the -+ kernel is traced (``trace``), and the traced address. The space also -+ sets the hand-off's scope: device for SCRATCH, system for PEER.""" -+ -+ region: str -+ space: Space -+ value: object -+ -+ -+def sreg(base, offset, region): -+ """Scratch region ``region`` at byte ``offset`` of ``base``.""" -+ return Addr(region, Space.SCRATCH, base + fx.Int64(offset)) -+ -+ -+def preg(base, offset, region): -+ """Peer-symmetric region ``region`` at byte ``offset`` of a rank's buffer -+ ``base``.""" -+ return Addr(region, Space.PEER, base + fx.Int64(offset)) -+ -+ -+def shift(a, delta): -+ """``a`` moved by ``delta`` bytes (a traced or Python integer), same region.""" -+ return Addr(a.region, a.space, a.value + delta) -+ -+ -+# the store / load cache policy of a hand-off in each space -+_SCOPE_CM = {Space.SCRATCH: CM_DEV, Space.PEER: CM_SYS} -+ -+ -+def _noted(a, kind): -+ note(a.region, a.space, kind) -+ return a.value -+ -+ -+class Mailbox: -+ """Tagged-pair hand-off between CTAs (device scope) and GPUs (system scope). -+ -+ Every 32-bit value is stored next to this launch's tag; a consumer -+ polls the payload until all tags match, so a hand-off costs one round trip -+ (no store drain, no separate flag). The tag is ``layer + 1``: it tells the -+ layers of one step apart, and the runner zeroes every mailbox at a step's -+ start (``StepMailboxes.begin_step``), so a pair from an earlier step never -+ carries a live tag. Every address is an ``Addr``: its space picks the scope. -+ """ -+ -+ def __init__(self, layer, health, alarm, region_ids): -+ """``health``: this rank's health words (FAILED, then a wait record a -+ CTA); ``alarm``: this rank's flag word on the output rank; -+ ``region_ids``: region name -> the id a record names it by.""" -+ self.tag = layer + 1 -+ self.health = health -+ self.alarm = alarm -+ self.region_ids = region_ids -+ -+ def put(self, a, i, v): -+ bits = v.bitcast(fx.Int32) if isinstance(v, fx.Float32) else fx.Int32(v) -+ bo.buffer_store( -+ fx.Vector.from_elements([bits, self.tag], fx.Int32), -+ rsrc(_noted(a, "put")), -+ i * 2, -+ cache_modifier=_SCOPE_CM[a.space], -+ ) -+ -+ def put_bf(self, a, i, vs): -+ """Elements i .. i+len(vs) (2 or 4, i aligned) as packed bf16 pairs, one -+ store.""" -+ words = [] -+ for j in range_constexpr(len(vs) // 2): -+ words += [bf16_pair(vs[2 * j], vs[2 * j + 1]).bitcast(fx.Int32), self.tag] -+ bo.buffer_store( -+ fx.Vector.from_elements(words, fx.Int32), -+ rsrc(_noted(a, "put")), -+ i, -+ cache_modifier=_SCOPE_CM[a.space], -+ ) -+ -+ def put_words(self, a, i, words): -+ """Raw 32-bit words (1 or 2) at pairs i.. as (word, tag). Integer payloads -+ (fp8 words) must come this way: through an f32 value a NaN bit pattern is -+ not kept.""" -+ vals = [] -+ for w in words: -+ vals += [fx.Int32(w), self.tag] -+ bo.buffer_store( -+ fx.Vector.from_elements(vals, fx.Int32), -+ rsrc(_noted(a, "put")), -+ i * 2, -+ cache_modifier=_SCOPE_CM[a.space], -+ ) -+ -+ def poll(self, specs, batch=POLL_MAX): -+ """Batched poll of pairs ``specs`` = [(Addr, pair index, npairs in {1, 2})], -+ all of one space. -+ -+ All pairs are re-loaded together while any tag is stale, so a batch costs -+ one round trip after its last producer lands. Returns the value words.""" -+ if const_expr(len(specs) == 0): -+ return [] -+ if const_expr(len(specs) > batch): -+ return self.poll(specs[:batch], batch) + self.poll(specs[batch:], batch) -+ spaces = {a.space for a, _, _ in specs} -+ assert len(spaces) == 1, f"one poll batch spans {spaces}" -+ cm = _SCOPE_CM[spaces.pop()] -+ bases = [_noted(a, "poll") for a, _, _ in specs] -+ tag = self.tag -+ -+ def load_all(): -+ words = [] -+ for b, (_, i, n) in zip(bases, specs): -+ w = fx.Vector( -+ bo.buffer_load( -+ rsrc(b), -+ fx.Int32(i) * 2, -+ vec_width=2 * n, -+ dtype=T.i32, -+ cache_modifier=cm, -+ ) -+ ) -+ words += [w[e] for e in range(2 * n)] -+ return fx.Vector.from_elements(words, fx.Int32) -+ -+ nw = sum(2 * n for _, _, n in specs) -+ -+ def pending(v): -+ bad = v[1] != tag -+ for e in range_constexpr(3, nw, 2): -+ bad = bad | (v[e] != tag) -+ return bad -+ -+ rid = self.region_ids.get(specs[0][0].region, -1) -+ v = _spin(load_all, pending, self.health, self.alarm, rid, specs[0][1], tag) -+ outs, e = [], 0 -+ for _, _, n in specs: -+ outs.append([v[e + 2 * q] for q in range(n)]) -+ e += 2 * n -+ return outs -+ -+ -+@traced -+def publish(put, a, i, v, who): -+ """This CTA's plain stores made readable by every CTA, then (thread ``who``) -+ the mailbox pair (``a``, i) := v. The stores are drained first: a flag may -+ not overtake the data it announces, and a workgroup-scope release fence does -+ not wait for them (LLVM's AMDGPU memory model omits vmcnt outside tgsplit -+ mode: the workgroup's waves share one CU, the readers here do not).""" -+ rocdl.s_waitcnt(vmcnt=0) -+ gpu.barrier() -+ if who: -+ put(a, i, v) -diff --git a/vllm/models/kimi_k3/amd/specialized_decode/mono/layer.py b/vllm/models/kimi_k3/amd/specialized_decode/mono/layer.py -new file mode 100644 -index 000000000..c89c2b309 ---- /dev/null -+++ b/vllm/models/kimi_k3/amd/specialized_decode/mono/layer.py -@@ -0,0 +1,573 @@ -+# SPDX-License-Identifier: Apache-2.0 -+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -+"""K2 ``layer_post``: one Kimi-K3 layer from the attention's core output (the -+o_proj input) to the layer's MoE output, in one launch -- K2a's o_proj, its -+all-reduce and the MLP AttnRes seam, then K2b's latent MoE (``moe``). -+ -+As ATOM's ``layer_post``, the two halves share the launch so that the MoE's -+static weights stream while the AttnRes chain (partials -> reduce -> slice -> -+gate -> mix) is still settling: the CTAs that will run the router, latent and -+shared gate/up GEMVs issue their weights before they first wait. -+ -+ CTAs 0 .. 55 K2a's reduce + AttnRes (slice, gate, mix -> x rows + XRDY); -+ then route (0 .. S-1) and the shared SiTU (8 .. 8+S-1) -+ CTAs 56 .. 159 o_proj (448 x 16 rows) -> partials to every rank; then -+ the shared gate/up (96 tasks, weights right after o_proj) -+ CTAs 160 .. 243 router parts (28) and latent parts (56): weights first -+ every CTA ug; down (0 .. 223) or the shared down (224 .. 255); -+ lnorm (56 .. 83), up (0 .. 55), reduce (200 .. 255) -+""" -+ -+from dataclasses import dataclass -+ -+import flydsl.compiler as flyc -+import flydsl.expr as fx -+import torch -+from flydsl.expr import const_expr, range_constexpr -+from flydsl.expr.typing import Int32, Int64 -+ -+from vllm.models.kimi_k3.amd.specialized_decode.mono.attention import back as k2a -+from vllm.models.kimi_k3.amd.specialized_decode.mono.attention.kda import ( -+ _SOURCES, -+ COLS, -+ HIDDEN, -+ NA, -+ stage_gate, -+ stage_mix, -+ stage_slice, -+ step_tag, -+) -+from vllm.models.kimi_k3.amd.specialized_decode.mono.common.debug import ( -+ region_ids, -+) -+from vllm.models.kimi_k3.amd.specialized_decode.mono.common.ops import ( -+ ld_i32, -+ load_ptr64, -+ traced, -+ uniform, -+) -+from vllm.models.kimi_k3.amd.specialized_decode.mono.common.plan import ( -+ BLOCKS, -+ THREADS, -+ WAVES, -+ KernelAbi, -+ key_tuple, -+ pair_layout, -+) -+from vllm.models.kimi_k3.amd.specialized_decode.mono.common.ranks import peer_bases -+from vllm.models.kimi_k3.amd.specialized_decode.mono.common.sync import ( -+ Mailbox, -+ preg, -+ shift, -+ sreg, -+ stall, -+) -+from vllm.models.kimi_k3.amd.specialized_decode.mono.stages import gemv as gemv -+from vllm.models.kimi_k3.amd.specialized_decode.mono.stages import moe as k2b -+ -+O0, O1 = NA, BLOCKS # o_proj CTAs: every CTA past AttnRes's -+O_N = O1 - O0 # 200 -+OTASKS = k2a.OTASKS # 448 -+O_ROUNDS = (OTASKS + O_N - 1) // O_N # 5 -+E0 = 160 # router / latent CTAs (their weights in flight under their o_proj) -+H0 = 8 # the shared SiTU's CTAs (after their AttnRes work) -+SD0 = k2b.DOWN_TASKS # 224: the shared down's CTAs -+ -+assert E0 + k2b.R_TASKS + k2b.L_TASKS <= BLOCKS and k2b.G_TASKS <= E0 - O0 -+ -+REGIONS = ("part", "wgt", "msq", "xrdy", "ar") + k2b.REGIONS -+ -+_STREAM = fx.Stream(None) -+ -+ -+@dataclass(frozen=True) -+class K2Build: -+ tokens: int -+ nblocks: int # the MLP AttnRes block sources -+ eps: float = 1e-5 # mlp_res_norm -+ out_eps: float = 1e-5 # post_attention_layernorm -+ ln_eps: float = 1e-5 # the latent RMSNorm -+ reduce: bool = True # the final reduce here (else the next layer's K1) -+ reset: bool = False # a block-write layer: the prefix restarts at the o_proj output -+ # the fault hook (sync.stall): rank fault_rank stalls fault_ticks once, in -+ # slot fault_layer's launch of the step at epoch fault_epoch -+ fault_rank: int = -1 -+ fault_epoch: int = 0 -+ fault_layer: int = 0 -+ fault_ticks: int = 0 -+ -+ -+ABI = KernelAbi( -+ ( -+ "core", -+ "w_o", -+ "prefix", -+ "blocks", -+ "blk_sm", -+ "blk_sr", -+ "ares_nw", -+ "ares_qk", -+ "in_nw", -+ "w_gate", -+ "bias", -+ "w_ld", -+ "w_sgu", -+ "w_sd", -+ "w13", -+ "w13s", -+ "w2", -+ "w2s", -+ "ln_w", -+ "w_up", -+ "out", -+ "scratch", -+ "peers", -+ "rank", -+ "epoch", -+ "layer", -+ "health", -+ "alarm", -+ ) -+) -+ -+ -+def peer_layout(s): -+ """K2a's AR region, then K2b's: one parity's (name -> (byte offset, bytes)).""" -+ pairs = {"ar": k2a.ar_half_pairs(s), **k2b.peer_pairs(s)} -+ return pair_layout(pairs.items()) -+ -+ -+def peer_bytes(s): -+ lay = peer_layout(s) -+ return 2 * max(o + n for o, n in lay.values()) -+ -+ -+def scratch_layout(key: K2Build) -> dict: -+ s = key.tokens -+ a = k2a.scratch_layout(s, key.nblocks) -+ end = max(o + n for o, n in a.values()) -+ b = k2b.scratch_layout(s) -+ lay = dict(a) -+ for name, (o, n) in b.items(): -+ lay[name] = ((end + 15) // 16 * 16 + o, n) -+ end = max(o + n for o, n in lay.values()) -+ lay["xrow"] = ((end + 15) // 16 * 16, s * HIDDEN * 2) -+ return lay -+ -+ -+def scratch_bytes(key: K2Build) -> int: -+ return max(o + n for o, n in scratch_layout(key).values()) -+ -+ -+@traced -+def stage_oproj(c, gi): -+ """o_proj tasks gi, gi + 200, gi + 400 (< 448): every task's weights in -+ flight at once (one memory latency, not three), then the MFMAs.""" -+ ops = [k2a.o_loads(c, fx.min(gi + O_N * j, OTASKS - 1)) for j in range(O_ROUNDS)] -+ k2a.load_core(c) -+ for j in range_constexpr(O_ROUNDS): -+ t = gi + O_N * j -+ if t < OTASKS: -+ k2a.stage_o(c, t, ops[j]) -+ -+ -+_BUILDS: dict = {} -+ -+ -+def build(key: K2Build): -+ if key in _BUILDS: -+ return _BUILDS[key] -+ s, ns = key.tokens, key.nblocks + 1 -+ assert 1 <= s <= 8 -+ lay = scratch_layout(key) -+ play = peer_layout(s) -+ half = peer_bytes(s) // 2 -+ -+ @fx.struct -+ class RouteLds: -+ route: fx.Array[fx.Int32, s * 2 * k2b.TOPK, 16] -+ flag: fx.Array[fx.Int32, k2b.E, 16] -+ expert: fx.Array[fx.Int32, k2b.MAX_U, 16] -+ kof: fx.Array[fx.Int32, k2b.MAX_U * 8, 16] -+ scan: fx.Array[fx.Int32, WAVES, 16] -+ -+ @fx.struct -+ class EarlyLds: -+ xl: fx.Array[fx.Int32, s * (HIDDEN + gemv.LDS_PAD) // 2, 16] -+ -+ @fx.struct -+ class AttnLds: -+ vals: fx.Array[fx.Float32, s * ns * COLS, 16] -+ dl: fx.Array[fx.Float32, s * COLS, 16] -+ -+ @fx.struct -+ class UgLds: -+ latq: fx.Array[fx.Int32, s * k2b.LAT // 8, 16] -+ latsc: fx.Array[fx.Int32, s * k2b.LAT // 32, 16] -+ av: fx.Array[fx.Float32, THREADS, 16] -+ -+ @fx.struct -+ class DownLds: -+ interq: fx.Array[fx.Int32, s * k2b.TOPK * 48, 16] -+ intsc: fx.Array[fx.Int32, s * k2b.TOPK * k2b.UG_GROUPS, 16] -+ contrib: fx.Array[fx.Float32, s * k2b.TOPK * k2b.ROWS, 16] -+ -+ @fx.union -+ class StageLds: -+ early: EarlyLds -+ attn: AttnLds -+ ug: UgLds -+ down: DownLds -+ -+ @fx.struct -+ class Smem: -+ red: fx.Array[fx.Float32, WAVES * 64 * 4, 16] -+ -+ keyed = key_tuple(key, _SOURCES) -+ name = f"k3_mono_k2_s{s}_nb{key.nblocks}_x{int(key.reset)}_r{int(key.reduce)}" -+ if key.fault_rank >= 0: -+ name += ( -+ f"_f{key.fault_rank}_{key.fault_epoch}_{key.fault_layer}_{key.fault_ticks}" -+ ) -+ -+ @flyc.kernel(name=name, known_block_size=[THREADS, 1, 1]) -+ def k2( -+ core: Int64, -+ w_o: Int64, -+ prefix: Int64, -+ blocks: Int64, -+ blk_sm: Int32, -+ blk_sr: Int32, -+ ares_nw: Int64, -+ ares_qk: Int64, -+ in_nw: Int64, -+ w_gate: Int64, -+ bias: Int64, -+ w_ld: Int64, -+ w_sgu: Int64, -+ w_sd: Int64, -+ w13: Int64, -+ w13s: Int64, -+ w2: Int64, -+ w2s: Int64, -+ ln_w: Int64, -+ w_up: Int64, -+ out: Int64, -+ scratch: Int64, -+ peers: Int64, -+ rank: Int32, -+ epoch: Int64, -+ layer: Int32, -+ health: Int64, -+ alarm: Int64, -+ ): -+ _ = keyed -+ tid = fx.thread_idx.x -+ bid = fx.block_idx.x -+ alloc = fx.SharedAllocator() -+ lds = alloc.allocate(Smem).peek() -+ rl = alloc.allocate(RouteLds).peek() -+ st = alloc.allocate(StageLds) -+ el, al, ul, dl = st.early.peek(), st.attn.peek(), st.ug.peek(), st.down.peek() -+ if const_expr(key.fault_rank >= 0): -+ stall( -+ key.fault_rank, -+ key.fault_epoch, -+ key.fault_layer, -+ key.fault_ticks, -+ rank, -+ epoch, -+ layer, -+ ) -+ mb = Mailbox(step_tag(epoch, layer) - 1, health, alarm, region_ids(REGIONS)) -+ par = uniform(ld_i32(epoch, 0)) % 2 -+ bases = peer_bases(peers, TP_) -+ own = load_ptr64(peers, rank) -+ peer = {} -+ for region in ("rlog", "lat", "latr", "final", "ar"): -+ off = play[region][0] -+ peer[region] = [ -+ shift(preg(b, off, region), fx.Int64(par) * half) for b in bases -+ ] -+ peer[region + "_own"] = shift(preg(own, off, region), fx.Int64(par) * half) -+ c = { -+ "S": s, -+ "NS": ns, -+ "NB": key.nblocks, -+ "tid": tid, -+ "bid": bid, -+ "lane": tid % 64, -+ "wave": tid // 64, -+ "rank": rank, -+ "delta": True, -+ "reset": key.reset, -+ "eps": key.eps, -+ "out_eps": key.out_eps, -+ "ln_eps": key.ln_eps, -+ "put": mb.put, -+ "put_bf": mb.put_bf, -+ "put_words": mb.put_words, -+ "poll": mb.poll, -+ "red": lds.red.ptr, -+ "peer": peer, -+ "ar_peer": peer["ar"], -+ "ar_own": peer["ar_own"], -+ "xl": el.xl.ptr, -+ "vals": al.vals.ptr, -+ "delta_lds": al.dl.ptr, -+ "latq": ul.latq.ptr, -+ "latsc": ul.latsc.ptr, -+ "av": ul.av.ptr, -+ "interq": dl.interq.ptr, -+ "intsc": dl.intsc.ptr, -+ "contrib": dl.contrib.ptr, -+ "scan": rl.scan.ptr, -+ "rt": { -+ "route": rl.route.ptr, -+ "flag": rl.flag.ptr, -+ "expert": rl.expert.ptr, -+ "kof": rl.kof.ptr, -+ }, -+ "xbuf": scratch + fx.Int64(lay["xrow"][0]), -+ "xflags": NA, -+ "ln": scratch + fx.Int64(lay["ln"][0]), -+ "shp": scratch + fx.Int64(lay["shp"][0]), -+ "args": { -+ "core": core, -+ "w_o": w_o, -+ "prefix": prefix, -+ "blocks": blocks, -+ "blk_sm": blk_sm, -+ "blk_sr": blk_sr, -+ "ares_nw": ares_nw, -+ "ares_qk": ares_qk, -+ "in_nw": in_nw, -+ "x_dbg": 0, -+ "w_gate": w_gate, -+ "bias": bias, -+ "w_ld": w_ld, -+ "w_sgu": w_sgu, -+ "w_sd": w_sd, -+ "w13": w13, -+ "w13s": w13s, -+ "w2": w2, -+ "w2s": w2s, -+ "ln_w": ln_w, -+ "w_up": w_up, -+ "out": out, -+ }, -+ } -+ for region in ( -+ "part", -+ "wgt", -+ "msq", -+ "xrdy", -+ "sgu", -+ "h", -+ "rpart", -+ "lpart", -+ "route", -+ "inter", -+ "isc", -+ "lmsq", -+ "lnrdy", -+ "shrdy", -+ ): -+ c[region] = sreg(scratch, lay[region][0], region) -+ # ---- o_proj on every CTA past AttnRes's -+ eb = bid - E0 -+ lb = eb - k2b.R_TASKS -+ is_r = (bid >= E0) & (eb < k2b.R_TASKS) -+ is_l = (bid >= E0 + k2b.R_TASKS) & (lb < k2b.L_TASKS) -+ # o_proj first (every AttnRes CTA waits on the slowest partial; loads -+ # retire in order), the early weights after: they land before x does -+ if is_r: -+ stage_oproj(c, bid - O0) -+ rops = k2b.early_loads(c, 0, eb // k2b.R_PARTS, eb % k2b.R_PARTS) -+ k2b.early_compute(c, 0, eb // k2b.R_PARTS, eb % k2b.R_PARTS, rops) -+ if is_l: -+ stage_oproj(c, bid - O0) -+ lops = k2b.early_loads(c, 1, lb // k2b.L_PARTS, lb % k2b.L_PARTS) -+ k2b.early_compute(c, 1, lb // k2b.L_PARTS, lb % k2b.L_PARTS, lops) -+ if (bid >= O0) & ~is_r & ~is_l: -+ stage_oproj(c, bid - O0) -+ gi = bid - O0 -+ if gi < k2b.G_TASKS: -+ k2b.early_compute(c, 2, gi, 0, k2b.early_loads(c, 2, gi, 0)) -+ # ---- AttnRes: the reduced o_proj -> x; then route / SiTU -+ if bid < NA: -+ k2a.stage_reduce(c, bid) -+ stage_slice(c, bid) -+ if bid < s: -+ stage_gate(c, bid) -+ stage_mix(c, bid) -+ if bid < s: -+ k2b.stage_route(c, bid) -+ if (bid >= H0) & (bid < H0 + s): -+ k2b.stage_h(c, bid - H0) -+ # ---- the routed experts -+ # the latent (ready ~10 us before the route) quantized first -+ k2b.load_latq(c) -+ nu = k2b.load_route(c) -+ total = nu * k2b.UG_GROUPS -+ last = total - 1 -+ # reversed placement: the round's leftover tasks go to the high CTAs, -+ # whose early work ends first (the AttnRes / route CTAs start ug last) -+ task = BLOCKS - 1 - bid -+ ops = k2b.ug_loads(c, fx.min(task, last)) -+ while task < total: -+ nxt = task + BLOCKS -+ ops_n = k2b.ug_loads(c, fx.min(nxt, last)) -+ k2b.stage_ug(c, task, ops) -+ task = nxt -+ ops = ops_n -+ if bid < k2b.DOWN_TASKS: -+ dops = k2b.down_loads(c, bid, nu, fx.min(c["wave"], nu - 1)) -+ k2b.load_inter(c) -+ k2b.stage_down(c, bid, nu, dops) -+ k2b.run_sdown_rest(c, bid, w_sd) -+ k2b.run_sdown(c, bid, w_sd) -+ if (bid >= k2b.LN0) & (bid < k2b.LN0 + k2b.LN_TASKS): -+ k2b.stage_lnorm(c, bid - k2b.LN0) -+ if bid < k2b.UP_TASKS: -+ k2b.stage_up(c, bid) -+ if const_expr(key.reduce): # noqa: SIM102 (traced: `and` does not combine traced values) -+ if (bid >= k2b.RED0) & (bid < k2b.RED0 + k2b.RED_TASKS): -+ k2b.stage_reduce(c, bid - k2b.RED0) -+ -+ @flyc.jit -+ def launch( -+ core: Int64, -+ w_o: Int64, -+ prefix: Int64, -+ blocks: Int64, -+ blk_sm: Int32, -+ blk_sr: Int32, -+ ares_nw: Int64, -+ ares_qk: Int64, -+ in_nw: Int64, -+ w_gate: Int64, -+ bias: Int64, -+ w_ld: Int64, -+ w_sgu: Int64, -+ w_sd: Int64, -+ w13: Int64, -+ w13s: Int64, -+ w2: Int64, -+ w2s: Int64, -+ ln_w: Int64, -+ w_up: Int64, -+ out: Int64, -+ scratch: Int64, -+ peers: Int64, -+ rank: Int32, -+ epoch: Int64, -+ layer: Int32, -+ health: Int64, -+ alarm: Int64, -+ stream: fx.Stream = _STREAM, -+ ): -+ _ = keyed -+ k2( -+ core, -+ w_o, -+ prefix, -+ blocks, -+ blk_sm, -+ blk_sr, -+ ares_nw, -+ ares_qk, -+ in_nw, -+ w_gate, -+ bias, -+ w_ld, -+ w_sgu, -+ w_sd, -+ w13, -+ w13s, -+ w2, -+ w2s, -+ ln_w, -+ w_up, -+ out, -+ scratch, -+ peers, -+ rank, -+ epoch, -+ layer, -+ health, -+ alarm, -+ ).launch(grid=(BLOCKS,), block=(THREADS,), stream=stream) -+ -+ ABI.check(k2, launch) -+ _BUILDS[key] = launch -+ return launch -+ -+ -+TP_ = k2b.TP -+ -+ -+def k2_launch( -+ key: K2Build, -+ *, -+ core, -+ w_o, -+ prefix, -+ blocks, -+ ares_nw, -+ ares_qk, -+ in_nw, -+ w_gate, -+ bias, -+ w_ld, -+ w_sgu, -+ w_sd, -+ w13, -+ w13s, -+ w2, -+ w2s, -+ ln_w, -+ w_up, -+ out, -+ scratch, -+ peers, -+ rank, -+ epoch, -+ layer, -+ health, -+ alarm, -+): -+ """Host launcher: ``core`` [S, 1536] (o_proj's input), ``prefix`` updated in -+ place, ``out`` the MoE output [S, 7168]; MoE weights as ``moe``.""" -+ f = build(key) -+ f( -+ core.data_ptr(), -+ w_o.data_ptr(), -+ prefix.data_ptr(), -+ blocks.data_ptr(), -+ blocks.stride(0), -+ blocks.stride(1), -+ ares_nw.data_ptr(), -+ ares_qk.data_ptr(), -+ in_nw.data_ptr(), -+ w_gate.data_ptr(), -+ bias.data_ptr(), -+ w_ld.data_ptr(), -+ w_sgu.data_ptr(), -+ w_sd.data_ptr(), -+ w13.data_ptr(), -+ w13s.data_ptr(), -+ w2.data_ptr(), -+ w2s.data_ptr(), -+ ln_w.data_ptr(), -+ w_up.data_ptr(), -+ out.data_ptr(), -+ scratch.data_ptr(), -+ peers.data_ptr(), -+ rank, -+ epoch.data_ptr(), -+ layer, -+ health.data_ptr(), -+ alarm, -+ stream=torch.cuda.current_stream(), -+ ) -diff --git a/vllm/models/kimi_k3/amd/specialized_decode/mono/stages/__init__.py b/vllm/models/kimi_k3/amd/specialized_decode/mono/stages/__init__.py -new file mode 100644 -index 000000000..208f01a7c ---- /dev/null -+++ b/vllm/models/kimi_k3/amd/specialized_decode/mono/stages/__init__.py -@@ -0,0 +1,2 @@ -+# SPDX-License-Identifier: Apache-2.0 -+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -diff --git a/vllm/models/kimi_k3/amd/specialized_decode/mono/stages/gemv.py b/vllm/models/kimi_k3/amd/specialized_decode/mono/stages/gemv.py -new file mode 100644 -index 000000000..26acddb5a ---- /dev/null -+++ b/vllm/models/kimi_k3/amd/specialized_decode/mono/stages/gemv.py -@@ -0,0 +1,126 @@ -+# SPDX-License-Identifier: Apache-2.0 -+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -+"""bf16 GEMV pieces the K3 mono kernels share: 16 weight rows a task against up -+to 16 token rows in LDS, one ``mfma_f32_16x16x32_bf16`` per 32 K. -+ -+K is walked in 128-wide chunks, wave w taking chunks w, w + 8, ...; a lane reads -+64 contiguous bytes of its row a chunk (4 dwordx4 at q 32 + g 8 elements: lane -+group g = lane / 16), and the token rows are read at the same K. The K order -+inside a chunk is a permutation both operands share, so the dot is unchanged. -+""" -+ -+import flydsl.expr as fx -+from aiter.ops.flydsl.kernels import buffer_ops as bo -+from flydsl.expr import gpu, range_constexpr -+from flydsl.expr.typing import T -+ -+from vllm.models.kimi_k3.amd.specialized_decode.mono.common.ops import ( -+ CM_DEV, -+ CM_NT, -+ mfma_bf16, -+ rsrc, -+ traced, -+) -+from vllm.models.kimi_k3.amd.specialized_decode.mono.common.plan import THREADS, WAVES -+ -+ROWS = 16 -+KCH = 128 -+LDS_PAD = 8 # bf16 a padded token row -+ -+ -+def chunks(k): -+ """A wave's K chunks for K = ``k``: (count a wave, total).""" -+ n = k // KCH -+ return (n + WAVES - 1) // WAVES, n -+ -+ -+def loads(lane, wave, w, row0, k, nt=True, ldk=None, k0=0): -+ """This wave's operands of weight rows ``row0`` .. + 16 (row-major, ``ldk`` -+ bf16 a row at ``w``, K columns ``k0`` .. ``k0 + k``): a list a chunk -+ (clamped into K: ``mfmas`` zeroes the extra).""" -+ per, n = chunks(k) -+ ldk = k if ldk is None else ldk -+ r_w = rsrc(w) -+ row = row0 + lane % ROWS -+ g = lane // ROWS -+ out = [] -+ for i in range(per): -+ ch = fx.min(wave + WAVES * i, n - 1) -+ base = (row * ldk + k0 + ch * KCH + g * 8) // 2 -+ out.append( -+ [ -+ fx.Vector( -+ bo.buffer_load( -+ r_w, -+ base + 16 * q, -+ vec_width=4, -+ dtype=T.i32, -+ cache_modifier=CM_NT if nt else 0, -+ ) -+ ) -+ for q in range(4) -+ ] -+ ) -+ return out -+ -+ -+def mfmas(lane, wave, xl, xrow, k, s, ops): -+ """``loads``' operands against token rows 0 .. s - 1 of the LDS rows ``xl`` -+ (``xrow`` bf16 a row) -> this wave's 16x16 accumulator.""" -+ per, n = chunks(k) -+ g = lane // ROWS -+ t = fx.min(lane % ROWS, s - 1) -+ acc = fx.Vector.filled(4, 0.0, fx.Float32) -+ for i in range_constexpr(per): -+ ch = wave + WAVES * i -+ live = ch < n -+ kk = fx.min(ch, n - 1) * KCH + g * 8 -+ for q in range_constexpr(4): -+ xb = fx.Vector( -+ fx.ptr_load( -+ xl + (t * xrow + kk + 32 * q) // 2, -+ result_type=fx.Vector.make_type(4, fx.Int32), -+ ) -+ ) -+ if const_live(per, n): -+ xb = fx.Vector.from_elements( -+ [live.select(xb[d], fx.Int32(0)) for d in range(4)], fx.Int32 -+ ) -+ acc = mfma_bf16( -+ ops[i][q].bitcast(fx.BFloat16), xb.bitcast(fx.BFloat16), acc -+ ) -+ return acc -+ -+ -+def const_live(per, n): -+ """Whether some wave's last chunk is past K (its MFMA must add zero).""" -+ return per * WAVES != n -+ -+ -+@traced -+def rows_to_lds(tid, src, k, s, xl, xrow, cm=0, ldk=None, k0=0): -+ """Columns k0 .. k0 + k of token rows 0 .. s - 1 of the bf16 [s, ldk] at -+ ``src`` -> LDS (``xrow`` a row), 16 B a lane.""" -+ ldk = k if ldk is None else ldk -+ words = s * k // 8 -+ for i in range_constexpr((words + THREADS - 1) // THREADS): -+ e = fx.min(tid + THREADS * i, words - 1) -+ t = e // (k // 8) -+ kk = e % (k // 8) * 8 -+ v = fx.Vector( -+ bo.buffer_load( -+ rsrc(src), -+ (t * ldk + k0 + kk) // 2, -+ vec_width=4, -+ dtype=T.i32, -+ cache_modifier=cm, -+ ) -+ ) -+ if tid + THREADS * i < words: -+ fx.ptr_store(v, xl + (t * xrow + kk) // 2) -+ gpu.barrier() -+ -+ -+def dev_rows_to_lds(tid, src, k, s, xl, xrow): -+ """``rows_to_lds`` at device scope: rows another CTA of this launch wrote.""" -+ rows_to_lds(tid, src, k, s, xl, xrow, CM_DEV) -diff --git a/vllm/models/kimi_k3/amd/specialized_decode/mono/stages/moe.py b/vllm/models/kimi_k3/amd/specialized_decode/mono/stages/moe.py -new file mode 100644 -index 000000000..01a0f861e ---- /dev/null -+++ b/vllm/models/kimi_k3/amd/specialized_decode/mono/stages/moe.py -@@ -0,0 +1,1025 @@ -+# SPDX-License-Identifier: Apache-2.0 -+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -+"""K2b ``moe``: one Kimi-K3 latent-MoE layer, from the MoE input row x to the -+layer's output, every TP reduction inside the kernel. -+ -+One launch a layer and rank, ``BLOCKS`` x ``THREADS``, S <= 8 rows (a -+spec-verify step). The original path (vLLM ROCm, see MOE_SPEC.md) replicates the -+router and the latent down projection on every rank; here each rank computes -+its eighth and the rows are all-gathered over the peer buffer. -+ -+ early (CTAs 0 .. 130, one task each, x from the input row; weights first): -+ router (7 x 16 rows of this rank's 112 experts): fp32 logits -> every -+ rank's RLOG -+ latent (28 x 16 rows of this rank's 448): bf16 -> every rank's LAT -+ sgu (96 x 16 rows of the shared gate | up): bf16 -> SGU -+ route (S, CTAs 131 ..): sigmoid (v_rcp), + bias, top-16, the pre-bias -+ sigmoids renormalized (IEEE 1 / sum) -> ROUTE (every rank alike) -+ sdown (448 x 16 rows, K 768, CTAs 131 ..): SiTU of SGU (vLLM's kernel) -> -+ bf16 -> GEMV -> bf16 partial; this rank's up columns -> SHP (plain, -+ SHRDY flags), the rest -> every rank's FINAL -+ ug (|U| x 12, every CTA): the union U of the step's picks; an expert's -+ 32 gate + 32 up rows against the MXFP4 latent (aiter's inline quant), -+ SiTUv2 (aiter's fp32 epilogue), MXFP4 -> INTER / ISC a pick -+ down (224 x 16 latent rows, every CTA): w2 of every expert of U against -+ its picks' INTER, bf16(w acc) folded in each token's top-k order with -+ a bf16 rounding an add -> every rank's LATR -+ lnorm (28 x 128 columns): the ranks' LATR in rank order, bf16, RMSNorm -+ (the fused AR + RMSNorm form) -> LN (plain, LNRDY flags) -+ up (56 x 16 rows of this rank's 896): LN GEMV + SHP -> bf16 -> every -+ rank's FINAL -+ reduce (56 x 128 columns, ``reduce``): FINAL in rank order -> bf16 out -+ -+Peer regions are double-buffered by the epoch's parity. -+""" -+ -+import math -+ -+import flydsl.expr as fx -+from aiter.ops.flydsl.kernels import buffer_ops as bo -+from flydsl.expr import const_expr, gpu, range_constexpr, rocdl -+from flydsl.expr.typing import T, as_ir_value -+ -+from vllm.models.kimi_k3.amd.specialized_decode.mono.common.mx import ( -+ FP4, -+ UNIT_SCALE, -+ fp4_tile_load, -+ mfma_scaled, -+ scale_index, -+) -+from vllm.models.kimi_k3.amd.specialized_decode.mono.common.ops import ( -+ CM_DEV, -+ ballot, -+ bf16_round, -+ bf_hi, -+ bf_lo, -+ block_excl_scan, -+ butterfly, -+ hw_exp2, -+ hw_rcp, -+ hw_rsq, -+ lanes_below, -+ popcount, -+ row_sum, -+ rsrc, -+ traced, -+ uniform, -+ wave_sum, -+) -+from vllm.models.kimi_k3.amd.specialized_decode.mono.common.plan import ( -+ BLOCKS, -+ THREADS, -+ WAVES, -+ pair_layout, -+) -+from vllm.models.kimi_k3.amd.specialized_decode.mono.common.ranks import sum_partials -+from vllm.models.kimi_k3.amd.specialized_decode.mono.common.sync import publish -+from vllm.models.kimi_k3.amd.specialized_decode.mono.stages import gemv as gemv -+ -+REGIONS = ( -+ "sgu", -+ "h", -+ "rpart", -+ "lpart", -+ "route", -+ "inter", -+ "isc", -+ "lmsq", -+ "lnrdy", -+ "shrdy", -+ "rlog", -+ "lat", -+ "latr", -+ "final", -+) -+ -+ -+TP = 8 -+HIDDEN = 7168 -+LAT = 3584 -+E = 896 -+TOPK = 16 -+RI = 384 # routed intermediate a rank -+SI = 768 # shared intermediate a rank -+ROWS = 16 -+LOG2E = 1.0 / math.log(2.0) -+BETA, LBETA = 4.0, 25.0 -+ -+R_GROUPS = E // TP // ROWS # 7 router row groups a rank -+R_PARTS = 4 # K quarters: the router is on the path to the route -+L_GROUPS = LAT // TP // ROWS # 28 latent row groups a rank -+L_PARTS = 2 # K halves: the latent is on the path to ug -+R_TASKS = R_GROUPS * R_PARTS # 28 -+L_TASKS = L_GROUPS * L_PARTS # 56 -+G_TASKS = 2 * SI // ROWS # 96 -+EARLY = R_TASKS + L_TASKS + G_TASKS # 180 -+LATE0 = EARLY # route from here -+H_CTA = LATE0 + 8 # the shared experts' SiTU -+SPARE0 = H_CTA + 8 # 196: CTAs with no early task run the shared down first -+SD_TASKS = HIDDEN // ROWS # 448 -+UP_N = HIDDEN // TP # 896 columns a rank up-projects -+UP_TASKS = UP_N // ROWS # 56 -+LN_TASKS = LAT // 128 # 28 -+LN0 = 56 # lnorm's CTAs -+RED_TASKS = HIDDEN // 128 # 56 -+RED0 = 200 -+UG_GROUPS = RI // 32 # 12 -+W13_K = LAT # 3584 -+W13_STEPS = W13_K // 128 # 28 -+W2_STEPS = RI // 128 # 3 -+DOWN_TASKS = LAT // ROWS # 224 -+DOWN_BATCH = 5 # a wave's experts whose w2 is in flight together -+MAX_U = 8 * TOPK -+ -+_STREAM = fx.Stream(None) -+ -+ -+def peer_pairs(s): -+ """One parity's peer regions, pairs: name -> count.""" -+ return { -+ "rlog": TP * s * (E // TP), -+ "lat": s * LAT // 2, -+ "latr": TP * s * LAT // 2, -+ "final": TP * s * HIDDEN // 2, -+ } -+ -+ -+def scratch_layout(s: int) -> dict: -+ pairs = { -+ "sgu": s * SI, -+ "h": s * SI // 2, -+ "rpart": R_GROUPS * R_PARTS * s * ROWS, -+ "lpart": L_GROUPS * L_PARTS * s * ROWS, -+ "route": s * 2 * TOPK, -+ "inter": s * TOPK * 48, -+ "isc": s * TOPK * UG_GROUPS, -+ "lmsq": LN_TASKS * s, -+ "lnrdy": LN_TASKS, -+ "shrdy": UP_TASKS, -+ } -+ lay = pair_layout(pairs.items()) -+ end = max(o + n for o, n in lay.values()) -+ lay["ln"] = ((end + 15) // 16 * 16, s * LAT * 2) -+ end = lay["ln"][0] + lay["ln"][1] -+ lay["shp"] = ((end + 15) // 16 * 16, s * UP_N * 2) -+ return lay -+ -+ -+def f32_of(w): -+ return w.bitcast(fx.Float32) -+ -+ -+def tanh_libm(x): -+ """The ``tanhf`` of x to a few ulp: (1 - e) / (1 + e), e = exp(-2|x|), x's sign.""" -+ e = hw_exp2(abs(x) * fx.Float32(-2.0 * LOG2E)) -+ t = (fx.Float32(1.0) - e) / (fx.Float32(1.0) + e) -+ return (x < fx.Float32(0.0)).select(-t, t) -+ -+ -+def situ_vllm(g, u): -+ """The ``situ_and_mul`` of vLLM on bf16 g, u: -+ ((2 tanh(g/4)) (1 + tanh(g/2))) (25 tanh(u/25)).""" -+ a = fx.Float32(0.5 * BETA) * tanh_libm(g * fx.Float32(1.0 / BETA)) -+ a = a * (fx.Float32(1.0) + tanh_libm(g * fx.Float32(0.5))) -+ return a * (fx.Float32(LBETA) * tanh_libm(u * fx.Float32(1.0 / LBETA))) -+ -+ -+def tanh_aiter(x): -+ """The ``act.tanh_f32`` of aiter: (1 - e) rcp(1 + e), e = exp2(|x| (-2 log2 e)).""" -+ e = hw_exp2(abs(x) * fx.Float32(-2.0 * LOG2E)) -+ t = (fx.Float32(1.0) - e) * hw_rcp(fx.Float32(1.0) + e) -+ return (x > fx.Float32(0.0)).select(t, -t) -+ -+ -+def situ_aiter(g, u): -+ """The situv2 epilogue of aiter on fp32 accumulators.""" -+ sg = hw_rcp(fx.Float32(1.0) + hw_exp2(g * fx.Float32(-LOG2E))) -+ a = (fx.Float32(BETA) * tanh_aiter(g * fx.Float32(1.0 / BETA))) * sg -+ return (a * fx.Float32(LBETA)) * tanh_aiter(u * fx.Float32(1.0 / LBETA)) -+ -+ -+def e8m0_roundup(amax): -+ """The ``_e8m0_roundup(amax, 6)`` of aiter: ceil code of amax * f32(1/6), -+ at most 254.""" -+ wi = (amax * fx.Float32(1.0 / 6.0)).bitcast(fx.Int32) -+ b = fx.Int32(fx.Uint32(wi + fx.Int32(0x7FFFFF)) >> fx.Uint32(23)) & 0xFF -+ return fx.min(b, fx.Int32(254)) -+ -+ -+def pk_fp4_f32(old, a, b, scale, j): -+ """Two f32 / scale -> E2M1 nibbles into byte j of ``old`` (RNE, saturating).""" -+ return fx.Int32( -+ rocdl.cvt_scalef32_pk_fp4_f32( -+ T.i32, -+ as_ir_value(fx.Int32(old)), -+ as_ir_value(fx.Float32(a)), -+ as_ir_value(fx.Float32(b)), -+ as_ir_value(fx.Float32(scale)), -+ j, -+ ) -+ ) -+ -+ -+def pk_fp4_bf16(old, word, scale, j): -+ """A packed bf16 pair word / scale -> E2M1 nibbles into byte j of ``old``.""" -+ v = fx.Vector.from_elements([fx.Int32(word)], fx.Int32).bitcast(fx.BFloat16) -+ return fx.Int32( -+ rocdl.cvt_scalef32_pk_fp4_bf16( -+ T.i32, -+ as_ir_value(fx.Int32(old)), -+ as_ir_value(v), -+ as_ir_value(fx.Float32(scale)), -+ j, -+ ) -+ ) -+ -+ -+# ---------------------------------------------------------------- early -+@traced -+def early_out(c, kind, i, q, red): -+ """Task (i, q) of ``kind`` (0 router, 1 latent, 2 sgu): its 16 rows. A K-split -+ kind's parts q > 0 publish fp32 partials; part 0 adds them in part order and -+ publishes the rows (router: fp32 to every rank's RLOG, latent: bf16 to every -+ rank's LAT); sgu bf16 to SGU.""" -+ s, tid = c["S"], c["tid"] -+ if const_expr(kind == 2): -+ if tid < (ROWS // 2) * s: -+ rp = tid % (ROWS // 2) -+ t = tid // (ROWS // 2) -+ v0 = bf16_round(row_sum(red, 2 * rp, t)) -+ v1 = bf16_round(row_sum(red, 2 * rp + 1, t)) -+ c["put_bf"](c["sgu"], t * 2 * SI + i * ROWS + 2 * rp, [v0, v1]) -+ else: -+ parts = R_PARTS if kind == 0 else L_PARTS -+ region = c["rpart"] if kind == 0 else c["lpart"] -+ if tid < ROWS * s: -+ r = tid % ROWS -+ t = tid // ROWS -+ v = row_sum(red, r, t) -+ at = lambda qq: ((i * parts + qq) * s + t) * ROWS + r # noqa: E731 -+ if q > 0: -+ c["put"](region, at(q), v) -+ if q == 0: -+ got = c["poll"]([(region, at(qq), 1) for qq in range(1, parts)]) -+ for qq in range_constexpr(parts - 1): -+ v = v + f32_of(got[qq][0]) -+ fx.ptr_store(v, c["av"] + tid) -+ gpu.barrier() -+ if const_expr(kind == 0): -+ if (q == 0) & (tid < ROWS * s): -+ r = tid % ROWS -+ t = tid // ROWS -+ v = fx.ptr_load(c["av"] + tid) -+ for p in range_constexpr(TP): -+ c["put"]( -+ c["peer"]["rlog"][p], -+ (c["rank"] * s + t) * (E // TP) + i * ROWS + r, -+ v, -+ ) -+ else: -+ if (q == 0) & (tid < (ROWS // 2) * s): -+ rp = tid % (ROWS // 2) -+ t = tid // (ROWS // 2) -+ v0 = bf16_round(fx.ptr_load(c["av"] + (t * ROWS + 2 * rp))) -+ v1 = bf16_round(fx.ptr_load(c["av"] + (t * ROWS + 2 * rp + 1))) -+ col = c["rank"] * (LAT // TP) + i * ROWS + 2 * rp -+ for p in range_constexpr(TP): -+ c["put_bf"](c["peer"]["lat"][p], t * LAT + col, [v0, v1]) -+ gpu.barrier() -+ -+ -+def early_shape(kind): -+ """(K a part, parts) of an early kind: router a quarter, latent a half, sgu all.""" -+ return { -+ 0: (HIDDEN // R_PARTS, R_PARTS), -+ 1: (HIDDEN // L_PARTS, L_PARTS), -+ 2: (HIDDEN, 1), -+ }[kind] -+ -+ -+def early_loads(c, kind, i, q): -+ """Early task (i, q) of ``kind``'s weights for this wave: issued as early as -+ the CTA can, then ``early_compute`` takes them.""" -+ a = c["args"] -+ klen, _ = early_shape(kind) -+ if kind == 0: -+ w, row0 = a["w_gate"], c["rank"] * (E // TP) + i * ROWS -+ elif kind == 1: -+ w, row0 = a["w_ld"], c["rank"] * (LAT // TP) + i * ROWS -+ else: -+ w, row0 = a["w_sgu"], i * ROWS -+ return gemv.loads(c["lane"], c["wave"], w, row0, klen, ldk=HIDDEN, k0=q * klen) -+ -+ -+@traced -+def early_compute(c, kind, i, q, ops): -+ """Early task (i, q): x's K slice (the input row, or -- ``xflags`` -- the -+ row this launch's AttnRes stage publishes) against ``ops`` -> ``early_out``.""" -+ s, tid, lane, wave, red = c["S"], c["tid"], c["lane"], c["wave"], c["red"] -+ klen, _ = early_shape(kind) -+ xrow = klen + gemv.LDS_PAD -+ if const_expr(c.get("xflags") is not None): -+ if tid < c["xflags"]: -+ c["poll"]([(c["xrdy"], tid, 1)]) -+ gpu.barrier() -+ gemv.rows_to_lds( -+ tid, c["xbuf"], klen, s, c["xl"], xrow, CM_DEV, ldk=HIDDEN, k0=q * klen -+ ) -+ else: -+ gemv.rows_to_lds( -+ tid, c["args"]["x"], klen, s, c["xl"], xrow, ldk=HIDDEN, k0=q * klen -+ ) -+ acc = gemv.mfmas(lane, wave, c["xl"], xrow, klen, s, ops) -+ fx.ptr_store(acc, red + (wave * 64 + lane) * 4) -+ gpu.barrier() -+ early_out(c, kind, i, q, red) -+ -+ -+@traced -+def stage_early(c, bid): -+ """K2b alone: CTA bid's early task (router parts, latent parts, sgu).""" -+ lb = bid - R_TASKS -+ gb = bid - R_TASKS - L_TASKS -+ if bid < R_TASKS: -+ early_compute( -+ c, -+ 0, -+ bid // R_PARTS, -+ bid % R_PARTS, -+ early_loads(c, 0, bid // R_PARTS, bid % R_PARTS), -+ ) -+ if (bid >= R_TASKS) & (bid < R_TASKS + L_TASKS): -+ early_compute( -+ c, -+ 1, -+ lb // L_PARTS, -+ lb % L_PARTS, -+ early_loads(c, 1, lb // L_PARTS, lb % L_PARTS), -+ ) -+ if bid >= R_TASKS + L_TASKS: -+ early_compute(c, 2, gb, 0, early_loads(c, 2, gb, 0)) -+ -+ -+# ---------------------------------------------------------------- route -+def _okey(v): -+ """An order-preserving uint32 key of an f32 (larger float, larger key).""" -+ b = v.bitcast(fx.Int32) -+ neg = b < fx.Int32(0) -+ return fx.Uint32(neg.select(~b, b ^ fx.Int32(-2147483648))) -+ -+ -+@traced -+def stage_route(c, t): -+ """Token t (wave 0): 896 logits, 14 a lane (expert lane + 64 i) -> sigmoid -+ (v_rcp) + bias -> the top-16 by a radix select (the 16th largest key by a -+ 32-step bit search of ballot counts; equal keys: the lower experts) -> the -+ 16 ranked by score, descending (aiter's slot order) -> the pre-bias -+ sigmoids renormalized in that order (IEEE 1 / sum).""" -+ lane, wave, red = c["lane"], c["wave"], c["red"] -+ a = c["args"] -+ if wave == 0: -+ s = c["S"] -+ per = E // 64 -+ specs = [] -+ for i in range_constexpr(per): -+ e = lane + 64 * i -+ src = e // (E // TP) -+ specs.append( -+ (c["peer"]["rlog_own"], (src * s + t) * (E // TP) + e % (E // TP), 1) -+ ) -+ got = c["poll"](specs) -+ keys, sigs = [], [] -+ for i in range_constexpr(per): -+ e = lane + 64 * i -+ lg = f32_of(got[i][0]) -+ sg = hw_rcp(fx.Float32(1.0) + hw_exp2(lg * fx.Float32(-LOG2E))) -+ sc = sg + bo_ld_f32(a["bias"], e) -+ # a NaN logit (a warmup's garbage row) ranks last -+ keys.append(_okey((sc == sc).select(sc, fx.Float32(-3.0e38)))) -+ sigs.append((sg == sg).select(sg, fx.Float32(0.0))) -+ -+ def count_ge(th): -+ n = fx.Int32(0) -+ for i in range_constexpr(per): -+ n = n + popcount(ballot(keys[i] >= th)) -+ return n -+ -+ th = fx.Uint32(0) -+ for bit in range_constexpr(31, -1, -1): -+ cand = th | fx.Uint32(1 << bit) -+ th = (count_ge(cand) >= TOPK).select(cand, th) -+ # selected: key > th, then the lowest experts with key == th -+ n_gt = fx.Int32(0) -+ for i in range_constexpr(per): -+ n_gt = n_gt + popcount(ballot(keys[i] > th)) -+ need = fx.Int32(TOPK) - n_gt -+ eq_before = fx.Int32(0) -+ sel_before = fx.Int32(0) -+ for i in range_constexpr(per): -+ eqb = ballot(keys[i] == th) -+ eq_rank = eq_before + lanes_below(eqb, lane) -+ sel = (keys[i] > th) | ((keys[i] == th) & (eq_rank < need)) -+ selb = ballot(sel) -+ pos = sel_before + lanes_below(selb, lane) -+ if sel: -+ fx.ptr_store(keys[i].bitcast(fx.Float32), red + pos * 3) -+ fx.ptr_store( -+ fx.Int32(lane + 64 * i).bitcast(fx.Float32), red + (pos * 3 + 1) -+ ) -+ fx.ptr_store(sigs[i], red + (pos * 3 + 2)) -+ eq_before = eq_before + popcount(eqb) -+ sel_before = sel_before + popcount(selb) -+ # rank the 16 by key descending, the lower expert first on a tie -+ q = fx.min(lane, TOPK - 1) -+ mk = fx.Uint32(fx.ptr_load(red + q * 3).bitcast(fx.Int32)) -+ me = fx.ptr_load(red + (q * 3 + 1)).bitcast(fx.Int32) -+ ms = fx.ptr_load(red + (q * 3 + 2)) -+ rank = fx.Int32(0) -+ for j in range_constexpr(TOPK): -+ ok = fx.Uint32(fx.ptr_load(red + j * 3).bitcast(fx.Int32)) -+ oe = fx.ptr_load(red + (j * 3 + 1)).bitcast(fx.Int32) -+ before = (ok > mk) | ((ok == mk) & (oe < me)) -+ rank = rank + before.select(fx.Int32(1), fx.Int32(0)) -+ # the sigmoids in rank order, summed serially (aiter's renorm) -+ if lane < TOPK: -+ fx.ptr_store(ms, red + (64 + rank)) -+ total = fx.Float32(0.0) -+ for j in range_constexpr(TOPK): -+ total = total + fx.ptr_load(red + (64 + j)) -+ inv = fx.Float32(1.0) / total -+ if lane < TOPK: -+ c["put"]( -+ c["route"], t * 2 * TOPK + rank, fx.min(fx.max(me, 0), fx.Int32(E - 1)) -+ ) -+ c["put"](c["route"], t * 2 * TOPK + TOPK + rank, ms * inv) -+ gpu.barrier() -+ -+ -+def bo_ld_f32(ptr, i): -+ return fx.Float32(bo.buffer_load(rsrc(ptr), i, vec_width=1, dtype=T.f32)) -+ -+ -+@traced -+def load_route(c): -+ """ROUTE -> LDS (ids, weights); the union U: ``slot_of[e]``, -+ ``expert_of[slot]``, ``kof[slot][t]`` (the token's top-k index, -1 if not -+ picked); returns |U|.""" -+ s, tid, lane, wave = c["S"], c["tid"], c["lane"], c["wave"] -+ rt = c["rt"] -+ n = s * 2 * TOPK -+ if tid < n: -+ v = c["poll"]([(c["route"], tid, 1)])[0][0] -+ fx.ptr_store(v, rt["route"] + tid) -+ for i in range_constexpr((E + THREADS - 1) // THREADS): -+ e = tid + THREADS * i -+ if e < E: -+ fx.ptr_store(fx.Int32(0), rt["flag"] + e) -+ for i in range_constexpr((MAX_U * 8 + THREADS - 1) // THREADS): -+ j = tid + THREADS * i -+ if j < MAX_U * 8: -+ fx.ptr_store(fx.Int32(-1), rt["kof"] + j) -+ gpu.barrier() -+ if tid < s * TOPK: -+ t = tid // TOPK -+ k = tid % TOPK -+ e = fx.max(fx.min(fx.ptr_load(rt["route"] + (t * 2 * TOPK + k)), E - 1), 0) -+ fx.ptr_store(fx.Int32(1), rt["flag"] + e) -+ gpu.barrier() -+ e0 = fx.min(2 * tid, E - 1) -+ f0 = (2 * tid < E).select(fx.ptr_load(rt["flag"] + e0), fx.Int32(0)) -+ f1 = (2 * tid + 1 < E).select( -+ fx.ptr_load(rt["flag"] + fx.min(2 * tid + 1, E - 1)), fx.Int32(0) -+ ) -+ off, tot = block_excl_scan(f0 + f1, lane, wave, c["scan"]) -+ gpu.barrier() -+ if f0 == 1: -+ fx.ptr_store(off, rt["flag"] + 2 * tid) -+ fx.ptr_store(fx.Int32(2 * tid), rt["expert"] + off) -+ if f1 == 1: -+ fx.ptr_store(off + f0, rt["flag"] + 2 * tid + 1) -+ fx.ptr_store(fx.Int32(2 * tid + 1), rt["expert"] + off + f0) -+ gpu.barrier() -+ if tid < s * TOPK: -+ t = tid // TOPK -+ k = tid % TOPK -+ e = fx.max(fx.min(fx.ptr_load(rt["route"] + (t * 2 * TOPK + k)), E - 1), 0) -+ slot = fx.ptr_load(rt["flag"] + e) -+ fx.ptr_store(fx.Int32(k), rt["kof"] + (slot * 8 + t)) -+ gpu.barrier() -+ return uniform(tot) -+ -+ -+@traced -+def load_latq(c): -+ """LAT (every rank's 448 columns) -> MXFP4 as aiter's inline quant: a -+ thread a (token, 32-block): amax of |bf16| bits, the ceil code of amax / 6, -+ v_cvt_scalef32_pk_fp4_bf16 -> latq words, latsc codes.""" -+ s, tid = c["S"], c["tid"] -+ nblk = s * (LAT // 32) -+ for i in range_constexpr((nblk + THREADS - 1) // THREADS): -+ b = fx.min(tid + THREADS * i, nblk - 1) -+ t = b // (LAT // 32) -+ kb = b % (LAT // 32) -+ base = (t * LAT + kb * 32) // 2 -+ got = c["poll"]([(c["peer"]["lat_own"], base + j, 1) for j in range(16)]) -+ ws = [got[j][0] for j in range(16)] -+ am = fx.Int32(0) -+ for w in ws: -+ am = fx.max(am, w & 0x7FFF) -+ am = fx.max(am, (w >> 16) & 0x7FFF) -+ amax = (am << 16).bitcast(fx.Float32) -+ code = e8m0_roundup(amax) -+ scale = (code << 23).bitcast(fx.Float32) -+ if tid + THREADS * i < nblk: -+ for q in range_constexpr(4): -+ pk = fx.Int32(0) -+ for j in range_constexpr(4): -+ pk = pk_fp4_bf16(pk, ws[4 * q + j], scale, j) -+ fx.ptr_store(pk, c["latq"] + (t * (LAT // 8) + kb * 4 + q)) -+ fx.ptr_store(code, c["latsc"] + (t * (LAT // 32) + kb)) -+ gpu.barrier() -+ -+ -+# ---------------------------------------------------------------- ug -+def w13_scale_word(e, rg, st, lane): -+ """The dword of aiter's gate_up=False ``shuffle_scale`` holding (row rg 16 + -+ lane % 16, K block 4 st + lane / 16) for st's pair and rg's pair: byte -+ (st % 2) 2 + rg % 2.""" -+ n1 = rg // 2 -+ k1 = st // 2 -+ return ( -+ e * (2 * RI * (W13_K // 32) // 4) -+ + ((n1 * (W13_STEPS // 2) + k1) * 4 + lane // 16) * 16 -+ + lane % 16 -+ ) -+ -+ -+def ug_loads(c, task): -+ """Task ``task``'s w13 operands for this wave (its row group, its K half) -+ and their scale words: issued a task ahead (``stage_ug`` takes them).""" -+ lane, wave = c["lane"], c["wave"] -+ a, rt = c["args"], c["rt"] -+ slot = task // UG_GROUPS -+ g = task % UG_GROUPS -+ e = uniform(fx.ptr_load(rt["expert"] + slot)) -+ rgl = wave % 4 # 0, 1 gate; 2, 3 up -+ half = wave // 4 -+ rg = (rgl // 2) * (RI // ROWS) + 2 * g + rgl % 2 -+ r_w = rsrc(a["w13"]) -+ base = e * (2 * RI * W13_K // 8) -+ steps = [half * (W13_STEPS // 2) + i for i in range(W13_STEPS // 2)] -+ av = [fp4_tile_load(r_w, base, rg, W13_K, st, lane * 4) for st in steps] -+ r_s = rsrc(a["w13s"]) -+ sw = [ -+ fx.Int32( -+ bo.buffer_load( -+ r_s, w13_scale_word(e, rg, st, lane), vec_width=1, dtype=T.i32 -+ ) -+ ) -+ for st in steps[::2] -+ ] -+ return av + sw -+ -+ -+@traced -+def stage_ug(c, task, ops): -+ """Slot task / 12, 32-column group task % 12: gate rows 32 g .. and up rows -+ 384 + 32 g .. of the slot's expert (``ops``: ``ug_loads``) against the MXFP4 -+ latent; SiTUv2; MXFP4 a 32-column group -> INTER / ISC of each picking -+ token.""" -+ s, tid, lane, wave, red = c["S"], c["tid"], c["lane"], c["wave"], c["red"] -+ rt = c["rt"] -+ slot = task // UG_GROUPS -+ g = task % UG_GROUPS -+ rgl = wave % 4 -+ half = wave // 4 -+ rg = (rgl // 2) * (RI // ROWS) + 2 * g + rgl % 2 -+ steps = [half * (W13_STEPS // 2) + i for i in range(W13_STEPS // 2)] -+ av = ops[: len(steps)] -+ sw = ops[len(steps) :] -+ t = fx.min(lane % 16, s - 1) -+ kb = lane // 16 -+ acc = fx.Vector.filled(4, 0.0, fx.Float32) -+ for i in range_constexpr(len(steps)): -+ st = steps[i] -+ bw = fx.Vector( -+ fx.ptr_load( -+ c["latq"] + (t * (LAT // 8) + st * 16 + kb * 4), -+ result_type=fx.Vector.make_type(4, fx.Int32), -+ ) -+ ) -+ sb = fx.ptr_load(c["latsc"] + (t * (LAT // 32) + st * 4 + kb)) -+ # the code's byte (st % 2) 2 + rg % 2 of the word; rg % 2 is the wave's -+ sa = (sw[i // 2] >> ((st % 2) * 16 + (rg % 2) * 8)) & 0xFF -+ acc = mfma_scaled(av[i], bw, acc, sa, sb, FP4, FP4, 0, 0) -+ fx.ptr_store(acc, red + (wave * 64 + lane) * 4) -+ gpu.barrier() -+ # thread (token, column): g / u summed over the K halves, SiTUv2 -+ tt = tid // 32 -+ col = tid % 32 -+ live = tt < s -+ tt = fx.min(tt, s - 1) -+ first = col < 16 # gate / up row group 0 (waves 0, 2) or 1 (1, 3); +4: K's 2nd half -+ r = col % 16 -+ gv = first.select( -+ row_sum(red, r, tt, waves=(0, 4)), row_sum(red, r, tt, waves=(1, 5)) -+ ) -+ uv = first.select( -+ row_sum(red, r, tt, waves=(2, 6)), row_sum(red, r, tt, waves=(3, 7)) -+ ) -+ v = situ_aiter(gv, uv) -+ amax = butterfly(abs(v), (16, 8, 4, 2, 1), fx.max) -+ code = e8m0_roundup(amax) -+ gpu.barrier() -+ fx.ptr_store(v, c["av"] + tid) -+ gpu.barrier() -+ k = fx.ptr_load(rt["kof"] + (slot * 8 + tt)) -+ if live & (col < 4) & (k >= 0): -+ scale = (code << 23).bitcast(fx.Float32) -+ pk = fx.Int32(0) -+ for j in range_constexpr(4): -+ x0 = fx.ptr_load(c["av"] + (tt * 32 + col * 8 + 2 * j)) -+ x1 = fx.ptr_load(c["av"] + (tt * 32 + col * 8 + 2 * j + 1)) -+ pk = pk_fp4_f32(pk, x0, x1, scale, j) -+ pick = tt * TOPK + k -+ c["put_words"](c["inter"], pick * 48 + g * 4 + col, [pk]) -+ if col == 0: -+ c["put"](c["isc"], pick * UG_GROUPS + g, code) -+ gpu.barrier() -+ -+ -+# ---------------------------------------------------------------- queues -+@traced -+def load_inter(c): -+ """Every pick's INTER words and ISC codes -> LDS.""" -+ s, tid = c["S"], c["tid"] -+ nw = s * TOPK * 48 -+ per = (nw + THREADS - 1) // THREADS -+ got = c["poll"]( -+ [(c["inter"], fx.min(tid + THREADS * i, nw - 1), 1) for i in range(per)] -+ ) -+ for i in range_constexpr(per): -+ if tid + THREADS * i < nw: -+ fx.ptr_store(got[i][0], c["interq"] + (tid + THREADS * i)) -+ nc = s * TOPK * UG_GROUPS -+ perc = (nc + THREADS - 1) // THREADS -+ gc = c["poll"]( -+ [(c["isc"], fx.min(tid + THREADS * i, nc - 1), 1) for i in range(perc)] -+ ) -+ for i in range_constexpr(perc): -+ if tid + THREADS * i < nc: -+ fx.ptr_store(gc[i][0], c["intsc"] + (tid + THREADS * i)) -+ gpu.barrier() -+ -+ -+def w2_scale(c, e, rgd, st, lane): -+ """The e8m0 code of w2 (row e 3584 + rgd 16 + lane % 16, K block 4 st + -+ lane / 16) in ``e8m0_shuffle``'s layout (16 codes a row).""" -+ row = e * LAT + rgd * ROWS + lane % 16 -+ idx = scale_index(row, st * 4 + lane // 16, 16) -+ w = fx.Int32( -+ bo.buffer_load(rsrc(c["args"]["w2s"]), idx // 4, vec_width=1, dtype=T.i32) -+ ) -+ return (w >> ((idx % 4) * 8)) & 0xFF -+ -+ -+def poll_inter(c, pick, kb): -+ """A lane's INTER words (3 steps x 4) and ISC codes of ``pick``, polled at -+ its use: a slot whose ug tasks are done is multiplied while later slots' -+ still run (no wait for all of U).""" -+ return c["poll"]( -+ [ -+ (c["inter"], pick * 48 + st * 16 + kb * 4 + d, 1) -+ for st in range(W2_STEPS) -+ for d in range(4) -+ ] -+ + [(c["isc"], pick * UG_GROUPS + st * 4 + kb, 1) for st in range(W2_STEPS)] -+ ) -+ -+ -+def down_loads(c, rgd, nu, b0): -+ """The w2 tiles and scale codes of a wave's batch from slot b0 (slots b0, -+ b0 + 8, ...; past |U| the last slot's): flat, ``stage_down`` unpacks.""" -+ lane = c["lane"] -+ r_w = rsrc(c["args"]["w2"]) -+ rt = c["rt"] -+ slots = [fx.min(b0 + WAVES * j, nu - 1) for j in range(DOWN_BATCH)] -+ es = [uniform(fx.ptr_load(rt["expert"] + sl)) for sl in slots] -+ av = [ -+ fp4_tile_load(r_w, es[j] * (LAT * RI // 8), rgd, RI, st, lane * 4) -+ for j in range(DOWN_BATCH) -+ for st in range(W2_STEPS) -+ ] -+ sa = [ -+ w2_scale(c, es[j], rgd, st, lane) -+ for j in range(DOWN_BATCH) -+ for st in range(W2_STEPS) -+ ] -+ return av + sa -+ -+ -+@traced -+def stage_down(c, rgd, nu, ops0): -+ """Latent rows 16 rgd ..: every slot's w2 against its picks (wave w the slots -+ w, w + 8, ...; a batch of DOWN_BATCH in flight, the next issued under this -+ one's MFMAs; ``ops0`` the first, issued before INTER was waited for) -> -+ bf16(w acc) a (token, k) in LDS -> folded in top-k order with a bf16 -+ rounding an add -> every rank's LATR.""" -+ s, tid, lane, wave = c["S"], c["tid"], c["lane"], c["wave"] -+ rt = c["rt"] -+ t = fx.min(lane % 16, s - 1) -+ kb = lane // 16 -+ nbs = DOWN_BATCH * W2_STEPS -+ b0 = wave -+ ops = ops0 -+ while b0 < nu: -+ nb = b0 + WAVES * DOWN_BATCH -+ ops_n = down_loads(c, rgd, nu, fx.min(nb, nu - 1)) -+ slots = [fx.min(b0 + WAVES * j, nu - 1) for j in range(DOWN_BATCH)] -+ for j in range_constexpr(DOWN_BATCH): -+ live = (b0 + WAVES * j) < nu -+ k = fx.ptr_load(rt["kof"] + (slots[j] * 8 + t)) -+ picked = (k >= 0) & (lane % 16 < s) -+ pick = t * TOPK + fx.max(k, 0) -+ acc = fx.Vector.filled(4, 0.0, fx.Float32) -+ # this slot's picks' INTER, polled here (not all of U first): a slot -+ # whose ug tasks are done is multiplied while later ones still run -+ for st in range_constexpr(W2_STEPS): -+ bw = fx.Vector( -+ fx.ptr_load( -+ c["interq"] + (pick * 48 + st * 16 + kb * 4), -+ result_type=fx.Vector.make_type(4, fx.Int32), -+ ) -+ ) -+ scode = fx.ptr_load(c["intsc"] + (pick * UG_GROUPS + st * 4 + kb)) -+ bw = fx.Vector.from_elements( -+ [picked.select(bw[d], fx.Int32(0)) for d in range(4)], fx.Int32 -+ ) -+ sb = picked.select(scode, fx.Int32(UNIT_SCALE)) -+ acc = mfma_scaled( -+ ops[j * W2_STEPS + st], -+ bw, -+ acc, -+ ops[nbs + j * W2_STEPS + st], -+ sb, -+ FP4, -+ FP4, -+ 0, -+ 0, -+ ) -+ # lane: rows 4 (lane / 16) + i of token lane % 16 -+ if live & picked: -+ wt = fx.ptr_load(rt["route"] + (t * 2 * TOPK + TOPK + k)).bitcast( -+ fx.Float32 -+ ) -+ for i in range_constexpr(4): -+ row = 4 * (lane // 16) + i -+ fx.ptr_store( -+ bf16_round(acc[i] * wt), -+ c["contrib"] + ((t * TOPK + k) * ROWS + row), -+ ) -+ b0 = nb -+ ops = ops_n -+ gpu.barrier() -+ if tid < (ROWS // 2) * s: -+ rp = tid % (ROWS // 2) -+ tt = tid // (ROWS // 2) -+ y0 = fx.Float32(0.0) -+ y1 = fx.Float32(0.0) -+ for k in range_constexpr(TOPK): -+ y0 = bf16_round( -+ y0 + fx.ptr_load(c["contrib"] + ((tt * TOPK + k) * ROWS + 2 * rp)) -+ ) -+ y1 = bf16_round( -+ y1 + fx.ptr_load(c["contrib"] + ((tt * TOPK + k) * ROWS + 2 * rp + 1)) -+ ) -+ for p in range_constexpr(TP): -+ c["put_bf"]( -+ c["peer"]["latr"][p], -+ (c["rank"] * s + tt) * LAT + rgd * ROWS + 2 * rp, -+ [y0, y1], -+ ) -+ gpu.barrier() -+ -+ -+# ---------------------------------------------------------------- shared down -+@traced -+def stage_h(c, t): -+ """Token t's shared h = SiTU(g, u) (vLLM's situ_and_mul, bf16) of SGU, once -+ -> H; a thread a column pair, its polls in one batch.""" -+ tid = c["tid"] -+ n = SI // 2 -+ if tid < n: -+ got = c["poll"]( -+ [ -+ (c["sgu"], (t * 2 * SI) // 2 + tid, 1), -+ (c["sgu"], (t * 2 * SI + SI) // 2 + tid, 1), -+ ] -+ ) -+ g0, g1 = bf_lo(got[0][0]), bf_hi(got[0][0]) -+ u0, u1 = bf_lo(got[1][0]), bf_hi(got[1][0]) -+ c["put_bf"](c["h"], t * SI + 2 * tid, [situ_vllm(g0, u0), situ_vllm(g1, u1)]) -+ gpu.barrier() -+ -+ -+@traced -+def load_h(c): -+ """H -> LDS rows (for the shared down GEMV), once a CTA.""" -+ s, tid = c["S"], c["tid"] -+ hrow = SI + gemv.LDS_PAD -+ n = s * SI // 2 -+ per = (n + THREADS - 1) // THREADS -+ got = c["poll"]([(c["h"], fx.min(tid + THREADS * i, n - 1), 1) for i in range(per)]) -+ for i in range_constexpr(per): -+ pidx = tid + THREADS * i -+ t = fx.min(pidx, n - 1) // (SI // 2) -+ cp = fx.min(pidx, n - 1) % (SI // 2) -+ if pidx < n: -+ fx.ptr_store(got[i][0], c["xl"] + (t * hrow + 2 * cp) // 2) -+ gpu.barrier() -+ -+ -+@traced -+def stage_sdown(c, j, ops): -+ """Shared down rows 16 j .. (K 768) against h in LDS (``load_h``): bf16 -+ partial -> SHP (this rank's up columns) or every rank's FINAL.""" -+ s, tid, lane, wave, red = c["S"], c["tid"], c["lane"], c["wave"], c["red"] -+ hrow = SI + gemv.LDS_PAD -+ acc = gemv.mfmas(lane, wave, c["xl"], hrow, SI, s, ops) -+ fx.ptr_store(acc, red + (wave * 64 + lane) * 4) -+ gpu.barrier() -+ lo = c["rank"] * UP_TASKS -+ mine = (j >= lo) & (j < lo + UP_TASKS) -+ if tid < (ROWS // 2) * s: -+ rp = tid % (ROWS // 2) -+ tt = tid // (ROWS // 2) -+ v0 = bf16_round(row_sum(red, 2 * rp, tt)) -+ v1 = bf16_round(row_sum(red, 2 * rp + 1, tt)) -+ if mine: -+ w = ( -+ fx.Vector.from_elements([v0, v1], fx.Float32) -+ .to(fx.BFloat16) -+ .bitcast(fx.Int32)[0] -+ ) -+ bo.buffer_store( -+ w, -+ rsrc(c["shp"]), -+ (tt * UP_N + (j - lo) * ROWS + 2 * rp) // 2, -+ cache_modifier=CM_DEV, -+ ) -+ else: -+ for p in range_constexpr(TP): -+ c["put_bf"]( -+ c["peer"]["final"][p], -+ (c["rank"] * s + tt) * HIDDEN + j * ROWS + 2 * rp, -+ [v0, v1], -+ ) -+ publish( -+ c["put"], -+ c["shrdy"], -+ fx.max(fx.min(j - lo, UP_TASKS - 1), 0), -+ fx.Int32(1), -+ (tid == 0) & mine, -+ ) -+ gpu.barrier() -+ -+ -+@traced -+def run_sdown(c, bid, w_sd): -+ """This rank's 56 shared-down tasks whose columns the up stage adds to (k -> -+ row group (k + 56 rank) % 448), on the CTAs past the down tasks: up waits -+ on them.""" -+ load_h(c) -+ lo = c["rank"] * UP_TASKS -+ if bid >= DOWN_TASKS: -+ for k in range(bid - DOWN_TASKS, UP_TASKS, BLOCKS - DOWN_TASKS): -+ j = (k + lo) % SD_TASKS -+ stage_sdown(c, j, gemv.loads(c["lane"], c["wave"], w_sd, j * ROWS, SI)) -+ -+ -+@traced -+def run_sdown_rest(c, bid, w_sd): -+ """The other 392 shared-down tasks (only the final reduce waits on them), on -+ the CTAs past lnorm / up's (84 ..), after their down: off the latent-norm -> -+ up critical path.""" -+ r0 = LN0 + LN_TASKS # 84: past the lnorm / up CTAs -+ if bid >= r0: -+ load_h(c) -+ lo = c["rank"] * UP_TASKS -+ for k in range(UP_TASKS + bid - r0, SD_TASKS, BLOCKS - r0): -+ j = (k + lo) % SD_TASKS -+ stage_sdown(c, j, gemv.loads(c["lane"], c["wave"], w_sd, j * ROWS, SI)) -+ -+ -+# ---------------------------------------------------------------- lnorm / up / reduce -+@traced -+def stage_lnorm(c, j): -+ """Columns 128 j ..: the ranks' LATR (rank order, fp32 -> bf16), every -+ task's sum of squares -> rstd -> bf16(acc rstd w) -> LN, LNRDY.""" -+ s, tid, lane, red = c["S"], c["tid"], c["lane"], c["red"] -+ a = c["args"] -+ mine = tid < s * 64 -+ t = fx.min(tid // 64, s - 1) -+ cp = tid % 64 -+ col = j * 128 + 2 * cp -+ lo, hi = sum_partials( -+ c["poll"], -+ c["peer"]["latr_own"], -+ lambda src: ((src * s + t) * LAT + col) // 2, -+ TP, -+ ) -+ lo, hi = bf16_round(lo), bf16_round(hi) -+ sq = wave_sum(lo * lo + hi * hi) -+ if mine & (lane == 0): -+ c["put"](c["lmsq"], j * s + t, sq) -+ gpu.barrier() -+ if mine & (lane < LN_TASKS): -+ v = c["poll"]([(c["lmsq"], lane * s + t, 1)])[0][0] -+ fx.ptr_store(f32_of(v), red + (t * 32 + lane)) -+ gpu.barrier() -+ tot = fx.ptr_load(red + t * 32) -+ for q in range_constexpr(1, LN_TASKS): -+ tot = tot + fx.ptr_load(red + (t * 32 + q)) -+ r = hw_rsq(tot * fx.Float32(1.0 / LAT) + fx.Float32(c["ln_eps"])) -+ w0 = fx.Float32( -+ fx.BFloat16(bo.buffer_load(rsrc(a["ln_w"]), col, vec_width=1, dtype=T.bf16)) -+ ) -+ w1 = fx.Float32( -+ fx.BFloat16(bo.buffer_load(rsrc(a["ln_w"]), col + 1, vec_width=1, dtype=T.bf16)) -+ ) -+ y = ( -+ fx.Vector.from_elements([lo * r * w0, hi * r * w1], fx.Float32) -+ .to(fx.BFloat16) -+ .bitcast(fx.Int32)[0] -+ ) -+ if mine: -+ bo.buffer_store(y, rsrc(c["ln"]), (t * LAT + col) // 2, cache_modifier=CM_DEV) -+ publish(c["put"], c["lnrdy"], j, fx.Int32(1), tid == 0) -+ gpu.barrier() -+ -+ -+@traced -+def stage_up(c, j): -+ """This rank's up rows 16 j .. of its shard (hidden rows 896 rank ..): LN GEMV -+ + SHP -> bf16 -> FINAL.""" -+ s, tid, lane, wave, red = c["S"], c["tid"], c["lane"], c["wave"], c["red"] -+ a = c["args"] -+ row0 = c["rank"] * UP_N + j * ROWS # the hidden column; w_up is the shard -+ ops = gemv.loads(lane, wave, a["w_up"], j * ROWS, LAT) -+ if tid < LN_TASKS: -+ c["poll"]([(c["lnrdy"], tid, 1)]) -+ if tid == 0: -+ c["poll"]([(c["shrdy"], j, 1)]) -+ gpu.barrier() -+ lrow = LAT + gemv.LDS_PAD -+ gemv.dev_rows_to_lds(tid, c["ln"], LAT, s, c["xl"], lrow) -+ acc = gemv.mfmas(lane, wave, c["xl"], lrow, LAT, s, ops) -+ fx.ptr_store(acc, red + (wave * 64 + lane) * 4) -+ gpu.barrier() -+ if tid < (ROWS // 2) * s: -+ rp = tid % (ROWS // 2) -+ tt = tid // (ROWS // 2) -+ w = fx.Int32( -+ bo.buffer_load( -+ rsrc(c["shp"]), -+ (tt * UP_N + j * ROWS + 2 * rp) // 2, -+ vec_width=1, -+ dtype=T.i32, -+ cache_modifier=CM_DEV, -+ ) -+ ) -+ v0 = bf16_round(bf_lo(w) + row_sum(red, 2 * rp, tt)) -+ v1 = bf16_round(bf_hi(w) + row_sum(red, 2 * rp + 1, tt)) -+ for p in range_constexpr(TP): -+ c["put_bf"]( -+ c["peer"]["final"][p], -+ (c["rank"] * s + tt) * HIDDEN + row0 + 2 * rp, -+ [v0, v1], -+ ) -+ gpu.barrier() -+ -+ -+@traced -+def stage_reduce(c, j): -+ """Columns 128 j ..: FINAL in rank order -> bf16 out.""" -+ s, tid = c["S"], c["tid"] -+ if tid < s * 64: -+ t = tid // 64 -+ cp = tid % 64 -+ col = j * 128 + 2 * cp -+ lo, hi = sum_partials( -+ c["poll"], -+ c["peer"]["final_own"], -+ lambda src: ((src * s + t) * HIDDEN + col) // 2, -+ TP, -+ ) -+ w = ( -+ fx.Vector.from_elements([lo, hi], fx.Float32) -+ .to(fx.BFloat16) -+ .bitcast(fx.Int32)[0] -+ ) -+ bo.buffer_store(w, rsrc(c["args"]["out"]), (t * HIDDEN + col) // 2) -+ gpu.barrier() -+ -+ -+# ---------------------------------------------------------------- kernel -diff --git a/vllm/models/kimi_k3/amd/specialized_decode/runtime/__init__.py b/vllm/models/kimi_k3/amd/specialized_decode/runtime/__init__.py -new file mode 100644 -index 000000000..e17a0e3ff ---- /dev/null -+++ b/vllm/models/kimi_k3/amd/specialized_decode/runtime/__init__.py -@@ -0,0 +1,219 @@ -+# SPDX-License-Identifier: Apache-2.0 -+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -+"""Kimi-K3 specialized decode on MI355X (layer mode), enabled by -+``additional_config={"decode_backend": "latency_specialized"}``. -+ -+``specialized_decode_config`` runs in the model's ``__init__``, before the weights -+load. It returns None for the stock decode, and raises with every reason from -+``unsupported`` when the specialized decode is requested for a deployment it cannot -+serve. Checks that need the loaded weights run when the executor is built -+(``decode.SpecializedDecode``), those that need the KV cache in ``prepare_kv``.""" -+ -+from dataclasses import dataclass -+ -+import torch -+ -+from vllm.logger import init_logger -+ -+logger = init_logger(__name__) -+ -+AUTO, SPECIALIZED = "auto", "latency_specialized" -+TP = 8 -+# a launch's x rows fit one CTA's LDS -+MAX_TOKENS = 8 -+WHOLE, FFN = "whole", "ffn" -+# Kimi-K3's text shapes the kernels are written for (hidden size, routed experts, -+# KDA heads, KDA head dim, KDA short-conv width) -+K3_SHAPES = (7168, 896, 96, 128, 4) -+ -+ -+@dataclass(frozen=True) -+class DecodeConfig: -+ # rows a request verifies: num_speculative_tokens + 1 -+ q: int -+ -+ @property -+ def widths(self) -> tuple[int, ...]: -+ """The rows a launch serves: whole requests, at most MAX_TOKENS rows. Every -+ uniform verify step up to MAX_TOKENS rows is one of them, so an eager step -+ never pads; a FULL replay pads whole requests the way vLLM's stock decode -+ does.""" -+ return tuple(range(self.q, MAX_TOKENS + 1, self.q)) -+ -+ -+def decode_backend(vllm_config) -> str: -+ """The requested decode backend: the one place that reads the switch.""" -+ extra = vllm_config.additional_config -+ return extra.get("decode_backend", AUTO) if isinstance(extra, dict) else AUTO -+ -+ -+def specialized_decode_config(vllm_config) -> DecodeConfig | None: -+ """The deployment's specialized decode settings, None for the stock decode. -+ Raises when the specialized decode is requested for an unsupported deployment, -+ listing every reason.""" -+ backend = decode_backend(vllm_config) -+ if backend == AUTO: -+ return None -+ if backend != SPECIALIZED: -+ raise ValueError( -+ f"decode_backend must be '{AUTO}' or '{SPECIALIZED}', got {backend!r}" -+ ) -+ reasons = unsupported(vllm_config) -+ if reasons: -+ raise ValueError( -+ "Kimi-K3 specialized decode (decode_backend=latency_specialized) does " -+ "not support: " + "; ".join(reasons) + f" (decode_backend={AUTO} runs " -+ "vLLM's decode)" -+ ) -+ if piecewise_graphs_only(vllm_config): -+ logger.warning( -+ "Kimi-K3 specialized decode: PIECEWISE-only CUDA graphs capture every " -+ "decode step as stock; use FULL decode graphs" -+ ) -+ spec = vllm_config.speculative_config -+ return DecodeConfig(q=spec.num_speculative_tokens + 1) -+ -+ -+def unsupported(vllm_config) -> list[str]: -+ """Every reason this deployment cannot run the specialized decode (empty: it -+ can). Config only: the loaded weights are checked when the executor is built. -+ Default-deny: every feature below is refused, the kernels serve none of them.""" -+ from vllm.config import CompilationMode -+ from vllm.model_executor.layers.mamba.mamba_utils import ( -+ MambaStateDtypeCalculator, -+ ) -+ -+ mc, pc, cc = ( -+ vllm_config.model_config, -+ vllm_config.parallel_config, -+ vllm_config.cache_config, -+ ) -+ spec, oc = vllm_config.speculative_config, vllm_config.offload_config -+ checks = ( -+ (pc.tensor_parallel_size != TP, f"TP {pc.tensor_parallel_size} (TP{TP} only)"), -+ (pc.pipeline_parallel_size != 1, "pipeline parallelism"), -+ (pc.data_parallel_size != 1, "data parallelism"), -+ (pc.enable_expert_parallel, "expert parallelism"), -+ (pc.enable_eplb, "EPLB"), -+ ( -+ pc.decode_context_parallel_size not in (1, pc.tensor_parallel_size), -+ "decode context parallelism other than 1 or TP", -+ ), -+ (pc.prefill_context_parallel_size != 1, "prefill context parallelism"), -+ # the kernels read one batch's attention metadata per step -+ (pc.use_ubatching, "DBO (dual-batch overlap)"), -+ (vllm_config.lora_config is not None, "LoRA"), -+ (vllm_config.kv_transfer_config is not None, "KV transfer connectors"), -+ ( -+ vllm_config.aux_output_config.enable_return_routed_experts, -+ "routed-expert capture", -+ ), -+ # the kernels read the loaded weights in place -+ (mc.enable_sleep_mode, "sleep mode"), -+ ( -+ oc.uva.cpu_offload_gb > 0 or oc.prefetch.offload_group_size > 0, -+ "CPU weight offload", -+ ), -+ (vllm_config.weight_transfer_config is not None, "weight updates"), -+ # the kernels serve spec-verify steps of whole uniform requests only -+ (spec is None, "decoding without speculative decoding"), -+ ( -+ spec is not None and spec.num_speculative_tokens + 1 > MAX_TOKENS, -+ f"more than {MAX_TOKENS - 1} speculative tokens", -+ ), -+ ( -+ spec is not None and spec.enable_adaptive_verification, -+ "adaptive verification", -+ ), -+ ( -+ spec is not None and spec.uses_dynamic_speculative_decoding(), -+ "dynamic speculative decoding", -+ ), -+ # K1 keeps a BF16 conv state and an FP32 recurrent state in place; it -+ # maintains no ReplaySSM / RecoverSSM records -+ ( -+ MambaStateDtypeCalculator.kda_state_dtype(mc.dtype, cc.mamba_cache_dtype) -+ != (torch.bfloat16, torch.float32), -+ "KDA cache dtypes other than a BF16 conv state and an FP32 recurrent " -+ "state (mamba cache dtype)", -+ ), -+ (cc.use_replayssm or cc.use_kda_recoverssm, "ReplaySSM / RecoverSSM"), -+ # the model state hooks build the executor and bind it to each KV cache -+ ( -+ not vllm_config.use_v2_model_runner, -+ "the V1 model runner (set VLLM_USE_V2_MODEL_RUNNER=1)", -+ ), -+ # plan() and the layer kernels run in forward; a compiled forward would -+ # trace them -+ ( -+ vllm_config.compilation_config.mode == CompilationMode.STOCK_TORCH_COMPILE, -+ "torch.compile (compilation mode STOCK_TORCH_COMPILE)", -+ ), -+ ) -+ reasons = [why for bad, why in checks if bad] -+ shapes = text_shapes(mc.hf_text_config) -+ if shapes != K3_SHAPES: -+ reasons.append( -+ f"model shapes {shapes} (hidden, experts, KDA heads, KDA head dim, conv " -+ f"width; Kimi-K3 is {K3_SHAPES})" -+ ) -+ if getattr(mc.hf_text_config, "attn_res_block_size", None) is None: -+ reasons.append("models without AttnRes") -+ return reasons + platform_reasons() -+ -+ -+def text_shapes(text) -> tuple: -+ kda = getattr(text, "linear_attn_config", None) or {} -+ return ( -+ getattr(text, "hidden_size", None), -+ getattr(text, "num_experts", None), -+ kda.get("num_heads"), -+ kda.get("head_dim"), -+ kda.get("short_conv_kernel_size"), -+ ) -+ -+ -+def platform_reasons() -> list[str]: -+ from vllm.platforms import current_platform -+ -+ if not current_platform.is_rocm(): -+ return ["this platform (ROCm only)"] -+ from vllm.platforms.rocm import on_gfx950 -+ -+ if not on_gfx950(): -+ return ["this GPU (gfx950 / MI355X only)"] -+ from vllm.models.kimi_k3.amd.specialized_decode.mono.common.plan import BLOCKS -+ -+ cus = torch.cuda.get_device_properties().multi_processor_count -+ if cus != BLOCKS: -+ # every CTA of a launch must be resident: a partitioned GPU deadlocks -+ return [f"{cus} CUs (a partitioned GPU; the kernels need {BLOCKS})"] -+ return [] -+ -+ -+def piecewise_graphs_only(vllm_config) -> bool: -+ """Decode steps captured as PIECEWISE graphs only: plan() runs them stock.""" -+ from vllm.config import CUDAGraphMode -+ -+ mode = vllm_config.compilation_config.cudagraph_mode -+ if mode is None or mode == CUDAGraphMode.NONE: -+ return False -+ return mode.decode_mode() != CUDAGraphMode.FULL -+ -+ -+def layer_kinds(config) -> dict[int, str]: -+ """Decoder layer -> kind, for every layer a kernel serves: the MoE layers with -+ AttnRes and at least one block before them, minus the MLA block-write layers -+ (none in Kimi-K3).""" -+ block = config.attn_res_block_size -+ if block is None or config.num_experts is None: -+ return {} -+ out = {} -+ for i in range(1, config.num_hidden_layers): -+ if i < config.first_k_dense_replace or i % config.moe_layer_freq: -+ continue -+ if config.is_kda_layer(i): -+ out[i] = WHOLE -+ elif i % block: -+ out[i] = FFN -+ return out -diff --git a/vllm/models/kimi_k3/amd/specialized_decode/runtime/decode.py b/vllm/models/kimi_k3/amd/specialized_decode/runtime/decode.py -new file mode 100644 -index 000000000..fa62155fd ---- /dev/null -+++ b/vllm/models/kimi_k3/amd/specialized_decode/runtime/decode.py -@@ -0,0 +1,198 @@ -+# SPDX-License-Identifier: Apache-2.0 -+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -+"""The Kimi-K3 specialized decode executor (layer mode): one routing decision per -+step, then the layer kernels inside the stock layer loop. -+ -+Built by the model state once the weights are loaded (``KimiK3ModelState``): checks -+the loaded weights, builds the per-rank launch state (``runner``) and the layer -+kernels of every width, and compiles every build, before memory profiling and graph -+capture. ``prepare_kv`` binds it to every KV cache vLLM allocates. Each step -+``plan`` returns a ``Step`` or None; ``Step.layer(i)`` is the kernel replacing -+decoder layer i (None for vLLM's layers). Weight reloads are not supported.""" -+ -+from collections.abc import Callable, Mapping -+from functools import partial -+from typing import Any, NamedTuple -+ -+import torch -+ -+from vllm.config import CUDAGraphMode -+from vllm.distributed import get_tp_group -+from vllm.forward_context import get_forward_context, is_forward_context_available -+from vllm.logger import init_logger -+from vllm.models.kimi_k3.amd.specialized_decode.mono.abi import FAILED -+from vllm.models.kimi_k3.amd.specialized_decode.runtime import ( -+ WHOLE, -+ DecodeConfig, -+ layer_kinds, -+) -+from vllm.models.kimi_k3.amd.specialized_decode.runtime.guards import ( -+ FailStop, -+ cache_reasons, -+ weight_reasons, -+) -+from vllm.models.kimi_k3.amd.specialized_decode.runtime.layers import ( -+ FfnLayer, -+ WholeLayer, -+ compile_build, -+) -+from vllm.models.kimi_k3.amd.specialized_decode.runtime.peer import tp_agree -+from vllm.models.kimi_k3.amd.specialized_decode.runtime.runner import MonoRunner -+ -+logger = init_logger(__name__) -+ -+Kernel = WholeLayer | FfnLayer -+ -+ -+class Step(NamedTuple): -+ """A planned step: ``width`` rows of whole uniform spec-verify requests.""" -+ -+ width: int -+ attn_metadata: Mapping[str, Any] # by layer name -+ kernels: Mapping[int, Kernel] # by decoder layer, at this width -+ -+ def layer(self, idx: int) -> Callable | None: -+ """The kernel replacing decoder layer idx in this step, or None (vLLM's).""" -+ kernel = self.kernels.get(idx) -+ return None if kernel is None else partial(kernel.launch, self) -+ -+ -+def step_reason(attn_metadata, kda_prefix: str, width: int) -> str | None: -+ """None if a step's KDA metadata (identical on every TP rank) is a pure -+ spec-verify step of whole uniform requests filling ``width`` rows, else why -+ not.""" -+ if not isinstance(attn_metadata, dict): -+ return "no metadata" -+ m: Any = attn_metadata.get(kda_prefix) -+ if m is None or m.spec_sequence_masks is None: -+ return "not spec-verify" -+ if m.num_prefills or m.num_decodes or m.num_spec_decodes < 1: -+ return "not pure decode" -+ q = m.spec_state_indices_tensor.size(-1) -+ if m.uniform_spec_sequence_length != q or m.num_spec_decodes * q != width: -+ return "not uniform" -+ return None -+ -+ -+def agree(reasons: list[str], group, what: str) -> None: -+ """Raises on every TP rank if any rank has ``reasons``: a rank that refused -+ alone would leave the others spinning in-kernel.""" -+ if not tp_agree(not reasons, group): -+ why = "; ".join(reasons) if reasons else "another TP rank refused" -+ raise ValueError( -+ f"Kimi-K3 specialized decode does not support {what}: {why} " -+ "(decode_backend=auto runs vLLM's decode)" -+ ) -+ -+ -+class SpecializedDecode: -+ """The executor: builds the layer kernels, binds them to each KV cache and -+ plans every step.""" -+ -+ def __init__(self, backbone, config: DecodeConfig): -+ tp = get_tp_group() -+ agree(weight_reasons(backbone), tp.cpu_group, "the loaded weights") -+ self.backbone, self.config = backbone, config -+ kinds = { -+ i: kind -+ for i, kind in layer_kinds(backbone.config).items() -+ if backbone.start_layer <= i < backbone.end_layer -+ } -+ runner = self.runner = MonoRunner(next(backbone.parameters()).device) -+ runner.num_spec = config.q - 1 -+ runner.first_layer = min(kinds) -+ whole = [i for i, kind in kinds.items() if kind == WHOLE] -+ runner.kda_prefix = backbone.layers[whole[0]].self_attn.prefix -+ self.kernels: dict[int, dict[int, Kernel]] = { -+ w: { -+ i: (WholeLayer if kind == WHOLE else FfnLayer)( -+ backbone.layers[i], runner, w -+ ) -+ for i, kind in kinds.items() -+ } -+ for w in config.widths -+ } -+ self.kda: list[WholeLayer] = [ -+ k for k in self.kernels[config.widths[0]].values() -+ if isinstance(k, WholeLayer) -+ ] # fmt: skip -+ builds = { -+ b for by in self.kernels.values() for k in by.values() for b in k.builds() -+ } -+ for build in builds: -+ compile_build(build) -+ self.failstop = FailStop( -+ runner.health[FAILED : FAILED + 1], runner.rank, tp.world_size, tp.cpu_group -+ ) -+ # (data_ptr, numel) of the first KDA layer's state prepare_kv bound -+ self._kv: tuple[int, int] | None = None -+ logger.info( -+ "Kimi-K3 specialized decode: %d layers, widths %s, %d builds; other " -+ "steps run vLLM's decode", -+ len(kinds), -+ list(config.widths), -+ len(builds), -+ ) -+ -+ def _kv_key(self) -> tuple[int, int] | None: -+ state = self.kda[0].layer.self_attn.kv_cache -+ if not isinstance(state, (list, tuple)) or len(state) != 2: -+ return None -+ if state[1].numel() == 0: -+ return None -+ return state[1].data_ptr(), state[1].numel() -+ -+ def prepare_kv(self) -> None: -+ """Checks and binds the KDA states vLLM just bound (the CUDA-graph memory -+ profiling cache, then the real one).""" -+ self._kv = None -+ if self._kv_key() is None: -+ return -+ caches = {k.layer.self_attn.prefix: k.caches() for k in self.kda} -+ agree(cache_reasons(caches), get_tp_group().cpu_group, "the KV caches") -+ self._kv = self._kv_key() -+ -+ def plan(self, positions: torch.Tensor, intermediate_tensors) -> Step | None: -+ """This step's plan, or None to run vLLM's layers. -+ -+ Steps that run vLLM's layers: -+ - pipeline-parallel stage inputs, and steps before a KV cache is bound -+ (profiling, dummy runs); -+ - PIECEWISE steps (captured without attention metadata); -+ - widths without kernels (more than MAX_TOKENS rows); -+ - steps that are not pure spec-verify steps of whole uniform requests -+ (``step_reason``). -+ -+ Eager steps run the kernels too. Every predicate is the same on all TP -+ ranks: the kernels' mailboxes need every rank on one path. Under FULL -+ graphs this runs at capture only and the plan is baked in; a replay with -+ fewer requests than its graph pads whole requests (NULL_BLOCK_ID state, -+ one accepted token), which the kernels serve like real ones. -+ """ -+ if intermediate_tensors is not None or self._kv is None: -+ return None -+ if not is_forward_context_available(): -+ return None -+ context = get_forward_context() -+ if context.cudagraph_runtime_mode == CUDAGraphMode.PIECEWISE: -+ return None -+ width = positions.size(-1) -+ kernels = self.kernels.get(width) -+ if kernels is None: -+ return None -+ md = context.attn_metadata -+ if step_reason(md, self.runner.kda_prefix, width) is not None: -+ return None -+ assert isinstance(md, dict) -+ if self._kv_key() != self._kv: -+ raise RuntimeError( -+ "Kimi-K3 specialized decode: the KV cache was re-bound without " -+ "prepare_kv (ModelState.initialize_kv_cache)" -+ ) -+ return Step(width, md, kernels) -+ -+ def after_step(self) -> None: -+ """Once per executed step, from compute_logits (eager on every rank, also -+ under graph replay). Raises on every rank once a rank gave up an in-kernel -+ wait.""" -+ self.failstop.after_step() -diff --git a/vllm/models/kimi_k3/amd/specialized_decode/runtime/guards.py b/vllm/models/kimi_k3/amd/specialized_decode/runtime/guards.py -new file mode 100644 -index 000000000..9d0808d72 ---- /dev/null -+++ b/vllm/models/kimi_k3/amd/specialized_decode/runtime/guards.py -@@ -0,0 +1,154 @@ -+# SPDX-License-Identifier: Apache-2.0 -+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -+"""The refusals that need the loaded model (``weight_reasons``) or the bound KV -+caches (``cache_reasons``), and the fail-stop watch on the kernels' error words. -+Each refusal reports every reason, not only the first.""" -+ -+from collections.abc import Mapping -+from typing import Any -+ -+import torch -+ -+from vllm.models.kimi_k3.amd.specialized_decode.runtime import WHOLE, layer_kinds -+ -+ -+def weight_reasons(model) -> list[str]: -+ """``model``: the loaded ``KimiLinearModel``; every kernel layer's weights as -+ the kernels read them, and the MoE backend that lays them out.""" -+ from vllm.model_executor.layers.fused_moe.oracle.mxfp4 import Mxfp4MoeBackend -+ from vllm.models.kimi_k3.amd.specialized_decode.mono.attention import back, kda -+ from vllm.models.kimi_k3.amd.specialized_decode.mono.stages import moe -+ -+ reasons: list[str] = [] -+ bf, f32 = torch.bfloat16, torch.float32 -+ -+ def need(name, t, shape, dtype): -+ if t is None: -+ reasons.append(f"{name}: missing") -+ elif tuple(t.shape) != shape or t.dtype != dtype or not t.is_contiguous(): -+ reasons.append( -+ f"{name}: {tuple(t.shape)} {t.dtype}" -+ f"{'' if t.is_contiguous() else ' strided'}," -+ f" kernel reads contiguous {shape} {dtype}" -+ ) -+ -+ kinds = layer_kinds(model.config) -+ if not kinds: -+ return ["no layer has an AttnRes latent MoE"] -+ for i, kind in kinds.items(): -+ if not model.start_layer <= i < model.end_layer: -+ continue -+ layer = model.layers[i] -+ p = f"layer {i}" -+ m = layer.mlp -+ if not getattr(m, "use_latent_moe", False) or m.shared_experts is None: -+ reasons.append(f"{p}: needs a latent MoE with shared experts") -+ continue -+ if m.routed_scaling_factor != 1.0: -+ reasons.append(f"{p}: routed_scaling_factor {m.routed_scaling_factor}") -+ ex = getattr(m.experts, "routed_experts", m.experts) -+ qm = getattr(ex, "quant_method", None) -+ if getattr( -+ qm, "mxfp4_backend", None -+ ) != Mxfp4MoeBackend.AITER_MXFP4_BF16 or not getattr( -+ qm, "is_k3_situ_aiter", False -+ ): -+ reasons.append(f"{p}: MoE backend must be AITER_MXFP4_BF16 (K3 SiTU)") -+ continue -+ H, LAT, E = moe.HIDDEN, moe.LAT, moe.E -+ need(f"{p} gate", m.gate.weight, (E, H), bf) -+ need(f"{p} gate bias", m.gate.e_score_correction_bias, (E,), f32) -+ need(f"{p} latent down", m.routed_expert_down_proj.weight, (LAT, H), bf) -+ need(f"{p} latent norm", m.routed_expert_norm.weight, (LAT,), bf) -+ need(f"{p} shared gate_up", m.shared_experts.gate_up_proj.weight, -+ (2 * moe.SI, H), bf) # fmt: skip -+ need(f"{p} shared down", m.shared_experts.down_proj.weight, (H, moe.SI), bf) -+ sharded = getattr(m.routed_output_transform, "row_sharded", False) -+ need(f"{p} up_proj", m.routed_expert_up_proj.weight, -+ (moe.UP_N if sharded else H, LAT), bf) # fmt: skip -+ for name in ("w13_weight", "w13_weight_scale", "w2_weight", "w2_weight_scale"): -+ if getattr(ex, name, None) is None: -+ reasons.append(f"{p}: experts have no {name}") -+ need(f"{p} mlp_res_norm", layer.mlp_res_norm.weight, (H,), bf) -+ need(f"{p} mlp_res_proj", layer.mlp_res_proj.weight, (1, H), bf) -+ need(f"{p} post_attention_layernorm", layer.post_attention_layernorm.weight, -+ (H,), bf) # fmt: skip -+ need(f"{p} o_proj", layer.self_attn.o_proj.weight, (H, back.OK), bf) -+ if kind != WHOLE: -+ continue -+ a = layer.self_attn -+ if a.local_num_heads != kda.NH or a.head_dim != kda.HD or a.conv_size != 4: -+ reasons.append(f"{p}: KDA heads/head_dim/conv_size") -+ continue -+ need(f"{p} in_proj", a.in_proj_qkvgfab.weight, (kda.NPROJ, H), bf) -+ need(f"{p} f_b_proj", a.f_b_proj.weight, (kda.PROJ, kda.HD), bf) -+ need(f"{p} o_norm", a.o_norm.weight, (kda.HD,), bf) -+ need(f"{p} conv1d", a.conv1d.weight.view(-1, 4), (kda.QKV, 4), f32) -+ need(f"{p} A_log", a.A_log.view(-1), (kda.NH,), f32) -+ need(f"{p} dt_bias", a.dt_bias.view(-1), (kda.PROJ,), f32) -+ need(f"{p} input_layernorm", layer.input_layernorm.weight, (H,), bf) -+ need(f"{p} self_attention_res_norm", layer.self_attention_res_norm.weight, -+ (H,), bf) # fmt: skip -+ need(f"{p} self_attention_res_proj", layer.self_attention_res_proj.weight, -+ (1, H), bf) # fmt: skip -+ return reasons -+ -+ -+def cache_reasons(caches: Mapping[str, Any]) -> list[str]: -+ """Every KDA layer's (conv, recurrent) state as K1 reads them.""" -+ from vllm.models.kimi_k3.amd.specialized_decode.mono.attention import kda -+ -+ reasons = [] -+ for name, (conv_state, rstate) in caches.items(): -+ if conv_state.dtype != torch.bfloat16: -+ reasons.append(f"{name}: conv state {conv_state.dtype}, needs bf16") -+ if rstate.dtype != torch.float32: -+ reasons.append(f"{name}: recurrent state {rstate.dtype}, needs fp32") -+ elif ( -+ rstate[0].numel() != kda.NH * kda.HD * kda.HD -+ or not rstate[0].is_contiguous() -+ ): -+ reasons.append(f"{name}: a slot's recurrent state must be dense") -+ return reasons -+ -+ -+class FailStop: -+ """Stops every TP rank at the same step once any of them gave up an in-kernel -+ wait (its health word FAILED, ``mono.abi``). Each step queues a D2H copy of -+ this rank's word and keeps the newest copy that has landed, without blocking; -+ every CHECK_EVERY steps the ranks agree on those over the TP CPU group and all -+ raise together. A rank that raised alone would leave the others in the next -+ collective (the drafter's) for ever. A failure is seen within about -+ CHECK_EVERY steps.""" -+ -+ CHECK_EVERY = 16 -+ -+ def __init__(self, failed: torch.Tensor, rank: int, world: int, group): -+ self.failed, self.rank, self.world, self.group = failed, rank, world, group -+ self.host = torch.zeros(1, dtype=torch.int32, pin_memory=True) -+ self.copied: torch.cuda.Event | None = None -+ self.seen = False -+ self.steps = 0 -+ -+ def after_step(self) -> None: -+ """Once per step on every TP rank (compute_logits, itself collective).""" -+ if torch.cuda.is_current_stream_capturing(): -+ return -+ if self.copied is None or self.copied.query(): -+ if self.copied is not None: -+ self.seen |= bool(self.host[0]) -+ self.host.copy_(self.failed, non_blocking=True) -+ self.copied = torch.cuda.Event() -+ self.copied.record() -+ self.steps += 1 -+ if self.steps % self.CHECK_EVERY: -+ return -+ votes = torch.zeros(self.world, dtype=torch.int32) -+ votes[self.rank] = int(self.seen) -+ torch.distributed.all_reduce(votes, group=self.group) -+ if votes.any(): -+ ranks = [r for r, v in enumerate(votes.tolist()) if v] -+ raise RuntimeError( -+ f"Kimi-K3 specialized decode: TP ranks {ranks} gave up an in-kernel " -+ "wait; stopping the engine (decode_backend=auto runs vLLM's decode)" -+ ) -diff --git a/vllm/models/kimi_k3/amd/specialized_decode/runtime/layers.py b/vllm/models/kimi_k3/amd/specialized_decode/runtime/layers.py -new file mode 100644 -index 000000000..843e66dc5 ---- /dev/null -+++ b/vllm/models/kimi_k3/amd/specialized_decode/runtime/layers.py -@@ -0,0 +1,221 @@ -+# SPDX-License-Identifier: Apache-2.0 -+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -+"""The layer kernels: each takes and returns what ``KimiDecoderLayer.forward`` -+does with AttnRes, ``(prefix_sum, block_residual, prefix_delta)``. -+ -+- ``WholeLayer`` (KDA layers): K1 (AttnRes, in_proj, f_b, spec conv, the KDA -+ recurrence, the gated norm) writes o_proj's input; K2 runs o_proj, its -+ all-reduce, the MLP AttnRes and the latent MoE. -+- ``FfnLayer`` (MLA layers): vLLM's AttnRes and MLA up to o_proj's input, then -+ K2. -+ -+Weights are views of the loaded parameters, taken at every launch. -+""" -+ -+from typing import TYPE_CHECKING, Any -+ -+import torch -+ -+from vllm.model_executor.layers.mamba.mamba_utils import is_conv_state_dim_first -+from vllm.models.kimi_k3.amd.specialized_decode.runtime.runner import MonoRunner -+ -+if TYPE_CHECKING: -+ from vllm.models.kimi_k3.amd.specialized_decode.runtime.decode import Step -+ -+ -+def _k2_key(layer, runner: MonoRunner, tokens: int, reset: bool): -+ from vllm.models.kimi_k3.amd.specialized_decode.mono import layer as k2 -+ -+ m = layer.mlp -+ return k2.K2Build( -+ tokens=tokens, -+ nblocks=layer.prev_valid_blocks + (1 if reset else 0), -+ reset=reset, -+ eps=layer.mlp_res_norm.variance_epsilon, -+ out_eps=layer.post_attention_layernorm.variance_epsilon, -+ ln_eps=m.routed_expert_norm.variance_epsilon, -+ **runner.fault, -+ ) -+ -+ -+def _k2(layer, runner: MonoRunner, core, prefix, blocks, reset=False): -+ """K2 from ``core`` (o_proj's input): ``prefix`` updated in place (written, -+ when ``reset``), the MoE output returned.""" -+ from vllm.models.kimi_k3.amd.specialized_decode.mono import layer as k2 -+ from vllm.models.kimi_k3.amd.specialized_decode.mono.attention.kda import HIDDEN -+ from vllm.models.kimi_k3.amd.specialized_decode.mono.stages import moe -+ -+ m = layer.mlp -+ ex = getattr(m.experts, "routed_experts", m.experts) -+ w_up = m.routed_expert_up_proj.weight -+ if not getattr(m.routed_output_transform, "row_sharded", False): -+ w_up = w_up.narrow(0, runner.rank * moe.UP_N, moe.UP_N) -+ out = torch.empty_like(prefix) -+ k2.k2_launch( -+ _k2_key(layer, runner, core.size(0), reset), -+ core=core, -+ w_o=layer.self_attn.o_proj.weight, -+ prefix=prefix, -+ blocks=blocks, -+ ares_nw=layer.mlp_res_norm.weight, -+ ares_qk=layer.mlp_res_proj.weight.view(HIDDEN), -+ in_nw=layer.post_attention_layernorm.weight, -+ w_gate=m.gate.weight, -+ bias=m.gate.e_score_correction_bias, -+ w_ld=m.routed_expert_down_proj.weight, -+ w_sgu=m.shared_experts.gate_up_proj.weight, -+ w_sd=m.shared_experts.down_proj.weight, -+ w13=ex.w13_weight, -+ w13s=ex.w13_weight_scale, -+ w2=ex.w2_weight, -+ w2s=ex.w2_weight_scale, -+ ln_w=m.routed_expert_norm.weight, -+ w_up=w_up, -+ out=out, -+ layer=128 + layer.layer_idx, -+ **runner.launch_args(), -+ **runner.peer_args(), -+ ) -+ return out -+ -+ -+def compile_build(key) -> None: -+ """Compiles one launch's kernel (a ``K2Build`` or ``KdaPreBuild``) without -+ running it: FlyDSL's compile-only mode returns before the launch.""" -+ import os -+ -+ from vllm.models.kimi_k3.amd.specialized_decode.mono import layer as k2 -+ from vllm.models.kimi_k3.amd.specialized_decode.mono.attention import kda -+ -+ mod = kda if isinstance(key, kda.KdaPreBuild) else k2 -+ launch = mod.build(key) -+ # FlyDSL's compile-only switch has no FLYDSL_ prefix -+ old = os.environ.get("COMPILE_ONLY") -+ os.environ["COMPILE_ONLY"] = "1" -+ try: -+ launch(*[0] * len(mod.ABI.names), stream=torch.cuda.current_stream()) -+ finally: -+ if old is None: -+ del os.environ["COMPILE_ONLY"] -+ else: -+ os.environ["COMPILE_ONLY"] = old -+ -+ -+class _Layer: -+ def __init__(self, layer, runner: MonoRunner, width: int): -+ self.layer = layer -+ self.runner = runner -+ self.width = width -+ # the step's first launch bumps the epoch (a graph captures it) -+ self.first = layer.layer_idx == runner.first_layer -+ -+ def builds(self) -> list: -+ """The builds this layer launches at its width, compiled before capture.""" -+ reset = getattr(self.layer, "is_block_write_layer", False) -+ return [_k2_key(self.layer, self.runner, self.width, reset)] -+ -+ -+class WholeLayer(_Layer): -+ def _k1_key(self, qlen: int, state_len: int, delta: bool): -+ from vllm.models.kimi_k3.amd.specialized_decode.mono.attention import kda -+ -+ layer = self.layer -+ attn = layer.self_attn -+ return kda.KdaPreBuild( -+ tokens=self.width, -+ qlen=qlen, -+ nblocks=layer.prev_valid_blocks, -+ delta=delta, -+ state_len=state_len, -+ eps=layer.self_attention_res_norm.variance_epsilon, -+ out_eps=layer.input_layernorm.variance_epsilon, -+ onorm_eps=attn.o_norm.eps, -+ lower_bound=attn.gate_lower_bound, -+ write_idx=layer.block_write_idx if layer.is_block_write_layer else -1, -+ ) -+ -+ def builds(self) -> list: -+ from vllm.models.kimi_k3.amd.specialized_decode.mono.attention import kda -+ -+ q = self.runner.num_spec + 1 -+ # the conv state keeps CONV_W - 1 + num_spec entries; a layer after -+ # another always has a pending delta -+ k1 = self._k1_key(q, kda.CONV_W - 1 + self.runner.num_spec, True) -+ return [k1, *super().builds()] -+ -+ def caches(self) -> tuple[torch.Tensor, torch.Tensor]: -+ """The conv state as [slots, dim, state_len] (a view) and the recurrent -+ state.""" -+ conv_state, rstate = self.layer.self_attn.kv_cache -+ if not is_conv_state_dim_first(): -+ conv_state = conv_state.transpose(-1, -2) -+ return conv_state, rstate -+ -+ def launch(self, step: "Step", positions, prefix_sum, block_residual, prefix_delta): -+ from vllm.models.kimi_k3.amd.specialized_decode.mono.attention import kda -+ -+ layer, runner = self.layer, self.runner -+ attn = layer.self_attn -+ m: Any = step.attn_metadata[attn.prefix] -+ if self.first: -+ runner.step_begin() -+ conv_state, rstate = self.caches() -+ s = prefix_sum.size(0) -+ core = prefix_sum.new_empty(s, kda.PROJ) -+ kda.kda_pre( -+ self._k1_key( -+ m.spec_state_indices_tensor.size(-1), -+ conv_state.size(-1), -+ prefix_delta is not None, -+ ), -+ prefix=prefix_sum, -+ delta=prefix_delta, -+ blocks=block_residual, -+ ares_nw=layer.self_attention_res_norm.weight, -+ ares_qk=layer.self_attention_res_proj.weight.view(kda.HIDDEN), -+ in_nw=layer.input_layernorm.weight, -+ w_in=attn.in_proj_qkvgfab.weight, -+ w_fb=attn.f_b_proj.weight, -+ conv_w=attn.conv1d.weight.view(kda.QKV, 4), -+ conv_state=conv_state, -+ a_log=attn.A_log, -+ dt_bias=attn.dt_bias, -+ on_w=attn.o_norm.weight, -+ rstate=rstate, -+ st_idx=m.spec_state_indices_tensor, -+ num_acc=m.num_accepted_tokens, -+ core_out=core, -+ layer=layer.layer_idx, -+ **runner.launch_args(), -+ ) -+ if layer.is_block_write_layer: -+ # the prefix restarts at the attention's output -+ prefix_sum = torch.empty_like(prefix_sum) -+ out = _k2(layer, runner, core, prefix_sum, block_residual, reset=True) -+ else: -+ out = _k2(layer, runner, core, prefix_sum, block_residual) -+ return prefix_sum, block_residual, out -+ -+ -+class FfnLayer(_Layer): -+ def launch(self, step: "Step", positions, prefix_sum, block_residual, prefix_delta): -+ from vllm.models.kimi_k3.amd.linear import _apply_attn_res -+ -+ layer = self.layer -+ if self.first: -+ self.runner.step_begin() -+ hidden_states = _apply_attn_res( -+ prefix_sum, -+ block_residual, -+ layer.self_attention_res_proj, -+ layer.self_attention_res_norm, -+ layer.prev_valid_blocks, -+ delta=prefix_delta, -+ output_norm=layer.input_layernorm, -+ quant_key=layer.self_attn.get_input_quant_key(), -+ ) -+ core = layer.self_attn.mla_attn.forward( -+ positions, hidden_states, skip_o_proj=True -+ ) -+ out = _k2(layer, self.runner, core, prefix_sum, block_residual) -+ return prefix_sum, block_residual, out -diff --git a/vllm/models/kimi_k3/amd/specialized_decode/runtime/peer.py b/vllm/models/kimi_k3/amd/specialized_decode/runtime/peer.py -new file mode 100644 -index 000000000..8e07955eb ---- /dev/null -+++ b/vllm/models/kimi_k3/amd/specialized_decode/runtime/peer.py -@@ -0,0 +1,112 @@ -+# SPDX-License-Identifier: Apache-2.0 -+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -+# Adapted from ROCm/ATOM (https://github.com/ROCm/ATOM) at a526f0d, MIT License, -+# Copyright (C) 2026, Advanced Micro Devices, Inc.: -+# atom/mono/runtime/peer_memory.py -+# atom/mono/runtime/consensus.py -+"""Peer memory for in-kernel TP collectives: one uncached buffer a rank, every -+rank holding every peer's address, and the TP agreements that must precede -+anything collective.""" -+ -+import torch -+import torch.distributed as dist -+ -+ -+class _DeviceBytes: -+ """``__cuda_array_interface__`` of ``nbytes`` bytes at a raw device address, for a -+ non-owning torch view.""" -+ -+ def __init__(self, ptr: int, nbytes: int): -+ self.__cuda_array_interface__ = { -+ "shape": (nbytes,), -+ "typestr": "|u1", -+ "data": (ptr, False), -+ "version": 3, -+ } -+ -+ -+class PeerBuffer: -+ """One uncached buffer per rank, every rank holding every peer's address -+ (``addresses``, int64, indexed by TP rank). What lives where, and why no rank -+ overwrites a peer's unread data, is the model's layout's to say. -+ -+ ``bytes`` is the whole buffer: the kernels tag every pair with the step's -+ epoch, so it is never zeroed between steps. -+ -+ The constructor is collective over ``group`` (the handle exchange, then an -+ agreement that every rank mapped every peer): every rank must reach it, so a -+ runner reaches it only after ``tp_agree``. A failed mapping on any rank raises -+ ``RuntimeError`` on all of them. -+ """ -+ -+ def __init__( -+ self, -+ nbytes: int, -+ group, -+ rank: int, -+ npes: int, -+ device: torch.device, -+ ): -+ from aiter.ops.flydsl.quick_allreduce_int4_ipc import UncachedIpcHeap -+ -+ self._heap = UncachedIpcHeap -+ self.local = UncachedIpcHeap.alloc_uncached(nbytes) -+ # a view, not an owner: close() frees the memory under it -+ self.bytes = torch.as_tensor(_DeviceBytes(self.local, nbytes), device=device) -+ self.rank = rank -+ self._opened: list[int] = [] -+ addresses = [self.local] -+ if npes > 1: -+ handles = UncachedIpcHeap.gather_object_list_via_broadcast( -+ group, UncachedIpcHeap.get_mem_handle_bytes(self.local) -+ ) -+ addresses, failure = [], None -+ try: -+ for peer, handle in enumerate(handles): -+ if peer == rank: -+ addresses.append(self.local) -+ else: -+ base = UncachedIpcHeap.open_mem_handle(handle) -+ self._opened.append(base) -+ addresses.append(base) -+ except RuntimeError as err: # every rank must still reach the agreement -+ failure = err -+ if not tp_agree(failure is None, group): -+ self.close() -+ raise RuntimeError( -+ "peer memory handshake failed on " -+ f"{'this' if failure else 'another'}" -+ " TP rank" -+ ) from failure -+ self.addresses = torch.tensor(addresses, dtype=torch.int64, device=device) -+ -+ def close(self) -> None: -+ """Close the peer mappings and free the local buffer; idempotent. Only once no -+ kernel of any rank can still touch it (the runner is being dropped).""" -+ opened, self._opened = self._opened, [] -+ for base in opened: -+ self._heap.close_mem_handle(base) -+ if self.local: -+ self.bytes = None -+ self._heap.free_device_mem(self.local) -+ self.local = 0 -+ -+ -+# Decisions every TP rank must take alike. -+# -+# A mono kernel waits in-kernel on every peer rank, so a rank that turned mono off -+# alone -- or failed to build what the others launch -- leaves the others spinning -+# on the GPU for ever. Each such decision is a local verdict agreed on over the TP -+# group's CPU group before anything collective (peer handshake, launch) follows. -+# -+ -+ -+def tp_agree(ok: bool, group) -> bool: -+ """True when every rank of ``group`` passed ``ok`` true. Every rank must call -+ it at the same point; ``group`` is a CPU (gloo) process group, or None for a -+ single rank.""" -+ if group is None: -+ return ok -+ flag = torch.tensor([int(ok)], dtype=torch.int32) -+ dist.all_reduce(flag, op=dist.ReduceOp.MIN, group=group) -+ return bool(flag.item()) -diff --git a/vllm/models/kimi_k3/amd/specialized_decode/runtime/runner.py b/vllm/models/kimi_k3/amd/specialized_decode/runtime/runner.py -new file mode 100644 -index 000000000..85119c7bd ---- /dev/null -+++ b/vllm/models/kimi_k3/amd/specialized_decode/runtime/runner.py -@@ -0,0 +1,76 @@ -+# SPDX-License-Identifier: Apache-2.0 -+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -+"""The specialized decode's per-rank state: one scratch buffer every launch shares -+(their mailbox tags differ), the device step epoch, the health words and the TP -+group's peer memory. -+ -+Mailbox tags carry the epoch the model bumps once a forward (``step_begin``, -+which a CUDA graph captures), so no buffer is zeroed between steps. -+ -+The peer buffer's first ``PEER_HEAD`` bytes hold the output rank's flags, a -+word a rank (``alarm``): a rank whose wait gives up sets its word there before it -+stores anything else. The kernels see the buffer past the head. -+""" -+ -+import torch -+ -+from vllm.distributed import get_tp_group -+from vllm.models.kimi_k3.amd.specialized_decode.mono.abi import PEER_HEAD -+from vllm.models.kimi_k3.amd.specialized_decode.mono.common.plan import HEALTH_WORDS -+from vllm.models.kimi_k3.amd.specialized_decode.runtime import MAX_TOKENS -+from vllm.models.kimi_k3.amd.specialized_decode.runtime.peer import PeerBuffer -+ -+ -+class MonoRunner: -+ """Collective: every TP rank builds it at the same point.""" -+ -+ def __init__(self, device: torch.device): -+ from vllm.models.kimi_k3.amd.specialized_decode.mono import layer -+ from vllm.models.kimi_k3.amd.specialized_decode.mono.attention import kda -+ -+ self.scratch = torch.zeros( -+ max( -+ kda.scratch_bytes( -+ kda.KdaPreBuild( -+ tokens=MAX_TOKENS, -+ qlen=MAX_TOKENS, -+ nblocks=8, -+ delta=True, -+ state_len=3 + MAX_TOKENS, -+ ) -+ ), -+ layer.scratch_bytes(layer.K2Build(tokens=MAX_TOKENS, nblocks=8)), -+ ), -+ dtype=torch.uint8, -+ device=device, -+ ) -+ self.epoch = torch.zeros(1, dtype=torch.int32, device=device) -+ self.health = torch.zeros(HEALTH_WORDS, dtype=torch.int32, device=device) -+ -+ tp = get_tp_group() -+ self.rank = tp.rank_in_group -+ nbytes = PEER_HEAD + layer.peer_bytes(MAX_TOKENS) -+ self.peers = PeerBuffer(nbytes, tp.cpu_group, self.rank, tp.world_size, device) -+ self.peers.bytes.zero_() -+ self.peer_table = self.peers.addresses + PEER_HEAD -+ self.alarm = int(self.peers.addresses[0].item()) + 4 * self.rank -+ self.flags = self.peers.bytes[: 4 * tp.world_size].view(torch.int32) -+ self.fault: dict = {} # K2Build's fault_* fields (tests only) -+ self.first_layer = -1 # the step's first kernel layer: bumps the epoch -+ self.kda_prefix = "" # a KDA layer's name: the step's metadata -+ self.num_spec = 0 # speculative tokens a request -+ -+ def step_begin(self) -> None: -+ """Once a forward, before its first mono launch (a graph captures it).""" -+ self.epoch.add_(1) -+ -+ def launch_args(self) -> dict: -+ return { -+ "scratch": self.scratch, -+ "epoch": self.epoch, -+ "health": self.health, -+ "alarm": self.alarm, -+ } -+ -+ def peer_args(self) -> dict: -+ return {"peers": self.peer_table, "rank": self.rank} diff --git a/inferencex-e2e/benchmarks/multi_node/srt-slurm-recipes/configs/patches/kimik3-nightly/02-vllm-58743-state.patch b/inferencex-e2e/benchmarks/multi_node/srt-slurm-recipes/configs/patches/kimik3-nightly/02-vllm-58743-state.patch deleted file mode 100644 index 1e94c67316..0000000000 --- a/inferencex-e2e/benchmarks/multi_node/srt-slurm-recipes/configs/patches/kimik3-nightly/02-vllm-58743-state.patch +++ /dev/null @@ -1,155 +0,0 @@ -From a884f79a30c723c7fbbe5a83f099e8c7ab97d1f9 Mon Sep 17 00:00:00 2001 -From: Yichao Zhu -Date: Sat, 26 Sep 2026 07:35:45 +0800 -Subject: [PATCH] [ROCm][Kimi-K3] Support BF16 KDA recurrent state - -Signed-off-by: Yichao Zhu -(cherry picked from commit e09e870c38649ac9db758133054483611349ab53) ---- - tests/models/kimi_k3/test_amd_kda_chunk.py | 23 ++++++++++++++++ - vllm/models/kimi_k3/amd/kda.py | 4 ++- - vllm/models/kimi_k3/amd/linear.py | 4 ++- - vllm/models/kimi_k3/amd/ops/kda_prefill.py | 31 +++++++++++++++------- - 4 files changed, 51 insertions(+), 11 deletions(-) - -diff --git a/tests/models/kimi_k3/test_amd_kda_chunk.py b/tests/models/kimi_k3/test_amd_kda_chunk.py -index 2dcf886c9..0d0944032 100644 ---- a/tests/models/kimi_k3/test_amd_kda_chunk.py -+++ b/tests/models/kimi_k3/test_amd_kda_chunk.py -@@ -701,6 +701,29 @@ def test_paged_state_matches_the_dense_pair( - ) - - -+@pytest.mark.parametrize("use_fused", [False, True]) -+def test_paged_bf16_state_matches_dense_pair(use_fused: bool) -> None: -+ """The FP32 walk may load and store a BF16 persistent cache.""" -+ if use_fused: -+ _requires_kernel() -+ seqlens = [513, 64, 1] -+ inp = _inputs(seqlens, seed=17) -+ rows = [2, 7, 11] -+ cache, idx, warm = _paged_cache(inp, [True] * len(seqlens), rows, slots=16) -+ cache = cache.to(torch.bfloat16) -+ -+ dense = dict(inp) -+ dense["h0"] = cache[idx.long()].float() -+ o_ref, ht_ref = _run(dense, use_fused=use_fused) -+ o_got, ht_got = _run_paged(inp, cache, idx, warm, use_fused=use_fused) -+ -+ assert ht_got is None -+ torch.testing.assert_close(o_got, o_ref, rtol=5e-3, atol=5e-3) -+ torch.testing.assert_close( -+ cache[idx.long()], ht_ref.to(torch.bfloat16), rtol=5e-3, atol=5e-3 -+ ) -+ -+ - @pytest.mark.parametrize("use_fused", [False, True]) - def test_cold_rows_start_from_zero_not_from_cache_junk(use_fused: bool) -> None: - """`has_initial_state=False` must ignore whatever the row already holds.""" -diff --git a/vllm/models/kimi_k3/amd/kda.py b/vllm/models/kimi_k3/amd/kda.py -index 73e357e44..4216c6f94 100644 ---- a/vllm/models/kimi_k3/amd/kda.py -+++ b/vllm/models/kimi_k3/amd/kda.py -@@ -78,7 +78,9 @@ class KimiK3DeltaAttention(GatedDeltaNetAttention): - if self.model_config is None or self.cache_config is None: - raise ValueError("model_config and cache_config must be set") - return MambaStateDtypeCalculator.kda_state_dtype( -- self.model_config.dtype, self.cache_config.mamba_cache_dtype -+ self.model_config.dtype, -+ self.cache_config.mamba_cache_dtype, -+ self.cache_config.mamba_ssm_cache_dtype, - ) - - def get_state_shape(self) -> tuple[tuple[int, ...], tuple[int, ...]]: -diff --git a/vllm/models/kimi_k3/amd/linear.py b/vllm/models/kimi_k3/amd/linear.py -index 029636146..d080c34e7 100644 ---- a/vllm/models/kimi_k3/amd/linear.py -+++ b/vllm/models/kimi_k3/amd/linear.py -@@ -1125,7 +1125,9 @@ class KimiLinearForCausalLM( - vllm_config: "VllmConfig", - ) -> tuple[torch.dtype, torch.dtype]: - return MambaStateDtypeCalculator.kda_state_dtype( -- vllm_config.model_config.dtype, vllm_config.cache_config.mamba_cache_dtype -+ vllm_config.model_config.dtype, -+ vllm_config.cache_config.mamba_cache_dtype, -+ vllm_config.cache_config.mamba_ssm_cache_dtype, - ) - - @classmethod -diff --git a/vllm/models/kimi_k3/amd/ops/kda_prefill.py b/vllm/models/kimi_k3/amd/ops/kda_prefill.py -index 040f3afe6..18d90b811 100644 ---- a/vllm/models/kimi_k3/amd/ops/kda_prefill.py -+++ b/vllm/models/kimi_k3/amd/ops/kda_prefill.py -@@ -86,9 +86,10 @@ def chunk_kda_prefill( - checkpoint_offsets: per-sequence token offset to snapshot at, ``0`` - for none. - checkpoint_state_indices: optional per-sequence destination row. -- state_cache: the paged recurrent state. When given, the fused backend -- reads and writes it in place and neither a gather nor a scatter is -- needed around this call; the returned final state is ``None``. -+ state_cache: the paged recurrent state. The fused backend addresses an -+ FP32 cache directly. Other cache dtypes are gathered into the FP32 -+ accumulator format and cast back after the walk. The returned final -+ state is ``None`` in either case. - state_indices: per-sequence cache row. - has_initial_state: per-sequence flag; false starts from a zero state. - -@@ -150,11 +151,11 @@ def chunk_kda_prefill( - has_initial_state=has_initial_state, - ) - -- # Only the fused walk addresses the paged rows directly. The Triton path -- # gathers the rows it needs and scatters the results back here, so callers -- # pass the cache the same way for either backend. -+ # Only the fused walk addresses FP32 paged rows directly. The Triton path -+ # and lower-precision persistent caches gather the active rows and scatter -+ # the final state back here, so callers use one cache API for every backend. - scatter_to: torch.Tensor | None = None -- if state_cache is not None and not fused: -+ if state_cache is not None and (not fused or state_cache.dtype != torch.float32): - initial_state = gather_initial_states( - state_cache, state_indices, has_initial_state - ) -@@ -165,6 +166,11 @@ def chunk_kda_prefill( - # Restated for the type checker; `fused` already implies all three. - assert cu_seqlens is not None and lower_bound is not None - assert g_bias is not None -+ # The fused HIP prefill kernel accumulates recurrent state in FP32. -+ # Convert only active rows; the persistent cache keeps its configured -+ # dtype and receives a cast-back below. -+ if initial_state is not None and initial_state.dtype != torch.float32: -+ initial_state = initial_state.to(torch.float32) - logger.info_once( - "Kimi-K3 KDA prefill: dispatching the fused ROCm chunk kernel." - ) -@@ -181,7 +187,7 @@ def chunk_kda_prefill( - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - ) -- return fused_kda_chunk( -+ o, final_state = fused_kda_chunk( - qg=ws["qg"], - w=ws["w"], - u=ws["u"], -@@ -208,6 +214,12 @@ def chunk_kda_prefill( - state_indices=state_indices, - has_initial_state=has_initial_state, - ) -+ if scatter_to is not None: -+ assert state_indices is not None -+ assert final_state is not None -+ scatter_to[state_indices.long()] = final_state.to(scatter_to.dtype) -+ return o, None -+ return o, final_state - - o, final_state = chunk_kda_with_fused_gate( - q=q, -@@ -231,6 +243,7 @@ def chunk_kda_prefill( - o = out - if scatter_to is not None: - assert state_indices is not None -- scatter_to[state_indices.long()] = final_state -+ assert final_state is not None -+ scatter_to[state_indices.long()] = final_state.to(scatter_to.dtype) - return o, None - return o, final_state diff --git a/inferencex-e2e/benchmarks/multi_node/srt-slurm-recipes/configs/patches/kimik3-nightly/03-vllm-58743-checkpoints.patch b/inferencex-e2e/benchmarks/multi_node/srt-slurm-recipes/configs/patches/kimik3-nightly/03-vllm-58743-checkpoints.patch deleted file mode 100644 index 1e1f5c075e..0000000000 --- a/inferencex-e2e/benchmarks/multi_node/srt-slurm-recipes/configs/patches/kimik3-nightly/03-vllm-58743-checkpoints.patch +++ /dev/null @@ -1,251 +0,0 @@ -From c8a4c72073306a2d2b7370e45bd91be11c0df796 Mon Sep 17 00:00:00 2001 -From: Yichao Zhu -Date: Sat, 17 Oct 2026 12:35:33 +0800 -Subject: [PATCH] fix(kimik3): support BF16 prefill checkpoint stores on ROCm - -Preserve FP32 accumulation and write valid checkpoint rows in their persistent cache dtype. Extend the existing checkpoint regressions for BF16 destinations and padded rows. - -Signed-off-by: Yichao Zhu -(cherry picked from commit 53bfc127fdc8e87d81e1a50dbb773f197e9df875) ---- - csrc/kimi_k3/fused_kda_chunk_kernel_rocm.cu | 31 +++++++++---- - tests/models/kimi_k3/test_amd_kda_chunk.py | 48 ++++++++++++++------- - 2 files changed, 55 insertions(+), 24 deletions(-) - -diff --git a/csrc/kimi_k3/fused_kda_chunk_kernel_rocm.cu b/csrc/kimi_k3/fused_kda_chunk_kernel_rocm.cu -index bb9485dd7..3ed350c85 100644 ---- a/csrc/kimi_k3/fused_kda_chunk_kernel_rocm.cu -+++ b/csrc/kimi_k3/fused_kda_chunk_kernel_rocm.cu -@@ -48,6 +48,7 @@ - - using bf16_t = __bf16; - using bf16x2 = __bf16 __attribute__((ext_vector_type(2))); -+using bf16x4 = __bf16 __attribute__((ext_vector_type(4))); - using bf16x8 = __bf16 __attribute__((ext_vector_type(8))); - using f32x4 = float __attribute__((ext_vector_type(4))); - using u32x2 = uint32_t __attribute__((ext_vector_type(2))); -@@ -151,7 +152,8 @@ struct Params { - // replaying the tail as a second model forward. `ckidx`, when present, - // redirects sequence n's snapshot to row `ckidx[n]`, which lets the store go - // straight into the paged state cache; a negative row disables it. -- float* ckpt; -+ void* ckpt; -+ bool ckpt_bf16; - const int* ckoff; - const int* ckidx; - int G; -@@ -842,11 +844,18 @@ __global__ __launch_bounds__(64 * WAVES * KS, - // Same store as the final state below, at an interior boundary instead of - // the last one. - auto store_checkpoint = [&]() { -- float* dst = p.ckpt + ck_dst + (v0 + lrow) * kK + 4 * lgrp; -+ const int64_t base = ck_dst + (v0 + lrow) * kK + 4 * lgrp; - for (int vt = 0; vt < kNV; ++vt) -- for (int kt = 0; kt < kNKTL; ++kt) -- *reinterpret_cast(dst + vt * 16 * kK + (kt0 + kt) * 16) = -- S[kt][vt]; -+ for (int kt = 0; kt < kNKTL; ++kt) { -+ const int64_t offset = base + vt * 16 * kK + (kt0 + kt) * 16; -+ if (p.ckpt_bf16) { -+ *reinterpret_cast(static_cast(p.ckpt) + offset) = -+ __builtin_convertvector(S[kt][vt], bf16x4); -+ } else { -+ *reinterpret_cast(static_cast(p.ckpt) + offset) = -+ S[kt][vt]; -+ } -+ } - }; - - const int tail = last_group ? T - (nt_all - 1) * kBT : kBT; -@@ -2112,7 +2121,7 @@ void fused_kda_chunk( - // Prefill checkpoint export. `checkpoint_state` is either a [N, H, V, K] - // staging buffer, as FlashKDA takes, or the paged state cache itself when - // `checkpoint_state_indices` says which row each sequence owns. -- float* ckpt_ptr = nullptr; -+ void* ckpt_ptr = nullptr; - int const* ckoff_ptr = nullptr; - int const* ckidx_ptr = nullptr; - if (checkpoint_state.has_value() || checkpoint_offsets.has_value()) { -@@ -2121,7 +2130,8 @@ void fused_kda_chunk( - "checkpoint_state and checkpoint_offsets must be given together"); - STD_TORCH_CHECK( - checkpoint_state->is_cuda() && -- checkpoint_state->scalar_type() == ScalarType::Float && -+ (checkpoint_state->scalar_type() == ScalarType::Float || -+ checkpoint_state->scalar_type() == ScalarType::BFloat16) && - checkpoint_state->dim() == 4 && - checkpoint_state->size(1) == num_heads && - checkpoint_state->size(2) == kV && -@@ -2129,7 +2139,8 @@ void fused_kda_chunk( - checkpoint_state->stride(3) == 1 && - checkpoint_state->stride(2) == kK && - checkpoint_state->stride(1) == kV * kK, -- "checkpoint_state must be an fp32 [rows, H, 128, 128] tensor whose " -+ "checkpoint_state must be an fp32 or bf16 [rows, H, 128, 128] tensor " -+ "whose " - "rows are dense; stride(0) may be padded, as the paged state cache " - "it usually is"); - STD_TORCH_CHECK(checkpoint_offsets->is_cuda() && -@@ -2154,7 +2165,7 @@ void fused_kda_chunk( - checkpoint_state->size(0), " rows for ", num_seqs, - " sequences"); - } -- ckpt_ptr = static_cast(checkpoint_state->data_ptr()); -+ ckpt_ptr = checkpoint_state->data_ptr(); - ckoff_ptr = static_cast(checkpoint_offsets->data_ptr()); - } - -@@ -2238,6 +2249,8 @@ void fused_kda_chunk( - params.ckrow = checkpoint_state.has_value() ? checkpoint_state->stride(0) - : num_heads * kV * kK; - params.ckpt = ckpt_ptr; -+ params.ckpt_bf16 = checkpoint_state.has_value() && -+ checkpoint_state->scalar_type() == ScalarType::BFloat16; - params.ckoff = ckoff_ptr; - params.ckidx = ckidx_ptr; - -diff --git a/tests/models/kimi_k3/test_amd_kda_chunk.py b/tests/models/kimi_k3/test_amd_kda_chunk.py -index 0d0944032..e3e000137 100644 ---- a/tests/models/kimi_k3/test_amd_kda_chunk.py -+++ b/tests/models/kimi_k3/test_amd_kda_chunk.py -@@ -447,13 +447,13 @@ def test_accepts_chunk_metadata_from_the_attention_builder( - torch.testing.assert_close(state_got, state_ref, rtol=1e-3, atol=1e-3) - - --def _checkpoint_buffer(rows: int) -> torch.Tensor: -+def _checkpoint_buffer(rows: int, dtype: torch.dtype = torch.float32) -> torch.Tensor: - """A destination filled with NaN, so an unwritten row is unmistakable.""" - return torch.full( - (rows, NUM_HEADS, HEAD_DIM, HEAD_DIM), - float("nan"), - device="cuda", -- dtype=torch.float32, -+ dtype=dtype, - ) - - -@@ -470,9 +470,13 @@ def _prefix_inputs(inp: dict, seq: int, num_tokens: int) -> dict: - - - def _run_with_checkpoints( -- inp: dict, offsets: list[int], rows: int | None = None, **kw -+ inp: dict, -+ offsets: list[int], -+ rows: int | None = None, -+ dtype: torch.dtype = torch.float32, -+ **kw, - ) -> tuple[torch.Tensor, tuple]: -- ckpt = _checkpoint_buffer(len(offsets) if rows is None else rows) -+ ckpt = _checkpoint_buffer(len(offsets) if rows is None else rows, dtype) - result = _run( - inp, - use_fused=True, -@@ -534,7 +538,10 @@ def test_checkpoint_leaves_the_output_and_final_state_alone() -> None: - - - @pytest.mark.parametrize("offset", [100, 63, 65, 1024]) --def test_unusable_checkpoint_offsets_write_nothing(offset: int) -> None: -+@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) -+def test_unusable_checkpoint_offsets_write_nothing( -+ offset: int, dtype: torch.dtype -+) -> None: - """Offsets the walk never reaches must leave the destination untouched. - - The recurrent state only exists at chunk boundaries, so an offset that is -@@ -550,12 +557,13 @@ def test_unusable_checkpoint_offsets_write_nothing(offset: int) -> None: - - assert offset % KDA_CHECKPOINT_ALIGNMENT != 0 or offset > 512 - inp = _inputs([512], seed=31) -- ckpt, _ = _run_with_checkpoints(inp, [offset]) -+ ckpt, _ = _run_with_checkpoints(inp, [offset], dtype=dtype) - - assert ckpt.isnan().all() - - --def test_checkpoint_state_indices_pick_the_destination_row() -> None: -+@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) -+def test_checkpoint_state_indices_pick_the_destination_row(dtype: torch.dtype) -> None: - """With indices the export lands straight in the paged state cache. - - FlashKDA writes one row per sequence into a staging buffer that the caller -@@ -574,6 +582,7 @@ def test_checkpoint_state_indices_pick_the_destination_row() -> None: - inp, - [64, 64, 128], - rows=rows, -+ dtype=dtype, - checkpoint_state_indices=torch.tensor( - indices, device="cuda", dtype=torch.int32 - ), -@@ -586,7 +595,7 @@ def test_checkpoint_state_indices_pick_the_destination_row() -> None: - continue - seq, offset = written[row] - _, truncated = _run(_prefix_inputs(inp, seq, offset), use_fused=True) -- torch.testing.assert_close(ckpt[row], truncated[0], rtol=0, atol=0) -+ torch.testing.assert_close(ckpt[row], truncated[0].to(dtype), rtol=0, atol=0) - - - @pytest.mark.parametrize("groups", [2, 4]) -@@ -811,11 +820,14 @@ def test_paged_state_survives_the_group_split(groups: int, monkeypatch) -> None: - torch.testing.assert_close(cache[idx.long()], ht_ref, rtol=5e-3, atol=5e-3) - - --def test_paged_state_and_checkpoint_export_together() -> None: -+@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) -+def test_paged_state_and_checkpoint_export_together(dtype: torch.dtype) -> None: - """The checkpoint lands in its own row while the walk owns the state row.""" - _requires_kernel() - inp = _inputs([1024], seed=14) -+ inp["h0"] = inp["h0"].to(dtype).float() - cache, idx, warm = _paged_cache(inp, [True], [3], slots=12) -+ cache = cache.to(dtype) - ckidx = torch.tensor([8], device="cuda", dtype=torch.int32) - - o_ref, ht_ref = _run(inp, use_fused=True) -@@ -833,8 +845,10 @@ def test_paged_state_and_checkpoint_export_together() -> None: - checkpoint_state_indices=ckidx, - ) - torch.testing.assert_close(o_got, o_ref, rtol=5e-3, atol=5e-3) -- torch.testing.assert_close(cache[idx.long()], ht_ref, rtol=5e-3, atol=5e-3) -- torch.testing.assert_close(cache[ckidx.long()], ck_ref, rtol=5e-3, atol=5e-3) -+ torch.testing.assert_close(cache[idx.long()].float(), ht_ref, rtol=5e-3, atol=5e-3) -+ torch.testing.assert_close( -+ cache[ckidx.long()].float(), ck_ref, rtol=5e-3, atol=5e-3 -+ ) - - - def test_paged_state_rejects_a_dense_initial_state() -> None: -@@ -898,7 +912,8 @@ def test_padded_cache_stride_is_honoured(use_fused: bool) -> None: - ) - - --def test_padded_checkpoint_stride_is_honoured() -> None: -+@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) -+def test_padded_checkpoint_stride_is_honoured(dtype: torch.dtype) -> None: - """The checkpoint destination is the paged cache, which pads stride(0). - - Every other checkpoint case here builds a contiguous buffer, so none of -@@ -908,9 +923,10 @@ def test_padded_checkpoint_stride_is_honoured() -> None: - _requires_kernel() - seqlens = [1024] - inp = _inputs(seqlens, seed=31) -+ inp["h0"] = inp["h0"].to(dtype).float() - row = NUM_HEADS * HEAD_DIM * HEAD_DIM - slots, pad = 16, 512 -- raw = torch.full((slots, row + pad), -7.0, device="cuda", dtype=torch.float32) -+ raw = torch.full((slots, row + pad), -7.0, device="cuda", dtype=dtype) - cache = raw.as_strided( - (slots, NUM_HEADS, HEAD_DIM, HEAD_DIM), - (row + pad, HEAD_DIM * HEAD_DIM, HEAD_DIM, 1), -@@ -942,8 +958,10 @@ def test_padded_checkpoint_stride_is_honoured() -> None: - ) - - torch.testing.assert_close(o_got, o_ref, rtol=5e-3, atol=5e-3) -- torch.testing.assert_close(cache[idx.long()], ht_ref, rtol=5e-3, atol=5e-3) -- torch.testing.assert_close(cache[ck_rows[0]], ck_ref[0], rtol=5e-3, atol=5e-3) -+ torch.testing.assert_close(cache[idx.long()].float(), ht_ref, rtol=5e-3, atol=5e-3) -+ torch.testing.assert_close( -+ cache[ck_rows[0]].float(), ck_ref[0], rtol=5e-3, atol=5e-3 -+ ) - - touched = torch.zeros(slots, dtype=torch.bool, device="cuda") - touched[state_rows + ck_rows] = True diff --git a/inferencex-e2e/benchmarks/multi_node/srt-slurm-recipes/configs/patches/kimik3-nightly/04-vllm-58743-native-prefill.patch b/inferencex-e2e/benchmarks/multi_node/srt-slurm-recipes/configs/patches/kimik3-nightly/04-vllm-58743-native-prefill.patch deleted file mode 100644 index c01df8737b..0000000000 --- a/inferencex-e2e/benchmarks/multi_node/srt-slurm-recipes/configs/patches/kimik3-nightly/04-vllm-58743-native-prefill.patch +++ /dev/null @@ -1,512 +0,0 @@ -From 7542922112f99a0af8a3fb636a2d1bb5476c61e6 Mon Sep 17 00:00:00 2001 -From: Yichao Zhu -Date: Mon, 19 Oct 2026 17:23:49 +0800 -Subject: [PATCH] Support native BF16 paged states in ROCm KDA prefill - -Read and write BF16 persistent states in the HIP chunk kernel while retaining FP32 accumulation and scan workspaces. Avoid gathering and scattering paged state, preserving inactive cache rows. Extend the existing tests for BF16 paging and checkpoint resume under graph replay. - -Signed-off-by: Yichao Zhu -(cherry picked from commit 85c7ecc6eae032f6af5ab44485ca42fb76ca31a3) ---- - csrc/kimi_k3/fused_kda_chunk_kernel_rocm.cu | 98 +++++++++------- - tests/models/kimi_k3/test_amd_kda_chunk.py | 123 +++++++++++++++++--- - vllm/models/kimi_k3/amd/ops/kda_chunk.py | 4 +- - vllm/models/kimi_k3/amd/ops/kda_prefill.py | 26 ++--- - 4 files changed, 172 insertions(+), 79 deletions(-) - -diff --git a/csrc/kimi_k3/fused_kda_chunk_kernel_rocm.cu b/csrc/kimi_k3/fused_kda_chunk_kernel_rocm.cu -index 3ed350c85..f6f9845dd 100644 ---- a/csrc/kimi_k3/fused_kda_chunk_kernel_rocm.cu -+++ b/csrc/kimi_k3/fused_kda_chunk_kernel_rocm.cu -@@ -121,11 +121,12 @@ struct Params { - // separate pointers so that a paged cache (both equal) and a dense - // [N, H, V, K] pair (two tensors, identity rows) share one code path. - // Either pointer may be null. -- const float* cache_in; -- float* cache_out; -+ const void* cache_in; -+ void* cache_out; -+ bool cache_bf16; - const int* sidx; - const bool* hh0; -- // Floats between consecutive rows: the tensor's stride(0), which a paged -+ // Elements between consecutive rows: the tensor's stride(0), which a paged - // cache may pad beyond H * kV * kK. Only each row itself is dense. - int64_t crow; - int64_t ckrow; // the same, for the checkpoint destination -@@ -165,6 +166,17 @@ struct Params { - // with a leading dimension of 0, so nothing inside the chunk loop changes. - [[maybe_unused]] __device__ bf16_t g_uzero[kK]; - -+__device__ __forceinline__ void store_state4(void* dst, int64_t offset, -+ f32x4 value, bool bf16) { -+ if (bf16) { -+ const bf16x4 packed = __builtin_convertvector(value, bf16x4); -+ __builtin_memcpy(static_cast(dst) + offset, &packed, -+ sizeof(packed)); -+ } else { -+ *reinterpret_cast(static_cast(dst) + offset) = value; -+ } -+} -+ - union Frag { - bf16x8 v; - uint32_t d[4]; -@@ -760,13 +772,22 @@ __global__ __launch_bounds__(64 * WAVES * KS, - // state the scan composed for it; in pass one they start at zero, which is - // the B_g the scan expects. Transfer blocks start at the identity. - if (HAS_H0 && g == 0 && warm && !from_sin) { -- const float* s0 = -- p.cache_in + h0_row + -- ((static_cast(h) * kV + v0 + lrow) * kK + 4 * lgrp); -+ const int64_t base = -+ h0_row + ((static_cast(h) * kV + v0 + lrow) * kK + 4 * lgrp); - for (int vt = 0; vt < kNV; ++vt) -- for (int kt = 0; kt < kNKTL; ++kt) -- S[kt][vt] = *reinterpret_cast(s0 + vt * 16 * kK + -- (kt0 + kt) * 16); -+ for (int kt = 0; kt < kNKTL; ++kt) { -+ const int64_t offset = base + vt * 16 * kK + (kt0 + kt) * 16; -+ if (p.cache_bf16) { -+ bf16x4 packed; -+ __builtin_memcpy(&packed, -+ static_cast(p.cache_in) + offset, -+ sizeof(packed)); -+ S[kt][vt] = __builtin_convertvector(packed, f32x4); -+ } else { -+ S[kt][vt] = *reinterpret_cast( -+ static_cast(p.cache_in) + offset); -+ } -+ } - } else if (PASS2 && (g > 0 || (from_sin && warm))) { - const float* s0 = p.sin_ + - (static_cast(g) * p.NH() + nh) * -@@ -848,13 +869,7 @@ __global__ __launch_bounds__(64 * WAVES * KS, - for (int vt = 0; vt < kNV; ++vt) - for (int kt = 0; kt < kNKTL; ++kt) { - const int64_t offset = base + vt * 16 * kK + (kt0 + kt) * 16; -- if (p.ckpt_bf16) { -- *reinterpret_cast(static_cast(p.ckpt) + offset) = -- __builtin_convertvector(S[kt][vt], bf16x4); -- } else { -- *reinterpret_cast(static_cast(p.ckpt) + offset) = -- S[kt][vt]; -- } -+ store_state4(p.ckpt, offset, S[kt][vt], p.ckpt_bf16); - } - }; - -@@ -1058,15 +1073,18 @@ __global__ __launch_bounds__(64 * WAVES * KS, - - // With G == 1 this is the final state. With G > 1 it is what the group - // reaches from a zero start, which the scan then composes. -- float* s1 = nullptr; -- int64_t st_vt = 16 * kK, st_kt = 16; // the cache row is [H][V][K] fp32 - if constexpr (PASS2) { -- // Only the last group holds the sequence's final state. The row's inner -- // layout is [V][K], so this is one f32x4 per lane per tile, strictly -- // inside the row this block owns. -- if (last_group && p.cache_out != nullptr) -- s1 = p.cache_out + h0_row + -- ((static_cast(h) * kV + v0 + lrow) * kK + 4 * lgrp); -+ // Only the last group writes the persistent state; group workspaces -+ // retain the accumulator precision regardless of the cache dtype. -+ if (last_group && p.cache_out != nullptr) { -+ const int64_t base = -+ h0_row + ((static_cast(h) * kV + v0 + lrow) * kK + 4 * lgrp); -+ for (int vt = 0; vt < kNV; ++vt) -+ for (int kt = 0; kt < kNKTL; ++kt) -+ store_state4(p.cache_out, base + vt * 16 * kK + (kt0 + kt) * 16, -+ S[kt][vt], p.cache_bf16); -+ } -+ return; - } else if (xfer) { - // Publish C_g = M_g - diag(A_g). A_g is broadcast through the - // double-buffer half the last chunk did not use, so the diagonal comes off -@@ -1095,18 +1113,15 @@ __global__ __launch_bounds__(64 * WAVES * KS, - } else { - if (slice == 0 && tid < kK) - p.ag[(static_cast(g) * p.NH() + nh) * kK + tid] = aacc; -- s1 = p.bg + -- (static_cast(g) * p.NH() + nh) * -- (static_cast(kV) * kK) + -- static_cast(v0 / 16) * kStTile + lane * 4; -- st_vt = kStTile; -- st_kt = kStKt; -- } -- if (s1 != nullptr) -+ float* const s1 = p.bg + -+ (static_cast(g) * p.NH() + nh) * -+ (static_cast(kV) * kK) + -+ static_cast(v0 / 16) * kStTile + lane * 4; - for (int vt = 0; vt < kNV; ++vt) - for (int kt = 0; kt < kNKTL; ++kt) -- *reinterpret_cast(s1 + vt * st_vt + (kt0 + kt) * st_kt) = -+ *reinterpret_cast(s1 + vt * kStTile + (kt0 + kt) * kStKt) = - S[kt][vt]; -+ } - #else - kda_unsupported_arch("kda_chunk_fused"); - #endif // ifdef ROCM_USE_FUSED_PREFILL_KDA -@@ -2173,8 +2188,8 @@ void fused_kda_chunk( - // reads and writes row `state_indices[n]` of a paged cache. Without it, the - // dense `initial_state` / `final_state` pair is used, which is the same - // layout with one row per sequence in order. -- float const* cache_in = h0_ptr; -- float* cache_out = ht_ptr; -+ void const* cache_in = h0_ptr; -+ void* cache_out = ht_ptr; - int const* sidx_ptr = nullptr; - bool const* hh0_ptr = nullptr; - if (state_cache.has_value()) { -@@ -2186,13 +2201,14 @@ void fused_kda_chunk( - "state_cache needs state_indices and has_initial_state"); - STD_TORCH_CHECK( - state_cache->is_cuda() && -- state_cache->scalar_type() == ScalarType::Float && -+ (state_cache->scalar_type() == ScalarType::Float || -+ state_cache->scalar_type() == ScalarType::BFloat16) && - state_cache->dim() == 4 && state_cache->size(1) == num_heads && - state_cache->size(2) == kV && state_cache->size(3) == kK && - state_cache->stride(3) == 1 && state_cache->stride(2) == kK && - state_cache->stride(1) == kV * kK, -- "state_cache must be an fp32 [slots, H, 128, 128] tensor whose rows " -- "are dense; stride(0) may be padded"); -+ "state_cache must be an fp32 or bf16 [slots, H, 128, 128] tensor whose " -+ "rows are dense; stride(0) may be padded"); - STD_TORCH_CHECK(state_indices->is_cuda() && - state_indices->is_contiguous() && - state_indices->scalar_type() == ScalarType::Int && -@@ -2208,8 +2224,8 @@ void fused_kda_chunk( - // A block reads row `state_indices[n]` at entry and writes it at exit, so - // `checkpoint_state_indices` must name no row in `state_indices`; a - // checkpoint aimed at one would be a write racing that read. -- cache_in = static_cast(state_cache->data_ptr()); -- cache_out = static_cast(state_cache->data_ptr()); -+ cache_in = state_cache->data_ptr(); -+ cache_out = state_cache->data_ptr(); - sidx_ptr = static_cast(state_indices->data_ptr()); - hh0_ptr = static_cast(has_initial_state->data_ptr()); - } -@@ -2242,6 +2258,8 @@ void fused_kda_chunk( - params.NHv = static_cast(nh_total); - params.cache_in = cache_in; - params.cache_out = cache_out; -+ params.cache_bf16 = state_cache.has_value() && -+ state_cache->scalar_type() == ScalarType::BFloat16; - params.sidx = sidx_ptr; - params.hh0 = hh0_ptr; - params.crow = -diff --git a/tests/models/kimi_k3/test_amd_kda_chunk.py b/tests/models/kimi_k3/test_amd_kda_chunk.py -index e3e000137..c7a2c511f 100644 ---- a/tests/models/kimi_k3/test_amd_kda_chunk.py -+++ b/tests/models/kimi_k3/test_amd_kda_chunk.py -@@ -645,13 +645,15 @@ def test_triton_fallback_refuses_a_checkpoint_export() -> None: - # --------------------------------------------------------------------------- - - --def _paged_cache(inp: dict, warm: list[bool], rows: list[int], slots: int): -+def _paged_cache( -+ inp: dict, warm: list[bool], rows: list[int], slots: int, dtype=torch.float32 -+): - """A cache holding each sequence's h0 at its row, a fill pattern elsewhere.""" - cache = torch.full( - (slots, NUM_HEADS, HEAD_DIM, HEAD_DIM), - -7.0, - device="cuda", -- dtype=torch.float32, -+ dtype=dtype, - ) - idx = torch.tensor(rows, device="cuda", dtype=torch.int32) - for n, row in enumerate(rows): -@@ -734,14 +736,18 @@ def test_paged_bf16_state_matches_dense_pair(use_fused: bool) -> None: - - - @pytest.mark.parametrize("use_fused", [False, True]) --def test_cold_rows_start_from_zero_not_from_cache_junk(use_fused: bool) -> None: -+@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) -+def test_cold_rows_start_from_zero_not_from_cache_junk( -+ use_fused: bool, dtype: torch.dtype -+) -> None: - """`has_initial_state=False` must ignore whatever the row already holds.""" - _requires_kernel() - seqlens = [256, 256, 256] - inp = _inputs(seqlens, seed=12) -+ inp["h0"] = inp["h0"].to(dtype).float() - warm = [True, False, True] - rows = [4, 9, 1] -- cache, idx, warm_t = _paged_cache(inp, warm, rows, slots=16) -+ cache, idx, warm_t = _paged_cache(inp, warm, rows, slots=16, dtype=dtype) - - # The dense reference zeroes exactly the cold sequences' initial states. - ref = dict(inp) -@@ -751,11 +757,14 @@ def test_cold_rows_start_from_zero_not_from_cache_junk(use_fused: bool) -> None: - o_got, _ = _run_paged(inp, cache, idx, warm_t, use_fused=use_fused) - - torch.testing.assert_close(o_got, o_ref, rtol=5e-3, atol=5e-3) -- torch.testing.assert_close(cache[idx.long()], ht_ref, rtol=5e-3, atol=5e-3) -+ torch.testing.assert_close(cache[idx.long()].float(), ht_ref, rtol=5e-3, atol=5e-3) - - - @pytest.mark.parametrize("use_fused", [False, True]) --def test_state_indices_accept_a_block_table_column(use_fused: bool) -> None: -+@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) -+def test_state_indices_accept_a_block_table_column( -+ use_fused: bool, dtype: torch.dtype -+) -> None: - # The layer passes block_table[:, 0], which is contiguous only while the - # table has one column. Speculative decoding reserves a state page per - # draft token, so the table widens and that column becomes strided. -@@ -764,7 +773,9 @@ def test_state_indices_accept_a_block_table_column(use_fused: bool) -> None: - inp = _inputs(seqlens, seed=43) - n = len(seqlens) - rows = [3 * i + 5 for i in range(n)] -- cache, idx, warm = _paged_cache(inp, [True] * n, rows, slots=3 * n + 12) -+ cache, idx, warm = _paged_cache( -+ inp, [True] * n, rows, slots=3 * n + 12, dtype=dtype -+ ) - ref_cache = cache.clone() - - table = torch.zeros(n, 9, dtype=torch.int32, device="cuda") -@@ -780,7 +791,10 @@ def test_state_indices_accept_a_block_table_column(use_fused: bool) -> None: - - - @pytest.mark.parametrize("use_fused", [False, True]) --def test_padded_zero_length_sequence_is_inert(use_fused: bool) -> None: -+@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) -+def test_padded_zero_length_sequence_is_inert( -+ use_fused: bool, dtype: torch.dtype, monkeypatch -+) -> None: - # A batch carrying speculative decodes pads its non-spec group with - # zero-length sequences whose cache row still belongs to a request, so - # consuming no tokens has to mean writing nothing and perturbing nobody. -@@ -791,11 +805,21 @@ def test_padded_zero_length_sequence_is_inert(use_fused: bool) -> None: - padded, unpadded = _inputs([*live, 0], seed=41), _inputs(live, seed=41) - rows = [3 * i + 5 for i in range(4)] - -- cache, idx, warm = _paged_cache(padded, [True] * 4, rows, slots=24) -+ cache, idx, warm = _paged_cache(padded, [True] * 4, rows, slots=24, dtype=dtype) - before = cache.clone() -- _run_paged(padded, cache, idx, warm, use_fused=use_fused) -- -- ref_cache, ref_idx, ref_warm = _paged_cache(unpadded, [True] * 3, rows[:3], 24) -+ # Poison unwritten outputs rather than relying on allocator contents. -+ deterministic = torch.are_deterministic_algorithms_enabled() -+ warn_only = torch.is_deterministic_algorithms_warn_only_enabled() -+ monkeypatch.setattr(torch.utils.deterministic, "fill_uninitialized_memory", True) -+ try: -+ torch.use_deterministic_algorithms(True) -+ _run_paged(padded, cache, idx, warm, use_fused=use_fused) -+ finally: -+ torch.use_deterministic_algorithms(deterministic, warn_only=warn_only) -+ -+ ref_cache, ref_idx, ref_warm = _paged_cache( -+ unpadded, [True] * 3, rows[:3], 24, dtype=dtype -+ ) - _run_paged(unpadded, ref_cache, ref_idx, ref_warm, use_fused=use_fused) - - assert torch.equal(cache[rows[3]], before[rows[3]]), ( -@@ -805,19 +829,23 @@ def test_padded_zero_length_sequence_is_inert(use_fused: bool) -> None: - - - @pytest.mark.parametrize("groups", [2, 4]) --def test_paged_state_survives_the_group_split(groups: int, monkeypatch) -> None: -+@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) -+def test_paged_state_survives_the_group_split( -+ groups: int, dtype: torch.dtype, monkeypatch -+) -> None: - """At G>1 the read and the write are different blocks of one launch.""" - _requires_kernel() - _force_groups(monkeypatch, groups) - inp = _inputs([1024, 512], seed=13) -+ inp["h0"] = inp["h0"].to(dtype).float() - rows = [6, 2] -- cache, idx, warm = _paged_cache(inp, [True, True], rows, slots=12) -+ cache, idx, warm = _paged_cache(inp, [True, True], rows, slots=12, dtype=dtype) - - o_ref, ht_ref = _run(inp, use_fused=True) - o_got, _ = _run_paged(inp, cache, idx, warm, use_fused=True) - - torch.testing.assert_close(o_got, o_ref, rtol=5e-3, atol=5e-3) -- torch.testing.assert_close(cache[idx.long()], ht_ref, rtol=5e-3, atol=5e-3) -+ torch.testing.assert_close(cache[idx.long()].float(), ht_ref, rtol=5e-3, atol=5e-3) - - - @pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) -@@ -867,8 +895,64 @@ def test_paged_state_rejects_a_dense_initial_state() -> None: - ) - - -+@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) -+@pytest.mark.parametrize("groups", [1, 2, 4]) -+def test_paged_checkpoint_can_resume_with_graph_replay( -+ dtype: torch.dtype, groups: int, monkeypatch -+) -> None: -+ """Resume from a stored checkpoint, including cold zero-token padding.""" -+ _requires_kernel() -+ _force_groups(monkeypatch, groups) -+ inp = _inputs([1024], seed=47) -+ cache, idx, warm = _paged_cache(inp, [True], [3], slots=12, dtype=dtype) -+ ckidx = torch.tensor([8], device="cuda", dtype=torch.int32) -+ _run_paged( -+ inp, -+ cache, -+ idx, -+ warm, -+ use_fused=True, -+ checkpoint_state=cache, -+ checkpoint_offsets=torch.tensor([512], device="cuda", dtype=torch.int32), -+ checkpoint_state_indices=ckidx, -+ ) -+ -+ # Compare against the same rounded checkpoint, not an uninterrupted FP32 walk. -+ tail = dict(inp) -+ for key in ("q", "k", "v", "raw_g", "raw_beta"): -+ tail[key] = inp[key][:, 512:].contiguous() -+ tail["cu"] = torch.tensor([0, 512], device="cuda", dtype=torch.int32) -+ tail["h0"] = cache[ckidx.long()].float() -+ o_ref, ht_ref = _run(tail, use_fused=True) -+ -+ tail["cu"] = torch.tensor([0, 512, 512], device="cuda", dtype=torch.int32) -+ resume_idx = torch.tensor([8, 10], device="cuda", dtype=torch.int32) -+ resume_warm = torch.tensor([True, False], device="cuda", dtype=torch.bool) -+ before = cache.clone() -+ stream = torch.cuda.Stream() -+ stream.wait_stream(torch.cuda.current_stream()) -+ with torch.cuda.stream(stream): -+ _run_paged(tail, cache, resume_idx, resume_warm, use_fused=True) -+ torch.cuda.current_stream().wait_stream(stream) -+ cache.copy_(before) -+ graph = torch.cuda.CUDAGraph() -+ with torch.cuda.graph(graph): -+ o_got, ht_got = _run_paged(tail, cache, resume_idx, resume_warm, use_fused=True) -+ for _ in range(2): -+ cache.copy_(before) -+ graph.replay() -+ assert ht_got is None -+ torch.testing.assert_close(o_got, o_ref, rtol=5e-3, atol=5e-3) -+ torch.testing.assert_close( -+ cache[ckidx.long()].float(), ht_ref, rtol=5e-3, atol=5e-3 -+ ) -+ untouched = torch.arange(cache.shape[0], device="cuda") != 8 -+ assert torch.equal(cache[untouched], before[untouched]) -+ -+ - @pytest.mark.parametrize("use_fused", [False, True]) --def test_padded_cache_stride_is_honoured(use_fused: bool) -> None: -+@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) -+def test_padded_cache_stride_is_honoured(use_fused: bool, dtype: torch.dtype) -> None: - """A mamba page may pad stride(0); only each state row is dense. - - The real `recurrent_state` is not a contiguous tensor, so a kernel that -@@ -879,10 +963,11 @@ def test_padded_cache_stride_is_honoured(use_fused: bool) -> None: - _requires_kernel() - seqlens = [1024, 512] - inp = _inputs(seqlens, seed=21) -+ inp["h0"] = inp["h0"].to(dtype).float() - n = len(seqlens) - row = NUM_HEADS * HEAD_DIM * HEAD_DIM -- slots, pad = 16, 512 -- raw = torch.full((slots, row + pad), -7.0, device="cuda", dtype=torch.float32) -+ slots, pad = 16, 513 if dtype == torch.bfloat16 else 512 -+ raw = torch.full((slots, row + pad), -7.0, device="cuda", dtype=dtype) - cache = raw.as_strided( - (slots, NUM_HEADS, HEAD_DIM, HEAD_DIM), - (row + pad, HEAD_DIM * HEAD_DIM, HEAD_DIM, 1), -@@ -901,7 +986,7 @@ def test_padded_cache_stride_is_honoured(use_fused: bool) -> None: - - assert ht_got is None - torch.testing.assert_close(o_got, o_ref, rtol=5e-3, atol=5e-3) -- torch.testing.assert_close(cache[idx.long()], ht_ref, rtol=5e-3, atol=5e-3) -+ torch.testing.assert_close(cache[idx.long()].float(), ht_ref, rtol=5e-3, atol=5e-3) - - # The padding between rows, and every unowned row, must be untouched. - untouched = torch.ones(slots, dtype=torch.bool, device="cuda") -diff --git a/vllm/models/kimi_k3/amd/ops/kda_chunk.py b/vllm/models/kimi_k3/amd/ops/kda_chunk.py -index a3fbada90..1562b69fc 100644 ---- a/vllm/models/kimi_k3/amd/ops/kda_chunk.py -+++ b/vllm/models/kimi_k3/amd/ops/kda_chunk.py -@@ -235,7 +235,7 @@ def fused_kda_chunk( - output_final_state: whether to return the final recurrent state. - chunk_offsets: int32 per-sequence first chunk index. - checkpoint_state: destination for the mid-prefill state snapshots, -- fp32 ``[rows, H, 128, 128]``. Without -+ fp32 or bf16 ``[rows, H, 128, 128]``. Without - ``checkpoint_state_indices`` it is a staging buffer indexed by - sequence; with them it can be the paged state cache itself. - checkpoint_offsets: int32 ``[N]`` token offsets, each relative to its -@@ -246,7 +246,7 @@ def fused_kda_chunk( - rather than rely on the store happening. - checkpoint_state_indices: optional int32 ``[N]`` destination row per - sequence. A negative row disables the export for that sequence. -- state_cache: the paged recurrent state, fp32 -+ state_cache: the paged recurrent state, fp32 or bf16 - ``[slots, H, 128, 128]``. When given, the walk reads sequence - ``n``'s initial state from row ``state_indices[n]`` and writes its - final state back to the same row, which removes the gather and the -diff --git a/vllm/models/kimi_k3/amd/ops/kda_prefill.py b/vllm/models/kimi_k3/amd/ops/kda_prefill.py -index 18d90b811..e1900b5e5 100644 ---- a/vllm/models/kimi_k3/amd/ops/kda_prefill.py -+++ b/vllm/models/kimi_k3/amd/ops/kda_prefill.py -@@ -86,10 +86,9 @@ def chunk_kda_prefill( - checkpoint_offsets: per-sequence token offset to snapshot at, ``0`` - for none. - checkpoint_state_indices: optional per-sequence destination row. -- state_cache: the paged recurrent state. The fused backend addresses an -- FP32 cache directly. Other cache dtypes are gathered into the FP32 -- accumulator format and cast back after the walk. The returned final -- state is ``None`` in either case. -+ state_cache: the paged recurrent state. The fused backend reads and -+ writes FP32 or BF16 cache rows directly, accumulating in FP32. -+ The returned final state is ``None``. - state_indices: per-sequence cache row. - has_initial_state: per-sequence flag; false starts from a zero state. - -@@ -151,11 +150,10 @@ def chunk_kda_prefill( - has_initial_state=has_initial_state, - ) - -- # Only the fused walk addresses FP32 paged rows directly. The Triton path -- # and lower-precision persistent caches gather the active rows and scatter -- # the final state back here, so callers use one cache API for every backend. -+ # Only the fused walk addresses paged rows directly. The Triton path -+ # gathers the active rows and scatters the final state back here. - scatter_to: torch.Tensor | None = None -- if state_cache is not None and (not fused or state_cache.dtype != torch.float32): -+ if state_cache is not None and not fused: - initial_state = gather_initial_states( - state_cache, state_indices, has_initial_state - ) -@@ -166,9 +164,7 @@ def chunk_kda_prefill( - # Restated for the type checker; `fused` already implies all three. - assert cu_seqlens is not None and lower_bound is not None - assert g_bias is not None -- # The fused HIP prefill kernel accumulates recurrent state in FP32. -- # Convert only active rows; the persistent cache keeps its configured -- # dtype and receives a cast-back below. -+ # The dense state API remains FP32; paged caches bypass this conversion. - if initial_state is not None and initial_state.dtype != torch.float32: - initial_state = initial_state.to(torch.float32) - logger.info_once( -@@ -187,7 +183,7 @@ def chunk_kda_prefill( - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - ) -- o, final_state = fused_kda_chunk( -+ return fused_kda_chunk( - qg=ws["qg"], - w=ws["w"], - u=ws["u"], -@@ -214,12 +210,6 @@ def chunk_kda_prefill( - state_indices=state_indices, - has_initial_state=has_initial_state, - ) -- if scatter_to is not None: -- assert state_indices is not None -- assert final_state is not None -- scatter_to[state_indices.long()] = final_state.to(scatter_to.dtype) -- return o, None -- return o, final_state - - o, final_state = chunk_kda_with_fused_gate( - q=q, diff --git a/inferencex-e2e/benchmarks/multi_node/srt-slurm-recipes/configs/patches/kimik3-nightly/05-vllm-58743-backends.patch b/inferencex-e2e/benchmarks/multi_node/srt-slurm-recipes/configs/patches/kimik3-nightly/05-vllm-58743-backends.patch deleted file mode 100644 index bb862b66ad..0000000000 --- a/inferencex-e2e/benchmarks/multi_node/srt-slurm-recipes/configs/patches/kimik3-nightly/05-vllm-58743-backends.patch +++ /dev/null @@ -1,181 +0,0 @@ -From 983548e8aa73a29b1d8b5543ac83d231350440f8 Mon Sep 17 00:00:00 2001 -From: Yichao Zhu -Date: Mon, 19 Oct 2026 18:34:00 +0800 -Subject: [PATCH] Select compatible ROCm KDA backends for BF16 states - -Keep FP32-only Gluon and HIP decode paths out of BF16 recurrent-state execution. Preserve FP32 backend priority and validate BF16 decode state writes and backend selection. - -Signed-off-by: Yichao Zhu -(cherry picked from commit 149262f3ca7c3a07837797942342d45f0abe9f64) ---- - tests/models/kimi_k3/test_amd_kda_decode.py | 59 ++++++++++++++++++++- - vllm/models/kimi_k3/amd/kda.py | 14 +++-- - vllm/models/kimi_k3/amd/ops/kda_decode.py | 2 + - 3 files changed, 71 insertions(+), 4 deletions(-) - -diff --git a/tests/models/kimi_k3/test_amd_kda_decode.py b/tests/models/kimi_k3/test_amd_kda_decode.py -index fc6039ee5..8b12d31c8 100644 ---- a/tests/models/kimi_k3/test_amd_kda_decode.py -+++ b/tests/models/kimi_k3/test_amd_kda_decode.py -@@ -14,6 +14,8 @@ Every output of each is compared against the unfused path the AMD layer falls - back to, with the argument layout the layer itself uses. - """ - -+from functools import partial -+ - import pytest - import torch - -@@ -335,7 +337,11 @@ def test_fused_kda_decode_skips_null_block_padding() -> None: - - - @requires_hip_kernel_arch --def test_decode_weights_staged_after_loading(dist_init) -> None: -+@pytest.mark.parametrize("state_dtype", ["float32", "bfloat16"]) -+@pytest.mark.parametrize("aiter_available", [False, True]) -+def test_decode_weights_staged_after_loading( -+ dist_init, monkeypatch, state_dtype: str, aiter_available: bool -+) -> None: - """The fused-decode weight copies must not depend on the weight loaders. - - Loaders that bypass ``weight_loader`` (e.g. ``--load-format dummy``) left -@@ -344,11 +350,20 @@ def test_decode_weights_staged_after_loading(dist_init) -> None: - """ - _requires_kernel() - from tests.v1.attention.utils import create_vllm_config -+ from vllm._aiter_ops import rocm_aiter_ops - from vllm.config import set_current_vllm_config - from vllm.model_executor.model_loader.utils import process_weights_after_loading - from vllm.models.kimi_k3.amd.kda import KimiK3DeltaAttention -+ from vllm.models.kimi_k3.amd.ops.kda_chunk import is_fused_kda_chunk_supported - from vllm.transformers_utils.configs.kimi_linear import KimiLinearConfig - -+ monkeypatch.setenv("VLLM_SSM_CONV_STATE_LAYOUT", "SD") -+ monkeypatch.setattr( -+ rocm_aiter_ops, "is_kda_gluon_decode_available", lambda: aiter_available -+ ) -+ monkeypatch.setattr( -+ rocm_aiter_ops, "is_kda_gluon_prefill_available", lambda: aiter_available -+ ) - config = KimiLinearConfig( - hidden_size=256, - linear_attn_config={ -@@ -362,10 +377,26 @@ def test_decode_weights_staged_after_loading(dist_init) -> None: - }, - ) - vllm_config = create_vllm_config(model_name="Qwen/Qwen3.5-0.8B", dtype=DTYPE) -+ vllm_config.cache_config.mamba_ssm_cache_dtype = state_dtype - with set_current_vllm_config(vllm_config), torch.device("cuda"): - layer = KimiK3DeltaAttention( - config, vllm_config, prefix="model.layers.0.self_attn" - ) -+ use_aiter = aiter_available and state_dtype == "float32" -+ assert layer.use_aiter_decode == use_aiter -+ expected_prefill = ( -+ "gluon" -+ if use_aiter -+ else "fused" -+ if is_fused_kda_chunk_supported() -+ else "triton" -+ ) -+ assert layer.kda_prefill_backend == expected_prefill -+ if state_dtype == "bfloat16": -+ # Neither FP32-only decode backend may stage weights or take its fast path. -+ assert layer.decode_conv1d_weight is None -+ assert layer.decode_norm_weight is None -+ return - assert layer.decode_conv1d_weight is not None - # The AITER Gluon decode path takes o_norm.weight directly, so only the HIP - # kernel gets an fp32 norm copy. -@@ -625,6 +656,32 @@ def _run_aiter_spec( - return core_attn_out, conv_state, recurrent_state - - -+@pytest.mark.parametrize("num_spec", [0, 3]) -+@torch.inference_mode() -+def test_triton_decode_with_bf16_recurrent_state(num_spec: int) -> None: -+ """BF16 fallback reads rounded states and preserves paged writeback semantics.""" -+ if num_spec: -+ batch = SpecBatch([4, 2, 1], [1, 2, 4], 12, num_spec, seed=5) -+ inp = batch.inp -+ run = partial(_run_spec_triton_chain, batch) -+ else: -+ inp = KdaDecodeInputs(3, 12, num_slots=8, seed=5) -+ run = partial(_run_triton_chain, inp) -+ -+ initial_state = inp.recurrent_state.bfloat16() -+ inp.recurrent_state = initial_state.float() -+ expected_out, expected_conv, expected_state = run() -+ inp.recurrent_state = initial_state -+ actual_out, actual_conv, actual_state = run() -+ -+ assert actual_state.dtype == torch.bfloat16 -+ torch.testing.assert_close(actual_out, expected_out, atol=3e-2, rtol=3e-2) -+ torch.testing.assert_close(actual_conv, expected_conv, atol=0, rtol=0) -+ torch.testing.assert_close( -+ actual_state, expected_state.bfloat16(), atol=2e-3, rtol=3e-2 -+ ) -+ -+ - @requires_aiter_gluon_kda - @pytest.mark.parametrize("num_heads", [12, 24]) - @pytest.mark.parametrize( -diff --git a/vllm/models/kimi_k3/amd/kda.py b/vllm/models/kimi_k3/amd/kda.py -index 4216c6f94..8460c09f7 100644 ---- a/vllm/models/kimi_k3/amd/kda.py -+++ b/vllm/models/kimi_k3/amd/kda.py -@@ -194,8 +194,11 @@ class KimiK3DeltaAttention(GatedDeltaNetAttention): - # ROCm can fuse the whole decode step (conv + recurrence + gated norm) - # into one kernel, which wants a width-major fp32 conv weight staged at - # load time. Everything else keeps the [channel, width] layout. -- conv_state_dtype, _ = self.get_state_dtype() -- self.use_aiter_decode = bool(rocm_aiter_ops.is_kda_gluon_decode_available()) -+ conv_state_dtype, recurrent_state_dtype = self.get_state_dtype() -+ # AITER's Gluon kernels currently require FP32 recurrent states. -+ self.use_aiter_decode = recurrent_state_dtype == torch.float32 and bool( -+ rocm_aiter_ops.is_kda_gluon_decode_available() -+ ) - use_hip_decode = not self.use_aiter_decode and is_fused_kda_decode_supported( - self.local_num_heads, - self.head_dim, -@@ -203,6 +206,7 @@ class KimiK3DeltaAttention(GatedDeltaNetAttention): - self.num_spec, - vllm_config.model_config.dtype, - conv_state_dtype, -+ recurrent_state_dtype, - ) - if use_hip_decode: - logger.info_once("Fused KDA decode kernel (conv+KDA+norm) is enabled.") -@@ -256,7 +260,11 @@ class KimiK3DeltaAttention(GatedDeltaNetAttention): - ) - self.kda_prefill_backend = backend - if backend == "auto": -- if self.use_safe_gate and rocm_aiter_ops.is_kda_gluon_prefill_available(): -+ if ( -+ recurrent_state_dtype == torch.float32 -+ and self.use_safe_gate -+ and rocm_aiter_ops.is_kda_gluon_prefill_available() -+ ): - self.kda_prefill_backend = "gluon" - elif is_fused_kda_chunk_supported(): - self.kda_prefill_backend = "fused" -diff --git a/vllm/models/kimi_k3/amd/ops/kda_decode.py b/vllm/models/kimi_k3/amd/ops/kda_decode.py -index d944e1eb0..236a95d72 100644 ---- a/vllm/models/kimi_k3/amd/ops/kda_decode.py -+++ b/vllm/models/kimi_k3/amd/ops/kda_decode.py -@@ -30,6 +30,7 @@ def is_fused_kda_decode_supported( - num_spec: int, - input_dtype: torch.dtype, - conv_state_dtype: torch.dtype, -+ recurrent_state_dtype: torch.dtype, - ) -> bool: - """Whether the fused decode kernel can serve this layer on this device.""" - from vllm.platforms.rocm import on_gfx942, on_gfx950 -@@ -41,6 +42,7 @@ def is_fused_kda_decode_supported( - or num_spec != 0 - or input_dtype != torch.bfloat16 - or conv_state_dtype != torch.bfloat16 -+ or recurrent_state_dtype != torch.float32 - or is_conv_state_dim_first() - or not hasattr(torch.ops._C, "fused_kda_decode") - ): diff --git a/inferencex-e2e/benchmarks/multi_node/srt-slurm-recipes/configs/patches/kimik3-nightly/06-vllm-58743-tests.patch b/inferencex-e2e/benchmarks/multi_node/srt-slurm-recipes/configs/patches/kimik3-nightly/06-vllm-58743-tests.patch deleted file mode 100644 index 330ad5bf00..0000000000 --- a/inferencex-e2e/benchmarks/multi_node/srt-slurm-recipes/configs/patches/kimik3-nightly/06-vllm-58743-tests.patch +++ /dev/null @@ -1,115 +0,0 @@ -From a02a204017e4e0ef88e2ae0d1eac76e8ac1b9332 Mon Sep 17 00:00:00 2001 -From: Yichao Zhu -Date: Mon, 19 Oct 2026 19:30:48 +0800 -Subject: [PATCH] Keep the KDA dense state contract and consolidate BF16 tests - -Remove the implicit dense initial-state conversion now that HIP prefill handles BF16 paged states natively. Extend the existing paged-state comparison instead of maintaining a separate BF16 case. - -Signed-off-by: Yichao Zhu -(cherry picked from commit ae6bb332600412fdf95fec94548e091a7638be99) ---- - tests/models/kimi_k3/test_amd_kda_chunk.py | 36 ++++++---------------- - vllm/models/kimi_k3/amd/ops/kda_prefill.py | 8 ++--- - 2 files changed, 13 insertions(+), 31 deletions(-) - -diff --git a/tests/models/kimi_k3/test_amd_kda_chunk.py b/tests/models/kimi_k3/test_amd_kda_chunk.py -index c7a2c511f..6865dea65 100644 ---- a/tests/models/kimi_k3/test_amd_kda_chunk.py -+++ b/tests/models/kimi_k3/test_amd_kda_chunk.py -@@ -687,15 +687,19 @@ def _run_paged(inp: dict, cache, idx, warm, use_fused: bool, **kw): - - @pytest.mark.parametrize("use_fused", [False, True]) - @pytest.mark.parametrize("seqlens", [[512], [2048], [320, 64], [513, 64, 1]]) -+@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) - def test_paged_state_matches_the_dense_pair( -- use_fused: bool, seqlens: list[int] -+ use_fused: bool, seqlens: list[int], dtype: torch.dtype - ) -> None: - """A paged cache must give the same result as the dense pair.""" - _requires_kernel() - inp = _inputs(seqlens, seed=11) -+ inp["h0"] = inp["h0"].to(dtype).float() - n = len(seqlens) - rows = [3 * i + 5 for i in range(n)] -- cache, idx, warm = _paged_cache(inp, [True] * n, rows, slots=3 * n + 12) -+ cache, idx, warm = _paged_cache( -+ inp, [True] * n, rows, slots=3 * n + 12, dtype=dtype -+ ) - before = cache.clone() - - o_ref, ht_ref = _run(inp, use_fused=use_fused) -@@ -703,7 +707,9 @@ def test_paged_state_matches_the_dense_pair( - - assert ht_got is None, "the final state is in the cache, not returned" - torch.testing.assert_close(o_got, o_ref, rtol=5e-3, atol=5e-3) -- torch.testing.assert_close(cache[idx.long()], ht_ref, rtol=5e-3, atol=5e-3) -+ torch.testing.assert_close( -+ cache[idx.long()], ht_ref.to(dtype), rtol=5e-3, atol=5e-3 -+ ) - - untouched = torch.ones(cache.shape[0], dtype=torch.bool, device="cuda") - untouched[idx.long()] = False -@@ -712,29 +718,6 @@ def test_paged_state_matches_the_dense_pair( - ) - - --@pytest.mark.parametrize("use_fused", [False, True]) --def test_paged_bf16_state_matches_dense_pair(use_fused: bool) -> None: -- """The FP32 walk may load and store a BF16 persistent cache.""" -- if use_fused: -- _requires_kernel() -- seqlens = [513, 64, 1] -- inp = _inputs(seqlens, seed=17) -- rows = [2, 7, 11] -- cache, idx, warm = _paged_cache(inp, [True] * len(seqlens), rows, slots=16) -- cache = cache.to(torch.bfloat16) -- -- dense = dict(inp) -- dense["h0"] = cache[idx.long()].float() -- o_ref, ht_ref = _run(dense, use_fused=use_fused) -- o_got, ht_got = _run_paged(inp, cache, idx, warm, use_fused=use_fused) -- -- assert ht_got is None -- torch.testing.assert_close(o_got, o_ref, rtol=5e-3, atol=5e-3) -- torch.testing.assert_close( -- cache[idx.long()], ht_ref.to(torch.bfloat16), rtol=5e-3, atol=5e-3 -- ) -- -- - @pytest.mark.parametrize("use_fused", [False, True]) - @pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) - def test_cold_rows_start_from_zero_not_from_cache_junk( -@@ -966,6 +949,7 @@ def test_padded_cache_stride_is_honoured(use_fused: bool, dtype: torch.dtype) -> - inp["h0"] = inp["h0"].to(dtype).float() - n = len(seqlens) - row = NUM_HEADS * HEAD_DIM * HEAD_DIM -+ # An odd BF16 row stride exercises unaligned vector loads and stores. - slots, pad = 16, 513 if dtype == torch.bfloat16 else 512 - raw = torch.full((slots, row + pad), -7.0, device="cuda", dtype=dtype) - cache = raw.as_strided( -diff --git a/vllm/models/kimi_k3/amd/ops/kda_prefill.py b/vllm/models/kimi_k3/amd/ops/kda_prefill.py -index e1900b5e5..ecdb36f4b 100644 ---- a/vllm/models/kimi_k3/amd/ops/kda_prefill.py -+++ b/vllm/models/kimi_k3/amd/ops/kda_prefill.py -@@ -150,8 +150,9 @@ def chunk_kda_prefill( - has_initial_state=has_initial_state, - ) - -- # Only the fused walk addresses paged rows directly. The Triton path -- # gathers the active rows and scatters the final state back here. -+ # Only the fused walk addresses the paged rows directly. The Triton path -+ # gathers the rows it needs and scatters the results back here, so callers -+ # pass the cache the same way for either backend. - scatter_to: torch.Tensor | None = None - if state_cache is not None and not fused: - initial_state = gather_initial_states( -@@ -164,9 +165,6 @@ def chunk_kda_prefill( - # Restated for the type checker; `fused` already implies all three. - assert cu_seqlens is not None and lower_bound is not None - assert g_bias is not None -- # The dense state API remains FP32; paged caches bypass this conversion. -- if initial_state is not None and initial_state.dtype != torch.float32: -- initial_state = initial_state.to(torch.float32) - logger.info_once( - "Kimi-K3 KDA prefill: dispatching the fused ROCm chunk kernel." - ) diff --git a/inferencex-e2e/benchmarks/multi_node/srt-slurm-recipes/configs/patches/kimik3-nightly/07-vllm-60808.patch b/inferencex-e2e/benchmarks/multi_node/srt-slurm-recipes/configs/patches/kimik3-nightly/07-vllm-60808.patch deleted file mode 100644 index c3d485cbee..0000000000 --- a/inferencex-e2e/benchmarks/multi_node/srt-slurm-recipes/configs/patches/kimik3-nightly/07-vllm-60808.patch +++ /dev/null @@ -1,569 +0,0 @@ -From a0ef2bdc9879a40cd1d6e6566e505622194fdf32 Mon Sep 17 00:00:00 2001 -From: Mir Mustafa Ali -Date: Wed, 21 Oct 2026 07:11:16 -0500 -Subject: [PATCH] [Bugfix][KV Connector][ROCm] MoRIIO READ: await every read - before the forward (#60808) - -Signed-off-by: Mir Mustafa Ali -Co-authored-by: Claude -(cherry picked from commit db3b0f0fae4bc98492b3863685e541e62bc59576) ---- - .../kv_connector/unit/test_moriio_tp_ack.py | 278 +++++++++++++++++- - .../v1/moriio/moriio_connector.py | 116 ++------ - .../kv_connector/v1/moriio/moriio_engine.py | 21 +- - 3 files changed, 316 insertions(+), 99 deletions(-) - -diff --git a/tests/v1/kv_connector/unit/test_moriio_tp_ack.py b/tests/v1/kv_connector/unit/test_moriio_tp_ack.py -index 34d193e27..72453f743 100644 ---- a/tests/v1/kv_connector/unit/test_moriio_tp_ack.py -+++ b/tests/v1/kv_connector/unit/test_moriio_tp_ack.py -@@ -3,9 +3,11 @@ - - import threading - from types import SimpleNamespace -+from unittest.mock import Mock - - import pytest - -+from vllm.distributed.kv_transfer.kv_connector.v1.moriio import moriio_engine - from vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_common import ( - MoRIIOMode, - MoRIIOTransferAck, -@@ -460,9 +462,9 @@ def test_failed_read_reports_blocks_only_without_hma(has_mamba, expected_invalid - assert worker.get_block_ids_with_load_errors() == expected_invalid - - --def test_hybrid_step_barrier_fails_closed(monkeypatch): -+def test_hybrid_step_barrier_fails_closed(): - class FailingWrapper: -- def waiting_for_transfer_complete(self, _statuses): -+ def waiting_for_transfer_complete(self, _statuses, fail_fast=True): - raise TransferError("failed") - - def shutdown(self): -@@ -470,18 +472,278 @@ def test_hybrid_step_barrier_fails_closed(monkeypatch): - - worker = MoRIIOConnectorWorker.__new__(MoRIIOConnectorWorker) - worker._has_mamba = True -- worker._reads_issued_this_step = [object()] -- worker._mamba_reads_this_step = [object()] -+ worker._reads_issued_this_step = [_ReadStatus("failed")] - worker.moriio_wrapper = FailingWrapper() -- monkeypatch.setattr( -- "vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_connector.get_forward_context", -- lambda: SimpleNamespace(cudagraph_runtime_mode=None), -- ) - - with pytest.raises(TransferError, match="failed"): - worker._await_reads_issued_this_step() - - -+def test_step_barrier_awaits_every_read_before_forward(): -+ class RecordingWrapper: -+ def __init__(self): -+ self.awaited = [] -+ -+ def waiting_for_transfer_complete(self, statuses, fail_fast=True): -+ self.awaited.append(list(statuses)) -+ -+ def shutdown(self): -+ pass -+ -+ attn, indexer, mtp_attn = object(), object(), object() -+ worker = MoRIIOConnectorWorker.__new__(MoRIIOConnectorWorker) -+ worker._has_mamba = False -+ worker._reads_issued_this_step = [attn, indexer, mtp_attn] -+ worker.moriio_wrapper = RecordingWrapper() -+ -+ worker._await_reads_issued_this_step() -+ worker._await_reads_issued_this_step() -+ -+ assert worker.moriio_wrapper.awaited == [[attn, indexer, mtp_attn]] -+ assert worker._reads_issued_this_step == [] -+ -+ -+def test_wait_for_layer_load_is_noop(): -+ class _Boom: -+ def __getattr__(self, name): -+ raise AssertionError("wait_for_layer_load must not touch the worker") -+ -+ connector = MoRIIOConnector.__new__(MoRIIOConnector) -+ connector.connector_worker = _Boom() -+ -+ assert connector.wait_for_layer_load("layer") is None -+ -+ -+def test_step_barrier_without_mamba_does_not_raise_on_failure(): -+ # Without HMA, failures go to get_block_ids_with_load_errors. -+ class FailingWrapper: -+ def waiting_for_transfer_complete(self, _statuses, fail_fast=True): -+ raise TransferError("failed") -+ -+ def shutdown(self): -+ pass -+ -+ worker = MoRIIOConnectorWorker.__new__(MoRIIOConnectorWorker) -+ worker._has_mamba = False -+ worker._reads_issued_this_step = [_ReadStatus("failed")] -+ worker.moriio_wrapper = FailingWrapper() -+ -+ worker._await_reads_issued_this_step() -+ assert worker._reads_issued_this_step == [] -+ -+ -+class _ReadStatus: -+ def __init__(self, state: str): -+ self.state = state -+ -+ def Succeeded(self): -+ return self.state == "success" -+ -+ def Failed(self): -+ return self.state == "failed" -+ -+ def Message(self): -+ return "SQ full" -+ -+ def Code(self): -+ return -1 -+ -+ -+class _PolledStatus(_ReadStatus): -+ """Lands on its succeed_after-th Succeeded() poll; never if None.""" -+ -+ def __init__(self, succeed_after: int | None = None): -+ super().__init__("in_progress") -+ self.succeed_after = succeed_after -+ self.polls = 0 -+ -+ def Succeeded(self): -+ self.polls += 1 -+ if self.polls == self.succeed_after: -+ self.state = "success" -+ return super().Succeeded() -+ -+ -+class _FakeClock: -+ def __init__(self): -+ self.now = 0.0 -+ self.sleeps = 0 -+ -+ def monotonic(self): -+ return self.now -+ -+ def sleep(self, seconds): -+ self.sleeps += 1 -+ self.now += seconds -+ -+ -+@pytest.fixture -+def fake_clock(monkeypatch): -+ clock = _FakeClock() -+ monkeypatch.setattr(moriio_engine, "time", clock) -+ return clock -+ -+ -+def _worker_with_failed_and_pending_read(has_mamba, pending_b, monkeypatch): -+ failed_a = _ReadStatus("failed") -+ worker = MoRIIOConnectorWorker.__new__(MoRIIOConnectorWorker) -+ worker._has_mamba = has_mamba -+ worker.world_size = 1 -+ worker.moriio_config = SimpleNamespace(recv_abort_timeout=600.0) -+ worker.moriio_wrapper = MoRIIOWrapper(moriio_engine=object(), transfer_timeout=5.0) -+ monkeypatch.setattr(worker.moriio_wrapper, "send_notify", lambda *a, **kw: None) -+ worker._reads_issued_this_step = [failed_a, pending_b] -+ worker._recving_transfers = { -+ "a": {"layer0": [failed_a]}, -+ "b": {"layer0": [pending_b]}, -+ } -+ worker._recving_transfers_callback_addr = { -+ "a": ("127.0.0.1", "7000", "tx-a"), -+ "b": ("127.0.0.1", "7000", "tx-b"), -+ } -+ worker._recving_transfers_start = {} -+ worker._recving_local_blocks = {"a": [1, 2], "b": [3, 4]} -+ worker._invalid_block_ids = set() -+ return worker -+ -+ -+def test_step_barrier_failure_still_awaits_other_requests(monkeypatch, fake_clock): -+ late_b = _PolledStatus(succeed_after=5) -+ worker = _worker_with_failed_and_pending_read(False, late_b, monkeypatch) -+ -+ worker._await_reads_issued_this_step() -+ -+ assert late_b.state == "success" -+ assert late_b.polls >= 5 -+ assert fake_clock.sleeps >= 4 -+ assert worker._pop_done_transfers() == {"tx-b"} -+ assert worker.get_block_ids_with_load_errors() == {1, 2} -+ -+ -+def test_hybrid_step_barrier_failure_does_not_wait(monkeypatch, fake_clock): -+ pending_b = _PolledStatus() -+ worker = _worker_with_failed_and_pending_read(True, pending_b, monkeypatch) -+ -+ with pytest.raises(TransferError, match="SQ full"): -+ worker._await_reads_issued_this_step() -+ -+ assert pending_b.state == "in_progress" -+ assert fake_clock.sleeps == 0 -+ -+ -+@pytest.mark.parametrize("has_mamba", [False, True]) -+def test_step_barrier_timeout_fails_closed_while_reads_in_flight(has_mamba, fake_clock): -+ worker = MoRIIOConnectorWorker.__new__(MoRIIOConnectorWorker) -+ worker._has_mamba = has_mamba -+ worker.moriio_wrapper = MoRIIOWrapper(moriio_engine=object(), transfer_timeout=0.05) -+ worker._reads_issued_this_step = [_PolledStatus(), _PolledStatus()] -+ -+ with pytest.raises(TransferError, match="timed out"): -+ worker._await_reads_issued_this_step() -+ -+ assert fake_clock.now > 0.05 -+ -+ -+@pytest.mark.skipif( -+ moriio_engine.StatusCode is None, reason="mori batched wait unavailable" -+) -+@pytest.mark.parametrize("has_mamba", [False, True]) -+def test_step_barrier_native_timeout_fails_closed_while_reads_in_flight(has_mamba): -+ engine = SimpleNamespace( -+ wait_all=Mock(return_value=moriio_engine.StatusCode.IN_PROGRESS) -+ ) -+ worker = MoRIIOConnectorWorker.__new__(MoRIIOConnectorWorker) -+ worker._has_mamba = has_mamba -+ worker.moriio_wrapper = MoRIIOWrapper(moriio_engine=engine, transfer_timeout=0.05) -+ worker._reads_issued_this_step = [_PolledStatus(), _PolledStatus()] -+ -+ with pytest.raises(TransferError, match="timed out"): -+ worker._await_reads_issued_this_step() -+ -+ engine.wait_all.assert_called_once() -+ -+ -+def test_step_barrier_noop_when_no_reads_posted(): -+ class RecordingWrapper: -+ def __init__(self): -+ self.called = False -+ -+ def waiting_for_transfer_complete(self, _statuses): -+ self.called = True -+ -+ def shutdown(self): -+ pass -+ -+ worker = MoRIIOConnectorWorker.__new__(MoRIIOConnectorWorker) -+ worker._has_mamba = False -+ worker._reads_issued_this_step = [] -+ worker.moriio_wrapper = RecordingWrapper() -+ -+ worker._await_reads_issued_this_step() -+ -+ assert worker.moriio_wrapper.called is False -+ -+ -+def test_write_mode_read_blocks_posts_no_read(): -+ class RecordingWrapper: -+ def __init__(self): -+ self.sent = [] -+ -+ def send_notify(self, *args, **kwargs): -+ self.sent.append((args, kwargs)) -+ -+ def shutdown(self): -+ pass -+ -+ worker = MoRIIOConnectorWorker.__new__(MoRIIOConnectorWorker) -+ worker.mode = MoRIIOMode.WRITE -+ worker._reads_issued_this_step = [] -+ worker.moriio_wrapper = RecordingWrapper() -+ -+ worker._read_blocks( -+ local_block_ids=[1, 2], -+ remote_block_ids=[[10, 11], [90]], -+ dst_engine_id="prefill", -+ request_id="req-write", -+ transfer_id="tx-write", -+ remote_host="127.0.0.1", -+ remote_notify_port=7000, -+ remote_tp_size=8, -+ ) -+ -+ assert worker._reads_issued_this_step == [] -+ assert worker.moriio_wrapper.sent == [] -+ -+ -+def test_producer_start_load_kv_posts_no_read(): -+ class RecordingWrapper: -+ def __init__(self): -+ self.waited = False -+ self.async_waited = False -+ -+ def async_wait_reqid(self): -+ self.async_waited = True -+ -+ def waiting_for_transfer_complete(self, _statuses): -+ self.waited = True -+ -+ def shutdown(self): -+ pass -+ -+ worker = MoRIIOConnectorWorker.__new__(MoRIIOConnectorWorker) -+ worker.is_producer = True -+ worker._reads_issued_this_step = [] -+ worker._consumer_notification_counts = {} -+ worker._completed_consumer_notifications = set() -+ worker.moriio_wrapper = RecordingWrapper() -+ -+ worker.start_load_kv(SimpleNamespace(transfer_id_to_request_id={})) -+ -+ assert worker._reads_issued_this_step == [] -+ assert worker.moriio_wrapper.async_waited is True -+ assert worker.moriio_wrapper.waited is False -+ -+ - def test_requested_cudagraph_mode_is_never_overridden(): - # The configured cudagraph mode is always honored: the barrier fires when - # the operator sets cudagraph_mode=PIECEWISE, and READ mode with full -diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py -index 4f19d73d9..d6039e09e 100644 ---- a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py -+++ b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py -@@ -17,7 +17,7 @@ import numpy as np - import torch - import zmq - --from vllm.config import CUDAGraphMode, KVTransferConfig, VllmConfig -+from vllm.config import KVTransferConfig, VllmConfig - from vllm.distributed.kv_transfer.kv_connector.utils import ( - BlockIds, - clip_ssm_state_blocks, -@@ -97,7 +97,7 @@ from vllm.distributed.parallel_state import ( - get_tp_group, - get_world_group, - ) --from vllm.forward_context import ForwardContext, get_forward_context -+from vllm.forward_context import ForwardContext - from vllm.logger import init_logger - from vllm.utils.math_utils import cdiv - from vllm.utils.network_utils import ( -@@ -416,8 +416,8 @@ class MoRIIOConnector(KVConnectorBase_V1, SupportsHMA): - self.connector_worker.start_load_kv(self._connector_metadata) - - def wait_for_layer_load(self, layer_name: str) -> None: -- assert self.connector_worker is not None -- self.connector_worker.wait_for_layer_load(layer_name) -+ """READs are awaited in start_load_kv, before the forward.""" -+ pass - - def save_kv_layer( - self, -@@ -833,8 +833,8 @@ class MoRIIOConnectorScheduler: - # - # The second element declares who waits for the KV, not whether it has - # arrived. False keeps the wait on this side: the forward runs in the -- # same step, and the connector blocks in wait_for_layer_load (or in -- # start_load_kv under full-graph capture) until the KV is really there. -+ # same step, and the connector blocks in start_load_kv, before the -+ # forward, until the KV is really there. - return max(num_external_tokens - 1, 0), False - - def on_new_request(self, request: "Request") -> None: -@@ -1907,13 +1907,8 @@ class MoRIIOConnectorWorker: - self._recving_local_blocks: dict[ReqId, list[int]] = {} - # Drained by get_block_ids_with_load_errors. - self._invalid_block_ids: set[int] = set() -- # This step's posted reads, awaited before the forward when the -- # per-layer barrier cannot run. -+ # This step's posted reads, awaited in start_load_kv before the forward. - self._reads_issued_this_step: list = [] -- # This step's posted recurrent-state reads. Mamba layers are not -- # wrapped by the attention KV-transfer decorator, so wait_for_layer_load -- # never fires for them and these must be awaited in start_load_kv. -- self._mamba_reads_this_step: list = [] - - # Track the expiration time of requests that are waiting to be sent. - self._reqs_to_send: dict[ReqId, float] = {} -@@ -2900,44 +2895,6 @@ class MoRIIOConnectorWorker: - failed_recving=failed_recving, - ) - -- def wait_for_layer_load(self, layer_name: str) -> None: -- """Block until all in-flight READs of this layer have landed. -- -- A host-side blocking wait must not run during full-graph capture. -- """ -- if self.is_producer or self.mode != MoRIIOMode.READ: -- return -- -- if get_forward_context().cudagraph_runtime_mode == CUDAGraphMode.FULL: -- return -- -- with self.moriio_wrapper.lock: -- pending = [ -- status -- for status_by_layer in self._recving_transfers.values() -- for status in status_by_layer.get(layer_name, ()) -- ] -- -- if not pending: -- return -- -- try: -- self.moriio_wrapper.waiting_for_transfer_complete(pending) -- except TransferError: -- if self._has_mamba: -- logger.exception( -- "MoRIIO hybrid READ failed before layer %s; aborting the " -- "forward because HMA load-error recovery is unsupported.", -- layer_name, -- ) -- raise -- logger.warning( -- "MoRIIO READ barrier did not complete for layer %s; proceeding " -- "to transfer cleanup.", -- layer_name, -- exc_info=True, -- ) -- - def _pop_done_transfers(self) -> set[str]: - done_req_ids: set[str] = set() - _xfer_timeout = self.moriio_config.recv_abort_timeout -@@ -3020,40 +2977,37 @@ class MoRIIOConnectorWorker: - return done_req_ids - - def _await_reads_issued_this_step(self) -> None: -- """Step-level READ barrier, for the reads no per-layer barrier covers. -- -- load_kv_async=False promises the KV is in place before the forward runs. -- wait_for_layer_load keeps that promise per layer, letting early layers -- overlap later transfers, but it only fires for layers that go through -- the attention KV-transfer decorator. Two cases fall through it: -- -- * mamba/KDA layers have no such hook at all, so their conv+ssm reads -- would otherwise be consumed while still in flight; -- * under CUDAGraphMode.FULL a host-side blocking wait is illegal -- anywhere in the graph, so every read is awaited here instead. -- -- start_load_kv runs outside the graph, so blocking here is safe. -- Attention reads keep their per-layer overlap in the non-FULL case. -- -- Attention-only failures are reported via -- get_block_ids_with_load_errors. Hybrid failures are re-raised because -- the scheduler's recovery path does not support multiple cache groups; -- this fails the step closed instead of running forward with incomplete -- recurrent state. -+ """Block until every READ posted this step has landed. -+ -+ load_kv_async=False promises the KV is in place before the forward. A -+ per-layer barrier can't keep it: illegal under FULL cudagraph, and absent -+ for mamba/KDA state and indexer caches. -+ -+ A READ still in flight when the wait gives up fails the step in both -+ modes: RDMA may still write into its blocks, so the forward must not -+ consume them. If every READ is terminal, attention-only failures are -+ reported via get_block_ids_with_load_errors. Hybrid failures are -+ re-raised because the scheduler's recovery path does not support -+ multiple cache groups. - """ -- all_statuses = self._reads_issued_this_step -+ statuses = self._reads_issued_this_step - self._reads_issued_this_step = [] -- mamba_statuses = self._mamba_reads_this_step -- self._mamba_reads_this_step = [] -- if get_forward_context().cudagraph_runtime_mode == CUDAGraphMode.FULL: -- statuses = all_statuses -- else: -- statuses = mamba_statuses - if not statuses: - return - try: -- self.moriio_wrapper.waiting_for_transfer_complete(statuses) -+ self.moriio_wrapper.waiting_for_transfer_complete( -+ statuses, fail_fast=self._has_mamba -+ ) - except TransferError: -+ in_flight = sum(not (s.Succeeded() or s.Failed()) for s in statuses) -+ if in_flight: -+ logger.exception( -+ "MoRIIO READ barrier failed with %d/%d reads still in " -+ "flight; failing the step.", -+ in_flight, -+ len(statuses), -+ ) -+ raise - if self._has_mamba: - logger.exception( - "MoRIIO hybrid READ failed before the forward; aborting " -@@ -4065,13 +4019,7 @@ class MoRIIOConnectorWorker: - self._recving_local_blocks.setdefault( - request_id, list(failed_recv_blocks) - ) -- # Awaited before the forward when the per-layer barrier cannot -- # run (see _await_reads_issued_this_step). - self._reads_issued_this_step.extend(statuses) -- if self._is_mamba_layer(layer_name): -- # No per-layer barrier exists for mamba layers, so these are -- # always awaited in start_load_kv. -- self._mamba_reads_this_step.extend(statuses) - self._recving_transfers_callback_addr[request_id] = ( - remote_host, - str( -diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_engine.py b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_engine.py -index 9929341b3..61e2c8d38 100644 ---- a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_engine.py -+++ b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_engine.py -@@ -893,7 +893,9 @@ class MoRIIOWrapper: - state = TransferBatchState.PENDING - return state - -- def waiting_for_transfer_complete(self, transfer_statuses: list[Any] | None = None): -+ def waiting_for_transfer_complete( -+ self, transfer_statuses: list[Any] | None = None, fail_fast: bool = True -+ ): - if transfer_statuses is None: - with self.lock: - transfers_to_wait = self.transfer_status[:] -@@ -906,9 +908,11 @@ class MoRIIOWrapper: - - wait_all = getattr(self.moriio_engine, "wait_all", None) - if wait_all is None or StatusCode is None: -- self._poll_transfers_until_done(transfers_to_wait) -+ self._poll_transfers_until_done(transfers_to_wait, fail_fast) - return - -+ # wait_all waits on every status until the deadline; a failure does not -+ # end it early, so fail_fast only applies to the polling fallback. - # Zero means a nonblocking progress call in Mori, not an immediate - # deadline. Round up so even sub-millisecond waits remain blocking. - timeout_ms = ceil(self._transfer_timeout * 1000) -@@ -927,12 +931,15 @@ class MoRIIOWrapper: - f"MoRIIO batch wait failed ({result}): " + "; ".join(details) - ) - -- def _poll_transfers_until_done(self, transfers_to_wait: list[Any]) -> None: -+ def _poll_transfers_until_done( -+ self, transfers_to_wait: list[Any], fail_fast: bool = True -+ ) -> None: - """Fallback for mori builds without the batched wait. - -- Gives failure precedence over pending: once any status has failed the -- request is lost, so waiting out the remaining ones would only delay the -- barrier by up to transfer_timeout. -+ With fail_fast, failure takes precedence over pending: once any status -+ has failed the request is lost, so waiting out the remaining ones would -+ only delay the barrier by up to transfer_timeout. Pass False when the -+ statuses span several requests. - """ - timeout = self._transfer_timeout - deadline = time.monotonic() + timeout -@@ -959,7 +966,7 @@ class MoRIIOWrapper: - ) - else: - still_waiting.append(status) -- if errors: -+ if errors and fail_fast: - break - remaining = still_waiting - if remaining: diff --git a/inferencex-e2e/benchmarks/multi_node/srt-slurm-recipes/configs/patches/kimik3-nightly/08-vllm-63031.patch b/inferencex-e2e/benchmarks/multi_node/srt-slurm-recipes/configs/patches/kimik3-nightly/08-vllm-63031.patch deleted file mode 100644 index 3a69f6047c..0000000000 --- a/inferencex-e2e/benchmarks/multi_node/srt-slurm-recipes/configs/patches/kimik3-nightly/08-vllm-63031.patch +++ /dev/null @@ -1,172 +0,0 @@ -From fcd71bb9a36adb63b9ce9f43c7e4170ffadb595c Mon Sep 17 00:00:00 2001 -From: Yichao Zhu -Date: Thu, 22 Oct 2026 12:08:07 +0800 -Subject: [PATCH] [ROCm][Kimi-K3] Support MoRIIO consumers with specialized - decode - -Allow the built-in MoRIIO consumer in READ and WRITE modes. Reuse the existing transfer completion guarantees, and keep other connectors and producer-side specialized execution unsupported. Add configuration coverage and document the supported deployment. - -Signed-off-by: Yichao Zhu -(cherry picked from commit e34bf5f4784569ca3c2381312f2e7e4d512434de) ---- - docs/features/moriio_connector_usage.md | 8 ++++ - .../amd_specialized_decode/test_runtime.py | 46 ++++++++++++++++++- - .../kimi_k3/amd/specialized_decode/README.md | 9 +++- - .../specialized_decode/runtime/__init__.py | 16 ++++++- - 4 files changed, 75 insertions(+), 4 deletions(-) - -diff --git a/docs/features/moriio_connector_usage.md b/docs/features/moriio_connector_usage.md -index 2fe1ffa63..7da33d6a4 100644 ---- a/docs/features/moriio_connector_usage.md -+++ b/docs/features/moriio_connector_usage.md -@@ -144,6 +144,14 @@ For DSpark, configure the same draft model on both instances through `--speculat - - Hybrid READ keeps its existing single-transferable-attention-group requirement. Plain-attention WRITE retains its existing notification format. - -+### Kimi-K3 specialized decode -+ -+Kimi-K3 on MI355X supports `--additional-config '{"decode_backend":"latency_specialized"}'` on a MoRIIO READ or WRITE consumer. Use `VLLM_USE_V2_MODEL_RUNNER=1`, TP8, DCP1 or DCP8, DSpark with up to seven speculative tokens, and BF16 convolution / FP32 recurrent state on both instances. Keep the producer on the default decode backend and use the same transfer mode and cache layouts on both instances. -+ -+The specialized kernels serve pure speculative verification steps of at most eight token rows. For example, with seven speculative tokens and `--max-num-seqs 1`, a verification step has eight rows; use `--compilation-config '{"cudagraph_mode":"FULL_DECODE_ONLY","cudagraph_capture_sizes":[1,8]}'` on the consumer. Prefill, mixed, and wider steps use the regular model layers. All transferred attention and recurrent states are ready before the forward, so FULL graph replay does not need per-layer transfer hooks. -+ -+Producers using the specialized backend and other connectors remain unsupported. -+ - ### Different prefill and decode TP sizes - - In WRITE mode, the prefill and decode instances can use different tensor-parallel sizes, as long as one is a multiple of the other. Each prefill rank writes its shard into the decode ranks that hold it: -diff --git a/tests/models/kimi_k3/amd_specialized_decode/test_runtime.py b/tests/models/kimi_k3/amd_specialized_decode/test_runtime.py -index 6cc5c793c..5fde9359c 100644 ---- a/tests/models/kimi_k3/amd_specialized_decode/test_runtime.py -+++ b/tests/models/kimi_k3/amd_specialized_decode/test_runtime.py -@@ -8,7 +8,7 @@ from types import SimpleNamespace - import pytest - import torch - --from vllm.config import CompilationMode, CUDAGraphMode -+from vllm.config import CompilationMode, CUDAGraphMode, KVTransferConfig - from vllm.models.kimi_k3.amd.specialized_decode import runtime - from vllm.models.kimi_k3.amd.specialized_decode.runtime import ( - FFN, -@@ -141,6 +141,41 @@ def test_supported_config_is_accepted(on_mi355x): - assert specialized_decode_config(_vllm_config()) == DecodeConfig(q=8) - - -+@pytest.mark.parametrize("read_mode", [False, True]) -+@pytest.mark.parametrize("backend", ["rdma", "xgmi"]) -+def test_moriio_consumer_is_supported(read_mode, backend): -+ config = KVTransferConfig( -+ kv_connector="MoRIIOConnector", -+ kv_role="kv_consumer", -+ kv_connector_extra_config={"read_mode": read_mode, "backend": backend}, -+ ) -+ assert runtime.supports_kv_transfer(config) -+ -+ -+@pytest.mark.parametrize( -+ "overrides", -+ [ -+ {"kv_connector": "NixlConnector"}, -+ {"kv_connector": "MultiConnector"}, -+ {"kv_connector_module_path": "custom.connector"}, -+ {"kv_role": "kv_producer"}, -+ {"kv_role": "kv_both"}, -+ ], -+) -+def test_unsupported_kv_transfer_is_refused(overrides): -+ config = KVTransferConfig( -+ **( -+ { -+ "kv_connector": "MoRIIOConnector", -+ "kv_role": "kv_consumer", -+ "kv_connector_extra_config": {"read_mode": True}, -+ } -+ | overrides -+ ) -+ ) -+ assert not runtime.supports_kv_transfer(config) -+ -+ - @pytest.mark.parametrize( - "over, reason", - [ -@@ -153,7 +188,14 @@ def test_supported_config_is_accepted(on_mi355x): - ), - ({"parallel_config.pipeline_parallel_size": 2}, "pipeline parallelism"), - ({"parallel_config.decode_context_parallel_size": 2}, "decode context"), -- ({"kv_transfer_config": object()}, "KV transfer connectors"), -+ ( -+ { -+ "kv_transfer_config": KVTransferConfig( -+ kv_connector="NixlConnector", kv_role="kv_consumer" -+ ) -+ }, -+ "KV transfer connectors", -+ ), - ({"speculative_config": None}, "without speculative decoding"), - ({"speculative_config.num_speculative_tokens": 8}, "more than 7"), - ({"speculative_config.enable_adaptive_verification": True}, "adaptive"), -diff --git a/vllm/models/kimi_k3/amd/specialized_decode/README.md b/vllm/models/kimi_k3/amd/specialized_decode/README.md -index 3e0e048b8..602fb2613 100644 ---- a/vllm/models/kimi_k3/amd/specialized_decode/README.md -+++ b/vllm/models/kimi_k3/amd/specialized_decode/README.md -@@ -31,11 +31,18 @@ TP8 on MI355X with all 256 CUs, decode context parallel 1 or 8, speculative - decoding with at most 7 speculative tokens, Model Runner V2, BF16 conv / FP32 - recurrent KDA states. - -+MoRIIO READ and WRITE consumers are supported. READ finishes all attention and -+recurrent-state loads before the forward; WRITE requests wait for transfer -+completion before admission. Neither requires per-layer transfer hooks during -+FULL graph replay. Keep the producer on the stock decode backend and use -+matching cache layouts. Other KV connector configurations remain unsupported. -+ - If you request the specialized decode for anything else, startup fails with one - `ValueError` that lists every unsupported setting: - - - **Config, in the model's `__init__`:** pipeline, data, expert or prefill -- context parallelism, EPLB, DBO, LoRA, KV connectors, routed-expert capture, -+ context parallelism, EPLB, DBO, LoRA, unsupported KV connectors, -+ routed-expert capture, - sleep mode, CPU weight offload, weight updates, no speculative decoding, - adaptive verification, dynamic speculative decoding, ReplaySSM / RecoverSSM, - the V1 model runner, `STOCK_TORCH_COMPILE`, another model's shapes, another GPU. -diff --git a/vllm/models/kimi_k3/amd/specialized_decode/runtime/__init__.py b/vllm/models/kimi_k3/amd/specialized_decode/runtime/__init__.py -index e17a0e3ff..4dd549911 100644 ---- a/vllm/models/kimi_k3/amd/specialized_decode/runtime/__init__.py -+++ b/vllm/models/kimi_k3/amd/specialized_decode/runtime/__init__.py -@@ -13,6 +13,7 @@ from dataclasses import dataclass - - import torch - -+from vllm.config import KVTransferConfig - from vllm.logger import init_logger - - logger = init_logger(__name__) -@@ -103,7 +104,10 @@ def unsupported(vllm_config) -> list[str]: - # the kernels read one batch's attention metadata per step - (pc.use_ubatching, "DBO (dual-batch overlap)"), - (vllm_config.lora_config is not None, "LoRA"), -- (vllm_config.kv_transfer_config is not None, "KV transfer connectors"), -+ ( -+ not supports_kv_transfer(vllm_config.kv_transfer_config), -+ "KV transfer connectors other than a MoRIIO consumer", -+ ), - ( - vllm_config.aux_output_config.enable_return_routed_experts, - "routed-expert capture", -@@ -162,6 +166,16 @@ def unsupported(vllm_config) -> list[str]: - return reasons + platform_reasons() - - -+def supports_kv_transfer(config: KVTransferConfig | None) -> bool: -+ # MoRIIO completes READs before forward and WRITE receives before admission; -+ # neither consumer needs the per-layer hooks specialized layers bypass. -+ return config is None or ( -+ config.kv_connector == "MoRIIOConnector" -+ and config.kv_connector_module_path is None -+ and config.kv_role == "kv_consumer" -+ ) -+ -+ - def text_shapes(text) -> tuple: - kda = getattr(text, "linear_attn_config", None) or {} - return ( diff --git a/inferencex-e2e/benchmarks/multi_node/srt-slurm-recipes/kimik3/vllm/mi355x-fp4/agentx/disagg-variants.yaml b/inferencex-e2e/benchmarks/multi_node/srt-slurm-recipes/kimik3/vllm/mi355x-fp4/agentx/disagg-variants.yaml index 3683082660..c024b6879e 100644 --- a/inferencex-e2e/benchmarks/multi_node/srt-slurm-recipes/kimik3/vllm/mi355x-fp4/agentx/disagg-variants.yaml +++ b/inferencex-e2e/benchmarks/multi_node/srt-slurm-recipes/kimik3/vllm/mi355x-fp4/agentx/disagg-variants.yaml @@ -1,9 +1,6 @@ base: schema: 2 name: kimik3-fp4-mi355x-vllm-disagg-agentic - setup_script: kimik3-nightly-kernels.sh - environment: - KIMIK3_KERNEL_BUILD_JOBS: '12' slurm: time_limit: '08:00:00' sbatch_directives: diff --git a/inferencex-e2e/perf-changelog.yaml b/inferencex-e2e/perf-changelog.yaml index f1b68aa05d..071876d2f6 100644 --- a/inferencex-e2e/perf-changelog.yaml +++ b/inferencex-e2e/perf-changelog.yaml @@ -11106,5 +11106,4 @@ description: - "Use the ROCm nightly image and BF16 SSM states for the throughput points." - "Replace 1P1D c56 and 1P3D c12 with 1P1D c1, K7 specialized decode, and FP32 SSM states." - - "Apply vLLM #60683, #58743, #60808, and #63031 from their source commits." pr-link: https://github.com/SemiAnalysisAI/InferenceX/pull/4391 From 9fc0eacf45d98892b6ea5478919052089bf45ee0 Mon Sep 17 00:00:00 2001 From: Yichao Zhu Date: Sat, 24 Oct 2026 23:40:40 +0800 Subject: [PATCH 4/6] fix(kimik3): pin ROCm nightly image by digest MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 固定包含所需上游功能的 nightly 提交和镜像 digest，确保各节点及重复运行使用同一镜像。 --- inferencex-e2e/configs/amd-master.yaml | 2 +- inferencex-e2e/perf-changelog.yaml | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/inferencex-e2e/configs/amd-master.yaml b/inferencex-e2e/configs/amd-master.yaml index ae3eddfa3d..008888e799 100644 --- a/inferencex-e2e/configs/amd-master.yaml +++ b/inferencex-e2e/configs/amd-master.yaml @@ -1682,7 +1682,7 @@ musespark-fp4-mi355x-sglang-agentic-dflash: # Native Kimi-K3 PD with TP8 workers and MoRIIO READ. kimik3-fp4-mi355x-vllm-disagg-agentic: - image: vllm/vllm-openai-rocm:nightly + image: vllm/vllm-openai-rocm:nightly-c50aca14416ef29f6ed0ed256302f30646726dcd@sha256:ad061ff860a530c324b0ebc5a33c18cd9da170eee2f55704bd52621beedc0b31 model: moonshotai/Kimi-K3 model-prefix: kimik3 runner: cluster:mi355x-amds diff --git a/inferencex-e2e/perf-changelog.yaml b/inferencex-e2e/perf-changelog.yaml index 071876d2f6..1c74206ac9 100644 --- a/inferencex-e2e/perf-changelog.yaml +++ b/inferencex-e2e/perf-changelog.yaml @@ -11104,6 +11104,6 @@ - config-keys: - kimik3-fp4-mi355x-vllm-disagg-agentic description: - - "Use the ROCm nightly image and BF16 SSM states for the throughput points." + - "Use ROCm nightly-c50aca14416ef29f6ed0ed256302f30646726dcd pinned by digest and BF16 SSM states for the throughput points." - "Replace 1P1D c56 and 1P3D c12 with 1P1D c1, K7 specialized decode, and FP32 SSM states." pr-link: https://github.com/SemiAnalysisAI/InferenceX/pull/4391 From a2555a52171cec3117a36a34630a5893de328f0a Mon Sep 17 00:00:00 2001 From: Yichao Zhu Date: Mon, 26 Oct 2026 17:38:04 +0800 Subject: [PATCH 5/6] fix(kimik3): use c8 decode and increase CPU offload budget --- .../vllm/mi355x-fp4/agentx/disagg-variants.yaml | 4 ++-- inferencex-e2e/configs/amd-master.yaml | 14 +++++++------- inferencex-e2e/perf-changelog.yaml | 3 ++- 3 files changed, 11 insertions(+), 10 deletions(-) diff --git a/inferencex-e2e/benchmarks/multi_node/srt-slurm-recipes/kimik3/vllm/mi355x-fp4/agentx/disagg-variants.yaml b/inferencex-e2e/benchmarks/multi_node/srt-slurm-recipes/kimik3/vllm/mi355x-fp4/agentx/disagg-variants.yaml index c024b6879e..75b0182560 100644 --- a/inferencex-e2e/benchmarks/multi_node/srt-slurm-recipes/kimik3/vllm/mi355x-fp4/agentx/disagg-variants.yaml +++ b/inferencex-e2e/benchmarks/multi_node/srt-slurm-recipes/kimik3/vllm/mi355x-fp4/agentx/disagg-variants.yaml @@ -72,7 +72,7 @@ base: # The native srt driver supplies the measured AL for the selected DSpark # draft length in throughput runs and keeps real block rejection for eval. speculative-config: '{"model":"/models/Inferact-Kimi-K3-DSpark","method":"dspark","num_speculative_tokens":4,"attention_backend":"ROCM_AITER_MLA","kv_cache_dtype":"fp8","draft_sample_method":"probabilistic","rejection_sample_method":"block"}' - kv-transfer-config: '{"kv_connector":"MultiConnector","kv_role":"kv_both","kv_load_failure_policy":"fail","kv_connector_extra_config":{"connectors":[{"kv_connector":"MoRIIOConnector","kv_role":"kv_producer","kv_load_failure_policy":"fail","kv_connector_extra_config":{"backend":"rdma"}},{"kv_connector":"SimpleCPUOffloadConnector","kv_role":"kv_both","kv_connector_extra_config":{"cpu_bytes_to_use":1799000000000,"lazy_offload":false}}]}}' + kv-transfer-config: '{"kv_connector":"MultiConnector","kv_role":"kv_both","kv_load_failure_policy":"fail","kv_connector_extra_config":{"connectors":[{"kv_connector":"MoRIIOConnector","kv_role":"kv_producer","kv_load_failure_policy":"fail","kv_connector_extra_config":{"backend":"rdma"}},{"kv_connector":"SimpleCPUOffloadConnector","kv_role":"kv_both","kv_connector_extra_config":{"cpu_bytes_to_use":2399000000000,"lazy_offload":false}}]}}' env: &common_env NCCL_IB_HCA: rdma0,rdma1,rdma2,rdma3,rdma4,rdma5,rdma6,rdma7 NCCL_SOCKET_IFNAME: eno0 @@ -118,7 +118,7 @@ base: placement: node: head -override_1p1d_c1: +override_1p1d_c8: roles: prefill: args: diff --git a/inferencex-e2e/configs/amd-master.yaml b/inferencex-e2e/configs/amd-master.yaml index d042e8c110..f68caf8d28 100644 --- a/inferencex-e2e/configs/amd-master.yaml +++ b/inferencex-e2e/configs/amd-master.yaml @@ -1678,7 +1678,7 @@ kimik3-fp4-mi355x-vllm-disagg-agentic: srt-recipe-dir: kimik3/vllm/mi355x-fp4/agentx scenarios: agentic-coding: - - dram-utilization: 0.6 + - dram-utilization: 0.8 search-space: - spec-decoding: mtp conc-list: [24] @@ -1692,7 +1692,7 @@ kimik3-fp4-mi355x-vllm-disagg-agentic: ep: 1 dp-attn: false additional-settings: - - TOTAL_CPU_DRAM_GB=1799 + - TOTAL_CPU_DRAM_GB=2399 decode: { num-worker: 3, tp: 8, dcp-size: 8, ep: 1, dp-attn: false } - spec-decoding: mtp conc-list: [24] @@ -1706,7 +1706,7 @@ kimik3-fp4-mi355x-vllm-disagg-agentic: ep: 1 dp-attn: false additional-settings: - - TOTAL_CPU_DRAM_GB=1799 + - TOTAL_CPU_DRAM_GB=2399 decode: { num-worker: 2, tp: 8, dcp-size: 8, ep: 1, dp-attn: false } - spec-decoding: mtp conc-list: [48] @@ -1721,7 +1721,7 @@ kimik3-fp4-mi355x-vllm-disagg-agentic: dp-attn: false additional-settings: # Recipe enables P-only scratch reclaim; D and GMU0.90 are unchanged. - - TOTAL_CPU_DRAM_GB=1799 + - TOTAL_CPU_DRAM_GB=2399 decode: { num-worker: 1, tp: 8, dcp-size: 8, ep: 1, dp-attn: false } - spec-decoding: mtp conc-list: [48] @@ -1736,12 +1736,12 @@ kimik3-fp4-mi355x-vllm-disagg-agentic: dp-attn: false additional-settings: # Inherits the same P-only scratch-reclaim setting as 1P1D c48. - - TOTAL_CPU_DRAM_GB=1799 + - TOTAL_CPU_DRAM_GB=2399 decode: { num-worker: 2, tp: 8, dcp-size: 8, ep: 1, dp-attn: false } - search-space: - spec-decoding: mtp - conc-list: [1] + conc-list: [8] kv-offloading: none - srt-recipe: disagg-variants.yaml:override_1p1d_c1 + srt-recipe: disagg-variants.yaml:override_1p1d_c8 prefill: { num-worker: 1, tp: 8, dcp-size: 8, ep: 1, dp-attn: false } decode: { num-worker: 1, tp: 8, dcp-size: 8, ep: 1, dp-attn: false } diff --git a/inferencex-e2e/perf-changelog.yaml b/inferencex-e2e/perf-changelog.yaml index 969ab36153..e94cf999fe 100644 --- a/inferencex-e2e/perf-changelog.yaml +++ b/inferencex-e2e/perf-changelog.yaml @@ -11215,5 +11215,6 @@ - kimik3-fp4-mi355x-vllm-disagg-agentic description: - "Use ROCm nightly-c50aca14416ef29f6ed0ed256302f30646726dcd pinned by digest and BF16 SSM states for the throughput points." - - "Replace 1P1D c56 and 1P3D c12 with 1P1D c1, K7 specialized decode, and FP32 SSM states." + - "Use 1P1D c8 with K7 specialized decode, max-num-seqs=1 on P/D, FP32 SSM states, and GPU-resident KV; retire c1, c56, and 1P3D c12." + - "Set prefill-only SimpleCPUOffload points to DRAM utilization 0.8 (2399 GB per P worker)." pr-link: https://github.com/SemiAnalysisAI/InferenceX/pull/4391 From da441c7d2093e856c1d7aedc02be0ecfae3cc081 Mon Sep 17 00:00:00 2001 From: Yichao Zhu Date: Wed, 28 Oct 2026 14:29:52 +0800 Subject: [PATCH 6/6] docs: correct retired Kimi-K3 points in changelog --- inferencex-e2e/perf-changelog.yaml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/inferencex-e2e/perf-changelog.yaml b/inferencex-e2e/perf-changelog.yaml index 29cbb2f5f6..f59872979b 100644 --- a/inferencex-e2e/perf-changelog.yaml +++ b/inferencex-e2e/perf-changelog.yaml @@ -11418,7 +11418,7 @@ - kimik3-fp4-mi355x-vllm-disagg-agentic description: - "Use ROCm nightly-c50aca14416ef29f6ed0ed256302f30646726dcd pinned by digest and BF16 SSM states for the throughput points." - - "Use 1P1D c8 with K7 specialized decode, max-num-seqs=1 on P/D, FP32 SSM states, and GPU-resident KV; retire c1, c56, and 1P3D c12." + - "Use 1P1D c8 with K7 specialized decode, max-num-seqs=1 on P/D, FP32 SSM states, and GPU-resident KV; retire c56 and 1P3D c12." - "Set prefill-only SimpleCPUOffload points to DRAM utilization 0.8 (2399 GB per P worker)." pr-link: https://github.com/SemiAnalysisAI/InferenceX/pull/4391