torch_npu.npu_flash_attention(q, k, v, ...)，一行代码出结果。但从这行调用到真正在昇腾910的达芬奇架构上跑起来，中间经过了图编排、算子加载、Tiling参数计算、UB内存分配、任务流水线调度。完成这些工作的，是CANN Runtime。

Runtime在CANN生态里的位置

CANN软件栈从下到上分四层：

Ascend 910 NPU硬件
 ↑
CANN基础层：驱动、CCE编译器、hccl集合通信库
 ↑
算子层：ops-transformer（FlashAttention实现）、catlass（算子模板）
 ↑
Runtime层：GE图引擎的核心调度子系统 ← 重点在这
 ↑
框架适配层：torch_npu（Python→ATen算子→Runtime）

Runtime不是一个独立程序，而是GE（图引擎）的核心调度子系统。GE负责整张计算图的编排和优化，Runtime负责把图里的每个算子调度到达芬奇内核上执行。

从功能上看，Runtime做三件事：

1. 算子管理：查找算子实现、加载编译产物、管理算子生命周期
2. 内存管理：UB分配与回收、HBM→UB→UB的数据搬运编排
3. 任务调度：计算单元和DMA引擎的流水线编排、Tiling参数下发

FlashAttention的所有特殊处理（Tiling计算、UB分区、causal mask块级跳过）都是Runtime和算子内核配合完成的——Runtime负责调度决策，算子内核负责具体计算。

调用链：从Python到Runtime

FlashAttention在昇腾NPU上的完整调用路径：

用户Python代码
 ↓ torch_npu.npu_flash_attention(...)
ATen算子注册（torch_npu/csrc/aten/.../flash_attention.cpp）
 ↓ op.feed_to_graph()
GE图编排（确定算子依赖、融合机会）
 ↓
Runtime调度（算子查找、Tiling计算、UB分配）
 ↓ DispatchToAicore()
CCE可执行文件（.cce格式，达芬奇微指令序列）
 ↓
达芬奇AICore执行（Cube+Vector+DMA）

用户感受到的调用是torch_npu.npu_flash_attention()，但实际执行链路要经过四层转换。Runtime的关键作用体现在中间两段：算子查找与加载和调度决策。

算子查找与加载

torch_npu调用npu_flash_attention时，先通过ATen算子注册找到对应的算子名：

// torch_npu/csrc/aten/flash_attention.cpp
at::Tensor npu_flash_attention(
 at::Tensor q, at::Tensor k, at::Tensor v,
 int64_t head_num, std::string input_layout, double scale
) {
 // 构建算子描述符
 OpDesc op_desc;
 op_desc.set_name("FlashAttentionScore"); // 算子名对应ops-transformer注册名
 op_desc.set_input("q", q);
 op_desc.set_input("k", k);
 op_desc.set_input("v", v);
 op_desc.set_attr("head_num", head_num);
 op_desc.set_attr("input_layout", input_layout);
 op_desc.set_attr("scale", scale);
 
 // 交付给GE
 return graph_engine::execute(op_desc);
}

Runtime收到"FlashAttentionScore"这个算子名后，在注册表里查找实现：

// GE/算子管理/operator_registry.cpp
Operator* OperatorRegistry::FindOperator(const std::string& op_name) {
 // 查已注册的算子
 auto it = registered_ops_.find(op_name);
 if (it != registered_ops_.end()) {
 return it->second;
 }
 
 // 查自定义算子路径（ASCEND_CUSTOM_OPP_PATH环境变量指定）
 return LoadFromCustomPath(op_name);
}

ops-transformer编译产物libflash_attention*.so被加载后，Runtime通过符号查找到算子入口。这就是为什么改了ops-transformer代码后必须重新编译替换so——符号找不到，Runtime就找不到算子实现。

Tiling参数：Runtime怎么算出分块大小

