[All] Refactor fused attention APIs with cuDNN-frontend support checks and opaque config/params handles - #2964
[All] Refactor fused attention APIs with cuDNN-frontend support checks and opaque config/params handles#2964cyanguwa wants to merge 95 commits into
Conversation
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
for more information, see https://pre-commit.ci
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
cyanguwa
changed the title
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
Greptile SummaryThis PR refactors the fused attention APIs across PyTorch and JAX backends to introduce opaque config/params handles (
Confidence Score: 3/5
Important Files Changed
Sequence DiagramsequenceDiagram
participant C as Caller (PyTorch/JAX)
participant API as fused_attn.cpp
participant Cfg as FusedAttnConfig
participant Cache as GraphCache
participant CDNN as cuDNN Frontend
Note over C,CDNN: Backend Query Phase
C->>API: nvte_get_fused_attn_backend_v2(config_handle)
API->>Cfg: FusedAttnConfig::derive()
API->>Cache: get_or_build_cached_graph(fwd_key)
Cache->>CDNN: validate_and_check_support()
CDNN-->>Cache: GraphAndTensors or UnsupportedGraph
Cache-->>API: "CachedGraph* or nullptr"
API->>Cache: get_or_build_cached_graph(bwd_key) [training only]
Cache-->>API: "CachedGraph* or nullptr"
API-->>C: "NVTE_FUSED_ATTN_* backend + message"
Note over C,CDNN: Forward Execution Phase
C->>API: nvte_fused_attn_fwd_v2(params_handle)
API->>Cfg: FusedAttnFwdParams::make_config()
API->>Cfg: FusedAttnConfig::derive()
API->>Cfg: FusedAttnConfig::make_cache_key()
API->>Cache: get_or_build_cached_graph()
Cache-->>API: "CachedGraph* (cache hit)"
API->>CDNN: ensure_plans_built() [lazy compile]
API->>CDNN: execute(variant_pack)
CDNN-->>C: output tensors
Note over C,CDNN: Backward Execution Phase
C->>API: nvte_fused_attn_bwd_v2(params_handle)
API->>Cfg: FusedAttnBwdParams::make_config()
API->>Cfg: FusedAttnConfig::derive()
API->>Cache: get_or_build_cached_graph()
Cache-->>API: "CachedGraph* (cache hit)"
API->>CDNN: ensure_plans_built()
API->>CDNN: execute(variant_pack)
CDNN-->>C: grad tensors
Reviews (41): Last reviewed commit: "address review comments" | Re-trigger Greptile |
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
for more information, see https://pre-commit.ci
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
for more information, see https://pre-commit.ci
|
/te-ci L1 |
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
for more information, see https://pre-commit.ci
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
for more information, see https://pre-commit.ci
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
for more information, see https://pre-commit.ci
|
/te-ci L1 |
| t_q, | ||
| num_tokens_kv * cp_size * s_kv // max_seqlen_kv if max_seqlen_kv else 0, | ||
| ) | ||
| for s_kv in dict.fromkeys([s_kv_chunk, max_seqlen_kv]) |
There was a problem hiding this comment.
this varies for every rank, would it make a difference?
|
| "<b>Note</b>\n", | ||
| " \n", | ||
| "Environment variables <code>NVTE_FLASH_ATTN</code>, <code>NVTE_UNFUSED_ATTN</code>, <code>NVTE_FUSED_ATTN_BACKEND</code>, and <code>NVTE_FUSED_ATTN_USE_FAv2_BWD</code> are supported in PyTorch. <code>NVTE_FUSED_ATTN</code> and <code>NVTE_ALLOW_NONDETERMINISTIC_ALGO</code> are supported in both PyTorch and JAX.\n", | ||
| "Environment variables <code>NVTE_FLASH_ATTN</code>, <code>NVTE_UNFUSED_ATTN</code>, and <code>NVTE_FUSED_ATTN_USE_FAv2_BWD</code> are supported in PyTorch. <code>NVTE_FUSED_ATTN</code> and <code>NVTE_ALLOW_NONDETERMINISTIC_ALGO</code> are supported in both PyTorch and JAX.\n", |
There was a problem hiding this comment.
No NVTE_FUSED_ATTN_USE_FAv2_BWD support in JAX ?
| // Restrict each direction's key to the fields its graph actually consumes, so | ||
| // no redundant graphs are built and no cache misses either |
There was a problem hiding this comment.
That's a good change to avoid redundant graphs if fwd only
| " bytes)"); | ||
| NVTE_CHECK(buf != nullptr, "Invalid buffer (got NULL)"); | ||
|
|
||
| auto &cfg = *get_fused_attn_config_mutable(config); |
There was a problem hiding this comment.
nit: for the setter and getter (nvte_set_fused_attn_config_attribute and nvte_get_fused_attn_config_attribute) - do you think in the future it might make sense to make this a bulk API in which one can request to set/get multiple attributes in in one call? Which would then reduce multiple calls to get_fused_attn_config_mutable() and get_fused_attn_config ?
| int64_t window_size_right, bool return_max_logit, bool cuda_graph, bool deterministic) { | ||
| namespace { | ||
|
|
||
| // The per-thread storage for the diagnostic string; it's re-used (cleared + re-populated) |
There was a problem hiding this comment.
nit:
| // The per-thread storage for the diagnostic string; it's re-used (cleared + re-populated) | |
| // The per-thread storage for the diagnostic string; it is re-used (cleared + re-populated) |
| // Only used when THD format is requested. | ||
| cudnnHandle_t handle = cudnnExecutionPlanManager::Instance().GetHandle(); | ||
| const NVTE_QKV_Format qkv_format = nvte_get_qkv_format(cfg.qkv_layout); | ||
| const NVTE_QKV_Layout_Group layout_group = nvte_get_qkv_layout_group(cfg.qkv_layout); |
There was a problem hiding this comment.
nit: Although having the explicit type like NVTE_QKV_Format and NVTE_QKV_Layout_Group is almost always the better option, consider using auto instead ? Especially since the var names are descriptive enough to understand the data type ?
It could help reduce code verbosity
| cache_hit = (it != cache.end()); | ||
| if (cache_hit) cached_graph = it->second; | ||
| } | ||
| graph_cache_debug::record_cache_lookup("fwd", cache_hit, cfg); |
There was a problem hiding this comment.
I think it makes sense to have the recording outside the mutex's scope so that the mutex is not held when performing (slow) I/o ops, however, this would mean that the recorded logs for the cache ops may not reflect exact wall clock ordering. I think it is vital to mention this in the docs/code if not already so that the users are aware of this
| bool is_ragged_q = (q_format == NVTE_QKV_Format::NVTE_THD); | ||
| bool is_ragged_kv = (kv_format == NVTE_QKV_Format::NVTE_THD); | ||
| bool is_ragged_q = cfg.is_ragged_q; | ||
| bool is_ragged_kv = cfg.is_ragged_kv; |
| bool is_causal_bottom_right = cfg.is_causal_bottom_right; | ||
| bool is_padding = cfg.is_padding; |
| graph_cfg.derive(); | ||
|
|
||
| size_t workspace_size = 0; | ||
| try { |
There was a problem hiding this comment.
Thanks for adding the try catch in here. I was hoping for it while reviewing the code for fused_attn_arbitrary_seqlen_fwd_impl()
| static thread_local CacheType sdpa_f16_bprop_cache; | ||
| using CacheType = std::map<FusedAttnConfig, graph_and_tensors>; | ||
| static CacheType sdpa_f16_bprop_cache; | ||
| static std::mutex sdpa_f16_bprop_cache_mutex; |
There was a problem hiding this comment.
maybe nit and overthinking on my part but would it better to reverse the order of creation ?
Because the order of destruction will be reverse of creation, we'd like to first destroy the resource (cache) and then the mutex guarding it, right ?
It may also be beneficial (to make it mistake proof) if we tie these together in a struct with the suggested new ordering above so that if ever anyone else touches the cache and mutex code in the future they do not need to worry about the individual object ordering (destroying the struct object is all they'd care about and we can take care of the reordering in the struct object)
|
Note to self: integrate these changes to this PR, thanks to @sudhakarsingh27. cc #3092 PR 2964 addresses the following points from the above PR: These points aren't fully addressed in 2964: |
| void* devActualSeqlenKV = static_cast<int8_t*>(devActualSeqlenQ) + b * sizeof(int32_t); | ||
| cu_seqlens_to_actual_seqlens<<<grid, nthreads_per_block, 0, stream>>>( | ||
| b, b, static_cast<const int32_t*>(devPtrcuSeqlensQ), // TODO(pass max_b) | ||
| b, b, static_cast<const int32_t*>(devPtrcuSeqlensQ), // TODO(pass bucketed_batch_size) |
There was a problem hiding this comment.
Is this TODO for the future ? if yes,
| b, b, static_cast<const int32_t*>(devPtrcuSeqlensQ), // TODO(pass bucketed_batch_size) | |
| b, b, static_cast<const int32_t*>(devPtrcuSeqlensQ), // TODO(<GH username>): pass bucketed_batch_size |
| bool is_bias = (bias_type == NVTE_Bias_Type::NVTE_POST_SCALE_BIAS); | ||
| bool is_alibi = (bias_type == NVTE_Bias_Type::NVTE_ALIBI); | ||
| bool is_causal = ((mask_type == NVTE_Mask_Type::NVTE_CAUSAL_MASK) || | ||
| (mask_type == NVTE_Mask_Type::NVTE_PADDING_CAUSAL_MASK)); | ||
| bool is_padding = ((mask_type == NVTE_Mask_Type::NVTE_PADDING_MASK) || | ||
| (mask_type == NVTE_Mask_Type::NVTE_PADDING_CAUSAL_MASK)); | ||
| bool is_causal_bottom_right = cfg.is_causal_bottom_right; | ||
| bool is_padding = cfg.is_padding; | ||
| bool is_dropout = (dropout_probability != 0.0f); | ||
| bool is_softmax_offset = (softmax_type != NVTE_Softmax_Type::NVTE_VANILLA_SOFTMAX); |
There was a problem hiding this comment.
outide this PR's scope but if possible: const ?
| graph_cache_debug::record_build("bwd"); | ||
| // Lock the insert. If another thread inserted a graph for the same key while we were building, | ||
| // use their graph (it's the same as ours) and discard our graph. | ||
| { |
There was a problem hiding this comment.
Maybe we already have this, but if not, it might be useful to add a cache specific test especially since that's a pretty imp component of graphing in TE attention.
Add a single-thread test that queries one config twice, then executes fused attention with matching params.
Assert the first query causes one miss/build, while the second query and execution are hits with no additional build. Maybe then modify one graph-defining field and assert exactly one new miss/build—this directly catches broken key normalization and unintended recompilation.
| try { | ||
| fused_attn::fused_attn_fp8_fwd_impl( | ||
| graph_cfg, | ||
| /*devPtrQ=*/nullptr, /*devPtrK=*/nullptr, /*devPtrV=*/nullptr, |
There was a problem hiding this comment.
thanks for adding the args as comments here
|
|
||
| // More readable, shorter thread IDs (0, 1, 2, ...). | ||
| inline unsigned thread_seq_id() { | ||
| static std::atomic<unsigned> next{0}; |
There was a problem hiding this comment.
This got me thinking about logging device id a bit more
IIUC, the cache key includes device ID, but the debug events omit it right ?.
Could we log descriptor.device_id and pass the normalized descriptor to the recorder? This would make multi-GPU cache behavior diagnosable.
| } // namespace fused_attn | ||
| } // namespace transformer_engine | ||
|
|
||
| #endif // TRANSFORMER_ENGINE_COMMON_FUSED_ATTN_GRAPH_CACHE_DEBUG_H_ |
There was a problem hiding this comment.
If it is not too much effort, would be nice to see an example of what this logs looks like for the different diagnostics enabled. This is a good to have only so feel free to skip over
| template <typename T> | ||
| FusedAttnFwdParamsWrapper &set_attr(NVTEFusedAttnFwdParamsAttribute attr, T val) noexcept { | ||
| nvte_set_fused_attn_fwd_params_attribute(params_, attr, &val, sizeof(val)); | ||
| return *this; |
There was a problem hiding this comment.
Chaining is a good addition to this 👍
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
|
/te-ci L0 L1 L2 L3 |
| namespace { | ||
|
|
||
| // The per-thread storage for the diagnostic string; it's re-used (cleared + re-populated) | ||
| // on every call to nvte_get_fused_attn_backend_v2 on the same thread. | ||
| thread_local std::string fused_attn_backend_message_buffer; | ||
|
|
||
| // Stash `reason` in the thread-local buffer and, if the caller asked for a diagnostic, | ||
| // publish a NUL-terminated pointer to it via `*message`. Safe to call with `message == nullptr`. | ||
| void set_message(const char **message, std::string reason) { | ||
| if (message == nullptr) return; | ||
| fused_attn_backend_message_buffer = std::move(reason); | ||
| *message = fused_attn_backend_message_buffer.c_str(); | ||
| } |
There was a problem hiding this comment.
message pointer is invalidated by any subsequent same-thread backend call
*message is set to .c_str() of a thread_local std::string. Any call to nvte_get_fused_attn_backend_v2 on the same thread (including the internal calls from nvte_fused_attn_fwd_v2 and nvte_fused_attn_bwd_v2) will std::move() a new string into fused_attn_backend_message_buffer, destroying the previous string object and making the pointer dangle. The internal probe calls currently pass nullptr so the buffer isn't clobbered by them, but any caller that stores the returned const char* and then makes a subsequent backend call will read freed memory. The API contract (e.g. "copy this string before calling anything else") should be clearly documented, or the v2 signature should return std::string instead of const char**.
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
…run, still build plans in probes Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
…be, dry-run, still build plans in probes" This reverts commit 8fdd81d. Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
…ad and not modify cfg Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
|
/te-ci L0 L1 L2 L3 |
| /************************************************************************* | ||
| * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. | ||
| * | ||
| * See LICENSE for license information. | ||
| ************************************************************************/ | ||
|
|
||
| // ============================================================================ | ||
| // The fused-attention graph cache: what a cache entry is, how one is looked up | ||
| // or built, and the frontend calls that make a constructed graph usable. | ||
| // | ||
| // Each of the four build sites (f16 and fp8, forward and backward) differs only | ||
| // in how it constructs its graph and which tensors it hands back. Everything | ||
| // after that -- the lookup, the locking, the once-per-entry plan build, the | ||
| // support check, and the remembering of what cuDNN refused -- is the same at all | ||
| // four, and lives here so it has one definition rather than four copies to keep | ||
| // in step. | ||
| // | ||
| // This header is deliberately not part of utils.h: it needs the cuDNN frontend, | ||
| // and utils.h is included by translation units (utils.cu) that otherwise do not. | ||
| // ============================================================================ | ||
|
|
||
| #ifndef TRANSFORMER_ENGINE_COMMON_FUSED_ATTN_GRAPH_CACHE_H_ | ||
| #define TRANSFORMER_ENGINE_COMMON_FUSED_ATTN_GRAPH_CACHE_H_ | ||
|
|
||
| #include <exception> | ||
| #include <map> | ||
| #include <memory> | ||
| #include <mutex> | ||
| #include <stdexcept> | ||
| #include <string> | ||
| #include <tuple> | ||
| #include <utility> | ||
|
|
||
| #include "../common.h" | ||
| #include "../cudnn_utils.h" | ||
| #include "config_and_params.h" | ||
| #include "graph_cache_debug.h" | ||
|
|
||
| namespace transformer_engine { | ||
| namespace fused_attn { | ||
|
|
||
| // cuDNN's refusal to run a graph, as opposed to a failure to try. The distinction is what makes | ||
| // the negative cache in get_or_build_cached_graph() safe: a refusal is a verdict on the | ||
| // configuration and reproducible for a given key, so it can be remembered and replayed, whereas | ||
| // a failure that came from the machine's state at that moment (an allocation that did not fit, a | ||
| // CUDA error left behind by unrelated work) could well succeed on the next attempt and must not | ||
| // be turned into a permanent answer. Only the four adjudicating frontend calls in | ||
| // validate_and_check_support() raise this; every other failure keeps its ordinary type and is | ||
| // re-attempted the next time the key comes around. | ||
| struct UnsupportedGraph : public std::runtime_error { | ||
| explicit UnsupportedGraph(const std::string &reason) : std::runtime_error(reason) {} | ||
| }; | ||
|
|
||
| // The reason string an is_supported_* helper reports for `e`: its message, or `fallback` if it | ||
| // has none. Those helpers signal support by returning the empty string, so a refusal that | ||
| // arrives without an explanation would be read as an endorsement and the caller would go on to | ||
| // run a graph cuDNN has just declined. Nothing raised through NVTE_ERROR can be empty, since it | ||
| // prefixes file and line, but that is a property of our macros rather than of every exception | ||
| // that can reach a catch clause, and it is not what the contract should rest on. | ||
| inline std::string refusal_reason(const std::exception &e, const char *fallback) { | ||
| const char *what = e.what(); | ||
| return (what != nullptr && what[0] != '\0') ? std::string(what) : std::string(fallback); | ||
| } | ||
|
|
||
| // A graph in the cache, plus the tensor attributes needed to bind runtime pointers to it. | ||
| // | ||
| // Entries are built only as far as check_support(), which is all it takes to decide whether | ||
| // a configuration is supported. build_plans() is the kernel-compilation step and the most | ||
| // expensive of the five frontend calls, so a support query stops short of it: the query never | ||
| // executes the graph, and many of the keys it builds are never executed by anything. The | ||
| // execution path finishes the build instead, the first time the graph is needed to run. | ||
| // | ||
| // plans_built guards that completion. It has to happen exactly once per entry, because the | ||
| // cached graph is shared across threads and build_plans() mutates it in place -- two threads | ||
| // reaching the same unfinished entry must not both build it. Keeping the flag inside the entry | ||
| // keeps it from drifting away from the graph it describes, and leaves unrelated keys free to | ||
| // build concurrently. A build that throws leaves the flag unset, so a later call retries | ||
| // rather than executing a graph with no plans. | ||
| template <typename GraphAndTensors> | ||
| struct CachedGraph { | ||
| explicit CachedGraph(GraphAndTensors tensors) : tensors(std::move(tensors)) {} | ||
|
|
||
| GraphAndTensors tensors; | ||
| std::once_flag plans_built; | ||
| }; | ||
|
|
||
| // One build site's cache. Process-wide rather than per-thread so that a graph is reused | ||
| // across threads instead of rebuilt by each: cuDNN >= 9.0 allows concurrent execution of a | ||
| // shared plan, and cudnn-frontend >= 1.25.0 has a thread-safe execute(). | ||
| // | ||
| // Refusals are cached alongside the graphs, under the same keys and the same lock. A support | ||
| // query for an unsupported configuration is otherwise the most expensive thing this cache sees: | ||
| // it builds the whole graph, spends the four frontend calls, and throws the result away, and it | ||
| // does so again on every query, because a rejection left nothing behind to find. `unsupported` | ||
| // is what it leaves behind -- cuDNN's own account of the refusal, which is the entire useful | ||
| // output of a failed query, so nothing is lost by answering from it. Reasons are short strings | ||
| // and there is one per refused key, so this grows far slower than the graphs beside it. | ||
| // Holding the lock and the maps together is also what fixes their relative lifetimes. Members | ||
| // are destroyed in reverse declaration order, so the mutex is declared first to be destroyed | ||
| // last: the maps go while their guard is still valid, rather than the other way round. Declaring | ||
| // a cache and its lock as two separate objects leaves that ordering to whoever writes the next | ||
| // one; declaring them here settles it once. | ||
| template <typename GraphAndTensors> | ||
| struct GraphCache { | ||
| std::mutex mutex; // guards both maps below | ||
| std::map<FusedAttnConfig, std::shared_ptr<CachedGraph<GraphAndTensors>>> supported; | ||
| std::map<FusedAttnConfig, std::string> unsupported; | ||
| }; | ||
|
|
||
| // Takes a constructed graph through the frontend calls that decide whether cuDNN can run it: | ||
| // validate, build_operation_graph, create_execution_plans, check_support. The sequence is | ||
| // identical for both passes and both backends, so it is defined once here; `pass` only selects | ||
| // which set of stage timers the calls are attributed to. | ||
| // | ||
| // Support is reported by throwing rather than by a return value. NVTE_CHECK_CUDNN_FE raises | ||
| // an exception carrying cuDNN's own explanation of the rejection, and that text is what the | ||
| // is_supported_* helpers return as the reason a backend was refused -- so a bool here would | ||
| // discard the one thing a support probe exists to produce. Callers that are about to execute | ||
| // the graph want the throw as well, since there is nothing useful to do with an unsupported | ||
| // graph but fail. | ||
| // | ||
| // The throw is re-raised as UnsupportedGraph, which is what marks it cacheable. These four calls | ||
| // are cuDNN adjudicating a graph it has been handed, so a failure among them is a statement about | ||
| // the graph rather than about the moment -- which is the property the negative cache needs, and | ||
| // the reason the boundary is drawn here rather than around a wider region. build_plans() and | ||
| // execute() sit outside it: they commit real resources and can fail for reasons that have nothing | ||
| // to do with the configuration. | ||
| // | ||
| // build_plans() is left out for a second reason as well: it belongs to whoever executes the graph, | ||
| // once, the first time it is needed. See CachedGraph. | ||
| inline void validate_and_check_support(const char *pass, cudnn_frontend::graph::Graph &graph, | ||
| cudnnHandle_t handle) { | ||
| try { | ||
| graph_cache_debug::timer(pass, graph_cache_debug::BuildStage::Validate, | ||
| [&] { NVTE_CHECK_CUDNN_FE(graph.validate()); }); | ||
| graph_cache_debug::timer(pass, graph_cache_debug::BuildStage::BuildOpGraph, | ||
| [&] { NVTE_CHECK_CUDNN_FE(graph.build_operation_graph(handle)); }); | ||
| graph_cache_debug::timer(pass, graph_cache_debug::BuildStage::CreatePlans, [&] { | ||
| NVTE_CHECK_CUDNN_FE(graph.create_execution_plans({cudnn_frontend::HeurMode_t::A})); | ||
| }); | ||
| graph_cache_debug::timer(pass, graph_cache_debug::BuildStage::CheckSupport, | ||
| [&] { NVTE_CHECK_CUDNN_FE(graph.check_support()); }); | ||
| } catch (const std::exception &e) { | ||
| throw UnsupportedGraph(e.what()); | ||
| } | ||
| } | ||
|
|
||
| // The cached entry for `key`, building and inserting it via `build` if absent. Throws | ||
| // UnsupportedGraph if cuDNN refuses the graph -- this time or on an earlier call, the two being | ||
| // indistinguishable to the caller by design. | ||
| // | ||
| // `build` only constructs a graph; this is what puts it through validate_and_check_support(), so | ||
| // the entries in the cache are exactly the graphs cuDNN has agreed to run. Those four calls sit | ||
| // on the miss path because they are part of building an entry rather than reading one: repeating | ||
| // them on a hit would redo the operation graph and the plan search for a graph that has already | ||
| // been through both. | ||
| // | ||
| // `key` must be a normalized key -- FusedAttnConfig::make_cache_key()'s output -- and not a | ||
| // raw execution config. Two configs that differ only in a field no graph reads (attn_scale, | ||
| // say) have to reach the same entry, which is what normalization is for; passing the raw | ||
| // config instead silently multiplies the cache by fields the graph never consumes. | ||
| // | ||
| // Only the map operations are locked, not `build`. A graph build is the expensive part and | ||
| // holding the lock across it would serialize builds of unrelated keys, so two threads racing | ||
| // on the same key may both build. That is a wasted build, not a correctness problem: the | ||
| // loser drops its own graph and takes the winner's, so every caller of a given key gets one | ||
| // shared entry and the once-flag inside it still governs the plan build. The wasted build is | ||
| // visible in diagnostics as a BUILD with no matching MISS of its own. The same race on a | ||
| // refused key is equally harmless, both threads storing the same reason. | ||
| template <typename GraphAndTensors, typename BuildFn> | ||
| std::shared_ptr<CachedGraph<GraphAndTensors>> get_or_build_cached_graph( | ||
| GraphCache<GraphAndTensors> &cache, const FusedAttnConfig &key, const char *pass, | ||
| cudnnHandle_t handle, BuildFn &&build) { | ||
| using Entry = CachedGraph<GraphAndTensors>; | ||
|
|
||
| std::shared_ptr<Entry> cached; | ||
| bool refused = false; | ||
| std::string reason; | ||
| { | ||
| std::lock_guard<std::mutex> lock(cache.mutex); | ||
| auto it = cache.supported.find(key); | ||
| if (it != cache.supported.end()) { | ||
| cached = it->second; | ||
| } else { | ||
| auto refusal = cache.unsupported.find(key); | ||
| refused = (refusal != cache.unsupported.end()); | ||
| if (refused) reason = refusal->second; | ||
| } | ||
| } | ||
| using graph_cache_debug::LookupResult; | ||
| LookupResult outcome = LookupResult::Miss; | ||
| if (cached != nullptr) { | ||
| outcome = LookupResult::Hit; | ||
| } else if (refused) { | ||
| outcome = LookupResult::Unsupported; | ||
| } | ||
| // Recorded after the lock is dropped, so that writing a trace line cannot hold up threads | ||
| // querying other keys. The counters are exact, but two lookups that raced on the lock can be | ||
| // recorded in the opposite order, so read a level-2 trace as the set of lookups that happened | ||
| // rather than as the sequence they happened in. | ||
| graph_cache_debug::record_cache_lookup(pass, outcome, key); | ||
|
|
||
| if (cached != nullptr) return cached; | ||
| // Raised rather than returned so that a replayed refusal is the same event as a fresh one: | ||
| // every caller already has to handle the build refusing, and none of them would have anything | ||
| // else to do with a second, quieter way of saying so. | ||
| if (refused) throw UnsupportedGraph(reason); | ||
|
|
||
| std::shared_ptr<Entry> entry; | ||
| try { | ||
| entry = std::make_shared<Entry>(build()); | ||
| // Every site's tensor tuple leads with its graph, which is the one thing all four have in | ||
| // common and the only element this needs. A tuple that stopped leading with it would fail to | ||
| // compile here rather than quietly validate the wrong object. | ||
| validate_and_check_support(pass, *std::get<0>(entry->tensors), handle); | ||
| } catch (const UnsupportedGraph &e) { | ||
| { | ||
| std::lock_guard<std::mutex> lock(cache.mutex); | ||
| cache.unsupported.insert({key, e.what()}); | ||
| } | ||
| graph_cache_debug::record_unsupported(pass); | ||
| throw; | ||
| } | ||
| graph_cache_debug::record_build(pass); | ||
| { | ||
| std::lock_guard<std::mutex> lock(cache.mutex); | ||
| return cache.supported.insert({key, std::move(entry)}).first->second; | ||
| } | ||
| } | ||
|
|
||
| // Runs the plan build that get_or_build_cached_graph() left undone, once per entry. | ||
| // | ||
| // Call this only when the graph is about to be executed, which is why it is a separate step | ||
| // rather than the tail of the lookup: a support query builds entries that nothing ever runs, and | ||
| // kernel compilation is the most expensive of the five frontend calls, so a query that paid for | ||
| // it would be paying for nothing. See CachedGraph for why the flag lives inside the entry and | ||
| // what a throw here leaves behind. | ||
| template <typename GraphAndTensors> | ||
| void ensure_plans_built(const char *pass, CachedGraph<GraphAndTensors> &entry) { | ||
| std::call_once(entry.plans_built, [&] { | ||
| cudnn_frontend::graph::Graph &graph = *std::get<0>(entry.tensors); | ||
| graph_cache_debug::timer(pass, graph_cache_debug::BuildStage::BuildPlans, | ||
| [&] { NVTE_CHECK_CUDNN_FE(graph.build_plans()); }); | ||
| graph_cache_debug::record_plans_built(pass); | ||
| }); | ||
| } | ||
|
|
||
| } // namespace fused_attn | ||
| } // namespace transformer_engine | ||
|
|
||
| #endif // TRANSFORMER_ENGINE_COMMON_FUSED_ATTN_GRAPH_CACHE_H_ |
There was a problem hiding this comment.
validate_and_check_support() wraps every thrown exception — including std::bad_alloc, CUDA OOM, or transient driver errors — as UnsupportedGraph and inserts it into the negative (unsupported) map. Once a config is stored there, no future call will ever retry it. A single low-memory spike during build_operation_graph will permanently mark a valid, supported config as unsupported for the lifetime of the process, silently falling back to a slower kernel without any indication of the root cause.
The fix is to distinguish truly-unsupported graphs (cuDNN returning CUDNN_STATUS_NOT_SUPPORTED or equivalent) from transient resource failures, and only cache the former.
5 participants
Description
TE currently hand-maintains the fused-attention backend-selection logic in
nvte_get_fused_attn_backend, duplicating cuDNN's support rules. This list drifts out of sync as cuDNN evolves, and the support check can disagree with what actually runs.This PR replaces that logic with cuDNN-frontend's production-grade support checks. The new
nvte_get_fused_attn_backend_v2builds the same graph cuDNN executes at runtime, so the probe and execution can no longer diverge. It caches the graph on success and returns a diagnostic message on failure, giving users actionable guidance (e.g. adjust the config, GPU architecture, or cuDNN version).This PR also reworks
nvte_fused_attn_fwd/nvte_fused_attn_bwdintonvte_fused_attn_fwd_v2/nvte_fused_attn_bwd_v2, which take opaque, attribute-based config/params handles instead of long flat argument lists — improving TE's API and ABI stability.Legacy APIs are retained as deprecated shims that route through the v2 APIs, so existing callers keep working.
Type of change
Changes
API rework (opaque config/params + v2 entry points)
common/fused_attn/config_and_params.{h,cpp},common/include/transformer_engine/fused_attn.h): newNVTEFusedAttnConfig/NVTEFusedAttnFwdParams/NVTEFusedAttnBwdParamswithcreate/destroy/get/setattribute accessors, for better API/ABI stability. The cache key, probe, and execution now all originate from one place viamake_config/derive/make_cache_key.common/fused_attn/fused_attn*.{cpp,cu}):nvte_get_fused_attn_backend_v2,nvte_fused_attn_fwd_v2, andnvte_fused_attn_bwd_v2. The F16 and FP8is_supported_*probes copy the config, set direction,derive(), and attempt a null-pointer graph build viacheck_support— i.e. the same graph cuDNN builds at runtime, so probe and execution can't diverge.nvte_get_fused_attn_backend/nvte_fused_attn_fwd/nvte_fused_attn_bwdare retained, routed through the v2 APIs.csrc/extensions/attention.cpp) and JAX (jax/csrc/extensions/attention.cpp).Correctness & backend selection
cp_per_step_configsprobes each context-parallel step instead of only the global, non-CP config.log2(0)guard: avoids UB when casting-inftosize_tinget_max_batch_size/get_max_tokens.Diagnostics
NVTE_DEBUG/NVTE_DEBUG_LEVELfor JAX (parity with PyTorch): level 1 reports the selected backend; level 2 adds a diagnostic message explaining why fused attention was rejected.NVTE_FUSED_ATTN_CACHE_DEBUG: opt-in instrumentation that reports cuDNN graph build-vs-execution counts and per-stage cudnn-frontend build timings, so cache hit/miss/build/exec behaviors and graph build time can be inspected. Off by default; available for both PyTorch and Jax.Cleanup / removals
NVTE_FUSED_ATTN_BACKEND— the two remaining backends (F16, FP8) are mutually exclusive now that max512 is gone.Q_ID/.../MASK_VAL_IDmacros (used only by the max512 backend).cudnn_frontend::xxxutility functions (used only byfp8_impl_v0and max512).fused_attn/headers.Tests
Checklist: