content: 
  ### blockwise attn
  
  LSE (Log-Sum-Exp) 更新是 FlashAttention 的核心
  
  https://github.com/zhuzilin/ring-flash-attention/pull/34#issuecomment-2076126795
  
  @torch.jit.script
  def _update_out_and_lse(
      out: torch.Tensor,
      lse: torch.Tensor,
      block_out: torch.Tensor,
      block_lse: torch.Tensor,
  ) -> Tuple[torch.Tensor, torch.Tensor]:
  
      block_out = block_out.to(torch.float32)
      block_lse = block_lse.transpose(-2, -1).unsqueeze(dim=-1)
  
      # 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
      out = out - F.sigmoid(block_lse - lse) * (out - block_out)
      lse = lse - F.logsigmoid(lse - block_lse)
  
      return out, lse
  
  ### ring attn
  
  def ring_flash_attn_forward
  	comm = RingComm(process_group)
  
      out = None
      lse = None
  
      next_k, next_v = None, None
  
      for step in range(comm.world_size):
          if step + 1 != comm.world_size:
              next_k, next_v = comm.send_recv_kv(k, v)
          if not causal or step <= comm.rank:
          	...
          	"causal": causal and step == 0,
          	fwd
          	...
          	out, lse = update_out_and_lse(out, lse, block_out, block_lse)
  
          if step + 1 != comm.world_size:
              comm.wait()
              k, v = next_k, next_v