[Perf] diffusion: Keep same-device copies out of the pipelined Ulysses exchange

python/sglang/kernels/ops/diffusion/__init__.py

@@ -543,6 +543,12 @@
         "layout.ulysses_qkv_triton:pack_qkv_destination_major",
         _CUDA,
     ),
+    (
+        "diffusion.copy_head_groups",
+        KernelBackend.TRITON,
+        "layout.head_group_copy_triton:copy_head_groups",
+        _CUDA,
+    ),
     (
         "diffusion.varlen_pack_qkv",
         KernelBackend.TRITON,
@@ -793,6 +799,7 @@
     "fused_causal_conv3d_cat_pad_cuda": "sglang.kernels.kda_kernels.causal_conv3d_cat_pad_jit",
     "fused_causal_conv3d_cat_pad": "layout.causal_conv3d_cat_pad_triton",
     "pack_qkv_destination_major": "layout.ulysses_qkv_triton",
+    "copy_head_groups": "layout.head_group_copy_triton",
     "fused_qknorm_rope_out_of_place": "rope.qknorm_rope_jit",
     "vdn_delta_factors": "attention.vdn_delta_factors_jit",
     "can_use_vdn_delta_factors": "attention.vdn_delta_factors_jit",

python/sglang/kernels/ops/diffusion/layout/head_group_copy_triton.py

@@ -0,0 +1,90 @@
+# SPDX-License-Identifier: Apache-2.0
+
+import torch
+import triton
+import triton.language as tl
+
+_MAX_INT32 = 2**31 - 1
+
+
+@triton.jit(do_not_specialize=["rows"])
+def _copy_head_groups_kernel(
+    src_ptr,
+    dst_ptr,
+    rows,
+    batch,
+    width,
+    stride_src_g,
+    stride_src_b,
+    stride_src_s,
+    stride_dst_g,
+    stride_dst_b,
+    stride_dst_s,
+    IDX64: tl.constexpr,
+    BLOCK_ROWS: tl.constexpr,
+    BLOCK_WIDTH: tl.constexpr,
+):
+    group_batch = tl.program_id(1)
+    g = group_batch // batch
+    b = group_batch % batch
+    row = tl.program_id(0) * BLOCK_ROWS + tl.arange(0, BLOCK_ROWS)
+    col = tl.arange(0, BLOCK_WIDTH)
+    if IDX64:
+        g = g.to(tl.int64)
+        b = b.to(tl.int64)
+        row = row.to(tl.int64)
+    mask = (row[:, None] < rows) & (col[None, :] < width)
+    src = (
+        src_ptr
+        + g * stride_src_g
+        + b * stride_src_b
+        + row[:, None] * stride_src_s
+        + col[None, :]
+    )
+    dst = (
+        dst_ptr
+        + g * stride_dst_g
+        + b * stride_dst_b
+        + row[:, None] * stride_dst_s
+        + col[None, :]
+    )
+    tl.store(dst, tl.load(src, mask=mask), mask=mask)
+
+
+def copy_head_groups(src: torch.Tensor, dst: torch.Tensor) -> None:
+    """``dst.copy_(src)`` for ``[group, batch, rows, width]`` views whose last dim
+    is contiguous and whose other strides are arbitrary: one pass that moves the
+    head groups of a sequence shard between per-group blocks and a token-major
+    tensor (a same-device copy per block costs a launch and runs at a fraction
+    of the bandwidth)."""
+    if src.dim() != 4 or src.shape != dst.shape:
+        raise ValueError(
+            f"src {tuple(src.shape)} and dst {tuple(dst.shape)} must be the same 4D shape"
+        )
+    if src.stride(-1) != 1 or dst.stride(-1) != 1:
+        raise ValueError("src and dst must be contiguous in the last dim")
+    if src.dtype != dst.dtype or src.device != dst.device:
+        raise ValueError("src and dst must share dtype and device")
+    if src.numel() == 0:
+        return
+    groups, batch, rows, width = src.shape
+    bound = lambda t: sum((n - 1) * s for n, s in zip(t.shape, t.stride())) + 1
+    block_width = triton.next_power_of_2(width)
+    block_rows = max(1, 8192 // block_width)
+    with torch.get_device_module().device(src.device):
+        _copy_head_groups_kernel[(triton.cdiv(rows, block_rows), groups * batch)](
+            src,
+            dst,
+            rows,
+            batch,
+            width,
+            *src.stride()[:3],
+            *dst.stride()[:3],
+            IDX64=max(bound(src), bound(dst)) > _MAX_INT32,
+            BLOCK_ROWS=block_rows,
+            BLOCK_WIDTH=block_width,
+            num_warps=8,
+        )
+
+
+__all__ = ["copy_head_groups"]

python/sglang/multimodal_gen/runtime/distributed/device_communicators/ipc_a2a_multi.py

@@ -27,6 +27,7 @@
 import torch.distributed as dist
 
 from sglang.kernels.ops.communication.ipc_a2a import load_ipc_a2a_sync
+from sglang.kernels.ops.diffusion import copy_head_groups
 from sglang.multimodal_gen import envs
 from sglang.multimodal_gen.runtime.distributed.device_communicators.ipc_a2a import (
     _Unsupported,
@@ -491,6 +492,7 @@ def pipelined_attention(
         signature=None,
         replicated=(None, None, None),
         replicated_first: bool = True,
+        attend_into=None,
     ):
         """Ulysses attention with the head exchange pipelined against it.
 
@@ -506,11 +508,12 @@ def pipelined_attention(
         ``attend(q, k, v)`` runs attention on one head group, ``[batch, rows,
         heads / (world * groups), head_dim]`` per part in the sequence order
         above, and must treat heads independently, which keeps the result
-        bit-identical to the sequential exchange. Group g's copies run on the
-        copy engine while group g-1 attends, and each group attends on its own
-        stream and sends its output from there, so no stream waits on a kernel
-        still running on another. Every copy, the merge into the output
-        included, is a pitched copy-engine copy, so the exchange takes no SMs.
+        bit-identical to the sequential exchange. Group g's peer copies run on
+        the copy engine while group g-1 attends, and each group attends on its
+        own stream and sends its output from there, so no stream waits on a
+        kernel still running on another. A same-device copy runs on the SMs and
+        would wait out a running attention, so this rank's own rows of every
+        group go into its slots in one pass before the first group attends.
         Returns the output rows of this rank, or None when this call cannot
         pipeline. ``groups < 0`` picks the count and steps aside, on every
         rank, when the buffers would not fit.
@@ -522,6 +525,13 @@ def pipelined_attention(
         go straight into its own receive slot, and v always moves straight
         from ``v``. It is only called once the call is known to pipeline.
 
+        ``attend_into(q, k, v, out)``, when given, writes what ``attend``
+        returns into ``out`` (any token stride); with ``out=None`` it only says
+        whether it serves those inputs. Each group then attends straight into a
+        token-major buffer of every row and head whose rows of this rank are the
+        output, so only the peers' rows move: out on the copy engine, back in
+        one pass.
+
         ``sequential()``, when given, returns what the sequential exchange
         gives for this call. The first call per plan and ``signature`` (the
         caller's attention setup) runs both and keeps the pipeline for them
@@ -658,7 +668,33 @@ def place_replicated(stream, part, g):
                     n,
                 )
 
+        def own_rows_views(part):
+            # my shard of one part as [batch, s_local, group, group heads * head_dim],
+            # and my own slots of it as [group, batch, s_local, ...]
+            gh = plan.group_heads(part)
+            width = gh * head_dim
+            head0 = start(r, 0, part)
+            shard = parts[part][:, :, head0 : head0 + groups * gh]
+            inb = pool.inb[r]
+            dst = inb.as_strided(
+                (groups, batch, s_local, width),
+                (slots.group_numel, slots.rows[part] * width, width, 1),
+                inb.storage_offset()
+                + slots.part_offset[part]
+                + my_shard_row[part] * width,
+            )
+            return shard.unflatten(2, (groups, gh)).flatten(3).permute(2, 0, 1, 3), dst
+
         ready = main.record_event()
+        # my own rows of every group, here and in one pass per part: a same-device
+        # copy runs on the SMs, so on the copy stream it would wait out the
+        # previous group's attention and hold up the peer sends queued behind it
+        for part in (2,) if fill is not None else range(3):
+            copy_head_groups(*own_rows_views(part))
+        for g in range(groups):
+            for part in range(3):
+                place_replicated(main, part, g)
+        placed = main.record_event()
         own, sent = [], []
         if fill is not None:
             # group by group: the peers' blocks first, so their copies start
@@ -674,19 +710,14 @@ def place_replicated(stream, part, g):
                 sent.append(main.record_event())
                 fill(start(r, g, 0), hg, own_rows(0, g), own_rows(1, g))
                 own.append(main.record_event())
-        own_cin = []
         cin.wait_event(ready)
         with torch.cuda.stream(cin):
             for g in range(groups):
-                for p in (r, *peers):
+                for p in peers:
                     send_rows(cin, 2, p, g)
                     if fill is None:
                         send_rows(cin, 0, p, g)
                         send_rows(cin, 1, p, g)
-                    if p == r:
-                        for part in range(3):
-                            place_replicated(cin, part, g)
-                        own_cin.append(cin.record_event())
                 if fill is not None:
                     cin.wait_event(sent[g])
                     chunk = s_local * row_bytes[0]
@@ -709,14 +740,32 @@ def place_replicated(stream, part, g):
         # heads. Each group stream writes its blocks as soon as they exist, so
         # the merge overlaps the groups still attending instead of trailing them.
         hq = plan.group_heads(0)
-        merged = torch.empty(
-            batch,
-            plan.out_rows,
-            plan.heads[0],
-            head_dim,
-            dtype=q.dtype,
-            device=q.device,
+        # zero-copy output: every row of every head, this rank's rows returned
+        into = (
+            attend_into is not None
+            and not squeeze
+            and not any(plan.replicated)
+            and attend_into(*(slots.slot(pool, r, 0, part) for part in range(3)), None)
         )
+        if into:
+            full = torch.empty(
+                batch,
+                world * s_local,
+                plan.heads[0],
+                head_dim,
+                dtype=q.dtype,
+                device=q.device,
+            )
+            merged = full[:, my_shard_row[0] : my_shard_row[0] + s_local]
+        else:
+            merged = torch.empty(
+                batch,
+                plan.out_rows,
+                plan.heads[0],
+                head_dim,
+                dtype=q.dtype,
+                device=q.device,
+            )
         merged_ptr = merged.data_ptr()
         merged_row = plan.heads[0] * head_bytes
         out_row = row_bytes[0]
@@ -726,13 +775,31 @@ def place_replicated(stream, part, g):
         done = []
         for g in range(groups):
             stream = self.pipe_group_streams[g]
-            if fill is not None:
-                stream.wait_event(own[g])
-            stream.wait_event(own_cin[g])
+            stream.wait_event(own[g] if fill is not None else placed)
             with torch.cuda.stream(stream):
                 for p in peers:
                     mem.wait(stream, pool.fin_ptr[r] + 4 * p, base + g + 1)
                 views = [slots.slot(pool, r, g, part) for part in range(3)]
+                if into:
+                    out = full[:, :, start(r, g, 0) : start(r, g, 0) + hq]
+                    attend_into(*views, out)
+                    for p in peers:
+                        # the peer's rows of this group, out of the token-major buffer
+                        src_rows = out[:, plan.shard_row(0, p) :]
+                        for b in range(batch):
+                            mem.copy2d(
+                                stream,
+                                slots.out_block_ptr(pool, p, r, g) + b * block_batch,
+                                out_row,
+                                src_rows[b].data_ptr(),
+                                merged_row,
+                                out_row,
+                                s_local,
+                            )
+                    for p in peers:
+                        mem.write(stream, pool.fout_ptr[p] + 4 * (g * world + r), call)
+                    done.append(stream.record_event())
+                    continue
                 if squeeze:
                     out = attend(*(t[0] for t in views)).contiguous()[None]
                 else:
@@ -789,6 +856,23 @@ def place_replicated(stream, part, g):
                 done.append(stream.record_event())
         for event in done:
             main.wait_event(event)
+        if into:
+            # the peers' blocks of this rank's rows, every group in one pass each
+            for g in range(groups):
+                for p in peers:
+                    mem.wait(main, pool.fout_ptr[r] + 4 * (g * world + p), call)
+            outb = pool.outb[r]
+            for p in peers:
+                blocks = outb.as_strided(
+                    (groups, batch, s_local, out_row // esz),
+                    (slots.out_numel, s_local * out_row // esz, out_row // esz, 1),
+                    outb.storage_offset() + p * groups * slots.out_numel,
+                )
+                heads = merged[:, :, start(p, 0, 0) : start(p, groups, 0)]
+                copy_head_groups(
+                    blocks,
+                    heads.unflatten(2, (groups, hq)).flatten(3).permute(2, 0, 1, 3),
+                )
         # the next call's fill rewrites `send` and the caller may free `v`, both
         # read by the copy stream
         main.wait_stream(cin)
@@ -974,6 +1058,7 @@ def ulysses_pipelined_attention(
     signature=None,
     replicated=(None, None, None),
     replicated_first: bool = True,
+    attend_into=None,
 ):
     """Pipelined Ulysses attention over the copy-engine transport, or None
     when the Ulysses group cannot use it (the caller keeps its NCCL path).
@@ -994,4 +1079,5 @@ def ulysses_pipelined_attention(
         signature=signature,
         replicated=replicated,
         replicated_first=replicated_first,
+        attend_into=attend_into,
     )

python/sglang/multimodal_gen/runtime/layers/attention/backends/attention_backend.py

@@ -157,6 +157,19 @@ def forward(
     ) -> torch.Tensor:
         raise NotImplementedError
 
+    def forward_into(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        out: torch.Tensor | None,
+        attn_metadata: T,
+    ) -> bool:
+        """Write what ``forward`` returns into ``out`` (any token stride), or with
+        ``out=None`` only answer; False, having written nothing, when this
+        backend cannot serve these inputs that way."""
+        return False
+
     def forward_varlen(
         self,
         query: torch.Tensor,

python/sglang/multimodal_gen/runtime/layers/attention/backends/flash_attn.py

@@ -500,6 +500,40 @@ def forward(
 
         raise ValueError(f"flash attention version {fa_ver} is not supported.")
 
+    def forward_into(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        out: torch.Tensor | None,
+        attn_metadata: AttentionMetadata = None,
+    ) -> bool:
+        # FA4 writes its epilogue through any token stride; same call as forward
+        if fa_ver != 4 or self._request_skip_softmax_threshold()[0]:
+            return False
+        if out is None:
+            return True
+        if attn_metadata is not None:
+            max_seqlen_q = attn_metadata.max_seqlen_q or query.shape[1]
+            max_seqlen_k = attn_metadata.max_seqlen_k or key.shape[1]
+        else:
+            max_seqlen_q, max_seqlen_k = query.shape[1], key.shape[1]
+        flash_attn_varlen_func(
+            q=query,
+            k=key,
+            v=value,
+            cu_seqlens_q=None,
+            cu_seqlens_k=None,
+            max_seqlen_q=max_seqlen_q,
+            max_seqlen_k=max_seqlen_k,
+            softmax_scale=self.softmax_scale,
+            causal=self.causal,
+            window_size=(-1, -1),
+            ver=fa_ver,
+            out=out,
+        )
+        return True
+
     def forward_varlen(
         self,
         query: torch.Tensor,