Skip to content

[Feature] VMI缺少面向 bf16 ↔ f32 或 e4m3 ↔ f32 的融合加宽/窄化内存操作 #1644

Closed

Description

@kuri780

Summary

VMI 目前没有面向 bf16 ↔ f32 或 e4m3 ↔ f32 的融合加宽/窄化内存操作。
VMI kernel 每处理 64 个逻辑值需要两条指令(vload + vcvt,或
vcvt + vstore);ASC 路径可以用一条融合的 vld(dist='UNPK_*') /
vsts(dist='PK*') 完成同样的搬运。PTO 后端已经支持这些融合形式,缺的是
等价的逻辑 VMI 接口。

Motivation / use case

SwiGLU forward/backward 从 ASC 迁移到 VMI
(https://github.com/PTO-ISA/TileKernels-PTO/pull/20)后,16 个 backward
level-0 config 低于 0.90 的验收门槛,全部是 bf16 或 e4m3 路径;8 个 fp32
backward config 全部 ≥ 0.90。它们正好是走分离 unpack/pack 的路径。

  • 当前 VMI 写法:
with T.SimdVF():
    mask = T.vmi.create_mask(64, size=64)
    x = T.vmi.vload(x_ub[0], size=64)      # bf16x64 或 e4m3x64
    xf = T.vmi.vcvt(x, 'float32')           # f32x64
    T.vmi.vstore(xf, out_ub[0], mask)
  • T.vmi.vload 只支持连续、group+stride、block_stride 三种寻址,
    没有加宽/转换模式;T.vmi.vcvt 不改变逻辑 lane 数。因此 bf16/FP8 必须先
    load 再 convert,store 方向同理。
  • 每 64 个逻辑值的预期物理指令数(需要 lower-level lowering 最终确认):
搬运方向 ASC 融合 VMI size=64 VMI size=128 VMI size=256
bf16 → f32 load 1 2 1.5 1.5
e4m3 → f32 load 1 2 1.5 1.25
f32 → bf16 store 1 2 1.5 1.5
f32 → e4m3 store 1 2 1.5 1.25

bf16 的收益在 128 lane 就饱和;FP8 的 256 lane load 可以到 1.25/64。
更大宽度下算术部分每值的物理指令数不变,收益主要在 load/store 路径
(以及 loop/地址计算开销)。

受影响的 backward config

bf16

h topk clamp weights recompute ASC_us VMI_us ratio
3584 8 0.5 否 否 745.0 917.5 0.812
3584 6 0.5 否 否 668.6 822.1 0.813
2048 8 0.5 否 否 445.4 533.1 0.835
2048 6 0.5 否 否 408.1 481.3 0.848
3584 8 — 是 是 1099.6 1246.4 0.882
3584 6 — 是 是 971.9 1101.5 0.882
2048 8 — 是 是 628.0 709.2 0.886
2048 6 — 是 是 566.7 639.2 0.887

e4m3

h topk clamp weights recompute ASC_us VMI_us ratio
3584 8 0.5 否 否 1213.8 1408.3 0.862
3584 6 0.5 否 否 1095.8 1265.4 0.866
2048 8 0.5 否 否 717.5 822.6 0.872
2048 6 0.5 否 否 647.6 742.3 0.872
3584 8 — 是 是 1525.3 1746.6 0.873
3584 6 — 是 是 1377.2 1571.5 0.876
2048 8 — 是 是 897.3 1021.2 0.879
2048 6 — 是 是 810.0 914.3 0.886

Proposed API / behavior

No response

Alternatives considered

No response

Additional context

No response

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    enhancementNew feature or request

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions