Add optimized Ulysses attention path

[STwangyingrui]

No description provided.

[gemini-code-assist]

Code Review

This pull request introduces an optimized Ulysses attention implementation (ulysses-opt) featuring fused pre- and post-communication stages and custom Triton layout kernels with FP8 support. The review feedback highlights two important improvements: first, explicitly creating the cu_seqlens tensors on the GPU to avoid runtime device mismatch errors during attention execution; second, optimizing the head-parallel path by consolidating multiple dist.all_gather calls into a single collective operation to minimize communication and synchronization overhead.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[gemini-code-assist]

The cu_seqlens tensors in _build_attention_cu_seqlens and _build_img_only_cu_seqlens are created on CPU by default because no device is specified. This will cause a runtime device mismatch error when calling the attention module (e.g., FlashAttention or SageAttention), which expects these tensors to be on the GPU. Since Ulysses attention is strictly for CUDA execution, these tensors should be explicitly created on "cuda".

Suggested change

	@staticmethod
	def _build_attention_cu_seqlens(txt_len, txt_mask_len, global_img_len):
	cu_seqlens = torch.zeros([2], dtype=torch.int32)
	cu_seqlens[1] = txt_len + global_img_len
	if txt_mask_len is not None:
	cu_seqlens = torch.cat((cu_seqlens, torch.tensor([txt_mask_len + global_img_len], dtype=torch.int32)))
	return cu_seqlens
	@staticmethod
	def _build_img_only_cu_seqlens(global_img_len):
	cu_seqlens = torch.zeros([2], dtype=torch.int32)
	cu_seqlens[1] = global_img_len
	return cu_seqlens
	@staticmethod
	def _build_attention_cu_seqlens(txt_len, txt_mask_len, global_img_len):
	cu_seqlens = torch.zeros([2], dtype=torch.int32, device="cuda")
	cu_seqlens[1] = txt_len + global_img_len
	if txt_mask_len is not None:
	cu_seqlens = torch.cat((cu_seqlens, torch.tensor([txt_mask_len + global_img_len], dtype=torch.int32, device="cuda")))
	return cu_seqlens
	@staticmethod
	def _build_img_only_cu_seqlens(global_img_len):
	cu_seqlens = torch.zeros([2], dtype=torch.int32, device="cuda")
	cu_seqlens[1] = global_img_len
	return cu_seqlens

[gemini-code-assist]

In _finish_head_parallel_attn_a2a, performing dist.all_gather inside the loop over records (which corresponds to shard_heads) introduces significant launch and synchronization overhead for each head. Since all txt_attn slices are available before the loop, we can concatenate them and perform a single all_gather collective operation, then slice and reshape the gathered tensor for each head. This significantly reduces collective communication overhead and improves performance.

        # Perform a single all_gather for all heads to avoid launch overhead of multiple collectives
        with record_function("ulysses_opt/head_text_gather"):
            all_txt_attn = torch.cat([record[0] for record in records], dim=1)
            gathered_all_txt_attn = [torch.empty_like(all_txt_attn) for _ in range(ctx.world_size)]
            dist.all_gather(gathered_all_txt_attn, all_txt_attn, group=ctx.seq_p_group)
            gathered_tensor = torch.stack(gathered_all_txt_attn, dim=1).reshape(all_txt_attn.shape[0], ctx.world_size, ctx.shard_heads, ctx.hidden_dims)
            gathered_tensor = gathered_tensor.permute(0, 2, 1, 3)

        head_outputs = []
        for local_head, record in enumerate(records):
            if ctx.use_fp8_comm:
                _, _payload, _scale, a2a_payload, a2a_scale, attn_shape, attn_scale_shape, _payload_work, _scale_work = record
            else:
                _, _input_buf, a2a_output, _unused0, _unused1, _work = record
            txt_attn = gathered_tensor[:, local_head].reshape(gathered_tensor.shape[0], ctx.world_size * ctx.hidden_dims)
            with record_function("ulysses_opt/head_attn_post"):
                if ctx.use_fp8_comm:
                    head_out = attn_post_fp8(a2a_payload, a2a_scale, attn_shape, attn_scale_shape, txt_attn, ctx.img_first)
                else:
                    head_out = attn_post(a2a_output, txt_attn, ctx.img_first)
            head_outputs.append(head_out.reshape(head_out.shape[0], ctx.world_size, ctx.hidden_dims))