torch_musa: silently wrong results past ~3e9 elements

[cheersluvs]

Component: torch_musa (specific aten op not yet isolated)

Environment

component	version
card	MTT S5000, 79.9 GiB
MUSA toolkit / mcc	4.3.5 (release_musa_4.3.0)
torch	2.9.0
torch_musa	2.9.0

Most serious of the three, because nothing raises.

Summary

A pure-PyTorch computation over a 3.22e9-element bf16 tensor returns results that
differ from the same computation performed on slices of the same input, even
though the computation is per-row independent and slicing is therefore exact.

Evidence

The operator under test is per-token independent: RMSNorm is per (token, head),
RoPE and quantisation are per token, and each token writes its own cache slot, so
no reduction crosses a token boundary. Slicing along tokens must be exact. On
identical inputs at 98304 tokens × 64 heads (q is 3.22e9 elements):

	vs the Triton kernel, whole input
torch reference, whole input	4.16% of elements differ, max abs diff 31.3
torch reference, 4 slices (each < 2^31)	0.0002% differ, max abs diff 0.031 (≈1 bf16 ULP)

As a control, the Triton kernel's own output whole-input vs 4-way-sliced was
bit-identical. So the kernel is size-independent and the whole-input torch
path is the one that is wrong.

Under the operator's tolerance (rtol=atol=1e-2) the sliced reference gives 0
violations out of 3.22e9; the whole-input reference gives 134 million.

Values of 31.3 where the correct magnitude is ~1 are wrong data, not imprecision.

What is not the cause

• Not addressing. A Triton kernel writing through computed 64-bit offsets is
  correct from 2^30 through 2^32 elements.
• Not the kernel under test. See the bit-identical control above.
• Not an old stack. Reproduced on the versions in the table above. We have
  also seen the failing threshold shift between toolkit releases (from "every
  shape at or above 2^31" to "somewhere between 2.15e9 and 3.22e9"), so this area
  appears to have changed at some point without the defect being fixed.

Not yet isolated

We have not narrowed this to a single aten operator; the reference performs
several dozen torch calls (elementwise mul/add, rsqrt, reductions along the
last dimension, abs/max, log2/ceil/exp2, an fp8 cast, and
scatter-style cache writes). We can bisect on request.

If you would rather narrow it internally, the fastest starting point is any
per-row reduction or elementwise chain over a bf16 tensor larger than 3e9
elements, compared against a chunked evaluation of the same expression.

Impact

Silent wrong results at a tensor size that is ordinary for LLM inference. Any
correctness check performed on MUSA with tensors this large may be validating
against a corrupted oracle — the failure looks like a bug in whatever is being
tested, which is exactly how it presented to us.

[cheersluvs]

Update — the summary above is broader than our evidence supports. Narrowing it, and flagging that we may withdraw it.

We tried to reduce this to a minimal reproduction and could not. On the same card and stack (torch 2.9.0, torch_musa 2.9.0+0680da1), with a [98304*64, 512] bf16 tensor — 3.22e9 elements, the same size as in the report — evaluating a per-row computation over the whole tensor and over four row-slices agrees exactly, zero differing elements, for every candidate we tried:

op	max abs diff	differing elements
x.float().pow(2).mean(-1)	0	0
x.float().sum(-1)	0	0
x.abs().amax(-1)	0	0
x.float().mul(2).sum(-1)	0	0
x[:, :1].float().squeeze(-1)	0	0

So this is not a general "per-row torch computation goes wrong past 3e9 elements", and the title and summary as written overstate what we can show.

The original observation stands as an observation — it was measured — but it came from a much longer chain: per-head RMSNorm, GPT-J RoPE, UE8M0 FP8 quantisation and scatter-style cache writes, several dozen torch calls, with in-place modification of the large tensor. Whatever is happening is somewhere in there rather than in an ordinary reduction.

We are re-running that chain whole versus sliced on the current stack to see whether it still reproduces at all; the card is occupied at the moment, so this is pending. If it does not reproduce we will close this issue. Flagging now rather than leaving the broad claim standing unqualified.