[tonyh168]

What

On the Iluvatar backend, flag_gems' *_non_inner reduction kernels read rows they should have masked off. This PR keeps those calls off that kernel: the reduced axis is moved to the back, the original entry point runs there on the (correct) inner-dim kernel, and the result is put back.

Fixes the Iluvatar half of #617 ("same input, repeated forward, different results" on BI-V150 with torch_fl 2.10.0+corex4.5.0).

Why

The tile those kernels load is (TILE_N, TILE_K) with mask = (n < N) & (k < K). On Iluvatar the row half of that mask is not applied — rows past N are read anyway:

• TILE_N > N ⇒ the reduction covers TILE_N rows instead of N (silently wrong), and
• the last group of rows reads past the end of the tensor into whatever the allocator last handed out ⇒ repeated identical calls return different numbers.

Minimal reproduction, a standalone Triton kernel with no flag_gems involved (same box, same Triton package, driven by torch_fl):

variant (N=4, TILE_N=8)	result	expected
tl.load(mask=(n<N)&(k<K), other=0.0)	28	6
load + tl.where(n < N, v, 0)	28	6
TILE_N == N (no masked row)	6	6
per-row loop	6	6
1-D mask / inner-axis mask only	correct	correct

The same kernel driven by vendor CoreX torch is correct, so the miscompile is in the path this package installs, not in the kernel source.

Not a corner case for a real model: Qwen3.5-VL's vision tower computes its bilinear position embedding as sum(dim=1) on a (556, 4, 1024) tensor. Tracing operator by operator (1666 ops, 3 runs):

• the sum is the first and only op whose output differs between two identical forwards — its inputs are bit-identical across runs; all 1064 later differences are downstream of it,
• it is off a CPU float64 reference by 4.68 at an input absmax of 3.14,
• six repeats of the vision tower gave four distinct results under torch_fl; vendor torch gave one, bit-identical across six.

prod, std, logsumexp and the non-inner softmax/log_softmax share the same tile and are wrong for the same reason, so all of them are wrapped.

How

torch_fl/compile/iluvatar_masked_reductions.py, installed from torch_fl/__init__.py right after the existing layer-norm workaround (same reason for the ordering: the generated kernel resolves the callable on first use), gated on _build_accelerator() == "iluvatar".

• Only a single non-inner int dim takes the workaround; inner-dim, multi-dim, None and non-int dims are forwarded unchanged, so the transpose is paid only on the calls that were wrong.
• keepdim is forced off for the inner call and re-applied by the reshape (positional for sum_dim/logsumexp, keyword-only for std), and the dim form (int vs single-element list, as sum.dim_IntList passes) is preserved, because flag_gems' own entry points reject the other form.
• softmax/log_softmax permute the result back; the kernels need a contiguous view, so one is materialised.

Verification (BI-V150, iluvatar-109, one card)

check	before	after
mask_blast equivalent vs CPU: sum, softmax, log_softmax, prod, std, logsumexp (N=4, K=1024, dim=1)	all wrong (rel 0.6 – 2.7e12)	all ≤ 6e-8
sum shape sweep, fp32 (N∈{1..16}, K∈{64..1024})	wrong for N=4, K≥256	all exact
vision tower × 6, fixed pixel_values	4 distinct results	6/6 bit-identical, mean 1.08693671 / std 77.26651001
full model, 12 samples (bf16)	4/12 out of range, 2 NaN	12/12 same side of threshold, max|Δ| = 7.54e-3, no NaN
two separate processes, all 12 samples	—	bit-identical

