correct ulysses-opt text layout and collectives#1239
Conversation
Preserve the full stored text tail for three-entry cumulative sequence lengths. Round Triton row tiles up to powers of two and reject non-positive values. Batch head-parallel text all-gather after image all-to-all completion while preserving overlap. Remove leftover profiler regions and the redundant apply wrapper.
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request optimizes the Ulysses attention implementation by removing profiling overhead, introducing a txt_storage_len property to handle masked text sequences, and batching the all_gather operation for text heads outside the loop. The review feedback suggests further optimizing the post-gather processing by stacking, permuting, and reshaping the gathered text heads in a single operation outside the loop to avoid multiple torch.cat operations and reduce GPU kernel launch 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.
| local_txt_heads = torch.stack([record[0] for record in records], dim=0) | ||
| if local_txt_heads.shape[1] == 0: | ||
| gathered_txt_heads = None | ||
| else: | ||
| gathered_txt_heads = [torch.empty_like(local_txt_heads) for _ in range(ctx.world_size)] | ||
| dist.all_gather(gathered_txt_heads, local_txt_heads, group=ctx.seq_p_group) | ||
|
|
||
| head_outputs = [] | ||
| for record in records: | ||
| for local_head, record in enumerate(records): | ||
| if ctx.use_fp8_comm: | ||
| txt_attn, _payload, _scale, a2a_payload, a2a_scale, attn_shape, attn_scale_shape, _payload_work, _scale_work = record | ||
| _txt_attn, _payload, _scale, a2a_payload, a2a_scale, attn_shape, attn_scale_shape, _payload_work, _scale_work = record | ||
| else: | ||
| txt_attn, _input_buf, a2a_output, _unused0, _unused1, _work = record | ||
| with record_function("ulysses_opt/head_text_gather"): | ||
| gathered_txt_attn = [torch.empty_like(txt_attn) for _ in range(ctx.world_size)] | ||
| dist.all_gather(gathered_txt_attn, txt_attn, group=ctx.seq_p_group) | ||
| txt_attn = torch.cat(gathered_txt_attn, dim=1) | ||
| 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) | ||
| _txt_attn, _input_buf, a2a_output, _unused0, _unused1, _work = record | ||
| if gathered_txt_heads is None: | ||
| txt_attn = local_txt_heads.new_empty((0, ctx.world_size * ctx.hidden_dims)) | ||
| else: | ||
| txt_attn = torch.cat([rank_txt[local_head] for rank_txt in gathered_txt_heads], dim=1) |
There was a problem hiding this comment.
Instead of performing shard_heads separate torch.cat operations inside the loop (which launches multiple GPU kernels and incurs overhead), we can stack, permute, and reshape the gathered text heads in a single highly-optimized operation outside the loop. This significantly improves performance by reducing GPU kernel launch overhead.
local_txt_heads = torch.stack([record[0] for record in records], dim=0)
if local_txt_heads.shape[1] == 0:
txt_attn_all = None
else:
gathered_txt_heads = [torch.empty_like(local_txt_heads) for _ in range(ctx.world_size)]
dist.all_gather(gathered_txt_heads, local_txt_heads, group=ctx.seq_p_group)
stacked = torch.stack(gathered_txt_heads, dim=0)
txt_attn_all = stacked.permute(1, 2, 0, 3).reshape(ctx.shard_heads, -1, ctx.world_size * ctx.hidden_dims)
head_outputs = []
for local_head, record in enumerate(records):
if ctx.use_fp8_comm:
_txt_attn, _payload, _scale, a2a_payload, a2a_scale, attn_shape, attn_scale_shape, _payload_work, _scale_work = record
else:
_txt_attn, _input_buf, a2a_output, _unused0, _unused1, _work = record
if txt_attn_all is None:
txt_attn = local_txt_heads.new_empty((0, ctx.world_size * ctx.hidden_dims))
else:
txt_attn = txt_attn_all[local_head]Materialize gathered text heads in one stack-permute-reshape pass and reuse per-head views, removing per-head cat launches while releasing intermediate buffers early.
2 participants
ulysses-opt usage:
"parallel": {
"seq_p_tensor_fusion": true,
"seq_p_attn_type": "ulysses-opt",
}