Merge pull request #194 from ausimian/feat/expr-compiler-sdpa-sinks

feat: lower SDPA-with-sinks blocks in the native Expr compiler

Author: ausimian

RELEASE.md

@@ -88,6 +88,16 @@
   seeding (every primitive in the expansion was already on the native
   path), bit-identical to the Evaluator.
 
+- **`Emily.Fast.Block.SDPAWithSinks` and `SDPAWithMaskAndSinks` lower
+  natively** — finishes the Expr op-coverage checklist (#188). Two new
+  opcodes (`fast_sdpa_sinks` / `fast_sdpa_mask_sinks`) call the same
+  `mx::fast::scaled_dot_product_attention` kernel the existing
+  `fast_sdpa{,_mask}` opcodes use, with the extra `sinks` operand
+  wired into the kernel's `sinks` arg — same call shape as
+  `Emily.Backend.fast_scaled_dot_product_attention_with{,_mask_and}_sinks`.
+  gpt-oss attention now compiles fully native (with or without an
+  additive mask), bit-identical to the Evaluator.
+
 - **`Nx.LinAlg.cholesky` / `solve` / `qr` / `eigh` / `lu` / `svd` /
   `determinant` lower natively** — closes the LinAlg-block cluster on
   #188. Two new infrastructure pieces enable this: a small extension to

c_src/emily/opcodes.hpp

@@ -241,9 +241,20 @@ enum class Opcode : int64_t {
   // operands [a]; arity 3 (U, S, V); `full_matrices` hard-coded `true` —
   // the IR slices for the thin case, matching Emily.Backend.native_svd/3.
   LinalgSVD = 120,
+  // SDPA with attention sinks (gpt-oss). operands [q, k, v, sinks];
+  // iattrs [[scale_bits], [causal]]. Mirrors
+  // Emily.Backend.fast_scaled_dot_product_attention_with_sinks/6 — same
+  // mx::fast::scaled_dot_product_attention entry point as FastSDPA, with
+  // the sinks operand wired into the kernel's `sinks` arg.
+  FastSDPASinks = 121,
+  // SDPA with both an explicit additive mask and attention sinks
+  // (gpt-oss). operands [q, k, v, mask, sinks]; iattrs [[scale_bits]];
+  // mask_mode = "array". Mirrors
+  // Emily.Backend.fast_scaled_dot_product_attention_with_mask_and_sinks/7.
+  FastSDPAMaskSinks = 122,
 };
 
-inline constexpr int64_t kOpcodeCount = 121;
+inline constexpr int64_t kOpcodeCount = 123;
 
 // Quant mode code (Emily.IR @quant_modes) -> MLX mode string.
 inline std::string qmode_from_code(int64_t code) {
@@ -562,6 +573,31 @@ inline mx::array dispatch_op(Opcode op, const std::vector<mx::array> &in,
         in[0], in[1], in[2], scale, "array", std::optional<mx::array>(in[3]),
         std::nullopt, s);
   }
+  case Opcode::FastSDPASinks: {
+    if (in.size() != 4) {
+      throw std::invalid_argument("fast_sdpa_sinks expects 4 operands, got " +
+                                  std::to_string(in.size()));
+    }
+    auto scale = static_cast<float>(
+        emily::f64_from_bits(scalar_at(iattrs, 0, "fast_sdpa_sinks")));
+    std::string mask_mode =
+        scalar_at(iattrs, 1, "fast_sdpa_sinks") != 0 ? "causal" : "";
+    return mx::fast::scaled_dot_product_attention(
+        in[0], in[1], in[2], scale, mask_mode, std::nullopt,
+        std::optional<mx::array>(in[3]), s);
+  }
+  case Opcode::FastSDPAMaskSinks: {
+    if (in.size() != 5) {
+      throw std::invalid_argument(
+          "fast_sdpa_mask_sinks expects 5 operands, got " +
+          std::to_string(in.size()));
+    }
+    auto scale = static_cast<float>(
+        emily::f64_from_bits(scalar_at(iattrs, 0, "fast_sdpa_mask_sinks")));
+    return mx::fast::scaled_dot_product_attention(
+        in[0], in[1], in[2], scale, "array", std::optional<mx::array>(in[3]),
+        std::optional<mx::array>(in[4]), s);
+  }
   case Opcode::QuantizedMatmul: {
     if (in.size() != 4) {
       throw std::invalid_argument("quantized_matmul expects 4 operands, got " +

lib/emily/ir.ex

@@ -208,7 +208,15 @@ defmodule Emily.IR do
     linalg_qr: 117,
     linalg_eigh: 118,
     linalg_lu: 119,
-    linalg_svd: 120
+    linalg_svd: 120,
+    # SDPA with attention sinks (gpt-oss). The mask-only and sinks-only
+    # variants of the existing `fast_sdpa{,_mask}` family — same
+    # mx::fast::scaled_dot_product_attention entry point, with the
+    # extra `sinks` operand wired through.
+    # operands [q, k, v, sinks]; iattrs [[scale_bits], [causal]]
+    fast_sdpa_sinks: 121,
+    # operands [q, k, v, mask, sinks]; iattrs [[scale_bits]]
+    fast_sdpa_mask_sinks: 122
   }
 
   # Quant mode string -> code; decoded by qmode_from_code in
@@ -1261,6 +1269,59 @@ defmodule Emily.IR do
     emit_coerced(state, :fast_sdpa_mask, [rq, rk, rv, rm], [[float_bits(scale)]], t.type)
   end
 
+  # SDPA with attention sinks (gpt-oss). Mirrors
+  # Emily.Backend.fast_scaled_dot_product_attention_with_sinks/6 —
+  # routes through the same mx::fast::scaled_dot_product_attention
+  # entry point as SDPA, with the sinks tensor as the kernel's `sinks`
+  # arg and the causal flag as the mask_mode discriminator
+  # ("causal" / "").
+  defp lower_block(
+         %FB.SDPAWithSinks{scale: scale, causal: causal},
+         [q, k, v, sinks],
+         _expr,
+         t,
+         state
+       ) do
+    {rq, state} = lower_node(q, state)
+    {rk, state} = lower_node(k, state)
+    {rv, state} = lower_node(v, state)
+    {rs, state} = lower_node(sinks, state)
+
+    emit_coerced(
+      state,
+      :fast_sdpa_sinks,
+      [rq, rk, rv, rs],
+      [[float_bits(scale)], [bool_int(causal)]],
+      t.type
+    )
+  end
+
+  # SDPA with an additive mask *and* attention sinks (gpt-oss). Mirrors
+  # Emily.Backend.fast_scaled_dot_product_attention_with_mask_and_sinks/7 —
+  # mask_mode is fixed "array" (the dispatcher hard-codes it), so only the
+  # scale crosses as an iattr.
+  defp lower_block(
+         %FB.SDPAWithMaskAndSinks{scale: scale},
+         [q, k, v, mask, sinks],
+         _expr,
+         t,
+         state
+       ) do
+    {rq, state} = lower_node(q, state)
+    {rk, state} = lower_node(k, state)
+    {rv, state} = lower_node(v, state)
+    {rm, state} = lower_node(mask, state)
+    {rs, state} = lower_node(sinks, state)
+
+    emit_coerced(
+      state,
+      :fast_sdpa_mask_sinks,
+      [rq, rk, rv, rm, rs],
+      [[float_bits(scale)]],
+      t.type
+    )
+  end
+
   defp lower_block(%QB.QuantizedMatmul{} = qb, [x, q, s, b], _expr, t, state) do
     {rx, state} = lower_node(x, state)
     {rq, state} = lower_node(q, state)

test/emily/compiler_equivalence_test.exs

@@ -682,6 +682,57 @@ defmodule Emily.CompilerEquivalenceTest do
         [q, k, v, mask]
       )
     end
+
+    # SDPA with attention sinks (gpt-oss). Backend dispatches both the
+    # plain sinks variant and the masked sinks variant through the same
+    # mx::fast::scaled_dot_product_attention kernel (with the `sinks` arg
+    # populated). The native IR routes through new fast_sdpa_sinks /
+    # fast_sdpa_mask_sinks opcodes that call the same kernel — so the two
+    # paths are bit-identical to the evaluator.
+    test "sdpa with sinks (non-causal) matches the fused kernel" do
+      shape = {1, 2, 4, 8}
+      q = Nx.iota(shape, type: :f32, backend: Emily.Backend) |> Nx.divide(64.0)
+      k = Nx.iota(shape, type: :f32, backend: Emily.Backend) |> Nx.divide(48.0)
+      v = Nx.iota(shape, type: :f32, backend: Emily.Backend) |> Nx.divide(32.0)
+      # Sinks is a per-head 1-D tensor (one extra logit per head).
+      sinks = et([0.5, -0.5])
+
+      assert_equiv(
+        fn q, k, v, s -> Emily.Fast.scaled_dot_product_attention(q, k, v, sinks: s) end,
+        [q, k, v, sinks]
+      )
+    end
+
+    test "sdpa with sinks (causal) matches the fused kernel" do
+      shape = {1, 2, 4, 8}
+      q = Nx.iota(shape, type: :f32, backend: Emily.Backend) |> Nx.divide(64.0)
+      k = Nx.iota(shape, type: :f32, backend: Emily.Backend) |> Nx.divide(48.0)
+      v = Nx.iota(shape, type: :f32, backend: Emily.Backend) |> Nx.divide(32.0)
+      sinks = et([0.0, 1.0])
+
+      assert_equiv(
+        fn q, k, v, s ->
+          Emily.Fast.scaled_dot_product_attention(q, k, v, sinks: s, causal: true)
+        end,
+        [q, k, v, sinks]
+      )
+    end
+
+    test "sdpa with an additive mask AND sinks matches the fused kernel" do
+      shape = {1, 2, 4, 8}
+      q = Nx.iota(shape, type: :f32, backend: Emily.Backend) |> Nx.divide(64.0)
+      k = Nx.iota(shape, type: :f32, backend: Emily.Backend) |> Nx.divide(48.0)
+      v = Nx.iota(shape, type: :f32, backend: Emily.Backend) |> Nx.divide(32.0)
+      mask = Nx.broadcast(Nx.tensor(0.0, backend: Emily.Backend), {1, 1, 4, 4})
+      sinks = et([-1.0, 1.0])
+
+      assert_equiv(
+        fn q, k, v, m, s ->
+          Emily.Fast.scaled_dot_product_attention_with_mask(q, k, v, m, sinks: s)
+        end,
+        [q, k, v, mask, sinks]
+      )
+    end
   end
 
   describe "quantized matmul block" do