[bug] NVFP4 + torch.compile: errors with 3D input #3626

Summary

Under torch.compile, an NVFP4-quantized te.pytorch.Linear fails with

AssertionError: wrong number of dimensions2 for op:
torch.ops.transformer_engine_compile.linear.default

whenever the input activation has 3 (e.g. the normal
(batch, seq, hidden) shape of any real transformer forward pass). A 2D
input (tokens, hidden) does not trigger it, which is why this can look
like it works in isolated/synthetic benchmarks and only surfaces once a
real model is compiled end-to-end.

Environment

• TransformerEngine: 2.21.0.dev0+9233c73
• PyTorch: 2.15.0a0+git8c12b70
• GPU: GB200
• Recipe: transformer_engine.common.recipe.NVFP4BlockScaling() with RHT
  enabled (the default)

Minimal repro

import torch
import transformer_engine.pytorch as te
from transformer_engine.common import recipe

hidden = 5120
model = te.Linear(hidden, hidden, params_dtype=torch.bfloat16, device="cuda")
x = torch.randn(1, 4096, hidden, dtype=torch.bfloat16, device="cuda", requires_grad=True)

def step(inp):
    with te.autocast(recipe=recipe.NVFP4BlockScaling()):
        return model(inp).sum()

compiled_step = torch.compile(step, fullgraph=False)
# loss = step(x)  # works
loss = compiled_step(x)
loss.backward()

Observed failure

AssertionError: wrong number of dimensions2 for op:
torch.ops.transformer_engine_compile.linear.default

  assert_tensor_metadata(buf17, (5120, 1, 2048), (2048, 2048, 1), torch.uint8,
  'torch.ops.transformer_engine_compile.linear.default')

Root cause (AI generated)

NVFP4Quantizer.get_columnwise_shape()
(transformer_engine/pytorch/tensor/nvfp4_tensor.py:309-325) computes the
shape of the RHT/columnwise-quantized buffer as:

colwise_shape = [shape[-1]]
for i in range(len(shape) - 1):
    colwise_shape.append(shape[i])
return tuple(colwise_shape)

i.e. it preserves the input's full rank, permuting each leading dimension
individually ((b, s, h) → (h, b, s)).

torch.compile's Inductor backend uses this function's result (via
inner_tensor_specs) as the compile-time-authoritative shape/stride for
the columnwise buffer, and bakes an assert_tensor_metadata /
assert_size_stride guard into the generated code against it.

But the real CUDA quantization kernel does not preserve rank: it always
collapses every leading dimension into a single M dimension and treats
the tensor as 2D (M, K), regardless of the logical input rank.

For a 2D input, there's only one leading dimension, so the two views
happen to coincide and the bug stays masked. For any ≥3D input — i.e. a
normal batched activation — the predicted (compile-time) shape and the
actual (runtime) shape diverge, and the Inductor-generated guard fails.

Proposed fix (AI generated)

Collapse the leading dimensions into one, matching what the real kernel
does:

--- a/transformer_engine/pytorch/tensor/nvfp4_tensor.py
+++ b/transformer_engine/pytorch/tensor/nvfp4_tensor.py
@@ -327,10 +327,19 @@ class NVFP4Quantizer(Quantizer):
         if len(shape) == 0:
             return tuple()
         # and then after AG, a reorganize kernel will be called to restore the shape
-        colwise_shape = [shape[-1]]
-        for i in range(len(shape) - 1):
-            colwise_shape.append(shape[i])
-        return tuple(colwise_shape)
+        # Collapse all leading dims into one -- the real RHT/columnwise CUDA
+        # kernel always treats the input as a 2D (M, K) matrix (M = prod of
+        # leading dims), never preserving the original tensor rank. Returning
+        # a full-rank shape here (one entry per leading dim) made this
+        # function disagree with the real kernel output whenever the input
+        # had more than one leading dim (e.g. a (batch, seq, hidden)
+        # activation): torch.compile's Inductor backend registers this
+        # function's result as the expected compile-time shape/stride for
+        # the columnwise buffer, so the mismatch surfaced as
+        # "AssertionError: wrong number of dimensions2 for op:
+        # torch.ops.transformer_engine_compile.linear.default" the first time
+        # a >= 3D activation reached an NVFP4 te.Linear under torch.compile.
+        return (shape[-1], math.prod(shape[:-1]))

Verification of the fix

Verified on a real GB200 node, forward + backward, across:

• The originally-failing case: NVFP4, 3D activation, torch.compile
  (fullgraph=False, default mode) — now passes
• NVFP4, 3D activation, reduce-overhead mode — passes
• NVFP4, 2D activation (pre-existing passing case) — still passes, no
  regression
• MXFP8, 3D activation (pre-existing passing case, different
  quantizer) — still passes, no regression

No other code path needed changes — an initial backward-pass error seen
mid-investigation (RuntimeError: mismatch in length of strides and shape) turned out to be a stale torch.compile/Inductor disk cache
(/tmp/torchinductor_*) left over from a pre-patch run, not a second bug;
clearing the cache resolved it with only the one fix above.

cc: @pggPL