Triton #11360 — Why a 128 MiB MXFP4 Layout Conversion Peaked at Nearly 3 GiB

Fused Hopper bit conversion, direct Blackwell tile shuffling, and the difference between faster weight preparation and faster matmul

The packed weights occupied 128 MiB.

Converting their layout temporarily allocated almost 3 GiB.

The largest recorded case was the Hopper MXFP4 unswizzle path:

Existing packed input:
128 MiB

Peak allocation during conversion:
2,944 MiB

Warm conversion time:
13.424 ms

After Triton PR #11360:

Peak allocation:
256 MiB

Warm conversion time:
3.170 ms

That was:

91.3% less peak allocated memory

and

4.24× faster warm conversion

On GB300, the measured Blackwell conversions were smaller in absolute memory use but faster by more than twenty times after warmup.

Blackwell swizzle:
1.887 ms
→
0.0887 ms
Blackwell unswizzle:
1.882 ms
→
0.0891 ms

The patch did not make the matrix multiplication itself twenty times faster.

It changed the preparation step that converts packed MXFP4 values into the physical layouts expected by Hopper and Blackwell kernels.

Packed checkpoint or quantized weight

        ↓

layout conversion

        ↓

hardware-oriented value storage

        ↓

later matmul

The old paths expressed bit packing, repacking, padding, transposition, and tile shuffling as sequences of PyTorch tensor operations.

Each operation was individually valid.

Several of them materialized another full-size tensor.

The new paths moved those transformations into Triton kernels that calculate the destination representation directly.

Old

source
→ full intermediate A
→ full intermediate B
→ padded intermediate
→ contiguous transpose
→ final destination
New

source bytes
→ register-level bit or index transformation
→ final destination

The patch’s reported measurements used a logical shape of:

[16, 4096, 4096]

containing 128 MiB of packed MXFP4 value data. Peak allocation included the conversion output and excluded the already-existing input; it was allocated tensor memory, not total reserved VRAM.

A four-bit format saves storage only when preparing that format does not repeatedly expand the complete weight tensor into wider temporary representations.


What was being converted

FP4 stores one value in four bits.

Two FP4 values can therefore be packed into one byte.

A logical matrix containing:

K × N FP4 values

may use canonical packed storage shaped approximately like:

[K, N / 2] bytes

But canonical packed order is not necessarily the physical order a particular GPU kernel wants.

A hardware-oriented kernel may need the values:

  • grouped by threads,

  • interleaved for a dequantization instruction sequence,

  • packed along another matrix dimension,

  • padded to tile boundaries,

  • or arranged into TMA-compatible tiles.

The logical values can remain unchanged while their byte-level representation changes.

Same FP4 values

+

different packed ordering

→ different physical layout

The patch addressed two such layout families.

Hopper

Hopper’s MXFP4 value layout interleaves bits so the later dequantization path can recover values through a compact instruction sequence.

Blackwell

Blackwell’s shuffled MXFP4 layout arranges packed values into a five-dimensional tiled storage format matching the TMA block shape expected by its matmul path.

The conversion is therefore not numerical quantization.

It is predominantly:

bit placement

+

byte placement

+

tile placement

Why the conversion existed outside the matmul

A matmul kernel benefits from receiving values in the representation it is designed to load efficiently.

It should not necessarily perform a complete checkpoint-layout transformation every time it multiplies the weights.

A separate preparation step can convert the values once:

canonical or checkpoint representation

        ↓

Hopper or Blackwell execution representation

The converted weights can then be reused by later operations.

This separates two costs.

Preparation cost

→ convert packed weights into execution layout
Matmul cost

→ consume the prepared representation

PR #11360 optimized the first cost.

It did not modify the mathematical GEMM computation or publish a matmul-latency improvement.

That scope is explicit in the PR:

This improves weight preparation, not the matmul itself.


Hopper: the old path widened the complete tensor repeatedly

The old Hopper conversion contained Torch implementations named:

_pack_bits()

_unpack_bits()

These functions used ordinary tensor-wide operations to transform packed bytes.

The operations included:

cast to int32

