Skip to content

Commit 353c4f6

Browse files
janpfeiferjanpfeifer
authored
Merge pull request #41 from gomlx/add-gqa-fusedattention
Support Grouped Query Attention (GQA) and update flash attention options
2 parents 10d891e + 3b49a55 commit 353c4f6

3 files changed

Lines changed: 163 additions & 29 deletions

File tree

‎compute/xla/flash.go‎

Lines changed: 144 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@ package xla
55
import (
66
"encoding/json"
77
"maps"
8+
"math"
89
"strconv"
910

1011
"github.com/gomlx/compute"
@@ -57,11 +58,12 @@ type fmhaVariant struct {
5758
// Bias routing: cfg.Bias non-nil selects the fmhaScaleBias targets. cuDNN's ScaleBias kernel does
5859
// not accept seqlen operands, so bias+seqlens returns ErrNotImplemented and the caller falls back
5960
// to the decomposed path.
60-
func selectFMHAVariant(op string, qkvDType dtypes.DType, causal bool,
61+
func selectFMHAVariant(op string, qkvDType dtypes.DType,
6162
cfg *compute.ScaledDotProductAttentionConfig) (fmhaVariant, error) {
6263
var v fmhaVariant
6364
hasSeqLens := cfg != nil && cfg.QuerySeqLen != nil && cfg.KeyValueSeqLen != nil
6465
hasBias := cfg != nil && cfg.Bias != nil
66+
causal := cfg != nil && cfg.Causal
6567

6668
switch qkvDType {
6769
case dtypes.BFloat16:
@@ -262,8 +264,7 @@ func formatScale(scale float64) string {
262264
// f16/bf16 (fp8 paused), BSHD-layout, equal-head, on a cuda plugin. Causality and
263265
// per-batch seqlen padding are supported (mask_type derives from them in selectFMHAVariant);
264266
// an explicit materialized mask is not (use seqlens instead). Anything else -> ErrNotImplemented.
265-
func (f *Function) flashSupported(op string, qkvDType dtypes.DType, mask compute.Value, numHeads, numKVHeads, featureDim int, axesLayout compute.AxesLayout, causal bool, options *compute.ScaledDotProductAttentionConfig) error {
266-
_ = causal
267+
func (f *Function) flashSupported(op string, qkvDType dtypes.DType, mask compute.Value, numHeads, numKVHeads, featureDim int, axesLayout compute.AxesLayout, options *compute.ScaledDotProductAttentionConfig) error {
267268
if !f.builder.backend.plugin.IsCUDA() {
268269
return errors.Wrapf(compute.ErrNotImplemented, "%s: cuDNN flash needs the cuda plugin, have %q", op, f.builder.backend.pluginName)
269270
}
@@ -283,10 +284,10 @@ func (f *Function) flashSupported(op string, qkvDType dtypes.DType, mask compute
283284
return errors.Wrapf(compute.ErrNotImplemented,
284285
"%s: cuDNN flash path supports only BSHD or BHSD layouts (got layout %v)", op, axesLayout)
285286
}
286-
if numKVHeads != numHeads {
287+
if numHeads%numKVHeads != 0 {
287288
return errors.Wrapf(compute.ErrNotImplemented,
288-
"%s: cuDNN flash path requires equal number of q/kv heads (got layout=%v heads=%d/%d)",
289-
op, axesLayout, numHeads, numKVHeads)
289+
"%s: cuDNN flash path requires query heads (%d) to be a multiple of kv heads (%d)",
290+
op, numHeads, numKVHeads)
290291
}
291292
// One of QuerySeqLen/KeyValueSeqLen set without the other is ambiguous.
292293
if options != nil && (options.QuerySeqLen != nil) != (options.KeyValueSeqLen != nil) {
@@ -374,8 +375,12 @@ var bwdResultLayouts = [][]int{{3, 1, 2, 0}, {3, 1, 2, 0}, {3, 1, 2, 0}, {0}}
374375
// (BSHD), bf16. It returns the [B,S,H,D] bf16 output and the [B,H,S] f32 softmax statistics
375376
// (log-sum-exp) the flash backward needs. The [B,H,S,S] scores never materialize. On non-cuda
376377
// plugins or unsupported option combinations it returns ErrNotImplemented.
377-
func (f *Function) FusedScaledDotProductAttention(query, key, value, mask compute.Value, _, _ int, axesLayout compute.AxesLayout, scale float64, causal bool, options *compute.ScaledDotProductAttentionConfig) (output compute.Value, statesForVJP []compute.Value, err error) {
378+
func (f *Function) FusedScaledDotProductAttention(query, key, value compute.Value, axesLayout compute.AxesLayout, options *compute.ScaledDotProductAttentionConfig) (output compute.Value, statesForVJP []compute.Value, err error) {
378379
op := compute.OpTypeFusedScaledDotProductAttention.String()
380+
var mask compute.Value
381+
if options != nil {
382+
mask = options.Mask
383+
}
379384
batchSize, seqLen, numHeads, featureDim, err := f.bshdDims(op, query, axesLayout)
380385
if err != nil {
381386
return nil, nil, err
@@ -388,16 +393,35 @@ func (f *Function) FusedScaledDotProductAttention(query, key, value, mask comput
388393
if err != nil {
389394
return nil, nil, err
390395
}
391-
if err = f.flashSupported(op, qDType, mask, numHeads, numKVHeads, featureDim, axesLayout, causal, options); err != nil {
396+
if err = f.flashSupported(op, qDType, mask, numHeads, numKVHeads, featureDim, axesLayout, options); err != nil {
392397
return nil, nil, err
393398
}
394-
variant, err := selectFMHAVariant(op, qDType, causal, options)
399+
variant, err := selectFMHAVariant(op, qDType, options)
395400
if err != nil {
396401
return nil, nil, err
397402
}
403+
var scale float64
404+
if options != nil {
405+
scale = options.Scale
406+
}
407+
if scale == 0 {
408+
scale = 1.0 / math.Sqrt(float64(featureDim))
409+
}
410+
broadcastedKey := key
411+
broadcastedValue := value
412+
if numHeads != numKVHeads {
413+
broadcastedKey, err = f.broadcastGQA(key, numHeads, numKVHeads, axesLayout)
414+
if err != nil {
415+
return nil, nil, err
416+
}
417+
broadcastedValue, err = f.broadcastGQA(value, numHeads, numKVHeads, axesLayout)
418+
if err != nil {
419+
return nil, nil, err
420+
}
421+
}
398422
// Operand order cuDNN expects: q, k, v, [bias], [seqQ, seqKV]. Bias goes before seqlens
399423
// (bias and seqlens are mutually exclusive here; see selectFMHAVariant).
400-
operands := []compute.Value{query, key, value}
424+
operands := []compute.Value{query, broadcastedKey, broadcastedValue}
401425
operandLayouts := [][]int{{3, 2, 1, 0}, {3, 2, 1, 0}, {3, 2, 1, 0}}
402426
if variant.hasBias {
403427
if err = validateBias("Bias", options.Bias, qDType, batchSize, numHeads, seqLen, seqLen); err != nil {
@@ -446,8 +470,12 @@ func (f *Function) FusedScaledDotProductAttention(query, key, value, mask comput
446470
// (statesForVJP[0], from the forward) plus the forward output and the output gradient dOutput into
447471
// the cuDNN backward custom-call, so the [B,H,S,S] scores never materialize in the backward either.
448472
// Returns dQuery, dKey, dValue as [B,S,H,D] bf16.
449-
func (f *Function) FusedScaledDotProductAttentionVJP(query, key, value, mask compute.Value, _, _ int, axesLayout compute.AxesLayout, scale float64, causal bool, options *compute.ScaledDotProductAttentionConfig, output compute.Value, statesForVJP []compute.Value, dOutput compute.Value) (dQuery, dKey, dValue compute.Value, err error) {
473+
func (f *Function) FusedScaledDotProductAttentionVJP(query, key, value compute.Value, axesLayout compute.AxesLayout, options *compute.ScaledDotProductAttentionConfig, output compute.Value, statesForVJP []compute.Value, dOutput compute.Value) (dQuery, dKey, dValue compute.Value, err error) {
450474
op := compute.OpTypeFusedScaledDotProductAttentionVJP.String()
475+
var mask compute.Value
476+
if options != nil {
477+
mask = options.Mask
478+
}
451479
batchSize, seqLen, numHeads, featureDim, err := f.bshdDims(op, query, axesLayout)
452480
if err != nil {
453481
return nil, nil, nil, err
@@ -464,16 +492,35 @@ func (f *Function) FusedScaledDotProductAttentionVJP(query, key, value, mask com
464492
if err != nil {
465493
return nil, nil, nil, err
466494
}
467-
if err = f.flashSupported(op, qDType, mask, numHeads, numKVHeads, featureDim, axesLayout, causal, options); err != nil {
495+
if err = f.flashSupported(op, qDType, mask, numHeads, numKVHeads, featureDim, axesLayout, options); err != nil {
468496
return nil, nil, nil, err
469497
}
470-
variant, err := selectFMHAVariant(op, qDType, causal, options)
498+
variant, err := selectFMHAVariant(op, qDType, options)
471499
if err != nil {
472500
return nil, nil, nil, err
473501
}
502+
var scale float64
503+
if options != nil {
504+
scale = options.Scale
505+
}
506+
if scale == 0 {
507+
scale = 1.0 / math.Sqrt(float64(featureDim))
508+
}
509+
broadcastedKey := key
510+
broadcastedValue := value
511+
if numHeads != numKVHeads {
512+
broadcastedKey, err = f.broadcastGQA(key, numHeads, numKVHeads, axesLayout)
513+
if err != nil {
514+
return nil, nil, nil, err
515+
}
516+
broadcastedValue, err = f.broadcastGQA(value, numHeads, numKVHeads, axesLayout)
517+
if err != nil {
518+
return nil, nil, nil, err
519+
}
520+
}
474521
// Operand order cuDNN expects: q, k, v, softmax_sum, dO, [bias], O. Bias goes at index 5,
475522
// before O (bias and seqlens are mutually exclusive; see selectFMHAVariant).
476-
operands := []compute.Value{query, key, value, softmaxStats, dOutput}
523+
operands := []compute.Value{query, broadcastedKey, broadcastedValue, softmaxStats, dOutput}
477524
operandLayouts := [][]int{{3, 2, 1, 0}, {3, 2, 1, 0}, {3, 2, 1, 0}, {2, 1, 0}, {3, 2, 1, 0}}
478525
if variant.hasBias {
479526
if err = validateBias("Bias", options.Bias, qDType, batchSize, numHeads, seqLen, seqLen); err != nil {
@@ -532,5 +579,88 @@ func (f *Function) FusedScaledDotProductAttentionVJP(query, key, value, mask com
532579
dKey = grads[1]
533580
dValue = grads[2]
534581
}
582+
if numHeads != numKVHeads {
583+
dKey, err = f.reduceSumGQA(dKey, numHeads, numKVHeads, axesLayout)
584+
if err != nil {
585+
return nil, nil, nil, err
586+
}
587+
dValue, err = f.reduceSumGQA(dValue, numHeads, numKVHeads, axesLayout)
588+
if err != nil {
589+
return nil, nil, nil, err
590+
}
591+
}
535592
return dQuery, dKey, dValue, nil
536593
}
594+
595+
func (f *Function) broadcastGQA(x compute.Value, numQueryHeads, numKVHeads int, layout compute.AxesLayout) (compute.Value, error) {
596+
groupSize := numQueryHeads / numKVHeads
597+
if groupSize == 1 {
598+
return x, nil
599+
}
600+
shape, err := f.Shape(x)
601+
if err != nil {
602+
return nil, err
603+
}
604+
dtype := shape.DType
605+
dims := shape.Dimensions
606+
if len(dims) != 4 {
607+
return nil, errors.Errorf("broadcastGQA: expected rank-4 tensor, got rank-%d (shape %v)", len(dims), dims)
608+
}
609+
610+
if layout == compute.AxesLayoutBSHD {
611+
b, s, _, d := dims[0], dims[1], dims[2], dims[3]
612+
reshaped, err := f.Reshape(x, b, s, numKVHeads, 1, d)
613+
if err != nil {
614+
return nil, err
615+
}
616+
broadcastShape := shapes.Make(dtype, b, s, numKVHeads, groupSize, d)
617+
broadcasted, err := f.BroadcastInDim(reshaped, broadcastShape, []int{0, 1, 2, 3, 4})
618+
if err != nil {
619+
return nil, err
620+
}
621+
return f.Reshape(broadcasted, b, s, numQueryHeads, d)
622+
} else {
623+
b, _, s, d := dims[0], dims[1], dims[2], dims[3]
624+
reshaped, err := f.Reshape(x, b, numKVHeads, 1, s, d)
625+
if err != nil {
626+
return nil, err
627+
}
628+
broadcastShape := shapes.Make(dtype, b, numKVHeads, groupSize, s, d)
629+
broadcasted, err := f.BroadcastInDim(reshaped, broadcastShape, []int{0, 1, 2, 3, 4})
630+
if err != nil {
631+
return nil, err
632+
}
633+
return f.Reshape(broadcasted, b, numQueryHeads, s, d)
634+
}
635+
}
636+
637+
func (f *Function) reduceSumGQA(dx compute.Value, numQueryHeads, numKVHeads int, layout compute.AxesLayout) (compute.Value, error) {
638+
groupSize := numQueryHeads / numKVHeads
639+
if groupSize == 1 {
640+
return dx, nil
641+
}
642+
shape, err := f.Shape(dx)
643+
if err != nil {
644+
return nil, err
645+
}
646+
dims := shape.Dimensions
647+
if len(dims) != 4 {
648+
return nil, errors.Errorf("reduceSumGQA: expected rank-4 tensor, got rank-%d (shape %v)", len(dims), dims)
649+
}
650+
651+
if layout == compute.AxesLayoutBSHD {
652+
b, s, _, d := dims[0], dims[1], dims[2], dims[3]
653+
reshaped, err := f.Reshape(dx, b, s, numKVHeads, groupSize, d)
654+
if err != nil {
655+
return nil, err
656+
}
657+
return f.ReduceSum(reshaped, 3)
658+
} else {
659+
b, _, s, d := dims[0], dims[1], dims[2], dims[3]
660+
reshaped, err := f.Reshape(dx, b, numKVHeads, groupSize, s, d)
661+
if err != nil {
662+
return nil, err
663+
}
664+
return f.ReduceSum(reshaped, 2)
665+
}
666+
}

‎compute/xla/flash_dispatch_test.go‎

Lines changed: 13 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,7 @@ import (
1414
)
1515

1616
func TestSelectFMHAVariant_StandardCausal(t *testing.T) {
17-
v, err := selectFMHAVariant("op", dtypes.BFloat16, true, nil)
17+
v, err := selectFMHAVariant("op", dtypes.BFloat16, &compute.ScaledDotProductAttentionConfig{Causal: true})
1818
if err != nil {
1919
t.Fatalf("unexpected err: %v", err)
2020
}
@@ -27,7 +27,7 @@ func TestSelectFMHAVariant_StandardCausal(t *testing.T) {
2727
}
2828

2929
func TestSelectFMHAVariant_NoMaskWhenNotCausal(t *testing.T) {
30-
v, err := selectFMHAVariant("op", dtypes.Float16, false, nil)
30+
v, err := selectFMHAVariant("op", dtypes.Float16, nil)
3131
if err != nil {
3232
t.Fatalf("unexpected err: %v", err)
3333
}
@@ -37,7 +37,7 @@ func TestSelectFMHAVariant_NoMaskWhenNotCausal(t *testing.T) {
3737
}
3838

3939
func TestSelectFMHAVariant_RejectsF32(t *testing.T) {
40-
_, err := selectFMHAVariant("op", dtypes.Float32, true, nil)
40+
_, err := selectFMHAVariant("op", dtypes.Float32, &compute.ScaledDotProductAttentionConfig{Causal: true})
4141
if !errors.Is(err, compute.ErrNotImplemented) {
4242
t.Errorf("err = %v, want ErrNotImplemented", err)
4343
}
@@ -46,9 +46,9 @@ func TestSelectFMHAVariant_RejectsF32(t *testing.T) {
4646
// TestSelectFMHAVariant_FP8NotImplemented pins the fp8-paused seam: fp8 dtypes must return
4747
// ErrNotImplemented so the caller falls back to the decomposed path. CPU, Mac-runnable.
4848
func TestSelectFMHAVariant_FP8NotImplemented(t *testing.T) {
49-
_, err := selectFMHAVariant("fmha", dtypes.F8E4M3FN, true, nil)
49+
_, err := selectFMHAVariant("fmha", dtypes.F8E4M3FN, &compute.ScaledDotProductAttentionConfig{Causal: true})
5050
require.True(t, compute.IsNotImplemented(err), "fp8 must be NotImplemented (paused), got %v", err)
51-
_, err = selectFMHAVariant("fmha", dtypes.F8E5M2, true, nil)
51+
_, err = selectFMHAVariant("fmha", dtypes.F8E5M2, &compute.ScaledDotProductAttentionConfig{Causal: true})
5252
require.True(t, compute.IsNotImplemented(err), "fp8 e5m2 must be NotImplemented (paused), got %v", err)
5353
}
5454

@@ -63,7 +63,7 @@ func TestSelectFMHAVariant_SeqLenPadding(t *testing.T) {
6363
KeyValueSeqLen: sent,
6464
}
6565

66-
v, err := selectFMHAVariant("op", dtypes.BFloat16, false, cfgBoth)
66+
v, err := selectFMHAVariant("op", dtypes.BFloat16, cfgBoth)
6767
if err != nil {
6868
t.Fatalf("unexpected err: %v", err)
6969
}
@@ -74,7 +74,8 @@ func TestSelectFMHAVariant_SeqLenPadding(t *testing.T) {
7474
t.Errorf("hasSeqLens = false, want true")
7575
}
7676

77-
v, err = selectFMHAVariant("op", dtypes.BFloat16, true, cfgBoth)
77+
cfgBoth.Causal = true
78+
v, err = selectFMHAVariant("op", dtypes.BFloat16, cfgBoth)
7879
if err != nil {
7980
t.Fatalf("unexpected err: %v", err)
8081
}
@@ -83,7 +84,7 @@ func TestSelectFMHAVariant_SeqLenPadding(t *testing.T) {
8384
}
8485

8586
// nil cfg still routes causal -> CAUSAL (no regression).
86-
v, err = selectFMHAVariant("op", dtypes.BFloat16, true, nil)
87+
v, err = selectFMHAVariant("op", dtypes.BFloat16, &compute.ScaledDotProductAttentionConfig{Causal: true})
8788
if err != nil {
8889
t.Fatalf("unexpected err: %v", err)
8990
}
@@ -98,7 +99,8 @@ func TestSelectFMHAVariant_Bias(t *testing.T) {
9899
var sent compute.Value = struct{}{}
99100
cfg := &compute.ScaledDotProductAttentionConfig{Bias: sent}
100101

101-
v, err := selectFMHAVariant("op", dtypes.BFloat16, true, cfg)
102+
cfg.Causal = true
103+
v, err := selectFMHAVariant("op", dtypes.BFloat16, cfg)
102104
if err != nil {
103105
t.Fatalf("unexpected err: %v", err)
104106
}
@@ -119,7 +121,7 @@ func TestSelectFMHAVariant_BiasSeqLensNotImplemented(t *testing.T) {
119121
var sent compute.Value = struct{}{}
120122
cfg := &compute.ScaledDotProductAttentionConfig{Bias: sent, QuerySeqLen: sent, KeyValueSeqLen: sent}
121123

122-
_, err := selectFMHAVariant("op", dtypes.BFloat16, false, cfg)
124+
_, err := selectFMHAVariant("op", dtypes.BFloat16, cfg)
123125
if !compute.IsNotImplemented(err) {
124126
t.Errorf("err = %v, want ErrNotImplemented", err)
125127
}
@@ -194,7 +196,7 @@ func TestFlashBackendConfigV_ElementTypeFromDtype(t *testing.T) {
194196
{dtypes.BFloat16, "BF16"},
195197
{dtypes.Float16, "F16"},
196198
} {
197-
v, err := selectFMHAVariant("op", tc.dtype, true, nil)
199+
v, err := selectFMHAVariant("op", tc.dtype, &compute.ScaledDotProductAttentionConfig{Causal: true})
198200
require.NoError(t, err)
199201
require.Equal(t, tc.want, v.elementType, "elementType for %s", tc.dtype)
200202
cfg := strings.ReplaceAll(flashBackendConfigV(1, 1, 8, 1.0, map[string]any{"x": 1}, v), " ", "")

‎compute/xla/flash_seqlen_cuda_test.go‎

Lines changed: 6 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -39,8 +39,8 @@ func getFusionBackend(t *testing.T) *Backend {
3939
if err != nil {
4040
t.Fatalf("probe: parameter v: %v", err)
4141
}
42-
out, _, err := fn.FusedScaledDotProductAttention(q, k, v, nil, H, H,
43-
compute.AxesLayoutBSHD, 1.0/math.Sqrt(float64(D)), true, nil)
42+
out, _, err := fn.FusedScaledDotProductAttention(q, k, v,
43+
compute.AxesLayoutBSHD, &compute.ScaledDotProductAttentionConfig{Scale: 1.0/math.Sqrt(float64(D)), Causal: true})
4444
if compute.IsNotImplemented(err) {
4545
t.Skipf("[cuda] cuDNN fMHA not supported on this host: %v", err)
4646
}
@@ -166,8 +166,10 @@ func TestFMHA_SeqLenPaddingCausal_cuda(t *testing.T) {
166166
QuerySeqLen: qSeq,
167167
KeyValueSeqLen: kvSeq,
168168
}
169-
out, _, err := fn.FusedScaledDotProductAttention(q, k, v, nil, H, H,
170-
compute.AxesLayoutBSHD, scale, true, cfg)
169+
cfg.Scale = scale
170+
cfg.Causal = true
171+
out, _, err := fn.FusedScaledDotProductAttention(q, k, v,
172+
compute.AxesLayoutBSHD, cfg)
171173
if err != nil {
172174
t.Fatalf("FusedScaledDotProductAttention: %v", err)
173175
}

0 commit comments

Comments
 (0)