[TRITON] Unit test fixes to la tests and gemm afp4wfp4 (ROCm#1494)

* Fix: Add debug mode for mismatch reporting in lean attention tests


* Fix: Adjust x_scales_shuffled slicing according to x_scales size

Author: Boss2002n

op_tests/triton_tests/test_gemm_afp4wfp4.py

@@ -122,7 +122,7 @@ def generate_gemm_afp4wfp4_inputs(
         w_shuffed,
         x_scales[:M],
         w_scales,
-        x_scales_shuffled,
+        x_scales_shuffled[:M],
         w_scales_shuffled,
         out_dtype,
         y,

op_tests/triton_tests/test_la.py

@@ -14,6 +14,8 @@
 import aiter.ops.triton.utils._triton.arch_info as arch_info
 import pytest
 
+DEBUG_MODE = False
+
 
 def get_lean_attn_inputs(
     batch: int,
@@ -422,8 +424,9 @@ def test_persistent_lean_attention(
     try:
         torch.testing.assert_close(ref_out, la_out, atol=atol, rtol=rtol)
     except AssertionError:
-        print("Assertion failed! Showing mismatches:")
-        print_mismatches(ref_out, la_out, atol, rtol)
+        if DEBUG_MODE:
+            print("Assertion failed! Showing mismatches:")
+            print_mismatches(ref_out, la_out, atol, rtol)
         raise  # Re-raise the exception after printing mismatches