Remove advanced path (#5205)

Advanced path is no longer functional for flash attention.

Closes #3286

Signed-off-by: Whitney Tsang <[email protected]>

Author: whitneywhtsang

.github/workflows/build-test-reusable.yml

@@ -342,7 +342,6 @@ jobs:
           07-extern-functions
           08-grouped-gemm
           10-experimental-block-pointer
-          10i-experimental-block-pointer
           EOF
 
       - name: Run Tutorials

benchmarks/triton_kernels_benchmark/flash_attention_benchmark.py

@@ -440,49 +440,25 @@ def forward(ctx, q, k, v, causal, sm_scale):
         assert Lq == Lk and Lk == Lv
         assert Lk in {16, 32, 64, 128}
         o = torch.empty_like(q)
-        BLOCK_M = 128
-        BLOCK_N = 64
-        num_stages = 3
-        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']), 1, q.shape[0] * q.shape[1])
         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':
-            # default pipeline
-            _attention.tune_attn_fwd[grid](  # pylint: disable=unsubscriptable-object
-                q, k, v, sm_scale, M, o,  #
-                q.stride(0), q.stride(1), q.stride(2), q.stride(3),  #
-                k.stride(0), k.stride(1), k.stride(2), k.stride(3),  #
-                v.stride(0), v.stride(1), v.stride(2), v.stride(3),  #
-                o.stride(0), o.stride(1), o.stride(2), o.stride(3),  #
-                q.shape[0], q.shape[1],  #
-                N_CTX=q.shape[2],  #
-                BLOCK_DMODEL=Lk,  #
-                STAGE=stage,  #
-                split_barriers_scope='None',  # possible scope value: 'Subgroup','Workgroup'
-            )
-        else:
-            _attention.attn_fwd[grid](  # pylint: disable=unsubscriptable-object
-                q, k, v, sm_scale, M, o,  #
-                q.stride(0), q.stride(1), q.stride(2), q.stride(3),  #
-                k.stride(0), k.stride(1), k.stride(2), k.stride(3),  #
-                v.stride(0), v.stride(1), v.stride(2), v.stride(3),  #
-                o.stride(0), o.stride(1), o.stride(2), o.stride(3),  #
-                q.shape[0], q.shape[1],  #
-                N_CTX=q.shape[2],  #
-                BLOCK_M=BLOCK_M,  #
-                BLOCK_N=BLOCK_N,  #
-                BLOCK_DMODEL=Lk,  #
-                STAGE=stage,  #
-                num_warps=num_warps,  #
-                num_stages=num_stages,  #
-                grf_mode='large',  #
-                advanced_path=True,  #
-            )
+        _attention.tune_attn_fwd[grid](  # pylint: disable=unsubscriptable-object
+            q, k, v, sm_scale, M, o,  #
+            q.stride(0), q.stride(1), q.stride(2), q.stride(3),  #
+            k.stride(0), k.stride(1), k.stride(2), k.stride(3),  #
+            v.stride(0), v.stride(1), v.stride(2), v.stride(3),  #
+            o.stride(0), o.stride(1), o.stride(2), o.stride(3),  #
+            q.shape[0], q.shape[1],  #
+            N_CTX=q.shape[2],  #
+            BLOCK_DMODEL=Lk,  #
+            STAGE=stage,  #
+            split_barriers_scope='None',  # possible scope value: 'Subgroup','Workgroup'
+        )
 
         ctx.save_for_backward(q, k, v, o, M)
         ctx.sm_scale = sm_scale

bin/RegisterTritonDialects.h

@@ -8,7 +8,6 @@
 #include "intel/include/TritonGENToLLVM/Passes.h"
 #include "intel/include/TritonGENToSPIRV/Passes.h"
 #include "intel/include/TritonIntelGPUToLLVM/Passes.h"
-#include "intel/include/TritonToTritonGPUWarp/Passes.h"
 
 #include "amd/include/Dialect/TritonAMDGPU/IR/Dialect.h"
 #include "amd/include/TritonAMDGPUTransforms/Passes.h"
@@ -89,7 +88,6 @@ inline void registerTritonDialects(mlir::DialectRegistry &registry) {
   mlir::test::registerTestAMDGPUMembarPass();
   mlir::test::registerTestTritonAMDGPURangeAnalysis();
   mlir::triton::registerConvertTritonToTritonGPUPass();
-  mlir::triton::intel::registerConvertTritonToTritonGPUWarpPass();
   mlir::triton::intel::registerTritonIntelTensorDescToBlockPointer();
   mlir::triton::intel::registerTritonIntelRemoveMasks();
   mlir::triton::registerRelayoutTritonGPUPass();

include/triton/Tools/Sys/GetEnv.hpp

@@ -44,17 +44,12 @@ inline const std::set<std::string> CACHE_INVALIDATING_ENV_VARS = {
     "ALLOW_LHS_TMEM_LAYOUT_CONVERSION",
     "TRITON_F32_DEFAULT",
     "TRITON_PREFER_TMEM_16x256_LAYOUT",
-    "TRITON_INTEL_ADVANCED_PATH",
     "TRITON_INTEL_AGGRESSIVE_DPAS_REUSE",
-    "TRITON_INTEL_DO_NOT_SINK_INSTR_ACROSS_RGN",
     "TRITON_INTEL_ENABLE_BLOCK_IO_ALL_LAYOUTS",
     "TRITON_INTEL_ENABLE_DPAS_FOR_WARP_SIZE_32",
-    "TRITON_INTEL_ENABLE_FIRST_LOAD_TO_SLM",
-    "TRITON_INTEL_ENABLE_INSTR_SCHED",
     "TRITON_INTEL_FAST_MATH",
     "TRITON_INTEL_ONE_MATRIX_PER_LOAD_BT",
     "TRITON_INTEL_PREDICATED",
-    "TRITON_INTEL_REDUCE_TRANSPOSE",
     // clang-format on
 };
 

lib/Analysis/Utility.cpp

@@ -105,11 +105,6 @@ unsigned ReduceOpHelper::getIntraWarpSizeWithUniqueData() {
 }
 
 bool ReduceOpHelper::isWarpSynchronous() {
-  // FIXME: In the default path tensors will always have a layout. Tensors do
-  // not have a layout only in the advanced path. We need to find a workaround
-  // in order to remove this change.
-  if (!srcEncoding)
-    return true;
   return getWarpsPerCTA(srcEncoding, srcShape)[axis] == 1;
 }
 

lib/Dialect/TritonGPU/IR/Dialect.cpp

@@ -1426,8 +1426,6 @@ void SliceEncodingAttr::print(mlir::AsmPrinter &printer) const {
 LogicalResult
 SliceEncodingAttr::verify(function_ref<InFlightDiagnostic()> emitError,
                           unsigned dim, DistributedEncodingTrait parent) {
-  if (mlir::triton::tools::getBoolEnv("TRITON_INTEL_ADVANCED_PATH"))
-    return success();
   unsigned rank = cast<LayoutEncodingTrait>(parent).getRank();
   if (rank <= 1)
     return emitError() << "parent layout must have at least rank >= 2";
@@ -2558,13 +2556,6 @@ LogicalResult DotOperandEncodingAttr::verify(
     return success();
   }
 
-  if (auto parentAttr = mlir::dyn_cast<intel::WarpEncodingAttr>(parent)) {
-    if (kWidth != 0)
-      return emitError() << "ttg.dot_op kWidth parameter is not supported "
-                            "when the parent is a warp layout";
-    return success();
-  }
-
   if (auto parentAttr = mlir::dyn_cast<BlockedEncodingAttr>(parent)) {
     if (kWidth != 0)
       return emitError() << "ttg.dot_op kWidth parameter is not supported "
@@ -2597,9 +2588,6 @@ class TritonGPUOpAsmInterface : public OpAsmDialectInterface {
     } else if (auto linearAttr = mlir::dyn_cast<LinearEncodingAttr>(attr)) {
       os << "linear";
       return AliasResult::FinalAlias;
-    } else if (auto warpAttr = mlir::dyn_cast<intel::WarpEncodingAttr>(attr)) {
-      os << "warp";
-      return AliasResult::FinalAlias;
     } /* else if (auto sliceAttr = dyn_cast<SliceEncodingAttr>(attr)) {
       os << "slice";
       return AliasResult::FinalAlias;
@@ -3298,8 +3286,6 @@ struct TritonGPUVerifyTensorLayoutInterface
     if (!distr)
       return makeErr()
              << "Non-distributed layout is not allowed in tensor type.";
-    if (mlir::triton::tools::getBoolEnv("TRITON_INTEL_ADVANCED_PATH"))
-      return success();
     auto rank = distr.getRepOrder().size();
     if (rank != rankedTy.getRank())
       return makeErr() << "Layout has rank " << rank

python/triton/knobs.py

@@ -548,9 +548,7 @@ class intel_knobs(base_knobs):
     dump_shader_info: env_bool = env_bool("TRITON_INTEL_ENABLE_IGC_SHADER_DUMP", False)
     gen_native_code: env_bool = env_bool("TRITON_XPU_GEN_NATIVE_CODE", False)
     tile_load_ll: env_bool = env_bool("TRITON_XPU_ENABLE_TILE_LOAD_LINEAR_LAYOUT", True)
-    advanced_path: env_bool = env_bool("TRITON_INTEL_ADVANCED_PATH", False)
     opt_reduction_locality: env_bool = env_bool("TRITON_INTEL_OPTIMIZE_REDUCTION_LOCALITY", False)
-    reduce_transpose: env_bool = env_bool("TRITON_INTEL_REDUCE_TRANSPOSE", False)
     disable_igc_opt: env_bool = env_bool("TRITON_INTEL_DISABLE_IGC_OPT", False)
 
     dump_spirv_kernel_args: env_opt_str = env_opt_str("TRITON_XPU_DUMP_SPIRV_KERNEL_ARGS")

python/tutorials/10-experimental-block-pointer.py

@@ -90,15 +90,11 @@
 # Final Result
 # ------------
 
-import os
-
 import torch
 
 import triton
 import triton.language as tl
 
-SMALL_GRF = os.getenv('TRITON_INTEL_ADVANCED_PATH', '0') == '0'
-
 
 @triton.autotune(
     configs=[
@@ -107,18 +103,14 @@
             num_stages=s, num_warps=32) for s in [1, 2, 3]
     ] + [
         triton.Config({'BLOCK_SIZE_M': 256, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 32, 'GROUP_SIZE_M': 4, 'grf_mode': m},
-                      num_stages=s, num_warps=w)
-        for s in [2, 3, 4]
-        for (m, w) in ([('large', 32), ('small', 64)] if SMALL_GRF else [('large', 32)])
+                      num_stages=s, num_warps=w) for s in [2, 3, 4] for (m, w) in ([('large', 32), ('small', 64)])
     ] + [
         triton.Config(
             {'BLOCK_SIZE_M': 64, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 32, 'GROUP_SIZE_M': 4, 'grf_mode': 'large'},
             num_stages=s, num_warps=32) for s in [2]
     ] + [
         triton.Config({'BLOCK_SIZE_M': 8, 'BLOCK_SIZE_N': 512, 'BLOCK_SIZE_K': 64, 'GROUP_SIZE_M': 1, 'grf_mode': m},
-                      num_stages=s, num_warps=w)
-        for s in [2, 3]
-        for (m, w) in ([('large', 32), ('small', 64)] if SMALL_GRF else [('large', 32)])
+                      num_stages=s, num_warps=w) for s in [2, 3] for (m, w) in ([('large', 32), ('small', 64)])
     ],
     key=['M', 'N', 'K'],
 )
@@ -349,9 +341,7 @@ def matmul(a, b, accum_dtype, res_dtype):
 FP8_TYPES = [(torch.float8_e4m3fn, torch.float32, torch.float16)]
 
 torch.manual_seed(0)
-for dtype, accum_dtype, res_dtype in FP16_TYPES + FP32_TYPES + INT8_TYPES + (FP8_TYPES if os.getenv(
-        'TRITON_INTEL_ADVANCED_PATH', '0') == '0' else []):
-
+for dtype, accum_dtype, res_dtype in FP16_TYPES + FP32_TYPES + INT8_TYPES + FP8_TYPES:
     for shape in [(512, 512), (4, 512, 512)]:
         assert shape[-1] == shape[-2], "Only square matrices are supported"
         if dtype.is_floating_point:

python/tutorials/10i-experimental-block-pointer.py

@@ -58,9 +58,6 @@ def run_tutorial(path: pathlib.Path) -> float:
     if not spec or not spec.loader:
         raise AssertionError(f'Failed to load module from {path}')
     module = importlib.util.module_from_spec(spec)
-    # Set __file__ to the absolute name, a workaround for 10i-experimental-block-pointer, which
-    # uses dirname of its location to find 10-experimental-block-pointer.
-    module.__file__ = path.resolve().as_posix()
     # Reset sys.argv because some tutorials, such as 09, parse their command line arguments.
     sys.argv = [str(path)]
     start_time = datetime.datetime.now()