left shift

right shift

bitwise AND

bitwise OR

stack

flatten

cast to int16

cast back to uint8

For a small tensor, this is straightforward and readable.

For a 128 MiB packed weight tensor, each tensor-wide intermediate is large.


The old packing path promoted byte lanes to int32

A simplified form of _pack_bits() was:

x = x.contiguous()

x = x.reshape(
    ...,
    x.shape[-1] // 4,
    4,
)

result =
    compress(x[..., 0])

result |=
    shifted_compress(x[..., 1])

result |=
    shifted_compress(x[..., 2])

result |=
    special_compress(x[..., 3])

The helper functions converted their inputs to:

torch.int32

and then performed several masks and shifts.

A byte-oriented input was therefore temporarily represented through much wider values.

uint8 source lane

→ int32 intermediate lane

The conversion was not wrong.

But each operation could create another tensor proportional to the full weight size.


The old unpacking path created even more live state

The inverse path was more expensive.

It began by viewing packed bytes as int32, then constructed several full-width expressions:

a = (x << 1) & mask_a

b = right_shift(x, 3) & mask_b

c = right_shift(x, 7) & mask_c

It then built four separate decoded components:

unpacked = [
    x & mask,
    (x << 3) & mask,
    (x << 6) & mask,
    (a | b) | c,
]

and materialized them through:

torch.stack(
    unpacked,
    dim=-1,
)

The stacked tensor was subsequently flattened and converted again through int16 and uint8 intermediates.

The old source makes this expansion directly visible.


A view is cheap; the operations around it were not

This line:

x = x.view(torch.int32)

does not necessarily allocate another complete buffer.

It changes how the existing bytes are interpreted.

But operations such as:

x << 3

x & mask

torch.stack(...)

.to(torch.int16)

.to(torch.uint8)

produce new tensors.

The distinction is:

view

→ metadata-only interpretation
shift / mask / cast / stack

→ materialized result

The old code mixed inexpensive views with full-tensor materializations.

Looking only at the original and final tensor sizes therefore understated the conversion’s actual live memory.


Why the peak could exceed the input by more than twenty times

A 128 MiB packed tensor does not mean every intermediate is also exactly 128 MiB.

Some operations temporarily use wider element types.

Some stack several decoded forms.

Some source tensors must remain alive while the destination is being produced.

Some subsequent operations begin before every prior temporary has been released by the allocator.

Conceptually:

128 MiB source

+

several int32-scale intermediates

+

four-way stack

+

narrowing outputs

+

final layout buffer

→ multi-GiB peak

The exact live-buffer schedule is determined by PyTorch operations and allocator lifetimes.

The measured result was the relevant evidence:

H100 direct unswizzle

2,944 MiB peak

above the pre-existing input allocation.

The article does not need to claim that one named temporary alone occupied the complete 2,944 MiB.

The old operation graph created enough simultaneous full-tensor state for the measured peak to reach that value.


The new Hopper kernel keeps the expansion in registers

PR #11360 introduced:

_convert_bits_kernel

for CUDA uint8 tensors.

Instead of constructing several tensor-wide integer expressions, one Triton program instance processes a local group of bytes.

Conceptually:

load four source bytes

        ↓

promote four scalar or lane values in registers

        ↓

perform masks and shifts

        ↓

store one transformed 32-bit group

The temporary bit patterns exist as Triton values local to the executing program.

They are not global tensors allocated through PyTorch.


Four byte loads replaced several complete tensors

The kernel calculates one output group from:

a = source[4 × offset + 0]

b = source[4 × offset + 1]

c = source[4 × offset + 2]

d = source[4 × offset + 3]

For the forward conversion, those four bytes are compressed and interleaved into one result.

For the inverse conversion, they are assembled into a 32-bit value, decoded through masks and shifts, and returned to packed FP4 bytes.

The same kernel uses a compile-time flag:

INVERSE

to select the direction.

INVERSE = false

→ pack / interleave
INVERSE = true

→ unpack / restore

The merged implementation and its current main form show the four byte loads, register-level bit operations, one output allocation, and INVERSE specialization.


