From 6e8ef213d896c36e5a27992f21e773d68b45a795 Mon Sep 17 00:00:00 2001 From: WeiQing Chen <40507679+david6666666@users.noreply.github.com> Date: Fri, 5 Jun 2026 06:25:01 +0800 Subject: [PATCH] [BugFix] Support ModelOpt FP8 under batched diffusion serving (#4155) Signed-off-by: lishunyang12 Co-authored-by: lishunyang12 --- vllm_omni/patch.py | 29 +++++++++++++++++++++++++++++ 1 file changed, 29 insertions(+) diff --git a/vllm_omni/patch.py b/vllm_omni/patch.py index e5dc955eb..29cc56e0f 100644 --- a/vllm_omni/patch.py +++ b/vllm_omni/patch.py @@ -188,3 +188,32 @@ def _patch_chat_template_registry(): _patch_chat_template_registry() + + +def _patch_scaled_mm_fp8_contiguous_activation(): + """Support batched diffusion activations on the ModelOpt FP8 (ScaledMM) path. + + The FP8 ScaledMM linear flattens its activation with ``x.view(-1, ...)``, which + needs a contiguous tensor. Under step-execution batching (``--max-num-seqs > 1``) + the sequence-packed diffusion activations can be non-contiguous, so we make the + activation contiguous before the GEMM (no-op when it already is). Mixed FP8/NVFP4 + routes through the CUTLASS NVFP4 path and is unaffected. + """ + try: + from vllm.model_executor.kernels.linear.scaled_mm.ScaledMMLinearKernel import ( + ScaledMMLinearKernel, + ) + except ImportError: + return + + _original_apply_weights = ScaledMMLinearKernel.apply_weights + + def _patched_apply_weights(self, layer, x, bias=None): + if not x.is_contiguous(): + x = x.contiguous() + return _original_apply_weights(self, layer, x, bias) + + ScaledMMLinearKernel.apply_weights = _patched_apply_weights + + +_patch_scaled_mm_fp8_contiguous_activation()