clean submission: comp168 base + max_model_len=100000

.dockerignore

@@ -0,0 +1,21 @@
+# Exclude everything not needed for the Docker image
+cccl_upstream/
+upstream_ref/
+vllm/
+muh/
+docs/
+optimizations/
+vllm_adapter/
+*.zip
+*.txt
+*.md
+*.json
+*.muh
+.git/
+.gitignore
+__pycache__/
+*.pyc
+debug_*.py
+verify_*.py
+# Keep: qwen3_6_scripts/, computility-run.yaml, Dockerfile, ex_engine/
+ex_engine/

Dockerfile

@@ -0,0 +1,13 @@
+FROM git.modelhub.org.cn:9443/enginex-iluvatar/bi100-3.2.3-x86-ubuntu20.04-py3.10-poc-llm-infer:v1.2.3
+
+RUN mkdir -p /workspace
+WORKDIR /workspace/
+
+# Copy all our engine patches
+COPY ./qwen3_6_scripts /workspace/qwen3_6_scripts
+COPY ./computility-run.yaml /workspace/computility-run.yaml
+
+# Make patch script executable and run it
+RUN chmod +x /workspace/qwen3_6_scripts/patch_ops.sh && \
+    bash /workspace/qwen3_6_scripts/patch_ops.sh 2>&1 | tee /workspace/patch_ops.log ; \
+    echo "[Dockerfile] patch_ops exit code: $?"

computility-run.yaml

@@ -0,0 +1,46 @@
+concurrency: 1
+command:
+    - python3
+    - -m
+    - vllm.entrypoints.openai.api_server
+    - --model
+    - /model
+    - --served-model-name
+    - llm
+    - --max-model-len
+    - '100000'
+    - --gpu-memory-utilization
+    - '0.90'
+    - --trust-remote-code
+    - -tp
+    - '4'
+    - --max-num-seqs
+    - '2'
+    - --disable-log-requests
+    - --disable-frontend-multiprocessing
+    - --enforce-eager
+    - --enable-auto-tool-choice
+    - --tool-call-parser
+    - qwen3_coder
+    - --reasoning-parser
+    - qwen3
+    - --enable-prefix-caching
+    - --max-seq-len-to-capture
+    - '8192'
+    - --dtype
+    - half
+env:
+    - name: VLLM_ENGINE_ITERATION_TIMEOUT_S
+      value: '3600'
+    - name: VLLM_ATTENTION_BACKEND
+      value: XFORMERS
+    - name: ENABLE_CUSTOM_IPC
+      value: '1'
+    - name: PYTHONPATH
+      value: /usr/local/corex/lib/python3/dist-packages:/usr/local/corex/lib64/python3/dist-packages
+    - name: LD_LIBRARY_PATH
+      value: /usr/local/corex/lib64:/usr/local/openmpi/lib:/usr/local/corex/lib64/python3/dist-packages/ixformer
+    - name: PYTORCH_CUDA_ALLOC_CONF
+      value: max_split_size_mb:512
+    - name: OMP_NUM_THREADS
+      value: '1'

qwen3_6_scripts/api_server.py

@@ -0,0 +1,595 @@
+import asyncio
+import importlib
+import inspect
+import multiprocessing
+import os
+import regex as re
+import signal
+import socket
+import tempfile
+from argparse import Namespace
+from contextlib import asynccontextmanager
+from functools import partial
+from http import HTTPStatus
+from typing import AsyncIterator, Set
+
+import uvloop
+from fastapi import APIRouter, FastAPI, Request
+from fastapi.exceptions import RequestValidationError
+from fastapi.middleware.cors import CORSMiddleware
+from fastapi.responses import JSONResponse, Response, StreamingResponse
+from starlette.datastructures import State
+from starlette.routing import Mount
+from typing_extensions import assert_never
+
+import vllm.envs as envs
+from vllm.config import ModelConfig
+from vllm.engine.arg_utils import AsyncEngineArgs
+from vllm.engine.async_llm_engine import AsyncLLMEngine
+from vllm.engine.multiprocessing.client import MQLLMEngineClient
+from vllm.engine.multiprocessing.engine import run_mp_engine
+from vllm.engine.protocol import EngineClient
+from vllm.entrypoints.launcher import serve_http
+from vllm.entrypoints.logger import RequestLogger
+from vllm.entrypoints.openai.cli_args import (make_arg_parser,
+                                              validate_parsed_serve_args)
+# yapf conflicts with isort for this block
+# yapf: disable
+from vllm.entrypoints.openai.protocol import (ChatCompletionRequest,
+                                              ChatCompletionResponse,
+                                              CompletionRequest,
+                                              CompletionResponse,
+                                              DetokenizeRequest,
+                                              DetokenizeResponse,
+                                              EmbeddingRequest,
+                                              EmbeddingResponse, ErrorResponse,
+                                              LoadLoraAdapterRequest,
+                                              TokenizeRequest,
+                                              TokenizeResponse,
+                                              UnloadLoraAdapterRequest)
+# yapf: enable
+from vllm.entrypoints.openai.serving_chat import OpenAIServingChat
+from vllm.entrypoints.openai.serving_completion import OpenAIServingCompletion
+from vllm.entrypoints.openai.serving_embedding import OpenAIServingEmbedding
+from vllm.entrypoints.openai.serving_engine import BaseModelPath
+from vllm.entrypoints.openai.serving_tokenization import (
+    OpenAIServingTokenization)
+from vllm.entrypoints.openai.tool_parsers import ToolParserManager
+from vllm.reasoning import ReasoningParserManager
+from vllm.logger import init_logger
+from vllm.usage.usage_lib import UsageContext
+from vllm.utils import FlexibleArgumentParser, get_open_zmq_ipc_path
+from vllm.version import __version__ as VLLM_VERSION
+
+TIMEOUT_KEEP_ALIVE = 5  # seconds
+
+prometheus_multiproc_dir: tempfile.TemporaryDirectory
+
+# Cannot use __name__ (https://github.com/vllm-project/vllm/pull/4765)
+logger = init_logger('vllm.entrypoints.openai.api_server')
+
+_running_tasks: Set[asyncio.Task] = set()
+
+
+@asynccontextmanager
+async def lifespan(app: FastAPI):
+    try:
+        if app.state.log_stats:
+            engine_client: EngineClient = app.state.engine_client
+
+            async def _force_log():
+                while True:
+                    await asyncio.sleep(10.)
+                    await engine_client.do_log_stats()
+
+            task = asyncio.create_task(_force_log())
+            _running_tasks.add(task)
+            task.add_done_callback(_running_tasks.remove)
+        else:
+            task = None
+        try:
+            yield
+        finally:
+            if task is not None:
+                task.cancel()
+    finally:
+        # Ensure app state including engine ref is gc'd
+        del app.state
+
+
+@asynccontextmanager
+async def build_async_engine_client(
+        args: Namespace) -> AsyncIterator[EngineClient]:
+
+    # Context manager to handle engine_client lifecycle
+    # Ensures everything is shutdown and cleaned up on error/exit
+    engine_args = AsyncEngineArgs.from_cli_args(args)
+
+    async with build_async_engine_client_from_engine_args(
+            engine_args, args.disable_frontend_multiprocessing) as engine:
+        yield engine
+
+
+@asynccontextmanager
+async def build_async_engine_client_from_engine_args(
+    engine_args: AsyncEngineArgs,
+    disable_frontend_multiprocessing: bool = False,
+) -> AsyncIterator[EngineClient]:
+    """
+    Create EngineClient, either:
+        - in-process using the AsyncLLMEngine Directly
+        - multiprocess using AsyncLLMEngine RPC
+
+    Returns the Client or None if the creation failed.
+    """
+
+    # Fall back
+    # TODO: fill out feature matrix.
+    if (MQLLMEngineClient.is_unsupported_config(engine_args)
+            or disable_frontend_multiprocessing):
+        engine_config = engine_args.create_engine_config()
+        uses_ray = getattr(AsyncLLMEngine._get_executor_cls(engine_config),
+                           "uses_ray", False)
+
+        build_engine = partial(AsyncLLMEngine.from_engine_args,
+                               engine_args=engine_args,
+                               engine_config=engine_config,
+                               usage_context=UsageContext.OPENAI_API_SERVER)
+        if uses_ray:
+            # Must run in main thread with ray for its signal handlers to work
+            engine_client = build_engine()
+        else:
+            engine_client = await asyncio.get_running_loop().run_in_executor(
+                None, build_engine)
+
+        yield engine_client
+        return
+
+    # Otherwise, use the multiprocessing AsyncLLMEngine.
+    else:
+        if "PROMETHEUS_MULTIPROC_DIR" not in os.environ:
+            # Make TemporaryDirectory for prometheus multiprocessing
+            # Note: global TemporaryDirectory will be automatically
+            #   cleaned up upon exit.
+            global prometheus_multiproc_dir
+            prometheus_multiproc_dir = tempfile.TemporaryDirectory()
+            os.environ[
+                "PROMETHEUS_MULTIPROC_DIR"] = prometheus_multiproc_dir.name
+        else:
+            logger.warning(
+                "Found PROMETHEUS_MULTIPROC_DIR was set by user. "
+                "This directory must be wiped between vLLM runs or "
+                "you will find inaccurate metrics. Unset the variable "
+                "and vLLM will properly handle cleanup.")
+
+        # Select random path for IPC.
+        ipc_path = get_open_zmq_ipc_path()
+        logger.info("Multiprocessing frontend to use %s for IPC Path.",
+                    ipc_path)
+
+        # Start RPCServer in separate process (holds the LLMEngine).
+        # the current process might have CUDA context,
+        # so we need to spawn a new process
+        context = multiprocessing.get_context("spawn")
+
+        engine_process = context.Process(target=run_mp_engine,
+                                         args=(engine_args,
+                                               UsageContext.OPENAI_API_SERVER,
+                                               ipc_path))
+        engine_process.start()
+        logger.info("Started engine process with PID %d", engine_process.pid)
+
+        # Build RPCClient, which conforms to EngineClient Protocol.
+        # NOTE: Actually, this is not true yet. We still need to support
+        # embedding models via RPC (see TODO above)
+        engine_config = engine_args.create_engine_config()
+        mp_engine_client = MQLLMEngineClient(ipc_path, engine_config)
+
+        try:
+            while True:
+                try:
+                    await mp_engine_client.setup()
+                    break
+                except TimeoutError:
+                    if not engine_process.is_alive():
+                        raise RuntimeError(
+                            "Engine process failed to start") from None
+
+            yield mp_engine_client  # type: ignore[misc]
+        finally:
+            # Ensure rpc server process was terminated
+            engine_process.terminate()
+
+            # Close all open connections to the backend
+            mp_engine_client.close()
+
+            # Wait for engine process to join
+            engine_process.join(4)
+            if engine_process.exitcode is None:
+                # Kill if taking longer than 5 seconds to stop
+                engine_process.kill()
+
+            # Lazy import for prometheus multiprocessing.
+            # We need to set PROMETHEUS_MULTIPROC_DIR environment variable
+            # before prometheus_client is imported.
+            # See https://prometheus.github.io/client_python/multiprocess/
+            from prometheus_client import multiprocess
+            multiprocess.mark_process_dead(engine_process.pid)
+
+
+router = APIRouter()
+
+
+def mount_metrics(app: FastAPI):
+    # Lazy import for prometheus multiprocessing.
+    # We need to set PROMETHEUS_MULTIPROC_DIR environment variable
+    # before prometheus_client is imported.
+    # See https://prometheus.github.io/client_python/multiprocess/
+    from prometheus_client import (CollectorRegistry, make_asgi_app,
+                                   multiprocess)
+
+    prometheus_multiproc_dir_path = os.getenv("PROMETHEUS_MULTIPROC_DIR", None)
+    if prometheus_multiproc_dir_path is not None:
+        logger.info("vLLM to use %s as PROMETHEUS_MULTIPROC_DIR",
+                    prometheus_multiproc_dir_path)
+        registry = CollectorRegistry()
+        multiprocess.MultiProcessCollector(registry)
+
+        # Add prometheus asgi middleware to route /metrics requests
+        metrics_route = Mount("/metrics", make_asgi_app(registry=registry))
+    else:
+        # Add prometheus asgi middleware to route /metrics requests
+        metrics_route = Mount("/metrics", make_asgi_app())
+
+    # Workaround for 307 Redirect for /metrics
+    metrics_route.path_regex = re.compile("^/metrics(?P<path>.*)$")
+    app.routes.append(metrics_route)
+
+
+def chat(request: Request) -> OpenAIServingChat:
+    return request.app.state.openai_serving_chat
+
+
+def completion(request: Request) -> OpenAIServingCompletion:
+    return request.app.state.openai_serving_completion
+
+
+def tokenization(request: Request) -> OpenAIServingTokenization:
+    return request.app.state.openai_serving_tokenization
+
+
+def embedding(request: Request) -> OpenAIServingEmbedding:
+    return request.app.state.openai_serving_embedding
+
+
+def engine_client(request: Request) -> EngineClient:
+    return request.app.state.engine_client
+
+
+@router.get("/health")
+async def health(raw_request: Request) -> Response:
+    """Health check."""
+    await engine_client(raw_request).check_health()
+    return Response(status_code=200)
+
+
+@router.post("/tokenize")
+async def tokenize(request: TokenizeRequest, raw_request: Request):
+    generator = await tokenization(raw_request).create_tokenize(request)
+    if isinstance(generator, ErrorResponse):
+        return JSONResponse(content=generator.model_dump(),
+                            status_code=generator.code)
+    elif isinstance(generator, TokenizeResponse):
+        return JSONResponse(content=generator.model_dump())
+
+    assert_never(generator)
+
+
+@router.post("/detokenize")
+async def detokenize(request: DetokenizeRequest, raw_request: Request):
+    generator = await tokenization(raw_request).create_detokenize(request)
+    if isinstance(generator, ErrorResponse):
+        return JSONResponse(content=generator.model_dump(),
+                            status_code=generator.code)
+    elif isinstance(generator, DetokenizeResponse):
+        return JSONResponse(content=generator.model_dump())
+
+    assert_never(generator)
+
+
+@router.get("/v1/models")
+async def show_available_models(raw_request: Request):
+    models = await completion(raw_request).show_available_models()
+    return JSONResponse(content=models.model_dump())
+
+
+@router.get("/version")
+async def show_version():
+    ver = {"version": VLLM_VERSION}
+    return JSONResponse(content=ver)
+
+
+@router.post("/v1/chat/completions")
+async def create_chat_completion(request: ChatCompletionRequest,
+                                 raw_request: Request):
+
+    generator = await chat(raw_request).create_chat_completion(
+        request, raw_request)
+
+    if isinstance(generator, ErrorResponse):
+        return JSONResponse(content=generator.model_dump(),
+                            status_code=generator.code)
+
+    elif isinstance(generator, ChatCompletionResponse):
+        return JSONResponse(content=generator.model_dump())
+
+    return StreamingResponse(content=generator, media_type="text/event-stream")
+
+
+@router.post("/v1/completions")
+async def create_completion(request: CompletionRequest, raw_request: Request):
+    generator = await completion(raw_request).create_completion(
+        request, raw_request)
+    if isinstance(generator, ErrorResponse):
+        return JSONResponse(content=generator.model_dump(),
+                            status_code=generator.code)
+    elif isinstance(generator, CompletionResponse):
+        return JSONResponse(content=generator.model_dump())
+
+    return StreamingResponse(content=generator, media_type="text/event-stream")
+
+
+@router.post("/v1/embeddings")
+async def create_embedding(request: EmbeddingRequest, raw_request: Request):
+    generator = await embedding(raw_request).create_embedding(
+        request, raw_request)
+    if isinstance(generator, ErrorResponse):
+        return JSONResponse(content=generator.model_dump(),
+                            status_code=generator.code)
+    elif isinstance(generator, EmbeddingResponse):
+        return JSONResponse(content=generator.model_dump())
+
+    assert_never(generator)
+
+
+if envs.VLLM_TORCH_PROFILER_DIR:
+    logger.warning(
+        "Torch Profiler is enabled in the API server. This should ONLY be "
+        "used for local development!")
+
+    @router.post("/start_profile")
+    async def start_profile(raw_request: Request):
+        logger.info("Starting profiler...")
+        await engine_client(raw_request).start_profile()
+        logger.info("Profiler started.")
+        return Response(status_code=200)
+
+    @router.post("/stop_profile")
+    async def stop_profile(raw_request: Request):
+        logger.info("Stopping profiler...")
+        await engine_client(raw_request).stop_profile()
+        logger.info("Profiler stopped.")
+        return Response(status_code=200)
+
+
+if envs.VLLM_ALLOW_RUNTIME_LORA_UPDATING:
+    logger.warning(
+        "Lora dynamic loading & unloading is enabled in the API server. "
+        "This should ONLY be used for local development!")
+
+    @router.post("/v1/load_lora_adapter")
+    async def load_lora_adapter(request: LoadLoraAdapterRequest,
+                                raw_request: Request):
+        response = await chat(raw_request).load_lora_adapter(request)
+        if isinstance(response, ErrorResponse):
+            return JSONResponse(content=response.model_dump(),
+                                status_code=response.code)
+
+        response = await completion(raw_request).load_lora_adapter(request)
+        if isinstance(response, ErrorResponse):
+            return JSONResponse(content=response.model_dump(),
+                                status_code=response.code)
+
+        return Response(status_code=200, content=response)
+
+    @router.post("/v1/unload_lora_adapter")
+    async def unload_lora_adapter(request: UnloadLoraAdapterRequest,
+                                  raw_request: Request):
+        response = await chat(raw_request).unload_lora_adapter(request)
+        if isinstance(response, ErrorResponse):
+            return JSONResponse(content=response.model_dump(),
+                                status_code=response.code)
+
+        response = await completion(raw_request).unload_lora_adapter(request)
+        if isinstance(response, ErrorResponse):
+            return JSONResponse(content=response.model_dump(),
+                                status_code=response.code)
+
+        return Response(status_code=200, content=response)
+
+
+def build_app(args: Namespace) -> FastAPI:
+    if args.disable_fastapi_docs:
+        app = FastAPI(openapi_url=None,
+                      docs_url=None,
+                      redoc_url=None,
+                      lifespan=lifespan)
+    else:
+        app = FastAPI(lifespan=lifespan)
+    app.include_router(router)
+    app.root_path = args.root_path
+
+    mount_metrics(app)
+
+    app.add_middleware(
+        CORSMiddleware,
+        allow_origins=args.allowed_origins,
+        allow_credentials=args.allow_credentials,
+        allow_methods=args.allowed_methods,
+        allow_headers=args.allowed_headers,
+    )
+
+    @app.exception_handler(RequestValidationError)
+    async def validation_exception_handler(_, exc):
+        chat = app.state.openai_serving_chat
+        err = chat.create_error_response(message=str(exc))
+        return JSONResponse(err.model_dump(),
+                            status_code=HTTPStatus.BAD_REQUEST)
+
+    if token := envs.VLLM_API_KEY or args.api_key:
+
+        @app.middleware("http")
+        async def authentication(request: Request, call_next):
+            root_path = "" if args.root_path is None else args.root_path
+            if request.method == "OPTIONS":
+                return await call_next(request)
+            if not request.url.path.startswith(f"{root_path}/v1"):
+                return await call_next(request)
+            if request.headers.get("Authorization") != "Bearer " + token:
+                return JSONResponse(content={"error": "Unauthorized"},
+                                    status_code=401)
+            return await call_next(request)
+
+    for middleware in args.middleware:
+        module_path, object_name = middleware.rsplit(".", 1)
+        imported = getattr(importlib.import_module(module_path), object_name)
+        if inspect.isclass(imported):
+            app.add_middleware(imported)
+        elif inspect.iscoroutinefunction(imported):
+            app.middleware("http")(imported)
+        else:
+            raise ValueError(f"Invalid middleware {middleware}. "
+                             f"Must be a function or a class.")
+
+    return app
+
+
+def init_app_state(
+    engine_client: EngineClient,
+    model_config: ModelConfig,
+    state: State,
+    args: Namespace,
+) -> None:
+    if args.served_model_name is not None:
+        served_model_names = args.served_model_name
+    else:
+        served_model_names = [args.model]
+
+    if args.disable_log_requests:
+        request_logger = None
+    else:
+        request_logger = RequestLogger(max_log_len=args.max_log_len)
+
+    base_model_paths = [
+        BaseModelPath(name=name, model_path=args.model)
+        for name in served_model_names
+    ]
+
+    state.engine_client = engine_client
+    state.log_stats = not args.disable_log_stats
+
+    state.openai_serving_chat = OpenAIServingChat(
+        engine_client,
+        model_config,
+        base_model_paths,
+        args.response_role,
+        lora_modules=args.lora_modules,
+        prompt_adapters=args.prompt_adapters,
+        request_logger=request_logger,
+        chat_template=args.chat_template,
+        return_tokens_as_token_ids=args.return_tokens_as_token_ids,
+        enable_auto_tools=args.enable_auto_tool_choice,
+        tool_parser=args.tool_call_parser,
+        reasoning_parser=getattr(args, 'reasoning_parser', None))
+    state.openai_serving_completion = OpenAIServingCompletion(
+        engine_client,
+        model_config,
+        base_model_paths,
+        lora_modules=args.lora_modules,
+        prompt_adapters=args.prompt_adapters,
+        request_logger=request_logger,
+        return_tokens_as_token_ids=args.return_tokens_as_token_ids,
+    )
+    state.openai_serving_embedding = OpenAIServingEmbedding(
+        engine_client,
+        model_config,
+        base_model_paths,
+        request_logger=request_logger,
+    )
+    state.openai_serving_tokenization = OpenAIServingTokenization(
+        engine_client,
+        model_config,
+        base_model_paths,
+        lora_modules=args.lora_modules,
+        request_logger=request_logger,
+        chat_template=args.chat_template,
+    )
+
+
+async def run_server(args, **uvicorn_kwargs) -> None:
+    logger.info("vLLM API server version %s", VLLM_VERSION)
+    logger.info("args: %s", args)
+
+    if args.tool_parser_plugin and len(args.tool_parser_plugin) > 3:
+        ToolParserManager.import_tool_parser(args.tool_parser_plugin)
+
+    valide_tool_parses = ToolParserManager.tool_parsers.keys()
+    if args.enable_auto_tool_choice \
+        and args.tool_call_parser not in valide_tool_parses:
+        raise KeyError(f"invalid tool call parser: {args.tool_call_parser} "
+                       f"(chose from {{ {','.join(valide_tool_parses)} }})")
+
+    reasoning_parser = getattr(args, 'reasoning_parser', None)
+    if reasoning_parser:
+        valid_reasoning = ReasoningParserManager.list_registered()
+        if reasoning_parser not in valid_reasoning:
+            raise KeyError(
+                f"invalid reasoning parser: {reasoning_parser} "
+                f"(chose from {{ {','.join(valid_reasoning)} }})")
+
+    # workaround to make sure that we bind the port before the engine is set up.
+    # This avoids race conditions with ray.
+    # see https://github.com/vllm-project/vllm/issues/8204
+    sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+    sock.bind(("", args.port))
+
+    def signal_handler(*_) -> None:
+        # Interrupt server on sigterm while initializing
+        raise KeyboardInterrupt("terminated")
+
+    signal.signal(signal.SIGTERM, signal_handler)
+
+    async with build_async_engine_client(args) as engine_client:
+        app = build_app(args)
+
+        model_config = await engine_client.get_model_config()
+        init_app_state(engine_client, model_config, app.state, args)
+
+        shutdown_task = await serve_http(
+            app,
+            host=args.host,
+            port=args.port,
+            log_level=args.uvicorn_log_level,
+            timeout_keep_alive=TIMEOUT_KEEP_ALIVE,
+            ssl_keyfile=args.ssl_keyfile,
+            ssl_certfile=args.ssl_certfile,
+            ssl_ca_certs=args.ssl_ca_certs,
+            ssl_cert_reqs=args.ssl_cert_reqs,
+            fd=sock.fileno(),
+            **uvicorn_kwargs,
+        )
+
+    # NB: Await server shutdown only after the backend context is exited
+    await shutdown_task
+
+
+if __name__ == "__main__":
+    # NOTE(simon):
+    # This section should be in sync with vllm/scripts.py for CLI entrypoints.
+    parser = FlexibleArgumentParser(
+        description="vLLM OpenAI-Compatible RESTful API server.")
+    parser = make_arg_parser(parser)
+    args = parser.parse_args()
+    validate_parsed_serve_args(args)
+
+    uvloop.run(run_server(args))

qwen3_6_scripts/chat_utils.py

