Experimental FlyDSL GEMM backend for TE PyTorch (BF16/FP16/FP32/FP8/MXFP8) - #676
Experimental FlyDSL GEMM backend for TE PyTorch (BF16/FP16/FP32/FP8/MXFP8)#676aris134 wants to merge 33 commits into
Conversation
aris134
changed the title
aris134
changed the title
Add dedicated FlyDSL FP8 NN and NT kernels alongside the existing TN path. * dispatch TN, NN, and NT to specialized kernels * select matching TE rowwise or columnwise storage without copies * preserve operand scales and independent FP8 dtypes * derive M/N/K from each kernel’s physical layout * preserve TE output shapes while flattening only for launch * validate unsupported layouts and shapes for controlled fallback
Route all tensorwise FP8 layouts through the common FP8 GEMM core after wrapper-side storage backing selection. Remove the redundant FP8 NN and NT kernel variants since columnwise FP8 storage already provides the required materialized transpose.
Replace the eager PyTorch MXFP8 scale-packing path with stride-aware FlyDSL kernels that directly convert TE E8M0 scales into the HK/MFMA-ready [K/128, dim] packed layout. The previous implementation composed packing from arange, indexing, casts, shifts, masks, ORs, transposes, and contiguous copies. PyTorch lowered these into dozens of small GPU kernels around every GEMM, which dominated end-to-end runtime despite the FlyDSL GEMMs themselves being faster. The new path: - launches one fused scale-pack kernel per GEMM operand - supports both rowwise and columnwise TE scale layouts - consumes non-contiguous scale views using their actual strides - eliminates the intermediate iteration-major scale tensor - removes eager transpose/contiguous preparation from the scale path - preserves the existing HK/MFMA-ready packed representation
aris134
changed the title
There was a problem hiding this comment.
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
![github-actions[bot]](https://avatars.githubusercontent.com/in/15368?s=60&v=4){kind=small}
| *args, | ||
| **kwargs, | ||
| ) | ||
| except FlyDSLUnsupportedError as exc: |
There was a problem hiding this comment.
Blocking: the fallback only catches one exception type, so most unsupported configs crash instead of falling back.
The PR description says "when ... a GEMM configuration is not supported by the FlyDSL backend, Transformer Engine continues to use the existing GEMM backend." That is only true for FlyDSLUnsupportedError. te_generic_gemm_flydsl rejects most unsupported configurations with other exception types, which propagate straight out of general_gemm and kill the training run:
| Condition | Raised as | Reachable from |
|---|---|---|
bias is not None |
NotImplementedError (gemm_wrappers.py:117) |
linear.py:518 — any te.Linear(bias=True) |
quantizer is not None |
NotImplementedError (gemm_wrappers.py:101) |
linear.py:516 quantization_params=output_quantizer |
accumulate=True |
NotImplementedError (gemm_wrappers.py:111) |
wgrad accumulation |
| TT layout | NotImplementedError (gemm_wrappers.py:1295) |
— |
unsupported output_dtype |
NotImplementedError (×5) |
— |
| mixed quantized/regular operands | ValueError (gemm_wrappers.py:1315, 1349) |
— |
| NVFP4 / blockwise operands | ValueError (gemm_wrappers.py:150) |
NVFP4 recipes |
| unsupported input dtype | NotImplementedError (gemm_wrappers.py:1444) |
fp64, mixed bf16/fp16 |
NotImplementedError derives from RuntimeError, not from FlyDSLUnsupportedError, so except FlyDSLUnsupportedError does not catch it.
Net effect: with NVTE_USE_FLYDSL=1 on gfx950, a plain te.Linear(..., bias=True) raises NotImplementedError on the first forward pass. That makes the backend unusable for most real models rather than gracefully degrading.
Two options:
- Make every config-dependent rejection in
gemm_wrappers.pyraiseFlyDSLUnsupportedError(keepTypeError/ValueErroronly for genuine programming errors, i.e. states TE itself should never produce), or - Broaden the catch here to
(FlyDSLUnsupportedError, NotImplementedError).
Option 1 is safer — a blanket catch would also swallow real bugs inside the kernels.
| def _validate_common_epilogue( | ||
| *, | ||
| quantizer, | ||
| bias, | ||
| gelu, | ||
| grad, | ||
| accumulate, | ||
| alpha, | ||
| beta, | ||
| ): | ||
| """Validate features not yet implemented by the FlyDSL GEMM backend.""" | ||
| if quantizer is not None: | ||
| raise NotImplementedError( | ||
| "FlyDSL GEMM output quantization is not implemented" | ||
| ) | ||
|
|
||
| if float(alpha) != 1.0 or float(beta) != 0.0: | ||
| raise NotImplementedError( | ||
| "FlyDSL GEMM currently supports only alpha=1 and beta=0" | ||
| ) | ||
|
|
||
| # TODO: Add accumulate option | ||
| if accumulate: | ||
| raise NotImplementedError( | ||
| "FlyDSL GEMM accumulation is not implemented" | ||
| ) | ||
|
|
||
| # TODO: Add fused bias and BGRADB epilogues | ||
| if bias is not None and bias.numel() != 0: | ||
| raise NotImplementedError( | ||
| "FlyDSL GEMM bias is not implemented" | ||
| ) |
There was a problem hiding this comment.
Blocking: gelu and grad are accepted but never validated — fused GELU is silently dropped.
_validate_common_epilogue takes gelu and grad in its signature but the body only checks quantizer, alpha/beta, accumulate, and bias. te_generic_gemm_flydsl then does del gelu_in and always returns gelu_input=None.
gelu=True is reachable in production: layernorm_mlp.py:611 passes gelu=gemm_gelu_fusion and layernorm_mlp.py:1288 passes gelu=fc2_dgrad_gemm_gelu_fusion to general_gemm. With NVTE_USE_FLYDSL=1 on gfx950, a LayerNormMLP with GEMM+GELU fusion enabled gets a plain un-activated GEMM result back, plus gelu_input=None for the backward pass. That is silent numerical corruption, not a crash — the worst failure mode, and none of the added tests would catch it (neither test exercises gelu=True).
Suggest rejecting both explicitly so the fallback path handles them (see the exception-type comment on cpp_extensions/gemm.py — these need to be FlyDSLUnsupportedError to actually fall back):
| def _validate_common_epilogue( | |
| *, | |
| quantizer, | |
| bias, | |
| gelu, | |
| grad, | |
| accumulate, | |
| alpha, | |
| beta, | |
| ): | |
| """Validate features not yet implemented by the FlyDSL GEMM backend.""" | |
| if quantizer is not None: | |
| raise NotImplementedError( | |
| "FlyDSL GEMM output quantization is not implemented" | |
| ) | |
| if float(alpha) != 1.0 or float(beta) != 0.0: | |
| raise NotImplementedError( | |
| "FlyDSL GEMM currently supports only alpha=1 and beta=0" | |
| ) | |
| # TODO: Add accumulate option | |
| if accumulate: | |
| raise NotImplementedError( | |
| "FlyDSL GEMM accumulation is not implemented" | |
| ) | |
| # TODO: Add fused bias and BGRADB epilogues | |
| if bias is not None and bias.numel() != 0: | |
| raise NotImplementedError( | |
| "FlyDSL GEMM bias is not implemented" | |
| ) | |
| def _validate_common_epilogue( | |
| *, | |
| quantizer, | |
| bias, | |
| gelu, | |
| grad, | |
| accumulate, | |
| alpha, | |
| beta, | |
| ): | |
| """Validate features not yet implemented by the FlyDSL GEMM backend.""" | |
| del grad | |
| if quantizer is not None: | |
| raise FlyDSLUnsupportedError( | |
| "FlyDSL GEMM output quantization is not implemented" | |
| ) | |
| if float(alpha) != 1.0 or float(beta) != 0.0: | |
| raise FlyDSLUnsupportedError( | |
| "FlyDSL GEMM currently supports only alpha=1 and beta=0" | |
| ) | |
| # TODO: Add accumulate option | |
| if accumulate: | |
| raise FlyDSLUnsupportedError( | |
| "FlyDSL GEMM accumulation is not implemented" | |
| ) | |
| # TODO: Add fused bias and BGRADB epilogues | |
| if bias is not None and bias.numel() != 0: | |
| raise FlyDSLUnsupportedError( | |
| "FlyDSL GEMM bias is not implemented" | |
| ) | |
| # TODO: Add fused GELU epilogue | |
| if gelu: | |
| raise FlyDSLUnsupportedError( | |
| "FlyDSL GEMM fused GELU is not implemented" | |
| ) |
| del bias_type | ||
| del gelu_in | ||
| del workspace | ||
| del workspaceSize | ||
| del use_split_accumulator | ||
| del comm_overlap | ||
| del comm_type | ||
| del extra_output | ||
| del bulk_overlap |
There was a problem hiding this comment.
comm_overlap / comm_type / extra_output are silently discarded — TP comm-overlap becomes a no-op.
linear.py:513 passes ub=ub_obj, ub_type=ub_type, extra_output=reduce_scatter_out into general_gemm, which forwards them as comm_overlap/comm_type/extra_output kwargs. Here they are del'd, and the function returns extra_output=None. In cpp_extensions/gemm.py that return value is assigned back over the caller's extra_output, so Linear gets reduce_scatter_out = None and the reduce-scatter never happens.
With tensor-parallel comm+GEMM overlap enabled this produces wrong results (or a None deref downstream) rather than an error. Same class of issue as gelu above.
Please reject rather than drop:
if comm_overlap is not None or comm_type is not None or extra_output is not None:
raise FlyDSLUnsupportedError(
"FlyDSL GEMM does not support comm+GEMM overlap"
)workspace / workspaceSize / bias_type / use_split_accumulator are genuinely fine to drop — those are hints or unused given the other guards. bulk_overlap should be covered by the comm_overlap check.
| A_arg = A.view(torch.uint8).view(-1) | ||
| B_arg = B.view(torch.uint8).view(-1) | ||
| C_arg = C.view(-1) |
There was a problem hiding this comment.
The int32 launch-argument overflow from the PR description originates here — and it is not recoverable.
The PR description lists three BF16 shapes that fail with struct.error: 'i' format requires -2147483648 <= number <= 2147483647. That comes from packing these flat byte views into FlyDSL's launch signature:
- Llama-3.1-405B MBS=4 (
M=32768, K=53248):Ais 32768x53248x2 = 3.49 GB >INT32_MAX - Qwen2.5-72B MBS=4 (
M=32768, N=59136):Cis 32768x59136x2 = 3.87 GB >INT32_MAX - Qwen2.5-7B MBS=4 (
M=32768, N=37888):Cis 32768x37888x2 = 2.48 GB >INT32_MAX
struct.error is neither a FlyDSLUnsupportedError nor a subclass of one, so it escapes general_gemm and aborts the run — these shapes cannot be "skipped", they hard-fail. Large-vocab / large-MLP layers at MBS=4 are exactly the shapes people will hit.
Please add a size guard next to the existing _BLOCK_* divisibility checks in doGemm, so oversized shapes fall back to the default backend instead:
_MAX_LAUNCH_BYTES = 2**31 - 1
for name, t in (("A", A), ("B", B), ("C", C)):
if t.numel() * t.element_size() > _MAX_LAUNCH_BYTES:
raise FlyDSLUnsupportedError(
f"FlyDSL BF16 GEMM operand {name} is "
f"{t.numel() * t.element_size()} bytes, which exceeds the "
f"int32 launch-argument limit of {_MAX_LAUNCH_BYTES}"
)The same guard is needed in fp16_gemm.py, fp32_gemm.py, fp8_gemm.py, and mxfp8_gemm.py.
| if not IS_HIP_EXTENSION: | ||
| pytest.skip("FlyDSL GEMM is only supported on HIP.") |
There was a problem hiding this comment.
Blocking: this test cannot fail — it compares the native backend against itself on every GPU except gfx950.
cpp_extensions/gemm.py:497 gates the FlyDSL path on get_device_compute_capability() == (9, 5). This test only checks IS_HIP_EXTENSION, so on gfx942 (the bulk of ROCm CI) setting NVTE_USE_FLYDSL=1 is a no-op and out_ref / out_flydsl are both produced by tex.generic_gemm. The three assert_close calls are then tautologies. test_numerics.py runs in CI (ci/pytorch.sh:80), so this adds 48 green-but-empty cases to a job that is already long.
Two further gaps that persist even on gfx950:
-
No assertion that the FlyDSL path was actually taken.
NVTE_FLYDSL_GEMM_WARN_FALLBACK=1is set but the warnings are never captured. Every GEMM could raiseFlyDSLUnsupportedErrorand fall back, and the test still passes. -
model="small"always falls back.smallishidden_size = 8 * 16 = 128, so the fprop GEMM hasK=128→ 2K64tiles, belowdoGemm'snum_k_tiles >= 4minimum (bf16_gemm.py:1396). Half the parametrization exercises nothing.126m(hidden_size=768) does satisfy the tile constraints.
The file already imports get_device_compute_capability (line 42) and uses this exact pattern at line 765. Suggest:
| if not IS_HIP_EXTENSION: | |
| pytest.skip("FlyDSL GEMM is only supported on HIP.") | |
| if not IS_HIP_EXTENSION or get_device_compute_capability() != (9, 5): | |
| pytest.skip("FlyDSL GEMM is only supported on gfx950.") | |
| pytest.importorskip("flydsl", reason="FlyDSL package is not installed.") |
and then asserting the path was hit, e.g. wrapping the FlyDSL forward/backward in warnings.catch_warnings(record=True) and failing if a [FLYDSL WARNING] fallback was emitted for a config the PR claims to support. Without that, a regression that disables FlyDSL entirely would go unnoticed.
The importorskip matters because flydsl is only installed when NVTE_USE_FLYDSL=1 is set at build time (setup.py); on a default gfx950 build the lazy import at cpp_extensions/gemm.py:502 raises an uncaught ImportError and this test errors rather than skipping.
| major, minor = torch.cuda.get_device_capability() | ||
|
|
||
| # The current FlyDSL MXFP8 implementation uses the gfx950 fp8-scaled MFMA. | ||
| has_mxfp8_support = major == 9 and minor >= 5 | ||
|
|
||
| requires_mxfp8_support = pytest.mark.skipif( | ||
| not has_mxfp8_support, | ||
| reason="FlyDSL MXFP8 requires gfx950+ fp8-scaled MFMA support", | ||
| ) |
There was a problem hiding this comment.
Only the MXFP8 tests are gated on gfx950 — the fp32/fp16/bf16/FP8 tests silently pass on every other GPU.
cpp_extensions/gemm.py:497 gates all FlyDSL dispatch on get_device_compute_capability() == (9, 5), not just MXFP8. On gfx942 call_gemm(..., use_flydsl=True) sets the env var but still runs tex.generic_gemm, so:
test_flydsl_vs_pytorch_regular/_fp8/_fp8_multidimbecome tests of hipBLASLt, not FlyDSL.test_flydsl_vs_cpp_regular/_fp8become literal tautologies — both sides aretex.generic_gemmwith identical inputs.
Please apply a module-level gate covering all tests, not just MXFP8:
| major, minor = torch.cuda.get_device_capability() | |
| # The current FlyDSL MXFP8 implementation uses the gfx950 fp8-scaled MFMA. | |
| has_mxfp8_support = major == 9 and minor >= 5 | |
| requires_mxfp8_support = pytest.mark.skipif( | |
| not has_mxfp8_support, | |
| reason="FlyDSL MXFP8 requires gfx950+ fp8-scaled MFMA support", | |
| ) | |
| major, minor = torch.cuda.get_device_capability() | |
| # All FlyDSL dispatch is gated on gfx950 in cpp_extensions/gemm.py. | |
| has_flydsl_support = (major, minor) == (9, 5) | |
| pytestmark = pytest.mark.skipif( | |
| not has_flydsl_support, | |
| reason="FlyDSL GEMM dispatch requires gfx950", | |
| ) | |
| # The current FlyDSL MXFP8 implementation uses the gfx950 fp8-scaled MFMA. | |
| has_mxfp8_support = has_flydsl_support | |
| requires_mxfp8_support = pytest.mark.skipif( | |
| not has_mxfp8_support, | |
| reason="FlyDSL MXFP8 requires gfx950+ fp8-scaled MFMA support", | |
| ) |
Two related points:
- Add
pytest.importorskip("flydsl")too.flydslis only installed whenNVTE_USE_FLYDSL=1at build time, so a default gfx950 build makes the lazy import atcpp_extensions/gemm.py:502raise an uncaughtImportErrorand every test here errors instead of skipping. torch.cuda.get_device_capability()at module scope runs at collection time and initialises CUDA during collection; on a CPU-only box this errors the whole module rather than skipping it. Deferring it into a cached helper (or apytest.mark.skipifontorch.cuda.is_available()) would be more robust.
| def get_shapes(layout, M, K, N): | ||
| """Return the A/B storage shapes used by TE's public GEMM tests.""" | ||
| if layout == "TN": | ||
| return (M, K), (N, K) | ||
| if layout == "NN": | ||
| return (M, K), (K, M) | ||
| if layout == "NT": | ||
| return (M, K), (M, K) | ||
| raise ValueError(f"Unsupported layout: {layout}") |
There was a problem hiding this comment.
N is ignored for NN and NT, and no shape in the suite has M != N — so the operand-swap logic is never tested asymmetrically.
The NN branch returns (K, M) and the NT branch returns (M, K); neither reads N. That is masked today only because every entry in FLYDSL_SHAPES and MXFP8_SHAPES happens to have N == M:
(512, 512, 512) (512, 1024, 512) (1024, 512, 1024)
Working through getGemmOutputShape (csrc/extensions/gemm.cpp:50), the actual GEMMs produced are:
| layout | A | B | output |
|---|---|---|---|
| TN | (M,K) | (N,K) | (N, M) |
| NN | (M,K) | (K,M) | (K, K) |
| NT | (M,K) | (M,K) | (K, K) |
So NN and NT only ever produce square outputs, and TN only produces square ones given the current shape list. gemm_wrappers.py does a lot of M/N ownership swapping (a_flydsl = B_data; b_flydsl = A_data, plus per-layout expected_a/expected_b contracts) — a transposed-M/N bug there would produce a correctly-shaped, numerically-wrong result and every test here would still pass.
Please make N actually independent (NN → (K, N), NT → (N, K), adjusting compute_pytorch_reference to match) and add at least one M != N shape, e.g. (512, 512, 1024).
| A_flat = A_fp8.dequantize().reshape(-1, K) | ||
| B_flat = B_fp8.dequantize().reshape(-1, K) | ||
| expected = torch.matmul(B_flat, A_flat.T) | ||
|
|
||
| assert_gemm_close(output, expected, atol=5e-3, rtol=1e-2) |
There was a problem hiding this comment.
This assertion looks shape-mismatched — please confirm this test actually runs green.
general_gemm preserves B's leading dims when transb is false. For A = (batch, M, K), B = (batch, N, K), TN:
getGemmOutputShape(csrc/extensions/gemm.cpp:62-78) unflattensB0→[batch, N], then appendsA0 = batch*M⇒(batch, N, batch*M), e.g.(2, 256, 512).gemm_wrappers.py:1202derives the same thing for the FlyDSL path:(*B_payload_shape[:-1], n).
But expected here is (batch*N, batch*M) = (512, 512). torch.testing.assert_close compares shape before values and raises on a mismatch, so this should fail on both the FlyDSL and native paths regardless of GPU.
The values are right — it is only the rank. One-line fix:
| A_flat = A_fp8.dequantize().reshape(-1, K) | |
| B_flat = B_fp8.dequantize().reshape(-1, K) | |
| expected = torch.matmul(B_flat, A_flat.T) | |
| assert_gemm_close(output, expected, atol=5e-3, rtol=1e-2) | |
| A_flat = A_fp8.dequantize().reshape(-1, K) | |
| B_flat = B_fp8.dequantize().reshape(-1, K) | |
| expected = torch.matmul(B_flat, A_flat.T).reshape(output.shape) | |
| assert_gemm_close(output, expected, atol=5e-3, rtol=1e-2) |
If this test is currently passing in your runs, that would suggest it is being skipped or the shapes differ from my reading — worth checking either way, since a green result here would mean the flatten semantics aren't what the test name claims.
| def _run_fp16_gemm( | ||
| A, | ||
| transa, | ||
| B, | ||
| transb, | ||
| D, | ||
| *, | ||
| output_dtype: torch.dtype, | ||
| ): |
There was a problem hiding this comment.
_run_fp16_gemm is a line-for-line copy of _run_bf16_gemm (lines 347-450).
The two functions differ only in the string "BF16" → "FP16", torch.bfloat16 → torch.float16, and bf16_matmul → fp16_matmul. The layout dispatch table, the expected_a/expected_b contracts, the kb != k check, the _product(output_shape) != m * n check, and the output allocation are all identical.
This repeats at module scale: bf16_gemm.py and fp16_gemm.py are 1430 lines each and diff reports only 102 differing lines (~96% identical). That's ~1350 lines of duplicated kernel-construction logic whose only real difference is the MFMA input element type.
The TODO at lines 22-27 already acknowledges cross-backend duplication with the Triton work, and @ipanfilo raised the same theme on test_gemm.py. This one is worth addressing before merge rather than deferring, because every future fix to the BF16 layout contract has to be applied twice (or silently won't be).
Minimal version that doesn't require restructuring the kernels: parametrise these two wrappers on (torch_dtype, matmul_fn, name) and keep one body. For the kernel modules themselves, if the FlyDSL element type is the only difference, a single module taking the element type as a compile-time specialisation key — which the @functools.lru_cache(maxsize=None) on _cached_launch already supports — would collapse both.
| from . import gemm | ||
|
|
||
| __all__ = ["gemm"] |
There was a problem hiding this comment.
Copyright audit: missing AMD header, and missing module docstring will fail CI lint.
This is the only new file in the PR with no copyright header at all. qa/L0_pytorch_lint/test.sh also runs pylint --recursive=y transformer_engine/pytorch, which now walks this package — a module with no docstring trips C0114 missing-module-docstring.
The sibling package transformer_engine/pytorch/triton_kernels/__init__.py is the established pattern here, and notably it does not eagerly import its submodules. The eager from . import gemm chains through gemm/__init__.py → gemm_wrappers → all five kernel modules → import flydsl, so anything that touches transformer_engine.pytorch.flydsl_kernels for any reason pulls in the whole FlyDSL stack. Dropping it keeps the "lazy import" property the design depends on:
| from . import gemm | |
| __all__ = ["gemm"] | |
| # Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. | |
| # | |
| # See LICENSE for license information. | |
| """FlyDSL kernels for ROCm gfx950 GEMM replacement.""" |
The existing from ..flydsl_kernels.gemm import (...) in cpp_extensions/gemm.py:502 keeps working unchanged.
| # SPDX-License-Identifier: Apache-2.0 | ||
| # Copyright (c) 2025 FlyDSL Project Contributors | ||
| """Byte-staging helpers for the BF16 four-wave GEMM.""" |
There was a problem hiding this comment.
Copyright audit: third-party header vendored into transformer_engine/ — needs a provenance decision.
This file and fp8_gemm_utils.py carry only:
# SPDX-License-Identifier: Apache-2.0
# Copyright (c) 2025 FlyDSL Project Contributors
Every other new file in this PR carries Copyright (c) 2026, Advanced Micro Devices, Inc.. Two things to sort out:
- Provenance. If these are copied verbatim from the upstream FlyDSL project, that should be stated (a
# Vendored from <repo> @ <commit>line), andqa/L0_license/config.jsonmay need an entry — the repo's convention is that third-party code lives under3rdparty/, not insidetransformer_engine/. If they were modified by AMD, an AMD copyright line should be added above the existing one per the repo's header rules. - Naming. The module is called
fp16_gemm_utils.pybut its docstring says "Byte-staging helpers for the BF16 four-wave GEMM", andbf16_gemm.py:35-42importscompute_global_bf16_transpose_swizzle/make_bf16_byte_buffer_tensorfrom it. A name that reflects what it actually holds (shared 16-bit staging helpers) would be less confusing.
Also worth noting: both this file and fp8_gemm_utils.py define cdiv, ceildiv, and a module-level def divmod(a, b) that shadows the Python builtin — pylint will flag W0622 redefined-builtin in the same CI lint job.
| if qk % 4 != 0: | ||
| raise ValueError( | ||
| f"Scale K/32 dimension={qk} must be divisible by 4" | ||
| ) | ||
| if dim % 64 != 0: | ||
| raise ValueError( | ||
| f"Scale outer dimension={dim} must be a multiple of 64" | ||
| ) |
There was a problem hiding this comment.
These shape rejections run before the FlyDSLUnsupportedError checks, so unsupported MXFP8 shapes crash instead of falling back.
pack_mx32_scales_for_hk is called from mxfp8_matmul at lines 1752, 1765, and 1778 — all before do_gemm at line 1796, where the FlyDSLUnsupportedError divisibility checks live (lines 1558-1578). So any MXFP8 GEMM whose K isn't a multiple of 128 (qk % 4) or whose M/N isn't a multiple of 64 hits this ValueError first and propagates out of general_gemm — the exact opposite of the intended fallback. Same for the k % SCALE_GROUP_SIZE != 0 check at line 1716.
This is the same root cause as the comment on cpp_extensions/gemm.py:512: only FlyDSLUnsupportedError triggers the fallback. Since these are shape-dependent rejections (not programming errors), they should be FlyDSLUnsupportedError:
| if qk % 4 != 0: | |
| raise ValueError( | |
| f"Scale K/32 dimension={qk} must be divisible by 4" | |
| ) | |
| if dim % 64 != 0: | |
| raise ValueError( | |
| f"Scale outer dimension={dim} must be a multiple of 64" | |
| ) | |
| if qk % 4 != 0: | |
| raise FlyDSLUnsupportedError( | |
| f"Scale K/32 dimension={qk} must be divisible by 4" | |
| ) | |
| if dim % 64 != 0: | |
| raise FlyDSLUnsupportedError( | |
| f"Scale outer dimension={dim} must be a multiple of 64" | |
| ) |
More generally, it's worth auditing the whole kernel layer for this. The same file/function is inconsistent with itself elsewhere: bf16_gemm.py:1292 correctly raises FlyDSLUnsupportedError for non-contiguous A/B, but bf16_gemm.py:1328 raises ValueError for non-contiguous C (same in fp16:1328, fp32:1036, fp8:1083). Unsupported output dtypes raise TypeError throughout (bf16/fp16:1318,1377, fp8:1078, mxfp8:1590,1665).
Relatedly, mxfp8_gemm.py:1607 does C_arg = C.contiguous().view(-1) — unlike fp8_gemm.py:1083 there's no contiguity check on the MXFP8 output, so a non-contiguous D would silently receive nothing (.contiguous() copies to a throwaway buffer). Today that's only accidentally blocked by D.view(m, n) raising in the wrapper.
| if bool(transa): | ||
| b_tn = A_flat.transpose(0, 1).contiguous() | ||
| else: | ||
| b_tn = A_flat |
There was a problem hiding this comment.
FP32 TN transposes the weight twice per GEMM, ending back on the original layout.
For the default TE TN case (transa=True), this materialises b_tn = A_flat.transpose(0, 1).contiguous() — one full copy of the weight. Then fp32_matmul immediately undoes it:
# fp32_gemm.py:1038
b_hk = b.transpose(0, 1).contiguous()
doGemm(a, b_hk, c, stream=stream)So every FP32 TN GEMM performs two full transposing copies of A and hands the kernel a buffer with the same layout A_flat already had. The PR notes FP32 perf wasn't a target, but this is pure waste (2 × M*K*4 bytes of extra traffic and allocation per call), and the comment above claims "only the operands whose BLAS transpose flags require it are materialized" — which isn't what happens.
Either pass A_flat straight through and drop the transpose in fp32_matmul, or drop it here. Worth also checking whether the unconditional b.transpose(0, 1).contiguous() in fp32_matmul is correct for the NN/NT paths, since those reach it having not been pre-transposed.
| pid_m, pid_n = _xcd_swizzle(num_blocks_m, num_blocks_n) | ||
| else: | ||
| pid_m, pid_n = divmod(fx.block_idx.x, num_blocks_n) |
There was a problem hiding this comment.
This uses the Python builtin divmod on DSL values, not the module's _divmod helper.
_divmod is defined at line 160 of this same file and is unused. fp8_gemm_utils.divmod's docstring states the reason it exists: "The builtin divmod rejects DSL scalar types." fx.block_idx.x is a DSL scalar.
Same pattern at fp16_gemm.py:474, fp32_gemm.py:277, fp8_gemm.py:312 — all four define _divmod and none of them call it.
This is latent rather than live: it only executes on the use_xcd_remap=False branch, and every caller today takes the use_xcd_remap=True default. But it means the non-swizzled path has almost certainly never been exercised, so it'll fail the moment someone flips that flag.
| pid_m, pid_n = _xcd_swizzle(num_blocks_m, num_blocks_n) | |
| else: | |
| pid_m, pid_n = divmod(fx.block_idx.x, num_blocks_n) | |
| pid_m, pid_n = _xcd_swizzle(num_blocks_m, num_blocks_n) | |
| else: | |
| pid_m, pid_n = _divmod(fx.block_idx.x, num_blocks_n) |
While in here: LDS_SYM_A0/A1/B0/B1, LDS_ALIAS_DOMAIN, and SCOPE_IDS (lines 84-89, and the equivalents at fp16:84-89, fp32:77-82, fp8:79-84) are defined but never referenced anywhere. Likewise fp32_gemm.py:93 make_fp32_inputs is a test-data generator sitting in a shipped kernel module, and fp8_gemm_utils.py ships seven unused symbols (preshuffle_b, compute_global_linear_128x128, G2STransposeLoader, StoreC, wait_barrier, Mfma16x16x128, cdiv/ceildiv/divmod). Worth pruning before merge.
| """Launch the wrapper-selected BF16 TN/NN/NT specialization.""" | ||
| if layout not in ("TN", "NN", "NT"): | ||
| raise ValueError(f"Unsupported FP16 layout: {layout}") | ||
| if a.ndim != 2 or b.ndim != 2: | ||
| raise ValueError( | ||
| f"FlyDSL BF16 expects rank-2 operands, got A{tuple(a.shape)} " |
There was a problem hiding this comment.
Copy-paste leftovers: user-facing FP16 error messages say "BF16".
Lines 1278, 1283, 1293, 1310, and 1374 in this file all say BF16, including line 1374's f"BF16 {layout} requires BF16 inputs, got {A.dtype} and {B.dtype}" — which is what a user actually sees when an FP16 GEMM is rejected. Line 493's comment has the same issue.
These are cosmetic on their own, but they're symptomatic of the wholesale file duplication flagged on gemm_wrappers.py:454 — diff bf16_gemm.py fp16_gemm.py reports only 102 differing lines out of 1430, and these are the strings the rename pass missed.
| # 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"]) |
There was a problem hiding this comment.
NVTE_USE_FLYDSL means two different things, and there's no runtime guard for the mismatch.
Here it is a build-time flag that decides whether flydsl lands in install_requires. In cpp_extensions/gemm.py:497 the same name is a runtime flag that decides whether to dispatch to FlyDSL. Nothing ties them together, so NVTE_USE_FLYDSL=1 at runtime on a wheel built without it makes the lazy import at cpp_extensions/gemm.py:502 raise a bare ImportError out of general_gemm — no fallback, no actionable message.
Either give the build-time flag a distinct name (e.g. NVTE_BUILD_WITH_FLYDSL), or catch ImportError alongside the unsupported-config path so a missing package degrades to the default backend with a clear warning.
Separately, this will fail the black pre-commit hook — the condition fits in 100 columns, so black collapses it:
| # 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"]) | |
| # 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.append("flydsl") |
| use_gemm_flydsl = (IS_HIP_EXTENSION | ||
| and get_device_compute_capability() == (9, 5) | ||
| and bool(int(os.environ.get("NVTE_USE_FLYDSL", "0")))) |
There was a problem hiding this comment.
Formatting: this will fail the black pre-commit hook.
.pre-commit-config.yaml runs black --line-length=100 --preview; continuation lines aligned to the opening paren get reformatted to a 4-space hanging indent:
| use_gemm_flydsl = (IS_HIP_EXTENSION | |
| and get_device_compute_capability() == (9, 5) | |
| and bool(int(os.environ.get("NVTE_USE_FLYDSL", "0")))) | |
| use_gemm_flydsl = ( | |
| IS_HIP_EXTENSION | |
| and get_device_compute_capability() == (9, 5) | |
| and bool(int(os.environ.get("NVTE_USE_FLYDSL", "0"))) | |
| ) |
The new files under flydsl_kernels/ and tests/pytorch/flydsl_kernels/ have the same problem more broadly — they don't look like they've been through qa/format.sh. There are also 9 trailing-whitespace hits that the trailing-whitespace hook will reject:
bf16_gemm.py:59, 1111fp16_gemm.py:59, 1111fp32_gemm.py:52, 842fp8_gemm.py:54, 853mxfp8_gemm.py:1348
Running bash qa/format.sh should clear all of it.
Separately on this gate: NVTE_USE_FLYDSL is re-read from os.environ on every GEMM call. That's a dict lookup plus int()/bool() per GEMM on the hot path. The surrounding code caches this kind of thing (functools is already imported in this module) — worth hoisting into an @functools.lru_cache helper.
| # Copyright (c) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. | ||
| # | ||
| # License for AMD contributions = MIT. See LICENSE for more information |
There was a problem hiding this comment.
Copyright audit: new file should carry a single current year, not a range.
2024-2026 implies the AMD copyright dates back to 2024, but this file is added by this PR. Per the repo's header rules, newly added files take the current year as a single value:
| # Copyright (c) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. | |
| # | |
| # License for AMD contributions = MIT. See LICENSE for more information | |
| # Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. | |
| # | |
| # License for AMD contributions = MIT. See LICENSE for more information |
(The License for AMD contributions = MIT line is fine — it's a valid alternative to See LICENSE for license information. for AMD-only files.)
Two structural notes on this file while I'm here:
- It's never run by CI.
ci/pytorch.shenumerates test files explicitly (lines 93-99 list eachtriton_kernels/test_*.py);flydsl_kernels/test_gemm.pyisn't added, so this 551-line suite won't execute in CI. tests/pytorch/flydsl_kernels/has no__init__.py. PR Experimental Triton GEMM backend for TE PyTorch (BF16/FP16/FP32/FP8/MXFP8) #667 also adds atests/pytorch/triton_kernels/test_gemm.py. With pytest's defaultprependimport mode and no__init__.pyin either directory, two same-named test modules under one rootdir collide withimport file mismatch. Adding__init__.pyto both directories (or switching toimportmode=importlib) avoids it — and is worth coordinating with @ipanfilo's comment about unifying these into one backend-parametrised suite.
The if __name__ == "__main__": smoke-test block at lines 522-551 also duplicates what pytest already does, bypasses the cleanup_env fixture, and leaks NVTE_USE_FLYDSL=1 into the process. Suggest dropping it.
Claude review — FlyDSL GEMM backendReviewed the full diff against the merge base ( Verdict: changes requested. The kernel work looks solid — I checked the MFMA intrinsics, LDS sizing, swizzles, and fp8 cbsz/blgp encoding against each module's dtype and found no divergence bugs despite the heavy file duplication. The problems are concentrated in the integration and test layers. Blocking
Non-blocking but worth addressing: This overlaps with @ipanfilo's point about a shared backend-parametrised test rather than per-backend Copyright headers: 3 files need attention — |
3 participants
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=1is set during installation,setup.pyadds theflydslpackage 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
Validation
Validation was added at both the kernel and model-integration levels.
Added user-facing GEMM tests under
tests/pytorch/flydsl_kernels/covering:torch.matmulreferences and the existing Transformer Engine GEMM backendAdded
test_linear_accuracy_flydsltotests/pytorch/test_numerics.pyto validate the FlyDSL backend through the public Transformer EngineLinearmodule. The test compares the native and FlyDSL execution paths for:dgrad)wgrad)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.
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
DelayedScalingrecipe 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.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.
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:
M=32768, N=16384, K=53248M=32768, N=37888, K=3584M=32768, N=59136, K=8192These cases failed before kernel execution while packing a FlyDSL launch argument into a signed 32-bit integer, producing:
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