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

[Diffusion][ERNIE] Fuse QKNorm with full-width RoPE (#34620)

Browse Source
This commit is contained in:
Xiaoyu Zhang
2026-08-13 23:23:21 +08:00
committed by GitHub
parent 82f7afb881
commit ebca0bbde4
6 changed files with 293 additions and 21 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/kernels/jit/csrc/diffusion/qknorm_rope.cuh
+24 -13
View File
@@ -172,6 +172,7 @@ template <
typename CacheDType,
bool kRoundNormBeforeRope,
bool kPackKV,
bool kCacheHasFullWidth,
typename IdType>
__global__ void fused_qknorm_rope_warp(const QKNormRopeParamsT<kPackKV> __grid_constant__ params) {
using namespace device;
@@ -185,7 +186,8 @@ __global__ void fused_qknorm_rope_warp(const QKNormRopeParamsT<kPackKV> __grid_c
constexpr uint32_t kRotaryLanes = kRopeDim / kElemsPerThread;
constexpr uint32_t kHalfRotaryLanes = kRotaryLanes / 2;
constexpr uint32_t kActiveMask = active_mask<kRotaryLanes>();
constexpr int64_t kCosSinStrideBytes = kRopeDim * sizeof(CacheDType);
constexpr int64_t kCacheRotaryDim = kCacheHasFullWidth ? 2 * kRopeDim : kRopeDim;
constexpr int64_t kCosSinStrideBytes = kCacheRotaryDim * sizeof(CacheDType);
static_assert(kElemsPerThread % 2 == 0, "Each lane must own an even number of elements");
static_assert(kRopeDim > 0 && kRopeDim <= kHeadDim, "Invalid rope dimension");
@@ -285,7 +287,7 @@ __global__ void fused_qknorm_rope_warp(const QKNormRopeParamsT<kPackKV> __grid_c
auto output_vec = norm::apply_norm_warp<kHeadDim>(input_vec, weight_vec, eps);
const auto pos = static_cast<int64_t>(static_cast<const IdType*>(positions)[token_id]);
const auto cos_ptr = static_cast<const CacheDType*>(pointer::offset(cos_sin_cache_ptr, pos * kCosSinStrideBytes));
const auto sin_ptr = cos_ptr + kRopeDim / 2;
const auto sin_ptr = cos_ptr + (kCacheHasFullWidth ? kRopeDim : kRopeDim / 2);
if constexpr (kIsNeox) {
if (lane_id < kRotaryLanes) {
@@ -301,9 +303,10 @@ __global__ void fused_qknorm_rope_warp(const QKNormRopeParamsT<kPackKV> __grid_c
const auto& partner_values = unpack(partner_vec);
#pragma unroll
for (uint32_t i = 0; i < 2; ++i) {
const auto half_idx = (lane_id % kHalfRotaryLanes) * kElemsPerThread + 2 * j + i;
const auto cos = load_cache_value(cos_ptr, half_idx);
const auto sin = load_cache_value(sin_ptr, half_idx);
const auto cache_idx =
(kCacheHasFullWidth ? lane_id : lane_id % kHalfRotaryLanes) * kElemsPerThread + 2 * j + i;
const auto cos = load_cache_value(cos_ptr, cache_idx);
const auto sin = load_cache_value(sin_ptr, cache_idx);
values[i] = lane_id < kHalfRotaryLanes ? rotary_sub(values[i], cos, partner_values[i], sin)
: rotary_add(values[i], cos, partner_values[i], sin);
}
@@ -354,7 +357,7 @@ __global__ void fused_qknorm_rope_warp(const QKNormRopeParamsT<kPackKV> __grid_c
const auto pos = static_cast<int64_t>(static_cast<const IdType*>(positions)[token_id]);
const auto cos_ptr =
static_cast<const CacheDType*>(pointer::offset(cos_sin_cache_ptr, pos * kCosSinStrideBytes));
const auto sin_ptr = cos_ptr + kRopeDim / 2;
const auto sin_ptr = cos_ptr + (kCacheHasFullWidth ? kRopeDim : kRopeDim / 2);
const auto partner_lane = lane_id < kHalfRotaryLanes ? lane_id + kHalfRotaryLanes : lane_id - kHalfRotaryLanes;
#pragma unroll
@@ -363,9 +366,9 @@ __global__ void fused_qknorm_rope_warp(const QKNormRopeParamsT<kPackKV> __grid_c
if (lane_id < kHalfRotaryLanes) {
swapped = -swapped;
}
const auto half_idx = (lane_id % kHalfRotaryLanes) * kElemsPerThread + i;
const float cos = cast<fp32_t>(load_cache_value(cos_ptr, half_idx));
const float sin = cast<fp32_t>(load_cache_value(sin_ptr, half_idx));
const auto cache_idx = (kCacheHasFullWidth ? lane_id : lane_id % kHalfRotaryLanes) * kElemsPerThread + i;
const float cos = cast<fp32_t>(load_cache_value(cos_ptr, cache_idx));
const float sin = cast<fp32_t>(load_cache_value(sin_ptr, cache_idx));
elems[i] = elems[i] * cos + swapped * sin;
}
}
@@ -374,7 +377,7 @@ __global__ void fused_qknorm_rope_warp(const QKNormRopeParamsT<kPackKV> __grid_c
const auto pos = static_cast<int64_t>(static_cast<const IdType*>(positions)[token_id]);
const auto cos_ptr =
static_cast<const CacheDType*>(pointer::offset(cos_sin_cache_ptr, pos * kCosSinStrideBytes));
const auto sin_ptr = cos_ptr + kRopeDim / 2;
const auto sin_ptr = cos_ptr + (kCacheHasFullWidth ? kRopeDim : kRopeDim / 2);
#pragma unroll
for (uint32_t i = 0; i < kElemsPerThread; i += 2) {
@@ -406,7 +409,8 @@ template <
bool kUsePDL,
typename DType,
typename CacheDType,
bool kRoundNormBeforeRope>
bool kRoundNormBeforeRope,
bool kCacheHasFullWidth>
struct QKNormRopeKernel {
static_assert(kHeadDim <= 256, "Only head_dim <= 256 is supported");
template <typename IdType>
@@ -419,6 +423,7 @@ struct QKNormRopeKernel {
CacheDType,
kRoundNormBeforeRope,
false,
kCacheHasFullWidth,
IdType>;
static void
@@ -448,7 +453,10 @@ struct QKNormRopeKernel {
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({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({-1, kCacheHasFullWidth ? 2 * kRopeDim : kRopeDim})
.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());
@@ -498,8 +506,10 @@ template <
bool kUsePDL,
typename DType,
typename CacheDType,
bool kRoundNormBeforeRope>
bool kRoundNormBeforeRope,
bool kCacheHasFullWidth>
struct QKNormRopePackKVKernel {
static_assert(!kCacheHasFullWidth, "KV packing does not support full-width cos/sin caches");
template <typename IdType>
static constexpr auto kernel = fused_qknorm_rope_warp<
kHeadDim,
@@ -510,6 +520,7 @@ struct QKNormRopePackKVKernel {
CacheDType,
kRoundNormBeforeRope,
true,
kCacheHasFullWidth,
IdType>;
static void
python/sglang/kernels/ops/diffusion/qknorm_rope.py
+19 -1
View File
@@ -32,6 +32,7 @@ def _jit_qknorm_rope_module(
cache_dtype: torch.dtype,
round_norm_before_rope: bool,
pack_kv: bool = False,
cache_has_full_width: bool = False,
) -> Module:
args = make_cpp_args(
head_dim,
@@ -41,6 +42,7 @@ def _jit_qknorm_rope_module(
dtype,
cache_dtype,
round_norm_before_rope,
cache_has_full_width,
)
op_name = "qknorm_rope_pack_kv" if pack_kv else "qknorm_rope"
kernel_name = "QKNormRopePackKVKernel" if pack_kv else "QKNormRopeKernel"
@@ -60,6 +62,7 @@ def _can_use_fused_qknorm_rope(
cache_dtype: torch.dtype,
round_norm_before_rope: bool,
pack_kv: bool,
cache_has_full_width: bool,
) -> bool:
if dtype not in _SUPPORTED_DTYPES or cache_dtype not in _SUPPORTED_CACHE_DTYPES:
logger.warning(
@@ -93,6 +96,12 @@ def _can_use_fused_qknorm_rope(
rotary_lanes,
)
return False
elif cache_has_full_width:
logger.warning("Full-width cos/sin caches are only supported for NeoX RoPE")
return False
if pack_kv and cache_has_full_width:
logger.warning("KV packing does not support full-width cos/sin caches")
return False
if round_norm_before_rope and cache_dtype != dtype:
logger.warning(
"Exact fused QKNorm+RoPE requires cache dtype %s to match activation dtype %s",
@@ -109,6 +118,7 @@ def _can_use_fused_qknorm_rope(
cache_dtype,
round_norm_before_rope,
pack_kv,
cache_has_full_width,
)
return True
except Exception as e:
@@ -127,6 +137,7 @@ def can_use_fused_inplace_qknorm_rope(
cache_dtype: torch.dtype = torch.float32,
round_norm_before_rope: bool = False,
pack_kv: bool = False,
cache_has_full_width: bool = False,
) -> bool:
return _can_use_fused_qknorm_rope(
head_dim,
@@ -136,6 +147,7 @@ def can_use_fused_inplace_qknorm_rope(
cache_dtype,
round_norm_before_rope,
pack_kv,
cache_has_full_width,
)
@@ -153,9 +165,12 @@ def fused_inplace_qknorm_rope(
head_dim: int = 0,
rope_dim: int = 0,
round_norm_before_rope: bool = False,
cache_has_full_width: bool = False,
) -> None:
head_dim = head_dim or q.size(-1)
rope_dim = rope_dim or cos_sin_cache.size(-1)
if not rope_dim:
cache_width = cos_sin_cache.size(-1)
rope_dim = cache_width // 2 if cache_has_full_width else cache_width
module = _jit_qknorm_rope_module(
head_dim,
rope_dim,
@@ -163,6 +178,8 @@ def fused_inplace_qknorm_rope(
q.dtype,
cos_sin_cache.dtype,
round_norm_before_rope,
False,
cache_has_full_width,
)
module.qknorm_rope(q, k, q_weight, k_weight, cos_sin_cache, positions, eps)
@@ -198,6 +215,7 @@ def fused_qknorm_rope_pack_kv(
cos_sin_cache.dtype,
round_norm_before_rope,
True,
False,
)
module.qknorm_rope_pack_kv(
q.view(-1, q.shape[-2], head_dim),
python/sglang/multimodal_gen/runtime/layers/layernorm.py
+24 -2
View File
@@ -973,11 +973,15 @@ def apply_qk_norm_rope(
position_offset: int = 0,
allow_inplace: bool = True,
allow_strided_qk: bool = False,
round_norm_before_rope: bool = False,
cache_has_full_width: bool = False,
) -> Tuple[torch.Tensor, torch.Tensor]:
"""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.
``cache_has_full_width`` describes ``[full cos, full sin]`` cache rows and
requires the fused CUDA path; the ordinary cache stores half-width cos/sin.
"""
from sglang.multimodal_gen.runtime.layers.rotary_embedding import (
@@ -1004,7 +1008,12 @@ def apply_qk_norm_rope(
batch_size, seq_len, _, _ = q.shape
q_eps = q_norm.variance_epsilon
k_eps = k_norm.variance_epsilon
rope_dim = cos_sin_cache.size(-1)
cache_width = cos_sin_cache.size(-1)
if cache_has_full_width and cache_width % 2:
raise ValueError(
f"full-width cos/sin cache must have even width, got {cache_width}"
)
rope_dim = cache_width // 2 if cache_has_full_width else cache_width
if rope_dim % 2 != 0 or rope_dim > head_dim:
raise ValueError(
f"cos_sin_cache width must be even and <= head_dim, got {rope_dim} vs {head_dim}"
@@ -1054,7 +1063,15 @@ def apply_qk_norm_rope(
and k_norm.weight.dtype == k.dtype
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)
and can_use_fused_inplace_qknorm_rope(
head_dim=head_dim,
rope_dim=rope_dim,
is_neox=is_neox,
dtype=q.dtype,
cache_dtype=cos_sin_cache.dtype,
round_norm_before_rope=round_norm_before_rope,
cache_has_full_width=cache_has_full_width,
)
):
fused_inplace_qknorm_rope(
q=q.view(-1, q.shape[-2], head_dim),
@@ -1067,9 +1084,14 @@ def apply_qk_norm_rope(
eps=q_eps,
head_dim=head_dim,
rope_dim=rope_dim,
round_norm_before_rope=round_norm_before_rope,
cache_has_full_width=cache_has_full_width,
)
return q, k
if cache_has_full_width:
raise RuntimeError("full-width cos/sin cache requires fused QKNorm+RoPE")
if (
_is_xpu
and allow_inplace
python/sglang/multimodal_gen/runtime/models/dits/ernie_image.py
+109 -5
View File
@@ -49,7 +49,11 @@ from sglang.multimodal_gen.runtime.layers.attention.layer import (
USPAttention,
build_varlen_mask_meta,
)
from sglang.multimodal_gen.runtime.layers.layernorm import RMSNorm, apply_qk_norm
from sglang.multimodal_gen.runtime.layers.layernorm import (
RMSNorm,
apply_qk_norm,
apply_qk_norm_rope,
)
from sglang.multimodal_gen.runtime.layers.linear import (
ColumnParallelLinear,
MergedColumnParallelLinear,
@@ -68,6 +72,7 @@ logger = init_logger(__name__)
_ERNIE_NORM = BitExactFusionGate("ERNIE fused-norm")
_ERNIE_GATED_NORM = BitExactFusionGate("ERNIE fused gated-norm")
_ERNIE_ROPE = BitExactFusionGate("ERNIE fused RoPE")
_ERNIE_QKNORM_ROPE = BitExactFusionGate("ERNIE fused QKNorm+RoPE")
_ERNIE_GEGLU = BitExactFusionGate("ERNIE fused GELU-mul")
@@ -284,6 +289,8 @@ class ErnieImageSelfAttention(nn.Module):
x: torch.Tensor,
rope_cos: torch.Tensor,
rope_sin: torch.Tensor,
rope_cache: torch.Tensor,
rope_positions: torch.Tensor,
attn_mask: torch.Tensor | None = None,
attn_mask_meta: dict | None = None,
) -> torch.Tensor:
@@ -298,16 +305,20 @@ class ErnieImageSelfAttention(nn.Module):
v = v.view(B, S, self.num_local_heads, self.head_dim)
if self.qk_layernorm:
q, k = apply_qk_norm(
q, k = _ernie_qknorm_rope(
q,
k,
self.norm_q,
self.norm_k,
self.head_dim,
rope_cos,
rope_sin,
rope_cache,
rope_positions,
)
q = _ernie_rope(q, rope_cos, rope_sin)
k = _ernie_rope(k, rope_cos, rope_sin)
else:
q = _ernie_rope(q, rope_cos, rope_sin)
k = _ernie_rope(k, rope_cos, rope_sin)
attn_out = self.attn(
q, k, v, attn_mask=attn_mask, attn_mask_meta=attn_mask_meta
@@ -378,6 +389,8 @@ class ErnieImageSharedAdaLNBlock(nn.Module):
x: torch.Tensor,
rope_cos: torch.Tensor,
rope_sin: torch.Tensor,
rope_cache: torch.Tensor,
rope_positions: torch.Tensor,
shift_msa: torch.Tensor,
scale_msa: torch.Tensor,
gate_msa: torch.Tensor,
@@ -393,6 +406,8 @@ class ErnieImageSharedAdaLNBlock(nn.Module):
x,
rope_cos,
rope_sin,
rope_cache,
rope_positions,
attn_mask=attn_mask,
attn_mask_meta=attn_mask_meta,
)
@@ -472,6 +487,89 @@ def _ernie_rope(
return _apply_rotary_bshd_eager(x, cos_, sin_)
def _ernie_qknorm_rope_reference(
q: torch.Tensor,
k: torch.Tensor,
q_norm: RMSNorm,
k_norm: RMSNorm,
head_dim: int,
rope_cos: torch.Tensor,
rope_sin: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor]:
q, k = apply_qk_norm(q, k, q_norm, k_norm, head_dim)
return _ernie_rope(q, rope_cos, rope_sin), _ernie_rope(k, rope_cos, rope_sin)
def _ernie_qknorm_rope(
q: torch.Tensor,
k: torch.Tensor,
q_norm: RMSNorm,
k_norm: RMSNorm,
head_dim: int,
rope_cos: torch.Tensor,
rope_sin: torch.Tensor,
rope_cache: torch.Tensor,
rope_positions: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor]:
"""Fuse ERNIE's QK RMSNorm and rotate-half RoPE without changing bits."""
verified = _ERNIE_QKNORM_ROPE.verified
if not _ERNIE_QKNORM_ROPE.disabled and (
verified or _ERNIE_QKNORM_ROPE.can_attempt_once()
):
q_input = q.clone() if not verified else q
k_input = k.clone() if not verified else k
try:
out = apply_qk_norm_rope(
q=q,
k=k,
q_norm=q_norm,
k_norm=k_norm,
head_dim=head_dim,
cos_sin_cache=rope_cache,
is_neox=True,
positions=rope_positions,
round_norm_before_rope=True,
cache_has_full_width=True,
)
except Exception as exc:
_ERNIE_QKNORM_ROPE.on_exception(exc, logger=logger)
return _ernie_qknorm_rope_reference(
q_input,
k_input,
q_norm,
k_norm,
head_dim,
rope_cos,
rope_sin,
)
else:
if verified:
return out
ref = _ernie_qknorm_rope_reference(
q_input,
k_input,
q_norm,
k_norm,
head_dim,
rope_cos,
rope_sin,
)
return _ERNIE_QKNORM_ROPE.accept_or_fallback(
out,
ref,
equal=tensors_equal,
logger=logger,
mismatch_msg=(
"ERNIE fused QKNorm+RoPE fast path is not bit-exact on "
"this platform; falling back to split kernels"
),
)
return _ernie_qknorm_rope_reference(
q, k, q_norm, k_norm, head_dim, rope_cos, rope_sin
)
def _eager_geglu(gate_up: torch.Tensor) -> torch.Tensor:
gate, up = gate_up.chunk(2, dim=-1)
return up * F.gelu(gate)
@@ -687,6 +785,10 @@ class ErnieImageTransformer2DModel(CachableDiT, LayerwiseOffloadableModuleMixin)
all_ids = torch.cat([image_ids, text_ids], dim=1)
rotary_pos_emb = self.pos_embed(all_ids)
rope_cos, rope_sin = _precompute_rope_cos_sin(rotary_pos_emb, dtype)
rope_cache = torch.cat((rope_cos, rope_sin), dim=-1).contiguous()
rope_positions = torch.arange(
rope_cache.shape[0], device=device, dtype=torch.long
)
attn_mask = attn_mask_meta = None
if encoder_hidden_states_mask is not None:
@@ -715,6 +817,8 @@ class ErnieImageTransformer2DModel(CachableDiT, LayerwiseOffloadableModuleMixin)
x,
rope_cos,
rope_sin,
rope_cache,
rope_positions,
shift_msa,
scale_msa,
gate_msa,