diff --git a/tests/e2e/accuracy/test_hunyuan_image3_pixel_accuracy.py b/tests/e2e/accuracy/test_hunyuan_image3_pixel_accuracy.py index 6ac930f46..3f799c038 100644 --- a/tests/e2e/accuracy/test_hunyuan_image3_pixel_accuracy.py +++ b/tests/e2e/accuracy/test_hunyuan_image3_pixel_accuracy.py @@ -48,6 +48,7 @@ _DEPLOY_CONFIG = { "trust_remote_code": True, "devices": "0,1,2,3", # set dynamically by _write_deploy_config "vae_use_slicing": False, + "moe_backend": "flashinfer_cutlass", "vae_use_tiling": False, "parallel_config": { "pipeline_parallel_size": 1, diff --git a/vllm_omni/config/omni_config.py b/vllm_omni/config/omni_config.py index dc5f4ef9c..256dcd5c7 100644 --- a/vllm_omni/config/omni_config.py +++ b/vllm_omni/config/omni_config.py @@ -469,6 +469,7 @@ class _DiffusionConfigProjection: diffusion_kv_cache_skip_layers: str | list[int] | tuple[int, ...] | set[int] | None = None diffusion_kv_cache_skip_step_indices: set[int] | None = None diffusion_kv_cache_skip_layer_indices: set[int] | None = None + moe_backend: str = "auto" force_cutlass_fp8: bool = False enable_diffusion_pipeline_profiler: bool = False step_execution: bool = False diff --git a/vllm_omni/diffusion/data.py b/vllm_omni/diffusion/data.py index 0feab5b28..e131b833b 100644 --- a/vllm_omni/diffusion/data.py +++ b/vllm_omni/diffusion/data.py @@ -664,6 +664,9 @@ class OmniDiffusionConfig: # STA_mode: STA_Mode = STA_Mode.STA_INFERENCE skip_time_steps: int = 15 + # MoE kernel backend selection + moe_backend: str = "auto" + # Compilation enforce_eager: bool = False diff --git a/vllm_omni/diffusion/worker/diffusion_worker.py b/vllm_omni/diffusion/worker/diffusion_worker.py index 8cce47897..006bbae48 100644 --- a/vllm_omni/diffusion/worker/diffusion_worker.py +++ b/vllm_omni/diffusion/worker/diffusion_worker.py @@ -288,6 +288,12 @@ class DiffusionWorker: # Since vLLM v0.20.0, IR wraps GPU ops. Set IR op priority preference to enforce GPU op fusion during wrapping. # Also need to log, because vLLM internally logs another line in VllmConfig.__post_init__. Avoid confusion. vllm_config.kernel_config.ir_op_priority = _resolve_ir_op_priority(self.od_config, vllm_config) + if self.od_config.moe_backend != "auto": + logger.warning( + "Overriding MoE backend from default 'auto' to '%s' per deploy config.", + self.od_config.moe_backend, + ) + vllm_config.kernel_config.moe_backend = self.od_config.moe_backend logger.info( "Final IR op priority after setting vLLM-Omni overrides: %s", vllm_config.kernel_config.ir_op_priority )