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

[ROCm] attn res + sigmoid_mul + conv fusions (#52968)

Browse Source 
Signed-off-by: Omar Muhammad <omar.muhammad@amd.com>
Signed-off-by: Nicholas Susanto <nicholas.susanto@amd.com>
Co-authored-by: Nicholas Susanto <nicholas.susanto@amd.com>
Co-authored-by: Douglas Lehr <91553416+dllehr-amd@users.noreply.github.com>
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
This commit is contained in:
Omar Muhammad
2026-10-17 00:54:54 +00:00
committed by GitHub
co-authored by Nicholas Susanto Douglas Lehr mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
parent 783380af39
commit f019e3e105
3 changed files with 219 additions and 35 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/models/kimi_k3/test_amd_kda_fused_conv1d.py
+163
View File
@@ -0,0 +1,163 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import pytest
import torch
import torch.nn.functional as F
from vllm.model_executor.layers.mamba.ops.causal_conv1d import causal_conv1d_fn
from vllm.platforms import current_platform
from vllm.v1.attention.backends.utils import NULL_BLOCK_ID
pytestmark = pytest.mark.skipif(
not current_platform.is_rocm(), reason="AITER fused conv1d is ROCm-only"
)
WIDTH = 4 # Kimi-K3 short_conv_kernel_size
STATE_LEN = WIDTH - 1
HEAD_DIM = 128
LOCAL_HEADS = 4 # per-TP-rank heads; projection = 512, packed dim = 1536
def _fused_fn():
gdn = pytest.importorskip("aiter.ops.triton.gated_delta_net")
return gdn.causal_conv1d_split_qkv_triton_tile_fn
def _make_inputs(seqlens, has_initial_state, seed=0, device="cuda"):
torch.manual_seed(seed)
proj = HEAD_DIM * LOCAL_HEADS
dim = 3 * proj
total = sum(seqlens)
n_seqs = len(seqlens)
# Token-major activations, as produced by the in_proj GEMM.
mixed_qkv = torch.randn(total, dim, dtype=torch.bfloat16, device=device)
weight = torch.randn(dim, WIDTH, dtype=torch.bfloat16, device=device) * 0.3
# More cache lines than sequences, assigned in a shuffled order. Line 0 is
# vLLM's null block (NULL_BLOCK_ID), which causal_conv1d_fn skips, so real
# sequences never map to it.
n_lines = n_seqs + 2
conv_state = torch.randn(
n_lines, dim, STATE_LEN, dtype=torch.bfloat16, device=device
)
cache_indices = (torch.randperm(n_lines - 1, device=device)[:n_seqs] + 1).to(
torch.int32
)
query_start_loc = torch.tensor(
[0] + list(torch.tensor(seqlens).cumsum(0).tolist()),
dtype=torch.int32,
device=device,
)
has_initial = torch.full(
(n_seqs,), has_initial_state, dtype=torch.bool, device=device
)
return mixed_qkv, weight, conv_state, cache_indices, query_start_loc, has_initial
def _reference(
mixed_qkv, weight, conv_state, cache_indices, query_start_loc, has_initial
):
"""fp32 per-sequence depthwise causal conv + SiLU.
Returns (out[T, dim], new_conv_state).
"""
x = mixed_qkv.float().t() # (dim, T)
w = weight.float()
dim = x.shape[0]
out = torch.empty_like(x)
new_state = conv_state.float().clone()
qsl = query_start_loc.tolist()
for s in range(len(qsl) - 1):
start, end = qsl[s], qsl[s + 1]
ci = int(cache_indices[s])
xs = x[:, start:end]
if bool(has_initial[s]):
xin = torch.cat([conv_state[ci].float(), xs], dim=1)
else:
xin = F.pad(xs, (STATE_LEN, 0))
y = F.conv1d(xin.unsqueeze(0), w.unsqueeze(1), None, groups=dim)[0]
out[:, start:end] = F.silu(y)
new_state[ci] = xin[:, -STATE_LEN:]
return out.t(), new_state
def _three_launch(
mixed_qkv, weight, conv_state, cache_indices, query_start_loc, has_initial
):
"""The pre-fusion KDA prefill path: one causal_conv1d_fn per projection."""
proj = mixed_qkv.shape[1] // 3
outs = []
for x, state, w in zip(
mixed_qkv.split(proj, dim=-1),
conv_state.split(proj, dim=-2),
weight.split(proj, dim=0),
):
outs.append(
causal_conv1d_fn(
x.transpose(0, 1),
w,
None,
activation="silu",
conv_states=state,
has_initial_state=has_initial,
cache_indices=cache_indices,
query_start_loc=query_start_loc,
).transpose(0, 1)
)
return outs
@pytest.mark.parametrize(
"seqlens", [[1024], [8, 1, 300], [64, 64, 64], [513, 64, 1, 1200]]
)
@pytest.mark.parametrize("has_initial_state", [False, True])
@pytest.mark.parametrize("block_m", [8, 128])
def test_fused_split_qkv_conv1d_matches_reference(seqlens, has_initial_state, block_m):
fused_fn = _fused_fn()
mixed_qkv, weight, conv_state, cache_indices, qsl, has_initial = _make_inputs(
seqlens, has_initial_state
)
proj = mixed_qkv.shape[1] // 3
ref_out, ref_state = _reference(
mixed_qkv, weight, conv_state, cache_indices, qsl, has_initial
)
state_3 = conv_state.clone()
q3, k3, v3 = _three_launch(
mixed_qkv, weight, state_3, cache_indices, qsl, has_initial
)
state_f = conv_state.clone()
# Exactly the call made by KimiK3DeltaAttention._forward.
qf, kf, vf = fused_fn(
mixed_qkv.transpose(0, 1),
weight,
None,
state_f,
qsl,
proj,
proj,
cache_indices=cache_indices,
has_initial_state=has_initial,
activation="silu",
block_m=block_m,
metadata=None,
preserve_input_dtype=True,
)
fused = torch.cat([qf, kf, vf], dim=-1)
assert fused.shape == mixed_qkv.shape and fused.dtype == mixed_qkv.dtype
assert all(t.is_contiguous() for t in (qf, kf, vf))
tol = dict(atol=2e-2, rtol=2e-2)
torch.testing.assert_close(fused.float(), ref_out, **tol)
torch.testing.assert_close(fused, torch.cat([q3, k3, v3], dim=-1), **tol)
# conv_state writeback: assigned lines updated, the others untouched.
torch.testing.assert_close(state_f.float(), ref_state, **tol)
torch.testing.assert_close(state_f, state_3, **tol)
untouched = torch.ones(conv_state.shape[0], dtype=torch.bool, device="cuda")
untouched[cache_indices.long()] = False
assert untouched[NULL_BLOCK_ID]
assert torch.equal(state_f[untouched], conv_state[untouched])
vllm/models/kimi_k3/amd/kda.py
+47 -32
View File
@@ -134,6 +134,14 @@ class KimiK3DeltaAttention(GatedDeltaNetAttention):
self.conv_size = kda_config["short_conv_kernel_size"]
self.use_full_rank_gate = True
self._fused_qkv_conv_fn = None
if rocm_aiter_ops.is_enabled():
from aiter.ops.triton.gated_delta_net import (
causal_conv1d_split_qkv_triton_tile_fn,
)
self._fused_qkv_conv_fn = causal_conv1d_split_qkv_triton_tile_fn
# Keep f_a before the narrow beta shard, then pad each TP-local row to
# select the aligned BF16 GEMM path. The padding also avoids an Inductor
# correctness issue seen with the row-strided G view.
@@ -580,35 +588,48 @@ class KimiK3DeltaAttention(GatedDeltaNetAttention):
if mixed_qkv_ns is not None:
assert g1_ns is not None and beta_ns is not None
if m.num_prefills > 0:
q_ns, k_ns, v_ns = mixed_qkv_ns.split(
self.local_projection_size, dim=-1
)
# Packed prefill conv would require copying V solely to make
# it dense for KDA. Separate calls accept the strided inputs
# and produce dense Q/K/V without that extra traffic.
# TODO: Use packed conv once every KDA prefill backend accepts
# row-strided Q/K/V directly.
def _prefill_conv(
x: torch.Tensor,
state: torch.Tensor,
weight: torch.Tensor,
) -> torch.Tensor:
return causal_conv1d_fn(
x.transpose(0, 1),
weight,
if self._fused_qkv_conv_fn is not None:
# One launch (fused qkv) instead of 3
q_ns, k_ns, v_ns = self._fused_qkv_conv_fn(
mixed_qkv_ns.transpose(0, 1),
conv_weights,
None,
activation="silu",
conv_states=state,
has_initial_state=has_initial_state,
conv_state,
non_spec_query_start_loc,
self.local_projection_size,
self.local_projection_size,
cache_indices=non_spec_state_indices_tensor,
query_start_loc=non_spec_query_start_loc,
has_initial_state=has_initial_state,
activation="silu",
block_m=8,
metadata=m,
).transpose(0, 1)
preserve_input_dtype=True,
)
else:
q_ns, k_ns, v_ns = mixed_qkv_ns.split(
self.local_projection_size, dim=-1
)
q_ns = _prefill_conv(q_ns, q_conv_state, q_conv_weight)
k_ns = _prefill_conv(k_ns, k_conv_state, k_conv_weight)
v_ns = _prefill_conv(v_ns, v_conv_state, v_conv_weight)
def _prefill_conv(
x: torch.Tensor,
state: torch.Tensor,
weight: torch.Tensor,
) -> torch.Tensor:
return causal_conv1d_fn(
x.transpose(0, 1),
weight,
None,
activation="silu",
conv_states=state,
has_initial_state=has_initial_state,
cache_indices=non_spec_state_indices_tensor,
query_start_loc=non_spec_query_start_loc,
metadata=m,
).transpose(0, 1)
q_ns = _prefill_conv(q_ns, q_conv_state, q_conv_weight)
k_ns = _prefill_conv(k_ns, k_conv_state, k_conv_weight)
v_ns = _prefill_conv(v_ns, v_conv_state, v_conv_weight)
q_ns, k_ns, v_ns = (
rearrange(x, "n (h d) -> 1 n h d", d=self.head_dim)
for x in (q_ns, k_ns, v_ns)
@@ -752,15 +773,9 @@ class KimiK3DeltaAttention(GatedDeltaNetAttention):
# ---------- merge spec and non-spec outputs ----------
if core_attn_out_spec is not None and core_attn_out_non_spec is not None:
# Mixed batches require indexed placement in the original order.
merged = torch.empty(
(1, num_actual_tokens, *core_attn_out_spec.shape[2:]),
dtype=core_attn_out_spec.dtype,
device=core_attn_out_spec.device,
)
merged = core_attn_out[:, :num_actual_tokens]
merged.index_copy_(1, spec_token_indx, core_attn_out_spec)
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:
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[
vllm/models/kimi_k3/amd/mla.py
+9 -3
View File
@@ -23,8 +23,8 @@ _OPT_MIN_SIZE = 2048
class KimiK3MultiHeadLatentAttentionWrapper(MultiHeadLatentAttentionWrapper):
"""Kimi-K3 MLA wrapper with eager AITER q/kv RMSNorm fusion and a fused
decode Q-prep path."""
"""Kimi-K3 MLA wrapper with eager AITER q/kv RMSNorm and output gate
fusions and a fused decode Q-prep path."""
def __init__(self, *args, **kwargs) -> None:
super().__init__(*args, **kwargs)
@@ -303,6 +303,12 @@ class KimiK3MultiHeadLatentAttentionWrapper(MultiHeadLatentAttentionWrapper):
)
if self.g_proj is not None:
attn_out = attn_out * self.g_proj(hidden_states)[0].sigmoid()
gate = self.g_proj(hidden_states)[0]
if rocm_aiter_ops.is_enabled():
from aiter.ops.triton.fusions.fused_sigmoid_mul import fused_sigmoid_mul
attn_out = fused_sigmoid_mul(attn_out, gate)
else:
attn_out = attn_out * gate.sigmoid()
return self.o_proj(attn_out)[0]