FlashAttention的分块大小不是写死在算子里的。Runtime根据输入shape、UB容量、头维度，实时计算Tiling参数，下发给算子：

// GE/Runtime/tiling/flash_attention_tiling.cc
struct TilingParam {
 uint32_t block_m; // Q块在seq维的大小
 uint32_t block_n; // K/V块在seq维的大小
 uint32_t block_k; // head_dim
 uint32_t num_q_blocks; // Q块数量
 uint32_t num_kv_blocks;// K/V块数量
};

TilingParam CalcFlashAttentionTiling(
 uint32_t seq_len, uint32_t head_dim, uint32_t batch, uint32_t heads
) {
 // UB容量（昇腾910约256KB）
 const uint32_t kUBSize = 256 * 1024;
 const uint32_t kDtypeBytes = 2; // FP16
 
 // 计算单块所需空间
 // Q块 + K块 + V块 + 输出块 + Softmax中间结果
 uint32_t elements_per_block = kUBSize / (5 * kDtypeBytes);
 
 // 反推seq维块大小
 uint32_t block_seq = elements_per_block / head_dim;
 
 // 向下对齐到16（达芬奇对齐约束）
 block_seq = (block_seq / 16) * 16;
 block_seq = std::min(block_seq, seq_len);
 
 // 计算块数量
 uint32_t num_blocks = (seq_len + block_seq - 1) / block_seq;
 
 return {
 block_seq, // block_m
 block_seq, // block_n
 head_dim, // block_k
 num_blocks, // num_q_blocks
 num_blocks // num_kv_blocks
 };
}

Tiling参数通过OpContext打包后传给算子内核。算子代码从context里取出这些参数，决定每次循环处理多大的块。Runtime不参与具体的分块计算逻辑，但它的Tiling策略直接决定算子能否在UB里跑起来。

UB内存管理：Runtime怎么分配Unified Buffer

Unified Buffer是达芬奇架构的高速片上存储，所有计算单元直接访问。FlashAttention的分块计算要求5个buffer同时待在UB里——Q块、K块、V块、输出块、softmax中间结果。任何一块超容量，整个算子就OOM。

// GE/Runtime/memory/ub_allocator.cc
class UBAllocator {
 static constexpr size_t kUBSize = 256 * 1024; // 256KB
 
 UBBuffer Allocate(const TilingParam& tiling) {
 // 计算所需空间
 size_t q_bytes = tiling.block_m * tiling.block_k * sizeof(float16);
 size_t k_bytes = tiling.block_n * tiling.block_k * sizeof(float16);
 size_t v_bytes = tiling.block_n * tiling.block_k * sizeof(float16);
 size_t out_bytes = tiling.block_m * tiling.block_k * sizeof(float16);
 size_t softmax_bytes = tiling.block_m * tiling.block_n * sizeof(float16);
 
 size_t total = q_bytes + k_bytes + v_bytes + out_bytes + softmax_bytes;
 
 if (total > kUBSize) {
 // 超过UB容量，自动降级到更小的分块
 auto reduced = ReduceBlockSize(tiling, kUBSize);
 return Allocate(reduced); // 递归重试
 }
 
 // 128字节对齐分配
 return UBBuffer{
 .q_offset = Align(q_bytes, 128),
 .k_offset = Align(k_bytes, 128),
 .v_offset = Align(v_bytes, 128),
 .out_offset = Align(out_bytes, 128),
 .softmax_offset = Align(softmax_bytes, 128)
 };
 }
};

Runtime的降级策略：检测到超容量后自动缩小分块大小，重新计算所需空间，直到能装下。这保证了下限能跑——哪怕性能差，至少不会OOM。但降级后的block太小会导致循环次数增加，性能可能退化。

任务调度：DMA和Cube怎么配合

Runtime的调度器把FlashAttention的计算组织成流水线。达芬奇架构有独立的DMA引擎和Cube计算单元，两者可以并行——这是性能的关键：

