[PyTorch] Support for cuDNN-backed flex attention#2984
Conversation
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
for more information, see https://pre-commit.ci
Greptile SummaryThis PR introduces a Python-only cuDNN frontend code path for
Confidence Score: 4/5The implementation is largely correct, but two of the four new integration tests will fail at runtime on any CUDA machine because score_mod tensors are created on CPU and passed into a cuDNN CUDA graph execution. The causal and softcap test cases in test_flex_attention.py pass CPU tensors (torch.full without a device argument) as score_mod_tensors to the cuDNN graph execution path. cuDNN CUDA kernels require all variant-pack tensors to reside on the compute device; the CPU tensors will trigger a device-mismatch error, causing those test cases to fail on any real CUDA runner. The core implementation in flex_attention.py is sound. tests/pytorch/attention/test_flex_attention.py — the causal and softcap score_mod tensor construction. Important Files Changed
Sequence DiagramsequenceDiagram
participant User
participant DPA as DotProductAttention
participant FA as FusedAttention
participant Func as FusedAttentionWithScoreModFunc
participant Cache as _cudnn_score_mod_graph_cache
participant cuDNN as cuDNN Frontend
User->>DPA: forward(q,k,v, score_mod, score_mod_tensors, ...)
DPA->>DPA: validate score_mod inputs
DPA->>DPA: get_attention_backend (has_score_mod filter)
DPA->>FA: forward(score_mod, score_mod_bprop, ...)
FA->>Func: apply(is_training, q, k, v, score_mod, ...)
Func->>Cache: _get_cudnn_score_mod_fwd_graph(key)
alt cache miss
Cache->>cuDNN: _build_cudnn_score_mod_fwd_graph()
cuDNN-->>Cache: CudnnScoreModFwdGraphEntry
Cache->>Cache: store entry
end
Cache-->>Func: entry (graph + tensor handles)
Func->>cuDNN: _execute_cudnn_graph(variant_pack)
cuDNN-->>Func: output tensor
Func-->>DPA: output
User->>Func: backward(d_out)
Func->>Cache: _get_cudnn_score_mod_bwd_graph(key)
alt cache miss
Cache->>cuDNN: _build_cudnn_score_mod_bwd_graph()
cuDNN-->>Cache: CudnnScoreModBwdGraphEntry
end
Func->>cuDNN: _execute_cudnn_graph(variant_pack)
cuDNN-->>Func: dq, dk, dv
Func-->>User: dq, dk, dv
Reviews (4): Last reviewed commit: "[pre-commit.ci] auto fixes from pre-comm..." | Re-trigger Greptile |
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
| def _score_mod_callback_cache_key(callback: Optional[Callable]) -> Optional[Tuple[Any, ...]]: | ||
| """Create a stable cache key for a score_mod callable.""" | ||
| if callback is None: | ||
| return None | ||
| self_obj = getattr(callback, "__self__", None) | ||
| func_obj = getattr(callback, "__func__", None) | ||
| if self_obj is not None and func_obj is not None: | ||
| return ("bound_method", id(self_obj), id(func_obj)) | ||
| return ("callable", id(callback)) |
There was a problem hiding this comment.
id()-based cache key is unsafe for parameterized bound-method score_mods
id(self_obj) identifies a Python object by its memory address. When a bound-method instance is garbage-collected, Python may immediately reuse that memory for a new instance. If the new instance belongs to the same class (same id(func_obj)), the cache key is identical, so _get_cudnn_score_mod_fwd_graph returns the old compiled graph even though the new instance might construct a structurally different computation — e.g., a score_mod class whose forward loops self.n_layers times. The wrong graph is executed without any error, silently producing incorrect attention outputs.
For stateless module-level functions this is fine (they're never GC'd), but any stateful class-based score_mod where different instances produce different graph topologies can hit this bug in long-running programs. Consider using type(self_obj) and a per-class sequence counter, or requiring callers to provide an explicit cache key.
| _flash_attn_varlen_fwd = None | ||
| _flash_attn_varlen_bwd = None |
There was a problem hiding this comment.
_cudnn_score_mod_graph_cache is a plain dict with no eviction policy. Cache keys encode tensor shapes, strides, dtype, and device, so every new (batch, seq, heads, dim) combination — extremely common in training with variable-length sequences or multi-task workloads — inserts a permanent entry. Each cached cuDNN graph holds compiled CUDA kernels and associated state, which can be several tens of MB. Over a long training run this will silently consume increasing GPU/CPU memory. Consider a bounded LRU cache (e.g., functools.lru_cache or a collections.OrderedDict with a size cap).
| fused_attention_backend = tex.get_fused_attn_backend( | ||
| self.training, | ||
| q_type, | ||
| q_type, | ||
| dpa_utils.QKVLayout["bshd_bshd_bshd"], | ||
| dpa_utils.AttnBiasType["no_bias"], | ||
| dpa_utils.AttnMaskType["no_mask"], | ||
| dpa_utils.SoftmaxType["vanilla"], |
There was a problem hiding this comment.
get_fused_attn_backend availability check always uses bshd_bshd_bshd regardless of actual format
The score_mod path hard-codes dpa_utils.QKVLayout["bshd_bshd_bshd"] for the backend probe, even when the user passes qkv_format="sbhd". The result is only used to gate on NVTE_No_Backend, so in practice it likely works today because backend availability for a given dtype is layout-independent. However, if a future cuDNN version makes SBHD/BSHD support diverge, this probe would give a false-positive (accepts sbhd even though no backend supports it) or false-negative (rejects sbhd when it is actually supported). Using the real layout for the probe would make the check self-documenting and future-proof.
| ) | ||
|
|
||
| if context_parallel: | ||
| if score_mod is not None: |
There was a problem hiding this comment.
I think this should be in the else branch, because it doesn't support context parallelism. Something like this:
if context_parallel: elif score_mod is not None: else:
| score_mod: Optional[Callable] = None, | ||
| score_mod_bprop: Optional[Callable] = None, | ||
| score_mod_tensors: Optional[Dict[str, torch.Tensor]] = None, | ||
| score_mod_bprop_tensors: Optional[Dict[str, torch.Tensor]] = None, |
There was a problem hiding this comment.
Do you think it'd be clearer if we add "fprop" to the names? i.e. score_mod_fprop, score_mod_bprop, score_mod_fprop_tensors, score_mod_bprop_tensors?
| isinstance(k, str) and isinstance(v, torch.Tensor) | ||
| for k, v in score_mod_bprop_tensors.items() | ||
| ), "score_mod_bprop_tensors must map string names to torch.Tensor instances!" | ||
|
|
There was a problem hiding this comment.
I think all these checks can go into dpa_utils.get_attention_backend(), and with score_mod_xxx args passed in (to AttentionParams), that utility function can return use_fused_attention=False if one of the checks if violated. dpa_utils.get_attention_backend() is used in the tests as well (by get_available_attention_backends()).
| raise ValueError( | ||
| "score_mod requires a cuDNN FusedAttention backend, but no fused " | ||
| "attention backend supports the provided inputs." | ||
| ) |
There was a problem hiding this comment.
For the score_mod path, I don't think we need to call tex.get_fused_attn_backend() and check if it's supported or not. If anything, we should add graph.validate() -> .... graph.build_plans() to dpa_utils.get_attention_backend(attention_params), but if that's too heavy-handed, we can only do the checks you had above (the asserts). Once those checks were added to dpa_utils.get_attention_backend, whether FusedAttention backend is run or not will be controlled by the following logic (just like with non-score_mod cases):
(
use_flash_attention,
flash_attention_backend,
use_fused_attention,
fused_attention_backend,
use_unfused_attention,
_,
) = dpa_utils.get_attention_backend(attention_params)
| else: | ||
| pad_between_seqs = False | ||
|
|
||
| if score_mod is None: |
There was a problem hiding this comment.
Please label this "experimental".
There was a problem hiding this comment.
Agreed.
Also please add a comment on top (similar to other sections in forward().
Something like : checking compatibility for flex attn
|
|
||
| def _build_cudnn_pygraph(dtype: torch.dtype, device: torch.device): | ||
| """Create a cuDNN frontend Python graph for F16/BF16 SDPA.""" | ||
| import cudnn # pylint: disable=import-outside-toplevel |
There was a problem hiding this comment.
Can you import the cudnn from 3rdparty/cudnn-frontend, instead of from the environment/system-wide installation? We have control over the version in 3rdparty/cudnn-frontend, but not the system one.
| @pytest.mark.parametrize("dtype", param_types) | ||
| @pytest.mark.parametrize("qkv_format", ["sbhd", "bshd"]) | ||
| @pytest.mark.parametrize("scalar_loss", [False, True]) | ||
| def test_dot_product_attention_score_mod(dtype, qkv_format, scalar_loss): |
There was a problem hiding this comment.
Would @pytest.mark.parameterize("score_mod", ["causal", "softcap", "post_scale_bias"]) simplify the tests a bit, so that we don't have 3 separate tests with a lot of repeated code?
| score_mod: Callable, | ||
| score_mod_tensors: Optional[Dict[str, torch.Tensor]], | ||
| output_layer: torch.Tensor, | ||
| stats_bhs1: Optional[torch.Tensor], |
There was a problem hiding this comment.
I think we can just call this stats, even though it might only support bhs1 shape right now. On the C++ side, cuDNN does support th1 (for THD format) as well. Could we leave the name generic for now in case we want to add more support to it in the future?
| return output.contiguous() | ||
|
|
||
|
|
||
| def _bhsd_dim_stride( |
There was a problem hiding this comment.
We have a lot of small utility functions here - is there a way to pack them up a bit or group them in some way, so the code is easier to read? I know this is Python and we probably do need more than 2 functions (fwd+bwd) but could you please have a look into this? Thanks.
There was a problem hiding this comment.
I agree with this and was my first thought too.
We should club these function into a couple classes that can sit in this file at the very least.
However, I think this approach is still not the right approach. We should have a separate flex_attention.py file similar to context_parallel.py and backends.py can import it similar to how it imports the CP functions right now.
I strongly recommend this for two reasons :
- When we refactored attention as a whole early last year, the idea was to modularize attention. That was the reason CP was moved out of attention. With Flex attention's functionality and code in here being fairly decoupled from vanilla DPA, it should be easier to move it out. Leaving this code in here would add ~1000 lines of code that is not related to the vanilla DPA and would practically be undoing the refactoring work we did early last year. The same reason for moving CP to it's own file should also apply to Flex attention.
- A developer/user of TE PyT DPA should not have to worry about the details of flex attention. Similarly someones modifying flex should not be bogged down by the details of vanilla fused attn. Hence, decoupling is important to aid with debugging as well
| ) | ||
|
|
||
|
|
||
| def _score_mod_relative_position(score_mod_graph, score_tensor, _tensors): |
There was a problem hiding this comment.
We can just call this "post_scale_bias" to be consistent with our nomenclature elsewhere.
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
for more information, see https://pre-commit.ci
| if ( | ||
| inspect.isfunction(callback) | ||
| and callback.__closure__ is None | ||
| and "<locals>" not in callback.__qualname__ | ||
| ): | ||
| return ("function", callback.__module__, callback.__qualname__) |
There was a problem hiding this comment.
__qualname__ = "<lambda>", so two different lambdas defined at module scope in the same file (e.g., sm1 = lambda g, s, t: s and sm2 = lambda g, s, t: g.neg(input=s)) would produce the identical cache key ("function", module, "<lambda>"). The second lambda would silently reuse the compiled graph from the first, computing wrong attention scores with no error. Named module-level functions are safe because their qualnames are unique, but lambdas are not. Excluding <lambda> from the cacheable path makes them _SCORE_MOD_UNCACHEABLE, which builds a fresh graph every call — the same safe fallback already used for closures and nested functions.
| if ( | |
| inspect.isfunction(callback) | |
| and callback.__closure__ is None | |
| and "<locals>" not in callback.__qualname__ | |
| ): | |
| return ("function", callback.__module__, callback.__qualname__) | |
| if ( | |
| inspect.isfunction(callback) | |
| and callback.__closure__ is None | |
| and "<locals>" not in callback.__qualname__ | |
| and "<lambda>" not in callback.__qualname__ | |
| ): | |
| return ("function", callback.__module__, callback.__qualname__) |
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
| score_mod_kwargs = { | ||
| "score_mod": _score_mod_causal, | ||
| "score_mod_bprop": _score_mod_causal_bprop, | ||
| "score_mod_tensors": {"neg_inf": torch.full((1, 1, 1, 1), -1e9)}, | ||
| "score_mod_bprop_tensors": {"zero": torch.full((1, 1, 1, 1), 0.0)}, | ||
| } |
There was a problem hiding this comment.
neg_inf and zero tensors are created on CPU (torch.full defaults to CPU), but the attention computation runs on CUDA. When cuDNN executes the graph it calls into CUDA kernels and expects all variant-pack tensors to reside on the compute device. Passing CPU tensors here will produce a device-mismatch error at graph execution time, causing both the "causal" test cases to fail.
| score_mod_kwargs = { | |
| "score_mod": _score_mod_causal, | |
| "score_mod_bprop": _score_mod_causal_bprop, | |
| "score_mod_tensors": {"neg_inf": torch.full((1, 1, 1, 1), -1e9)}, | |
| "score_mod_bprop_tensors": {"zero": torch.full((1, 1, 1, 1), 0.0)}, | |
| } | |
| score_mod_kwargs = { | |
| "score_mod": _score_mod_causal, | |
| "score_mod_bprop": _score_mod_causal_bprop, | |
| "score_mod_tensors": {"neg_inf": torch.full((1, 1, 1, 1), -1e9, device="cuda")}, | |
| "score_mod_bprop_tensors": {"zero": torch.full((1, 1, 1, 1), 0.0, device="cuda")}, | |
| } |
| score_mod : bool, default = False | ||
| Whether a score_mod callback was provided. | ||
| score_mod_bprop : bool, default = False | ||
| Whether a score_mod bprop callback was provided. |
There was a problem hiding this comment.
nit: If this is a bool, to match has_attention_mask, consider has_score_mod and has_score_mod_bprop instead ?
| logger.debug("Disabling all backends for max_logit with FP8 attention") | ||
|
|
||
| # Filter: score_mod | ||
| if score_mod_bprop and not score_mod: |
There was a problem hiding this comment.
What happens (is expected to happen) if score_mod_bprop=False and score_mod=True ?
There was a problem hiding this comment.
It's a perfectly legal case, if, for instance, score_mod is used only for masking.
| use_flash_attention = False | ||
| use_flash_attention_2 = False | ||
| use_flash_attention_3 = False | ||
| use_flash_attention_4 = False | ||
| use_fused_attention = False | ||
| use_unfused_attention = False |
There was a problem hiding this comment.
nit: Outside the scope of this PR but would be good to do in this or subsequent PR: having a function or something similar for when performing an action/query on all flash_attention vars
| if use_flash_attention_2 or use_flash_attention_3 or use_flash_attention_4: | ||
| logger.debug("Disabling FlashAttention for score_mod") | ||
| use_flash_attention = False | ||
| use_flash_attention_2 = False | ||
| use_flash_attention_3 = False | ||
| use_flash_attention_4 = False |
There was a problem hiding this comment.
Consider this maybe ?:
| if use_flash_attention_2 or use_flash_attention_3 or use_flash_attention_4: | |
| logger.debug("Disabling FlashAttention for score_mod") | |
| use_flash_attention = False | |
| use_flash_attention_2 = False | |
| use_flash_attention_3 = False | |
| use_flash_attention_4 = False | |
| if use_flash_attention_2 or use_flash_attention_3 or use_flash_attention_4: | |
| logger.debug("Disabling FlashAttention for score_mod") | |
| use_flash_attention = False | |
| use_flash_attention_2 = False | |
| use_flash_attention_3 = False | |
| use_flash_attention_4 = False |
unless there's a good reason to do otherwise ?
| if use_unfused_attention: | ||
| logger.debug("Disabling UnfusedDotProductAttention for score_mod") | ||
| use_unfused_attention = False |
There was a problem hiding this comment.
Consider this maybe ?
| if use_unfused_attention: | |
| logger.debug("Disabling UnfusedDotProductAttention for score_mod") | |
| use_unfused_attention = False | |
| if use_unfused_attention: | |
| logger.debug("Disabling UnfusedDotProductAttention for score_mod") | |
| use_unfused_attention = False |
unless there's a good reason to do otherwise ?
| if score_mod is None: | ||
| assert score_mod_bprop is None, "score_mod_bprop requires score_mod!" | ||
| assert score_mod_tensors is None, "score_mod_tensors requires score_mod!" | ||
| assert ( | ||
| score_mod_bprop_tensors is None | ||
| ), "score_mod_bprop_tensors requires score_mod!" | ||
| else: | ||
| assert callable(score_mod), "score_mod must be callable!" | ||
| assert score_mod_bprop is None or callable( | ||
| score_mod_bprop | ||
| ), "score_mod_bprop must be callable when provided!" | ||
| assert query_layer.dtype in [ | ||
| torch.float16, | ||
| torch.bfloat16, | ||
| ], "score_mod only supports FP16 and BF16 tensors!" | ||
| assert ( | ||
| key_layer.dtype == query_layer.dtype and value_layer.dtype == query_layer.dtype | ||
| ), "score_mod requires Q, K and V tensors to have the same dtype!" | ||
| assert ( | ||
| type(query_layer) is torch.Tensor | ||
| and type(key_layer) is torch.Tensor | ||
| and type(value_layer) is torch.Tensor | ||
| ), "score_mod only supports unquantized torch.Tensor Q, K and V inputs!" | ||
| assert not self.fp8, "score_mod is not supported with FP8 DotProductAttention!" | ||
| assert not fp8_output, "score_mod is not supported with fp8_output!" | ||
| assert not context_parallel, "score_mod is not supported with context parallelism!" | ||
| assert qkv_format != "thd", "score_mod is not supported with qkv_format='thd'!" | ||
| assert ( | ||
| not user_supplied_seqlens | ||
| ), "score_mod is mutually exclusive with explicit sequence length metadata!" | ||
| assert not pad_between_seqs, "score_mod is not supported with pad_between_seqs!" | ||
| assert ( | ||
| attention_mask is None | ||
| ), "score_mod is mutually exclusive with attention_mask!" | ||
| assert attn_mask_type == "no_mask", "score_mod requires attn_mask_type='no_mask'!" | ||
| assert window_size is None or window_size == ( | ||
| -1, | ||
| -1, | ||
| ), "score_mod is mutually exclusive with sliding window attention!" | ||
| assert ( | ||
| core_attention_bias_type == "no_bias" and core_attention_bias is None | ||
| ), "score_mod is mutually exclusive with attention bias!" | ||
| assert alibi_slopes is None, "score_mod is mutually exclusive with ALiBi!" | ||
| assert ( | ||
| self.softmax_type == "vanilla" | ||
| ), "score_mod is mutually exclusive with sink attention!" | ||
| assert ( | ||
| self.attention_dropout == 0.0 | ||
| ), "score_mod is not supported with attention dropout!" | ||
| assert ( | ||
| not self.return_max_logit | ||
| ), "score_mod is not supported with return_max_logit!" | ||
| assert ( | ||
| not checkpoint_core_attention | ||
| ), "score_mod is not supported with checkpoint_core_attention!" | ||
| assert ( | ||
| not is_graph_capturing() | ||
| ), "score_mod is not supported with CUDA graph capture!" | ||
| assert num_splits == 1, "score_mod is not supported with num_splits != 1!" | ||
| assert q_format in ["sbhd", "bshd"] and kv_format in [ | ||
| "sbhd", | ||
| "bshd", | ||
| ], "score_mod only supports SBHD/BSHD QKV formats!" | ||
| if score_mod_tensors is not None: | ||
| assert isinstance(score_mod_tensors, dict), "score_mod_tensors must be a dict!" | ||
| assert all( | ||
| isinstance(k, str) and isinstance(v, torch.Tensor) | ||
| for k, v in score_mod_tensors.items() | ||
| ), "score_mod_tensors must map string names to torch.Tensor instances!" | ||
| if score_mod_bprop_tensors is not None: |
There was a problem hiding this comment.
Aren't these checks duplicates of the checks in get_attention_backend() ?
get_attention_backend() alone can be the source of truth and this duplications should not be needed here. The only checks that should be added in here are those that cannot be / should not be added in the get_attention_backend()
I'd suggest to change this to accommodate for the same
| ) | ||
| global _attention_backends | ||
| if is_in_onnx_export_mode(): | ||
| if is_in_onnx_export_mode() and score_mod is None: |
There was a problem hiding this comment.
Is this necessary here if dpa_utils.get_attention_backend(attention_params) does get called in the else block below ?
The flash, fused, unfused would be set in there anyways rgiht ?
Or am I missing something ?
cc: @cyanguwa
| return output.contiguous() | ||
|
|
||
|
|
||
| def _bhsd_dim_stride( |
There was a problem hiding this comment.
I agree with this and was my first thought too.
We should club these function into a couple classes that can sit in this file at the very least.
However, I think this approach is still not the right approach. We should have a separate flex_attention.py file similar to context_parallel.py and backends.py can import it similar to how it imports the CP functions right now.
I strongly recommend this for two reasons :
- When we refactored attention as a whole early last year, the idea was to modularize attention. That was the reason CP was moved out of attention. With Flex attention's functionality and code in here being fairly decoupled from vanilla DPA, it should be easier to move it out. Leaving this code in here would add ~1000 lines of code that is not related to the vanilla DPA and would practically be undoing the refactoring work we did early last year. The same reason for moving CP to it's own file should also apply to Flex attention.
- A developer/user of TE PyT DPA should not have to worry about the details of flex attention. Similarly someones modifying flex should not be bogged down by the details of vanilla fused attn. Hence, decoupling is important to aid with debugging as well
| def _import_cudnn_frontend(): | ||
| """Import the vendored cuDNN frontend if built, otherwise use the installed package.""" | ||
| cudnn_frontend_path = str(_CUDNN_FRONTEND_PYTHON_PATH) | ||
| cudnn_frontend_package = _CUDNN_FRONTEND_PYTHON_PATH / "cudnn" | ||
| if ( | ||
| any(cudnn_frontend_package.glob("_compiled_module*")) | ||
| and cudnn_frontend_path not in sys.path | ||
| ): | ||
| sys.path.insert(0, cudnn_frontend_path) | ||
| return importlib.import_module("cudnn") | ||
|
|
There was a problem hiding this comment.
How about this?:
def _import_cudnn_frontend():
cudnn_frontend_path = str(_CUDNN_FRONTEND_PYTHON_PATH)
cudnn_frontend_package = _CUDNN_FRONTEND_PYTHON_PATH / "cudnn"
if (
any(cudnn_frontend_package.glob("_compiled_module*"))
and cudnn_frontend_path not in sys.path
):
sys.path.insert(0, cudnn_frontend_path)
return importlib.import_module("cudnn")
# Fall back
if importlib.util.find_spec("cudnn") is not None:
return importlib.import_module("cudnn")
# Fail with a message
raise ImportError(
"cuDNN Frontend Python package not found. "
"Install it with: pip install nvidia-cudnn-frontend"
)
| return out, max_logit, (None, None, None, d_softmax_offset) | ||
|
|
||
|
|
||
| def _score_mod_causal(score_mod_graph, score_tensor, tensors): |
There was a problem hiding this comment.
I would strongly recommend that similar to the CP tests we have a separate Flex attention test file. Firstly for modularization and secondly because the Flex attention tests do not really end up using the test_dot_product_attention() base test like other DPA tests in the file do so there's no code reuse reasons for it either.
These isolated ~800 lines of code can sit in it's own file if it isn't really using of the funtions in here directly but writing the flex tests as "new" tests or else the flex tests must reuse the DPA setup in here and integrate into that.
I've also shared more details on this in my comment in the backends.py file
cc: @cyanguwa
|
Thanks for creating this PR @vcherepanov-nv I was curious about:
|
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
for more information, see https://pre-commit.ci
|
Thanks for the thorough review!
I haven't done any benchmarking. Reportedly (from a Slack thread) score_mod can lead to significant perf gains if it allows to avoid mask materialization. For causal, I think I observed cuDNN choosing exactly the same kernel with score_mod and the explicit causal flag.
Sure, thanks for linking! |
3 participants
Description
This PR introduces an alternative, Python-only code path for the FusedAttention backend for PyTorch.
The user can specify score_mod and score_mod_bprop functions, which get routed to the corresponding parameters of the sdpa and sdpa_backward calls to cuDNN FE.
Fixes # (issue)
#2492
Type of change
Changes
Checklist: