Experimental FlyDSL GEMM backend for TE PyTorch (BF16/FP16/FP32/FP8/MXFP8)

[aris134]

Summary

This PR introduces an experimental FlyDSL GEMM backend for the Transformer Engine PyTorch frontend, following the integration approach established in #667. The initial implementation currently targets gfx950.

The implementation connects FlyDSL to the existing GEMM dispatch path and adds the infrastructure required to invoke FlyDSL kernels from Transformer Engine. FlyDSL is integrated as an optional PyTorch dependency for ROCm builds. When NVTE_USE_FLYDSL=1 is set during installation, setup.py adds the flydsl package to the Python installation requirements.

The supplied kernels are built around a common four-wave GEMM design modeled after the core scheduling and data-movement structure used by HipKittens-style kernels, with datatype- and layout-specific adaptations layered on top.

The backend remains opt-in at runtime and is selected with NVTE_USE_FLYDSL=1. FlyDSL modules are imported lazily only when this path is selected. When the variable is unset, the active device is unsupported, or a GEMM configuration is not supported by the FlyDSL backend, Transformer Engine continues to use the existing GEMM backend.

This PR provides preliminary FlyDSL GEMM support for FP32, FP16, BF16, tensor-wise FP8, mixed E4M3/E5M2 FP8 inputs, and MXFP8.

Supported Data Types

Data Type	Functional Support	Performance Results
FP32	Enabled	Not targeted
FP16	Enabled	Preliminary; rerun pending
BF16	Enabled	Available
Mixed FP8 E4M3/E5M2	Enabled	Available
MXFP8	Enabled	Available

Validation

Validation was added at both the kernel and model-integration levels.

• Added user-facing GEMM tests under tests/pytorch/flydsl_kernels/ covering:

  • FP32, FP16, and BF16 inputs
  • Tensor-wise FP8, including mixed E4M3/E5M2 input formats
  • MXFP8
  • TN, NN, and NT GEMM layouts
  • Comparison against both torch.matmul references and the existing Transformer Engine GEMM backend
  • Multidimensional FP8 inputs and flattening behavior

• Added test_linear_accuracy_flydsl to tests/pytorch/test_numerics.py to validate the FlyDSL backend through the public Transformer Engine Linear module. The test compares the native and FlyDSL execution paths for:

  • Forward output
  • Input gradient (dgrad)
  • Weight gradient (wgrad)
  • FP32, FP16, and BF16 parameter types
  • Unquantized, tensor-wise FP8, and MXFP8 execution
  • Multiple batch sizes and model configurations

The model-level test uses identical parameters and inputs for the native and FlyDSL paths, resets FP8 state between executions, and verifies that unsupported FlyDSL configurations fall back through the existing Transformer Engine dispatch path.

Performance

All performance was measured across an 81-shape LLM GEMM suite covering Llama 2, Llama 3.1, Qwen 2.5, and Mistral dimensions at microbatch sizes 1, 2, and 4. The complete shape list and benchmark methodology are included in the attached benchmark script.

MXFP8

MXFP8 performance was evaluated against the existing TE HipKittens MXFP8 GEMM backend. All 81 MXFP8 configurations completed successfully with no skips or validation failures.

Metric	Median Speedup	Geometric Mean	Wins
End-to-end forward	1.064x	1.042x	69 / 81
End-to-end backward	1.018x	1.000x	45 / 81
Forward GEMM	1.084x	1.122x	79 / 81
Dgrad GEMM	1.125x	1.185x	78 / 81
Wgrad GEMM	1.132x	1.177x	80 / 81
Combined backward GEMMs	1.128x	1.189x	80 / 81

The FlyDSL MXFP8 kernels outperform HipKittens on nearly the entire kernel-level sweep. Forward GEMM wins 79 of 81 shapes, while the combined dgrad and wgrad path wins 80 of 81 shapes. Median kernel-level speedups range from approximately 1.08x for forward GEMM to 1.13x for the backward GEMMs.

End-to-end results include quantization, tensor preparation, dispatch, and other framework overheads in addition to GEMM execution. Under this measurement, FlyDSL achieves a 1.064x median forward speedup and approximately matches HipKittens in backward by geometric mean, while retaining a 1.018x median backward speedup.

FP8

Tensor-wise FP8 performance was evaluated using the HYBRID DelayedScaling recipe against the default Transformer Engine GEMM backend across the same 81-shape LLM suite. All 81 configurations completed successfully with no skips or validation failures.

Metric	Median Speedup	Geometric Mean	Wins
End-to-end forward	0.950x	0.978x	21 / 81
End-to-end backward	0.926x	0.924x	13 / 81
Forward GEMM	0.955x	0.996x	20 / 81
Dgrad GEMM	0.968x	0.976x	29 / 81
Wgrad GEMM	0.966x	0.995x	29 / 81
Combined backward GEMMs	0.968x	0.986x	27 / 81

The default Transformer Engine backend outperforms the FlyDSL FP8 implementation on most shapes. At the kernel level, however, FlyDSL remains close to baseline: forward GEMM and wgrad are approximately at parity by geometric mean, while dgrad and the combined backward GEMMs reach approximately 0.98x and 0.99x baseline performance, respectively.

End-to-end results include FP8 quantization, scaling-state updates, tensor preparation, dispatch, and other framework overheads in addition to GEMM execution. Under this measurement, FlyDSL reaches approximately 0.95x baseline performance for forward execution and 0.93x for backward execution.

BF16

BF16 performance was evaluated against the default Transformer Engine GEMM backend across the same 81-shape LLM suite. Of the 81 configurations, 78 completed successfully and 3 were skipped.

Metric	Median Speedup	Geometric Mean	Wins
End-to-end forward	0.907x	0.915x	1 / 78
End-to-end backward	0.916x	0.903x	0 / 78
Forward GEMM	0.922x	0.928x	3 / 78
Dgrad GEMM	0.979x	0.981x	15 / 78
Wgrad GEMM	0.934x	0.933x	4 / 78
Combined backward GEMMs	0.958x	0.957x	5 / 78

The default Transformer Engine backend outperforms the FlyDSL BF16 implementation on most shapes. FlyDSL is closest to baseline on dgrad, reaching approximately 0.98x baseline performance by both median and geometric mean. Combined backward GEMM performance is approximately 0.96x baseline, while forward GEMM and wgrad reach approximately 0.92x and 0.93x baseline, respectively.

End-to-end results include framework overheads in addition to GEMM execution. Under this measurement, FlyDSL reaches approximately 0.91x baseline performance for both forward and backward execution.

The three skipped configurations were the largest shapes in the sweep:

• Llama-3.1-405B MBS=4: M=32768, N=16384, K=53248
• Qwen2.5-7B MBS=4: M=32768, N=37888, K=3584
• Qwen2.5-72B MBS=4: M=32768, N=59136, K=8192

These cases failed before kernel execution while packing a FlyDSL launch argument into a signed 32-bit integer, producing:

struct.error: 'i' format requires -2147483648 <= number <= 2147483647

FP32

Functional. FP32 optimization was not a target for this PR.

Artifacts

Performance Scripts

benchmark_mxfp8_flydsl_vs_baseline_suite.py

benchmark_bf16_flydsl_vs_baseline_suite.py

benchmark_fp8_flydsl_vs_baseline_suite.py

[ipanfilo]

Bot this PR and #667 create specialized test_gemm.py. It should rather be generic test that is called with different backends selected with env and it should also be single backend selection env instead of env per backend