[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 97 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 replaces TE's hand-maintained cuDNN support matrix with cuDNN-frontend's production-grade support checks and introduces opaque config/params handles (
Confidence Score: 3/5
Important Files Changed
Sequence DiagramsequenceDiagram
participant Caller as PyTorch/JAX Caller
participant v2 as nvte_get_fused_attn_backend_v2
participant Probe as is_supported_f16/fp8_fwd/bwd
participant Cache as Process-wide GraphCache
participant cuDNN as cuDNN Frontend
Caller->>v2: config (NVTEFusedAttnConfig)
v2->>v2: cfg.derive()
v2->>Probe: is_supported_f16_fwd(cfg, handle)
Probe->>Probe: make_cache_key()
Probe->>Cache: get_or_build_cached_graph(key)
alt Cache Miss
Cache->>cuDNN: build_sdpa_f16_fwd_graph()
cuDNN-->>Cache: graph
Cache->>cuDNN: validate + build_op_graph + create_plans + check_support
cuDNN-->>Cache: UnsupportedGraph OR ok
Cache-->>Probe: cached entry OR throw UnsupportedGraph
else Cache Hit (supported)
Cache-->>Probe: cached entry
else Cache Hit (unsupported)
Cache-->>Probe: throw UnsupportedGraph (replayed)
end
Probe-->>v2: "" (supported) OR reason string
v2-->>Caller: NVTE_F16_arbitrary_seqlen / NVTE_FP8 / NVTE_No_Backend + message
Note over Caller,cuDNN: Execution path (nvte_fused_attn_fwd_v2)
Caller->>v2: "re-probe with check_for_forward_support=true"
v2->>Cache: get_or_build_cached_graph (same key → HIT)
Cache-->>v2: cached entry
v2->>Cache: ensure_plans_built() [once per entry]
Cache->>cuDNN: build_plans()
cuDNN-->>Cache: compiled kernel
v2->>cuDNN: graph.execute(runtime tensors)
Reviews (42): 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}
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
|
/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 |
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
|
/te-ci L0 L1 L2 L3 |
for more information, see https://pre-commit.ci
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: