EXLA: support custom call MLIR attributes by humdrum00001010 · Pull Request #1824 · elixir-nx/nx

humdrum00001010 · 2026-08-28T08:39:25Z

EXLA.CustomCall.Spec.attributes is serialized into backend_config, so external custom calls cannot attach compiler-visible metadata to stablehlo.custom_call.

Add mlir_attributes as {name, value} strings emitted at the operation level, while preserving attributes as handler-visible backend configuration. Duplicate and EXLA-owned attribute names are rejected. Tests cover validation, StableHLO emission, the existing backend configuration path, and the compiled execution path.

The external FlashAttention-3 experiment uses this API for layout constraints and a Shardy rule. Its performance measurements do not represent an Nx optimization: this change only makes compiler metadata expressible, and the measured throughput differences are not attributed to Nx or mlir_attributes.

Closes #1823

humdrum00001010 · 2026-08-29T17:34:23Z

One subtle issue: I think the LLM needs some documentation around the EXLA.NIF module.

After context compaction, it seems to lose track of how the native library is supposed to be loaded. In particular, it refuses to use EXLA.NIF.load_dylib() because, as it puts it, "it's private," and instead tries to call :erlang.load_nif() directly.

polvalente · 2026-08-29T17:50:12Z

One subtle issue: I think the LLM needs some documentation around the EXLA.NIF module.

After context compaction, it seems to lose track of how the native library is supposed to be loaded. In particular, it refuses to use EXLA.NIF.load_dylib() because, as it puts it, "it's private," and instead tries to call :erlang.load_nif() directly.

The LLM is correct in this case. That function is an auxiliary for EXLA.CustomCallTest.load_plugin!/0.

We could expose it in a separate PR in EXLA, but EXLA.NIF is meant to be a private module.

polvalente · 2026-08-29T17:51:28Z

      `name = ` (for example `{"k", "42 : i64"}`). An empty list omits the dictionary
      from the op.

+    * **`operation_attributes`** — Optional `{name, attr}` pairs, default `[]`, emitted


Suggested change

* **`operation_attributes`** — Optional `{name, attr}` pairs, default `[]`, emitted

* **`mlir_attributes`** — Optional `{name, attr}` pairs, default `[]`, emitted

Let's use mlir_attributes instead of the original suggestion? This is very coupled to the stablehlo implementation, so adding a disconnected name would be more confusing

polvalente · 2026-08-29T17:51:56Z

+    attributes =
+      Enum.map(pairs, fn
+        {name, attr} when is_binary(name) and is_binary(attr) ->
+          {String.to_atom(name), attr}


We shouldn't use String.to_atom blindly. This opens an attack surface for crashing the BEAM.
We should instead compare the strings against the reserved attributes

wow. I'm surprised by this comment

iex(1)> Enum.map(1..1_000_000, fn x -> String.to_atom("#{x}") end) [:"1", :"2", :"3", :"4", :"5", :"6", :"7", :"8", :"9", :"10", :"11", :"12", :"13", :"14", :"15", :"16", :"17", :"18", :"19", :"20", :"21", :"22", :"23", :"24", :"25", :"26", :"27", :"28", :"29", :"30", :"31", :"32", :"33", :"34", :"35", :"36", :"37", :"38", :"39", :"40", :"41", :"42", :"43", :"44", :"45", :"46", :"47", :"48", :"49", :"50", :"51", :"52", :"53", :"54", :"55", :"56", :"57", :"58", :"59", :"60", :"61", :"62", :"63", :"64", :"65", :"66", :"67", :"68", :"69", :"70", :"71", :"72", :"73", :"74", :"75", :"76", :"77", :"78", :"79", :"80", :"81", :"82", :"83", :"84", :"85", :"86", :"87", :"88", :"89", :"90", :"91", :"92", :"93", :"94", :"95", :"96", :"97", :"98", :"99", :"100", :"101", :"102", :"103", :"104", :"105", :"106", :"107", :"108", :"109", :"110", :"111", :"112", :"113", :"114", :"115", :"116", :"117", :"118", :"119", :"120", :"121", :"122", :"123", :"124", :"125", :"126", :"127", :"128", :"129", :"130", :"131", :"132", :"133", :"134", :"135", :"136", :"137", :"138", :"139", :"140", :"141", :"142", :"143", :"144", :"145", :"146", :"147", :"148", :"149", :"150", :"151", :"152", :"153", :"154", :"155", :"156", :"157", :"158", :"159", :"160", :"161", :"162", :"163", :"164", :"165", :"166", :"167", :"168", :"169", :"170", :"171", :"172", :"173", :"174", :"175", :"176", :"177", :"178", :"179", :"180", :"181", :"182", :"183", :"184", :"185", :"186", :"187", :"188", :"189", :"190", :"191", :"192", :"193", :"194", :"195", :"196", :"197", :"198", :"199", :"200", ...] iex(2)> Enum.map(1..10_000_000, fn x -> String.to_atom("#{x}") end) no more index entries in atom_tab (max=1048576) Crash dump is being written to: erl_crash.dump...done

You can look this up as "BEAM Atom Table Exhaustion"

polvalente · 2026-08-29T17:53:23Z

+          case spec.operation_attributes do
+            list when is_list(list) ->
+              list
+
+            other ->
+              raise ArgumentError,
+                    "EXLA.CustomCall.Spec operation_attributes must be a list of {binary_key, binary_attr} pairs, got: #{inspect(other)}"
+          end


if is_list(spec.operation_attributes)

polvalente · 2026-08-29T17:53:45Z

  import Nx.Defn

  alias EXLA.Test.QRAliasBlock
+  alias EXLA.MLIR.{Function, Module, Value}


Let's avoid single-line alias expansion in favor of 3 separate aliases

polvalente

Looks good overall! I just left a couple quick comments and we can merge after they're dealt with

polvalente · 2026-08-29T18:46:31Z

    |> attr_dict()
  end

+  @reserved_custom_call_attrs ~w(call_target_name api_version backend_config)


Suggested change

@reserved_custom_call_attrs ~w(call_target_name api_version backend_config)

@reserved_custom_call_attrs ["call_target_name", "api_version", "backend_config"]

polvalente · 2026-08-29T18:47:32Z

+            "custom_call MLIR attribute names must be unique and cannot override EXLA attributes"
+    end
+
+    Enum.map(attributes, fn {name, attr} -> {String.to_atom(name), attr} end)


String.to_atom is still something that can't go on here. If there's code requiring atoms that uses the results here, it needs to be adapted to accept only strings instead.

I changed exla.cc to accept string

This was referenced Aug 28, 2026

Expose FlashAttention-3 through an Nx custom-call boundary humdrum00001010/ex_flashattention3#1

Closed

Expose FlashAttention-3 through an Nx custom-call boundary humdrum00001010/ex_flashattention3#2

Merged

polvalente reviewed Aug 29, 2026

View reviewed changes

humdrum00001010 force-pushed the exla/custom-call-operation-attributes branch from 9f26fac to 9118b6a Compare August 29, 2026 18:19

humdrum00001010 mentioned this pull request Aug 29, 2026

EXLA: provide a public custom-call library loader #1825

Open

polvalente reviewed Aug 29, 2026

View reviewed changes

humdrum00001010 force-pushed the exla/custom-call-operation-attributes branch from 9118b6a to 1d42bc6 Compare August 29, 2026 19:12

humdrum00001010 changed the title ~~EXLA: support custom call operation attributes~~ EXLA: support custom call MLIR attributes Aug 29, 2026

humdrum00001010 added 2 commits August 30, 2026 13:02

Emit custom call MLIR attributes

4ec40b3

Wire custom call MLIR attributes

62ecad5

humdrum00001010 force-pushed the exla/custom-call-operation-attributes branch from 1d42bc6 to 62ecad5 Compare August 30, 2026 04:10

This was referenced Aug 30, 2026

Follow the upstream rename to mlir_attributes humdrum00001010/ex_flashattention3#3

Merged

Follow the upstream rename to mlir_attributes humdrum00001010/ex_flashattention3#5

Closed

Track EXLA's custom-call API humdrum00001010/ex_flashattention3#7

Merged

Navigation Menu

Uh oh!

EXLA: support custom call MLIR attributes - #1824

EXLA: support custom call MLIR attributes#1824
humdrum00001010 wants to merge 2 commits into
elixir-nx:mainfrom
humdrum00001010:exla/custom-call-operation-attributes

humdrum00001010 commented Aug 28, 2026 •

edited

Loading

Uh oh!

humdrum00001010 commented Aug 29, 2026

Uh oh!

polvalente commented Aug 29, 2026

Uh oh!

polvalente Aug 29, 2026

Uh oh!

polvalente Aug 29, 2026 •

edited

Loading

Uh oh!

humdrum00001010 Aug 29, 2026

Uh oh!

polvalente Aug 29, 2026

Uh oh!

polvalente Aug 29, 2026

Uh oh!

polvalente Aug 29, 2026

Uh oh!

polvalente left a comment

Uh oh!

polvalente Aug 29, 2026

Uh oh!

polvalente Aug 29, 2026

Uh oh!

humdrum00001010 Aug 29, 2026 •

edited

Loading

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

	* `operation_attributes` — Optional `{name, attr}` pairs, default `[]`, emitted
	* `mlir_attributes` — Optional `{name, attr}` pairs, default `[]`, emitted

	@reserved_custom_call_attrs ~w(call_target_name api_version backend_config)
	@reserved_custom_call_attrs ["call_target_name", "api_version", "backend_config"]

Uh oh!

Conversation

humdrum00001010 commented Aug 28, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

humdrum00001010 commented Aug 29, 2026

Uh oh!

polvalente commented Aug 29, 2026

Uh oh!

polvalente Aug 29, 2026

Choose a reason for hiding this comment

Uh oh!

polvalente Aug 29, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

humdrum00001010 Aug 29, 2026

Choose a reason for hiding this comment

Uh oh!

polvalente Aug 29, 2026

Choose a reason for hiding this comment

Uh oh!

polvalente Aug 29, 2026

Choose a reason for hiding this comment

Uh oh!

polvalente Aug 29, 2026

Choose a reason for hiding this comment

Uh oh!

polvalente left a comment

Choose a reason for hiding this comment

Uh oh!

polvalente Aug 29, 2026

Choose a reason for hiding this comment

Uh oh!

polvalente Aug 29, 2026

Choose a reason for hiding this comment

Uh oh!

humdrum00001010 Aug 29, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

humdrum00001010 commented Aug 28, 2026 •

edited

Loading

polvalente Aug 29, 2026 •

edited

Loading

humdrum00001010 Aug 29, 2026 •

edited

Loading