The algorithm still performs the same bit transformation

The patch did not replace the Hopper bit contract with an approximate or numerically different layout.

The Torch implementation remains available as a reference and fallback.

The CUDA test compares:

Triton-generated bytes

against

Torch-generated bytes

for combinations including:

  • several ranks,

  • transposed and non-transposed inputs,

  • contiguous and strided views,

  • two MX axes,

  • and Hopper MMA versions 2 and 3.

It also converts the result back and requires equality with the original packed data.

The optimization therefore targets the execution of the transformation, not its specified output.


The retained allocation is bounded differently for Hopper

The Hopper regression does not require literally one destination allocation in every case.

The conversion can still need:

  • a contiguous byte representation,

  • the final transformed buffer,

  • and small launch-related overhead.

Its peak-allocation assertion allows:

up to two overlapping output-sized byte buffers

+

1 MiB tolerance

but rejects the former full-tensor int32 intermediates.

peak
<=
2 × swizzled.nbytes
+
1 MiB

The test comment states the intended boundary:

Allow overlapping byte buffers,
but not whole-tensor int32 intermediates.

This is more precise than claiming that the conversion became allocation-free.

It still allocates the representation it needs to return.

It stops allocating the large widened representations that do not survive the function boundary.


Blackwell: the old path materialized every layout stage

The old Blackwell shuffled conversion began from canonical storage:

[..., K, N_packed]

where the final dimension packs two FP4 values into each byte.

Blackwell’s physical representation instead packs along K and arranges values into tiled storage:

[
  E,
  num_tiles_k,
  num_tiles_n,
  tile_n,
  tile_k_packed
]

The old forward path performed several explicit stages.

canonical N-packed bytes

        ↓

repack to K-packed physical matrix

        ↓

pad K and N to tile boundaries

        ↓

transpose

        ↓

make transpose contiguous

        ↓

reshape into tiles

        ↓

permute tile axes

        ↓

make final permutation contiguous

Each materializing step could require another tensor comparable to the complete weight representation.

The old source shows the separate repack, padded tensor, contiguous transpose, and final contiguous permutation.


The old forward conversion first changed the packing dimension

Canonical storage packed adjacent N values into bytes.

The Blackwell physical path needed the corresponding packed structure along K.

The old code allocated a full output tensor and called:

repack(
    source,
    from axis -1,
    to axis -2
)

This produced:

[E, K_packed, N]

as a complete intermediate.

That representation was useful for expressing the next transformations.

It was not the final five-dimensional shuffled layout.


Padding created another complete representation

The TMA-oriented layout required K and N extents rounded to tile-compatible sizes.

When the actual dimensions did not fill the final tile, the old path allocated:

[E, padded_K_packed, padded_N]

initialized it to zero, and copied the unpadded physical matrix into its leading region.

physical intermediate

→ padded physical intermediate

The padding was semantically required.

Materializing a complete standalone padded matrix was an implementation choice.


Transpose and final permutation each requested contiguity

The code then converted:

[E, K_packed, N]

to:

[E, N, K_packed]

through:

data.transpose(1, 2).contiguous()

The transpose itself can be a view.

Calling:

contiguous()

materializes the reordered bytes.

After reshaping into tiles, the code performed another permutation:

[E, num_tiles_n, tile_n, num_tiles_k, tile_k_packed]

        ↓

[E, num_tiles_k, num_tiles_n, tile_n, tile_k_packed]

and again called:

contiguous()

The final call was necessary to produce the required physical storage.

The earlier full physical and transposed representations were temporary steps toward it.


The inverse path repeated the sequence backward

Blackwell unswizzle reversed the transformations.

shuffled 5D storage

        ↓

inverse tile permutation

        ↓

contiguous buffer

        ↓

view as padded transposed matrix

        ↓

transpose to K-packed physical form

        ↓

contiguous buffer

        ↓

trim padding

        ↓

contiguous buffer

        ↓

repack from K-packed to canonical N-packed

Again, every individual operation expressed a clear part of the layout mapping.

Together, they created several complete weight-sized buffers.


