[Perf] diffusion: Write Wan's normed q/k straight into the Ulysses exchange

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

@@ -775,6 +775,8 @@
     "fused_inplace_qknorm_rope": "rope.qknorm_rope_jit",
     "can_use_fused_rmsnorm_interleaved_rope": "rope.rmsnorm_interleaved_rope_cutedsl",
     "fused_rmsnorm_interleaved_rope_": "rope.rmsnorm_interleaved_rope_cutedsl",
+    "fused_rmsnorm_interleaved_rope_into_blocks": "rope.rmsnorm_interleaved_rope_cutedsl",
+    "rmsnorm_interleaved_rope_block_cols": "rope.rmsnorm_interleaved_rope_cutedsl",
     "fused_qknorm_rope_pack_kv": "rope.qknorm_rope_jit",
     "try_fused_qwen_qkv_epilogue": "rope.qwen_qkv_epilogue_jit",
     "can_use_fused_rope_rotate_half_fp32": "rope.rope_rotate_half_fp32",

python/sglang/kernels/ops/diffusion/rope/rmsnorm_interleaved_rope_cutedsl.py

@@ -39,10 +39,13 @@ def _threads_per_row(H: int) -> int:
 
 
 class RMSNormInterleavedRope:
-    def __init__(self, dtype: cutlass.Numeric, H: int, head_dim: int):
+    def __init__(self, dtype: cutlass.Numeric, H: int, head_dim: int, blocked: bool):
         self.dtype = dtype
         self.H = H
         self.head_dim = head_dim
+        # store each thread column run (block_cols columns) to its own [M, block_cols]
+        # destination block instead of to ``mY``
+        self.blocked = blocked
         elem_bytes = dtype.width // 8
         self.vec_size = min(H & (-H), COPY_BITS // 8 // elem_bytes)
         self.copy_bits = self.vec_size * dtype.width
@@ -74,6 +77,7 @@ def __call__(
         mW: cute.Tensor,
         mY: cute.Tensor,
         mCS: cute.Tensor,
+        mDst: cute.Tensor,
         M: Int64,
         seq_len: Int64,
         eps: Float32,
@@ -92,7 +96,7 @@ def __call__(
         )
         tv_layout = cute.make_layout(tv_shape, stride=tv_stride)
         tiler_mn = (self.rows_per_block, self.cols_per_tile)
-        self.kernel(mX, mW, mY, mCS, M, seq_len, eps, tv_layout, tiler_mn).launch(
+        self.kernel(mX, mW, mY, mCS, mDst, M, seq_len, eps, tv_layout, tiler_mn).launch(
             grid=[cute.ceil_div(M, self.rows_per_block), 1, 1],
             block=[self.num_threads, 1, 1],
             smem=self._smem_size_in_bytes(),
@@ -106,6 +110,7 @@ def kernel(
         mW: cute.Tensor,
         mY: cute.Tensor,
         mCS: cute.Tensor,
+        mDst: cute.Tensor,
         M: Int64,
         seq_len: Int64,
         eps: Float32,
@@ -242,7 +247,27 @@ def kernel(
                 tXrO[(v, blk), 0, 0] = (a * c + (-b) * s).to(mY.element_type)
                 tXrO[(v + 1, blk), 0, 0] = (b * c + a * s).to(mY.element_type)
 
-        if row_in_bounds:
+        if cutlass.const_expr(self.blocked):
+            # column run ``blk`` of row ``row`` goes to row ``row`` of block mDst[blk]
+            block_cols = self.vec_size * threads_per_row
+            elem_bytes = mY.element_type.width // 8
+            if row_in_bounds:
+                offset = (
+                    Int64(row_coord[0]) * block_cols
+                    + tXcX[(0, 0), 0, 0][1] % block_cols
+                ) * elem_bytes
+                for blk in cutlass.range_constexpr(self.num_vec_blocks):
+                    dst = cute.make_tensor(
+                        cute.make_ptr(
+                            mY.element_type,
+                            mDst[blk] + offset,
+                            cute.AddressSpace.gmem,
+                            assumed_align=16,
+                        ),
+                        cute.make_layout((self.vec_size,)),
+                    )
+                    cute.autovec_copy(tXrO[(None, blk), 0, 0], dst)
+        elif row_in_bounds:
             cute.copy(copy_atom_store, tXrO, tXgO, pred=tXpX)
 
 
@@ -275,20 +300,15 @@ def can_use_fused_rmsnorm_interleaved_rope(
     )
 
 
-def fused_rmsnorm_interleaved_rope_(
-    x: torch.Tensor,
-    weight: torch.Tensor,
-    cos_sin_cache: torch.Tensor,
-    eps: float,
-    head_dim: int,
-) -> torch.Tensor:
-    """In place: x [M, H] <- RoPE(RMSNorm(x) * weight), row m at position m % len(cos_sin_cache).
+def rmsnorm_interleaved_rope_block_cols(H: int) -> int:
+    """Columns of one destination block of ``fused_rmsnorm_interleaved_rope_into_blocks``"""
+    return 8 * _threads_per_row(H)
 
-    ``cos_sin_cache`` is [seq_len, head_dim] fp32 (cos | sin of each pair); bit-identical
-    to ``flashinfer.norm.rmsnorm`` then FlashInfer's interleaved cos/sin-cache RoPE.
-    """
+
+def _launch(x, weight, cos_sin_cache, eps, head_dim, out, dst_table) -> None:
     M, H = x.shape
-    key = (x.dtype, H, head_dim)
+    blocked = dst_table is not None
+    key = (x.dtype, H, head_dim, blocked)
     compiled = _COMPILE_CACHE.get(key)
     if compiled is None:
         dtype = _TORCH_TO_CUTE[x.dtype]
@@ -300,18 +320,66 @@ def fused_rmsnorm_interleaved_rope_(
         cs_fake = cute.runtime.make_fake_compact_tensor(
             Float32, (cute.sym_int(64), head_dim), stride_order=(1, 0), assumed_align=16
         )
+        dst_fake = cute.runtime.make_fake_compact_tensor(
+            Int64, (cute.sym_int(64),), assumed_align=8
+        )
         compiled = _COMPILE_CACHE[key] = cute.compile(
-            RMSNormInterleavedRope(dtype, H, head_dim),
+            RMSNormInterleavedRope(dtype, H, head_dim, blocked),
             x_fake,
             w_fake,
             x_fake,
             cs_fake,
+            dst_fake,
             Int64(1),
             Int64(1),
             Float32(1e-6),
             cuda.CUstream(cuda.CUstream_flags.CU_STREAM_DEFAULT),
             options="--enable-tvm-ffi",
         )
+    if dst_table is None:
+        dst_table = _no_blocks(x.device)
     stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream)
-    compiled(x, weight, x, cos_sin_cache, M, cos_sin_cache.shape[0], eps, stream)
+    compiled(
+        x, weight, out, cos_sin_cache, dst_table, M, cos_sin_cache.shape[0], eps, stream
+    )
+
+
+_NO_BLOCKS: dict[torch.device, torch.Tensor] = {}
+
+
+def _no_blocks(device: torch.device) -> torch.Tensor:
+    # the in-place variant never reads its destination table
+    table = _NO_BLOCKS.get(device)
+    if table is None:
+        table = _NO_BLOCKS[device] = torch.zeros(1, dtype=torch.int64, device=device)
+    return table
+
+
+def fused_rmsnorm_interleaved_rope_(
+    x: torch.Tensor,
+    weight: torch.Tensor,
+    cos_sin_cache: torch.Tensor,
+    eps: float,
+    head_dim: int,
+) -> torch.Tensor:
+    """In place: x [M, H] <- RoPE(RMSNorm(x) * weight), row m at position m % len(cos_sin_cache).
+
+    ``cos_sin_cache`` is [seq_len, head_dim] fp32 (cos | sin of each pair); bit-identical
+    to ``flashinfer.norm.rmsnorm`` then FlashInfer's interleaved cos/sin-cache RoPE.
+    """
+    _launch(x, weight, cos_sin_cache, eps, head_dim, x, None)
     return x
+
+
+def fused_rmsnorm_interleaved_rope_into_blocks(
+    x: torch.Tensor,
+    weight: torch.Tensor,
+    cos_sin_cache: torch.Tensor,
+    eps: float,
+    head_dim: int,
+    dst_table: torch.Tensor,
+) -> None:
+    """``fused_rmsnorm_interleaved_rope_`` written elsewhere, ``x`` untouched: columns
+    [i * B, (i + 1) * B) go to the contiguous [M, B] block at address ``dst_table[i]``
+    (int64, CUDA), B = ``rmsnorm_interleaved_rope_block_cols(H)``."""
+    _launch(x, weight, cos_sin_cache, eps, head_dim, x, dst_table)

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

@@ -493,6 +493,7 @@ def pipelined_attention(
         replicated=(None, None, None),
         replicated_first: bool = True,
         attend_into=None,
+        fill_all=None,
     ):
         """Ulysses attention with the head exchange pipelined against it.
 
@@ -525,6 +526,10 @@ 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.
 
+        ``fill_all(blocks)`` is ``fill`` given every block at once, as a list of
+        ``(head_start, head_count, q_dst, k_dst)`` (for a norm over all heads of a row,
+        which one pass over q and k then writes everywhere).
+
         ``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
@@ -541,6 +546,7 @@ def pipelined_attention(
         pool released, so that call never holds both at once.
         """
         world, r = self.world, self.rank
+        filled = fill is not None or fill_all is not None
         squeeze = q.dim() == 3
         if squeeze:
             q, k, v = q[None], k[None], v[None]
@@ -555,11 +561,11 @@ def pipelined_attention(
             tuple(None if t is None else t.shape for t in replicated),
             replicated_first,
             groups,
-            fill is not None,
+            filled,
         )
         if shapes not in self.plans:
             self.plans[shapes] = self._plan(
-                q, k, v, replicated, replicated_first, groups, fill is not None
+                q, k, v, replicated, replicated_first, groups, filled
             )
         plan = self.plans[shapes]
         if plan is None:
@@ -568,7 +574,7 @@ def pipelined_attention(
 
         # each token's heads of one group must be one contiguous run to move
         # with a pitched copy
-        copied = [v] if fill is not None else [q, k, v]
+        copied = [v] if filled else [q, k, v]
         copied += [t for t in replicated if t is not None]
         if not all(t.stride(-1) == 1 and t.stride(-2) == head_dim for t in copied):
             return None
@@ -592,7 +598,7 @@ def pipelined_attention(
             esz = q.element_size()
             held = 0 if pool is None else pool.nbytes(esz)
             need = (sizes[0] + sizes[1]) * esz
-            if fill is not None:
+            if filled:
                 need += plan.world * groups * slots.send_numel * esz
             if auto and not self._fits_on_every_rank(need, reclaimable=held):
                 logger.info(
@@ -689,14 +695,27 @@ def own_rows_views(part):
         # 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):
+        for part in (2,) if filled 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:
+        if fill_all is not None:
+            own_rows = lambda part, g: slots.slot(pool, r, g, part)[
+                0, my_shard_row[part] : my_shard_row[part] + s_local
+            ]
+            hg = plan.group_heads(0)
+            blocks = []
+            for g in range(groups):
+                for p in peers:
+                    qk = slots.send_block(pool, p, g)
+                    blocks.append((start(p, g, 0), hg, qk[0], qk[1]))
+                blocks.append((start(r, g, 0), hg, own_rows(0, g), own_rows(1, g)))
+            fill_all(blocks)
+            sent = own = [main.record_event()] * groups
+        elif fill is not None:
             # group by group: the peers' blocks first, so their copies start
             # while this rank fills its own block straight into its slot
             own_rows = lambda part, g: slots.slot(pool, r, g, part)[
@@ -715,10 +734,10 @@ def own_rows_views(part):
             for g in range(groups):
                 for p in peers:
                     send_rows(cin, 2, p, g)
-                    if fill is None:
+                    if not filled:
                         send_rows(cin, 0, p, g)
                         send_rows(cin, 1, p, g)
-                if fill is not None:
+                if filled:
                     cin.wait_event(sent[g])
                     chunk = s_local * row_bytes[0]
                     for p in peers:
@@ -775,7 +794,7 @@ def own_rows_views(part):
         done = []
         for g in range(groups):
             stream = self.pipe_group_streams[g]
-            stream.wait_event(own[g] if fill is not None else placed)
+            stream.wait_event(own[g] if filled else placed)
             with torch.cuda.stream(stream):
                 for p in peers:
                     mem.wait(stream, pool.fin_ptr[r] + 4 * p, base + g + 1)
@@ -1059,6 +1078,7 @@ def ulysses_pipelined_attention(
     replicated=(None, None, None),
     replicated_first: bool = True,
     attend_into=None,
+    fill_all=None,
 ):
     """Pipelined Ulysses attention over the copy-engine transport, or None
     when the Ulysses group cannot use it (the caller keeps its NCCL path).
@@ -1080,4 +1100,5 @@ def ulysses_pipelined_attention(
         replicated=replicated,
         replicated_first=replicated_first,
         attend_into=attend_into,
+        fill_all=fill_all,
     )

python/sglang/multimodal_gen/runtime/layers/attention/layer.py

@@ -1587,6 +1587,7 @@ def _pipelined_ulysses(
         replicated=(None, None, None),
         replicated_first: bool = True,
         attend_into=None,
+        fill_all=None,
     ) -> torch.Tensor | None:
         """Ulysses attention with the exchange pipelined over head groups.
 
@@ -1647,7 +1648,7 @@ def reference():
             v,
             attend,
             groups,
-            sequential=reference,
+            sequential=None if sequential is None else reference,
             signature=(
                 type(self.attn_impl).__name__,
                 self.causal,
@@ -1657,11 +1658,51 @@ def reference():
             replicated=replicated,
             replicated_first=replicated_first,
             attend_into=attend_into,
+            fill_all=fill_all,
         )
         if out is None:
             self._pipeline_declined.add(key)
         return out
 
+    def forward_qk_filled(
+        self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, fill_all
+    ) -> torch.Tensor | None:
+        """Dense Ulysses attention whose q/k exchange blocks ``fill_all`` writes.
+
+        ``q``/``k`` are the caller's raw projections: ``fill_all`` (see
+        ``ulysses_pipelined_attention``) writes what ``forward`` would read into
+        every exchange block, this rank's own slot rows included, so no own-row
+        copy runs. Returns None, having written nothing, when the call would not
+        take the pipelined dense path; the caller then runs ``forward`` itself.
+        """
+        if (
+            self.skip_sequence_parallel
+            or get_sequence_parallel_world_size() == 1
+            or get_ring_parallel_world_size() > 1
+            or self.enable_packed_qkv_input_a2a
+            or get_request_skip_softmax_params() is not None
+        ):
+            return None
+        ctx_attn_metadata = get_forward_context().attn_metadata
+        return self._pipelined_ulysses(
+            q,
+            k,
+            v,
+            lambda q_group, k_group, v_group: self.attn_impl.forward(
+                q_group, k_group, v_group, ctx_attn_metadata
+            ),
+            # no first-sight reference: the pipelined backends attend a head group
+            # exactly as inside the full call (PIPELINED_ATTENTION_BACKENDS)
+            None,
+            layout="dense-filled",
+            attend_into=lambda q_group, k_group, v_group, out: (
+                self.attn_impl.forward_into(
+                    q_group, k_group, v_group, out, ctx_attn_metadata
+                )
+            ),
+            fill_all=fill_all,
+        )
+
     def _forward_ring_tail_pad(
         self,
         q: torch.Tensor,

python/sglang/multimodal_gen/runtime/managers/forward_context.py

@@ -12,6 +12,10 @@
     from sglang.multimodal_gen.runtime.layers.attention import AttentionMetadata
     from sglang.multimodal_gen.runtime.pipelines_core import Req
 
+# Req.extra key for what models keep for one denoising loop (Wan's text K/V): the denoising
+# stage drops it when the loop ends; the leading "_" keeps it out of request transfers
+DENOISING_LOOP_STATE = "_denoising_loop_state"
+
 
 @dataclass
 class ForwardContext: