[SM120] Support mixed FP8 and FP4 scaled dot by sjoerdmeijer · Pull Request #11154 · triton-lang/triton

sjoerdmeijer · 2026-08-03T10:36:24Z

Enable mixed-precision tt.dot_scaled on SM120 for FP8 and packed FP4 operands in either order.

As described in the PTX doc, the MMA instruction expects 4 bits to be placed in an unpacked byte as 00xxxx00. The new fp4 mode in ldmatrix was introduced to help create such register layouts. With .src_fmt .b4x16_p64 and .dst_fmt .b8x16, it allows copying packed fp4 values in smem into registers in an unpacked form. So a byte xxxxyyyy in smem is copied into two bytes in registers as 0000xxxx, 0000yyyy. Enabling this variant of ldmatrix is the key contribution of this work.

The approach relies on the following:

The src SMEM format .b4x16_p64 is what we call fp4Padded, originally introduced for sm100 mixed precision. We can reuse the existing infrastructure to efficiently copy packed fp4 values in global memory into padded SMEM via TMA.
To model the unpacked layout 0000xxxx in registers, we introduce a new attribute fp4Unpacked to DotOperandEncoding. Similarly to fp4Padded for SMEM, the logical shape is packed, but the physical storage in registers is doubled. Support for this was added to some linear layout utilities. With fp4Padded source and fp4Unpacked destination, we can model the new fp4 variant of ldmatrix exactly.
The MMA instruction expects 4 bits to be unpacked as 00xxxx00 while the result of ldmatrix is 0000xxxx. We resolve this quirk at the lowest layer, MMAV2.cpp, by manual shifting, so that the rest of Triton can work with the MMA op as if its operand layout was fp4Unpacked, for which we have a LinearLayout representation.

In terms of lowering and passes, we roughly go through the following stages:

AccelerateMatmul: introduces the MMA layouts and FP4 shared-memory staging; the fp4 operand is staged through fp4Padded shared memory and loaded into an fp4Unpacked register dot operand.
OptimizeDotOperands: Extend FuseTransMMAV3Plus for sm120 so that OptimizeDescriptorEncoding can select an fp4Padded encoding for tensor descriptors, just like how it works for sm100 mixed-precision today.
LocalLoadOpConversion: if we find fp4Padded source and fp4Unpacked destination, we emit the new fp4 ldmatrix instruction. If we cannot select ldmatrix for some reason, we fall back to scalar loads.
DotOpToLLVM/MMAv2: Take care of the manual shifting, and emit the new mixed-precision mma.sync instruction.

Performance on 8k x 8k x 8k matmul, comparing against two baselines:

Path	TFLOPs	Best configuration
BF16 dequantization (main)	217	128×128×64, 4 warps, 3 stages, no TMA
FP4→FP8 promotion (PR 10746)	488	128×128×64, 4 warps, 2 stages, TMA for A/B only
Native mixed-precision	691	128×128×128, 4 warps, 2 stages, TMA for A/B and scales

For reproducibility, we have used the following autotuning to script to find the optimal configurations:
https://gist.github.com/sjoerdmeijer/62efdf1db84c88ae246c0a92d3726599

New contributor declaration

I am not making a trivial change, such as fixing a typo in a comment.
I have written a PR description following these
rules.
I have run pre-commit run --from-ref origin/main --to-ref HEAD.
Select one of the following.
- I have added tests.
  - /test for lit tests
  - /unittest for C++ tests
  - /python/test for end-to-end tests
- This PR does not need a test because FILL THIS IN.
Select one of the following.
- I have not added any lit tests.
- The lit tests I have added follow these best practices,
  including the "tests should be minimal" section. (Usually running Python code
  and using the instructions it generates is not minimal.)

masahi · 2026-08-03T11:14:37Z

cc @guoriyue @mobicham

masahi · 2026-08-03T11:24:52Z

@ThomasRaoux @lezcano We admit that the non-trivial changes to LL and ldmatrix lowering in this PR might pass the complexity threshold, but I hope that the new fp4Unpacked concept and its application to fp4 ldmatrix are genuinely interesting. Other diffs are fairly local / test only.

