minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[diffusion] optimize: fuse cosmos qk norm, rope, and kv packing (#34275)

Browse Source
This commit is contained in:
Mick
2026-08-12 23:18:16 +08:00
committed by GitHub
parent f28bc5a6de
commit ad47dde65c
5 changed files with 570 additions and 30 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/kernels/jit/csrc/diffusion/qknorm_rope.cuh
+219 -7
View File
@@ -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 <bool kPackKV>
using QKNormRopeParamsT = std::conditional_t<kPackKV, QKNormRopePackKVParams, QKNormRopeParams>;
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<kPackKV> __grid_constant__ params) {
using namespace device;
static_assert(std::is_same_v<DType, fp16_t> || std::is_same_v<DType, bf16_t>);
@@ -181,23 +201,81 @@ __global__ void fused_qknorm_rope_warp(const QKNormRopeParams __grid_constant__
using Storage = AlignedVector<Packed, kVecSize>;
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<const QKNormRopeParams&>(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<kUsePDL>();
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<Storage>(input, lane_id);
store_as<Storage>(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<void*>(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<Storage>(input, lane_id);
@@ -246,7 +324,7 @@ __global__ void fused_qknorm_rope_warp(const QKNormRopeParams __grid_constant__
}
}
}
store_as<Storage>(const_cast<void*>(input), output_vec, lane_id);
store_as<Storage>(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<Packed, fp32x2_t>({elems[2 * j], elems[2 * j + 1]});
}
store_as<Storage>(const_cast<void*>(input), input_vec, lane_id);
store_as<Storage>(output, input_vec, lane_id);
}
PDLTriggerSecondary<kUsePDL>();
@@ -332,8 +410,16 @@ template <
struct QKNormRopeKernel {
static_assert(kHeadDim <= 256, "Only head_dim <= 256 is supported");
template <typename IdType>
static constexpr auto kernel =
fused_qknorm_rope_warp<kHeadDim, kRopeDim, kIsNeox, kUsePDL, DType, CacheDType, kRoundNormBeforeRope, IdType>;
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 <typename IdType>
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<kDLCUDA>();
TensorMatcher({N, Q, D}).with_strides({Dq, Dd, 1}).with_dtype<DType>().with_device(device).verify(q);
TensorMatcher({N, K, D}).with_strides({Dk, Dd, 1}).with_dtype<DType>().with_device(device).verify(k);
TensorMatcher({N, K, D}).with_strides({Dv, Dd, 1}).with_dtype<DType>().with_device(device).verify(v);
TensorMatcher({NP, K, D}).with_strides({Dkp, Dd, 1}).with_dtype<DType>().with_device(device).verify(k_prefix);
TensorMatcher({NP, K, D}).with_strides({Dvp, Dd, 1}).with_dtype<DType>().with_device(device).verify(v_prefix);
TensorMatcher({B, T, K, D}).with_dtype<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<DType>().with_device(device).verify(q_weight).verify(k_weight);
TensorMatcher({-1, R}).with_dtype<CacheDType>().with_device(device).verify(cos_sin_cache);
TensorMatcher({N}).with_dtype<int32_t, int64_t>(id_type).with_device(device).verify(positions);
const auto num_tokens = static_cast<uint32_t>(N.unwrap());
const auto num_qo_heads = static_cast<uint32_t>(Q.unwrap());
const auto num_kv_heads = static_cast<uint32_t>(K.unwrap());
if (num_tokens == 0 || (num_qo_heads == 0 && num_kv_heads == 0)) return;
const auto head_stride_bytes = static_cast<int64_t>(Dd.unwrap() * sizeof(DType));
const int64_t k_offset = static_cast<int64_t>(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<int64_t>(Dq.unwrap() * sizeof(DType));
params.k_stride_bytes = static_cast<int64_t>(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<int64_t>(Dv.unwrap() * sizeof(DType));
params.k_prefix_stride_bytes = static_cast<int64_t>(Dkp.unwrap() * sizeof(DType));
params.v_prefix_stride_bytes = static_cast<int64_t>(Dvp.unwrap() * sizeof(DType));
params.packed_token_stride_bytes = static_cast<int64_t>(num_kv_heads * kHeadDim * sizeof(DType));
params.packed_head_stride_bytes = static_cast<int64_t>(kHeadDim * sizeof(DType));
params.batch_size = static_cast<uint32_t>(batch_size);
params.prefix_tokens = static_cast<uint32_t>(prefix_tokens);
params.suffix_tokens = static_cast<uint32_t>(suffix_tokens);
const auto is_int32 = id_type.is_type<int32_t>();
const auto selected_kernel = is_int32 ? kernel<int32_t> : kernel<int64_t>;
const uint32_t kNumSM = runtime::get_sm_count(device.unwrap().device_id);
static const uint32_t kOccupancyTable[2] = {
runtime::get_blocks_per_sm(kernel<int32_t>, kThreadsPerBlock),
runtime::get_blocks_per_sm(kernel<int64_t>, kThreadsPerBlock),
};
const auto max_blocks = kOccupancyTable[is_int32 ? 0 : 1] * kNumSM;
const uint32_t num_prefix_works = static_cast<uint32_t>(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
python/sglang/kernels/ops/diffusion/qknorm_rope.py
+85 -8
View File
@@ -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,
)
python/sglang/multimodal_gen/runtime/layers/layernorm.py
+22 -7
View File
@@ -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
python/sglang/multimodal_gen/runtime/models/dits/cosmos3video.py
+86 -8
View File
@@ -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
test/registered/kernels/ops/diffusion/test_qknorm_rope.py
+158
View File
@@ -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