The new Blackwell kernel computes final addresses directly

PR #11360 replaced the CUDA uint8 path with:

_convert_shuffled_mxfp4

The new kernel receives:

canonical pointer

shuffled pointer

canonical shape and strides

shuffled shape and strides

tile dimensions

conversion direction

For each block of logical positions, it calculates both address systems.

Where is this FP4 pair
in canonical packed storage?
Where must that pair appear
in shuffled tiled storage?

It then loads the source bytes, rearranges their nibbles in registers, and stores directly to the destination address.

The intermediate physical matrix no longer exists as a tensor.

Neither does the full padded matrix or contiguous transposed matrix.

The current source retains this direct mapping.


A two-by-two nibble transpose changes the packing axis

The kernel’s core bit operation can be understood using four FP4 nibbles.

Suppose two canonical bytes contain:

byte a:
[A0, A1]

byte b:
[B0, B1]

where each bracket contains two four-bit values.

The shuffled representation needs:

output low byte:
[A0, B0]
output high byte:
[A1, B1]

The code constructs:

lo =
low_nibble(a)
|
low_nibble(b) shifted into the high half

and:

hi =
high_nibble(a)
|
high_nibble(b)

in their destination positions.

Conceptually:

[A0, A1]      [A0, B0]
          → 
[B0, B1]      [A1, B1]

This changes which dimension is packed into each byte.

The operation is its own inverse.

Applying the same two-by-two transpose again restores:

[A0, A1]

[B0, B1]

The source comment explicitly records:

Transposing a 2x2 group of FP4 nibbles
is its own inverse.

Padding became masked destination writes

The Blackwell destination still needs padded tile extents.

The patch did not remove that layout requirement.

It changed how the padding is produced.

Old:

allocate complete padded matrix

copy source into it

then shuffle it

New:

calculate padded destination coordinates

write transformed source where source coordinates are valid

write or preserve zero-valued padded regions through masked conversion logic

The destination exists because it is the final storage.

A separate padded intermediate does not.


Noncontiguous leading dimensions were preserved

The new kernel does not assume that all leading input dimensions can be flattened through one contiguous view.

It receives canonical strides and reconstructs the leading batch offset axis by axis.

Conceptually:

flattened batch index

        ↓

recover coordinate for each leading dimension

        ↓

multiply by that dimension’s actual stride

        ↓

sum canonical source offset

This allows the CUDA path to match the Torch reference for strided and sliced inputs.

The tests deliberately construct noncontiguous views by taking every second element and compare the resulting bytes with the CPU Torch implementation.


The Blackwell allocation test requires only the destination

For the direct Blackwell swizzle and unswizzle methods, the peak test uses a stronger bound than the Hopper test.

peak
<=
actual.nbytes
+
1 MiB

Its comment states:

Shuffling should allocate its output,
not another whole weight tensor.

That assertion protects the central optimization.

A future implementation may use small overhead allocations.

It may not silently restore another complete canonical or physical intermediate.


The measured direct-conversion results

The author compared the upstream implementation with the patch for 128 MiB of packed MXFP4 values.

Direct layout transformation

GPU and conversionPeak beforePeak afterReductionWarm time beforeWarm time afterSpeedup
H100 Hopper swizzle896 MiB256 MiB71.4%9.459 ms2.159 ms4.38×
H100 Hopper unswizzle2,944 MiB256 MiB91.3%13.424 ms3.170 ms4.24×
GB300 Blackwell swizzle256 MiB128 MiB50.0%1.887 ms0.0887 ms21.26×
GB300 Blackwell unswizzle256 MiB128 MiB50.0%1.882 ms0.0891 ms21.11×

These figures belong to the documented logical shape, hardware, layouts, software versions, and synchronized warm-call methodology.

They should not be read as:

every MXFP4 conversion on every GPU
becomes twenty-one times faster

The largest speedups were the reported GB300 cases.

The H100 cases were approximately four times faster.


Direct transformation and public convert_layout() were not identical measurements

The PR reports a second table for the public:

convert_layout()

API.

