Skip to content

Conversation

@bangtianliu
Copy link
Contributor

@bangtianliu bangtianliu commented Oct 23, 2025

On top of PR #22466, #22490 and #22694 , try to enable O2 for llama 8b fp16 quality tests

ci-extra: test_torch

@bangtianliu bangtianliu marked this pull request as draft October 23, 2025 17:56
@bangtianliu bangtianliu force-pushed the bangtian_debug_argmax branch 5 times, most recently from 55654f5 to caf8fa0 Compare October 24, 2025 16:37
@bangtianliu
Copy link
Contributor Author

For this CI error: https://github.com/iree-org/iree/actions/runs/18786121779/job/53605513613?pr=22394#step:7:111

The issue is with this arg_compare operation:

    %29:2 = iree_linalg_ext.arg_compare dimension(2) ins(%28 : tensor<4x1x128256xf16>) outs(%11, %12 : tensor<4x1xf16>, tensor<4x1xi32>) {
    ^bb0(%arg5: f16, %arg6: f16):
      %31 = arith.cmpf ogt, %arg5, %arg6 : f16
      iree_linalg_ext.yield %31 : i1
    } -> tensor<4x1xf16>, tensor<4x1xi32>

The root cause is that split reduction is not enabled for the arg_compare operation.
Currently, the reduction along dimension 2 (128,256 elements) is processed in a single dispatch, attempting to allocate the entire input to shared memory. This exceeds the 64KB limit.

@bangtianliu bangtianliu force-pushed the bangtian_debug_argmax branch 2 times, most recently from d126efd to d966aea Compare October 30, 2025 22:41
@bangtianliu bangtianliu force-pushed the bangtian_debug_argmax branch 4 times, most recently from 0b1d87d to 40fbc22 Compare November 14, 2025 01:59
@bangtianliu bangtianliu force-pushed the bangtian_debug_argmax branch 12 times, most recently from d998ac6 to 9eb733c Compare November 25, 2025 21:12
@bangtianliu bangtianliu force-pushed the bangtian_debug_argmax branch from 9eb733c to aac4781 Compare November 28, 2025 21:11
bangtianliu and others added 6 commits November 28, 2025 13:15
Signed-off-by: Bangtian Liu <[email protected]>
Signed-off-by: Bangtian Liu <[email protected]>
Signed-off-by: Bangtian Liu <[email protected]>
Signed-off-by: Bangtian Liu <[email protected]>
ci-extra: test_torch
Signed-off-by: Bangtian Liu <[email protected]>

Signed-off-by: Bangtian Liu <[email protected]>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants