Triton #11171 — Why a Larger block_n Exhausted Blackwell’s 512-Column Tensor Memory

Triton Blackwell NVFP4 × Dense Matmul, Part 1 of 3 — When a persistent accumulator and an expanded operand competed for the same TMEM budget

Larger matrix-multiplication tiles are often associated with better GPU performance.

A larger tile can:

reduce the number of program instances

increase arithmetic work per tile

improve data reuse

amortize scheduling and launch overhead

But a tile is not free.

Every increase in tile size also increases the amount of state that the kernel must keep alive.

On NVIDIA Blackwell, some persistent matrix-multiplication paths place more than the accumulator in Tensor Memory, or TMEM.

They may also stage an expanded operand there.

That creates a hard resource question:

Can the accumulator tile and the additional operand representation fit inside the same 512-column TMEM budget?

Triton PR #11171 fixed a case where the answer was no.

The affected path combined:

Blackwell GPU

persistent matmul

NVFP4-scaled left-hand activation

dense BF16 or FP16 right-hand operand

activation-scale-swizzled layout

The planner could select:

block_n = 256

for an N = 256 matmul.

But this persistent path needed TMEM for both:

the output accumulator

and

an expanded operand used by the dot operation

The resulting kernels exceeded Blackwell’s 512-column Tensor Memory limit.

In the validation matrix added with the patch:

8 executions requested 640 TMEM columns

4 executions requested 576 TMEM columns

All twelve failed with Tensor Memory exhaustion.

The fix did not redesign the matmul.

It changed the default planning boundary:

For persistent Blackwell NVFP4 × dense FP16/BF16 matmul:

default block_n
→ capped at 128

That smaller N tile left enough TMEM headroom for the expanded operand.

The same twelve cases then passed.

The important lesson is not merely that 256 was too large.

It is that:

A planner cannot choose a matrix tile from matrix dimensions alone. It must also account for every hidden representation that the lowering keeps alive in limited on-chip memory.

What `block_n` controls

A matrix multiplication can be written as:

C[M, N]
=
A[M, K]
×
B[K, N]

A tiled kernel does not calculate the full output matrix in one program instance.

It divides the output into smaller regions.

block_m
→ number of output rows handled by one tile

block_n
→ number of output columns handled by one tile

block_k
→ reduction depth processed during one K step

For example:

block_m = 128

block_n = 256

means one output tile covers:

128 rows
×
256 columns

The accumulator for that tile must preserve partial sums for all of those output positions while the kernel traverses K.

A larger block_n therefore means:

more output columns per tile

fewer tiles across N

larger accumulator state

The first two properties may improve performance.

The third consumes more on-chip resources.

The planner naturally preferred `block_n = 256`

Before the fix, Triton’s NVIDIA matmul planner used a default N-tile rule that could select:

block_n = 256

when:

N > 128

and the relevant precision configuration allowed the ordinary larger tile.

The newly added regression case used:

M = 256

N = 256

K = 128

For this shape, a default planner seeing N = 256 had an obvious candidate:

block_n = 256

That would let one N tile span the full output width.

Conceptually:

N = 256

block_n = 256

number of N tiles
=
1

Compared with:

N = 256

block_n = 128

number of N tiles
=
2

the larger tile appears attractive.

But the number of output tiles was not the only constraint.

The activation-scale layout also pushed the kernel into the persistent path

The new tests did not use only a plain NVFP4 tensor.

They exercised a routed or ragged matmul with activation-scale HBM swizzling.

The relevant path uses a Blackwell activation-scale layout rather than an ordinary strided scale tensor.

That layout requires the persistent kernel.

The planner’s logic is broadly:

Scale layout is ordinary and strided
→ persistent execution may be optional

Scale layout requires Blackwell-specific handling
→ persistent execution becomes required

Once the scale layout is recognized as the Blackwell activation-scale-swizzled form, the planner also enforces:

block_m = 128

to align the computation with that swizzled scale layout.

The affected default tile therefore became approximately:

block_m = 128

block_n = 256

block_k determined separately

This was not simply a large ordinary GEMM tile.

It was a persistent Blackwell tile with scale-layout obligations and additional TMEM staging.

What a persistent matmul changes

A non-persistent kernel generally launches program instances over a tile grid and lets each instance handle its assigned output region.

A persistent design tries to keep a controlled number of program instances resident and have them process multiple tiles.

Conceptually:

Ordinary tiled execution

Program 0 → tile 0
Program 1 → tile 1
Program 2 → tile 2
...

