Skip to content

[WIP] Experimental FlyDSL GEMM backend for TE PyTorch (BF16/FP16/FP32/FP8/MXFP8) #676

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
aris134 wants to merge 17 commits into
base: dev
Choose a base branch
from
Draft

[WIP] Experimental FlyDSL GEMM backend for TE PyTorch (BF16/FP16/FP32/FP8/MXFP8) #676

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
17 commits
Select commit Hold shift + click to select a range
cf54ce6
Add optional FlyDSL dependency for ROCm PyTorch builds via NVTE_USE_F…
aris134 Jul 22, 2026
2ed8285
Wire FlyDSL into the TransformerEngine MXFP8 GEMM dispatch path
aris134 Jul 22, 2026
17b9b74
Add initial support for TN BF16 FlyDSL GEMM backend
aris134 Jul 22, 2026
155e81e
Add initial support for TN FP16 FlyDSL GEMM backend
aris134 Jul 22, 2026
492a29c
Add initial support for TN FP8 FlyDSL GEMM backend
aris134 Jul 22, 2026
6ef76b3
Add initial support for TN FP32 FlyDSL GEMM backend
aris134 Jul 22, 2026
9a9462e
add layout support for NN/NT FlyDSL GEMM
aris134 Jul 22, 2026
4c9543f
add support for bf16/fp32 output types for flydsl mxfp8 gemm
aris134 Jul 22, 2026
b9eaa8a
add support for bf16/fp32 output types for flydsl fp8 gemm
aris134 Jul 22, 2026
71c4ef4
add broad output dtype support for flydsl fp8/fp16/bf16 gemms
aris134 Jul 22, 2026
aa19610
add mixed e4m3/e5m2 fp8 dtype support for flydsl fp8 GEMM
aris134 Jul 22, 2026
2ee10d3
add mixed e4m3/e5m2 fp8 dtype support for flydsl mxfp8 GEMM
aris134 Jul 22, 2026
979f38c
add pytorch flydsl gemm tests
aris134 Jul 22, 2026
e1896cd
add the actual pytorch flydsl gemm tests
aris134 Jul 22, 2026
838e64f
add e2e flydsl test in test_numerics
aris134 Jul 22, 2026
9670650
add flydsl gemm fallback support
aris134 Jul 22, 2026
ede7ace
Add direct FP8 NN/NT FlyDSL GEMM specializations
aris134 Jul 23, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 8 additions & 0 deletions setup.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -156,6 +156,14 @@ def setup_requirements() -> Tuple[List[str], List[str]]:
]
test_reqs: List[str] = ["pytest>=8.2.1"]

# Optional FlyDSL dependency for ROCm PyTorch builds.
if (
rocm_build()
and "pytorch" in frameworks
and bool(int(os.getenv("NVTE_USE_FLYDSL", "0")))
):
install_reqs.extend(["flydsl"])

# Framework-specific requirements
if not bool(int(os.getenv("NVTE_RELEASE_BUILD", "0"))):
if "pytorch" in frameworks:
Expand Down
Loading
Loading