diff --git a/.buildkite/pipeline.yml b/.buildkite/pipeline.yml
index a2d0cdfe0..402e62518 100644
--- a/.buildkite/pipeline.yml
+++ b/.buildkite/pipeline.yml
@@ -233,43 +233,43 @@ steps:
         volumes:
         - "/fsx/hf_cache:/fsx/hf_cache"
 
-  - label: "Omni Model Test with H100"
-    timeout_in_minutes: 30
-    depends_on: image-build
-    commands:
-      - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-      - export VLLM_TEST_CLEAN_GPU_MEMORY="1"
-      - pytest -s -v tests/e2e/offline_inference/test_qwen3_omni.py
-      - pytest -s -v tests/e2e/online_serving/test_qwen3_omni.py
-      - pytest -s -v tests/e2e/online_serving/test_async_omni.py
-    agents:
-      queue: "mithril-h100-pool"
-    plugins:
-      - kubernetes:
-          podSpec:
-            containers:
-              - image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
-                resources:
-                  limits:
-                    nvidia.com/gpu: 2
-                volumeMounts:
-                  - name: devshm
-                    mountPath: /dev/shm
-                  - name: hf-cache
-                    mountPath: /root/.cache/huggingface
-                env:
-                  - name: HF_HOME
-                    value: /root/.cache/huggingface
-            nodeSelector:
-              node.kubernetes.io/instance-type: gpu-h100-sxm
-            volumes:
-              - name: devshm
-                emptyDir:
-                  medium: Memory
-              - name: hf-cache
-                hostPath:
-                  path: /mnt/hf-cache
-                  type: DirectoryOrCreate
+  # - label: "Omni Model Test with H100"
+  #   timeout_in_minutes: 30
+  #   depends_on: image-build
+  #   commands:
+  #     - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+  #     - export VLLM_TEST_CLEAN_GPU_MEMORY="1"
+  #     - pytest -s -v tests/e2e/offline_inference/test_qwen3_omni.py
+  #     - pytest -s -v tests/e2e/online_serving/test_qwen3_omni.py
+  #     - pytest -s -v tests/e2e/online_serving/test_async_omni.py
+  #   agents:
+  #     queue: "mithril-h100-pool"
+  #   plugins:
+  #     - kubernetes:
+  #         podSpec:
+  #           containers:
+  #             - image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+  #               resources:
+  #                 limits:
+  #                   nvidia.com/gpu: 2
+  #               volumeMounts:
+  #                 - name: devshm
+  #                   mountPath: /dev/shm
+  #                 - name: hf-cache
+  #                   mountPath: /root/.cache/huggingface
+  #               env:
+  #                 - name: HF_HOME
+  #                   value: /root/.cache/huggingface
+  #           nodeSelector:
+  #             node.kubernetes.io/instance-type: gpu-h100-sxm
+  #           volumes:
+  #             - name: devshm
+  #               emptyDir:
+  #                 medium: Memory
+  #             - name: hf-cache
+  #               hostPath:
+  #                 path: /mnt/hf-cache
+  #                 type: DirectoryOrCreate
 
   - label: "Diffusion Image Edit Test with H100 (1 GPU)"
     timeout_in_minutes: 20
@@ -305,3 +305,38 @@ steps:
                 hostPath:
                   path: /mnt/hf-cache
                   type: DirectoryOrCreate
+
+  # - label: "Bagel Text2Img Model Test with H100"
+  #   timeout_in_minutes: 30
+  #   depends_on: image-build
+  #   commands:
+  #     - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+  #     - pytest -s -v tests/e2e/offline_inference/test_bagel_text2img.py
+  #   agents:
+  #     queue: "mithril-h100-pool"
+  #   plugins:
+  #     - kubernetes:
+  #         podSpec:
+  #           containers:
+  #             - image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+  #               resources:
+  #                 limits:
+  #                   nvidia.com/gpu: 1
+  #               volumeMounts:
+  #                 - name: devshm
+  #                   mountPath: /dev/shm
+  #                 - name: hf-cache
+  #                   mountPath: /root/.cache/huggingface
+  #               env:
+  #                 - name: HF_HOME
+  #                   value: /root/.cache/huggingface
+  #           nodeSelector:
+  #             node.kubernetes.io/instance-type: gpu-h100-sxm
+  #           volumes:
+  #             - name: devshm
+  #               emptyDir:
+  #                 medium: Memory
+  #             - name: hf-cache
+  #               hostPath:
+  #                 path: /mnt/hf-cache
+  #                 type: DirectoryOrCreate
diff --git a/README.md b/README.md
index c4a71f126..e55025126 100644
--- a/README.md
+++ b/README.md
@@ -15,6 +15,7 @@ Easy, fast, and cheap omni-modality model serving for everyone
 ---
 
 *Latest News* 🔥
