基于 SOC 职业分类
用 Codex 或 Claude 帮你安装 复制这段 Prompt，粘贴到 Codex、Claude 或其他助手里，让它检查 Skill 页面并帮你完成安装。
直接命令不会经过审查 Prompt；运行前请先检查来源。
npx skills add https://github.com/mindspore-ai/akg --skill pypto-case-matvec命令会保持在同一行。复制前请横向滚动并检查完整内容。
想先保存到本地？可下载 SkillsMP 当前能够提供的文件。
正在显示 SKILL.md
矩阵乘法矩阵乘法 A[M, K] @ B[K, N] = C[M, N]中，大K维度矩阵乘法(K>>M,N)优化：针对M/N较小但K极大(如M=N=256,K=131072)的场景，Split-K切分K维度并行化、Workspace+Reduce替代全局同步，实现显著性能提升
Triton Ascend hard API restrictions and forbidden syntax. MUST-follow rules that apply to every kernel: forbidden control flow (return/break/continue/lambda/while), tensor slice/index restrictions, scalar conversion rules, BLOCK_SIZE upper bound. Violating any of these produces a compile or runtime error on Ascend.
Triton Ascend 性能优化通用策略: BLOCK_SIZE 选择 (1024-2048 for elementwise, must be <65536), grid configuration (use VEC_CORE_NUM / CUBE_CORE_NUM, 2D/3D grid for matmul / conv / reduce, 1D grid + inner loop for elementwise / pointwise), 256B alignment for memory transfers, autotune block-size patterns, fp16 / fp32 precision conversion. Bind via keywords like matmul, elementwise, reduce, block_size, grid, autotune, alignment, fp16, fp32, tile, interleaved-loop, cube-core, vec-core.
| name | pypto-case-matvec |
| description | 矩阵-向量乘法：K > 65535 时用 elementwise mul + sum 替代 matmul |
| category | case |
| version | 1.0.0 |
| metadata | {"backend":"ascend","dsl":"pypto","operator_patterns":"matrix_vector,matvec,large_k"} |
A: (256, 131072), B: (131072, 1) -> C: (256, 1)
K=131072 超过 pypto.matmul 限制（最后一维 <= 65535），用 sum(a * b_row, dim=1) 替代。
def create_matvec_sum_kernel(a_shape, b_shape):
out_shape = (a_shape[0], 1)
@pypto.frontend.jit(...)
def matvec_sum_kernel(
a: pypto.Tensor(a_shape, pypto.DT_FP32),
b_row: pypto.Tensor(b_shape, pypto.DT_FP32),
) -> pypto.Tensor(out_shape, pypto.DT_FP32):
output = pypto.tensor(list(out_shape), pypto.DT_FP32)
pypto.set_vec_tile_shapes(1, 8192)
output[:] = pypto.sum(a * b_row, dim=1, keepdim=True)
return output
return matvec_sum_kernel
class ModelNew(torch.nn.Module):
def forward(self, A, B):
assert A.dim() == 2
assert tuple(A.shape) == (256, 131072)
assert B.dim() == 2
assert tuple(B.shape) == (131072, 1)
A = A.contiguous()
# B: (K, 1) -> (1, K) 用于广播乘法
B_row = B.contiguous().reshape(1, -1)
return create_matvec_sum_kernel(tuple(A.shape), tuple(B_row.shape))(A, B_row)
关键点：forward 中 B.reshape(1, -1) 将列向量转为行向量，使 a * b_row 可广播。