@@ -0,0 +1,603 @@
+import asyncio
+import codecs
+import json
+from abc import ABC, abstractmethod
+from collections import defaultdict
+from functools import lru_cache, partial
+from pathlib import Path
+from typing import (Any, Awaitable, Dict, Generic, Iterable, List, Literal,
+                    Mapping, Optional, Tuple, TypeVar, Union, cast)
+
+# yapf conflicts with isort for this block
+# yapf: disable
+from openai.types.chat import (ChatCompletionAssistantMessageParam,
+                               ChatCompletionContentPartImageParam)
+from openai.types.chat import (
+    ChatCompletionContentPartParam as OpenAIChatCompletionContentPartParam)
+from openai.types.chat import (ChatCompletionContentPartRefusalParam,
+                               ChatCompletionContentPartTextParam)
+from openai.types.chat import (
+    ChatCompletionMessageParam as OpenAIChatCompletionMessageParam)
+from openai.types.chat import (ChatCompletionMessageToolCallParam,
+                               ChatCompletionToolMessageParam)
+# yapf: enable
+# pydantic needs the TypedDict from typing_extensions
+from pydantic import ConfigDict
+from transformers import PreTrainedTokenizer, PreTrainedTokenizerFast
+from typing_extensions import Required, TypeAlias, TypedDict
+
+from vllm.config import ModelConfig
+from vllm.logger import init_logger
+from vllm.multimodal import MultiModalDataDict
+from vllm.multimodal.utils import (async_get_and_parse_audio,
+                                   async_get_and_parse_image,
+                                   get_and_parse_audio, get_and_parse_image)
+from vllm.transformers_utils.tokenizer import AnyTokenizer, MistralTokenizer
+
+logger = init_logger(__name__)
+
+
+class AudioURL(TypedDict, total=False):
+    url: Required[str]
+    """
+    Either a URL of the audio or a data URL with base64 encoded audio data.
+    """
+
+
+class ChatCompletionContentPartAudioParam(TypedDict, total=False):
+    audio_url: Required[AudioURL]
+
+    type: Required[Literal["audio_url"]]
+    """The type of the content part."""
+
+
+class CustomChatCompletionContentPartParam(TypedDict, total=False):
+    __pydantic_config__ = ConfigDict(extra="allow")  # type: ignore
+
+    type: Required[str]
+    """The type of the content part."""
+
+
+ChatCompletionContentPartParam: TypeAlias = Union[
+    OpenAIChatCompletionContentPartParam, ChatCompletionContentPartAudioParam,
+    ChatCompletionContentPartRefusalParam,
+    CustomChatCompletionContentPartParam]
+
+
+class CustomChatCompletionMessageParam(TypedDict, total=False):
+    """Enables custom roles in the Chat Completion API."""
+    role: Required[str]
+    """The role of the message's author."""
+
+    content: Union[str, List[ChatCompletionContentPartParam]]
+    """The contents of the message."""
+
+    name: str
+    """An optional name for the participant.
+
+    Provides the model information to differentiate between participants of the
+    same role.
+    """
+
+    tool_call_id: Optional[str]
+    """Tool call that this message is responding to."""
+
+    tool_calls: Optional[Iterable[ChatCompletionMessageToolCallParam]]
+    """The tool calls generated by the model, such as function calls."""
+
+    reasoning_content: Optional[str]
+    """Reasoning / thinking content for assistant messages (vLLM extension).
+    When present in a previous assistant turn, it is rendered as
+    <think>...</think> before the main content so the model sees its own
+    chain-of-thought in subsequent turns."""
+
+
+ChatCompletionMessageParam = Union[OpenAIChatCompletionMessageParam,
+                                   CustomChatCompletionMessageParam]
+
+
+# TODO: Make fields ReadOnly once mypy supports it
+class ConversationMessage(TypedDict, total=False):
+    role: Required[str]
+    """The role of the message's author."""
+
+    content: Optional[str]
+    """The contents of the message"""
+
+    tool_call_id: Optional[str]
+    """Tool call that this message is responding to."""
+
+    name: Optional[str]
+    """The name of the function to call"""
+
+    tool_calls: Optional[Iterable[ChatCompletionMessageToolCallParam]]
+    """The tool calls generated by the model, such as function calls."""
+
+    reasoning_content: Optional[str]
+    """Reasoning / thinking content for assistant messages.
+    Passed directly to the chat template (Qwen3 reads message.reasoning_content
+    natively) instead of being manually wrapped in <think>...</think>."""
+
+
+ModalityStr = Literal["image", "audio", "video"]
+_T = TypeVar("_T")
+
+
+class BaseMultiModalItemTracker(ABC, Generic[_T]):
+    """
+    Tracks multi-modal items in a given request and ensures that the number
+    of multi-modal items in a given request does not exceed the configured
+    maximum per prompt.
+    """
+
+    def __init__(self, model_config: ModelConfig, tokenizer: AnyTokenizer):
+        super().__init__()
+
+        self._model_config = model_config
+        self._tokenizer = tokenizer
+        self._allowed_items = (model_config.multimodal_config.limit_per_prompt
+                               if model_config.multimodal_config else {})
+        self._consumed_items = {k: 0 for k in self._allowed_items}
+
+        self._items: List[_T] = []
+
+    @staticmethod
+    @lru_cache(maxsize=None)
+    def _cached_token_str(tokenizer: AnyTokenizer, token_index: int) -> str:
+        return tokenizer.decode(token_index)
+
+    def _placeholder_str(self, modality: ModalityStr,
+                         current_count: int) -> Optional[str]:
+        # TODO: Let user specify how to insert image tokens into prompt
+        # (similar to chat template)
+        hf_config = self._model_config.hf_config
+        model_type = hf_config.model_type
+
+        if modality == "image":
+            if model_type == "phi3_v":
+                # Workaround since this token is not defined in the tokenizer
+                return f"<|image_{current_count}|>"
+            if model_type == "minicpmv":
+                return "(<image>./</image>)"
+            if model_type in ("blip-2", "chatglm", "fuyu", "paligemma",
+                              "pixtral"):
+                # These models do not use image tokens in the prompt
+                return None
+            if model_type == "qwen":
+                return f"Picture {current_count}: <img></img>"
+            if model_type.startswith("llava"):
+                return self._cached_token_str(self._tokenizer,
+                                              hf_config.image_token_index)
+            if model_type in ("chameleon", "internvl_chat", "NVLM_D"):
+                return "<image>"
+            if model_type == "mllama":
+                return "<|image|>"
+            if model_type in ("qwen2_vl", "qwen2_5_vl",
+                              "qwen3_5", "qwen3_5_moe"):
+                return "<|vision_start|><|image_pad|><|vision_end|>"
+            if model_type == "molmo":
+                return ""
+
+            raise TypeError(f"Unknown model type: {model_type}")
+        elif modality == "audio":
+            if model_type == "ultravox":
+                return "<|reserved_special_token_0|>"
+            raise TypeError(f"Unknown model type: {model_type}")
+        elif modality == "video":
+            if model_type in ("qwen2_vl", "qwen2_5_vl",
+                              "qwen3_5", "qwen3_5_moe"):
+                return "<|vision_start|><|video_pad|><|vision_end|>"
+            raise TypeError(f"Unknown model type: {model_type}")
+        else:
+            raise TypeError(f"Unknown modality: {modality}")
+
+    @staticmethod
+    def _combine(items: List[MultiModalDataDict]) -> MultiModalDataDict:
+        mm_lists: Mapping[str, List[object]] = defaultdict(list)
+
+        # Merge all the multi-modal items
+        for single_mm_data in items:
+            for mm_key, mm_item in single_mm_data.items():
+                if isinstance(mm_item, list):
+                    mm_lists[mm_key].extend(mm_item)
+                else:
+                    mm_lists[mm_key].append(mm_item)
+
+        # Unpack any single item lists for models that don't expect multiple.
+        return {
+            mm_key: mm_list[0] if len(mm_list) == 1 else mm_list
+            for mm_key, mm_list in mm_lists.items()
+        }
+
+    def add(self, modality: ModalityStr, item: _T) -> Optional[str]:
+        """
+        Add a multi-modal item to the current prompt and returns the
+        placeholder string to use, if any.
+        """
+        allowed_count = self._allowed_items.get(modality, 1)
+        current_count = self._consumed_items.get(modality, 0) + 1
+        if current_count > allowed_count:
+            raise ValueError(
+                f"At most {allowed_count} {modality}(s) may be provided in "
+                "one request.")
+
+        self._consumed_items[modality] = current_count
+        self._items.append(item)
+
+        return self._placeholder_str(modality, current_count)
+
+    @abstractmethod
+    def create_parser(self) -> "BaseMultiModalContentParser":
+        raise NotImplementedError
+
+
+class MultiModalItemTracker(BaseMultiModalItemTracker[MultiModalDataDict]):
+
+    def all_mm_data(self) -> Optional[MultiModalDataDict]:
+        return self._combine(self._items) if self._items else None
+
+    def create_parser(self) -> "BaseMultiModalContentParser":
+        return MultiModalContentParser(self)
+
+
+class AsyncMultiModalItemTracker(
+        BaseMultiModalItemTracker[Awaitable[MultiModalDataDict]]):
+
+    async def all_mm_data(self) -> Optional[MultiModalDataDict]:
+        if self._items:
+            items = await asyncio.gather(*self._items)
+            return self._combine(items)
+
+        return None
+
+    def create_parser(self) -> "BaseMultiModalContentParser":
+        return AsyncMultiModalContentParser(self)
+
+
+class BaseMultiModalContentParser(ABC):
+
+    def __init__(self) -> None:
+        super().__init__()
+
+        # multimodal placeholder_string : count
+        self._placeholder_counts: Dict[str, int] = defaultdict(lambda: 0)
+
+    def _add_placeholder(self, placeholder: Optional[str]):
+        if placeholder:
+            self._placeholder_counts[placeholder] += 1
+
+    def mm_placeholder_counts(self) -> Dict[str, int]:
+        return dict(self._placeholder_counts)
+
+    @abstractmethod
+    def parse_image(self, image_url: str) -> None:
+        raise NotImplementedError
+
+    @abstractmethod
+    def parse_audio(self, audio_url: str) -> None:
+        raise NotImplementedError
+
+
+class MultiModalContentParser(BaseMultiModalContentParser):
+
+    def __init__(self, tracker: MultiModalItemTracker) -> None:
+        super().__init__()
+
+        self._tracker = tracker
+
+    def parse_image(self, image_url: str) -> None:
+        image = get_and_parse_image(image_url)
+
+        placeholder = self._tracker.add("image", image)
+        self._add_placeholder(placeholder)
+
+    def parse_audio(self, audio_url: str) -> None:
+        audio = get_and_parse_audio(audio_url)
+
+        placeholder = self._tracker.add("audio", audio)
+        self._add_placeholder(placeholder)
+
+
+class AsyncMultiModalContentParser(BaseMultiModalContentParser):
+
+    def __init__(self, tracker: AsyncMultiModalItemTracker) -> None:
+        super().__init__()
+
+        self._tracker = tracker
+
+    def parse_image(self, image_url: str) -> None:
+        image_coro = async_get_and_parse_image(image_url)
+
+        placeholder = self._tracker.add("image", image_coro)
+        self._add_placeholder(placeholder)
+
+    def parse_audio(self, audio_url: str) -> None:
+        audio_coro = async_get_and_parse_audio(audio_url)
+
+        placeholder = self._tracker.add("audio", audio_coro)
+        self._add_placeholder(placeholder)
+
+
+def validate_chat_template(chat_template: Optional[Union[Path, str]]):
+    """Raises if the provided chat template appears invalid."""
+    if chat_template is None:
+        return
+
+    elif isinstance(chat_template, Path) and not chat_template.exists():
+        raise FileNotFoundError(
+            "the supplied chat template path doesn't exist")
+
+    elif isinstance(chat_template, str):
+        JINJA_CHARS = "{}\n"
+        if not any(c in chat_template
+                   for c in JINJA_CHARS) and not Path(chat_template).exists():
+            raise ValueError(
+                f"The supplied chat template string ({chat_template}) "
+                f"appears path-like, but doesn't exist!")
+
+    else:
+        raise TypeError(
+            f"{type(chat_template)} is not a valid chat template type")
+
+
+def load_chat_template(
+        chat_template: Optional[Union[Path, str]]) -> Optional[str]:
+    if chat_template is None:
+        return None
+    try:
+        with open(chat_template, "r") as f:
+            resolved_chat_template = f.read()
+    except OSError as e:
+        if isinstance(chat_template, Path):
+            raise
+
+        JINJA_CHARS = "{}\n"
+        if not any(c in chat_template for c in JINJA_CHARS):
+            msg = (f"The supplied chat template ({chat_template}) "
+                   f"looks like a file path, but it failed to be "
+                   f"opened. Reason: {e}")
+            raise ValueError(msg) from e
+
+        # If opening a file fails, set chat template to be args to
+        # ensure we decode so our escape are interpreted correctly
+        resolved_chat_template = codecs.decode(chat_template, "unicode_escape")
+
+    logger.info("Using supplied chat template:\n%s", resolved_chat_template)
+    return resolved_chat_template
+
+
+# TODO: Let user specify how to insert multimodal tokens into prompt
+# (similar to chat template)
+def _get_full_multimodal_text_prompt(placeholder_counts: Dict[str, int],
+                                     text_prompt: str) -> str:
+    """Combine multimodal prompts for a multimodal language model."""
+
+    # Look through the text prompt to check for missing placeholders
+    missing_placeholders: List[str] = []
+    for placeholder in placeholder_counts:
+
+        # For any existing placeholder in the text prompt, we leave it as is
+        placeholder_counts[placeholder] -= text_prompt.count(placeholder)
+
+        if placeholder_counts[placeholder] < 0:
+            raise ValueError(
+                f"Found more '{placeholder}' placeholders in input prompt than "
+                "actual multimodal data items.")
+
+        missing_placeholders.extend([placeholder] *
+                                    placeholder_counts[placeholder])
+
+    # NOTE: For now we always add missing placeholders at the front of
+    # the prompt. This may change to be customizable in the future.
+    return "\n".join(missing_placeholders + [text_prompt])
+
+
+# No need to validate using Pydantic again
+_TextParser = partial(cast, ChatCompletionContentPartTextParam)
+_ImageParser = partial(cast, ChatCompletionContentPartImageParam)
+_AudioParser = partial(cast, ChatCompletionContentPartAudioParam)
+_RefusalParser = partial(cast, ChatCompletionContentPartRefusalParam)
+MODEL_KEEP_MULTI_MODAL_CONTENT = {'mllama'}
+
+
+def _parse_chat_message_content_parts(
+    role: str,
+    parts: Iterable[ChatCompletionContentPartParam],
+    mm_tracker: BaseMultiModalItemTracker,
+) -> List[ConversationMessage]:
+    texts: List[str] = []
+
+    mm_parser = mm_tracker.create_parser()
+    keep_multimodal_content = \
+        mm_tracker._model_config.hf_config.model_type in \
+            MODEL_KEEP_MULTI_MODAL_CONTENT
+
+    has_image = False
+    for part in parts:
+        part_type = part["type"]
+        if part_type == "text":
+            text = _TextParser(part)["text"]
+            texts.append(text)
+        elif part_type == "image_url":
+            image_url = _ImageParser(part)["image_url"]
+
+            if image_url.get("detail", "auto") != "auto":
+                logger.warning(
+                    "'image_url.detail' is currently not supported and "
+                    "will be ignored.")
+
+            mm_parser.parse_image(image_url["url"])
+            has_image = True
+        elif part_type == "audio_url":
+            audio_url = _AudioParser(part)["audio_url"]
+
+            mm_parser.parse_audio(audio_url["url"])
+        elif part_type == "refusal":
+            text = _RefusalParser(part)["refusal"]
+            texts.append(text)
+        else:
+            raise NotImplementedError(f"Unknown part type: {part_type}")
+
+    text_prompt = "\n".join(texts)
+    if keep_multimodal_content:
+        text_prompt = "\n".join(texts)
+        role_content = [{'type': 'text', 'text': text_prompt}]
+
+        if has_image:
+            role_content = [{'type': 'image'}] + role_content
+        return [ConversationMessage(role=role,
+                                    content=role_content)]  # type: ignore
+    else:
+        mm_placeholder_counts = mm_parser.mm_placeholder_counts()
+        if mm_placeholder_counts:
+            text_prompt = _get_full_multimodal_text_prompt(
+                mm_placeholder_counts, text_prompt)
+        return [ConversationMessage(role=role, content=text_prompt)]
+
+
+# No need to validate using Pydantic again
+_AssistantParser = partial(cast, ChatCompletionAssistantMessageParam)
+_ToolParser = partial(cast, ChatCompletionToolMessageParam)
+
+
+def _parse_chat_message_content(
+    message: ChatCompletionMessageParam,
+    mm_tracker: BaseMultiModalItemTracker,
+) -> List[ConversationMessage]:
+    role = message["role"]
+    content = message.get("content")
+
+    if content is None:
+        content = []
+    elif isinstance(content, str):
+        content = [
+            ChatCompletionContentPartTextParam(type="text", text=content)
+        ]
+
+    result = _parse_chat_message_content_parts(
+        role,
+        content,  # type: ignore
+        mm_tracker,
+    )
+
+    for result_msg in result:
+        if role == 'assistant':
+            parsed_msg = _AssistantParser(message)
+
+            if "tool_calls" in parsed_msg:
+                result_msg["tool_calls"] = list(parsed_msg["tool_calls"])
+
+            # Pass reasoning content as a dedicated field so the chat template
+            # can render it natively (Qwen3: message.reasoning_content branch).
+            # Accept both "reasoning" (new vllm) and "reasoning_content" (ours).
+            reasoning = (message.get("reasoning")  # type: ignore[arg-type]
+                         or message.get("reasoning_content"))  # type: ignore[arg-type]
+            if reasoning and isinstance(reasoning, str):
+                result_msg["reasoning_content"] = reasoning
+
+        elif role == "tool":
+            parsed_msg = _ToolParser(message)
+            if "tool_call_id" in parsed_msg:
+                result_msg["tool_call_id"] = parsed_msg["tool_call_id"]
+
+        if "name" in message and isinstance(message["name"], str):
+            result_msg["name"] = message["name"]
+
+    return result
+
+
+def _postprocess_messages(messages: List[ConversationMessage]) -> None:
+    # per the Transformers docs & maintainers, tool call arguments in
+    # assistant-role messages with tool_calls need to be dicts not JSON str -
+    # this is how tool-use chat templates will expect them moving forwards
+    # so, for messages that have tool_calls, parse the string (which we get
+    # from openAI format) to dict
+    for message in messages:
+        if (message["role"] == "assistant" and "tool_calls" in message
+                and isinstance(message["tool_calls"], list)):
+
+            for item in message["tool_calls"]:
+                item["function"]["arguments"] = json.loads(
+                    item["function"]["arguments"])
+
+
+def parse_chat_messages(
+    messages: List[ChatCompletionMessageParam],
+    model_config: ModelConfig,
+    tokenizer: AnyTokenizer,
+) -> Tuple[List[ConversationMessage], Optional[MultiModalDataDict]]:
+    conversation: List[ConversationMessage] = []
+    mm_tracker = MultiModalItemTracker(model_config, tokenizer)
+
+    for msg in messages:
+        sub_messages = _parse_chat_message_content(msg, mm_tracker)
+
+        conversation.extend(sub_messages)
+
+    _postprocess_messages(conversation)
+
+    return conversation, mm_tracker.all_mm_data()
+
+
+def parse_chat_messages_futures(
+    messages: List[ChatCompletionMessageParam],
+    model_config: ModelConfig,
+    tokenizer: AnyTokenizer,
+) -> Tuple[List[ConversationMessage], Awaitable[Optional[MultiModalDataDict]]]:
+    conversation: List[ConversationMessage] = []
+    mm_tracker = AsyncMultiModalItemTracker(model_config, tokenizer)
+
+    for msg in messages:
+        sub_messages = _parse_chat_message_content(msg, mm_tracker)
+
+        conversation.extend(sub_messages)
+
+    _postprocess_messages(conversation)
+
+    return conversation, mm_tracker.all_mm_data()
+
+
+def apply_hf_chat_template(
+    tokenizer: Union[PreTrainedTokenizer, PreTrainedTokenizerFast],
+    conversation: List[ConversationMessage],
+    chat_template: Optional[str],
+    *,
+    tokenize: bool = False,  # Different from HF's default
+    **kwargs: Any,
+) -> str:
+    if chat_template is None and tokenizer.chat_template is None:
+        raise ValueError(
+            "As of transformers v4.44, default chat template is no longer "
+            "allowed, so you must provide a chat template if the tokenizer "
+            "does not define one.")
+
+    return tokenizer.apply_chat_template(
+        conversation=conversation,  # type: ignore[arg-type]
+        chat_template=chat_template,
+        tokenize=tokenize,
+        **kwargs,
+    )
+
+
+def apply_mistral_chat_template(
+    tokenizer: MistralTokenizer,
+    messages: List[ChatCompletionMessageParam],
+    chat_template: Optional[str] = None,
+    **kwargs: Any,
+) -> List[int]:
+    if chat_template is not None:
+        logger.warning(
+            "'chat_template' cannot be overridden for mistral tokenizer.")
+    if "add_generation_prompt" in kwargs:
+        logger.warning(
+            "'add_generation_prompt' is not supported for mistral tokenizer, "
+            "so it will be ignored.")
+    if "continue_final_message" in kwargs:
+        logger.warning(
+            "'continue_final_message' is not supported for mistral tokenizer, "
+            "so it will be ignored.")
+
+    return tokenizer.apply_chat_template(
+        messages=messages,
+        **kwargs,
+    )

qwen3_6_scripts/cli_args.py

@@ -0,0 +1,261 @@
+"""
+This file contains the command line arguments for the vLLM's
+OpenAI-compatible server. It is kept in a separate file for documentation
+purposes.
+"""
+
+import argparse
+import json
+import ssl
+from typing import List, Optional, Sequence, Union
+
+from vllm.engine.arg_utils import AsyncEngineArgs, nullable_str
+from vllm.entrypoints.chat_utils import validate_chat_template
+from vllm.entrypoints.openai.serving_engine import (LoRAModulePath,
+                                                    PromptAdapterPath)
+from vllm.entrypoints.openai.tool_parsers import ToolParserManager
+from vllm.utils import FlexibleArgumentParser
+
+
+class LoRAParserAction(argparse.Action):
+
+    def __call__(
+        self,
+        parser: argparse.ArgumentParser,
+        namespace: argparse.Namespace,
+        values: Optional[Union[str, Sequence[str]]],
+        option_string: Optional[str] = None,
+    ):
+        if values is None:
+            values = []
+        if isinstance(values, str):
+            raise TypeError("Expected values to be a list")
+
+        lora_list: List[LoRAModulePath] = []
+        for item in values:
+            if item in [None, '']:  # Skip if item is None or empty string
+                continue
+            if '=' in item and ',' not in item:  # Old format: name=path
+                name, path = item.split('=')
+                lora_list.append(LoRAModulePath(name, path))
+            else:  # Assume JSON format
+                try:
+                    lora_dict = json.loads(item)
+                    lora = LoRAModulePath(**lora_dict)
+                    lora_list.append(lora)
+                except json.JSONDecodeError:
+                    parser.error(
+                        f"Invalid JSON format for --lora-modules: {item}")
+                except TypeError as e:
+                    parser.error(
+                        f"Invalid fields for --lora-modules: {item} - {str(e)}"
+                    )
+        setattr(namespace, self.dest, lora_list)
+
+
+class PromptAdapterParserAction(argparse.Action):
+
+    def __call__(
+        self,
+        parser: argparse.ArgumentParser,
+        namespace: argparse.Namespace,
+        values: Optional[Union[str, Sequence[str]]],
+        option_string: Optional[str] = None,
+    ):
+        if values is None:
+            values = []
+        if isinstance(values, str):
+            raise TypeError("Expected values to be a list")
+
+        adapter_list: List[PromptAdapterPath] = []
+        for item in values:
+            name, path = item.split('=')
+            adapter_list.append(PromptAdapterPath(name, path))
+        setattr(namespace, self.dest, adapter_list)
+
+
+def make_arg_parser(parser: FlexibleArgumentParser) -> FlexibleArgumentParser:
+    parser.add_argument("--host",
+                        type=nullable_str,
+                        default=None,
+                        help="host name")
+    parser.add_argument("--port", type=int, default=8000, help="port number")
+    parser.add_argument(
+        "--uvicorn-log-level",
+        type=str,
+        default="info",
+        choices=['debug', 'info', 'warning', 'error', 'critical', 'trace'],
+        help="log level for uvicorn")
+    parser.add_argument("--allow-credentials",
+                        action="store_true",
+                        help="allow credentials")
+    parser.add_argument("--allowed-origins",
+                        type=json.loads,
+                        default=["*"],
+                        help="allowed origins")
+    parser.add_argument("--allowed-methods",
+                        type=json.loads,
+                        default=["*"],
+                        help="allowed methods")
+    parser.add_argument("--allowed-headers",
+                        type=json.loads,
+                        default=["*"],
+                        help="allowed headers")
+    parser.add_argument("--api-key",
+                        type=nullable_str,
+                        default=None,
+                        help="If provided, the server will require this key "
+                        "to be presented in the header.")
+    parser.add_argument(
+        "--lora-modules",
+        type=nullable_str,
+        default=None,
+        nargs='+',
+        action=LoRAParserAction,
+        help="LoRA module configurations in either 'name=path' format"
+        "or JSON format. "
+        "Example (old format): 'name=path' "
+        "Example (new format): "
+        "'{\"name\": \"name\", \"local_path\": \"path\", "
+        "\"base_model_name\": \"id\"}'")
+    parser.add_argument(
+        "--prompt-adapters",
+        type=nullable_str,
+        default=None,
+        nargs='+',
+        action=PromptAdapterParserAction,
+        help="Prompt adapter configurations in the format name=path. "
+        "Multiple adapters can be specified.")
+    parser.add_argument("--chat-template",
+                        type=nullable_str,
+                        default=None,
+                        help="The file path to the chat template, "
+                        "or the template in single-line form "
+                        "for the specified model")
+    parser.add_argument("--response-role",
+                        type=nullable_str,
+                        default="assistant",
+                        help="The role name to return if "
+                        "`request.add_generation_prompt=true`.")
+    parser.add_argument("--ssl-keyfile",
+                        type=nullable_str,
+                        default=None,
+                        help="The file path to the SSL key file")
+    parser.add_argument("--ssl-certfile",
+                        type=nullable_str,
+                        default=None,
+                        help="The file path to the SSL cert file")
+    parser.add_argument("--ssl-ca-certs",
+                        type=nullable_str,
+                        default=None,
+                        help="The CA certificates file")
+    parser.add_argument(
+        "--ssl-cert-reqs",
+        type=int,
+        default=int(ssl.CERT_NONE),
+        help="Whether client certificate is required (see stdlib ssl module's)"
+    )
+    parser.add_argument(
+        "--root-path",
+        type=nullable_str,
+        default=None,
+        help="FastAPI root_path when app is behind a path based routing proxy")
+    parser.add_argument(
+        "--middleware",
+        type=nullable_str,
+        action="append",
+        default=[],
+        help="Additional ASGI middleware to apply to the app. "
+        "We accept multiple --middleware arguments. "
+        "The value should be an import path. "
+        "If a function is provided, vLLM will add it to the server "
+        "using @app.middleware('http'). "
+        "If a class is provided, vLLM will add it to the server "
+        "using app.add_middleware(). ")
+    parser.add_argument(
+        "--return-tokens-as-token-ids",
+        action="store_true",
+        help="When --max-logprobs is specified, represents single tokens as "
+        "strings of the form 'token_id:{token_id}' so that tokens that "
+        "are not JSON-encodable can be identified.")
+    parser.add_argument(
+        "--disable-frontend-multiprocessing",
+        action="store_true",
+        help="If specified, will run the OpenAI frontend server in the same "
+        "process as the model serving engine.")
+
+    parser.add_argument(
+        "--enable-auto-tool-choice",
+        action="store_true",
+        default=False,
+        help=
+        "Enable auto tool choice for supported models. Use --tool-call-parser"
+        "to specify which parser to use")
+
+    valid_tool_parsers = ToolParserManager.tool_parsers.keys()
+    parser.add_argument(
+        "--tool-call-parser",
+        type=str,
+        metavar="{" + ",".join(valid_tool_parsers) + "} or name registered in "
+        "--tool-parser-plugin",
+        default=None,
+        help=
+        "Select the tool call parser depending on the model that you're using."
+        " This is used to parse the model-generated tool call into OpenAI API "
+        "format. Required for --enable-auto-tool-choice.")
+
+    parser.add_argument(
+        "--tool-parser-plugin",
+        type=str,
+        default="",
+        help=
+        "Special the tool parser plugin write to parse the model-generated tool"
+        " into OpenAI API format, the name register in this plugin can be used "
+        "in --tool-call-parser.")
+
+    parser.add_argument(
+        "--reasoning-parser",
+        type=str,
+        default=None,
+        help=
+        "Select the reasoning parser to split <think>...</think> content into "
+        "reasoning_content vs content in the response. "
+        "Supported: qwen3")
+
+    parser = AsyncEngineArgs.add_cli_args(parser)
+
+    parser.add_argument('--max-log-len',
+                        type=int,
+                        default=None,
+                        help='Max number of prompt characters or prompt '
+                        'ID numbers being printed in log.'
+                        '\n\nDefault: Unlimited')
+
+    parser.add_argument(
+        "--disable-fastapi-docs",
+        action='store_true',
+        default=False,
+        help="Disable FastAPI's OpenAPI schema, Swagger UI, and ReDoc endpoint"
+    )
+
+    return parser
+
+
+def validate_parsed_serve_args(args: argparse.Namespace):
+    """Quick checks for model serve args that raise prior to loading."""
+    if hasattr(args, "subparser") and args.subparser != "serve":
+        return
+
+    # Ensure that the chat template is valid; raises if it likely isn't
+    validate_chat_template(args.chat_template)
+
+    # Enable auto tool needs a tool call parser to be valid
+    if args.enable_auto_tool_choice and not args.tool_call_parser:
+        raise TypeError("Error: --enable-auto-tool-choice requires "
+                        "--tool-call-parser")
+
+
+def create_parser_for_docs() -> FlexibleArgumentParser:
+    parser_for_docs = FlexibleArgumentParser(
+        prog="-m vllm.entrypoints.openai.api_server")
+    return make_arg_parser(parser_for_docs)

qwen3_6_scripts/flash_qla_sm70/build/.ninja_log

@@ -0,0 +1,5 @@
+# ninja log v5
+0	61739	1786467036068204659	gdn_forward.cuda.o	4fbd18c8f06e5181
+61739	62033	1786467036388208334	flash_qla_sm70_gdn_strided.so	a5d04d69a8ccfcee
+0	60985	1786469746403441679	gdn_forward.cuda.o	15f5cb32976bd0b3
+60985	61271	1786469746711445255	flash_qla_sm70_gdn_strided.so	a5d04d69a8ccfcee

qwen3_6_scripts/flash_qla_sm70/build/build.ninja

@@ -0,0 +1,31 @@
+ninja_required_version = 1.3
+cxx = c++
+nvcc = /usr/local/corex/bin/clang++
+
+cflags = -DTORCH_EXTENSION_NAME=flash_qla_sm70_gdn_strided -DTORCH_API_INCLUDE_EXTENSION_H -DPYBIND11_COMPILER_TYPE=\"_gcc\" -DPYBIND11_STDLIB=\"_libstdcpp\" -DPYBIND11_BUILD_ABI=\"_cxxabi1011\" -isystem /usr/local/corex/lib64/python3/dist-packages/torch/include -isystem /usr/local/corex/lib64/python3/dist-packages/torch/include/torch/csrc/api/include -isystem /usr/local/corex/lib64/python3/dist-packages/torch/include/TH -isystem /usr/local/corex/lib64/python3/dist-packages/torch/include/THC -isystem /usr/local/corex/include -isystem /usr/local/include/python3.10 -D_GLIBCXX_USE_CXX11_ABI=0 -fPIC -std=c++17 -O3
+post_cflags = 
+cuda_cflags = -DTORCH_EXTENSION_NAME=flash_qla_sm70_gdn_strided -DTORCH_API_INCLUDE_EXTENSION_H -DPYBIND11_COMPILER_TYPE=\"_gcc\" -DPYBIND11_STDLIB=\"_libstdcpp\" -DPYBIND11_BUILD_ABI=\"_cxxabi1011\" -isystem /usr/local/corex/lib64/python3/dist-packages/torch/include -isystem /usr/local/corex/lib64/python3/dist-packages/torch/include/torch/csrc/api/include -isystem /usr/local/corex/lib64/python3/dist-packages/torch/include/TH -isystem /usr/local/corex/lib64/python3/dist-packages/torch/include/THC -isystem /usr/local/corex/include -isystem /usr/local/include/python3.10 -D_GLIBCXX_USE_CXX11_ABI=0 -D__CUDA_NO_HALF_OPERATORS__ -D__CUDA_NO_HALF_CONVERSIONS__ -D__CUDA_NO_BFLOAT16_CONVERSIONS__ -D__CUDA_NO_HALF2_OPERATORS__ -D__ILUVATAR__ -D__ILUVATAR_WORKAROUND__ -D__ILUVATAR_DIAG__ -cl-single-precision-constant -fPIC -mllvm --bonus-inst-threshold=0 -O3 --cuda-gpu-arch=ivcore10 --cuda-path=/usr/local/corex -std=c++17
+cuda_post_cflags = 
+cuda_dlink_post_cflags = 
+ldflags = -shared -L/usr/local/corex/lib64/python3/dist-packages/torch/lib -lc10 -lc10_cuda -ltorch_cpu -ltorch_cuda -ltorch -ltorch_python -L/usr/local/corex/lib64 -lcudart
+
+rule compile
+  command = $cxx -MMD -MF $out.d $cflags -c $in -o $out $post_cflags
+  depfile = $out.d
+  deps = gcc
+
+rule cuda_compile
+  command = $nvcc  $cuda_cflags -c $in -o $out $cuda_post_cflags
+
+
+
+rule link
+  command = $cxx $in $ldflags -o $out
+
+build gdn_forward.cuda.o: cuda_compile /workspace/qwen3_6_scripts/flash_qla_sm70/csrc/gdn_forward.cu
+
+
+
+build flash_qla_sm70_gdn_strided.so: link gdn_forward.cuda.o
+
+default flash_qla_sm70_gdn_strided.so

qwen3_6_scripts/logits_processor.py

@@ -0,0 +1,158 @@
+"""A layer that compute logits from hidden_stats."""
+import inspect
+from typing import Optional
+
+import torch
+import torch.nn as nn
+
+from vllm.distributed import (tensor_model_parallel_all_gather,
+                              tensor_model_parallel_gather)
+from vllm.model_executor.layers.vocab_parallel_embedding import (
+    VocabParallelEmbedding)
+from vllm.model_executor.sampling_metadata import SamplingMetadata
+from vllm.platforms import current_platform
+
+
+class LogitsProcessor(nn.Module):
+    """Process logits and apply logits processors from sampling metadata.
+
+    This layer does the following:
+    1. Gather logits from model hidden_states.
+    2. Scale logits if needed.
+    3. Apply logits processors (if any).
+    """
+
+    def __init__(self,
+                 vocab_size: int,
+                 org_vocab_size: Optional[int] = None,
+                 scale: float = 1.0,
+                 logits_as_input: bool = False,
+                 soft_cap: Optional[float] = None) -> None:
+        """
+        Args:
+            scale: A scaling factor to apply to the logits.
+        """
+        super().__init__()
+        self.scale = scale
+        self.vocab_size = vocab_size
+        # Whether the input is logits (default is hidden states).
+        self.logits_as_input = logits_as_input
+        # original vocabulary size (without LoRA).
+        self.org_vocab_size = org_vocab_size or vocab_size
+        # Soft cap the logits. Used in Gemma 2.
+        self.soft_cap = soft_cap
+        # Whether to use gather or all-gather to gather the logits.
+        self.use_gather = not current_platform.is_tpu()
+
+    def forward(
+        self,
+        lm_head: VocabParallelEmbedding,
+        hidden_states: torch.Tensor,
+        sampling_metadata: SamplingMetadata,
+        embedding_bias: Optional[torch.Tensor] = None,
+    ) -> Optional[torch.Tensor]:
+        if self.logits_as_input:
+            logits = hidden_states
+        else:
+            hidden_states = _prune_hidden_states(hidden_states,
+                                                 sampling_metadata)
+
+            # Get the logits for the next tokens.
+            if hidden_states.shape[0] > 0:
+                logits = self._get_logits(hidden_states, lm_head, embedding_bias)
+            else:
+                logits = torch.empty([0, lm_head.weight.shape[0]], device=hidden_states.device, dtype=hidden_states.dtype)
+        if logits is not None:
+            if self.soft_cap is not None:
+                logits = logits / self.soft_cap
+                logits = torch.tanh(logits)
+                logits = logits * self.soft_cap
+
+            if self.scale != 1.0:
+                logits *= self.scale
+
+            # Apply logits processors (if any).
+            logits = _apply_logits_processors(logits, sampling_metadata)
+
+        return logits
+
+    def _get_logits(
+        self,
+        hidden_states: torch.Tensor,
+        lm_head: VocabParallelEmbedding,
+        embedding_bias: Optional[torch.Tensor],
+    ) -> Optional[torch.Tensor]:
+        # Get the logits for the next tokens.
+        logits = lm_head.linear_method.apply(lm_head,
+                                             hidden_states,
+                                             bias=embedding_bias)
+        if self.use_gather:
+            # None may be returned for rank > 0
+            logits = tensor_model_parallel_gather(logits)
+        else:
+            # Gather is not supported for some devices such as TPUs.
+            # Use all-gather instead.
+            # NOTE(woosuk): Here, the outputs of every device should not be None
+            # because XLA requires strict SPMD among all devices. Every device
+            # should execute the same operations after gathering the logits.
+            logits = tensor_model_parallel_all_gather(logits)
+        # Remove paddings in vocab (if any).
+        if logits is not None:
+            logits = logits[..., :self.org_vocab_size]
+        return logits
+
+    def extra_repr(self) -> str:
+        s = f"vocab_size={self.vocab_size}"
+        s += f", forg_vocab_size={self.org_vocab_size}"
+        s += f", scale={self.scale}, logits_as_input={self.logits_as_input}"
+        return s
+
+
+def _prune_hidden_states(
+    hidden_states: torch.Tensor,
+    sampling_metadata: SamplingMetadata,
+) -> torch.Tensor:
+    return hidden_states.index_select(0,
+                                      sampling_metadata.selected_token_indices)
+
+
+def _apply_logits_processors(
+    logits: torch.Tensor,
+    sampling_metadata: SamplingMetadata,
+) -> torch.Tensor:
+    if sampling_metadata.seq_groups is None:  # intermediate chunked-prefill chunk
+        return logits
+    found_logits_processors = False
+    logits_processed = 0
+    for seq_group in sampling_metadata.seq_groups:
+        seq_ids = seq_group.seq_ids
+        sampling_params = seq_group.sampling_params
+        logits_processors = sampling_params.logits_processors
+        if logits_processors:
+            found_logits_processors = True
+
+            for seq_id, logits_row_idx in zip(seq_ids,
+                                              seq_group.sample_indices):
+                logits_row = logits[logits_row_idx]
+                past_tokens_ids = seq_group.seq_data[seq_id].output_token_ids
+                prompt_tokens_ids = seq_group.seq_data[seq_id].prompt_token_ids
+
+                for logits_processor in logits_processors:
+                    parameters = inspect.signature(logits_processor).parameters
+                    if len(parameters) == 3:
+                        logits_row = logits_processor(prompt_tokens_ids,
+                                                      past_tokens_ids,
+                                                      logits_row)
+                    else:
+                        logits_row = logits_processor(past_tokens_ids,
+                                                      logits_row)
+
+                logits[logits_row_idx] = logits_row
+
+        logits_processed += len(seq_group.sample_indices) + len(
+            seq_group.prompt_logprob_indices)
+
+    if found_logits_processors:
+        # verifies that no rows in logits were missed unexpectedly
+        assert logits_processed == logits.shape[0]
+    return logits

