[Ascend] Tensor.exponential_() produces 1024 NaNs for large tensors (n >= 65536) — breaks LLM random sampling

[Al2O3Al2O3]

[Ascend] Tensor.exponential_() produces 1024 NaNs for large tensors (n ≥ 65536) — breaks all LLM random sampling

Repo: https://github.com/flagos-ai/FlagGems

Summary

FlagGems' exponential_ op (flag_gems/ops/exponential_.py, Triton kernel) on
the Ascend backend writes exactly 1024 NaN values into large output tensors.
For tensors with a trailing dimension n ≥ 65536 the last (or one) 1024-element
block comes out as NaN instead of an exponential-distributed value; smaller
tensors are fine.

This silently breaks all random sampling in vLLM on Ascend. vLLM's
random_sample does q = torch.empty_like(probs); q.exponential_(); probs.div_(q).argmax(-1).
With a vocab of 151936, q gets NaNs, probs / NaN = NaN, and argmax over an
all-NaN row returns index 0 — so every sampled token is id 0 and completions
degenerate to "!!!!". Greedy decoding (pure argmax, no exponential_) is
unaffected.

Environment

Component	Version
flag_gems	5.4.0rc2.post3+gad89ec477
flagtree	0.7.0rc2+ascend3.5
vllm	0.20.2+flagos
torch / torch_npu	2.10.0
CANN	9.0.0
Device	Ascend910_9382 (910C)

Minimal reproduction

import torch, torch_npu, flag_gems
flag_gems.enable()
d = torch.device("npu")

for n in [2048, 4096, 8192, 16384, 32768, 65536, 131072, 151936]:
    q = torch.empty((1, n), dtype=torch.float32, device=d)
    q.exponential_()
    print(n, "nan=", int(q.isnan().sum()))

Output:

2048   nan=0
4096   nan=0
8192   nan=0
16384  nan=0
32768  nan=0        # OK up to here
65536  nan=1024     # NaNs appear
131072 nan=1024
151936 nan=1024     # real Qwen3 vocab size

The NaN count is a constant 1024 regardless of tensor size once the
threshold (somewhere in 32768 < n ≤ 65536) is crossed — pointing at a
tiling / block-boundary bug in the kernel (one 1024-element block's philox /
safe_fast_log output is not written correctly). uniform_() on the identical
tensor is fine, so this is specific to exponential_, not a general RNG or
memory-layout issue.

How it was found (in a live vLLM server)

Instrumenting random_sample in a running TP=2 server, for a
temperature=1.0 request:

[DIAG] q_dtype=float32 q_contig=True q_isnan=1024 exp_nan=1024 uni_nan=0 \
       pmax=0.41 psum=1.0 pnan=0 arg=[0]

i.e. probs is correct (sums to 1, no NaN), but q from exponential_() has
1024 NaNs, so the sampled token collapses to 0. uniform_() on the same tensor
in the same process returns valid values.

Impact

No usable random sampling on Ascend for any model with a large vocab (Qwen3,
etc.). Combined with the separate Triton top-k/top-p compile crash
(FlagTree issue: _topk_topp_kernel / ConvertTritonIRToLinalgIR PassManager),
there is currently no working sampling path on this Ascend stack — only
greedy decoding works.

Ask

Fix the block-boundary handling in the Ascend exponential_ Triton kernel so
large tensors don't emit NaN. Happy to provide the dumped kernel or run further
diagnostics.

[bin913]

Root cause identified and fixed in #6539 (the PR carries Resolves #6446, so merging it closes this issue).

Why exponential_ was the visible victim

exponential_ is listed in the Ascend backend's CUSTOMIZED_UNUSED_OPS, so gems does not take it over and the torch_npu composite implementation runs instead: uniform_() followed by log_(). The corruption came from the log_ half.

Root cause

triton-ascend does not queue programs past the number of vector cores the device has. A launch that uses more CTAs than that is padded up to the next multiple of the core count, and every padding CTA reports program_id == 0. Program 0 therefore runs many times over.

The Ascend codegen config capped max_grid_size[0] at 48, which is neither <= the 40 vector cores of an Ascend910 nor a multiple of it, and the generic log_ launched cdiv(n_elements, 1024) programs with no cap at all. For n = 41984 that is 41 programs, padded to 80, with program 0 executed 33 times. Because log_ is in-place, the first block ended up as log(log(x)) — NaN for x < 1, i.e. exactly the 1024 NaNs in [0..1023] reported here.

Reproduced directly with an atomic counter kernel on Ascend910B (tl.atomic_add(counter + pid, 1)):

grid	program 0 executions
32	1
40	1
48	33
64	17
80	1

uniform_() uses cdiv(N, 4096) tiles and stays under the core count, which is why only the log_() half of the composite corrupted the result — and why exponential_ broke only for n >= 41984, while n = 40960 (exactly 40 tiles) was fine.

Fix in #6539

1. src/flag_gems/utils/codegen_config_utils.py: cap the Ascend grid at the backend's real core count, so num_ctas = min(max_grid_size[0], num_tiles) stays inside the range triton-ascend launches as a single, duplicate-free wave.
2. src/flag_gems/ops/log_.py: make the kernel grid-stride and cap its launch at max_grid_size[0], so an in-place kernel stays correct on a backend that limits the grid.
3. tests/test_log_.py: add test_log_large_inplace over sizes that straddle the 40-tile boundary.

Verification (Ascend910B, CANN 9.0.0, triton-ascend 3.5)

input	log_ NaNs before	log_ NaNs after	exponential_ NaNs before	exponential_ NaNs after
32768	0	0	0	0
40960	0	0	0	0
41984	1024	0	1024	0
49152	1024	0	1024	0
65536	1024	0	1024	0
151936	1024	0	1024	0
(2, 151936)	—	—	1024	0
(64, 151936)	—	—	1024	0

Before the fix exponential_'s max() was NaN, which is why vLLM random sampling always returned token id 0; after the fix max() is a normal sample (11.4–14.1 across these sizes). tests/test_log_.py goes from 10 failed / 20 passed on master to 1 failed / 33 passed, the single remaining failure being the native int16 path, which fails identically on master. The new regression test is not vacuous: run against master it fails on 3 of its 4 sizes.

Related

This is the same root cause as the symptom in #542 ([Ascend] Random sampling (temperature > 0) always returns token 0), which lives on the same code path, so #6539 should resolve that one too.