系统版本 / 镜像
| 项目 | 信息 |
|---|---|
| 系统 | Bianbu 4.0.1 (Resolute Raccoon) |
| 内核 | 6.18.3-generic |
| 架构 | RISC-V 64 (rv64gc) |
| SoC | SpacemiT Key Stone K3 |
| AI 核 | 8× A100, VLEN=1024, 60 TOPS |
| 相关库 | Spine-Triton (SpacemiT 定制版) |
| 开发板 | SpacemiT K3 Pico ITX |
使用场景或测试目标
使用 Spine-Triton 在 A100 AI 核上进行GEMM 运算。
GEMM 矩阵规模为：
col: 518400 × 27（或 144）weight: 16 × 27（或 144 × 16）- 总元素数：
518400 × 27 = 13,996,800> 1M
为了在 A100 上执行，采用了 K 维分块策略：
python
分块策略 for k_start in range(0, K, max_K): block = a[:, k_start:k_start+max_K] @ b[k_start:k_start+max_K, :] # max_K = min(1048576 // N, K) # 确保分块后 <=1M
这样每个 make_block_ptr 只指向当前 K 分块（≤1M 元素），而非整个 tensor。
操作步骤
1. 编写 Spine-Triton kernel
python
@triton.jit def gemm_kernel(a_ptr, b_ptr, c_ptr, M, N, K, …): # 获取当前 program 的块索引 pid_m = tl.program_id(0) pid_n = tl.program_id(1) pid_k = tl.program_id(2) # K 维分块 # 使用 make_block_ptr 指向分块后的数据 a_block = tl.make_block_ptr( base=a_ptr, shape=(M, K_block), # K_block ≤ 1M strides=(K, 1), # stride 基于原始 K offsets=(m_off, k_off), block_shape=(BLOCK_M, BLOCK_K), order=(1, 0) ) # …
2. 编译 kernel
bash
python3 -c " import triton import triton.language as tl @triton.jit def test_kernel(…): # 包含 make_block_ptr 的 kernel pass # 尝试编译 test_kernel(1,1,1) "
3. 观察编译器错误
text
RuntimeError: make_block_ptr cannot be used on tensors with numel > 1048576
结果数据 / 截图 / 日志
编译器错误
text
RuntimeError: make_block_ptr cannot be used on tensors with numel > 1048576
编译器检查逻辑（推测）
python
当前检查逻辑（在 spine-triton 编译器中） def check_make_block_ptr(tensor): if tensor.numel() > MAX_TENSOR_NUMEL: # 1048576 raise RuntimeError(“make_block_ptr cannot be used on tensors with numel > 1048576”)
分块前后的对比
| 场景 | K 分块后 K_block | 元素数 | 编译器检查结果 |
|---|---|---|---|
| 无分块 | 518400 | 13,996,800 |  失败 |
| K 分块 (max_K=1048576//N) | 65536 | 1,769,472 | 仍 >1M |
| K 分块 (max_K=1048576//N) | 32768 | 884,736 |  <1M 但 N 太小 |
问题在于：编译器检查的是基 tensor 的总元素数 ，而不是 make_block_ptr 实际指向的分块大小 。
实际内存占用
| 数据 | 位置 | 说明 |
|---|---|---|
| 原始 tensor | CPU 内存 | 由 Python numpy 管理 |
| 分块数据 | A100 Scratchpad | 仅加载当前 K 分块 (≤1M) |
编译器的一刀切检查阻止了用户手动分块的意图——虽然分块后的数据可以放入 Scratchpad，但编译器仍拒绝编译。
遇到的问题
核心问题 ：Spine-Triton 编译器对所有 make_block_ptr 做一刀切检查，没有区分“这个 tensor 是否需要完全驻留在 Scratchpad 上”和“这个 tensor 只是被分块访问”。
| 检查项 | 预期 | 实际 |
|---|---|---|
| 分块后 tensor 元素数 ≤1M | 是 |
是 |
| 编译器检查通过 | 应该通过 |
拒绝编译 |
| 编译器检查的对象 | 分块后的视图 | 原始 tensor |
影响范围
- K 维分块 A100 GEMM 无法工作
你的判断 / 改进建议
判断
- 编译器检查过于激进 ：
make_block_ptr的检查应该基于实际 pointer 指向的视图大小 ，而非基 tensor 的总大小。用户已经通过分块策略确保了每个make_block_ptr指向的数据 ≤1M。 - Triton 设计意图 ：在标准 Triton (CUDA) 中，
make_block_ptr的shape参数就是实际分块形状，编译器应基于shape而非基 tensor 的numel()做检查。SpacemiT 版本可能修改了此行为。 - 性能影响 ：如果此问题不修复，A100 在 K3 上的可用性将严重受限——任何需要处理大 tensor 的 GEMM 都无法使用 A100，只能回退到 CPU。