Skip to content

[FIX] Concretize reductions in data-dependent addresses - #535

Open
mark14wu wants to merge 1 commit into
mainfrom
fix/concretize-reductions
Open

mark14wu wants to merge 1 commit into
mainfrom
fix/concretize-reductions

Conversation

@mark14wu

@mark14wu mark14wu commented Nov 7, 2026

Copy link
Copy Markdown
Collaborator

Summary

An address built from a reduction of loaded values fails under the sanitizer:

seq = idx - tl.max(tl.where(cumsum_lens <= idx, cumsum_lens, 0))
NotImplementedError: Concretize for op max is not implemented

When an address depends on loaded values, the sanitizer concretizes the address's subtree (replace_subtree), and ReduceSymbolicExpr had no concretize. This PR adds one for all of its ops: sum, max, min, argmax, argmin, xor_sum, and reduce_or, computed with NumPy in the reduction's result dtype. max and min use fmax/fmin on floats, so they ignore NaNs as Triton's maxnumf/minnumf do; argmax and argmin return the first index, as with tie_break_left=True.

Tests

  • New in tests/end_to_end/test_sanitizer.py: test_sanitizer_concretizes_reductions_in_addresses, a kernel that loads from offsets given by each of the seven reductions of loaded indices. With a large enough source tensor it expects no report; with smaller ones it expects reports at exactly the offsets whose reduction value is out of bounds. All 3 cases fail on main with the error above.
  • pytest tests/: 366 passed, 6 skipped on main; 369 passed, 6 skipped with this PR (Python 3.12, Triton 3.8.0).
  • ruff check and ruff format --check (v0.1.9) pass on the changed files.

🤖 Generated with Claude Code

`ReduceSymbolicExpr` had no `concretize`, so an address built from a
reduction of loaded values, such as `idx - tl.max(tl.where(...))` or
`tl.argmax(vals, 0)`, raised "Concretize for op max is not implemented"
in the sanitizer. It now computes `sum`, `max`, `min`, `argmax`,
`argmin`, `xor_sum`, and `reduce_or` with NumPy, in the reduction's
result dtype; `max` and `min` ignore NaNs, as Triton's do.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
@github-actions

github-actions Bot commented Nov 7, 2026

Copy link
Copy Markdown

Performance Benchmark

Benchmark main (min) PR (min) Change Samples
gemm 0.101s 0.100s -0.4% 20 / 20
gemm_oob 0.113s 0.112s -0.8% 20 / 20
indirect_load 0.021s 0.021s -0.9% 20 / 20
nested_loop 0.225s 0.223s -1.1% 20 / 20
block_pointer_loop_advance 0.120s 0.122s +1.3% 20 / 20
liger_jsd 0.134s 0.133s -0.2% 20 / 20
flaggems_layernorm 0.379s 0.377s -0.4% 20 / 20
swiglu 0.164s 0.164s -0.1% 20 / 20
cross_entropy 0.944s 0.936s -0.8% 20 / 20
fused_linear_jsd 0.203s 0.202s -0.2% 20 / 20
Total 2.404s 2.392s -0.5% N/A

Iterations: 1 warmup + 20 measured
Samples are shown as main / PR; long pytest benchmarks may use fewer samples.

This branch has not been deployed

No deployments
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.

1 participant