[Fix] Correct dense FP8 Marlin bias ordering (#35020)

2026-08-17 03:43:36 +00:00
parent b6d7602914
commit 92b1d382c7
2 changed files with 58 additions and 2 deletions
@@ -186,10 +186,12 @@ def prepare_fp8_layer_for_marlin(
    marlin_scales = fp8_fused_exponent_bias_into_scales(marlin_scales)
    layer.weight_scale = torch.nn.Parameter(marlin_scales, requires_grad=False)

+    # The dense FP8 Marlin wrapper adds bias after the kernel returns, so the
+    # bias must remain in logical output-channel order. Only scales need the
+    # Marlin tile permutation.
    if hasattr(layer, "bias") and layer.bias is not None:
        assert layer.bias.shape == (part_size_n,)
-        bias = marlin_permute_bias(layer.bias)
-        layer.bias = torch.nn.Parameter(bias, requires_grad=False)
+        layer.bias = torch.nn.Parameter(layer.bias.detach(), requires_grad=False)


 def prepare_moe_fp8_layer_for_marlin(