使用 TileLang 在 AMD MI300X 上实现高性能 FlashMLA从 Hopper 到 CDNA3 的架构适配与优化实践【免费下载链接】tilelangDomain-specific language designed to streamline the development of high-performance GPU/CPU/Accelerators kernels项目地址: https://gitcode.com/GitHub_Trending/ti/tilelang本指南基于 TileLang 开源仓库中 AMD 版 FlashMLA 实现文档系统讲解如何将 TileLang 编写的 FlashMLAMulti-Head Latent Attention解码内核从 NVIDIA Hopper 迁移至 AMD MI300XCDNA3加速器。文章对比两大平台的指令集、共享内存与访存特性差异给出对应的内核改造策略寄存器缓存、流水线级数缩减、Warp 策略调整等并结合仓库内可运行的 TileLang / aiter / Triton 三套基准测试脚本与真实性能对比结果帮助读者掌握在 AMD ROCm 平台上使用 TileLang 编写高性能 MLA 解码内核的完整方法论。一、背景从 Hopper 到 MI300X 的移植动机DeepSeek 的 MLA 是一种以硬件效率著称的注意力机制其核心特点是巨大的头维度query/key 的头维度为 576512 64其中 64 为 RoPE 位置编码维度value 的头维度为 512。在 Hopper 版 TileLang FlashMLA 实现 中内核大量依赖 Hopper 特有的硬件特性包括 TMATensor Memory Accelerator、WGMMA 异步矩阵指令、warp specialization 等。而 AMD MI300X 采用 CDNA3 架构其硬件能力与编程模型和 Hopper 存在本质差异。本文档所在仓库在 examples/deepseek_mla/amd/ 目录下提供了面向 MI300X 的移植实现并用三份独立的基准脚本分别对 TileLang、AMD 官方汇编级内核库 aitermla_decode_fwd以及 Triton 进行对比评测。二、架构差异与对应的优化策略原文档将 Hopper 与 MI300X 的关键差异归纳为四点下面逐条展开并结合仓库源码印证。1. 指令集差异无显式 TMA 与 Warp SpecializationMI300X 不需要显式的 Tensor Memory AccessTMA指令也不需要手工编写 warp specialization 的 producer/consumer 逻辑。在 Hopper 上这些能力由编译器/驱动层自动处理因此移植到 MI300X 后内核源码层面几乎看不到差异——同一个T.copy/T.gemm前端写法在两种目标上都会被自动降级为各自平台的指令。对照两份内核实现可以直观看到这一点Hopper 版 example_mla_decode.py 使用T.Pipelined(loop_range, num_stages2)的多级流水线而 AMD 版 benchmark_mla_decode_amd_tilelang.py 使用T.Pipelined(loop_range, num_stages0)但核心的T.gemm、T.copy、T.reduce_max等前端原语完全一致。2. 共享内存约束64KB vs 228KB这是 MI300X 移植中影响最大的硬件约束。Hopper 每 SM 提供 228KB 共享内存而 MI300X 仅有 64KB。原文档给出的核心策略是缩减软件流水线级数多级流水线需要多份 shared memory 缓冲副本将 Q 矩阵从共享内存改为寄存器缓存T.alloc_fragment从而把宝贵的 shared memory 留给 K/V 分块。原文档给出的前后对比代码为# 原始方案共享内存分配 Q_shared T.alloc_shared([block_H, dim], dtype) Q_pe_shared T.alloc_shared([block_H, pe_dim], dtype) # 优化方案寄存器分配 Q_local T.alloc_fragment([block_H, dim], dtype) Q_pe_local T.alloc_fragment([block_H, pe_dim], dtype)这在仓库的 AMD 内核实现中得到了完整落地见 benchmark_mla_decode_amd_tilelang.pyQ_local T.alloc_fragment([block_H, dim], dtype) Q_pe_local T.alloc_fragment([block_H, pe_dim], dtype) KV_shared T.alloc_shared([block_N, dim], dtype) K_pe_shared T.alloc_shared([block_N, pe_dim], dtype) acc_s T.alloc_fragment([block_H, block_N], accum_dtype) acc_s_cast T.alloc_fragment([block_H, block_N], dtype) acc_o T.alloc_fragment([block_H, dim], accum_dtype)可以看到Q与Q_pe被放在寄存器fragment中而动态分块加载的KV_shared与K_pe_shared使用共享内存acc_s/acc_o这类高频累加缓冲也全部使用寄存器。这种分配策略在 64KB 共享内存预算下将空间让给了每轮迭代都要重新加载的 K/V 数据。3. Tile Size 灵活性摆脱 block_m 为 64 的倍数约束Hopper 的wgmma.mma_async指令要求最小 M 维度为 64导致block_m必须是 64 的倍数而 MI300X 没有 WGMMA 指令因此 Tile 尺寸的选择更加灵活不再受此约束。这使得开发者可以针对具体的 batch/head 数量选取更匹配的block_H例如 AMD 内核在 batch128、heads128 的默认配置下使用BLOCK_H 64、BLOCK_N 32、threads 128见 benchmark_mla_decode_amd_tilelang.py。4. 内存 Bank 冲突的 Swizzle 策略MI300X 与 NVIDIA 拥有不同的共享内存 bank 冲突规则因此需要采用不同的 swizzling 策略。TileLang 会自动为 AMD 目标选择适配的 swizzle 方案因此在用户代码层面同样没有可见差异。这也是文档强调的由 TileLang 自动处理的能力之一。三、AMD 内核实现细节以 benchmark_mla_decode_amd_tilelang.py 为例仓库中 benchmark_mla_decode_amd_tilelang.py 是 AMD 版 TileLang 内核的完整实现含自动调优配置、split 与 no-split 两个变体、参考实现与基准测试核心要点如下。1. 结构Splitnum_split 1与 No-split 双内核与 Hopper 版一致AMD 版也实现了类似 FlashDecoding 的 Split-KV 优化当 batch 较小时SM 并行度不足可以把kv_ctx维度切分成多份交由多个 SM 并行计算部分 logsumglse与部分输出Output_partial最后用一个轻量 combine 内核按 logsum 权重合并。main_splitT.Kernel(batch, heads // block_H, num_split)将 seqlen 均分到num_split个 block 上写出glse[batch, heads, num_split]与Output_partial[batch, heads, num_split, dim]main_no_split单块直接计算并写出最终Output两者的选择逻辑见 benchmark_mla_decode_amd_tilelang.pynum_split 1时返回main_split否则返回main_no_split。combine 内核使用T.serial(num_split)扫描各分块的 logsum先求全局最大值lse_max_local再以T.exp2(lse_local_split - lse_max_local)加权累加各分块输出L111-L136与标准 FlashDecoding 的合并逻辑一致。2. 数值技巧基于 log2(e) 的 exp2 软max内核把 softmax 的exp(x)改写为exp2(x * scale)其中scale (1.0 / (dim pe_dim)) ** 0.5 * 1.44269504 # log2(e)1.44269504即log2(e)。乘法缩放并入exp2之后编译器可以直接映射到硬件的exp2指令配合TL_ENABLE_FAST_MATH快速路径进一步压低指令开销。对应代码见 benchmark_mla_decode_amd_tilelang.py。3. 关键 TileLang 原语在本内核中的用法原语AMD 内核中的用法说明T.alloc_fragmentQ、Q_pe、acc_s、acc_o、scores 系列寄存器缓冲缓解共享内存压力T.alloc_sharedKV_shared、K_pe_shared仅 K/V 分块驻留共享内存T.gemm(..., transpose_BTrue, policyT.GemmWarpPolicy.FullRow)QK、Q_pe·K_pe、P·V 三次 GEMMFullRow 表示所有 warp 沿行方向划分见下文T.Pipelined(loop_range, num_stages0)K/V 加载与计算循环MI300X 上缩减流水线级数以节省共享内存T.use_swizzle(10)线程块调度通过数学映射调整 threadblock 执行顺序提升 L2 命中T.reduce_max/T.reduce_sum行方向归约维护在线 softmax 的 running max/sumT.exp2softmax 概率与缩放配合 log2(e) 系数关于GemmWarpPolicy仓库 tilelang/tileop/base.py 给出了明确定义Square为均衡的方形划分FullRow将所有 warp 沿行方向分配FullCol则沿列方向分配。AMD 内核统一使用FullRow与 Hopper 版使用的FullCol见 example_mla_decode.py形成对照——这正是 TileLang 允许开发者按目标架构灵活指定 warp 划分策略的体现。4. 自动调优配置内核通过tilelang.autotune遍历配置空间get_configsL9-L27BLOCK_N [16, 32, 64, 128] BLOCK_H [16, 32, 64, 128] num_split [1, 2, 4, 8, 16, 32] threads [128, 256]共4 × 4 × 6 × 2 192组候选配置配合tilelang.jit(out_idx[6], pass_configs{TL_ENABLE_FAST_MATH: True})编译。命令行通过--autotune开关决定是否启用自动调优默认使用BLOCK_N32, BLOCK_H64, num_split4, threads128。四、性能评估与 aiter-asm 持平、显著优于 Triton原文档给出了 float16 精度、batch size 为 64 与 128 条件下跨框架的计算吞吐对比结论见前文配图原图路径 examples/deepseek_mla/figures/flashmla-amd.pngTileLang vs aiter-asmAMD 官方手写汇编内核在绝大多数测试用例中达到性能对齐相对吞吐为0.73x 到 1.21xTileLang vs Triton显著胜出最高快6.5x如此性能仅由约70 行 Python 代码实现Hopper 版约为 80 行见主 README充分体现 TileLang 在 AMD 平台上低代码量 高性能的优势。需要说明的是这些数值是仓库文档在特定软硬件环境下AMD MI300X、float16、batch64/128、特定 seqlen 集合的实验结果实际表现会随设备、驱动与 ROCm 版本变化读者应以本地复测为准。五、基准测试脚本与复现方法仓库在 examples/deepseek_mla/amd/ 下提供了三份脚本用于在不同框架间横向对比脚本对比对象说明benchmark_mla_decode_amd_tilelang.pyTileLang 内核自身含正确性校验与ref_program对照rtol/atol0.01、自动调优、延迟与 TFlops 统计benchmark_mla_decode_amd_aiter.pytorchvsmla_aiter基于 DeepSeek FlashMLA 基准改造调用aiter.mla.mla_decode_fwd输出 CSVbenchmark_mla_decode_amd_triton.pytorchvsflash_mla_tritonTriton 版 MLA 解码含 Split-KV 两阶段 kernel输出 CSV三个脚本共享相同的评测形状集batch ∈ {64, 128}、seqlen ∈ {1024, 2048, 4096, 8192, 16384}、h_q128、h_kv1、d57651264、dv512。aiter 脚本使用 bfloat16Triton 脚本使用 float16。典型运行方式需在 AMD MI300X ROCm 已安装对应依赖的环境中执行# 验证并评测 TileLang 内核默认配置batch128 python benchmark_mla_decode_amd_tilelang.py # 开启自动调优 python benchmark_mla_decode_amd_tilelang.py --autotune # aiter 与 torch 对比--compare 模式输出 CSV python benchmark_mla_decode_amd_aiter.py --baseline torch --target mla_aiter --compare # Triton 与 torch 对比 python benchmark_mla_decode_amd_triton.py --baseline torch --target flash_mla_triton --compareTileLang 脚本还支持通过命令行参数调整问题规模--batch、--heads、--kv_heads、--kv_ctx、--dim、--pe_dim默认值分别为 128、128、1、8192、512、64L241-L251。正确性校验方面TileLang 脚本以ref_program基于 einops 与torch.nn.functional.softmax的稠密参考实现为基准并利用torch.testing.assert_close以rtol0.01, atol0.01验证输出aiter 与 Triton 脚本则通过compare_ab双向校验输出与 logsum。注意 aiter 脚本在未安装aiter库时会打印提示而非报错退出评测时需确保该库可用。六、未来优化方向原文档提出了两个明确的后续研究方向仓库亦在持续演进如 Hopper 侧已有 persistent 版本内核 作为参考。1. 内存 Bank 冲突的进一步缓解当前实现主要借助 TileLang 的自动优化解决 NT 布局下的 bank 冲突针对其他内存布局非 NT研究更精细的 swizzling 技术仍是开放的优化方向。感兴趣的读者可从 TileLang 的 swizzle 布局工具入手例如T.annotate_layout与make_swizzled_layout用法见主 README。2. 维度并行化对大型 MLA 维度例如 576 元素可研究头维度head dimension的分片策略预期收益包括降低共享内存压力大dim分片后每块驻留数据更少改善计算访存比更小的输出块带来更高的复用效率提升并行度通过维度级任务分发让更多 SM 参与同一 batch/head 的计算。七、结语与致谢本文围绕 TileLang 在 AMD MI300X 上的 FlashMLA 移植实践从指令集差异、共享内存约束、Tile 尺寸灵活性与 bank 冲突 swizzle 四个维度展开并结合仓库内可运行的 TileLang 内核与 aiter / Triton 基准脚本完整复现了架构分析 → 内核改造 → 性能验证的移植路径。实验表明借助 TileLang 的自动布局推断与后端适配能力一份约 70 行的 Python 内核即可在 MI300X 上达到与手写汇编内核aiter-asm相当、且显著优于 Triton 的吞吐表现。原文档特别对 AMD ROCm 与 Composable Kernel 团队的贡献表达了感谢认为从中受益良多这也提示读者在 AMD 平台做内核优化时ROCm 软件栈包含 aiter、Composable Kernel 等是值得对标与学习的权威参考。【免费下载链接】tilelangDomain-specific language designed to streamline the development of high-performance GPU/CPU/Accelerators kernels项目地址: https://gitcode.com/GitHub_Trending/ti/tilelang创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考
使用 TileLang 在 AMD MI300X 上实现高性能 FlashMLA：从 Hopper 到 CDNA3 的架构适配与优化实践
使用 TileLang 在 AMD MI300X 上实现高性能 FlashMLA从 Hopper 到 CDNA3 的架构适配与优化实践【免费下载链接】tilelangDomain-specific language designed to streamline the development of high-performance GPU/CPU/Accelerators kernels项目地址: https://gitcode.com/GitHub_Trending/ti/tilelang本指南基于 TileLang 开源仓库中 AMD 版 FlashMLA 实现文档系统讲解如何将 TileLang 编写的 FlashMLAMulti-Head Latent Attention解码内核从 NVIDIA Hopper 迁移至 AMD MI300XCDNA3加速器。文章对比两大平台的指令集、共享内存与访存特性差异给出对应的内核改造策略寄存器缓存、流水线级数缩减、Warp 策略调整等并结合仓库内可运行的 TileLang / aiter / Triton 三套基准测试脚本与真实性能对比结果帮助读者掌握在 AMD ROCm 平台上使用 TileLang 编写高性能 MLA 解码内核的完整方法论。一、背景从 Hopper 到 MI300X 的移植动机DeepSeek 的 MLA 是一种以硬件效率著称的注意力机制其核心特点是巨大的头维度query/key 的头维度为 576512 64其中 64 为 RoPE 位置编码维度value 的头维度为 512。在 Hopper 版 TileLang FlashMLA 实现 中内核大量依赖 Hopper 特有的硬件特性包括 TMATensor Memory Accelerator、WGMMA 异步矩阵指令、warp specialization 等。而 AMD MI300X 采用 CDNA3 架构其硬件能力与编程模型和 Hopper 存在本质差异。本文档所在仓库在 examples/deepseek_mla/amd/ 目录下提供了面向 MI300X 的移植实现并用三份独立的基准脚本分别对 TileLang、AMD 官方汇编级内核库 aitermla_decode_fwd以及 Triton 进行对比评测。二、架构差异与对应的优化策略原文档将 Hopper 与 MI300X 的关键差异归纳为四点下面逐条展开并结合仓库源码印证。1. 指令集差异无显式 TMA 与 Warp SpecializationMI300X 不需要显式的 Tensor Memory AccessTMA指令也不需要手工编写 warp specialization 的 producer/consumer 逻辑。在 Hopper 上这些能力由编译器/驱动层自动处理因此移植到 MI300X 后内核源码层面几乎看不到差异——同一个T.copy/T.gemm前端写法在两种目标上都会被自动降级为各自平台的指令。对照两份内核实现可以直观看到这一点Hopper 版 example_mla_decode.py 使用T.Pipelined(loop_range, num_stages2)的多级流水线而 AMD 版 benchmark_mla_decode_amd_tilelang.py 使用T.Pipelined(loop_range, num_stages0)但核心的T.gemm、T.copy、T.reduce_max等前端原语完全一致。2. 共享内存约束64KB vs 228KB这是 MI300X 移植中影响最大的硬件约束。Hopper 每 SM 提供 228KB 共享内存而 MI300X 仅有 64KB。原文档给出的核心策略是缩减软件流水线级数多级流水线需要多份 shared memory 缓冲副本将 Q 矩阵从共享内存改为寄存器缓存T.alloc_fragment从而把宝贵的 shared memory 留给 K/V 分块。原文档给出的前后对比代码为# 原始方案共享内存分配 Q_shared T.alloc_shared([block_H, dim], dtype) Q_pe_shared T.alloc_shared([block_H, pe_dim], dtype) # 优化方案寄存器分配 Q_local T.alloc_fragment([block_H, dim], dtype) Q_pe_local T.alloc_fragment([block_H, pe_dim], dtype)这在仓库的 AMD 内核实现中得到了完整落地见 benchmark_mla_decode_amd_tilelang.pyQ_local T.alloc_fragment([block_H, dim], dtype) Q_pe_local T.alloc_fragment([block_H, pe_dim], dtype) KV_shared T.alloc_shared([block_N, dim], dtype) K_pe_shared T.alloc_shared([block_N, pe_dim], dtype) acc_s T.alloc_fragment([block_H, block_N], accum_dtype) acc_s_cast T.alloc_fragment([block_H, block_N], dtype) acc_o T.alloc_fragment([block_H, dim], accum_dtype)可以看到Q与Q_pe被放在寄存器fragment中而动态分块加载的KV_shared与K_pe_shared使用共享内存acc_s/acc_o这类高频累加缓冲也全部使用寄存器。这种分配策略在 64KB 共享内存预算下将空间让给了每轮迭代都要重新加载的 K/V 数据。3. Tile Size 灵活性摆脱 block_m 为 64 的倍数约束Hopper 的wgmma.mma_async指令要求最小 M 维度为 64导致block_m必须是 64 的倍数而 MI300X 没有 WGMMA 指令因此 Tile 尺寸的选择更加灵活不再受此约束。这使得开发者可以针对具体的 batch/head 数量选取更匹配的block_H例如 AMD 内核在 batch128、heads128 的默认配置下使用BLOCK_H 64、BLOCK_N 32、threads 128见 benchmark_mla_decode_amd_tilelang.py。4. 内存 Bank 冲突的 Swizzle 策略MI300X 与 NVIDIA 拥有不同的共享内存 bank 冲突规则因此需要采用不同的 swizzling 策略。TileLang 会自动为 AMD 目标选择适配的 swizzle 方案因此在用户代码层面同样没有可见差异。这也是文档强调的由 TileLang 自动处理的能力之一。三、AMD 内核实现细节以 benchmark_mla_decode_amd_tilelang.py 为例仓库中 benchmark_mla_decode_amd_tilelang.py 是 AMD 版 TileLang 内核的完整实现含自动调优配置、split 与 no-split 两个变体、参考实现与基准测试核心要点如下。1. 结构Splitnum_split 1与 No-split 双内核与 Hopper 版一致AMD 版也实现了类似 FlashDecoding 的 Split-KV 优化当 batch 较小时SM 并行度不足可以把kv_ctx维度切分成多份交由多个 SM 并行计算部分 logsumglse与部分输出Output_partial最后用一个轻量 combine 内核按 logsum 权重合并。main_splitT.Kernel(batch, heads // block_H, num_split)将 seqlen 均分到num_split个 block 上写出glse[batch, heads, num_split]与Output_partial[batch, heads, num_split, dim]main_no_split单块直接计算并写出最终Output两者的选择逻辑见 benchmark_mla_decode_amd_tilelang.pynum_split 1时返回main_split否则返回main_no_split。combine 内核使用T.serial(num_split)扫描各分块的 logsum先求全局最大值lse_max_local再以T.exp2(lse_local_split - lse_max_local)加权累加各分块输出L111-L136与标准 FlashDecoding 的合并逻辑一致。2. 数值技巧基于 log2(e) 的 exp2 软max内核把 softmax 的exp(x)改写为exp2(x * scale)其中scale (1.0 / (dim pe_dim)) ** 0.5 * 1.44269504 # log2(e)1.44269504即log2(e)。乘法缩放并入exp2之后编译器可以直接映射到硬件的exp2指令配合TL_ENABLE_FAST_MATH快速路径进一步压低指令开销。对应代码见 benchmark_mla_decode_amd_tilelang.py。3. 关键 TileLang 原语在本内核中的用法原语AMD 内核中的用法说明T.alloc_fragmentQ、Q_pe、acc_s、acc_o、scores 系列寄存器缓冲缓解共享内存压力T.alloc_sharedKV_shared、K_pe_shared仅 K/V 分块驻留共享内存T.gemm(..., transpose_BTrue, policyT.GemmWarpPolicy.FullRow)QK、Q_pe·K_pe、P·V 三次 GEMMFullRow 表示所有 warp 沿行方向划分见下文T.Pipelined(loop_range, num_stages0)K/V 加载与计算循环MI300X 上缩减流水线级数以节省共享内存T.use_swizzle(10)线程块调度通过数学映射调整 threadblock 执行顺序提升 L2 命中T.reduce_max/T.reduce_sum行方向归约维护在线 softmax 的 running max/sumT.exp2softmax 概率与缩放配合 log2(e) 系数关于GemmWarpPolicy仓库 tilelang/tileop/base.py 给出了明确定义Square为均衡的方形划分FullRow将所有 warp 沿行方向分配FullCol则沿列方向分配。AMD 内核统一使用FullRow与 Hopper 版使用的FullCol见 example_mla_decode.py形成对照——这正是 TileLang 允许开发者按目标架构灵活指定 warp 划分策略的体现。4. 自动调优配置内核通过tilelang.autotune遍历配置空间get_configsL9-L27BLOCK_N [16, 32, 64, 128] BLOCK_H [16, 32, 64, 128] num_split [1, 2, 4, 8, 16, 32] threads [128, 256]共4 × 4 × 6 × 2 192组候选配置配合tilelang.jit(out_idx[6], pass_configs{TL_ENABLE_FAST_MATH: True})编译。命令行通过--autotune开关决定是否启用自动调优默认使用BLOCK_N32, BLOCK_H64, num_split4, threads128。四、性能评估与 aiter-asm 持平、显著优于 Triton原文档给出了 float16 精度、batch size 为 64 与 128 条件下跨框架的计算吞吐对比结论见前文配图原图路径 examples/deepseek_mla/figures/flashmla-amd.pngTileLang vs aiter-asmAMD 官方手写汇编内核在绝大多数测试用例中达到性能对齐相对吞吐为0.73x 到 1.21xTileLang vs Triton显著胜出最高快6.5x如此性能仅由约70 行 Python 代码实现Hopper 版约为 80 行见主 README充分体现 TileLang 在 AMD 平台上低代码量 高性能的优势。需要说明的是这些数值是仓库文档在特定软硬件环境下AMD MI300X、float16、batch64/128、特定 seqlen 集合的实验结果实际表现会随设备、驱动与 ROCm 版本变化读者应以本地复测为准。五、基准测试脚本与复现方法仓库在 examples/deepseek_mla/amd/ 下提供了三份脚本用于在不同框架间横向对比脚本对比对象说明benchmark_mla_decode_amd_tilelang.pyTileLang 内核自身含正确性校验与ref_program对照rtol/atol0.01、自动调优、延迟与 TFlops 统计benchmark_mla_decode_amd_aiter.pytorchvsmla_aiter基于 DeepSeek FlashMLA 基准改造调用aiter.mla.mla_decode_fwd输出 CSVbenchmark_mla_decode_amd_triton.pytorchvsflash_mla_tritonTriton 版 MLA 解码含 Split-KV 两阶段 kernel输出 CSV三个脚本共享相同的评测形状集batch ∈ {64, 128}、seqlen ∈ {1024, 2048, 4096, 8192, 16384}、h_q128、h_kv1、d57651264、dv512。aiter 脚本使用 bfloat16Triton 脚本使用 float16。典型运行方式需在 AMD MI300X ROCm 已安装对应依赖的环境中执行# 验证并评测 TileLang 内核默认配置batch128 python benchmark_mla_decode_amd_tilelang.py # 开启自动调优 python benchmark_mla_decode_amd_tilelang.py --autotune # aiter 与 torch 对比--compare 模式输出 CSV python benchmark_mla_decode_amd_aiter.py --baseline torch --target mla_aiter --compare # Triton 与 torch 对比 python benchmark_mla_decode_amd_triton.py --baseline torch --target flash_mla_triton --compareTileLang 脚本还支持通过命令行参数调整问题规模--batch、--heads、--kv_heads、--kv_ctx、--dim、--pe_dim默认值分别为 128、128、1、8192、512、64L241-L251。正确性校验方面TileLang 脚本以ref_program基于 einops 与torch.nn.functional.softmax的稠密参考实现为基准并利用torch.testing.assert_close以rtol0.01, atol0.01验证输出aiter 与 Triton 脚本则通过compare_ab双向校验输出与 logsum。注意 aiter 脚本在未安装aiter库时会打印提示而非报错退出评测时需确保该库可用。六、未来优化方向原文档提出了两个明确的后续研究方向仓库亦在持续演进如 Hopper 侧已有 persistent 版本内核 作为参考。1. 内存 Bank 冲突的进一步缓解当前实现主要借助 TileLang 的自动优化解决 NT 布局下的 bank 冲突针对其他内存布局非 NT研究更精细的 swizzling 技术仍是开放的优化方向。感兴趣的读者可从 TileLang 的 swizzle 布局工具入手例如T.annotate_layout与make_swizzled_layout用法见主 README。2. 维度并行化对大型 MLA 维度例如 576 元素可研究头维度head dimension的分片策略预期收益包括降低共享内存压力大dim分片后每块驻留数据更少改善计算访存比更小的输出块带来更高的复用效率提升并行度通过维度级任务分发让更多 SM 参与同一 batch/head 的计算。七、结语与致谢本文围绕 TileLang 在 AMD MI300X 上的 FlashMLA 移植实践从指令集差异、共享内存约束、Tile 尺寸灵活性与 bank 冲突 swizzle 四个维度展开并结合仓库内可运行的 TileLang 内核与 aiter / Triton 基准脚本完整复现了架构分析 → 内核改造 → 性能验证的移植路径。实验表明借助 TileLang 的自动布局推断与后端适配能力一份约 70 行的 Python 内核即可在 MI300X 上达到与手写汇编内核aiter-asm相当、且显著优于 Triton 的吞吐表现。原文档特别对 AMD ROCm 与 Composable Kernel 团队的贡献表达了感谢认为从中受益良多这也提示读者在 AMD 平台做内核优化时ROCm 软件栈包含 aiter、Composable Kernel 等是值得对标与学习的权威参考。【免费下载链接】tilelangDomain-specific language designed to streamline the development of high-performance GPU/CPU/Accelerators kernels项目地址: https://gitcode.com/GitHub_Trending/ti/tilelang创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考