[BUG] Metal scaled_dot_product_attention is ~300x less accurate for 2-8 query tokens at 33-128 stacked rows since #4723 (sdpa_split)

Chiusa
  • Qualcosa non va?
#4,808 2 commenti 0 reazioni 0 assegnatari Vedi su GitHub

I maintainer di solito rispondono entro 3 giorni

Nessuno ha ancora preso questa issue.

  • #1039 di @STRML — chiusa senza merge

Valutazione

Difficoltà
4/5
Tempo stimato
3-5 giorni
Idoneità per principianti
22/100
Tipo di issue
Bug
Chiarezza
Specificata chiaramente
Stato di attività
Ferma
Stack tecnologico
cpp
Ambito
machine-learning

Direzione di ricerca

Start with use_split in mlx/backend/metal/scaled_dot_product_attention.cpp, the branch that routes 33-128 stacked rows to sdpa_split. Build the mlx-c repro in the issue and run ./sdpa_repro 1000 causal 256, which currently prints about 0.035 against about 0.0001 on v0.32.3. Done means the error returns to about 1e-4 across the listed head dims and kv lengths, with the change verified on Metal. Open PR #1039 already targets this, so check it before starting.

Scritto dal modello di indicizzazione a partire dal testo della issue.

Descrizione

bug

Since #4723, mx.fast.scaled_dot_product_attention on Metal returns results about 300x less accurate when a call has 2 to 8 query tokens, 33 to 128 stacked rows (query tokens times query heads per KV head) and at least 128 keys. The error against a float64 reference goes from 1e-4 to 3e-2 to 4e-2, causal or unmasked. Switching use_split off restores the v0.32.3 numbers to the last digit.

