« prev ^ index » next coverage.py v7.13.1, created at 2026-11-15 07:18 +0800
1# Copyright 2026 Huawei Technologies Co., Ltd
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7# http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14# ============================================================================
15"""RD pointwise kernels; matrix multiplication stays on the existing path."""
16import triton
17import triton.language as tl
20@triton.jit
21def kda_rd_finish_round_kernel(
22 product: tl.tensor, previous: tl.tensor, matrix: tl.tensor,
23 row_count: tl.constexpr, state_width: tl.constexpr, packed_width: tl.constexpr,
24 block_rows: tl.constexpr, save_matrix: tl.constexpr,
25) -> None:
26 """Add the previous state and write a compact transfer for later rounds.
28 Args:
29 product: Fresh contiguous FP32 packed S/M product pointer.
30 previous: Previous packed S/M pointer; only S is read.
31 matrix: Compact M output pointer, unused when save_matrix is false.
32 row_count: Number of flattened matrix rows.
33 state_width: Columns in each compact S or M matrix.
34 packed_width: Columns in the concatenated S/M row.
35 block_rows: Flattened matrix rows per program, unrelated to token chunk size.
36 save_matrix: Whether a later round consumes compact M.
37 """
38 rows = tl.program_id(0) * block_rows + tl.arange(0, block_rows)
39 cols = tl.arange(0, state_width)
40 offsets = rows[:, None] * packed_width + cols[None, :]
41 left = tl.load(product + offsets, rows[:, None] < row_count, 0)
42 old = tl.load(previous + offsets, rows[:, None] < row_count, 0)
43 tl.store(product + offsets, left + old, rows[:, None] < row_count)
44 if save_matrix:
45 right = tl.load(product + offsets + state_width, rows[:, None] < row_count, 0)
46 tl.store(matrix + rows[:, None] * state_width + cols[None, :], right, rows[:, None] < row_count)
49@triton.jit
50def kda_rd_finish_to_state_kernel(
51 product: tl.tensor, previous: tl.tensor, matrix: tl.tensor, state: tl.tensor,
52 row_count: tl.constexpr, state_width: tl.constexpr, packed_width: tl.constexpr,
53 block_rows: tl.constexpr, save_matrix: tl.constexpr,
54) -> None:
55 """Also write the compact state consumed by the final communication stage.
57 Args:
58 product: Fresh contiguous FP32 packed S/M product pointer.
59 previous: Previous packed S/M pointer; only S is read.
60 matrix: Compact M output pointer, unused when save_matrix is false.
61 state: Compact S output pointer.
62 row_count: Number of flattened matrix rows.
63 state_width: Columns in each compact S or M matrix.
64 packed_width: Columns in the concatenated S/M row.
65 block_rows: Flattened matrix rows per program, unrelated to token chunk size.
66 save_matrix: Whether a later round consumes compact M.
67 """
68 rows = tl.program_id(0) * block_rows + tl.arange(0, block_rows)
69 cols = tl.arange(0, state_width)
70 packed = rows[:, None] * packed_width + cols[None, :]
71 compact = rows[:, None] * state_width + cols[None, :]
72 left = tl.load(product + packed, rows[:, None] < row_count, 0)
73 old = tl.load(previous + packed, rows[:, None] < row_count, 0)
74 result = left + old
75 tl.store(product + packed, result, rows[:, None] < row_count)
76 tl.store(state + compact, result, rows[:, None] < row_count)
77 if save_matrix:
78 right = tl.load(product + packed + state_width, rows[:, None] < row_count, 0)
79 tl.store(matrix + compact, right, rows[:, None] < row_count)