Add FP8/FP4 MX dtype support in task_interface

[yanghaoran29]

Register the MX block-scale host dtypes so torch FP8/FP4 tensors can flow through the task interface (host↔device transfer and native MX matmul).

Changes

• DataType::FP8E4M3FN / FP8E8M0 (1 byte each) and FP4E2M1 (packed E2M1 ×2 / byte, matches torch.float4_e2m1fn_x2) — A5 only.
• FP8E5M2 intentionally omitted (MX paths use E4M3FN + E8M0 scale, not E5M2).
• get_element_size / get_dtype_name + nanobind bindings.
• torch_interop maps float8_e4m3fn / float8_e8m0fnu / float4_e2m1fn_x2 (getattr-guarded).
• UT
  • test_task_interface.py: enum / sizes / names / make_tensor_arg (skip only when none of the optional torch dtypes exist).
  • test_mx_fp_numeric.py + simpler_setup/mx_fp_numeric.py: E4M3/E2M1/E8M0 decode, MXFP8/MXFP4 golden with E8M0 block scales and ZZ/NN packing helpers.
• ST tests/st/a5/tensormap_and_ringbuffer/mx_fp_gemm/ (a5 only): AIC kernel calls pto-isa TMATMUL_MX for host-prequant MXFP8 / MXFP4 (E8M0 scales, MX_A_ZZ / MX_B_NN) and checks FP32 output against the host golden.

a5sim leftover

CPU stub has no GetScaleAddr, and TLOAD does not support MX_A_ZZ / MX_B_NN, so this ST is onboard-a5 only for now. Unifying the sim path belongs in pto-isa (or a dual kernel), not this PR.

Verified (torch 2.13.0 ≥ 2.7)

• Related UT passed (E8M0/FP4 mappings executed, not skipped).
• mx_fp_gemm a5 (device 0): MXFP8_TMATMUL_MX / MXFP4_TMATMUL_MX PASSED

[coderabbitai]

Important

Review skipped

Auto incremental reviews are disabled on this repository.

Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro Plus

Run ID: af48b836-8690-4be9-bfb0-0394ca401cfb

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

• 🔍 Trigger review

📝 Walkthrough

Walkthrough

The runtime and Python bindings add four FP8/FP4 DataType values, assign their sizes and names, map available Torch dtypes, and extend unit coverage for enum exposure, conversion, tensor arguments, and packed sizes.

Changes

FP8 and FP4 DataType support

Layer / File(s)	Summary
Runtime DataType contract and Python exposure src/common/task_interface/data_type.h, python/bindings/task_interface.cpp	Adds FP8E4M3FN, FP8E5M2, FP8E8M0, and FP4E2M1 to the runtime and Python enums, with 1-byte sizes and explicit names.
Torch interop mapping and validation simpler_setup/torch_interop.py, tests/ut/py/test_task_interface.py	Registers available Torch FP8/MXFP4 dtypes and tests mappings, tensor arguments, enum values, names, element sizes, and packed byte counts.

Estimated code review effort: 2 (Simple) | ~15 minutes

Sequence Diagram(s)

sequenceDiagram
  participant Torch
  participant EnsureTorchMap
  participant torch_dtype_to_datatype
  participant make_tensor_arg
  Torch->>EnsureTorchMap: expose available dtype attributes
  EnsureTorchMap->>torch_dtype_to_datatype: register DataType mappings
  torch_dtype_to_datatype->>make_tensor_arg: provide converted DataType
  make_tensor_arg-->>Torch: create tensor argument and report nbytes

Loading

Poem

I’m a rabbit hopping through types,
FP eights and fours in tidy stripes.
Torch maps softly, tests take flight,
One-byte formats shine bright.
Hop, hop—bindings now feel right!

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Title check	✅ Passed	The title clearly summarizes the main change: adding FP8/FP4 MX dtype support in task_interface.
Description check	✅ Passed	The description is directly related to the dtype support and test updates in this pull request.

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}

[coderabbitai]

🧹 Nitpick comments (1)

tests/ut/py/test_task_interface.py (1)

184-201: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Make the all-unavailable case an explicit skip.

If every optional dtype is absent, this loop executes zero assertions and the test passes without validating any mapping or tensor construction. Collect supported cases and call pytest.skip(...) when none are available so CI reports an inapplicable test rather than silently showing coverage.

Suggested adjustment

+        supported_cases = []
         for attr, expected in cases:
             torch_dt = getattr(torch, attr, None)
             if torch_dt is None:
                 continue
