fix: bypass FP8 linear replacement for unaligned layers to prevent ROCm scaled MM kernel crashes

Signed-off-by: akshatvishu <akshatnayak197@gmail.com>

Author: akshatvishu

vllm_omni/diffusion/models/utils.py

@@ -87,7 +87,15 @@ def _recursive_replace(module: nn.Module, prefix: str):
             # Replace modules as needed
             if isinstance(child_module, nn.Linear):
                 style = "replicate"
-                new_module = replace_linear_class(child_module, style, quant_config, prefix=qual_name)
+                # AMD ROCm FP8 kernels require K (in_features) and N (out_features) to be divisible by 16.
+                # If they are not divisible, bypass FP8 replacement for this layer to avoid runtime errors.
+                is_fp8 = (
+                    quant_config is not None and hasattr(quant_config, "get_name") and quant_config.get_name() == "fp8"
+                )
+                layer_quant_config = quant_config
+                if is_fp8 and (child_module.in_features % 16 != 0 or child_module.out_features % 16 != 0):
+                    layer_quant_config = None
+                new_module = replace_linear_class(child_module, style, layer_quant_config, prefix=qual_name)
             else:
                 _recursive_replace(child_module, prefix=qual_name)
             if new_module is not child_module: