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

[Diffusion] Fuse Qwen-Image FP8 norm and activation quantization (#37156)

Browse Source 
Co-authored-by: Cursor <cursoragent@cursor.com>
This commit is contained in:
Xiaoyu Zhang
2026-08-31 18:19:42 +08:00
committed by GitHub
co-authored by Cursor
parent a874b83c2c
commit 1ed9bfac2c
7 changed files with 761 additions and 17 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/kernels/jit/csrc/diffusion/norm_scale_shift.cuh
+123 -2
View File
@@ -17,6 +17,7 @@
#include <sgl_kernel/tensor.h> // For TensorMatcher, SymbolicSize, SymbolicDevice
#include <sgl_kernel/math.cuh> // For device::math::rsqrt
#include <sgl_kernel/type.cuh> // For DTypeTrait
#include <sgl_kernel/utils.cuh> // For SGL_DEVICE, bf16_t, LaunchKernel
#include <sgl_kernel/vec.cuh> // For AlignedVector
#include <sgl_kernel/warp.cuh> // For warp::reduce_sum
@@ -39,12 +40,14 @@ static_assert(kWarps == 6);
struct NormScaleShiftParams {
void* y;
void* res_out;
void* quantized;
const void* x;
const void* input_bias;
const void* residual;
const void* gate;
const void* scale;
const void* shift;
const void* input_scale;
float eps;
};
@@ -66,7 +69,16 @@ SGL_DEVICE float cta_reduce_sum(float v, int warp, int lane, float* scratch) {
return scratch[kWarps];
}
template <bool kHasResidual, bool kHasInputBias = false>
SGL_DEVICE float triton_scale_reciprocal(float scale) {
float reciprocal;
// Triton's static FP8 quantizer lowers `1.0 / scale` to div.full.f32.
// Match it exactly because a one-ULP difference at an E4M3 midpoint can
// change the quantized byte.
asm("div.full.f32 %0, %1, %2;" : "=f"(reciprocal) : "f"(1.0f), "f"(scale));
return reciprocal;
}
template <bool kHasResidual, bool kHasInputBias = false, bool kQuantizeFp8 = false>
__global__ void norm_scale_shift_kernel(const NormScaleShiftParams __grid_constant__ params) {
using namespace device;
using Vec = AlignedVector<bf16_t, kVecElems>;
@@ -135,15 +147,31 @@ __global__ void norm_scale_shift_kernel(const NormScaleShiftParams __grid_consta
Vec scv;
Vec shv;
Vec yv;
AlignedVector<fp8_e4m3_t, kVecElems> qv;
scv.load(static_cast<const bf16_t*>(params.scale) + elem_offset);
shv.load(static_cast<const bf16_t*>(params.shift) + elem_offset);
float input_scale_inv = 0.0f;
if constexpr (kQuantizeFp8) {
input_scale_inv = triton_scale_reciprocal(*static_cast<const float*>(params.input_scale));
}
#pragma unroll
for (int i = 0; i < kVecElems; ++i) {
const float norm = static_cast<float>(static_cast<bf16_t>((v[i] - mean) * factor));
yv[i] = static_cast<bf16_t>(norm * (1.0f + static_cast<float>(scv[i])) + static_cast<float>(shv[i]));
const bf16_t rounded = static_cast<bf16_t>(norm * (1.0f + static_cast<float>(scv[i])) + static_cast<float>(shv[i]));
yv[i] = rounded;
if constexpr (kQuantizeFp8) {
const float scaled = static_cast<float>(rounded) * input_scale_inv;
const float clamped =
math::min(math::max(scaled, -DTypeTrait<fp8_e4m3_t>::kFloatMax), DTypeTrait<fp8_e4m3_t>::kFloatMax);
qv[i] = static_cast<fp8_e4m3_t>(clamped);
}
}
yv.store(static_cast<bf16_t*>(params.y) + row_offset + elem_offset);
if constexpr (kQuantizeFp8) {
qv.store(static_cast<fp8_e4m3_t*>(params.quantized) + row_offset + elem_offset);
}
}
__global__ void bias_mul_add_kernel(const NormScaleShiftParams __grid_constant__ params) {
@@ -197,12 +225,14 @@ struct NormScaleShiftKernel {
const auto params = NormScaleShiftParams{
.y = y.data_ptr(),
.res_out = nullptr,
.quantized = nullptr,
.x = x.data_ptr(),
.input_bias = nullptr,
.residual = nullptr,
.gate = nullptr,
.scale = scale.data_ptr(),
.shift = shift.data_ptr(),
.input_scale = nullptr,
.eps = static_cast<float>(eps),
};
LaunchKernel(grid, kThreads, device.unwrap())(norm_scale_shift_kernel<false>, params);
@@ -237,18 +267,105 @@ struct ScaleResidualNormScaleShiftKernel {
const auto params = NormScaleShiftParams{
.y = y.data_ptr(),
.res_out = res_out.data_ptr(),
.quantized = nullptr,
.x = x.data_ptr(),
.input_bias = nullptr,
.residual = residual.data_ptr(),
.gate = gate.data_ptr(),
.scale = scale.data_ptr(),
.shift = shift.data_ptr(),
.input_scale = nullptr,
.eps = static_cast<float>(eps),
};
LaunchKernel(grid, kThreads, device.unwrap())(norm_scale_shift_kernel<true>, params);
}
};
/** \brief Fuse Qwen LayerNorm/modulation with static E4M3 activation quantization. */
struct NormScaleShiftFp8Kernel {
static void
run(tvm::ffi::TensorView y,
tvm::ffi::TensorView quantized,
tvm::ffi::TensorView x,
tvm::ffi::TensorView scale,
tvm::ffi::TensorView shift,
tvm::ffi::TensorView input_scale,
double eps) {
using namespace host;
auto N = SymbolicSize{"num_rows"};
auto device = SymbolicDevice{};
device.set_options<kDLCUDA>();
TensorMatcher({N, kHidden}).with_dtype<bf16_t>().with_device(device).verify(x).verify(y);
TensorMatcher({N, kHidden}).with_dtype<fp8_e4m3_t>().with_device(device).verify(quantized);
TensorMatcher({kHidden}).with_dtype<bf16_t>().with_device(device).verify(scale).verify(shift);
TensorMatcher({1}).with_dtype<fp32_t>().with_device(device).verify(input_scale);
const uint32_t grid = verify_nss_geometry(N);
const auto params = NormScaleShiftParams{
.y = y.data_ptr(),
.res_out = nullptr,
.quantized = quantized.data_ptr(),
.x = x.data_ptr(),
.input_bias = nullptr,
.residual = nullptr,
.gate = nullptr,
.scale = scale.data_ptr(),
.shift = shift.data_ptr(),
.input_scale = input_scale.data_ptr(),
.eps = static_cast<float>(eps),
};
LaunchKernel(grid, kThreads, device.unwrap())(norm_scale_shift_kernel<false, false, true>, params);
}
};
/** \brief Fuse Qwen residual LayerNorm/modulation with static E4M3 activation quantization. */
struct ScaleResidualNormScaleShiftFp8Kernel {
static void
run(tvm::ffi::TensorView y,
tvm::ffi::TensorView quantized,
tvm::ffi::TensorView res_out,
tvm::ffi::TensorView residual,
tvm::ffi::TensorView x,
tvm::ffi::TensorView gate,
tvm::ffi::TensorView scale,
tvm::ffi::TensorView shift,
tvm::ffi::TensorView input_scale,
double eps) {
using namespace host;
auto N = SymbolicSize{"num_rows"};
auto device = SymbolicDevice{};
device.set_options<kDLCUDA>();
TensorMatcher({N, kHidden})
.with_dtype<bf16_t>()
.with_device(device)
.verify(x)
.verify(residual)
.verify(y)
.verify(res_out);
TensorMatcher({N, kHidden}).with_dtype<fp8_e4m3_t>().with_device(device).verify(quantized);
TensorMatcher({kHidden}).with_dtype<bf16_t>().with_device(device).verify(gate).verify(scale).verify(shift);
TensorMatcher({1}).with_dtype<fp32_t>().with_device(device).verify(input_scale);
const uint32_t grid = verify_nss_geometry(N);
const auto params = NormScaleShiftParams{
.y = y.data_ptr(),
.res_out = res_out.data_ptr(),
.quantized = quantized.data_ptr(),
.x = x.data_ptr(),
.input_bias = nullptr,
.residual = residual.data_ptr(),
.gate = gate.data_ptr(),
.scale = scale.data_ptr(),
.shift = shift.data_ptr(),
.input_scale = input_scale.data_ptr(),
.eps = static_cast<float>(eps),
};
LaunchKernel(grid, kThreads, device.unwrap())(norm_scale_shift_kernel<true, false, true>, params);
}
};
struct BiasScaleResidualNormScaleShiftKernel {
static void
run(tvm::ffi::TensorView y,
@@ -284,12 +401,14 @@ struct BiasScaleResidualNormScaleShiftKernel {
const auto params = NormScaleShiftParams{
.y = y.data_ptr(),
.res_out = res_out.data_ptr(),
.quantized = nullptr,
.x = x.data_ptr(),
.input_bias = input_bias.data_ptr(),
.residual = residual.data_ptr(),
.gate = gate.data_ptr(),
.scale = scale.data_ptr(),
.shift = shift.data_ptr(),
.input_scale = nullptr,
.eps = static_cast<float>(eps),
};
LaunchKernel(grid, kThreads, device.unwrap())(norm_scale_shift_kernel<true, true>, params);
@@ -315,12 +434,14 @@ struct BiasMulAddKernel {
const auto params = NormScaleShiftParams{
.y = y.data_ptr(),
.res_out = nullptr,
.quantized = nullptr,
.x = x.data_ptr(),
.input_bias = input_bias.data_ptr(),
.residual = residual.data_ptr(),
.gate = gate.data_ptr(),
.scale = nullptr,
.shift = nullptr,
.input_scale = nullptr,
.eps = 0.0f,
};
LaunchKernel(grid, kThreads, device.unwrap())(bias_mul_add_kernel, params);
python/sglang/kernels/ops/diffusion/__init__.py
+4
View File
@@ -382,6 +382,10 @@ _EXPORTS: dict[str, str] = {
"fused_scale_residual_rmsnorm_scale_shift_bitexact": "norm.rmsnorm_scale_shift_bitexact",
"fused_norm_scale_shift": "norm.scale_residual_norm_cutedsl",
"fused_scale_residual_norm_scale_shift": "norm.scale_residual_norm_cutedsl",
"fused_norm_scale_shift_fp8": "norm.norm_scale_shift_jit",
"fused_scale_residual_norm_scale_shift_fp8": "norm.norm_scale_shift_jit",
"try_fused_norm_scale_shift_fp8": "norm.norm_scale_shift_jit",
"try_fused_scale_residual_norm_scale_shift_fp8": "norm.norm_scale_shift_jit",
"validate_scale_shift": "norm.scale_residual_norm_cutedsl",
"can_use_wan_rmsnorm_silu": "norm.wan_rmsnorm_silu_triton",
"wan_rmsnorm_silu": "norm.wan_rmsnorm_silu_triton",
python/sglang/kernels/ops/diffusion/norm/norm_scale_shift_jit.py
+112
View File
@@ -67,6 +67,11 @@ def _row_bf16(t, device: torch.device):
@cache_once
def norm_scale_shift_module() -> Module:
device = torch.device("cuda", torch.cuda.current_device())
if not _blackwell_or_newer(device):
raise RuntimeError(
"Qwen-Image norm-scale-shift JIT kernels require NVIDIA Blackwell or newer"
)
return load_jit(
"norm_scale_shift_native",
cuda_files=["diffusion/norm_scale_shift.cuh"],
@@ -79,6 +84,14 @@ def norm_scale_shift_module() -> Module:
"srnss_bf16_row",
"norm_scale_shift::ScaleResidualNormScaleShiftKernel::run",
),
(
"nss_fp8_row",
"norm_scale_shift::NormScaleShiftFp8Kernel::run",
),
(
"srnss_fp8_row",
"norm_scale_shift::ScaleResidualNormScaleShiftFp8Kernel::run",
),
(
"bias_srnss_bf16_row",
"norm_scale_shift::BiasScaleResidualNormScaleShiftKernel::run",
@@ -94,6 +107,44 @@ def norm_scale_shift_module() -> Module:
_module = norm_scale_shift_module
def fused_norm_scale_shift_fp8(x, scale, shift, input_scale, eps):
"""Return exact BF16 modulation output and its static E4M3 quantization."""
normalized = torch.empty_like(x)
quantized = torch.empty_like(x, dtype=torch.float8_e4m3fn)
_module().nss_fp8_row(
normalized.view(-1, _HIDDEN),
quantized.view(-1, _HIDDEN),
x.view(-1, _HIDDEN),
scale,
shift,
input_scale.reshape(1),
float(eps),
)
return normalized, quantized
def fused_scale_residual_norm_scale_shift_fp8(
residual, x, gate, scale, shift, input_scale, eps
):
"""Return exact BF16 residual/modulation outputs and E4M3 quantization."""
normalized = torch.empty_like(x)
quantized = torch.empty_like(x, dtype=torch.float8_e4m3fn)
residual_out = torch.empty_like(x)
_module().srnss_fp8_row(
normalized.view(-1, _HIDDEN),
quantized.view(-1, _HIDDEN),
residual_out.view(-1, _HIDDEN),
residual.view(-1, _HIDDEN),
x.view(-1, _HIDDEN),
gate,
scale,
shift,
input_scale.reshape(1),
float(eps),
)
return normalized, quantized, residual_out
def try_fused_norm_scale_shift(x, weight, bias, scale, shift, norm_type, eps):
if norm_type != "layer" or weight is not None or bias is not None:
return None
@@ -145,6 +196,67 @@ def try_fused_scale_residual_norm_scale_shift(
return y, residual_out
def try_fused_norm_scale_shift_fp8(
x, weight, bias, scale, shift, input_scale, norm_type, eps
):
if norm_type != "layer" or weight is not None or bias is not None:
return None
if not _nss_activation(x) or not _blackwell_or_newer(x.device):
return None
scale = _row_bf16(scale, x.device)
shift = _row_bf16(shift, x.device)
if scale is None or shift is None:
return None
if not _fp8_input_scale(input_scale, x.device):
return None
return fused_norm_scale_shift_fp8(x, scale, shift, input_scale, eps)
def try_fused_scale_residual_norm_scale_shift_fp8(
residual,
x,
gate,
weight,
bias,
scale,
shift,
input_scale,
norm_type,
eps,
):
if norm_type != "layer" or weight is not None or bias is not None:
return None
if not (
_nss_activation(x)
and _nss_activation(residual, x)
and _blackwell_or_newer(x.device)
):
return None
gate = _row_bf16(gate, x.device)
scale = _row_bf16(scale, x.device)
shift = _row_bf16(shift, x.device)
if gate is None or scale is None or shift is None:
return None
if not _fp8_input_scale(input_scale, x.device):
return None
return fused_scale_residual_norm_scale_shift_fp8(
residual, x, gate, scale, shift, input_scale, eps
)
def _fp8_input_scale(t, device: torch.device) -> bool:
return (
isinstance(t, torch.Tensor)
and t.is_cuda
and t.device == device
and t.dtype == torch.float32
and t.numel() == 1
and t.is_contiguous()
)
def try_fused_bias_scale_residual_norm_scale_shift(
residual, x, input_bias, gate, weight, bias, scale, shift, norm_type, eps
):
python/sglang/multimodal_gen/runtime/models/dits/qwen_image.py
+255 -15
View File
@@ -22,6 +22,8 @@ from sglang.kernels.ops.diffusion import (
mark_fused_gelu_site,
try_fused_bias_mul_add,
try_fused_bias_scale_residual_norm_scale_shift,
try_fused_norm_scale_shift_fp8,
try_fused_scale_residual_norm_scale_shift_fp8,
)
from sglang.multimodal_gen.configs.models.dits.qwenimage import QwenImageDitConfig
from sglang.multimodal_gen.configs.models.fsdp import is_transformer_block
@@ -70,6 +72,9 @@ from sglang.multimodal_gen.runtime.layers.quantization.configs.nunchaku_config i
NunchakuConfig,
is_nunchaku_available,
)
from sglang.multimodal_gen.runtime.layers.quantization.modelopt_quant import (
ModelOptFp8LinearMethod,
)
from sglang.multimodal_gen.runtime.layers.rotary_embedding import (
apply_flashinfer_rope_qk_inplace,
)
@@ -1125,6 +1130,65 @@ class QwenImageTransformerBlock(nn.Module):
self.img_mlp = NunchakuFeedForward(self.img_mlp, **nunchaku_kwargs)
self.txt_mlp = NunchakuFeedForward(self.txt_mlp, **nunchaku_kwargs)
self._fp8_img_attn_norm_quant = False
self._fp8_txt_attn_norm_quant = False
self._fp8_img_mlp_norm_quant = False
self._fp8_txt_mlp_norm_quant = False
@staticmethod
def _valid_modelopt_fp8_linear(linear: nn.Module) -> bool:
input_scale = getattr(linear, "input_scale", None)
return (
isinstance(getattr(linear, "quant_method", None), ModelOptFp8LinearMethod)
and isinstance(input_scale, torch.Tensor)
and input_scale.is_cuda
and input_scale.dtype == torch.float32
and input_scale.numel() == 1
and input_scale.is_contiguous()
and bool(torch.isfinite(input_scale).all().item())
and bool((input_scale > 0).all().item())
)
@classmethod
def _shared_modelopt_fp8_scale(cls, linears: list[nn.Module]) -> bool:
if not all(cls._valid_modelopt_fp8_linear(linear) for linear in linears):
return False
reference = linears[0].input_scale
return all(torch.equal(reference, linear.input_scale) for linear in linears[1:])
def configure_fp8_norm_quant(self) -> None:
"""Enable exact norm+quant paths after checkpoint scales are materialized."""
if not torch.cuda.is_available():
return
capability = torch.cuda.get_device_capability()
if self.dim != 3072 or capability[0] < 10 or self.zero_cond_t:
return
if self.attn.use_fused_qkv:
self._fp8_img_attn_norm_quant = self._valid_modelopt_fp8_linear(
self.attn.to_qkv
)
else:
self._fp8_img_attn_norm_quant = self._shared_modelopt_fp8_scale(
[self.attn.to_q, self.attn.to_k, self.attn.to_v]
)
if self.attn.added_kv_proj_dim is not None:
if self.attn.use_fused_added_qkv:
self._fp8_txt_attn_norm_quant = self._valid_modelopt_fp8_linear(
self.attn.to_added_qkv
)
else:
self._fp8_txt_attn_norm_quant = self._shared_modelopt_fp8_scale(
[self.attn.add_q_proj, self.attn.add_k_proj, self.attn.add_v_proj]
)
if isinstance(self.img_mlp, QwenImageFeedForward):
self._fp8_img_mlp_norm_quant = self._valid_modelopt_fp8_linear(
self.img_mlp.net[0].proj
)
if isinstance(self.txt_mlp, QwenImageFeedForward):
self._fp8_txt_mlp_norm_quant = self._valid_modelopt_fp8_linear(
self.txt_mlp.net[0].proj
)
def _norm_scale_shift(
self,
norm_module: LayerNormScaleShift,
@@ -1200,6 +1264,68 @@ class QwenImageTransformerBlock(nn.Module):
)
return img_mod_params, txt_mod_params
def _try_fp8_norm_quant(
self,
norm_module: LayerNormScaleShift,
*,
x: torch.Tensor,
mod_params: torch.Tensor,
input_scale: Optional[torch.Tensor],
enabled: bool,
modulate_index: Optional[torch.Tensor],
use_bcg_helpers: bool,
) -> Optional[tuple[torch.Tensor, torch.Tensor, torch.Tensor]]:
if not enabled or modulate_index is not None or use_bcg_helpers:
return None
shift, scale, gate = mod_params.chunk(3, dim=-1)
result = try_fused_norm_scale_shift_fp8(
x,
getattr(norm_module.norm, "weight", None),
getattr(norm_module.norm, "bias", None),
scale,
shift,
input_scale,
norm_module.norm_type,
norm_module.eps,
)
if result is None:
return None
normalized, quantized = result
return quantized, gate.unsqueeze(1), normalized
def _try_fp8_residual_norm_quant(
self,
norm_module: ScaleResidualLayerNormScaleShift,
*,
residual: torch.Tensor,
x: torch.Tensor,
residual_gate: torch.Tensor,
mod_params: torch.Tensor,
input_scale: Optional[torch.Tensor],
enabled: bool,
modulate_index: Optional[torch.Tensor],
use_bcg_helpers: bool,
) -> Optional[tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]]:
if not enabled or modulate_index is not None or use_bcg_helpers:
return None
shift, scale, gate = mod_params.chunk(3, dim=-1)
result = try_fused_scale_residual_norm_scale_shift_fp8(
residual,
x,
residual_gate,
getattr(norm_module.norm, "weight", None),
getattr(norm_module.norm, "bias", None),
scale,
shift,
input_scale,
norm_module.norm_type,
norm_module.eps,
)
if result is None:
return None
normalized, quantized, residual_out = result
return quantized, residual_out, gate.unsqueeze(1), normalized
def _modulate(
self,
x: torch.Tensor,
@@ -1361,16 +1487,57 @@ class QwenImageTransformerBlock(nn.Module):
use_bcg_helpers = is_in_breakable_cuda_graph()
# Process image stream - norm1 + modulation
img_modulated, img_gate1 = self._modulate(
hidden_states,
img_mod1,
img_fp8 = self._try_fp8_norm_quant(
self.img_norm1,
modulate_index,
x=hidden_states,
mod_params=img_mod1,
input_scale=(
(
self.attn.to_qkv.input_scale
if self.attn.use_fused_qkv
else self.attn.to_q.input_scale
)
if self._fp8_img_attn_norm_quant
else None
),
enabled=self._fp8_img_attn_norm_quant,
modulate_index=modulate_index,
use_bcg_helpers=use_bcg_helpers,
)
if img_fp8 is None:
img_modulated, img_gate1 = self._modulate(
hidden_states,
img_mod1,
self.img_norm1,
modulate_index,
use_bcg_helpers=use_bcg_helpers,
)
img_modulated_bf16 = None
else:
img_modulated, img_gate1, img_modulated_bf16 = img_fp8
# Process text stream - norm1 + modulation
txt_fp8 = self._try_fp8_norm_quant(
self.txt_norm1,
x=encoder_hidden_states,
mod_params=txt_mod1,
input_scale=(
(
self.attn.to_added_qkv.input_scale
if self.attn.use_fused_added_qkv
else self.attn.add_q_proj.input_scale
)
if self._fp8_txt_attn_norm_quant
else None
),
enabled=self._fp8_txt_attn_norm_quant,
modulate_index=modulate_index,
use_bcg_helpers=use_bcg_helpers,
)
txt_shift1, txt_scale1, txt_gate1_raw = txt_mod1.chunk(3, dim=-1)
if use_bcg_helpers:
if txt_fp8 is not None:
txt_modulated, txt_gate1, txt_modulated_bf16 = txt_fp8
elif use_bcg_helpers:
txt_modulated_bf16 = None
txt_modulated = self._norm_scale_shift(
self.txt_norm1,
encoder_hidden_states,
@@ -1378,10 +1545,12 @@ class QwenImageTransformerBlock(nn.Module):
scale=txt_scale1,
)
else:
txt_modulated_bf16 = None
txt_modulated = self.txt_norm1(
encoder_hidden_states, shift=txt_shift1, scale=txt_scale1
)
txt_gate1 = txt_gate1_raw.unsqueeze(1)
if txt_fp8 is None:
txt_gate1 = txt_gate1_raw.unsqueeze(1)
# Use QwenAttnProcessor2_0 for joint attention computation
# This directly implements the DoubleStreamLayerMegatron logic:
@@ -1399,6 +1568,7 @@ class QwenImageTransformerBlock(nn.Module):
image_rotary_emb=image_rotary_emb,
**joint_attention_kwargs,
)
del img_modulated_bf16, txt_modulated_bf16
# QwenAttnProcessor2_0 returns (img_output, txt_output) when encoder_hidden_states is provided
(
@@ -1408,16 +1578,40 @@ class QwenImageTransformerBlock(nn.Module):
txt_attn_bias,
) = attn_output
# Process image stream - norm2 + MLP
img_modulated2, hidden_states, img_gate2 = self._modulate(
img_attn_output,
img_mod2,
img_fp8_mlp = self._try_fp8_residual_norm_quant(
self.img_norm2,
modulate_index,
gate_x=img_gate1,
residual_x=hidden_states,
x_bias=img_attn_bias,
residual=hidden_states,
x=img_attn_output,
residual_gate=img_gate1,
mod_params=img_mod2,
input_scale=(
self.img_mlp.net[0].proj.input_scale
if self._fp8_img_mlp_norm_quant
else None
),
enabled=self._fp8_img_mlp_norm_quant,
modulate_index=modulate_index,
use_bcg_helpers=use_bcg_helpers,
)
if img_fp8_mlp is None:
img_modulated2, hidden_states, img_gate2 = self._modulate(
img_attn_output,
img_mod2,
self.img_norm2,
modulate_index,
gate_x=img_gate1,
residual_x=hidden_states,
x_bias=img_attn_bias,
use_bcg_helpers=use_bcg_helpers,
)
img_modulated2_bf16 = None
else:
(
img_modulated2,
hidden_states,
img_gate2,
img_modulated2_bf16,
) = img_fp8_mlp
if isinstance(self.img_mlp, QwenImageFeedForward):
img_mlp_output, img_mlp_bias = self.img_mlp.forward_with_bias(
img_modulated2
@@ -1425,6 +1619,7 @@ class QwenImageTransformerBlock(nn.Module):
else:
img_mlp_output = self.img_mlp(img_modulated2)
img_mlp_bias = None
del img_modulated2_bf16
if img_mlp_output.dim() == 2:
img_mlp_output = img_mlp_output.unsqueeze(0)
@@ -1438,7 +1633,30 @@ class QwenImageTransformerBlock(nn.Module):
# Process text stream - norm2 + MLP
txt_shift2, txt_scale2, txt_gate2_raw = txt_mod2.chunk(3, dim=-1)
if use_bcg_helpers:
txt_fp8_mlp = self._try_fp8_residual_norm_quant(
self.txt_norm2,
residual=encoder_hidden_states,
x=txt_attn_output,
residual_gate=txt_gate1,
mod_params=txt_mod2,
input_scale=(
self.txt_mlp.net[0].proj.input_scale
if self._fp8_txt_mlp_norm_quant
else None
),
enabled=self._fp8_txt_mlp_norm_quant,
modulate_index=modulate_index,
use_bcg_helpers=use_bcg_helpers,
)
if txt_fp8_mlp is not None:
(
txt_modulated2,
encoder_hidden_states,
txt_gate2,
txt_modulated2_bf16,
) = txt_fp8_mlp
elif use_bcg_helpers:
txt_modulated2_bf16 = None
if txt_attn_bias is not None:
txt_attn_output = txt_attn_output + txt_attn_bias
(
@@ -1453,6 +1671,7 @@ class QwenImageTransformerBlock(nn.Module):
scale=txt_scale2,
)
elif txt_attn_bias is not None:
txt_modulated2_bf16 = None
txt_modulated2, encoder_hidden_states, _ = self._modulate(
txt_attn_output,
txt_mod2,
@@ -1462,6 +1681,7 @@ class QwenImageTransformerBlock(nn.Module):
x_bias=txt_attn_bias,
)
else:
txt_modulated2_bf16 = None
txt_modulated2, encoder_hidden_states = self.txt_norm2(
residual=encoder_hidden_states,
x=txt_attn_output,
@@ -1469,7 +1689,8 @@ class QwenImageTransformerBlock(nn.Module):
shift=txt_shift2,
scale=txt_scale2,
)
txt_gate2 = txt_gate2_raw.unsqueeze(1)
if txt_fp8_mlp is None:
txt_gate2 = txt_gate2_raw.unsqueeze(1)
if isinstance(self.txt_mlp, QwenImageFeedForward):
txt_mlp_output, txt_mlp_bias = self.txt_mlp.forward_with_bias(
txt_modulated2
@@ -1477,6 +1698,7 @@ class QwenImageTransformerBlock(nn.Module):
else:
txt_mlp_output = self.txt_mlp(txt_modulated2)
txt_mlp_bias = None
del txt_modulated2_bf16
if txt_mlp_output.dim() == 2:
txt_mlp_output = txt_mlp_output.unsqueeze(0)
@@ -1632,6 +1854,24 @@ class QwenImageTransformer2DModel(CachableDiT, LayerwiseOffloadableModuleMixin):
self.layer_names = ["transformer_blocks"]
def post_load_weights(self) -> None:
super().post_load_weights()
for block in self.transformer_blocks:
block.configure_fp8_norm_quant()
enabled = sum(
block._fp8_img_attn_norm_quant
+ block._fp8_txt_attn_norm_quant
+ block._fp8_img_mlp_norm_quant
+ block._fp8_txt_mlp_norm_quant
for block in self.transformer_blocks
)
if enabled:
logger.info(
"Enabled Qwen FP8 norm+quant fusion for %d/%d block paths",
enabled,
4 * len(self.transformer_blocks),
)
@functools.lru_cache(maxsize=50)
def build_modulate_index(self, img_shapes: tuple[int, int, int], device):
sp_world_size = get_sp_world_size()