diff --git a/apex/transformer/functional/fused_rope.py b/apex/transformer/functional/fused_rope.py index e74906151..3be97285e 100644 --- a/apex/transformer/functional/fused_rope.py +++ b/apex/transformer/functional/fused_rope.py @@ -111,6 +111,12 @@ def forward( freqs: torch.Tensor, transpose_output_memory: bool = False, ) -> torch.Tensor: + # The aiter kernel requires the head-dim (last) stride to be 1. The + # input may be non-contiguous (e.g. transposed layout), so enforce + # contiguity here. + if t.stride(-1) != 1: t = t.contiguous() + # t = t.contiguous() + s = t.shape[0] b = t.shape[1] h = t.shape[2] @@ -137,6 +143,11 @@ def backward( ) -> Tuple[Union[torch.Tensor, None], ...]: (freqs,) = ctx.saved_tensors + # The aiter kernel requires the head-dim (last) stride to be 1. The + # incoming gradient may be non-contiguous (e.g. broadcast/expanded from + # a reduction such as `.sum()`), so enforce contiguity here. + if grad_output.stride(-1) != 1: grad_output = grad_output.contiguous() + s = grad_output.shape[0] b = grad_output.shape[1] h = grad_output.shape[2] @@ -238,6 +249,11 @@ def forward( sin_: torch.Tensor, transpose_output_memory: bool = False, ) -> torch.Tensor: + # The aiter kernel requires the head-dim (last) stride to be 1. The + # input may be non-contiguous (e.g. transposed layout), so enforce + # contiguity here. + if t.stride(-1) != 1: t = t.contiguous() + s = t.shape[0] b = t.shape[1] h = t.shape[2] @@ -263,6 +279,11 @@ def backward( ) -> Tuple[Union[torch.Tensor, None], ...]: cos_, sin_ = ctx.saved_tensors + # The aiter kernel requires the head-dim (last) stride to be 1. The + # incoming gradient may be non-contiguous (e.g. broadcast/expanded from + # a reduction such as `.sum()`), so enforce contiguity here. + if grad_output.stride(-1) != 1: grad_output = grad_output.contiguous() + s = grad_output.shape[0] b = grad_output.shape[1] h = grad_output.shape[2] @@ -360,6 +381,11 @@ def forward( cu_seqlens: torch.Tensor, freqs: torch.Tensor, ) -> torch.Tensor: + # The aiter kernel requires the head-dim (last) stride to be 1. The + # input may be non-contiguous (e.g. transposed layout), so enforce + # contiguity here. + if t.stride(-1) != 1: t = t.contiguous() + t1 = t.shape[0] h = t.shape[1] d = t.shape[2] @@ -379,6 +405,11 @@ def backward( ) -> Tuple[Union[torch.Tensor, None], ...]: cu_seqlens, freqs = ctx.saved_tensors + # The aiter kernel requires the head-dim (last) stride to be 1. The + # incoming gradient may be non-contiguous (e.g. broadcast/expanded from + # a reduction such as `.sum()`), so enforce contiguity here. + if grad_output.stride(-1) != 1: grad_output = grad_output.contiguous() + t = grad_output.shape[0] h = grad_output.shape[1] d = grad_output.shape[2] @@ -488,7 +519,12 @@ def forward( cos_w: torch.Tensor, sin_w: torch.Tensor, ) -> torch.Tensor: - + + # The aiter kernel requires the head-dim (last) stride to be 1. The + # input may be non-contiguous (e.g. transposed layout), so enforce + # contiguity here. + if t.stride(-1) != 1: t = t.contiguous() + s = t.shape[0] h = t.shape[2] d = t.shape[3] @@ -509,7 +545,12 @@ def backward( ) -> Tuple[Union[torch.Tensor, None], ...]: cos_h, sin_h, cos_w, sin_w = ctx.saved_tensors - + + # The aiter kernel requires the head-dim (last) stride to be 1. The + # incoming gradient may be non-contiguous (e.g. broadcast/expanded from + # a reduction such as `.sum()`), so enforce contiguity here. + if grad_output.stride(-1) != 1: grad_output = grad_output.contiguous() + s = grad_output.shape[0] h = grad_output.shape[2] d = grad_output.shape[3]