diff --git a/python/sglang/kernels/jit/csrc/diffusion/qknorm_rope.cuh b/python/sglang/kernels/jit/csrc/diffusion/qknorm_rope.cuh index 1d2774f6f..a514ea82c 100644 --- a/python/sglang/kernels/jit/csrc/diffusion/qknorm_rope.cuh +++ b/python/sglang/kernels/jit/csrc/diffusion/qknorm_rope.cuh @@ -31,6 +31,25 @@ struct QKNormRopeParams { float eps; }; +struct QKNormRopePackKVParams : QKNormRopeParams { + const void* __restrict__ v_ptr; + const void* __restrict__ k_prefix_ptr; + const void* __restrict__ v_prefix_ptr; + void* __restrict__ packed_k_ptr; + void* __restrict__ packed_v_ptr; + int64_t v_stride_bytes; + int64_t k_prefix_stride_bytes; + int64_t v_prefix_stride_bytes; + int64_t packed_token_stride_bytes; + int64_t packed_head_stride_bytes; + uint32_t batch_size; + uint32_t prefix_tokens; + uint32_t suffix_tokens; +}; + +template +using QKNormRopeParamsT = std::conditional_t; + constexpr uint32_t kThreadsPerBlock = 256; constexpr uint32_t kWarpsPerBlock = kThreadsPerBlock / device::kWarpThreads; @@ -152,8 +171,9 @@ template < typename DType, typename CacheDType, bool kRoundNormBeforeRope, + bool kPackKV, typename IdType> -__global__ void fused_qknorm_rope_warp(const QKNormRopeParams __grid_constant__ params) { +__global__ void fused_qknorm_rope_warp(const QKNormRopeParamsT __grid_constant__ params) { using namespace device; static_assert(std::is_same_v || std::is_same_v); @@ -181,23 +201,81 @@ __global__ void fused_qknorm_rope_warp(const QKNormRopeParams __grid_constant__ using Storage = AlignedVector; const auto& [q_ptr, k_ptr, q_weight_ptr, k_weight_ptr, cos_sin_cache_ptr, positions, q_stride_bytes, k_stride_bytes, head_stride_bytes, num_qo_heads, num_kv_heads, num_tokens, eps] = - params; + static_cast(params); const uint32_t lane_id = threadIdx.x % kWarpThreads; const uint32_t warp_id = threadIdx.x / kWarpThreads; const uint32_t start_worker_id = blockIdx.x * kWarpsPerBlock + warp_id; const uint32_t num_workers = gridDim.x * kWarpsPerBlock; const uint32_t num_qk_heads = num_qo_heads + num_kv_heads; - const uint32_t num_works = num_qk_heads * num_tokens; + const uint32_t num_qk_works = num_qk_heads * num_tokens; + uint32_t num_prefix_works = 0; + uint32_t num_works = num_qk_works; + if constexpr (kPackKV) { + num_prefix_works = params.batch_size * params.prefix_tokens * num_kv_heads; + num_works += 2 * num_prefix_works + num_tokens * num_kv_heads; + } PDLWaitPrimary(); for (uint32_t idx = start_worker_id; idx < num_works; idx += num_workers) { + if constexpr (kPackKV) { + if (idx >= num_qk_works) { + const uint32_t copy_idx = idx - num_qk_works; + const bool copy_k_prefix = copy_idx < num_prefix_works; + const bool copy_v_prefix = copy_idx >= num_prefix_works && copy_idx < 2 * num_prefix_works; + const uint32_t local_idx = + copy_k_prefix ? copy_idx : (copy_v_prefix ? copy_idx - num_prefix_works : copy_idx - 2 * num_prefix_works); + const uint32_t token_id = local_idx / num_kv_heads; + const uint32_t head_id = local_idx % num_kv_heads; + const bool copy_prefix = copy_k_prefix || copy_v_prefix; + const uint32_t batch_id = token_id / (copy_prefix ? params.prefix_tokens : params.suffix_tokens); + const uint32_t sequence_id = token_id % (copy_prefix ? params.prefix_tokens : params.suffix_tokens); + const uint32_t packed_token_id = batch_id * (params.prefix_tokens + params.suffix_tokens) + + (copy_prefix ? sequence_id : params.prefix_tokens + sequence_id); + const void* input = nullptr; + void* output = nullptr; + if (copy_k_prefix) { + input = pointer::offset( + params.k_prefix_ptr, token_id * params.k_prefix_stride_bytes, head_id * head_stride_bytes); + output = pointer::offset( + params.packed_k_ptr, + packed_token_id * params.packed_token_stride_bytes, + head_id * params.packed_head_stride_bytes); + } else { + const void* v_ptr = copy_v_prefix ? params.v_prefix_ptr : params.v_ptr; + const int64_t v_stride = copy_v_prefix ? params.v_prefix_stride_bytes : params.v_stride_bytes; + input = pointer::offset(v_ptr, token_id * v_stride, head_id * head_stride_bytes); + output = pointer::offset( + params.packed_v_ptr, + packed_token_id * params.packed_token_stride_bytes, + head_id * params.packed_head_stride_bytes); + } + const auto copy_vec = load_as(input, lane_id); + store_as(output, copy_vec, lane_id); + continue; + } + } + const uint32_t token_id = idx / num_qk_heads; const uint32_t head_id = idx % num_qk_heads; const bool load_q = head_id < num_qo_heads; const void* input = load_q ? pointer::offset(q_ptr, token_id * q_stride_bytes, head_id * head_stride_bytes) : pointer::offset(k_ptr, token_id * k_stride_bytes, head_id * head_stride_bytes); + void* output = const_cast(input); + if constexpr (kPackKV) { + if (!load_q) { + const uint32_t batch_id = token_id / params.suffix_tokens; + const uint32_t sequence_id = token_id % params.suffix_tokens; + const uint32_t kv_head_id = head_id - num_qo_heads; + const uint32_t packed_token_id = + batch_id * (params.prefix_tokens + params.suffix_tokens) + params.prefix_tokens + sequence_id; + output = pointer::offset( + params.packed_k_ptr, + packed_token_id * params.packed_token_stride_bytes, + kv_head_id * params.packed_head_stride_bytes); + } + } const void* weight_ptr = load_q ? q_weight_ptr : k_weight_ptr; auto input_vec = load_as(input, lane_id); @@ -246,7 +324,7 @@ __global__ void fused_qknorm_rope_warp(const QKNormRopeParams __grid_constant__ } } } - store_as(const_cast(input), output_vec, lane_id); + store_as(output, output_vec, lane_id); continue; } @@ -315,7 +393,7 @@ __global__ void fused_qknorm_rope_warp(const QKNormRopeParams __grid_constant__ for (uint32_t j = 0; j < kVecSize; ++j) { input_vec[j] = cast({elems[2 * j], elems[2 * j + 1]}); } - store_as(const_cast(input), input_vec, lane_id); + store_as(output, input_vec, lane_id); } PDLTriggerSecondary(); @@ -332,8 +410,16 @@ template < struct QKNormRopeKernel { static_assert(kHeadDim <= 256, "Only head_dim <= 256 is supported"); template - static constexpr auto kernel = - fused_qknorm_rope_warp; + static constexpr auto kernel = fused_qknorm_rope_warp< + kHeadDim, + kRopeDim, + kIsNeox, + kUsePDL, + DType, + CacheDType, + kRoundNormBeforeRope, + false, + IdType>; static void run(const tvm::ffi::TensorView q, @@ -405,4 +491,130 @@ struct QKNormRopeKernel { } }; +template < + int64_t kHeadDim, + int64_t kRopeDim, + bool kIsNeox, + bool kUsePDL, + typename DType, + typename CacheDType, + bool kRoundNormBeforeRope> +struct QKNormRopePackKVKernel { + template + static constexpr auto kernel = fused_qknorm_rope_warp< + kHeadDim, + kRopeDim, + kIsNeox, + kUsePDL, + DType, + CacheDType, + kRoundNormBeforeRope, + true, + IdType>; + + static void + run(const tvm::ffi::TensorView q, + const tvm::ffi::TensorView k, + const tvm::ffi::TensorView v, + const tvm::ffi::TensorView k_prefix, + const tvm::ffi::TensorView v_prefix, + const tvm::ffi::TensorView packed_k, + const tvm::ffi::TensorView packed_v, + const tvm::ffi::TensorView q_weight, + const tvm::ffi::TensorView k_weight, + const tvm::ffi::TensorView cos_sin_cache, + const tvm::ffi::TensorView positions, + int64_t batch_size, + int64_t prefix_tokens, + int64_t suffix_tokens, + float eps) { + using namespace host; + + auto N = SymbolicSize{"num_tokens"}; + auto NP = SymbolicSize{"num_prefix_tokens"}; + auto B = SymbolicSize{"batch_size"}; + auto T = SymbolicSize{"packed_tokens"}; + auto Q = SymbolicSize{"num_qo_heads"}; + auto K = SymbolicSize{"num_kv_heads"}; + auto D = SymbolicSize{"head_dim"}; + auto R = SymbolicSize{"rope_dim"}; + auto Dq = SymbolicSize{"q_stride"}; + auto Dk = SymbolicSize{"k_stride"}; + auto Dv = SymbolicSize{"v_stride"}; + auto Dkp = SymbolicSize{"k_prefix_stride"}; + auto Dvp = SymbolicSize{"v_prefix_stride"}; + auto Dd = SymbolicSize{"head_stride"}; + auto device = SymbolicDevice{}; + auto id_type = SymbolicDType{}; + N.set_value(batch_size * suffix_tokens); + NP.set_value(batch_size * prefix_tokens); + B.set_value(batch_size); + T.set_value(prefix_tokens + suffix_tokens); + D.set_value(kHeadDim); + R.set_value(kRopeDim); + device.set_options(); + + TensorMatcher({N, Q, D}).with_strides({Dq, Dd, 1}).with_dtype().with_device(device).verify(q); + TensorMatcher({N, K, D}).with_strides({Dk, Dd, 1}).with_dtype().with_device(device).verify(k); + TensorMatcher({N, K, D}).with_strides({Dv, Dd, 1}).with_dtype().with_device(device).verify(v); + TensorMatcher({NP, K, D}).with_strides({Dkp, Dd, 1}).with_dtype().with_device(device).verify(k_prefix); + TensorMatcher({NP, K, D}).with_strides({Dvp, Dd, 1}).with_dtype().with_device(device).verify(v_prefix); + TensorMatcher({B, T, K, D}).with_dtype().with_device(device).verify(packed_k).verify(packed_v); + RuntimeCheck(packed_k.is_contiguous(), "packed_k must be contiguous"); + RuntimeCheck(packed_v.is_contiguous(), "packed_v must be contiguous"); + TensorMatcher({D}).with_dtype().with_device(device).verify(q_weight).verify(k_weight); + TensorMatcher({-1, R}).with_dtype().with_device(device).verify(cos_sin_cache); + TensorMatcher({N}).with_dtype(id_type).with_device(device).verify(positions); + + const auto num_tokens = static_cast(N.unwrap()); + const auto num_qo_heads = static_cast(Q.unwrap()); + const auto num_kv_heads = static_cast(K.unwrap()); + if (num_tokens == 0 || (num_qo_heads == 0 && num_kv_heads == 0)) return; + const auto head_stride_bytes = static_cast(Dd.unwrap() * sizeof(DType)); + const int64_t k_offset = static_cast(num_qo_heads) * head_stride_bytes; + QKNormRopePackKVParams params{}; + params.q_ptr = q.data_ptr(); + params.k_ptr = pointer::offset(k.data_ptr(), -k_offset); + params.q_weight_ptr = q_weight.data_ptr(); + params.k_weight_ptr = k_weight.data_ptr(); + params.cos_sin_cache_ptr = cos_sin_cache.data_ptr(); + params.positions = positions.data_ptr(); + params.q_stride_bytes = static_cast(Dq.unwrap() * sizeof(DType)); + params.k_stride_bytes = static_cast(Dk.unwrap() * sizeof(DType)); + params.head_stride_bytes = head_stride_bytes; + params.num_qo_heads = num_qo_heads; + params.num_kv_heads = num_kv_heads; + params.num_tokens = num_tokens; + params.eps = eps; + params.v_ptr = v.data_ptr(); + params.k_prefix_ptr = k_prefix.data_ptr(); + params.v_prefix_ptr = v_prefix.data_ptr(); + params.packed_k_ptr = packed_k.data_ptr(); + params.packed_v_ptr = packed_v.data_ptr(); + params.v_stride_bytes = static_cast(Dv.unwrap() * sizeof(DType)); + params.k_prefix_stride_bytes = static_cast(Dkp.unwrap() * sizeof(DType)); + params.v_prefix_stride_bytes = static_cast(Dvp.unwrap() * sizeof(DType)); + params.packed_token_stride_bytes = static_cast(num_kv_heads * kHeadDim * sizeof(DType)); + params.packed_head_stride_bytes = static_cast(kHeadDim * sizeof(DType)); + params.batch_size = static_cast(batch_size); + params.prefix_tokens = static_cast(prefix_tokens); + params.suffix_tokens = static_cast(suffix_tokens); + + const auto is_int32 = id_type.is_type(); + const auto selected_kernel = is_int32 ? kernel : kernel; + const uint32_t kNumSM = runtime::get_sm_count(device.unwrap().device_id); + static const uint32_t kOccupancyTable[2] = { + runtime::get_blocks_per_sm(kernel, kThreadsPerBlock), + runtime::get_blocks_per_sm(kernel, kThreadsPerBlock), + }; + const auto max_blocks = kOccupancyTable[is_int32 ? 0 : 1] * kNumSM; + const uint32_t num_prefix_works = static_cast(batch_size * prefix_tokens) * num_kv_heads; + const uint32_t num_works = + (num_qo_heads + num_kv_heads) * num_tokens + 2 * num_prefix_works + num_tokens * num_kv_heads; + const auto needed_blocks = div_ceil(num_works, kWarpsPerBlock); + const auto num_blocks = std::min(max_blocks, needed_blocks); + LaunchKernel(num_blocks, kThreadsPerBlock, device.unwrap()).enable_pdl(kUsePDL)(selected_kernel, params); + } +}; + } // namespace sglang diff --git a/python/sglang/kernels/ops/diffusion/qknorm_rope.py b/python/sglang/kernels/ops/diffusion/qknorm_rope.py index cbf67640d..37b0ccb05 100644 --- a/python/sglang/kernels/ops/diffusion/qknorm_rope.py +++ b/python/sglang/kernels/ops/diffusion/qknorm_rope.py @@ -31,6 +31,7 @@ def _jit_qknorm_rope_module( dtype: torch.dtype, cache_dtype: torch.dtype, round_norm_before_rope: bool, + pack_kv: bool = False, ) -> Module: args = make_cpp_args( head_dim, @@ -41,23 +42,24 @@ def _jit_qknorm_rope_module( cache_dtype, round_norm_before_rope, ) + op_name = "qknorm_rope_pack_kv" if pack_kv else "qknorm_rope" + kernel_name = "QKNormRopePackKVKernel" if pack_kv else "QKNormRopeKernel" return load_jit( - "qknorm_rope", + op_name, *args, cuda_files=["diffusion/qknorm_rope.cuh"], - cuda_wrappers=[("qknorm_rope", f"QKNormRopeKernel<{args}>::run")], + cuda_wrappers=[(op_name, f"{kernel_name}<{args}>::run")], ) -@torch.compiler.assume_constant_result -@cache_once -def can_use_fused_inplace_qknorm_rope( +def _can_use_fused_qknorm_rope( head_dim: int, rope_dim: int, is_neox: bool, dtype: torch.dtype, - cache_dtype: torch.dtype = torch.float32, - round_norm_before_rope: bool = False, + cache_dtype: torch.dtype, + round_norm_before_rope: bool, + pack_kv: bool, ) -> bool: if dtype not in _SUPPORTED_DTYPES or cache_dtype not in _SUPPORTED_CACHE_DTYPES: logger.warning( @@ -106,13 +108,37 @@ def can_use_fused_inplace_qknorm_rope( dtype, cache_dtype, round_norm_before_rope, + pack_kv, ) return True except Exception as e: - logger.warning(f"Failed to load JIT fused QKNorm+RoPE kernel: {e}") + suffix = "+KV pack" if pack_kv else "" + logger.warning(f"Failed to load JIT fused QKNorm+RoPE{suffix} kernel: {e}") return False +@torch.compiler.assume_constant_result +@cache_once +def can_use_fused_inplace_qknorm_rope( + head_dim: int, + rope_dim: int, + is_neox: bool, + dtype: torch.dtype, + cache_dtype: torch.dtype = torch.float32, + round_norm_before_rope: bool = False, + pack_kv: bool = False, +) -> bool: + return _can_use_fused_qknorm_rope( + head_dim, + rope_dim, + is_neox, + dtype, + cache_dtype, + round_norm_before_rope, + pack_kv, + ) + + @register_custom_op(mutates_args=["q", "k"]) def fused_inplace_qknorm_rope( q: torch.Tensor, @@ -139,3 +165,54 @@ def fused_inplace_qknorm_rope( round_norm_before_rope, ) module.qknorm_rope(q, k, q_weight, k_weight, cos_sin_cache, positions, eps) + + +@register_custom_op(mutates_args=["q", "packed_kv"]) +def fused_qknorm_rope_pack_kv( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + k_prefix: torch.Tensor, + v_prefix: torch.Tensor, + packed_kv: torch.Tensor, + q_weight: torch.Tensor, + k_weight: torch.Tensor, + cos_sin_cache: torch.Tensor, + positions: torch.Tensor, + *, + is_neox: bool, + eps: float = 1e-6, + head_dim: int = 0, + rope_dim: int = 0, + round_norm_before_rope: bool = False, +) -> None: + head_dim = head_dim or q.size(-1) + rope_dim = rope_dim or cos_sin_cache.size(-1) + batch_size, suffix_tokens = q.shape[:2] + prefix_tokens = k_prefix.shape[1] + module = _jit_qknorm_rope_module( + head_dim, + rope_dim, + is_neox, + q.dtype, + cos_sin_cache.dtype, + round_norm_before_rope, + True, + ) + module.qknorm_rope_pack_kv( + q.view(-1, q.shape[-2], head_dim), + k.view(-1, k.shape[-2], head_dim), + v.view(-1, v.shape[-2], head_dim), + k_prefix.view(-1, k_prefix.shape[-2], head_dim), + v_prefix.view(-1, v_prefix.shape[-2], head_dim), + packed_kv[0], + packed_kv[1], + q_weight, + k_weight, + cos_sin_cache, + positions, + batch_size, + prefix_tokens, + suffix_tokens, + eps, + ) diff --git a/python/sglang/multimodal_gen/runtime/layers/layernorm.py b/python/sglang/multimodal_gen/runtime/layers/layernorm.py index d778e1553..a2202bec4 100755 --- a/python/sglang/multimodal_gen/runtime/layers/layernorm.py +++ b/python/sglang/multimodal_gen/runtime/layers/layernorm.py @@ -972,8 +972,13 @@ def apply_qk_norm_rope( positions: Optional[torch.Tensor] = None, position_offset: int = 0, allow_inplace: bool = True, + allow_strided_qk: bool = False, ) -> Tuple[torch.Tensor, torch.Tensor]: - """Apply QK RMSNorm followed by RoPE, fusing both on supported CUDA/XPU shapes.""" + """Apply QK RMSNorm followed by RoPE, fusing supported CUDA/XPU shapes. + + Strided packed-QKV views require an explicit opt-in because selecting the fused + kernel changes the numerical path for models that historically used the fallback. + """ from sglang.multimodal_gen.runtime.layers.rotary_embedding import ( apply_flashinfer_rope_qk_inplace, @@ -1010,6 +1015,17 @@ def apply_qk_norm_rope( "off", "no", } + q_has_supported_layout = q.is_contiguous() + k_has_supported_layout = k.is_contiguous() + if allow_strided_qk: + q_has_supported_layout = ( + q.stride(-1) == 1 + and q.stride(-2) == k.stride(-2) + and q.stride(0) == seq_len * q.stride(1) + ) + k_has_supported_layout = k.stride(-1) == 1 and k.stride( + 0 + ) == seq_len * k.stride(1) if positions is None: pos_1d = torch.arange( @@ -1033,15 +1049,16 @@ def apply_qk_norm_rope( and allow_inplace and (q_eps == k_eps) and q.dtype in (torch.float16, torch.bfloat16) + and k.dtype == q.dtype and q_norm.weight.dtype == q.dtype and k_norm.weight.dtype == k.dtype - and q.is_contiguous() - and k.is_contiguous() + and q_has_supported_layout + and k_has_supported_layout and can_use_fused_inplace_qknorm_rope(head_dim, rope_dim, is_neox, q.dtype) ): fused_inplace_qknorm_rope( - q=q.reshape(-1, q.shape[-2], head_dim), - k=k.reshape(-1, k.shape[-2], head_dim), + q=q.view(-1, q.shape[-2], head_dim), + k=k.view(-1, k.shape[-2], head_dim), q_weight=q_norm.weight, k_weight=k_norm.weight, cos_sin_cache=cos_sin_cache, @@ -1053,8 +1070,6 @@ def apply_qk_norm_rope( ) return q, k - # TODO: Once CUDA fused_inplace_qknorm_rope supports last-dimension-contiguous q/k, - # merge this path with the CUDA fused qknorm+rope branch. if ( _is_xpu and allow_inplace diff --git a/python/sglang/multimodal_gen/runtime/models/dits/cosmos3video.py b/python/sglang/multimodal_gen/runtime/models/dits/cosmos3video.py index ee0916727..6ff98af59 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/cosmos3video.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/cosmos3video.py @@ -14,6 +14,10 @@ import torch import torch.nn as nn import torch.nn.functional as F +from sglang.kernels.ops.diffusion.qknorm_rope import ( + can_use_fused_inplace_qknorm_rope, + fused_qknorm_rope_pack_kv, +) from sglang.multimodal_gen.configs.models.dits.cosmos3video import Cosmos3VideoConfig from sglang.multimodal_gen.configs.models.fsdp import is_module_list_entry_in from sglang.multimodal_gen.runtime.distributed import ( @@ -229,17 +233,60 @@ def _apply_qwen3_qk_norm_rope( rope_cache_positions: torch.Tensor, ) -> tuple[torch.Tensor, torch.Tensor]: return apply_qk_norm_rope( - q=q.contiguous(), - k=k.contiguous(), + q=q, + k=k, q_norm=q_norm, k_norm=k_norm, head_dim=head_dim, cos_sin_cache=cos_sin_cache, is_neox=True, positions=rope_cache_positions, + allow_strided_qk=True, ) +def _apply_qwen3_qk_norm_rope_pack_kv( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + k_prefix: torch.Tensor, + v_prefix: torch.Tensor, + q_norm: RMSNorm, + k_norm: RMSNorm, + head_dim: int, + cos_sin_cache: torch.Tensor, + rope_cache_positions: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + batch_size, suffix_tokens, _, _ = q.shape + prefix_tokens = k_prefix.shape[1] + packed_kv = torch.empty( + 2, + batch_size, + prefix_tokens + suffix_tokens, + k.shape[2], + head_dim, + dtype=q.dtype, + device=q.device, + ) + fused_qknorm_rope_pack_kv( + q, + k, + v, + k_prefix, + v_prefix, + packed_kv, + q_norm.weight, + k_norm.weight, + cos_sin_cache, + rope_cache_positions, + is_neox=True, + eps=q_norm.variance_epsilon, + head_dim=head_dim, + rope_dim=cos_sin_cache.shape[-1], + ) + return q, packed_kv[0], packed_kv[1] + + def _apply_qwen3_rope_from_cache( q: torch.Tensor, k: torch.Tensor, cos_sin_cache: torch.Tensor ) -> tuple[torch.Tensor, torch.Tensor]: @@ -733,7 +780,39 @@ class Cosmos3CrossAttention(nn.Module): :, ] - if use_fused_qk_norm_rope: + use_fused_kv_pack = ( + use_fused_qk_norm_rope + and q.device.type == "cuda" + and not torch.compiler.is_compiling() + and get_sp_world_size() == 1 + and q.dtype == k.dtype == v.dtype == k_und.dtype == v_und.dtype + and self.norm_q.weight.dtype == q.dtype + and self.norm_k.weight.dtype == k.dtype + and self.norm_q.variance_epsilon == self.norm_k.variance_epsilon + and can_use_fused_inplace_qknorm_rope( + self.head_dim, + cos_sin_cache.shape[-1], + True, + q.dtype, + cos_sin_cache.dtype, + pack_kv=True, + ) + ) + if use_fused_kv_pack: + q, packed_k, packed_v = _apply_qwen3_qk_norm_rope_pack_kv( + q, + k, + v, + k_und, + v_und, + self.norm_q, + self.norm_k, + self.head_dim, + cos_sin_cache, + rope_cache_positions, + ) + out = self.attn.forward(q, packed_k, packed_v) + elif use_fused_qk_norm_rope: q, k = _apply_qwen3_qk_norm_rope( q, k, @@ -747,11 +826,10 @@ class Cosmos3CrossAttention(nn.Module): q, k = _apply_qwen3_qk_norm_rope_split( q, k, self.norm_q, self.norm_k, self.head_dim, cos_sin_cache ) - - # K/V = [text (replicated on every SP rank) | image (sharded same as Q)]. - # USPAttention routes through the registered attention backend (FA, sage, - # …) and handles the Ulysses all-to-all when SP > 1. - out = self.attn.forward_with_replicated_kv_prefix(q, k_und, v_und, k, v) + if not use_fused_kv_pack: + # K/V = [UND prefix (replicated on SP ranks) | GEN suffix]. + # USPAttention applies the configured backend and SP collectives. + out = self.attn.forward_with_replicated_kv_prefix(q, k_und, v_und, k, v) out = out.reshape(batch_size, seq_len_gen, -1) out, _ = self.to_out(out) return out diff --git a/test/registered/kernels/ops/diffusion/test_qknorm_rope.py b/test/registered/kernels/ops/diffusion/test_qknorm_rope.py index 264d1c0bd..87e3cc462 100644 --- a/test/registered/kernels/ops/diffusion/test_qknorm_rope.py +++ b/test/registered/kernels/ops/diffusion/test_qknorm_rope.py @@ -216,6 +216,164 @@ def test_qknorm_rope_preserves_split_bf16_rounding() -> None: assert torch.equal(k_ref, k_fused) +def test_qknorm_rope_requires_opt_in_for_strided_packed_gqa() -> None: + from sglang.kernels.ops.diffusion.qknorm_rope import ( + fused_inplace_qknorm_rope, + ) + from sglang.multimodal_gen.runtime.layers.layernorm import ( + RMSNorm, + apply_qk_norm_rope, + ) + + num_tokens, num_q_heads, num_kv_heads, head_dim = 257, 32, 8, 128 + num_heads = num_q_heads + 2 * num_kv_heads + qkv = torch.randn(1, num_tokens, num_heads, head_dim, device=DEVICE, dtype=DTYPE) + q_weight = torch.randn(head_dim, device=DEVICE, dtype=DTYPE) + k_weight = torch.randn(head_dim, device=DEVICE, dtype=DTYPE) + positions = torch.arange(num_tokens, device=DEVICE, dtype=torch.int64) + cos_sin_cache = create_cos_sin_cache(head_dim, num_tokens) + + q_ref = qkv[:, :, :num_q_heads].contiguous() + k_ref = qkv[:, :, num_q_heads : num_q_heads + num_kv_heads].contiguous() + q_norm = RMSNorm(head_dim, eps=1e-6).to(device=DEVICE, dtype=DTYPE) + k_norm = RMSNorm(head_dim, eps=1e-6).to(device=DEVICE, dtype=DTYPE) + q_norm.weight.data.copy_(q_weight) + k_norm.weight.data.copy_(k_weight) + qkv_default = qkv.clone() + q_default = qkv_default[:, :, :num_q_heads] + k_default = qkv_default[:, :, num_q_heads : num_q_heads + num_kv_heads] + q_default_out, k_default_out = apply_qk_norm_rope( + q=q_default, + k=k_default, + q_norm=q_norm, + k_norm=k_norm, + head_dim=head_dim, + cos_sin_cache=cos_sin_cache, + is_neox=True, + positions=positions, + ) + assert q_default_out.data_ptr() != q_default.data_ptr() + assert k_default_out.data_ptr() != k_default.data_ptr() + + qkv_fused = qkv.clone() + q_fused = qkv_fused[:, :, :num_q_heads] + k_fused = qkv_fused[:, :, num_q_heads : num_q_heads + num_kv_heads] + v_before = qkv_fused[:, :, num_q_heads + num_kv_heads :].clone() + + fused_inplace_qknorm_rope( + q_ref.view(-1, num_q_heads, head_dim), + k_ref.view(-1, num_kv_heads, head_dim), + q_weight, + k_weight, + cos_sin_cache, + positions, + is_neox=True, + rope_dim=head_dim, + ) + q_out, k_out = apply_qk_norm_rope( + q=q_fused, + k=k_fused, + q_norm=q_norm, + k_norm=k_norm, + head_dim=head_dim, + cos_sin_cache=cos_sin_cache, + is_neox=True, + positions=positions, + allow_strided_qk=True, + ) + + assert q_out.data_ptr() == q_fused.data_ptr() + assert k_out.data_ptr() == k_fused.data_ptr() + assert torch.equal(q_ref, q_out) + assert torch.equal(k_ref, k_out) + assert torch.equal(v_before, qkv_fused[:, :, num_q_heads + num_kv_heads :]) + + +def test_qknorm_rope_pack_kv_matches_separate_ops() -> None: + from sglang.kernels.ops.diffusion.qknorm_rope import ( + fused_inplace_qknorm_rope, + fused_qknorm_rope_pack_kv, + ) + + batch_size = 2 + prefix_tokens, suffix_tokens = 17, 257 + num_q_heads, num_kv_heads, head_dim = 32, 8, 128 + num_heads = num_q_heads + 2 * num_kv_heads + qkv = torch.randn( + batch_size, + suffix_tokens, + num_heads, + head_dim, + device=DEVICE, + dtype=DTYPE, + ) + prefix_qkv = torch.randn( + batch_size, + prefix_tokens, + num_heads, + head_dim, + device=DEVICE, + dtype=DTYPE, + ) + k_prefix = prefix_qkv[:, :, num_q_heads : num_q_heads + num_kv_heads] + v_prefix = prefix_qkv[:, :, num_q_heads + num_kv_heads :] + q_weight = torch.randn(head_dim, device=DEVICE, dtype=DTYPE) + k_weight = torch.randn(head_dim, device=DEVICE, dtype=DTYPE) + positions = torch.arange( + batch_size * suffix_tokens, device=DEVICE, dtype=torch.int64 + ) + cos_sin_cache = create_cos_sin_cache(head_dim, batch_size * suffix_tokens) + + qkv_ref = qkv.clone() + q_ref = qkv_ref[:, :, :num_q_heads] + k_ref = qkv_ref[:, :, num_q_heads : num_q_heads + num_kv_heads] + v_ref = qkv_ref[:, :, num_q_heads + num_kv_heads :] + fused_inplace_qknorm_rope( + q_ref.view(-1, num_q_heads, head_dim), + k_ref.view(-1, num_kv_heads, head_dim), + q_weight, + k_weight, + cos_sin_cache, + positions, + is_neox=True, + rope_dim=head_dim, + ) + packed_k_ref = torch.cat([k_prefix, k_ref], dim=1) + packed_v_ref = torch.cat([v_prefix, v_ref], dim=1) + + qkv_fused = qkv.clone() + q_fused = qkv_fused[:, :, :num_q_heads] + k_fused = qkv_fused[:, :, num_q_heads : num_q_heads + num_kv_heads] + v_fused = qkv_fused[:, :, num_q_heads + num_kv_heads :] + packed_kv = torch.empty( + 2, + batch_size, + prefix_tokens + suffix_tokens, + num_kv_heads, + head_dim, + device=DEVICE, + dtype=DTYPE, + ) + fused_qknorm_rope_pack_kv( + q_fused, + k_fused, + v_fused, + k_prefix, + v_prefix, + packed_kv, + q_weight, + k_weight, + cos_sin_cache, + positions, + is_neox=True, + rope_dim=head_dim, + ) + + assert torch.equal(q_ref, q_fused) + assert torch.equal(packed_k_ref, packed_kv[0]) + assert torch.equal(packed_v_ref, packed_kv[1]) + + def test_qknorm_rope_accepts_empty_token_dimension() -> None: from sglang.kernels.ops.diffusion.qknorm_rope import fused_inplace_qknorm_rope