That API can materialize canonical storage as part of the broader transition between tensor layout objects.

The observed peak can therefore include work outside the direct:

transformation.swizzle_data()

or

transformation.unswizzle_data()

method.

Public API results

GPU and conversionPeak beforePeak afterReductionWarm time beforeWarm time afterSpeedup
H100 Hopper swizzle1,024 MiB384 MiB62.5%9.565 ms2.247 ms4.26×
H100 Hopper unswizzle2,944 MiB256 MiB91.3%13.499 ms3.263 ms4.14×
GB300 Blackwell swizzle384 MiB256 MiB33.3%1.925 ms0.1223 ms15.74×
GB300 Blackwell unswizzle256 MiB256 MiB0%1.928 ms0.1412 ms13.65×

The Blackwell public unswizzle remained at:

256 MiB

before and after.

That does not mean the new kernel failed to remove its internal intermediate.

The public API still had to materialize the canonical destination required by its contract.

Its measured warm time nevertheless dropped from:

1.928 ms

to

0.1412 ms

The PR explicitly separates this case from the direct conversion measurements.


Peak memory and speed did not move in identical proportions

Removing a complete intermediate can lower both:

  • allocator pressure,

  • and memory traffic.

But API-required output storage remains.

This produces cases such as:

Blackwell public unswizzle

peak reduction:
0%

warm speedup:
13.65×

The final canonical output still had to exist.

The new kernel reached it without executing the previous chain of repack, transpose, contiguous-copy, and related operations.

Same required output bytes

+

less transformation work

→ large time reduction without a lower measured API peak

This is why performance results should distinguish:

required output memory

from

temporary conversion memory

Fewer tensor operations also meant fewer launches

The old implementations expressed transformations through several PyTorch calls.

On CUDA, those calls can correspond to multiple kernels and allocator operations.

The new paths concentrate much of the work into one specialized Triton launch.

Old

cast launch

shift launch

mask launch

stack launch

repack launch

transpose copy

...
New

one direct conversion kernel

The source structure is consistent with both lower memory traffic and lower launch overhead.

The PR does not publish a component-by-component attribution showing exactly how much of each speedup came from:

  • reduced bytes moved,

  • fewer allocations,

  • fewer launches,

  • or more efficient bit instructions.

The directly supported conclusion is narrower:

the fused implementations produced
the measured combined improvement

under the reported setup.


JIT compilation created a cold-start trade-off

The old Torch path used operations whose required GPU kernels and runtime support were already available through PyTorch.

The new converter is itself a Triton kernel.

Before its first execution, Triton may need to:

  • initialize runtime state,

  • specialize the kernel,

  • compile it,

  • load the result,

  • and populate its cache.

That one-time work can dominate a conversion lasting less than a millisecond after warmup.


Empty cache: the new path was slower on its first call

The PR reports direct first-call medians for several Triton states.

Fresh process and empty Triton cache

H100

swizzle:
880.9 ms

unswizzle:
900.5 ms
GB300

swizzle:
1,079.7 ms

unswizzle:
1,082.2 ms

The upstream path remained approximately:

43–45 ms on H100

65–67 ms on GB300

For a one-off conversion in a fresh process with no cached converter:

new fused kernel
→ lower memory

but

→ much higher first-call latency

This is a real trade-off, not an omitted benchmark detail.


Runtime initialization and kernel cache were separate costs

The PR measured four states.

Triton stateH100 swizzle / unswizzleGB300 swizzle / unswizzle
Fresh process, empty cache880.9 / 900.5 ms1,079.7 / 1,082.2 ms
Prior Triton call, converter uncached124.0 / 141.4 ms250.8 / 253.1 ms
Fresh process, converter cached397.2 / 399.0 ms421.9 / 420.1 ms
Runtime initialized, converter cached35.8 / 36.0 ms4.68 / 4.64 ms

These states show two distinct forms of startup work.

Triton runtime initialization

and

converter-specific compilation/cache state

A disk-cached converter does not eliminate all fresh-process initialization.

A previously executed Triton operation does not mean this particular converter is already compiled.