Persistent execution

Resident program 0
→ tile 0
→ tile 4
→ tile 8

Resident program 1
→ tile 1
→ tile 5
→ tile 9

Persistent execution can improve scheduling and data-movement behavior.

It also tends to make resource planning stricter.

The kernel must keep the state required by its persistent dot implementation within the available on-chip budget.

For the path addressed by #11171, that state included more than the output accumulator.

NVFP4 storage does not remain four-bit at every internal stage

NVFP4 represents activation values using a compact FP4 payload together with scaling metadata.

At the external data-format level:

activation payload
→ FP4

block scaling
→ NVFP4 scale structure

It is tempting to assume that a four-bit input must consume very little internal storage.

But the dot operation may not keep every representation in the same packed form throughout lowering.

The PR’s planner comment states that the Blackwell NVFP4-left-hand-side plus dense-right-hand-side persistent dot stages an expanded operand in TMEM in addition to the accumulator.

The resource path is therefore closer to:

Packed or scaled NVFP4 activation
        ↓
internal expanded operand representation
        ↓
stage expanded operand in TMEM
        ↓
perform dot operation

At the same time:

partial output accumulator
→ also occupies TMEM

The total demand is not:

accumulator only

It is:

accumulator columns

+

expanded-operand columns

+

any additional path-specific TMEM requirements

FP4 reduces external data size.

It does not guarantee that every internal representation remains four-bit.

Two valid allocations exceeded one physical budget

The accumulator allocation was valid by itself.

The expanded operand allocation was also required by the chosen lowering.

The failure came from their sum.

Accumulator
→ requires part of TMEM

Expanded operand
→ requires additional TMEM

Combined request
→ exceeds 512 columns

The validation results reported two failing totals:

576 columns

640 columns

Both exceed:

512-column limit

The public PR does not break every 576- or 640-column request into a complete per-object allocation table.

It does establish the structural cause:

persistent accumulator
+
expanded operand
→ Tensor Memory exhaustion

The precise split should not be invented beyond that evidence.

This was not an HBM out-of-memory failure

The phrase “memory exhaustion” can suggest that the model used too much HBM.

That was not the reported boundary.

The failure concerned Blackwell Tensor Memory used by the kernel’s dot implementation.

HBM
→ stores large tensors and model data

TMEM
→ limited on-chip tensor-compute resource

A system can have ample device HBM and still fail to compile or execute a kernel whose tile requires more TMEM than the architecture permits.

The relevant question was not:

Does the GPU have enough total memory for the matrices?

It was:

Does one persistent program’s live tensor state
fit inside the 512-column TMEM allocation limit?

The failure happened because the planner selected an impossible kernel

The input tensors themselves were valid.

The matrix dimensions were valid.

The NVFP4 and BF16 or FP16 data types were supported.

The reference calculation was not the problem.

The selected tile configuration requested more TMEM than the hardware allocation could provide.

Valid operation

+

valid shape

+

valid data formats

+

invalid resource plan

=

kernel construction failure

This is a planner correctness issue.

A planner is not correct merely because it chooses a theoretically high-performance tile.

It must choose a tile that the compiler and hardware can instantiate.

A larger tile can cross a discontinuous boundary

Performance tuning often appears continuous.

block_n = 64
→ certain throughput

block_n = 128
→ perhaps more throughput

block_n = 256
→ perhaps still more throughput

Resource validity is not continuous in the same way.

TMEM request <= 512
→ kernel can be allocated

TMEM request > 512
→ kernel cannot be allocated

There is no partial success at:

576 columns

The configuration does not merely become slower.

It becomes invalid.

That creates a hard planning boundary:

fast-looking tile below limit
→ candidate

fast-looking tile above limit
→ impossible

Why the planner needed an NVFP4-specific predicate

Triton already had a related safeguard for another Blackwell path:

MX-scaled left-hand operand
×
dense FP16/BF16 right-hand operand

That route also stages an expanded operand in TMEM during persistent execution.

Its default block_n was already capped to leave headroom.

It may seem that NVFP4 could simply be classified as part of that existing MX condition.

The patch deliberately did not do that.

Instead, it added a separate predicate:

is_blackwell_nvfp4_lhs_dense_rhs(...)

The reason is that the existing MX predicate controls more than the N-tile safeguard.

It also participates in tuning decisions such as:

block_k

num_warps

Treating NVFP4 as ordinary MX input everywhere could therefore alter unrelated planner decisions.