// GE/Runtime/scheduler/flash_attention_scheduler.cc
void ScheduleFlashAttention(const OpContext& ctx) {
 for (uint32_t qi = 0; qi < ctx.num_q_blocks; qi++) {
 // DMA：加载Q块到UB
 dma_submit(DMAJob::LOAD, q_base + qi * ctx.block_m, ub_q);
 
 for (uint32_t ki = 0; ki < ctx.num_kv_blocks; ki++) {
 // 关键：K块的DMA加载和当前块计算并行
 dma_submit(DMAJob::LOAD, k_base + ki * ctx.block_n, ub_k);
 
 // 等待K块加载完成
 dma_wait(DMAJob::LOAD, ub_k);
 
 // Cube：Q × K^T（与DMA并行）
 cube_submit(CubeJob::GEMM, ub_q, ub_k, ub_scores);
 cube_wait(CubeJob::GEMM);
 
 // Scale + 在线softmax
 cube_submit(CubeJob::SCALE, ub_scores, ctx.scale);
 cube_submit(CubeJob::REDUCE_MAX, ub_scores, ub_row_max);
 cube_submit(CubeJob::REDUCE_SUM, ub_scores, ub_row_sum);
 
 // DMA：加载V块
 dma_submit(DMAJob::LOAD, v_base + ki * ctx.block_n, ub_v);
 dma_wait(DMAJob::LOAD, ub_v);
 
 // Cube：softmax累加
 cube_submit(CubeJob::SOFTMAX_ACC, ub_scores, ub_v, ub_out);
 }
 
 // Cube：归一化
 cube_submit(CubeJob::SCALE, ub_out, 1.0f / ctx.row_sum);
 
 // DMA：写回输出
 dma_submit(DMAJob::STORE, ub_out, out_base + qi * ctx.block_m);
 dma_wait(DMAJob::STORE, ub_out);
 }
}

关键设计是DMA和Cube流水线化：K块和V块的DMA加载与Cube计算overlap。当前块做完Q×K^T后，K块已经加载完成，直接开始下一块。Cube不是等着DMA，而是DMA在Cube空档时偷偷搬运数据。

错误处理：Runtime怎么兜底

Runtime提供了多层错误处理：

// GE/Runtime/error/op_error_handler.cc
enum class OpErrorCode {
 SUCCESS = 0,
 SHAPE_MISMATCH = 1, // layout不匹配
 TILING_FAILED = 2, // Tiling参数算不出
 UB_OOM = 3, // UB超容量
 COMPILATION_FAILED = 4,
 DMA_TIMEOUT = 5,
};

OpResult ExecuteWithFallback(OpNode* node) {
 auto err = ValidateInput(node);
 if (err != SUCCESS) {
 return OpResult::Error(err);
 }
 
 auto tiling = CalcTiling(node->input());
 if (tiling.ub_required > kMaxUBSize) {
 // 降级策略1：减小分块
 tiling = TryWithSmallerBlock(tiling);
 if (tiling.ub_required > kMaxUBSize) {
 // 降级策略2：强制padding到对齐长度
 tiling = ForcePadAndRetile(node->input());
 }
 }
 
 return DispatchToAicore(node, tiling);
}

三个常见Runtime错误：

SHAPE_MISMATCH：传入的tensor layout（BNSD/BSND）和算子不匹配，Runtime在ValidateInput阶段检测到，返回错误。检查input_layout参数和tensor实际格式。

TILING_FAILED：序列太大UB装不下，典型场景是seq=10000且dim=256。手动padding到更小的seq_len，或调整head_dim。

UB_OOM：Runtime自动降级后仍超容量，强制报错。这种情况需要调整输入shape或分块参数。

跟GE的关系：Runtime是GE的调度引擎

Runtime不是一个独立的进程或服务，它是GE的核心子系统。GE在图编排阶段做的事：

// GE/GraphEngine/
...(truncated)...