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

[Diffusion] Preserve BF16 rounding in Hopper LTX QKNorm and RoPE fusion (#38533)

Browse Source
This commit is contained in:
Xiaoyu Zhang
2026-09-11 23:02:54 +08:00
committed by GitHub
parent dd67a42634
commit 7f09fbcd25
2 changed files with 35 additions and 14 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/kernels/kda_kernels/csrc/diffusion/ltx2_qknorm_split_rope.cuh
+28 -11
View File
@@ -4,8 +4,8 @@
// Developed with MIT HAN Lab Kernel Design Agents:
// https://github.com/mit-han-lab/kernel-design-agents
//
// This mirrors the LTX2 eager oracle: RMSNorm and split RoPE both run in
// fp32, rounding to bf16 only once at the final attention input.
// The original FP32 path rounds only at the final attention input. The
// Hopper variant preserves the eager BF16 RMSNorm and cosine-product rounding.
#pragma once
@@ -71,17 +71,29 @@ SGL_DEVICE float compute_rstd(
return *s_rstd;
}
template <bool kRoundIntermediates>
SGL_DEVICE float norm_value(float x, float weight, float rstd) {
return weight * (rstd * x);
const float value = weight * (rstd * x);
if constexpr (kRoundIntermediates) {
return __bfloat162float(__float2bfloat16_rn(value));
}
return value;
}
template <bool kRoundIntermediates>
SGL_DEVICE void rope_pair(float x0, float x1, float cos, float sin, float& y0, float& y1) {
const float p0 = x0 * cos;
const float p1 = x1 * cos;
float p0 = x0 * cos;
float p1 = x1 * cos;
if constexpr (kRoundIntermediates) {
// Hopper eager stores the BF16 cosine product before its addcmul update.
p0 = __bfloat162float(__float2bfloat16_rn(p0));
p1 = __bfloat162float(__float2bfloat16_rn(p1));
}
y0 = fmaf(-sin, x1, p0);
y1 = fmaf(sin, x0, p1);
}
template <bool kRoundIntermediates>
__global__ void ltx2_qknorm_split_rope_kernel(
const bf16_t* __restrict__ x,
const bf16_t* __restrict__ cos,
@@ -119,19 +131,23 @@ __global__ void ltx2_qknorm_split_rope_kernel(
const int64_t offset = pair - head * half_dim;
const int64_t idx0 = head * head_dim + offset;
const int64_t idx1 = idx0 + half_dim;
const float n0 = norm_value(__bfloat162float(xrow[idx0]), __bfloat162float(weight[idx0]), rstd);
const float n1 = norm_value(__bfloat162float(xrow[idx1]), __bfloat162float(weight[idx1]), rstd);
const float n0 =
norm_value<kRoundIntermediates>(__bfloat162float(xrow[idx0]), __bfloat162float(weight[idx0]), rstd);
const float n1 =
norm_value<kRoundIntermediates>(__bfloat162float(xrow[idx1]), __bfloat162float(weight[idx1]), rstd);
const int64_t cos_offset = batch * stride_cos_b + head * stride_cos_h + token * stride_cos_t + offset;
const int64_t sin_offset = batch * stride_sin_b + head * stride_sin_h + token * stride_sin_t + offset;
float y0;
float y1;
rope_pair(n0, n1, __bfloat162float(cos[cos_offset]), __bfloat162float(sin[sin_offset]), y0, y1);
rope_pair<kRoundIntermediates>(
n0, n1, __bfloat162float(cos[cos_offset]), __bfloat162float(sin[sin_offset]), y0, y1);
outrow[idx0] = __float2bfloat16_rn(y0);
outrow[idx1] = __float2bfloat16_rn(y1);
}
}
template <bool kRoundIntermediates>
inline void launch_one(
const tvm::ffi::TensorView& x,
const tvm::ffi::TensorView& cos,
@@ -155,7 +171,7 @@ inline void launch_one(
}
host::RuntimeCheck(num_rows <= static_cast<int64_t>(UINT32_MAX), "LTX2 QKNorm split-RoPE grid is too large");
host::LaunchKernel(dim3(static_cast<uint32_t>(num_rows)), dim3(32, 4), device)(
ltx2_qknorm_split_rope_kernel,
ltx2_qknorm_split_rope_kernel<kRoundIntermediates>,
reinterpret_cast<const bf16_t*>(data_ptr(x)),
reinterpret_cast<const bf16_t*>(data_ptr(cos)),
reinterpret_cast<const bf16_t*>(data_ptr(sin)),
@@ -174,6 +190,7 @@ inline void launch_one(
}
struct LTX2QKNormSplitRopeKernel {
template <bool kRoundIntermediates>
static void
run(tvm::ffi::TensorView q_out,
tvm::ffi::TensorView k_out,
@@ -233,7 +250,7 @@ struct LTX2QKNormSplitRopeKernel {
const int64_t batch_size = batch.unwrap();
const DLDevice dl_device = device.unwrap();
launch_one(
launch_one<kRoundIntermediates>(
q,
q_cos,
q_sin,
@@ -251,7 +268,7 @@ struct LTX2QKNormSplitRopeKernel {
q_sin.stride(1),
q_sin.stride(2),
dl_device);
launch_one(
launch_one<kRoundIntermediates>(
k,
k_cos,
k_sin,
python/sglang/kernels/kda_kernels/ltx2_qknorm_split_rope_jit.py
+7 -3
View File
@@ -13,14 +13,15 @@ if TYPE_CHECKING:
@cache_once
def _jit_ltx2_qknorm_split_rope_module() -> Module:
def _jit_ltx2_qknorm_split_rope_module(round_intermediates: bool = False) -> Module:
return load_jit(
"diffusion_ltx2_qknorm_split_rope",
cuda_files=[_cuda_source("diffusion/ltx2_qknorm_split_rope.cuh")],
cuda_wrappers=[
(
"ltx2_qknorm_split_rope_pair",
"ltx2_qknorm_split_rope::LTX2QKNormSplitRopeKernel::run",
"ltx2_qknorm_split_rope::LTX2QKNormSplitRopeKernel::run"
f"<{str(round_intermediates).lower()}>",
)
],
)
@@ -38,6 +39,7 @@ def _fake_impl(
eps: float,
num_heads: int,
head_dim: int,
round_intermediates: bool = False,
) -> tuple[torch.Tensor, torch.Tensor]:
return torch.empty_like(q, dtype=torch.bfloat16), torch.empty_like(
k, dtype=torch.bfloat16
@@ -61,10 +63,11 @@ def _ltx2_qknorm_split_rope_custom_op(
eps: float,
num_heads: int,
head_dim: int,
round_intermediates: bool = False,
) -> tuple[torch.Tensor, torch.Tensor]:
q_out = torch.empty_like(q, dtype=torch.bfloat16)
k_out = torch.empty_like(k, dtype=torch.bfloat16)
module = _jit_ltx2_qknorm_split_rope_module()
module = _jit_ltx2_qknorm_split_rope_module(round_intermediates)
module.ltx2_qknorm_split_rope_pair(
q_out,
k_out,
@@ -215,4 +218,5 @@ def ltx2_qknorm_split_rope_cuda(
float(eps),
int(num_heads),
int(head_dim),
round_intermediates=allow_sm90 and _is_sm90(q),
)