(mean/std differs from vendor torch's 1.08774793 / 77.35216522 by 0.15%, which is the bf16 GEMM gap, not this defect.)

Unit tests cover the part that would silently drift: which calls take the path at all, which axis moves, keepdim round-trip, dim-form preservation, that *_out bindings are left alone (including the aliased case), the bool→int64 out= dtype rule, and idempotency. ruff check / ruff format --check clean under the repo's pinned 0.15.12.

Review follow-ups (8ebd0c7)

out= entry points are left alone. The alias sweep no longer rewires any *_out binding. An out= entry point is not the function it is named after — sum_dim_out passes the caller's keepdim through and reads the result dtype off the input, so torch.sum(bool_input, dim=1, out=int64_out) lands in an int64 buffer; the wrapper forces keepdim=False and returns a fresh tensor. out= calls therefore keep the workaround's documented scope (unmasked, still wrong on this backend) instead of acquiring a second silent change. Two tests pin it. On flag_gems 5.4.0 and 5.5.0.dev0 the two are distinct objects, so no caller saw the changed behaviour — but the sweep had no guard, and now it does.

The new tests run in CI as manifest step iluvatar_reduction_contract, on every platform. Not in the platform-agnostic job (no torch there by design — the file would be collected and skipped) and not in unit_tests_per_file (widening that group to Iluvatar still surfaces five unrelated failures in the rest of tests/unit: test_profiler_privateuse1 crashes, the ascend/gcu qwenimage rope files fail, test_flagcicd_coverage/reports fail). The file resolves torch_fl through the installed wheel first, so it runs in the wheel-only workspace and in a source checkout alike — 19 passed in each on iluvatar-109.

The two Iluvatar-lane failures in tests/integration/ops/test_dtype_cast_copy.py were the test's premise, not a regression, and are scoped:

case	why it failed on Iluvatar	change
test_cuda_abi_cast_does_not_route_through_host ×4	reads csrc/aten/copy_ops.cc; the wheel-only workspace has no csrc/	skips with that reason (10 passed / 6 skipped with no csrc/ present)
test_bf16_fp32_cast_is_transfer_free ×2	copy_ops.cc has no DtypeCast for USE_ILUVATAR, so the same-device dtype branch takes the documented CPU-round-trip fallback; backends_iluvatar.conf agrees (_to_copy.out = none)	skips for the four builds with no on-device cast (bpu, iluvatar, mps, tsingmicro); unchanged everywhere else

An Iluvatar on-device cast is a real performance follow-up, but it is a new feature in copy_ops.cc, not a fix to this wrapper.

End-to-end is unchanged by the follow-up commit — vision tower 3/3 bit-identical, 12 samples 12/12 same side, max|Δ| = 7.54e-3, no NaN, and bit-identical to the pre-follow-up run.

🤖 Generated with Claude Code

[lvyufeng]

Overview

Conclusion: request changes. Moving the reduced axis to the last dimension is a reasonable workaround for the Iluvatar masked-load defect. The flag_gems 5.4.0 kernels select the inner-dimension path when K == 1, so this approach uses the existing correct implementation instead of duplicating a kernel. I found no error in axis movement or keepdim handling for the checked reduction paths.

There are two blocking code and test-coverage findings, plus an Iluvatar integration failure that must be resolved before a green lane can demonstrate the fix end to end.

Blocking findings

1. sum_dim_out is changed by the alias sweep. It has a different argument behavior from sum_dim: with an unspecified dtype, the original passes keepdim through and converts a bool input to int64; the wrapped path forces keepdim=False, so torch.sum(bool_input, dim=1, out=bool_output) can write bool output instead of the expected int64 result. Please retain the original sum_dim_out behavior and add a regression test for this out= case.

2. The new unit test file does not run in this PR's CI. It is missing from .github/workflows/agnostic-checks.yml, and the manifest's per-file unit test group names only DCU. The Iluvatar lane ran integration tests but did not execute tests/unit/test_iluvatar_masked_reductions.py, so there is no CI evidence for the wrapper or alias-sweep behavior.

3. The Iluvatar lane failed. Its operator group reported six failures in the shared copy test:

   • test_bf16_fp32_cast_is_transfer_free observed four host/device transfers in each direction.
   • Four test_cuda_abi_cast_does_not_route_through_host cases failed because the wheel-only workspace has no csrc/aten/copy_ops.cc. That test is marked anyplatform and is parameterized by all four CUDA-ABI platform names, so it runs with platform=cuda on Iluvatar.

   These failures are outside the files changed here, but they block a green Iluvatar check. The failure occurs before the shared suite can demonstrate end-to-end correctness of this change.

Verification performed

I inspected all three changed files, the FlagGems 5.4.0 reduction and softmax entry points, the Iluvatar registration/conf route, import ordering, and the CI manifest. The focused Ruff 0.15.12 check and format check pass on both changed Python files. The new tests were not run locally because the available host lacks the required PyTorch 2.10 / FlagGems runtime.

Current Iluvatar job: 120245425740. The operator group ended with 6 failed, 694 passed, 1573 skipped.

Landscape

Area	Assessment
Axis movement and keepdim restoration for the tested reductions	Reasonable
Preserving int vs list dimension form	Covered by synthetic tests
out= typing and argument behavior	Needs a fix and regression check
CI execution of the new tests	Missing
Iluvatar integration evidence	Failing at the time of review

Out of scope (a separate PR would be too big)

• Fixing FlagGems backward kernels or the Iluvatar masked-load miscompile itself.
• Running the full cross-platform hardware matrix after the blocking items are fixed.

[lvyufeng]

P1: The alias sweep also replaces sum_dim_out, whose argument behavior differs from sum_dim: with an unspecified dtype, the original passes keepdim through and converts a bool input to int64; this wrapper forces keepdim=False, which can make torch.sum(bool_input, dim=1, out=bool_output) write bool output instead of int64. Please retain the original behavior for the out= entry point and add a regression test.

[tonyh168]

Fixed in 8ebd0c7. The alias sweep now skips every *_out binding, and the reason is in the module docstring next to the sweep: an out= entry point is not the function it is named after — sum_dim_out passes the caller's keepdim through and reads the result dtype off the input, so torch.sum(bool_input, dim=1, out=int64_out) lands in an int64 buffer. The wrapper forces keepdim=False and returns a fresh tensor, which would leave the caller's buffer untouched.

Two regression tests cover it: test_out_entry_points_are_never_rewired (including the worst case where a build aliases sum_dim_out to sum_dim, which is exactly what an identity sweep cannot tell apart) and test_bool_sum_out_dtype_survives_the_patch (mirrors flag_gems' dtype rule and asserts the int64 out= result).

For the record, on flag_gems 5.4.0 and 5.5.0.dev0 sum_dim_out is a distinct object from sum_dim, so no caller saw the changed behaviour — but the sweep had no guard, so the fix is a guard rather than a note. out= calls now keep the workaround's documented scope (unmasked, therefore still wrong on this backend) instead of acquiring a second silent change.

[lvyufeng]

P1: The new tests are absent from the platform-agnostic test list, while the integration manifest's per-file unit group still names only DCU. The Iluvatar lane in this PR ran integration tests but did not execute this file. Please wire the new tests into an Iluvatar job and confirm that they run there; the synthetic wrapper tests cannot detect the changed sum_dim_out typing path.

[tonyh168]

Fixed in 8ebd0c7 as a manifest step of its own: iluvatar_reduction_contract runs tests/unit/test_iluvatar_masked_reductions.py on every platform, Iluvatar included.

Why not the two lists named in the review:

• .github/workflows/agnostic-checks.yml installs no torch by design (its header says so, and the bpu/ exclusion records the same constraint). The file needs torch, so it would be collected and skipped there — collected, but not evidence.
• unit_tests_per_file currently names only dcu because the rest have gaps. I widened it to Iluvatar and ran the whole directory on iluvatar-109 first: 95/100 files pass, and the five failures are unrelated — test_profiler_privateuse1.py crashes (exit -11), test_ascend_qwenimage_rope.py and test_gcu_qwenimage_rope.py fail, test_flagcicd_coverage.py (exit 2) and test_flagcicd_reports.py fail. Turning that group on for Iluvatar in this PR would make the lane red for reasons this change does not touch, so it stays a separate triage.

The file resolves torch_fl through the installed wheel first (same arrangement as tests/integration/ops/conftest.py::_platform_path), so it runs from the wheel-only workspace and from a source checkout. Verified both on iluvatar-109: 19 passed in each.

On the coverage point: test_out_entry_points_are_never_rewired now asserts the sweep skips *_out bindings, including the case where a build aliases sum_dim_out to sum_dim, and test_bool_sum_out_dtype_survives_the_patch reproduces flag_gems' bool→int64 dtype rule and checks the out= buffer.

[lvyufeng]

P1: The Iluvatar lane has not yet provided green evidence for the fix. It failed test_bf16_fp32_cast_is_transfer_free with four host/device transfers, plus four test_cuda_abi_cast_does_not_route_through_host cases that could not find csrc/aten/copy_ops.cc in the wheel-only workspace. That test is marked anyplatform and runs with the CUDA platform parameter on Iluvatar. Please resolve or scope this shared copy test and obtain a green Iluvatar lane before merging.

[tonyh168]

Both failures in the Iluvatar operator group were the test's premise rather than a regression, and are scoped in 8ebd0c7.

test_cuda_abi_cast_does_not_route_through_host (×4). The case asserts on csrc/aten/copy_ops.cc, which the wheel-only workspace does not have — parents[3] pointed at a tree with no csrc/, hence the missing-file error. It now skips with that reason instead of raising. Verified by running the file from a directory with no csrc/ on iluvatar-109: 10 passed, 6 skipped (the 4 source cases plus the 2 below), no errors.

test_bf16_fp32_cast_is_transfer_free (×2). This one asserts a runtime property that only holds where an on-device cast exists, and Iluvatar is not one of those builds. In copy_ops.cc's same-device flagos→flagos dtype branch, USE_ILUVATAR is in the group that reaches the // Fallback: CPU round-trip when no on-device cast is available arm — Ascend (aclnnCast), GCU (topsatenToCopy), Kunlunxin (cast) and MUSA (MudnnCopy) each implement a DtypeCast; Iluvatar does not. backends_iluvatar.conf agrees: _to_copy.out = none. So four transfers per direction is the documented design on this backend, not a silent host round trip to catch. The case is now skipped for the four builds with no DtypeCast (bpu, iluvatar, mps, tsingmicro) and unchanged everywhere else — an Iluvatar on-device cast is a real performance follow-up, but it is a new feature in copy_ops.cc, not a fix to this wrapper, and folding it in here would make this PR about something else.

With both scoped, the operator group's 6 failures are gone: on iluvatar-109 the file reports 14 passed / 2 skipped in the checkout. I cannot re-run the self-hosted Iluvatar lane myself; the new commit should produce the green lane this was blocked on, and I will follow it.

[tonyh168]

it