Merge remote-tracking branch 'origin/main' into dev/vllm-align

.buildkite/pipeline.yml

@@ -176,7 +176,7 @@ steps:
   #                 type: DirectoryOrCreate
 
   - label: "Omni Model Test"
-    timeout_in_minutes: 15
+    timeout_in_minutes: 20
     depends_on: image-build
     commands:
       - export VLLM_LOGGING_LEVEL=DEBUG

.github/scripts/pr_reviewer.py

@@ -0,0 +1,629 @@
+#!/usr/bin/env python3
+"""
+PR Reviewer using GLM API for vllm-omni project.
+"""
+
+import json
+import logging
+import os
+import sys
+import time
+from dataclasses import dataclass
+from typing import Any, TypedDict
+
+import requests
+
+
+# Type definitions for API responses
+class PRDetails(TypedDict):
+    """Type definition for GitHub PR details response."""
+
+    title: str
+    body: str
+    number: int
+    state: str
+    user: dict[str, Any]
+
+
+class GLMMessage(TypedDict):
+    """Type definition for GLM API message."""
+
+    role: str
+    content: str
+
+
+class GLMChoice(TypedDict):
+    """Type definition for GLM API choice."""
+
+    message: GLMMessage
+    finish_reason: str
+
+
+class GLMResponse(TypedDict):
+    """Type definition for GLM API response."""
+
+    choices: list[GLMChoice]
+    usage: dict[str, int] | None
+
+
+class GitHubComment(TypedDict):
+    """Type definition for GitHub comment."""
+
+    id: int
+    body: str
+    created_at: str
+    user: dict[str, Any]
+
+
+# Configuration
+TRIGGER_PHRASE: str = "@vllm-omni-reviewer"
+DEFAULT_GLM_API_URL: str = "https://open.bigmodel.cn/api/paas/v4/chat/completions"  # noqa: E501
+DEFAULT_GLM_MODEL: str = "glm-5"
+DEFAULT_COOLDOWN_MINUTES: int = 5
+DEFAULT_MAX_RETRIES: int = 3
+DEFAULT_RETRY_DELAY: float = 1.0
+MAX_DIFF_SIZE: int = 100_000  # Maximum diff size in characters
+
+
+@dataclass
+class Config:
+    """Configuration for the PR reviewer."""
+
+    glm_api_url: str
+    glm_model: str
+    cooldown_minutes: int
+    max_retries: int
+    retry_delay: float
+    max_diff_size: int
+
+
+# Setup logging
+logging.basicConfig(
+    level=logging.INFO,
+    format="[PR Reviewer] %(asctime)s - %(levelname)s - %(message)s",
+    datefmt="%Y-%m-%d %H:%M:%S",
+)
+logger: logging.Logger = logging.getLogger(__name__)
+
+
+def get_config() -> Config:
+    """Load configuration from environment variables with defaults."""
+    return Config(
+        glm_api_url=os.getenv("GLM_API_URL", DEFAULT_GLM_API_URL),
+        glm_model=os.getenv("GLM_MODEL", DEFAULT_GLM_MODEL),
+        cooldown_minutes=int(
+            os.getenv(
+                "PR_REVIEWER_COOLDOWN_MINUTES",
+                str(DEFAULT_COOLDOWN_MINUTES),
+            )
+        ),
+        max_retries=int(
+            os.getenv(
+                "PR_REVIEWER_MAX_RETRIES",
+                str(DEFAULT_MAX_RETRIES),
+            )
+        ),
+        retry_delay=float(os.getenv("PR_REVIEWER_RETRY_DELAY", str(DEFAULT_RETRY_DELAY))),
+        max_diff_size=int(os.getenv("PR_REVIEWER_MAX_DIFF_SIZE", str(MAX_DIFF_SIZE))),  # noqa: E501
+    )
+
+
+def get_env_var(name: str) -> str:
+    """
+    Get an environment variable or raise an error.
+
+    Args:
+        name: Name of the environment variable.
+
+    Returns:
+        The value of the environment variable.
+
+    Raises:
+        SystemExit: If the environment variable is not set.
+    """
+    value = os.environ.get(name)
+    if not value:
+        logger.error(f"Environment variable {name} is not set")
+        sys.exit(1)
+    return value
+
+
+def check_trigger(comment_body: str) -> bool:
+    """
+    Check if the comment contains the trigger phrase.
+
+    Args:
+        comment_body: The body of the comment to check.
+
+    Returns:
+        True if the trigger phrase is found, False otherwise.
+    """
+    return TRIGGER_PHRASE in comment_body
+
+
+def fetch_pr_diff(
+    repo_name: str,
+    pr_number: int,
+    token: str,
+    max_size: int = MAX_DIFF_SIZE,
+) -> str | None:
+    """
+    Fetch the diff for a pull request.
+
+    Args:
+        repo_name: The repository name in format "owner/repo".
+        pr_number: The pull request number.
+        token: GitHub authentication token.
+        max_size: Maximum diff size in characters.
+
+    Returns:
+        The diff content as a string, or None if fetching failed.
+        Returns empty string if diff is larger than max_size.
+    """
+    url: str = f"https://api.github.com/repos/{repo_name}/pulls/{pr_number}"
+    headers: dict[str, str] = {
+        "Authorization": f"Bearer {token}",
+        "Accept": "application/vnd.github.v3.diff",
+    }
+
+    logger.info(f"Fetching PR diff from {url}")
+    response = requests.get(url, headers=headers, timeout=30)
+
+    if response.status_code == 200:
+        diff: str = response.text
+        if len(diff) > max_size:
+            logger.warning(
+                f"Diff size ({len(diff)} bytes) exceeds maximum "
+                f"({max_size} bytes), truncating to first "
+                f"{max_size} bytes"
+            )
+            return diff[:max_size] + "\n\n... [Diff truncated due to size] ..."
+        logger.info(f"Successfully fetched diff ({len(diff)} bytes)")
+        return diff
+    else:
+        logger.error(f"Failed to fetch PR diff: {response.status_code}")
+        logger.error(f"Response: {response.text}")
+        return None
+
+
+def fetch_pr_details(
+    repo_name: str,
+    pr_number: int,
+    token: str,
+) -> PRDetails | None:
+    """
+    Fetch PR details including title and description.
+
+    Args:
+        repo_name: The repository name in format "owner/repo".
+        pr_number: The pull request number.
+        token: GitHub authentication token.
+
+    Returns:
+        A dictionary containing PR details, or None if fetching failed.
+    """
+    url: str = f"https://api.github.com/repos/{repo_name}/pulls/{pr_number}"
+    headers: dict[str, str] = {
+        "Authorization": f"Bearer {token}",
+        "Accept": "application/vnd.github.v3+json",
+    }
+
+    logger.info(f"Fetching PR details from {url}")
+    response = requests.get(url, headers=headers, timeout=30)
+
+    if response.status_code == 200:
+        return response.json()
+    else:
+        logger.error(f"Failed to fetch PR details: {response.status_code}")
+        return None
+
+
+def build_review_prompt(pr_title: str, pr_description: str, diff: str) -> str:
+    """
+    Build the prompt for the GLM-4.7 API.
+
+    Args:
+        pr_title: The title of the pull request.
+        pr_description: The description/body of the pull request.
+        diff: The diff content of the pull request.
+
+    Returns:
+        The formatted prompt string for the API.
+    """
+    return f"""You are an expert code reviewer for the VLLM-Omni project. \
+Please review the following pull request:
+
+## Pull Request Details
+**Title:** {pr_title}
+
+**Description:**
+{pr_description if pr_description else "No description provided."}
+
+## Code Changes (Diff)
+{diff}
+
+## Review Guidelines
+
+Please provide a comprehensive code review with the following sections:
+
+### 1. Overview
+- Brief summary of the changes
+- Overall assessment (positive, neutral, or concerns)
+
+### 2. Code Quality
+- Code style and consistency
+- Potential bugs or edge cases
+- Performance considerations
+- Error handling
+
+### 3. Architecture & Design
+- Integration with existing codebase
+- Design patterns and best practices
+- Potential improvements
+
+### 4. Security & Safety
+- Security concerns (if any)
+- Resource management
+- Input validation
+
+### 5. Testing & Documentation
+- Test coverage considerations
+- Documentation completeness
+- Examples and usage clarity
+
+### 6. Specific Suggestions
+- Line-by-line specific feedback (use `file:line` format)
+- Concrete actionable suggestions
+- Code examples for improvements (if applicable)
+
+### 7. Approval Status
+- **LGTM** (Looks Good To Me) if the PR is ready to merge
+- **LGTM with suggestions** if the PR is good but has minor suggestions
+- **Changes requested** if significant changes are needed
+
+## Important Notes
+- Be constructive and helpful
+- Focus on objective technical feedback
+- Acknowledge good practices when you see them
+- Prioritize critical issues over nitpicks
+- If the diff is empty or minimal, acknowledge this and provide
+  any relevant context-specific guidance
+
+Please format your response in Markdown with clear section headers.
+"""
+
+
+def validate_glm_response(data: dict[str, Any]) -> str | None:
+    """
+    Validate and extract content from GLM API response.
+
+    Args:
+        data: The response data from GLM API.
+
+    Returns:
+        The review content string if valid, None otherwise.
+    """
+    # Check if choices exists and is a non-empty list
+    if "choices" not in data:
+        logger.error("GLM API response missing 'choices' field")
+        logger.error(f"Response structure: {json.dumps(data, indent=2)}")
+        return None
+
+    choices = data["choices"]
+    if not isinstance(choices, list):
+        logger.error(f"GLM API 'choices' is not a list: {type(choices)}")
+        return None
+
+    if len(choices) == 0:
+        logger.error("GLM API 'choices' is an empty list")
+        return None
+
+    # Check if first choice has message
+    try:
+        first_choice = choices[0]
+        if not isinstance(first_choice, dict):
+            logger.error(f"GLM API choice is not a dict: {type(first_choice)}")
+            return None
+
+        if "message" not in first_choice:
+            logger.error("GLM API choice missing 'message' field")
+            logger.error(f"Choice structure: {json.dumps(first_choice, indent=2)}")  # noqa: E501
+            return None
+
+        message = first_choice["message"]
+        if not isinstance(message, dict):
+            logger.error(f"GLM API message is not a dict: {type(message)}")
+            return None
+
+        if "content" not in message:
+            logger.error("GLM API message missing 'content' field")
+            logger.error(f"Message structure: {json.dumps(message, indent=2)}")
+            return None
+
+        content = message["content"]
+        if not isinstance(content, str):
+            logger.error(f"GLM API content is not a string: {type(content)}")
+            return None
+
+        return content
+
+    except (KeyError, IndexError, TypeError) as e:
+        logger.error(f"Failed to parse GLM API response: {e}")
+        logger.error(f"Response: {json.dumps(data, indent=2)}")
+        return None
+
+
+def call_glm_api(prompt: str, api_key: str, config: Config) -> str | None:
+    """
+    Call the GLM-4.7 API to get code review with retry logic.
+
+    Args:
+        prompt: The prompt to send to the API.
+        api_key: The GLM API key.
+        config: Configuration object.
+
+    Returns:
+        The review content as a string, or None if all retries failed.
+    """
+    headers: dict[str, str] = {
+        "Authorization": f"Bearer {api_key}",
+        "Content-Type": "application/json",
+    }
+
+    payload: dict[str, Any] = {
+        "model": config.glm_model,
+        "messages": [{"role": "user", "content": prompt}],
+        "temperature": 0.3,
+        "max_tokens": 32000,
+        "top_p": 0.9,
+    }
+
+    last_error: str | None = None
+
+    for attempt in range(config.max_retries):
+        try:
+            logger.info(f"Calling GLM API ({config.glm_model}) - Attempt {attempt + 1}/{config.max_retries}")
+            response = requests.post(
+                config.glm_api_url,
+                headers=headers,
+                json=payload,
+                timeout=120,
+            )
+
+            if response.status_code == 200:
+                data = response.json()
+                review = validate_glm_response(data)
+                if review:
+                    logger.info(f"Successfully received review ({len(review)} chars)")  # noqa: E501
+                    return review
+                else:
+                    last_error = "Failed to validate API response structure"
+                    logger.error(last_error)
+            else:
+                last_error = f"GLM API request failed: {response.status_code} - {response.text}"
+                logger.error(last_error)
+
+        except requests.exceptions.Timeout:
+            last_error = f"GLM API request timed out (attempt {attempt + 1})"
+            logger.error(last_error)
+        except requests.exceptions.RequestException as e:
+            last_error = f"GLM API request exception: {e}"
+            logger.error(last_error)
+        except json.JSONDecodeError as e:
+            last_error = f"Failed to decode GLM API response as JSON: {e}"
+            logger.error(last_error)
+
+        # Exponential backoff before retry
+        if attempt < config.max_retries - 1:
+            wait_time: float = config.retry_delay * (2**attempt)
+            logger.info(f"Waiting {wait_time}s before retry...")  # noqa: E501
+            time.sleep(wait_time)
+
+    logger.error(
+        f"All {config.max_retries} attempts failed. Last error: {last_error}"  # noqa: E501
+    )
+    return None
+
+
+def check_cooldown(  # noqa: E501
+    repo_name: str,
+    pr_number: int,
+    token: str,
+    cooldown_minutes: int,
+) -> bool:
+    """
+    Check if the PR is within the cooldown period.
+
+    Args:
+        repo_name: The repository name in format "owner/repo".
+        pr_number: The pull request number.
+        token: GitHub authentication token.
+        cooldown_minutes: Cooldown period in minutes.
+
+    Returns:
+        True if within cooldown period (should skip), False otherwise.
+    """
+    from datetime import datetime, timedelta
+
+    url: str = (
+        f"https://api.github.com/repos/{repo_name}/issues/"
+        f"{pr_number}/comments"  # noqa: E501
+    )
+    headers: dict[str, str] = {
+        "Authorization": f"Bearer {token}",
+        "Accept": "application/vnd.github.v3+json",
+    }
+
+    logger.info(f"Checking cooldown period ({cooldown_minutes} minutes)")
+    response = requests.get(url, headers=headers, timeout=30)
+
+    if response.status_code != 200:
+        logger.warning(f"Failed to check cooldown: {response.status_code}, proceeding with review")
+        return False
+
+    comments: list[dict[str, Any]] = response.json()
+    cutoff_time: datetime = datetime.utcnow() - timedelta(minutes=cooldown_minutes)  # noqa: E501
+
+    for comment in reversed(comments):
+        # Check if this is a bot comment
+        body: str = comment.get("body", "")
+        if "VLLM-Omni PR Review" in body or "PR Reviewer Bot" in body:
+            created_at_str: str = comment.get("created_at", "")
+            try:
+                # Parse GitHub timestamp format
+                created_at = datetime.fromisoformat(created_at_str.replace("Z", "+00:00"))
+                created_at = created_at.replace(tzinfo=None)
+                if created_at > cutoff_time:
+                    logger.info(f"PR is within cooldown period (last review: {created_at_str})")
+                    return True
+            except ValueError:
+                logger.warning(f"Failed to parse comment timestamp: {created_at_str}")  # noqa: E501
+                continue
+
+    logger.info("PR is outside cooldown period, proceeding with review")
+    return False
+
+
+def post_review_comment(  # noqa: E501
+    repo_name: str,
+    pr_number: int,
+    token: str,
+    review: str,
+) -> bool:
+    """
+    Post the review as a comment on the PR.
+
+    Args:
+        repo_name: The repository name in format "owner/repo".
+        pr_number: The pull request number.
+        token: GitHub authentication token.
+        review: The review content to post.
+
+    Returns:
+        True if posting succeeded, False otherwise.
+    """
+    url: str = (
+        f"https://api.github.com/repos/{repo_name}/issues/"
+        f"{pr_number}/comments"  # noqa: E501
+    )
+    headers: dict[str, str] = {
+        "Authorization": f"Bearer {token}",
+        "Accept": "application/vnd.github.v3+json",
+    }
+
+    # Format the review comment
+    comment_body: str = f"""## 🤖 VLLM-Omni PR Review
+
+{review}
+
+---
+*This review was generated automatically by the VLLM-Omni PR Reviewer Bot
+using {os.getenv("GLM_MODEL", DEFAULT_GLM_MODEL)}.*
+"""
+
+    payload: dict[str, str] = {"body": comment_body}
+
+    logger.info(f"Posting review comment to PR #{pr_number}")
+    response = requests.post(url, headers=headers, json=payload, timeout=30)
+
+    if response.status_code == 201:
+        logger.info("Successfully posted review comment")
+        return True
+    else:
+        logger.error(f"Failed to post comment: {response.status_code}")
+        logger.error(f"Response: {response.text}")
+        return False
+
+
+def main() -> int:
+    """
+    Main entry point for the PR reviewer bot.
+
+    Returns:
+        0 on success, 1 on error.
+    """
+    logger.info("VLLM-Omni PR Reviewer Bot starting...")
+
+    # Load configuration
+    config: Config = get_config()
+    logger.info(
+        f"Configuration: model={config.glm_model}, "
+        f"cooldown={config.cooldown_minutes}min, "
+        f"max_retries={config.max_retries}"
+    )
+
+    # Get environment variables
+    token: str = get_env_var("GITHUB_TOKEN")
+    api_key: str = get_env_var("GLM_API_KEY")
+    repo_name: str = get_env_var("REPO_NAME")
+    pr_number_str: str = get_env_var("PR_NUMBER")
+    comment_body: str = get_env_var("COMMENT_BODY")
+
+    try:
+        pr_number: int = int(pr_number_str)
+    except ValueError:
+        logger.error(f"Invalid PR number: {pr_number_str}")
+        return 1
+
+    logger.info(f"Repository: {repo_name}")
+    logger.info(f"PR Number: {pr_number}")
+
+    # Check if the comment contains the trigger phrase
+    if not check_trigger(comment_body):
+        logger.info(
+            f"Comment does not contain trigger phrase '{TRIGGER_PHRASE}', exiting"  # noqa: E501
+        )
+        return 0
+
+    logger.info("Trigger phrase detected! Starting review process...")
+
+    # Check cooldown period
+    if check_cooldown(repo_name, pr_number, token, config.cooldown_minutes):
+        logger.info("Skipping review due to cooldown period")
+        return 0
+
+    # Fetch PR details
+    logger.info("Step 1/4: Fetching PR details...")
+    pr_details: PRDetails | None = fetch_pr_details(repo_name, pr_number, token)  # noqa: E501
+    if not pr_details:
+        logger.error("Failed to fetch PR details")
+        return 1
+
+    pr_title: str = pr_details.get("title", "Unknown")
+    pr_description: str = pr_details.get("body", "")
+
+    logger.info(f"PR Title: {pr_title}")
+
+    # Fetch PR diff
+    logger.info("Step 2/4: Fetching PR diff...")
+    diff: str | None = fetch_pr_diff(repo_name, pr_number, token, config.max_diff_size)
+    if diff is None:
+        logger.error("Failed to fetch PR diff")
+        return 1
+
+    if not diff:
+        logger.warning("Warning: Empty diff - this might be a draft PR or no code changes")
+
+    # Build prompt
+    logger.info("Step 3/4: Building review prompt...")
+    prompt: str = build_review_prompt(pr_title, pr_description, diff)
+
+    # Call GLM API
+    logger.info("Step 4/4: Calling GLM API...")
+    review: str | None = call_glm_api(prompt, api_key, config)
+    if not review:
+        logger.error("Failed to get review from GLM API")
+        return 1
+
+    # Post review comment
+    logger.info("Posting review comment...")
+    if not post_review_comment(repo_name, pr_number, token, review):
+        logger.error("Failed to post review comment")
+        return 1
+
+    logger.info("PR review completed successfully!")
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())

