flash_attention: support also cross attention. #8427

dudulightricks · 2024-11-28T16:26:00Z

In case that q and kv have different shapes (cross attention) flash attention with spmd fails since it does not support it.

dudulightricks · 2024-11-28T16:26:18Z

@JackCaoG

In case that q and kv have different shapes (cross attention) flash attention with spmd fails since it does not support it.

dudulightricks · 2024-12-01T10:48:50Z

Rerun the tests when you can. @JackCaoG

dudulightricks · 2024-12-02T13:38:23Z

@JackCaoG when using XLA_DISABLE_FUNCTIONALIZATION=1, the flash attention backward tests are failing with error (unrelated to this PR specifically):


ERROR: test_flash_attention_backward_spmd_data_parallel (__main__.PallasTest)
----------------------------------------------------------------------
Traceback (most recent call last):
  File "/opt/xla/test/test_pallas_spmd.py", line 86, in test_flash_attention_backward_spmd_data_parallel
    loss.backward()
  File "/home/ubuntu/.local/lib/python3.10/site-packages/torch/_tensor.py", line 624, in backward
    torch.autograd.backward(
  File "/home/ubuntu/.local/lib/python3.10/site-packages/torch/autograd/__init__.py", line 347, in backward
    _engine_run_backward(
  File "/home/ubuntu/.local/lib/python3.10/site-packages/torch/autograd/graph.py", line 825, in _engine_run_backward
    return Variable._execution_engine.run_backward(  # Calls into the C++ engine to run the backward pass
RuntimeError: Bad StatusOr access: INTERNAL: RET_CHECK failure (third_party/tensorflow/compiler/xla/service/sharding_propagation.cc:1466) instruction->has_sharding() Sharding instruction must have a sharding attribute

JackCaoG · 2024-12-02T18:29:33Z

hmm, not the first time I saw this issue, let me see if I can do anything before I return my laptop...

JackCaoG · 2024-12-02T18:43:09Z

I can repo the issue with

XLA_DISABLE_FUNCTIONALIZATION=1, python test/test_pallas_spmd.py 

======================================================================
ERROR: test_flash_attention_backward_spmd_data_parallel (__main__.PallasTest)
----------------------------------------------------------------------
Traceback (most recent call last):
  File "/workspaces/dk3/pytorch/xla/test/test_pallas_spmd.py", line 85, in test_flash_attention_backward_spmd_data_parallel
    loss.backward()
  File "/workspaces/dk3/pytorch/torch/_tensor.py", line 626, in backward
    torch.autograd.backward(
  File "/workspaces/dk3/pytorch/torch/autograd/__init__.py", line 347, in backward
    _engine_run_backward(
  File "/workspaces/dk3/pytorch/torch/autograd/graph.py", line 825, in _engine_run_backward
    return Variable._execution_engine.run_backward(  # Calls into the C++ engine to run the backward pass
RuntimeError: Bad StatusOr access: INTERNAL: RET_CHECK failure (third_party/tensorflow/compiler/xla/service/sharding_propagation.cc:1464) instruction->has_sharding() Sharding instruction must have a sharding attribute

so this has nothing to do with segment ID.

I check XLA it is from
https://github.com/openxla/xla/blob/89358a8c397a9c6fc3ec5e4faad17d11da946a94/xla/service/sharding_propagation.cc#L1463-L1466
and this happens during compilation. The issue seems to be there is a HLO has the custom call Sharding but does not have sharding..

JackCaoG · 2024-12-02T19:20:31Z

there is something more fundemenrtal, it seems like during the backward there is a graph break, I saw

  %custom-call.58 = f32[1,2,128,8]{3,2,1,0} custom-call(f32[1,2,128,8]{3,2,1,0} %custom-call.57, f32[1,2,128,8]{3,2,1,0} %custom-call.56, f32[1,2,128,8]{3,2,1,0} %custom-call.55, f32[1,2,128,128]{3,2,1,0} %custom-call.54, f32[1,2,128,128]{3,2,1,0} %custom-call.44, /*index=5*/f32[1,2,128,8]{3,2,1,0} %custom-call.34, f32[1,2,128,128]{3,2,1,0} %custom-call.32), custom_call_target="tpu_custom_call", operand_layout_constraints={f32[1,2,128,8]{3,2,1,0}, f32[1,2,128,8]{3,2,1,0}, f32[1,2,128,8]{3,2,1,0}, f32[1,2,128,128]{3,2,1,0}, f32[1,2,128,128]{3,2,1,0}, f32[1,2,128,8]{3,2,1,0}, f32[1,2,128,128]{3,2,1,0}}, metadata={op_type="xla__tpu_custom_call" op_name="xla__tpu_custom_call" source_file="/workspaces/dk3/pytorch/xla/torch_xla/experimental/custom_kernel.py" source_line=423}, backend_config={"custom_call_config": {"body": "TUzvUgFNTElSMjAuMC4wZ2l0AAE9CQEDBQcBAwkDKQsNDxETFRcZGx0fISMlJykrLS8xA3YEAgQdAfkHFwsXFwsLCwsLCwsPCxsLExMTExMTE5MXDw8PEw8LEw8PkxMLCyMPCxcTDx8PxQ8PCwsLCwsLCwsTCwsTCwsLExMTCxcTLwsLCwsTCxMTDwsLExMTCxcPCxMPExsLD0MLGwuFC5MLCwsLKxsLGwsbCxsLGwsbCxsLGwsbBQ2NYXmRKgIqAgH5GxsbGxsbGw8TExcXCxcXFwsXFxcLFxcLGwsPExcLFwsXExMXFxMTExcbCw8TExcTExcTExcTExcTExcTExcTExcfExMXCwsPExcLDxcLExMXFx8LExMXDxMXDxMTFw8TFw8PExcXHy8TC1MTExMXDxMXExcfUw8TFxcLFwcFWVkBHQ8HHxsHKw8vIwsnI0MvAooYHwMDF7YCBTMDAxcmAhWCAo4CBTUFNwU5BTsFPQ0bBT8dPakFQQMDMgP+AwVDHRu6Ah2HxgIdh9ICHRveAh0b6gIdG/YCHRsCAyMNCUEBAAAAAAAAAAEAAAAAAAAAgAAAAAAAAAAIAAAAAAAAABU2AkICHX9HHX8/ERMAFZYDCR09OQVFFcoDMR09sx09tyMNCUEBAAAAAAAAAAEAAAAAAAAAgAAAAAAAAACAAAAAAAAAABV2AjEFRwVJAwUeA08iA08RAQUFSx1WA1oDFYYDCR0bOQMFobIDo6UdG6lhZmZpbmVfbWFwPChkMCwgZDEsIGQyLCBkMykgLT4gKGQwLCBkMSwgZDIsIGQzKT4AEQ0BEQ0FBU0FTwVRBVMFVQVXBVkFWx11KgIFXQVfAwN7XwVhBWMFZQMDe2EdfgIJHUmiAgVnAwMXDgMdSxIDAwmPAgKR/ZM3lTcFaQVrBW0FbxUqAwkFcRVqAwkdS3YDHUs5BXMFdSMBAQEddboDFdIDPwV3AwMX4gMdsbMFeRXmAz8dsbcV7gNHAwW7vQ0fBXsRDQ0DD8HDD8XJy81fz2EN0dPVBX0BCfn5+f8NGWFmZmluZV9tYXA8KGQwLCBkMSkgLT4gKGQwLCBkMSk+AAV/Iw0JQQEAAAAAAAAAAgAAAAAAAAABAAAAAAAAAAEAAAAAAAAABYEFgwWFBYcBEdfb3+Pn6+/zAwUR2RMvCWMDBRHdEy8JZQMFEeETLwlnAwUR5RNFCWkDBRHpE0UJawMFEe0TLwltAwUR8RNFCW8DBRH1Ey8JcQMFDxUNYyN0cHUuZGltZW5zaW9uX3NlbWFudGljczxwYXJhbGxlbD4AI3RwdS5tZW1vcnlfc3BhY2U8dm1lbT4AI3RwdS5jb250cmFjdF9wcmVjaXNpb248ZnAzMj4AI3RwdS5kaW1lbnNpb25fc2VtYW50aWNzPGFyYml0cmFyeT4AI3RwdS5kb3RfZGltZW5zaW9uX251bWJlcnM8WzFdLCBbMV0sIFswXSwgWzBdLCBbMCwgMCwgMSwgMF0sIFtdLCBbXT4AI3RwdS5kb3RfZGltZW5zaW9uX251bWJlcnM8WzFdLCBbMF0sIFswXSwgWzFdLCBbMCwgMCwgMSwgMV0sIFtdLCBbXT4AAwUPFQ1lAwUPFQ1nAwUPFQ1pAwUPFQ1rAwUPFQ1tAwUPFQ1vAwUPFQ1xEQEBFS4CMR0fMgIXBZISAR06Aj4CBYkXBR4XARVGAlICHUoCTgIFixd3NgIBFVYCYgIdWgJeAgWNF3cCBgEdZgJqAgWPF24CzgQBBZEdfUcdH3oCFwWWEgEFkx2GAooCBZUXBUYUARWSAjEdH5YCFwVCFAEDAxeeAhEBAgQVpgIJHQuqAhcFphIBAwOyAvoDBZcRAwEVvgIJHQvCAhcFqhIBFcoCCR0LzgIXBa4SARXWAgkdC9oCFwW6EgEV4gIJHQvmAhcFxhIBFe4CCR0L8gIXBcoSARX6AgkdC/4CFwXOEgEVBgMJHQsKAxcF0hIBJQUJAAAAABUWAwkdCxoDFwXaEgEFmQWbHVGXHQsuAxcFhhMBBZ0dmZcdPgNCAwWfFUYDCR0LSgMXBYITAQMDF1IDEwkQAADgDwWhFV4DCR0LYgMXBZITAR1Rmx0LbgMXBY4TAR1JmxV6AwkdC34DFwWqEwEdUVUdC4oDFwXCEwEdmVUdSVUdC5oDFwUCFAEDAxeiAyUHCQAAAAADCY8GApH9kzeVNx2uAzkFoyMBCSEBAAAAAQAAAAIAAAAAAAAAAwMXTxW+AzEdH8IDFwViFAEdfT8dH84DFwVmFAEdq9YDFwVqFAEDBaHeA6OlIwEJIQEAAAABAAAABAAAAAAAAAATCQEdq+oDFwVuFAEd8gP2AwWlFwWaEgEjYXJpdGgub3ZlcmZsb3c8bm9uZT4AI2FyaXRoLmZhc3RtYXRoPG5vbmU+AAECAgMnBQIEAgQJJwUCBCEJCxf7CQUFAgQhCV0BAgQX+wkFBQIEAgQJXScJBQUCBCEJAQknCQUFAgQCBAkX+wUCBCEJxwUbAQEBAQsLCw8PCw8LFwEFCQEBAQEJAQEBAQQ+EgUBEQG5BwMBJQcRAb8HA7umAhsBAQEBAQEBAQsBCwELAQ8BDwELAQ8BCwEXAQMDcwcDAQ0Hc3kDEwUHGxkGcgIDAQMdAwMzBwMBDQczgQMTBR8hGxQzAyMJAwsdAwO1rQMJFwa1AwcDuwMDQwMDAwMDQwMDAwUGQwMHBxm/wQ8FQ1kJvRm/wREAMwMBBREAMwMDgwcDAQMDhZoCAwEjB4WuAgMBBSUnAwMhAwMDAwMhAwMDAwMhAwMDAwMhAwMDBQYhAxELCSstLzELBiEDBwMzAwMjAwMDAwMjAwMDHQYjAwMDKQMDIwMDAwUGIwMRCws3OTs9CwYjAwcDPwMDJQMDAwMDJQMDAx0GJQMDAykDAyUDAwMFBiUDEQsNQ0VHSQsGJQMHA0sDAycDAwMDAycDAwMDAycDAwMDAycDAwMFBicDFQsPT1FTVQsGJwMFA1cDAykDAwMDAykDAwMDAykDAwMDAykDAwMFBikDFQsRW11fYQsGKQMFA2MDAysDAwMDAysDAwMDAysDAwMDAysDAwMFBisDEQsTZ2lrbQsGKwMHA28DAy0DAwMDAy0DAwMDAy0DAwMDAy0DAwMFBi0DFQsVc3V3eQsGLQMFA3sDA4uJAwUTB4uNAwUHNUF/FQcmA00DBQNlHwc2Ax0DBQWBgyUHOgMdAwUDhQMDU04DAwkXBlMDBQOJJwdTHQMFBYtZFQdmA00DBQONIQdyAx0DBQWHjwMDnYkDBRMHnY0DBQdxTZMVB4IDTQMFA30fB44DHQMFBZWXIQeSAx0DBQWZkQMDVwMDAwMDVwMDAwUGVwMHBxmdnwMDn54DAwcTB5+mAwMHB5tBoykHqgMdAwcFoaUDAzsDAwMDAzsDAwMFBjsDBwcZqasPBTtZCacZqasDA4O2AwMBAwOnBwMBDQeneQMTBQexGQbGAwMBA7MDAzUHAwENBzWBAxMFtbcbFDUDuQkDH0kDA1sDAwMDA1sDAwMFBlsDBwcZu70DAxkDAwMDAxkDAwMDAxkDAwMDAxkDAwMFBhkDEQsXwcPFxwsGGQMHA8kLBhkDEQO/DwUZ2gMNzRfBw8XHAwOvrQMJFwavAwcDzwMDQQMDAwMDQQMDAwUGQQMHBxnT1Q8FQVkJ0RnT1REANQMBBREANQkAAQcRAfcHAw0PCQEBAQEBAQEBAwMBBwMBAwMBBwMBCQQBCQEDBQkHEQEKAgcDDQ8JAQEBAQEBAQEDAwEHAwEDAwEHAwEJBAEJAQMHCQcRAQ4CBwMNDwkBAQEBAQEBAQMDAQcDAQMDAQcDAQkEAQkBAwcJBxEBEgIHAw0PCQEBAQEBAQEBAwMBBwMBAwMBBwMBCQQBCQEDBQkHEQEWAgcDDQ8JAQEBAQEBAQEDAwEHAwEDAwEHAwEJBAEJAQMFCQcRARoCBwMNDwkBAQEBAQEBAQMDAQcDAQMDAQcDAQkEAQkBAwUJBxEBHgIHAw0PCQEBAQEBAQEBAwMBBwMBAwMBBwMBCQQBCQEDBQkHEQEiAgcDDQ8JAQEBAQEBAQEDAwEHAwEDAwEHAwEJBAEJAQMFCQYDAQUBADoTpycLCwsTDRUdCQ1nDRMbMR0LIyEjKS0lJxEpCx0dFSUbDS0ViQkZGRkZGRkZGREbCw03Cw0dJR0TC7cXFxMXFxcjDxkjFxcVIxclGRUZHw8NCR0RYnVpbHRpbgBzdGFibGVfbW9zYWljAHRwdQBhcml0aABtb2R1bGUAYXJpdGguY29uc3RhbnQAdmVjdG9yLmxvYWQAZnVuYy5mdW5jAGZ1bmMucmV0dXJuAHZlY3Rvci5zaGFwZV9jYXN0AGFyaXRoLmNtcGkAdHB1LnZlY3Rvcl9zdG9yZQBzY2YueWllbGQAdHB1Lm1hdG11bAB0cHUucmVwZWF0AHZlY3Rvci5icm9hZGNhc3QAYXJpdGguZXh0dWkAc2NmLmlmAGFyaXRoLmluZGV4X2Nhc3QAYXJpdGguc3ViZgBhcml0aC5tdWxmAGFyaXRoLm11bGkAbWF0aC5leHAAYXJpdGguZGl2ZgBhcml0aC5hZGRmAC91c3IvbG9jYWwvbGliL3B5dGhvbjMuMTAvc2l0ZS1wYWNrYWdlcy9qYXgvZXhwZXJpbWVudGFsL3BhbGxhcy9vcHMvdHB1L2ZsYXNoX2F0dGVudGlvbi5weQBib2R5AHN5bV9uYW1lAGZ1bmN0aW9uX3R5cGUAdHJhbnNmb3JtX2luZGljZXMAd2luZG93X2JvdW5kcwB2YWx1ZQAvZ2V0AF9mbGFzaF9hdHRlbnRpb25fZHFfa2VybmVsAC9zd2FwAC9tdWwAL2RvdF9nZW5lcmFsAC9yZXBlYXQAdHJhbnNmb3JtXzAAdHJhbnNmb3JtXzEAdHJhbnNmb3JtXzIAdHJhbnNmb3JtXzMAdHJhbnNmb3JtXzQAdHJhbnNmb3JtXzUAdHJhbnNmb3JtXzYAdHJhbnNmb3JtXzcAL2VxAC93b3Jrc3BhY2VzL2RrMy9weXRvcmNoL3hsYS90b3JjaF94bGEvZXhwZXJpbWVudGFsL2N1c3RvbV9rZXJuZWwucHkAcHJlZGljYXRlAC9jb252ZXJ0X2VsZW1lbnRfdHlwZQAvY29uZAAvbWFza2VkX2xvYWQAZGltZW5zaW9uX251bWJlcnMAcHJlY2lzaW9uAHRyYW5zcG9zZV9saHMAdHJhbnNwb3NlX3JocwAvc3ViAG9wZXJhbmRTZWdtZW50U2l6ZXMAc3RyaWRlcwBlbmRfb2Zfa3Zfc2VxdWVuY2UAL2Jyb2FkY2FzdF9pbl9kaW0Ac3RhYmxlX21vc2FpYy52ZXJzaW9uAGRpbWVuc2lvbl9zZW1hbnRpY3MAaXRlcmF0aW9uX2JvdW5kcwBzY2FsYXJfcHJlZmV0Y2gAc2NyYXRjaF9vcGVyYW5kcwBtYWluAHdpbmRvd19wYXJhbXMAX2ZsYXNoX2F0dGVudGlvbl9id2RfZHEAdHJhY2VfcGFsbGFzAGJhY2t3YXJkAGFwcGx5AC93b3Jrc3BhY2VzL2RrMy9weXRvcmNoL3RvcmNoL2F1dG9ncmFkL2Z1bmN0aW9uLnB5AC9zY2FuAHJ1bgBvdmVyZmxvd0ZsYWdzAGRpbWVuc2lvbgB0aW1lcwBmYXN0bWF0aAAvZXhwAC9kaXYAL2FkZABzdGFydF9uZXdfc2VxdWVuY2UA", "serialization_format": 1, "needs_layout_passes": true}}
  %custom-call.59 = f32[1,2,128,8]{3,2,1,0} custom-call(f32[1,2,128,8]{3,2,1,0} %custom-call.58), custom_call_target="Sharding", sharding={manual}, metadata={op_type="xla__custom_sharding" op_name="xla__custom_sharding" source_file="/workspaces/dk3/pytorch/xla/torch_xla/distributed/spmd/xla_sharding.py" source_line=537}
  %custom-call.60 = f32[4,2,128,8]{3,2,1,0} custom-call(f32[1,2,128,8]{3,2,1,0} %custom-call.59), custom_call_target="SPMDShardToFullShape", sharding={devices=[4,1,1,1]0,1,2,3}, metadata={op_type="xla__custom_sharding" op_name="xla__custom_sharding" source_file="/workspaces/dk3/pytorch/xla/torch_xla/distributed/spmd/xla_sharding.py" source_line=570}
  ROOT %tuple.61 = (f32[4,2,128,8]{3,2,1,0}) tuple(f32[4,2,128,8]{3,2,1,0} %custom-call.60)
}

before any mark_step, this is the cause of the crash above.. I tried to revert this pr but I still see the same graph break/ Need to first figure out where this is coming from.

JackCaoG · 2024-12-02T19:33:43Z

ok I found an even better repo, if I just mark_step in the fwd it will crash too

HloModule IrToHlo.45, entry_computation_layout={(f32[4,2,128,8]{2,3,0,1}, f32[4,2,128,8]{2,3,0,1}, f32[4,2,128,8]{2,3,0,1})->(f32[4,2,128,8]{3,2,1,0}, f32[4,2,128]{2,1,0}, f32[4,2,128]{2,1,0}, f32[])}

%AddComputation.39 (x.40: f32[], y.41: f32[]) -> f32[] {
  %x.40 = f32[] parameter(0)
  %y.41 = f32[] parameter(1)
  ROOT %add.42 = f32[] add(f32[] %x.40, f32[] %y.41)
}

ENTRY %IrToHlo.45 (p0.1: f32[4,2,128,8], p1.3: f32[4,2,128,8], p2.5: f32[4,2,128,8]) -> (f32[4,2,128,8], f32[4,2,128], f32[4,2,128], f32[]) {
  %constant.38 = s32[] constant(8192), metadata={op_type="aten__sum" op_name="aten__sum" source_file="/workspaces/dk3/pytorch/xla/test/test_pallas_spmd.py" source_line=84}
  %p2.5 = f32[4,2,128,8]{2,3,0,1} parameter(2), sharding={devices=[4,1,1,1]0,1,2,3}, metadata={op_type="xla__device_data" op_name="xla__device_data" source_file="/workspaces/dk3/pytorch/xla/torch_xla/distributed/spmd/xla_sharding.py" source_line=625}
  %custom-call.6 = f32[1,2,128,8]{3,2,1,0} custom-call(f32[4,2,128,8]{2,3,0,1} %p2.5), custom_call_target="SPMDFullToShardShape", sharding={manual}, metadata={op_type="xla__custom_sharding" op_name="xla__custom_sharding" source_file="/workspaces/dk3/pytorch/xla/torch_xla/distributed/spmd/xla_sharding.py" source_line=555}
  %p1.3 = f32[4,2,128,8]{2,3,0,1} parameter(1), sharding={devices=[4,1,1,1]0,1,2,3}, metadata={op_type="xla__device_data" op_name="xla__device_data" source_file="/workspaces/dk3/pytorch/xla/torch_xla/distributed/spmd/xla_sharding.py" source_line=625}
  %custom-call.4 = f32[1,2,128,8]{3,2,1,0} custom-call(f32[4,2,128,8]{2,3,0,1} %p1.3), custom_call_target="SPMDFullToShardShape", sharding={manual}, metadata={op_type="xla__custom_sharding" op_name="xla__custom_sharding" source_file="/workspaces/dk3/pytorch/xla/torch_xla/distributed/spmd/xla_sharding.py" source_line=555}
  %p0.1 = f32[4,2,128,8]{2,3,0,1} parameter(0), sharding={devices=[4,1,1,1]0,1,2,3}, metadata={op_type="xla__device_data" op_name="xla__device_data" source_file="/workspaces/dk3/pytorch/xla/torch_xla/distributed/spmd/xla_sharding.py" source_line=625}
  %custom-call.2 = f32[1,2,128,8]{3,2,1,0} custom-call(f32[4,2,128,8]{2,3,0,1} %p0.1), custom_call_target="SPMDFullToShardShape", sharding={manual}, metadata={op_type="xla__custom_sharding" op_name="xla__custom_sharding" source_file="/workspaces/dk3/pytorch/xla/torch_xla/distributed/spmd/xla_sharding.py" source_line=555}
  %custom-call.7 = (f32[1,2,128,8]{3,2,1,0}, f32[1,2,128,128]{3,2,1,0}, f32[1,2,128,128]{3,2,1,0}) custom-call(f32[1,2,128,8]{3,2,1,0} %custom-call.6, f32[1,2,128,8]{3,2,1,0} %custom-call.4, f32[1,2,128,8]{3,2,1,0} %custom-call.2), custom_call_target="tpu_custom_call", operand_layout_constraints={f32[1,2,128,8]{3,2,1,0}, f32[1,2,128,8]{3,2,1,0}, f32[1,2,128,8]{3,2,1,0}}, metadata={op_type="xla__tpu_custom_call" op_name="xla__tpu_custom_call" source_file="/workspaces/dk3/pytorch/xla/torch_xla/experimental/custom_kernel.py" source_line=309}, backend_config={"custom_call_config": {"body": "TUzvUgFNTElSMjAuMC4wZ2l0AAEvCwEDBQcJAQMLAxkNDxETFRcZGx0fISMDYgPiAh8B+QcTCxMLDwsPDw8LCwsLCw8TE5MPCwsLGwsrC8UPCwsLCwuTCwsTCwsLCwsTEwsLMxMTExMXDxMPEw8TGwsPQwsbCwuTCwsLCyMbCxsLGwsbCxsLGwsbGxsbGxsPDw8PFw8XDw8LFw8PCxcPDwsXDw8LFwsPDwsXDw8LExcFDY1heZEqAioCAWcLFxMTFxcfExMXLxcfCxMXHw8LExcLEwsXHwsTFx8PCxMTFxMXC1MLExMXExMXFx8TFy8HBV1JCQNZAR8PBx8rBxsPFyMvGyc3LwsCghMfAwMduQUlAwMdtwUnFcHFBSkdMW0dMXEdMXUFKwUtBS8NGwUxHSu7HSsSAh0rtgIjDQlBAQAAAAAAAAABAAAAAAAAAIAAAAAAAAAACAAAAAAAAAARHQAFMwU1BTcDA14C3gIFOQMFngKiAqYCqgIFO2FmZmluZV9tYXA8KGQwLCBkMSwgZDIsIGQzKSAtPiAoZDAsIGQxLCBkMiwgZDMpPgARDQEFPQU/BUEFQwVFIw0JQQEAAAAAAAAAAQAAAAAAAACAAAAAAAAAAIAAAAAAAAAABUcFSR1NJgIFSwVNBU8FUQVTHT4CWRVCAgsFVQVXIw0DEQEAAAAAAAAAHVICYxVWAgsdcgJnFXYCCx2GAooCHSltFZYCCx0pcRWuAgsdTXUVygILAwV5ew01BVkRDQ0DD3+BFYOFh4k5izkNjY+RBVsBCfn5+f8NGQVdIw0JQQEAAAAAAAAAAgAAAAAAAAABAAAAAAAAAAEAAAAAAAAABV8FYQVjBWUBDZOXm5+jpwMFF5UZJQk7AwUXmRklCT0DBRedGSUJPwMFF6EZJQlBAwUXpRlFCUMDBRepGUUJRwMFFRsNOwMFFRsNPQMFFRsNPwMFFRsNQQMFFRsNQwMFFRsNRxEBAREDARW9Cx0JvxcFCggBHTXDFwVOBQEVx80dycsFZxcF/gsBFc/VHdHTBWkXLTYCARXX3R3Z2wVrFy12BAEV3+cd4eMFbRfl/ggBBW8V6e8d6+0FcRcttgcBFfH3HfP1BXMXSacBHQoCDgIjdHB1LmRpbWVuc2lvbl9zZW1hbnRpY3M8cGFyYWxsZWw+ACN0cHUubWVtb3J5X3NwYWNlPHZtZW0+ACN0cHUuY29udHJhY3RfcHJlY2lzaW9uPGZwMzI+ACN0cHUuZGltZW5zaW9uX3NlbWFudGljczxhcmJpdHJhcnk+ACN0cHUuZG90X2RpbWVuc2lvbl9udW1iZXJzPFsxXSwgWzFdLCBbMF0sIFswXSwgWzAsIDAsIDEsIDBdLCBbXSwgW10+ACN0cHUuZG90X2RpbWVuc2lvbl9udW1iZXJzPFsxXSwgWzBdLCBbMF0sIFsxXSwgWzAsIDAsIDEsIDFdLCBbXSwgW10+AAV1F0mCBQEVFgILHQkaAhcFDggBAwMdIgIlBQkAAAAAFSoCCx0JLgIXBRIIAQMJTwICUf1TJ1UnAwMdOgIlDwkAAID/BXcdCUYCFwWmCAEDBVvWAl1fHSlZBXkdCVoCFwWqCAEFex1mAmMFfQMDHW4CJQ8JAAAAAAV/HQl6AhcFrggBAwVb2gJdXx0pZwWBFY4CCx0JkgIXBbIIAR0JmgIXBb4IAQWDIwEJIQEAAAABAAAABAAAAAAAAAAFhSMBAQEdCbICFwXGCAEVugILHQm+AhcFzggBAwMdxgIlCwkAAAAAHQnOAhcF0ggBAwlPBgJR/VMnVScjdmVjdG9yLmtpbmQ8bWF4aW11bWY+ACN2ZWN0b3Iua2luZDxhZGQ+ACNhcml0aC5mYXN0bWF0aDxub25lPgABAgIDJwUCBAIECRf7CQUFAgQhCTcLJwUCBCEJAQIEJwMCBAknCQUFAgQhCRf7CQUFAgQCBAk3JwUCBAUJJwkFBQIEAgQJBRUBAQEBBwcHBxMTAQUJAQEBAQkBAQEBAQkEUgsFAREBdwcDAR0HEQF9BwOJ+xUBAQEBAQEBAQcBBwEHAQcBEwETAQMDHwMDAwMDHwMDAwMDHwMDAwMDHwMDAwsGHwMRCwkVFxkbBQYfAwsDHQMDIQMDAwMDIQMDAwMDIQMDAwMDIQMDAwsGIQMRCwshIyUnBQYhAwsDKQMDSx4CAwURB0syAgMFBx8rLQMDVzYCAw8TB1dKAgMPBS8xBQZOAgMVAzMNBmEDBQM1FQdhLwMFBS83FwdiAi8DBQM5AwNlagIDDxMHZX4CAw8FOz0FBoICAxUDPw0GaQMFA0EZB2kvAwUFO0MFBmsDFQM1DQZrAwUDRwMDDwMDAwMDDwMDAwMDDwMDAwMDDwMDAwsGDwMXCxNLTU9RBQYPAwUDUwUGDwMXA0kPBQ8zDVcTS01PUQUGbwMVA0ENBm8DBQNZAwMRAwMDAwMRAwMDAwMRAwMDAwMRAwMDCwYRAxcLEV1fYWMFBhEDBQNlBQYRAxcDWw8FETMNaRFdX2FjAwMjAwMDAwMjAwMDAwMjAwMDAwMjAwMDCwYjAxELDWttb3EFBiMDCwNzAwNzwgIDCxEHc9ICAwsHRXV3AwMTAwMDAwMTAwMDAwMTAwMDAwMTAwMDCwYTAxELD3t9f4EFBhMDCwODBQYTAxEDeQ8FEzMNhw97fX+BCQABBxEBqwcDDQ8JAQEBAQEBAQEDAwEHAwEDAwEHAwEJBAEJAQMFCQcRAa0HAw0PCQEBAQEBAQEBAwMBBwMBAwMBBwMBCQQBCQEDBwkHEQGvBwMNDwkBAQEBAQEBAQMDAQcDAQMDAQcDAQkEAQkBAwcJBxEBsQcDDQ8JAQEBAQEBAQEDAwEHAwEDAwEHAwEJBAEJAQMFCQcRAbMHAw0PCQEBAQEBAQEBAwMBBwMBAwMBBwMBCQQBCQEDBQkHEQG1BwMNDwkBAQEBAQEBAQMDAQcDAQMDAQcDAQkEAQkBAwUJBgMBBQEAZhKHESkLGQsTCxkTYyFnDREbLR0LIyEjKS0fCx0dFSUbaxkZGRkZGTENiQslDR0lHRNjtxcTFy8XIyMZGRUlHw8NDwkdEWJ1aWx0aW4Ac3RhYmxlX21vc2FpYwB0cHUAdmVjdG9yAGFyaXRoAG1vZHVsZQBhcml0aC5jb25zdGFudAB2ZWN0b3Iuc2hhcGVfY2FzdABmdW5jLmZ1bmMAZnVuYy5yZXR1cm4AdmVjdG9yLmxvYWQAdmVjdG9yLmJyb2FkY2FzdAB0cHUudmVjdG9yX3N0b3JlAHRwdS5tYXRtdWwAdmVjdG9yLm11bHRpX3JlZHVjdGlvbgBhcml0aC5zdWJmAG1hdGguZXhwAGFyaXRoLmRpdmYAL3Vzci9sb2NhbC9saWIvcHl0aG9uMy4xMC9zaXRlLXBhY2thZ2VzL2pheC9leHBlcmltZW50YWwvcGFsbGFzL29wcy90cHUvZmxhc2hfYXR0ZW50aW9uLnB5AF9mbGFzaF9hdHRlbnRpb25fa2VybmVsX3NpbmdsZV9iYXRjaF9zaW5nbGVfc3RlcABzeW1fbmFtZQBmdW5jdGlvbl90eXBlAHRyYW5zZm9ybV9pbmRpY2VzAHdpbmRvd19ib3VuZHMAdmFsdWUAL2Jyb2FkY2FzdF9pbl9kaW0AL2dldAAvd29ya3NwYWNlcy9kazMvcHl0b3JjaC94bGEvdG9yY2hfeGxhL2V4cGVyaW1lbnRhbC9jdXN0b21fa2VybmVsLnB5AC9zd2FwAF9mbGFzaF9hdHRlbnRpb25fa2VybmVsAHRyYW5zZm9ybV8wAHRyYW5zZm9ybV8xAHRyYW5zZm9ybV8yAHRyYW5zZm9ybV8zAHRyYW5zZm9ybV80AHRyYW5zZm9ybV81AC93b3Jrc3BhY2VzL2RrMy9weXRvcmNoL3hsYS90ZXN0L3Rlc3RfcGFsbGFzX3NwbWQucHkAL2RvdF9nZW5lcmFsAGRpbWVuc2lvbl9udW1iZXJzAHByZWNpc2lvbgB0cmFuc3Bvc2VfbGhzAHRyYW5zcG9zZV9yaHMAa2luZAByZWR1Y3Rpb25fZGltcwBzdGFibGVfbW9zYWljLnZlcnNpb24AZGltZW5zaW9uX3NlbWFudGljcwBpdGVyYXRpb25fYm91bmRzAHNjYWxhcl9wcmVmZXRjaABzY3JhdGNoX29wZXJhbmRzAG1haW4Ad2luZG93X3BhcmFtcwBfZmxhc2hfYXR0ZW50aW9uX2ltcGwAdHJhY2VfcGFsbGFzAGZvcndhcmQAYXBwbHkAL3dvcmtzcGFjZXMvZGszL3B5dG9yY2gvdG9yY2gvYXV0b2dyYWQvZnVuY3Rpb24ucHkAZmxhc2hfYXR0ZW50aW9uAHRlc3RfZmxhc2hfYXR0ZW50aW9uX2JhY2t3YXJkX3NwbWRfZGF0YV9wYXJhbGxlbAA8bW9kdWxlPgAvcmVkdWNlX21heAAvc3ViAGZhc3RtYXRoAC9leHAAL3JlZHVjZV9zdW0AL2RpdgBvcGVyYW5kU2VnbWVudFNpemVzAHN0cmlkZXMA", "cost_estimate": {"flops": 1114112, "transcendentals": 32768, "bytes_accessed": 294912}, "serialization_format": 1, "needs_layout_passes": true}}
  %get-tuple-element.8 = f32[1,2,128,8]{3,2,1,0} get-tuple-element((f32[1,2,128,8]{3,2,1,0}, f32[1,2,128,128]{3,2,1,0}, f32[1,2,128,128]{3,2,1,0}) %custom-call.7), index=0, metadata={op_type="xla__tpu_custom_call" op_name="xla__tpu_custom_call" source_file="/workspaces/dk3/pytorch/xla/torch_xla/experimental/custom_kernel.py" source_line=309}
  %custom-call.11 = f32[1,2,128,8]{3,2,1,0} custom-call(f32[1,2,128,8]{3,2,1,0} %get-tuple-element.8), custom_call_target="Sharding", sharding={manual}, metadata={op_type="xla__custom_sharding" op_name="xla__custom_sharding" source_file="/workspaces/dk3/pytorch/xla/torch_xla/distributed/spmd/xla_sharding.py" source_line=537}
  %custom-call.12 = f32[4,2,128,8]{3,2,1,0} custom-call(f32[1,2,128,8]{3,2,1,0} %custom-call.11), custom_call_target="SPMDShardToFullShape", sharding={devices=[4,1,1,1]0,1,2,3}, metadata={op_type="xla__custom_sharding" op_name="xla__custom_sharding" source_file="/workspaces/dk3/pytorch/xla/torch_xla/distributed/spmd/xla_sharding.py" source_line=570}
  %get-tuple-element.9 = f32[1,2,128,128]{3,2,1,0} get-tuple-element((f32[1,2,128,8]{3,2,1,0}, f32[1,2,128,128]{3,2,1,0}, f32[1,2,128,128]{3,2,1,0}) %custom-call.7), index=1, metadata={op_type="xla__tpu_custom_call" op_name="xla__tpu_custom_call" source_file="/workspaces/dk3/pytorch/xla/torch_xla/experimental/custom_kernel.py" source_line=309}
  %slice.13 = f32[1,2,128,1]{3,2,1,0} slice(f32[1,2,128,128]{3,2,1,0} %get-tuple-element.9), slice={[0:1], [0:2], [0:128], [0:1]}, metadata={op_type="xla__generic_slice" op_name="xla__generic_slice" source_file="/workspaces/dk3/pytorch/xla/torch_xla/distributed/spmd/xla_sharding.py" source_line=537}
  %reshape.14 = f32[1,2,128]{2,1,0} reshape(f32[1,2,128,1]{3,2,1,0} %slice.13), metadata={op_type="aten__view" op_name="aten__view" source_file="/workspaces/dk3/pytorch/xla/torch_xla/distributed/spmd/xla_sharding.py" source_line=537}
  %custom-call.15 = f32[1,2,128]{2,1,0} custom-call(f32[1,2,128]{2,1,0} %reshape.14), custom_call_target="Sharding", metadata={op_type="xla__custom_sharding" op_name="xla__custom_sharding" source_file="/workspaces/dk3/pytorch/xla/torch_xla/distributed/spmd/xla_sharding.py" source_line=537}
  %reshape.16 = f32[1,2,128,1]{3,2,1,0} reshape(f32[1,2,128]{2,1,0} %custom-call.15), metadata={op_type="aten__view" op_name="aten__view" source_file="/workspaces/dk3/pytorch/xla/torch_xla/distributed/spmd/xla_sharding.py" source_line=537}
  %constant.17 = s64[] constant(0), metadata={op_type="xla__update_slice" op_name="xla__update_slice" source_file="/workspaces/dk3/pytorch/xla/torch_xla/distributed/spmd/xla_sharding.py" source_line=537}
  %constant.18 = s64[] constant(0), metadata={op_type="xla__update_slice" op_name="xla__update_slice" source_file="/workspaces/dk3/pytorch/xla/torch_xla/distributed/spmd/xla_sharding.py" source_line=537}
  %constant.19 = s64[] constant(0), metadata={op_type="xla__update_slice" op_name="xla__update_slice" source_file="/workspaces/dk3/pytorch/xla/torch_xla/distributed/spmd/xla_sharding.py" source_line=537}
  %constant.20 = s64[] constant(0), metadata={op_type="xla__update_slice" op_name="xla__update_slice" source_file="/workspaces/dk3/pytorch/xla/torch_xla/distributed/spmd/xla_sharding.py" source_line=537}
  %dynamic-update-slice.21 = f32[1,2,128,128]{3,2,1,0} dynamic-update-slice(f32[1,2,128,128]{3,2,1,0} %get-tuple-element.9, f32[1,2,128,1]{3,2,1,0} %reshape.16, s64[] %constant.17, s64[] %constant.18, s64[] %constant.19, /*index=5*/s64[] %constant.20), metadata={op_type="xla__update_slice" op_name="xla__update_slice" source_file="/workspaces/dk3/pytorch/xla/torch_xla/distributed/spmd/xla_sharding.py" source_line=537}
  %slice.22 = f32[1,2,128,1]{3,2,1,0} slice(f32[1,2,128,128]{3,2,1,0} %dynamic-update-slice.21), slice={[0:1], [0:2], [0:128], [0:1]}, metadata={op_type="xla__generic_slice" op_name="xla__generic_slice" source_file="/workspaces/dk3/pytorch/xla/torch_xla/distributed/spmd/xla_sharding.py" source_line=537}
  %reshape.23 = f32[1,2,128]{2,1,0} reshape(f32[1,2,128,1]{3,2,1,0} %slice.22), sharding={manual}, metadata={op_type="aten__view" op_name="aten__view" source_file="/workspaces/dk3/pytorch/xla/torch_xla/distributed/spmd/xla_sharding.py" source_line=537}
  %custom-call.24 = f32[4,2,128]{2,1,0} custom-call(f32[1,2,128]{2,1,0} %reshape.23), custom_call_target="SPMDShardToFullShape", sharding={devices=[4,1,1]0,1,2,3}, metadata={op_type="xla__custom_sharding" op_name="xla__custom_sharding" source_file="/workspaces/dk3/pytorch/xla/torch_xla/distributed/spmd/xla_sharding.py" source_line=570}
  %get-tuple-element.10 = f32[1,2,128,128]{3,2,1,0} get-tuple-element((f32[1,2,128,8]{3,2,1,0}, f32[1,2,128,128]{3,2,1,0}, f32[1,2,128,128]{3,2,1,0}) %custom-call.7), index=2, metadata={op_type="xla__tpu_custom_call" op_name="xla__tpu_custom_call" source_file="/workspaces/dk3/pytorch/xla/torch_xla/experimental/custom_kernel.py" source_line=309}
  %slice.25 = f32[1,2,128,1]{3,2,1,0} slice(f32[1,2,128,128]{3,2,1,0} %get-tuple-element.10), slice={[0:1], [0:2], [0:128], [0:1]}, metadata={op_type="xla__generic_slice" op_name="xla__generic_slice" source_file="/workspaces/dk3/pytorch/xla/torch_xla/distributed/spmd/xla_sharding.py" source_line=537}
  %reshape.26 = f32[1,2,128]{2,1,0} reshape(f32[1,2,128,1]{3,2,1,0} %slice.25), metadata={op_type="aten__view" op_name="aten__view" source_file="/workspaces/dk3/pytorch/xla/torch_xla/distributed/spmd/xla_sharding.py" source_line=537}
  %custom-call.27 = f32[1,2,128]{2,1,0} custom-call(f32[1,2,128]{2,1,0} %reshape.26), custom_call_target="Sharding", metadata={op_type="xla__custom_sharding" op_name="xla__custom_sharding" source_file="/workspaces/dk3/pytorch/xla/torch_xla/distributed/spmd/xla_sharding.py" source_line=537}
  %reshape.28 = f32[1,2,128,1]{3,2,1,0} reshape(f32[1,2,128]{2,1,0} %custom-call.27), metadata={op_type="aten__view" op_name="aten__view" source_file="/workspaces/dk3/pytorch/xla/torch_xla/distributed/spmd/xla_sharding.py" source_line=537}
  %constant.29 = s64[] constant(0), metadata={op_type="xla__update_slice" op_name="xla__update_slice" source_file="/workspaces/dk3/pytorch/xla/torch_xla/distributed/spmd/xla_sharding.py" source_line=537}
  %constant.30 = s64[] constant(0), metadata={op_type="xla__update_slice" op_name="xla__update_slice" source_file="/workspaces/dk3/pytorch/xla/torch_xla/distributed/spmd/xla_sharding.py" source_line=537}
  %constant.31 = s64[] constant(0), metadata={op_type="xla__update_slice" op_name="xla__update_slice" source_file="/workspaces/dk3/pytorch/xla/torch_xla/distributed/spmd/xla_sharding.py" source_line=537}
  %constant.32 = s64[] constant(0), metadata={op_type="xla__update_slice" op_name="xla__update_slice" source_file="/workspaces/dk3/pytorch/xla/torch_xla/distributed/spmd/xla_sharding.py" source_line=537}
  %dynamic-update-slice.33 = f32[1,2,128,128]{3,2,1,0} dynamic-update-slice(f32[1,2,128,128]{3,2,1,0} %get-tuple-element.10, f32[1,2,128,1]{3,2,1,0} %reshape.28, s64[] %constant.29, s64[] %constant.30, s64[] %constant.31, /*index=5*/s64[] %constant.32), metadata={op_type="xla__update_slice" op_name="xla__update_slice" source_file="/workspaces/dk3/pytorch/xla/torch_xla/distributed/spmd/xla_sharding.py" source_line=537}
  %slice.34 = f32[1,2,128,1]{3,2,1,0} slice(f32[1,2,128,128]{3,2,1,0} %dynamic-update-slice.33), slice={[0:1], [0:2], [0:128], [0:1]}, metadata={op_type="xla__generic_slice" op_name="xla__generic_slice" source_file="/workspaces/dk3/pytorch/xla/torch_xla/distributed/spmd/xla_sharding.py" source_line=537}
  %reshape.35 = f32[1,2,128]{2,1,0} reshape(f32[1,2,128,1]{3,2,1,0} %slice.34), sharding={manual}, metadata={op_type="aten__view" op_name="aten__view" source_file="/workspaces/dk3/pytorch/xla/torch_xla/distributed/spmd/xla_sharding.py" source_line=537}
  %custom-call.36 = f32[4,2,128]{2,1,0} custom-call(f32[1,2,128]{2,1,0} %reshape.35), custom_call_target="SPMDShardToFullShape", sharding={devices=[4,1,1]0,1,2,3}, metadata={op_type="xla__custom_sharding" op_name="xla__custom_sharding" source_file="/workspaces/dk3/pytorch/xla/torch_xla/distributed/spmd/xla_sharding.py" source_line=570}
  %constant.37 = f32[] constant(0), metadata={op_type="aten__sum" op_name="aten__sum" source_file="/workspaces/dk3/pytorch/xla/test/test_pallas_spmd.py" source_line=84}
  %reduce.43 = f32[] reduce(f32[4,2,128,8]{3,2,1,0} %custom-call.12, f32[] %constant.37), dimensions={0,1,2,3}, to_apply=%AddComputation.39, metadata={op_type="aten__sum" op_name="aten__sum" source_file="/workspaces/dk3/pytorch/xla/test/test_pallas_spmd.py" source_line=84}
  ROOT %tuple.44 = (f32[4,2,128,8]{3,2,1,0}, f32[4,2,128]{2,1,0}, f32[4,2,128]{2,1,0}, f32[]) tuple(f32[4,2,128,8]{3,2,1,0} %custom-call.12, f32[4,2,128]{2,1,0} %custom-call.24, f32[4,2,128]{2,1,0} %custom-call.36, f32[] %reduce.43)
}

with functionizaiton

HloModule IrToHlo.29, entry_computation_layout={(f32[4,2,128,8]{2,3,0,1}, f32[4,2,128,8]{2,3,0,1}, f32[4,2,128,8]{2,3,0,1})->(f32[4,2,128,8]{3,2,1,0}, f32[4,2,128]{2,1,0}, f32[4,2,128]{2,1,0}, f32[])}

%AddComputation.23 (x.24: f32[], y.25: f32[]) -> f32[] {
  %x.24 = f32[] parameter(0)
  %y.25 = f32[] parameter(1)
  ROOT %add.26 = f32[] add(f32[] %x.24, f32[] %y.25)
}

ENTRY %IrToHlo.29 (p0.1: f32[4,2,128,8], p1.3: f32[4,2,128,8], p2.5: f32[4,2,128,8]) -> (f32[4,2,128,8], f32[4,2,128], f32[4,2,128], f32[]) {
  %constant.22 = s32[] constant(8192), metadata={op_type="aten__sum" op_name="aten__sum" source_file="/workspaces/dk3/pytorch/xla/test/test_pallas_spmd.py" source_line=84}
  %p2.5 = f32[4,2,128,8]{2,3,0,1} parameter(2), sharding={devices=[4,1,1,1]0,1,2,3}, metadata={op_type="xla__device_data" op_name="xla__device_data" source_file="/workspaces/dk3/pytorch/xla/torch_xla/distributed/spmd/xla_sharding.py" source_line=625}
  %custom-call.6 = f32[1,2,128,8]{3,2,1,0} custom-call(f32[4,2,128,8]{2,3,0,1} %p2.5), custom_call_target="SPMDFullToShardShape", sharding={manual}, metadata={op_type="xla__custom_sharding" op_name="xla__custom_sharding" source_file="/workspaces/dk3/pytorch/xla/torch_xla/distributed/spmd/xla_sharding.py" source_line=555}
  %p1.3 = f32[4,2,128,8]{2,3,0,1} parameter(1), sharding={devices=[4,1,1,1]0,1,2,3}, metadata={op_type="xla__device_data" op_name="xla__device_data" source_file="/workspaces/dk3/pytorch/xla/torch_xla/distributed/spmd/xla_sharding.py" source_line=625}
  %custom-call.4 = f32[1,2,128,8]{3,2,1,0} custom-call(f32[4,2,128,8]{2,3,0,1} %p1.3), custom_call_target="SPMDFullToShardShape", sharding={manual}, metadata={op_type="xla__custom_sharding" op_name="xla__custom_sharding" source_file="/workspaces/dk3/pytorch/xla/torch_xla/distributed/spmd/xla_sharding.py" source_line=555}
  %p0.1 = f32[4,2,128,8]{2,3,0,1} parameter(0), sharding={devices=[4,1,1,1]0,1,2,3}, metadata={op_type="xla__device_data" op_name="xla__device_data" source_file="/workspaces/dk3/pytorch/xla/torch_xla/distributed/spmd/xla_sharding.py" source_line=625}
  %custom-call.2 = f32[1,2,128,8]{3,2,1,0} custom-call(f32[4,2,128,8]{2,3,0,1} %p0.1), custom_call_target="SPMDFullToShardShape", sharding={manual}, metadata={op_type="xla__custom_sharding" op_name="xla__custom_sharding" source_file="/workspaces/dk3/pytorch/xla/torch_xla/distributed/spmd/xla_sharding.py" source_line=555}
  %custom-call.7 = (f32[1,2,128,8]{3,2,1,0}, f32[1,2,128,128]{3,2,1,0}, f32[1,2,128,128]{3,2,1,0}) custom-call(f32[1,2,128,8]{3,2,1,0} %custom-call.6, f32[1,2,128,8]{3,2,1,0} %custom-call.4, f32[1,2,128,8]{3,2,1,0} %custom-call.2), custom_call_target="tpu_custom_call", operand_layout_constraints={f32[1,2,128,8]{3,2,1,0}, f32[1,2,128,8]{3,2,1,0}, f32[1,2,128,8]{3,2,1,0}}, metadata={op_type="xla__tpu_custom_call" op_name="xla__tpu_custom_call" source_file="/workspaces/dk3/pytorch/xla/torch_xla/experimental/custom_kernel.py" source_line=309}, backend_config={"custom_call_config": {"body": "TUzvUgFNTElSMjAuMC4wZ2l0AAEvCwEDBQcJAQMLAxkNDxETFRcZGx0fISMDYgPiAh8B+QcTCxMLDwsPDw8LCwsLCw8TE5MPCwsLGwsrC8UPCwsLCwuTCwsTCwsLCwsTEwsLMxMTExMXDxMPEw8TGwsPQwsbCwuTCwsLCyMbCxsLGwsbCxsLGwsbGxsbGxsPDw8PFw8XDw8LFw8PCxcPDwsXDw8LFwsPDwsXDw8LExcFDY1heZEqAioCAWcLFxMTFxcfExMXLxcfCxMXHw8LExcLEwsXHwsTFx8PCxMTFxMXC1MLExMXExMXFx8TFy8HBV1JCQNZAR8PBx8rBxsPFyMvGyc3LwsCghMfAwMduQUlAwMdtwUnFcHFBSkdMW0dMXEdMXUFKwUtBS8NGwUxHSu7HSsSAh0rtgIjDQlBAQAAAAAAAAABAAAAAAAAAIAAAAAAAAAACAAAAAAAAAARHQAFMwU1BTcDA14C3gIFOQMFngKiAqYCqgIFO2FmZmluZV9tYXA8KGQwLCBkMSwgZDIsIGQzKSAtPiAoZDAsIGQxLCBkMiwgZDMpPgARDQEFPQU/BUEFQwVFIw0JQQEAAAAAAAAAAQAAAAAAAACAAAAAAAAAAIAAAAAAAAAABUcFSR1NJgIFSwVNBU8FUQVTHT4CWRVCAgsFVQVXIw0DEQEAAAAAAAAAHVICYxVWAgsdcgJnFXYCCx2GAooCHSltFZYCCx0pcRWuAgsdTXUVygILAwV5ew01BVkRDQ0DD3+BFYOFh4k5izkNjY+RBVsBCfn5+f8NGQVdIw0JQQEAAAAAAAAAAgAAAAAAAAABAAAAAAAAAAEAAAAAAAAABV8FYQVjBWUBDZOXm5+jpwMFF5UZJQk7AwUXmRklCT0DBRedGSUJPwMFF6EZJQlBAwUXpRlFCUMDBRepGUUJRwMFFRsNOwMFFRsNPQMFFRsNPwMFFRsNQQMFFRsNQwMFFRsNRxEBAREDARW9Cx0JvxcFCggBHTXDFwVOBQEVx80dycsFZxcF/gsBFc/VHdHTBWkXLTYCARXX3R3Z2wVrFy12BAEV3+cd4eMFbRfl/ggBBW8V6e8d6+0FcRcttgcBFfH3HfP1BXMXSacBHQoCDgIjdHB1LmRpbWVuc2lvbl9zZW1hbnRpY3M8cGFyYWxsZWw+ACN0cHUubWVtb3J5X3NwYWNlPHZtZW0+ACN0cHUuY29udHJhY3RfcHJlY2lzaW9uPGZwMzI+ACN0cHUuZGltZW5zaW9uX3NlbWFudGljczxhcmJpdHJhcnk+ACN0cHUuZG90X2RpbWVuc2lvbl9udW1iZXJzPFsxXSwgWzFdLCBbMF0sIFswXSwgWzAsIDAsIDEsIDBdLCBbXSwgW10+ACN0cHUuZG90X2RpbWVuc2lvbl9udW1iZXJzPFsxXSwgWzBdLCBbMF0sIFsxXSwgWzAsIDAsIDEsIDFdLCBbXSwgW10+AAV1F0mCBQEVFgILHQkaAhcFDggBAwMdIgIlBQkAAAAAFSoCCx0JLgIXBRIIAQMJTwICUf1TJ1UnAwMdOgIlDwkAAID/BXcdCUYCFwWmCAEDBVvWAl1fHSlZBXkdCVoCFwWqCAEFex1mAmMFfQMDHW4CJQ8JAAAAAAV/HQl6AhcFrggBAwVb2gJdXx0pZwWBFY4CCx0JkgIXBbIIAR0JmgIXBb4IAQWDIwEJIQEAAAABAAAABAAAAAAAAAAFhSMBAQEdCbICFwXGCAEVugILHQm+AhcFzggBAwMdxgIlCwkAAAAAHQnOAhcF0ggBAwlPBgJR/VMnVScjdmVjdG9yLmtpbmQ8bWF4aW11bWY+ACN2ZWN0b3Iua2luZDxhZGQ+ACNhcml0aC5mYXN0bWF0aDxub25lPgABAgIDJwUCBAIECRf7CQUFAgQhCTcLJwUCBCEJAQIEJwMCBAknCQUFAgQhCRf7CQUFAgQCBAk3JwUCBAUJJwkFBQIEAgQJBRUBAQEBBwcHBxMTAQUJAQEBAQkBAQEBAQkEUgsFAREBdwcDAR0HEQF9BwOJ+xUBAQEBAQEBAQcBBwEHAQcBEwETAQMDHwMDAwMDHwMDAwMDHwMDAwMDHwMDAwsGHwMRCwkVFxkbBQYfAwsDHQMDIQMDAwMDIQMDAwMDIQMDAwMDIQMDAwsGIQMRCwshIyUnBQYhAwsDKQMDSx4CAwURB0syAgMFBx8rLQMDVzYCAw8TB1dKAgMPBS8xBQZOAgMVAzMNBmEDBQM1FQdhLwMFBS83FwdiAi8DBQM5AwNlagIDDxMHZX4CAw8FOz0FBoICAxUDPw0GaQMFA0EZB2kvAwUFO0MFBmsDFQM1DQZrAwUDRwMDDwMDAwMDDwMDAwMDDwMDAwMDDwMDAwsGDwMXCxNLTU9RBQYPAwUDUwUGDwMXA0kPBQ8zDVcTS01PUQUGbwMVA0ENBm8DBQNZAwMRAwMDAwMRAwMDAwMRAwMDAwMRAwMDCwYRAxcLEV1fYWMFBhEDBQNlBQYRAxcDWw8FETMNaRFdX2FjAwMjAwMDAwMjAwMDAwMjAwMDAwMjAwMDCwYjAxELDWttb3EFBiMDCwNzAwNzwgIDCxEHc9ICAwsHRXV3AwMTAwMDAwMTAwMDAwMTAwMDAwMTAwMDCwYTAxELD3t9f4EFBhMDCwODBQYTAxEDeQ8FEzMNhw97fX+BCQABBxEBqwcDDQ8JAQEBAQEBAQEDAwEHAwEDAwEHAwEJBAEJAQMFCQcRAa0HAw0PCQEBAQEBAQEBAwMBBwMBAwMBBwMBCQQBCQEDBwkHEQGvBwMNDwkBAQEBAQEBAQMDAQcDAQMDAQcDAQkEAQkBAwcJBxEBsQcDDQ8JAQEBAQEBAQEDAwEHAwEDAwEHAwEJBAEJAQMFCQcRAbMHAw0PCQEBAQEBAQEBAwMBBwMBAwMBBwMBCQQBCQEDBQkHEQG1BwMNDwkBAQEBAQEBAQMDAQcDAQMDAQcDAQkEAQkBAwUJBgMBBQEAZhKHESkLGQsTCxkTYyFnDREbLR0LIyEjKS0fCx0dFSUbaxkZGRkZGTENiQslDR0lHRNjtxcTFy8XIyMZGRUlHw8NDwkdEWJ1aWx0aW4Ac3RhYmxlX21vc2FpYwB0cHUAdmVjdG9yAGFyaXRoAG1vZHVsZQBhcml0aC5jb25zdGFudAB2ZWN0b3Iuc2hhcGVfY2FzdABmdW5jLmZ1bmMAZnVuYy5yZXR1cm4AdmVjdG9yLmxvYWQAdmVjdG9yLmJyb2FkY2FzdAB0cHUudmVjdG9yX3N0b3JlAHRwdS5tYXRtdWwAdmVjdG9yLm11bHRpX3JlZHVjdGlvbgBhcml0aC5zdWJmAG1hdGguZXhwAGFyaXRoLmRpdmYAL3Vzci9sb2NhbC9saWIvcHl0aG9uMy4xMC9zaXRlLXBhY2thZ2VzL2pheC9leHBlcmltZW50YWwvcGFsbGFzL29wcy90cHUvZmxhc2hfYXR0ZW50aW9uLnB5AF9mbGFzaF9hdHRlbnRpb25fa2VybmVsX3NpbmdsZV9iYXRjaF9zaW5nbGVfc3RlcABzeW1fbmFtZQBmdW5jdGlvbl90eXBlAHRyYW5zZm9ybV9pbmRpY2VzAHdpbmRvd19ib3VuZHMAdmFsdWUAL2Jyb2FkY2FzdF9pbl9kaW0AL2dldAAvd29ya3NwYWNlcy9kazMvcHl0b3JjaC94bGEvdG9yY2hfeGxhL2V4cGVyaW1lbnRhbC9jdXN0b21fa2VybmVsLnB5AC9zd2FwAF9mbGFzaF9hdHRlbnRpb25fa2VybmVsAHRyYW5zZm9ybV8wAHRyYW5zZm9ybV8xAHRyYW5zZm9ybV8yAHRyYW5zZm9ybV8zAHRyYW5zZm9ybV80AHRyYW5zZm9ybV81AC93b3Jrc3BhY2VzL2RrMy9weXRvcmNoL3hsYS90ZXN0L3Rlc3RfcGFsbGFzX3NwbWQucHkAL2RvdF9nZW5lcmFsAGRpbWVuc2lvbl9udW1iZXJzAHByZWNpc2lvbgB0cmFuc3Bvc2VfbGhzAHRyYW5zcG9zZV9yaHMAa2luZAByZWR1Y3Rpb25fZGltcwBzdGFibGVfbW9zYWljLnZlcnNpb24AZGltZW5zaW9uX3NlbWFudGljcwBpdGVyYXRpb25fYm91bmRzAHNjYWxhcl9wcmVmZXRjaABzY3JhdGNoX29wZXJhbmRzAG1haW4Ad2luZG93X3BhcmFtcwBfZmxhc2hfYXR0ZW50aW9uX2ltcGwAdHJhY2VfcGFsbGFzAGZvcndhcmQAYXBwbHkAL3dvcmtzcGFjZXMvZGszL3B5dG9yY2gvdG9yY2gvYXV0b2dyYWQvZnVuY3Rpb24ucHkAZmxhc2hfYXR0ZW50aW9uAHRlc3RfZmxhc2hfYXR0ZW50aW9uX2JhY2t3YXJkX3NwbWRfZGF0YV9wYXJhbGxlbAA8bW9kdWxlPgAvcmVkdWNlX21heAAvc3ViAGZhc3RtYXRoAC9leHAAL3JlZHVjZV9zdW0AL2RpdgBvcGVyYW5kU2VnbWVudFNpemVzAHN0cmlkZXMA", "cost_estimate": {"flops": 1114112, "transcendentals": 32768, "bytes_accessed": 294912}, "serialization_format": 1, "needs_layout_passes": true}}
  %get-tuple-element.8 = f32[1,2,128,8]{3,2,1,0} get-tuple-element((f32[1,2,128,8]{3,2,1,0}, f32[1,2,128,128]{3,2,1,0}, f32[1,2,128,128]{3,2,1,0}) %custom-call.7), index=0, metadata={op_type="xla__tpu_custom_call" op_name="xla__tpu_custom_call" source_file="/workspaces/dk3/pytorch/xla/torch_xla/experimental/custom_kernel.py" source_line=309}
  %custom-call.11 = f32[1,2,128,8]{3,2,1,0} custom-call(f32[1,2,128,8]{3,2,1,0} %get-tuple-element.8), custom_call_target="Sharding", sharding={manual}, metadata={op_type="xla__custom_sharding" op_name="xla__custom_sharding" source_file="/workspaces/dk3/pytorch/xla/torch_xla/distributed/spmd/xla_sharding.py" source_line=537}
  %custom-call.12 = f32[4,2,128,8]{3,2,1,0} custom-call(f32[1,2,128,8]{3,2,1,0} %custom-call.11), custom_call_target="SPMDShardToFullShape", sharding={devices=[4,1,1,1]0,1,2,3}, metadata={op_type="xla__custom_sharding" op_name="xla__custom_sharding" source_file="/workspaces/dk3/pytorch/xla/torch_xla/distributed/spmd/xla_sharding.py" source_line=570}
  %get-tuple-element.9 = f32[1,2,128,128]{3,2,1,0} get-tuple-element((f32[1,2,128,8]{3,2,1,0}, f32[1,2,128,128]{3,2,1,0}, f32[1,2,128,128]{3,2,1,0}) %custom-call.7), index=1, metadata={op_type="xla__tpu_custom_call" op_name="xla__tpu_custom_call" source_file="/workspaces/dk3/pytorch/xla/torch_xla/experimental/custom_kernel.py" source_line=309}
  %slice.13 = f32[1,2,128,1]{3,2,1,0} slice(f32[1,2,128,128]{3,2,1,0} %get-tuple-element.9), slice={[0:1], [0:2], [0:128], [0:1]}, metadata={op_type="xla__generic_slice" op_name="xla__generic_slice" source_file="/workspaces/dk3/pytorch/xla/torch_xla/experimental/custom_kernel.py" source_line=319}
  %reshape.14 = f32[1,2,128]{2,1,0} reshape(f32[1,2,128,1]{3,2,1,0} %slice.13), metadata={op_type="aten__view" op_name="aten__view" source_file="/workspaces/dk3/pytorch/xla/torch_xla/experimental/custom_kernel.py" source_line=319}
  %custom-call.15 = f32[1,2,128]{2,1,0} custom-call(f32[1,2,128]{2,1,0} %reshape.14), custom_call_target="Sharding", sharding={manual}, metadata={op_type="xla__custom_sharding" op_name="xla__custom_sharding" source_file="/workspaces/dk3/pytorch/xla/torch_xla/distributed/spmd/xla_sharding.py" source_line=537}
  %custom-call.16 = f32[4,2,128]{2,1,0} custom-call(f32[1,2,128]{2,1,0} %custom-call.15), custom_call_target="SPMDShardToFullShape", sharding={devices=[4,1,1]0,1,2,3}, metadata={op_type="xla__custom_sharding" op_name="xla__custom_sharding" source_file="/workspaces/dk3/pytorch/xla/torch_xla/distributed/spmd/xla_sharding.py" source_line=570}
  %get-tuple-element.10 = f32[1,2,128,128]{3,2,1,0} get-tuple-element((f32[1,2,128,8]{3,2,1,0}, f32[1,2,128,128]{3,2,1,0}, f32[1,2,128,128]{3,2,1,0}) %custom-call.7), index=2, metadata={op_type="xla__tpu_custom_call" op_name="xla__tpu_custom_call" source_file="/workspaces/dk3/pytorch/xla/torch_xla/experimental/custom_kernel.py" source_line=309}
  %slice.17 = f32[1,2,128,1]{3,2,1,0} slice(f32[1,2,128,128]{3,2,1,0} %get-tuple-element.10), slice={[0:1], [0:2], [0:128], [0:1]}, metadata={op_type="xla__generic_slice" op_name="xla__generic_slice" source_file="/workspaces/dk3/pytorch/xla/torch_xla/experimental/custom_kernel.py" source_line=319}
  %reshape.18 = f32[1,2,128]{2,1,0} reshape(f32[1,2,128,1]{3,2,1,0} %slice.17), metadata={op_type="aten__view" op_name="aten__view" source_file="/workspaces/dk3/pytorch/xla/torch_xla/experimental/custom_kernel.py" source_line=319}
  %custom-call.19 = f32[1,2,128]{2,1,0} custom-call(f32[1,2,128]{2,1,0} %reshape.18), custom_call_target="Sharding", sharding={manual}, metadata={op_type="xla__custom_sharding" op_name="xla__custom_sharding" source_file="/workspaces/dk3/pytorch/xla/torch_xla/distributed/spmd/xla_sharding.py" source_line=537}
  %custom-call.20 = f32[4,2,128]{2,1,0} custom-call(f32[1,2,128]{2,1,0} %custom-call.19), custom_call_target="SPMDShardToFullShape", sharding={devices=[4,1,1]0,1,2,3}, metadata={op_type="xla__custom_sharding" op_name="xla__custom_sharding" source_file="/workspaces/dk3/pytorch/xla/torch_xla/distributed/spmd/xla_sharding.py" source_line=570}
  %constant.21 = f32[] constant(0), metadata={op_type="aten__sum" op_name="aten__sum" source_file="/workspaces/dk3/pytorch/xla/test/test_pallas_spmd.py" source_line=84}
  %reduce.27 = f32[] reduce(f32[4,2,128,8]{3,2,1,0} %custom-call.12, f32[] %constant.21), dimensions={0,1,2,3}, to_apply=%AddComputation.23, metadata={op_type="aten__sum" op_name="aten__sum" source_file="/workspaces/dk3/pytorch/xla/test/test_pallas_spmd.py" source_line=84}
  ROOT %tuple.28 = (f32[4,2,128,8]{3,2,1,0}, f32[4,2,128]{2,1,0}, f32[4,2,128]{2,1,0}, f32[]) tuple(f32[4,2,128,8]{3,2,1,0} %custom-call.12, f32[4,2,128]{2,1,0} %custom-call.16, f32[4,2,128]{2,1,0} %custom-call.20, f32[] %reduce.27)
}

JackCaoG · 2024-12-02T20:07:34Z

Ok I narrowed it down to this line

xla/torch_xla/experimental/custom_kernel.py

Line 319 in 591c397

l, m = (v[..., 0] for v in aux[-2:])

if I disable the functionizaiton and print l it will crash right away. This also make sense as it is trying to do some view op and likely functionizaiton has some issue..

JackCaoG · 2024-12-02T22:26:46Z

here is my diff.

diff --git a/torch_xla/experimental/custom_kernel.py b/torch_xla/experimental/custom_kernel.py
index 5e30ffba2..09cbc755e 100644
--- a/torch_xla/experimental/custom_kernel.py
+++ b/torch_xla/experimental/custom_kernel.py
@@ -317,6 +317,8 @@ class FlashAttention(torch.autograd.Function):
         return o
       o, *aux = o
       l, m = (v[..., 0] for v in aux[-2:])
+      print(f"before view {torch_xla._XLAC._get_xla_tensor_debug_info(l)}")
+
 
     # SPMD integration
     if partition_spec is not None:
@@ -324,6 +326,11 @@ class FlashAttention(torch.autograd.Function):
           o, partition_spec, ctx.full_shape, mesh=mesh).global_tensor
       l = xs.disable_manual_sharding(
           l, partition_spec[0:3], ctx.full_shape[0:3], mesh=mesh).global_tensor
+      print(f"after view {torch_xla._XLAC._get_xla_tensor_debug_info(l)}")    
+      breakpoint()
+      print(torch_xla._XLAC._get_xla_tensors_hlo([l]))
+      # this will crash without functionization
+      print(l)
       m = xs.disable_manual_sharding(
           m, partition_spec[0:3], ctx.full_shape[0:3], mesh=mesh).global_tensor

when I compare the HLO printed, I found out that in the failing case there is a weird HLO

  %custom-call.13 = f32[1,2,128]{2,1,0} custom-call(f32[1,2,128]{2,1,0} %reshape.12), custom_call_target="Sharding", metadata={op_type="xla__custom_sharding" op_name="xla__custom_sharding" source_file="/workspaces/dk3/pytorch/xla/torch_xla/distributed/spmd/xla_sharding.py" source_line=537}

....



  %reshape.21 = f32[1,2,128]{2,1,0} reshape(f32[1,2,128,1]{3,2,1,0} %slice.20), sharding={manual}, metadata={op_type="aten__view" op_name="aten__view" source_file="/workspaces/dk3/pytorch/xla/torch_xla/distributed/spmd/xla_sharding.py" source_line=537}

where the passing case the custom_call_target="Sharding" and sharding={manual} is in the same HLO


  %custom-call.13 = f32[1,2,128]{2,1,0} custom-call(f32[1,2,128]{2,1,0} %reshape.12), custom_call_target="Sharding", sharding={manual}, metadata={op_type="xla__custom_sharding" op_name="xla__custom_sharding"

I think this is where the issue is from

JackCaoG · 2024-12-02T22:43:06Z

ok I think the issue is in disable_manual_sharding , more specifically

xla/torch_xla/distributed/spmd/xla_sharding.py

Line 569 in 591c397

t = _mark_manual_sharding(unwrap_sharded_tensor(t))

the _mark_manual_sharding should annotate that current tensor is sharded with manual so we can bring it to the right shape. However with XLA_DISABLE_FUNCTIONIZATION=1, the pytorch/xla's view logic kicks in and apply some of the view replay logic and move the sharding away from the custom_call = Sharding HLO.

IR

  %8 = f32[1,2,128]{2,1,0} aten::view(%7), location=forward@custom_kernel.py:320, xla_shape=f32[1,2,128]{2,1,0}
  %9 = f32[1,2,128]{2,1,0} xla::custom_sharding(%8), location=_mark_manual_sharding@xla_sharding.py:537, xla_shape=f32[1,2,128]{2,1,0}

next step is to look into where view logic kick in. @miladm I will let you find someone to follow up.

dudulightricks · 2024-12-10T11:38:58Z

@JackCaoG Hey! Can you merge this PR? is it fine?
Also, is there any progress with the issue with XLA_DISABLE_FUNCTIONIZATION?

JackCaoG · 2024-12-10T17:16:24Z

I am no longer with the team, you should check with @miladm

JackCaoG added the tpuci label Nov 28, 2024

dudulightricks force-pushed the bug-fix/support-cross-flash-attn branch 3 times, most recently from 840e3fe to 1bc1fb4 Compare December 1, 2024 10:44

flash_attention: support also cross attention.

9e1e24b

In case that q and kv have different shapes (cross attention) flash attention with spmd fails since it does not support it.

dudulightricks force-pushed the bug-fix/support-cross-flash-attn branch from 1bc1fb4 to 9e1e24b Compare December 1, 2024 10:47

JackCaoG mentioned this pull request Dec 2, 2024

TPU memory use increased significantly in torch/xla - 2.6.0.dev20241107 #8423

Open

JackCaoG approved these changes Dec 10, 2024

View reviewed changes

JackCaoG merged commit 5790092 into pytorch:master Dec 10, 2024
12 checks passed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

flash_attention: support also cross attention. #8427

flash_attention: support also cross attention. #8427

dudulightricks commented Nov 28, 2024

dudulightricks commented Nov 28, 2024

dudulightricks commented Dec 1, 2024 •

edited

Loading

dudulightricks commented Dec 2, 2024 •

edited

Loading

JackCaoG commented Dec 2, 2024

JackCaoG commented Dec 2, 2024

JackCaoG commented Dec 2, 2024

JackCaoG commented Dec 2, 2024

JackCaoG commented Dec 2, 2024

JackCaoG commented Dec 2, 2024

JackCaoG commented Dec 2, 2024

dudulightricks commented Dec 10, 2024 •

edited

Loading

JackCaoG commented Dec 10, 2024

flash_attention: support also cross attention. #8427

flash_attention: support also cross attention. #8427

Conversation

dudulightricks commented Nov 28, 2024

dudulightricks commented Nov 28, 2024

dudulightricks commented Dec 1, 2024 • edited Loading

dudulightricks commented Dec 2, 2024 • edited Loading

JackCaoG commented Dec 2, 2024

JackCaoG commented Dec 2, 2024

JackCaoG commented Dec 2, 2024

JackCaoG commented Dec 2, 2024

JackCaoG commented Dec 2, 2024

JackCaoG commented Dec 2, 2024

JackCaoG commented Dec 2, 2024

dudulightricks commented Dec 10, 2024 • edited Loading

JackCaoG commented Dec 10, 2024

dudulightricks commented Dec 1, 2024 •

edited

Loading

dudulightricks commented Dec 2, 2024 •

edited

Loading

dudulightricks commented Dec 10, 2024 •

edited

Loading