[BugFix][Argreduce] Fix MACA tiled argmax/argmin always returning index 0

[Lfan-ke]

Problem

The MACA tiled argmax/argmin always returns 0.

argreduce_maca.py exists because MACA gives a block 64 KiB of shared memory, so wide rows have to be reduced tile by tile. The per-tile extremum and its index are computed correctly; the bug is in the merge step that folds each tile into the running row extremum:

for i in T.Parallel(block_m):
    row_extreme[i] = T.if_then_else(
        tile_extreme[i] > row_extreme[i],
        tile_extreme[i],
        row_extreme[i],
    )
    out_idx[i] = T.if_then_else(
        tile_extreme[i] > row_extreme[i],   # row_extreme[i] was just overwritten
        tile_idx[i],
        out_idx[i],
    )

row_extreme[i] is written first, so by the time out_idx is decided, row_extreme[i] >= tile_extreme[i] holds by construction and the predicate is always false. out_idx is therefore never written and keeps the 0 it was filled with. Both the argmax branch and the argmin branch have it.

Only the tiled path is affected, and only MACA reaches it — the single-tile kernel scans the whole row in one shared buffer and is fine, and on an NVIDIA card the larger shared-memory budget keeps almost everything single-tile. So this is a silent wrong answer confined to MetaX, on the wide-row argmax the manifest declares (lm-head).

Fix

Decide out_idx against the previous row extremum, i.e. write it before row_extreme is overwritten. Two statements swapped, in each of the two branches. The comparison stays strict (> / <), which is what keeps the leftmost-index semantics that torch.argmax / torch.argmin also have.

Test node delta

File                           Base    HEAD    Delta
----------------------------------------------------
tests/ops/test_argreduce.py      84      88       +4
----------------------------------------------------
TOTAL                            84      88       +4

Growth: +4.8%

The four new cells are regression coverage for this defect. They are the only cases in the file that reach the tiled path at all: with a 64 KiB budget and a per-column cost of elem_bytes + 4, compute_tile_n keeps everything below N ≈ 11008 (fp16 / bf16) on the single-tile path, and the widest existing fixture is N = 4096. So the patched lines had no coverage whatsoever — which is how "always returns index 0" shipped in the first place. The test asserts kernel.config["tile_n"] > 0 so it cannot silently regress back to the single-tile path, and it puts the extremum in the last tile so a merge that never writes out_idx is visible rather than accidentally right. It is skipif(not is_maca()), since the tiled kernel is MACA-only.

Test Result

MetaX C500, MACA 3.5.3.20.

On dev, the new cases fail:

4 failed, 84 deselected

With this patch:

4 passed, 84 deselected

Whole file, with this patch:

88 passed

ruff check --config pyproject.toml is clean on both files.

One thing this PR does not fix

float32 at N = 16384 picks tile_n = 8192, i.e. 8192 × (4 + 4) = 65536 bytes, exactly device_smem_budget(), and the launch is rejected:

RuntimeError: MACALaunch mcErrorInvalidValue
grid=(4,1,1), block=(128,1,1) dyn_smem_bytes=65536

The same kernel launches at tile_n = 8064 (64512 bytes), and softmax at N = 32768 requests exactly 65536 bytes (2 × 16384 × 2) yet launches and passes today. Survival at exactly device_smem_budget() thus depends on the kernel, and the budget sits on the edge by construction: compute_tile_n returns tile_n_max, consuming the budget to the byte, whenever N_padded is a multiple of it, as every power-of-two hidden size is. Sweeping 256-aligned N_padded up to 40960 for argreduce fp32, 71 of 128 tiled cases land exactly on 65536.

I have not isolated the cause, so I am not moving the budget here. The device reports sharedMemPerBlock = 65536 with no reserve, and a bare mxcc kernel launches at 65536, as does the driver API, so the shortfall is not a hardware limit. This is separate from the merge bug and predates it; the new cases use fp16 / bf16 to stay clear. Happy to follow up once the cause is known.

[gemini-code-assist]

Code Review

This pull request corrects the execution order in the argreduce MACA kernel to ensure that out_idx is updated before row_extreme is overwritten, resolving a bug where the index was evaluated against the updated extreme value instead of the previous one. As there are no review comments, I have no feedback to provide.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[Lfan-ke]

@codex review please. read code at pr and repositories.