fused_ops: typed flash backward (FusedSDPA + VJP), not a generic CustomCall by guygrigsby · Pull Request #13 · gomlx/compute

guygrigsby · 2026-06-23T20:19:03Z

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 (compute/xla: cuDNN flash as typed FusedSDPA + VJP 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.

CustomCall is a backend-specific escape hatch to a target the backend recognizes by name (e.g. cuDNN flash attention, "__cudnn$fmhaSoftmax"), bypassing the normal op set. It is multi-output. Backends without custom-call support return ErrNotImplemented so callers can fall back to a decomposed implementation.

janpfeifer · 2026-06-24T06:04:13Z

So, while I expected compute.Backend to eventually hold escape hatches for specific backends, I'm thinking we don't need one in this case: can we implement the FusedScaledDotProductAttention (internally using a CustomCallV2) in the xla backend when it is a CUDA plugin ? You can change the signature to return the extra softmax. The go-xla implementation can always return a NotImplemented error if some option is not supported, and core.Attention should fallback to use the decomposed version.

Also, define the Backend.FusedScaledDotProductAttentionVJP (for the back-propagation operation).

In case we actually wanted the XLA escape hatch, I would suggest creating it on a separate xla.go file, with type specifications custom to the xla backend, and prefix the function with XLA or something (XLACustomCallV2?). But I'm thinking we don't need it here: the logic/knowledge of specific custom operations should not (if possible) bubble up to the GoMLX level, this mapping ("__cudnn$fmhaSoftmax" implements FusedSPDA/FlashAttention) is better left in the go-xla package.

…e hatch Drop CustomCall/CustomCallSpec from the Function interface (the cuDNN target mapping belongs in the xla backend, not the cross-backend surface). Instead: - FusedScaledDotProductAttention now returns (output, softmaxStats, err); softmaxStats is optional (nil) and carries the log-sum-exp the fused backward needs. - new FusedScaledDotProductAttentionVJP(query,key,value,mask,...,output,softmaxStats,dOutput) -> (dQuery,dKey,dValue); ErrNotImplemented when the backend has no fused backward, so the caller differentiates the decomposed attention. Regenerated opsregistration; hand-stubbed the multi-return notimplemented methods (the generator only templates (Value,error)). Go backend returns nil stats (decomposed gradient).

janpfeifer · 2026-06-28T06:12:33Z

The PR looks all good to me. I'll wait for your thoughts on extending the FusedSDPA to accept optional keyValueSeqLen and querySeqLen (or something similar), in which case this could be done on the same PR.

…p loop limits - equalOptions: drop Value-typed QuerySeqLen/KeyValueSeqLen from == comparison (non-comparable concrete types panic; input operand identity handles dedup) - execFusedScaledDotProductAttention: use comma-ok type assertion for seqlen tensors; return clear error on wrong dtype or mismatched length vs batch - sdpaMultiHeadGeneric: clamp per-batch qLimit into [0, seqLen] and kvLimit into [0, kvLen] before using as loop bounds - fused_ops.go: clarify ScaledDotProductAttentionConfig doc: correctness fields require ErrNotImplemented when unsupported; optimization hints may be ignored Tests: TestSDPADirect_OptionEqualityNoValuePanic, TestSDPADirect_SeqLenWrongDtype, TestSDPADirect_SeqLenClamped (7/7 pass)

Stage 1: seqlen config fields + CPU reference + fp8 gate

guygrigsby · 2026-07-01T02:29:31Z

Stage 1 of the fused-attention work just landed here. Terminology: Stage 1 is the config fields and CPU reference for sequence-length padding masks, this PR. Stage 2 is the additive-bias config field and reference, a separate cumulative PR: #15.

What the review asked for:

FusedSDPA and its VJP live in the xla backend, not a generic escape hatch. The new parameters ride inside *ScaledDotProductAttentionConfig so the FusedOps method signatures do not change.
QuerySeqLen / KeyValueSeqLen added; the go-backend reference implements the sequence-length padding mask combined with causal. The go-backend SDPA capability stays false on purpose (the SIMD matmul plus decomposed path is faster), so the reference is exercised by direct unit tests rather than the capability-gated backend.
fp8 input dtype returns ErrNotImplemented (paused, no local hardware). The bias config field is Stage 2 (Stage 2: additive attention bias (cumulative on #13) #15).

janpfeifer

Thanks @guygrigsby !!

There is one larger request: softmaxStats Value -> statesForVJP []Value. I think this will be more stable if other fusion SDPA are implemented (at least for the Go backend there will be a few with SIMD support).

And then a few smaller ones, that if you are changing the API, maybe it's a good time to do it.

Let me know what you think, and if it's ok.

janpfeifer · 2026-07-01T04:30:44Z

+	//     (log-sum-exp), shape [batch, numHeads, seqLen] f32, which the VJP needs. Backends
+	//     without a fused backward return nil here and ErrNotImplemented from the VJP, so the
+	//     caller differentiates the decomposed attention instead.
 	FusedScaledDotProductAttention(


Since we are changing this a couple of suggestions (but I can do them later):

I think the numHeads and numKVHeads are redundant, since they can be inferred from the shapes, given the axesLayout. Maybe we just remove them ? (the Go backend can trivially be changed to read these from the shape, and the XLA doesn't need them I think)

Move mask, causal, scale and eventually bias to the Scaled.DotProductAttentionConfig since they are also optional. And make it such that scale == 0 (the default) means no scale.

Deferring these to a follow-up PR to keep the statesForVJP refactor isolated: drop numHeads/numKVHeads (infer from shape + axesLayout), and move mask/causal/scale/bias into ScaledDotProductAttentionConfig with scale == 0 meaning the default. No behavior change intended, just the ergonomics pass you suggested.

janpfeifer · 2026-07-01T05:36:08Z

+
+	// FusedScaledDotProductAttentionVJP computes the gradients (dQuery, dKey, dValue) of
+	// FusedScaledDotProductAttention given the forward output, the softmaxStats it returned, and
+	// the output gradient dOutput. The query/key/value/mask and numHeads..options parameters are


Nit: "the adjoint output gradient dOutput" (that's how I've been calling the V in VJP, following the convention I found in the related literature).

Renamed: the VJP doc now calls it "the adjoint output gradient dOutput". guygrigsby@e815ee8

janpfeifer · 2026-07-01T05:46:38Z

-	// Output: same shape as query.
+	// Outputs:
+	//   - output: same shape as query.
+	//   - softmaxStats: optional (may be nil). When the backend supports a fused backward


Apologies, but I just realized something: the softmaxStats is very specific to the cudnn implementation. Actually, I'm surprised it also doesn't want the max-exp of the attention logits, since having it could speed-up the online softmax, which I assume it re-calculates in the VJP (?).

But, I was going to suggest to change softmaxStats Value to statesForVJP []Value. And the FusedScaledDotProductAttentionVJP would take that as input.

I think the cudnn implementation works as a gradient checkpoint, it regenerates some forward values to save temporary memory. But other implementations (for CPU, where memory is cheaper) may find a different balance and decide to keep more state in temporary memory to avoid having to regenerate them on the VJP. Hence the change from a single Value to a []Value.

Actually, I asked the AI, it mentions that some version of the cudnn FMHA also requires a byte tensor called workspace as additional output that needs to be passed back to the backward (VJP) function.

Any thoughts ?

Double checked with the AI and looks like this is the right call. I'll add statesForVJP []Value in place of softmaxStats Value.

Done: softmaxStats Value -> statesForVJP []Value on the forward return and the VJP input, threaded through the interface, the notimplemented default, and the go backend reference. An empty slice keeps the old nil meaning (no fused backward -> decomposed VJP). guygrigsby@502e5de

FusedScaledDotProductAttention/VJP hardcoded a single cuDNN-shaped softmax-stats tensor. Widen it to a slice so other backends (and cuDNN FMHA variants needing a workspace tensor) can carry whatever backward state they need. nil/empty slice keeps the old "no fused backward" meaning. Regenerated gen_opsregistration.go via go generate ./internal/gobackend/...; no behavior change.

… nit)

janpfeifer

Thank you Guy! All looks good. Going to review the other 2 PRs and I'll merge when those are also ready!

…aramters.

…Attention tests.

janpfeifer · 2026-07-03T13:22:06Z

Following up on the fixes in gomlx/gomlx#427 I did a few changes here:

Added the missing OpType constant for the Fused VJP.
Skip fused ops tests if they return a "not implemented" error, in the generic support/backendtest.
Added BFloat16 and Float16 tests for Fused SDPA tests (also with featureDim=8, required) in the support/backendtest that tests XLA CUDA.

This allows the 3 PRs to be committed.

Stage 2 of the fused-attention work: additive attention bias. This is cumulative on top of Stage 1 (#13), which is the prerequisite, so the Stage 1 commits show here too until #13 merges to main, after which this shrinks to just the bias delta. Adds ScaledDotProductAttentionConfig.Bias plus the go-backend CPU reference: bias added to the scores before softmax, combining with causal and seqlen masking. The bias is deduped via operand identity (equality compares the presence flag, not the Value), validated for shape (clear error, not a panic). Capability stays off; exercised by direct unit tests. Validated on an RTX 3070 Ti through the go-xla and gomlx bias PRs: cuDNN ScaleBias applies the bias forward and backward.

This was referenced Jun 23, 2026

compute/xla: cuDNN flash as typed FusedSDPA + VJP gomlx/go-xla#37

Merged

attention: flash via typed FusedSDPA + backend VJP gomlx/gomlx#427

Merged

guygrigsby mentioned this pull request Jun 24, 2026

Expose OptimizationBarrier for gradient checkpointing / rematerialization gomlx/gomlx#422

Closed

guygrigsby changed the title ~~Add CustomCall to the backend Function interface~~ fused_ops: typed flash backward (FusedSDPA + VJP), not a generic CustomCall Jun 24, 2026

janpfeifer and others added 7 commits June 28, 2026 08:26

Merged with main.

725aad0

fused_ops: add seqlen fields to attention config and direct test harness

8940377

fused_ops: implement seqlen padding mask in go backend

67997e2

fused_ops: zero scores tail when seqlen masking restricts kvLen

e8b8dbe

fused_ops: return ErrNotImplemented for float8 attention dtype

45cebd1

Merge pull request #1 from guygrigsby/stage1-seqlen-config

d69f2d0

Stage 1: seqlen config fields + CPU reference + fp8 gate

This was referenced Jul 1, 2026

Stage 2: additive attention bias (cumulative on #13) #15

Merged

Stage 2: additive attention bias (config field + CPU reference) guygrigsby/compute#4

Closed

janpfeifer requested changes Jul 1, 2026

View reviewed changes

guygrigsby added 2 commits July 2, 2026 10:54

fused_ops: call dOutput the "adjoint output gradient" in VJP doc (jan…

e815ee8

… nit)

janpfeifer approved these changes Jul 3, 2026

View reviewed changes

janpfeifer added 4 commits July 3, 2026 12:15

Added OpTypeFusedScaledDotProductAttentionVJP.

8eb259d

Skip fused ops tests when they are not implemented for the specific p…

20331a9

…aramters.

Added BF16 with featureDim=8 tests for XLA CUDA FusedScaledDotProduct…

0058d22

…Attention tests.

Added Float16 tests for Fused SDPA.

e738459

janpfeifer merged commit 8164cab into gomlx:main Jul 3, 2026

Uh oh!

Conversation

guygrigsby commented Jun 23, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

janpfeifer commented Jun 24, 2026

Uh oh!

janpfeifer commented Jun 28, 2026

Uh oh!

guygrigsby commented Jul 1, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

janpfeifer left a comment

Choose a reason for hiding this comment

Uh oh!

janpfeifer Jul 1, 2026

Choose a reason for hiding this comment

Uh oh!

guygrigsby Jul 2, 2026

Choose a reason for hiding this comment

Uh oh!

janpfeifer Jul 1, 2026

Choose a reason for hiding this comment

Uh oh!

guygrigsby Jul 2, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

janpfeifer Jul 1, 2026

Choose a reason for hiding this comment

Uh oh!

guygrigsby Jul 1, 2026

Choose a reason for hiding this comment

Uh oh!

guygrigsby Jul 2, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

janpfeifer left a comment

Choose a reason for hiding this comment

Uh oh!

janpfeifer commented Jul 3, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

guygrigsby commented Jun 23, 2026 •

edited

Loading

guygrigsby commented Jul 1, 2026 •

edited

Loading

guygrigsby Jul 2, 2026 •

edited

Loading

guygrigsby Jul 2, 2026 •

edited

Loading