Enable mixed-precision tt.dot_scaled on SM120 for FP8 and packed FP4 operands in either order. As described in [the PTX doc](https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#warp-level-matrix-instructions-mma), the MMA instruction expects 4 bits to be placed in an unpacked byte as 00xxxx00. The new fp4 mode in [ldmatrix](https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#warp-level-matrix-instructions-ldmatrix) was introduced to help create such register layouts. With `.src_fmt` .b4x16_p64 and `.dst_fmt` .b8x16, it allows copying packed fp4 values in smem into registers in an unpacked form. So a byte xxxxyyyy in smem is copied into two bytes in registers as 0000xxxx, 0000yyyy. Enabling this variant of ldmatrix is the key contribution of this work. The approach relies on the following: - The src SMEM format .b4x16_p64 is what we call `fp4Padded`, originally introduced for sm100 mixed precision. We can reuse the existing infrastructure to efficiently copy packed fp4 values in global memory into padded SMEM via TMA. - To model the unpacked layout 0000xxxx in registers, we introduce a new attribute `fp4Unpacked` to `DotOperandEncoding`. Similarly to `fp4Padded` for SMEM, the logical shape is packed, but the physical storage in registers is doubled. Support for this was added to some linear layout utilities. With `fp4Padded` source and `fp4Unpacked` destination, we can model the new fp4 variant of ldmatrix exactly. - The MMA instruction expects 4 bits to be unpacked as 00xxxx00 while the result of ldmatrix is 0000xxxx. We resolve this quirk at the lowest layer, `MMAV2.cpp`, by manual shifting, so that the rest of Triton can work with the MMA op as if its operand layout was `fp4Unpacked`, for which we have a LinearLayout representation. In terms of lowering and passes, we roughly go through the following stages: - `AccelerateMatmul`: introduces the MMA layouts and FP4 shared-memory staging; the fp4 operand is staged through `fp4Padded` shared memory and loaded into an `fp4Unpacked` register dot operand. - `OptimizeDotOperands`: Extend `FuseTransMMAV3Plus` for sm120 so that `OptimizeDescriptorEncoding` can select an `fp4Padded` encoding for tensor descriptors, just like how it works for sm100 mixed-precision today. - `LocalLoadOpConversion`: if we find `fp4Padded` source and `fp4Unpacked` destination, we emit the new fp4 ldmatrix instruction. If we cannot select ldmatrix for some reason, we fall back to scalar loads. - `DotOpToLLVM/MMAv2`: Take care of the manual shifting, and emit the new mixed-precision mma.sync instruction. Performance on 8k x 8k x 8k matmul, comparing against two baselines: | Path | TFLOPs | Best configuration | |---|---:|---| | BF16 dequantization (main) | **217** | 128×128×64, 4 warps, 3 stages, no TMA | | FP4→FP8 promotion ([PR 10746](triton-lang#10746)) | **488** | 128×128×64, 4 warps, 2 stages, TMA for A/B only | | Native mixed-precision | **691** | 128×128×128, 4 warps, 2 stages, TMA for A/B and scales | For reproducibility, we have used the following autotuning to script to find the optimal configurations: https://gist.github.com/sjoerdmeijer/62efdf1db84c88ae246c0a92d3726599

sjoerdmeijer · 2026-08-06T10:50:34Z

Rebased

lezcano

Why do you need toRegisterElementLinearLayout? Seems quite hacky.

Also stagePackedFp4OperandThroughSmem seems a bit iffy. I would like @ThomasRaoux to have a look at it as well.

lezcano · 2026-08-06T12:23:18Z

+    layout = nvidiaDotToLinearLayout(shape, *this);
  }
