Merge branch 'main' into gregory/windows-support

Author: gshimansky

benchmarks/triton_kernels_benchmark/flash_attention_fwd_benchmark.py

@@ -78,6 +78,9 @@ def _attn_fwd(Q, K, V, sm_scale, M, Out,  #
     start_m = tl.program_id(2)
     off_z = tl.program_id(0)
     off_h = tl.program_id(1)
+    if N_CTX <= 512:
+        start_m = tl.program_id(0)
+        off_z = tl.program_id(2)
     qvk_offset = off_z.to(tl.int64) * stride_qz + off_h.to(tl.int64) * stride_qh
 
     # block pointers
@@ -176,6 +179,9 @@ def forward(q, k, v, causal, sm_scale):
     num_warps = 8 if Lq == 64 else 16
     stage = 3 if causal else 1
     grid = lambda args: (q.shape[0], q.shape[1], triton.cdiv(q.shape[2], args['BLOCK_M']))
+    n_ctx = q.shape[2]
+    if n_ctx <= 512:
+        grid = lambda args: (triton.cdiv(q.shape[2], args['BLOCK_M']), q.shape[1], q.shape[0])
     M = torch.empty((q.shape[0], q.shape[1], q.shape[2]), device=q.device, dtype=torch.float32)
 
     if os.getenv('TRITON_INTEL_ADVANCED_PATH', '0') == '0':

python/test/unit/language/test_core.py

@@ -1490,18 +1490,30 @@ def kernel(X):
                           for shape in [(2, 2), (2, 8), (8, 2), (8, 8), (32, 32), (64, 64)]
                           for axis in [0, 1]
                           for num_ctas in num_ctas_list
-                          for dtype_x_str in ['float32', 'uint64', 'int64', 'float64']])
+                          for dtype_x_str in ['float16', 'float32', 'uint64', 'int64', 'float64']])
 def test_tensor_atomic_rmw(shape, axis, num_ctas, dtype_x_str, device):
     check_type_supported(dtype_x_str, device)
+    if is_interpreter() and dtype_x_str == 'float16':
+        pytest.skip('float16 atomic_add does not work in the interpreter mode')
     shape0, shape1 = shape
     # triton kernel
 
     @triton.jit
-    def kernel(Z, X, OLD, AXIS: tl.constexpr, SHAPE0: tl.constexpr, SHAPE1: tl.constexpr):
+    def kernel(Z, X, OLD, AXIS: tl.constexpr, SHAPE0: tl.constexpr, SHAPE1: tl.constexpr, DTYPE: tl.constexpr):
         off0 = tl.arange(0, SHAPE0)
         off1 = tl.arange(0, SHAPE1)
         x = tl.load(X + off0[:, None] * SHAPE1 + off1[None, :])
+
+        if DTYPE == tl.float16:
+            # sum can have bad numerics when accumulating in float16.
+            # if we're dealing with float16, do the sum in float32.
+            x = x.to(tl.float32)
+
         z = tl.sum(x, axis=AXIS)
+
+        if DTYPE == tl.float16:
+            z = z.to(DTYPE)
+
         if AXIS == 1:
             old = tl.atomic_add(Z + off0, z)
             tl.store(OLD + off0, old)
@@ -1515,13 +1527,23 @@ def kernel(Z, X, OLD, AXIS: tl.constexpr, SHAPE0: tl.constexpr, SHAPE1: tl.const
     z = numpy_random(z_shape, dtype_str=dtype_x_str, rs=rs)
     old = np.zeros(z_shape, dtype=getattr(np, dtype_x_str))
     # reference results
-    z_ref = z + np.sum(x, axis=axis, keepdims=False)
+    if x.dtype == np.float16:
+        # do the sum in float32 to reduce numerical variation
+        z_ref = z + np.sum(x.astype(np.float32), axis=axis, keepdims=False).astype(x.dtype)
+    else:
+        z_ref = z + np.sum(x, axis=axis, keepdims=False)
     old_ref = np.copy(z)
     # triton result
     x_tri = to_triton(x, device=device)
     z_tri = to_triton(z, device=device)
     old_tri = to_triton(old, device=device)
-    kernel[(1, )](z_tri, x_tri, old_tri, axis, shape0, shape1, num_ctas=num_ctas)
+
+    def torch_to_triton_dtype(t):
+        if t == torch.float16:
+            return tl.float16
+        return None
+
+    kernel[(1, )](z_tri, x_tri, old_tri, axis, shape0, shape1, torch_to_triton_dtype(x_tri.dtype), num_ctas=num_ctas)
     np.testing.assert_allclose(z_ref, to_numpy(z_tri), rtol=1e-4)
     np.testing.assert_equal(old_ref, to_numpy(old_tri))
 

scripts/skiplist/a770/language.txt

@@ -1,5 +1,166 @@
 # https://github.com/intel/intel-xpu-backend-for-triton/issues/1434
 test/unit/language/test_core.py::test_precise_math[1-tl.math.sqrt_rn(x)-tl.math.sqrt(x.to(tl.float64)).to(tl.float32)]
+test/unit/language/test_core.py::test_dot3d[1-1-32-32-32-32-32-float16-float16]
+test/unit/language/test_core.py::test_dot3d[1-1-32-32-32-32-32-float16-float32]
+test/unit/language/test_core.py::test_dot3d[1-1-32-32-32-32-32-float32-float32]
+test/unit/language/test_core.py::test_dot3d[1-1-32-32-32-32-32-int8-int8]
+test/unit/language/test_core.py::test_dot3d[1-1-64-64-64-32-32-float16-float16]
+test/unit/language/test_core.py::test_dot3d[1-1-64-64-64-32-32-float16-float32]
+test/unit/language/test_core.py::test_dot3d[1-1-64-64-64-32-32-float32-float32]
+test/unit/language/test_core.py::test_dot3d[1-1-64-64-64-32-32-int8-int8]
+test/unit/language/test_core.py::test_dot3d[1-16-32-32-32-32-32-float16-float16]
+test/unit/language/test_core.py::test_dot3d[1-16-32-32-32-32-32-float16-float32]
+test/unit/language/test_core.py::test_dot3d[1-16-32-32-32-32-32-float32-float32]
+test/unit/language/test_core.py::test_dot3d[1-16-32-32-32-32-32-int8-int8]
+test/unit/language/test_core.py::test_dot3d[1-16-64-64-64-32-32-float16-float16]
+test/unit/language/test_core.py::test_dot3d[1-16-64-64-64-32-32-float16-float32]
+test/unit/language/test_core.py::test_dot3d[1-16-64-64-64-32-32-float32-float32]
+test/unit/language/test_core.py::test_dot3d[1-16-64-64-64-32-32-int8-int8]
+test/unit/language/test_core.py::test_dot3d[1-2-32-32-32-32-32-float16-float16]
+test/unit/language/test_core.py::test_dot3d[1-2-32-32-32-32-32-float16-float32]
+test/unit/language/test_core.py::test_dot3d[1-2-32-32-32-32-32-float32-float32]
+test/unit/language/test_core.py::test_dot3d[1-2-32-32-32-32-32-int8-int8]
+test/unit/language/test_core.py::test_dot3d[1-2-64-64-64-32-32-float16-float16]
+test/unit/language/test_core.py::test_dot3d[1-2-64-64-64-32-32-float16-float32]
+test/unit/language/test_core.py::test_dot3d[1-2-64-64-64-32-32-float32-float32]
+test/unit/language/test_core.py::test_dot3d[1-2-64-64-64-32-32-int8-int8]
+test/unit/language/test_core.py::test_dot3d[1-4-32-32-32-32-32-float16-float16]
+test/unit/language/test_core.py::test_dot3d[1-4-32-32-32-32-32-float16-float32]
+test/unit/language/test_core.py::test_dot3d[1-4-32-32-32-32-32-float32-float32]
+test/unit/language/test_core.py::test_dot3d[1-4-32-32-32-32-32-int8-int8]
+test/unit/language/test_core.py::test_dot3d[1-4-64-64-64-32-32-float16-float16]
+test/unit/language/test_core.py::test_dot3d[1-4-64-64-64-32-32-float16-float32]
+test/unit/language/test_core.py::test_dot3d[1-4-64-64-64-32-32-float32-float32]
+test/unit/language/test_core.py::test_dot3d[1-4-64-64-64-32-32-int8-int8]
+test/unit/language/test_core.py::test_dot3d[1-8-32-32-32-32-32-float16-float16]
+test/unit/language/test_core.py::test_dot3d[1-8-32-32-32-32-32-float16-float32]
+test/unit/language/test_core.py::test_dot3d[1-8-32-32-32-32-32-float32-float32]
+test/unit/language/test_core.py::test_dot3d[1-8-32-32-32-32-32-int8-int8]
+test/unit/language/test_core.py::test_dot3d[1-8-64-64-64-32-32-float16-float16]
+test/unit/language/test_core.py::test_dot3d[1-8-64-64-64-32-32-float16-float32]
+test/unit/language/test_core.py::test_dot3d[1-8-64-64-64-32-32-float32-float32]
+test/unit/language/test_core.py::test_dot3d[1-8-64-64-64-32-32-int8-int8]
+test/unit/language/test_core.py::test_dot3d[2-1-32-32-32-32-32-float16-float16]
+test/unit/language/test_core.py::test_dot3d[2-1-32-32-32-32-32-float16-float32]
+test/unit/language/test_core.py::test_dot3d[2-1-32-32-32-32-32-float32-float32]
+test/unit/language/test_core.py::test_dot3d[2-1-32-32-32-32-32-int8-int8]
+test/unit/language/test_core.py::test_dot3d[2-1-64-64-64-32-32-float16-float16]
+test/unit/language/test_core.py::test_dot3d[2-1-64-64-64-32-32-float16-float32]
+test/unit/language/test_core.py::test_dot3d[2-1-64-64-64-32-32-float32-float32]
+test/unit/language/test_core.py::test_dot3d[2-1-64-64-64-32-32-int8-int8]
+test/unit/language/test_core.py::test_dot3d[2-16-32-32-32-32-32-float16-float16]
+test/unit/language/test_core.py::test_dot3d[2-16-32-32-32-32-32-float16-float32]
+test/unit/language/test_core.py::test_dot3d[2-16-32-32-32-32-32-float32-float32]
+test/unit/language/test_core.py::test_dot3d[2-16-32-32-32-32-32-int8-int8]
+test/unit/language/test_core.py::test_dot3d[2-16-64-64-64-32-32-float16-float16]
+test/unit/language/test_core.py::test_dot3d[2-16-64-64-64-32-32-float16-float32]
+test/unit/language/test_core.py::test_dot3d[2-16-64-64-64-32-32-float32-float32]
+test/unit/language/test_core.py::test_dot3d[2-16-64-64-64-32-32-int8-int8]
+test/unit/language/test_core.py::test_dot3d[2-2-32-32-32-32-32-float16-float16]
+test/unit/language/test_core.py::test_dot3d[2-2-32-32-32-32-32-float16-float32]
+test/unit/language/test_core.py::test_dot3d[2-2-32-32-32-32-32-float32-float32]
+test/unit/language/test_core.py::test_dot3d[2-2-32-32-32-32-32-int8-int8]
+test/unit/language/test_core.py::test_dot3d[2-2-64-64-64-32-32-float16-float16]
+test/unit/language/test_core.py::test_dot3d[2-2-64-64-64-32-32-float16-float32]
+test/unit/language/test_core.py::test_dot3d[2-2-64-64-64-32-32-float32-float32]
+test/unit/language/test_core.py::test_dot3d[2-2-64-64-64-32-32-int8-int8]
+test/unit/language/test_core.py::test_dot3d[2-4-32-32-32-32-32-float16-float16]
+test/unit/language/test_core.py::test_dot3d[2-4-32-32-32-32-32-float16-float32]
+test/unit/language/test_core.py::test_dot3d[2-4-32-32-32-32-32-float32-float32]
+test/unit/language/test_core.py::test_dot3d[2-4-32-32-32-32-32-int8-int8]
+test/unit/language/test_core.py::test_dot3d[2-4-64-64-64-32-32-float16-float16]
+test/unit/language/test_core.py::test_dot3d[2-4-64-64-64-32-32-float16-float32]
+test/unit/language/test_core.py::test_dot3d[2-4-64-64-64-32-32-float32-float32]
+test/unit/language/test_core.py::test_dot3d[2-4-64-64-64-32-32-int8-int8]
+test/unit/language/test_core.py::test_dot3d[2-8-32-32-32-32-32-float16-float16]
+test/unit/language/test_core.py::test_dot3d[2-8-32-32-32-32-32-float16-float32]
+test/unit/language/test_core.py::test_dot3d[2-8-32-32-32-32-32-float32-float32]
+test/unit/language/test_core.py::test_dot3d[2-8-32-32-32-32-32-int8-int8]
+test/unit/language/test_core.py::test_dot3d[2-8-64-64-64-32-32-float16-float16]
+test/unit/language/test_core.py::test_dot3d[2-8-64-64-64-32-32-float16-float32]
+test/unit/language/test_core.py::test_dot3d[2-8-64-64-64-32-32-float32-float32]
+test/unit/language/test_core.py::test_dot3d[2-8-64-64-64-32-32-int8-int8]
+test/unit/language/test_core.py::test_dot3d[4-1-32-32-32-32-32-float16-float16]
+test/unit/language/test_core.py::test_dot3d[4-1-32-32-32-32-32-float16-float32]
+test/unit/language/test_core.py::test_dot3d[4-1-32-32-32-32-32-float32-float32]
+test/unit/language/test_core.py::test_dot3d[4-1-32-32-32-32-32-int8-int8]
+test/unit/language/test_core.py::test_dot3d[4-1-64-64-64-32-32-float16-float16]
+test/unit/language/test_core.py::test_dot3d[4-1-64-64-64-32-32-float16-float32]
+test/unit/language/test_core.py::test_dot3d[4-1-64-64-64-32-32-float32-float32]
+test/unit/language/test_core.py::test_dot3d[4-1-64-64-64-32-32-int8-int8]
+test/unit/language/test_core.py::test_dot3d[4-16-32-32-32-32-32-float16-float16]
+test/unit/language/test_core.py::test_dot3d[4-16-32-32-32-32-32-float16-float32]
+test/unit/language/test_core.py::test_dot3d[4-16-32-32-32-32-32-float32-float32]
+test/unit/language/test_core.py::test_dot3d[4-16-32-32-32-32-32-int8-int8]
+test/unit/language/test_core.py::test_dot3d[4-16-64-64-64-32-32-float16-float16]
+test/unit/language/test_core.py::test_dot3d[4-16-64-64-64-32-32-float16-float32]
+test/unit/language/test_core.py::test_dot3d[4-16-64-64-64-32-32-float32-float32]
+test/unit/language/test_core.py::test_dot3d[4-16-64-64-64-32-32-int8-int8]
+test/unit/language/test_core.py::test_dot3d[4-2-32-32-32-32-32-float16-float16]
+test/unit/language/test_core.py::test_dot3d[4-2-32-32-32-32-32-float16-float32]
+test/unit/language/test_core.py::test_dot3d[4-2-32-32-32-32-32-float32-float32]
+test/unit/language/test_core.py::test_dot3d[4-2-32-32-32-32-32-int8-int8]
+test/unit/language/test_core.py::test_dot3d[4-2-64-64-64-32-32-float16-float16]
+test/unit/language/test_core.py::test_dot3d[4-2-64-64-64-32-32-float16-float32]
+test/unit/language/test_core.py::test_dot3d[4-2-64-64-64-32-32-float32-float32]
+test/unit/language/test_core.py::test_dot3d[4-2-64-64-64-32-32-int8-int8]
+test/unit/language/test_core.py::test_dot3d[4-4-128-128-64-64-64-float16-float16]
+test/unit/language/test_core.py::test_dot3d[4-4-32-32-32-32-32-float16-float16]
+test/unit/language/test_core.py::test_dot3d[4-4-32-32-32-32-32-float16-float32]
+test/unit/language/test_core.py::test_dot3d[4-4-32-32-32-32-32-float32-float32]
+test/unit/language/test_core.py::test_dot3d[4-4-32-32-32-32-32-int8-int8]
+test/unit/language/test_core.py::test_dot3d[4-4-64-64-64-32-32-float16-float16]
+test/unit/language/test_core.py::test_dot3d[4-4-64-64-64-32-32-float16-float32]
+test/unit/language/test_core.py::test_dot3d[4-4-64-64-64-32-32-float32-float32]
+test/unit/language/test_core.py::test_dot3d[4-4-64-64-64-32-32-int8-int8]
+test/unit/language/test_core.py::test_dot3d[4-8-32-32-32-32-32-float16-float16]
+test/unit/language/test_core.py::test_dot3d[4-8-32-32-32-32-32-float16-float32]
+test/unit/language/test_core.py::test_dot3d[4-8-32-32-32-32-32-float32-float32]
+test/unit/language/test_core.py::test_dot3d[4-8-32-32-32-32-32-int8-int8]
+test/unit/language/test_core.py::test_dot3d[4-8-64-64-64-32-32-float16-float16]
+test/unit/language/test_core.py::test_dot3d[4-8-64-64-64-32-32-float16-float32]
+test/unit/language/test_core.py::test_dot3d[4-8-64-64-64-32-32-float32-float32]
+test/unit/language/test_core.py::test_dot3d[4-8-64-64-64-32-32-int8-int8]
+test/unit/language/test_core.py::test_dot3d[8-1-32-32-32-32-32-float16-float16]
+test/unit/language/test_core.py::test_dot3d[8-1-32-32-32-32-32-float16-float32]
+test/unit/language/test_core.py::test_dot3d[8-1-32-32-32-32-32-float32-float32]
+test/unit/language/test_core.py::test_dot3d[8-1-32-32-32-32-32-int8-int8]
+test/unit/language/test_core.py::test_dot3d[8-1-64-64-64-32-32-float16-float16]
+test/unit/language/test_core.py::test_dot3d[8-1-64-64-64-32-32-float16-float32]
+test/unit/language/test_core.py::test_dot3d[8-1-64-64-64-32-32-float32-float32]
+test/unit/language/test_core.py::test_dot3d[8-1-64-64-64-32-32-int8-int8]
+test/unit/language/test_core.py::test_dot3d[8-16-32-32-32-32-32-float16-float16]
+test/unit/language/test_core.py::test_dot3d[8-16-32-32-32-32-32-float16-float32]
+test/unit/language/test_core.py::test_dot3d[8-16-32-32-32-32-32-float32-float32]
+test/unit/language/test_core.py::test_dot3d[8-16-32-32-32-32-32-int8-int8]
+test/unit/language/test_core.py::test_dot3d[8-16-64-64-64-32-32-float16-float16]
+test/unit/language/test_core.py::test_dot3d[8-16-64-64-64-32-32-float16-float32]
+test/unit/language/test_core.py::test_dot3d[8-16-64-64-64-32-32-float32-float32]
+test/unit/language/test_core.py::test_dot3d[8-16-64-64-64-32-32-int8-int8]
+test/unit/language/test_core.py::test_dot3d[8-2-32-32-32-32-32-float16-float16]
+test/unit/language/test_core.py::test_dot3d[8-2-32-32-32-32-32-float16-float32]
+test/unit/language/test_core.py::test_dot3d[8-2-32-32-32-32-32-float32-float32]
+test/unit/language/test_core.py::test_dot3d[8-2-32-32-32-32-32-int8-int8]
+test/unit/language/test_core.py::test_dot3d[8-2-64-64-64-32-32-float16-float16]
+test/unit/language/test_core.py::test_dot3d[8-2-64-64-64-32-32-float16-float32]
+test/unit/language/test_core.py::test_dot3d[8-2-64-64-64-32-32-float32-float32]
+test/unit/language/test_core.py::test_dot3d[8-2-64-64-64-32-32-int8-int8]
+test/unit/language/test_core.py::test_dot3d[8-4-32-32-32-32-32-float16-float16]
+test/unit/language/test_core.py::test_dot3d[8-4-32-32-32-32-32-float16-float32]
+test/unit/language/test_core.py::test_dot3d[8-4-32-32-32-32-32-float32-float32]
+test/unit/language/test_core.py::test_dot3d[8-4-32-32-32-32-32-int8-int8]
+test/unit/language/test_core.py::test_dot3d[8-4-64-64-64-32-32-float16-float16]
+test/unit/language/test_core.py::test_dot3d[8-4-64-64-64-32-32-float16-float32]
+test/unit/language/test_core.py::test_dot3d[8-4-64-64-64-32-32-float32-float32]
+test/unit/language/test_core.py::test_dot3d[8-4-64-64-64-32-32-int8-int8]
+test/unit/language/test_core.py::test_dot3d[8-8-32-32-32-32-32-float16-float16]
+test/unit/language/test_core.py::test_dot3d[8-8-32-32-32-32-32-float16-float32]
+test/unit/language/test_core.py::test_dot3d[8-8-32-32-32-32-32-float32-float32]
+test/unit/language/test_core.py::test_dot3d[8-8-32-32-32-32-32-int8-int8]
+test/unit/language/test_core.py::test_dot3d[8-8-64-64-64-32-32-float16-float16]
+test/unit/language/test_core.py::test_dot3d[8-8-64-64-64-32-32-float16-float32]
+test/unit/language/test_core.py::test_dot3d[8-8-64-64-64-32-32-float32-float32]
+test/unit/language/test_core.py::test_dot3d[8-8-64-64-64-32-32-int8-int8]
 # https://github.com/intel/intel-xpu-backend-for-triton/issues/983
 test/unit/language/test_core.py::test_noinline[shared]
 test/unit/language/test_core.py::test_dot[1-128-128-64-2-True-True-none-tf32-int8-int8-1_0]