EXLA.CustomCall: expose compiler-visible operation attributes

[humdrum00001010]

EXLA.CustomCall.Spec.attributes is emitted only as stablehlo.custom_call's opaque backend_config. Consequently, an external layout-sensitive, sharded custom call cannot express compiler-visible metadata such as operand_layouts, result_layouts, or an sdy.sharding_rule.

Would EXLA accept a generic way for Spec to attach validated top-level attributes to the emitted stablehlo.custom_call? This would keep native targets in external dependencies while allowing XLA layout assignment and Shardy to reason about them.

[humdrum00001010]

sorry. I made typo; this one belongs to xla!

[polvalente]

This actually belongs here! I believe we should add whatever attributes are needed for sharding to work. Do you have an example of what's needed?

[humdrum00001010]

I'm not sure on this abstraction but:

%EXLA.CustomCall.Spec{
  call_target_name: "flash_attention_3_forward",
  attributes: [{"causal", "true"}],
  operand_element_types: [{:bf, 16}, {:bf, 16}, {:bf, 16}, {:s, 32}],

  # Local native-buffer ABI; D is contiguous for BSHD tensors.
  operand_layouts: [[3, 2, 1, 0], [3, 2, 1, 0], [3, 2, 1, 0], [0]],
  result_layouts: [[3, 2, 1, 0], [2, 0, 1]],

  # Operation-level dimension relationship, independent of the concrete mesh.
  sharding_rule: "#sdy.op_sharding_rule<...>"
}

The FA3 implementation needs the three metadata fields, to control physical layout!

[humdrum00001010]

An example would be FA3 implementation. I tried to implement it in native, and it required physical layout control + shardy. There was no abstraction with these in custom_call and shard_jit.

[humdrum00001010]

Expected usage:

defimpl EXLA.CustomCall, for: FlashAttention3.Forward do
  def call(%{mode: mode}, _outputs, inputs, %{platform: :cuda}) do
    result_layouts =
      case mode do
        :padded -> [[3, 2, 1, 0], [2, 1, 0]]
        :varlen -> [[3, 2, 1, 0], [2, 0, 1]]
      end

    {:ok,
     %EXLA.CustomCall.Spec{
       call_target_name: "flash_attention_3_forward",
       operand_layouts: operand_layouts(inputs, mode),
       result_layouts: result_layouts,
       sharding_rule: sharding_rule(inputs)
     }}
  end
end

[polvalente]

I think we should add an :operation_attributes that is a list of {string, string} for name and value of the attribute.
This allows for any extensions that the StableHLO compiler needs (we do have a way to emit the stablehlo, so maybe another compiler has other annotations we don't know of).

[humdrum00001010]

I re-read your sentence. Um. why not fusing them in "attributes"? I'm not familiar with this custom_call, just asking.

[polvalente]

I re-read your sentence. Um. why not fusing them in "attributes"? I'm not familiar with this custom_call, just asking.

attributes goes into backend config and becomes part of the attributes the custom call can access from within the native implementation.

operation attributes (or a better name perhaps) would be the attributes annotated at the top level of the stablehlo.custom_call node

[humdrum00001010]

Would it be okay if I worked on this issue, or would you prefer to leave it open? If it’s available, I’d be happy to give it a try.

[polvalente]

Yeah, please send a PR!