DID loop split for SDPA

[Priya2698]

In this PR, I explicitly parallelize the outputs attn, log_sumexp of sdpfa_fwd. Sharding propagation for loop split does not work correctly in this case at the moment.

[Priya2698]

!test

[wujingyue]

Is this needed? I believe it'll get overwritten by set_allocation_as_loop so we can simply start with the default stride_order.

[github-actions]

PR Reviewer Guide 🔍

(Review updated until commit 026229d)

Here are some key observations to aid the review process:

⏱️ Estimated effort to review: 4 🔵🔵🔵🔵⚪

🧪 PR contains tests

⚡ Recommended focus areas for review Potential Logic Change The introduction of QkvFormat parameter in the Model class and its usage in definition method might alter the expected behavior of the SDPA loop split functionality. Reviewers should verify that the added logic for QkvFormat.BHSE and QkvFormat.BSHE correctly handles the stride ordering and tensor reformatting. @pytest.mark.skipif( utils.is_pre_ampere(), reason="Flash Attention is only supported on Ampere and newer devices.", ) @pytest.mark.parametrize("qkv_format", [QkvFormat.BHSE, QkvFormat.BSHE]) @pytest.mark.mpi def test_sdpa_loop_split(multidevice_test, qkv_format: QkvFormat): d, b, s, h, e = multidevice_test.size, 2, 1024, 12, 768 if h % d != 0: pytest.skip(f"We only support even split, so {h} has to be divisible by {d}.") mesh = nvfuser.DeviceMesh(range(d)) class Model(FusionDefinition): def __init__(self, qkv_format: QkvFormat): super().__init__() self._qkv_format = qkv_format def definition(self) -> None: match self._qkv_format: case QkvFormat.BHSE: stride_order = [3, 2, 1, 0] case QkvFormat.BSHE: stride_order = [3, 1, 2, 0] self.q, self.k, self.v, self.out_grad = [ self.define_tensor( shape=[b, h, s, e // h], dtype=DataType.BFloat16, stride_order=stride_order, ) for _ in range(4) ] # TODO(#3123): support sharded dropout and change this to a # positive probability. dropout_p = self.define_scalar(0.0, dtype=DataType.Double) is_causal = self.define_scalar(True, dtype=DataType.Bool) self.attn, self.log_sumexp, seed, offset = self.ops.sdpfa_fwd( self.q, self.k, self.v, dropout_p, is_causal, scale=None ) self.q_grad, self.k_grad, self.v_grad = self.ops.sdpfa_bwd( self.out_grad, self.q, self.k, self.v, self.attn, self.log_sumexp, dropout_p, is_causal, seed, offset, scale=None, ) self.add_output(self.attn) for grad in [self.q_grad, self.k_grad, self.v_grad]: self.add_output(grad) def multidevice_schedule(self) -> None: for t in [ self.q, self.k, self.v, self.attn, self.log_sumexp, self.out_grad, self.q_grad, self.k_grad, self.v_grad, ]: self.sched._set_device_mesh(t, mesh) self.sched.split(t, 1, d, False) self.sched.parallelize(t, 1, nvfuser.ParallelType.mesh_x) if self._qkv_format == QkvFormat.BSHE: # The loop domain is: {i{B}, i{DIDx}, i{H//D}, i{S}, i{E//H}} # Reorder i{S} in the allocation domain for BHSE: {i{DIDx}, i{B}, i{S}, i{H//D}, i{E//H}} self.sched.reorder(t, {2: 3, 3: 2}) self.sched.set_allocation_as_loop(t) Device Compatibility The test is skipped for pre-Ampere devices due to Flash Attention support. Ensure that this restriction is correctly implemented and documented, and consider adding tests for Ampere and newer devices to cover the functionality. @pytest.mark.skipif( utils.is_pre_ampere(), reason="Flash Attention is only supported on Ampere and newer devices.", )