Use jax.lax.scaled_dot for MXFP dot_general in Qwix instead of jax.nn.scaled_matmul that binds to cuDNN

qwix/_src/core/mxfp_dot.py

@@ -21,19 +21,26 @@
 from qwix._src.core import qarray
 
 
+# Dimension numbers of the 3D form produced by _flatten_to_3d, i.e.
+# (B, M, K) x (B, N, K) -> (B, M, N).
+_DIMENSION_NUMBERS_3D = (((2,), (2,)), ((0,), (0,)))
+
+
 def mxfp_dot_general(
     lhs: qarray.MaybeQArray,
     rhs: qarray.MaybeQArray,
     dimension_numbers: jax.lax.DotDimensionNumbers,
     preferred_element_type: jax.typing.DTypeLike | None = None,
 ) -> jax.Array | None:
-  """Handles MXFP dot_general using `jax.nn.scaled_matmul`.
+  """Handles MXFP dot_general using `jax.lax.scaled_dot`.
 
   This dispatcher attempts to accelerate or decompose OCP/NVIDIA microscaled
-  matmuls using `jax.nn.scaled_matmul` (hardware Tensor Cores on Blackwell, or
-  JAX emulation/decomposition on TPUs, CPUs, and legacy GPUs).
+  matmuls using `jax.lax.scaled_dot`, which emits an `xla.scaled_dot` composite.
+  Today only the XLA GPU backend rewrites that composite into a native
+  block-scaled dot; every other backend decomposes it into a regular
+  dot_general.
 
-  Note that `scaled_matmul` is supported when BOTH operands are microscaled
+  Note that `scaled_dot` is supported when BOTH operands are microscaled
   formats (MXFP8, MXFP4, or NVFP4) with matching batch and contracting scale
   dimensions. One-sided microscaled operations or mismatched scale dimensions
   will cleanly return `None` to fall back to standard float emulation.
@@ -61,16 +68,16 @@ def mxfp_dot_general(
       lhs, rhs, preferred_element_type=preferred_element_type
   )
 
-  # TODO(b/538686860): Migrate to new API once it's available.
-  # jax.nn.scaled_matmul is natively accelerated on Blackwell GPUs (via cuDNN
-  # scaled matmul kernels) and emulated/decomposed on TPUs, CPUs, and legacy
-  # GPUs (like H100).
+  # TODO(b/538686860): Only the XLA GPU backend rewrites the `xla.scaled_dot`
+  # composite into a native block-scaled dot. Revisit once TPUs with native
+  # MXFP MXUs lower it natively too.
   try:
