diff --git a/vllm_omni/worker/gpu_diffusion_model_runner.py b/vllm_omni/worker/gpu_diffusion_model_runner.py
index d30cf99c5..2cf308895 100644
--- a/vllm_omni/worker/gpu_diffusion_model_runner.py
+++ b/vllm_omni/worker/gpu_diffusion_model_runner.py
@@ -10,7 +10,6 @@ import logging
 import torch
 
 from vllm.v1.worker.gpu_model_runner import (
-    GPUModelRunner,
     EMPTY_MODEL_RUNNER_OUTPUT,
     IntermediateTensors,
     get_pp_group,
@@ -30,7 +29,7 @@ from vllm_omni.worker.gpu_model_runner import OmniGPUModelRunner
 logger = logging.getLogger(__name__)
 
 
-class DiffusionModelRunner(OmniGPUModelRunner):
+class GPUDiffusionModelRunner(OmniGPUModelRunner):
     """Diffusion model runner for vLLM-omni (non-autoregressive).
 
     - Reuses GPUModelRunner preparation, multimodal handling, and TP/PP/DP glue.
@@ -47,11 +46,9 @@ class DiffusionModelRunner(OmniGPUModelRunner):
         self._update_states(scheduler_output)
 
         if not scheduler_output.total_num_scheduled_tokens:
-            if not has_kv_transfer_group():
-                return EMPTY_MODEL_RUNNER_OUTPUT
-            return self.kv_connector_no_forward(scheduler_output,
-                                                self.vllm_config)
 
+            return EMPTY_MODEL_RUNNER_OUTPUT
+           
         # Prepare decoder inputs and attention metadata (for batch/order mapping)
         (attn_metadata, attention_cuda_graphs, logits_indices,
          spec_decode_metadata, num_scheduled_tokens_np,
diff --git a/vllm_omni/worker/gpu_diffusion_worker.py b/vllm_omni/worker/gpu_diffusion_worker.py
index a9767f013..be4c2047d 100644
--- a/vllm_omni/worker/gpu_diffusion_worker.py
+++ b/vllm_omni/worker/gpu_diffusion_worker.py
@@ -9,10 +9,10 @@ from vllm.v1.worker.gpu_worker import init_worker_distributed_environment
 from vllm.model_executor import set_random_seed
 from vllm.v1.utils import report_usage_stats
 
-from vllm_omni.worker.diffusion_model_runner import DiffusionModelRunner
+from vllm_omni.worker.gpu_diffusion_model_runner import GPUDiffusionModelRunner
 
 
-class DiffusionGPUWorker(GPUWorker):
+class GPUDiffusionWorker(GPUWorker):
     def init_device(self):
         if self.device_config.device.type == "cuda":
             # torch.distributed.all_reduce does not free the input tensor until
@@ -59,7 +59,7 @@ class DiffusionGPUWorker(GPUWorker):
         set_random_seed(self.model_config.seed)
 
         # Construct the model runner
-        self.model_runner: DiffusionModelRunner = DiffusionModelRunner(
+        self.model_runner: GPUDiffusionModelRunner = GPUDiffusionModelRunner(
             self.vllm_config, self.device)
 
         if self.rank == 0: