[mark14wu]

Summary

A softmax over fp16 data fails under the sanitizer and the race detector, although it runs natively and under TRITON_INTERPRET=1:

x = tl.load(x_ptr + offs)          # fp16
e = tl.exp(x - tl.max(x, 0))

ValueError: Expected dtype ['fp32', 'fp64'] but got fp16

tl.max, tl.min, and tl.sum convert their input before reducing (triton/language/standard.py):

• tl.max and tl.min convert bf16 to fp32 and, without return_indices, any float narrower than 32 bits to fp32 and any narrower integer to int32.
• tl.sum converts signed integers narrower than 32 bits to int32 and unsigned ones (including int1) to uint32, unless dtype is given.

TileLens replaces these functions for symbolic clients, and ReduceSymbolicExpr gave every reduction its input's dtype; _op_reduce_sum_overrider also ignored dtype. So tl.max(x, 0) stayed fp16 and tl.exp's dtype check failed. The overriders now compute Triton's result dtype and pass it to ReduceSymbolicExpr, which takes an optional dtype.

Tests

• New in tests/end_to_end/test_sanitizer.py:
  • test_reduce_max_promotes_fp16_like_triton: the softmax above. It fails on main with the error above.
  • test_reduce_result_dtypes_match_triton: tl.static_asserts on the result dtypes of tl.sum, tl.max, and tl.min for int8, uint8, fp16, and bf16 inputs, with and without return_indices and dtype. It fails on main. The same assertions hold natively and under TRITON_INTERPRET=1 (Triton 3.8.0 and 3.7.1).
• pytest tests/: 366 passed, 6 skipped on main; 368 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

[github-actions]

Performance Benchmark

Benchmark	main (min)	PR (min)	Change	Samples
gemm	0.101s	0.101s	+0.5%	20 / 20
gemm_oob	0.113s	0.113s	+0.5%	20 / 20
indirect_load	0.021s	0.021s	+0.3%	20 / 20
nested_loop	0.224s	0.224s	-0.1%	20 / 20
block_pointer_loop_advance	0.119s	0.121s	+1.3%	20 / 20
liger_jsd	0.133s	0.134s	+0.7%	20 / 20
flaggems_layernorm	0.375s	0.378s	+0.6%	20 / 20
swiglu	0.164s	0.163s	-0.3%	20 / 20
cross_entropy	0.937s	0.942s	+0.5%	20 / 20
fused_linear_jsd	0.201s	0.203s	+0.9%	20 / 20
Total	2.389s	2.400s	+0.5%	N/A

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