The patch needed to share one resource rule:

persistent expanded operand
→ cap default block_n

without sharing every tuning rule.

Shared resource behavior did not imply identical tuning behavior

The two paths have one common property:

Both need TMEM headroom beyond the accumulator.

They do not necessarily share:

the same microblock size

the same K tile

the same warp count

the same data-format semantics

the same performance optimum

The implementation therefore uses:

MX predicate
OR
NVFP4 predicate

only at the block_n safeguard.

Elsewhere, the predicates remain distinct.

This is a precise abstraction boundary.

Shared fact:
expanded TMEM operand exists

Backend policy:
other tuning parameters may differ

What identifies the affected NVFP4 path

The new predicate recognizes a specific configuration.

Conceptually, it requires:

Blackwell-or-later NVIDIA target

precision configuration exists

left-hand activation scale exists

NVFP4 microblock size is used

right-hand scale is absent

output scale is absent

left-hand data type is FP4

right-hand data type is FP16 or BF16

This describes:

scaled NVFP4 activation
×
dense FP16/BF16 operand

It does not classify every FP4 matmul as the same path.

Examples outside the predicate include:

FP4 × FP4

scaled right-hand operand

different scale configuration

pre-Blackwell targets

non-persistent execution

The safeguard is targeted at the resource contract that actually requires it.

The fix caps only the default `block_n`

The new planner condition includes:

constraints.get("block_n", None) is None

That means the cap applies only when Triton is selecting the tile automatically.

No explicit block_n constraint
→ planner may cap default at 128

Explicit block_n supplied
→ planner preserves that constraint

The patch does not silently rewrite a caller’s explicit tile request.

This preserves a useful distinction:

Planner-owned default
→ planner must make it safe

Caller-owned override
→ caller retains responsibility for the requested constraint

An explicit oversized constraint is not guaranteed to fit merely because the automatic path was repaired.

The patch simply refuses to change its meaning silently.

Why silently overriding explicit constraints could be misleading

Suppose a caller requests:

block_n = 256

That request may be intended for:

a controlled experiment,
a planner test,
a future hardware configuration,
a different lowering route,
or a deliberate failure check.

If the planner silently converted it into:

block_n = 128

the kernel might run, but the caller would not be testing the requested configuration.

Requested plan:
256

Executed plan:
128

The result could look successful while violating the explicit experiment’s contract.

PR #11171 therefore repairs automatic planning without erasing explicit intent.

The actual code change is small

The central planner change is effectively:

if (
    is_persistent
    and no_explicit_block_n
    and (
        is_blackwell_mx_lhs_dense_rhs(...)
        or is_blackwell_nvfp4_lhs_dense_rhs(...)
    )
):
    block_n = min(block_n, 128)

Before the fix, only the MX condition entered that safeguard.

After the fix, the NVFP4 activation × dense FP16/BF16 condition does as well.

The numerical change is one line:

256
→ 128

The restored contract is larger:

An automatically planned persistent Blackwell tile must leave enough TMEM for every live operand representation, not only the accumulator.

Why `128` creates headroom

The old default could use:

block_n = 256

The fix limits the default to:

block_n = 128

Reducing the output-column tile reduces accumulator demand.

Conceptually:

Before

128 output rows
×
256 output columns
→ larger accumulator

After

128 output rows
×
128 output columns
→ smaller accumulator

The expanded operand still needs TMEM.

But the accumulator no longer consumes so much of the 512-column budget that the combined allocation reaches 576 or 640 columns.

The patch does not state that every possible block_n = 128 configuration is universally safe.

It establishes that this cap makes the targeted default Blackwell NVFP4 × dense path fit in the tested cases.

The fix trades one large N tile for more smaller tiles

For:

N = 256

the planning change is approximately:

Before

block_n = 256
→ one N tile

After

block_n = 128
→ two N tiles

This may increase the number of tile tasks.

But the comparison is not:

one fast tile

versus

two slower tiles

because the one-tile configuration could not be allocated.

The real comparison is:

one impossible tile

versus

two valid tiles

Resource feasibility comes before performance comparison.

The newly added cases exposed the missing planner boundary

Before #11171, the test matrix already contained simple NVFP4 × dense cases such as:

M = 128

N = 128

K = 128

plain layout

Those cases naturally fit a smaller tile.

They did not require the planner to confront:

N = 256

persistent execution

activation-scale swizzling

routed or ragged variants

The patch added:

M = 256

N = 256

