[kuri780]

quant: add SwiGLU forward/backward VMI kernels

概述

把 TileKernels 的 SwiGLU 前向/反向 ASC kernel 迁移到 VMI（PTO）后端。外层
T.Kernel / T.Persistent / T.copy / tiling / num_stages / buffer versions /
SF 布局 / token 调度与 ASC 源码保持一致，只把 with T.SimdVF() 内的向量数学改写为
逻辑 T.vmi.*；fast-certificate → exact 重算路径完整保留。

实现要点

• 不保留 ASC 实现作为 fallback；不引入 host 端预处理，不按 shape/dtype 新增特判。
• 不使用 ASC 的物理 lane 布局技巧（dist=、PART、PK、INTLV、E2B_B32、ONEPT），
  改用 VMI 的逻辑 vload/vcvt/vstore、vcmax(group=...) 和 compact scale。
• 可省略的 mask 全部去掉；API 强制要求 mask 的算子
  （vadds/vmuls/vmaxs/vmins/vshrs/vshls、vexpdif、vmula、vcmp/vsel、
  vcadd/vcmax/vcmin、vgather 以及 exact vdiv）保留。
• 前向 amax 使用宽向量 vcmax：hidden=2048 用 256-lane/group=8，hidden=3584 用
  128-lane/group=4；带 clamp 计数和非量化输出保持 64-lane/group=2。
  vcmax 的 group 是输出个数，因此 256-lane 不能用 group=4（会得到每 64 通道一组，
  与 sf_block=(1,32) 不符）。
• ptoas 0.66 不支持 256-lane 的 full reduction（vcmin/vcmax/vcadd 的 group=1 报
  VMI-RESIDUAL-OP），certificate 与 clamp 的最终归约使用 grouped reduction + 标量合并。
• packed SF 的 GM 拷贝通过 uint8 view 完成：当前固定的 tilelang 快照仍要求 MTE 源/目标
  元素宽度一致。上游 tilelang PR 11 Port mhc post kernels to VMI #14（codex/unify-asc-mte-pointer-lowering）合入并升级
  快照后可以删除该 workaround。
• FP8/bf16 输出 cast 显式 NOSAT：复现 ASC 当前行为并保留 NaN。
• exact 除法：forward 的 tie-break 用 lt，backward 用 le，与 ASC 源码一致；
  weight_grad 使用 vcadd(reassoc=True)，按仓库 1e-4 标准验收。

验证

• level-0 正确性（forward 24 + backward 24）：48/48 通过，输出与 ASC 逐位一致
  （weight_grad 按 1e-4）。
• 对抗性用例：review_adversarial_ab.py ALL_OK，覆盖 f32/bf16 次正规数、
  amax=inf 的 packed UE8M0、forward/backward 非有限输入。
• hidden=384（非 256/128 对齐）的 64-lane 回退路径单独通过。

性能

测试条件：num_send_tokens=4001、alignment=256、with_psum_layout=True。
ASC 取 5 次运行 per-file 中位数；VMI 为同一公共 API 的实测值。

方向	配置数	geomean ASC/VMI	min	max	≥0.90
前向	24	1.061	0.901	1.122	24
反向	24	0.907	0.815	1.030	8

前向 24 配置明细（按 ratio 降序）

方向	tokens	h	hidden_block	fmt	experts/topk	weights/rsf	clamp	ASC_us	VMI_us	ASC/VMI	ASC GB/s	VMI GB/s	结论
前向	4001	2048	—	bf16	32/6	—	0.5	340.2	303.3	1.122	994	1113	✅ ≥0.90
前向	4001	3584	—	bf16	4/8	—	0.5	701.9	629.4	1.115	1105	1232	✅ ≥0.90
前向	4001	2048	—	bf16	4/8	—	0.5	427.6	386.8	1.106	1031	1138	✅ ≥0.90
前向	4001	2048	—	bf16	32/6	yes/2.5	—	260.3	236.8	1.099	1301	1430	✅ ≥0.90
前向	4001	3584	—	bf16	32/6	—	0.5	538.5	490.1	1.099	1104	1213	✅ ≥0.90
前向	4001	2048	—	fp32	4/8	yes/2.5	—	325.6	296.4	1.098	1817	1995	✅ ≥0.90
前向	4001	3584	—	fp32	4/8	yes/2.5	—	541.9	494.3	1.096	1915	2099	✅ ≥0.90
前向	4001	3584	—	bf16	32/6	yes/2.5	—	437.8	399.9	1.095	1362	1492	✅ ≥0.90
前向	4001	2048	—	bf16	4/8	yes/2.5	—	336.4	307.7	1.093	1311	1433	✅ ≥0.90
前向	4001	3584	—	bf16	4/8	yes/2.5	—	563.9	515.9	1.093	1370	1498	✅ ≥0.90
前向	4001	3584	—	e4m3	4/8	yes/2.5	—	539.1	493.8	1.092	1202	1313	✅ ≥0.90
前向	4001	3584	—	fp32	32/6	yes/2.5	—	412.3	378.2	1.090	1901	2072	✅ ≥0.90
前向	4001	3584	—	e4m3	32/6	yes/2.5	—	416.2	385.0	1.081	1192	1289	✅ ≥0.90
前向	4001	2048	—	fp32	32/6	yes/2.5	—	246.7	229.3	1.076	1837	1977	✅ ≥0.90
前向	4001	3584	—	fp32	4/8	—	0.5	664.1	620.9	1.070	1553	1660	✅ ≥0.90
前向	4001	2048	—	fp32	4/8	—	0.5	410.8	384.6	1.068	1446	1544	✅ ≥0.90
前向	4001	2048	—	fp32	32/6	—	0.5	316.1	298.2	1.060	1414	1498	✅ ≥0.90
前向	4001	2048	—	e4m3	4/8	yes/2.5	—	339.7	321.3	1.057	1099	1161	✅ ≥0.90
前向	4001	3584	—	fp32	32/6	—	0.5	511.2	483.9	1.056	1543	1630	✅ ≥0.90
前向	4001	2048	—	e4m3	32/6	yes/2.5	—	259.4	248.8	1.043	1092	1139	✅ ≥0.90
前向	4001	3584	—	e4m3	32/6	—	0.5	467.8	464.0	1.008	1070	1077	✅ ≥0.90
前向	4001	3584	—	e4m3	4/8	—	0.5	587.4	594.2	0.988	1110	1097	✅ ≥0.90
前向	4001	2048	—	e4m3	32/6	—	0.5	309.3	341.3	0.906	927	840	✅ ≥0.90
前向	4001	2048	—	e4m3	4/8	—	0.5	388.2	430.8	0.901	956	853	✅ ≥0.90