Warm execution and first-call execution answer different questions

The warm conversion table asks:

Once the kernel exists,
how efficiently does it rearrange the weights?

The first-call table asks:

How long until the first converted result is available
under this runtime and cache state?

Both matter.

A long-lived process that converts many weights may amortize the compilation cost.

A command-line tool converting one already-quantized checkpoint and exiting may not.

The patch explicitly states:

The one-off uncached 128 MiB conversion is still slower.

It also avoids claiming that complete model loading became faster.


Existing Triton work can partially prewarm the path

The PR notes that some quantize-then-pack workflows already execute Triton before the layout conversion.

In those flows:

earlier Triton work

→ runtime initialized

→ part of the cold cost already paid

Already-quantized checkpoints may reach layout conversion without any prior Triton execution.

Their first-use behavior can differ.

This means “model loading” is not one uniform benchmark category.

Quantize during load

and

load already-packed values

can enter different runtime states before conversion starts.


Memory savings remained even when latency was cold

JIT compilation takes time.

It does not require the old full-size tensor intermediates to return.

The PR reports that the memory reductions remained across the measured startup states.

That can matter even when one-off latency is worse.

A process may be able to complete a conversion under its available memory only because the peak fell.

Cold latency:
higher
Peak temporary memory:
still lower

Time and capacity are independent resource boundaries.


The kernel had to follow the input device

Adding a Triton launch introduced a device-ownership obligation that the old Torch operations largely handled automatically.

Consider:

current CUDA device:
cuda:0

while the input tensor resides on:

cuda:1

The converter must launch for the input tensor’s device.

It must not silently use whatever device happens to be current on the host thread.

The new implementations wrap their launch in:

with torch.cuda.device(data.device):
    ...

This binds the Triton launch to the source tensor’s CUDA device.


The active stream on that device also had to be preserved

The regression creates a non-default stream on device 1.

It then enters a context where:

active stream for cuda:1
→ custom stream

while:

torch.cuda.current_device()
→ cuda:0

The input is placed on:

cuda:1

The conversion must satisfy all three facts.

Do not change the caller’s current device permanently.

Launch on the input tensor’s device.

Use the active stream belonging to that input device.

The test requires:

current device remains 0

current stream on device 1 remains the selected stream

output device is cuda:1

output bytes match the reference

It covers Hopper and Blackwell, forward and inverse conversion.


Correct arithmetic on the wrong device is still incorrect execution

A conversion kernel can have perfect index and bit logic and still fail if launched against the wrong device context.

Possible outcomes include:

  • invalid pointer use,

  • device mismatch errors,

  • work issued to an unintended stream,

  • or reading the result before the intended stream has completed.

The patch therefore treats launch ownership as part of layout correctness.

Value mapping

+

device mapping

+

stream mapping

→ complete conversion contract

CPU, meta, and non-uint8 paths retained the Torch implementation

The Triton fast paths are not used universally.

For Hopper:

if device is not CUDA
or dtype is not uint8:

    use Torch pack/unpack

Blackwell follows the same broad policy.

CUDA uint8 packed values
→ Triton conversion
CPU

meta tensors

non-uint8 representations

→ Torch fallback

This preserves functionality outside the kernel’s supported execution domain.

The patch does not claim:

one Triton converter handles every tensor type and device

The old implementation remains both:

  • a fallback,

  • and a reference oracle for tests.


Standalone fake-CUDA execution remained unsupported

The PR explicitly excludes standalone fake-CUDA execution from the supported fast path.

That means the patch should not be described as universally compatible with every tensor object that reports CUDA-like metadata.

The direct supported path is real CUDA uint8 storage under the documented Triton runtime.

This is another reason the fallback and test matrix matter.


Empty tensors remained valid

An optimized kernel often fails at a zero-sized edge because launch dimensions or compile-time divisors become zero.

The Blackwell path computes:

grid_k

grid_n

from packed tensor dimensions.

For an empty tensor, either can be zero.

The actual launch grid can therefore be empty.

But compile-time expressions inside the kernel still contain divisions involving:

GRID_K

GRID_N

