Skip to content

fix(vpto): preserve MXFP4 load offsets and scale strides - #1055

Open
erhsh wants to merge 1 commit into
hw-native-sys:mainfrom
erhsh:fix/mxfp4-load-controls
Open

fix(vpto): preserve MXFP4 load offsets and scale strides#1055
erhsh wants to merge 1 commit into
hw-native-sys:mainfrom
erhsh:fix/mxfp4-load-controls

Conversation

@erhsh

@erhsh erhsh commented Jul 29, 2026

Copy link
Copy Markdown
Contributor

概要

  • 在 L1 到 L0 的 wrapper 展开过程中保留用户指定的 packed-FP4 kStart
  • 为 MX scale load 增加可选的源端和目标端 stride 覆盖参数
  • 在 PTODSL 中开放 stride 覆盖能力,同时保留原有的自动推导默认行为
  • 补充 PTODSL 与 VPTO 回归测试,覆盖非零 FP4 偏移和显式 MX stride

问题背景

packed-FP4 的 L1 到 L0 展开逻辑此前会同时换算 kStepkStart。因此,用户指定的 start_col = 4 最终会被错误生成为原始指令中的 kStart = 2,导致相邻 K 子块的数据范围发生重叠。

MX scale wrapper 指令此前只能根据逻辑 tile shape 推导源端和目标端 stride,无法表达某些合法布局。例如,需要 src_stride = 8dst_stride = 4 的布局会被静默生成为 4/4

这两类问题都可能正常编译,但在运行时读取错误的数据或 scale 元数据,进而产生静默数值错误。

修改内容

  • 保持 kStart 的原始语义,仅对 packed-FP4 的 kStep 进行宽度换算
  • mte_l1_l0a_mxmte_l1_l0b_mx 增加可选的 src_stridedst_stride 参数
  • wrapper 展开时优先使用显式 stride;未指定时继续采用原有的自动推导逻辑
  • 在 PTODSL 中增加对应的可选关键字参数
  • 增加回归测试,验证非零 FP4 起始偏移和显式 MX stride 能正确保留到 VPTO

未提供 stride 覆盖参数的现有调用行为保持不变。

影响范围

本次修改仅修正 packed-FP4 load 控制参数,并增强 MX scale load 对不同内存布局的表达能力。默认参数和已有调用接口保持兼容,不会改变未使用新参数的代码生成结果。

验证情况

  • PTOAS 全量构建通过
  • python ptodsl/tests/test_jit_compile.py 通过
  • VPTO wrapper round-trip 检查通过
  • VPTO 左、右 MX load wrapper 展开检查通过
  • packed-FP4 的 kStart 已验证生成为 4
  • 显式 MX stride 已验证生成为 8/4
  • git diff --check 通过

相关工作

#1047 处理的是 MX scale 目标地址编码问题,与本 PR 修复的 load 控制参数问题相互独立。本分支不包含 #1047 的修改。

fix: #1060

- preserve packed FP4 kStart during L1-to-L0 expansion
- support explicit source and destination strides for MX scale loads
- expose optional stride controls through PTODSL
- add DSL and VPTO expansion regression tests
@erhsh
erhsh marked this pull request as ready for review July 30, 2026 06:21
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[Bug] MXFP4 L1 到 L0 展开生成了错误的非零 K 偏移和 scale stride【阻塞】

1 participant