+
+  if (getFp4Unpacked()) {


but this is just used in nvidiaDotToLinearLayout, why modify everything?

It is currently only used for nvidiaDot, but we think the concept of fp4Unpacked is generic. So we chose not to over-specialize for the NV path.

lezcano · 2026-08-06T12:47:32Z

+// Lower a local_load whose result uses the fp4Unpacked mxfp4 dot encoding. The
+// source must use the fp4Padded NVMMA shared-memory encoding. First try the fp4
+// ldmatrix instruction; if the source and register layouts are not compatible
+// with that instruction, fall back to scalar loads.
+static LogicalResult lowerFp4UnpackedLocalLoad(
+    triton::gpu::LocalLoadOp op, triton::gpu::LocalLoadOp::Adaptor adaptor,
+    const LLVMTypeConverter *typeConverter, ConversionPatternRewriter &rewriter,
+    const NVIDIA::TargetInfo &targetInfo) {


Can't we do this in the generic lowering where we try to use ldmatrix and otherwise we fallback to ld.shared?

lezcano · 2026-08-06T12:48:44Z

+    // fp4Unpacked mxfp4 operand: expand the packed bytes into one 8-bit field
+    // per e2m1 value as part of the load.
+    if (auto dotEnc = dyn_cast<DotOperandEncodingAttr>(dstTy.getEncoding()))
+      if (dotEnc.getFp4Unpacked())
+        return lowerFp4UnpackedLocalLoad(op, adaptor, getTypeConverter(),
+                                         rewriter, targetInfo);


this feels a bit too specific. Fine if we carry the dotEnc.getFp4Unpacked() through the main path, but we shouldn't write a full different path that will be barely tested as that's too risky.

The trade off is that the main path will get sm120-specific complexities that you don't want to deal with otherwise. If you are fine with that, we should be able to merge the paths @sjoerdmeijer

sjoerdmeijer · 2026-08-06T13:53:31Z

I will try to answer the high level question first before addressing the inlined ones, and see if you're happy with this.

Why do you need toRegisterElementLinearLayout?

Each logical i8 tensor element initially contains two FP4 values, in the low and high nibble of a byte. fp4Unpacked is the register representation where the low/high nibble are placed in separate bytes and needs two layout views:

toLinearLayout maps both unpacked nibble registers to the same coordinate,
toRegisterElementLinearLayout provides the register view: it maps the two registers to virtual consecutive K coordinates. This prevents LLVM struct sizing and pack/unpack logic from discarding the high nibble as a broadcast.

Also stagePackedFp4OperandThroughSmem seems a bit iffy. I would like @ThomasRaoux to have a look at it as well.

The short summary and rationale for this function is that we lower the fp4 unpack sequence into a hardware matrix load, and as a bonus we expose it to sw pipelining and TMA. Or in other words, the staging doesn't have pipelining as a goal, but creates this abstraction or transition of fp4Padded shared mem to a local_load with fp4Unpacked=true.
We also keep things general. I.e., at this point we have defined that the fp4 values live in bits [3:0]. But due to the architecture quirk, we need to shift this to [5:2], which we don't do here, but later, and we made this specific for the MMAv2.

ThomasRaoux · 2026-08-06T15:53:36Z

+// This helper creates a `local_alloc` followed by a `local_load` for a packed
+// mxfp4 input. Thus, it stages a packed input through shared memory, and then
+// creates a fp4Unpacked dot operand. This is an optimisation for SM120.
+static Value stagePackedFp4OperandThroughSmem(Value v, int opIdx,
+                                              RankedTensorType newRetType,
+                                              PatternRewriter &rewriter) {
+  // Keep the packed (K) axis contiguous in SMEM for the dot operand.
+  Value smem = getSharedMemoryMMAOperand(v, rewriter, opIdx,
+                                         /*allowTranspose=*/false,
+                                         /*fp4Padded=*/true);
+  auto vType = cast<RankedTensorType>(v.getType());
+  auto dotEnc = DotOperandEncodingAttr::get(v.getContext(), opIdx,
+                                            newRetType.getEncoding(),
+                                            /*kWidth=*/2u,
+                                            /*fp4Unpacked=*/true);
+  auto regType = vType.cloneWithEncoding(dotEnc);
+  return LocalLoadOp::create(rewriter, v.getLoc(), regType, smem);
+}


Why are we not just doing a convert layout like we do for other cases where the input is in register?

Does my message in #11154 (comment) answer this?

ThomasRaoux · 2026-08-06T15:56:47Z

+// where user is:
+//  - WGMMA/MMAv5, or
+//  - `ttg.local_load` if the allocation uses an `fp4Padded` shared memory
+//    encoding.


I think mixing the two is going to be messy, I don't get why this doesn't work the same way other mmav2 cases do

This is for propagating fp4Padded = true alloc from the MMA operand across transpose, so that OptimizeDescriptorEncoding can select fp4Padded desc encoding. This is exactly what sm100 mixed precision lowering depends on, and we are just extending this flow for local_load user of fp4Padded alloc.

ThomasRaoux · 2026-08-06T15:58:02Z

+fp4Unpacked specifies that each byte containing two packed fp4 elements is
+unpacked into two bytes, with each fp4 value occupying bits [3:0] of its byte.


the data type is not fp4 though? so what does that mean in practice for the layout?

It is conceptually similar to SMEM fp4Padded in that the logical tensor is packed in an i8 container but the physical storage is doubled. In this case a byte yyyyxxxx in a logical shape maps to two i8 registers, 0000yyyy and 0000xxxx.

ThomasRaoux · 2026-08-06T16:00:49Z

+            if not is_mixed:
+                pytest.skip("Mixed fp8 x fp4 operands expected here")
+            if not PACK_B_ALONG_K:
+                pytest.skip("Mixed fp8 x fp4 requires K-packed fp4")


why the non-mixed are not being tested? so we don't support the non K-packed at all?

The ISA requires K packing. We can run this test with MN-pack but we'd be testing the dequantized path. @sjoerdmeijer Let's test that anyway here.

ThomasRaoux · 2026-08-06T16:03:04Z

+    parser.add_argument("--block-m", type=int, help="override the M tile size")
+    parser.add_argument("--block-n", type=int, help="override the N tile size")
+    parser.add_argument("--block-k", type=int, help="override the K tile size")
+    parser.add_argument("--num-warps", type=int, help="override the number of warps")
+    parser.add_argument("--num-stages", type=int, help="override the number of pipeline stages")


can we skip the complexity, tutorials are meant to stay as simple as possible

ThomasRaoux · 2026-08-06T16:04:29Z

@@ -1,4 +1,4 @@
-// RUN: triton-opt --split-input-file %s --verify-diagnostics
+// RUN: triton-opt %s --split-input-file --allocate-shared-memory-nv='compute-capability=120' --convert-triton-gpu-to-llvm='compute-capability=120' --verify-diagnostics


why would we need to run those passes here, this seems off. This is meant to check the verifier based on some ops

ThomasRaoux · 2026-08-06T16:06:09Z

  bool supportLdStMatrixB8() const override {
    return targetFeatures.supportLdStMatrixB8();
  }
+  bool supportsFp4Ldmatrix() const {


maybe the ldmatrix changes can be split in a separate PR?

ThomasRaoux · 2026-08-06T16:09:59Z

The short summary and rationale for this function is that we lower the fp4 unpack sequence into a hardware matrix load, and as a bonus we expose it to sw pipelining and TMA. Or in other words, the staging doesn't have pipelining as a goal, but creates this abstraction or transition of fp4Padded shared mem to a local_load with fp4Unpacked=true.
We also keep things general. I.e., at this point we have defined that the fp4 values live in bits [3:0]. But due to the architecture quirk, we need to shift this to [5:2], which we don't do here, but later, and we made this specific for the MMAv2.

I think this is not the right way to look at it. We shouldn't be writing the code based on how we expect the code to look like for some specific kernels but we should be thinking about the right way to layer the transformations in the compiler without assuming what the kernel would do. For instance there may be cases where the input of the mma is already in register and we don't want it to go through smem (think A100 flash attention cases). I think making this case go through a different flow than other mmav2 cases is not a good direction.

masahi · 2026-08-06T19:08:20Z

To answer why we are allocating SMEM for fp4 early and why the lowering cannot work like "other mmav2 cases":

Our perf results in the PR description demonstrates why using ldmatrix with the new fp4 mode is important for sm120 mixed-prec mma. We started this work based on an assumption: If we get sm120 mixed-precision dot_scaled workload, we always want to use ldmatrix.

The primary complexity of this ldmatrix mode is that the input SMEM needs to be fp4Padded. If we don't go through SMEM in AccelerateMatmul, SWP will allocate SMEM but for packed operands. TMA layout selection also chooses a packed descriptor encoding. We need to unpack those operands at some point, but

Not being able to use the specialized ldmatrix is inefficient
Introducing fp4Padded after the fact to replace the existing packed encodings would be extremely messy - SWP, convert_layout, and OptimizeDescriptorEncoding etc need to be aware of the sm120 mixed-precision requirement

So, in terms of "the right way to layer the transformations in the compiler", we believe our solution is the simplest and the least invasive one - AccelerateMatmul knows exactly what kind of layout the fp4 operand should have, and SWP or OptimizeDescriptorEncoding can directly work on the padded layout without SM120-specific knowledge.

Granted, our solution targets the most common use of SM120 mixed-precision dot where the input comes straight from global memory. It remains functionally correct when the input is already in registers. We are not aware of motivating use cases where a packed fp4 operand in mixed-precision MMA is produced in registers, but we can potentially consider an alternative path for it, for example, the fp4 -> fp8 promotion solution from #10746. Since that is more of an efficiency problem, we want to defer that and keep our changes to the compiler minimal for this first work.

sjoerdmeijer requested review from lezcano, peterbell10 and ptillet as code owners August 3, 2026 10:36

sjoerdmeijer force-pushed the sm120-mixed-fp4fp8-mma branch from 2461fc7 to f1de86f Compare August 3, 2026 11:00

sjoerdmeijer force-pushed the sm120-mixed-fp4fp8-mma branch from f1de86f to e1d5676 Compare August 6, 2026 10:49

lezcano reviewed Aug 6, 2026

View reviewed changes

ThomasRaoux reviewed Aug 6, 2026

View reviewed changes

		fp4Unpacked specifies that each byte containing two packed fp4 elements is
		unpacked into two bytes, with each fp4 value occupying bits [3:0] of its byte.

		@@ -1,4 +1,4 @@
		// RUN: triton-opt --split-input-file %s --verify-diagnostics
		// RUN: triton-opt %s --split-input-file --allocate-shared-memory-nv='compute-capability=120' --convert-triton-gpu-to-llvm='compute-capability=120' --verify-diagnostics

Uh oh!

Conversation

sjoerdmeijer commented Aug 3, 2026 • edited by masahi Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

New contributor declaration

Uh oh!

masahi commented Aug 3, 2026

Uh oh!

masahi commented Aug 3, 2026

Uh oh!

sjoerdmeijer commented Aug 6, 2026

Uh oh!

lezcano left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

sjoerdmeijer commented Aug 6, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

ThomasRaoux commented Aug 6, 2026

Uh oh!

masahi commented Aug 6, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

sjoerdmeijer commented Aug 3, 2026 •

edited by masahi

Loading