K = 128

ragged routing

NVFP4 activation

activation-scale HBM swizzling

for both:

dense BF16 right-hand side

dense FP16 right-hand side

These cases reached the path where the default 256-column tile became invalid.

A passing small case did not prove the larger planner region

The old matrix established that:

NVFP4 × dense matmul can work

for a plain 128 × 128 × 128 case

It did not establish:

the automatic planner remains valid

when N grows to 256

and scale layout requires persistence

and routed variants expand the test matrix

The distinction is similar to other dispatch gaps.

Same high-level operator

different shape and layout

different planner decision

different resource contract

Testing the data type did not automatically test every tile-selection region.

Main plus the new cases failed twelve times

The new two top-level cases expanded through the existing test machinery into multiple variants.

On the unmodified main revision, the validation recorded:

12 failures

all due to Tensor Memory exhaustion

The reported required allocations were:

8 cases
→ 640 columns

4 cases
→ 576 columns

This consistency is important.

The failures were not a mixture of:

numerical mismatches,
unsupported data types,
routing errors,
scale-layout errors,
and random runtime failures.

They converged on one planner boundary:

TMEM request > 512

The fixed branch passed the same cases

With the block_n safeguard in place, the identical generated cases produced:

12 passed

84 skipped

The validation included variants involving:

scatter

gamma

K-ragged behavior

The patch also recorded:

NVIDIA planner and split-K tests:
16 passed, 4 skipped

and:

related MXFP/NVFP4 × dense matmul matrix:
444 passed, 516 skipped

The larger matrix matters because the change touched shared planner logic.

A local success would not be enough if the new NVFP4 predicate accidentally altered MXFP tuning or unrelated matmul routes.

Why preserving MXFP behavior mattered

The safeguard already applied to MX-scaled inputs.

Adding NVFP4 through a broad shared predicate could have changed:

block_k selection

warp-count selection

other MX-specific tuning decisions

The validation over hundreds of related cases checked that the planner change did not broadly destabilize the existing matrix.

The code also documents the intended separation directly:

Keep NVFP4 separate from the MX predicate

because the MX predicate also controls block_k and num_warps.

This is not merely a comment about code organization.

It identifies which behavior is shared and which must remain backend-format-specific.

The patch repaired a default, not a universal theorem about tile size

The following statement is supported:

For the tested persistent Blackwell NVFP4-activation
× dense BF16/FP16 path,
the automatic 256-column N tile exceeded TMEM,
and capping the default at 128 fixed the tested failures.

The following broader claims are not established:

block_n = 256 is always invalid on Blackwell

block_n = 128 is optimal for every NVFP4 matmul

every FP4 kernel stages the same expanded operand

every persistent kernel uses the same TMEM layout

all explicit block_n = 256 constraints must fail

The resource rule belongs to one planner region.

The right tile depends on hidden lowering state

At the high-level matmul interface, the planner sees:

M

N

K

input dtypes

output dtype

The valid tile also depends on lower-level facts:

persistent or non-persistent execution

scale layout

whether an operand is expanded

where that expanded operand is staged

accumulator representation

target architecture

TMEM allocation rules

Two operations with the same matrix dimensions can require different resource plans because their lowerings keep different state alive.

Same M, N, K

different precision path

different TMEM obligations

different valid block_n

This is why planner predicates cannot stop at matrix shape.

Compression can move the cost rather than eliminate it

NVFP4 reduces the external activation representation.

But the path may need to expand that representation internally.

HBM traffic or storage
→ reduced by FP4

Internal compute representation
→ may require expansion

TMEM pressure
→ can increase because both accumulator and expansion coexist

The benefit of compression remains real.

Its resource effect is not uniform across every layer of the execution stack.

A format can reduce one bottleneck while creating a new planning constraint elsewhere.

More on-chip reuse can require more on-chip state

Persistent kernels aim to keep work close to the compute units and reduce repeated scheduling or movement.

That can improve efficiency.

It also increases the value of accurate resource accounting.

More persistent reuse

+

more live internal representations

→ greater on-chip allocation pressure

The performance technique and the resource problem are two sides of the same design.

A planner is part of correctness

Tile selection is sometimes treated as performance policy.

In this case, it determined whether the kernel was constructible.

Correct arithmetic implementation

+

invalid tile plan

→ operation still fails

The planner therefore participates in the correctness boundary.

Its obligations include:

select a supported route

respect explicit constraints

fit register and shared-memory limits

fit TMEM limits

