[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 151 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 SummaryThe PR refactors fused-attention backend selection and execution around opaque configuration/parameter handles and cuDNN-frontend support checks.
Confidence Score: 5/5The PR appears safe to merge because no blocking failure remains established at the current head. No blocking failure remains. Important Files Changed
Sequence DiagramsequenceDiagram
participant Framework as PyTorch/JAX
participant Config as Opaque config/params
participant Selector as Backend selector
participant Cache as Process graph cache
participant cuDNN as cuDNN frontend
Framework->>Config: Populate request attributes
Config->>Selector: Derive normalized configuration
Selector->>Cache: Lookup device-specific graph key
alt cached supported graph
Cache-->>Selector: Reuse graph
else cache miss
Selector->>cuDNN: Build and check graph support
cuDNN-->>Cache: Cache supported graph
end
Selector-->>Framework: Backend and diagnostic
Framework->>Cache: Execute matching graph
Reviews (65): Last reviewed commit: "minor tweaks to common" | 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 L0 L1 L2 L3 |
| graph_cache_debug::record_hit_miss(backend, pass, LookupResult::Miss, key); | ||
|
|
||
| auto entry = std::make_shared<CacheEntry<GraphAndTensors>>(build()); | ||
| graph_cache_debug::record_create_graph(backend, pass); |
There was a problem hiding this comment.
I think recording unconditionally can cause a. mismatch right ?
record_cache_graph() runs before we know whether a given thread won insertion. With concurrent same-key misses, both builders increment CACHE_GRAPH, although only one entry is inserted and the losing graph is discarded.
IIUC CACHE_GRAPH is supposed to be number of graphs cached so in the above case it might increase that number by 1 incorrectly (this is relevant for what is reported in logs and does not affect the feature code)
Could we inspect insert(...).second under the mutex, then call record_cache_graph() outside the lock only for the winner? CREATE_GRAPH can still count both builds because it is creating those 2 graphs but only 1 is being cached.
There was a problem hiding this comment.
Fixed as you suggested: the insert result is captured under the mutex, the lock is
released, and record_cache_graph() fires outside the lock only when
result.second is true. CREATE_GRAPH still counts every build, since multiple graphs
really were built.
|
|
||
| // Per-thread counters, so the summary can break every column down by thread and backend. | ||
| struct ThreadCounters { | ||
| unsigned tid = 0; |
There was a problem hiding this comment.
The per thread structure has only one device, which should be okay for most cases, but I'm wondering if we have:
GPU0: 10 fwd events
GPU1: 4 fwd events
Summary might say: tid=2 dev=1 hit=14 which would imply that all events belong to GPU1 when then came from GPU 0
There was a problem hiding this comment.
Fixed. ThreadCounters now carries a multi_device flag set by a note_device() helper,
and the summary prints dev=mixed instead of naming a device when a thread's events
span more than one. The per-thread block are kept un-split by device, and the accurate per-device attribution can be found out by using level 2 diagnostics: the event lines do carry the live device ID.
| const char *fallback = nullptr; | ||
| try { | ||
| get_graph<kBackend, kPass, kCreateGraphFn>(cfg, handle); | ||
| return ""; | ||
| } catch (const std::exception &e) { | ||
| if (e.what()[0] != '\0') return e.what(); | ||
| fallback = "rejected without a reason."; | ||
| } catch (...) { | ||
| fallback = "unknown failure."; | ||
| } |
There was a problem hiding this comment.
This seems to treat every thrown exception as a support rejected which is okay for when cuDNN is unsupported but it would also potentially incorrectly report other failures like RT error, bad_alloc or some cuda error as No backend.
Would be good if we could classify this ?
There was a problem hiding this comment.
Fixed. Now query_support() throws a dedicated
detail::UnsupportedByCudnn, support_verdict() catches only that, and the catch (...)
is gone, so a non-empty return means cuDNN said no and anything genuinely broken surfaces
as an error.
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
…cord cache_graph after insert on winner thread 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 |
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>
|
/te-ci L0 L1 L2 L3 |
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>
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>
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 FP8supported_verdict_*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=<level>[:<ranks>]: 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;[:<ranks>]specifies select ranks for diagnostics; 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
test_fused_attn_graph_cacheto test graph cache's behavior withNVTE_FUSED_ATTN_CACHE_DEBUG=2on, andtest_fused_attn_backend_messageto test the surfacing of TE-specific or cuDNN-related error messages to users.Checklist: