diff --git a/python/sglang/kernels/ops/diffusion/__init__.py b/python/sglang/kernels/ops/diffusion/__init__.py index 6c5eeb542f27..edcd75f6a567 100644 --- a/python/sglang/kernels/ops/diffusion/__init__.py +++ b/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", diff --git a/python/sglang/kernels/ops/diffusion/rope/rmsnorm_interleaved_rope_cutedsl.py b/python/sglang/kernels/ops/diffusion/rope/rmsnorm_interleaved_rope_cutedsl.py index f075665eefbb..b23412faaeb2 100644 --- a/python/sglang/kernels/ops/diffusion/rope/rmsnorm_interleaved_rope_cutedsl.py +++ b/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) diff --git a/python/sglang/multimodal_gen/runtime/distributed/device_communicators/ipc_a2a_multi.py b/python/sglang/multimodal_gen/runtime/distributed/device_communicators/ipc_a2a_multi.py index 84955bbde241..e900a90895ee 100644 --- a/python/sglang/multimodal_gen/runtime/distributed/device_communicators/ipc_a2a_multi.py +++ b/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, ) diff --git a/python/sglang/multimodal_gen/runtime/layers/attention/layer.py b/python/sglang/multimodal_gen/runtime/layers/attention/layer.py index 83f8b37981b0..bb50b4a923f1 100644 --- a/python/sglang/multimodal_gen/runtime/layers/attention/layer.py +++ b/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, diff --git a/python/sglang/multimodal_gen/runtime/managers/forward_context.py b/python/sglang/multimodal_gen/runtime/managers/forward_context.py index 02b992cb8a5c..f2676ec48a98 100644 --- a/python/sglang/multimodal_gen/runtime/managers/forward_context.py +++ b/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: diff --git a/python/sglang/multimodal_gen/runtime/models/dits/wanvideo.py b/python/sglang/multimodal_gen/runtime/models/dits/wanvideo.py index 6a4a03815f46..cd137cca4183 100755 --- a/python/sglang/multimodal_gen/runtime/models/dits/wanvideo.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/wanvideo.py @@ -17,7 +17,9 @@ can_use_linear_gelu, fused_linear_gelu_tanh, fused_rmsnorm_interleaved_rope_, + fused_rmsnorm_interleaved_rope_into_blocks, fused_temb_table_slices, + rmsnorm_interleaved_rope_block_cols, tensors_equal, ) from sglang.multimodal_gen.configs.models.dits import WanVideoConfig @@ -45,6 +47,8 @@ from sglang.multimodal_gen.runtime.layers.linear import ( ColumnParallelLinear, RowParallelLinear, + UnquantizedLinearMethod, + apply_unquantized_linear, ) from sglang.multimodal_gen.runtime.layers.mlp import MLP from sglang.multimodal_gen.runtime.layers.quantization.configs.base_config import ( @@ -63,7 +67,10 @@ PatchEmbed, TimestepEmbedder, ) -from sglang.multimodal_gen.runtime.managers.forward_context import get_forward_context +from sglang.multimodal_gen.runtime.managers.forward_context import ( + DENOISING_LOOP_STATE, + get_forward_context, +) from sglang.multimodal_gen.runtime.managers.memory_managers.layerwise_offload import ( LayerwiseOffloadableModuleMixin, ) @@ -264,12 +271,13 @@ class WanT2VCrossAttention(WanSelfAttention): def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs, is_cross_attention=True) - def forward(self, x, context, context_lens): + def forward(self, x, context, context_lens, context_kv=None): r""" Args: x(Tensor): Shape [B, L1, C] context(Tensor): Shape [B, L2, C] context_lens(Tensor): Shape [B] + context_kv: ``project_context_kv(context)``, when already projected """ q, _ = self.to_q(x) if self.tp_rmsnorm: @@ -278,15 +286,7 @@ def forward(self, x, context, context_lens): q = self.norm_q(q) q = q.unflatten(2, (self.local_num_heads, self.head_dim)) - k, _ = self.to_k(context) - if self.tp_rmsnorm: - k = tensor_parallel_rms_norm(k, self.norm_k) - else: - k = self.norm_k(k) - k = k.unflatten(2, (self.local_num_heads, self.head_dim)) - - v, _ = self.to_v(context) - v = v.unflatten(2, (self.local_num_heads, self.head_dim)) + k, v = self.project_context_kv(context) if context_kv is None else context_kv # compute attention x = self.attn(q, k, v) @@ -296,6 +296,20 @@ def forward(self, x, context, context_lens): x, _ = self.to_out(x) return x + def project_context_kv( + self, context: torch.Tensor + ) -> tuple[torch.Tensor, torch.Tensor]: + """The normed K and the V of ``context``, as [B, L2, heads, head_dim]""" + k, _ = self.to_k(context) + if self.tp_rmsnorm: + k = tensor_parallel_rms_norm(k, self.norm_k) + else: + k = self.norm_k(k) + k = k.unflatten(2, (self.local_num_heads, self.head_dim)) + v, _ = self.to_v(context) + v = v.unflatten(2, (self.local_num_heads, self.head_dim)) + return k, v + class WanI2VCrossAttention(WanSelfAttention): def __init__( @@ -420,6 +434,26 @@ def fused(): return _eager_temb_table_slices(table, temb) +def _rows_of_one_buffer(parts: list[torch.Tensor]) -> torch.Tensor | None: + """``torch.cat(parts)`` without a copy, when the parts are consecutive equal-shape + rows of one buffer""" + first = parts[0] + start, nbytes = first.data_ptr(), first.nbytes + storage = first.untyped_storage() + if start + len(parts) * nbytes > storage.data_ptr() + storage.nbytes(): + return None + for i, part in enumerate(parts): + if ( + part.data_ptr() != start + i * nbytes + or part.shape != first.shape + or not part.is_contiguous() + ): + return None + return first.as_strided( + (len(parts) * first.shape[0], *first.shape[1:]), first.stride() + ) + + class WanTransformerBlock(nn.Module): def __init__( self, @@ -608,8 +642,62 @@ def __init__( quant_config=quant_config, ) self.mlp_residual = MulAdd() + # norm1 that first adds the previous block's gated FFN output (its MulAdd) + self.ffn_residual_norm1 = ScaleResidualLayerNormScaleShift( + dim, + eps=eps, + elementwise_affine=False, + dtype=torch.float32, + ) self.scale_shift_table = nn.Parameter(torch.randn(1, 6, dim) / dim**0.5) + # set by fuse_qkv_weights: to_q / to_k / to_v are plain unquantized linears + self.qkv_fusable = False + # exchange-block addresses -> (q, k) destination tables of _attend_qk_filled + self._qk_tables: dict[tuple, tuple[torch.Tensor, torch.Tensor]] = {} + + def fuse_qkv_weights(self) -> None: + """Move to_q / to_k / to_v into consecutive rows of one weight and one bias, so a + single GEMM projects all three (``_fused_qkv``)""" + layers = (self.to_q, self.to_k, self.to_v) + self.qkv_fusable = all( + type(layer) is ColumnParallelLinear + and isinstance(layer.quant_method, UnquantizedLinearMethod) + and layer.bias is not None + and not layer.skip_bias_add + and not layer.gather_output + for layer in layers + ) + if not self.qkv_fusable: + return + with torch.no_grad(): + weight = torch.cat([layer.weight for layer in layers]) + bias = torch.cat([layer.bias for layer in layers]) + rows = self.to_q.weight.shape[0] + for i, layer in enumerate(layers): + layer.weight.data = weight[i * rows : (i + 1) * rows] + layer.bias.data = bias[i * rows : (i + 1) * rows] + + def _fused_qkv( + self, device: torch.device + ) -> tuple[torch.Tensor, torch.Tensor] | None: + """The [3 * dim, dim] weight and [3 * dim] bias of ``fuse_qkv_weights``, read back + from the layers while they still share them on ``device``. Only the layers hold + the shared buffers, so moving the module or offloading or swapping a layer's + weight frees them.""" + layers = (self.to_q, self.to_k, self.to_v) + if ( + not self.qkv_fusable + or torch.compiler.is_compiling() + or not all(type(layer) is ColumnParallelLinear for layer in layers) + or self.to_q.weight.device != device + ): + return None + weight = _rows_of_one_buffer([layer.weight for layer in layers]) + bias = _rows_of_one_buffer([layer.bias for layer in layers]) + if weight is None or bias is None: + return None + return weight, bias def _qkv( self, norm_hidden_states: torch.Tensor, local_heads: int @@ -624,11 +712,29 @@ def _qkv_rope( ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, bool]: """q, k, v as [B, S, heads, head_dim], and whether q and k already carry the RoPE of rope_cos_sin_cache (applied in the q/k norm pass)""" + return self._norm_qk( + *self._project_qkv(norm_hidden_states), local_heads, rope_cos_sin_cache + ) + + def _project_qkv( + self, norm_hidden_states: torch.Tensor + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + fused = self._fused_qkv(norm_hidden_states.device) + if fused is not None: + qkv = apply_unquantized_linear(norm_hidden_states, *fused) + return qkv.split(self.to_q.weight.shape[0], dim=-1) query, _ = self.to_q(norm_hidden_states) key, _ = self.to_k(norm_hidden_states) value, _ = self.to_v(norm_hidden_states) - qk_norms = ((query, self.norm_q), (key, self.norm_k)) - roped = ( + return query, key, value + + def _fuses_qk_norm_rope( + self, + query: torch.Tensor, + key: torch.Tensor, + rope_cos_sin_cache: torch.Tensor | None, + ) -> bool: + return ( rope_cos_sin_cache is not None and self.fuse_qk_norm_rope and not torch.compiler.is_compiling() @@ -639,9 +745,106 @@ def _qkv_rope( rope_cos_sin_cache, self.dim_head, ) - for x, norm in qk_norms + for x, norm in ((query, self.norm_q), (key, self.norm_k)) ) ) + + def _attend_qk_filled( + self, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + rope_cos_sin_cache: torch.Tensor | None, + ) -> torch.Tensor | None: + """Self-attention whose q/k norm-RoPE pass writes straight into the Ulysses + exchange blocks (this rank's own slot rows and the peers' send blocks), so no + own-row copy of q and k runs; None, with q/k untouched, when it cannot.""" + if ( + type(self.attn1) is not USPAttention + or self.use_offline_qk_rotation + or not self._fuses_qk_norm_rope(query, key, rope_cos_sin_cache) + ): + return None + q2d, k2d = (x.view(-1, x.shape[-1]) for x in (query, key)) + block_cols = rmsnorm_interleaved_rope_block_cols(q2d.shape[-1]) + + def fill_all(blocks) -> None: + tables = self._qk_block_tables(blocks, block_cols, q2d.shape[-1]) + if tables is None: + # blocks the kernel cannot address: norm in place, then copy them + for x, norm in ((q2d, self.norm_q), (k2d, self.norm_k)): + fused_rmsnorm_interleaved_rope_( + x, + norm.weight, + rope_cos_sin_cache, + norm.variance_epsilon, + self.dim_head, + ) + for head_start, head_count, q_dst, k_dst in blocks: + cols = slice( + head_start * self.dim_head, + (head_start + head_count) * self.dim_head, + ) + q_dst.copy_(q2d[:, cols].view_as(q_dst)) + k_dst.copy_(k2d[:, cols].view_as(k_dst)) + return + for x, norm, table in zip((q2d, k2d), (self.norm_q, self.norm_k), tables): + fused_rmsnorm_interleaved_rope_into_blocks( + x, + norm.weight, + rope_cos_sin_cache, + norm.variance_epsilon, + self.dim_head, + table, + ) + + return self.attn1.forward_qk_filled( + *( + x.squeeze(1).unflatten(2, (self.local_num_heads, self.dim_head)) + for x in (query, key, value) + ), + fill_all, + ) + + def _qk_block_tables( + self, blocks, block_cols: int, hidden: int + ) -> tuple[torch.Tensor, torch.Tensor] | None: + # destination address of every block_cols-wide column run, for q and for k; + # the exchange blocks keep their addresses, so the tables are built once + if ( + any( + head_count * self.dim_head != block_cols + or (head_start * self.dim_head) % block_cols + or not (q_dst.is_contiguous() and k_dst.is_contiguous()) + for head_start, head_count, q_dst, k_dst in blocks + ) + or len(blocks) * block_cols != hidden + ): + return None + ordered = sorted(blocks, key=lambda block: block[0]) + key = tuple((q.data_ptr(), k.data_ptr()) for _, _, q, k in ordered) + tables = self._qk_tables.get(key) + if tables is None: + tables = self._qk_tables[key] = tuple( + torch.tensor( + [block[i].data_ptr() for block in ordered], + dtype=torch.int64, + device=ordered[0][2].device, + ) + for i in (2, 3) + ) + return tables + + def _norm_qk( + self, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + local_heads: int, + rope_cos_sin_cache: torch.Tensor | None, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, bool]: + qk_norms = ((query, self.norm_q), (key, self.norm_k)) + roped = self._fuses_qk_norm_rope(query, key, rope_cos_sin_cache) if roped: for x, norm in qk_norms: fused_rmsnorm_interleaved_rope_( @@ -680,7 +883,9 @@ def _cross_and_ffn( c_scale_msa: torch.Tensor, c_gate_msa: torch.Tensor, orig_dtype: torch.dtype, - ) -> torch.Tensor: + defer_ffn_residual: bool = False, + text_kv: tuple[torch.Tensor, torch.Tensor] | None = None, + ) -> torch.Tensor | tuple[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]: attn_output, _ = self.to_out(attn_output) attn_output = attn_output.squeeze(1) null_shift = null_scale = torch.zeros( @@ -691,17 +896,36 @@ def _cross_and_ffn( ) norm_hidden_states = norm_hidden_states.to(orig_dtype) hidden_states = hidden_states.to(orig_dtype) - attn_output = self.attn2( - norm_hidden_states, context=encoder_hidden_states, context_lens=None - ) + if text_kv is None: + attn_output = self.attn2( + norm_hidden_states, context=encoder_hidden_states, context_lens=None + ) + else: + attn_output = self.attn2( + norm_hidden_states, + context=encoder_hidden_states, + context_lens=None, + context_kv=text_kv, + ) norm_hidden_states, hidden_states = self.cross_attn_residual_norm( hidden_states, attn_output, 1, c_shift_msa, c_scale_msa ) norm_hidden_states = norm_hidden_states.to(orig_dtype) hidden_states = hidden_states.to(orig_dtype) ff_output = self.ffn(norm_hidden_states) + if defer_ffn_residual: + return hidden_states, (ff_output, c_gate_msa) return self.mlp_residual(ff_output, c_gate_msa, hidden_states).to(orig_dtype) + @property + def chains_ffn_residual(self) -> bool: + """Whether the FFN residual (MulAdd) and the next norm1 run as one kernel with + the same bits: both sides on their CUDA kernels""" + return all( + op._forward_method == op.forward_cuda + for op in (self.mlp_residual, self.norm1, self.ffn_residual_norm1) + ) + def forward( self, hidden_states: torch.Tensor, @@ -709,7 +933,13 @@ def forward( temb: torch.Tensor, freqs_cis: tuple[torch.Tensor, torch.Tensor], rope_cos_sin_cache: torch.Tensor | None = None, - ) -> torch.Tensor: + ffn_residual: tuple[torch.Tensor, torch.Tensor] | None = None, + defer_ffn_residual: bool = False, + text_kv: tuple[torch.Tensor, torch.Tensor] | None = None, + ) -> torch.Tensor | tuple[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]: + """``ffn_residual`` is the previous block's (ff_output, gate), still to be added + to ``hidden_states``; ``defer_ffn_residual`` hands this block's back the same way. + ``text_kv`` is the cross-attention's projected (K, V) of ``encoder_hidden_states``.""" if hidden_states.dim() == 4: hidden_states = hidden_states.squeeze(1) bs, seq_length, _ = hidden_states.shape @@ -739,9 +969,43 @@ def forward( assert shift_msa.dtype == torch.float32 # 1. Self-attention - norm_hidden_states = self.norm1(hidden_states, shift_msa, scale_msa) - query, key, value, roped = self._qkv_rope( - norm_hidden_states, self.local_num_heads, rope_cos_sin_cache + if ffn_residual is None: + norm_hidden_states = self.norm1(hidden_states, shift_msa, scale_msa) + else: + ff_output, ff_gate = ffn_residual + norm_hidden_states, hidden_states = self.ffn_residual_norm1( + hidden_states, ff_output, ff_gate, shift_msa, scale_msa + ) + query, key, value = self._project_qkv(norm_hidden_states) + attn_output = self._attend_qk_filled(query, key, value, rope_cos_sin_cache) + if attn_output is None: + attn_output = self._self_attention( + query, key, value, freqs_cis, rope_cos_sin_cache + ) + attn_output = attn_output.flatten(2) + return self._cross_and_ffn( + hidden_states, + attn_output, + encoder_hidden_states, + gate_msa, + c_shift_msa, + c_scale_msa, + c_gate_msa, + orig_dtype, + defer_ffn_residual, + text_kv, + ) + + def _self_attention( + self, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + freqs_cis: tuple[torch.Tensor, torch.Tensor], + rope_cos_sin_cache: torch.Tensor | None, + ) -> torch.Tensor: + query, key, value, roped = self._norm_qk( + query, key, value, self.local_num_heads, rope_cos_sin_cache ) # Apply rotary embeddings @@ -800,18 +1064,7 @@ def forward( query = torch.matmul(query, self.q_rot) key = torch.matmul(key, self.k_rot) - attn_output = self.attn1(query, key, value) - attn_output = attn_output.flatten(2) - return self._cross_and_ffn( - hidden_states, - attn_output, - encoder_hidden_states, - gate_msa, - c_shift_msa, - c_scale_msa, - c_gate_msa, - orig_dtype, - ) + return self.attn1(query, key, value) class WanTransformerBlock_VSA(nn.Module): @@ -1067,6 +1320,10 @@ def forward( return hidden_states +# DENOISING_LOOP_STATE entry of the text K/V (WanTransformer3DModel._text_kvs) +_TEXT_KV_STATE = "wan_text_kv" + + class WanTransformer3DModel(CachableDiT, LayerwiseOffloadableModuleMixin): _fsdp_shard_conditions = [is_block] _compile_conditions = [is_block] @@ -1158,6 +1415,84 @@ def __init__( self.layer_names = ["blocks"] + def _text_kvs( + self, source: torch.Tensor, context: torch.Tensor + ) -> list[tuple[torch.Tensor, torch.Tensor]] | None: + """Every block's cross-attention (K, V) of ``context``, the embedding of ``source``. + + The denoising stage hands the same prompt tensor to every step of a loop, so the + K/V are projected once per (model, prompt tensor) and kept in the request's + DENOISING_LOOP_STATE, which the stage drops when the loop ends; the module holds + none of them. Holding the tensor keeps its id from being reused while cached. + Off when the projections could differ step to step (LoRA wrappers, offloaded + weights) or would be captured. + """ + req = get_forward_context().forward_batch + if ( + req is None + or not isinstance(source, torch.Tensor) + or self.layerwise_offload_managers + or torch.compiler.is_compiling() + or torch.cuda.is_current_stream_capturing() + or not self._plain_blocks() + or not all( + type(block.attn2) is WanT2VCrossAttention + and type(block.attn2.to_k) is ColumnParallelLinear + and type(block.attn2.to_v) is ColumnParallelLinear + for block in self.blocks + ) + ): + return None + # (id(model), id(prompt tensor)) -> (that tensor, every block's K/V) + cached = req.extra.setdefault(DENOISING_LOOP_STATE, {}).setdefault( + _TEXT_KV_STATE, {} + ) + key = (id(self), id(source)) + entry = cached.pop(key, None) + if ( + entry is not None + and entry[0] is source + and entry[1][0][0].device == context.device + ): + cached[key] = entry + return entry[1] + # the other expert is done with its K/V; keep this model's other CFG branch + mine = [k for k in cached if k[0] == id(self)] + for stale in [k for k in cached if k[0] != id(self)] + mine[:-1]: + del cached[stale] + del entry + kvs = [block.attn2.project_context_kv(context) for block in self.blocks] + cached[key] = (source, kvs) + return kvs + + def _plain_blocks(self) -> bool: + # plain Wan blocks in a plain list: no cache adapter or block hook replaced them + return type(self.blocks) is nn.ModuleList and all( + type(block) is WanTransformerBlock and "forward" not in vars(block) + for block in self.blocks + ) + + def _chains_ffn_residuals(self) -> bool: + return ( + len(self.blocks) > 1 + and self._plain_blocks() + and all(block.chains_ffn_residual for block in self.blocks) + ) + + def post_load_weights(self) -> None: + super().post_load_weights() + for block in self.blocks: + if isinstance(block, WanTransformerBlock): + block.fuse_qkv_weights() + fused = sum( + isinstance(block, WanTransformerBlock) + and block._fused_qkv(block.to_q.weight.device) is not None + for block in self.blocks + ) + logger.info( + "Fused q/k/v projection weights in %d/%d blocks", fused, len(self.blocks) + ) + def _make_block( self, *, @@ -1210,6 +1545,7 @@ def forward( guidance=None, **kwargs, ) -> torch.Tensor: + text_source = encoder_hidden_states forward_batch = get_forward_context().forward_batch if forward_batch is not None: sequence_shard_enabled = ( @@ -1381,14 +1717,44 @@ def forward( ], dim=-1, ) - for block in self.blocks: - hidden_states = block( + text_kvs = self._text_kvs(text_source, encoder_hidden_states) + text_kv_kwargs = [ + {} if text_kvs is None else {"text_kv": kv} + for kv in (text_kvs or [None] * len(self.blocks)) + ] + if self._chains_ffn_residuals(): + # each block's FFN residual is added by the next block's norm1 + ffn_residual = None + for block, kv_kwargs in zip(self.blocks[:-1], text_kv_kwargs): + hidden_states, ffn_residual = block( + hidden_states, + encoder_hidden_states, + timestep_proj, + freqs_cis, + rope_cos_sin_cache=rope_cos_sin_cache, + ffn_residual=ffn_residual, + defer_ffn_residual=True, + **kv_kwargs, + ) + hidden_states = self.blocks[-1]( hidden_states, encoder_hidden_states, timestep_proj, freqs_cis, rope_cos_sin_cache=rope_cos_sin_cache, + ffn_residual=ffn_residual, + **text_kv_kwargs[-1], ) + else: + for block, kv_kwargs in zip(self.blocks, text_kv_kwargs): + hidden_states = block( + hidden_states, + encoder_hidden_states, + timestep_proj, + freqs_cis, + rope_cos_sin_cache=rope_cos_sin_cache, + **kv_kwargs, + ) # if teacache is enabled, we need to cache the original hidden states if self.enable_teacache: self.maybe_cache_states(hidden_states, original_hidden_states) diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py index 222de08bf8f7..a4c63bc18d8e 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py @@ -105,7 +105,10 @@ configure_sta, save_mask_search_results, ) -from sglang.multimodal_gen.runtime.managers.forward_context import set_forward_context +from sglang.multimodal_gen.runtime.managers.forward_context import ( + DENOISING_LOOP_STATE, + set_forward_context, +) from sglang.multimodal_gen.runtime.managers.memory_managers.component_manager import ( ComponentUse, ) @@ -1751,6 +1754,8 @@ def _post_denoising_loop( *args, **kwargs, ): + # what the models kept for this loop is freed before decode + batch.extra.pop(DENOISING_LOOP_STATE, None) # Gather results if using sequence parallelism if trajectory_latents: trajectory_tensor = torch.stack(trajectory_latents, dim=1) diff --git a/python/sglang/multimodal_gen/test/unit/test_wan_temb_table_slices.py b/python/sglang/multimodal_gen/test/unit/test_wan_temb_table_slices.py index 34359c7c2a9a..87f1317609dc 100644 --- a/python/sglang/multimodal_gen/test/unit/test_wan_temb_table_slices.py +++ b/python/sglang/multimodal_gen/test/unit/test_wan_temb_table_slices.py @@ -146,7 +146,10 @@ def run(start_frames, forward_batch): @torch.no_grad() def test_blocks_receive_fp32_modulation(self): forward_batch = SimpleNamespace( - enable_sequence_shard=False, enable_teacache=False, enable_spectrum=False + enable_sequence_shard=False, + enable_teacache=False, + enable_spectrum=False, + extra={}, ) with ( tiny_dit_runtime(torch.bfloat16), diff --git a/test/registered/kernels/ops/diffusion/test_model_fast_paths.py b/test/registered/kernels/ops/diffusion/test_model_fast_paths.py index 2ad6d95d6557..91b06d6b41c4 100644 --- a/test/registered/kernels/ops/diffusion/test_model_fast_paths.py +++ b/test/registered/kernels/ops/diffusion/test_model_fast_paths.py @@ -19,10 +19,13 @@ without it a test still passes when the fast path never engaged at all. """ +import gc import sys import unittest +import weakref +from contextlib import ExitStack, contextmanager from copy import deepcopy -from unittest.mock import patch +from unittest.mock import MagicMock, patch import pytest import torch @@ -62,6 +65,12 @@ from sglang.multimodal_gen.configs.models.vaes.stablediffusion3 import ( StableDiffusion3VAEConfig, ) +from sglang.multimodal_gen.configs.sample.sampling_params import SamplingParams +from sglang.multimodal_gen.runtime.distributed.parallel_state import ( + maybe_init_distributed_environment_and_model_parallel, + model_parallel_is_initialized, +) +from sglang.multimodal_gen.runtime.layers.attention import layer as attn_layer from sglang.multimodal_gen.runtime.layers.attention.backends import sdpa as sdpa_backend from sglang.multimodal_gen.runtime.layers.layernorm import ( RMSNorm, @@ -72,6 +81,10 @@ from sglang.multimodal_gen.runtime.layers.rotary_embedding.utils import ( _apply_rotary_emb, ) +from sglang.multimodal_gen.runtime.managers.forward_context import ( + DENOISING_LOOP_STATE, + set_forward_context, +) from sglang.multimodal_gen.runtime.models.dits.ernie_image import ( _ernie_gated_norm_scale_shift, _ernie_norm_scale_shift, @@ -115,6 +128,11 @@ from sglang.multimodal_gen.runtime.models.dits.sana import ( sana_ln_modulate, ) +from sglang.multimodal_gen.runtime.models.dits.wanvideo import ( + _TEXT_KV_STATE, + WanTransformer3DModel, + WanTransformerBlock, +) from sglang.multimodal_gen.runtime.models.vaes import ( autoencoder_kl_qwenimage as qwen_vae, ) @@ -141,13 +159,18 @@ VaeFastPathGate, ) from sglang.multimodal_gen.runtime.models.vaes.wanvae import WanRMS_norm +from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req +from sglang.multimodal_gen.runtime.pipelines_core.stages.denoising import DenoisingStage from sglang.multimodal_gen.runtime.platforms.interface import DeviceCapability +from sglang.multimodal_gen.test.single_test_file.component_accuracy.utils import ( + ensure_distributed_env_defaults, +) from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci from sglang.test.test_utils import CustomTestCase -register_cuda_ci(est_time=100, stage="base-b-kernel-unit", runner_config="1-gpu-large") +register_cuda_ci(est_time=140, stage="base-b-kernel-unit", runner_config="1-gpu-large") # backend-specific: preserve ROCm eager fallback coverage for CUDA-only fusions. -register_amd_ci(est_time=6, suite="nightly-amd-kernel-1-gpu", nightly=True) +register_amd_ci(est_time=8, suite="nightly-amd-kernel-1-gpu", nightly=True) pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") @@ -2405,5 +2428,220 @@ def test_moss_image_selected_kernel_rejects_mismatched_operands(monkeypatch, ope assert not moss_image._RMSNORM_FUSION.disabled +# ------------------------------------------------------------------------- +# Wan -- fused q/k/v GEMM, FFN residual chain, text K/V once per loop +# ------------------------------------------------------------------------- + + +@contextmanager +def _wan_block_env(): + """Single-rank model parallel state and the attention patches a Wan block + needs to be built and run outside a pipeline.""" + if not model_parallel_is_initialized(): + ensure_distributed_env_defaults() + maybe_init_distributed_environment_and_model_parallel(tp_size=1, sp_size=1) + with ExitStack() as stack: + for name, value in ( + ("get_attn_backend", sdpa_backend.SDPABackend), + ("get_compute_dtype", torch.bfloat16), + ("_resolve_sp_attention_mode", ("ulysses", False)), + ("_maybe_install_backend_autotune", None), + ): + stack.enter_context(patch.object(attn_layer, name, return_value=value)) + stack.enter_context(set_forward_context(0, None)) + yield + + +def _wan_block(dim=5120, heads=40): + block = WanTransformerBlock( + dim, + 1024, + heads, + qk_norm="rms_norm_across_heads", + cross_attn_norm=True, + eps=1e-6, + supported_attention_backends=set(), + ).cuda() + with torch.no_grad(): + for p in block.parameters(): + p.normal_(0, 0.02) + return block.to(torch.bfloat16) + + +def _wan_block_inputs(seq=512, dim=5120): + """x, (context, temb, freqs) as the block's positional args, and the + rope cos/sin cache""" + x = torch.randn(1, seq, dim, device="cuda", dtype=torch.bfloat16) + context = torch.randn(1, 64, dim, device="cuda", dtype=torch.bfloat16) + temb = torch.randn(1, 6, dim, device="cuda") + theta = torch.rand(seq, 64, device="cuda") * 50 + freqs = (theta.cos(), theta.sin()) + return x, (context, temb, freqs), torch.cat(freqs, dim=-1).contiguous() + + +def _bits(x): + return x.view(torch.int16) + + +@requires_inline_ptx +@torch.no_grad() +def test_wan_fused_qkv_keeps_bits_and_frees_with_its_layers(): + """q/k/v share one weight after loading and run as one GEMM with the same + bits; a layer that stops reading the shared weight sends the block back to + its three projections, and moving the block takes the shared weight along.""" + with _wan_block_env(): + block = _wan_block() + x, args, cache = _wan_block_inputs() + run = lambda: _bits(block(x, *args, rope_cos_sin_cache=cache)) + projs = (block.to_q, block.to_k, block.to_v) + weights = [proj.weight.clone() for proj in projs] + ref = run() + assert block._fused_qkv(x.device) is None + block.fuse_qkv_weights() + assert block._fused_qkv(x.device) is not None + for proj, w in zip(projs, weights): + assert torch.equal(proj.weight, w) + assert torch.equal(run(), ref) + + # a layer whose weight moved (offload, a weight swap) stops the fused GEMM + block.to_k.weight.data = block.to_k.weight.data.clone() + assert block._fused_qkv(x.device) is None + assert torch.equal(run(), ref) + + # moving the block takes the shared weight with it: nothing else holds it + block.fuse_qkv_weights() + on_gpu = [*block.parameters(), *block.buffers()] + gpu_bytes = sum(t.numel() * t.element_size() for t in on_gpu) + del on_gpu + allocated = torch.cuda.memory_allocated() + block.to("cpu") + assert allocated - torch.cuda.memory_allocated() >= gpu_bytes + assert block._fused_qkv(torch.device("cpu")) is None + block.fuse_qkv_weights() + fused_weight, fused_bias = block._fused_qkv(torch.device("cpu")) + assert torch.equal(fused_weight, torch.cat(weights).cpu()) + assert torch.equal(fused_bias, torch.cat([p.bias for p in projs])) + + block.to("cuda") + assert block._fused_qkv(x.device) is None + assert torch.equal(run(), ref) + block.fuse_qkv_weights() + assert block._fused_qkv(x.device) is not None + assert torch.equal(run(), ref) + + +@requires_inline_ptx +@torch.no_grad() +def test_wan_chained_blocks_match_blocks_one_by_one(): + """Each block hands its FFN residual (MulAdd) to the next block's norm1, which + adds it in the same kernel; the chain keeps the bits of blocks run one by one.""" + with _wan_block_env(): + blocks = [_wan_block() for _ in range(3)] + assert all(block.chains_ffn_residual for block in blocks) + x, args, cache = _wan_block_inputs() + x = x * 4 + + ref = x + for block in blocks: + ref = block(ref, *args, rope_cos_sin_cache=cache) + + hidden, residual = x, None + for block in blocks[:-1]: + hidden, residual = block( + hidden, + *args, + rope_cos_sin_cache=cache, + ffn_residual=residual, + defer_ffn_residual=True, + ) + out = blocks[-1](hidden, *args, rope_cos_sin_cache=cache, ffn_residual=residual) + assert torch.equal(_bits(out), _bits(ref)) + + +@requires_inline_ptx +@torch.no_grad() +def test_wan_projected_text_kv_keeps_bits(): + """A block given the cross-attention K/V projected ahead of the step gives the + bits of the block projecting them itself.""" + with _wan_block_env(): + block = _wan_block() + x, args, cache = _wan_block_inputs() + kv = block.attn2.project_context_kv(args[0]) + out = block(x, *args, rope_cos_sin_cache=cache, text_kv=kv) + ref = block(x, *args, rope_cos_sin_cache=cache) + assert torch.equal(_bits(out), _bits(ref)) + + +class _WanTextKVModel(nn.Module): + """Two blocks and the model methods that project and keep their text K/V""" + + _text_kvs = WanTransformer3DModel._text_kvs + _plain_blocks = WanTransformer3DModel._plain_blocks + + def __init__(self): + super().__init__() + self.blocks = nn.ModuleList(_wan_block(1024, 8) for _ in range(2)) + self.layerwise_offload_managers = [] + + +@requires_inline_ptx +@torch.no_grad() +def test_wan_text_kv_lives_for_one_loop_of_one_request(): + """The text K/V projected once per loop live in the request for that loop only: + the next expert, the end of the loop and the end of the request free them.""" + with _wan_block_env(): + expert, expert_2 = _WanTextKVModel(), _WanTextKVModel() + context = torch.randn(1, 64, 1024, device="cuda", dtype=torch.bfloat16) + positive, negative = torch.randn(2, 1, 64, 4096, device="cuda") + new_req = lambda: Req(sampling_params=SamplingParams(prompt="test")) + + def text_kvs(model, source, req): + with set_forward_context(0, None, req): + return model._text_kvs(source, context) + + def first_k(model, source, req): + return weakref.ref(text_kvs(model, source, req)[0][0]) + + # every step of the loop reads the K/V projected on its first, per CFG branch + req = new_req() + kvs = text_kvs(expert, positive, req) + neg_kvs = text_kvs(expert, negative, req) + assert text_kvs(expert, positive, req) is kvs + assert text_kvs(expert, negative, req) is neg_kvs + for (k, v), block in zip(kvs, expert.blocks): + ref_k, ref_v = block.attn2.project_context_kv(context) + assert torch.equal(k, ref_k) and torch.equal(v, ref_v) + + # another request projects its own; K/V left on another device are projected again + assert text_kvs(expert, positive, new_req()) is not kvs + req.extra[DENOISING_LOOP_STATE][_TEXT_KV_STATE][(id(expert), id(positive))] = ( + positive, + [(k.cpu(), v.cpu()) for k, v in kvs], + ) + moved = text_kvs(expert, positive, req) + assert moved[0][0].is_cuda and torch.equal(moved[0][0], kvs[0][0]) + + # the second expert's first step frees the first expert's K/V + gone = [weakref.ref(t) for t in (moved[0][0], neg_kvs[0][0])] + del kvs, neg_kvs, moved, ref_k, ref_v, k, v + text_kvs(expert_2, positive, req) + assert all(ref() is None for ref in gone) + + # the end of the loop frees the rest, before decode + gone = first_k(expert_2, positive, req) + stage = MagicMock() + stage._postprocess_sp_latents.return_value = (None, None) + stage._sp_world_size.return_value = 1 + DenoisingStage._post_denoising_loop(stage, req, None, [], [], MagicMock(), True) + assert gone() is None and DENOISING_LOOP_STATE not in req.extra + + # a request that ends mid-loop takes its K/V with it: the model holds none + req = new_req() + gone = first_k(expert, positive, req) + del req + gc.collect() + assert gone() is None + + if __name__ == "__main__": sys.exit(pytest.main([__file__, "-v"])) diff --git a/test/registered/kernels/ops/diffusion/test_rmsnorm_interleaved_rope.py b/test/registered/kernels/ops/diffusion/test_rmsnorm_interleaved_rope.py index a0177e8be44c..563f59da6123 100644 --- a/test/registered/kernels/ops/diffusion/test_rmsnorm_interleaved_rope.py +++ b/test/registered/kernels/ops/diffusion/test_rmsnorm_interleaved_rope.py @@ -13,6 +13,8 @@ from sglang.kernels.ops.diffusion import ( can_use_fused_rmsnorm_interleaved_rope, fused_rmsnorm_interleaved_rope_, + fused_rmsnorm_interleaved_rope_into_blocks, + rmsnorm_interleaved_rope_block_cols, ) from sglang.multimodal_gen.runtime.distributed.parallel_state import ( maybe_init_distributed_environment_and_model_parallel, @@ -76,6 +78,32 @@ def test_matches_rmsnorm_then_rope_bitwise(batch, seq, hidden, head_dim, row_pad assert buf[:, hidden:].isnan().all() +def test_into_blocks_matches_in_place(): + """Each column run lands in its own block, with the in-place result's bits.""" + _require_sm90() + g = torch.Generator(device="cuda").manual_seed(2) + eps, hidden, head_dim, rows = 1e-6, 5120, 128, 300 + x = (torch.randn(rows, 3 * hidden, generator=g, device="cuda") * 3).bfloat16() + x = x[:, hidden : 2 * hidden] # a strided column slice, as in a fused qkv output + weight = (1 + 0.1 * torch.randn(hidden, generator=g, device="cuda")).bfloat16() + theta = torch.rand(rows, head_dim // 2, generator=g, device="cuda") + cache = torch.cat([theta.cos(), theta.sin()], dim=-1).contiguous() + cols = rmsnorm_interleaved_rope_block_cols(hidden) + blocks = [ + torch.full((rows, cols), float("nan"), device="cuda").bfloat16() + for _ in range(hidden // cols) + ] + table = torch.tensor([b.data_ptr() for b in reversed(blocks)], device="cuda") + raw = x.clone() + fused_rmsnorm_interleaved_rope_into_blocks(x, weight, cache, eps, head_dim, table) + assert torch.equal(x.view(torch.int16), raw.view(torch.int16)) + ref = fused_rmsnorm_interleaved_rope_(x.clone(), weight, cache, eps, head_dim) + for i, block in enumerate(reversed(blocks)): + assert torch.equal( + block.view(torch.int16), ref[:, i * cols : (i + 1) * cols].view(torch.int16) + ) + + def test_runs_on_the_current_stream(): """Launched on torch's current stream: its input may still be in flight there.""" _require_sm90()