[Kimi-K3][AMD] Fuse AttnRes state updates and norms (#50593)
Signed-off-by: Liuyinfeng01 <199041580+LiuYinfeng01@users.noreply.github.com> Co-authored-by: Liuyinfeng01 <199041580+LiuYinfeng01@users.noreply.github.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
yinfengLiu committed
Aug 5, 2026 at 00:32 UTC
7ac2ec758208a27bbe0c4155a136969463e8e6c8
3 files changed
+208
-49
tests/models/kimi_k3/test_amd_attn_res.py
+91
@@ -88,11 +88,15 @@ def test_amd_attn_res_matches_reference(
88
89
actual = attn_res(
90
prefix,
91
+ None,
92
blocks,
93
norm_weight,
94
qk_weight,
95
+ None,
96
num_blocks,
97
+ -1,
98
eps,
99
+ 0.0,
100
)
101
102
torch.testing.assert_close(actual, expected, atol=8e-2, rtol=3e-2)
@@ -100,3 +104,90 @@ def test_amd_attn_res_matches_reference(
104
torch.testing.assert_close(blocks, original_blocks, atol=0, rtol=0)
105
assert actual.shape == prefix.shape
106
assert actual.is_contiguous()
107
+
108
+
109
+@pytest.mark.parametrize(
110
+ (
111
+ "num_tokens",
112
+ "num_blocks",
113
+ "hidden_size",
114
+ "has_delta",
115
+ "write_block",
116
+ "apply_output_norm",
117
+ ),
118
+ [
119
+ pytest.param(1, 0, 128, False, True, True, id="empty-write-norm"),
120
+ pytest.param(7, 1, 1024, True, False, True, id="single-add-norm"),
121
+ pytest.param(17, 5, 7168, True, True, True, id="padded-write-add"),
122
+ pytest.param(3, 8, 7168, True, False, True, id="full-add-norm"),
123
+ pytest.param(320, 4, 7168, True, False, False, id="prefill-add"),
124
+ ],
125
+)
126
+def test_amd_attn_res_fused_contract(
127
+ num_tokens: int,
128
+ num_blocks: int,
129
+ hidden_size: int,
130
+ has_delta: bool,
131
+ write_block: bool,
132
+ apply_output_norm: bool,
133
+) -> None:
134
+ torch.manual_seed(42)
135
+ eps = 1e-5
136
+ output_eps = 2e-5
137
+ block_capacity = 9
138
+ prefix = _randn_with_row_padding(num_tokens, hidden_size, padding=7)
139
+ delta = (
140
+ _randn_with_row_padding(num_tokens, hidden_size, padding=11)
141
+ if has_delta
142
+ else None
143
+ )
144
+ blocks = _randn_with_row_padding(
145
+ num_tokens, block_capacity, hidden_size, padding=13
146
+ )
147
+ norm_weight = 1 + 0.1 * torch.randn(
148
+ hidden_size, device="cuda", dtype=torch.bfloat16
149
+ )
150
+ qk_weight = (
151
+ torch.randn(hidden_size, device="cuda", dtype=torch.bfloat16) / hidden_size**0.5
152
+ )
153
+ output_norm_weight = (
154
+ 1 + 0.1 * torch.randn(hidden_size, device="cuda", dtype=torch.bfloat16)
155
+ if apply_output_norm
156
+ else None
157
+ )
158
+ expected_prefix = prefix.clone()
159
+ if delta is not None:
160
+ expected_prefix = expected_prefix + delta
161
+ values = torch.cat(
162
+ (blocks[:, :num_blocks].clone(), expected_prefix.unsqueeze(1)), dim=1
163
+ )
164
+ keys = F.rms_norm(values.float(), (hidden_size,), norm_weight.float(), eps)
165
+ probs = (keys @ qk_weight.float()).softmax(dim=-1)
166
+ expected = torch.matmul(probs.unsqueeze(1), values.float()).squeeze(1)
167
+ if output_norm_weight is not None:
168
+ expected = F.rms_norm(
169
+ expected, (hidden_size,), output_norm_weight.float(), output_eps
170
+ )
171
+ expected = expected.to(prefix.dtype)
172
+ original_blocks = blocks.clone()
173
+ block_write_idx = num_blocks if write_block else -1
174
+
175
+ actual = attn_res(
176
+ prefix,
177
+ delta,
178
+ blocks,
179
+ norm_weight,
180
+ qk_weight,
181
+ output_norm_weight,
182
+ num_blocks,
183
+ block_write_idx,
184
+ eps,
185
+ output_eps,
186
+ )
187
+
188
+ torch.testing.assert_close(actual, expected, atol=8e-2, rtol=3e-2)
189
+ torch.testing.assert_close(prefix, expected_prefix, atol=0, rtol=0)
190
+ if write_block:
191
+ original_blocks[:, block_write_idx].copy_(expected_prefix)
192
+ torch.testing.assert_close(blocks, original_blocks, atol=0, rtol=0)
193
+ assert actual.is_contiguous()
vllm/models/kimi_k3/amd/linear.py
+16
-9
@@ -141,17 +141,22 @@ def _apply_attn_res(
141
proj: ReplicatedLinear,
142
norm: RMSNorm,
143
num_valid_blocks: int,
144
+ *,
145
+ delta: torch.Tensor | None = None,
146
+ output_norm: RMSNorm | None = None,
147
+ block_write_idx: int = -1,
148
) -> torch.Tensor:
145
- if num_valid_blocks <= 0:
146
- return prefix_sum
147
-
149
return attn_res(
150
prefix_sum,
151
+ delta,
152
block_residual,
153
norm.weight,
154
proj.weight.squeeze(0),
155
+ None if output_norm is None else output_norm.weight,
156
num_valid_blocks,
157
+ block_write_idx,
158
norm.variance_epsilon,
159
+ 0.0 if output_norm is None else output_norm.variance_epsilon,
160
)
161
162
@@ -610,19 +615,20 @@ class KimiDecoderLayer(nn.Module):
615
self.self_attention_res_proj,
616
self.self_attention_res_norm,
617
self.prev_valid_blocks,
618
+ output_norm=self.input_layernorm,
619
+ block_write_idx=(self.block_write_idx if self.is_block_write_layer else -1),
620
)
621
622
if self.is_block_write_layer:
616
- block_residual[:, self.block_write_idx, :].copy_(prefix_sum)
623
prefix_sum = None
624
619
- hidden_states = self.input_layernorm(hidden_states)
625
hidden_states = self._run_self_attn(positions, hidden_states)
626
622
- if prefix_sum is not None:
623
- prefix_sum = prefix_sum + hidden_states
624
- else:
627
+ if prefix_sum is None:
628
prefix_sum = hidden_states
629
+ prefix_delta = None
630
+ else:
631
+ prefix_delta = hidden_states
632
633
mlp_valid_blocks = self.prev_valid_blocks + (
634
1 if self.is_block_write_layer else 0
@@ -633,9 +639,10 @@ class KimiDecoderLayer(nn.Module):
639
self.mlp_res_proj,
640
self.mlp_res_norm,
641
mlp_valid_blocks,
642
+ delta=prefix_delta,
643
+ output_norm=self.post_attention_layernorm,
644
)
645
638
- hidden_states = self.post_attention_layernorm(hidden_states)
646
hidden_states = self.mlp(hidden_states)
647
prefix_sum = prefix_sum + hidden_states
648
return prefix_sum, block_residual
vllm/models/kimi_k3/amd/ops/attn_res.py
+101
-40
@@ -7,6 +7,7 @@
7
# the following copyright notice:
8
# Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li
9
10
+
11
import torch
12
13
from vllm.triton_utils import tl, triton
@@ -15,17 +16,25 @@ from vllm.triton_utils import tl, triton
16
@triton.jit
17
def _attn_res_kernel(
18
prefix_ptr,
19
+ delta_ptr,
20
blocks_ptr,
21
norm_weight_ptr,
22
qk_weight_ptr,
23
+ output_norm_weight_ptr,
24
output_ptr,
25
stride_prefix_m: tl.constexpr,
26
+ stride_delta_m: tl.constexpr,
27
stride_block_m: tl.constexpr,
28
stride_block_r: tl.constexpr,
29
stride_output_m: tl.constexpr,
30
num_blocks: tl.constexpr,
31
hidden_size: tl.constexpr,
32
+ block_write_idx: tl.constexpr,
33
eps: tl.constexpr,
34
+ output_norm_eps: tl.constexpr,
35
+ HAS_DELTA: tl.constexpr,
36
+ WRITE_BLOCK: tl.constexpr,
37
+ APPLY_OUTPUT_NORM: tl.constexpr,
38
BLOCK_L: tl.constexpr,
39
BLOCK_D: tl.constexpr,
40
):
@@ -33,51 +42,93 @@ def _attn_res_kernel(
42
d_offsets = tl.max_contiguous(tl.arange(0, BLOCK_D), BLOCK_D)
43
d_mask = d_offsets < hidden_size
44
36
- prefix = tl.load(
45
+ updated_prefix = tl.load(
46
prefix_ptr + row_idx * stride_prefix_m + d_offsets,
47
mask=d_mask,
48
other=0.0,
49
).to(tl.float32)
41
- input_qk_weight = tl.load(norm_weight_ptr + d_offsets, mask=d_mask, other=0.0).to(
42
- tl.float32
43
- ) * tl.load(qk_weight_ptr + d_offsets, mask=d_mask, other=0.0).to(tl.float32)
44
-
45
- max_logit = tl.full((), -float("inf"), tl.float32)
46
- denominator = tl.zeros((), tl.float32)
47
- mixed = tl.zeros((BLOCK_D,), tl.float32)
48
- num_sources = num_blocks + 1
49
-
50
- for source_tile in range(tl.cdiv(num_sources, BLOCK_L)):
51
- source_offsets = source_tile * BLOCK_L + tl.arange(0, BLOCK_L)
52
- source_mask = source_offsets < num_sources
53
- is_prefix = source_offsets == num_blocks
54
- block_ptrs = (
50
+ if HAS_DELTA:
51
+ delta = tl.load(
52
+ delta_ptr + row_idx * stride_delta_m + d_offsets,
53
+ mask=d_mask,
54
+ other=0.0,
55
+ ).to(tl.float32)
56
+ updated_prefix += delta
57
+ # Match the BF16 prefix-add result before using it as a residual source.
58
+ updated_prefix = updated_prefix.to(prefix_ptr.dtype.element_ty).to(tl.float32)
59
+ tl.store(
60
+ prefix_ptr + row_idx * stride_prefix_m + d_offsets,
61
+ updated_prefix,
62
+ mask=d_mask,
63
+ )
64
+ if WRITE_BLOCK:
65
+ tl.store(
66
blocks_ptr
67
+ row_idx * stride_block_m
57
- + source_offsets[:, None] * stride_block_r
58
- + d_offsets[None, :]
68
+ + block_write_idx * stride_block_r
69
+ + d_offsets,
70
+ updated_prefix,
71
+ mask=d_mask,
72
)
60
- block_values = tl.load(
61
- block_ptrs,
62
- mask=(source_mask[:, None] & ~is_prefix[:, None] & d_mask[None, :]),
63
- other=0.0,
64
- eviction_policy="evict_first",
73
+ # With only the prefix source, the AttnRes softmax is exactly one.
74
+ if num_blocks == 0:
75
+ mixed = updated_prefix
76
+ else:
77
+ # Reloading avoids keeping the full prefix vector live across the loop.
78
+ if HAS_DELTA:
79
+ tl.debug_barrier()
80
+ input_qk_weight = tl.load(
81
+ norm_weight_ptr + d_offsets, mask=d_mask, other=0.0
82
+ ).to(tl.float32) * tl.load(
83
+ qk_weight_ptr + d_offsets, mask=d_mask, other=0.0
84
).to(tl.float32)
66
- values = tl.where(is_prefix[:, None], prefix[None, :], block_values)
67
- reciprocal_std = tl.rsqrt(
68
- tl.sum(values * values, axis=1) * (1.0 / hidden_size) + eps
69
- )
70
- logits = tl.sum(values * input_qk_weight[None, :], axis=1) * reciprocal_std
71
- scores = tl.where(source_mask, logits, -float("inf"))
85
+ max_logit = tl.full((), -float("inf"), tl.float32)
86
+ denominator = tl.zeros((), tl.float32)
87
+ mixed = tl.zeros((BLOCK_D,), tl.float32)
88
+
89
+ num_sources = num_blocks + 1
90
+ for source_tile in range(tl.cdiv(num_sources, BLOCK_L)):
91
+ source_offsets = source_tile * BLOCK_L + tl.arange(0, BLOCK_L)
92
+ source_mask = source_offsets < num_sources
93
+ is_prefix = source_offsets == num_blocks
94
+ block_ptrs = (
95
+ blocks_ptr
96
+ + row_idx * stride_block_m
97
+ + source_offsets[:, None] * stride_block_r
98
+ + d_offsets[None, :]
99
+ )
100
+ block_values = tl.load(
101
+ block_ptrs,
102
+ mask=(source_mask[:, None] & ~is_prefix[:, None] & d_mask[None, :]),
103
+ other=0.0,
104
+ eviction_policy="evict_first",
105
+ ).to(tl.float32)
106
+ values = tl.where(is_prefix[:, None], updated_prefix[None, :], block_values)
107
+ reciprocal_std = tl.rsqrt(
108
+ tl.sum(values * values, axis=1) * (1.0 / hidden_size) + eps
109
+ )
110
+ logits = tl.sum(values * input_qk_weight[None, :], axis=1) * reciprocal_std
111
+ scores = tl.where(source_mask, logits, -float("inf"))
112
+
113
+ new_max_logit = tl.maximum(max_logit, tl.max(scores, axis=0))
114
+ old_scale = tl.exp(max_logit - new_max_logit)
115
+ block_scales = tl.exp(scores - new_max_logit)
116
+ denominator = denominator * old_scale + tl.sum(block_scales, axis=0)
117
+ mixed = mixed * old_scale + tl.sum(block_scales[:, None] * values, axis=0)
118
+ max_logit = new_max_logit
119
73
- new_max_logit = tl.maximum(max_logit, tl.max(scores, axis=0))
74
- old_scale = tl.exp(max_logit - new_max_logit)
75
- block_scales = tl.exp(scores - new_max_logit)
76
- denominator = denominator * old_scale + tl.sum(block_scales, axis=0)
77
- mixed = mixed * old_scale + tl.sum(block_scales[:, None] * values, axis=0)
78
- max_logit = new_max_logit
120
+ mixed /= denominator
121
+ output = mixed
122
80
- output = mixed / denominator
123
+ if APPLY_OUTPUT_NORM:
124
+ output_reciprocal_std = tl.rsqrt(
125
+ tl.sum(tl.where(d_mask, mixed * mixed, 0.0), axis=0) * (1.0 / hidden_size)
126
+ + output_norm_eps
127
+ )
128
+ output_norm_weight = tl.load(
129
+ output_norm_weight_ptr + d_offsets, mask=d_mask, other=0.0
130
+ ).to(tl.float32)
131
+ output = mixed * output_reciprocal_std * output_norm_weight
132
tl.store(
133
output_ptr + row_idx * stride_output_m + d_offsets,
134
output,
@@ -87,43 +138,53 @@ def _attn_res_kernel(
138
139
def attn_res(
140
prefix: torch.Tensor,
141
+ delta: torch.Tensor | None,
142
blocks: torch.Tensor,
143
norm_weight: torch.Tensor,
144
qk_weight: torch.Tensor,
145
+ output_norm_weight: torch.Tensor | None,
146
num_blocks: int,
147
+ block_write_idx: int,
148
eps: float,
149
+ output_norm_eps: float,
150
) -> torch.Tensor:
151
num_tokens, hidden_size = prefix.shape
97
- assert 0 < num_blocks <= blocks.shape[1]
98
- assert blocks.shape[0] == num_tokens
99
- assert norm_weight.numel() == hidden_size
100
- assert qk_weight.numel() == hidden_size
152
assert prefix.stride(-1) == 1
153
+ assert delta is None or delta.stride(-1) == 1
154
assert blocks.stride(-1) == 1
155
assert norm_weight.stride(-1) == 1
156
assert qk_weight.stride(-1) == 1
105
-
157
+ assert output_norm_weight is None or output_norm_weight.stride(-1) == 1
158
output = prefix.new_empty(prefix.shape)
159
if num_tokens == 0:
160
return output
161
162
+ # Source tiling helps decode, while one-source tiles scale better for prefill.
163
if num_tokens >= 256 or num_blocks <= 1:
164
block_l, num_warps = 1, 4
165
else:
166
block_l, num_warps = 4, 8
167
_attn_res_kernel[(num_tokens,)](
168
prefix,
169
+ delta,
170
blocks,
171
norm_weight,
172
qk_weight,
173
+ output_norm_weight,
174
output,
175
prefix.stride(0),
176
+ 0 if delta is None else delta.stride(0),
177
blocks.stride(0),
178
blocks.stride(1),
179
output.stride(0),
180
num_blocks,
181
hidden_size,
182
+ block_write_idx,
183
eps,
184
+ output_norm_eps,
185
+ HAS_DELTA=delta is not None,
186
+ WRITE_BLOCK=block_write_idx >= 0,
187
+ APPLY_OUTPUT_NORM=output_norm_weight is not None,
188
BLOCK_L=block_l,
189
BLOCK_D=triton.next_power_of_2(hidden_size),
190
num_warps=num_warps,