From bb5e619860f3f58752b82c9e296dcea9b53b3210 Mon Sep 17 00:00:00 2001 From: Xiaoyu Zhang <1182563586@qq.com> Date: Mon, 31 Aug 2026 08:25:37 +0800 Subject: [PATCH] [diffusion] perf: absorb Qwen-Image output projection biases (#37116) --- .../jit/csrc/diffusion/norm_scale_shift.cuh | 118 +++++++++++++- .../sglang/kernels/ops/diffusion/__init__.py | 2 + .../diffusion/norm/norm_scale_shift_jit.py | 70 ++++++++ .../runtime/models/dits/qwen_image.py | 152 +++++++++++++++--- .../test_qwen_output_bias_absorption.py | 123 ++++++++++++++ 5 files changed, 442 insertions(+), 23 deletions(-) create mode 100644 test/registered/kernels/ops/diffusion/test_qwen_output_bias_absorption.py diff --git a/python/sglang/kernels/jit/csrc/diffusion/norm_scale_shift.cuh b/python/sglang/kernels/jit/csrc/diffusion/norm_scale_shift.cuh index ec3673203..3f19bddf8 100644 --- a/python/sglang/kernels/jit/csrc/diffusion/norm_scale_shift.cuh +++ b/python/sglang/kernels/jit/csrc/diffusion/norm_scale_shift.cuh @@ -40,6 +40,7 @@ struct NormScaleShiftParams { void* y; void* res_out; const void* x; + const void* input_bias; const void* residual; const void* gate; const void* scale; @@ -65,7 +66,7 @@ SGL_DEVICE float cta_reduce_sum(float v, int warp, int lane, float* scratch) { return scratch[kWarps]; } -template +template __global__ void norm_scale_shift_kernel(const NormScaleShiftParams __grid_constant__ params) { using namespace device; using Vec = AlignedVector; @@ -89,6 +90,16 @@ __global__ void norm_scale_shift_kernel(const NormScaleShiftParams __grid_consta v[i] = static_cast(xv[i]); } + if constexpr (kHasInputBias) { + Vec bv; + bv.load(static_cast(params.input_bias) + elem_offset); +#pragma unroll + for (int i = 0; i < kVecElems; ++i) { + // Match the standalone BF16 output-projection bias addition. + v[i] = static_cast(static_cast(v[i] + static_cast(bv[i]))); + } + } + if constexpr (kHasResidual) { Vec gv; Vec rv; @@ -135,6 +146,31 @@ __global__ void norm_scale_shift_kernel(const NormScaleShiftParams __grid_consta yv.store(static_cast(params.y) + row_offset + elem_offset); } +__global__ void bias_mul_add_kernel(const NormScaleShiftParams __grid_constant__ params) { + using namespace device; + using Vec = AlignedVector; + + const int row_offset = blockIdx.x * kHidden; + const int elem_offset = threadIdx.x * kVecElems; + + Vec xv; + Vec bv; + Vec gv; + Vec rv; + Vec yv; + xv.load(static_cast(params.x) + row_offset + elem_offset); + bv.load(static_cast(params.input_bias) + elem_offset); + gv.load(static_cast(params.gate) + elem_offset); + rv.load(static_cast(params.residual) + row_offset + elem_offset); + +#pragma unroll + for (int i = 0; i < kVecElems; ++i) { + const bf16_t biased = static_cast(static_cast(xv[i]) + static_cast(bv[i])); + yv[i] = __hfma(biased, gv[i], rv[i]); + } + yv.store(static_cast(params.y) + row_offset + elem_offset); +} + inline uint32_t verify_nss_geometry(host::SymbolicSize& num_rows) { using namespace host; RuntimeCheck(num_rows.unwrap() > 0, "num_rows must be positive"); @@ -162,6 +198,7 @@ struct NormScaleShiftKernel { .y = y.data_ptr(), .res_out = nullptr, .x = x.data_ptr(), + .input_bias = nullptr, .residual = nullptr, .gate = nullptr, .scale = scale.data_ptr(), @@ -201,6 +238,7 @@ struct ScaleResidualNormScaleShiftKernel { .y = y.data_ptr(), .res_out = res_out.data_ptr(), .x = x.data_ptr(), + .input_bias = nullptr, .residual = residual.data_ptr(), .gate = gate.data_ptr(), .scale = scale.data_ptr(), @@ -211,6 +249,84 @@ struct ScaleResidualNormScaleShiftKernel { } }; +struct BiasScaleResidualNormScaleShiftKernel { + static void + run(tvm::ffi::TensorView y, + tvm::ffi::TensorView res_out, + tvm::ffi::TensorView residual, + tvm::ffi::TensorView x, + tvm::ffi::TensorView input_bias, + tvm::ffi::TensorView gate, + tvm::ffi::TensorView scale, + tvm::ffi::TensorView shift, + double eps) { + using namespace host; + auto N = SymbolicSize{"num_rows"}; + auto device = SymbolicDevice{}; + device.set_options(); + + TensorMatcher({N, kHidden}) + .with_dtype() + .with_device(device) + .verify(x) + .verify(residual) + .verify(y) + .verify(res_out); + TensorMatcher({kHidden}) + .with_dtype() + .with_device(device) + .verify(input_bias) + .verify(gate) + .verify(scale) + .verify(shift); + + const uint32_t grid = verify_nss_geometry(N); + const auto params = NormScaleShiftParams{ + .y = y.data_ptr(), + .res_out = res_out.data_ptr(), + .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(), + .eps = static_cast(eps), + }; + LaunchKernel(grid, kThreads, device.unwrap())(norm_scale_shift_kernel, params); + } +}; + +struct BiasMulAddKernel { + static void + run(tvm::ffi::TensorView y, + tvm::ffi::TensorView x, + tvm::ffi::TensorView input_bias, + tvm::ffi::TensorView gate, + tvm::ffi::TensorView residual) { + using namespace host; + auto N = SymbolicSize{"num_rows"}; + auto device = SymbolicDevice{}; + device.set_options(); + + TensorMatcher({N, kHidden}).with_dtype().with_device(device).verify(x).verify(residual).verify(y); + TensorMatcher({kHidden}).with_dtype().with_device(device).verify(input_bias).verify(gate); + + const uint32_t grid = verify_nss_geometry(N); + const auto params = NormScaleShiftParams{ + .y = y.data_ptr(), + .res_out = nullptr, + .x = x.data_ptr(), + .input_bias = input_bias.data_ptr(), + .residual = residual.data_ptr(), + .gate = gate.data_ptr(), + .scale = nullptr, + .shift = nullptr, + .eps = 0.0f, + }; + LaunchKernel(grid, kThreads, device.unwrap())(bias_mul_add_kernel, params); + } +}; + } // namespace norm_scale_shift } // namespace sglang diff --git a/python/sglang/kernels/ops/diffusion/__init__.py b/python/sglang/kernels/ops/diffusion/__init__.py index ea8486b5d..c13824a3e 100644 --- a/python/sglang/kernels/ops/diffusion/__init__.py +++ b/python/sglang/kernels/ops/diffusion/__init__.py @@ -373,6 +373,8 @@ _EXPORTS: dict[str, str] = { "rmsnorm_tanh_residual": "norm.native_bf16_rmsnorm_triton", "norm_infer": "norm.norm_triton", "rms_norm_fn": "norm.norm_triton", + "try_fused_bias_mul_add": "norm.norm_scale_shift_jit", + "try_fused_bias_scale_residual_norm_scale_shift": "norm.norm_scale_shift_jit", "triton_one_pass_rms_norm": "norm.rmsnorm_onepass_triton", "can_use_fused_rmsnorm_scale_shift": "norm.rmsnorm_scale_shift_bitexact", "can_use_fused_scale_residual_rmsnorm_scale_shift": "norm.rmsnorm_scale_shift_bitexact", diff --git a/python/sglang/kernels/ops/diffusion/norm/norm_scale_shift_jit.py b/python/sglang/kernels/ops/diffusion/norm/norm_scale_shift_jit.py index 5b1712112..5d0881a6e 100644 --- a/python/sglang/kernels/ops/diffusion/norm/norm_scale_shift_jit.py +++ b/python/sglang/kernels/ops/diffusion/norm/norm_scale_shift_jit.py @@ -24,6 +24,13 @@ def _blackwell_or_newer(device: torch.device) -> bool: ) +def _sm103(device: torch.device) -> bool: + return torch.cuda.is_available() and torch.cuda.get_device_capability(device) == ( + 10, + 3, + ) + + def _nss_activation(t, like=None) -> bool: return ( isinstance(t, torch.Tensor) @@ -72,6 +79,14 @@ def norm_scale_shift_module() -> Module: "srnss_bf16_row", "norm_scale_shift::ScaleResidualNormScaleShiftKernel::run", ), + ( + "bias_srnss_bf16_row", + "norm_scale_shift::BiasScaleResidualNormScaleShiftKernel::run", + ), + ( + "bias_mul_add_bf16_row", + "norm_scale_shift::BiasMulAddKernel::run", + ), ], ) @@ -128,3 +143,58 @@ def try_fused_scale_residual_norm_scale_shift( float(eps), ) return y, residual_out + + +def try_fused_bias_scale_residual_norm_scale_shift( + residual, x, input_bias, gate, weight, bias, scale, shift, norm_type, eps +): + if torch.compiler.is_compiling(): + return None + 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 _sm103(x.device)): + return None + + input_bias = _row_bf16(input_bias, x.device) + gate = _row_bf16(gate, x.device) + scale = _row_bf16(scale, x.device) + shift = _row_bf16(shift, x.device) + if any(tensor is None for tensor in (input_bias, gate, scale, shift)): + return None + + y = torch.empty_like(x) + residual_out = torch.empty_like(x) + _module().bias_srnss_bf16_row( + y.view(-1, _HIDDEN), + residual_out.view(-1, _HIDDEN), + residual.view(-1, _HIDDEN), + x.view(-1, _HIDDEN), + input_bias, + gate, + scale, + shift, + float(eps), + ) + return y, residual_out + + +def try_fused_bias_mul_add(x, input_bias, gate, residual): + if torch.compiler.is_compiling(): + return None + if not (_nss_activation(x) and _nss_activation(residual, x) and _sm103(x.device)): + return None + + input_bias = _row_bf16(input_bias, x.device) + gate = _row_bf16(gate, x.device) + if input_bias is None or gate is None: + return None + + y = torch.empty_like(x) + _module().bias_mul_add_bf16_row( + y.view(-1, _HIDDEN), + x.view(-1, _HIDDEN), + input_bias, + gate, + residual.view(-1, _HIDDEN), + ) + return y diff --git a/python/sglang/multimodal_gen/runtime/models/dits/qwen_image.py b/python/sglang/multimodal_gen/runtime/models/dits/qwen_image.py index 3414f00a2..f2bf082d3 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/qwen_image.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/qwen_image.py @@ -20,6 +20,8 @@ from sglang.kernels.ops.diffusion import ( fused_gelu_active, fused_linear_gelu_tanh, mark_fused_gelu_site, + try_fused_bias_mul_add, + try_fused_bias_scale_residual_norm_scale_shift, ) from sglang.multimodal_gen.configs.models.dits.qwenimage import QwenImageDitConfig from sglang.multimodal_gen.configs.models.fsdp import is_transformer_block @@ -76,7 +78,10 @@ from sglang.multimodal_gen.runtime.managers.memory_managers.layerwise_offload im ) from sglang.multimodal_gen.runtime.models.dits.base import CachableDiT from sglang.multimodal_gen.runtime.models.dits.common import get_qkv_projections -from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum +from sglang.multimodal_gen.runtime.platforms import ( + AttentionBackendEnum, + current_platform, +) from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger from sglang.srt.model_executor.runner_backend_utils.breakable_cuda_graph import ( is_in_breakable_cuda_graph, @@ -110,6 +115,27 @@ def _local_seq_len(seq_len: int, sp_world_size: int) -> int: _get_qkv_projections = get_qkv_projections +def _can_defer_modelopt_output_bias( + quant_config: Optional[QuantizationConfig], capability: Any +) -> bool: + # Absorbing the bias moves a BF16 rounding point. The resulting image + # quality has only been validated on SM103, so other GPUs keep the GEMM + # bias epilogue used before this optimization. + return ( + quant_config is not None + and hasattr(quant_config, "get_name") + and quant_config.get_name() in {"modelopt_fp8", "modelopt_fp4"} + and capability is not None + and (capability.major, capability.minor) == (10, 3) + ) + + +def _defer_modelopt_output_bias(quant_config: Optional[QuantizationConfig]) -> bool: + return _can_defer_modelopt_output_bias( + quant_config, current_platform.get_device_capability() + ) + + def _safe_tensor_version(tensor: torch.Tensor) -> Optional[int]: """Read a tensor version counter without rejecting inference tensors.""" return None if tensor.is_inference() else tensor._version @@ -579,6 +605,7 @@ class QwenImageCrossAttention(nn.Module): self.parallel_attention = parallel_attention self.added_kv_proj_dim = added_kv_proj_dim self.prefix = prefix + self.defer_output_bias = _defer_modelopt_output_bias(quant_config) self.use_fused_qkv = isinstance(quant_config, NunchakuConfig) @@ -672,6 +699,7 @@ class QwenImageCrossAttention(nn.Module): self.dim, bias=out_bias, input_is_parallel=True, + skip_bias_add=self.defer_output_bias, quant_config=quant_config, prefix=f"{prefix}.to_add_out", ) @@ -686,6 +714,7 @@ class QwenImageCrossAttention(nn.Module): self.dim, bias=out_bias, input_is_parallel=True, + skip_bias_add=self.defer_output_bias, quant_config=quant_config, prefix=f"{prefix}.to_out.0", ) @@ -857,13 +886,13 @@ class QwenImageCrossAttention(nn.Module): ) # Apply output projections - img_attn_output, _ = self.to_out[0](img_attn_output) + img_attn_output, img_attn_bias = self.to_out[0](img_attn_output) if len(self.to_out) > 1: (img_attn_output,) = self.to_out[1](img_attn_output) # dropout - txt_attn_output, _ = self.to_add_out(txt_attn_output) + txt_attn_output, txt_attn_bias = self.to_add_out(txt_attn_output) - return img_attn_output, txt_attn_output + return img_attn_output, txt_attn_output, img_attn_bias, txt_attn_bias class QwenImageGELU(nn.Module): @@ -918,6 +947,7 @@ class QwenImageFeedForward(nn.Module): ) -> None: super().__init__() inner_dim = dim * mult + self.defer_output_bias = _defer_modelopt_output_bias(quant_config) if replicated: # Keep the whole FFN resident on every rank: no per-block # all-reduce. Only worth it when the branch's token count is small @@ -926,6 +956,7 @@ class QwenImageFeedForward(nn.Module): inner_dim, dim_out, bias=True, + skip_bias_add=self.defer_output_bias, quant_config=quant_config, prefix=f"{prefix}.net.2", ) @@ -935,6 +966,7 @@ class QwenImageFeedForward(nn.Module): dim_out, bias=True, input_is_parallel=True, + skip_bias_add=self.defer_output_bias, quant_config=quant_config, prefix=f"{prefix}.net.2", ) @@ -952,11 +984,16 @@ class QwenImageFeedForward(nn.Module): ] ) - def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + def forward_with_bias( + self, hidden_states: torch.Tensor + ) -> tuple[torch.Tensor, Optional[torch.Tensor]]: hidden_states = self.net[0](hidden_states) hidden_states = self.net[1](hidden_states) - hidden_states, _ = self.net[2](hidden_states) - return hidden_states + return self.net[2](hidden_states) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + hidden_states, bias = self.forward_with_bias(hidden_states) + return hidden_states if bias is None else hidden_states + bias class QwenImageTransformerBlock(nn.Module): @@ -1120,6 +1157,25 @@ class QwenImageTransformerBlock(nn.Module): ) -> torch.Tensor: return self.fuse_mul_add(a, b, c, k) + def _bias_mul_add( + self, + a: torch.Tensor, + bias: Optional[torch.Tensor], + b: torch.Tensor, + c: torch.Tensor, + *, + use_bcg_helpers: bool, + ) -> torch.Tensor: + if bias is not None and not use_bcg_helpers: + fused = try_fused_bias_mul_add(a, bias, b, c) + if fused is not None: + return fused + if bias is not None: + a = a + bias + if use_bcg_helpers: + return self._mul_add(a, b, c) + return self.fuse_mul_add(a, b, c) + def _get_modulation_params( self, temb_img_silu: torch.Tensor, @@ -1152,6 +1208,7 @@ class QwenImageTransformerBlock(nn.Module): index: Optional[torch.Tensor] = None, gate_x: Optional[torch.Tensor] = None, residual_x: Optional[torch.Tensor] = None, + x_bias: Optional[torch.Tensor] = None, use_bcg_helpers: bool = False, ) -> Union[ Tuple[torch.Tensor, torch.Tensor], @@ -1164,6 +1221,8 @@ class QwenImageTransformerBlock(nn.Module): shift, scale, gate = mod_params.chunk(3, dim=-1) if index is not None: + if x_bias is not None: + x = x + x_bias actual_batch = x.shape[0] shift0, shift1 = ( shift[:actual_batch], @@ -1214,6 +1273,24 @@ class QwenImageTransformerBlock(nn.Module): scale_result = scale.unsqueeze(1) gate_result = gate.unsqueeze(1) if is_scale_residual: + if x_bias is not None and not use_bcg_helpers: + fused = try_fused_bias_scale_residual_norm_scale_shift( + residual_x, + x, + x_bias, + gate_x, + getattr(norm_module.norm, "weight", None), + getattr(norm_module.norm, "bias", None), + scale_result, + shift_result, + norm_module.norm_type, + norm_module.eps, + ) + if fused is not None: + modulated, residual_out = fused + return modulated, residual_out, gate_result + if x_bias is not None: + x = x + x_bias if use_bcg_helpers: modulated, residual_out = self._scale_residual_norm_scale_shift( norm_module, @@ -1324,7 +1401,12 @@ class QwenImageTransformerBlock(nn.Module): ) # QwenAttnProcessor2_0 returns (img_output, txt_output) when encoder_hidden_states is provided - img_attn_output, txt_attn_output = attn_output + ( + img_attn_output, + txt_attn_output, + img_attn_bias, + txt_attn_bias, + ) = attn_output # Process image stream - norm2 + MLP img_modulated2, hidden_states, img_gate2 = self._modulate( img_attn_output, @@ -1333,20 +1415,32 @@ class QwenImageTransformerBlock(nn.Module): modulate_index, gate_x=img_gate1, residual_x=hidden_states, + x_bias=img_attn_bias, use_bcg_helpers=use_bcg_helpers, ) - img_mlp_output = self.img_mlp(img_modulated2) + if isinstance(self.img_mlp, QwenImageFeedForward): + img_mlp_output, img_mlp_bias = self.img_mlp.forward_with_bias( + img_modulated2 + ) + else: + img_mlp_output = self.img_mlp(img_modulated2) + img_mlp_bias = None if img_mlp_output.dim() == 2: img_mlp_output = img_mlp_output.unsqueeze(0) - if use_bcg_helpers: - hidden_states = self._mul_add(img_mlp_output, img_gate2, hidden_states) - else: - hidden_states = self.fuse_mul_add(img_mlp_output, img_gate2, hidden_states) + hidden_states = self._bias_mul_add( + img_mlp_output, + img_mlp_bias, + img_gate2, + hidden_states, + use_bcg_helpers=use_bcg_helpers, + ) # Process text stream - norm2 + MLP txt_shift2, txt_scale2, txt_gate2_raw = txt_mod2.chunk(3, dim=-1) if use_bcg_helpers: + if txt_attn_bias is not None: + txt_attn_output = txt_attn_output + txt_attn_bias ( txt_modulated2, encoder_hidden_states, @@ -1358,6 +1452,15 @@ class QwenImageTransformerBlock(nn.Module): shift=txt_shift2, scale=txt_scale2, ) + elif txt_attn_bias is not None: + txt_modulated2, encoder_hidden_states, _ = self._modulate( + txt_attn_output, + txt_mod2, + self.txt_norm2, + gate_x=txt_gate1, + residual_x=encoder_hidden_states, + x_bias=txt_attn_bias, + ) else: txt_modulated2, encoder_hidden_states = self.txt_norm2( residual=encoder_hidden_states, @@ -1367,18 +1470,23 @@ class QwenImageTransformerBlock(nn.Module): scale=txt_scale2, ) txt_gate2 = txt_gate2_raw.unsqueeze(1) - txt_mlp_output = self.txt_mlp(txt_modulated2) + if isinstance(self.txt_mlp, QwenImageFeedForward): + txt_mlp_output, txt_mlp_bias = self.txt_mlp.forward_with_bias( + txt_modulated2 + ) + else: + txt_mlp_output = self.txt_mlp(txt_modulated2) + txt_mlp_bias = None if txt_mlp_output.dim() == 2: txt_mlp_output = txt_mlp_output.unsqueeze(0) - if use_bcg_helpers: - encoder_hidden_states = self._mul_add( - txt_mlp_output, txt_gate2, encoder_hidden_states - ) - else: - encoder_hidden_states = self.fuse_mul_add( - txt_mlp_output, txt_gate2, encoder_hidden_states - ) + encoder_hidden_states = self._bias_mul_add( + txt_mlp_output, + txt_mlp_bias, + txt_gate2, + encoder_hidden_states, + use_bcg_helpers=use_bcg_helpers, + ) # Clip to prevent overflow for fp16 if encoder_hidden_states.dtype == torch.float16: diff --git a/test/registered/kernels/ops/diffusion/test_qwen_output_bias_absorption.py b/test/registered/kernels/ops/diffusion/test_qwen_output_bias_absorption.py new file mode 100644 index 000000000..0bafa4737 --- /dev/null +++ b/test/registered/kernels/ops/diffusion/test_qwen_output_bias_absorption.py @@ -0,0 +1,123 @@ +import sys +from unittest.mock import patch + +import pytest +import torch + +import sglang.multimodal_gen.runtime.models.dits.qwen_image as qwen_image +from sglang.kernels.ops.diffusion import ( + try_fused_bias_mul_add, + try_fused_bias_scale_residual_norm_scale_shift, +) +from sglang.multimodal_gen.runtime.layers.elementwise import MulAdd +from sglang.multimodal_gen.runtime.layers.layernorm import ( + ScaleResidualLayerNormScaleShift, +) +from sglang.multimodal_gen.runtime.platforms.interface import DeviceCapability +from sglang.test.ci.ci_register import register_cuda_ci + +register_cuda_ci(est_time=15, stage="base-b-kernel-unit", runner_config="4-gpu-b200") + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") + +requires_sm103 = pytest.mark.skipif( + not torch.cuda.is_available() or torch.cuda.get_device_capability() != (10, 3), + reason="Qwen-Image output-bias fusion is validated on SM103", +) + + +@pytest.fixture(autouse=True) +def _seed_cuda(): + torch.cuda.manual_seed(0) + + +@pytest.mark.parametrize("quant_name", ["modelopt_fp8", "modelopt_fp4"]) +@pytest.mark.parametrize( + "capability,expected", + [ + (DeviceCapability(10, 0), False), + (DeviceCapability(10, 3), True), + (DeviceCapability(12, 0), False), + (None, False), + ], +) +def test_qwen_output_bias_absorption_is_sm103_only(quant_name, capability, expected): + class _QuantConfig: + def get_name(self): + return quant_name + + assert ( + qwen_image._can_defer_modelopt_output_bias(_QuantConfig(), capability) + is expected + ) + + +@requires_sm103 +def test_qwen_output_bias_absorption_is_bit_exact(): + hidden = 3072 + shape = (1, 64, hidden) + x = torch.randn(shape, device="cuda", dtype=torch.bfloat16) + residual = torch.randn_like(x) + bias = torch.randn(hidden, device="cuda", dtype=torch.bfloat16) + gate = torch.randn(1, 1, hidden, device="cuda", dtype=torch.bfloat16) + scale = torch.randn_like(gate) + shift = torch.randn_like(gate) + norm = ScaleResidualLayerNormScaleShift( + hidden, eps=1e-6, elementwise_affine=False + ).cuda() + + expected_norm, expected_residual = norm( + residual=residual, + x=x + bias, + gate=gate, + scale=scale, + shift=shift, + ) + actual_norm, actual_residual = try_fused_bias_scale_residual_norm_scale_shift( + residual, + x, + bias, + gate, + None, + None, + scale, + shift, + "layer", + 1e-6, + ) + expected_final = MulAdd().cuda()(x + bias, gate, residual) + actual_final = try_fused_bias_mul_add(x, bias, gate, residual) + + assert torch.equal(actual_norm, expected_norm) + assert torch.equal(actual_residual, expected_residual) + assert torch.equal(actual_final, expected_final) + + # A small residual must still break a BF16 product-rounding tie. Casting + # through FP32 before the final BF16 store loses this information at large + # magnitudes, so keep this case as a guard for the native BF16 FMA. + x.fill_(-24576.0) + bias.fill_(-0.01055908203125) + gate.fill_(206.0) + residual.fill_(-0.2080078125) + expected_tie = MulAdd().cuda()(x + bias, gate, residual) + actual_tie = try_fused_bias_mul_add(x, bias, gate, residual) + assert torch.equal(actual_tie, expected_tie) + + +def test_qwen_output_bias_absorption_rejects_unsupported_inputs(): + hidden = 3072 + x = torch.randn(2, 17, hidden, device="cuda", dtype=torch.bfloat16) + residual = torch.randn_like(x) + row = torch.randn(hidden, device="cuda", dtype=torch.bfloat16) + assert try_fused_bias_mul_add(x, row, row, residual) is None + + x = x[:1].contiguous() + residual = residual[:1].contiguous() + if torch.cuda.get_device_capability() != (10, 3): + assert try_fused_bias_mul_add(x, row, row, residual) is None + with patch("torch.compiler.is_compiling", return_value=True): + assert try_fused_bias_mul_add(x, row, row, residual) is None + + +if __name__ == "__main__": + sys.exit(pytest.main([__file__, "-v"]))