Unsupported ROCm Tests `tests/pallas/triton_pallas_test.py`

[JehandadKhan]

Description

The tests in tests/pallas/triton_pallas_test.py are being skipped and need to be triaged to assess if they can be enabled. If they cant due to missing ROCm features, then we need to implement such a feature.

• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_array_atomic_add0
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_array_atomic_add1
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_atomic_cas0
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_atomic_cas1
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_atomic_cas2
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_atomic_cas3
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_atomic_cas4
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_atomic_cas5
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_atomic_counter0
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_atomic_counter1
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_atomic_counter2
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_atomic_counter3
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_atomic_counter4
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_fp_dtype_cast0
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_fp_dtype_cast1
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_fp_dtype_cast10
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_fp_dtype_cast11
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_fp_dtype_cast12
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_fp_dtype_cast13
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_fp_dtype_cast14
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_fp_dtype_cast15
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_fp_dtype_cast16
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_fp_dtype_cast17
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_fp_dtype_cast18
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_fp_dtype_cast19
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_fp_dtype_cast2
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_fp_dtype_cast20
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_fp_dtype_cast21
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_fp_dtype_cast22
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_fp_dtype_cast23
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_fp_dtype_cast24
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_fp_dtype_cast3
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_fp_dtype_cast4
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_fp_dtype_cast5
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_fp_dtype_cast6
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_fp_dtype_cast7
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_fp_dtype_cast8
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_fp_dtype_cast9
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_invalid_broadcasted_load0
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_invalid_broadcasted_load1
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_invalid_broadcasted_load2
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_masked_load_store0
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_masked_load_store1
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_masked_load_store10
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_masked_load_store11
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_masked_load_store12
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_masked_load_store13
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_masked_load_store14
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_masked_load_store15
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_masked_load_store16
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_masked_load_store17
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_masked_load_store18
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_masked_load_store19
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_masked_load_store2
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_masked_load_store20
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_masked_load_store21
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_masked_load_store22
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_masked_load_store23
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_masked_load_store24
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_masked_load_store3
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_masked_load_store4
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_masked_load_store5
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_masked_load_store6
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_masked_load_store7
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_masked_load_store8
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_masked_load_store9
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_masked_oob_load_store_slice
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_scalar_atomic_add_f16
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_scalar_atomic_add_f32
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_scalar_atomic_add_i32
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_scalar_atomic_max_f32
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_scalar_atomic_max_i32
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_scalar_atomic_min_f32
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_scalar_atomic_min_i32
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_softmax_1_1_128_float32
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_softmax_1_1_2_float32
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_softmax_1_1_256_float32
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_softmax_1_1_32_float32
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_softmax_1_1_64_float32
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_softmax_1_129_256_float32
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_softmax_1_2_128_float32
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_softmax_1_2_256_float32
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_softmax_1_2_32_float32
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_softmax_1_2_64_float32
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_softmax_1_255_256_float32
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_softmax_2_1_128_float32
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_softmax_2_1_2_float32
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_softmax_2_1_256_float32
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_softmax_2_1_32_float32
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_softmax_2_1_64_float32
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_softmax_2_129_256_float32
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_softmax_2_2_128_float32
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_softmax_2_2_256_float32
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_softmax_2_2_32_float32
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_softmax_2_2_64_float32
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_softmax_2_255_256_float32
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_softmax_23_1_128_float32
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_softmax_23_1_2_float32
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_softmax_23_1_256_float32
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_softmax_23_1_32_float32
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_softmax_23_1_64_float32
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_softmax_23_129_256_float32
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_softmax_23_2_128_float32
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_softmax_23_2_256_float32
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_softmax_23_2_32_float32
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_softmax_23_2_64_float32
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_softmax_23_255_256_float32
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_softmax_4_1_128_float32
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_softmax_4_1_2_float32
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_softmax_4_1_256_float32
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_softmax_4_1_32_float32
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_softmax_4_1_64_float32
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_softmax_4_129_256_float32
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_softmax_4_2_128_float32
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_softmax_4_2_256_float32
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_softmax_4_2_32_float32
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_softmax_4_2_64_float32
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_softmax_4_255_256_float32

