diff --git a/tests/e2e/nightly/single_node/ops/singlecard_ops/triton/test_kimi_k3_fusions.py b/tests/e2e/nightly/single_node/ops/singlecard_ops/triton/test_kimi_k3_fusions.py new file mode 100644 index 000000000000..67d37349aaa8 --- /dev/null +++ b/tests/e2e/nightly/single_node/ops/singlecard_ops/triton/test_kimi_k3_fusions.py @@ -0,0 +1,97 @@ +# Copyright (c) 2026 Huawei Technologies Co., Ltd. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Numerical regression coverage for Kimi K3 attention residual fusion.""" + +from types import SimpleNamespace + +import pytest +import torch +import torch_npu # noqa: F401 +from vllm.triton_utils import HAS_TRITON + +if HAS_TRITON: + from vllm_ascend.ops.triton.kimi_k3.attention_residual import apply_attn_res + + +pytestmark = [ + pytest.mark.skipif(not HAS_TRITON, reason="Triton is not available"), + pytest.mark.skipif(not torch.npu.is_available(), reason="NPU required"), + pytest.mark.skip_global_cleanup, +] + + +@torch.inference_mode() +@pytest.mark.parametrize( + ("num_tokens", "num_blocks", "block_capacity"), + [ + pytest.param(7, 4, 7, id="partial-capacity"), + pytest.param(512, 8, 8, id="profile-shape"), + ], +) +def test_kimi_k3_attention_residual_triton_matches_reference( + num_tokens, + num_blocks, + block_capacity, +): + torch.manual_seed(1) + hidden_size = 7168 + eps = 1e-6 + prefix_sum = torch.randn( + (num_tokens, hidden_size), + dtype=torch.bfloat16, + device="npu", + ) + block_residual = torch.randn( + (num_tokens, block_capacity, hidden_size), + dtype=torch.bfloat16, + device="npu", + ) + projection = SimpleNamespace( + weight=torch.randn( + (1, hidden_size), + dtype=torch.bfloat16, + device="npu", + ) + ) + norm = SimpleNamespace( + weight=torch.randn( + (hidden_size,), + dtype=torch.bfloat16, + device="npu", + ), + variance_epsilon=eps, + ) + + actual = apply_attn_res( + prefix_sum, + block_residual, + projection, + norm, + num_blocks, + ) + + values = torch.cat((block_residual[:, :num_blocks, :], prefix_sum.unsqueeze(1)), dim=1).float() + normalized = values * torch.rsqrt(values.square().mean(dim=-1, keepdim=True) + eps) + score_weight = norm.weight.float() * projection.weight.squeeze(0).float() + scores = (normalized * score_weight).sum(dim=-1) + probabilities = scores.softmax(-1).unsqueeze(1) + expected = torch.matmul(probabilities, values).squeeze(1).to(prefix_sum.dtype) + + torch.testing.assert_close( + actual.cpu(), + expected.cpu(), + rtol=1e-2, + atol=1e-2, + ) diff --git a/vllm_ascend/ops/triton/kimi_k3/__init__.py b/vllm_ascend/ops/triton/kimi_k3/__init__.py new file mode 100644 index 000000000000..43e799bd0bd0 --- /dev/null +++ b/vllm_ascend/ops/triton/kimi_k3/__init__.py @@ -0,0 +1 @@ +"""Triton fusion kernels specific to Kimi K3.""" diff --git a/vllm_ascend/ops/triton/kimi_k3/attention_residual.md b/vllm_ascend/ops/triton/kimi_k3/attention_residual.md new file mode 100644 index 000000000000..5bc4a4cc608d --- /dev/null +++ b/vllm_ascend/ops/triton/kimi_k3/attention_residual.md @@ -0,0 +1,37 @@ +# Kimi K3 Attention Residual 算子说明 + +## 功能 + +`apply_attn_res` 将 Kimi K3 每个 token 的有效 block residual 与 +`prefix_sum` residual 做可学习的 softmax 加权融合。实现位于 +`vllm_ascend/ops/triton/kimi_k3/attention_residual.py`。 + +对每条 residual stream `v_s`,算子先计算 RMSNorm,再通过 +`norm.weight * proj.weight` 得到标量分数: + +```text +score_s = sum(RMSNorm(v_s) * norm.weight * proj.weight) +weight_s = softmax(score)_s +output = sum(weight_s * v_s) +``` + +## 输入与输出 + +| 参数 | 形状 | 说明 | +| --- | --- | --- | +| `prefix_sum` | `[num_tokens, hidden_size]` | 每个 token 的 prefix-sum residual,也是最后一条参与融合的 stream。 | +| `block_residual` | `[num_tokens, block_capacity, hidden_size]` | vLLM 预分配的 residual buffer。只有前 `num_valid_blocks` 个 block 已初始化。 | +| `proj` | `[1, hidden_size]` | 将归一化 residual 投影为标量分数的线性层。 | +| `norm` | `[hidden_size]` | RMSNorm 权重及 epsilon。 | +| `num_valid_blocks` | `int` | `block_residual` 中有效 block 的数量。 | +| 返回值 | `[num_tokens, hidden_size]` | 所有有效 residual stream 的加权和。 | + +## 实现约束 + +- kernel 的 `B` 等于 `num_valid_blocks`,`BLOCK_CAPACITY` 来自预分配 + buffer 的第二维;kernel 只读取 `[0, B)`,不会读取未初始化容量。 +- `prefix_sum` 使用逻辑索引 `s == B`。启动端将 stream 数设置为 + `next_power_of_2(B + 1)`,因此 `NB` 始终覆盖 `B` 个 block residual + 加一条 prefix stream。 +- softmax 计算使用 FP32,输出再转换为 `prefix_sum` 的 dtype。 +- 启动 grid 使用设备 vector core 数,每个 program 处理一段连续 token。 diff --git a/vllm_ascend/ops/triton/kimi_k3/attention_residual.py b/vllm_ascend/ops/triton/kimi_k3/attention_residual.py new file mode 100644 index 000000000000..0a666fd766f4 --- /dev/null +++ b/vllm_ascend/ops/triton/kimi_k3/attention_residual.py @@ -0,0 +1,116 @@ +"""Fused Kimi K3 attention-residual mixture. + +For every token, the operator RMS-normalizes each valid block residual and the +prefix-sum residual, projects them to scalar scores, applies a softmax across +those streams, and returns their weighted sum. ``block_residual`` follows +vLLM's preallocated ``[num_tokens, block_capacity, hidden_size]`` contract; +``num_valid_blocks`` identifies the initialized prefix of that capacity. +""" + +import torch +from vllm.triton_utils import tl, triton + +from vllm_ascend.ops.triton.triton_utils import ( + get_vectorcore_num, + init_device_properties_triton, +) + + +@triton.jit +def _apply_attn_res_kernel( + block_residual_ptr, + prefix_sum_ptr, + norm_w_ptr, + proj_w_ptr, + out_ptr, + N: tl.constexpr, + H: tl.constexpr, + B: tl.constexpr, + BLOCK_CAPACITY: tl.constexpr, + EPS: tl.constexpr, + NUM_CORES: tl.constexpr, + NB: tl.constexpr, +): + tl.static_assert(NB >= B + 1, "NB must include all block residuals and prefix_sum") + block_size = (N - 1) // NUM_CORES + 1 + pid = tl.program_id(0) + tok0 = pid * block_size + if tok0 >= N: + return + tok1 = tl.minimum(tok0 + block_size, N) + + cols = tl.arange(0, H) + s_idx = tl.arange(0, NB) + block_residual_stride = BLOCK_CAPACITY * H + + norm_w = tl.load(norm_w_ptr + cols).to(tl.float32) + proj_w = tl.load(proj_w_ptr + cols).to(tl.float32) + w = norm_w * proj_w + + for tok in range(tok0, tok1): + scores = tl.full([NB], -float("inf"), dtype=tl.float32) + for s in range(B + 1): + if s < B: + v = tl.load(block_residual_ptr + tok * block_residual_stride + s * H + cols).to(tl.float32) + else: + v = tl.load(prefix_sum_ptr + tok * H + cols).to(tl.float32) + ms = tl.sum(v * v) / H + rstd = tl.rsqrt(ms + EPS) + k = v * rstd + scores = tl.where(s_idx == s, tl.sum(k * w), scores) + + scores_max = tl.max(scores) + exp_scores = tl.exp(scores - scores_max) + weights = exp_scores / tl.sum(exp_scores) + + out = tl.zeros([H], dtype=tl.float32) + for s in range(B + 1): + if s < B: + v = tl.load(block_residual_ptr + tok * block_residual_stride + s * H + cols).to(tl.float32) + else: + v = tl.load(prefix_sum_ptr + tok * H + cols).to(tl.float32) + w_s = tl.sum(tl.where(s_idx == s, weights, 0.0)) + out += w_s * v + + tl.store(out_ptr + tok * H + cols, out.to(out_ptr.dtype.element_ty)) + + +def apply_attn_res( + prefix_sum: torch.Tensor, + block_residual: torch.Tensor, + proj: torch.nn.Module, + norm: torch.nn.Module, + num_valid_blocks: int, +) -> torch.Tensor: + """Return K3's learned softmax mixture of residual streams.""" + num_tokens, hidden_size = prefix_sum.shape + block_capacity = block_residual.shape[1] + proj_w = proj.weight.squeeze(0) + norm_w = norm.weight + eps = norm.variance_epsilon + + out = torch.empty( + (num_tokens, hidden_size), + dtype=prefix_sum.dtype, + device=prefix_sum.device, + ) + # The extra stream is prefix_sum, so NB must cover num_valid_blocks + 1. + num_streams = triton.next_power_of_2(num_valid_blocks + 1) + init_device_properties_triton() + num_vectorcore = get_vectorcore_num() + _apply_attn_res_kernel[(num_vectorcore,)]( + block_residual, + prefix_sum, + norm_w, + proj_w, + out, + N=num_tokens, + H=hidden_size, + B=num_valid_blocks, + BLOCK_CAPACITY=block_capacity, + EPS=eps, + NUM_CORES=num_vectorcore, + NB=num_streams, + multibuffer=True, + ) + return out