rows = qL * (H / Hkv)   with qL <= 8:  33..128  -> sdpa_split (new in #4723)
                                        <= 32   -> sdpa_vector (unchanged, accurate)

Setup: bf16, 24 query heads over 2 KV heads (12 rows per query token), 4 query tokens, inputs uniform in [-0.5, 0.5], scale 1/16, max abs error of the output against a float64 reference on the CPU. Apple M5 Max, macOS 27.0.1, MLX built with NAX (deployment target 26.2).

head_dim kL mask v0.32.3 4ca8d5d5 (#4568, 159 commits later)
256 104 causal 4.1e-4 4.1e-4 (kL < 128, not split)
256 1004 causal 1.1e-4 3.5e-2
256 1004 none 9.3e-5 3.5e-2
256 2055 causal 8.5e-5 2.6e-2
64 1004 causal 1.0e-4 3.8e-2
80 1004 causal 1.2e-4 1.2e-4 (not split)
96 1004 causal 9.3e-5 3.4e-2
128 1004 causal 1.1e-4 3.3e-2
192 1004 causal 1.4e-4 4.3e-2

The call hits use_split in mlx/backend/metal/scaled_dot_product_attention.cpp. A model with 8 KV heads and 40 query heads (5 rows per token) reaches it at 7 or 8 verified tokens; one with 12 rows per token reaches it at 3 tokens, so speculative decoding verification is the case #4723 targets and the case it breaks.

I have not tested other chips, the non-NAX sdpa_split path, float16, or the Python binding. The repro below uses mlx-c.

// clang -O1 sdpa_repro.c -I <mlx-c include> -L <mlx lib> -lmlxc -Wl,-rpath,<mlx lib> -o sdpa_repro
// ./sdpa_repro <kv> [causal|none] [head_dim]
#include <math.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include "mlx/c/mlx.h"

static uint64_t state = 0x5D9A;
static float rnd(void) {
  state = state * 6364136223846793005ULL + 1442695040888963407ULL;
  return ((state >> 40) & 0xFFFFFF) / (float)0x1000000 - 0.5f;
}

// A bf16 array of uniform(-0.5, 0.5) values, plus the bf16-rounded values on the host.
static mlx_array make_bf16(int n, const int* shape, int dim, mlx_stream s, float** host) {
  float* d = malloc(n * sizeof(float));
  for (int i = 0; i < n; i++) d[i] = rnd();
  mlx_array f32 = mlx_array_new_data(d, shape, dim, MLX_FLOAT32);
  mlx_array out = mlx_array_new();
  mlx_astype(&out, f32, MLX_BFLOAT16, s);
  mlx_array back = mlx_array_new();
  mlx_astype(&back, out, MLX_FLOAT32, s);
  mlx_array_eval(back);
  *host = malloc(n * sizeof(float));
  memcpy(*host, mlx_array_data_float32(back), n * sizeof(float));
  mlx_array_free(f32);
  mlx_array_free(back);
  free(d);
  return out;
}

int main(int argc, char** argv) {
  const int H = 24, HKV = 2, L = 4;
  const int D = argc > 3 ? atoi(argv[3]) : 256;
  const int kv = argc > 1 ? atoi(argv[1]) : 1000;
  const int causal = !(argc > 2 && strcmp(argv[2], "none") == 0);
  const int KL = kv + L;
  mlx_stream s = mlx_default_gpu_stream_new();
  const int qs[4] = {1, H, L, D}, ks[4] = {1, HKV, KL, D};
  float *hq, *hk, *hv;
  mlx_array q = make_bf16(H * L * D, qs, 4, s, &hq);
  mlx_array k = make_bf16(HKV * KL * D, ks, 4, s, &hk);
  mlx_array v = make_bf16(HKV * KL * D, ks, 4, s, &hv);
  const float scale = 1.0f / 16;
  mlx_array none = {0};
  mlx_array out = mlx_array_new();
  if (mlx_fast_scaled_dot_product_attention(&out, q, k, v, scale, causal ? "causal" : "", none, none, false, s)) return 1;
  mlx_array o32 = mlx_array_new();
  mlx_astype(&o32, out, MLX_FLOAT32, s);
  mlx_array_eval(o32);
  const float* o = mlx_array_data_float32(o32);

  double maxd = 0;
  double* p = malloc(KL * sizeof(double));
  for (int h = 0; h < H; h++) {
    const int kvh = h / (H / HKV);
    for (int r = 0; r < L; r++) {
      const int lim = causal ? KL - L + r : KL - 1; // the row sees keys 0..lim (bottom-right causal)
      double mx = -1e300;
      for (int j = 0; j <= lim; j++) {
        double dot = 0;
        for (int d = 0; d < D; d++) dot += (double)hq[(h * L + r) * D + d] * (double)hk[(kvh * KL + j) * D + d];
        p[j] = dot * scale;
        if (p[j] > mx) mx = p[j];
      }
      double sum = 0;
      for (int j = 0; j <= lim; j++) {
        p[j] = exp(p[j] - mx);
        sum += p[j];
      }
      for (int d = 0; d < D; d++) {
        double acc = 0;
        for (int j = 0; j <= lim; j++) acc += p[j] * (double)hv[(kvh * KL + j) * D + d];
        const double diff = fabs(acc / sum - (double)o[(h * L + r) * D + d]);
        if (diff > maxd) maxd = diff;
      }
    }
  }
  printf("kv=%d %s  max |mlx - float64 reference| = %g\n", kv, causal ? "causal" : "no mask", maxd);
  return 0;
}

./sdpa_repro 1000 causal 256 prints 0.000111444 on v0.32.3 and 0.0354094 on 4ca8d5d5.

Lingua principale
C++
Stelle
29.1k
Fork
2.4k
Merge medio
3g 23h
PR unite (30g)
90

Preparare l'ambiente

Come iniziare

  1. Leggi tutta la issue e poi la guida ai contributi del progetto.
  2. Commenta sulla issue per dire che te ne occupi tu — evita che due persone facciano lo stesso lavoro.
  3. Fai un fork del repository e lavora su un branch.
  4. Apri una pull request che faccia riferimento al numero della issue.

Altre issue di ml-explore/mlx

Tutte le issue di ml-explore/mlx

Issue simili

Altre issue su C++

Ricevi le nuove issue nella tua casella

Un breve riepilogo di issue GitHub adatte ai principianti.