From 78e584ffb92cd72f0697d59012a2e4907412cc64 Mon Sep 17 00:00:00 2001 From: Bruno Volpato Date: Mon, 2 Nov 2026 19:55:00 -0500 Subject: [PATCH] [PyTorch] Keep Flex Attention score_mod graphs on the cuDNN backend Only the cuDNN backend implements score_mod. With cuDNN frontend 1.31.0 to 1.32.1 from PyPI, the frontend does not see the callback on sdpa(), so its Python SDPA engine can serve the graph as plain attention. The score modifier is dropped without an error, which gives wrong results. Exclude the Python SDPA engines from these graphs. Raise an error if a Python engine is still selected, so that a later frontend cannot silently drop the modifier again. Signed-off-by: Bruno Volpato --- .../attention/dot_product_attention/flex_attention.py | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/transformer_engine/pytorch/attention/dot_product_attention/flex_attention.py b/transformer_engine/pytorch/attention/dot_product_attention/flex_attention.py index 8ba48682b42..c122b8faa43 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/flex_attention.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/flex_attention.py @@ -270,12 +270,21 @@ def _finalize_cudnn_graph(graph) -> int: graph.validate() graph.build_operation_graph() + # Only the cuDNN backend implements score_mod. cuDNN frontend 1.31 can offer its Python SDPA + # engines for these graphs, and those engines run plain attention without the modifier. + graph.deselect_engines(["sdpa_fwd", "sdpa_bwd"]) try: graph.create_execution_plans([cudnn.heur_mode.A, cudnn.heur_mode.FALLBACK]) graph.check_support() except cudnn.cudnnGraphNotSupportedError as exc: raise RuntimeError(f"cuDNN Flex Attention SDPA graph is not supported: {exc}") from exc graph.build_plans(cudnn.build_plan_policy.HEURISTICS_CHOICE) + selected_engine = getattr(graph, "selected_engine", None) + if selected_engine is not None: + raise RuntimeError( + f"cuDNN frontend {cudnn.__version__} selected {selected_engine} for a score_mod graph. " + "Only the cuDNN backend implements score_mod." + ) return max(graph.get_workspace_size(), 1)