.github/workflows/pr-reviewer.yml

@@ -0,0 +1,62 @@
+name: VLLM-Omni PR Reviewer
+
+on:
+  issue_comment:
+    types: [created]
+
+permissions:
+  contents: read
+  pull-requests: write
+  issues: write
+
+jobs:
+  pr-reviewer:
+    name: Review Pull Request
+    runs-on: ubuntu-latest
+    timeout-minutes: 10
+    # Only run when the comment is from a collaborator/owner/member
+    if: |
+      github.event_name == 'issue_comment' &&
+      github.event.issue.pull_request != null &&
+      contains(github.event.comment.body, '@vllm-omni-reviewer') &&
+      (github.event.comment.author_association == 'MEMBER' ||
+       github.event.comment.author_association == 'COLLABORATOR' ||
+       github.event.comment.author_association == 'OWNER')
+    steps:
+      - name: Checkout repository
+        uses: actions/checkout@v4.2.2
+        with:
+          fetch-depth: 0
+
+      - name: Set up Python
+        uses: actions/setup-python@v5.3.0
+        with:
+          python-version: '3.11'
+          cache: 'pip'
+
+      - name: Install dependencies
+        run: |
+          python -m pip install --upgrade pip==24.0
+          pip install requests==2.31.0 pyyaml==6.0.1
+
+      - name: Run PR Reviewer
+        id: reviewer
+        env:
+          GLM_API_KEY: ${{ secrets.GLM_API_KEY }}
+          GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
+          PR_NUMBER: ${{ github.event.issue.number || github.event.pull_request.number }}
+          COMMENT_BODY: ${{ github.event.comment.body }}
+          REPO_NAME: ${{ github.repository }}
+          PR_REVIEWER_COOLDOWN_MINUTES: 5
+          PR_REVIEWER_MAX_RETRIES: 3
+        run: |
+          python .github/scripts/pr_reviewer.py 2>&1 | tee "/tmp/pr_review_${PR_NUMBER}.log"
+
+      - name: Upload review logs
+        if: always()
+        uses: actions/upload-artifact@v4.4.3
+        with:
+          name: pr-review-logs-${{ github.event.issue.number || github.event.pull_request.number }}
+          path: /tmp/pr_review_*.log
+          retention-days: 7
+          if-no-files-found: ignore

docker/Dockerfile.xpu

@@ -0,0 +1,139 @@
+# Argument to configure vllm base image if pre-built
+ARG VLLM_BASE=vllm-base
+
+FROM intel/deep-learning-essentials:2025.3.2-0-devel-ubuntu24.04 AS vllm-base
+
+WORKDIR /workspace/
+
+ARG PYTHON_VERSION=3.12
+ARG PIP_EXTRA_INDEX_URL="https://download.pytorch.org/whl/xpu"
+
+RUN wget -O- https://apt.repos.intel.com/intel-gpg-keys/GPG-PUB-KEY-INTEL-SW-PRODUCTS.PUB | gpg --dearmor | tee /usr/share/keyrings/oneapi-archive-keyring.gpg > /dev/null && \
+    echo "deb [signed-by=/usr/share/keyrings/oneapi-archive-keyring.gpg] https://apt.repos.intel.com/oneapi all main" | tee /etc/apt/sources.list.d/oneAPI.list && \
+    add-apt-repository -y ppa:kobuk-team/intel-graphics
+
+RUN apt clean && apt-get update -y && \
+    apt-get install -y --no-install-recommends --fix-missing \
+    curl \
+    ffmpeg \
+    git \
+    libsndfile1 \
+    libsm6 \
+    libxext6 \
+    libgl1 \
+    lsb-release \
+    libaio-dev \
+    numactl \
+    wget \
+    vim \
+    python3.12 \
+    python3.12-dev \
+    python3-pip
+
+RUN apt update && apt upgrade -y && \
+    apt install -y libze1 libze-dev libze-intel-gpu1 intel-opencl-icd libze-intel-gpu-raytracing intel-ocloc && \
+    apt install -y intel-oneapi-compiler-dpcpp-cpp-2025.3
+
+ENV PATH="/root/.local/bin:$PATH"
+ENV VIRTUAL_ENV="/opt/venv"
+ENV UV_PYTHON_INSTALL_DIR=/opt/uv/python
+RUN curl -LsSf https://astral.sh/uv/install.sh | sh
+RUN uv venv --python ${PYTHON_VERSION} --seed ${VIRTUAL_ENV}
+ENV PATH="$VIRTUAL_ENV/bin:$PATH"
+
+# This oneccl contains the BMG support which is not the case for default version of oneapi 2025.2.
+ARG ONECCL_INSTALLER="intel-oneccl-2021.15.7.8_offline.sh"
+RUN wget "https://github.com/uxlfoundation/oneCCL/releases/download/2021.15.7/${ONECCL_INSTALLER}" && \
+    bash "${ONECCL_INSTALLER}" -a --silent --eula accept && \
+    rm "${ONECCL_INSTALLER}" && \
+    echo "source /opt/intel/oneapi/setvars.sh --force" >> /root/.bashrc && \
+    echo "source /opt/intel/oneapi/ccl/2021.15/env/vars.sh --force" >> /root/.bashrc
+RUN rm -f /opt/intel/oneapi/ccl/latest && \
+    ln -s /opt/intel/oneapi/ccl/2021.15 /opt/intel/oneapi/ccl/latest
+
+SHELL ["bash", "-c"]
+CMD ["bash", "-c", "source /root/.bashrc && exec bash"]
+
+WORKDIR /workspace/
+ENV UV_HTTP_TIMEOUT=500
+
+# Configure package index for XPU
+ENV PIP_EXTRA_INDEX_URL=${PIP_EXTRA_INDEX_URL}
+ENV UV_EXTRA_INDEX_URL=${PIP_EXTRA_INDEX_URL}
+ENV UV_INDEX_STRATEGY="unsafe-best-match"
+ENV UV_LINK_MODE="copy"
+
+ARG VLLM_VERSION=v0.16.0
+RUN git clone -b ${VLLM_VERSION} https://github.com/vllm-project/vllm
+WORKDIR /workspace/vllm
+
+RUN --mount=type=cache,target=/root/.cache/uv \
+    uv pip install --upgrade pip && \
+    uv pip install -r requirements/xpu.txt
+
+ # used for suffix method speculative decoding
+ # build deps for proto + nanobind-based extensions to set up the build environment
+RUN --mount=type=cache,target=/root/.cache/uv \
+    uv pip install grpcio-tools protobuf nanobind
+ # arctic-inference is built from source which needs torch-xpu properly installed first
+RUN --mount=type=cache,target=/root/.cache/uv \
+    source /opt/intel/oneapi/setvars.sh --force && \
+    source /opt/intel/oneapi/ccl/2021.15/env/vars.sh --force && \
+    export CMAKE_PREFIX_PATH="$(python -c 'import site; print(site.getsitepackages()[0])'):${CMAKE_PREFIX_PATH}" && \
+    uv pip install --no-build-isolation arctic-inference==0.1.1
+
+ENV LD_LIBRARY_PATH="$LD_LIBRARY_PATH:/usr/local/lib/"
+
+ENV VLLM_TARGET_DEVICE=xpu
+ENV VLLM_WORKER_MULTIPROC_METHOD=spawn
+
+RUN --mount=type=cache,target=/root/.cache/uv \
+    uv pip install --no-build-isolation .
+
+CMD ["/bin/bash"]
+
+FROM vllm-base AS vllm-openai
+
+# install additional dependencies for openai api server
+RUN --mount=type=cache,target=/root/.cache/uv \
+    uv pip install accelerate hf_transfer pytest pytest_asyncio lm_eval[api] modelscope
+
+# install development dependencies (for testing)
+RUN uv pip install -e tests/vllm_test_utils
+
+# install nixl from source code
+ENV NIXL_VERSION=0.7.0
+RUN python /workspace/vllm/tools/install_nixl_from_source_ubuntu.py
+
+# ensure vllm is properly installed
+RUN python -c "import vllm, inspect; print(vllm.__file__)"
+RUN uv pip show vllm
+
+CMD ["/bin/bash"]
+
+ENTRYPOINT []
+
+FROM ${VLLM_BASE} AS vllm-omni
+
+WORKDIR /workspace/vllm-omni
+COPY . .
+
+ENV VLLM_OMNI_TARGET_DEVICE=xpu
+RUN uv pip install --no-cache-dir ".[dev]" --no-build-isolation
+
+# FIX triton
+RUN --mount=type=cache,target=/root/.cache/uv \
+    uv pip uninstall triton triton-xpu && \
+    uv pip install triton-xpu==3.6.0
+
+# remove torch bundled oneccl to avoid conflicts
+RUN --mount=type=cache,target=/root/.cache/uv \
+    uv pip uninstall oneccl oneccl-devel
+
+FROM vllm-omni AS vllm-omni-openai
+
+RUN ln -sf /usr/bin/python3 /usr/bin/python
+
+ENV VLLM_WORKER_MULTIPROC_METHOD=spawn
+
+ENTRYPOINT ["vllm", "serve", "--omni"]

docs/getting_started/installation/README.md

@@ -5,4 +5,5 @@ vLLM-Omni supports the following hardware platforms:
 - [GPU](gpu.md)
     - [NVIDIA CUDA](gpu.md)
     - [AMD ROCm](gpu.md)
+    - [Intel XPU](gpu.md)
 - [NPU](npu.md)

docs/getting_started/installation/gpu.md

@@ -18,6 +18,10 @@ vLLM-Omni is a Python library that supports the following GPU variants. The libr
 
     --8<-- "docs/getting_started/installation/gpu/rocm.inc.md:requirements"
 
+=== "Intel XPU"
+
+    --8<-- "docs/getting_started/installation/gpu/xpu.inc.md:requirements"
+
 ## Set up using Python
 
 ### Create a new Python environment
@@ -32,11 +36,14 @@ Note: Pre-built wheels are currently only available for vLLM-Omni 0.11.0rc1, 0.1
 
     --8<-- "docs/getting_started/installation/gpu/cuda.inc.md:pre-built-wheels"
 
-
 === "AMD ROCm"
 
     --8<-- "docs/getting_started/installation/gpu/rocm.inc.md:pre-built-wheels"
 