qwen3_6_scripts/mamba_cache.py

@@ -0,0 +1,229 @@
+from typing import Dict, List, Optional
+
+import torch
+
+from vllm.attention.backends.abstract import AttentionMetadata
+
+
+class MambaCacheManager:
+
+    def __init__(self, dtype, num_mamba_layers, max_batch_size,
+                 conv_state_shape, temporal_state_shape):
+
+        conv_state = torch.empty(size=(num_mamba_layers, max_batch_size) +
+                                 conv_state_shape,
+                                 dtype=dtype,
+                                 device="cuda")
+        temporal_state = torch.zeros(size=(num_mamba_layers, max_batch_size) +
+                                     temporal_state_shape,
+                                     dtype=dtype,
+                                     device="cuda")
+
+        self.mamba_cache = (conv_state, temporal_state)
+
+        # Maps between the request id and a dict that maps between the seq_id
+        # and its index inside the self.mamba_cache
+        self.mamba_cache_indices_mapping: Dict[str, Dict[int, int]] = {}
+
+    def current_run_tensors(self, input_ids: torch.Tensor,
+                            attn_metadata: AttentionMetadata, **kwargs):
+        """
+        Return the tensors for the current run's conv and ssm state.
+        """
+        if "seqlen_agnostic_capture_inputs" not in kwargs:
+            # We get here only on Prefill/Eager mode runs
+            request_ids_to_seq_ids = kwargs["request_ids_to_seq_ids"]
+            finished_requests_ids = kwargs["finished_requests_ids"]
+
+            self._release_finished_requests(finished_requests_ids)
+            mamba_cache_tensors = self._prepare_current_run_mamba_cache(
+                request_ids_to_seq_ids, finished_requests_ids)
+
+        else:
+            # CUDA graph capturing runs
+            mamba_cache_tensors = kwargs["seqlen_agnostic_capture_inputs"]
+
+        return mamba_cache_tensors
+
+    def copy_inputs_before_cuda_graphs(self, input_buffers, **kwargs):
+        """
+        Copy the relevant Mamba cache into the CUDA graph input buffer
+        that was provided during the capture runs
+        (JambaForCausalLM.mamba_gc_cache_buffer).
+        """
+        assert all(
+            key in kwargs
+            for key in ["request_ids_to_seq_ids", "finished_requests_ids"])
+        finished_requests_ids = kwargs["finished_requests_ids"]
+        request_ids_to_seq_ids = kwargs["request_ids_to_seq_ids"]
+
+        self._release_finished_requests(finished_requests_ids)
+        self._prepare_current_run_mamba_cache(request_ids_to_seq_ids,
+                                              finished_requests_ids)
+
+    def get_seqlen_agnostic_capture_inputs(self, batch_size: int):
+        """
+        Provide the CUDA graph capture runs with a buffer in adjusted size.
+        The buffer is used to maintain the Mamba Cache during the CUDA graph
+        replay runs.
+        """
+        return tuple(buffer[:, :batch_size] for buffer in self.mamba_cache)
+
+    def _swap_mamba_cache(self, from_index: int, to_index: int):
+        # CCCL DeviceCopy::Batched uses separate src/dst buffers — never
+        # in-place scatter. PyTorch advanced indexing assignment
+        # cache[:, [a,b]] = cache[:, [b,a]] has undefined evaluation order.
+        # Use explicit temp clone for correctness.
+        assert len(self.mamba_cache) > 0
+        for cache_t in self.mamba_cache:
+            tmp = cache_t[:, from_index].clone()
+            cache_t[:, from_index].copy_(cache_t[:, to_index])
+            cache_t[:, to_index].copy_(tmp)
+
+    def _copy_mamba_cache(self, from_index: int, to_index: int):
+        assert len(self.mamba_cache) > 0
+        for cache_t in self.mamba_cache:
+            cache_t[:, to_index].copy_(cache_t[:, from_index],
+                                       non_blocking=True)
+
+    def _move_out_if_already_occupied(self, index: int,
+                                      all_occupied_indices: List[int]):
+        if index in all_occupied_indices:
+            first_free_index = self._first_free_index_in_mamba_cache()
+            # In case occupied, move the occupied to a new empty block
+            self._move_cache_index_and_mappings(from_index=index,
+                                                to_index=first_free_index)
+
+    def _assign_seq_id_to_mamba_cache_in_specific_dest(self, cur_rid: str,
+                                                       seq_id: int,
+                                                       destination_index: int):
+        """
+        Assign (req_id,seq_id) pair to a `destination_index` index, if
+        already occupied, move the occupying index to a free index.
+        """
+        all_occupied_indices = self._get_all_occupied_indices()
+        if cur_rid not in self.mamba_cache_indices_mapping:
+            self._move_out_if_already_occupied(
+                index=destination_index,
+                all_occupied_indices=all_occupied_indices)
+            for cache_t in self.mamba_cache:
+                cache_t[:, destination_index].zero_()
+            self.mamba_cache_indices_mapping[cur_rid] = {
+                seq_id: destination_index
+            }
+        elif seq_id not in (seq_ids2indices :=
+                            self.mamba_cache_indices_mapping[cur_rid]):
+            # parallel sampling , where n > 1, assume prefill have
+            # already happened now we only need to copy the already
+            # existing cache into the siblings seq_ids caches
+            self._move_out_if_already_occupied(
+                index=destination_index,
+                all_occupied_indices=all_occupied_indices)
+            index_exists = list(seq_ids2indices.values())[0]
+            # case of decoding n>1, copy prefill cache to decoding indices
+            self._copy_mamba_cache(from_index=index_exists,
+                                   to_index=destination_index)
+            self.mamba_cache_indices_mapping[cur_rid][
+                seq_id] = destination_index
+        else:
+            # already exists
+            cache_index_already_exists = self.mamba_cache_indices_mapping[
+                cur_rid][seq_id]
+            if cache_index_already_exists != destination_index:
+                # In case the seq id already exists but not in
+                # the right destination, swap it with what's occupying it
+                self._swap_pair_indices_and_mappings(
+                    from_index=cache_index_already_exists,
+                    to_index=destination_index)
+
+    def _prepare_current_run_mamba_cache(
+            self, request_ids_to_seq_ids: Dict[str, list[int]],
+            finished_requests_ids: List[str]):
+        running_indices = []
+        request_ids_to_seq_ids_flatten = [
+            (req_id, seq_id)
+            for req_id, seq_ids in request_ids_to_seq_ids.items()
+            for seq_id in seq_ids
+        ]
+        batch_size = len(request_ids_to_seq_ids_flatten)
+        for dest_index, (request_id,
+                         seq_id) in enumerate(request_ids_to_seq_ids_flatten):
+            if request_id in finished_requests_ids:
+                # Do not allocate cache index for requests that run
+                # and finish right after
+                continue
+            self._assign_seq_id_to_mamba_cache_in_specific_dest(
+                request_id, seq_id, dest_index)
+            running_indices.append(dest_index)
+
+        self._clean_up_first_bs_blocks(batch_size, running_indices)
+        conv_state = self.mamba_cache[0][:, :batch_size]
+        temporal_state = self.mamba_cache[1][:, :batch_size]
+
+        return (conv_state, temporal_state)
+
+    def _get_all_occupied_indices(self):
+        return [
+            cache_idx
+            for seq_ids2indices in self.mamba_cache_indices_mapping.values()
+            for cache_idx in seq_ids2indices.values()
+        ]
+
+    def _clean_up_first_bs_blocks(self, batch_size: int,
+                                  indices_for_current_run: List[int]):
+        # move out all of the occupied but currently not running blocks
+        # outside of the first n blocks
+        destination_indices = range(batch_size)
+        max_possible_batch_size = self.mamba_cache[0].shape[1]
+        for destination_index in destination_indices:
+            if destination_index in self._get_all_occupied_indices() and  \
+               destination_index not in indices_for_current_run:
+                # move not running indices outside of the batch
+                all_other_indices = list(
+                    range(batch_size, max_possible_batch_size))
+                first_avail_index = self._first_free_index_in_mamba_cache(
+                    all_other_indices)
+                self._swap_indices(from_index=destination_index,
+                                   to_index=first_avail_index)
+
+    def _move_cache_index_and_mappings(self, from_index: int, to_index: int):
+        self._copy_mamba_cache(from_index=from_index, to_index=to_index)
+        self._update_mapping_index(from_index=from_index, to_index=to_index)
+
+    def _swap_pair_indices_and_mappings(self, from_index: int, to_index: int):
+        self._swap_mamba_cache(from_index=from_index, to_index=to_index)
+        self._swap_mapping_index(from_index=from_index, to_index=to_index)
+
+    def _swap_mapping_index(self, from_index: int, to_index: int):
+        for seq_ids2index in self.mamba_cache_indices_mapping.values():
+            for seq_id, index in seq_ids2index.items():
+                if from_index == index:
+                    seq_ids2index.update({seq_id: to_index})
+                elif to_index == index:
+                    seq_ids2index.update({seq_id: from_index})
+
+    def _update_mapping_index(self, from_index: int, to_index: int):
+        for seq_ids2index in self.mamba_cache_indices_mapping.values():
+            for seq_id, index in seq_ids2index.items():
+                if from_index == index:
+                    seq_ids2index.update({seq_id: to_index})
+                    return
+
+    def _release_finished_requests(self,
+                                   finished_seq_groups_req_ids: List[str]):
+        for req_id in finished_seq_groups_req_ids:
+            if req_id in self.mamba_cache_indices_mapping:
+                self.mamba_cache_indices_mapping.pop(req_id)
+
+    def _first_free_index_in_mamba_cache(
+            self, indices_range: Optional[List[int]] = None) -> int:
+        assert self.mamba_cache is not None
+        if indices_range is None:
+            max_possible_batch_size = self.mamba_cache[0].shape[1]
+            indices_range = list(range(max_possible_batch_size))
+        all_occupied_indices = self._get_all_occupied_indices()
+        for i in indices_range:
+            if i not in all_occupied_indices:
+                return i
+        raise Exception("Couldn't find a free spot in the mamba cache! This"
+                        "should never happen")

qwen3_6_scripts/paged_attention_v2_pytorch.py

@@ -0,0 +1,343 @@
+"""
+paged_attention_v2_pytorch.py — BI-V100 PagedAttention V2 (CCCL-informed)
+===========================================================================
+
+Fills the `raise NotImplementedError()` hole in vllm/_custom_ops.py.
+
+Algorithm: Partitioned attention with log-sum-exp reduction.
+Architecture informed by CCCL patterns:
+  - summary_statistics.cu: fuse multiple statistics in a single reduction pass
+  - warp_reduce_shfl.cuh: accumulate (max, sum, weighted_output) as one compound type
+  - block_reduce_warp_reductions.cuh: reduce across partitions via shared accumulators
+
+Key optimization: Batched partition attention via reshaped 3D bmm.
+  Instead of looping over P partitions with P × torch.bmm calls,
+  reshape KV into [H, P*part_len, d] and Q into [H, 1, d], then
+  slice scores into [H, P, part_len] for partition-wise softmax.
+  This gives ONE bmm launch for all partitions.
+
+  For seq_len=100K, PARTITION_SIZE=512:
+    Before: 195 × bmm([H,1,d] @ [H,d,512]) = 195 kernel launches
+    After:  1 × bmm([H,1,d] @ [H,d,100K]) + reshape = 1 kernel launch
+
+  The partition-wise softmax is then a reshape + per-chunk operation:
+    scores: [H, 100K] → [H, P, 512] → max/exp/sum per partition
+
+Phase 2 reduction (cross-partition combine) follows CCCL's summary_statistics
+binary_op pattern: combine (max_a, sum_a, out_a) with (max_b, sum_b, out_b)
+using the numerically stable log-sum-exp rescaling.
+"""
+
+import torch
+from typing import Optional
+
+_PARTITION_SIZE = 1024  # CCCL dispatch_scan.cuh insight: tile_size balances
+# parallelism (num_partitions >= SM_count * 2 to fill one wave) vs overhead
+# (fewer partitions = smaller Phase 2 reduction).
+# BI-V100: 16 SMs, max ~32 concurrent CTAs.
+# For 100K tokens: 1024 → 98 partitions (3 waves), 512 → 195 (6 waves).
+# 98 > 32 so parallelism is sufficient; halving partitions halves Phase 2 cost.
+
+# CCCL dispatch_reduce.cuh GridEvenShare formula (line ~180):
+#   max_blocks = sm_occupancy * sm_count * subscription_factor
+#   subscription_factor = 5 (default in cub/util_device.cuh)
+# For BI-V100: sm_count=16, sm_occupancy ~= 2 (limited by registers/SMEM)
+#   → max_blocks = 2 * 16 * 5 = 160
+# If seq_len=100K with PARTITION_SIZE=1024 → 98 partitions < 160 → fine.
+# Threshold for V1→V2 handoff: when single-tile can't hold all tokens.
+#   CCCL single_tile threshold = threads * items_per_thread
+#   = 512 * 24 = 12288 tokens → V1 handles ≤12288, V2 handles >12288.
+# This aligns with BI-V100 paged_attn.py _PARTITION_SIZE=512:
+#   V2 triggers when seq_len > 512 * (max_blocks_per_seq_for_v1).
+_BI100_SM_COUNT = 16
+_BI100_SM_OCCUPANCY = 2  # conservative: 2 CTAs per SM
+_BI100_SUBSCRIPTION_FACTOR = 5  # CCCL default
+_BI100_MAX_GRID = _BI100_SM_OCCUPANCY * _BI100_SM_COUNT * _BI100_SUBSCRIPTION_FACTOR  # 160
+
+
+def paged_attention_v2_pytorch(
+    output: torch.Tensor,          # [num_seqs, num_heads, head_size]
+    exp_sums: torch.Tensor,        # [num_seqs, num_heads, max_num_partitions]
+    max_logits: torch.Tensor,      # [num_seqs, num_heads, max_num_partitions]
+    tmp_output: torch.Tensor,      # [num_seqs, num_heads, max_num_partitions, head_size]
+    query: torch.Tensor,           # [num_seqs, num_heads, head_size]
+    key_cache: torch.Tensor,       # [num_blocks, num_kv_heads, head_size/x, block_size, x]
+    value_cache: torch.Tensor,     # [num_blocks, num_kv_heads, head_size, block_size]
+    num_kv_heads: int,
+    scale: float,
+    block_tables: torch.Tensor,    # [num_seqs, max_blocks_per_seq]
+    seq_lens: torch.Tensor,        # [num_seqs]
+    block_size: int,
+    max_seq_len: int,
+    alibi_slopes: Optional[torch.Tensor],
+    kv_cache_dtype: str = "auto",
+    k_scale: float = 1.0,
+    v_scale: float = 1.0,
+    tp_rank: int = 0,
+    blocksparse_local_blocks: int = 0,
+    blocksparse_vert_stride: int = 0,
+    blocksparse_block_size: int = 64,
+    blocksparse_head_sliding_step: int = 0,
+) -> None:
+    num_seqs, num_heads, head_size = query.shape
+    gqa_ratio = num_heads // num_kv_heads
+    max_num_partitions = tmp_output.shape[2]
+
+    # Initialize unused slots
+    max_logits.fill_(float('-inf'))
+    exp_sums.zero_()
+    tmp_output.zero_()
+
+    # CCCL kernel_reduce.cuh SingleTile fast path (line ~270):
+    #   if (num_items <= threads_per_block * items_per_thread)
+    #     → InvokeSingleTile() — one CTA, no temp buffer, no Phase 2
+    # PyTorch translation: if seq_len fits in one partition, skip Phase 2 entirely.
+    # This avoids the partition/reshape/bmm overhead for short decode sequences.
+    # Qwen3.6 typical decode: seq_len grows from 1 to 100K over generation.
+    # Early tokens (seq_len < 1024) hit this fast path every step.
+    _SINGLE_TILE_THRESHOLD = _PARTITION_SIZE  # sequences this short skip partitioning
+
+    # ─── CCCL SmemResource pre-allocation (warpspeed/resource/smem_resource.cuh) ──
+    # Instead of torch.full/torch.zeros inside the loop (which allocates new GPU
+    # tensors every decode step → OOM after thousands of steps), pre-allocate
+    # staging buffers sized for the worst case and reuse them via .fill_()/.zero_().
+    # This mirrors CCCL SmemResource's stageCount-based buffer pool pattern.
+    _max_padded = max_num_partitions * _PARTITION_SIZE
+    _staging_scores = torch.full(
+        (num_heads, _max_padded), float('-inf'),
+        dtype=torch.float32, device=query.device)
+    if gqa_ratio > 1:
+        _staging_v_kv = torch.zeros(
+            (num_kv_heads, _max_padded, head_size),
+            dtype=torch.float32, device=query.device)
+    else:
+        _staging_v = torch.zeros(
+            (num_heads, _max_padded, head_size),
+            dtype=torch.float32, device=query.device)
+    # ─── End pre-allocation ──────────────────────────────────────────────────────
+
+    for seq_idx in range(num_seqs):
+        seq_len = int(seq_lens[seq_idx].item())
+        if seq_len == 0:
+            output[seq_idx].zero_()
+            continue
+
+        num_blocks_seq = (seq_len + block_size - 1) // block_size
+        num_partitions = (seq_len + _PARTITION_SIZE - 1) // _PARTITION_SIZE
+
+        # ─── CCCL SingleTile fast path ───────────────────────────
+        # From kernel_reduce.cuh: when everything fits in one tile,
+        # do a single-pass attention without partition overhead.
+        # agent_reduce.cuh ConsumeRange → BlockReduce → done.
+        if num_partitions == 1:
+            blk_ids = block_tables[seq_idx, :num_blocks_seq]
+            q = query[seq_idx].float()  # [H, d]
+
+            # Gather KV (same as below but no partition reshape)
+            k_gathered = key_cache[blk_ids]
+            k_flat = (k_gathered
+                      .permute(0, 3, 1, 2, 4)
+                      .reshape(-1, num_kv_heads, head_size))[:seq_len]
+            v_flat = (value_cache[blk_ids]
+                      .permute(0, 3, 1, 2)
+                      .reshape(-1, num_kv_heads, head_size))[:seq_len]
+
+            if k_scale != 1.0:
+                k_flat = k_flat.float().mul_(k_scale)
+            if v_scale != 1.0:
+                v_flat = v_flat.float().mul_(v_scale)
+
+            if gqa_ratio > 1:
+                k_kv = k_flat.permute(1, 2, 0).float().contiguous()
+                v_kv = v_flat.permute(1, 0, 2).float().contiguous()
+                q_grouped = q.view(num_kv_heads, gqa_ratio, 1, head_size)
+                scores = torch.matmul(q_grouped, k_kv.unsqueeze(1)).squeeze(2)
+                scores = scores.reshape(num_heads, seq_len) * scale
+            else:
+                k_t = k_flat.permute(1, 2, 0).float().contiguous()
+                scores = torch.bmm(q.unsqueeze(1), k_t).squeeze(1) * scale
+
+            if alibi_slopes is not None:
+                positions = torch.arange(seq_len, device=query.device, dtype=torch.float32)
+                scores = scores + alibi_slopes.unsqueeze(1) * positions.unsqueeze(0)
+
+            # Direct softmax + V weighted sum — no partition overhead
+            weights = torch.softmax(scores, dim=-1)  # [H, seq_len]
+            if gqa_ratio > 1:
+                w_grouped = weights.view(num_kv_heads, gqa_ratio, 1, seq_len)
+                result = torch.matmul(w_grouped, v_kv.unsqueeze(1)).squeeze(2)
+                output[seq_idx] = result.reshape(num_heads, head_size).to(output.dtype)
+            else:
+                v_perm = v_flat.permute(1, 0, 2).float().contiguous()
+                result = torch.bmm(weights.unsqueeze(1), v_perm).squeeze(1)
+                output[seq_idx] = result.to(output.dtype)
+
+            # Store dummy partition values for compatibility
+            max_logits[seq_idx, :, 0] = scores.max(dim=-1).values
+            exp_sums[seq_idx, :, 0] = weights.sum(dim=-1)
+            tmp_output[seq_idx, :, 0, :] = output[seq_idx].float()
+            continue
+        # ─── End SingleTile fast path ────────────────────────────
+
+        # =============================================================
+        # Batched KV gather: ONE index_select, ONE reshape
+        # Pattern: avoid per-block Python loop (CCCL does this via
+        # block-cooperative load, we do it via batched indexing)
+        # =============================================================
+        blk_ids = block_tables[seq_idx, :num_blocks_seq]
+
+        # Key: [nblk, kv_h, d/x, blk_sz, x] → [nblk*blk_sz, kv_h, d]
+        k_gathered = key_cache[blk_ids]
+        k_flat = (k_gathered
+                  .permute(0, 3, 1, 2, 4)
+                  .reshape(-1, num_kv_heads, head_size))[:seq_len]
+
+        # Value: [nblk, kv_h, d, blk_sz] → [nblk*blk_sz, kv_h, d]
+        v_flat = (value_cache[blk_ids]
+                  .permute(0, 3, 1, 2)
+                  .reshape(-1, num_kv_heads, head_size))[:seq_len]
+
+        if k_scale != 1.0:
+            k_flat = k_flat.float().mul_(k_scale)
+        if v_scale != 1.0:
+            v_flat = v_flat.float().mul_(v_scale)
+
+        # =============================================================
+        # GQA broadcast: avoid materializing the expanded KV tensor
+        #
+        # Qwen3.6: H=24, kv_h=4, gqa_ratio=6, head_dim=256
+        # Old: expand kv_h→H then contiguous → allocates seq_len×H×d (1.2GB at 100K)
+        # New: reshape Q as [kv_h, gqa, 1, d], K as [kv_h, 1, d, seq_len]
+        #      → bmm with broadcasting → [kv_h, gqa, 1, seq_len]
+        #      → reshape to [H, seq_len]
+        # Saves: gqa_ratio × memory (6x for Qwen3.6 = 1GB per decode step)
+        # =============================================================
+        q = query[seq_idx].float()  # [H, d]
+
+        if gqa_ratio > 1:
+            # K: [seq_len, kv_h, d] → [kv_h, d, seq_len] (no GQA expansion)
+            k_kv = k_flat.permute(1, 2, 0).float().contiguous()  # [kv_h, d, seq_len]
+            v_kv = v_flat.permute(1, 0, 2).float().contiguous()  # [kv_h, seq_len, d]
+
+            # Q: [H, d] → [kv_h, gqa, 1, d]
+            q_grouped = q.view(num_kv_heads, gqa_ratio, 1, head_size)
+
+            # Scores: [kv_h, gqa, 1, d] @ [kv_h, 1, d, seq_len] → [kv_h, gqa, 1, seq_len]
+            scores_all = torch.matmul(q_grouped, k_kv.unsqueeze(1)).squeeze(2)  # [kv_h, gqa, seq_len]
+            scores_all = scores_all.reshape(num_heads, seq_len) * scale  # [H, seq_len]
+        else:
+            k_t = k_flat.permute(1, 2, 0).float().contiguous()  # [H, d, seq_len]
+            scores_all = torch.bmm(q.unsqueeze(1), k_t).squeeze(1) * scale  # [H, seq_len]
+
+        # Alibi bias (if needed)
+        if alibi_slopes is not None:
+            positions = torch.arange(seq_len, device=query.device, dtype=torch.float32)
+            scores_all = scores_all + alibi_slopes.unsqueeze(1) * positions.unsqueeze(0)
+
+        # Pad to exact multiple of _PARTITION_SIZE for clean reshape
+        # CCCL SmemResource: reuse staging buffer instead of allocating
+        padded_len = num_partitions * _PARTITION_SIZE
+        if padded_len > seq_len:
+            scores_padded = _staging_scores[:, :padded_len]
+            scores_padded.fill_(float('-inf'))
+            scores_padded[:, :seq_len] = scores_all
+        else:
+            scores_padded = scores_all
+
+        # Reshape: [H, padded_len] → [H, P, part_sz]
+        scores_parts = scores_padded.view(num_heads, num_partitions, _PARTITION_SIZE)
+
+        # Per-partition online softmax (vectorized over H and P simultaneously)
+        # Pattern from CCCL summary_statistics: compute (max, sum) in one pass
+        part_max = scores_parts.max(dim=-1).values          # [H, P]
+        scores_exp = torch.exp(scores_parts - part_max.unsqueeze(-1))  # [H, P, part_sz]
+        part_sum = scores_exp.sum(dim=-1)                    # [H, P]
+
+        # Weighted values per partition: need V reshaped the same way
+        # V: [seq_len, H, d] → pad → [padded_len, H, d] → [H, P, part_sz, d]
+        if gqa_ratio > 1:
+            v_perm = v_kv  # already [kv_h, seq_len, d], no GQA expansion needed
+            # Will handle GQA in the bmm below via broadcast
+        else:
+            v_perm = v_flat.permute(1, 0, 2).float().contiguous()  # [H, seq_len, d]
+        # Weighted V sum per partition
+        # NOTE: v_perm shape differs by GQA mode:
+        #   GQA: v_perm = v_kv = [kv_h, seq_len, d]
+        #   No GQA: v_perm = [H, seq_len, d]
+        # scores_exp: [H, P, part_sz] → [kv_h, gqa, P, part_sz]
+        # v_perm: [kv_h, seq_len, d] → [kv_h, P, part_sz, d]
+        if gqa_ratio > 1:
+            se_grouped = scores_exp.view(num_kv_heads, gqa_ratio, num_partitions, _PARTITION_SIZE)
+            # V: pad and reshape to [kv_h, P, part_sz, d]
+            # CCCL SmemResource: reuse staging buffer
+            if padded_len > seq_len:
+                v_padded_kv = _staging_v_kv[:, :padded_len, :]
+                v_padded_kv.zero_()
+                v_padded_kv[:, :seq_len, :] = v_kv
+            else:
+                v_padded_kv = v_kv
+            v_parts_kv = v_padded_kv.view(num_kv_heads, num_partitions, _PARTITION_SIZE, head_size)
+            # Broadcast: [kv_h, gqa, P, 1, part_sz] @ [kv_h, 1, P, part_sz, d]
+            #          → [kv_h, gqa, P, 1, d]
+            part_out_grouped = torch.matmul(
+                se_grouped.unsqueeze(3),       # [kv_h, gqa, P, 1, part_sz]
+                v_parts_kv.unsqueeze(1)        # [kv_h, 1, P, part_sz, d]
+            ).squeeze(3)                        # [kv_h, gqa, P, d]
+            part_out = part_out_grouped.reshape(num_heads, num_partitions, head_size)
+        else:
+            # Non-GQA: v_perm is [H, seq_len, d], pad and reshape normally
+            # CCCL SmemResource: reuse staging buffer
+            if padded_len > seq_len:
+                v_padded = _staging_v[:, :padded_len, :]
+                v_padded.zero_()
+                v_padded[:, :seq_len, :] = v_perm
+            else:
+                v_padded = v_perm
+            v_parts = v_padded.view(num_heads, num_partitions, _PARTITION_SIZE, head_size)
+            HP = num_heads * num_partitions
+            scores_exp_flat = scores_exp.reshape(HP, 1, _PARTITION_SIZE)
+            v_parts_flat = v_parts.reshape(HP, _PARTITION_SIZE, head_size)
+            part_out_flat = torch.bmm(scores_exp_flat, v_parts_flat)  # [HP, 1, d]
+            part_out = part_out_flat.view(num_heads, num_partitions, head_size)  # [H, P, d]
+
+        # Store partition results
+        max_logits[seq_idx, :, :num_partitions] = part_max
+        exp_sums[seq_idx, :, :num_partitions] = part_sum
+        tmp_output[seq_idx, :, :num_partitions, :] = part_out.to(tmp_output.dtype)
+
+        # =============================================================
+        # Phase 2: Cross-partition reduction (CCCL binary_op pattern)
+        #
+        # CCCL kernel_reduce.cuh insight: when grid_size fits in a single
+        # tile (num_partitions <= threads * items_per_thread), the reduce
+        # uses SingleTile path — one CTA, no temp buffer, no pass 2 kernel.
+        #
+        # For BI-V100 with 98 partitions (100K tokens / 1024 partition_size):
+        #   SingleTile threshold = 512 * 24 = 12288 >> 98 → always SingleTile
+        #   This means Phase 2 is never the bottleneck.
+        #
+        # CCCL single_pass_scan_operators.cuh insight: delay() has a
+        # GridThreshold=500 gate. BI-V100 scan grids are always < 500 blocks,
+        # so ALL delay strategies (no_delay, fixed_delay, exponential_backon)
+        # collapse to __threadfence_block(). Delay tuning is irrelevant here.
+        #
+        # Phase 2 follows summary_statistics.cu binary_op: combine
+        # (max_a, sum_a, out_a) ⊕ (max_b, sum_b, out_b) via log-sum-exp.
+        # Fully vectorized — no loop over partitions.
+        # =============================================================
+        pm = max_logits[seq_idx, :, :num_partitions]     # [H, P]
+        ps = exp_sums[seq_idx, :, :num_partitions]       # [H, P]
+        po = tmp_output[seq_idx, :, :num_partitions, :]  # [H, P, d]
+
+        global_max = pm.max(dim=-1).values               # [H]
+        rescale = torch.exp(pm - global_max.unsqueeze(-1)) * ps  # [H, P]
+        total = rescale.sum(dim=-1, keepdim=True)         # [H, 1]
+
+        # CCCL norm.cu principle: fuse transform with reduce to minimize traversals.
+        # Instead of: weights = rescale/total; final = bmm(weights, po)
+        # Do: final = bmm(rescale, po) / total
+        # Saves one element-wise division kernel launch (rescale/total → H*P elements).
+        # The division moves to the output (H*d elements, typically smaller than H*P).
+        # [H, 1, P] @ [H, P, d] → [H, 1, d] → [H, d]
+        final = torch.bmm(rescale.unsqueeze(1), po.float()).squeeze(1) / total  # [H, d]
+        output[seq_idx] = final.to(output.dtype)