The code supplies:

max(grid_k, 1)

max(grid_n, 1)

as compile-time divisors while retaining the actual zero-sized launch count.

Its comment states:

Keep indexing divisors valid
even when the launch grid is empty.

The kernel performs no data work, while the compiled indexing expressions remain valid.


Zero size did not waive the FP4 packing contract

Two FP4 values share one byte.

A dimension being used as the packing axis must therefore have an even logical extent.

The patch adds explicit rejection for:

odd K packing dimension

and, during conversion between packing orientations:

odd N packing dimension

The tests include combinations such as:

odd K

while N = 0

and:

odd N

while K = 0

This is deliberate.

An empty orthogonal dimension does not make an invalid packed representation well-defined.

No elements are executed

does not imply

the descriptor may violate its packing contract

The test suite protects both valid empty tensors and invalid odd packing independently.


The patch expanded empty-case coverage substantially

The PR reports:

204 valid-empty cases

within:

434 passing tests

8 existing skips

These cases covered relevant CPU, meta, and CUDA layout round trips.

The test files verify:

  • logical shape preservation,

  • physical storage shape,

  • canonical unpacked shape,

  • and round-trip storage shape.

The empty result is not accepted merely because no exception occurred.

Its metadata and layout contract are checked.


Retiling had to preserve exact packed bytes

Blackwell layouts can be parameterized by values such as:

block_k

block_n

Converting between two differently parameterized shuffled layouts may require:

old shuffled layout

→ canonical representation

→ new shuffled layout

The patch tests:

  • multiple block sizes,

  • padded and non-padded dimensions,

  • higher-rank inputs,

  • noncontiguous source views,

  • forward swizzle,

  • inverse unswizzle,

  • and retiling.

The result must match the existing Torch transformation byte for byte.

This matters because two layouts can represent the same logical FP4 values while using different physical padding or tile order.

A numerical comparison after dequantization could miss a byte-layout incompatibility that breaks the later matmul loader.

The tests therefore compare the packed storage itself.


The patch did not replace all MXFP4 layout code

PR #11360 modified the CUDA value-conversion paths for:

Hopper MXFP4 values

Blackwell shuffled MXFP4 values

It did not redesign every neighboring component.

The PR states that:

default Blackwell value layouts

and

scale layouts

were unchanged.

The patch also did not remove the Torch implementations.

They remain for unsupported devices or types and for reference behavior.

The correct scope is:

selected CUDA packed-value layout conversions

not:

the complete MXFP4 software stack was replaced

What the patch directly changed

The merged patch directly added or changed:

  • a Triton Hopper bit-conversion kernel,

  • one forward/inverse CUDA dispatch helper for Hopper values,

  • direct Blackwell canonical-to-shuffled and shuffled-to-canonical conversion,

  • register-level two-by-two FP4 nibble transposition,

  • masked Blackwell padding writes,

  • stride-aware indexing for noncontiguous leading dimensions,

  • input-device-scoped Triton launches,

  • odd-packing validation,

  • CUDA empty-tensor handling,

  • peak-allocation regressions,

  • exact CPU-versus-CUDA byte comparisons,

  • retiling and padding tests,

  • and a two-GPU/non-default-stream regression.

The PR was merged into Triton main on August 21, 2026, as commit:

2ddf835cf63ab52507ddd513489915b5b039fd60

What the patch does not establish

The public evidence does not establish that:

  • MXFP4 matmul itself became 4–21× faster,

  • every model loads faster end to end,

  • a fresh-process uncached conversion is faster,

  • every Hopper or Blackwell layout has the same memory reduction,

  • every input shape reaches the reported speedup,

  • AMD or non-CUDA backends use these new kernels,

  • all FP4 formats share this layout path,

  • reserved VRAM falls by the same amount as allocated-memory peak,

  • or one model’s total memory requirement falls by 91%.

The supported conclusion is narrower:

For the documented CUDA MXFP4 value-layout conversions, replacing tensor-wide Torch intermediate graphs with fused Triton kernels reduced warm conversion time and removed large temporary allocations.


