« prev ^ index » next coverage.py v7.13.1, created at 2026-10-25 05:45 +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# Portions copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li.
16# FLA's MIT-licensed CP merge recurrence is adapted to reuse local M when gathering G.
17# See FLA_LICENSE in this directory.
18"""FLA-style one-kernel prefix/suffix with separate transfer storage and fused transpose."""
19import triton
20import triton.language as tl
23@triton.jit(do_not_specialize=["count", "rank"])
24def kda_cached_affine_merge_kernel(
25 output: tl.tensor, states: tl.tensor, matrices: tl.tensor, count: tl.tensor, rank: tl.tensor,
26 state_rank_stride: tl.constexpr, state_head_stride: tl.constexpr, state_row_stride: tl.constexpr,
27 matrix_rank_stride: tl.constexpr, matrix_head_stride: tl.constexpr, matrix_row_stride: tl.constexpr,
28 forward: tl.constexpr, block_v: tl.constexpr,
29) -> None:
30 """Evaluate a zero-entry prefix/suffix state, without materializing a composed M.
32 Local transfer matrices are produced before this kernel. Each iteration
33 evaluates M @ state + S (or M.T @ gradient + G) in chronological order.
35 Args:
36 output: Contiguous FP32 recurrent-state output pointer.
37 states: Chronological affine offsets [P,B,H,128,128].
38 matrices: Matching chronological FP32 transfer matrices.
39 count: Number of chronological offsets to merge.
40 rank: Chronological rank within the CP group.
41 state_rank_stride: Element stride between source-rank state slabs.
42 state_head_stride: Element stride between flattened batch/head state slabs.
43 state_row_stride: Element stride between key rows in state storage.
44 matrix_rank_stride: Element stride between source-rank transfer slabs.
45 matrix_head_stride: Element stride between flattened batch/head transfer slabs.
46 matrix_row_stride: Element stride between rows in M.
47 forward: Compile-time selection of prefix or transposed suffix recurrence.
48 block_v: Value-column tile width.
49 """
50 column, head = tl.program_id(0), tl.program_id(1)
51 rows = tl.arange(0, 128)
52 cols = column*block_v + tl.arange(0, block_v)
53 value = tl.zeros((128, block_v), tl.float32)
54 for index in range(count):
55 source = index if forward else rank+count-index
56 offset = states + source*state_rank_stride + head*state_head_stride
57 state = tl.load(offset + rows[:, None]*state_row_stride + cols[None, :])
58 offset_m = matrices + source*matrix_rank_stride + head*matrix_head_stride
59 if forward:
60 offsets_m = rows[:, None]*matrix_row_stride + rows[None, :]
61 else:
62 offsets_m = rows[None, :]*matrix_row_stride + rows[:, None]
63 matrix = tl.load(offset_m+offsets_m)
64 value = tl.dot(matrix.to(tl.float32), value, input_precision="ieee") + state.to(tl.float32)
65 tl.store(output+head*16384+rows[:, None]*128+cols[None, :], value)
68@triton.jit(do_not_specialize=["count", "start", "prefix_count"])
69def kda_affine_scan_kernel(
70 output: tl.tensor, prefix_output: tl.tensor, maps: tl.tensor, initial: tl.tensor,
71 count: tl.tensor, start: tl.tensor, prefix_count: tl.tensor,
72 rank_stride: tl.constexpr, head_stride: tl.constexpr, row_stride: tl.constexpr,
73 forward: tl.constexpr, save_prefix: tl.constexpr, block_v: tl.constexpr,
74) -> None:
75 """Scan packed summaries from a nonzero state, optionally saving an owner prefix.
77 Args:
78 output: Contiguous FP32 final state pointer.
79 prefix_output: Optional contiguous FP32 intermediate-state pointer.
80 maps: Packed offsets and M, shaped [P,B,H,128,256].
81 initial: Contiguous incoming state shaped [B,H,128,128].
82 count: Number of ordered affine applications.
83 start: First source index in the packed summaries.
84 prefix_count: Number of applications before saving the intermediate state.
85 rank_stride: Element stride between source ranks.
86 head_stride: Element stride between flattened batch/head slabs.
87 row_stride: Element stride between key rows.
88 forward: Select forward order and M instead of reverse order and M.T.
89 save_prefix: Write an intermediate state during the scan.
90 block_v: Value-column tile width.
91 """
92 column, head = tl.program_id(0), tl.program_id(1)
93 rows = tl.arange(0, 128)
94 cols = column * block_v + tl.arange(0, block_v)
95 output_offsets = head * 16384 + rows[:, None] * 128 + cols[None, :]
96 value = tl.load(initial + output_offsets).to(tl.float32)
97 for index in range(count):
98 source = start + index if forward else start - index
99 base = maps + source * rank_stride + head * head_stride
100 offset = tl.load(base + rows[:, None] * row_stride + cols[None, :])
101 if forward:
102 matrix_offsets = rows[:, None] * row_stride + rows[None, :]
103 else:
104 matrix_offsets = rows[None, :] * row_stride + rows[:, None]
105 matrix = tl.load(base + 128 + matrix_offsets).to(tl.float32)
106 value = tl.dot(matrix, value, input_precision="ieee") + offset.to(tl.float32)
107 if save_prefix:
108 if index + 1 == prefix_count:
109 tl.store(prefix_output + output_offsets, value)
110 tl.store(output + output_offsets, value)