fix(ascend): correct flattened argmax semantics

Two issues on the dim=None path:

1. keepdim. The output was allocated with torch.empty([]) unconditionally, so
   keepdim=True was dropped. PyTorch returns shape [1] * inp.dim().
2. Block size. block_size was derived only from sqrt(numel) with no upper
   bound. For large inputs it went past what the device handles on this path
   and the returned index pointed at an element that was not the maximum.
   It is now capped at 1024, the same bound argmin uses.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

Author: zheng1

src/flag_gems/runtime/backend/_ascend/ops/argmax.py

@@ -113,13 +113,14 @@ def argmax(inp, dim=None, keepdim=False, *, dtype=None):
         M = inp.numel()
         if dtype is None:
             dtype = inp.dtype
-        block_size = triton.next_power_of_2(math.ceil(math.sqrt(M)))
+        block_size = min(1024, triton.next_power_of_2(math.ceil(math.sqrt(M))))
         mid_size = triton.cdiv(M, block_size)
         block_mid = triton.next_power_of_2(mid_size)
 
         mid_value = torch.empty((mid_size,), dtype=dtype, device=inp.device)
         mid_index = torch.empty((mid_size,), dtype=torch.int64, device=inp.device)
-        out = torch.empty([], dtype=torch.int64, device=inp.device)
+        out_shape = [1] * inp.dim() if keepdim else []
+        out = torch.empty(out_shape, dtype=torch.int64, device=inp.device)
 
         with torch_device_fn.device(inp.device):
             argmax_kernel_1[(mid_size, 1, 1)](

tests/test_argmax.py

@@ -16,6 +16,7 @@
 import torch
 
 import flag_gems
+from flag_gems.runtime import torch_device_fn
 
 from . import accuracy_utils as utils
 from . import conftest as cfg
@@ -60,9 +61,12 @@ def test_argmax(shape, dim, keepdim, dtype):
     else:
         inp = torch.randn(shape, dtype=dtype, device=flag_gems.device)
 
-    ref_inp = utils.to_reference(inp)
+    use_cpu_ref = flag_gems.vendor_name == "ascend" and dim is None and keepdim
+    ref_inp = inp.cpu() if use_cpu_ref else utils.to_reference(inp)
 
     ref_out = torch.argmax(ref_inp, dim=dim, keepdim=keepdim)
+    if use_cpu_ref and not cfg.TO_CPU:
+        ref_out = ref_out.to(inp.device)
     with flag_gems.use_gems():
         res_out = torch.argmax(inp, dim=dim, keepdim=keepdim)
 
@@ -87,3 +91,22 @@ def test_argmax_full_reduction_noncontiguous(dtype):
         res_out = torch.argmax(inp)
 
     utils.gems_assert_equal(res_out, ref_out)
+
+
+@pytest.mark.argmax
+@pytest.mark.skipif(
+    flag_gems.vendor_name != "ascend", reason="Ascend-specific regression test"
+)
+def test_argmax_large_flattened_block_size():
+    # block_size was derived from sqrt(numel) with no upper bound. For an input
+    # this large it went past what the device handles on the flattened path, and
+    # the returned index pointed at an element that was not the maximum.
+    for seed in (0, 2, 3, 9):
+        torch_device_fn.manual_seed_all(seed)
+        inp = torch.randn((200, 2560, 3), dtype=torch.bfloat16, device=flag_gems.device)
+        ref_out = torch.argmax(inp.cpu())
+
+        with flag_gems.use_gems():
+            res_out = torch.argmax(inp)
+
+        torch.testing.assert_close(res_out.cpu(), ref_out, atol=0, rtol=0)