feat(core): add FP8E8M0 dtype for MX block-scale exponents - #2144
Conversation
📝 WalkthroughWalkthroughAdds FP8E8M0 to the C++ and Python dtype APIs, PyTorch and annotation conversions, specialization, runtime and MLIR code generation, and DataType unit tests. ChangesFP8E8M0 dtype support
Estimated code review effort: 2 (Simple) | ~10 minutes Sequence Diagram(s)sequenceDiagram
participant Torch
participant DtypeInterop
participant TypeResolver
participant Specializer
participant Codegen
Torch->>DtypeInterop: Provide optional float8/MX dtype
DtypeInterop->>TypeResolver: Resolve FP8E8M0 annotation
TypeResolver->>Specializer: Pass DataType.FP8E8M0
Specializer->>Codegen: Emit runtime and MLIR dtype strings
Possibly related PRs
Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
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. Comment |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 63174728b0
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
- Map IR fp4 → torch.float4_e2m1fn_x2 in _DATATYPE_TO_TORCH (Codex P2)
- Document pl.FP8E8M0 in en/zh language guide dtype tables (Codex P2)
Interface alignment (this PR's dtype surface):
pto-isa (A5 / MX):
float8_e8m0_t — MXTypes.hpp (E8M0, bias=127); A5 TypeGet → vector_f8e8m0
float8_e4m3_t / float8_e5m2_t — MXFP8 data element types
float4_e2m1x2_t — packed MXFP4 E2M1x2
TileType::ScaleLeft / ScaleRight — L0A/L0B MX scale buffers (type.hpp);
used with float8_e8m0_t (e.g. TLoad MX_A_ZZ / MX_B_NN)
PTOAS (v0.48+ dialect / EmitC):
!pto.f8E8M0 → float8_e8m0_t
f8E4M3FN → float8_e4m3_t
f8E5M2 → float8_e5m2_t
!pto.f4E2M1x2 → float4_e2m1x2_t
loc=scaling + !pto.f8E8M0 → TileType::ScaleLeft|ScaleRight
(ui8 + scaling wrongly maps to Fixpipe TileType::Scaling)
PyPTO ↔ interface:
DataType::FP8E8M0 ↔ !pto.f8E8M0 / float8_e8m0_t / torch.float8_e8m0fnu
DataType::FP8E4M3FN ↔ f8E4M3FN / float8_e4m3_t / torch.float8_e4m3fn
DataType::FP8E5M2 ↔ f8E5M2 / float8_e5m2_t / torch.float8_e5m2
DataType::FP4 ↔ !pto.f4E2M1x2 / float4_e2m1x2_t / torch.float4_e2m1fn_x2
Introduce DataType::FP8E8M0 (1-byte E8M0 MX scale exponent) through the
dtype system, EmitC/MLIR spellings, Python bindings/stubs/DSL resolver,
JIT torch maps, and language-guide dtype tables. Also fills previously
missing EmitC/MLIR mappings for FP8E4M3FN, FP8E5M2, and packed FP4, and
maps IR fp4 → torch.float4_e2m1fn_x2 in _DATATYPE_TO_TORCH. No new
memory spaces or ops.
Interface alignment (dtype surface):
pto-isa (A5 / MX):
float8_e8m0_t — MXTypes.hpp (E8M0, bias=127); A5 TypeGet → vector_f8e8m0
float8_e4m3_t / float8_e5m2_t — MXFP8 data element types
float4_e2m1x2_t — packed MXFP4 E2M1x2
TileType::ScaleLeft / ScaleRight — L0A/L0B MX scale buffers; used with
float8_e8m0_t (e.g. TLoad MX_A_ZZ / MX_B_NN)
PTOAS (v0.48+ dialect / EmitC):
!pto.f8E8M0 → float8_e8m0_t
f8E4M3FN → float8_e4m3_t
f8E5M2 → float8_e5m2_t
!pto.f4E2M1x2 → float4_e2m1x2_t
loc=scaling + !pto.f8E8M0 → TileType::ScaleLeft|ScaleRight
(ui8 + scaling wrongly maps to Fixpipe TileType::Scaling)
PyPTO ↔ interface:
FP8E8M0 ↔ !pto.f8E8M0 / float8_e8m0_t / torch.float8_e8m0fnu
FP8E4M3FN ↔ f8E4M3FN / float8_e4m3_t / torch.float8_e4m3fn
FP8E5M2 ↔ f8E5M2 / float8_e5m2_t / torch.float8_e5m2
FP4 ↔ !pto.f4E2M1x2 / float4_e2m1x2_t / torch.float4_e2m1fn_x2
2df907d to
1ea80eb
Compare
Add Ascend950 LeftScale/RightScale memory spaces (L0A/L0B MX scale
buffers, SoC capacity, loc=scaling) and the scale data path: tile.load
mx_layout, Mat↔Scale move with target_shape and deferred fill until
tget_scale_addr, FP4/FP8E8M0 reinterpret_view aliases, and ptoas ND-flat
→ mx_a_zz preprocess. No MX quant/matmul ops yet.
Interface alignment (memspace + scale data path):
pto-isa (A5 / MX):
TileType::ScaleLeft / ScaleRight — L0A/L0B MX scale buffers
(include/pto/common/type.hpp); payload float8_e8m0_t
Layout::MX_A_ZZ / MX_A_ND / MX_A_DN — activation scale GM layouts;
TLoadMxCube* (AZZ2ZZ / AND2ZZ / …); SFractal=32 row-major ZZ
Layout::MX_B_NN / MX_B_ND / MX_B_DN — weight scale layouts;
col-major NN for MX_B_*
TMov CommonCheckMX — uint8_t Mat → float8_e8m0 ScaleLeft/Right
GetScaleAddr(Left/Right) — bind Scale tile to L0 sidecar before fill
(sidecar bind-then-fill; Mat→Mat tmov illegal on A5)
PTOAS (v0.48+ dialect / EmitC):
loc=scaling + !pto.f8E8M0 → TileType::ScaleLeft|ScaleRight
(ui8 + scaling wrongly maps to Fixpipe TileType::Scaling;
LeftScale/RightScale share one MLIR loc `scaling` until distinct)
#pto.layout<mx_a_zz|mx_a_nd|…|mx_b_nn|…> on make_tensor_view / tload
→ Layout::MX_*; InferPTOLayout prefers mx_a_zz/mx_b_nn for
!pto.f8E8M0 Mat+fractal=32
pto.tmov Mat→scaling — must follow tget_scale_addr (matches
PTOA5NormalizeTMovPass); PyPTO defers fill until then
EmitC preprocess: float8_e8m0/uint8 MX_A_ZZ → MX_A_ND so TLoad uses
AND2ZZ (activation scales are ND-flat on GM); MTE3 drain before
Vec TPUSH; pipe_barrier before MX_A_ND TLOAD after AIV store
PyPTO ↔ interface:
MemorySpace::LeftScale ↔ loc=scaling + !pto.f8E8M0 ↔ ScaleLeft
MemorySpace::RightScale ↔ loc=scaling + !pto.f8E8M0 ↔ ScaleRight
tile.load(..., mx_layout=mx_a_zz|…) ↔ #pto.layout<…> ↔ Layout::MX_*
tile.move Mat→LeftScale/RightScale ↔ deferred tmov after GetScaleAddr
FP8E8M0 / UINT8 scale payload ↔ float8_e8m0_t (1-byte MX exp)
Dependencies:
- PyPTO: FP8E8M0 dtype (hw-native-sys#2144 / DataType::FP8E8M0)
- PTOAS: v0.48+ (loc=scaling, MX layouts, NormalizeTMov)
- pto-isa: A5 MX TileType + Layout + TLoadMx / TMov / GetScaleAddr
Add Ascend950 LeftScale/RightScale memory spaces (L0A/L0B MX scale buffers, SoC capacity, loc=scaling) and the scale data path: tile.load mx_layout, Mat↔Scale move with target_shape and deferred fill until tget_scale_addr, FP4/FP8E8M0 reinterpret_view aliases, and ptoas ND-flat → mx_a_zz preprocess. No MX quant/matmul ops yet. Phase 2 of the hw-native-sys#2117 split (includes former Phase 1.2). Depends on FP8E8M0 (hw-native-sys#2144). Refs hw-native-sys#1975. Interface alignment (memspace + scale data path): pto-isa (A5 / MX): TileType::ScaleLeft / ScaleRight — L0A/L0B MX scale buffers (include/pto/common/type.hpp); payload float8_e8m0_t Layout::MX_A_ZZ / MX_A_ND / MX_A_DN — activation scale GM layouts; TLoadMxCube* (AZZ2ZZ / AND2ZZ / …); SFractal=32 row-major ZZ Layout::MX_B_NN / MX_B_ND / MX_B_DN — weight scale layouts; col-major NN for MX_B_* TMov CommonCheckMX — uint8_t Mat → float8_e8m0 ScaleLeft/Right GetScaleAddr(Left/Right) — bind Scale tile to L0 sidecar before fill (sidecar bind-then-fill; Mat→Mat tmov illegal on A5) PTOAS (v0.48+ dialect / EmitC): loc=scaling + !pto.f8E8M0 → TileType::ScaleLeft|ScaleRight (ui8 + scaling wrongly maps to Fixpipe TileType::Scaling; LeftScale/RightScale share one MLIR loc `scaling` until distinct) #pto.layout<mx_a_zz|mx_a_nd|…|mx_b_nn|…> on make_tensor_view / tload → Layout::MX_*; InferPTOLayout prefers mx_a_zz/mx_b_nn for !pto.f8E8M0 Mat+fractal=32 pto.tmov Mat→scaling — must follow tget_scale_addr (matches PTOA5NormalizeTMovPass); PyPTO defers fill until then EmitC preprocess: float8_e8m0/uint8 MX_A_ZZ → MX_A_ND so TLoad uses AND2ZZ (activation scales are ND-flat on GM); MTE3 drain before Vec TPUSH; pipe_barrier before MX_A_ND TLOAD after AIV store PyPTO ↔ interface: MemorySpace::LeftScale ↔ loc=scaling + !pto.f8E8M0 ↔ ScaleLeft MemorySpace::RightScale ↔ loc=scaling + !pto.f8E8M0 ↔ ScaleRight tile.load(..., mx_layout=mx_a_zz|…) ↔ #pto.layout<…> ↔ Layout::MX_* tile.move Mat→LeftScale/RightScale ↔ deferred tmov after GetScaleAddr FP8E8M0 / UINT8 scale payload ↔ float8_e8m0_t (1-byte MX exp) Dependencies: - PyPTO: FP8E8M0 dtype (hw-native-sys#2144 / DataType::FP8E8M0) - PTOAS: v0.48+ (loc=scaling, MX layouts, NormalizeTMov) - pto-isa: A5 MX TileType + Layout + TLoadMx / TMov / GetScaleAddr Also includes review/CI fixes for hw-native-sys#2147: - Scope tile.reshape layout inheritance to MX scale fractal=32 only - Stamp memory_space_ on tile.move only for LeftScale/RightScale - FIFO-sticky: registered tpush_to_*/tpop_from_* + Scale moves only - Drop drive-by Bias from IsCubeMemorySpace; do not change .gitignore - Default mx_layout loads to Mat; tighten mx_layout / flat-store checks - Ascend950 LeftScale/RightScale capacity and Mat→Scale path UTs
Add Ascend950 LeftScale/RightScale memory spaces (L0A/L0B MX scale buffers, SoC capacity, loc=scaling) and the scale data path: tile.load mx_layout, Mat↔Scale move with target_shape and deferred fill until tget_scale_addr, FP4/FP8E8M0 reinterpret_view aliases, and ptoas ND-flat → mx_a_zz preprocess. No MX quant/matmul ops yet. Phase 2 of the hw-native-sys#2117 split (includes former Phase 1.2). Depends on FP8E8M0 (hw-native-sys#2144). Refs hw-native-sys#1975. Also includes review fixes for hw-native-sys#2147: - Drop unused TakePending*/DeferTFree/TakeDeferredTFreeCount (flush API already lives in feat/mx-quant-ops and feat/mx-matmul-ops) - Assert no unflushed deferred scale fills at GenerateFunction end - Use INTERNAL_UNREACHABLE_SPAN for the unsupported f8 Mat→Scale path - Align _barrier_before_mx_a_nd_tload comment with no-distance-cutoff code Co-authored-by: Cursor <cursoragent@cursor.com>
Add Ascend950 LeftScale/RightScale memory spaces (L0A/L0B MX scale buffers, SoC capacity, loc=scaling) and the scale data path: tile.load mx_layout, Mat↔Scale move with target_shape and deferred fill until tget_scale_addr, FP4/FP8E8M0 reinterpret_view aliases, and ptoas ND-flat → mx_a_zz preprocess. No MX quant/matmul ops yet. Phase 2 of the hw-native-sys#2117 split (includes former Phase 1.2). Depends on FP8E8M0 (hw-native-sys#2144). Refs hw-native-sys#1975. Also includes review/CI fixes for hw-native-sys#2147: - Drop unused TakePending*/DeferTFree/TakeDeferredTFreeCount (flush API already lives in feat/mx-quant-ops and feat/mx-matmul-ops) - Assert no unflushed deferred scale fills at GenerateFunction end - Use INTERNAL_UNREACHABLE_SPAN for the unsupported f8 Mat→Scale path - Align _barrier_before_mx_a_nd_tload comment with no-distance-cutoff code - Narrow TileType.shape dims via isinstance(ConstInt) for pyright Co-authored-by: Cursor <cursoragent@cursor.com>
Add Ascend950 LeftScale/RightScale memory spaces (L0A/L0B MX scale buffers, SoC capacity, loc=scaling) and the scale data path: tile.load mx_layout, Mat↔Scale move with target_shape and deferred fill until tget_scale_addr, FP4/FP8E8M0 reinterpret_view aliases, and ptoas ND-flat → mx_a_zz preprocess. No MX quant/matmul ops yet. Phase 2 of the hw-native-sys#2117 split (includes former Phase 1.2). Depends on FP8E8M0 (hw-native-sys#2144). Refs hw-native-sys#1975. Also includes review/CI fixes for hw-native-sys#2147: - Drop unused TakePending*/DeferTFree/TakeDeferredTFreeCount (flush API already lives in feat/mx-quant-ops and feat/mx-matmul-ops) - Assert no unflushed deferred scale fills at GenerateFunction end - Use INTERNAL_UNREACHABLE_SPAN for the unsupported f8 Mat→Scale path - Align _barrier_before_mx_a_nd_tload comment with no-distance-cutoff code - Narrow TileType.shape dims via isinstance(ConstInt) for pyright
Add Ascend950 LeftScale/RightScale memory spaces (L0A/L0B MX scale buffers, SoC capacity, loc=scaling) and the scale data path: tile.load mx_layout, Mat↔Scale move with target_shape and deferred fill until tget_scale_addr, FP4/FP8E8M0 reinterpret_view aliases, and ptoas ND-flat → mx_a_zz preprocess. No MX quant/matmul ops yet. Phase 2 of the hw-native-sys#2117 split (includes former Phase 1.2). Depends on FP8E8M0 (hw-native-sys#2144). Refs hw-native-sys#1975. Also includes review/CI fixes for hw-native-sys#2147: - Drop unused TakePending*/DeferTFree/TakeDeferredTFreeCount (flush API already lives in feat/mx-quant-ops and feat/mx-matmul-ops) - Assert no unflushed deferred scale fills at GenerateFunction end - Use INTERNAL_UNREACHABLE_SPAN for the unsupported f8 Mat→Scale path - Align _barrier_before_mx_a_nd_tload comment with no-distance-cutoff code - Narrow TileType.shape dims via isinstance(ConstInt) for pyright after rewrite; raise if PTOAS changed its layout spelling. count equals expected Vec-TPUSH count.
Add Ascend950 LeftScale/RightScale memory spaces (L0A/L0B MX scale buffers, SoC capacity, loc=scaling) and the scale data path: tile.load mx_layout, Mat↔Scale move with target_shape and deferred fill until tget_scale_addr, and FP8E8M0 reinterpret_view aliases. Reject FP8E5M2 reinterpret on Ascend950. No MX quant/matmul ops or MXFP4 aliases yet (those land in hw-native-sys#2117 with ND rewrite / AIC PIPE_ALL barriers). Phase 2 of the hw-native-sys#2117 split (includes former Phase 1.2). Depends on FP8E8M0 (hw-native-sys#2144). Refs hw-native-sys#1975. Also includes review/CI fixes for hw-native-sys#2147: - Drop unused TakePending*/DeferTFree/TakeDeferredTFreeCount (flush API already lives in feat/mx-quant-ops and feat/mx-matmul-ops) - Assert no unflushed deferred scale fills at GenerateFunction end - Use INTERNAL_UNREACHABLE_SPAN for the unsupported f8 Mat→Scale path - Narrow TileType.shape dims via isinstance(ConstInt) for pyright - Drop fragile PTOAS C++ regex post-process (MX_A_ZZ→ND, MTE3 drain, pipe_barrier); defer sync/layout fixes to hw-native-sys#2117
Add Ascend950 LeftScale/RightScale memory spaces and the scale data path (tile.load mx_layout, Mat→Scale move with deferred fill), plus tile.tget_scale_addr (flush PendingScaleFill) and tile.matmul_mx only (pto.tmatmul.mx). Host-prequant FP8E4M3FN data + FP8E8M0 scales; reject FP8E5M2 for reinterpret and matmul_mx data dtype. Document pto-isa and PTOAS constraints in operators docs (en/zh). Defer to hw-native-sys#2117: tquant/tdequant, MXFP4, matmul_mx_acc/bias, ND rewrite / AIC PIPE_ALL, ExpandMixed FIFO sticky. Depends on FP8E8M0 (hw-native-sys#2144). Refs hw-native-sys#1975.
Add Ascend950 LeftScale/RightScale memory spaces and the scale data path (tile.load mx_layout, Mat→Scale move with deferred fill), plus tile.tget_scale_addr (flush PendingScaleFill) and tile.matmul_mx (pto.tmatmul.mx). Host-prequant FP8E4M3FN data + FP8E8M0 scales; reject FP8E5M2 for reinterpret and matmul_mx data dtype. Document pto-isa and PTOAS constraints in operators docs (en/zh). Depends on FP8E8M0 (hw-native-sys#2144). Refs hw-native-sys#1975.
Summary
Introduce
DataType::FP8E8M0(1-byte E8M0 MX scale exponent) through the dtype system, EmitC/MLIR spellings (float8_e8m0_t/!pto.f8E8M0), and Python bindings/stubs/DSL resolver/JIT torch maps. Also fills previously missing EmitC/MLIR mappings for FP8E4M3FN, FP8E5M2, and packed FP4. No new memory spaces or ops.pto-isa / PTOAS interface alignment
This PR only adds the dtype / type-string surface that later MX PRs wire into ops. Mapping:
pto-isa (A5 MX)
float8_e8m0_tinclude/pto/cpu/MXTypes.hpp; A5TypeGet→vector_f8e8m0float8_e4m3_t/float8_e5m2_tfloat4_e2m1x2_tTileType::ScaleLeft/ScaleRightinclude/pto/common/type.hpp); scale tiles usefloat8_e8m0_t(e.g.TLoadwithMX_A_ZZ/MX_B_NN)PTOAS (v0.48+ dialect + EmitC)
!pto.f8E8M0float8_e8m0_tf8E4M3FNfloat8_e4m3_tf8E5M2float8_e5m2_t!pto.f4E2M1x2float4_e2m1x2_tloc=scaling+!pto.f8E8M0TileType::ScaleLeft/ScaleRight(not FixpipeTileType::Scaling; bareui8+scalingmaps wrongly)Bare keywords
f8E8M0/f4E2M1x2do not parse in PTOAS — dialect types are required.PyPTO
DataTypeFP8E8M0!pto.f8E8M0float8_e8m0_tfloat8_e8m0fnuFP8E4M3FNf8E4M3FNfloat8_e4m3_tfloat8_e4m3fnFP8E5M2f8E5M2float8_e5m2_tfloat8_e5m2FP4!pto.f4E2M1x2float4_e2m1x2_tfloat4_e2m1fn_x2Test plan
pytest tests/ut/core/test_dtype.py -qpl.FP8E8M0(en + zh)_DATATYPE_TO_TORCHincludesfp4→float4_e2m1fn_x2when torch exposes it