-    out_3d = jax.nn.scaled_matmul(
+    out_3d = jax.lax.scaled_dot(
         lhs_val_3d,
         rhs_val_3d,
-        lhs_scale_3d,
-        rhs_scale_3d,
+        lhs_scale=lhs_scale_3d,
+        rhs_scale=rhs_scale_3d,
+        dimension_numbers=_DIMENSION_NUMBERS_3D,
         preferred_element_type=result_type,
     )
   except Exception:  # pylint: disable=broad-except
@@ -95,21 +102,44 @@ def _inputs_compatible(
     lhs_scale_3d: jax.Array,
     rhs_scale_3d: jax.Array,
 ) -> bool:
-  """Checks 3D value and scale tensors are compatible for scaled_matmul."""
+  """Checks 3D value and scale tensors are compatible for scaled_dot."""
   return (
       lhs_val_3d.shape[0] == rhs_val_3d.shape[0]
       and lhs_val_3d.shape[2] == rhs_val_3d.shape[2]
       and lhs_scale_3d.shape[0] == rhs_scale_3d.shape[0]
       and lhs_scale_3d.shape[2] == rhs_scale_3d.shape[2]
+      and _subchannel_supported(lhs_val_3d, lhs_scale_3d)
+      and _subchannel_supported(rhs_val_3d, rhs_scale_3d)
   )
 
 
+def _subchannel_supported(val_3d: jax.Array, scale_3d: jax.Array) -> bool:
+  """Checks scaled_dot's subchannel constraints on the contracting dim.
+
+  `jax.lax.scaled_dot` requires the contracting dim to be a multiple of, and at
+  least twice as large as, the scale's contracting dim. Degenerate contractions
+  (e.g. outer products, which flatten to a contracting size of 1) don't qualify.
+
+  Args:
+    val_3d: The 3D value tensor.
+    scale_3d: The 3D scale tensor.
+
+  Returns:
+    Whether the contracting dim satisfies scaled_dot's subchannel constraints.
+  """
+  contracting_size = val_3d.shape[2]
+  scale_size = scale_3d.shape[2]
+  if contracting_size % scale_size != 0:
+    return False
+  return contracting_size // scale_size >= 2
+
+
 def _flatten_to_3d(
     operand: qarray.QArray,
     ca: Sequence[int],
     ba: Sequence[int],
 ) -> tuple[jax.Array, jax.Array]:
-  """Flattens a QArray operand and its scale to 3D for scaled_matmul."""
+  """Flattens a QArray operand and its scale to 3D for scaled_dot."""
   val = operand.qvalue
   scale = operand.scale
   ndim = operand.ndim
@@ -145,7 +175,7 @@ def _unflatten_from_3d(
     rhs: qarray.QArray,
     dimension_numbers: jax.lax.DotDimensionNumbers,
 ) -> jax.Array:
-  """Reshapes the 3D scaled_matmul output back to the expected target shape."""
+  """Reshapes the 3D scaled_dot output back to the expected target shape."""
   (lhs_ca, rhs_ca), (lhs_ba, rhs_ba) = dimension_numbers
   batch_shape = [lhs.shape[i] for i in lhs_ba]
   lhs_free_shape = [

tests/_src/core/dot_general_test.py

@@ -145,10 +145,12 @@ def test_output_dtype(self, lhs_dtype, rhs_dtype, expected_output_dtype):
         self.assertEqual(einsum_output.dtype, expected_output_dtype)
         self.assertEqual(dot_general_output.dtype, expected_output_dtype)
 
-  @mock.patch.object(jax.nn, 'scaled_matmul')
-  def test_outer_product(self, mock_scaled_matmul):
-    mock_scaled_matmul.return_value = jnp.ones((1, 10, 10), jnp.float32)
-
+  @mock.patch.object(jax.lax, 'scaled_dot')
+  def test_outer_product(self, mock_scaled_dot):
+    # An outer product has no contracting dims, so the flattened contracting
+    # size is 1, which doesn't satisfy jax.lax.scaled_dot's requirement that the
+    # contracting dim be at least twice the scale's. The mxfp fast path must
+    # decline and let the regular emulation handle it.
     lhs = qarray.QArray(
         jnp.ones((10, 1), jnp.float8_e4m3fn),
         jnp.ones((10, 1), jnp.bfloat16),
@@ -162,15 +164,9 @@ def test_outer_product(self, mock_scaled_matmul):
     dnums = (((), ()), ((), ()))
     res = dot_general.dot_general(lhs, rhs, dnums)
     self.assertEqual(res.shape, (10, 1, 1, 10))
+    self.assertTrue(jnp.array_equal(res, jnp.ones(res.shape, res.dtype)))
 
-    mock_scaled_matmul.assert_called_once()
-
-    args, _ = mock_scaled_matmul.call_args
-    lhs_3d, rhs_3d, lhs_scale_3d, rhs_scale_3d = args
-    self.assertEqual(lhs_3d.shape, (1, 10, 1))
-    self.assertEqual(rhs_3d.shape, (1, 10, 1))
-    self.assertEqual(lhs_scale_3d.shape, (1, 10, 1))
-    self.assertEqual(rhs_scale_3d.shape, (1, 10, 1))
+    mock_scaled_dot.assert_not_called()
 
   def test_innermost_tiling_heuristic(self):
     """Verifies that multi-dimensional dot_general picks the innermost contracting reduction axis."""

tests/_src/core/mxfp_dot_test.py

@@ -334,6 +334,45 @@ def test_mxfp_dot_general_returns_array(self):
     )
     self.assertIsNotNone(res)
 
+  def test_mxfp_dot_general_emits_scaled_dot_composite(self):
+    lhs = qarray.QArray(
+        qvalue=jnp.ones((2, 32), jnp.float8_e4m3fn),
+        scale=jnp.ones((2, 1)),
+        qtype="mxfp8",
+    )
+    rhs = qarray.QArray(
+        qvalue=jnp.ones((2, 32), jnp.float8_e4m3fn),
+        scale=jnp.ones((2, 1)),
+        qtype="mxfp8",
+    )
+
+    lowered = jax.jit(
+        functools.partial(
+            mxfp_dot.mxfp_dot_general,
+            dimension_numbers=(((1,), (1,)), ((), ())),
+        )
+    ).lower(lhs, rhs)
+    self.assertIn("xla.scaled_dot", lowered.as_text())
+
+  def test_unsupported_subchannel_ratio_fallback(self):
+    # jax.lax.scaled_dot requires the contracting dim to be at least twice the
+    # scale's contracting dim, which a per-element scale doesn't satisfy.
+    lhs = qarray.QArray(
+        qvalue=jnp.ones((2, 32), jnp.float8_e4m3fn),
+        scale=jnp.ones((2, 32)),
+        qtype="mxfp8",
+    )
+    rhs = qarray.QArray(
+        qvalue=jnp.ones((2, 32), jnp.float8_e4m3fn),
+        scale=jnp.ones((2, 32)),
+        qtype="mxfp8",
+    )
+
+    res = mxfp_dot.mxfp_dot_general(
+        lhs, rhs, dimension_numbers=(((1,), (1,)), ((), ()))
+    )
+    self.assertIsNone(res)
+
   def test_one_side_mxfp_fallback(self):
     lhs = qarray.QArray(
         qvalue=jnp.ones((2, 32), jnp.float8_e4m3fn),