Why the title uses “nearly 3 GiB”

The largest reported peak was:

2,944 MiB

Converting to GiB:

2,944 / 1,024

=

2.875 GiB

That peak:

  • included the conversion output,

  • excluded the existing 128 MiB input,

  • and measured allocated tensor memory rather than total reserved VRAM.

The article is not claiming:

a 128 MiB checkpoint permanently occupied 3 GiB

It is claiming:

the old H100 direct unswizzle conversion
temporarily reached a measured 2,944 MiB allocation peak
while processing the documented 128 MiB packed input

That distinction is essential.


The complete old and new paths

Hopper before

packed CUDA bytes
        ↓
layout permutation and contiguous representation
        ↓
tensor-wide int32 casts
        ↓
multiple shifts and masks
        ↓
four-way stack
        ↓
int16 / uint8 narrowing
        ↓
destination layout

Hopper after

packed CUDA bytes
        ↓
four byte loads per local output group
        ↓
register-level masks and shifts
        ↓
one transformed store
        ↓
destination layout

Blackwell before

canonical N-packed bytes
        ↓
full K-packed intermediate
        ↓
full padded intermediate
        ↓
contiguous transpose
        ↓
tile reshape and contiguous permutation
        ↓
shuffled destination

Blackwell after

canonical source coordinates
        ↓
direct shuffled destination coordinates
        ↓
2x2 nibble transpose in registers
        ↓
masked destination store

The final lesson is that intermediate layouts are not free

A layout conversion can be mathematically simple.

No values are learned.

No matrix multiplication occurs.

No new numerical approximation is introduced.

But an implementation can still consume several times the weight size when every conceptual stage is represented as another tensor.

repack

transpose

pad

stack

cast

permute

Each word sounds like a metadata operation.

Several of them become full memory operations once contiguity or a new dtype is required.

The important question is:

Does this intermediate representation need to survive
outside the conversion step?

If the answer is no, the conversion kernel can often calculate the final destination directly.

When every destination byte can be derived from a small source neighborhood, a full-tensor intermediate is not part of the data contract. It is only an implementation cost.

Triton #11360 removed that cost from the supported CUDA MXFP4 paths.

It made weight preparation substantially lighter after warmup.

It did not change what the later matmul computes.

That boundary is what makes the improvement both large and precisely scoped.


Related material


Patch status: Merged into Triton main
Affected stage: MXFP4 weight preparation and layout conversion
Affected CUDA layouts: Hopper MXFP4 values and Blackwell shuffled MXFP4 values
Measured packed input: 128 MiB, logical shape [16,4096,4096]
Largest old peak: 2,944 MiB on H100 direct unswizzle
Corresponding new peak: 256 MiB
Largest reported warm speedup: 21.26× on GB300 direct swizzle
H100 warm speedup: Approximately 4.2–4.4× in the documented cases
Core repair: Replace full-tensor Torch intermediates with fused Triton conversion kernels
Hopper mechanism: Four-byte local bit conversion in registers
Blackwell mechanism: Direct canonical-to-shuffled addressing and self-inverse 2×2 nibble transpose
Fallback: CPU, meta, and non-uint8 paths retain Torch implementations
Launch ownership: Input CUDA device and its active stream
Cold boundary: Empty-cache first call can take approximately 0.88–1.08 seconds
End-to-end claim: No complete model-loading or matmul speedup claimed
Validation: Author-reported 434 passed, eight existing skips, ten integration jobs passed

This is a standalone Resonetic Lab Code & Patch Analysis article and the final article in the current review batch.

#Triton #NVIDIA #Hopper #Blackwell #MXFP4 #FP4 #GPUProgramming #MemoryOptimization #LayoutConversion #WeightPreparation #CompilerCorrectness #CodeAnalysis

Popular posts from this blog

AMD Is Trying to Turn Kernel Diversity Into One Software Asset

PyTorch #188031 — Why Did Forward Use 64-Bit Indexing While Backward Still Used uint32_t?

NVIDIA CUTLASS #3017 — Why load() and store() Mapped the Same Tile Coordinate to Different Addresses