[Performance][Ops] Add Kimi K3 attention residual fusion

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,
+    )

vllm_ascend/ops/triton/kimi_k3/__init__.py

@@ -0,0 +1 @@
+"""Triton fusion kernels specific to Kimi K3."""

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。

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