preserve layout contracts

produce a runnable kernel

A high-performance invalid tile is not a valid planner result.

The patch’s direct scope

PR #11171 directly changed:

automatic tile selection for persistent Blackwell NVFP4-left-hand-side × dense FP16/BF16 matmul,
the default block_n safeguard,
detection of the NVFP4 microblock configuration,
separation between NVFP4 and broader MXFP tuning predicates,
and test coverage for routed, activation-scale-swizzled M=256, N=256, K=128 cases.

It retained:

explicit block_n constraints,
existing MXFP block_k tuning,
existing MXFP warp tuning,
other FP4 combinations,
and non-persistent behavior outside the predicate.

What this patch confirms

The source and validation confirm that:

Blackwell persistent NVFP4 activation × dense FP16/BF16
stages an expanded operand in TMEM.

The previous automatic planner could choose block_n = 256.

The new test matrix produced 12 TMEM-exhaustion failures
requiring 576 or 640 columns against a 512-column budget.

Capping the automatic block_n at 128
made those 12 cases pass.

The fix remains present in Triton main.

What remains unproven

The public evidence does not establish:

a complete allocation table for every 576- and 640-column failure,
that block_n = 128 is the fastest feasible configuration,
the performance difference between valid 128- and hypothetical valid 256-column routes,
that every explicit oversized constraint is rejected at the same stage,
identical behavior on GB200, B200, or every Blackwell derivative,
that every routed model workload uses this exact planner configuration,
or that the safeguard covers all future expanded-operand paths automatically.

The patch was validated on:

NVIDIA GB300

SM 10.3

The result should remain tied to the tested planner and hardware path.

The deeper failure was incomplete resource accounting

The old planner knew:

output tile dimensions

It did not apply the existing expanded-operand TMEM safeguard to this NVFP4 configuration.

The missing relationship was:

NVFP4 activation × dense operand

+

persistent Blackwell lowering

→ expanded operand also occupies TMEM

Once that relationship was represented explicitly, the planner could make a safe default decision.

Larger tiles are useful only after they are legal

The optimization intuition was:

N = 256

block_n = 256

→ cover all output columns at once

The hardware contract answered:

Accumulator plus expanded operand

→ 576 or 640 TMEM columns

→ maximum 512

→ invalid

The fix changed the question from:

How wide can this tile be?

to:

How wide can the accumulator remain
while leaving room for every additional live operand?

That is the correct planning question.

The fastest theoretical tile is irrelevant if the internal state required by that tile cannot coexist on the hardware.

Part 2: why the fix changed only one default

The next article examines why the patch:

caps default block_n at 128

but

preserves explicit block_n constraints

and why NVFP4 received a separate predicate instead of being folded into the existing MXFP condition.

The central question will be:

How can two data paths share one resource limit without being forced to share every tuning policy?

Part 3: why the old test matrix missed the failure

The final article examines the regression design.

Existing case

M = 128
N = 128
K = 128
plain layout
→ passed

New case

M = 256
N = 256
K = 128
ragged
activation-scale swizzled
→ exposed impossible default tile

It will also separate:

two top-level cases

from:

twelve generated failing executions

and explain why the same matrix had to cover fast paths, routed variants, scale layouts, and neighboring MXFP behavior.

Related material

Patch status: Merged into Triton main
Affected path: Persistent Blackwell NVFP4 activation × dense BF16/FP16 matmul
Hardware validation: NVIDIA GB300, SM 10.3
TMEM limit: 512 columns
Pre-fix validation: 12 Tensor Memory exhaustion failures; eight requested 640 columns and four requested 576
Planner repair: Cap the automatic block_n at 128
Explicit constraint behavior: User-supplied block_n remains unchanged
Neighboring policy: MXFP block_k and warp tuning remain separate
Post-fix validation: The same 12 generated cases passed

This is Part 1 of a three-part series on Triton’s Blackwell NVFP4 × dense matmul planning boundary.

Part 2 examines why the patch modified only the planner-owned N-tile default, preserved explicit constraints, and introduced a separate NVFP4 predicate rather than reusing every MXFP tuning rule.

Part 3 examines why the previous plain 128×128×128 tests did not expose the persistent routed path, and how the expanded test matrix converted twelve TMEM failures into a regression contract.

#Triton #NVIDIA #Blackwell #GB300 #NVFP4 #FP4 #TensorMemory #TMEM #MatrixMultiplication #GPUProgramming #KernelPlanning #CodeAnalysis