+=== "Intel XPU"
+
+    --8<-- "docs/getting_started/installation/gpu/xpu.inc.md:pre-built-wheels"
+
 [](){ #build-from-source }
 
 ### Build wheel from source
@@ -49,6 +56,10 @@ Note: Pre-built wheels are currently only available for vLLM-Omni 0.11.0rc1, 0.1
 
     --8<-- "docs/getting_started/installation/gpu/rocm.inc.md:build-wheel-from-source"
 
+=== "Intel XPU"
+
+    --8<-- "docs/getting_started/installation/gpu/xpu.inc.md:build-wheel-from-source"
+
 ## Set up using Docker
 
 ### Pre-built images
@@ -61,8 +72,16 @@ Note: Pre-built wheels are currently only available for vLLM-Omni 0.11.0rc1, 0.1
 
     --8<-- "docs/getting_started/installation/gpu/rocm.inc.md:pre-built-images"
 
+=== "Intel XPU"
+
+    --8<-- "docs/getting_started/installation/gpu/xpu.inc.md:pre-built-images"
+
 ### Build your own docker image
 
 === "AMD ROCm"
 
     --8<-- "docs/getting_started/installation/gpu/rocm.inc.md:build-docker"
+
+=== "Intel XPU"
+
+    --8<-- "docs/getting_started/installation/gpu/xpu.inc.md:build-docker"

docs/getting_started/installation/gpu/xpu.inc.md

@@ -0,0 +1,48 @@
+# --8<-- [start:requirements]
+
+- GPU: Validated on Intel® Arc™ B-Series.
+
+# --8<-- [end:requirements]
+# --8<-- [start:set-up-using-python]
+
+vLLM-Omni currently recommends using the Docker image setup steps below.
+
+# --8<-- [start:pre-built-wheels]
+
+# --8<-- [end:pre-built-wheels]
+
+# --8<-- [start:build-wheel-from-source]
+
+# --8<-- [end:build-wheel-from-source]
+
+# --8<-- [start:build-docker]
+
+#### Build docker image
+
+```bash
+DOCKER_BUILDKIT=1 docker build -f docker/Dockerfile.xpu -t vllm-omni-xpu --shm-size=4g .
+```
+
+#### Launch the docker image
+
+##### Launch with OpenAI API Server
+
+```
+docker run -it -d --shm-size 10g \
+  --name {container_name} \
+  --net=host \
+  --ipc=host \
+  --privileged \
+  -v /dev/dri/by-path:/dev/dri/by-path \
+  --device /dev/dri:/dev/dri \
+  -v ~/.cache/huggingface:/root/.cache/huggingface \
+  --env "HF_TOKEN=$HF_TOKEN" \
+  vllm-omni-xpu \
+  --model Qwen/Qwen2.5-Omni-3B --port 8091
+```
+
+# --8<-- [end:build-docker]
+
+# --8<-- [start:pre-built-images]
+
+# --8<-- [end:pre-built-images]

docs/pr_reviewer.md

@@ -0,0 +1,255 @@
+# VLLM-Omni PR Reviewer
+
+## Overview
+
+The VLLM-Omni PR Reviewer is an automated code review bot powered by GLM-4.7 AI model. It helps maintain code quality by providing intelligent feedback on pull requests.
+
+## Features
+
+- **Intelligent Code Analysis**: Leverages GLM-4.7 for understanding code context and providing meaningful feedback
+- **Comprehensive Reviews**: Covers code quality, architecture, security, testing, and documentation
+- **Structured Output**: Provides well-formatted reviews with clear sections and actionable suggestions
+- **Rate Limiting**: Built-in cooldown mechanism to prevent excessive API usage
+- **Retry Logic**: Automatic retries with exponential backoff for transient API failures
+- **Defensive Parsing**: Robust validation of API responses to handle malformed data
+- **Cost Control**: Only repository members/collaborators/owners can trigger reviews
+
+## How to Use
+
+### Triggering a Review
+
+To trigger an automated PR review, mention the bot in a PR comment:
+
+```
+@vllm-omni-reviewer please review
+```
+
+Or include in your PR description:
+
+```
+@vllm-omni-reviewer
+```
+
+The bot will automatically review your changes and post a detailed comment.
+
+## What Gets Reviewed
+
+- **vLLM Architecture Compatibility**: Ensures changes align with vLLM's design patterns
+- **Multi-modal Integration**: Reviews audio, vision, and text processing implementations
+- **Performance Implications**: Analyzes impact on inference latency and throughput
+- **Code Quality**: Checks Python best practices, type hints, and documentation
+- **Security Considerations**: Identifies potential security vulnerabilities
+- **Testing Coverage**: Recommends additional test cases when needed
+
+## Review Output
+
+The bot posts a structured review comment with:
+
+- **Overview**: Brief summary of the PR's purpose
+- **Critical Issues (Must Fix)**: Blocking issues that need to be addressed
+- **Important Issues (Should Fix)**: Significant concerns that should be resolved
+- **Minor Issues & Suggestions**: Small improvements and optional suggestions
+- **Positive Aspects**: Highlights well-implemented features
+- **Performance Considerations**: Analysis of performance impact
+- **Testing Recommendations**: Suggestions for additional tests
+- **Overall Assessment**: Final recommendation (Approve/Request Changes/Needs Major Work)
+
+## Rate Limiting and Cooldown
+
+The bot includes a cooldown mechanism to prevent excessive API usage:
+
+- **Default cooldown**: 5 minutes between reviews per PR
+- **Configurable**: Can be adjusted via `PR_REVIEWER_COOLDOWN_MINUTES` environment variable
+- **Smart detection**: Checks for previous bot comments before starting a review
+
+If you trigger a review within the cooldown period, the bot will log a message and skip the review.
+
+## Architecture
+
+```
+┌─────────────────┐
+│  PR Comment     │
+│  @vllm-omni-    │
+│  reviewer       │
+└────────┬────────┘
+         │
+         ▼
+┌─────────────────────────────────┐
+│  GitHub Actions Workflow        │
+│  (.github/workflows/            │
+│   pr-reviewer.yml)              │
+│                                 │
+│  - Python 3.11                  │
+│  - requests==2.31.0             │
+│  - pyyaml==6.0.1                │
+└────────┬────────────────────────┘
+         │
+         ▼
+┌─────────────────────────────────┐
+│  PR Reviewer Script             │
+│  (.github/scripts/              │
+│   pr_reviewer.py)               │
+│                                 │
+│  1. Check cooldown              │
+│  2. Fetch PR details & diff     │
+│  3. Build review prompt         │
+│  4. Call GLM-4.7 API            │
+│     (with retry logic)          │
+│  5. Validate response           │
+│  6. Post review comment         │
+└────────┬────────────────────────┘
+         │
+         ▼
+┌─────────────────────────────────┐
+│  GLM-4.7 API                    │
+│  (open.bigmodel.cn)             │
+└─────────────────────────────────┘
+```
+
+1. **GitHub Actions Workflow** (`.github/workflows/pr-reviewer.yml`): Triggers on @mention
+2. **Python Script** (`.github/scripts/pr_reviewer.py`): Fetches PR data and calls GLM-4.7 API
+3. **GLM-4.7 API**: Provides intelligent code analysis
+
+## Testing
+
+### Testing the PR Reviewer Bot
+
+To test the PR reviewer bot before deploying to production:
+
+1. **Create a test PR** - Make a small, safe change (e.g., documentation update)
+2. **Open the PR** - Create a pull request with a descriptive title
+3. **Trigger the review** - Comment `@vllm-omni-reviewer` on the PR
+4. **Monitor results** - Check the Actions tab for workflow execution logs
+
+### Running Unit Tests
+
+The bot includes comprehensive unit tests that can be run locally:
+
+```bash
+# Run all tests
+pytest .github/tests/test_pr_reviewer.py -v
+
+# Run specific test
+pytest .github/tests/test_pr_reviewer.py::TestCheckTrigger -v
+
+# Run with coverage
+pytest .github/tests/test_pr_reviewer.py --cov=.github/scripts/pr_reviewer.py --cov-report=term-missing
+```
+
+### What to Look For
+
+When testing, verify that:
+- [ ] The workflow triggers on the `@vllm-omni-reviewer` comment
+- [ ] The cooldown mechanism works correctly
+- [ ] The GLM API call completes without errors (with retry if needed)
+- [ ] A review comment is posted to the PR
+- [ ] The review content is meaningful and well-structured
+- [ ] The cost is within the expected range (0.50-5 CNY)
+
+### Safe Test Changes
+
+For testing, consider making these types of safe changes:
+- Documentation updates (like adding this Testing section)
+- Comment improvements
+- README enhancements
+- Non-functional file additions
+
+### Example Test PR
+
+A good test PR might:
+- Update a documentation file
+- Add explanatory comments
+- Improve code formatting
+- Fix a minor typo
+
+These changes are safe to merge if the test is successful and won't affect functionality.
+
+## Troubleshooting
+
+### Bot Doesn't Respond
+
+1. **Check permissions** - Verify you have Owner/Member/Collaborator access
+2. **Check Actions tab** - Look for workflow execution and view logs
+3. **Check cooldown** - If another review was posted recently, wait for the cooldown period
+4. **Check API key** - Ensure `GLM_API_KEY` is configured in repository secrets
+
+### API Errors
+
+If the GLM API call fails:
+- Check the Actions tab for detailed error logs
+- Verify the `GLM_API_KEY` secret is correctly configured
+- Ensure sufficient API quota is available
+- The bot will automatically retry up to 3 times with exponential backoff
+
+### Review Seems Truncated
+
+If the review appears incomplete:
+- Large diffs may be truncated at 100,000 characters
+- Check the logs for truncation warnings
+- Consider breaking large PRs into smaller chunks
+
+## Configuration
+
+### Required Secrets
+
+The following secret must be configured in the repository settings:
+
+- `GLM_API_KEY` - Your GLM (BigModel) API key for accessing the GLM-4.7 API
+
+To add the secret:
+1. Go to repository Settings → Secrets and variables → Actions
+2. Click "New repository secret"
+3. Name: `GLM_API_KEY`
+4. Value: Your GLM API key
+
+### Optional Configuration
+
+The following optional environment variables can be set in the workflow file:
+
+| Variable | Default | Description |
+|----------|---------|-------------|
+| `GLM_API_URL` | `https://open.bigmodel.cn/api/paas/v4/chat/completions` | GLM API endpoint |
+| `GLM_MODEL` | `glm-4.7` | Model to use for reviews |
+| `PR_REVIEWER_COOLDOWN_MINUTES` | `5` | Cooldown period between reviews |
+| `PR_REVIEWER_MAX_RETRIES` | `3` | Maximum API retry attempts |
+| `PR_REVIEWER_RETRY_DELAY` | `1.0` | Base delay for retry backoff (seconds) |
+| `PR_REVIEWER_MAX_DIFF_SIZE` | `100000` | Maximum diff size before truncation |
+
+### Workflow Customization
+
+The workflow can be customized in `.github/workflows/pr-reviewer.yml`:
+- Change Python version (default: 3.11)
+- Adjust timeout value (default: 10 minutes)
+- Modify trigger conditions
+- Add additional dependencies
+
+## Code Quality
+
+The PR reviewer script follows vllm-omni coding standards:
+
+- **Type hints**: All functions have complete type hints following mypy strict mode
+- **Logging**: Uses Python's logging module for structured logging
+- **Testing**: Comprehensive unit tests with pytest
+- **Pre-commit**: Script is checked by pre-commit hooks (flake8)
+
+## Cost Estimate
+
+| Component | Cost |
+|-----------|------|
+| GitHub Actions (public repo) | Free |
+| GLM API (glm-4.7) | ~0.50-5 CNY per PR (varies by size) |
+| Total (20 PRs/month) | ~10-100 CNY/month (~$2-15 USD) |
+
+## Contributing
+
+To improve the PR reviewer bot:
+
+1. Edit `.github/scripts/pr_reviewer.py` for logic changes
+2. Edit `.github/workflows/pr-reviewer.yml` for workflow changes
+3. Add tests to `.github/tests/test_pr_reviewer.py`
+4. Run `pre-commit run --files .github/scripts/pr_reviewer.py` to check code quality
+5. Test thoroughly with a test PR before deploying to production
+
+## License
+
+This bot is part of the VLLM-Omni project and follows the same license terms.

examples/offline_inference/qwen3_tts/end2end.py

@@ -4,18 +4,21 @@ Provides single and batch sample inputs for CustomVoice, VoiceDesign, and Base
 tasks, then runs Omni generation and saves output wav files.
 """
 
+import logging
 import os
-from typing import NamedTuple
+from typing import Any, NamedTuple
 
 import soundfile as sf
+import torch
 
 os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
 
-from vllm import SamplingParams
 from vllm.utils.argparse_utils import FlexibleArgumentParser
 
 from vllm_omni import Omni
 
+logger = logging.getLogger(__name__)
+
 
 class QueryResult(NamedTuple):
     """Container for a prepared Omni request."""
@@ -24,6 +27,44 @@ class QueryResult(NamedTuple):
     model_name: str
 
 
+def _estimate_prompt_len(
+    additional_information: dict[str, Any],
+    model_name: str,
+    _cache: dict[str, Any] = {},
+) -> int:
+    """Estimate prompt_token_ids placeholder length for the Talker stage.
+
+    The AR Talker replaces all input embeddings via ``preprocess``, so the
+    placeholder values are irrelevant but the **length** must match the
+    embeddings that ``preprocess`` will produce.
+    """
+    try:
+        from vllm_omni.model_executor.models.qwen3_tts.configuration_qwen3_tts import Qwen3TTSConfig
+        from vllm_omni.model_executor.models.qwen3_tts.qwen3_tts_talker import (
+            Qwen3TTSTalkerForConditionalGeneration,
+        )
+
+        if model_name not in _cache:
+            from transformers import AutoTokenizer
+
+            tok = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True, padding_side="left")
+            cfg = Qwen3TTSConfig.from_pretrained(model_name, trust_remote_code=True)
+            _cache[model_name] = (tok, getattr(cfg, "talker_config", None))
+
+        tok, tcfg = _cache[model_name]
+        task_type = (additional_information.get("task_type") or ["CustomVoice"])[0]
+        return Qwen3TTSTalkerForConditionalGeneration.estimate_prompt_len_from_additional_information(
+            additional_information=additional_information,
+            task_type=task_type,
+            tokenize_prompt=lambda t: tok(t, padding=False)["input_ids"],
+            codec_language_id=getattr(tcfg, "codec_language_id", None),
+            spk_is_dialect=getattr(tcfg, "spk_is_dialect", None),
+        )
+    except Exception as exc:
+        logger.warning("Failed to estimate prompt length, using fallback 2048: %s", exc)
+        return 2048
+
+
 def get_custom_voice_query(use_batch_sample: bool = False) -> QueryResult:
     """Build CustomVoice sample inputs.
 
@@ -34,6 +75,7 @@ def get_custom_voice_query(use_batch_sample: bool = False) -> QueryResult:
         QueryResult with Omni inputs and the CustomVoice model path.
     """
     task_type = "CustomVoice"
+    model_name = "Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice"
     if use_batch_sample:
         texts = ["其实我真的有发现，我是一个特别善于观察别人情绪的人。", "She said she would be here by noon."]
         instructs = ["", "Very happy."]
@@ -41,18 +83,18 @@ def get_custom_voice_query(use_batch_sample: bool = False) -> QueryResult:
         speakers = ["Vivian", "Ryan"]
         inputs = []
         for text, instruct, language, speaker in zip(texts, instructs, languages, speakers):
-            prompt = f"<|im_start|>assistant\n{text}<|im_end|>\n<|im_start|>assistant\n"
+            additional_information = {
+                "task_type": [task_type],
+                "text": [text],
+                "instruct": [instruct],
+                "language": [language],
+                "speaker": [speaker],
+                "max_new_tokens": [2048],
+            }
             inputs.append(
                 {
-                    "prompt": prompt,
-                    "additional_information": {
-                        "task_type": [task_type],
-                        "text": [text],
-                        "instruct": [instruct],
-                        "language": [language],
-                        "speaker": [speaker],
-                        "max_new_tokens": [2048],
-                    },
+                    "prompt_token_ids": [0] * _estimate_prompt_len(additional_information, model_name),
+                    "additional_information": additional_information,
                 }
             )
     else:
@@ -60,21 +102,21 @@ def get_custom_voice_query(use_batch_sample: bool = False) -> QueryResult:
         language = "Chinese"
         speaker = "Vivian"
         instruct = "用特别愤怒的语气说"
-        prompts = f"<|im_start|>assistant\n{text}<|im_end|>\n<|im_start|>assistant\n"
+        additional_information = {
+            "task_type": [task_type],
+            "text": [text],
+            "language": [language],
+            "speaker": [speaker],
+            "instruct": [instruct],
+            "max_new_tokens": [2048],
+        }
         inputs = {
-            "prompt": prompts,
-            "additional_information": {
-                "task_type": [task_type],
-                "text": [text],
-                "language": [language],
-                "speaker": [speaker],
-                "instruct": [instruct],
-                "max_new_tokens": [2048],
-            },
+            "prompt_token_ids": [0] * _estimate_prompt_len(additional_information, model_name),
+            "additional_information": additional_information,
         }
     return QueryResult(
         inputs=inputs,
-        model_name="Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice",
+        model_name=model_name,
     )
 
 
@@ -88,6 +130,7 @@ def get_voice_design_query(use_batch_sample: bool = False) -> QueryResult:
         QueryResult with Omni inputs and the VoiceDesign model path.
     """
     task_type = "VoiceDesign"
+    model_name = "Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign"
     if use_batch_sample:
         texts = [
             "哥哥，你回来啦，人家等了你好久好久了，要抱抱！",
@@ -100,39 +143,39 @@ def get_voice_design_query(use_batch_sample: bool = False) -> QueryResult:
         languages = ["Chinese", "English"]
         inputs = []
         for text, instruct, language in zip(texts, instructs, languages):
-            prompt = f"<|im_start|>assistant\n{text}<|im_end|>\n<|im_start|>assistant\n"
-            inputs.append(
-                {
-                    "prompt": prompt,
-                    "additional_information": {
-                        "task_type": [task_type],
-                        "text": [text],
-                        "language": [language],
-                        "instruct": [instruct],
-                        "max_new_tokens": [2048],
-                        "non_streaming_mode": [True],
-                    },
-                }
-            )
-    else:
-        text = "哥哥，你回来啦，人家等了你好久好久了，要抱抱！"
-        instruct = "体现撒娇稚嫩的萝莉女声，音调偏高且起伏明显，营造出黏人、做作又刻意卖萌的听觉效果。"
-        language = "Chinese"
-        prompt = f"<|im_start|>assistant\n{text}<|im_end|>\n<|im_start|>assistant\n"
-        inputs = {
-            "prompt": prompt,
-            "additional_information": {
+            additional_information = {
                 "task_type": [task_type],
                 "text": [text],
                 "language": [language],
                 "instruct": [instruct],
                 "max_new_tokens": [2048],
                 "non_streaming_mode": [True],
-            },
+            }
+            inputs.append(
+                {
+                    "prompt_token_ids": [0] * _estimate_prompt_len(additional_information, model_name),
+                    "additional_information": additional_information,
+                }
+            )
+    else:
+        text = "哥哥，你回来啦，人家等了你好久好久了，要抱抱！"
+        instruct = "体现撒娇稚嫩的萝莉女声，音调偏高且起伏明显，营造出黏人、做作又刻意卖萌的听觉效果。"
+        language = "Chinese"
+        additional_information = {
+            "task_type": [task_type],
+            "text": [text],
+            "language": [language],
+            "instruct": [instruct],
+            "max_new_tokens": [2048],
+            "non_streaming_mode": [True],
+        }
+        inputs = {
+            "prompt_token_ids": [0] * _estimate_prompt_len(additional_information, model_name),
+            "additional_information": additional_information,
         }
     return QueryResult(
         inputs=inputs,
-        model_name="Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign",
+        model_name=model_name,
     )
 
 
@@ -147,6 +190,7 @@ def get_base_query(use_batch_sample: bool = False, mode_tag: str = "icl") -> Que
         QueryResult with Omni inputs and the Base model path.
     """
     task_type = "Base"
+    model_name = "Qwen/Qwen3-TTS-12Hz-1.7B-Base"
     ref_audio_path_1 = "https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen3-TTS-Repo/clone_2.wav"
     ref_audio_single = ref_audio_path_1
     ref_text_single = (
@@ -163,38 +207,38 @@ def get_base_query(use_batch_sample: bool = False, mode_tag: str = "icl") -> Que
         syn_lang_batch = ["Chinese", "English"]
         inputs = []
         for text, language in zip(syn_text_batch, syn_lang_batch):
-            prompt = f"<|im_start|>assistant\n{text}<|im_end|>\n<|im_start|>assistant\n"
-            inputs.append(
-                {
-                    "prompt": prompt,
-                    "additional_information": {
-                        "task_type": [task_type],
-                        "ref_audio": [ref_audio_single],
-                        "ref_text": [ref_text_single],
-                        "text": [text],
-                        "language": [language],
-                        "x_vector_only_mode": [x_vector_only_mode],
-                        "max_new_tokens": [2048],
-                    },
-                }
-            )
-    else:
-        prompt = f"<|im_start|>assistant\n{syn_text_single}<|im_end|>\n<|im_start|>assistant\n"
-        inputs = {
-            "prompt": prompt,
-            "additional_information": {
+            additional_information = {
                 "task_type": [task_type],
                 "ref_audio": [ref_audio_single],
                 "ref_text": [ref_text_single],
-                "text": [syn_text_single],
-                "language": [syn_lang_single],
+                "text": [text],
+                "language": [language],
                 "x_vector_only_mode": [x_vector_only_mode],
                 "max_new_tokens": [2048],
-            },
+            }
+            inputs.append(
+                {
+                    "prompt_token_ids": [0] * _estimate_prompt_len(additional_information, model_name),
+                    "additional_information": additional_information,
+                }
+            )
+    else:
+        additional_information = {
+            "task_type": [task_type],
+            "ref_audio": [ref_audio_single],
+            "ref_text": [ref_text_single],
+            "text": [syn_text_single],
+            "language": [syn_lang_single],
+            "x_vector_only_mode": [x_vector_only_mode],
+            "max_new_tokens": [2048],
+        }
+        inputs = {
+            "prompt_token_ids": [0] * _estimate_prompt_len(additional_information, model_name),
+            "additional_information": additional_information,
         }
     return QueryResult(
         inputs=inputs,
-        model_name="Qwen/Qwen3-TTS-12Hz-1.7B-Base",
+        model_name=model_name,
     )
 
 
@@ -223,30 +267,22 @@ def main(args):
         stage_init_timeout=args.stage_init_timeout,
     )
 
-    sampling_params = SamplingParams(
-        temperature=0.9,
-        top_p=1.0,
-        top_k=50,
-        max_tokens=2048,
-        seed=42,
-        detokenize=False,
-        repetition_penalty=1.05,
-    )
-
-    sampling_params_list = [
-        sampling_params,
-    ]
-
     output_dir = args.output_dir if getattr(args, "output_dir", None) else args.output_wav
     os.makedirs(output_dir, exist_ok=True)
 
-    omni_generator = omni.generate(query_result.inputs, sampling_params_list)
+    omni_generator = omni.generate(query_result.inputs, sampling_params_list=None)
     for stage_outputs in omni_generator:
         for output in stage_outputs.request_output:
             request_id = output.request_id
-            audio_tensor = output.outputs[0].multimodal_output["audio"]
+            audio_data = output.outputs[0].multimodal_output["audio"]
+            # async_chunk mode returns a list of chunks; concatenate them.
+            if isinstance(audio_data, list):
+                audio_tensor = torch.cat(audio_data, dim=-1)
+            else:
+                audio_tensor = audio_data
             output_wav = os.path.join(output_dir, f"output_{request_id}.wav")
-            audio_samplerate = output.outputs[0].multimodal_output["sr"].item()
+            sr_val = output.outputs[0].multimodal_output["sr"]
+            audio_samplerate = sr_val.item() if hasattr(sr_val, "item") else int(sr_val[-1])
             # Convert to numpy array and ensure correct format
             audio_numpy = audio_tensor.float().detach().cpu().numpy()
 

examples/online_serving/qwen2_5_omni/gradio_demo.py

@@ -15,38 +15,43 @@ from PIL import Image
 
 SEED = 42
 
+DEFAULT_SAMPLING_PARAMS = {
+    "thinker": {
+        "temperature": 0.0,
+        "top_p": 1.0,
+        "top_k": -1,
+        "max_tokens": 2048,
+        "seed": SEED,
+        "detokenize": True,
+        "repetition_penalty": 1.1,
+    },
+    "talker": {
+        "temperature": 0.9,
+        "top_p": 0.8,
+        "top_k": 40,
+        "max_tokens": 2048,
+        "seed": SEED,
+        "detokenize": True,
+        "repetition_penalty": 1.05,
+        "stop_token_ids": [8294],
+    },
+    "code2wav": {
+        "temperature": 0.0,
+        "top_p": 1.0,
+        "top_k": -1,
+        "max_tokens": 2048,
+        "seed": SEED,
+        "detokenize": True,
+        "repetition_penalty": 1.1,
+    },
+}
+
 SUPPORTED_MODELS: dict[str, dict[str, Any]] = {
+    "Qwen/Qwen2.5-Omni-3B": {
+        "sampling_params": DEFAULT_SAMPLING_PARAMS,
+    },
     "Qwen/Qwen2.5-Omni-7B": {
-        "sampling_params": {
-            "thinker": {
-                "temperature": 0.0,
-                "top_p": 1.0,
-                "top_k": -1,
-                "max_tokens": 2048,
-                "seed": SEED,
-                "detokenize": True,
-                "repetition_penalty": 1.1,
-            },
-            "talker": {
-                "temperature": 0.9,
-                "top_p": 0.8,
-                "top_k": 40,
-                "max_tokens": 2048,
-                "seed": SEED,
-                "detokenize": True,
-                "repetition_penalty": 1.05,
-                "stop_token_ids": [8294],
-            },
-            "code2wav": {
-                "temperature": 0.0,
-                "top_p": 1.0,
-                "top_k": -1,
-                "max_tokens": 2048,
-                "seed": SEED,
-                "detokenize": True,
-                "repetition_penalty": 1.1,
-            },
-        },
+        "sampling_params": DEFAULT_SAMPLING_PARAMS,
     },
 }
 # Ensure deterministic behavior across runs.

tests/entrypoints/openai_api/test_serving_speech.py

@@ -310,10 +310,12 @@ class TestTTSMethods:
         speech_server.engine_client.stage_list = [mock_stage]
         assert speech_server._is_tts_model() is True
 
-    def test_build_tts_prompt(self, speech_server):
-        """Test TTS prompt format."""
-        prompt = speech_server._build_tts_prompt("Hello")
-        assert prompt == "<|im_start|>assistant\nHello<|im_end|>\n<|im_start|>assistant\n"
+    def test_estimate_prompt_len_fallback(self, speech_server):
+        """Test prompt length estimation falls back to 2048 when model is unavailable."""
+        tts_params = {"text": ["Hello"], "task_type": ["CustomVoice"]}
+        result = speech_server._estimate_prompt_len(tts_params)
+        # Without a real model, it should fall back to 2048.
+        assert result == 2048
 
     def test_validate_tts_request_basic(self, speech_server):
         """Test basic validation cases."""

tests/worker/test_omni_gpu_model_runner.py

@@ -69,6 +69,7 @@ def _make_runner(req_ids=("r1", "r2"), hidden_size=4):
     runner.text_step = DummyBuffer(torch.zeros((bsz, hidden_size), dtype=torch.float32))
 
     runner.talker_mtp = DummyTalkerMTP()
+    runner.model = SimpleNamespace(talker_mtp_output_key="code_predictor_codes")
     runner.vllm_config = object()
 
     # Provide a minimal implementation that returns the expected 4-tuple.

vllm_omni/config/model.py

@@ -58,6 +58,7 @@ class OmniModelConfig(ModelConfig):
         }
     )
     omni_kv_config: dict | None = None
+    codec_frame_rate_hz: float | None = None
 
     @property
     def registry(self):
@@ -130,6 +131,21 @@ class OmniModelConfig(ModelConfig):
             video_pruning_rate=video_pruning_rate,
         )
 
+        # Qwen3-TTS: infer codec frame rate from the model config for online serving.
+        if self.codec_frame_rate_hz is None and self.model_arch == "Qwen3TTSTalkerForConditionalGenerationARVLLM":
+            talker_cfg = getattr(self.hf_config, "talker_config", None)
+            if isinstance(talker_cfg, dict):
+                pos_per_sec = talker_cfg.get("position_id_per_seconds")
+            else:
+                pos_per_sec = getattr(talker_cfg, "position_id_per_seconds", None)
+            if pos_per_sec is not None:
+                try:
+                    fps = float(pos_per_sec)
+                except Exception:
+                    fps = None
+                if fps is not None and fps > 0:
+                    self.codec_frame_rate_hz = fps
+
         # Override hf_text_config with omni-specific logic for multi-stage models
         # (e.g., thinker_config, talker_config)
         new_hf_text_config = self.draw_hf_text_config()

vllm_omni/core/sched/omni_generation_scheduler.py

@@ -62,7 +62,6 @@ class OmniGenerationScheduler(VLLMScheduler):
         while req_index < len(self.running) and token_budget > 0:
             request = self.running[req_index]
             # OMNI: Skip requests that are not in self.requests
-            # This can happen when connector marks request as finished and it's removed from requests
             if request.request_id not in self.requests or (
                 self.chunk_transfer_adapter is None and request.status == RequestStatus.FINISHED_STOPPED
             ):
@@ -71,7 +70,25 @@ class OmniGenerationScheduler(VLLMScheduler):
                 continue
 
             num_computed_tokens = request.num_computed_tokens
-            required_tokens = max(len(request.prompt_token_ids) - num_computed_tokens, 1)
+            required_tokens = len(request.prompt_token_ids) - num_computed_tokens
+            # async_chunk: don't schedule placeholder tokens when no new chunk is available.
+            if required_tokens <= 0:
+                if (
+                    self.chunk_transfer_adapter is not None
+                    and request.request_id in self.chunk_transfer_adapter.finished_requests
+                ):
+                    request.status = RequestStatus.FINISHED_STOPPED
+                    # Upstream may finish with no terminal tokens; append one pad token so we can emit FINISHED.
+                    if len(request.prompt_token_ids) <= num_computed_tokens:
+                        request.prompt_token_ids.append(0)
+                        try:
+                            request._all_token_ids.append(0)  # type: ignore[attr-defined]
+                        except Exception:
+                            pass
+                    required_tokens = len(request.prompt_token_ids) - num_computed_tokens
+                else:
+                    req_index += 1
+                    continue
             num_new_tokens = min(required_tokens, token_budget)
             new_blocks = self.kv_cache_manager.allocate_slots(
                 request,
@@ -109,6 +126,20 @@ class OmniGenerationScheduler(VLLMScheduler):
                 self.waiting.pop_request()
                 continue
 
+            # async_chunk: wait for the first upstream chunk (don't start with placeholders).
+            if self.chunk_transfer_adapter is not None and len(request.prompt_token_ids) == 0:
+                if request.request_id in self.chunk_transfer_adapter.finished_requests:
+                    request.status = RequestStatus.FINISHED_STOPPED
+                    request.prompt_token_ids.append(0)
+                    try:
+                        request._all_token_ids.append(0)  # type: ignore[attr-defined]
+                    except Exception:
+                        pass
+                else:
+                    self.waiting.pop_request()
+                    skipped_waiting_requests.prepend_request(request)
+                    continue
+
             # Uniformly treat as diffusion. A feature flag can be added later
             # via config or request tag.
 
@@ -145,11 +176,13 @@ class OmniGenerationScheduler(VLLMScheduler):
 
         # If fast path scheduled none, fall back to the original scheduling
         if not num_scheduled_tokens:
-            res = super().schedule()
             if self.chunk_transfer_adapter:
+                # Don't fall back: base scheduler doesn't handle async_chunk
+                # requests with empty prompt_token_ids.
                 self.chunk_transfer_adapter.restore_queues(self.waiting, self.running)
-                self.chunk_transfer_adapter.postprocess_scheduler_output(res)
-            return res
+            else:
+                res = super().schedule()
+                return res
 
         # Compute common prefix blocks (aligned with v1)
         num_common_prefix_blocks = [0] * len(self.kv_cache_config.kv_cache_groups)
@@ -318,9 +351,8 @@ class OmniGenerationScheduler(VLLMScheduler):
                 continue
             request = self.requests.get(req_id)
             if request is None or request.is_finished():
-                # The request is already finished. This can happen if the
-                # request is aborted while the model is executing it (e.g.,
-                # in pipeline parallelism or async scheduling).
+                # Request may already be finished (e.g., aborted during
+                # execution / pipeline parallelism / async scheduling).
                 continue
 
             req_index = model_runner_output.req_id_to_index[req_id]
@@ -360,8 +392,14 @@ class OmniGenerationScheduler(VLLMScheduler):
             routed_experts = None
 
             # Diffusion request: completes in one step; mark finished and free resources
-            if request.status == RequestStatus.FINISHED_STOPPED or (
-                self.chunk_transfer_adapter is None and request.num_computed_tokens >= request.num_prompt_tokens
+            if (
+                request.status == RequestStatus.FINISHED_STOPPED
+                or (self.chunk_transfer_adapter is None and request.num_computed_tokens >= request.num_prompt_tokens)
+                or (
+                    self.chunk_transfer_adapter is not None
+                    and request.request_id in self.chunk_transfer_adapter.finished_requests
+                    and request.num_computed_tokens >= len(request.prompt_token_ids)
+                )
             ):
                 request.status = RequestStatus.FINISHED_STOPPED
                 # Optional: set a stop_reason for front-end clarity
@@ -375,15 +413,11 @@ class OmniGenerationScheduler(VLLMScheduler):
                 finished = self._handle_stopped_request(request)
                 if finished:
                     kv_transfer_params = self._free_request(request)
-                if status_before_stop == RequestStatus.RUNNING:
-                    stopped_running_reqs.add(request)
-                elif status_before_stop == RequestStatus.WAITING_FOR_CHUNK:
-                    # In async chunk mode, request may be in either queue.
-                    # Remove from both to avoid stale queue entries.
+                if status_before_stop == RequestStatus.WAITING_FOR_CHUNK:
                     stopped_running_reqs.add(request)
                     stopped_preempted_reqs.add(request)
                 else:
-                    stopped_preempted_reqs.add(request)
+                    stopped_running_reqs.add(request)
 
             # Extract sample logprobs if needed.
             if request.sampling_params is not None and request.sampling_params.logprobs is not None and logprobs:

vllm_omni/diffusion/models/hunyuan_image_3/hunyuan_image_3_transformer.py

@@ -64,6 +64,7 @@ from vllm.v1.attention.backend import AttentionType
 
 from vllm_omni.diffusion.attention.layer import Attention
 from vllm_omni.diffusion.distributed.parallel_state import get_pp_group
+from vllm_omni.diffusion.distributed.utils import get_local_device
 from vllm_omni.diffusion.layers.rope import RotaryEmbedding
 
 logger = logging.getLogger(__name__)
@@ -1737,6 +1738,7 @@ class HunyuanImage3Model(nn.Module):
         lora_config = None
         self.num_redundant_experts = 0
         self.config = config
+        self.device = get_local_device()
         self.quant_config = quant_config
         self.padding_idx = config.pad_token_id
         lora_vocab = (lora_config.lora_extra_vocab_size * (lora_config.max_loras or 1)) if lora_config else 0
@@ -2430,7 +2432,7 @@ class HunyuanImage3Text2ImagePipeline(DiffusionPipeline):
                     **model_kwargs,
                 )
 
-                with torch.autocast(device_type="cuda", dtype=torch.bfloat16, enabled=True):
+                with torch.autocast(device_type=self.device.type, dtype=torch.bfloat16, enabled=True):
                     model_output = self.model.forward_call(**model_inputs, first_step=(i == 0))
                     pred = model_output["diffusion_prediction"]
                 pred = pred.to(dtype=torch.float32)
@@ -2477,7 +2479,7 @@ class HunyuanImage3Text2ImagePipeline(DiffusionPipeline):
         if hasattr(self.vae, "ffactor_temporal"):
             latents = latents.unsqueeze(2)
 
-        with torch.autocast(device_type="cuda", dtype=torch.float16, enabled=True):
+        with torch.autocast(device_type=self.device.type, dtype=torch.float16, enabled=True):
             image = self.vae.decode(latents, return_dict=False, generator=generator)[0]
 
         if hasattr(self.vae, "ffactor_temporal"):

vllm_omni/diffusion/models/hunyuan_image_3/pipeline_hunyuan_image_3.py

@@ -140,7 +140,7 @@ class HunyuanImage3Pipeline(HunyuanImage3PreTrainedModel, GenerationMixin):
             if hasattr(self, prefix.split(".")[0]):
                 module = dict(self.named_modules()).get(prefix)
                 if module:
-                    module.to(f"cuda:{tp_rank}")
+                    module.to(f"{self.model.device.type}:{tp_rank}")
 
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
         self.pre_load()
@@ -369,7 +369,7 @@ class HunyuanImage3Pipeline(HunyuanImage3PreTrainedModel, GenerationMixin):
     def vae_encode(self, image, cfg_factor=1):
         config = self.vae.config
 
-        with torch.autocast(device_type="cuda", dtype=torch.float16, enabled=True):
+        with torch.autocast(device_type=self.model.device.type, dtype=torch.float16, enabled=True):
             vae_encode_result = self.vae.encode(image)
             if isinstance(vae_encode_result, torch.Tensor):
                 latents = vae_encode_result

vllm_omni/diffusion/models/z_image/pipeline_z_image.py

@@ -564,14 +564,14 @@ class ZImagePipeline(nn.Module):
 
             # Run CFG only if configured AND scale is non-zero
             apply_cfg = self.do_classifier_free_guidance and current_guidance_scale > 0
+            latents_typed = latents.to(self.od_config.dtype)
 
             if apply_cfg:
-                latents_typed = latents.to(self.transformer.dtype)
                 latent_model_input = latents_typed.repeat(2, 1, 1, 1)
                 prompt_embeds_model_input = prompt_embeds + negative_prompt_embeds
                 timestep_model_input = timestep.repeat(2)
             else:
-                latent_model_input = latents.to(self.transformer.dtype)
+                latent_model_input = latents_typed
                 prompt_embeds_model_input = prompt_embeds
                 timestep_model_input = timestep
 

vllm_omni/diffusion/models/z_image/z_image_transformer.py

@@ -614,7 +614,6 @@ class ZImageTransformer2DModel(CachedTransformer):
         quant_config: "QuantizationConfig | None" = None,
     ) -> None:
         super().__init__()
-        self.dtype = torch.bfloat16
         self.in_channels = in_channels
         self.out_channels = in_channels
         self.all_patch_size = all_patch_size

vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py

@@ -179,11 +179,17 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase):
             else:
                 if payload_data.get("finished"):
                     self.finished_requests.add(req_id)
-                    req.status = RequestStatus.FINISHED_STOPPED
 
-                req.prompt_token_ids = payload_data.get("code_predictor_codes", [])
+                # req.prompt_token_ids = payload_data.get("code_predictor_codes", [])
+                # req.num_computed_tokens = 0
+                new_ids = payload_data.get("code_predictor_codes", [])
+                req.prompt_token_ids = new_ids
                 req.num_computed_tokens = 0
 
+                # Empty chunk with more data expected: keep polling.
+                if not new_ids and not payload_data.get("finished"):
+                    return
+
             # Mark as finished for consumption
             with self.lock:
                 self._finished_load_reqs.add(req_id)
@@ -308,7 +314,6 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase):
                     # of schedule, but have not scheduled
                     continue
                 if request.request_id in self.finished_requests:
-                    request.additional_information = {}
                     continue
                 # Requests that waiting for chunk
                 self.load_async(request)

vllm_omni/entrypoints/async_omni.py

@@ -390,7 +390,8 @@ class AsyncOmni(OmniBase):
                     submit_flag = False
                     prompt_token_ids = engine_outputs.prompt_token_ids
                     engine_input = copy.deepcopy(prompt)
-                    engine_input["prompt_token_ids"] = [0] * compute_talker_prompt_ids_length(prompt_token_ids)
+                    next_prompt_len = max(1, compute_talker_prompt_ids_length(prompt_token_ids))
+                    engine_input["prompt_token_ids"] = [0] * next_prompt_len
                     engine_input["multi_modal_data"] = engine_input["mm_processor_kwargs"] = None
                     for i in range(1, len(self.stage_list)):
                         task = {

vllm_omni/entrypoints/openai/serving_speech.py

@@ -1,6 +1,14 @@
 import asyncio
+import base64
+import io
+import ipaddress
+import socket
 from typing import Any
+from urllib.parse import urlparse
+from urllib.request import urlopen
 
+import numpy as np
+import soundfile as sf
 from fastapi import Request
 from fastapi.responses import Response
 from vllm.entrypoints.openai.engine.serving import OpenAIServing
@@ -17,6 +25,19 @@ from vllm_omni.outputs import OmniRequestOutput
 
 logger = init_logger(__name__)
 
+_REF_AUDIO_TIMEOUT_S = 15
+_REF_AUDIO_MAX_BYTES = 50 * 1024 * 1024  # 50 MB
+_REF_AUDIO_BLOCKED_NETWORKS = [
+    ipaddress.ip_network("127.0.0.0/8"),
+    ipaddress.ip_network("10.0.0.0/8"),
+    ipaddress.ip_network("172.16.0.0/12"),
+    ipaddress.ip_network("192.168.0.0/16"),
+    ipaddress.ip_network("169.254.0.0/16"),
+    ipaddress.ip_network("::1/128"),
+    ipaddress.ip_network("fc00::/7"),
+    ipaddress.ip_network("fe80::/10"),
+]
+
 # TTS Configuration (currently supports Qwen3-TTS)
 _TTS_MODEL_STAGES: set[str] = {"qwen3_tts"}
 _TTS_LANGUAGES: set[str] = {
@@ -43,6 +64,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
         # Load supported speakers
         self.supported_speakers = self._load_supported_speakers()
         logger.info(f"Loaded {len(self.supported_speakers)} supported speakers: {sorted(self.supported_speakers)}")
+        self._tts_tokenizer = None
 
     def _load_supported_speakers(self) -> set[str]:
         """Load supported speakers (case-insensitive) from the model configuration."""
@@ -62,6 +84,36 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
 
         return set()
 
+    def _estimate_prompt_len(self, tts_params: dict[str, Any]) -> int:
+        """Estimate prompt length so the placeholder matches model-side embeddings."""
+        try:
+            from vllm_omni.model_executor.models.qwen3_tts.qwen3_tts_talker import (
+                Qwen3TTSTalkerForConditionalGeneration,
+            )
+
+            if self._tts_tokenizer is None:
+                from transformers import AutoTokenizer
+
+                model_name = self.engine_client.model_config.model
+                self._tts_tokenizer = AutoTokenizer.from_pretrained(
+                    model_name,
+                    trust_remote_code=True,
+                    padding_side="left",
+                )
+            hf_config = self.engine_client.model_config.hf_config
+            talker_config = hf_config.talker_config
+            task_type = (tts_params.get("task_type") or ["CustomVoice"])[0]
+            return Qwen3TTSTalkerForConditionalGeneration.estimate_prompt_len_from_additional_information(
+                additional_information=tts_params,
+                task_type=task_type,
+                tokenize_prompt=lambda t: self._tts_tokenizer(t, padding=False)["input_ids"],
+                codec_language_id=getattr(talker_config, "codec_language_id", None),
+                spk_is_dialect=getattr(talker_config, "spk_is_dialect", None),
+            )
+        except Exception as e:
+            logger.warning("Failed to estimate TTS prompt length, using fallback 2048: %s", e)
+            return 2048
+
     def _is_tts_model(self) -> bool:
         """Check if the current model is a supported TTS model."""
         stage_list = getattr(self.engine_client, "stage_list", None)
@@ -125,9 +177,44 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
 
         return None
 
-    def _build_tts_prompt(self, text: str) -> str:
-        """Build TTS prompt from input text."""
-        return f"<|im_start|>assistant\n{text}<|im_end|>\n<|im_start|>assistant\n"
+    @staticmethod
+    async def _resolve_ref_audio(ref_audio_str: str) -> tuple[list[float], int]:
+        """Resolve ref_audio URL/base64 to (wav_samples, sample_rate)."""
+        parsed = urlparse(ref_audio_str)
+
+        def _check_ssrf(url: str) -> None:
+            host = urlparse(url).hostname
+            if not host:
+                raise ValueError("ref_audio URL must include a hostname")
+            for info in socket.getaddrinfo(host, None):
+                ip_str = str(info[4][0]).split("%", 1)[0]
+                addr = ipaddress.ip_address(ip_str)
+                if any(addr in net for net in _REF_AUDIO_BLOCKED_NETWORKS):
+                    raise ValueError(f"ref_audio URL resolves to blocked address: {addr}")
+
+        def _fetch_sync() -> tuple[np.ndarray, int]:
+            if parsed.scheme in ("http", "https"):
+                _check_ssrf(ref_audio_str)
+                with urlopen(ref_audio_str, timeout=_REF_AUDIO_TIMEOUT_S) as resp:
+                    data = resp.read(_REF_AUDIO_MAX_BYTES + 1)
+                    if len(data) > _REF_AUDIO_MAX_BYTES:
+                        raise ValueError(f"ref_audio URL exceeds {_REF_AUDIO_MAX_BYTES} bytes")
+                buf = io.BytesIO(data)
+            elif ref_audio_str.startswith("data:"):
+                b64 = ref_audio_str
+                if "," in b64:
+                    b64 = b64.split(",", 1)[1]
+                buf = io.BytesIO(base64.b64decode(b64))
+            else:
+                raise ValueError("ref_audio must be an http(s) URL or data: base64 URI")
+            audio, sr = sf.read(buf, dtype="float32", always_2d=False)
+            if isinstance(audio, np.ndarray) and audio.ndim > 1:
+                audio = np.mean(audio, axis=-1)
+            return np.asarray(audio, dtype=np.float32), int(sr)
+
+        loop = asyncio.get_running_loop()
+        wav_np, sr = await loop.run_in_executor(None, _fetch_sync)
+        return wav_np.tolist(), sr
 
     def _build_tts_params(self, request: OpenAICreateSpeechRequest) -> dict[str, Any]:
         """Build TTS parameters from request.
@@ -164,9 +251,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
         else:
             params["instruct"] = [""]
 
-        # Voice clone parameters (used with Base task)
-        if request.ref_audio is not None:
-            params["ref_audio"] = [request.ref_audio]
+        # Voice clone: ref_audio resolved in create_speech(), not here.
         if request.ref_text is not None:
             params["ref_text"] = [request.ref_text]
         if request.x_vector_only_mode is not None:
@@ -221,11 +306,19 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
                 if validation_error:
                     return self.create_error_response(validation_error)
 
-                # Build TTS parameters and prompt
+                # Must use prompt_token_ids (not text prompt): the AR Talker
+                # operates on codec tokens; text token IDs exceed codec vocab.
+                # model.preprocess replaces all embeddings, so placeholder value
+                # is irrelevant -- but length must match to avoid excess padding.
                 tts_params = self._build_tts_params(request)
-                prompt_text = self._build_tts_prompt(request.input)
+
+                if request.ref_audio is not None:
+                    wav_list, sr = await self._resolve_ref_audio(request.ref_audio)
+                    tts_params["ref_audio"] = [[wav_list, sr]]
+
+                ph_len = self._estimate_prompt_len(tts_params)
                 prompt = {
-                    "prompt": prompt_text,
+                    "prompt_token_ids": [1] * ph_len,
                     "additional_information": tts_params,
                 }
             else:
@@ -282,6 +375,11 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
             if hasattr(sample_rate, "item"):
                 sample_rate = sample_rate.item()
 
+            # Streaming accumulates chunks as a list; concat first.
+            if isinstance(audio_tensor, list):
+                import torch
+
+                audio_tensor = torch.cat(audio_tensor, dim=-1)
             # Convert tensor to numpy
             if hasattr(audio_tensor, "float"):
                 audio_tensor = audio_tensor.float().detach().cpu().numpy()

vllm_omni/model_executor/models/qwen3_tts/configuration_qwen3_tts.py

@@ -504,19 +504,39 @@ class Qwen3TTSConfig(PretrainedConfig):
         self.tts_bos_token_id = tts_bos_token_id
         self.tts_eos_token_id = tts_eos_token_id
 
-        # TODO: remove these dummy values after
-        self.image_token_id = 0  # dummy image token id
-        self.video_token_id = 0  # dummy video token id
-        self.vision_start_token_id = 0  # dummy vision start token id
+        # Dummy vision token IDs that must never collide with real codec tokens.
+        # mrope scans prompt_token_ids for these; using -1 ensures no false match.
+        self.image_token_id = -1
+        self.video_token_id = -1
+        self.vision_start_token_id = -1
         self.vision_config = PretrainedConfig()  # dummy vision config
         self.vision_config.spatial_merge_size = 1
 
+    @property
+    def codec_frame_rate_hz(self) -> float | None:
+        pos_per_sec = getattr(self.talker_config, "position_id_per_seconds", None)
+        if pos_per_sec is None:
+            return None
+        try:
+            fps = float(pos_per_sec)
+        except (TypeError, ValueError):
+            return None
+        return fps if fps > 0 else None
+
     def get_text_config(self, **kwargs):
         # vLLM expects text config to expose hidden_size/num_attention_heads.
         # For Qwen3 TTS, the talker config is the text model config.
         config = self.talker_config
-        # if hasattr(config, "rope_parameters"):
-        #     delattr(config, "rope_parameters")
+        # Code2Wav is a pure convolutional waveform decoder; it does NOT use
+        # rotary position embeddings.  When hf_overrides sets architectures
+        # to [Qwen3TTSCode2Wav], strip rope_parameters so that the model
+        # runner sees uses_mrope == False and skips mrope position computation
+        # on codec tokens.  Each stage loads its own config instance, so this
+        # in-place mutation does not affect the Talker stage.
+        archs = getattr(self, "architectures", []) or []
+        if any("Code2Wav" in str(a) for a in archs):
+            if hasattr(config, "rope_parameters"):
+                delattr(config, "rope_parameters")
         return config
 
 

vllm_omni/model_executor/models/qwen3_tts/processing_qwen3_tts.py

@@ -1,102 +0,0 @@
-# Copyright 2026 The Qwen team, Alibaba Group and the HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-from transformers.feature_extraction_utils import BatchFeature
-from transformers.processing_utils import ProcessingKwargs, ProcessorMixin
-
-
-class Qwen3TTSProcessorKwargs(ProcessingKwargs, total=False):
-    _defaults = {
-        "text_kwargs": {
-            "padding": False,
-            "padding_side": "left",
-        }
-    }
-
-
-class Qwen3TTSProcessor(ProcessorMixin):
-    r"""
-    Constructs a Qwen3TTS processor.
-
-    Args:
-        tokenizer ([`Qwen2TokenizerFast`], *optional*):
-            The text tokenizer.
-        chat_template (`Optional[str]`, *optional*):
-            The Jinja template to use for formatting the conversation.
-            If not provided, the default chat template is used.
-    """
-
-    attributes = ["tokenizer"]
-    tokenizer_class = ("Qwen2Tokenizer", "Qwen2TokenizerFast")
-
-    def __init__(self, tokenizer=None, chat_template=None):
-        super().__init__(tokenizer, chat_template=chat_template)
-
-    def __call__(self, text=None, **kwargs) -> BatchFeature:
-        """
-        Main method to prepare for the model one or several sequences(s) and audio(s).
-        This method forwards the `text` and `kwargs` arguments to
-        Qwen2TokenizerFast's [`~Qwen2TokenizerFast.__call__`] if `text` is not `None`
-        to encode the text.
-
-        Args:
-            text (`str`, `List[str]`, `List[List[str]]`):
-                The sequence or batch of sequences to be encoded. Each sequence can be a string or a list of strings
-                (pretokenized string). If the sequences are provided as list of strings (pretokenized), you must set
-                `is_split_into_words=True` (to lift the ambiguity with a batch of sequences).
-        """
-
-        if text is None:
-            raise ValueError("You need to specify either a `text` input to process.")
-
-        output_kwargs = self._merge_kwargs(
-            Qwen3TTSProcessorKwargs,
-            tokenizer_init_kwargs=self.tokenizer.init_kwargs,
-            **kwargs,
-        )
-        if not isinstance(text, list):
-            text = [text]
-
-        texts_inputs = self.tokenizer(text, **output_kwargs["text_kwargs"])
-
-        return BatchFeature(
-            data={**texts_inputs},
-            tensor_type=kwargs.get("return_tensors"),
-        )
-
-    def batch_decode(self, *args, **kwargs):
-        """
-        This method forwards all its arguments to Qwen2TokenizerFast's [`~PreTrainedTokenizer.batch_decode`]. Please
-        refer to the docstring of this method for more information.
-        """
-        return self.tokenizer.batch_decode(*args, **kwargs)
-
-    def decode(self, *args, **kwargs):
-        """
-        This method forwards all its arguments to Qwen2TokenizerFast's [`~PreTrainedTokenizer.decode`]. Please refer to
-        the docstring of this method for more information.
-        """
-        return self.tokenizer.decode(*args, **kwargs)
-
-    def apply_chat_template(self, conversations, chat_template=None, **kwargs):
-        if isinstance(conversations[0], dict):
-            conversations = [conversations]
-        return super().apply_chat_template(conversations, chat_template, **kwargs)
-
-    @property
-    def model_input_names(self):
-        tokenizer_input_names = self.tokenizer.model_input_names
-        return list(dict.fromkeys(tokenizer_input_names))
-
-
-__all__ = ["Qwen3TTSProcessor"]

vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_code2wav.py

@@ -0,0 +1,239 @@
+from __future__ import annotations
+
+import os
+from collections.abc import Iterable
+from typing import Any
+
+import numpy as np
+import torch
+import torch.nn as nn
+from transformers.utils.hub import cached_file
+from vllm.config import VllmConfig
+from vllm.logger import init_logger
+
+from vllm_omni.model_executor.models.output_templates import OmniOutput
+
+from .qwen3_tts_tokenizer import Qwen3TTSTokenizer
+
+logger = init_logger(__name__)
+
+
+class Qwen3TTSCode2Wav(nn.Module):
+    """Stage-1 code2wav model for Qwen3-TTS (GenerationModelRunner).
+    Consumes frame-aligned codec tokens from input_ids and decodes waveform via SpeechTokenizer."""
+
+    input_modalities = "audio"
+
+    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
+        super().__init__()
+        self.vllm_config = vllm_config
+        self.model_path = vllm_config.model_config.model
+
+        self.have_multimodal_outputs = True
+        self.has_preprocess = False
+        self.has_postprocess = False
+        # Generation-only stage (no logits / sampling).
+        self.requires_raw_input_tokens = True
+
+        self._speech_tokenizer: Qwen3TTSTokenizer | None = None
+        self._num_quantizers: int | None = None
+        self._decode_upsample_rate: int | None = None
+        self._output_sample_rate: int | None = None
+        self._logged_codec_stats = False
+
+    @staticmethod
+    def _module_device(module: nn.Module) -> torch.device:
+        try:
+            return next(module.parameters()).device
+        except StopIteration:
+            for _, buf in module.named_buffers(recurse=True):
+                return buf.device
+            return torch.device("cpu")
+
+    def _ensure_speech_tokenizer_loaded(self) -> Qwen3TTSTokenizer:
+        if self._speech_tokenizer is not None:
+            return self._speech_tokenizer
+
+        # Locate speech_tokenizer dir from HF cache (or local path).
+        cfg_path = cached_file(self.model_path, "speech_tokenizer/config.json")
+        if cfg_path is None:
+            raise ValueError(f"{self.model_path}/speech_tokenizer/config.json not found")
+        speech_tokenizer_dir = os.path.dirname(cfg_path)
+
+        # Stage-1 only needs decode; skip HF feature extractor to avoid heavy optional deps.
+        # Still require preprocessor_config.json (use cached_file so online runs can fetch it).
+        prep_cfg = cached_file(self.model_path, "speech_tokenizer/preprocessor_config.json")
+        if prep_cfg is None:
+            raise ValueError(
+                f"{self.model_path}/speech_tokenizer/preprocessor_config.json not found. "
+                "Please make sure the checkpoint contains the required HF preprocessing files."
+            )
+
+        tok = Qwen3TTSTokenizer.from_pretrained(
+            speech_tokenizer_dir,
+            torch_dtype=torch.bfloat16,
+            load_feature_extractor=False,
+        )
+
+        # Align device with vLLM worker, then read back from module.
+        if tok.model is not None:
+            tok.model.to(device=self.vllm_config.device_config.device)
+            tok.device = self._module_device(tok.model)
+
+        # Derive codec group count and rates from tokenizer config.
+        dec_cfg = getattr(tok.model.config, "decoder_config", None)
+        num_q = getattr(dec_cfg, "num_quantizers", None) if dec_cfg is not None else None
+        if num_q is None:
+            raise ValueError("speech_tokenizer decoder_config.num_quantizers not found")
+        num_q = int(num_q)
+        if num_q <= 0:
+            raise ValueError(f"Invalid speech_tokenizer num_quantizers={num_q}")
+
+        try:
+            upsample = int(tok.get_decode_upsample_rate())
+        except Exception as e:
+            raise ValueError(f"Failed to get decode upsample rate: {e}") from e
+        if upsample <= 0:
+            raise ValueError(f"Invalid decode upsample rate: {upsample}")
+
+        try:
+            out_sr = int(tok.get_output_sample_rate())
+        except Exception as e:
+            raise ValueError(f"Failed to get output sample rate: {e}") from e
+
+        self._speech_tokenizer = tok
+        self._num_quantizers = num_q
+        self._decode_upsample_rate = upsample
+        self._output_sample_rate = out_sr
+        return tok
+
+    def embed_input_ids(self, input_ids: torch.Tensor, **_: Any) -> torch.Tensor:
+        # This stage ignores token embeddings. Keep a stable dummy embedding for vLLM runner.
+        if input_ids.numel() == 0:
+            return torch.empty((0, 1), device=input_ids.device, dtype=torch.float32)
+        return torch.zeros((input_ids.shape[0], 1), device=input_ids.device, dtype=torch.float32)
+
+    def compute_logits(self, hidden_states: torch.Tensor | OmniOutput, sampling_metadata: Any = None) -> None:
+        return None
+
+    @torch.no_grad()
+    def forward(
+        self,
+        input_ids: torch.Tensor | None = None,
+        positions: torch.Tensor | None = None,
+        intermediate_tensors: Any = None,
+        inputs_embeds: torch.Tensor | None = None,
+        **kwargs: Any,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """Decode codec codes into audio waveform.
+
+        input_ids layout: [codec_context_frames, *flat_codes]
+        where flat_codes is codebook-major [q*F].
+        """
+        tok = self._ensure_speech_tokenizer_loaded()
+        assert self._num_quantizers is not None
+        assert self._output_sample_rate is not None
+
+        sr_val = self._output_sample_rate
+        empty_ret = (
+            torch.zeros((0,), dtype=torch.float32),
+            torch.tensor(sr_val, dtype=torch.int32),
+        )
+
+        if input_ids is None:
+            return empty_ret
+
+        q = int(self._num_quantizers)
+        ids = input_ids.reshape(-1).to(dtype=torch.long)
+        n_tokens = ids.numel()
+
+        if n_tokens == 0:
+            return empty_ret
+
+        # input_ids[0] = codec_context_frames (prepended by stage_input_processor).
+        ctx_frames = int(ids[0].item())
+        ids = ids[1:]
+        n_tokens = ids.numel()
+
+        if n_tokens == 0:
+            return empty_ret
+
+        # Warmup / dummy_run: not divisible by num_quantizers.
+        if n_tokens % q != 0:
+            logger.warning(
+                "Code2Wav input_ids length %d not divisible by num_quantizers %d, "
+                "likely a warmup run; returning empty audio.",
+                n_tokens,
+                q,
+            )
+            return empty_ret
+
+        total_frames = n_tokens // q
+
+        # Reshape codebook-major flat [q*F] -> [q, F] -> [F, q] for SpeechTokenizer.
+        codes_fq = ids.reshape(q, total_frames).transpose(0, 1).contiguous()
+
+        if not self._logged_codec_stats and total_frames > 1:
+            self._logged_codec_stats = True
+            try:
+                uniq = int(torch.unique(codes_fq).numel())
+                cmin = int(codes_fq.min().item())
+                cmax = int(codes_fq.max().item())
+                head = codes_fq[: min(2, total_frames), : min(8, q)].cpu().tolist()
+                logger.info(
+                    "Code2Wav codec: frames=%d q=%d uniq=%d range=[%d,%d] head=%s",
+                    total_frames,
+                    q,
+                    uniq,
+                    cmin,
+                    cmax,
+                    head,
+                )
+            except Exception:
+                pass
+
+        wavs, sr = tok.decode({"audio_codes": codes_fq})
+        if not wavs:
+            raise ValueError("SpeechTokenizer code2wav produced empty waveform list.")
+        audio_np = wavs[0].astype(np.float32, copy=False)
+
+        # Trim left-context waveform samples (streaming sliding window).
+        if ctx_frames > 0:
+            upsample = self._decode_upsample_rate
+            if upsample is None or upsample <= 0:
+                raise ValueError(f"Invalid decode upsample rate: {upsample}")
+            cut = ctx_frames * upsample
+            if cut < audio_np.shape[0]:
+                audio_np = audio_np[cut:]
+            else:
+                logger.warning(
+                    "Context trim %d >= decoded length %d; returning empty audio.",
+                    cut,
+                    audio_np.shape[0],
+                )
+                return empty_ret
+
+        audio_tensor = torch.from_numpy(audio_np).to(dtype=torch.float32).reshape(-1)
+        sr_tensor = torch.tensor(int(sr), dtype=torch.int32)
+        return audio_tensor, sr_tensor
+
+    def make_omni_output(self, model_outputs: torch.Tensor | OmniOutput, **kwargs: Any) -> OmniOutput:
+        if isinstance(model_outputs, OmniOutput):
+            return model_outputs
+
+        if not (isinstance(model_outputs, tuple) and len(model_outputs) == 2):
+            raise TypeError(f"Qwen3TTSCode2Wav expected (audio_tensor, sr) outputs, got {type(model_outputs)}")
+
+        audio_tensor, sr = model_outputs
+        return OmniOutput(
+            text_hidden_states=None,
+            multimodal_outputs={
+                "model_outputs": audio_tensor,
+                "sr": sr,
+            },
+        )
+
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+        # SpeechTokenizer weights live under `speech_tokenizer/` and are loaded
+        # lazily from that directory. Ignore main checkpoint weights.
+        return set()

vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_code_predictor_vllm.py

@@ -0,0 +1,495 @@
+from __future__ import annotations
+
+from collections.abc import Iterable
+from typing import Any
+
+import torch
+import torch.nn as nn
+from vllm.config import VllmConfig
+from vllm.config.vllm import set_current_vllm_config
+from vllm.forward_context import set_forward_context
+from vllm.model_executor.layers.layernorm import RMSNorm
+from vllm.model_executor.model_loader.weight_utils import (
+    default_weight_loader,
+    maybe_remap_kv_scale_name,
+)
+from vllm.model_executor.models.qwen3 import Qwen3DecoderLayer
+from vllm.model_executor.models.utils import is_pp_missing_parameter
+from vllm.v1.kv_cache_interface import KVCacheConfig, KVCacheGroupSpec, KVCacheSpec, KVCacheTensor
+from vllm.v1.worker.gpu import attn_utils
+
+from .configuration_qwen3_tts import Qwen3TTSTalkerCodePredictorConfig, Qwen3TTSTalkerConfig
+
+
+class _LocalPredictorKVCache:
+    """Minimal local KV cache + attention metadata for running
+    code_predictor inside one worker (independent of engine KV)."""
+
+    def __init__(
+        self,
+        *,
+        vllm_config: VllmConfig,
+        max_seq_len: int,
+        max_batch_size: int,
+        device: torch.device,
+    ) -> None:
+        self.vllm_config = vllm_config
+        self.device = device
+
+        # Collect attention layers registered in this vllm_config.
+        kv_cache_spec_by_layer = attn_utils.get_kv_cache_spec(vllm_config)
+        if not kv_cache_spec_by_layer:
+            raise RuntimeError("Local predictor KVCache requires vLLM Attention layers to be registered.")
+
+        # We only need enough blocks for a tiny per-frame sequence (<= max_seq_len).
+        any_spec = next(iter(kv_cache_spec_by_layer.values()))
+        block_size = int(any_spec.block_size)
+        blocks_per_seq = (int(max_seq_len) + block_size - 1) // block_size
+        num_blocks = max(1, int(max_batch_size) * int(blocks_per_seq))
+
+        # Allocate per-layer KV caches (small, independent).
+        kv_cache_tensors: list[KVCacheTensor] = []
+        for layer_name, spec in kv_cache_spec_by_layer.items():
+            kv_cache_tensors.append(KVCacheTensor(size=int(spec.page_size_bytes) * num_blocks, shared_by=[layer_name]))
+
+        merged_spec: KVCacheSpec = KVCacheSpec.merge(list(kv_cache_spec_by_layer.values()))
+        self.kv_cache_config = KVCacheConfig(
+            num_blocks=num_blocks,
+            kv_cache_tensors=kv_cache_tensors,
+            kv_cache_groups=[
+                KVCacheGroupSpec(layer_names=list(kv_cache_spec_by_layer.keys()), kv_cache_spec=merged_spec)
+            ],
+        )
+
+        # Init backend + bind KV cache tensors to attention modules.
+        self.attn_backends, self.attn_metadata_builders = attn_utils.init_attn_backend(
+            self.kv_cache_config, vllm_config, device
+        )
+        self.runner_kv_caches: list[torch.Tensor] = []
+        attn_utils.init_kv_cache(
+            self.runner_kv_caches,
+            vllm_config.compilation_config.static_forward_context,
+            self.kv_cache_config,
+            self.attn_backends,
+            device,
+        )
+
+        # Precompute a fixed block table mapping for the maximum batch.
+        self.block_size = block_size
+        self.blocks_per_seq = blocks_per_seq
+        self.max_batch_size = int(max_batch_size)
+
+        bt = torch.full((self.max_batch_size, self.blocks_per_seq), -1, dtype=torch.int32, device=device)
+        for i in range(self.max_batch_size):
+            for j in range(self.blocks_per_seq):
+                bt[i, j] = i * self.blocks_per_seq + j
+        self._block_table = bt
+
+    def build_attn_metadata(
+        self,
+        *,
+        num_reqs: int,
+        query_lens: torch.Tensor,  # (num_reqs,) int32 on cpu
+        seq_lens: torch.Tensor,  # (num_reqs,) int32 on cpu
+    ) -> tuple[dict[str, Any], torch.Tensor, dict[str, torch.Tensor]]:
+        """Build attention metadata, positions, and slot_mapping dict.
+
+        Returns:
+            (attn_metadata, positions, slot_mappings_by_layer)
+            - attn_metadata: per-layer attention metadata for attn backends.
+            - positions: (num_tokens,) position IDs on device.
+            - slot_mappings_by_layer: {layer_name: slot_mapping_tensor} for
+              set_forward_context so that unified_kv_cache_update can write
+              the KV cache correctly.
+        """
+        num_reqs = int(num_reqs)
+        if num_reqs <= 0:
+            return {}, torch.empty((0,), dtype=torch.int64, device=self.device), {}
+        if num_reqs > self.max_batch_size:
+            raise ValueError(f"num_reqs={num_reqs} exceeds local predictor max_batch_size={self.max_batch_size}")
+
+        query_lens_i32 = query_lens.to(dtype=torch.int32, device="cpu")
+        seq_lens_i32 = seq_lens.to(dtype=torch.int32, device="cpu")
+
+        # query_start_loc: prefix sums of query_lens.
+        qsl = torch.zeros((num_reqs + 1,), dtype=torch.int32, device="cpu")
+        qsl[1:] = torch.cumsum(query_lens_i32, dim=0)
+        num_tokens = int(qsl[-1].item())
+        if num_tokens <= 0:
+            return {}, torch.empty((0,), dtype=torch.int64, device=self.device), {}
+
+        # positions: for each request i, emit positions [seq_len-query_len .. seq_len-1]
+        pos_list: list[torch.Tensor] = []
+        for i in range(num_reqs):
+            ql = int(query_lens_i32[i].item())
+            sl = int(seq_lens_i32[i].item())
+            start = sl - ql
+            pos_list.append(torch.arange(start, sl, dtype=torch.int64))
+        positions_cpu = torch.cat(pos_list, dim=0)
+
+        # slot_mapping: map each query token to a physical slot in the paged KV cache.
+        # We allocate per-request contiguous blocks; slot = base + position.
+        slot_mapping = torch.empty((num_tokens,), dtype=torch.int64, device="cpu")
+        cursor = 0
+        for i in range(num_reqs):
+            ql = int(query_lens_i32[i].item())
+            sl = int(seq_lens_i32[i].item())
+            start = sl - ql
+            for p in range(start, sl):
+                block_idx = p // self.block_size
+                offset = p % self.block_size
+                block_id = int(self._block_table[i, block_idx].item())
+                slot_mapping[cursor] = block_id * self.block_size + offset
+                cursor += 1
+
+        max_seq_len = int(seq_lens_i32[:num_reqs].max().item())
+        query_start_loc_gpu = qsl.to(device=self.device)
+        seq_lens_gpu = seq_lens_i32.to(device=self.device)
+        block_table = self._block_table[:num_reqs].contiguous()
+        slot_mapping_gpu = slot_mapping.to(device=self.device)
+
+        attn_metadata = attn_utils.build_attn_metadata(
+            self.attn_metadata_builders,
+            num_reqs=num_reqs,
+            num_tokens=num_tokens,
+            query_start_loc_gpu=query_start_loc_gpu,
+            query_start_loc_cpu=qsl,
+            seq_lens=seq_lens_gpu,
+            max_seq_len=max_seq_len,
+            block_tables=[block_table],
+            slot_mappings=[slot_mapping_gpu],
+            kv_cache_config=self.kv_cache_config,
+        )
+
+        # Build slot_mappings_by_layer for set_forward_context.
+        # Fix for vllm 0.15.0
+        slot_mappings_by_layer: dict[str, torch.Tensor] = {}
+        for kv_cache_group in self.kv_cache_config.kv_cache_groups:
+            for layer_name in kv_cache_group.layer_names:
+                slot_mappings_by_layer[layer_name] = slot_mapping_gpu
+
+        return attn_metadata, positions_cpu.to(device=self.device), slot_mappings_by_layer
+
+
+class Qwen3TTSTalkerCodePredictorModelVLLM(nn.Module):
+    def __init__(
+        self,
+        config: Qwen3TTSTalkerCodePredictorConfig,
+        *,
+        talker_hidden_size: int | None = None,
+        cache_config=None,
+        quant_config=None,
+        prefix: str = "",
+    ) -> None:
+        super().__init__()
+        self.config = config
+        self.quant_config = quant_config
+
+        self.layers = nn.ModuleList(
+            [
+                Qwen3DecoderLayer(
+                    config, cache_config=cache_config, quant_config=quant_config, prefix=f"{prefix}.layers.{i}"
+                )
+                for i in range(config.num_hidden_layers)
+            ]
+        )
+        self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+
+        # Official code_predictor uses one embedding table per residual group.
+        # Some Qwen3-TTS checkpoints store codec embeddings in the talker hidden
+        # space, even when `code_predictor_config.hidden_size` is smaller.
+        # We keep the embedding dim aligned with the checkpoint and project down
+        # via `small_to_mtp_projection` in the wrapper module.
+        emb_dim = int(talker_hidden_size) if talker_hidden_size is not None else int(config.hidden_size)
+        self.codec_embedding = nn.ModuleList(
+            [nn.Embedding(config.vocab_size, emb_dim) for _ in range(config.num_code_groups - 1)]
+        )
+
+    def get_input_embeddings(self) -> nn.ModuleList:
+        return self.codec_embedding
+
+    def forward(self, positions: torch.Tensor, inputs_embeds: torch.Tensor) -> torch.Tensor:
+        # Token-major: [num_tokens, hidden]
+        hidden_states = inputs_embeds
+        residual = None
+        for layer in self.layers:
+            hidden_states, residual = layer(positions, hidden_states, residual)
+        hidden_states, _ = self.norm(hidden_states, residual)
+        return hidden_states
+
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+        # Match vLLM Qwen2/Qwen3 packing conventions: q_proj/k_proj/v_proj -> qkv_proj,
+        # gate_proj/up_proj -> gate_up_proj.
+        stacked_params_mapping = [
+            ("qkv_proj", "q_proj", "q"),
+            ("qkv_proj", "k_proj", "k"),
+            ("qkv_proj", "v_proj", "v"),
+            ("gate_up_proj", "gate_proj", 0),
+            ("gate_up_proj", "up_proj", 1),
+        ]
+        params_dict = dict(self.named_parameters(remove_duplicate=False))
+        loaded_params: set[str] = set()
+        for name, loaded_weight in weights:
+            if "rotary_emb.inv_freq" in name:
+                continue
+            if self.quant_config is not None and (scale_name := self.quant_config.get_cache_scale(name)):
+                param = params_dict[scale_name]
+                weight_loader = getattr(param, "weight_loader", default_weight_loader)
+                loaded_weight = loaded_weight if loaded_weight.dim() == 0 else loaded_weight[0]
+                weight_loader(param, loaded_weight)
+                loaded_params.add(scale_name)
+                continue
+
+            for param_name, weight_name, shard_id in stacked_params_mapping:
+                if weight_name not in name:
+                    continue
+                mapped = name.replace(weight_name, param_name)
+                if mapped.endswith(".bias") and mapped not in params_dict:
+                    continue
+                if is_pp_missing_parameter(mapped, self):
+                    continue
+                if mapped.endswith("scale"):
+                    mapped = maybe_remap_kv_scale_name(mapped, params_dict)
+                    if mapped is None:
+                        continue
+                param = params_dict.get(mapped)
+                if param is None:
+                    continue
+                weight_loader = getattr(param, "weight_loader", default_weight_loader)
+                if weight_loader == default_weight_loader:
+                    weight_loader(param, loaded_weight)
+                else:
+                    weight_loader(param, loaded_weight, shard_id)
+                loaded_params.add(mapped)
+                break
+            else:
+                mapped = maybe_remap_kv_scale_name(name, params_dict)
+                if mapped is None:
+                    continue
+                if name.endswith(".bias") and mapped not in params_dict:
+                    continue
+                if is_pp_missing_parameter(mapped, self):
+                    continue
+                param = params_dict.get(mapped)
+                if param is None:
+                    continue
+                weight_loader = getattr(param, "weight_loader", default_weight_loader)
+                weight_loader(param, loaded_weight)
+                loaded_params.add(mapped)
+        return loaded_params
+
+
+class Qwen3TTSTalkerCodePredictorForConditionalGenerationVLLM(nn.Module):
+    """vLLM-native code_predictor used by the AR talker (residual codebooks)."""
+
+    def __init__(
+        self,
+        *,
+        vllm_config: VllmConfig,
+        config: Qwen3TTSTalkerCodePredictorConfig,
+        talker_config: Qwen3TTSTalkerConfig,
+        prefix: str = "code_predictor",
+    ) -> None:
+        super().__init__()
+        self._vllm_config = vllm_config
+        self.config = config
+        self.talker_config = talker_config
+
+        # Keep module/weight names aligned with official checkpoint (talker.code_predictor.model.*).
+        self.model = Qwen3TTSTalkerCodePredictorModelVLLM(
+            config,
+            talker_hidden_size=int(talker_config.hidden_size),
+            cache_config=vllm_config.cache_config,
+            quant_config=vllm_config.quant_config,
+            prefix=f"{prefix}.model",
+        )
+
+        # One head per residual group.
+        self.lm_head = nn.ModuleList(
+            [nn.Linear(config.hidden_size, config.vocab_size, bias=False) for _ in range(config.num_code_groups - 1)]
+        )
+
+        if config.hidden_size != talker_config.hidden_size:
+            self.small_to_mtp_projection = nn.Linear(talker_config.hidden_size, config.hidden_size, bias=True)
+        else:
+            self.small_to_mtp_projection = nn.Identity()
+
+        self._kv_cache: _LocalPredictorKVCache | None = None
+
+    def get_input_embeddings(self) -> nn.ModuleList:
+        return self.model.get_input_embeddings()
+
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+        # Ensure all vLLM custom layers consult the predictor vllm_config
+        # (esp. for Attention static_forward_context).
+        with set_current_vllm_config(self._vllm_config):
+            loaded: set[str] = set()
+            model_weights: list[tuple[str, torch.Tensor]] = []
+            other_weights: list[tuple[str, torch.Tensor]] = []
+            for name, w in weights:
+                if name.startswith("model."):
+                    model_weights.append((name[len("model.") :], w))
+                else:
+                    other_weights.append((name, w))
+
+            loaded_model = self.model.load_weights(model_weights)
+            loaded |= {f"model.{n}" for n in loaded_model}
+
+            params = dict(self.named_parameters(remove_duplicate=False))
+            for name, w in other_weights:
+                if name not in params:
+                    continue
+                default_weight_loader(params[name], w)
+                loaded.add(name)
+            return loaded
+
+    def _maybe_init_kv_cache(self, device: torch.device) -> None:
+        if self._kv_cache is not None:
+            return
+        max_seq_len = int(getattr(self.config, "num_code_groups", 16) or 16)
+        # Upper bound on batch size: vLLM scheduler max_num_seqs (fallback 8).
+        max_batch = int(getattr(self._vllm_config.scheduler_config, "max_num_seqs", 8) or 8)
+        max_batch = max(1, max_batch)
+        self._kv_cache = _LocalPredictorKVCache(
+            vllm_config=self._vllm_config,
+            max_seq_len=max_seq_len,
+            max_batch_size=max_batch,
+            device=device,
+        )
+
+    @torch.inference_mode()
+    def reset_cache(self) -> None:
+        # We reuse a fixed kv cache buffer and overwrite starting at slot 0.
+        # No action required here (seq_lens controls what is read).
+        return
+
+    @torch.inference_mode()
+    def prefill_logits(self, inputs_embeds: torch.Tensor) -> torch.Tensor:
+        """Prefill with 2 tokens: [past_hidden, layer0_embed]. Returns logits for residual group 0."""
+        self._maybe_init_kv_cache(inputs_embeds.device)
+        assert self._kv_cache is not None
+
+        bsz = int(inputs_embeds.shape[0])
+        qlen = 2
+        # Flatten to token-major.
+        hs = inputs_embeds.to(dtype=torch.bfloat16).reshape(bsz * qlen, -1)
+        hs = self.small_to_mtp_projection(hs)
+
+        query_lens = torch.full((bsz,), qlen, dtype=torch.int32)
+        seq_lens = query_lens.clone()
+        attn_metadata, positions, slot_mappings = self._kv_cache.build_attn_metadata(
+            num_reqs=bsz, query_lens=query_lens, seq_lens=seq_lens
+        )
+
+        with (
+            set_current_vllm_config(self._vllm_config),
+            set_forward_context(
+                attn_metadata,
+                self._vllm_config,
+                num_tokens=int(hs.shape[0]),
+                slot_mapping=slot_mappings,
+            ),
+        ):
+            out = self.model(positions=positions, inputs_embeds=hs)
+
+        # Gather last token per request.
+        last_idx = torch.arange(qlen - 1, bsz * qlen, step=qlen, device=out.device, dtype=torch.long)
+        last_h = out.index_select(0, last_idx)
+        logits = self.lm_head[0](last_h)
+        return logits
+
+    @torch.inference_mode()
+    def decode_logits(self, input_ids: torch.Tensor, *, generation_step: int, past_seq_len: int) -> torch.Tensor:
+        """Decode one new token for residual group `generation_step` (1..Q-1)."""
+        self._maybe_init_kv_cache(input_ids.device)
+        assert self._kv_cache is not None
+        bsz = int(input_ids.shape[0])
+        if generation_step <= 0:
+            raise ValueError("generation_step must be >= 1 for decode_logits")
+
+        embed_idx = generation_step - 1
+        hs = self.model.get_input_embeddings()[embed_idx](input_ids.to(dtype=torch.long).reshape(bsz, 1))
+        hs = self.small_to_mtp_projection(hs.reshape(bsz, -1))
+
+        query_lens = torch.ones((bsz,), dtype=torch.int32)
+        seq_lens = torch.full((bsz,), int(past_seq_len) + 1, dtype=torch.int32)
+        attn_metadata, positions, slot_mappings = self._kv_cache.build_attn_metadata(
+            num_reqs=bsz, query_lens=query_lens, seq_lens=seq_lens
+        )
+
+        with (
+            set_current_vllm_config(self._vllm_config),
+            set_forward_context(
+                attn_metadata,
+                self._vllm_config,
+                num_tokens=int(hs.shape[0]),
+                slot_mapping=slot_mappings,
+            ),
+        ):
+            out = self.model(positions=positions, inputs_embeds=hs)
+
+        logits = self.lm_head[generation_step](out)
+        return logits
+
+    @torch.inference_mode()
+    def forward(
+        self,
+        layer0_code: torch.Tensor,
+        layer0_embed: torch.Tensor,
+        last_talker_hidden: torch.Tensor,
+        do_sample: bool = True,
+        temperature: float = 1.0,
+        top_k: int = 50,
+        top_p: float = 1.0,
+    ) -> torch.Tensor:
+        """Full autoregressive prediction of residual codebooks 1..Q-1.
+
+        Args:
+            layer0_code: [B, 1] first-layer codec token ids.
+            layer0_embed: [B, 1, H] embedding of layer0_code.
+            last_talker_hidden: [B, 1, H] hidden state from the talker.
+            do_sample: whether to sample or take argmax.
+            temperature: sampling temperature.
+            top_k: top-k filtering.
+            top_p: top-p (nucleus) filtering.
+
+        Returns:
+            audio_codes: [B, Q] all codebook tokens (layer0 + residuals).
+        """
+        bsz = int(layer0_code.shape[0])
+        num_groups = int(self.config.num_code_groups)
+        max_steps = num_groups - 1
+
+        # Reset KV cache for a fresh sequence.
+        self.reset_cache()
+
+        # Prefill: feed [last_talker_hidden, layer0_embed] → logits for group 1.
+        prefill_input = torch.cat([last_talker_hidden, layer0_embed], dim=1)  # [B, 2, H]
+        logits = self.prefill_logits(prefill_input)  # [B, vocab]
+
+        all_codes = [layer0_code.reshape(bsz, 1)]
+        past_seq_len = 2
+
+        for step in range(1, num_groups):
+            # Sample or argmax from logits.
+            if do_sample and temperature > 0:
+                scaled = logits / temperature
+                if top_k > 0:
+                    topk_vals, _ = scaled.topk(top_k, dim=-1)
+                    scaled = scaled.masked_fill(scaled < topk_vals[:, -1:], float("-inf"))
+                probs = torch.softmax(scaled, dim=-1)
+                next_ids = torch.multinomial(probs, num_samples=1)  # [B, 1]
+            else:
+                next_ids = logits.argmax(dim=-1, keepdim=True)  # [B, 1]
+            all_codes.append(next_ids)
+
+            # If not the last step, decode one more token.
+            if step < max_steps:
+                logits = self.decode_logits(
+                    next_ids.reshape(bsz),
+                    generation_step=step,
+                    past_seq_len=past_seq_len,
+                )
+                past_seq_len += 1
+
+        return torch.cat(all_codes, dim=1)  # [B, Q]

vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_tokenizer.py

@@ -80,16 +80,21 @@ class Qwen3TTSTokenizer:
         """
         inst = cls()
 
-        AutoConfig.register("qwen3_tts_tokenizer_25hz", Qwen3TTSTokenizerV1Config)
-        AutoModel.register(Qwen3TTSTokenizerV1Config, Qwen3TTSTokenizerV1Model)
+        load_feature_extractor = bool(kwargs.pop("load_feature_extractor", True))
 
         AutoConfig.register("qwen3_tts_tokenizer_12hz", Qwen3TTSTokenizerV2Config)
         AutoModel.register(Qwen3TTSTokenizerV2Config, Qwen3TTSTokenizerV2Model)
 
-        inst.feature_extractor = AutoFeatureExtractor.from_pretrained(pretrained_model_name_or_path)
+        AutoConfig.register("qwen3_tts_tokenizer_25hz", Qwen3TTSTokenizerV1Config)
+        AutoModel.register(Qwen3TTSTokenizerV1Config, Qwen3TTSTokenizerV1Model)
+
         inst.model = AutoModel.from_pretrained(pretrained_model_name_or_path, **kwargs)
         inst.config = inst.model.config
 
+        inst.feature_extractor = (
+            AutoFeatureExtractor.from_pretrained(pretrained_model_name_or_path) if load_feature_extractor else None
+        )
+
         inst.device = getattr(inst.model, "device", None)
         if inst.device is None:
             # fallback: infer from first parameter device
@@ -212,12 +217,7 @@ class Qwen3TTSTokenizer:
         audios: AudioInput,
         sr: int | None = None,
         return_dict: bool = True,
-    ) -> (
-        Qwen3TTSTokenizerV1EncoderOutput
-        | Qwen3TTSTokenizerV2EncoderOutput
-        | tuple[list[torch.Tensor], list[torch.Tensor] | None, list[torch.Tensor] | None]
-        | tuple[list[torch.Tensor]]
-    ):
+    ) -> Qwen3TTSTokenizerV1EncoderOutput | Qwen3TTSTokenizerV2EncoderOutput | tuple:
         """
         Batch-encode audio into discrete codes (and optional conditioning, depending on 25Hz/12Hz).
 

vllm_omni/model_executor/models/registry.py

@@ -50,8 +50,18 @@ _OMNI_MODELS = {
     ),
     "Qwen3TTSForConditionalGeneration": (
         "qwen3_tts",
+        "qwen3_tts_talker",
+        "Qwen3TTSTalkerForConditionalGeneration",
+    ),
+    "Qwen3TTSTalkerForConditionalGeneration": (
         "qwen3_tts",
-        "Qwen3TTSModelForGeneration",
+        "qwen3_tts_talker",
+        "Qwen3TTSTalkerForConditionalGeneration",
+    ),
+    "Qwen3TTSCode2Wav": (
+        "qwen3_tts",
+        "qwen3_tts_code2wav",
+        "Qwen3TTSCode2Wav",
     ),
 }
 

vllm_omni/model_executor/stage_configs/qwen3_tts.yaml

@@ -1,22 +1,101 @@
+async_chunk: true
 stage_args:
   - stage_id: 0
-    stage_type: llm  # Use llm stage type to launch OmniLLM
+    stage_type: llm
     runtime:
       devices: "0"
       max_batch_size: 1
     engine_args:
       model_stage: qwen3_tts
-      model_arch: Qwen3TTSForConditionalGeneration
+      model_arch: Qwen3TTSTalkerForConditionalGeneration
+      # Force stage-specific registered architecture.
+      hf_overrides:
+        architectures: [Qwen3TTSTalkerForConditionalGeneration]
+      worker_type: ar
+      scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
+      enforce_eager: false
+      trust_remote_code: true
+      async_scheduling: false
+      enable_prefix_caching: false
+      engine_output_type: latent
+      gpu_memory_utilization: 0.3
+      distributed_executor_backend: "mp"
+      max_num_batched_tokens: 512
+      max_model_len: 4096
+      custom_process_next_stage_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_tts.talker2code2wav_async_chunk
+    # Use named connector to apply runtime.connectors.extra.
+    output_connectors:
+      to_stage_1: connector_of_shared_memory
+    default_sampling_params:
+      temperature: 0.9
+      top_k: 50
+      max_tokens: 4096
+      seed: 42
+      detokenize: false
+      repetition_penalty: 1.05
+      stop_token_ids: [2150]
+
+  - stage_id: 1
+    stage_type: llm
+    runtime:
+      devices: "0"
+      max_batch_size: 1
+    engine_args:
+      model_stage: code2wav
+      model_arch: Qwen3TTSCode2Wav
+      # Force stage-specific registered architecture.
+      hf_overrides:
+        architectures: [Qwen3TTSCode2Wav]
       worker_type: generation
       scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
       enforce_eager: true
       trust_remote_code: true
       async_scheduling: false
       enable_prefix_caching: false
-      engine_output_type: audio  # Final output: audio waveform
-      gpu_memory_utilization: 0.1
+      engine_output_type: audio
+      gpu_memory_utilization: 0.2
       distributed_executor_backend: "mp"
-      max_num_batched_tokens: 1000000
-
+      # Must be divisible by num_code_groups and cover (left_context + chunk).
+      max_num_batched_tokens: 8192
+      # async_chunk appends windows per step; max_model_len must cover accumulated stream.
+      max_model_len: 32768
+    engine_input_source: [0]
     final_output: true
     final_output_type: audio
+    # Distributed connector configuration
+    input_connectors:
+      from_stage_0: connector_of_shared_memory
+    default_sampling_params:
+      temperature: 0.0
+      top_p: 1.0
+      top_k: -1
+      max_tokens: 65536
+      seed: 42
+      detokenize: true
+      repetition_penalty: 1.0
+
+runtime:
+  enabled: true
+  defaults:
+    window_size: -1
+    max_inflight: 1
+
+  connectors:
+    connector_of_shared_memory:
+      name: SharedMemoryConnector
+      extra:
+        shm_threshold_bytes: 65536
+        # Frame-aligned codec streaming transport.
+        codec_streaming: true
+        # Connector polling / timeout (unit: loop count, sleep interval in seconds).
+        connector_get_sleep_s: 0.01
+        connector_get_max_wait_first_chunk: 3000
+        connector_get_max_wait: 300
+        # Align with Omni: small chunks with sufficient context overlap.
+        codec_chunk_frames: 25
+        codec_left_context_frames: 25
+
+  edges:
+    - from: 0
+      to: 1
+      window_size: -1

vllm_omni/model_executor/stage_input_processors/qwen3_tts.py

@@ -0,0 +1,76 @@
+"""Stage input processor for Qwen3-TTS: Talker -> Code2Wav."""
+
+from typing import Any
+
+import torch
+
+
+def _extract_last_frame(pooling_output: dict[str, Any]) -> torch.Tensor | None:
+    audio_codes = pooling_output.get("audio_codes")
+    if not isinstance(audio_codes, torch.Tensor) or audio_codes.numel() == 0:
+        return None
+    if audio_codes.ndim == 2:
+        frame = audio_codes[-1]
+        if frame.numel() == 0 or not bool(frame.any().item()):
+            return None
+        return frame.to(torch.long).reshape(-1)
+    if audio_codes.ndim == 1:
+        return audio_codes.to(torch.long).reshape(-1)
+    raise ValueError(f"Invalid audio_codes shape for Qwen3-TTS async_chunk: {tuple(audio_codes.shape)}")
+
+
+def talker2code2wav_async_chunk(
+    transfer_manager: Any,
+    pooling_output: dict[str, Any],
+    request: Any,
+) -> dict[str, Any] | None:
+    if not isinstance(pooling_output, dict):
+        return None
+
+    request_id = request.external_req_id
+
+    connector = getattr(transfer_manager, "connector", None)
+    raw_cfg = getattr(connector, "config", {}) or {}
+    cfg = raw_cfg.get("extra", raw_cfg) if isinstance(raw_cfg, dict) else {}
+    chunk_size = int(cfg.get("codec_chunk_frames", 25))
+    left_context_size = int(cfg.get("codec_left_context_frames", 25))
+    if chunk_size <= 0 or left_context_size < 0:
+        raise ValueError(
+            f"Invalid codec chunk config: codec_chunk_frames={chunk_size}, "
+            f"codec_left_context_frames={left_context_size}"
+        )
+
+    finished = bool(request.is_finished())
+
+    frame = _extract_last_frame(pooling_output)
+    if frame is not None:
+        codec_codes = frame.cpu().tolist()
+        transfer_manager.code_prompt_token_ids[request_id].append(codec_codes)
+
+    length = len(transfer_manager.code_prompt_token_ids[request_id])
+    chunk_length = length % chunk_size
+
+    if chunk_length != 0 and not finished:
+        return None
+
+    context_length = chunk_length if chunk_length != 0 else chunk_size
+
+    if length <= 0:
+        return {
+            "code_predictor_codes": [],
+            "finished": torch.tensor(bool(finished), dtype=torch.bool),
+        }
+
+    end_index = min(length, left_context_size + context_length)
+    ctx_frames = max(0, int(end_index - context_length))
+    window_frames = transfer_manager.code_prompt_token_ids[request_id][-end_index:]
+
+    # Pack context + chunk into codebook-major flat codes for adapter.
+    code_predictor_codes = torch.tensor(window_frames).transpose(0, 1).reshape(-1).tolist()
+
+    # Build final prompt_token_ids with ctx_frames header for Qwen3-TTS Code2Wav.
+    # The model expects input_ids layout: [ctx_frames, *flat_codes].
+    return {
+        "code_predictor_codes": [int(ctx_frames)] + code_predictor_codes,
+        "finished": torch.tensor(bool(finished), dtype=torch.bool),
+    }

vllm_omni/platforms/__init__.py

@@ -84,8 +84,6 @@ def xpu_omni_platform_plugin() -> str | None:
     is_xpu = False
     logger.debug("Checking if XPU OmniPlatform is available.")
     try:
-        # installed IPEX if the machine has XPUs.
-        import intel_extension_for_pytorch  # noqa: F401
         import torch
 
         if supports_xccl():

vllm_omni/platforms/npu/stage_configs/qwen3_tts.yaml

@@ -1,22 +1,92 @@
+async_chunk: true
 stage_args:
   - stage_id: 0
-    stage_type: llm  # Use llm stage type to launch OmniLLM
+    stage_type: llm
     runtime:
       devices: "0"
       max_batch_size: 1
     engine_args:
       model_stage: qwen3_tts
-      model_arch: Qwen3TTSForConditionalGeneration
+      model_arch: Qwen3TTSTalkerForConditionalGeneration
+      hf_overrides:
+        architectures: [Qwen3TTSTalkerForConditionalGeneration]
+      worker_type: ar
+      scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
+      enforce_eager: true
+      trust_remote_code: true
+      async_scheduling: false
+      enable_prefix_caching: false
+      engine_output_type: latent
+      gpu_memory_utilization: 0.3
+      distributed_executor_backend: "mp"
+      max_num_batched_tokens: 512
+      max_model_len: 4096
+      custom_process_next_stage_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_tts.talker2code2wav_async_chunk
+    output_connectors:
+      to_stage_1: connector_of_shared_memory
+    default_sampling_params:
+      temperature: 0.9
+      top_k: 50
+      max_tokens: 4096
+      seed: 42
+      detokenize: false
+      repetition_penalty: 1.05
+      stop_token_ids: [2150]
+
+  - stage_id: 1
+    stage_type: llm
+    runtime:
+      devices: "0"
+      max_batch_size: 1
+    engine_args:
+      model_stage: code2wav
+      model_arch: Qwen3TTSCode2Wav
+      hf_overrides:
+        architectures: [Qwen3TTSCode2Wav]
       worker_type: generation
       scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
       enforce_eager: true
       trust_remote_code: true
       async_scheduling: false
       enable_prefix_caching: false
-      engine_output_type: audio  # Final output: audio waveform
-      gpu_memory_utilization: 0.1
+      engine_output_type: audio
+      gpu_memory_utilization: 0.2
       distributed_executor_backend: "mp"
-      max_num_batched_tokens: 1000000
-
+      max_num_batched_tokens: 8192
+      max_model_len: 32768
+    engine_input_source: [0]
     final_output: true
     final_output_type: audio
+    input_connectors:
+      from_stage_0: connector_of_shared_memory
+    default_sampling_params:
+      temperature: 0.0
+      top_p: 1.0
+      top_k: -1
+      max_tokens: 65536
+      seed: 42
+      detokenize: true
+      repetition_penalty: 1.0
+
+runtime:
+  enabled: true
+  defaults:
+    window_size: -1
+    max_inflight: 1
+
+  connectors:
+    connector_of_shared_memory:
+      name: SharedMemoryConnector
+      extra:
+        shm_threshold_bytes: 65536
+        codec_streaming: true
+        connector_get_sleep_s: 0.01
+        connector_get_max_wait_first_chunk: 3000
+        connector_get_max_wait: 300
+        codec_chunk_frames: 25
+        codec_left_context_frames: 25
+
+  edges:
+    - from: 0
+      to: 1
+      window_size: -1

vllm_omni/platforms/npu/worker/npu_generation_model_runner.py

@@ -238,9 +238,12 @@ class NPUGenerationModelRunner(OmniNPUModelRunner):
             pooler_output.append(mm_payload)
         else:
             raise RuntimeError("Unsupported diffusion output type")
+        # [Omni] Copy req_id mappings to avoid async scheduling mutation.
+        req_ids_output_copy = self.input_batch.req_ids.copy()
+        req_id_to_index_output_copy = self.input_batch.req_id_to_index.copy()
         output = OmniModelRunnerOutput(
-            req_ids=self.input_batch.req_ids,
-            req_id_to_index=self.input_batch.req_id_to_index,
+            req_ids=req_ids_output_copy,
+            req_id_to_index=req_id_to_index_output_copy,
             sampled_token_ids=[],
             logprobs=None,
             prompt_logprobs_dict={},

vllm_omni/worker/gpu_generation_model_runner.py

@@ -380,9 +380,12 @@ class GPUGenerationModelRunner(OmniGPUModelRunner):
                 pooler_output.append(mm_payload)
         else:
             raise RuntimeError("Unsupported diffusion output type")
+        # [Omni] Copy req_id mappings to avoid async scheduling mutation.
+        req_ids_output_copy = self.input_batch.req_ids.copy()
+        req_id_to_index_output_copy = self.input_batch.req_id_to_index.copy()
         output = OmniModelRunnerOutput(
-            req_ids=self.input_batch.req_ids,
-            req_id_to_index=self.input_batch.req_id_to_index,
+            req_ids=req_ids_output_copy,
+            req_id_to_index=req_id_to_index_output_copy,
             sampled_token_ids=[],
             logprobs=None,
             prompt_logprobs_dict={},

vllm_omni/worker/gpu_model_runner.py

@@ -70,16 +70,24 @@ class OmniGPUModelRunner(GPUModelRunner):
     @instrument(span_name="Loading (GPU)")
     def load_model(self, *args, **kwargs) -> None:
         super().load_model(*args, **kwargs)
+
         # TODO move this model specific logic to a separate class
-        if hasattr(self.model, "talker_mtp") and self.model.talker is not None:
-            self.talker_mtp = self.model.talker_mtp
+        # TTS model IS the talker (no .talker sub-attr); use getattr to support both Omni and TTS.
+        talker_mtp = getattr(self.model, "talker_mtp", None)
+        if talker_mtp is not None:
+            self.talker_mtp = talker_mtp  # type: ignore[assignment]
             cudagraph_mode = self.compilation_config.cudagraph_mode
             assert cudagraph_mode is not None
-            if cudagraph_mode.has_full_cudagraphs():
-                self.talker_mtp = CUDAGraphWrapper(
-                    self.model.talker_mtp, self.vllm_config, runtime_mode=CUDAGraphMode.FULL
-                )
-            hidden_size = self.model_config.hf_config.talker_config.text_config.hidden_size
+            # Only wrap talker_mtp in CUDAGraphWrapper for Omni models that
+            # have a separate .talker sub-module.  TTS models' code predictor
+            # has internal AR loops / torch.multinomial — not graph-safe.
+            has_separate_talker = getattr(self.model, "talker", None) is not None
+            if cudagraph_mode.has_full_cudagraphs() and has_separate_talker:
+                self.talker_mtp = CUDAGraphWrapper(talker_mtp, self.vllm_config, runtime_mode=CUDAGraphMode.FULL)
+            # TTS exposes mtp_hidden_size; Omni uses hf_text_config.hidden_size.
+            hidden_size = int(
+                getattr(self.model, "mtp_hidden_size", 0) or getattr(self.model_config.hf_text_config, "hidden_size")
+            )
             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(
@@ -644,6 +652,11 @@ class OmniGPUModelRunner(GPUModelRunner):
                 input_ids = None
                 inputs_embeds = self.inputs_embeds.gpu[:num_tokens_padded]
                 model_kwargs = self._init_model_kwargs()
+            elif getattr(getattr(self, "model", None), "has_preprocess", False):
+                # Capture CUDA graph with inputs_embeds path so replay reads
+                # from the same buffer that _preprocess writes into.
+                input_ids = self.input_ids.gpu[:num_tokens_padded]
+                inputs_embeds = self.inputs_embeds.gpu[:num_tokens_padded]
             else:
                 input_ids = self.input_ids.gpu[:num_tokens_padded]
                 inputs_embeds = None
@@ -986,6 +999,11 @@ class OmniGPUModelRunner(GPUModelRunner):
             inputs_embeds = self.inputs_embeds.gpu[:num_input_tokens]
             model_kwargs = self._init_model_kwargs()
             input_ids = self.input_ids.gpu[:num_input_tokens]
+        elif getattr(self.model, "has_preprocess", False):
+            # Use pre-allocated buffer for CUDA graph compatibility.
+            input_ids = self.input_ids.gpu[:num_input_tokens]
+            inputs_embeds = self.inputs_embeds.gpu[:num_input_tokens]
+            model_kwargs = self._init_model_kwargs()
         else:
             # For text-only models, we use token ids as input.
             # While it is possible to use embeddings as input just like the
@@ -1047,9 +1065,16 @@ class OmniGPUModelRunner(GPUModelRunner):
                 span_len = int(e) - int(s)
 
                 # call the custom process function
+                embed_slice = inputs_embeds[s:e] if inputs_embeds is not None else None
                 req_input_ids, req_embeds, update_dict = self.model.preprocess(
-                    input_ids=input_ids[s:e], input_embeds=inputs_embeds[s:e], **req_infos
+                    input_ids=input_ids[s:e], input_embeds=embed_slice, **req_infos
                 )
+                if inputs_embeds is None:
+                    inputs_embeds = torch.empty(
+                        (input_ids.shape[0], req_embeds.shape[-1]),
+                        device=req_embeds.device,
+                        dtype=req_embeds.dtype,
+                    )
 
                 if hasattr(self.model, "talker_mtp") and span_len == 1:
                     last_talker_hidden, text_step = update_dict.pop("mtp_inputs")
@@ -1093,6 +1118,9 @@ class OmniGPUModelRunner(GPUModelRunner):
             max_num_scheduled_tokens=1,
             use_cascade_attn=False,
         )
+        # Force eager for unwrapped code predictors (AR loops / multinomial).
+        if not isinstance(self.talker_mtp, CUDAGraphWrapper):
+            _cudagraph_mode = CUDAGraphMode.NONE
         num_tokens_padded = batch_desc.num_tokens
         req_input_ids = self.talker_mtp_input_ids.gpu[:num_tokens_padded]
         req_embeds = self.talker_mtp_inputs_embeds.gpu[:num_tokens_padded]
@@ -1104,11 +1132,12 @@ class OmniGPUModelRunner(GPUModelRunner):
             req_embeds, code_predictor_codes = self.talker_mtp(req_input_ids, req_embeds, last_talker_hidden, text_step)
         # update the inputs_embeds and code_predictor_codes
         code_predictor_codes_cpu = code_predictor_codes.detach().to("cpu").contiguous()
+        out_key = getattr(self.model, "talker_mtp_output_key", "code_predictor_codes")
         for idx, req_id in enumerate(decode_req_ids):
             req_index = self.input_batch.req_ids.index(req_id)
             start_offset = int(self.query_start_loc.cpu[req_index])
             inputs_embeds[start_offset : start_offset + 1] = req_embeds[idx : idx + 1]
-            update_dict = {"code_predictor_codes": code_predictor_codes_cpu[idx : idx + 1]}
+            update_dict = {out_key: code_predictor_codes_cpu[idx : idx + 1]}
             self._merge_additional_information_update(req_id, update_dict)
 
     def _model_forward(
@@ -1142,7 +1171,9 @@ class OmniGPUModelRunner(GPUModelRunner):
         self._omni_last_model_output = model_output
         return model_output
 
-    def _merge_additional_information_update(self, req_id: str, upd: dict) -> None:
+    def _merge_additional_information_update(self, req_id: str, upd: dict | None) -> None:
+        if not isinstance(upd, dict):
+            return
         req_state = self.requests.get(req_id)
         if req_state is None:
             return