反向 24 配置明细（按 ratio 降序）

方向	tokens	h	hidden_block	fmt	experts/topk	weights/rsf	clamp	recompute	ASC_us	VMI_us	ASC/VMI	ASC GB/s	VMI GB/s	结论
反向	4001	2048	2048	fp32	32/6	—	0.5	no	583.8	566.9	1.030	2299	2371	✅ ≥0.90
反向	4001	3584	3584	fp32	32/6	—	0.5	no	900.7	891.8	1.010	2608	2637	✅ ≥0.90
反向	4001	3584	3584	fp32	32/6	yes/2.5	—	yes	1044.7	1034.7	1.010	2698	2725	✅ ≥0.90
反向	4001	2048	2048	fp32	4/8	—	0.5	no	621.4	622.8	0.998	2401	2393	✅ ≥0.90
反向	4001	2048	2048	fp32	32/6	yes/2.5	—	yes	659.7	661.4	0.998	2442	2436	✅ ≥0.90
反向	4001	2048	2048	fp32	4/8	yes/2.5	—	yes	731.3	733.3	0.997	2480	2472	✅ ≥0.90
反向	4001	3584	3584	fp32	4/8	—	0.5	no	974.2	996.5	0.978	2694	2636	✅ ≥0.90
反向	4001	3584	3584	fp32	4/8	yes/2.5	—	yes	1126.4	1160.3	0.971	2780	2696	✅ ≥0.90
反向	4001	2048	2048	bf16	32/6	yes/2.5	—	yes	571.1	642.2	0.889	1411	1254	⚠️ <0.90
反向	4001	2048	2048	bf16	4/8	yes/2.5	—	yes	632.0	711.7	0.888	1414	1256	⚠️ <0.90
反向	4001	3584	3584	bf16	32/6	yes/2.5	—	yes	973.4	1102.6	0.883	1448	1278	⚠️ <0.90
反向	4001	3584	3584	bf16	4/8	yes/2.5	—	yes	1101.9	1248.1	0.883	1449	1279	⚠️ <0.90
反向	4001	2048	2048	e4m3	4/8	yes/2.5	—	yes	897.9	1022.8	0.878	1028	903	⚠️ <0.90
反向	4001	2048	2048	e4m3	32/6	yes/2.5	—	yes	809.2	922.5	0.877	1013	889	⚠️ <0.90
反向	4001	2048	2048	e4m3	32/6	—	0.5	no	649.2	742.9	0.874	1066	932	⚠️ <0.90
反向	4001	2048	2048	e4m3	4/8	—	0.5	no	719.6	825.9	0.871	1074	936	⚠️ <0.90
反向	4001	3584	3584	e4m3	4/8	yes/2.5	—	yes	1523.1	1749.1	0.871	1046	911	⚠️ <0.90
反向	4001	3584	3584	e4m3	32/6	yes/2.5	—	yes	1377.3	1585.3	0.869	1050	912	⚠️ <0.90
反向	4001	3584	3584	e4m3	4/8	—	0.5	no	1216.0	1409.7	0.863	1113	960	⚠️ <0.90
反向	4001	3584	3584	e4m3	32/6	—	0.5	no	1096.8	1272.0	0.862	1104	952	⚠️ <0.90
反向	4001	2048	2048	bf16	32/6	—	0.5	no	415.6	485.5	0.856	1615	1382	⚠️ <0.90
反向	4001	2048	2048	bf16	4/8	—	0.5	no	450.8	543.8	0.829	1652	1370	⚠️ <0.90
反向	4001	3584	3584	bf16	4/8	—	0.5	no	751.0	919.6	0.817	1747	1427	⚠️ <0.90
反向	4001	3584	3584	bf16	32/6	—	0.5	no	672.4	825.3	0.815	1747	1423	⚠️ <0.90

• 反向低于 0.90 的配置集中在 bf16/e4m3，试验测得性能差距原因在exact div,去掉exact之后反向性能也全部达标。

[Zhendong404]

能省略的mask尽量省掉

[Zhendong404]

这个PR应该能解决T.copy src/dst类型不一致的问题，可以试试：codex/unify-asc-mte-pointer-lowering

[Zhendong404]

这里可以用256长度group=4的vreg表达

[kuri780]

已改。这里 vcmax 的 group 是输出组数，256-lane f32 / group=4 会得到 64 通道一组，与 sf_block=(1,32) 不符，所以 2048 用 256-lane/group=8；3584 用 128-lane/group=4（256-lane 时 num_vecs=14，tile 不满会触发 ptoas 的 layout 冲突）。clamp 计数保持 64-lane，避免宽计数器溢出寄存器。