chore: import upstream snapshot with attribution

swift/sequence_parallel/__init__.py

@@ -0,0 +1,5 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+
+from .sequence_parallel import SequenceParallel, sequence_parallel
+from .utils import (ChunkedCrossEntropyLoss, GatherLoss, GatherTensor, SequenceParallelDispatcher,
+                    SequenceParallelSampler)

swift/sequence_parallel/ring_utils.py

@@ -0,0 +1,58 @@
+# Some code borrowed from the awesome work: https://github.com/zhuzilin/ring-flash-attention
+# Copyright (c) ModelScope Contributors. All rights reserved.
+import torch
+import torch.distributed as dist
+from typing import Optional, Tuple
+
+
+class RingComm:
+
+    def __init__(self, process_group: dist.ProcessGroup):
+        self._process_group = process_group
+        self._ops = []
+        self.rank = dist.get_rank(self._process_group)
+        self.world_size = dist.get_world_size(self._process_group)
+        self._reqs = None
+
+        self.send_rank = (self.rank + 1) % self.world_size
+        self.recv_rank = (self.rank - 1) % self.world_size
+
+        if process_group is not None:
+            self.send_rank = dist.get_global_rank(self._process_group, self.send_rank)
+            self.recv_rank = dist.get_global_rank(self._process_group, self.recv_rank)
+
+    def send_recv(self, to_send: torch.Tensor, recv_tensor: Optional[torch.Tensor] = None) -> torch.Tensor:
+        if recv_tensor is None:
+            res = torch.empty_like(to_send)
+        else:
+            res = recv_tensor
+
+        send_op = dist.P2POp(dist.isend, to_send, self.send_rank, group=self._process_group)
+        recv_op = dist.P2POp(dist.irecv, res, self.recv_rank, group=self._process_group)
+        self._ops.append(send_op)
+        self._ops.append(recv_op)
+        return res
+
+    def commit(self):
+        if self._reqs is not None:
+            raise RuntimeError('commit called twice')
+        self._reqs = dist.batch_isend_irecv(self._ops)
+
+    def wait(self):
+        if self._reqs is None:
+            raise RuntimeError('wait called before commit')
+        for req in self._reqs:
+            req.wait()
+        self._reqs = None
+        self._ops = []
+
+    def send_recv_kv(
+        self,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        k_buffer: Optional[torch.Tensor] = None,
+        v_buffer: Optional[torch.Tensor] = None,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        next_k, next_v = self.send_recv(k, k_buffer), self.send_recv(v, v_buffer)
+        self.commit()
+        return next_k, next_v

swift/sequence_parallel/sequence_parallel.py

@@ -0,0 +1,723 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+import inspect
+import math
+import torch
+import torch.distributed as dist
+from functools import lru_cache, partial
+from torch.distributed import init_device_mesh
+from transformers import PreTrainedTokenizer
+from types import SimpleNamespace
+from typing import Optional
+
+from swift.utils import HfConfigFactory, get_cu_seqlens_from_position_ids, get_device, get_dist_setting
+from .ulysses import DistributedAttention
+from .zigzag_ring_attn import zigzag_ring_flash_attn_varlen_func
+
+
+@lru_cache(maxsize=None)
+def get_signature_parameters(fn):
+    try:
+        return inspect.signature(fn).parameters
+    except (TypeError, ValueError):
+        return None
+
+
+def has_signature_parameter(fn, parameter: str) -> bool:
+    parameters = get_signature_parameters(fn)
+    return parameters is not None and parameter in parameters
+
+
+def call_with_supported_kwargs(fn, *args, **kwargs):
+    parameters = get_signature_parameters(fn)
+    if parameters is None:
+        return fn(*args, **kwargs)
+    if not any(param.kind == inspect.Parameter.VAR_KEYWORD for param in parameters.values()):
+        kwargs = {key: value for key, value in kwargs.items() if key in parameters}
+    return fn(*args, **kwargs)
+
+
+def _call_create_causal_mask(fn, config, input_embeds, attention_mask, cache_position_or_past_key_values, *args,
+                             **kwargs):
+    if has_signature_parameter(fn, 'cache_position'):
+        return call_with_supported_kwargs(
+            fn,
+            config,
+            input_embeds,
+            attention_mask,
+            cache_position_or_past_key_values,
+            *args,
+            **kwargs,
+        )
+    if cache_position_or_past_key_values is None and 'past_key_values' in kwargs:
+        return call_with_supported_kwargs(fn, config, input_embeds, attention_mask, *args, **kwargs)
+    return call_with_supported_kwargs(
+        fn,
+        config,
+        input_embeds,
+        attention_mask,
+        cache_position_or_past_key_values,
+        *args,
+        **kwargs,
+    )
+
+
+class SequenceParallel:
+
+    _global_inited: bool = False
+
+    def __init__(self):
+        self.sp_world_size = None
+        self.dp_world_size = None
+        self.rp_world_size = None
+        self.world_size = None
+        self.model_dtype = None
+        self.tokenizer = None
+        self.device_mesh = None
+        self.num_heads = None
+        self.causal_mask_func = None
+        self.extra_kwargs = {}
+
+    @property
+    def real_position_ids(self) -> torch.Tensor:
+        """The real position ids, this is different from the position_ids in mrope"""
+        return self.extra_kwargs.get('text_position_ids')
+
+    def _prepare_flash_attn(self, base_model: torch.nn.Module):
+        try:
+            from transformers import masking_utils
+
+            _origin_flash_attention_mask = masking_utils.flash_attention_mask
+
+            def flash_attention_mask(*args, **kwargs):
+                if self.world_size == 1:
+                    return _origin_flash_attention_mask(*args, **kwargs)
+                attention_mask = kwargs.get('attention_mask')
+                if attention_mask is not None:
+                    if attention_mask.all():
+                        attention_mask = None
+
+                return attention_mask
+
+            masking_utils.flash_attention_mask = flash_attention_mask
+            masking_utils.ALL_MASK_ATTENTION_FUNCTIONS._global_mapping['flash_attention_2'] = flash_attention_mask
+
+            def sdpa_mask(*args, **kwargs):
+                if self.world_size == 1:
+                    return masking_utils.ALL_MASK_ATTENTION_FUNCTIONS._global_mapping['sdpa_origin'](*args, **kwargs)
+                if 'cache_position' in kwargs:
+                    device = kwargs['cache_position'].device
+                else:
+                    # transformers>=5.4.0
+                    device = kwargs['device']
+                cache_position = self.real_position_ids[0]
+                cache_position = self.pad(cache_position, padding_value=-1, position_ids=self.real_position_ids, dim=0)
+                cache_position = torch.arange(0, cache_position.shape[0], device=device)
+                kwargs['kv_length'] = cache_position.shape[0]
+                if 'cache_position' in kwargs:
+                    kwargs['cache_position'] = cache_position
+                else:
+                    kwargs['q_length'] = kwargs['kv_length']
+                return masking_utils.ALL_MASK_ATTENTION_FUNCTIONS._global_mapping['sdpa_origin'](*args, **kwargs)
+
+            masking_utils.ALL_MASK_ATTENTION_FUNCTIONS._global_mapping[
+                'sdpa_origin'] = masking_utils.ALL_MASK_ATTENTION_FUNCTIONS._global_mapping['sdpa']
+            masking_utils.ALL_MASK_ATTENTION_FUNCTIONS._global_mapping['sdpa'] = sdpa_mask
+
+            def create_causal_mask(config,
+                                   input_embeds,
+                                   attention_mask,
+                                   cache_position_or_past_key_values=None,
+                                   *args,
+                                   **kwargs):
+                if self.world_size == 1:
+                    return _call_create_causal_mask(masking_utils.origin_create_causal_mask, config, input_embeds,
+                                                    attention_mask, cache_position_or_past_key_values, *args, **kwargs)
+                input_embeds = torch.ones(
+                    (input_embeds.shape[0], input_embeds.shape[1] * self.sp_world_size, input_embeds.shape[2]),
+                    dtype=input_embeds.dtype,
+                    device=input_embeds.device)
+                if has_signature_parameter(masking_utils.origin_create_causal_mask, 'cache_position'):
+                    cache_position_or_past_key_values = torch.arange(
+                        0,
+                        input_embeds.shape[1],
+                        device=input_embeds.device,
+                    )
+                return _call_create_causal_mask(masking_utils.origin_create_causal_mask, config, input_embeds,
+                                                attention_mask, cache_position_or_past_key_values, *args, **kwargs)
+
+            masking_utils.origin_create_causal_mask = masking_utils.create_causal_mask
+            masking_utils.create_causal_mask = create_causal_mask
+        except ImportError:
+            pass
+
+        if hasattr(base_model, 'language_model'):
+            text_model = base_model.language_model
+        else:
+            text_model = base_model
+
+        from transformers.modeling_flash_attention_utils import is_flash_attn_available
+        if is_flash_attn_available():
+            # TODO this works for multi-modal models like qwen2.5-vl
+            # SDPA is not supported, because we need to copy the code to our project, which will bring
+            # more works for maintaining.
+            from transformers import modeling_flash_attention_utils
+            from transformers.modeling_flash_attention_utils import _flash_attention_forward
+            _distributed_flash_attention = DistributedAttention(_flash_attention_forward, self)
+
+            modeling_flash_attention_utils._flash_attention_forward_origin = _flash_attention_forward
+
+            def flash_attention_forward(query_states: torch.Tensor, key_states: torch.Tensor,
+                                        value_states: torch.Tensor, attention_mask: Optional[torch.Tensor], q_len,
+                                        *args, **kwargs):
+                if self.world_size == 1:
+                    return _flash_attention_forward(query_states, key_states, value_states, attention_mask, q_len,
+                                                    *args, **kwargs)
+                return _distributed_flash_attention(query_states, key_states, value_states, attention_mask,
+                                                    q_len * self.sp_world_size, *args, **kwargs)
+
+            modeling_flash_attention_utils._flash_attention_forward = flash_attention_forward
+
+        from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS
+
+        def local_flash_attn(module: torch.nn.Module, query_states, key_states, value_states, attention_mask, *args,
+                             dist_attn, **kwargs):
+            if self.world_size == 1 or module.__class__ not in [m.__class__ for m in text_model.modules()]:
+                return ALL_ATTENTION_FUNCTIONS['flash_attention_2_origin'](module, query_states, key_states,
+                                                                           value_states, attention_mask, *args,
+                                                                           **kwargs)
+            if dist_attn.local_attn is None:
+
+                def _attention(query, key, value, *args, **kwargs):
+                    query = query.transpose(1, 2)
+                    key = key.transpose(1, 2)
+                    value = value.transpose(1, 2)
+                    if self.rp_world_size is not None and self.rp_world_size > 1:
+                        position_ids = kwargs['position_ids']
+                        cu_seqlens = get_cu_seqlens_from_position_ids(position_ids)
+                        max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max().item()
+                        position_ids = self._split_packed(position_ids, cu_seqlens)
+                        mask = position_ids != -1
+                        query = query.transpose(1, 2)
+                        key = key.transpose(1, 2)
+                        value = value.transpose(1, 2)
+                        # this is important
+                        # the length is correct, but the value is wrong, by mask qkv, we do:
+                        # mask the padded values of q and v to zero
+                        # mask the padded values of k to -1e5
+                        # to keep the attention correct
+                        query, key, value = self._mask_qkv(query, key, value, mask)
+                        output = zigzag_ring_flash_attn_varlen_func(
+                            query,
+                            key,
+                            value,
+                            cu_seqlens=cu_seqlens,
+                            max_seqlen=max_seqlen,
+                            causal=module.is_causal,
+                            dropout_p=kwargs.get('dropout', 0.0),
+                            softmax_scale=kwargs.get('scaling', 0.0),
+                            window_size=kwargs.get('sliding_window') or (-1, -1),
+                            group=self.rp_group)
+                        return output
+                    else:
+                        if 'cu_seq_lens_q' in kwargs:
+                            position_ids = kwargs.get('position_ids')
+                            if self.real_position_ids is not None:
+                                position_ids = self.real_position_ids
+                            position_ids = self.pad(position_ids, padding_value=-1, position_ids=position_ids)
+                            cu_seqlens = get_cu_seqlens_from_position_ids(position_ids)
+                            max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max().item()
+                            assert query.shape[2] == cu_seqlens[-1]
+                            kwargs['cu_seq_lens_q'] = cu_seqlens
+                            kwargs['cu_seq_lens_k'] = cu_seqlens
+                            kwargs['max_length_q'] = max_seqlen
+                            kwargs['max_length_k'] = max_seqlen
+                        return ALL_ATTENTION_FUNCTIONS['flash_attention_2_origin'](module, query, key, value, *args,
+                                                                                   **kwargs)[0]
+
+                dist_attn.local_attn = _attention
+
+            return dist_attn(
+                query_states.transpose(1, 2), key_states.transpose(1, 2), value_states.transpose(1, 2), attention_mask,
+                *args, **kwargs), None
+
+        def local_sdpa_attn(module: torch.nn.Module, query_states, key_states, value_states, attention_mask, *args,
+                            dist_attn, **kwargs):
+            # Bypass SP logic when world_size == 1 (SP disabled) or module not in text_model
+            if self.world_size == 1 or module.__class__ not in [m.__class__ for m in text_model.modules()]:
+                return ALL_ATTENTION_FUNCTIONS['sdpa_origin'](module, query_states, key_states, value_states,
+                                                              attention_mask, *args, **kwargs)
+            if dist_attn.local_attn is None:
+
+                def _attention(query, key, value, *args, **kwargs):
+                    query = query.transpose(1, 2)
+                    key = key.transpose(1, 2)
+                    value = value.transpose(1, 2)
+                    if self.rp_world_size > 1:
+                        raise NotImplementedError('SDPA does not support Ring attention.')
+                    return ALL_ATTENTION_FUNCTIONS['sdpa_origin'](module, query, key, value, *args, **kwargs)[0]
+
+                dist_attn.local_attn = _attention
+            return dist_attn(
+                query_states.transpose(1, 2), key_states.transpose(1, 2), value_states.transpose(1, 2), attention_mask,
+                *args, **kwargs), None
+
+        ALL_ATTENTION_FUNCTIONS['flash_attention_2_origin'] = ALL_ATTENTION_FUNCTIONS['flash_attention_2']
+        ALL_ATTENTION_FUNCTIONS['sdpa_origin'] = ALL_ATTENTION_FUNCTIONS['sdpa']
+        ALL_ATTENTION_FUNCTIONS['flash_attention_2'] = partial(
+            local_flash_attn, dist_attn=DistributedAttention(None, self))
+        ALL_ATTENTION_FUNCTIONS['sdpa'] = partial(local_sdpa_attn, dist_attn=DistributedAttention(None, self))
+
+    def _prepare_forward_hook(self, base_model: torch.nn.Module):
+
+        def pre_forward_split_hook(_self, args, kwargs):
+            if self.world_size == 1:
+                return args, kwargs
+            input_ids = kwargs.get('input_ids', None)
+            inputs_embeds = kwargs.get('inputs_embeds', None)
+            position_ids = kwargs['position_ids']
+            attention_mask = kwargs.get('attention_mask', None)
+            if hasattr(_self, 'language_model'):
+                embed_tokens = getattr(_self.language_model, 'embed_tokens', None)
+            else:
+                embed_tokens = getattr(_self, 'embed_tokens', None)
+            input_ids, inputs_embeds, _, position_ids, attention_mask, _, _ = self.pad_and_split_inputs(
+                input_ids,
+                inputs_embeds,
+                None,
+                position_ids,
+                attention_mask,
+                None,
+                embed_tokens=embed_tokens,
+                real_position_ids=self.real_position_ids)
+            kwargs['input_ids'] = input_ids
+            kwargs['inputs_embeds'] = inputs_embeds
+            kwargs['position_ids'] = position_ids
+            kwargs['attention_mask'] = attention_mask
+            return args, kwargs
+
+        base_model.register_forward_pre_hook(pre_forward_split_hook, with_kwargs=True)
+
+    def _prepare_moe_aux_loss(self, base_model: torch.nn.Module):
+        from .utils import GatherLoss
+
+        def moe_aux_loss_hook(module, args, kwargs, output):
+            router_logits = getattr(output, 'router_logits', None)
+            if router_logits is None:
+                return output
+
+            attention_mask = kwargs['attention_mask']
+            if attention_mask is None:
+                batch_size = 1
+            else:
+                batch_size = attention_mask.shape[0]
+
+            assert router_logits[0].shape[0] % batch_size == 0
+            seq_len = router_logits[0].shape[0] // batch_size
+
+            _gathered_logits = []
+            for i in range(batch_size):
+                _slice = slice(i * seq_len, (i + 1) * seq_len)
+                _bs_logits = [logit[_slice] for logit in router_logits]
+                compute_device = _bs_logits[0].device
+                _bs_logits = torch.stack([layer_gate.to(compute_device) for layer_gate in _bs_logits], dim=0)
+                _bs_logits, _ = GatherLoss.apply(_bs_logits, None, 1, self.real_position_ids)
+                _gathered_logits.append(_bs_logits)
+            router_logits = torch.stack(_gathered_logits, dim=0)
+            if self.real_position_ids is not None:
+                router_logits = router_logits[:, :, :self.real_position_ids.shape[1], :]
+            output['router_logits'] = tuple(
+                [logit.reshape(-1, logit.shape[-1]) for logit in router_logits.split(1, dim=1)])
+            return output
+
+        base_model.register_forward_hook(moe_aux_loss_hook, with_kwargs=True)
+
+    def prepare(self, sp_size: int, model: torch.nn.Module, tokenizer: PreTrainedTokenizer, padding_free: bool):
+        from swift.model import get_llm_model
+        self.num_heads = HfConfigFactory.get_config_attr(model.config, 'num_key_value_heads')
+        if self.num_heads is None:
+            self.num_heads = HfConfigFactory.get_config_attr(model.config, 'num_attention_heads')
+        assert self.num_heads is not None, 'Cannot find num_heads config in config.json'
+        self.padding_free = padding_free
+        self.world_size = sp_size
+
+        llm_model = get_llm_model(model)
+
+        if hasattr(llm_model, 'language_model'):
+            if hasattr(llm_model.language_model, '_update_causal_mask'):
+                self.causal_mask_func = llm_model.language_model._update_causal_mask
+        else:
+            if hasattr(llm_model, '_update_causal_mask'):
+                self.causal_mask_func = llm_model._update_causal_mask
+
+        if not SequenceParallel._global_inited:
+            # these operations are global initializations and patches
+            self._init_device_mesh()
+            self._prepare_flash_attn(llm_model)
+            SequenceParallel._global_inited = True
+
+        self._prepare_forward_hook(llm_model)
+
+        if model.model_info.is_moe_model:
+            self._prepare_moe_aux_loss(llm_model)
+
+        self.model_dtype = next(model.parameters()).dtype
+        self.tokenizer = tokenizer
+        if self.rp_world_size > 1 and not self.padding_free:
+            raise NotImplementedError(
+                f'The world_size {self.world_size} needs ulysses/ring-attention, which needs --padding_free true')
+
+    def _mask_qkv(self, query, key, value, mask):
+        mask = mask.unsqueeze(2).unsqueeze(3)
+        query = query * mask
+        value = value * mask
+        mask = (~mask) * -1e5  # for bf16
+        key = key + mask.to(key.dtype)
+        return query, key, value
+
+    def pad(self, tensor, padding_value, position_ids=None, dim=1):
+        """Pad tensor for sequence parallel"""
+        if self.rp_world_size > 1:
+            world_size = self.world_size * 2
+        else:
+            world_size = self.world_size
+
+        def _do_pad(tensor):
+            length = tensor.shape[dim]
+            pad_num = world_size - (length % world_size)
+            if pad_num == 0 or pad_num == world_size:
+                return tensor
+            if not isinstance(padding_value, torch.Tensor):
+                # ids
+                pad_shape = ((*tensor.shape[:dim], pad_num, *tensor.shape[dim + 1:]) if dim != -1 else
+                             (*tensor.shape[:dim], pad_num))
+                pad = torch.full(pad_shape, padding_value, dtype=tensor.dtype, device=tensor.device)
+                tensor = torch.cat([tensor, pad], dim=dim)
+            else:
+                # For embeddings
+                tensor = torch.cat([tensor, padding_value.unsqueeze(0).repeat(tensor.shape[0], pad_num, 1)], dim=dim)
+            return tensor
+
+        if position_ids is not None and self.rp_world_size > 1:
+            cu_seqlens = get_cu_seqlens_from_position_ids(position_ids)
+            all_tensors = []
+            for i in range(len(cu_seqlens) - 1):
+                if dim == 1:
+                    sub_tensor = tensor[:, cu_seqlens[i]:cu_seqlens[i + 1]]
+                elif dim == -1:
+                    sub_tensor = tensor[..., cu_seqlens[i]:cu_seqlens[i + 1]]
+                else:
+                    raise NotImplementedError()
+                all_tensors.append(_do_pad(sub_tensor))
+            tensor = torch.cat(all_tensors, dim=dim)
+
+        return _do_pad(tensor)
+
+    def gather(self, local_output, dim: int, position_ids=None):
+        """Gather tensor for sequence parallel - reverse of split"""
+        if self.world_size == 1:
+            return local_output
+
+        if self.rp_world_size > 1:
+            input_dim = local_output.dim()
+            assert input_dim >= 2
+
+            if position_ids is not None:
+                position_ids = self.pad(position_ids, padding_value=-1, position_ids=position_ids)
+
+            # Step 1: Gather from all sequence parallel ranks
+            # Each sp_rank has its own piece, we need to gather them first
+            gathered_sp = [torch.zeros_like(local_output) for _ in range(self.sp_world_size)]
+            torch.distributed.all_gather(gathered_sp, local_output.contiguous(), group=self.sp_group)
+
+            # Concatenate the sp pieces to form the complete chunk for this rp_rank
+            rp_chunk = torch.cat(gathered_sp, dim=dim)
+
+            # Step 2: Gather all rp chunks
+            gathered_rp = [torch.zeros_like(rp_chunk) for _ in range(self.rp_world_size)]
+            torch.distributed.all_gather(gathered_rp, rp_chunk, group=self.rp_group)
+
+            cu_seqlens = get_cu_seqlens_from_position_ids(position_ids)
+            all_tensor_length = []
+            for i in range(len(cu_seqlens) - 1):
+                length = cu_seqlens[i + 1] - cu_seqlens[i]
+                padding_length = math.ceil(length / (self.world_size * 2)) * (self.world_size * 2)
+                all_tensor_length.append(padding_length)
+
+            full_output = torch.zeros(
+                [local_output.shape[0], sum(all_tensor_length), *local_output.shape[2:]], device=local_output.device)
+            for idx_rp, rp_tensor in enumerate(gathered_rp):  # rp world size
+                # re-group the zigzag to the correct order
+                accumulated_length = 0
+                for idx_seq, length in enumerate(all_tensor_length):  # sequence number
+                    local_length = length // self.rp_world_size
+                    local_tensor = rp_tensor[:, accumulated_length:local_length + accumulated_length]
+                    chunk_size = local_length // 2
+                    left_idx = accumulated_length * self.rp_world_size + idx_rp * chunk_size
+                    right_idx = accumulated_length * self.rp_world_size + (idx_rp + 1) * chunk_size
+                    full_output[:, left_idx:right_idx] = local_tensor[:, :chunk_size]
+                    left_idx = accumulated_length * self.rp_world_size + (2 * self.rp_world_size - idx_rp
+                                                                          - 1) * chunk_size
+                    right_idx = accumulated_length * self.rp_world_size + (2 * self.rp_world_size - idx_rp) * chunk_size
+                    full_output[:, left_idx:right_idx] = local_tensor[:, chunk_size:]
+                    accumulated_length += local_length
+
+            return full_output.contiguous()
+        else:
+            gathered_sp = torch.empty(
+                [local_output.shape[0] * self.sp_world_size] + list(local_output.shape[1:]),
+                dtype=local_output.dtype,
+                device=local_output.device)
+            dist.all_gather_into_tensor(gathered_sp, local_output, group=self.sp_group)
+            gathered_sp = torch.cat(gathered_sp.split(local_output.shape[0], dim=0), dim=dim)
+            return gathered_sp.contiguous()
+
+    def _split_packed(self, value, cu_seqlens, dim=1):
+        """Split and re-group in zigzag"""
+        local_values = []
+        for i in range(len(cu_seqlens) - 1):
+            start, end = cu_seqlens[i], cu_seqlens[i + 1]
+            if dim == 1:
+                sub_value = value[:, start:end]
+            elif dim == -1:
+                sub_value = value[..., start:end]
+            else:
+                raise NotImplementedError()
+            local_value = sub_value.chunk(2 * self.rp_world_size, dim=dim)
+            local_values.extend([
+                local_value[self.rp_rank],
+                local_value[2 * self.rp_world_size - 1 - self.rp_rank],
+            ])
+        return torch.cat(local_values, dim=dim).contiguous()
+
+    def split(self, input, dim: int, position_ids=None):
+        """Split tensor for sequence parallel"""
+        if self.world_size == 1:
+            return input
+
+        if self.rp_world_size > 1:
+            input_dim = input.dim()
+            assert input_dim >= 2
+            cu_seqlens = get_cu_seqlens_from_position_ids(position_ids)
+            assert torch.all(cu_seqlens % (2 * self.rp_world_size) == 0)
+            value_chunks = self._split_packed(input, cu_seqlens, dim=dim)
+            local_value = value_chunks.chunk(self.sp_world_size, dim=dim)[self.sp_rank].contiguous()
+            return local_value
+        else:
+            rank = self.sp_rank
+            dim_size = input.size(dim)
+            assert dim_size % self.sp_world_size == 0, (
+                f'The dimension to split ({dim_size}) is not a multiple of '
+                f'world size ({self.sp_world_size}), cannot split tensor evenly')
+
+            tensor_list = torch.split(input, dim_size // self.sp_world_size, dim=dim)
+            output = tensor_list[rank].contiguous()
+            return output
+
+    def pad_and_split_mm_tokens(self, visual_mask, mm_embeds):
+        input_ids = self.extra_kwargs['input_ids']
+        empty_embeds = torch.empty(
+            (input_ids.shape[0], input_ids.shape[1], mm_embeds.shape[-1])).to(mm_embeds.device).to(mm_embeds.dtype)
+        empty_embeds[visual_mask] = mm_embeds
+
+        embeds = SimpleNamespace(weight=mm_embeds)
+
+        _, split_input_embeds, _, _, _, _, extra_values = self.pad_and_split_inputs(
+            None,
+            empty_embeds,
+            None,
+            None,
+            None,
+            None,
+            embeds,
+            self.real_position_ids,
+            extra_split_values=[(visual_mask, 0, -1)])
+        visual_mask = extra_values[0]
+        return visual_mask, split_input_embeds[visual_mask]
+
+    def pad_and_split_inputs(self,
+                             input_ids,
+                             input_embeds,
+                             labels,
+                             position_ids,
+                             attention_mask,
+                             loss_scale,
+                             embed_tokens=None,
+                             real_position_ids=None,
+                             extra_split_values=None):
+        """Common implementation for padding and splitting inputs
+
+        When a sequence comes, it will be split into rp_world_size * 2 sub tensors, and group them as the
+        zigzag order. So we get rp_world_size tensors, then we split each tensor to sp_world_size ones.
+        So, we should first pad the original sequence to the length can be divided by 2 * world_size, then re-group it.
+
+        Only support padding_free for ring-attention, because non-padding_free mode needs another pad/split workflow
+        man, that's a lot of work...
+
+        Args:
+            input_ids: input_ids
+            input_embeds: input_embeds
+            labels: labels
+            position_ids: position_ids or, position_ids for mrope
+            attention_mask: attention_mask
+            loss_scale: loss_scale
+            embed_tokens: embed_tokens
+            real_position_ids: the real position_ids to represent the seq length information
+            extra_split_values: List of Tuples for extra split values, e.g.: (tensor, pad_value, split_dim)
+        """
+        tokenizer = self.tokenizer
+        real_position_ids = real_position_ids if real_position_ids is not None else position_ids
+        extra_values = []
+        batch_size = input_ids.shape[
+            0] if input_ids is not None else input_embeds.shape[0] if input_embeds is not None else None
+        if real_position_ids is not None and batch_size is not None and real_position_ids.shape[0] == batch_size:
+            # TODO clone everytime, but the position_ids is a small tensor
+            self.extra_kwargs['text_position_ids'] = real_position_ids.clone()
+        if input_ids is not None:
+            input_ids = self.pad(input_ids, padding_value=tokenizer.pad_token_id, position_ids=real_position_ids)
+            self.extra_kwargs['input_ids'] = input_ids.clone()
+        if input_embeds is not None:
+            pad_emb = torch.zeros(
+                (1, embed_tokens.weight.shape[-1])).to(embed_tokens.weight.device).to(embed_tokens.weight.dtype)
+            input_embeds = self.pad(input_embeds, padding_value=pad_emb, position_ids=real_position_ids)
+        batch_size = input_ids.shape[
+            0] if input_ids is not None else input_embeds.shape[0] if input_embeds is not None else 1
+        if position_ids is not None:
+            position_ids = self.pad(position_ids, padding_value=-1, position_ids=real_position_ids, dim=-1)
+        if labels is not None:
+            labels = self.pad(labels, padding_value=-100, position_ids=real_position_ids)
+        if loss_scale is not None:
+            loss_scale = self.pad(loss_scale, padding_value=0., position_ids=real_position_ids)
+        if real_position_ids is not None:
+            real_position_ids = self.pad(real_position_ids, padding_value=-1, position_ids=real_position_ids)
+        if (input_ids is not None or input_embeds is not None) and batch_size > 1:
+            # not padding_free, so not ring-attention
+            inputs = input_ids if input_ids is not None else input_embeds
+            attn_shape = inputs.shape[1]  # The sequence length
+            if attention_mask is None:
+                attention_mask = torch.ones_like(real_position_ids)
+            # no need position_ids here, because padding_free does not need attention_mask,
+            # so this is not ring-attention
+            attention_mask = self.pad(attention_mask, padding_value=0)
+            cache_position = torch.arange(0, attn_shape, device=inputs.device)
+            # pad attention mask to 4d to avoid calculation errors
+            if hasattr(self, 'causal_mask_func') and self.causal_mask_func is not None:
+                attention_mask = self.causal_mask_func(attention_mask, inputs.to(self.model_dtype), cache_position,
+                                                       None, None)
+        if extra_split_values is not None:
+            for (tensor, pad_value, split_dim) in extra_split_values:
+                extra_values.append(
+                    self.pad(tensor, padding_value=pad_value, position_ids=real_position_ids, dim=split_dim))
+        if input_ids is not None:
+            input_ids = self.split(input_ids, dim=1, position_ids=real_position_ids)
+        if input_embeds is not None:
+            input_embeds = self.split(input_embeds, dim=1, position_ids=real_position_ids)
+        if labels is not None:
+            labels = torch.roll(labels, shifts=-1, dims=-1)
+            labels = self.split(labels, dim=-1, position_ids=real_position_ids)
+        if loss_scale is not None:
+            loss_scale = torch.roll(loss_scale, shifts=-1, dims=-1)
+            loss_scale = self.split(loss_scale, dim=-1, position_ids=real_position_ids)
+
+        if position_ids is not None:
+            position_ids = self.split(position_ids, dim=-1, position_ids=real_position_ids)
+        if extra_split_values is not None:
+            for i in range(len(extra_values)):
+                extra_values[i] = self.split(
+                    extra_values[i], dim=extra_split_values[i][2], position_ids=real_position_ids)
+        return input_ids, input_embeds, labels, position_ids, attention_mask, loss_scale, extra_values
+
+    def _gather_object_dp(self, input_data):
+        """Gather object for data parallel"""
+        input_data_list = [None] * self.dp_world_size
+        dist.all_gather_object(input_data_list, input_data, group=self.dp_group)
+        return [x for y in input_data_list for x in y]
+
+    def _init_device_mesh(self):
+        """Initialize device mesh for sequence and ring parallel.
+
+        The logic is unified:
+        1. Determine the Sequence Parallel (SP) size first based on GCD to satisfy constraints.
+        2. Allocate all remaining model parallelism to Ring Parallel (RP).
+        """
+        _, _, world_size, _ = get_dist_setting()
+        self.dp_world_size = world_size // self.world_size
+        # SP size is the GCD of num_heads and world_size, guaranteeing it divides both.
+        self.sp_world_size = math.gcd(self.num_heads, self.world_size)
+        # RP takes the remaining factor so all model-parallel GPUs are used.
+        self.rp_world_size = self.world_size // self.sp_world_size
+
+        if self.rp_world_size > 1:
+            mesh_shape = (self.dp_world_size, self.rp_world_size, self.sp_world_size)
+            mesh_dim_names = ('data', 'ring', 'sequence')
+        else:
+            mesh_shape = (self.dp_world_size, self.sp_world_size)
+            mesh_dim_names = ('data', 'sequence')
+        self.device_mesh = init_device_mesh(
+            get_device().split(':')[0], mesh_shape=mesh_shape, mesh_dim_names=mesh_dim_names)
+
+    def _dim_group(self, dim_name: str):
+        """Return the process group of the given mesh dim, or None if absent."""
+        if not self.device_mesh or dim_name not in self.device_mesh.mesh_dim_names:
+            return None
+        return self.device_mesh[dim_name].get_group()
+
+    def _dim_rank(self, dim_name: str, default: int) -> int:
+        """Return the rank within the given mesh dim, or `default` if absent."""
+        group = self._dim_group(dim_name)
+        return dist.get_rank(group) if group is not None else default
+
+    @property
+    def sp_group(self):
+        """Return the sequence parallel group"""
+        return self._dim_group('sequence')
+
+    def enabled(self):
+        return self.world_size is not None and self.world_size > 1
+
+    @property
+    def sp_rank(self):
+        """Return the sequence parallel rank"""
+        return self._dim_rank('sequence', default=0)
+
+    @property
+    def dp_group(self):
+        """Return the data parallel group"""
+        return self._dim_group('data')
+
+    @property
+    def dp_rank(self):
+        """Return the data parallel rank"""
+        return self._dim_rank('data', default=0)
+
+    @property
+    def rp_group(self):
+        """Return the ring parallel group"""
+        return self._dim_group('ring')
+
+    @property
+    def rp_rank(self):
+        """Return the ring parallel rank"""
+        return self._dim_rank('ring', default=-1)
+
+    def prepare_inputs(self, inputs):
+        """Prepare inputs
+
+        1. set extra_kwargs['text_position_ids']
+        2. split labels
+        """
+        position_ids = inputs.get('text_position_ids')
+        input_ids = inputs.get('input_ids')
+        if position_ids is None:
+            position_ids = inputs.get('position_ids')
+        if position_ids is not None and input_ids is not None and position_ids.shape[0] == input_ids.shape[0]:
+            self.extra_kwargs['text_position_ids'] = position_ids.clone()
+        if input_ids is not None:
+            self.extra_kwargs['input_ids'] = input_ids.clone()
+        if 'labels' in inputs:
+            labels = inputs['labels']
+            _, _, labels, _, _, _, _ = self.pad_and_split_inputs(
+                None, None, labels, None, None, None, real_position_ids=position_ids)
+            inputs['labels'] = labels
+
+
+sequence_parallel = SequenceParallel()

swift/sequence_parallel/ulysses.py

@@ -0,0 +1,147 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+import torch
+import torch.distributed as dist
+from typing import Any, Tuple
+
+
+# Code borrowed from deepspeed, here is why:
+# 1. Reduce the dependency
+# 2. The original code is complex
+def _generate_layout_params(scatter_idx, seq_world_size, input):
+    if scatter_idx < 2:
+        bs, global_seq_len, num_local_head, head_dim = input.shape
+        pre_all2all_inp_shape = [bs, seq_world_size, global_seq_len // seq_world_size, num_local_head, head_dim]
+        pre_all2all_permute_idx = (1, 0, 2, 3, 4)
+
+        post_all2all_permute_idx = (1, 2, 0, 3, 4)
+        post_all2all_res_shape = [bs, global_seq_len // seq_world_size, seq_world_size * num_local_head, head_dim]
+    else:
+        bs, local_seq_len, num_total_head, head_dim = input.shape
+        assert num_total_head % seq_world_size == 0, (f'Number of heads ({num_total_head}) must be divisible '
+                                                      f'by the sequence parallel size ({seq_world_size})!')
+        pre_all2all_inp_shape = [bs, local_seq_len, seq_world_size, num_total_head // seq_world_size, head_dim]
+        pre_all2all_permute_idx = (2, 0, 1, 3, 4)
+
+        post_all2all_permute_idx = (1, 0, 2, 3, 4)
+        post_all2all_res_shape = [bs, seq_world_size * local_seq_len, num_total_head // seq_world_size, head_dim]
+
+    return pre_all2all_permute_idx, pre_all2all_inp_shape, post_all2all_permute_idx, post_all2all_res_shape
+
+
+def post_all2all(permute_idx, res_shape):
+    """
+    Post-processing function for `all2all` communication.
+    """
+
+    def post_func(input):
+        if permute_idx is not None:
+            input = input.permute(permute_idx).contiguous()
+        output = input.reshape(res_shape).contiguous()
+
+        return output
+
+    return post_func
+
+
+def pre_all2all_fun(permute_idx, inp_shape, input):
+    """
+    Pre-processing function for `all2all` communication.
+    """
+    input_t = input.reshape(inp_shape).contiguous()
+    if permute_idx is not None:
+        input_t = input_t.permute(permute_idx).contiguous()
+    return input_t
+
+
+def single_all_to_all(input, scatter_idx, gather_idx, group, **kwargs):
+    seq_world_size = dist.get_world_size(group)
+    num_heads = input.shape[2]
+    if num_heads % seq_world_size != 0 and not scatter_idx < 2:
+        raise NotImplementedError(f'num_heads {num_heads} cannot be split by sp world size {seq_world_size}')
+    pre_all2all_permute_idx, pre_all2all_inp_shape, post_all2all_permute_idx, post_all2all_res_shape = (
+        _generate_layout_params(scatter_idx, seq_world_size, input))
+
+    input_t = pre_all2all_fun(pre_all2all_permute_idx, pre_all2all_inp_shape, input)
+
+    post_all2all_fun = post_all2all(post_all2all_permute_idx, post_all2all_res_shape)
+    output = torch.empty_like(input_t)
+    dist.all_to_all_single(output, input_t, group=group)
+
+    res = post_all2all_fun(output)
+    return res
+
+
+class _SeqAllToAll(torch.autograd.Function):
+
+    @staticmethod
+    def forward(
+        ctx: Any,
+        group: dist.ProcessGroup,
+        input: torch.Tensor,
+        scatter_idx: int,
+        gather_idx: int,
+    ) -> torch.Tensor:
+        ctx.group = group
+        ctx.scatter_idx = scatter_idx
+        ctx.gather_idx = gather_idx
+        res = single_all_to_all(input, scatter_idx, gather_idx, group)
+        return res
+
+    @staticmethod
+    def backward(ctx: Any, *grad_output: torch.Tensor) -> Tuple[None, torch.Tensor, None, None]:
+        return None, _SeqAllToAll.apply(ctx.group, *grad_output, ctx.gather_idx, ctx.scatter_idx), None, None
+
+
+class DistributedAttention(torch.nn.Module):
+
+    def __init__(
+        self,
+        local_attention,
+        sequence_parallel,
+        scatter_idx: int = 2,
+        gather_idx: int = 1,
+    ) -> None:
+        super(DistributedAttention, self).__init__()
+        self.local_attn = local_attention
+        self.sequence_parallel = sequence_parallel
+        self.scatter_idx = scatter_idx
+        self.gather_idx = gather_idx
+
+    def forward(self, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, attention_mask: torch.Tensor, *args:
+                Any, **kwargs) -> torch.Tensor:
+        if self.sequence_parallel.world_size == 1:
+            return self.local_attn(query, key, value, attention_mask, *args, **kwargs)
+
+        # gather ulysses first, ring-attention next
+        if self.sequence_parallel.sp_world_size > 1:
+            query_layer = _SeqAllToAll.apply(self.sequence_parallel.sp_group, query, self.scatter_idx, self.gather_idx)
+            key_layer = _SeqAllToAll.apply(self.sequence_parallel.sp_group, key, self.scatter_idx, self.gather_idx)
+            value_layer = _SeqAllToAll.apply(self.sequence_parallel.sp_group, value, self.scatter_idx, self.gather_idx)
+        else:
+            query_layer, key_layer, value_layer = query, key, value
+
+        if self.sequence_parallel.rp_world_size > 1:
+            # if need ring-attention
+            kwargs.pop('position_ids', None)
+            # Get the real position ids, this is filled by `sequence_parallel.prepare_inputs`
+            # real position ids is different from the position_ids when model uses mrope
+            position_ids = self.sequence_parallel.real_position_ids
+            # pad and split it by zigzag method
+            position_ids = self.sequence_parallel.pad(position_ids, padding_value=-1, position_ids=position_ids)
+        else:
+            # only ulysses
+            position_ids = kwargs.pop('position_ids')
+            if position_ids is not None:
+                # Reuse the generic gather path so 2D and 3D position_ids share the same SP behavior.
+                position_ids = self.sequence_parallel.gather(position_ids.contiguous(), dim=-1, position_ids=None)
+
+        context_layer = self.local_attn(
+            query_layer, key_layer, value_layer, attention_mask, *args, position_ids=position_ids, **kwargs)
+
+        if self.sequence_parallel.sp_world_size > 1:
+            output = _SeqAllToAll.apply(self.sequence_parallel.sp_group, context_layer, self.gather_idx,
+                                        self.scatter_idx)
+        else:
+            output = context_layer
+
+        return output

swift/sequence_parallel/utils.py

@@ -0,0 +1,171 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+import math
+import torch
+import torch.distributed as dist
+from torch.nn import CrossEntropyLoss
+from torch.utils.data import Sampler
+from typing import Any, Iterator
+
+from swift.dataloader import DataLoaderDispatcher
+from .sequence_parallel import sequence_parallel
+
+
+class GatherTensor(torch.autograd.Function):
+    """Gather tensor from sequence group (autograd supported)"""
+
+    @staticmethod
+    def forward(ctx, tensor, dim=0, position_ids=None):
+        ctx.dim = dim
+        if position_ids is not None:
+            position_ids = sequence_parallel.pad(position_ids, padding_value=-1, position_ids=position_ids)
+        ctx.position_ids = position_ids
+        return sequence_parallel.gather(tensor, dim=dim, position_ids=position_ids)
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        grad_input = sequence_parallel.split(grad_output, dim=ctx.dim, position_ids=ctx.position_ids)
+        return grad_input, None, None
+
+
+class GatherLoss(torch.autograd.Function):
+    """Gather loss from sequence group"""
+
+    @staticmethod
+    def forward(ctx, loss, labels, gather_idx=None, position_ids=None):
+        """
+        Args:
+            loss: loss tensor after splitting
+            labels: labels tensor after splitting
+            gather_idx: gather the tensors on this dim
+        """
+        # change from label.shape to loss, because label may be None
+        ctx.scatter_shape = loss.shape[gather_idx or 0]
+        ctx.gather_idx = gather_idx or 0
+        if position_ids is not None:
+            position_ids = sequence_parallel.pad(position_ids, padding_value=-1, position_ids=position_ids)
+        ctx.position_ids = position_ids
+        output = sequence_parallel.gather(loss, dim=ctx.gather_idx, position_ids=position_ids)
+        if labels is not None:
+            labels_output = sequence_parallel.gather(labels, dim=ctx.gather_idx, position_ids=position_ids)
+        else:
+            labels_output = None
+        return output, labels_output
+
+    @staticmethod
+    def backward(ctx, *grad_output):
+        _grad = grad_output[0] * sequence_parallel.world_size
+        if sequence_parallel.rp_world_size > 1:
+            _grad = sequence_parallel.split(_grad, dim=ctx.gather_idx, position_ids=ctx.position_ids).contiguous()
+        else:
+            _grad = _grad.split(
+                ctx.scatter_shape, dim=ctx.gather_idx)[dist.get_rank(sequence_parallel.sp_group)].contiguous()
+        return _grad, None, None, None
+
+
+class ChunkedCrossEntropyLoss(torch.autograd.Function):
+
+    @staticmethod
+    def forward(ctx, logits, labels, chunk_size):
+        ctx.save_for_backward(logits, labels)
+        ctx.chunk_size = chunk_size
+
+        losses = []
+        for i in range(math.ceil(logits.shape[0] / chunk_size)):
+            l_start = i * chunk_size
+            l_end = min((i + 1) * chunk_size, logits.shape[0])
+            logits_chunk = logits[l_start:l_end]
+            labels_chunk = labels[l_start:l_end]
+            loss_fct = CrossEntropyLoss(reduction='none')
+            loss_chunk = loss_fct(logits_chunk, labels_chunk)
+            losses.append(loss_chunk)
+            del logits_chunk
+            del labels_chunk
+        all_losses = torch.cat(losses)
+        return all_losses
+
+    @staticmethod
+    def backward(ctx: Any, *grad_outputs: Any):
+        logits, labels = ctx.saved_tensors
+        chunk_size = ctx.chunk_size
+
+        for i in range(math.ceil(logits.shape[0] / chunk_size)):
+            l_start = i * chunk_size
+            l_end = min((i + 1) * chunk_size, logits.shape[0])
+            logits_chunk = logits[l_start:l_end].detach().requires_grad_(True)
+            labels_chunk = labels[l_start:l_end]
+            loss_fct = CrossEntropyLoss(reduction='none')
+            with torch.enable_grad():
+                loss_chunk = loss_fct(logits_chunk, labels_chunk)
+                grad_output_chunk = grad_outputs[0][l_start:l_end]
+                _loss_chunk = (loss_chunk * grad_output_chunk).sum()
+                grad_chunk = torch.autograd.grad(_loss_chunk, logits_chunk, retain_graph=False)[0]
+                logits[l_start:l_end] = grad_chunk
+
+        return logits, None, None
+
+
+class SequenceParallelSampler(Sampler):
+    """Sampler for sequence parallel training"""
+
+    def __init__(self, sp_instance, dataset, shuffle: bool = True, seed=None, round_up: bool = True) -> None:
+        self.sp_instance = sp_instance
+        rank = dist.get_rank(sp_instance.device_mesh['data'].get_group())
+        world_size = sp_instance.device_mesh['data'].size()
+        self.rank = rank
+        self.world_size = world_size
+
+        self.dataset = dataset
+        self.shuffle = shuffle
+        assert seed is not None
+        self.seed = seed
+        self.epoch = 0
+        self.round_up = round_up
+
+        if self.round_up:
+            self.num_samples = math.ceil(len(self.dataset) / world_size)
+            self.total_size = self.num_samples * self.world_size
+        else:
+            self.num_samples = math.ceil((len(self.dataset) - rank) / world_size)
+            self.total_size = len(self.dataset)
+
+    def __iter__(self) -> Iterator[int]:
+        if self.shuffle:
+            g = torch.Generator()
+            g.manual_seed(self.seed + self.epoch)
+            indices = torch.randperm(len(self.dataset), generator=g).tolist()
+        else:
+            indices = torch.arange(len(self.dataset)).tolist()
+
+        if self.round_up:
+            indices = (indices * int(self.total_size / len(indices) + 1))[:self.total_size]
+
+        indices = indices[self.rank:self.total_size:self.world_size]
+        return iter(indices)
+
+    def __len__(self) -> int:
+        return self.num_samples
+
+    def set_epoch(self, epoch: int) -> None:
+        self.epoch = epoch
+
+
+class SequenceParallelDispatcher(DataLoaderDispatcher):
+    """Dispatcher for sequence parallel training"""
+
+    def __init__(self, dataloader, sp_instance, device=None, skip_batches: int = 0):
+        super().__init__(dataloader)
+        self.sp_instance = sp_instance
+        self.device = device
+        self.skip_batches = skip_batches
+
+    @property
+    def rank(self):
+        return self.sp_instance.dp_rank if dist.is_initialized() else 0
+
+    @property
+    def world_size(self):
+        return self.sp_instance.dp_world_size if dist.is_initialized() else 1
+
+    @property
+    def group(self):
+        return self.sp_instance.dp_group if dist.is_initialized() else 1

swift/sequence_parallel/zigzag_ring_attn.py

@@ -0,0 +1,711 @@
+# Some code borrowed from the awesome work: https://github.com/zhuzilin/ring-flash-attention
+# Copyright (c) ModelScope Contributors. All rights reserved.
+import inspect
+import torch
+import torch.distributed as dist
+import torch.nn.functional as F
+from functools import cache
+
+from .ring_utils import RingComm
+from .zigzag_ring_attn_npu import is_npu_tensor, npu_backward, npu_forward
+
+
+def get_half_index(cu_seqlens, *, front: bool):
+    """Get half of the index
+
+    Args:
+        cu_seqlens: The cu_seqlens passed into flash_attn
+        front: The head part or the tail part
+
+    Returns:
+        The slice or the tensor mask.
+    """
+
+    if len(cu_seqlens) == 2:
+        if front:
+            return slice(None, cu_seqlens[-1] // 2)
+        else:
+            return slice(cu_seqlens[-1] // 2, None)
+
+    index = torch.zeros((cu_seqlens[-1].item(), ), dtype=torch.bool)
+    for i in range(len(cu_seqlens) - 1):
+        start, end = cu_seqlens[i], cu_seqlens[i + 1]
+        if front:
+            end = (start + end) // 2
+        else:
+            start = (start + end) // 2
+        index[start:end] = True
+    return index
+
+
+@torch.jit.script
+def get_half_lse(lse, cu_seqlens, *, front: bool):
+    """Get half of the lse
+
+    Args:
+        lse: The input lse, with shape [num_heads, seqlen]
+        cu_seqlens: The cu_seqlens passed into flash_attn
+        front: The head part or the tail part
+
+    Returns:
+        The filtered lse with the same shape as lse
+    """
+    new_lse = torch.empty(
+        (lse.shape[0], lse.shape[1] // 2),
+        dtype=lse.dtype,
+        device=lse.device,
+    )
+    for i in range(len(cu_seqlens) - 1):
+        start, end = cu_seqlens[i].item(), cu_seqlens[i + 1].item()
+        new_start, new_end = start // 2, end // 2
+        if front:
+            end -= (end - start) // 2
+        else:
+            start += (end - start) // 2
+        new_lse[:, new_start:new_end] = lse[:, start:end]
+    return new_lse
+
+
+def update_out_and_lse(out, lse, block_out, block_lse):
+    """Update output and lse:
+    new_lse = lse + torch.log(1 + torch.exp(block_lse - lse))
+    torch.exp(lse - new_lse) * out + torch.exp(block_lse - new_lse) * block_out
+    # For additional context and discussion, please refer to:
+    # https://github.com/zhuzilin/ring-flash-attention/pull/34#issuecomment-2076126795
+
+    Args:
+        out: The accumulated output of shape [seqlen, num_heads, hidden_size]
+        lse: The accumulated lse of shape [num_heads, seqlen]
+        block_out: The current block output of shape [seqlen, num_heads, hidden_size]
+        block_lse: The current block lse of shape [num_heads, seqlen]
+
+    Returns:
+        The updated output[seqlen, num_heads, hidden_size] and lse (shape: [seqlen, num_heads, 1]) and
+            the intermediate value of torch.sigmoid(block_lse - lse) (shape: [seqlen, num_heads, 1])
+    """
+    if out is None:
+        out = block_out.to(torch.float32)
+        lse = block_lse.transpose(-2, -1).unsqueeze(dim=-1)
+        sig_diff = None
+    else:
+        block_out = block_out.to(torch.float32)
+        block_lse = block_lse.transpose(-2, -1).unsqueeze(dim=-1)
+
+        diff = block_lse - lse
+        sig_diff = torch.sigmoid(diff)
+
+        out = out - sig_diff * (out - block_out)  # (..., D)
+        lse = lse - F.logsigmoid(lse - block_lse)  # (..., 1)
+    return out, lse, sig_diff
+
+
+@cache
+def _get_default_args(func):
+    spec = inspect.getfullargspec(func)
+    defaults = spec.defaults if spec.defaults is not None else ()
+    padded_defaults = (None, ) * (len(spec.args) - len(defaults)) + defaults
+    args = dict(zip(spec.args, padded_defaults))
+    if 'softcap' in args:
+        args['softcap'] = 0.0
+    return args
+
+
+def get_default_args(func):
+    if inspect.isfunction(func):
+        return _get_default_args(func)
+    else:
+        # Use the origin _init_fn in CustomOpDef
+        return _get_default_args(func._init_fn)
+
+
+def squeeze_batch(*t):
+    """Squeeze the batch dim of tensors"""
+    tensors = []
+    for sub in t:
+        if sub.shape[0] == 1:
+            tensors.append(sub.squeeze(0))
+        else:
+            tensors.append(sub)
+    return tuple(tensors)
+
+
+def padding(tensor, cu_seqlens, padding_value, front):
+    """Pad the tensor according to the cu_seqlens
+
+    Args:
+        tensor: The input tensor of shape [seqlen, *]
+        cu_seqlens: The cu_seqlens
+        padding_value: The padding value
+        front: tensor is the head or tail part
+    """
+    if len(cu_seqlens) == 2:
+        if front:
+            return torch.cat((tensor, torch.full_like(tensor, padding_value).to(tensor.dtype).to(tensor.device)), dim=0)
+        else:
+            return torch.cat((torch.full_like(tensor, padding_value).to(tensor.dtype).to(tensor.device), tensor), dim=0)
+
+    output = []
+    acc = 0
+    for i in range(len(cu_seqlens) - 1):
+        start, end = cu_seqlens[i], cu_seqlens[i + 1]
+        half_len = (end - start) // 2
+        acc += half_len
+        half_start = start // 2
+        local_tensor = tensor[half_start:half_start + half_len]
+        if front:
+            output.append(local_tensor)
+            output.append(torch.full_like(local_tensor, padding_value).to(local_tensor.dtype).to(local_tensor.device))
+        else:
+            output.append(torch.full_like(local_tensor, padding_value).to(local_tensor.dtype).to(local_tensor.device))
+            output.append(local_tensor)
+    assert acc == tensor.shape[0]
+    return torch.cat(output)
+
+
+def forward(q, k, v, causal, cu_seqlens, max_seqlen, block_seq_len, dropout_p, softmax_scale, alibi_slopes,
+            window_size):
+    seqlen_q = q.shape[0]
+    seqlen_kv = k.shape[0]
+    half_cu_seqlens = cu_seqlens // 2
+    half_max_seqlen = max_seqlen // 2
+    cu_seqlens_q = half_cu_seqlens if seqlen_q == block_seq_len else cu_seqlens
+    max_seqlen_q = half_max_seqlen if seqlen_q == block_seq_len else max_seqlen
+    cu_seqlens_kv = half_cu_seqlens if seqlen_kv == block_seq_len else cu_seqlens
+    max_seqlen_kv = half_max_seqlen if seqlen_kv == block_seq_len else max_seqlen
+    if is_npu_tensor(q):
+        # Keep the ring schedule in this file unchanged; only the per-block
+        # flash-attn call is swapped to Ascend's TND varlen attention kernel.
+        return npu_forward(
+            q,
+            k,
+            v,
+            causal,
+            cu_seqlens_q,
+            cu_seqlens_kv,
+            dropout_p,
+            softmax_scale,
+            deterministic=False,
+            window_size=window_size,
+        )
+    from flash_attn.flash_attn_interface import _flash_attn_varlen_forward
+    params = get_default_args(_flash_attn_varlen_forward).copy()
+    params.update({
+        'q': q,
+        'k': k,
+        'v': v,
+        # the first half and the second half are the same
+        'cu_seqlens_q': cu_seqlens_q,
+        'cu_seqlens_k': cu_seqlens_kv,
+        'max_seqlen_q': max_seqlen_q,
+        'max_seqlen_k': max_seqlen_kv,
+        'dropout_p': dropout_p,
+        'softmax_scale': softmax_scale,
+        'causal': causal,
+        'alibi_slopes': alibi_slopes,
+        'return_softmax': True and dropout_p > 0,
+    })
+    if 'window_size' in params:
+        params.update({'window_size': window_size})
+    else:
+        params.update({
+            'window_size_left': window_size[0],
+            'window_size_right': window_size[1],
+        })
+    assert k.shape[-0] == cu_seqlens_kv[-1]
+    assert q.shape[-0] == cu_seqlens_q[-1]
+    assert max_seqlen_q == (cu_seqlens_q[1:] - cu_seqlens_q[:-1]).max().item()
+    assert max_seqlen_kv == (cu_seqlens_kv[1:] - cu_seqlens_kv[:-1]).max().item()
+    outputs = _flash_attn_varlen_forward(**params)
+    if len(outputs) == 8:
+        block_out, _, _, _, _, block_lse, _, _ = outputs
+    else:
+        assert len(outputs) == 4
+        block_out, block_lse, _, _ = outputs
+    return block_out, block_lse
+
+
+def backward(dout, q, k, v, out, softmax_lse, causal, cu_seqlens, max_seqlen, block_seq_len, dq_buffer, dk_buffer,
+             dv_buffer, dropout_p, softmax_scale, alibi_slopes, deterministic, window_size):
+    seqlen_q = q.shape[0]
+    seqlen_kv = k.shape[0]
+
+    half_cu_seqlens = cu_seqlens // 2
+    half_max_seqlen = max_seqlen // 2
+    cu_seqlens_q = half_cu_seqlens if seqlen_q == block_seq_len else cu_seqlens
+    max_seqlen_q = half_max_seqlen if seqlen_q == block_seq_len else max_seqlen
+    cu_seqlens_kv = half_cu_seqlens if seqlen_kv == block_seq_len else cu_seqlens
+    max_seqlen_kv = half_max_seqlen if seqlen_kv == block_seq_len else max_seqlen
+    from flash_attn.flash_attn_interface import _flash_attn_varlen_backward
+    params = get_default_args(_flash_attn_varlen_backward).copy()
+    params.update({
+        'dout': dout,
+        'q': q,
+        'k': k,
+        'v': v,
+        'out': out,
+        'softmax_lse': softmax_lse,
+        'dq': dq_buffer[:seqlen_q],
+        'dk': dk_buffer[:seqlen_kv],
+        'dv': dv_buffer[:seqlen_kv],
+        # the first half and the second half are the same
+        'cu_seqlens_q': cu_seqlens_q,
+        'cu_seqlens_k': cu_seqlens_kv,
+        'max_seqlen_q': max_seqlen_q,
+        'max_seqlen_k': max_seqlen_kv,
+        'dropout_p': dropout_p,
+        'softmax_scale': softmax_scale,
+        'causal': causal,
+        'alibi_slopes': alibi_slopes,
+        'deterministic': deterministic,
+    })
+    assert dout.shape[0] == q.shape[0]
+    assert dout.shape[0] == out.shape[0]
+    assert softmax_lse.shape[1] == q.shape[0]
+    assert k.shape[0] == cu_seqlens_kv[-1]
+    assert q.shape[0] == cu_seqlens_q[-1]
+    assert max_seqlen_q == (cu_seqlens_q[1:] - cu_seqlens_q[:-1]).max().item()
+    assert max_seqlen_kv == (cu_seqlens_kv[1:] - cu_seqlens_kv[:-1]).max().item()
+    if 'window_size' in params:
+        params.update({'window_size': window_size})
+    else:
+        params.update({
+            'window_size_left': window_size[0],
+            'window_size_right': window_size[1],
+        })
+    _flash_attn_varlen_backward(**params)
+
+
+def lse_grad(out, lse, block_out, block_lse, sig, grad_out, grad_lse):
+    """Calculate the grad of each block.
+
+    Args:
+        out: The accumulated output of shape [seqlen, num_heads, hidden_size]
+        lse: The accumulated lse of shape [num_heads, seqlen, 1]
+        block_out: The current block output of shape [seqlen, num_heads, hidden_size]
+        block_lse: The current block lse of shape [num_heads, seqlen, 1]
+        grad_out: The input grad of output of the current block shape [seqlen, num_heads, hidden_size]
+        grad_lse: The input grad of lse of the current block shape [num_heads, seqlen, 1]
+
+    Returns:
+        The accumulated grad of out and lse, and the grad of out and lse of the current block
+    """
+    grad_out_input = grad_out * (1 - sig)
+
+    grad_block_out = grad_out * sig
+
+    d_new_out_d_lse = (out - block_out) * (sig * (1 - sig))
+    grad_lse_input = (grad_out * d_new_out_d_lse).sum(dim=-1, keepdim=True)
+    grad_lse_input_final = grad_lse_input + grad_lse * torch.sigmoid(lse - block_lse)
+
+    grad_block_lse = -grad_lse_input_final + grad_lse
+
+    return grad_out_input, grad_lse_input_final, grad_block_out, grad_block_lse
+
+
+def zigzag_ring_flash_attn_varlen_forward(
+        process_group,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        cu_seqlens,
+        max_seqlen,
+        half_index0,
+        half_index1,
+        softmax_scale,
+        dropout_p=0,
+        causal=True,
+        window_size=(-1, -1),
+        alibi_slopes=None,
+        deterministic=False,
+):
+    assert causal, 'zigzag ring is meaningless for causal=False'
+    comm = RingComm(process_group)
+    q, k, v = squeeze_batch(q, k, v)
+    q1 = q[half_index1]
+    # Input cu_seqlens is the total length, divided by world_size to fit the split ones
+    cu_seqlens = cu_seqlens // comm.world_size
+    # Same with above
+    max_seqlen = max_seqlen // comm.world_size
+    block_seq_len = q.shape[0] // 2
+    out = None
+    lse = None
+    next_k, next_v = None, None
+
+    for step in range(comm.world_size):
+        # from step 0 to the last
+        if step + 1 != comm.world_size:
+            next_k, next_v = comm.send_recv_kv(k, v)
+        """
+        world_size = 4, total 8 parts
+        0/7 is group0
+        1/6 is group1
+        2/5 is group2
+        3/4 is group3
+        consider 1/6，take the query as the left axis, key as the top axis:
+        step 0:
+           1   6
+        1  ✅  ❎
+        6  ✅  ✅
+        all needed, causal=True
+        step 1(step <= comm.rank):
+           0   7
+        1  ✅  ❎
+        6  ✅  ❎
+        the first part of kv is needed, causal=False
+        step 2(step > comm.rank):
+           3   4
+        1  ❎  ❎
+        6  ✅  ✅
+        the second part of q is needed, causal=False
+        """
+        # Here block_lse shape: [num_heads, seqlen]
+        # lse shape: [seqlen, num_heads, 1]
+        if step == 0:
+            block_out, block_lse = forward(q, k, v, True, cu_seqlens, max_seqlen, block_seq_len, dropout_p,
+                                           softmax_scale, alibi_slopes, window_size)
+            out, lse, sig_diff = update_out_and_lse(out, lse, block_out, block_lse)
+        elif step <= comm.rank:
+            k0 = k[half_index0]
+            v0 = v[half_index0]
+            block_out, block_lse = forward(q, k0, v0, False, cu_seqlens, max_seqlen, block_seq_len, dropout_p,
+                                           softmax_scale, alibi_slopes, window_size)
+            out, lse, sig_diff = update_out_and_lse(out, lse, block_out, block_lse)
+        else:
+            block_out, block_lse = forward(q1, k, v, False, cu_seqlens, max_seqlen, block_seq_len, dropout_p,
+                                           softmax_scale, alibi_slopes, window_size)
+            out[half_index1], lse[half_index1], sig_diff = update_out_and_lse(out[half_index1], lse[half_index1],
+                                                                              block_out, block_lse)
+
+        if step + 1 != comm.world_size:
+            comm.wait()
+            k, v = next_k, next_v
+
+    out = out.to(q.dtype)
+    lse = lse.squeeze(dim=-1).transpose(0, 1)  # [num_heads, seqlen]
+    return out.unsqueeze(0), lse.unsqueeze(0)
+
+
+def zigzag_ring_flash_attn_varlen_backward(
+        process_group,
+        dout,
+        q,
+        k,
+        v,
+        out,
+        softmax_lse,
+        cu_seqlens,
+        max_seqlen,
+        half_index0,
+        half_index1,
+        softmax_scale,
+        dropout_p=0,
+        causal=True,
+        window_size=(-1, -1),
+        alibi_slopes=None,
+        deterministic=False,
+):
+    assert causal, 'zigzag ring is meaningless for causal=False'
+    if is_npu_tensor(q):
+        # NPU backward uses native flash-attn grad with the final ring out/lse
+        # patched into each block ctx. Missing kernel support should fail loudly.
+        return npu_backward(
+            process_group,
+            dout,
+            q,
+            k,
+            v,
+            out,
+            softmax_lse,
+            cu_seqlens,
+            max_seqlen,
+            half_index0,
+            half_index1,
+            softmax_scale,
+            dropout_p=dropout_p,
+            window_size=window_size,
+            deterministic=deterministic,
+        )
+    kv_comm = RingComm(process_group)
+    d_kv_comm = RingComm(process_group)
+    dk_comm_buffer = dv_comm_buffer = None
+    dq, dk, dv = None, None, None
+    next_dk, next_dv = None, None
+    next_k, next_v = None, None
+
+    # squeeze the axis of batch
+    dout, q, k, v, out, softmax_lse = squeeze_batch(dout, q, k, v, out, softmax_lse)
+    q1 = q[half_index1]
+    # Input cu_seqlens is the total length, divided by world_size to fit the split ones
+    cu_seqlens = cu_seqlens // kv_comm.world_size
+    # Same as above
+    max_seqlen = max_seqlen // kv_comm.world_size
+    # half of the part
+    block_seq_len = q.shape[0] // 2
+
+    # repeatly allocating buffer may be slow...
+    dq_buffer = torch.empty(q.shape, dtype=q.dtype, device=q.device)
+    dk_buffer = torch.empty(k.shape, dtype=k.dtype, device=k.device)
+    dv_buffer = torch.empty(v.shape, dtype=v.dtype, device=v.device)
+    origin_q, origin_k, origin_v = q, k, v
+
+    out_lse = []
+    fout = None
+    flse = None
+    # Recalculate forward with the same qkv to generate out_lse, used to calculate the grad
+    for step in range(kv_comm.world_size):
+        if step + 1 != kv_comm.world_size:
+            next_k, next_v = kv_comm.send_recv_kv(k, v)
+
+        if step == 0:
+            block_out, block_lse = forward(q, k, v, True, cu_seqlens, max_seqlen, block_seq_len, dropout_p,
+                                           softmax_scale, alibi_slopes, window_size)
+            fout, flse, sig_diff = update_out_and_lse(fout, flse, block_out, block_lse)
+        elif step <= kv_comm.rank:
+            k0 = k[half_index0]
+            v0 = v[half_index0]
+            block_out, block_lse = forward(q, k0, v0, False, cu_seqlens, max_seqlen, block_seq_len, dropout_p,
+                                           softmax_scale, alibi_slopes, window_size)
+            fout, flse, sig_diff = update_out_and_lse(fout, flse, block_out, block_lse)
+        else:
+            block_out, block_lse = forward(q1, k, v, False, cu_seqlens, max_seqlen, block_seq_len, dropout_p,
+                                           softmax_scale, alibi_slopes, window_size)
+            fout[half_index1], flse[half_index1], sig_diff = update_out_and_lse(fout[half_index1], flse[half_index1],
+                                                                                block_out, block_lse)
+
+        block_lse = block_lse.transpose(0, 1).unsqueeze(-1)
+        if step > kv_comm.rank:
+            # cat zeros because there are may be a half of the out/lse
+            block_out = padding(block_out, cu_seqlens, 0, front=False)
+            block_lse = padding(block_lse, cu_seqlens, -1e5, front=False)
+            sig_diff = padding(sig_diff, cu_seqlens, 0, front=False)
+
+        # save to out_lse
+        out_lse.append((fout, flse, block_out, block_lse, sig_diff))
+
+        if step + 1 != kv_comm.world_size:
+            kv_comm.wait()
+            k, v = next_k, next_v
+
+    current_dout = dout
+    current_dlse = torch.zeros_like(softmax_lse.transpose(0, 1).unsqueeze(-1))
+    block_gradients = {}
+
+    for i in reversed(range(len(out_lse))):
+        if i == 0:
+            # the first step does not need
+            continue
+        stored_out, stored_lse, stored_block_out, stored_block_lse, stored_sig = out_lse[i]
+        grad_out_input, grad_lse_input, grad_block_out, grad_block_lse = lse_grad(stored_out, stored_lse,
+                                                                                  stored_block_out, stored_block_lse,
+                                                                                  stored_sig, current_dout,
+                                                                                  current_dlse)
+        current_dout = grad_out_input
+        current_dlse = grad_lse_input
+        block_gradients[i] = {'grad_block_out': grad_block_out, 'grad_block_lse': grad_block_lse}
+
+    q, k, v = origin_q, origin_k, origin_v
+
+    for step in range(kv_comm.world_size):
+        _, _, block_out, block_lse, _ = out_lse[step]
+        if block_out.isnan().any() or block_lse.isnan().any():
+            raise
+        block_lse = block_lse.transpose(0, 1).squeeze(2)
+
+        if step + 1 != kv_comm.world_size:
+            next_k, next_v = kv_comm.send_recv_kv(k, v)
+
+        if step == 0:
+            # if step == 0, use the final current_dout
+            block_dout = current_dout
+        else:
+            # else use the grad in the block_gradients
+            block_dout = block_gradients[step]['grad_block_out']
+
+        if block_dout.isnan().any():
+            raise
+
+        if step == 0:
+            backward(
+                block_dout.to(dout.dtype), q, k, v, block_out, block_lse, True, cu_seqlens, max_seqlen, block_seq_len,
+                dq_buffer, dk_buffer, dv_buffer, dropout_p, softmax_scale, alibi_slopes, deterministic, window_size)
+            dq = dq_buffer.to(torch.float32)
+            dk = dk_buffer.to(torch.float32)
+            dv = dv_buffer.to(torch.float32)
+            if dq.isnan().any() or dk.isnan().any() or dv.isnan().any():
+                raise
+        else:
+            if step <= kv_comm.rank:
+                k0 = k[half_index0]
+                v0 = v[half_index0]
+                backward(
+                    block_dout.to(dout.dtype), q, k0, v0, block_out, block_lse, False, cu_seqlens, max_seqlen,
+                    block_seq_len, dq_buffer, dk_buffer, dv_buffer, dropout_p, softmax_scale, alibi_slopes,
+                    deterministic, window_size)
+                dq += dq_buffer
+            else:
+                backward(block_dout[half_index1].to(dout.dtype), q1, k, v, block_out[half_index1],
+                         get_half_lse(block_lse, cu_seqlens,
+                                      front=False), False, cu_seqlens, max_seqlen, block_seq_len, dq_buffer, dk_buffer,
+                         dv_buffer, dropout_p, softmax_scale, alibi_slopes, deterministic, window_size)
+                # only need to add to the tail half, because the head half does not match the causal condition
+                dq[half_index1] += dq_buffer[:block_seq_len]
+
+            d_kv_comm.wait()
+            # dk_comm_buffer, dv_comm_buffer = dk, dv
+            # avoid d_kv_comm.send_recv_kv causing dk_comm_buffer reuse the same memory with next_dk and dk
+            dk_comm_buffer = torch.empty_like(dk)
+            dv_comm_buffer = torch.empty_like(dv)
+            dk_comm_buffer.copy_(dk)
+            dv_comm_buffer.copy_(dv)
+            # next_dk, next_dv comes from a previous gpu, add kv grad to them, and pass them to the next gpu
+            dk, dv = next_dk, next_dv
+
+            if step <= kv_comm.rank:
+                # only need to add to the head part, because the tail part does not match the causal condition
+                dk[half_index0] += dk_buffer[:block_seq_len]
+                dv[half_index0] += dv_buffer[:block_seq_len]
+            else:
+                dk += dk_buffer
+                dv += dv_buffer
+            if dq.isnan().any() or dk.isnan().any() or dv.isnan().any():
+                raise
+        if step + 1 != kv_comm.world_size:
+            kv_comm.wait()
+            k, v = next_k, next_v
+
+        next_dk, next_dv = d_kv_comm.send_recv_kv(dk, dv, dk_comm_buffer, dv_comm_buffer)
+
+    d_kv_comm.wait()
+
+    return dq.to(q.dtype).unsqueeze(0), next_dk.to(q.dtype).unsqueeze(0), next_dv.to(q.dtype).unsqueeze(0)
+
+
+class ZigZagRingFlashAttnVarlenFunc(torch.autograd.Function):
+
+    @staticmethod
+    def forward(
+        ctx,
+        q,
+        k,
+        v,
+        cu_seqlens,
+        max_seqlen,
+        dropout_p,
+        softmax_scale,
+        causal,
+        window_size,
+        alibi_slopes,
+        deterministic,
+        return_softmax,
+        group,
+    ):
+        if softmax_scale is None:
+            softmax_scale = q.shape[-1]**(-0.5)
+
+        assert alibi_slopes is None
+        k = k.contiguous()
+        v = v.contiguous()
+        rp_world_size = dist.get_world_size(group)
+        half_index0 = get_half_index(cu_seqlens // rp_world_size, front=True)
+        half_index1 = get_half_index(cu_seqlens // rp_world_size, front=False)
+        out, softmax_lse = zigzag_ring_flash_attn_varlen_forward(
+            group,
+            q,
+            k,
+            v,
+            cu_seqlens,
+            max_seqlen,
+            half_index0,
+            half_index1,
+            softmax_scale=softmax_scale,
+            dropout_p=dropout_p,
+            causal=causal,
+            window_size=window_size,
+            alibi_slopes=alibi_slopes,
+            deterministic=False,
+        )
+        # this should be out_padded
+        is_half_index_tensor = isinstance(half_index0, torch.Tensor)
+        ctx.is_half_index_tensor = is_half_index_tensor
+        if is_half_index_tensor:
+            """
+            Shapes:
+            qkv: [1, seqlen, num_heads, hidden_size]
+            out: [1, seqlen, num_heads, hidden_size]
+            softmax_lse: [1, num_heads, seqlen]
+            """
+            ctx.save_for_backward(q, k, v, out, softmax_lse, cu_seqlens, half_index0, half_index1)
+        else:
+            ctx.save_for_backward(q, k, v, out, softmax_lse, cu_seqlens)
+            ctx.half_index0 = half_index0
+            ctx.half_index1 = half_index1
+        ctx.max_seqlen = max_seqlen
+        ctx.dropout_p = dropout_p
+        ctx.softmax_scale = softmax_scale
+        ctx.causal = causal
+        ctx.window_size = window_size
+        ctx.alibi_slopes = alibi_slopes
+        ctx.deterministic = deterministic
+        ctx.group = group
+        return out if not return_softmax else (out, softmax_lse, None)
+
+    @staticmethod
+    def backward(ctx, dout, *args):
+        if ctx.is_half_index_tensor:
+            (q, k, v, out, softmax_lse, cu_seqlens, half_index0, half_index1) = (ctx.saved_tensors)
+        else:
+            q, k, v, out, softmax_lse, cu_seqlens = ctx.saved_tensors
+            half_index0 = ctx.half_index0
+            half_index1 = ctx.half_index1
+        dq, dk, dv = zigzag_ring_flash_attn_varlen_backward(
+            ctx.group,
+            dout,
+            q,
+            k,
+            v,
+            out,
+            softmax_lse,
+            cu_seqlens,
+            ctx.max_seqlen,
+            half_index0,
+            half_index1,
+            softmax_scale=ctx.softmax_scale,
+            dropout_p=ctx.dropout_p,
+            causal=ctx.causal,
+            window_size=ctx.window_size,
+            alibi_slopes=ctx.alibi_slopes,
+            deterministic=ctx.deterministic,
+        )
+        return dq, dk, dv, None, None, None, None, None, None, None, None, None, None
+
+
+def zigzag_ring_flash_attn_varlen_func(
+    q,
+    k,
+    v,
+    cu_seqlens,
+    max_seqlen,
+    dropout_p=0.0,
+    softmax_scale=None,
+    causal=False,
+    window_size=(-1, -1),  # -1 means infinite context window
+    alibi_slopes=None,
+    deterministic=False,
+    return_attn_probs=False,
+    group=None,
+):
+    return ZigZagRingFlashAttnVarlenFunc.apply(
+        q,
+        k,
+        v,
+        cu_seqlens,
+        max_seqlen,
+        dropout_p,
+        softmax_scale,
+        causal,
+        window_size,
+        alibi_slopes,
+        deterministic,
+        return_attn_probs,
+        group,
+    )

swift/sequence_parallel/zigzag_ring_attn_npu.py

@@ -0,0 +1,459 @@
+# Some code borrowed from the awesome work: https://github.com/zhuzilin/ring-flash-attention
+# Copyright (c) ModelScope Contributors. All rights reserved.
+import torch
+from functools import cache
+
+from .ring_utils import RingComm
+
+_NPU_BLOCK_MASK_SIZE = 2048
+_NPU_FULL_TOKENS = 2147483647
+_NPU_TND_SOFTMAX_STAT_REPEAT = 8
+
+
+def is_npu_tensor(tensor: torch.Tensor) -> bool:
+    return tensor.device.type == 'npu'
+
+
+def _cu_seqlens_to_actual_seq(cu_seqlens: torch.Tensor) -> tuple[int, ...]:
+    return tuple(int(x) for x in cu_seqlens[1:].detach().cpu().tolist())
+
+
+@cache
+def _get_npu_causal_mask_cpu() -> torch.Tensor:
+    return torch.triu(torch.ones((_NPU_BLOCK_MASK_SIZE, _NPU_BLOCK_MASK_SIZE), dtype=torch.bool), diagonal=1)
+
+
+def _get_npu_causal_mask(device: torch.device) -> torch.Tensor:
+    return _get_npu_causal_mask_cpu().to(device=device)
+
+
+def _normalize_window_size(window_size):
+    if window_size is None:
+        return -1, -1
+    return window_size
+
+
+def _get_npu_sparse_params(causal: bool, window_size, device: torch.device) -> dict:
+    window_size = _normalize_window_size(window_size)
+    if window_size != (-1, -1):
+        left, right = window_size
+        left = _NPU_FULL_TOKENS if left < 0 else int(left)
+        right = _NPU_FULL_TOKENS if right < 0 else int(right)
+        if causal:
+            right = 0
+        return {
+            'atten_mask': _get_npu_causal_mask(device),
+            'sparse_mode': 4,
+            'pre_tockens': left,
+            'next_tockens': right,
+        }
+    if causal:
+        return {
+            'atten_mask': _get_npu_causal_mask(device),
+            'sparse_mode': 3,
+            'pre_tockens': _NPU_FULL_TOKENS,
+            'next_tockens': _NPU_FULL_TOKENS,
+        }
+    return {
+        'atten_mask': None,
+        'sparse_mode': 0,
+        'pre_tockens': _NPU_FULL_TOKENS,
+        'next_tockens': _NPU_FULL_TOKENS,
+    }
+
+
+def _reshape_npu_lse(lse: torch.Tensor, seqlen_q: int, num_heads: int) -> torch.Tensor:
+    """Normalize Ascend softmax stats to flash-attn's [num_heads, seqlen] layout."""
+    if lse.dim() == 2:
+        if lse.shape == (num_heads, seqlen_q):
+            return lse.contiguous()
+        if lse.shape == (seqlen_q, num_heads):
+            return lse.transpose(0, 1).contiguous()
+    elif lse.dim() == 3:
+        # Some CANN versions return an extra trailing size-8 axis with repeated
+        # stats. Ring merge only needs one copy of each token/head lse.
+        if lse.shape[-1] == 8:
+            lse = lse[..., 0]
+            if lse.shape == (seqlen_q, num_heads):
+                return lse.transpose(0, 1).contiguous()
+            if lse.shape == (num_heads, seqlen_q):
+                return lse.contiguous()
+        if lse.shape[0] == seqlen_q:
+            return lse.permute(1, 2, 0).reshape(num_heads, seqlen_q).contiguous()
+        if lse.shape[1] == seqlen_q:
+            return lse.permute(0, 2, 1).reshape(num_heads, seqlen_q).contiguous()
+    raise RuntimeError(f'Unexpected NPU lse shape {tuple(lse.shape)} for seqlen_q={seqlen_q}, num_heads={num_heads}')
+
+
+def _get_npu_attention_common_kwargs(
+    q: torch.Tensor,
+    *,
+    cu_seqlens_q: torch.Tensor,
+    cu_seqlens_kv: torch.Tensor,
+    softmax_scale: float,
+    dropout_p: float,
+    causal: bool,
+    window_size,
+    deterministic: bool,
+) -> dict:
+    sparse_params = _get_npu_sparse_params(causal, window_size, q.device)
+    return {
+        'head_num': q.shape[1],
+        'input_layout': 'TND',
+        'scale_value': softmax_scale or q.shape[-1]**(-0.5),
+        'keep_prob': 1. - dropout_p,
+        'actual_seq_qlen': _cu_seqlens_to_actual_seq(cu_seqlens_q),
+        'actual_seq_kvlen': _cu_seqlens_to_actual_seq(cu_seqlens_kv),
+        'sync': bool(deterministic and dropout_p > 0),
+        **sparse_params,
+    }
+
+
+def _call_npu_fusion_attention(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    *,
+    cu_seqlens_q: torch.Tensor,
+    cu_seqlens_kv: torch.Tensor,
+    softmax_scale: float,
+    dropout_p: float,
+    causal: bool,
+    window_size,
+    deterministic: bool,
+):
+    import torch_npu
+
+    common_kwargs = _get_npu_attention_common_kwargs(
+        q,
+        cu_seqlens_q=cu_seqlens_q,
+        cu_seqlens_kv=cu_seqlens_kv,
+        softmax_scale=softmax_scale,
+        dropout_p=dropout_p,
+        causal=causal,
+        window_size=window_size,
+        deterministic=deterministic,
+    )
+    params = {
+        'query': q,
+        'key': k,
+        'value': v,
+        'scale': common_kwargs['scale_value'],
+        'softmax_layout': 'TND',
+    }
+    params.update(common_kwargs)
+    params.pop('scale_value')
+    return torch_npu.npu_fusion_attention(**params)
+
+
+def _call_npu_fusion_attention_grad(
+    dout: torch.Tensor,
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    *,
+    attention_out: torch.Tensor,
+    softmax_lse: torch.Tensor,
+    cu_seqlens_q: torch.Tensor,
+    cu_seqlens_kv: torch.Tensor,
+    softmax_scale: float,
+    dropout_p: float,
+    causal: bool,
+    window_size,
+    deterministic: bool,
+):
+    import torch_npu
+
+    if not hasattr(torch_npu, 'npu_fusion_attention_grad'):
+        raise AttributeError('torch_npu.npu_fusion_attention_grad is not available')
+    # Dropout backward needs the exact seed/offset from the original forward,
+    # which this ring ctx does not save. Fail instead of using a wrong mask.
+    if dropout_p != 0.0:
+        raise NotImplementedError('NPU ring attention native backward currently requires dropout_p=0.')
+
+    common_kwargs = _get_npu_attention_common_kwargs(
+        q,
+        cu_seqlens_q=cu_seqlens_q,
+        cu_seqlens_kv=cu_seqlens_kv,
+        softmax_scale=softmax_scale,
+        dropout_p=dropout_p,
+        causal=causal,
+        window_size=window_size,
+        deterministic=deterministic,
+    )
+    softmax_max, softmax_sum = _npu_softmax_stats_from_global_lse(
+        softmax_lse,
+        q_tokens=q.shape[0],
+        num_heads=q.shape[1],
+    )
+
+    params = {
+        'query': q,
+        'key': k,
+        'value': v,
+        'dy': dout,
+        'head_num': common_kwargs['head_num'],
+        'input_layout': common_kwargs['input_layout'],
+        'atten_mask': common_kwargs['atten_mask'],
+        'softmax_max': softmax_max,
+        'softmax_sum': softmax_sum,
+        'softmax_in': None,
+        'attention_in': (attention_out if torch.is_tensor(attention_out) and attention_out.numel() > 0 else None),
+        'scale_value': common_kwargs['scale_value'],
+        'keep_prob': common_kwargs['keep_prob'],
+        'pre_tockens': common_kwargs['pre_tockens'],
+        'next_tockens': common_kwargs['next_tockens'],
+        'seed': 0,
+        'offset': 0,
+        'numels': 0,
+        'actual_seq_qlen': common_kwargs['actual_seq_qlen'],
+        'actual_seq_kvlen': common_kwargs['actual_seq_kvlen'],
+        'sparse_mode': common_kwargs['sparse_mode'],
+        'sync': common_kwargs['sync'],
+        'softmax_layout': 'TND',
+    }
+    return torch_npu.npu_fusion_attention_grad(**params)
+
+
+def _normalize_flash_attn_lse(softmax_lse: torch.Tensor, total_len: int) -> torch.Tensor:
+    """Normalize flash-attn lse to [num_heads, total_len]."""
+    lse = softmax_lse
+    if lse.dim() == 3 and lse.shape[0] == 1:
+        lse = lse.squeeze(0)
+    if lse.dim() != 2:
+        raise RuntimeError(f'Unexpected softmax_lse shape: {tuple(softmax_lse.shape)}')
+    if lse.shape[1] != total_len:
+        lse = lse.transpose(0, 1).contiguous()
+    if lse.shape[1] != total_len:
+        raise RuntimeError(f'Unexpected softmax_lse shape: {tuple(softmax_lse.shape)} for total_len={total_len}')
+    return lse
+
+
+def _npu_softmax_stats_from_global_lse(
+    softmax_lse_global: torch.Tensor,
+    q_tokens: int,
+    num_heads: int,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    lse_h_t = _normalize_flash_attn_lse(softmax_lse_global, q_tokens)
+    if lse_h_t.shape[0] != num_heads:
+        raise RuntimeError(f'Unexpected global lse shape: {tuple(softmax_lse_global.shape)} '
+                           f'for q_tokens={q_tokens}, num_heads={num_heads}')
+
+    # With softmax_layout='TND', Ascend returns softmax stats as [T, N, 8].
+    # The split-attention backward only needs logsumexp; max=lse and sum=1
+    # encode the same value without replaying the block forward.
+    lse_t_h = lse_h_t.transpose(0, 1).contiguous().to(torch.float32)
+    softmax_max = lse_t_h.unsqueeze(-1).expand(
+        q_tokens,
+        num_heads,
+        _NPU_TND_SOFTMAX_STAT_REPEAT,
+    ).contiguous()
+    return softmax_max, torch.ones_like(softmax_max)
+
+
+def _get_second_half_lse(softmax_lse: torch.Tensor, cu_seqlens: torch.Tensor) -> torch.Tensor:
+    total_len = int(cu_seqlens[-1].item())
+    lse = _normalize_flash_attn_lse(softmax_lse, total_len)
+
+    # The step > rank branch only differentiates q[half_index1]. Slice the final
+    # merged lse per sequence so the native grad ctx sees the same query span.
+    second_half_lse = torch.empty((lse.shape[0], lse.shape[1] // 2), dtype=lse.dtype, device=lse.device)
+    for i in range(len(cu_seqlens) - 1):
+        start, end = cu_seqlens[i].item(), cu_seqlens[i + 1].item()
+        new_start, new_end = start // 2, end // 2
+        start += (end - start) // 2
+        second_half_lse[:, new_start:new_end] = lse[:, start:end]
+    return second_half_lse
+
+
+def _npu_block_backward_with_global_stats(
+    block_dout,
+    block_q,
+    block_k,
+    block_v,
+    block_out_global,
+    block_lse_global,
+    block_causal,
+    cu_seqlens_q,
+    cu_seqlens_kv,
+    softmax_scale,
+    dropout_p,
+    window_size,
+    deterministic,
+):
+    """Run one native NPU block backward using the final merged ring stats."""
+    return _call_npu_fusion_attention_grad(
+        block_dout.to(block_q.dtype),
+        block_q,
+        block_k,
+        block_v,
+        attention_out=block_out_global,
+        softmax_lse=block_lse_global,
+        cu_seqlens_q=cu_seqlens_q,
+        cu_seqlens_kv=cu_seqlens_kv,
+        softmax_scale=softmax_scale,
+        dropout_p=dropout_p,
+        causal=block_causal,
+        window_size=window_size,
+        deterministic=deterministic,
+    )[:3]
+
+
+def _squeeze_batch(*tensors):
+    squeezed = []
+    for tensor in tensors:
+        if tensor.shape[0] == 1:
+            squeezed.append(tensor.squeeze(0))
+        else:
+            squeezed.append(tensor)
+    return tuple(squeezed)
+
+
+def npu_backward(
+        process_group,
+        dout,
+        q,
+        k,
+        v,
+        out,
+        softmax_lse,
+        cu_seqlens,
+        max_seqlen,
+        half_index0,
+        half_index1,
+        softmax_scale,
+        dropout_p=0.0,
+        window_size=(-1, -1),
+        deterministic=False,
+):
+    kv_comm = RingComm(process_group)
+    d_kv_comm = RingComm(process_group)
+    dout, q, k, v, out, softmax_lse = _squeeze_batch(dout, q, k, v, out, softmax_lse)
+    cu_seqlens = cu_seqlens // kv_comm.world_size
+    del max_seqlen
+
+    half_cu_seqlens = cu_seqlens // 2
+    q1 = q[half_index1]
+    dout1 = dout[half_index1]
+    out1 = out[half_index1]
+    softmax_lse1 = _get_second_half_lse(softmax_lse, cu_seqlens)
+
+    dq = torch.zeros_like(q, dtype=torch.float32)
+    current_step_dk = torch.empty_like(k, dtype=torch.float32)
+    current_step_dv = torch.empty_like(v, dtype=torch.float32)
+    next_dk = next_dv = None
+
+    for step in range(kv_comm.world_size):
+        current_step_dk.zero_()
+        current_step_dv.zero_()
+        if step == 0:
+            bdq, bdk, bdv = _npu_block_backward_with_global_stats(
+                dout,
+                q,
+                k,
+                v,
+                out,
+                softmax_lse,
+                True,
+                cu_seqlens,
+                cu_seqlens,
+                softmax_scale,
+                dropout_p,
+                window_size,
+                deterministic,
+            )
+            dq += bdq.to(torch.float32)
+            current_step_dk += bdk.to(torch.float32)
+            current_step_dv += bdv.to(torch.float32)
+        elif step <= kv_comm.rank:
+            k0 = k[half_index0]
+            v0 = v[half_index0]
+            bdq, bdk, bdv = _npu_block_backward_with_global_stats(
+                dout,
+                q,
+                k0,
+                v0,
+                out,
+                softmax_lse,
+                False,
+                cu_seqlens,
+                half_cu_seqlens,
+                softmax_scale,
+                dropout_p,
+                window_size,
+                deterministic,
+            )
+            dq += bdq.to(torch.float32)
+            current_step_dk[half_index0] += bdk.to(torch.float32)
+            current_step_dv[half_index0] += bdv.to(torch.float32)
+        else:
+            bdq, bdk, bdv = _npu_block_backward_with_global_stats(
+                dout1,
+                q1,
+                k,
+                v,
+                out1,
+                softmax_lse1,
+                False,
+                half_cu_seqlens,
+                cu_seqlens,
+                softmax_scale,
+                dropout_p,
+                window_size,
+                deterministic,
+            )
+            dq[half_index1] += bdq.to(torch.float32)
+            current_step_dk += bdk.to(torch.float32)
+            current_step_dv += bdv.to(torch.float32)
+
+        # K/V gradients are owned by the rank that originally held that shard.
+        # Rotate the accumulated gradients in the opposite ring direction until
+        # each owner receives its final dk/dv.
+        if step == 0:
+            dk = current_step_dk
+            dv = current_step_dv
+        else:
+            dk = next_dk
+            dv = next_dv
+            dk += current_step_dk
+            dv += current_step_dv
+
+        next_dk, next_dv = d_kv_comm.send_recv_kv(dk, dv)
+        d_kv_comm.wait()
+
+        if step + 1 != kv_comm.world_size:
+            next_k, next_v = kv_comm.send_recv_kv(k, v)
+            kv_comm.wait()
+            k, v = next_k, next_v
+
+    return dq.to(q.dtype).unsqueeze(0), next_dk.to(q.dtype).unsqueeze(0), next_dv.to(q.dtype).unsqueeze(0)
+
+
+def npu_forward(
+    q,
+    k,
+    v,
+    causal,
+    cu_seqlens_q,
+    cu_seqlens_kv,
+    dropout_p,
+    softmax_scale,
+    deterministic,
+    window_size,
+):
+    outputs = _call_npu_fusion_attention(
+        q,
+        k,
+        v,
+        cu_seqlens_q=cu_seqlens_q,
+        cu_seqlens_kv=cu_seqlens_kv,
+        softmax_scale=softmax_scale,
+        dropout_p=dropout_p,
+        causal=causal,
+        window_size=window_size,
+        deterministic=deterministic,
+    )
+    block_out, softmax_max, softmax_sum = outputs[:3]
+    block_lse = softmax_max.to(torch.float32) + torch.log(softmax_sum.to(torch.float32))
+    block_lse = _reshape_npu_lse(block_lse, q.shape[0], q.shape[1])
+    return block_out, block_lse