+            supported_cases.append((torch_dt, expected))
+
+        if not supported_cases:
+            pytest.skip("No FP8/FP4 dtypes are available in this PyTorch build")
+
+        for torch_dt, expected in supported_cases:
             assert torch_dtype_to_datatype(torch_dt) == expected

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@tests/ut/py/test_task_interface.py` around lines 184 - 201, The
test_torch_dtype_fp8_fp4_and_make_tensor_arg method should collect cases whose
torch dtype attributes are available, then explicitly call pytest.skip(...) when
the collection is empty. Iterate over the collected supported cases afterward so
the test still validates mappings and tensor construction whenever at least one
optional dtype exists.

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Nitpick comments:
In `@tests/ut/py/test_task_interface.py`:
- Around line 184-201: The test_torch_dtype_fp8_fp4_and_make_tensor_arg method
should collect cases whose torch dtype attributes are available, then explicitly
call pytest.skip(...) when the collection is empty. Iterate over the collected
supported cases afterward so the test still validates mappings and tensor
construction whenever at least one optional dtype exists.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro Plus

Run ID: 1af726db-0526-4c95-b7f9-e58a55449c7d

📥 Commits

Reviewing files that changed from the base of the PR and between bdf6fed and 13f9b99.

📒 Files selected for processing (4)

• python/bindings/task_interface.cpp
• simpler_setup/torch_interop.py
• src/common/task_interface/data_type.h
• tests/ut/py/test_task_interface.py

[ChaoZheng109]

这几个 dtype 是 A5-only(被 A5 独有的 pto.tquant.mx / tmatmul.mx 消费),但新增的测试全在 tests/ut/py/test_task_interface.py,是纯 host 侧、arch 无关的断言(枚举值 / element_size / 名称 / torch→DataType 映射)。CI 里它们只被 pytest tests/ut -m "not requires_hardware"(github-hosted、无 --platform、无硬件)跑一次,在 a2a3 语境下也照样通过。

问题:没有任何 a5 / a5sim 的 scene test 真正让 FP8/FP4 张量走一遍 A5 上的 host↔device 往返。也就是说,这些 UT 证明的是"枚举注册了、host 映射对了",而不是"这些 dtype 在 A5 上真能用"。想确认两点:

1. 是否应该补一个 a5sim(至少)的 st,让 FP8/FP4 经过 task interface 的 host↔device 传输并做校验?这样能先在本地验证通过,CI 恢复后也能及时发现回归。
2. float8_e8m0fnu 需要 torch ≥ 2.7、float4_e2m1fn_x2 也要较新的 torch;当前用 getattr(...) is None: continue 守卫,torch 偏旧时这两个 case 会被静默跳过(本机 torch 2.5.1 实测就是如此)。请确认 PR 描述里 "118 passed" 是在 torch ≥ 2.7 上跑的,否则 E8M0 / FP4 连 host 映射断言都没被执行到。

[yanghaoran29]

针对 @ChaoZheng109 的两点确认，已在 torch 2.13.0（≥2.7，含 float8_e8m0fnu / float4_e2m1fn_x2）下实测：

2) host 映射是否被静默跳过

不再静默跳过：UT 在三个可选 dtype 全部缺失时才会 pytest.skip；本机三者均可用，test_torch_dtype_fp8_fp4_and_make_tensor_arg 实际执行了 E4M3FN / E8M0 / FP4 映射 + make_tensor_arg。

torch 2.13.0+cpu
OK mapped float8_e4m3fn -> FP8E4M3FN
OK mapped float8_e8m0fnu -> FP8E8M0
OK mapped float4_e2m1fn_x2 -> FP4E2M1
tests/ut/py/test_task_interface.py  →  116 passed

（FP8E5M2 在 A5 上无完整支持，已删除。）

1) a5 / a5sim ST host↔device 往返

新增 tests/st/a5/tensormap_and_ringbuffer/mx_dtype_copy/：orch 用 Tensor::nbytes()（host 侧 get_element_size）算字节数，AIC 按字节拷贝。

platform	FP8E4M3FN	FP8E8M0	FP4E2M1
a5sim	PASSED	PASSED	PASSED
a5 (device 0 via task-submit)	PASSED	PASSED	PASSED

tip: 3e42362a

[ChaoZheng109]

关于新增的 simpler_setup/mx_fp_numeric.py:它只被两个测试文件引用(a5 mx_fp_gemm ST + test_mx_fp_numeric.py UT),不被任何 production 代码 import,也没在 simpler_setup/__init__.py 导出——本质是 ST golden + 造数的测试脚手架。

建议挪到 simpler_setup/goldens/ 子包(如 simpler_setup/goldens/mx_fp_gemm.py),理由:

• simpler_setup/goldens/ 就是既有的"scene test / example 共享 golden 参考"约定,__init__.py 原话:"Shared golden reference compute modules for scene tests and examples." 现有的 paged_attention.py / qwen3_14b_decode.py 就在那里,被约 30 个 test/example 以 from simpler_setup.goldens.xxx import ... 引用。
• 放在 simpler_setup 顶层会让这个纯测试用的数值库跟 KernelCompiler / SceneTestCase / platform_info 这些真正的用户 API 平级,语义上不对。
• 挪进 goldens/ 后仍随 wheel 一起打包(wheel.packages 已含 simpler_setup),不影响 ST/UT import,只是层级和现有 golden 保持一致。

(这个文件本身是合理的——之前的 dtype torch 原生就能算 golden,MX FP8/FP4 因为 torch 的 fp8/fp4 支持有限 + E8M0 scale 的 ZZ/NN 打包是 pto-isa 特有布局,才第一次需要手写 decode + 参考 matmul。只是归位问题。)

[yanghaoran29]

@ChaoZheng109 已按建议调整：

• 将 simpler_setup/mx_fp_numeric.py 挪到 simpler_setup/goldens/mx_fp_gemm.py，与现有 paged_attention / qwen3_14b_decode golden 约定对齐。
• UT / ST 的 import 已改为 from simpler_setup.goldens.mx_fp_gemm import ...；顶层不再保留该测试脚手架。

感谢指正。