qwen3_6_scripts/paged_attn.py

@@ -0,0 +1,807 @@
+from dataclasses import dataclass
+from typing import List, Optional, Tuple
+import sys
+import torch
+import traceback
+from vllm import _custom_ops as ops
+
+# from vllm.attention.ops.prefix_prefill import context_attention_fwd
+# NOTE: context_attention_fwd (Triton kernel from prefix_prefill.py) is NOT
+# imported here.  On Iluvatar BI-V100 that kernel hangs the GPU card
+# permanently.  Chunked-prefill / prefix-caching attention is handled by
+# _forward_prefix_pytorch below (pure PyTorch, no Triton dependency).
+
+# Should be the same as PARTITION_SIZE in `paged_attention_v2_launcher`.
+_PARTITION_SIZE = 512
+
+
+@dataclass
+class PagedAttentionMetadata:
+    """Metadata for PagedAttention."""
+    # (batch_size,). The length of sequences (entire tokens seen so far) per
+    # sequence.
+    seq_lens_tensor: Optional[torch.Tensor]
+    # Maximum sequence length in the batch. 0 if it is prefill-only batch.
+    max_decode_seq_len: int
+    # (batch_size, max_blocks_per_seq).
+    # Block addresses per sequence. (Seq id -> list of physical block)
+    # E.g., [0, 1, 2] means tokens are stored in 0th, 1st, and 2nd blocks
+    # in the kv cache. Each block can contain up to block_size tokens.
+    # 2nd dimensions are padded up to max_blocks_per_seq if it is cuda-graph
+    # captured.
+    block_tables: Optional[torch.Tensor]
+
+
+class PagedAttention:
+
+    @staticmethod
+    def get_supported_head_sizes() -> List[int]:
+        return [64, 80, 96, 112, 120, 128, 192, 256]
+
+    @staticmethod
+    def get_kv_cache_shape(
+        num_blocks: int,
+        block_size: int,
+        num_kv_heads: int,
+        head_size: int,
+    ) -> Tuple[int, ...]:
+        return (2, num_blocks, block_size * num_kv_heads * head_size)
+
+    @staticmethod
+    def split_kv_cache(
+        kv_cache: torch.Tensor,
+        num_kv_heads: int,
+        head_size: int,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        x = 16 // kv_cache.element_size()
+        num_blocks = kv_cache.shape[1]
+
+        key_cache = kv_cache[0]
+        key_cache = key_cache.view(num_blocks, num_kv_heads, head_size // x,
+                                   -1, x)
+        value_cache = kv_cache[1]
+        value_cache = value_cache.view(num_blocks, num_kv_heads, head_size, -1)
+        return key_cache, value_cache
+
+    @staticmethod
+    def write_to_paged_cache(
+        key: torch.Tensor,
+        value: torch.Tensor,
+        key_cache: torch.Tensor,
+        value_cache: torch.Tensor,
+        slot_mapping: torch.Tensor,
+        kv_cache_dtype: str,
+        k_scale: float,
+        v_scale: float,
+    ) -> None:
+        ops.reshape_and_cache(
+            key,
+            value,
+            key_cache,
+            value_cache,
+            slot_mapping.flatten(),
+            kv_cache_dtype,
+            k_scale,
+            v_scale,
+        )
+
+    @staticmethod
+    def _forward_decode_pytorch(
+        query: torch.Tensor,
+        key_cache: torch.Tensor,
+        value_cache: torch.Tensor,
+        block_tables: torch.Tensor,
+        seq_lens: torch.Tensor,
+        scale: float,
+    ) -> torch.Tensor:
+        """Pure-PyTorch decode attention for long contexts (no hardware kernel).
+
+        Architecture mirrors CCCL's three-layer reduce:
+          dispatch_reduce.cuh  →  kernel_reduce.cuh  →  agent_reduce.cuh
+          (work distribution)     (kernel entry)         (tile consumption)
+
+        CCCL agent_reduce.cuh has two key patterns we translate here:
+
+        1. ConsumeFullTile vectorized path: data loaded as VectorT in striped
+           access (no BlockLoad staging → no SMEM for data, only for BlockReduce
+           scratch). PyTorch equivalent: single reshape+view without .contiguous()
+           when possible; fall back to one .contiguous() per K/V gather.
+
+        2. ConsumeTiles with GridEvenShare STRIP_MINE: each CTA strides across
+           the input with stride = grid_size * tile_items. For decode (q_len=1),
+           we tile over KV blocks with adaptive tile_sz per the same
+           GridEvenShare formula: max_tiles = sm_count * subscription_factor.
+
+        3. summary_statistics.cu compound reduce: accumulator = {m, l, o}.
+           unary_op: score_tile → (max, sum_exp, weighted_V).
+           binary_op: online softmax merge with correction factor.
+           This is the Flash Attention online softmax — identical structure.
+
+        For decode, q_len=1 per sequence. The attention weight is [H, 1, seq_len]
+        which is small (~5 MB at 50K tokens). We tile over KV blocks to control
+        peak memory and apply online softmax (Flash Attention Algorithm 1) per tile.
+
+        Shapes
+        ------
+        query       : [num_seqs, num_heads, head_dim]
+        key_cache   : [num_blocks, num_kv_heads, head_dim//x, block_size, x]
+        value_cache : [num_blocks, num_kv_heads, head_dim,    block_size]
+        block_tables: [num_seqs, max_blocks_per_seq]
+        seq_lens    : [num_seqs]
+        """
+        num_seqs, num_heads, head_dim = query.shape
+        num_kv_heads = key_cache.shape[1]
+        block_size = value_cache.shape[3]
+        gqa_ratio = num_heads // num_kv_heads
+        orig_dtype = query.dtype
+        dev = query.device
+
+        output = torch.empty_like(query)
+
+        # ================================================================
+        # CCCL spread_out_items_per_thread adaptive tile sizing for decode
+        #
+        # Ported from dispatch_transform.cuh::spread_out_items_per_thread
+        # and dispatch_reduce.cuh::InvokePasses GridEvenShare.
+        #
+        # CCCL formula (dispatch_transform.cuh line 183):
+        #   items = min(max_items,
+        #               ceil_div(num_items, sm_count * threads * max_occupancy))
+        #   items = clamp(items, min_items, max_items)
+        #
+        # Our translation for PyTorch decode:
+        #   "items" = KV blocks per tile (how much work per matmul call)
+        #   "num_items" = total KV blocks in the sequence
+        #   "sm_count * max_occupancy" = target number of tiles (~4-8)
+        #   Fewer tiles = fewer Python loop iterations = less launch overhead
+        #
+        # For decode (q_len=1), score tensor per tile is tiny:
+        #   kv_h × gqa × 1 × (tile_blocks × block_size) × 4 bytes
+        #   = 4 × 6 × 1 × 16384 × 4 = 1.5 MB  (even at kv_h=4, safe)
+        # So the constraint is NOT memory — it's minimizing loop iterations.
+        #
+        # CCCL grid_even_share.cuh DispatchInit logic:
+        #   total_tiles = ceil_div(num_items, tile_size)
+        #   grid_size = min(total_tiles, max_grid_size)
+        #   big_shares = total_tiles - (avg_tiles * grid_size)
+        # Our target: ~4 tiles max (Python overhead >> kernel launch overhead)
+        # ================================================================
+        # CCCL GridEvenShare: max_blocks = sm_occupancy * sm_count * subscription_factor
+        # BI-V100: 1 * 16 * 5 = 80 max CTAs for CUDA kernels.
+        # But this is Python (PyTorch ops), not CUDA launches — Python loop
+        # overhead dominates. Each iteration = 1 torch.matmul launch + online
+        # softmax update. Target 2 iterations (not 4): the matmul itself is
+        # already parallelized across SMs, so fewer Python loops = less overhead.
+        # For seq_len=100K with block_size=16: 6250 blocks / 2 = 3125 blocks/tile.
+        # Score tensor: 4 kv_heads × 6 gqa × 1 × 50000 × 4B = 4.8 MB — fits.
+        _BI100_TARGET_TILES = 2   # 2 iterations: minimize Python loop overhead
+        _MIN_TILE_BLOCKS = 128    # floor: ensure matmul is large enough to saturate 16 SMs
+        _MAX_TILE_BLOCKS = 8192   # ceiling: 8192 × 16 = 128K tokens per tile — fits in memory
+
+        try:
+            for i in range(num_seqs):
+                seq_len = int(seq_lens[i].item())
+                if seq_len == 0:
+                    output[i].zero_()
+                    continue
+
+                num_blocks_i = (seq_len + block_size - 1) // block_size
+                blk_ids = block_tables[i, :num_blocks_i]
+
+                # Q reshaped once: [kv_h, gqa, 1, d] fp32 — tiny for decode
+                q_grouped = (query[i].float()
+                             .view(num_kv_heads, gqa_ratio, head_dim)
+                             .unsqueeze(2)
+                             .mul_(scale))
+
+                # Online softmax accumulators (CCCL summary_stats_data pattern)
+                # accumulator = {m (running max), l (running sum_exp), o (running output)}
+                m = torch.full((num_kv_heads, gqa_ratio, 1),
+                               float('-inf'), dtype=torch.float32, device=dev)
+                l = torch.zeros_like(m)
+                o = torch.zeros((num_kv_heads, gqa_ratio, 1, head_dim),
+                                dtype=torch.float32, device=dev)
+
+                # Tile over KV blocks — CCCL spread_out_items_per_thread pattern
+                # Adaptive: tile_blocks = ceil(num_blocks / target_tiles)
+                # clamped to [_MIN_TILE_BLOCKS, _MAX_TILE_BLOCKS]
+                tile_blocks = max(_MIN_TILE_BLOCKS,
+                                  min(_MAX_TILE_BLOCKS,
+                                      (num_blocks_i + _BI100_TARGET_TILES - 1)
+                                      // _BI100_TARGET_TILES))
+                for tile_start in range(0, num_blocks_i, tile_blocks):
+                    tile_end = min(tile_start + tile_blocks, num_blocks_i)
+                    tile_blk_ids = blk_ids[tile_start:tile_end]
+
+                    # Valid tokens in this tile
+                    tile_token_start = tile_start * block_size
+                    tile_token_end = min(tile_end * block_size, seq_len)
+                    valid_tokens = tile_token_end - tile_token_start
+
+                    # --------------------------------------------------------
+                    # KV gather — agent_reduce.cuh ConsumeFullTile pattern
+                    #
+                    # agent_reduce loads VectorT in striped access when possible.
+                    # PyTorch equivalent: reshape the 5D cache layout to 3D in
+                    # one permute+contiguous, avoiding the double-contiguous
+                    # pattern of the old code.
+                    #
+                    # key_cache shape: [num_blocks, kv_h, d//x, blk_sz, x]
+                    # Target: [kv_h, d, valid_tokens] for Q@K^T
+                    #
+                    # Optimized path: permute(1,2,4,0,3) → [kv_h, d//x, x, n_blk, blk_sz]
+                    # → reshape to [kv_h, d, n_blk*blk_sz] → slice [:valid_tokens]
+                    # This is ONE contiguous() call instead of TWO.
+                    # --------------------------------------------------------
+                    k_gathered = key_cache[tile_blk_ids]  # [n, kv_h, d//x, blk_sz, x]
+                    k_t = (k_gathered
+                           .permute(1, 2, 4, 0, 3)       # [kv_h, d//x, x, n, blk_sz]
+                           .contiguous()
+                           .view(num_kv_heads, head_dim, -1)  # [kv_h, d, n*blk_sz]
+                           [:, :, :valid_tokens]
+                           .unsqueeze(1)                  # [kv_h, 1, d, valid]
+                           .float())
+                    del k_gathered
+
+                    v_gathered = value_cache[tile_blk_ids]  # [n, kv_h, d, blk_sz]
+                    v_t = (v_gathered
+                           .permute(1, 2, 0, 3)           # [kv_h, d, n, blk_sz]
+                           .contiguous()
+                           .view(num_kv_heads, head_dim, -1)  # [kv_h, d, n*blk_sz]
+                           [:, :, :valid_tokens]
+                           .transpose(1, 2)               # [kv_h, valid, d]
+                           .unsqueeze(1)                  # [kv_h, 1, valid, d]
+                           .float())
+                    del v_gathered
+
+                    # --------------------------------------------------------
+                    # Scores + online softmax — summary_statistics.cu pattern
+                    #
+                    # unary_op: score_tile → (max, sum_exp, weighted_V)
+                    # binary_op: merge with correction factor
+                    #
+                    # CCCL summary_stats_binary_op merges:
+                    #   result.mean = x.mean + delta * y.n / n
+                    #   result.M2  = x.M2 + y.M2 + delta² * x.n * y.n / n
+                    #
+                    # Online softmax merge:
+                    #   m_new  = max(m_old, m_tile)
+                    #   corr   = exp(m_old - m_new)        ← rescale factor
+                    #   l_new  = l_old * corr + l_tile
+                    #   o_new  = o_old * corr + tile_exp @ V
+                    #
+                    # Structurally identical: m↔max, l↔n, o↔mean×n.
+                    # --------------------------------------------------------
+
+                    # [kv_h, gqa, 1, valid_tokens]
+                    s = torch.matmul(q_grouped, k_t)
+                    del k_t
+
+                    # Online softmax update (Flash Attention Algorithm 1)
+                    m_tile = s.amax(dim=-1, keepdim=True)    # [kv_h, gqa, 1, 1]
+                    m_new = torch.maximum(m, m_tile.squeeze(-1))
+                    corr = torch.exp(m - m_new)              # rescale old accum
+
+                    exp_s = torch.exp(s - m_new.unsqueeze(-1))
+                    del s
+
+                    m.copy_(m_new)
+                    l.mul_(corr).add_(exp_s.sum(dim=-1))
+                    o.mul_(corr.unsqueeze(-1)).add_(torch.matmul(exp_s, v_t))
+                    del exp_s, v_t, corr, m_new, m_tile
+
+                # Finalize: normalize
+                o.div_(l.unsqueeze(-1))
+                output[i] = (o.view(num_heads, head_dim)
+                              .to(orig_dtype))
+
+        except Exception as e:
+            print(f"[decode_pytorch ERROR] {type(e).__name__}: {e}",
+                  file=sys.stderr, flush=True)
+            traceback.print_exc(file=sys.stderr)
+            raise
+
+        return output
+
+    # ================================================================
+    # CCCL Design Pattern: summary_statistics.cu transform_reduce
+    #
+    # CCCL packs {n, min, max, mean, M2, M3, M4} into one struct and
+    # computes ALL statistics in a single pass via transform_reduce.
+    # The binary_op merges two partial results (Welford parallel algo).
+    #
+    # Our online softmax is the same pattern:
+    #   accumulator = {m (running max), l (running sum_exp), o (running output)}
+    #   unary_op: score_tile → {max(tile), sum(exp(tile-max)), exp(tile-max) @ V}
+    #   binary_op: merge two accumulators with correction factor
+    #
+    # Key insight: kv_heads are INDEPENDENT — no cross-head dependency.
+    # Current code already batches via [kv_h, gqa, q_len, tile_sz] tensor ops.
+    # The CCCL pattern validates this is optimal: one matmul per tile across
+    # all heads simultaneously, not per-head iteration.
+    #
+    # Future optimization: if we ever get Triton/CUDA access, the binary_op
+    # merge step ({m,l,o} update) could be fused with the matmul via a
+    # custom epilogue — this is what FlashAttention-2/3 does at the CUDA level.
+    # ================================================================
+
+    # paged_attention_v1 on BI-V100: ixformer native kernel handles long contexts.
+    # PyTorch fallback is only for emergency (kernel crash at extreme lengths).
+    # CCCL GridEvenShare principle: each work unit (decode step) must complete
+    # within bounded time — Python fallback is too slow for seq_len > 32K
+    # (causes HTTP timeout → service crash). Native V1 kernel is O(1) per step.
+    # Threshold raised to avoid fallback during normal operation.
+    _PYTORCH_DECODE_THRESHOLD = 999999
+
+    @staticmethod
+    def forward_decode(
+        query: torch.Tensor,
+        key_cache: torch.Tensor,
+        value_cache: torch.Tensor,
+        block_tables: torch.Tensor,
+        seq_lens: torch.Tensor,
+        max_seq_len: int,
+        kv_cache_dtype: str,
+        num_kv_heads: int,
+        scale: float,
+        alibi_slopes: Optional[torch.Tensor],
+        k_scale: float,
+        v_scale: float,
+        tp_rank: int = 0,
+        blocksparse_local_blocks: int = 0,
+        blocksparse_vert_stride: int = 0,
+        blocksparse_block_size: int = 64,
+        blocksparse_head_sliding_step: int = 0,
+    ) -> torch.Tensor:
+        actual_max = int(seq_lens.max().item()) if seq_lens.numel() > 0 else max_seq_len
+        if actual_max > PagedAttention._PYTORCH_DECODE_THRESHOLD:
+            return PagedAttention._forward_decode_pytorch(
+                query, key_cache, value_cache, block_tables, seq_lens, scale)
+
+        if blocksparse_vert_stride is not None and blocksparse_vert_stride > 1:
+            # use blocksparse paged attention
+            block_size = value_cache.size(-1)
+            assert (blocksparse_block_size > 0 and
+                    blocksparse_block_size % block_size == 0), \
+                (f"{blocksparse_block_size=} needs to be a multiple of"
+                 f"{block_size=} used in block_tables.")
+
+        output = torch.empty_like(query)
+        block_size = value_cache.shape[3]
+        num_seqs, num_heads, head_size = query.shape
+        max_num_partitions = ((max_seq_len + _PARTITION_SIZE - 1) //
+                              _PARTITION_SIZE)
+        # NOTE(woosuk): We use a simple heuristic to decide whether to use
+        # PagedAttention V1 or V2. If the number of partitions is 1, we use
+        # V1 to avoid the overhead of reduction. Also, if the number of
+        # sequences or heads is large, we use V1 since there is enough work
+        # to parallelize.
+        # TODO(woosuk): Tune this heuristic.
+        # For context len > 8192, use V2 kernel to avoid shared memory shortage.
+        # CCCL dispatch_reduce.cuh two-path dispatch architecture:
+        #   single-tile:  num_items ≤ threads × items  →  one CTA, zero temp buffer
+        #   multi-tile:   GridEvenShare partitions across sm_count × occupancy CTAs
+        #
+        # Paged attention equivalent:
+        #   V1 = single-pass: one CTA iterates ALL KV blocks (like DeviceReduceSingleTileKernel)
+        #   V2 = partitioned: KV blocks split into PARTITION_SIZE chunks across CTAs,
+        #        then a second kernel merges partition results (like InvokePasses two-phase)
+        #
+        # V1 is optimal when seq_len fits in one CTA's tile (small context).
+        # V2 is optimal when seq_len >> PARTITION_SIZE (long context) — parallelism
+        # across partitions compensates for the merge overhead.
+        #
+        # CCCL's GridEvenShare formula:
+        #   max_blocks = sm_occupancy × sm_count × subscription_factor
+        #   BI-V100: ~1 × 16 × 5 = 80 max blocks
+        #   V2 becomes worthwhile when max_num_partitions > 1 AND the partition
+        #   parallelism exceeds the sequence×head parallelism.
+        #
+        # Original heuristic (before hardcode): V1 when max_seq_len ≤ 8192 OR
+        # when batch×heads already saturates the GPU (num_seqs*num_heads > 512).
+        # Restored with BI-V100 SM count awareness.
+        bi100_sm_count = 16
+        bi100_saturation = bi100_sm_count * 32  # ~512 concurrent warps
+        use_v1 = (max_num_partitions == 1
+                  or max_seq_len <= 8192
+                  or num_seqs * num_heads > bi100_saturation)
+        if use_v1:
+            # Run PagedAttention V1.
+            ops.paged_attention_v1(
+                output,
+                query,
+                key_cache,
+                value_cache,
+                num_kv_heads,
+                scale,
+                block_tables,
+                seq_lens,
+                block_size,
+                max_seq_len,
+                alibi_slopes,
+            )
+        else:
+            # Run PagedAttention V2.
+            assert _PARTITION_SIZE % block_size == 0
+            # CCCL agent_merge_sort.cuh union _TempStorage pattern:
+            # agent_merge_sort shares a single SMEM allocation across
+            # load_keys, load_items, store_keys, and block_merge ops
+            # (they don't execute concurrently, so one buffer suffices).
+            # Our equivalent: cache V2 temp tensors across decode steps.
+            # For max_num_seqs=1 (competition config), these shapes are
+            # stable across all decode steps for the same sequence.
+            _v2_key = ("v2_tmp", num_seqs, num_heads, max_num_partitions,
+                       head_size, output.dtype, output.device)
+            _v2_cached = getattr(PagedAttention, '_v2_cache', {}).get(_v2_key)
+            if _v2_cached is not None:
+                tmp_output, exp_sums, max_logits = _v2_cached
+            else:
+                tmp_output = torch.empty(
+                    size=(num_seqs, num_heads, max_num_partitions, head_size),
+                    dtype=output.dtype,
+                    device=output.device,
+                )
+                exp_sums = torch.empty(
+                    size=(num_seqs, num_heads, max_num_partitions),
+                    dtype=torch.float32,
+                    device=output.device,
+                )
+                max_logits = torch.empty_like(exp_sums)
+                if not hasattr(PagedAttention, '_v2_cache'):
+                    PagedAttention._v2_cache = {}
+                PagedAttention._v2_cache[_v2_key] = (tmp_output, exp_sums, max_logits)
+            ops.paged_attention_v2(
+                output,
+                exp_sums,
+                max_logits,
+                tmp_output,
+                query,
+                key_cache,
+                value_cache,
+                num_kv_heads,
+                scale,
+                block_tables,
+                seq_lens,
+                block_size,
+                max_seq_len,
+                alibi_slopes,
+                kv_cache_dtype,
+                k_scale,
+                v_scale,
+                tp_rank,
+                blocksparse_local_blocks,
+                blocksparse_vert_stride,
+                blocksparse_block_size,
+                blocksparse_head_sliding_step,
+            )
+        return output
+
+    @staticmethod
+    def forward_prefix(
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        kv_cache_dtype: str,
+        key_cache: torch.Tensor,
+        value_cache: torch.Tensor,
+        block_tables: torch.Tensor,
+        query_start_loc: torch.Tensor,
+        seq_lens_tensor: torch.Tensor,
+        context_lens: torch.Tensor,
+        max_query_len: int,
+        alibi_slopes: Optional[torch.Tensor],
+        sliding_window: Optional[int],
+        k_scale: float,
+        v_scale: float,
+    ) -> torch.Tensor:
+        # NOTE: The Triton context_attention_fwd kernel hangs on Iluvatar
+        # BI-V100 hardware (same class of issue as cudnnFlashAttnForward).
+        # Use a pure-PyTorch fallback that reads the paged KV cache directly.
+        return PagedAttention._forward_prefix_pytorch(
+            query, key, value,
+            key_cache, value_cache,
+            block_tables, query_start_loc,
+            seq_lens_tensor, context_lens,
+        )
+
+    @staticmethod
+    def _forward_prefix_pytorch(
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        key_cache: torch.Tensor,
+        value_cache: torch.Tensor,
+        block_tables: torch.Tensor,
+        query_start_loc: torch.Tensor,
+        seq_lens_tensor: torch.Tensor,
+        context_lens: torch.Tensor,
+    ) -> torch.Tensor:
+        """Pure-PyTorch prefix-attention with K-tiling (Flash-Attention online softmax).
+
+        Memory complexity: O(q_len), independent of kv_len.
+        With chunked prefill (q_len ≤ max_num_batched_tokens = 4096) peak
+        per layer ≈ 96 MB regardless of context length.
+
+        Algorithm: Flash Attention online softmax.
+        Q is reshaped once to [kv_h, gqa, q_len, d] (24 MB) and held for all
+        K-tiles.  For each tile a running (m, l, o) accumulator is updated —
+        the [q_len × kv_len] attention matrix is NEVER materialised in full.
+
+        Tile budget (kv_h=1, gqa=6, q_len=4096, tile=256 tokens):
+            q_seq   [1, 6, 4096, 256] fp32  24 MB  (held all tiles)
+            o_acc   same shape               24 MB  (held all tiles)
+            s       same shape               24 MB  (per tile, freed before exp_s)
+            exp_s   same shape               24 MB  (per tile, brief overlap with s)
+            Peak ≈ 96 MB  (s and exp_s briefly coexist during update).
+
+        Shapes
+        ------
+        query          : [total_q_tokens, num_q_heads,  head_dim]
+        key            : [total_q_tokens, num_kv_heads, head_dim]
+        value          : [total_q_tokens, num_kv_heads, head_dim]
+        key_cache      : [num_blocks, num_kv_heads, head_dim//x, block_size, x]
+        value_cache    : [num_blocks, num_kv_heads, head_dim,    block_size]
+        block_tables   : [batch_size, max_blocks_per_seq]
+        query_start_loc: [batch_size + 1]
+        seq_lens_tensor: [batch_size]  total length (context + query)
+        context_lens   : [batch_size]  tokens already in KV cache
+        """
+        try:
+            # ================================================================
+            # Tile sizing strategy — ported from CCCL dispatch_reduce.cuh
+            #
+            # CCCL's GridEvenShare computes:
+            #   max_blocks = sm_occupancy × sm_count × subscription_factor
+            #   tile_size  = num_items / max_blocks  (evenly distributed)
+            #
+            # For BI-V100 (16 SMs), fixed _BLOCKS_PER_TILE=32 wastes memory
+            # on short contexts and underutilizes on long ones.
+            #
+            # Key insight from kernel_reduce.cuh:
+            #   StableReductionOrder=false uses atomicAdd → single kernel pass.
+            #   For online softmax (our case), we accumulate (m, l, o) per tile
+            #   then merge — this IS a multi-pass reduce. Larger tiles = fewer
+            #   merge steps = less numerical drift + less Python loop overhead.
+            #
+            # CCCL subscription_factor = CUB_SUBSCRIPTION_FACTOR(0) = 5
+            # Effective: 16 SM × 1 CTA/SM × 5 = 80 concurrent tiles max.
+            # But Python loop overhead dominates, so we want FEWER, LARGER tiles.
+            #
+            # Strategy: target ~4-8 tiles per context phase.
+            #   Fewer tiles → fewer matmul calls → less launch overhead.
+            #   SMEM constraint: score tensor [kv_h, gqa, q_len, tile_sz] fp32
+            #     must not cause OOM. With q_len=4096, kv_h=1, gqa=6:
+            #     tile_sz=1024 → 1×6×4096×1024×4 = 96 MB (too much)
+            #     tile_sz=512  → 48 MB (borderline)
+            #     tile_sz=256  → 24 MB (safe)
+            #   For decode (q_len=1): tile_sz=4096 → only 96 KB (always safe)
+            # ================================================================
+            _SMEM_BUDGET_BYTES = 256 * 1024 * 1024  # 256 MB score tensor budget
+            # CCCL GridEvenShare: fewer tiles = fewer iterations = less overhead
+            # BI-V100 has 32 GB HBM per card; 256 MB temporary is safe.
+
+            batch_size   = seq_lens_tensor.shape[0]
+            num_q_heads  = query.shape[1]
+            num_kv_heads = key_cache.shape[1]
+            head_dim     = query.shape[2]
+            gqa_ratio    = num_q_heads // num_kv_heads
+            block_size   = value_cache.shape[3]
+            scale        = head_dim ** -0.5
+            orig_dtype   = query.dtype
+            output       = torch.empty_like(query)
+            dev          = query.device
+
+            for i in range(batch_size):
+                ctx_len = int(context_lens[i].item())
+                q_start = int(query_start_loc[i].item())
+                q_end   = int(query_start_loc[i + 1].item())
+                q_len   = q_end - q_start
+
+                q_i = query[q_start:q_end]   # [q_len, q_h,  d]
+                k_i = key  [q_start:q_end]   # [q_len, kv_h, d]
+                v_i = value[q_start:q_end]
+
+                # CCCL spread_out_items_per_thread adaptive tile sizing.
+                #
+                # Two constraints compete:
+                # 1. Memory: score tensor [kv_h, gqa, q_len, tile_sz] × 4 ≤ budget
+                # 2. Iteration count: want ~4-8 tiles to minimize Python overhead
+                #
+                # CCCL dispatch_transform.cuh::spread_out_items_per_thread:
+                #   items = ceil_div(num_items, sm_count * threads * occupancy)
+                #   items = clamp(items, min_items, max_items)
+                #
+                # Our translation: tile_sz = max context tokens / target_tiles,
+                # then clamp by memory budget.
+                score_row_bytes = num_kv_heads * gqa_ratio * q_len * 4
+                if score_row_bytes > 0:
+                    mem_max_tokens = _SMEM_BUDGET_BYTES // score_row_bytes
+                    mem_max_tokens = (mem_max_tokens // block_size) * block_size
+                else:
+                    mem_max_tokens = block_size * 256
+
+                total_kv_tokens = ctx_len + q_len
+                # spread_out: target 4 tiles for context, 4 for current chunk
+                spread_tile = max(block_size,
+                                  (total_kv_tokens + 3) // 4)
+                # Round to block_size
+                spread_tile = (spread_tile // block_size) * block_size
+                spread_tile = max(spread_tile, block_size)
+                # Clamp by memory budget
+                tile_sz = min(spread_tile, mem_max_tokens)
+                tile_sz = max(tile_sz, block_size)  # floor
+
+                # Q reshaped and scaled once; held for all K-tiles.
+                # [kv_h, gqa, q_len, d] fp32  —  24 MB for q_len=4096, d=256
+                q_seq = (q_i.permute(1, 0, 2)
+                           .float()
+                           .view(num_kv_heads, gqa_ratio, q_len, head_dim)
+                           .mul_(scale))
+
+                # Flash-Attention online-softmax accumulators.
+                # m, l : [kv_h, gqa, q_len]     fp32  — <0.1 MB
+                # o    : [kv_h, gqa, q_len, d]  fp32  — 24 MB
+                m = torch.full((num_kv_heads, gqa_ratio, q_len),
+                               float('-inf'), dtype=torch.float32, device=dev)
+                l = torch.zeros_like(m)
+                o = torch.zeros((num_kv_heads, gqa_ratio, q_len, head_dim),
+                                dtype=torch.float32, device=dev)
+
+                # --------------------------------------------------------------
+                # Phase 1 — context tokens (positions 0 … ctx_len-1).
+                #
+                # Every context key has absolute position < ctx_len; every
+                # query has position ≥ ctx_len.  k_pos < q_pos is always True
+                # → no causal mask needed for pure context tiles.
+                # --------------------------------------------------------------
+                # Convert token-based tile_sz to block count for iteration
+                blocks_per_tile = tile_sz // block_size
+
+                if ctx_len > 0:
+                    num_ctx_blocks = (ctx_len + block_size - 1) // block_size
+                    if num_ctx_blocks > block_tables.shape[1]:
+                        print(
+                            f"[paged_attn WARNING] seq {i}: num_ctx_blocks={num_ctx_blocks} "
+                            f"> block_tables.shape[1]={block_tables.shape[1]}, ctx_len={ctx_len}. "
+                            "Block table is undersized (prefix_cache_hit bug). "
+                            "Capping context to available blocks — attention may be incorrect.",
+                            file=sys.stderr, flush=True)
+                        num_ctx_blocks = block_tables.shape[1]
+                    for tile_blk in range(0, num_ctx_blocks, blocks_per_tile):
+                        blk_end = min(tile_blk + blocks_per_tile, num_ctx_blocks)
+                        blk_ids = block_tables[i, tile_blk:blk_end]
+
+                        # Gather K/V for this tile.
+                        # key_cache  [blk_ids]: [n, kv_h, d//x, blk_sz, x]
+                        # value_cache[blk_ids]: [n, kv_h, d,    blk_sz]
+                        k_tile = (key_cache[blk_ids]
+                                  .permute(0, 3, 1, 2, 4)
+                                  .contiguous()
+                                  .view(-1, num_kv_heads, head_dim))
+                        v_tile = (value_cache[blk_ids]
+                                  .permute(0, 3, 1, 2)
+                                  .contiguous()
+                                  .view(-1, num_kv_heads, head_dim))
+
+                        # Trim padding in the last block of the tile.
+                        valid = (min(blk_end * block_size, ctx_len)
+                                 - tile_blk * block_size)
+                        k_tile = k_tile[:valid]   # [valid, kv_h, d]
+                        v_tile = v_tile[:valid]
+
+                        # k_t: [kv_h, 1, d, valid]   (broadcast over gqa_ratio)
+                        # v_t: [kv_h, 1, valid, d]
+                        k_t = (k_tile.permute(1, 0, 2)
+                                     .unsqueeze(1)
+                                     .transpose(-1, -2)
+                                     .float())
+                        v_t = (v_tile.permute(1, 0, 2)
+                                     .unsqueeze(1)
+                                     .float())
+                        del k_tile, v_tile
+
+                        # Scores: [kv_h, gqa, q_len, valid]
+                        s = torch.matmul(q_seq, k_t)
+                        del k_t
+                        # No causal mask: all context keys precede all queries.
+
+                        # Online softmax update — Flash-Attention Algorithm 1.
+                        # exp_s = s - new_max  (in-place exp after del s)
+                        m_blk = s.amax(dim=-1)
+                        m_new = torch.maximum(m, m_blk)
+                        exp_s = s - m_new.unsqueeze(-1)
+                        del s
+                        exp_s.exp_()
+                        corr  = torch.exp(m - m_new)
+                        m.copy_(m_new)
+                        del m_blk, m_new
+                        l.mul_(corr).add_(exp_s.sum(dim=-1))
+                        o.mul_(corr.unsqueeze(-1)).add_(
+                            torch.matmul(exp_s, v_t))
+                        del exp_s, v_t, corr
+
+                # --------------------------------------------------------------
+                # Phase 2 — current-chunk tokens (positions ctx_len … ctx_len+q_len-1).
+                #
+                # Causal mask: query at relative position j sees key at relative
+                # position k only when k ≤ j.  Tiles of tile_sz tokens each.
+                # --------------------------------------------------------------
+                for kc_start in range(0, q_len, tile_sz):
+                    kc_end = min(kc_start + tile_sz, q_len)
+                    kc_len = kc_end - kc_start
+
+                    k_blk = k_i[kc_start:kc_end]   # [kc_len, kv_h, d]
+                    v_blk = v_i[kc_start:kc_end]
+
+                    k_t = (k_blk.permute(1, 0, 2)
+                                 .unsqueeze(1)
+                                 .transpose(-1, -2)
+                                 .float())   # [kv_h, 1, d, kc_len]
+                    v_t = (v_blk.permute(1, 0, 2)
+                                 .unsqueeze(1)
+                                 .float())   # [kv_h, 1, kc_len, d]
+
+                    s = torch.matmul(q_seq, k_t)   # [kv_h, gqa, q_len, kc_len]
+                    del k_t
+
+                    # Causal mask: key at (kc_start+k) must not exceed query j.
+                    k_rel = torch.arange(kc_start, kc_end, device=dev)
+                    q_rel = torch.arange(q_len, device=dev)
+                    mask  = k_rel.unsqueeze(0) > q_rel.unsqueeze(1)  # [q_len, kc_len]
+                    s.masked_fill_(mask.unsqueeze(0).unsqueeze(0), float('-inf'))
+                    del mask, k_rel, q_rel
+
+                    # Online softmax update (identical to context phase).
+                    m_blk = s.amax(dim=-1)
+                    m_new = torch.maximum(m, m_blk)
+                    exp_s = s - m_new.unsqueeze(-1)
+                    del s
+                    exp_s.exp_()
+                    corr  = torch.exp(m - m_new)
+                    m.copy_(m_new)
+                    del m_blk, m_new
+                    l.mul_(corr).add_(exp_s.sum(dim=-1))
+                    o.mul_(corr.unsqueeze(-1)).add_(
+                        torch.matmul(exp_s, v_t))
+                    del exp_s, v_t, corr
+
+                # --------------------------------------------------------------
+                # Finalize: normalize running output by normalization factor.
+                # o: [kv_h, gqa, q_len, d]  →  [q_len, q_h, d]
+                # --------------------------------------------------------------
+                o.div_(l.unsqueeze(-1))
+                output[q_start:q_end] = (
+                    o.view(num_q_heads, q_len, head_dim)
+                     .permute(1, 0, 2)
+                     .to(orig_dtype)
+                )
+
+        except Exception as e:
+            print(f"[paged_attn ERROR] {type(e).__name__}: {e}",
+                  file=sys.stderr, flush=True)
+            traceback.print_exc(file=sys.stderr)
+            raise
+        return output
+
+    @staticmethod
+    def swap_blocks(
+        src_kv_cache: torch.Tensor,
+        dst_kv_cache: torch.Tensor,
+        src_to_dst: torch.Tensor,
+    ) -> None:
+        src_key_cache = src_kv_cache[0]
+        dst_key_cache = dst_kv_cache[0]
+        ops.swap_blocks(src_key_cache, dst_key_cache, src_to_dst)
+
+        src_value_cache = src_kv_cache[1]
+        dst_value_cache = dst_kv_cache[1]
+        ops.swap_blocks(src_value_cache, dst_value_cache, src_to_dst)
+
+    @staticmethod
+    def copy_blocks(
+        kv_caches: List[torch.Tensor],
+        src_to_dists: torch.Tensor,
+    ) -> None:
+        key_caches = [kv_cache[0] for kv_cache in kv_caches]
+        value_caches = [kv_cache[1] for kv_cache in kv_caches]
+        ops.copy_blocks(key_caches, value_caches, src_to_dists)

qwen3_6_scripts/patch_ops.sh

@@ -0,0 +1,141 @@
+#!/bin/bash
+set -eo pipefail
+# BI-V100 engine patches for Qwen3.6-35B-A3B (Qwen3_5 architecture)
+#
+# All modifications are FULL FILE REPLACEMENTS — no AST patch scripts.
+# Each file was read in full from the base image vllm source, modified
+# with the necessary fixes, and placed here as a complete copy.
+#
+# Base image: git.modelhub.org.cn:9443/enginex-iluvatar/bi100-3.2.3-x86-ubuntu20.04-py3.10-poc-llm-infer:v1.2.3
+# vllm install path: /usr/local/corex/lib/python3/dist-packages/vllm/
+
+# CRITICAL: cd into this script's directory so all ./relative paths work
+# regardless of WORKDIR in Dockerfile or caller's cwd.
+cd "$(dirname "$0")"
+echo "[patch_ops] working directory: $(pwd)"
+
+VLLM=/usr/local/corex/lib/python3/dist-packages/vllm
+VLLM64=/usr/local/corex/lib64/python3/dist-packages/vllm
+
+# Deploy to ALL existing vllm paths — Python may load from either one
+# depending on PYTHONPATH ordering and namespace package resolution.
+TARGETS=()
+if [ -d "$VLLM" ]; then
+    TARGETS+=("$VLLM")
+fi
+if [ -d "$VLLM64" ]; then
+    TARGETS+=("$VLLM64")
+fi
+
+if [ ${#TARGETS[@]} -eq 0 ]; then
+    echo "[patch_ops] ERROR: vllm not found at lib or lib64 path"
+    exit 1
+fi
+
+echo "[patch_ops] vllm paths found: ${TARGETS[*]}"
+
+# Helper: copy file to all target vllm roots
+deploy() {
+    local src="$1"
+    local rel_dst="$2"  # relative path within vllm, e.g. "attention/ops/paged_attn.py"
+    for V in "${TARGETS[@]}"; do
+        local dst="$V/$rel_dst"
+        mkdir -p "$(dirname "$dst")"
+        cp "$src" "$dst"
+    done
+}
+
+# --- _custom_ops.py: SMEM 48KB fix + hardware ops bindings -------------------
+# Base image returns 32KB (32768) for get_max_shared_memory_per_block, but
+# BI-V100 actually has 48KB (49152) confirmed via ixsmi. This limits Triton
+# tile sizes and ixformer internal allocations if not corrected.
+# CCCL GridEvenShare test (catch2_test_grid_even_share.cu) validates that
+# work distribution depends on correct hardware parameters — wrong SMEM
+# means wrong tile_size means wrong grid_size.
+# FULL FILE REPLACEMENT.
+deploy ./_custom_ops.py "_custom_ops.py"
+echo "[patch_ops] _custom_ops.py → / (SMEM 32KB→48KB fix)"
+
+# --- paged_attn.py: pure-PyTorch attention fallback --------------------------
+deploy ./paged_attn.py "attention/ops/paged_attn.py"
+echo "[patch_ops] paged_attn.py → attention/ops/"
+
+# --- prefix_prefill.py: Triton-free prefix attention -------------------------
+deploy ./prefix_prefill.py "attention/ops/prefix_prefill.py"
+echo "[patch_ops] prefix_prefill.py → attention/ops/"
+
+# --- model_runner.py: prefix_cache_hit fix -----------------------------------
+deploy ./model_runner.py "worker/model_runner.py"
+echo "[patch_ops] model_runner.py → worker/"
+
+# --- xformers.py: head_dim>128 fallback + Q-tiling --------------------------
+deploy ./xformers.py "attention/backends/xformers.py"
+echo "[patch_ops] xformers.py → attention/backends/"
+
+# --- arg_utils.py: disable auto chunked-prefill for 32K+ --------------------
+deploy ./arg_utils.py "engine/arg_utils.py"
+echo "[patch_ops] arg_utils.py → engine/"
+
+# --- logits_processor.py: seq_groups=None guard ------------------------------
+deploy ./logits_processor.py "model_executor/layers/logits_processor.py"
+echo "[patch_ops] logits_processor.py → model_executor/layers/"
+
+# --- sampler.py: CCCL-ported top-k fast path for sampling --------------------
+deploy ./sampler.py "model_executor/layers/sampler.py"
+echo "[patch_ops] sampler.py → model_executor/layers/"
+
+# --- transformers: Qwen3_5 tokenizer / model files --------------------------
+# NOTE: patch_transformers_qwen3_5.py is the ONLY remaining patch script.
+# It modifies pip-installed transformers' configuration_auto.py and __init__.py
+# to register qwen3_5/qwen3_5_moe. These files come from pip (version-specific)
+# so we can't pre-copy them — the patch script inserts lines after known anchors.
+pip install transformers==4.55.3 -i https://pypi.tuna.tsinghua.edu.cn/simple 2>/dev/null || \
+pip install transformers==4.55.3 2>/dev/null || \
+echo "[patch_ops] WARNING: pip install transformers failed, using pre-installed version"
+cp -r ./qwen3_5 /usr/local/lib/python3.10/site-packages/transformers/models/
+cp -r ./qwen3_5_moe /usr/local/lib/python3.10/site-packages/transformers/models/
+python3 ./patch_transformers_qwen3_5.py
+echo "[patch_ops] transformers Qwen3_5 models installed"
+
+# --- vllm model: Qwen3.6 (Qwen3_5 arch) ------------------------------------
+for V in "${TARGETS[@]}"; do
+    cp ./mamba_cache.py "$V/model_executor/models/"
+done
+deploy ./qwen3_5.py "model_executor/models/qwen3_5.py"
+deploy ./registry.py "model_executor/models/registry.py"
+echo "[patch_ops] qwen3_5.py + registry.py deployed"
+
+# --- paged_attention_v2_pytorch.py: PyTorch V2 attention fallback ------------
+for V in "${TARGETS[@]}"; do
+    cp ./paged_attention_v2_pytorch.py "$V/paged_attention_v2_pytorch.py"
+done
+cp ./paged_attention_v2_pytorch.py /workspace/paged_attention_v2_pytorch.py
+echo "[patch_ops] paged_attention_v2_pytorch.py → all paths + /workspace/"
+
+# --- sequence.py: fix completion_tokens inflation ----------------------------
+deploy ./sequence.py "sequence.py"
+echo "[patch_ops] sequence.py → /"
+
+# --- scheduler.py: record num_cached_tokens ---------------------------------
+deploy ./scheduler.py "core/scheduler.py"
+echo "[patch_ops] scheduler.py → core/"
+
+# --- tool parser: Qwen3 XML tool call format --------------------------------
+for V in "${TARGETS[@]}"; do
+    cp ./qwen3coder_tool_parser.py "$V/entrypoints/openai/tool_parsers/"
+    cp ./tool_parsers_init.py "$V/entrypoints/openai/tool_parsers/__init__.py"
+done
+echo "[patch_ops] qwen3_coder tool parser deployed"
+
+# --- reasoning parser: Qwen3 <think>...</think> split -----------------------
+for V in "${TARGETS[@]}"; do
+    cp -r ./reasoning "$V/"
+    cp ./protocol.py "$V/entrypoints/openai/protocol.py"
+    cp ./cli_args.py "$V/entrypoints/openai/cli_args.py"
+    cp ./serving_chat.py "$V/entrypoints/openai/serving_chat.py"
+    cp ./api_server.py "$V/entrypoints/openai/api_server.py"
+    cp ./chat_utils.py "$V/entrypoints/chat_utils.py"
+done
+echo "[patch_ops] reasoning parser + serving files installed"
+
+echo "[patch_ops] DONE — all patches applied via full file replacement"

qwen3_6_scripts/patch_transformers_qwen3_5.py

@@ -0,0 +1,117 @@
+"""
+Patches transformers 4.55.3 to register qwen3_5 and qwen3_5_moe model types.
+
+Deploy steps on the remote machine:
+  1. cp -r modified_scripts/qwen3_5     /usr/local/lib/python3.10/site-packages/transformers/models/qwen3_5
+  2. cp -r modified_scripts/qwen3_5_moe /usr/local/lib/python3.10/site-packages/transformers/models/qwen3_5_moe
+  3. python3 modified_scripts/patch_transformers_qwen3_5.py
+
+Target: pip-installed transformers at /usr/local/lib/python3.10/site-packages/transformers/
+(Not the corex pre-installed path at /usr/local/corex/lib64/python3/dist-packages/)
+"""
+
+import sys
+
+TRANSFORMERS_ROOT = "/usr/local/lib/python3.10/site-packages/transformers"
+AUTO_CONFIG = f"{TRANSFORMERS_ROOT}/models/auto/configuration_auto.py"
+MODELS_INIT = f"{TRANSFORMERS_ROOT}/models/__init__.py"
+
+
+def patch_file(path, replacements):
+    with open(path, "r") as f:
+        content = f.read()
+
+    patched = False
+    for old, new in replacements:
+        if new in content:
+            print(f"  [skip] already patched: {repr(new[:60])}")
+            continue
+        if old not in content:
+            print(f"  [warn] anchor not found: {repr(old[:60])}")
+            continue
+        content = content.replace(old, new, 1)
+        patched = True
+        print(f"  [ok]   inserted after: {repr(old[:60])}")
+
+    if patched:
+        with open(path, "w") as f:
+            f.write(content)
+
+
+def main():
+    print(f"=== Patching {AUTO_CONFIG} ===")
+    patch_file(AUTO_CONFIG, [
+        # CONFIG_MAPPING_NAMES: insert qwen3_5 + qwen3_5_moe right after qwen3
+        (
+            '("qwen3", "Qwen3Config"),',
+            '("qwen3", "Qwen3Config"),\n        ("qwen3_5", "Qwen3_5Config"),\n        ("qwen3_5_moe", "Qwen3_5MoeConfig"),',
+        ),
+        (
+            '("qwen3", "Qwen3Config")\n',
+            '("qwen3", "Qwen3Config"),\n        ("qwen3_5", "Qwen3_5Config"),\n        ("qwen3_5_moe", "Qwen3_5MoeConfig"),\n',
+        ),
+        # MODEL_NAMES_MAPPING (model_type -> human readable name)
+        (
+            '("qwen3", "Qwen3"),',
+            '("qwen3", "Qwen3"),\n        ("qwen3_5", "Qwen3_5"),\n        ("qwen3_5_moe", "Qwen3_5_MoE"),',
+        ),
+        (
+            '("qwen3", "Qwen3")\n',
+            '("qwen3", "Qwen3"),\n        ("qwen3_5", "Qwen3_5"),\n        ("qwen3_5_moe", "Qwen3_5_MoE"),\n',
+        ),
+    ])
+
+    print(f"\n=== Patching {MODELS_INIT} ===")
+    patch_file(MODELS_INIT, [
+        (
+            "from .qwen3 import *\n",
+            "from .qwen3 import *\n    from .qwen3_5 import *\n    from .qwen3_5_moe import *\n",
+        ),
+    ])
+
+    # Verification
+    print("\n=== Verification ===")
+    try:
+        import importlib.util, types
+
+        def _load_config_mod(module_name, file_path):
+            spec = importlib.util.spec_from_file_location(module_name, file_path)
+            mod = importlib.util.module_from_spec(spec)
+            mod.__package__ = ".".join(module_name.split(".")[:-1])
+            pkg = sys.modules.setdefault("transformers", types.ModuleType("transformers"))
+            pkg.__path__ = [TRANSFORMERS_ROOT]
+            cu = sys.modules.setdefault(
+                "transformers.configuration_utils", types.ModuleType("transformers.configuration_utils"))
+            class _PC:
+                def __init__(self, **kwargs): pass
+            cu.PretrainedConfig = _PC
+            for sub in ("transformers.models", f"transformers.models.{module_name.split('.')[-2]}"):
+                m = sys.modules.setdefault(sub, types.ModuleType(sub))
+                m.__path__ = [TRANSFORMERS_ROOT]
+            spec.loader.exec_module(mod)
+            return mod
+
+        mod27 = _load_config_mod(
+            "transformers.models.qwen3_5.configuration_qwen3_5",
+            f"{TRANSFORMERS_ROOT}/models/qwen3_5/configuration_qwen3_5.py",
+        )
+        cfg = mod27.Qwen3_5Config()
+        print(f"  Qwen3_5Config() smoke-test OK     (model_type={cfg.model_type})")
+
+        mod35 = _load_config_mod(
+            "transformers.models.qwen3_5_moe.configuration_qwen3_5_moe",
+            f"{TRANSFORMERS_ROOT}/models/qwen3_5_moe/configuration_qwen3_5_moe.py",
+        )
+        moe_cfg = mod35.Qwen3_5MoeConfig()
+        print(f"  Qwen3_5MoeConfig() smoke-test OK  (model_type={moe_cfg.model_type})")
+        t = moe_cfg.text_config
+        print(f"    num_experts={t.num_experts}, top_k={t.num_experts_per_tok}, "
+              f"shared={t.shared_expert_intermediate_size}, layers={t.num_hidden_layers}")
+    except Exception as e:
+        print(f"  [warn] smoke-test failed (may be fine at runtime): {e}")
+
+    print("\nDone.")
+
+
+if __name__ == "__main__":
+    main()

qwen3_6_scripts/prefix_prefill.py

@@ -0,0 +1,900 @@
+# The kernels in this file are adapted from LightLLM's context_attention_fwd:
+# https://github.com/ModelTC/lightllm/blob/main/lightllm/models/llama/triton_kernel/context_flashattention_nopad.py
+
+import torch
+import triton
+import triton.language as tl
+
+from vllm.platforms import current_platform
+
+if triton.__version__ >= "2.1.0":
+
+    @triton.jit
+    def _fwd_kernel(
+        Q,
+        K,
+        V,
+        K_cache,
+        V_cache,
+        B_Loc,
+        sm_scale,
+        k_scale,
+        v_scale,
+        B_Start_Loc,
+        B_Seqlen,
+        B_Ctxlen,
+        block_size,
+        x,
+        Out,
+        stride_b_loc_b,
+        stride_b_loc_s,
+        stride_qbs,
+        stride_qh,
+        stride_qd,
+        stride_kbs,
+        stride_kh,
+        stride_kd,
+        stride_vbs,
+        stride_vh,
+        stride_vd,
+        stride_obs,
+        stride_oh,
+        stride_od,
+        stride_k_cache_bs,
+        stride_k_cache_h,
+        stride_k_cache_d,
+        stride_k_cache_bl,
+        stride_k_cache_x,
+        stride_v_cache_bs,
+        stride_v_cache_h,
+        stride_v_cache_d,
+        stride_v_cache_bl,
+        num_queries_per_kv: int,
+        BLOCK_M: tl.constexpr,
+        BLOCK_DMODEL: tl.constexpr,  # head size
+        BLOCK_DMODEL_PADDED: tl.constexpr,  # head size padded to a power of 2
+        BLOCK_N: tl.constexpr,
+        SLIDING_WINDOW: tl.constexpr,
+    ):
+        cur_batch = tl.program_id(0)
+        cur_head = tl.program_id(1)
+        start_m = tl.program_id(2)
+
+        cur_kv_head = cur_head // num_queries_per_kv
+
+        cur_batch_ctx_len = tl.load(B_Ctxlen + cur_batch)
+        cur_batch_seq_len = tl.load(B_Seqlen + cur_batch)
+        cur_batch_in_all_start_index = tl.load(B_Start_Loc + cur_batch)
+        cur_batch_query_len = cur_batch_seq_len - cur_batch_ctx_len
+
+        # start position inside of the query
+        # generally, N goes over kv, while M goes over query_len
+        block_start_loc = BLOCK_M * start_m
+
+        # initialize offsets
+        # [N]; starts at 0
+        offs_n = tl.arange(0, BLOCK_N)
+        # [D]; starts at 0
+        offs_d = tl.arange(0, BLOCK_DMODEL_PADDED)
+        # [M]; starts at current position in query
+        offs_m = start_m * BLOCK_M + tl.arange(0, BLOCK_M)
+        # [M,D]
+        off_q = (
+            (cur_batch_in_all_start_index + offs_m[:, None]) * stride_qbs +
+            cur_head * stride_qh + offs_d[None, :] * stride_qd)
+
+        dim_mask = tl.where(
+            tl.arange(0, BLOCK_DMODEL_PADDED) < BLOCK_DMODEL, 1,
+            0).to(tl.int1)  # [D]
+
+        q = tl.load(Q + off_q,
+                    mask=dim_mask[None, :] &
+                    (offs_m[:, None] < cur_batch_query_len),
+                    other=0.0)  # [M,D]
+
+        # initialize pointer to m and l
+        m_i = tl.zeros([BLOCK_M], dtype=tl.float32) - float("inf")  # [M]
+        l_i = tl.zeros([BLOCK_M], dtype=tl.float32)  # [M]
+        acc = tl.zeros([BLOCK_M, BLOCK_DMODEL_PADDED],
+                       dtype=tl.float32)  # [M,D]
+
+        # compute query against context (no causal mask here)
+        for start_n in range(0, cur_batch_ctx_len, BLOCK_N):
+            start_n = tl.multiple_of(start_n, BLOCK_N)
+            # -- compute qk ----
+            bn = tl.load(B_Loc + cur_batch * stride_b_loc_b +
+                         ((start_n + offs_n) // block_size) * stride_b_loc_s,
+                         mask=(start_n + offs_n) < cur_batch_ctx_len,
+                         other=0)  # [N]
+            # [D,N]
+            off_k = (bn[None, :] * stride_k_cache_bs +
+                     cur_kv_head * stride_k_cache_h +
+                     (offs_d[:, None] // x) * stride_k_cache_d +
+                     ((start_n + offs_n[None, :]) % block_size) *
+                     stride_k_cache_bl +
+                     (offs_d[:, None] % x) * stride_k_cache_x)
+            # [N,D]
+            off_v = (
+                bn[:, None] * stride_v_cache_bs +
+                cur_kv_head * stride_v_cache_h +
+                offs_d[None, :] * stride_v_cache_d +
+                (start_n + offs_n[:, None]) % block_size * stride_v_cache_bl)
+            k_load = tl.load(K_cache + off_k,
+                             mask=dim_mask[:, None] &
+                             ((start_n + offs_n[None, :]) < cur_batch_ctx_len),
+                             other=0.0)  # [D,N]
+
+            if k_load.dtype.is_fp8():
+                k = (k_load.to(tl.float32) * k_scale).to(q.dtype)
+            else:
+                k = k_load
+
+            qk = tl.zeros([BLOCK_M, BLOCK_N], dtype=tl.float32)  # [M,N]
+            qk += tl.dot(q, k)
+            qk = tl.where((start_n + offs_n[None, :]) < cur_batch_ctx_len, qk,
+                          float("-inf"))
+            qk *= sm_scale
+            if SLIDING_WINDOW > 0:
+                # (cur_batch_ctx_len + offs_m[:, None]) are the positions of
+                # Q entries in sequence
+                # (start_n + offs_n[None, :]) are the positions of
+                # KV entries in sequence
+                # So the condition makes sure each entry in Q only attends
+                # to KV entries not more than SLIDING_WINDOW away.
+                #
+                # We can't use -inf here, because the
+                # sliding window may lead to the entire row being masked.
+                # This then makes m_ij contain -inf, which causes NaNs in
+                # exp().
+                qk = tl.where((cur_batch_ctx_len + offs_m[:, None]) -
+                              (start_n + offs_n[None, :]) < SLIDING_WINDOW, qk,
+                              -10000)
+
+            # -- compute m_ij, p, l_ij
+            m_ij = tl.max(qk, 1)  # [M]
+            p = tl.exp(qk - m_ij[:, None])  # [M,N]
+            l_ij = tl.sum(p, 1)  # [M]
+            # -- update m_i and l_i
+            m_i_new = tl.maximum(m_i, m_ij)  # [M]
+            alpha = tl.exp(m_i - m_i_new)  # [M]
+            beta = tl.exp(m_ij - m_i_new)  # [M]
+            l_i_new = alpha * l_i + beta * l_ij  # [M]
+
+            # -- update output accumulator --
+            # scale p
+            p_scale = beta / l_i_new
+            p = p * p_scale[:, None]
+            # scale acc
+            acc_scale = l_i / l_i_new * alpha
+            acc = acc * acc_scale[:, None]
+            # update acc
+            v_load = tl.load(V_cache + off_v,
+                             mask=dim_mask[None, :] &
+                             ((start_n + offs_n[:, None]) < cur_batch_ctx_len),
+                             other=0.0)  # [N,D]
+            if v_load.dtype.is_fp8():
+                v = (v_load.to(tl.float32) * v_scale).to(q.dtype)
+            else:
+                v = v_load
+            p = p.to(v.dtype)
+
+            acc += tl.dot(p, v)
+            # # update m_i and l_i
+            l_i = l_i_new
+            m_i = m_i_new
+
+        off_k = (offs_n[None, :] * stride_kbs + cur_kv_head * stride_kh +
+                 offs_d[:, None] * stride_kd)
+        off_v = (offs_n[:, None] * stride_vbs + cur_kv_head * stride_vh +
+                 offs_d[None, :] * stride_vd)
+        k_ptrs = K + off_k
+        v_ptrs = V + off_v
+
+        # block_mask is 0 when we're already past the current query length
+        block_mask = tl.where(block_start_loc < cur_batch_query_len, 1, 0)
+
+        # compute query against itself (with causal mask)
+        for start_n in range(0, block_mask * (start_m + 1) * BLOCK_M, BLOCK_N):
+            start_n = tl.multiple_of(start_n, BLOCK_N)
+            # -- compute qk ----
+            k = tl.load(k_ptrs +
+                        (cur_batch_in_all_start_index + start_n) * stride_kbs,
+                        mask=dim_mask[:, None] &
+                        ((start_n + offs_n[None, :]) < cur_batch_query_len),
+                        other=0.0)
+
+            qk = tl.zeros([BLOCK_M, BLOCK_N], dtype=tl.float32)
+            qk += tl.dot(q, k)
+            qk *= sm_scale
+            # apply causal mask
+            qk = tl.where(offs_m[:, None] >= (start_n + offs_n[None, :]), qk,
+                          float("-inf"))
+            if SLIDING_WINDOW > 0:
+                qk = tl.where(
+                    offs_m[:, None] -
+                    (start_n + offs_n[None, :]) < SLIDING_WINDOW, qk, -10000)
+
+            # -- compute m_ij, p, l_ij
+            m_ij = tl.max(qk, 1)
+            p = tl.exp(qk - m_ij[:, None])
+            l_ij = tl.sum(p, 1)
+            # -- update m_i and l_i
+            m_i_new = tl.maximum(m_i, m_ij)
+            alpha = tl.exp(m_i - m_i_new)
+            beta = tl.exp(m_ij - m_i_new)
+            l_i_new = alpha * l_i + beta * l_ij
+            # -- update output accumulator --
+            # scale p
+            p_scale = beta / l_i_new
+            p = p * p_scale[:, None]
+            # scale acc
+            acc_scale = l_i / l_i_new * alpha
+            acc = acc * acc_scale[:, None]
+            # update acc
+            v = tl.load(v_ptrs +
+                        (cur_batch_in_all_start_index + start_n) * stride_vbs,
+                        mask=dim_mask[None, :] &
+                        ((start_n + offs_n[:, None]) < cur_batch_query_len),
+                        other=0.0)
+            p = p.to(v.dtype)
+
+            acc += tl.dot(p, v)
+            # update m_i and l_i
+            l_i = l_i_new
+            m_i = m_i_new
+        # initialize pointers to output
+        off_o = (
+            (cur_batch_in_all_start_index + offs_m[:, None]) * stride_obs +
+            cur_head * stride_oh + offs_d[None, :] * stride_od)
+        out_ptrs = Out + off_o
+        tl.store(out_ptrs,
+                 acc,
+                 mask=dim_mask[None, :] &
+                 (offs_m[:, None] < cur_batch_query_len))
+        return
+
+    @triton.jit
+    def _fwd_kernel_flash_attn_v2(
+        Q,
+        K,
+        V,
+        K_cache,
+        V_cache,
+        B_Loc,
+        sm_scale,
+        B_Start_Loc,
+        B_Seqlen,
+        B_Ctxlen,
+        block_size,
+        x,
+        Out,
+        stride_b_loc_b,
+        stride_b_loc_s,
+        stride_qbs,
+        stride_qh,
+        stride_qd,
+        stride_kbs,
+        stride_kh,
+        stride_kd,
+        stride_vbs,
+        stride_vh,
+        stride_vd,
+        stride_obs,
+        stride_oh,
+        stride_od,
+        stride_k_cache_bs,
+        stride_k_cache_h,
+        stride_k_cache_d,
+        stride_k_cache_bl,
+        stride_k_cache_x,
+        stride_v_cache_bs,
+        stride_v_cache_h,
+        stride_v_cache_d,
+        stride_v_cache_bl,
+        num_queries_per_kv: int,
+        BLOCK_M: tl.constexpr,
+        BLOCK_DMODEL: tl.constexpr,
+        BLOCK_N: tl.constexpr,
+    ):
+        cur_batch = tl.program_id(0)
+        cur_head = tl.program_id(1)
+        start_m = tl.program_id(2)
+
+        cur_kv_head = cur_head // num_queries_per_kv
+
+        cur_batch_ctx_len = tl.load(B_Ctxlen + cur_batch)
+        cur_batch_seq_len = tl.load(B_Seqlen + cur_batch)
+        cur_batch_in_all_start_index = tl.load(B_Start_Loc + cur_batch)
+
+        block_start_loc = BLOCK_M * start_m
+
+        # initialize offsets
+        offs_n = tl.arange(0, BLOCK_N)
+        offs_d = tl.arange(0, BLOCK_DMODEL)
+        offs_m = start_m * BLOCK_M + tl.arange(0, BLOCK_M)
+        off_q = (
+            (cur_batch_in_all_start_index + offs_m[:, None]) * stride_qbs +
+            cur_head * stride_qh + offs_d[None, :] * stride_qd)
+
+        q = tl.load(
+            Q + off_q,
+            mask=offs_m[:, None] < cur_batch_seq_len - cur_batch_ctx_len,
+            other=0.0)
+
+        # # initialize pointer to m and l
+        m_i = tl.zeros([BLOCK_M], dtype=tl.float32) - float("inf")
+        l_i = tl.zeros([BLOCK_M], dtype=tl.float32)
+        acc = tl.zeros([BLOCK_M, BLOCK_DMODEL], dtype=tl.float32)
+
+        for start_n in range(0, cur_batch_ctx_len, BLOCK_N):
+            start_n = tl.multiple_of(start_n, BLOCK_N)
+            # -- compute qk ----
+            bn = tl.load(B_Loc + cur_batch * stride_b_loc_b +
+                         ((start_n + offs_n) // block_size) * stride_b_loc_s,
+                         mask=(start_n + offs_n) < cur_batch_ctx_len,
+                         other=0)
+            off_k = (bn[None, :] * stride_k_cache_bs +
+                     cur_kv_head * stride_k_cache_h +
+                     (offs_d[:, None] // x) * stride_k_cache_d +
+                     ((start_n + offs_n[None, :]) % block_size) *
+                     stride_k_cache_bl +
+                     (offs_d[:, None] % x) * stride_k_cache_x)
+            off_v = (
+                bn[:, None] * stride_v_cache_bs +
+                cur_kv_head * stride_v_cache_h +
+                offs_d[None, :] * stride_v_cache_d +
+                (start_n + offs_n[:, None]) % block_size * stride_v_cache_bl)
+            k = tl.load(K_cache + off_k,
+                        mask=(start_n + offs_n[None, :]) < cur_batch_ctx_len,
+                        other=0.0)
+            qk = tl.zeros([BLOCK_M, BLOCK_N], dtype=tl.float32)
+            qk += tl.dot(q, k)
+            qk = tl.where((start_n + offs_n[None, :]) < cur_batch_ctx_len, qk,
+                          float("-inf"))
+            qk *= sm_scale
+
+            # -- compute m_ij, p, l_ij
+            m_ij = tl.max(qk, 1)
+            m_i_new = tl.maximum(m_i, m_ij)
+            p = tl.math.exp(qk - m_i_new[:, None])
+            l_ij = tl.sum(p, 1)
+            # -- update m_i and l_i
+
+            alpha = tl.math.exp(m_i - m_i_new)
+            l_i_new = alpha * l_i + l_ij
+            # -- update output accumulator --
+            # scale p
+            # scale acc
+            acc_scale = alpha
+            # acc_scale = l_i / l_i_new * alpha
+            acc = acc * acc_scale[:, None]
+            # update acc
+            v = tl.load(V_cache + off_v,
+                        mask=(start_n + offs_n[:, None]) < cur_batch_ctx_len,
+                        other=0.0)
+
+            p = p.to(v.dtype)
+            acc += tl.dot(p, v)
+            # update m_i and l_i
+            l_i = l_i_new
+            m_i = m_i_new
+
+        off_k = (offs_n[None, :] * stride_kbs + cur_kv_head * stride_kh +
+                 offs_d[:, None] * stride_kd)
+        off_v = (offs_n[:, None] * stride_vbs + cur_kv_head * stride_vh +
+                 offs_d[None, :] * stride_vd)
+        k_ptrs = K + off_k
+        v_ptrs = V + off_v
+
+        block_mask = tl.where(
+            block_start_loc < cur_batch_seq_len - cur_batch_ctx_len, 1, 0)
+
+        for start_n in range(0, block_mask * (start_m + 1) * BLOCK_M, BLOCK_N):
+            start_n = tl.multiple_of(start_n, BLOCK_N)
+            # -- compute qk ----
+            k = tl.load(k_ptrs +
+                        (cur_batch_in_all_start_index + start_n) * stride_kbs,
+                        mask=(start_n + offs_n[None, :]) <
+                        cur_batch_seq_len - cur_batch_ctx_len,
+                        other=0.0)
+
+            qk = tl.zeros([BLOCK_M, BLOCK_N], dtype=tl.float32)
+            qk += tl.dot(q, k)
+            qk *= sm_scale
+            qk = tl.where(offs_m[:, None] >= (start_n + offs_n[None, :]), qk,
+                          float("-inf"))
+
+            # -- compute m_ij, p, l_ij
+            m_ij = tl.max(qk, 1)
+            m_i_new = tl.maximum(m_i, m_ij)
+            p = tl.math.exp(qk - m_i_new[:, None])
+            l_ij = tl.sum(p, 1)
+            # -- update m_i and l_i
+
+            alpha = tl.math.exp(m_i - m_i_new)
+            l_i_new = alpha * l_i + l_ij
+            # -- update output accumulator --
+            # scale p
+            # scale acc
+            acc_scale = alpha
+            # acc_scale = l_i / l_i_new * alpha
+            acc = acc * acc_scale[:, None]
+            # update acc
+            v = tl.load(v_ptrs +
+                        (cur_batch_in_all_start_index + start_n) * stride_vbs,
+                        mask=(start_n + offs_n[:, None]) <
+                        cur_batch_seq_len - cur_batch_ctx_len,
+                        other=0.0)
+
+            p = p.to(v.dtype)
+            acc += tl.dot(p, v)
+            # update m_i and l_i
+            l_i = l_i_new
+            m_i = m_i_new
+
+        # BUG FIX: v2 kernel accumulates unnormalized softmax weights.
+        # Without this final division, output = sum(softmax_unnorm * V)
+        # instead of the correct sum(softmax_normalized * V).
+        # v1 kernel does online normalization inside the loop (p_scale/acc_scale).
+        # v2 defers normalization — it MUST happen here.
+        acc = acc / l_i[:, None]
+        # initialize pointers to output
+        off_o = (
+            (cur_batch_in_all_start_index + offs_m[:, None]) * stride_obs +
+            cur_head * stride_oh + offs_d[None, :] * stride_od)
+        out_ptrs = Out + off_o
+        tl.store(out_ptrs,
+                 acc,
+                 mask=offs_m[:, None] < cur_batch_seq_len - cur_batch_ctx_len)
+        return
+
+    @triton.jit
+    def _fwd_kernel_alibi(
+        Q,
+        K,
+        V,
+        K_cache,
+        V_cache,
+        B_Loc,
+        sm_scale,
+        k_scale,
+        v_scale,
+        B_Start_Loc,
+        B_Seqlen,
+        B_Ctxlen,
+        Alibi_slopes,
+        block_size,
+        x,
+        Out,
+        stride_b_loc_b,
+        stride_b_loc_s,
+        stride_qbs,
+        stride_qh,
+        stride_qd,
+        stride_kbs,
+        stride_kh,
+        stride_kd,
+        stride_vbs,
+        stride_vh,
+        stride_vd,
+        stride_obs,
+        stride_oh,
+        stride_od,
+        stride_k_cache_bs,
+        stride_k_cache_h,
+        stride_k_cache_d,
+        stride_k_cache_bl,
+        stride_k_cache_x,
+        stride_v_cache_bs,
+        stride_v_cache_h,
+        stride_v_cache_d,
+        stride_v_cache_bl,
+        num_queries_per_kv: int,
+        BLOCK_M: tl.constexpr,
+        BLOCK_DMODEL: tl.constexpr,  # head size
+        BLOCK_DMODEL_PADDED: tl.constexpr,  # head size padded to a power of 2
+        BLOCK_N: tl.constexpr,
+    ):
+        # attn_bias[]
+        cur_batch = tl.program_id(0)
+        cur_head = tl.program_id(1)
+        start_m = tl.program_id(2)
+
+        cur_kv_head = cur_head // num_queries_per_kv
+
+        # cur_batch_seq_len: the length of prompts
+        # cur_batch_ctx_len: the length of prefix
+        # cur_batch_in_all_start_index: the start id of the dim=0
+        cur_batch_ctx_len = tl.load(B_Ctxlen + cur_batch)
+        cur_batch_seq_len = tl.load(B_Seqlen + cur_batch)
+        cur_batch_in_all_start_index = tl.load(B_Start_Loc + cur_batch)
+
+        block_start_loc = BLOCK_M * start_m
+
+        # initialize offsets
+        offs_n = tl.arange(0, BLOCK_N)
+        offs_d = tl.arange(0, BLOCK_DMODEL_PADDED)
+        offs_m = start_m * BLOCK_M + tl.arange(0, BLOCK_M)
+        off_q = (
+            (cur_batch_in_all_start_index + offs_m[:, None]) * stride_qbs +
+            cur_head * stride_qh + offs_d[None, :] * stride_qd)
+
+        dim_mask = tl.where(
+            tl.arange(0, BLOCK_DMODEL_PADDED) < BLOCK_DMODEL, 1, 0).to(tl.int1)
+
+        q = tl.load(Q + off_q,
+                    mask=dim_mask[None, :] &
+                    (offs_m[:, None] < cur_batch_seq_len - cur_batch_ctx_len),
+                    other=0.0)
+
+        # # initialize pointer to m and l
+        m_i = tl.zeros([BLOCK_M], dtype=tl.float32) - float("inf")
+        l_i = tl.zeros([BLOCK_M], dtype=tl.float32)
+        acc = tl.zeros([BLOCK_M, BLOCK_DMODEL_PADDED], dtype=tl.float32)
+
+        alibi_slope = tl.load(Alibi_slopes + cur_head)
+        alibi_start_q = tl.arange(
+            0, BLOCK_M) + block_start_loc + cur_batch_ctx_len
+        alibi_start_k = 0
+        for start_n in range(0, cur_batch_ctx_len, BLOCK_N):
+            start_n = tl.multiple_of(start_n, BLOCK_N)
+            # -- compute qk ----
+            bn = tl.load(B_Loc + cur_batch * stride_b_loc_b +
+                         ((start_n + offs_n) // block_size) * stride_b_loc_s,
+                         mask=(start_n + offs_n) < cur_batch_ctx_len,
+                         other=0)
+            off_k = (bn[None, :] * stride_k_cache_bs +
+                     cur_kv_head * stride_k_cache_h +
+                     (offs_d[:, None] // x) * stride_k_cache_d +
+                     ((start_n + offs_n[None, :]) % block_size) *
+                     stride_k_cache_bl +
+                     (offs_d[:, None] % x) * stride_k_cache_x)
+            off_v = (
+                bn[:, None] * stride_v_cache_bs +
+                cur_kv_head * stride_v_cache_h +
+                offs_d[None, :] * stride_v_cache_d +
+                (start_n + offs_n[:, None]) % block_size * stride_v_cache_bl)
+            k_load = tl.load(K_cache + off_k,
+                             mask=dim_mask[:, None] &
+                             ((start_n + offs_n[None, :]) < cur_batch_ctx_len),
+                             other=0.0)  # [D,N]
+
+            if k_load.dtype.is_fp8():
+                k = (k_load.to(tl.float32) * k_scale).to(q.dtype)
+            else:
+                k = k_load
+
+            qk = tl.zeros([BLOCK_M, BLOCK_N], dtype=tl.float32)
+            qk += tl.dot(q, k)
+            qk = tl.where((start_n + offs_n[None, :]) < cur_batch_ctx_len, qk,
+                          float("-inf"))
+            qk *= sm_scale
+
+            # load alibi
+            alibi = (tl.arange(0, BLOCK_N)[None, :] + alibi_start_k -
+                     alibi_start_q[:, None]) * alibi_slope
+            alibi = tl.where(
+                (alibi <= 0) & (alibi_start_q[:, None] < cur_batch_seq_len),
+                alibi, float("-inf"))
+            qk += alibi
+            alibi_start_k += BLOCK_N
+
+            # -- compute m_ij, p, l_ij
+            m_ij = tl.max(qk, 1)
+            m_i_new = tl.maximum(m_i, m_ij)
+            p = tl.math.exp(qk - m_i_new[:, None])
+            l_ij = tl.sum(p, 1)
+            # -- update m_i and l_i
+
+            alpha = tl.math.exp(m_i - m_i_new)
+            l_i_new = alpha * l_i + l_ij
+            # -- update output accumulator --
+            # scale p
+            # scale acc
+            acc_scale = alpha
+            # acc_scale = l_i / l_i_new * alpha
+            acc = acc * acc_scale[:, None]
+            # update acc
+            v_load = tl.load(V_cache + off_v,
+                             mask=dim_mask[None, :] &
+                             ((start_n + offs_n[:, None]) < cur_batch_ctx_len),
+                             other=0.0)
+            if v_load.dtype.is_fp8():
+                v = (v_load.to(tl.float32) * v_scale).to(q.dtype)
+            else:
+                v = v_load
+            p = p.to(v.dtype)
+
+            acc += tl.dot(p, v, allow_tf32=False)
+            # update m_i and l_i
+            l_i = l_i_new
+            m_i = m_i_new
+
+        off_k = (offs_n[None, :] * stride_kbs + cur_kv_head * stride_kh +
+                 offs_d[:, None] * stride_kd)
+        off_v = (offs_n[:, None] * stride_vbs + cur_kv_head * stride_vh +
+                 offs_d[None, :] * stride_vd)
+        k_ptrs = K + off_k
+        v_ptrs = V + off_v
+
+        block_mask = tl.where(
+            block_start_loc < cur_batch_seq_len - cur_batch_ctx_len, 1, 0)
+
+        # init alibi
+        alibi_slope = tl.load(Alibi_slopes + cur_head)
+        alibi_start_q = tl.arange(
+            0, BLOCK_M) + block_start_loc + cur_batch_ctx_len
+        alibi_start_k = cur_batch_ctx_len
+        # # init debugger
+        # offset_db_q = tl.arange(0, BLOCK_M) + block_start_loc
+        # offset_db_k = tl.arange(0, BLOCK_N)
+        # calc q[BLOCK_M, BLOCK_MODEL] mul k[prefix_len: , BLOCK_DMODEL]
+        for start_n in range(0, block_mask * (start_m + 1) * BLOCK_M, BLOCK_N):
+            start_n = tl.multiple_of(start_n, BLOCK_N)
+            # -- compute qk ----
+            k = tl.load(k_ptrs +
+                        (cur_batch_in_all_start_index + start_n) * stride_kbs,
+                        mask=dim_mask[:, None] &
+                        ((start_n + offs_n[None, :]) <
+                         cur_batch_seq_len - cur_batch_ctx_len),
+                        other=0.0)
+
+            qk = tl.zeros([BLOCK_M, BLOCK_N], dtype=tl.float32)
+            qk += tl.dot(q, k, allow_tf32=False)
+            qk *= sm_scale
+            qk = tl.where(offs_m[:, None] >= (start_n + offs_n[None, :]), qk,
+                          float("-inf"))
+
+            # load alibi
+            alibi = (tl.arange(0, BLOCK_N)[None, :] + alibi_start_k -
+                     alibi_start_q[:, None]) * alibi_slope
+            alibi = tl.where(
+                (alibi <= 0) & (alibi_start_q[:, None] < cur_batch_seq_len),
+                alibi, float("-inf"))
+            qk += alibi
+            alibi_start_k += BLOCK_N
+
+            # -- compute m_ij, p, l_ij
+            m_ij = tl.max(qk, 1)
+            m_i_new = tl.maximum(m_i, m_ij)
+            p = tl.math.exp(qk - m_i_new[:, None])
+            l_ij = tl.sum(p, 1)
+            # -- update m_i and l_i
+
+            alpha = tl.math.exp(m_i - m_i_new)
+            l_i_new = alpha * l_i + l_ij
+            # -- update output accumulator --
+            # scale p
+            # scale acc
+            acc_scale = alpha
+            # acc_scale = l_i / l_i_new * alpha
+            acc = acc * acc_scale[:, None]
+            # update acc
+            v = tl.load(v_ptrs +
+                        (cur_batch_in_all_start_index + start_n) * stride_vbs,
+                        mask=dim_mask[None, :] &
+                        ((start_n + offs_n[:, None]) <
+                         cur_batch_seq_len - cur_batch_ctx_len),
+                        other=0.0)
+            p = p.to(v.dtype)
+
+            acc += tl.dot(p, v, allow_tf32=False)
+            # update m_i and l_i
+            l_i = l_i_new
+            m_i = m_i_new
+
+        acc = acc / l_i[:, None]
+
+        # initialize pointers to output
+        off_o = (
+            (cur_batch_in_all_start_index + offs_m[:, None]) * stride_obs +
+            cur_head * stride_oh + offs_d[None, :] * stride_od)
+        out_ptrs = Out + off_o
+        tl.store(out_ptrs,
+                 acc,
+                 mask=dim_mask[None, :] &
+                 (offs_m[:, None] < cur_batch_seq_len - cur_batch_ctx_len))
+        return
+
+    @torch.inference_mode()
+    def context_attention_fwd(q,
+                              k,
+                              v,
+                              o,
+                              kv_cache_dtype: str,
+                              k_cache,
+                              v_cache,
+                              b_loc,
+                              b_start_loc,
+                              b_seq_len,
+                              b_ctx_len,
+                              max_input_len,
+                              k_scale: float = 1.0,
+                              v_scale: float = 1.0,
+                              alibi_slopes=None,
+                              sliding_window=None):
+
+        # CCCL-informed block size selection for BI-V100 (SM=16, 48KB SMEM)
+        #
+        # Key insight from CCCL AgentReduce (agent_reduce.cuh):
+        #   - Q tile stays resident in registers/SMEM across the K/V loop
+        #   - K/V tiles stream through: each iteration loads a new BLOCK_N chunk
+        #   - Therefore BLOCK_N can be larger than BLOCK_M (asymmetric tiling)
+        #   - Larger BLOCK_N = fewer loop iterations = fewer kernel barriers
+        #
+        # SMEM budget (peak, not simultaneous - Triton pipelines K/V loads):
+        #   Q resident: BLOCK_M * head_dim * elem_size  (stays across all iters)
+        #   K per iter: head_dim * BLOCK_N * elem_size   (loaded, consumed, freed)
+        #   softmax:    BLOCK_M * 4 * 2                  (m_i + l_i, fp32)
+        #   Total peak: Q + K + softmax_state
+        #
+        # For BI-V100 with head_dim=128, fp16 (2B):
+        #   BLOCK_M=32, BLOCK_N=64: Q=8KB + K=16KB + ss=256B = 24.25KB (49%)
+        #   BLOCK_M=64, BLOCK_N=64: Q=16KB + K=16KB + ss=512B = 32.5KB (66%)
+        #   BLOCK_M=32, BLOCK_N=128: Q=8KB + K=32KB + ss=256B = 40.25KB (82%)
+        #
+        # CCCL scan tuning reference (tuning_scan.cuh):
+        #   SM100 best: ipt=22, tpb=384 → tile = 8448 elements
+        #   BI-V100 bench best: ipt=22, tpb=384, no_delay → 1.038x
+        #   Maps to: moderate tile, no inter-CTA delay (16 SMs = low contention)
+        #
+        # Strategy: BLOCK_M=32 (small Q tile, high occupancy) +
+        #           BLOCK_N=64 (moderate K sweep, fits SMEM easily)
+        #           This gives 2 CTAs per SM occupancy with 16 SMs = 32 CTAs
+        _is_bi_v100 = not current_platform.has_device_capability(80)
+        if _is_bi_v100:
+            BLOCK = 64       # BLOCK_M for Q tile
+            BLOCK_N = 64     # BLOCK_N for K/V sweep (can differ from BLOCK_M)
+            NUM_WARPS = 4
+        else:
+            BLOCK = 128
+            BLOCK_N = BLOCK  # symmetric for NVIDIA GPUs
+            NUM_WARPS = 8
+
+        # need to reduce num. blocks when using fp32
+        # due to increased use of GPU shared memory
+        if q.dtype is torch.float32:
+            BLOCK = BLOCK // 2
+
+        # Conversion of FP8 Tensor from uint8 storage to
+        # appropriate torch.dtype for interpretation by Triton
+        if "fp8" in kv_cache_dtype:
+            assert (k_cache.dtype == torch.uint8)
+            assert (v_cache.dtype == torch.uint8)
+
+            if kv_cache_dtype in ("fp8", "fp8_e4m3"):
+                target_dtype = torch.float8_e4m3fn
+            elif kv_cache_dtype == "fp8_e5m2":
+                target_dtype = torch.float8_e5m2
+            else:
+                raise ValueError("Unsupported FP8 dtype:", kv_cache_dtype)
+
+            k_cache = k_cache.view(target_dtype)
+            v_cache = v_cache.view(target_dtype)
+
+        if (k_cache.dtype == torch.uint8
+                or v_cache.dtype == torch.uint8 and kv_cache_dtype == "auto"):
+            raise ValueError("kv_cache_dtype='auto' unsupported for\
+                FP8 KV Cache prefill kernel")
+
+        # shape constraints
+        Lq, Lk, Lv = q.shape[-1], k.shape[-1], v.shape[-1]
+        assert Lq == Lk and Lk == Lv
+        # round up Lk to a power of 2 - this is required for Triton block size
+        Lk_padded = triton.next_power_of_2(Lk)
+
+        sm_scale = 1.0 / (Lq**0.5)
+        batch, head = b_seq_len.shape[0], q.shape[1]
+        num_queries_per_kv = q.shape[1] // k.shape[1]
+
+        grid = (batch, head, triton.cdiv(max_input_len, BLOCK))  # batch, head,
+
+        # 0 means "disable"
+        if sliding_window is None or sliding_window <= 0:
+            sliding_window = 0
+
+        if alibi_slopes is not None:
+            _fwd_kernel_alibi[grid](
+                q,
+                k,
+                v,
+                k_cache,
+                v_cache,
+                b_loc,
+                sm_scale,
+                k_scale,
+                v_scale,
+                b_start_loc,
+                b_seq_len,
+                b_ctx_len,
+                alibi_slopes,
+                v_cache.shape[3],
+                k_cache.shape[4],
+                o,
+                b_loc.stride(0),
+                b_loc.stride(1),
+                q.stride(0),
+                q.stride(1),
+                q.stride(2),
+                k.stride(0),
+                k.stride(1),
+                k.stride(2),
+                v.stride(0),
+                v.stride(1),
+                v.stride(2),
+                o.stride(0),
+                o.stride(1),
+                o.stride(2),
+                k_cache.stride(0),
+                k_cache.stride(1),
+                k_cache.stride(2),
+                k_cache.stride(3),
+                k_cache.stride(
+                    4
+                ),  #[num_blocks, num_kv_heads, head_size/x, block_size, x]
+                v_cache.stride(0),
+                v_cache.stride(1),
+                v_cache.stride(2),
+                v_cache.stride(
+                    3),  #[num_blocks, num_kv_heads, head_size, block_size]
+                num_queries_per_kv=num_queries_per_kv,
+                BLOCK_M=BLOCK,
+                BLOCK_DMODEL=Lk,
+                BLOCK_DMODEL_PADDED=Lk_padded,
+                BLOCK_N=BLOCK_N,
+                num_warps=NUM_WARPS,
+                num_stages=1,
+            )
+            return
+
+        _fwd_kernel[grid](
+            q,
+            k,
+            v,
+            k_cache,
+            v_cache,
+            b_loc,
+            sm_scale,
+            k_scale,
+            v_scale,
+            b_start_loc,
+            b_seq_len,
+            b_ctx_len,
+            v_cache.shape[3],
+            k_cache.shape[4],
+            o,
+            b_loc.stride(0),
+            b_loc.stride(1),
+            q.stride(0),
+            q.stride(1),
+            q.stride(2),
+            k.stride(0),
+            k.stride(1),
+            k.stride(2),
+            v.stride(0),
+            v.stride(1),
+            v.stride(2),
+            o.stride(0),
+            o.stride(1),
+            o.stride(2),
+            k_cache.stride(0),
+            k_cache.stride(1),
+            k_cache.stride(2),
+            k_cache.stride(3),
+            k_cache.stride(
+                4),  #[num_blocks, num_kv_heads, head_size/x, block_size, x]
+            v_cache.stride(0),
+            v_cache.stride(1),
+            v_cache.stride(2),
+            v_cache.stride(
+                3),  #[num_blocks, num_kv_heads, head_size, block_size]
+            num_queries_per_kv=num_queries_per_kv,
+            BLOCK_M=BLOCK,
+            BLOCK_DMODEL=Lk,
+            BLOCK_DMODEL_PADDED=Lk_padded,
+            BLOCK_N=BLOCK_N,
+            SLIDING_WINDOW=sliding_window,
+            num_warps=NUM_WARPS,
+            num_stages=1,
+        )
+        return

qwen3_6_scripts/qwen3_5/__init__.py

@@ -0,0 +1,3 @@
+from .configuration_qwen3_5 import Qwen3_5Config, Qwen3_5TextConfig, Qwen3_5VisionConfig
+
+__all__ = ["Qwen3_5Config", "Qwen3_5TextConfig", "Qwen3_5VisionConfig"]

qwen3_6_scripts/qwen3_5/configuration_qwen3_5.py

@@ -0,0 +1,188 @@
+# Adapted from transformers 5.2.0 for compatibility with transformers 4.55.3 + torch 2.1.0
+# Stubs layer_type_validation and RopeParameters which do not exist in 4.55.3
+
+from typing import Optional, List
+
+from ...configuration_utils import PretrainedConfig as PreTrainedConfig
+
+# --- Local stubs for APIs not present in transformers 4.55.3 ---
+# Always use these definitions; do NOT import from the older transformers
+# as same-named functions there have incompatible signatures.
+
+def layer_type_validation(layer_types, num_hidden_layers=None, attention=True):
+    allowed = {"full_attention", "linear_attention"}
+    if not all(lt in allowed for lt in layer_types):
+        raise ValueError(f"layer_types entries must be in {allowed}, got {layer_types}")
+    if num_hidden_layers is not None and num_hidden_layers != len(layer_types):
+        raise ValueError(
+            f"num_hidden_layers ({num_hidden_layers}) != len(layer_types) ({len(layer_types)})"
+        )
+
+try:
+    from typing import TypedDict
+    class RopeParameters(TypedDict, total=False):
+        rope_theta: float
+        rope_type: str
+        partial_rotary_factor: float
+        factor: float
+except Exception:
+    RopeParameters = dict
+
+# --- End stubs ---
+
+
+class Qwen3_5TextConfig(PreTrainedConfig):
+    r"""
+    Configuration for the text backbone of Qwen3.5 / Qwen3.6-27B models.
+    model_type is "qwen3_5_text" (used internally by the nested config).
+    """
+
+    model_type = "qwen3_5_text"
+    keys_to_ignore_at_inference = ["past_key_values"]
+
+    def __init__(
+        self,
+        vocab_size=248320,
+        hidden_size=4096,
+        intermediate_size=12288,
+        num_hidden_layers=32,
+        num_attention_heads=16,
+        num_key_value_heads=4,
+        hidden_act="silu",
+        max_position_embeddings=32768,
+        initializer_range=0.02,
+        rms_norm_eps=1e-6,
+        use_cache=True,
+        tie_word_embeddings=False,
+        rope_parameters=None,
+        attention_bias=False,
+        attention_dropout=0.0,
+        head_dim=256,
+        linear_conv_kernel_dim=4,
+        linear_key_head_dim=128,
+        linear_value_head_dim=128,
+        linear_num_key_heads=16,
+        linear_num_value_heads=32,
+        layer_types=None,
+        pad_token_id=None,
+        bos_token_id=None,
+        eos_token_id=None,
+        **kwargs,
+    ):
+        self.pad_token_id = pad_token_id
+        self.bos_token_id = bos_token_id
+        self.eos_token_id = eos_token_id
+        self.tie_word_embeddings = tie_word_embeddings
+        self.vocab_size = vocab_size
+        self.max_position_embeddings = max_position_embeddings
+        self.hidden_size = hidden_size
+        self.intermediate_size = intermediate_size
+        self.num_hidden_layers = num_hidden_layers
+        self.num_attention_heads = num_attention_heads
+        self.num_key_value_heads = num_key_value_heads
+        self.hidden_act = hidden_act
+        self.initializer_range = initializer_range
+        self.rms_norm_eps = rms_norm_eps
+        self.use_cache = use_cache
+        self.attention_bias = attention_bias
+        self.attention_dropout = attention_dropout
+        self.head_dim = head_dim
+        self.rope_parameters = rope_parameters
+        kwargs.setdefault("partial_rotary_factor", 0.25)
+
+        self.layer_types = layer_types
+        if self.layer_types is None:
+            interval_pattern = kwargs.get("full_attention_interval", 4)
+            self.layer_types = [
+                "linear_attention" if bool((i + 1) % interval_pattern) else "full_attention"
+                for i in range(self.num_hidden_layers)
+            ]
+        layer_type_validation(self.layer_types, self.num_hidden_layers)
+
+        self.linear_conv_kernel_dim = linear_conv_kernel_dim
+        self.linear_key_head_dim = linear_key_head_dim
+        self.linear_value_head_dim = linear_value_head_dim
+        self.linear_num_key_heads = linear_num_key_heads
+        self.linear_num_value_heads = linear_num_value_heads
+        super().__init__(**kwargs)
+
+
+class Qwen3_5VisionConfig(PreTrainedConfig):
+    model_type = "qwen3_5_vision"
+
+    def __init__(
+        self,
+        depth=27,
+        hidden_size=1152,
+        hidden_act="gelu_pytorch_tanh",
+        intermediate_size=4304,
+        num_heads=16,
+        in_channels=3,
+        patch_size=16,
+        spatial_merge_size=2,
+        temporal_patch_size=2,
+        out_hidden_size=3584,
+        num_position_embeddings=2304,
+        initializer_range=0.02,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.depth = depth
+        self.hidden_size = hidden_size
+        self.hidden_act = hidden_act
+        self.intermediate_size = intermediate_size
+        self.num_heads = num_heads
+        self.in_channels = in_channels
+        self.patch_size = patch_size
+        self.spatial_merge_size = spatial_merge_size
+        self.temporal_patch_size = temporal_patch_size
+        self.out_hidden_size = out_hidden_size
+        self.num_position_embeddings = num_position_embeddings
+        self.initializer_range = initializer_range
+
+
+class Qwen3_5Config(PreTrainedConfig):
+    r"""
+    Top-level configuration for Qwen3.5 / Qwen3.6-27B.
+    model_type = "qwen3_5" matches the model card / config.json.
+    Wraps Qwen3_5TextConfig (and optionally Qwen3_5VisionConfig for multimodal use).
+    For vLLM text-only inference only text_config is consumed.
+    """
+
+    model_type = "qwen3_5"
+    keys_to_ignore_at_inference = ["past_key_values"]
+
+    def __init__(
+        self,
+        text_config=None,
+        vision_config=None,
+        image_token_id=248056,
+        video_token_id=248057,
+        vision_start_token_id=248053,
+        vision_end_token_id=248054,
+        tie_word_embeddings=False,
+        **kwargs,
+    ):
+        if isinstance(text_config, dict):
+            self.text_config = Qwen3_5TextConfig(**text_config)
+        elif text_config is None:
+            self.text_config = Qwen3_5TextConfig()
+        else:
+            self.text_config = text_config
+
+        if isinstance(vision_config, dict):
+            self.vision_config = Qwen3_5VisionConfig(**vision_config)
+        elif vision_config is None:
+            self.vision_config = Qwen3_5VisionConfig()
+        else:
+            self.vision_config = vision_config
+
+        self.image_token_id = image_token_id
+        self.video_token_id = video_token_id
+        self.vision_start_token_id = vision_start_token_id
+        self.vision_end_token_id = vision_end_token_id
+        self.tie_word_embeddings = tie_word_embeddings
+        super().__init__(**kwargs)
+
+
+__all__ = ["Qwen3_5Config", "Qwen3_5TextConfig", "Qwen3_5VisionConfig"]

qwen3_6_scripts/qwen3_5_moe/__init__.py

@@ -0,0 +1,3 @@
+from .configuration_qwen3_5_moe import Qwen3_5MoeConfig, Qwen3_5MoeTextConfig
+
+__all__ = ["Qwen3_5MoeConfig", "Qwen3_5MoeTextConfig"]

qwen3_6_scripts/qwen3_5_moe/configuration_qwen3_5_moe.py

@@ -0,0 +1,198 @@
+# Adapted from transformers 5.2.0 for compatibility with transformers 4.55.3 + torch 2.1.0
+# Source: transformers/models/qwen3_5_moe/configuration_qwen3_5_moe.py
+# Stubs layer_type_validation and RopeParameters which do not exist in 4.55.3
+# Removes ignore_keys_at_rope_validation / base_model_tp_plan / base_model_pp_plan
+# which are 5.x-only and irrelevant for vLLM inference.
+
+from typing import Optional
+
+from ...configuration_utils import PretrainedConfig as PreTrainedConfig
+
+# --- Local stubs for APIs not present in transformers 4.55.3 ---
+def layer_type_validation(layer_types, num_hidden_layers=None, attention=True):
+    allowed = {"full_attention", "linear_attention"}
+    if not all(lt in allowed for lt in layer_types):
+        raise ValueError(f"layer_types entries must be in {allowed}, got {layer_types}")
+    if num_hidden_layers is not None and num_hidden_layers != len(layer_types):
+        raise ValueError(
+            f"num_hidden_layers ({num_hidden_layers}) != len(layer_types) ({len(layer_types)})"
+        )
+
+try:
+    from typing import TypedDict
+    class RopeParameters(TypedDict, total=False):
+        rope_theta: float
+        rope_type: str
+        partial_rotary_factor: float
+        factor: float
+except Exception:
+    RopeParameters = dict
+
+# --- End stubs ---
+
+
+class Qwen3_5MoeTextConfig(PreTrainedConfig):
+    r"""
+    Configuration for the text backbone of Qwen3.5-MoE / Qwen3.6-35B-A3B models.
+    model_type is "qwen3_5_moe_text" (used internally by the nested config).
+    """
+
+    model_type = "qwen3_5_moe_text"
+    keys_to_ignore_at_inference = ["past_key_values"]
+
+    def __init__(
+        self,
+        vocab_size=248320,
+        hidden_size=2048,
+        num_hidden_layers=40,
+        num_attention_heads=16,
+        num_key_value_heads=2,
+        hidden_act="silu",
+        max_position_embeddings=32768,
+        initializer_range=0.02,
+        rms_norm_eps=1e-6,
+        use_cache=True,
+        tie_word_embeddings=False,
+        rope_parameters=None,
+        attention_bias=False,
+        attention_dropout=0.0,
+        head_dim=256,
+        linear_conv_kernel_dim=4,
+        linear_key_head_dim=128,
+        linear_value_head_dim=128,
+        linear_num_key_heads=16,
+        linear_num_value_heads=32,
+        moe_intermediate_size=512,
+        shared_expert_intermediate_size=512,
+        num_experts_per_tok=8,
+        num_experts=256,
+        output_router_logits=False,
+        router_aux_loss_coef=0.001,
+        layer_types=None,
+        pad_token_id=None,
+        bos_token_id=None,
+        eos_token_id=None,
+        **kwargs,
+    ):
+        self.pad_token_id = pad_token_id
+        self.bos_token_id = bos_token_id
+        self.eos_token_id = eos_token_id
+        self.tie_word_embeddings = tie_word_embeddings
+        self.vocab_size = vocab_size
+        self.max_position_embeddings = max_position_embeddings
+        self.hidden_size = hidden_size
+        self.num_hidden_layers = num_hidden_layers
+        self.num_attention_heads = num_attention_heads
+        self.num_key_value_heads = num_key_value_heads
+        self.hidden_act = hidden_act
+        self.initializer_range = initializer_range
+        self.rms_norm_eps = rms_norm_eps
+        self.use_cache = use_cache
+        self.attention_bias = attention_bias
+        self.attention_dropout = attention_dropout
+        self.head_dim = head_dim
+        self.rope_parameters = rope_parameters
+        kwargs.setdefault("partial_rotary_factor", 0.25)
+
+        self.layer_types = layer_types
+        if self.layer_types is None:
+            interval_pattern = kwargs.get("full_attention_interval", 4)
+            self.layer_types = [
+                "linear_attention" if bool((i + 1) % interval_pattern) else "full_attention"
+                for i in range(self.num_hidden_layers)
+            ]
+        layer_type_validation(self.layer_types, self.num_hidden_layers)
+
+        self.linear_conv_kernel_dim = linear_conv_kernel_dim
+        self.linear_key_head_dim = linear_key_head_dim
+        self.linear_value_head_dim = linear_value_head_dim
+        self.linear_num_key_heads = linear_num_key_heads
+        self.linear_num_value_heads = linear_num_value_heads
+        self.moe_intermediate_size = moe_intermediate_size
+        self.shared_expert_intermediate_size = shared_expert_intermediate_size
+        self.num_experts_per_tok = num_experts_per_tok
+        self.num_experts = num_experts
+        self.output_router_logits = output_router_logits
+        self.router_aux_loss_coef = router_aux_loss_coef
+        super().__init__(**kwargs)
+
+
+class Qwen3_5MoeVisionConfig(PreTrainedConfig):
+    model_type = "qwen3_5_moe"
+
+    def __init__(
+        self,
+        depth=27,
+        hidden_size=1152,
+        hidden_act="gelu_pytorch_tanh",
+        intermediate_size=4304,
+        num_heads=16,
+        in_channels=3,
+        patch_size=16,
+        spatial_merge_size=2,
+        temporal_patch_size=2,
+        out_hidden_size=3584,
+        num_position_embeddings=2304,
+        initializer_range=0.02,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.depth = depth
+        self.hidden_size = hidden_size
+        self.hidden_act = hidden_act
+        self.intermediate_size = intermediate_size
+        self.num_heads = num_heads
+        self.in_channels = in_channels
+        self.patch_size = patch_size
+        self.spatial_merge_size = spatial_merge_size
+        self.temporal_patch_size = temporal_patch_size
+        self.out_hidden_size = out_hidden_size
+        self.num_position_embeddings = num_position_embeddings
+        self.initializer_range = initializer_range
+
+
+class Qwen3_5MoeConfig(PreTrainedConfig):
+    r"""
+    Top-level configuration for Qwen3.5-MoE / Qwen3.6-35B-A3B.
+    model_type = "qwen3_5_moe" matches the model card / config.json.
+    Wraps Qwen3_5MoeTextConfig (and optionally Qwen3_5MoeVisionConfig).
+    For vLLM text-only inference only text_config is consumed.
+    """
+
+    model_type = "qwen3_5_moe"
+    keys_to_ignore_at_inference = ["past_key_values"]
+
+    def __init__(
+        self,
+        text_config=None,
+        vision_config=None,
+        image_token_id=248056,
+        video_token_id=248057,
+        vision_start_token_id=248053,
+        vision_end_token_id=248054,
+        tie_word_embeddings=False,
+        **kwargs,
+    ):
+        if isinstance(text_config, dict):
+            self.text_config = Qwen3_5MoeTextConfig(**text_config)
+        elif text_config is None:
+            self.text_config = Qwen3_5MoeTextConfig()
+        else:
+            self.text_config = text_config
+
+        if isinstance(vision_config, dict):
+            self.vision_config = Qwen3_5MoeVisionConfig(**vision_config)
+        elif vision_config is None:
+            self.vision_config = Qwen3_5MoeVisionConfig()
+        else:
+            self.vision_config = vision_config
+
+        self.image_token_id = image_token_id
+        self.video_token_id = video_token_id
+        self.vision_start_token_id = vision_start_token_id
+        self.vision_end_token_id = vision_end_token_id
+        self.tie_word_embeddings = tie_word_embeddings
+        super().__init__(**kwargs)
+
+
+__all__ = ["Qwen3_5MoeConfig", "Qwen3_5MoeTextConfig"]

qwen3_6_scripts/qwen3coder_tool_parser.py

@@ -0,0 +1,531 @@
+import ast
+import json
+import uuid
+from typing import Any, Dict, List, Optional, Sequence, Union
+
+import regex as re
+
+from vllm.entrypoints.openai.protocol import (ChatCompletionRequest,
+                                              ChatCompletionToolsParam,
+                                              DeltaFunctionCall, DeltaMessage,
+                                              DeltaToolCall,
+                                              ExtractedToolCallInformation,
+                                              FunctionCall, ToolCall)
+from vllm.entrypoints.openai.tool_parsers.abstract_tool_parser import (
+    ToolParser, ToolParserManager)
+from vllm.logger import init_logger
+from vllm.transformers_utils.tokenizer import AnyTokenizer
+
+logger = init_logger(__name__)
+
+
+@ToolParserManager.register_module("qwen3_coder")
+class Qwen3CoderToolParser(ToolParser):
+    """
+    Tool parser for Qwen3 models using XML-style tool call format:
+      <tool_call><function=name><parameter=key>
+      value
+      </parameter></function></tool_call>
+
+    Port of vllm-original qwen3coder_tool_parser.py to vllm 0.6.3 API.
+    """
+
+    def __init__(self, tokenizer: AnyTokenizer):
+        super().__init__(tokenizer)
+
+        self.current_tool_name_sent: bool = False
+        self.prev_tool_call_arr: List[Dict] = []
+        # Base class uses int; we override with string IDs
+        self.current_tool_id: Optional[str] = None  # type: ignore[assignment]
+        self.streamed_args_for_tool: List[str] = []
+
+        self.tool_call_start_token: str = "<tool_call>"
+        self.tool_call_end_token: str = "</tool_call>"
+        self.tool_call_prefix: str = "<function="
+        self.function_end_token: str = "</function>"
+        self.parameter_prefix: str = "<parameter="
+        self.parameter_end_token: str = "</parameter>"
+        self.is_tool_call_started: bool = False
+
+        self._reset_streaming_state()
+
+        self.tool_call_complete_regex = re.compile(
+            r"<tool_call>(.*?)</tool_call>", re.DOTALL)
+        self.tool_call_regex = re.compile(
+            r"<tool_call>(.*?)</tool_call>|<tool_call>(.*?)$", re.DOTALL)
+        self.tool_call_function_regex = re.compile(
+            r"<function=(.*?)</function>|<function=(.*)$", re.DOTALL)
+        self.tool_call_parameter_regex = re.compile(
+            r"<parameter=(.*?)(?:</parameter>|(?=<parameter=)|(?=</function>)|$)",
+            re.DOTALL)
+
+        if not self.model_tokenizer:
+            raise ValueError(
+                "The model tokenizer must be passed to the ToolParser "
+                "constructor during construction.")
+
+        self.tool_call_start_token_id = self.vocab.get(
+            self.tool_call_start_token)
+        self.tool_call_end_token_id = self.vocab.get(self.tool_call_end_token)
+
+        if (self.tool_call_start_token_id is None
+                or self.tool_call_end_token_id is None):
+            raise RuntimeError(
+                "Qwen3 XML Tool parser could not locate tool call start/end "
+                "tokens in the tokenizer!")
+
+        logger.debug("vLLM Successfully imported tool parser %s !",
+                     self.__class__.__name__)
+
+    def adjust_request(
+            self, request: "ChatCompletionRequest") -> "ChatCompletionRequest":
+        """Disable thinking when tools are active with auto choice.
+
+        On BI-V100 hardware, the model's <think>...</think> phase can consume
+        the entire max_tokens budget, leaving no room for the <tool_call> XML.
+        Competition reference (sub168) completes d03_tool_call in 2.12s with
+        tools=1; our sub509 took 49s with tools=0 because thinking ate the
+        budget.  Disabling thinking for tool-call requests ensures the model
+        emits tool XML within the token budget.
+        """
+        if (request.tools and request.tool_choice in ("auto", None)
+                and not isinstance(request.tool_choice,
+                                   type(None).__class__)):
+            # Only override if thinking was not explicitly requested
+            ctk = request.chat_template_kwargs or {}
+            if "enable_thinking" not in ctk:
+                ctk = dict(ctk)  # shallow copy
+                ctk["enable_thinking"] = False
+                request.chat_template_kwargs = ctk
+        return request
+
+
+    def _generate_tool_call_id(self) -> str:
+        return f"call_{uuid.uuid4().hex[:24]}"
+
+    def _reset_streaming_state(self) -> None:
+        self.current_tool_index = 0
+        self.is_tool_call_started = False
+        self.header_sent = False
+        self.current_tool_id = None
+        self.current_function_name: Optional[str] = None
+        self.current_param_name: Optional[str] = None
+        self.current_param_value: str = ""
+        self.param_count = 0
+        self.in_param = False
+        self.in_function = False
+        self.accumulated_text: str = ""
+        self.json_started = False
+        self.json_closed = False
+        self.accumulated_params: Dict[str, Any] = {}
+        self.streaming_request: Optional[ChatCompletionRequest] = None
+
+    def _get_arguments_config(
+            self, func_name: str,
+            tools: Optional[List[ChatCompletionToolsParam]]) -> Dict:
+        if tools is None:
+            return {}
+        for config in tools:
+            if not hasattr(config, "type") or not (
+                    hasattr(config, "function")
+                    and hasattr(config.function, "name")):
+                continue
+            if config.type == "function" and config.function.name == func_name:
+                if not hasattr(config.function, "parameters"):
+                    return {}
+                params = config.function.parameters
+                if isinstance(params, dict) and "properties" in params:
+                    return params["properties"]
+                elif isinstance(params, dict):
+                    return params
+                else:
+                    return {}
+        logger.debug("Tool '%s' is not defined in the tools list.", func_name)
+        return {}
+
+    def _convert_param_value(self, param_value: str, param_name: str,
+                             param_config: Dict, func_name: str) -> Any:
+        if param_value.lower() == "null":
+            return None
+
+        if param_name not in param_config:
+            if param_config != {}:
+                logger.debug(
+                    "Parsed parameter '%s' is not defined in tool '%s', "
+                    "returning string value.", param_name, func_name)
+            return param_value
+
+        if (isinstance(param_config[param_name], dict)
+                and "type" in param_config[param_name]):
+            param_type = str(
+                param_config[param_name]["type"]).strip().lower()
+        else:
+            param_type = "string"
+
+        if param_type in ["string", "str", "text", "varchar", "char", "enum"]:
+            return param_value
+        elif (param_type.startswith("int") or param_type.startswith("uint")
+              or param_type.startswith("long")
+              or param_type.startswith("short")
+              or param_type.startswith("unsigned")):
+            try:
+                return int(param_value)
+            except (ValueError, TypeError):
+                return param_value
+        elif param_type.startswith("num") or param_type.startswith("float"):
+            try:
+                v = float(param_value)
+                return int(v) if v - int(v) == 0 else v
+            except (ValueError, TypeError):
+                return param_value
+        elif param_type in ["boolean", "bool", "binary"]:
+            lower = param_value.lower()
+            if lower not in ["true", "false"]:
+                logger.debug(
+                    "Parameter '%s' value '%s' is not boolean in tool '%s'.",
+                    param_name, param_value, func_name)
+            return lower == "true"
+        else:
+            if (param_type in ["object", "array", "arr"]
+                    or param_type.startswith("dict")
+                    or param_type.startswith("list")):
+                try:
+                    return json.loads(param_value)
+                except (json.JSONDecodeError, TypeError, ValueError):
+                    pass
+            try:
+                return ast.literal_eval(param_value)
+            except (ValueError, SyntaxError, TypeError):
+                pass
+            return param_value
+
+    def _parse_xml_function_call(
+            self, function_call_str: str,
+            tools: Optional[List[ChatCompletionToolsParam]]) -> ToolCall:
+        end_index = function_call_str.index(">")
+        function_name = function_call_str[:end_index]
+        param_config = self._get_arguments_config(function_name, tools)
+        parameters = function_call_str[end_index + 1:]
+        param_dict: Dict[str, Any] = {}
+        for match_text in self.tool_call_parameter_regex.findall(parameters):
+            idx = match_text.index(">")
+            param_name = match_text[:idx]
+            param_value = str(match_text[idx + 1:])
+            if param_value.startswith("\n"):
+                param_value = param_value[1:]
+            if param_value.endswith("\n"):
+                param_value = param_value[:-1]
+            param_dict[param_name] = self._convert_param_value(
+                param_value, param_name, param_config, function_name)
+        return ToolCall(
+            type="function",
+            function=FunctionCall(
+                name=function_name,
+                arguments=json.dumps(param_dict, ensure_ascii=False)))
+
+    def _get_function_calls(self, model_output: str) -> List[str]:
+        matched_ranges = self.tool_call_regex.findall(model_output)
+        raw_tool_calls = [
+            match[0] if match[0] else match[1] for match in matched_ranges
+        ]
+        if not raw_tool_calls:
+            raw_tool_calls = [model_output]
+        raw_function_calls: List[tuple] = []
+        for tool_call in raw_tool_calls:
+            raw_function_calls.extend(
+                self.tool_call_function_regex.findall(tool_call))
+        return [match[0] if match[0] else match[1]
+                for match in raw_function_calls]
+
+    def extract_tool_calls(
+            self, model_output: str,
+            request: ChatCompletionRequest) -> ExtractedToolCallInformation:
+        if self.tool_call_prefix not in model_output:
+            return ExtractedToolCallInformation(tools_called=False,
+                                                tool_calls=[],
+                                                content=model_output)
+        try:
+            function_calls = self._get_function_calls(model_output)
+            if not function_calls:
+                return ExtractedToolCallInformation(tools_called=False,
+                                                    tool_calls=[],
+                                                    content=model_output)
+
+            tool_calls = [
+                self._parse_xml_function_call(fc, request.tools)
+                for fc in function_calls
+            ]
+
+            self.prev_tool_call_arr.clear()
+            for tc in tool_calls:
+                self.prev_tool_call_arr.append({
+                    "name": tc.function.name,
+                    "arguments": tc.function.arguments,
+                })
+
+            content_index = model_output.find(self.tool_call_start_token)
+            idx = model_output.find(self.tool_call_prefix)
+            content_index = content_index if content_index >= 0 else idx
+            content = model_output[:content_index]
+
+            return ExtractedToolCallInformation(
+                tools_called=bool(tool_calls),
+                tool_calls=tool_calls,
+                content=content if content else None,
+            )
+        except Exception:
+            logger.exception("Error extracting tool call from response.")
+            return ExtractedToolCallInformation(tools_called=False,
+                                                tool_calls=[],
+                                                content=model_output)
+
+    def extract_tool_calls_streaming(
+        self,
+        previous_text: str,
+        current_text: str,
+        delta_text: str,
+        previous_token_ids: Sequence[int],
+        current_token_ids: Sequence[int],
+        delta_token_ids: Sequence[int],
+        request: ChatCompletionRequest,
+    ) -> Union[DeltaMessage, None]:
+        if not previous_text:
+            self._reset_streaming_state()
+            self.streaming_request = request
+
+        if not delta_text:
+            if delta_token_ids and self.tool_call_end_token_id not in delta_token_ids:
+                complete_calls = len(
+                    self.tool_call_complete_regex.findall(current_text))
+                if complete_calls > 0 and self.prev_tool_call_arr:
+                    open_calls = (
+                        current_text.count(self.tool_call_start_token) -
+                        current_text.count(self.tool_call_end_token))
+                    if open_calls == 0:
+                        return DeltaMessage(content="")
+                elif not self.is_tool_call_started and current_text:
+                    return DeltaMessage(content="")
+            return None
+
+        self.accumulated_text = current_text
+
+        if self.json_closed and not self.in_function:
+            tool_ends = current_text.count(self.tool_call_end_token)
+            if tool_ends > self.current_tool_index:
+                self.current_tool_index += 1
+                self.header_sent = False
+                self.param_count = 0
+                self.json_started = False
+                self.json_closed = False
+                self.accumulated_params = {}
+                tool_starts = current_text.count(self.tool_call_start_token)
+                if self.current_tool_index >= tool_starts:
+                    self.is_tool_call_started = False
+                return None
+
+        if not self.is_tool_call_started:
+            if (self.tool_call_start_token_id in delta_token_ids
+                    or self.tool_call_start_token in delta_text):
+                self.is_tool_call_started = True
+                if self.tool_call_start_token in delta_text:
+                    content_before = delta_text[:delta_text.index(
+                        self.tool_call_start_token)]
+                    if content_before:
+                        return DeltaMessage(content=content_before)
+                return None
+            else:
+                if (current_text.rstrip().endswith(self.tool_call_end_token)
+                        and delta_text.strip() == ""):
+                    return None
+                return DeltaMessage(content=delta_text)
+
+        tool_starts_count = current_text.count(self.tool_call_start_token)
+        if self.current_tool_index >= tool_starts_count:
+            return None
+
+        # Locate the current tool call's text slice
+        tool_start_positions: List[int] = []
+        search = 0
+        while True:
+            search = current_text.find(self.tool_call_start_token, search)
+            if search == -1:
+                break
+            tool_start_positions.append(search)
+            search += len(self.tool_call_start_token)
+
+        if self.current_tool_index >= len(tool_start_positions):
+            return None
+
+        tool_start_idx = tool_start_positions[self.current_tool_index]
+        tool_end_idx = current_text.find(self.tool_call_end_token,
+                                         tool_start_idx)
+        if tool_end_idx == -1:
+            tool_text = current_text[tool_start_idx:]
+        else:
+            tool_text = current_text[tool_start_idx:tool_end_idx +
+                                     len(self.tool_call_end_token)]
+
+        if not self.header_sent:
+            if self.tool_call_prefix in tool_text:
+                func_start = (tool_text.find(self.tool_call_prefix) +
+                              len(self.tool_call_prefix))
+                func_end = tool_text.find(">", func_start)
+                if func_end != -1:
+                    self.current_function_name = tool_text[func_start:func_end]
+                    self.current_tool_id = self._generate_tool_call_id()
+                    self.header_sent = True
+                    self.in_function = True
+                    self.prev_tool_call_arr.append({
+                        "name": self.current_function_name,
+                        "arguments": "{}",
+                    })
+                    self.streamed_args_for_tool.append("")
+                    return DeltaMessage(tool_calls=[
+                        DeltaToolCall(
+                            index=self.current_tool_index,
+                            id=self.current_tool_id,
+                            function=DeltaFunctionCall(
+                                name=self.current_function_name,
+                                arguments=""),
+                            type="function",
+                        )
+                    ])
+            return None
+
+        if self.in_function:
+            if not self.json_started:
+                self.json_started = True
+                self.streamed_args_for_tool[self.current_tool_index] += "{"
+                return DeltaMessage(tool_calls=[
+                    DeltaToolCall(
+                        index=self.current_tool_index,
+                        function=DeltaFunctionCall(arguments="{"),
+                    )
+                ])
+
+            # Collect all complete parameters in one pass (speculative-decode safe)
+            param_starts: List[int] = []
+            search = 0
+            while True:
+                search = tool_text.find(self.parameter_prefix, search)
+                if search == -1:
+                    break
+                param_starts.append(search)
+                search += len(self.parameter_prefix)
+
+            json_fragments: List[str] = []
+            while not self.in_param and self.param_count < len(param_starts):
+                param_idx = param_starts[self.param_count]
+                param_start = param_idx + len(self.parameter_prefix)
+                remaining = tool_text[param_start:]
+
+                if ">" not in remaining:
+                    break
+
+                name_end = remaining.find(">")
+                current_param_name = remaining[:name_end]
+                value_start = param_start + name_end + 1
+                value_text = tool_text[value_start:]
+                if value_text.startswith("\n"):
+                    value_text = value_text[1:]
+
+                param_end_idx = value_text.find(self.parameter_end_token)
+                if param_end_idx == -1:
+                    next_param = value_text.find(self.parameter_prefix)
+                    func_end = value_text.find(self.function_end_token)
+                    if next_param != -1 and (func_end == -1
+                                              or next_param < func_end):
+                        param_end_idx = next_param
+                    elif func_end != -1:
+                        param_end_idx = func_end
+                    else:
+                        tool_end_in_value = value_text.find(
+                            self.tool_call_end_token)
+                        if tool_end_in_value != -1:
+                            param_end_idx = tool_end_in_value
+                        else:
+                            break
+
+                if param_end_idx == -1:
+                    break
+
+                param_value = value_text[:param_end_idx]
+                if param_value.endswith("\n"):
+                    param_value = param_value[:-1]
+
+                self.accumulated_params[current_param_name] = param_value
+                param_config = self._get_arguments_config(
+                    self.current_function_name or "",
+                    self.streaming_request.tools
+                    if self.streaming_request else None)
+                converted = self._convert_param_value(
+                    param_value, current_param_name, param_config,
+                    self.current_function_name or "")
+                serialized = json.dumps(converted, ensure_ascii=False)
+
+                sep = "" if self.param_count == 0 else ", "
+                json_fragments.append(
+                    f'{sep}"{current_param_name}": {serialized}')
+                self.param_count += 1
+
+            if json_fragments:
+                combined = "".join(json_fragments)
+                if self.current_tool_index < len(self.streamed_args_for_tool):
+                    self.streamed_args_for_tool[
+                        self.current_tool_index] += combined
+                else:
+                    logger.warning(
+                        "streamed_args_for_tool out of sync: index=%d len=%d",
+                        self.current_tool_index,
+                        len(self.streamed_args_for_tool))
+                return DeltaMessage(tool_calls=[
+                    DeltaToolCall(
+                        index=self.current_tool_index,
+                        function=DeltaFunctionCall(arguments=combined),
+                    )
+                ])
+
+            # Emit closing brace when </function> is seen (after params are done)
+            if not self.json_closed and self.function_end_token in tool_text:
+                self.json_closed = True
+                func_start = (tool_text.find(self.tool_call_prefix) +
+                              len(self.tool_call_prefix))
+                func_content_end = tool_text.find(self.function_end_token,
+                                                   func_start)
+                if func_content_end != -1:
+                    try:
+                        parsed_tool = self._parse_xml_function_call(
+                            tool_text[func_start:func_content_end],
+                            self.streaming_request.tools
+                            if self.streaming_request else None)
+                        if self.current_tool_index < len(
+                                self.prev_tool_call_arr):
+                            self.prev_tool_call_arr[
+                                self.current_tool_index]["arguments"] = (
+                                    parsed_tool.function.arguments)
+                    except Exception:
+                        logger.debug("Failed to parse tool call during "
+                                     "streaming: %s",
+                                     tool_text,
+                                     exc_info=True)
+
+                if self.current_tool_index < len(self.streamed_args_for_tool):
+                    self.streamed_args_for_tool[
+                        self.current_tool_index] += "}"
+                else:
+                    logger.warning(
+                        "streamed_args_for_tool out of sync: index=%d len=%d",
+                        self.current_tool_index,
+                        len(self.streamed_args_for_tool))
+
+                result = DeltaMessage(tool_calls=[
+                    DeltaToolCall(
+                        index=self.current_tool_index,
+                        function=DeltaFunctionCall(arguments="}"),
+                    )
+                ])
+                self.in_function = False
+                self.accumulated_params = {}
+                return result
+
+        return None

qwen3_6_scripts/reasoning/__init__.py

@@ -0,0 +1,16 @@
+"""
+Reasoning parser module for vLLM 0.6.3 (BI-V100 / Qwen3.6-27B adaptation).
+
+Usage: --reasoning-parser qwen3
+"""
+
+from vllm.reasoning.abs_reasoning_parsers import ReasoningParser, ReasoningParserManager
+
+__all__ = ["ReasoningParser", "ReasoningParserManager"]
+
+# Lazy-register Qwen3 parser; imported on first get_reasoning_parser("qwen3").
+ReasoningParserManager.register_lazy(
+    "qwen3",
+    "vllm.reasoning.qwen3_reasoning_parser",
+    "Qwen3ReasoningParser",
+)

qwen3_6_scripts/reasoning/abs_reasoning_parsers.py

@@ -0,0 +1,243 @@
+"""
+Abstract reasoning parser base classes for vLLM 0.6.3.
+Adapted from vllm-original/vllm/reasoning/abs_reasoning_parsers.py:
+  - Removed vllm.entrypoints.mcp, vllm.utils.collection_utils, import_utils
+  - DeltaMessage from vllm 0.6.3 protocol path
+  - TokenizerLike -> AnyTokenizer
+  - ReasoningParserManager: simplified eager + lazy registration
+"""
+
+import importlib
+from abc import abstractmethod
+from collections.abc import Iterable, Sequence
+from functools import cached_property
+from typing import Any, Optional, TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from vllm.entrypoints.openai.protocol import DeltaMessage
+    from vllm.transformers_utils.tokenizer import AnyTokenizer
+else:
+    DeltaMessage = Any
+    AnyTokenizer = Any
+
+
+class ReasoningParser:
+    """Abstract base for all reasoning parsers."""
+
+    def __init__(self, tokenizer: "AnyTokenizer", *args, **kwargs):
+        self.model_tokenizer = tokenizer
+
+    @cached_property
+    def vocab(self) -> dict:
+        return self.model_tokenizer.get_vocab()
+
+    @abstractmethod
+    def is_reasoning_end(self, input_ids: Sequence[int]) -> bool:
+        """Return True once the reasoning block has closed in input_ids."""
+
+    def is_reasoning_end_streaming(
+        self, input_ids: Sequence[int], delta_ids: Iterable[int]
+    ) -> bool:
+        return self.is_reasoning_end(input_ids)
+
+    @abstractmethod
+    def extract_content_ids(self, input_ids: list) -> list:
+        """Return token ids that belong to the content (post-reasoning) part."""
+
+    def count_reasoning_tokens(self, token_ids: Sequence[int]) -> int:
+        return 0
+
+    @abstractmethod
+    def extract_reasoning(
+        self, model_output: str, request: Any
+    ) -> "tuple[Optional[str], Optional[str]]":
+        """
+        Split a complete model output into (reasoning_text, content_text).
+        Either part may be None.
+        """
+
+    @abstractmethod
+    def extract_reasoning_streaming(
+        self,
+        previous_text: str,
+        current_text: str,
+        delta_text: str,
+        previous_token_ids: Sequence[int],
+        current_token_ids: Sequence[int],
+        delta_token_ids: Sequence[int],
+    ) -> Optional["DeltaMessage"]:
+        """
+        Extract reasoning from a streaming delta.
+        Returns a DeltaMessage with reasoning_content and/or content set,
+        or None if this delta should be suppressed (control token).
+        """
+
+
+class BaseThinkingReasoningParser(ReasoningParser):
+    """
+    Base for parsers that use <start_token>...</end_token> delimiters.
+    Subclasses define start_token / end_token properties.
+    """
+
+    @property
+    @abstractmethod
+    def start_token(self) -> str:
+        raise NotImplementedError
+
+    @property
+    @abstractmethod
+    def end_token(self) -> str:
+        raise NotImplementedError
+
+    def __init__(self, tokenizer: "AnyTokenizer", *args, **kwargs):
+        super().__init__(tokenizer, *args, **kwargs)
+
+        if not self.model_tokenizer:
+            raise ValueError("Tokenizer must be passed to ReasoningParser.")
+        if not self.start_token or not self.end_token:
+            raise ValueError("start_token and end_token must be defined.")
+
+        self.start_token_id: Optional[int] = self.vocab.get(self.start_token)
+        self.end_token_id: Optional[int] = self.vocab.get(self.end_token)
+        if self.start_token_id is None or self.end_token_id is None:
+            raise RuntimeError(
+                f"{self.__class__.__name__}: could not find think tokens "
+                f"'{self.start_token}'/'{self.end_token}' in tokenizer vocab."
+            )
+
+    def is_reasoning_end(self, input_ids: Sequence[int]) -> bool:
+        for token_id in reversed(input_ids):
+            if token_id == self.start_token_id:
+                return False
+            if token_id == self.end_token_id:
+                return True
+        return False
+
+    def is_reasoning_end_streaming(
+        self, input_ids: Sequence[int], delta_ids: Iterable[int]
+    ) -> bool:
+        return self.end_token_id in delta_ids
+
+    def extract_content_ids(self, input_ids: list) -> list:
+        if self.end_token_id not in input_ids[:-1]:
+            return []
+        return input_ids[input_ids.index(self.end_token_id) + 1:]
+
+    def count_reasoning_tokens(self, token_ids: Sequence[int]) -> int:
+        count = 0
+        depth = 0
+        for tid in token_ids:
+            if tid == self.start_token_id:
+                depth += 1
+            elif tid == self.end_token_id:
+                if depth > 0:
+                    depth -= 1
+            elif depth > 0:
+                count += 1
+        return count
+
+    def extract_reasoning(
+        self, model_output: str, request: Any
+    ) -> "tuple[Optional[str], Optional[str]]":
+        # Strip <think> if the model generated it (old-style template).
+        parts = model_output.partition(self.start_token)
+        model_output = parts[2] if parts[1] else parts[0]
+
+        if self.end_token not in model_output:
+            return model_output, None
+        reasoning, _, content = model_output.partition(self.end_token)
+        return reasoning, content or None
+
+    def extract_reasoning_streaming(
+        self,
+        previous_text: str,
+        current_text: str,
+        delta_text: str,
+        previous_token_ids: Sequence[int],
+        current_token_ids: Sequence[int],
+        delta_token_ids: Sequence[int],
+    ) -> Optional["DeltaMessage"]:
+        from vllm.entrypoints.openai.protocol import DeltaMessage as _DeltaMessage
+
+        # Suppress lone control tokens.
+        if len(delta_token_ids) == 1 and delta_token_ids[0] in (
+            self.start_token_id, self.end_token_id
+        ):
+            return None
+
+        start_in_prev = self.start_token_id in previous_token_ids
+        start_in_delta = self.start_token_id in delta_token_ids
+        end_in_prev = self.end_token_id in previous_token_ids
+        end_in_delta = self.end_token_id in delta_token_ids
+
+        if start_in_prev:
+            if end_in_delta:
+                end_idx = delta_text.find(self.end_token)
+                reasoning = delta_text[:end_idx] if end_idx >= 0 else ""
+                content = delta_text[end_idx + len(self.end_token):] if end_idx >= 0 else None
+                return _DeltaMessage(
+                    reasoning_content=reasoning or None,
+                    content=content or None,
+                )
+            elif end_in_prev:
+                return _DeltaMessage(content=delta_text)
+            else:
+                return _DeltaMessage(reasoning_content=delta_text)
+
+        elif start_in_delta:
+            if end_in_delta:
+                start_idx = delta_text.find(self.start_token)
+                end_idx = delta_text.find(self.end_token)
+                reasoning = delta_text[start_idx + len(self.start_token):end_idx]
+                content = delta_text[end_idx + len(self.end_token):]
+                return _DeltaMessage(
+                    reasoning_content=reasoning or None,
+                    content=content or None,
+                )
+            else:
+                return _DeltaMessage(reasoning_content=delta_text)
+
+        else:
+            return _DeltaMessage(content=delta_text)
+
+
+class ReasoningParserManager:
+    """
+    Registry for ReasoningParser implementations.
+    Supports eager and lazy registration.
+    """
+
+    _parsers: dict = {}           # name -> class (eager)
+    _lazy: dict = {}              # name -> (module_path, class_name)
+
+    @classmethod
+    def register_module(cls, name: str, parser_cls: type) -> None:
+        """Eagerly register a ReasoningParser class."""
+        if not issubclass(parser_cls, ReasoningParser):
+            raise TypeError(f"{parser_cls} is not a ReasoningParser subclass.")
+        cls._parsers[name] = parser_cls
+
+    @classmethod
+    def register_lazy(cls, name: str, module_path: str, class_name: str) -> None:
+        """Register a parser for deferred import."""
+        cls._lazy[name] = (module_path, class_name)
+
+    @classmethod
+    def get_reasoning_parser(cls, name: str) -> type:
+        if name in cls._parsers:
+            return cls._parsers[name]
+        if name in cls._lazy:
+            module_path, class_name = cls._lazy[name]
+            mod = importlib.import_module(module_path)
+            parser_cls = getattr(mod, class_name)
+            cls._parsers[name] = parser_cls
+            return parser_cls
+        registered = sorted(set(cls._parsers) | set(cls._lazy))
+        raise KeyError(
+            f"Reasoning parser '{name}' not found. "
+            f"Available: {registered}"
+        )
+
+    @classmethod
+    def list_registered(cls) -> list:
+        return sorted(set(cls._parsers) | set(cls._lazy))

qwen3_6_scripts/reasoning/qwen3_reasoning_parser.py

@@ -0,0 +1,110 @@
+"""
+Reasoning parser for Qwen3 / Qwen3.5 / Qwen3.6 model family.
+Adapted from vllm-original/vllm/reasoning/qwen3_reasoning_parser.py.
+
+The model uses <think>...</think> to wrap chain-of-thought output.
+For Qwen3.5+ the chat template injects <think> into the prompt, so only
+</think> appears in the generated tokens; older templates generate <think>
+themselves.  Both styles are handled.
+"""
+
+from typing import Optional, Sequence, Any
+
+from vllm.reasoning.abs_reasoning_parsers import (
+    BaseThinkingReasoningParser,
+    ReasoningParserManager,
+)
+
+
+class Qwen3ReasoningParser(BaseThinkingReasoningParser):
+
+    def __init__(self, tokenizer: Any, *args, **kwargs):
+        super().__init__(tokenizer, *args, **kwargs)
+        chat_kwargs = kwargs.get("chat_template_kwargs", {}) or {}
+        self.thinking_enabled = chat_kwargs.get("enable_thinking", True)
+
+    @property
+    def start_token(self) -> str:
+        return "<think>"
+
+    @property
+    def end_token(self) -> str:
+        return "</think>"
+
+    def extract_reasoning(
+        self, model_output: str, request: Any
+    ) -> "tuple[Optional[str], Optional[str]]":
+        # Strip <think> if the model generated it (old template / edge case).
+        parts = model_output.partition(self.start_token)
+        model_output = parts[2] if parts[1] else parts[0]
+
+        if self.end_token not in model_output:
+            if not self.thinking_enabled:
+                return None, model_output
+            # Thinking enabled but output truncated before </think>.
+            # All output is reasoning; content is None.
+            return model_output, None
+
+        reasoning, _, content = model_output.partition(self.end_token)
+        content = content.strip() if content else ""
+        return reasoning or None, content if content else None
+
+    def count_reasoning_tokens(self, token_ids: Sequence[int]) -> int:
+        token_ids = list(token_ids)
+        if self.start_token_id in token_ids:
+            # Old-style template: model generates <think> itself.
+            # Use depth-counting from the base class.
+            return super().count_reasoning_tokens(token_ids)
+        elif self.end_token_id in token_ids:
+            # New-style template (Qwen3.5+): <think> is injected into the
+            # prompt, so output starts already inside the thinking block.
+            # Every token before </think> is a reasoning token.
+            return token_ids.index(self.end_token_id)
+        else:
+            # No </think> in output: either truncated (all reasoning)
+            # or thinking disabled (none).
+            return len(token_ids) if self.thinking_enabled else 0
+
+    def extract_reasoning_streaming(
+        self,
+        previous_text: str,
+        current_text: str,
+        delta_text: str,
+        previous_token_ids: Sequence[int],
+        current_token_ids: Sequence[int],
+        delta_token_ids: Sequence[int],
+    ):
+        from vllm.entrypoints.openai.protocol import DeltaMessage
+
+        if not self.thinking_enabled:
+            return DeltaMessage(content=delta_text) if delta_text else None
+
+        # Strip <think> from delta if the model generates it itself.
+        if self.start_token_id in delta_token_ids:
+            start_idx = delta_text.find(self.start_token)
+            if start_idx >= 0:
+                delta_text = delta_text[start_idx + len(self.start_token):]
+
+        if self.end_token_id in delta_token_ids:
+            end_idx = delta_text.find(self.end_token)
+            if end_idx >= 0:
+                reasoning = delta_text[:end_idx]
+                content = delta_text[end_idx + len(self.end_token):]
+                if not reasoning and not content:
+                    return None
+                return DeltaMessage(
+                    reasoning_content=reasoning or None,
+                    content=content or None,
+                )
+            return None
+
+        if not delta_text:
+            return None
+        elif self.end_token_id in previous_token_ids:
+            return DeltaMessage(content=delta_text)
+        else:
+            return DeltaMessage(reasoning_content=delta_text)
+
+
+# Register immediately when this module is imported.
+ReasoningParserManager.register_module("qwen3", Qwen3ReasoningParser)

qwen3_6_scripts/registry.py

@@ -0,0 +1,456 @@
+import importlib
+import pickle
+import subprocess
+import sys
+import tempfile
+from abc import ABC, abstractmethod
+from dataclasses import dataclass, field
+from functools import lru_cache
+from typing import Callable, Dict, List, Optional, Tuple, Type, TypeVar, Union
+
+import cloudpickle
+import torch.nn as nn
+
+from vllm.logger import init_logger
+from vllm.utils import is_hip
+
+from .interfaces import (has_inner_state, is_attention_free,
+                         supports_multimodal, supports_pp)
+from .interfaces_base import is_embedding_model, is_text_generation_model
+
+logger = init_logger(__name__)
+
+# yapf: disable
+_TEXT_GENERATION_MODELS = {
+    # [Decoder-only]
+    "AquilaModel": ("llama", "LlamaForCausalLM"),
+    "AquilaForCausalLM": ("llama", "LlamaForCausalLM"),  # AquilaChat2
+    "ArcticForCausalLM": ("arctic", "ArcticForCausalLM"),
+    "BaiChuanForCausalLM": ("baichuan", "BaiChuanForCausalLM"),  # baichuan-7b
+    "BaichuanForCausalLM": ("baichuan", "BaichuanForCausalLM"),  # baichuan-13b
+    "BloomForCausalLM": ("bloom", "BloomForCausalLM"),
+    # ChatGLMModel supports multimodal
+    "CohereForCausalLM": ("commandr", "CohereForCausalLM"),
+    "DbrxForCausalLM": ("dbrx", "DbrxForCausalLM"),
+    "DeciLMForCausalLM": ("decilm", "DeciLMForCausalLM"),
+    "DeepseekForCausalLM": ("deepseek", "DeepseekForCausalLM"),
+    "DeepseekV2ForCausalLM": ("deepseek_v2", "DeepseekV2ForCausalLM"),
+    "ExaoneForCausalLM": ("exaone", "ExaoneForCausalLM"),
+    "FalconForCausalLM": ("falcon", "FalconForCausalLM"),
+    "GemmaForCausalLM": ("gemma", "GemmaForCausalLM"),
+    "Gemma2ForCausalLM": ("gemma2", "Gemma2ForCausalLM"),
+    "Glm4ForCausalLM": ("glm4", "Glm4ForCausalLM"),
+    "GPT2LMHeadModel": ("gpt2", "GPT2LMHeadModel"),
+    "GPTBigCodeForCausalLM": ("gpt_bigcode", "GPTBigCodeForCausalLM"),
+    "GPTJForCausalLM": ("gpt_j", "GPTJForCausalLM"),
+    "GPTNeoXForCausalLM": ("gpt_neox", "GPTNeoXForCausalLM"),
+    "GraniteForCausalLM": ("granite", "GraniteForCausalLM"),
+    "GraniteMoeForCausalLM": ("granitemoe", "GraniteMoeForCausalLM"),
+    "InternLMForCausalLM": ("llama", "LlamaForCausalLM"),
+    "InternLM2ForCausalLM": ("internlm2", "InternLM2ForCausalLM"),
+    "JAISLMHeadModel": ("jais", "JAISLMHeadModel"),
+    "JambaForCausalLM": ("jamba", "JambaForCausalLM"),
+    "LlamaForCausalLM": ("llama", "LlamaForCausalLM"),
+    # For decapoda-research/llama-*
+    "LLaMAForCausalLM": ("llama", "LlamaForCausalLM"),
+    "MambaForCausalLM": ("mamba", "MambaForCausalLM"),
+    "MistralForCausalLM": ("llama", "LlamaForCausalLM"),
+    "MixtralForCausalLM": ("mixtral", "MixtralForCausalLM"),
+    "QuantMixtralForCausalLM": ("mixtral_quant", "MixtralForCausalLM"),
+    # transformers's mpt class has lower case
+    "MptForCausalLM": ("mpt", "MPTForCausalLM"),
+    "MPTForCausalLM": ("mpt", "MPTForCausalLM"),
+    "MiniCPMForCausalLM": ("minicpm", "MiniCPMForCausalLM"),
+    "MiniCPM3ForCausalLM": ("minicpm3", "MiniCPM3ForCausalLM"),
+    "NemotronForCausalLM": ("nemotron", "NemotronForCausalLM"),
+    "OlmoForCausalLM": ("olmo", "OlmoForCausalLM"),
+    "OlmoeForCausalLM": ("olmoe", "OlmoeForCausalLM"),
+    "OPTForCausalLM": ("opt", "OPTForCausalLM"),
+    "OrionForCausalLM": ("orion", "OrionForCausalLM"),
+    "PersimmonForCausalLM": ("persimmon", "PersimmonForCausalLM"),
+    "PhiForCausalLM": ("phi", "PhiForCausalLM"),
+    "Phi3ForCausalLM": ("phi3", "Phi3ForCausalLM"),
+    "Phi3SmallForCausalLM": ("phi3_small", "Phi3SmallForCausalLM"),
+    "PhiMoEForCausalLM": ("phimoe", "PhiMoEForCausalLM"),
+    # QWenLMHeadModel supports multimodal
+    "Qwen2ForCausalLM": ("qwen2", "Qwen2ForCausalLM"),
+    "Qwen2MoeForCausalLM": ("qwen2_moe", "Qwen2MoeForCausalLM"),
+    "Qwen3ForCausalLM": ("qwen3", "Qwen3ForCausalLM"),
+    "Qwen3MoeForCausalLM": ("qwen3_moe", "Qwen3MoeForCausalLM"),
+    "Qwen3_5ForCausalLM": ("qwen3_5", "Qwen3_5ForCausalLM"),
+    "Qwen3_5MoeForCausalLM": ("qwen3_5", "Qwen3_5MoeForCausalLM"),
+    "RWForCausalLM": ("falcon", "FalconForCausalLM"),
+    "StableLMEpochForCausalLM": ("stablelm", "StablelmForCausalLM"),
+    "StableLmForCausalLM": ("stablelm", "StablelmForCausalLM"),
+    "Starcoder2ForCausalLM": ("starcoder2", "Starcoder2ForCausalLM"),
+    "SolarForCausalLM": ("solar", "SolarForCausalLM"),
+    "XverseForCausalLM": ("xverse", "XverseForCausalLM"),
+    # [Encoder-decoder]
+    "BartModel": ("bart", "BartForConditionalGeneration"),
+    "BartForConditionalGeneration": ("bart", "BartForConditionalGeneration"),
+}
+
+_EMBEDDING_MODELS = {
+    "MistralModel": ("llama_embedding", "LlamaEmbeddingModel"),
+    "Qwen2ForRewardModel": ("qwen2_rm", "Qwen2ForRewardModel"),
+    "Gemma2Model": ("gemma2_embedding", "Gemma2EmbeddingModel"),
+}
+
+_MULTIMODAL_MODELS = {
+    # [Decoder-only]
+    "Blip2ForConditionalGeneration": ("blip2", "Blip2ForConditionalGeneration"),
+    "ChameleonForConditionalGeneration": ("chameleon", "ChameleonForConditionalGeneration"),  # noqa: E501
+    "ChatGLMModel": ("chatglm", "ChatGLMForCausalLM"),
+    "ChatGLMForConditionalGeneration": ("chatglm", "ChatGLMForCausalLM"),
+    "FuyuForCausalLM": ("fuyu", "FuyuForCausalLM"),
+    "InternVLChatModel": ("internvl", "InternVLChatModel"),
+    "LlavaForConditionalGeneration": ("llava", "LlavaForConditionalGeneration"),
+    "LlavaNextForConditionalGeneration": ("llava_next", "LlavaNextForConditionalGeneration"),  # noqa: E501
+    "LlavaNextVideoForConditionalGeneration": ("llava_next_video", "LlavaNextVideoForConditionalGeneration"),  # noqa: E501
+    "LlavaOnevisionForConditionalGeneration": ("llava_onevision", "LlavaOnevisionForConditionalGeneration"),  # noqa: E501
+    "MiniCPMV": ("minicpmv", "MiniCPMV"),
+    "MolmoForCausalLM": ("molmo", "MolmoForCausalLM"),
+    "NVLM_D": ("nvlm_d", "NVLM_D_Model"),
+    "PaliGemmaForConditionalGeneration": ("paligemma", "PaliGemmaForConditionalGeneration"),  # noqa: E501
+    "Phi3VForCausalLM": ("phi3v", "Phi3VForCausalLM"),
+    "PixtralForConditionalGeneration": ("pixtral", "PixtralForConditionalGeneration"),  # noqa: E501
+    "QWenLMHeadModel": ("qwen", "QWenLMHeadModel"),
+    "Qwen2VLForConditionalGeneration": ("qwen2_vl", "Qwen2VLForConditionalGeneration"),  # noqa: E501
+    "Qwen2_5_VLForConditionalGeneration": ("qwen2_5_vl", "Qwen2_5_VLForConditionalGeneration"),  # noqa: E501
+    "UltravoxModel": ("ultravox", "UltravoxModel"),
+    # [Encoder-decoder]
+    "MllamaForConditionalGeneration": ("mllama", "MllamaForConditionalGeneration"),  # noqa: E501
+}
+
+_SPECULATIVE_DECODING_MODELS = {
+    "EAGLEModel": ("eagle", "EAGLE"),
+    "MedusaModel": ("medusa", "Medusa"),
+    "MLPSpeculatorPreTrainedModel": ("mlp_speculator", "MLPSpeculator"),
+}
+# yapf: enable
+
+_VLLM_MODELS = {
+    **_TEXT_GENERATION_MODELS,
+    **_EMBEDDING_MODELS,
+    **_MULTIMODAL_MODELS,
+    **_SPECULATIVE_DECODING_MODELS,
+}
+
+# Models not supported by ROCm.
+_ROCM_UNSUPPORTED_MODELS: List[str] = []
+
+# Models partially supported by ROCm.
+# Architecture -> Reason.
+_ROCM_SWA_REASON = ("Sliding window attention (SWA) is not yet supported in "
+                    "Triton flash attention. For half-precision SWA support, "
+                    "please use CK flash attention by setting "
+                    "`VLLM_USE_TRITON_FLASH_ATTN=0`")
+_ROCM_PARTIALLY_SUPPORTED_MODELS: Dict[str, str] = {
+    "Qwen2ForCausalLM":
+    _ROCM_SWA_REASON,
+    "MistralForCausalLM":
+    _ROCM_SWA_REASON,
+    "MixtralForCausalLM":
+    _ROCM_SWA_REASON,
+    "PaliGemmaForConditionalGeneration":
+    ("ROCm flash attention does not yet "
+     "fully support 32-bit precision on PaliGemma"),
+    "Phi3VForCausalLM":
+    ("ROCm Triton flash attention may run into compilation errors due to "
+     "excessive use of shared memory. If this happens, disable Triton FA "
+     "by setting `VLLM_USE_TRITON_FLASH_ATTN=0`")
+}
+
+
+@dataclass(frozen=True)
+class _ModelInfo:
+    is_text_generation_model: bool
+    is_embedding_model: bool
+    supports_multimodal: bool
+    supports_pp: bool
+    has_inner_state: bool
+    is_attention_free: bool
+
+    @staticmethod
+    def from_model_cls(model: Type[nn.Module]) -> "_ModelInfo":
+        return _ModelInfo(
+            is_text_generation_model=is_text_generation_model(model),
+            is_embedding_model=is_embedding_model(model),
+            supports_multimodal=supports_multimodal(model),
+            supports_pp=supports_pp(model),
+            has_inner_state=has_inner_state(model),
+            is_attention_free=is_attention_free(model),
+        )
+
+
+class _BaseRegisteredModel(ABC):
+
+    @abstractmethod
+    def inspect_model_cls(self) -> _ModelInfo:
+        raise NotImplementedError
+
+    @abstractmethod
+    def load_model_cls(self) -> Type[nn.Module]:
+        raise NotImplementedError
+
+
+@dataclass(frozen=True)
+class _RegisteredModel(_BaseRegisteredModel):
+    """
+    Represents a model that has already been imported in the main process.
+    """
+
+    interfaces: _ModelInfo
+    model_cls: Type[nn.Module]
+
+    @staticmethod
+    def from_model_cls(model_cls: Type[nn.Module]):
+        return _RegisteredModel(
+            interfaces=_ModelInfo.from_model_cls(model_cls),
+            model_cls=model_cls,
+        )
+
+    def inspect_model_cls(self) -> _ModelInfo:
+        return self.interfaces
+
+    def load_model_cls(self) -> Type[nn.Module]:
+        return self.model_cls
+
+
+@dataclass(frozen=True)
+class _LazyRegisteredModel(_BaseRegisteredModel):
+    """
+    Represents a model that has not been imported in the main process.
+    """
+    module_name: str
+    class_name: str
+
+    # Performed in another process to avoid initializing CUDA
+    def inspect_model_cls(self) -> _ModelInfo:
+        return _run_in_subprocess(
+            lambda: _ModelInfo.from_model_cls(self.load_model_cls()))
+
+    def load_model_cls(self) -> Type[nn.Module]:
+        mod = importlib.import_module(self.module_name)
+        return getattr(mod, self.class_name)
+
+
+@lru_cache(maxsize=128)
+def _try_load_model_cls(
+    model_arch: str,
+    model: _BaseRegisteredModel,
+) -> Optional[Type[nn.Module]]:
+    if is_hip():
+        if model_arch in _ROCM_UNSUPPORTED_MODELS:
+            raise ValueError(f"Model architecture '{model_arch}' is not "
+                             "supported by ROCm for now.")
+
+        if model_arch in _ROCM_PARTIALLY_SUPPORTED_MODELS:
+            msg = _ROCM_PARTIALLY_SUPPORTED_MODELS[model_arch]
+            logger.warning(
+                "Model architecture '%s' is partially "
+                "supported by ROCm: %s", model_arch, msg)
+
+    try:
+        return model.load_model_cls()
+    except Exception:
+        logger.exception("Error in loading model architecture '%s'",
+                         model_arch)
+        return None
+
+
+@lru_cache(maxsize=128)
+def _try_inspect_model_cls(
+    model_arch: str,
+    model: _BaseRegisteredModel,
+) -> Optional[_ModelInfo]:
+    try:
+        return model.inspect_model_cls()
+    except Exception:
+        logger.exception("Error in inspecting model architecture '%s'",
+                         model_arch)
+        return None
+
+
+@dataclass
+class _ModelRegistry:
+    # Keyed by model_arch
+    models: Dict[str, _BaseRegisteredModel] = field(default_factory=dict)
+
+    def get_supported_archs(self) -> List[str]:
+        return list(self.models.keys())
+
+    def register_model(
+        self,
+        model_arch: str,
+        model_cls: Union[Type[nn.Module], str],
+    ) -> None:
+        """
+        Register an external model to be used in vLLM.
+
+        :code:`model_cls` can be either:
+
+        - A :class:`torch.nn.Module` class directly referencing the model.
+        - A string in the format :code:`<module>:<class>` which can be used to
+          lazily import the model. This is useful to avoid initializing CUDA
+          when importing the model and thus the related error
+          :code:`RuntimeError: Cannot re-initialize CUDA in forked subprocess`.
+        """
+        if model_arch in self.models:
+            logger.warning(
+                "Model architecture %s is already registered, and will be "
+                "overwritten by the new model class %s.", model_arch,
+                model_cls)
+
+        if isinstance(model_cls, str):
+            split_str = model_cls.split(":")
+            if len(split_str) != 2:
+                msg = "Expected a string in the format `<module>:<class>`"
+                raise ValueError(msg)
+
+            model = _LazyRegisteredModel(*split_str)
+        else:
+            model = _RegisteredModel.from_model_cls(model_cls)
+
+        self.models[model_arch] = model
+
+    def _raise_for_unsupported(self, architectures: List[str]):
+        all_supported_archs = self.get_supported_archs()
+
+        raise ValueError(
+            f"Model architectures {architectures} are not supported for now. "
+            f"Supported architectures: {all_supported_archs}")
+
+    def _try_load_model_cls(self,
+                            model_arch: str) -> Optional[Type[nn.Module]]:
+        if model_arch not in self.models:
+            return None
+
+        return _try_load_model_cls(model_arch, self.models[model_arch])
+
+    def _try_inspect_model_cls(self, model_arch: str) -> Optional[_ModelInfo]:
+        if model_arch not in self.models:
+            return None
+
+        return _try_inspect_model_cls(model_arch, self.models[model_arch])
+
+    def _normalize_archs(
+        self,
+        architectures: Union[str, List[str]],
+    ) -> List[str]:
+        if isinstance(architectures, str):
+            architectures = [architectures]
+        if not architectures:
+            logger.warning("No model architectures are specified")
+
+        return architectures
+
+    def inspect_model_cls(
+        self,
+        architectures: Union[str, List[str]],
+    ) -> _ModelInfo:
+        architectures = self._normalize_archs(architectures)
+
+        for arch in architectures:
+            model_info = self._try_inspect_model_cls(arch)
+            if model_info is not None:
+                return model_info
+
+        return self._raise_for_unsupported(architectures)
+
+    def resolve_model_cls(
+        self,
+        architectures: Union[str, List[str]],
+    ) -> Tuple[Type[nn.Module], str]:
+        architectures = self._normalize_archs(architectures)
+
+        for arch in architectures:
+            model_cls = self._try_load_model_cls(arch)
+            if model_cls is not None:
+                return (model_cls, arch)
+
+        return self._raise_for_unsupported(architectures)
+
+    def is_text_generation_model(
+        self,
+        architectures: Union[str, List[str]],
+    ) -> bool:
+        return self.inspect_model_cls(architectures).is_text_generation_model
+
+    def is_embedding_model(
+        self,
+        architectures: Union[str, List[str]],
+    ) -> bool:
+        return self.inspect_model_cls(architectures).is_embedding_model
+
+    def is_multimodal_model(
+        self,
+        architectures: Union[str, List[str]],
+    ) -> bool:
+        return self.inspect_model_cls(architectures).supports_multimodal
+
+    def is_pp_supported_model(
+        self,
+        architectures: Union[str, List[str]],
+    ) -> bool:
+        return self.inspect_model_cls(architectures).supports_pp
+
+    def model_has_inner_state(self, architectures: Union[str,
+                                                         List[str]]) -> bool:
+        return self.inspect_model_cls(architectures).has_inner_state
+
+    def is_attention_free_model(self, architectures: Union[str,
+                                                           List[str]]) -> bool:
+        return self.inspect_model_cls(architectures).is_attention_free
+
+
+ModelRegistry = _ModelRegistry({
+    model_arch: _LazyRegisteredModel(
+        module_name=f"vllm.model_executor.models.{mod_relname}",
+        class_name=cls_name,
+    )
+    for model_arch, (mod_relname, cls_name) in _VLLM_MODELS.items()
+})
+
+_T = TypeVar("_T")
+
+
+def _run_in_subprocess(fn: Callable[[], _T]) -> _T:
+    with tempfile.NamedTemporaryFile() as output_file:
+        # `cloudpickle` allows pickling lambda functions directly
+        input_bytes = cloudpickle.dumps((fn, output_file.name))
+
+        # cannot use `sys.executable __file__` here because the script
+        # contains relative imports
+        returned = subprocess.run(
+            [sys.executable, "-m", "vllm.model_executor.models.registry"],
+            input=input_bytes,
+            capture_output=True)
+
+        # check if the subprocess is successful
+        try:
+            returned.check_returncode()
+        except Exception as e:
+            # wrap raised exception to provide more information
+            raise RuntimeError(f"Error raised in subprocess:\n"
+                               f"{returned.stderr.decode()}") from e
+
+        with open(output_file.name, "rb") as f:
+            return pickle.load(f)
+
+
+def _run() -> None:
+    # Setup plugins
+    from vllm.plugins import load_general_plugins
+    load_general_plugins()
+
+    fn, output_file = pickle.loads(sys.stdin.buffer.read())
+
+    result = fn()
+
+    with open(output_file, "wb") as f:
+        f.write(pickle.dumps(result))
+
+
+if __name__ == "__main__":
+    _run()

qwen3_6_scripts/tool_parsers_init.py

@@ -0,0 +1,12 @@
+from .abstract_tool_parser import ToolParser, ToolParserManager
+from .hermes_tool_parser import Hermes2ProToolParser
+from .internlm2_tool_parser import Internlm2ToolParser
+from .llama_tool_parser import Llama3JsonToolParser
+from .mistral_tool_parser import MistralToolParser
+from .qwen3coder_tool_parser import Qwen3CoderToolParser
+
+__all__ = [
+    "ToolParser", "ToolParserManager", "Hermes2ProToolParser",
+    "MistralToolParser", "Internlm2ToolParser", "Llama3JsonToolParser",
+    "Qwen3CoderToolParser"
+]