compute/xla: cuDNN flash as typed FusedSDPA + VJP#37
Conversation
Exposes stablehlo.custom_call (the op was in the optype enum but had no builder). Sets call_target_name, api_version, raw backend_config, and pre-rendered operand/result layout attributes; multi-output for the result + workspace tuple. This is the primitive needed to emit the cuDNN fused-attention custom-calls (__cudnn$fmhaSoftmax / -Backward). Test emits the flash forward and checks the rendered StableHLO matches what JAX emits.
Emits __cudnn$fmhaSoftmax via stablehlo.CustomCall, compiles through PJRT, and runs it. Follows the existing convention: a normal test that skips unless -plugin cuda (so `go test ./...` on the cpu plugin skips it). With q=k=v=ones the uniform-softmax output must be all 1.0; verified on an RTX 3070 Ti (sm_86, cuDNN 9.23). Confirms XLA lowers our emitted custom_call to cuDNN flash. Run: go test -plugin cuda -run TestFMHA ./pjrt/
Chains the training forward (which also returns the softmax stats f32[B,H,S]) through a transpose into __cudnn$fmhaSoftmaxBackward, all in one compiled graph, and runs it: dQ/dK/dV come back finite. Confirms XLA lowers the flash backward we emit. Plus buffer/exec cleanup so the forward and backward tests don't starve each other's device memory when run in the same process. Both skip unless -plugin cuda.
Wires the backend Function.CustomCall through to stablehlo.CustomCall, converting the compute output shapes to the StableHLO shape type. Multi-output: one result Value per spec.OutputShapes entry. This is what lets the gomlx graph layer emit cuDNN flash-attention custom_calls on the XLA backend.
A __cudnn$ custom_call only lowers on the cuda plugin; on cpu/rocm it built fine and then failed at compile time with an unknown-target error. Return ErrNotImplemented at build time instead, so callers (flash attention) fall back to a decomposed path rather than hitting an uncatchable compile failure.
| // output plus a scratch workspace buffer). | ||
| // | ||
| // Returns one output Value per outputShape, in order. | ||
| func CustomCall(target string, apiVersion int, backendConfig, operandLayouts, resultLayouts string, |
There was a problem hiding this comment.
What about the following API instead:
func CustomCallV2(
target string,
backendConfig string,
operands []*Value, // Grouped inputs
operandLayouts [][]int,
outputShapes []shapes.Shape, // Grouped outputs
outputLayouts [][]int,
) ([]*Value, error) {
// ...
}- API version hard-coded in the signature: the current V4 has a different signature, it cannot easily be the same function (https://openxla.org/xla/custom_call).
- Operands come before outputs.
- Layouts inlined with paired operand/result -- they are coupled by the length if defined.
- Layout as a list of int: if they are nil, they default to row-major, that is, a decreasing list of ints with the length of the corresponding parameter's rank.
Example call:
out, err := CustomCallV2(
"__cudnn$fmhaSoftmax", cfg,
[]*Value{q, k, v}, nil, // inputs and their layouts
outShapes, nil, // outputs and their layouts
)There was a problem hiding this comment.
Adopted this shape in 0015b77. CustomCallV2(target, backendConfig, operands []*Value, operandLayouts [][]int, outputShapes []shapes.Shape, outputLayouts [][]int) ([]*Value, error) — operands before outputs, layouts paired by position, nil layout defaults to row-major. Version stays out of the signature (V2, not a hard-coded V4). The nil-slice vs nil-entry default is documented in 8bd845e.
| // 1.0 (the weighted sum of all-ones values), independent of the causal mask. Needs | ||
| // the cuda plugin (cuDNN); skipped otherwise. | ||
| func TestFMHAForwardExecute(t *testing.T) { | ||
| if *FlagPluginName != "cuda" { |
There was a problem hiding this comment.
Instead, use client.Plugin().IsCUDA()
There was a problem hiding this comment.
Switched to client.Plugin().IsCUDA() in 9350689 (fmha tests also moved onto CustomCallV2 there).
|
From digging (the AI did the digging) into the Jax code I gathered there are the following custom calls:
A |
|
Also of interest, some performance metrics: jax-ml/jax#24934 |
…mCall Implement FusedScaledDotProductAttention (forward) and FusedScaledDotProductAttentionVJP (backward) in flash.go via the cuDNN __cudnn$fmhaSoftmax / __cudnn$fmhaSoftmaxBackward custom-calls. The backend_config + layout knowledge moves here from gomlx; it never leaves go-xla. The forward returns the f32 softmax stats (log-sum-exp) the VJP threads into the backward. Both return ErrNotImplemented off-cuda or for unsupported combos (non-causal, explicit mask, non-BSHD, unequal heads) so the graph layer differentiates decomposed. CustomCall drops from the public Function surface (interface no longer has it) and becomes a private customCall helper. FusedSDPA capability is gated on the cuda plugin at construction. The low-level stablehlo.CustomCall primitive and its tests are unchanged.
guygrigsby
changed the title
|
Thanks for the changes @guygrigsby ! A few initial comments:
Ugh ... this is a long list of asks, I apologize. Please, let me know if it's beyond what you have time to contribute, I can work on this next week if that is the case. Btw, I attempted to run a benchmark on other examples of sentence embedders and generators, but because of masking they never triggered the cuDNN custom calls. See last entries in the Sentence Embed benchmarks tab in this spreadsheet. |
|
Btw, I merged some changes I did in main into your branch -- I hope it's ok. You probably will need to pull those changes into your local clone. It includes adding more (and better) options support for the XLA backend, including the options Or, to enable internal XLA loggign (it uses an old tensorflow logging framework): |
I will review the list closely tonight or tomorrow. I don't think it will be too much. I don't mind. ❤️ |
flashSupported no longer rejects non-causal or missing seqlens: any combination of causal/seqlens is valid; selectFMHAVariant picks the mask_type (NO_MASK, CAUSAL, PADDING, PADDING_CAUSAL). Materialized mask and mismatched seqlens (one without the other) remain rejected. Add TestSelectFMHAVariant_SeqLenPadding (CPU, confirms PADDING / PADDING_CAUSAL routing and causal regression). Add TestFMHA_SeqLenPaddingCausal_cuda ([cuda], self-skips on Mac).
Replace 3 stablehlo.CustomCall sites with CustomCallV2 ([][]int layouts). Switch cuda guard from FlagPluginName string compare to client.Plugin().IsCUDA(); move skip after getPJRTClient so the plugin is available. Both tests compile and SKIP cleanly on CPU/Mac; cuDNN execution path unchanged for cuda host.
Add readFlatBF16 and assertSeqLenMasksOutput helpers. The latter compares the padding positions (s >= S/2) of the short batch element against the same positions of the full-length element: with q=k=v=ones, cuDNN PADDING_CAUSAL zeros padding-query outputs, so they must differ from the active (non-zero) outputs of batch 0. The test now fails if cuDNN ignores the seqlens. Finiteness assertion retained.
|
Reviewed the list and it'll take me a couple days between other things, but it's no problem. All your suggestions are good. I can handle the API changes, padding/masking, most variant work and |
…istinction compute/xla: validate seqlen operand dtype and shape before cuDNN custom-call validateSeqLen helper checks int32 rank-1 [B] for QuerySeqLen/KeyValueSeqLen in both the forward and VJP paths; wrong dtype or length now returns a clear error before the operand reaches XLA lowering. Add CPU-runnable unit tests covering wrong dtype, wrong rank, wrong length, non-Node value, and the happy path (TestValidateSeqLen).
Thanks!! No worry about the FP8, I can add it later. |
| if !f.builder.backend.plugin.IsCUDA() { | ||
| return errors.Wrapf(compute.ErrNotImplemented, "%s: cuDNN flash needs the cuda plugin, have %q", op, f.builder.backend.pluginName) | ||
| } | ||
| if !causal || mask != nil || axesLayout != compute.AxesLayoutBSHD || numKVHeads != numHeads { |
There was a problem hiding this comment.
nit: What about adding to the error exactly which things were not supported, so one can easily track down why some model can't use flash attention ? With the logging output I added earlier, one can output the failure to use a fusion op, and easily investigate whether/why their model is not using fusion (flash attention).
Also I'm missing the check for the dtypes supported (I think only dtypes.Float16 and dtype.BFloat16, which can be checked with dtype.IsHalfPrecision()).
There was a problem hiding this comment.
Done in the incoming update: flashSupported now returns a distinct ErrNotImplemented per failing condition (cuda plugin, dtype, materialized mask, non-BSHD/unequal heads, mismatched seqlens), and gates dtype via IsHalfPrecision (f16/bf16 only). Pairs with your V(2) fusion logging.
There was a problem hiding this comment.
Landed in 7ebc65e: flashSupported returns a distinct ErrNotImplemented reason per failing condition (non-cuda plugin, unsupported dtype, materialized mask, non-BSHD/unequal heads, mismatched seqlens), and gates dtype via IsHalfPrecision (f16/bf16 only). Pairs with your V(2) fusion logging for tracing why a model isn't fusing.
|
After a bit more investigation, the Jax But also, dropout in attention models is not usually used -- it doesn't help and it's pretty expensive (it regenerates the random mask all the time). In other words, probably we don't want to support dropout. In some future we may want to integrate NVIDIA Transformer Engine into GoMLX, also for adding supprot for Blackwells (>= RTX 5090) FP4 format, but it's lots of work, and I would expect eventually that support to be added to the standard JAX path. |
… clearer reason Per review: consolidate the 'why no fusion' reasons in flashSupported and add the missing dtype check (float16/bfloat16 via IsHalfPrecision) with an explicit message.
|
On dropout: agreed, and thanks for confirming from the JAX side. Cutting dropout (and bias+dropout) from scope for exactly that reason. fp8 is likewise paused (no local sm_8.9+ hardware). Both are left as NotImplemented seams that fall back to the decomposed path. |
Stage 1: CustomCallV2, fmha target seam, seqlen padding masks
|
Stage 1 of the fused-attention work just landed here. Terminology: Stage 1 is CustomCallV2, the config-driven fmha variant seams, and sequence-length padding masks, which is this PR. Stage 2 is additive attention bias via __cudnn$fmhaScaleBiasSoftmax, a separate cumulative PR: #39. What the review asked for:
|
…mentation to the PJRT plugin. Properly set the `FusedScaledDotProductAttentionVJP` capability.
…use json encoding to generate the backend config, to facilitate reconfigurability.
BackendConfig now generated from arbitrary `map[string]any`, easier to manipulate and configure.
Added support for Float16.
|
Following the updates to tests done in gomlx/gomlx/#427, I updated the implementation to match:
I'll merge this as is, because this frees you to move to the stage 2. I apologize for the main this may be merging into the other PR... But hopefully the refactored |
…omCall (#13) Per the design discussion in gomlx/gomlx#422: keep the cuDNN target mapping off the cross-backend surface and express flash attention as typed fused ops instead of a generic escape hatch. - `FusedScaledDotProductAttention` now returns `(output, softmaxStats, err)`. `softmaxStats` is optional (may be nil); a backend with a fused backward returns the log-sum-exp the backward needs, others return nil. - New `FusedScaledDotProductAttentionVJP(query, key, value, mask, ..., output, softmaxStats, dOutput) (dQuery, dKey, dValue, err)`. Returns `ErrNotImplemented` when the backend has no fused backward, so the graph layer differentiates the decomposed attention instead. - Removed the generic `CustomCall` / `CustomCallSpec` from the `Function` interface. The cuDNN custom-call mapping now lives in the xla backend (gomlx/go-xla#37), not on the shared interface. Regenerated opsregistration; the two multi-return methods are hand-stubbed in `notimplemented` (the generator only templates `(Value, error)`). The Go backend returns nil stats, so its attention gradient stays decomposed. Base of the flash stack: gomlx/go-xla#37 and gomlx/gomlx#427 build on it. Context: gomlx/gomlx#422. --------- Co-authored-by: Jan Pfeifer <pfeifer@gmail.com>
|
Thanks for the help. No worries about the other merges. I'll sort it out. |
Stage 2: cuDNN fmhaScaleBias variant (additive attention bias, cumulative on #37)
2 participants
Implements the typed flash ops from gomlx/compute#13 for the XLA backend, so the cuDNN custom-call mapping stays inside go-xla (per gomlx/gomlx#422):
FusedScaledDotProductAttention(forward) andFusedScaledDotProductAttentionVJP(backward) in a newflash.go, via the__cudnn$fmhaSoftmax/__cudnn$fmhaSoftmaxBackwardcustom-calls. The forward returns the f32 softmax stats (log-sum-exp); the VJP threads them into the backward, so neither pass materializes the[B,H,S,S]scores.IsCUDA). Off-cuda and unsupported combos (non-causal, explicit mask, non-BSHD, unequal heads) returnErrNotImplementedfor a clean fallback to decomposed attention.CustomCallis now a privatecustomCallhelper rather than a method on theFunctioninterface. The low-levelstablehlo.CustomCallemitter and its tests are unchanged.Depends on gomlx/compute#13.
Context: gomlx/gomlx#422.