Fix flex_attention test

Differential Revision: D81257355

Pull Request resolved: #373

Author: xuzhao9

test/test_gpu/main.py

@@ -89,7 +89,8 @@ def _run_one_operator(args: List[str]):
         del op
         tb_args.mode = "bwd"
         if tb_args.op in BWD_ARGS_OPS:
-            extra_args.extend(BWD_ARGS_OPS[tb_args.op])
+            args.extend(BWD_ARGS_OPS[tb_args.op])
+            tb_args, extra_args = parser.parse_known_args(args)
         op = Operator(tb_args=tb_args, extra_args=extra_args)
         op.run()
         check_ci_output(op)

tritonbench/operators/flex_attention/operator.py

@@ -415,8 +415,7 @@ def sdpa_fn():
                 with sdpa_kernel([SDPBackend.CUDNN_ATTENTION]):
                     return sdpa(q, k, v, is_causal=is_causal)
             except RuntimeError as e:
-                print(f"[SKIP] cuDNN backend failed: {e}")
-                return None
+                raise NotImplementedError(str(e))
 
         return sdpa_fn
 

tritonbench/utils/triton_op.py

@@ -1507,8 +1507,8 @@ def _init_extra_metrics() -> Dict[str, Any]:
                 self.dump_ir(input_id, fn)
         except torch.cuda.OutOfMemoryError:
             metrics.error_msg = "CUDA OOM"
-        except NotImplementedError:
-            metrics.error_msg = "not supported"
+        except NotImplementedError as e:
+            metrics.error_msg = str(e)
         except Exception as e:
             if not self.tb_args.keep_going:
                 raise