+- [2026/02] We released [0.14.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.14.0) - This is the first **stable release** of vLLM-Omni that expands Omni’s diffusion / image-video generation and audio / TTS stack, improves distributed execution and memory efficiency, and broadens platform/backend coverage (GPU/ROCm/NPU/XPU). It also brings meaningful upgrades to serving APIs, profiling & benchmarking, and overall stability.
 - [2026/01] We released [0.12.0rc1](https://github.com/vllm-project/vllm-omni/releases/tag/v0.12.0rc1) - a major RC milestone focused on maturing the diffusion stack, strengthening OpenAI-compatible serving, expanding omni-model coverage, and improving stability across platforms (GPU/NPU/ROCm), please check our latest [design](https://docs.google.com/presentation/d/1qv4qMW1rKAqDREMXiUDLIgqqHQe7TDPj/edit?usp=sharing&ouid=110473603432222024453&rtpof=true&sd=true).
 - [2025/11] vLLM community officially released [vllm-project/vllm-omni](https://github.com/vllm-project/vllm-omni) in order to support omni-modality models serving.
 
diff --git a/docs/.nav.yml b/docs/.nav.yml
index 18e9c907e..428fd16b5 100644
--- a/docs/.nav.yml
+++ b/docs/.nav.yml
@@ -7,6 +7,7 @@ nav:
   - Serving:
     - OpenAI-Compatible API:
       - Image Generation: serving/image_generation_api.md
+      - Image Edit: serving/image_edit_api.md
   - Examples:
     - examples/README.md
     - Offline Inference:
diff --git a/docs/contributing/model/adding_diffusion_model.md b/docs/contributing/model/adding_diffusion_model.md
index e5257749a..80212360f 100644
--- a/docs/contributing/model/adding_diffusion_model.md
+++ b/docs/contributing/model/adding_diffusion_model.md
@@ -204,7 +204,151 @@ Key point for writing the example:
 
 + Save or display the generated results so users can validate the integration.
 
-## Step 5: Open a Pull Request
+## Step 5: TeaCache Coefficient Estimation (Optional)
+
+If your model supports TeaCache acceleration, you need to estimate the polynomial coefficients for optimal caching performance.
+
+### 5.1 Add Extractor Function
+
+First, implement an extractor function in `vllm_omni/diffusion/cache/teacache/extractors.py`. The extractor extracts the modulated input and defines how to run transformer blocks:
+
+```python
+def extract_your_model_context(
+    module: nn.Module,
+    hidden_states: torch.Tensor,
+    timestep: torch.Tensor,
+    **kwargs: Any,
+) -> CacheContext:
+    # 1. Preprocessing
+    temb = module.time_embed(timestep)
+
+    # 2. Extract modulated input (for cache decision)
+    modulated_input = module.transformer_blocks[0].norm1(hidden_states, temb)
+
+    # 3. Define transformer execution
+    def run_transformer_blocks():
+        h = hidden_states
+        for block in module.transformer_blocks:
+            h = block(h, temb=temb)
+        return (h,)
+
+    # 4. Define postprocessing
+    def postprocess(h):
+        return module.proj_out(module.norm_out(h, temb))
+
+    return CacheContext(
+        modulated_input=modulated_input,
+        hidden_states=hidden_states,
+        encoder_hidden_states=None,
+        temb=temb,
+        run_transformer_blocks=run_transformer_blocks,
+        postprocess=postprocess,
+    )
+```
+
+Register it in `EXTRACTOR_REGISTRY`:
+```python
+EXTRACTOR_REGISTRY = {
+    ...
+    "YourTransformer2DModel": extract_your_model_context,
+}
+```
+
+### 5.2 Add Adapter for Coefficient Estimation
+
+Add an adapter in `vllm_omni/diffusion/cache/teacache/coefficient_estimator.py`:
+
+```python
+class YourModelAdapter:
+    @staticmethod
+    def load_pipeline(model_path: str, device: str, dtype: torch.dtype) -> Any:
+        # Load your pipeline
+        ...
+
+    @staticmethod
+    def get_transformer(pipeline: Any) -> tuple[Any, str]:
+        return pipeline.transformer, "YourTransformer2DModel"
+
+    @staticmethod
+    def install_hook(transformer: Any, hook: DataCollectionHook) -> None:
+        registry = HookRegistry.get_or_create(transformer)
+        registry.register_hook(hook._HOOK_NAME, hook)
+
+_MODEL_ADAPTERS["YourModel"] = YourModelAdapter
+```
+
+### 5.3 Run Coefficient Estimation
+
+Use the provided script to estimate coefficients:
+
+```python
+from vllm_omni.diffusion.cache.teacache.coefficient_estimator import (
+    TeaCacheCoefficientEstimator,
+)
+from datasets import load_dataset
+from tqdm import tqdm
+
+# Load model
+estimator = TeaCacheCoefficientEstimator(
+    model_path="/path/to/model",
+    model_type="Bagel",  # Your model type
+    device="cuda",
+)
+
+# Load prompts (paper suggests ~70 prompts)
+dataset = load_dataset("nateraw/parti-prompts", split="train")
+prompts = dataset["Prompt"][:70]
+
+# Collect data
+for prompt in tqdm(prompts):
+    estimator.collect_from_prompt(prompt, num_inference_steps=50)
+
+# Estimate coefficients
+coeffs = estimator.estimate(poly_order=4)
+print(f"Coefficients: {coeffs}")
+```
+
+### 5.4 Interpreting Coefficient Estimation Results
+
+The estimator outputs statistics and polynomial coefficients. Here's how to interpret them:
+
+**Example Output:**
+```
+Data statistics:
+Count: 48
+Input Diffs (x): min=1.1089e-02, max=5.2555e-02, mean=2.8435e-02
+Output Diffs (y): min=2.8242e-02, max=2.9792e-01, mean=7.0312e-02
+Coefficients: [1333131.29, -168644.23, 7950.51, -163.75, 1.26]
+```
+
+**What to Check:**
+- **Count**: Number of timestep pairs analyzed. Should be at least 30-50 for reliable estimation. Low count suggests insufficient prompts or inference steps.
+- **Input/Output Ranges**: Verify output differences correlate with input differences. If ranges seem unusual, check your prompt diversity.
+- **Coefficient Magnitude**: Extremely large values (>1e8) may indicate numerical instability - try collecting more diverse data.
+
+**Troubleshooting:**
+- If results seem unreliable, try:
+  - Increasing number of prompts (100+ recommended)
+  - Using more diverse prompts from multiple datasets
+  - Adjusting `num_inference_steps` (try 20, 50, 100)
+
+### 5.5 Add Coefficients to Config
+
+Add the estimated coefficients to `vllm_omni/diffusion/cache/teacache/config.py`:
+
+```python
+_MODEL_COEFFICIENTS = {
+    ...
+    "YourTransformer2DModel": [
+        1.04730573e+06,  # a4
+        -1.34150749e+05, # a3
+        6.51517806e+03,  # a2
+        -1.41209108e+02, # a1
+        1.17241808e+00,  # a0
+    ],
+}
+```
+## Step 6: Open a Pull Request
 
 When submitting a pull request to add support for a new model, please include the following information in the PR description:
 
diff --git a/docs/getting_started/installation/gpu/cuda.inc.md b/docs/getting_started/installation/gpu/cuda.inc.md
index c4b34ab15..3996a0243 100644
--- a/docs/getting_started/installation/gpu/cuda.inc.md
+++ b/docs/getting_started/installation/gpu/cuda.inc.md
@@ -17,7 +17,7 @@ Therefore, it is recommended to install vLLM and vLLM-Omni with a **fresh new**
 # --8<-- [start:pre-built-wheels]
 
 #### Installation of vLLM
-Note: Pre-built wheels are currently only available for vLLM-Omni 0.11.0rc1, 0.12.0rc1, 0.14.0rc1. For the latest version, please [build from source](https://docs.vllm.ai/projects/vllm-omni/en/latest/getting_started/installation/gpu/#build-wheel-from-source).
+Note: Pre-built wheels are currently only available for vLLM-Omni 0.11.0rc1, 0.12.0rc1, 0.14.0rc1, 0.14.0. For the latest version, please [build from source](https://docs.vllm.ai/projects/vllm-omni/en/latest/getting_started/installation/gpu/#build-wheel-from-source).
 
 
 vLLM-Omni is built based on vLLM. Please install it with command below.
@@ -93,7 +93,7 @@ docker run --runtime nvidia --gpus 2 \
     --env "HF_TOKEN=$HF_TOKEN" \
     -p 8091:8091 \
     --ipc=host \
-    vllm/vllm-omni:v0.14.0rc1 \
+    vllm/vllm-omni:v0.14.0 \
     --model Qwen/Qwen3-Omni-30B-A3B-Instruct --port 8091
 ```
 
diff --git a/docs/getting_started/installation/gpu/rocm.inc.md b/docs/getting_started/installation/gpu/rocm.inc.md
index 60ad1f2b2..1a8ffb612 100644
--- a/docs/getting_started/installation/gpu/rocm.inc.md
+++ b/docs/getting_started/installation/gpu/rocm.inc.md
@@ -80,7 +80,7 @@ docker run --rm \
   -v ~/.cache/huggingface:/root/.cache/huggingface \
   --env "HF_TOKEN=$HF_TOKEN" \
   -p 8091:8091 \
-  vllm/vllm-omni-rocm:v0.14.0rc1 \
+  vllm/vllm-omni-rocm:v0.14.0 \
   --model Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091
 ```
 
@@ -99,7 +99,7 @@ docker run --rm -it \
   -v ~/.cache/huggingface:/root/.cache/huggingface \
   --env "HF_TOKEN=$HF_TOKEN" \
   --entrypoint bash \
-  vllm/vllm-omni-rocm:v0.14.0rc1
+  vllm/vllm-omni-rocm:v0.14.0
 ```
 
 # --8<-- [end:pre-built-images]
diff --git a/docs/serving/image_edit_api.md b/docs/serving/image_edit_api.md
new file mode 100644
index 000000000..d254ac06a
--- /dev/null
+++ b/docs/serving/image_edit_api.md
@@ -0,0 +1,205 @@
+# Image Edit API
+
+vLLM-Omni provides an OpenAI DALL-E compatible API for image edit using diffusion models.
+
+Each server instance runs a single model (specified at startup via `vllm serve <model> --omni`).
+
+## Quick Start
+
+### Start the Server
+
+For example...
+
+```bash
+# Qwen-Image
+vllm serve Qwen/Qwen-Image-Edit-2511 --omni --port 8000
+
+
+### Generate Images
+
+**Using curl:**
+
+```bash
+curl -s -D >(grep -i x-request-id >&2) \
+  -o >(jq -r '.data[0].b64_json' | base64 --decode > gift-basket.png) \
+  -X POST "http://localhost:8000/v1/images/edits" \
+  -F "model=xxx" \
+  -F "image=@./xx.png" \
+  -F "prompt='this bear is wearing sportwear. holding a basketball, and bending one leg.'" \
+  -F "size=1024x1024" \
+  -F "output_format=png"
+```
+
+
+**Using OpenAI SDK:**
+
+```python
+import base64
+from openai import OpenAI
+from pathlib import Path
+client = OpenAI(
+    api_key="None",
+    base_url="http://localhost:8000/v1"
+)
+
+input_image_url = "https://vllm-public-assets.s3.us-west-2.amazonaws.com/omni-assets/qwen-bear.png"
+
+result = client.images.edit(
+    image=[],
+    model="Qwen-Image-Edit-2511",
+    prompt="Change the bears in the two input images into walking together.",
+    size='512x512',
+    stream=False,
+    output_format='jpeg',
+    # url格式
+    extra_body={
+        "url": [input_image_url1,input_image_url],
+        "num_inference_steps": 50,
+        "guidance_scale": 1,
+        "seed": 777,
+    }
+)
+
+image_base64 = result.data[0].b64_json
+image_bytes = base64.b64decode(image_base64)
+
+# Save the image to a file
+with open("edit_out_http.jpeg", "wb") as f:
+    f.write(image_bytes)
+```
+
+## API Reference
+
+### Endpoint
+
+```
+POST /v1/images/edits
+Content-Type: multipart/form-data
+```
+
+### Request Parameters
+
+#### OpenAI Standard Parameters
+
+| Parameter | Type | Default | Description |
+|-----------|------|---------|-------------|
+| `prompt` | string | **required** | A text description of the desired image |
+| `model` | string | server's model | Model to use (optional, should match server if specified) |
+| `image` | string or array | **required** | The image(s) to edit. |
+| `n` | integer | 1 | Number of images to generate (1-10) |
+| `size` | string | "auto" | Image dimensions in WxH format (e.g., "1024x1024", "512x512"), when set to auto, it decide size from first input image. |
+| `response_format` | string | "b64_json" | Response format (only "b64_json" supported) |
+| `user` | string | null | User identifier for tracking |
+| `output_format` | string | "png" | The format in which the generated images are returned. Must be one of "png", "jpg", "jpeg", "webp". |
+| `output_compression` | integer | 100 | The compression level (0-100%) for the generated images. |
+| `background` | string or null | "auto" | Allows to set transparency for the background of the generated image(s).
+
+#### vllm-omni Extension Parameters
+
+| Parameter | Type | Default | Description |
+|-----------|------|---------|-------------|
+| `url` | string or array | None | The image(s) to edit. |
+| `negative_prompt` | string | null | Text describing what to avoid in the image |
+| `num_inference_steps` | integer | model defaults | Number of diffusion steps |
+| `guidance_scale` | float | model defaults | Classifier-free guidance scale (typically 0.0-20.0) |
+| `true_cfg_scale` | float | model defaults | True CFG scale (model-specific parameter, may be ignored if not supported) |
+| `seed` | integer | null | Random seed for reproducibility |
+
+### Response Format
+
+```json
+{
+  "created": 1701234567,
+  "data": [
+    {
+      "b64_json": "<base64-encoded PNG>",
+      "url": null,
+      "revised_prompt": null
+    }
+  ],
+  "output_format": null,
+  "size": null,
+}
+```
+
+## Examples
+
+### Multiple Images input
+
+```bash
+curl -s -D >(grep -i x-request-id >&2) \
+  -o >(jq -r '.data[0].b64_json' | base64 --decode > gift-basket.png) \
+  -X POST "http://localhost:8000/v1/images/edits" \
+  -F "model=xxx" \
+  -F "image=@xx.png" \
+  -F "image=@xx.png"
+  -F "prompt='this bear is wearing sportwear. holding a basketball, and bending one leg.'" \
+  -F "size=1024x1024" \
+  -F "output_format=png"
+```
+
+
+## Parameter Handling
+
+The API passes parameters directly to the diffusion pipeline without model-specific transformation:
+
+- **Default values**: When parameters are not specified, the underlying model uses its own defaults
+- **Pass-through design**: User-provided values are forwarded directly to the diffusion engine
+- **Minimal validation**: Only basic type checking and range validation at the API level
+
+### Parameter Compatibility
+
+The API passes parameters directly to the diffusion pipeline without model-specific validation.
+
+- Unsupported parameters may be silently ignored by the model
+- Incompatible values will result in errors from the underlying pipeline
+- Recommended values vary by model - consult model documentation
+
+**Best Practice:** Start with the model's recommended parameters, then adjust based on your needs.
+
+## Error Responses
+
+### 400 Bad Request
+
+Invalid parameters (e.g., model mismatch):
+
+```json
+{
+  "detail": "Invalid size format: '1024x'. Expected format: 'WIDTHxHEIGHT' (e.g., '1024x1024')."
+}
+```
+
+### 422 Unprocessable Entity
+
+Validation errors (missing required fields):
+
+```json
+{
+  "detail": "Field 'image' or 'url' is required"
+}
+```
+
+## Troubleshooting
+
+### Server Not Running
+
+```bash
+# Check if server is responding
+curl -X http://localhost:8000/v1/images/edit \
+  -F "prompt='test'"
+```
+
+### Out of Memory
+
+If you encounter OOM errors:
+1. Reduce image size: `"size": "512x512"`
+2. Reduce inference steps: `"num_inference_steps": 25`
+
+## Development
+
+Enable debug logging to see prompts and generation details:
+
+```bash
+vllm serve Qwen/Qwen-Image-Edit-2511 --omni \
+  --uvicorn-log-level debug
+```
diff --git a/docs/user_guide/diffusion/parallelism_acceleration.md b/docs/user_guide/diffusion/parallelism_acceleration.md
index 7334df1d4..6e2c18d64 100644
--- a/docs/user_guide/diffusion/parallelism_acceleration.md
+++ b/docs/user_guide/diffusion/parallelism_acceleration.md
@@ -18,20 +18,19 @@ The following table shows which models are currently supported by parallelism me
 
 ### ImageGen
 
-| Model | Model Identifier | Ulysses-SP | Ring-SP | CFG-Parallel | Tensor-Parallel |
-|-------|------------------|------------|---------|--------------|--------------------------|
-| **LongCat-Image** | `meituan-longcat/LongCat-Image` | ✅ | ✅ | ❌ | ❌ |
-| **LongCat-Image-Edit** | `meituan-longcat/LongCat-Image-Edit` | ✅ | ✅ | ❌ | ❌ |
-| **Ovis-Image** | `OvisAI/Ovis-Image` | ❌ | ❌ | ❌ | ❌ |
-| **Qwen-Image** | `Qwen/Qwen-Image` | ✅ | ✅ | ✅ | ✅ |
-| **Qwen-Image-Edit** | `Qwen/Qwen-Image-Edit` | ✅ | ✅ | ✅ | ✅ |
-| **Qwen-Image-Edit-2509** | `Qwen/Qwen-Image-Edit-2509` | ✅ | ✅ | ✅ | ✅ |
-| **Qwen-Image-Layered** | `Qwen/Qwen-Image-Layered` | ✅ | ✅ | ✅ | ✅ |
-| **Z-Image** | `Tongyi-MAI/Z-Image-Turbo` | ✅ | ✅ | ❌ | ✅ (TP=2 only) |
-| **Stable-Diffusion3.5** | `stabilityai/stable-diffusion-3.5` | ❌ | ❌ | ❌ | ❌ |
-| **FLUX.2-klein** | `black-forest-labs/FLUX.2-klein-4B` | ❌ | ❌ | ❌ | ✅ |
-| **FLUX.1-dev** | `black-forest-labs/FLUX.1-dev` | ❌ | ❌ | ❌ | ✅ |
-
+| Model                    | Model Identifier                     | Ulysses-SP | Ring-SP | CFG-Parallel | Tensor-Parallel |
+|--------------------------|--------------------------------------|:----------:|:-------:|:------------:|:---------------:|
+| **LongCat-Image**        | `meituan-longcat/LongCat-Image`      |     ✅      |    ✅    |      ❌       |        ✅        |
+| **LongCat-Image-Edit**   | `meituan-longcat/LongCat-Image-Edit` |     ✅      |    ✅    |      ❌       |        ✅        |
+| **Ovis-Image**           | `OvisAI/Ovis-Image`                  |     ❌      |    ❌    |      ❌       |        ❌        |
+| **Qwen-Image**           | `Qwen/Qwen-Image`                    |     ✅      |    ✅    |      ✅       |        ✅        |
+| **Qwen-Image-Edit**      | `Qwen/Qwen-Image-Edit`               |     ✅      |    ✅    |      ✅       |        ✅        |
+| **Qwen-Image-Edit-2509** | `Qwen/Qwen-Image-Edit-2509`          |     ✅      |    ✅    |      ✅       |        ✅        |
+| **Qwen-Image-Layered**   | `Qwen/Qwen-Image-Layered`            |     ✅      |    ✅    |      ✅       |        ✅        |
+| **Z-Image**              | `Tongyi-MAI/Z-Image-Turbo`           |     ✅      |    ✅    |      ❌       |  ✅ (TP=2 only)  |
+| **Stable-Diffusion3.5**  | `stabilityai/stable-diffusion-3.5`   |     ❌      |    ❌    |      ❌       |        ❌        |
+| **FLUX.2-klein**         | `black-forest-labs/FLUX.2-klein-4B`  |     ❌      |    ❌    |      ❌       |        ✅        |
+| **FLUX.1-dev**           | `black-forest-labs/FLUX.1-dev`       |     ❌      |    ❌    |      ❌       |        ✅        |
 
 !!! note "TP Limitations for Diffusion Models"
     We currently implement Tensor Parallelism (TP) only for the DiT (Diffusion Transformer) blocks. This is because the `text_encoder` component in vLLM-Omni uses the original Transformers implementation, which does not yet support TP.
diff --git a/pyproject.toml b/pyproject.toml
index 3e1149099..25454102f 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "vllm-omni"
-version = "0.14.0rc1"
+version = "0.14.0"
 description = "A framework for efficient model inference with omni-modality models"
 readme = "README.md"
 requires-python = ">=3.10,<3.14"
@@ -59,7 +59,8 @@ dev = [
     "psutil>=7.2.0",
     "soundfile>=0.13.1",
     "imageio[ffmpeg]>=0.6.0",
-    "opencv-python>=4.12.0.88"
+    "opencv-python>=4.12.0.88",
+    "mooncake-transfer-engine==0.3.8.post1"
 ]
 
 docs = [
diff --git a/tests/e2e/offline_inference/stage_configs/bagel_mooncake_ci.yaml b/tests/e2e/offline_inference/stage_configs/bagel_mooncake_ci.yaml
new file mode 100644
index 000000000..8c170dffd
--- /dev/null
+++ b/tests/e2e/offline_inference/stage_configs/bagel_mooncake_ci.yaml
@@ -0,0 +1,85 @@
+# stage config for running BAGEL with Mooncake connector for CI e2e tests.
+# This config is optimized for single GPU tests with Mooncake inter-stage communication.
+
+stage_args:
+  - stage_id: 0
+    stage_type: llm
+    runtime:
+      devices: "0"
+      max_batch_size: 1
+    engine_args:
+      model_stage: thinker
+      model_arch: BagelForConditionalGeneration
+      worker_type: ar
+      scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
+      gpu_memory_utilization: 0.35
+      enforce_eager: true
+      trust_remote_code: true
+      engine_output_type: text
+      distributed_executor_backend: mp
+      enable_prefix_caching: false
+      max_num_batched_tokens: 32768
+      tensor_parallel_size: 1
+      omni_kv_config:
+        need_send_cache: true
+        kv_transfer_criteria:
+          type: prefill_finished
+    final_output: true
+    final_output_type: text
+    is_comprehension: true
+    default_sampling_params:
+      temperature: 0.4
+      top_p: 0.9
+      top_k: 1
+      max_tokens: 2048
+      seed: 52
+      detokenize: true
+      repetition_penalty: 1.05
+    output_connectors:
+      to_stage_1: mooncake_connector
+  - stage_id: 1
+    stage_type: diffusion
+    runtime:
+      devices: "0"
+      max_batch_size: 1
+    engine_args:
+      model_stage: dit
+      gpu_memory_utilization: 0.55
+      enforce_eager: true
+      trust_remote_code: true
+      engine_output_type: image
+      distributed_executor_backend: mp
+      enable_prefix_caching: false
+      max_num_batched_tokens: 32768
+      tensor_parallel_size: 1
+      omni_kv_config:
+        need_recv_cache: true
+    engine_input_source: [0]
+    final_output: true
+    final_output_type: image
+    is_comprehension: false
+    default_sampling_params:
+      seed: 52
+    input_connectors:
+      from_stage_0: mooncake_connector
+
+# Top-level runtime config with Mooncake connector
+runtime:
+  enabled: true
+  defaults:
+    window_size: -1
+    max_inflight: 1
+  connectors:
+    mooncake_connector:
+      name: MooncakeConnector
+      extra:
+        host: "${MOONCAKE_HOST}"
+        metadata_server: "http://${MOONCAKE_HOST}:${MOONCAKE_HTTP_PORT}/metadata"
+        master: "${MOONCAKE_HOST}:${MOONCAKE_RPC_PORT}"
+        segment: 64000000
+        localbuf: 64000000
+        proto: tcp
+  edges:
+    - from: 0
+      to: 1
+      window_size: -1
diff --git a/tests/e2e/offline_inference/stage_configs/bagel_sharedmemory_ci.yaml b/tests/e2e/offline_inference/stage_configs/bagel_sharedmemory_ci.yaml
new file mode 100644
index 000000000..ef605a3a3
--- /dev/null
+++ b/tests/e2e/offline_inference/stage_configs/bagel_sharedmemory_ci.yaml
@@ -0,0 +1,83 @@
+# stage config for running BAGEL with SharedMemory connector for CI e2e tests.
+# This config is optimized for single GPU tests with SharedMemory inter-stage communication.
+
+stage_args:
+  - stage_id: 0
+    stage_type: llm
+    runtime:
+      devices: "0"
+      max_batch_size: 1
+    engine_args:
+      model_stage: thinker
+      model_arch: BagelForConditionalGeneration
+      worker_type: ar
+      scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
+      gpu_memory_utilization: 0.35
+      enforce_eager: true
+      trust_remote_code: true
+      engine_output_type: text
+      distributed_executor_backend: "mp"
+      enable_prefix_caching: false
+      max_num_batched_tokens: 32768
+      tensor_parallel_size: 1
+      omni_kv_config:
+        need_send_cache: true
+        kv_transfer_criteria:
+          type: prefill_finished #or special token generated
+    final_output: true
+    final_output_type: text
+    is_comprehension: true
+    default_sampling_params:
+      temperature: 0.4
+      top_p: 0.9
+      top_k: 1
+      max_tokens: 2048
+      seed: 52
+      detokenize: True
+      repetition_penalty: 1.05
+
+  - stage_id: 1
+    stage_type: diffusion
+    runtime:
+      devices: "0"
+      max_batch_size: 1
+    engine_args:
+      model_stage: dit
+      gpu_memory_utilization: 0.55
+      enforce_eager: true
+      trust_remote_code: true
+      engine_output_type: image
+      distributed_executor_backend: "mp"
+      enable_prefix_caching: false
+      max_num_batched_tokens: 32768
+      tensor_parallel_size: 1
+      omni_kv_config:
+        need_recv_cache: true
+    engine_input_source: [0]
+
+    final_output: true
+    final_output_type: image
+    is_comprehension: false
+    default_sampling_params:
+      seed: 52
+
+# Runtime edges
+runtime:
+  enabled: true
+  defaults:
+    window_size: -1
+    max_inflight: 1
+
+  # Distributed connectors configuration (optional)
+  # More connectors will be supported in the future.
+  connectors:
+    shared_memory_connector:
+      name: SharedMemoryConnector
+      extra:
+        shm_threshold_bytes: 65536 # 64KB threshold
+
+
+  edges:
+    - from: 0
+      to: 1
+      window_size: -1
diff --git a/tests/e2e/offline_inference/test_bagel_text2img.py b/tests/e2e/offline_inference/test_bagel_text2img.py
new file mode 100644
index 000000000..9aef47a74
--- /dev/null
+++ b/tests/e2e/offline_inference/test_bagel_text2img.py
@@ -0,0 +1,308 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""
+End-to-end test for Bagel text2img generation.
+
+This test validates that the Bagel model generates images that match
+expected reference pixel values within a ±5 tolerance.
+
+Equivalent to running:
+    python3 examples/offline_inference/bagel/end2end.py \
+        --prompts "A futuristic city skyline at twilight, cyberpunk style" \
+        --modality text2img --step 15
+"""
+
+import os
+
+os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
+os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "1"
+
+import signal
+import socket
+import subprocess
+import tempfile
+import time
+from pathlib import Path
+from typing import Any
+
+import pytest
+from PIL import Image
+
+from tests.utils import hardware_test
+from vllm_omni.entrypoints.omni import Omni
+
+# Reference pixel data extracted from the known-good output image
+# Each entry contains (x, y) position and expected (R, G, B) values
+# "Generated with seed=52, num_inference_steps=15,
+# prompt='A futuristic city skyline at twilight, cyberpunk style'"
+REFERENCE_PIXELS = [
+    {"position": (100, 100), "rgb": (68, 107, 134)},
+    {"position": (400, 50), "rgb": (95, 139, 166)},
+    {"position": (700, 100), "rgb": (99, 122, 151)},
+    {"position": (150, 400), "rgb": (111, 125, 153)},
+    {"position": (512, 512), "rgb": (97, 107, 131)},
+    {"position": (700, 400), "rgb": (48, 64, 98)},
+    {"position": (100, 700), "rgb": (79, 63, 84)},
+    {"position": (400, 700), "rgb": (40, 58, 79)},
+    {"position": (700, 700), "rgb": (60, 75, 103)},
+    {"position": (256, 256), "rgb": (97, 128, 156)},
+]
+
+# Maximum allowed difference per color channel
+PIXEL_TOLERANCE = 5
+
+# Default test prompt
+DEFAULT_PROMPT = "<|im_start|>A futuristic city skyline at twilight, cyberpunk style<|im_end|>"
+
+
+def _find_free_port() -> int:
+    """Find and return a free ephemeral port by binding to port 0."""
+    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
+        s.bind(("127.0.0.1", 0))
+        s.listen(1)
+        port = s.getsockname()[1]
+    return port
+
+
+def _configure_sampling_params(omni: Omni, max_tokens: int = 1, num_inference_steps: int = 15) -> list:
+    """Configure sampling parameters for Bagel text2img generation.
+
+    Args:
+        omni: The Omni instance to get default params from.
+        max_tokens: Maximum tokens for the first stage.
+        num_inference_steps: Number of inference steps for the diffusion stage.
+
+    Returns:
+        Configured sampling params list.
+    """
+    params_list = omni.default_sampling_params_list
+    params_list[0].max_tokens = max_tokens  # type: ignore
+    if len(params_list) > 1:
+        params_list[1].num_inference_steps = num_inference_steps  # type: ignore
+    return params_list
+
+
+def _extract_generated_image(omni_outputs: list) -> Image.Image | None:
+    """Extract the generated image from Omni outputs.
+
+    Args:
+        omni_outputs: List of outputs from omni.generate().
+
+    Returns:
+        The first generated PIL Image, or None if no image found.
+    """
+    for req_output in omni_outputs:
+        if images := getattr(req_output, "images", None):
+            return images[0]
+        if hasattr(req_output, "request_output") and req_output.request_output:
+            for stage_out in req_output.request_output:
+                if hasattr(stage_out, "images") and stage_out.images:
+                    return stage_out.images[0]
+    return None
+
+
+def _validate_pixels(
+    image: Image.Image,
+    reference_pixels: list[dict[str, Any]] = REFERENCE_PIXELS,
+    tolerance: int = PIXEL_TOLERANCE,
+) -> None:
+    """Validate that image pixels match expected reference values.
+
+    Args:
+        image: The PIL Image to validate.
+        reference_pixels: List of dicts with 'position' (x, y) and 'rgb' (R, G, B).
+        tolerance: Maximum allowed difference per color channel.
+
+    Raises:
+        AssertionError: If any pixel differs beyond tolerance.
+    """
+    for ref in reference_pixels:
+        x, y = ref["position"]
+        expected = ref["rgb"]
+        actual = image.getpixel((x, y))[:3]
+        assert all(abs(a - e) <= tolerance for a, e in zip(actual, expected)), (
+            f"Pixel mismatch at ({x}, {y}): expected {expected}, got {actual}"
+        )
+
+
+def _generate_bagel_image(omni: Omni, prompt: str = DEFAULT_PROMPT) -> Image.Image:
+    """Generate an image using Bagel model with configured parameters.
+
+    Args:
+        omni: The Omni instance to use for generation.
+        prompt: The text prompt for image generation.
+
+    Returns:
+        The generated PIL Image.
+
+    Raises:
+        AssertionError: If no image is generated or size is incorrect.
+    """
+    params_list = _configure_sampling_params(omni)
+
+    omni_outputs = list(
+        omni.generate(
+            prompts=[{"prompt": prompt, "modalities": ["image"]}],
+            sampling_params_list=params_list,
+        )
+    )
+
+    generated_image = _extract_generated_image(omni_outputs)
+    assert generated_image is not None, "No images generated"
+    assert generated_image.size == (1024, 1024), f"Expected 1024x1024, got {generated_image.size}"
+
+    return generated_image
+
+
+@pytest.mark.core_model
+@pytest.mark.diffusion
+@hardware_test(res={"cuda": "H100"})
+def test_bagel_text2img_shared_memory_connector():
+    """Test Bagel text2img with shared memory connector."""
+    config_path = str(Path(__file__).parent / "stage_configs" / "bagel_sharedmemory_ci.yaml")
+    omni = Omni(model="ByteDance-Seed/BAGEL-7B-MoT", stage_configs_path=config_path, stage_init_timeout=300)
+
+    try:
+        generated_image = _generate_bagel_image(omni)
+        _validate_pixels(generated_image)
+    finally:
+        omni.close()
+
+
+def _wait_for_port(host: str, port: int, timeout: int = 30) -> bool:
+    """Wait for a port to become available.
+
+    Args:
+        host: The host address.
+        port: The port number.
+        timeout: Maximum seconds to wait.
+
+    Returns:
+        True if port becomes available, False otherwise.
+    """
+    for _ in range(timeout):
+        try:
+            with socket.create_connection((host, port), timeout=1):
+                return True
+        except (TimeoutError, ConnectionRefusedError):
+            time.sleep(1)
+    return False
+
+
+def _cleanup_mooncake_processes(timeout_secs: int = 5) -> None:
+    """Clean up any existing mooncake_master processes.
+
+    Args:
+        timeout_secs: Maximum seconds to wait for graceful termination.
+    """
+    subprocess.run(
+        ["pkill", "-f", "mooncake_master"],
+        stdout=subprocess.DEVNULL,
+        stderr=subprocess.DEVNULL,
+    )
+    start_time = time.time()
+    while time.time() - start_time < timeout_secs:
+        result = subprocess.run(
+            ["pgrep", "-f", "mooncake_master"],
+            stdout=subprocess.DEVNULL,
+            stderr=subprocess.DEVNULL,
+        )
+        if result.returncode != 0:
+            break
+        time.sleep(0.5)
+    else:
+        subprocess.run(
+            ["pkill", "-9", "-f", "mooncake_master"],
+            stdout=subprocess.DEVNULL,
+            stderr=subprocess.DEVNULL,
+        )
+    time.sleep(1)
+
+
+def _load_mooncake_config(host: str, rpc_port: int, http_port: int) -> str:
+    """Load Mooncake config from YAML and substitute placeholders.
+
+    Args:
+        host: Mooncake host address.
+        rpc_port: RPC port for Mooncake master.
+        http_port: HTTP metadata server port.
+
+    Returns:
+        Path to the temporary config file with substituted values.
+    """
+    config_path = str(Path(__file__).parent / "stage_configs" / "bagel_mooncake_ci.yaml")
+    with open(config_path) as f:
+        config_content = f.read()
+
+    # Substitute placeholders
+    config_content = config_content.replace("${MOONCAKE_HOST}", host)
+    config_content = config_content.replace("${MOONCAKE_RPC_PORT}", str(rpc_port))
+    config_content = config_content.replace("${MOONCAKE_HTTP_PORT}", str(http_port))
+
+    # Write to temp file
+    temp_file = tempfile.NamedTemporaryFile(mode="w", suffix=".yaml", delete=False)
+    temp_file.write(config_content)
+    temp_file.close()
+    return temp_file.name
+
+
+@pytest.mark.core_model
+@pytest.mark.diffusion
+@hardware_test(res={"cuda": "H100"})
+def test_bagel_text2img_mooncake_connector():
+    """Test Bagel text2img with Mooncake connector for inter-stage communication."""
+    MOONCAKE_HOST = "127.0.0.1"
+    MOONCAKE_RPC_PORT = _find_free_port()
+    MOONCAKE_HTTP_PORT = _find_free_port()
+    MOONCAKE_METRICS_PORT = _find_free_port()
+
+    mooncake_master_proc = None
+    temp_config_file = None
+    omni = None
+
+    try:
+        _cleanup_mooncake_processes()
+
+        # Start mooncake_master
+        mooncake_master_proc = subprocess.Popen(
+            [
+                "mooncake_master",
+                f"--rpc_port={MOONCAKE_RPC_PORT}",
+                "--enable_http_metadata_server=true",
+                "--http_metadata_server_host=0.0.0.0",
+                f"--http_metadata_server_port={MOONCAKE_HTTP_PORT}",
+                f"--metrics_port={MOONCAKE_METRICS_PORT}",
+            ],
+            stdout=subprocess.DEVNULL,
+            stderr=subprocess.DEVNULL,
+            preexec_fn=os.setsid,
+        )
+
+        assert _wait_for_port(MOONCAKE_HOST, MOONCAKE_RPC_PORT), "mooncake_master failed to start"
+
+        # Create temp config and initialize Omni
+        temp_config_file = _load_mooncake_config(
+            host=MOONCAKE_HOST,
+            rpc_port=MOONCAKE_RPC_PORT,
+            http_port=MOONCAKE_HTTP_PORT,
+        )
+
+        omni = Omni(model="ByteDance-Seed/BAGEL-7B-MoT", stage_configs_path=temp_config_file, stage_init_timeout=300)
+
+        generated_image = _generate_bagel_image(omni)
+        _validate_pixels(generated_image)
+
+    finally:
+        if omni:
+            omni.close()
+        if temp_config_file:
+            try:
+                os.unlink(temp_config_file)
+            except OSError:
+                pass
+        if mooncake_master_proc:
+            try:
+                os.killpg(os.getpgid(mooncake_master_proc.pid), signal.SIGKILL)
+            except OSError:
+                pass
diff --git a/tests/entrypoints/openai_api/test_image_server.py b/tests/entrypoints/openai_api/test_image_server.py
index 575dfaa33..3df484b1e 100644
--- a/tests/entrypoints/openai_api/test_image_server.py
+++ b/tests/entrypoints/openai_api/test_image_server.py
@@ -9,6 +9,7 @@ OpenAI-compatible async text-to-image generation API endpoints in api_server.py.
 
 import base64
 import io
+from argparse import Namespace
 from unittest.mock import AsyncMock, Mock
 
 import pytest
@@ -115,9 +116,11 @@ class FakeAsyncOmni:
         self.stage_list = ["llm", "diffusion"]
         self.default_sampling_params_list = [SamplingParams(temperature=0.1), OmniDiffusionSamplingParams()]
         self.captured_sampling_params_list = None
+        self.captured_prompt = None
 
     async def generate(self, prompt, request_id, sampling_params_list):
         self.captured_sampling_params_list = sampling_params_list
+        self.captured_prompt = prompt
         images = [Image.new("RGB", (64, 64), color="green")]
         yield MockGenerationResult(images)
 
@@ -133,6 +136,8 @@ def mock_async_diffusion():
         # Return n PIL images wrapped in result object
         print("!!!!!!!!!!!!!!!!!!!!! kwargs", kwargs)
         n = kwargs["sampling_params_list"][0].num_outputs_per_prompt
+        mock.captured_sampling_params_list = kwargs["sampling_params_list"]
+        mock.captured_prompt = kwargs["prompt"]
         images = [Image.new("RGB", (64, 64), color="blue") for _ in range(n)]
         return MockGenerationResult(images)
 
@@ -155,6 +160,10 @@ def test_client(mock_async_diffusion):
     app.state.diffusion_engine = mock_async_diffusion  # Also set for health endpoint
     app.state.stage_configs = [{"stage_type": "diffusion"}]
     app.state.diffusion_model_name = "Qwen/Qwen-Image"  # For models endpoint
+    app.state.args = Namespace(
+        default_sampling_params='{"0": {"num_inference_steps":4, "guidance_scale":7.5}}',
+        max_generated_image_size=4096,  # 64*64
+    )
 
     return TestClient(app)
 
@@ -171,7 +180,10 @@ def async_omni_test_client():
 
     app.state.engine_client = FakeAsyncOmni()
     app.state.stage_configs = [{"stage_type": "llm"}, {"stage_type": "diffusion"}]
-
+    app.state.args = Namespace(
+        default_sampling_params='{"1": {"num_inference_steps":4, "guidance_scale":7.5}}',
+        max_generated_image_size=4096,  # 64*64
+    )
     return TestClient(app)
 
 
@@ -526,3 +538,279 @@ def test_model_field_omitted_works(test_client):
         },
     )
     assert response.status_code == 200
+
+
+def make_test_image_bytes(size=(64, 64)) -> bytes:
+    img = Image.new(
+        "RGB",
+        size,
+    )
+    buf = io.BytesIO()
+    img.save(buf, format="PNG")
+    return buf.getvalue()
+
+
+def test_image_edit_images_processing(async_omni_test_client):
+    img_bytes_1 = make_test_image_bytes((16, 16))
+    img_bytes_2 = make_test_image_bytes((32, 32))
+
+    # uploadfile with image key
+    response = async_omni_test_client.post(
+        "/v1/images/edits",
+        files=[
+            ("image", img_bytes_1),
+            ("image", img_bytes_2),
+        ],
+        data={"prompt": "hello world."},
+    )
+    assert response.status_code == 200
+    engine = async_omni_test_client.app.state.engine_client
+    captured_prompt = engine.captured_prompt
+    processed_images = captured_prompt["multi_modal_data"]["image"]
+    assert len(processed_images) == 2
+    assert isinstance(processed_images[0], Image.Image)
+    assert isinstance(processed_images[1], Image.Image)
+    assert processed_images[0].size == (16, 16)
+    assert processed_images[1].size == (32, 32)
+
+    # uploadfile with image[] key
+    response = async_omni_test_client.post(
+        "/v1/images/edits",
+        files=[
+            ("image[]", img_bytes_2),
+            ("image[]", img_bytes_1),
+        ],
+        data={"prompt": "hello world."},
+    )
+
+    assert response.status_code == 200
+    engine = async_omni_test_client.app.state.engine_client
+    captured_prompt = engine.captured_prompt
+    processed_images = captured_prompt["multi_modal_data"]["image"]
+    assert len(processed_images) == 2
+    assert isinstance(processed_images[0], Image.Image)
+    assert isinstance(processed_images[1], Image.Image)
+    assert processed_images[0].size == (32, 32)
+    assert processed_images[1].size == (16, 16)
+
+    # base64 url
+    buf1 = io.BytesIO()
+    img1 = Image.new("RGB", (16, 16))
+    img1.save(buf1, format="PNG")
+    b64_1 = "data:image/png;base64," + base64.b64encode(buf1.getvalue()).decode()
+
+    buf2 = io.BytesIO()
+    img2 = Image.new("RGB", (24, 24))
+    img2.save(buf2, format="PNG")
+    b64_2 = "data:image/png;base64," + base64.b64encode(buf2.getvalue()).decode()
+
+    response = async_omni_test_client.post(
+        "/v1/images/edits",
+        data={
+            "prompt": "hello from base64",
+            "url": [b64_1, b64_2],
+        },
+    )
+    assert response.status_code == 200
+    processed_images = engine.captured_prompt["multi_modal_data"]["image"]
+    assert len(processed_images) == 2
+    assert isinstance(processed_images[0], Image.Image)
+    assert isinstance(processed_images[1], Image.Image)
+    assert processed_images[0].size == (16, 16)
+    assert processed_images[1].size == (24, 24)
+
+
+def test_image_edit_parameter_pass(async_omni_test_client):
+    img_bytes_1 = make_test_image_bytes((16, 16))
+
+    # uploadfile with image key
+    response = async_omni_test_client.post(
+        "/v1/images/edits",
+        files=[("image", img_bytes_1)],
+        data={
+            "prompt": "hello world.",
+            "size": "16x24",
+            "output_format": "jpeg",
+            "num_inference_steps": 20,
+            "guidance_scale": 8.0,
+            "seed": 1234,
+            "negative_prompt": "negative",
+            "n": 2,
+        },
+    )
+    assert response.status_code == 200
+    engine = async_omni_test_client.app.state.engine_client
+    captured_prompt = engine.captured_prompt
+    captured_sampling_params = engine.captured_sampling_params_list[-1]
+
+    assert captured_prompt["prompt"] == "hello world."
+    assert captured_prompt["negative_prompt"] == "negative"
+    assert captured_sampling_params.num_inference_steps == 20
+    assert captured_sampling_params.guidance_scale == 8.0
+    assert captured_sampling_params.seed == 1234
+    assert captured_sampling_params.num_outputs_per_prompt == 2
+    assert captured_sampling_params.width == 16
+    assert captured_sampling_params.height == 24
+
+    data = response.json()
+    # All images should be valid
+    for img_data in data["data"]:
+        assert "b64_json" in img_data
+        img_bytes = base64.b64decode(img_data["b64_json"])
+        img = Image.open(io.BytesIO(img_bytes))
+        assert img.format.lower() == "jpeg"
+        assert data["output_format"] == "jpeg"
+        assert data["size"] == "16x24"
+
+
+def test_image_edit_parameter_default(async_omni_test_client):
+    img_bytes_1 = make_test_image_bytes((24, 16))
+
+    # uploadfile with image key
+    response = async_omni_test_client.post(
+        "/v1/images/edits",
+        files=[("image", img_bytes_1)],
+        data={
+            "prompt": "hello world.",
+            "size": "auto",
+        },
+    )
+    assert response.status_code == 200
+    engine = async_omni_test_client.app.state.engine_client
+    captured_sampling_params = engine.captured_sampling_params_list[-1]
+
+    assert captured_sampling_params.width == 24
+    assert captured_sampling_params.height == 16
+    assert captured_sampling_params.num_outputs_per_prompt == 1
+    assert captured_sampling_params.num_inference_steps == 4
+    assert captured_sampling_params.guidance_scale == 7.5
+
+    response = async_omni_test_client.post(
+        "/v1/images/edits",
+        files=[("image", img_bytes_1)],
+        data={
+            "prompt": "hello world.",
+            "size": "96x96",
+        },
+    )
+    assert response.status_code == 400
+
+
+def test_image_edit_parameter_default_single_stage(test_client):
+    img_bytes_1 = make_test_image_bytes((24, 16))
+
+    # uploadfile with image key
+    response = test_client.post(
+        "/v1/images/edits",
+        files=[("image", img_bytes_1)],
+        data={
+            "prompt": "hello world.",
+        },
+    )
+    assert response.status_code == 200
+    engine = test_client.app.state.engine_client
+    captured_sampling_params = engine.captured_sampling_params_list[0]
+
+    assert captured_sampling_params.width == 24
+    assert captured_sampling_params.height == 16
+    assert captured_sampling_params.num_outputs_per_prompt == 1
+    assert captured_sampling_params.num_inference_steps == 4
+    assert captured_sampling_params.guidance_scale == 7.5
+
+    response = test_client.post(
+        "/v1/images/edits",
+        files=[("image", img_bytes_1)],
+        data={
+            "prompt": "hello world.",
+            "size": "96x96",
+        },
+    )
+    assert response.status_code == 400
+
+
+def test_image_edit_compression_jpeg(test_client):
+    img_bytes_1 = make_test_image_bytes((16, 16))
+    # uploadfile with image key
+    response = test_client.post(
+        "/v1/images/edits",
+        files=[("image", img_bytes_1)],
+        data={"prompt": "hello world.", "output_format": "jpeg", "output_compression": 100},
+    )
+    assert response.status_code == 200
+    data = response.json()
+    img_bytes_100 = base64.b64decode(data["data"][0]["b64_json"])
+    img = Image.open(io.BytesIO(img_bytes_100))
+    assert img.format.lower() == "jpeg"
+
+    response = test_client.post(
+        "/v1/images/edits",
+        files=[("image", img_bytes_1)],
+        data={
+            "prompt": "hello world.",
+            "output_format": "jpeg",
+            "output_compression": 50,
+        },
+    )
+    assert response.status_code == 200
+    data = response.json()
+    img_bytes_50 = base64.b64decode(data["data"][0]["b64_json"])
+
+    response = test_client.post(
+        "/v1/images/edits",
+        files=[("image", img_bytes_1)],
+        data={
+            "prompt": "hello world.",
+            "output_format": "jpeg",
+            "output_compression": 10,
+        },
+    )
+    assert response.status_code == 200
+    data = response.json()
+    img_bytes_10 = base64.b64decode(data["data"][0]["b64_json"])
+
+    assert len(img_bytes_10) < len(img_bytes_50)
+    assert len(img_bytes_50) < len(img_bytes_100)
+
+
+def test_image_edit_compression_png(async_omni_test_client):
+    img_bytes_1 = make_test_image_bytes((16, 16))
+    # uploadfile with image key
+    response = async_omni_test_client.post(
+        "/v1/images/edits",
+        files=[("image", img_bytes_1)],
+        data={"prompt": "hello world.", "output_format": "PNG", "output_compression": 100},
+    )
+    assert response.status_code == 200
+    data = response.json()
+    img_bytes_100 = base64.b64decode(data["data"][0]["b64_json"])
+    img = Image.open(io.BytesIO(img_bytes_100))
+    assert img.format.lower() == "png"
+
+    response = async_omni_test_client.post(
+        "/v1/images/edits",
+        files=[("image", img_bytes_1)],
+        data={
+            "prompt": "hello world.",
+            "output_format": "PNG",
+            "output_compression": 50,
+        },
+    )
+    assert response.status_code == 200
+    data = response.json()
+    img_bytes_50 = base64.b64decode(data["data"][0]["b64_json"])
+
+    response = async_omni_test_client.post(
+        "/v1/images/edits",
+        files=[("image", img_bytes_1)],
+        data={
+            "prompt": "hello world.",
+            "output_format": "PNG",
+            "output_compression": 10,
+        },
+    )
+    assert response.status_code == 200
+    data = response.json()
+    img_bytes_10 = base64.b64decode(data["data"][0]["b64_json"])
+
+    assert len(img_bytes_10) < len(img_bytes_50)
+    assert len(img_bytes_50) < len(img_bytes_100)
diff --git a/vllm_omni/diffusion/attention/backends/flash_attn.py b/vllm_omni/diffusion/attention/backends/flash_attn.py
index dc54acf04..45158e8f8 100644
--- a/vllm_omni/diffusion/attention/backends/flash_attn.py
+++ b/vllm_omni/diffusion/attention/backends/flash_attn.py
@@ -10,26 +10,8 @@ from vllm_omni.diffusion.attention.backends.abstract import (
     AttentionMetadata,
 )
 
-# Import flash attention functions with fallback chain from utils/fa.py
-# FA3 (fa3_fwd_interface) -> FA3 (flash_attn_interface) -> FA2 (flash_attn)
-from vllm_omni.diffusion.attention.backends.utils.fa import (
-    HAS_FLASH_ATTN,
-    _pad_input,
-    _unpad_input,
-    _upad_input,
-    flash_attn_func,
-    flash_attn_varlen_func,
-)
-
 logger = init_logger(__name__)
 
-if not HAS_FLASH_ATTN:
-    raise ImportError(
-        "FlashAttentionBackend requires Flash Attention. "
-        "Please install one of: fa3-fwd, flash-attention, or flash-attn. "
-        "Otherwise, use SDPA backend by setting DIFFUSION_ATTENTION_BACKEND=TORCH_SDPA"
-    )
-
 
 class FlashAttentionBackend(AttentionBackend):
     accept_output_buffer: bool = True
@@ -74,6 +56,24 @@ class FlashAttentionImpl(AttentionImpl):
         attn_metadata: AttentionMetadata = None,
     ) -> torch.Tensor:
         """CUDA/ROCm flash attention implementation."""
+        # Import flash attention functions with fallback chain from utils/fa.py
+        # FA3 (fa3_fwd_interface) -> FA3 (flash_attn_interface) -> FA2 (flash_attn)
+        from vllm_omni.diffusion.attention.backends.utils.fa import (
+            HAS_FLASH_ATTN,
+            _pad_input,
+            _unpad_input,
+            _upad_input,
+            flash_attn_func,
+            flash_attn_varlen_func,
+        )
+
+        if not HAS_FLASH_ATTN:
+            raise ImportError(
+                "FlashAttentionBackend requires Flash Attention. "
+                "Please install one of: fa3-fwd, flash-attention, or flash-attn. "
+                "Otherwise, use SDPA backend by setting DIFFUSION_ATTENTION_BACKEND=TORCH_SDPA"
+            )
+
         query_length = query.size(1)
         attention_mask = attn_metadata.attn_mask if attn_metadata is not None else None
         #  Contains at least one padding token in the sequence
@@ -122,7 +122,15 @@ class FlashAttentionImpl(AttentionImpl):
         attn_metadata: AttentionMetadata = None,
     ) -> torch.Tensor:
         """NPU attention implementation using mindiesd."""
-        from mindiesd import attention_forward
+        try:
+            from mindiesd import attention_forward
+        except ImportError:
+            raise ImportError(
+                "FlashAttentionBackend NPU implementation requires MindIE-SD. "
+                "Please install MindIE-SD to enable NPU attention support. "
+                "For installation details, see https://gitcode.com/Ascend/MindIE-SD"
+                "Otherwise, use SDPA backend by setting DIFFUSION_ATTENTION_BACKEND=TORCH_SDPA"
+            )
 
         attention_mask = attn_metadata.attn_mask if attn_metadata else None
         output = attention_forward(
diff --git a/vllm_omni/diffusion/attention/backends/sdpa.py b/vllm_omni/diffusion/attention/backends/sdpa.py
index 16bc5fcf8..cfe460e3d 100644
--- a/vllm_omni/diffusion/attention/backends/sdpa.py
+++ b/vllm_omni/diffusion/attention/backends/sdpa.py
@@ -13,6 +13,29 @@ from vllm_omni.diffusion.attention.backends.abstract import (
 logger = init_logger(__name__)
 
 
+def _maybe_reshape_attn_mask(query: torch.Tensor, key: torch.Tensor, attn_mask: torch.Tensor | None = None):
+    """
+    Reshape Attention Mask
+    [batch_size, seq_len_k] -> [batch_size, 1, seq_len_q, seq_len_k]
+    """
+    # Skip Attention Mask if all values are 1, `None` mask can speedup the computation
+    if attn_mask is not None and torch.all(attn_mask != 0):
+        attn_mask = None
+
+    # Reshape Attention Mask
+    # [batch_size, seq_len_k] -> [batch_size, 1, seq_len_q, seq_len_k]
+    if (
+        attn_mask is not None
+        and attn_mask.ndim == 2
+        and attn_mask.shape[0] == query.shape[0]
+        and attn_mask.shape[1] == key.shape[1]
+    ):
+        B, Sq, Skv = attn_mask.shape[0], query.shape[1], key.shape[1]
+        attn_mask = attn_mask.to(torch.bool)
+        attn_mask = attn_mask.unsqueeze(1).expand(B, Sq, Skv).unsqueeze(1).contiguous()
+    return attn_mask
+
+
 class SDPABackend(AttentionBackend):
     accept_output_buffer: bool = True
 
@@ -47,16 +70,15 @@ class SDPAImpl(AttentionImpl):
         self.causal = causal
         self.softmax_scale = softmax_scale
 
-    def forward(
+    def forward_cuda(
         self,
         query: torch.Tensor,
         key: torch.Tensor,
         value: torch.Tensor,
-        attn_metadata: AttentionMetadata = None,
+        attn_metadata: AttentionMetadata | None = None,
     ) -> torch.Tensor:
         query, key, value = (x.permute(0, 2, 1, 3) for x in (query, key, value))
         attention_mask = attn_metadata.attn_mask if attn_metadata else None
-
         output = torch.nn.functional.scaled_dot_product_attention(
             query,
             key,
@@ -68,3 +90,33 @@ class SDPAImpl(AttentionImpl):
         )
         out = output.permute(0, 2, 1, 3)
         return out
+
+    def forward_xpu(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        attn_metadata: AttentionMetadata | None = None,
+    ) -> torch.Tensor:
+        return self.forward_cuda(query, key, value, attn_metadata)
+
+    def forward_hip(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        attn_metadata: AttentionMetadata | None = None,
+    ) -> torch.Tensor:
+        return self.forward_cuda(query, key, value, attn_metadata)
+
+    def forward_npu(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        attn_metadata: AttentionMetadata | None = None,
+    ) -> torch.Tensor:
+        if attn_metadata:
+            attention_mask = _maybe_reshape_attn_mask(query, key, attn_metadata.attn_mask)
+            setattr(attn_metadata, "attn_mask", attention_mask)
+        return self.forward_cuda(query, key, value, attn_metadata)
diff --git a/vllm_omni/diffusion/cache/teacache/coefficient_estimator.py b/vllm_omni/diffusion/cache/teacache/coefficient_estimator.py
new file mode 100644
index 000000000..cdacf49dc
--- /dev/null
+++ b/vllm_omni/diffusion/cache/teacache/coefficient_estimator.py
@@ -0,0 +1,197 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import types
+from typing import Any
+
+import numpy as np
+import torch
+from vllm.config import LoadConfig
+
+from vllm_omni.diffusion.cache.teacache.extractors import get_extractor
+from vllm_omni.diffusion.data import OmniDiffusionConfig
+from vllm_omni.diffusion.hooks import HookRegistry, ModelHook
+from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
+from vllm_omni.diffusion.models.bagel.pipeline_bagel import BagelPipeline
+
+
+class DataCollectionHook(ModelHook):
+    """Hook to collect modulated inputs and model outputs for TeaCache coefficient estimation."""
+
+    _HOOK_NAME = "teacache_collector"
+
+    def __init__(self, transformer_type: str):
+        super().__init__()
+        self.transformer_type = transformer_type
+        self.extractor_fn = None
+        self.current_trajectory: list[tuple[np.ndarray, np.ndarray]] = []
+
+    def initialize_hook(self, module: torch.nn.Module) -> torch.nn.Module:
+        self.extractor_fn = get_extractor(self.transformer_type)
+        return module
+
+    def new_forward(self, module: torch.nn.Module, *args: Any, **kwargs: Any) -> Any:
+        ctx = self.extractor_fn(module, *args, **kwargs)
+        modulated_input_cpu = ctx.modulated_input.detach().cpu().numpy()
+
+        outputs = ctx.run_transformer_blocks()
+        ctx.hidden_states = outputs[0]
+        if len(outputs) > 1 and ctx.encoder_hidden_states is not None:
+            ctx.encoder_hidden_states = outputs[1]
+
+        model_output_cpu = ctx.hidden_states.detach().cpu().numpy()
+        self.current_trajectory.append((modulated_input_cpu, model_output_cpu))
+
+        return ctx.postprocess(ctx.hidden_states)
+
+    def start_collection(self):
+        self.current_trajectory = []
+
+    def stop_collection(self) -> list[tuple[np.ndarray, np.ndarray]]:
+        return list(self.current_trajectory)
+
+
+class BagelAdapter:
+    """Adapter for Bagel model."""
+
+    @staticmethod
+    def load_pipeline(model_path: str, device: str = "cuda", dtype: torch.dtype = torch.bfloat16) -> BagelPipeline:
+        od_config = OmniDiffusionConfig.from_kwargs(model=model_path, dtype=dtype)
+        od_config.model_class_name = "BagelPipeline"
+
+        pipeline = BagelPipeline(od_config=od_config)
+        loader = DiffusersPipelineLoader(LoadConfig())
+        loader.load_weights(pipeline)
+        pipeline.to(device)
+        return pipeline
+
+    @staticmethod
+    def get_transformer(pipeline: Any) -> tuple[Any, str]:
+        return pipeline.bagel, "Bagel"
+
+    @staticmethod
+    def install_hook(transformer: Any, hook: DataCollectionHook) -> None:
+        original_forward_flow = transformer._forward_flow
+
+        def forward_alias(self, *args, **kwargs):
+            return original_forward_flow(*args, **kwargs)
+
+        transformer.forward = types.MethodType(forward_alias, transformer)
+        registry = HookRegistry.get_or_create(transformer)
+        registry.register_hook(hook._HOOK_NAME, hook)
+        transformer._forward_flow = transformer.forward
+
+
+class DefaultAdapter:
+    """Default adapter for standard diffusers pipelines."""
+
+    @staticmethod
+    def load_pipeline(model_path: str, device: str, dtype: torch.dtype) -> Any:
+        raise NotImplementedError("DefaultAdapter.load_pipeline not implemented")
+
+    @staticmethod
+    def get_transformer(pipeline: Any) -> tuple[Any, str]:
+        return pipeline.transformer, pipeline.transformer.__class__.__name__
+
+    @staticmethod
+    def install_hook(transformer: Any, hook: DataCollectionHook) -> None:
+        registry = HookRegistry.get_or_create(transformer)
+        registry.register_hook(hook._HOOK_NAME, hook)
+
+
+_MODEL_ADAPTERS: dict[str, type] = {
+    "Bagel": BagelAdapter,
+}
+
+_EPSILON = 1e-6
+
+
+def calculate_relative_l1(tensor_current: np.ndarray, tensor_next: np.ndarray) -> float:
+    """Calculate relative L1 distance (Eq. 4 from TeaCache paper)."""
+    diff = np.abs(tensor_current - tensor_next).sum()
+    norm = np.abs(tensor_current).sum() + _EPSILON
+    return diff / norm
+
+
+def estimate_teacache_coefficients(
+    collected_data: list[list[tuple[np.ndarray, np.ndarray]]], poly_order: int = 4
+) -> list[float]:
+    """Estimate polynomial coefficients for TeaCache using np.polyfit."""
+    input_diffs, output_diffs = [], []
+
+    for sample in collected_data:
+        for t in range(len(sample) - 1):
+            feat_in_curr, feat_out_curr = sample[t]
+            feat_in_next, feat_out_next = sample[t + 1]
+            input_diffs.append(calculate_relative_l1(feat_in_curr, feat_in_next))
+            output_diffs.append(calculate_relative_l1(feat_out_curr, feat_out_next))
+
+    x = np.array(input_diffs, dtype=np.float64)
+    y = np.array(output_diffs, dtype=np.float64)
+
+    print("Data statistics:")
+    print(f"  Count: {len(x)}")
+    print(f"  Input Diffs (x): min={x.min():.4e}, max={x.max():.4e}, mean={x.mean():.4e}")
+    print(f"  Output Diffs (y): min={y.min():.4e}, max={y.max():.4e}, mean={y.mean():.4e}")
+
+    return np.polyfit(x, y, poly_order).tolist()
+
+
+class TeaCacheCoefficientEstimator:
+    """Model-agnostic helper class to collect data and estimate TeaCache coefficients."""
+
+    def __init__(
+        self,
+        model_path: str,
+        model_type: str = "Bagel",
+        device: str = "cuda",
+        dtype: torch.dtype = torch.bfloat16,
+    ):
+        # Add validation here ⬇️
+        if model_type not in _MODEL_ADAPTERS:
+            available_types = list(_MODEL_ADAPTERS.keys())
+            raise ValueError(
+                f"Unsupported model_type: '{model_type}'. "
+                f"Available types: {available_types}. "
+                f"To add support for a new model, add an entry to _MODEL_ADAPTERS."
+            )
+
+        adapter = _MODEL_ADAPTERS.get(model_type, DefaultAdapter)
+        self.pipeline = adapter.load_pipeline(model_path, device, dtype)
+        self.transformer, self.transformer_type = adapter.get_transformer(self.pipeline)
+        self.hook = DataCollectionHook(self.transformer_type)
+        self.collected_data: list[list[tuple[np.ndarray, np.ndarray]]] = []
+        adapter.install_hook(self.transformer, self.hook)
+
+    def collect_from_prompt(self, prompt: str, **generate_kwargs):
+        self.hook.start_collection()
+        from vllm_omni.diffusion.request import OmniDiffusionRequest
+
+        req = OmniDiffusionRequest(
+            prompt=prompt,
+            num_inference_steps=generate_kwargs.get("num_inference_steps", 20),
+            seed=generate_kwargs.get("seed", 42),
+        )
+        self.pipeline.forward(req)
+        trajectory = self.hook.stop_collection()
+        if trajectory:
+            self.collected_data.append(trajectory)
+
+    def estimate(self, poly_order: int = 4) -> list[float]:
+        """Estimate polynomial coefficients from collected data.
+
+        Args:
+            poly_order: Order of polynomial fit (default: 4)
+
+        Returns:
+            List of polynomial coefficients [a_n, a_{n-1}, ..., a_1, a_0]
+
+        Raises:
+            RuntimeError: If no data has been collected
+        """
+        if not self.collected_data:
+            raise RuntimeError(
+                "No data collected for coefficient estimation. "
+                "Call collect_from_prompt() at least once before calling estimate()."
+            )
+        return estimate_teacache_coefficients(self.collected_data, poly_order)
diff --git a/vllm_omni/diffusion/models/longcat_image/longcat_image_transformer.py b/vllm_omni/diffusion/models/longcat_image/longcat_image_transformer.py
index 9d216363d..2f6d54aef 100644
--- a/vllm_omni/diffusion/models/longcat_image/longcat_image_transformer.py
+++ b/vllm_omni/diffusion/models/longcat_image/longcat_image_transformer.py
@@ -9,10 +9,11 @@ import torch.nn as nn
 from diffusers.models.embeddings import TimestepEmbedding, Timesteps, apply_rotary_emb, get_1d_rotary_pos_embed
 from diffusers.models.modeling_outputs import Transformer2DModelOutput
 from diffusers.models.normalization import AdaLayerNormContinuous, AdaLayerNormZero, AdaLayerNormZeroSingle
+from vllm.distributed import get_tensor_model_parallel_world_size, tensor_model_parallel_all_gather
 from vllm.logger import init_logger
 from vllm.model_executor.layers.activation import get_act_fn
 from vllm.model_executor.layers.layernorm import RMSNorm
-from vllm.model_executor.layers.linear import QKVParallelLinear, ReplicatedLinear
+from vllm.model_executor.layers.linear import ColumnParallelLinear, QKVParallelLinear, RowParallelLinear
 from vllm.model_executor.model_loader.weight_utils import default_weight_loader
 
 from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
@@ -35,9 +36,9 @@ class FeedForward(nn.Module):
         inner_dim = int(dim * mult)
         dim_out = dim_out if dim_out is not None else dim
 
-        self.w_in = ReplicatedLinear(dim, inner_dim, bias=bias, return_bias=False)
+        self.w_in = ColumnParallelLinear(dim, inner_dim, bias=bias, return_bias=False)
         self.act = get_act_fn("gelu_pytorch_tanh")
-        self.w_out = ReplicatedLinear(inner_dim, dim_out, bias=bias, return_bias=False)
+        self.w_out = RowParallelLinear(inner_dim, dim_out, bias=bias, return_bias=False)
 
     def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
         hidden_states = self.w_in(hidden_states)
@@ -83,12 +84,11 @@ class LongCatImageAttention(nn.Module):
             hidden_size=query_dim,
             head_size=self.head_dim,
             total_num_heads=self.heads,
-            disable_tp=True,
             bias=bias,
         )
 
         if not self.pre_only:
-            self.to_out = torch.nn.Linear(self.inner_dim, self.out_dim, bias=out_bias)
+            self.to_out = RowParallelLinear(self.inner_dim, self.out_dim, bias=out_bias)
 
         if self.added_kv_proj_dim is not None:
             self.norm_added_q = RMSNorm(dim_head, eps=eps)
@@ -98,11 +98,10 @@ class LongCatImageAttention(nn.Module):
                 hidden_size=self.added_kv_proj_dim,
                 head_size=self.head_dim,
                 total_num_heads=self.heads,
-                disable_tp=True,
                 bias=added_proj_bias,
             )
 
-            self.to_add_out = ReplicatedLinear(self.inner_dim, query_dim, bias=out_bias)
+            self.to_add_out = RowParallelLinear(self.inner_dim, query_dim, bias=out_bias)
 
         self.attn = Attention(
             num_heads=heads,
@@ -185,22 +184,26 @@ class LongCatImageAttention(nn.Module):
         """
         qkv, _ = self.to_qkv(hidden_states)
 
-        query, key, value = qkv.chunk(3, dim=-1)
+        q_size = self.to_qkv.num_heads * self.head_dim
+        kv_size = self.to_qkv.num_kv_heads * self.head_dim
+        query, key, value = qkv.split([q_size, kv_size, kv_size], dim=-1)
 
-        query = query.unflatten(-1, (self.heads, -1))
-        key = key.unflatten(-1, (self.heads, -1))
-        value = value.unflatten(-1, (self.heads, -1))
+        query = query.unflatten(-1, (self.to_qkv.num_heads, -1))
+        key = key.unflatten(-1, (self.to_qkv.num_kv_heads, -1))
+        value = value.unflatten(-1, (self.to_qkv.num_kv_heads, -1))
 
         query = self.norm_q(query)
         key = self.norm_k(key)
 
         if self.added_kv_proj_dim is not None:
             encoder_qkv, _ = self.add_kv_proj(encoder_hidden_states)
-            encoder_query, encoder_key, encoder_value = encoder_qkv.chunk(3, dim=-1)
+            q_size = self.add_kv_proj.num_heads * self.head_dim
+            kv_size = self.add_kv_proj.num_kv_heads * self.head_dim
+            encoder_query, encoder_key, encoder_value = encoder_qkv.split([q_size, kv_size, kv_size], dim=-1)
 
-            encoder_query = encoder_query.unflatten(-1, (self.heads, -1))
-            encoder_key = encoder_key.unflatten(-1, (self.heads, -1))
-            encoder_value = encoder_value.unflatten(-1, (self.heads, -1))
+            encoder_query = encoder_query.unflatten(-1, (self.add_kv_proj.num_heads, -1))
+            encoder_key = encoder_key.unflatten(-1, (self.add_kv_proj.num_kv_heads, -1))
+            encoder_value = encoder_value.unflatten(-1, (self.add_kv_proj.num_kv_heads, -1))
 
             # Apply RMSNorm to text Q/K
             encoder_query = self.norm_added_q(encoder_query)
@@ -284,11 +287,14 @@ class LongCatImageAttention(nn.Module):
             encoder_hidden_states, hidden_states = hidden_states.split_with_sizes(
                 [encoder_hidden_states.shape[1], hidden_states.shape[1] - encoder_hidden_states.shape[1]], dim=1
             )
-            hidden_states = self.to_out(hidden_states)
+            hidden_states, _ = self.to_out(hidden_states)
             encoder_hidden_states, _ = self.to_add_out(encoder_hidden_states)
 
             return hidden_states, encoder_hidden_states
         else:
+            # For single-stream blocks, there's no to_out (RowParallelLinear) to handle the reduction
+            if get_tensor_model_parallel_world_size() > 1:
+                hidden_states = tensor_model_parallel_all_gather(hidden_states, dim=-1)
             return hidden_states
 
 
diff --git a/vllm_omni/entrypoints/cli/serve.py b/vllm_omni/entrypoints/cli/serve.py
index 083114ee4..540d6d3b4 100644
--- a/vllm_omni/entrypoints/cli/serve.py
+++ b/vllm_omni/entrypoints/cli/serve.py
@@ -232,6 +232,22 @@ class OmniServeCommand(CLISubcommand):
             help="Number of devices for CFG parallel computation for diffusion models. "
             "Equivalent to setting DiffusionParallelConfig.cfg_parallel_size.",
         )
+
+        # Default sampling parameters
+        omni_config_group.add_argument(
+            "--default-sampling-params",
+            type=str,
+            help="Json str for Default sampling parameters, \n"
+            'Structure: {"<stage_id>": {<sampling_param>: value, ...}, ...}\n'
+            'e.g., \'{"0": {"num_inference_steps":50, "guidance_scale":1}}\'. '
+            "Currently only supports diffusion models.",
+        )
+        # Diffusion model mixed precision
+        omni_config_group.add_argument(
+            "--max-generated-image-size",
+            type=float,
+            help="The max size of generate image (height * width).",
+        )
         return serve_parser
 
 
diff --git a/vllm_omni/entrypoints/openai/api_server.py b/vllm_omni/entrypoints/openai/api_server.py
index e428a6276..2e057901a 100644
--- a/vllm_omni/entrypoints/openai/api_server.py
+++ b/vllm_omni/entrypoints/openai/api_server.py
@@ -1,5 +1,8 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import base64
+import io
+import json
 import multiprocessing
 import multiprocessing.forkserver as forkserver
 import os
@@ -12,11 +15,13 @@ from argparse import Namespace
 from collections.abc import AsyncIterator
 from contextlib import asynccontextmanager
 from http import HTTPStatus
-from typing import Any, cast
+from typing import Annotated, Any, cast
 
+import httpx
 import vllm.envs as envs
-from fastapi import Depends, HTTPException, Request
+from fastapi import Depends, File, Form, HTTPException, Request, UploadFile
 from fastapi.responses import JSONResponse, StreamingResponse
+from PIL import Image
 from starlette.datastructures import State
 from starlette.routing import Route
 from vllm import SamplingParams
@@ -890,6 +895,251 @@ async def generate_images(request: ImageGenerationRequest, raw_request: Request)
     Raises:
         HTTPException: For validation errors, missing engine, or generation failures
     """
+    # Get engine client (AsyncOmni) from app state
+    engine_client, model_name, stage_types = _get_engine_and_model(raw_request)
+
+    # Validate model field (warn if mismatch, don't error)
+    if request.model is not None and request.model != model_name:
+        logger.warning(
+            f"Model mismatch: request specifies '{request.model}' but "
+            f"server is running '{model_name}'. Using server model."
+        )
+
+    try:
+        # Build params - pass through user values directly
+        prompt: OmniTextPrompt = {"prompt": request.prompt}
+        if request.negative_prompt is not None:
+            prompt["negative_prompt"] = request.negative_prompt
+        gen_params = OmniDiffusionSamplingParams(num_outputs_per_prompt=request.n)
+
+        # Parse per-request LoRA (compatible with chat's extra_body.lora shape).
+        lora_request, lora_scale = _parse_lora_request(request.lora)
+        _update_if_not_none(gen_params, "lora_request", lora_request)
+        _update_if_not_none(gen_params, "lora_scale", lora_scale)
+
+        # Parse and add size if provided
+        width, height = None, None
+        if request.size:
+            width, height = parse_size(request.size)
+            size_str = f"{width}x{height}"
+        else:
+            size_str = "model default"
+        _update_if_not_none(gen_params, "width", width)
+        _update_if_not_none(gen_params, "height", height)
+
+        # 3.3 Add optional parameters ONLY if provided
+        _update_if_not_none(gen_params, "num_inference_steps", request.num_inference_steps)
+        _update_if_not_none(gen_params, "guidance_scale", request.guidance_scale)
+        _update_if_not_none(gen_params, "true_cfg_scale", request.true_cfg_scale)
+        # If seed is not provided, generate a random one to ensure
+        # a proper generator is initialized in the backend.
+        # This fixes issues where using the default global generator
+        # might produce blurry images in some environments.
+        _update_if_not_none(gen_params, "seed", random.randint(0, 2**32 - 1) if request.seed is None else request.seed)
+
+        request_id = f"img_gen_{uuid.uuid4().hex}"
+
+        logger.info(f"Generating {request.n} image(s) {size_str}")
+
+        # Generate images using AsyncOmni (multi-stage mode)
+        result = await _generate_with_async_omni(
+            engine_client=engine_client,
+            gen_params=gen_params,
+            stage_types=stage_types,
+            prompt=prompt,
+            request_id=request_id,
+        )
+
+        if result is None:
+            raise HTTPException(
+                status_code=HTTPStatus.INTERNAL_SERVER_ERROR.value,
+                detail="No output generated from multi-stage pipeline.",
+            )
+
+        # Extract images from result
+        images = _extract_images_from_result(result)
+
+        logger.info(f"Successfully generated {len(images)} image(s)")
+
+        # Encode images to base64
+        image_data = [ImageData(b64_json=encode_image_base64(img), revised_prompt=None) for img in images]
+
+        return ImageGenerationResponse(
+            created=int(time.time()),
+            data=image_data,
+        )
+
+    except HTTPException:
+        # Re-raise HTTPExceptions as-is
+        raise
+    except ValueError as e:
+        logger.error(f"Validation error: {e}")
+        raise HTTPException(status_code=HTTPStatus.BAD_REQUEST.value, detail=str(e))
+    except Exception as e:
+        logger.exception(f"Image generation failed: {e}")
+        raise HTTPException(
+            status_code=HTTPStatus.INTERNAL_SERVER_ERROR.value, detail=f"Image generation failed: {str(e)}"
+        )
+
+
+@router.post(
+    "/v1/images/edits",
+    responses={
+        HTTPStatus.OK.value: {"model": ImageGenerationResponse},
+        HTTPStatus.BAD_REQUEST.value: {"model": ErrorResponse},
+        HTTPStatus.SERVICE_UNAVAILABLE.value: {"model": ErrorResponse},
+        HTTPStatus.INTERNAL_SERVER_ERROR.value: {"model": ErrorResponse},
+    },
+)
+async def edit_images(
+    raw_request: Request,
+    image: list[UploadFile] | None = File(None),
+    image_array: list[UploadFile] | None = File(None, alias="image[]"),
+    url: list[str] | None = Form(None),
+    url_array: list[str] | None = Form(None, alias="url[]"),
+    prompt: str = Form(...),
+    model: str = Form(None),
+    n: int = Form(1),
+    size: str = Form("auto"),
+    response_format: str = Form("b64_json"),
+    output_format: str | None = Form("png"),
+    background: str | None = Form("auto"),
+    output_compression: Annotated[int, Form(ge=0, le=100)] = 100,
+    user: str | None = Form(None),  # unused now
+    # vllm-omni extensions for diffusion control
+    negative_prompt: str | None = Form(None),
+    num_inference_steps: int | None = Form(None),
+    guidance_scale: float | None = Form(None),
+    true_cfg_scale: float | None = Form(None),
+    seed: int | None = Form(None),
+    # vllm-omni extension for per-request LoRA.
+    lora: str | None = Form(None),  # Json string
+) -> ImageGenerationResponse:
+    """
+    OpenAI-compatible image edit endpoint.
+    """
+    # 1. get engine and model
+    engine_client, model_name, stage_types = _get_engine_and_model(raw_request)
+    if model is not None and model != model_name:
+        logger.warning(
+            f"Model mismatch: request specifies '{model}' but server is running '{model_name}'. Using server model."
+        )
+    # 2. get output format & compression
+    output_format = _choose_output_format(output_format, background)
+    if response_format != "b64_json":
+        raise HTTPException(
+            status_code=HTTPStatus.BAD_REQUEST.value,
+            detail="Only response_format 'b64_json' is supported now.",
+        )
+    try:
+        # 2. Build prompt & images params
+        prompt: OmniTextPrompt = {"prompt": prompt}
+        if negative_prompt is not None:
+            prompt["negative_prompt"] = negative_prompt
+        input_images_list = []
+        images = image or image_array
+        urls = url or url_array
+        if images:
+            input_images_list.extend(images)
+        if urls:
+            input_images_list.extend(urls)
+        if not input_images_list:
+            raise HTTPException(status_code=422, detail="Field 'image' or 'url' is required")
+        pil_images = await _load_input_images(input_images_list)
+        prompt["multi_modal_data"] = {}
+        prompt["multi_modal_data"]["image"] = pil_images
+
+        # 3 Build sample params
+        gen_params = OmniDiffusionSamplingParams()
+        # 3.0 Init with system default values
+        app_state_args = getattr(raw_request.app.state, "args", None)
+        default_sample_param = getattr(app_state_args, "default_sampling_params", None)
+        # Currently only have one diffusion stage
+        diffusion_stage_id = [i for i, t in enumerate(stage_types) if t == "diffusion"][0]
+        apply_stage_default_sampling_params(
+            default_sample_param,
+            gen_params,
+            str(diffusion_stage_id),
+        )
+        _update_if_not_none(gen_params, "num_outputs_per_prompt", n)
+        # 3.1 Parse per-request LoRA (compatible with chat's extra_body.lora shape).
+        lora_dict = _get_lora_from_json_str(lora)
+        lora_request, lora_scale = _parse_lora_request(lora_dict)
+        _update_if_not_none(gen_params, "lora_request", lora_request)
+        _update_if_not_none(gen_params, "lora_scale", lora_scale)
+        # 3.2 Parse and add size if provided
+        max_generated_image_size = getattr(app_state_args, "max_generated_image_size", None)
+        width, height = None, None
+        if size.lower() == "auto":
+            width, height = pil_images[0].size  # Use first image size
+        else:
+            width, height = parse_size(size)
+        if max_generated_image_size is not None and (width * height > max_generated_image_size):
+            raise HTTPException(
+                status_code=HTTPStatus.BAD_REQUEST.value,
+                detail=f"Requested image size {width}x{height} exceeds the maximum allowed "
+                f"size of {max_generated_image_size} pixels.",
+            )
+
+        size_str = f"{width}x{height}"
+        _update_if_not_none(gen_params, "width", width)
+        _update_if_not_none(gen_params, "height", height)
+
+        # 3.3 Add optional parameters ONLY if provided
+        _update_if_not_none(gen_params, "num_inference_steps", num_inference_steps)
+        _update_if_not_none(gen_params, "guidance_scale", guidance_scale)
+        _update_if_not_none(gen_params, "true_cfg_scale", true_cfg_scale)
+        # If seed is not provided, generate a random one to ensure
+        # a proper generator is initialized in the backend.
+        # This fixes issues where using the default global generator
+        # might produce blurry images in some environments.
+        _update_if_not_none(gen_params, "seed", seed or random.randint(0, 2**32 - 1))
+
+        # 4. Generate images using AsyncOmni (multi-stage mode)
+        request_id = f"img_edit_{int(time.time())}"
+        logger.info(f"Generating {n} image(s) {size_str}")
+        result = await _generate_with_async_omni(
+            engine_client=engine_client,
+            gen_params=gen_params,
+            stage_types=stage_types,
+            prompt=prompt,
+            request_id=request_id,
+        )
+
+        # 5. Extract images from result
+        images = _extract_images_from_result(result)
+        logger.info(f"Successfully generated {len(images)} image(s)")
+
+        # Encode images to base64
+        image_data = [
+            ImageData(
+                b64_json=_encode_image_base64_with_compression(
+                    img, format=output_format, output_compression=output_compression
+                ),
+                revised_prompt=None,
+            )
+            for img in images
+        ]
+
+        return ImageGenerationResponse(
+            created=int(time.time()),
+            data=image_data,
+            output_format=output_format,
+            size=size_str,
+        )
+
+    except HTTPException:
+        # Re-raise HTTPExceptions as-is
+        raise
+    except ValueError as e:
+        logger.error(f"Validation error: {e}")
+        raise HTTPException(status_code=HTTPStatus.BAD_REQUEST.value, detail=str(e))
+    except Exception as e:
+        logger.exception(f"Image edit failed: {e}")
+        raise HTTPException(status_code=HTTPStatus.INTERNAL_SERVER_ERROR.value, detail=f"Image edit failed: {str(e)}")
+
+
+def _get_engine_and_model(raw_request: Request):
     # Get engine client (AsyncOmni) from app state
     engine_client: EngineClient | AsyncOmni | None = getattr(raw_request.app.state, "engine_client", None)
     if engine_client is None or not hasattr(engine_client, "stage_list"):
@@ -942,157 +1192,229 @@ async def generate_images(request: ImageGenerationRequest, raw_request: Request)
     else:
         model_name = "unknown"
 
-    # Validate model field (warn if mismatch, don't error)
-    if request.model is not None and request.model != model_name:
-        logger.warning(
-            f"Model mismatch: request specifies '{request.model}' but "
-            f"server is running '{model_name}'. Using server model."
-        )
+    return engine_client, model_name, stage_types
 
+
+def _get_lora_from_json_str(lora_body):
+    if lora_body is None:
+        return None
     try:
-        # Build params - pass through user values directly
-        prompt: OmniTextPrompt = {"prompt": request.prompt}
-        gen_params = OmniDiffusionSamplingParams(num_outputs_per_prompt=request.n)
+        lora_dict = json.loads(lora_body)
+    except json.JSONDecodeError:
+        raise HTTPException(status_code=400, detail="Invalid LoRA JSON string")
 
-        # Parse per-request LoRA (compatible with chat's extra_body.lora shape).
-        if request.lora is not None:
-            if not isinstance(request.lora, dict):
-                raise HTTPException(
-                    status_code=HTTPStatus.BAD_REQUEST.value,
-                    detail="Invalid lora field: expected an object.",
-                )
-            lora_body = request.lora
-            lora_name = lora_body.get("name") or lora_body.get("lora_name") or lora_body.get("adapter")
-            lora_path = (
-                lora_body.get("local_path")
-                or lora_body.get("path")
-                or lora_body.get("lora_path")
-                or lora_body.get("lora_local_path")
-            )
-            lora_scale = lora_body.get("scale")
-            if lora_scale is None:
-                lora_scale = lora_body.get("lora_scale")
-            lora_int_id = lora_body.get("int_id")
-            if lora_int_id is None:
-                lora_int_id = lora_body.get("lora_int_id")
-            if lora_int_id is None and lora_path:
-                lora_int_id = stable_lora_int_id(str(lora_path))
+    if not isinstance(lora_dict, dict):
+        raise HTTPException(status_code=400, detail="LoRA must be a JSON object")
 
-            if not lora_name or not lora_path:
-                raise HTTPException(
-                    status_code=HTTPStatus.BAD_REQUEST.value,
-                    detail="Invalid lora object: both name and path are required.",
-                )
+    return lora_dict
 
-            gen_params.lora_request = LoRARequest(str(lora_name), int(lora_int_id), str(lora_path))
-            if lora_scale is not None:
-                gen_params.lora_scale = float(lora_scale)
 
-        # Parse and add size if provided
-        if request.size:
-            width, height = parse_size(request.size)
-            gen_params.height = height
-            gen_params.width = width
-            size_str = f"{width}x{height}"
-        else:
-            size_str = "model default"
-
-        # Add optional parameters ONLY if provided
-        if request.num_inference_steps is not None:
-            gen_params.num_inference_steps = request.num_inference_steps
-        if request.negative_prompt is not None:
-            prompt["negative_prompt"] = request.negative_prompt
-        if request.guidance_scale is not None:
-            gen_params.guidance_scale = request.guidance_scale
-        if request.true_cfg_scale is not None:
-            gen_params.true_cfg_scale = request.true_cfg_scale
-        if request.seed is not None:
-            gen_params.seed = request.seed
-        else:
-            # If seed is not provided, generate a random one to ensure
-            # a proper generator is initialized in the backend.
-            # This fixes issues where using the default global generator
-            # might produce blurry images in some environments.
-            gen_params.seed = random.randint(0, 2**32 - 1)
-
-        request_id = f"img_gen_{uuid.uuid4().hex}"
-
-        logger.info(f"Generating {request.n} image(s) {size_str}")
-
-        # Generate images using AsyncOmni (multi-stage mode)
-        engine_client = cast(AsyncOmni, engine_client)
-        result = None
-        stage_list = getattr(engine_client, "stage_list", None)
-        if isinstance(stage_list, list):
-            default_params_list: list[OmniSamplingParams] | None = getattr(
-                engine_client, "default_sampling_params_list", None
-            )
-            if not isinstance(default_params_list, list):
-                default_params_list = [
-                    OmniDiffusionSamplingParams() if st == "diffusion" else SamplingParams() for st in stage_types
-                ]
-            else:
-                default_params_list = list(default_params_list)
-            if len(default_params_list) != len(stage_types):
-                default_params_list = (
-                    default_params_list
-                    + [OmniDiffusionSamplingParams() if st == "diffusion" else SamplingParams() for st in stage_types]
-                )[: len(stage_types)]
-
-            sampling_params_list: list[OmniSamplingParams] = []
-            for idx, stage_type in enumerate(stage_types):
-                if stage_type == "diffusion":
-                    sampling_params_list.append(gen_params)
-                else:
-                    base_params = default_params_list[idx]
-                    sampling_params_list.append(base_params)
-
-            async for output in engine_client.generate(
-                prompt=prompt,
-                request_id=request_id,
-                sampling_params_list=sampling_params_list,
-            ):
-                result = output
-        else:
-            result = await engine_client.generate(
-                prompt=prompt, request_id=request_id, sampling_params_list=[gen_params]
-            )
-
-        if result is None:
+def _parse_lora_request(lora_body: dict[str, Any]):
+    if lora_body is not None:
+        if not isinstance(lora_body, dict):
             raise HTTPException(
-                status_code=HTTPStatus.INTERNAL_SERVER_ERROR.value,
-                detail="No output generated from multi-stage pipeline.",
+                status_code=HTTPStatus.BAD_REQUEST.value,
+                detail="Invalid lora field: expected an object.",
+            )
+        lora_name = lora_body.get("name") or lora_body.get("lora_name") or lora_body.get("adapter")
+        lora_path = (
+            lora_body.get("local_path")
+            or lora_body.get("path")
+            or lora_body.get("lora_path")
+            or lora_body.get("lora_local_path")
+        )
+        lora_scale = lora_body.get("scale")
+        if lora_scale is None:
+            lora_scale = lora_body.get("lora_scale")
+        lora_int_id = lora_body.get("int_id")
+        if lora_int_id is None:
+            lora_int_id = lora_body.get("lora_int_id")
+        if lora_int_id is None and lora_path:
+            lora_int_id = stable_lora_int_id(str(lora_path))
+
+        if not lora_name or not lora_path:
+            raise HTTPException(
+                status_code=HTTPStatus.BAD_REQUEST.value,
+                detail="Invalid lora object: both name and path are required.",
             )
 
-        # Extract images from result
-        images = []
-        if hasattr(result, "images") and result.images:
-            images = result.images
-        elif hasattr(result, "request_output"):
-            request_output = result.request_output
-            if isinstance(request_output, dict) and request_output.get("images"):
-                images = request_output["images"]
-            elif hasattr(request_output, "images") and request_output.images:
-                images = request_output.images
+        return LoRARequest(str(lora_name), int(lora_int_id), str(lora_path)), lora_scale
+    return None, None
 
-        logger.info(f"Successfully generated {len(images)} image(s)")
 
-        # Encode images to base64
-        image_data = [ImageData(b64_json=encode_image_base64(img), revised_prompt=None) for img in images]
+async def _generate_with_async_omni(
+    engine_client: AsyncOmni | Any,
+    gen_params: Any,
+    stage_types: list[str],
+    **kwargs,
+):
+    engine_client = cast(AsyncOmni, engine_client)
+    result = None
+    stage_list = getattr(engine_client, "stage_list", None)
+    if isinstance(stage_list, list):
+        default_params_list: list[OmniSamplingParams] | None = getattr(
+            engine_client, "default_sampling_params_list", None
+        )
+        if not isinstance(default_params_list, list):
+            default_params_list = [
+                OmniDiffusionSamplingParams() if st == "diffusion" else SamplingParams() for st in stage_types
+            ]
+        else:
+            default_params_list = list(default_params_list)
+        if len(default_params_list) != len(stage_types):
+            default_params_list = (
+                default_params_list
+                + [OmniDiffusionSamplingParams() if st == "diffusion" else SamplingParams() for st in stage_types]
+            )[: len(stage_types)]
 
-        return ImageGenerationResponse(
-            created=int(time.time()),
-            data=image_data,
+        sampling_params_list: list[OmniSamplingParams] = []
+        for idx, stage_type in enumerate(stage_types):
+            if stage_type == "diffusion":
+                sampling_params_list.append(gen_params)
+            else:
+                base_params = default_params_list[idx]
+                sampling_params_list.append(base_params)
+
+        async for output in engine_client.generate(
+            sampling_params_list=sampling_params_list,
+            **kwargs,
+        ):
+            result = output
+    else:
+        result = await engine_client.generate(
+            sampling_params_list=[gen_params],
+            **kwargs,
         )
 
-    except HTTPException:
-        # Re-raise HTTPExceptions as-is
-        raise
-    except ValueError as e:
-        logger.error(f"Validation error: {e}")
-        raise HTTPException(status_code=HTTPStatus.BAD_REQUEST.value, detail=str(e))
-    except Exception as e:
-        logger.exception(f"Image generation failed: {e}")
+    if result is None:
         raise HTTPException(
-            status_code=HTTPStatus.INTERNAL_SERVER_ERROR.value, detail=f"Image generation failed: {str(e)}"
+            status_code=HTTPStatus.INTERNAL_SERVER_ERROR.value,
+            detail="No output generated from multi-stage pipeline.",
         )
+    return result
+
+
+def _update_if_not_none(object: any, key: str, val: any) -> None:
+    if val is not None:
+        setattr(object, key, val)
+
+
+def _extract_images_from_result(result: Any) -> list[Any]:
+    images = []
+    if hasattr(result, "images") and result.images:
+        images = result.images
+    elif hasattr(result, "request_output"):
+        request_output = result.request_output
+        if isinstance(request_output, dict) and request_output.get("images"):
+            images = request_output["images"]
+        elif hasattr(request_output, "images") and request_output.images:
+            images = request_output.images
+    return images
+
+
+async def _load_input_images(
+    inputs: list[str],
+) -> list[Image.Image]:
+    """
+    convert to PIL.Image.Image list
+    """
+    if isinstance(inputs, str):
+        inputs = [inputs]
+
+    images: list[Image.Image] = []
+
+    for inp in inputs:
+        # 1. URL + base64
+        if isinstance(inp, str) and inp.startswith("data:image"):
+            try:
+                _, b64_data = inp.split(",", 1)
+                image_bytes = base64.b64decode(b64_data)
+                img = Image.open(io.BytesIO(image_bytes))
+                images.append(img)
+            except Exception as e:
+                raise ValueError(f"Invalid base64 image: {e}")
+
+        # 2. URL
+        elif isinstance(inp, str) and inp.startswith("http"):
+            async with httpx.AsyncClient(timeout=60) as client:
+                try:
+                    resp = await client.get(inp)
+                    resp.raise_for_status()
+                    img = Image.open(io.BytesIO(resp.content))
+                    images.append(img)
+                except Exception as e:
+                    raise ValueError(f"Failed to download image from URL {inp}: {e}")
+
+        # 3. UploadFile
+        elif hasattr(inp, "file"):
+            try:
+                img_data = await inp.read()
+                img = Image.open(io.BytesIO(img_data))
+                images.append(img)
+            except Exception as e:
+                raise ValueError(f"Failed to open uploaded file: {e}")
+        else:
+            raise ValueError(f"Unsupported input: {inp}")
+
+    if not images:
+        raise ValueError("No valid input images found")
+
+    return images
+
+
+def _choose_output_format(output_format: str | None, background: str | None) -> str:
+    # Normalize and choose extension
+    fmt = (output_format or "").lower()
+    if fmt in {"jpg", "png", "webp", "jpeg"}:
+        return fmt
+    # If transparency requested, prefer png
+    if (background or "auto").lower() == "transparent":
+        return "png"
+    # Default
+    return "jpeg"
+
+
+def _encode_image_base64_with_compression(
+    image: Image.Image, format: str = "png", output_compression: int = 100
+) -> str:
+    """Encode PIL Image to base64 PNG string.
+
+    Args:
+        image: PIL Image object
+        format: Output image format (e.g., "PNG", "JPEG", "WEBP")
+        output_compression: Compression level (0-100%), 100 for best quality
+    Returns:
+        Base64-encoded image as string
+    """
+    buffer = io.BytesIO()
+    save_kwargs = {}
+    if format in ("jpg", "jpeg", "webp"):
+        save_kwargs["quality"] = output_compression
+    elif format == "png":
+        save_kwargs["compress_level"] = max(0, min(9, 9 - output_compression // 11))  # Map 0-100 to 9-0
+
+    image.save(buffer, format=format, **save_kwargs)
+    buffer.seek(0)
+    return base64.b64encode(buffer.read()).decode("utf-8")
+
+
+def apply_stage_default_sampling_params(
+    default_params_json: str | None,
+    sampling_params: any,
+    stage_key: str,
+) -> None:
+    """
+    Update a stage's sampling parameters with vLLM-Omni defaults.
+
+    Args:
+        default_params_json: JSON string of stage-keyed default parameters
+        sampling_params: The sampling parameters object to update
+        stage_key: The stage ID/key in the pipeline
+    """
+    if default_params_json is not None:
+        default_params_dict = json.loads(default_params_json)
+        if stage_key in default_params_dict:
+            stage_defaults = default_params_dict[stage_key]
+            for param_name, param_value in stage_defaults.items():
+                if hasattr(sampling_params, param_name):
+                    setattr(sampling_params, param_name, param_value)
diff --git a/vllm_omni/entrypoints/openai/protocol/images.py b/vllm_omni/entrypoints/openai/protocol/images.py
index 093ecaddf..63514b191 100644
--- a/vllm_omni/entrypoints/openai/protocol/images.py
+++ b/vllm_omni/entrypoints/openai/protocol/images.py
@@ -123,3 +123,5 @@ class ImageGenerationResponse(BaseModel):
 
     created: int = Field(..., description="Unix timestamp of when the generation completed")
     data: list[ImageData] = Field(..., description="Array of generated images")
+    output_format: str = Field(None, description="The output format of the image generation")
+    size: str = Field(None, description="The size of the image generated")
diff --git a/vllm_omni/entrypoints/openai/serving_chat.py b/vllm_omni/entrypoints/openai/serving_chat.py
index d8e387b5b..ae2fb6552 100644
--- a/vllm_omni/entrypoints/openai/serving_chat.py
+++ b/vllm_omni/entrypoints/openai/serving_chat.py
@@ -287,8 +287,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
 
                 generators.append(generator)
         except ValueError as e:
-            # TODO: Use a vllm-specific Validation Error
-            return self.create_error_response(str(e))
+            return self.create_error_response(e)
 
         assert len(generators) == 1
         (result_generator,) = generators
@@ -316,8 +315,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
                 request_metadata,
             )
         except ValueError as e:
-            # TODO: Use a vllm-specific Validation Error
-            return self.create_error_response(str(e))
+            return self.create_error_response(e)
 
     async def _preprocess_chat(
         self,
@@ -634,7 +632,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
                 )
         except RuntimeError as e:
             logger.exception("Error in reasoning parser creation.")
-            data = self.create_streaming_error_response(str(e))
+            data = self.create_streaming_error_response(e)
             yield f"data: {data}\n\n"
             yield "data: [DONE]\n\n"
             return
@@ -646,7 +644,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
                 tool_parsers = [None] * num_choices
         except Exception as e:
             logger.exception("Error in tool parser creation.")
-            data = self.create_streaming_error_response(str(e))
+            data = self.create_streaming_error_response(e)
             yield f"data: {data}\n\n"
             yield "data: [DONE]\n\n"
             return
@@ -1286,9 +1284,8 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
                     )
 
         except Exception as e:
-            # TODO: Use a vllm-specific Validation Error
             logger.exception("Error in chat completion stream generator.")
-            data = self.create_streaming_error_response(str(e))
+            data = self.create_streaming_error_response(e)
             yield f"data: {data}\n\n"
         # Send the final done message after all response.n are finished
         yield "data: [DONE]\n\n"
@@ -1313,8 +1310,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
         except asyncio.CancelledError:
             return self.create_error_response("Client disconnected")
         except ValueError as e:
-            # TODO: Use a vllm-specific Validation Error
-            return self.create_error_response(str(e))
+            return self.create_error_response(e)
 
         assert final_outputs is not None
 
@@ -1487,7 +1483,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
                     )
                 except RuntimeError as e:
                     logger.exception("Error in reasoning parser creation.")
-                    return self.create_error_response(str(e))
+                    return self.create_error_response(e)
                 # If the reasoning parser is enabled,
                 # tool calls are extracted exclusively from the content.
                 reasoning_content, content = reasoning_parser.extract_reasoning(output.text, request=request)
@@ -1572,7 +1568,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
                     tool_parser = self.tool_parser(tokenizer)
                 except RuntimeError as e:
                     logger.exception("Error in tool parser creation.")
-                    return self.create_error_response(str(e))
+                    return self.create_error_response(e)
 
                 tool_call_info = tool_parser.extract_tool_calls(content if content is not None else "", request=request)
                 # In the OpenAI API the finish_reason is "tools_called"
diff --git a/vllm_omni/entrypoints/openai/serving_speech.py b/vllm_omni/entrypoints/openai/serving_speech.py
index cb9f6f8a8..5e1767700 100644
--- a/vllm_omni/entrypoints/openai/serving_speech.py
+++ b/vllm_omni/entrypoints/openai/serving_speech.py
@@ -296,7 +296,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
         except asyncio.CancelledError:
             return self.create_error_response("Client disconnected")
         except ValueError as e:
-            return self.create_error_response(str(e))
+            return self.create_error_response(e)
         except Exception as e:
             logger.exception("Speech generation failed: %s", e)
             return self.create_error_response(f"Speech generation failed: {e}")
diff --git a/vllm_omni/model_executor/stage_configs/bagel.yaml b/vllm_omni/model_executor/stage_configs/bagel.yaml
index ee2b17729..c60f3b700 100644
--- a/vllm_omni/model_executor/stage_configs/bagel.yaml
+++ b/vllm_omni/model_executor/stage_configs/bagel.yaml
@@ -11,7 +11,7 @@ stage_args:
       model_arch: BagelForConditionalGeneration
       worker_type: ar
       scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
-      gpu_memory_utilization: 0.4
+      gpu_memory_utilization: 0.35
       enforce_eager: true
       trust_remote_code: true
       engine_output_type: text
@@ -42,7 +42,7 @@ stage_args:
       max_batch_size: 1
     engine_args:
       model_stage: dit
-      gpu_memory_utilization: 0.4
+      gpu_memory_utilization: 0.55
       enforce_eager: true
       trust_remote_code: true
       engine_output_type: image
diff --git a/vllm_omni/model_executor/stage_configs/bagel_multiconnector.yaml b/vllm_omni/model_executor/stage_configs/bagel_multiconnector.yaml
index eddc844d5..c19302c76 100644
--- a/vllm_omni/model_executor/stage_configs/bagel_multiconnector.yaml
+++ b/vllm_omni/model_executor/stage_configs/bagel_multiconnector.yaml
@@ -11,7 +11,7 @@ stage_args:
       model_arch: BagelForConditionalGeneration
       worker_type: ar
       scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
-      gpu_memory_utilization: 0.4
+      gpu_memory_utilization: 0.35
       enforce_eager: true
       trust_remote_code: true
       engine_output_type: text
@@ -41,11 +41,11 @@ stage_args:
   - stage_id: 1
     stage_type: diffusion
     runtime:
-      devices: "1"
+      devices: "0"
       max_batch_size: 1
     engine_args:
       model_stage: dit
-      gpu_memory_utilization: 0.4
+      gpu_memory_utilization: 0.55
       enforce_eager: true
       trust_remote_code: true
       engine_output_type: image
diff --git a/vllm_omni/platforms/npu/stage_configs/qwen2_5_omni.yaml b/vllm_omni/platforms/npu/stage_configs/qwen2_5_omni.yaml
index a02d84ed6..d021ef218 100644
--- a/vllm_omni/platforms/npu/stage_configs/qwen2_5_omni.yaml
+++ b/vllm_omni/platforms/npu/stage_configs/qwen2_5_omni.yaml
@@ -39,7 +39,7 @@ stage_args:
       worker_type: ar
       scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
       gpu_memory_utilization: 0.8
-      enforce_eager: false
+      enforce_eager: true # haven't supported talker ACL graph on NPU
       trust_remote_code: true
       enable_prefix_caching: false
       engine_output_type: latent
diff --git a/vllm_omni/platforms/npu/stage_configs/qwen3_omni_moe.yaml b/vllm_omni/platforms/npu/stage_configs/qwen3_omni_moe.yaml
index fdf817239..f99ed22e8 100644
--- a/vllm_omni/platforms/npu/stage_configs/qwen3_omni_moe.yaml
+++ b/vllm_omni/platforms/npu/stage_configs/qwen3_omni_moe.yaml
@@ -44,7 +44,7 @@ stage_args:
        worker_type: ar
        scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
        gpu_memory_utilization: 0.2
-       enforce_eager: false
+       enforce_eager: true # haven't supported talker ACL graph on NPU
        trust_remote_code: true
        engine_output_type: latent  # Output codec codes for code2wav
        # tensor_parallel_size: 2
diff --git a/vllm_omni/platforms/npu/worker/npu_ar_model_runner.py b/vllm_omni/platforms/npu/worker/npu_ar_model_runner.py
index 0ba57e038..fbf9e632f 100644
--- a/vllm_omni/platforms/npu/worker/npu_ar_model_runner.py
+++ b/vllm_omni/platforms/npu/worker/npu_ar_model_runner.py
@@ -35,7 +35,7 @@ from vllm_ascend.compilation.acl_graph import update_full_graph_params
 from vllm_ascend.ops.rotary_embedding import update_cos_sin
 from vllm_ascend.utils import ProfileExecuteDuration
 
-from vllm_omni.core.sched.omni_ar_scheduler import KVCacheTransferData
+from vllm_omni.distributed.omni_connectors.kv_transfer_manager import OmniKVTransferManager
 from vllm_omni.outputs import OmniModelRunnerOutput
 from vllm_omni.platforms.npu.worker.npu_model_runner import OmniNPUModelRunner
 
@@ -65,7 +65,8 @@ class NPUARModelRunner(OmniNPUModelRunner):
         # each model stage has their own hidden size
         self.hidden_size = self.model_config.hf_text_config.hidden_size
         self.inputs_embeds = self._make_buffer(self.max_num_tokens, self.hidden_size, dtype=self.dtype, numpy=False)
-        self.omni_connector = None
+        # Initialize KV cache manager (preserve vllm_config fallback behavior)
+        self.kv_transfer_manager = OmniKVTransferManager.from_vllm_config(self.vllm_config, self.model_config)
 
     def _make_buffer(self, *size, dtype, numpy=True):
         # Prevent ray from pinning the buffer due to large size
@@ -92,7 +93,13 @@ class NPUARModelRunner(OmniNPUModelRunner):
 
         #  -------------------------------------- Omni-new -------------------------------------------------
         # [Omni] Handle KV transfer BEFORE updating states (which removes finished requests)
-        self.kv_extracted_req_ids = self._handle_finished_requests_kv_transfer(scheduler_output)
+        self.kv_extracted_req_ids = self.kv_transfer_manager.handle_finished_requests_kv_transfer(
+            finished_reqs=getattr(scheduler_output, "finished_requests_needing_kv_transfer", {}),
+            kv_caches=self.kv_caches,
+            block_size=self.cache_config.block_size,
+            cache_dtype=str(self.cache_config.cache_dtype),
+            request_id_resolver=self._resolve_global_request_id,
+        )
         #  -------------------------------------- Omni-new -------------------------------------------------
 
         with ProfileExecuteDuration().capture_async("prepare input"):
@@ -499,161 +506,6 @@ class NPUARModelRunner(OmniNPUModelRunner):
         return hidden_states if self.pcp_size == 1 else self.pcp_manager.get_restore_hidden_states(
             hidden_states)
 
-    def _handle_finished_requests_kv_transfer(self, scheduler_output: SchedulerOutput) -> list[str]:
-        """Handle KV cache transfer for finished requests.
-
-        Returns list of request IDs that were processed (for Scheduler to free blocks).
-        """
-        finished_reqs = getattr(scheduler_output, "finished_requests_needing_kv_transfer", {})
-        if not finished_reqs:
-            return []
-
-        logger.debug(f"Processing KV transfer for {len(finished_reqs)} requests")
-
-        extracted_ids = []
-        for req_id, data in finished_reqs.items():
-            try:
-                seq_len = data.get("seq_len", 0)
-                block_ids = data.get("block_ids", [])
-                if not block_ids:
-                    logger.warning(f"Request {req_id} has no block IDs, skipping")
-                    continue
-
-                # Extract KV cache from GPU blocks -> CPU tensors
-                kv_data = self._extract_kv_cache(req_id, block_ids, seq_len)
-                if kv_data:
-                    # Transfer to downstream stage via connector
-                    self._transfer_kv_cache(kv_data)
-
-            except Exception as e:
-                logger.error(f"Failed KV transfer for {req_id}: {e}")
-            finally:
-                extracted_ids.append(req_id)
-
-        return extracted_ids
-
-    def _extract_kv_cache(self, req_id: str, block_ids: list[int], seq_len: int) -> KVCacheTransferData | None:
-        """Extract KV cache from GPU blocks for a single request."""
-        num_layers = len(self.kv_caches)
-        key_cache = [None] * num_layers
-        value_cache = [None] * num_layers
-
-        for layer_idx, kv_tensor in enumerate(self.kv_caches):
-            # Validate block IDs
-            max_block = kv_tensor.shape[1] - 1
-            valid_ids = [bid for bid in block_ids if 0 <= bid <= max_block]
-            if not valid_ids:
-                continue
-
-            # Extract and reshape: [2, n_blocks, block_size, n_heads, head_dim]
-            # -> [2, seq_len, n_heads, head_dim]
-            selected = kv_tensor[:, valid_ids]  # [2, n_valid, block_size, n_heads, head_dim]
-            n_kv, n_blks, blk_sz, n_heads, d_head = selected.shape
-            flat = selected.reshape(n_kv, n_blks * blk_sz, n_heads, d_head)
-            if seq_len < flat.shape[1]:
-                flat = flat[:, :seq_len]
-
-            # Move to CPU
-            flat_cpu = flat.detach().cpu().contiguous()
-            key_cache[layer_idx] = flat_cpu[0]
-            value_cache[layer_idx] = flat_cpu[1]
-
-        if not any(k is not None for k in key_cache):
-            return None
-
-        return KVCacheTransferData(
-            request_id=req_id,
-            layer_blocks={"key_cache": key_cache, "value_cache": value_cache},
-            block_ids=block_ids,
-            metadata={
-                "block_size": self.cache_config.block_size,
-                "num_layers": num_layers,
-                "dtype": str(self.cache_config.cache_dtype),
-                "seq_len": seq_len,
-            },
-        )
-
-    def _transfer_kv_cache(self, kv_data: KVCacheTransferData) -> None:
-        """Transfer KV cache data to downstream stage via OmniConnector."""
-        connector = self._get_or_create_connector()
-        if not connector:
-            return
-
-        # Resolve global request ID if available
-        transfer_req_id = self._resolve_global_request_id(kv_data.request_id)
-        from_stage, to_stage = self._detect_transfer_stages()
-
-        # Prepare data and transfer with retry
-        data_dict = kv_data.to_dict()
-        data_dict["request_id"] = transfer_req_id
-
-        success, size, _ = self._transfer_with_retry(
-            connector, from_stage, to_stage, f"kv_cache_{transfer_req_id}", data_dict
-        )
-
-        if success:
-            logger.info(f"KV transfer OK: {transfer_req_id}, {size} bytes")
-        else:
-            logger.error(f"KV transfer FAILED: {transfer_req_id}")
-
-    def _get_or_create_connector(self) -> Any | None:
-        """Get existing connector or create one from config."""
-        if self.omni_connector:
-            return self.omni_connector
-
-        from vllm_omni.distributed.omni_connectors.factory import OmniConnectorFactory
-        from vllm_omni.distributed.omni_connectors.utils.config import ConnectorSpec
-
-        config = self._get_omni_connector_config()
-        if not config or not isinstance(config, dict):
-            logger.warning("No valid OmniConnector config found")
-            return None
-
-        c_type = config.get("type")
-        if not c_type:
-            logger.error("OmniConnector config missing 'type' field")
-            return None
-
-        c_extra = {k: v for k, v in config.items() if k != "type"}
-        self.omni_connector = OmniConnectorFactory.create_connector(ConnectorSpec(name=c_type, extra=c_extra))
-        return self.omni_connector
-
-    def _get_omni_connector_config(self) -> dict[str, Any] | None:
-        """Get OmniConnector configuration from model config."""
-        # Primary: omni_kv_config from YAML
-        omni_kv = getattr(self.model_config, "omni_kv_config", None)
-        if isinstance(omni_kv, dict):
-            cfg = omni_kv.get("connector_config")
-            if isinstance(cfg, dict) and cfg:
-                return cfg
-
-        # Fallback: kv_transfer_config
-        kv_cfg = getattr(self.vllm_config, "kv_transfer_config", None)
-        if kv_cfg:
-            direct = getattr(kv_cfg, "omni_connector_config", None)
-            if isinstance(direct, dict) and direct:
-                return direct
-            extra = getattr(kv_cfg, "kv_connector_extra_config", None)
-            if isinstance(extra, dict):
-                omni = extra.get("omni_connector_config")
-                if isinstance(omni, dict) and omni:
-                    return omni
-
-        return None
-
-    def _detect_transfer_stages(self) -> tuple[str, str]:
-        """Detect source and target stages for KV transfer."""
-        omni_kv = getattr(self.model_config, "omni_kv_config", None)
-        if isinstance(omni_kv, dict):
-            from_s = omni_kv.get("omni_from_stage")
-            to_s = omni_kv.get("omni_to_stage")
-            if from_s and to_s:
-                return str(from_s), str(to_s)
-
-        raise ValueError(
-            "KV transfer stages not configured. Please set 'omni_from_stage' and 'omni_to_stage' in omni_kv_config."
-        )
-
     def _resolve_global_request_id(self, req_id: str) -> str:
         """Resolve global request ID from request state."""
         req_state = self.requests.get(req_id)
@@ -669,32 +521,3 @@ class NPUARModelRunner(OmniNPUModelRunner):
                 return global_id.decode("utf-8")
             return str(global_id)
         return req_id
-
-    def _transfer_with_retry(
-        self,
-        connector: Any,
-        from_stage: str,
-        to_stage: str,
-        request_id: str,
-        data: dict[str, Any],
-        max_retries: int = 3,
-    ) -> tuple[bool, int, dict[str, Any] | None]:
-        """Transfer data with retry and exponential backoff."""
-        import time
-
-        for attempt in range(max_retries):
-            try:
-                put_key = f"omni_{from_stage}_to_{to_stage}_{request_id}"
-                success, size, metadata = connector.put(
-                    from_stage=from_stage, to_stage=to_stage, put_key=put_key, data=data
-                )
-                if success:
-                    return success, size, metadata
-                logger.warning(f"Transfer attempt {attempt + 1} failed for {request_id}")
-            except Exception as e:
-                logger.warning(f"Transfer attempt {attempt + 1} exception: {e}")
-
-            if attempt < max_retries - 1:
-                time.sleep(0.1 * (2**attempt))
-
-        return False, 0, None
diff --git a/vllm_omni/platforms/npu/worker/npu_model_runner.py b/vllm_omni/platforms/npu/worker/npu_model_runner.py
index b0fbf7a4e..a9a376f95 100644
--- a/vllm_omni/platforms/npu/worker/npu_model_runner.py
+++ b/vllm_omni/platforms/npu/worker/npu_model_runner.py
@@ -54,12 +54,13 @@ class OmniNPUModelRunner(NPUModelRunner):
                     self.model.talker_mtp, self.vllm_config, runtime_mode=CUDAGraphMode.FULL
                 )
             hidden_size = self.model_config.hf_config.talker_config.text_config.hidden_size
-            self.talker_mtp_input_ids = self._make_buffer(self.max_num_reqs, dtype=torch.int32)
+            max_batch_size = max(self.max_num_reqs, self.compilation_config.max_cudagraph_capture_size)
+            self.talker_mtp_input_ids = self._make_buffer(max_batch_size, dtype=torch.int32)
             self.talker_mtp_inputs_embeds = self._make_buffer(
-                self.max_num_reqs, hidden_size, dtype=self.dtype, numpy=False
+                max_batch_size, hidden_size, dtype=self.dtype, numpy=False
             )
-            self.last_talker_hidden = self._make_buffer(self.max_num_reqs, hidden_size, dtype=self.dtype, numpy=False)
-            self.text_step = self._make_buffer(self.max_num_reqs, hidden_size, dtype=self.dtype, numpy=False)
+            self.last_talker_hidden = self._make_buffer(max_batch_size, hidden_size, dtype=self.dtype, numpy=False)
+            self.text_step = self._make_buffer(max_batch_size, hidden_size, dtype=self.dtype, numpy=False)
 
     def _init_mrope_positions(self, req_state: CachedRequestState):
         image_grid_thw = []
@@ -590,12 +591,16 @@ class OmniNPUModelRunner(NPUModelRunner):
                 model_instance=self.model,
             ):
                 if getattr(self.model, "talker", None) is not None and hasattr(self.model, "talker_mtp"):
+                    num_tokens_padded_talker_mtp = num_tokens_padded
+                    if num_tokens_padded_talker_mtp == self.max_num_tokens:
+                        num_tokens_padded_talker_mtp = self.talker_mtp_input_ids.gpu.shape[0]
                     hidden_states = self.talker_mtp(
-                        self.talker_mtp_input_ids.gpu[:num_tokens_padded],
-                        self.talker_mtp_inputs_embeds.gpu[:num_tokens_padded],
-                        self.last_talker_hidden.gpu[:num_tokens_padded],
-                        self.text_step.gpu[:num_tokens_padded],
+                        self.talker_mtp_input_ids.gpu[:num_tokens_padded_talker_mtp],
+                        self.talker_mtp_inputs_embeds.gpu[:num_tokens_padded_talker_mtp],
+                        self.last_talker_hidden.gpu[:num_tokens_padded_talker_mtp],
+                        self.text_step.gpu[:num_tokens_padded_talker_mtp],
                     )
+                    self.compilation_config.cache_dir = None
                 hidden_states = self._generate_dummy_run_hidden_states(
                     input_ids, positions, num_tokens_padded, intermediate_tensors, inputs_embeds
                 )
diff --git a/vllm_omni/version.py b/vllm_omni/version.py
index 0a231e9bd..15253fe6f 100644
--- a/vllm_omni/version.py
+++ b/vllm_omni/version.py
@@ -1,3 +1,3 @@
-__version__ = "0.14.0rc1"
-__version_tuple__ = (0, 14, "0rc1")
+__version__ = "0.14.0"
+__version_tuple__ = (0, 14, 0)
 # TODO: add auto version generation