iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[ROCm][Perf] Kimi-K3 Fused kernels for KDA prefill (#52606)

Browse Source 
Signed-off-by: kliuae <kuanfu.liu@embeddedllm.com>
Signed-off-by: tjtanaa <tunjian.tan@embeddedllm.com>
Co-authored-by: Claude <noreply@anthropic.com>
Co-authored-by: TJian <tunjian.tan@embeddedllm.com>
This commit is contained in:
kliuae
2026-08-21 09:07:13 -07:00
committed by GitHub
co-authored by Claude TJian
parent d53b1c2efc
commit 463aa5e30f
9 changed files with 2716 additions and 18 deletions
Download Patch File Download Diff File Expand all files Collapse all files
CMakeLists.txt
+21 -5
View File
@@ -1180,11 +1180,13 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
endif()
# HIP counterpart of the FUSED_KDA_DECODE block above: same `fused_kda_decode`
# op, same VLLM_ENABLE_FUSED_KDA_DECODE guard in torch_bindings.cpp. Built for
# the CDNA archs the kernel supports (gfx942 / gfx950), so it is only emitted
# when one of those is in the arch list. A multi-arch build that includes them
# still emits this source for every arch in that list. The runtime gate in
# kda_decode.py keeps the kernel off unsupported devices.
# op, same VLLM_ENABLE_FUSED_KDA_DECODE guard in torch_bindings.cpp
# Built for the CDNA archs the kernel supports
# (gfx942 / gfx950), so it is only emitted when one of those is in the arch
# list. A multi-arch build that includes them still emits this source for
# every arch in that list. The runtime gate in kda_decode.py keeps the kernel
# off unsupported devices. Fused prefill KDA kernels is Only measured on
# gfx950, so it is only built when gfx950 is in the arch list.
if(VLLM_GPU_LANG STREQUAL "HIP")
set(FUSED_KDA_DECODE_HIP_ARCHS ${VLLM_GPU_ARCHES})
list(FILTER FUSED_KDA_DECODE_HIP_ARCHS INCLUDE REGEX "gfx942|gfx950")
@@ -1195,6 +1197,16 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
message(STATUS
"Building fused KDA decode for HIP archs: ${FUSED_KDA_DECODE_HIP_ARCHS}")
endif()
set(FUSED_KDA_CHUNK_HIP_ARCHS ${VLLM_GPU_ARCHES})
list(FILTER FUSED_KDA_CHUNK_HIP_ARCHS INCLUDE REGEX "gfx950")
if(FUSED_KDA_CHUNK_HIP_ARCHS)
set(FUSED_KDA_CHUNK_HIP_SRC
"csrc/libtorch_stable/kimi_k3/fused_kda_chunk_kernel_rocm.cu")
list(APPEND VLLM_STABLE_EXT_SRC "${FUSED_KDA_CHUNK_HIP_SRC}")
message(STATUS
"Building fused KDA chunk for HIP archs: ${FUSED_KDA_CHUNK_HIP_ARCHS}")
endif()
endif()
message(STATUS "Enabling C_stable extension.")
@@ -1245,6 +1257,10 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
target_compile_definitions(_C_stable_libtorch PRIVATE
VLLM_ENABLE_FUSED_KDA_DECODE=1)
endif()
if(FUSED_KDA_CHUNK_HIP_SRC)
target_compile_definitions(_C_stable_libtorch PRIVATE
VLLM_ENABLE_FUSED_KDA_CHUNK=1)
endif()
endif()
# On ROCm, _C_stable_libtorch calls raw HIP APIs (e.g. hipGetDevice in
csrc/libtorch_stable/kimi_k3/fused_kda_chunk_kernel_rocm.cu
+1790
View File
File diff suppressed because it is too large Load Diff
csrc/libtorch_stable/ops.h
+27
View File
@@ -406,6 +406,33 @@ void fused_gdn_decode_post_conv_mtp(
#endif
#ifdef VLLM_ENABLE_FUSED_KDA_CHUNK
void fused_kda_prologue(
torch::stable::Tensor const& q, torch::stable::Tensor const& k,
torch::stable::Tensor const& v, torch::stable::Tensor const& raw_g,
torch::stable::Tensor const& raw_beta, torch::stable::Tensor const& a_log,
torch::stable::Tensor const& dt_bias, torch::stable::Tensor& qg,
torch::stable::Tensor& w, torch::stable::Tensor& u,
torch::stable::Tensor& kg_t, torch::stable::Tensor& aqk,
torch::stable::Tensor& decay, torch::stable::Tensor const& cu_seqlens,
torch::stable::Tensor const& chunk_indices,
std::optional<torch::stable::Tensor> conv_weight,
std::optional<torch::stable::Tensor> conv_state,
std::optional<torch::stable::Tensor> conv_state_indices,
std::optional<torch::stable::Tensor> conv_has_initial_state, double scale,
double lower_bound);
void fused_kda_chunk(
torch::stable::Tensor const& qg, torch::stable::Tensor const& w,
torch::stable::Tensor const& u, torch::stable::Tensor const& kg_t,
torch::stable::Tensor const& aqk, torch::stable::Tensor const& decay,
std::optional<torch::stable::Tensor> initial_state,
std::optional<torch::stable::Tensor> final_state,
torch::stable::Tensor& out, torch::stable::Tensor const& cu_seqlens,
torch::stable::Tensor const& chunk_offsets, double scale,
std::optional<torch::stable::Tensor> group_state, int64_t groups);
#endif
#ifdef VLLM_ENABLE_KIMI_K3_ATTN_RES
void kimi_k3_attn_res(torch::stable::Tensor& prefix,
torch::stable::Tensor const& delta,
csrc/libtorch_stable/torch_bindings.cpp
+22
View File
@@ -538,6 +538,23 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
"float scale, float norm_eps=1e-5) -> ()");
#endif
#ifdef VLLM_ENABLE_FUSED_KDA_CHUNK
ops.def(
"fused_kda_prologue("
"Tensor q, Tensor k, Tensor v, Tensor raw_g, Tensor raw_beta, "
"Tensor A_log, Tensor dt_bias, Tensor! qg, Tensor! w, Tensor! u, "
"Tensor! kg_t, Tensor! aqk, Tensor! decay, Tensor cu_seqlens, "
"Tensor chunk_indices, Tensor? conv_weight, Tensor(e!)? conv_state, "
"Tensor? conv_state_indices, Tensor? conv_has_initial_state, "
"float scale, float lower_bound) -> ()");
ops.def(
"fused_kda_chunk("
"Tensor qg, Tensor w, Tensor u, Tensor kg_t, Tensor aqk, Tensor decay, "
"Tensor? initial_state, Tensor(a!)? final_state, Tensor! out, "
"Tensor cu_seqlens, Tensor chunk_offsets, float scale, "
"Tensor(b!)? group_state, int groups) -> ()");
#endif
#ifdef VLLM_ENABLE_KIMI_K3_ATTN_RES
ops.def(
"kimi_k3_attn_res("
@@ -804,6 +821,11 @@ STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) {
TORCH_BOX(&fused_gdn_decode_post_conv_mtp));
#endif
#ifdef VLLM_ENABLE_FUSED_KDA_CHUNK
ops.impl("fused_kda_prologue", TORCH_BOX(&fused_kda_prologue));
ops.impl("fused_kda_chunk", TORCH_BOX(&fused_kda_chunk));
#endif
#ifdef VLLM_ENABLE_KIMI_K3_ATTN_RES
ops.impl("kimi_k3_attn_res", TORCH_BOX(&kimi_k3_attn_res));
#endif
tests/models/kimi_k3/test_amd_kda_chunk.py
+426
View File
@@ -0,0 +1,426 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""The fused ROCm KDA chunk kernel must match the Triton chunk path.
The fused kernel reassociates nothing, but it keeps the chunk states in
registers and rounds them to bfloat16 only where the Triton path does, so the
outputs agree to bfloat16 rounding and the final states nearly exactly.
"""
import pytest
import torch
from vllm.platforms import current_platform
def _on_gfx950() -> bool:
if not current_platform.is_rocm():
return False
from vllm.platforms.rocm import on_gfx950
return on_gfx950()
pytestmark = pytest.mark.skipif(
not _on_gfx950(),
reason="The fused KDA chunk kernel is only built for gfx950",
)
HEAD_DIM = 128
NUM_HEADS = 12
LOWER_BOUND = -5.0
def _requires_kernel() -> None:
if not hasattr(torch.ops._C, "fused_kda_chunk"):
pytest.skip("vLLM was built without the fused KDA chunk kernel")
def _inputs(seqlens: list[int], seed: int) -> dict:
torch.manual_seed(seed)
dev, dtype = "cuda", torch.bfloat16
total = sum(seqlens)
cu = torch.tensor(
[0] + torch.tensor(seqlens).cumsum(0).tolist(), device=dev, dtype=torch.int32
)
shape = (1, total, NUM_HEADS, HEAD_DIM)
return dict(
q=torch.randn(shape, device=dev, dtype=dtype) * 0.5,
k=torch.randn(shape, device=dev, dtype=dtype) * 0.5,
v=torch.randn(shape, device=dev, dtype=dtype) * 0.5,
raw_g=torch.randn(shape, device=dev, dtype=dtype),
raw_beta=torch.randn(1, total, NUM_HEADS, device=dev, dtype=dtype),
A_log=torch.randn(NUM_HEADS, device=dev, dtype=torch.float32) * 0.5,
dt_bias=torch.randn(NUM_HEADS * HEAD_DIM, device=dev, dtype=torch.float32)
* 0.5,
h0=torch.randn(
len(seqlens), NUM_HEADS, HEAD_DIM, HEAD_DIM, device=dev, dtype=torch.float32
)
* 0.1,
cu=cu,
)
def _run(inp: dict, use_fused: bool, out: torch.Tensor | None = None, **kw):
from vllm.models.kimi_k3.amd.ops.kda_prefill import chunk_kda_prefill
return chunk_kda_prefill(
q=inp["q"].clone(),
k=inp["k"].clone(),
v=inp["v"].clone(),
raw_g=inp["raw_g"].clone(),
raw_beta=inp["raw_beta"],
A_log=inp["A_log"],
g_bias=inp["dt_bias"],
lower_bound=LOWER_BOUND,
initial_state=inp["h0"].clone(),
output_final_state=True,
use_qk_l2norm_in_kernel=True,
cu_seqlens=inp["cu"],
use_fused_chunk=use_fused,
out=out,
**kw,
)
def _force_groups(monkeypatch, groups: int) -> None:
"""Pin the group count, which is otherwise a function of the device."""
import vllm.models.kimi_k3.amd.ops.kda_chunk as kda_chunk
monkeypatch.setattr(kda_chunk, "_chunk_groups", lambda *_: groups)
@pytest.mark.parametrize(
"seqlens",
[
[1], # shorter than a chunk
[64], # exactly one chunk
[1024],
[4096], # many chunks, one sequence
[513, 64, 1, 1200], # ragged, including a partial tail chunk
[512] * 6,
],
)
def test_fused_chunk_matches_triton(seqlens: list[int]) -> None:
_requires_kernel()
inp = _inputs(seqlens, seed=len(seqlens))
o_ref, state_ref = _run(inp, use_fused=False)
o_fused, state_fused = _run(inp, use_fused=True)
torch.testing.assert_close(o_fused.float(), o_ref.float(), rtol=2e-2, atol=2e-2)
torch.testing.assert_close(state_fused, state_ref, rtol=1e-3, atol=1e-3)
@pytest.mark.parametrize("use_fused", [False, True])
def test_out_argument_receives_the_result(use_fused: bool) -> None:
"""``out`` is a promise about where the result lands.
The layer passes a slice of its own output buffer so neither path has to
copy afterwards, and it then skips the copy based on the returned tensor
aliasing that buffer. Both must hold for the fused path and for the Triton
fallback, which allocates its own buffer internally.
"""
if use_fused:
_requires_kernel()
seqlens = [321, 64, 700]
inp = _inputs(seqlens, seed=13)
o_ref, _ = _run(inp, use_fused=use_fused)
# A non-zero offset into a larger buffer, as the layer produces when a
# mixed batch puts decode tokens first.
total = sum(seqlens)
buf = torch.empty(
1, total + 8, NUM_HEADS, HEAD_DIM, device="cuda", dtype=o_ref.dtype
)
view = buf[:, 8:]
o_got, _ = _run(inp, use_fused=use_fused, out=view)
assert o_got.data_ptr() == view.data_ptr()
torch.testing.assert_close(view.float(), o_ref.float(), rtol=2e-2, atol=2e-2)
def test_fused_chunk_accepts_projection_views() -> None:
"""`beta` and `g` arrive as last-dim slices of the fused QKVGFAB projection.
Those views carry the projection's row stride, so a kernel that indexes them
densely reads the wrong elements or rejects the tensor outright. Build the
inputs the way the layer does rather than as standalone tensors.
"""
_requires_kernel()
seqlens = [777, 64, 391]
inp = _inputs(seqlens, seed=7)
total = sum(seqlens)
# One wide projection, split on the last dim exactly like KimiK3DeltaAttention.
qkv_width = 3 * NUM_HEADS * HEAD_DIM
projection = torch.randn(
total, qkv_width + NUM_HEADS, device="cuda", dtype=torch.bfloat16
)
beta_view = projection.split([qkv_width, NUM_HEADS], dim=-1)[1].unsqueeze(0)
assert not beta_view.is_contiguous()
inp["raw_beta"] = beta_view.contiguous()
o_ref, state_ref = _run(inp, use_fused=False)
inp["raw_beta"] = beta_view
o_fused, state_fused = _run(inp, use_fused=True)
torch.testing.assert_close(o_fused.float(), o_ref.float(), rtol=2e-2, atol=2e-2)
torch.testing.assert_close(state_fused, state_ref, rtol=1e-3, atol=1e-3)
def test_fused_chunk_survives_correlated_keys_and_real_gate_bias() -> None:
"""The input regime the served model actually produces.
This is what forces the intra-chunk inverse to keep fp32 MFMA operands: at
bf16 the Neumann doubling's intermediate powers cancel catastrophically
here, diverging by orders of magnitude and driving the final state to inf.
Three properties of real traffic combine here and in none of the tests
above, all measured from tensors captured off a live TP8 server:
* repetitive text gives near-duplicate L2-normalized keys, so ``k_i . k_j``
approaches 1 and the Neumann terms stop shrinking;
* the checkpoint's ``dt_bias`` lies in ``[-7.8, -1.4]``, which makes the
within-chunk gate cumsum span ~300 powers of two;
* ``beta`` saturates near 0.95.
``_inputs`` draws all three from ``randn``, landing in the benign regime.
"""
_requires_kernel()
torch.manual_seed(11)
dev, dtype = "cuda", torch.bfloat16
total = 1088
shape = (1, total, NUM_HEADS, HEAD_DIM)
base = torch.randn(1, 1, NUM_HEADS, HEAD_DIM, device=dev, dtype=dtype)
def correlated():
return (base + torch.randn(shape, device=dev, dtype=dtype) * 0.02).contiguous()
inp = dict(
q=correlated(),
k=correlated(),
v=torch.randn(shape, device=dev, dtype=dtype) * 0.5,
raw_g=torch.randn(shape, device=dev, dtype=dtype),
raw_beta=torch.full((1, total, NUM_HEADS), 3.0, device=dev, dtype=dtype),
A_log=torch.rand(NUM_HEADS, device=dev, dtype=torch.float32) * 0.6 - 0.42,
dt_bias=torch.rand(NUM_HEADS * HEAD_DIM, device=dev, dtype=torch.float32) * 6.4
- 7.8,
h0=torch.zeros(
1, NUM_HEADS, HEAD_DIM, HEAD_DIM, device=dev, dtype=torch.float32
),
cu=torch.tensor([0, total], device=dev, dtype=torch.int32),
)
o_ref, state_ref = _run(inp, use_fused=False)
o_fused, state_fused = _run(inp, use_fused=True)
# Looser than the tests above: near-identical keys plus a saturated beta
# are past anything real text produces, and the inverse is rounded to bf16
# before the w/u GEMMs. The tolerance still catches divergence, which in
# this regime is orders of magnitude rather than a few times the bound.
torch.testing.assert_close(o_fused.float(), o_ref.float(), rtol=5e-2, atol=5e-2)
torch.testing.assert_close(state_fused, state_ref, rtol=5e-2, atol=5e-2)
@pytest.mark.parametrize(
"seqlens",
[
[1], # every conv tap comes from the cache
[2], # cache and batch both contribute
[64],
[1024],
[513, 64, 1, 1200], # ragged, tail chunk shorter than the conv width
],
)
@pytest.mark.parametrize("has_state", [False, True])
def test_fused_conv_matches_reference_conv(seqlens: list[int], has_state: bool) -> None:
"""The prologue's optional fused convolution must match a separate conv.
The layer does not use this path -- recomputing the conv once per load
phase costs more than the ``causal_conv1d_fn`` launches it saves -- so it
is driven through the ops directly rather than through
``chunk_kda_prefill``, which does not expose it.
The reference is computed in torch rather than with ``causal_conv1d_fn``,
which is not self-consistent for these shapes on ROCm (it leaves part of
its output unwritten, so repeated identical calls disagree).
The cases that matter are the ones touching the sequence boundary, where
the fused path takes taps from the conv cache instead of the batch:
sequences shorter than the kernel width, and a cache that is or is not
populated. The rolled-forward cache is checked too, since a wrong shift
only shows up on the next step.
"""
import torch.nn.functional as F
from vllm.models.kimi_k3.amd.ops.kda_chunk import (
fused_kda_chunk,
fused_kda_prologue,
)
_requires_kernel()
dev, dtype = "cuda", torch.bfloat16
lp = NUM_HEADS * HEAD_DIM
width, n = 4, len(seqlens)
inp = _inputs(seqlens, seed=7 + len(seqlens) + int(has_state))
total = int(inp["cu"][-1])
torch.manual_seed(11 + int(has_state))
# The layer hands the prologue a strided view of the QKVGFAB projection.
projection = torch.randn(total, 3 * lp + 64, device=dev, dtype=dtype) * 0.5
mixed = projection[:, : 3 * lp]
# conv1d is built with bias=False and fp32 weights; the fused kernel takes
# the width-major [qkv, width, channel] mirror the decode path already
# keeps.
weight = torch.randn(3, width, lp, device=dev, dtype=torch.float32) * 0.3
store = torch.zeros(n, width - 1, 3 * lp, device=dev, dtype=dtype)
if has_state:
store.normal_(0.0, 0.5)
state = store.transpose(-1, -2)
idx = torch.arange(n, device=dev, dtype=torch.int32)
has_init = torch.full((n,), has_state, device=dev, dtype=torch.bool)
# Reference conv, per sequence, in fp32 then rounded like a bf16 output.
conv_ref = torch.empty(total, 3 * lp, device=dev, dtype=dtype)
state_ref = torch.empty_like(store)
flat_w = weight.permute(0, 2, 1).reshape(3 * lp, width)
for i, (b, e) in enumerate(zip(inp["cu"][:-1].tolist(), inp["cu"][1:].tolist())):
pre = (
store[i].float()
if has_state
else torch.zeros(width - 1, 3 * lp, device=dev, dtype=torch.float32)
)
xs = torch.cat([pre, mixed[b:e].float()], 0)
acc = torch.zeros(e - b, 3 * lp, device=dev, dtype=torch.float32)
for w in range(width):
acc += xs[w : w + (e - b)] * flat_w[:, w].unsqueeze(0)
conv_ref[b:e] = F.silu(acc).to(dtype)
state_ref[i] = xs[-(width - 1) :].to(dtype)
def _bands(x, contiguous):
out = [
b.unsqueeze(0).unflatten(-1, (NUM_HEADS, HEAD_DIM))
for b in x.split(lp, dim=-1)
]
return [b.contiguous() for b in out] if contiguous else out
# Compare through both kernels rather than the prologue workspaces: `aqk`
# is deliberately left unwritten past each chunk's valid rows, so those
# elements are whatever the allocator last held.
scale = HEAD_DIM**-0.5
def _through_kernels(q, k, v, conv: dict | None):
ws = fused_kda_prologue(
q=q,
k=k,
v=v,
raw_g=inp["raw_g"],
raw_beta=inp["raw_beta"],
A_log=inp["A_log"],
dt_bias=inp["dt_bias"],
scale=scale,
lower_bound=LOWER_BOUND,
cu_seqlens=inp["cu"],
**(conv or {}),
)
return fused_kda_chunk(
qg=ws["qg"],
w=ws["w"],
u=ws["u"],
kg_t=ws["kg_t"],
aqk=ws["aqk"],
decay=ws["decay"],
out=torch.empty_like(ws["u"]),
scale=scale,
cu_seqlens=inp["cu"],
initial_state=inp["h0"].clone(),
output_final_state=True,
)
qr, kr, vr = _bands(conv_ref, True)
o_ref, s_ref = _through_kernels(qr, kr, vr, None)
got_state = state.clone()
qf, kf, vf = _bands(mixed, False)
o_got, s_got = _through_kernels(
qf,
kf,
vf,
dict(
conv_weight=weight,
conv_state=got_state,
conv_state_indices=idx,
conv_has_initial_state=has_init,
),
)
torch.testing.assert_close(o_got.float(), o_ref.float(), rtol=2e-2, atol=2e-2)
torch.testing.assert_close(s_got, s_ref, rtol=1e-3, atol=1e-3)
torch.testing.assert_close(
got_state.transpose(-1, -2).float(), state_ref.float(), rtol=0, atol=0
)
@pytest.mark.parametrize("groups", [2, 4, 8])
@pytest.mark.parametrize("seqlens", [[64], [64, 64], [100, 64], [513, 64, 1]])
def test_chunk_group_scan_matches_serial_walk(
groups: int, seqlens: list[int], monkeypatch
) -> None:
"""Splitting the chunk walk into parallel groups must not change the result.
Two bugs lived here, both invisible while the path was gated to >= 128
chunks per sequence:
* a group with no chunks returned before writing its ``B_g``, so the scan
folded whatever the workspace last held into the final state — these
shapes have fewer chunks than groups, so most groups are empty;
* the second pass started every group after the first from a zero state
instead of the one the scan composed, which only looked right because a
strongly decaying gate makes the dropped term small.
"""
_requires_kernel()
inp = _inputs(seqlens, seed=5 + groups)
_force_groups(monkeypatch, 1)
o_ref, state_ref = _run(inp, use_fused=True)
_force_groups(monkeypatch, groups)
o_got, state_got = _run(inp, use_fused=True)
torch.testing.assert_close(o_got.float(), o_ref.float(), rtol=2e-2, atol=2e-2)
torch.testing.assert_close(state_got, state_ref, rtol=1e-3, atol=1e-3)
@pytest.mark.parametrize("seqlens", [[320], [513, 64, 1, 1200], [64, 64, 64]])
def test_accepts_chunk_metadata_from_the_attention_builder(
seqlens: list[int],
) -> None:
"""The layer passes precomputed chunk metadata, not None.
``KimiK3ROCmKDAMetadataBuilder`` builds these on device to avoid a blocking
device->host copy. Passing ``None`` instead takes a different path that
builds them locally, so drive the op with the builder's own tensors.
"""
_requires_kernel()
from vllm.models.kimi_k3.amd.kda_metadata import prepare_chunk_metadata_device
from vllm.third_party.flash_linear_attention.ops.utils import FLA_CHUNK_SIZE
inp = _inputs(seqlens, seed=17)
cu = inp["cu"]
chunk_indices, chunk_offsets = prepare_chunk_metadata_device(
cu, cu.cpu(), FLA_CHUNK_SIZE
)
o_ref, state_ref = _run(inp, use_fused=True)
o_got, state_got = _run(
inp, use_fused=True, chunk_indices=chunk_indices, chunk_offsets=chunk_offsets
)
torch.testing.assert_close(o_got.float(), o_ref.float(), rtol=2e-2, atol=2e-2)
torch.testing.assert_close(state_got, state_ref, rtol=1e-3, atol=1e-3)
vllm/engine/arg_utils.py
+15 -3
View File
@@ -764,7 +764,7 @@ class EngineArgs:
fail_on_environ_validation: bool = False
gdn_prefill_backend: Literal["flashinfer", "triton", "cutedsl"] | None = None
kda_prefill_backend: Literal["auto", "triton", "flashkda"] | None = None
kda_prefill_backend: Literal["auto", "triton", "flashkda", "fused"] | None = None
def __post_init__(self):
# support `EngineArgs(compilation_config={...})`
@@ -1720,9 +1720,10 @@ class EngineArgs:
parser.add_argument(
"--kda-prefill-backend",
dest="kda_prefill_backend",
choices=["auto", "triton", "flashkda"],
choices=["auto", "triton", "flashkda", "fused"],
default=None,
help="Select KDA prefill backend.",
help="Select KDA prefill backend. 'flashkda' is CUDA-only and "
"'fused' is ROCm-only; 'auto' picks a supported backend.",
)
return parser
@@ -2523,6 +2524,17 @@ class EngineArgs:
if self.gdn_prefill_backend is not None:
self.additional_config["gdn_prefill_backend"] = self.gdn_prefill_backend
if self.kda_prefill_backend is not None:
if (
self.kda_prefill_backend == "flashkda"
and not current_platform.is_cuda()
):
raise ValueError(
"--kda-prefill-backend=flashkda is only available on CUDA."
)
if self.kda_prefill_backend == "fused" and not current_platform.is_rocm():
raise ValueError(
"--kda-prefill-backend=fused is only available on ROCm."
)
self.additional_config["kda_prefill_backend"] = self.kda_prefill_backend
config = VllmConfig(
vllm/models/kimi_k3/amd/kda.py
+38 -10
View File
@@ -41,13 +41,16 @@ from vllm.model_executor.layers.mamba.ops.gather_initial_states import (
from vllm.model_executor.model_loader.weight_utils import sharded_weight_loader
from vllm.model_executor.utils import set_weight_attrs
from vllm.models.kimi_k3.amd.kda_metadata import KimiK3ROCmKDABackend
from vllm.models.kimi_k3.amd.ops.kda_chunk import (
is_fused_kda_chunk_supported,
)
from vllm.models.kimi_k3.amd.ops.kda_decode import (
is_fused_kda_decode_supported,
make_decode_conv1d_weight_loader,
make_decode_norm_weight_loader,
)
from vllm.models.kimi_k3.amd.ops.kda_prefill import chunk_kda_prefill
from vllm.models.kimi_k3.amd.ops.third_party.kda import (
chunk_kda_with_fused_gate,
fused_recurrent_kda,
fused_recurrent_kda_packed_decode,
)
@@ -210,10 +213,21 @@ class KimiK3DeltaAttention(GatedDeltaNetAttention):
if isinstance(additional_config, dict)
else "auto"
)
backend = "triton" if backend == "auto" else backend
assert backend == "triton", (
"The ROCm Kimi-K3 KDA layer only supports the Triton KDA prefill "
f"backend, got {backend!r}."
assert backend in ("auto", "triton", "fused"), (
"The ROCm Kimi-K3 KDA prefill backend must be one of "
f"'auto', 'triton' or 'fused', got {backend!r}."
)
if backend == "fused" and not is_fused_kda_chunk_supported():
raise RuntimeError(
"The fused KDA chunk kernel requires gfx950 and a build that "
"includes it."
)
self.use_fused_chunk = backend == "fused" or (
backend == "auto" and is_fused_kda_chunk_supported()
)
logger.info_once(
"Kimi-K3 KDA prefill backend: %s",
"fused" if self.use_fused_chunk else "triton",
)
self.o_norm = FusedRMSNormGated(self.head_dim, activation="sigmoid")
@@ -496,6 +510,12 @@ class KimiK3DeltaAttention(GatedDeltaNetAttention):
# prefill tail only.
core_attn_out_decode = None
split_non_spec = spec_sequence_masks is None and m.num_decodes > 0
# Without a spec split the non-spec tokens are exactly
# core_attn_out[:, :num_actual_tokens] in order, so both kernels
# can write their slice straight into it. With one, the tokens
# are interleaved and have to be scattered by index afterwards.
direct = spec_sequence_masks is None
use_fused_chunk = self.use_fused_chunk and spec_sequence_masks is None
if split_non_spec:
assert non_spec_query_start_loc is not None
nd_tok = m.num_decode_tokens
@@ -513,6 +533,7 @@ class KimiK3DeltaAttention(GatedDeltaNetAttention):
ssm_state_indices=non_spec_state_indices_tensor[
: m.num_decodes
],
out=core_attn_out[:, :nd_tok] if direct else None,
)
q_ns = q_ns[:, nd_tok:]
k_ns = k_ns[:, nd_tok:]
@@ -529,6 +550,7 @@ class KimiK3DeltaAttention(GatedDeltaNetAttention):
prefill_query_start_loc = non_spec_query_start_loc
prefill_state_indices = non_spec_state_indices_tensor
prefill_has_initial_state = has_initial_state
nd_tok = 0
initial_state = gather_initial_states(
recurrent_state,
@@ -538,7 +560,7 @@ class KimiK3DeltaAttention(GatedDeltaNetAttention):
(
core_attn_out_non_spec,
last_recurrent_state,
) = chunk_kda_with_fused_gate(
) = chunk_kda_prefill(
q=q_ns,
k=k_ns,
v=v_ns,
@@ -553,11 +575,16 @@ class KimiK3DeltaAttention(GatedDeltaNetAttention):
cu_seqlens=prefill_query_start_loc,
chunk_indices=m.chunk_indices,
chunk_offsets=m.chunk_offsets,
use_fused_chunk=use_fused_chunk,
out=core_attn_out[:, nd_tok:num_actual_tokens] if direct else None,
)
# Init cache
recurrent_state[prefill_state_indices] = last_recurrent_state
if split_non_spec:
if direct:
# Both slices already landed in core_attn_out.
core_attn_out_non_spec = core_attn_out[:, :num_actual_tokens]
elif split_non_spec:
# Restore decode-first token order for the merge below.
core_attn_out_non_spec = torch.cat(
[core_attn_out_decode, core_attn_out_non_spec], dim=1
@@ -609,9 +636,10 @@ class KimiK3DeltaAttention(GatedDeltaNetAttention):
merged.index_copy_(1, non_spec_token_indx, core_attn_out_non_spec)
core_attn_out[0, :num_actual_tokens] = merged[0, :num_actual_tokens]
elif core_attn_out_non_spec is not None:
core_attn_out[0, :num_actual_tokens] = core_attn_out_non_spec[
0, :num_actual_tokens
]
if core_attn_out_non_spec.data_ptr() != core_attn_out.data_ptr():
core_attn_out[0, :num_actual_tokens] = core_attn_out_non_spec[
0, :num_actual_tokens
]
else:
assert core_attn_out_spec is not None
core_attn_out.copy_(self.o_norm(core_attn_out, g2))
vllm/models/kimi_k3/amd/ops/kda_chunk.py
+251
View File
@@ -0,0 +1,251 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""ROCm entry point for the fused Kimi-K3 KDA chunk kernel.
The kernel in ``csrc/libtorch_stable/kimi_k3/fused_kda_chunk_kernel_rocm.cu``
replaces the chunk-state recurrence and the output GEMM of the Triton chunk
path with a single launch that keeps the per-chunk state in registers, so the
``[chunks, H, V, K]`` state tensor and the recomputed values never reach HBM.
"""
from functools import cache
import torch
from vllm.third_party.flash_linear_attention.ops.index import (
prepare_chunk_indices,
prepare_chunk_offsets,
)
CHUNK_SIZE = 64
HEAD_DIM = 128
# Chunk-group split policy constants; see _chunk_groups.
_BLOCKS_P1 = 2 # (kV + kK) / kBV
_BLOCKS_P2 = 1 # kV / kBV
_MIN_LEN = 4
_DEEP_LEN = 12
_MAX_GROUPS = 32
@cache
def is_fused_kda_chunk_supported() -> bool:
from vllm.platforms.rocm import on_gfx950
if not hasattr(torch.ops._C, "fused_kda_chunk"):
return False
# TODO: Verify on other archs, currently only validated on gfx950
return on_gfx950()
def can_use_fused_kda_chunk(
head_k_dim: int,
head_v_dim: int,
dtype: torch.dtype,
chunk_size: int,
) -> bool:
return (
head_k_dim == HEAD_DIM
and head_v_dim == HEAD_DIM
and chunk_size == CHUNK_SIZE
and dtype == torch.bfloat16
and is_fused_kda_chunk_supported()
)
def fused_kda_prologue(
q: torch.Tensor,
k: torch.Tensor,
v: torch.Tensor,
raw_g: torch.Tensor,
raw_beta: torch.Tensor,
A_log: torch.Tensor,
dt_bias: torch.Tensor,
scale: float,
lower_bound: float,
cu_seqlens: torch.Tensor,
conv_weight: torch.Tensor | None = None,
conv_state: torch.Tensor | None = None,
conv_state_indices: torch.Tensor | None = None,
conv_has_initial_state: torch.Tensor | None = None,
chunk_indices: torch.Tensor | None = None,
) -> dict[str, torch.Tensor]:
"""Run the whole chunk-path prologue in one launch.
Replaces the two L2 norms, the gate cumsum, both intra-chunk passes and the
w/u recompute. Returns the operands the fused chunk kernel consumes.
The kernel can also apply the depthwise conv and its silu, in which case
``q``/``k``/``v`` are the three raw bands of the QKV projection and
``conv_state`` is updated in place. This path is currently not used.
"""
# `beta` and `g` reach the layer as last-dim slices of the fused QKVGFAB
# projection and carry its row stride. The kernel reads both with an
# explicit per-token stride, so neither is copied here.
_, t_total, num_heads, _ = q.shape
if chunk_indices is None:
chunk_indices = prepare_chunk_indices(cu_seqlens, CHUNK_SIZE)
chunk_indices = chunk_indices.to(torch.int32)
num_chunks = chunk_indices.shape[0]
dev = q.device
# q may be a strided band view, so the workspaces are sized rather than
# cloned from it.
def _like(last: int) -> torch.Tensor:
return torch.empty(1, t_total, num_heads, last, dtype=q.dtype, device=q.device)
ws = dict(
qg=_like(HEAD_DIM),
w=_like(HEAD_DIM),
u=_like(HEAD_DIM),
kg_t=torch.empty(
num_chunks, num_heads, HEAD_DIM, CHUNK_SIZE, dtype=q.dtype, device=dev
),
aqk=torch.empty(1, t_total, num_heads, CHUNK_SIZE, dtype=q.dtype, device=dev),
decay=torch.empty(
num_chunks, num_heads, HEAD_DIM, dtype=torch.float32, device=dev
),
)
torch.ops._C.fused_kda_prologue(
q,
k,
v,
raw_g,
raw_beta,
A_log.reshape(-1),
dt_bias.reshape(-1),
ws["qg"],
ws["w"],
ws["u"],
ws["kg_t"],
ws["aqk"],
ws["decay"],
cu_seqlens.to(torch.int32),
chunk_indices,
conv_weight,
conv_state,
conv_state_indices,
conv_has_initial_state,
scale,
lower_bound,
)
return ws
@cache
def _num_cus(device: int) -> int:
from vllm.platforms import current_platform
return current_platform.get_device_properties(device).multi_processor_count
@cache
def _chunk_groups(chunks_per_seq: int, num_seqs: int, num_heads: int) -> int:
"""How many parallel chunk groups to cut each sequence into.
The scan composes the group operators with the exact transfer
``M_g = prod_c (diag(d_c) - w_c^T kg_c)``, so the composition error does
not grow with group length and the choice is purely a machine fit:
* ``fill1``/``fill2``: the largest G whose pass-one / pass-two grid still
fits in one scheduling batch. Rows are launched in blocks of at most
``kBV = 128``, so pass one (``kV + kK`` rows) needs two blocks per
``(n, h, g)`` and pass two one. Past a pass's fill point its block count
grows as fast as its depth shrinks and the rung stops paying.
* ``_DEEP_LEN``: unless the groups are still long, where pass two's halving
is worth letting pass one spill to a second batch.
* ``_MIN_LEN``: below this many chunks per group the per-group fixed cost
exceeds what the shorter walk saves.
``G == 2`` is depth-neutral -- two passes of ``nt/2`` -- and only doubles
staging, so the split is taken only from ``G >= 4``.
"""
cus = _num_cus(torch.accelerator.current_device_index())
nh = max(num_seqs * num_heads, 1)
fill1 = cus // (_BLOCKS_P1 * nh)
fill2 = cus // (_BLOCKS_P2 * nh)
cand = fill1
if fill2 > 0 and chunks_per_seq // fill2 >= _DEEP_LEN:
cand = fill2
cand = min(cand, chunks_per_seq // _MIN_LEN, _MAX_GROUPS)
return cand if cand >= 4 else 1
def _kda_group_workspace(
groups: int, nh: int, device: torch.device
) -> torch.Tensor | None:
"""One fp32 buffer the kernel carves into bg / sin_ / ag / mgT.
Kept caller-owned so the memory stays inside the caching allocator and
vLLM's memory profiling accounts for it; the kernel's carve order must
match the layout here.
"""
if groups <= 1:
return None
planes = groups * nh
plane = HEAD_DIM * HEAD_DIM
floats = 2 * planes * plane + planes * HEAD_DIM # bg, sin_, ag
floats += (planes * plane + 1) // 2 # mgT, bf16
return torch.empty(floats, dtype=torch.float32, device=device)
def fused_kda_chunk(
qg: torch.Tensor,
w: torch.Tensor,
u: torch.Tensor,
kg_t: torch.Tensor,
aqk: torch.Tensor,
decay: torch.Tensor,
out: torch.Tensor,
scale: float,
cu_seqlens: torch.Tensor,
initial_state: torch.Tensor | None,
output_final_state: bool,
chunk_offsets: torch.Tensor | None = None,
) -> tuple[torch.Tensor, torch.Tensor | None]:
"""Run the chunk recurrence and the output projection in one launch.
Args:
qg: ``q * exp2(gk_cumsum)``, ``[1, T, H, 128]``.
kg_t: chunk-major transposed gated keys, ``[chunks, H, 128, 64]``.
decay: ``exp2`` of each chunk's last gate row, ``[chunks, H, 128]``.
out: output buffer, ``[1, T, H, 128]``; may alias ``u``'s source.
"""
num_seqs = cu_seqlens.numel() - 1
final_state = None
if output_final_state:
final_state = torch.empty(
num_seqs,
u.shape[2],
u.shape[3],
qg.shape[3],
dtype=torch.float32,
device=u.device,
)
cu_seqlens = cu_seqlens.to(torch.int32)
if chunk_offsets is None:
chunk_offsets = prepare_chunk_offsets(cu_seqlens, CHUNK_SIZE)
# The chunk walk is serial in chunks. Splitting each sequence into `groups`
# stretches that run in parallel costs a second pass and a scan, and only
# pays once the chain is long enough to amortise both.
chunks_per_seq = kg_t.shape[0] // max(num_seqs, 1)
groups = _chunk_groups(chunks_per_seq, num_seqs, u.shape[2])
group_state = _kda_group_workspace(groups, num_seqs * u.shape[2], u.device)
torch.ops._C.fused_kda_chunk(
qg,
w,
u,
kg_t,
aqk,
decay,
initial_state,
final_state,
out,
cu_seqlens,
chunk_offsets,
scale,
group_state,
groups,
)
return out, final_state
vllm/models/kimi_k3/amd/ops/kda_prefill.py
+126
View File
@@ -0,0 +1,126 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""KDA prefill backend selection for ROCm.
The Kimi-K3 KDA layer calls :func:`chunk_kda_prefill`, which either runs the
fused HIP kernels in ``kda_chunk`` or falls back to the vendored Triton chunk
path.
"""
import torch
from vllm.logger import init_logger
from vllm.models.kimi_k3.amd.ops.kda_chunk import (
can_use_fused_kda_chunk,
fused_kda_chunk,
fused_kda_prologue,
)
from vllm.models.kimi_k3.amd.ops.third_party.kda import chunk_kda_with_fused_gate
from vllm.third_party.flash_linear_attention.ops.utils import FLA_CHUNK_SIZE
logger = init_logger(__name__)
def chunk_kda_prefill(
q: torch.Tensor,
k: torch.Tensor,
v: torch.Tensor,
raw_g: torch.Tensor,
raw_beta: torch.Tensor,
A_log: torch.Tensor,
g_bias: torch.Tensor | None = None,
scale: float | None = None,
initial_state: torch.Tensor | None = None,
output_final_state: bool = False,
lower_bound: float | None = None,
use_qk_l2norm_in_kernel: bool = False,
cu_seqlens: torch.Tensor | None = None,
chunk_indices: torch.Tensor | None = None,
chunk_offsets: torch.Tensor | None = None,
use_fused_chunk: bool = False,
out: torch.Tensor | None = None,
) -> tuple[torch.Tensor, torch.Tensor | None]:
"""Run chunk KDA from raw gate and beta projections.
Args:
use_fused_chunk: request the two-kernel ROCm path. It is used only when
every one of its preconditions holds; otherwise the Triton path
runs unchanged.
out: buffer the result must land in. Honoured by both backends, so the
caller can hand in a slice of its own output and skip a copy.
Returns:
The output and, when requested, the final recurrent state.
"""
if scale is None:
scale = k.shape[-1] ** -0.5
# The fused prologue folds the q/k L2 norm and the gate activation in, so it
# needs the raw projections and a bounded gate rather than the pre-normalized
# tensors the Triton path takes.
fused = (
use_fused_chunk
and use_qk_l2norm_in_kernel
and cu_seqlens is not None
and lower_bound is not None
and g_bias is not None
and can_use_fused_kda_chunk(k.shape[-1], v.shape[-1], k.dtype, FLA_CHUNK_SIZE)
)
if fused:
# 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
logger.info_once(
"Kimi-K3 KDA prefill: dispatching the fused ROCm chunk kernel."
)
ws = fused_kda_prologue(
q=q,
k=k,
v=v,
raw_g=raw_g,
raw_beta=raw_beta,
A_log=A_log,
dt_bias=g_bias,
scale=scale,
lower_bound=lower_bound,
cu_seqlens=cu_seqlens,
chunk_indices=chunk_indices,
)
return fused_kda_chunk(
qg=ws["qg"],
w=ws["w"],
u=ws["u"],
kg_t=ws["kg_t"],
aqk=ws["aqk"],
decay=ws["decay"],
# v is dead by this point, so it doubles as the output buffer.
out=out if out is not None else v,
scale=scale,
cu_seqlens=cu_seqlens,
initial_state=initial_state,
output_final_state=output_final_state,
chunk_offsets=chunk_offsets,
)
o, final_state = chunk_kda_with_fused_gate(
q=q,
k=k,
v=v,
raw_g=raw_g,
raw_beta=raw_beta,
A_log=A_log,
g_bias=g_bias,
scale=scale,
initial_state=initial_state,
output_final_state=output_final_state,
lower_bound=lower_bound,
use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel,
cu_seqlens=cu_seqlens,
chunk_indices=chunk_indices,
chunk_offsets=chunk_offsets,
)
if out is not None and o.data_ptr() != out.data_ptr():
out.copy_(o)
o = out
return o, final_state