diff --git a/emlx/c_src/emlx_compiler.cpp b/emlx/c_src/emlx_compiler.cpp index 46995fc..bcd37f8 100644 --- a/emlx/c_src/emlx_compiler.cpp +++ b/emlx/c_src/emlx_compiler.cpp @@ -1500,33 +1500,105 @@ static const std::unordered_map op_registry = { // rope (per-batch offset array): operands = [a, position_ids]; // attrs = [dims, traditional, base_bits, scale_bits]. Offsets are // position_ids[:, 0] (matches EMLX.Fast.rope_with_positions_fast_callback). + // + // H>1 (Bumblebee {B,T,H,D} convention) falls back to the hand + // cos/sin/rotate composition (rope_rotate_from_angles) — see + // emlx_fast.cpp's multi-head-safe RoPE helpers comment / + // https://github.com/elixir-nx/emlx/issues/121. H==1 keeps the fused + // mlx::core::fast::rope path. The full position_ids operand is already + // available here, so no offset+arange reconstruction is needed. {"fast_rope_ids", [](const auto &ops, const auto &attrs) { + const auto &a = ops[0]; const auto &pos = ops[1]; + int dims = static_cast(attrs[0]); + bool traditional = attrs[1] != 0; + + if (a.ndim() == 4 && a.shape(2) > 1) { + if (traditional) { + throw std::invalid_argument( + "fast_rope_ids: traditional=true not supported for " + "multi-head (H>1) input"); + } + float base = attr_to_float(attrs[2]); + float scale = attr_to_float(attrs[3]); + int half = dims / 2; + + std::vector inv_freq_host(half); + for (int i = 0; i < half; ++i) { + float expo = (2.0f * static_cast(i)) / static_cast(dims); + inv_freq_host[i] = 1.0f / std::pow(base, expo); + } + auto inv_freq = + mlx::core::array(inv_freq_host.begin(), {half}, mlx::core::float32); + + int B = pos.shape(0); + int T = pos.shape(1); + auto pos_f = mlx::core::astype(pos, mlx::core::float32); + auto pos_bt1 = mlx::core::reshape(pos_f, to_shape({B, T, 1})); + auto inv_11h = mlx::core::reshape(inv_freq, to_shape({1, 1, half})); + auto angles = mlx::core::multiply( + mlx::core::multiply(pos_bt1, inv_11h), + mlx::core::array(scale, mlx::core::float32)); + + return rope_rotate_from_angles(a, angles, dims); + } + int B = pos.shape(0); auto offsets = mlx::core::reshape( mlx::core::slice(pos, to_shape({0, 0}), to_shape({B, 1}), to_shape({1, 1})), to_shape({B})); return mlx::core::fast::rope( - ops[0], static_cast(attrs[0]), attrs[1] != 0, + a, dims, traditional, attr_to_float(attrs[2]), attr_to_float(attrs[3]), offsets, std::nullopt); }}, // rope (precomputed freqs): operands = [a, position_ids, freqs]; // attrs = [dims, traditional, scale_bits]. base = nullopt (freqs supplied). + // + // H>1 falls back to the hand cos/sin/rotate composition — see the + // "fast_rope_ids" comment above and emlx_fast.cpp's multi-head-safe RoPE + // helpers comment. H==1 keeps the fused mlx::core::fast::rope path. {"fast_rope_with_freqs", [](const auto &ops, const auto &attrs) { + const auto &a = ops[0]; const auto &pos = ops[1]; + const auto &freqs = ops[2]; + int dims = static_cast(attrs[0]); + bool traditional = attrs[1] != 0; + + if (a.ndim() == 4 && a.shape(2) > 1) { + if (traditional) { + throw std::invalid_argument( + "fast_rope_with_freqs: traditional=true not supported for " + "multi-head (H>1) input"); + } + float scale = attr_to_float(attrs[2]); + int half = dims / 2; + + int B = pos.shape(0); + int T = pos.shape(1); + auto pos_f = mlx::core::astype(pos, mlx::core::float32); + auto pos_bt1 = mlx::core::reshape(pos_f, to_shape({B, T, 1})); + auto inv_freq = mlx::core::reciprocal(mlx::core::astype(freqs, mlx::core::float32)); + auto inv_11h = mlx::core::reshape(inv_freq, to_shape({1, 1, half})); + auto angles = mlx::core::multiply( + mlx::core::multiply(pos_bt1, inv_11h), + mlx::core::array(scale, mlx::core::float32)); + + return rope_rotate_from_angles(a, angles, dims); + } + int B = pos.shape(0); auto offsets = mlx::core::reshape( mlx::core::slice(pos, to_shape({0, 0}), to_shape({B, 1}), to_shape({1, 1})), to_shape({B})); return mlx::core::fast::rope( - ops[0], static_cast(attrs[0]), attrs[1] != 0, std::nullopt, - attr_to_float(attrs[2]), offsets, ops[2]); + a, dims, traditional, std::nullopt, + attr_to_float(attrs[2]), offsets, freqs); }}, // rope (arbitrary per-token position_ids, base-derived inv_freq): diff --git a/emlx/c_src/emlx_fast.cpp b/emlx/c_src/emlx_fast.cpp index 705133b..3bb1f55 100644 --- a/emlx/c_src/emlx_fast.cpp +++ b/emlx/c_src/emlx_fast.cpp @@ -9,6 +9,87 @@ // plus a one-line `FINE_ASYNC_NIF(NAME)` — the registered NIF name, arity, // and the ASYNC_NIF/nif_funcs[] wiring in emlx_nif.cpp are all unchanged. +// ── Multi-head-safe RoPE helpers ──────────────────────────────────────────── +// +// mlx::core::fast::rope only recognizes Bumblebee's {B, T, H, D} +// (heads-not-yet-transposed) convention via a stride-detection special case +// meant to catch a *transposed view* of {B, H, T, D}. A plain, freshly +// allocated {B, T, H, D} tensor is row-contiguous, so that guard never fires +// and the kernel falls into its row_contiguous branch, which treats axis -2 +// (H) as the sequence axis — rotating head `h` at angle `position + h` +// instead of `position` for every h > 0. +// See https://github.com/elixir-nx/emlx/issues/121. +// +// H==1 has no such ambiguity (there's only one "row" per position), so the +// fused Metal path stays correct and fast there. For H>1 we compose the same +// split-half rotation from primitive ops instead, broadcasting cos/sin across +// every head — this mirrors mlx::fast::rope's non-traditional formula +// bit-for-bit: +// out[..half] = a[..half]*cos(angles) - a[half..dims]*sin(angles) +// out[half..] = a[half..dims]*cos(angles) + a[..half]*sin(angles) +// (equivalently: rotate_half = concat(-a[half..dims], a[..half]); out = +// a[..dims]*cos_full + rotate_half*sin_full — the form used below). +// +// `a` is {B, T, H, D}; `angles` is {B, T, half} (already position*inv_freq*scale). +static mlx::core::array rope_rotate_from_angles(const mlx::core::array &a, + const mlx::core::array &angles, + int dims, + mlx::core::Device device) { + int B = a.shape(0); + int T = a.shape(1); + int H = a.shape(2); + int D = a.shape(3); + int half = dims / 2; + + auto cos_bt1h = astype(reshape(cos(angles, device), {B, T, 1, half}, device), a.dtype(), device); + auto sin_bt1h = astype(reshape(sin(angles, device), {B, T, 1, half}, device), a.dtype(), device); + + auto cos_full = concatenate(std::vector{cos_bt1h, cos_bt1h}, 3, device); + auto sin_full = concatenate(std::vector{sin_bt1h, sin_bt1h}, 3, device); + + auto x1 = slice(a, to_shape({0, 0, 0, 0}), to_shape({B, T, H, half}), device); + auto x2 = slice(a, to_shape({0, 0, 0, half}), to_shape({B, T, H, dims}), device); + auto rotated = concatenate(std::vector{negative(x2, device), x1}, 3, device); + + auto a_head = slice(a, to_shape({0, 0, 0, 0}), to_shape({B, T, H, dims}), device); + auto rope_head = add(multiply(a_head, cos_full, device), multiply(rotated, sin_full, device), device); + + if (dims == D) { + return rope_head; + } + auto tail = slice(a, to_shape({0, 0, 0, dims}), to_shape({B, T, H, D}), device); + return concatenate(std::vector{rope_head, tail}, 3, device); +} + +// inv_freq[i] = 1 / base^(2i/dims), i in [0, dims/2). +static mlx::core::array rope_inv_freq_from_base(int dims, double base) { + int half = dims / 2; + std::vector inv_freq_host(half); + float base_f = static_cast(base); + for (int i = 0; i < half; ++i) { + float expo = (2.0f * static_cast(i)) / static_cast(dims); + inv_freq_host[i] = 1.0f / std::pow(base_f, expo); + } + return array(inv_freq_host.begin(), {half}, float32); +} + +// Reconstructs per-token angles = (offset[b] + t) * inv_freq * scale for the +// assumed-sequential-positions contract of fast_rope_ids/fast_rope_with_freqs +// (offset is {B}, one starting position per batch example; T is a's shape(1)). +static mlx::core::array rope_sequential_angles(const mlx::core::array &offset, + const mlx::core::array &inv_freq, + int B, int T, int half, + double scale, + mlx::core::Device device) { + auto offset_b1 = reshape(astype(offset, float32, device), {B, 1}, device); + auto steps_1t = reshape(astype(arange(T, int32, device), float32, device), {1, T}, device); + auto pos_bt = add(offset_b1, steps_1t, device); + auto pos_bt1 = reshape(pos_bt, {B, T, 1}, device); + auto inv_11h = reshape(inv_freq, {1, 1, half}, device); + auto scale_arr = array(static_cast(scale), float32); + return multiply(multiply(pos_bt1, inv_11h, device), scale_arr, device); +} + // fast_rms_norm — fused RMS normalisation // MLX: mlx::fast::rms_norm(x, weight, eps, stream) → array, same shape as x // weight is optional in the C++ API; a plain TensorArg implicitly @@ -86,10 +167,31 @@ FINE_ASYNC_NIF(fast_layer_norm_no_bias) // Calls the array-offset overload of mlx::fast::rope. // offset must be shape {B} — one starting position per batch example. // Assumes positions are sequential within each example: [offset[b], offset[b]+1, ..., offset[b]+T-1]. +// +// H>1 (Bumblebee {B,T,H,D} convention — see the multi-head-safe RoPE helpers +// comment above) falls back to the hand cos/sin/rotate composition; H==1 +// keeps the fused Metal path. mlx::core::array fast_rope_ids_impl(ErlNifEnv *env, TensorArg a, int dims, bool traditional, double base, double scale, TensorArg offset, mlx::core::Device device) { + if (a->ndim() == 4 && a->shape(2) > 1) { + if (traditional) { + throw std::invalid_argument( + "fast_rope_ids: traditional=true not supported for multi-head " + "(H>1) input"); + } + + int B = a->shape(0); + int T = a->shape(1); + int half = dims / 2; + + auto inv_freq = rope_inv_freq_from_base(dims, base); + auto angles = rope_sequential_angles(*offset, inv_freq, B, T, half, scale, device); + + return rope_rotate_from_angles(*a, angles, dims, device); + } + return fast::rope(*a, dims, traditional, (float)base, (float)scale, *offset, std::nullopt, device); } @@ -99,11 +201,34 @@ FINE_ASYNC_NIF(fast_rope_ids) // Calls the freqs overload of mlx::fast::rope (base=nullopt, freqs supplied). // offset must be shape {B} — one starting position per batch example. // freqs must be shape {dims/2} — precomputed inverse frequencies. +// +// H>1 falls back to the hand cos/sin/rotate composition (see the +// multi-head-safe RoPE helpers comment above); H==1 keeps the fused Metal +// path. mlx::fast::rope's freqs overload uses inv_freq = reciprocal(freqs) +// internally, so we replicate that reciprocal here to stay bit-for-bit with +// the fused path. mlx::core::array fast_rope_with_freqs_impl(ErlNifEnv *env, TensorArg a, int dims, bool traditional, double scale, TensorArg offset, TensorArg freqs, mlx::core::Device device) { + if (a->ndim() == 4 && a->shape(2) > 1) { + if (traditional) { + throw std::invalid_argument( + "fast_rope_with_freqs: traditional=true not supported for " + "multi-head (H>1) input"); + } + + int B = a->shape(0); + int T = a->shape(1); + int half = dims / 2; + + auto inv_freq = reciprocal(astype(*freqs, float32, device), device); + auto angles = rope_sequential_angles(*offset, inv_freq, B, T, half, scale, device); + + return rope_rotate_from_angles(*a, angles, dims, device); + } + return fast::rope(*a, dims, traditional, std::nullopt, (float)scale, *offset, *freqs, device); } @@ -140,7 +265,6 @@ mlx::core::array fast_rope_positions_impl(ErlNifEnv *env, TensorArg a, int B = a->shape(0); int T = a->shape(1); - int H = a->shape(2); int D = a->shape(3); if (position_ids->shape(0) != B || position_ids->shape(1) != T) { @@ -158,14 +282,7 @@ mlx::core::array fast_rope_positions_impl(ErlNifEnv *env, TensorArg a, int half = dims / 2; - std::vector inv_freq_host(half); - float base_f = static_cast(base); - for (int i = 0; i < half; ++i) { - float expo = (2.0f * static_cast(i)) / static_cast(dims); - inv_freq_host[i] = 1.0f / std::pow(base_f, expo); - } - - auto inv_freq = array(inv_freq_host.begin(), {half}, float32); + auto inv_freq = rope_inv_freq_from_base(dims, base); auto pos_f = astype(*position_ids, float32, device); auto pos_bt1 = reshape(pos_f, {B, T, 1}, device); @@ -173,25 +290,7 @@ mlx::core::array fast_rope_positions_impl(ErlNifEnv *env, TensorArg a, auto scale_arr = array(static_cast(scale), float32); auto angles = multiply(multiply(pos_bt1, inv_11h, device), scale_arr, device); - auto cos_bt1h = astype(reshape(cos(angles, device), {B, T, 1, half}, device), a->dtype(), device); - auto sin_bt1h = astype(reshape(sin(angles, device), {B, T, 1, half}, device), a->dtype(), device); - - auto cos_full = concatenate(std::vector{cos_bt1h, cos_bt1h}, 3, device); - auto sin_full = concatenate(std::vector{sin_bt1h, sin_bt1h}, 3, device); - - auto x1 = slice(*a, to_shape({0, 0, 0, 0}), to_shape({B, T, H, half}), device); - auto x2 = slice(*a, to_shape({0, 0, 0, half}), to_shape({B, T, H, dims}), device); - auto rotated = concatenate(std::vector{negative(x2, device), x1}, 3, device); - - auto a_head = slice(*a, to_shape({0, 0, 0, 0}), to_shape({B, T, H, dims}), device); - auto rope_head = add(multiply(a_head, cos_full, device), multiply(rotated, sin_full, device), device); - - if (dims == D) { - return rope_head; - } else { - auto tail = slice(*a, to_shape({0, 0, 0, dims}), to_shape({B, T, H, D}), device); - return concatenate(std::vector{rope_head, tail}, 3, device); - } + return rope_rotate_from_angles(*a, angles, dims, device); } FINE_ASYNC_NIF(fast_rope_positions) diff --git a/emlx/lib/emlx.ex b/emlx/lib/emlx.ex index fdb5e07..9846ef0 100644 --- a/emlx/lib/emlx.ex +++ b/emlx/lib/emlx.ex @@ -588,7 +588,16 @@ defmodule EMLX do Single Metal shader. Applies RoPE with a scalar position `offset`. - - `a` — input `{B, ..., T, D}` + - `a` — input `{B, ..., T, D}`; **`T` must be the second-to-last + axis** — e.g. `{B, T, D}` or, with a heads axis, `{B, H, T, D}` + (heads-transposed). If `H` (or any other middle axis) is placed *before* + `T` instead — e.g. `{B, T, H, D}` — `mlx::core::fast::rope` silently + rotates row `h` at angle `position + h` instead of `position` for every + `h > 0` (only row 0 is correct); see + [elixir-nx/emlx#121](https://github.com/elixir-nx/emlx/issues/121). + For Bumblebee's heads-not-yet-transposed `{B, T, H, D}` convention, use + `EMLX.fast_rope_ids/6` or `EMLX.fast_rope_positions/6` instead, which + guard against this. - `dims` — number of feature dims to rotate (≤ last-axis size, must be even) - `traditional` — `false` for split-half (Qwen3); `true` for interleaved - `base` — angular frequency base (e.g. 10_000 or 1_000_000) diff --git a/emlx/lib/emlx/fast.ex b/emlx/lib/emlx/fast.ex index 973149f..87bb6d0 100644 --- a/emlx/lib/emlx/fast.ex +++ b/emlx/lib/emlx/fast.ex @@ -461,7 +461,12 @@ defmodule EMLX.Fast do @doc """ Fused rotary position embedding (`mlx::fast::rope`). - - `a` — input `{B, ..., T, D}`; `...` dims are passed through. + - `a` — input `{B, ..., T, D}`; `...` dims are passed through, but + **`T` must be the second-to-last axis** (e.g. `{B, H, T, D}`, + heads-transposed) — see `EMLX.fast_rope/6`'s `@doc` for why a heads-axis + placed *before* `T` (Bumblebee's `{B, T, H, D}`) silently miscomputes for + more than one head ([elixir-nx/emlx#121](https://github.com/elixir-nx/emlx/issues/121)). + Use `rope_with_positions/6` or `rope_with_freqs/6` for that layout instead. - `dims` — number of feature dims to rotate (≤ last-axis size, must be even). - `traditional` — `false` for split-half (Qwen3); `true` for interleaved. - `base` — angular frequency base (e.g. `10_000` or `1_000_000`). diff --git a/emlx/test/emlx/fast_test.exs b/emlx/test/emlx/fast_test.exs index 5a4444a..a3544d4 100644 --- a/emlx/test/emlx/fast_test.exs +++ b/emlx/test/emlx/fast_test.exs @@ -123,6 +123,32 @@ defmodule EMLX.FastTest do assert_all_close(out2, expected, atol: 1.0e-4, rtol: 1.0e-4) end + + # https://github.com/elixir-nx/emlx/issues/121 — mlx::core::fast::rope's + # head_seq_transpose stride-detection guard doesn't trigger for a + # freshly-allocated, contiguous {B,T,H,D} tensor, so its row_contiguous + # fallback used to rotate head h at angle `position + h` instead of + # `position`, for every h > 0. Every head, computed jointly, must equal + # that same head computed alone (same position, no cross-head leakage). + test "T=1 decode, H>1: every head matches its single-head computation" do + dims = 8 + half = div(dims, 2) + freqs = Nx.iota({half}, type: :f32) |> Nx.add(0.1) |> Nx.divide(20) |> gpu() + + a = Nx.iota({1, 1, 3, dims}, type: :f32) |> Nx.divide(20) |> gpu() + pos = Nx.tensor([[6]], type: :s32) |> gpu() + + joint = Fast.rope_with_freqs(a, pos, dims, false, 1.0, freqs) |> Nx.backend_transfer() + + for head <- 0..2 do + a_head = a[[.., .., head..head, ..]] + + alone = + Fast.rope_with_freqs(a_head, pos, dims, false, 1.0, freqs) |> Nx.backend_transfer() + + assert_all_close(joint[[.., .., head..head, ..]], alone, atol: 1.0e-4, rtol: 1.0e-4) + end + end end describe "EMLX.Fast.rope_with_positions/6" do @@ -151,6 +177,32 @@ defmodule EMLX.FastTest do assert_all_close(out, expected, atol: 1.0e-4, rtol: 1.0e-4) end + + # https://github.com/elixir-nx/emlx/issues/121 — low-base (< 1.0e5) T=1 + # decode routes through fast_rope_ids (mlx::core::fast::rope's + # array-offset overload), which had the same multi-head bug as + # fast_rope_with_freqs (see the regression test in the describe block + # above). Every head, computed jointly, must equal that same head + # computed alone. + test "low-base decode (T=1), H>1: every head matches its single-head computation" do + dims = 8 + base = 10_000.0 + + a = Nx.iota({1, 1, 3, dims}, type: :f32) |> Nx.divide(50) |> gpu() + pos = Nx.tensor([[7]], type: :s32) |> gpu() + + joint = Fast.rope_with_positions(a, pos, dims, false, base, 1.0) |> Nx.backend_transfer() + + for head <- 0..2 do + a_head = a[[.., .., head..head, ..]] + + alone = + Fast.rope_with_positions(a_head, pos, dims, false, base, 1.0) + |> Nx.backend_transfer() + + assert_all_close(joint[[.., .., head..head, ..]], alone, atol: 1.0e-4, rtol: 1.0e-4) + end + end end # ── scaled_dot_product_attention ───────────────────────────────────────────── diff --git a/emlx/test/emlx/native/expr_test.exs b/emlx/test/emlx/native/expr_test.exs index fa45edc..e07c1a4 100644 --- a/emlx/test/emlx/native/expr_test.exs +++ b/emlx/test/emlx/native/expr_test.exs @@ -2657,6 +2657,56 @@ defmodule EMLX.Native.ExprTest do assert_all_close(native, eager, tol: 1.0e-3) end + # https://github.com/elixir-nx/emlx/issues/121 — mlx::core::fast::rope's + # head_seq_transpose stride-detection guard doesn't trigger for a + # freshly-allocated, contiguous {B,T,H,D} tensor, so its row_contiguous + # fallback used to rotate head h at angle `position + h` instead of + # `position`, for every h > 0. "fused replay matches eager" above only + # proves native and eager *agree*, which they trivially did even while + # both called the same buggy primitive underneath — these tests instead + # check each head against an independent single-head computation (every + # head, computed alone, must use only its own position). + test "rope_with_positions (decode T=1, H>1): every head matches its single-head computation" do + dims = 64 + base = 10_000.0 + fun = fn a, pos -> EMLX.Fast.rope_with_positions(a, pos, dims, false, base, 1.0) end + + a = Nx.iota({1, 1, 3, dims}, type: :f32) |> Nx.divide(100) |> gpu_t() + pos = Nx.tensor([[6]], type: :s32) |> gpu_t() + + native = Nx.Defn.jit(fun, compiler: EMLX, device: :gpu).(a, pos) + eager = Nx.Defn.jit(fun, compiler: Nx.Defn.Evaluator).(a, pos) + + for head <- 0..2 do + a_head = a[[.., .., head..head, ..]] + alone = EMLX.Fast.rope_with_positions(a_head, pos, dims, false, base, 1.0) + + assert_all_close(native[[.., .., head..head, ..]], alone, tol: 1.0e-3) + assert_all_close(eager[[.., .., head..head, ..]], alone, tol: 1.0e-3) + end + end + + test "rope_with_freqs (decode T=1, H>1): every head matches its single-head computation" do + dims = 64 + half = div(dims, 2) + fun = fn a, pos, freqs -> EMLX.Fast.rope_with_freqs(a, pos, dims, false, 1.0, freqs) end + + a = Nx.iota({1, 1, 3, dims}, type: :f32) |> Nx.divide(100) |> gpu_t() + pos = Nx.tensor([[6]], type: :s32) |> gpu_t() + freqs = Nx.iota({half}, type: :f32) |> Nx.add(1) |> Nx.divide(1000) |> gpu_t() + + native = Nx.Defn.jit(fun, compiler: EMLX, device: :gpu).(a, pos, freqs) + eager = Nx.Defn.jit(fun, compiler: Nx.Defn.Evaluator).(a, pos, freqs) + + for head <- 0..2 do + a_head = a[[.., .., head..head, ..]] + alone = EMLX.Fast.rope_with_freqs(a_head, pos, dims, false, 1.0, freqs) + + assert_all_close(native[[.., .., head..head, ..]], alone, tol: 1.0e-3) + assert_all_close(eager[[.., .., head..head, ..]], alone, tol: 1.0e-3) + end + end + test "decode-shaped block: fused path improves over primitive replay" do # A small attention+norm decode step: RMSNorm → causal SDPA → RMSNorm. scale = 0.125 @@ -3495,11 +3545,15 @@ defmodule EMLX.Native.ExprTest do # Pure-Nx primitive reference for prefill RoPE against a precomputed `freqs` # tensor (mirrors emlx_compiler.cpp's fast_rope_with_freqs_positions lambda: - # inv_freq = reciprocal(freqs), half-rotate cos/sin blend). Needed for H>1 - # cases because the eager EMLX.Fast.rope_with_freqs reference calls - # mlx::core::fast::rope directly, which miscomputes non-head-0 rotations for - # multi-head input — see mlx-fast-rope-multihead-bugreport.md. `a` is - # {B, T, H, D}; `pos` is {B, T}; `freqs` is {dims/2}. + # inv_freq = reciprocal(freqs), half-rotate cos/sin blend). Kept as an + # independent oracle for H>1 prefill: the T>1 lowering path never calls + # mlx::core::fast::rope directly (it always uses the hand cos/sin/rotate + # composition, broadcast across heads), but a hand-written primitive here + # still guards against regressions in that composition without relying on + # the code under test to check itself. (T=1 decode used to have a related, + # now-fixed, multi-head bug — see elixir-nx/emlx#121 and the "multi-head + # correctness" tests below.) `a` is {B, T, H, D}; `pos` is {B, T}; `freqs` + # is {dims/2}. defp rope_freqs_prim(a, pos, freqs, dims, scale) do {b, t, _h, d} = Nx.shape(a) half = div(dims, 2) diff --git a/emlx_axon/lib/emlx_axon/qwen3/attention.ex b/emlx_axon/lib/emlx_axon/qwen3/attention.ex index 9b07f85..132a54b 100644 --- a/emlx_axon/lib/emlx_axon/qwen3/attention.ex +++ b/emlx_axon/lib/emlx_axon/qwen3/attention.ex @@ -7,8 +7,11 @@ defmodule EMLXAxon.Qwen3.Attention do the cache without an extra transpose, and feeds directly into `EMLX.Fast.scaled_dot_product_attention` which expects `{B, N, T, D}`. - RoPE is computed by `EMLX.Fast.rope/6` (single Metal shader, no precomputed - cos/sin needed). The `offset` is the current cache fill length. + `forward/12` delegates to the native `EMLX.Native.Qwen3.kv_cache_attention`/ + `attention_block` NIFs (`qwen3_plugin.cpp`), which transpose Q/K to + `{B, N, T, D}` and apply RoPE internally via `mlx::core::fast::rope` + (single Metal shader, no precomputed cos/sin needed) before the fused + cache-update + SDPA. The `offset` is the current cache fill length. """ alias EMLXAxon.Qwen3.Layers diff --git a/emlx_axon/lib/emlx_axon/qwen3/layers.ex b/emlx_axon/lib/emlx_axon/qwen3/layers.ex index e988b3f..87d2fa8 100644 --- a/emlx_axon/lib/emlx_axon/qwen3/layers.ex +++ b/emlx_axon/lib/emlx_axon/qwen3/layers.ex @@ -3,8 +3,10 @@ defmodule EMLXAxon.Qwen3.Layers do Stateless layer primitives: RMSNorm, SwiGLU. `rms_norm/3` delegates to `EMLX.Fast.rms_norm` (single fused Metal shader). - RoPE is no longer computed here — `EMLX.Fast.rope/6` is called directly in - `Attention.forward/10` after projecting and transposing to `{B, N, T, D}`. + RoPE is no longer computed here — `Attention.forward/12` delegates to the + native `EMLX.Native.Qwen3.kv_cache_attention/attention_block` NIFs, which + transpose to `{B, N, T, D}` and apply RoPE internally via + `mlx::core::fast::rope` (see `qwen3_plugin.cpp`). """ import Nx.Defn