Is there a supported FP8 storage dtype on TRN2? F8E4M3FN is rejected even when never used as an operand

[dadmobile]

Describe the bug

Our checkpoint ships natively in FP8 with 128x128 blockwise scales. We keep the
expert weights in torch.float8_e4m3fn purely as storage and dequantise them to
BF16 inside forward, so the hardware never sees an e4m3 operand. The compiler
rejects the tensor regardless:

[ERROR] [NCC_EVRF051] Data type F8E4M3FN is not supported on TRN1/TRN2.
Target TRN3 or later hardware

For a 313B-parameter MoE this is the difference between 19.03 GB and 9.51 GB of
expert weight per rank, which in turn decides how much context and batch fit.

What we are missing is what to do instead:

1. Is there an FP8 storage dtype that is accepted in a graph on TRN2?
2. If not, is uint8 storage plus a 256-entry lookup table the intended pattern?
   That is exact, since e4m3 has exactly 256 values, but it materialises BF16 on the
   read, so it saves resident HBM rather than bandwidth.
3. torch._scaled_mm does compile on TRN2 and NeuronCore-v3 does e4m3 arithmetic.
   Is a true FP8 GEMM the only supported route to FP8 weights on this hardware?

Model Name

GLM-5.3-Flash (glm5_next), 313B parameters, 288 routed experts, FP8 checkpoint
with 128x128 blockwise scales.

Describe the workload type

Inference. Tensor parallel 32, expert parallel 32. The FP8 tensors are expert
weights only, held as storage and dequantised per layer inside forward.

Instance Type

trn2.48xlarge, logical-neuroncore-config: 2

Release version

neuronx-cc                2.27.5334.0+f702b353
aws-neuronx-runtime-lib   2.34.10.0-ac18d186d
aws-neuronx-dkms          2.30.2.0
torch                     2.11.0
torch-xla                 2.11.0

Reproduction Steps

Any graph containing an f8e4m3fn tensor reproduces it. Minimally:

import torch, torch_xla.core.xla_model as xm
dev = xm.xla_device()
w8 = torch.zeros(256, 256, dtype=torch.float8_e4m3fn, device=dev)
scale = torch.ones(2, 2, device=dev)
# dequantise before any arithmetic: the FP8 value is never an operand to a matmul
w = (w8.to(torch.float32) * scale.repeat_interleave(128, 0).repeat_interleave(128, 1))
x = torch.randn(8, 256, device=dev)
y = x @ w.to(torch.bfloat16).t()
xm.mark_step()
print(y.sum().cpu())

Expected: compiles, since no e4m3 operand ever reaches the hardware.
Actual: NCC_EVRF051.

Regression Issue

• Select this option if this issue appears to be a regression.

Possible Solution

If f8e4m3fn cannot be permitted as an inert storage dtype, documenting the intended
alternative would be enough. "The checkpoint is FP8, so keep it FP8" is the obvious
first thing to try with any recent MoE checkpoint, and there is currently nothing
pointing to the supported path.

Logs/Context/Additional Information

We fell back to BF16 storage at 19.59 GB/rank, which does fit on one trn2.48xlarge,
so this is not blocking us. It cost us the headroom we had planned to spend on
longer context and larger batches.

[github-actions]

Hi @dadmobile, Thank you for filing the issue! We will take a look and get back to you.

[Michellejia]

Hi @dadmobile, thank you for filing the issue. Below is the response to your questions. Please let us know if it makes sense.

1. Is there an FP8 storage dtype accepted on TRN2?
   No. f8e4m3fn and the other FP8 variants are rejected during graph ingest on TRN1/TRN2, even if the tensor is only being stored and never used as an operand. There isn't another FP8 dtype we can use just for storage. If FP8 shows up directly in the graph, that requires TRN3.

2. Is uint8 storage + a 256-entry LUT the intended workaround?
   Yes. If the goal is just to reduce storage, using uint8 is the workaround, and it can be done entirely on the framework side without any compiler changes.

torch.float8_e4m3fn is one byte per element, so we can reinterpret the FP8 data as uint8 on CPU with .view(torch.uint8). The XLA graph then only sees uint8, so NCC_EVRF051 isn't triggered.

# Build the LUT once on CPU: all 256 e4m3 codes -> bf16
codes = torch.arange(256, dtype=torch.uint8)
lut_bf16 = codes.view(torch.float8_e4m3fn).to(torch.bfloat16)

# Store weights as uint8. Still 1 byte/element in HBM.
w_u8 = w_fp8.view(torch.uint8).to(dev)
lut = lut_bf16.to(dev)

# Convert back to bf16 before the matmul.
w_bf16 = (
    lut[w_u8.long()]
    * scale.repeat_interleave(128, 0).repeat_interleave(128, 1)
)
y = x @ w_bf16.t()

This keeps the resident HBM footprint at the 9.51 GB/rank you're targeting.

The limitation is that this only helps with HBM capacity. We still materialize BF16 before the GEMM, so DMA traffic into SBUF and the working set there are still BF16. In other words, this gives back the HBM headroom for a larger context/batch, but not the bandwidth savings we'd get from a native FP8 path.

3. Is torch._scaled_mm the only route to a true FP8 GEMM?
   On TRN2, yes, that's the supported path if we want the FP8 operand to remain FP8 all the way to the tensor engine and get the bandwidth benefit as well.

The uint8 approach above is mainly a storage workaround. If the goal is an actual FP8 GEMM rather than just storing the weights in one byte per element, _scaled_mm is the path to use.

[dadmobile]

Thanks very much for this! Your HBM comment on 2 aligns with what we saw in our tests.

On answer 3, we cannot get _scaled_mm to reach the device. It returns correct numbers on an xla:0 tensor, but:

torch_xla counter after one call:   aten::_scaled_mm = 1
HLO for the result: a bare parameter fed to the output, no dot

​
That counter is torch-xla's CPU-fallback counter, so the GEMM ran on the host. Same under VLLM_NEURON_BACKEND=neuron_native. A normal matmul on the same tensors lowers fine.
Are we calling it wrong, or does the XLA lowering not exist in 2.27.5334.0? Happy to send a repro if it helps.