System info (python version, jaxlib version, accelerator, etc.)

JAX version 0.8.0
MI300 Accelerator
Independent of OS and Python version

[Shahzeb-AMD]

Fixed via #34575

Resolved UTs

• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_array_atomic_add0
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_array_atomic_add1
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_atomic_cas0
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_atomic_cas1
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_atomic_cas2
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_atomic_cas3
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_atomic_cas4
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_atomic_cas5
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_atomic_counter0
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_atomic_counter1
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_atomic_counter2
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_atomic_counter3
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_atomic_counter4
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_fp_dtype_cast0
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_fp_dtype_cast1
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_fp_dtype_cast10
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_fp_dtype_cast11
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_fp_dtype_cast12
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_fp_dtype_cast13
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_fp_dtype_cast14
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_fp_dtype_cast15
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_fp_dtype_cast16
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_fp_dtype_cast17
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_fp_dtype_cast18
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_fp_dtype_cast19
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_fp_dtype_cast2
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_fp_dtype_cast20
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_fp_dtype_cast21
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_fp_dtype_cast22
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_fp_dtype_cast23
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_fp_dtype_cast24
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_fp_dtype_cast3
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_fp_dtype_cast4
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_fp_dtype_cast5
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_fp_dtype_cast6
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_fp_dtype_cast7
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_fp_dtype_cast8
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_fp_dtype_cast9
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_invalid_broadcasted_load0
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_invalid_broadcasted_load1
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_invalid_broadcasted_load2
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_masked_load_store0
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_masked_load_store1
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_masked_load_store10
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_masked_load_store11
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_masked_load_store12
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_masked_load_store13
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_masked_load_store14
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_masked_load_store15
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_masked_load_store16
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_masked_load_store17
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_masked_load_store18
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_masked_load_store19
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_masked_load_store2
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_masked_load_store20
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_masked_load_store21
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_masked_load_store22
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_masked_load_store23
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_masked_load_store24
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_masked_load_store3
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_masked_load_store4
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_masked_load_store5
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_masked_load_store6
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_masked_load_store7
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_masked_load_store8
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_masked_load_store9
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_masked_oob_load_store_slice
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_scalar_atomic_add_f16
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_scalar_atomic_add_f32
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_scalar_atomic_add_i32
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_scalar_atomic_max_i32
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_scalar_atomic_min_i32
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_softmax_1_1_128_float32
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_softmax_1_1_2_float32
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_softmax_1_1_256_float32
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_softmax_1_1_32_float32
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_softmax_1_1_64_float32
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_softmax_1_129_256_float32
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_softmax_1_2_128_float32
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_softmax_1_2_256_float32
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_softmax_1_2_32_float32
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_softmax_1_2_64_float32
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_softmax_1_255_256_float32
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_softmax_2_1_128_float32
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_softmax_2_1_2_float32
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_softmax_2_1_256_float32
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_softmax_2_1_32_float32
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_softmax_2_1_64_float32
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_softmax_2_129_256_float32
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_softmax_2_2_128_float32
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_softmax_2_2_256_float32
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_softmax_2_2_32_float32
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_softmax_2_2_64_float32
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_softmax_2_255_256_float32
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_softmax_23_1_128_float32
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_softmax_23_1_2_float32
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_softmax_23_1_256_float32
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_softmax_23_1_32_float32
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_softmax_23_1_64_float32
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_softmax_23_129_256_float32
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_softmax_23_2_128_float32
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_softmax_23_2_256_float32
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_softmax_23_2_32_float32
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_softmax_23_2_64_float32
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_softmax_23_255_256_float32
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_softmax_4_1_128_float32
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_softmax_4_1_2_float32
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_softmax_4_1_256_float32
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_softmax_4_1_32_float32
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_softmax_4_1_64_float32
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_softmax_4_129_256_float32
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_softmax_4_2_128_float32
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_softmax_4_2_256_float32
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_softmax_4_2_32_float32
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_softmax_4_2_64_float32
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_softmax_4_255_256_float32

Blocked on ROCm

• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_scalar_atomic_max_f32
• tests/pallas/triton_pallas_test.py:::TritonPallasTest:::test_scalar_atomic_min_f32